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.
2026-09-17 22:09:41 +02:00
parent 602997bcdb
commit 4c8e1066d4
53 changed files with 2911 additions and 1016 deletions
@@ -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")
-		if err != nil {
-			t.Logf("Warning: failed to drop procedure %s: %v", proc, err)
-		}
-	}
-
-	// Drop tables
-	for _, table := range tables {
-		_, err := db.Exec("DROP TABLE IF EXISTS " + table + " CASCADE")
-		if err != nil {
-			t.Logf("Warning: failed to drop table %s: %v", table, err)
-		}
+	_, err := db.Exec("DROP SCHEMA IF EXISTS broker CASCADE")
+	if err != nil {
+		t.Logf("Warning: failed to drop broker schema: %v", err)
 	}
 }