wdevs/pgsql-broker
5
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 3 Wiki Activity

feat(broker): migrations-based install, roles, RLS, and job dependency groups

Browse Source 
Replace the ad-hoc tables/procedures install layout with versioned,
ordered SQL migrations tracked in broker_schema_migrations. Add
optional least-privilege role provisioning (--with-roles), multi-tenant
row-level security, lease-based job claiming with stale-lease recovery,
and job dependencies -- both by job id and by fan-in job group. Add
Docker/Compose support for running the broker and its test suite.
This commit is contained in:
warkanum
2026-09-17 22:09:41 +02:00
parent 602997bcdb
commit 4c8e1066d4
53 changed files with 2911 additions and 1016 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/broker/adapter/database.go
+5
View File
@@ -19,6 +19,11 @@ type DBAdapter interface {
// Begin starts a new transaction
Begin(ctx context.Context) (DBTransaction, error)
// Conn returns a single physical connection pinned out of the pool, for
// session-scoped state (e.g. advisory locks) that must survive across
// calls. The caller owns it and must Close() it when done.
Conn(ctx context.Context) (*sql.Conn, error)
// Exec executes a query without returning rows
Exec(ctx context.Context, query string, args ...interface{}) (sql.Result, error)
pkg/broker/adapter/postgres.go
+24 -5
View File
@@ -192,8 +192,11 @@ func (p *PostgresAdapter) Listen(ctx context.Context, channel string, handler No
p.logger.Info("listening on channel", "channel", channel)
// Start notification handler in goroutine
go func() {
// Start notification handler in a supervised goroutine: it must keep
// running for the life of the process, so a panic (e.g. from a
// misbehaving handler) is logged and the loop restarted rather than
// silently dying.
SupervisedGo(p.logger, "listener-"+channel, func() {
for {
select {
case n := <-listener.Notify:
@@ -208,10 +211,10 @@ func (p *PostgresAdapter) Listen(ctx context.Context, channel string, handler No
p.logger.Info("stopping listener", "channel", channel)
return
case <-time.After(90 * time.Second):
go listener.Ping()
SafeGo(p.logger, "listener-ping-"+channel, func() { listener.Ping() })
}
}
}()
})
return nil
}
@@ -237,7 +240,7 @@ func (p *PostgresAdapter) buildConnectionString() string {
}
return fmt.Sprintf(
"host=%s port=%d user=%s password=%s dbname=%s sslmode=%s",
"host=%s port=%d user=%s password=%s dbname=%s sslmode=%s options='-c search_path=broker,public'",
p.config.Host,
p.config.Port,
p.config.User,
@@ -247,6 +250,22 @@ func (p *PostgresAdapter) buildConnectionString() string {
)
}
// Conn returns a single physical connection pinned out of the pool, for
// session-scoped operations (e.g. pg_try_advisory_lock) that must survive
// across calls and must not be silently reaped or handed to another caller
// by the pool. The caller owns its lifecycle and must Close() it.
func (p *PostgresAdapter) Conn(ctx context.Context) (*sql.Conn, error) {
p.mu.RLock()
db := p.db
p.mu.RUnlock()
if db == nil {
return nil, fmt.Errorf("database connection not established")
}
return db.Conn(ctx)
}
// postgresTransaction implements DBTransaction
type postgresTransaction struct {
tx *sql.Tx
pkg/broker/adapter/safego.go
+56
View File
@@ -0,0 +1,56 @@
package adapter
import (
"runtime/debug"
"time"
)
// RecoverAndLog recovers a panic (if any) and logs it with a stack trace.
// Call it via `defer adapter.RecoverAndLog(logger, "name")` at the top of
// any goroutine body that must never be allowed to crash the process.
func RecoverAndLog(logger Logger, name string) {
if r := recover(); r != nil {
logger.Error("recovered from panic", "component", name, "panic", r, "stack", string(debug.Stack()))
}
}
// SafeGo runs fn in a new goroutine, recovering any panic so it can never
// crash the process. Use for one-shot/fire-and-forget goroutines; a panic is
// logged (with its stack trace) and the goroutine simply ends.
func SafeGo(logger Logger, name string, fn func()) {
go func() {
defer RecoverAndLog(logger, name)
fn()
}()
}
// SupervisedGo runs fn in a new goroutine. If fn panics, the panic is logged
// and fn is restarted (after a short backoff) instead of letting the
// goroutine die permanently. Use for long-running loops (ticker routines,
// notification listeners, worker loops) that must keep running for the life
// of the process. fn must return normally, without panicking, once its own
// shutdown/context-done condition is met -- a clean return is not restarted.
func SupervisedGo(logger Logger, name string, fn func()) {
go func() {
for {
if runSupervised(logger, name, fn) {
return
}
time.Sleep(time.Second)
}
}()
}
// runSupervised runs fn once, recovering a panic if it occurs. It returns
// true if fn returned normally (no restart needed) and false if it panicked
// (caller should restart it).
func runSupervised(logger Logger, name string, fn func()) (clean bool) {
defer func() {
if r := recover(); r != nil {
logger.Error("recovered from panic, restarting", "component", name, "panic", r, "stack", string(debug.Stack()))
clean = false
}
}()
fn()
return true
}
pkg/broker/broker.go
+1
View File
@@ -100,6 +100,7 @@ func (b *Broker) stopInstances() {
wg.Add(1)
go func(inst *DatabaseInstance) {
defer wg.Done()
defer adapter.RecoverAndLog(b.logger, "stop-instance-"+inst.DatabaseName)
if err := inst.Stop(); err != nil {
b.logger.Error("failed to stop instance", "name", inst.DatabaseName, "error", err)
}
pkg/broker/config/config.go
+19
View File
@@ -29,6 +29,15 @@ type DatabaseConfig struct {
ConnMaxLifetime time.Duration `mapstructure:"conn_max_lifetime"`
ConnMaxIdleTime time.Duration `mapstructure:"conn_max_idle_time"`
QueueCount int `mapstructure:"queue_count"`
// TenantID is the RLS tenant this instance's own workers operate as
// (via broker_set_tenant) when claiming/running jobs. Defaults to
// "default" so single-tenant deployments are unaffected.
TenantID string `mapstructure:"tenant_id"`
// AutoMigrate, when true, applies any pending embedded migrations on
// connect during normal `start`. When false (default), startup fails
// fast if the schema is behind, naming the missing migrations and
// pointing at `pgsql-broker install`.
AutoMigrate bool `mapstructure:"auto_migrate"`
}
// BrokerConfig holds broker-specific settings
@@ -40,6 +49,11 @@ type BrokerConfig struct {
WorkerIdleTimeoutSec int `mapstructure:"worker_idle_timeout_sec"`
NotifyRetrySeconds time.Duration `mapstructure:"notify_retry_seconds"`
EnableDebug bool `mapstructure:"enable_debug"`
// LeaseSeconds is how long a claimed job's lease is valid for before
// broker_recover_stale_jobs considers it abandoned.
LeaseSeconds int `mapstructure:"lease_seconds"`
// StaleJobRecoverySec is the interval between broker_recover_stale_jobs sweeps.
StaleJobRecoverySec int `mapstructure:"stale_job_recovery_sec"`
}
// LoggingConfig holds logging settings
@@ -103,6 +117,8 @@ func setDefaults(v *viper.Viper) {
v.SetDefault("broker.worker_idle_timeout_sec", 10)
v.SetDefault("broker.notify_retry_seconds", 30*time.Second)
v.SetDefault("broker.enable_debug", false)
v.SetDefault("broker.lease_seconds", 60)
v.SetDefault("broker.stale_job_recovery_sec", 30)
// Logging defaults
v.SetDefault("logging.level", "info")
@@ -160,6 +176,9 @@ func applyDatabaseDefaults(config *Config) {
if db.QueueCount == 0 {
db.QueueCount = 4
}
if db.TenantID == "" {
db.TenantID = "default"
}
}
}
pkg/broker/database_instance.go
+169 -72
View File
@@ -2,15 +2,17 @@ package broker
import (
"context"
"database/sql" // Import sql package
"database/sql"
"encoding/json"
"fmt"
"os"
"strings"
"sync"
"time"
"git.warky.dev/wdevs/pgsql-broker/pkg/broker/adapter"
"git.warky.dev/wdevs/pgsql-broker/pkg/broker/config"
"git.warky.dev/wdevs/pgsql-broker/pkg/broker/install"
"git.warky.dev/wdevs/pgsql-broker/pkg/broker/models"
"git.warky.dev/wdevs/pgsql-broker/pkg/broker/queue"
)
@@ -34,6 +36,12 @@ type DatabaseInstance struct {
shutdownMu sync.RWMutex
jobsHandled int64
startTime time.Time
// sessionConn holds the pg_try_advisory_lock acquired by
// registerInstance. The lock is scoped to this one physical connection,
// so it must be kept open (never returned to the pool) for the life of
// the process and explicitly unlocked on Stop().
sessionConn *sql.Conn
}
// NewDatabaseInstance creates a new database instance
@@ -70,6 +78,11 @@ func (i *DatabaseInstance) Start() error {
return fmt.Errorf("failed to connect to database: %w", err)
}
// Ensure the schema is up to date before touching any broker objects.
if err := i.ensureSchema(); err != nil {
return err
}
// Register instance in database
if err := i.registerInstance(); err != nil {
return fmt.Errorf("failed to register instance: %w", err)
@@ -87,13 +100,48 @@ func (i *DatabaseInstance) Start() error {
return fmt.Errorf("failed to start listener: %w", err)
}
// Start ping routine
go i.pingRoutine()
// Start ping routine (auto-restarted on panic; must run for the life of
// the process)
adapter.SupervisedGo(i.logger, "ping-routine", i.pingRoutine)
// Start stale/expired-lease job recovery routine (auto-restarted on
// panic; must run for the life of the process)
adapter.SupervisedGo(i.logger, "stale-job-recovery-routine", i.staleJobRecoveryRoutine)
i.logger.Info("database instance started successfully")
return nil
}
// ensureSchema checks the embedded migration set against the database and,
// depending on dbConfig.AutoMigrate, either applies pending migrations or
// fails startup fast rather than running against a stale/missing schema.
func (i *DatabaseInstance) ensureSchema() error {
installer := install.New(i.db, i.logger)
pending, err := installer.PendingMigrations(i.ctx)
if err != nil {
return fmt.Errorf("failed to check schema migrations: %w", err)
}
if len(pending) == 0 {
return nil
}
if !i.dbConfig.AutoMigrate {
return fmt.Errorf(
"schema is missing or behind: %d migration(s) not applied (%s); either run `pgsql-broker install` or set databases[].auto_migrate: true",
len(pending), strings.Join(pending, ", "),
)
}
i.logger.Info("auto-migrating database schema", "pending", pending)
if err := installer.ApplyMigrations(i.ctx); err != nil {
return fmt.Errorf("auto-migration failed: %w", err)
}
return nil
}
// Stop gracefully stops the database instance
func (i *DatabaseInstance) Stop() error {
i.shutdownMu.Lock()
@@ -116,10 +164,26 @@ func (i *DatabaseInstance) Stop() error {
}
i.queuesMu.Unlock()
// Update instance status in database
if err := i.shutdownInstance(); err != nil {
// Update instance status in database. i.ctx may already be canceled by
// the parent broker's Stop(), so use a fresh short-lived context here.
shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
if err := i.shutdownInstance(shutdownCtx); err != nil {
i.logger.Error("failed to shutdown instance in database", "error", err)
}
cancel()
// Release the advisory lock and close the pinned session connection.
if i.sessionConn != nil {
if _, err := i.sessionConn.ExecContext(context.Background(),
"SELECT pg_advisory_unlock(hashtextextended($1, 0))", "broker:"+i.Name,
); err != nil {
i.logger.Error("failed to release advisory lock", "error", err)
}
if err := i.sessionConn.Close(); err != nil {
i.logger.Error("failed to close session connection", "error", err)
}
i.sessionConn = nil
}
// Close database connection
if err := i.db.Close(); err != nil {
@@ -130,72 +194,48 @@ func (i *DatabaseInstance) Stop() error {
return nil
}
// registerInstance registers the instance in the database
// registerInstance registers the instance in the database. The advisory
// lock taken by broker_register_instance is session-scoped, so this runs on
// a connection pinned out of the pool (i.sessionConn) that is kept open for
// the life of the process rather than returned after this call.
func (i *DatabaseInstance) registerInstance() error {
conn, err := i.db.Conn(i.ctx)
if err != nil {
return fmt.Errorf("failed to acquire session connection: %w", err)
}
var retval int
var errmsg string
var nullableInstanceID sql.NullInt64 // Change to nullable type
var nullableInstanceID sql.NullInt64
i.logger.Debug("registering instance", "name", i.Name, "hostname", i.Hostname, "pid", i.PID, "version", i.Version, "queue_count", i.dbConfig.QueueCount)
err := i.db.QueryRow(i.ctx,
"SELECT p_retval, p_errmsg, p_instance_id FROM broker_register_instance($1, $2, $3, $4, $5)",
err = conn.QueryRowContext(i.ctx,
"SELECT p_retval, p_errmsg, p_instance_id FROM broker.broker_register_instance($1, $2, $3, $4, $5)",
i.Name, i.Hostname, i.PID, i.Version, i.dbConfig.QueueCount,
).Scan(&retval, &errmsg, &nullableInstanceID)
if err != nil {
conn.Close()
i.logger.Error("query error during instance registration", "error", err)
return fmt.Errorf("query error: %w", err)
}
if retval == 3 {
i.logger.Warn("another broker instance is already active, attempting to retrieve ID", "error", errmsg)
// Try to retrieve the ID of the active instance
var activeID int64
err := i.db.QueryRow(i.ctx,
"SELECT id_broker_queueinstance FROM broker_queueinstance WHERE name = $1 AND hostname = $2 AND status = 'active' ORDER BY started_at DESC LIMIT 1",
i.Name, i.Hostname,
).Scan(&activeID)
if err != nil {
i.logger.Error("failed to retrieve ID of active instance", "error", err)
return fmt.Errorf("failed to retrieve ID of active instance: %w", err)
}
i.ID = activeID
i.logger.Info("retrieved active instance ID", "id", i.ID)
return nil
} else if retval > 0 {
if retval > 0 {
conn.Close()
i.logger.Error("broker_register_instance error", "retval", retval, "errmsg", errmsg)
return fmt.Errorf("broker_register_instance error: %s", errmsg)
}
// If successfully registered, nullableInstanceID.Valid will be true
if nullableInstanceID.Valid {
i.ID = nullableInstanceID.Int64
i.logger.Info("registered new instance", "id", i.ID)
// Debug logging: Retrieve all entries from broker_queueinstance
rows, err := i.db.Query(i.ctx, "SELECT id_broker_queueinstance, name, hostname, status FROM broker_queueinstance")
if err != nil {
i.logger.Error("debug query failed", "error", err)
} else {
defer rows.Close()
for rows.Next() {
var id int64
var name, hostname, status string
if err := rows.Scan(&id, &name, &hostname, &status); err != nil {
i.logger.Error("debug scan failed", "error", err)
break
}
i.logger.Debug("broker_queueinstance entry", "id", id, "name", name, "hostname", hostname, "status", status)
}
}
} else {
// This case should ideally not happen if retval is 0 (success)
// but if it does, it means p_instance_id was NULL despite success.
// This would be an unexpected scenario.
if !nullableInstanceID.Valid {
conn.Close()
i.logger.Error("broker_register_instance returned success but no instance ID", "retval", retval, "errmsg", errmsg)
return fmt.Errorf("broker_register_instance returned success but no instance ID")
}
i.ID = nullableInstanceID.Int64
i.sessionConn = conn
i.logger.Info("registered new instance", "id", i.ID)
return nil
}
@@ -204,6 +244,11 @@ func (i *DatabaseInstance) startQueues() error {
i.queuesMu.Lock()
defer i.queuesMu.Unlock()
leaseSeconds := i.config.Broker.LeaseSeconds
if leaseSeconds <= 0 {
leaseSeconds = 60
}
for queueNum := 1; queueNum <= i.dbConfig.QueueCount; queueNum++ {
queueCfg := queue.Config{
Number: queueNum,
@@ -214,6 +259,8 @@ func (i *DatabaseInstance) startQueues() error {
BufferSize: i.config.Broker.QueueBufferSize,
TimerSeconds: i.config.Broker.QueueTimerSec,
FetchSize: i.config.Broker.FetchQueryQueSize,
TenantID: i.dbConfig.TenantID,
LeaseSeconds: leaseSeconds,
}
q := queue.New(queueCfg)
@@ -241,42 +288,38 @@ func (i *DatabaseInstance) startListener() error {
return nil
}
// handleNotification processes incoming job notifications
// handleNotification processes incoming wake-up notifications. The payload
// only carries the queue number (plus a job id kept for logging) --
// NOTIFY is wake-only, never a hand-off of a job to execute directly, so the
// woken worker always re-claims via broker_get.
func (i *DatabaseInstance) handleNotification(n *adapter.Notification) {
defer adapter.RecoverAndLog(i.logger, "handle-notification")
if i.config.Broker.EnableDebug {
i.logger.Debug("received notification", "channel", n.Channel, "payload", n.Payload)
}
var job models.Job
if err := json.Unmarshal([]byte(n.Payload), &job); err != nil {
var wake models.WakeNotification
if err := json.Unmarshal([]byte(n.Payload), &wake); err != nil {
i.logger.Error("failed to unmarshal notification", "error", err, "payload", n.Payload)
return
}
if job.ID <= 0 {
i.logger.Warn("notification missing job ID", "payload", n.Payload)
if wake.Queue <= 0 {
i.logger.Warn("notification missing queue number", "payload", n.Payload)
return
}
if job.JobQueue <= 0 {
i.logger.Warn("notification missing queue number", "job_id", job.ID)
return
}
// Get the queue
i.queuesMu.RLock()
q, exists := i.queues[job.JobQueue]
q, exists := i.queues[wake.Queue]
i.queuesMu.RUnlock()
if !exists {
i.logger.Warn("queue not found for job", "job_id", job.ID, "queue", job.JobQueue)
i.logger.Warn("queue not found for notification", "queue", wake.Queue, "job_id", wake.JobID)
return
}
// Add job to queue
if err := q.AddJob(job); err != nil {
i.logger.Error("failed to add job to queue", "job_id", job.ID, "queue", job.JobQueue, "error", err)
}
q.Wake()
}
// pingRoutine periodically updates the instance status in the database
@@ -304,13 +347,67 @@ func (i *DatabaseInstance) pingRoutine() {
}
}
// staleJobRecoveryRoutine periodically requeues (or dead-letters) jobs whose
// lease has expired while still running.
func (i *DatabaseInstance) staleJobRecoveryRoutine() {
interval := time.Duration(i.config.Broker.StaleJobRecoverySec) * time.Second
if interval <= 0 {
interval = 30 * time.Second
}
ticker := time.NewTicker(interval)
defer ticker.Stop()
for {
select {
case <-ticker.C:
i.shutdownMu.RLock()
if i.shutdown {
i.shutdownMu.RUnlock()
return
}
i.shutdownMu.RUnlock()
if err := i.recoverStaleJobs(); err != nil {
i.logger.Error("stale job recovery failed", "error", err)
}
case <-i.ctx.Done():
return
}
}
}
// recoverStaleJobs invokes broker_recover_stale_jobs.
func (i *DatabaseInstance) recoverStaleJobs() error {
var retval int
var errmsg string
var recoveredCount int
err := i.db.QueryRow(i.ctx, "SELECT p_retval, p_errmsg, p_recovered_count FROM broker.broker_recover_stale_jobs()").
Scan(&retval, &errmsg, &recoveredCount)
if err != nil {
return fmt.Errorf("query error: %w", err)
}
if retval > 0 {
return fmt.Errorf("broker_recover_stale_jobs error: %s", errmsg)
}
if recoveredCount > 0 {
i.logger.Info("recovered stale jobs", "count", recoveredCount)
}
return nil
}
// ping updates the instance ping timestamp
func (i *DatabaseInstance) ping() error {
var retval int
var errmsg string
err := i.db.QueryRow(i.ctx,
"SELECT p_retval, p_errmsg FROM broker_ping_instance($1, $2)",
"SELECT p_retval, p_errmsg FROM broker.broker_ping_instance($1, $2)",
i.ID, i.jobsHandled,
).Scan(&retval, &errmsg)
@@ -326,12 +423,12 @@ func (i *DatabaseInstance) ping() error {
}
// shutdownInstance marks the instance as shutdown in the database
func (i *DatabaseInstance) shutdownInstance() error {
func (i *DatabaseInstance) shutdownInstance(ctx context.Context) error {
var retval int
var errmsg string
err := i.db.QueryRow(i.ctx,
"SELECT p_retval, p_errmsg FROM broker_shutdown_instance($1)",
err := i.db.QueryRow(ctx,
"SELECT p_retval, p_errmsg FROM broker.broker_shutdown_instance($1)",
i.ID,
).Scan(&retval, &errmsg)
@@ -370,4 +467,4 @@ func (i *DatabaseInstance) GetStats() map[string]interface{} {
stats["queues"] = queueStats
return stats
}
}
pkg/broker/install/install.go
+377 -147
View File
@@ -4,17 +4,47 @@ import (
"context"
"embed"
"fmt"
"io/fs"
"regexp"
"sort"
"strconv"
"strings"
"github.com/lib/pq"
"git.warky.dev/wdevs/pgsql-broker/pkg/broker/adapter"
)
//go:embed all:sql
var sqlFS embed.FS
//go:embed all:sql/migrations
var migrationsFS embed.FS
// Installer handles database schema installation
//go:embed all:sql/roles
var rolesFS embed.FS
const migrationsDir = "sql/migrations"
const rolesDir = "sql/roles"
// migrationsTableSQL creates the version-tracking table itself. It is applied
// unconditionally (idempotently) before any numbered migration file, and is
// not itself a numbered migration.
const migrationsTableSQL = `
CREATE SCHEMA IF NOT EXISTS broker;
CREATE TABLE IF NOT EXISTS broker.broker_schema_migrations (
version BIGINT PRIMARY KEY,
name TEXT NOT NULL,
applied_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);
`
var migrationFileRe = regexp.MustCompile(`^(\d+)_(.+)\.sql$`)
// migrationFile describes one embedded migration.
type migrationFile struct {
version int64
name string
path string
}
// Installer handles database schema installation via versioned migrations.
type Installer struct {
db adapter.DBAdapter
logger adapter.Logger
@@ -28,217 +58,417 @@ func New(db adapter.DBAdapter, logger adapter.Logger) *Installer {
}
}
// InstallSchema installs the complete database schema
func (i *Installer) InstallSchema(ctx context.Context) error {
i.logger.Info("starting schema installation")
// Install tables first
if err := i.installTables(ctx); err != nil {
return fmt.Errorf("failed to install tables: %w", err)
// loadMigrations reads and sorts every embedded migration file by numeric prefix.
func loadMigrations() ([]migrationFile, error) {
entries, err := migrationsFS.ReadDir(migrationsDir)
if err != nil {
return nil, fmt.Errorf("failed to read migrations directory: %w", err)
}
// Then install procedures
if err := i.installProcedures(ctx); err != nil {
return fmt.Errorf("failed to install procedures: %w", err)
var migrations []migrationFile
for _, e := range entries {
if e.IsDir() {
continue
}
m := migrationFileRe.FindStringSubmatch(e.Name())
if m == nil {
continue
}
version, err := strconv.ParseInt(m[1], 10, 64)
if err != nil {
return nil, fmt.Errorf("invalid migration filename %s: %w", e.Name(), err)
}
migrations = append(migrations, migrationFile{
version: version,
name: m[2],
path: migrationsDir + "/" + e.Name(),
})
}
i.logger.Info("schema installation completed successfully")
sort.Slice(migrations, func(i, j int) bool { return migrations[i].version < migrations[j].version })
return migrations, nil
}
// ensureMigrationsTable creates the broker schema and the migrations
// tracking table if they don't already exist. This is DDL and requires
// CREATE privilege on the database -- only ApplyMigrations (run by an
// admin-privileged connection, e.g. `pgsql-broker install`) calls it.
func (i *Installer) ensureMigrationsTable(ctx context.Context) error {
if _, err := i.db.Exec(ctx, migrationsTableSQL); err != nil {
return fmt.Errorf("failed to ensure migrations table: %w", err)
}
return nil
}
// installTables installs all table definitions
func (i *Installer) installTables(ctx context.Context) error {
i.logger.Info("installing tables")
files, err := sqlFS.ReadDir("sql/tables")
// migrationsTableExists reports whether the migrations tracking table is
// present, without creating it -- a read-only check safe to run with a
// least-privilege runtime role (e.g. broker_runtime) that has no CREATE
// privilege on the database.
func (i *Installer) migrationsTableExists(ctx context.Context) (bool, error) {
var exists bool
err := i.db.QueryRow(ctx, "SELECT to_regclass('broker.broker_schema_migrations') IS NOT NULL").Scan(&exists)
if err != nil {
return fmt.Errorf("failed to read tables directory: %w", err)
return false, fmt.Errorf("failed to check migrations table: %w", err)
}
return exists, nil
}
// appliedVersions returns the set of migration versions already recorded.
func (i *Installer) appliedVersions(ctx context.Context) (map[int64]bool, error) {
rows, err := i.db.Query(ctx, "SELECT version FROM broker.broker_schema_migrations")
if err != nil {
return nil, fmt.Errorf("failed to query applied migrations: %w", err)
}
defer rows.Close()
applied := make(map[int64]bool)
for rows.Next() {
var v int64
if err := rows.Scan(&v); err != nil {
return nil, fmt.Errorf("failed to scan migration version: %w", err)
}
applied[v] = true
}
return applied, rows.Err()
}
// PendingMigrations returns the names of embedded migrations that have not
// yet been applied to the database, without applying them or creating the
// migrations table -- safe to call with a least-privilege runtime role.
func (i *Installer) PendingMigrations(ctx context.Context) ([]string, error) {
migrations, err := loadMigrations()
if err != nil {
return nil, err
}
// Filter and sort SQL files
sqlFiles := filterAndSortSQLFiles(files)
exists, err := i.migrationsTableExists(ctx)
if err != nil {
return nil, err
}
if !exists {
pending := make([]string, len(migrations))
for idx, m := range migrations {
pending[idx] = fmt.Sprintf("%04d_%s", m.version, m.name)
}
return pending, nil
}
for _, file := range sqlFiles {
// Skip install script
if file == "00_install.sql" {
applied, err := i.appliedVersions(ctx)
if err != nil {
return nil, err
}
var pending []string
for _, m := range migrations {
if !applied[m.version] {
pending = append(pending, fmt.Sprintf("%04d_%s", m.version, m.name))
}
}
return pending, nil
}
// ApplyMigrations applies every embedded migration that has not yet been
// recorded in broker.broker_schema_migrations, each inside its own transaction.
func (i *Installer) ApplyMigrations(ctx context.Context) error {
i.logger.Info("applying migrations")
if err := i.ensureMigrationsTable(ctx); err != nil {
return err
}
migrations, err := loadMigrations()
if err != nil {
return err
}
applied, err := i.appliedVersions(ctx)
if err != nil {
return err
}
appliedCount := 0
for _, m := range migrations {
if applied[m.version] {
continue
}
i.logger.Info("executing table script", "file", file)
content, err := sqlFS.ReadFile("sql/tables/" + file)
content, err := migrationsFS.ReadFile(m.path)
if err != nil {
return fmt.Errorf("failed to read file %s: %w", file, err)
return fmt.Errorf("failed to read migration %s: %w", m.path, err)
}
if err := i.executeSQL(ctx, string(content)); err != nil {
return fmt.Errorf("failed to execute %s: %w", file, err)
i.logger.Info("applying migration", "version", m.version, "name", m.name)
tx, err := i.db.Begin(ctx)
if err != nil {
return fmt.Errorf("failed to begin transaction for migration %s: %w", m.name, err)
}
if err := execStatements(ctx, tx, string(content)); err != nil {
tx.Rollback()
return fmt.Errorf("failed to apply migration %s: %w", m.name, err)
}
if _, err := tx.Exec(ctx,
"INSERT INTO broker.broker_schema_migrations (version, name) VALUES ($1, $2)",
m.version, m.name,
); err != nil {
tx.Rollback()
return fmt.Errorf("failed to record migration %s: %w", m.name, err)
}
if err := tx.Commit(); err != nil {
return fmt.Errorf("failed to commit migration %s: %w", m.name, err)
}
appliedCount++
}
i.logger.Info("tables installed successfully")
if appliedCount == 0 {
i.logger.Info("no pending migrations")
} else {
i.logger.Info("migrations applied successfully", "count", appliedCount)
}
return nil
}
// installProcedures installs all stored procedures
func (i *Installer) installProcedures(ctx context.Context) error {
i.logger.Info("installing procedures")
// RolePasswords holds the login passwords for the reference broker_admin,
// broker_runtime, and broker_enqueue roles created by InstallRoles. All
// three are required -- there is no placeholder/default fallback, since
// these roles carry real database privileges.
type RolePasswords struct {
AdminPassword string
RuntimePassword string
EnqueuePassword string
}
files, err := sqlFS.ReadDir("sql/procedures")
// InstallRoles applies the embedded role/grant scripts (sql/roles), which
// create (or, if already present, rotate the password of) broker_admin,
// broker_runtime, and broker_enqueue, then grant them the appropriate
// schema/table/function privileges. The caller must connect as a superuser
// or a role with CREATEROLE -- this is intentionally separate from the
// migration-running connection.
func (i *Installer) InstallRoles(ctx context.Context, passwords RolePasswords) error {
if passwords.AdminPassword == "" || passwords.RuntimePassword == "" || passwords.EnqueuePassword == "" {
return fmt.Errorf("all three role passwords (admin, runtime, enqueue) are required")
}
entries, err := rolesFS.ReadDir(rolesDir)
if err != nil {
return fmt.Errorf("failed to read procedures directory: %w", err)
return fmt.Errorf("failed to read roles directory: %w", err)
}
// Filter and sort SQL files
sqlFiles := filterAndSortSQLFiles(files)
for _, file := range sqlFiles {
// Skip install script
if file == "00_install.sql" {
continue
var names []string
for _, e := range entries {
if !e.IsDir() {
names = append(names, e.Name())
}
}
sort.Strings(names)
i.logger.Info("executing procedure script", "file", file)
replacer := strings.NewReplacer(
"__BROKER_ADMIN_PASSWORD__", pq.QuoteLiteral(passwords.AdminPassword),
"__BROKER_RUNTIME_PASSWORD__", pq.QuoteLiteral(passwords.RuntimePassword),
"__BROKER_ENQUEUE_PASSWORD__", pq.QuoteLiteral(passwords.EnqueuePassword),
)
content, err := sqlFS.ReadFile("sql/procedures/" + file)
for _, name := range names {
content, err := rolesFS.ReadFile(rolesDir + "/" + name)
if err != nil {
return fmt.Errorf("failed to read file %s: %w", file, err)
return fmt.Errorf("failed to read roles script %s: %w", name, err)
}
if err := i.executeSQL(ctx, string(content)); err != nil {
return fmt.Errorf("failed to execute %s: %w", file, err)
i.logger.Info("applying roles script", "name", name)
tx, err := i.db.Begin(ctx)
if err != nil {
return fmt.Errorf("failed to begin transaction for roles script %s: %w", name, err)
}
// pq.QuoteLiteral already produces a safely quoted SQL string
// literal (doubling embedded quotes, or switching to E'...' escape
// syntax if the password contains a backslash), so this is a plain
// textual substitution, not string concatenation of untrusted input
// into SQL syntax.
rendered := replacer.Replace(string(content))
if err := execStatements(ctx, tx, rendered); err != nil {
tx.Rollback()
return fmt.Errorf("failed to apply roles script %s: %w", name, err)
}
if err := tx.Commit(); err != nil {
return fmt.Errorf("failed to commit roles script %s: %w", name, err)
}
}
i.logger.Info("procedures installed successfully")
i.logger.Info("roles installed successfully")
return nil
}
// executeSQL executes SQL statements
func (i *Installer) executeSQL(ctx context.Context, sql string) error {
// Remove comments and split by statement
statements := splitSQLStatements(sql)
// execStatements runs every statement in sql within tx.
func execStatements(ctx context.Context, tx adapter.DBTransaction, sqlText string) error {
statements := splitSQLStatements(sqlText)
for _, stmt := range statements {
stmt = strings.TrimSpace(stmt)
if stmt == "" {
if stmt == "" || strings.HasPrefix(stmt, "\\") {
continue
}
// Skip psql-specific commands
if strings.HasPrefix(stmt, "\\") {
continue
}
if _, err := i.db.Exec(ctx, stmt); err != nil {
if _, err := tx.Exec(ctx, stmt); err != nil {
return fmt.Errorf("failed to execute statement: %w\nStatement: %s", err, stmt)
}
}
return nil
}
// filterAndSortSQLFiles filters and sorts SQL files
func filterAndSortSQLFiles(files []fs.DirEntry) []string {
var sqlFiles []string
for _, file := range files {
if !file.IsDir() && strings.HasSuffix(file.Name(), ".sql") {
sqlFiles = append(sqlFiles, file.Name())
}
}
sort.Strings(sqlFiles)
return sqlFiles
}
// splitSQLStatements splits SQL into individual statements
func splitSQLStatements(sql string) []string {
// Simple split by semicolon
// This doesn't handle all edge cases (strings with semicolons, dollar-quoted strings, etc.)
// but works for our use case
statements := strings.Split(sql, ";")
// splitSQLStatements splits SQL into individual statements, keeping
// $$-quoted function bodies intact.
// splitSQLStatements splits a SQL script into individual statements on
// top-level semicolons, ignoring semicolons that appear inside single-quoted
// strings ('...', with '' as an escaped quote), double-quoted identifiers,
// line comments (--), and dollar-quoted bodies ($$...$$ or $tag$...$tag$).
func splitSQLStatements(sqlText string) []string {
var result []string
var buffer string
var buffer strings.Builder
for _, stmt := range statements {
stmt = strings.TrimSpace(stmt)
if stmt == "" {
runes := []rune(sqlText)
n := len(runes)
i := 0
for i < n {
c := runes[i]
switch {
case c == '-' && i+1 < n && runes[i+1] == '-':
// Line comment: copy through end of line.
for i < n && runes[i] != '\n' {
buffer.WriteRune(runes[i])
i++
}
continue
}
buffer += stmt + ";"
case c == '\'':
buffer.WriteRune(c)
i++
for i < n {
buffer.WriteRune(runes[i])
if runes[i] == '\'' {
if i+1 < n && runes[i+1] == '\'' {
buffer.WriteRune(runes[i+1])
i += 2
continue
}
i++
break
}
i++
}
continue
// Check if we're inside a function definition ($$)
dollarCount := strings.Count(buffer, "$$")
if dollarCount%2 == 0 {
// Even number of $$ means we're outside function definitions
result = append(result, buffer)
buffer = ""
} else {
// Odd number means we're inside a function, keep accumulating
buffer += " "
case c == '"':
buffer.WriteRune(c)
i++
for i < n {
buffer.WriteRune(runes[i])
if runes[i] == '"' {
i++
break
}
i++
}
continue
case c == '$':
if tag, ok := matchDollarTag(runes, i); ok {
closer := tag
buffer.WriteString(closer)
i += len(closer)
end := indexOfRunes(runes, i, closer)
if end == -1 {
buffer.WriteString(string(runes[i:]))
i = n
} else {
buffer.WriteString(string(runes[i:end]))
buffer.WriteString(closer)
i = end + len(closer)
}
continue
}
buffer.WriteRune(c)
i++
case c == ';':
stmt := strings.TrimSpace(buffer.String())
if stmt != "" {
result = append(result, stmt+";")
}
buffer.Reset()
i++
default:
buffer.WriteRune(c)
i++
}
}
// Add any remaining buffered content
if buffer != "" {
result = append(result, buffer)
if stmt := strings.TrimSpace(buffer.String()); stmt != "" {
result = append(result, stmt)
}
return result
}
// VerifyInstallation checks if the schema is properly installed
// matchDollarTag checks whether runes[pos:] begins a dollar-quote tag
// ($$ or $tag$) and returns that tag if so.
func matchDollarTag(runes []rune, pos int) (string, bool) {
if runes[pos] != '$' {
return "", false
}
j := pos + 1
for j < len(runes) && (runes[j] == '_' || isAlnum(runes[j])) {
j++
}
if j < len(runes) && runes[j] == '$' {
return string(runes[pos : j+1]), true
}
return "", false
}
func isAlnum(r rune) bool {
return (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') || (r >= '0' && r <= '9')
}
// indexOfRunes returns the index of the first occurrence of sub in
// runes[from:], or -1 if not found.
func indexOfRunes(runes []rune, from int, sub string) int {
subRunes := []rune(sub)
for i := from; i+len(subRunes) <= len(runes); i++ {
match := true
for j, r := range subRunes {
if runes[i+j] != r {
match = false
break
}
}
if match {
return i
}
}
return -1
}
// VerifyInstallation checks that every embedded migration has been applied.
func (i *Installer) VerifyInstallation(ctx context.Context) error {
i.logger.Info("verifying installation")
tables := []string{"broker_queueinstance", "broker_jobs", "broker_schedule"}
procedures := []string{
"broker_get",
"broker_run",
"broker_set",
"broker_add_job",
"broker_register_instance",
"broker_ping_instance",
"broker_shutdown_instance",
pending, err := i.PendingMigrations(ctx)
if err != nil {
return fmt.Errorf("failed to check pending migrations: %w", err)
}
// Check tables
for _, table := range tables {
var exists bool
err := i.db.QueryRow(ctx,
"SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_name = $1)",
table,
).Scan(&exists)
if err != nil {
return fmt.Errorf("failed to check table %s: %w", table, err)
}
if !exists {
return fmt.Errorf("table %s does not exist", table)
}
i.logger.Info("table verified", "table", table)
}
// Check procedures
for _, proc := range procedures {
var exists bool
err := i.db.QueryRow(ctx,
"SELECT EXISTS (SELECT FROM pg_proc WHERE proname = $1)",
proc,
).Scan(&exists)
if err != nil {
return fmt.Errorf("failed to check procedure %s: %w", proc, err)
}
if !exists {
return fmt.Errorf("procedure %s does not exist", proc)
}
i.logger.Info("procedure verified", "procedure", proc)
if len(pending) > 0 {
return fmt.Errorf("schema is behind: %d migration(s) not applied: %s", len(pending), strings.Join(pending, ", "))
}
i.logger.Info("installation verified successfully")
pkg/broker/install/sql/migrations/0001_schema.sql
+6
View File
@@ -0,0 +1,6 @@
-- Dedicated schema for all broker objects.
CREATE SCHEMA IF NOT EXISTS broker;
REVOKE ALL ON SCHEMA broker FROM PUBLIC;
-- gen_random_uuid() for lease tokens.
CREATE EXTENSION IF NOT EXISTS pgcrypto;
pkg/broker/install/sql/migrations/0002_broker_queueinstance.sql
+25
View File
@@ -0,0 +1,25 @@
-- broker.broker_queueinstance
-- Tracks active and historical broker queue instances.
CREATE TABLE IF NOT EXISTS broker.broker_queueinstance (
id_broker_queueinstance BIGSERIAL PRIMARY KEY,
name VARCHAR(255) NOT NULL,
hostname VARCHAR(255) NOT NULL,
pid INTEGER NOT NULL,
version VARCHAR(50) NOT NULL,
status VARCHAR(20) NOT NULL DEFAULT 'active',
started_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
last_ping_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
shutdown_at TIMESTAMP WITH TIME ZONE,
queue_count INTEGER NOT NULL DEFAULT 0,
jobs_handled BIGINT NOT NULL DEFAULT 0,
CONSTRAINT broker_queueinstance_status_check CHECK (status IN ('active', 'inactive', 'shutdown'))
);
CREATE INDEX IF NOT EXISTS idx_broker_queueinstance_status ON broker.broker_queueinstance(status);
CREATE INDEX IF NOT EXISTS idx_broker_queueinstance_hostname ON broker.broker_queueinstance(hostname);
CREATE INDEX IF NOT EXISTS idx_broker_queueinstance_last_ping ON broker.broker_queueinstance(last_ping_at);
COMMENT ON TABLE broker.broker_queueinstance IS 'Tracks broker queue instances (active and historical). Single-active-instance-per-name is enforced via a pg_try_advisory_lock in broker_register_instance, not by this status column, which is observational only.';
COMMENT ON COLUMN broker.broker_queueinstance.status IS 'Observational status: active, inactive, or shutdown. Ownership is enforced via advisory lock, not by reading this column.';
pkg/broker/install/sql/migrations/0003_broker_schedule.sql
+41
View File
@@ -0,0 +1,41 @@
-- broker.broker_schedule
-- Stores scheduled jobs (cron-like functionality).
CREATE TABLE IF NOT EXISTS broker.broker_schedule (
id_broker_schedule BIGSERIAL PRIMARY KEY,
name VARCHAR(255) NOT NULL UNIQUE,
cron_expr VARCHAR(100) NOT NULL,
enabled BOOLEAN NOT NULL DEFAULT true,
job_name VARCHAR(255) NOT NULL,
job_priority INTEGER NOT NULL DEFAULT 0,
job_queue INTEGER NOT NULL DEFAULT 1,
job_language VARCHAR(50) NOT NULL DEFAULT 'sql',
execute_str TEXT NOT NULL,
run_as VARCHAR(100),
created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
last_run_at TIMESTAMP WITH TIME ZONE,
next_run_at TIMESTAMP WITH TIME ZONE,
CONSTRAINT broker_schedule_job_queue_check CHECK (job_queue > 0)
);
CREATE INDEX IF NOT EXISTS idx_broker_schedule_enabled ON broker.broker_schedule(enabled);
CREATE INDEX IF NOT EXISTS idx_broker_schedule_next_run ON broker.broker_schedule(next_run_at) WHERE enabled = true;
CREATE INDEX IF NOT EXISTS idx_broker_schedule_name ON broker.broker_schedule(name);
COMMENT ON TABLE broker.broker_schedule IS 'Scheduled jobs (cron-like functionality)';
CREATE OR REPLACE FUNCTION broker.tf_broker_schedule_update_timestamp()
RETURNS TRIGGER AS $$
BEGIN
NEW.updated_at = NOW();
RETURN NEW;
END;
$$ LANGUAGE plpgsql;
DROP TRIGGER IF EXISTS t_broker_schedule_updated_at ON broker.broker_schedule;
CREATE TRIGGER t_broker_schedule_updated_at
BEFORE UPDATE ON broker.broker_schedule
FOR EACH ROW
EXECUTE FUNCTION broker.tf_broker_schedule_update_timestamp();
pkg/broker/install/sql/migrations/0004_broker_jobs.sql
+93
View File
@@ -0,0 +1,93 @@
-- broker.broker_jobs
-- Job queue for broker execution.
-- tenant_id / RLS: rows are only visible/writable when tenant_id matches
-- current_setting('broker.tenant_id', true) for the current transaction.
-- Callers must invoke broker.broker_set_tenant(...) before enqueue/claim;
-- if they don't, tenant_id defaults to 'default' and current_setting
-- also defaults to NULL -> broker_add_job coalesces to 'default' so
-- single-tenant use keeps working unmodified.
CREATE TABLE IF NOT EXISTS broker.broker_jobs (
id_broker_jobs BIGSERIAL PRIMARY KEY,
job_name VARCHAR(255) NOT NULL,
job_priority INTEGER NOT NULL DEFAULT 0,
job_queue INTEGER NOT NULL DEFAULT 1,
job_language VARCHAR(50) NOT NULL DEFAULT 'sql',
execute_str TEXT NOT NULL,
execute_result TEXT,
error_msg TEXT,
complete_status INTEGER NOT NULL DEFAULT 0,
run_as VARCHAR(100),
rid_broker_schedule BIGINT,
rid_broker_queueinstance BIGINT,
tenant_id TEXT NOT NULL DEFAULT 'default',
-- Lease / retry / idempotency
attempt_count INTEGER NOT NULL DEFAULT 0,
max_attempts INTEGER NOT NULL DEFAULT 1,
available_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
leased_at TIMESTAMP WITH TIME ZONE,
lease_expires_at TIMESTAMP WITH TIME ZONE,
lease_token UUID,
idempotency_key TEXT,
created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
started_at TIMESTAMP WITH TIME ZONE,
completed_at TIMESTAMP WITH TIME ZONE,
CONSTRAINT broker_jobs_complete_status_check CHECK (complete_status IN (0, 1, 2, 3, 4)),
CONSTRAINT broker_jobs_job_queue_check CHECK (job_queue > 0),
CONSTRAINT fk_schedule FOREIGN KEY (rid_broker_schedule) REFERENCES broker.broker_schedule(id_broker_schedule) ON DELETE SET NULL,
CONSTRAINT fk_instance FOREIGN KEY (rid_broker_queueinstance) REFERENCES broker.broker_queueinstance(id_broker_queueinstance) ON DELETE SET NULL
);
-- General-purpose indexes
CREATE INDEX IF NOT EXISTS idx_broker_jobs_status ON broker.broker_jobs(complete_status);
CREATE INDEX IF NOT EXISTS idx_broker_jobs_schedule ON broker.broker_jobs(rid_broker_schedule);
CREATE INDEX IF NOT EXISTS idx_broker_jobs_instance ON broker.broker_jobs(rid_broker_queueinstance);
CREATE INDEX IF NOT EXISTS idx_broker_jobs_created ON broker.broker_jobs(created_at);
CREATE INDEX IF NOT EXISTS idx_broker_jobs_name ON broker.broker_jobs(job_name, complete_status);
CREATE INDEX IF NOT EXISTS idx_broker_jobs_tenant ON broker.broker_jobs(tenant_id);
-- Claim index: exactly what broker_get's WHERE/ORDER BY needs, partial on pending rows only.
CREATE INDEX IF NOT EXISTS idx_broker_jobs_claim
ON broker.broker_jobs (job_queue, job_priority DESC, created_at, id_broker_jobs)
WHERE complete_status = 0;
-- Idempotency: at most one pending/any job per (queue, key) when a key is supplied.
CREATE UNIQUE INDEX IF NOT EXISTS idx_broker_jobs_idempotency
ON broker.broker_jobs (job_queue, idempotency_key)
WHERE idempotency_key IS NOT NULL;
COMMENT ON TABLE broker.broker_jobs IS 'Job queue for broker execution';
COMMENT ON COLUMN broker.broker_jobs.complete_status IS '0=pending, 1=running, 2=completed, 3=failed (terminal or dead-lettered once attempt_count>=max_attempts), 4=cancelled';
COMMENT ON COLUMN broker.broker_jobs.tenant_id IS 'RLS tenant scaffold; defaults to ''default'' for single-tenant use';
COMMENT ON COLUMN broker.broker_jobs.attempt_count IS 'Number of times this job has been claimed/executed';
COMMENT ON COLUMN broker.broker_jobs.max_attempts IS 'Job is dead-lettered (failed) once attempt_count reaches this value';
COMMENT ON COLUMN broker.broker_jobs.available_at IS 'Job is not claimable until now() >= available_at (used for retry backoff)';
COMMENT ON COLUMN broker.broker_jobs.lease_token IS 'Token handed out by broker_get; broker_run requires a matching token to execute, so an expired/reclaimed lease cannot be double-processed';
COMMENT ON COLUMN broker.broker_jobs.idempotency_key IS 'Optional caller-supplied key; unique per (job_queue, idempotency_key)';
CREATE OR REPLACE FUNCTION broker.tf_broker_jobs_update_timestamp()
RETURNS TRIGGER AS $$
BEGIN
NEW.updated_at = NOW();
RETURN NEW;
END;
$$ LANGUAGE plpgsql;
DROP TRIGGER IF EXISTS t_broker_jobs_updated_at ON broker.broker_jobs;
CREATE TRIGGER t_broker_jobs_updated_at
BEFORE UPDATE ON broker.broker_jobs
FOR EACH ROW
EXECUTE FUNCTION broker.tf_broker_jobs_update_timestamp();
-- Row Level Security: tenant isolation
ALTER TABLE broker.broker_jobs ENABLE ROW LEVEL SECURITY;
ALTER TABLE broker.broker_jobs FORCE ROW LEVEL SECURITY;
DROP POLICY IF EXISTS broker_jobs_tenant_isolation ON broker.broker_jobs;
CREATE POLICY broker_jobs_tenant_isolation ON broker.broker_jobs
USING (tenant_id = COALESCE(NULLIF(current_setting('broker.tenant_id', true), ''), 'default'))
WITH CHECK (tenant_id = COALESCE(NULLIF(current_setting('broker.tenant_id', true), ''), 'default'));
pkg/broker/install/sql/migrations/0005_broker_job_dependency.sql
+27
View File
@@ -0,0 +1,27 @@
-- broker.broker_job_dependency
-- Replaces the old broker_jobs.depends_on text[] column: job_id is only
-- claimable once every row it depends on has complete_status = 2 (completed).
CREATE TABLE IF NOT EXISTS broker.broker_job_dependency (
job_id BIGINT NOT NULL REFERENCES broker.broker_jobs(id_broker_jobs) ON DELETE CASCADE,
depends_on_job_id BIGINT NOT NULL REFERENCES broker.broker_jobs(id_broker_jobs) ON DELETE CASCADE,
tenant_id TEXT NOT NULL DEFAULT 'default',
created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
PRIMARY KEY (job_id, depends_on_job_id),
CHECK (job_id <> depends_on_job_id)
);
CREATE INDEX IF NOT EXISTS idx_broker_job_dependency_reverse
ON broker.broker_job_dependency (depends_on_job_id, job_id);
CREATE INDEX IF NOT EXISTS idx_broker_job_dependency_tenant
ON broker.broker_job_dependency (tenant_id);
COMMENT ON TABLE broker.broker_job_dependency IS 'job_id is not claimable until every depends_on_job_id row has complete_status = 2 (completed)';
ALTER TABLE broker.broker_job_dependency ENABLE ROW LEVEL SECURITY;
ALTER TABLE broker.broker_job_dependency FORCE ROW LEVEL SECURITY;
DROP POLICY IF EXISTS broker_job_dependency_tenant_isolation ON broker.broker_job_dependency;
CREATE POLICY broker_job_dependency_tenant_isolation ON broker.broker_job_dependency
USING (tenant_id = COALESCE(NULLIF(current_setting('broker.tenant_id', true), ''), 'default'))
WITH CHECK (tenant_id = COALESCE(NULLIF(current_setting('broker.tenant_id', true), ''), 'default'));
pkg/broker/install/sql/migrations/0006_broker_set_tenant.sql
+13
View File
@@ -0,0 +1,13 @@
-- broker.broker_set_tenant
-- Sets the RLS tenant context for the current transaction (SET LOCAL semantics
-- via set_config(..., true)). Callers must invoke this before enqueue/claim
-- if they are not using the 'default' tenant.
CREATE OR REPLACE FUNCTION broker.broker_set_tenant(p_tenant_id TEXT)
RETURNS VOID
LANGUAGE SQL
AS $$
SELECT set_config('broker.tenant_id', p_tenant_id, true);
$$;
COMMENT ON FUNCTION broker.broker_set_tenant IS 'Sets broker.tenant_id for the current transaction only (SET LOCAL semantics)';
pkg/broker/install/sql/migrations/0007_broker_get.sql
+83
View File
@@ -0,0 +1,83 @@
-- broker.broker_get
-- Claims the next eligible job from a queue: pending, available (backoff
-- elapsed), no incomplete dependency, and visible under the caller's RLS
-- tenant. Grants a lease (lease_token) that broker_run must present back.
-- Returns: p_retval (0=success, >0=infra error), p_errmsg, p_job_id, p_lease_token.
CREATE OR REPLACE FUNCTION broker.broker_get(
p_queue_number INTEGER,
p_instance_id BIGINT DEFAULT NULL,
p_lease_seconds INTEGER DEFAULT 60,
OUT p_retval INTEGER,
OUT p_errmsg TEXT,
OUT p_job_id BIGINT,
OUT p_lease_token UUID
)
RETURNS RECORD
LANGUAGE plpgsql
AS $$
DECLARE
v_job_id BIGINT;
v_lease_token UUID;
BEGIN
p_retval := 0;
p_errmsg := '';
p_job_id := NULL;
p_lease_token := NULL;
IF p_queue_number IS NULL OR p_queue_number <= 0 THEN
p_retval := 1;
p_errmsg := 'Invalid queue number';
RETURN;
END IF;
IF p_lease_seconds IS NULL OR p_lease_seconds <= 0 THEN
p_lease_seconds := 60;
END IF;
SELECT candidate.id_broker_jobs
INTO v_job_id
FROM broker.broker_jobs candidate
WHERE candidate.job_queue = p_queue_number
AND candidate.complete_status = 0
AND candidate.available_at <= NOW()
AND NOT EXISTS (
SELECT 1
FROM broker.broker_job_dependency d
JOIN broker.broker_jobs dep ON dep.id_broker_jobs = d.depends_on_job_id
WHERE d.job_id = candidate.id_broker_jobs
AND dep.complete_status <> 2
)
ORDER BY candidate.job_priority DESC, candidate.created_at ASC, candidate.id_broker_jobs ASC
LIMIT 1
FOR UPDATE SKIP LOCKED;
IF NOT FOUND THEN
RETURN;
END IF;
v_lease_token := gen_random_uuid();
UPDATE broker.broker_jobs
SET complete_status = 1, -- running
started_at = NOW(),
rid_broker_queueinstance = p_instance_id,
attempt_count = attempt_count + 1,
lease_token = v_lease_token,
leased_at = NOW(),
lease_expires_at = NOW() + make_interval(secs => p_lease_seconds),
updated_at = NOW()
WHERE id_broker_jobs = v_job_id;
p_job_id := v_job_id;
p_lease_token := v_lease_token;
EXCEPTION
WHEN OTHERS THEN
p_retval := 2;
p_errmsg := SQLERRM;
RAISE WARNING 'broker_get error: %', SQLERRM;
END;
$$;
COMMENT ON FUNCTION broker.broker_get IS 'Claims the next eligible job from a queue and grants a lease';
pkg/broker/install/sql/migrations/0008_broker_run.sql
+135
View File
@@ -0,0 +1,135 @@
-- broker.broker_run
-- Executes a job by its ID, presenting the lease token it was claimed with.
--
-- p_retval is reserved for infra failures (bad job id, job not found, wrong
-- state, lease mismatch/expired, DB error). An executed-and-caught job
-- failure is a *successful* invocation: p_retval stays 0 and the outcome is
-- reported via p_job_status (0=requeued for retry, 2=completed, 3=dead-lettered)
-- so the caller commits the terminal/retry state instead of rolling it back.
--
-- On failure, if attempt_count < max_attempts the job is reset to pending
-- with exponential backoff (base 5s, capped at 300s); otherwise it is
-- dead-lettered as complete_status = 3.
CREATE OR REPLACE FUNCTION broker.broker_run(
p_job_id BIGINT,
p_lease_token UUID,
OUT p_retval INTEGER,
OUT p_errmsg TEXT,
OUT p_job_status INTEGER
)
RETURNS RECORD
LANGUAGE plpgsql
AS $$
DECLARE
v_job_record RECORD;
v_execute_result TEXT;
v_error_occurred BOOLEAN := false;
v_backoff_base CONSTANT INTEGER := 5;
v_backoff_cap CONSTANT INTEGER := 300;
v_backoff_secs INTEGER;
BEGIN
p_retval := 0;
p_errmsg := '';
p_job_status := NULL;
v_execute_result := '';
IF p_job_id IS NULL OR p_job_id <= 0 THEN
p_retval := 1;
p_errmsg := 'Invalid job ID';
RETURN;
END IF;
SELECT id_broker_jobs, execute_str, job_language, complete_status, attempt_count, max_attempts, lease_token
INTO v_job_record
FROM broker.broker_jobs
WHERE id_broker_jobs = p_job_id
FOR UPDATE;
IF NOT FOUND THEN
p_retval := 2;
p_errmsg := 'Job not found';
RETURN;
END IF;
IF v_job_record.complete_status != 1 THEN
p_retval := 3;
p_errmsg := format('Job is not in running state (status: %s)', v_job_record.complete_status);
RETURN;
END IF;
IF v_job_record.lease_token IS DISTINCT FROM p_lease_token THEN
p_retval := 4;
p_errmsg := 'Lease token mismatch or expired; job was reclaimed by another worker';
RETURN;
END IF;
-- Execute the job
BEGIN
IF v_job_record.job_language IN ('sql', 'plpgsql') THEN
EXECUTE v_job_record.execute_str;
v_execute_result := 'Success';
ELSE
v_error_occurred := true;
v_execute_result := format('Unsupported job language: %s', v_job_record.job_language);
END IF;
EXCEPTION
WHEN OTHERS THEN
v_error_occurred := true;
v_execute_result := format('Error: %s', SQLERRM);
END;
IF v_error_occurred THEN
IF v_job_record.attempt_count < v_job_record.max_attempts THEN
v_backoff_secs := LEAST(POWER(2, v_job_record.attempt_count)::INTEGER * v_backoff_base, v_backoff_cap);
UPDATE broker.broker_jobs
SET complete_status = 0, -- pending, retry
available_at = NOW() + make_interval(secs => v_backoff_secs),
error_msg = v_execute_result,
execute_result = v_execute_result,
lease_token = NULL,
leased_at = NULL,
lease_expires_at = NULL,
updated_at = NOW()
WHERE id_broker_jobs = p_job_id;
p_job_status := 0;
ELSE
UPDATE broker.broker_jobs
SET complete_status = 3, -- failed (dead-letter, attempts exhausted)
error_msg = v_execute_result,
execute_result = v_execute_result,
lease_token = NULL,
leased_at = NULL,
lease_expires_at = NULL,
completed_at = NOW(),
updated_at = NOW()
WHERE id_broker_jobs = p_job_id;
p_job_status := 3;
END IF;
ELSE
UPDATE broker.broker_jobs
SET complete_status = 2, -- completed
execute_result = v_execute_result,
error_msg = NULL,
lease_token = NULL,
leased_at = NULL,
lease_expires_at = NULL,
completed_at = NOW(),
updated_at = NOW()
WHERE id_broker_jobs = p_job_id;
p_job_status := 2;
END IF;
EXCEPTION
WHEN OTHERS THEN
p_retval := 6;
p_errmsg := SQLERRM;
RAISE WARNING 'broker_run error: %', SQLERRM;
END;
$$;
COMMENT ON FUNCTION broker.broker_run IS 'Executes a leased job; reports outcome via p_job_status without forcing a rollback of the terminal/retry state';
pkg/broker/install/sql/migrations/0009_broker_set.sql
+65
View File
@@ -0,0 +1,65 @@
-- broker.broker_set
-- Minimal whitelist of session options. The previous SET SESSION AUTHORIZATION
-- and search_path branches were removed: they let a caller assume an arbitrary
-- Postgres role or schema search order from inside a plpgsql function with no
-- identity model behind it, which is unsafe and was unused.
CREATE OR REPLACE FUNCTION broker.broker_set(
p_option_name TEXT,
p_option_value TEXT,
OUT p_retval INTEGER,
OUT p_errmsg TEXT
)
RETURNS RECORD
LANGUAGE plpgsql
AS $$
DECLARE
v_sql TEXT;
BEGIN
p_retval := 0;
p_errmsg := '';
IF p_option_name IS NULL OR p_option_name = '' THEN
p_retval := 1;
p_errmsg := 'Option name is required';
RETURN;
END IF;
CASE LOWER(p_option_name)
WHEN 'application_name' THEN
BEGIN
v_sql := format('SET LOCAL application_name TO %L', p_option_value);
EXECUTE v_sql;
EXCEPTION
WHEN OTHERS THEN
p_retval := 3;
p_errmsg := format('Failed to set application_name: %s', SQLERRM);
RETURN;
END;
WHEN 'timezone' THEN
BEGIN
v_sql := format('SET LOCAL timezone TO %L', p_option_value);
EXECUTE v_sql;
EXCEPTION
WHEN OTHERS THEN
p_retval := 5;
p_errmsg := format('Failed to set timezone: %s', SQLERRM);
RETURN;
END;
ELSE
p_retval := 10;
p_errmsg := format('Unknown option: %s', p_option_name);
RETURN;
END CASE;
EXCEPTION
WHEN OTHERS THEN
p_retval := 99;
p_errmsg := SQLERRM;
RAISE WARNING 'broker_set error: %', SQLERRM;
END;
$$;
COMMENT ON FUNCTION broker.broker_set IS 'Sets a whitelisted session-local option (application_name, timezone)';
pkg/broker/install/sql/migrations/0010_broker_register_instance.sql
+64
View File
@@ -0,0 +1,64 @@
-- broker.broker_register_instance
-- Registers a broker instance, using a session-scoped advisory lock keyed by
-- name to guarantee only one active instance per name -- no race window, no
-- "check COUNT(*) then insert" gap. The caller MUST run this on a pinned
-- connection it keeps open for the process lifetime (Go: db.Conn(ctx)) and
-- release the lock (pg_advisory_unlock) itself on shutdown, since the lock
-- lives with the backend session, not with the row.
CREATE OR REPLACE FUNCTION broker.broker_register_instance(
p_name TEXT,
p_hostname TEXT,
p_pid INTEGER,
p_version TEXT,
p_queue_count INTEGER,
OUT p_retval INTEGER,
OUT p_errmsg TEXT,
OUT p_instance_id BIGINT
)
RETURNS RECORD
LANGUAGE plpgsql
AS $$
DECLARE
v_lock_key BIGINT;
BEGIN
p_retval := 0;
p_errmsg := '';
p_instance_id := NULL;
IF p_name IS NULL OR p_name = '' THEN
p_retval := 1;
p_errmsg := 'Instance name is required';
RETURN;
END IF;
IF p_hostname IS NULL OR p_hostname = '' THEN
p_retval := 2;
p_errmsg := 'Hostname is required';
RETURN;
END IF;
v_lock_key := hashtextextended('broker:' || p_name, 0);
IF NOT pg_try_advisory_lock(v_lock_key) THEN
p_retval := 3;
p_errmsg := 'Another broker instance is already active for this name (advisory lock held). Only one broker instance per name is allowed.';
RETURN;
END IF;
INSERT INTO broker.broker_queueinstance (
name, hostname, pid, version, status, queue_count, started_at, last_ping_at
) VALUES (
p_name, p_hostname, p_pid, p_version, 'active', p_queue_count, NOW(), NOW()
)
RETURNING id_broker_queueinstance INTO p_instance_id;
EXCEPTION
WHEN OTHERS THEN
p_retval := 99;
p_errmsg := SQLERRM;
RAISE WARNING 'broker_register_instance error: %', SQLERRM;
END;
$$;
COMMENT ON FUNCTION broker.broker_register_instance IS 'Registers a broker instance; caller must hold the connection open for the process lifetime (advisory lock is session-scoped)';
pkg/broker/install/sql/migrations/0011_broker_add_job.sql
+137
View File
@@ -0,0 +1,137 @@
-- broker.broker_add_job
-- Adds a new job (optionally with dependencies and an idempotency key) and
-- sends a wake-only NOTIFY -- the payload carries only the queue number
-- (job id kept solely for logging); workers re-claim via broker_get rather
-- than executing the notified row directly, so a notification can never
-- hand a job to a worker before it's actually claimable.
CREATE OR REPLACE FUNCTION broker.broker_add_job(
p_job_name TEXT,
p_execute_str TEXT,
p_job_queue INTEGER DEFAULT 1,
p_job_priority INTEGER DEFAULT 0,
p_job_language TEXT DEFAULT 'sql',
p_run_as TEXT DEFAULT NULL,
p_schedule_id BIGINT DEFAULT NULL,
p_depends_on_job_ids BIGINT[] DEFAULT NULL,
p_idempotency_key TEXT DEFAULT NULL,
p_max_attempts INTEGER DEFAULT 1,
OUT p_retval INTEGER,
OUT p_errmsg TEXT,
OUT p_job_id BIGINT
)
RETURNS RECORD
LANGUAGE plpgsql
AS $$
DECLARE
v_notification_payload JSON;
v_tenant_id TEXT;
v_dep_id BIGINT;
v_cycle_exists BOOLEAN;
BEGIN
p_retval := 0;
p_errmsg := '';
p_job_id := NULL;
IF p_job_name IS NULL OR p_job_name = '' THEN
p_retval := 1;
p_errmsg := 'Job name is required';
RETURN;
END IF;
IF p_execute_str IS NULL OR p_execute_str = '' THEN
p_retval := 2;
p_errmsg := 'Execute string is required';
RETURN;
END IF;
IF p_job_queue IS NULL OR p_job_queue <= 0 THEN
p_retval := 3;
p_errmsg := 'Invalid job queue number';
RETURN;
END IF;
IF p_max_attempts IS NULL OR p_max_attempts <= 0 THEN
p_max_attempts := 1;
END IF;
-- Falls back to 'default' when the caller never called broker_set_tenant,
-- so single-tenant use (and the RLS WITH CHECK on insert) keeps working.
v_tenant_id := COALESCE(NULLIF(current_setting('broker.tenant_id', true), ''), 'default');
IF p_idempotency_key IS NOT NULL THEN
SELECT id_broker_jobs INTO p_job_id
FROM broker.broker_jobs
WHERE job_queue = p_job_queue
AND idempotency_key = p_idempotency_key
AND tenant_id = v_tenant_id;
IF FOUND THEN
p_errmsg := 'Job with this idempotency key already exists; returning existing job id';
RETURN;
END IF;
END IF;
INSERT INTO broker.broker_jobs (
job_name, job_priority, job_queue, job_language, execute_str, run_as,
rid_broker_schedule, tenant_id, max_attempts, idempotency_key, complete_status
) VALUES (
p_job_name, p_job_priority, p_job_queue, p_job_language, p_execute_str, p_run_as,
p_schedule_id, v_tenant_id, p_max_attempts, p_idempotency_key, 0
)
RETURNING id_broker_jobs INTO p_job_id;
IF p_depends_on_job_ids IS NOT NULL THEN
FOREACH v_dep_id IN ARRAY p_depends_on_job_ids LOOP
IF v_dep_id IS NULL THEN
CONTINUE;
END IF;
IF v_dep_id = p_job_id THEN
p_retval := 20;
p_errmsg := 'Invalid dependency: a job cannot depend on itself';
RETURN;
END IF;
SELECT EXISTS (
SELECT 1 FROM broker.broker_job_dependency
WHERE job_id = v_dep_id AND depends_on_job_id = p_job_id
) INTO v_cycle_exists;
IF v_cycle_exists THEN
p_retval := 21;
p_errmsg := format('Invalid dependency: job %s already depends on %s (would create a cycle)', v_dep_id, p_job_id);
RETURN;
END IF;
INSERT INTO broker.broker_job_dependency (job_id, depends_on_job_id, tenant_id)
VALUES (p_job_id, v_dep_id, v_tenant_id)
ON CONFLICT (job_id, depends_on_job_id) DO NOTHING;
END LOOP;
END IF;
v_notification_payload := json_build_object(
'queue', p_job_queue,
'job_id', p_job_id
);
PERFORM pg_notify('broker.event', v_notification_payload::text);
EXCEPTION
WHEN unique_violation THEN
-- Concurrent insert raced us to the same idempotency key.
p_retval := 0;
p_errmsg := 'Job with this idempotency key already exists; returning existing job id';
SELECT id_broker_jobs INTO p_job_id
FROM broker.broker_jobs
WHERE job_queue = p_job_queue
AND idempotency_key = p_idempotency_key
AND tenant_id = v_tenant_id;
WHEN OTHERS THEN
p_retval := 99;
p_errmsg := SQLERRM;
RAISE WARNING 'broker_add_job error: %', SQLERRM;
END;
$$;
COMMENT ON FUNCTION broker.broker_add_job IS 'Adds a job (with optional dependencies/idempotency key) and sends a wake-only NOTIFY';
pkg/broker/install/sql/procedures/06_broker_ping_instance.sql → pkg/broker/install/sql/migrations/0012_broker_ping_shutdown_instance.sql
+10 -25
View File
@@ -1,8 +1,6 @@
-- broker_ping_instance function
-- Updates the last_ping_at timestamp for a broker instance
-- Returns: p_retval (0=success, >0=error), p_errmsg (error message)
-- broker.broker_ping_instance / broker.broker_shutdown_instance
CREATE OR REPLACE FUNCTION broker_ping_instance(
CREATE OR REPLACE FUNCTION broker.broker_ping_instance(
p_instance_id BIGINT,
p_jobs_handled BIGINT DEFAULT NULL,
OUT p_retval INTEGER,
@@ -15,26 +13,22 @@ BEGIN
p_retval := 0;
p_errmsg := '';
-- Validate instance ID
IF p_instance_id IS NULL OR p_instance_id <= 0 THEN
p_retval := 1;
p_errmsg := 'Invalid instance ID';
RETURN;
END IF;
-- Update ping timestamp
IF p_jobs_handled IS NOT NULL THEN
UPDATE broker_queueinstance
SET last_ping_at = NOW(),
jobs_handled = p_jobs_handled
UPDATE broker.broker_queueinstance
SET last_ping_at = NOW(), jobs_handled = p_jobs_handled
WHERE id_broker_queueinstance = p_instance_id;
ELSE
UPDATE broker_queueinstance
UPDATE broker.broker_queueinstance
SET last_ping_at = NOW()
WHERE id_broker_queueinstance = p_instance_id;
END IF;
-- Check if instance was found
IF NOT FOUND THEN
p_retval := 2;
p_errmsg := 'Instance not found';
@@ -49,11 +43,7 @@ EXCEPTION
END;
$$;
-- broker_shutdown_instance function
-- Marks a broker instance as shutdown
-- Returns: p_retval (0=success, >0=error), p_errmsg (error message)
CREATE OR REPLACE FUNCTION broker_shutdown_instance(
CREATE OR REPLACE FUNCTION broker.broker_shutdown_instance(
p_instance_id BIGINT,
OUT p_retval INTEGER,
OUT p_errmsg TEXT
@@ -65,20 +55,16 @@ BEGIN
p_retval := 0;
p_errmsg := '';
-- Validate instance ID
IF p_instance_id IS NULL OR p_instance_id <= 0 THEN
p_retval := 1;
p_errmsg := 'Invalid instance ID';
RETURN;
END IF;
-- Update instance status
UPDATE broker_queueinstance
SET status = 'shutdown',
shutdown_at = NOW()
UPDATE broker.broker_queueinstance
SET status = 'shutdown', shutdown_at = NOW()
WHERE id_broker_queueinstance = p_instance_id;
-- Check if instance was found
IF NOT FOUND THEN
p_retval := 2;
p_errmsg := 'Instance not found';
@@ -93,6 +79,5 @@ EXCEPTION
END;
$$;
-- Comments
COMMENT ON FUNCTION broker_ping_instance IS 'Updates the last ping timestamp for an instance';
COMMENT ON FUNCTION broker_shutdown_instance IS 'Marks an instance as shutdown';
COMMENT ON FUNCTION broker.broker_ping_instance IS 'Updates the last ping timestamp for an instance';
COMMENT ON FUNCTION broker.broker_shutdown_instance IS 'Marks an instance as shutdown (does not release the advisory lock -- caller must pg_advisory_unlock on its pinned connection)';
pkg/broker/install/sql/migrations/0013_broker_recover_stale_jobs.sql
+66
View File
@@ -0,0 +1,66 @@
-- broker.broker_recover_stale_jobs
-- Recovers jobs whose lease has expired while still 'running' (a worker died
-- or was killed mid-execution without updating status). Applies the same
-- retry/backoff rule as broker_run: retry while attempts remain, otherwise
-- dead-letter. Marked SECURITY DEFINER so the sweep runs across all tenants
-- regardless of caller: this requires the function's owner (whichever role
-- runs the migrations, intended to be broker_admin) to have BYPASSRLS --
-- broker_jobs/broker_job_dependency use FORCE ROW LEVEL SECURITY, so without
-- BYPASSRLS on the owner this would silently only ever see the empty/no
-- tenant context.
CREATE OR REPLACE FUNCTION broker.broker_recover_stale_jobs(
OUT p_retval INTEGER,
OUT p_errmsg TEXT,
OUT p_recovered_count INTEGER
)
RETURNS RECORD
LANGUAGE plpgsql
SECURITY DEFINER
SET search_path = broker, pg_temp
AS $$
DECLARE
v_backoff_base CONSTANT INTEGER := 5;
v_backoff_cap CONSTANT INTEGER := 300;
BEGIN
p_retval := 0;
p_errmsg := '';
p_recovered_count := 0;
WITH stale AS (
SELECT id_broker_jobs, attempt_count, max_attempts
FROM broker.broker_jobs
WHERE complete_status = 1
AND lease_expires_at IS NOT NULL
AND lease_expires_at < NOW()
FOR UPDATE SKIP LOCKED
),
recovered AS (
UPDATE broker.broker_jobs j
SET complete_status = CASE WHEN s.attempt_count < s.max_attempts THEN 0 ELSE 3 END,
available_at = CASE
WHEN s.attempt_count < s.max_attempts
THEN NOW() + make_interval(secs => LEAST(POWER(2, s.attempt_count)::INTEGER * v_backoff_base, v_backoff_cap))
ELSE j.available_at
END,
error_msg = CASE WHEN s.attempt_count >= s.max_attempts THEN COALESCE(j.error_msg, 'Lease expired and max attempts exhausted') ELSE j.error_msg END,
completed_at = CASE WHEN s.attempt_count >= s.max_attempts THEN NOW() ELSE j.completed_at END,
lease_token = NULL,
leased_at = NULL,
lease_expires_at = NULL,
updated_at = NOW()
FROM stale s
WHERE j.id_broker_jobs = s.id_broker_jobs
RETURNING j.id_broker_jobs
)
SELECT COUNT(*) INTO p_recovered_count FROM recovered;
EXCEPTION
WHEN OTHERS THEN
p_retval := 99;
p_errmsg := SQLERRM;
RAISE WARNING 'broker_recover_stale_jobs error: %', SQLERRM;
END;
$$;
COMMENT ON FUNCTION broker.broker_recover_stale_jobs IS 'Requeues (or dead-letters) jobs whose lease expired while still running';
pkg/broker/install/sql/migrations/0014_job_groups.sql
+41
View File
@@ -0,0 +1,41 @@
-- Adds job groups: every job belongs to a job_group (defaults to its own
-- job_name when not given explicitly, set by broker_add_job). A dependency
-- can now target a whole group instead of a single job id -- the dependent
-- job is claimable once every job tagged with that group has completed
-- (complete_status = 2); already-completed group members simply drop out of
-- the gating check, they don't need to have existed at any particular time.
-- The existing id-based dependency (broker_job_dependency.depends_on_job_id)
-- is kept as-is; each dependency row targets exactly one of an id or a group.
ALTER TABLE broker.broker_jobs ADD COLUMN IF NOT EXISTS job_group TEXT;
UPDATE broker.broker_jobs SET job_group = job_name WHERE job_group IS NULL;
ALTER TABLE broker.broker_jobs ALTER COLUMN job_group SET NOT NULL;
CREATE INDEX IF NOT EXISTS idx_broker_jobs_group
ON broker.broker_jobs (tenant_id, job_group, complete_status);
ALTER TABLE broker.broker_job_dependency DROP CONSTRAINT IF EXISTS broker_job_dependency_pkey;
ALTER TABLE broker.broker_job_dependency ALTER COLUMN depends_on_job_id DROP NOT NULL;
ALTER TABLE broker.broker_job_dependency ADD COLUMN IF NOT EXISTS depends_on_group TEXT;
ALTER TABLE broker.broker_job_dependency ADD COLUMN IF NOT EXISTS id_broker_job_dependency BIGSERIAL;
ALTER TABLE broker.broker_job_dependency ADD PRIMARY KEY (id_broker_job_dependency);
ALTER TABLE broker.broker_job_dependency DROP CONSTRAINT IF EXISTS broker_job_dependency_one_target;
ALTER TABLE broker.broker_job_dependency ADD CONSTRAINT broker_job_dependency_one_target CHECK (
(depends_on_job_id IS NOT NULL AND depends_on_group IS NULL) OR
(depends_on_job_id IS NULL AND depends_on_group IS NOT NULL)
);
-- Plain (non-partial) unique constraint, matching the old PK's guarantee --
-- NULLs in depends_on_job_id (the group-dependency rows) are never
-- considered equal by a standard unique constraint, so this only constrains
-- id-based rows, and keeps "ON CONFLICT (job_id, depends_on_job_id)" (no
-- predicate needed) working for existing callers.
ALTER TABLE broker.broker_job_dependency DROP CONSTRAINT IF EXISTS broker_job_dependency_unique_id;
ALTER TABLE broker.broker_job_dependency ADD CONSTRAINT broker_job_dependency_unique_id
UNIQUE (job_id, depends_on_job_id);
CREATE UNIQUE INDEX IF NOT EXISTS idx_broker_job_dependency_unique_group
ON broker.broker_job_dependency (job_id, depends_on_group) WHERE depends_on_group IS NOT NULL;
COMMENT ON COLUMN broker.broker_jobs.job_group IS 'Group tag for this job; defaults to job_name. Other jobs can depend on the whole group.';
COMMENT ON COLUMN broker.broker_job_dependency.depends_on_group IS 'Alternative to depends_on_job_id: job_id is not claimable until every job with job_group = depends_on_group has completed';
pkg/broker/install/sql/migrations/0015_broker_add_job_groups.sql
+184
View File
@@ -0,0 +1,184 @@
-- broker.broker_add_job: adds job groups.
-- p_job_group defaults to p_job_name when not given. p_depends_on_groups is
-- the group-based counterpart to the existing p_depends_on_job_ids: the new
-- job is not claimable until every job tagged with each named group has
-- completed. Both dependency kinds can be combined on the same job.
-- New parameters are appended after existing ones with defaults, so every
-- existing positional call (however many args it passes) keeps working
-- unchanged. Appending arguments changes the function's identity though --
-- CREATE OR REPLACE would create a second, ambiguous overload rather than
-- replacing -- so the old 10-arg signature is dropped explicitly first.
DROP FUNCTION IF EXISTS broker.broker_add_job(
TEXT, TEXT, INTEGER, INTEGER, TEXT, TEXT, BIGINT, BIGINT[], TEXT, INTEGER
);
CREATE OR REPLACE FUNCTION broker.broker_add_job(
p_job_name TEXT,
p_execute_str TEXT,
p_job_queue INTEGER DEFAULT 1,
p_job_priority INTEGER DEFAULT 0,
p_job_language TEXT DEFAULT 'sql',
p_run_as TEXT DEFAULT NULL,
p_schedule_id BIGINT DEFAULT NULL,
p_depends_on_job_ids BIGINT[] DEFAULT NULL,
p_idempotency_key TEXT DEFAULT NULL,
p_max_attempts INTEGER DEFAULT 1,
p_job_group TEXT DEFAULT NULL,
p_depends_on_groups TEXT[] DEFAULT NULL,
OUT p_retval INTEGER,
OUT p_errmsg TEXT,
OUT p_job_id BIGINT
)
RETURNS RECORD
LANGUAGE plpgsql
AS $$
DECLARE
v_notification_payload JSON;
v_tenant_id TEXT;
v_job_group TEXT;
v_dep_id BIGINT;
v_dep_group TEXT;
v_cycle_exists BOOLEAN;
BEGIN
p_retval := 0;
p_errmsg := '';
p_job_id := NULL;
IF p_job_name IS NULL OR p_job_name = '' THEN
p_retval := 1;
p_errmsg := 'Job name is required';
RETURN;
END IF;
IF p_execute_str IS NULL OR p_execute_str = '' THEN
p_retval := 2;
p_errmsg := 'Execute string is required';
RETURN;
END IF;
IF p_job_queue IS NULL OR p_job_queue <= 0 THEN
p_retval := 3;
p_errmsg := 'Invalid job queue number';
RETURN;
END IF;
IF p_max_attempts IS NULL OR p_max_attempts <= 0 THEN
p_max_attempts := 1;
END IF;
v_job_group := COALESCE(NULLIF(p_job_group, ''), p_job_name);
-- Falls back to 'default' when the caller never called broker_set_tenant,
-- so single-tenant use (and the RLS WITH CHECK on insert) keeps working.
v_tenant_id := COALESCE(NULLIF(current_setting('broker.tenant_id', true), ''), 'default');
IF p_idempotency_key IS NOT NULL THEN
SELECT id_broker_jobs INTO p_job_id
FROM broker.broker_jobs
WHERE job_queue = p_job_queue
AND idempotency_key = p_idempotency_key
AND tenant_id = v_tenant_id;
IF FOUND THEN
p_errmsg := 'Job with this idempotency key already exists; returning existing job id';
RETURN;
END IF;
END IF;
INSERT INTO broker.broker_jobs (
job_name, job_group, job_priority, job_queue, job_language, execute_str, run_as,
rid_broker_schedule, tenant_id, max_attempts, idempotency_key, complete_status
) VALUES (
p_job_name, v_job_group, p_job_priority, p_job_queue, p_job_language, p_execute_str, p_run_as,
p_schedule_id, v_tenant_id, p_max_attempts, p_idempotency_key, 0
)
RETURNING id_broker_jobs INTO p_job_id;
IF p_depends_on_job_ids IS NOT NULL THEN
FOREACH v_dep_id IN ARRAY p_depends_on_job_ids LOOP
IF v_dep_id IS NULL THEN
CONTINUE;
END IF;
IF v_dep_id = p_job_id THEN
p_retval := 20;
p_errmsg := 'Invalid dependency: a job cannot depend on itself';
RETURN;
END IF;
SELECT EXISTS (
SELECT 1 FROM broker.broker_job_dependency
WHERE job_id = v_dep_id AND depends_on_job_id = p_job_id
) INTO v_cycle_exists;
IF v_cycle_exists THEN
p_retval := 21;
p_errmsg := format('Invalid dependency: job %s already depends on %s (would create a cycle)', v_dep_id, p_job_id);
RETURN;
END IF;
INSERT INTO broker.broker_job_dependency (job_id, depends_on_job_id, tenant_id)
VALUES (p_job_id, v_dep_id, v_tenant_id)
ON CONFLICT (job_id, depends_on_job_id) DO NOTHING;
END LOOP;
END IF;
IF p_depends_on_groups IS NOT NULL THEN
FOREACH v_dep_group IN ARRAY p_depends_on_groups LOOP
IF v_dep_group IS NULL OR v_dep_group = '' THEN
CONTINUE;
END IF;
IF v_dep_group = v_job_group THEN
p_retval := 22;
p_errmsg := format('Invalid dependency: a job cannot depend on its own group (%s)', v_job_group);
RETURN;
END IF;
SELECT EXISTS (
SELECT 1
FROM broker.broker_jobs j
JOIN broker.broker_job_dependency d ON d.job_id = j.id_broker_jobs
WHERE j.tenant_id = v_tenant_id
AND j.job_group = v_dep_group
AND d.depends_on_group = v_job_group
) INTO v_cycle_exists;
IF v_cycle_exists THEN
p_retval := 23;
p_errmsg := format('Invalid dependency: group %s already depends on %s (would create a cycle)', v_dep_group, v_job_group);
RETURN;
END IF;
INSERT INTO broker.broker_job_dependency (job_id, depends_on_group, tenant_id)
VALUES (p_job_id, v_dep_group, v_tenant_id)
ON CONFLICT (job_id, depends_on_group) WHERE depends_on_group IS NOT NULL DO NOTHING;
END LOOP;
END IF;
v_notification_payload := json_build_object(
'queue', p_job_queue,
'job_id', p_job_id
);
PERFORM pg_notify('broker.event', v_notification_payload::text);
EXCEPTION
WHEN unique_violation THEN
-- Concurrent insert raced us to the same idempotency key.
p_retval := 0;
p_errmsg := 'Job with this idempotency key already exists; returning existing job id';
SELECT id_broker_jobs INTO p_job_id
FROM broker.broker_jobs
WHERE job_queue = p_job_queue
AND idempotency_key = p_idempotency_key
AND tenant_id = v_tenant_id;
WHEN OTHERS THEN
p_retval := 99;
p_errmsg := SQLERRM;
RAISE WARNING 'broker_add_job error: %', SQLERRM;
END;
$$;
COMMENT ON FUNCTION broker.broker_add_job IS 'Adds a job (with optional id/group dependencies, job group, and idempotency key) and sends a wake-only NOTIFY';
pkg/broker/install/sql/migrations/0016_broker_get_groups.sql
+91
View File
@@ -0,0 +1,91 @@
-- broker.broker_get: also gates claiming on group-based dependencies
-- (broker_job_dependency.depends_on_group) alongside the existing id-based
-- ones. Signature is unchanged, only the eligibility query grows a second
-- NOT EXISTS clause.
CREATE OR REPLACE FUNCTION broker.broker_get(
p_queue_number INTEGER,
p_instance_id BIGINT DEFAULT NULL,
p_lease_seconds INTEGER DEFAULT 60,
OUT p_retval INTEGER,
OUT p_errmsg TEXT,
OUT p_job_id BIGINT,
OUT p_lease_token UUID
)
RETURNS RECORD
LANGUAGE plpgsql
AS $$
DECLARE
v_job_id BIGINT;
v_lease_token UUID;
BEGIN
p_retval := 0;
p_errmsg := '';
p_job_id := NULL;
p_lease_token := NULL;
IF p_queue_number IS NULL OR p_queue_number <= 0 THEN
p_retval := 1;
p_errmsg := 'Invalid queue number';
RETURN;
END IF;
IF p_lease_seconds IS NULL OR p_lease_seconds <= 0 THEN
p_lease_seconds := 60;
END IF;
SELECT candidate.id_broker_jobs
INTO v_job_id
FROM broker.broker_jobs candidate
WHERE candidate.job_queue = p_queue_number
AND candidate.complete_status = 0
AND candidate.available_at <= NOW()
AND NOT EXISTS (
SELECT 1
FROM broker.broker_job_dependency d
JOIN broker.broker_jobs dep ON dep.id_broker_jobs = d.depends_on_job_id
WHERE d.job_id = candidate.id_broker_jobs
AND dep.complete_status <> 2
)
AND NOT EXISTS (
SELECT 1
FROM broker.broker_job_dependency d
JOIN broker.broker_jobs dep ON dep.tenant_id = candidate.tenant_id
AND dep.job_group = d.depends_on_group
WHERE d.job_id = candidate.id_broker_jobs
AND dep.id_broker_jobs <> candidate.id_broker_jobs
AND dep.complete_status <> 2
)
ORDER BY candidate.job_priority DESC, candidate.created_at ASC, candidate.id_broker_jobs ASC
LIMIT 1
FOR UPDATE SKIP LOCKED;
IF NOT FOUND THEN
RETURN;
END IF;
v_lease_token := gen_random_uuid();
UPDATE broker.broker_jobs
SET complete_status = 1, -- running
started_at = NOW(),
rid_broker_queueinstance = p_instance_id,
attempt_count = attempt_count + 1,
lease_token = v_lease_token,
leased_at = NOW(),
lease_expires_at = NOW() + make_interval(secs => p_lease_seconds),
updated_at = NOW()
WHERE id_broker_jobs = v_job_id;
p_job_id := v_job_id;
p_lease_token := v_lease_token;
EXCEPTION
WHEN OTHERS THEN
p_retval := 2;
p_errmsg := SQLERRM;
RAISE WARNING 'broker_get error: %', SQLERRM;
END;
$$;
COMMENT ON FUNCTION broker.broker_get IS 'Claims the next eligible job from a queue (honoring id- and group-based dependencies) and grants a lease';
pkg/broker/install/sql/migrations/0017_broker_add_job_simple_groups.sql
+41
View File
@@ -0,0 +1,41 @@
-- broker.broker_add_job_simple
-- Convenience wrapper around broker.broker_add_job for the common case: a
-- job named p_job_name running p_execute_str at p_job_priority, depending on
-- other jobs by group name (p_depends_on_groups) -- every job's group
-- defaults to its own job_name, so passing job names here just works.
-- Pure pass-through: no name-to-id resolution needed since dependencies are
-- resolved live, by group, inside broker_get.
CREATE OR REPLACE FUNCTION broker.broker_add_job_simple(
p_job_name TEXT,
p_execute_str TEXT,
p_job_priority INTEGER DEFAULT 0,
p_depends_on_groups TEXT[] DEFAULT NULL,
OUT p_retval INTEGER,
OUT p_errmsg TEXT,
OUT p_job_id BIGINT
)
RETURNS RECORD
LANGUAGE plpgsql
AS $$
BEGIN
SELECT r.p_retval, r.p_errmsg, r.p_job_id
INTO p_retval, p_errmsg, p_job_id
FROM broker.broker_add_job(
p_job_name,
p_execute_str,
1, -- p_job_queue
p_job_priority,
'sql', -- p_job_language
NULL, -- p_run_as
NULL, -- p_schedule_id
NULL, -- p_depends_on_job_ids
NULL, -- p_idempotency_key
1, -- p_max_attempts
NULL, -- p_job_group (defaults to p_job_name)
p_depends_on_groups
) AS r;
END;
$$;
COMMENT ON FUNCTION broker.broker_add_job_simple IS 'Shortcut for broker_add_job: name, execute string, priority, and dependencies by group name (defaults to job name)';
pkg/broker/install/sql/procedures/00_install.sql
-13
View File
@@ -1,13 +0,0 @@
-- PostgreSQL Broker Procedures Installation Script
-- Run this script to create all required stored procedures
\echo 'Installing PostgreSQL Broker procedures...'
\i 01_broker_get.sql
\i 02_broker_run.sql
\i 03_broker_set.sql
\i 04_broker_register_instance.sql
\i 05_broker_add_job.sql
\i 06_broker_ping_instance.sql
\echo 'PostgreSQL Broker procedures installed successfully!'
pkg/broker/install/sql/procedures/01_broker_get.sql
-76
View File
@@ -1,76 +0,0 @@
-- broker_get function
-- Fetches the next job from the queue for a given queue number
-- Returns: p_retval (0=success, >0=error), p_errmsg (error message), p_job_id (job ID if found)
CREATE OR REPLACE FUNCTION broker_get(
p_queue_number INTEGER,
p_instance_id BIGINT DEFAULT NULL,
OUT p_retval INTEGER,
OUT p_errmsg TEXT,
OUT p_job_id BIGINT
)
RETURNS RECORD
LANGUAGE plpgsql
AS $$
DECLARE
v_job_record RECORD;
BEGIN
p_retval := 0;
p_errmsg := '';
p_job_id := NULL;
-- Validate queue number
IF p_queue_number IS NULL OR p_queue_number <= 0 THEN
p_retval := 1;
p_errmsg := 'Invalid queue number';
RETURN;
END IF;
-- Find and lock the next pending job for this queue
-- Uses SKIP LOCKED to avoid blocking on jobs being processed by other workers
-- Skip jobs with pending dependencies
SELECT id_broker_jobs, job_name, job_priority, execute_str
INTO v_job_record
FROM broker_jobs
WHERE job_queue = p_queue_number
AND complete_status = 0 -- pending
AND (
depends_on IS NULL -- no dependencies
OR depends_on = '{}' -- empty dependencies
OR NOT EXISTS ( -- all dependencies completed
SELECT 1
FROM broker_jobs dep
WHERE dep.job_name = ANY(broker_jobs.depends_on)
AND dep.complete_status = 0 -- pending dependency
)
)
ORDER BY job_priority DESC, created_at ASC
LIMIT 1
FOR UPDATE SKIP LOCKED;
-- If no job found, return success with NULL job_id
IF NOT FOUND THEN
RETURN;
END IF;
-- Update job status to running
UPDATE broker_jobs
SET complete_status = 1, -- running
started_at = NOW(),
rid_broker_queueinstance = p_instance_id,
updated_at = NOW()
WHERE id_broker_jobs = v_job_record.id_broker_jobs;
-- Return the job ID
p_job_id := v_job_record.id_broker_jobs;
EXCEPTION
WHEN OTHERS THEN
p_retval := 2;
p_errmsg := SQLERRM;
RAISE WARNING 'broker_get error: %', SQLERRM;
END;
$$;
-- Comments
COMMENT ON FUNCTION broker_get IS 'Fetches the next pending job from the specified queue';
pkg/broker/install/sql/procedures/02_broker_run.sql
-113
View File
@@ -1,113 +0,0 @@
-- broker_run function
-- Executes a job by its ID
-- Returns: p_retval (0=success, >0=error), p_errmsg (error message)
CREATE OR REPLACE FUNCTION broker_run(
p_job_id BIGINT,
OUT p_retval INTEGER,
OUT p_errmsg TEXT
)
RETURNS RECORD
LANGUAGE plpgsql
AS $$
DECLARE
v_job_record RECORD;
v_execute_result TEXT;
v_error_occurred BOOLEAN := false;
BEGIN
p_retval := 0;
p_errmsg := '';
v_execute_result := '';
-- Validate job ID
IF p_job_id IS NULL OR p_job_id <= 0 THEN
p_retval := 1;
p_errmsg := 'Invalid job ID';
RETURN;
END IF;
-- Get job details
SELECT id_broker_jobs, execute_str, job_language, run_as, complete_status
INTO v_job_record
FROM broker_jobs
WHERE id_broker_jobs = p_job_id
FOR UPDATE;
-- Check if job exists
IF NOT FOUND THEN
p_retval := 2;
p_errmsg := 'Job not found';
RETURN;
END IF;
-- Check if job is in running state
IF v_job_record.complete_status != 1 THEN
p_retval := 3;
p_errmsg := format('Job is not in running state (status: %s)', v_job_record.complete_status);
RETURN;
END IF;
-- Execute the job
BEGIN
-- For SQL/PLPGSQL jobs, execute directly
IF v_job_record.job_language IN ('sql', 'plpgsql') THEN
EXECUTE v_job_record.execute_str;
v_execute_result := 'Success';
ELSE
-- Other languages would need external execution
p_retval := 4;
p_errmsg := format('Unsupported job language: %s', v_job_record.job_language);
v_error_occurred := true;
END IF;
EXCEPTION
WHEN OTHERS THEN
v_error_occurred := true;
p_retval := 5;
p_errmsg := SQLERRM;
v_execute_result := format('Error: %s', SQLERRM);
END;
-- Update job with results
IF v_error_occurred THEN
UPDATE broker_jobs
SET complete_status = 3, -- failed
error_msg = p_errmsg,
execute_result = v_execute_result,
completed_at = NOW(),
updated_at = NOW()
WHERE id_broker_jobs = p_job_id;
ELSE
UPDATE broker_jobs
SET complete_status = 2, -- completed
execute_result = v_execute_result,
error_msg = NULL,
completed_at = NOW(),
updated_at = NOW()
WHERE id_broker_jobs = p_job_id;
END IF;
EXCEPTION
WHEN OTHERS THEN
p_retval := 6;
p_errmsg := SQLERRM;
RAISE WARNING 'broker_run error: %', SQLERRM;
-- Try to update job status to failed
BEGIN
UPDATE broker_jobs
SET complete_status = 3, -- failed
error_msg = SQLERRM,
completed_at = NOW(),
updated_at = NOW()
WHERE id_broker_jobs = p_job_id;
EXCEPTION
WHEN OTHERS THEN
-- Ignore update errors
NULL;
END;
END;
$$;
-- Comments
COMMENT ON FUNCTION broker_run IS 'Executes a job by its ID and updates the status';
pkg/broker/install/sql/procedures/03_broker_set.sql
-95
View File
@@ -1,95 +0,0 @@
-- broker_set function
-- Sets broker runtime options and context
-- Supports: user, application_name, and custom settings
-- Returns: p_retval (0=success, >0=error), p_errmsg (error message)
CREATE OR REPLACE FUNCTION broker_set(
p_option_name TEXT,
p_option_value TEXT,
OUT p_retval INTEGER,
OUT p_errmsg TEXT
)
RETURNS RECORD
LANGUAGE plpgsql
AS $$
DECLARE
v_sql TEXT;
BEGIN
p_retval := 0;
p_errmsg := '';
-- Validate inputs
IF p_option_name IS NULL OR p_option_name = '' THEN
p_retval := 1;
p_errmsg := 'Option name is required';
RETURN;
END IF;
-- Handle different option types
CASE LOWER(p_option_name)
WHEN 'user' THEN
-- Set session user context
-- This is useful for audit trails and permissions
BEGIN
v_sql := format('SET SESSION AUTHORIZATION %I', p_option_value);
EXECUTE v_sql;
EXCEPTION
WHEN OTHERS THEN
p_retval := 2;
p_errmsg := format('Failed to set user: %s', SQLERRM);
RETURN;
END;
WHEN 'application_name' THEN
-- Set application name (visible in pg_stat_activity)
BEGIN
v_sql := format('SET application_name TO %L', p_option_value);
EXECUTE v_sql;
EXCEPTION
WHEN OTHERS THEN
p_retval := 3;
p_errmsg := format('Failed to set application_name: %s', SQLERRM);
RETURN;
END;
WHEN 'search_path' THEN
-- Set schema search path
BEGIN
v_sql := format('SET search_path TO %s', p_option_value);
EXECUTE v_sql;
EXCEPTION
WHEN OTHERS THEN
p_retval := 4;
p_errmsg := format('Failed to set search_path: %s', SQLERRM);
RETURN;
END;
WHEN 'timezone' THEN
-- Set timezone
BEGIN
v_sql := format('SET timezone TO %L', p_option_value);
EXECUTE v_sql;
EXCEPTION
WHEN OTHERS THEN
p_retval := 5;
p_errmsg := format('Failed to set timezone: %s', SQLERRM);
RETURN;
END;
ELSE
-- Unknown option
p_retval := 10;
p_errmsg := format('Unknown option: %s', p_option_name);
RETURN;
END CASE;
EXCEPTION
WHEN OTHERS THEN
p_retval := 99;
p_errmsg := SQLERRM;
RAISE WARNING 'broker_set error: %', SQLERRM;
END;
$$;
-- Comments
COMMENT ON FUNCTION broker_set IS 'Sets broker runtime options and session context (user, application_name, search_path, timezone)';
pkg/broker/install/sql/procedures/04_broker_register_instance.sql
-82
View File
@@ -1,82 +0,0 @@
-- broker_register_instance function
-- Registers a new broker instance in the database
-- Returns: p_retval (0=success, >0=error), p_errmsg (error message), p_instance_id (new instance ID)
CREATE OR REPLACE FUNCTION broker_register_instance(
p_name TEXT,
p_hostname TEXT,
p_pid INTEGER,
p_version TEXT,
p_queue_count INTEGER,
OUT p_retval INTEGER,
OUT p_errmsg TEXT,
OUT p_instance_id BIGINT
)
RETURNS RECORD
LANGUAGE plpgsql
AS $$
DECLARE
v_active_count INTEGER;
BEGIN
p_retval := 0;
p_errmsg := '';
p_instance_id := NULL;
-- Validate inputs
IF p_name IS NULL OR p_name = '' THEN
p_retval := 1;
p_errmsg := 'Instance name is required';
RETURN;
END IF;
IF p_hostname IS NULL OR p_hostname = '' THEN
p_retval := 2;
p_errmsg := 'Hostname is required';
RETURN;
END IF;
-- Check for existing active instances
-- Only one broker instance should be active per database
SELECT COUNT(*)
INTO v_active_count
FROM broker_queueinstance
WHERE status = 'active';
IF v_active_count > 0 THEN
p_retval := 3;
p_errmsg := 'Another broker instance is already active in this database. Only one broker instance per database is allowed.';
RETURN;
END IF;
-- Insert new instance
INSERT INTO broker_queueinstance (
name,
hostname,
pid,
version,
status,
queue_count,
started_at,
last_ping_at
) VALUES (
p_name,
p_hostname,
p_pid,
p_version,
'active',
p_queue_count,
NOW(),
NOW()
)
RETURNING id_broker_queueinstance INTO p_instance_id;
EXCEPTION
WHEN OTHERS THEN
p_retval := 99;
p_errmsg := SQLERRM;
RAISE WARNING 'broker_register_instance error: %', SQLERRM;
END;
$$;
-- Comments
COMMENT ON FUNCTION broker_register_instance IS 'Registers a new broker instance';
pkg/broker/install/sql/procedures/05_broker_add_job.sql
-91
View File
@@ -1,91 +0,0 @@
-- broker_add_job function
-- Adds a new job to the broker queue and sends a notification
-- Returns: p_retval (0=success, >0=error), p_errmsg (error message), p_job_id (new job ID)
CREATE OR REPLACE FUNCTION broker_add_job(
p_job_name TEXT,
p_execute_str TEXT,
p_job_queue INTEGER DEFAULT 1,
p_job_priority INTEGER DEFAULT 0,
p_job_language TEXT DEFAULT 'sql',
p_run_as TEXT DEFAULT NULL,
p_schedule_id BIGINT DEFAULT NULL,
p_depends_on TEXT[] DEFAULT NULL,
OUT p_retval INTEGER,
OUT p_errmsg TEXT,
OUT p_job_id BIGINT
)
RETURNS RECORD
LANGUAGE plpgsql
AS $$
DECLARE
v_notification_payload JSON;
BEGIN
p_retval := 0;
p_errmsg := '';
p_job_id := NULL;
-- Validate inputs
IF p_job_name IS NULL OR p_job_name = '' THEN
p_retval := 1;
p_errmsg := 'Job name is required';
RETURN;
END IF;
IF p_execute_str IS NULL OR p_execute_str = '' THEN
p_retval := 2;
p_errmsg := 'Execute string is required';
RETURN;
END IF;
IF p_job_queue IS NULL OR p_job_queue <= 0 THEN
p_retval := 3;
p_errmsg := 'Invalid job queue number';
RETURN;
END IF;
-- Insert new job
INSERT INTO broker_jobs (
job_name,
job_priority,
job_queue,
job_language,
execute_str,
run_as,
rid_broker_schedule,
depends_on,
complete_status
) VALUES (
p_job_name,
p_job_priority,
p_job_queue,
p_job_language,
p_execute_str,
p_run_as,
p_schedule_id,
p_depends_on,
0 -- pending
)
RETURNING id_broker_jobs INTO p_job_id;
-- Create notification payload
v_notification_payload := json_build_object(
'id', p_job_id,
'job_name', p_job_name,
'job_queue', p_job_queue,
'job_priority', p_job_priority
);
-- Send notification to broker
PERFORM pg_notify('broker.event', v_notification_payload::text);
EXCEPTION
WHEN OTHERS THEN
p_retval := 99;
p_errmsg := SQLERRM;
RAISE WARNING 'broker_add_job error: %', SQLERRM;
END;
$$;
-- Comments
COMMENT ON FUNCTION broker_add_job IS 'Adds a new job to the broker queue and sends a NOTIFY event';
pkg/broker/install/sql/roles/0001_roles.sql
+60
View File
@@ -0,0 +1,60 @@
-- Reference role/grant setup for pgsql-broker.
--
-- Applied via `pgsql-broker install --with-roles` (run once per cluster;
-- the schema/grant statements are safe to re-run per database). The
-- __BROKER_*_PASSWORD__ placeholders are substituted by the installer at
-- render time -- never edit this file to hardcode a real password. Each
-- CREATE ROLE is guarded so re-running this (e.g. against a second
-- configured database) rotates the password via ALTER ROLE instead of
-- failing on an already-existing role.
--
-- Roles:
-- broker_admin -- schema owner, runs migrations (`pgsql-broker install`).
-- Needs BYPASSRLS so broker_recover_stale_jobs (SECURITY
-- DEFINER, owned by this role) can sweep all tenants.
-- broker_runtime -- the role the running broker process connects as.
-- No BYPASSRLS, no ownership, SEARCH_PATH=broker so the
-- broker's unqualified table/function references resolve.
-- broker_enqueue -- narrow role for services that only need to add jobs.
DO $$
BEGIN
IF NOT EXISTS (SELECT FROM pg_roles WHERE rolname = 'broker_admin') THEN
CREATE ROLE broker_admin LOGIN PASSWORD __BROKER_ADMIN_PASSWORD__ BYPASSRLS;
ELSE
ALTER ROLE broker_admin WITH LOGIN PASSWORD __BROKER_ADMIN_PASSWORD__ BYPASSRLS;
END IF;
IF NOT EXISTS (SELECT FROM pg_roles WHERE rolname = 'broker_runtime') THEN
CREATE ROLE broker_runtime LOGIN PASSWORD __BROKER_RUNTIME_PASSWORD__;
ELSE
ALTER ROLE broker_runtime WITH LOGIN PASSWORD __BROKER_RUNTIME_PASSWORD__;
END IF;
IF NOT EXISTS (SELECT FROM pg_roles WHERE rolname = 'broker_enqueue') THEN
CREATE ROLE broker_enqueue LOGIN PASSWORD __BROKER_ENQUEUE_PASSWORD__;
ELSE
ALTER ROLE broker_enqueue WITH LOGIN PASSWORD __BROKER_ENQUEUE_PASSWORD__;
END IF;
END
$$;
ALTER ROLE broker_runtime SET search_path = broker, public;
ALTER ROLE broker_enqueue SET search_path = broker, public;
-- Run once broker.broker_jobs etc. already exist (i.e. after `pgsql-broker install`
-- as broker_admin), so schema ownership/grants land on the right objects.
ALTER SCHEMA broker OWNER TO broker_admin;
GRANT USAGE ON SCHEMA broker TO broker_runtime, broker_enqueue;
GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA broker TO broker_runtime;
GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA broker TO broker_runtime;
GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA broker TO broker_runtime;
REVOKE ALL ON ALL FUNCTIONS IN SCHEMA broker FROM broker_enqueue;
GRANT EXECUTE ON FUNCTION broker.broker_add_job TO broker_enqueue;
GRANT EXECUTE ON FUNCTION broker.broker_add_job_simple TO broker_enqueue;
GRANT EXECUTE ON FUNCTION broker.broker_set_tenant TO broker_enqueue;
GRANT INSERT, SELECT ON broker.broker_jobs, broker.broker_job_dependency TO broker_enqueue;
GRANT USAGE ON broker.broker_jobs_id_broker_jobs_seq TO broker_enqueue;
pkg/broker/install/sql/tables/00_install.sql
-10
View File
@@ -1,10 +0,0 @@
-- PostgreSQL Broker Tables Installation Script
-- Run this script to create all required tables
\echo 'Installing PostgreSQL Broker tables...'
\i 01_broker_queueinstance.sql
\i 02_broker_schedule.sql
\i 03_broker_jobs.sql
\echo 'PostgreSQL Broker tables installed successfully!'
pkg/broker/install/sql/tables/01_broker_queueinstance.sql
-31
View File
@@ -1,31 +0,0 @@
-- broker_queueinstance table
-- Tracks active and historical broker queue instances
CREATE TABLE IF NOT EXISTS broker_queueinstance (
id_broker_queueinstance BIGSERIAL PRIMARY KEY,
name VARCHAR(255) NOT NULL,
hostname VARCHAR(255) NOT NULL,
pid INTEGER NOT NULL,
version VARCHAR(50) NOT NULL,
status VARCHAR(20) NOT NULL DEFAULT 'active',
started_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
last_ping_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
shutdown_at TIMESTAMP WITH TIME ZONE,
queue_count INTEGER NOT NULL DEFAULT 0,
jobs_handled BIGINT NOT NULL DEFAULT 0,
CONSTRAINT broker_queueinstance_status_check CHECK (status IN ('active', 'inactive', 'shutdown'))
);
-- Indexes
CREATE INDEX IF NOT EXISTS idx_broker_queueinstance_status ON broker_queueinstance(status);
CREATE INDEX IF NOT EXISTS idx_broker_queueinstance_hostname ON broker_queueinstance(hostname);
CREATE INDEX IF NOT EXISTS idx_broker_queueinstance_last_ping ON broker_queueinstance(last_ping_at);
-- Comments
COMMENT ON TABLE broker_queueinstance IS 'Tracks broker queue instances (active and historical)';
COMMENT ON COLUMN broker_queueinstance.name IS 'Human-readable name of the broker instance';
COMMENT ON COLUMN broker_queueinstance.hostname IS 'Hostname where the broker is running';
COMMENT ON COLUMN broker_queueinstance.pid IS 'Process ID of the broker';
COMMENT ON COLUMN broker_queueinstance.status IS 'Current status: active, inactive, or shutdown';
COMMENT ON COLUMN broker_queueinstance.jobs_handled IS 'Total number of jobs handled by this instance';
pkg/broker/install/sql/tables/02_broker_schedule.sql
-50
View File
@@ -1,50 +0,0 @@
-- broker_schedule table
-- Stores scheduled jobs (cron-like functionality)
CREATE TABLE IF NOT EXISTS broker_schedule (
id_broker_schedule BIGSERIAL PRIMARY KEY,
name VARCHAR(255) NOT NULL UNIQUE,
cron_expr VARCHAR(100) NOT NULL,
enabled BOOLEAN NOT NULL DEFAULT true,
job_name VARCHAR(255) NOT NULL,
job_priority INTEGER NOT NULL DEFAULT 0,
job_queue INTEGER NOT NULL DEFAULT 1,
job_language VARCHAR(50) NOT NULL DEFAULT 'sql',
execute_str TEXT NOT NULL,
run_as VARCHAR(100),
created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
last_run_at TIMESTAMP WITH TIME ZONE,
next_run_at TIMESTAMP WITH TIME ZONE,
CONSTRAINT broker_schedule_job_queue_check CHECK (job_queue > 0)
);
-- Indexes
CREATE INDEX IF NOT EXISTS idx_broker_schedule_enabled ON broker_schedule(enabled);
CREATE INDEX IF NOT EXISTS idx_broker_schedule_next_run ON broker_schedule(next_run_at) WHERE enabled = true;
CREATE INDEX IF NOT EXISTS idx_broker_schedule_name ON broker_schedule(name);
-- Comments
COMMENT ON TABLE broker_schedule IS 'Scheduled jobs (cron-like functionality)';
COMMENT ON COLUMN broker_schedule.name IS 'Unique name for the schedule';
COMMENT ON COLUMN broker_schedule.cron_expr IS 'Cron expression for scheduling';
COMMENT ON COLUMN broker_schedule.enabled IS 'Whether the schedule is active';
COMMENT ON COLUMN broker_schedule.job_name IS 'Name of the job to create';
COMMENT ON COLUMN broker_schedule.execute_str IS 'SQL or code to execute';
COMMENT ON COLUMN broker_schedule.last_run_at IS 'Last time the job was executed';
COMMENT ON COLUMN broker_schedule.next_run_at IS 'Next scheduled execution time';
-- Trigger to update updated_at
CREATE OR REPLACE FUNCTION tf_broker_schedule_update_timestamp()
RETURNS TRIGGER AS $$
BEGIN
NEW.updated_at = NOW();
RETURN NEW;
END;
$$ LANGUAGE plpgsql;
CREATE TRIGGER t_broker_schedule_updated_at
BEFORE UPDATE ON broker_schedule
FOR EACH ROW
EXECUTE FUNCTION tf_broker_schedule_update_timestamp();
pkg/broker/install/sql/tables/03_broker_jobs.sql
-62
View File
@@ -1,62 +0,0 @@
-- broker_jobs table
-- Stores jobs to be executed by the broker
CREATE TABLE IF NOT EXISTS broker_jobs (
id_broker_jobs BIGSERIAL PRIMARY KEY,
job_name VARCHAR(255) NOT NULL,
job_priority INTEGER NOT NULL DEFAULT 0,
job_queue INTEGER NOT NULL DEFAULT 1,
job_language VARCHAR(50) NOT NULL DEFAULT 'sql',
execute_str TEXT NOT NULL,
execute_result TEXT,
error_msg TEXT,
complete_status INTEGER NOT NULL DEFAULT 0,
run_as VARCHAR(100),
rid_broker_schedule BIGINT,
rid_broker_queueinstance BIGINT,
depends_on TEXT[],
created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
started_at TIMESTAMP WITH TIME ZONE,
completed_at TIMESTAMP WITH TIME ZONE,
CONSTRAINT broker_jobs_complete_status_check CHECK (complete_status IN (0, 1, 2, 3, 4)),
CONSTRAINT broker_jobs_job_queue_check CHECK (job_queue > 0),
CONSTRAINT fk_schedule FOREIGN KEY (rid_broker_schedule) REFERENCES broker_schedule(id_broker_schedule) ON DELETE SET NULL,
CONSTRAINT fk_instance FOREIGN KEY (rid_broker_queueinstance) REFERENCES broker_queueinstance(id_broker_queueinstance) ON DELETE SET NULL
);
-- Indexes
CREATE INDEX IF NOT EXISTS idx_broker_jobs_status ON broker_jobs(complete_status);
CREATE INDEX IF NOT EXISTS idx_broker_jobs_queue ON broker_jobs(job_queue, complete_status, job_priority);
CREATE INDEX IF NOT EXISTS idx_broker_jobs_schedule ON broker_jobs(rid_broker_schedule);
CREATE INDEX IF NOT EXISTS idx_broker_jobs_instance ON broker_jobs(rid_broker_queueinstance);
CREATE INDEX IF NOT EXISTS idx_broker_jobs_created ON broker_jobs(created_at);
CREATE INDEX IF NOT EXISTS idx_broker_jobs_name ON broker_jobs(job_name, complete_status);
-- Comments
COMMENT ON TABLE broker_jobs IS 'Job queue for broker execution';
COMMENT ON COLUMN broker_jobs.job_name IS 'Name/description of the job';
COMMENT ON COLUMN broker_jobs.job_priority IS 'Job priority (higher = more important)';
COMMENT ON COLUMN broker_jobs.job_queue IS 'Queue number (allows parallel processing)';
COMMENT ON COLUMN broker_jobs.job_language IS 'Execution language (sql, plpgsql, etc.)';
COMMENT ON COLUMN broker_jobs.execute_str IS 'SQL or code to execute';
COMMENT ON COLUMN broker_jobs.complete_status IS '0=pending, 1=running, 2=completed, 3=failed, 4=cancelled';
COMMENT ON COLUMN broker_jobs.run_as IS 'User context to run the job as';
COMMENT ON COLUMN broker_jobs.rid_broker_schedule IS 'Reference to schedule if job was scheduled';
COMMENT ON COLUMN broker_jobs.rid_broker_queueinstance IS 'Instance that processed this job';
COMMENT ON COLUMN broker_jobs.depends_on IS 'Array of job names that must be completed before this job can run';
-- Trigger to update updated_at
CREATE OR REPLACE FUNCTION tf_broker_jobs_update_timestamp()
RETURNS TRIGGER AS $$
BEGIN
NEW.updated_at = NOW();
RETURN NEW;
END;
$$ LANGUAGE plpgsql;
CREATE TRIGGER t_broker_jobs_updated_at
BEFORE UPDATE ON broker_jobs
FOR EACH ROW
EXECUTE FUNCTION tf_broker_jobs_update_timestamp();
pkg/broker/models/models.go
+14
View File
@@ -16,10 +16,24 @@ type Job struct {
RunAs string `json:"run_as"`
UserLogin string `json:"user_login"`
ScheduleID int64 `json:"schedule_id"`
TenantID string `json:"tenant_id"`
AttemptCount int `json:"attempt_count"`
MaxAttempts int `json:"max_attempts"`
LeaseToken string `json:"lease_token,omitempty"`
IdempotencyKey string `json:"idempotency_key,omitempty"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
}
// WakeNotification is the payload sent over pg_notify('broker.event', ...).
// It carries only what a worker needs to decide whether to wake: the queue
// number. job_id is included solely for logging -- workers always re-claim
// via broker_get rather than executing the notified id directly.
type WakeNotification struct {
Queue int `json:"queue"`
JobID int64 `json:"job_id,omitempty"`
}
// Instance represents a broker instance
type Instance struct {
ID int64 `json:"id"`
pkg/broker/queue/queue.go
+9 -14
View File
@@ -6,7 +6,6 @@ import (
"sync"
"git.warky.dev/wdevs/pgsql-broker/pkg/broker/adapter"
"git.warky.dev/wdevs/pgsql-broker/pkg/broker/models"
"git.warky.dev/wdevs/pgsql-broker/pkg/broker/worker"
)
@@ -33,6 +32,8 @@ type Config struct {
BufferSize int
TimerSeconds int
FetchSize int
TenantID string
LeaseSeconds int
}
// New creates a new queue manager
@@ -70,6 +71,8 @@ func (q *Queue) Start(cfg Config) error {
BufferSize: cfg.BufferSize,
TimerSeconds: cfg.TimerSeconds,
FetchSize: cfg.FetchSize,
TenantID: cfg.TenantID,
LeaseSeconds: cfg.LeaseSeconds,
})
if err := w.Start(q.ctx); err != nil {
@@ -109,24 +112,16 @@ func (q *Queue) Stop() error {
return nil
}
// AddJob adds a job to the least busy worker
func (q *Queue) AddJob(job models.Job) error {
// Wake signals every worker in the queue to check for available jobs
// immediately. There is no job hand-off: fetching is always DB-driven via
// broker_get, so waking a worker that finds nothing is harmless.
func (q *Queue) Wake() {
q.mu.RLock()
defer q.mu.RUnlock()
if len(q.workers) == 0 {
return fmt.Errorf("no workers available")
}
// Simple round-robin: use first available worker
// Could be enhanced with load balancing
for _, w := range q.workers {
if err := w.AddJob(job); err == nil {
return nil
}
w.Wake()
}
return fmt.Errorf("all workers are busy")
}
// GetStats returns statistics for all workers in the queue
pkg/broker/worker/worker.go
+113 -50
View File
@@ -2,13 +2,13 @@ package worker
import (
"context"
"database/sql" // Import sql package
"database/sql"
"fmt"
"runtime/debug"
"sync"
"time"
"git.warky.dev/wdevs/pgsql-broker/pkg/broker/adapter"
"git.warky.dev/wdevs/pgsql-broker/pkg/broker/models"
)
// Worker represents a single job processing worker
@@ -18,7 +18,7 @@ type Worker struct {
InstanceID int64
db adapter.DBAdapter
logger adapter.Logger
jobChan chan models.Job
wakeChan chan struct{}
shutdown chan struct{}
wg *sync.WaitGroup
running bool
@@ -27,6 +27,8 @@ type Worker struct {
jobsHandled int64
timerSeconds int
fetchSize int
tenantID string
leaseSeconds int
}
// Stats holds worker statistics
@@ -46,21 +48,30 @@ type Config struct {
BufferSize int
TimerSeconds int
FetchSize int
TenantID string
LeaseSeconds int
}
// New creates a new worker
func New(cfg Config) *Worker {
leaseSeconds := cfg.LeaseSeconds
if leaseSeconds <= 0 {
leaseSeconds = 60
}
return &Worker{
ID: cfg.ID,
QueueNumber: cfg.QueueNumber,
InstanceID: cfg.InstanceID,
db: cfg.DBAdapter,
logger: cfg.Logger.With("worker_id", cfg.ID).With("queue", cfg.QueueNumber),
jobChan: make(chan models.Job, cfg.BufferSize),
wakeChan: make(chan struct{}, 1),
shutdown: make(chan struct{}),
wg: &sync.WaitGroup{},
timerSeconds: cfg.TimerSeconds,
fetchSize: cfg.FetchSize,
tenantID: cfg.TenantID,
leaseSeconds: leaseSeconds,
}
}
@@ -77,11 +88,48 @@ func (w *Worker) Start(ctx context.Context) error {
w.logger.Info("worker starting")
w.wg.Add(1)
go w.processLoop(ctx)
go w.superviseProcessLoop(ctx)
return nil
}
// superviseProcessLoop runs processLoop for the life of the worker,
// restarting it (after a short backoff) if it ever panics, so a bug in job
// processing can never permanently kill this worker's goroutine.
func (w *Worker) superviseProcessLoop(ctx context.Context) {
defer w.wg.Done()
for {
select {
case <-w.shutdown:
return
case <-ctx.Done():
return
default:
}
if w.runProcessLoopOnce(ctx) {
return
}
time.Sleep(time.Second)
}
}
// runProcessLoopOnce runs processLoop, recovering any panic. It returns true
// if processLoop returned normally (shutdown/context done, no restart
// needed) and false if it panicked (caller should restart it).
func (w *Worker) runProcessLoopOnce(ctx context.Context) (clean bool) {
defer func() {
if r := recover(); r != nil {
w.logger.Error("worker panic recovered, restarting", "panic", r, "stack", string(debug.Stack()))
clean = false
}
}()
w.processLoop(ctx)
return true
}
// Stop gracefully stops the worker
func (w *Worker) Stop() error {
w.mu.Lock()
@@ -103,35 +151,33 @@ func (w *Worker) Stop() error {
return nil
}
// AddJob adds a job to the worker's queue
func (w *Worker) AddJob(job models.Job) error {
// Wake signals the worker to check for available jobs immediately, instead
// of waiting for the next timer tick. Fetching is always DB-driven (via
// broker_get with FOR UPDATE SKIP LOCKED), so a redundant or coalesced wake
// is harmless.
func (w *Worker) Wake() {
select {
case w.jobChan <- job:
return nil
case w.wakeChan <- struct{}{}:
default:
return fmt.Errorf("worker %d job channel is full", w.ID)
}
}
// processLoop is the main worker processing loop
func (w *Worker) processLoop(ctx context.Context) {
defer w.wg.Done()
defer w.recoverPanic()
timer := time.NewTimer(time.Duration(w.timerSeconds) * time.Second)
defer timer.Stop()
for {
select {
case job := <-w.jobChan:
case <-w.wakeChan:
w.updateActivity()
w.processJobs(ctx, &job)
w.processJobs(ctx)
case <-timer.C:
// Timer expired - fetch jobs from database
if w.timerSeconds > 0 {
w.updateActivity()
w.processJobs(ctx, nil)
w.processJobs(ctx)
}
timer.Reset(time.Duration(w.timerSeconds) * time.Second)
@@ -147,84 +193,101 @@ func (w *Worker) processLoop(ctx context.Context) {
}
// processJobs processes jobs from the queue within a transaction
func (w *Worker) processJobs(ctx context.Context, specificJob *models.Job) {
func (w *Worker) processJobs(ctx context.Context) {
defer w.recoverPanic()
for i := 0; i < w.fetchSize; i++ {
tx, err := w.db.Begin(ctx) // Start transaction
tx, err := w.db.Begin(ctx)
if err != nil {
w.logger.Error("failed to begin transaction", "error", err)
return
}
var jobID int64
if err := w.setTenantTx(ctx, tx); err != nil {
tx.Rollback()
w.logger.Error("failed to set tenant", "error", err)
return
}
if specificJob != nil && specificJob.ID > 0 {
jobID = specificJob.ID
specificJob = nil // Only process once
} else {
jobID, err = w.fetchNextJobTx(ctx, tx) // Use transaction
if err != nil {
tx.Rollback() // Rollback on fetch error
w.logger.Error("failed to fetch job", "error", err)
return
}
jobID, leaseToken, err := w.fetchNextJobTx(ctx, tx)
if err != nil {
tx.Rollback()
w.logger.Error("failed to fetch job", "error", err)
return
}
if jobID <= 0 {
tx.Rollback() // No job found, rollback
return // No more jobs
return // No more jobs
}
// Run the job
if err := w.runJobTx(ctx, tx, jobID); err != nil { // Use transaction
tx.Rollback() // Rollback on job execution error
if err := w.runJobTx(ctx, tx, jobID, leaseToken); err != nil {
tx.Rollback() // Rollback on genuine infra failure
w.logger.Error("failed to run job", "job_id", jobID, "error", err)
} else {
tx.Commit() // Commit if job successful
if err := tx.Commit(); err != nil {
w.logger.Error("failed to commit job", "job_id", jobID, "error", err)
continue
}
w.jobsHandled++
}
}
}
// fetchNextJobTx fetches the next job from the queue within a transaction
func (w *Worker) fetchNextJobTx(ctx context.Context, tx adapter.DBTransaction) (int64, error) {
// setTenantTx applies this worker's RLS tenant for the duration of tx.
func (w *Worker) setTenantTx(ctx context.Context, tx adapter.DBTransaction) error {
tenantID := w.tenantID
if tenantID == "" {
tenantID = "default"
}
_, err := tx.Exec(ctx, "SELECT broker.broker_set_tenant($1)", tenantID)
return err
}
// fetchNextJobTx fetches the next job from the queue within a transaction,
// claiming it with a lease that must be presented back to broker_run.
func (w *Worker) fetchNextJobTx(ctx context.Context, tx adapter.DBTransaction) (int64, string, error) {
var retval int
var errmsg string
var nullableJobID sql.NullInt64
var nullableLeaseToken sql.NullString
err := tx.QueryRow(ctx,
"SELECT p_retval, p_errmsg, p_job_id FROM broker_get($1, $2)",
w.QueueNumber, w.InstanceID,
).Scan(&retval, &errmsg, &nullableJobID)
"SELECT p_retval, p_errmsg, p_job_id, p_lease_token FROM broker.broker_get($1, $2, $3)",
w.QueueNumber, w.InstanceID, w.leaseSeconds,
).Scan(&retval, &errmsg, &nullableJobID, &nullableLeaseToken)
if err != nil {
return 0, fmt.Errorf("query error: %w", err)
return 0, "", fmt.Errorf("query error: %w", err)
}
if retval > 0 {
return 0, fmt.Errorf("broker_get error: %s", errmsg)
return 0, "", fmt.Errorf("broker_get error: %s", errmsg)
}
if !nullableJobID.Valid {
return 0, nil
return 0, "", nil
}
return nullableJobID.Int64, nil
return nullableJobID.Int64, nullableLeaseToken.String, nil
}
// runJobTx executes a job within a transaction
func (w *Worker) runJobTx(ctx context.Context, tx adapter.DBTransaction, jobID int64) error {
// runJobTx executes a leased job within a transaction. It only returns an
// error (triggering a rollback of the claim) on a genuine infra failure --
// job outcomes reported via p_job_status (requeued/completed/dead-lettered)
// are always committed.
func (w *Worker) runJobTx(ctx context.Context, tx adapter.DBTransaction, jobID int64, leaseToken string) error {
w.logger.Debug("running job", "job_id", jobID)
var retval int
var errmsg string
var jobStatus int
err := tx.QueryRow(ctx,
"SELECT p_retval, p_errmsg FROM broker_run($1)",
jobID,
).Scan(&retval, &errmsg)
"SELECT p_retval, p_errmsg, p_job_status FROM broker.broker_run($1, $2)",
jobID, leaseToken,
).Scan(&retval, &errmsg, &jobStatus)
if err != nil {
return fmt.Errorf("query error: %w", err)
@@ -234,7 +297,7 @@ func (w *Worker) runJobTx(ctx context.Context, tx adapter.DBTransaction, jobID i
return fmt.Errorf("broker_run error: %s", errmsg)
}
w.logger.Debug("job completed", "job_id", jobID)
w.logger.Debug("job finished", "job_id", jobID, "job_status", jobStatus)
return nil
}
@@ -255,6 +318,6 @@ func (w *Worker) GetStats() (lastActivity time.Time, jobsHandled int64, running
// recoverPanic recovers from panics in the worker
func (w *Worker) recoverPanic() {
if r := recover(); r != nil {
w.logger.Error("worker panic recovered", "panic", r)
w.logger.Error("worker panic recovered", "panic", r, "stack", string(debug.Stack()))
}
}