diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..32dc051 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,6 @@ +.git +bin +broker.log +broker.pid +plan +*.md diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..fb0b595 --- /dev/null +++ b/.env.example @@ -0,0 +1,11 @@ +# Copy to .env (gitignored) and fill in real values. +# POSTGRES_PASSWORD is the postgres superuser password (used only by the +# one-shot `migrate` service to run `install --with-roles`). +# The three BROKER_*_PASSWORD values must match the passwords you put in +# broker.docker.yaml (copied from broker.docker.example.yaml) for the +# corresponding role -- BROKER_RUNTIME_PASSWORD in particular must match +# the `password:` field the broker service itself connects with. +POSTGRES_PASSWORD=change_me +BROKER_ADMIN_PASSWORD=change_me +BROKER_RUNTIME_PASSWORD=change_me +BROKER_ENQUEUE_PASSWORD=change_me diff --git a/.gitignore b/.gitignore index d98dac4..f435bc3 100644 --- a/.gitignore +++ b/.gitignore @@ -32,7 +32,9 @@ bin/ broker.yaml broker.yml broker.json +broker.docker.yaml !broker.example.yaml +!broker.docker.example.yaml # IDE .vscode/ diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..3d57cac --- /dev/null +++ b/Dockerfile @@ -0,0 +1,40 @@ +# syntax=docker/dockerfile:1 + +# ---- Build stage ---- +FROM golang:1.26-alpine AS builder + +WORKDIR /src + +COPY go.mod go.sum ./ +RUN go mod download + +COPY . . + +ARG VERSION=dev +ARG COMMIT=unknown +ARG BUILD_TIME=unknown + +RUN CGO_ENABLED=0 GOOS=linux go build \ + -trimpath \ + -ldflags "-w -s -X main.Version=${VERSION} -X main.Commit=${COMMIT} -X main.BuildTime=${BUILD_TIME}" \ + -o /out/pgsql-broker \ + ./cmd/broker + +# ---- Runtime stage ---- +FROM alpine:3.22 + +RUN apk add --no-cache ca-certificates tzdata && \ + addgroup -S broker && adduser -S -G broker -h /home/broker broker && \ + mkdir -p /etc/pgsql-broker && \ + chown -R broker:broker /etc/pgsql-broker + +COPY --from=builder /out/pgsql-broker /usr/local/bin/pgsql-broker + +USER broker +WORKDIR /home/broker + +# broker.yaml is expected to be mounted at /etc/pgsql-broker/broker.yaml +# (config.LoadConfig's default search path) -- no config is baked into the +# image, since it holds database credentials. +ENTRYPOINT ["pgsql-broker"] +CMD ["start"] diff --git a/README.md b/README.md index 2c8df75..81f63cd 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # PostgreSQL Broker -A robust, event-driven job processing system for PostgreSQL that uses LISTEN/NOTIFY for real-time job execution. It supports multiple queues, priority-based scheduling, and can be used both as a standalone service or as a Go library. +A robust, event-driven job processing system for PostgreSQL that uses LISTEN/NOTIFY for real-time job execution. It supports multiple queues, priority-based scheduling, multi-tenant row-level security, and can be used both as a standalone service or as a Go library. ## Features @@ -8,20 +8,23 @@ A robust, event-driven job processing system for PostgreSQL that uses LISTEN/NOT - **Event-Driven**: Uses PostgreSQL LISTEN/NOTIFY for instant job notifications - **Multiple Queues**: Support for concurrent job processing across multiple queues per database - **Priority Scheduling**: Jobs can be prioritized for execution order -- **Job Dependencies**: Jobs can depend on other jobs being completed first +- **Job Dependencies**: Jobs can depend on other jobs completing first (cycle-checked, gated in `broker_get`) +- **Multi-Tenant Row-Level Security**: Jobs are isolated per `tenant_id` via FORCE RLS policies (`broker.broker_set_tenant`) +- **Lease-Based Claiming**: Jobs are claimed with an expiring lease and recovered automatically if a worker dies mid-job +- **Least-Privilege Roles**: Optional `broker_admin` / `broker_runtime` / `broker_enqueue` role separation +- **Versioned Migrations**: Embedded, ordered SQL migrations with an applied-versions tracking table - **Adapter Pattern**: Clean interfaces for database and logging (easy to extend) - **Standalone or Library**: Use as a CLI tool or integrate into your Go application - **Configuration Management**: Viper-based config with support for YAML, JSON, and environment variables - **Graceful Shutdown**: Proper cleanup and job completion on shutdown -- **Instance Tracking**: Monitor active broker instances through the database - **Single Instance Per Database**: Enforces one broker instance per database to prevent conflicts -- **Embedded SQL Installer**: Database schema embedded in binary with built-in install command +- **Docker Support**: Production `Dockerfile`/`docker-compose.yml` and a test-runner `Dockerfile.test`/`docker-compose.test.yml` ## Architecture The broker supports multi-database architecture where a single broker process can manage multiple database connections. Each database has its own instance with dedicated queues, but only ONE broker instance is allowed per database. -``` +```text ┌─────────────────────────────────────────────────────────────────┐ │ Broker Process │ ├─────────────────────────────────────────────────────────────────┤ @@ -46,17 +49,20 @@ The broker supports multi-database architecture where a single broker process ca ┌──────────────────────┐ ┌──────────────────────┐ │ PostgreSQL (DB1) │ │ PostgreSQL (DB2) │ │ - broker_jobs │ │ - broker_jobs │ + │ - broker_job_dependency│ │ - broker_job_dependency│ │ - broker_queueinstance│ │ - broker_queueinstance│ │ - broker_schedule │ │ - broker_schedule │ └──────────────────────┘ └──────────────────────┘ ``` **Key Points**: + - One broker process can manage multiple databases - Each database has exactly ONE active broker instance - Each database instance has its own queues and workers - Validation prevents multiple broker processes from connecting to the same database - Different databases can have different queue counts +- Jobs are isolated by `tenant_id` via FORCE RLS; `broker_get` only claims jobs with no incomplete dependency ## Installation @@ -76,27 +82,57 @@ The binary will be available in `bin/pgsql-broker`. go get git.warky.dev/wdevs/pgsql-broker ``` +### With Docker + +See [Docker](#docker) below for the production `Dockerfile`/`docker-compose.yml` and the `Dockerfile.test`/`docker-compose.test.yml` test runner. + ## Quick Start ### 1. Setup Database -Install the required tables and stored procedures: +Migrations are embedded in the binary and applied in order, tracked in `broker.broker_schema_migrations`. ```bash -# Using the CLI (recommended) +# Apply all pending migrations, then verify ./bin/pgsql-broker install --config broker.yaml +# Verify only, no changes +./bin/pgsql-broker install --verify-only --config broker.yaml + # Or with make make sql-install - -# Verify installation -./bin/pgsql-broker install --verify-only --config broker.yaml - -# Or manually with psql: -psql -f pkg/broker/install/sql/tables/00_install.sql -psql -f pkg/broker/install/sql/procedures/00_install.sql ``` +#### Optional: least-privilege roles + +`install --with-roles` additionally creates/rotates three roles (idempotent — safe to re-run): + +| Role | Purpose | +| ---- | ------- | +| `broker_admin` | Schema owner, `BYPASSRLS`, runs migrations | +| `broker_runtime` | The role the running broker connects as; table/sequence/function privileges only, no `BYPASSRLS` | +| `broker_enqueue` | Narrow role for enqueuing jobs only (`broker_add_job`, `broker_set_tenant`, insert on `broker_jobs`/`broker_job_dependency`) | + +```bash +./bin/pgsql-broker install --with-roles \ + --admin-user postgres \ + --broker-admin-password ... \ + --broker-runtime-password ... \ + --broker-enqueue-password ... +``` + +Credential resolution order for each value: CLI flag → environment variable → interactive masked prompt (TTY only) → error. + +| Value | Flag | Env fallback(s) | +| ----- | ---- | ---------------- | +| Admin user | `--admin-user` | `PGUSER`, `PG_USER` | +| Admin password | `--admin-password` | `PGPASSWORD`, `PG_PASS` | +| `broker_admin` password | `--broker-admin-password` | `BROKER_ADMIN_PASSWORD` | +| `broker_runtime` password | `--broker-runtime-password` | `BROKER_RUNTIME_PASSWORD` | +| `broker_enqueue` password | `--broker-enqueue-password` | `BROKER_ENQUEUE_PASSWORD` | + +`--with-roles` always connects as the admin login (not the config file's own user) since `broker_admin` must own the schema; it cannot be combined with `--verify-only`. + ### 2. Configure Create a configuration file `broker.yaml`: @@ -111,6 +147,7 @@ databases: password: your_password sslmode: disable queue_count: 4 + tenant_id: default # optional; sets broker.tenant_id for RLS # Optional: add more databases - name: db2 @@ -125,13 +162,15 @@ databases: broker: name: pgsql-broker enable_debug: false + lease_seconds: 60 # job lease duration before it's reclaimable + stale_job_recovery_sec: 30 # how often expired leases are recovered logging: level: info format: json ``` -**Note**: Each database requires a unique `name` identifier and can have its own `queue_count` configuration. +**Note**: Each database requires a unique `name` identifier and can have its own `queue_count` configuration. See [broker.example.yaml](./broker.example.yaml) for the full field set. ### 3. Run the Broker @@ -149,18 +188,47 @@ make run ### 4. Add a Job ```sql -SELECT broker_add_job( - 'My Job', -- job_name - 'SELECT do_something()', -- execute_str - 1, -- job_queue (default: 1) - 0, -- job_priority (default: 0) - 'sql', -- job_language (default: 'sql') - NULL, -- run_as - NULL, -- user_login - NULL -- schedule_id +SELECT * FROM broker.broker_add_job( + 'My Job', -- p_job_name + 'SELECT do_something()', -- p_execute_str + 1, -- p_job_queue (default: 1) + 0, -- p_job_priority (default: 0) + 'sql', -- p_job_language (default: 'sql') + NULL, -- p_run_as + NULL, -- p_schedule_id + ARRAY[123]::BIGINT[], -- p_depends_on_job_ids (NULL = no deps) + NULL, -- p_idempotency_key + 1, -- p_max_attempts (default: 1) + NULL, -- p_job_group (default: p_job_name) + ARRAY['extract'] -- p_depends_on_groups (NULL = no deps) ); ``` +Or, for the common case (name, execute string, priority, dependencies), use the shortcut: + +```sql +SELECT * FROM broker.broker_add_job_simple('transform', 'SELECT do_something()', 5, ARRAY['extract']); +``` + +#### Dependencies: job groups vs. job ids + +Every job belongs to a **job group**, defaulting to its own `job_name` when `p_job_group` isn't given. A dependency can target either a whole group (`p_depends_on_groups`) or a single job id (`p_depends_on_job_ids`): + +- **Group dependency** (fan-in): the dependent job is only claimable once **every** job currently tagged with that group has `complete_status = 2`. Multiple jobs can share a group (e.g. several parallel `extract` jobs feeding one `transform`); already-completed group members simply drop out of the check. +- **Id dependency**: waits on one specific job row, as before. + +Both `broker_add_job` and `broker_add_job_simple` reject a job depending on its own group/id, and a one-hop cycle (A depends on B, B already depends on A). + +When calling as `broker_enqueue` (or any RLS-scoped role), set the tenant first and wrap in a single transaction so `SET LOCAL` persists across statements: + +```sql +BEGIN; +SET ROLE broker_enqueue; +SELECT broker.broker_set_tenant('my_tenant'); +SELECT * FROM broker.broker_add_job_simple('extract', 'SELECT 1'); +COMMIT; +``` + ## Usage as a Library ```go @@ -194,34 +262,65 @@ func main() { See the [examples](./examples/) directory for complete examples. +## Docker + +### Production: `Dockerfile` + `docker-compose.yml` + +Multi-stage build (`golang:1.26-alpine` → `alpine:3.22`, static binary, non-root user). The compose stack runs Postgres, a one-shot `migrate` service (`install --with-roles`), then the `broker` service connecting as `broker_runtime`. + +```bash +cp broker.docker.example.yaml broker.docker.yaml # set broker_runtime password +cp .env.example .env # set POSTGRES_PASSWORD + 3 BROKER_*_PASSWORD +docker-compose up --build -d +``` + +`broker.docker.yaml` and `.env` are gitignored — never commit real credentials. `BROKER_RUNTIME_PASSWORD` in `.env` must match the `password:` field in `broker.docker.yaml`. + +### Test runner: `Dockerfile.test` + `docker-compose.test.yml` + +Builds the module and runs the full Go integration suite (`tests/integration/...`) against a disposable Postgres container: + +```bash +docker-compose -f docker-compose.test.yml up --build --abort-on-container-exit --exit-code-from tests +docker-compose -f docker-compose.test.yml down -v +``` + ## Database Schema +Schema objects live under `broker.*` and are applied via ordered, embedded migrations (`pkg/broker/install/sql/migrations/`), tracked in `broker.broker_schema_migrations`. + ### Tables - **broker_queueinstance**: Tracks active broker queue instances (one per database) -- **broker_jobs**: Job queue with status tracking +- **broker_jobs**: Job queue with status, lease, attempt, tenant, and `job_group` tracking +- **broker_job_dependency**: Dependency edges, each targeting either a `depends_on_job_id` or a `depends_on_group`, enforced at claim time - **broker_schedule**: Scheduled jobs (cron-like functionality) ### Stored Procedures -- **broker_get**: Fetch the next job from a queue +- **broker_get**: Fetch and lease the next claimable job from a queue (skips jobs with an incomplete id- or group-dependency) - **broker_run**: Execute a job - **broker_set**: Set runtime options (user, application_name, etc.) -- **broker_add_job**: Add a new job to the queue -- **broker_register_instance**: Register a broker instance -- **broker_ping_instance**: Update instance heartbeat -- **broker_shutdown_instance**: Mark instance as shutdown +- **broker_set_tenant**: Set the `broker.tenant_id` GUC used by RLS policies +- **broker_add_job**: Add a new job to the queue, with an optional `job_group` and dependencies by id and/or group +- **broker_add_job_simple**: Shortcut for `broker_add_job` — name, execute string, priority, and dependencies by group name +- **broker_register_instance / broker_ping_instance / broker_shutdown_instance**: Instance lifecycle tracking +- **broker_recover_stale_jobs**: Reclaim jobs whose lease has expired + +### Roles (optional, via `install --with-roles`) + +See `pkg/broker/install/sql/roles/0001_roles.sql` — `broker_admin`, `broker_runtime`, `broker_enqueue` (see table above). ## Configuration Reference -See [broker.example.yaml](./broker.example.yaml) for a complete configuration example. +See [broker.example.yaml](./broker.example.yaml) for a complete configuration example, or [broker.docker.example.yaml](./broker.docker.example.yaml) for the Docker Compose variant. ### Database Settings The `databases` array can contain multiple database configurations. Each entry supports: | Setting | Description | Default | -|---------|-------------|---------| +| ------- | ----------- | ------- | | `name` | Unique identifier for this database | **Required** | | `host` | PostgreSQL host | `localhost` | | `port` | PostgreSQL port | `5432` | @@ -234,13 +333,14 @@ The `databases` array can contain multiple database configurations. Each entry s | `conn_max_lifetime` | Connection max lifetime | `5m` | | `conn_max_idle_time` | Connection max idle time | `10m` | | `queue_count` | Number of queues for this database | `4` | +| `tenant_id` | Tenant id set for this connection (RLS) | `default` | ### Broker Settings Global settings applied to all database instances: | Setting | Description | Default | -|---------|-------------|---------| +| ------- | ----------- | ------- | | `name` | Broker instance name | `pgsql-broker` | | `fetch_query_que_size` | Jobs per fetch cycle | `100` | | `queue_timer_sec` | Seconds between polls | `10` | @@ -248,6 +348,8 @@ Global settings applied to all database instances: | `worker_idle_timeout_sec` | Worker idle timeout | `10` | | `notify_retry_seconds` | NOTIFY retry interval | `30s` | | `enable_debug` | Enable debug logging | `false` | +| `lease_seconds` | Job lease duration before reclaimable | - | +| `stale_job_recovery_sec` | Interval for reclaiming expired leases | - | ## Development @@ -262,9 +364,19 @@ make vet # Run go vet make test # Run tests ``` +### Testing + +```bash +go test ./... # unit tests +go test -v ./tests/integration/... # integration tests (needs local Postgres, see below) +docker-compose -f docker-compose.test.yml up --build --abort-on-container-exit --exit-code-from tests +``` + +Integration tests expect Postgres reachable at `localhost:5433` (see `tests/integration/`), including `rls_test.go` (multi-tenant isolation) and `stage5_test.go`. + ### Project Structure -``` +```text pgsql-broker/ ├── cmd/broker/ # CLI application ├── pkg/broker/ # Core broker package @@ -274,20 +386,26 @@ pgsql-broker/ │ ├── queue/ # Queue management │ ├── worker/ # Worker implementation │ └── install/ # Database installer with embedded SQL -│ └── sql/ # SQL schema (embedded in binary) -│ ├── tables/ # Table definitions -│ └── procedures/ # Stored procedures +│ └── sql/ +│ ├── migrations/ # Ordered, versioned schema migrations (embedded in binary) +│ └── roles/ # Optional least-privilege role DDL (--with-roles) +├── tests/integration/ # Go integration test suite ├── examples/ # Usage examples +├── Dockerfile # Production broker image +├── Dockerfile.test # Integration test runner image +├── docker-compose.yml # Postgres + migrate + broker +├── docker-compose.test.yml # Postgres + test runner └── Makefile # Build automation ``` ## Contributing Contributions are welcome! Please ensure: + - Code is formatted with `go fmt` - Tests pass with `go test ./...` - Documentation is updated ## License -See [LICENSE](./LICENSE) file for details. \ No newline at end of file +See [LICENSE](./LICENSE) file for details. diff --git a/broker.docker.example.yaml b/broker.docker.example.yaml new file mode 100644 index 0000000..7115826 --- /dev/null +++ b/broker.docker.example.yaml @@ -0,0 +1,33 @@ +# Config for the docker-compose.yml stack. Copy to broker.docker.yaml +# (already gitignored), fill in a real password, and it is mounted +# read-only into the broker/migrate containers at +# /etc/pgsql-broker/broker.yaml. + +databases: + - name: primary + host: postgres + port: 5432 + database: broker + user: broker_runtime + password: change_me + sslmode: disable + max_open_conns: 25 + max_idle_conns: 5 + conn_max_lifetime: 5m + conn_max_idle_time: 10m + queue_count: 4 + +broker: + name: pgsql-broker + fetch_query_que_size: 100 + queue_timer_sec: 10 + queue_buffer_size: 50 + worker_idle_timeout_sec: 10 + notify_retry_seconds: 30s + enable_debug: false + lease_seconds: 60 + stale_job_recovery_sec: 30 + +logging: + level: info + format: json diff --git a/cmd/broker/main.go b/cmd/broker/main.go index c7e2287..b81bbce 100644 --- a/cmd/broker/main.go +++ b/cmd/broker/main.go @@ -6,9 +6,12 @@ import ( "log/slog" "os" "os/signal" + "runtime/debug" "syscall" "github.com/spf13/cobra" + "golang.org/x/term" + "git.warky.dev/wdevs/pgsql-broker/pkg/broker" "git.warky.dev/wdevs/pgsql-broker/pkg/broker/adapter" "git.warky.dev/wdevs/pgsql-broker/pkg/broker/config" @@ -25,6 +28,13 @@ var ( cfgFile string logLevel string verifyOnly bool + + withRoles bool + adminUser string + adminPassword string + brokerAdminPassword string + brokerRuntimePassword string + brokerEnqueuePassword string ) func main() { @@ -83,9 +93,59 @@ func init() { // Install command flags installCmd.Flags().BoolVar(&verifyOnly, "verify-only", false, "only verify installation without installing") + installCmd.Flags().BoolVar(&withRoles, "with-roles", false, "also create/update the broker_admin, broker_runtime, and broker_enqueue roles") + installCmd.Flags().StringVar(&adminUser, "admin-user", "", "superuser/CREATEROLE login used only for --with-roles (falls back to PGUSER/PG_USER env)") + installCmd.Flags().StringVar(&adminPassword, "admin-password", "", "password for --admin-user (falls back to PGPASSWORD/PG_PASS env, then an interactive prompt)") + installCmd.Flags().StringVar(&brokerAdminPassword, "broker-admin-password", "", "password to set for broker_admin (falls back to BROKER_ADMIN_PASSWORD env, then an interactive prompt)") + installCmd.Flags().StringVar(&brokerRuntimePassword, "broker-runtime-password", "", "password to set for broker_runtime (falls back to BROKER_RUNTIME_PASSWORD env, then an interactive prompt)") + installCmd.Flags().StringVar(&brokerEnqueuePassword, "broker-enqueue-password", "", "password to set for broker_enqueue (falls back to BROKER_ENQUEUE_PASSWORD env, then an interactive prompt)") } -func runBroker() error { +// resolveCredential returns the first non-empty value among the flag value, +// the given environment variables (checked in order), and -- if none are +// set -- an interactive masked prompt. It errors rather than prompting when +// stdin is not a terminal, since a hang in a non-interactive context (CI, +// systemd) is worse than a clear failure. +func resolveCredential(flagVal string, envNames []string, promptLabel string) (string, error) { + if flagVal != "" { + return flagVal, nil + } + for _, name := range envNames { + if v := os.Getenv(name); v != "" { + return v, nil + } + } + + fd := int(os.Stdin.Fd()) + if !term.IsTerminal(fd) { + return "", fmt.Errorf( + "%s not provided and stdin is not a terminal to prompt on; pass it via flag or one of %v", + promptLabel, envNames, + ) + } + + fmt.Fprintf(os.Stderr, "%s: ", promptLabel) + b, err := term.ReadPassword(fd) + fmt.Fprintln(os.Stderr) + if err != nil { + return "", fmt.Errorf("failed to read %s: %w", promptLabel, err) + } + if len(b) == 0 { + return "", fmt.Errorf("%s must not be empty", promptLabel) + } + return string(b), nil +} + +func runBroker() (err error) { + // Top-level safety net: an unrecovered panic anywhere in startup or the + // shutdown wait must not crash the process with a raw trace -- log it + // and return a normal error instead. + defer func() { + if r := recover(); r != nil { + err = fmt.Errorf("recovered from panic in runBroker: %v\n%s", r, debug.Stack()) + } + }() + // Load configuration cfg, err := config.LoadConfig(cfgFile) if err != nil { @@ -136,12 +196,50 @@ func runInstall() error { ctx := context.Background() + var rolePasswords install.RolePasswords + var adminUserVal, adminPasswordVal string + if withRoles { + if verifyOnly { + return fmt.Errorf("--with-roles cannot be combined with --verify-only") + } + + var err error + adminUserVal, err = resolveCredential(adminUser, []string{"PGUSER", "PG_USER"}, "admin user (superuser/CREATEROLE login for --with-roles)") + if err != nil { + return err + } + adminPasswordVal, err = resolveCredential(adminPassword, []string{"PGPASSWORD", "PG_PASS"}, "admin password") + if err != nil { + return err + } + rolePasswords.AdminPassword, err = resolveCredential(brokerAdminPassword, []string{"BROKER_ADMIN_PASSWORD"}, "broker_admin password") + if err != nil { + return err + } + rolePasswords.RuntimePassword, err = resolveCredential(brokerRuntimePassword, []string{"BROKER_RUNTIME_PASSWORD"}, "broker_runtime password") + if err != nil { + return err + } + rolePasswords.EnqueuePassword, err = resolveCredential(brokerEnqueuePassword, []string{"BROKER_ENQUEUE_PASSWORD"}, "broker_enqueue password") + if err != nil { + return err + } + } + // Install/verify on all configured databases for i, dbCfg := range cfg.Databases { logger.Info("processing database", "index", i, "name", dbCfg.Name, "host", dbCfg.Host, "database", dbCfg.Database) - // Create database adapter - dbAdapter := adapter.NewPostgresAdapter(dbCfg.ToPostgresConfig(), logger) + // Create database adapter. With --with-roles, the config file's own + // user (typically the least-privilege broker_runtime) may not exist + // yet on a fresh cluster -- migrations and role creation both run as + // the admin login instead, since broker_admin must own the schema. + pgCfg := dbCfg.ToPostgresConfig() + if withRoles { + pgCfg.User = adminUserVal + pgCfg.Password = adminPasswordVal + } + dbAdapter := adapter.NewPostgresAdapter(pgCfg, logger) // Connect to database if err := dbAdapter.Connect(ctx); err != nil { @@ -161,9 +259,9 @@ func runInstall() error { } logger.Info("database schema verified successfully", "database", dbCfg.Name) } else { - // Install schema - logger.Info("installing database schema", "database", dbCfg.Name) - if err := installer.InstallSchema(ctx); err != nil { + // Apply migrations + logger.Info("applying database migrations", "database", dbCfg.Name) + if err := installer.ApplyMigrations(ctx); err != nil { dbAdapter.Close() logger.Error("installation failed", "database", dbCfg.Name, "error", err) return fmt.Errorf("installation failed for %s: %w", dbCfg.Name, err) @@ -180,6 +278,15 @@ func runInstall() error { logger.Info("database schema installed and verified successfully", "database", dbCfg.Name) } + if withRoles && !verifyOnly { + logger.Info("applying roles", "database", dbCfg.Name) + if err := installer.InstallRoles(ctx, rolePasswords); err != nil { + dbAdapter.Close() + return fmt.Errorf("failed to install roles for %s: %w", dbCfg.Name, err) + } + logger.Info("roles installed successfully", "database", dbCfg.Name) + } + dbAdapter.Close() } diff --git a/docker-compose.test.yml b/docker-compose.test.yml new file mode 100644 index 0000000..0b42e8c --- /dev/null +++ b/docker-compose.test.yml @@ -0,0 +1,26 @@ +services: + postgres: + image: docker.io/library/postgres:16-alpine + # Tests hardcode host=localhost port=5433, so postgres listens on 5433 + # internally and the tests container joins its network namespace below. + command: ["postgres", "-p", "5433"] + environment: + POSTGRES_DB: broker_test + POSTGRES_USER: user + POSTGRES_PASSWORD: password + ports: + - "5434:5433" + healthcheck: + test: ["CMD-SHELL", "pg_isready -U user -d broker_test -p 5433"] + interval: 2s + timeout: 3s + retries: 30 + + tests: + build: + context: . + dockerfile: Dockerfile.test + network_mode: "service:postgres" + depends_on: + postgres: + condition: service_healthy diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..9cdac40 --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,50 @@ +services: + postgres: + image: docker.io/library/postgres:16-alpine + environment: + POSTGRES_DB: broker + POSTGRES_USER: postgres + POSTGRES_PASSWORD: ${POSTGRES_PASSWORD:?set POSTGRES_PASSWORD in .env} + volumes: + - postgres-data:/var/lib/postgresql/data + healthcheck: + test: ["CMD-SHELL", "pg_isready -U postgres -d broker"] + interval: 2s + timeout: 3s + retries: 30 + restart: unless-stopped + + # One-shot: applies migrations and creates/rotates the least-privilege + # broker_admin/broker_runtime/broker_enqueue roles, then exits. The + # broker service below only starts once this completes successfully. + migrate: + build: + context: . + dockerfile: Dockerfile + depends_on: + postgres: + condition: service_healthy + volumes: + - ./broker.docker.yaml:/etc/pgsql-broker/broker.yaml:ro + environment: + PGUSER: postgres + PGPASSWORD: ${POSTGRES_PASSWORD:?set POSTGRES_PASSWORD in .env} + BROKER_ADMIN_PASSWORD: ${BROKER_ADMIN_PASSWORD:?set BROKER_ADMIN_PASSWORD in .env} + BROKER_RUNTIME_PASSWORD: ${BROKER_RUNTIME_PASSWORD:?set BROKER_RUNTIME_PASSWORD in .env} + BROKER_ENQUEUE_PASSWORD: ${BROKER_ENQUEUE_PASSWORD:?set BROKER_ENQUEUE_PASSWORD in .env} + command: ["install", "--with-roles"] + restart: "no" + + broker: + build: + context: . + dockerfile: Dockerfile + depends_on: + migrate: + condition: service_completed_successfully + volumes: + - ./broker.docker.yaml:/etc/pgsql-broker/broker.yaml:ro + restart: unless-stopped + +volumes: + postgres-data: diff --git a/go.mod b/go.mod index 8f2f9dd..63a8274 100644 --- a/go.mod +++ b/go.mod @@ -1,12 +1,13 @@ module git.warky.dev/wdevs/pgsql-broker -go 1.25.5 +go 1.26.0 require ( github.com/lib/pq v1.10.9 github.com/spf13/cobra v1.10.2 github.com/spf13/viper v1.21.0 github.com/stretchr/testify v1.11.1 + golang.org/x/term v0.46.0 ) require ( @@ -23,7 +24,7 @@ require ( github.com/spf13/pflag v1.0.10 // indirect github.com/subosito/gotenv v1.6.0 // indirect go.yaml.in/yaml/v3 v3.0.4 // indirect - golang.org/x/sys v0.29.0 // indirect + golang.org/x/sys v0.48.0 // indirect golang.org/x/text v0.28.0 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/go.sum b/go.sum index 3fa1763..9d1a74f 100644 --- a/go.sum +++ b/go.sum @@ -45,8 +45,10 @@ github.com/subosito/gotenv v1.6.0 h1:9NlTDc1FTs4qu0DDq7AEtTPNw6SVm7uBMsUCUjABIf8 github.com/subosito/gotenv v1.6.0/go.mod h1:Dk4QP5c2W3ibzajGcXpNraDfq2IrhjMIvMSWPKKo0FU= go.yaml.in/yaml/v3 v3.0.4 h1:tfq32ie2Jv2UxXFdLJdh3jXuOzWiL1fo0bu/FbuKpbc= go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg= -golang.org/x/sys v0.29.0 h1:TPYlXGxvx1MGTn2GiZDhnjPA9wZzZeGKHHmKhHYvgaU= -golang.org/x/sys v0.29.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/sys v0.48.0 h1:bbX/i/6MgT9BVLM9RT1thmxL04yeTAhbEz4SyadbXoo= +golang.org/x/sys v0.48.0/go.mod h1:hNLxWAXmnKAxqDtdwIYC4bM9oQPEecfsnNMuSxOs3og= +golang.org/x/term v0.46.0 h1:3+OXuTbaKDgwk8jTi3aSLHRlmWqHEUDUtxnbFigO4YE= +golang.org/x/term v0.46.0/go.mod h1:+K02xbkittuwc0Am4abfA3Fc+XRGXkvBXNO88NCXPoc= golang.org/x/text v0.28.0 h1:rhazDwis8INMIwQ4tpjLDzUhx6RlXqZNPEM0huQojng= golang.org/x/text v0.28.0/go.mod h1:U8nCwOR8jO/marOQ0QbDiOngZVEBB7MAiitBuMjXiNU= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= diff --git a/pkg/broker/adapter/database.go b/pkg/broker/adapter/database.go index cbcdef9..9825662 100644 --- a/pkg/broker/adapter/database.go +++ b/pkg/broker/adapter/database.go @@ -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) diff --git a/pkg/broker/adapter/postgres.go b/pkg/broker/adapter/postgres.go index 32b7a4a..abed6cf 100644 --- a/pkg/broker/adapter/postgres.go +++ b/pkg/broker/adapter/postgres.go @@ -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 diff --git a/pkg/broker/adapter/safego.go b/pkg/broker/adapter/safego.go new file mode 100644 index 0000000..65f3e40 --- /dev/null +++ b/pkg/broker/adapter/safego.go @@ -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 +} diff --git a/pkg/broker/broker.go b/pkg/broker/broker.go index 38d788b..fc6f3fc 100644 --- a/pkg/broker/broker.go +++ b/pkg/broker/broker.go @@ -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) } diff --git a/pkg/broker/config/config.go b/pkg/broker/config/config.go index c1f3c5f..7b26603 100644 --- a/pkg/broker/config/config.go +++ b/pkg/broker/config/config.go @@ -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" + } } } diff --git a/pkg/broker/database_instance.go b/pkg/broker/database_instance.go index a6ed81f..de1433c 100644 --- a/pkg/broker/database_instance.go +++ b/pkg/broker/database_instance.go @@ -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 -} \ No newline at end of file +} diff --git a/pkg/broker/install/install.go b/pkg/broker/install/install.go index 7e71479..43cb532 100644 --- a/pkg/broker/install/install.go +++ b/pkg/broker/install/install.go @@ -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") diff --git a/pkg/broker/install/sql/migrations/0001_schema.sql b/pkg/broker/install/sql/migrations/0001_schema.sql new file mode 100644 index 0000000..edb1f0c --- /dev/null +++ b/pkg/broker/install/sql/migrations/0001_schema.sql @@ -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; diff --git a/pkg/broker/install/sql/migrations/0002_broker_queueinstance.sql b/pkg/broker/install/sql/migrations/0002_broker_queueinstance.sql new file mode 100644 index 0000000..5b6ec7f --- /dev/null +++ b/pkg/broker/install/sql/migrations/0002_broker_queueinstance.sql @@ -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.'; diff --git a/pkg/broker/install/sql/migrations/0003_broker_schedule.sql b/pkg/broker/install/sql/migrations/0003_broker_schedule.sql new file mode 100644 index 0000000..7c6cd6c --- /dev/null +++ b/pkg/broker/install/sql/migrations/0003_broker_schedule.sql @@ -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(); diff --git a/pkg/broker/install/sql/migrations/0004_broker_jobs.sql b/pkg/broker/install/sql/migrations/0004_broker_jobs.sql new file mode 100644 index 0000000..45c4c2c --- /dev/null +++ b/pkg/broker/install/sql/migrations/0004_broker_jobs.sql @@ -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')); diff --git a/pkg/broker/install/sql/migrations/0005_broker_job_dependency.sql b/pkg/broker/install/sql/migrations/0005_broker_job_dependency.sql new file mode 100644 index 0000000..e808681 --- /dev/null +++ b/pkg/broker/install/sql/migrations/0005_broker_job_dependency.sql @@ -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')); diff --git a/pkg/broker/install/sql/migrations/0006_broker_set_tenant.sql b/pkg/broker/install/sql/migrations/0006_broker_set_tenant.sql new file mode 100644 index 0000000..dc9a9cb --- /dev/null +++ b/pkg/broker/install/sql/migrations/0006_broker_set_tenant.sql @@ -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)'; diff --git a/pkg/broker/install/sql/migrations/0007_broker_get.sql b/pkg/broker/install/sql/migrations/0007_broker_get.sql new file mode 100644 index 0000000..80eff35 --- /dev/null +++ b/pkg/broker/install/sql/migrations/0007_broker_get.sql @@ -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'; diff --git a/pkg/broker/install/sql/migrations/0008_broker_run.sql b/pkg/broker/install/sql/migrations/0008_broker_run.sql new file mode 100644 index 0000000..ef13d09 --- /dev/null +++ b/pkg/broker/install/sql/migrations/0008_broker_run.sql @@ -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'; diff --git a/pkg/broker/install/sql/migrations/0009_broker_set.sql b/pkg/broker/install/sql/migrations/0009_broker_set.sql new file mode 100644 index 0000000..3785130 --- /dev/null +++ b/pkg/broker/install/sql/migrations/0009_broker_set.sql @@ -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)'; diff --git a/pkg/broker/install/sql/migrations/0010_broker_register_instance.sql b/pkg/broker/install/sql/migrations/0010_broker_register_instance.sql new file mode 100644 index 0000000..7836d76 --- /dev/null +++ b/pkg/broker/install/sql/migrations/0010_broker_register_instance.sql @@ -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)'; diff --git a/pkg/broker/install/sql/migrations/0011_broker_add_job.sql b/pkg/broker/install/sql/migrations/0011_broker_add_job.sql new file mode 100644 index 0000000..9cc034e --- /dev/null +++ b/pkg/broker/install/sql/migrations/0011_broker_add_job.sql @@ -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'; diff --git a/pkg/broker/install/sql/procedures/06_broker_ping_instance.sql b/pkg/broker/install/sql/migrations/0012_broker_ping_shutdown_instance.sql similarity index 59% rename from pkg/broker/install/sql/procedures/06_broker_ping_instance.sql rename to pkg/broker/install/sql/migrations/0012_broker_ping_shutdown_instance.sql index 4037bd4..b849dae 100644 --- a/pkg/broker/install/sql/procedures/06_broker_ping_instance.sql +++ b/pkg/broker/install/sql/migrations/0012_broker_ping_shutdown_instance.sql @@ -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)'; diff --git a/pkg/broker/install/sql/migrations/0013_broker_recover_stale_jobs.sql b/pkg/broker/install/sql/migrations/0013_broker_recover_stale_jobs.sql new file mode 100644 index 0000000..7c35500 --- /dev/null +++ b/pkg/broker/install/sql/migrations/0013_broker_recover_stale_jobs.sql @@ -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'; diff --git a/pkg/broker/install/sql/migrations/0014_job_groups.sql b/pkg/broker/install/sql/migrations/0014_job_groups.sql new file mode 100644 index 0000000..060a7db --- /dev/null +++ b/pkg/broker/install/sql/migrations/0014_job_groups.sql @@ -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'; diff --git a/pkg/broker/install/sql/migrations/0015_broker_add_job_groups.sql b/pkg/broker/install/sql/migrations/0015_broker_add_job_groups.sql new file mode 100644 index 0000000..d69ac9e --- /dev/null +++ b/pkg/broker/install/sql/migrations/0015_broker_add_job_groups.sql @@ -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'; diff --git a/pkg/broker/install/sql/migrations/0016_broker_get_groups.sql b/pkg/broker/install/sql/migrations/0016_broker_get_groups.sql new file mode 100644 index 0000000..183da77 --- /dev/null +++ b/pkg/broker/install/sql/migrations/0016_broker_get_groups.sql @@ -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'; diff --git a/pkg/broker/install/sql/migrations/0017_broker_add_job_simple_groups.sql b/pkg/broker/install/sql/migrations/0017_broker_add_job_simple_groups.sql new file mode 100644 index 0000000..e550e9c --- /dev/null +++ b/pkg/broker/install/sql/migrations/0017_broker_add_job_simple_groups.sql @@ -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)'; diff --git a/pkg/broker/install/sql/procedures/00_install.sql b/pkg/broker/install/sql/procedures/00_install.sql deleted file mode 100644 index 414831f..0000000 --- a/pkg/broker/install/sql/procedures/00_install.sql +++ /dev/null @@ -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!' diff --git a/pkg/broker/install/sql/procedures/01_broker_get.sql b/pkg/broker/install/sql/procedures/01_broker_get.sql deleted file mode 100644 index 9d6dd13..0000000 --- a/pkg/broker/install/sql/procedures/01_broker_get.sql +++ /dev/null @@ -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'; diff --git a/pkg/broker/install/sql/procedures/02_broker_run.sql b/pkg/broker/install/sql/procedures/02_broker_run.sql deleted file mode 100644 index 9c12f6d..0000000 --- a/pkg/broker/install/sql/procedures/02_broker_run.sql +++ /dev/null @@ -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'; diff --git a/pkg/broker/install/sql/procedures/03_broker_set.sql b/pkg/broker/install/sql/procedures/03_broker_set.sql deleted file mode 100644 index d177f77..0000000 --- a/pkg/broker/install/sql/procedures/03_broker_set.sql +++ /dev/null @@ -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)'; diff --git a/pkg/broker/install/sql/procedures/04_broker_register_instance.sql b/pkg/broker/install/sql/procedures/04_broker_register_instance.sql deleted file mode 100644 index 18049d8..0000000 --- a/pkg/broker/install/sql/procedures/04_broker_register_instance.sql +++ /dev/null @@ -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'; diff --git a/pkg/broker/install/sql/procedures/05_broker_add_job.sql b/pkg/broker/install/sql/procedures/05_broker_add_job.sql deleted file mode 100644 index 0de3fc7..0000000 --- a/pkg/broker/install/sql/procedures/05_broker_add_job.sql +++ /dev/null @@ -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'; diff --git a/pkg/broker/install/sql/roles/0001_roles.sql b/pkg/broker/install/sql/roles/0001_roles.sql new file mode 100644 index 0000000..e76949e --- /dev/null +++ b/pkg/broker/install/sql/roles/0001_roles.sql @@ -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; diff --git a/pkg/broker/install/sql/tables/00_install.sql b/pkg/broker/install/sql/tables/00_install.sql deleted file mode 100644 index 1a529f2..0000000 --- a/pkg/broker/install/sql/tables/00_install.sql +++ /dev/null @@ -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!' \ No newline at end of file diff --git a/pkg/broker/install/sql/tables/01_broker_queueinstance.sql b/pkg/broker/install/sql/tables/01_broker_queueinstance.sql deleted file mode 100644 index dbd1462..0000000 --- a/pkg/broker/install/sql/tables/01_broker_queueinstance.sql +++ /dev/null @@ -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'; diff --git a/pkg/broker/install/sql/tables/02_broker_schedule.sql b/pkg/broker/install/sql/tables/02_broker_schedule.sql deleted file mode 100644 index 78ce5a7..0000000 --- a/pkg/broker/install/sql/tables/02_broker_schedule.sql +++ /dev/null @@ -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(); diff --git a/pkg/broker/install/sql/tables/03_broker_jobs.sql b/pkg/broker/install/sql/tables/03_broker_jobs.sql deleted file mode 100644 index 8475439..0000000 --- a/pkg/broker/install/sql/tables/03_broker_jobs.sql +++ /dev/null @@ -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(); diff --git a/pkg/broker/models/models.go b/pkg/broker/models/models.go index a413360..0e55cd6 100644 --- a/pkg/broker/models/models.go +++ b/pkg/broker/models/models.go @@ -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"` diff --git a/pkg/broker/queue/queue.go b/pkg/broker/queue/queue.go index fd547b8..a7b4681 100644 --- a/pkg/broker/queue/queue.go +++ b/pkg/broker/queue/queue.go @@ -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 diff --git a/pkg/broker/worker/worker.go b/pkg/broker/worker/worker.go index 4926612..f87c822 100644 --- a/pkg/broker/worker/worker.go +++ b/pkg/broker/worker/worker.go @@ -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())) } } diff --git a/plan/docs/recommendations.md b/plan/docs/recommendations.md index 33b88a3..1d8b824 100644 --- a/plan/docs/recommendations.md +++ b/plan/docs/recommendations.md @@ -1,5 +1,5 @@ # PostgreSQL Schema and Function Recommendations - +/home/warkanum/.claude/plans/binary-puzzling-mountain.md **Review date:** 2026-09-14 **Scope:** `pkg/broker/install/sql` tables and functions, including their interaction with the Go worker. diff --git a/tests/integration/rls_test.go b/tests/integration/rls_test.go new file mode 100644 index 0000000..2c851da --- /dev/null +++ b/tests/integration/rls_test.go @@ -0,0 +1,145 @@ +package integration + +import ( + "context" + "database/sql" + "testing" + + _ "github.com/lib/pq" + "github.com/stretchr/testify/require" +) + +// TestRLSTenantIsolation verifies that a non-superuser role without +// BYPASSRLS, connected via broker_set_tenant, only ever sees jobs and +// dependency rows for its own tenant -- the core guarantee behind the +// broker_jobs/broker_job_dependency FORCE ROW LEVEL SECURITY policies. +func TestRLSTenantIsolation(t *testing.T) { + ctx := context.Background() + adminDB := setupStage5Schema(t) + + // A restricted role, no BYPASSRLS, mirroring sql/roles/0001_roles.sql's + // broker_runtime (superuser test DB roles already bypass RLS entirely, + // so this test would be meaningless against the "user" role). + _, err := adminDB.Exec("DROP ROLE IF EXISTS test_broker_runtime") + require.NoError(t, err) + _, err = adminDB.Exec("CREATE ROLE test_broker_runtime LOGIN PASSWORD 'test-pass' NOSUPERUSER NOBYPASSRLS") + require.NoError(t, err) + t.Cleanup(func() { + if _, err := adminDB.Exec("REASSIGN OWNED BY test_broker_runtime TO CURRENT_USER"); err != nil { + t.Logf("warning: failed to reassign objects owned by test_broker_runtime: %v", err) + } + if _, err := adminDB.Exec("DROP OWNED BY test_broker_runtime"); err != nil { + t.Logf("warning: failed to drop grants owned by test_broker_runtime: %v", err) + } + if _, err := adminDB.Exec("DROP ROLE IF EXISTS test_broker_runtime"); err != nil { + t.Logf("warning: failed to drop role test_broker_runtime: %v", err) + } + }) + + for _, stmt := range []string{ + "GRANT USAGE ON SCHEMA broker TO test_broker_runtime", + "GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA broker TO test_broker_runtime", + "GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA broker TO test_broker_runtime", + "GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA broker TO test_broker_runtime", + } { + _, err = adminDB.Exec(stmt) + require.NoError(t, err) + } + + runtimeDB, err := sql.Open("postgres", + "user=test_broker_runtime password=test-pass dbname=broker_test host=localhost port=5433 sslmode=disable options='-c search_path=broker,public'") + require.NoError(t, err) + defer runtimeDB.Close() + require.NoError(t, runtimeDB.Ping()) + + // broker_set_tenant uses SET LOCAL semantics (set_config(..., true)), so + // it only takes effect for the remainder of the transaction it runs in. + // Callers must set the tenant and perform the tenant-scoped operation in + // the same explicit transaction (as worker.go's processJobs does) -- two + // separate autocommitted statements would each run in their own + // transaction and the tenant setting would not carry over. + addJobAsTenant := func(tenant, name string) int64 { + conn, err := runtimeDB.Conn(ctx) + require.NoError(t, err) + defer conn.Close() + + tx, err := conn.BeginTx(ctx, nil) + require.NoError(t, err) + + _, err = tx.ExecContext(ctx, "SELECT broker.broker_set_tenant($1)", tenant) + require.NoError(t, err) + + var retval int + var errmsg string + var jobID int64 + err = tx.QueryRowContext(ctx, ` + SELECT p_retval, p_errmsg, p_job_id FROM broker.broker_add_job( + $1, 'SELECT 1', 1, 0, 'sql', NULL, NULL, NULL, NULL, 1 + )`, name, + ).Scan(&retval, &errmsg, &jobID) + require.NoError(t, err) + require.Equal(t, 0, retval, errmsg) + + require.NoError(t, tx.Commit()) + return jobID + } + + tenantAJob := addJobAsTenant("tenant-a", "tenant-a-job") + tenantBJob := addJobAsTenant("tenant-b", "tenant-b-job") + require.NotEqual(t, tenantAJob, tenantBJob) + + // As tenant-a, only tenant-a's job must be visible. All of tenant-a's + // checks (including the broker_get claim below) share one explicit + // transaction so the SET LOCAL tenant context stays in effect throughout. + connA, err := runtimeDB.Conn(ctx) + require.NoError(t, err) + defer connA.Close() + txA, err := connA.BeginTx(ctx, nil) + require.NoError(t, err) + defer txA.Rollback() + + _, err = txA.ExecContext(ctx, "SELECT broker.broker_set_tenant($1)", "tenant-a") + require.NoError(t, err) + + var visibleCount int + err = txA.QueryRowContext(ctx, "SELECT COUNT(*) FROM broker.broker_jobs WHERE id_broker_jobs IN ($1, $2)", + tenantAJob, tenantBJob).Scan(&visibleCount) + require.NoError(t, err) + require.Equal(t, 1, visibleCount, "tenant-a must see only its own job, not tenant-b's") + + var visibleName string + err = txA.QueryRowContext(ctx, "SELECT job_name FROM broker.broker_jobs WHERE id_broker_jobs = $1", tenantAJob).Scan(&visibleName) + require.NoError(t, err) + require.Equal(t, "tenant-a-job", visibleName) + + // As tenant-b, only tenant-b's job must be visible. + connB, err := runtimeDB.Conn(ctx) + require.NoError(t, err) + defer connB.Close() + txB, err := connB.BeginTx(ctx, nil) + require.NoError(t, err) + defer txB.Rollback() + + _, err = txB.ExecContext(ctx, "SELECT broker.broker_set_tenant($1)", "tenant-b") + require.NoError(t, err) + + err = txB.QueryRowContext(ctx, "SELECT COUNT(*) FROM broker.broker_jobs WHERE id_broker_jobs IN ($1, $2)", + tenantAJob, tenantBJob).Scan(&visibleCount) + require.NoError(t, err) + require.Equal(t, 1, visibleCount, "tenant-b must see only its own job, not tenant-a's") + require.NoError(t, txB.Commit()) + + // broker_get run under tenant-a's context must never be able to claim + // tenant-b's job. + var claimedID sql.NullInt64 + var getRetval int + var getErrmsg string + err = txA.QueryRowContext(ctx, + "SELECT p_retval, p_errmsg, p_job_id FROM broker.broker_get($1, NULL, $2)", 1, 60, + ).Scan(&getRetval, &getErrmsg, &claimedID) + require.NoError(t, err) + require.Equal(t, 0, getRetval, getErrmsg) + require.True(t, claimedID.Valid) + require.Equal(t, tenantAJob, claimedID.Int64, "tenant-a's broker_get must only ever claim tenant-a's own job") + require.NoError(t, txA.Commit()) +} diff --git a/tests/integration/stage5_test.go b/tests/integration/stage5_test.go new file mode 100644 index 0000000..5ffe3ca --- /dev/null +++ b/tests/integration/stage5_test.go @@ -0,0 +1,341 @@ +package integration + +import ( + "context" + "database/sql" + "log/slog" + "testing" + "time" + + _ "github.com/lib/pq" + "github.com/stretchr/testify/require" + + "git.warky.dev/wdevs/pgsql-broker/pkg/broker/adapter" + "git.warky.dev/wdevs/pgsql-broker/pkg/broker/install" +) + +const stage5ConnStr = "user=user password=password dbname=broker_test host=localhost port=5433 sslmode=disable" + +func newStage5Adapter(logger adapter.Logger) *adapter.PostgresAdapter { + return adapter.NewPostgresAdapter(adapter.PostgresConfig{ + Host: "localhost", Port: 5433, Database: "broker_test", + User: "user", Password: "password", SSLMode: "disable", + MaxOpenConns: 10, MaxIdleConns: 2, + ConnMaxLifetime: 5 * time.Minute, ConnMaxIdleTime: 10 * time.Minute, + }, logger) +} + +// setupStage5Schema drops and re-installs a clean broker schema, returning a +// superuser *sql.DB for direct SQL against it. +func setupStage5Schema(t *testing.T) *sql.DB { + t.Helper() + + db, err := connectWithRetry(stage5ConnStr, 10, 2*time.Second) + require.NoError(t, err) + + cleanupSchema(t, db) + + logger := adapter.NewSlogLogger(slog.LevelWarn) + dbAdapter := newStage5Adapter(logger) + require.NoError(t, dbAdapter.Connect(context.Background())) + defer dbAdapter.Close() + + installer := install.New(dbAdapter, logger) + require.NoError(t, installer.ApplyMigrations(context.Background())) + + t.Cleanup(func() { db.Close() }) + return db +} + +// TestRepeatMigrationRunIsNoOp verifies applying the migration set a second +// time (with nothing pending) makes no changes and reports no error. +func TestRepeatMigrationRunIsNoOp(t *testing.T) { + ctx := context.Background() + setupStage5Schema(t) + + logger := adapter.NewSlogLogger(slog.LevelWarn) + dbAdapter := newStage5Adapter(logger) + require.NoError(t, dbAdapter.Connect(ctx)) + defer dbAdapter.Close() + + installer := install.New(dbAdapter, logger) + + pending, err := installer.PendingMigrations(ctx) + require.NoError(t, err) + require.Empty(t, pending, "no migrations should be pending right after install") + + require.NoError(t, installer.ApplyMigrations(ctx)) + require.NoError(t, installer.VerifyInstallation(ctx)) +} + +// TestDuplicateInstanceStartFails verifies that broker_register_instance +// rejects a second registration under the same instance name while the +// first holds the advisory lock, and succeeds again once it's released. +func TestDuplicateInstanceStartFails(t *testing.T) { + ctx := context.Background() + db := setupStage5Schema(t) + + connA, err := db.Conn(ctx) + require.NoError(t, err) + defer connA.Close() + + var retval int + var errmsg string + var instanceID sql.NullInt64 + + err = connA.QueryRowContext(ctx, + "SELECT p_retval, p_errmsg, p_instance_id FROM broker.broker_register_instance($1,$2,$3,$4,$5)", + "dup-test", "host-a", 111, "test", 2, + ).Scan(&retval, &errmsg, &instanceID) + require.NoError(t, err) + require.Equal(t, 0, retval, "first registration should succeed: %s", errmsg) + require.True(t, instanceID.Valid) + + // Second registration under the same name, on a different connection, + // must fail while connA still holds the advisory lock. + connB, err := db.Conn(ctx) + require.NoError(t, err) + defer connB.Close() + + err = connB.QueryRowContext(ctx, + "SELECT p_retval, p_errmsg, p_instance_id FROM broker.broker_register_instance($1,$2,$3,$4,$5)", + "dup-test", "host-b", 222, "test", 2, + ).Scan(&retval, &errmsg, &instanceID) + require.NoError(t, err) + require.Equal(t, 3, retval, "second registration must be rejected by the advisory lock") + require.False(t, instanceID.Valid) + + // Release the lock (as registerInstance's caller would on shutdown) and + // confirm a fresh registration then succeeds. + _, err = connA.ExecContext(ctx, "SELECT pg_advisory_unlock(hashtextextended($1, 0))", "broker:dup-test") + require.NoError(t, err) + require.NoError(t, connA.Close()) + + connC, err := db.Conn(ctx) + require.NoError(t, err) + defer connC.Close() + + err = connC.QueryRowContext(ctx, + "SELECT p_retval, p_errmsg, p_instance_id FROM broker.broker_register_instance($1,$2,$3,$4,$5)", + "dup-test", "host-c", 333, "test", 2, + ).Scan(&retval, &errmsg, &instanceID) + require.NoError(t, err) + require.Equal(t, 0, retval, "registration should succeed again once the lock is released: %s", errmsg) +} + +// TestJobDependencies covers dependency ordering (a job is not claimable +// while an incomplete dependency exists), rejection of a direct +// self-dependency at the table level, and idempotent duplicate-dependency +// inserts. +func TestJobDependencies(t *testing.T) { + ctx := context.Background() + db := setupStage5Schema(t) + + addJob := func(name string, deps string) int64 { + var retval int + var errmsg string + var jobID int64 + depsArg := sql.NullString{String: deps, Valid: deps != ""} + err := db.QueryRowContext(ctx, ` + SELECT p_retval, p_errmsg, p_job_id FROM broker.broker_add_job( + $1, 'SELECT 1', 1, 0, 'sql', NULL, NULL, $2::BIGINT[], NULL, 3 + )`, name, depsArg, + ).Scan(&retval, &errmsg, &jobID) + require.NoError(t, err) + require.Equal(t, 0, retval, "broker_add_job(%s) failed: %s", name, errmsg) + return jobID + } + + base := addJob("base", "") + dependent := addJob("dependent", "{"+sqlItoa(base)+"}") + require.NotZero(t, dependent) + + // broker_get must not return "dependent" while "base" is still pending; + // it should return "base" instead. + var jobID sql.NullInt64 + var getRetval int + var getErrmsg string + err := db.QueryRowContext(ctx, + "SELECT p_retval, p_errmsg, p_job_id FROM broker.broker_get($1, NULL, $2)", 1, 60, + ).Scan(&getRetval, &getErrmsg, &jobID) + require.NoError(t, err) + require.Equal(t, 0, getRetval, getErrmsg) + require.True(t, jobID.Valid) + require.Equal(t, base, jobID.Int64, "dependency-free job must be claimed before its dependent") + + // A self-dependency must be rejected by the table's CHECK constraint, + // regardless of caller. + target := addJob("self-target", "") + _, err = db.ExecContext(ctx, + "INSERT INTO broker.broker_job_dependency (job_id, depends_on_job_id) VALUES ($1, $1)", target, + ) + require.Error(t, err, "self-dependency must violate the CHECK constraint") + + // Re-adding the same dependency pair must be a no-op, not an error + // (ON CONFLICT DO NOTHING on the (job_id, depends_on_job_id) pair). + other := addJob("other-dep-target", "") + dependentTwo := addJob("dependent-two", "{"+sqlItoa(other)+"}") + _, err = db.ExecContext(ctx, + "INSERT INTO broker.broker_job_dependency (job_id, depends_on_job_id) VALUES ($1, $2) ON CONFLICT (job_id, depends_on_job_id) DO NOTHING", + dependentTwo, other, + ) + require.NoError(t, err) + + var depCount int + err = db.QueryRowContext(ctx, + "SELECT COUNT(*) FROM broker.broker_job_dependency WHERE job_id = $1 AND depends_on_job_id = $2", + dependentTwo, other, + ).Scan(&depCount) + require.NoError(t, err) + require.Equal(t, 1, depCount, "duplicate dependency insert must not create a second row") +} + +// TestStaleLeaseRecovery verifies that a job whose lease has expired while +// still marked running is requeued (attempts remain) by +// broker_recover_stale_jobs, and dead-lettered once attempts are exhausted. +func TestStaleLeaseRecovery(t *testing.T) { + ctx := context.Background() + db := setupStage5Schema(t) + + var jobID int64 + var retval int + var errmsg string + err := db.QueryRowContext(ctx, ` + SELECT p_retval, p_errmsg, p_job_id FROM broker.broker_add_job( + 'stale-job', 'SELECT 1', 1, 0, 'sql', NULL, NULL, NULL, NULL, 2 + )`, + ).Scan(&retval, &errmsg, &jobID) + require.NoError(t, err) + require.Equal(t, 0, retval, errmsg) + + // Simulate a worker having claimed the job with a lease that has + // already expired, without ever calling broker_run. + _, err = db.ExecContext(ctx, ` + UPDATE broker.broker_jobs + SET complete_status = 1, attempt_count = 1, + lease_token = gen_random_uuid(), leased_at = NOW() - INTERVAL '2 minutes', + lease_expires_at = NOW() - INTERVAL '1 minute' + WHERE id_broker_jobs = $1`, jobID) + require.NoError(t, err) + + var recRetval, recovered int + var recErrmsg string + err = db.QueryRowContext(ctx, + "SELECT p_retval, p_errmsg, p_recovered_count FROM broker.broker_recover_stale_jobs()", + ).Scan(&recRetval, &recErrmsg, &recovered) + require.NoError(t, err) + require.Equal(t, 0, recRetval, recErrmsg) + require.Equal(t, 1, recovered) + + var status int + var leaseToken sql.NullString + err = db.QueryRowContext(ctx, + "SELECT complete_status, lease_token FROM broker.broker_jobs WHERE id_broker_jobs = $1", jobID, + ).Scan(&status, &leaseToken) + require.NoError(t, err) + require.Equal(t, 0, status, "job with attempts remaining must be requeued as pending") + require.False(t, leaseToken.Valid, "lease must be cleared on recovery") + + // Exhaust attempts (attempt_count already 1, max_attempts 2) then + // simulate one more stale lease -- this time it must dead-letter. + _, err = db.ExecContext(ctx, ` + UPDATE broker.broker_jobs + SET complete_status = 1, attempt_count = 2, + lease_token = gen_random_uuid(), leased_at = NOW() - INTERVAL '2 minutes', + lease_expires_at = NOW() - INTERVAL '1 minute' + WHERE id_broker_jobs = $1`, jobID) + require.NoError(t, err) + + err = db.QueryRowContext(ctx, + "SELECT p_retval, p_errmsg, p_recovered_count FROM broker.broker_recover_stale_jobs()", + ).Scan(&recRetval, &recErrmsg, &recovered) + require.NoError(t, err) + require.Equal(t, 0, recRetval, recErrmsg) + require.Equal(t, 1, recovered) + + err = db.QueryRowContext(ctx, + "SELECT complete_status FROM broker.broker_jobs WHERE id_broker_jobs = $1", jobID, + ).Scan(&status) + require.NoError(t, err) + require.Equal(t, 3, status, "job with attempts exhausted must be dead-lettered") +} + +// TestFailedJobRetriesThenCompletesWithoutStranding verifies the original +// bug fix: a job whose execution fails is requeued for retry (not stranded +// in the running state) and, once it succeeds, ends up completed. +func TestFailedJobRetriesThenCompletesWithoutStranding(t *testing.T) { + ctx := context.Background() + db := setupStage5Schema(t) + + var jobID int64 + var retval int + var errmsg string + err := db.QueryRowContext(ctx, ` + SELECT p_retval, p_errmsg, p_job_id FROM broker.broker_add_job( + 'flaky-job', 'SELECT 1/0', 1, 0, 'sql', NULL, NULL, NULL, NULL, 2 + )`, + ).Scan(&retval, &errmsg, &jobID) + require.NoError(t, err) + require.Equal(t, 0, retval, errmsg) + + var claimedID sql.NullInt64 + var leaseToken sql.NullString + var getRetval int + var getErrmsg string + err = db.QueryRowContext(ctx, + "SELECT p_retval, p_errmsg, p_job_id, p_lease_token FROM broker.broker_get($1, NULL, $2)", 1, 60, + ).Scan(&getRetval, &getErrmsg, &claimedID, &leaseToken) + require.NoError(t, err) + require.Equal(t, 0, getRetval, getErrmsg) + require.True(t, claimedID.Valid) + require.Equal(t, jobID, claimedID.Int64) + + var runRetval, jobStatus int + var runErrmsg string + err = db.QueryRowContext(ctx, + "SELECT p_retval, p_errmsg, p_job_status FROM broker.broker_run($1, $2)", jobID, leaseToken.String, + ).Scan(&runRetval, &runErrmsg, &jobStatus) + require.NoError(t, err) + require.Equal(t, 0, runRetval, "broker_run must report success (retval=0) even when the job itself failed: %s", runErrmsg) + require.Equal(t, 0, jobStatus, "failing job with attempts remaining must be requeued (job_status=0), not stranded") + + var status int + err = db.QueryRowContext(ctx, + "SELECT complete_status FROM broker.broker_jobs WHERE id_broker_jobs = $1", jobID, + ).Scan(&status) + require.NoError(t, err) + require.Equal(t, 0, status, "job must be pending again, not stuck at running (1)") + + // Presenting a stale lease token after the job was already reset must + // be rejected rather than silently re-running. broker_run returns early + // on the lease mismatch without ever setting p_job_status, so it comes + // back NULL here. + var staleJobStatus sql.NullInt64 + err = db.QueryRowContext(ctx, + "SELECT p_retval, p_errmsg, p_job_status FROM broker.broker_run($1, $2)", jobID, leaseToken.String, + ).Scan(&runRetval, &runErrmsg, &staleJobStatus) + require.NoError(t, err) + require.NotEqual(t, 0, runRetval, "broker_run must reject a stale/mismatched lease token") +} + +func sqlItoa(v int64) string { + if v == 0 { + return "0" + } + neg := v < 0 + if neg { + v = -v + } + var buf [20]byte + i := len(buf) + for v > 0 { + i-- + buf[i] = byte('0' + v%10) + v /= 10 + } + if neg { + i-- + buf[i] = '-' + } + return string(buf[i:]) +} diff --git a/tests/integration/workflow_test.go b/tests/integration/workflow_test.go index 227b243..7c061c2 100644 --- a/tests/integration/workflow_test.go +++ b/tests/integration/workflow_test.go @@ -65,8 +65,8 @@ func TestBrokerWorkflow(t *testing.T) { // Install schema t.Log("Installing database schema...") installer := install.New(dbAdapter, logger) - err = installer.InstallSchema(ctx) - require.NoError(t, err, "Failed to install schema") + err = installer.ApplyMigrations(ctx) + require.NoError(t, err, "Failed to apply migrations") // Verify installation t.Log("Verifying schema installation...") @@ -127,7 +127,7 @@ func TestBrokerWorkflow(t *testing.T) { var errmsg string var jobID int64 err = db.QueryRowContext(ctx, ` - SELECT * FROM broker_add_job( + SELECT * FROM broker.broker_add_job( $1, -- job_name $2, -- execute_str $3, -- job_queue @@ -135,7 +135,9 @@ func TestBrokerWorkflow(t *testing.T) { $5, -- job_language NULL, -- run_as NULL, -- schedule_id - NULL -- depends_on + NULL, -- depends_on_job_ids + NULL, -- idempotency_key + NULL -- max_attempts ) `, "Test Job", @@ -165,7 +167,7 @@ func TestBrokerWorkflow(t *testing.T) { SELECT id_broker_jobs, job_name, job_priority, job_queue, job_language, execute_str, execute_result, error_msg, complete_status, created_at, updated_at - FROM broker_jobs + FROM broker.broker_jobs WHERE id_broker_jobs = $1 `, jobID).Scan( &job.ID, @@ -243,32 +245,10 @@ func connectWithRetry(connStr string, maxRetries int, retryInterval time.Duratio return nil, err } -// cleanupSchema removes all broker tables and functions for a clean test +// cleanupSchema drops the entire broker schema for a clean test run. func cleanupSchema(t *testing.T, db *sql.DB) { - tables := []string{"broker_jobs", "broker_queueinstance", "broker_schedule"} - procedures := []string{ - "broker_get", - "broker_run", - "broker_set", - "broker_add_job", - "broker_register_instance", - "broker_ping_instance", - "broker_shutdown_instance", - } - - // Drop procedures - for _, proc := range procedures { - _, err := db.Exec("DROP FUNCTION IF EXISTS " + proc + " CASCADE") - if err != nil { - t.Logf("Warning: failed to drop procedure %s: %v", proc, err) - } - } - - // Drop tables - for _, table := range tables { - _, err := db.Exec("DROP TABLE IF EXISTS " + table + " CASCADE") - if err != nil { - t.Logf("Warning: failed to drop table %s: %v", table, err) - } + _, err := db.Exec("DROP SCHEMA IF EXISTS broker CASCADE") + if err != nil { + t.Logf("Warning: failed to drop broker schema: %v", err) } }