feat(broker): migrations-based install, roles, RLS, and job dependency groups
Replace the ad-hoc tables/procedures install layout with versioned, ordered SQL migrations tracked in broker_schema_migrations. Add optional least-privilege role provisioning (--with-roles), multi-tenant row-level security, lease-based job claiming with stale-lease recovery, and job dependencies -- both by job id and by fan-in job group. Add Docker/Compose support for running the broker and its test suite.
This commit is contained in:
@@ -0,0 +1,6 @@
|
||||
.git
|
||||
bin
|
||||
broker.log
|
||||
broker.pid
|
||||
plan
|
||||
*.md
|
||||
@@ -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
|
||||
@@ -32,7 +32,9 @@ bin/
|
||||
broker.yaml
|
||||
broker.yml
|
||||
broker.json
|
||||
broker.docker.yaml
|
||||
!broker.example.yaml
|
||||
!broker.docker.example.yaml
|
||||
|
||||
# IDE
|
||||
.vscode/
|
||||
|
||||
+40
@@ -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"]
|
||||
@@ -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,16 +386,22 @@ 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
|
||||
|
||||
@@ -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
|
||||
+113
-6
@@ -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()
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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:
|
||||
@@ -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
|
||||
)
|
||||
|
||||
@@ -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=
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+168
-71
@@ -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)
|
||||
|
||||
|
||||
+375
-145
@@ -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
|
||||
}
|
||||
|
||||
// Filter and sort SQL files
|
||||
sqlFiles := filterAndSortSQLFiles(files)
|
||||
// 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()
|
||||
|
||||
for _, file := range sqlFiles {
|
||||
// Skip install script
|
||||
if file == "00_install.sql" {
|
||||
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
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
i.logger.Info("tables installed successfully")
|
||||
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++
|
||||
}
|
||||
|
||||
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
|
||||
|
||||
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
|
||||
|
||||
buffer += stmt + ";"
|
||||
case c == '"':
|
||||
buffer.WriteRune(c)
|
||||
i++
|
||||
for i < n {
|
||||
buffer.WriteRune(runes[i])
|
||||
if runes[i] == '"' {
|
||||
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 = ""
|
||||
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 {
|
||||
// Odd number means we're inside a function, keep accumulating
|
||||
buffer += " "
|
||||
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",
|
||||
}
|
||||
|
||||
// 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)
|
||||
|
||||
pending, err := i.PendingMigrations(ctx)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to check table %s: %w", table, err)
|
||||
return fmt.Errorf("failed to check pending migrations: %w", 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")
|
||||
|
||||
@@ -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;
|
||||
@@ -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.';
|
||||
@@ -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();
|
||||
@@ -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'));
|
||||
@@ -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'));
|
||||
@@ -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)';
|
||||
@@ -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';
|
||||
@@ -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';
|
||||
@@ -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)';
|
||||
@@ -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)';
|
||||
@@ -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';
|
||||
+10
-25
@@ -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)';
|
||||
@@ -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';
|
||||
@@ -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';
|
||||
@@ -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';
|
||||
@@ -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';
|
||||
@@ -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)';
|
||||
@@ -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!'
|
||||
@@ -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';
|
||||
@@ -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';
|
||||
@@ -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)';
|
||||
@@ -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';
|
||||
@@ -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';
|
||||
@@ -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;
|
||||
@@ -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!'
|
||||
@@ -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';
|
||||
@@ -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();
|
||||
@@ -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();
|
||||
@@ -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"`
|
||||
|
||||
@@ -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,26 +112,18 @@ 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
|
||||
func (q *Queue) GetStats() map[int]worker.Stats {
|
||||
q.mu.RLock()
|
||||
|
||||
+111
-48
@@ -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,29 +193,27 @@ 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 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)
|
||||
if err := w.setTenantTx(ctx, tx); err != nil {
|
||||
tx.Rollback()
|
||||
w.logger.Error("failed to set tenant", "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 {
|
||||
@@ -178,53 +222,72 @@ func (w *Worker) processJobs(ctx context.Context, specificJob *models.Job) {
|
||||
}
|
||||
|
||||
// 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()))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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())
|
||||
}
|
||||
@@ -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:])
|
||||
}
|
||||
@@ -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")
|
||||
_, err := db.Exec("DROP SCHEMA IF EXISTS broker 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)
|
||||
}
|
||||
t.Logf("Warning: failed to drop broker schema: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user