From 15a1e5d5303024a3e1f19f1f4279614c36002dab Mon Sep 17 00:00:00 2001 From: rlindgren Date: Fri, 24 Apr 2026 13:53:45 -0400 Subject: [PATCH] refactor(db): introduce DBTX interface for query abstraction Replace *pgxpool.Pool with a DBTX interface in all ~60 package-level db functions. DBTX is satisfied by both *pgxpool.Pool and pgx.Tx, following the sqlc convention. This enables future session variable scoping via SET LOCAL in transactions without changing callers. Zero behavior change: all Client methods continue to pass c.pool, which satisfies DBTX. No new dependencies, no test changes required. --- internal/tigerfs/db/constraints.go | 21 +++--- internal/tigerfs/db/dbtx.go | 28 +++++++ internal/tigerfs/db/dbtx_test.go | 20 +++++ internal/tigerfs/db/export.go | 13 ++-- internal/tigerfs/db/indexes.go | 61 ++++++++------- internal/tigerfs/db/keys.go | 13 ++-- internal/tigerfs/db/permissions.go | 13 ++-- internal/tigerfs/db/pipeline.go | 13 ++-- internal/tigerfs/db/query.go | 49 ++++++------ internal/tigerfs/db/schema.go | 117 ++++++++++++++--------------- 10 files changed, 194 insertions(+), 154 deletions(-) create mode 100644 internal/tigerfs/db/dbtx.go create mode 100644 internal/tigerfs/db/dbtx_test.go 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..f5acf1c --- /dev/null +++ b/internal/tigerfs/db/dbtx.go @@ -0,0 +1,28 @@ +package db + +import ( + "context" + + "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. +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) +) 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..534a7b2 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) } @@ -80,7 +79,7 @@ func (c *Client) GetAllRows(ctx context.Context, schema, table string, limit int // 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 +91,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) } @@ -138,7 +137,7 @@ func (c *Client) GetFirstNRowsWithData(ctx context.Context, schema, table string // 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 +149,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) } diff --git a/internal/tigerfs/db/indexes.go b/internal/tigerfs/db/indexes.go index f7b3bdf..92e6723 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 (pool 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) } @@ -126,15 +125,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 (pool 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 } @@ -161,13 +160,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 (pool 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 } @@ -195,13 +194,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 (pool 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 } @@ -229,7 +228,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 (pool or transaction) // - schema: Schema name // - table: Table name // - column: Column name to get distinct values for @@ -237,7 +236,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 +249,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) } @@ -298,7 +297,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 +315,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) } @@ -355,7 +354,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 (pool or transaction) // - schema: Schema name // - table: Table name // - column: Indexed column to query @@ -364,7 +363,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,7 +377,7 @@ 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. @@ -391,7 +390,7 @@ func (c *Client) GetRowsByIndexValue(ctx context.Context, schema, table, column, // 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,7 +409,7 @@ 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. @@ -423,8 +422,8 @@ func (c *Client) GetRowsByIndexValueOrdered(ctx context.Context, schema, table, // 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 +464,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 (pool or transaction) // - schema: Schema name // - table: Table name // - targetColumn: Column to get distinct values for (the "next" column in navigation) @@ -475,7 +474,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 +504,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) } @@ -558,7 +557,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 (pool or transaction) // - schema: Schema name // - table: Table name // - columns: Indexed columns to query (in index order) @@ -567,7 +566,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,7 +595,7 @@ 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. diff --git a/internal/tigerfs/db/keys.go b/internal/tigerfs/db/keys.go index 458e794..4505427 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,7 +80,7 @@ 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 @@ -103,7 +102,7 @@ func (c *Client) ListRows(ctx context.Context, schema, table string, pkColumns [ // 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 +113,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) } diff --git a/internal/tigerfs/db/permissions.go b/internal/tigerfs/db/permissions.go index 8e6a854..7d646cd 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 (pool 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, @@ -120,12 +119,12 @@ func (c *Client) GetTablePermissions(ctx context.Context, schema, table string) // // Parameters: // - ctx: Context for cancellation -// - pool: PostgreSQL connection pool +// - dbtx: Database connection (pool 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 +142,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) } diff --git a/internal/tigerfs/db/pipeline.go b/internal/tigerfs/db/pipeline.go index 1001877..6cf4e4f 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 (pool 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) } @@ -175,11 +174,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 (pool 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 +193,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) } diff --git a/internal/tigerfs/db/query.go b/internal/tigerfs/db/query.go index 6da2f54..4bad62e 100644 --- a/internal/tigerfs/db/query.go +++ b/internal/tigerfs/db/query.go @@ -5,7 +5,6 @@ import ( "fmt" "strings" - "github.com/jackc/pgx/v5/pgxpool" "github.com/timescale/tigerfs/internal/tigerfs/format" "github.com/timescale/tigerfs/internal/tigerfs/logging" "go.uber.org/zap" @@ -58,7 +57,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), @@ -71,7 +70,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) } @@ -129,7 +128,7 @@ func (c *Client) GetRow(ctx context.Context, schema, table string, pk *PKMatch) } // 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), @@ -143,7 +142,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) } @@ -167,7 +166,7 @@ func (c *Client) GetColumn(ctx context.Context, schema, table string, pk *PKMatc // 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), @@ -193,7 +192,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) } @@ -262,7 +261,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), @@ -297,7 +296,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) } @@ -345,7 +344,7 @@ func (c *Client) InsertRow(ctx context.Context, schema, table string, columns [] } // 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), @@ -379,7 +378,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) } @@ -406,7 +405,7 @@ func (c *Client) UpdateRow(ctx context.Context, schema, table string, pk *PKMatc } // 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), @@ -420,7 +419,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) } @@ -448,7 +447,7 @@ func (c *Client) DeleteRow(ctx context.Context, schema, table string, pk *PKMatc // 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), @@ -460,7 +459,7 @@ 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 @@ -473,7 +472,7 @@ func (c *Client) GetFirstNRows(ctx context.Context, schema, table string, pkColu // 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), @@ -485,7 +484,7 @@ 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 @@ -512,7 +511,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), @@ -544,7 +543,7 @@ 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 @@ -566,7 +565,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), @@ -579,7 +578,7 @@ 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 @@ -601,7 +600,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), @@ -614,7 +613,7 @@ 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 @@ -627,8 +626,8 @@ func (c *Client) GetLastNRowsOrdered(ctx context.Context, schema, table string, // 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) } diff --git a/internal/tigerfs/db/schema.go b/internal/tigerfs/db/schema.go index 951c707..0e91fc4 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: Database connection (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) } @@ -42,7 +41,7 @@ func (c *Client) GetCurrentSchema(ctx context.Context) (string, error) { } // 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 +51,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 +78,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 +96,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) } @@ -143,7 +142,7 @@ func (c *Client) GetTables(ctx context.Context, schema string) ([]string, error) // 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 +160,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) } @@ -200,7 +199,7 @@ func (c *Client) GetViews(ctx context.Context, schema string) ([]string, error) // 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 +211,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) @@ -250,7 +249,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 +263,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) } @@ -302,7 +301,7 @@ func (c *Client) GetColumns(ctx context.Context, schema, table string) ([]Column } // 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 +312,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) } @@ -348,12 +347,12 @@ const SmallTableThreshold = 100000 // // Parameters: // - ctx: Context for cancellation -// - pool: PostgreSQL connection pool +// - dbtx: Database connection (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 +365,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", @@ -403,18 +402,18 @@ func (c *Client) GetTableRowCountEstimate(ctx context.Context, schema, table str // // Parameters: // - ctx: Context for cancellation -// - pool: PostgreSQL connection pool +// - dbtx: Database connection (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 +423,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 +432,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 @@ -462,7 +461,7 @@ func (c *Client) GetRowCountSmart(ctx context.Context, schema, table string) (in // 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 +477,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) } @@ -515,7 +514,7 @@ func (c *Client) GetRowCountEstimates(ctx context.Context, schema string, 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 +532,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 +591,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) } @@ -668,7 +667,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 +686,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) } @@ -719,7 +718,7 @@ func (c *Client) GetIndexDDL(ctx context.Context, schema, table string) (string, } // 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 +736,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) } @@ -769,7 +768,7 @@ func (c *Client) GetForeignKeyDDL(ctx context.Context, schema, table string) (st } // 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 +786,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) } @@ -819,7 +818,7 @@ func (c *Client) GetCheckConstraintDDL(ctx context.Context, schema, table string } // 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 +834,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) } @@ -867,7 +866,7 @@ func (c *Client) GetTriggerDDL(ctx context.Context, schema, table string) (strin } // 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 +878,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 +902,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) } @@ -941,7 +940,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 +948,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 +957,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 +968,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 +979,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 +990,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 +1001,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) } @@ -1041,7 +1040,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 +1069,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) } @@ -1114,7 +1113,7 @@ func (c *Client) GetReferencingForeignKeys(ctx context.Context, schema, table st // 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 +1126,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) } @@ -1146,7 +1145,7 @@ func (c *Client) GetSchemaTableCount(ctx context.Context, schema string) (int, e // 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 +1160,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 @@ -1182,7 +1181,7 @@ func (c *Client) GetViewComment(ctx context.Context, schema, view string) (strin // 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 +1194,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) } @@ -1232,7 +1231,7 @@ func (c *Client) GetViewCommentsBatch(ctx context.Context, schema string) (map[s } // 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 +1246,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) } @@ -1264,7 +1263,7 @@ func (c *Client) GetViewDefinition(ctx context.Context, schema, view string) (st } // 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 +1281,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) }