Add schema-scoped persistence

go.mod

@@ -40,7 +40,7 @@ require (
 )
 
 require (
-	github.com/flanksource/commons-db v0.1.26
+	github.com/flanksource/commons-db v0.1.27-0.20260806180738-8eb0bca01132
 	github.com/gliderlabs/ssh v0.3.8
 	github.com/pelletier/go-toml/v2 v2.4.3
 )

go.sum

@@ -280,8 +280,8 @@ github.com/flanksource/clicky/aichat v1.21.48 h1:f8Kvl96Lfp1qcqPuve1zsjaYN8ZcK1/
 github.com/flanksource/clicky/aichat v1.21.48/go.mod h1:PGN/lVAgxpchRctciUCpR4YIuqWoDwRxDh339A6wi3w=
 github.com/flanksource/commons v1.55.0 h1:gj9zBY3V1qgAAnEiLaeGbkqCmNK0p1tJVQCDurdTZ2k=
 github.com/flanksource/commons v1.55.0/go.mod h1:gupTCRqGpgD8dd2ooE7bMDJxkfcVKvkPVuHg3cWkW+Q=
-github.com/flanksource/commons-db v0.1.26 h1:NXAP0WvMs4ufyDfl1L2ryRxBv5qxV67GiI1nINd4YIw=
-github.com/flanksource/commons-db v0.1.26/go.mod h1:i378WIxy8g9xOeLBvhh1y3FO99oCumUqNmfhuDF79kM=
+github.com/flanksource/commons-db v0.1.27-0.20260806180738-8eb0bca01132 h1:dbBt8TmT2tEE9y3zEUJh6bGbB7gH5JisP4BbhXoYvmU=
+github.com/flanksource/commons-db v0.1.27-0.20260806180738-8eb0bca01132/go.mod h1:i378WIxy8g9xOeLBvhh1y3FO99oCumUqNmfhuDF79kM=
 github.com/flanksource/gomplate/v3 v3.24.84 h1:UOE0yCJsczTIKRaHUvhD6tjCYrbNvOugAizuy0FVlhE=
 github.com/flanksource/gomplate/v3 v3.24.84/go.mod h1:NMMZkFsjbLy/8iY8Fip5N86Y0PP6lZeq+kmPwpVVIL0=
 github.com/flanksource/is-healthy v1.0.88 h1:ATQuKoNdp8Qfzf41/eMFajmT0qzOmZlZNG5eLK41RFo=

migrations/concurrency_integration_test.go

@@ -17,7 +17,7 @@ func TestConcurrentApplySerializesCaptainMigrations(t *testing.T) {
 	// Hold the same session lock before releasing a group of Apply calls. This
 	// proves every caller enters through the advisory-lock boundary rather than
 	// racing the Atlas inspect/diff/apply window.
-	blocker, err := acquireMigrationLock(t.Context(), dsn)
+	blocker, err := acquireMigrationLock(t.Context(), applyRequest{Connection: dsn, Schema: DefaultSchema})
 	require.NoError(t, err)
 	t.Cleanup(func() { require.NoError(t, blocker.Close()) })
 
@@ -63,7 +63,7 @@ func TestConcurrentApplySerializesCaptainMigrations(t *testing.T) {
 	// Bound the reacquisition to catch a leaked dedicated connection cleanly.
 	reacquireCtx, cancel := context.WithTimeout(t.Context(), 2*time.Second)
 	defer cancel()
-	reacquired, err := acquireMigrationLock(reacquireCtx, dsn)
+	reacquired, err := acquireMigrationLock(reacquireCtx, applyRequest{Connection: dsn, Schema: DefaultSchema})
 	require.NoError(t, err)
 	require.NoError(t, reacquired.Close())
 }

migrations/migrations.go

@@ -5,8 +5,10 @@ package migrations
 
 import (
 	"context"
+	"crypto/sha256"
 	"database/sql"
 	"embed"
+	"encoding/binary"
 	"errors"
 	"fmt"
 	"strings"
@@ -18,6 +20,7 @@ import (
 )
 
 const Scope = "captain"
+const DefaultSchema = "public"
 
 const (
 	// captainMigrationLockNamespace and captainMigrationLockKey are stable,
@@ -42,16 +45,38 @@ type migrationLockHandle interface {
 }
 
 type applyDependencies struct {
-	acquireLock func(context.Context, string) (migrationLockHandle, error)
-	migrate     func(context.Context, string) error
-	verify      func(context.Context, string) error
+	acquireLock func(context.Context, applyRequest) (migrationLockHandle, error)
+	migrate     func(context.Context, applyRequest) error
+	verify      func(context.Context, applyRequest) error
+}
+
+type applyRequest struct {
+	Connection string
+	Schema     string
+}
+
+type options struct {
+	schema string
+}
+
+// Option configures Captain's migration bundle.
+type Option func(*options)
+
+// WithSchema selects the schema that owns Captain's migration bundle.
+func WithSchema(name string) Option {
+	return func(options *options) { options.schema = name }
 }
 
 var defaultApplyDependencies = applyDependencies{
 	acquireLock: acquireMigrationLock,
-	migrate: func(ctx context.Context, connection string) error {
-		return commonsmigrate.Apply(ctx, connection, schemaFS,
+	migrate: func(ctx context.Context, request applyRequest) error {
+		filesystem, err := schemaFilesystem(request.Schema)
+		if err != nil {
+			return err
+		}
+		return commonsmigrate.Apply(ctx, request.Connection, filesystem,
 			commonsmigrate.WithName(Scope),
+			commonsmigrate.WithSchema(request.Schema),
 			commonsmigrate.WithExclude("todo_*"),
 		)
 	},
@@ -63,16 +88,25 @@ var defaultApplyDependencies = applyDependencies{
 // migration bundle across processes. It is safe to call repeatedly and uses a
 // stable scope so Captain can share a database with other independently
 // migrated applications.
-func Apply(ctx context.Context, connection string) error {
-	return apply(ctx, connection, defaultApplyDependencies)
+func Apply(ctx context.Context, connection string, optionFns ...Option) error {
+	config := options{schema: DefaultSchema}
+	for _, option := range optionFns {
+		if option != nil {
+			option(&config)
+		}
+	}
+	return apply(ctx, applyRequest{Connection: connection, Schema: config.schema}, defaultApplyDependencies)
 }
 
-func apply(ctx context.Context, connection string, deps applyDependencies) (resultErr error) {
-	if strings.TrimSpace(connection) == "" {
+func apply(ctx context.Context, request applyRequest, deps applyDependencies) (resultErr error) {
+	if strings.TrimSpace(request.Connection) == "" {
 		return errors.New("captain migration connection string is empty")
 	}
+	if err := commonsmigrate.ValidateSchemaName(request.Schema); err != nil {
+		return fmt.Errorf("captain migration schema: %w", err)
+	}
 
-	lock, err := deps.acquireLock(ctx, connection)
+	lock, err := deps.acquireLock(ctx, request)
 	if err != nil {
 		return fmt.Errorf("acquire Captain migration lock: %w", err)
 	}
@@ -82,16 +116,20 @@ func apply(ctx context.Context, connection string, deps applyDependencies) (resu
 		}
 	}()
 
-	if err := deps.migrate(ctx, connection); err != nil {
+	if err := deps.migrate(ctx, request); err != nil {
 		return fmt.Errorf("migrate Captain database: %w", err)
 	}
-	if err := deps.verify(ctx, connection); err != nil {
+	if err := deps.verify(ctx, request); err != nil {
 		return fmt.Errorf("verify Captain database: %w", err)
 	}
 	return nil
 }
 
-func verifyToolApprovalIdentity(ctx context.Context, connection string) (resultErr error) {
+func verifyToolApprovalIdentity(ctx context.Context, request applyRequest) (resultErr error) {
+	connection, err := commonsmigrate.ConnectionForSchema(request.Connection, request.Schema)
+	if err != nil {
+		return fmt.Errorf("scope schema verification database: %w", err)
+	}
 	db, err := commonsdb.NewDB(connection)
 	if err != nil {
 		return fmt.Errorf("open schema verification database: %w", err)
@@ -109,11 +147,11 @@ func verifyToolApprovalIdentity(ctx context.Context, connection string) (resultE
 		FROM pg_constraint c
 		JOIN pg_class relation ON relation.oid = c.conrelid
 		JOIN pg_namespace namespace ON namespace.oid = relation.relnamespace
-		WHERE namespace.nspname = 'public'
+		WHERE namespace.nspname = $1
 		  AND relation.relname = 'captain_turn_requests'
 		  AND c.conname = 'captain_turn_requests_tool_approval_identity'
 		  AND c.contype = 'c'
-	`).Scan(&validated, &definition)
+	`, request.Schema).Scan(&validated, &definition)
 	if errors.Is(err, sql.ErrNoRows) {
 		return errors.New("captain_turn_requests_tool_approval_identity constraint is missing")
 	}
@@ -138,13 +176,14 @@ func verifyToolApprovalIdentity(ctx context.Context, connection string) (resultE
 type migrationLock struct {
 	db   *sql.DB
 	conn *sql.Conn
+	key  int32
 
 	once sync.Once
 	err  error
 }
 
-func acquireMigrationLock(ctx context.Context, connection string) (migrationLockHandle, error) {
-	db, err := commonsdb.NewDB(connection)
+func acquireMigrationLock(ctx context.Context, request applyRequest) (migrationLockHandle, error) {
+	db, err := commonsdb.NewDB(request.Connection)
 	if err != nil {
 		return nil, fmt.Errorf("open advisory-lock database: %w", err)
 	}
@@ -154,12 +193,12 @@ func acquireMigrationLock(ctx context.Context, connection string) (migrationLock
 		return nil, fmt.Errorf("reserve advisory-lock connection: %w", err)
 	}
 	if _, err := conn.ExecContext(ctx, `SELECT pg_advisory_lock($1, $2)`,
-		captainMigrationLockNamespace, captainMigrationLockKey); err != nil {
+		captainMigrationLockNamespace, migrationLockKey(request.Schema)); err != nil {
 		_ = conn.Close()
 		_ = db.Close()
 		return nil, fmt.Errorf("lock Captain migration scope: %w", err)
 	}
-	return &migrationLock{db: db, conn: conn}, nil
+	return &migrationLock{db: db, conn: conn, key: migrationLockKey(request.Schema)}, nil
 }
 
 func (lock *migrationLock) Close() error {
@@ -172,7 +211,7 @@ func (lock *migrationLock) Close() error {
 			ctx, cancel := context.WithTimeout(context.Background(), migrationUnlockTimeout)
 			var unlocked bool
 			if err := lock.conn.QueryRowContext(ctx, `SELECT pg_advisory_unlock($1, $2)`,
-				captainMigrationLockNamespace, captainMigrationLockKey).Scan(&unlocked); err != nil {
+				captainMigrationLockNamespace, lock.key).Scan(&unlocked); err != nil {
 				cleanupErrors = append(cleanupErrors, fmt.Errorf("unlock Captain migration scope: %w", err))
 			} else if !unlocked {
 				cleanupErrors = append(cleanupErrors, errors.New("captain migration advisory lock was not held"))
@@ -191,3 +230,11 @@ func (lock *migrationLock) Close() error {
 	})
 	return lock.err
 }
+
+func migrationLockKey(schemaName string) int32 {
+	if schemaName == DefaultSchema {
+		return captainMigrationLockKey
+	}
+	digest := sha256.Sum256([]byte(schemaName))
+	return int32(binary.BigEndian.Uint32(digest[:4]))
+}

migrations/migrations_test.go

@@ -231,16 +231,16 @@ func TestApplyHoldsMigrationLockAcrossMigration(t *testing.T) {
 	t.Parallel()
 
 	var events []string
-	err := apply(t.Context(), "postgres://captain", applyDependencies{
-		acquireLock: func(context.Context, string) (migrationLockHandle, error) {
+	err := apply(t.Context(), applyRequest{Connection: "postgres://captain", Schema: DefaultSchema}, applyDependencies{
+		acquireLock: func(context.Context, applyRequest) (migrationLockHandle, error) {
 			events = append(events, "lock")
 			return &recordingMigrationLock{events: &events}, nil
 		},
-		migrate: func(context.Context, string) error {
+		migrate: func(context.Context, applyRequest) error {
 			events = append(events, "migrate")
 			return nil
 		},
-		verify: func(context.Context, string) error {
+		verify: func(context.Context, applyRequest) error {
 			events = append(events, "verify")
 			return nil
 		},
@@ -256,16 +256,16 @@ func TestApplyReleasesMigrationLockOnVerificationFailure(t *testing.T) {
 
 	var events []string
 	verificationErr := errors.New("constraint drifted")
-	err := apply(t.Context(), "postgres://captain", applyDependencies{
-		acquireLock: func(context.Context, string) (migrationLockHandle, error) {
+	err := apply(t.Context(), applyRequest{Connection: "postgres://captain", Schema: DefaultSchema}, applyDependencies{
+		acquireLock: func(context.Context, applyRequest) (migrationLockHandle, error) {
 			events = append(events, "lock")
 			return &recordingMigrationLock{events: &events}, nil
 		},
-		migrate: func(context.Context, string) error {
+		migrate: func(context.Context, applyRequest) error {
 			events = append(events, "migrate")
 			return nil
 		},
-		verify: func(context.Context, string) error {
+		verify: func(context.Context, applyRequest) error {
 			events = append(events, "verify")
 			return verificationErr
 		},
@@ -281,12 +281,12 @@ func TestApplyReleasesMigrationLockOnMigrationFailure(t *testing.T) {
 
 	var events []string
 	migrationErr := errors.New("atlas failed")
-	err := apply(t.Context(), "postgres://captain", applyDependencies{
-		acquireLock: func(context.Context, string) (migrationLockHandle, error) {
+	err := apply(t.Context(), applyRequest{Connection: "postgres://captain", Schema: DefaultSchema}, applyDependencies{
+		acquireLock: func(context.Context, applyRequest) (migrationLockHandle, error) {
 			events = append(events, "lock")
 			return &recordingMigrationLock{events: &events}, nil
 		},
-		migrate: func(context.Context, string) error {
+		migrate: func(context.Context, applyRequest) error {
 			events = append(events, "migrate")
 			return migrationErr
 		},
@@ -303,8 +303,8 @@ func TestApplyReportsLockAcquisitionAndReleaseErrors(t *testing.T) {
 	t.Run("acquire", func(t *testing.T) {
 		t.Parallel()
 		wantErr := errors.New("lock unavailable")
-		err := apply(t.Context(), "postgres://captain", applyDependencies{
-			acquireLock: func(context.Context, string) (migrationLockHandle, error) {
+		err := apply(t.Context(), applyRequest{Connection: "postgres://captain", Schema: DefaultSchema}, applyDependencies{
+			acquireLock: func(context.Context, applyRequest) (migrationLockHandle, error) {
 				return nil, wantErr
 			},
 		})
@@ -317,11 +317,11 @@ func TestApplyReportsLockAcquisitionAndReleaseErrors(t *testing.T) {
 		t.Parallel()
 		migrationErr := errors.New("migration failed")
 		releaseErr := errors.New("unlock failed")
-		err := apply(t.Context(), "postgres://captain", applyDependencies{
-			acquireLock: func(context.Context, string) (migrationLockHandle, error) {
+		err := apply(t.Context(), applyRequest{Connection: "postgres://captain", Schema: DefaultSchema}, applyDependencies{
+			acquireLock: func(context.Context, applyRequest) (migrationLockHandle, error) {
 				return &recordingMigrationLock{err: releaseErr}, nil
 			},
-			migrate: func(context.Context, string) error { return migrationErr },
+			migrate: func(context.Context, applyRequest) error { return migrationErr },
 		})
 		if !errors.Is(err, migrationErr) || !errors.Is(err, releaseErr) {
 			t.Fatalf("apply error = %v, want joined migration and release errors", err)

migrations/schema.go

@@ -0,0 +1,42 @@
+package migrations
+
+import (
+	"fmt"
+	"io/fs"
+	"path"
+	"strings"
+	"testing/fstest"
+
+	commonsmigrate "github.com/flanksource/commons-db/migrate"
+)
+
+func schemaFilesystem(schemaName string) (fs.FS, error) {
+	if err := commonsmigrate.ValidateSchemaName(schemaName); err != nil {
+		return nil, fmt.Errorf("captain migration schema: %w", err)
+	}
+	if schemaName == DefaultSchema {
+		return schemaFS, nil
+	}
+	files := fstest.MapFS{}
+	err := fs.WalkDir(schemaFS, ".", func(name string, entry fs.DirEntry, walkErr error) error {
+		if walkErr != nil {
+			return walkErr
+		}
+		if entry.IsDir() {
+			return nil
+		}
+		content, err := fs.ReadFile(schemaFS, name)
+		if err != nil {
+			return fmt.Errorf("read Captain migration %s: %w", name, err)
+		}
+		if strings.EqualFold(path.Ext(name), ".sql") {
+			content = []byte(strings.ReplaceAll(string(content), DefaultSchema+".", schemaName+"."))
+		}
+		files[name] = &fstest.MapFile{Data: content}
+		return nil
+	})
+	if err != nil {
+		return nil, fmt.Errorf("render Captain migration schema: %w", err)
+	}
+	return files, nil
+}

migrations/schema_ginkgo_test.go

@@ -0,0 +1,46 @@
+package migrations
+
+import (
+	"io/fs"
+	"strings"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("schema-scoped Captain migrations", func() {
+	It("leaves the public bundle unchanged", func() {
+		filesystem, err := schemaFilesystem(DefaultSchema)
+		Expect(err).NotTo(HaveOccurred())
+		content, err := fs.ReadFile(filesystem, "51_state_triggers.sql")
+		Expect(err).NotTo(HaveOccurred())
+		Expect(string(content)).To(ContainSubstring("public.captain_sessions"))
+	})
+
+	It("qualifies SQL objects with the selected schema while retaining portable HCL", func() {
+		const schemaName = "agent_namespace_context"
+		filesystem, err := schemaFilesystem(schemaName)
+		Expect(err).NotTo(HaveOccurred())
+
+		sqlContent, err := fs.ReadFile(filesystem, "51_state_triggers.sql")
+		Expect(err).NotTo(HaveOccurred())
+		Expect(string(sqlContent)).To(ContainSubstring(schemaName + ".captain_sessions"))
+		Expect(string(sqlContent)).NotTo(ContainSubstring(DefaultSchema + ".captain_"))
+
+		hclContent, err := fs.ReadFile(filesystem, "10_sessions.pg.hcl")
+		Expect(err).NotTo(HaveOccurred())
+		Expect(string(hclContent)).To(ContainSubstring("schema.public"))
+		Expect(string(hclContent)).NotTo(ContainSubstring(schemaName))
+	})
+
+	It("rejects invalid schemas", func() {
+		_, err := schemaFilesystem(strings.Repeat("x", 64))
+		Expect(err).To(HaveOccurred())
+	})
+
+	It("uses a stable schema-specific advisory lock", func() {
+		Expect(migrationLockKey(DefaultSchema)).To(Equal(captainMigrationLockKey))
+		Expect(migrationLockKey("agent_namespace_one")).To(Equal(migrationLockKey("agent_namespace_one")))
+		Expect(migrationLockKey("agent_namespace_one")).NotTo(Equal(migrationLockKey("agent_namespace_two")))
+	})
+})