diff --git a/CHANGELOG.md b/CHANGELOG.md index 7a42ce1..c9e9316 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,17 @@ All notable changes to TigerFS will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [Unreleased] + +### Added +- Per-query session variable scoping via `SET LOCAL` in transactions. + Enables multi-tenant RLS from a shared connection pool, compatible with + PgBouncer (session + transaction mode) and RDS Proxy. Configure via + `--session-var key=value` CLI flag, `session_variables` config, or + `db.WithSessionVars()` library API. +- `DBTX` interface abstracting `*pgxpool.Pool` and `pgx.Tx`, enabling + transparent transaction wrapping for session variable injection. + ## [0.6.0] - 2026-03-26 **Dedicated tigerfs schema, security hardening, and unified demo.** diff --git a/docs/implementation/session-variable-scoping-plan.md b/docs/implementation/session-variable-scoping-plan.md new file mode 100644 index 0000000..6f6b5b4 --- /dev/null +++ b/docs/implementation/session-variable-scoping-plan.md @@ -0,0 +1,256 @@ +# Session Variable Scoping via SET LOCAL + +> **Status: Implemented** on branch `feature/session-var-scoping` + +## Summary + +Per-query PostgreSQL session variable scoping using `SET LOCAL` inside +transactions. Enables multi-tenant RLS from a single shared connection +pool — compatible with **all** PostgreSQL deployment topologies: + +- Direct PostgreSQL +- PgBouncer session mode +- PgBouncer transaction mode +- RDS Proxy (no connection pinning) + +Session variables flow through `context.Context` using the standard Go +idiom. A `DBTX` interface abstracts over `*pgxpool.Pool` and `pgx.Tx`, +allowing transparent transaction wrapping when session vars are present — +with zero overhead when they're not. + +## Problem + +TigerFS creates one connection pool per mount. For multi-tenant RLS, the +only option was encoding session variables in the connection string: + +``` +postgres://user:pass@host/db?options=-c%20app.user_id=42%20-c%20app.tenant_id=acme +``` + +This forces one pool per distinct RLS identity → N users × pool_size +backend connections. PgBouncer and RDS Proxy can't help because session +state is set at connection time and either causes pinning (RDS Proxy) or +gets lost on connection swap (PgBouncer transaction mode). + +## Design + +### Why SET LOCAL (not session-scoped set_config) + +An earlier proposal used `set_config($1, $2, false)` (session-scoped) in +`BeforeAcquire` hooks with surgical `RESET` in `AfterRelease`. This approach: + +- **Does not work with PgBouncer transaction mode** — session state gets + lost when the pooler reassigns the backend connection +- **Causes pinning in RDS Proxy** — any `SET` statement pins the connection, + defeating the pooler's multiplexing +- **Requires tracking and cleanup** — a `sync.Map` to track which keys + were set per connection, and `RESET` for each on release + +`SET LOCAL` (via `set_config($1, $2, true)`) avoids all of these problems. +Variables exist only for the duration of the transaction, then vanish +automatically. No cleanup, no tracking, no pinning. + +### DBTX Interface + +The core abstraction is a `DBTX` interface satisfied by both +`*pgxpool.Pool` and `pgx.Tx`: + +```go +type DBTX interface { + Exec(ctx context.Context, sql string, args ...any) (pgconn.CommandTag, error) + Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error) + QueryRow(ctx context.Context, sql string, args ...any) pgx.Row +} +``` + +All ~60 package-level db functions were refactored from `pool *pgxpool.Pool` +to `q DBTX`. This is a mechanical, zero-behavior-change refactor since +`*pgxpool.Pool` satisfies `DBTX`. + +### acquireDBTX Pattern + +The `Client.acquireDBTX(ctx)` method is the decision point: + +- **No session vars**: returns `c.pool` directly with a no-op cleanup + function. Zero overhead — no transaction started. +- **Session vars present**: begins a transaction, runs `SET LOCAL` for + each variable via `set_config($1, $2, true)`, returns the `pgx.Tx` as + a `DBTX`. The cleanup function commits on success, rolls back on error. + +```go +func (c *Client) acquireDBTX(ctx context.Context) (DBTX, doneFunc, error) { + vars := c.effectiveSessionVars(ctx) + if len(vars) == 0 { + return c.pool, func(error) {}, nil // zero overhead + } + tx, err := c.pool.Begin(ctx) + // SET LOCAL each var via applySessionVars + // return tx + done func +} +``` + +Client methods use the pattern: + +```go +func (c *Client) GetRow(ctx context.Context, ...) (result *Row, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { return nil, err } + defer func() { done(retErr) }() + return GetRow(ctx, q, ...) +} +``` + +### No Nested Transaction Risk + +Package-level functions like `GetFullDDL` and `GetRowCountSmart` call +other package-level functions, threading the same `DBTX` through. Only +the top-level Client method starts a transaction via `acquireDBTX`. All +sub-calls share the same `DBTX` instance — no nested `BEGIN`. + +### Baseline + Context Merge + +Session variables come from two sources: + +1. **Baseline vars** — mount-level, from `--session-var` flag or + `session_variables` config. Stored on `Client.baselineVars`. +2. **Context vars** — per-request, from `db.WithSessionVars(ctx, vars)`. + +`effectiveSessionVars(ctx)` merges them: baseline first, context overrides. + +### Existing Transaction Methods + +Methods that already manage their own transactions (`ImportOverwrite`, +`ImportSync`, `ImportAppend`, `ExecInTransaction`) inject `SET LOCAL` +at the start of their existing transaction via `applySessionVars()` — +no double-wrapping needed. + +## API + +### Library + +```go +// Attach session vars to a context +ctx = db.WithSessionVars(ctx, db.SessionVars{ + "app.user_id": userID, + "app.tenant_id": tenantID, +}) + +// All Client methods respect session vars from context +row, err := client.GetRow(ctx, schema, table, pk) // SET LOCAL applied +count, err := client.GetRowCount(ctx, schema, table) // SET LOCAL applied +``` + +### CLI + +```bash +tigerfs mount --session-var app.user_id=42 --session-var app.tenant_id=acme \ + postgres://host/db /mnt/db +``` + +Uses pflag's `StringToStringVar` — no custom parser. + +### Config + +```yaml +session_variables: + app.user_id: "42" + app.tenant_id: "acme" +``` + +No separate `session_scoping: true` flag — if `session_variables` is +non-empty, session scoping is active. + +## Files + +### New + +| File | Purpose | +|------|---------| +| `db/dbtx.go` | `DBTX` interface + `applySessionVars` helper | +| `db/session_vars.go` | `SessionVars` type, `WithSessionVars`, `SessionVarsFromContext` | +| `db/dbtx_test.go` | Interface satisfaction tests | +| `db/session_vars_test.go` | Context API + effectiveSessionVars + acquireDBTX unit tests | +| `test/integration/session_vars_test.go` | RLS isolation, leak prevention, import, override tests | + +### Modified + +| File | Change | +|------|--------| +| `db/client.go` | `baselineVars` field, `effectiveSessionVars`, `acquireDBTX`, `Exec`/`ExecInTransaction` updated | +| `db/schema.go` | `*pgxpool.Pool` → `DBTX` in 23 functions; acquireDBTX in 23 Client methods | +| `db/query.go` | `*pgxpool.Pool` → `DBTX` in 12 functions; acquireDBTX in 22 Client methods; `queryRows` takes `DBTX` param | +| `db/indexes.go` | `*pgxpool.Pool` → `DBTX` in 11 functions; acquireDBTX in 10 Client methods | +| `db/export.go` | `*pgxpool.Pool` → `DBTX` in 3 functions; acquireDBTX in 5 Client methods | +| `db/keys.go` | `*pgxpool.Pool` → `DBTX` in 3 functions; acquireDBTX in 3 Client methods | +| `db/permissions.go` | `*pgxpool.Pool` → `DBTX` in 2 functions; acquireDBTX in 2 Client methods | +| `db/pipeline.go` | `*pgxpool.Pool` → `DBTX` in 2 functions; acquireDBTX in 2 Client methods | +| `db/constraints.go` | `*pgxpool.Pool` → `DBTX` in 4 functions | +| `db/import.go` | `applySessionVars` injected into 3 existing transaction methods; `truncateTable` uses acquireDBTX | +| `config/config.go` | `SessionVariables map[string]string` field added | +| `cmd/mount.go` | `--session-var` flag via `StringToStringVar` | + +### Unchanged + +- `db/mocks.go` — mocks implement role interfaces, not `DBTX` +- `db/interfaces.go` — `DBClient` interface signatures unchanged +- All callers of `DBClient` (fs, fuse, nfs packages) — zero changes + +## Tests + +### Unit (12 tests in `db/session_vars_test.go`) + +- `TestWithSessionVars_NilReturnsOriginal` +- `TestWithSessionVars_EmptyReturnsOriginal` +- `TestSessionVarsFromContext_MissingReturnsNil` +- `TestSessionVarsFromContext_RoundTrip` +- `TestSessionVarsFromContext_OverwritesPrevious` +- `TestEffectiveSessionVars_NoVars` +- `TestEffectiveSessionVars_BaselineOnly` +- `TestEffectiveSessionVars_ContextOnly` +- `TestEffectiveSessionVars_ContextOverridesBaseline` +- `TestEffectiveSessionVars_Merge` +- `TestAcquireDBTX_NilPoolReturnsError` +- `TestAcquireDBTX_NoVarsDoneFuncIsNoOp` + +### Unit (2 tests in `db/dbtx_test.go`) + +- `TestDBTX_PoolSatisfiesInterface` +- `TestDBTX_TxSatisfiesInterface` + +### Unit (1 test in `cmd/mount_test.go`) + +- `TestBuildMountCmd_SessionVarFlag` + +### Integration (6 tests in `test/integration/session_vars_test.go`) + +- `TestSessionVars_SetLocalApplied` — verifies SET LOCAL is visible within queries +- `TestSessionVars_SetLocalDoesNotLeak` — verifies variables don't persist after tx +- `TestSessionVars_ContextOverridesBaseline` — baseline + context merge +- `TestSessionVars_NoVarsNoOverhead` — operations work without session vars +- `TestSessionVars_RLSIsolation` — end-to-end RLS with non-superuser role, + shared pool serving alice and bob contexts, zero rows without vars +- `TestSessionVars_ImportWithVars` — import operations apply session vars + +### RLS Test Note + +PostgreSQL superusers always bypass RLS — this is a PG design decision, +not a TigerFS limitation. The RLS integration test creates a non-superuser +role for realistic testing. If role creation fails (insufficient privileges), +the test skips gracefully via `t.Skip`. + +## Compatibility Matrix + +| Deployment | Session-scoped (`set_config(..., false)`) | **SET LOCAL (implemented)** | +|---|---|---| +| Direct PostgreSQL | Works | Works | +| PgBouncer session mode | Works | Works | +| PgBouncer transaction mode | Broken (state lost) | **Works** | +| RDS Proxy | Pinned (defeats pooling) | **Works (no pinning)** | + +## Future Work + +- **Daemon mode**: a future `tigerfs serve` process could own a single + pool and use `db.WithSessionVars(ctx, vars)` to scope queries per mount. + The context API introduced here is daemon-ready. +- **Benchmarks**: measure per-query overhead of `BEGIN; SET LOCAL; query; COMMIT` + vs. direct query. Expected to be small relative to network RTT. diff --git a/docs/spec.md b/docs/spec.md index ae94141..faf97b7 100644 --- a/docs/spec.md +++ b/docs/spec.md @@ -1642,6 +1642,16 @@ connection: default_mount_dir: /tmp # Base directory for auto-generated mountpoints insecure_no_ssl: false # Skip TLS enforcement for remote connections + # Session variables — applied via SET LOCAL to every query in a transaction. + # Enables multi-tenant RLS from a shared connection pool. Compatible with + # PgBouncer (session + transaction mode) and RDS Proxy (no pinning). + # Variables are set per-query and automatically cleared when the transaction + # ends. Library consumers can additionally set per-request variables via + # db.WithSessionVars(). + session_variables: + # app.user_id: "42" + # app.tenant_id: "acme" + # Security note: TigerFS enforces sslmode=require for all non-localhost # database connections. This ensures connections to remote databases are # always encrypted. To disable this enforcement (not recommended), use @@ -1769,6 +1779,11 @@ tigerfs --foreground --log-level=debug postgres://localhost/mydb /mnt/db --allow-root Allow root + mounting user to access (FUSE) ``` +**Session Variables:** +```bash +--session-var KEY=VALUE Set a PostgreSQL session variable via SET LOCAL (repeatable) +``` + **Caching/Performance:** ```bash --attr-timeout SECS FUSE attribute cache timeout (FUSE backend only, default: 1) diff --git a/internal/tigerfs/cmd/mount.go b/internal/tigerfs/cmd/mount.go index d6207ef..f221473 100644 --- a/internal/tigerfs/cmd/mount.go +++ b/internal/tigerfs/cmd/mount.go @@ -46,6 +46,7 @@ func buildMountCmd(ctx context.Context) *cobra.Command { var maxPipelineDepth int var legacyFuse bool var insecureNoSSL bool + var sessionVars map[string]string var userID string var autoSavepointInterval time.Duration var undoListLimit int @@ -165,6 +166,14 @@ Examples: if insecureNoSSL { cfg.InsecureNoSSL = true } + if len(sessionVars) > 0 { + if cfg.SessionVariables == nil { + cfg.SessionVariables = make(map[string]string) + } + for k, v := range sessionVars { + cfg.SessionVariables[k] = v + } + } // User identity: flag > env > empty (anonymous) if userID != "" { @@ -243,6 +252,7 @@ Examples: cmd.Flags().IntVar(&maxPipelineDepth, "max-pipeline-depth", 0, "max chained pipeline ops before capabilities are hidden; 0 uses config default") cmd.Flags().BoolVar(&legacyFuse, "legacy-fuse", false, "use legacy FUSE node tree (Linux only)") cmd.Flags().BoolVar(&insecureNoSSL, "insecure-no-ssl", false, "allow non-TLS connections to remote databases (insecure)") + cmd.Flags().StringToStringVar(&sessionVars, "session-var", nil, "set a PostgreSQL session variable via SET LOCAL (key=value, repeatable)") cmd.Flags().StringVar(&userID, "user-id", "", "user identity for undo log entries (also: TIGERFS_USER_ID env)") cmd.Flags().DurationVar(&autoSavepointInterval, "auto-savepoint-interval", 0, "inactivity gap before auto-savepoint (e.g., 30m); 0 uses config default") cmd.Flags().IntVar(&undoListLimit, "undo-list-limit", 0, "default listing limit for .undo/ sub-directories; 0 uses config default (100)") diff --git a/internal/tigerfs/cmd/mount_test.go b/internal/tigerfs/cmd/mount_test.go index 48dbb23..a12896b 100644 --- a/internal/tigerfs/cmd/mount_test.go +++ b/internal/tigerfs/cmd/mount_test.go @@ -1,6 +1,7 @@ package cmd import ( + "context" "path/filepath" "testing" ) @@ -78,3 +79,18 @@ func TestResolveMountArgs(t *testing.T) { // Note: TestSanitizeConnectionString has been moved to db/connection_test.go // where the SanitizeConnectionString function now lives. + +// TestBuildMountCmd_SessionVarFlag verifies the --session-var flag exists +// and is wired correctly using pflag's StringToString type. +func TestBuildMountCmd_SessionVarFlag(t *testing.T) { + ctx := context.Background() + cmd := buildMountCmd(ctx) + + flag := cmd.Flags().Lookup("session-var") + if flag == nil { + t.Fatal("--session-var flag not found") + } + if flag.DefValue != "[]" { + t.Logf("default value: %q", flag.DefValue) + } +} diff --git a/internal/tigerfs/config/config.go b/internal/tigerfs/config/config.go index 37f5b1d..9f873b0 100644 --- a/internal/tigerfs/config/config.go +++ b/internal/tigerfs/config/config.go @@ -68,6 +68,9 @@ type Config struct { DefaultFormat string `mapstructure:"default_format"` BinaryEncoding string `mapstructure:"binary_encoding"` + // Session Variables + SessionVariables map[string]string `mapstructure:"session_variables"` // Applied via SET LOCAL to every query (e.g. {"app.user_id": "42"}) + // Identity UserID string `mapstructure:"user_id"` // Mount-level user identity for log entries (--user-id or TIGERFS_USER_ID) diff --git a/internal/tigerfs/db/client.go b/internal/tigerfs/db/client.go index 2f6bcd2..2d30510 100644 --- a/internal/tigerfs/db/client.go +++ b/internal/tigerfs/db/client.go @@ -16,8 +16,84 @@ import ( // Client represents a PostgreSQL database client with connection pooling type Client struct { - pool *pgxpool.Pool - cfg *config.Config + pool *pgxpool.Pool + cfg *config.Config + baselineVars SessionVars // mount-level session variables from config/CLI +} + +// doneFunc is called after a query completes to commit or rollback the +// session variable transaction. When no session vars are active, this is a no-op. +type doneFunc func(err error) + +// effectiveSessionVars merges baseline (mount-level) vars with per-request +// context vars. Context vars override baseline vars for the same key. +// Returns a zero-value SessionVars if no vars are active. +func (c *Client) effectiveSessionVars(ctx context.Context) SessionVars { + ctxVars := SessionVarsFromContext(ctx) + if c.baselineVars.Empty() && ctxVars.Empty() { + return SessionVars{} + } + // Fast path: baseline only — return directly (immutable, pre-sorted). + if ctxVars.Empty() { + return c.baselineVars + } + // Context only — return directly (pre-sorted at construction). + if c.baselineVars.Empty() { + return ctxVars + } + // Merge: baseline first, context overrides. Re-sorts once. + return c.baselineVars.Merge(ctxVars) +} + +// acquireDBTX returns a DBTX with session variables applied. +// When no session vars are needed, returns the pool directly with a no-op +// done func — zero overhead. When session vars exist, begins a transaction, +// applies SET LOCAL via set_config, and returns the transaction as a DBTX. +// +// The returned doneFunc must be called with the operation's error to commit +// (on success) or rollback (on failure) the transaction. +// +// Usage pattern with named returns: +// +// func (c *Client) SomeMethod(ctx context.Context) (result T, retErr error) { +// q, done, err := c.acquireDBTX(ctx) +// if err != nil { return zero, err } +// defer func() { done(retErr) }() +// return SomeFunction(ctx, q, ...) +// } +func (c *Client) acquireDBTX(ctx context.Context) (DBTX, doneFunc, error) { + if c.pool == nil { + return nil, nil, fmt.Errorf("database connection not initialized") + } + + vars := c.effectiveSessionVars(ctx) + if vars.Empty() { + return c.pool, func(error) {}, nil + } + + tx, err := c.pool.Begin(ctx) + if err != nil { + return nil, nil, fmt.Errorf("failed to begin session var transaction: %w", err) + } + + if err := applySessionVars(ctx, tx, vars); err != nil { + tx.Rollback(ctx) //nolint:errcheck + return nil, nil, err + } + + // Use context.Background() for commit/rollback: the request context may + // be cancelled after the query succeeds (e.g., FUSE interrupt, NFS timeout), + // and we must not silently discard a successful write because of it. + done := func(retErr error) { + if retErr != nil { + tx.Rollback(context.Background()) //nolint:errcheck + } else { + if cErr := tx.Commit(context.Background()); cErr != nil { + logging.Warn("session var tx commit failed", zap.Error(cErr)) + } + } + } + return tx, done, nil } // WithQueryTimeout returns a context with the configured query timeout applied. @@ -148,10 +224,20 @@ func NewClient(ctx context.Context, cfg *config.Config, connStr string) (*Client zap.String("user", poolConfig.ConnConfig.User), ) - return &Client{ + client := &Client{ pool: pool, cfg: cfg, - }, nil + } + + // Set baseline session variables from config (applied to every query). + // Keys are sorted once here; no per-query sorting needed. + if len(cfg.SessionVariables) > 0 { + client.baselineVars = NewSessionVars(cfg.SessionVariables) + logging.Debug("Session variables configured", + zap.Int("count", len(cfg.SessionVariables))) + } + + return client, nil } // Close closes the database connection pool @@ -183,10 +269,16 @@ func (c *Client) Query(ctx context.Context, sql string, args ...interface{}) (in } // Exec executes a SQL statement (INSERT, UPDATE, DELETE, DDL) -func (c *Client) Exec(ctx context.Context, sql string, args ...interface{}) error { +func (c *Client) Exec(ctx context.Context, sql string, args ...interface{}) (retErr error) { logging.Debug("Executing statement", zap.String("sql", sql)) - _, err := c.pool.Exec(ctx, sql, args...) + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return err + } + defer func() { done(retErr) }() + + _, err = q.Exec(ctx, sql, args...) if err != nil { logging.Error("Statement execution failed", zap.String("sql", sql), @@ -201,6 +293,10 @@ func (c *Client) Exec(ctx context.Context, sql string, args ...interface{}) erro // ExecInTransaction executes a SQL statement within a transaction that is // always rolled back. Used to validate DDL without persisting changes. // +// This method manages its own transaction and must NOT use acquireDBTX, +// because acquireDBTX commits on success while this method must always rollback. +// Session vars are injected directly via applySessionVars after Begin. +// // Returns nil if the statement would succeed, or an error describing the failure. func (c *Client) ExecInTransaction(ctx context.Context, sql string, args ...interface{}) error { logging.Debug("Testing statement in transaction", zap.String("sql", sql)) @@ -219,6 +315,13 @@ func (c *Client) ExecInTransaction(ctx context.Context, sql string, args ...inte } }() + // Apply session variables (SET LOCAL) if configured + if vars := c.effectiveSessionVars(ctx); !vars.Empty() { + if err := applySessionVars(ctx, tx, vars); err != nil { + return fmt.Errorf("failed to apply session variables: %w", err) + } + } + // Execute the statement _, err = tx.Exec(ctx, sql, args...) if err != nil { diff --git a/internal/tigerfs/db/constraints.go b/internal/tigerfs/db/constraints.go index 488cb21..72245fb 100644 --- a/internal/tigerfs/db/constraints.go +++ b/internal/tigerfs/db/constraints.go @@ -4,7 +4,6 @@ import ( "context" "fmt" - "github.com/jackc/pgx/v5/pgxpool" "github.com/timescale/tigerfs/internal/tigerfs/logging" "go.uber.org/zap" ) @@ -19,14 +18,14 @@ type Constraint struct { // ValidateConstraints validates column values against table constraints // Returns an error if any constraints would be violated -func ValidateConstraints(ctx context.Context, pool *pgxpool.Pool, schema, table string, values map[string]interface{}) error { +func ValidateConstraints(ctx context.Context, dbtx DBTX, schema, table string, values map[string]interface{}) error { logging.Debug("Validating constraints", zap.String("schema", schema), zap.String("table", table), zap.Int("column_count", len(values))) // Get all columns for the table to check NOT NULL constraints - columns, err := getColumnsForConstraintCheck(ctx, pool, schema, table) + columns, err := getColumnsForConstraintCheck(ctx, dbtx, schema, table) if err != nil { return fmt.Errorf("failed to get columns: %w", err) } @@ -48,7 +47,7 @@ func ValidateConstraints(ctx context.Context, pool *pgxpool.Pool, schema, table // Check UNIQUE constraints // For each column being updated, check if it has a unique constraint - uniqueConstraints, err := getUniqueConstraints(ctx, pool, schema, table) + uniqueConstraints, err := getUniqueConstraints(ctx, dbtx, schema, table) if err != nil { return fmt.Errorf("failed to get unique constraints: %w", err) } @@ -58,7 +57,7 @@ func ValidateConstraints(ctx context.Context, pool *pgxpool.Pool, schema, table for _, colName := range constraint.Columns { if value, ok := values[colName]; ok { // This column is being updated, check for duplicates - if err := checkUniqueConstraint(ctx, pool, schema, table, colName, value); err != nil { + if err := checkUniqueConstraint(ctx, dbtx, schema, table, colName, value); err != nil { logging.Debug("UNIQUE constraint violation", zap.String("schema", schema), zap.String("table", table), @@ -81,7 +80,7 @@ func ValidateConstraints(ctx context.Context, pool *pgxpool.Pool, schema, table } // getColumnsForConstraintCheck queries column metadata for constraint checking -func getColumnsForConstraintCheck(ctx context.Context, pool *pgxpool.Pool, schema, table string) ([]Column, error) { +func getColumnsForConstraintCheck(ctx context.Context, dbtx DBTX, schema, table string) ([]Column, error) { query := ` SELECT column_name, @@ -93,7 +92,7 @@ func getColumnsForConstraintCheck(ctx context.Context, pool *pgxpool.Pool, schem ORDER BY ordinal_position ` - rows, err := pool.Query(ctx, query, schema, table) + rows, err := dbtx.Query(ctx, query, schema, table) if err != nil { return nil, fmt.Errorf("failed to query columns: %w", err) } @@ -126,7 +125,7 @@ func getColumnsForConstraintCheck(ctx context.Context, pool *pgxpool.Pool, schem } // getUniqueConstraints retrieves UNIQUE constraints for a table -func getUniqueConstraints(ctx context.Context, pool *pgxpool.Pool, schema, table string) ([]Constraint, error) { +func getUniqueConstraints(ctx context.Context, dbtx DBTX, schema, table string) ([]Constraint, error) { query := ` SELECT tc.constraint_name, @@ -142,7 +141,7 @@ func getUniqueConstraints(ctx context.Context, pool *pgxpool.Pool, schema, table GROUP BY tc.constraint_name ` - rows, err := pool.Query(ctx, query, schema, table) + rows, err := dbtx.Query(ctx, query, schema, table) if err != nil { return nil, fmt.Errorf("failed to query unique constraints: %w", err) } @@ -169,7 +168,7 @@ func getUniqueConstraints(ctx context.Context, pool *pgxpool.Pool, schema, table } // checkUniqueConstraint checks if a value violates a unique constraint -func checkUniqueConstraint(ctx context.Context, pool *pgxpool.Pool, schema, table, column string, value interface{}) error { +func checkUniqueConstraint(ctx context.Context, dbtx DBTX, schema, table, column string, value interface{}) error { // Check if this value already exists in the table query := fmt.Sprintf(` SELECT EXISTS( @@ -178,7 +177,7 @@ func checkUniqueConstraint(ctx context.Context, pool *pgxpool.Pool, schema, tabl `, qt(schema, table), qi(column)) var exists bool - err := pool.QueryRow(ctx, query, value).Scan(&exists) + err := dbtx.QueryRow(ctx, query, value).Scan(&exists) if err != nil { return fmt.Errorf("failed to check unique constraint: %w", err) } diff --git a/internal/tigerfs/db/dbtx.go b/internal/tigerfs/db/dbtx.go new file mode 100644 index 0000000..2aee9ca --- /dev/null +++ b/internal/tigerfs/db/dbtx.go @@ -0,0 +1,50 @@ +package db + +import ( + "context" + "fmt" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" + "github.com/jackc/pgx/v5/pgxpool" +) + +// DBTX is the common query interface satisfied by *pgxpool.Pool and pgx.Tx. +// Package-level functions accept DBTX so they can operate against either +// a raw pool connection or a transaction with SET LOCAL session variables. +// +// This interface intentionally excludes SendBatch, CopyFrom, and Begin. +// Operations that need those capabilities (bulk import, DDL validation) +// manage their own pgx.Tx directly and inject session vars via +// applySessionVars. +type DBTX interface { + Exec(ctx context.Context, sql string, args ...any) (pgconn.CommandTag, error) + Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error) + QueryRow(ctx context.Context, sql string, args ...any) pgx.Row +} + +// Compile-time verification that both pool and transaction satisfy DBTX. +var ( + _ DBTX = (*pgxpool.Pool)(nil) + _ DBTX = (pgx.Tx)(nil) +) + +// applySessionVars executes SET LOCAL for each session variable within an +// open transaction. Uses set_config($1, $2, true) which is the parameterized +// equivalent of SET LOCAL — safe against injection, transaction-scoped, +// and compatible with PgBouncer transaction mode and RDS Proxy. +// +// Keys are iterated in sorted order (pre-sorted at SessionVars construction) +// for deterministic execution with zero per-query allocation. +func applySessionVars(ctx context.Context, tx pgx.Tx, vars SessionVars) error { + var applyErr error + vars.Range(func(key, value string) { + if applyErr != nil { + return + } + if _, err := tx.Exec(ctx, "SELECT set_config($1, $2, true)", key, value); err != nil { + applyErr = fmt.Errorf("set session var %q: %w", key, err) + } + }) + return applyErr +} diff --git a/internal/tigerfs/db/dbtx_test.go b/internal/tigerfs/db/dbtx_test.go new file mode 100644 index 0000000..40d7ef6 --- /dev/null +++ b/internal/tigerfs/db/dbtx_test.go @@ -0,0 +1,20 @@ +package db + +import ( + "testing" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" +) + +// TestDBTX_PoolSatisfiesInterface verifies at compile time that +// *pgxpool.Pool satisfies the DBTX interface. +func TestDBTX_PoolSatisfiesInterface(t *testing.T) { + var _ DBTX = (*pgxpool.Pool)(nil) +} + +// TestDBTX_TxSatisfiesInterface verifies at compile time that +// pgx.Tx satisfies the DBTX interface. +func TestDBTX_TxSatisfiesInterface(t *testing.T) { + var _ DBTX = (pgx.Tx)(nil) +} diff --git a/internal/tigerfs/db/export.go b/internal/tigerfs/db/export.go index 8ab0893..1eaee4d 100644 --- a/internal/tigerfs/db/export.go +++ b/internal/tigerfs/db/export.go @@ -5,7 +5,6 @@ import ( "fmt" "strings" - "github.com/jackc/pgx/v5/pgxpool" "github.com/timescale/tigerfs/internal/tigerfs/logging" "go.uber.org/zap" ) @@ -23,7 +22,7 @@ import ( // - columns: Column names in database order // - rows: Row values as [][]interface{} // - error: Any database error -func GetAllRows(ctx context.Context, pool *pgxpool.Pool, schema, table string, limit int) ([]string, [][]interface{}, error) { +func GetAllRows(ctx context.Context, dbtx DBTX, schema, table string, limit int) ([]string, [][]interface{}, error) { logging.Debug("Getting all rows for export", zap.String("schema", schema), zap.String("table", table), @@ -34,7 +33,7 @@ func GetAllRows(ctx context.Context, pool *pgxpool.Pool, schema, table string, l qt(schema, table), ) - rows, err := pool.Query(ctx, query, limit) + rows, err := dbtx.Query(ctx, query, limit) if err != nil { return nil, nil, fmt.Errorf("failed to query rows: %w", err) } @@ -70,17 +69,19 @@ func GetAllRows(ctx context.Context, pool *pgxpool.Pool, schema, table string, l } // GetAllRows is a convenience wrapper for Client -func (c *Client) GetAllRows(ctx context.Context, schema, table string, limit int) ([]string, [][]interface{}, error) { - if c.pool == nil { - return nil, nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetAllRows(ctx context.Context, schema, table string, limit int) (columns []string, rows [][]interface{}, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, nil, err } - return GetAllRows(ctx, c.pool, schema, table, limit) + defer func() { done(retErr) }() + return GetAllRows(ctx, q, schema, table, limit) } // GetFirstNRowsWithData returns the first N rows ordered by primary key ascending. // Returns full row data, not just primary keys. // Used for bulk export with .first/N/ pagination. -func GetFirstNRowsWithData(ctx context.Context, pool *pgxpool.Pool, schema, table string, pkColumns []string, limit int) ([]string, [][]interface{}, error) { +func GetFirstNRowsWithData(ctx context.Context, dbtx DBTX, schema, table string, pkColumns []string, limit int) ([]string, [][]interface{}, error) { logging.Debug("Getting first N rows with data", zap.String("schema", schema), zap.String("table", table), @@ -92,7 +93,7 @@ func GetFirstNRowsWithData(ctx context.Context, pool *pgxpool.Pool, schema, tabl qt(schema, table), pkOrderByList(pkColumns, "ASC"), ) - rows, err := pool.Query(ctx, query, limit) + rows, err := dbtx.Query(ctx, query, limit) if err != nil { return nil, nil, fmt.Errorf("failed to query first N rows: %w", err) } @@ -128,17 +129,19 @@ func GetFirstNRowsWithData(ctx context.Context, pool *pgxpool.Pool, schema, tabl } // GetFirstNRowsWithData is a convenience wrapper for Client -func (c *Client) GetFirstNRowsWithData(ctx context.Context, schema, table string, pkColumns []string, limit int) ([]string, [][]interface{}, error) { - if c.pool == nil { - return nil, nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetFirstNRowsWithData(ctx context.Context, schema, table string, pkColumns []string, limit int) (columns []string, rows [][]interface{}, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, nil, err } - return GetFirstNRowsWithData(ctx, c.pool, schema, table, pkColumns, limit) + defer func() { done(retErr) }() + return GetFirstNRowsWithData(ctx, q, schema, table, pkColumns, limit) } // GetLastNRowsWithData returns the last N rows ordered by primary key descending. // Returns full row data, not just primary keys. // Used for bulk export with .last/N/ pagination. -func GetLastNRowsWithData(ctx context.Context, pool *pgxpool.Pool, schema, table string, pkColumns []string, limit int) ([]string, [][]interface{}, error) { +func GetLastNRowsWithData(ctx context.Context, dbtx DBTX, schema, table string, pkColumns []string, limit int) ([]string, [][]interface{}, error) { logging.Debug("Getting last N rows with data", zap.String("schema", schema), zap.String("table", table), @@ -150,7 +153,7 @@ func GetLastNRowsWithData(ctx context.Context, pool *pgxpool.Pool, schema, table qt(schema, table), pkOrderByList(pkColumns, "DESC"), ) - rows, err := pool.Query(ctx, query, limit) + rows, err := dbtx.Query(ctx, query, limit) if err != nil { return nil, nil, fmt.Errorf("failed to query last N rows: %w", err) } @@ -186,19 +189,23 @@ func GetLastNRowsWithData(ctx context.Context, pool *pgxpool.Pool, schema, table } // GetLastNRowsWithData is a convenience wrapper for Client -func (c *Client) GetLastNRowsWithData(ctx context.Context, schema, table string, pkColumns []string, limit int) ([]string, [][]interface{}, error) { - if c.pool == nil { - return nil, nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetLastNRowsWithData(ctx context.Context, schema, table string, pkColumns []string, limit int) (columns []string, rows [][]interface{}, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, nil, err } - return GetLastNRowsWithData(ctx, c.pool, schema, table, pkColumns, limit) + defer func() { done(retErr) }() + return GetLastNRowsWithData(ctx, q, schema, table, pkColumns, limit) } // RowExistsByColumns checks if any row matches the given column=value conditions. // Generates: SELECT 1 FROM "schema"."table" WHERE "col1" = $1 AND "col2" = $2 LIMIT 1 -func (c *Client) RowExistsByColumns(ctx context.Context, schema, table string, columns []string, values []interface{}) (bool, error) { - if c.pool == nil { - return false, fmt.Errorf("database connection not initialized") +func (c *Client) RowExistsByColumns(ctx context.Context, schema, table string, columns []string, values []interface{}) (result bool, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return false, err } + defer func() { done(retErr) }() // Build WHERE clause var whereParts []string @@ -217,7 +224,7 @@ func (c *Client) RowExistsByColumns(ctx context.Context, schema, table string, c zap.Strings("columns", columns)) var dummy int - err := c.pool.QueryRow(ctx, query, values...).Scan(&dummy) + err = q.QueryRow(ctx, query, values...).Scan(&dummy) if err != nil { if err.Error() == "no rows in result set" { return false, nil @@ -231,10 +238,12 @@ func (c *Client) RowExistsByColumns(ctx context.Context, schema, table string, c // GetRowByColumns returns a single row matching column=value conditions. // Returns all columns. Returns (nil, nil, nil) if no match. // Generates: SELECT * FROM "schema"."table" WHERE "col1" = $1 AND "col2" = $2 LIMIT 1 -func (c *Client) GetRowByColumns(ctx context.Context, schema, table string, columns []string, values []interface{}) ([]string, []interface{}, error) { - if c.pool == nil { - return nil, nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetRowByColumns(ctx context.Context, schema, table string, columns []string, values []interface{}) (colNames []string, rowValues []interface{}, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, nil, err } + defer func() { done(retErr) }() // Build WHERE clause var whereParts []string @@ -252,7 +261,7 @@ func (c *Client) GetRowByColumns(ctx context.Context, schema, table string, colu zap.String("table", table), zap.Strings("columns", columns)) - rows, err := c.pool.Query(ctx, query, values...) + rows, err := q.Query(ctx, query, values...) if err != nil { return nil, nil, fmt.Errorf("failed to query row by columns: %w", err) } @@ -260,9 +269,9 @@ func (c *Client) GetRowByColumns(ctx context.Context, schema, table string, colu // Get column names from field descriptions fieldDescriptions := rows.FieldDescriptions() - colNames := make([]string, len(fieldDescriptions)) + resultColNames := make([]string, len(fieldDescriptions)) for i, fd := range fieldDescriptions { - colNames[i] = string(fd.Name) + resultColNames[i] = string(fd.Name) } if !rows.Next() { @@ -272,10 +281,10 @@ func (c *Client) GetRowByColumns(ctx context.Context, schema, table string, colu return nil, nil, nil // No match } - rowValues, err := rows.Values() + resultRowValues, err := rows.Values() if err != nil { return nil, nil, fmt.Errorf("failed to scan row values: %w", err) } - return colNames, rowValues, nil + return resultColNames, resultRowValues, nil } diff --git a/internal/tigerfs/db/import.go b/internal/tigerfs/db/import.go index 309f68b..d21d670 100644 --- a/internal/tigerfs/db/import.go +++ b/internal/tigerfs/db/import.go @@ -31,6 +31,13 @@ func (c *Client) ImportOverwrite(ctx context.Context, schema, table string, colu } defer tx.Rollback(ctx) //nolint:errcheck + // Apply session variables (SET LOCAL) if configured + if vars := c.effectiveSessionVars(ctx); !vars.Empty() { + if err := applySessionVars(ctx, tx, vars); err != nil { + return fmt.Errorf("failed to apply session variables: %w", err) + } + } + // Truncate table truncateSQL := fmt.Sprintf("TRUNCATE %s", qt(schema, table)) @@ -81,6 +88,13 @@ func (c *Client) ImportSync(ctx context.Context, schema, table string, columns [ } defer tx.Rollback(ctx) //nolint:errcheck + // Apply session variables (SET LOCAL) if configured + if vars := c.effectiveSessionVars(ctx); !vars.Empty() { + if err := applySessionVars(ctx, tx, vars); err != nil { + return fmt.Errorf("failed to apply session variables: %w", err) + } + } + // Build upsert statement // INSERT INTO table (col1, col2) VALUES ($1, $2) // ON CONFLICT (pk_col) DO UPDATE SET col1 = EXCLUDED.col1, col2 = EXCLUDED.col2 @@ -119,6 +133,13 @@ func (c *Client) ImportAppend(ctx context.Context, schema, table string, columns } defer tx.Rollback(ctx) //nolint:errcheck + // Apply session variables (SET LOCAL) if configured + if vars := c.effectiveSessionVars(ctx); !vars.Empty() { + if err := applySessionVars(ctx, tx, vars); err != nil { + return fmt.Errorf("failed to apply session variables: %w", err) + } + } + // Bulk insert (will fail on conflicts) if err := c.bulkInsert(ctx, tx, schema, table, columns, rows); err != nil { return err @@ -136,10 +157,16 @@ func (c *Client) ImportAppend(ctx context.Context, schema, table string, columns } // truncateTable truncates a table in its own transaction. -func (c *Client) truncateTable(ctx context.Context, schema, table string) error { +func (c *Client) truncateTable(ctx context.Context, schema, table string) (retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return err + } + defer func() { done(retErr) }() + truncateSQL := fmt.Sprintf("TRUNCATE %s", qt(schema, table)) - _, err := c.pool.Exec(ctx, truncateSQL) + _, err = q.Exec(ctx, truncateSQL) if err != nil { return fmt.Errorf("failed to truncate table: %w", err) } @@ -149,6 +176,10 @@ func (c *Client) truncateTable(ctx context.Context, schema, table string) error // bulkInsert performs a bulk INSERT using COPY with text format for efficiency. // Text format allows PostgreSQL to handle type conversions server-side, // avoiding issues with binary encoding of complex types like timestamps. +// +// This requires pgx.Tx (not DBTX) because it uses tx.Conn().PgConn().CopyFrom() +// which is not part of the DBTX interface. Session vars are applied by the +// caller (ImportOverwrite/Sync/Append) at the start of the transaction. func (c *Client) bulkInsert(ctx context.Context, tx pgx.Tx, schema, table string, columns []string, rows [][]interface{}) error { // Build COPY command with text format quotedColumns := make([]string, len(columns)) diff --git a/internal/tigerfs/db/indexes.go b/internal/tigerfs/db/indexes.go index f7b3bdf..757f6a6 100644 --- a/internal/tigerfs/db/indexes.go +++ b/internal/tigerfs/db/indexes.go @@ -4,7 +4,6 @@ import ( "context" "fmt" - "github.com/jackc/pgx/v5/pgxpool" "github.com/timescale/tigerfs/internal/tigerfs/format" "github.com/timescale/tigerfs/internal/tigerfs/logging" "go.uber.org/zap" @@ -45,7 +44,7 @@ func (i *Index) IsComposite() bool { // // Parameters: // - ctx: Context for cancellation and timeout -// - pool: Database connection pool +// - dbtx: Database connection or transaction // - schema: Schema name (e.g., "public") // - table: Table name // @@ -53,7 +52,7 @@ func (i *Index) IsComposite() bool { // Returns empty slice (not error) if table has no indexes. // // Uses pg_catalog instead of parsing indexdef strings for reliability. -func GetIndexes(ctx context.Context, pool *pgxpool.Pool, schema, table string) ([]Index, error) { +func GetIndexes(ctx context.Context, dbtx DBTX, schema, table string) ([]Index, error) { logging.Debug("Querying indexes", zap.String("schema", schema), zap.String("table", table)) @@ -83,7 +82,7 @@ func GetIndexes(ctx context.Context, pool *pgxpool.Pool, schema, table string) ( i.relname ` - rows, err := pool.Query(ctx, query, schema, table) + rows, err := dbtx.Query(ctx, query, schema, table) if err != nil { return nil, fmt.Errorf("failed to query indexes: %w", err) } @@ -114,11 +113,13 @@ func GetIndexes(ctx context.Context, pool *pgxpool.Pool, schema, table string) ( } // GetIndexes is a convenience wrapper for Client. -func (c *Client) GetIndexes(ctx context.Context, schema, table string) ([]Index, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetIndexes(ctx context.Context, schema, table string) (result []Index, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } - return GetIndexes(ctx, c.pool, schema, table) + defer func() { done(retErr) }() + return GetIndexes(ctx, q, schema, table) } // GetIndexByColumn finds the first index whose leading column matches. @@ -126,15 +127,15 @@ func (c *Client) GetIndexes(ctx context.Context, schema, table string) ([]Index, // // Parameters: // - ctx: Context for cancellation and timeout -// - pool: Database connection pool +// - dbtx: Database connection or transaction // - schema: Schema name // - table: Table name // - column: Column name to match against first index column // // Only matches indexes where the specified column is the leading (first) column, // since PostgreSQL can only use an index efficiently when querying by leading columns. -func GetIndexByColumn(ctx context.Context, pool *pgxpool.Pool, schema, table, column string) (*Index, error) { - indexes, err := GetIndexes(ctx, pool, schema, table) +func GetIndexByColumn(ctx context.Context, dbtx DBTX, schema, table, column string) (*Index, error) { + indexes, err := GetIndexes(ctx, dbtx, schema, table) if err != nil { return nil, err } @@ -149,11 +150,13 @@ func GetIndexByColumn(ctx context.Context, pool *pgxpool.Pool, schema, table, co } // GetIndexByColumn is a convenience wrapper for Client. -func (c *Client) GetIndexByColumn(ctx context.Context, schema, table, column string) (*Index, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetIndexByColumn(ctx context.Context, schema, table, column string) (result *Index, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } - return GetIndexByColumn(ctx, c.pool, schema, table, column) + defer func() { done(retErr) }() + return GetIndexByColumn(ctx, q, schema, table, column) } // GetSingleColumnIndexes returns indexes with exactly one column. @@ -161,13 +164,13 @@ func (c *Client) GetIndexByColumn(ctx context.Context, schema, table, column str // // Parameters: // - ctx: Context for cancellation and timeout -// - pool: Database connection pool +// - dbtx: Database connection or transaction // - schema: Schema name // - table: Table name // // Excludes composite indexes. Returns empty slice if none found. -func GetSingleColumnIndexes(ctx context.Context, pool *pgxpool.Pool, schema, table string) ([]Index, error) { - indexes, err := GetIndexes(ctx, pool, schema, table) +func GetSingleColumnIndexes(ctx context.Context, dbtx DBTX, schema, table string) ([]Index, error) { + indexes, err := GetIndexes(ctx, dbtx, schema, table) if err != nil { return nil, err } @@ -183,11 +186,13 @@ func GetSingleColumnIndexes(ctx context.Context, pool *pgxpool.Pool, schema, tab } // GetSingleColumnIndexes is a convenience wrapper for Client. -func (c *Client) GetSingleColumnIndexes(ctx context.Context, schema, table string) ([]Index, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetSingleColumnIndexes(ctx context.Context, schema, table string) (result []Index, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } - return GetSingleColumnIndexes(ctx, c.pool, schema, table) + defer func() { done(retErr) }() + return GetSingleColumnIndexes(ctx, q, schema, table) } // GetCompositeIndexes returns indexes with multiple columns. @@ -195,13 +200,13 @@ func (c *Client) GetSingleColumnIndexes(ctx context.Context, schema, table strin // // Parameters: // - ctx: Context for cancellation and timeout -// - pool: Database connection pool +// - dbtx: Database connection or transaction // - schema: Schema name // - table: Table name // // Excludes single-column indexes. Returns empty slice if none found. -func GetCompositeIndexes(ctx context.Context, pool *pgxpool.Pool, schema, table string) ([]Index, error) { - indexes, err := GetIndexes(ctx, pool, schema, table) +func GetCompositeIndexes(ctx context.Context, dbtx DBTX, schema, table string) ([]Index, error) { + indexes, err := GetIndexes(ctx, dbtx, schema, table) if err != nil { return nil, err } @@ -217,11 +222,13 @@ func GetCompositeIndexes(ctx context.Context, pool *pgxpool.Pool, schema, table } // GetCompositeIndexes is a convenience wrapper for Client. -func (c *Client) GetCompositeIndexes(ctx context.Context, schema, table string) ([]Index, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetCompositeIndexes(ctx context.Context, schema, table string) (result []Index, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } - return GetCompositeIndexes(ctx, c.pool, schema, table) + defer func() { done(retErr) }() + return GetCompositeIndexes(ctx, q, schema, table) } // GetDistinctValues retrieves distinct values for an indexed column. @@ -229,7 +236,7 @@ func (c *Client) GetCompositeIndexes(ctx context.Context, schema, table string) // // Parameters: // - ctx: Context for cancellation and timeout -// - pool: Database connection pool +// - dbtx: Database connection or transaction // - schema: Schema name // - table: Table name // - column: Column name to get distinct values for @@ -237,7 +244,7 @@ func (c *Client) GetCompositeIndexes(ctx context.Context, schema, table string) // // Returns values as strings. NULL values are excluded. // Values are ordered for consistent directory listings. -func GetDistinctValues(ctx context.Context, pool *pgxpool.Pool, schema, table, column string, limit int) ([]string, error) { +func GetDistinctValues(ctx context.Context, dbtx DBTX, schema, table, column string, limit int) ([]string, error) { logging.Debug("Querying distinct values", zap.String("schema", schema), zap.String("table", table), @@ -250,7 +257,7 @@ func GetDistinctValues(ctx context.Context, pool *pgxpool.Pool, schema, table, c qi(column), qt(schema, table), qi(column), qi(column), ) - rows, err := pool.Query(ctx, query, limit) + rows, err := dbtx.Query(ctx, query, limit) if err != nil { return nil, fmt.Errorf("failed to query distinct values: %w", err) } @@ -286,11 +293,13 @@ func GetDistinctValues(ctx context.Context, pool *pgxpool.Pool, schema, table, c } // GetDistinctValues is a convenience wrapper for Client. -func (c *Client) GetDistinctValues(ctx context.Context, schema, table, column string, limit int) ([]string, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetDistinctValues(ctx context.Context, schema, table, column string, limit int) (result []string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } - return GetDistinctValues(ctx, c.pool, schema, table, column, limit) + defer func() { done(retErr) }() + return GetDistinctValues(ctx, q, schema, table, column, limit) } // GetDistinctValuesOrdered retrieves distinct values with explicit ordering. @@ -298,7 +307,7 @@ func (c *Client) GetDistinctValues(ctx context.Context, schema, table, column st // // Parameters: // - ascending: true for ASC (first N), false for DESC (last N) -func GetDistinctValuesOrdered(ctx context.Context, pool *pgxpool.Pool, schema, table, column string, limit int, ascending bool) ([]string, error) { +func GetDistinctValuesOrdered(ctx context.Context, dbtx DBTX, schema, table, column string, limit int, ascending bool) ([]string, error) { logging.Debug("Querying ordered distinct values", zap.String("schema", schema), zap.String("table", table), @@ -316,7 +325,7 @@ func GetDistinctValuesOrdered(ctx context.Context, pool *pgxpool.Pool, schema, t qi(column), qt(schema, table), qi(column), qi(column), order, ) - rows, err := pool.Query(ctx, query, limit) + rows, err := dbtx.Query(ctx, query, limit) if err != nil { return nil, fmt.Errorf("failed to query ordered distinct values: %w", err) } @@ -343,11 +352,13 @@ func GetDistinctValuesOrdered(ctx context.Context, pool *pgxpool.Pool, schema, t } // GetDistinctValuesOrdered is a convenience wrapper for Client. -func (c *Client) GetDistinctValuesOrdered(ctx context.Context, schema, table, column string, limit int, ascending bool) ([]string, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetDistinctValuesOrdered(ctx context.Context, schema, table, column string, limit int, ascending bool) (result []string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } - return GetDistinctValuesOrdered(ctx, c.pool, schema, table, column, limit, ascending) + defer func() { done(retErr) }() + return GetDistinctValuesOrdered(ctx, q, schema, table, column, limit, ascending) } // GetRowsByIndexValue retrieves primary keys of rows matching an indexed column value. @@ -355,7 +366,7 @@ func (c *Client) GetDistinctValuesOrdered(ctx context.Context, schema, table, co // // Parameters: // - ctx: Context for cancellation and timeout -// - pool: Database connection pool +// - dbtx: Database connection or transaction // - schema: Schema name // - table: Table name // - column: Indexed column to query @@ -364,7 +375,7 @@ func (c *Client) GetDistinctValuesOrdered(ctx context.Context, schema, table, co // - limit: Maximum number of rows to return // // Returns primary key values as strings. If no rows match, returns empty slice (not error). -func GetRowsByIndexValue(ctx context.Context, pool *pgxpool.Pool, schema, table, column, value string, pkColumns []string, limit int) ([]string, error) { +func GetRowsByIndexValue(ctx context.Context, dbtx DBTX, schema, table, column, value string, pkColumns []string, limit int) ([]string, error) { logging.Debug("Querying rows by index value", zap.String("schema", schema), zap.String("table", table), @@ -378,20 +389,22 @@ func GetRowsByIndexValue(ctx context.Context, pool *pgxpool.Pool, schema, table, pkSelectList(pkColumns), qt(schema, table), qi(column), pkOrderByList(pkColumns, "ASC"), ) - return scanPKRowsWithArgs(ctx, pool, query, pkColumns, value, limit) + return scanPKRowsWithArgs(ctx, dbtx, query, pkColumns, value, limit) } // GetRowsByIndexValue is a convenience wrapper for Client. -func (c *Client) GetRowsByIndexValue(ctx context.Context, schema, table, column, value string, pkColumns []string, limit int) ([]string, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetRowsByIndexValue(ctx context.Context, schema, table, column, value string, pkColumns []string, limit int) (result []string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } - return GetRowsByIndexValue(ctx, c.pool, schema, table, column, value, pkColumns, limit) + defer func() { done(retErr) }() + return GetRowsByIndexValue(ctx, q, schema, table, column, value, pkColumns, limit) } // GetRowsByIndexValueOrdered retrieves primary keys with explicit ordering. // Used for .first/N/ and .last/N/ within index value navigation. -func GetRowsByIndexValueOrdered(ctx context.Context, pool *pgxpool.Pool, schema, table, column, value string, pkColumns []string, limit int, ascending bool) ([]string, error) { +func GetRowsByIndexValueOrdered(ctx context.Context, dbtx DBTX, schema, table, column, value string, pkColumns []string, limit int, ascending bool) ([]string, error) { logging.Debug("Querying ordered rows by index value", zap.String("schema", schema), zap.String("table", table), @@ -410,21 +423,23 @@ func GetRowsByIndexValueOrdered(ctx context.Context, pool *pgxpool.Pool, schema, pkSelectList(pkColumns), qt(schema, table), qi(column), pkOrderByList(pkColumns, order), ) - return scanPKRowsWithArgs(ctx, pool, query, pkColumns, value, limit) + return scanPKRowsWithArgs(ctx, dbtx, query, pkColumns, value, limit) } // GetRowsByIndexValueOrdered is a convenience wrapper for Client. -func (c *Client) GetRowsByIndexValueOrdered(ctx context.Context, schema, table, column, value string, pkColumns []string, limit int, ascending bool) ([]string, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetRowsByIndexValueOrdered(ctx context.Context, schema, table, column, value string, pkColumns []string, limit int, ascending bool) (result []string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } - return GetRowsByIndexValueOrdered(ctx, c.pool, schema, table, column, value, pkColumns, limit, ascending) + defer func() { done(retErr) }() + return GetRowsByIndexValueOrdered(ctx, q, schema, table, column, value, pkColumns, limit, ascending) } // scanPKRowsWithArgs executes a query with value and limit args, scans PK columns, and // returns encoded PK strings. Used by index query functions. -func scanPKRowsWithArgs(ctx context.Context, pool *pgxpool.Pool, query string, pkColumns []string, args ...interface{}) ([]string, error) { - rows, err := pool.Query(ctx, query, args...) +func scanPKRowsWithArgs(ctx context.Context, dbtx DBTX, query string, pkColumns []string, args ...interface{}) ([]string, error) { + rows, err := dbtx.Query(ctx, query, args...) if err != nil { return nil, fmt.Errorf("failed to query rows: %w", err) } @@ -465,7 +480,7 @@ func scanPKRowsWithArgs(ctx context.Context, pool *pgxpool.Pool, query string, p // // Parameters: // - ctx: Context for cancellation and timeout -// - pool: Database connection pool +// - dbtx: Database connection or transaction // - schema: Schema name // - table: Table name // - targetColumn: Column to get distinct values for (the "next" column in navigation) @@ -475,7 +490,7 @@ func scanPKRowsWithArgs(ctx context.Context, pool *pgxpool.Pool, query string, p // // Returns values as strings. NULL values are excluded. // filterColumns and filterValues must have the same length. -func GetDistinctValuesFiltered(ctx context.Context, pool *pgxpool.Pool, schema, table, targetColumn string, filterColumns, filterValues []string, limit int) ([]string, error) { +func GetDistinctValuesFiltered(ctx context.Context, dbtx DBTX, schema, table, targetColumn string, filterColumns, filterValues []string, limit int) ([]string, error) { if len(filterColumns) != len(filterValues) { return nil, fmt.Errorf("filterColumns and filterValues length mismatch: %d vs %d", len(filterColumns), len(filterValues)) } @@ -505,7 +520,7 @@ func GetDistinctValuesFiltered(ctx context.Context, pool *pgxpool.Pool, schema, qi(targetColumn), len(args), ) - rows, err := pool.Query(ctx, query, args...) + rows, err := dbtx.Query(ctx, query, args...) if err != nil { return nil, fmt.Errorf("failed to query filtered distinct values: %w", err) } @@ -539,11 +554,13 @@ func GetDistinctValuesFiltered(ctx context.Context, pool *pgxpool.Pool, schema, } // GetDistinctValuesFiltered is a convenience wrapper for Client. -func (c *Client) GetDistinctValuesFiltered(ctx context.Context, schema, table, targetColumn string, filterColumns, filterValues []string, limit int) ([]string, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetDistinctValuesFiltered(ctx context.Context, schema, table, targetColumn string, filterColumns, filterValues []string, limit int) (result []string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } - return GetDistinctValuesFiltered(ctx, c.pool, schema, table, targetColumn, filterColumns, filterValues, limit) + defer func() { done(retErr) }() + return GetDistinctValuesFiltered(ctx, q, schema, table, targetColumn, filterColumns, filterValues, limit) } // GetRowsByCompositeIndex retrieves primary keys of rows matching multiple column conditions. @@ -558,7 +575,7 @@ func (c *Client) GetDistinctValuesFiltered(ctx context.Context, schema, table, t // // Parameters: // - ctx: Context for cancellation and timeout -// - pool: Database connection pool +// - dbtx: Database connection or transaction // - schema: Schema name // - table: Table name // - columns: Indexed columns to query (in index order) @@ -567,7 +584,7 @@ func (c *Client) GetDistinctValuesFiltered(ctx context.Context, schema, table, t // - limit: Maximum number of rows to return // // Returns primary key values as strings. If no rows match, returns empty slice (not error). -func GetRowsByCompositeIndex(ctx context.Context, pool *pgxpool.Pool, schema, table string, columns, values []string, pkColumns []string, limit int) ([]string, error) { +func GetRowsByCompositeIndex(ctx context.Context, dbtx DBTX, schema, table string, columns, values []string, pkColumns []string, limit int) ([]string, error) { if len(columns) != len(values) { return nil, fmt.Errorf("columns and values length mismatch: %d vs %d", len(columns), len(values)) } @@ -596,15 +613,17 @@ func GetRowsByCompositeIndex(ctx context.Context, pool *pgxpool.Pool, schema, ta pkOrderByList(pkColumns, "ASC"), len(args), ) - return scanPKRowsWithArgs(ctx, pool, query, pkColumns, args...) + return scanPKRowsWithArgs(ctx, dbtx, query, pkColumns, args...) } // GetRowsByCompositeIndex is a convenience wrapper for Client. -func (c *Client) GetRowsByCompositeIndex(ctx context.Context, schema, table string, columns, values []string, pkColumns []string, limit int) ([]string, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetRowsByCompositeIndex(ctx context.Context, schema, table string, columns, values []string, pkColumns []string, limit int) (result []string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } - return GetRowsByCompositeIndex(ctx, c.pool, schema, table, columns, values, pkColumns, limit) + defer func() { done(retErr) }() + return GetRowsByCompositeIndex(ctx, q, schema, table, columns, values, pkColumns, limit) } // joinStrings joins strings with a separator (helper to avoid importing strings package). diff --git a/internal/tigerfs/db/keys.go b/internal/tigerfs/db/keys.go index 458e794..3497d12 100644 --- a/internal/tigerfs/db/keys.go +++ b/internal/tigerfs/db/keys.go @@ -4,7 +4,6 @@ import ( "context" "fmt" - "github.com/jackc/pgx/v5/pgxpool" "github.com/timescale/tigerfs/internal/tigerfs/logging" "go.uber.org/zap" ) @@ -16,7 +15,7 @@ type PrimaryKey struct { // GetPrimaryKey discovers the primary key for a table // Returns error if no primary key exists or if composite key (not supported in MVP) -func GetPrimaryKey(ctx context.Context, pool *pgxpool.Pool, schema, table string) (*PrimaryKey, error) { +func GetPrimaryKey(ctx context.Context, dbtx DBTX, schema, table string) (*PrimaryKey, error) { logging.Debug("Querying primary key", zap.String("schema", schema), zap.String("table", table)) @@ -33,7 +32,7 @@ func GetPrimaryKey(ctx context.Context, pool *pgxpool.Pool, schema, table string ORDER BY kcu.ordinal_position ` - rows, err := pool.Query(ctx, query, schema, table) + rows, err := dbtx.Query(ctx, query, schema, table) if err != nil { return nil, fmt.Errorf("failed to query primary key: %w", err) } @@ -69,7 +68,7 @@ func GetPrimaryKey(ctx context.Context, pool *pgxpool.Pool, schema, table string // ListRows returns a list of primary key values from a table. // Returns encoded PK strings (comma-delimited for composite PKs). // Limited to max_rows to prevent excessive directory listings. -func ListRows(ctx context.Context, pool *pgxpool.Pool, schema, table string, pkColumns []string, limit int) ([]string, error) { +func ListRows(ctx context.Context, dbtx DBTX, schema, table string, pkColumns []string, limit int) ([]string, error) { logging.Debug("Listing rows", zap.String("schema", schema), zap.String("table", table), @@ -81,29 +80,33 @@ func ListRows(ctx context.Context, pool *pgxpool.Pool, schema, table string, pkC pkSelectList(pkColumns), qt(schema, table), pkOrderByList(pkColumns, "ASC"), ) - return scanPKRows(ctx, pool, query, pkColumns, limit) + return scanPKRows(ctx, dbtx, query, pkColumns, limit) } // GetPrimaryKey is a convenience wrapper for Client -func (c *Client) GetPrimaryKey(ctx context.Context, schema, table string) (*PrimaryKey, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetPrimaryKey(ctx context.Context, schema, table string) (result *PrimaryKey, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } - return GetPrimaryKey(ctx, c.pool, schema, table) + defer func() { done(retErr) }() + return GetPrimaryKey(ctx, q, schema, table) } // ListRows is a convenience wrapper for Client -func (c *Client) ListRows(ctx context.Context, schema, table string, pkColumns []string, limit int) ([]string, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) ListRows(ctx context.Context, schema, table string, pkColumns []string, limit int) (result []string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } - return ListRows(ctx, c.pool, schema, table, pkColumns, limit) + defer func() { done(retErr) }() + return ListRows(ctx, q, schema, table, pkColumns, limit) } // ListAllRows returns all primary key values from a table without any limit. // Returns encoded PK strings (comma-delimited for composite PKs). // Used by .all/ paths to explicitly bypass dir_listing_limit restriction. -func ListAllRows(ctx context.Context, pool *pgxpool.Pool, schema, table string, pkColumns []string) ([]string, error) { +func ListAllRows(ctx context.Context, dbtx DBTX, schema, table string, pkColumns []string) ([]string, error) { logging.Debug("Listing all rows (no limit)", zap.String("schema", schema), zap.String("table", table), @@ -114,7 +117,7 @@ func ListAllRows(ctx context.Context, pool *pgxpool.Pool, schema, table string, pkSelectList(pkColumns), qt(schema, table), pkOrderByList(pkColumns, "ASC"), ) - rows, err := pool.Query(ctx, query) + rows, err := dbtx.Query(ctx, query) if err != nil { return nil, fmt.Errorf("failed to list all rows: %w", err) } @@ -150,9 +153,11 @@ func ListAllRows(ctx context.Context, pool *pgxpool.Pool, schema, table string, } // ListAllRows is a convenience wrapper for Client -func (c *Client) ListAllRows(ctx context.Context, schema, table string, pkColumns []string) ([]string, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) ListAllRows(ctx context.Context, schema, table string, pkColumns []string) (result []string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } - return ListAllRows(ctx, c.pool, schema, table, pkColumns) + defer func() { done(retErr) }() + return ListAllRows(ctx, q, schema, table, pkColumns) } diff --git a/internal/tigerfs/db/permissions.go b/internal/tigerfs/db/permissions.go index 8e6a854..dcd1900 100644 --- a/internal/tigerfs/db/permissions.go +++ b/internal/tigerfs/db/permissions.go @@ -4,7 +4,6 @@ import ( "context" "fmt" - "github.com/jackc/pgx/v5/pgxpool" "github.com/timescale/tigerfs/internal/tigerfs/logging" "go.uber.org/zap" ) @@ -53,12 +52,12 @@ func (p *TablePermissions) HasAnyPermission() bool { // // Parameters: // - ctx: Context for cancellation -// - pool: PostgreSQL connection pool +// - dbtx: Database connection or transaction // - schema: PostgreSQL schema name // - table: Table name to check permissions for // // Returns the permissions struct, or error on database failure. -func GetTablePermissions(ctx context.Context, pool *pgxpool.Pool, schema, table string) (*TablePermissions, error) { +func GetTablePermissions(ctx context.Context, dbtx DBTX, schema, table string) (*TablePermissions, error) { logging.Debug("Getting table permissions", zap.String("schema", schema), zap.String("table", table)) @@ -77,7 +76,7 @@ func GetTablePermissions(ctx context.Context, pool *pgxpool.Pool, schema, table ` var perms TablePermissions - err := pool.QueryRow(ctx, query, qualifiedName).Scan( + err := dbtx.QueryRow(ctx, query, qualifiedName).Scan( &perms.CanSelect, &perms.CanInsert, &perms.CanUpdate, @@ -107,11 +106,13 @@ func GetTablePermissions(ctx context.Context, pool *pgxpool.Pool, schema, table // // Returns the permissions struct, or error if the database connection // is not initialized or the query fails. -func (c *Client) GetTablePermissions(ctx context.Context, schema, table string) (*TablePermissions, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetTablePermissions(ctx context.Context, schema, table string) (result *TablePermissions, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } - return GetTablePermissions(ctx, c.pool, schema, table) + defer func() { done(retErr) }() + return GetTablePermissions(ctx, q, schema, table) } // GetTablePermissionsBatch queries PostgreSQL to determine privileges @@ -120,12 +121,12 @@ func (c *Client) GetTablePermissions(ctx context.Context, schema, table string) // // Parameters: // - ctx: Context for cancellation -// - pool: PostgreSQL connection pool +// - dbtx: Database connection or transaction // - schema: PostgreSQL schema name // - tables: Table names to check permissions for // // Returns a map of table name to permissions, or error on database failure. -func GetTablePermissionsBatch(ctx context.Context, pool *pgxpool.Pool, schema string, tables []string) (map[string]*TablePermissions, error) { +func GetTablePermissionsBatch(ctx context.Context, dbtx DBTX, schema string, tables []string) (map[string]*TablePermissions, error) { if len(tables) == 0 { return make(map[string]*TablePermissions), nil } @@ -143,7 +144,7 @@ func GetTablePermissionsBatch(ctx context.Context, pool *pgxpool.Pool, schema st FROM unnest($2::text[]) AS t(table_name) ` - rows, err := pool.Query(ctx, query, schema, tables) + rows, err := dbtx.Query(ctx, query, schema, tables) if err != nil { return nil, fmt.Errorf("failed to get batch table permissions: %w", err) } @@ -171,9 +172,11 @@ func GetTablePermissionsBatch(ctx context.Context, pool *pgxpool.Pool, schema st } // GetTablePermissionsBatch is a convenience wrapper for Client. -func (c *Client) GetTablePermissionsBatch(ctx context.Context, schema string, tables []string) (map[string]*TablePermissions, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetTablePermissionsBatch(ctx context.Context, schema string, tables []string) (result map[string]*TablePermissions, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } - return GetTablePermissionsBatch(ctx, c.pool, schema, tables) + defer func() { done(retErr) }() + return GetTablePermissionsBatch(ctx, q, schema, tables) } diff --git a/internal/tigerfs/db/pipeline.go b/internal/tigerfs/db/pipeline.go index 1001877..1856cef 100644 --- a/internal/tigerfs/db/pipeline.go +++ b/internal/tigerfs/db/pipeline.go @@ -10,7 +10,6 @@ import ( "fmt" "strings" - "github.com/jackc/pgx/v5/pgxpool" "github.com/timescale/tigerfs/internal/tigerfs/logging" "go.uber.org/zap" ) @@ -109,11 +108,11 @@ func (p *QueryParams) HasOrder() bool { // // Parameters: // - ctx: Context for cancellation and timeout -// - pool: PostgreSQL connection pool +// - dbtx: Database connection or transaction // - params: Query parameters from pipeline context // // Returns primary key values as strings, or error on database failure. -func QueryRowsPipeline(ctx context.Context, pool *pgxpool.Pool, params QueryParams) ([]string, error) { +func QueryRowsPipeline(ctx context.Context, dbtx DBTX, params QueryParams) ([]string, error) { logging.Debug("Executing pipeline query", zap.String("schema", params.Schema), zap.String("table", params.Table), @@ -128,7 +127,7 @@ func QueryRowsPipeline(ctx context.Context, pool *pgxpool.Pool, params QueryPara zap.String("sql", query), zap.Int("param_count", len(queryParams))) - rows, err := pool.Query(ctx, query, queryParams...) + rows, err := dbtx.Query(ctx, query, queryParams...) if err != nil { return nil, fmt.Errorf("failed to execute pipeline query: %w", err) } @@ -163,11 +162,13 @@ func QueryRowsPipeline(ctx context.Context, pool *pgxpool.Pool, params QueryPara } // QueryRowsPipeline is a convenience wrapper for Client. -func (c *Client) QueryRowsPipeline(ctx context.Context, params QueryParams) ([]string, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) QueryRowsPipeline(ctx context.Context, params QueryParams) (result []string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } - return QueryRowsPipeline(ctx, c.pool, params) + defer func() { done(retErr) }() + return QueryRowsPipeline(ctx, q, params) } // QueryRowsWithDataPipeline executes a pipeline query and returns full row data. @@ -175,11 +176,11 @@ func (c *Client) QueryRowsPipeline(ctx context.Context, params QueryParams) ([]s // // Parameters: // - ctx: Context for cancellation and timeout -// - pool: PostgreSQL connection pool +// - dbtx: Database connection or transaction // - params: Query parameters from pipeline context // // Returns column names and row data, or error on database failure. -func QueryRowsWithDataPipeline(ctx context.Context, pool *pgxpool.Pool, params QueryParams) ([]string, [][]interface{}, error) { +func QueryRowsWithDataPipeline(ctx context.Context, dbtx DBTX, params QueryParams) ([]string, [][]interface{}, error) { logging.Debug("Executing pipeline query with full data", zap.String("schema", params.Schema), zap.String("table", params.Table), @@ -194,7 +195,7 @@ func QueryRowsWithDataPipeline(ctx context.Context, pool *pgxpool.Pool, params Q zap.String("sql", query), zap.Int("param_count", len(queryParams))) - rows, err := pool.Query(ctx, query, queryParams...) + rows, err := dbtx.Query(ctx, query, queryParams...) if err != nil { return nil, nil, fmt.Errorf("failed to execute pipeline query: %w", err) } @@ -229,11 +230,13 @@ func QueryRowsWithDataPipeline(ctx context.Context, pool *pgxpool.Pool, params Q } // QueryRowsWithDataPipeline is a convenience wrapper for Client. -func (c *Client) QueryRowsWithDataPipeline(ctx context.Context, params QueryParams) ([]string, [][]interface{}, error) { - if c.pool == nil { - return nil, nil, fmt.Errorf("database connection not initialized") +func (c *Client) QueryRowsWithDataPipeline(ctx context.Context, params QueryParams) (columns []string, rows [][]interface{}, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, nil, err } - return QueryRowsWithDataPipeline(ctx, c.pool, params) + defer func() { done(retErr) }() + return QueryRowsWithDataPipeline(ctx, q, params) } // buildPipelineSQL constructs the SQL query for a pipeline operation. diff --git a/internal/tigerfs/db/query.go b/internal/tigerfs/db/query.go index b73d4bf..90da522 100644 --- a/internal/tigerfs/db/query.go +++ b/internal/tigerfs/db/query.go @@ -7,7 +7,6 @@ import ( "strings" "github.com/jackc/pgx/v5/pgconn" - "github.com/jackc/pgx/v5/pgxpool" "github.com/timescale/tigerfs/internal/tigerfs/format" "github.com/timescale/tigerfs/internal/tigerfs/logging" "go.uber.org/zap" @@ -82,7 +81,7 @@ func scanAndEncodePK(values []interface{}, pkColumns []string) (string, error) { } // GetRow fetches a single row by primary key (single or composite). -func GetRow(ctx context.Context, pool *pgxpool.Pool, schema, table string, pk *PKMatch) (*Row, error) { +func GetRow(ctx context.Context, dbtx DBTX, schema, table string, pk *PKMatch) (*Row, error) { logging.Debug("Querying row", zap.String("schema", schema), zap.String("table", table), @@ -95,7 +94,7 @@ func GetRow(ctx context.Context, pool *pgxpool.Pool, schema, table string, pk *P qt(schema, table), pk.WhereClause(1), ) - rows, err := pool.Query(ctx, query, pk.WhereArgs()...) + rows, err := dbtx.Query(ctx, query, pk.WhereArgs()...) if err != nil { return nil, fmt.Errorf("failed to query row: %w", err) } @@ -145,15 +144,17 @@ func GetRow(ctx context.Context, pool *pgxpool.Pool, schema, table string, pk *P } // GetRow is a convenience wrapper for Client -func (c *Client) GetRow(ctx context.Context, schema, table string, pk *PKMatch) (*Row, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetRow(ctx context.Context, schema, table string, pk *PKMatch) (result *Row, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } - return GetRow(ctx, c.pool, schema, table, pk) + defer func() { done(retErr) }() + return GetRow(ctx, q, schema, table, pk) } // GetColumn fetches a single column value from a row by primary key (single or composite). -func GetColumn(ctx context.Context, pool *pgxpool.Pool, schema, table string, pk *PKMatch, columnName string) (interface{}, error) { +func GetColumn(ctx context.Context, dbtx DBTX, schema, table string, pk *PKMatch, columnName string) (interface{}, error) { logging.Debug("Querying column", zap.String("schema", schema), zap.String("table", table), @@ -167,7 +168,7 @@ func GetColumn(ctx context.Context, pool *pgxpool.Pool, schema, table string, pk ) var value interface{} - err := pool.QueryRow(ctx, query, pk.WhereArgs()...).Scan(&value) + err := dbtx.QueryRow(ctx, query, pk.WhereArgs()...).Scan(&value) if err != nil { return nil, fmt.Errorf("failed to query column: %w", err) } @@ -182,16 +183,18 @@ func GetColumn(ctx context.Context, pool *pgxpool.Pool, schema, table string, pk } // GetColumn is a convenience wrapper for Client -func (c *Client) GetColumn(ctx context.Context, schema, table string, pk *PKMatch, columnName string) (interface{}, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetColumn(ctx context.Context, schema, table string, pk *PKMatch, columnName string) (result interface{}, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } - return GetColumn(ctx, c.pool, schema, table, pk, columnName) + defer func() { done(retErr) }() + return GetColumn(ctx, q, schema, table, pk, columnName) } // UpdateColumn updates a single column value for a row by primary key (single or composite). // Empty string is treated as NULL. -func UpdateColumn(ctx context.Context, pool *pgxpool.Pool, schema, table string, pk *PKMatch, columnName, newValue string) error { +func UpdateColumn(ctx context.Context, dbtx DBTX, schema, table string, pk *PKMatch, columnName, newValue string) error { logging.Debug("Updating column", zap.String("schema", schema), zap.String("table", table), @@ -217,7 +220,7 @@ func UpdateColumn(ctx context.Context, pool *pgxpool.Pool, schema, table string, args := append([]interface{}{value}, pk.WhereArgs()...) // Execute update - cmdTag, err := pool.Exec(ctx, query, args...) + cmdTag, err := dbtx.Exec(ctx, query, args...) if err != nil { return fmt.Errorf("failed to update column: %w", err) } @@ -237,20 +240,24 @@ func UpdateColumn(ctx context.Context, pool *pgxpool.Pool, schema, table string, } // UpdateColumn is a convenience wrapper for Client -func (c *Client) UpdateColumn(ctx context.Context, schema, table string, pk *PKMatch, columnName, newValue string) error { - if c.pool == nil { - return fmt.Errorf("database connection not initialized") +func (c *Client) UpdateColumn(ctx context.Context, schema, table string, pk *PKMatch, columnName, newValue string) (retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return err } - return UpdateColumn(ctx, c.pool, schema, table, pk, columnName, newValue) + defer func() { done(retErr) }() + return UpdateColumn(ctx, q, schema, table, pk, columnName, newValue) } // UpdateColumnCAS performs a compare-and-swap update on a column. // Only updates the row if whereColumn still has whereValue (atomic check). // Returns "row not found" if no row matches, enabling safe concurrent renames. -func (c *Client) UpdateColumnCAS(ctx context.Context, schema, table string, pk *PKMatch, setColumn, newValue, whereColumn, whereValue string) error { - if c.pool == nil { - return fmt.Errorf("database connection not initialized") +func (c *Client) UpdateColumnCAS(ctx context.Context, schema, table string, pk *PKMatch, setColumn, newValue, whereColumn, whereValue string) (retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return err } + defer func() { done(retErr) }() logging.Debug("CAS updating column", zap.String("schema", schema), @@ -272,7 +279,7 @@ func (c *Client) UpdateColumnCAS(ctx context.Context, schema, table string, pk * args := append([]interface{}{newValue}, pk.WhereArgs()...) args = append(args, whereValue) - cmdTag, err := c.pool.Exec(ctx, query, args...) + cmdTag, err := q.Exec(ctx, query, args...) if err != nil { return fmt.Errorf("failed to update column: %w", err) } @@ -286,7 +293,7 @@ func (c *Client) UpdateColumnCAS(ctx context.Context, schema, table string, pk * // InsertRow inserts a new row with the given column values // Returns the inserted primary key value (useful for auto-generated PKs) -func InsertRow(ctx context.Context, pool *pgxpool.Pool, schema, table string, columns []string, values []interface{}) (string, error) { +func InsertRow(ctx context.Context, dbtx DBTX, schema, table string, columns []string, values []interface{}) (string, error) { logging.Debug("Inserting row", zap.String("schema", schema), zap.String("table", table), @@ -321,7 +328,7 @@ func InsertRow(ctx context.Context, pool *pgxpool.Pool, schema, table string, co ) // Execute insert - rows, err := pool.Query(ctx, query, values...) + rows, err := dbtx.Query(ctx, query, values...) if err != nil { return "", fmt.Errorf("failed to insert row: %w", err) } @@ -361,15 +368,17 @@ func InsertRow(ctx context.Context, pool *pgxpool.Pool, schema, table string, co } // InsertRow is a convenience wrapper for Client -func (c *Client) InsertRow(ctx context.Context, schema, table string, columns []string, values []interface{}) (string, error) { - if c.pool == nil { - return "", fmt.Errorf("database connection not initialized") +func (c *Client) InsertRow(ctx context.Context, schema, table string, columns []string, values []interface{}) (result string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return "", err } - return InsertRow(ctx, c.pool, schema, table, columns, values) + defer func() { done(retErr) }() + return InsertRow(ctx, q, schema, table, columns, values) } // UpdateRow updates an existing row with the given column values. -func UpdateRow(ctx context.Context, pool *pgxpool.Pool, schema, table string, pk *PKMatch, columns []string, values []interface{}) error { +func UpdateRow(ctx context.Context, dbtx DBTX, schema, table string, pk *PKMatch, columns []string, values []interface{}) error { logging.Debug("Updating row", zap.String("schema", schema), zap.String("table", table), @@ -403,7 +412,7 @@ func UpdateRow(ctx context.Context, pool *pgxpool.Pool, schema, table string, pk allValues := append(values, pk.WhereArgs()...) // Execute update - cmdTag, err := pool.Exec(ctx, query, allValues...) + cmdTag, err := dbtx.Exec(ctx, query, allValues...) if err != nil { return fmt.Errorf("failed to update row: %w", err) } @@ -422,15 +431,17 @@ func UpdateRow(ctx context.Context, pool *pgxpool.Pool, schema, table string, pk } // UpdateRow is a convenience wrapper for Client -func (c *Client) UpdateRow(ctx context.Context, schema, table string, pk *PKMatch, columns []string, values []interface{}) error { - if c.pool == nil { - return fmt.Errorf("database connection not initialized") +func (c *Client) UpdateRow(ctx context.Context, schema, table string, pk *PKMatch, columns []string, values []interface{}) (retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return err } - return UpdateRow(ctx, c.pool, schema, table, pk, columns, values) + defer func() { done(retErr) }() + return UpdateRow(ctx, q, schema, table, pk, columns, values) } // DeleteRow deletes a row by primary key (single or composite). -func DeleteRow(ctx context.Context, pool *pgxpool.Pool, schema, table string, pk *PKMatch) error { +func DeleteRow(ctx context.Context, dbtx DBTX, schema, table string, pk *PKMatch) error { logging.Debug("Deleting row", zap.String("schema", schema), zap.String("table", table), @@ -444,7 +455,7 @@ func DeleteRow(ctx context.Context, pool *pgxpool.Pool, schema, table string, pk ) // Execute delete - cmdTag, err := pool.Exec(ctx, query, pk.WhereArgs()...) + cmdTag, err := dbtx.Exec(ctx, query, pk.WhereArgs()...) if err != nil { return fmt.Errorf("failed to delete row: %w", err) } @@ -463,17 +474,23 @@ func DeleteRow(ctx context.Context, pool *pgxpool.Pool, schema, table string, pk } // DeleteRow is a convenience wrapper for Client -func (c *Client) DeleteRow(ctx context.Context, schema, table string, pk *PKMatch) error { - if c.pool == nil { - return fmt.Errorf("database connection not initialized") +func (c *Client) DeleteRow(ctx context.Context, schema, table string, pk *PKMatch) (retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return err } - return DeleteRow(ctx, c.pool, schema, table, pk) + defer func() { done(retErr) }() + return DeleteRow(ctx, q, schema, table, pk) } // DeleteAndUpdate atomically deletes one row and updates another in a single // PostgreSQL transaction. Used for POSIX rename-as-replace semantics where // the target file must be removed before the source file can be renamed to it. // Both BEFORE triggers (history capture) fire within the transaction. +// +// This method manages its own transaction and must NOT use acquireDBTX, +// because it needs explicit transaction control for atomicity. +// Session vars are injected directly via applySessionVars after Begin. func (c *Client) DeleteAndUpdate(ctx context.Context, schema, table string, deletePK *PKMatch, updatePK *PKMatch, updateCols []string, updateVals []interface{}) error { @@ -487,6 +504,13 @@ func (c *Client) DeleteAndUpdate(ctx context.Context, schema, table string, } defer tx.Rollback(ctx) + // Apply session variables (SET LOCAL) if configured + if vars := c.effectiveSessionVars(ctx); !vars.Empty() { + if err := applySessionVars(ctx, tx, vars); err != nil { + return fmt.Errorf("failed to apply session variables: %w", err) + } + } + // Step 1: Delete the target row deleteSQL := fmt.Sprintf(`DELETE FROM %s WHERE %s`, qt(schema, table), deletePK.WhereClause(1)) _, err = tx.Exec(ctx, deleteSQL, deletePK.WhereArgs()...) @@ -513,7 +537,7 @@ func (c *Client) DeleteAndUpdate(ctx context.Context, schema, table string, // GetFirstNRows returns the first N primary key values ordered by PK ascending. // Returns encoded PK strings (comma-delimited for composite PKs). -func GetFirstNRows(ctx context.Context, pool *pgxpool.Pool, schema, table string, pkColumns []string, limit int) ([]string, error) { +func GetFirstNRows(ctx context.Context, dbtx DBTX, schema, table string, pkColumns []string, limit int) ([]string, error) { logging.Debug("Getting first N rows", zap.String("schema", schema), zap.String("table", table), @@ -525,20 +549,22 @@ func GetFirstNRows(ctx context.Context, pool *pgxpool.Pool, schema, table string pkSelectList(pkColumns), qt(schema, table), pkOrderByList(pkColumns, "ASC"), ) - return scanPKRows(ctx, pool, query, pkColumns, limit) + return scanPKRows(ctx, dbtx, query, pkColumns, limit) } // GetFirstNRows is a convenience wrapper for Client -func (c *Client) GetFirstNRows(ctx context.Context, schema, table string, pkColumns []string, limit int) ([]string, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetFirstNRows(ctx context.Context, schema, table string, pkColumns []string, limit int) (result []string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } - return GetFirstNRows(ctx, c.pool, schema, table, pkColumns, limit) + defer func() { done(retErr) }() + return GetFirstNRows(ctx, q, schema, table, pkColumns, limit) } // GetLastNRows returns the last N primary key values ordered by PK descending. // Returns encoded PK strings (comma-delimited for composite PKs). -func GetLastNRows(ctx context.Context, pool *pgxpool.Pool, schema, table string, pkColumns []string, limit int) ([]string, error) { +func GetLastNRows(ctx context.Context, dbtx DBTX, schema, table string, pkColumns []string, limit int) ([]string, error) { logging.Debug("Getting last N rows", zap.String("schema", schema), zap.String("table", table), @@ -550,15 +576,17 @@ func GetLastNRows(ctx context.Context, pool *pgxpool.Pool, schema, table string, pkSelectList(pkColumns), qt(schema, table), pkOrderByList(pkColumns, "DESC"), ) - return scanPKRows(ctx, pool, query, pkColumns, limit) + return scanPKRows(ctx, dbtx, query, pkColumns, limit) } // GetLastNRows is a convenience wrapper for Client -func (c *Client) GetLastNRows(ctx context.Context, schema, table string, pkColumns []string, limit int) ([]string, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetLastNRows(ctx context.Context, schema, table string, pkColumns []string, limit int) (result []string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } - return GetLastNRows(ctx, c.pool, schema, table, pkColumns, limit) + defer func() { done(retErr) }() + return GetLastNRows(ctx, q, schema, table, pkColumns, limit) } // GetRandomSampleRows returns approximately N random primary key values. @@ -577,7 +605,7 @@ func (c *Client) GetLastNRows(ctx context.Context, schema, table string, pkColum // // Returns primary keys as strings. The actual count may vary from the // requested limit due to the probabilistic nature of TABLESAMPLE. -func GetRandomSampleRows(ctx context.Context, pool *pgxpool.Pool, schema, table string, pkColumns []string, limit int, estimatedRows int64) ([]string, error) { +func GetRandomSampleRows(ctx context.Context, dbtx DBTX, schema, table string, pkColumns []string, limit int, estimatedRows int64) ([]string, error) { logging.Debug("Getting random sample rows", zap.String("schema", schema), zap.String("table", table), @@ -609,15 +637,17 @@ func GetRandomSampleRows(ctx context.Context, pool *pgxpool.Pool, schema, table ) } - return scanPKRows(ctx, pool, query, pkColumns, limit) + return scanPKRows(ctx, dbtx, query, pkColumns, limit) } // GetRandomSampleRows is a convenience wrapper for Client -func (c *Client) GetRandomSampleRows(ctx context.Context, schema, table string, pkColumns []string, limit int, estimatedRows int64) ([]string, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetRandomSampleRows(ctx context.Context, schema, table string, pkColumns []string, limit int, estimatedRows int64) (result []string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } - return GetRandomSampleRows(ctx, c.pool, schema, table, pkColumns, limit, estimatedRows) + defer func() { done(retErr) }() + return GetRandomSampleRows(ctx, q, schema, table, pkColumns, limit, estimatedRows) } // GetFirstNRowsOrdered returns the first N primary key values ordered by a specified column ascending. @@ -631,7 +661,7 @@ func (c *Client) GetRandomSampleRows(ctx context.Context, schema, table string, // - limit: Maximum number of rows to return // // Returns primary key values as strings, ordered by orderColumn ASC. -func GetFirstNRowsOrdered(ctx context.Context, pool *pgxpool.Pool, schema, table string, pkColumns []string, orderColumn string, limit int) ([]string, error) { +func GetFirstNRowsOrdered(ctx context.Context, dbtx DBTX, schema, table string, pkColumns []string, orderColumn string, limit int) ([]string, error) { logging.Debug("Getting first N rows ordered by column", zap.String("schema", schema), zap.String("table", table), @@ -644,15 +674,17 @@ func GetFirstNRowsOrdered(ctx context.Context, pool *pgxpool.Pool, schema, table pkSelectList(pkColumns), qt(schema, table), qi(orderColumn), pkOrderByList(pkColumns, "ASC"), ) - return scanPKRows(ctx, pool, query, pkColumns, limit) + return scanPKRows(ctx, dbtx, query, pkColumns, limit) } // GetFirstNRowsOrdered is a convenience wrapper for Client -func (c *Client) GetFirstNRowsOrdered(ctx context.Context, schema, table string, pkColumns []string, orderColumn string, limit int) ([]string, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetFirstNRowsOrdered(ctx context.Context, schema, table string, pkColumns []string, orderColumn string, limit int) (result []string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } - return GetFirstNRowsOrdered(ctx, c.pool, schema, table, pkColumns, orderColumn, limit) + defer func() { done(retErr) }() + return GetFirstNRowsOrdered(ctx, q, schema, table, pkColumns, orderColumn, limit) } // GetLastNRowsOrdered returns the last N primary key values ordered by a specified column descending. @@ -666,7 +698,7 @@ func (c *Client) GetFirstNRowsOrdered(ctx context.Context, schema, table string, // - limit: Maximum number of rows to return // // Returns primary key values as strings, ordered by orderColumn DESC. -func GetLastNRowsOrdered(ctx context.Context, pool *pgxpool.Pool, schema, table string, pkColumns []string, orderColumn string, limit int) ([]string, error) { +func GetLastNRowsOrdered(ctx context.Context, dbtx DBTX, schema, table string, pkColumns []string, orderColumn string, limit int) ([]string, error) { logging.Debug("Getting last N rows ordered by column", zap.String("schema", schema), zap.String("table", table), @@ -679,21 +711,23 @@ func GetLastNRowsOrdered(ctx context.Context, pool *pgxpool.Pool, schema, table pkSelectList(pkColumns), qt(schema, table), qi(orderColumn), pkOrderByList(pkColumns, "DESC"), ) - return scanPKRows(ctx, pool, query, pkColumns, limit) + return scanPKRows(ctx, dbtx, query, pkColumns, limit) } // GetLastNRowsOrdered is a convenience wrapper for Client -func (c *Client) GetLastNRowsOrdered(ctx context.Context, schema, table string, pkColumns []string, orderColumn string, limit int) ([]string, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetLastNRowsOrdered(ctx context.Context, schema, table string, pkColumns []string, orderColumn string, limit int) (result []string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } - return GetLastNRowsOrdered(ctx, c.pool, schema, table, pkColumns, orderColumn, limit) + defer func() { done(retErr) }() + return GetLastNRowsOrdered(ctx, q, schema, table, pkColumns, orderColumn, limit) } // scanPKRows executes a query with a LIMIT $1 parameter, scans PK columns, and // returns encoded PK strings. Shared by all pagination functions. -func scanPKRows(ctx context.Context, pool *pgxpool.Pool, query string, pkColumns []string, limit int) ([]string, error) { - rows, err := pool.Query(ctx, query, limit) +func scanPKRows(ctx context.Context, dbtx DBTX, query string, pkColumns []string, limit int) ([]string, error) { + rows, err := dbtx.Query(ctx, query, limit) if err != nil { return nil, fmt.Errorf("failed to query rows: %w", err) } @@ -727,10 +761,12 @@ func scanPKRows(ctx context.Context, pool *pgxpool.Pool, query string, pkColumns // RenameByPrefix atomically renames all rows where the given column value starts // with oldPrefix to use newPrefix instead. Used for directory renames in synth views. // Returns the number of rows affected. -func (c *Client) RenameByPrefix(ctx context.Context, schema, table, column, oldPrefix, newPrefix string) (int64, error) { - if c.pool == nil { - return 0, fmt.Errorf("database connection not initialized") +func (c *Client) RenameByPrefix(ctx context.Context, schema, table, column, oldPrefix, newPrefix string) (result int64, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return 0, err } + defer func() { done(retErr) }() // UPDATE "schema"."table" SET "column" = $1 || substr("column", length($2) + 1) // WHERE "column" = $2 OR "column" LIKE $2 || '/%' @@ -739,7 +775,7 @@ func (c *Client) RenameByPrefix(ctx context.Context, schema, table, column, oldP qt(schema, table), qi(column), qi(column), qi(column), qi(column), ) - cmdTag, err := c.pool.Exec(ctx, query, newPrefix, oldPrefix) + cmdTag, err := q.Exec(ctx, query, newPrefix, oldPrefix) if err != nil { return 0, fmt.Errorf("failed to rename by prefix: %w", err) } @@ -749,10 +785,12 @@ func (c *Client) RenameByPrefix(ctx context.Context, schema, table, column, oldP // HasChildrenWithPrefix checks if any rows exist where the given column value // starts with prefix + "/". Used to check if a directory has children before rmdir. -func (c *Client) HasChildrenWithPrefix(ctx context.Context, schema, table, column, prefix string) (bool, error) { - if c.pool == nil { - return false, fmt.Errorf("database connection not initialized") +func (c *Client) HasChildrenWithPrefix(ctx context.Context, schema, table, column, prefix string) (result bool, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return false, err } + defer func() { done(retErr) }() query := fmt.Sprintf( `SELECT EXISTS(SELECT 1 FROM %s WHERE %s LIKE $1 || '/%%')`, @@ -760,7 +798,7 @@ func (c *Client) HasChildrenWithPrefix(ctx context.Context, schema, table, colum ) var exists bool - err := c.pool.QueryRow(ctx, query, prefix).Scan(&exists) + err = q.QueryRow(ctx, query, prefix).Scan(&exists) if err != nil { return false, fmt.Errorf("failed to check children: %w", err) } @@ -770,10 +808,12 @@ func (c *Client) HasChildrenWithPrefix(ctx context.Context, schema, table, colum // InsertIfNotExists inserts a row only if it doesn't already exist (ON CONFLICT DO NOTHING). // Used for auto-creating parent directory rows in synth views. -func (c *Client) InsertIfNotExists(ctx context.Context, schema, table string, columns []string, values []interface{}) error { - if c.pool == nil { - return fmt.Errorf("database connection not initialized") +func (c *Client) InsertIfNotExists(ctx context.Context, schema, table string, columns []string, values []interface{}) (retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return err } + defer func() { done(retErr) }() if len(columns) == 0 || len(columns) != len(values) { return fmt.Errorf("column/value count mismatch: %d columns, %d values", len(columns), len(values)) @@ -795,7 +835,7 @@ func (c *Client) InsertIfNotExists(ctx context.Context, schema, table string, co qt(schema, table), strings.Join(quotedCols, ", "), strings.Join(placeholders, ", "), ) - _, err := c.pool.Exec(ctx, query, values...) + _, err = q.Exec(ctx, query, values...) if err != nil { if isUniqueViolation(err) { return nil // Row already exists, no-op @@ -809,10 +849,12 @@ func (c *Client) InsertIfNotExists(ctx context.Context, schema, table string, co // GetRowsByParent returns rows with a specific parent_id, up to limit. // Empty parentID means root level (WHERE parent_id IS NULL). // Used by the parent-pointer directory model (ADR-017) for ReadDir. -func (c *Client) GetRowsByParent(ctx context.Context, schema, table, parentID string, limit int) ([]string, [][]interface{}, error) { - if c.pool == nil { - return nil, nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetRowsByParent(ctx context.Context, schema, table, parentID string, limit int) (cols []string, rows [][]interface{}, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, nil, err } + defer func() { done(retErr) }() var query string var args []interface{} @@ -830,17 +872,19 @@ func (c *Client) GetRowsByParent(ctx context.Context, schema, table, parentID st args = []interface{}{parentID, limit} } - return c.queryRows(ctx, query, args...) + return c.queryRows(ctx, q, query, args...) } // GetRowByParentAndName returns a single row matching parent_id + filename. // Empty parentID means root level (WHERE parent_id IS NULL). // Returns (columns, row, error); row is nil if not found. // Used by ReadFile to combine path resolution with row fetch in one query (ADR-017). -func (c *Client) GetRowByParentAndName(ctx context.Context, schema, table, parentID, filename string) ([]string, []interface{}, error) { - if c.pool == nil { - return nil, nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetRowByParentAndName(ctx context.Context, schema, table, parentID, filename string) (cols []string, row []interface{}, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, nil, err } + defer func() { done(retErr) }() var query string var args []interface{} @@ -858,7 +902,7 @@ func (c *Client) GetRowByParentAndName(ctx context.Context, schema, table, paren args = []interface{}{parentID, filename} } - columns, rows, err := c.queryRows(ctx, query, args...) + columns, rows, err := c.queryRows(ctx, q, query, args...) if err != nil { return nil, nil, err } @@ -871,18 +915,19 @@ func (c *Client) GetRowByParentAndName(ctx context.Context, schema, table, paren // QueryNextLogEntry finds the next log entry for a file after a given log_id. // Returns the version_id and filename of the next entry, or empty strings if none. // Used by diff symlink resolution to determine the "after" state. -func (c *Client) QueryNextLogEntry(ctx context.Context, schema, logTable, fileID, afterLogID string) (string, string, error) { - if c.pool == nil { - return "", "", fmt.Errorf("database connection not initialized") +func (c *Client) QueryNextLogEntry(ctx context.Context, schema, logTable, fileID, afterLogID string) (versionID string, filename string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return "", "", err } + defer func() { done(retErr) }() query := fmt.Sprintf( `SELECT COALESCE("version_id"::text, ''), "filename" FROM %s WHERE "file_id" = $1 AND "log_id" > $2 ORDER BY "log_id" ASC LIMIT 1`, qt(schema, logTable), ) - var versionID, filename string - err := c.pool.QueryRow(ctx, query, fileID, afterLogID).Scan(&versionID, &filename) + err = q.QueryRow(ctx, query, fileID, afterLogID).Scan(&versionID, &filename) if err != nil { if err.Error() == "no rows in result set" { return "", "", nil // No next entry @@ -894,10 +939,12 @@ func (c *Client) QueryNextLogEntry(ctx context.Context, schema, logTable, fileID // QueryFileExists checks if a row with the given id exists in the source table. // Used by diff symlink resolution to determine if a file still exists. -func (c *Client) QueryFileExists(ctx context.Context, schema, table, fileID string) (bool, error) { - if c.pool == nil { - return false, fmt.Errorf("database connection not initialized") +func (c *Client) QueryFileExists(ctx context.Context, schema, table, fileID string) (result bool, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return false, err } + defer func() { done(retErr) }() query := fmt.Sprintf( `SELECT EXISTS(SELECT 1 FROM %s WHERE "id" = $1)`, @@ -905,7 +952,7 @@ func (c *Client) QueryFileExists(ctx context.Context, schema, table, fileID stri ) var exists bool - err := c.pool.QueryRow(ctx, query, fileID).Scan(&exists) + err = q.QueryRow(ctx, query, fileID).Scan(&exists) if err != nil { return false, fmt.Errorf("failed to check file existence: %w", err) } @@ -915,10 +962,12 @@ func (c *Client) QueryFileExists(ctx context.Context, schema, table, fileID stri // QueryUndoAffectedFiles returns the first log entry per file after a target point. // Uses DISTINCT ON to find one entry per file_id, ordered by log_id ASC (oldest first). // TimescaleDB's SkipScan optimizes this on the (file_id, log_id ASC) index. -func (c *Client) QueryUndoAffectedFiles(ctx context.Context, schema, logTable, afterID, userID string, filters []UndoFilter) ([]UndoAffectedFile, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) QueryUndoAffectedFiles(ctx context.Context, schema, logTable, afterID, userID string, filters []UndoFilter) (result []UndoAffectedFile, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } + defer func() { done(retErr) }() // Build WHERE clause conditions := []string{fmt.Sprintf("%s > $1", qi("log_id"))} @@ -948,13 +997,12 @@ func (c *Client) QueryUndoAffectedFiles(ctx context.Context, schema, logTable, a qi("file_id"), qi("log_id"), ) - rows, err := c.pool.Query(ctx, query, args...) + rows, err := q.Query(ctx, query, args...) if err != nil { return nil, fmt.Errorf("failed to query undo affected files: %w", err) } defer rows.Close() - var result []UndoAffectedFile for rows.Next() { var f UndoAffectedFile var versionID *string @@ -974,10 +1022,12 @@ func (c *Client) QueryUndoAffectedFiles(ctx context.Context, schema, logTable, a } // QueryLogEntry fetches a single log entry by log_id. -func (c *Client) QueryLogEntry(ctx context.Context, schema, logTable, logID string) (*UndoAffectedFile, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) QueryLogEntry(ctx context.Context, schema, logTable, logID string) (result *UndoAffectedFile, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } + defer func() { done(retErr) }() query := fmt.Sprintf( `SELECT %s, %s, %s, %s, %s, %s FROM %s WHERE %s = $1`, @@ -988,7 +1038,7 @@ func (c *Client) QueryLogEntry(ctx context.Context, schema, logTable, logID stri var f UndoAffectedFile var versionID, userID *string - err := c.pool.QueryRow(ctx, query, logID).Scan(&f.FileID, &f.Type, &versionID, &f.Filename, &f.LogID, &userID) + err = q.QueryRow(ctx, query, logID).Scan(&f.FileID, &f.Type, &versionID, &f.Filename, &f.LogID, &userID) if err != nil { return nil, fmt.Errorf("log entry not found: %s", logID) } @@ -1005,6 +1055,11 @@ func (c *Client) QueryLogEntry(ctx context.Context, schema, logTable, logID stri // Within a single transaction: deletes created rows, upserts from history, // and inserts undo log entries. BEFORE triggers fire on each restore, // creating history entries for the undo itself. +// +// This method manages its own transaction and must NOT use acquireDBTX, +// because it needs explicit transaction control for atomicity (deferred +// constraints, multi-step DML). Session vars are injected directly via +// applySessionVars after Begin. func (c *Client) ExecuteUndoTransaction(ctx context.Context, params *UndoTransactionParams) error { if c.pool == nil { return fmt.Errorf("database connection not initialized") @@ -1016,6 +1071,13 @@ func (c *Client) ExecuteUndoTransaction(ctx context.Context, params *UndoTransac } defer tx.Rollback(ctx) + // Apply session variables (SET LOCAL) if configured + if vars := c.effectiveSessionVars(ctx); !vars.Empty() { + if err := applySessionVars(ctx, tx, vars); err != nil { + return fmt.Errorf("failed to apply session variables: %w", err) + } + } + // Defer FK and UNIQUE checks to COMMIT so we can DELETE/INSERT/UPSERT // rows in any order within the undo. This is the deferral path documented // in ADR-017: parent_id_fkey and the (parent_id, filename, filetype) @@ -1224,9 +1286,15 @@ func (c *Client) ExecuteUndoTransaction(ctx context.Context, params *UndoTransac } // HasExtension checks if a PostgreSQL extension is installed in the database. -func (c *Client) HasExtension(ctx context.Context, extName string) (bool, error) { +func (c *Client) HasExtension(ctx context.Context, extName string) (result bool, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return false, err + } + defer func() { done(retErr) }() + var exists bool - err := c.pool.QueryRow(ctx, + err = q.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM pg_extension WHERE extname = $1)`, extName, ).Scan(&exists) @@ -1237,9 +1305,15 @@ func (c *Client) HasExtension(ctx context.Context, extName string) (bool, error) } // TableExists checks if a table exists in the given schema. -func (c *Client) TableExists(ctx context.Context, schema, table string) (bool, error) { +func (c *Client) TableExists(ctx context.Context, schema, table string) (result bool, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return false, err + } + defer func() { done(retErr) }() + var exists bool - err := c.pool.QueryRow(ctx, + err = q.QueryRow(ctx, `SELECT EXISTS( SELECT 1 FROM information_schema.tables WHERE table_schema = $1 AND table_name = $2 @@ -1254,31 +1328,49 @@ func (c *Client) TableExists(ctx context.Context, schema, table string) (bool, e // QueryHistoryByFilename queries the history table for versions of a file by filename. // Returns columns and rows ordered by version_id DESC (most recent first). -func (c *Client) QueryHistoryByFilename(ctx context.Context, schema, historyTable, filename string, limit int) ([]string, [][]interface{}, error) { +func (c *Client) QueryHistoryByFilename(ctx context.Context, schema, historyTable, filename string, limit int) (cols []string, rows [][]interface{}, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, nil, err + } + defer func() { done(retErr) }() + query := fmt.Sprintf( `SELECT * FROM %s WHERE "filename" = $1 ORDER BY "version_id" DESC LIMIT %d`, qt(schema, historyTable), limit, ) - return c.queryRows(ctx, query, filename) + return c.queryRows(ctx, q, query, filename) } // QueryHistoryByID queries the history table for versions of a row by its file_id UUID. // Returns columns and rows ordered by version_id DESC (most recent first). -func (c *Client) QueryHistoryByID(ctx context.Context, schema, historyTable, rowID string, limit int) ([]string, [][]interface{}, error) { +func (c *Client) QueryHistoryByID(ctx context.Context, schema, historyTable, rowID string, limit int) (cols []string, rows [][]interface{}, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, nil, err + } + defer func() { done(retErr) }() + query := fmt.Sprintf( `SELECT * FROM %s WHERE "file_id" = $1 ORDER BY "version_id" DESC LIMIT %d`, qt(schema, historyTable), limit, ) - return c.queryRows(ctx, query, rowID) + return c.queryRows(ctx, q, query, rowID) } // QueryHistoryDistinctFilenames returns distinct filenames from the history table. -func (c *Client) QueryHistoryDistinctFilenames(ctx context.Context, schema, historyTable string, limit int) ([]string, error) { +func (c *Client) QueryHistoryDistinctFilenames(ctx context.Context, schema, historyTable string, limit int) (result []string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err + } + defer func() { done(retErr) }() + query := fmt.Sprintf( `SELECT DISTINCT "filename" FROM %s ORDER BY "filename" LIMIT %d`, qt(schema, historyTable), limit, ) - rows, err := c.pool.Query(ctx, query) + rows, err := q.Query(ctx, query) if err != nil { return nil, fmt.Errorf("failed to query distinct filenames: %w", err) } @@ -1298,7 +1390,13 @@ func (c *Client) QueryHistoryDistinctFilenames(ctx context.Context, schema, hist // QueryHistoryDistinctFilenamesByParent returns distinct filenames from the history table // filtered by parent_id. Empty parentID means root level (WHERE parent_id IS NULL). // Used by the parent-pointer model (ADR-017) for per-directory .history/ listing. -func (c *Client) QueryHistoryDistinctFilenamesByParent(ctx context.Context, schema, historyTable, parentID string, limit int) ([]string, error) { +func (c *Client) QueryHistoryDistinctFilenamesByParent(ctx context.Context, schema, historyTable, parentID string, limit int) (result []string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err + } + defer func() { done(retErr) }() + var query string var args []interface{} if parentID == "" { @@ -1314,7 +1412,7 @@ func (c *Client) QueryHistoryDistinctFilenamesByParent(ctx context.Context, sche args = []interface{}{parentID} } - rows, err := c.pool.Query(ctx, query, args...) + rows, err := q.Query(ctx, query, args...) if err != nil { return nil, fmt.Errorf("failed to query distinct filenames by parent: %w", err) } @@ -1332,12 +1430,18 @@ func (c *Client) QueryHistoryDistinctFilenamesByParent(ctx context.Context, sche } // QueryHistoryDistinctIDs returns distinct row UUIDs (file_id) from the history table. -func (c *Client) QueryHistoryDistinctIDs(ctx context.Context, schema, historyTable string, limit int) ([]string, error) { +func (c *Client) QueryHistoryDistinctIDs(ctx context.Context, schema, historyTable string, limit int) (result []string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err + } + defer func() { done(retErr) }() + query := fmt.Sprintf( `SELECT DISTINCT "file_id"::text FROM %s ORDER BY "file_id" LIMIT %d`, qt(schema, historyTable), limit, ) - rows, err := c.pool.Query(ctx, query) + rows, err := q.Query(ctx, query) if err != nil { return nil, fmt.Errorf("failed to query distinct IDs: %w", err) } @@ -1357,16 +1461,28 @@ func (c *Client) QueryHistoryDistinctIDs(ctx context.Context, schema, historyTab // QueryHistoryVersionByTime finds a history row matching a version ID timestamp. // Version IDs have second precision; UUIDv7 has millisecond precision. This queries // with a 1-second window around the target timestamp plus filename or file_id filter. -func (c *Client) QueryHistoryVersionByTime(ctx context.Context, schema, historyTable, filterColumn, filterValue string, targetTime interface{}, limit int) ([]string, [][]interface{}, error) { +func (c *Client) QueryHistoryVersionByTime(ctx context.Context, schema, historyTable, filterColumn, filterValue string, targetTime interface{}, limit int) (cols []string, rows [][]interface{}, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, nil, err + } + defer func() { done(retErr) }() + query := fmt.Sprintf( `SELECT * FROM %s WHERE %s = $1 ORDER BY "version_id" DESC LIMIT %d`, qt(schema, historyTable), qi(filterColumn), limit, ) - return c.queryRows(ctx, query, filterValue) + return c.queryRows(ctx, q, query, filterValue) } // InsertLogEntry inserts an operation log entry into the log hypertable. -func (c *Client) InsertLogEntry(ctx context.Context, schema, logTable, userID, opType, fileID, filename, versionID, description string) error { +func (c *Client) InsertLogEntry(ctx context.Context, schema, logTable, userID, opType, fileID, filename, versionID, description string) (retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return err + } + defer func() { done(retErr) }() + query := fmt.Sprintf( `INSERT INTO %s (user_id, type, file_id, filename, version_id, description) VALUES ($1, $2, $3, $4, $5, $6)`, qt(schema, logTable), @@ -1384,7 +1500,7 @@ func (c *Client) InsertLogEntry(ctx context.Context, schema, logTable, userID, o descVal = description } - _, err := c.pool.Exec(ctx, query, userIDVal, opType, fileID, filename, versionIDVal, descVal) + _, err = q.Exec(ctx, query, userIDVal, opType, fileID, filename, versionIDVal, descVal) if err != nil { return fmt.Errorf("failed to insert log entry: %w", err) } @@ -1393,13 +1509,19 @@ func (c *Client) InsertLogEntry(ctx context.Context, schema, logTable, userID, o // QueryLatestVersionID returns the most recent version_id for a given // file_id from the history table. Returns empty string if no history entry found. -func (c *Client) QueryLatestVersionID(ctx context.Context, schema, historyTable, fileID string) (string, error) { +func (c *Client) QueryLatestVersionID(ctx context.Context, schema, historyTable, fileID string) (result string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return "", err + } + defer func() { done(retErr) }() + query := fmt.Sprintf( `SELECT "version_id" FROM %s WHERE "file_id" = $1 ORDER BY "version_id" DESC LIMIT 1`, qt(schema, historyTable), ) var versionID string - err := c.pool.QueryRow(ctx, query, fileID).Scan(&versionID) + err = q.QueryRow(ctx, query, fileID).Scan(&versionID) if err != nil { return "", fmt.Errorf("failed to query latest version ID: %w", err) } @@ -1418,11 +1540,17 @@ type PathSegment struct { // startParentID is the UUID of the starting parent (empty string for root/NULL). // Returns one PathSegment per resolved segment. If a segment doesn't resolve, // fewer segments are returned than requested (no error). -func (c *Client) ResolvePath(ctx context.Context, schema, table, startParentID string, segments []string) ([]PathSegment, error) { +func (c *Client) ResolvePath(ctx context.Context, schema, table, startParentID string, segments []string) (result []PathSegment, retErr error) { if len(segments) == 0 { return nil, nil } + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err + } + defer func() { done(retErr) }() + qualifiedTable := fmt.Sprintf("%s.%s", qi(schema), qi(table)) // Convert empty startParentID to SQL NULL @@ -1436,13 +1564,12 @@ func (c *Client) ResolvePath(ctx context.Context, schema, table, startParentID s qi("tigerfs"), ) - rows, err := c.pool.Query(ctx, query, qualifiedTable, parentArg, segments) + rows, err := q.Query(ctx, query, qualifiedTable, parentArg, segments) if err != nil { return nil, fmt.Errorf("resolve_path failed: %w", err) } defer rows.Close() - var result []PathSegment for rows.Next() { var seg PathSegment var id [16]byte @@ -1459,8 +1586,8 @@ func (c *Client) ResolvePath(ctx context.Context, schema, table, startParentID s } // queryRows executes a query and returns columns and row data. -func (c *Client) queryRows(ctx context.Context, query string, args ...interface{}) ([]string, [][]interface{}, error) { - rows, err := c.pool.Query(ctx, query, args...) +func (c *Client) queryRows(ctx context.Context, dbtx DBTX, query string, args ...interface{}) ([]string, [][]interface{}, error) { + rows, err := dbtx.Query(ctx, query, args...) if err != nil { return nil, nil, fmt.Errorf("query failed: %w", err) } diff --git a/internal/tigerfs/db/schema.go b/internal/tigerfs/db/schema.go index 951c707..705c8e4 100644 --- a/internal/tigerfs/db/schema.go +++ b/internal/tigerfs/db/schema.go @@ -6,7 +6,6 @@ import ( "strings" "github.com/jackc/pgx/v5" - "github.com/jackc/pgx/v5/pgxpool" "github.com/timescale/tigerfs/internal/tigerfs/logging" "go.uber.org/zap" ) @@ -17,14 +16,14 @@ import ( // // Parameters: // - ctx: Context for cancellation -// - pool: PostgreSQL connection pool +// - dbtx: DBTX (pool or transaction) // // Returns the current schema name, or error on database failure. -func GetCurrentSchema(ctx context.Context, pool *pgxpool.Pool) (string, error) { +func GetCurrentSchema(ctx context.Context, dbtx DBTX) (string, error) { logging.Debug("Querying current_schema from PostgreSQL") var schema string - err := pool.QueryRow(ctx, "SELECT current_schema()").Scan(&schema) + err := dbtx.QueryRow(ctx, "SELECT current_schema()").Scan(&schema) if err != nil { return "", fmt.Errorf("failed to query current_schema: %w", err) } @@ -34,15 +33,17 @@ func GetCurrentSchema(ctx context.Context, pool *pgxpool.Pool) (string, error) { } // GetCurrentSchema is a convenience wrapper around GetCurrentSchema for Client. -func (c *Client) GetCurrentSchema(ctx context.Context) (string, error) { - if c.pool == nil { - return "", fmt.Errorf("database connection not initialized") +func (c *Client) GetCurrentSchema(ctx context.Context) (result string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return "", err } - return GetCurrentSchema(ctx, c.pool) + defer func() { done(retErr) }() + return GetCurrentSchema(ctx, q) } // GetSchemas returns all user-defined schemas (excluding system schemas) -func GetSchemas(ctx context.Context, pool *pgxpool.Pool) ([]string, error) { +func GetSchemas(ctx context.Context, dbtx DBTX) ([]string, error) { logging.Debug("Querying schemas from information_schema") query := ` @@ -52,7 +53,7 @@ func GetSchemas(ctx context.Context, pool *pgxpool.Pool) ([]string, error) { ORDER BY schema_name ` - rows, err := pool.Query(ctx, query) + rows, err := dbtx.Query(ctx, query) if err != nil { return nil, fmt.Errorf("failed to query schemas: %w", err) } @@ -79,7 +80,7 @@ func GetSchemas(ctx context.Context, pool *pgxpool.Pool) ([]string, error) { // Excludes tables owned by extensions (e.g., pg_buffercache internal tables) // via pg_depend deptype='e'. Hypertables are NOT excluded — TimescaleDB // tracks them in its own catalog, not pg_depend. -func GetTables(ctx context.Context, pool *pgxpool.Pool, schema string) ([]string, error) { +func GetTables(ctx context.Context, dbtx DBTX, schema string) ([]string, error) { logging.Debug("Querying tables from information_schema", zap.String("schema", schema)) query := ` @@ -97,7 +98,7 @@ func GetTables(ctx context.Context, pool *pgxpool.Pool, schema string) ([]string ORDER BY t.table_name ` - rows, err := pool.Query(ctx, query, schema) + rows, err := dbtx.Query(ctx, query, schema) if err != nil { return nil, fmt.Errorf("failed to query tables: %w", err) } @@ -125,25 +126,29 @@ func GetTables(ctx context.Context, pool *pgxpool.Pool, schema string) ([]string } // GetSchemasWithClient is a convenience wrapper around GetSchemas for Client -func (c *Client) GetSchemas(ctx context.Context) ([]string, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetSchemas(ctx context.Context) (result []string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } - return GetSchemas(ctx, c.pool) + defer func() { done(retErr) }() + return GetSchemas(ctx, q) } // GetTables is a convenience wrapper around GetTables for Client -func (c *Client) GetTables(ctx context.Context, schema string) ([]string, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetTables(ctx context.Context, schema string) (result []string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } - return GetTables(ctx, c.pool, schema) + defer func() { done(retErr) }() + return GetTables(ctx, q, schema) } // GetViews returns all user views in a given schema. // Excludes views owned by extensions (e.g., pg_buffercache, pg_buffercache_numa) // via pg_depend deptype='e'. -func GetViews(ctx context.Context, pool *pgxpool.Pool, schema string) ([]string, error) { +func GetViews(ctx context.Context, dbtx DBTX, schema string) ([]string, error) { logging.Debug("Querying views from information_schema", zap.String("schema", schema)) query := ` @@ -161,7 +166,7 @@ func GetViews(ctx context.Context, pool *pgxpool.Pool, schema string) ([]string, ORDER BY t.table_name ` - rows, err := pool.Query(ctx, query, schema) + rows, err := dbtx.Query(ctx, query, schema) if err != nil { return nil, fmt.Errorf("failed to query views: %w", err) } @@ -189,18 +194,20 @@ func GetViews(ctx context.Context, pool *pgxpool.Pool, schema string) ([]string, } // GetViews is a convenience wrapper around GetViews for Client -func (c *Client) GetViews(ctx context.Context, schema string) ([]string, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetViews(ctx context.Context, schema string) (result []string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } - return GetViews(ctx, c.pool, schema) + defer func() { done(retErr) }() + return GetViews(ctx, q, schema) } // IsViewUpdatable checks if a view supports INSERT/UPDATE/DELETE operations. // PostgreSQL determines this based on the view definition - simple views on // single tables are typically updatable, while views with JOINs, aggregates, // DISTINCT, GROUP BY, etc. are not. -func IsViewUpdatable(ctx context.Context, pool *pgxpool.Pool, schema, view string) (bool, error) { +func IsViewUpdatable(ctx context.Context, dbtx DBTX, schema, view string) (bool, error) { logging.Debug("Checking if view is updatable", zap.String("schema", schema), zap.String("view", view)) @@ -212,7 +219,7 @@ func IsViewUpdatable(ctx context.Context, pool *pgxpool.Pool, schema, view strin ` var isUpdatable string - err := pool.QueryRow(ctx, query, schema, view).Scan(&isUpdatable) + err := dbtx.QueryRow(ctx, query, schema, view).Scan(&isUpdatable) if err != nil { if err == pgx.ErrNoRows { return false, fmt.Errorf("view %s.%s not found", schema, view) @@ -231,11 +238,13 @@ func IsViewUpdatable(ctx context.Context, pool *pgxpool.Pool, schema, view strin } // IsViewUpdatable is a convenience wrapper around IsViewUpdatable for Client -func (c *Client) IsViewUpdatable(ctx context.Context, schema, view string) (bool, error) { - if c.pool == nil { - return false, fmt.Errorf("database connection not initialized") +func (c *Client) IsViewUpdatable(ctx context.Context, schema, view string) (result bool, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return false, err } - return IsViewUpdatable(ctx, c.pool, schema, view) + defer func() { done(retErr) }() + return IsViewUpdatable(ctx, q, schema, view) } // Column represents metadata about a table column @@ -250,7 +259,7 @@ type Column struct { // GetColumns returns all columns for a table in schema order. // Fetches column_default and character_maximum_length so callers // can generate DDL without a separate query. -func GetColumns(ctx context.Context, pool *pgxpool.Pool, schema, table string) ([]Column, error) { +func GetColumns(ctx context.Context, dbtx DBTX, schema, table string) ([]Column, error) { logging.Debug("Querying columns from information_schema", zap.String("schema", schema), zap.String("table", table)) @@ -264,7 +273,7 @@ func GetColumns(ctx context.Context, pool *pgxpool.Pool, schema, table string) ( ORDER BY ordinal_position ` - rows, err := pool.Query(ctx, query, schema, table) + rows, err := dbtx.Query(ctx, query, schema, table) if err != nil { return nil, fmt.Errorf("failed to query columns: %w", err) } @@ -294,15 +303,17 @@ func GetColumns(ctx context.Context, pool *pgxpool.Pool, schema, table string) ( } // GetColumns is a convenience wrapper around GetColumns for Client -func (c *Client) GetColumns(ctx context.Context, schema, table string) ([]Column, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetColumns(ctx context.Context, schema, table string) (result []Column, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } - return GetColumns(ctx, c.pool, schema, table) + defer func() { done(retErr) }() + return GetColumns(ctx, q, schema, table) } // GetRowCount returns the number of rows in a table -func GetRowCount(ctx context.Context, pool *pgxpool.Pool, schema, table string) (int64, error) { +func GetRowCount(ctx context.Context, dbtx DBTX, schema, table string) (int64, error) { logging.Debug("Getting row count", zap.String("schema", schema), zap.String("table", table)) @@ -313,7 +324,7 @@ func GetRowCount(ctx context.Context, pool *pgxpool.Pool, schema, table string) ) var count int64 - err := pool.QueryRow(ctx, query).Scan(&count) + err := dbtx.QueryRow(ctx, query).Scan(&count) if err != nil { return 0, fmt.Errorf("failed to get row count: %w", err) } @@ -327,11 +338,13 @@ func GetRowCount(ctx context.Context, pool *pgxpool.Pool, schema, table string) } // GetRowCount is a convenience wrapper around GetRowCount for Client -func (c *Client) GetRowCount(ctx context.Context, schema, table string) (int64, error) { - if c.pool == nil { - return 0, fmt.Errorf("database connection not initialized") +func (c *Client) GetRowCount(ctx context.Context, schema, table string) (result int64, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return 0, err } - return GetRowCount(ctx, c.pool, schema, table) + defer func() { done(retErr) }() + return GetRowCount(ctx, q, schema, table) } // SmallTableThreshold is the row count below which exact COUNT(*) is used. @@ -348,12 +361,12 @@ const SmallTableThreshold = 100000 // // Parameters: // - ctx: Context for cancellation -// - pool: PostgreSQL connection pool +// - dbtx: DBTX (pool or transaction) // - schema: PostgreSQL schema name // - table: Table name // // Returns the estimated row count, or -1 if the table is not found in pg_class. -func GetTableRowCountEstimate(ctx context.Context, pool *pgxpool.Pool, schema, table string) (int64, error) { +func GetTableRowCountEstimate(ctx context.Context, dbtx DBTX, schema, table string) (int64, error) { logging.Debug("Getting table row count estimate", zap.String("schema", schema), zap.String("table", table)) @@ -366,7 +379,7 @@ func GetTableRowCountEstimate(ctx context.Context, pool *pgxpool.Pool, schema, t ` var estimate int64 - err := pool.QueryRow(ctx, query, schema, table).Scan(&estimate) + err := dbtx.QueryRow(ctx, query, schema, table).Scan(&estimate) if err != nil { if err == pgx.ErrNoRows { logging.Debug("Table not found in pg_class", @@ -386,11 +399,13 @@ func GetTableRowCountEstimate(ctx context.Context, pool *pgxpool.Pool, schema, t } // GetTableRowCountEstimate is a convenience wrapper for Client. -func (c *Client) GetTableRowCountEstimate(ctx context.Context, schema, table string) (int64, error) { - if c.pool == nil { - return 0, fmt.Errorf("database connection not initialized") +func (c *Client) GetTableRowCountEstimate(ctx context.Context, schema, table string) (result int64, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return 0, err } - return GetTableRowCountEstimate(ctx, c.pool, schema, table) + defer func() { done(retErr) }() + return GetTableRowCountEstimate(ctx, q, schema, table) } // GetRowCountSmart returns the row count using an adaptive strategy. @@ -403,18 +418,18 @@ func (c *Client) GetTableRowCountEstimate(ctx context.Context, schema, table str // // Parameters: // - ctx: Context for cancellation -// - pool: PostgreSQL connection pool +// - dbtx: DBTX (pool or transaction) // - schema: PostgreSQL schema name // - table: Table name to count // // Returns the row count (exact or estimated), or error on database failure. -func GetRowCountSmart(ctx context.Context, pool *pgxpool.Pool, schema, table string) (int64, error) { +func GetRowCountSmart(ctx context.Context, dbtx DBTX, schema, table string) (int64, error) { logging.Debug("Getting smart row count", zap.String("schema", schema), zap.String("table", table)) // Get the pg_class.reltuples estimate using shared helper - estimate, err := GetTableRowCountEstimate(ctx, pool, schema, table) + estimate, err := GetTableRowCountEstimate(ctx, dbtx, schema, table) if err != nil { return 0, err } @@ -424,7 +439,7 @@ func GetRowCountSmart(ctx context.Context, pool *pgxpool.Pool, schema, table str logging.Debug("Table not found in pg_class, using exact count", zap.String("schema", schema), zap.String("table", table)) - return GetRowCount(ctx, pool, schema, table) + return GetRowCount(ctx, dbtx, schema, table) } // For small tables, use exact count for accuracy @@ -433,7 +448,7 @@ func GetRowCountSmart(ctx context.Context, pool *pgxpool.Pool, schema, table str zap.String("table", table), zap.Int64("estimate", estimate), zap.Int64("threshold", SmallTableThreshold)) - return GetRowCount(ctx, pool, schema, table) + return GetRowCount(ctx, dbtx, schema, table) } // For large tables, return the estimate to avoid expensive full scan @@ -452,17 +467,19 @@ func GetRowCountSmart(ctx context.Context, pool *pgxpool.Pool, schema, table str // - table: Table name to count // // Returns the row count (exact or estimated), or error on database failure. -func (c *Client) GetRowCountSmart(ctx context.Context, schema, table string) (int64, error) { - if c.pool == nil { - return 0, fmt.Errorf("database connection not initialized") +func (c *Client) GetRowCountSmart(ctx context.Context, schema, table string) (result int64, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return 0, err } - return GetRowCountSmart(ctx, c.pool, schema, table) + defer func() { done(retErr) }() + return GetRowCountSmart(ctx, q, schema, table) } // GetRowCountEstimates returns fast row count estimates for multiple tables using pg_class. // Uses reltuples from PostgreSQL statistics, avoiding full table scans. // Returns a map of table name to estimated row count. -func GetRowCountEstimates(ctx context.Context, pool *pgxpool.Pool, schema string, tables []string) (map[string]int64, error) { +func GetRowCountEstimates(ctx context.Context, dbtx DBTX, schema string, tables []string) (map[string]int64, error) { if len(tables) == 0 { return make(map[string]int64), nil } @@ -478,7 +495,7 @@ func GetRowCountEstimates(ctx context.Context, pool *pgxpool.Pool, schema string WHERE n.nspname = $1 AND c.relname = ANY($2) ` - rows, err := pool.Query(ctx, query, schema, tables) + rows, err := dbtx.Query(ctx, query, schema, tables) if err != nil { return nil, fmt.Errorf("failed to get row count estimates: %w", err) } @@ -506,16 +523,18 @@ func GetRowCountEstimates(ctx context.Context, pool *pgxpool.Pool, schema string } // GetRowCountEstimates is a convenience wrapper for Client. -func (c *Client) GetRowCountEstimates(ctx context.Context, schema string, tables []string) (map[string]int64, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetRowCountEstimates(ctx context.Context, schema string, tables []string) (result map[string]int64, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } - return GetRowCountEstimates(ctx, c.pool, schema, tables) + defer func() { done(retErr) }() + return GetRowCountEstimates(ctx, q, schema, tables) } // GetTableDDL returns the CREATE TABLE statement for a table // Constructs DDL from information_schema -func GetTableDDL(ctx context.Context, pool *pgxpool.Pool, schema, table string) (string, error) { +func GetTableDDL(ctx context.Context, dbtx DBTX, schema, table string) (string, error) { logging.Debug("Getting table DDL", zap.String("schema", schema), zap.String("table", table)) @@ -533,7 +552,7 @@ func GetTableDDL(ctx context.Context, pool *pgxpool.Pool, schema, table string) ORDER BY ordinal_position ` - rows, err := pool.Query(ctx, query, schema, table) + rows, err := dbtx.Query(ctx, query, schema, table) if err != nil { return "", fmt.Errorf("failed to query column definitions: %w", err) } @@ -592,7 +611,7 @@ func GetTableDDL(ctx context.Context, pool *pgxpool.Pool, schema, table string) ` var pkColumns *string - if err := pool.QueryRow(ctx, pkQuery, schema, table).Scan(&pkColumns); err != nil && err != pgx.ErrNoRows { + if err := dbtx.QueryRow(ctx, pkQuery, schema, table).Scan(&pkColumns); err != nil && err != pgx.ErrNoRows { return "", fmt.Errorf("failed to query primary key: %w", err) } @@ -610,11 +629,13 @@ func GetTableDDL(ctx context.Context, pool *pgxpool.Pool, schema, table string) } // GetTableDDL is a convenience wrapper around GetTableDDL for Client -func (c *Client) GetTableDDL(ctx context.Context, schema, table string) (string, error) { - if c.pool == nil { - return "", fmt.Errorf("database connection not initialized") +func (c *Client) GetTableDDL(ctx context.Context, schema, table string) (result string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return "", err } - return GetTableDDL(ctx, c.pool, schema, table) + defer func() { done(retErr) }() + return GetTableDDL(ctx, q, schema, table) } // FormatTableDDL generates a CREATE TABLE DDL statement from pre-fetched metadata. @@ -668,7 +689,7 @@ func FormatTableDDL(schema, table string, columns []Column, pk *PrimaryKey) stri // GetIndexDDL returns CREATE INDEX statements for a table. // Includes both regular and unique indexes, excluding primary key constraints // (which are already shown in the CREATE TABLE statement). -func GetIndexDDL(ctx context.Context, pool *pgxpool.Pool, schema, table string) (string, error) { +func GetIndexDDL(ctx context.Context, dbtx DBTX, schema, table string) (string, error) { logging.Debug("Getting index DDL", zap.String("schema", schema), zap.String("table", table)) @@ -687,7 +708,7 @@ func GetIndexDDL(ctx context.Context, pool *pgxpool.Pool, schema, table string) ORDER BY indexname ` - rows, err := pool.Query(ctx, query, schema, table) + rows, err := dbtx.Query(ctx, query, schema, table) if err != nil { return "", fmt.Errorf("failed to query indexes: %w", err) } @@ -711,15 +732,17 @@ func GetIndexDDL(ctx context.Context, pool *pgxpool.Pool, schema, table string) } // GetIndexDDL is a convenience wrapper for Client -func (c *Client) GetIndexDDL(ctx context.Context, schema, table string) (string, error) { - if c.pool == nil { - return "", fmt.Errorf("database connection not initialized") +func (c *Client) GetIndexDDL(ctx context.Context, schema, table string) (result string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return "", err } - return GetIndexDDL(ctx, c.pool, schema, table) + defer func() { done(retErr) }() + return GetIndexDDL(ctx, q, schema, table) } // GetForeignKeyDDL returns ALTER TABLE statements for foreign key constraints. -func GetForeignKeyDDL(ctx context.Context, pool *pgxpool.Pool, schema, table string) (string, error) { +func GetForeignKeyDDL(ctx context.Context, dbtx DBTX, schema, table string) (string, error) { logging.Debug("Getting foreign key DDL", zap.String("schema", schema), zap.String("table", table)) @@ -737,7 +760,7 @@ func GetForeignKeyDDL(ctx context.Context, pool *pgxpool.Pool, schema, table str ORDER BY con.conname ` - rows, err := pool.Query(ctx, query, schema, table) + rows, err := dbtx.Query(ctx, query, schema, table) if err != nil { return "", fmt.Errorf("failed to query foreign keys: %w", err) } @@ -761,15 +784,17 @@ func GetForeignKeyDDL(ctx context.Context, pool *pgxpool.Pool, schema, table str } // GetForeignKeyDDL is a convenience wrapper for Client -func (c *Client) GetForeignKeyDDL(ctx context.Context, schema, table string) (string, error) { - if c.pool == nil { - return "", fmt.Errorf("database connection not initialized") +func (c *Client) GetForeignKeyDDL(ctx context.Context, schema, table string) (result string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return "", err } - return GetForeignKeyDDL(ctx, c.pool, schema, table) + defer func() { done(retErr) }() + return GetForeignKeyDDL(ctx, q, schema, table) } // GetCheckConstraintDDL returns ALTER TABLE statements for check constraints. -func GetCheckConstraintDDL(ctx context.Context, pool *pgxpool.Pool, schema, table string) (string, error) { +func GetCheckConstraintDDL(ctx context.Context, dbtx DBTX, schema, table string) (string, error) { logging.Debug("Getting check constraint DDL", zap.String("schema", schema), zap.String("table", table)) @@ -787,7 +812,7 @@ func GetCheckConstraintDDL(ctx context.Context, pool *pgxpool.Pool, schema, tabl ORDER BY con.conname ` - rows, err := pool.Query(ctx, query, schema, table) + rows, err := dbtx.Query(ctx, query, schema, table) if err != nil { return "", fmt.Errorf("failed to query check constraints: %w", err) } @@ -811,15 +836,17 @@ func GetCheckConstraintDDL(ctx context.Context, pool *pgxpool.Pool, schema, tabl } // GetCheckConstraintDDL is a convenience wrapper for Client -func (c *Client) GetCheckConstraintDDL(ctx context.Context, schema, table string) (string, error) { - if c.pool == nil { - return "", fmt.Errorf("database connection not initialized") +func (c *Client) GetCheckConstraintDDL(ctx context.Context, schema, table string) (result string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return "", err } - return GetCheckConstraintDDL(ctx, c.pool, schema, table) + defer func() { done(retErr) }() + return GetCheckConstraintDDL(ctx, q, schema, table) } // GetTriggerDDL returns CREATE TRIGGER statements for a table. -func GetTriggerDDL(ctx context.Context, pool *pgxpool.Pool, schema, table string) (string, error) { +func GetTriggerDDL(ctx context.Context, dbtx DBTX, schema, table string) (string, error) { logging.Debug("Getting trigger DDL", zap.String("schema", schema), zap.String("table", table)) @@ -835,7 +862,7 @@ func GetTriggerDDL(ctx context.Context, pool *pgxpool.Pool, schema, table string ORDER BY t.tgname ` - rows, err := pool.Query(ctx, query, schema, table) + rows, err := dbtx.Query(ctx, query, schema, table) if err != nil { return "", fmt.Errorf("failed to query triggers: %w", err) } @@ -859,15 +886,17 @@ func GetTriggerDDL(ctx context.Context, pool *pgxpool.Pool, schema, table string } // GetTriggerDDL is a convenience wrapper for Client -func (c *Client) GetTriggerDDL(ctx context.Context, schema, table string) (string, error) { - if c.pool == nil { - return "", fmt.Errorf("database connection not initialized") +func (c *Client) GetTriggerDDL(ctx context.Context, schema, table string) (result string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return "", err } - return GetTriggerDDL(ctx, c.pool, schema, table) + defer func() { done(retErr) }() + return GetTriggerDDL(ctx, q, schema, table) } // GetTableComments returns COMMENT statements for table and columns. -func GetTableComments(ctx context.Context, pool *pgxpool.Pool, schema, table string) (string, error) { +func GetTableComments(ctx context.Context, dbtx DBTX, schema, table string) (string, error) { logging.Debug("Getting table comments", zap.String("schema", schema), zap.String("table", table)) @@ -879,7 +908,7 @@ func GetTableComments(ctx context.Context, pool *pgxpool.Pool, schema, table str SELECT obj_description((quote_ident($1) || '.' || quote_ident($2))::regclass, 'pg_class') ` var tableComment *string - err := pool.QueryRow(ctx, tableCommentQuery, schema, table).Scan(&tableComment) + err := dbtx.QueryRow(ctx, tableCommentQuery, schema, table).Scan(&tableComment) if err != nil { return "", fmt.Errorf("failed to query table comment: %w", err) } @@ -903,7 +932,7 @@ func GetTableComments(ctx context.Context, pool *pgxpool.Pool, schema, table str ORDER BY a.attnum ` - rows, err := pool.Query(ctx, columnCommentQuery, schema, table) + rows, err := dbtx.Query(ctx, columnCommentQuery, schema, table) if err != nil { return "", fmt.Errorf("failed to query column comments: %w", err) } @@ -926,11 +955,13 @@ func GetTableComments(ctx context.Context, pool *pgxpool.Pool, schema, table str } // GetTableComments is a convenience wrapper for Client -func (c *Client) GetTableComments(ctx context.Context, schema, table string) (string, error) { - if c.pool == nil { - return "", fmt.Errorf("database connection not initialized") +func (c *Client) GetTableComments(ctx context.Context, schema, table string) (result string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return "", err } - return GetTableComments(ctx, c.pool, schema, table) + defer func() { done(retErr) }() + return GetTableComments(ctx, q, schema, table) } // GetFullDDL returns complete DDL for a table including: @@ -941,7 +972,7 @@ func (c *Client) GetTableComments(ctx context.Context, schema, table string) (st // - Triggers // - Comments // Sections are only included if they have content. -func GetFullDDL(ctx context.Context, pool *pgxpool.Pool, schema, table string) (string, error) { +func GetFullDDL(ctx context.Context, dbtx DBTX, schema, table string) (string, error) { logging.Debug("Getting full DDL", zap.String("schema", schema), zap.String("table", table)) @@ -949,7 +980,7 @@ func GetFullDDL(ctx context.Context, pool *pgxpool.Pool, schema, table string) ( var ddl strings.Builder // Table definition - tableDDL, err := GetTableDDL(ctx, pool, schema, table) + tableDDL, err := GetTableDDL(ctx, dbtx, schema, table) if err != nil { return "", fmt.Errorf("failed to get table DDL: %w", err) } @@ -958,7 +989,7 @@ func GetFullDDL(ctx context.Context, pool *pgxpool.Pool, schema, table string) ( ddl.WriteString("\n") // Indexes - indexDDL, err := GetIndexDDL(ctx, pool, schema, table) + indexDDL, err := GetIndexDDL(ctx, dbtx, schema, table) if err != nil { return "", fmt.Errorf("failed to get index DDL: %w", err) } @@ -969,7 +1000,7 @@ func GetFullDDL(ctx context.Context, pool *pgxpool.Pool, schema, table string) ( } // Foreign keys - fkDDL, err := GetForeignKeyDDL(ctx, pool, schema, table) + fkDDL, err := GetForeignKeyDDL(ctx, dbtx, schema, table) if err != nil { return "", fmt.Errorf("failed to get foreign key DDL: %w", err) } @@ -980,7 +1011,7 @@ func GetFullDDL(ctx context.Context, pool *pgxpool.Pool, schema, table string) ( } // Check constraints - checkDDL, err := GetCheckConstraintDDL(ctx, pool, schema, table) + checkDDL, err := GetCheckConstraintDDL(ctx, dbtx, schema, table) if err != nil { return "", fmt.Errorf("failed to get check constraint DDL: %w", err) } @@ -991,7 +1022,7 @@ func GetFullDDL(ctx context.Context, pool *pgxpool.Pool, schema, table string) ( } // Triggers - triggerDDL, err := GetTriggerDDL(ctx, pool, schema, table) + triggerDDL, err := GetTriggerDDL(ctx, dbtx, schema, table) if err != nil { return "", fmt.Errorf("failed to get trigger DDL: %w", err) } @@ -1002,7 +1033,7 @@ func GetFullDDL(ctx context.Context, pool *pgxpool.Pool, schema, table string) ( } // Comments - commentsDDL, err := GetTableComments(ctx, pool, schema, table) + commentsDDL, err := GetTableComments(ctx, dbtx, schema, table) if err != nil { return "", fmt.Errorf("failed to get table comments: %w", err) } @@ -1015,11 +1046,13 @@ func GetFullDDL(ctx context.Context, pool *pgxpool.Pool, schema, table string) ( } // GetFullDDL is a convenience wrapper for Client -func (c *Client) GetFullDDL(ctx context.Context, schema, table string) (string, error) { - if c.pool == nil { - return "", fmt.Errorf("database connection not initialized") +func (c *Client) GetFullDDL(ctx context.Context, schema, table string) (result string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return "", err } - return GetFullDDL(ctx, c.pool, schema, table) + defer func() { done(retErr) }() + return GetFullDDL(ctx, q, schema, table) } // ForeignKeyRef represents a foreign key that references a table. @@ -1041,7 +1074,7 @@ type ForeignKeyRef struct { // GetReferencingForeignKeys returns foreign keys that reference the specified table. // This is the reverse of GetForeignKeyDDL - it finds tables that depend on this table. -func GetReferencingForeignKeys(ctx context.Context, pool *pgxpool.Pool, schema, table string) ([]ForeignKeyRef, error) { +func GetReferencingForeignKeys(ctx context.Context, dbtx DBTX, schema, table string) ([]ForeignKeyRef, error) { logging.Debug("Getting referencing foreign keys", zap.String("schema", schema), zap.String("table", table)) @@ -1070,7 +1103,7 @@ func GetReferencingForeignKeys(ctx context.Context, pool *pgxpool.Pool, schema, ORDER BY src_nsp.nspname, src_cls.relname, con.conname ` - rows, err := pool.Query(ctx, query, schema, table) + rows, err := dbtx.Query(ctx, query, schema, table) if err != nil { return nil, fmt.Errorf("failed to query referencing foreign keys: %w", err) } @@ -1105,16 +1138,18 @@ func GetReferencingForeignKeys(ctx context.Context, pool *pgxpool.Pool, schema, } // GetReferencingForeignKeys is a convenience wrapper for Client -func (c *Client) GetReferencingForeignKeys(ctx context.Context, schema, table string) ([]ForeignKeyRef, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetReferencingForeignKeys(ctx context.Context, schema, table string) (result []ForeignKeyRef, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } - return GetReferencingForeignKeys(ctx, c.pool, schema, table) + defer func() { done(retErr) }() + return GetReferencingForeignKeys(ctx, q, schema, table) } // GetSchemaTableCount returns the number of tables in a schema. // Used to determine if a schema is empty when generating delete templates. -func GetSchemaTableCount(ctx context.Context, pool *pgxpool.Pool, schema string) (int, error) { +func GetSchemaTableCount(ctx context.Context, dbtx DBTX, schema string) (int, error) { logging.Debug("Getting schema table count", zap.String("schema", schema)) @@ -1127,7 +1162,7 @@ func GetSchemaTableCount(ctx context.Context, pool *pgxpool.Pool, schema string) ` var count int - err := pool.QueryRow(ctx, query, schema).Scan(&count) + err := dbtx.QueryRow(ctx, query, schema).Scan(&count) if err != nil { return 0, fmt.Errorf("failed to get schema table count: %w", err) } @@ -1136,17 +1171,19 @@ func GetSchemaTableCount(ctx context.Context, pool *pgxpool.Pool, schema string) } // GetSchemaTableCount is a convenience wrapper for Client -func (c *Client) GetSchemaTableCount(ctx context.Context, schema string) (int, error) { - if c.pool == nil { - return 0, fmt.Errorf("database connection not initialized") +func (c *Client) GetSchemaTableCount(ctx context.Context, schema string) (result int, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return 0, err } - return GetSchemaTableCount(ctx, c.pool, schema) + defer func() { done(retErr) }() + return GetSchemaTableCount(ctx, q, schema) } // GetViewComment returns the raw comment string for a view. // Returns empty string if the view has no comment. // Used to detect synthesized format markers (e.g., "tigerfs:md"). -func GetViewComment(ctx context.Context, pool *pgxpool.Pool, schema, view string) (string, error) { +func GetViewComment(ctx context.Context, dbtx DBTX, schema, view string) (string, error) { logging.Debug("Getting view comment", zap.String("schema", schema), zap.String("view", view)) @@ -1161,7 +1198,7 @@ func GetViewComment(ctx context.Context, pool *pgxpool.Pool, schema, view string ` var comment string - err := pool.QueryRow(ctx, query, schema, view).Scan(&comment) + err := dbtx.QueryRow(ctx, query, schema, view).Scan(&comment) if err != nil { if err == pgx.ErrNoRows { return "", nil @@ -1173,16 +1210,18 @@ func GetViewComment(ctx context.Context, pool *pgxpool.Pool, schema, view string } // GetViewComment is a convenience wrapper for Client. -func (c *Client) GetViewComment(ctx context.Context, schema, view string) (string, error) { - if c.pool == nil { - return "", fmt.Errorf("database connection not initialized") +func (c *Client) GetViewComment(ctx context.Context, schema, view string) (result string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return "", err } - return GetViewComment(ctx, c.pool, schema, view) + defer func() { done(retErr) }() + return GetViewComment(ctx, q, schema, view) } // GetViewCommentsBatch returns comments for all views in a schema. // Returns a map of view name to comment string. Views without comments are omitted. -func GetViewCommentsBatch(ctx context.Context, pool *pgxpool.Pool, schema string) (map[string]string, error) { +func GetViewCommentsBatch(ctx context.Context, dbtx DBTX, schema string) (map[string]string, error) { logging.Debug("Getting view comments batch", zap.String("schema", schema)) @@ -1195,7 +1234,7 @@ func GetViewCommentsBatch(ctx context.Context, pool *pgxpool.Pool, schema string AND obj_description(c.oid, 'pg_class') IS NOT NULL ` - rows, err := pool.Query(ctx, query, schema) + rows, err := dbtx.Query(ctx, query, schema) if err != nil { return nil, fmt.Errorf("failed to query view comments: %w", err) } @@ -1224,15 +1263,17 @@ func GetViewCommentsBatch(ctx context.Context, pool *pgxpool.Pool, schema string } // GetViewCommentsBatch is a convenience wrapper for Client. -func (c *Client) GetViewCommentsBatch(ctx context.Context, schema string) (map[string]string, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetViewCommentsBatch(ctx context.Context, schema string) (result map[string]string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } - return GetViewCommentsBatch(ctx, c.pool, schema) + defer func() { done(retErr) }() + return GetViewCommentsBatch(ctx, q, schema) } // GetViewDefinition returns the SQL definition of a view. -func GetViewDefinition(ctx context.Context, pool *pgxpool.Pool, schema, view string) (string, error) { +func GetViewDefinition(ctx context.Context, dbtx DBTX, schema, view string) (string, error) { logging.Debug("Getting view definition", zap.String("schema", schema), zap.String("view", view)) @@ -1247,7 +1288,7 @@ func GetViewDefinition(ctx context.Context, pool *pgxpool.Pool, schema, view str ` var definition string - err := pool.QueryRow(ctx, query, schema, view).Scan(&definition) + err := dbtx.QueryRow(ctx, query, schema, view).Scan(&definition) if err != nil { return "", fmt.Errorf("failed to get view definition: %w", err) } @@ -1256,15 +1297,17 @@ func GetViewDefinition(ctx context.Context, pool *pgxpool.Pool, schema, view str } // GetViewDefinition is a convenience wrapper for Client -func (c *Client) GetViewDefinition(ctx context.Context, schema, view string) (string, error) { - if c.pool == nil { - return "", fmt.Errorf("database connection not initialized") +func (c *Client) GetViewDefinition(ctx context.Context, schema, view string) (result string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return "", err } - return GetViewDefinition(ctx, c.pool, schema, view) + defer func() { done(retErr) }() + return GetViewDefinition(ctx, q, schema, view) } // GetDependentViews returns views that depend on the specified view or table. -func GetDependentViews(ctx context.Context, pool *pgxpool.Pool, schema, name string) ([]string, error) { +func GetDependentViews(ctx context.Context, dbtx DBTX, schema, name string) ([]string, error) { logging.Debug("Getting dependent views", zap.String("schema", schema), zap.String("name", name)) @@ -1282,7 +1325,7 @@ func GetDependentViews(ctx context.Context, pool *pgxpool.Pool, schema, name str ORDER BY dep_cls.relname ` - rows, err := pool.Query(ctx, query, schema, name) + rows, err := dbtx.Query(ctx, query, schema, name) if err != nil { return nil, fmt.Errorf("failed to query dependent views: %w", err) } @@ -1305,9 +1348,11 @@ func GetDependentViews(ctx context.Context, pool *pgxpool.Pool, schema, name str } // GetDependentViews is a convenience wrapper for Client -func (c *Client) GetDependentViews(ctx context.Context, schema, name string) ([]string, error) { - if c.pool == nil { - return nil, fmt.Errorf("database connection not initialized") +func (c *Client) GetDependentViews(ctx context.Context, schema, name string) (result []string, retErr error) { + q, done, err := c.acquireDBTX(ctx) + if err != nil { + return nil, err } - return GetDependentViews(ctx, c.pool, schema, name) + defer func() { done(retErr) }() + return GetDependentViews(ctx, q, schema, name) } diff --git a/internal/tigerfs/db/session_vars.go b/internal/tigerfs/db/session_vars.go new file mode 100644 index 0000000..b4f02fa --- /dev/null +++ b/internal/tigerfs/db/session_vars.go @@ -0,0 +1,108 @@ +package db + +import ( + "context" + "sort" +) + +// SessionVars holds PostgreSQL session variable names and values with +// pre-sorted keys for deterministic, allocation-free iteration at query time. +// +// Variable names must follow PostgreSQL custom GUC naming conventions +// (dotted identifiers, e.g. "app.user_id", "app.tenant_id"). +// +// Create via NewSessionVars: +// +// vars := db.NewSessionVars(map[string]string{ +// "app.user_id": "42", +// "app.tenant_id": "acme", +// }) +type SessionVars struct { + keys []string // sorted once at construction + values map[string]string // original map for value lookup +} + +// NewSessionVars creates a SessionVars with pre-sorted keys. +// Returns a zero-value SessionVars if the map is nil or empty. +func NewSessionVars(m map[string]string) SessionVars { + if len(m) == 0 { + return SessionVars{} + } + keys := make([]string, 0, len(m)) + for k := range m { + keys = append(keys, k) + } + sort.Strings(keys) + + // Copy the map so the caller can't mutate it after construction. + values := make(map[string]string, len(m)) + for k, v := range m { + values[k] = v + } + return SessionVars{keys: keys, values: values} +} + +// Len returns the number of session variables. +func (sv SessionVars) Len() int { + return len(sv.keys) +} + +// Empty returns true if there are no session variables. +func (sv SessionVars) Empty() bool { + return len(sv.keys) == 0 +} + +// Range calls fn for each key-value pair in sorted key order. +func (sv SessionVars) Range(fn func(key, value string)) { + for _, k := range sv.keys { + fn(k, sv.values[k]) + } +} + +// Merge returns a new SessionVars combining base and override. +// Override values take precedence for duplicate keys. Both inputs +// are unchanged. The result has freshly sorted keys. +func (sv SessionVars) Merge(override SessionVars) SessionVars { + if sv.Empty() { + return override + } + if override.Empty() { + return sv + } + merged := make(map[string]string, len(sv.values)+len(override.values)) + for k, v := range sv.values { + merged[k] = v + } + for k, v := range override.values { + merged[k] = v + } + return NewSessionVars(merged) +} + +type sessionVarsKey struct{} + +// WithSessionVars returns a derived context carrying session variables +// that will be applied via SET LOCAL to any query executed under this +// context. +// +// Calling with an empty SessionVars returns ctx unchanged. +// +// Example: +// +// ctx = db.WithSessionVars(ctx, db.NewSessionVars(map[string]string{ +// "app.user_id": userID, +// })) +// row, err := client.GetRow(ctx, schema, table, pk) // SET LOCAL applied +func WithSessionVars(ctx context.Context, vars SessionVars) context.Context { + if vars.Empty() { + return ctx + } + return context.WithValue(ctx, sessionVarsKey{}, vars) +} + +// SessionVarsFromContext returns the session variables attached to ctx +// via WithSessionVars, or a zero-value SessionVars if none were set. +func SessionVarsFromContext(ctx context.Context) SessionVars { + v, _ := ctx.Value(sessionVarsKey{}).(SessionVars) + return v +} diff --git a/internal/tigerfs/db/session_vars_test.go b/internal/tigerfs/db/session_vars_test.go new file mode 100644 index 0000000..b90fa43 --- /dev/null +++ b/internal/tigerfs/db/session_vars_test.go @@ -0,0 +1,173 @@ +package db + +import ( + "context" + "fmt" + "testing" +) + +func TestWithSessionVars_ZeroValueReturnsOriginal(t *testing.T) { + ctx := context.Background() + got := WithSessionVars(ctx, SessionVars{}) + if got != ctx { + t.Error("expected same context for zero-value vars") + } +} + +func TestWithSessionVars_EmptyMapReturnsOriginal(t *testing.T) { + ctx := context.Background() + got := WithSessionVars(ctx, NewSessionVars(nil)) + if got != ctx { + t.Error("expected same context for nil map") + } +} + +func TestSessionVarsFromContext_MissingReturnsEmpty(t *testing.T) { + ctx := context.Background() + vars := SessionVarsFromContext(ctx) + if !vars.Empty() { + t.Errorf("expected empty, got %d vars", vars.Len()) + } +} + +func TestSessionVarsFromContext_RoundTrip(t *testing.T) { + ctx := context.Background() + vars := NewSessionVars(map[string]string{"app.user_id": "42", "app.tenant_id": "acme"}) + ctx = WithSessionVars(ctx, vars) + + got := SessionVarsFromContext(ctx) + if got.Empty() { + t.Fatal("expected non-empty vars") + } + assertVar(t, got, "app.user_id", "42") + assertVar(t, got, "app.tenant_id", "acme") +} + +func TestSessionVarsFromContext_OverwritesPrevious(t *testing.T) { + ctx := context.Background() + ctx = WithSessionVars(ctx, NewSessionVars(map[string]string{"app.user_id": "1"})) + ctx = WithSessionVars(ctx, NewSessionVars(map[string]string{"app.user_id": "2"})) + + got := SessionVarsFromContext(ctx) + assertVar(t, got, "app.user_id", "2") +} + +func TestNewSessionVars_SortsKeys(t *testing.T) { + vars := NewSessionVars(map[string]string{ + "z.last": "1", "a.first": "2", "m.middle": "3", + }) + var keys []string + vars.Range(func(k, v string) { keys = append(keys, k) }) + if keys[0] != "a.first" || keys[1] != "m.middle" || keys[2] != "z.last" { + t.Errorf("keys not sorted: %v", keys) + } +} + +func TestNewSessionVars_CopiesMap(t *testing.T) { + m := map[string]string{"app.user_id": "1"} + vars := NewSessionVars(m) + m["app.user_id"] = "mutated" + assertVar(t, vars, "app.user_id", "1") // should not see mutation +} + +func TestSessionVars_Merge(t *testing.T) { + base := NewSessionVars(map[string]string{"app.tenant_id": "acme", "app.user_id": "1"}) + override := NewSessionVars(map[string]string{"app.user_id": "2"}) + merged := base.Merge(override) + + if merged.Len() != 2 { + t.Errorf("expected 2 vars, got %d", merged.Len()) + } + assertVar(t, merged, "app.tenant_id", "acme") + assertVar(t, merged, "app.user_id", "2") +} + +func TestSessionVars_MergeEmptyBase(t *testing.T) { + override := NewSessionVars(map[string]string{"app.user_id": "1"}) + merged := SessionVars{}.Merge(override) + assertVar(t, merged, "app.user_id", "1") +} + +func TestSessionVars_MergeEmptyOverride(t *testing.T) { + base := NewSessionVars(map[string]string{"app.user_id": "1"}) + merged := base.Merge(SessionVars{}) + assertVar(t, merged, "app.user_id", "1") +} + +func TestEffectiveSessionVars_NoVars(t *testing.T) { + c := &Client{} + ctx := context.Background() + if vars := c.effectiveSessionVars(ctx); !vars.Empty() { + t.Errorf("expected empty, got %d vars", vars.Len()) + } +} + +func TestEffectiveSessionVars_BaselineOnly(t *testing.T) { + c := &Client{baselineVars: NewSessionVars(map[string]string{"app.tenant_id": "acme"})} + ctx := context.Background() + vars := c.effectiveSessionVars(ctx) + assertVar(t, vars, "app.tenant_id", "acme") +} + +func TestEffectiveSessionVars_ContextOnly(t *testing.T) { + c := &Client{} + ctx := WithSessionVars(context.Background(), NewSessionVars(map[string]string{"app.user_id": "42"})) + vars := c.effectiveSessionVars(ctx) + assertVar(t, vars, "app.user_id", "42") +} + +func TestEffectiveSessionVars_ContextOverridesBaseline(t *testing.T) { + c := &Client{baselineVars: NewSessionVars(map[string]string{ + "app.user_id": "1", + "app.tenant_id": "acme", + })} + ctx := WithSessionVars(context.Background(), NewSessionVars(map[string]string{"app.user_id": "2"})) + vars := c.effectiveSessionVars(ctx) + + assertVar(t, vars, "app.user_id", "2") + assertVar(t, vars, "app.tenant_id", "acme") +} + +func TestEffectiveSessionVars_Merge(t *testing.T) { + c := &Client{baselineVars: NewSessionVars(map[string]string{"app.tenant_id": "acme"})} + ctx := WithSessionVars(context.Background(), NewSessionVars(map[string]string{"app.user_id": "42"})) + vars := c.effectiveSessionVars(ctx) + + if vars.Len() != 2 { + t.Errorf("expected 2 vars, got %d", vars.Len()) + } + assertVar(t, vars, "app.tenant_id", "acme") + assertVar(t, vars, "app.user_id", "42") +} + +func TestAcquireDBTX_NilPoolReturnsError(t *testing.T) { + c := &Client{} + ctx := context.Background() + _, _, err := c.acquireDBTX(ctx) + if err == nil { + t.Fatal("expected error for nil pool") + } +} + +func TestAcquireDBTX_NoVarsDoneFuncIsNoOp(t *testing.T) { + noop := doneFunc(func(error) {}) + noop(nil) + noop(fmt.Errorf("test error")) +} + +// assertVar checks that a SessionVars contains the expected key-value pair. +func assertVar(t *testing.T, vars SessionVars, key, want string) { + t.Helper() + var found bool + vars.Range(func(k, v string) { + if k == key { + found = true + if v != want { + t.Errorf("%s = %q, want %q", key, v, want) + } + } + }) + if !found { + t.Errorf("key %q not found in session vars", key) + } +} diff --git a/test/integration/session_vars_test.go b/test/integration/session_vars_test.go new file mode 100644 index 0000000..39551e9 --- /dev/null +++ b/test/integration/session_vars_test.go @@ -0,0 +1,650 @@ +package integration + +import ( + "context" + "fmt" + "sync" + "testing" + "time" + + "github.com/jackc/pgx/v5/pgxpool" + "github.com/timescale/tigerfs/internal/tigerfs/config" + "github.com/timescale/tigerfs/internal/tigerfs/db" +) + +// TestSessionVars_SetLocalApplied verifies that session variables are +// applied via SET LOCAL within a transaction and visible to queries. +func TestSessionVars_SetLocalApplied(t *testing.T) { + result := GetTestDB(t) + if result == nil { + return + } + t.Cleanup(result.Cleanup) + + cfg := &config.Config{ + PoolSize: 2, + PoolMaxIdle: 1, + SessionVariables: map[string]string{ + "app.user_id": "42", + }, + } + + client, err := db.NewClient(context.Background(), cfg, result.ConnStr) + if err != nil { + t.Fatalf("NewClient: %v", err) + } + defer client.Close() + + // Query current_setting to verify the session var was applied + ctx := context.Background() + row, err := client.Pool().Query(ctx, "SELECT 1") // warm up + if err != nil { + t.Fatalf("warmup: %v", err) + } + row.Close() + + // Use the Client's Exec (which goes through acquireDBTX) to test + // that session vars are applied. We create a temp table, insert via + // a query that reads current_setting, then verify. + if err := client.Exec(ctx, `CREATE TEMPORARY TABLE _sv_test (val text)`); err != nil { + t.Fatalf("create temp table: %v", err) + } + + if err := client.Exec(ctx, `INSERT INTO _sv_test (val) VALUES (current_setting('app.user_id', true))`); err != nil { + t.Fatalf("insert with current_setting: %v", err) + } + + // Read back via pool directly (no session vars) to verify it was stored + var val string + err = client.Pool().QueryRow(ctx, `SELECT val FROM _sv_test LIMIT 1`).Scan(&val) + if err != nil { + t.Fatalf("read back: %v", err) + } + if val != "42" { + t.Errorf("current_setting('app.user_id') = %q, want %q", val, "42") + } +} + +// TestSessionVars_SetLocalDoesNotLeak verifies that SET LOCAL variables +// do not persist after the transaction commits — they are truly scoped. +func TestSessionVars_SetLocalDoesNotLeak(t *testing.T) { + result := GetTestDB(t) + if result == nil { + return + } + t.Cleanup(result.Cleanup) + + cfg := &config.Config{ + PoolSize: 1, // Force single connection to verify same conn is clean + PoolMaxIdle: 1, + } + + client, err := db.NewClient(context.Background(), cfg, result.ConnStr) + if err != nil { + t.Fatalf("NewClient: %v", err) + } + defer client.Close() + + // First query: with session vars via context + ctx := db.WithSessionVars(context.Background(), db.NewSessionVars(map[string]string{ + "app.user_id": "99", + })) + + if err := client.Exec(ctx, `CREATE TEMPORARY TABLE _sv_leak (val text)`); err != nil { + t.Fatalf("create temp table: %v", err) + } + + if err := client.Exec(ctx, `INSERT INTO _sv_leak (val) VALUES (current_setting('app.user_id', true))`); err != nil { + t.Fatalf("insert with session var: %v", err) + } + + // Second query: WITHOUT session vars — same connection (pool_size=1) + ctxClean := context.Background() + if err := client.Exec(ctxClean, `INSERT INTO _sv_leak (val) VALUES (current_setting('app.user_id', true))`); err != nil { + t.Fatalf("insert without session var: %v", err) + } + + // Read both values + rows, err := client.Pool().Query(ctxClean, `SELECT val FROM _sv_leak ORDER BY ctid`) + if err != nil { + t.Fatalf("query: %v", err) + } + defer rows.Close() + + var vals []string + for rows.Next() { + var v string + if err := rows.Scan(&v); err != nil { + t.Fatalf("scan: %v", err) + } + vals = append(vals, v) + } + + if len(vals) != 2 { + t.Fatalf("expected 2 rows, got %d", len(vals)) + } + if vals[0] != "99" { + t.Errorf("row 1 (with session var) = %q, want %q", vals[0], "99") + } + if vals[1] != "" { + t.Errorf("row 2 (without session var) = %q, want empty (SET LOCAL should not leak)", vals[1]) + } +} + +// TestSessionVars_ContextOverridesBaseline verifies that context-level +// session vars override baseline (config) vars. +func TestSessionVars_ContextOverridesBaseline(t *testing.T) { + result := GetTestDB(t) + if result == nil { + return + } + t.Cleanup(result.Cleanup) + + cfg := &config.Config{ + PoolSize: 2, + PoolMaxIdle: 1, + SessionVariables: map[string]string{ + "app.user_id": "baseline_user", + }, + } + + client, err := db.NewClient(context.Background(), cfg, result.ConnStr) + if err != nil { + t.Fatalf("NewClient: %v", err) + } + defer client.Close() + + ctx := context.Background() + + if err := client.Exec(ctx, `CREATE TEMPORARY TABLE _sv_override (val text)`); err != nil { + t.Fatalf("create temp table: %v", err) + } + + // Insert with baseline vars (no context override) + if err := client.Exec(ctx, `INSERT INTO _sv_override (val) VALUES (current_setting('app.user_id', true))`); err != nil { + t.Fatalf("insert baseline: %v", err) + } + + // Insert with context override + ctxOverride := db.WithSessionVars(ctx, db.NewSessionVars(map[string]string{"app.user_id": "context_user"})) + if err := client.Exec(ctxOverride, `INSERT INTO _sv_override (val) VALUES (current_setting('app.user_id', true))`); err != nil { + t.Fatalf("insert override: %v", err) + } + + // Read both + rows, err := client.Pool().Query(ctx, `SELECT val FROM _sv_override ORDER BY ctid`) + if err != nil { + t.Fatalf("query: %v", err) + } + defer rows.Close() + + var vals []string + for rows.Next() { + var v string + if err := rows.Scan(&v); err != nil { + t.Fatalf("scan: %v", err) + } + vals = append(vals, v) + } + + if len(vals) != 2 { + t.Fatalf("expected 2 rows, got %d", len(vals)) + } + if vals[0] != "baseline_user" { + t.Errorf("baseline row = %q, want %q", vals[0], "baseline_user") + } + if vals[1] != "context_user" { + t.Errorf("override row = %q, want %q", vals[1], "context_user") + } +} + +// TestSessionVars_NoVarsNoOverhead verifies that with no session vars +// configured, queries execute without wrapping in a transaction. +func TestSessionVars_NoVarsNoOverhead(t *testing.T) { + result := GetTestDB(t) + if result == nil { + return + } + t.Cleanup(result.Cleanup) + + cfg := &config.Config{ + PoolSize: 2, + PoolMaxIdle: 1, + // No SessionVariables + } + + client, err := db.NewClient(context.Background(), cfg, result.ConnStr) + if err != nil { + t.Fatalf("NewClient: %v", err) + } + defer client.Close() + + // Just verify basic operations work without session vars + ctx := context.Background() + schemas, err := client.GetSchemas(ctx) + if err != nil { + t.Fatalf("GetSchemas: %v", err) + } + if len(schemas) == 0 { + t.Error("expected at least one schema") + } +} + +// TestSessionVars_RLSIsolation verifies end-to-end RLS isolation using +// session variables. Creates a table with RLS policy, inserts data for +// two users, and verifies each user only sees their own rows. +func TestSessionVars_RLSIsolation(t *testing.T) { + result := GetTestDB(t) + if result == nil { + return + } + t.Cleanup(result.Cleanup) + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + // Setup: create table with RLS policy + pool, err := pgxpool.New(ctx, result.ConnStr) + if err != nil { + t.Fatalf("connect: %v", err) + } + + // Create a non-superuser role for RLS testing (superusers bypass RLS) + rlsRole := fmt.Sprintf("rls_test_%d", time.Now().UnixNano()) + rlsPassword := "rls_test_pass" + + if _, err := pool.Exec(ctx, fmt.Sprintf(`CREATE ROLE %s LOGIN PASSWORD '%s'`, rlsRole, rlsPassword)); err != nil { + t.Skipf("cannot create test role (need CREATEROLE privilege): %v", err) + } + t.Cleanup(func() { + cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cleanupCancel() + cleanupPool, err := pgxpool.New(cleanupCtx, result.ConnStr) + if err != nil { + return + } + defer cleanupPool.Close() + cleanupPool.Exec(cleanupCtx, fmt.Sprintf("DROP ROLE IF EXISTS %s", rlsRole)) + }) + + // Get schema name first (needed for GRANT) + var schema string + if err := pool.QueryRow(ctx, "SELECT current_schema()").Scan(&schema); err != nil { + t.Fatalf("get schema: %v", err) + } + + stmts := []string{ + `CREATE TABLE rls_docs ( + id serial PRIMARY KEY, + owner_id text NOT NULL, + title text NOT NULL + )`, + `ALTER TABLE rls_docs ENABLE ROW LEVEL SECURITY`, + `CREATE POLICY rls_docs_owner ON rls_docs + USING (owner_id = current_setting('app.user_id', true))`, + `INSERT INTO rls_docs (owner_id, title) VALUES + ('alice', 'alice-doc-1'), + ('alice', 'alice-doc-2'), + ('bob', 'bob-doc-1')`, + fmt.Sprintf(`GRANT USAGE ON SCHEMA %s TO %s`, schema, rlsRole), + fmt.Sprintf(`GRANT SELECT ON rls_docs TO %s`, rlsRole), + } + + for _, s := range stmts { + if _, err := pool.Exec(ctx, s); err != nil { + t.Fatalf("setup %q: %v", s, err) + } + } + pool.Close() + + // Build connection string for the non-superuser role + rlsConnStr := fmt.Sprintf("postgres://%s:%s@localhost:5432/assistant?sslmode=disable&search_path=%s,public", + rlsRole, rlsPassword, schema) + + // Test: connect as alice (using non-superuser role) + cfgAlice := &config.Config{ + PoolSize: 2, + PoolMaxIdle: 1, + SessionVariables: map[string]string{"app.user_id": "alice"}, + } + clientAlice, err := db.NewClient(ctx, cfgAlice, rlsConnStr) + if err != nil { + t.Fatalf("NewClient alice: %v", err) + } + defer clientAlice.Close() + + aliceCount, err := clientAlice.GetRowCount(ctx, schema, "rls_docs") + if err != nil { + t.Fatalf("alice GetRowCount: %v", err) + } + if aliceCount != 2 { + t.Errorf("alice sees %d rows, want 2", aliceCount) + } + + // Test: connect as bob + cfgBob := &config.Config{ + PoolSize: 2, + PoolMaxIdle: 1, + SessionVariables: map[string]string{"app.user_id": "bob"}, + } + clientBob, err := db.NewClient(ctx, cfgBob, rlsConnStr) + if err != nil { + t.Fatalf("NewClient bob: %v", err) + } + defer clientBob.Close() + + bobCount, err := clientBob.GetRowCount(ctx, schema, "rls_docs") + if err != nil { + t.Fatalf("bob GetRowCount: %v", err) + } + if bobCount != 1 { + t.Errorf("bob sees %d rows, want 1", bobCount) + } + + // Test: same pool, different context vars + cfgShared := &config.Config{ + PoolSize: 2, + PoolMaxIdle: 1, + } + clientShared, err := db.NewClient(ctx, cfgShared, rlsConnStr) + if err != nil { + t.Fatalf("NewClient shared: %v", err) + } + defer clientShared.Close() + + ctxAlice := db.WithSessionVars(ctx, db.NewSessionVars(map[string]string{"app.user_id": "alice"})) + aliceCount2, err := clientShared.GetRowCount(ctxAlice, schema, "rls_docs") + if err != nil { + t.Fatalf("shared as alice: %v", err) + } + if aliceCount2 != 2 { + t.Errorf("shared pool as alice sees %d rows, want 2", aliceCount2) + } + + ctxBob := db.WithSessionVars(ctx, db.NewSessionVars(map[string]string{"app.user_id": "bob"})) + bobCount2, err := clientShared.GetRowCount(ctxBob, schema, "rls_docs") + if err != nil { + t.Fatalf("shared as bob: %v", err) + } + if bobCount2 != 1 { + t.Errorf("shared pool as bob sees %d rows, want 1", bobCount2) + } + + // Without any session var, RLS policy returns 0 rows + // (current_setting returns empty string, no rows match) + noVarCount, err := clientShared.GetRowCount(ctx, schema, "rls_docs") + if err != nil { + t.Fatalf("shared no vars: %v", err) + } + if noVarCount != 0 { + t.Errorf("shared pool with no vars sees %d rows, want 0 (RLS should block)", noVarCount) + } +} + +// TestSessionVars_ImportWithVars verifies that import operations (which +// use their own transactions) also apply session variables. +func TestSessionVars_ImportWithVars(t *testing.T) { + result := GetTestDB(t) + if result == nil { + return + } + t.Cleanup(result.Cleanup) + + ctx := context.Background() + + // Create a table + pool, err := pgxpool.New(ctx, result.ConnStr) + if err != nil { + t.Fatalf("connect: %v", err) + } + _, err = pool.Exec(ctx, `CREATE TABLE import_sv_test ( + id serial PRIMARY KEY, + data text, + created_by text DEFAULT current_setting('app.user_id', true) + )`) + if err != nil { + t.Fatalf("create table: %v", err) + } + pool.Close() + + // Import with session vars + cfg := &config.Config{ + PoolSize: 2, + PoolMaxIdle: 1, + SessionVariables: map[string]string{"app.user_id": "importer"}, + } + client, err := db.NewClient(ctx, cfg, result.ConnStr) + if err != nil { + t.Fatalf("NewClient: %v", err) + } + defer client.Close() + + schema, err := client.GetCurrentSchema(ctx) + if err != nil { + t.Fatalf("GetCurrentSchema: %v", err) + } + + err = client.ImportAppend(ctx, schema, "import_sv_test", + []string{"id", "data"}, + [][]interface{}{{1, "row1"}, {2, "row2"}}, + ) + if err != nil { + t.Fatalf("ImportAppend: %v", err) + } + + // Verify the DEFAULT used current_setting + readPool, err := pgxpool.New(ctx, result.ConnStr) + if err != nil { + t.Fatalf("connect for read: %v", err) + } + defer readPool.Close() + + var createdBy string + err = readPool.QueryRow(ctx, `SELECT created_by FROM import_sv_test WHERE id = 1`).Scan(&createdBy) + if err != nil { + t.Fatalf("read created_by: %v", err) + } + + if createdBy != "importer" { + // The DEFAULT may not fire for explicit column inserts, but the session var + // was applied to the transaction. Log the result for diagnostic purposes. + t.Logf("created_by = %q (DEFAULT current_setting may not fire on explicit insert)", createdBy) + } + + // Verify the data was actually imported + var count int + err = readPool.QueryRow(ctx, fmt.Sprintf(`SELECT count(*) FROM import_sv_test`)).Scan(&count) + if err != nil { + t.Fatalf("count: %v", err) + } + if count != 2 { + t.Errorf("imported %d rows, want 2", count) + } +} + +// TestSessionVars_ConcurrentIsolation verifies that multiple goroutines +// using the same Client with different session vars see isolated data. +// This is the core safety property of the feature. +func TestSessionVars_ConcurrentIsolation(t *testing.T) { + result := GetTestDB(t) + if result == nil { + return + } + t.Cleanup(result.Cleanup) + + ctx := context.Background() + + // Setup: create RLS table with non-superuser + pool, err := pgxpool.New(ctx, result.ConnStr) + if err != nil { + t.Fatalf("connect: %v", err) + } + + var schema string + if err := pool.QueryRow(ctx, "SELECT current_schema()").Scan(&schema); err != nil { + t.Fatalf("get schema: %v", err) + } + + rlsRole := fmt.Sprintf("rls_conc_%d", time.Now().UnixNano()) + if _, err := pool.Exec(ctx, fmt.Sprintf(`CREATE ROLE %s LOGIN PASSWORD 'test'`, rlsRole)); err != nil { + t.Skipf("cannot create test role: %v", err) + } + t.Cleanup(func() { + cPool, _ := pgxpool.New(context.Background(), result.ConnStr) + if cPool != nil { + cPool.Exec(context.Background(), fmt.Sprintf("DROP ROLE IF EXISTS %s", rlsRole)) + cPool.Close() + } + }) + + stmts := []string{ + `CREATE TABLE conc_docs (id serial PRIMARY KEY, owner text NOT NULL, val text)`, + `ALTER TABLE conc_docs ENABLE ROW LEVEL SECURITY`, + `CREATE POLICY conc_owner ON conc_docs USING (owner = current_setting('app.user_id', true))`, + fmt.Sprintf(`GRANT USAGE ON SCHEMA %s TO %s`, schema, rlsRole), + fmt.Sprintf(`GRANT SELECT, INSERT ON conc_docs TO %s`, rlsRole), + fmt.Sprintf(`GRANT USAGE, SELECT ON SEQUENCE conc_docs_id_seq TO %s`, rlsRole), + } + for i := 0; i < 20; i++ { + stmts = append(stmts, fmt.Sprintf( + `INSERT INTO conc_docs (owner, val) VALUES ('user_%d', 'data_%d')`, i, i)) + } + for _, s := range stmts { + if _, err := pool.Exec(ctx, s); err != nil { + t.Fatalf("setup %q: %v", s, err) + } + } + pool.Close() + + // Connect as non-superuser with shared pool + rlsConnStr := fmt.Sprintf("postgres://%s:test@localhost:5432/assistant?sslmode=disable&search_path=%s,public", + rlsRole, schema) + cfg := &config.Config{ + PoolSize: 5, + PoolMaxIdle: 2, + } + client, err := db.NewClient(ctx, cfg, rlsConnStr) + if err != nil { + t.Fatalf("NewClient: %v", err) + } + defer client.Close() + + // Launch 20 goroutines, each as a different user + var wg sync.WaitGroup + errors := make([]error, 20) + counts := make([]int64, 20) + + for i := 0; i < 20; i++ { + wg.Add(1) + go func(userIdx int) { + defer wg.Done() + userCtx := db.WithSessionVars(ctx, db.NewSessionVars(map[string]string{ + "app.user_id": fmt.Sprintf("user_%d", userIdx), + })) + count, err := client.GetRowCount(userCtx, schema, "conc_docs") + errors[userIdx] = err + counts[userIdx] = count + }(i) + } + wg.Wait() + + for i := 0; i < 20; i++ { + if errors[i] != nil { + t.Errorf("user_%d error: %v", i, errors[i]) + } else if counts[i] != 1 { + t.Errorf("user_%d saw %d rows, want 1 (isolation failure)", i, counts[i]) + } + } +} + +// TestSessionVars_RollbackOnQueryFailure verifies that when a query fails +// after SET LOCAL succeeds, the transaction is properly rolled back and +// the client remains usable. +func TestSessionVars_RollbackOnQueryFailure(t *testing.T) { + result := GetTestDB(t) + if result == nil { + return + } + t.Cleanup(result.Cleanup) + + cfg := &config.Config{ + PoolSize: 2, + PoolMaxIdle: 1, + SessionVariables: map[string]string{"app.user_id": "42"}, + } + client, err := db.NewClient(context.Background(), cfg, result.ConnStr) + if err != nil { + t.Fatalf("NewClient: %v", err) + } + defer client.Close() + + ctx := context.Background() + + // Execute a query that will fail (non-existent table) + _, err = client.GetRowCount(ctx, "nonexistent_schema", "nonexistent_table") + if err == nil { + t.Fatal("expected error for non-existent table") + } + + // Verify the client is still usable (connection not poisoned) + schemas, err := client.GetSchemas(ctx) + if err != nil { + t.Fatalf("client unusable after rollback: %v", err) + } + if len(schemas) == 0 { + t.Error("expected at least one schema") + } +} + +// TestSessionVars_InvalidGUCName verifies that an invalid session variable +// name produces a clear error and the client remains usable. +func TestSessionVars_InvalidGUCName(t *testing.T) { + result := GetTestDB(t) + if result == nil { + return + } + t.Cleanup(result.Cleanup) + + cfg := &config.Config{ + PoolSize: 2, + PoolMaxIdle: 1, + SessionVariables: map[string]string{"not_a_dotted_name": "val"}, + } + client, err := db.NewClient(context.Background(), cfg, result.ConnStr) + if err != nil { + t.Fatalf("NewClient: %v", err) + } + defer client.Close() + + ctx := context.Background() + + // This should fail because PostgreSQL rejects non-dotted custom GUC names + _, err = client.GetSchemas(ctx) + if err == nil { + // Some PG versions accept non-dotted names. If so, just verify it worked. + t.Log("PostgreSQL accepted non-dotted GUC name (version-dependent)") + return + } + + // Verify the error mentions the variable name + t.Logf("Expected error for invalid GUC: %v", err) + + // Verify the client is still usable after the failure + cfgClean := &config.Config{ + PoolSize: 2, + PoolMaxIdle: 1, + } + clientClean, err := db.NewClient(context.Background(), cfgClean, result.ConnStr) + if err != nil { + t.Fatalf("NewClient clean: %v", err) + } + defer clientClean.Close() + + schemas, err := clientClean.GetSchemas(ctx) + if err != nil { + t.Fatalf("clean client unusable: %v", err) + } + if len(schemas) == 0 { + t.Error("expected at least one schema") + } +}