diff --git a/kvdb/postgres/db.go b/kvdb/postgres/db.go index 5d2b482a2..6aca0276d 100644 --- a/kvdb/postgres/db.go +++ b/kvdb/postgres/db.go @@ -16,12 +16,10 @@ var sqliteCmdReplacements = sqlbase.SQLiteCmdReplacements{ "INTEGER PRIMARY KEY": "BIGSERIAL PRIMARY KEY", } -// newPostgresBackend returns a db object initialized with the passed backend -// config. If postgres connection cannot be established, then returns error. -func newPostgresBackend(ctx context.Context, config *Config, prefix string) ( - walletdb.DB, error) { - - cfg := &sqlbase.Config{ +// newSQLBaseConfig builds the shared sqlbase config used by both the regular +// and migration Postgres backends from the passed backend config and prefix. +func newSQLBaseConfig(config *Config, prefix string) *sqlbase.Config { + return &sqlbase.Config{ DriverName: "pgx", Dsn: config.Dsn, Timeout: config.Timeout, @@ -30,6 +28,22 @@ func newPostgresBackend(ctx context.Context, config *Config, prefix string) ( SQLiteCmdReplacements: sqliteCmdReplacements, WithTxLevelLock: config.WithGlobalLock, } - - return sqlbase.NewSqlBackend(ctx, cfg) +} + +// newPostgresBackend returns a db object initialized with the passed backend +// config. If postgres connection cannot be established, then returns error. +func newPostgresBackend(ctx context.Context, config *Config, prefix string) ( + walletdb.DB, error) { + + return sqlbase.NewSqlBackend(ctx, newSQLBaseConfig(config, prefix)) +} + +// NewMigrationBackend returns a Postgres backend that explicitly exposes the +// migration-only bulk KV interface. +func NewMigrationBackend(ctx context.Context, config *Config, prefix string) ( + sqlbase.MigrationBackend, error) { + + return sqlbase.NewPostgresBackend( + ctx, newSQLBaseConfig(config, prefix), + ) } diff --git a/kvdb/sqlbase/migration_bulk_postgres.go b/kvdb/sqlbase/migration_bulk_postgres.go new file mode 100644 index 000000000..a64373e34 --- /dev/null +++ b/kvdb/sqlbase/migration_bulk_postgres.go @@ -0,0 +1,428 @@ +//go:build kvdb_postgres + +package sqlbase + +import ( + "context" + "database/sql" + "errors" + "fmt" + "strings" + "sync" + + "github.com/btcsuite/btcwallet/walletdb" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/stdlib" +) + +// postgresDB adds Postgres-only capabilities to the shared SQL backend. +type postgresDB struct { + *db +} + +var ( + _ walletdb.DB = (*postgresDB)(nil) + _ MigrationBulkKVStore = (*postgresDB)(nil) +) + +// bulkLeafCols is the leaf-row projection of the shared KV table schema +// defined in schema.go. The id column is database-generated. Sequence is +// walletdb bucket metadata copied separately by InsertBucket, so leaf rows do +// not include either column in the COPY operation. +var bulkLeafCols = []string{"parent_id", "key", "value"} + +// NewPostgresBackend returns a shared SQL backend with Postgres-only +// capabilities, including migration bulk loading. +func NewPostgresBackend(ctx context.Context, cfg *Config) ( + MigrationBackend, error) { + + db, err := NewSqlBackend(ctx, cfg) + if err != nil { + return nil, err + } + + return &postgresDB{db: db}, nil +} + +// CheckEmpty returns whether the underlying KV table has no rows. +func (p *postgresDB) CheckEmpty(ctx context.Context) (bool, error) { + locker := p.bulkLocker(true) + locker.Lock() + defer locker.Unlock() + + var count int64 + err := p.db.db.QueryRowContext( + ctx, "SELECT COUNT(*) FROM "+p.table, + ).Scan(&count) + if err != nil { + return false, err + } + + return count == 0, nil +} + +// TruncateTargetTable unconditionally and irreversibly removes every row from +// the underlying KV table. It is only intended for fresh migration recovery +// where the caller owns the whole target table. +func (p *postgresDB) TruncateTargetTable(ctx context.Context) error { + locker := p.bulkLocker(false) + locker.Lock() + defer locker.Unlock() + + _, err := p.db.db.ExecContext(ctx, "TRUNCATE TABLE "+p.table) + + return err +} + +// BeginBulk opens a write transaction for bulk loading. It uses a dedicated +// *sql.Conn so InsertLeaves can reach the underlying pgx connection and COPY +// into the same transaction. Callers must defer Rollback immediately after a +// successful open so the lock and connection are released on all exits. +func (p *postgresDB) BeginBulk(ctx context.Context) (MigrationBulkKVTx, error) { + locker := p.bulkLocker(false) + locker.Lock() + + conn, err := p.db.db.Conn(ctx) + if err != nil { + locker.Unlock() + return nil, err + } + + tx, err := conn.BeginTx(ctx, &sql.TxOptions{ + Isolation: sql.LevelSerializable, + }) + if err != nil { + locker.Unlock() + _ = conn.Close() + return nil, err + } + + return &postgresBulkKVTx{ + db: p.db, + conn: conn, + tx: tx, + locker: locker, + active: true, + }, nil +} + +// BeginBulkVerify opens a read-only transaction for batched verification. +// Callers must defer Rollback immediately after a successful open so the read +// transaction lock is always released. +func (p *postgresDB) BeginBulkVerify( + ctx context.Context) (MigrationBulkKVVerifier, error) { + + locker := p.bulkLocker(true) + locker.Lock() + + tx, err := p.db.db.BeginTx(ctx, &sql.TxOptions{ + ReadOnly: true, + Isolation: sql.LevelRepeatableRead, + }) + if err != nil { + locker.Unlock() + return nil, err + } + + return &postgresBulkKVVerifier{ + db: p.db, + tx: tx, + locker: locker, + active: true, + }, nil +} + +// bulkLocker returns the same optional global lock used by regular sqlbase +// transactions so migration-only transactions respect WithTxLevelLock. +func (p *postgresDB) bulkLocker(readOnly bool) sync.Locker { + if !p.cfg.WithTxLevelLock { + return newNoopLocker() + } + if readOnly { + return p.lock.RLocker() + } + + return &p.lock +} + +// postgresBulkKVTx is a migration-only Postgres transaction for loading the SQL +// KV table directly. +type postgresBulkKVTx struct { + db *db + conn *sql.Conn + tx *sql.Tx + locker sync.Locker + active bool +} + +// InsertBucket inserts a bucket row and returns its generated id. +func (p *postgresBulkKVTx) InsertBucket(ctx context.Context, + parentID *int64, key []byte, seq uint64) (int64, error) { + + if !p.active { + return 0, walletdb.ErrTxClosed + } + if len(key) == 0 { + return 0, walletdb.ErrBucketNameRequired + } + + keyCopy := cloneBulkBytes(key) + + var id int64 + if seq != 0 { + err := p.tx.QueryRowContext( + ctx, "INSERT INTO "+p.db.table+ + " (parent_id, key, sequence) "+ + "VALUES ($1,$2,$3) RETURNING id", + parentID, keyCopy, int64(seq), + ).Scan(&id) + + return id, err + } + + err := p.tx.QueryRowContext( + ctx, "INSERT INTO "+p.db.table+" (parent_id, key) "+ + "VALUES ($1,$2) RETURNING id", + parentID, keyCopy, + ).Scan(&id) + + return id, err +} + +// InsertLeaves inserts leaf rows with Postgres COPY. +func (p *postgresBulkKVTx) InsertLeaves(ctx context.Context, + leaves []MigrationBulkLeaf) error { + + if !p.active { + return walletdb.ErrTxClosed + } + if len(leaves) == 0 { + return nil + } + + rows := make([][]any, len(leaves)) + for i := range leaves { + if leaves[i].ParentID <= 0 { + return fmt.Errorf( + "bulk leaf %d has invalid parent id %d", i, + leaves[i].ParentID, + ) + } + if len(leaves[i].Key) == 0 { + return fmt.Errorf( + "bulk leaf %d: %w", i, walletdb.ErrKeyRequired, + ) + } + + value := cloneBulkBytes(leaves[i].Value) + if value == nil { + value = []byte{} + } + + rows[i] = []any{ + leaves[i].ParentID, + cloneBulkBytes(leaves[i].Key), + value, + } + } + + var copied int64 + err := p.conn.Raw(func(driverConn any) error { + pgxConn, ok := driverConn.(*stdlib.Conn) + if !ok { + return fmt.Errorf("driver conn is %T, not "+ + "pgx/v5/stdlib.Conn", driverConn) + } + + var copyErr error + // The shared schema and normal SQL paths use unquoted + // identifiers, which Postgres folds to lowercase. CopyFrom + // quotes its identifier, so fold it explicitly to resolve the + // same physical table. + copied, copyErr = pgxConn.Conn().CopyFrom( + ctx, pgx.Identifier{strings.ToLower(p.db.table)}, + bulkLeafCols, + pgx.CopyFromRows(rows), + ) + + return copyErr + }) + if err != nil { + return err + } + if copied != int64(len(leaves)) { + return fmt.Errorf("bulk leaf copy count mismatch: got=%d "+ + "want=%d", copied, len(leaves)) + } + + return nil +} + +// Commit commits the bulk transaction and releases its dedicated connection. +func (p *postgresBulkKVTx) Commit() error { + if !p.active { + return walletdb.ErrTxClosed + } + + err := p.tx.Commit() + p.active = false + p.locker.Unlock() + closeErr := p.conn.Close() + if err != nil { + if closeErr != nil { + log.Warnf( + "Could not close bulk migration connection: %v", + closeErr, + ) + } + + return err + } + if closeErr != nil { + log.Warnf("Could not close bulk migration connection after "+ + "commit: %v", closeErr) + } + + return nil +} + +// Rollback rolls back the bulk transaction and releases its dedicated +// connection. It is idempotent for already-closed transactions. +func (p *postgresBulkKVTx) Rollback() error { + if !p.active { + return nil + } + + err := p.tx.Rollback() + p.active = false + p.locker.Unlock() + closeErr := p.conn.Close() + if err != nil && !errors.Is(err, sql.ErrTxDone) { + return err + } + + return closeErr +} + +// postgresBulkKVVerifier is a read-only Postgres transaction for batched SQL KV +// verification. +type postgresBulkKVVerifier struct { + db *db + tx *sql.Tx + locker sync.Locker + active bool +} + +// FetchTopLevel returns all top-level rows ordered by key. +func (p *postgresBulkKVVerifier) FetchTopLevel( + ctx context.Context) ([]MigrationBulkChild, error) { + + if !p.active { + return nil, walletdb.ErrTxClosed + } + + // The table name is constructed internally from the configured prefix. + //nolint:gosec + rows, err := p.tx.QueryContext(ctx, "SELECT id, parent_id, key, "+ + "value, sequence, CASE WHEN value IS NULL THEN 1 ELSE 0 END "+ + "FROM "+p.db.table+" WHERE parent_id IS NULL ORDER BY key") + if err != nil { + return nil, err + } + defer rows.Close() + + return scanBulkChildren(rows) +} + +// FetchChildren returns direct children for parentIDs ordered by parent id and +// key. +func (p *postgresBulkKVVerifier) FetchChildren(ctx context.Context, + parentIDs []int64) ([]MigrationBulkChild, error) { + + if !p.active { + return nil, walletdb.ErrTxClosed + } + if len(parentIDs) == 0 { + return nil, nil + } + + // parentIDs is passed as a native []int64; the pgx stdlib driver + // encodes it as a Postgres bigint array for the ANY($1) match. + // + // The table name is constructed internally from the configured prefix. + //nolint:gosec + rows, err := p.tx.QueryContext(ctx, "SELECT id, parent_id, key, "+ + "value, sequence, CASE WHEN value IS NULL THEN 1 ELSE 0 END "+ + "FROM "+p.db.table+" WHERE parent_id = ANY($1) "+ + "ORDER BY parent_id, key", parentIDs) + if err != nil { + return nil, err + } + defer rows.Close() + + return scanBulkChildren(rows) +} + +// Rollback closes the verifier read transaction. +func (p *postgresBulkKVVerifier) Rollback() error { + if !p.active { + return nil + } + + err := p.tx.Rollback() + p.active = false + p.locker.Unlock() + if err != nil && !errors.Is(err, sql.ErrTxDone) { + return err + } + + return nil +} + +// scanBulkChildren scans verifier rows and preserves an explicit IsBucket flag +// so empty leaf values are not confused with SQL NULL bucket markers. +func scanBulkChildren(rows *sql.Rows) ([]MigrationBulkChild, error) { + var children []MigrationBulkChild + for rows.Next() { + var ( + child MigrationBulkChild + parentID sql.NullInt64 + sequence sql.NullInt64 + bucketFlag int + ) + if err := rows.Scan( + &child.ID, &parentID, &child.Key, &child.Value, + &sequence, &bucketFlag, + ); err != nil { + return nil, err + } + if parentID.Valid { + id := parentID.Int64 + child.ParentID = &id + } + if sequence.Valid { + child.Sequence = uint64(sequence.Int64) + } + child.IsBucket = bucketFlag == 1 + if !child.IsBucket && child.Value == nil { + child.Value = []byte{} + } + + children = append(children, child) + } + + return children, rows.Err() +} + +// cloneBulkBytes copies driver-owned byte slices before they are buffered or +// returned to callers. +func cloneBulkBytes(b []byte) []byte { + if b == nil { + return nil + } + + out := make([]byte, len(b)) + copy(out, b) + + return out +}