From 6106a18fa05e32aa3d3d923141d6d68d3e04dbf0 Mon Sep 17 00:00:00 2001 From: Wei Cao Date: Mon, 7 Sep 2026 11:26:43 +0800 Subject: [PATCH] feat: persist DuckLake object-table writes with rollback cleanup Replay accepted object-storage candidate 450fe6f0 onto current main, plus FlightSQL idempotent ClosePreparedStatement. Usage boundaries: 1. MySQL CREATE TABLE is an implicit commit and is not undone by ROLLBACK (local and object tables alike). 2. One transaction cannot write both a local table and an object table (1105). 3. The local DuckLake catalog must persist; losing it cannot be recovered from the object bucket alone. Does not update PR #489 / task75-object-table. Tree matches previous head d73a66f0. --- adapter/adapter.go | 632 +++- adapter/adapter_test.go | 583 ++++ backend/completion_registry_test.go | 182 + backend/ducklake_operation.go | 531 +++ backend/ducklake_operation_test.go | 571 +++ backend/engine.go | 71 +- backend/engine_test.go | 34 + backend/execution_lease.go | 17 + backend/execution_lease_test.go | 144 + backend/executor.go | 309 +- backend/executor_test.go | 41 + backend/handler.go | 31 +- backend/loaddata.go | 34 +- backend/loaddata_test.go | 155 + backend/mysql_handler_lifecycle_test.go | 134 + backend/raw_session_lifecycle_test.go | 126 + backend/relation_rewrite.go | 382 ++ backend/relation_rewrite_test.go | 106 + backend/session.go | 1528 +++++++- backend/session_implicit_test.go | 353 ++ backend/session_operation_test.go | 261 ++ binlogreplication/binlog_replica_applier.go | 21 +- catalog/completion_reentrancy_test.go | 389 +++ catalog/completion_registry_test.go | 232 ++ catalog/connpool.go | 3096 ++++++++++++++++- catalog/connpool_test.go | 923 +++++ catalog/database.go | 305 +- catalog/database_ddl_test.go | 38 + catalog/database_storage_failure_test.go | 501 +++ catalog/ducklake.go | 1581 ++++++++- catalog/ducklake_error_test.go | 1316 +++++++ catalog/execution_lease_callsites_test.go | 387 +++ catalog/external_session_test.go | 854 +++++ catalog/identity_lifecycle_test.go | 94 + catalog/inserter.go | 97 +- catalog/object_storage.go | 663 ++++ catalog/operation_barrier_test.go | 825 +++++ catalog/provider.go | 1529 +++++++- catalog/provider_ducklake_test.go | 201 ++ catalog/provider_generation_test.go | 487 +++ catalog/rollback_cleanup_test.go | 349 ++ catalog/table.go | 174 +- catalog/table_storage.go | 423 +++ catalog/table_storage_test.go | 215 ++ configuration/ducklake.go | 109 +- configuration/ducklake_test.go | 9 +- docker/Dockerfile | 1 + docs/design-task77-rollback-repair.md | 57 + docs/object-storage.md | 23 + flightsqlserver/managed_lifecycle.go | 809 +++++ flightsqlserver/managed_lifecycle_test.go | 314 ++ flightsqlserver/sqlite_server.go | 291 +- .../sqlite_server_ducklake_test.go | 6 +- .../sqlite_tables_schema_batch_reader.go | 15 +- main.go | 9 +- pgserver/arrow_lifecycle_test.go | 143 + pgserver/arrowloader.go | 222 +- pgserver/arrowwriter.go | 178 +- pgserver/connection_data.go | 34 +- pgserver/connection_handler.go | 1482 +++++++- pgserver/connection_handler_copy_test.go | 536 +++ pgserver/copy_binding.go | 280 ++ pgserver/copy_binding_test.go | 105 + pgserver/copy_operation_lifecycle_test.go | 235 ++ pgserver/copy_query.go | 93 + pgserver/copy_query_test.go | 78 + pgserver/copy_stdout_error_test.go | 30 + pgserver/dataloader.go | 499 ++- pgserver/dataloader_lifecycle_test.go | 206 ++ pgserver/datawriter.go | 122 +- pgserver/discard_to_sync_test.go | 258 ++ pgserver/duck_handler.go | 872 ++++- pgserver/ducklake_operation_owner_test.go | 77 + pgserver/implicit_transaction_test.go | 115 + pgserver/in_place_handler.go | 31 +- pgserver/logrepl/replication.go | 125 +- pgserver/logrepl/subscription.go | 19 +- pgserver/object_routing.go | 438 +++ pgserver/pending_copy_release_test.go | 81 + pgserver/prepared_returning_test.go | 68 + pgserver/ready_for_query_test.go | 202 ++ pgserver/returning_test.go | 19 + pgserver/rollback_cleanup.go | 191 + pgserver/rollback_cleanup_test.go | 213 ++ pgserver/subscription_handler.go | 87 +- pgserver/table_storage.go | 148 + pgserver/table_storage_wire_test.go | 64 + pgserver/transaction.go | 837 +++++ pgserver/transaction_test.go | 388 +++ 89 files changed, 30078 insertions(+), 966 deletions(-) create mode 100644 adapter/adapter_test.go create mode 100644 backend/completion_registry_test.go create mode 100644 backend/ducklake_operation.go create mode 100644 backend/ducklake_operation_test.go create mode 100644 backend/execution_lease.go create mode 100644 backend/execution_lease_test.go create mode 100644 backend/loaddata_test.go create mode 100644 backend/mysql_handler_lifecycle_test.go create mode 100644 backend/raw_session_lifecycle_test.go create mode 100644 backend/relation_rewrite.go create mode 100644 backend/relation_rewrite_test.go create mode 100644 backend/session_implicit_test.go create mode 100644 backend/session_operation_test.go create mode 100644 catalog/completion_reentrancy_test.go create mode 100644 catalog/completion_registry_test.go create mode 100644 catalog/database_ddl_test.go create mode 100644 catalog/database_storage_failure_test.go create mode 100644 catalog/ducklake_error_test.go create mode 100644 catalog/execution_lease_callsites_test.go create mode 100644 catalog/external_session_test.go create mode 100644 catalog/identity_lifecycle_test.go create mode 100644 catalog/object_storage.go create mode 100644 catalog/operation_barrier_test.go create mode 100644 catalog/provider_generation_test.go create mode 100644 catalog/rollback_cleanup_test.go create mode 100644 catalog/table_storage.go create mode 100644 catalog/table_storage_test.go create mode 100644 docs/design-task77-rollback-repair.md create mode 100644 docs/object-storage.md create mode 100644 flightsqlserver/managed_lifecycle.go create mode 100644 flightsqlserver/managed_lifecycle_test.go create mode 100644 pgserver/arrow_lifecycle_test.go create mode 100644 pgserver/connection_handler_copy_test.go create mode 100644 pgserver/copy_binding.go create mode 100644 pgserver/copy_binding_test.go create mode 100644 pgserver/copy_operation_lifecycle_test.go create mode 100644 pgserver/copy_query.go create mode 100644 pgserver/copy_query_test.go create mode 100644 pgserver/copy_stdout_error_test.go create mode 100644 pgserver/dataloader_lifecycle_test.go create mode 100644 pgserver/discard_to_sync_test.go create mode 100644 pgserver/ducklake_operation_owner_test.go create mode 100644 pgserver/implicit_transaction_test.go create mode 100644 pgserver/object_routing.go create mode 100644 pgserver/pending_copy_release_test.go create mode 100644 pgserver/prepared_returning_test.go create mode 100644 pgserver/ready_for_query_test.go create mode 100644 pgserver/rollback_cleanup.go create mode 100644 pgserver/rollback_cleanup_test.go create mode 100644 pgserver/table_storage.go create mode 100644 pgserver/table_storage_wire_test.go create mode 100644 pgserver/transaction.go create mode 100644 pgserver/transaction_test.go diff --git a/adapter/adapter.go b/adapter/adapter.go index 14bffd10..d09c1365 100644 --- a/adapter/adapter.go +++ b/adapter/adapter.go @@ -3,11 +3,21 @@ package adapter import ( "context" stdsql "database/sql" + "errors" "fmt" + "strings" + "sync" "github.com/dolthub/go-mysql-server/sql" ) +// ErrTransactionBindingChanged reports that a finalizer was asked to operate +// on a physical transaction which is no longer the current binding for the +// logical session. A stale callback must be visible to the caller: treating it +// as a successful commit can make the protocol advertise success while the +// transaction was already rolled back or replaced. +var ErrTransactionBindingChanged = errors.New("transaction binding is no longer current") + type ConnectionHolder interface { GetConn(ctx context.Context) (*stdsql.Conn, error) GetTxn(ctx context.Context, options *stdsql.TxOptions) (*stdsql.Tx, error) @@ -20,6 +30,191 @@ type ConnectionHolder interface { CloseConn() } +// TransactionBindingHolder is an optional extension implemented by session +// types that can return the physical connection and transaction as one +// lifecycle snapshot. Keeping this separate from ConnectionHolder preserves +// compatibility with the small test/session implementations supplied by GMS. +type TransactionBindingHolder interface { + GetTxnBinding() (*stdsql.Conn, *stdsql.Tx) +} + +// ReleaseTransactionBindingHolder is the release-side counterpart to +// TransactionBindingHolder. Implementations may admit this snapshot while a +// pool shutdown is waiting for ordinary readers to drain. Callers must use it +// only to capture state for a subsequent identity-safe release/finalizer. +type ReleaseTransactionBindingHolder interface { + GetTxnBindingForRelease() (*stdsql.Conn, *stdsql.Tx) +} + +// ReleaseTransactionHolder exposes a transaction-only release snapshot for +// legacy call sites that do not need the physical owner. Production sessions +// implement this alongside ReleaseTransactionBindingHolder. +type ReleaseTransactionHolder interface { + TryGetTxnForRelease() *stdsql.Tx +} + +// ExecutionSnapshotHolder can acquire the connection/transaction pair under +// the session's lifecycle lock. The catalog flag selects a connection without +// schema switching for metadata work; ordinary execution selects the current +// logical schema. +type ExecutionSnapshotHolder interface { + GetExecutionSnapshot(context.Context, bool) (*stdsql.Conn, *stdsql.Tx, error) +} + +// ExecutionSnapshotLeaseHolder is the stronger execution-snapshot surface. +// Implementations keep the pool generation alive until release is called. The +// release callback must be safe to invoke more than once; adapter callers also +// wrap it defensively so a failed construction path cannot leak a lifecycle +// reader. +// +// The interface is optional to preserve compatibility with the small session +// doubles used by GMS and by downstream embedders. Those holders receive a +// no-op release from GetExecutionSnapshotWithLease. +type ExecutionSnapshotLeaseHolder interface { + GetExecutionSnapshotLease(context.Context, bool) (*stdsql.Conn, *stdsql.Tx, func(), error) +} + +// IdentityTransactionCloser removes a transaction mapping only when it still +// refers to the supplied transaction. This prevents an old commit/rollback +// defer from deleting a newer transaction created for the same session ID. +type IdentityTransactionCloser interface { + CloseTxnIf(*stdsql.Tx) +} + +// IdentityConnectionCloser removes and closes a connection only when it still +// refers to the supplied physical connection. +type IdentityConnectionCloser interface { + CloseConnIf(*stdsql.Conn) error +} + +// IdentityConnectionBindingCloser removes and closes a connection only when +// both the physical connection and its expected transaction binding still +// match. This closes the same-connection replacement race that a connection- +// only identity check cannot see. +type IdentityConnectionBindingCloser interface { + CloseConnIfBinding(*stdsql.Conn, *stdsql.Tx) error +} + +// TransactionFinalizer performs the driver operation and lifecycle-map +// update as one session-scoped operation. Implementations must keep a stale +// transaction from evicting a replacement binding for the same session. +type TransactionFinalizer interface { + CommitTxn(*stdsql.Tx) error + RollbackTxn(*stdsql.Tx) (*stdsql.Conn, bool, error) +} + +// RollbackCleanupFinalizer extends rollback finalization with an optional +// callback that runs while the session still owns its lifecycle lock. This +// closes the gap in which a replacement transaction could otherwise be +// admitted before post-rollback maintenance has finished. +type RollbackCleanupFinalizer interface { + RollbackTxnWithCleanup(*stdsql.Tx, func(*stdsql.Conn) error) (bool, error) +} + +// SQLExecutor is the common database/sql surface used by statement execution. +// Both *sql.Conn and *sql.Tx satisfy it. Keeping the interface here lets query +// paths bind to an already active session transaction without giving up the +// connection-specific setup needed by callers that acquire a *sql.Conn. +type SQLExecutor interface { + ExecContext(context.Context, string, ...any) (stdsql.Result, error) + QueryContext(context.Context, string, ...any) (*stdsql.Rows, error) + QueryRowContext(context.Context, string, ...any) *stdsql.Row + PrepareContext(context.Context, string) (*stdsql.Stmt, error) +} + +// LeasedRows keeps an execution snapshot alive until the result set is closed +// or exhausted. database/sql normally closes Rows at EOF; the explicit wrapper +// also releases the pool lifecycle reader at that same terminal boundary. +type LeasedRows struct { + *stdsql.Rows + release func() + once sync.Once +} + +func newLeasedRows(rows *stdsql.Rows, release func()) *LeasedRows { + return &LeasedRows{Rows: rows, release: idempotentRelease(release)} +} + +func (r *LeasedRows) releaseLease() { + if r == nil { + return + } + r.once.Do(r.release) +} + +func (r *LeasedRows) Close() (err error) { + if r == nil { + return nil + } + defer r.releaseLease() + if r.Rows == nil { + return nil + } + return r.Rows.Close() +} + +func (r *LeasedRows) Next() bool { + if r == nil || r.Rows == nil { + if r != nil { + r.releaseLease() + } + return false + } + if r.Rows.Next() { + return true + } + _ = r.Close() + return false +} + +func (r *LeasedRows) NextResultSet() bool { + if r == nil || r.Rows == nil { + if r != nil { + r.releaseLease() + } + return false + } + if r.Rows.NextResultSet() { + return true + } + _ = r.Close() + return false +} + +// LeasedRow retains a snapshot until Scan consumes the database/sql row. A Row +// has no Close method, so Scan is its only normal terminal operation. +type LeasedRow struct { + *stdsql.Row + release func() + err error + once sync.Once +} + +func NewLeasedRow(row *stdsql.Row, release func()) *LeasedRow { + return &LeasedRow{Row: row, release: idempotentRelease(release)} +} + +// NewLeasedRowError preserves an error that occurred before database/sql could +// construct a Row. QueryRow-style APIs surface errors through Scan, so callers +// must receive the original admission or lifecycle error there as well. +func NewLeasedRowError(err error) *LeasedRow { + return &LeasedRow{release: idempotentRelease(nil), err: err} +} + +func (r *LeasedRow) Scan(dest ...any) (err error) { + if r == nil { + return fmt.Errorf("query row is unavailable") + } + defer r.once.Do(r.release) + if r.err != nil { + return r.err + } + if r.Row == nil { + return fmt.Errorf("query row is unavailable") + } + return r.Row.Scan(dest...) +} + func GetConn(ctx *sql.Context) (*stdsql.Conn, error) { return ctx.Session.(ConnectionHolder).GetConn(ctx) } @@ -44,6 +239,42 @@ func TryGetTxn(ctx *sql.Context) *stdsql.Tx { return ctx.Session.(ConnectionHolder).TryGetTxn() } +// TryGetTxnBinding returns an atomic (connection, transaction) snapshot when +// the session supplies that stronger lifecycle surface. Older holders retain +// the transaction-only fallback; callers that require connection affinity can +// treat a nil connection as unavailable rather than guessing another pool +// connection. +func TryGetTxnBinding(ctx *sql.Context) (*stdsql.Conn, *stdsql.Tx) { + holder := ctx.Session.(ConnectionHolder) + if binding, ok := holder.(TransactionBindingHolder); ok { + return binding.GetTxnBinding() + } + return nil, holder.TryGetTxn() +} + +// TryGetTxnBindingForRelease captures a binding from a release-aware session +// surface. The production pool admits this read while shutdown is pending, +// which prevents a close/finalizer path from deadlocking behind the writer it +// is about to help drain. Legacy holders fall back to the ordinary snapshot. +func TryGetTxnBindingForRelease(ctx *sql.Context) (*stdsql.Conn, *stdsql.Tx) { + holder := ctx.Session.(ConnectionHolder) + if binding, ok := holder.(ReleaseTransactionBindingHolder); ok { + return binding.GetTxnBindingForRelease() + } + return TryGetTxnBinding(ctx) +} + +// TryGetTxnForRelease is the transaction-only form of +// TryGetTxnBindingForRelease. +func TryGetTxnForRelease(ctx *sql.Context) *stdsql.Tx { + holder := ctx.Session.(ConnectionHolder) + if release, ok := holder.(ReleaseTransactionHolder); ok { + return release.TryGetTxnForRelease() + } + _, tx := TryGetTxnBindingForRelease(ctx) + return tx +} + func GetCurrentCatalog(ctx *sql.Context) string { return ctx.Session.(ConnectionHolder).GetCurrentCatalog() } @@ -56,81 +287,436 @@ func CloseTxn(ctx *sql.Context) { ctx.Session.(ConnectionHolder).CloseTxn() } -func Query(ctx *sql.Context, query string, args ...any) (*stdsql.Rows, error) { - conn, err := GetConn(ctx) +// CloseTxnIf performs an identity-safe transaction mapping removal when the +// session supports it, with the legacy unconditional behavior as a fallback. +func CloseTxnIf(ctx *sql.Context, tx *stdsql.Tx) { + if tx == nil { + return + } + holder := ctx.Session.(ConnectionHolder) + if closer, ok := holder.(IdentityTransactionCloser); ok { + closer.CloseTxnIf(tx) + return + } + // Legacy holders expose no identity-aware removal primitive. Compare the + // current pointer before removing it so a delayed callback does not evict a + // replacement transaction in the common, non-concurrent case. + if TryGetTxnForRelease(ctx) == tx { + holder.CloseTxn() + } +} + +// CloseConnIf performs an identity-safe connection removal when available. +// The fallback closes the supplied connection but cannot update a legacy +// holder's map; production sessions implement the stronger surface. +func CloseConnIf(ctx *sql.Context, conn *stdsql.Conn) error { + if conn == nil { + return nil + } + holder := ctx.Session.(ConnectionHolder) + if closer, ok := holder.(IdentityConnectionCloser); ok { + return closer.CloseConnIf(conn) + } + return conn.Close() +} + +// CloseConnIfBinding performs an atomic connection/transaction identity check +// when the session supports it. Legacy holders fall back to a strict snapshot +// check before closing the supplied connection. +func CloseConnIfBinding(ctx *sql.Context, conn *stdsql.Conn, tx *stdsql.Tx) error { + if conn == nil { + return nil + } + holder := ctx.Session.(ConnectionHolder) + if closer, ok := holder.(IdentityConnectionBindingCloser); ok { + return closer.CloseConnIfBinding(conn, tx) + } + currentConn, currentTx := TryGetTxnBindingForRelease(ctx) + if currentConn != conn || currentTx != tx { + return nil + } + return conn.Close() +} + +// FinalizeCommit commits tx and removes its session binding. Production +// sessions implement TransactionFinalizer so the commit and identity check +// are serialized by the connection pool. The fallback preserves compatibility +// with small test/session holders. +func FinalizeCommit(ctx *sql.Context, tx *stdsql.Tx) error { + if tx == nil { + return nil + } + holder := ctx.Session.(ConnectionHolder) + if finalizer, ok := holder.(TransactionFinalizer); ok { + return finalizer.CommitTxn(tx) + } + conn, current := TryGetTxnBindingForRelease(ctx) + // A fallback holder cannot atomically finalize a transaction that is no + // longer its current binding. In particular, a nil current transaction is + // not permission to operate on the caller's stale *sql.Tx: doing so can + // commit or roll back a transaction whose owner has already been replaced + // or evicted. Production holders use TransactionFinalizer above; this strict + // check keeps legacy/test holders fail-closed as well. + if current != tx { + return ErrTransactionBindingChanged + } + err := tx.Commit() + CloseTxnIf(ctx, tx) + if err != nil { + if conn != nil { + err = errors.Join(err, CloseConnIf(ctx, conn)) + } + } + return err +} + +// FinalizeRollback rolls back tx, removes its session binding, and returns the +// owning physical connection plus whether the driver proved the transaction +// inactive. The connection is nil when the transaction was stale and no +// longer owned by this session, which prevents cleanup from touching a +// replacement transaction. +func FinalizeRollback(ctx *sql.Context, tx *stdsql.Tx) (*stdsql.Conn, bool, error) { + if tx == nil { + return nil, true, nil + } + holder := ctx.Session.(ConnectionHolder) + if finalizer, ok := holder.(TransactionFinalizer); ok { + return finalizer.RollbackTxn(tx) + } + conn, current := TryGetTxnBindingForRelease(ctx) + // Do not guess at ownership when the mapping disappeared. A stale rollback + // must not touch the transaction or acquire/reuse a connection that may now + // belong to another session lifecycle. + if current != tx { + return nil, false, nil + } + err := tx.Rollback() + inactive := transactionInactiveError(err) + CloseTxnIf(ctx, tx) + if !inactive && conn != nil { + err = errors.Join(err, CloseConnIf(ctx, conn)) + } + return conn, inactive, err +} + +// FinalizeRollbackWithCleanup rolls back tx and, when supported by the +// session, runs cleanup while the old transaction's physical owner remains +// reserved. Stale transactions never invoke the callback. +func FinalizeRollbackWithCleanup( + ctx *sql.Context, + tx *stdsql.Tx, + cleanup func(*stdsql.Conn) error, +) (bool, error) { + if tx == nil { + return true, nil + } + holder := ctx.Session.(ConnectionHolder) + if finalizer, ok := holder.(RollbackCleanupFinalizer); ok { + return finalizer.RollbackTxnWithCleanup(tx, cleanup) + } + conn, inactive, err := FinalizeRollback(ctx, tx) + if inactive && conn != nil && cleanup != nil { + err = errors.Join(err, cleanup(conn)) + } + return inactive, err +} + +func transactionInactiveError(err error) bool { + if err == nil { + return true + } + return errors.Is(err, stdsql.ErrTxDone) || + strings.Contains(strings.ToLower(err.Error()), "no transaction is active") +} + +// IsTransactionInactiveError reports the narrow set of driver outcomes that +// prove a transaction has already ended. Callers that perform replication or +// connection teardown can ignore those outcomes while still surfacing an +// unexpected rollback failure. +func IsTransactionInactiveError(err error) bool { + return transactionInactiveError(err) +} + +// SQLExecutorForConn returns the active session transaction when one exists; +// otherwise it returns the already acquired connection. Callers that need +// driver-connection operations (for example, DuckLake ATTACH) should continue +// using the original *sql.Conn directly. +func SQLExecutorForConn(ctx *sql.Context, conn *stdsql.Conn) SQLExecutor { + if _, tx := TryGetTxnBinding(ctx); tx != nil { + return tx + } + return conn +} + +// GetSQLExecutor acquires the normal session connection only when no active +// transaction owns it. It is intended for statement helpers whose work must +// participate in the session transaction when present. +func GetSQLExecutor(ctx *sql.Context) (SQLExecutor, error) { + execer, _, _, err := GetExecutionSnapshot(ctx) + return execer, err +} + +// GetCatalogExecutor is the catalog equivalent of GetSQLExecutor. Metadata +// reads and writes therefore observe the same transaction as object-table DML. +func GetCatalogExecutor(ctx *sql.Context) (SQLExecutor, error) { + execer, _, _, err := getExecutionSnapshot(ctx, true) + return execer, err +} + +// GetExecutionSnapshot acquires the executor used for statement work together +// with its underlying physical connection and transaction. Session +// implementations with an atomic binding avoid a Load -> acquire -> Load +// race; legacy holders still receive the best available snapshot. +func GetExecutionSnapshot(ctx *sql.Context) (SQLExecutor, *stdsql.Conn, *stdsql.Tx, error) { + return getExecutionSnapshot(ctx, false) +} + +// GetExecutionSnapshotWithLease is the lifetime-safe form of +// GetExecutionSnapshot. The returned release function keeps the selected +// physical connection and transaction generation usable until the caller has +// finished executing and consuming the result. Callers must release it on +// every terminal path, including iterator construction errors. +func GetExecutionSnapshotWithLease(ctx *sql.Context) (SQLExecutor, *stdsql.Conn, *stdsql.Tx, func(), error) { + return getExecutionSnapshotWithLease(ctx, false) +} + +// GetCatalogExecutionSnapshot is the metadata variant of +// GetExecutionSnapshot. It preserves the active transaction binding while +// avoiding an implicit current-schema switch on a newly acquired connection. +func GetCatalogExecutionSnapshot(ctx *sql.Context) (SQLExecutor, *stdsql.Conn, *stdsql.Tx, error) { + return getExecutionSnapshot(ctx, true) +} + +// GetCatalogExecutionSnapshotWithLease is the metadata variant of +// GetExecutionSnapshotWithLease. +func GetCatalogExecutionSnapshotWithLease(ctx *sql.Context) (SQLExecutor, *stdsql.Conn, *stdsql.Tx, func(), error) { + return getExecutionSnapshotWithLease(ctx, true) +} + +// GetTxnExecutionSnapshotWithLease opens or observes the session transaction, +// then retains the generation containing that exact transaction for subsequent +// driver work. A transition between the two snapshots is reported as a binding +// change; callers never execute on the stale pointer returned by GetTxn. +func GetTxnExecutionSnapshotWithLease(ctx *sql.Context, options *stdsql.TxOptions) (*stdsql.Conn, *stdsql.Tx, func(), error) { + return getTxnExecutionSnapshotWithLease(ctx, options, false) +} + +// GetCatalogTxnExecutionSnapshotWithLease is the catalog-only counterpart to +// GetTxnExecutionSnapshotWithLease. +func GetCatalogTxnExecutionSnapshotWithLease(ctx *sql.Context, options *stdsql.TxOptions) (*stdsql.Conn, *stdsql.Tx, func(), error) { + return getTxnExecutionSnapshotWithLease(ctx, options, true) +} + +func getTxnExecutionSnapshotWithLease( + ctx *sql.Context, + options *stdsql.TxOptions, + catalogOnly bool, +) (*stdsql.Conn, *stdsql.Tx, func(), error) { + holder := ctx.Session.(ConnectionHolder) + var ( + tx *stdsql.Tx + err error + ) + if catalogOnly { + tx, err = holder.GetCatalogTxn(ctx, options) + } else { + tx, err = holder.GetTxn(ctx, options) + } + if err != nil { + return nil, nil, func() {}, err + } + _, conn, current, release, err := getExecutionSnapshotWithLease(ctx, catalogOnly) + if err != nil { + return nil, nil, release, err + } + if current == nil || current != tx { + release() + return nil, nil, func() {}, ErrTransactionBindingChanged + } + return conn, current, release, nil +} + +func idempotentRelease(release func()) func() { + if release == nil { + return func() {} + } + var once sync.Once + return func() { + once.Do(release) + } +} + +func getExecutionSnapshotWithLease(ctx *sql.Context, catalogOnly bool) (SQLExecutor, *stdsql.Conn, *stdsql.Tx, func(), error) { + holder := ctx.Session.(ConnectionHolder) + if leased, ok := holder.(ExecutionSnapshotLeaseHolder); ok { + conn, tx, release, err := leased.GetExecutionSnapshotLease(ctx, catalogOnly) + release = idempotentRelease(release) + if err != nil { + // A holder may have admitted a lease before discovering an invalid + // snapshot. Do not make every caller remember to release on error. + release() + return nil, nil, tx, func() {}, err + } + if tx != nil { + if conn == nil { + release() + return nil, nil, tx, func() {}, fmt.Errorf("transaction execution snapshot has no connection") + } + return tx, conn, tx, release, nil + } + if conn == nil { + release() + return nil, nil, nil, func() {}, fmt.Errorf("execution snapshot has no connection") + } + return conn, conn, nil, release, nil + } + + // Legacy holders cannot retain a pool lifecycle reader, so preserve the + // existing best-effort snapshot and make the lease a no-op. + execer, conn, tx, err := getExecutionSnapshot(ctx, catalogOnly) + return execer, conn, tx, func() {}, err +} + +func getExecutionSnapshot(ctx *sql.Context, catalogOnly bool) (SQLExecutor, *stdsql.Conn, *stdsql.Tx, error) { + holder := ctx.Session.(ConnectionHolder) + getConn := holder.GetConn + if catalogOnly { + getConn = holder.GetCatalogConn + } + if snapshot, ok := holder.(ExecutionSnapshotHolder); ok { + conn, tx, err := snapshot.GetExecutionSnapshot(ctx, catalogOnly) + if err != nil { + return nil, nil, tx, err + } + if tx != nil { + if conn == nil { + return nil, nil, tx, fmt.Errorf("transaction execution snapshot has no connection") + } + return tx, conn, tx, nil + } + if conn == nil { + return nil, nil, nil, fmt.Errorf("execution snapshot has no connection") + } + return conn, conn, nil, nil + } + conn, tx := TryGetTxnBinding(ctx) + if tx != nil { + if conn == nil { + // A legacy holder exposes only the transaction pointer. Reacquiring + // a session connection here would guess at the transaction's physical + // owner and could execute against a replacement connection. Refuse the + // snapshot instead; production holders provide the atomic owner map. + return nil, nil, tx, fmt.Errorf("active transaction has no physical owner") + } + return tx, conn, tx, nil + } + if conn == nil { + var err error + conn, err = getConn(ctx) + if err != nil { + return nil, nil, nil, err + } + } + // A transaction may have been established while the connection was being + // acquired by a legacy holder. Recheck before selecting the executor. + if reboundConn, reboundTx := TryGetTxnBinding(ctx); reboundTx != nil { + if reboundConn == nil { + return nil, nil, reboundTx, fmt.Errorf("active transaction has no physical owner") + } + conn = reboundConn + return reboundTx, conn, reboundTx, nil + } + return conn, conn, nil, nil +} + +func Query(ctx *sql.Context, query string, args ...any) (*LeasedRows, error) { + execer, _, _, release, err := GetExecutionSnapshotWithLease(ctx) + if err != nil { + return nil, err + } + rows, err := execer.QueryContext(ctx, query, args...) if err != nil { + release() return nil, err } - return conn.QueryContext(ctx, query, args...) + return newLeasedRows(rows, release), nil } -func QueryRow(ctx *sql.Context, query string, args ...any) *stdsql.Row { - conn, err := GetConn(ctx) +func QueryRow(ctx *sql.Context, query string, args ...any) *LeasedRow { + execer, _, _, release, err := GetExecutionSnapshotWithLease(ctx) if err != nil { - return nil + return NewLeasedRowError(err) } - return conn.QueryRowContext(ctx, query, args...) + return NewLeasedRow(execer.QueryRowContext(ctx, query, args...), release) } // QueryCatalog is a helper function to query the catalog, such as information_schema. // Unlike QueryContext, this function does not require a schema name to be set on the connection, // and the current schema of the connection does not matter. -func QueryCatalog(ctx *sql.Context, query string, args ...any) (*stdsql.Rows, error) { - conn, err := ctx.Session.(ConnectionHolder).GetCatalogConn(ctx) +func QueryCatalog(ctx *sql.Context, query string, args ...any) (*LeasedRows, error) { + execer, _, _, release, err := GetCatalogExecutionSnapshotWithLease(ctx) if err != nil { return nil, err } - return conn.QueryContext(ctx, query, args...) + rows, err := execer.QueryContext(ctx, query, args...) + if err != nil { + release() + return nil, err + } + return newLeasedRows(rows, release), nil } -func QueryRowCatalog(ctx *sql.Context, query string, args ...any) *stdsql.Row { - conn, err := ctx.Session.(ConnectionHolder).GetCatalogConn(ctx) +func QueryRowCatalog(ctx *sql.Context, query string, args ...any) *LeasedRow { + execer, _, _, release, err := GetCatalogExecutionSnapshotWithLease(ctx) if err != nil { - return nil + return NewLeasedRowError(err) } - return conn.QueryRowContext(ctx, query, args...) + return NewLeasedRow(execer.QueryRowContext(ctx, query, args...), release) } func Exec(ctx *sql.Context, query string, args ...any) (stdsql.Result, error) { - conn, err := GetConn(ctx) + execer, _, _, release, err := GetExecutionSnapshotWithLease(ctx) if err != nil { return nil, err } - return conn.ExecContext(ctx, query, args...) + defer release() + return execer.ExecContext(ctx, query, args...) } // ExecCatalog is a helper function to execute a catalog modification query, such as creating a database. // Unlike ExecContext, this function does not require a schema name to be set on the connection, // and the current schema of the connection does not matter. func ExecCatalog(ctx *sql.Context, query string, args ...any) (stdsql.Result, error) { - conn, err := ctx.Session.(ConnectionHolder).GetCatalogConn(ctx) + execer, _, _, release, err := GetCatalogExecutionSnapshotWithLease(ctx) if err != nil { return nil, err } - return conn.ExecContext(ctx, query, args...) + defer release() + return execer.ExecContext(ctx, query, args...) } func ExecCatalogInTxn(ctx *sql.Context, query string, args ...any) (stdsql.Result, error) { - tx, err := ctx.Session.(ConnectionHolder).GetCatalogTxn(ctx, nil) + _, tx, release, err := GetCatalogTxnExecutionSnapshotWithLease(ctx, nil) if err != nil { return nil, err } + defer release() return tx.ExecContext(ctx, query, args...) } func ExecInTxn(ctx *sql.Context, query string, args ...any) (stdsql.Result, error) { - tx, err := GetTxn(ctx, nil) + _, tx, release, err := GetTxnExecutionSnapshotWithLease(ctx, nil) if err != nil { return nil, err } + defer release() return tx.ExecContext(ctx, query, args...) } func CommitAndCloseTxn(sqlCtx *sql.Context) error { - tx := TryGetTxn(sqlCtx) + tx := TryGetTxnForRelease(sqlCtx) if tx != nil { - defer CloseTxn(sqlCtx) - if err := tx.Commit(); err != nil { + if err := FinalizeCommit(sqlCtx, tx); err != nil { return fmt.Errorf("failed to commit transaction: %w", err) } } diff --git a/adapter/adapter_test.go b/adapter/adapter_test.go new file mode 100644 index 00000000..3b040a82 --- /dev/null +++ b/adapter/adapter_test.go @@ -0,0 +1,583 @@ +package adapter_test + +import ( + "context" + stdsql "database/sql" + "database/sql/driver" + "errors" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/DATA-DOG/go-sqlmock" + "github.com/apecloud/myduckserver/adapter" + "github.com/dolthub/go-mysql-server/memory" + "github.com/dolthub/go-mysql-server/sql" + "github.com/stretchr/testify/require" +) + +type snapshotHolder struct { + *memory.Session + conn *stdsql.Conn + queries []bool + activeTx *stdsql.Tx + currentCat string + currentSch string +} + +// legacyHolder intentionally implements only ConnectionHolder. It models the +// small session doubles that predate the atomic transaction-owner snapshot. +type legacyHolder struct { + *memory.Session + conn *stdsql.Conn + activeTx *stdsql.Tx + getConnCalls int + onGetConn func() +} + +type leaseHolder struct { + *snapshotHolder + leasedTx *stdsql.Tx + releaseCalls atomic.Int32 + onRelease func() + leaseErr error +} + +func (h *leaseHolder) GetExecutionSnapshotLease(_ context.Context, catalogOnly bool) (*stdsql.Conn, *stdsql.Tx, func(), error) { + h.queries = append(h.queries, catalogOnly) + tx := h.activeTx + if h.leasedTx != nil { + tx = h.leasedTx + } + release := func() { + h.releaseCalls.Add(1) + if h.onRelease != nil { + h.onRelease() + } + } + return h.conn, tx, release, h.leaseErr +} + +type leaseTransitionWaiter struct { + released chan struct{} + started chan struct{} + done chan struct{} + once sync.Once +} + +func newLeaseTransitionWaiter() *leaseTransitionWaiter { + return &leaseTransitionWaiter{ + released: make(chan struct{}), + started: make(chan struct{}), + done: make(chan struct{}), + } +} + +func (w *leaseTransitionWaiter) release() { + w.once.Do(func() { close(w.released) }) +} + +func (w *leaseTransitionWaiter) start() { + go func() { + close(w.started) + <-w.released + close(w.done) + }() + <-w.started +} + +type blockingDriverArgument struct { + started chan struct{} + allow <-chan struct{} + once sync.Once +} + +func (a *blockingDriverArgument) Match(driver.Value) bool { + a.once.Do(func() { close(a.started) }) + <-a.allow + return true +} + +type blockingQueryMatcher struct { + started chan struct{} + allow <-chan struct{} + delegate sqlmock.QueryMatcher + once sync.Once +} + +func (m *blockingQueryMatcher) Match(expectedSQL, actualSQL string) error { + m.once.Do(func() { close(m.started) }) + <-m.allow + return m.delegate.Match(expectedSQL, actualSQL) +} + +var _ adapter.ConnectionHolder = (*snapshotHolder)(nil) +var _ adapter.ExecutionSnapshotHolder = (*snapshotHolder)(nil) +var _ adapter.ExecutionSnapshotLeaseHolder = (*leaseHolder)(nil) + +func (h *snapshotHolder) GetConn(context.Context) (*stdsql.Conn, error) { return h.conn, nil } + +func (h *snapshotHolder) GetTxn(context.Context, *stdsql.TxOptions) (*stdsql.Tx, error) { + return h.activeTx, nil +} + +func (h *snapshotHolder) GetCatalogConn(context.Context) (*stdsql.Conn, error) { return h.conn, nil } + +func (h *snapshotHolder) GetCatalogTxn(context.Context, *stdsql.TxOptions) (*stdsql.Tx, error) { + return h.activeTx, nil +} + +func (h *snapshotHolder) TryGetTxn() *stdsql.Tx { return h.activeTx } + +func (h *snapshotHolder) GetCurrentCatalog() string { return h.currentCat } + +func (h *snapshotHolder) GetCurrentSchema() string { return h.currentSch } + +func (h *snapshotHolder) CloseTxn() {} + +func (h *snapshotHolder) CloseConn() {} + +func (h *snapshotHolder) GetExecutionSnapshot(_ context.Context, catalogOnly bool) (*stdsql.Conn, *stdsql.Tx, error) { + h.queries = append(h.queries, catalogOnly) + return h.conn, h.activeTx, nil +} + +func (h *legacyHolder) GetConn(context.Context) (*stdsql.Conn, error) { + h.getConnCalls++ + if h.onGetConn != nil { + h.onGetConn() + } + return h.conn, nil +} + +func (h *legacyHolder) GetTxn(context.Context, *stdsql.TxOptions) (*stdsql.Tx, error) { + return h.activeTx, nil +} + +func (h *legacyHolder) GetCatalogConn(context.Context) (*stdsql.Conn, error) { return h.conn, nil } + +func (h *legacyHolder) GetCatalogTxn(context.Context, *stdsql.TxOptions) (*stdsql.Tx, error) { + return h.activeTx, nil +} + +func (h *legacyHolder) TryGetTxn() *stdsql.Tx { return h.activeTx } + +func (h *legacyHolder) GetCurrentCatalog() string { return "memory" } + +func (h *legacyHolder) GetCurrentSchema() string { return "main" } + +func (h *legacyHolder) CloseTxn() {} + +func (h *legacyHolder) CloseConn() {} + +func newLeaseTestContext( + t *testing.T, + matcher sqlmock.QueryMatcher, +) (*sql.Context, *leaseHolder, sqlmock.Sqlmock) { + t.Helper() + if matcher == nil { + matcher = sqlmock.QueryMatcherEqual + } + db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(matcher)) + require.NoError(t, err) + t.Cleanup(func() { _ = db.Close() }) + conn, err := db.Conn(context.Background()) + require.NoError(t, err) + t.Cleanup(func() { _ = conn.Close() }) + + holder := &leaseHolder{snapshotHolder: &snapshotHolder{ + Session: memory.NewSession(sql.NewBaseSession(), nil), + conn: conn, + currentCat: "memory", + currentSch: "main", + }} + ctx := sql.NewContext(context.Background(), sql.WithSession(holder)) + return ctx, holder, mock +} + +func newLeaseTxnTestContext( + t *testing.T, + matcher sqlmock.QueryMatcher, +) (*sql.Context, *leaseHolder, sqlmock.Sqlmock, *stdsql.Tx) { + t.Helper() + ctx, holder, mock := newLeaseTestContext(t, matcher) + mock.ExpectBegin() + tx, err := holder.conn.BeginTx(context.Background(), nil) + require.NoError(t, err) + t.Cleanup(func() { _ = tx.Rollback() }) + holder.activeTx = tx + return ctx, holder, mock, tx +} + +func requireTransitionWaiting(t *testing.T, waiter *leaseTransitionWaiter) { + t.Helper() + select { + case <-waiter.done: + t.Fatal("lifecycle transition completed before the execution lease was released") + default: + } +} + +func requireSignal(t *testing.T, signal <-chan struct{}, description string) { + t.Helper() + timer := time.NewTimer(5 * time.Second) + defer timer.Stop() + select { + case <-signal: + case <-timer.C: + t.Fatalf("timed out waiting for %s", description) + } +} + +func TestGetCatalogExecutorRequestsCatalogOnlySnapshot(t *testing.T) { + db, _, err := sqlmock.New() + require.NoError(t, err) + t.Cleanup(func() { _ = db.Close() }) + conn, err := db.Conn(context.Background()) + require.NoError(t, err) + t.Cleanup(func() { _ = conn.Close() }) + + holder := &snapshotHolder{ + Session: memory.NewSession(sql.NewBaseSession(), nil), + conn: conn, + currentCat: "memory", + currentSch: "main", + } + ctx := sql.NewContext(context.Background(), sql.WithSession(holder)) + execer, err := adapter.GetCatalogExecutor(ctx) + require.NoError(t, err) + require.Same(t, conn, execer) + require.Equal(t, []bool{true}, holder.queries) + + _, err = adapter.GetSQLExecutor(ctx) + require.NoError(t, err) + require.Equal(t, []bool{true, false}, holder.queries) +} + +func TestGetExecutionSnapshotRejectsLegacyOwnerlessTransaction(t *testing.T) { + holder := &legacyHolder{ + Session: memory.NewSession(sql.NewBaseSession(), nil), + activeTx: new(stdsql.Tx), + } + ctx := sql.NewContext(context.Background(), sql.WithSession(holder)) + + execer, conn, tx, err := adapter.GetExecutionSnapshot(ctx) + require.ErrorContains(t, err, "active transaction has no physical owner") + require.Nil(t, execer) + require.Nil(t, conn) + require.Same(t, holder.activeTx, tx) + require.Zero(t, holder.getConnCalls) +} + +func TestGetExecutionSnapshotRejectsOwnerlessLegacyTransactionAfterAcquire(t *testing.T) { + holder := &legacyHolder{Session: memory.NewSession(sql.NewBaseSession(), nil)} + holder.onGetConn = func() { holder.activeTx = new(stdsql.Tx) } + ctx := sql.NewContext(context.Background(), sql.WithSession(holder)) + + execer, conn, tx, err := adapter.GetExecutionSnapshot(ctx) + require.ErrorContains(t, err, "active transaction has no physical owner") + require.Nil(t, execer) + require.Nil(t, conn) + require.Same(t, holder.activeTx, tx) + require.Equal(t, 1, holder.getConnCalls) +} + +func TestGetExecutionSnapshotWithLeaseReleasesExactlyOnce(t *testing.T) { + db, _, err := sqlmock.New() + require.NoError(t, err) + t.Cleanup(func() { _ = db.Close() }) + conn, err := db.Conn(context.Background()) + require.NoError(t, err) + t.Cleanup(func() { _ = conn.Close() }) + + holder := &leaseHolder{snapshotHolder: &snapshotHolder{ + Session: memory.NewSession(sql.NewBaseSession(), nil), + conn: conn, + currentCat: "memory", + currentSch: "main", + }} + ctx := sql.NewContext(context.Background(), sql.WithSession(holder)) + execer, owner, tx, release, err := adapter.GetExecutionSnapshotWithLease(ctx) + require.NoError(t, err) + require.Same(t, conn, execer) + require.Same(t, conn, owner) + require.Nil(t, tx) + require.NotNil(t, release) + release() + release() + require.Equal(t, int32(1), holder.releaseCalls.Load()) + require.Equal(t, []bool{false}, holder.queries) +} + +func TestGetExecutionSnapshotWithLeaseReleasesInvalidLeaseOnError(t *testing.T) { + leaseErr := errors.New("snapshot failed") + holder := &leaseHolder{ + snapshotHolder: &snapshotHolder{Session: memory.NewSession(sql.NewBaseSession(), nil)}, + leaseErr: leaseErr, + } + ctx := sql.NewContext(context.Background(), sql.WithSession(holder)) + execer, owner, tx, release, err := adapter.GetExecutionSnapshotWithLease(ctx) + require.ErrorIs(t, err, leaseErr) + require.Nil(t, execer) + require.Nil(t, owner) + require.Nil(t, tx) + require.NotNil(t, release) + release() + require.Equal(t, int32(1), holder.releaseCalls.Load(), "the adapter must release an admitted lease on error") +} + +func TestExecInTxnKeepsLeaseUntilDriverCallReturns(t *testing.T) { + ctx, holder, mock, _ := newLeaseTxnTestContext(t, nil) + waiter := newLeaseTransitionWaiter() + holder.onRelease = waiter.release + + driverStarted := make(chan struct{}) + allowDriver := make(chan struct{}) + var allowOnce sync.Once + t.Cleanup(func() { allowOnce.Do(func() { close(allowDriver) }) }) + mock.ExpectExec("UPDATE lease_test SET value = ?"). + WithArgs(&blockingDriverArgument{started: driverStarted, allow: allowDriver}). + WillReturnResult(sqlmock.NewResult(0, 1)) + + type execResult struct { + result stdsql.Result + err error + } + execDone := make(chan execResult, 1) + go func() { + result, err := adapter.ExecInTxn(ctx, "UPDATE lease_test SET value = ?", 7) + execDone <- execResult{result: result, err: err} + }() + + requireSignal(t, driverStarted, "driver ExecContext entry") + waiter.start() + requireTransitionWaiting(t, waiter) + require.Zero(t, holder.releaseCalls.Load()) + + allowOnce.Do(func() { close(allowDriver) }) + var completed execResult + select { + case completed = <-execDone: + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for ExecInTxn to finish") + } + require.NoError(t, completed.err) + require.NotNil(t, completed.result) + affected, err := completed.result.RowsAffected() + require.NoError(t, err) + require.Equal(t, int64(1), affected) + requireSignal(t, waiter.done, "lifecycle transition after ExecContext") + require.Equal(t, int32(1), holder.releaseCalls.Load()) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestLeasedRowsKeepLeaseUntilEOFOrExplicitClose(t *testing.T) { + tests := []struct { + name string + catalogOnly bool + finish func(*testing.T, *adapter.LeasedRows) + }{ + { + name: "eof", + finish: func(t *testing.T, rows *adapter.LeasedRows) { + var values []int + for rows.Next() { + var value int + require.NoError(t, rows.Scan(&value)) + values = append(values, value) + } + require.NoError(t, rows.Err()) + require.Equal(t, []int{1, 2}, values) + }, + }, + { + name: "explicit close", + catalogOnly: true, + finish: func(t *testing.T, rows *adapter.LeasedRows) { + require.NoError(t, rows.Close()) + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + ctx, holder, mock := newLeaseTestContext(t, nil) + waiter := newLeaseTransitionWaiter() + holder.onRelease = waiter.release + const query = "SELECT value FROM lease_rows" + mock.ExpectQuery(query).WillReturnRows( + sqlmock.NewRows([]string{"value"}).AddRow(1).AddRow(2), + ) + + var ( + rows *adapter.LeasedRows + err error + ) + if test.catalogOnly { + rows, err = adapter.QueryCatalog(ctx, query) + } else { + rows, err = adapter.Query(ctx, query) + } + require.NoError(t, err) + require.NotNil(t, rows) + t.Cleanup(func() { _ = rows.Close() }) + + waiter.start() + requireTransitionWaiting(t, waiter) + require.Zero(t, holder.releaseCalls.Load()) + + test.finish(t, rows) + requireSignal(t, waiter.done, "lifecycle transition after rows termination") + require.NoError(t, rows.Close()) + require.Equal(t, int32(1), holder.releaseCalls.Load()) + require.Equal(t, []bool{test.catalogOnly}, holder.queries) + require.NoError(t, mock.ExpectationsWereMet()) + }) + } +} + +func TestLeasedRowKeepsLeaseUntilScan(t *testing.T) { + ctx, holder, mock := newLeaseTestContext(t, nil) + waiter := newLeaseTransitionWaiter() + holder.onRelease = waiter.release + const query = "SELECT value FROM lease_row" + mock.ExpectQuery(query).WillReturnRows(sqlmock.NewRows([]string{"value"}).AddRow(42)) + + row := adapter.QueryRow(ctx, query) + require.NotNil(t, row) + waiter.start() + requireTransitionWaiting(t, waiter) + require.Zero(t, holder.releaseCalls.Load()) + + var value int + require.NoError(t, row.Scan(&value)) + require.Equal(t, 42, value) + requireSignal(t, waiter.done, "lifecycle transition after Row.Scan") + require.Equal(t, int32(1), holder.releaseCalls.Load()) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestLeasedRowPreservesExecutionLeaseAdmissionError(t *testing.T) { + for _, test := range []struct { + name string + catalogOnly bool + queryRow func(*sql.Context) *adapter.LeasedRow + }{ + { + name: "ordinary query", + queryRow: func(ctx *sql.Context) *adapter.LeasedRow { + return adapter.QueryRow(ctx, "SELECT must_not_execute") + }, + }, + { + name: "catalog query", + catalogOnly: true, + queryRow: func(ctx *sql.Context) *adapter.LeasedRow { + return adapter.QueryRowCatalog(ctx, "SELECT must_not_execute") + }, + }, + } { + t.Run(test.name, func(t *testing.T) { + ctx, holder, mock := newLeaseTestContext(t, nil) + admissionErr := errors.New("execution lease admission failed") + holder.leaseErr = admissionErr + + row := test.queryRow(ctx) + require.NotNil(t, row) + var value int + require.ErrorIs(t, row.Scan(&value), admissionErr) + require.ErrorIs(t, row.Scan(&value), admissionErr) + require.Equal(t, int32(1), holder.releaseCalls.Load()) + require.Equal(t, []bool{test.catalogOnly}, holder.queries) + // There is deliberately no Query expectation. Any target driver call + // would replace the admission sentinel with sqlmock's unexpected-call error. + require.NoError(t, mock.ExpectationsWereMet()) + }) + } +} + +func TestTxnExecutionLeaseCoversPrepareContext(t *testing.T) { + driverStarted := make(chan struct{}) + allowDriver := make(chan struct{}) + var allowOnce sync.Once + t.Cleanup(func() { allowOnce.Do(func() { close(allowDriver) }) }) + matcher := &blockingQueryMatcher{ + started: driverStarted, + allow: allowDriver, + delegate: sqlmock.QueryMatcherEqual, + } + ctx, holder, mock, activeTx := newLeaseTxnTestContext(t, matcher) + waiter := newLeaseTransitionWaiter() + holder.onRelease = waiter.release + const query = "INSERT INTO lease_prepare VALUES (?)" + mock.ExpectPrepare(query) + + type prepareResult struct { + owner *stdsql.Conn + tx *stdsql.Tx + err error + } + prepareDone := make(chan prepareResult, 1) + go func() { + owner, tx, release, err := adapter.GetTxnExecutionSnapshotWithLease(ctx, nil) + if err != nil { + prepareDone <- prepareResult{owner: owner, tx: tx, err: err} + return + } + stmt, prepareErr := tx.PrepareContext(ctx, query) + if stmt != nil { + prepareErr = errors.Join(prepareErr, stmt.Close()) + } + release() + prepareDone <- prepareResult{owner: owner, tx: tx, err: prepareErr} + }() + + requireSignal(t, driverStarted, "driver PrepareContext entry") + waiter.start() + requireTransitionWaiting(t, waiter) + require.Zero(t, holder.releaseCalls.Load()) + + allowOnce.Do(func() { close(allowDriver) }) + var completed prepareResult + select { + case completed = <-prepareDone: + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for PrepareContext to finish") + } + require.NoError(t, completed.err) + require.Same(t, holder.conn, completed.owner) + require.Same(t, activeTx, completed.tx) + requireSignal(t, waiter.done, "lifecycle transition after PrepareContext") + require.Equal(t, int32(1), holder.releaseCalls.Load()) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestExecInTxnAdmissionFailureStartsNoDriverCall(t *testing.T) { + ctx, holder, mock, _ := newLeaseTxnTestContext(t, nil) + admissionErr := errors.New("execution lease admission failed") + holder.leaseErr = admissionErr + + result, err := adapter.ExecInTxn(ctx, "UPDATE must_not_execute SET value = 1") + require.ErrorIs(t, err, admissionErr) + require.Nil(t, result) + require.Equal(t, int32(1), holder.releaseCalls.Load()) + // There is deliberately no Exec expectation. Any target driver call would + // replace the admission sentinel with sqlmock's unexpected-call error. + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestExecInTxnBindingMismatchReleasesLeaseBeforeDriverCall(t *testing.T) { + ctx, holder, mock := newLeaseTestContext(t, nil) + holder.activeTx = new(stdsql.Tx) + holder.leasedTx = new(stdsql.Tx) + require.NotSame(t, holder.activeTx, holder.leasedTx) + + result, err := adapter.ExecInTxn(ctx, "UPDATE must_not_execute SET value = 1") + require.ErrorIs(t, err, adapter.ErrTransactionBindingChanged) + require.Nil(t, result) + require.Equal(t, int32(1), holder.releaseCalls.Load()) + // Identity validation precedes driver dispatch, so no target expectation is + // installed or consumed when the transaction generation has changed. + require.NoError(t, mock.ExpectationsWereMet()) +} diff --git a/backend/completion_registry_test.go b/backend/completion_registry_test.go new file mode 100644 index 00000000..805407ed --- /dev/null +++ b/backend/completion_registry_test.go @@ -0,0 +1,182 @@ +package backend + +import ( + stdsql "database/sql" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func newPostPhysicalRegistrySessionForTest() *Session { + return &Session{postPhysical: make(map[*stdsql.Tx]*postTransactionCleanup)} +} + +func TestPostPhysicalTombstonesBoundedAndPendingRetained(t *testing.T) { + sess := newPostPhysicalRegistrySessionForTest() + + pendingTx := new(stdsql.Tx) + pendingState := &postTransactionCleanup{} + sess.postPhysical[pendingTx] = pendingState + var pendingCalls atomic.Int32 + require.True(t, pendingState.register(func(bool) { + pendingCalls.Add(1) + })) + + completed := make([]*stdsql.Tx, maxPostPhysicalTombstones+8) + for i := range completed { + completed[i] = new(stdsql.Tx) + sess.completePhysicalPostCleanup(completed[i], i%2 == 0) + } + + sess.postPhysicalMu.Lock() + mapLen := len(sess.postPhysical) + orderLen := len(sess.postPhysicalOrder) + _, pendingRetained := sess.postPhysical[pendingTx] + _, oldestRetained := sess.postPhysical[completed[0]] + _, newestRetained := sess.postPhysical[completed[len(completed)-1]] + sess.postPhysicalMu.Unlock() + + require.LessOrEqual(t, mapLen, maxPostPhysicalTombstones+1) + require.LessOrEqual(t, orderLen, maxPostPhysicalTombstones) + require.True(t, pendingRetained, "pending callback state must never be evicted") + require.False(t, oldestRetained, "old completed tombstones should be evicted") + require.True(t, newestRetained) + require.Zero(t, pendingCalls.Load()) +} + +func TestPostPhysicalTTLAndLateRegistration(t *testing.T) { + sess := newPostPhysicalRegistrySessionForTest() + + recent := new(stdsql.Tx) + sess.completePhysicalPostCleanup(recent, true) + var recentSuccess atomic.Bool + require.True(t, sess.RegisterPostPhysicalTransactionCleanup(nil, recent, func(success bool) { + recentSuccess.Store(success) + })) + require.True(t, recentSuccess.Load(), "a recent tombstone must preserve its outcome") + + expired := new(stdsql.Tx) + sess.completePhysicalPostCleanup(expired, true) + sess.postPhysicalMu.Lock() + state := sess.postPhysical[expired] + sess.postPhysicalMu.Unlock() + require.NotNil(t, state) + state.mu.Lock() + state.completedAt = time.Now().Add(-postPhysicalTombstoneTTL - time.Second) + state.mu.Unlock() + + var expiredCalls atomic.Int32 + require.False(t, sess.RegisterPostPhysicalTransactionCleanup(nil, expired, func(success bool) { + expiredCalls.Add(1) + require.False(t, success) + })) + require.Zero(t, expiredCalls.Load(), "false registration leaves fallback release to the caller") + + sess.postPhysicalMu.Lock() + _, stillRetained := sess.postPhysical[expired] + sess.postPhysicalMu.Unlock() + require.False(t, stillRetained) +} + +func TestPostPhysicalRegistrationWithoutFinalizerFailsClosed(t *testing.T) { + sess := newPostPhysicalRegistrySessionForTest() + tx := new(stdsql.Tx) + var calls atomic.Int32 + require.False(t, sess.RegisterPostPhysicalTransactionCleanup(nil, tx, func(success bool) { + calls.Add(1) + require.False(t, success) + })) + // A false registration result leaves the caller responsible for fallback + // release; the callback itself must not be invoked as well. + require.Zero(t, calls.Load()) + sess.postPhysicalMu.Lock() + _, retained := sess.postPhysical[tx] + sess.postPhysicalMu.Unlock() + require.False(t, retained) +} + +func TestPostTransactionCleanupDrainsCallbacksBeforeRepanic(t *testing.T) { + state := new(postTransactionCleanup) + var firstCalls atomic.Int32 + var secondCalls atomic.Int32 + require.True(t, state.register(func(bool) { + firstCalls.Add(1) + panic("post-transaction cleanup panic") + })) + require.True(t, state.register(func(bool) { + secondCalls.Add(1) + })) + + require.PanicsWithValue(t, "post-transaction cleanup panic", func() { + state.complete(false) + }) + require.Equal(t, int32(1), firstCalls.Load()) + require.Equal(t, int32(1), secondCalls.Load()) + require.True(t, state.isCompleted()) +} + +func TestPostPhysicalTombstoneWaitsForCallbackDrain(t *testing.T) { + tx := new(stdsql.Tx) + state := new(postTransactionCleanup) + sess := &Session{postPhysical: map[*stdsql.Tx]*postTransactionCleanup{tx: state}} + + callbackStarted := make(chan struct{}) + releaseCallback := make(chan struct{}) + defer func() { + select { + case <-releaseCallback: + default: + close(releaseCallback) + } + }() + require.True(t, state.register(func(bool) { + close(callbackStarted) + <-releaseCallback + })) + + completionDone := make(chan struct{}) + go func() { + defer close(completionDone) + sess.completePhysicalPostCleanup(tx, true) + }() + select { + case <-callbackStarted: + case <-time.After(time.Second): + t.Fatal("physical completion callback did not start") + } + + var lateCalls atomic.Int32 + require.True(t, sess.RegisterPostPhysicalTransactionCleanup(nil, tx, func(success bool) { + require.True(t, success) + lateCalls.Add(1) + })) + require.Equal(t, int32(1), lateCalls.Load()) + + // The terminal outcome is visible to late registration, but the original + // callback batch is still in flight. It is not a pruneable tombstone yet. + sess.postPhysicalMu.Lock() + sess.prunePostPhysicalCompletionsLocked(time.Now().Add(postPhysicalTombstoneTTL + time.Second)) + _, retained := sess.postPhysical[tx] + orderLenBeforeDrain := len(sess.postPhysicalOrder) + sess.postPhysicalMu.Unlock() + require.True(t, retained) + require.Zero(t, orderLenBeforeDrain) + + close(releaseCallback) + select { + case <-completionDone: + case <-time.After(time.Second): + t.Fatal("physical completion did not finish") + } + + sess.postPhysicalMu.Lock() + recorded := sess.postPhysical[tx] + orderLenAfterDrain := len(sess.postPhysicalOrder) + sess.postPhysicalMu.Unlock() + require.Same(t, state, recorded) + completed, _ := state.completionSnapshot() + require.True(t, completed) + require.Equal(t, 1, orderLenAfterDrain) +} diff --git a/backend/ducklake_operation.go b/backend/ducklake_operation.go new file mode 100644 index 00000000..2bae7aed --- /dev/null +++ b/backend/ducklake_operation.go @@ -0,0 +1,531 @@ +package backend + +import ( + stdsql "database/sql" + "errors" + "io" + "sync" + "sync/atomic" + + "github.com/dolthub/go-mysql-server/sql" +) + +// duckLakeOperationIter keeps statement resources alive until the wrapped +// child is finished. Snapshot and provider-operation leases are separate: +// the snapshot lease ends as soon as the child closes, while an operation lease +// may remain held until an implicit transaction's commit/rollback completes. +type duckLakeOperationIter struct { + state *duckLakeOperationState +} + +// GMS commits autocommit transactions from TransactionCommittingIter.Close. +// An iterator abandoned before EOF must therefore surface an error from Close, +// otherwise a short-circuited operation can be committed after its child has +// already been retired. +var errDuckLakeOperationClosedBeforeEOF = errors.New("ducklake operation iterator closed before EOF") + +var errDuckLakeOperationValueIterUnavailable = errors.New("ducklake operation child does not support ValueRowIter") + +// duckLakeOperationState is shared by every mutable alias returned by +// WithChildIter. In particular, aliases never retain their own child pointer: +// GMS may call Next on an older alias after a rewrite replaced the child. +type duckLakeOperationState struct { + mu sync.Mutex + current sql.RowIter + + snapshotRelease func() + operationRelease func() + snapshotOnce sync.Once + operationOnce sync.Once + closeOnce sync.Once + closeErr error + closeCtx *sql.Context + terminalErr error + + closed atomic.Bool + terminal atomic.Bool + failed atomic.Bool + + // cleanupHandle belongs to the pre-finalization implicit cleanup hook. It + // is deliberately called with false by finish so the hook remains armed + // until the session's commit/rollback code decides the final outcome. + cleanupHandle func(bool) + operationDeferred atomic.Bool +} + +func newDuckLakeOperationIter(child sql.RowIter, release func()) *duckLakeOperationIter { + return newDuckLakeOperationIterWithLeases(child, release, nil) +} + +func newDuckLakeOperationIterWithLeases( + child sql.RowIter, + snapshotRelease func(), + operationRelease func(), +) *duckLakeOperationIter { + return &duckLakeOperationIter{ + state: &duckLakeOperationState{ + current: child, + snapshotRelease: snapshotRelease, + operationRelease: operationRelease, + }, + } +} + +func (i *duckLakeOperationIter) currentChild() sql.RowIter { + if i == nil || i.state == nil || i.state.closed.Load() { + return nil + } + i.state.mu.Lock() + child := i.state.current + i.state.mu.Unlock() + return child +} + +func (i *duckLakeOperationIter) releaseSnapshot() { + if i == nil || i.state == nil || i.state.snapshotRelease == nil { + return + } + i.state.snapshotOnce.Do(i.state.snapshotRelease) +} + +func (i *duckLakeOperationIter) releaseOperation() { + if i == nil || i.state == nil || i.state.operationRelease == nil { + return + } + i.state.operationOnce.Do(i.state.operationRelease) +} + +// done is retained for callers of the old wrapper. It now means release all +// remaining leases, and is only used on construction/fallback paths where no +// transaction post-finalization hook owns the operation lease. +func (i *duckLakeOperationIter) done() { + if i == nil || i.state == nil { + return + } + i.releaseSnapshot() + if !i.state.operationDeferred.Load() { + i.releaseOperation() + } +} + +// closeChild closes the shared current child at most once. The current child +// is captured and cleared while holding the state mutex so GetChildIter and +// every alias observe the same ownership transition. +func (i *duckLakeOperationIter) closeChild(ctx *sql.Context) error { + if i == nil || i.state == nil { + return nil + } + i.state.closeOnce.Do(func() { + i.state.mu.Lock() + i.state.closed.Store(true) + i.state.closeCtx = ctx + child := i.state.current + i.state.current = nil + i.state.mu.Unlock() + + var closeErr error + if child != nil { + closeErr = child.Close(ctx) + } + i.state.mu.Lock() + i.state.closeErr = closeErr + i.state.mu.Unlock() + }) + i.state.mu.Lock() + defer i.state.mu.Unlock() + return i.state.closeErr +} + +func (i *duckLakeOperationIter) setTerminalError(err error) { + if i == nil || i.state == nil || err == nil || errors.Is(err, io.EOF) { + return + } + i.state.mu.Lock() + i.state.terminalErr = errors.Join(i.state.terminalErr, err) + i.state.mu.Unlock() + i.state.failed.Store(true) +} + +func (i *duckLakeOperationIter) terminalError() error { + if i == nil || i.state == nil { + return nil + } + i.state.mu.Lock() + defer i.state.mu.Unlock() + return i.state.terminalErr +} + +// finish closes/releases the child and updates the implicit-transaction +// registration. A false result deliberately leaves the registration armed so +// GMS's SetTransaction(nil) error boundary can run the callback and rollback. +func (i *duckLakeOperationIter) finish(ctx *sql.Context, success bool) error { + if i == nil || i.state == nil { + return nil + } + if !success { + i.state.failed.Store(true) + } + closeErr := i.closeChild(ctx) + if closeErr != nil { + // A child close failure is a statement failure. Keeping failed=true makes + // TransactionCommittingIter skip an implicit commit and take rollback. + i.state.failed.Store(true) + } + i.releaseSnapshot() + if i.state.cleanupHandle != nil { + i.state.cleanupHandle(false) + } + if !i.state.operationDeferred.Load() { + i.releaseOperation() + } + return errors.Join(i.terminalError(), closeErr) +} + +func (i *duckLakeOperationIter) Next(ctx *sql.Context) (sql.Row, error) { + if i == nil || i.state == nil { + return nil, io.EOF + } + if terminalErr := i.terminalError(); terminalErr != nil { + return nil, terminalErr + } + if i.state.terminal.Load() { + return nil, io.EOF + } + child := i.currentChild() + if child == nil { + _ = i.finish(ctx, true) + return nil, io.EOF + } + row, err := child.Next(ctx) + if err != nil { + isEOF := errors.Is(err, io.EOF) + if !isEOF { + i.setTerminalError(err) + } + i.state.terminal.Store(isEOF) + if closeErr := i.finish(ctx, isEOF); closeErr != nil { + err = errors.Join(err, closeErr) + } + } + return row, err +} + +func (i *duckLakeOperationIter) Close(ctx *sql.Context) error { + if i == nil || i.state == nil { + return nil + } + // A close before the child reported EOF is an aborted/short-circuited + // execution. Mark it failed so a GMS implicit transaction cannot commit it. + if !i.state.terminal.Load() && !i.state.failed.Load() { + i.setTerminalError(errDuckLakeOperationClosedBeforeEOF) + } + success := i.state.terminal.Load() && !i.state.failed.Load() + return i.finish(ctx, success) +} + +// GetChildIter and WithChildIter preserve the mutable-iterator contract used +// by GMS projection rewrites. The returned child is always the shared current +// child, so an older alias cannot continue driving an abandoned iterator. +func (i *duckLakeOperationIter) GetChildIter() sql.RowIter { + if i == nil || i.state == nil { + return nil + } + i.state.mu.Lock() + defer i.state.mu.Unlock() + return i.state.current +} + +func (i *duckLakeOperationIter) WithChildIter(child sql.RowIter) sql.RowIter { + if i == nil || i.state == nil { + return child + } + i.state.mu.Lock() + if i.state.closed.Load() { + // A rewrite after terminal close cannot become the owner of this lease. + // Close the supplied child immediately so a late rewrite cannot leak a + // resource after the shared lease has already finished. + closeCtx := i.state.closeCtx + i.state.mu.Unlock() + if child != nil { + _ = child.Close(closeCtx) + } + return child + } + i.state.current = child + i.state.mu.Unlock() + return i.aliasForChild(child) +} + +func (i *duckLakeOperationIter) aliasForChild(child sql.RowIter) sql.RowIter { + if _, ok := child.(sql.ValueRowIter); ok { + return &duckLakeOperationValueIter{ + duckLakeOperationIter: &duckLakeOperationIter{state: i.state}, + } + } + return &duckLakeOperationIter{state: i.state} +} + +// duckLakeOperationValueIter is selected only when the wrapped child supports +// sql.ValueRowIter. Its NextValueRow method dynamically resolves the shared +// current child; it intentionally stores no child pointer of its own. +type duckLakeOperationValueIter struct { + *duckLakeOperationIter +} + +func (i *duckLakeOperationValueIter) currentValueChild() sql.ValueRowIter { + if i == nil || i.duckLakeOperationIter == nil { + return nil + } + child := i.currentChild() + valueChild, _ := child.(sql.ValueRowIter) + return valueChild +} + +func (i *duckLakeOperationValueIter) NextValueRow(ctx *sql.Context) (sql.ValueRow, error) { + if i == nil || i.state == nil { + return nil, io.EOF + } + if terminalErr := i.terminalError(); terminalErr != nil { + return nil, terminalErr + } + if i.state.terminal.Load() { + return nil, io.EOF + } + valueChild := i.currentValueChild() + if valueChild == nil { + if i.currentChild() != nil { + err := errDuckLakeOperationValueIterUnavailable + i.setTerminalError(err) + _ = i.finish(ctx, false) + return nil, err + } + _ = i.finish(ctx, true) + return nil, io.EOF + } + row, err := valueChild.NextValueRow(ctx) + if err != nil { + isEOF := errors.Is(err, io.EOF) + if !isEOF { + i.setTerminalError(err) + } + i.state.terminal.Store(isEOF) + if closeErr := i.finish(ctx, isEOF); closeErr != nil { + err = errors.Join(err, closeErr) + } + } + return row, err +} + +func (i *duckLakeOperationValueIter) IsValueRowIter(ctx *sql.Context) bool { + valueChild := i.currentValueChild() + return valueChild != nil && valueChild.IsValueRowIter(ctx) +} + +func (i *duckLakeOperationValueIter) WithChildIter(child sql.RowIter) sql.RowIter { + if i == nil || i.state == nil { + return child + } + return i.duckLakeOperationIter.WithChildIter(child) +} + +func registerPostTransactionCleanup( + ctx *sql.Context, + tx sql.Transaction, + release func(success bool), +) bool { + if ctx == nil || ctx.Session == nil || tx == nil || release == nil { + return false + } + registrar, ok := ctx.Session.(PostTransactionCleanupRegistrar) + if !ok { + return false + } + return registrar.RegisterPostTransactionCleanup(ctx, tx, release) +} + +// RegisterPostPhysicalTransactionCleanupForContext lets protocol adapters +// defer an operation release until a raw database/sql transaction has been +// finalized. It is deliberately optional so narrow session doubles retain the +// legacy iterator-close behavior. +func RegisterPostPhysicalTransactionCleanupForContext( + ctx *sql.Context, + tx *stdsql.Tx, + release func(success bool), +) bool { + if ctx == nil || ctx.Session == nil || tx == nil || release == nil { + return false + } + registrar, ok := ctx.Session.(PhysicalPostTransactionCleanupRegistrar) + if !ok { + return false + } + return registrar.RegisterPostPhysicalTransactionCleanup(ctx, tx, release) +} + +// implicitTransactionCleanupState is an optional session-side probe used to +// distinguish a GMS implicit autocommit transaction from an explicit or +// adapter-owned transaction. RegisterImplicitTransactionCleanup runs a +// callback immediately when its marker is already gone, which is correct for +// a racing implicit rollback but would incorrectly retire a child belonging +// to an explicit transaction (or a MySQL DDL transaction that deliberately +// has no implicit marker). +type implicitTransactionCleanupState interface { + HasImplicitTransactionCleanup(sql.Transaction) bool +} + +func hasImplicitTransactionCleanup(ctx *sql.Context, tx sql.Transaction) bool { + if ctx == nil || ctx.Session == nil || tx == nil { + return false + } + if state, ok := ctx.Session.(implicitTransactionCleanupState); ok { + return state.HasImplicitTransactionCleanup(tx) + } + // Legacy session implementations do not expose marker state. Preserve the + // historical registration behavior for those adapters. + return true +} + +func wrapDuckLakeOperationIter(iter sql.RowIter, release func()) sql.RowIter { + return wrapDuckLakeOperationIterWithLeases(nil, iter, release, nil) +} + +// wrapDuckLakeOperationIterWithContext preserves the historical one-release +// contract. New execution paths that acquire separate leases should call +// wrapDuckLakeOperationIterWithLeases instead. +func wrapDuckLakeOperationIterWithContext(ctx *sql.Context, iter sql.RowIter, release func()) sql.RowIter { + return wrapDuckLakeOperationIterWithLeases(ctx, iter, release, nil) +} + +// wrapDuckLakeOperationIterWithLeases wraps an iterator with independent +// snapshot and provider-operation releases. For an implicit transaction, the +// operation release is handed to the session's post-finalization hook while +// snapshot release remains tied to child close. If no such hook accepts the +// transaction, both releases retain the legacy child-close boundary. +func wrapDuckLakeOperationIterWithLeases( + ctx *sql.Context, + iter sql.RowIter, + snapshotRelease func(), + operationRelease func(), +) sql.RowIter { + return wrapDuckLakeOperationIterWithLeasesAndTx( + ctx, + iter, + nil, + snapshotRelease, + operationRelease, + ) +} + +// wrapDuckLakeOperationIterWithLeasesAndTx is the implementation shared by +// the legacy and binding-aware wrappers. Keeping the old four-argument entry +// point above preserves existing GMS/test adapters while execution callers can +// carry the atomic physical transaction explicitly. The physical transaction +// is passed through from that atomic snapshot instead of being looked up again +// while snapshotRelease still holds the pool lifecycle reader. A second +// ordinary GetTxnBinding call can block forever once pool shutdown has queued +// its writer (the lifecycle gate rejects new readers during shutdown). +// +// physicalTx may be nil when the snapshot selected an idle connection. In that +// case a logical GMS transaction, when present, remains the fallback callback +// boundary. Callers using this entry point must pass the exact transaction +// returned by adapter.GetExecutionSnapshotWithLease; the function never guesses +// a replacement binding. +func wrapDuckLakeOperationIterWithLeasesAndTx( + ctx *sql.Context, + iter sql.RowIter, + physicalTx *stdsql.Tx, + snapshotRelease func(), + operationRelease func(), +) sql.RowIter { + if iter == nil { + if snapshotRelease != nil { + snapshotRelease() + } + if operationRelease != nil { + operationRelease() + } + return nil + } + base := newDuckLakeOperationIterWithLeases(iter, snapshotRelease, operationRelease) + + if ctx != nil { + if operationRelease != nil { + // A transaction-backed session can expose the raw physical owner as + // well as the logical GMS transaction. The owner is supplied by the + // execution snapshot, so registration never re-enters the pool's + // ordinary lifecycle reader while snapshotRelease is held. + deferred := false + if physicalTx != nil { + deferred = RegisterPostPhysicalTransactionCleanupForContext(ctx, physicalTx, func(bool) { + base.releaseOperation() + }) + } + if !deferred && ctx.GetTransaction() != nil { + deferred = registerPostTransactionCleanup(ctx, ctx.GetTransaction(), func(bool) { + base.releaseOperation() + }) + } + base.state.operationDeferred.Store(deferred) + } + // This pre-finalization callback closes the child and snapshot if GMS + // clears an implicit transaction before the iterator reaches its normal + // terminal path. It intentionally does not release a deferred operation. + transaction := ctx.GetTransaction() + if transaction != nil && hasImplicitTransactionCleanup(ctx, transaction) { + base.state.cleanupHandle = RegisterImplicitTransactionCleanupForContext(ctx, transaction, func() error { + err := base.closeChild(ctx) + if err == nil { + base.releaseSnapshot() + } else { + base.state.failed.Store(true) + base.releaseSnapshot() + } + if !base.state.operationDeferred.Load() { + base.releaseOperation() + } + return errors.Join(err, base.terminalError()) + }) + } + } + + if _, ok := iter.(sql.ValueRowIter); ok { + return &duckLakeOperationValueIter{duckLakeOperationIter: base} + } + return base +} + +// WrapRowIterWithLeases is the legacy explicit split entry point. The snapshot +// lease is released at child termination; operationRelease is deferred until a +// logical transaction's post-finalization hook when one is available. Callers +// that already hold an atomic physical transaction should use +// WrapRowIterWithLeasesAndTx so wrapper construction never performs a binding +// lookup. +func WrapRowIterWithLeases( + ctx *sql.Context, + iter sql.RowIter, + snapshotRelease func(), + operationRelease func(), +) sql.RowIter { + return wrapDuckLakeOperationIterWithLeases(ctx, iter, snapshotRelease, operationRelease) +} + +// WrapRowIterWithLeasesAndTx is the binding-aware exported entry point for +// protocol adapters that already hold an atomic execution snapshot. Passing +// physicalTx avoids a second lifecycle lookup while the snapshot lease is +// active; the legacy four-argument function remains available for adapters +// that do not expose a physical transaction. +func WrapRowIterWithLeasesAndTx( + ctx *sql.Context, + iter sql.RowIter, + physicalTx *stdsql.Tx, + snapshotRelease func(), + operationRelease func(), +) sql.RowIter { + return wrapDuckLakeOperationIterWithLeasesAndTx( + ctx, + iter, + physicalTx, + snapshotRelease, + operationRelease, + ) +} diff --git a/backend/ducklake_operation_test.go b/backend/ducklake_operation_test.go new file mode 100644 index 00000000..43f8aed5 --- /dev/null +++ b/backend/ducklake_operation_test.go @@ -0,0 +1,571 @@ +package backend + +import ( + "context" + stdsql "database/sql" + "errors" + "io" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/apecloud/myduckserver/adapter" + "github.com/apecloud/myduckserver/mycontext" + "github.com/dolthub/go-mysql-server/memory" + "github.com/dolthub/go-mysql-server/sql" + "github.com/dolthub/go-mysql-server/sql/rowexec" + "github.com/stretchr/testify/require" +) + +var errOperationChild = errors.New("operation child failed") + +type orderedOperationIter struct { + events *[]string + nextErr error + valueErr error + closeErr error + closeCalls int + nextCalls int + valueCalls int +} + +func (i *orderedOperationIter) Next(*sql.Context) (sql.Row, error) { + i.nextCalls++ + *i.events = append(*i.events, "next") + if i.nextErr != nil { + return nil, i.nextErr + } + return nil, io.EOF +} + +func (i *orderedOperationIter) NextValueRow(*sql.Context) (sql.ValueRow, error) { + i.valueCalls++ + *i.events = append(*i.events, "next-value") + if i.valueErr != nil { + return nil, i.valueErr + } + return nil, io.EOF +} + +func (i *orderedOperationIter) IsValueRowIter(*sql.Context) bool { return true } + +func (i *orderedOperationIter) Close(*sql.Context) error { + i.closeCalls++ + *i.events = append(*i.events, "close") + return i.closeErr +} + +type postCleanupTestSession struct { + *Session + postCleanup func(bool) +} + +// deterministicShutdownGate models the lifecycle gate's important property: +// once an exclusive teardown writer is waiting, ordinary readers cannot enter +// again. The explicit channels make that state deterministic in a unit test; +// a second GetTxnBinding call would remain blocked until the retained snapshot +// reader is released. +type deterministicShutdownGate struct { + mu sync.Mutex + cond *sync.Cond + readers int + writerWaiting bool + writerActive bool + writerReady chan struct{} + writerGot chan struct{} + writerRelease chan struct{} + writerDone chan struct{} + releaseOnce sync.Once +} + +func newDeterministicShutdownGate() *deterministicShutdownGate { + gate := &deterministicShutdownGate{ + writerReady: make(chan struct{}), + writerGot: make(chan struct{}), + writerRelease: make(chan struct{}), + writerDone: make(chan struct{}), + } + gate.cond = sync.NewCond(&gate.mu) + return gate +} + +func (gate *deterministicShutdownGate) RLock() { + gate.mu.Lock() + for gate.writerWaiting || gate.writerActive { + gate.cond.Wait() + } + gate.readers++ + gate.mu.Unlock() +} + +func (gate *deterministicShutdownGate) RUnlock() { + gate.mu.Lock() + if gate.readers > 0 { + gate.readers-- + } + gate.cond.Broadcast() + gate.mu.Unlock() +} + +func (gate *deterministicShutdownGate) startWriter() { + gate.mu.Lock() + gate.writerWaiting = true + close(gate.writerReady) + for gate.readers > 0 { + gate.cond.Wait() + } + gate.writerWaiting = false + gate.writerActive = true + close(gate.writerGot) + gate.mu.Unlock() + + <-gate.writerRelease + + gate.mu.Lock() + gate.writerActive = false + gate.cond.Broadcast() + gate.mu.Unlock() + close(gate.writerDone) +} + +func (gate *deterministicShutdownGate) releaseWriter() { + gate.releaseOnce.Do(func() { close(gate.writerRelease) }) +} + +type shutdownGateSession struct { + *memory.Session + gate *deterministicShutdownGate + conn *stdsql.Conn + tx *stdsql.Tx + + bindingCalls atomic.Int32 + postMu sync.Mutex + postTx *stdsql.Tx + postCallback func(bool) +} + +var _ adapter.ConnectionHolder = (*shutdownGateSession)(nil) +var _ adapter.ExecutionSnapshotHolder = (*shutdownGateSession)(nil) +var _ adapter.ExecutionSnapshotLeaseHolder = (*shutdownGateSession)(nil) +var _ PhysicalPostTransactionCleanupRegistrar = (*shutdownGateSession)(nil) + +func (s *shutdownGateSession) GetConn(context.Context) (*stdsql.Conn, error) { + return s.conn, nil +} + +func (s *shutdownGateSession) GetTxn(context.Context, *stdsql.TxOptions) (*stdsql.Tx, error) { + return s.tx, nil +} + +func (s *shutdownGateSession) GetCatalogConn(context.Context) (*stdsql.Conn, error) { + return s.conn, nil +} + +func (s *shutdownGateSession) GetCatalogTxn(context.Context, *stdsql.TxOptions) (*stdsql.Tx, error) { + return s.tx, nil +} + +func (s *shutdownGateSession) TryGetTxn() *stdsql.Tx { return s.tx } + +func (s *shutdownGateSession) GetCurrentCatalog() string { return "memory" } + +func (s *shutdownGateSession) GetCurrentSchema() string { return "main" } + +func (s *shutdownGateSession) CloseTxn() {} + +func (s *shutdownGateSession) CloseConn() {} + +func (s *shutdownGateSession) GetExecutionSnapshot( + context.Context, + bool, +) (*stdsql.Conn, *stdsql.Tx, error) { + return s.conn, s.tx, nil +} + +func (s *shutdownGateSession) GetExecutionSnapshotLease( + context.Context, + bool, +) (*stdsql.Conn, *stdsql.Tx, func(), error) { + s.gate.RLock() + var once sync.Once + return s.conn, s.tx, func() { + once.Do(s.gate.RUnlock) + }, nil +} + +func (s *shutdownGateSession) GetTxnBinding() (*stdsql.Conn, *stdsql.Tx) { + s.bindingCalls.Add(1) + s.gate.RLock() + defer s.gate.RUnlock() + return s.conn, s.tx +} + +func (s *shutdownGateSession) RegisterPostPhysicalTransactionCleanup( + _ *sql.Context, + tx *stdsql.Tx, + callback func(bool), +) bool { + if tx == nil || callback == nil { + return false + } + s.postMu.Lock() + s.postTx = tx + s.postCallback = callback + s.postMu.Unlock() + return true +} + +func (s *postCleanupTestSession) RegisterPostTransactionCleanup( + _ *sql.Context, + _ sql.Transaction, + callback func(bool), +) bool { + s.postCleanup = callback + return true +} + +func TestDuckLakeOperationIterClosesChildBeforeReleaseOnTerminalNext(t *testing.T) { + ctx := sql.NewEmptyContext() + events := []string{} + child := &orderedOperationIter{events: &events, nextErr: errOperationChild} + iter := wrapDuckLakeOperationIter(child, func() { events = append(events, "release") }) + + _, err := iter.Next(ctx) + require.ErrorIs(t, err, errOperationChild) + require.Equal(t, []string{"next", "close", "release"}, events) + + // The engine commonly calls Close after a terminal error. It must not close + // the child or release the operation a second time, but it must preserve the + // terminal failure so a transaction wrapper cannot mistake the close for a + // successful statement. + require.ErrorIs(t, iter.Close(ctx), errOperationChild) + require.Equal(t, 1, child.closeCalls) + require.Equal(t, []string{"next", "close", "release"}, events) +} + +func TestDuckLakeOperationIterEarlyCloseAbortsBeforeCommit(t *testing.T) { + ctx := sql.NewEmptyContext() + events := []string{} + child := &orderedOperationIter{events: &events} + iter := wrapDuckLakeOperationIter(child, func() { events = append(events, "release") }) + + require.ErrorIs(t, iter.Close(ctx), errDuckLakeOperationClosedBeforeEOF) + require.Equal(t, []string{"close", "release"}, events) + require.Equal(t, 1, child.closeCalls) + require.ErrorIs(t, iter.Close(ctx), errDuckLakeOperationClosedBeforeEOF) + require.Equal(t, 1, child.closeCalls) +} + +func TestDuckLakeOperationIterGMSDoesNotCommitAfterChildError(t *testing.T) { + sess, ctx, transaction := newFrontendImplicitSession(t) + child := &orderedOperationIter{events: new([]string), nextErr: errOperationChild} + iter := wrapDuckLakeOperationIterWithContext(ctx, child, nil) + committing, err := rowexec.AddTransactionCommittingIter(ctx, nil, iter) + require.NoError(t, err) + + _, err = committing.Next(ctx) + require.ErrorIs(t, err, errOperationChild) + // TransactionCommittingIter must observe the preserved child error from + // Close and skip Session.CommitTransaction entirely. + require.ErrorIs(t, committing.Close(ctx), errOperationChild) + require.Same(t, transaction, sess.Session.GetTransaction()) + require.NotNil(t, sess.implicit) + + // The engine's clearAutocommitOnError boundary now performs the callback and + // rollback, after the child has already been closed exactly once. + ctx.SetTransaction(nil) + require.Nil(t, sess.Session.GetTransaction()) + require.Nil(t, sess.implicit) +} + +func TestDuckLakeOperationIterGMSCommitsAfterEOF(t *testing.T) { + sess, ctx, _ := newFrontendImplicitSession(t) + child := &orderedOperationIter{events: new([]string)} + iter := wrapDuckLakeOperationIterWithContext(ctx, child, nil) + committing, err := rowexec.AddTransactionCommittingIter(ctx, nil, iter) + require.NoError(t, err) + + _, err = committing.Next(ctx) + require.ErrorIs(t, err, io.EOF) + require.NoError(t, committing.Close(ctx)) + require.Nil(t, sess.Session.GetTransaction()) + require.Nil(t, sess.implicit) +} + +func TestDuckLakeOperationIterDoesNotTreatExplicitTransactionAsImplicit(t *testing.T) { + sess, ctx, transaction := newFrontendImplicitSession(t) + // Remove only the frontend autocommit marker. The logical transaction stays + // current, modeling an explicit/DDL transaction with no implicit cleanup + // boundary. + sess.removeImplicitState(sess.implicit) + require.False(t, sess.HasImplicitTransactionCleanup(transaction)) + + events := []string{} + child := &orderedOperationIter{events: &events} + iter := wrapDuckLakeOperationIterWithLeases( + ctx, + child, + func() { events = append(events, "snapshot-release") }, + nil, + ) + // Registration must not run the callback immediately just because the + // current transaction has no implicit marker. + require.Equal(t, []string{}, events) + + _, err := iter.Next(ctx) + require.ErrorIs(t, err, io.EOF) + require.Equal(t, []string{"next", "close", "snapshot-release"}, events) + require.Equal(t, 1, child.closeCalls) +} + +func TestDuckLakeOperationValueIterClosesChildBeforeReleaseOnEOF(t *testing.T) { + ctx := sql.NewEmptyContext() + events := []string{} + child := &orderedOperationIter{events: &events} + iter := wrapDuckLakeOperationIter(child, func() { events = append(events, "release") }) + + valueIter, ok := iter.(sql.ValueRowIter) + require.True(t, ok) + _, err := valueIter.NextValueRow(ctx) + require.ErrorIs(t, err, io.EOF) + require.Equal(t, []string{"next-value", "close", "release"}, events) + require.NoError(t, valueIter.Close(ctx)) + require.Equal(t, 1, child.closeCalls) +} + +func TestDuckLakeOperationIterPreservesChildCloseErrorForExplicitClose(t *testing.T) { + ctx := sql.NewEmptyContext() + events := []string{} + child := &orderedOperationIter{ + events: &events, + closeErr: errOperationChild, + } + iter := wrapDuckLakeOperationIter(child, func() { events = append(events, "release") }) + + _, err := iter.Next(ctx) + require.ErrorIs(t, err, io.EOF) + // EOF must stay EOF even if the eager child close reports an error. + require.Equal(t, []string{"next", "close", "release"}, events) + require.ErrorIs(t, iter.Close(ctx), errOperationChild) + require.Equal(t, 1, child.closeCalls) +} + +func TestDuckLakeOperationIterSharesCloseStateAcrossMutableAliases(t *testing.T) { + ctx := sql.NewEmptyContext() + events := []string{} + first := &orderedOperationIter{events: &events} + second := &orderedOperationIter{events: &events} + iter := wrapDuckLakeOperationIter(first, func() { events = append(events, "release") }) + + mutable, ok := iter.(sql.MutableRowIter) + require.True(t, ok) + replacement := mutable.WithChildIter(second) + _, err := replacement.Next(ctx) + require.ErrorIs(t, err, io.EOF) + require.Equal(t, []string{"next", "close", "release"}, events) + + // The original alias must not close its abandoned child after the + // replacement has already released the shared lease. + require.NoError(t, iter.Close(ctx)) + require.Zero(t, first.closeCalls) + require.Equal(t, 1, second.closeCalls) +} + +func TestDuckLakeOperationIterMutableAliasDispatchesCurrentRowChild(t *testing.T) { + ctx := sql.NewEmptyContext() + events := []string{} + first := &orderedOperationIter{events: &events} + second := &orderedOperationIter{events: &events} + iter := wrapDuckLakeOperationIter(first, nil) + + mutable, ok := iter.(sql.MutableRowIter) + require.True(t, ok) + replacement := mutable.WithChildIter(second) + require.Same(t, second, iter.(sql.MutableRowIter).GetChildIter()) + + // Both the original alias and the alias returned by WithChildIter must read + // the shared current child. Calling the original alias here is the regression + // that used to continue driving `first` through its stale child field. + _, err := iter.Next(ctx) + require.ErrorIs(t, err, io.EOF) + require.Zero(t, first.nextCalls) + require.Equal(t, 1, second.nextCalls) + require.Equal(t, 1, second.closeCalls) + _, err = replacement.Next(ctx) + require.ErrorIs(t, err, io.EOF) + require.Equal(t, 1, second.nextCalls) +} + +func TestDuckLakeOperationValueAliasDispatchesCurrentValueChild(t *testing.T) { + ctx := sql.NewEmptyContext() + events := []string{} + first := &orderedOperationIter{events: &events} + second := &orderedOperationIter{events: &events} + iter := wrapDuckLakeOperationIter(first, nil) + valueAlias, ok := iter.(sql.ValueRowIter) + require.True(t, ok) + mutable, ok := iter.(sql.MutableRowIter) + require.True(t, ok) + mutable.WithChildIter(second) + require.True(t, valueAlias.IsValueRowIter(ctx)) + + _, err := valueAlias.NextValueRow(ctx) + require.ErrorIs(t, err, io.EOF) + require.Zero(t, first.valueCalls) + require.Equal(t, 1, second.valueCalls) + require.Equal(t, 1, second.closeCalls) +} + +func TestDuckLakeOperationIterSplitLeasesReleaseInOrder(t *testing.T) { + ctx := sql.NewEmptyContext() + events := []string{} + child := &orderedOperationIter{events: &events} + iter := wrapDuckLakeOperationIterWithLeases( + nil, + child, + func() { events = append(events, "snapshot-release") }, + func() { events = append(events, "operation-release") }, + ) + + _, err := iter.Next(ctx) + require.ErrorIs(t, err, io.EOF) + require.Equal(t, + []string{"next", "close", "snapshot-release", "operation-release"}, + events, + ) + require.NoError(t, iter.Close(ctx)) + require.Equal(t, 1, child.closeCalls) +} + +func TestDuckLakeOperationIterDefersOperationLeaseToPostFinalization(t *testing.T) { + backendSession := NewSession(memory.NewSession(sql.NewBaseSession(), nil), nil) + sess := &postCleanupTestSession{Session: backendSession} + ctx := sql.NewContext( + mycontext.WithFrontendQuery(context.Background()), + sql.WithSession(sess), + ) + tx, err := sess.StartTransaction(ctx, sql.ReadWrite) + require.NoError(t, err) + ctx.SetTransaction(tx) + + events := []string{} + child := &orderedOperationIter{events: &events} + iter := wrapDuckLakeOperationIterWithLeases( + ctx, + child, + func() { events = append(events, "snapshot-release") }, + func() { events = append(events, "operation-release") }, + ) + require.NotNil(t, sess.postCleanup) + + _, err = iter.Next(ctx) + require.ErrorIs(t, err, io.EOF) + require.Equal(t, []string{"next", "close", "snapshot-release"}, events) + + // The session invokes the post hook only after commit or rollback has + // completed. It is idempotent even if both a normal terminal path and a + // later cleanup callback try to release the same operation. + sess.postCleanup(true) + sess.postCleanup(false) + require.Equal(t, []string{"next", "close", "snapshot-release", "operation-release"}, events) +} + +func TestDuckLakeOperationIterUsesCapturedTxDuringPendingShutdown(t *testing.T) { + gate := newDeterministicShutdownGate() + holder := &shutdownGateSession{ + Session: memory.NewSession(sql.NewBaseSession(), nil), + gate: gate, + conn: new(stdsql.Conn), + tx: new(stdsql.Tx), + } + ctx := sql.NewContext(context.Background(), sql.WithSession(holder)) + logicalTx, err := holder.Session.StartTransaction(ctx, sql.ReadWrite) + require.NoError(t, err) + ctx.SetTransaction(logicalTx) + + _, owner, physicalTx, snapshotRelease, err := adapter.GetExecutionSnapshotWithLease(ctx) + require.NoError(t, err) + require.Same(t, holder.conn, owner) + require.Same(t, holder.tx, physicalTx) + t.Cleanup(func() { + snapshotRelease() + gate.releaseWriter() + select { + case <-gate.writerDone: + case <-time.After(time.Second): + } + }) + + go gate.startWriter() + select { + case <-gate.writerReady: + case <-time.After(time.Second): + t.Fatal("shutdown writer did not become pending") + } + select { + case <-gate.writerGot: + t.Fatal("shutdown writer acquired before the retained snapshot was released") + default: + } + + var operationCalls atomic.Int32 + child := &orderedOperationIter{events: new([]string)} + wrappedDone := make(chan sql.RowIter, 1) + go func() { + wrappedDone <- wrapDuckLakeOperationIterWithLeasesAndTx( + ctx, + child, + physicalTx, + snapshotRelease, + func() { operationCalls.Add(1) }, + ) + }() + + var wrapped sql.RowIter + select { + case wrapped = <-wrappedDone: + case <-time.After(time.Second): + t.Fatal("iterator wrapper re-entered the pending shutdown reader gate") + } + require.Zero(t, holder.bindingCalls.Load(), "wrapper must not perform a second GetTxnBinding lookup") + holder.postMu.Lock() + registeredTx, callback := holder.postTx, holder.postCallback + holder.postMu.Unlock() + require.Same(t, physicalTx, registeredTx) + require.NotNil(t, callback) + + // Close releases the retained reader, allowing the pending teardown writer + // to acquire. The deferred operation remains held until physical finalization. + require.ErrorIs(t, wrapped.Close(ctx), errDuckLakeOperationClosedBeforeEOF) + select { + case <-gate.writerGot: + case <-time.After(time.Second): + t.Fatal("shutdown writer did not acquire after snapshot release") + } + require.Zero(t, operationCalls.Load()) + + gate.releaseWriter() + select { + case <-gate.writerDone: + case <-time.After(time.Second): + t.Fatal("shutdown writer did not finish") + } + callback(false) + require.Equal(t, int32(1), operationCalls.Load()) +} + +func TestDuckLakeOperationIterNilChildReleasesSplitLeasesOnce(t *testing.T) { + snapshotCalls, operationCalls := 0, 0 + iter := wrapDuckLakeOperationIterWithLeases( + nil, + nil, + func() { snapshotCalls++ }, + func() { operationCalls++ }, + ) + require.Nil(t, iter) + // A failed iterator construction has no child boundary, so both leases are + // released immediately and never need a later Close call. + require.Equal(t, 1, snapshotCalls) + require.Equal(t, 1, operationCalls) +} diff --git a/backend/engine.go b/backend/engine.go index dfe1e74e..0aeabb33 100644 --- a/backend/engine.go +++ b/backend/engine.go @@ -22,6 +22,7 @@ import ( sqle "github.com/dolthub/go-mysql-server" "github.com/dolthub/go-mysql-server/sql" "github.com/dolthub/go-mysql-server/sql/analyzer" + "github.com/dolthub/go-mysql-server/sql/plan" "github.com/dolthub/go-mysql-server/sql/types" "github.com/dolthub/vitess/go/vt/sqlparser" ) @@ -33,6 +34,11 @@ func NewEngine(provider *catalog.DatabaseProvider) (*sqle.Engine, *DuckBuilder) parser := &mysqlParser{Parser: sql.NewMysqlParser()} overrides := sql.EngineOverrides{ Builder: sql.BuilderOverrides{Parser: parser}, + Hooks: sql.ExecutionHooks{ + CreateTable: sql.CreateTable{ + PreSQLExecution: prepareMySQLCreateTableStorage, + }, + }, } engine := sqle.New(analyzer.NewBuilder(provider).AddOverrides(overrides).Build(), nil) builder := NewDuckBuilder(engine.Analyzer.ExecBuilder, provider) @@ -40,6 +46,25 @@ func NewEngine(provider *catalog.DatabaseProvider) (*sqle.Engine, *DuckBuilder) return engine, builder } +// prepareMySQLCreateTableStorage bridges the planner's table-option map to +// the catalog's request-scoped storage selector. The catalog consumes the +// selector while creating the table and persists it in the managed comment; +// no credentials, endpoints, or object paths are accepted from SQL. +func prepareMySQLCreateTableStorage(ctx *sql.Context, _ sql.StatementRunner, node sql.Node) (sql.Node, error) { + create, ok := node.(*plan.CreateTable) + if !ok { + return node, nil + } + selection, err := catalog.ResolveMySQLTableStorage(create.TableOpts) + if err != nil { + return nil, err + } + if err := catalog.SetTableStorageSelection(ctx, selection); err != nil { + return nil, err + } + return create, nil +} + // registerMySQLCompatibilitySystemVariables keeps MyDuck's advertised SQL // compatibility level stable across GMS upgrades. Clients such as MySQL Shell // branch on @@version and otherwise probe newer variables MyDuck does not @@ -67,13 +92,21 @@ type mysqlParser struct { func (p *mysqlParser) ParseSimple(query string) (sqlparser.Statement, error) { compat := rewriteMySQLCompatibility(query) stmt, err := p.Parser.ParseSimple(compat.query) - return normalizeMySQLStatement(stmt, compat.replacements), err + stmt = normalizeMySQLStatement(stmt, compat.replacements) + if err == nil { + err = validateMySQLTableStorageStatement(stmt) + } + return stmt, err } func (p *mysqlParser) Parse(ctx *sql.Context, query string, multi bool) (sqlparser.Statement, string, string, error) { compat := rewriteMySQLCompatibility(query) stmt, parsed, remainder, err := p.Parser.Parse(ctx, compat.query, multi) - return normalizeMySQLStatement(stmt, compat.replacements), compat.restoreParsedQuery(parsed), remainder, err + stmt = normalizeMySQLStatement(stmt, compat.replacements) + if err == nil { + err = validateMySQLTableStorageStatement(stmt) + } + return stmt, compat.restoreParsedQuery(parsed), remainder, err } func (p *mysqlParser) ParseWithOptions( @@ -85,7 +118,11 @@ func (p *mysqlParser) ParseWithOptions( ) (sqlparser.Statement, string, string, error) { compat := rewriteMySQLCompatibility(query) stmt, parsed, remainder, err := p.Parser.ParseWithOptions(ctx, compat.query, delimiter, multi, options) - return normalizeMySQLStatement(stmt, compat.replacements), compat.restoreParsedQuery(parsed), remainder, err + stmt = normalizeMySQLStatement(stmt, compat.replacements) + if err == nil { + err = validateMySQLTableStorageStatement(stmt) + } + return stmt, compat.restoreParsedQuery(parsed), remainder, err } func (p *mysqlParser) ParseOneWithOptions( @@ -95,7 +132,33 @@ func (p *mysqlParser) ParseOneWithOptions( ) (sqlparser.Statement, int, error) { compat := rewriteMySQLCompatibility(query) stmt, index, err := p.Parser.ParseOneWithOptions(ctx, compat.query, options) - return normalizeMySQLStatement(stmt, compat.replacements), compat.originalOffset(index), err + stmt = normalizeMySQLStatement(stmt, compat.replacements) + if err == nil { + err = validateMySQLTableStorageStatement(stmt) + } + return stmt, compat.originalOffset(index), err +} + +// validateMySQLTableStorageStatement runs before the planner turns table +// options into a map. That preserves duplicate ENGINE/myduck_storage +// declarations, which would otherwise be silently overwritten by the map. +func validateMySQLTableStorageStatement(stmt sqlparser.Statement) error { + ddl, ok := stmt.(*sqlparser.DDL) + if !ok || ddl.TableSpec == nil || len(ddl.TableSpec.TableOpts) == 0 { + return nil + } + options := make([]catalog.TableStorageOption, 0, len(ddl.TableSpec.TableOpts)) + for _, option := range ddl.TableSpec.TableOpts { + if option == nil { + continue + } + options = append(options, catalog.TableStorageOption{ + Name: option.Name, + Value: option.Value, + }) + } + _, err := catalog.NormalizeTableStorageOptions(options) + return err } type mysqlOptionReplacement struct { diff --git a/backend/engine_test.go b/backend/engine_test.go index 5ef3778d..bab68544 100644 --- a/backend/engine_test.go +++ b/backend/engine_test.go @@ -18,6 +18,7 @@ import ( "strings" "testing" + "github.com/apecloud/myduckserver/catalog" "github.com/dolthub/go-mysql-server/sql" "github.com/dolthub/vitess/go/vt/sqlparser" "github.com/stretchr/testify/require" @@ -154,3 +155,36 @@ func TestMySQLParserRestoresDatabaseFiltersInMultiQuery(t *testing.T) { require.NoError(t, err) require.Equal(t, strings.Index(query, " SELECT 1"), index) } + +func TestMySQLParserTableStorageOptions(t *testing.T) { + parser := &mysqlParser{Parser: sql.NewMysqlParser()} + for _, query := range []string{ + "CREATE TABLE object_table (id INT) ENGINE=DUCKLAKE", + "CREATE TABLE local_table (id INT) ENGINE=InnoDB", + } { + stmt, _, _, err := parser.ParseWithOptions(context.Background(), query, ';', false, sqlparser.ParserOptions{}) + require.NoError(t, err, query) + ddl, ok := stmt.(*sqlparser.DDL) + require.True(t, ok, query) + require.NotNil(t, ddl.TableSpec, query) + require.NotEmpty(t, ddl.TableSpec.TableOpts, query) + for _, option := range ddl.TableSpec.TableOpts { + t.Logf("%s => name=%q value=%q", query, option.Name, option.Value) + } + } +} + +func TestMySQLParserRejectsConflictingTableStorageOptions(t *testing.T) { + parser := &mysqlParser{Parser: sql.NewMysqlParser()} + for _, test := range []struct { + query string + want error + }{ + {query: "CREATE TABLE duplicate_engine (id INT) ENGINE=DUCKLAKE ENGINE=DUCKLAKE", want: catalog.ErrTableStorageDuplicate}, + {query: "CREATE TABLE conflicting_engine (id INT) ENGINE=DUCKLAKE ENGINE=LOCAL", want: catalog.ErrTableStorageConflict}, + } { + _, _, _, err := parser.ParseWithOptions(context.Background(), test.query, ';', false, sqlparser.ParserOptions{}) + require.Error(t, err, test.query) + require.ErrorIs(t, err, test.want, test.query) + } +} diff --git a/backend/execution_lease.go b/backend/execution_lease.go new file mode 100644 index 00000000..eeb4970e --- /dev/null +++ b/backend/execution_lease.go @@ -0,0 +1,17 @@ +package backend + +import "github.com/dolthub/go-mysql-server/sql" + +// WrapRowIterWithRelease keeps a caller-owned lifecycle lease until the +// wrapped iterator has been fully consumed or explicitly closed. The release +// callback is invoked at most once, after the child iterator is closed, and +// the wrapper preserves the mutable-child and ValueRowIter behavior used by +// GMS execution nodes. +// +// This is intentionally a thin export of the same wrapper used for the +// provider's DuckLake operation lease. PostgreSQL protocol paths can therefore +// retain an execution snapshot without introducing a second iterator lifetime +// implementation. +func WrapRowIterWithRelease(iter sql.RowIter, release func()) sql.RowIter { + return wrapDuckLakeOperationIter(iter, release) +} diff --git a/backend/execution_lease_test.go b/backend/execution_lease_test.go new file mode 100644 index 00000000..67bc8d16 --- /dev/null +++ b/backend/execution_lease_test.go @@ -0,0 +1,144 @@ +package backend + +import ( + "context" + "io" + "runtime" + "testing" + "time" + + "github.com/apecloud/myduckserver/adapter" + "github.com/apecloud/myduckserver/catalog" + "github.com/dolthub/go-mysql-server/memory" + "github.com/dolthub/go-mysql-server/sql" + "github.com/stretchr/testify/require" +) + +func TestWrapRowIterWithReleaseKeepsExistingTerminalOrdering(t *testing.T) { + ctx := sql.NewEmptyContext() + events := []string{} + child := &orderedOperationIter{events: &events} + wrapped := WrapRowIterWithRelease(child, func() { + events = append(events, "release") + }) + + _, err := wrapped.Next(ctx) + require.ErrorIs(t, err, io.EOF) + require.Equal(t, []string{"next", "close", "release"}, events) + require.NoError(t, wrapped.Close(ctx)) + require.Equal(t, 1, child.closeCalls) +} + +type sessionLeaseFixture struct { + provider *catalog.DatabaseProvider + session *Session + ctx *sql.Context +} + +func newSessionLeaseFixture(t *testing.T) *sessionLeaseFixture { + t.Helper() + provider := catalog.NewInMemoryDBProvider() + session := NewSession(memory.NewSession(sql.NewBaseSession(), provider), provider) + ctx := sql.NewContext(context.Background(), sql.WithSession(session)) + return &sessionLeaseFixture{provider: provider, session: session, ctx: ctx} +} + +func startProviderCloseAfterStorageShutdown( + t *testing.T, + fixture *sessionLeaseFixture, +) <-chan error { + t.Helper() + closeDone := make(chan error, 1) + go func() { closeDone <- fixture.provider.Close() }() + + deadline := time.Now().Add(5 * time.Second) + for fixture.provider.Storage().PingContext(context.Background()) == nil { + if time.Now().After(deadline) { + t.Fatal("provider close did not announce storage shutdown") + } + runtime.Gosched() + } + return closeDone +} + +func requireProviderCloseWaiting(t *testing.T, closeDone <-chan error) { + t.Helper() + select { + case err := <-closeDone: + t.Fatalf("provider closed before the query lease reached its terminal boundary: %v", err) + default: + } +} + +func requireProviderClosed(t *testing.T, closeDone <-chan error) { + t.Helper() + select { + case err := <-closeDone: + require.NoError(t, err) + case <-time.After(5 * time.Second): + t.Fatal("provider close remained blocked after query lease release") + } +} + +func TestSessionLeasedRowsKeepProviderCloseWaitingUntilEOF(t *testing.T) { + fixture := newSessionLeaseFixture(t) + rows, err := adapter.QueryCatalog(fixture.ctx, "SELECT value FROM (VALUES (1), (2)) AS lease_rows(value)") + require.NoError(t, err) + t.Cleanup(func() { _ = rows.Close() }) + + closeDone := startProviderCloseAfterStorageShutdown(t, fixture) + requireProviderCloseWaiting(t, closeDone) + + var values []int + for rows.Next() { + var value int + require.NoError(t, rows.Scan(&value)) + values = append(values, value) + } + require.NoError(t, rows.Err()) + require.Equal(t, []int{1, 2}, values) + requireProviderClosed(t, closeDone) +} + +func TestSessionLeasedRowsKeepProviderCloseWaitingUntilExplicitClose(t *testing.T) { + fixture := newSessionLeaseFixture(t) + rows, err := adapter.QueryCatalog(fixture.ctx, "SELECT value FROM (VALUES (1), (2)) AS lease_rows(value)") + require.NoError(t, err) + + closeDone := startProviderCloseAfterStorageShutdown(t, fixture) + requireProviderCloseWaiting(t, closeDone) + require.NoError(t, rows.Close()) + requireProviderClosed(t, closeDone) +} + +func TestSessionLeasedRowKeepsProviderCloseWaitingUntilScan(t *testing.T) { + fixture := newSessionLeaseFixture(t) + row := fixture.session.QueryRow(context.Background(), "SELECT 42") + require.NotNil(t, row) + + closeDone := startProviderCloseAfterStorageShutdown(t, fixture) + requireProviderCloseWaiting(t, closeDone) + + var value int + require.NoError(t, row.Scan(&value)) + require.Equal(t, 42, value) + requireProviderClosed(t, closeDone) +} + +func TestSessionLeasedRowPreservesExecutionLeaseAdmissionError(t *testing.T) { + fixture := newSessionLeaseFixture(t) + require.NoError(t, fixture.provider.Close()) + + _, _, release, expectedErr := fixture.session.GetExecutionSnapshotLease(context.Background(), true) + release() + require.Error(t, expectedErr) + + row := fixture.session.QueryRow(context.Background(), "SELECT must_not_execute") + require.NotNil(t, row) + var value int + require.EqualError(t, row.Scan(&value), expectedErr.Error()) +} + +var _ interface { + QueryRow(context.Context, string, ...any) *adapter.LeasedRow +} = (*Session)(nil) diff --git a/backend/executor.go b/backend/executor.go index 410d9901..5d983956 100644 --- a/backend/executor.go +++ b/backend/executor.go @@ -16,6 +16,7 @@ package backend import ( stdsql "database/sql" "fmt" + "strings" "github.com/apecloud/myduckserver/adapter" "github.com/apecloud/myduckserver/catalog" @@ -74,16 +75,32 @@ func (b *DuckBuilder) Provider() *catalog.DatabaseProvider { } func (b *DuckBuilder) Build(ctx *sql.Context, root sql.Node, r sql.Row) (iter sql.RowIter, err error) { - // The engine passes a nil input row for the top-level result iterator. - // Subquery iterators must not count intermediate rows against the limit. - if r == nil { - defer func() { - if err == nil && iter != nil { - iter = ApplyQueryRowLimit(ctx, root.Schema(ctx), iter) + topLevel := r == nil + var snapshotRelease, operationRelease func() + var physicalTx *stdsql.Tx + defer func() { + // Query-row limiting is a top-level concern. Keep the two lifecycle leases + // outside that wrapper so an early limit/error closes the child before the + // snapshot is released, while an admitted operation can remain held until + // transaction finalization. + if topLevel && err == nil && iter != nil { + iter = ApplyQueryRowLimit(ctx, root.Schema(ctx), iter) + if snapshotRelease != nil || operationRelease != nil { + iter = wrapDuckLakeOperationIterWithLeasesAndTx(ctx, iter, physicalTx, snapshotRelease, operationRelease) + snapshotRelease = nil + operationRelease = nil } - }() - } + } + if snapshotRelease != nil { + snapshotRelease() + } + if operationRelease != nil { + operationRelease() + } + }() + // The engine passes a nil input row for the top-level result iterator. + // Subquery iterators must not count intermediate rows against the limit. // Flush the delta buffer before executing the query. Replication controller // commands must remain able to stop an applier while it is connecting and // therefore cannot wait on the applier-owned query barrier. @@ -101,18 +118,51 @@ func (b *DuckBuilder) Build(ctx *sql.Context, root sql.Node, r sql.Row) (iter sq }).Traceln("Building node:", n) } + // Admit every top-level frontend execution before any fallback + // classification. GMS-owned CALL/DDL/constraint paths can still touch + // DuckLake through their base builder, and classifying them first leaves a + // cleanup window in which rollback can overtake that work. Transaction + // controls are explicitly excluded because their job is to drain/release the + // barrier and must remain usable while cleanup is active. + if topLevel && b.shouldBeginDuckLakeOperation(ctx, root) { + operationRelease, err = b.beginDuckLakeOperationWithError(ctx) + if err != nil { + return nil, err + } + } + // TODO; find a better way to fallback to the base builder switch n.(type) { + case *plan.Call: + // GMS buildCall temporarily clears the session transaction and restores + // it in a defer. Mark that synchronous scope so Session can suspend an + // implicit marker instead of treating the detach as a real rollback. + if sess, ok := ctx.Session.(*Session); ok { + endProcedureScope := sess.beginProcedureTransactionScope() + defer endProcedureScope() + } + return b.base.Build(ctx, root, r) case *plan.CreateDB, *plan.DropDB, *plan.DropTable, *plan.RenameTable, *plan.CreateTable, *plan.AddColumn, *plan.RenameColumn, *plan.DropColumn, *plan.ModifyColumn, *plan.Truncate, *plan.CreateIndex, *plan.DropIndex, *plan.AlterIndex, *plan.ShowIndexes, *plan.ShowTables, *plan.ShowCreateTable, *plan.ShowColumns, *plan.ShowBinlogs, *plan.ShowBinlogStatus, *plan.ShowWarnings, - *plan.StartTransaction, *plan.Commit, *plan.Rollback, + *plan.Commit, *plan.Rollback, *plan.Set, *plan.ShowVariables, *plan.AlterDefaultSet, *plan.AlterDefaultDrop: return b.base.Build(ctx, root, r) + case *plan.StartTransaction: + // GMS sets IgnoreAutoCommit only after calling Session.StartTransaction. + // An explicit BEGIN must therefore opt in before delegating so the + // session can bind the transaction to DuckDB even when @@autocommit=1. + ignoreAutoCommit := ctx.GetIgnoreAutoCommit() + ctx.SetIgnoreAutoCommit(true) + iter, err := b.base.Build(ctx, root, r) + if err != nil { + ctx.SetIgnoreAutoCommit(ignoreAutoCommit) + } + return iter, err case *plan.InsertInto: insert := n.(*plan.InsertInto) @@ -158,8 +208,11 @@ func (b *DuckBuilder) Build(ctx *sql.Context, root sql.Node, r sql.Row) (iter sq ctx.GetLogger().Traceln("Falling back to the base builder") return b.base.Build(ctx, root, r) } + if !b.needsExecutionSnapshot(n) { + return b.base.Build(ctx, n, r) + } - conn, err := b.provider.Pool().GetConnForSchema(ctx, ctx.ID(), ctx.GetCurrentDatabase()) + execer, conn, physicalTx, snapshotRelease, err := b.executionSnapshotWithLease(ctx) if err != nil { return nil, err } @@ -167,19 +220,23 @@ func (b *DuckBuilder) Build(ctx *sql.Context, root sql.Node, r sql.Row) (iter sq switch node := n.(type) { case *plan.Use: useStmt := "USE " + catalog.FullSchemaName(adapter.GetCurrentCatalog(ctx), node.Database().Name()) - if _, err := conn.ExecContext(ctx.Context, useStmt); err != nil { + if _, err := execer.ExecContext(ctx.Context, useStmt); err != nil { if catalog.IsDuckDBSetSchemaNotFoundError(err) { return nil, sql.ErrDatabaseNotFound.New(node.Database().Name()) } return nil, err } + // USE has finished its only physical operation. Do not retain a pool + // snapshot while the native GMS result iterator is finalized. + snapshotRelease() + snapshotRelease = nil return b.base.Build(ctx, root, r) // ResolvedTable is for `SELECT * FROM table` and `TABLE table` // SubqueryAlias is for `SELECT * FROM view` case *plan.ResolvedTable, *plan.SubqueryAlias, *plan.TableAlias: - return b.executeQuery(ctx, node, conn) + return b.executeQuery(ctx, node, execer, conn) case *plan.Distinct, *plan.OrderedDistinct: - return b.executeQuery(ctx, node, conn) + return b.executeQuery(ctx, node, execer, conn) case *plan.TableCopier: // We preserve the table schema in a best-effort manner. // For simple `CREATE TABLE t AS SELECT * FROM t`, @@ -188,22 +245,125 @@ func (b *DuckBuilder) Build(ctx *sql.Context, root sql.Node, r sql.Row) (iter sq if _, ok := node.Source.(*plan.ResolvedTable); ok { return b.base.Build(ctx, root, r) } - return b.executeDML(ctx, node, conn) + return b.executeDML(ctx, node, execer, conn) case *plan.DeleteFrom: if node.Returning != nil { - return b.executeQuery(ctx, node, conn) + return b.executeQuery(ctx, node, execer, conn) } node.Child = withMySQLOkResultSchema(node.Child) - return b.executeDML(ctx, node, conn) + return b.executeDML(ctx, node, execer, conn) case sql.Expressioner: - return b.executeExpressioner(ctx, node, conn) + return b.executeExpressioner(ctx, node, execer, conn) case *plan.Truncate: - return b.executeDML(ctx, node, conn) + return b.executeDML(ctx, node, execer, conn) default: return b.base.Build(ctx, n, r) } } +// needsExecutionSnapshot identifies the branches that actually execute SQL +// through the DuckDB connection captured by executionSnapshotWithLease. GMS +// fallback plans must retain their native iterator contract; taking a snapshot +// for them merely to release it from a wrapper can cause finalization to rewrite +// an otherwise valid child to nil. +func (b *DuckBuilder) needsExecutionSnapshot(root sql.Node) bool { + // DDL and SHOW plans are intentionally executed by GMS's catalog builder. + // Several of them also satisfy sql.Expressioner, so this guard must precede + // the interface case below. + if plan.IsDDLNode(root) || plan.IsShowNode(root) { + return false + } + switch node := root.(type) { + case *plan.Use, + *plan.ResolvedTable, + *plan.SubqueryAlias, + *plan.TableAlias, + *plan.Distinct, + *plan.OrderedDistinct, + *plan.DeleteFrom, + sql.Expressioner: + return true + case *plan.TableCopier: + _, simple := node.Source.(*plan.ResolvedTable) + return !simple + default: + return false + } +} + +func (b *DuckBuilder) shouldBeginDuckLakeOperation(ctx *sql.Context, root sql.Node) bool { + if b == nil || b.provider == nil || ctx == nil || ctx.Session == nil { + return false + } + // Do not install a lifecycle wrapper for ordinary local/catalog queries when + // the service-managed DuckLake catalog is disabled. Besides avoiding needless + // iterator layers, this preserves GMS's native fallback iterator contracts + // (some DDL paths intentionally return an unwrapped result iterator). + if !b.provider.DuckLakeObjectStorageEnabled() { + return false + } + // Transaction controls can invoke rollback/commit from iterator Close. They + // must not hold the same provider operation lease that rollback cleanup waits + // on, or a GMS-only control would wait on itself. + switch root.(type) { + case *plan.StartTransaction, *plan.Commit, *plan.Rollback: + return false + } + if _, ok := ctx.Session.(adapter.ConnectionHolder); !ok { + return false + } + // Admission is intentionally independent of the current transaction + // pointer. A transaction can be created immediately after this check; taking + // the operation lease first closes that race, while the provider's + // transaction barrier continues to account for the physical transaction. + return true +} + +func (b *DuckBuilder) beginDuckLakeOperation(ctx *sql.Context) func() { + release, _ := b.beginDuckLakeOperationWithError(ctx) + return release +} + +func (b *DuckBuilder) beginDuckLakeOperationWithError(ctx *sql.Context) (func(), error) { + if b == nil || b.provider == nil || ctx == nil { + return func() {}, nil + } + if sess, ok := ctx.Session.(*Session); ok { + return sess.BeginDuckLakeOperationWithError(ctx) + } + return b.provider.BeginDuckLakeOperationWithError(ctx) +} + +// executionSnapshot obtains the executor and its physical owner atomically +// from the session lifecycle. The small legacy fallback is retained for unit +// contexts that expose a provider pool but no adapter.ConnectionHolder. +func (b *DuckBuilder) executionSnapshot(ctx *sql.Context) (adapter.SQLExecutor, *stdsql.Conn, *stdsql.Tx, error) { + execer, conn, tx, release, err := b.executionSnapshotWithLease(ctx) + if release != nil { + release() + } + return execer, conn, tx, err +} + +// executionSnapshotWithLease is the lifetime-safe execution form. The +// returned release callback must remain held until all iterator work using the +// captured executor and physical owner has completed. +func (b *DuckBuilder) executionSnapshotWithLease(ctx *sql.Context) (adapter.SQLExecutor, *stdsql.Conn, *stdsql.Tx, func(), error) { + if ctx != nil && ctx.Session != nil { + if _, ok := ctx.Session.(adapter.ConnectionHolder); ok { + return adapter.GetExecutionSnapshotWithLease(ctx) + } + } + if b == nil || b.provider == nil { + return nil, nil, nil, func() {}, fmt.Errorf("database provider is unavailable") + } + conn, err := b.provider.Pool().GetConnForSchema(ctx, ctx.ID(), ctx.GetCurrentDatabase()) + if err != nil { + return nil, nil, nil, func() {}, err + } + return conn, conn, nil, func() {}, nil +} + func (b *DuckBuilder) flushDeltaBuffer(ctx *sql.Context, root sql.Node) error { if b.FlushDeltaBuffer == nil || !shouldFlushDeltaBuffer(root) { return nil @@ -239,25 +399,25 @@ func isStandaloneVectorConversion(ctx *sql.Context, n sql.Node) bool { return found } -func (b *DuckBuilder) executeExpressioner(ctx *sql.Context, n sql.Expressioner, conn *stdsql.Conn) (sql.RowIter, error) { +func (b *DuckBuilder) executeExpressioner(ctx *sql.Context, n sql.Expressioner, execer adapter.SQLExecutor, conn *stdsql.Conn) (sql.RowIter, error) { node := n.(sql.Node) switch n := n.(type) { case *plan.InsertInto: if n.Returning != nil { - return b.executeQuery(ctx, node, conn) + return b.executeQuery(ctx, node, execer, conn) } - return b.executeDML(ctx, node, conn) + return b.executeDML(ctx, node, execer, conn) case *plan.Update: if n.Returning == nil { n.Child = withMySQLOkResultSchema(n.Child) } - return b.executeDML(ctx, node, conn) + return b.executeDML(ctx, node, execer, conn) default: - return b.executeQuery(ctx, node, conn) + return b.executeQuery(ctx, node, execer, conn) } } -func (b *DuckBuilder) executeQuery(ctx *sql.Context, n sql.Node, conn *stdsql.Conn) (sql.RowIter, error) { +func (b *DuckBuilder) executeQuery(ctx *sql.Context, n sql.Node, execer adapter.SQLExecutor, owners ...*stdsql.Conn) (sql.RowIter, error) { ctx.GetLogger().Trace("Executing Query...") var ( @@ -279,6 +439,14 @@ func (b *DuckBuilder) executeQuery(ctx *sql.Context, n sql.Node, conn *stdsql.Co return nil, catalog.ErrTranspiler.New(err) } duckSQL = QueryForJSONScan(duckSQL, n.Schema(ctx)) + var conn *stdsql.Conn + if len(owners) > 0 { + conn = owners[0] + } + duckSQL, err = b.rewriteObjectRelationsWithSnapshot(ctx, n, duckSQL, execer, conn) + if err != nil { + return nil, err + } if log := ctx.GetLogger(); log.Logger.IsLevelEnabled(logrus.TraceLevel) { log.WithFields(logrus.Fields{ @@ -288,7 +456,7 @@ func (b *DuckBuilder) executeQuery(ctx *sql.Context, n sql.Node, conn *stdsql.Co } // Execute the DuckDB query - rows, err := conn.QueryContext(ctx.Context, duckSQL) + rows, err := execer.QueryContext(ctx.Context, duckSQL) if err != nil { return nil, err } @@ -296,12 +464,20 @@ func (b *DuckBuilder) executeQuery(ctx *sql.Context, n sql.Node, conn *stdsql.Co return NewSQLRowIter(rows, n.Schema(ctx)) } -func (b *DuckBuilder) executeDML(ctx *sql.Context, n sql.Node, conn *stdsql.Conn) (sql.RowIter, error) { +func (b *DuckBuilder) executeDML(ctx *sql.Context, n sql.Node, execer adapter.SQLExecutor, owners ...*stdsql.Conn) (sql.RowIter, error) { // Translate the MySQL query to a DuckDB query duckSQL, err := transpiler.TranslateWithSQLGlot(queryForTranslation(ctx)) if err != nil { return nil, catalog.ErrTranspiler.New(err) } + var conn *stdsql.Conn + if len(owners) > 0 { + conn = owners[0] + } + duckSQL, err = b.rewriteObjectRelationsWithSnapshot(ctx, n, duckSQL, execer, conn) + if err != nil { + return nil, err + } if log := ctx.GetLogger(); log.Logger.IsLevelEnabled(logrus.TraceLevel) { log.WithFields(logrus.Fields{ @@ -314,10 +490,10 @@ func (b *DuckBuilder) executeDML(ctx *sql.Context, n sql.Node, conn *stdsql.Conn if _, ok := n.(*plan.TableCopier); ok { // DuckDB returns the CTAS insert count as a one-row result. Its C API // rows-changed value is defined only for INSERT, UPDATE, and DELETE. - err = conn.QueryRowContext(ctx.Context, duckSQL).Scan(&affected) + err = execer.QueryRowContext(ctx.Context, duckSQL).Scan(&affected) } else { var result stdsql.Result - result, err = conn.ExecContext(ctx.Context, duckSQL) + result, err = execer.ExecContext(ctx.Context, duckSQL) if err == nil { affected, err = result.RowsAffected() } @@ -349,6 +525,83 @@ func (b *DuckBuilder) executeDML(ctx *sql.Context, n sql.Node, conn *stdsql.Conn })), nil } +// rewriteObjectRelations keeps GMS's logical shadow tables available for +// analysis while routing the executed SQL to the durable DuckLake relation. +// The plan is the source of truth for relation identity, so a bare table name +// is only routed when it is unambiguous within this statement. +func (b *DuckBuilder) rewriteObjectRelations(ctx *sql.Context, root sql.Node, query string) (string, error) { + execer, conn, _, err := b.executionSnapshot(ctx) + if err != nil { + return "", err + } + return b.rewriteObjectRelationsWithSnapshot(ctx, root, query, execer, conn) +} + +// rewriteObjectRelationsWithSnapshot is the transaction-affine variant used +// by statement execution. It never asks the pool for another executor while +// resolving durable object-table metadata. +func (b *DuckBuilder) rewriteObjectRelationsWithSnapshot( + ctx *sql.Context, + root sql.Node, + query string, + execer adapter.SQLExecutor, + conn *stdsql.Conn, +) (string, error) { + if b.provider == nil || root == nil { + return query, nil + } + collector := &tableAndFuncCollector{ctx: ctx} + transform.Walk(collector, root) + routes := make(map[string]string) + bareRoutes := make(map[string]string) + bareHasLocal := make(map[string]bool) + ambiguousBare := make(map[string]bool) + for _, node := range collector.tables { + var table *catalog.Table + switch underlying := node.UnderlyingTable().(type) { + case *catalog.Table: + table = underlying + case *catalog.IndexedTable: + table = underlying.Table + } + if table == nil { + continue + } + physical, object, err := b.provider.PhysicalTableNameForTableWithExecutor(ctx, table, execer, conn) + if err != nil { + return "", err + } + if !object { + bareHasLocal[strings.ToLower(table.Name())] = true + continue + } + schema := "" + if db := node.Database(); db != nil { + schema = db.Name() + } + name := strings.ToLower(table.Name()) + if schema != "" { + routes[strings.ToLower(schema)+"."+name] = physical + } + if ambiguousBare[name] { + continue + } + if previous, exists := bareRoutes[name]; exists && previous != physical { + delete(bareRoutes, name) + ambiguousBare[name] = true + continue + } + bareRoutes[name] = physical + } + for name, physical := range bareRoutes { + if !bareHasLocal[name] && !ambiguousBare[name] { + routes[name] = physical + } + } + rewritten, _ := RewriteSQLRelations(query, routes) + return rewritten, nil +} + // queryForTranslation canonicalizes ANSI-quoted identifiers before SQLGlot // reparses the original MySQL query with its default dialect settings. func queryForTranslation(ctx *sql.Context) string { diff --git a/backend/executor_test.go b/backend/executor_test.go index 5e531bd2..caf744e7 100644 --- a/backend/executor_test.go +++ b/backend/executor_test.go @@ -1,10 +1,15 @@ package backend import ( + "context" + "errors" "testing" + "github.com/apecloud/myduckserver/catalog" + "github.com/dolthub/go-mysql-server/memory" "github.com/dolthub/go-mysql-server/sql" "github.com/dolthub/go-mysql-server/sql/plan" + "github.com/dolthub/go-mysql-server/sql/rowexec" "github.com/dolthub/go-mysql-server/sql/types" "github.com/stretchr/testify/require" ) @@ -67,3 +72,39 @@ func TestQueryForTranslationLeavesDoubleQuotedStringsWhenAnsiQuotesDisabled(t *t ctx = ctx.WithQuery(query) require.Equal(t, query, queryForTranslation(ctx)) } + +func TestStartTransactionBindsUnderlyingDuckDBTransactionBeforeGMSStart(t *testing.T) { + provider := catalog.NewInMemoryDBProvider() + t.Cleanup(func() { require.NoError(t, provider.Close()) }) + + session := NewSession(memory.NewSession(sql.NewBaseSession(), provider), provider) + ctx := sql.NewContext(context.Background(), sql.WithSession(session)) + + base := rowexec.NewBuilder(nil, sql.EngineOverrides{}) + builder := &DuckBuilder{base: base} + iter, err := builder.Build(ctx, plan.NewStartTransaction(sql.ReadWrite), sql.Row{}) + require.NoError(t, err) + require.NoError(t, iter.Close(ctx)) + require.True(t, ctx.GetIgnoreAutoCommit()) + require.NotNil(t, ctx.GetTransaction()) + require.NotNil(t, session.TryGetTxn()) + + require.NoError(t, session.Rollback(ctx, ctx.GetTransaction())) + ctx.SetTransaction(nil) + ctx.SetIgnoreAutoCommit(false) +} + +type startTransactionErrorBuilder struct{} + +func (startTransactionErrorBuilder) Build(*sql.Context, sql.Node, sql.Row) (sql.RowIter, error) { + return nil, errors.New("start transaction sentinel") +} + +func TestStartTransactionRestoresIgnoreAutoCommitOnBuildError(t *testing.T) { + ctx := sql.NewEmptyContext() + builder := &DuckBuilder{base: startTransactionErrorBuilder{}} + + _, err := builder.Build(ctx, plan.NewStartTransaction(sql.ReadWrite), sql.Row{}) + require.EqualError(t, err, "start transaction sentinel") + require.False(t, ctx.GetIgnoreAutoCommit()) +} diff --git a/backend/handler.go b/backend/handler.go index f52cb887..873583d0 100644 --- a/backend/handler.go +++ b/backend/handler.go @@ -16,6 +16,7 @@ package backend import ( "context" + stdsql "database/sql" "errors" "fmt" @@ -41,8 +42,34 @@ type MyHandler struct { } func (h *MyHandler) ConnectionClosed(c *mysql.Conn) { - h.provider.Pool().CloseConn(c.ConnectionID) - h.Handler.ConnectionClosed(c) + var expectedConn *stdsql.Conn + var expectedTx *stdsql.Tx + if h != nil && h.provider != nil && h.provider.Pool() != nil && c != nil { + expectedConn, expectedTx = h.provider.Pool().GetTxnBindingForRelease(c.ConnectionID) + } + // Let GMS remove the session first. Its LifecycleAwareSession callback + // (SessionEnd) must still see the generic pool connection while it rolls back + // a logical autocommit transaction and performs any same-connection cleanup; + // closing the pool entry first could make that cleanup acquire and strand a + // replacement connection. The final pool close then retires the exact entry. + if h != nil && h.Handler != nil { + h.Handler.ConnectionClosed(c) + } + if h != nil && h.provider != nil && h.provider.Pool() != nil && c != nil && expectedConn != nil { + pool := h.provider.Pool() + if expectedTx != nil { + // First honor the complete pair. If SessionEnd already removed the old + // transaction, the nil-transaction binding close below may retire the + // generic connection, but only if it is still idle and identical. + _ = pool.CloseConnIfBinding(c.ConnectionID, expectedConn, expectedTx) + _ = pool.CloseConnIfBinding(c.ConnectionID, expectedConn, nil) + } else { + // The nil transaction is part of the identity check. A replacement + // transaction admitted after the initial snapshot therefore blocks this + // close instead of being rolled back by an unconditional connection close. + _ = pool.CloseConnIfBinding(c.ConnectionID, expectedConn, nil) + } + } } func (h *MyHandler) ComInitDB(c *mysql.Conn, schemaName string) error { diff --git a/backend/loaddata.go b/backend/loaddata.go index 50a86822..ea122fae 100644 --- a/backend/loaddata.go +++ b/backend/loaddata.go @@ -9,7 +9,6 @@ import ( "strconv" "strings" - "github.com/apecloud/myduckserver/adapter" "github.com/apecloud/myduckserver/catalog" "github.com/dolthub/go-mysql-server/sql" "github.com/dolthub/go-mysql-server/sql/plan" @@ -137,6 +136,27 @@ func (db *DuckBuilder) executeLoadData(ctx *sql.Context, insert *plan.InsertInto b.WriteString(" INTO ") qualifiedTableName := catalog.ConnectIdentifiersANSI(insert.Database().Name(), dst.Name()) + // Keep the physical snapshot alive through metadata routing, the INSERT, + // and consumption of the returned OK row. LOAD DATA is built before the + // generic DuckBuilder snapshot path, so releasing here would let a pool + // reset retire the selected connection between any of those steps. + execer, conn, _, snapshotRelease, err := db.executionSnapshotWithLease(ctx) + if err != nil { + return nil, err + } + release := snapshotRelease + defer func() { + if release != nil { + release() + } + }() + if db.provider != nil { + if physical, object, err := db.provider.PhysicalTableNameForTableWithExecutor(ctx, dst, execer, conn); err != nil { + return nil, err + } else if object { + qualifiedTableName = physical + } + } b.WriteString(qualifiedTableName) if len(load.ColNames) > 0 { @@ -205,7 +225,7 @@ func (db *DuckBuilder) executeLoadData(ctx *sql.Context, insert *plan.InsertInto duckSQL := b.String() ctx.GetLogger().Trace(duckSQL) - result, err := adapter.Exec(ctx, duckSQL) + result, err := execer.ExecContext(ctx, duckSQL) if err != nil { return nil, err } @@ -220,10 +240,16 @@ func (db *DuckBuilder) executeLoadData(ctx *sql.Context, insert *plan.InsertInto return nil, err } - return sql.RowsToRowIter(sql.NewRow(types.OkResult{ + iter := sql.RowsToRowIter(sql.NewRow(types.OkResult{ RowsAffected: uint64(affected), InsertID: uint64(insertId), - })), nil + })) + // The statement has no further synchronous work after this point. Transfer + // ownership of the snapshot lease to the result iterator and release it only + // after that iterator closes (including an error or an early close). + iter = wrapDuckLakeOperationIterWithContext(ctx, iter, snapshotRelease) + release = nil + return iter, nil } func singleQuotedDuckChar(s string) string { diff --git a/backend/loaddata_test.go b/backend/loaddata_test.go new file mode 100644 index 00000000..93d33e98 --- /dev/null +++ b/backend/loaddata_test.go @@ -0,0 +1,155 @@ +package backend + +import ( + "context" + stdsql "database/sql" + "io" + "sync/atomic" + "testing" + + "github.com/DATA-DOG/go-sqlmock" + "github.com/apecloud/myduckserver/adapter" + "github.com/dolthub/go-mysql-server/memory" + "github.com/dolthub/go-mysql-server/sql" + "github.com/dolthub/go-mysql-server/sql/plan" + "github.com/dolthub/go-mysql-server/sql/types" + "github.com/stretchr/testify/require" +) + +// loadDataLeaseSession is a small production-shaped session double. It +// exposes the atomic execution snapshot surface while keeping the test +// independent of a live DuckLake extension. +type loadDataLeaseSession struct { + *memory.Session + conn *stdsql.Conn + releaseCalls atomic.Int32 +} + +var _ adapter.ConnectionHolder = (*loadDataLeaseSession)(nil) +var _ adapter.ExecutionSnapshotLeaseHolder = (*loadDataLeaseSession)(nil) + +func (s *loadDataLeaseSession) GetConn(context.Context) (*stdsql.Conn, error) { + return s.conn, nil +} + +func (s *loadDataLeaseSession) GetCatalogConn(context.Context) (*stdsql.Conn, error) { + return s.conn, nil +} + +func (s *loadDataLeaseSession) GetTxn(context.Context, *stdsql.TxOptions) (*stdsql.Tx, error) { + return nil, nil +} + +func (s *loadDataLeaseSession) GetCatalogTxn(context.Context, *stdsql.TxOptions) (*stdsql.Tx, error) { + return nil, nil +} + +func (s *loadDataLeaseSession) TryGetTxn() *stdsql.Tx { return nil } + +func (s *loadDataLeaseSession) GetCurrentCatalog() string { return "memory" } + +func (s *loadDataLeaseSession) GetCurrentSchema() string { return "main" } + +func (s *loadDataLeaseSession) CloseTxn() {} + +func (s *loadDataLeaseSession) CloseConn() {} + +func (s *loadDataLeaseSession) GetExecutionSnapshotLease( + context.Context, + bool, +) (*stdsql.Conn, *stdsql.Tx, func(), error) { + return s.conn, nil, func() { s.releaseCalls.Add(1) }, nil +} + +type loadDataTestTable struct { + name string + schema sql.PrimaryKeySchema +} + +var _ sql.InsertableTable = (*loadDataTestTable)(nil) + +func (t *loadDataTestTable) Name() string { return t.name } + +func (t *loadDataTestTable) String() string { return t.name } + +func (t *loadDataTestTable) Schema(*sql.Context) sql.Schema { return t.schema.Schema } + +func (t *loadDataTestTable) Collation() sql.CollationID { return sql.Collation_Default } + +func (t *loadDataTestTable) Partitions(*sql.Context) (sql.PartitionIter, error) { + return sql.PartitionsToPartitionIter(), nil +} + +func (t *loadDataTestTable) PartitionRows(*sql.Context, sql.Partition) (sql.RowIter, error) { + return nil, nil +} + +func (t *loadDataTestTable) Inserter(*sql.Context) sql.RowInserter { return nil } + +func newLoadDataLeaseFixture(t *testing.T) (*sql.Context, *loadDataLeaseSession, sqlmock.Sqlmock, func()) { + t.Helper() + db, mock, err := sqlmock.New() + require.NoError(t, err) + conn, err := db.Conn(context.Background()) + require.NoError(t, err) + + session := &loadDataLeaseSession{ + Session: memory.NewSession(sql.NewBaseSession(), nil), + conn: conn, + } + session.SetCurrentDatabase("main") + ctx := sql.NewContext(context.Background(), sql.WithSession(session)) + cleanup := func() { + _ = conn.Close() + _ = db.Close() + } + return ctx, session, mock, cleanup +} + +func newLoadDataTestInputs() (*plan.InsertInto, sql.InsertableTable, *plan.LoadData) { + schema := sql.NewPrimaryKeySchema(sql.Schema{ + &sql.Column{Name: "id", Type: types.Int64}, + }) + table := &loadDataTestTable{name: "load_target", schema: schema} + insert := plan.NewInsertInto(memory.NewDatabase("main"), nil, nil, false, nil, nil, false) + load := plan.NewLoadData(false, "input.csv", schema.Schema, nil, nil, 0, "") + return insert, table, load +} + +func expectLoadDataExec(mock sqlmock.Sqlmock) { + mock.ExpectExec(`(?s).*INSERT INTO.*read_csv.*`).WillReturnResult(sqlmock.NewResult(17, 3)) +} + +func TestExecuteLoadDataRetainsSnapshotLeaseUntilResultEOF(t *testing.T) { + ctx, session, mock, cleanup := newLoadDataLeaseFixture(t) + defer cleanup() + expectLoadDataExec(mock) + insert, table, load := newLoadDataTestInputs() + + iter, err := (&DuckBuilder{}).executeLoadData(ctx, insert, table, load, "input.csv") + require.NoError(t, err) + require.Zero(t, session.releaseCalls.Load(), "the snapshot must cover the returned result iterator") + + _, err = iter.Next(ctx) + require.NoError(t, err) + require.Zero(t, session.releaseCalls.Load(), "the lease must remain while the OK row is being consumed") + + _, err = iter.Next(ctx) + require.ErrorIs(t, err, io.EOF) + require.Equal(t, int32(1), session.releaseCalls.Load()) + require.NoError(t, iter.Close(ctx)) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestExecuteLoadDataReleasesSnapshotLeaseOnEarlyClose(t *testing.T) { + ctx, session, mock, cleanup := newLoadDataLeaseFixture(t) + defer cleanup() + expectLoadDataExec(mock) + insert, table, load := newLoadDataTestInputs() + + iter, err := (&DuckBuilder{}).executeLoadData(ctx, insert, table, load, "input.csv") + require.NoError(t, err) + require.ErrorIs(t, iter.Close(ctx), errDuckLakeOperationClosedBeforeEOF) + require.Equal(t, int32(1), session.releaseCalls.Load()) + require.NoError(t, mock.ExpectationsWereMet()) +} diff --git a/backend/mysql_handler_lifecycle_test.go b/backend/mysql_handler_lifecycle_test.go new file mode 100644 index 00000000..5ea0fa91 --- /dev/null +++ b/backend/mysql_handler_lifecycle_test.go @@ -0,0 +1,134 @@ +package backend + +import ( + "context" + stdsql "database/sql" + "net" + "sync/atomic" + "testing" + + "github.com/apecloud/myduckserver/catalog" + "github.com/dolthub/go-mysql-server/server" + "github.com/dolthub/go-mysql-server/sql" + "github.com/dolthub/vitess/go/mysql" + "github.com/stretchr/testify/require" +) + +// replacementOnSessionEnd models SessionManager replacing a connection's GMS +// session while MyHandler.ConnectionClosed is between its initial binding +// snapshot and the final pool close. It uses only public pool operations: the +// old raw transaction is finalized, its generic entry is retired, and a new +// connection is acquired under the same numeric session ID. +type replacementOnSessionEnd struct { + *Session + pool *catalog.ConnectionPool + id uint32 + replacement *stdsql.Conn + err error +} + +func (s *replacementOnSessionEnd) SessionEnd() { + s.Session.SessionEnd() + if s.err != nil { + return + } + if s.err = s.pool.CloseConn(s.id); s.err != nil { + return + } + s.replacement, s.err = s.pool.GetConn(context.Background(), s.id) +} + +// MyHandler must close the exact transaction/owner pair it observed before +// GMS removes the session. If GMS installs a replacement generic connection +// during that callback, the stale pair cleanup must leave the replacement +// usable instead of closing it by numeric connection ID. +func TestMyHandlerConnectionClosedPreservesReplacementConnection(t *testing.T) { + provider := catalog.NewInMemoryDBProvider() + t.Cleanup(func() { require.NoError(t, provider.Close()) }) + engine, _ := NewEngine(provider) + + listener, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + + const connectionID uint32 = 313 + baseBuilder := NewSessionBuilder(provider) + var replacementSession *replacementOnSessionEnd + builder := func(ctx context.Context, conn *mysql.Conn, addr string) (sql.Session, error) { + built, buildErr := baseBuilder(ctx, conn, addr) + if buildErr != nil { + return nil, buildErr + } + session, ok := built.(*Session) + if !ok { + t.Fatalf("session builder returned %T, want *Session", built) + } + replacementSession = &replacementOnSessionEnd{ + Session: session, + pool: provider.Pool(), + id: connectionID, + } + return replacementSession, nil + } + + var handler *MyHandler + srv, err := server.NewServerWithHandler( + server.Config{ + Listener: listener, + Protocol: "tcp", + Address: listener.Addr().String(), + DisableConnectionWatcher: true, + }, + engine, + sql.NewContext, + builder, + nil, + func(h mysql.Handler) (mysql.Handler, error) { + wrapped, wrapErr := WrapHandler(provider, engine, nil, false)(h) + if wrapErr == nil { + handler = wrapped.(*MyHandler) + } + return wrapped, wrapErr + }, + ) + require.NoError(t, err) + t.Cleanup(func() { + _ = srv.Close() + _ = listener.Close() + }) + require.NotNil(t, handler) + + clientConn, peerConn := net.Pipe() + t.Cleanup(func() { + _ = clientConn.Close() + _ = peerConn.Close() + }) + mysqlConn := &mysql.Conn{Conn: clientConn, ConnectionID: connectionID} + sm := srv.SessionManager() + sm.AddConn(mysqlConn) + require.NoError(t, sm.NewSession(context.Background(), mysqlConn)) + require.NotNil(t, replacementSession) + + tx, err := replacementSession.GetCatalogTxn(context.Background(), nil) + require.NoError(t, err) + owner, bound := replacementSession.GetTxnBindingForRelease() + require.NotNil(t, owner) + require.Same(t, tx, bound) + var callbacks atomic.Int32 + require.True(t, replacementSession.RegisterPostPhysicalTransactionCleanup(nil, tx, func(success bool) { + require.False(t, success) + callbacks.Add(1) + })) + + handler.ConnectionClosed(mysqlConn) + + require.NoError(t, replacementSession.err) + require.NotNil(t, replacementSession.replacement) + require.NotSame(t, owner, replacementSession.replacement) + require.Equal(t, int32(1), callbacks.Load()) + currentConn, currentTx := provider.Pool().GetTxnBindingForRelease(connectionID) + require.Same(t, replacementSession.replacement, currentConn) + require.Nil(t, currentTx) + var got int + require.NoError(t, replacementSession.replacement.QueryRowContext(context.Background(), "SELECT 1").Scan(&got)) + require.Equal(t, 1, got) +} diff --git a/backend/raw_session_lifecycle_test.go b/backend/raw_session_lifecycle_test.go new file mode 100644 index 00000000..93901479 --- /dev/null +++ b/backend/raw_session_lifecycle_test.go @@ -0,0 +1,126 @@ +package backend + +import ( + "context" + stdsql "database/sql" + "sync/atomic" + "testing" + + "github.com/apecloud/myduckserver/catalog" + "github.com/dolthub/go-mysql-server/memory" + "github.com/dolthub/go-mysql-server/sql" + "github.com/stretchr/testify/require" +) + +func newRawLifecycleSession(t *testing.T, id uint32) (*catalog.DatabaseProvider, *Session, *stdsql.Conn, *stdsql.Tx) { + t.Helper() + provider := catalog.NewInMemoryDBProvider() + session := NewSession(memory.NewSession(sql.NewBaseSession(), provider), provider) + session.SetConnectionId(id) + tx, err := session.GetCatalogTxn(context.Background(), nil) + require.NoError(t, err) + owner, bound := session.GetTxnBindingForRelease() + require.NotNil(t, owner) + require.Same(t, tx, bound) + return provider, session, owner, tx +} + +// PostgreSQL transaction controls own a raw database/sql transaction without +// placing a GMS transaction in Session.GetTransaction. SessionEnd must still +// roll it back while retaining the generic pool connection for subsequent +// lifecycle cleanup. +func TestSessionEndRawTransactionRollsBackAndRetainsGenericConnection(t *testing.T) { + provider, session, owner, tx := newRawLifecycleSession(t, 311) + t.Cleanup(func() { require.NoError(t, provider.Close()) }) + + var callbacks atomic.Int32 + var success atomic.Bool + require.True(t, session.RegisterPostPhysicalTransactionCleanup(nil, tx, func(ok bool) { + callbacks.Add(1) + success.Store(ok) + })) + // Registration on an active raw transaction must only arm the bridge. A + // reversed pool-result branch would complete this callback immediately and + // release a COPY/operation lease before the physical transaction ends. + require.Zero(t, callbacks.Load()) + + // Do not call SetTransaction: this is the raw PostgreSQL path. + session.SessionEnd() + + require.Equal(t, int32(1), callbacks.Load()) + require.False(t, success.Load()) + currentConn, currentTx := session.GetTxnBindingForRelease() + require.Same(t, owner, currentConn) + require.Nil(t, currentTx) + + // The generic connection remains usable after raw rollback and is not + // replaced merely because the transaction binding was finalized. + conn, err := session.GetCatalogConn(context.Background()) + require.NoError(t, err) + require.Same(t, owner, conn) + var got int + require.NoError(t, conn.QueryRowContext(context.Background(), "SELECT 1").Scan(&got)) + require.Equal(t, 1, got) +} + +// A stale SessionEnd may run after another lifecycle path has already removed +// the raw binding and installed a replacement generic connection. It must +// finish only its captured transaction and never close the replacement. +func TestSessionEndRawBindingMismatchKeepsReplacementConnection(t *testing.T) { + provider, session, owner, tx := newRawLifecycleSession(t, 312) + t.Cleanup(func() { require.NoError(t, provider.Close()) }) + + var callbacks atomic.Int32 + require.True(t, session.RegisterPostPhysicalTransactionCleanup(nil, tx, func(bool) { + callbacks.Add(1) + })) + + // Model an earlier finalizer that completed the raw transaction and removed + // its pool mapping, but did not clear this stale Session object's snapshot. + require.NoError(t, tx.Rollback()) + provider.Pool().CloseTxn(session.ID()) + require.NoError(t, provider.Pool().CloseConn(session.ID())) + replacement, err := provider.Pool().GetConn(context.Background(), session.ID()) + require.NoError(t, err) + require.NotNil(t, replacement) + require.NotSame(t, owner, replacement) + require.Equal(t, int32(1), callbacks.Load()) + + session.SessionEnd() + + currentConn, currentTx := provider.Pool().GetTxnBindingForRelease(session.ID()) + require.Same(t, replacement, currentConn) + require.Nil(t, currentTx) + require.Equal(t, int32(1), callbacks.Load(), "stale SessionEnd must not complete the callback twice") +} + +func TestPoolBridgeDrainsSessionCallbacksBeforeRepanic(t *testing.T) { + provider, session, owner, tx := newRawLifecycleSession(t, 313) + t.Cleanup(func() { require.NoError(t, provider.Close()) }) + + var firstCalls atomic.Int32 + var secondCalls atomic.Int32 + require.True(t, session.RegisterPostPhysicalTransactionCleanup(nil, tx, func(bool) { + firstCalls.Add(1) + panic("pool bridge cleanup panic") + })) + require.True(t, session.RegisterPostPhysicalTransactionCleanup(nil, tx, func(bool) { + secondCalls.Add(1) + })) + + require.PanicsWithValue(t, "pool bridge cleanup panic", func() { + _, _, _ = session.RollbackTxn(tx) + }) + require.Equal(t, int32(1), firstCalls.Load()) + require.Equal(t, int32(1), secondCalls.Load()) + currentConn, currentTx := provider.Pool().GetTxnBindingForRelease(session.ID()) + require.Same(t, owner, currentConn) + require.Nil(t, currentTx) + + var lateCalls atomic.Int32 + require.True(t, session.RegisterPostPhysicalTransactionCleanup(nil, tx, func(success bool) { + require.False(t, success) + lateCalls.Add(1) + })) + require.Equal(t, int32(1), lateCalls.Load()) +} diff --git a/backend/relation_rewrite.go b/backend/relation_rewrite.go new file mode 100644 index 00000000..9c708648 --- /dev/null +++ b/backend/relation_rewrite.go @@ -0,0 +1,382 @@ +package backend + +import ( + "sort" + "strings" + "unicode" +) + +// RewriteSQLRelations replaces table references in relation positions. It is +// intentionally a small lexer rather than strings.Replace: quoted literals, +// dollar-quoted bodies, and comments may contain table-looking text and must +// remain untouched. +// +// Routes are keyed by an unquoted, case-insensitive relation name. A +// qualified key (for example, "app.orders") takes precedence over a bare +// table key. Values are emitted verbatim and are expected to be trusted, +// already-quoted DuckDB identifiers supplied by the catalog layer. +func RewriteSQLRelations(query string, routes map[string]string) (string, bool) { + if query == "" || len(routes) == 0 { + return query, false + } + tokens := lexSQLForRelations(query) + if len(tokens) == 0 { + return query, false + } + spans := relationSpans(tokens) + if len(spans) == 0 { + return query, false + } + + // A statement can expose the same relation through more than one context + // (for example an INSERT ... SELECT). Keep spans ordered and reject an + // accidental overlap before constructing the output. + sort.SliceStable(spans, func(i, j int) bool { return spans[i].start < spans[j].start }) + var out strings.Builder + out.Grow(len(query)) + last := 0 + changed := false + for _, span := range spans { + if span.start < last || span.start < 0 || span.end > len(query) || span.end <= span.start { + continue + } + key := strings.ToLower(strings.Join(span.parts, ".")) + replacement, ok := routes[key] + if !ok && len(span.parts) > 1 { + replacement, ok = routes[strings.ToLower(span.parts[len(span.parts)-1])] + } + if !ok || replacement == "" { + continue + } + out.WriteString(query[last:span.start]) + out.WriteString(replacement) + last = span.end + changed = true + } + if !changed { + return query, false + } + out.WriteString(query[last:]) + return out.String(), true +} + +type relationTokenKind uint8 + +const ( + relationWord relationTokenKind = iota + relationIdentifier + relationDot + relationOther +) + +type relationToken struct { + kind relationTokenKind + value string + start, end int +} + +type relationSpan struct { + start, end int + parts []string +} + +func relationSpans(tokens []relationToken) []relationSpan { + var spans []relationSpan + for i := 0; i < len(tokens); i++ { + token := tokens[i] + if token.kind != relationWord { + continue + } + switch token.value { + case "from", "join", "using", "update", "into": + span, next := relationList(tokens, i+1, true) + spans = append(spans, span...) + if next > i+1 { + i = next - 1 + } + case "truncate": + j := i + 1 + j = skipWords(tokens, j, "table", "only") + span, next := relationList(tokens, j, false) + spans = append(spans, span...) + if next > j { + i = next - 1 + } + case "drop": + j := i + 1 + if !wordAt(tokens, j, "table") { + continue + } + j++ + j = skipWords(tokens, j, "if", "exists", "only") + span, next := relationList(tokens, j, false) + spans = append(spans, span...) + if next > j { + i = next - 1 + } + case "alter": + j := i + 1 + if !wordAt(tokens, j, "table") { + continue + } + j++ + j = skipWords(tokens, j, "if", "exists", "only") + span, _ := relationList(tokens, j, false) + if len(span) > 0 { + spans = append(spans, span[0]) + } + case "copy": + // COPY (SELECT ...) ... is a query source, not a table target. + j := i + 1 + if tokenAt(tokens, j).kind == relationOther && tokenAt(tokens, j).value == "(" { + continue + } + span, _ := relationList(tokens, j, false) + if len(span) > 0 { + spans = append(spans, span[0]) + } + case "create": + // CREATE INDEX ... ON table (...). Restrict this to INDEX so an + // ON clause in another CREATE statement is not mistaken for a + // relation target. + if !wordAt(tokens, i+1, "unique") && !wordAt(tokens, i+1, "index") { + continue + } + j := i + 1 + if wordAt(tokens, j, "unique") { + j++ + } + for j < len(tokens) && !(tokens[j].kind == relationWord && tokens[j].value == "on") { + j++ + } + if j < len(tokens) { + span, _ := relationList(tokens, j+1, false) + if len(span) > 0 { + spans = append(spans, span[0]) + } + } + } + } + return spans +} + +func relationList(tokens []relationToken, start int, aliases bool) ([]relationSpan, int) { + var spans []relationSpan + i := start + for { + i = skipWords(tokens, i, "only", "lateral") + nameStart := i + if !isRelationNameToken(tokenAt(tokens, i)) { + break + } + parts := []string{tokens[i].value} + nameEnd := i + 1 + for nameEnd+1 < len(tokens) && tokens[nameEnd].kind == relationDot && isRelationNameToken(tokens[nameEnd+1]) { + parts = append(parts, tokens[nameEnd+1].value) + nameEnd += 2 + } + spans = append(spans, relationSpan{ + start: tokens[nameStart].start, + end: tokens[nameEnd-1].end, + parts: parts, + }) + + // Skip an optional alias, but leave it outside the replacement span. + listEnd := nameEnd + if aliases && wordAt(tokens, listEnd, "as") { + if isRelationNameToken(tokenAt(tokens, listEnd+1)) { + listEnd += 2 + } + } else if aliases && isRelationNameToken(tokenAt(tokens, listEnd)) { + listEnd++ + } + if !commaAt(tokens, listEnd) { + return spans, listEnd + } + i = listEnd + 1 + } + return spans, i +} + +func skipWords(tokens []relationToken, i int, words ...string) int { + for _, word := range words { + if wordAt(tokens, i, word) { + i++ + } + } + return i +} + +func tokenAt(tokens []relationToken, i int) relationToken { + if i < 0 || i >= len(tokens) { + return relationToken{kind: relationOther} + } + return tokens[i] +} + +func wordAt(tokens []relationToken, i int, value string) bool { + token := tokenAt(tokens, i) + return token.kind == relationWord && token.value == value +} + +func commaAt(tokens []relationToken, i int) bool { + token := tokenAt(tokens, i) + return token.kind == relationOther && token.value == "," +} + +func isRelationNameToken(token relationToken) bool { + if token.kind == relationIdentifier { + return true + } + if token.kind != relationWord { + return false + } + // An unquoted word is an identifier unless it begins the next clause. + // Quoted keywords are relationIdentifier tokens and remain valid names. + switch token.value { + case "select", "with", "where", "on", "join", "left", "right", "full", "inner", "cross", "natural", "group", "order", "limit", "having", "union", "returning", "set", "values", "from", "using", "into", "table", "if", "exists", "only", "as", "cascade", "restrict", "add", "drop", "alter", "rename", "column", "to": + return false + default: + return true + } +} + +func lexSQLForRelations(query string) []relationToken { + var tokens []relationToken + for i := 0; i < len(query); { + if isSQLSpace(query[i]) { + i++ + continue + } + start := i + if query[i] == '$' { + if end, ok := dollarQuoteEnd(query, i); ok { + i = end + tokens = append(tokens, relationToken{kind: relationOther, start: start, end: i}) + continue + } + } + switch query[i] { + case '\'', '"', '`': + quote := query[i] + i++ + for i < len(query) { + if query[i] == '\\' && quote != '`' && i+1 < len(query) { + i += 2 + continue + } + if query[i] == quote { + i++ + if i < len(query) && query[i] == quote { + i++ + continue + } + break + } + i++ + } + if quote == '"' || quote == '`' { + end := i + if end > start+1 && query[end-1] == quote { + end-- + } + value := strings.ReplaceAll(query[start+1:end], string(quote)+string(quote), string(quote)) + tokens = append(tokens, relationToken{kind: relationIdentifier, value: value, start: start, end: i}) + } else { + tokens = append(tokens, relationToken{kind: relationOther, start: start, end: i}) + } + case '-', '#': + if query[i] == '#' || (i+1 < len(query) && query[i+1] == '-') { + for i < len(query) && query[i] != '\n' { + i++ + } + continue + } + i++ + case '/': + if i+1 < len(query) && query[i+1] == '*' { + i += 2 + depth := 1 + for i < len(query) && depth > 0 { + if i+1 < len(query) && query[i] == '/' && query[i+1] == '*' { + depth++ + i += 2 + continue + } + if i+1 < len(query) && query[i] == '*' && query[i+1] == '/' { + depth-- + i += 2 + continue + } + i++ + } + continue + } + fallthrough + case '.': + i++ + kind := relationOther + if query[start] == '.' { + kind = relationDot + } + value := "" + if kind == relationOther { + value = query[start:i] + } + tokens = append(tokens, relationToken{kind: kind, value: value, start: start, end: i}) + default: + if isSQLIdentStart(query[i]) { + i++ + for i < len(query) && isSQLIdentPart(query[i]) { + i++ + } + value := strings.ToLower(query[start:i]) + tokens = append(tokens, relationToken{kind: relationWord, value: value, start: start, end: i}) + continue + } + i++ + value := query[start:i] + tokens = append(tokens, relationToken{kind: relationOther, value: value, start: start, end: i}) + } + } + return tokens +} + +func dollarQuoteEnd(query string, start int) (int, bool) { + if start >= len(query) || query[start] != '$' { + return 0, false + } + i := start + 1 + if i < len(query) && query[i] == '$' { + i++ + } else { + if i >= len(query) || !(query[i] == '_' || query[i] >= 'a' && query[i] <= 'z' || query[i] >= 'A' && query[i] <= 'Z') { + return 0, false + } + i++ + for i < len(query) && (query[i] == '_' || query[i] >= 'a' && query[i] <= 'z' || query[i] >= 'A' && query[i] <= 'Z' || query[i] >= '0' && query[i] <= '9') { + i++ + } + if i >= len(query) || query[i] != '$' { + return 0, false + } + i++ + } + tag := query[start:i] + if end := strings.Index(query[i:], tag); end >= 0 { + return i + end + len(tag), true + } + // An unterminated dollar quote is treated as a non-relation token so + // identifiers inside it cannot be rewritten accidentally. + return len(query), true +} + +func isSQLSpace(c byte) bool { return unicode.IsSpace(rune(c)) } + +func isSQLIdentStart(c byte) bool { + return c == '_' || c == '$' || c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z' +} + +func isSQLIdentPart(c byte) bool { + return isSQLIdentStart(c) || c >= '0' && c <= '9' +} diff --git a/backend/relation_rewrite_test.go b/backend/relation_rewrite_test.go new file mode 100644 index 00000000..0d4d05d8 --- /dev/null +++ b/backend/relation_rewrite_test.go @@ -0,0 +1,106 @@ +package backend + +import "testing" + +func TestRewriteSQLRelationsSkipsLiteralsAndComments(t *testing.T) { + query := "SELECT 'FROM object_table', \"FROM object_table\" AS text FROM object_table AS o /* JOIN object_table */ JOIN local_table l ON l.id = o.id -- FROM object_table\n" + want := "SELECT 'FROM object_table', \"FROM object_table\" AS text FROM \"__myduck_ducklake\".\"app\".\"object_table\" AS o /* JOIN object_table */ JOIN local_table l ON l.id = o.id -- FROM object_table\n" + routes := map[string]string{ + "object_table": `"__myduck_ducklake"."app"."object_table"`, + } + got, changed := RewriteSQLRelations(query, routes) + if !changed { + t.Fatal("expected relation rewrite") + } + if got != want { + t.Fatalf("rewritten query = %q, want %q", got, want) + } +} + +func TestRewriteSQLRelationsPrefersQualifiedRouteAndPreservesLocal(t *testing.T) { + query := `SELECT * FROM "app"."orders" o JOIN "orders" local ON local.id = o.id` + routes := map[string]string{ + "app.orders": `"__myduck_ducklake"."app"."orders"`, + } + got, changed := RewriteSQLRelations(query, routes) + want := `SELECT * FROM "__myduck_ducklake"."app"."orders" o JOIN "orders" local ON local.id = o.id` + if !changed || got != want { + t.Fatalf("rewrite = (%q, %v), want (%q, true)", got, changed, want) + } +} + +func TestRewriteSQLRelationsHandlesCommaSeparatedRelations(t *testing.T) { + query := `SELECT * FROM object_table AS object, other_table WHERE object.id = other_table.id` + routes := map[string]string{ + "object_table": `"__myduck_ducklake"."app"."object_table"`, + "other_table": `"__myduck_ducklake"."app"."other_table"`, + } + want := `SELECT * FROM "__myduck_ducklake"."app"."object_table" AS object, "__myduck_ducklake"."app"."other_table" WHERE object.id = other_table.id` + got, changed := RewriteSQLRelations(query, routes) + if !changed || got != want { + t.Fatalf("rewrite = (%q, %v), want (%q, true)", got, changed, want) + } +} + +func TestRewriteSQLRelationsHandlesDMLTargets(t *testing.T) { + routes := map[string]string{"object_table": `"__myduck_ducklake"."app"."object_table"`} + for _, query := range []string{ + `INSERT INTO object_table VALUES (1)`, + `UPDATE object_table SET value = 2`, + `DELETE FROM object_table WHERE id = 1`, + `SELECT * FROM object_table JOIN other_table ON true`, + } { + got, changed := RewriteSQLRelations(query, routes) + if !changed || got == query { + t.Errorf("query %q was not rewritten: %q, %v", query, got, changed) + } + } +} + +func TestRewriteSQLRelationsHandlesDDLAndCopyTargets(t *testing.T) { + routes := map[string]string{ + "object_table": `"__myduck_ducklake"."app"."object_table"`, + "other_table": `"__myduck_ducklake"."app"."other_table"`, + } + tests := []struct { + query string + want string + }{ + { + query: `TRUNCATE TABLE ONLY object_table, other_table`, + want: `TRUNCATE TABLE ONLY "__myduck_ducklake"."app"."object_table", "__myduck_ducklake"."app"."other_table"`, + }, + { + query: `DROP TABLE IF EXISTS "app"."object_table", other_table CASCADE`, + want: `DROP TABLE IF EXISTS "__myduck_ducklake"."app"."object_table", "__myduck_ducklake"."app"."other_table" CASCADE`, + }, + { + query: `ALTER TABLE IF EXISTS ONLY object_table ADD COLUMN n INT`, + want: `ALTER TABLE IF EXISTS ONLY "__myduck_ducklake"."app"."object_table" ADD COLUMN n INT`, + }, + { + query: `COPY object_table FROM STDIN`, + want: `COPY "__myduck_ducklake"."app"."object_table" FROM STDIN`, + }, + { + query: `CREATE INDEX idx ON object_table (id)`, + want: `CREATE INDEX idx ON "__myduck_ducklake"."app"."object_table" (id)`, + }, + } + for _, test := range tests { + got, changed := RewriteSQLRelations(test.query, routes) + if !changed || got != test.want { + t.Errorf("rewrite(%q) = (%q, %v), want (%q, true)", test.query, got, changed, test.want) + } + } +} + +func TestRewriteSQLRelationsSkipsDollarQuotedBodies(t *testing.T) { + query := `DO $$ BEGIN INSERT INTO object_table VALUES (1); END $$; SELECT * FROM object_table` + routes := map[string]string{"object_table": `"__myduck_ducklake"."app"."object_table"`} + want := `DO $$ BEGIN INSERT INTO object_table VALUES (1); END $$; SELECT * FROM "__myduck_ducklake"."app"."object_table"` + got, changed := RewriteSQLRelations(query, routes) + if !changed || got != want { + t.Fatalf("rewrite = (%q, %v), want (%q, true)", got, changed, want) + } +} diff --git a/backend/session.go b/backend/session.go index bc108a6a..88b089a1 100644 --- a/backend/session.go +++ b/backend/session.go @@ -16,15 +16,19 @@ package backend import ( "context" stdsql "database/sql" + "errors" "fmt" "strconv" "strings" + "sync" "sync/atomic" + "time" "github.com/sirupsen/logrus" adapter "github.com/apecloud/myduckserver/adapter" "github.com/apecloud/myduckserver/catalog" + "github.com/apecloud/myduckserver/mycontext" "github.com/dolthub/go-mysql-server/memory" "github.com/dolthub/go-mysql-server/sql" "github.com/dolthub/go-mysql-server/sql/plan" @@ -36,6 +40,329 @@ type Session struct { db *catalog.DatabaseProvider queryRowLimit uint64 mysqlImplicitDDL atomic.Bool + // physicalBinding is the last connection/transaction identity acquired by + // this Session object. SessionManager can invoke SessionEnd on an old + // session after a connection ID has been reused; retaining the identity lets + // lifecycle cleanup fail closed instead of rolling back the replacement. + bindingMu sync.Mutex + boundConn *stdsql.Conn + boundTx *stdsql.Tx + bindingSeen bool + // sessionEndMu serializes lifecycle callbacks from SessionManager. GMS may + // call SessionEnd while replacing a connection session and can call it again + // during disconnect cleanup; only one callback may consume the current + // transaction/implicit marker at a time. + sessionEndMu sync.Mutex + implicitMu sync.Mutex + implicit *implicitTransactionCleanup + // postPhysical retains callbacks for protocol paths (notably PostgreSQL) + // that own a database/sql transaction without creating a GMS Transaction. + // Entries stay addressable after completion so a registration racing the + // finalizer can be invoked immediately with the already-known outcome. + postPhysicalMu sync.Mutex + postPhysical map[*stdsql.Tx]*postTransactionCleanup + postPhysicalOrder []*stdsql.Tx + // procedureScopes holds markers detached by GMS's stored-procedure + // transaction swap. The newest procedure scope is always at the end; this + // lets nested CALLs restore their exact outer transaction without allowing + // an inner marker to overwrite it. + procedureScopes []*procedureTransactionScope + cleanupErr error +} + +// maxPostPhysicalTombstones bounds the raw-transaction callback registry. A +// short recent window preserves late registration while preventing a long-lived +// PostgreSQL session from retaining every completed *sql.Tx forever. +const maxPostPhysicalTombstones = 1024 + +// postPhysicalTombstoneTTL retains a completed raw transaction long enough for +// a late protocol finalizer to observe its outcome, without keeping every +// database/sql transaction alive for the lifetime of a session. +const postPhysicalTombstoneTTL = 5 * time.Minute + +// ImplicitTransactionCleanupRegistrar is implemented by sessions that can +// retain a callback for the GMS autocommit transaction. GMS clears an +// implicitly-created transaction with SetTransaction(nil), which has no error +// return and no context argument; the callback gives the session a chance to +// close the statement iterator and release its operation/connection leases +// before it performs the compensating rollback. +type ImplicitTransactionCleanupRegistrar interface { + RegisterImplicitTransactionCleanup(*sql.Context, sql.Transaction, func() error) func(success bool) +} + +// PostTransactionCleanupRegistrar is implemented by sessions that can retain +// a callback until a logical GMS transaction has physically and logically +// finished. The callback receives true only for a successful commit; rollback +// and every failed finalization receive false. +type PostTransactionCleanupRegistrar interface { + RegisterPostTransactionCleanup(*sql.Context, sql.Transaction, func(success bool)) bool +} + +// PhysicalPostTransactionCleanupRegistrar is the protocol-facing counterpart +// for a raw database/sql transaction which is finalized outside GMS's +// TransactionSession callbacks. +type PhysicalPostTransactionCleanupRegistrar interface { + RegisterPostPhysicalTransactionCleanup(*sql.Context, *stdsql.Tx, func(success bool)) bool +} + +// postTransactionCleanup is a small one-shot callback registry. Registration +// and completion synchronize on the same mutex, so a callback can never be +// stranded between a finalizer's callback snapshot and its completion bit. +type postTransactionCleanup struct { + mu sync.Mutex + callbacks []func(bool) + completed bool + success bool + tombstone bool + completedAt time.Time + poolBridgeRegistered bool +} + +func completedPostTransactionCleanup(success bool) *postTransactionCleanup { + return &postTransactionCleanup{completed: true, success: success, completedAt: time.Now()} +} + +func (c *postTransactionCleanup) register(callback func(bool)) bool { + if c == nil || callback == nil { + return false + } + c.mu.Lock() + if !c.completed { + c.callbacks = append(c.callbacks, callback) + c.mu.Unlock() + return true + } + success := c.success + c.mu.Unlock() + // Completion already happened. Invoke outside the lock so callbacks may + // safely register/release other lifecycle state. + callback(success) + return true +} + +func (c *postTransactionCleanup) complete(success bool) { + panicValue, panicked := c.completeCollectPanic(success) + if panicked { + panic(panicValue) + } +} + +// completeCollectPanic drains every callback even when one of them panics. A +// physical pool completion reaches Session through one bridge callback, so the +// pool cannot see or recover the remaining session callbacks individually. +// Preserve the first panic for the caller after all local resources are +// released. +func (c *postTransactionCleanup) completeCollectPanic(success bool) (panicValue any, panicked bool) { + if c == nil { + return nil, false + } + c.mu.Lock() + if c.completed { + c.mu.Unlock() + return nil, false + } + c.completed = true + c.success = success + c.completedAt = time.Now() + callbacks := append([]func(bool){}, c.callbacks...) + c.callbacks = nil + c.mu.Unlock() + for _, callback := range callbacks { + if callback == nil { + continue + } + value, didPanic := callPostTransactionCleanupCallback(func() { + callback(success) + }) + if didPanic && !panicked { + panicValue, panicked = value, true + } + } + return panicValue, panicked +} + +func callPostTransactionCleanupCallback(callback func()) (panicValue any, panicked bool) { + defer func() { + if recovered := recover(); recovered != nil { + panicValue, panicked = recovered, true + } + }() + callback() + return nil, false +} + +func (c *postTransactionCleanup) markTombstone() bool { + if c == nil { + return false + } + c.mu.Lock() + defer c.mu.Unlock() + if !c.completed || c.tombstone { + return false + } + if c.completedAt.IsZero() { + c.completedAt = time.Now() + } + c.tombstone = true + return true +} + +func (c *postTransactionCleanup) isCompleted() bool { + if c == nil { + return false + } + c.mu.Lock() + completed := c.completed + c.mu.Unlock() + return completed +} + +func (c *postTransactionCleanup) completionSnapshot() (bool, time.Time) { + if c == nil { + return false, time.Time{} + } + c.mu.Lock() + completed, completedAt := c.tombstone, c.completedAt + c.mu.Unlock() + return completed, completedAt +} + +func (sess *Session) recordPostPhysicalCompletion(tx *stdsql.Tx, state *postTransactionCleanup) { + if sess == nil || tx == nil || state == nil { + return + } + sess.postPhysicalMu.Lock() + defer sess.postPhysicalMu.Unlock() + if sess.postPhysical == nil || sess.postPhysical[tx] != state { + return + } + // Keep the registry's lock order consistent (session map, then state). A + // concurrent registration may inspect the state while holding the map lock; + // taking state first here would create an avoidable cycle. + if !state.markTombstone() { + return + } + sess.postPhysicalOrder = append(sess.postPhysicalOrder, tx) + sess.prunePostPhysicalCompletionsLocked(time.Now()) +} + +// prunePostPhysicalCompletionsLocked removes only completed raw transaction +// tombstones that have aged out or exceeded the FIFO bound. Pending callback +// states remain in the map so a provider teardown cannot silently strand a +// lease. The caller holds postPhysicalMu. +func (sess *Session) prunePostPhysicalCompletionsLocked(now time.Time) { + if sess == nil { + return + } + if len(sess.postPhysical) == 0 { + sess.postPhysicalOrder = nil + return + } + for tx, state := range sess.postPhysical { + completed, completedAt := state.completionSnapshot() + if completed && !completedAt.IsZero() && now.Sub(completedAt) >= postPhysicalTombstoneTTL { + delete(sess.postPhysical, tx) + } + } + order := make([]*stdsql.Tx, 0, len(sess.postPhysicalOrder)) + seen := make(map[*stdsql.Tx]struct{}, len(sess.postPhysicalOrder)) + for _, tx := range sess.postPhysicalOrder { + if _, duplicate := seen[tx]; duplicate { + continue + } + state, present := sess.postPhysical[tx] + if !present { + continue + } + completed, _ := state.completionSnapshot() + if !completed { + // Never remove a pending state merely because it appeared in a stale + // order slice. + continue + } + seen[tx] = struct{}{} + order = append(order, tx) + } + for tx, state := range sess.postPhysical { + completed, completedAt := state.completionSnapshot() + if !completed { + continue + } + if !completedAt.IsZero() && now.Sub(completedAt) >= postPhysicalTombstoneTTL { + delete(sess.postPhysical, tx) + continue + } + if _, present := seen[tx]; !present { + seen[tx] = struct{}{} + order = append(order, tx) + } + } + for len(order) > maxPostPhysicalTombstones { + oldest := order[0] + order = order[1:] + if state := sess.postPhysical[oldest]; state != nil { + completed, _ := state.completionSnapshot() + if completed { + delete(sess.postPhysical, oldest) + } + } + } + sess.postPhysicalOrder = order +} + +type implicitTransactionCleanup struct { + tx *Transaction + ctx *sql.Context + + mu sync.Mutex + callbacks []func() error + cleanupStarted bool + cleanupOnce sync.Once + cleanupErr error + suspended bool + procedureOwned bool +} + +type procedureTransactionScope struct { + suspended *implicitTransactionCleanup +} + +func (c *implicitTransactionCleanup) add(callback func() error) bool { + if c == nil || callback == nil { + return false + } + c.mu.Lock() + defer c.mu.Unlock() + if c.cleanupStarted { + return false + } + c.callbacks = append(c.callbacks, callback) + return true +} + +func (c *implicitTransactionCleanup) run() error { + if c == nil { + return nil + } + c.cleanupOnce.Do(func() { + c.mu.Lock() + c.cleanupStarted = true + callbacks := append([]func() error(nil), c.callbacks...) + c.mu.Unlock() + + var cleanupErr error + for _, callback := range callbacks { + if callback != nil { + cleanupErr = errors.Join(cleanupErr, callback()) + } + } + c.mu.Lock() + c.cleanupErr = cleanupErr + c.mu.Unlock() + }) + c.mu.Lock() + defer c.mu.Unlock() + return c.cleanupErr } type SessionOption func(*Session) @@ -61,6 +388,671 @@ func (sess *Session) Provider() *catalog.DatabaseProvider { return sess.db } +func (sess *Session) rememberBinding(conn *stdsql.Conn, tx *stdsql.Tx) { + if sess == nil || (conn == nil && tx == nil) { + return + } + sess.bindingMu.Lock() + sess.boundConn = conn + sess.boundTx = tx + sess.bindingSeen = true + sess.bindingMu.Unlock() +} + +func (sess *Session) trackedBinding() (*stdsql.Conn, *stdsql.Tx, bool) { + if sess == nil { + return nil, nil, false + } + sess.bindingMu.Lock() + conn, tx, seen := sess.boundConn, sess.boundTx, sess.bindingSeen + sess.bindingMu.Unlock() + return conn, tx, seen +} + +func (sess *Session) clearTrackedBinding(conn *stdsql.Conn, tx *stdsql.Tx) { + if sess == nil { + return + } + sess.bindingMu.Lock() + // A nil value is an actual part of the binding, not a wildcard. Callers + // that only want to clear a transaction use clearTrackedTransaction below. + // Treating nil as a wildcard let an idle snapshot erase a newer transaction + // tracked on the same physical connection. + if sess.bindingSeen && sess.boundConn == conn && sess.boundTx == tx { + sess.boundConn = nil + sess.boundTx = nil + sess.bindingSeen = false + } + sess.bindingMu.Unlock() +} + +// clearTrackedTransaction retires only the transaction half of a remembered +// binding. The physical connection remains tracked until its own identity-safe +// close; this prevents a later stale callback from falling back to an ID-only +// pool lookup. +func (sess *Session) clearTrackedTransaction(tx *stdsql.Tx) { + if sess == nil || tx == nil { + return + } + sess.bindingMu.Lock() + if sess.bindingSeen && sess.boundTx == tx { + sess.boundTx = nil + if sess.boundConn == nil { + sess.bindingSeen = false + } + } + sess.bindingMu.Unlock() +} + +// clearTrackedTransactionIfChanged retires a remembered transaction only after +// the pool proves that the exact transaction is no longer active. In +// particular, a malformed owner mapping must keep the identity available for a +// later recovery/teardown path. +func (sess *Session) clearTrackedTransactionIfChanged(tx *stdsql.Tx) { + if sess == nil || tx == nil || sess.db == nil || sess.db.Pool() == nil { + return + } + _, trackedTx, seen := sess.trackedBinding() + if !seen || trackedTx != tx { + return + } + _, current, active, err := sess.db.Pool().GetTxnBindingForReleaseStatus(sess.ID()) + if err != nil || (active && current == tx) { + return + } + sess.clearTrackedTransaction(tx) +} + +// clearTrackedBindingIfChanged drops an old remembered pair only after the +// pool confirms that the pair is no longer current. A malformed active mapping +// is deliberately retained so a later recovery path still has its identity. +func (sess *Session) clearTrackedBindingIfChanged(conn *stdsql.Conn, tx *stdsql.Tx) { + if sess == nil || sess.db == nil || sess.db.Pool() == nil { + return + } + trackedConn, trackedTx, seen := sess.trackedBinding() + if !seen || trackedConn != conn || trackedTx != tx { + return + } + currentConn, current, _, err := sess.db.Pool().GetTxnBindingForReleaseStatus(sess.ID()) + if err != nil { + return + } + if currentConn == conn && current == tx { + return + } + sess.clearTrackedBinding(conn, tx) +} + +// RegisterImplicitTransactionCleanup attaches a resource-release callback to +// the currently tracked GMS implicit transaction. The returned handle is a +// success-path disarmer: false is intentionally a no-op (the callback carries +// terminal failure state), while true disarms the marker after a successful +// commit. Resource callbacks are run by CommitTransaction or the rollback +// path, never by the success handle itself. +// +// The callback is intentionally allowed to return an error, but the handle has +// no error result because it is normally called from iterator/finalizer paths +// that cannot return one. Such errors are retained by the session and exposed +// through TakeAutocommitCleanupError. SetTransaction(nil) still performs the +// rollback when a callback fails. +func (sess *Session) RegisterImplicitTransactionCleanup( + ctx *sql.Context, + tx sql.Transaction, + callback func() error, +) func(success bool) { + noOp := func(bool) {} + if sess == nil || callback == nil { + return noOp + } + transaction, ok := tx.(*Transaction) + if !ok || transaction == nil { + return noOp + } + + // Check the marker and append the callback while holding implicitMu. The + // previous split check/add sequence allowed SetTransaction(nil) or a newer + // armImplicitState call to detach the marker between those operations, + // leaving the callback on an unreachable state and leaking its leases. + sess.implicitMu.Lock() + state := sess.implicit + if state == nil || state.tx != transaction || !state.add(callback) { + sess.implicitMu.Unlock() + // The marker may have been detached/finalized concurrently. The callback + // still owns a resource, so run it immediately instead of silently + // leaking that resource. Preserve callback failures for the protocol + // layer, which cannot receive an error from this registration API. + if callbackErr := callback(); callbackErr != nil { + sess.recordAutocommitCleanupError(callbackErr) + } + return noOp + } + sess.implicitMu.Unlock() + + return func(success bool) { + if !success { + // Failure is reported by the registered callback. The wrapper has + // already closed its child and released its leases before reaching + // this handle, so there is no work to do on the false path. + return + } + // The success handle is called only after Session.CommitTransaction has + // completed. Do not run callbacks here: doing so before the physical + // commit would make a later commit failure look like a clean statement. + sess.removeImplicitState(state) + } +} + +// RegisterPostTransactionCleanup retains a callback until the supplied logical +// transaction has completed. Unlike the pre-finalization implicit hook, this +// callback is deliberately not run by iterator close: an operation lease must +// remain admitted while commit/rollback is still deciding the transaction's +// outcome. If registration races a completed transaction, the callback is +// invoked immediately with the recorded outcome. +func (sess *Session) RegisterPostTransactionCleanup( + _ *sql.Context, + tx sql.Transaction, + callback func(success bool), +) bool { + if sess == nil || callback == nil { + return false + } + transaction, ok := tx.(*Transaction) + if !ok || transaction == nil { + return false + } + return transaction.registerPostCleanup(callback) +} + +// RegisterPostPhysicalTransactionCleanup retains a callback for a raw +// database/sql transaction finalized directly by a protocol adapter. This is +// used by PostgreSQL's extended/COPY paths, whose transaction controls bypass +// GMS's logical TransactionSession methods. +func (sess *Session) RegisterPostPhysicalTransactionCleanup( + _ *sql.Context, + tx *stdsql.Tx, + callback func(success bool), +) bool { + if sess == nil || tx == nil || callback == nil { + return false + } + sess.postPhysicalMu.Lock() + sess.prunePostPhysicalCompletionsLocked(time.Now()) + if sess.postPhysical == nil { + sess.postPhysical = make(map[*stdsql.Tx]*postTransactionCleanup) + } + state := sess.postPhysical[tx] + statePresent := state != nil + if state == nil { + state = &postTransactionCleanup{} + sess.postPhysical[tx] = state + } + state.mu.Lock() + completed := state.completed + bridgeNeeded := !state.poolBridgeRegistered && !completed + if bridgeNeeded { + state.poolBridgeRegistered = true + } + state.mu.Unlock() + sess.postPhysicalMu.Unlock() + pool := (*catalog.ConnectionPool)(nil) + if sess.db != nil { + pool = sess.db.Pool() + } + if bridgeNeeded && pool == nil { + // Without a pool (or another physical finalizer) there is no trustworthy + // proof that a newly-created raw transaction is still active. Fail closed + // instead of installing a callback that could wait forever. Existing + // completed tombstones take the path above and still support late delivery. + if !statePresent { + sess.postPhysicalMu.Lock() + if sess.postPhysical[tx] == state { + delete(sess.postPhysical, tx) + } + sess.postPhysicalMu.Unlock() + } + state.complete(false) + return false + } + // The pool may finalize this raw transaction directly during provider + // Close/Reset. Bridge that terminal event into the session registry as well + // as the normal Session adapter finalizers; the local state remains the + // single once-guard for user callbacks. + if bridgeNeeded { + if pool.RegisterTransactionCompletion(tx, func(success bool) { + sess.completePhysicalPostCleanup(tx, success) + }) { + // The pool has an active binding or a retained recent tombstone. + // Continue with local registration below. + } else { + // An evicted/inactive physical transaction cannot safely accept a new + // pending callback. Resolve callbacks already retained for this state and + // return false so COPY/operation callers take their fallback release path. + sess.completePhysicalPostCleanup(tx, false) + return false + } + } + return state.register(callback) +} + +func (transaction *Transaction) registerPostCleanup(callback func(bool)) bool { + if transaction == nil || callback == nil { + return false + } + transaction.postMu.Lock() + if transaction.post == nil { + transaction.post = &postTransactionCleanup{} + } + state := transaction.post + transaction.postMu.Unlock() + return state.register(callback) +} + +func (transaction *Transaction) completePostCleanup(success bool) { + if transaction == nil { + return + } + transaction.postMu.Lock() + state := transaction.post + if state == nil { + // Keep a completion tombstone so a callback registered after the + // finalizer still observes the terminal outcome instead of being queued + // on an unreachable transaction. + state = completedPostTransactionCleanup(success) + transaction.post = state + } + transaction.postMu.Unlock() + state.complete(success) +} + +func (sess *Session) completePhysicalPostCleanup(tx *stdsql.Tx, success bool) { + if sess == nil || tx == nil { + return + } + sess.postPhysicalMu.Lock() + sess.prunePostPhysicalCompletionsLocked(time.Now()) + if sess.postPhysical == nil { + sess.postPhysical = make(map[*stdsql.Tx]*postTransactionCleanup) + } + state := sess.postPhysical[tx] + if state == nil { + // Pool/connection teardown can complete a raw transaction before the + // protocol callback has registered. Retain the outcome under the same + // registry lock so late registration is immediately resolved. + state = completedPostTransactionCleanup(success) + sess.postPhysical[tx] = state + } + sess.postPhysicalMu.Unlock() + panicValue, panicked := state.completeCollectPanic(success) + sess.recordPostPhysicalCompletion(tx, state) + if panicked { + panic(panicValue) + } +} + +// completePhysicalPostCleanupIfUnbound resolves a callback after an +// identity-safe teardown call. A close operation may remove the transaction +// mapping and still return a physical close error; that is nevertheless a +// terminal lifecycle event for the callback. Conversely, retain the callback +// when the exact transaction is still mapped, because its driver state remains +// live/unknown and must not be released early. +func (sess *Session) completePhysicalPostCleanupIfUnbound(tx *stdsql.Tx) { + if sess == nil || tx == nil || sess.db == nil { + return + } + _, current, active, err := sess.db.Pool().GetTxnBindingForReleaseStatus(sess.ID()) + if err != nil || (active && current == tx) { + // A malformed owner/transaction mapping is still active even though the + // pool cannot return a typed pointer. Keep the callback armed until a real + // teardown path resolves that exact physical state. + return + } + sess.completePhysicalPostCleanup(tx, false) +} + +// TakeAutocommitCleanupError returns and clears errors retained by an implicit +// rollback/lease cleanup path. SetTransaction cannot return an error, so GMS +// or a protocol adapter may consume this value after the original statement +// error has been recorded. +func (sess *Session) TakeAutocommitCleanupError() error { + if sess == nil { + return nil + } + sess.implicitMu.Lock() + defer sess.implicitMu.Unlock() + err := sess.cleanupErr + sess.cleanupErr = nil + return err +} + +// AutocommitCleanupError observes the retained cleanup error without clearing +// it. It is useful to include the error in a surrounding protocol response +// before a later owner consumes it. +func (sess *Session) AutocommitCleanupError() error { + if sess == nil { + return nil + } + sess.implicitMu.Lock() + defer sess.implicitMu.Unlock() + return sess.cleanupErr +} + +// ConsumeAutocommitCleanupError is an explicit alias for callers that prefer +// the verb used by other session error queues. +func (sess *Session) ConsumeAutocommitCleanupError() error { + return sess.TakeAutocommitCleanupError() +} + +func (sess *Session) recordAutocommitCleanupError(err error) { + if sess == nil || err == nil { + return + } + sess.implicitMu.Lock() + sess.cleanupErr = errors.Join(sess.cleanupErr, err) + sess.implicitMu.Unlock() + if sess.Session != nil { + sess.GetLogger().WithError(err).Error("implicit autocommit cleanup failed") + } +} + +// beginProcedureTransactionScope brackets the synchronous GMS buildCall +// transaction swap. GMS clears the session transaction before running a stored +// procedure and restores it in a defer; SetTransaction has no context +// argument, so the executor supplies this small, session-local scope marker. +// The returned release is idempotent because a builder may abandon a plan +// after constructing an iterator. A missing restore is treated as a real clear +// at scope end, so a failed procedure cannot strand a marker or its leases. +func (sess *Session) beginProcedureTransactionScope() func() { + if sess == nil { + return func() {} + } + scope := &procedureTransactionScope{} + sess.implicitMu.Lock() + sess.procedureScopes = append(sess.procedureScopes, scope) + sess.implicitMu.Unlock() + + var once sync.Once + return func() { + once.Do(func() { + sess.implicitMu.Lock() + idx := -1 + for i := len(sess.procedureScopes) - 1; i >= 0; i-- { + if sess.procedureScopes[i] == scope { + idx = i + break + } + } + if idx < 0 { + sess.implicitMu.Unlock() + return + } + // Scopes are expected to unwind LIFO. If an unusual builder path + // releases an outer scope first, leave inner scopes intact and only + // remove this exact record. + state := scope.suspended + scope.suspended = nil + copy(sess.procedureScopes[idx:], sess.procedureScopes[idx+1:]) + sess.procedureScopes[len(sess.procedureScopes)-1] = nil + sess.procedureScopes = sess.procedureScopes[:len(sess.procedureScopes)-1] + if state != nil { + state.suspended = false + } + sess.implicitMu.Unlock() + if state != nil { + sess.resolveImplicitState(state) + } + }) + } +} + +func (sess *Session) removeImplicitState(state *implicitTransactionCleanup) { + if sess == nil || state == nil { + return + } + sess.implicitMu.Lock() + if sess.implicit == state { + sess.implicit = nil + state.suspended = false + } + for _, scope := range sess.procedureScopes { + if scope == nil || scope.suspended != state { + continue + } + scope.suspended = nil + state.suspended = false + } + sess.implicitMu.Unlock() +} + +func (sess *Session) armImplicitState(ctx *sql.Context, tx *Transaction) { + if sess == nil || tx == nil || ctx == nil { + return + } + state := &implicitTransactionCleanup{tx: tx, ctx: ctx} + sess.implicitMu.Lock() + old := sess.implicit + state.procedureOwned = len(sess.procedureScopes) > 0 + sess.implicit = state + sess.implicitMu.Unlock() + if old != nil && old != state { + // A new transaction should never inherit an old iterator lease. Do not + // rollback the stale logical transaction here (it may be a temporary GMS + // transaction swap), but release resources that are no longer usable. + if err := old.run(); err != nil { + sess.recordAutocommitCleanupError(err) + } + } +} + +func (sess *Session) takeImplicitForCurrent() (*implicitTransactionCleanup, sql.Transaction, bool) { + if sess == nil { + return nil, nil, false + } + sess.implicitMu.Lock() + defer sess.implicitMu.Unlock() + state := sess.implicit + if state == nil { + return nil, nil, false + } + current := sql.Transaction(nil) + if sess.Session != nil { + current = sess.Session.GetTransaction() + } + matched := current == state.tx + // A mismatch can be a stored-procedure/event temporary swap or a newer + // transaction that replaced the tracked one. Leave the marker intact and + // never clear the replacement from this stale nil setter. + if matched { + sess.implicit = nil + } + return state, current, matched +} + +// suspendImplicitForProcedure detaches the current implicit marker while a +// GMS stored procedure owns the session transaction slot. The marker remains +// available for the deferred SetTransaction(oldTx) restore. +func (sess *Session) suspendImplicitForProcedure() bool { + if sess == nil || sess.Session == nil { + return false + } + sess.implicitMu.Lock() + state := sess.implicit + current := sess.Session.GetTransaction() + if len(sess.procedureScopes) == 0 || state == nil || current != state.tx { + sess.implicitMu.Unlock() + return false + } + scope := sess.procedureScopes[len(sess.procedureScopes)-1] + if scope == nil || scope.suspended != nil { + // This scope already detached its outer transaction. Any nil clear now + // belongs to a child statement and must retain the ordinary rollback + // semantics. + sess.implicitMu.Unlock() + return false + } + sess.implicit = nil + state.suspended = true + scope.suspended = state + sess.implicitMu.Unlock() + + // Keep the base session's transaction slot in the state GMS expects while + // the procedure's first child statement starts its own transaction. + sess.Session.SetTransaction(nil) + return true +} + +// resolveImplicitState performs the real clear for a marker that was +// speculatively detached. It restores the exact logical transaction while the +// callback runs, then uses the normal rollback path, and finally clears only +// that identity from the base session. +func (sess *Session) resolveImplicitState(state *implicitTransactionCleanup) { + if sess == nil || state == nil || sess.Session == nil { + return + } + sess.Session.SetTransaction(state.tx) + cleanupErr := state.run() + rollbackErr := sess.Rollback(state.ctx, state.tx) + if err := errors.Join(cleanupErr, rollbackErr); err != nil { + sess.recordAutocommitCleanupError(err) + } + current := sess.Session.GetTransaction() + if current == state.tx || current == nil { + sess.Session.SetTransaction(nil) + } +} + +// restoreSuspendedImplicit recognizes the exact transaction object captured +// by GMS before its temporary nil assignment. Any newer detached markers are +// resolved first; they belong to an inner procedure that unwound without +// restoring its own transaction. +func (sess *Session) restoreSuspendedImplicit(tx *Transaction) bool { + if sess == nil || tx == nil || sess.Session == nil { + return false + } + sess.implicitMu.Lock() + if len(sess.procedureScopes) == 0 { + sess.implicitMu.Unlock() + return false + } + scope := sess.procedureScopes[len(sess.procedureScopes)-1] + if scope == nil || scope.suspended == nil || scope.suspended.tx != tx { + sess.implicitMu.Unlock() + return false + } + current := sess.Session.GetTransaction() + if current != nil && current != tx { + // A procedure child may have left a transaction active. Resolve it only + // when its marker was created inside a procedure scope; an unrelated + // replacement must remain untouched. + child := sess.implicit + if child == nil || child.tx != current || !child.procedureOwned { + sess.implicitMu.Unlock() + return false + } + sess.implicitMu.Unlock() + sess.resolveImplicitState(child) + sess.implicitMu.Lock() + if len(sess.procedureScopes) == 0 || sess.procedureScopes[len(sess.procedureScopes)-1] != scope || + scope.suspended == nil || scope.suspended.tx != tx { + sess.implicitMu.Unlock() + return false + } + current = sess.Session.GetTransaction() + if current != nil && current != tx { + sess.implicitMu.Unlock() + return false + } + } + target := scope.suspended + scope.suspended = nil + target.suspended = false + sess.implicit = target + sess.implicitMu.Unlock() + sess.Session.SetTransaction(tx) + return true +} + +func (sess *Session) peekImplicit(tx *Transaction) *implicitTransactionCleanup { + if sess == nil || tx == nil { + return nil + } + sess.implicitMu.Lock() + defer sess.implicitMu.Unlock() + if sess.implicit == nil || sess.implicit.tx != tx { + return nil + } + return sess.implicit +} + +// HasImplicitTransactionCleanup reports whether tx is currently backed by the +// frontend implicit-autocommit marker. Iterator wrappers use this distinction +// before registering a pre-finalization callback: explicit transactions and +// MySQL's DDL transaction wrapper intentionally have no such marker, and a +// failed registration must not close their live child iterator immediately. +func (sess *Session) HasImplicitTransactionCleanup(tx sql.Transaction) bool { + transaction, ok := tx.(*Transaction) + if !ok || sess == nil || sess.Session == nil { + return false + } + sess.implicitMu.Lock() + defer sess.implicitMu.Unlock() + return sess.implicit != nil && sess.implicit.tx == transaction && sess.Session.GetTransaction() == transaction +} + +func (sess *Session) disarmImplicitTransaction(tx *Transaction, runCleanup bool) error { + state := sess.peekImplicit(tx) + if state == nil { + return nil + } + sess.removeImplicitState(state) + if !runCleanup { + return nil + } + err := state.run() + if err != nil { + sess.recordAutocommitCleanupError(err) + } + return err +} + +// RegisterImplicitTransactionCleanupForContext lets wrappers avoid a direct +// type assertion while keeping the concrete implementation in this package. +func RegisterImplicitTransactionCleanupForContext( + ctx *sql.Context, + tx sql.Transaction, + callback func() error, +) func(success bool) { + if ctx == nil || ctx.Session == nil { + return func(bool) {} + } + if registrar, ok := ctx.Session.(ImplicitTransactionCleanupRegistrar); ok { + return registrar.RegisterImplicitTransactionCleanup(ctx, tx, callback) + } + return func(bool) {} +} + +// BeginDuckLakeOperation reserves provider-wide admission for a top-level +// frontend statement. Keeping this small forwarding method on the session +// lets the execution builder use the session lifecycle without reaching into +// provider internals or changing the existing ConnectionHolder contract. +func (sess *Session) BeginDuckLakeOperation(ctx context.Context) func() { + if sess == nil || sess.db == nil { + return func() {} + } + return sess.db.BeginDuckLakeOperation(ctx) +} + +// BeginDuckLakeOperationWithError is the production admission surface used by +// execution builders that can return an error. Keeping the error-bearing form +// here prevents a poisoned provider generation from being represented as the +// legacy no-op release closure. +func (sess *Session) BeginDuckLakeOperationWithError(ctx context.Context) (func(), error) { + if sess == nil || sess.db == nil { + return func() {}, nil + } + return sess.db.BeginDuckLakeOperationWithError(ctx) +} + // QueryRowLimit returns the maximum number of rows a query may return. Zero // means unlimited. func (sess *Session) QueryRowLimit() uint64 { @@ -87,11 +1079,13 @@ func (sess *Session) SetSessionVariable(ctx *sql.Context, name string, value int if !ok || transaction.tx != nil { return nil } - if err := sess.Session.Rollback(ctx, &transaction.Transaction); err != nil { - return err - } + // The SET itself is the boundary that changes autocommit semantics. It is + // not a failed user operation, so discard the temporary GMS marker and its + // resource callback without invoking the DuckLake orphan sweep here. + cleanupErr := sess.disarmImplicitTransaction(transaction, true) + rollbackErr := sess.Session.Rollback(ctx, &transaction.Transaction) ctx.SetTransaction(nil) - return nil + return errors.Join(cleanupErr, rollbackErr) } func (sess *Session) CurrentSchemaOfUnderlyingConn() string { @@ -125,11 +1119,113 @@ func NewSessionBuilder(provider *catalog.DatabaseProvider, opts ...SessionOption var _ sql.TransactionSession = (*Session)(nil) var _ sql.PersistableSession = (*Session)(nil) +var _ sql.LifecycleAwareSession = (*Session)(nil) var _ adapter.ConnectionHolder = (*Session)(nil) +var _ adapter.ExecutionSnapshotHolder = (*Session)(nil) +var _ adapter.ExecutionSnapshotLeaseHolder = (*Session)(nil) +var _ adapter.ReleaseTransactionBindingHolder = (*Session)(nil) +var _ adapter.ReleaseTransactionHolder = (*Session)(nil) +var _ adapter.TransactionFinalizer = (*Session)(nil) +var _ adapter.RollbackCleanupFinalizer = (*Session)(nil) +var _ PostTransactionCleanupRegistrar = (*Session)(nil) +var _ PhysicalPostTransactionCleanupRegistrar = (*Session)(nil) type Transaction struct { memory.Transaction - tx *stdsql.Tx + tx *stdsql.Tx + postMu sync.Mutex + post *postTransactionCleanup +} + +// CommandBegin and CommandEnd opt the session into GMS lifecycle callbacks. +// MyDuck does not need per-command bookkeeping, but implementing the optional +// interface is important because SessionEnd is the only callback GMS gives an +// integrator when a MySQL session is reset or disconnected. +func (sess *Session) CommandBegin() error { return nil } + +func (sess *Session) CommandEnd() {} + +// SessionEnd rolls back the exact logical transaction still owned by this +// session. GMS invokes this callback from SessionManager while replacing or +// removing a connection; the request context that started the transaction may +// already be canceled, so use a fresh, non-cancelable frontend context. The +// embedded setter is used only after rollback and only when the identity is +// unchanged, preserving a replacement transaction installed by a concurrent +// lifecycle path. +func (sess *Session) SessionEnd() { + if sess == nil || sess.Session == nil { + return + } + sess.sessionEndMu.Lock() + defer sess.sessionEndMu.Unlock() + + transaction, ok := sess.Session.GetTransaction().(*Transaction) + ctx := sql.NewContext( + mycontext.WithFrontendQuery(context.Background()), + sql.WithSession(sess), + ) + if ok && transaction != nil { + if err := sess.Rollback(ctx, transaction); err != nil { + sess.recordAutocommitCleanupError(err) + } + + // Rollback intentionally does not clear GMS's logical transaction slot: the + // normal rowexec finalizer does that through SetTransaction(nil). Do the same + // here without re-entering the custom setter, and never clear a replacement. + if current := sess.Session.GetTransaction(); current == transaction { + sess.Session.SetTransaction(nil) + } + return + } + + // PostgreSQL transaction controls bypass GMS's TransactionSession slot. A + // reset/disconnect can therefore leave only the raw pool binding behind. Use + // the exact pair captured when this Session acquired it; an ID-only lookup is + // unsafe because SessionManager may already have installed a replacement. + trackedConn, trackedTx, seen := sess.trackedBinding() + if !seen || trackedTx == nil || sess.db == nil || sess.db.Pool() == nil { + return + } + currentConn, currentTx := sess.db.Pool().GetTxnBindingForRelease(sess.ID()) + if currentTx == trackedTx && currentConn == trackedConn { + var cleanup func(*stdsql.Conn) error + var releaseCleanup func() + var reservationErr error + if sess.db.DuckLakeObjectStorageEnabled() { + releaseCleanup, reservationErr = sess.db.BeginDuckLakeRollbackCleanupForOwnerWithError(trackedTx, trackedTx) + if reservationErr == nil { + defer releaseCleanup() + cleanup = func(conn *stdsql.Conn) error { + cleanupBase := context.WithoutCancel(ctx) + cleanupBase = mycontext.WithQueryOrigin(cleanupBase, mycontext.MaintenanceQueryOrigin) + cleanupBase = catalog.WithRollbackCleanupOwner(cleanupBase, conn) + cleanupBase = catalog.WithDuckLakeCleanupLease(cleanupBase) + return sess.db.CleanupDuckLakeOrphansOnConn(ctx.WithContext(cleanupBase), conn) + } + } + } + _, finalizerErr := sess.RollbackTxnWithCleanup(trackedTx, cleanup) + if err := errors.Join(reservationErr, finalizerErr); err != nil { + sess.recordAutocommitCleanupError(err) + } + sess.clearTrackedBindingIfChanged(trackedConn, trackedTx) + return + } + + // The mapping disappeared or was replaced before SessionEnd. Roll back the + // captured raw transaction itself so database/sql can return its owner, but + // never close the connection pointer: it may now be serving the replacement. + if err := trackedTx.Rollback(); err != nil && !adapter.IsTransactionInactiveError(err) { + sess.recordAutocommitCleanupError(err) + } + // The pool normally reports this terminal event when it removes the old + // mapping. A stale SessionEnd can nevertheless arrive after an external + // owner removed the mapping without going through the pool finalizer. The raw + // rollback above is terminal for database/sql, so resolve this exact session + // callback even in that path; the registry is one-shot and tolerates the + // normal pool callback racing with it. + sess.completePhysicalPostCleanup(trackedTx, false) + sess.clearTrackedBindingIfChanged(trackedConn, trackedTx) } var _ sql.Transaction = (*Transaction)(nil) @@ -162,35 +1258,199 @@ func (sess *Session) StartTransaction(ctx *sql.Context, tCharacteristic sql.Tran return nil, err } } - return &Transaction{*base.(*memory.Transaction), tx}, nil + transaction := &Transaction{Transaction: *base.(*memory.Transaction), tx: tx} + // GMS creates a logical transaction for ordinary autocommit statements even + // though no physical DuckDB transaction is opened. Only frontend statements + // receive the compensating rollback hook; explicit controls, DDL wrappers, + // replication, maintenance, and recovery contexts must retain their own + // lifecycle semantics. + if tx == nil && !sess.mysqlImplicitDDL.Load() && !ctx.GetIgnoreAutoCommit() && + mycontext.QueryOrigin(ctx) == mycontext.FrontendQueryOrigin { + sess.armImplicitState(ctx, transaction) + } + return transaction, nil +} + +// SetTransaction overrides the embedded BaseSession setter so GMS's +// clearAutocommitOnError path becomes a real rollback boundary. The setter has +// no error return, therefore callback/rollback failures are retained for the +// protocol layer to consume through TakeAutocommitCleanupError. +func (sess *Session) SetTransaction(tx sql.Transaction) { + if sess == nil || sess.Session == nil { + return + } + if tx != nil { + // GMS restores the transaction captured before a stored-procedure swap + // through this setter. Recognize only the exact suspended identity; all + // other assignments retain ordinary replacement semantics. + if transaction, ok := tx.(*Transaction); ok && sess.restoreSuspendedImplicit(transaction) { + return + } + sess.Session.SetTransaction(tx) + return + } + + // The nil assignment at the start of GMS buildCall is a temporary detach, + // not a statement boundary. The executor brackets that synchronous call so + // this path can suspend the marker without invoking rollback/cleanup. + if sess.suspendImplicitForProcedure() { + return + } + + state, current, matched := sess.takeImplicitForCurrent() + if state == nil { + // There is no implicit marker to resolve, so preserve the ordinary GMS + // setter behavior. + sess.Session.SetTransaction(nil) + return + } + if !matched { + // A different transaction may be a stored-procedure temporary scope or a + // replacement installed by a newer request. This nil setter cannot prove + // that the tracked transaction is the one being cleared, so leave the + // marker and replacement untouched. When the underlying value is already + // nil, forwarding the nil assignment is harmless and keeps BaseSession's + // state canonical. + if current == nil { + sess.Session.SetTransaction(nil) + } + return + } + + // Keep the logical transaction visible to the rollback implementation until + // its provider/session cleanup has completed. This preserves the saved + // frontend context and prevents a replacement from racing the cleanup path. + cleanupErr := state.run() + rollbackErr := sess.Rollback(state.ctx, state.tx) + if err := errors.Join(cleanupErr, rollbackErr); err != nil { + sess.recordAutocommitCleanupError(err) + } + + // A concurrent lifecycle path may have installed a replacement transaction + // while cleanup ran. Clear only the exact transaction captured above; never + // evict a newer binding from this stale setter. + current = sess.Session.GetTransaction() + if current == state.tx || current == nil { + sess.Session.SetTransaction(nil) + } } // CommitTransaction implements sql.TransactionSession. func (sess *Session) CommitTransaction(ctx *sql.Context, tx sql.Transaction) error { sess.GetLogger().Trace("CommitTransaction") - transaction := tx.(*Transaction) + transaction, ok := tx.(*Transaction) + if !ok || transaction == nil { + return fmt.Errorf("cannot commit transaction of type %T", tx) + } + // Release any iterator/operation callback before committing the logical + // transaction. If resource release fails, leave the marker armed so GMS's + // subsequent SetTransaction(nil) can roll the transaction back instead of + // reporting a successful commit with live resources. + state := sess.peekImplicit(transaction) + if state != nil { + if err := state.run(); err != nil { + sess.recordAutocommitCleanupError(err) + return err + } + } if transaction.tx != nil { sess.GetLogger().Trace("CommitDuckTransaction") - defer sess.CloseTxn() - if err := transaction.tx.Commit(); err != nil { + if err := adapter.FinalizeCommit(ctx, transaction.tx); err != nil { return err } } - return sess.Session.CommitTransaction(ctx, &transaction.Transaction) + if err := sess.Session.CommitTransaction(ctx, &transaction.Transaction); err != nil { + // Physical commit may already have completed, but the logical GMS + // transaction did not. Release post callbacks as a failed overall + // finalization rather than leaving operation admission stranded. + transaction.completePostCleanup(false) + return err + } + // Both physical and logical commit have completed. Only now may operation + // leases registered against this logical transaction be released. + transaction.completePostCleanup(true) + // A successful logical commit must not be mistaken for a later failed + // autocommit clear. Remove only the exact marker; a replacement transaction + // installed concurrently remains untouched. + sess.removeImplicitState(sess.peekImplicit(transaction)) + return nil } // Rollback implements sql.TransactionSession. func (sess *Session) Rollback(ctx *sql.Context, tx sql.Transaction) error { sess.GetLogger().Trace("Rollback") transaction := tx.(*Transaction) + var implicitCleanupErr error + if state := sess.peekImplicit(transaction); state != nil { + // Explicit rollback paths (including protocol ROLLBACK) also need to + // release a retained iterator/operation lease before touching provider + // cleanup. The marker is removed first so the later SetTransaction(nil) + // is idempotent and cannot rollback twice. + sess.removeImplicitState(state) + implicitCleanupErr = state.run() + } + var rollbackErr error + underlyingInactive := transaction.tx == nil + var rollbackConn *stdsql.Conn + var cleanupErr error + cleanupEligible := ctx != nil && + mycontext.QueryOrigin(ctx) == mycontext.FrontendQueryOrigin && + sess.db != nil && sess.db.DuckLakeObjectStorageEnabled() + // Reserve the provider cleanup barrier before entering the pool/session + // finalizer. The finalizer holds a release-side lifecycle lock; acquiring + // the barrier from inside its callback can otherwise deadlock against an + // operation that still needs that lock to finish. The reservation also + // covers GMS-only (no physical *sql.Tx) rollback scopes. + var releaseCleanup func() + var reservationErr error + if cleanupEligible { + releaseCleanup, reservationErr = sess.db.BeginDuckLakeRollbackCleanupForOwnerWithError(transaction.tx, transaction) + if reservationErr == nil { + defer releaseCleanup() + } else { + // A poisoned generation must not prevent the physical/logical rollback; + // skip maintenance and return the admission error alongside rollback + // results. + releaseCleanup = nil + } + } + cleanup := func(conn *stdsql.Conn) error { + if !cleanupEligible || reservationErr != nil || conn == nil { + return nil + } + cleanupBase := context.WithoutCancel(ctx) + cleanupBase = mycontext.WithQueryOrigin(cleanupBase, mycontext.MaintenanceQueryOrigin) + cleanupBase = catalog.WithRollbackCleanupOwner(cleanupBase, conn) + cleanupBase = catalog.WithDuckLakeCleanupLease(cleanupBase) + cleanupSQLCtx := ctx.WithContext(cleanupBase) + return sess.db.CleanupDuckLakeOrphansOnConn(cleanupSQLCtx, conn) + } if transaction.tx != nil { sess.GetLogger().Trace("RollbackDuckTransaction") - defer sess.CloseTxn() - if err := transaction.tx.Rollback(); err != nil { - return err + underlyingInactive, rollbackErr = adapter.FinalizeRollbackWithCleanup(ctx, transaction.tx, cleanup) + } + + // For the GMS-only transaction case there is no physical transaction in the + // pool to hold the session lifecycle lock. Select and run post-rollback + // maintenance while the logical transaction is still owned by this session; + // clearing it first would let a replacement transaction or connection-close + // path race the cleanup operation. + if transaction.tx == nil && underlyingInactive && cleanupEligible && reservationErr == nil { + cleanupBase := context.WithoutCancel(ctx) + cleanupBase = mycontext.WithQueryOrigin(cleanupBase, mycontext.MaintenanceQueryOrigin) + cleanupBase = catalog.WithDuckLakeCleanupLease(cleanupBase) + cleanupSQLCtx := ctx.WithContext(cleanupBase) + rollbackConn, cleanupErr = sess.GetConn(cleanupSQLCtx) + if cleanupErr == nil && reservationErr == nil { + cleanupErr = sess.db.CleanupDuckLakeOrphansOnConn(cleanupSQLCtx, rollbackConn) } } - return sess.Session.Rollback(ctx, &transaction.Transaction) + sessionErr := sess.Session.Rollback(ctx, &transaction.Transaction) + // Rollback has reached its terminal point even when one component reports an + // error. A post-finalization callback must not keep the provider barrier held + // forever after the physical owner has been retired or proven inactive. + transaction.completePostCleanup(false) + return errors.Join(implicitCleanupErr, reservationErr, rollbackErr, sessionErr, cleanupErr) } // PersistGlobal implements sql.PersistableSession. @@ -253,29 +1513,117 @@ func (sess *Session) GetPersistedValue(k string) (interface{}, error) { // GetConn implements adapter.ConnectionHolder. func (sess *Session) GetConn(ctx context.Context) (*stdsql.Conn, error) { - return sess.db.Pool().GetConnForSchema(ctx, sess.ID(), sess.GetCurrentDatabase()) + if sess == nil || sess.db == nil || sess.db.Pool() == nil { + return nil, fmt.Errorf("database provider is unavailable") + } + conn, tx, err := sess.db.Pool().GetConnForSchemaWithBinding( + ctx, sess.ID(), sess.GetCurrentDatabase(), + ) + if err == nil && conn != nil { + // The pool selected the connection and transaction under one + // lifecycle/session critical section. Keep that exact pair; a second + // binding lookup here could observe a replacement transaction. + sess.rememberBinding(conn, tx) + } + return conn, err } // GetCatalogConn implements adapter.ConnectionHolder. func (sess *Session) GetCatalogConn(ctx context.Context) (*stdsql.Conn, error) { - return sess.db.Pool().GetConn(ctx, sess.ID()) + if sess == nil || sess.db == nil || sess.db.Pool() == nil { + return nil, fmt.Errorf("database provider is unavailable") + } + conn, tx, err := sess.db.Pool().GetConnWithBinding(ctx, sess.ID()) + if err == nil && conn != nil { + sess.rememberBinding(conn, tx) + } + return conn, err } // GetTxn implements adapter.ConnectionHolder. func (sess *Session) GetTxn(ctx context.Context, options *stdsql.TxOptions) (*stdsql.Tx, error) { - return sess.db.Pool().GetTxn(ctx, sess.ID(), sess.GetCurrentDatabase(), options) + if sess == nil || sess.db == nil || sess.db.Pool() == nil { + return nil, fmt.Errorf("database provider is unavailable") + } + conn, tx, err := sess.db.Pool().GetTxnWithBinding( + ctx, sess.ID(), sess.GetCurrentDatabase(), options, + ) + if err == nil && conn != nil && tx != nil { + sess.rememberBinding(conn, tx) + } + return tx, err } // GetCatalogTxn implements adapter.ConnectionHolder. func (sess *Session) GetCatalogTxn(ctx context.Context, options *stdsql.TxOptions) (*stdsql.Tx, error) { - return sess.db.Pool().GetTxn(ctx, sess.ID(), "", options) + if sess == nil || sess.db == nil || sess.db.Pool() == nil { + return nil, fmt.Errorf("database provider is unavailable") + } + conn, tx, err := sess.db.Pool().GetTxnWithBinding(ctx, sess.ID(), "", options) + if err == nil && conn != nil && tx != nil { + sess.rememberBinding(conn, tx) + } + return tx, err } // TryGetTxn implements adapter.ConnectionHolder. func (sess *Session) TryGetTxn() *stdsql.Tx { + if sess == nil || sess.db == nil || sess.db.Pool() == nil { + return nil + } return sess.db.Pool().TryGetTxn(sess.ID()) } +// GetTxnBinding implements adapter.TransactionBindingHolder. The pool returns +// the physical connection and transaction from one lifecycle snapshot. +func (sess *Session) GetTxnBinding() (*stdsql.Conn, *stdsql.Tx) { + if sess == nil || sess.db == nil || sess.db.Pool() == nil { + return nil, nil + } + return sess.db.Pool().GetTxnBinding(sess.ID()) +} + +// GetTxnBindingForRelease implements adapter.ReleaseTransactionBindingHolder. +// Close/finalizer paths may still capture this snapshot after pool shutdown +// has been announced; the pool admits it through the release lifecycle gate. +func (sess *Session) GetTxnBindingForRelease() (*stdsql.Conn, *stdsql.Tx) { + if sess == nil || sess.db == nil || sess.db.Pool() == nil { + return nil, nil + } + return sess.db.Pool().GetTxnBindingForRelease(sess.ID()) +} + +// TryGetTxnForRelease implements adapter.ReleaseTransactionHolder. +func (sess *Session) TryGetTxnForRelease() *stdsql.Tx { + if sess == nil || sess.db == nil || sess.db.Pool() == nil { + return nil + } + return sess.db.Pool().TryGetTxnForRelease(sess.ID()) +} + +// GetExecutionSnapshot implements adapter.ExecutionSnapshotHolder. The pool +// selects the physical owner and transaction under one lifecycle lock; catalog +// snapshots deliberately avoid changing the current schema. +func (sess *Session) GetExecutionSnapshot(ctx context.Context, catalogOnly bool) (*stdsql.Conn, *stdsql.Tx, error) { + if sess == nil || sess.db == nil || sess.db.Pool() == nil { + return nil, nil, fmt.Errorf("database provider is unavailable") + } + return sess.db.Pool().GetExecutionSnapshot(ctx, sess.ID(), sess.GetCurrentDatabase(), catalogOnly) +} + +// GetExecutionSnapshotLease implements adapter.ExecutionSnapshotLeaseHolder. +// The returned release callback keeps the selected pool generation alive until +// the caller has finished consuming its iterator. +func (sess *Session) GetExecutionSnapshotLease( + ctx context.Context, + catalogOnly bool, +) (*stdsql.Conn, *stdsql.Tx, func(), error) { + if sess == nil || sess.db == nil || sess.db.Pool() == nil { + return nil, nil, func() {}, fmt.Errorf("database provider is unavailable") + } + return sess.db.Pool().GetExecutionSnapshotLease(ctx, sess.ID(), sess.GetCurrentDatabase(), catalogOnly) +} + // GetCurrentCatalog implements adapter.ConnectionHolder. func (sess *Session) GetCurrentCatalog() string { return sess.db.Pool().CurrentCatalog(sess.ID()) @@ -288,26 +1636,156 @@ func (sess *Session) GetCurrentSchema() string { // CloseTxn implements adapter.ConnectionHolder. func (sess *Session) CloseTxn() { - sess.db.Pool().CloseTxn(sess.ID()) + if sess == nil || sess.db == nil || sess.db.Pool() == nil { + return + } + // Prefer the exact transaction this Session observed. An unconditional + // session-ID close can retire a replacement transaction after a reset has + // reused the numeric ID. Only sessions that never observed a binding use the + // no identity means this Session cannot prove ownership, so it must not + // touch a transaction that may belong to a replacement session. + _, trackedTx, seen := sess.trackedBinding() + if seen { + if trackedTx == nil { + return + } + sess.db.Pool().CloseTxnIf(sess.ID(), trackedTx) + sess.completePhysicalPostCleanupIfUnbound(trackedTx) + sess.clearTrackedTransactionIfChanged(trackedTx) + return + } + return +} + +// CloseTxnIf implements adapter.IdentityTransactionCloser. +func (sess *Session) CloseTxnIf(tx *stdsql.Tx) { + if sess == nil || sess.db == nil || sess.db.Pool() == nil || tx == nil { + return + } + sess.db.Pool().CloseTxnIf(sess.ID(), tx) + // Identity-safe stale cleanup can be the only finalization path for a raw + // physical transaction (for example, when PostgreSQL discovers an ownerless + // or replaced binding). Complete any callbacks retained for that exact + // transaction so deferred COPY/operation leases cannot remain stranded. + sess.completePhysicalPostCleanupIfUnbound(tx) + sess.clearTrackedTransactionIfChanged(tx) +} + +// CommitTxn implements adapter.TransactionFinalizer. The pool performs the +// driver operation and identity check while holding its lifecycle lock, so a +// delayed GMS callback cannot finalize or evict a replacement transaction. +func (sess *Session) CommitTxn(tx *stdsql.Tx) error { + matched, err := sess.db.Pool().CommitTxnWithOutcome(sess.ID(), tx) + if !matched { + if err != nil { + // A malformed/missing physical owner leaves the transaction mapping + // active and therefore is not a terminal callback boundary. Keep the + // post-cleanup registration armed until a real finalizer or teardown + // resolves that mapping. + return err + } + // The expected transaction was already finalized or replaced. Surface + // that stale identity instead of reporting a successful commit, and + // complete only this old transaction's callbacks with failure. + sess.completePhysicalPostCleanupIfUnbound(tx) + sess.clearTrackedTransactionIfChanged(tx) + return adapter.ErrTransactionBindingChanged + } + sess.completePhysicalPostCleanup(tx, err == nil) + sess.clearTrackedTransactionIfChanged(tx) + return err +} + +// RollbackTxn implements adapter.TransactionFinalizer. The returned connection +// is the exact physical owner of tx and may be reused for post-rollback +// maintenance only when inactive is true. +func (sess *Session) RollbackTxn(tx *stdsql.Tx) (*stdsql.Conn, bool, error) { + conn, inactive, err := sess.db.Pool().RollbackTxn(sess.ID(), tx) + // A binding error leaves the exact transaction mapped and potentially live; + // only complete the callback after the pool proves that mapping is gone (or a + // replacement is current). This keeps an operation lease from being released + // while a malformed owner is still active. + sess.completePhysicalPostCleanupIfUnbound(tx) + sess.clearTrackedTransactionIfChanged(tx) + return conn, inactive, err +} + +// RollbackTxnWithCleanup keeps the session lifecycle lock held while the +// caller performs same-connection post-rollback maintenance. +func (sess *Session) RollbackTxnWithCleanup(tx *stdsql.Tx, cleanup func(*stdsql.Conn) error) (bool, error) { + inactive, err := sess.db.Pool().RollbackTxnWithCleanup(sess.ID(), tx, cleanup) + sess.completePhysicalPostCleanupIfUnbound(tx) + sess.clearTrackedTransactionIfChanged(tx) + return inactive, err } // CloseConn implements adapter.ConnectionHolder. func (sess *Session) CloseConn() { - sess.db.Pool().CloseConn(sess.ID()) + if sess == nil || sess.db == nil || sess.db.Pool() == nil { + return + } + // Prefer the identity this Session acquired. If no acquisition was observed, + // take one release-side snapshot and close only that exact connection; never + // issue an unconditional ID-based close that could retire a replacement. + conn, tx, seen := sess.trackedBinding() + if !seen || conn == nil { + return + } + _ = sess.db.Pool().CloseConnIfBinding(sess.ID(), conn, tx) + sess.clearTrackedBindingIfChanged(conn, tx) +} + +// CloseConnIf implements adapter.IdentityConnectionCloser. +func (sess *Session) CloseConnIf(conn *stdsql.Conn) error { + if sess == nil || sess.db == nil || sess.db.Pool() == nil { + return nil + } + currentConn, tx := sess.db.Pool().GetTxnBindingForRelease(sess.ID()) + if currentConn != conn { + return nil + } + err := sess.db.Pool().CloseConnIfBinding(sess.ID(), conn, tx) + sess.completePhysicalPostCleanupIfUnbound(tx) + sess.clearTrackedBindingIfChanged(conn, tx) + return err +} + +// CloseConnIfBinding implements adapter.IdentityConnectionBindingCloser. +// Pool teardown verifies the exact physical connection and transaction pair +// while holding the session lifecycle lock. +func (sess *Session) CloseConnIfBinding(conn *stdsql.Conn, tx *stdsql.Tx) error { + if sess == nil || sess.db == nil || sess.db.Pool() == nil { + return nil + } + err := sess.db.Pool().CloseConnIfBinding(sess.ID(), conn, tx) + sess.completePhysicalPostCleanupIfUnbound(tx) + sess.clearTrackedBindingIfChanged(conn, tx) + return err } func (sess *Session) ExecContext(ctx context.Context, query string, args ...any) (stdsql.Result, error) { - conn, err := sess.GetCatalogConn(ctx) + // Persistable-session and other database/sql callers may supply either a GMS + // context or a plain context. In both cases retain the selected pool generation + // through the driver call; a transaction pointer observed before Restart is an + // identity, not permission to execute after teardown begins. + conn, tx, release, err := sess.GetExecutionSnapshotLease(ctx, true) if err != nil { return nil, err } + defer release() + if tx != nil { + return tx.ExecContext(ctx, query, args...) + } return conn.ExecContext(ctx, query, args...) } -func (sess *Session) QueryRow(ctx context.Context, query string, args ...any) *stdsql.Row { - conn, err := sess.GetCatalogConn(ctx) +func (sess *Session) QueryRow(ctx context.Context, query string, args ...any) *adapter.LeasedRow { + conn, tx, release, err := sess.GetExecutionSnapshotLease(ctx, true) if err != nil { - return nil + return adapter.NewLeasedRowError(err) + } + if tx != nil { + return adapter.NewLeasedRow(tx.QueryRowContext(ctx, query, args...), release) } - return conn.QueryRowContext(ctx, query, args...) + return adapter.NewLeasedRow(conn.QueryRowContext(ctx, query, args...), release) } diff --git a/backend/session_implicit_test.go b/backend/session_implicit_test.go new file mode 100644 index 00000000..4b85e028 --- /dev/null +++ b/backend/session_implicit_test.go @@ -0,0 +1,353 @@ +package backend + +import ( + "context" + "errors" + "runtime" + "sync/atomic" + "testing" + "time" + + "github.com/apecloud/myduckserver/mycontext" + "github.com/dolthub/go-mysql-server/memory" + "github.com/dolthub/go-mysql-server/sql" + "github.com/stretchr/testify/require" +) + +func newFrontendImplicitSession(t *testing.T) (*Session, *sql.Context, *Transaction) { + t.Helper() + sess := NewSession(memory.NewSession(sql.NewBaseSession(), nil), nil) + ctx := sql.NewContext( + mycontext.WithFrontendQuery(context.Background()), + sql.WithSession(sess), + ) + tx, err := sess.StartTransaction(ctx, sql.ReadWrite) + require.NoError(t, err) + transaction, ok := tx.(*Transaction) + require.True(t, ok) + ctx.SetTransaction(transaction) + require.Same(t, transaction, sess.Session.GetTransaction()) + require.NotNil(t, sess.implicit) + return sess, ctx, transaction +} + +func TestImplicitCleanupSetTransactionRollsBackAfterCallback(t *testing.T) { + sess, ctx, transaction := newFrontendImplicitSession(t) + callbackCalls := 0 + var callbackTx sql.Transaction + sess.RegisterImplicitTransactionCleanup(ctx, transaction, func() error { + callbackCalls++ + callbackTx = sess.Session.GetTransaction() + return nil + }) + + ctx.SetTransaction(nil) + + require.Equal(t, 1, callbackCalls) + require.Same(t, transaction, callbackTx, "callback must run while the logical transaction is current") + require.Nil(t, sess.Session.GetTransaction()) + require.Nil(t, sess.implicit) + ctx.SetTransaction(nil) + require.Equal(t, 1, callbackCalls, "cleanup must be exactly once") +} + +func TestImplicitCleanupRetainsCallbackErrorAndStillRollsBack(t *testing.T) { + sess, ctx, transaction := newFrontendImplicitSession(t) + callbackErr := errors.New("iterator close failed") + callbackCalls := 0 + sess.RegisterImplicitTransactionCleanup(ctx, transaction, func() error { + callbackCalls++ + return callbackErr + }) + + ctx.SetTransaction(nil) + + require.Equal(t, 1, callbackCalls) + require.ErrorIs(t, sess.AutocommitCleanupError(), callbackErr) + require.ErrorIs(t, sess.TakeAutocommitCleanupError(), callbackErr) + require.Nil(t, sess.Session.GetTransaction()) + require.Nil(t, sess.AutocommitCleanupError()) +} + +func TestImplicitCleanupCommitRunsCallbackBeforeCommitAndDisarms(t *testing.T) { + sess, ctx, transaction := newFrontendImplicitSession(t) + callbackCalls := 0 + var callbackTx sql.Transaction + derived := ctx.WithContext(context.WithValue(ctx.Context, struct{}{}, "derived")) + sess.RegisterImplicitTransactionCleanup(derived, transaction, func() error { + callbackCalls++ + callbackTx = sess.Session.GetTransaction() + return nil + }) + + require.NoError(t, sess.CommitTransaction(ctx, transaction)) + require.Equal(t, 1, callbackCalls) + require.Same(t, transaction, callbackTx) + require.Nil(t, sess.implicit) + ctx.SetTransaction(nil) + require.Equal(t, 1, callbackCalls) +} + +func TestImplicitCleanupFailedHandleRejectsCommitUntilRollback(t *testing.T) { + sess, ctx, transaction := newFrontendImplicitSession(t) + callbackCalls := 0 + callbackErr := errors.New("terminal iterator failure") + handle := sess.RegisterImplicitTransactionCleanup(ctx, transaction, func() error { + callbackCalls++ + return callbackErr + }) + handle(false) + handle(false) + + err := sess.CommitTransaction(ctx, transaction) + require.ErrorIs(t, err, callbackErr) + require.Same(t, transaction, sess.Session.GetTransaction()) + require.NotNil(t, sess.implicit, "failed marker must remain armed") + require.Equal(t, 1, callbackCalls, "commit probes the callback before deciding whether to commit") + + ctx.SetTransaction(nil) + require.Equal(t, 1, callbackCalls, "rollback callback is idempotent") + require.Nil(t, sess.Session.GetTransaction()) + require.Nil(t, sess.implicit) +} + +func TestImplicitCleanupMismatchLeavesReplacementAndMarker(t *testing.T) { + sess, ctx, transaction := newFrontendImplicitSession(t) + replacement := &Transaction{} + callbackCalls := 0 + sess.RegisterImplicitTransactionCleanup(ctx, transaction, func() error { + callbackCalls++ + return nil + }) + + // Simulate a GMS temporary transaction swap. Use the embedded setter to + // avoid changing the marker while constructing the replacement. + sess.Session.SetTransaction(replacement) + ctx.SetTransaction(nil) + require.Same(t, replacement, sess.Session.GetTransaction()) + require.NotNil(t, sess.implicit, "a mismatched nil setter must not consume the marker") + require.Zero(t, callbackCalls) + + // Once the original transaction is current again, the nil setter can safely + // consume its marker and perform the rollback. + sess.Session.SetTransaction(transaction) + ctx.SetTransaction(nil) + require.Equal(t, 1, callbackCalls) + require.Nil(t, sess.Session.GetTransaction()) + require.Nil(t, sess.implicit) +} + +func TestImplicitCleanupRegistrationAndDetachAreAtomic(t *testing.T) { + // Hold the state mutex so registration and rollback contend at the exact + // check/append boundary. With a split implicitMu check followed by add, the + // detach path can finish first and registration then appends to an unreachable + // state. The atomic implementation keeps implicitMu held until add completes. + for iteration := 0; iteration < 100; iteration++ { + sess, ctx, transaction := newFrontendImplicitSession(t) + state := sess.implicit + require.NotNil(t, state) + + var callbackCalls atomic.Int32 + state.mu.Lock() + registerDone := make(chan struct{}) + go func() { + sess.RegisterImplicitTransactionCleanup(ctx, transaction, func() error { + callbackCalls.Add(1) + return nil + }) + close(registerDone) + }() + // Let the registration goroutine reach the state-level append lock before + // starting the competing detach. The bounded wait below keeps the test + // deterministic without depending on a scheduler-specific sleep. + for i := 0; i < 8; i++ { + runtime.Gosched() + } + + detachDone := make(chan struct{}) + go func() { + ctx.SetTransaction(nil) + close(detachDone) + }() + select { + case <-detachDone: + case <-time.After(5 * time.Millisecond): + } + state.mu.Unlock() + + select { + case <-registerDone: + case <-time.After(time.Second): + t.Fatal("implicit cleanup registration did not finish") + } + select { + case <-detachDone: + case <-time.After(time.Second): + t.Fatal("implicit cleanup detach did not finish") + } + + state.mu.Lock() + callbackCount := len(state.callbacks) + cleanupStarted := state.cleanupStarted + state.mu.Unlock() + if callbackCalls.Load() == 0 { + // A registration that loses the race must be rejected completely; it + // must not remain queued on the detached marker. + require.Zero(t, callbackCount) + } else { + require.Equal(t, int32(1), callbackCalls.Load()) + require.True(t, cleanupStarted) + require.Equal(t, 1, callbackCount) + } + } +} + +func TestImplicitCleanupStoredProcedureDetachRestoresOuterMarker(t *testing.T) { + sess, ctx, outer := newFrontendImplicitSession(t) + callbackCalls := 0 + var callbackTx sql.Transaction + sess.RegisterImplicitTransactionCleanup(ctx, outer, func() error { + callbackCalls++ + callbackTx = sess.Session.GetTransaction() + return nil + }) + + // This is the sequence in GMS rowexec/proc.go: buildCall saves oldTx, + // clears it before running the body, then restores oldTx in a defer. + endProcedure := sess.beginProcedureTransactionScope() + ctx.SetTransaction(nil) + require.Nil(t, sess.Session.GetTransaction()) + require.Nil(t, sess.implicit) + require.Zero(t, callbackCalls, "temporary procedure detach must not clean up outer state") + + innerSQLTx, err := sess.StartTransaction(ctx, sql.ReadWrite) + require.NoError(t, err) + inner, ok := innerSQLTx.(*Transaction) + require.True(t, ok) + ctx.SetTransaction(inner) + + // The deferred restore must identify the original transaction, resolve any + // procedure-owned child, and re-arm the original marker. + ctx.SetTransaction(outer) + require.Same(t, outer, sess.Session.GetTransaction()) + require.Same(t, outer, sess.implicit.tx) + require.Zero(t, callbackCalls) + + endProcedure() + ctx.SetTransaction(nil) + require.Equal(t, 1, callbackCalls) + require.Same(t, outer, callbackTx) + require.Nil(t, sess.Session.GetTransaction()) + require.Nil(t, sess.implicit) +} + +func TestImplicitCleanupStoredProcedureNestedScopesKeepMarkerIdentity(t *testing.T) { + sess, ctx, outer := newFrontendImplicitSession(t) + outerCalls, innerCalls := 0, 0 + sess.RegisterImplicitTransactionCleanup(ctx, outer, func() error { + outerCalls++ + return nil + }) + + endOuter := sess.beginProcedureTransactionScope() + ctx.SetTransaction(nil) + + innerTx, err := sess.StartTransaction(ctx, sql.ReadWrite) + require.NoError(t, err) + inner := innerTx.(*Transaction) + ctx.SetTransaction(inner) + sess.RegisterImplicitTransactionCleanup(ctx, inner, func() error { + innerCalls++ + return nil + }) + + endInner := sess.beginProcedureTransactionScope() + ctx.SetTransaction(nil) + require.Nil(t, sess.Session.GetTransaction()) + require.Nil(t, sess.implicit) + + // A child transaction created inside the nested scope is cleaned before the + // nested outer marker is restored. + childTx, err := sess.StartTransaction(ctx, sql.ReadWrite) + require.NoError(t, err) + child := childTx.(*Transaction) + ctx.SetTransaction(child) + ctx.SetTransaction(inner) + require.Same(t, inner, sess.Session.GetTransaction()) + require.Same(t, inner, sess.implicit.tx) + require.Zero(t, innerCalls, "restoring the nested outer marker must not clean it up") + + endInner() + ctx.SetTransaction(outer) + require.Same(t, outer, sess.Session.GetTransaction()) + require.Same(t, outer, sess.implicit.tx) + require.Equal(t, 1, innerCalls, "the nested marker is cleaned before restoring the outer marker") + require.Zero(t, outerCalls) + + endOuter() + ctx.SetTransaction(nil) + require.Equal(t, 1, outerCalls) + require.Nil(t, sess.Session.GetTransaction()) + require.Nil(t, sess.implicit) +} + +func TestImplicitCleanupOnlyArmsForFrontendOrigin(t *testing.T) { + origins := []mycontext.QueryOriginKind{ + mycontext.UnknownQueryOrigin, + mycontext.InternalQueryOrigin, + mycontext.MySQLReplicationQueryOrigin, + mycontext.PostgresReplicationQueryOrigin, + mycontext.MaintenanceQueryOrigin, + mycontext.RecoveryQueryOrigin, + } + for _, origin := range origins { + t.Run(originName(origin), func(t *testing.T) { + sess := NewSession(memory.NewSession(sql.NewBaseSession(), nil), nil) + base := context.Background() + if origin != mycontext.UnknownQueryOrigin { + base = mycontext.WithQueryOrigin(base, origin) + } + ctx := sql.NewContext(base, sql.WithSession(sess)) + tx, err := sess.StartTransaction(ctx, sql.ReadWrite) + require.NoError(t, err) + require.Nil(t, sess.implicit) + ctx.SetTransaction(tx) + sess.Session.SetTransaction(nil) + }) + } +} + +func TestSetAutocommitZeroDisarmsImplicitCleanupOnce(t *testing.T) { + sess, ctx, transaction := newFrontendImplicitSession(t) + callbackCalls := 0 + sess.RegisterImplicitTransactionCleanup(ctx, transaction, func() error { + callbackCalls++ + return nil + }) + + require.NoError(t, sess.SetSessionVariable(ctx, sql.AutoCommitSessionVar, int8(0))) + require.Equal(t, 1, callbackCalls) + require.Nil(t, sess.implicit) + require.Nil(t, sess.Session.GetTransaction()) + ctx.SetTransaction(nil) + require.Equal(t, 1, callbackCalls) +} + +func originName(origin mycontext.QueryOriginKind) string { + switch origin { + case mycontext.UnknownQueryOrigin: + return "unknown" + case mycontext.InternalQueryOrigin: + return "internal" + case mycontext.MySQLReplicationQueryOrigin: + return "mysql-replication" + case mycontext.PostgresReplicationQueryOrigin: + return "postgres-replication" + case mycontext.MaintenanceQueryOrigin: + return "maintenance" + case mycontext.RecoveryQueryOrigin: + return "recovery" + default: + return "origin" + } +} diff --git a/backend/session_operation_test.go b/backend/session_operation_test.go new file mode 100644 index 00000000..72fae133 --- /dev/null +++ b/backend/session_operation_test.go @@ -0,0 +1,261 @@ +package backend + +import ( + "context" + stdsql "database/sql" + "errors" + "sync/atomic" + "testing" + + "github.com/apecloud/myduckserver/adapter" + "github.com/apecloud/myduckserver/catalog" + "github.com/apecloud/myduckserver/mycontext" + "github.com/dolthub/go-mysql-server/memory" + "github.com/dolthub/go-mysql-server/sql" + "github.com/stretchr/testify/require" +) + +func TestSessionBeginDuckLakeOperationNoopsWithoutProvider(t *testing.T) { + // A nil provider/session must always return an idempotent no-op release. + var nilSession *Session + release := nilSession.BeginDuckLakeOperation(mycontext.WithFrontendQuery(context.Background())) + release() + release() + session := &Session{} + release = session.BeginDuckLakeOperation(mycontext.WithFrontendQuery(context.Background())) + release() + + // Keep the forwarding contract compile-checked even when no runtime is + // configured in this package-level test. + var _ interface { + BeginDuckLakeOperation(context.Context) func() + } = (*Session)(nil) +} + +func TestSessionEndRollsBackLogicalImplicitTransactionOnce(t *testing.T) { + session, ctx, transaction := newFrontendImplicitSession(t) + var callbackCalls atomic.Int32 + require.True(t, session.HasImplicitTransactionCleanup(transaction)) + session.RegisterImplicitTransactionCleanup(ctx, transaction, func() error { + callbackCalls.Add(1) + return nil + }) + + // SessionManager invokes this callback when COM_RESET_CONNECTION replaces a + // GMS session. A second invocation can occur during disconnect cleanup and + // must not run the operation-release callback twice. + session.SessionEnd() + session.SessionEnd() + + require.Equal(t, int32(1), callbackCalls.Load()) + require.Nil(t, session.Session.GetTransaction()) + require.Nil(t, session.implicit) +} + +func TestSessionEndRollsBackPhysicalTransactionAndKeepsConnection(t *testing.T) { + provider := catalog.NewInMemoryDBProvider() + t.Cleanup(func() { require.NoError(t, provider.Close()) }) + session := NewSession(memory.NewSession(sql.NewBaseSession(), provider), provider) + ctx := sql.NewContext( + mycontext.WithFrontendQuery(context.Background()), + sql.WithSession(session), + ) + ctx.SetIgnoreAutoCommit(true) + transactionValue, err := session.StartTransaction(ctx, sql.ReadWrite) + require.NoError(t, err) + transaction, ok := transactionValue.(*Transaction) + require.True(t, ok) + require.NotNil(t, transaction.tx) + ctx.SetTransaction(transaction) + + ownerBefore, txBefore := session.GetTxnBinding() + require.NotNil(t, ownerBefore) + require.Same(t, transaction.tx, txBefore) + var callbackCalls atomic.Int32 + var callbackSuccess atomic.Bool + require.True(t, session.RegisterPostPhysicalTransactionCleanup(nil, transaction.tx, func(success bool) { + callbackCalls.Add(1) + callbackSuccess.Store(success) + })) + + session.SessionEnd() + + require.Equal(t, int32(1), callbackCalls.Load()) + require.False(t, callbackSuccess.Load()) + require.Nil(t, session.Session.GetTransaction()) + ownerAfter, txAfter := session.GetTxnBinding() + require.Same(t, ownerBefore, ownerAfter, "SessionEnd must not close the generic connection") + require.Nil(t, txAfter) +} + +func TestSessionCloseTxnIfCompletesPhysicalPostCleanup(t *testing.T) { + provider := catalog.NewInMemoryDBProvider() + t.Cleanup(func() { require.NoError(t, provider.Close()) }) + session := NewSession(memory.NewSession(sql.NewBaseSession(), provider), provider) + + tx, err := session.GetCatalogTxn(context.Background(), nil) + require.NoError(t, err) + require.NotNil(t, tx) + + var callbackCalls atomic.Int32 + var callbackSuccess atomic.Bool + require.True(t, session.RegisterPostPhysicalTransactionCleanup(nil, tx, func(success bool) { + callbackCalls.Add(1) + callbackSuccess.Store(success) + })) + + // CloseTxnIf is used when the protocol has lost a trustworthy owner and + // cannot run the normal commit/rollback finalizer. It must still release the + // callback associated with this exact physical transaction. + session.CloseTxnIf(tx) + require.Equal(t, int32(1), callbackCalls.Load()) + require.False(t, callbackSuccess.Load()) + require.Nil(t, session.TryGetTxnForRelease()) + + // The raw transaction itself is still owned by the caller after an + // identity-only mapping removal; finish it so the test does not hold a + // database/sql connection through provider teardown. + require.NoError(t, tx.Rollback()) +} + +func TestSessionCloseTxnIfDoesNotCompleteReplacementCallback(t *testing.T) { + provider := catalog.NewInMemoryDBProvider() + t.Cleanup(func() { require.NoError(t, provider.Close()) }) + session := NewSession(memory.NewSession(sql.NewBaseSession(), provider), provider) + + first, err := session.GetCatalogTxn(context.Background(), nil) + require.NoError(t, err) + require.NotNil(t, first) + var callbackCalls atomic.Int32 + require.True(t, session.RegisterPostPhysicalTransactionCleanup(nil, first, func(bool) { + callbackCalls.Add(1) + })) + + // A stale callback carrying another transaction must not remove the current + // mapping or complete the current transaction's callback. + session.CloseTxnIf(new(stdsql.Tx)) + require.Same(t, first, session.TryGetTxnForRelease()) + require.Zero(t, callbackCalls.Load()) + + session.CloseTxnIf(first) + require.Equal(t, int32(1), callbackCalls.Load()) + require.Nil(t, session.TryGetTxnForRelease()) + require.NoError(t, first.Rollback()) +} + +func TestPostPhysicalCleanupLateRegistrationSeesCompletion(t *testing.T) { + provider := catalog.NewInMemoryDBProvider() + t.Cleanup(func() { require.NoError(t, provider.Close()) }) + session := NewSession(memory.NewSession(sql.NewBaseSession(), provider), provider) + + tx, err := session.GetCatalogTxn(context.Background(), nil) + require.NoError(t, err) + require.NotNil(t, tx) + + // Model pool teardown/finalization winning the race before COPY has + // registered its deferred operation callback. + session.completePhysicalPostCleanup(tx, false) + var calls atomic.Int32 + var outcome atomic.Bool + require.True(t, session.RegisterPostPhysicalTransactionCleanup(nil, tx, func(success bool) { + calls.Add(1) + outcome.Store(success) + })) + require.Equal(t, int32(1), calls.Load()) + require.False(t, outcome.Load()) + + // The mapping is still owned by the test session; finish the raw handle so + // provider teardown does not retain a database/sql connection. + require.NoError(t, tx.Rollback()) +} + +func TestPostLogicalCleanupLateRegistrationSeesCompletion(t *testing.T) { + transaction := &Transaction{} + transaction.completePostCleanup(false) + + var calls atomic.Int32 + var outcome atomic.Bool + require.True(t, transaction.registerPostCleanup(func(success bool) { + calls.Add(1) + outcome.Store(success) + })) + require.Equal(t, int32(1), calls.Load()) + require.False(t, outcome.Load()) +} + +func TestCloseConnIfBindingCompletesCallbackWhenRollbackAlreadyDone(t *testing.T) { + provider := catalog.NewInMemoryDBProvider() + t.Cleanup(func() { require.NoError(t, provider.Close()) }) + session := NewSession(memory.NewSession(sql.NewBaseSession(), provider), provider) + + started, err := session.GetCatalogTxn(context.Background(), nil) + require.NoError(t, err) + require.NotNil(t, started) + owner, tx := session.GetTxnBindingForRelease() + require.NotNil(t, owner) + require.NotNil(t, tx) + var calls atomic.Int32 + require.True(t, session.RegisterPostPhysicalTransactionCleanup(nil, tx, func(success bool) { + require.False(t, success) + calls.Add(1) + })) + + // Force CloseConnIfBinding to observe an already-inactive driver + // transaction. The pool still removes the exact mapping, but returns the + // driver's ErrTxDone; callback completion must follow the mapping removal, + // not the returned close error. + require.NoError(t, tx.Rollback()) + _ = session.CloseConnIfBinding(owner, tx) + require.Equal(t, int32(1), calls.Load()) + require.Nil(t, session.TryGetTxnForRelease()) +} + +func TestProviderCloseCompletesPhysicalPostCleanup(t *testing.T) { + provider := catalog.NewInMemoryDBProvider() + session := NewSession(memory.NewSession(sql.NewBaseSession(), provider), provider) + started, err := session.GetCatalogTxn(context.Background(), nil) + require.NoError(t, err) + require.NotNil(t, started) + _, tx := session.GetTxnBindingForRelease() + require.NotNil(t, tx) + + var calls atomic.Int32 + require.True(t, session.RegisterPostPhysicalTransactionCleanup(nil, tx, func(success bool) { + require.False(t, success) + calls.Add(1) + })) + + // Provider teardown finalizes through ConnectionPool.closeLocked rather + // than Session.CloseTxn; the pool completion bridge must still resolve the + // protocol callback exactly once. + require.NoError(t, provider.Close()) + require.Equal(t, int32(1), calls.Load()) +} + +func TestSessionCommitTxnReportsStaleBindingAndFailureOutcome(t *testing.T) { + provider := catalog.NewInMemoryDBProvider() + t.Cleanup(func() { _ = provider.Close() }) + session := NewSession(memory.NewSession(sql.NewBaseSession(), provider), provider) + + tx, err := session.GetCatalogTxn(context.Background(), nil) + require.NoError(t, err) + var calls atomic.Int32 + var success atomic.Bool + require.True(t, session.RegisterPostPhysicalTransactionCleanup(nil, tx, func(ok bool) { + calls.Add(1) + success.Store(ok) + })) + + // Retire the binding through the normal session close path, then deliver a + // delayed commit callback for the old physical transaction. The stale call + // must be visible to its caller and must not be reported as success. + require.NoError(t, tx.Rollback()) + session.CloseTxn() + require.Equal(t, int32(1), calls.Load()) + require.False(t, success.Load()) + + err = session.CommitTxn(tx) + require.Error(t, err) + require.True(t, errors.Is(err, adapter.ErrTransactionBindingChanged)) + require.Equal(t, int32(1), calls.Load()) +} diff --git a/binlogreplication/binlog_replica_applier.go b/binlogreplication/binlog_replica_applier.go index 7922afd1..ef46b1ca 100644 --- a/binlogreplication/binlog_replica_applier.go +++ b/binlogreplication/binlog_replica_applier.go @@ -16,7 +16,6 @@ package binlogreplication import ( "context" - stdsql "database/sql" "encoding/binary" "errors" "fmt" @@ -654,11 +653,10 @@ func (a *binlogReplicaApplier) discardOngoingTxn(ctx *sql.Context) error { // Most replication transactions are owned by the GMS transaction above. // Keep a fallback for direct writers that opened only the DuckDB transaction. - if txn := adapter.TryGetTxn(ctx); txn != nil { - if err := txn.Rollback(); err != nil && err != stdsql.ErrTxDone { + if txn := adapter.TryGetTxnForRelease(ctx); txn != nil { + if _, _, err := adapter.FinalizeRollback(ctx, txn); err != nil && !adapter.IsTransactionInactiveError(err) { rollbackErr = errors.Join(rollbackErr, err) } - adapter.CloseTxn(ctx) } a.tableWriterProvider.DiscardDeltaBuffer(ctx) @@ -1003,11 +1001,10 @@ func (a *binlogReplicaApplier) commitOngoingTxn(ctx *sql.Context, engine *gms.En } // The session manager does not start an actual transaction in autocommit=1 mode, // but there may be a transaction in progress if we have started it manually. - if tx := adapter.TryGetTxn(ctx); tx != nil { - if err := tx.Commit(); err != nil && err != stdsql.ErrTxDone { + if tx := adapter.TryGetTxnForRelease(ctx); tx != nil { + if err := adapter.FinalizeCommit(ctx, tx); err != nil && !adapter.IsTransactionInactiveError(err) { return err } - adapter.CloseTxn(ctx) } // --- Update the in-memory states --- // @@ -1307,10 +1304,11 @@ func (a *binlogReplicaApplier) writeChanges( dataRows = append(dataRows, dataRow) } - txn, err := adapter.GetTxn(ctx, nil) + _, txn, release, err := adapter.GetTxnExecutionSnapshotWithLease(ctx, nil) if err != nil { return err } + defer release() tableWriter, err := a.tableWriterProvider.GetTableWriter( ctx, txn, @@ -1463,14 +1461,11 @@ func (a *binlogReplicaApplier) appendRowFormatChanges( } func (a *binlogReplicaApplier) flushDeltaBuffer(ctx *sql.Context, reason delta.FlushReason) error { - conn, err := adapter.GetCatalogConn(ctx) - if err != nil { - return err - } - tx, err := adapter.GetCatalogTxn(ctx, nil) + conn, tx, release, err := adapter.GetCatalogTxnExecutionSnapshotWithLease(ctx, nil) if err != nil { return err } + defer release() defer a.deltaBufSize.Store(0) diff --git a/catalog/completion_reentrancy_test.go b/catalog/completion_reentrancy_test.go new file mode 100644 index 00000000..b36d54ff --- /dev/null +++ b/catalog/completion_reentrancy_test.go @@ -0,0 +1,389 @@ +package catalog + +import ( + "context" + stdsql "database/sql" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/DATA-DOG/go-sqlmock" + "github.com/stretchr/testify/require" +) + +const completionReentryTimeout = 3 * time.Second + +type completionObservation struct { + calls atomic.Int32 + outcome atomic.Int32 +} + +func registerReentrantCompletion( + t *testing.T, + p *ConnectionPool, + id uint32, + tx *stdsql.Tx, + beforeReturn func() bool, +) *completionObservation { + t.Helper() + observation := new(completionObservation) + require.True(t, p.RegisterTransactionCompletion(tx, func(success bool) { + // Re-enter both the lifecycle release gate and the same session mutex. A + // callback dispatched by a *Locked helper would deadlock here. + _, _ = p.GetTxnBindingForRelease(id) + if beforeReturn != nil && !beforeReturn() { + observation.outcome.Store(-2) + } else if success { + observation.outcome.Store(1) + } else { + observation.outcome.Store(-1) + } + observation.calls.Add(1) + })) + return observation +} + +func runCompletionFinalizer(t *testing.T, call func() error) { + t.Helper() + done := make(chan error, 1) + go func() { + done <- call() + }() + select { + case err := <-done: + require.NoError(t, err) + case <-time.After(completionReentryTimeout): + t.Fatal("transaction completion callback did not return after pool re-entry") + } +} + +func requireCompletionObservation(t *testing.T, observation *completionObservation, wantOutcome int32) { + t.Helper() + require.Equal(t, int32(1), observation.calls.Load(), "completion must run exactly once before the finalizer returns") + require.Equal(t, wantOutcome, observation.outcome.Load()) +} + +func newReentrantTransactionPool( + t *testing.T, + id uint32, +) (*ConnectionPool, *stdsql.DB, sqlmock.Sqlmock, *stdsql.Conn, *stdsql.Tx) { + t.Helper() + db, mock, err := sqlmock.New() + require.NoError(t, err) + p := NewConnectionPool(nil, db, "memory") + p.registerMySQLUDFsOnce.Do(func() {}) + mock.ExpectBegin() + tx, err := p.GetTxn(context.Background(), id, "", nil) + require.NoError(t, err) + conn, bound := p.GetTxnBindingForRelease(id) + require.NotNil(t, conn) + require.Same(t, tx, bound) + return p, db, mock, conn, tx +} + +func TestConnectionPoolMappingFinalizersDispatchCompletionAfterUnlock(t *testing.T) { + t.Run("CloseTxnAndGet", func(t *testing.T) { + p := NewConnectionPool(nil, nil, "memory") + const id uint32 = 401 + tx := new(stdsql.Tx) + p.txns.Store(id, tx) + p.markTransactionActive(tx) + observation := registerReentrantCompletion(t, p, id, tx, nil) + + var removed *stdsql.Tx + runCompletionFinalizer(t, func() error { + removed = p.CloseTxnAndGet(id) + return nil + }) + require.Same(t, tx, removed) + requireCompletionObservation(t, observation, -1) + }) + + t.Run("CloseTxnIf", func(t *testing.T) { + p := NewConnectionPool(nil, nil, "memory") + const id uint32 = 402 + tx := new(stdsql.Tx) + p.txns.Store(id, tx) + p.markTransactionActive(tx) + observation := registerReentrantCompletion(t, p, id, tx, nil) + + runCompletionFinalizer(t, func() error { + p.CloseTxnIf(id, tx) + return nil + }) + requireCompletionObservation(t, observation, -1) + }) +} + +func TestConnectionPoolDriverFinalizersDispatchCompletionAfterUnlock(t *testing.T) { + t.Run("CloseConn", func(t *testing.T) { + const id uint32 = 409 + p, _, mock, _, tx := newReentrantTransactionPool(t, id) + observation := registerReentrantCompletion(t, p, id, tx, nil) + mock.ExpectRollback() + mock.ExpectClose() + + runCompletionFinalizer(t, func() error { + return p.CloseConn(id) + }) + requireCompletionObservation(t, observation, -1) + require.NoError(t, mock.ExpectationsWereMet()) + }) + + t.Run("CloseConnIf", func(t *testing.T) { + const id uint32 = 410 + p, _, mock, conn, tx := newReentrantTransactionPool(t, id) + observation := registerReentrantCompletion(t, p, id, tx, nil) + mock.ExpectRollback() + mock.ExpectClose() + + runCompletionFinalizer(t, func() error { + return p.CloseConnIf(id, conn) + }) + requireCompletionObservation(t, observation, -1) + require.NoError(t, mock.ExpectationsWereMet()) + }) + + t.Run("CommitTxnWithOutcome", func(t *testing.T) { + const id uint32 = 403 + p, _, mock, _, tx := newReentrantTransactionPool(t, id) + observation := registerReentrantCompletion(t, p, id, tx, nil) + mock.ExpectCommit() + + runCompletionFinalizer(t, func() error { + matched, err := p.CommitTxnWithOutcome(id, tx) + if !matched { + return fmt.Errorf("commit did not match the current transaction") + } + return err + }) + requireCompletionObservation(t, observation, 1) + mock.ExpectClose() + require.NoError(t, p.Close()) + require.NoError(t, mock.ExpectationsWereMet()) + }) + + t.Run("RollbackTxn", func(t *testing.T) { + const id uint32 = 404 + p, _, mock, _, tx := newReentrantTransactionPool(t, id) + observation := registerReentrantCompletion(t, p, id, tx, nil) + mock.ExpectRollback() + + runCompletionFinalizer(t, func() error { + _, inactive, err := p.RollbackTxn(id, tx) + if !inactive { + return fmt.Errorf("rollback did not prove the transaction inactive") + } + return err + }) + requireCompletionObservation(t, observation, -1) + mock.ExpectClose() + require.NoError(t, p.Close()) + require.NoError(t, mock.ExpectationsWereMet()) + }) + + t.Run("RollbackTxnWithCleanup", func(t *testing.T) { + const id uint32 = 405 + p, _, mock, _, tx := newReentrantTransactionPool(t, id) + var cleanupDone atomic.Bool + observation := registerReentrantCompletion(t, p, id, tx, cleanupDone.Load) + mock.ExpectRollback() + + runCompletionFinalizer(t, func() error { + inactive, err := p.RollbackTxnWithCleanup(id, tx, func(*stdsql.Conn) error { + cleanupDone.Store(true) + return nil + }) + if !inactive { + return fmt.Errorf("rollback did not prove the transaction inactive") + } + return err + }) + requireCompletionObservation(t, observation, -1) + mock.ExpectClose() + require.NoError(t, p.Close()) + require.NoError(t, mock.ExpectationsWereMet()) + }) + + t.Run("RollbackTxnWithCleanup late registration", func(t *testing.T) { + const id uint32 = 408 + p, _, mock, _, tx := newReentrantTransactionPool(t, id) + var cleanupDone atomic.Bool + observation := new(completionObservation) + var registered bool + mock.ExpectRollback() + + runCompletionFinalizer(t, func() error { + inactive, err := p.RollbackTxnWithCleanup(id, tx, func(*stdsql.Conn) error { + registered = p.RegisterTransactionCompletion(tx, func(success bool) { + // Registration happens after the terminal outcome is recorded but + // before the pool/session locks are released. It must join the + // deferred batch instead of running inline under those locks. + _, _ = p.GetTxnBindingForRelease(id) + if !cleanupDone.Load() { + observation.outcome.Store(-2) + } else if success { + observation.outcome.Store(1) + } else { + observation.outcome.Store(-1) + } + observation.calls.Add(1) + }) + cleanupDone.Store(true) + return nil + }) + if !inactive { + return fmt.Errorf("rollback did not prove the transaction inactive") + } + return err + }) + require.True(t, registered) + requireCompletionObservation(t, observation, -1) + mock.ExpectClose() + require.NoError(t, p.Close()) + require.NoError(t, mock.ExpectationsWereMet()) + }) + + t.Run("CloseConnIfBinding", func(t *testing.T) { + const id uint32 = 406 + p, _, mock, conn, tx := newReentrantTransactionPool(t, id) + observation := registerReentrantCompletion(t, p, id, tx, nil) + mock.ExpectRollback() + mock.ExpectClose() + + runCompletionFinalizer(t, func() error { + return p.CloseConnIfBinding(id, conn, tx) + }) + requireCompletionObservation(t, observation, -1) + require.NoError(t, mock.ExpectationsWereMet()) + }) +} + +func TestConnectionPoolGenerationFinalizersDispatchInOrderAfterUnlock(t *testing.T) { + for _, reset := range []bool{false, true} { + name := "Close" + if reset { + name = "Reset" + } + t.Run(name, func(t *testing.T) { + const id uint32 = 407 + p, _, mock, _, tx := newReentrantTransactionPool(t, id) + var orderMu sync.Mutex + order := make([]string, 0, 2) + p.SetTransactionLifecycleHooks(nil, func(*stdsql.Tx) { + _, _ = p.GetTxnBindingForRelease(id) + orderMu.Lock() + order = append(order, "finished") + orderMu.Unlock() + }) + require.True(t, p.RegisterTransactionCompletion(tx, func(bool) { + _, _ = p.GetTxnBindingForRelease(id) + orderMu.Lock() + order = append(order, "callback") + orderMu.Unlock() + })) + mock.ExpectRollback() + mock.ExpectClose() + + var replacementDB *stdsql.DB + var replacementMock sqlmock.Sqlmock + if reset { + var err error + replacementDB, replacementMock, err = sqlmock.New() + require.NoError(t, err) + } + runCompletionFinalizer(t, func() error { + if reset { + return p.Reset(nil, replacementDB) + } + return p.Close() + }) + + orderMu.Lock() + gotOrder := append([]string(nil), order...) + orderMu.Unlock() + require.Equal(t, []string{"callback", "finished"}, gotOrder) + require.NoError(t, mock.ExpectationsWereMet()) + + if reset { + replacementMock.ExpectClose() + require.NoError(t, p.Close()) + require.NoError(t, replacementMock.ExpectationsWereMet()) + } + }) + } +} + +func TestTransactionCompletionBatchDrainsAfterCallbackPanic(t *testing.T) { + var eventsMu sync.Mutex + events := make([]string, 0, 4) + appendEvent := func(event string) { + eventsMu.Lock() + events = append(events, event) + eventsMu.Unlock() + } + + batch := transactionCompletionBatch{dispatches: []*transactionCompletionDispatch{ + { + callbacks: []func(bool){func(bool) { + appendEvent("first-callback") + panic("completion panic") + }}, + finished: func(*stdsql.Tx) { appendEvent("first-finished") }, + }, + { + callbacks: []func(bool){func(bool) { appendEvent("second-callback") }}, + finished: func(*stdsql.Tx) { appendEvent("second-finished") }, + }, + }} + + require.PanicsWithValue(t, "completion panic", batch.run) + eventsMu.Lock() + gotEvents := append([]string(nil), events...) + eventsMu.Unlock() + require.Equal(t, []string{ + "first-callback", + "first-finished", + "second-callback", + "second-finished", + }, gotEvents) + require.Empty(t, batch.dispatches) +} + +func TestTransactionCompletionBatchDrainsAfterFinishedHookPanic(t *testing.T) { + var eventsMu sync.Mutex + events := make([]string, 0, 4) + appendEvent := func(event string) { + eventsMu.Lock() + events = append(events, event) + eventsMu.Unlock() + } + + batch := transactionCompletionBatch{dispatches: []*transactionCompletionDispatch{ + { + callbacks: []func(bool){func(bool) { appendEvent("first-callback") }}, + finished: func(*stdsql.Tx) { + appendEvent("first-finished") + panic("finished hook panic") + }, + }, + { + callbacks: []func(bool){func(bool) { appendEvent("second-callback") }}, + finished: func(*stdsql.Tx) { appendEvent("second-finished") }, + }, + }} + + require.PanicsWithValue(t, "finished hook panic", batch.run) + eventsMu.Lock() + gotEvents := append([]string(nil), events...) + eventsMu.Unlock() + require.Equal(t, []string{ + "first-callback", + "first-finished", + "second-callback", + "second-finished", + }, gotEvents) + require.Empty(t, batch.dispatches) +} diff --git a/catalog/completion_registry_test.go b/catalog/completion_registry_test.go new file mode 100644 index 00000000..736e8b28 --- /dev/null +++ b/catalog/completion_registry_test.go @@ -0,0 +1,232 @@ +package catalog + +import ( + stdsql "database/sql" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func newCompletionRegistryPoolForTest() *ConnectionPool { + return &ConnectionPool{ + transactionCompletions: make(map[*stdsql.Tx]*transactionCompletion), + activeTransactions: make(map[*stdsql.Tx]struct{}), + } +} + +func finishTransactionForTest(p *ConnectionPool, tx *stdsql.Tx, success bool) { + p.finishTransactionWithOutcomeDeferred(tx, success).run() +} + +func TestTransactionCompletionTombstonesBoundedAndPendingRetained(t *testing.T) { + p := newCompletionRegistryPoolForTest() + + pendingTx := new(stdsql.Tx) + p.markTransactionActive(pendingTx) + var pendingCalls atomic.Int32 + require.True(t, p.RegisterTransactionCompletion(pendingTx, func(bool) { + pendingCalls.Add(1) + })) + + completed := make([]*stdsql.Tx, maxTransactionCompletionTombstones+8) + for i := range completed { + completed[i] = new(stdsql.Tx) + finishTransactionForTest(p, completed[i], i%2 == 0) + } + + p.transactionCompletionMu.Lock() + mapLen := len(p.transactionCompletions) + orderLen := len(p.transactionCompletionOrder) + _, pendingRetained := p.transactionCompletions[pendingTx] + _, oldestRetained := p.transactionCompletions[completed[0]] + _, newestRetained := p.transactionCompletions[completed[len(completed)-1]] + p.transactionCompletionMu.Unlock() + + require.LessOrEqual(t, mapLen, maxTransactionCompletionTombstones+1) + require.LessOrEqual(t, orderLen, maxTransactionCompletionTombstones) + require.True(t, pendingRetained, "pending callback state must never be evicted") + require.False(t, oldestRetained, "old completed tombstones should be evicted") + require.True(t, newestRetained) + require.Zero(t, pendingCalls.Load()) +} + +func TestTransactionCompletionTTLAndLateRegistration(t *testing.T) { + p := newCompletionRegistryPoolForTest() + + recent := new(stdsql.Tx) + finishTransactionForTest(p, recent, true) + var recentSuccess atomic.Bool + require.True(t, p.RegisterTransactionCompletion(recent, func(success bool) { + recentSuccess.Store(success) + })) + require.True(t, recentSuccess.Load(), "a recent tombstone must preserve its outcome") + + expired := new(stdsql.Tx) + finishTransactionForTest(p, expired, true) + p.transactionCompletionMu.Lock() + state := p.transactionCompletions[expired] + p.transactionCompletionMu.Unlock() + require.NotNil(t, state) + state.mu.Lock() + state.completedAt = time.Now().Add(-transactionCompletionTombstoneTTL - time.Second) + state.mu.Unlock() + + var expiredCalls atomic.Int32 + require.False(t, p.RegisterTransactionCompletion(expired, func(success bool) { + expiredCalls.Add(1) + require.False(t, success) + })) + require.Equal(t, int32(1), expiredCalls.Load(), "expired registration must fail closed and resolve fallback") + + p.transactionCompletionMu.Lock() + _, stillRetained := p.transactionCompletions[expired] + p.transactionCompletionMu.Unlock() + require.False(t, stillRetained) +} + +func TestClearCompletedTransactionCompletionsKeepsPending(t *testing.T) { + p := newCompletionRegistryPoolForTest() + completed := new(stdsql.Tx) + finishTransactionForTest(p, completed, false) + pending := new(stdsql.Tx) + p.markTransactionActive(pending) + require.True(t, p.RegisterTransactionCompletion(pending, func(bool) {})) + + p.clearCompletedTransactionCompletions() + + p.transactionCompletionMu.Lock() + _, completedRetained := p.transactionCompletions[completed] + _, pendingRetained := p.transactionCompletions[pending] + orderLen := len(p.transactionCompletionOrder) + p.transactionCompletionMu.Unlock() + require.False(t, completedRetained) + require.True(t, pendingRetained) + require.Zero(t, orderLen) +} + +func TestTransactionCompletionRegisterAndFinishRace(t *testing.T) { + p := newCompletionRegistryPoolForTest() + tx := new(stdsql.Tx) + p.markTransactionActive(tx) + start := make(chan struct{}) + var calls atomic.Int32 + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + <-start + p.RegisterTransactionCompletion(tx, func(bool) { + calls.Add(1) + }) + }() + go func() { + defer wg.Done() + <-start + finishTransactionForTest(p, tx, false) + }() + close(start) + wg.Wait() + require.Equal(t, int32(1), calls.Load()) +} + +func TestTransactionCompletionTombstoneWaitsForFinishedHook(t *testing.T) { + p := newCompletionRegistryPoolForTest() + tx := new(stdsql.Tx) + p.markTransactionActive(tx) + + finishedStarted := make(chan struct{}) + releaseFinished := make(chan struct{}) + defer func() { + select { + case <-releaseFinished: + default: + close(releaseFinished) + } + }() + p.SetTransactionLifecycleHooks(nil, func(*stdsql.Tx) { + close(finishedStarted) + <-releaseFinished + }) + + dispatch := p.finishTransactionWithOutcomeDeferred(tx, true) + require.NotNil(t, dispatch) + dispatchDone := make(chan struct{}) + go func() { + defer close(dispatchDone) + dispatch.run() + }() + + select { + case <-finishedStarted: + case <-time.After(time.Second): + t.Fatal("transaction finished hook did not start") + } + + var lateCalls atomic.Int32 + require.True(t, p.RegisterTransactionCompletion(tx, func(success bool) { + require.True(t, success) + lateCalls.Add(1) + })) + require.Equal(t, int32(1), lateCalls.Load()) + + // dispatchComplete is already true, but the owner has not completed its + // finished hook. Neither TTL pruning nor lifecycle tombstone clearing may + // remove this in-flight state. + p.transactionCompletionMu.Lock() + p.pruneTransactionCompletionsLocked(time.Now().Add(transactionCompletionTombstoneTTL + time.Second)) + _, retainedAfterPrune := p.transactionCompletions[tx] + orderLenBeforeFinished := len(p.transactionCompletionOrder) + p.transactionCompletionMu.Unlock() + require.True(t, retainedAfterPrune) + require.Zero(t, orderLenBeforeFinished) + + p.clearCompletedTransactionCompletions() + p.transactionCompletionMu.Lock() + _, retainedAfterClear := p.transactionCompletions[tx] + p.transactionCompletionMu.Unlock() + require.True(t, retainedAfterClear) + + close(releaseFinished) + select { + case <-dispatchDone: + case <-time.After(time.Second): + t.Fatal("transaction completion dispatch did not finish") + } + + p.transactionCompletionMu.Lock() + state := p.transactionCompletions[tx] + orderLenAfterFinished := len(p.transactionCompletionOrder) + p.transactionCompletionMu.Unlock() + require.NotNil(t, state) + recorded, _ := state.completionSnapshot() + require.True(t, recorded) + require.Equal(t, 1, orderLenAfterFinished) +} + +func TestTransactionCompletionDuplicateFinishHasSingleOwner(t *testing.T) { + p := newCompletionRegistryPoolForTest() + tx := new(stdsql.Tx) + p.markTransactionActive(tx) + + var callbackCalls atomic.Int32 + var finishedCalls atomic.Int32 + p.SetTransactionLifecycleHooks(nil, func(*stdsql.Tx) { + finishedCalls.Add(1) + }) + require.True(t, p.RegisterTransactionCompletion(tx, func(success bool) { + require.True(t, success) + callbackCalls.Add(1) + })) + + owner := p.finishTransactionWithOutcomeDeferred(tx, true) + duplicate := p.finishTransactionWithOutcomeDeferred(tx, false) + require.NotNil(t, owner) + require.Nil(t, duplicate) + owner.run() + + require.Equal(t, int32(1), callbackCalls.Load()) + require.Equal(t, int32(1), finishedCalls.Load()) +} diff --git a/catalog/connpool.go b/catalog/connpool.go index f45bcbab..95ee2b2b 100644 --- a/catalog/connpool.go +++ b/catalog/connpool.go @@ -21,6 +21,8 @@ import ( "fmt" "strings" "sync" + "sync/atomic" + "time" "github.com/dolthub/go-mysql-server/sql" "github.com/duckdb/duckdb-go/v2" @@ -29,25 +31,523 @@ import ( type ConnectionPool struct { *stdsql.DB - connector *duckdb.Connector - defaultCatalog string + connector *duckdb.Connector + defaultCatalog string + // lifecycleMu protects pool generation changes (Close/Reset). Individual + // session state is serialized by sessionLocks so a blocked acquisition for + // one logical session cannot prevent another session from releasing its + // connection. + lifecycleMu lifecycleGate + // dbPtr is an atomic shutdown signal. Close/Reset use it to close the + // current database before waiting on lifecycleMu, which releases pending + // database/sql connection requests that may already hold lifecycleMu.RLock. + dbPtr atomic.Pointer[stdsql.DB] + sessionLocks sync.Map // map[uint32]*sync.Mutex conns sync.Map // concurrent-safe map[uint32]*stdsql.Conn txns sync.Map // concurrent-safe map[uint32]*stdsql.Tx + txnConns sync.Map // concurrent-safe map[uint32]*stdsql.Conn closedConns sync.Map // connection IDs that completed their lifecycle + externalSessions sync.Map // map[*ExternalSession]struct{} + externalGeneration uint64 // protected by lifecycleMu + externalPanicMu sync.Mutex + externalPanicErr error initializerMu sync.RWMutex connectionInitializer func(context.Context, *stdsql.Conn) error - registerMySQLUDFsOnce sync.Once - registerMySQLUDFsErr error + transactionHooksMu sync.RWMutex + transactionAdmission func() func(*stdsql.Tx) + // transactionAdmissionWithError is the provider-aware admission surface. + // The legacy transactionAdmission hook cannot report a poisoned generation + // because its nil/empty completion callback historically meant only that a + // driver BeginTx failed. Keep both fields so standalone pools and older test + // doubles retain their original contract while production providers fail + // closed before opening a new physical transaction. + transactionAdmissionWithError func() (func(*stdsql.Tx), error) + // transactionAdmissionGuard is an optional provider-owned check used to + // reject new physical transactions after a storage generation has become + // unusable. It is kept separate from transactionAdmission because the latter + // predates error-returning admission and a nil completion callback historically + // meant only that the attempted BeginTx failed. + transactionAdmissionGuard func() error + // connectionAdmissionGuard rejects ordinary connection/snapshot work while + // the owning provider storage generation is transitioning or poisoned. It is + // context-aware so provider-owned rollback cleanup and recovery can reuse an + // explicitly reserved connection without opening the gate to frontend work. + connectionAdmissionGuard func(context.Context) error + transactionFinished func(*stdsql.Tx) + transactionCompletionMu sync.Mutex + transactionCompletions map[*stdsql.Tx]*transactionCompletion + transactionCompletionOrder []*stdsql.Tx + activeTransactions map[*stdsql.Tx]struct{} + registerMySQLUDFsOnce sync.Once + registerMySQLUDFsErr error +} + +var ( + ErrExternalSessionClosed = errors.New("external session is closed") + ErrExternalSessionStale = errors.New("external session belongs to a stale pool generation") + ErrExternalSessionGenerationPoisoned = errors.New("external session driver panic poisoned the storage generation") + ErrExternalSessionTransactionMismatch = errors.New( + "external session transaction does not match the active transaction", + ) +) + +// ExternalSession is an opaque, pool-owned connection identity for protocol +// surfaces whose handles do not share the uint32 namespace used by MySQL and +// PostgreSQL sessions. Driver access is exposed only through ExecutionLease so +// pool generation teardown can wait for admitted work before retiring the exact +// connection. +type ExternalSession struct { + pool *ConnectionPool + generation uint64 + + opMu sync.Mutex + + conn *stdsql.Conn + tx *stdsql.Tx + + cleanupHook func() + cleanupStarted bool + closed bool + stale bool + teardownDone chan struct{} + teardownErr error +} + +type externalSessionTeardown struct { + hook func() + conn *stdsql.Conn + tx *stdsql.Tx + quarantineErr error +} + +type externalSessionPanicError struct { + operation string + value any +} + +func (e *externalSessionPanicError) Error() string { + if e == nil { + return "external session driver operation panicked" + } + return fmt.Sprintf("external session %s panicked: %v", e.operation, e.value) +} + +// maxTransactionCompletionTombstones bounds the strong references retained by +// the late-registration bridge. Only completed entries are evicted; active +// transactions are never dropped from the admission map. +const maxTransactionCompletionTombstones = 1024 + +// transactionCompletionTombstoneTTL keeps a completed transaction addressable +// long enough for a protocol finalizer and its callback registration to cross +// threads, while ensuring an idle pool does not retain every *sql.Tx forever. +const transactionCompletionTombstoneTTL = 5 * time.Minute + +// transactionCompletion is the pool-side counterpart to backend's physical +// transaction callback registry. It lets generation teardown notify a session +// even when the pool finalizes a transaction directly, outside Session's +// adapter methods. Completed entries remain as tombstones so late registration +// receives the terminal outcome immediately. +type transactionCompletion struct { + mu sync.Mutex + callbacks []func(bool) + completed bool + dispatchComplete bool + success bool + tombstone bool + completedAt time.Time +} + +// transactionCompletionDispatch carries terminal callbacks out of the pool's +// lifecycle critical section. Pool finalizers update their identity maps while +// holding a release/session lock; invoking a callback there lets the callback +// re-enter the same pool and deadlock (and can also invert the provider barrier +// lock order). The dispatch is deliberately idempotent because a helper may +// discover the same terminal transaction through more than one cleanup path. +type transactionCompletionDispatch struct { + once sync.Once + callbacks []func(bool) + completion *transactionCompletion + pool *ConnectionPool + finished func(*stdsql.Tx) + tx *stdsql.Tx + success bool +} + +func (d *transactionCompletionDispatch) run() { + if d == nil { + return + } + panicValue, panicked := d.runCollectPanic() + if panicked { + panic(panicValue) + } +} + +// runCollectPanic always executes every terminal hook in this dispatch. A +// session callback is outside the pool's control and may panic; skipping the +// provider finished hook in that case can strand an active transaction in a +// generation transition forever. Preserve the first panic for the caller, but +// finish the required bookkeeping before it is rethrown. +func (d *transactionCompletionDispatch) runCollectPanic() (panicValue any, panicked bool) { + if d == nil { + return nil, false + } + d.once.Do(func() { + runCallbacks := func(callbacks []func(bool), success bool) { + for _, callback := range callbacks { + if callback == nil { + continue + } + value, didPanic := callTransactionCompletionHook(func() { + callback(success) + }) + if didPanic && !panicked { + panicValue, panicked = value, true + } + } + } + + runCallbacks(d.callbacks, d.success) + for d.completion != nil { + callbacks, success, complete := d.completion.takeCallbacksForDispatch() + runCallbacks(callbacks, success) + if complete { + break + } + } + if d.finished != nil { + value, didPanic := callTransactionCompletionHook(func() { + d.finished(d.tx) + }) + if didPanic && !panicked { + panicValue, panicked = value, true + } + } + if d.pool != nil && d.completion != nil { + d.pool.recordTransactionCompletion(d.tx, d.completion) + } + }) + return panicValue, panicked +} + +func callTransactionCompletionHook(hook func()) (panicValue any, panicked bool) { + defer func() { + if recovered := recover(); recovered != nil { + panicValue, panicked = recovered, true + } + }() + hook() + return nil, false +} + +type transactionCompletionBatch struct { + dispatches []*transactionCompletionDispatch + after []func() +} + +func (b *transactionCompletionBatch) add(dispatch *transactionCompletionDispatch) { + if b == nil || dispatch == nil { + return + } + b.dispatches = append(b.dispatches, dispatch) +} + +func (b *transactionCompletionBatch) addAfter(callback func()) { + if b == nil || callback == nil { + return + } + b.after = append(b.after, callback) +} + +func (b *transactionCompletionBatch) run() { + if b == nil { + return + } + dispatches := b.dispatches + b.dispatches = nil + after := b.after + b.after = nil + var firstPanic any + var panicked bool + for _, dispatch := range dispatches { + panicValue, dispatchPanicked := dispatch.runCollectPanic() + if dispatchPanicked && !panicked { + firstPanic, panicked = panicValue, true + } + } + for _, callback := range after { + if callback == nil { + continue + } + panicValue, callbackPanicked := callTransactionCompletionHook(callback) + if callbackPanicked && !panicked { + firstPanic, panicked = panicValue, true + } + } + if panicked { + panic(firstPanic) + } +} + +func (c *transactionCompletion) register(callback func(bool)) bool { + if c == nil || callback == nil { + return false + } + c.mu.Lock() + if !c.completed || !c.dispatchComplete { + c.callbacks = append(c.callbacks, callback) + c.mu.Unlock() + return true + } + success := c.success + c.mu.Unlock() + callback(success) + return true +} + +// completeDeferred records the terminal outcome without making callbacks +// inline-dispatchable. The owning public finalizer still holds pool locks at +// this point; its deferred dispatch marks dispatchComplete only after those +// locks have been released. +func (c *transactionCompletion) completeDeferred(success bool) bool { + if c == nil { + return false + } + c.mu.Lock() + defer c.mu.Unlock() + if c.completed { + return false + } + c.completed = true + c.success = success + c.completedAt = time.Now() + return true +} + +// takeCallbacksForDispatch drains callbacks registered before or during the +// lock-free dispatch. Returning complete=true is the linearization point after +// which a genuinely late registration may invoke its callback inline. +func (c *transactionCompletion) takeCallbacksForDispatch() (callbacks []func(bool), success bool, complete bool) { + if c == nil { + return nil, false, true + } + c.mu.Lock() + defer c.mu.Unlock() + success = c.success + if len(c.callbacks) > 0 { + callbacks = append([]func(bool){}, c.callbacks...) + c.callbacks = nil + return callbacks, success, false + } + c.dispatchComplete = true + return nil, success, true +} + +func (c *transactionCompletion) markTombstone() bool { + if c == nil { + return false + } + c.mu.Lock() + defer c.mu.Unlock() + if !c.completed || !c.dispatchComplete || c.tombstone { + return false + } + if c.completedAt.IsZero() { + c.completedAt = time.Now() + } + c.tombstone = true + return true +} + +func (c *transactionCompletion) isCompleted() bool { + if c == nil { + return false + } + c.mu.Lock() + completed := c.completed + c.mu.Unlock() + return completed +} + +// completionSnapshot reports a pruneable completion only after the owning +// dispatch has drained callbacks and the provider finished hook, then recorded +// the tombstone. dispatchComplete alone is not sufficient: that bit becomes +// visible just before the finished hook runs so genuinely late callbacks may +// receive the terminal result without joining the completed batch. +func (c *transactionCompletion) completionSnapshot() (bool, time.Time) { + if c == nil { + return false, time.Time{} + } + c.mu.Lock() + completed, completedAt := c.tombstone, c.completedAt + c.mu.Unlock() + return completed, completedAt +} + +// lifecycleGate coordinates ordinary pool users with generation teardown. +// +// A plain sync.RWMutex is tempting here, but it has an important shutdown +// starvation property: once Close queues for the write lock, new readers are +// blocked even when one of those readers is a release/finalizer that could +// free the resource keeping an existing reader blocked. The gate explicitly +// distinguishes those release operations. Close first marks the gate as +// shutting down, then waits for both classes of in-flight operations; release +// operations may still enter while shutdown is pending, but not after the +// exclusive teardown section starts. +type lifecycleGate struct { + mu sync.Mutex + cond *sync.Cond + readers int + releasers int + writer bool + shuttingDown bool +} + +func (g *lifecycleGate) initLocked() { + if g.cond == nil { + g.cond = sync.NewCond(&g.mu) + } +} + +// RLock enters an ordinary pool operation. New ordinary readers are rejected +// once shutdown has been announced so teardown cannot starve behind traffic. +func (g *lifecycleGate) RLock() { + g.mu.Lock() + g.initLocked() + for g.shuttingDown || g.writer { + g.cond.Wait() + } + g.readers++ + g.mu.Unlock() +} + +func (g *lifecycleGate) RUnlock() { + g.mu.Lock() + if g.readers > 0 { + g.readers-- + } + if g.cond != nil { + g.cond.Broadcast() + } + g.mu.Unlock() +} + +// ReleaseLock enters a transaction/connection finalizer. Finalizers are +// allowed to proceed while a writer is waiting, which lets them release a +// checked-out database/sql resource needed by an existing reader. They still +// wait if the exclusive teardown section has already started. +func (g *lifecycleGate) ReleaseLock() { + g.mu.Lock() + g.initLocked() + for g.writer { + g.cond.Wait() + } + g.releasers++ + g.mu.Unlock() +} + +func (g *lifecycleGate) ReleaseUnlock() { + g.mu.Lock() + if g.releasers > 0 { + g.releasers-- + } + if g.cond != nil { + g.cond.Broadcast() + } + g.mu.Unlock() +} + +// BeginShutdown prevents new ordinary readers while allowing release +// operations to drain. It is separate from Lock because callers must be able +// to close database/sql.DB before waiting for a blocked acquisition to return. +func (g *lifecycleGate) BeginShutdown() { + g.mu.Lock() + g.initLocked() + for g.shuttingDown || g.writer { + g.cond.Wait() + } + g.shuttingDown = true + g.mu.Unlock() +} + +// LockAfterShutdown waits for all operations admitted before shutdown and then +// enters the exclusive generation teardown section. BeginShutdown must have +// been called by the caller first. +func (g *lifecycleGate) LockAfterShutdown() { + g.mu.Lock() + g.initLocked() + if !g.shuttingDown { + g.shuttingDown = true + } + for g.writer || g.readers > 0 || g.releasers > 0 { + g.cond.Wait() + } + g.writer = true + g.mu.Unlock() +} + +// Lock retains the conventional API for any future exclusive callers. +func (g *lifecycleGate) Lock() { + g.BeginShutdown() + g.LockAfterShutdown() +} + +func (g *lifecycleGate) Unlock() { + g.mu.Lock() + g.writer = false + g.shuttingDown = false + if g.cond != nil { + g.cond.Broadcast() + } + g.mu.Unlock() } func NewConnectionPool(connector *duckdb.Connector, db *stdsql.DB, defaultCatalog string) *ConnectionPool { - return &ConnectionPool{ - DB: db, - connector: connector, - defaultCatalog: defaultCatalog, + p := &ConnectionPool{ + DB: db, + connector: connector, + defaultCatalog: defaultCatalog, + transactionCompletions: make(map[*stdsql.Tx]*transactionCompletion), + activeTransactions: make(map[*stdsql.Tx]struct{}), + } + p.dbPtr.Store(db) + return p +} + +// closeDatabaseForShutdown prevents a blocked database/sql acquisition from +// holding the pool's read lock forever. DB.Close is concurrency-safe and +// rejects pending connection requests; the lifecycle-locked closeLocked call +// below still performs identity-safe transaction/connection cleanup. +func (p *ConnectionPool) closeDatabaseForShutdown() error { + p.lifecycleMu.BeginShutdown() + db := p.dbPtr.Load() + if db == nil { + return nil + } + return db.Close() +} + +func (p *ConnectionPool) sessionLock(id uint32) *sync.Mutex { + for { + entry, _ := p.sessionLocks.LoadOrStore(id, &sync.Mutex{}) + if mu, ok := entry.(*sync.Mutex); ok && mu != nil { + return mu + } + // A malformed entry must not panic every lifecycle operation. Replace it + // atomically so concurrent callers still converge on one session lock. + replacement := &sync.Mutex{} + if p.sessionLocks.CompareAndSwap(id, entry, replacement) { + return replacement + } } } +func (p *ConnectionPool) lockSession(id uint32) func() { + mu := p.sessionLock(id) + mu.Lock() + return mu.Unlock +} + func (p *ConnectionPool) Connector() *duckdb.Connector { return p.connector } @@ -78,249 +578,2479 @@ func (p *ConnectionPool) initializeConnection(ctx context.Context, conn *stdsql. return initializer(ctx, conn) } -// CurrentSchema retrieves the current schema of the connection. -// Returns an empty string if the connection is not established -// or the schema cannot be retrieved. -func (p *ConnectionPool) CurrentSchema(id uint32) string { - entry, ok := p.conns.Load(id) - if !ok { - return "" +// OpenExternalSession checks out one exact physical connection and registers +// an opaque identity with the pool generation that owns it. +func (p *ConnectionPool) OpenExternalSession(ctx context.Context) (*ExternalSession, error) { + if p == nil { + return nil, fmt.Errorf("connection pool is unavailable") } - conn := entry.(*stdsql.Conn) - var schema string - if err := conn.QueryRowContext(context.Background(), "SELECT CURRENT_SCHEMA").Scan(&schema); err != nil { - logrus.WithError(err).Error("Failed to get current schema") - return "" + if ctx == nil { + ctx = context.Background() } - return schema + if err := p.connectionAdmissionError(ctx); err != nil { + return nil, err + } + + p.lifecycleMu.RLock() + defer p.lifecycleMu.RUnlock() + if err := p.connectionAdmissionError(ctx); err != nil { + return nil, err + } + + db := p.dbPtr.Load() + if db == nil { + return nil, fmt.Errorf("connection pool database is unavailable") + } + conn, err := db.Conn(ctx) + if err != nil { + return nil, err + } + if err := p.initializeExternalConnection(ctx, conn); err != nil { + return nil, errors.Join(err, callExternalSessionDriver(func() error { + return p.closePhysicalConnLocked(conn) + })) + } + + session := &ExternalSession{ + pool: p, + generation: p.externalGeneration, + conn: conn, + } + p.externalSessions.Store(session, struct{}{}) + return session, nil } -// CurrentCatalog retrieves the current catalog of the connection. Before the -// first connection, it returns the owning provider's catalog so GMS can resolve -// fully qualified names. Closed or broken connections still return empty. -func (p *ConnectionPool) CurrentCatalog(id uint32) string { - entry, ok := p.conns.Load(id) - if !ok { - if _, closed := p.closedConns.Load(id); closed { - return "" - } - return p.defaultCatalog +// externalStateErrorLocked validates a handle while the caller holds both the +// pool lifecycle lock and opMu. A stale generation is reported before an +// explicit close so Reset remains distinguishable even if Close is later called +// on the old handle. +func (s *ExternalSession) externalStateErrorLocked() error { + if s == nil || s.pool == nil { + return ErrExternalSessionClosed } - conn := entry.(*stdsql.Conn) - var catalog string - if err := conn.QueryRowContext(context.Background(), "SELECT CURRENT_CATALOG").Scan(&catalog); err != nil { - logrus.WithError(err).Error("Failed to get current catalog") - return "" + if s.stale || s.generation != s.pool.externalGeneration { + return ErrExternalSessionStale } - return catalog + if s.closed || s.conn == nil { + return ErrExternalSessionClosed + } + return nil } -func (p *ConnectionPool) GetConn(ctx context.Context, id uint32) (*stdsql.Conn, error) { +// ExecutionLease returns the exact connection and transaction owned by this +// handle. The release function is idempotent and must be held through all +// synchronous driver work and result consumption. +func (s *ExternalSession) ExecutionLease( + ctx context.Context, +) (*stdsql.Conn, *stdsql.Tx, func(), error) { + if s == nil || s.pool == nil { + return nil, nil, func() {}, ErrExternalSessionClosed + } if ctx == nil { ctx = context.Background() } - var conn *stdsql.Conn - entry, ok := p.conns.Load(id) - if !ok { - c, err := p.DB.Conn(ctx) - if err != nil { - return nil, err - } - if _, transactionActive := p.txns.Load(id); !transactionActive { - if err := p.initializeConnection(ctx, c); err != nil { - _ = c.Close() - return nil, err - } - } - if err := p.registerMySQLUDFs(c); err != nil { - _ = c.Close() - return nil, err - } - p.closedConns.Delete(id) - p.conns.Store(id, c) - conn = c - } else { - conn = entry.(*stdsql.Conn) - // A session transaction owns this connection's transaction-scoped - // state. Re-running the initializer here would execute LOAD/CREATE - // SECRET inside that transaction and could alter or roll back with user - // work. GetTxn initializes before BeginTx; keep the state stable until - // the transaction is closed. - if _, transactionActive := p.txns.Load(id); !transactionActive { - if err := p.initializeConnection(ctx, conn); err != nil { - // Do not leave a failed or partially initialized connection available - // to a later request. CompareAndDelete avoids removing a replacement - // installed by a concurrent recovery path. - p.conns.CompareAndDelete(id, conn) - p.closedConns.Store(id, struct{}{}) - _ = conn.Close() - return nil, err - } - } + p := s.pool + var completions transactionCompletionBatch + defer completions.run() + if err := p.connectionAdmissionError(ctx); err != nil { + return nil, nil, func() {}, err } - return conn, nil -} -// DuckDB stores registered scalar UDFs in the database catalog shared by all connections. -func (p *ConnectionPool) registerMySQLUDFs(conn *stdsql.Conn) error { - p.registerMySQLUDFsOnce.Do(func() { - if err := registerMySQLRand(conn); err != nil { - p.registerMySQLUDFsErr = fmt.Errorf("register mysql_rand: %w", err) - return - } - if err := registerMySQLRandomBytes(conn); err != nil { - p.registerMySQLUDFsErr = fmt.Errorf("register mysql_random_bytes: %w", err) - return - } - if err := registerMySQLStringToVector(conn); err != nil { - p.registerMySQLUDFsErr = fmt.Errorf("register string_to_vector: %w", err) + p.lifecycleMu.RLock() + s.opMu.Lock() + var releaseOnce sync.Once + release := func() { + releaseOnce.Do(func() { + s.opMu.Unlock() + p.lifecycleMu.RUnlock() + }) + } + if err := p.connectionAdmissionError(ctx); err != nil { + release() + return nil, nil, func() {}, err + } + if err := s.externalStateErrorLocked(); err != nil { + release() + return nil, nil, func() {}, err + } + if s.tx == nil { + if err := p.initializeExternalConnection(ctx, s.conn); err != nil { + teardown, owner := s.beginTeardownLocked(false) + if owner { + err = s.finishTeardownLocked(teardown, err, &completions) + } else { + err = errors.Join(err, s.waitForTeardownLocked()) + } + release() + return nil, nil, func() {}, err } - }) - return p.registerMySQLUDFsErr + } + return s.conn, s.tx, release, nil } -func (p *ConnectionPool) GetConnForSchema(ctx context.Context, id uint32, schemaName string) (*stdsql.Conn, error) { +// BeginTx opens a transaction on the handle's exact connection. Provider +// transaction admission is reserved without holding a lifecycle reader so a +// generation transition can make progress while admission waits. +func (s *ExternalSession) BeginTx(ctx context.Context, opts *stdsql.TxOptions) (*stdsql.Tx, error) { + if s == nil || s.pool == nil { + return nil, ErrExternalSessionClosed + } if ctx == nil { ctx = context.Background() } - conn, err := p.GetConn(ctx, id) - if err != nil { + p := s.pool + var completions transactionCompletionBatch + defer completions.run() + if err := p.connectionAdmissionError(ctx); err != nil { return nil, err } - if schemaName != "" { - // Schema selection is session state, but it should retain the origin - // value while avoiding cancellation of a request that is already being - // serviced. - schemaCtx := context.WithoutCancel(ctx) - var currentSchema string - if err := conn.QueryRowContext(schemaCtx, "SELECT CURRENT_SCHEMA").Scan(¤tSchema); err != nil { - logrus.WithError(err).Error("Failed to get current schema") - return nil, err - } else if currentSchema != schemaName { - if _, err := conn.ExecContext(schemaCtx, "USE "+FullSchemaName(p.CurrentCatalog(id), schemaName)); err != nil { - if IsDuckDBSetSchemaNotFoundError(err) { - return nil, sql.ErrDatabaseNotFound.New(schemaName) - } - logrus.WithField("schema", schemaName).WithError(err).Error("Failed to switch schema") - return nil, err - } - } + // Preserve the existing-transaction fast path without reserving another + // provider transaction slot. + p.lifecycleMu.RLock() + s.opMu.Lock() + if err := p.connectionAdmissionError(ctx); err != nil { + s.opMu.Unlock() + p.lifecycleMu.RUnlock() + return nil, err } + if err := s.externalStateErrorLocked(); err != nil { + s.opMu.Unlock() + p.lifecycleMu.RUnlock() + return nil, err + } + if s.tx != nil { + tx := s.tx + s.opMu.Unlock() + p.lifecycleMu.RUnlock() + return tx, nil + } + s.opMu.Unlock() + p.lifecycleMu.RUnlock() - return conn, nil -} - -func (p *ConnectionPool) CloseConn(id uint32) error { - defer p.conns.Delete(id) - defer p.closedConns.Store(id, struct{}{}) - var lastErr error - if entry, ok := p.txns.LoadAndDelete(id); ok { - if err := entry.(*stdsql.Tx).Rollback(); err != nil && - !errors.Is(err, stdsql.ErrTxDone) && - !strings.Contains(err.Error(), "no transaction is active") { - logrus.WithError(err).Warn("Failed to rollback transaction") - lastErr = err - } + if err := p.transactionAdmissionError(); err != nil { + return nil, err } - entry, ok := p.conns.Load(id) - if ok { - conn := entry.(*stdsql.Conn) - if err := conn.Raw(func(driverConn any) error { - // When driver.ErrBadConn is returned here, - // the connection will not be put back into - // the pool and will be closed instead. - return driver.ErrBadConn - }); err != nil && !errors.Is(err, driver.ErrBadConn) { - logrus.WithError(err).Warn("Failed to close connection during Raw function call") - return errors.Join(lastErr, err) - } - if err := conn.Close(); err != nil && !errors.Is(err, stdsql.ErrConnDone) { - logrus.WithError(err).Warn("Failed to close connection") - return errors.Join(lastErr, err) + completeAdmission, err := p.beginTransactionAdmissionWithError() + if err != nil { + return nil, err + } + admissionPending := true + finishAdmission := func(tx *stdsql.Tx) { + if !admissionPending { + return + } + admissionPending = false + completeAdmission(tx) + } + + p.lifecycleMu.RLock() + defer p.lifecycleMu.RUnlock() + s.opMu.Lock() + defer s.opMu.Unlock() + if err := p.connectionAdmissionError(ctx); err != nil { + finishAdmission(nil) + return nil, err + } + if err := p.transactionAdmissionError(); err != nil { + finishAdmission(nil) + return nil, err + } + if err := s.externalStateErrorLocked(); err != nil { + finishAdmission(nil) + return nil, err + } + if s.tx != nil { + finishAdmission(nil) + return s.tx, nil + } + if err := p.initializeExternalConnection(ctx, s.conn); err != nil { + finishAdmission(nil) + teardown, owner := s.beginTeardownLocked(false) + if owner { + teardownErr := s.finishTeardownLocked(teardown, err, &completions) + if teardownErr != nil { + err = teardownErr + } + } else { + err = errors.Join(err, s.waitForTeardownLocked()) + } + return nil, err + } + tx, err := callExternalSessionBeginTx(s.conn, context.WithoutCancel(ctx), opts) + if err != nil { + if externalSessionDriverPanicked(err) { + p.recordExternalSessionPanic(err) + } + finishAdmission(nil) + if externalSessionDriverPanicked(err) { + teardown, owner := s.beginTeardownLocked(false) + if owner { + teardown.quarantineErr = err + err = s.finishTeardownLocked(teardown, err, &completions) + } else { + err = errors.Join(err, s.waitForTeardownLocked()) + } + } + return nil, err + } + s.tx = tx + p.markTransactionActive(tx) + finishAdmission(tx) + return tx, nil +} + +func callExternalSessionCleanupHook(hook func()) (err error) { + if hook == nil { + return nil + } + defer func() { + if recovered := recover(); recovered != nil { + err = fmt.Errorf("external session cleanup hook panicked: %v", recovered) + } + }() + hook() + return nil +} + +func callExternalSessionDriver(fn func() error) (err error) { + if fn == nil { + return nil + } + defer func() { + if recovered := recover(); recovered != nil { + err = fmt.Errorf("external session driver operation panicked: %v", recovered) + } + }() + return fn() +} + +func callExternalSessionTransactionDriver(operation string, fn func() error) (err error) { + if fn == nil { + return nil + } + defer func() { + if recovered := recover(); recovered != nil { + err = &externalSessionPanicError{operation: operation, value: recovered} + } + }() + return fn() +} + +func externalSessionDriverPanicked(err error) bool { + var panicErr *externalSessionPanicError + return errors.As(err, &panicErr) +} + +func callExternalSessionBeginTx( + conn *stdsql.Conn, + ctx context.Context, + opts *stdsql.TxOptions, +) (tx *stdsql.Tx, err error) { + defer func() { + if recovered := recover(); recovered != nil { + err = &externalSessionPanicError{operation: "begin transaction", value: recovered} + tx = nil + } + }() + if conn == nil { + return nil, fmt.Errorf("external session connection is unavailable") + } + return conn.BeginTx(ctx, opts) +} + +func (p *ConnectionPool) initializeExternalConnection(ctx context.Context, conn *stdsql.Conn) error { + if err := callExternalSessionDriver(func() error { + return p.initializeConnection(ctx, conn) + }); err != nil { + return err + } + return callExternalSessionDriver(func() error { + return p.registerMySQLUDFs(conn) + }) +} + +func (s *ExternalSession) beginTeardownLocked(stale bool) (*externalSessionTeardown, bool) { + if stale { + s.stale = true + } else { + s.closed = true + } + if s.pool != nil { + s.pool.externalSessions.Delete(s) + } + if s.teardownDone != nil { + return nil, false + } + s.teardownDone = make(chan struct{}) + s.cleanupStarted = true + teardown := &externalSessionTeardown{ + hook: s.cleanupHook, + conn: s.conn, + tx: s.tx, + } + s.cleanupHook = nil + s.conn = nil + s.tx = nil + return teardown, true +} + +// runCleanupOutsideOperationLock invokes the installed hook without opMu while +// retaining lifecycle admission. The caller holds opMu on entry and again on +// return. SetCleanupHook serializes late registration on opMu, so a hook that +// wins the race after invalidation is invoked inline instead of being lost. +func (s *ExternalSession) runCleanupOutsideOperationLock(teardown *externalSessionTeardown) error { + if teardown == nil || teardown.hook == nil { + return nil + } + s.opMu.Unlock() + err := callExternalSessionCleanupHook(teardown.hook) + s.opMu.Lock() + return err +} + +func (s *ExternalSession) closeResourcesLocked( + teardown *externalSessionTeardown, + completions *transactionCompletionBatch, +) error { + if s == nil || s.pool == nil || teardown == nil { + return nil + } + var lastErr error + quarantineErr := teardown.quarantineErr + if teardown.tx != nil { + tx := teardown.tx + rollbackErr := callExternalSessionTransactionDriver("rollback transaction", tx.Rollback) + if rollbackErr != nil && !transactionInactiveError(rollbackErr) { + lastErr = errors.Join(lastErr, rollbackErr) + } + if externalSessionDriverPanicked(rollbackErr) { + quarantineErr = errors.Join(quarantineErr, rollbackErr) + s.pool.recordExternalSessionPanic(rollbackErr) + } + completions.add(s.pool.finishTransactionWithOutcomeDeferred(tx, false)) + } + if teardown.conn != nil { + if quarantineErr != nil { + lastErr = errors.Join( + lastErr, + s.pool.quarantineExternalConnectionAfterPanic(teardown.conn, quarantineErr), + ) + } else { + lastErr = errors.Join(lastErr, callExternalSessionDriver(func() error { + return s.pool.closePhysicalConnLocked(teardown.conn) + })) } } return lastErr } -func (p *ConnectionPool) GetTxn(ctx context.Context, id uint32, schemaName string, options *stdsql.TxOptions) (*stdsql.Tx, error) { - if ctx == nil { - ctx = context.Background() - } - var tx *stdsql.Tx - entry, ok := p.txns.Load(id) - if !ok { - conn, err := p.GetConnForSchema(ctx, id, schemaName) +// quarantineExternalConnectionAfterPanic discards the physical driver +// connection whose BeginTx, Commit, or Rollback call panicked. Those +// database/sql paths transfer a Conn read lock to the transaction without a +// panic-side release; ordinary Conn.Raw/Close retirement then blocks forever +// trying to acquire the corresponding write lock. Closing driver.Conn inside +// the Raw callback avoids that write lock. The callback must always return nil: +// returning ErrBadConn (or propagating a panic) would make database/sql enter +// the same blocked retirement path after the callback. +// +// The database/sql wrapper remains permanently checked out and detached, so it +// is never returned to the reusable pool. The generation is poisoned and must +// be rebuilt by process restart rather than Reset. +func (p *ConnectionPool) quarantineExternalConnectionAfterPanic( + conn *stdsql.Conn, + panicErr error, +) error { + if p == nil || conn == nil { + return nil + } + p.recordExternalSessionPanic(panicErr) + var physicalCloseErr error + rawErr := callExternalSessionDriver(func() error { + return conn.Raw(func(raw any) error { + physical, ok := raw.(driver.Conn) + if !ok || physical == nil { + physicalCloseErr = fmt.Errorf( + "external session driver connection has unexpected type %T", + raw, + ) + return nil + } + physicalCloseErr = callExternalSessionDriver(physical.Close) + return nil + }) + }) + err := errors.Join(rawErr, physicalCloseErr) + if err != nil && !errors.Is(err, stdsql.ErrConnDone) { + logrus.WithError(err).Warn("Failed to quarantine external session connection") + } + return err +} + +func (p *ConnectionPool) recordExternalSessionPanic(err error) { + if p == nil || err == nil { + return + } + p.externalPanicMu.Lock() + if p.externalPanicErr == nil { + p.externalPanicErr = errors.Join(ErrExternalSessionGenerationPoisoned, err) + } + p.externalPanicMu.Unlock() +} + +func (p *ConnectionPool) externalSessionPanicError() error { + if p == nil { + return nil + } + p.externalPanicMu.Lock() + err := p.externalPanicErr + p.externalPanicMu.Unlock() + return err +} + +func (s *ExternalSession) completeTeardownLocked(err error) { + if s == nil || s.teardownDone == nil { + return + } + s.teardownErr = err + close(s.teardownDone) +} + +func (s *ExternalSession) completeTeardown(err error) { + if s == nil { + return + } + s.opMu.Lock() + s.completeTeardownLocked(err) + s.opMu.Unlock() +} + +func (s *ExternalSession) waitForTeardownLocked() error { + done := s.teardownDone + s.opMu.Unlock() + <-done + s.opMu.Lock() + return s.teardownErr +} + +func (s *ExternalSession) finishTeardownLocked( + teardown *externalSessionTeardown, + baseErr error, + completions *transactionCompletionBatch, +) (err error) { + defer func() { + if recovered := recover(); recovered != nil { + err = errors.Join(err, fmt.Errorf("external session teardown panicked: %v", recovered)) + } + if completions == nil { + s.completeTeardownLocked(err) + } else { + teardownErr := err + completions.addAfter(func() { s.completeTeardown(teardownErr) }) + } + }() + err = errors.Join(baseErr, s.runCleanupOutsideOperationLock(teardown)) + err = errors.Join(err, s.closeResourcesLocked(teardown, completions)) + return err +} + +// SetCleanupHook installs the protocol-owned prepared-state cleanup for this +// handle. Only one hook may be installed. If teardown won the race, the hook is +// invoked inline so protocol state is never left behind; the returned state +// error still tells the caller that the handle can no longer be used. +// +// Cleanup hooks must not re-enter this ExternalSession. Both the normal +// teardown path and an already-invalid late registration invoke hooks outside +// opMu; a late registration waits for the teardown owner first so hooks cannot +// overlap or race connection retirement. +func (s *ExternalSession) SetCleanupHook(hook func()) error { + if hook == nil { + return nil + } + if s == nil || s.pool == nil { + return errors.Join(ErrExternalSessionClosed, callExternalSessionCleanupHook(hook)) + } + + p := s.pool + // Registration is release-side work: while shutdown is waiting, admitting it + // lets the writer observe the hook and run it before connection retirement. + // Once the writer is active, ReleaseLock waits and this method instead takes + // the already-stale inline-cleanup path below. + p.lifecycleMu.ReleaseLock() + s.opMu.Lock() + var stateErr error + if s.stale || s.generation != p.externalGeneration { + stateErr = ErrExternalSessionStale + } else if s.closed || s.cleanupStarted { + stateErr = ErrExternalSessionClosed + } else if s.cleanupHook != nil { + s.opMu.Unlock() + p.lifecycleMu.ReleaseUnlock() + return fmt.Errorf("external session cleanup hook is already set") + } else { + s.cleanupHook = hook + s.opMu.Unlock() + p.lifecycleMu.ReleaseUnlock() + return nil + } + done := s.teardownDone + s.opMu.Unlock() + p.lifecycleMu.ReleaseUnlock() + if done != nil { + <-done + } + return errors.Join(stateErr, callExternalSessionCleanupHook(hook)) +} + +// Commit commits only the exact active transaction supplied by the caller. +// A stale or mismatched pointer is rejected before any driver operation. +func (s *ExternalSession) Commit(expected *stdsql.Tx) error { + if s == nil || s.pool == nil { + return ErrExternalSessionClosed + } + p := s.pool + var completions transactionCompletionBatch + defer completions.run() + p.lifecycleMu.ReleaseLock() + defer p.lifecycleMu.ReleaseUnlock() + s.opMu.Lock() + defer s.opMu.Unlock() + + if err := s.externalStateErrorLocked(); err != nil { + return err + } + if expected == nil || s.tx == nil || s.tx != expected { + return ErrExternalSessionTransactionMismatch + } + err := callExternalSessionTransactionDriver("commit transaction", expected.Commit) + if externalSessionDriverPanicked(err) { + p.recordExternalSessionPanic(err) + } + s.tx = nil + completions.add(p.finishTransactionWithOutcomeDeferred(expected, err == nil)) + if err == nil { + return nil + } + + teardown, owner := s.beginTeardownLocked(false) + if !owner { + return errors.Join(err, s.waitForTeardownLocked()) + } + if externalSessionDriverPanicked(err) { + teardown.quarantineErr = err + } + return s.finishTeardownLocked(teardown, err, &completions) +} + +// Rollback rolls back only the exact active transaction supplied by the +// caller. cleanup runs on the same physical connection after the driver proves +// the transaction inactive and before another operation can enter the handle. +// The caller must reserve any provider rollback-maintenance barrier first. +func (s *ExternalSession) Rollback( + expected *stdsql.Tx, + cleanup func(*stdsql.Conn) error, +) error { + if s == nil || s.pool == nil { + return ErrExternalSessionClosed + } + p := s.pool + var completions transactionCompletionBatch + defer completions.run() + p.lifecycleMu.ReleaseLock() + defer p.lifecycleMu.ReleaseUnlock() + s.opMu.Lock() + defer s.opMu.Unlock() + + if err := s.externalStateErrorLocked(); err != nil { + return err + } + if expected == nil || s.tx == nil || s.tx != expected { + return ErrExternalSessionTransactionMismatch + } + err := callExternalSessionTransactionDriver("rollback transaction", expected.Rollback) + if externalSessionDriverPanicked(err) { + p.recordExternalSessionPanic(err) + } + inactive := transactionInactiveError(err) + s.tx = nil + completions.add(p.finishTransactionWithOutcomeDeferred(expected, false)) + if !inactive { + teardown, owner := s.beginTeardownLocked(false) + if !owner { + return errors.Join(err, s.waitForTeardownLocked()) + } + if externalSessionDriverPanicked(err) { + teardown.quarantineErr = err + } + return s.finishTeardownLocked(teardown, err, &completions) + } + if cleanup != nil { + err = errors.Join(err, callExternalSessionDriver(func() error { return cleanup(s.conn) })) + } + return err +} + +// Close releases an external handle. It is idempotent and uses release-side +// lifecycle admission so callers can close handles while pool shutdown waits +// for older operations to drain. +func (s *ExternalSession) Close() error { + if s == nil || s.pool == nil { + return nil + } + p := s.pool + var completions transactionCompletionBatch + defer completions.run() + p.lifecycleMu.ReleaseLock() + defer p.lifecycleMu.ReleaseUnlock() + s.opMu.Lock() + defer s.opMu.Unlock() + + if s.teardownDone != nil { + return s.waitForTeardownLocked() + } + teardown, _ := s.beginTeardownLocked(false) + return s.finishTeardownLocked(teardown, nil, &completions) +} + +// SetTransactionLifecycleHooks installs optional hooks for providers that need +// to coordinate transaction admission with shared catalog maintenance. The +// admission hook is called only when a new transaction is about to be opened; +// it returns a completion function that must be called with the newly opened +// transaction, or nil when opening failed. The finished hook is called whenever +// a transaction mapping is removed by any pool lifecycle path. +// +// Hooks are deliberately optional so standalone test pools and callers that do +// not own a provider retain the original pool behavior. +func (p *ConnectionPool) SetTransactionLifecycleHooks( + admission func() func(*stdsql.Tx), + finished func(*stdsql.Tx), +) { + p.transactionHooksMu.Lock() + p.transactionAdmission = admission + p.transactionAdmissionWithError = nil + p.transactionFinished = finished + p.transactionHooksMu.Unlock() +} + +// SetTransactionLifecycleHooksWithError installs the provider-aware +// transaction admission hook. A non-nil error prevents the pool from calling +// BeginTx; the returned completion callback is invoked only for an admission +// that actually reserved a transaction slot. The legacy setter above remains +// available for test pools and downstream embedders. +func (p *ConnectionPool) SetTransactionLifecycleHooksWithError( + admission func() (func(*stdsql.Tx), error), + finished func(*stdsql.Tx), +) { + if p == nil { + return + } + p.transactionHooksMu.Lock() + p.transactionAdmission = nil + p.transactionAdmissionWithError = admission + p.transactionFinished = finished + p.transactionHooksMu.Unlock() +} + +// SetTransactionAdmissionGuard installs an optional provider-owned guard for +// new physical transaction admission. Existing transactions still use the +// ordinary fast path so they can be finalized or repaired after a generation +// failure; only attempts to open a new transaction are rejected. +func (p *ConnectionPool) SetTransactionAdmissionGuard(guard func() error) { + if p == nil { + return + } + p.transactionHooksMu.Lock() + p.transactionAdmissionGuard = guard + p.transactionHooksMu.Unlock() +} + +// SetConnectionAdmissionGuard installs the provider-owned guard used by +// ordinary connection, execution-snapshot, and metadata entry points. The +// guard runs once before lifecycle admission and once after the pool/session +// locks are held. The second check closes the gap where shutdown completed and +// the lifecycle gate reopened while the provider transition was still active. +func (p *ConnectionPool) SetConnectionAdmissionGuard(guard func(context.Context) error) { + if p == nil { + return + } + p.transactionHooksMu.Lock() + p.connectionAdmissionGuard = guard + p.transactionHooksMu.Unlock() +} + +func (p *ConnectionPool) transactionAdmissionError() error { + if p == nil { + return nil + } + if err := p.externalSessionPanicError(); err != nil { + return err + } + p.transactionHooksMu.RLock() + guard := p.transactionAdmissionGuard + p.transactionHooksMu.RUnlock() + if guard == nil { + return nil + } + return guard() +} + +func (p *ConnectionPool) connectionAdmissionError(ctx context.Context) error { + if p == nil { + return nil + } + if err := p.externalSessionPanicError(); err != nil { + return err + } + p.transactionHooksMu.RLock() + guard := p.connectionAdmissionGuard + p.transactionHooksMu.RUnlock() + if guard == nil { + return nil + } + return guard(ctx) +} + +func (p *ConnectionPool) beginTransactionAdmission() func(*stdsql.Tx) { + completion, _ := p.beginTransactionAdmissionWithError() + return completion +} + +func (p *ConnectionPool) beginTransactionAdmissionWithError() (func(*stdsql.Tx), error) { + p.transactionHooksMu.RLock() + admissionWithError := p.transactionAdmissionWithError + admission := p.transactionAdmission + p.transactionHooksMu.RUnlock() + if admissionWithError != nil { + completion, err := admissionWithError() + if err != nil { + return func(*stdsql.Tx) {}, err + } + if completion == nil { + return func(*stdsql.Tx) {}, nil + } + return completion, nil + } + if admission == nil { + return func(*stdsql.Tx) {}, nil + } + completion := admission() + if completion == nil { + return func(*stdsql.Tx) {}, nil + } + return completion, nil +} + +func (p *ConnectionPool) markTransactionActive(tx *stdsql.Tx) { + if p == nil || tx == nil { + return + } + p.transactionCompletionMu.Lock() + if p.activeTransactions == nil { + p.activeTransactions = make(map[*stdsql.Tx]struct{}) + } + p.activeTransactions[tx] = struct{}{} + p.transactionCompletionMu.Unlock() +} + +func (p *ConnectionPool) transactionIsMapped(tx *stdsql.Tx) bool { + if p == nil || tx == nil { + return false + } + if _, active := p.activeTransactions[tx]; active { + return true + } + found := false + p.txns.Range(func(_, value any) bool { + if mapped, ok := value.(*stdsql.Tx); ok && mapped == tx { + found = true + return false + } + return true + }) + return found +} + +// pruneTransactionCompletionsLocked removes only completed tombstones that +// have exceeded the retention window or the FIFO bound. Callers hold +// transactionCompletionMu. Pending callbacks and active transactions are never +// removed by this method. +func (p *ConnectionPool) pruneTransactionCompletionsLocked(now time.Time) { + if p == nil { + return + } + if len(p.transactionCompletions) == 0 { + p.transactionCompletionOrder = nil + return + } + for tx, state := range p.transactionCompletions { + completed, completedAt := state.completionSnapshot() + if completed && !completedAt.IsZero() && now.Sub(completedAt) >= transactionCompletionTombstoneTTL { + delete(p.transactionCompletions, tx) + } + } + + // Rebuild the order from live map entries, removing stale/duplicate keys + // left behind by an earlier bounded prune. Completed entries created by a + // concurrent finalizer are appended below if they have not reached the FIFO. + order := make([]*stdsql.Tx, 0, len(p.transactionCompletionOrder)) + seen := make(map[*stdsql.Tx]struct{}, len(p.transactionCompletionOrder)) + for _, tx := range p.transactionCompletionOrder { + if _, duplicate := seen[tx]; duplicate { + continue + } + state, present := p.transactionCompletions[tx] + if !present { + continue + } + completed, _ := state.completionSnapshot() + if !completed { + // This should not normally occur because entries enter the FIFO only + // after completion, but retaining a pending map entry is safer than + // deleting it if a future caller violates that assumption. + continue + } + seen[tx] = struct{}{} + order = append(order, tx) + } + for tx, state := range p.transactionCompletions { + completed, completedAt := state.completionSnapshot() + if !completed { + continue + } + if !completedAt.IsZero() && now.Sub(completedAt) >= transactionCompletionTombstoneTTL { + delete(p.transactionCompletions, tx) + continue + } + if _, present := seen[tx]; !present { + seen[tx] = struct{}{} + order = append(order, tx) + } + } + for len(order) > maxTransactionCompletionTombstones { + oldest := order[0] + order = order[1:] + if state := p.transactionCompletions[oldest]; state != nil { + completed, _ := state.completionSnapshot() + if completed { + delete(p.transactionCompletions, oldest) + } + } + } + p.transactionCompletionOrder = order +} + +// recordTransactionCompletion keeps a bounded FIFO of completed states. The +// map lock is acquired before the state lock so it has the same lock ordering +// as registration/pruning and cannot deadlock with a concurrent late register. +func (p *ConnectionPool) recordTransactionCompletion(tx *stdsql.Tx, state *transactionCompletion) { + if p == nil || tx == nil || state == nil { + return + } + p.transactionCompletionMu.Lock() + defer p.transactionCompletionMu.Unlock() + if p.transactionCompletions == nil { + return + } + if current := p.transactionCompletions[tx]; current != state { + return + } + if !state.markTombstone() { + return + } + p.transactionCompletionOrder = append(p.transactionCompletionOrder, tx) + p.pruneTransactionCompletionsLocked(time.Now()) +} + +func (p *ConnectionPool) clearCompletedTransactionCompletions() { + if p == nil { + return + } + p.transactionCompletionMu.Lock() + for tx, state := range p.transactionCompletions { + if completed, _ := state.completionSnapshot(); completed { + delete(p.transactionCompletions, tx) + } + } + p.transactionCompletionOrder = nil + p.transactionCompletionMu.Unlock() +} + +// finishTransactionWithOutcomeDeferred records the terminal outcome and +// returns a dispatch token. The caller must run the token after releasing any +// pool lifecycle and session locks. Provider teardown can still own its outer +// generation locks while the pool drains this token, so callbacks are limited +// to terminal notification and resource release; they must not re-enter +// provider DDL, cleanup, Close, or Restart. +func (p *ConnectionPool) finishTransactionWithOutcomeDeferred(tx *stdsql.Tx, success bool) *transactionCompletionDispatch { + if tx == nil { + return nil + } + p.transactionCompletionMu.Lock() + p.pruneTransactionCompletionsLocked(time.Now()) + if p.transactionCompletions == nil { + p.transactionCompletions = make(map[*stdsql.Tx]*transactionCompletion) + } + if p.activeTransactions != nil { + delete(p.activeTransactions, tx) + } + completion := p.transactionCompletions[tx] + if completion == nil { + completion = &transactionCompletion{} + p.transactionCompletions[tx] = completion + } + p.transactionCompletionMu.Unlock() + ownsCompletionDispatch := completion.completeDeferred(success) + if !ownsCompletionDispatch { + return nil + } + + p.transactionHooksMu.RLock() + finished := p.transactionFinished + p.transactionHooksMu.RUnlock() + return &transactionCompletionDispatch{ + completion: completion, + pool: p, + finished: finished, + tx: tx, + success: success, + } +} + +// RegisterTransactionCompletion attaches a callback to a raw transaction's +// terminal lifecycle. It is intentionally small and optional; backend.Session +// uses it to bridge provider/pool teardown into its protocol callback registry. +// +// A recent terminal outcome is delivered inline before this method returns. +// Ordinary finalizers also deliver callbacks synchronously, but only after +// releasing the pool lifecycle and session locks, so callbacks may re-enter +// pool release/read APIs. Callbacks must remain terminal notification/resource +// release hooks and must not re-enter provider DDL, cleanup, Close, or Restart; +// a provider generation transition may still own those outer locks. +func (p *ConnectionPool) RegisterTransactionCompletion(tx *stdsql.Tx, callback func(bool)) bool { + if p == nil || tx == nil || callback == nil { + return false + } + p.transactionCompletionMu.Lock() + // Prune completed tombstones before admitting a new registration. If the + // transaction is no longer active and its tombstone has aged out, fail closed + // and let the caller release its resource through the fallback path. + p.pruneTransactionCompletionsLocked(time.Now()) + if p.transactionCompletions == nil { + p.transactionCompletions = make(map[*stdsql.Tx]*transactionCompletion) + } + completion := p.transactionCompletions[tx] + if completion == nil { + if !p.transactionIsMapped(tx) { + p.transactionCompletionMu.Unlock() + callback(false) + return false + } + completion = &transactionCompletion{} + p.transactionCompletions[tx] = completion + } + p.transactionCompletionMu.Unlock() + return completion.register(callback) +} + +// CurrentSchema retrieves the current schema of the connection. +// Returns an empty string if the connection is not established +// or the schema cannot be retrieved. +func (p *ConnectionPool) CurrentSchema(id uint32) string { + if err := p.connectionAdmissionError(context.Background()); err != nil { + return "" + } + p.lifecycleMu.RLock() + defer p.lifecycleMu.RUnlock() + unlock := p.lockSession(id) + defer unlock() + if err := p.connectionAdmissionError(context.Background()); err != nil { + return "" + } + return p.currentSchemaLocked(id) +} + +func (p *ConnectionPool) currentSchemaLocked(id uint32) string { + if tx, active, err := p.transactionEntryLocked(id); active { + if err != nil { + return "" + } + // An active transaction is usable only with the owner captured at + // BeginTx. Do not fall back to the generic session connection when that + // owner mapping is absent or malformed. + if _, err := p.activeTxnOwnerLocked(id); err != nil { + return "" + } + var schema string + if err := tx.QueryRowContext(context.Background(), "SELECT CURRENT_SCHEMA").Scan(&schema); err != nil { + logrus.WithError(err).Error("Failed to get current schema") + return "" + } + return schema + } + conn, present, err := p.connectionEntryLocked(id) + if !present || err != nil { + return "" + } + var schema string + if err := conn.QueryRowContext(context.Background(), "SELECT CURRENT_SCHEMA").Scan(&schema); err != nil { + logrus.WithError(err).Error("Failed to get current schema") + return "" + } + return schema +} + +// CurrentCatalog retrieves the current catalog of the connection. Before the +// first connection, it returns the owning provider's catalog so GMS can resolve +// fully qualified names. Closed or broken connections still return empty. +func (p *ConnectionPool) CurrentCatalog(id uint32) string { + if err := p.connectionAdmissionError(context.Background()); err != nil { + return "" + } + p.lifecycleMu.RLock() + defer p.lifecycleMu.RUnlock() + unlock := p.lockSession(id) + defer unlock() + if err := p.connectionAdmissionError(context.Background()); err != nil { + return "" + } + return p.currentCatalogLocked(id) +} + +func (p *ConnectionPool) currentCatalogLocked(id uint32) string { + if tx, active, err := p.transactionEntryLocked(id); active { + if err != nil { + return "" + } + // See currentSchemaLocked: never use p.conns as an owner substitute for + // an active transaction. + if _, err := p.activeTxnOwnerLocked(id); err != nil { + return "" + } + var catalog string + if err := tx.QueryRowContext(context.Background(), "SELECT CURRENT_CATALOG").Scan(&catalog); err != nil { + logrus.WithError(err).Error("Failed to get current catalog") + return "" + } + return catalog + } + conn, present, err := p.connectionEntryLocked(id) + if !present { + if _, closed := p.closedConns.Load(id); closed { + return "" + } + return p.defaultCatalog + } + if err != nil { + return "" + } + var catalog string + if err := conn.QueryRowContext(context.Background(), "SELECT CURRENT_CATALOG").Scan(&catalog); err != nil { + logrus.WithError(err).Error("Failed to get current catalog") + return "" + } + return catalog +} + +func (p *ConnectionPool) GetConn(ctx context.Context, id uint32) (*stdsql.Conn, error) { + conn, _, err := p.GetConnWithBinding(ctx, id) + return conn, err +} + +// GetConnWithBinding acquires a catalog connection and returns the transaction +// binding observed in the same lifecycle/session critical section. Callers that +// need to remember an identity for a later close must use this method instead +// of doing a second GetTxnBinding lookup after GetConn returns. +func (p *ConnectionPool) GetConnWithBinding(ctx context.Context, id uint32) (*stdsql.Conn, *stdsql.Tx, error) { + return p.getConnWithBinding(ctx, id, "") +} + +func (p *ConnectionPool) getConnWithBinding(ctx context.Context, id uint32, schemaName string) (*stdsql.Conn, *stdsql.Tx, error) { + if ctx == nil { + ctx = context.Background() + } + if err := p.connectionAdmissionError(ctx); err != nil { + return nil, nil, err + } + p.lifecycleMu.RLock() + defer p.lifecycleMu.RUnlock() + unlock := p.lockSession(id) + defer unlock() + if err := p.connectionAdmissionError(ctx); err != nil { + return nil, nil, err + } + var ( + conn *stdsql.Conn + err error + ) + if schemaName == "" { + conn, err = p.getConnLocked(ctx, id) + } else { + conn, err = p.getConnForSchemaLocked(ctx, id, schemaName) + } + if err != nil { + return nil, nil, err + } + boundConn, tx := p.txnBindingLocked(id) + if tx != nil && boundConn != conn { + // The helper should never return a transaction with a different owner. A + // mismatch indicates lifecycle corruption; do not let a caller remember a + // pair that was not selected atomically. + return nil, nil, fmt.Errorf("transaction execution binding changed during connection acquisition") + } + if tx == nil && boundConn != conn { + // A detached owner marker makes an idle connection unsafe to associate with + // this session. Preserve the connection result for normal callers only when + // no owner marker is present; identity-aware callers receive a nil binding. + if _, ownerPresent := p.txnConns.Load(id); ownerPresent { + boundConn = nil + } + } + return conn, tx, nil +} + +func (p *ConnectionPool) getConnLocked(ctx context.Context, id uint32) (*stdsql.Conn, error) { + // An active transaction is bound to the physical connection captured at + // BeginTx. Never consult the generic session map for this path: after a + // reconnect or a partially failed lifecycle operation it may point at a + // different connection. A missing/invalid owner is fail-closed rather than + // guessed from that generic entry. + if owner, err := p.activeTxnOwnerLocked(id); err != nil { + return nil, err + } else if owner != nil { + return owner, nil + } + var conn *stdsql.Conn + entry, present, err := p.connectionEntryLocked(id) + if err != nil { + return nil, err + } + if !present { + if p.DB == nil { + return nil, fmt.Errorf("connection pool database is unavailable") + } + c, err := p.DB.Conn(ctx) + if err != nil { + return nil, err + } + if _, transactionActive := p.txns.Load(id); !transactionActive { + if err := p.initializeConnection(ctx, c); err != nil { + _ = c.Close() + return nil, err + } + } + if err := p.registerMySQLUDFs(c); err != nil { + _ = c.Close() + return nil, err + } + p.closedConns.Delete(id) + p.conns.Store(id, c) + conn = c + } else { + conn = entry + // A session transaction owns this connection's transaction-scoped + // state. Re-running the initializer here would execute LOAD/CREATE + // SECRET inside that transaction and could alter or roll back with user + // work. GetTxn initializes before BeginTx; keep the state stable until + // the transaction is closed. + if _, transactionActive := p.txns.Load(id); !transactionActive { + if err := p.initializeConnection(ctx, conn); err != nil { + // Do not leave a failed or partially initialized connection available + // to a later request. CompareAndDelete avoids removing a replacement + // installed by a concurrent recovery path. + p.conns.CompareAndDelete(id, conn) + p.closedConns.Store(id, struct{}{}) + _ = conn.Close() + return nil, err + } + } + } + return conn, nil +} + +// activeTxnOwnerLocked returns the physical owner for an active transaction. +// A nil owner with a nil error means that no transaction is active; callers +// must hold the per-session lifecycle lock while using this helper. +func (p *ConnectionPool) activeTxnOwnerLocked(id uint32) (*stdsql.Conn, error) { + _, active, err := p.transactionEntryLocked(id) + if !active { + return nil, nil + } + if err != nil { + return nil, err + } + owner, present, err := p.transactionOwnerEntryLocked(id) + if !present { + return nil, fmt.Errorf("active transaction has no physical owner") + } + if err != nil { + return nil, err + } + return owner, nil +} + +// connectionEntryLocked and transactionEntryLocked distinguish a missing map +// entry from a malformed one. The latter is a lifecycle corruption and must be +// surfaced to callers that can return an error instead of being guessed away. +func (p *ConnectionPool) connectionEntryLocked(id uint32) (*stdsql.Conn, bool, error) { + entry, ok := p.conns.Load(id) + if !ok { + return nil, false, nil + } + conn, valid := entry.(*stdsql.Conn) + if !valid || conn == nil { + return nil, true, fmt.Errorf("connection mapping is invalid") + } + return conn, true, nil +} + +func (p *ConnectionPool) transactionEntryLocked(id uint32) (*stdsql.Tx, bool, error) { + entry, ok := p.txns.Load(id) + if !ok { + return nil, false, nil + } + tx, valid := entry.(*stdsql.Tx) + if !valid || tx == nil { + return nil, true, fmt.Errorf("active transaction mapping is invalid") + } + return tx, true, nil +} + +func (p *ConnectionPool) transactionOwnerEntryLocked(id uint32) (*stdsql.Conn, bool, error) { + entry, ok := p.txnConns.Load(id) + if !ok || entry == nil { + return nil, false, nil + } + conn, valid := entry.(*stdsql.Conn) + if !valid || conn == nil { + return nil, true, fmt.Errorf("active transaction has invalid physical owner") + } + return conn, true, nil +} + +// DuckDB stores registered scalar UDFs in the database catalog shared by all connections. +func (p *ConnectionPool) registerMySQLUDFs(conn *stdsql.Conn) error { + p.registerMySQLUDFsOnce.Do(func() { + if err := registerMySQLRand(conn); err != nil { + p.registerMySQLUDFsErr = fmt.Errorf("register mysql_rand: %w", err) + return + } + if err := registerMySQLRandomBytes(conn); err != nil { + p.registerMySQLUDFsErr = fmt.Errorf("register mysql_random_bytes: %w", err) + return + } + if err := registerMySQLStringToVector(conn); err != nil { + p.registerMySQLUDFsErr = fmt.Errorf("register string_to_vector: %w", err) + } + }) + return p.registerMySQLUDFsErr +} + +func (p *ConnectionPool) GetConnForSchema(ctx context.Context, id uint32, schemaName string) (*stdsql.Conn, error) { + conn, _, err := p.GetConnForSchemaWithBinding(ctx, id, schemaName) + return conn, err +} + +// GetConnForSchemaWithBinding is the schema-selecting counterpart to +// GetConnWithBinding. The connection, schema selection, and transaction +// identity are observed under one lifecycle/session lock. +func (p *ConnectionPool) GetConnForSchemaWithBinding(ctx context.Context, id uint32, schemaName string) (*stdsql.Conn, *stdsql.Tx, error) { + return p.getConnWithBinding(ctx, id, schemaName) +} + +// GetExecutionSnapshotLease acquires the connection and observes the +// transaction binding while holding the lifecycle lock. The returned release +// function keeps that lifecycle reader held until the caller has finished all +// driver work and consumed any result iterator. This prevents Close, Reset, +// and provider Restart from retiring the selected physical generation in the +// gap between snapshot acquisition and execution. +// +// The release function is idempotent. On an acquisition error the lease is +// released before returning and the returned callback is a no-op, so callers +// do not need a special error cleanup branch. +func (p *ConnectionPool) GetExecutionSnapshotLease( + ctx context.Context, + id uint32, + schemaName string, + catalogOnly bool, +) (*stdsql.Conn, *stdsql.Tx, func(), error) { + if ctx == nil { + ctx = context.Background() + } + if err := p.connectionAdmissionError(ctx); err != nil { + return nil, nil, func() {}, err + } + p.lifecycleMu.RLock() + var releaseOnce sync.Once + release := func() { + releaseOnce.Do(func() { + p.lifecycleMu.RUnlock() + }) + } + unlock := p.lockSession(id) + defer unlock() + if err := p.connectionAdmissionError(ctx); err != nil { + release() + return nil, nil, func() {}, err + } + + if tx, active, txErr := p.transactionEntryLocked(id); active { + if txErr != nil { + release() + return nil, nil, func() {}, txErr + } + conn, connOK, connErr := p.transactionOwnerEntryLocked(id) + if connErr != nil { + release() + return nil, tx, func() {}, connErr + } + if !connOK { + release() + return nil, tx, func() {}, fmt.Errorf("transaction execution snapshot has no physical owner") + } + return conn, tx, release, nil + } + var ( + conn *stdsql.Conn + err error + ) + if catalogOnly { + conn, err = p.getConnLocked(ctx, id) + } else { + conn, err = p.getConnForSchemaLocked(ctx, id, schemaName) + } + if err != nil { + release() + return nil, nil, func() {}, err + } + return conn, nil, release, nil +} + +// GetExecutionSnapshot acquires the connection and observes the transaction +// binding while holding the lifecycle lock. It retains the historical +// snapshot-only contract by releasing the lifecycle lease before returning; +// callers that execute or consume results after this point should use +// GetExecutionSnapshotLease through the adapter helper instead. +func (p *ConnectionPool) GetExecutionSnapshot(ctx context.Context, id uint32, schemaName string, catalogOnly bool) (*stdsql.Conn, *stdsql.Tx, error) { + conn, tx, release, err := p.GetExecutionSnapshotLease(ctx, id, schemaName, catalogOnly) + if release != nil { + release() + } + return conn, tx, err +} + +func (p *ConnectionPool) CloseConn(id uint32) error { + var completions transactionCompletionBatch + defer completions.run() + p.lifecycleMu.ReleaseLock() + defer p.lifecycleMu.ReleaseUnlock() + unlock := p.lockSession(id) + defer unlock() + return p.closeConnLockedWithBatch(id, nil, &completions) +} + +func (p *ConnectionPool) GetTxn(ctx context.Context, id uint32, schemaName string, options *stdsql.TxOptions) (*stdsql.Tx, error) { + _, tx, err := p.GetTxnWithBinding(ctx, id, schemaName, options) + return tx, err +} + +// GetTxnWithBinding opens (or observes) a transaction and returns its exact +// physical owner from the same lifecycle/session critical section. Session +// adapters use this to remember a close identity without a second, racy pool +// lookup. +func (p *ConnectionPool) GetTxnWithBinding(ctx context.Context, id uint32, schemaName string, options *stdsql.TxOptions) (*stdsql.Conn, *stdsql.Tx, error) { + if ctx == nil { + ctx = context.Background() + } + // New-transaction admission may wait for a normal provider generation + // transition. Do that wait + // before entering lifecycleMu.RLock; otherwise a caller that is blocked on + // cleanupActive would itself keep pool teardown from acquiring its writer. + // Existing transactions remain available to the dedicated release/finalizer + // accessors while cleanup is waiting. This execution-facing method checks + // connection admission before and after its lifecycle/session snapshot so an + // active mapping cannot be used to start more work in the retiring generation. + for { + if _, alreadyActive := p.txns.Load(id); alreadyActive { + if err := p.connectionAdmissionError(ctx); err != nil { + return nil, nil, err + } + p.lifecycleMu.RLock() + unlock := p.lockSession(id) + tx, active, err := p.transactionEntryLocked(id) + if active { + if err != nil { + unlock() + p.lifecycleMu.RUnlock() + return nil, nil, err + } + if err := p.connectionAdmissionError(ctx); err != nil { + unlock() + p.lifecycleMu.RUnlock() + return nil, nil, err + } + owner, err := p.activeTxnOwnerLocked(id) + if err != nil { + unlock() + p.lifecycleMu.RUnlock() + return nil, nil, err + } + unlock() + p.lifecycleMu.RUnlock() + return owner, tx, nil + } + // The observed transaction closed between the lock-free probe and + // the lifecycle snapshot. Retry so a new transaction gets a fresh + // admission reservation outside lifecycleMu. + unlock() + p.lifecycleMu.RUnlock() + continue + } + + if err := p.transactionAdmissionError(); err != nil { + return nil, nil, err + } + completeAdmission, admissionErr := p.beginTransactionAdmissionWithError() + if admissionErr != nil { + return nil, nil, admissionErr + } + p.lifecycleMu.RLock() + unlock := p.lockSession(id) + // A generation can become poisoned while the admission hook waits for a + // cleanup barrier. Recheck after the hook so its historical nil-completion + // convention cannot accidentally allow a new BeginTx into a retired + // generation. + if err := p.transactionAdmissionError(); err != nil { + completeAdmission(nil) + unlock() + p.lifecycleMu.RUnlock() + return nil, nil, err + } + + if tx, active, err := p.transactionEntryLocked(id); active { + var ownerErr error + var owner *stdsql.Conn + if err == nil { + owner, ownerErr = p.activeTxnOwnerLocked(id) + } + completeAdmission(nil) + unlock() + p.lifecycleMu.RUnlock() + if err != nil { + return nil, nil, err + } + if ownerErr != nil { + return nil, nil, ownerErr + } + return owner, tx, nil + } + + conn, err := p.getConnForSchemaLocked(ctx, id, schemaName) if err != nil { - return nil, err + completeAdmission(nil) + unlock() + p.lifecycleMu.RUnlock() + return nil, nil, err } // A session transaction can span multiple protocol requests (for example, // after SET autocommit=0), so a request-scoped cancellation must not end it. - t, err := conn.BeginTx(context.WithoutCancel(ctx), options) + tx, err := conn.BeginTx(context.WithoutCancel(ctx), options) if err != nil { - return nil, err + completeAdmission(nil) + unlock() + p.lifecycleMu.RUnlock() + return nil, nil, err } - p.txns.Store(id, t) - tx = t - } else { - tx = entry.(*stdsql.Tx) + // Keep the transaction and its physical owner as one lifecycle binding while + // holding the pool mutex. Readers can therefore never observe a transaction + // without the connection it owns. + p.txns.Store(id, tx) + p.txnConns.Store(id, conn) + p.markTransactionActive(tx) + completeAdmission(tx) + unlock() + p.lifecycleMu.RUnlock() + return conn, tx, nil } - return tx, nil } func (p *ConnectionPool) TryGetTxn(id uint32) *stdsql.Tx { - entry, ok := p.txns.Load(id) - if !ok { - return nil - } - return entry.(*stdsql.Tx) + p.lifecycleMu.RLock() + defer p.lifecycleMu.RUnlock() + unlock := p.lockSession(id) + defer unlock() + return p.tryGetTxnLocked(id) +} + +// TryGetTxnForRelease is the release-side transaction snapshot. Unlike the +// ordinary lookup, it remains admissible while pool shutdown is waiting for +// readers to drain, so a close/finalizer path can still discover the binding it +// needs to release. The returned pointer is only a snapshot; callers must use +// an identity-safe finalizer for the subsequent driver operation. +func (p *ConnectionPool) TryGetTxnForRelease(id uint32) *stdsql.Tx { + p.lifecycleMu.ReleaseLock() + defer p.lifecycleMu.ReleaseUnlock() + unlock := p.lockSession(id) + defer unlock() + return p.tryGetTxnLocked(id) +} + +func (p *ConnectionPool) tryGetTxnLocked(id uint32) *stdsql.Tx { + tx, _, _ := p.transactionEntryLocked(id) + return tx } func (p *ConnectionPool) CloseTxn(id uint32) { + _ = p.CloseTxnAndGet(id) +} + +// CloseTxnAndGet removes the transaction currently bound to id and returns the +// exact mapping that was removed. Capturing, identity removal, and completion +// notification happen under one release/session lifecycle section; callers do +// not need to perform a racy lookup followed by an unconditional CloseTxn. +// A nil result means that no valid transaction mapping was removed. +func (p *ConnectionPool) CloseTxnAndGet(id uint32) *stdsql.Tx { + var completions transactionCompletionBatch + defer completions.run() + p.lifecycleMu.ReleaseLock() + defer p.lifecycleMu.ReleaseUnlock() + unlock := p.lockSession(id) + defer unlock() + tx := p.closeTxnLocked(id, nil) + completions.add(p.finishTransactionWithOutcomeDeferred(tx, false)) + return tx +} + +// GetTxnBinding returns the physical connection and transaction currently +// associated with a logical session. The connection is returned even when no +// transaction is active, which lets statement callers use one atomic snapshot +// instead of acquiring a connection and then racing a transaction opener. +func (p *ConnectionPool) GetTxnBinding(id uint32) (*stdsql.Conn, *stdsql.Tx) { + p.lifecycleMu.RLock() + defer p.lifecycleMu.RUnlock() + unlock := p.lockSession(id) + defer unlock() + return p.txnBindingLocked(id) +} + +// GetTxnBindingForRelease captures a transaction binding from a release-side +// lifecycle section. It is intended for close/finalizer paths that may run +// after shutdown has been announced; ordinary execution must continue using +// GetTxnBinding or GetExecutionSnapshot so new work remains excluded during +// teardown. +func (p *ConnectionPool) GetTxnBindingForRelease(id uint32) (*stdsql.Conn, *stdsql.Tx) { + p.lifecycleMu.ReleaseLock() + defer p.lifecycleMu.ReleaseUnlock() + unlock := p.lockSession(id) + defer unlock() + return p.txnBindingLocked(id) +} + +// GetTxnBindingForReleaseStatus is the error-preserving form used by lifecycle +// callback bridges. A malformed active transaction mapping must remain +// distinguishable from an absent mapping; otherwise a callback could release a +// lease while the physical transaction is still live but uninspectable. +func (p *ConnectionPool) GetTxnBindingForReleaseStatus(id uint32) (*stdsql.Conn, *stdsql.Tx, bool, error) { + p.lifecycleMu.ReleaseLock() + defer p.lifecycleMu.ReleaseUnlock() + unlock := p.lockSession(id) + defer unlock() + tx, active, txErr := p.transactionEntryLocked(id) + if active { + if txErr != nil { + return nil, nil, true, txErr + } + conn, present, ownerErr := p.transactionOwnerEntryLocked(id) + if ownerErr != nil { + return nil, tx, true, ownerErr + } + if !present { + return nil, tx, true, fmt.Errorf("active transaction has no physical owner") + } + return conn, tx, true, nil + } + conn, present, connErr := p.connectionEntryLocked(id) + if connErr != nil { + return nil, nil, false, connErr + } + if !present { + // A transaction owner marker without a transaction is a lifecycle + // corruption, not an idle session. Treat it as active so callback bridges + // cannot release leases or close a checked-out owner before recovery has + // resolved the marker. + if _, ownerPresent := p.txnConns.Load(id); ownerPresent { + _, _, ownerErr := p.transactionOwnerEntryLocked(id) + if ownerErr != nil { + return nil, nil, true, ownerErr + } + return nil, nil, true, fmt.Errorf("transaction owner mapping exists without active transaction") + } + return nil, nil, false, nil + } + return conn, nil, false, nil +} + +func (p *ConnectionPool) txnBindingLocked(id uint32) (*stdsql.Conn, *stdsql.Tx) { + tx, active, txErr := p.transactionEntryLocked(id) + if active { + if txErr != nil { + // An invalid active mapping must not be mistaken for an idle session + // and then fall back to p.conns. + return nil, nil + } + // An active transaction is only usable with the exact owner captured at + // BeginTx. Never substitute the session's generic connection mapping: + // after a reset/replacement that entry may belong to a different + // transaction lifecycle. + if conn, present, err := p.transactionOwnerEntryLocked(id); present && err == nil { + return conn, tx + } + return nil, tx + } + // A detached owner marker is not a usable idle connection. Returning nil + // prevents release callers from falling back to p.conns and retiring a + // physical handle whose transaction identity is still unknown. + if _, ownerPresent := p.txnConns.Load(id); ownerPresent { + return nil, nil + } + if conn, present, err := p.connectionEntryLocked(id); present && err == nil { + return conn, nil + } + return nil, nil +} + +func (p *ConnectionPool) closeTxnLocked(id uint32, expected *stdsql.Tx) *stdsql.Tx { + tx, present, err := p.transactionEntryLocked(id) + if !present { + return nil + } + if err != nil { + // An unconditional close is allowed to discard a poisoned mapping so it + // cannot block future work. Identity-aware cleanup leaves malformed state + // untouched because it cannot prove that the caller owns it. + if expected == nil { + p.txns.Delete(id) + p.txnConns.Delete(id) + } + return nil + } + if expected != nil && tx != expected { + return nil + } p.txns.Delete(id) + p.txnConns.Delete(id) + return tx } -func (p *ConnectionPool) Close() error { - var txns []*stdsql.Tx - p.txns.Range(func(_, value any) bool { - txns = append(txns, value.(*stdsql.Tx)) +// CloseTxnIf removes a transaction mapping only if it still refers to tx. +// Older deferred cleanup must not delete a newer transaction for the same +// logical session ID. +func (p *ConnectionPool) CloseTxnIf(id uint32, tx *stdsql.Tx) { + if tx == nil { + return + } + var completions transactionCompletionBatch + defer completions.run() + p.lifecycleMu.ReleaseLock() + defer p.lifecycleMu.ReleaseUnlock() + unlock := p.lockSession(id) + defer unlock() + closed := p.closeTxnLocked(id, tx) + completions.add(p.finishTransactionWithOutcomeDeferred(closed, false)) +} + +// CommitTxn commits the transaction only while it is still the transaction +// bound to id. The driver operation, identity check, and bad-connection +// eviction are serialized under lifecycleMu, so a delayed commit cannot +// remove or close a replacement binding. +func (p *ConnectionPool) CommitTxn(id uint32, expected *stdsql.Tx) error { + _, err := p.CommitTxnWithOutcome(id, expected) + return err +} + +// CommitTxnWithOutcome is the identity-aware commit surface used by the +// production Session adapter. matched is false when the expected transaction +// is no longer the current binding (for example, it was already rolled back or +// replaced). The legacy CommitTxn wrapper intentionally preserves its historic +// no-op behavior for that stale case, while callers that need to report the +// distinction can turn matched=false into a binding error without guessing. +func (p *ConnectionPool) CommitTxnWithOutcome(id uint32, expected *stdsql.Tx) (matched bool, err error) { + var completions transactionCompletionBatch + defer completions.run() + p.lifecycleMu.ReleaseLock() + defer p.lifecycleMu.ReleaseUnlock() + unlock := p.lockSession(id) + defer unlock() + + tx, conn, matched, bindingErr := p.txnBindingForFinalizeLocked(id, expected) + if bindingErr != nil { + return false, bindingErr + } + if !matched { + // Another owner already finalized this transaction. Treat a delayed + // callback as a no-op rather than touching the replacement transaction. + return false, nil + } + err = tx.Commit() + closed := p.closeTxnLocked(id, tx) + completions.add(p.finishTransactionWithOutcomeDeferred(closed, err == nil)) + if err != nil && conn != nil { + // A failed commit leaves the driver state unknown. Evict this exact + // physical owner before it can be reused. The generic session mapping may + // already point at a replacement connection; closeExactOwnerLocked keeps + // that replacement untouched while still retiring the captured owner. + err = errors.Join(err, p.closeExactOwnerLockedWithBatch(id, conn, &completions)) + } + return true, err +} + +// RollbackTxn rolls back the transaction only while it is still the binding +// owned by id. It returns the physical owner and whether the driver proved the +// transaction inactive. Unexpected rollback errors invalidate and evict that +// owner; a stale callback returns no owner so it cannot clean a replacement. +func (p *ConnectionPool) RollbackTxn(id uint32, expected *stdsql.Tx) (*stdsql.Conn, bool, error) { + var conn *stdsql.Conn + var inactive bool + var err error + var completions transactionCompletionBatch + defer completions.run() + p.lifecycleMu.ReleaseLock() + defer p.lifecycleMu.ReleaseUnlock() + unlock := p.lockSession(id) + defer unlock() + + tx, owner, matched, bindingErr := p.txnBindingForFinalizeLocked(id, expected) + if bindingErr != nil { + return nil, false, bindingErr + } + if !matched { + return nil, false, nil + } + conn = owner + err = tx.Rollback() + inactive = transactionInactiveError(err) + closed := p.closeTxnLocked(id, tx) + completions.add(p.finishTransactionWithOutcomeDeferred(closed, false)) + if !inactive && conn != nil { + // The physical connection is not trustworthy after an unexpected + // rollback result. Retire the exact owner even when the generic session + // mapping was replaced while this transaction was active. + err = errors.Join(err, p.closeExactOwnerLockedWithBatch(id, conn, &completions)) + } + return conn, inactive, err +} + +// RollbackTxnWithCleanup performs rollback, unbinds the transaction, and runs +// cleanup before releasing the session lock. A replacement transaction for the +// same logical session therefore cannot slip between rollback and maintenance. +// The callback is invoked only for the still-current transaction and only when +// the driver proves that it is inactive. +func (p *ConnectionPool) RollbackTxnWithCleanup( + id uint32, + expected *stdsql.Tx, + cleanup func(*stdsql.Conn) error, +) (bool, error) { + var completions transactionCompletionBatch + defer completions.run() + p.lifecycleMu.ReleaseLock() + defer p.lifecycleMu.ReleaseUnlock() + unlock := p.lockSession(id) + defer unlock() + + tx, conn, matched, bindingErr := p.txnBindingForFinalizeLocked(id, expected) + if bindingErr != nil { + return false, bindingErr + } + if !matched { + return false, nil + } + err := tx.Rollback() + inactive := transactionInactiveError(err) + closed := p.closeTxnLocked(id, tx) + completions.add(p.finishTransactionWithOutcomeDeferred(closed, false)) + if !inactive && conn != nil { + // The physical connection is not trustworthy after an unexpected + // rollback result. Retire the exact owner even when the generic session + // mapping was replaced while this transaction was active. + err = errors.Join(err, p.closeExactOwnerLockedWithBatch(id, conn, &completions)) + return inactive, err + } + if inactive && conn != nil && cleanup != nil { + err = errors.Join(err, cleanup(conn)) + } + return inactive, err +} + +func (p *ConnectionPool) txnBindingForFinalizeLocked(id uint32, expected *stdsql.Tx) (*stdsql.Tx, *stdsql.Conn, bool, error) { + tx, present, err := p.transactionEntryLocked(id) + if !present { + return nil, nil, false, nil + } + if err != nil { + return nil, nil, false, err + } + if expected != nil && tx != expected { + return nil, nil, false, nil + } + conn, present, ownerErr := p.transactionOwnerEntryLocked(id) + if ownerErr != nil { + return nil, nil, false, ownerErr + } + if !present { + return nil, nil, false, fmt.Errorf("active transaction has no physical owner") + } + return tx, conn, true, nil +} + +func (p *ConnectionPool) CloseConnIf(id uint32, conn *stdsql.Conn) error { + if conn == nil { + return nil + } + var completions transactionCompletionBatch + defer completions.run() + p.lifecycleMu.ReleaseLock() + defer p.lifecycleMu.ReleaseUnlock() + unlock := p.lockSession(id) + defer unlock() + current, present, err := p.connectionEntryLocked(id) + if err != nil { + return err + } + if !present || current != conn { + return nil + } + return p.closeConnLockedWithBatch(id, conn, &completions) +} + +// CloseConnIfBinding closes a connection only while both the generic session +// mapping and the transaction binding still match the caller's snapshot. The +// session lock keeps a replacement transaction from being admitted between the +// identity check and closeConnLocked, which may itself roll back an active tx. +func (p *ConnectionPool) CloseConnIfBinding(id uint32, conn *stdsql.Conn, expectedTx *stdsql.Tx) error { + if conn == nil { + return nil + } + var completions transactionCompletionBatch + defer completions.run() + p.lifecycleMu.ReleaseLock() + defer p.lifecycleMu.ReleaseUnlock() + unlock := p.lockSession(id) + defer unlock() + if expectedTx != nil { + // For a transaction-bound close, the transaction-owner map is the + // authoritative identity. Do not require the generic p.conns entry to + // still point at conn: a reconnect may have replaced that entry while the + // old transaction remains active, and the old owner can still be retired + // safely in isolation. + currentTx, active, txErr := p.transactionEntryLocked(id) + if txErr != nil { + return txErr + } + if !active || currentTx != expectedTx { + return nil + } + owner, ownerOK, ownerErr := p.transactionOwnerEntryLocked(id) + if ownerErr != nil { + return ownerErr + } + if !ownerOK { + return fmt.Errorf("active transaction has no physical owner") + } + if owner != conn { + return nil + } + return p.closeExactBindingLockedWithBatch(id, conn, expectedTx, &completions) + } + + current, present, err := p.connectionEntryLocked(id) + if err != nil { + return err + } + if !present || current != conn { + return nil + } + _, active, txErr := p.transactionEntryLocked(id) + if txErr != nil { + return txErr + } + if active { + return nil + } + if _, ownerMapped := p.txnConns.Load(id); ownerMapped { + return nil + } + return p.closeConnLockedWithBatch(id, conn, &completions) +} + +func transactionInactiveError(err error) bool { + if err == nil { + return true + } + return errors.Is(err, stdsql.ErrTxDone) || + strings.Contains(strings.ToLower(err.Error()), "no transaction is active") +} + +func (p *ConnectionPool) getConnForSchemaLocked(ctx context.Context, id uint32, schemaName string) (*stdsql.Conn, error) { + conn, err := p.getConnLocked(ctx, id) + if err != nil { + return nil, err + } + // GetTxn performs schema selection before BeginTx. Never issue session-state + // SQL through *sql.Conn while an active *sql.Tx owns this connection. + if _, active, txErr := p.transactionEntryLocked(id); active { + if txErr != nil { + return nil, txErr + } + return conn, nil + } + if schemaName == "" { + return conn, nil + } + schemaCtx := context.WithoutCancel(ctx) + var currentSchema string + if err := conn.QueryRowContext(schemaCtx, "SELECT CURRENT_SCHEMA").Scan(¤tSchema); err != nil { + logrus.WithError(err).Error("Failed to get current schema") + return nil, err + } + if currentSchema == schemaName { + return conn, nil + } + if _, err := conn.ExecContext(schemaCtx, "USE "+FullSchemaName(p.currentCatalogLocked(id), schemaName)); err != nil { + if IsDuckDBSetSchemaNotFoundError(err) { + return nil, sql.ErrDatabaseNotFound.New(schemaName) + } + logrus.WithField("schema", schemaName).WithError(err).Error("Failed to switch schema") + return nil, err + } + return conn, nil +} + +// closeExternalSessionsLocked invalidates every handle in the retiring pool +// generation. The lifecycle writer guarantees all admitted execution leases +// have drained; opMu additionally serializes teardown with hook registration. +// Cleanup and driver failures are accumulated so one broken handle cannot keep +// later sessions from being rolled back and closed. +func (p *ConnectionPool) closeExternalSessionsLocked(completions *transactionCompletionBatch) error { + var sessions []*ExternalSession + p.externalSessions.Range(func(key, _ any) bool { + sessions = append(sessions, key.(*ExternalSession)) return true }) + var lastErr error - for _, tx := range txns { - if err := tx.Rollback(); err != nil && !strings.Contains(err.Error(), "no transaction is active") { - logrus.WithError(err).Warn("Failed to rollback transaction") - lastErr = err + for _, session := range sessions { + session.opMu.Lock() + teardown, owner := session.beginTeardownLocked(true) + if !owner { + lastErr = errors.Join(lastErr, session.waitForTeardownLocked()) + session.opMu.Unlock() + continue + } + teardownErr := session.finishTeardownLocked(teardown, nil, completions) + session.opMu.Unlock() + lastErr = errors.Join(lastErr, teardownErr) + } + return lastErr +} + +func (p *ConnectionPool) Close() error { + var completions transactionCompletionBatch + defer p.clearCompletedTransactionCompletions() + defer completions.run() + preCloseErr := p.closeDatabaseForShutdown() + p.lifecycleMu.LockAfterShutdown() + defer p.lifecycleMu.Unlock() + err := errors.Join( + preCloseErr, + p.closeExternalSessionsLocked(&completions), + p.closeLockedWithBatch(&completions), + p.externalSessionPanicError(), + ) + return err +} + +// closePhysicalConnLocked retires exactly the supplied database/sql +// connection. The caller must hold the pool lifecycle/session locks and must +// have already dealt with any transaction that owns the connection. Marking +// the driver handle bad prevents database/sql from returning a connection with +// unknown state to the shared pool. +func (p *ConnectionPool) closePhysicalConnLocked(conn *stdsql.Conn) error { + if conn == nil { + return nil + } + var lastErr error + if err := conn.Raw(func(any) error { return driver.ErrBadConn }); err != nil && + !errors.Is(err, driver.ErrBadConn) && !errors.Is(err, stdsql.ErrConnDone) { + logrus.WithError(err).Warn("Failed to close connection during Raw function call") + lastErr = errors.Join(lastErr, err) + } + if err := conn.Close(); err != nil && !errors.Is(err, stdsql.ErrConnDone) { + logrus.WithError(err).Warn("Failed to close connection") + lastErr = errors.Join(lastErr, err) + } + return lastErr +} + +// closeExactOwnerLocked retires a captured transaction owner even when the +// generic session connection mapping has since been replaced. The replacement +// mapping is deliberately left intact; only the exact owner is bad/closed. +// This helper is used after a failed commit/rollback, when the transaction +// mapping has already been removed and the old owner must not leak. +func (p *ConnectionPool) closeExactOwnerLockedWithBatch( + id uint32, + owner *stdsql.Conn, + completions *transactionCompletionBatch, +) error { + if owner == nil { + return nil + } + current, present, mappingErr := p.connectionEntryLocked(id) + if mappingErr != nil { + // A malformed generic entry cannot identify a replacement safely. Retire + // only the captured owner, preserve the malformed entry for diagnostics, + // and report both lifecycle failures to the caller. + return errors.Join(mappingErr, p.closePhysicalConnLocked(owner)) + } + if present && current == owner { + return p.closeConnLockedWithBatch(id, owner, completions) + } + + // The transaction map was removed by the caller, but CompareAndDelete also + // clears a stale owner marker without touching a newer owner installed by a + // recovery path. + p.txnConns.CompareAndDelete(id, owner) + if !present { + p.closedConns.Store(id, struct{}{}) + } + return p.closePhysicalConnLocked(owner) +} + +// closeExactBindingLocked rolls back and retires an exact transaction/owner +// pair. It intentionally does not require p.conns[id] to point at owner: a +// reconnect may have installed a replacement generic connection while the +// original transaction is still active. In that case only owner is closed and +// the replacement mapping remains available to its owner. +func (p *ConnectionPool) closeExactBindingLockedWithBatch( + id uint32, + owner *stdsql.Conn, + expectedTx *stdsql.Tx, + completions *transactionCompletionBatch, +) error { + if owner == nil || expectedTx == nil { + return nil + } + tx, active, txErr := p.transactionEntryLocked(id) + if txErr != nil { + return txErr + } + if !active || tx != expectedTx { + return nil + } + boundOwner, ownerPresent, ownerErr := p.transactionOwnerEntryLocked(id) + if ownerErr != nil { + return ownerErr + } + if !ownerPresent { + return fmt.Errorf("active transaction has no physical owner") + } + if boundOwner != owner { + return nil + } + + rollbackErr := tx.Rollback() + closed := p.closeTxnLocked(id, tx) + completions.add(p.finishTransactionWithOutcomeDeferred(closed, false)) + closeErr := p.closePhysicalConnLocked(owner) + p.txnConns.CompareAndDelete(id, owner) + + // Remove the generic map only when it still points at this exact owner. A + // replacement connection must remain untouched; if no generic mapping + // remains, remember the completed lifecycle so CurrentCatalog does not + // incorrectly fall back to the default catalog. + current, present, mappingErr := p.connectionEntryLocked(id) + if mappingErr != nil { + return errors.Join(rollbackErr, closeErr, mappingErr) + } + if present && current == owner { + p.conns.Delete(id) + p.closedConns.Store(id, struct{}{}) + } else if !present { + p.closedConns.Store(id, struct{}{}) + } + return errors.Join(rollbackErr, closeErr) +} + +// closeDetachedSessionStateLocked retires transaction state for an +// unconditional session teardown when the generic p.conns entry is already +// absent (or malformed). The transaction-owner map is authoritative for the +// physical handle; dropping the logical maps without closing that handle +// leaks a checked-out database/sql connection and can leave an active DuckDB +// transaction behind. +// +// The caller must hold the lifecycle and session locks. Identity-sensitive +// callers do not use this helper because they cannot prove ownership of a +// detached mapping. +func (p *ConnectionPool) closeDetachedSessionStateLockedWithBatch( + id uint32, + completions *transactionCompletionBatch, +) error { + var lastErr error + // A malformed transaction marker may still hide a live *sql.Tx. Its + // transaction-held connection close lock makes Raw/Close potentially block + // forever, so an unconditional teardown must report the corruption and leave + // the physical owner untouched for an operator-driven recovery path. + if _, present, txErr := p.transactionEntryLocked(id); present && txErr != nil { + return txErr + } + + // Capture the owner before closeTxnLocked removes the owner mapping. Even a + // malformed transaction marker may have a valid owner that still needs to be + // retired during an unconditional teardown. + owner, ownerPresent, ownerErr := p.transactionOwnerEntryLocked(id) + if ownerErr != nil { + lastErr = errors.Join(lastErr, ownerErr) + } + + tx, txPresent, txErr := p.transactionEntryLocked(id) + if txErr != nil { + lastErr = errors.Join(lastErr, txErr) + } else if txPresent { + if rollbackErr := tx.Rollback(); rollbackErr != nil && !transactionInactiveError(rollbackErr) { + logrus.WithError(rollbackErr).Warn("Failed to rollback transaction") + lastErr = errors.Join(lastErr, rollbackErr) + } + } + + // An unconditional close is allowed to discard a poisoned transaction + // marker. closeTxnLocked also removes txnConns for a valid transaction. + closed := p.closeTxnLocked(id, nil) + completions.add(p.finishTransactionWithOutcomeDeferred(closed, false)) + + if ownerPresent && ownerErr == nil && owner != nil { + lastErr = errors.Join(lastErr, p.closePhysicalConnLocked(owner)) + p.txnConns.CompareAndDelete(id, owner) + } else if ownerPresent { + // Do not leave an invalid owner marker to block a future lifecycle. There + // is no type-safe physical handle to close in this branch, so preserve the + // validation error while discarding only the map entry. + p.txnConns.Delete(id) + } + + return lastErr +} + +func (p *ConnectionPool) closeConnLockedWithBatch( + id uint32, + expected *stdsql.Conn, + completions *transactionCompletionBatch, +) error { + // Never retire a physical connection while the transaction marker for this + // session is malformed. The marker may conceal a live *sql.Tx whose + // database/sql close lock would make Raw/Close wait indefinitely. + if _, present, txErr := p.transactionEntryLocked(id); present && txErr != nil { + return txErr + } + conn, present, connErr := p.connectionEntryLocked(id) + if !present { + if expected == nil { + lastErr := p.closeDetachedSessionStateLockedWithBatch(id, completions) + p.closedConns.Store(id, struct{}{}) + return lastErr + } + return nil + } + if connErr != nil { + if expected != nil { + return connErr + } + // An unconditional close still has to retire any transaction owner even + // when the generic entry itself is poisoned. Preserve the mapping error + // for diagnostics, but do not strand the detached physical state. + lastErr := errors.Join(connErr, p.closeDetachedSessionStateLockedWithBatch(id, completions)) + p.closedConns.Store(id, struct{}{}) + return lastErr + } + if expected != nil && conn != expected { + return nil + } + var lastErr error + if tx, txOK, txErr := p.transactionEntryLocked(id); txOK { + if txErr != nil { + // The connection is explicitly being retired, so remove the poisoned + // transaction marker but never attempt a rollback through an unknown + // value. The owner cannot be closed safely while an untyped live + // database/sql transaction may still hold its close mutex; closeLocked's + // error report preserves this lifecycle corruption for the caller. + lastErr = txErr + p.txns.Delete(id) + p.txnConns.Delete(id) + } else { + owner, ownerPresent, ownerErr := p.transactionOwnerEntryLocked(id) + if ownerErr != nil { + if expected != nil { + // An invalid owner is not evidence that the generic connection is + // safe to retire. Leave the active transaction untouched and fail + // closed so callers cannot accidentally close a replacement. + return ownerErr + } + // CloseConn is unconditional. Roll back the transaction through its + // own handle, discard the malformed owner marker, and continue to + // retire the generic connection while returning the corruption error. + lastErr = ownerErr + if err := tx.Rollback(); err != nil && !transactionInactiveError(err) { + lastErr = errors.Join(lastErr, err) + } + closed := p.closeTxnLocked(id, tx) + completions.add(p.finishTransactionWithOutcomeDeferred(closed, false)) + } else if !ownerPresent { + if expected != nil { + return fmt.Errorf("active transaction has no physical owner") + } + // The transaction object is still usable for rollback even though + // its owner identity was lost. Unconditional teardown must not leave + // that transaction open behind the generic connection. + if err := tx.Rollback(); err != nil && !transactionInactiveError(err) { + lastErr = err + } + closed := p.closeTxnLocked(id, tx) + completions.add(p.finishTransactionWithOutcomeDeferred(closed, false)) + } else if owner == conn || expected == nil { + if err := tx.Rollback(); err != nil && + !errors.Is(err, stdsql.ErrTxDone) && + !strings.Contains(strings.ToLower(err.Error()), "no transaction is active") { + logrus.WithError(err).Warn("Failed to rollback transaction") + lastErr = err + } + closed := p.closeTxnLocked(id, tx) + completions.add(p.finishTransactionWithOutcomeDeferred(closed, false)) + if owner != conn { + // CloseConn(id) is unconditional and must retire both the active + // transaction owner and a stale/replacement generic connection. + lastErr = errors.Join(lastErr, p.closePhysicalConnLocked(owner)) + } + } + // When expected != nil and owner != conn, the active transaction is + // deliberately left mapped. The generic connection below is still the + // caller-supplied identity and can be retired independently. + } + } else { + if expected == nil { + // A stale owner marker without a transaction is also a checked-out + // physical handle. Sweep it before closing the generic mapping, while + // avoiding a second close when both entries point at the same object. + owner, ownerPresent, ownerErr := p.transactionOwnerEntryLocked(id) + if ownerErr != nil { + lastErr = errors.Join(lastErr, ownerErr) + p.txnConns.Delete(id) + } else if ownerPresent { + p.txnConns.Delete(id) + if owner != nil && owner != conn { + lastErr = errors.Join(lastErr, p.closePhysicalConnLocked(owner)) + } + } } + closed := p.closeTxnLocked(id, nil) + completions.add(p.finishTransactionWithOutcomeDeferred(closed, false)) + } + lastErr = errors.Join(lastErr, p.closePhysicalConnLocked(conn)) + p.conns.Delete(id) + p.closedConns.Store(id, struct{}{}) + return lastErr +} + +func (p *ConnectionPool) closeLockedWithBatch(completions *transactionCompletionBatch) error { + type transactionEntry struct { + id uint32 + tx *stdsql.Tx } - var conns []*stdsql.Conn - p.conns.Range(func(_, value any) bool { - conns = append(conns, value.(*stdsql.Conn)) + var ( + txns []transactionEntry + conns = make(map[*stdsql.Conn]struct{}) + protectedConns = make(map[*stdsql.Conn]struct{}) + validConnByID = make(map[uint32]*stdsql.Conn) + validOwnerByID = make(map[uint32]*stdsql.Conn) + validTxnByID = make(map[uint32]*stdsql.Tx) + malformedTxnIDs = make(map[uint32]struct{}) + malformedConnIDs = make(map[uint32]struct{}) + malformedOwnerIDs = make(map[uint32]struct{}) + lastErr error + ) + + // Inspect every transaction entry first. A malformed value cannot be + // interpreted as a live *sql.Tx, but it must not prevent unrelated valid + // sessions from being rolled back and closed. Keep its ID marked so any + // physical owner associated with that session remains protected below. + p.txns.Range(func(key, value any) bool { + id, keyOK := key.(uint32) + if !keyOK { + lastErr = errors.Join(lastErr, fmt.Errorf("active transaction mapping key is invalid")) + return true + } + tx, ok := value.(*stdsql.Tx) + if !ok || tx == nil { + lastErr = errors.Join(lastErr, fmt.Errorf("active transaction mapping is invalid for session %d", id)) + malformedTxnIDs[id] = struct{}{} + return true + } + validTxnByID[id] = tx + txns = append(txns, transactionEntry{id: id, tx: tx}) + return true + }) + + // Snapshot valid transaction owners. Owners attached to malformed transaction + // IDs are deliberately protected: closing a database/sql connection while an + // untyped transaction may still hold its internal close lock can hang the + // entire teardown. Invalid owner entries remain available for diagnostics. + p.txnConns.Range(func(key, value any) bool { + id, keyOK := key.(uint32) + if !keyOK { + lastErr = errors.Join(lastErr, fmt.Errorf("transaction owner mapping key is invalid")) + return true + } + conn, ok := value.(*stdsql.Conn) + if !ok || conn == nil { + lastErr = errors.Join(lastErr, fmt.Errorf("active transaction owner mapping is invalid for session %d", id)) + malformedOwnerIDs[id] = struct{}{} + return true + } + validOwnerByID[id] = conn + if _, malformed := malformedTxnIDs[id]; malformed { + protectedConns[conn] = struct{}{} + } else { + conns[conn] = struct{}{} + } + return true + }) + + // Generic connections are also captured before any close. Preserve entries + // whose session has a malformed transaction/owner mapping; all other valid + // connections can be retired independently of transaction processing. + p.conns.Range(func(key, value any) bool { + id, keyOK := key.(uint32) + if !keyOK { + lastErr = errors.Join(lastErr, fmt.Errorf("connection mapping key is invalid")) + return true + } + conn, ok := value.(*stdsql.Conn) + if !ok || conn == nil { + lastErr = errors.Join(lastErr, fmt.Errorf("connection mapping is invalid for session %d", id)) + malformedConnIDs[id] = struct{}{} + return true + } + validConnByID[id] = conn + if _, malformed := malformedTxnIDs[id]; malformed { + protectedConns[conn] = struct{}{} + } else if _, malformed := malformedOwnerIDs[id]; malformed { + protectedConns[conn] = struct{}{} + } else { + conns[conn] = struct{}{} + } return true }) - for _, conn := range conns { - if err := conn.Close(); err != nil && !errors.Is(err, stdsql.ErrConnDone) { - logrus.WithError(err).Warn("Failed to close connection") - lastErr = err + + // Roll back each distinct transaction once, then remove its identity + // binding. The owner was captured above, so it remains in the close set even + // though closeTxnLocked deletes txnConns. + seenTxns := make(map[*stdsql.Tx]struct{}, len(txns)) + for _, entry := range txns { + if _, seen := seenTxns[entry.tx]; seen { + continue + } + seenTxns[entry.tx] = struct{}{} + if err := entry.tx.Rollback(); err != nil && !transactionInactiveError(err) { + logrus.WithError(err).Warn("Failed to rollback transaction") + lastErr = errors.Join(lastErr, err) } + closed := p.closeTxnLocked(entry.id, entry.tx) + completions.add(p.finishTransactionWithOutcomeDeferred(closed, false)) + } + + // closePhysicalConnLocked is identity-safe and marks each handle bad before + // returning it to database/sql. The set deduplicates a normal owner/generic + // pair and also covers owners whose generic map entry was already removed. + for conn := range conns { + if _, protected := protectedConns[conn]; protected { + continue + } + lastErr = errors.Join(lastErr, p.closePhysicalConnLocked(conn)) + } + + // Remove only entries that were validated and processed. Malformed entries + // stay in place so an operator/recovery path can identify and repair them; + // using Clear here would erase the evidence and could strand a live owner. + for id, tx := range validTxnByID { + p.txns.CompareAndDelete(id, tx) + } + for id, conn := range validOwnerByID { + if _, malformed := malformedTxnIDs[id]; malformed { + continue + } + p.txnConns.CompareAndDelete(id, conn) + } + for id, conn := range validConnByID { + if _, malformed := malformedTxnIDs[id]; malformed { + continue + } + if _, malformed := malformedOwnerIDs[id]; malformed { + continue + } + p.conns.CompareAndDelete(id, conn) + } + // A connection-less session lock has no lifecycle state after a successful + // sweep. Keep locks for malformed IDs so a later repair can synchronize with + // the preserved mapping; clearing unrelated locks avoids retaining a whole + // generation after shutdown. + p.sessionLocks.Range(func(key, value any) bool { + id, ok := key.(uint32) + if !ok { + return true + } + if _, bad := malformedTxnIDs[id]; bad { + return true + } + if _, bad := malformedConnIDs[id]; bad { + return true + } + if _, bad := malformedOwnerIDs[id]; bad { + return true + } + p.sessionLocks.CompareAndDelete(id, value) + return true + }) + if p.DB == nil { + return lastErr } - p.conns.Clear() - p.txns.Clear() - p.closedConns.Clear() return errors.Join(lastErr, p.DB.Close()) } func (p *ConnectionPool) Reset(connector *duckdb.Connector, db *stdsql.DB) error { - err := p.Close() + var completions transactionCompletionBatch + defer p.clearCompletedTransactionCompletions() + defer completions.run() + preCloseErr := p.closeDatabaseForShutdown() + p.lifecycleMu.LockAfterShutdown() + defer p.lifecycleMu.Unlock() + err := errors.Join( + preCloseErr, + p.closeExternalSessionsLocked(&completions), + p.closeLockedWithBatch(&completions), + p.externalSessionPanicError(), + ) if err != nil { return fmt.Errorf("failed to close connection pool: %w", err) } + p.transactionCompletionMu.Lock() + p.activeTransactions = make(map[*stdsql.Tx]struct{}) + p.transactionCompletionMu.Unlock() p.conns.Clear() p.txns.Clear() + p.txnConns.Clear() p.closedConns.Clear() + p.sessionLocks.Clear() + p.externalGeneration++ p.DB = db + p.dbPtr.Store(db) p.connector = connector p.registerMySQLUDFsOnce = sync.Once{} p.registerMySQLUDFsErr = nil diff --git a/catalog/connpool_test.go b/catalog/connpool_test.go index ad2629ab..6e99c359 100644 --- a/catalog/connpool_test.go +++ b/catalog/connpool_test.go @@ -4,9 +4,15 @@ import ( "context" stdsql "database/sql" "database/sql/driver" + "errors" + "fmt" + "io" "os" + "sync" "testing" + "time" + "github.com/DATA-DOG/go-sqlmock" "github.com/apecloud/myduckserver/mycontext" "github.com/apecloud/myduckserver/testutil" gmsql "github.com/dolthub/go-mysql-server/sql" @@ -187,6 +193,31 @@ func TestConnectionPoolInitializerSkipsActiveTransaction(t *testing.T) { }, origins) } +func TestGetConnForSchemaDoesNotSwitchActiveTransaction(t *testing.T) { + connector, err := duckdb.NewConnector("", nil) + require.NoError(t, err) + db := stdsql.OpenDB(connector) + pool := NewConnectionPool(connector, db, "memory") + t.Cleanup(func() { + require.NoError(t, pool.Close()) + require.NoError(t, connector.Close()) + }) + + conn, err := pool.GetConn(context.Background(), 77) + require.NoError(t, err) + tx, err := pool.GetTxn(context.Background(), 77, "", nil) + require.NoError(t, err) + + // The requested schema does not exist. While the transaction owns the + // connection, GetConnForSchema must return it without issuing USE through + // the separate *sql.Conn surface. + got, err := pool.GetConnForSchema(context.Background(), 77, "schema_that_does_not_exist") + require.NoError(t, err) + require.Same(t, conn, got) + require.NoError(t, tx.Rollback()) + pool.CloseTxn(77) +} + func TestCloseConnRollsBackTransaction(t *testing.T) { connector, err := duckdb.NewConnector("", nil) require.NoError(t, err) @@ -216,6 +247,810 @@ func TestCloseConnRollsBackTransaction(t *testing.T) { require.Zero(t, count) } +func TestConnectionPoolStaleFinalizerCannotCommitReplacement(t *testing.T) { + connector, err := duckdb.NewConnector("", nil) + require.NoError(t, err) + db := stdsql.OpenDB(connector) + pool := NewConnectionPool(connector, db, "memory") + t.Cleanup(func() { + require.NoError(t, pool.Close()) + require.NoError(t, connector.Close()) + }) + + first, err := pool.GetTxn(context.Background(), 91, "", nil) + require.NoError(t, err) + require.NoError(t, first.Rollback()) + pool.CloseTxn(91) + second, err := pool.GetTxn(context.Background(), 91, "", nil) + require.NoError(t, err) + + // A delayed callback for the old transaction must not finalize the new one. + require.NoError(t, pool.CommitTxn(91, first)) + require.Same(t, second, pool.TryGetTxn(91)) + require.NoError(t, second.Rollback()) + pool.CloseTxn(91) +} + +func TestConnectionPoolRollbackErrorEvictsPhysicalOwner(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + pool := NewConnectionPool(nil, db, "memory") + pool.registerMySQLUDFsOnce.Do(func() {}) + t.Cleanup(func() { _ = db.Close() }) + + mock.ExpectBegin() + tx, err := pool.GetTxn(context.Background(), 92, "", nil) + require.NoError(t, err) + rollbackErr := errors.New("driver rollback state is unknown") + mock.ExpectRollback().WillReturnError(rollbackErr) + + _, inactive, err := pool.RollbackTxn(92, tx) + require.False(t, inactive) + require.ErrorIs(t, err, rollbackErr) + require.Nil(t, pool.TryGetTxn(92)) + _, stillMapped := pool.conns.Load(uint32(92)) + require.False(t, stillMapped, "rollback errors must evict the physical owner") + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestConnectionPoolActiveTransactionDoesNotFallbackToSessionConnection(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + p := NewConnectionPool(nil, db, "memory") + p.registerMySQLUDFsOnce.Do(func() {}) + t.Cleanup(func() { _ = db.Close() }) + + const sessionID = uint32(93) + mock.ExpectBegin() + tx, err := p.GetTxn(context.Background(), sessionID, "", nil) + require.NoError(t, err) + connEntry, ok := p.conns.Load(sessionID) + require.True(t, ok) + require.NotNil(t, connEntry) + + // Simulate a partially lost transaction-owner mapping while the generic + // session connection is still present. That generic entry is not proof of + // ownership and must never be used for active transaction work. + p.txnConns.Delete(sessionID) + owner, bound := p.GetTxnBinding(sessionID) + require.Nil(t, owner) + require.Same(t, tx, bound) + + // Every connection-facing lifecycle path must reject the ownerless active + // transaction instead of falling back to the generic session connection. + conn, err := p.GetConn(context.Background(), sessionID) + require.ErrorContains(t, err, "no physical owner") + require.Nil(t, conn) + conn, err = p.GetConnForSchema(context.Background(), sessionID, "main") + require.ErrorContains(t, err, "no physical owner") + require.Nil(t, conn) + gotTx, err := p.GetTxn(context.Background(), sessionID, "", nil) + require.ErrorContains(t, err, "no physical owner") + require.Nil(t, gotTx) + require.Empty(t, p.CurrentSchema(sessionID)) + require.Empty(t, p.CurrentCatalog(sessionID)) + + snapshotConn, snapshotTx, err := p.GetExecutionSnapshot( + context.Background(), sessionID, "", false, + ) + require.ErrorContains(t, err, "no physical owner") + require.Nil(t, snapshotConn) + require.Same(t, tx, snapshotTx) + + mock.ExpectRollback() + mock.ExpectClose() + require.NoError(t, p.Close()) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestConnectionPoolCloseConnSweepsOwnerWhenGenericMappingIsAbsent(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + p := NewConnectionPool(nil, db, "memory") + p.registerMySQLUDFsOnce.Do(func() {}) + t.Cleanup(func() { _ = db.Close() }) + + const sessionID = uint32(99) + mock.ExpectBegin() + _, err = p.GetTxn(context.Background(), sessionID, "", nil) + require.NoError(t, err) + ownerEntry, ok := p.txnConns.Load(sessionID) + require.True(t, ok) + owner, ok := ownerEntry.(*stdsql.Conn) + require.True(t, ok) + require.NotNil(t, owner) + + // A failed reconnect/identity-safe close can remove the generic session map + // while the active transaction still owns its physical connection. An + // unconditional session close must roll back and retire that owner instead + // of only deleting the logical transaction marker. + p.conns.Delete(sessionID) + mock.ExpectRollback() + mock.ExpectClose() + require.NoError(t, p.CloseConn(sessionID)) + + _, present := p.conns.Load(sessionID) + require.False(t, present) + _, present = p.txns.Load(sessionID) + require.False(t, present) + _, present = p.txnConns.Load(sessionID) + require.False(t, present) + require.Nil(t, p.TryGetTxn(sessionID)) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestConnectionPoolCloseSweepsOwnerWithoutGenericMapping(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + p := NewConnectionPool(nil, db, "memory") + p.registerMySQLUDFsOnce.Do(func() {}) + t.Cleanup(func() { _ = db.Close() }) + + const sessionID = uint32(100) + mock.ExpectBegin() + _, err = p.GetTxn(context.Background(), sessionID, "", nil) + require.NoError(t, err) + ownerEntry, ok := p.txnConns.Load(sessionID) + require.True(t, ok) + owner, ok := ownerEntry.(*stdsql.Conn) + require.True(t, ok) + require.NotNil(t, owner) + p.conns.Delete(sessionID) + + // closeLocked must sweep txnConns independently of conns. Otherwise this + // checked-out owner survives pool teardown and can retain an open driver + // transaction. + mock.ExpectRollback() + mock.ExpectClose() + require.NoError(t, p.Close()) + _, present := p.conns.Load(sessionID) + require.False(t, present) + _, present = p.txns.Load(sessionID) + require.False(t, present) + _, present = p.txnConns.Load(sessionID) + require.False(t, present) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestConnectionPoolActiveTransactionUsesPhysicalOwnerOverReplacementSessionConn(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + p := NewConnectionPool(nil, db, "memory") + p.registerMySQLUDFsOnce.Do(func() {}) + t.Cleanup(func() { _ = db.Close() }) + + const sessionID = uint32(94) + mock.ExpectBegin() + tx, err := p.GetTxn(context.Background(), sessionID, "", nil) + require.NoError(t, err) + ownerEntry, ok := p.txnConns.Load(sessionID) + require.True(t, ok) + owner, ok := ownerEntry.(*stdsql.Conn) + require.True(t, ok) + require.NotNil(t, owner) + + // Simulate a stale generic session mapping left behind by a reconnect. The + // active transaction must continue to use the owner captured at BeginTx. + replacement, err := db.Conn(context.Background()) + require.NoError(t, err) + p.conns.Store(sessionID, replacement) + got, err := p.GetConn(context.Background(), sessionID) + require.NoError(t, err) + require.Same(t, owner, got) + got, err = p.GetConnForSchema(context.Background(), sessionID, "schema_that_does_not_exist") + require.NoError(t, err) + require.Same(t, owner, got) + + // Restore the normal map before pool teardown and release the synthetic + // replacement separately; it is not part of the transaction binding. + p.conns.Store(sessionID, owner) + require.NoError(t, replacement.Close()) + mock.ExpectRollback() + require.NoError(t, tx.Rollback()) + p.CloseTxn(sessionID) + mock.ExpectClose() + mock.ExpectClose() + require.NoError(t, p.Close()) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestConnectionPoolMalformedMappingsFailClosed(t *testing.T) { + connector, err := duckdb.NewConnector("", nil) + require.NoError(t, err) + db := stdsql.OpenDB(connector) + p := NewConnectionPool(connector, db, "memory") + + // A malformed generic entry must never be treated as an uninitialized + // session (which would re-open a second connection) or panic a metadata + // lookup. CurrentCatalog/CurrentSchema have no error return, so they fail + // closed with an empty value; connection acquisition reports the corruption. + const sessionID = uint32(95) + p.conns.Store(sessionID, "not-a-connection") + require.NotPanics(t, func() { + require.Empty(t, p.CurrentSchema(sessionID)) + require.Empty(t, p.CurrentCatalog(sessionID)) + }) + _, err = p.GetConn(context.Background(), sessionID) + require.ErrorContains(t, err, "connection mapping is invalid") + + // A malformed active transaction mapping is likewise not idle and must not + // fall back to the generic connection path. CloseTxn can discard the poisoned + // marker without attempting a type-unsafe rollback. + const poisonedID = uint32(96) + p.txns.Store(poisonedID, "not-a-transaction") + p.txnConns.Store(poisonedID, "not-a-connection") + require.Nil(t, p.TryGetTxn(poisonedID)) + _, err = p.GetConn(context.Background(), poisonedID) + require.ErrorContains(t, err, "active transaction mapping is invalid") + _, _, err = p.GetExecutionSnapshot(context.Background(), poisonedID, "", false) + require.ErrorContains(t, err, "active transaction mapping is invalid") + p.CloseTxn(poisonedID) + _, stillPoisoned := p.txns.Load(poisonedID) + require.False(t, stillPoisoned) + + // Close helpers return a concrete error for malformed generic entries and + // never close an unrelated caller-supplied connection. + conn, err := db.Conn(context.Background()) + require.NoError(t, err) + const closeID = uint32(97) + p.conns.Store(closeID, "not-a-connection") + require.ErrorContains(t, p.CloseConnIf(closeID, conn), "connection mapping is invalid") + require.ErrorContains(t, p.CloseConnIfBinding(closeID, conn, nil), "connection mapping is invalid") + require.NoError(t, conn.Close()) + + // Pool teardown remains panic-free even when a map was poisoned. It reports + // the malformed entry so the caller can surface the lifecycle corruption. + const teardownID = uint32(98) + p.conns.Store(teardownID, "not-a-connection") + require.ErrorContains(t, p.Close(), "connection mapping is invalid") + require.NoError(t, connector.Close()) +} + +func TestConnectionPoolMalformedTransactionDoesNotCloseLiveOwner(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + p := NewConnectionPool(nil, db, "memory") + p.registerMySQLUDFsOnce.Do(func() {}) + + const sessionID = uint32(109) + mock.ExpectBegin() + tx, err := p.GetTxn(context.Background(), sessionID, "", nil) + require.NoError(t, err) + ownerEntry, ok := p.txnConns.Load(sessionID) + require.True(t, ok) + owner, ok := ownerEntry.(*stdsql.Conn) + require.True(t, ok) + require.NotNil(t, owner) + + // Replace only the logical marker. The physical transaction is still live, + // so closing its owner would wait forever on database/sql's close mutex. + p.txns.Store(sessionID, "corrupt transaction marker") + closeDone := make(chan error, 1) + go func() { closeDone <- p.CloseConn(sessionID) }() + select { + case closeErr := <-closeDone: + require.ErrorContains(t, closeErr, "active transaction mapping is invalid") + case <-time.After(time.Second): + t.Fatal("CloseConn hung while transaction mapping was malformed") + } + currentOwner, ok := p.txnConns.Load(sessionID) + require.True(t, ok) + require.Same(t, owner, currentOwner) + currentConn, ok := p.conns.Load(sessionID) + require.True(t, ok) + require.Same(t, owner, currentConn) + + // Restore a typed marker so the test can finish the live transaction and + // release the sqlmock connection normally. + p.txns.Store(sessionID, tx) + mock.ExpectRollback() + require.NoError(t, tx.Rollback()) + p.CloseTxn(sessionID) + mock.ExpectClose() + require.NoError(t, p.CloseConn(sessionID)) + require.NoError(t, mock.ExpectationsWereMet()) + _ = db.Close() +} + +func TestConnectionPoolCloseReportsMalformedTransactionWithoutBlocking(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + p := NewConnectionPool(nil, db, "memory") + p.registerMySQLUDFsOnce.Do(func() {}) + + const sessionID = uint32(110) + mock.ExpectBegin() + tx, err := p.GetTxn(context.Background(), sessionID, "", nil) + require.NoError(t, err) + ownerEntry, ok := p.txnConns.Load(sessionID) + require.True(t, ok) + owner, ok := ownerEntry.(*stdsql.Conn) + require.True(t, ok) + + p.txns.Store(sessionID, "corrupt transaction marker") + closeDone := make(chan error, 1) + go func() { closeDone <- p.Close() }() + select { + case closeErr := <-closeDone: + require.ErrorContains(t, closeErr, "active transaction mapping is invalid") + case <-time.After(time.Second): + t.Fatal("pool Close hung while transaction mapping was malformed") + } + currentOwner, ok := p.txnConns.Load(sessionID) + require.True(t, ok) + require.Same(t, owner, currentOwner) + + // The failed teardown leaves the live owner available for an explicit + // recovery, after which a normal close can retire it. + p.txns.Store(sessionID, tx) + mock.ExpectRollback() + require.NoError(t, tx.Rollback()) + p.CloseTxn(sessionID) + mock.ExpectClose() + require.NoError(t, p.CloseConn(sessionID)) + require.NoError(t, mock.ExpectationsWereMet()) + _ = db.Close() +} + +func TestConnectionPoolSingleConnectionReleaseDoesNotDeadlockOtherSession(t *testing.T) { + connector, err := duckdb.NewConnector("", nil) + require.NoError(t, err) + db := stdsql.OpenDB(connector) + db.SetMaxOpenConns(1) + p := NewConnectionPool(connector, db, "memory") + t.Cleanup(func() { + require.NoError(t, p.Close()) + require.NoError(t, connector.Close()) + }) + + conn1, err := p.GetConn(context.Background(), 101) + require.NoError(t, err) + tx, err := p.GetTxn(context.Background(), 101, "", nil) + require.NoError(t, err) + + acquired := make(chan *stdsql.Conn, 1) + acquireErr := make(chan error, 1) + go func() { + conn2, getErr := p.GetConn(context.Background(), 102) + if getErr != nil { + acquireErr <- getErr + return + } + acquired <- conn2 + }() + select { + case <-acquired: + t.Fatal("second session acquired the only connection before release") + case <-time.After(20 * time.Millisecond): + } + + require.NoError(t, tx.Rollback()) + p.CloseTxn(101) + require.NoError(t, p.CloseConn(101)) + select { + case err := <-acquireErr: + require.NoError(t, err) + case conn2 := <-acquired: + require.NotNil(t, conn2) + require.NoError(t, p.CloseConn(102)) + case <-time.After(time.Second): + t.Fatal("second session remained blocked after first session release") + } + _ = conn1 +} + +func TestConnectionPoolCloseUnblocksPendingConnectionAcquire(t *testing.T) { + connector, err := duckdb.NewConnector("", nil) + require.NoError(t, err) + db := stdsql.OpenDB(connector) + db.SetMaxOpenConns(1) + p := NewConnectionPool(connector, db, "memory") + + // Keep the sole physical connection checked out. The second acquisition + // enters database/sql's pending request queue while the pool read lock is + // held, which is the shutdown deadlock regression this test guards. + _, err = p.GetConn(context.Background(), 111) + require.NoError(t, err) + + acquireDone := make(chan error, 1) + go func() { + _, acquireErr := p.GetConn(context.Background(), 112) + acquireDone <- acquireErr + }() + // Wait after starting the goroutine so this test does not race the setup + // connection with the pending request. + deadline := time.Now().Add(time.Second) + for db.Stats().WaitCount < 1 && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + require.GreaterOrEqual(t, db.Stats().WaitCount, int64(1), "second acquisition did not reach database/sql wait queue") + + closeDone := make(chan error, 1) + go func() { closeDone <- p.Close() }() + select { + case closeErr := <-closeDone: + require.NoError(t, closeErr) + case <-time.After(time.Second): + t.Fatal("pool close remained blocked behind a pending connection acquire") + } + select { + case acquireErr := <-acquireDone: + require.Error(t, acquireErr) + case <-time.After(time.Second): + t.Fatal("pending connection acquire was not canceled by pool close") + } +} + +func TestConnectionPoolExecutionSnapshotLeaseBlocksClose(t *testing.T) { + connector, err := duckdb.NewConnector("", nil) + require.NoError(t, err) + db := stdsql.OpenDB(connector) + p := NewConnectionPool(connector, db, "memory") + + conn, tx, release, err := p.GetExecutionSnapshotLease( + context.Background(), 113, "", true, + ) + require.NoError(t, err) + require.NotNil(t, conn) + require.Nil(t, tx) + require.NotNil(t, release) + t.Cleanup(release) + + closeDone := make(chan error, 1) + go func() { closeDone <- p.Close() }() + waitForPoolShutdown(t, p) + + // Shutdown is announced, but the physical generation must remain alive while + // the snapshot holder still owns its lease. + select { + case closeErr := <-closeDone: + t.Fatalf("pool closed before snapshot lease release: %v", closeErr) + case <-time.After(100 * time.Millisecond): + } + _, mapped := p.conns.Load(uint32(113)) + require.True(t, mapped, "shutdown must not retire a leased connection") + + // Release is intentionally safe to call more than once. The first call lets + // teardown proceed; the second must not unlock the lifecycle gate again. + release() + release() + select { + case closeErr := <-closeDone: + require.NoError(t, closeErr) + case <-time.After(time.Second): + t.Fatal("pool close did not finish after snapshot lease release") + } + _, mapped = p.conns.Load(uint32(113)) + require.False(t, mapped) + require.NoError(t, connector.Close()) +} + +func TestConnectionPoolExecutionSnapshotLeaseBlocksReset(t *testing.T) { + oldConnector, err := duckdb.NewConnector("", nil) + require.NoError(t, err) + oldDB := stdsql.OpenDB(oldConnector) + p := NewConnectionPool(oldConnector, oldDB, "memory") + + conn, tx, release, err := p.GetExecutionSnapshotLease( + context.Background(), 114, "", true, + ) + require.NoError(t, err) + require.NotNil(t, conn) + require.Nil(t, tx) + t.Cleanup(release) + + newConnector, err := duckdb.NewConnector("", nil) + require.NoError(t, err) + newDB := stdsql.OpenDB(newConnector) + resetDone := make(chan error, 1) + go func() { resetDone <- p.Reset(newConnector, newDB) }() + waitForPoolShutdown(t, p) + + select { + case resetErr := <-resetDone: + t.Fatalf("pool reset retired a leased generation: %v", resetErr) + case <-time.After(100 * time.Millisecond): + } + _, mapped := p.conns.Load(uint32(114)) + require.True(t, mapped, "reset must wait for the leased connection") + + release() + release() + select { + case resetErr := <-resetDone: + require.NoError(t, resetErr) + case <-time.After(time.Second): + t.Fatal("pool reset did not finish after snapshot lease release") + } + + // The replacement generation is usable after the old leased generation has + // drained. + newConn, err := p.GetConn(context.Background(), 115) + require.NoError(t, err) + var got int + require.NoError(t, newConn.QueryRowContext(context.Background(), "SELECT 1").Scan(&got)) + require.Equal(t, 1, got) + require.NoError(t, p.CloseConn(115)) + require.NoError(t, p.Close()) + require.NoError(t, oldConnector.Close()) + require.NoError(t, newConnector.Close()) +} + +func TestConnectionPoolResetUnblocksPendingConnectionAcquire(t *testing.T) { + oldConnector, err := duckdb.NewConnector("", nil) + require.NoError(t, err) + oldDB := stdsql.OpenDB(oldConnector) + oldDB.SetMaxOpenConns(1) + p := NewConnectionPool(oldConnector, oldDB, "memory") + + _, err = p.GetConn(context.Background(), 121) + require.NoError(t, err) + acquireDone := make(chan error, 1) + go func() { + _, acquireErr := p.GetConn(context.Background(), 122) + acquireDone <- acquireErr + }() + deadline := time.Now().Add(time.Second) + for oldDB.Stats().WaitCount < 1 && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + require.GreaterOrEqual(t, oldDB.Stats().WaitCount, int64(1), "second acquisition did not reach database/sql wait queue") + + newConnector, err := duckdb.NewConnector("", nil) + require.NoError(t, err) + newDB := stdsql.OpenDB(newConnector) + resetDone := make(chan error, 1) + go func() { resetDone <- p.Reset(newConnector, newDB) }() + select { + case resetErr := <-resetDone: + require.NoError(t, resetErr) + case <-time.After(time.Second): + t.Fatal("pool reset remained blocked behind a pending connection acquire") + } + select { + case acquireErr := <-acquireDone: + require.Error(t, acquireErr) + case <-time.After(time.Second): + t.Fatal("pending connection acquire was not canceled by pool reset") + } + + // Reset must leave the replacement generation usable after the old DB was + // pre-closed to release the blocked request. + conn, err := p.GetConn(context.Background(), 123) + require.NoError(t, err) + var got int + require.NoError(t, conn.QueryRowContext(context.Background(), "SELECT 1").Scan(&got)) + require.Equal(t, 1, got) + require.NoError(t, p.CloseConn(123)) + require.NoError(t, p.Close()) +} + +func TestConnectionPoolCloseRollsBackActiveTransactionAfterPreClose(t *testing.T) { + connector, err := duckdb.NewConnector("", nil) + require.NoError(t, err) + db := stdsql.OpenDB(connector) + p := NewConnectionPool(connector, db, "memory") + conn, err := p.GetConn(context.Background(), 131) + require.NoError(t, err) + _, err = conn.ExecContext(context.Background(), "CREATE TABLE close_preclose_tx (id INTEGER)") + require.NoError(t, err) + tx, err := p.GetTxn(context.Background(), 131, "", nil) + require.NoError(t, err) + _, err = tx.ExecContext(context.Background(), "INSERT INTO close_preclose_tx VALUES (1)") + require.NoError(t, err) + + require.NoError(t, p.Close()) + _, present := p.txns.Load(uint32(131)) + require.False(t, present) +} + +func TestConnectionPoolShutdownAllowsReleaseWhileSchemaReadIsBlocked(t *testing.T) { + driverState := newBlockingSchemaDriver() + db := stdsql.OpenDB(driverState) + db.SetMaxOpenConns(2) + p := NewConnectionPool(nil, db, "memory") + p.registerMySQLUDFsOnce.Do(func() {}) + + conn1, err := p.GetConn(context.Background(), 141) + require.NoError(t, err) + conn2, err := p.GetConn(context.Background(), 142) + require.NoError(t, err) + require.NotNil(t, conn1) + require.NotNil(t, conn2) + + schemaDone := make(chan struct{}) + go func() { + _ = p.CurrentSchema(141) + close(schemaDone) + }() + select { + case <-driverState.schemaStarted: + case <-time.After(time.Second): + t.Fatal("schema query did not start") + } + + closeDone := make(chan error, 1) + go func() { closeDone <- p.Close() }() + waitForPoolShutdown(t, p) + + // Close is waiting for the blocked schema reader. A session release must + // still be admitted while that writer is pending; an RWMutex would block it + // behind the queued writer and leave both goroutines waiting. + releaseDone := make(chan error, 1) + go func() { releaseDone <- p.CloseConn(142) }() + select { + case releaseErr := <-releaseDone: + require.NoError(t, releaseErr) + case <-time.After(250 * time.Millisecond): + t.Fatal("connection release remained blocked by pending pool shutdown") + } + + close(driverState.releaseSchema) + select { + case <-schemaDone: + case <-time.After(time.Second): + t.Fatal("schema reader did not finish after release") + } + select { + case closeErr := <-closeDone: + require.NoError(t, closeErr) + case <-time.After(time.Second): + t.Fatal("pool close did not finish after schema reader") + } +} + +func TestConnectionPoolResetAllowsReleaseWhileSchemaReadIsBlocked(t *testing.T) { + oldDriver := newBlockingSchemaDriver() + oldDB := stdsql.OpenDB(oldDriver) + oldDB.SetMaxOpenConns(2) + p := NewConnectionPool(nil, oldDB, "memory") + p.registerMySQLUDFsOnce.Do(func() {}) + + _, err := p.GetConn(context.Background(), 151) + require.NoError(t, err) + _, err = p.GetConn(context.Background(), 152) + require.NoError(t, err) + + schemaDone := make(chan struct{}) + go func() { + _ = p.CurrentSchema(151) + close(schemaDone) + }() + select { + case <-oldDriver.schemaStarted: + case <-time.After(time.Second): + t.Fatal("schema query did not start") + } + + newDriver := newBlockingSchemaDriver() + newDB := stdsql.OpenDB(newDriver) + resetDone := make(chan error, 1) + go func() { resetDone <- p.Reset(nil, newDB) }() + waitForPoolShutdown(t, p) + + releaseDone := make(chan error, 1) + go func() { releaseDone <- p.CloseConn(152) }() + select { + case releaseErr := <-releaseDone: + require.NoError(t, releaseErr) + case <-time.After(250 * time.Millisecond): + t.Fatal("connection release remained blocked by pending pool reset") + } + + close(oldDriver.releaseSchema) + select { + case <-schemaDone: + case <-time.After(time.Second): + t.Fatal("schema reader did not finish after release") + } + select { + case resetErr := <-resetDone: + require.NoError(t, resetErr) + case <-time.After(time.Second): + t.Fatal("pool reset did not finish after schema reader") + } + + p.registerMySQLUDFsOnce.Do(func() {}) + _, err = p.GetConn(context.Background(), 153) + require.NoError(t, err) + require.NoError(t, p.CloseConn(153)) + require.NoError(t, p.Close()) +} + +func TestConnectionPoolReleaseBindingSnapshotAllowsShutdownProgress(t *testing.T) { + driverState := newBlockingSchemaDriver() + db := stdsql.OpenDB(driverState) + db.SetMaxOpenConns(3) + p := NewConnectionPool(nil, db, "memory") + p.registerMySQLUDFsOnce.Do(func() {}) + + // Keep one reader blocked in a driver call while the second session is the + // release target. The release snapshot must be admitted after shutdown has + // been announced; using the ordinary lifecycle read lock here would queue it + // behind Close and deadlock the pool teardown. + _, err := p.GetConn(context.Background(), 161) + require.NoError(t, err) + releaseConn, err := p.GetConn(context.Background(), 162) + require.NoError(t, err) + require.NotNil(t, releaseConn) + + schemaDone := make(chan struct{}) + go func() { + _ = p.CurrentSchema(161) + close(schemaDone) + }() + select { + case <-driverState.schemaStarted: + case <-time.After(time.Second): + t.Fatal("schema query did not start") + } + + closeDone := make(chan error, 1) + go func() { closeDone <- p.Close() }() + waitForPoolShutdown(t, p) + ordinarySnapshotDone := make(chan struct{}) + go func() { + _, _ = p.GetTxnBinding(162) + close(ordinarySnapshotDone) + }() + select { + case <-ordinarySnapshotDone: + t.Fatal("ordinary binding snapshot entered after pool shutdown") + case <-time.After(100 * time.Millisecond): + } + + releaseSnapshotDone := make(chan error, 1) + go func() { + conn, tx := p.GetTxnBindingForRelease(162) + if conn == nil || tx != nil { + releaseSnapshotDone <- fmt.Errorf("unexpected release binding conn=%v tx=%v", conn, tx) + return + } + releaseSnapshotDone <- p.CloseConnIfBinding(162, conn, nil) + }() + select { + case releaseErr := <-releaseSnapshotDone: + require.NoError(t, releaseErr) + case <-time.After(250 * time.Millisecond): + t.Fatal("release binding snapshot remained blocked by pending pool shutdown") + } + + close(driverState.releaseSchema) + select { + case <-schemaDone: + case <-time.After(time.Second): + t.Fatal("schema reader did not finish after release") + } + select { + case closeErr := <-closeDone: + require.NoError(t, closeErr) + case <-time.After(time.Second): + t.Fatal("pool close did not finish after schema reader") + } + select { + case <-ordinarySnapshotDone: + case <-time.After(time.Second): + t.Fatal("ordinary binding snapshot did not resume after pool shutdown") + } +} + +func waitForPoolShutdown(t *testing.T, p *ConnectionPool) { + t.Helper() + deadline := time.Now().Add(time.Second) + for { + p.lifecycleMu.mu.Lock() + shuttingDown := p.lifecycleMu.shuttingDown + p.lifecycleMu.mu.Unlock() + if shuttingDown { + return + } + if time.Now().After(deadline) { + t.Fatal("pool lifecycle did not announce shutdown") + } + time.Sleep(time.Millisecond) + } +} + func TestFirstMySQLQueryAfterServerStart(t *testing.T) { originalWorkingDir, err := os.Getwd() require.NoError(t, err) @@ -253,3 +1088,91 @@ func assertMySQLUDFsCallable(t *testing.T, conn *stdsql.Conn) { require.NoError(t, err) require.Equal(t, []float32{1.5, -2}, decoded) } + +// blockingSchemaDriver gives the shutdown regression a deterministic blocked +// driver query and a separate connection that can be released while Close is +// waiting for the blocked reader. +type blockingSchemaDriver struct { + mu sync.Mutex + nextID int + schemaStarted chan struct{} + releaseSchema chan struct{} + startedOnce sync.Once + connections map[int]*blockingSchemaConn +} + +func newBlockingSchemaDriver() *blockingSchemaDriver { + return &blockingSchemaDriver{ + schemaStarted: make(chan struct{}), + releaseSchema: make(chan struct{}), + connections: make(map[int]*blockingSchemaConn), + } +} + +func (d *blockingSchemaDriver) Connect(context.Context) (driver.Conn, error) { + d.mu.Lock() + defer d.mu.Unlock() + id := d.nextID + d.nextID++ + conn := &blockingSchemaConn{driver: d, id: id} + d.connections[id] = conn + return conn, nil +} + +func (d *blockingSchemaDriver) Driver() driver.Driver { return d } + +func (d *blockingSchemaDriver) Open(string) (driver.Conn, error) { + return d.Connect(context.Background()) +} + +type blockingSchemaConn struct { + driver *blockingSchemaDriver + id int + mu sync.Mutex + closed bool +} + +func (c *blockingSchemaConn) Prepare(string) (driver.Stmt, error) { + return nil, fmt.Errorf("prepare is unsupported") +} + +func (c *blockingSchemaConn) Close() error { + c.mu.Lock() + c.closed = true + c.mu.Unlock() + return nil +} + +func (c *blockingSchemaConn) Begin() (driver.Tx, error) { + return nil, fmt.Errorf("transactions are unsupported") +} + +func (c *blockingSchemaConn) QueryContext(ctx context.Context, query string, _ []driver.NamedValue) (driver.Rows, error) { + if query != "SELECT CURRENT_SCHEMA" { + return nil, fmt.Errorf("unexpected query %q", query) + } + c.driver.startedOnce.Do(func() { close(c.driver.schemaStarted) }) + select { + case <-c.driver.releaseSchema: + return &blockingSchemaRows{}, nil + case <-ctx.Done(): + return nil, ctx.Err() + } +} + +type blockingSchemaRows struct { + read bool +} + +func (*blockingSchemaRows) Columns() []string { return []string{"current_schema"} } + +func (r *blockingSchemaRows) Close() error { return nil } + +func (r *blockingSchemaRows) Next(dest []driver.Value) error { + if r.read { + return io.EOF + } + r.read = true + dest[0] = "main" + return nil +} diff --git a/catalog/database.go b/catalog/database.go index 72aa56c9..ffa38f61 100644 --- a/catalog/database.go +++ b/catalog/database.go @@ -2,6 +2,7 @@ package catalog import ( stdsql "database/sql" + "errors" "fmt" "strings" "sync" @@ -18,9 +19,10 @@ import ( ) type Database struct { - mu *sync.RWMutex - catalog string - name string + mu *sync.RWMutex + catalog string + name string + provider *DatabaseProvider } type ExtraViewInfo struct { @@ -52,10 +54,23 @@ func vectorGeneratedExpression(col *sql.Column, typ types.VectorType) (string, e } func NewDatabase(name string, catalogName string) *Database { + return newDatabase(name, catalogName, nil) +} + +// NewDatabaseWithProvider constructs a logical database bound to the owning +// provider. Protocol adapters that bypass DatabaseProvider.Database (notably +// PostgreSQL's direct parser path) must use this form so object-table +// materialization and physical-name resolution retain the DuckLake runtime. +func NewDatabaseWithProvider(name string, catalogName string, provider *DatabaseProvider) *Database { + return newDatabase(name, catalogName, provider) +} + +func newDatabase(name string, catalogName string, provider *DatabaseProvider) *Database { return &Database{ - mu: &sync.RWMutex{}, - name: name, - catalog: catalogName, + mu: &sync.RWMutex{}, + name: name, + catalog: catalogName, + provider: provider, } } @@ -146,7 +161,39 @@ func (d *Database) Name() string { return d.name } -func (d *Database) createAllTable(ctx *sql.Context, name string, schema sql.PrimaryKeySchema, collation sql.CollationID, comment string, temporary bool) error { +func (d *Database) createAllTable(ctx *sql.Context, name string, schema sql.PrimaryKeySchema, collation sql.CollationID, comment string, storage TableStorageSelection, temporary bool) error { + if err := storage.Validate(); err != nil { + return err + } + if temporary && storage.IsObjectStorage() { + return fmt.Errorf("%w: temporary tables cannot use object storage", ErrInvalidTableStorage) + } + if storage.Kind == "" { + storage = DefaultTableStorageSelection() + } + if storage.IsObjectStorage() { + return d.createObjectTable(ctx, name, schema, collation, comment, storage) + } + return d.createLocalTable(ctx, name, schema, collation, comment, storage, temporary) +} + +// createLocalTable materializes the managed shadow table. Object tables use +// the same path after their physical DuckLake relation has been created; the +// shadow is intentionally empty and exists only for GMS schema/comment +// discovery. +func (d *Database) createLocalTable(ctx *sql.Context, name string, schema sql.PrimaryKeySchema, collation sql.CollationID, comment string, storage TableStorageSelection, temporary bool) error { + execer, _, _, release, err := adapter.GetExecutionSnapshotWithLease(ctx) + if err != nil { + return err + } + defer release() + return d.createLocalTableWithExecutor(ctx, name, schema, collation, comment, storage, temporary, execer) +} + +func (d *Database) createLocalTableWithExecutor(ctx *sql.Context, name string, schema sql.PrimaryKeySchema, collation sql.CollationID, comment string, storage TableStorageSelection, temporary bool, execer adapter.SQLExecutor) error { + if execer == nil { + return fmt.Errorf("table creation executor is unavailable") + } var columns []string var columnCommentSQLs []string var fullTableName string @@ -269,7 +316,13 @@ func (d *Database) createAllTable(ctx *sql.Context, name string, schema sql.Prim b.WriteString(")") // Add comment to the table - info := ExtraTableInfo{schema.PkOrdinals, withoutIndex, fullSequenceName, nil} + info := ExtraTableInfo{ + PkOrdinals: schema.PkOrdinals, + Replicated: withoutIndex, + Sequence: fullSequenceName, + Checks: nil, + Storage: storage.Kind, + } b.WriteString(fmt.Sprintf( "; COMMENT ON TABLE %s IS '%s'", fullTableName, @@ -288,7 +341,7 @@ func (d *Database) createAllTable(ctx *sql.Context, name string, schema sql.Prim logger.WithField("DuckSQL", ddl).Debug("Executing DDL") } - _, err := adapter.Exec(ctx, ddl) + _, err := execer.ExecContext(ctx, ddl) if err != nil { if IsDuckDBTableAlreadyExistsError(err) { return sql.ErrTableAlreadyExists.New(name) @@ -322,14 +375,162 @@ func isIndexCreationDisabled(ctx *sql.Context) bool { func (d *Database) CreateTable(ctx *sql.Context, name string, schema sql.PrimaryKeySchema, collation sql.CollationID, comment string) error { d.mu.Lock() defer d.mu.Unlock() - return d.createAllTable(ctx, name, schema, collation, comment, false) + storage := DefaultTableStorageSelection() + if selected, ok := TableStorageSelectionFromContext(ctx); ok { + storage = selected + } + return d.createAllTable(ctx, name, schema, collation, comment, storage, false) +} + +// CreateTableWithStorage is the explicit catalog boundary for protocol +// adapters that already normalized a table selector. It is intentionally +// limited to selection propagation and metadata; object-table physical +// routing is owned by the follow-up storage implementation. +func (d *Database) CreateTableWithStorage(ctx *sql.Context, name string, schema sql.PrimaryKeySchema, collation sql.CollationID, comment string, storage TableStorageSelection) error { + if err := storage.Validate(); err != nil { + return err + } + if err := SetTableStorageSelection(ctx, storage); err != nil { + return err + } + d.mu.Lock() + defer d.mu.Unlock() + return d.createAllTable(ctx, name, schema, collation, comment, storage, false) +} + +// RecordTableStorageSelection updates the managed table metadata for a table +// created by a protocol path that bypasses sql.TableCreator (currently the +// PostgreSQL handler). It preserves the user-visible table comment and makes +// the selection available after a fresh catalog reload. +func (d *Database) RecordTableStorageSelection(ctx *sql.Context, name string, storage TableStorageSelection) error { + if storage.IsObjectStorage() { + _, activeTx := adapter.TryGetTxnBinding(ctx) + if activeTx == nil { + conn, tx, release, err := adapter.GetCatalogTxnExecutionSnapshotWithLease(ctx, nil) + if err != nil { + return err + } + defer release() + execer := adapter.SQLExecutor(tx) + operationErr := d.recordTableStorageSelectionWithExecutor(ctx, name, storage, execer, conn) + if operationErr != nil { + _, _, rollbackErr := adapter.FinalizeRollback(ctx, tx) + return errors.Join(operationErr, rollbackErr) + } + return adapter.FinalizeCommit(ctx, tx) + } + } + execer, conn, _, release, err := adapter.GetCatalogExecutionSnapshotWithLease(ctx) + if err != nil { + return err + } + defer release() + return d.recordTableStorageSelectionWithExecutor(ctx, name, storage, execer, conn) +} + +// RecordTableStorageSelectionWithExecutor updates managed table metadata +// through an executor/connection pair that the caller already captured. PG +// protocol execution holds that snapshot across the statement (and often its +// result iterator), so reacquiring it here can deadlock behind a pending pool +// shutdown writer. Callers using this method own any surrounding transaction; +// this method deliberately does not open a second scoped transaction. +func (d *Database) RecordTableStorageSelectionWithExecutor( + ctx *sql.Context, + name string, + storage TableStorageSelection, + execer adapter.SQLExecutor, + conn *stdsql.Conn, +) error { + return d.recordTableStorageSelectionWithExecutor(ctx, name, storage, execer, conn) +} + +func (d *Database) recordTableStorageSelectionWithExecutor( + ctx *sql.Context, + name string, + storage TableStorageSelection, + execer adapter.SQLExecutor, + conn *stdsql.Conn, +) error { + if err := storage.Validate(); err != nil { + return err + } + if d.catalog == "temp" && storage.IsObjectStorage() { + return fmt.Errorf("%w: temporary tables cannot use object storage", ErrInvalidTableStorage) + } + // Keep the PostgreSQL parser bridge's disabled-service contract aligned + // with the ordinary MySQL object-table path. The bridge runs after the + // logical shadow CREATE, so it cannot rely on createObjectTable's earlier + // configuration guard. + if storage.IsObjectStorage() && (d.provider == nil || d.provider.duckLake == nil) { + return fmt.Errorf("%w: DuckLake service configuration is disabled", ErrInvalidTableStorage) + } + d.mu.Lock() + defer d.mu.Unlock() + if execer == nil { + return fmt.Errorf("table storage metadata executor is unavailable") + } + + if storage.IsObjectStorage() { + // Both the ordinary catalog boundary and the PostgreSQL parser bridge must + // create the physical relation. The bridge supplies an executor-affine + // snapshot and deliberately skips opening a second scoped transaction, but + // that must not turn metadata publication into a shadow-only operation. + if err := d.materializeObjectTableWithExecutor(ctx, name, execer, conn); err != nil { + return err + } + if err := d.provider.EnsureDuckLakeConnectionWithExecutor(ctx, execer, conn); err != nil { + return err + } + } + + rows, err := execer.QueryContext(ctx, ` + SELECT comment + FROM duckdb_tables() + WHERE database_name = ? AND schema_name = ? AND table_name = ? + `, d.catalog, d.name, name) + if err != nil { + return ErrDuckDB.New(err) + } + if !rows.Next() { + if err := rows.Err(); err != nil { + _ = rows.Close() + return ErrDuckDB.New(err) + } + _ = rows.Close() + return sql.ErrTableNotFound.New(name) + } + + var rawComment stdsql.NullString + if err := rows.Scan(&rawComment); err != nil { + _ = rows.Close() + return ErrDuckDB.New(err) + } + if err := rows.Close(); err != nil { + return ErrDuckDB.New(err) + } + comment := DecodeComment[ExtraTableInfo](rawComment.String) + info := comment.Meta + info.Storage = storage.Kind + encoded := NewCommentWithMeta(comment.Text, info).Encode() + _, err = execer.ExecContext(ctx, fmt.Sprintf(`COMMENT ON TABLE %s IS '%s'`, FullTableName(d.catalog, d.name, name), encoded)) + if err != nil { + if storage.IsObjectStorage() { + d.dropLakeRelation(nil, nil, ctx, d.objectPhysicalTableName(name)) + } + return ErrDuckDB.New(err) + } + return nil } // CreateTemporaryTable implements sql.CreateTemporaryTable. func (d *Database) CreateTemporaryTable(ctx *sql.Context, name string, schema sql.PrimaryKeySchema, collation sql.CollationID) error { d.mu.Lock() defer d.mu.Unlock() - return d.createAllTable(ctx, name, schema, collation, "", true) + storage := DefaultTableStorageSelection() + if selected, ok := TableStorageSelectionFromContext(ctx); ok { + storage = selected + } + return d.createAllTable(ctx, name, schema, collation, "", storage, true) } // DropTable implements sql.TableDropper. @@ -337,7 +538,25 @@ func (d *Database) DropTable(ctx *sql.Context, name string) error { d.mu.Lock() defer d.mu.Unlock() - _, err := adapter.Exec(ctx, fmt.Sprintf(`DROP TABLE %s`, FullTableName(d.catalog, d.name, name))) + execer, conn, _, release, err := adapter.GetExecutionSnapshotWithLease(ctx) + if err != nil { + return err + } + defer release() + physical := FullTableName(d.catalog, d.name, name) + object := false + if d.provider != nil && d.provider.duckLake != nil { + if resolved, found, err := d.provider.ObjectTableNameWithExecutor(ctx, execer, d.catalog, d.name, name); err != nil { + return err + } else if found { + physical, object = resolved, true + if err := d.provider.EnsureDuckLakeConnectionWithExecutor(ctx, execer, conn); err != nil { + return err + } + } + } + + _, err = execer.ExecContext(ctx, fmt.Sprintf(`DROP TABLE %s`, physical)) if err != nil { if IsDuckDBTableNotFoundError(err) { @@ -345,6 +564,13 @@ func (d *Database) DropTable(ctx *sql.Context, name string) error { } return ErrDuckDB.New(err) } + if object { + // The local relation is a schema/comment shadow only. Remove it after + // the lake relation so a failed physical drop never hides live data. + if _, err := execer.ExecContext(ctx, fmt.Sprintf(`DROP TABLE %s`, FullTableName(d.catalog, d.name, name))); err != nil && !IsDuckDBTableNotFoundError(err) { + return ErrDuckDB.New(err) + } + } return nil } @@ -353,7 +579,25 @@ func (d *Database) RenameTable(ctx *sql.Context, oldName string, newName string) d.mu.Lock() defer d.mu.Unlock() - _, err := adapter.Exec(ctx, fmt.Sprintf(`ALTER TABLE %s RENAME TO "%s"`, FullTableName(d.catalog, d.name, oldName), newName)) + execer, conn, _, release, err := adapter.GetExecutionSnapshotWithLease(ctx) + if err != nil { + return err + } + defer release() + physical := FullTableName(d.catalog, d.name, oldName) + object := false + if d.provider != nil && d.provider.duckLake != nil { + if resolved, found, err := d.provider.ObjectTableNameWithExecutor(ctx, execer, d.catalog, d.name, oldName); err != nil { + return err + } else if found { + physical, object = resolved, true + if err := d.provider.EnsureDuckLakeConnectionWithExecutor(ctx, execer, conn); err != nil { + return err + } + } + } + + _, err = execer.ExecContext(ctx, fmt.Sprintf(`ALTER TABLE %s RENAME TO %s`, physical, QuoteIdentifierANSI(newName))) if err != nil { if IsDuckDBTableNotFoundError(err) { return sql.ErrTableNotFound.New(oldName) @@ -363,6 +607,11 @@ func (d *Database) RenameTable(ctx *sql.Context, oldName string, newName string) } return ErrDuckDB.New(err) } + if object { + if _, err := execer.ExecContext(ctx, fmt.Sprintf(`ALTER TABLE %s RENAME TO %s`, FullTableName(d.catalog, d.name, oldName), QuoteIdentifierANSI(newName))); err != nil { + return ErrDuckDB.New(err) + } + } return nil } @@ -539,10 +788,36 @@ func (d *Database) CopyTableData(ctx *sql.Context, sourceTable string, destinati d.mu.Lock() defer d.mu.Unlock() + execer, conn, _, release, err := adapter.GetExecutionSnapshotWithLease(ctx) + if err != nil { + return 0, err + } + defer release() + sourceName := FullTableName(d.catalog, d.name, sourceTable) + destinationName := FullTableName(d.catalog, d.name, destinationTable) + if d.provider != nil && d.provider.duckLake != nil { + if resolved, found, err := d.provider.ObjectTableNameWithExecutor(ctx, execer, d.catalog, d.name, sourceTable); err != nil { + return 0, err + } else if found { + sourceName = resolved + if err := d.provider.EnsureDuckLakeConnectionWithExecutor(ctx, execer, conn); err != nil { + return 0, err + } + } + if resolved, found, err := d.provider.ObjectTableNameWithExecutor(ctx, execer, d.catalog, d.name, destinationTable); err != nil { + return 0, err + } else if found { + destinationName = resolved + if err := d.provider.EnsureDuckLakeConnectionWithExecutor(ctx, execer, conn); err != nil { + return 0, err + } + } + } + // Use INSERT INTO ... SELECT to copy data - sql := `INSERT INTO ` + FullTableName(d.catalog, d.name, destinationTable) + ` FROM ` + FullTableName(d.catalog, d.name, sourceTable) + sql := `INSERT INTO ` + destinationName + ` SELECT * FROM ` + sourceName - res, err := adapter.Exec(ctx, sql) + res, err := execer.ExecContext(ctx, sql) if err != nil { return 0, ErrDuckDB.New(err) } diff --git a/catalog/database_ddl_test.go b/catalog/database_ddl_test.go new file mode 100644 index 00000000..a8e054f4 --- /dev/null +++ b/catalog/database_ddl_test.go @@ -0,0 +1,38 @@ +package catalog + +import ( + "context" + "os" + "path/filepath" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +// Database-file DDL must use the connection already owned by the caller. A +// max-open-connections value of one makes an accidental second acquisition +// deterministic: the operation would block while its owner is checked out. +func TestCreateCatalogOnConnUsesOwnedConnection(t *testing.T) { + provider, err := NewDBProvider("", t.TempDir(), "myduck") + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, provider.Close()) }) + + provider.Storage().SetMaxOpenConns(1) + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + conn, err := provider.Storage().Conn(ctx) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, conn.Close()) }) + + require.NoError(t, provider.CreateCatalogOnConn(ctx, conn, "owned_conn_db", false)) + var currentCatalog string + require.NoError(t, conn.QueryRowContext(ctx, "SELECT current_catalog").Scan(¤tCatalog)) + require.Equal(t, "myduck", currentCatalog) + _, err = os.Stat(filepath.Join(provider.DataDir(), "owned_conn_db.db")) + require.NoError(t, err) + + require.NoError(t, provider.DropCatalogOnConn(ctx, conn, "owned_conn_db", false)) + _, err = os.Stat(filepath.Join(provider.DataDir(), "owned_conn_db.db")) + require.ErrorIs(t, err, os.ErrNotExist) +} diff --git a/catalog/database_storage_failure_test.go b/catalog/database_storage_failure_test.go new file mode 100644 index 00000000..5c87c010 --- /dev/null +++ b/catalog/database_storage_failure_test.go @@ -0,0 +1,501 @@ +package catalog + +import ( + "context" + stdsql "database/sql" + "database/sql/driver" + "errors" + "strings" + "sync" + "testing" + + "github.com/apecloud/myduckserver/adapter" + "github.com/apecloud/myduckserver/configuration" + "github.com/apecloud/myduckserver/mycontext" + "github.com/dolthub/go-mysql-server/memory" + "github.com/dolthub/go-mysql-server/sql" + "github.com/duckdb/duckdb-go/v2" + "github.com/stretchr/testify/require" +) + +var errMetadataPublicationInjected = errors.New("metadata publication failure injected by test") + +// metadataFailureConnector delegates to DuckDB while giving the test a single +// deterministic failure point after physical relation creation. +type metadataFailureConnector struct { + inner driver.Connector + fail bool +} + +func (c *metadataFailureConnector) Connect(ctx context.Context) (driver.Conn, error) { + conn, err := c.inner.Connect(ctx) + if err != nil { + return nil, err + } + return &metadataFailureConn{inner: conn, owner: c}, nil +} + +func (c *metadataFailureConnector) Driver() driver.Driver { + return c.inner.Driver() +} + +type metadataFailureConn struct { + inner driver.Conn + owner *metadataFailureConnector +} + +func (c *metadataFailureConn) Prepare(query string) (driver.Stmt, error) { + return c.inner.Prepare(query) +} + +func (c *metadataFailureConn) Close() error { + return c.inner.Close() +} + +func (c *metadataFailureConn) Begin() (driver.Tx, error) { + return c.inner.Begin() +} + +func (c *metadataFailureConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) { + if begin, ok := c.inner.(driver.ConnBeginTx); ok { + return begin.BeginTx(ctx, opts) + } + return c.inner.Begin() +} + +func (c *metadataFailureConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) { + if c.shouldFail(query) { + return nil, errMetadataPublicationInjected + } + execer, ok := c.inner.(driver.ExecerContext) + if !ok { + return nil, driver.ErrSkip + } + return execer.ExecContext(ctx, query, args) +} + +func (c *metadataFailureConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) { + queryer, ok := c.inner.(driver.QueryerContext) + if !ok { + return nil, driver.ErrSkip + } + return queryer.QueryContext(ctx, query, args) +} + +func (c *metadataFailureConn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) { + if preparer, ok := c.inner.(driver.ConnPrepareContext); ok { + return preparer.PrepareContext(ctx, query) + } + return c.inner.Prepare(query) +} + +func (c *metadataFailureConn) CheckNamedValue(value *driver.NamedValue) error { + if checker, ok := c.inner.(driver.NamedValueChecker); ok { + return checker.CheckNamedValue(value) + } + return nil +} + +func (c *metadataFailureConn) Ping(ctx context.Context) error { + if pinger, ok := c.inner.(driver.Pinger); ok { + return pinger.Ping(ctx) + } + return nil +} + +func (c *metadataFailureConn) ResetSession(ctx context.Context) error { + if resetter, ok := c.inner.(driver.SessionResetter); ok { + return resetter.ResetSession(ctx) + } + return nil +} + +func (c *metadataFailureConn) IsValid() bool { + if validator, ok := c.inner.(driver.Validator); ok { + return validator.IsValid() + } + return true +} + +func (c *metadataFailureConn) shouldFail(query string) bool { + if c.owner == nil || !c.owner.fail { + return false + } + if !strings.HasPrefix(strings.ToUpper(strings.TrimSpace(query)), "COMMENT ON TABLE") { + return false + } + c.owner.fail = false + return true +} + +// metadataFailureSession is the smallest real session surface needed by the +// catalog bridge. It retains one physical connection and exposes its active +// transaction as an atomic pair to adapter helpers. +type metadataFailureSession struct { + *memory.Session + mu sync.Mutex + conn *stdsql.Conn + tx *stdsql.Tx +} + +var _ adapter.ConnectionHolder = (*metadataFailureSession)(nil) +var _ adapter.TransactionBindingHolder = (*metadataFailureSession)(nil) +var _ adapter.ExecutionSnapshotHolder = (*metadataFailureSession)(nil) +var _ adapter.IdentityTransactionCloser = (*metadataFailureSession)(nil) + +func (s *metadataFailureSession) GetConn(context.Context) (*stdsql.Conn, error) { + s.mu.Lock() + defer s.mu.Unlock() + return s.conn, nil +} + +func (s *metadataFailureSession) GetCatalogConn(context.Context) (*stdsql.Conn, error) { + s.mu.Lock() + defer s.mu.Unlock() + return s.conn, nil +} + +func (s *metadataFailureSession) GetTxn(ctx context.Context, options *stdsql.TxOptions) (*stdsql.Tx, error) { + s.mu.Lock() + defer s.mu.Unlock() + if s.tx != nil { + return s.tx, nil + } + tx, err := s.conn.BeginTx(ctx, options) + if err != nil { + return nil, err + } + s.tx = tx + return tx, nil +} + +func (s *metadataFailureSession) GetCatalogTxn(ctx context.Context, options *stdsql.TxOptions) (*stdsql.Tx, error) { + return s.GetTxn(ctx, options) +} + +func (s *metadataFailureSession) TryGetTxn() *stdsql.Tx { + s.mu.Lock() + defer s.mu.Unlock() + return s.tx +} + +func (s *metadataFailureSession) GetTxnBinding() (*stdsql.Conn, *stdsql.Tx) { + s.mu.Lock() + defer s.mu.Unlock() + return s.conn, s.tx +} + +func (s *metadataFailureSession) GetExecutionSnapshot(context.Context, bool) (*stdsql.Conn, *stdsql.Tx, error) { + s.mu.Lock() + defer s.mu.Unlock() + return s.conn, s.tx, nil +} + +func (s *metadataFailureSession) GetCurrentCatalog() string { return "memory" } + +func (s *metadataFailureSession) GetCurrentSchema() string { return "main" } + +func (s *metadataFailureSession) CloseTxn() { + s.mu.Lock() + s.tx = nil + s.mu.Unlock() +} + +func (s *metadataFailureSession) CloseTxnIf(tx *stdsql.Tx) { + s.mu.Lock() + if s.tx == tx { + s.tx = nil + } + s.mu.Unlock() +} + +func (s *metadataFailureSession) CloseConn() {} + +func TestRecordTableStorageSelectionRollsBackPhysicalRelationOnMetadataFailure(t *testing.T) { + inner, err := duckdb.NewConnector("", nil) + require.NoError(t, err) + failing := &metadataFailureConnector{inner: inner} + db := stdsql.OpenDB(failing) + t.Cleanup(func() { + require.NoError(t, db.Close()) + require.NoError(t, inner.Close()) + }) + + ctx := context.Background() + conn, err := db.Conn(ctx) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, conn.Close()) }) + + // Use a normal attached DuckDB catalog as a transaction-aware stand-in for + // the DuckLake relation. The provider runtime cache is marked initialized so + // this test exercises catalog publication/rollback without extension files. + _, err = conn.ExecContext(ctx, `ATTACH ':memory:' AS "__myduck_ducklake"`) + require.NoError(t, err) + provider := &DatabaseProvider{duckLake: &duckLakeRuntime{config: configuration.DuckLakeConfig{ + MetadataPath: t.TempDir() + "/catalog.ducklake", + DataPath: t.TempDir(), + }}} + physicalSchema := provider.LakeSchemaName("memory", "main") + _, err = conn.ExecContext(ctx, `CREATE SCHEMA IF NOT EXISTS "__myduck_ducklake"."`+physicalSchema+`"`) + require.NoError(t, err) + + const tableName = "metadata_failure_table" + shadowName := FullTableName("memory", "main", tableName) + _, err = conn.ExecContext(ctx, `CREATE TABLE `+shadowName+` (id INTEGER NULL)`) + require.NoError(t, err) + shadowComment := NewCommentWithMeta("keep this comment", ExtraTableInfo{ + Storage: TableStorageLocal, + }).Encode() + _, err = conn.ExecContext(ctx, `COMMENT ON TABLE `+shadowName+` IS '`+shadowComment+`'`) + require.NoError(t, err) + + runtime := provider.duckLake + require.NoError(t, conn.Raw(func(raw any) error { + runtime.initialized.Store(raw, struct{}{}) + runtime.attached.Store(raw, struct{}{}) + return nil + })) + session := &metadataFailureSession{ + Session: memory.NewSession(sql.NewBaseSession(), nil), + conn: conn, + } + sqlCtx := sql.NewContext(mycontext.WithFrontendQuery(ctx), sql.WithSession(session)) + database := NewDatabaseWithProvider("main", "memory", provider) + failing.fail = true + + err = database.RecordTableStorageSelection(sqlCtx, tableName, TableStorageSelection{ + Kind: TableStorageObject, + Explicit: true, + Source: "test", + }) + require.ErrorContains(t, err, errMetadataPublicationInjected.Error()) + require.Nil(t, session.TryGetTxn(), "scoped transaction must be released after publication failure") + + // The physical CREATE and metadata publication share the scoped transaction. + // A rollback must remove the relation rather than leave a routable object + // behind when publication fails. + var physicalCount int + require.NoError(t, conn.QueryRowContext(ctx, ` + SELECT count(*) + FROM duckdb_tables() + WHERE database_name = ? AND schema_name = ? AND table_name = ? + `, DuckLakeCatalogName, physicalSchema, tableName).Scan(&physicalCount)) + require.Zero(t, physicalCount) + + var rawComment string + require.NoError(t, conn.QueryRowContext(ctx, ` + SELECT comment + FROM duckdb_tables() + WHERE database_name = ? AND schema_name = ? AND table_name = ? + `, "memory", "main", tableName).Scan(&rawComment)) + require.Equal(t, TableStorageLocal, DecodeComment[ExtraTableInfo](rawComment).Meta.StorageKind()) +} + +// The PostgreSQL parser bridge supplies an already-captured executor and must +// still materialize the DuckLake relation. This exercises the executor-affine +// entry point directly; unlike RecordTableStorageSelection, it cannot open a +// second scoped transaction and therefore must not silently skip the physical +// CREATE. +func TestRecordTableStorageSelectionWithExecutorMaterializesPhysicalRelation(t *testing.T) { + inner, err := duckdb.NewConnector("", nil) + require.NoError(t, err) + db := stdsql.OpenDB(inner) + t.Cleanup(func() { + require.NoError(t, db.Close()) + require.NoError(t, inner.Close()) + }) + + ctx := context.Background() + conn, err := db.Conn(ctx) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, conn.Close()) }) + _, err = conn.ExecContext(ctx, `ATTACH ':memory:' AS "__myduck_ducklake"`) + require.NoError(t, err) + provider := &DatabaseProvider{duckLake: &duckLakeRuntime{config: configuration.DuckLakeConfig{ + MetadataPath: t.TempDir() + "/catalog.ducklake", + DataPath: t.TempDir(), + }}} + physicalSchema := provider.LakeSchemaName("memory", "main") + _, err = conn.ExecContext(ctx, `CREATE SCHEMA IF NOT EXISTS "__myduck_ducklake"."`+physicalSchema+`"`) + require.NoError(t, err) + + const tableName = "pg_executor_object" + shadowName := FullTableName("memory", "main", tableName) + _, err = conn.ExecContext(ctx, `CREATE TABLE `+shadowName+` (id INTEGER NULL)`) + require.NoError(t, err) + shadowComment := NewCommentWithMeta("pg bridge", ExtraTableInfo{Storage: TableStorageLocal}).Encode() + _, err = conn.ExecContext(ctx, `COMMENT ON TABLE `+shadowName+` IS '`+shadowComment+`'`) + require.NoError(t, err) + + runtime := provider.duckLake + require.NoError(t, conn.Raw(func(raw any) error { + runtime.initialized.Store(raw, struct{}{}) + runtime.attached.Store(raw, struct{}{}) + return nil + })) + session := &metadataFailureSession{ + Session: memory.NewSession(sql.NewBaseSession(), nil), + conn: conn, + } + sqlCtx := sql.NewContext(mycontext.WithFrontendQuery(ctx), sql.WithSession(session)) + database := NewDatabaseWithProvider("main", "memory", provider) + execer := adapter.SQLExecutorForConn(sqlCtx, conn) + + err = database.RecordTableStorageSelectionWithExecutor(sqlCtx, tableName, TableStorageSelection{ + Kind: TableStorageObject, + Explicit: true, + Source: "postgres", + }, execer, conn) + require.NoError(t, err) + + var physicalCount int + require.NoError(t, conn.QueryRowContext(ctx, ` + SELECT count(*) + FROM duckdb_tables() + WHERE database_name = ? AND schema_name = ? AND table_name = ? + `, DuckLakeCatalogName, physicalSchema, tableName).Scan(&physicalCount)) + require.Equal(t, 1, physicalCount) +} + +func TestDuckDBRejectsWritesToSecondAttachedDatabaseInOneTransaction(t *testing.T) { + inner, err := duckdb.NewConnector("", nil) + require.NoError(t, err) + db := stdsql.OpenDB(inner) + t.Cleanup(func() { + require.NoError(t, db.Close()) + require.NoError(t, inner.Close()) + }) + ctx := context.Background() + conn, err := db.Conn(ctx) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, conn.Close()) }) + _, err = conn.ExecContext(ctx, `ATTACH ':memory:' AS "__myduck_ducklake"`) + require.NoError(t, err) + + tx, err := conn.BeginTx(ctx, nil) + require.NoError(t, err) + defer func() { _ = tx.Rollback() }() + _, err = tx.ExecContext(ctx, `CREATE TABLE local_shadow (id INTEGER)`) + require.NoError(t, err) + _, err = tx.ExecContext(ctx, `CREATE SCHEMA IF NOT EXISTS "__myduck_ducklake"."main"`) + require.Error(t, err) + require.Contains(t, err.Error(), `already modified database`) +} + +func TestDedicatedConnectionWritesLakeAfterSessionTransactionTouchedLocal(t *testing.T) { + inner, err := duckdb.NewConnector("", nil) + require.NoError(t, err) + db := stdsql.OpenDB(inner) + t.Cleanup(func() { + require.NoError(t, db.Close()) + require.NoError(t, inner.Close()) + }) + ctx := context.Background() + sessionConn, err := db.Conn(ctx) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, sessionConn.Close()) }) + lakeConn, err := db.Conn(ctx) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, lakeConn.Close()) }) + + // DuckDB refuses to attach the same file path on two connections of one + // process. The product lake uses ATTACH IF NOT EXISTS ducklake:... per + // connection; this stand-in only needs a second connection whose writes + // are not in the session transaction. + _, err = lakeConn.ExecContext(ctx, `ATTACH ':memory:' AS "__myduck_ducklake"`) + require.NoError(t, err) + + tx, err := sessionConn.BeginTx(ctx, nil) + require.NoError(t, err) + _, err = tx.ExecContext(ctx, `CREATE TABLE local_shadow (id INTEGER)`) + require.NoError(t, err) + _, err = lakeConn.ExecContext(ctx, `CREATE SCHEMA IF NOT EXISTS "__myduck_ducklake"."main"`) + require.NoError(t, err) + _, err = lakeConn.ExecContext(ctx, `CREATE TABLE "__myduck_ducklake"."main"."pg_object" (id INTEGER)`) + require.NoError(t, err) + require.NoError(t, tx.Rollback()) + + var shadowCount int + require.NoError(t, sessionConn.QueryRowContext(ctx, ` + SELECT count(*) FROM duckdb_tables() + WHERE database_name = current_database() AND table_name = 'local_shadow' + `).Scan(&shadowCount)) + require.Zero(t, shadowCount, "user rollback must remove the local shadow") + + var lakeCount int + require.NoError(t, lakeConn.QueryRowContext(ctx, ` + SELECT count(*) FROM duckdb_tables() + WHERE database_name = '__myduck_ducklake' AND schema_name = 'main' AND table_name = 'pg_object' + `).Scan(&lakeCount)) + require.Equal(t, 1, lakeCount, "dedicated lake create survives session rollback until compensation") + + _, err = lakeConn.ExecContext(ctx, `DROP TABLE IF EXISTS "__myduck_ducklake"."main"."pg_object"`) + require.NoError(t, err) + require.NoError(t, lakeConn.QueryRowContext(ctx, ` + SELECT count(*) FROM duckdb_tables() + WHERE database_name = '__myduck_ducklake' AND schema_name = 'main' AND table_name = 'pg_object' + `).Scan(&lakeCount)) + require.Zero(t, lakeCount) +} + +func TestUncommittedLakeRelationCompensationRegistry(t *testing.T) { + provider := &DatabaseProvider{} + tx := new(stdsql.Tx) + conn := new(stdsql.Conn) + const physical = `"__myduck_ducklake"."main"."t"` + provider.registerUncommittedLakeRelation(tx, conn, physical) + require.Equal(t, []string{physical}, provider.takeUncommittedLakeRelationsForConn(conn)) + require.Empty(t, provider.takeUncommittedLakeRelationsForConn(conn)) + + provider.registerUncommittedLakeRelation(tx, conn, physical) + provider.forgetUncommittedLakeRelation(tx, physical) + require.Empty(t, provider.takeUncommittedLakeRelationsForConn(conn)) +} + +func TestMaterializeObjectTableFailsWhenSessionTxnAlreadyWroteLocal(t *testing.T) { + inner, err := duckdb.NewConnector("", nil) + require.NoError(t, err) + db := stdsql.OpenDB(inner) + t.Cleanup(func() { + require.NoError(t, db.Close()) + require.NoError(t, inner.Close()) + }) + ctx := context.Background() + conn, err := db.Conn(ctx) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, conn.Close()) }) + _, err = conn.ExecContext(ctx, `ATTACH ':memory:' AS "__myduck_ducklake"`) + require.NoError(t, err) + provider := &DatabaseProvider{duckLake: &duckLakeRuntime{config: configuration.DuckLakeConfig{ + MetadataPath: t.TempDir() + "/catalog.ducklake", + DataPath: t.TempDir(), + }}} + physicalSchema := provider.LakeSchemaName("memory", "main") + _, err = conn.ExecContext(ctx, `CREATE SCHEMA IF NOT EXISTS "__myduck_ducklake"."`+physicalSchema+`"`) + require.NoError(t, err) + + tx, err := conn.BeginTx(ctx, nil) + require.NoError(t, err) + defer func() { _ = tx.Rollback() }() + const tableName = "pg_same_txn_object" + shadowName := FullTableName("memory", "main", tableName) + _, err = tx.ExecContext(ctx, `CREATE TABLE `+shadowName+` (id INTEGER NULL)`) + require.NoError(t, err) + + runtime := provider.duckLake + require.NoError(t, conn.Raw(func(raw any) error { + runtime.initialized.Store(raw, struct{}{}) + runtime.attached.Store(raw, struct{}{}) + return nil + })) + session := &metadataFailureSession{ + Session: memory.NewSession(sql.NewBaseSession(), nil), + conn: conn, + tx: tx, + } + sqlCtx := sql.NewContext(mycontext.WithFrontendQuery(ctx), sql.WithSession(session)) + database := NewDatabaseWithProvider("main", "memory", provider) + err = database.materializeObjectTableWithExecutor(sqlCtx, tableName, tx, conn) + require.Error(t, err) + require.Contains(t, err.Error(), "create DuckLake schema failed") +} diff --git a/catalog/ducklake.go b/catalog/ducklake.go index 9270539d..ea074b41 100644 --- a/catalog/ducklake.go +++ b/catalog/ducklake.go @@ -5,6 +5,7 @@ import ( "crypto/sha256" "database/sql/driver" "encoding/hex" + stderrors "errors" "fmt" "io" "os" @@ -16,6 +17,7 @@ import ( "github.com/apecloud/myduckserver/configuration" "github.com/apecloud/myduckserver/mycontext" + "github.com/duckdb/duckdb-go/v2" ) // DuckDBExtensionVersion is the ABI/version used by the Go DuckDB binding in @@ -30,6 +32,11 @@ const DuckDBExtensionABI = DuckDBExtensionVersion // mentions any SECRET object is rejected before it reaches DuckDB. const DuckLakeSecretName = "__myduckserver_ducklake_service" +// DuckLakeCatalogName is the one service-owned catalog attached to each +// eligible physical connection. Logical MyDuck databases and schemas are +// mapped into this catalog by the provider; clients never select it directly. +const DuckLakeCatalogName = "__myduck_ducklake" + // ExtensionArtifact is one fixed, decompressed DuckDB extension binary. The // filename is deliberately fixed so service configuration cannot select an // arbitrary extension or trigger an INSTALL/network fallback. @@ -209,6 +216,1465 @@ type duckLakeRuntime struct { // statement while the provider clears it whenever the pool generation is // replaced (Reset/Restart). initialized sync.Map // map[driver.Conn]struct{} + attached sync.Map // map[driver.Conn]struct{} +} + +// duckLakeInitStage identifies the service-owned SQL step that failed. Keep +// these values stable: they are the only stage detail exposed to a protocol or +// log consumer when a driver error is returned. +type duckLakeInitStage string + +const ( + duckLakeStageLoad duckLakeInitStage = "load" + duckLakeStageCreateSecret duckLakeInitStage = "create_secret" + duckLakeStageAttach duckLakeInitStage = "attach" + duckLakeStageUnknown duckLakeInitStage = "unknown" +) + +type duckLakeFailureReason string + +const ( + duckLakeReasonContextCanceled duckLakeFailureReason = "context_canceled" + duckLakeReasonDeadlineExceeded duckLakeFailureReason = "deadline_exceeded" + duckLakeReasonBadConnection duckLakeFailureReason = "bad_connection" + duckLakeReasonPermissionDenied duckLakeFailureReason = "permission_denied" + duckLakeReasonHTTPFailure duckLakeFailureReason = "http_request_failed" + duckLakeReasonNetworkFailure duckLakeFailureReason = "network_failure" + duckLakeReasonMissingObject duckLakeFailureReason = "missing_object" + duckLakeReasonTLSFailure duckLakeFailureReason = "tls_failure" + duckLakeReasonIOFailure duckLakeFailureReason = "io_failure" + duckLakeReasonInvalidConfiguration duckLakeFailureReason = "invalid_configuration" + duckLakeReasonExtensionUnavailable duckLakeFailureReason = "extension_unavailable" + duckLakeReasonResourceExhaustion duckLakeFailureReason = "resource_exhaustion" + duckLakeReasonResourceUnavailable duckLakeFailureReason = "resource_unavailable" + duckLakeReasonInvalidStatement duckLakeFailureReason = "invalid_statement" + duckLakeReasonCatalogFailure duckLakeFailureReason = "catalog_failure" + duckLakeReasonInterrupted duckLakeFailureReason = "interrupted" + duckLakeReasonUnexpectedEOF duckLakeFailureReason = "unexpected_eof" + duckLakeReasonTimeout duckLakeFailureReason = "timeout" + duckLakeReasonDriverError duckLakeFailureReason = "driver_error" +) + +// These aliases make the minimum ATTACH diagnostic contract explicit while +// retaining the more descriptive reason names used by the existing protocol +// output. The underlying values are deliberately stable and closed. +const ( + duckLakeReasonPermission = duckLakeReasonPermissionDenied + duckLakeReasonHTTP = duckLakeReasonHTTPFailure + duckLakeReasonNetwork = duckLakeReasonNetworkFailure + duckLakeReasonExtension = duckLakeReasonExtensionUnavailable + duckLakeReasonGenericIO = duckLakeReasonIOFailure +) + +// duckLakeInitError keeps the original driver error private while exposing a +// small, stable diagnostic. A raw Unwrap is deliberately not provided: the +// PostgreSQL handler uses errors.As to replace an error's safe text with a +// PgError.Message, which would otherwise re-expose SQL, paths, or credentials. +// Is preserves sentinel matching (for example context cancellation and +// driver.ErrBadConn) without making the raw error reachable through protocol +// error formatting or errors.As. +type duckLakeInitError struct { + stage duckLakeInitStage + extension string + reason duckLakeFailureReason + cause error +} + +func newDuckLakeInitError(stage duckLakeInitStage, extension string, cause error) error { + if cause == nil { + return nil + } + stage, extension = normalizeDuckLakeInitErrorContext(stage, extension) + return &duckLakeInitError{ + stage: stage, + extension: extension, + reason: classifyDuckLakeFailureForStage(stage, cause), + cause: cause, + } +} + +func normalizeDuckLakeInitErrorContext(stage duckLakeInitStage, extension string) (duckLakeInitStage, string) { + switch stage { + case duckLakeStageLoad: + if extension == "httpfs" || extension == "ducklake" { + return stage, extension + } + return stage, "" + case duckLakeStageCreateSecret, duckLakeStageAttach: + return stage, "" + default: + return duckLakeStageUnknown, "" + } +} + +func (err *duckLakeInitError) operation() string { + switch err.stage { + case duckLakeStageLoad: + if err.extension != "" { + return "load duckdb extension " + err.extension + " failed" + } + return "load duckdb extension failed" + case duckLakeStageCreateSecret: + return "create ducklake service secret failed" + case duckLakeStageAttach: + return "attach ducklake catalog failed" + default: + return "ducklake initialization failed" + } +} + +func (err *duckLakeInitError) Error() string { + if err == nil { + return "" + } + message := err.operation() + if err.extension != "" { + return fmt.Sprintf("%s (stage=%s extension=%s reason=%s)", message, err.stage, err.extension, err.reason) + } + return fmt.Sprintf("%s (stage=%s reason=%s)", message, err.stage, err.reason) +} + +// Format keeps every fmt verb on the same redacted representation. In +// particular, %+v must not accidentally acquire a driver-specific dump. +func (err *duckLakeInitError) Format(state fmt.State, verb rune) { + if verb == 'q' { + _, _ = fmt.Fprintf(state, "%q", err.Error()) + return + } + _, _ = io.WriteString(state, err.Error()) +} + +func (err *duckLakeInitError) Is(target error) bool { + return err != nil && !isNilDuckLakeError(err.cause) && !isNilDuckLakeError(target) && stderrors.Is(err.cause, target) +} + +// rawCause is intentionally package-private. Diagnostics in this package may +// inspect the exact driver value, while protocol callers can only observe the +// fixed Error/Format output above. +func (err *duckLakeInitError) rawCause() error { + if err == nil { + return nil + } + return err.cause +} + +func classifyDuckLakeFailure(err error) duckLakeFailureReason { + return classifyDuckLakeFailureForStage(duckLakeStageUnknown, err) +} + +// classifyDuckLakeFailureForStage keeps the reason set closed while allowing +// ATTACH to distinguish a missing remote object from a generic HTTP or I/O +// failure. The driver category is authoritative for broad classes, but a few +// high-confidence markers (HTTP 404/NoSuchKey, for example) refine a typed +// DuckDB error before its broad category is applied. +func classifyDuckLakeFailureForStage(stage duckLakeInitStage, err error) duckLakeFailureReason { + if err == nil || isNilDuckLakeError(err) { + return duckLakeReasonDriverError + } + if stderrors.Is(err, context.Canceled) { + return duckLakeReasonContextCanceled + } + if stderrors.Is(err, context.DeadlineExceeded) { + return duckLakeReasonDeadlineExceeded + } + if stderrors.Is(err, driver.ErrBadConn) { + return duckLakeReasonBadConnection + } + + // A joined cause can contain more than one typed DuckDB error. Do not let + // errors.As choose whichever branch happened to appear first: the stable + // reason precedence below makes classification independent of join order. + var typedReasons []duckLakeFailureReason + for _, duckErr := range duckLakeTypedErrors(err) { + if reason, ok := classifyDuckLakeTypedError(stage, duckErr); ok { + typedReasons = append(typedReasons, reason) + } + } + if reason, ok := chooseDuckLakeFailureReason(typedReasons); ok { + return reason + } + + // Generic wrappers do not carry a typed DuckDB category. Classification + // examines only the message internally; Error/Format below never includes + // it. This lexical precedence is deliberately narrow and stable. + if reason, ok := classifyDuckLakeFailureMessage(stage, err.Error()); ok { + return reason + } + return duckLakeReasonDriverError +} + +func classifyDuckLakeTypedError(stage duckLakeInitStage, duckErr *duckdb.Error) (duckLakeFailureReason, bool) { + if isNilDuckLakeError(duckErr) { + return "", false + } + switch duckErr.Type { + case duckdb.ErrorTypePermission: + return duckLakeReasonPermissionDenied, true + case duckdb.ErrorTypeNetwork, duckdb.ErrorTypeConnection: + return duckLakeReasonNetworkFailure, true + case duckdb.ErrorTypeMissingExtension, duckdb.ErrorTypeAutoLoad: + return duckLakeReasonExtensionUnavailable, true + case duckdb.ErrorTypeHTTP: + // Refine only high-confidence HTTP statuses. Other statuses retain + // the driver's typed HTTP category. + if reason, ok := classifyDuckLakeHTTPStatus(stage, duckErr.Msg); ok { + return reason, true + } + return duckLakeReasonHTTPFailure, true + case duckdb.ErrorTypeIO: + // DuckDB often types HTTP, TLS, permission, and network failures as + // IO. Refine only those closed subclasses from the private message. + // Do not run the full untyped classifier: local "no such file" is + // resource_unavailable there, but a typed IO local miss must stay + // generic IO. The public Error/Format text never includes this message. + if reason, ok := classifyDuckLakeTypedIORefinement(stage, duckErr.Msg); ok { + return reason, true + } + return duckLakeReasonIOFailure, true + case duckdb.ErrorTypeInvalidConfiguration, + duckdb.ErrorTypeSettings, + duckdb.ErrorTypeInvalidInput, + duckdb.ErrorTypeParameterNotAllowed, + duckdb.ErrorTypeParameterNotResolved: + return duckLakeReasonInvalidConfiguration, true + case duckdb.ErrorTypeOutOfMemory, duckdb.ErrorTypeObjectSize: + return duckLakeReasonResourceExhaustion, true + case duckdb.ErrorTypeParser, duckdb.ErrorTypeSyntax: + return duckLakeReasonInvalidStatement, true + case duckdb.ErrorTypeCatalog: + return duckLakeReasonCatalogFailure, true + case duckdb.ErrorTypeInterrupt: + return duckLakeReasonInterrupted, true + default: + return "", false + } +} + +func classifyDuckLakeTypedIORefinement(stage duckLakeInitStage, message string) (duckLakeFailureReason, bool) { + if isDuckLakeMissingObjectMessage(stage, message) { + return duckLakeReasonMissingObject, true + } + if reason, ok := classifyDuckLakeHTTPStatus(stage, message); ok { + return reason, true + } + message = strings.ToLower(message) + tokens := duckLakeMessageTokens(message) + switch { + case isDuckLakePermissionMessage(message): + return duckLakeReasonPermissionDenied, true + case duckLakeTokenSequenceOutsideURL(message, tokens, "tls"), + duckLakeTokenSequenceOutsideURL(message, tokens, "ssl"), + duckLakeTokenSequenceOutsideURL(message, tokens, "x509"), + duckLakeTokenSequenceOutsideURL(message, tokens, "certificate"): + return duckLakeReasonTLSFailure, true + case duckLakeTokenSequenceOutsideURL(message, tokens, "timeout"), + duckLakeTokenSequenceOutsideURL(message, tokens, "timed", "out"), + duckLakeTokenSequenceOutsideURL(message, tokens, "deadline", "exceeded"): + return duckLakeReasonTimeout, true + case isDuckLakeNetworkMessage(message): + return duckLakeReasonNetworkFailure, true + case isDuckLakeHTTPMessage(message): + return duckLakeReasonHTTPFailure, true + default: + return "", false + } +} + +// duckLakeTypedErrors walks both ordinary wrappers and errors.Join trees. The +// visitor deliberately does not expose this tree through duckLakeInitError; +// it is an internal classification aid only. +func duckLakeTypedErrors(err error) []*duckdb.Error { + var typed []*duckdb.Error + seen := make(map[error]struct{}) + seenTyped := make(map[*duckdb.Error]struct{}) + appendTyped := func(duckErr *duckdb.Error) { + if isNilDuckLakeError(duckErr) { + return + } + if _, ok := seenTyped[duckErr]; ok { + return + } + seenTyped[duckErr] = struct{}{} + typed = append(typed, duckErr) + } + var visit func(error) + visit = func(current error) { + if current == nil || isNilDuckLakeError(current) { + return + } + if typ := reflect.TypeOf(current); typ != nil && typ.Comparable() { + if _, ok := seen[current]; ok { + return + } + seen[current] = struct{}{} + } + // Preserve the parent classifier's errors.As compatibility for opaque + // wrappers that expose a typed DuckDB error without an Unwrap method. + // Each join branch is visited independently below, so precedence stays + // deterministic even when the root As method returns only one branch. + var asDuckErr *duckdb.Error + if stderrors.As(current, &asDuckErr) { + appendTyped(asDuckErr) + } + if duckErr, ok := current.(*duckdb.Error); ok { + appendTyped(duckErr) + return + } + switch unwrapped := current.(type) { + case interface{ Unwrap() []error }: + for _, child := range unwrapped.Unwrap() { + visit(child) + } + case interface{ Unwrap() error }: + visit(unwrapped.Unwrap()) + } + } + visit(err) + return typed +} + +func chooseDuckLakeFailureReason(reasons []duckLakeFailureReason) (duckLakeFailureReason, bool) { + var best duckLakeFailureReason + bestPriority := int(^uint(0) >> 1) + for _, reason := range reasons { + priority := duckLakeFailureReasonPriority(reason) + if priority < bestPriority { + best = reason + bestPriority = priority + } + } + return best, bestPriority != int(^uint(0)>>1) +} + +func duckLakeFailureReasonPriority(reason duckLakeFailureReason) int { + // Keep the long-standing lexical precedence for joined typed causes. A + // specific ATTACH missing-object signal therefore beats broad network/HTTP + // causes regardless of errors.Join operand order. + for priority, candidate := range []duckLakeFailureReason{ + duckLakeReasonPermissionDenied, + duckLakeReasonMissingObject, + duckLakeReasonExtensionUnavailable, + duckLakeReasonTLSFailure, + duckLakeReasonUnexpectedEOF, + duckLakeReasonTimeout, + duckLakeReasonNetworkFailure, + duckLakeReasonHTTPFailure, + duckLakeReasonResourceExhaustion, + duckLakeReasonInvalidConfiguration, + duckLakeReasonInvalidStatement, + duckLakeReasonCatalogFailure, + duckLakeReasonResourceUnavailable, + duckLakeReasonIOFailure, + duckLakeReasonInterrupted, + duckLakeReasonContextCanceled, + duckLakeReasonDeadlineExceeded, + duckLakeReasonBadConnection, + } { + if reason == candidate { + return priority + } + } + return int(^uint(0) >> 1) +} + +func classifyDuckLakeFailureMessage(stage duckLakeInitStage, message string) (duckLakeFailureReason, bool) { + message = strings.ToLower(message) + tokens := duckLakeMessageTokens(message) + switch { + case isDuckLakePermissionMessage(message): + return duckLakeReasonPermissionDenied, true + case isDuckLakeMissingObjectMessage(stage, message): + return duckLakeReasonMissingObject, true + case isDuckLakeExtensionMessage(stage, message): + return duckLakeReasonExtensionUnavailable, true + case duckLakeTokenSequenceOutsideURL(message, tokens, "tls"), + duckLakeTokenSequenceOutsideURL(message, tokens, "ssl"), + duckLakeTokenSequenceOutsideURL(message, tokens, "x509"), + duckLakeTokenSequenceOutsideURL(message, tokens, "certificate"): + return duckLakeReasonTLSFailure, true + case duckLakeTokenSequenceOutsideURL(message, tokens, "unexpected", "eof"), + duckLakeTokenSequenceOutsideURL(message, tokens, "unexpected", "end", "of", "file"), + duckLakeTokenSequenceOutsideURL(message, tokens, "end", "of", "file"): + return duckLakeReasonUnexpectedEOF, true + case duckLakeTokenSequenceOutsideURL(message, tokens, "timeout"), + duckLakeTokenSequenceOutsideURL(message, tokens, "timed", "out"), + duckLakeTokenSequenceOutsideURL(message, tokens, "deadline", "exceeded"): + return duckLakeReasonTimeout, true + case isDuckLakeNetworkMessage(message): + return duckLakeReasonNetworkFailure, true + case isDuckLakeHTTPMessage(message): + return duckLakeReasonHTTPFailure, true + case duckLakeTokenSequenceOutsideURL(message, tokens, "out", "of", "memory"), + duckLakeLiteralOutsideURL(message, "out-of-memory"), + duckLakeTokenSequenceOutsideURL(message, tokens, "cannot", "allocate", "memory"), + duckLakeTokenSequenceOutsideURL(message, tokens, "no", "space", "left", "on", "device"), + duckLakeTokenSequenceOutsideURL(message, tokens, "disk", "full"), + duckLakeTokenSequenceOutsideURL(message, tokens, "quota", "exceeded"), + duckLakeTokenSequenceOutsideURL(message, tokens, "too", "many", "open", "files"), + duckLakeTokenSequenceOutsideURL(message, tokens, "resource", "exhausted"): + return duckLakeReasonResourceExhaustion, true + case duckLakeTokenSequenceOutsideURL(message, tokens, "invalid", "configuration"), + duckLakeTokenSequenceOutsideURL(message, tokens, "configuration", "error"), + duckLakeTokenSequenceOutsideURL(message, tokens, "invalid", "parameter"), + duckLakeTokenSequenceOutsideURL(message, tokens, "missing", "secret"), + duckLakeTokenSequenceOutsideURL(message, tokens, "mandatory", "setting"), + duckLakeTokenSequenceOutsideURL(message, tokens, "settings", "error"): + return duckLakeReasonInvalidConfiguration, true + case duckLakeTokenSequenceOutsideURL(message, tokens, "syntax"), + duckLakeTokenSequenceOutsideURL(message, tokens, "parser"), + duckLakeTokenSequenceOutsideURL(message, tokens, "invalid", "statement"): + return duckLakeReasonInvalidStatement, true + case duckLakeTokenSequenceOutsideURL(message, tokens, "catalog", "error"): + return duckLakeReasonCatalogFailure, true + case isDuckLakeResourceUnavailableMessage(message): + return duckLakeReasonResourceUnavailable, true + case duckLakeLiteralOutsideURL(message, "i/o error"), + duckLakeTokenSequenceOutsideURL(message, tokens, "io", "error"), + duckLakeLiteralOutsideURL(message, "input/output error"), + duckLakeLiteralOutsideURL(message, "read-only file system"), + duckLakeTokenSequenceOutsideURL(message, tokens, "read", "only", "file", "system"), + duckLakeTokenSequenceOutsideURL(message, tokens, "failed", "to", "open"), + duckLakeTokenSequenceOutsideURL(message, tokens, "failed", "to", "read"), + duckLakeTokenSequenceOutsideURL(message, tokens, "failed", "to", "write"), + duckLakeTokenSequenceOutsideURL(message, tokens, "bad", "file", "descriptor"): + return duckLakeReasonIOFailure, true + case duckLakeTokenSequenceOutsideURL(message, tokens, "interrupt"), + duckLakeTokenSequenceOutsideURL(message, tokens, "interrupted"), + duckLakeTokenSequenceOutsideURL(message, tokens, "interruption"), + duckLakeTokenSequenceOutsideURL(message, tokens, "cancel"), + duckLakeTokenSequenceOutsideURL(message, tokens, "canceled"), + duckLakeTokenSequenceOutsideURL(message, tokens, "cancelled"), + duckLakeTokenSequenceOutsideURL(message, tokens, "cancellation"): + return duckLakeReasonInterrupted, true + default: + return "", false + } +} + +func isDuckLakeMissingObjectMessage(stage duckLakeInitStage, message string) bool { + if stage != duckLakeStageAttach { + return false + } + message = strings.ToLower(message) + tokens := duckLakeMessageTokens(message) + for _, sequence := range [][]string{ + {"nosuchkey"}, + {"nosuchobject"}, + {"keynotfound"}, + {"objectnotfound"}, + {"no", "such", "key"}, + {"no", "such", "object"}, + {"key", "not", "found"}, + {"key", "does", "not", "exist"}, + {"specified", "key", "does", "not", "exist"}, + {"object", "not", "found"}, + {"object", "does", "not", "exist"}, + {"missing", "object"}, + {"object", "missing"}, + {"missing", "metadata"}, + {"metadata", "missing"}, + {"metadata", "not", "found"}, + {"catalog", "not", "found"}, + {"catalog", "does", "not", "exist"}, + {"database", "does", "not", "exist"}, + {"read", "only", "mode"}, + } { + if duckLakeTokenSequenceOutsideURL(message, tokens, sequence...) { + return true + } + } + // A numeric 404 is only treated as an object miss when it is clearly an + // HTTP status. This avoids classifying an arbitrary path containing "404". + if hasDuckLakeHTTPStatus(message, "404") { + return true + } + // A remote no-such-file response is an object miss only when the message + // names a remote/object context. Bare local-file errors remain generic I/O. + for _, sequence := range [][]string{ + {"no", "such", "file"}, + {"file", "not", "found"}, + {"file", "does", "not", "exist"}, + } { + if duckLakeRemoteFileSequenceOutsideURL(message, tokens, sequence...) { + return true + } + } + return false +} + +func classifyDuckLakeHTTPStatus(stage duckLakeInitStage, message string) (duckLakeFailureReason, bool) { + message = strings.ToLower(message) + // A typed HTTP error may put the status after an object code (for example, + // "NoSuchKey 403"). Exact status tokens therefore take precedence without + // depending on their position relative to the usual "HTTP Error" marker. + if hasDuckLakeHTTPStatusToken(message, "401") || hasDuckLakeHTTPStatusToken(message, "403") { + return duckLakeReasonPermissionDenied, true + } + if stage == duckLakeStageAttach && (hasDuckLakeHTTPStatusToken(message, "404") || isDuckLakeMissingObjectMessage(stage, message)) { + return duckLakeReasonMissingObject, true + } + return "", false +} + +func hasDuckLakeHTTPStatusToken(message, code string) bool { + tokens := duckLakeMessageTokens(message) + for index, token := range tokens { + if token.value != code || duckLakeTokenInURL(message, token) || duckLakeTokenIsHyphenated(message, token) { + continue + } + if duckLakeHTTPStatusTokenHasPathContext(message, tokens, index) { + continue + } + if duckLakeHTTPStatusContext(message, tokens, index) { + return true + } + } + return false +} + +func duckLakeHTTPStatusTokenHasPathContext(message string, tokens []duckLakeMessageToken, index int) bool { + if duckLakeTokenIsPathLike(message, tokens[index]) { + return true + } + // Query/identifier fields may contain a status-looking suffix (for + // example, foo=status=404). Only the known status field spellings are + // eligible for refinement; an unknown prefix must remain broad HTTP. + if duckLakeHTTPStatusTokenHasUnknownIdentifierPrefix(message, tokens, index) { + return true + } + for cursor := index - 1; cursor >= 0 && index-cursor <= 3; cursor-- { + token := tokens[cursor] + if token.value == "path" { + // A path assignment such as path=404 is not an HTTP status. The + // same applies when path appears before a status/code marker. + return true + } + if token.value != "status" && token.value != "response" && token.value != "code" { + continue + } + if duckLakeTokenIsPathComponent(message, token) || duckLakeTokenIsHyphenated(message, token) || duckLakeTokenInURL(message, token) { + return true + } + } + return false +} + +func duckLakeHTTPStatusTokenHasUnknownIdentifierPrefix(message string, tokens []duckLakeMessageToken, index int) bool { + if index <= 0 || index >= len(tokens) { + return false + } + // Only an equals sign can bind a numeric status to a named field. A dot or + // underscore makes the number part of an identifier instead (for example, + // status_code_404 or status_code.404). + numericSeparator := strings.TrimSpace(message[tokens[index-1].end:tokens[index].start]) + if numericSeparator == "" { + return false + } + if numericSeparator != "=" { + return strings.ContainsAny(numericSeparator, "_.=") + } + return !duckLakeHTTPStatusTokenHasKnownIdentifierPrefix(message, tokens, index) +} + +func duckLakeHTTPStatusTokenHasKnownIdentifierPrefix(message string, tokens []duckLakeMessageToken, index int) bool { + if index <= 0 || index >= len(tokens) || + strings.TrimSpace(message[tokens[index-1].end:tokens[index].start]) != "=" { + return false + } + // Walk an exact snake_case field to the left of the assignment. Repeated + // underscores and nested prefixes are identifiers, not known status fields. + fieldStart := index - 1 + for fieldStart > 0 && message[tokens[fieldStart-1].end:tokens[fieldStart].start] == "_" { + fieldStart-- + } + if fieldStart > 0 { + prefixSeparator := strings.TrimSpace(message[tokens[fieldStart-1].end:tokens[fieldStart].start]) + if strings.Contains(prefixSeparator, "=") { + return false + } + } + if start := tokens[fieldStart].start; start > 0 { + switch message[start-1] { + case '_', '.', '=': + return false + } + } + values := make([]string, 0, index-fieldStart) + for cursor := fieldStart; cursor < index; cursor++ { + values = append(values, tokens[cursor].value) + } + field := strings.Join(values, "_") + switch field { + case "http", "status", "response", "code", "returned", "return", + "status_code", "response_code", "http_status", "http_code", "error_code": + return true + default: + return false + } +} + +func hasDuckLakeHTTPStatus(message, code string) bool { + message = strings.ToLower(message) + tokens := duckLakeMessageTokens(message) + for index, token := range tokens { + if token.value == code && duckLakeHTTPStatusContext(message, tokens, index) { + return true + } + } + return false +} + +func duckLakeHTTPStatusContext(message string, tokens []duckLakeMessageToken, index int) bool { + if index < 0 || index >= len(tokens) || duckLakeTokenInURL(message, tokens[index]) { + return false + } + if duckLakeTokenIsHyphenated(message, tokens[index]) { + return false + } + if duckLakeTokenIsPathLike(message, tokens[index]) { + return false + } + if duckLakeTokenHasDottedIdentifierContext(message, tokens[index]) { + return false + } + if index == 0 { + if len(tokens) == 1 || duckLakeHTTPReasonPhraseFollows(message, tokens, index) { + if index+1 < len(tokens) && !duckLakeTokenIsMarker(message, tokens[index+1]) { + return false + } + return true + } + return duckLakeLeadingHTTPStatusContext(message, tokens, index) + } + for cursor := index - 1; cursor >= 0 && index-cursor <= 5; cursor-- { + token := tokens[cursor] + switch token.value { + case "http", "status", "response", "code", "returned", "return": + if duckLakeHTTPStatusTokenHasPathContext(message, tokens, index) || duckLakeTokenIsPathComponent(message, token) { + return false + } + if duckLakeTokenIsHyphenated(message, token) { + return false + } + if token.value == "http" && duckLakeHTTPVersionPrecedesStatus(message, token, tokens[index]) { + return duckLakeTokenIsPlain(message, token) && + !duckLakeTokenIsPathComponent(message, token) && + !duckLakeTokenIsIdentifierLike(message, token) + } + if !duckLakeTokenIsMarker(message, token) && + !duckLakeHTTPStatusTokenHasKnownIdentifierPrefix(message, tokens, index) { + return false + } + // A literal http:// URL is not an HTTP error marker. + if token.value == "http" && strings.HasPrefix(message[token.end:], "://") { + return false + } + if !duckLakeTokenInURL(message, token) { + return true + } + return false + case "nosuchkey", "nosuchobject", "keynotfound", "objectnotfound": + // Object-store response codes may precede the numeric status (for + // example, "NoSuchKey 403"). The object marker itself is enough to + // establish status context; URL/path markers were rejected above. + if !duckLakeTokenIsMarker(message, token) { + return false + } + return true + case "error", "request": + if !duckLakeTokenIsMarker(message, token) { + return false + } + continue + case "1", "2", "3", "0": + continue + default: + return false + } + } + return false +} + +func duckLakeHTTPVersionPrecedesStatus(message string, httpToken, statusToken duckLakeMessageToken) bool { + if httpToken.end > statusToken.start || statusToken.start > len(message) { + return false + } + version := strings.TrimSpace(message[httpToken.end:statusToken.start]) + if len(version) < 2 || version[0] != '/' { + return false + } + digitSeen := false + for index := 1; index < len(version); index++ { + switch value := version[index]; { + case value >= '0' && value <= '9': + digitSeen = true + case value == '.' && digitSeen && index+1 < len(version): + digitSeen = false + default: + return false + } + } + return digitSeen +} + +func duckLakeLeadingHTTPStatusContext(message string, tokens []duckLakeMessageToken, index int) bool { + if index+1 >= len(tokens) { + return false + } + if !duckLakeTokenIsMarker(message, tokens[index+1]) { + return false + } + switch tokens[index+1].value { + case "at", "from", "for", "on", "while", "during", "request", "response", "status", "error", "object", "objects", "key", "path", "catalog", "metadata", "resource", "url", "uri", "s3", "http", "https", "nosuchkey", "nosuchobject": + return true + default: + return false + } +} + +func duckLakeHTTPReasonPhraseFollows(message string, tokens []duckLakeMessageToken, index int) bool { + if index+1 >= len(tokens) { + return false + } + switch tokens[index].value { + case "401": + return tokens[index+1].value == "unauthorized" && duckLakeTokenIsMarker(message, tokens[index+1]) + case "403": + return tokens[index+1].value == "forbidden" && duckLakeTokenIsMarker(message, tokens[index+1]) + case "404": + return tokens[index+1].value == "not" && index+2 < len(tokens) && + tokens[index+2].value == "found" && duckLakeTokenIsMarker(message, tokens[index+1]) && + duckLakeTokenIsMarker(message, tokens[index+2]) + default: + return (tokens[index+1].value == "error" || tokens[index+1].value == "unavailable") && + duckLakeTokenIsMarker(message, tokens[index+1]) + } +} + +func isDuckLakePermissionMessage(message string) bool { + tokens := duckLakeMessageTokens(message) + if hasDuckLakeHTTPStatus(message, "401") || hasDuckLakeHTTPStatus(message, "403") { + return true + } + for _, sequence := range [][]string{{"permission", "denied"}, {"access", "denied"}} { + if duckLakeTokenSequenceOutsideURL(message, tokens, sequence...) { + return true + } + } + for index, token := range tokens { + if token.value != "unauthorized" && token.value != "forbidden" { + continue + } + if !duckLakeTokenIsMarker(message, token) { + continue + } + if len(tokens) == 1 || duckLakePermissionContext(message, tokens, index) { + return true + } + } + return false +} + +func duckLakePermissionContext(message string, tokens []duckLakeMessageToken, index int) bool { + if index > 0 { + previous := tokens[index-1] + if !duckLakeTokenIsMarker(message, previous) { + previous = duckLakeMessageToken{} + } + switch previous.value { + case "permission", "access", "operation", "http", "status", "code", "response", "returned", "return", "error", "request": + return true + } + } + if index+1 < len(tokens) { + next := tokens[index+1] + if !duckLakeTokenIsMarker(message, next) { + next = duckLakeMessageToken{} + } + switch next.value { + case "error", "response", "status", "code", "request", "access", "operation", "by", "from": + return true + } + } + return false +} + +func isDuckLakeNetworkMessage(message string) bool { + tokens := duckLakeMessageTokens(message) + for _, sequence := range [][]string{ + {"connection", "refused"}, + {"connection", "reset"}, + {"connection", "aborted"}, + {"connection", "closed"}, + {"no", "such", "host"}, + {"no", "route", "to", "host"}, + {"broken", "pipe"}, + {"dial", "tcp"}, + {"name", "resolution"}, + {"network", "error"}, + {"network", "failure"}, + {"network", "request"}, + {"network", "connection"}, + {"network", "unavailable"}, + {"network", "unreachable"}, + {"network", "operation"}, + {"dns", "error"}, + {"dns", "failure"}, + {"dns", "resolution"}, + } { + if duckLakeTokenSequenceOutsideURL(message, tokens, sequence...) { + return true + } + } + for index, token := range tokens { + if token.value != "network" && token.value != "dns" { + continue + } + if !duckLakeTokenIsMarker(message, token) { + continue + } + if len(tokens) == 1 || + (index > 0 && tokens[index-1].value == "error" && duckLakeTokenIsMarker(message, tokens[index-1])) || + (index+1 < len(tokens) && tokens[index+1].value == "error" && duckLakeTokenIsMarker(message, tokens[index+1])) { + return true + } + } + return false +} + +func isDuckLakeHTTPMessage(message string) bool { + tokens := duckLakeMessageTokens(message) + if hasDuckLakeHTTPStatusCode(message, tokens) { + return true + } + for _, sequence := range [][]string{ + {"http", "error"}, + {"http", "failure"}, + {"http", "failed"}, + {"http", "request"}, + {"http", "response"}, + {"http", "status"}, + } { + if duckLakeTokenSequenceOutsideURL(message, tokens, sequence...) { + return true + } + } + return false +} + +func hasDuckLakeHTTPStatusCode(message string, tokens []duckLakeMessageToken) bool { + for index, token := range tokens { + if !isDuckLakeHTTPStatusToken(token.value) { + continue + } + if duckLakeHTTPStatusContext(message, tokens, index) { + return true + } + } + return false +} + +func isDuckLakeHTTPStatusToken(value string) bool { + if len(value) != 3 { + return false + } + for index := range value { + if value[index] < '0' || value[index] > '9' { + return false + } + } + return value[0] >= '1' && value[0] <= '5' +} + +func isDuckLakeExtensionMessage(stage duckLakeInitStage, message string) bool { + tokens := duckLakeMessageTokens(message) + for _, sequence := range [][]string{ + {"cannot", "load", "extension"}, + {"cannot", "find", "extension"}, + {"failed", "to", "load", "extension"}, + {"autoload", "extension"}, + {"autoloading", "extension"}, + } { + if duckLakeTokenSequenceOutsideURL(message, tokens, sequence...) { + return true + } + } + // Ordinary extension diagnostics are trusted during the service-owned + // LOAD stage. At service stages after LOAD, require an explicit load/autoload + // phrase so an object path such as "remote extension missing" cannot hijack + // the ATTACH subclass. Unknown-stage generic classification retains the + // historical lexical behavior. + if stage == duckLakeStageLoad { + for _, sequence := range [][]string{ + {"missing", "extension"}, + } { + if duckLakeTokenSequenceOutsideURL(message, tokens, sequence...) { + return true + } + } + } + for index, token := range tokens { + if token.value != "extension" || !duckLakeTokenIsMarker(message, token) { + continue + } + if (stage == duckLakeStageAttach || stage == duckLakeStageCreateSecret) && !duckLakeExtensionHasExplicitLoadContext(message, tokens, index) { + continue + } + // An explicit load/autoload phrase is itself sufficient context. This + // covers word orders such as "failed to load ducklake extension" where + // the failure verb precedes the extension token. + if duckLakeExtensionHasExplicitLoadContext(message, tokens, index) { + return true + } + for cursor := index + 1; cursor < len(tokens) && cursor <= index+5; cursor++ { + candidate := tokens[cursor] + if !duckLakeTokenIsMarker(message, candidate) { + continue + } + switch candidate.value { + case "not": + if cursor+1 < len(tokens) && + duckLakeTokenIsMarker(message, tokens[cursor+1]) && + (tokens[cursor+1].value == "loaded" || tokens[cursor+1].value == "found") { + return true + } + case "unavailable", "autoload", "autoloading", "missing": + return true + case "cannot": + if cursor+1 < len(tokens) && + duckLakeTokenIsMarker(message, tokens[cursor+1]) && + (tokens[cursor+1].value == "load" || tokens[cursor+1].value == "find") { + return true + } + case "failed": + return true + } + } + } + return false +} + +func duckLakeExtensionHasExplicitLoadContext(message string, tokens []duckLakeMessageToken, index int) bool { + for cursor := index + 1; cursor < len(tokens) && cursor <= index+5; cursor++ { + if !duckLakeTokenIsMarker(message, tokens[cursor]) { + continue + } + switch tokens[cursor].value { + case "load", "loaded", "autoload", "autoloading": + return true + case "not", "cannot", "failed": + if cursor+1 < len(tokens) && duckLakeTokenIsMarker(message, tokens[cursor+1]) && + (tokens[cursor+1].value == "loaded" || tokens[cursor+1].value == "load" || tokens[cursor+1].value == "find") { + return true + } + } + } + for cursor := index - 1; cursor >= 0 && index-cursor <= 5; cursor-- { + if !duckLakeTokenIsMarker(message, tokens[cursor]) { + continue + } + switch tokens[cursor].value { + case "load", "loaded", "autoload", "autoloading": + return true + case "cannot", "failed": + if cursor+1 < len(tokens) && duckLakeTokenIsMarker(message, tokens[cursor+1]) && + (tokens[cursor+1].value == "load" || tokens[cursor+1].value == "find") { + return true + } + } + } + return false +} + +func isDuckLakeResourceUnavailableMessage(message string) bool { + tokens := duckLakeMessageTokens(message) + for _, sequence := range [][]string{ + {"resource", "temporarily", "unavailable"}, + {"resource", "unavailable"}, + {"no", "such", "file"}, + {"file", "not", "found"}, + {"file", "does", "not", "exist"}, + {"resource", "not", "found"}, + {"resource", "does", "not", "exist"}, + {"object", "not", "found"}, + {"object", "does", "not", "exist"}, + {"metadata", "not", "found"}, + {"catalog", "not", "found"}, + {"key", "not", "found"}, + } { + if duckLakeTokenSequenceOutsideURL(message, tokens, sequence...) { + return true + } + } + if (len(tokens) == 2 && tokens[0].value == "not" && tokens[1].value == "found" && + duckLakeTokenIsMarker(message, tokens[0]) && duckLakeTokenIsMarker(message, tokens[1])) || + (len(tokens) == 3 && tokens[0].value == "does" && tokens[1].value == "not" && tokens[2].value == "exist" && + duckLakeTokenIsMarker(message, tokens[0]) && duckLakeTokenIsMarker(message, tokens[1]) && + duckLakeTokenIsMarker(message, tokens[2])) { + return true + } + for index, token := range tokens { + if token.value != "unavailable" && token.value != "missing" { + continue + } + if !duckLakeTokenIsMarker(message, token) { + continue + } + if len(tokens) == 1 || + (index > 0 && duckLakeTokenIsMarker(message, tokens[index-1]) && isDuckLakeResourceContextWord(tokens[index-1].value)) || + (index+1 < len(tokens) && duckLakeTokenIsMarker(message, tokens[index+1]) && isDuckLakeResourceContextWord(tokens[index+1].value)) { + return true + } + } + return false +} + +func isDuckLakeResourceContextWord(value string) bool { + switch value { + case "resource", "file", "object", "metadata", "catalog", "key", "service", "path": + return true + default: + return false + } +} + +func duckLakeRemoteFileSequenceOutsideURL(message string, tokens []duckLakeMessageToken, sequence ...string) bool { + if len(sequence) == 0 || len(sequence) > len(tokens) { + return false + } + for index := 0; index <= len(tokens)-len(sequence); index++ { + if !duckLakeTokenSequenceMatchesAt(message, tokens, index, sequence...) { + continue + } + if duckLakeHasRemoteObjectContextNear(message, tokens, index, index+len(sequence)) { + return true + } + } + return false +} + +func duckLakeHasRemoteObjectContextNear(message string, tokens []duckLakeMessageToken, start, end int) bool { + forwardLink := false + for cursor := end; cursor < len(tokens) && cursor-end < 4; cursor++ { + if duckLakeTokenIsRemoteObjectContext(message, tokens[cursor]) { + return true + } + if !duckLakeTokenIsRemoteContextLink(message, tokens[cursor]) { + if forwardLink && duckLakeTokenIsExplicitLocalTarget(message, tokens[cursor]) { + return false + } + break + } + forwardLink = true + } + for cursor := start - 1; cursor >= 0 && start-cursor <= 4; cursor-- { + if duckLakeTokenIsRemoteObjectContext(message, tokens[cursor]) { + return true + } + if !duckLakeTokenIsRemoteContextLink(message, tokens[cursor]) { + break + } + } + return false +} + +func duckLakeTokenIsExplicitLocalTarget(message string, token duckLakeMessageToken) bool { + return !duckLakeTokenInRemoteURL(message, token) && + (duckLakeTokenIsPathLike(message, token) || duckLakeTokenIsIdentifierLike(message, token)) +} + +func duckLakeTokenIsRemoteObjectContext(message string, token duckLakeMessageToken) bool { + if duckLakeTokenInRemoteURL(message, token) { + return true + } + if token.value == "s3" && strings.HasPrefix(message[token.end:], "://") { + return false + } + return (token.value == "remote" || token.value == "s3") && duckLakeTokenIsMarker(message, token) +} + +func duckLakeTokenIsRemoteContextLink(message string, token duckLakeMessageToken) bool { + switch token.value { + case "at", "from", "on", "in", "for": + return duckLakeTokenIsMarker(message, token) + case "error", "io", "storage": + return duckLakeTokenIsMarker(message, token) + case "url", "uri", "endpoint", "path", "object", "key", "catalog", "metadata", "file": + if duckLakeTokenIsMarker(message, token) { + return true + } + // A known target field may bind directly to the remote URI. Keep the + // exception local to this proximity scan so arbitrary assignments remain + // opaque markers. + if !duckLakeTokenHasSingleAssignmentSuffix(message, token) || + !duckLakeTokenIsPlain(message, token) || duckLakeTokenIsPathComponent(message, token) { + return false + } + if token.start > 0 { + switch message[token.start-1] { + case '_', '.', '=': + return false + } + } + return true + default: + return false + } +} + +func duckLakeTokenHasSingleAssignmentSuffix(message string, token duckLakeMessageToken) bool { + if token.end >= len(message) || message[token.end] != '=' { + return false + } + for cursor := token.end + 1; cursor < len(message); cursor++ { + switch message[cursor] { + case ' ', '\t', '\r', '\n', '\'', '"', '`': + continue + default: + return message[cursor] != '=' + } + } + return true +} + +func duckLakeTokenInRemoteURL(message string, token duckLakeMessageToken) bool { + for _, scheme := range []string{"s3://", "http://", "https://", "azure://", "gs://", "gcs://"} { + for offset := 0; offset < len(message); { + relative := strings.Index(message[offset:], scheme) + if relative < 0 { + break + } + start := offset + relative + end := strings.IndexAny(message[start+len(scheme):], " \t\r\n") + if end < 0 { + end = len(message) + } else { + end += start + len(scheme) + } + if token.start >= start && token.start < end && duckLakeRemoteURLPrefixIsValid(message, start) { + return true + } + offset = start + len(scheme) + } + } + return false +} + +func duckLakeRemoteURLPrefixIsValid(message string, start int) bool { + if start <= 0 || start > len(message) { + return start == 0 + } + prefixTokens := duckLakeMessageTokens(message[:start]) + if len(prefixTokens) == 0 { + return true + } + field := prefixTokens[len(prefixTokens)-1] + separator := message[field.end:start] + compact := strings.Map(func(value rune) rune { + switch value { + case ' ', '\t', '\r', '\n', '\'', '"', '`': + return -1 + default: + return value + } + }, separator) + if strings.Contains(compact, "=") { + if compact != "=" { + return false + } + switch field.value { + case "url", "uri", "endpoint", "path", "object", "key", "catalog", "metadata", "file": + default: + return false + } + if field.start > 0 { + switch message[field.start-1] { + case '_', '.', '=': + return false + } + } + if duckLakeTokenHasDottedIdentifierContext(message[:start], field) { + return false + } + if len(prefixTokens) > 1 { + outerSeparator := strings.TrimSpace(message[prefixTokens[len(prefixTokens)-2].end:field.start]) + if strings.Contains(outerSeparator, "=") { + return false + } + } + return true + } + previous := message[start-1] + return !isDuckLakeTokenByte(previous) && !strings.ContainsRune("_-./\\?&#=", rune(previous)) +} + +type duckLakeMessageToken struct { + value string + start int + end int +} + +func duckLakeMessageTokens(message string) []duckLakeMessageToken { + message = strings.ToLower(message) + tokens := make([]duckLakeMessageToken, 0, 8) + for index := 0; index < len(message); { + for index < len(message) && !isDuckLakeTokenByte(message[index]) { + index++ + } + start := index + for index < len(message) && isDuckLakeTokenByte(message[index]) { + index++ + } + if start < index { + tokens = append(tokens, duckLakeMessageToken{value: message[start:index], start: start, end: index}) + } + } + return tokens +} + +func isDuckLakeTokenByte(value byte) bool { + return (value >= 'a' && value <= 'z') || (value >= '0' && value <= '9') +} + +// duckLakeTokenSequenceOutsideURL finds any complete sequence whose tokens are +// outside a URL. A message can contain a diagnostic-looking object URL before +// the actual error phrase; returning only the first match would hide the valid +// later phrase. +func duckLakeTokenSequenceOutsideURL(message string, tokens []duckLakeMessageToken, sequence ...string) bool { + if len(sequence) == 0 || len(sequence) > len(tokens) { + return false + } + for index := 0; index <= len(tokens)-len(sequence); index++ { + if duckLakeTokenSequenceMatchesAt(message, tokens, index, sequence...) { + return true + } + } + return false +} + +func duckLakeTokenSequenceMatchesAt(message string, tokens []duckLakeMessageToken, index int, sequence ...string) bool { + if index < 0 || len(sequence) == 0 || index+len(sequence) > len(tokens) { + return false + } + for offset, value := range sequence { + candidate := tokens[index+offset] + if candidate.value != value || !duckLakeTokenIsMarker(message, candidate) { + return false + } + } + return true +} + +// duckLakeLiteralOutsideURL handles the few canonical diagnostics whose own +// spelling contains a slash or hyphen. Identifier, query and path adjacency is +// still rejected; a trailing period remains ordinary sentence punctuation. +func duckLakeLiteralOutsideURL(message, literal string) bool { + for offset := 0; offset <= len(message)-len(literal); { + relative := strings.Index(message[offset:], literal) + if relative < 0 { + return false + } + start := offset + relative + end := start + len(literal) + if !duckLakeLiteralHasIdentifierBoundary(message, start, end) && + !duckLakeTokenInURL(message, duckLakeMessageToken{start: start, end: end}) { + return true + } + offset = start + 1 + } + return false +} + +func duckLakeLiteralHasIdentifierBoundary(message string, start, end int) bool { + left := start + skippedBoundary := false + +leftScan: + for left > 0 { + switch message[left-1] { + case ' ', '\t', '\r', '\n', '\'', '"', '`': + left-- + skippedBoundary = true + default: + break leftScan + } + } + if left > 0 { + previous := message[left-1] + if (skippedBoundary && previous == '=') || + (!skippedBoundary && (isDuckLakeTokenByte(previous) || + strings.ContainsRune("_-./\\?&#=", rune(previous)))) { + return true + } + } + if end >= len(message) { + return false + } + next := message[end] + if isDuckLakeTokenByte(next) || strings.ContainsRune("_-/\\?&#=", rune(next)) { + return true + } + return next == '.' && duckLakeDotsLeadToIdentifier(message, end) +} + +func duckLakeTokenInURL(message string, token duckLakeMessageToken) bool { + if token.start > len(message) { + return false + } + scheme := strings.LastIndex(message[:token.start], "://") + if scheme < 0 { + return false + } + end := strings.IndexAny(message[scheme+3:], " \t\r\n") + if end < 0 { + end = len(message) - (scheme + 3) + } + return token.start < scheme+3+end +} + +func duckLakeTokenIsHyphenated(message string, token duckLakeMessageToken) bool { + return (token.start > 0 && message[token.start-1] == '-') || + (token.end < len(message) && message[token.end] == '-') +} + +// duckLakeTokenIsPlain keeps lexical markers out of URLs and hyphenated +// identifiers. Path-component checks are intentionally separate because some +// callers use ordinary path words as explicit context. +func duckLakeTokenIsPlain(message string, token duckLakeMessageToken) bool { + return !duckLakeTokenInURL(message, token) && !duckLakeTokenIsHyphenated(message, token) +} + +func duckLakeTokenIsMarker(message string, token duckLakeMessageToken) bool { + return duckLakeTokenIsPlain(message, token) && !duckLakeTokenIsPathLike(message, token) && !duckLakeTokenIsIdentifierLike(message, token) +} + +func duckLakeTokenIsIdentifierLike(message string, token duckLakeMessageToken) bool { + if token.start > 0 { + switch message[token.start-1] { + case '_', '=': + return true + } + for cursor := token.start; cursor > 0; { + cursor-- + switch message[cursor] { + case ' ', '\t', '\r', '\n', '\'', '"', '`': + continue + case '=': + return true + default: + cursor = 0 + } + } + } + if token.end < len(message) { + switch message[token.end] { + case '_', '=': + return true + } + } + return duckLakeTokenHasDottedIdentifierContext(message, token) +} + +func duckLakeTokenHasDottedIdentifierContext(message string, token duckLakeMessageToken) bool { + if duckLakeSpanHasDottedIdentifierContext(message, token.start, token.end) { + return true + } + for _, delimiters := range [][2]byte{{'"', '"'}, {'`', '`'}, {'[', ']'}} { + if duckLakeDelimitedTokenHasDottedIdentifierContext(message, token, delimiters[0], delimiters[1]) { + return true + } + } + return false +} + +func duckLakeSpanHasDottedIdentifierContext(message string, start, end int) bool { + if start > 0 && message[start-1] == '.' { + return true + } + return end < len(message) && message[end] == '.' && + duckLakeDotsLeadToIdentifier(message, end) +} + +func duckLakeDelimitedTokenHasDottedIdentifierContext( + message string, + token duckLakeMessageToken, + openingDelimiter, closingDelimiter byte, +) bool { + for offset := 0; offset < len(message); { + openingOffset := strings.IndexByte(message[offset:], openingDelimiter) + if openingOffset < 0 { + return false + } + opening := offset + openingOffset + closing := opening + 1 + for closing < len(message) { + closingOffset := strings.IndexByte(message[closing:], closingDelimiter) + if closingOffset < 0 { + return false + } + closing += closingOffset + if closing+1 < len(message) && message[closing+1] == closingDelimiter { + closing += 2 + continue + } + break + } + if closing >= len(message) { + return false + } + if token.start > opening && token.end <= closing && + duckLakeSpanHasDottedIdentifierContext(message, opening, closing+1) { + return true + } + offset = closing + 1 + } + return false +} + +func duckLakeDotsLeadToIdentifier(message string, start int) bool { + for start < len(message) && message[start] == '.' { + start++ + } + return start < len(message) && + (isDuckLakeTokenByte(message[start]) || strings.ContainsRune("_=/\\?&#-", rune(message[start]))) +} + +func duckLakeTokenIsPathComponent(message string, token duckLakeMessageToken) bool { + if token.start > 0 { + switch message[token.start-1] { + case '/', '\\', '?', '&', '#': + return true + } + } + return false +} + +func duckLakeTokenIsPathLike(message string, token duckLakeMessageToken) bool { + if duckLakeTokenIsPathComponent(message, token) { + return true + } + if token.end >= len(message) { + return false + } + switch message[token.end] { + case '/', '\\', '?', '&', '#': + return true + default: + return false + } +} + +func isNilDuckLakeError(err error) bool { + if err == nil { + return false + } + value := reflect.ValueOf(err) + switch value.Kind() { + case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Ptr, reflect.Slice: + return value.IsNil() + default: + return false + } } func newDuckLakeRuntime(config configuration.DuckLakeConfig) (*duckLakeRuntime, error) { @@ -273,12 +1739,25 @@ func (rt *duckLakeRuntime) initializeForConn(ctx context.Context, conn driver.Co if err := VerifyDuckLakeExtensionsForTarget(rt.config.ExtensionDir, rt.manifest, runtime.GOOS, runtime.GOARCH); err != nil { return fmt.Errorf("ducklake extension integrity check failed") } + if err := rt.initializeSQLLocked(ctx, key, execer); err != nil { + return err + } + if key != nil { + rt.initialized.Store(key, struct{}{}) + } + return nil +} +// initializeSQLLocked executes the fixed service setup sequence. The caller +// must hold initializeMu and must have completed extension verification. It is +// kept separate so the sequence can be tested with a scripted executor without +// weakening the production integrity gate. +func (rt *duckLakeRuntime) initializeSQLLocked(ctx context.Context, key any, execer driver.ExecerContext) error { for _, artifact := range rt.manifest { path := filepath.Join(rt.config.ExtensionDir, artifact.FileName) query := "LOAD '" + strings.ReplaceAll(path, "'", "''") + "'" if _, err := execer.ExecContext(ctx, query, nil); err != nil && !isAlreadyLoadedError(err) { - return fmt.Errorf("load duckdb extension %s failed", artifact.Name) + return newDuckLakeInitError(duckLakeStageLoad, artifact.Name, err) } } @@ -293,12 +1772,83 @@ func (rt *duckLakeRuntime) initializeForConn(ctx context.Context, conn driver.Co } query += ")" if _, err := execer.ExecContext(ctx, query, nil); err != nil { - // Do not wrap the driver error: DuckDB may echo SQL fragments and a - // service credential in an extension/provider error. - return fmt.Errorf("create ducklake service secret failed") + return newDuckLakeInitError(duckLakeStageCreateSecret, "", err) + } + // A deployment may enable only the extension/S3 service layer (the #71 + // configuration). Attach the lake lazily in that case; object-table callers + // use EnsureAttached, which fails closed if either path is absent. When both + // paths are configured, attaching here guarantees that every eligible + // physical connection is ready before a transaction can begin. + if rt.config.MetadataPath != "" && rt.config.DataPath != "" { + if err := rt.attachLocked(ctx, key, execer); err != nil { + return err + } + } + return nil +} + +// EnsureAttached initializes and attaches the service lake on an eligible +// physical connection. It is the object-table boundary: incomplete service +// paths are rejected here rather than guessed from SQL or process state. +func (rt *duckLakeRuntime) EnsureAttached(ctx context.Context, conn driver.Conn, execer driver.ExecerContext) error { + if rt == nil { + return fmt.Errorf("%w: DuckLake service configuration is disabled", ErrInvalidTableStorage) + } + if !mycontext.IsDuckLakeEligibleQuery(ctx) { + return fmt.Errorf("ducklake is unavailable for this query origin") + } + if strings.TrimSpace(rt.config.MetadataPath) == "" || strings.TrimSpace(rt.config.DataPath) == "" { + return fmt.Errorf("ducklake metadata and data paths are required for object tables") + } + if err := rt.initializeForConn(ctx, conn, execer); err != nil { + return err } + // initializeForConn attaches when both paths are configured. The fallback + // below covers a nil/non-comparable driver identity used by unit tests. + var key any + if conn != nil { + typ := reflect.TypeOf(conn) + if typ != nil && typ.Comparable() { + key = conn + } + } + rt.initializeMu.Lock() + defer rt.initializeMu.Unlock() if key != nil { - rt.initialized.Store(key, struct{}{}) + if _, ok := rt.attached.Load(key); ok { + return nil + } + } + if err := rt.attachLocked(ctx, key, execer); err != nil { + return err + } + return nil +} + +func (rt *duckLakeRuntime) attachLocked(ctx context.Context, key any, execer driver.ExecerContext) error { + if key != nil { + if _, ok := rt.attached.Load(key); ok { + return nil + } + } + metadata := strings.TrimSpace(rt.config.MetadataPath) + dataPath := strings.TrimSpace(rt.config.DataPath) + if metadata == "" || dataPath == "" { + return fmt.Errorf("ducklake metadata and data paths are required for object tables") + } + if duckLakeRemoteCatalogURI(metadata) { + return newDuckLakeInitError(duckLakeStageAttach, "", fmt.Errorf("invalid configuration: ducklake catalog must be a local path")) + } + // Both values have already passed configuration validation. SQL-literal + // quoting is still required because service paths can contain apostrophes. + attach := "ATTACH IF NOT EXISTS " + duckDBStringLiteral("ducklake:"+metadata) + + " AS " + QuoteIdentifierANSI(DuckLakeCatalogName) + + " (DATA_PATH " + duckDBStringLiteral(dataPath) + ", DATA_INLINING_ROW_LIMIT 0, CREATE_IF_NOT_EXISTS true)" + if _, err := execer.ExecContext(ctx, attach, nil); err != nil { + return newDuckLakeInitError(duckLakeStageAttach, "", err) + } + if key != nil { + rt.attached.Store(key, struct{}{}) } return nil } @@ -312,9 +1862,30 @@ func (rt *duckLakeRuntime) resetInitialized() { } rt.initializeMu.Lock() rt.initialized.Clear() + rt.attached.Clear() rt.initializeMu.Unlock() } +func duckLakeRemoteCatalogURI(path string) bool { + lower := strings.ToLower(strings.TrimSpace(path)) + return strings.HasPrefix(lower, "s3://") || strings.HasPrefix(lower, "http://") || strings.HasPrefix(lower, "https://") +} + +func localDuckLakeCatalogPath() (string, error) { + dir := strings.TrimSpace(os.Getenv("DATA_PATH")) + if dir == "" { + home := strings.TrimSpace(os.Getenv("HOME")) + if home == "" { + return "", fmt.Errorf("ducklake local catalog directory is unavailable") + } + dir = filepath.Join(home, "data") + } + if err := os.MkdirAll(dir, 0o755); err != nil { + return "", err + } + return filepath.Join(dir, "ducklake-catalog.duckdb"), nil +} + func duckDBStringLiteral(value string) string { // Ordinary DuckDB string literals preserve backslashes; only a single quote // terminates the literal and therefore needs SQL-standard doubling. diff --git a/catalog/ducklake_error_test.go b/catalog/ducklake_error_test.go new file mode 100644 index 00000000..4ad25a92 --- /dev/null +++ b/catalog/ducklake_error_test.go @@ -0,0 +1,1316 @@ +package catalog + +import ( + "context" + "database/sql/driver" + stderrors "errors" + "fmt" + "strings" + "testing" + + "github.com/apecloud/myduckserver/configuration" + "github.com/duckdb/duckdb-go/v2" + "github.com/jackc/pgx/v5/pgconn" + "github.com/stretchr/testify/require" +) + +func TestDuckLakeInitErrorKeepsStageAndRedactsDriverDetails(t *testing.T) { + const ( + metadata = "s3://private-bucket/catalog.ducklake" + dataPath = "s3://private-bucket/data" + endpoint = "https://private-storage.example:9443" + access = "access-key-77" + secret = "secret-key-77" + query = "ATTACH 'ducklake:" + metadata + "' AS \"__myduck_ducklake\" (DATA_PATH '" + dataPath + "')" + ) + + tests := []struct { + name string + stage duckLakeInitStage + extension string + cause error + reason duckLakeFailureReason + operation string + }{ + { + name: "load io", + stage: duckLakeStageLoad, + extension: "httpfs", + cause: &duckdb.Error{ + Type: duckdb.ErrorTypeIO, + Msg: `IO Error: failed to open '/usr/local/lib/myduck/httpfs.duckdb_extension'`, + }, + reason: duckLakeReasonIOFailure, + operation: "load duckdb extension httpfs failed", + }, + { + name: "create secret permission", + stage: duckLakeStageCreateSecret, + cause: &duckdb.Error{ + Type: duckdb.ErrorTypePermission, + Msg: "Permission Error: endpoint=" + endpoint + " KEY_ID='" + access + "' SECRET='" + secret + "'", + }, + reason: duckLakeReasonPermissionDenied, + operation: "create ducklake service secret failed", + }, + { + name: "attach http", + stage: duckLakeStageAttach, + cause: &duckdb.Error{ + Type: duckdb.ErrorTypeHTTP, + Msg: "HTTP Error: request failed for " + endpoint + " while executing " + query, + }, + reason: duckLakeReasonHTTPFailure, + operation: "attach ducklake catalog failed", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := newDuckLakeInitError(tt.stage, tt.extension, tt.cause) + require.Error(t, err) + + var stageErr *duckLakeInitError + require.ErrorAs(t, err, &stageErr) + require.Equal(t, tt.stage, stageErr.stage) + require.Equal(t, tt.reason, stageErr.reason) + require.Same(t, tt.cause, stageErr.rawCause()) + require.True(t, stderrors.Is(err, tt.cause)) + require.Contains(t, err.Error(), "stage="+string(tt.stage)) + require.Contains(t, err.Error(), "reason="+string(tt.reason)) + require.True(t, strings.HasPrefix(err.Error(), tt.operation+" ")) + require.Equal(t, err.Error(), fmt.Sprintf("%+v", err)) + require.Equal(t, err.Error(), fmt.Sprintf("%s", err)) + require.Equal(t, fmt.Sprintf("%q", err.Error()), fmt.Sprintf("%q", err)) + require.Equal(t, err.Error(), fmt.Sprintf("%#v", err)) + var rawDuck *duckdb.Error + require.False(t, stderrors.As(err, &rawDuck)) + + for _, forbidden := range []string{metadata, dataPath, endpoint, access, secret, query, "/usr/local/lib/myduck"} { + require.NotContains(t, err.Error(), forbidden) + require.NotContains(t, fmt.Sprintf("%+v", err), forbidden) + } + }) + } +} + +func TestDuckLakeInitErrorDoesNotExposeProtocolErrorCause(t *testing.T) { + raw := &pgconn.PgError{ + Severity: "ERROR", + Code: "XX000", + Message: "remote endpoint=https://private-storage.example:9443 SECRET=secret-key-77", + } + err := newDuckLakeInitError(duckLakeStageAttach, "", raw) + + var exposed *pgconn.PgError + require.False(t, stderrors.As(err, &exposed)) + require.Nil(t, stderrors.Unwrap(err)) + require.True(t, stderrors.Is(err, raw)) + wrapper := fmt.Errorf("outer wrapper: %w", err) + require.Contains(t, wrapper.Error(), err.Error()) + require.True(t, stderrors.Is(wrapper, raw)) + require.False(t, stderrors.As(wrapper, &exposed)) + require.NotContains(t, err.Error(), "private-storage.example") + require.NotContains(t, err.Error(), "secret-key-77") +} + +func TestDuckLakeInitErrorClassifiesSentinelsWithoutRawFormatting(t *testing.T) { + tests := []struct { + name string + cause error + reason duckLakeFailureReason + }{ + {name: "canceled", cause: context.Canceled, reason: duckLakeReasonContextCanceled}, + {name: "deadline", cause: context.DeadlineExceeded, reason: duckLakeReasonDeadlineExceeded}, + {name: "bad connection", cause: driver.ErrBadConn, reason: duckLakeReasonBadConnection}, + {name: "unknown details", cause: fmt.Errorf("driver exploded at /var/lib/myduck with SECRET=secret-key-77"), reason: duckLakeReasonDriverError}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := newDuckLakeInitError(duckLakeStageAttach, "", tt.cause) + require.Contains(t, err.Error(), "reason="+string(tt.reason)) + require.True(t, stderrors.Is(err, tt.cause)) + require.NotContains(t, err.Error(), "secret-key-77") + require.NotContains(t, err.Error(), "/var/lib/myduck") + }) + } +} + +func TestDuckLakeAttachFailureUsesStageError(t *testing.T) { + const ( + metadata = "/var/lib/myduck/catalog.ducklake" + leaked = "s3://private-bucket/catalog.ducklake" + dataPath = "s3://private-bucket/data" + endpoint = "https://private-storage.example:9443" + secret = "secret-key-77" + ) + raw := &duckdb.Error{ + Type: duckdb.ErrorTypePermission, + Msg: "Permission Error: endpoint=" + endpoint + " metadata=" + leaked + " data=" + dataPath + " SECRET='" + secret + "'", + } + runtime := &duckLakeRuntime{config: configuration.DuckLakeConfig{ + MetadataPath: metadata, + DataPath: dataPath, + }} + execer := &recordingDuckLakeExecer{err: raw} + + err := runtime.attachLocked(context.Background(), nil, execer) + require.Error(t, err) + require.ErrorIs(t, err, raw) + require.ErrorContains(t, err, "attach ducklake catalog failed") + require.ErrorContains(t, err, "stage=attach") + require.ErrorContains(t, err, "reason=permission_denied") + require.Len(t, execer.queries, 1) + for _, forbidden := range []string{leaked, dataPath, endpoint, secret} { + require.NotContains(t, err.Error(), forbidden) + } +} + +func TestDuckLakeAttachFailureClosedSubclasses(t *testing.T) { + const ( + metadata = "s3://private-bucket/catalog.ducklake" + secret = "secret-key-77" + endpoint = "https://private-storage.example:9443" + query = "ATTACH 'ducklake:" + metadata + "' AS \"__myduck_ducklake\"" + ) + tests := []struct { + name string + cause error + reason duckLakeFailureReason + }{ + { + name: "permission", + cause: &duckdb.Error{ + Type: duckdb.ErrorTypePermission, + Msg: "Permission Error: endpoint=" + endpoint + " SECRET=" + secret, + }, + reason: duckLakeReasonPermission, + }, + { + name: "permission from HTTP status", + cause: &duckdb.Error{ + Type: duckdb.ErrorTypeHTTP, + Msg: "HTTP Error: 403 for " + metadata, + }, + reason: duckLakeReasonPermission, + }, + { + name: "http", + cause: &duckdb.Error{ + Type: duckdb.ErrorTypeHTTP, + Msg: "HTTP Error: status 503 while executing " + query, + }, + reason: duckLakeReasonHTTP, + }, + { + name: "network", + cause: &duckdb.Error{ + Type: duckdb.ErrorTypeNetwork, + Msg: "Network Error: dial tcp " + endpoint, + }, + reason: duckLakeReasonNetwork, + }, + { + name: "missing object from HTTP 404", + cause: &duckdb.Error{ + Type: duckdb.ErrorTypeHTTP, + Msg: "HTTP Error: 404 Not Found for " + metadata, + }, + reason: duckLakeReasonMissingObject, + }, + { + name: "missing object from S3 code", + cause: &duckdb.Error{ + Type: duckdb.ErrorTypeIO, + Msg: "IO Error: S3 NoSuchKey for " + metadata, + }, + reason: duckLakeReasonMissingObject, + }, + { + name: "typed IO read-only missing catalog is missing object", + cause: &duckdb.Error{ + Type: duckdb.ErrorTypeIO, + Msg: `IO Error: Failed to attach DuckLake MetaData in read-only mode: database does not exist`, + }, + reason: duckLakeReasonMissingObject, + }, + { + name: "extension", + cause: &duckdb.Error{ + Type: duckdb.ErrorTypeMissingExtension, + Msg: "Missing Extension Error: ducklake", + }, + reason: duckLakeReasonExtension, + }, + { + name: "generic IO", + cause: &duckdb.Error{ + Type: duckdb.ErrorTypeIO, + Msg: "IO Error: failed to read remote catalog at " + endpoint, + }, + reason: duckLakeReasonGenericIO, + }, + { + name: "typed IO local file remains generic", + cause: &duckdb.Error{ + Type: duckdb.ErrorTypeIO, + Msg: "IO Error: no such file or directory", + }, + reason: duckLakeReasonGenericIO, + }, + { + name: "typed IO HTTP 403 is permission", + cause: &duckdb.Error{ + Type: duckdb.ErrorTypeIO, + Msg: "IO Error: HTTP 403 Forbidden for " + metadata + " SECRET=" + secret, + }, + reason: duckLakeReasonPermission, + }, + { + name: "typed IO connection refused is network", + cause: &duckdb.Error{ + Type: duckdb.ErrorTypeIO, + Msg: "IO Error: Connection refused connecting to " + endpoint, + }, + reason: duckLakeReasonNetwork, + }, + { + name: "typed IO TLS handshake is tls", + cause: &duckdb.Error{ + Type: duckdb.ErrorTypeIO, + Msg: "IO Error: TLS handshake failed with " + endpoint, + }, + reason: duckLakeReasonTLSFailure, + }, + { + name: "typed IO timeout is timeout", + cause: &duckdb.Error{ + Type: duckdb.ErrorTypeIO, + Msg: "IO Error: network request timed out contacting " + endpoint, + }, + reason: duckLakeReasonTimeout, + }, + } + + seen := make(map[duckLakeFailureReason]struct{}, len(tests)) + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := newDuckLakeInitError(duckLakeStageAttach, "", tt.cause) + require.Error(t, err) + + var initErr *duckLakeInitError + require.ErrorAs(t, err, &initErr) + require.Equal(t, duckLakeStageAttach, initErr.stage) + require.Equal(t, tt.reason, initErr.reason) + require.Contains(t, err.Error(), "stage=attach") + require.Contains(t, err.Error(), "reason="+string(tt.reason)) + require.NotContains(t, err.Error(), metadata) + require.NotContains(t, err.Error(), endpoint) + require.NotContains(t, err.Error(), secret) + require.NotContains(t, err.Error(), query) + for _, format := range []string{"%v", "%+v", "%#v", "%q"} { + formatted := fmt.Sprintf(format, err) + require.NotContains(t, formatted, metadata) + require.NotContains(t, formatted, endpoint) + require.NotContains(t, formatted, secret) + require.NotContains(t, formatted, query) + } + seen[tt.reason] = struct{}{} + }) + } + require.Len(t, seen, 8, "typed IO must keep permission, network, TLS, timeout, HTTP, missing-object, extension, and generic IO distinct") +} + +func TestDuckLakeAttachFailureClassifiesHTTPStatusesAndStage(t *testing.T) { + tests := []struct { + name string + stage duckLakeInitStage + cause error + want duckLakeFailureReason + }{ + { + name: "bare 404 is missing object for attach", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "HTTP Error: 404"}, + want: duckLakeReasonMissingObject, + }, + { + name: "404 not found is missing object for attach", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "404 Not Found"}, + want: duckLakeReasonMissingObject, + }, + { + name: "terminal period after 404 is punctuation", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "HTTP Error: 404."}, + want: duckLakeReasonMissingObject, + }, + { + name: "dotted 404 suffix remains HTTP", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "HTTP Error: 404.txt"}, + want: duckLakeReasonHTTP, + }, + { + name: "identifier-like HTTP marker remains broad HTTP", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "HTTP_Error:404"}, + want: duckLakeReasonHTTP, + }, + { + name: "identifier-like HTTP version remains broad HTTP", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "foo_http/1.1 404"}, + want: duckLakeReasonHTTP, + }, + { + name: "status assignment with terminal period is missing object", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "status=404."}, + want: duckLakeReasonMissingObject, + }, + { + name: "status assignment with terminal ellipsis is missing object", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "status=404..."}, + want: duckLakeReasonMissingObject, + }, + { + name: "status assignment after a sentence is missing object", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "request failed. status=404"}, + want: duckLakeReasonMissingObject, + }, + { + name: "status assignment with dotted suffix remains HTTP", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "status=404.txt"}, + want: duckLakeReasonHTTP, + }, + { + name: "status assignment with multi-dot suffix remains HTTP", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "status=404..txt"}, + want: duckLakeReasonHTTP, + }, + { + name: "status assignment with dot path suffix remains HTTP", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "status=404../path"}, + want: duckLakeReasonHTTP, + }, + { + name: "bare typed 404 is missing object for attach", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "404"}, + want: duckLakeReasonMissingObject, + }, + { + name: "numeric status in an unrelated identifier remains HTTP", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "zipcode 404"}, + want: duckLakeReasonHTTP, + }, + { + name: "numeric status in a path assignment remains HTTP", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "path=404"}, + want: duckLakeReasonHTTP, + }, + { + name: "numeric status in a path status segment remains HTTP", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "path/status=404"}, + want: duckLakeReasonHTTP, + }, + { + name: "numeric permission status in a path code segment remains HTTP", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "path/code=403"}, + want: duckLakeReasonHTTP, + }, + { + name: "numeric status in an unknown assignment prefix remains HTTP", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "foo=status=404"}, + want: duckLakeReasonHTTP, + }, + { + name: "numeric permission status in an unknown assignment prefix remains HTTP", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "foo=code=403"}, + want: duckLakeReasonHTTP, + }, + { + name: "spaced unknown status assignment prefix remains HTTP", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "foo = status=404"}, + want: duckLakeReasonHTTP, + }, + { + name: "spaced unknown permission assignment prefix remains HTTP", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "foo = code=403"}, + want: duckLakeReasonHTTP, + }, + { + name: "status code joined to number by underscore remains HTTP", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "status_code_404"}, + want: duckLakeReasonHTTP, + }, + { + name: "status code joined to number by dot remains HTTP", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "status_code.404"}, + want: duckLakeReasonHTTP, + }, + { + name: "repeated status field separator remains HTTP", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "status__code=404"}, + want: duckLakeReasonHTTP, + }, + { + name: "prefixed status field remains HTTP", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "_status_code=404"}, + want: duckLakeReasonHTTP, + }, + { + name: "compact object marker in an identifier remains HTTP", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "foo_nosuchkey"}, + want: duckLakeReasonHTTP, + }, + { + name: "compact object marker assignment remains HTTP", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "foo=nosuchkey"}, + want: duckLakeReasonHTTP, + }, + { + name: "compact object marker identifier before status remains HTTP", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "foo_nosuchkey 403"}, + want: duckLakeReasonHTTP, + }, + { + name: "standalone code assignment is permission", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "code=403"}, + want: duckLakeReasonPermission, + }, + { + name: "leading 404 with remote payload is missing object for attach", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "404 for s3://private/catalog.ducklake"}, + want: duckLakeReasonMissingObject, + }, + { + name: "hyphenated 404 identifier remains HTTP", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "404-not-found"}, + want: duckLakeReasonHTTP, + }, + { + name: "NoSuchKey is missing object for attach", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "NoSuchKey"}, + want: duckLakeReasonMissingObject, + }, + { + name: "NoSuchKey path component remains HTTP", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "/NoSuchKey"}, + want: duckLakeReasonHTTP, + }, + { + name: "NoSuchKey path prefix remains HTTP", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "NoSuchKey/object"}, + want: duckLakeReasonHTTP, + }, + { + name: "NoSuchKey dot path suffix remains HTTP", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "NoSuchKey../path"}, + want: duckLakeReasonHTTP, + }, + { + name: "404 path prefix remains HTTP", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "HTTP Error: 404/object"}, + want: duckLakeReasonHTTP, + }, + { + name: "bare 404 remains HTTP outside attach", + stage: duckLakeStageCreateSecret, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "HTTP Error: 404"}, + want: duckLakeReasonHTTP, + }, + { + name: "bare typed 404 remains HTTP outside attach", + stage: duckLakeStageCreateSecret, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "404"}, + want: duckLakeReasonHTTP, + }, + { + name: "NoSuchKey remains HTTP outside attach", + stage: duckLakeStageLoad, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "NoSuchKey"}, + want: duckLakeReasonHTTP, + }, + { + name: "bare 403 is permission", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "HTTP Error: 403"}, + want: duckLakeReasonPermission, + }, + { + name: "403 after NoSuchKey is still permission", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "HTTP Error: NoSuchKey 403"}, + want: duckLakeReasonPermission, + }, + { + name: "401 after NoSuchKey is still permission", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "HTTP Error: NoSuchKey 401"}, + want: duckLakeReasonPermission, + }, + { + name: "remote typed IO no such file is missing object", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeIO, Msg: "IO Error: no such file at s3://private/catalog.ducklake"}, + want: duckLakeReasonMissingObject, + }, + { + name: "assigned remote typed IO target is missing object", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeIO, Msg: "IO Error: no such file at url=s3://private/catalog.ducklake"}, + want: duckLakeReasonMissingObject, + }, + { + name: "repeated remote typed IO assignment stays generic IO", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeIO, Msg: "IO Error: no such file at url==s3://private/catalog.ducklake"}, + want: duckLakeReasonGenericIO, + }, + { + name: "typed S3 error context is missing object", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeIO, Msg: "S3 Error: no such file"}, + want: duckLakeReasonMissingObject, + }, + { + name: "local typed IO no such file is generic IO", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeIO, Msg: "IO Error: no such file or directory"}, + want: duckLakeReasonGenericIO, + }, + { + name: "local path words do not imply remote object", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeIO, Msg: "IO Error: no such file at /var/lib/object/catalog"}, + want: duckLakeReasonGenericIO, + }, + { + name: "unrelated documentation URL does not make local file remote", + stage: duckLakeStageAttach, + cause: &duckdb.Error{Type: duckdb.ErrorTypeIO, Msg: "IO Error: no such file at /var/lib/catalog; see https://docs.example/help"}, + want: duckLakeReasonGenericIO, + }, + { + name: "extension missing lexical marker", + stage: duckLakeStageLoad, + cause: stderrors.New("extension httpfs missing"), + want: duckLakeReasonExtension, + }, + { + name: "extension cannot find lexical marker", + stage: duckLakeStageLoad, + cause: stderrors.New("cannot find extension httpfs"), + want: duckLakeReasonExtension, + }, + { + name: "attach extension wording is not trusted without load context", + stage: duckLakeStageAttach, + cause: stderrors.New("remote extension missing"), + want: duckLakeReasonDriverError, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := newDuckLakeInitError(tt.stage, "httpfs", tt.cause) + var initErr *duckLakeInitError + require.ErrorAs(t, err, &initErr) + require.Equal(t, tt.want, initErr.reason) + require.Contains(t, err.Error(), "reason="+string(tt.want)) + }) + } +} + +func TestDuckLakeAttachFailureRefinesWrappedTypedMissingObject(t *testing.T) { + const ( + metadata = "s3://private-bucket/catalog.ducklake" + dataPath = "s3://private-bucket/data" + secret = "secret-key-77" + ) + raw := &duckdb.Error{ + Type: duckdb.ErrorTypeHTTP, + Msg: "HTTP Error: status code: 404 (Not Found) metadata=" + metadata + " SECRET=" + secret, + } + wrapped := fmt.Errorf("driver wrapper: %w", raw) + err := newDuckLakeInitError(duckLakeStageAttach, "", wrapped) + + var initErr *duckLakeInitError + require.ErrorAs(t, err, &initErr) + require.Equal(t, duckLakeReasonMissingObject, initErr.reason) + require.Same(t, wrapped, initErr.rawCause()) + require.True(t, stderrors.Is(err, raw)) + require.Nil(t, stderrors.Unwrap(err)) + var exposed *duckdb.Error + require.False(t, stderrors.As(err, &exposed)) + for _, forbidden := range []string{metadata, dataPath, secret, "driver wrapper", raw.Msg} { + require.NotContains(t, err.Error(), forbidden) + require.NotContains(t, fmt.Sprintf("%+v", err), forbidden) + } +} + +func TestDuckLakeAttachFailureKeepsLexicalMarkersOutOfURLs(t *testing.T) { + tests := []struct { + name string + message string + want duckLakeFailureReason + }{ + {name: "zipcode is not a status marker", message: "zipcode 404", want: duckLakeReasonDriverError}, + {name: "hyphenated identifier is not a status marker", message: "zip-code 404", want: duckLakeReasonDriverError}, + {name: "hyphenated permission phrase is an identifier", message: "access-denied", want: duckLakeReasonDriverError}, + {name: "hyphenated missing phrase is an identifier", message: "object-not-found", want: duckLakeReasonDriverError}, + {name: "hyphenated network phrase is an identifier", message: "connection-refused", want: duckLakeReasonDriverError}, + {name: "hyphenated network error identifier is opaque", message: "network-error", want: duckLakeReasonDriverError}, + {name: "reversed hyphenated network identifier is opaque", message: "error-network", want: duckLakeReasonDriverError}, + {name: "hyphenated extension identifier is opaque", message: "extension-not-loaded", want: duckLakeReasonDriverError}, + {name: "hyphenated extension status suffix is opaque", message: "extension not-loaded", want: duckLakeReasonDriverError}, + {name: "reversed hyphenated extension identifier is opaque", message: "failed-to-load-extension", want: duckLakeReasonDriverError}, + {name: "hyphenated autoload identifier is opaque", message: "autoload-extension", want: duckLakeReasonDriverError}, + {name: "hyphenated resource identifier is opaque", message: "resource-unavailable", want: duckLakeReasonDriverError}, + {name: "hyphenated bare not-found phrase is opaque", message: "not-found", want: duckLakeReasonDriverError}, + {name: "single missing code with leading hyphen is opaque", message: "foo-nosuchkey", want: duckLakeReasonDriverError}, + {name: "single missing code with trailing hyphen is opaque", message: "nosuchkey-foo", want: duckLakeReasonDriverError}, + {name: "single object code with leading hyphen is opaque", message: "foo-objectnotfound", want: duckLakeReasonDriverError}, + {name: "single key code with trailing hyphen is opaque", message: "keynotfound-foo", want: duckLakeReasonDriverError}, + {name: "status code marker", message: "status_code=404", want: duckLakeReasonMissingObject}, + {name: "response code marker", message: "response_code=404", want: duckLakeReasonMissingObject}, + {name: "HTTP version marker", message: "HTTP/1.1 404 Not Found", want: duckLakeReasonMissingObject}, + {name: "leading status with payload", message: "404 for s3://private/catalog.ducklake", want: duckLakeReasonMissingObject}, + {name: "hyphenated leading status is opaque", message: "404-not-found", want: duckLakeReasonDriverError}, + {name: "plain status phrase", message: "404 Not Found", want: duckLakeReasonMissingObject}, + {name: "plain HTTP failure", message: "HTTP 503", want: duckLakeReasonHTTPFailure}, + {name: "network hostname is not a network error", message: "GET https://network.example/object", want: duckLakeReasonDriverError}, + {name: "permission hostname is not a permission error", message: "GET https://forbidden.example/object", want: duckLakeReasonDriverError}, + {name: "permission hostname with hyphen is not a permission error", message: "GET https://access-denied.example/object", want: duckLakeReasonDriverError}, + {name: "missing hostname is not an unavailable resource", message: "GET https://missing.example/object", want: duckLakeReasonDriverError}, + {name: "URL permission marker does not hide later permission", message: "GET https://access-denied.example/object; access denied while reading catalog", want: duckLakeReasonPermissionDenied}, + {name: "URL missing marker does not hide later missing object", message: "GET https://object-not-found.example/object; object not found", want: duckLakeReasonMissingObject}, + {name: "URL network marker does not hide later network error", message: "GET https://connection-refused.example/object; network error while reading catalog", want: duckLakeReasonNetworkFailure}, + {name: "network phrase", message: "network error while reading catalog", want: duckLakeReasonNetworkFailure}, + {name: "permission phrase", message: "access denied while reading catalog", want: duckLakeReasonPermissionDenied}, + {name: "unauthorized access phrase", message: "Unauthorized access", want: duckLakeReasonPermissionDenied}, + {name: "operation forbidden phrase", message: "Operation forbidden", want: duckLakeReasonPermissionDenied}, + {name: "remote extension wording is an object miss", message: "remote extension missing object", want: duckLakeReasonMissingObject}, + {name: "remote extension wording alone is opaque", message: "remote extension missing", want: duckLakeReasonDriverError}, + {name: "known remote extension wording alone is opaque", message: "remote ducklake extension unavailable", want: duckLakeReasonDriverError}, + {name: "known remote extension URL wording alone is opaque", message: "remote extension https://foo/httpfs missing", want: duckLakeReasonDriverError}, + {name: "explicit attach extension load wording is classified", message: "failed to load ducklake extension", want: duckLakeReasonExtensionUnavailable}, + {name: "local path words do not imply missing object", message: "no such file at /var/lib/object/catalog", want: duckLakeReasonResourceUnavailable}, + {name: "plain HTTP status remains an object miss", message: "404 Not Found", want: duckLakeReasonMissingObject}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := newDuckLakeInitError(duckLakeStageAttach, "", stderrors.New(tt.message)) + var initErr *duckLakeInitError + require.ErrorAs(t, err, &initErr) + require.Equal(t, tt.want, initErr.reason) + }) + } +} + +func TestDuckLakeLexicalMarkerHyphenBoundaries(t *testing.T) { + tests := []struct { + name string + stage duckLakeInitStage + message string + want duckLakeFailureReason + }{ + {name: "404 reason phrase joined by hyphen", stage: duckLakeStageAttach, message: "404 not-found", want: duckLakeReasonDriverError}, + {name: "401 permission phrase joined by hyphen", stage: duckLakeStageAttach, message: "401 unauthorized-error", want: duckLakeReasonDriverError}, + {name: "403 permission phrase joined by hyphen", stage: duckLakeStageAttach, message: "403 forbidden-error", want: duckLakeReasonDriverError}, + {name: "500 error context joined by hyphen", stage: duckLakeStageAttach, message: "500 error-code", want: duckLakeReasonDriverError}, + {name: "500 unavailable context joined by hyphen", stage: duckLakeStageAttach, message: "500 unavailable-now", want: duckLakeReasonDriverError}, + {name: "leading status context joined by hyphen", stage: duckLakeStageAttach, message: "404 for-s3", want: duckLakeReasonDriverError}, + {name: "leading object context joined by hyphen", stage: duckLakeStageAttach, message: "404 object-not-found", want: duckLakeReasonDriverError}, + {name: "hyphenated object code before status", stage: duckLakeStageAttach, message: "foo-nosuchkey 403", want: duckLakeReasonDriverError}, + {name: "object code in URL before status", stage: duckLakeStageAttach, message: "https://example.test/nosuchkey 403", want: duckLakeReasonDriverError}, + {name: "permission context suffix joined by hyphen", stage: duckLakeStageAttach, message: "unauthorized error-code", want: duckLakeReasonDriverError}, + {name: "permission context prefix joined by hyphen", stage: duckLakeStageAttach, message: "error-code unauthorized", want: duckLakeReasonDriverError}, + {name: "permission marker in an absolute path is opaque", stage: duckLakeStageAttach, message: "/unauthorized", want: duckLakeReasonDriverError}, + {name: "permission marker in a trailing path is opaque", stage: duckLakeStageAttach, message: "forbidden/", want: duckLakeReasonDriverError}, + {name: "permission marker in a query is opaque", stage: duckLakeStageAttach, message: "?unauthorized", want: duckLakeReasonDriverError}, + {name: "permission marker in an identifier is opaque", stage: duckLakeStageAttach, message: "foo_unauthorized", want: duckLakeReasonDriverError}, + {name: "status marker in an unknown assignment is opaque", stage: duckLakeStageAttach, message: "foo=status=404", want: duckLakeReasonDriverError}, + {name: "permission status in an unknown assignment is opaque", stage: duckLakeStageAttach, message: "foo=code=403", want: duckLakeReasonDriverError}, + {name: "compact missing marker in an identifier is opaque", stage: duckLakeStageAttach, message: "foo_nosuchkey", want: duckLakeReasonDriverError}, + {name: "compact missing marker in an assignment is opaque", stage: duckLakeStageAttach, message: "foo=nosuchkey", want: duckLakeReasonDriverError}, + {name: "compact missing marker before permission status is opaque", stage: duckLakeStageAttach, message: "foo_nosuchkey 403", want: duckLakeReasonDriverError}, + {name: "permission access context joined by hyphen", stage: duckLakeStageAttach, message: "forbidden access-denied", want: duckLakeReasonDriverError}, + {name: "permission actor context joined by hyphen", stage: duckLakeStageAttach, message: "unauthorized by-user", want: duckLakeReasonDriverError}, + {name: "load extension status joined by hyphen", stage: duckLakeStageLoad, message: "extension unavailable-now", want: duckLakeReasonDriverError}, + {name: "load extension autoload identifier", stage: duckLakeStageLoad, message: "extension autoload-extension", want: duckLakeReasonDriverError}, + {name: "load extension failure identifier", stage: duckLakeStageLoad, message: "extension failed-to-load", want: duckLakeReasonDriverError}, + {name: "load extension cannot identifier", stage: duckLakeStageLoad, message: "extension cannot-load", want: duckLakeReasonDriverError}, + {name: "load token in URL is not extension context", stage: duckLakeStageLoad, message: "extension https://example.test/load", want: duckLakeReasonDriverError}, + {name: "attach extension status joined by hyphen", stage: duckLakeStageAttach, message: "extension not-loaded", want: duckLakeReasonDriverError}, + {name: "remote identifier is not object context", stage: duckLakeStageAttach, message: "remote-object no such file", want: duckLakeReasonResourceUnavailable}, + {name: "s3 identifier is not object context", stage: duckLakeStageAttach, message: "s3-bucket no such file", want: duckLakeReasonResourceUnavailable}, + {name: "plain permission context remains classified", stage: duckLakeStageAttach, message: "unauthorized access", want: duckLakeReasonPermissionDenied}, + {name: "plain status reason remains classified", stage: duckLakeStageAttach, message: "404 not found", want: duckLakeReasonMissingObject}, + {name: "plain extension status remains classified", stage: duckLakeStageLoad, message: "extension unavailable", want: duckLakeReasonExtensionUnavailable}, + {name: "plain remote context remains missing object", stage: duckLakeStageAttach, message: "remote no such file", want: duckLakeReasonMissingObject}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := newDuckLakeInitError(tt.stage, "", stderrors.New(tt.message)) + var initErr *duckLakeInitError + require.ErrorAs(t, err, &initErr) + require.Equal(t, tt.want, initErr.reason) + }) + } +} + +func TestDuckLakeLexicalMarkerPathBoundaries(t *testing.T) { + tests := []struct { + name string + stage duckLakeInitStage + message string + want duckLakeFailureReason + }{ + {name: "leading path single missing code", stage: duckLakeStageAttach, message: "/nosuchkey", want: duckLakeReasonDriverError}, + {name: "trailing path single missing code", stage: duckLakeStageAttach, message: "nosuchkey/object", want: duckLakeReasonDriverError}, + {name: "leading path compact object code", stage: duckLakeStageAttach, message: "/objectnotfound", want: duckLakeReasonDriverError}, + {name: "relative object path", stage: duckLakeStageAttach, message: "object/not/found", want: duckLakeReasonDriverError}, + {name: "absolute object path", stage: duckLakeStageAttach, message: "/object/not/found", want: duckLakeReasonDriverError}, + {name: "windows object path", stage: duckLakeStageAttach, message: `C:\object\not\found`, want: duckLakeReasonDriverError}, + {name: "permission path", stage: duckLakeStageAttach, message: "access/denied", want: duckLakeReasonDriverError}, + {name: "network path", stage: duckLakeStageAttach, message: "connection/refused", want: duckLakeReasonDriverError}, + {name: "reversed network path", stage: duckLakeStageAttach, message: "error/network", want: duckLakeReasonDriverError}, + {name: "extension path", stage: duckLakeStageAttach, message: "cannot/load/extension", want: duckLakeReasonDriverError}, + {name: "load extension path", stage: duckLakeStageLoad, message: "extension/not/loaded", want: duckLakeReasonDriverError}, + {name: "resource path", stage: duckLakeStageAttach, message: "resource/unavailable", want: duckLakeReasonDriverError}, + {name: "remote path word is not object context", stage: duckLakeStageAttach, message: "remote/path no such file", want: duckLakeReasonResourceUnavailable}, + {name: "s3 path word is not object context", stage: duckLakeStageAttach, message: "s3/path no such file", want: duckLakeReasonResourceUnavailable}, + {name: "path marker does not hide later object marker", stage: duckLakeStageAttach, message: "/object/not/found; object not found", want: duckLakeReasonMissingObject}, + {name: "path marker does not hide later network marker", stage: duckLakeStageAttach, message: "/connection/refused; network error", want: duckLakeReasonNetworkFailure}, + {name: "status assignment remains classified", stage: duckLakeStageAttach, message: "status_code=404", want: duckLakeReasonMissingObject}, + {name: "HTTP version remains classified", stage: duckLakeStageAttach, message: "HTTP/1.1 404 Not Found", want: duckLakeReasonMissingObject}, + {name: "remote URI remains classified", stage: duckLakeStageAttach, message: "no such file at s3://private/catalog.ducklake", want: duckLakeReasonMissingObject}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := newDuckLakeInitError(tt.stage, "", stderrors.New(tt.message)) + var initErr *duckLakeInitError + require.ErrorAs(t, err, &initErr) + require.Equal(t, tt.want, initErr.reason) + }) + } +} + +func TestDuckLakeLexicalMarkerPunctuationBoundaries(t *testing.T) { + tests := []struct { + name string + stage duckLakeInitStage + message string + want duckLakeFailureReason + }{ + {name: "permission sentence period", stage: duckLakeStageAttach, message: "access denied.", want: duckLakeReasonPermissionDenied}, + {name: "missing object sentence period", stage: duckLakeStageAttach, message: "object not found.", want: duckLakeReasonMissingObject}, + {name: "compact object sentence period", stage: duckLakeStageAttach, message: "NoSuchKey.", want: duckLakeReasonMissingObject}, + {name: "standalone quoted compact object marker", stage: duckLakeStageAttach, message: `"NoSuchKey"`, want: duckLakeReasonMissingObject}, + {name: "standalone quoted permission phrase", stage: duckLakeStageAttach, message: `"access denied"`, want: duckLakeReasonPermissionDenied}, + {name: "standalone quoted missing phrase", stage: duckLakeStageAttach, message: `"object not found"`, want: duckLakeReasonMissingObject}, + {name: "standalone bracketed network phrase", stage: duckLakeStageAttach, message: `[network error]`, want: duckLakeReasonNetworkFailure}, + {name: "standalone backtick extension phrase", stage: duckLakeStageLoad, message: "`extension unavailable`", want: duckLakeReasonExtensionUnavailable}, + {name: "standalone escaped quoted missing phrase", stage: duckLakeStageAttach, message: `"prefix "" object not found "" suffix"`, want: duckLakeReasonMissingObject}, + {name: "standalone nested delimited missing phrase", stage: duckLakeStageAttach, message: `[prefix "object not found" suffix]`, want: duckLakeReasonMissingObject}, + {name: "network sentence period", stage: duckLakeStageAttach, message: "network error.", want: duckLakeReasonNetworkFailure}, + {name: "connection sentence period", stage: duckLakeStageAttach, message: "connection refused.", want: duckLakeReasonNetworkFailure}, + {name: "extension sentence period", stage: duckLakeStageLoad, message: "extension unavailable.", want: duckLakeReasonExtensionUnavailable}, + {name: "resource sentence period", stage: duckLakeStageAttach, message: "resource unavailable.", want: duckLakeReasonResourceUnavailable}, + {name: "ellipsis is punctuation", stage: duckLakeStageAttach, message: "NoSuchKey...", want: duckLakeReasonMissingObject}, + {name: "dot path suffix is identifier syntax", stage: duckLakeStageAttach, message: "NoSuchKey../path", want: duckLakeReasonDriverError}, + {name: "dot query suffix is identifier syntax", stage: duckLakeStageAttach, message: "network error..?query", want: duckLakeReasonDriverError}, + {name: "period before next sentence is punctuation", stage: duckLakeStageAttach, message: "NoSuchKey. Details follow", want: duckLakeReasonMissingObject}, + {name: "leading dot is identifier syntax", stage: duckLakeStageAttach, message: ".nosuchkey", want: duckLakeReasonDriverError}, + {name: "dotted object prefix is identifier syntax", stage: duckLakeStageAttach, message: "foo.nosuchkey", want: duckLakeReasonDriverError}, + {name: "dotted object suffix is identifier syntax", stage: duckLakeStageAttach, message: "nosuchkey.foo", want: duckLakeReasonDriverError}, + {name: "dotted underscore suffix is identifier syntax", stage: duckLakeStageAttach, message: "nosuchkey._suffix", want: duckLakeReasonDriverError}, + {name: "quoted object marker with dotted suffix is identifier syntax", stage: duckLakeStageAttach, message: `"NoSuchKey".field`, want: duckLakeReasonDriverError}, + {name: "quoted object marker with dotted prefix is identifier syntax", stage: duckLakeStageAttach, message: `schema."NoSuchKey"`, want: duckLakeReasonDriverError}, + {name: "quoted dotted permission phrase is identifier syntax", stage: duckLakeStageAttach, message: `"access"."denied"`, want: duckLakeReasonDriverError}, + {name: "backtick dotted missing phrase is identifier syntax", stage: duckLakeStageAttach, message: "`object`.`not`.`found`", want: duckLakeReasonDriverError}, + {name: "bracket dotted network phrase is identifier syntax", stage: duckLakeStageAttach, message: "[network].[error]", want: duckLakeReasonDriverError}, + {name: "quoted dotted extension phrase is identifier syntax", stage: duckLakeStageLoad, message: `"extension"."unavailable"`, want: duckLakeReasonDriverError}, + {name: "quoted multiword permission identifier", stage: duckLakeStageAttach, message: `schema."access denied"`, want: duckLakeReasonDriverError}, + {name: "quoted multiword missing identifier", stage: duckLakeStageAttach, message: `"object not found".field`, want: duckLakeReasonDriverError}, + {name: "bracketed multiword network identifier", stage: duckLakeStageAttach, message: `[network error].field`, want: duckLakeReasonDriverError}, + {name: "backtick multiword extension identifier", stage: duckLakeStageLoad, message: "schema.`extension unavailable`", want: duckLakeReasonDriverError}, + {name: "escaped quoted multiword missing identifier", stage: duckLakeStageAttach, message: `"prefix "" object not found "" suffix".field`, want: duckLakeReasonDriverError}, + {name: "nested delimited multiword missing identifier", stage: duckLakeStageAttach, message: `[prefix "object not found" suffix].field`, want: duckLakeReasonDriverError}, + {name: "dotted permission phrase is identifier syntax", stage: duckLakeStageAttach, message: "access.denied", want: duckLakeReasonDriverError}, + {name: "dotted missing phrase is identifier syntax", stage: duckLakeStageAttach, message: "object.not.found", want: duckLakeReasonDriverError}, + {name: "dotted network phrase is identifier syntax", stage: duckLakeStageAttach, message: "network.error", want: duckLakeReasonDriverError}, + {name: "dotted extension phrase is identifier syntax", stage: duckLakeStageLoad, message: "extension.unavailable", want: duckLakeReasonDriverError}, + {name: "dotted resource phrase is identifier syntax", stage: duckLakeStageAttach, message: "resource.unavailable", want: duckLakeReasonDriverError}, + {name: "dotted numeric reason is identifier syntax", stage: duckLakeStageAttach, message: "404.not.found", want: duckLakeReasonDriverError}, + {name: "leading dotted numeric status is identifier syntax", stage: duckLakeStageAttach, message: ".404", want: duckLakeReasonDriverError}, + {name: "status assignment terminal period", stage: duckLakeStageAttach, message: "status=404.", want: duckLakeReasonMissingObject}, + {name: "status assignment terminal ellipsis", stage: duckLakeStageAttach, message: "status=404...", want: duckLakeReasonMissingObject}, + {name: "status assignment after sentence period", stage: duckLakeStageAttach, message: "request failed. status=404", want: duckLakeReasonMissingObject}, + {name: "status assignment dotted suffix", stage: duckLakeStageAttach, message: "status=404.txt", want: duckLakeReasonDriverError}, + {name: "status assignment multi-dot suffix", stage: duckLakeStageAttach, message: "status=404..txt", want: duckLakeReasonDriverError}, + {name: "status assignment dot path suffix", stage: duckLakeStageAttach, message: "status=404../path", want: duckLakeReasonDriverError}, + {name: "identifier-like HTTP version", stage: duckLakeStageAttach, message: "foo_http/1.1 404", want: duckLakeReasonDriverError}, + {name: "assigned HTTP version", stage: duckLakeStageAttach, message: "foo=http/1.1 403", want: duckLakeReasonDriverError}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := newDuckLakeInitError(tt.stage, "", stderrors.New(tt.message)) + var initErr *duckLakeInitError + require.ErrorAs(t, err, &initErr) + require.Equal(t, tt.want, initErr.reason) + }) + } +} + +func TestDuckLakeGenericMarkerBoundaries(t *testing.T) { + tests := []struct { + name string + message string + want duckLakeFailureReason + }{ + {name: "plain TLS diagnostic", message: "TLS error.", want: duckLakeReasonTLSFailure}, + {name: "plain timeout diagnostic", message: "request timed out.", want: duckLakeReasonTimeout}, + {name: "plain syntax diagnostic", message: "syntax error.", want: duckLakeReasonInvalidStatement}, + {name: "plain cancellation diagnostic", message: "operation canceled.", want: duckLakeReasonInterrupted}, + {name: "plain interruption diagnostic", message: "operation interrupted.", want: duckLakeReasonInterrupted}, + {name: "hyphenated TLS identifier", message: "tls-error", want: duckLakeReasonDriverError}, + {name: "hyphenated timeout identifier", message: "timeout-token", want: duckLakeReasonDriverError}, + {name: "hyphenated syntax identifier", message: "syntax-error", want: duckLakeReasonDriverError}, + {name: "hyphenated parser identifier", message: "parser-error", want: duckLakeReasonDriverError}, + {name: "hyphenated interruption identifier", message: "interrupt-token", want: duckLakeReasonDriverError}, + {name: "hyphenated cancellation identifier", message: "cancel-token", want: duckLakeReasonDriverError}, + {name: "timeout query identifier", message: "request_timeout", want: duckLakeReasonDriverError}, + {name: "timeout path", message: "path=/timeout", want: duckLakeReasonDriverError}, + {name: "timeout URL", message: "https://timeout.example", want: duckLakeReasonDriverError}, + {name: "certificate URL", message: "https://certificate.example", want: duckLakeReasonDriverError}, + {name: "unexpected EOF identifier", message: "foo_unexpected eof_bar", want: duckLakeReasonDriverError}, + {name: "resource exhaustion identifier", message: "foo_no space left on device_bar", want: duckLakeReasonDriverError}, + {name: "quota identifier", message: "quota exceeded-token", want: duckLakeReasonDriverError}, + {name: "configuration identifier", message: "foo_missing secret", want: duckLakeReasonDriverError}, + {name: "catalog identifier", message: "foo_catalog error_bar", want: duckLakeReasonDriverError}, + {name: "IO path", message: "path=/failed to open", want: duckLakeReasonDriverError}, + {name: "canonical out of memory spelling", message: "out-of-memory", want: duckLakeReasonResourceExhaustion}, + {name: "canonical out of memory in sentence", message: "driver reports out-of-memory", want: duckLakeReasonResourceExhaustion}, + {name: "canonical out of memory after sentence period", message: "request failed. out-of-memory", want: duckLakeReasonResourceExhaustion}, + {name: "canonical read only spelling", message: "read-only file system", want: duckLakeReasonIOFailure}, + {name: "canonical input output spelling", message: "input/output error", want: duckLakeReasonIOFailure}, + {name: "canonical input output in sentence", message: "write failed input/output error", want: duckLakeReasonIOFailure}, + {name: "canonical input output after sentence period", message: "request failed. input/output error", want: duckLakeReasonIOFailure}, + {name: "canonical IO spelling", message: "i/o error", want: duckLakeReasonIOFailure}, + {name: "out of memory identifier", message: "foo_out-of-memory", want: duckLakeReasonDriverError}, + {name: "read only URL", message: "https://example.test/read-only file system", want: duckLakeReasonDriverError}, + {name: "input output path", message: "path=/input/output error", want: duckLakeReasonDriverError}, + {name: "quoted read only path", message: `path="read-only file system"`, want: duckLakeReasonDriverError}, + {name: "quoted input output path", message: `path='input/output error'`, want: duckLakeReasonDriverError}, + {name: "spaced assigned compact marker", message: "foo = nosuchkey", want: duckLakeReasonDriverError}, + {name: "quoted assigned compact marker", message: `foo = "nosuchkey"`, want: duckLakeReasonDriverError}, + {name: "spaced assigned permission phrase", message: "foo = access denied", want: duckLakeReasonDriverError}, + {name: "quoted assigned network phrase", message: `foo = "network error"`, want: duckLakeReasonDriverError}, + {name: "spaced assigned syntax phrase", message: "foo = syntax error", want: duckLakeReasonDriverError}, + {name: "multi-dot out of memory identifier", message: "out-of-memory..txt", want: duckLakeReasonDriverError}, + {name: "dot path out of memory identifier", message: "out-of-memory../path", want: duckLakeReasonDriverError}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := newDuckLakeInitError(duckLakeStageAttach, "", stderrors.New(tt.message)) + var initErr *duckLakeInitError + require.ErrorAs(t, err, &initErr) + require.Equal(t, tt.want, initErr.reason) + }) + } +} + +func TestDuckLakeRemoteFileContextMustBeProximate(t *testing.T) { + tests := []struct { + name string + message string + want duckLakeFailureReason + }{ + {name: "remote URL is the missing file target", message: "no such file at s3://private/catalog.ducklake", want: duckLakeReasonMissingObject}, + {name: "assigned remote URL is the missing file target", message: "no such file at url=s3://private/catalog.ducklake", want: duckLakeReasonMissingObject}, + {name: "assigned remote path is the missing file target", message: "no such file at path=https://storage.example/catalog.ducklake", want: duckLakeReasonMissingObject}, + {name: "quoted assigned remote URL is the missing file target", message: `"url"=s3://private/catalog.ducklake: no such file`, want: duckLakeReasonMissingObject}, + {name: "remote URL precedes the missing phrase", message: "s3://private/catalog.ducklake: no such file", want: duckLakeReasonMissingObject}, + {name: "remote word precedes the missing phrase", message: "remote object no such file", want: duckLakeReasonMissingObject}, + {name: "remote storage precedes the missing phrase", message: "remote storage: no such file", want: duckLakeReasonMissingObject}, + {name: "S3 error precedes the missing phrase", message: "S3 Error: no such file", want: duckLakeReasonMissingObject}, + {name: "S3 error remains remote before ordinary attach context", message: "S3 Error: no such file at attach time", want: duckLakeReasonMissingObject}, + {name: "S3 error remains remote before ordinary key context", message: "S3 Error: no such file at key foo", want: duckLakeReasonMissingObject}, + {name: "local file with unrelated later URL", message: "no such file at /var/lib/catalog; see https://docs.example/help", want: duckLakeReasonResourceUnavailable}, + {name: "local file with unrelated earlier URL", message: "see https://docs.example/help for details; no such file at /var/lib/catalog", want: duckLakeReasonResourceUnavailable}, + {name: "adjacent documentation URL does not override local target", message: "see https://docs.example/help: no such file at /var/lib/catalog", want: duckLakeReasonResourceUnavailable}, + {name: "local file with unrelated later remote word", message: "no such file at /var/lib/catalog; see remote docs", want: duckLakeReasonResourceUnavailable}, + {name: "nested remote assignment remains unrelated", message: "no such file at foo_url=s3://private/catalog.ducklake", want: duckLakeReasonResourceUnavailable}, + {name: "repeated remote assignment remains unrelated", message: "no such file at url==s3://private/catalog.ducklake", want: duckLakeReasonResourceUnavailable}, + {name: "spaced repeated remote assignment remains unrelated", message: "no such file at url = = s3://private/catalog.ducklake", want: duckLakeReasonResourceUnavailable}, + {name: "reverse nested remote identifier remains unrelated", message: "foo_url=s3://private/catalog.ducklake: no such file", want: duckLakeReasonResourceUnavailable}, + {name: "reverse nested remote assignment remains unrelated", message: "foo=url=s3://private/catalog.ducklake: no such file", want: duckLakeReasonResourceUnavailable}, + {name: "reverse quoted dotted remote assignment remains unrelated", message: `foo."url"=s3://private/catalog.ducklake: no such file`, want: duckLakeReasonResourceUnavailable}, + {name: "reverse quoted multiword remote assignment remains unrelated", message: `foo."remote url"=s3://private/catalog.ducklake: no such file`, want: duckLakeReasonResourceUnavailable}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := newDuckLakeInitError(duckLakeStageAttach, "", stderrors.New(tt.message)) + var initErr *duckLakeInitError + require.ErrorAs(t, err, &initErr) + require.Equal(t, tt.want, initErr.reason) + }) + } +} + +type duckLakeAsErrorWrapper struct { + duckErr *duckdb.Error +} + +func (wrapper duckLakeAsErrorWrapper) Error() string { + return "opaque wrapper" +} + +func (wrapper duckLakeAsErrorWrapper) As(target any) bool { + duckTarget, ok := target.(**duckdb.Error) + if !ok || duckTarget == nil || wrapper.duckErr == nil { + return false + } + *duckTarget = wrapper.duckErr + return true +} + +func TestDuckLakeFailureClassificationPreservesAsWrapperCompatibility(t *testing.T) { + raw := &duckdb.Error{ + Type: duckdb.ErrorTypePermission, + Msg: "Permission Error: opaque secret", + } + err := newDuckLakeInitError(duckLakeStageAttach, "", duckLakeAsErrorWrapper{duckErr: raw}) + + var initErr *duckLakeInitError + require.ErrorAs(t, err, &initErr) + require.Equal(t, duckLakeReasonPermissionDenied, initErr.reason) + require.True(t, stderrors.Is(err, initErr.rawCause())) + var exposed *duckdb.Error + require.False(t, stderrors.As(err, &exposed)) +} + +func TestDuckLakeInitErrorRedactsJoinedCauseAcrossFormats(t *testing.T) { + const ( + metadata = "s3://private-bucket/catalog.ducklake" + endpoint = "https://private-storage.example:9443" + secret = "secret-key-77" + ) + rawMissing := &duckdb.Error{ + Type: duckdb.ErrorTypeHTTP, + Msg: "HTTP Error: 404 Not Found metadata=" + metadata + " SECRET=" + secret, + } + rawNetwork := fmt.Errorf("network failure endpoint=%s", endpoint) + joined := stderrors.Join(rawNetwork, rawMissing) + err := newDuckLakeInitError(duckLakeStageAttach, "", joined) + + var initErr *duckLakeInitError + require.ErrorAs(t, err, &initErr) + require.Equal(t, duckLakeReasonMissingObject, initErr.reason) + require.Nil(t, stderrors.Unwrap(err)) + var exposed *duckdb.Error + require.False(t, stderrors.As(err, &exposed)) + require.True(t, stderrors.Is(err, rawMissing)) + require.True(t, stderrors.Is(err, rawNetwork)) + for _, format := range []string{"%v", "%+v", "%#v", "%q"} { + formatted := fmt.Sprintf(format, err) + require.NotContains(t, formatted, metadata) + require.NotContains(t, formatted, endpoint) + require.NotContains(t, formatted, secret) + require.NotContains(t, formatted, rawMissing.Msg) + } +} + +func TestDuckLakeInitErrorJoinedTypedCausesUseStableSpecificity(t *testing.T) { + rawMissing := &duckdb.Error{ + Type: duckdb.ErrorTypeHTTP, + Msg: "HTTP Error: 404 Not Found metadata=s3://private/catalog.ducklake", + } + rawNetwork := &duckdb.Error{ + Type: duckdb.ErrorTypeNetwork, + Msg: "Network Error: dial tcp https://private-storage.example", + } + for _, tc := range []struct { + name string + cause error + }{ + {name: "network then missing", cause: stderrors.Join(rawNetwork, rawMissing)}, + {name: "missing then network", cause: stderrors.Join(rawMissing, rawNetwork)}, + } { + t.Run(tc.name, func(t *testing.T) { + err := newDuckLakeInitError(duckLakeStageAttach, "", tc.cause) + var initErr *duckLakeInitError + require.ErrorAs(t, err, &initErr) + require.Equal(t, duckLakeReasonMissingObject, initErr.reason) + require.True(t, stderrors.Is(err, rawMissing)) + require.True(t, stderrors.Is(err, rawNetwork)) + require.Nil(t, stderrors.Unwrap(err)) + require.NotContains(t, err.Error(), rawMissing.Msg) + require.NotContains(t, err.Error(), rawNetwork.Msg) + }) + } +} + +func TestDuckLakeInitErrorSanitizesStageAndExtensionContext(t *testing.T) { + tests := []struct { + name string + stage duckLakeInitStage + extension string + want string + }{ + { + name: "unknown load extension is omitted", + stage: duckLakeStageLoad, + extension: "httpfs'; SECRET='leaked", + want: "load duckdb extension failed (stage=load reason=driver_error)", + }, + { + name: "non-load extension is omitted", + stage: duckLakeStageCreateSecret, + extension: "ducklake", + want: "create ducklake service secret failed (stage=create_secret reason=driver_error)", + }, + { + name: "unknown stage is fixed", + stage: duckLakeInitStage("attach s3://private/path SECRET=leaked"), + extension: "ducklake", + want: "ducklake initialization failed (stage=unknown reason=driver_error)", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := newDuckLakeInitError(tt.stage, tt.extension, stderrors.New("opaque driver detail SECRET=private")) + require.Equal(t, tt.want, err.Error()) + }) + } +} + +func TestDuckLakeFailureReasonPrecedenceAndGenericCoverage(t *testing.T) { + tests := []struct { + name string + stage duckLakeInitStage + message string + want duckLakeFailureReason + }{ + {name: "permission precedes http", stage: duckLakeStageAttach, message: "HTTP 403 forbidden", want: duckLakeReasonPermissionDenied}, + {name: "tls precedes timeout", stage: duckLakeStageAttach, message: "TLS handshake timed out", want: duckLakeReasonTLSFailure}, + {name: "eof precedes network", stage: duckLakeStageAttach, message: "network stream ended with unexpected EOF", want: duckLakeReasonUnexpectedEOF}, + {name: "timeout precedes network", stage: duckLakeStageAttach, message: "network request timed out", want: duckLakeReasonTimeout}, + {name: "network precedes http", stage: duckLakeStageAttach, message: "HTTP connection refused", want: duckLakeReasonNetworkFailure}, + {name: "http request", stage: duckLakeStageAttach, message: "HTTP request returned status 503", want: duckLakeReasonHTTPFailure}, + {name: "extension precedes unavailable", stage: duckLakeStageLoad, message: "ducklake extension httpfs is unavailable", want: duckLakeReasonExtensionUnavailable}, + {name: "resource exhaustion precedes io", stage: duckLakeStageAttach, message: "I/O error: no space left on device", want: duckLakeReasonResourceExhaustion}, + {name: "configuration precedes missing", stage: duckLakeStageAttach, message: "configuration error: missing secret", want: duckLakeReasonInvalidConfiguration}, + {name: "resource unavailable precedes io", stage: duckLakeStageAttach, message: "I/O error: no such file or directory", want: duckLakeReasonResourceUnavailable}, + {name: "read-only filesystem io", stage: duckLakeStageAttach, message: "read-only file system", want: duckLakeReasonIOFailure}, + {name: "resource temporarily unavailable", stage: duckLakeStageAttach, message: "resource temporarily unavailable", want: duckLakeReasonResourceUnavailable}, + {name: "interrupted", stage: duckLakeStageAttach, message: "operation canceled by caller", want: duckLakeReasonInterrupted}, + {name: "opaque fallback", stage: duckLakeStageAttach, message: "driver exploded at /private/path SECRET=private", want: duckLakeReasonDriverError}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := newDuckLakeInitError(tt.stage, "", stderrors.New(tt.message)) + require.Contains(t, err.Error(), "reason="+string(tt.want)) + }) + } +} + +func TestDuckLakeFailureReasonDuckDBTypeFallbacks(t *testing.T) { + tests := []struct { + name string + typ duckdb.ErrorType + msg string + want duckLakeFailureReason + }{ + {name: "settings is configuration", typ: duckdb.ErrorTypeSettings, msg: "Settings Error: invalid setting", want: duckLakeReasonInvalidConfiguration}, + {name: "dependency remains opaque", typ: duckdb.ErrorTypeDependency, msg: "Dependency Error: internal dependency", want: duckLakeReasonDriverError}, + {name: "fatal remains opaque", typ: duckdb.ErrorTypeFatal, msg: "FATAL Error: internal failure", want: duckLakeReasonDriverError}, + {name: "internal remains opaque", typ: duckdb.ErrorTypeInternal, msg: "INTERNAL Error: internal failure", want: duckLakeReasonDriverError}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := newDuckLakeInitError(duckLakeStageAttach, "", &duckdb.Error{Type: tt.typ, Msg: tt.msg}) + require.Contains(t, err.Error(), "reason="+string(tt.want)) + }) + } +} + +func TestDuckLakeInitErrorHandlesTypedNilDuckDBError(t *testing.T) { + var raw *duckdb.Error + var cause error = raw + var err error + require.NotPanics(t, func() { + err = newDuckLakeInitError(duckLakeStageAttach, "", cause) + }) + require.Error(t, err) + require.Contains(t, err.Error(), "reason=driver_error") + var exposed *duckdb.Error + require.False(t, stderrors.As(err, &exposed)) + require.NotPanics(t, func() { + require.False(t, stderrors.Is(err, context.Canceled)) + }) + var nilTarget *duckdb.Error + var target error = nilTarget + require.NotPanics(t, func() { + require.False(t, stderrors.Is(err, target)) + }) +} + +type scriptedDuckLakeExecer struct { + queries []string + errors []error +} + +func (e *scriptedDuckLakeExecer) ExecContext(_ context.Context, query string, _ []driver.NamedValue) (driver.Result, error) { + index := len(e.queries) + e.queries = append(e.queries, query) + if index < len(e.errors) && e.errors[index] != nil { + return nil, e.errors[index] + } + return driver.RowsAffected(0), nil +} + +func newScriptedDuckLakeRuntime() *duckLakeRuntime { + return &duckLakeRuntime{ + config: configuration.DuckLakeConfig{ + ExtensionDir: "/opt/myduck/extensions", + MetadataPath: "/var/lib/myduck/catalog.ducklake", + DataPath: "s3://test-bucket/data", + S3: configuration.DuckLakeS3Config{ + Endpoint: "https://storage.example", + Region: "us-east-1", + UseSSL: true, + }, + }, + manifest: []ExtensionArtifact{ + {Name: "httpfs", FileName: "httpfs.duckdb_extension"}, + {Name: "ducklake", FileName: "ducklake.duckdb_extension"}, + }, + } +} + +func TestDuckLakeInitSQLStopsAtFirstFailedStage(t *testing.T) { + tests := []struct { + name string + failAt int + cause error + stage duckLakeInitStage + reason duckLakeFailureReason + queryWant int + }{ + { + name: "first load", + failAt: 0, + cause: &duckdb.Error{Type: duckdb.ErrorTypeIO, Msg: "I/O Error: /private/httpfs.duckdb_extension"}, + stage: duckLakeStageLoad, + reason: duckLakeReasonIOFailure, + queryWant: 1, + }, + { + name: "second load", + failAt: 1, + cause: &duckdb.Error{Type: duckdb.ErrorTypeIO, Msg: "I/O Error: /private/ducklake.duckdb_extension"}, + stage: duckLakeStageLoad, + reason: duckLakeReasonIOFailure, + queryWant: 2, + }, + { + name: "create secret", + failAt: 2, + cause: &duckdb.Error{Type: duckdb.ErrorTypePermission, Msg: "Permission Error: SECRET=private"}, + stage: duckLakeStageCreateSecret, + reason: duckLakeReasonPermissionDenied, + queryWant: 3, + }, + { + name: "attach", + failAt: 3, + cause: &duckdb.Error{Type: duckdb.ErrorTypeHTTP, Msg: "HTTP Error: https://private.example/catalog"}, + stage: duckLakeStageAttach, + reason: duckLakeReasonHTTPFailure, + queryWant: 4, + }, + { + name: "all stages succeed", + failAt: -1, + queryWant: 4, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + runtime := newScriptedDuckLakeRuntime() + execer := &scriptedDuckLakeExecer{} + if tt.failAt >= 0 { + execer.errors = make([]error, 5) + execer.errors[tt.failAt] = tt.cause + } + + runtime.initializeMu.Lock() + err := runtime.initializeSQLLocked(context.Background(), nil, execer) + runtime.initializeMu.Unlock() + + if tt.failAt < 0 { + require.NoError(t, err) + } else { + require.Error(t, err) + require.Contains(t, err.Error(), "stage="+string(tt.stage)) + require.Contains(t, err.Error(), "reason="+string(tt.reason)) + } + require.Len(t, execer.queries, tt.queryWant) + for index, query := range execer.queries { + switch { + case index < 2: + require.True(t, strings.HasPrefix(query, "LOAD '")) + case index == 2: + require.True(t, strings.HasPrefix(query, "CREATE OR REPLACE SECRET ")) + case index == 3: + require.True(t, strings.HasPrefix(query, "ATTACH IF NOT EXISTS ")) + } + } + }) + } +} diff --git a/catalog/execution_lease_callsites_test.go b/catalog/execution_lease_callsites_test.go new file mode 100644 index 00000000..d9e32f86 --- /dev/null +++ b/catalog/execution_lease_callsites_test.go @@ -0,0 +1,387 @@ +package catalog + +import ( + "context" + stdsql "database/sql" + "errors" + "fmt" + "regexp" + "sync" + "sync/atomic" + "testing" + + "github.com/DATA-DOG/go-sqlmock" + "github.com/apecloud/myduckserver/adapter" + "github.com/apecloud/myduckserver/configuration" + "github.com/apecloud/myduckserver/mycontext" + "github.com/dolthub/go-mysql-server/memory" + "github.com/dolthub/go-mysql-server/sql" + "github.com/dolthub/go-mysql-server/sql/types" + "github.com/stretchr/testify/require" +) + +// callsiteLeaseSession models the atomic session surfaces used by catalog +// callsites without involving ConnectionPool internals. The release callback +// lets each test inspect driver expectations at the exact lease boundary. +type callsiteLeaseSession struct { + *memory.Session + + mu sync.Mutex + conn *stdsql.Conn + tx *stdsql.Tx + release func() +} + +var _ adapter.ConnectionHolder = (*callsiteLeaseSession)(nil) +var _ adapter.TransactionBindingHolder = (*callsiteLeaseSession)(nil) +var _ adapter.ReleaseTransactionBindingHolder = (*callsiteLeaseSession)(nil) +var _ adapter.ExecutionSnapshotHolder = (*callsiteLeaseSession)(nil) +var _ adapter.ExecutionSnapshotLeaseHolder = (*callsiteLeaseSession)(nil) +var _ adapter.IdentityTransactionCloser = (*callsiteLeaseSession)(nil) + +func (s *callsiteLeaseSession) GetConn(context.Context) (*stdsql.Conn, error) { + s.mu.Lock() + defer s.mu.Unlock() + return s.conn, nil +} + +func (s *callsiteLeaseSession) GetCatalogConn(ctx context.Context) (*stdsql.Conn, error) { + return s.GetConn(ctx) +} + +func (s *callsiteLeaseSession) GetTxn(ctx context.Context, options *stdsql.TxOptions) (*stdsql.Tx, error) { + s.mu.Lock() + defer s.mu.Unlock() + if s.tx != nil { + return s.tx, nil + } + tx, err := s.conn.BeginTx(ctx, options) + if err != nil { + return nil, err + } + s.tx = tx + return tx, nil +} + +func (s *callsiteLeaseSession) GetCatalogTxn(ctx context.Context, options *stdsql.TxOptions) (*stdsql.Tx, error) { + return s.GetTxn(ctx, options) +} + +func (s *callsiteLeaseSession) TryGetTxn() *stdsql.Tx { + s.mu.Lock() + defer s.mu.Unlock() + return s.tx +} + +func (s *callsiteLeaseSession) GetTxnBinding() (*stdsql.Conn, *stdsql.Tx) { + s.mu.Lock() + defer s.mu.Unlock() + return s.conn, s.tx +} + +func (s *callsiteLeaseSession) GetTxnBindingForRelease() (*stdsql.Conn, *stdsql.Tx) { + return s.GetTxnBinding() +} + +func (s *callsiteLeaseSession) GetExecutionSnapshot(context.Context, bool) (*stdsql.Conn, *stdsql.Tx, error) { + s.mu.Lock() + defer s.mu.Unlock() + return s.conn, s.tx, nil +} + +func (s *callsiteLeaseSession) GetExecutionSnapshotLease(context.Context, bool) (*stdsql.Conn, *stdsql.Tx, func(), error) { + s.mu.Lock() + defer s.mu.Unlock() + return s.conn, s.tx, s.release, nil +} + +func (s *callsiteLeaseSession) GetCurrentCatalog() string { return "memory" } + +func (s *callsiteLeaseSession) GetCurrentSchema() string { return "main" } + +func (s *callsiteLeaseSession) CloseTxn() { + s.mu.Lock() + s.tx = nil + s.mu.Unlock() +} + +func (s *callsiteLeaseSession) CloseTxnIf(tx *stdsql.Tx) { + s.mu.Lock() + if s.tx == tx { + s.tx = nil + } + s.mu.Unlock() +} + +func (s *callsiteLeaseSession) CloseConn() {} + +type callsiteLeaseObserver struct { + count atomic.Int32 + + mu sync.Mutex + err error + check func() error +} + +func (o *callsiteLeaseObserver) release() { + o.count.Add(1) + var err error + if o.check != nil { + err = o.check() + } + o.mu.Lock() + o.err = errors.Join(o.err, err) + o.mu.Unlock() +} + +func (o *callsiteLeaseObserver) requireReleasedOnce(t *testing.T) { + t.Helper() + require.Equal(t, int32(1), o.count.Load()) + o.mu.Lock() + defer o.mu.Unlock() + require.NoError(t, o.err, "lease released before all owned work completed") +} + +func newCallsiteLeaseContext(t *testing.T, id uint32) (*sql.Context, *callsiteLeaseSession, sqlmock.Sqlmock, *callsiteLeaseObserver) { + t.Helper() + db, mock, err := sqlmock.New() + require.NoError(t, err) + conn, err := db.Conn(context.Background()) + require.NoError(t, err) + t.Cleanup(func() { + require.NoError(t, conn.Close()) + _ = db.Close() + }) + + base := memory.NewSession(sql.NewBaseSession(), nil) + base.SetConnectionId(id) + base.SetCurrentDatabase("memory") + observer := &callsiteLeaseObserver{} + session := &callsiteLeaseSession{ + Session: base, + conn: conn, + release: observer.release, + } + sqlCtx := sql.NewContext( + mycontext.WithFrontendQuery(context.Background()), + sql.WithSession(session), + ) + return sqlCtx, session, mock, observer +} + +func newCallsiteRowInserter() *rowInserter { + return &rowInserter{ + db: "main", + table: "lease_rows", + schema: sql.Schema{ + &sql.Column{Name: "id", Type: types.Int32, Nullable: false}, + &sql.Column{Name: "payload", Type: types.Text, Nullable: true}, + }, + } +} + +func expectRowInserterCreate(mock sqlmock.Sqlmock, id uint32) { + tmpTable := fmt.Sprintf("main_lease_rows_%d", id) + mock.ExpectExec(regexp.QuoteMeta( + `CREATE TEMP TABLE IF NOT EXISTS ` + QuoteIdentifierANSI(tmpTable) + + ` AS FROM "main"."lease_rows" LIMIT 0`, + )).WillReturnResult(sqlmock.NewResult(0, 0)) +} + +func expectRowInserterDrop(mock sqlmock.Sqlmock, id uint32) { + tmpTable := fmt.Sprintf("main_lease_rows_%d", id) + mock.ExpectExec(regexp.QuoteMeta( + `DROP TABLE IF EXISTS temp.main.` + QuoteIdentifierANSI(tmpTable), + )).WillReturnResult(sqlmock.NewResult(0, 0)) +} + +func TestRowInserterExecutionLeaseTerminalPaths(t *testing.T) { + const sessionID = uint32(701) + insertSQL := `INSERT INTO "main_lease_rows_701" VALUES (?, ?)` + flushSQL := `INSERT INTO "main"."lease_rows" FROM "main_lease_rows_701"` + + t.Run("initialization error", func(t *testing.T) { + sqlCtx, _, mock, observer := newCallsiteLeaseContext(t, sessionID) + ri := newCallsiteRowInserter() + initErr := errors.New("prepare failed") + + expectRowInserterCreate(mock, sessionID) + mock.ExpectPrepare(regexp.QuoteMeta(insertSQL)).WillReturnError(initErr) + expectRowInserterDrop(mock, sessionID) + observer.check = func() error { + if ri.stmt != nil { + return errors.New("prepared statement remained live at lease release") + } + return mock.ExpectationsWereMet() + } + + ri.StatementBegin(sqlCtx) + require.ErrorIs(t, ri.StatementComplete(sqlCtx), initErr) + require.ErrorIs(t, ri.DiscardChanges(sqlCtx, initErr), initErr) + ri.releaseSnapshot() + ri.releaseSnapshot() + + observer.requireReleasedOnce(t) + require.NoError(t, mock.ExpectationsWereMet()) + }) + + t.Run("insert and close", func(t *testing.T) { + sqlCtx, _, mock, observer := newCallsiteLeaseContext(t, sessionID) + ri := newCallsiteRowInserter() + + expectRowInserterCreate(mock, sessionID) + prepared := mock.ExpectPrepare(regexp.QuoteMeta(insertSQL)).WillBeClosed() + prepared.ExpectExec().WithArgs(int64(7), "kept").WillReturnResult(sqlmock.NewResult(0, 1)) + mock.ExpectExec(regexp.QuoteMeta(flushSQL)).WillReturnResult(sqlmock.NewResult(0, 1)) + expectRowInserterDrop(mock, sessionID) + observer.check = func() error { + if ri.stmt != nil { + return errors.New("prepared statement remained live at lease release") + } + return mock.ExpectationsWereMet() + } + + require.NoError(t, ri.Insert(sqlCtx, sql.Row{int32(7), "kept"})) + require.NoError(t, ri.Close(sqlCtx)) + ri.releaseSnapshot() + + observer.requireReleasedOnce(t) + require.NoError(t, mock.ExpectationsWereMet()) + }) + + t.Run("discard changes", func(t *testing.T) { + sqlCtx, _, mock, observer := newCallsiteLeaseContext(t, sessionID) + ri := newCallsiteRowInserter() + + expectRowInserterCreate(mock, sessionID) + prepared := mock.ExpectPrepare(regexp.QuoteMeta(insertSQL)).WillBeClosed() + prepared.ExpectExec().WithArgs(int64(8), "discarded").WillReturnResult(sqlmock.NewResult(0, 1)) + expectRowInserterDrop(mock, sessionID) + observer.check = func() error { + if ri.stmt != nil { + return errors.New("prepared statement remained live at lease release") + } + return mock.ExpectationsWereMet() + } + + require.NoError(t, ri.Insert(sqlCtx, sql.Row{int32(8), "discarded"})) + require.NoError(t, ri.DiscardChanges(sqlCtx, errors.New("statement failed"))) + ri.releaseSnapshot() + + observer.requireReleasedOnce(t) + require.NoError(t, mock.ExpectationsWereMet()) + }) + + t.Run("cleanup after transaction rollback", func(t *testing.T) { + sqlCtx, session, mock, observer := newCallsiteLeaseContext(t, sessionID) + ri := newCallsiteRowInserter() + + mock.ExpectBegin() + tx, err := session.GetTxn(sqlCtx, nil) + require.NoError(t, err) + expectRowInserterCreate(mock, sessionID) + prepared := mock.ExpectPrepare(regexp.QuoteMeta(insertSQL)).WillBeClosed() + prepared.ExpectExec().WithArgs(int64(9), "rolled back").WillReturnResult(sqlmock.NewResult(0, 1)) + mock.ExpectRollback() + expectRowInserterDrop(mock, sessionID) + observer.check = func() error { + if ri.stmt != nil { + return errors.New("prepared statement remained live at lease release") + } + if session.TryGetTxn() != nil { + return errors.New("rolled-back transaction remained bound at lease release") + } + return mock.ExpectationsWereMet() + } + + require.NoError(t, ri.Insert(sqlCtx, sql.Row{int32(9), "rolled back"})) + owner, inactive, rollbackErr := adapter.FinalizeRollback(sqlCtx, tx) + require.Same(t, session.conn, owner) + require.True(t, inactive) + require.NoError(t, rollbackErr) + require.NoError(t, ri.DiscardChanges(sqlCtx, errors.New("transaction rolled back"))) + ri.releaseSnapshot() + + observer.requireReleasedOnce(t) + require.NoError(t, mock.ExpectationsWereMet()) + }) +} + +func TestRecordTableStorageSelectionScopedTransactionLease(t *testing.T) { + const ( + sessionID = uint32(702) + tableName = "lease_storage" + comment = "metadata consumed before release" + ) + + for _, test := range []struct { + name string + metadataErr error + }{ + {name: "commit"}, + {name: "rollback after metadata write failure", metadataErr: errors.New("metadata write failed")}, + } { + t.Run(test.name, func(t *testing.T) { + sqlCtx, session, mock, observer := newCallsiteLeaseContext(t, sessionID) + provider := &DatabaseProvider{ + defaultCatalogName: "memory", + duckLake: &duckLakeRuntime{config: configuration.DuckLakeConfig{ + MetadataPath: "/tmp/task77/catalog.ducklake", + DataPath: "/tmp/task77/data", + }}, + } + database := NewDatabaseWithProvider("main", "memory", provider) + storage := TableStorageSelection{Kind: TableStorageObject, Explicit: true, Source: "test"} + physical := database.objectPhysicalTableName(tableName) + encoded := NewCommentWithMeta(comment, ExtraTableInfo{Storage: TableStorageObject}).Encode() + + mock.ExpectBegin() + mock.ExpectQuery(`FROM duckdb_columns\(\)`). + WithArgs("memory", "main", tableName). + WillReturnRows(sqlmock.NewRows([]string{ + "column_name", "data_type", "is_nullable", "column_default", + }).AddRow("id", "INTEGER", true, nil)). + RowsWillBeClosed() + mock.ExpectQuery(`FROM duckdb_constraints\(\)`). + WithArgs("memory", "main", tableName). + WillReturnRows(sqlmock.NewRows([]string{"constraint_type"})). + RowsWillBeClosed() + mock.ExpectExec(regexp.QuoteMeta( + `CREATE SCHEMA IF NOT EXISTS "__myduck_ducklake"."main"`, + )).WillReturnResult(sqlmock.NewResult(0, 0)) + mock.ExpectExec(regexp.QuoteMeta( + `CREATE TABLE ` + physical + ` ("id" INTEGER NULL)`, + )).WillReturnResult(sqlmock.NewResult(0, 0)) + mock.ExpectQuery(`SELECT comment FROM duckdb_tables\(\)`). + WithArgs("memory", "main", tableName). + WillReturnRows(sqlmock.NewRows([]string{"comment"}).AddRow(comment)). + RowsWillBeClosed() + metadata := mock.ExpectExec(regexp.QuoteMeta( + `COMMENT ON TABLE ` + FullTableName("memory", "main", tableName) + ` IS '` + encoded + `'`, + )) + if test.metadataErr != nil { + metadata.WillReturnError(test.metadataErr) + mock.ExpectRollback() + } else { + metadata.WillReturnResult(sqlmock.NewResult(0, 0)) + mock.ExpectCommit() + } + observer.check = func() error { + if session.TryGetTxn() != nil { + return errors.New("scoped transaction remained bound at lease release") + } + return mock.ExpectationsWereMet() + } + + err := database.RecordTableStorageSelection(sqlCtx, tableName, storage) + if test.metadataErr != nil { + require.ErrorContains(t, err, test.metadataErr.Error()) + } else { + require.NoError(t, err) + } + + observer.requireReleasedOnce(t) + require.NoError(t, mock.ExpectationsWereMet()) + }) + } +} diff --git a/catalog/external_session_test.go b/catalog/external_session_test.go new file mode 100644 index 00000000..94b68a09 --- /dev/null +++ b/catalog/external_session_test.go @@ -0,0 +1,854 @@ +package catalog + +import ( + "context" + stdsql "database/sql" + "database/sql/driver" + "errors" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/DATA-DOG/go-sqlmock" + "github.com/stretchr/testify/require" +) + +func newExternalSessionTestPool( + t *testing.T, +) (*ConnectionPool, *stdsql.DB, sqlmock.Sqlmock) { + t.Helper() + db, mock, err := sqlmock.New() + require.NoError(t, err) + t.Cleanup(func() { _ = db.Close() }) + p := NewConnectionPool(nil, db, "memory") + // UDF registration is DuckDB-specific and orthogonal to these lifecycle + // tests. Mark it complete so sqlmock sees only the driver calls under test. + p.registerMySQLUDFsOnce.Do(func() {}) + return p, db, mock +} + +func TestExternalSessionExecutionLeaseUsesExactConnection(t *testing.T) { + p, _, mock := newExternalSessionTestPool(t) + var initializerCalls atomic.Int32 + p.SetConnectionInitializer(func(_ context.Context, conn *stdsql.Conn) error { + require.NotNil(t, conn) + initializerCalls.Add(1) + return nil + }) + + session, err := p.OpenExternalSession(context.Background()) + require.NoError(t, err) + require.Equal(t, int32(1), initializerCalls.Load()) + + conn, tx, release, err := session.ExecutionLease(context.Background()) + require.NoError(t, err) + require.Same(t, session.conn, conn) + require.Nil(t, tx) + require.Equal(t, int32(2), initializerCalls.Load()) + release() + release() + + var cleanupCalls atomic.Int32 + require.NoError(t, session.SetCleanupHook(func() { cleanupCalls.Add(1) })) + mock.ExpectClose() + require.NoError(t, session.Close()) + require.NoError(t, session.Close()) + require.Equal(t, int32(1), cleanupCalls.Load()) + + _, _, _, err = session.ExecutionLease(context.Background()) + require.ErrorIs(t, err, ErrExternalSessionClosed) + require.Equal(t, int32(2), initializerCalls.Load(), "closed handles must fail before initialization") + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestExternalSessionAdmissionFailsBeforeDriverWork(t *testing.T) { + t.Run("Open rechecks after lifecycle admission", func(t *testing.T) { + p, db, _ := newExternalSessionTestPool(t) + admissionErr := errors.New("external admission rejected") + var guardCalls atomic.Int32 + var initializerCalls atomic.Int32 + p.SetConnectionInitializer(func(context.Context, *stdsql.Conn) error { + initializerCalls.Add(1) + return nil + }) + p.SetConnectionAdmissionGuard(func(context.Context) error { + if guardCalls.Add(1) == 2 { + return admissionErr + } + return nil + }) + + session, err := p.OpenExternalSession(context.Background()) + require.Nil(t, session) + require.ErrorIs(t, err, admissionErr) + require.Equal(t, int32(2), guardCalls.Load()) + require.Zero(t, initializerCalls.Load()) + require.Zero(t, db.Stats().InUse) + }) + + t.Run("Lease rejects before initializer", func(t *testing.T) { + p, _, mock := newExternalSessionTestPool(t) + var initializerCalls atomic.Int32 + p.SetConnectionInitializer(func(context.Context, *stdsql.Conn) error { + initializerCalls.Add(1) + return nil + }) + session, err := p.OpenExternalSession(context.Background()) + require.NoError(t, err) + require.Equal(t, int32(1), initializerCalls.Load()) + + admissionErr := errors.New("generation is unavailable") + p.SetConnectionAdmissionGuard(func(context.Context) error { return admissionErr }) + _, _, _, err = session.ExecutionLease(context.Background()) + require.ErrorIs(t, err, admissionErr) + require.Equal(t, int32(1), initializerCalls.Load()) + + p.SetConnectionAdmissionGuard(nil) + mock.ExpectClose() + require.NoError(t, session.Close()) + require.NoError(t, mock.ExpectationsWereMet()) + }) +} + +func TestExternalSessionTransactionIdentityAndLifecycleHooks(t *testing.T) { + p, _, mock := newExternalSessionTestPool(t) + var admissionCalls atomic.Int32 + var admittedTransactions atomic.Int32 + var releasedAdmissions atomic.Int32 + var finishedCalls atomic.Int32 + var finishedMu sync.Mutex + var finished []*stdsql.Tx + p.SetTransactionLifecycleHooks(func() func(*stdsql.Tx) { + admissionCalls.Add(1) + return func(tx *stdsql.Tx) { + if tx == nil { + releasedAdmissions.Add(1) + return + } + admittedTransactions.Add(1) + } + }, func(tx *stdsql.Tx) { + finishedCalls.Add(1) + finishedMu.Lock() + finished = append(finished, tx) + finishedMu.Unlock() + }) + + session, err := p.OpenExternalSession(context.Background()) + require.NoError(t, err) + mock.ExpectBegin() + tx, err := session.BeginTx(context.Background(), nil) + require.NoError(t, err) + require.NotNil(t, tx) + require.Equal(t, int32(1), admissionCalls.Load()) + require.Equal(t, int32(1), admittedTransactions.Load()) + + wrong := new(stdsql.Tx) + require.ErrorIs(t, session.Commit(wrong), ErrExternalSessionTransactionMismatch) + require.Zero(t, finishedCalls.Load()) + conn, leasedTx, release, err := session.ExecutionLease(context.Background()) + require.NoError(t, err) + require.Same(t, session.conn, conn) + require.Same(t, tx, leasedTx) + release() + + var commitCompletions atomic.Int32 + require.True(t, p.RegisterTransactionCompletion(tx, func(success bool) { + require.True(t, success) + commitCompletions.Add(1) + })) + mock.ExpectCommit() + require.NoError(t, session.Commit(tx)) + require.Equal(t, int32(1), commitCompletions.Load()) + require.Equal(t, int32(1), finishedCalls.Load()) + + mock.ExpectBegin() + tx, err = session.BeginTx(context.Background(), nil) + require.NoError(t, err) + require.Equal(t, int32(2), admissionCalls.Load()) + require.Equal(t, int32(2), admittedTransactions.Load()) + require.Zero(t, releasedAdmissions.Load()) + + var rollbackCompletions atomic.Int32 + require.True(t, p.RegisterTransactionCompletion(tx, func(success bool) { + require.False(t, success) + rollbackCompletions.Add(1) + })) + mock.ExpectRollback() + var cleanupConn *stdsql.Conn + require.NoError(t, session.Rollback(tx, func(conn *stdsql.Conn) error { + cleanupConn = conn + return nil + })) + require.Same(t, session.conn, cleanupConn) + require.Equal(t, int32(1), rollbackCompletions.Load()) + require.Equal(t, int32(2), finishedCalls.Load()) + + finishedMu.Lock() + require.Equal(t, 2, len(finished)) + finishedMu.Unlock() + mock.ExpectClose() + require.NoError(t, session.Close()) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestExternalSessionResetWaitsForLeaseAndInvalidatesHandle(t *testing.T) { + p, _, oldMock := newExternalSessionTestPool(t) + var initializerCalls atomic.Int32 + p.SetConnectionInitializer(func(context.Context, *stdsql.Conn) error { + initializerCalls.Add(1) + return nil + }) + session, err := p.OpenExternalSession(context.Background()) + require.NoError(t, err) + oldGeneration := session.generation + var cleanupCalls atomic.Int32 + require.NoError(t, session.SetCleanupHook(func() { cleanupCalls.Add(1) })) + + _, _, release, err := session.ExecutionLease(context.Background()) + require.NoError(t, err) + require.Equal(t, int32(2), initializerCalls.Load()) + + newDB, newMock, err := sqlmock.New() + require.NoError(t, err) + t.Cleanup(func() { _ = newDB.Close() }) + oldMock.ExpectClose() + resetDone := make(chan error, 1) + go func() { resetDone <- p.Reset(nil, newDB) }() + waitForPoolShutdown(t, p) + select { + case err := <-resetDone: + t.Fatalf("Reset completed before the execution lease released: %v", err) + case <-time.After(100 * time.Millisecond): + } + + release() + select { + case err := <-resetDone: + require.NoError(t, err) + case <-time.After(3 * time.Second): + t.Fatal("Reset did not finish after the execution lease released") + } + require.Equal(t, int32(1), cleanupCalls.Load()) + require.NoError(t, oldMock.ExpectationsWereMet()) + + _, _, _, err = session.ExecutionLease(context.Background()) + require.ErrorIs(t, err, ErrExternalSessionStale) + require.Equal(t, int32(2), initializerCalls.Load(), "stale handles must fail before initialization") + var lateCleanupCalls atomic.Int32 + err = session.SetCleanupHook(func() { lateCleanupCalls.Add(1) }) + require.ErrorIs(t, err, ErrExternalSessionStale) + require.Equal(t, int32(1), lateCleanupCalls.Load()) + + // Reset owns a fresh UDF-registration generation. Keep this sqlmock-focused + // test independent of DuckDB's UDF implementation. + p.registerMySQLUDFsOnce.Do(func() {}) + newSession, err := p.OpenExternalSession(context.Background()) + require.NoError(t, err) + require.Equal(t, oldGeneration+1, newSession.generation) + newMock.ExpectClose() + require.NoError(t, newSession.Close()) + require.NoError(t, p.Close()) + require.NoError(t, newMock.ExpectationsWereMet()) +} + +func TestExternalSessionPoolCloseDrainsLeaseThenCleansBeforeRollback(t *testing.T) { + p, _, mock := newExternalSessionTestPool(t) + session, err := p.OpenExternalSession(context.Background()) + require.NoError(t, err) + mock.ExpectBegin() + tx, err := session.BeginTx(context.Background(), nil) + require.NoError(t, err) + + conn := session.conn + var cleanupComplete atomic.Bool + hookErr := make(chan error, 1) + require.NoError(t, session.SetCleanupHook(func() { + err := conn.Raw(func(any) error { + cleanupComplete.Store(true) + return nil + }) + hookErr <- err + })) + conn, leasedTx, release, err := session.ExecutionLease(context.Background()) + require.NoError(t, err) + require.Same(t, tx, leasedTx) + require.Same(t, session.conn, conn) + var completionCalls atomic.Int32 + require.True(t, p.RegisterTransactionCompletion(tx, func(success bool) { + require.False(t, success) + require.True(t, cleanupComplete.Load(), "cleanup must finish before transaction completion") + completionCalls.Add(1) + })) + + mock.ExpectRollback() + mock.ExpectClose() + closeDone := make(chan error, 1) + go func() { closeDone <- p.Close() }() + waitForPoolShutdown(t, p) + select { + case err := <-closeDone: + t.Fatalf("Close completed before the execution lease released: %v", err) + case <-time.After(100 * time.Millisecond): + } + release() + + select { + case err := <-closeDone: + require.NoError(t, err) + case <-time.After(3 * time.Second): + t.Fatal("Close did not finish after the execution lease released") + } + require.NoError(t, <-hookErr) + require.True(t, cleanupComplete.Load()) + require.Equal(t, int32(1), completionCalls.Load()) + require.NoError(t, mock.ExpectationsWereMet()) + _, _, _, err = session.ExecutionLease(context.Background()) + require.ErrorIs(t, err, ErrExternalSessionStale) +} + +func TestExternalSessionCleanupPanicDoesNotSkipLaterSessions(t *testing.T) { + p, _, mock := newExternalSessionTestPool(t) + first, err := p.OpenExternalSession(context.Background()) + require.NoError(t, err) + second, err := p.OpenExternalSession(context.Background()) + require.NoError(t, err) + + var firstCalls atomic.Int32 + var secondCalls atomic.Int32 + require.NoError(t, first.SetCleanupHook(func() { + firstCalls.Add(1) + panic("first cleanup failed") + })) + require.NoError(t, second.SetCleanupHook(func() { secondCalls.Add(1) })) + mock.ExpectClose() + mock.ExpectClose() + err = p.Close() + require.ErrorContains(t, err, "first cleanup failed") + require.Equal(t, int32(1), firstCalls.Load()) + require.Equal(t, int32(1), secondCalls.Load()) + require.NoError(t, mock.ExpectationsWereMet()) + require.Nil(t, first.conn) + require.Nil(t, second.conn) +} + +func TestExternalSessionConcurrentCloseWaitsForCleanup(t *testing.T) { + p, _, mock := newExternalSessionTestPool(t) + session, err := p.OpenExternalSession(context.Background()) + require.NoError(t, err) + hookStarted := make(chan struct{}) + releaseHook := make(chan struct{}) + var hookCalls atomic.Int32 + require.NoError(t, session.SetCleanupHook(func() { + hookCalls.Add(1) + close(hookStarted) + <-releaseHook + })) + mock.ExpectClose() + + firstDone := make(chan error, 1) + go func() { firstDone <- session.Close() }() + select { + case <-hookStarted: + case <-time.After(time.Second): + t.Fatal("cleanup hook did not start") + } + + secondDone := make(chan error, 1) + go func() { secondDone <- session.Close() }() + select { + case err := <-secondDone: + t.Fatalf("concurrent Close returned while cleanup was still running: %v", err) + case <-time.After(100 * time.Millisecond): + } + close(releaseHook) + require.NoError(t, <-firstDone) + require.NoError(t, <-secondDone) + require.Equal(t, int32(1), hookCalls.Load()) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestExternalSessionTeardownWaitsForCompletionDispatch(t *testing.T) { + p, _, mock := newExternalSessionTestPool(t) + session, err := p.OpenExternalSession(context.Background()) + require.NoError(t, err) + mock.ExpectBegin() + tx, err := session.BeginTx(context.Background(), nil) + require.NoError(t, err) + var callbackStarted atomic.Bool + callbackRelease := make(chan struct{}) + require.True(t, p.RegisterTransactionCompletion(tx, func(bool) { + callbackStarted.Store(true) + <-callbackRelease + })) + _, _, leaseRelease, err := session.ExecutionLease(context.Background()) + require.NoError(t, err) + mock.ExpectRollback() + mock.ExpectClose() + + closeDone := make(chan error, 1) + go func() { closeDone <- p.Close() }() + waitForPoolShutdown(t, p) + leaseRelease() + select { + case <-closeDone: + t.Fatal("pool Close returned before completion dispatch was released") + case <-time.After(100 * time.Millisecond): + } + // The callback runs after the pool writer unlocks. A concurrent release-side + // Close must still wait for the teardown completion signal, rather than + // returning while the terminal callback is in flight. + secondDone := make(chan error, 1) + go func() { secondDone <- session.Close() }() + select { + case err := <-secondDone: + t.Fatalf("external Close returned while completion callback was running: %v", err) + case <-time.After(100 * time.Millisecond): + } + require.Eventually(t, callbackStarted.Load, time.Second, time.Millisecond) + close(callbackRelease) + require.NoError(t, <-closeDone) + require.NoError(t, <-secondDone) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestExternalSessionInitializerFailureRetiresHandle(t *testing.T) { + p, _, mock := newExternalSessionTestPool(t) + initializerErr := errors.New("initializer failed") + var calls atomic.Int32 + p.SetConnectionInitializer(func(context.Context, *stdsql.Conn) error { + if calls.Add(1) == 2 { + return initializerErr + } + return nil + }) + session, err := p.OpenExternalSession(context.Background()) + require.NoError(t, err) + mock.ExpectClose() + _, _, _, err = session.ExecutionLease(context.Background()) + require.ErrorIs(t, err, initializerErr) + require.Equal(t, int32(2), calls.Load()) + _, _, _, err = session.ExecutionLease(context.Background()) + require.ErrorIs(t, err, ErrExternalSessionClosed) + require.NoError(t, mock.ExpectationsWereMet()) +} + +type panicBeginConnector struct { + closes *atomic.Int32 + closeErr error + closePanic bool +} + +func (c panicBeginConnector) Connect(context.Context) (driver.Conn, error) { + return &panicBeginConn{closes: c.closes, closeErr: c.closeErr, closePanic: c.closePanic}, nil +} + +func (c panicBeginConnector) Driver() driver.Driver { + return panicBeginDriver{closes: c.closes, closeErr: c.closeErr, closePanic: c.closePanic} +} + +type panicBeginDriver struct { + closes *atomic.Int32 + closeErr error + closePanic bool +} + +func (d panicBeginDriver) Open(string) (driver.Conn, error) { + return &panicBeginConn{closes: d.closes, closeErr: d.closeErr, closePanic: d.closePanic}, nil +} + +type panicBeginConn struct { + closes *atomic.Int32 + closeErr error + closePanic bool +} + +func (*panicBeginConn) Prepare(string) (driver.Stmt, error) { + return nil, errors.New("prepare unsupported") +} +func (c *panicBeginConn) Close() error { + if c.closes != nil { + c.closes.Add(1) + } + if c.closePanic { + panic("close panic") + } + return c.closeErr +} +func (*panicBeginConn) Begin() (driver.Tx, error) { panic("begin panic") } +func (*panicBeginConn) BeginTx(context.Context, driver.TxOptions) (driver.Tx, error) { + panic("begin tx panic") +} + +func TestExternalSessionBeginPanicRetiresHandle(t *testing.T) { + var physicalCloses atomic.Int32 + db := stdsql.OpenDB(panicBeginConnector{closes: &physicalCloses}) + t.Cleanup(func() { _ = db.Close() }) + db.SetMaxOpenConns(1) + p := NewConnectionPool(nil, db, "memory") + p.registerMySQLUDFsOnce.Do(func() {}) + session, err := p.OpenExternalSession(context.Background()) + require.NoError(t, err) + var cleanupCalls atomic.Int32 + require.NoError(t, session.SetCleanupHook(func() { cleanupCalls.Add(1) })) + + beginDone := make(chan error, 1) + go func() { + _, beginErr := session.BeginTx(context.Background(), nil) + beginDone <- beginErr + }() + select { + case err = <-beginDone: + require.ErrorContains(t, err, "begin transaction panicked") + case <-time.After(time.Second): + t.Fatal("BeginTx did not return after the driver panicked") + } + require.Equal(t, int32(1), cleanupCalls.Load()) + require.Eventually(t, func() bool { + return physicalCloses.Load() == 1 + }, time.Second, time.Millisecond) + _, _, _, err = session.ExecutionLease(context.Background()) + require.ErrorIs(t, err, ErrExternalSessionGenerationPoisoned) + _, err = p.OpenExternalSession(context.Background()) + require.ErrorIs(t, err, ErrExternalSessionGenerationPoisoned) + + // database/sql cannot release the Conn read lock lost by its panic path. + // The exact connection must remain checked out rather than becoming reusable. + acquireCtx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + _, err = db.Conn(acquireCtx) + require.ErrorIs(t, err, context.DeadlineExceeded) + + replacementDB, _, err := sqlmock.New() + require.NoError(t, err) + t.Cleanup(func() { _ = replacementDB.Close() }) + resetDone := make(chan error, 1) + go func() { resetDone <- p.Reset(nil, replacementDB) }() + select { + case err = <-resetDone: + require.ErrorIs(t, err, ErrExternalSessionGenerationPoisoned) + case <-time.After(time.Second): + t.Fatal("pool Reset waited for a quarantined external connection") + } + require.Same(t, db, p.DB, "Reset must not install a new in-process generation after a driver panic") + + closeDone := make(chan error, 1) + go func() { closeDone <- p.Close() }() + select { + case err = <-closeDone: + require.ErrorIs(t, err, ErrExternalSessionGenerationPoisoned) + case <-time.After(time.Second): + t.Fatal("pool Close waited for a quarantined external connection") + } + require.Equal(t, int32(1), physicalCloses.Load()) +} + +func TestExternalSessionBeginPanicPhysicalCloseFailureIsBounded(t *testing.T) { + tests := []struct { + name string + closeErr error + closePanic bool + expectedError string + }{ + { + name: "driver ErrBadConn", + closeErr: driver.ErrBadConn, + expectedError: driver.ErrBadConn.Error(), + }, + { + name: "driver close panic", + closePanic: true, + expectedError: "driver operation panicked: close panic", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var physicalCloses atomic.Int32 + db := stdsql.OpenDB(panicBeginConnector{ + closes: &physicalCloses, + closeErr: tt.closeErr, + closePanic: tt.closePanic, + }) + t.Cleanup(func() { _ = db.Close() }) + db.SetMaxOpenConns(1) + p := NewConnectionPool(nil, db, "memory") + p.registerMySQLUDFsOnce.Do(func() {}) + session, err := p.OpenExternalSession(context.Background()) + require.NoError(t, err) + + beginDone := make(chan error, 1) + go func() { + _, beginErr := session.BeginTx(context.Background(), nil) + beginDone <- beginErr + }() + select { + case err = <-beginDone: + require.ErrorContains(t, err, "begin transaction panicked") + require.ErrorContains(t, err, tt.expectedError) + case <-time.After(time.Second): + t.Fatal("BeginTx waited for database/sql retirement after driver Close failed") + } + require.Equal(t, int32(1), physicalCloses.Load()) + + closeDone := make(chan error, 1) + go func() { closeDone <- p.Close() }() + select { + case err = <-closeDone: + require.ErrorIs(t, err, ErrExternalSessionGenerationPoisoned) + case <-time.After(time.Second): + t.Fatal("pool Close waited after driver Close failed") + } + require.Equal(t, int32(1), physicalCloses.Load()) + }) + } +} + +func TestExternalSessionBeginPanicPoisonsBeforeCleanupCompletes(t *testing.T) { + var physicalCloses atomic.Int32 + db := stdsql.OpenDB(panicBeginConnector{closes: &physicalCloses}) + t.Cleanup(func() { _ = db.Close() }) + db.SetMaxOpenConns(1) + p := NewConnectionPool(nil, db, "memory") + p.registerMySQLUDFsOnce.Do(func() {}) + session, err := p.OpenExternalSession(context.Background()) + require.NoError(t, err) + + hookStarted := make(chan struct{}) + releaseHook := make(chan struct{}) + require.NoError(t, session.SetCleanupHook(func() { + close(hookStarted) + <-releaseHook + })) + beginDone := make(chan error, 1) + go func() { + _, beginErr := session.BeginTx(context.Background(), nil) + beginDone <- beginErr + }() + select { + case <-hookStarted: + case <-time.After(time.Second): + t.Fatal("cleanup did not start after the BeginTx driver panic") + } + + openCtx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + _, err = p.OpenExternalSession(openCtx) + require.ErrorIs(t, err, ErrExternalSessionGenerationPoisoned) + close(releaseHook) + select { + case err = <-beginDone: + require.ErrorContains(t, err, "begin transaction panicked") + case <-time.After(time.Second): + t.Fatal("BeginTx did not finish after cleanup was released") + } + require.Equal(t, int32(1), physicalCloses.Load()) + require.ErrorIs(t, p.Close(), ErrExternalSessionGenerationPoisoned) +} + +type panicTransactionConnector struct { + operation string + closes *atomic.Int32 +} + +func (c panicTransactionConnector) Connect(context.Context) (driver.Conn, error) { + return &panicTransactionConn{operation: c.operation, closes: c.closes}, nil +} + +func (c panicTransactionConnector) Driver() driver.Driver { + return panicTransactionDriver{operation: c.operation, closes: c.closes} +} + +type panicTransactionDriver struct { + operation string + closes *atomic.Int32 +} + +func (d panicTransactionDriver) Open(string) (driver.Conn, error) { + return &panicTransactionConn{operation: d.operation, closes: d.closes}, nil +} + +type panicTransactionConn struct { + operation string + closes *atomic.Int32 +} + +func (*panicTransactionConn) Prepare(string) (driver.Stmt, error) { + return nil, errors.New("prepare unsupported") +} + +func (c *panicTransactionConn) Close() error { + if c.closes != nil { + c.closes.Add(1) + } + return nil +} + +func (c *panicTransactionConn) Begin() (driver.Tx, error) { + return &panicTransactionTx{operation: c.operation}, nil +} + +func (c *panicTransactionConn) BeginTx(context.Context, driver.TxOptions) (driver.Tx, error) { + return &panicTransactionTx{operation: c.operation}, nil +} + +type panicTransactionTx struct { + operation string +} + +func (tx *panicTransactionTx) Commit() error { + if tx.operation == "commit" { + panic("commit panic") + } + return nil +} + +func (tx *panicTransactionTx) Rollback() error { + if tx.operation == "rollback" { + panic("rollback panic") + } + return nil +} + +func TestExternalSessionTransactionPanicQuarantinesConnection(t *testing.T) { + tests := []struct { + name string + driverPanic string + terminalError string + terminal func(*ExternalSession, *stdsql.Tx) error + }{ + { + name: "commit", + driverPanic: "commit", + terminalError: "commit transaction panicked", + terminal: func(session *ExternalSession, tx *stdsql.Tx) error { return session.Commit(tx) }, + }, + { + name: "rollback", + driverPanic: "rollback", + terminalError: "rollback transaction panicked", + terminal: func(session *ExternalSession, tx *stdsql.Tx) error { + return session.Rollback(tx, nil) + }, + }, + { + name: "close with active transaction", + driverPanic: "rollback", + terminalError: "rollback transaction panicked", + terminal: func(session *ExternalSession, _ *stdsql.Tx) error { return session.Close() }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var physicalCloses atomic.Int32 + db := stdsql.OpenDB(panicTransactionConnector{ + operation: tt.driverPanic, + closes: &physicalCloses, + }) + t.Cleanup(func() { _ = db.Close() }) + db.SetMaxOpenConns(1) + p := NewConnectionPool(nil, db, "memory") + p.registerMySQLUDFsOnce.Do(func() {}) + + var finishedCalls atomic.Int32 + p.SetTransactionLifecycleHooks(nil, func(*stdsql.Tx) { finishedCalls.Add(1) }) + session, err := p.OpenExternalSession(context.Background()) + require.NoError(t, err) + var cleanupCalls atomic.Int32 + require.NoError(t, session.SetCleanupHook(func() { cleanupCalls.Add(1) })) + tx, err := session.BeginTx(context.Background(), nil) + require.NoError(t, err) + var completionCalls atomic.Int32 + require.True(t, p.RegisterTransactionCompletion(tx, func(success bool) { + require.False(t, success) + completionCalls.Add(1) + })) + + terminalDone := make(chan error, 1) + go func() { terminalDone <- tt.terminal(session, tx) }() + select { + case err = <-terminalDone: + require.ErrorContains(t, err, tt.terminalError) + case <-time.After(time.Second): + t.Fatal("transaction terminal operation did not return after the driver panicked") + } + require.Equal(t, int32(1), cleanupCalls.Load()) + require.Equal(t, int32(1), completionCalls.Load()) + require.Equal(t, int32(1), finishedCalls.Load()) + require.Eventually(t, func() bool { + return physicalCloses.Load() == 1 + }, time.Second, time.Millisecond) + + _, _, _, err = session.ExecutionLease(context.Background()) + require.ErrorIs(t, err, ErrExternalSessionGenerationPoisoned) + acquireCtx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + _, err = db.Conn(acquireCtx) + require.ErrorIs(t, err, context.DeadlineExceeded) + + closeDone := make(chan error, 1) + go func() { closeDone <- p.Close() }() + select { + case err = <-closeDone: + require.ErrorIs(t, err, ErrExternalSessionGenerationPoisoned) + case <-time.After(time.Second): + t.Fatal("pool Close waited for a quarantined external connection") + } + require.Equal(t, int32(1), physicalCloses.Load()) + }) + } +} + +func TestExternalSessionSetCleanupHookAdmittedBeforeShutdownWriter(t *testing.T) { + p, _, mock := newExternalSessionTestPool(t) + session, err := p.OpenExternalSession(context.Background()) + require.NoError(t, err) + _, _, release, err := session.ExecutionLease(context.Background()) + require.NoError(t, err) + closeDone := make(chan error, 1) + go func() { closeDone <- p.Close() }() + waitForPoolShutdown(t, p) + var hookCalls atomic.Int32 + setDone := make(chan error, 1) + go func() { + setDone <- session.SetCleanupHook(func() { hookCalls.Add(1) }) + }() + select { + case err := <-setDone: + t.Fatalf("cleanup hook registration unexpectedly completed while lease was held: %v", err) + case <-time.After(100 * time.Millisecond): + } + mock.ExpectClose() + release() + require.NoError(t, <-setDone) + require.NoError(t, <-closeDone) + require.Equal(t, int32(1), hookCalls.Load()) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestExternalSessionUnexpectedRollbackErrorRetiresHandle(t *testing.T) { + p, _, mock := newExternalSessionTestPool(t) + var finishedCalls atomic.Int32 + p.SetTransactionLifecycleHooks(nil, func(*stdsql.Tx) { finishedCalls.Add(1) }) + session, err := p.OpenExternalSession(context.Background()) + require.NoError(t, err) + var cleanupCalls atomic.Int32 + require.NoError(t, session.SetCleanupHook(func() { cleanupCalls.Add(1) })) + mock.ExpectBegin() + tx, err := session.BeginTx(context.Background(), nil) + require.NoError(t, err) + + rollbackErr := errors.New("rollback state is unknown") + mock.ExpectRollback().WillReturnError(rollbackErr) + mock.ExpectClose() + err = session.Rollback(tx, nil) + require.ErrorIs(t, err, rollbackErr) + require.Equal(t, int32(1), cleanupCalls.Load()) + require.Equal(t, int32(1), finishedCalls.Load()) + require.NoError(t, mock.ExpectationsWereMet()) + + _, _, _, err = session.ExecutionLease(context.Background()) + require.ErrorIs(t, err, ErrExternalSessionClosed) +} diff --git a/catalog/identity_lifecycle_test.go b/catalog/identity_lifecycle_test.go new file mode 100644 index 00000000..26b94e64 --- /dev/null +++ b/catalog/identity_lifecycle_test.go @@ -0,0 +1,94 @@ +package catalog + +import ( + "database/sql" + "sync/atomic" + "testing" + + "github.com/DATA-DOG/go-sqlmock" + "github.com/stretchr/testify/require" +) + +// CloseTxnAndGet must resolve the binding at the point it acquires the +// release/session section. A caller may have captured an older snapshot before +// a replacement transaction was installed; the replacement is the only entry +// that this close operation is allowed to remove and complete. +func TestCloseTxnAndGetUsesCurrentReplacementAndReturnsNilWhenAbsent(t *testing.T) { + p := NewConnectionPool(nil, nil, "memory") + const sessionID uint32 = 301 + first := new(sql.Tx) + replacement := new(sql.Tx) + + var firstCompletions atomic.Int32 + var replacementCompletions atomic.Int32 + p.txns.Store(sessionID, first) + require.True(t, p.RegisterTransactionCompletion(first, func(bool) { + firstCompletions.Add(1) + })) + // This is the stale read that the old lookup-then-close sequence could use. + staleSnapshot := p.TryGetTxnForRelease(sessionID) + require.Same(t, first, staleSnapshot) + + p.txns.Store(sessionID, replacement) + require.True(t, p.RegisterTransactionCompletion(replacement, func(success bool) { + require.False(t, success) + replacementCompletions.Add(1) + })) + + removed := p.CloseTxnAndGet(sessionID) + require.Same(t, replacement, removed) + require.Zero(t, firstCompletions.Load(), "the stale snapshot must not be completed") + require.Equal(t, int32(1), replacementCompletions.Load()) + require.Nil(t, p.TryGetTxnForRelease(sessionID)) + + // A nil result is an observation, not permission to remove a later binding. + require.Nil(t, p.CloseTxnAndGet(sessionID)) + require.Nil(t, p.TryGetTxnForRelease(sessionID)) +} + +// A malformed/missing physical owner is a binding error, not a terminal +// transaction event. The completion callback must stay pending until a later +// finalizer can remove the exact mapping. +func TestRollbackTxnBindingErrorLeavesCompletionPending(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + p := NewConnectionPool(nil, db, "memory") + p.registerMySQLUDFsOnce.Do(func() {}) + t.Cleanup(func() { _ = db.Close() }) + const sessionID uint32 = 302 + mock.ExpectBegin() + tx, err := p.GetTxn(nil, sessionID, "", nil) + require.NoError(t, err) + owner, bound := p.GetTxnBindingForRelease(sessionID) + require.NotNil(t, owner) + require.Same(t, tx, bound) + + var completions atomic.Int32 + require.True(t, p.RegisterTransactionCompletion(tx, func(success bool) { + require.False(t, success) + completions.Add(1) + })) + + // Corrupt only the owner side. Rollback must fail closed and leave tx mapped + // so its deferred operation lease cannot be released early. + p.txnConns.Delete(sessionID) + conn, inactive, err := p.RollbackTxn(sessionID, tx) + require.Nil(t, conn) + require.False(t, inactive) + require.ErrorContains(t, err, "physical owner") + require.Zero(t, completions.Load()) + require.Same(t, tx, p.TryGetTxnForRelease(sessionID)) + + // Repair the binding and let the real finalizer close it. + p.txnConns.Store(sessionID, owner) + mock.ExpectRollback() + conn, inactive, err = p.RollbackTxn(sessionID, tx) + require.Same(t, owner, conn) + require.True(t, inactive) + require.NoError(t, err) + require.Equal(t, int32(1), completions.Load()) + require.Nil(t, p.TryGetTxnForRelease(sessionID)) + mock.ExpectClose() + require.NoError(t, p.Close()) + require.NoError(t, mock.ExpectationsWereMet()) +} diff --git a/catalog/inserter.go b/catalog/inserter.go index 882dbb9d..d01d3d90 100644 --- a/catalog/inserter.go +++ b/catalog/inserter.go @@ -12,14 +12,19 @@ import ( ) type rowInserter struct { - db string - table string - schema sql.Schema - hasPK bool - replace bool + catalog string + db string + table string + schema sql.Schema + hasPK bool + replace bool + provider *DatabaseProvider once sync.Once conn *stdsql.Conn + tx *stdsql.Tx + execer adapter.SQLExecutor + release func() tmpTable string stmt *stdsql.Stmt err error @@ -39,17 +44,30 @@ func (ri *rowInserter) Delete(ctx *sql.Context, row sql.Row) error { func (ri *rowInserter) init(ctx *sql.Context) { ri.tmpTable = fmt.Sprintf("%s_%s_%d", ri.db, ri.table, ctx.ID()) - ri.conn, ri.err = ctx.Session.(adapter.ConnectionHolder).GetConn(ctx) + ri.execer, ri.conn, ri.tx, ri.release, ri.err = adapter.GetExecutionSnapshotWithLease(ctx) if ri.err != nil { return } + target := ConnectIdentifiersANSI(ri.db, ri.table) + if ri.provider != nil && ri.catalog != "" && ri.provider.DuckLakeEnabled() { + if err := ri.provider.EnsureDuckLakeConnectionWithExecutor(ctx, ri.execer, ri.conn); err != nil { + ri.err = err + return + } + if physical, found, err := ri.provider.ObjectTableNameWithExecutor(ctx, ri.execer, ri.catalog, ri.db, ri.table); err != nil { + ri.err = err + return + } else if found { + target = physical + } + } ctx.GetLogger().WithField("db", ri.db).WithField("table", ri.table).Infoln("Creating temp table", ri.tmpTable) createTable := fmt.Sprintf( "CREATE TEMP TABLE IF NOT EXISTS %s AS FROM %s LIMIT 0", QuoteIdentifierANSI(ri.tmpTable), - ConnectIdentifiersANSI(ri.db, ri.table), + target, ) - if _, ri.err = ri.conn.ExecContext(ctx, createTable); ri.err != nil { + if _, ri.err = ri.execer.ExecContext(ctx, createTable); ri.err != nil { return } @@ -64,7 +82,7 @@ func (ri *rowInserter) init(ctx *sql.Context) { insert.WriteString(", ?") } insert.WriteByte(')') - ri.stmt, ri.err = ri.conn.PrepareContext(ctx, insert.String()) + ri.stmt, ri.err = ri.execer.PrepareContext(ctx, insert.String()) if ri.err != nil { return } @@ -75,7 +93,7 @@ func (ri *rowInserter) init(ctx *sql.Context) { insert.WriteString(" OR REPLACE") } insert.WriteString(" INTO ") - insert.WriteString(ConnectIdentifiersANSI(ri.db, ri.table)) + insert.WriteString(target) insert.WriteString(" FROM ") insert.WriteString(QuoteIdentifierANSI(ri.tmpTable)) ri.flushSQL = insert.String() @@ -103,11 +121,13 @@ func (ri *rowInserter) StatementComplete(ctx *sql.Context) error { func (ri *rowInserter) Close(ctx *sql.Context) error { ri.StatementBegin(ctx) - defer ri.clear(ctx) if ri.err == nil { - _, ri.err = ri.conn.ExecContext(ctx, ri.flushSQL) + _, ri.err = ri.execer.ExecContext(ctx, ri.flushSQL) } - return ri.err + // clear closes the prepared statement and drops the temporary table. Return + // its error as well as the write error so an owner/lifecycle change cannot be + // silently discarded behind a successful-looking INSERT. + return ri.clear(ctx) } func (ri *rowInserter) Insert(ctx *sql.Context, row sql.Row) error { @@ -137,12 +157,57 @@ func (ri *rowInserter) Insert(ctx *sql.Context, row sql.Row) error { } func (ri *rowInserter) clear(ctx *sql.Context) error { + defer ri.releaseSnapshot() if ri.stmt != nil { - ri.err = errors.Join(ri.err, ri.stmt.Close()) + closeErr := ri.stmt.Close() + // A statement prepared on the session transaction reports ErrTxDone when + // GMS invokes cleanup after ROLLBACK. The transaction is already inactive + // in that case, so it is not a product failure and must not mask the + // rollback result. + if !errors.Is(closeErr, stdsql.ErrTxDone) { + ri.err = errors.Join(ri.err, closeErr) + } + ri.stmt = nil } if ri.conn != nil { - _, err := ri.conn.ExecContext(ctx, "DROP TABLE IF EXISTS temp.main."+QuoteIdentifierANSI(ri.tmpTable)) - ri.err = errors.Join(ri.err, err) + dropper := ri.execer + if ri.tx != nil { + currentConn, current := adapter.TryGetTxnBindingForRelease(ctx) + sameOwner := current == ri.tx && (currentConn == nil || currentConn == ri.conn) + sessionIdle := current == nil && (currentConn == nil || currentConn == ri.conn) + switch { + case sameOwner: + // Keep the transaction executor while its owner is still active. + case sessionIdle: + // Once the owner has been rolled back, reuse that same physical + // connection rather than acquiring a second pooled connection. + dropper = ri.conn + default: + // A replacement transaction owns this connection (or the old + // connection was evicted). Issuing DROP through ri.conn would race + // the replacement and could corrupt its transaction. The original + // transaction's temporary objects are transaction-scoped; surface + // the lifecycle change and leave the replacement untouched. + return errors.Join(ri.err, fmt.Errorf("temporary table cleanup owner changed")) + } + } + if dropper != nil { + _, err := dropper.ExecContext(ctx, "DROP TABLE IF EXISTS temp.main."+QuoteIdentifierANSI(ri.tmpTable)) + if errors.Is(err, stdsql.ErrTxDone) && dropper != ri.conn { + // This is the narrow race where rollback happened after the + // transaction check above. Retry on the now-inactive connection. + _, err = ri.conn.ExecContext(ctx, "DROP TABLE IF EXISTS temp.main."+QuoteIdentifierANSI(ri.tmpTable)) + } + ri.err = errors.Join(ri.err, err) + } } return ri.err } + +func (ri *rowInserter) releaseSnapshot() { + if ri.release == nil { + return + } + ri.release() + ri.release = nil +} diff --git a/catalog/object_storage.go b/catalog/object_storage.go new file mode 100644 index 00000000..2a62a18d --- /dev/null +++ b/catalog/object_storage.go @@ -0,0 +1,663 @@ +package catalog + +import ( + "context" + stdsql "database/sql" + "fmt" + "strings" + + "github.com/apecloud/myduckserver/adapter" + "github.com/apecloud/myduckserver/mycontext" + "github.com/dolthub/go-mysql-server/sql" +) + +// objectColumnDefinition is deliberately a small, trusted representation of a +// column read from the local shadow table. DuckLake's first slice supports +// ordinary columns/defaults/nullability, but not generated columns, keys, or +// checks. +type objectColumnDefinition struct { + name string + dataType string + nullable bool + defaultValue string +} + +func (d *Database) createObjectTable(ctx *sql.Context, name string, schema sql.PrimaryKeySchema, _ sql.CollationID, comment string, storage TableStorageSelection) error { + if d.provider == nil || d.provider.duckLake == nil { + return fmt.Errorf("%w: DuckLake service configuration is disabled", ErrInvalidTableStorage) + } + if err := validateObjectSchema(schema); err != nil { + return err + } + lakeCtx, lakeConn, err := d.acquireDedicatedLakeConn() + if err != nil { + return err + } + defer lakeConn.Close() + physical := d.objectPhysicalTableName(name) + if err := d.ensureObjectSchema(lakeCtx, lakeConn); err != nil { + return err + } + ddl, err := objectCreateSQL(physical, schema) + if err != nil { + return err + } + if _, err := lakeConn.ExecContext(lakeCtx, ddl); err != nil { + if IsDuckDBTableAlreadyExistsError(err) { + return sql.ErrTableAlreadyExists.New(name) + } + return ErrDuckDB.New(err) + } + d.registerUncommittedLakeRelation(ctx, physical) + + execer, _, _, release, err := adapter.GetCatalogExecutionSnapshotWithLease(ctx) + if err != nil { + d.dropLakeRelation(lakeCtx, lakeConn, ctx, physical) + return err + } + defer release() + // Keep a managed local shadow for GMS catalog discovery. If shadow creation + // fails, remove the just-created lake relation so no unregistered files are + // left behind. Lake DDL uses a dedicated connection because DuckDB allows + // one transaction to write only one attached database. + if err := d.createLocalTableWithExecutor(ctx, name, schema, sql.Collation_Default, comment, storage, false, execer); err != nil { + d.dropLakeRelation(lakeCtx, lakeConn, ctx, physical) + return err + } + return nil +} + +func validateObjectSchema(schema sql.PrimaryKeySchema) error { + if len(schema.PkOrdinals) > 0 { + return fmt.Errorf("%w: object tables do not support primary keys", ErrInvalidTableStorage) + } + for _, col := range schema.Schema { + if col == nil { + continue + } + if col.PrimaryKey || col.AutoIncrement || col.Generated != nil { + return fmt.Errorf("%w: object tables do not support key, generated, or auto-increment columns", ErrInvalidTableStorage) + } + } + return nil +} + +func objectCreateSQL(physical string, schema sql.PrimaryKeySchema) (string, error) { + columns := make([]string, 0, len(schema.Schema)) + for _, col := range schema.Schema { + typ, err := DuckdbDataType(col.Type) + if err != nil { + return "", err + } + part := QuoteIdentifierANSI(col.Name) + " " + typ.name + if col.Nullable { + part += " NULL" + } else { + part += " NOT NULL" + } + if col.Default != nil { + expr, err := parseDefaultValue(col.Default.String()) + if err != nil { + return "", err + } + part += " DEFAULT " + expr + } + columns = append(columns, part) + } + return "CREATE TABLE " + physical + " (" + strings.Join(columns, ", ") + ")", nil +} + +func (d *Database) ensureObjectSchema(ctx context.Context, execer adapter.SQLExecutor) error { + // This interface is intentionally satisfied by *sql.Conn and keeps this + // helper independent from a concrete database/sql wrapper in tests. + schema := d.provider.LakeSchemaName(d.catalog, d.name) + query := "CREATE SCHEMA IF NOT EXISTS " + FullSchemaName(DuckLakeCatalogName, schema) + _, err := execer.ExecContext(ctx, query) + if err != nil { + return fmt.Errorf("create DuckLake schema failed") + } + return nil +} + +func lakeTransactionInactive(err error) bool { + if err == nil { + return true + } + lower := strings.ToLower(err.Error()) + return strings.Contains(lower, "no transaction") || + (strings.Contains(lower, "transaction") && strings.Contains(lower, "not active")) || + adapter.IsTransactionInactiveError(err) +} + +func (d *Database) acquireDedicatedLakeConn() (context.Context, *stdsql.Conn, error) { + if d == nil || d.provider == nil { + return nil, nil, fmt.Errorf("ducklake storage is unavailable") + } + storageDB := d.provider.Storage() + if storageDB == nil { + return nil, nil, fmt.Errorf("ducklake storage is unavailable") + } + lakeCtx := mycontext.WithMaintenanceQuery(context.Background()) + lakeConn, err := storageDB.Conn(lakeCtx) + if err != nil { + return nil, nil, err + } + if err := d.provider.EnsureDuckLakeConnectionWithExecutor(lakeCtx, lakeConn, lakeConn); err != nil { + _ = lakeConn.Close() + return nil, nil, err + } + if _, err := lakeConn.ExecContext(lakeCtx, "COMMIT"); err != nil && !lakeTransactionInactive(err) { + _ = lakeConn.Close() + return nil, nil, err + } + return lakeCtx, lakeConn, nil +} + +func sessionTxnBinding(ctx *sql.Context) (*stdsql.Conn, *stdsql.Tx) { + if ctx == nil || ctx.Session == nil { + return nil, nil + } + if _, ok := ctx.Session.(adapter.ConnectionHolder); !ok { + return nil, nil + } + return adapter.TryGetTxnBinding(ctx) +} + +func (d *Database) registerUncommittedLakeRelation(ctx *sql.Context, physical string) { + if d == nil || d.provider == nil { + return + } + conn, tx := sessionTxnBinding(ctx) + d.provider.registerUncommittedLakeRelation(tx, conn, physical) +} + +func (d *Database) dropLakeRelation(lakeCtx context.Context, lakeConn *stdsql.Conn, sessionCtx *sql.Context, physical string) { + if strings.TrimSpace(physical) == "" { + return + } + if sessionCtx != nil { + _, tx := sessionTxnBinding(sessionCtx) + if d != nil && d.provider != nil { + d.provider.forgetUncommittedLakeRelation(tx, physical) + } + } + if lakeConn != nil && lakeCtx != nil { + _, _ = lakeConn.ExecContext(lakeCtx, "DROP TABLE IF EXISTS "+physical) + return + } + if d == nil { + return + } + dropCtx, dropConn, err := d.acquireDedicatedLakeConn() + if err != nil { + return + } + defer dropConn.Close() + _, _ = dropConn.ExecContext(dropCtx, "DROP TABLE IF EXISTS "+physical) +} + +// objectPhysicalTableName returns the fixed lake-catalog relation for a table +// in this logical database. Callers must have validated/attached the runtime +// before executing it. +func (d *Database) objectPhysicalTableName(name string) string { + return FullTableName(DuckLakeCatalogName, d.provider.LakeSchemaName(d.catalog, d.name), name) +} + +// physicalTableName resolves the durable object-table relation for a table +// handle and initializes DuckLake on the same session connection that will +// execute the caller's statement. Local tables continue to use their managed +// catalog relation unchanged. +func (t *Table) physicalTableName(ctx *sql.Context) (string, error) { + if t == nil || t.db == nil { + return "", fmt.Errorf("table is unavailable") + } + execer, conn, _, release, err := adapter.GetExecutionSnapshotWithLease(ctx) + if err != nil { + return "", err + } + defer release() + return t.physicalTableNameWithExecutor(ctx, execer, conn) +} + +// physicalTableNameWithExecutor resolves a table using an already captured +// execution snapshot. Object metadata lookup and the statement that follows +// must use the same transaction/connection pair; acquiring a fresh snapshot +// here can otherwise observe a different catalog generation while a request is +// in flight. +func (t *Table) physicalTableNameWithExecutor(ctx *sql.Context, execer adapter.SQLExecutor, conn *stdsql.Conn) (string, error) { + if t == nil || t.db == nil { + return "", fmt.Errorf("table is unavailable") + } + if t.ExtraTableInfo().StorageKind() != TableStorageObject { + return FullTableName(t.db.catalog, t.db.name, t.name), nil + } + if t.db.provider == nil || t.db.provider.duckLake == nil { + return "", fmt.Errorf("%w: DuckLake service configuration is disabled", ErrInvalidTableStorage) + } + if execer == nil { + return "", fmt.Errorf("object-table metadata executor is unavailable") + } + if err := t.db.provider.EnsureDuckLakeConnectionWithExecutor(ctx, execer, conn); err != nil { + return "", err + } + physical, found, err := t.db.provider.ObjectTableNameWithExecutor(ctx, execer, t.db.catalog, t.db.name, t.name) + if err != nil { + return "", err + } + if !found { + return "", sql.ErrTableNotFound.New(t.name) + } + return physical, nil +} + +func (t *Table) shadowTableName() string { + return FullTableName(t.db.catalog, t.db.name, t.name) +} + +func (t *Table) objectStorage() bool { + return t != nil && t.ExtraTableInfo().StorageKind() == TableStorageObject +} + +func objectColumnSQL(column *sql.Column) (string, error) { + if column == nil { + return "", fmt.Errorf("column is unavailable") + } + if column.PrimaryKey || column.AutoIncrement || column.Generated != nil { + return "", fmt.Errorf("%w: object tables do not support key, generated, or auto-increment columns", ErrInvalidTableStorage) + } + typ, err := DuckdbDataType(column.Type) + if err != nil { + return "", err + } + part := QuoteIdentifierANSI(column.Name) + " " + typ.name + if column.Nullable { + part += " NULL" + } else { + part += " NOT NULL" + } + if column.Default != nil { + expr, err := parseDefaultValue(column.Default.String()) + if err != nil { + return "", err + } + part += " DEFAULT " + expr + } + return part, nil +} + +func (t *Table) addObjectColumn(ctx *sql.Context, column *sql.Column) error { + part, err := objectColumnSQL(column) + if err != nil { + return err + } + execer, conn, _, release, err := adapter.GetExecutionSnapshotWithLease(ctx) + if err != nil { + return err + } + defer release() + physical, err := t.physicalTableNameWithExecutor(ctx, execer, conn) + if err != nil { + return err + } + physicalSQL := "ALTER TABLE " + physical + " ADD COLUMN " + part + if _, err := execer.ExecContext(ctx, physicalSQL); err != nil { + return ErrDuckDB.New(err) + } + // Keep the shadow schema in lockstep; it is never used as the object table's + // data source but remains the durable GMS catalog representation. + shadowSQL := "ALTER TABLE " + t.shadowTableName() + " ADD COLUMN " + part + if _, err := execer.ExecContext(ctx, shadowSQL); err != nil { + return ErrDuckDB.New(err) + } + if !column.Nullable { + if _, err := execer.ExecContext(ctx, "ALTER TABLE "+physical+" ALTER COLUMN "+QuoteIdentifierANSI(column.Name)+" SET NOT NULL"); err != nil { + return ErrDuckDB.New(err) + } + if _, err := execer.ExecContext(ctx, "ALTER TABLE "+t.shadowTableName()+" ALTER COLUMN "+QuoteIdentifierANSI(column.Name)+" SET NOT NULL"); err != nil { + return ErrDuckDB.New(err) + } + } + comment := NewCommentWithMeta(column.Comment, MySQLType{}) + if _, err := execer.ExecContext(ctx, "COMMENT ON COLUMN "+FullColumnName(t.db.catalog, t.db.name, t.name, column.Name)+" IS '"+comment.Encode()+"'"); err != nil { + return ErrDuckDB.New(err) + } + return t.withSchemaWithExecutor(ctx, execer) +} + +func (t *Table) dropObjectColumn(ctx *sql.Context, columnName string) error { + execer, conn, _, release, err := adapter.GetExecutionSnapshotWithLease(ctx) + if err != nil { + return err + } + defer release() + physical, err := t.physicalTableNameWithExecutor(ctx, execer, conn) + if err != nil { + return err + } + if _, err := execer.ExecContext(ctx, "ALTER TABLE "+physical+" DROP COLUMN "+QuoteIdentifierANSI(columnName)); err != nil { + return ErrDuckDB.New(err) + } + if _, err := execer.ExecContext(ctx, "ALTER TABLE "+t.shadowTableName()+" DROP COLUMN "+QuoteIdentifierANSI(columnName)); err != nil { + return ErrDuckDB.New(err) + } + return t.withSchemaWithExecutor(ctx, execer) +} + +func (t *Table) modifyObjectColumn(ctx *sql.Context, columnName string, column *sql.Column) error { + if column == nil || column.PrimaryKey || column.AutoIncrement || column.Generated != nil { + return fmt.Errorf("%w: object tables do not support key, generated, or auto-increment columns", ErrInvalidTableStorage) + } + var old *sql.Column + for _, candidate := range t.schema.Schema { + if strings.EqualFold(candidate.Name, columnName) { + old = candidate + break + } + } + if old == nil { + return sql.ErrColumnNotFound.New(columnName) + } + typ, err := DuckdbDataType(column.Type) + if err != nil { + return err + } + execer, conn, _, release, err := adapter.GetExecutionSnapshotWithLease(ctx) + if err != nil { + return err + } + defer release() + physical, err := t.physicalTableNameWithExecutor(ctx, execer, conn) + if err != nil { + return err + } + build := func(tableName string) ([]string, error) { + base := "ALTER TABLE " + tableName + " ALTER COLUMN " + QuoteIdentifierANSI(columnName) + var statements []string + if !old.Type.Equals(column.Type) { + statements = append(statements, base+" TYPE "+typ.name) + } + if old.Nullable && !column.Nullable { + statements = append(statements, base+" SET NOT NULL") + } else if !old.Nullable && column.Nullable { + statements = append(statements, base+" DROP NOT NULL") + } + if column.Default != nil { + expr, err := parseDefaultValue(column.Default.String()) + if err != nil { + return nil, err + } + statements = append(statements, base+" SET DEFAULT "+expr) + } else if old.Default != nil { + statements = append(statements, base+" DROP DEFAULT") + } + if columnName != column.Name { + statements = append(statements, "ALTER TABLE "+tableName+" RENAME "+QuoteIdentifierANSI(columnName)+" TO "+QuoteIdentifierANSI(column.Name)) + } + return statements, nil + } + physicalStatements, err := build(physical) + if err != nil { + return err + } + if len(physicalStatements) > 0 { + if _, err := execer.ExecContext(ctx, strings.Join(physicalStatements, "; ")); err != nil { + return ErrDuckDB.New(err) + } + } + shadowStatements, err := build(t.shadowTableName()) + if err != nil { + return err + } + if len(shadowStatements) > 0 { + if _, err := execer.ExecContext(ctx, strings.Join(shadowStatements, "; ")); err != nil { + return ErrDuckDB.New(err) + } + } + commentName := column.Name + if _, err := execer.ExecContext(ctx, "COMMENT ON COLUMN "+FullColumnName(t.db.catalog, t.db.name, t.name, commentName)+" IS '"+NewCommentWithMeta(column.Comment, MySQLType{}).Encode()+"'"); err != nil { + return ErrDuckDB.New(err) + } + return t.withSchemaWithExecutor(ctx, execer) +} + +// materializeObjectTable is used by the PostgreSQL handler, which has already +// executed its parser-normalized CREATE TABLE against the local shadow. The +// local catalog is authoritative for the accepted DuckDB column types; the +// resulting lake relation deliberately omits constraints unsupported by +// DuckLake. +func (d *Database) materializeObjectTable(ctx *sql.Context, name string) error { + execer, conn, _, release, err := adapter.GetCatalogExecutionSnapshotWithLease(ctx) + if err != nil { + return err + } + defer release() + return d.materializeObjectTableWithExecutor(ctx, name, execer, conn) +} + +// materializeObjectTableWithExecutor is the executor-affine form used by +// metadata publication. The caller owns the snapshot (and, when present, its +// transaction), so the physical CREATE cannot silently move to another pool +// connection between the local-shadow lookup and the durable marker update. +func (d *Database) materializeObjectTableWithExecutor( + ctx *sql.Context, + name string, + execer adapter.SQLExecutor, + conn *stdsql.Conn, +) error { + if d.provider == nil || d.provider.duckLake == nil { + return fmt.Errorf("%w: DuckLake service configuration is disabled", ErrInvalidTableStorage) + } + if execer == nil { + return fmt.Errorf("object-table metadata executor is unavailable") + } + if conn == nil { + return fmt.Errorf("object-table metadata connection is unavailable") + } + if d.provider.Storage() == nil { + if err := d.provider.EnsureDuckLakeConnectionWithExecutor(ctx, execer, conn); err != nil { + return err + } + } + defs, err := d.localObjectColumnsWithExecutor(ctx, name, execer) + if err != nil { + return err + } + if len(defs) == 0 { + return fmt.Errorf("%w: object table has no columns", ErrInvalidTableStorage) + } + if err := d.rejectLocalObjectConstraintsWithExecutor(ctx, name, execer); err != nil { + return err + } + physical := d.objectPhysicalTableName(name) + parts := make([]string, 0, len(defs)) + for _, def := range defs { + if !validObjectTypeText(def.dataType) || strings.Contains(strings.ToLower(def.defaultValue), "nextval(") { + return fmt.Errorf("%w: unsupported column definition", ErrInvalidTableStorage) + } + part := QuoteIdentifierANSI(def.name) + " " + def.dataType + if def.nullable { + part += " NULL" + } else { + part += " NOT NULL" + } + if strings.TrimSpace(def.defaultValue) != "" { + part += " DEFAULT " + def.defaultValue + } + parts = append(parts, part) + } + createSQL := "CREATE TABLE " + physical + " (" + strings.Join(parts, ", ") + ")" + lakeExecer := execer + lakeCtx := context.Context(ctx) + var lakeConn *stdsql.Conn + var lakeDDLCtx context.Context + if d.provider.Storage() != nil { + // PostgreSQL already wrote the local shadow in the session transaction. + // DuckDB cannot also write the attached lake catalog in that transaction. + lakeDDLCtx, lakeConn, err = d.acquireDedicatedLakeConn() + if err != nil { + return err + } + defer lakeConn.Close() + lakeExecer = lakeConn + lakeCtx = lakeDDLCtx + } + if err := d.ensureObjectSchema(lakeCtx, lakeExecer); err != nil { + return err + } + if _, err := lakeExecer.ExecContext(lakeCtx, createSQL); err != nil { + if IsDuckDBTableAlreadyExistsError(err) { + return nil + } + return ErrDuckDB.New(err) + } + if lakeConn != nil { + d.registerUncommittedLakeRelation(ctx, physical) + } + return nil +} + +// MaterializeObjectTable exposes the PostgreSQL parser bridge without +// exposing the internal shadow-table representation. +func (d *Database) MaterializeObjectTable(ctx *sql.Context, name string) error { + return d.materializeObjectTable(ctx, name) +} + +func (d *Database) localObjectColumns(ctx *sql.Context, name string) ([]objectColumnDefinition, error) { + execer, _, _, release, err := adapter.GetCatalogExecutionSnapshotWithLease(ctx) + if err != nil { + return nil, err + } + defer release() + return d.localObjectColumnsWithExecutor(ctx, name, execer) +} + +func (d *Database) localObjectColumnsWithExecutor(ctx *sql.Context, name string, execer adapter.SQLExecutor) ([]objectColumnDefinition, error) { + if execer == nil { + return nil, fmt.Errorf("object-table metadata executor is unavailable") + } + rows, err := execer.QueryContext(ctx, ` + SELECT column_name, data_type, is_nullable, column_default + FROM duckdb_columns() + WHERE database_name = ? AND schema_name = ? AND table_name = ? + ORDER BY column_index + `, d.catalog, d.name, name) + if err != nil { + return nil, err + } + defer rows.Close() + var defs []objectColumnDefinition + for rows.Next() { + var def objectColumnDefinition + var nullable bool + var defaultValue stdsql.NullString + if err := rows.Scan(&def.name, &def.dataType, &nullable, &defaultValue); err != nil { + return nil, err + } + def.nullable = nullable + if defaultValue.Valid { + def.defaultValue = defaultValue.String + } + defs = append(defs, def) + } + if err := rows.Err(); err != nil { + return nil, err + } + return defs, nil +} + +func (d *Database) rejectLocalObjectConstraints(ctx *sql.Context, name string) error { + execer, _, _, release, err := adapter.GetCatalogExecutionSnapshotWithLease(ctx) + if err != nil { + return err + } + defer release() + return d.rejectLocalObjectConstraintsWithExecutor(ctx, name, execer) +} + +func (d *Database) rejectLocalObjectConstraintsWithExecutor(ctx *sql.Context, name string, execer adapter.SQLExecutor) error { + if execer == nil { + return fmt.Errorf("object-table metadata executor is unavailable") + } + rows, err := execer.QueryContext(ctx, ` + SELECT constraint_type + FROM duckdb_constraints() + WHERE database_name = ? AND schema_name = ? AND table_name = ? + `, d.catalog, d.name, name) + if err != nil { + return err + } + defer rows.Close() + for rows.Next() { + var kind string + if err := rows.Scan(&kind); err != nil { + return err + } + if strings.EqualFold(kind, "PRIMARY KEY") || strings.EqualFold(kind, "UNIQUE") || strings.EqualFold(kind, "CHECK") || strings.EqualFold(kind, "FOREIGN KEY") { + return fmt.Errorf("%w: object tables do not support %s constraints", ErrInvalidTableStorage, kind) + } + } + return rows.Err() +} + +func validObjectTypeText(value string) bool { + value = strings.TrimSpace(value) + if value == "" || strings.ContainsAny(value, "\x00;\r\n") { + return false + } + return true +} + +// PhysicalTableNameForTable is used by protocol COPY helpers and by callers +// that already resolved a GMS table. Local tables return their normal fully +// qualified name; object tables return the fixed lake relation. +func (prov *DatabaseProvider) PhysicalTableNameForTable(ctx *sql.Context, table sql.Table) (string, bool, error) { + execer, conn, _, release, err := adapter.GetExecutionSnapshotWithLease(ctx) + if err != nil { + return "", false, err + } + defer release() + return prov.PhysicalTableNameForTableWithExecutor(ctx, table, execer, conn) +} + +// PhysicalTableNameForTableWithExecutor resolves a table through the supplied +// execution snapshot. Metadata lookup and the eventual statement execution +// therefore remain on one transaction/connection pair. +func (prov *DatabaseProvider) PhysicalTableNameForTableWithExecutor( + ctx *sql.Context, + table sql.Table, + execer adapter.SQLExecutor, + conn *stdsql.Conn, +) (string, bool, error) { + var t *Table + switch typed := table.(type) { + case *Table: + t = typed + case *IndexedTable: + if typed != nil { + t = typed.Table + } + } + if t == nil { + return "", false, nil + } + info := t.ExtraTableInfo() + if info.StorageKind() != TableStorageObject { + return FullTableName(t.db.catalog, t.db.name, t.name), false, nil + } + if prov == nil || prov.duckLake == nil { + return "", false, fmt.Errorf("%w: DuckLake service configuration is disabled", ErrInvalidTableStorage) + } + if execer == nil { + return "", false, fmt.Errorf("object-table metadata executor is unavailable") + } + if conn != nil { + if err := prov.EnsureDuckLakeConnectionWithExecutor(ctx, execer, conn); err != nil { + return "", false, err + } + } + physical, found, err := prov.ObjectTableNameWithExecutor(ctx, execer, t.db.catalog, t.db.name, t.name) + if err != nil || !found { + return "", false, err + } + return physical, true, nil +} diff --git a/catalog/operation_barrier_test.go b/catalog/operation_barrier_test.go new file mode 100644 index 00000000..38adcc51 --- /dev/null +++ b/catalog/operation_barrier_test.go @@ -0,0 +1,825 @@ +package catalog + +import ( + "context" + stdsql "database/sql" + "errors" + "io" + "sync" + "testing" + "time" + + "github.com/apecloud/myduckserver/configuration" + "github.com/apecloud/myduckserver/mycontext" + "github.com/dolthub/go-mysql-server/sql" + "github.com/duckdb/duckdb-go/v2" + "github.com/stretchr/testify/require" +) + +func newDuckLakeOperationTestProvider() *DatabaseProvider { + return &DatabaseProvider{duckLake: &duckLakeRuntime{config: configuration.DuckLakeConfig{ + MetadataPath: "/tmp/task77/catalog.ducklake", + DataPath: "/tmp/task77/data", + }}} +} + +func duckLakeOperationCount(provider *DatabaseProvider) int { + provider.duckLakeTxnMu.Lock() + defer provider.duckLakeTxnMu.Unlock() + return provider.duckLakeActiveOperations +} + +func TestDuckLakeOperationAdmissionBlocksCleanup(t *testing.T) { + provider := newDuckLakeOperationTestProvider() + release := provider.BeginDuckLakeOperation(mycontext.WithFrontendQuery(context.Background())) + require.Equal(t, 1, duckLakeOperationCount(provider)) + + cleanupDone := make(chan struct{}) + go func() { + cleanupRelease := provider.beginDuckLakeCleanup() + cleanupRelease() + close(cleanupDone) + }() + + select { + case <-cleanupDone: + t.Fatal("cleanup entered while a logical operation was active") + case <-time.After(25 * time.Millisecond): + } + + release() + release() + select { + case <-cleanupDone: + case <-time.After(time.Second): + t.Fatal("cleanup did not resume after logical operation release") + } + require.Zero(t, duckLakeOperationCount(provider)) +} + +func TestDuckLakeRollbackCleanupExcludesOwnedOperationLease(t *testing.T) { + provider := newDuckLakeOperationTestProvider() + owner := new(int) + ctx := mycontext.WithFrontendQuery(context.Background()) + operationRelease := provider.BeginDuckLakeOperationForOwner(ctx, owner) + defer operationRelease() + + reservationDone := make(chan func(), 1) + go func() { + reservationDone <- provider.BeginDuckLakeRollbackCleanupForOwner(nil, owner) + }() + waitForDuckLakeCleanupActive(t, provider) + + var reservationRelease func() + select { + case reservationRelease = <-reservationDone: + case <-time.After(time.Second): + t.Fatal("rollback cleanup did not exempt the caller's operation lease") + } + defer reservationRelease() + + // The operation remains admitted while its owner holds the reservation; the + // exemption applies only to this caller, not by dropping the global count. + require.Equal(t, 1, duckLakeOperationCount(provider)) + operationRelease() + require.Zero(t, duckLakeOperationCount(provider)) +} + +func TestDuckLakeRollbackCleanupBlocksOtherOperationOwner(t *testing.T) { + provider := newDuckLakeOperationTestProvider() + ownerA := new(int) + ownerB := new(int) + ctx := mycontext.WithFrontendQuery(context.Background()) + operationRelease := provider.BeginDuckLakeOperationForOwner(ctx, ownerA) + defer operationRelease() + + reservationDone := make(chan func(), 1) + go func() { + reservationDone <- provider.BeginDuckLakeRollbackCleanupForOwner(nil, ownerB) + }() + waitForDuckLakeCleanupActive(t, provider) + select { + case release := <-reservationDone: + release() + t.Fatal("rollback cleanup entered while another owner's operation was active") + case <-time.After(50 * time.Millisecond): + } + + operationRelease() + var reservationRelease func() + select { + case reservationRelease = <-reservationDone: + case <-time.After(time.Second): + t.Fatal("rollback cleanup did not resume after the other owner released") + } + reservationRelease() +} + +func TestDuckLakeRollbackCleanupBlocksAnonymousOperationLease(t *testing.T) { + provider := newDuckLakeOperationTestProvider() + ctx := mycontext.WithFrontendQuery(context.Background()) + // A plain context has no logical GMS transaction identity, so the legacy + // entry point must retain a conservative anonymous lease. + operationRelease := provider.BeginDuckLakeOperation(ctx) + defer operationRelease() + owner := new(int) + + reservationDone := make(chan func(), 1) + go func() { + reservationDone <- provider.BeginDuckLakeRollbackCleanupForOwner(nil, owner) + }() + waitForDuckLakeCleanupActive(t, provider) + select { + case release := <-reservationDone: + release() + t.Fatal("rollback cleanup entered while an anonymous operation was active") + case <-time.After(50 * time.Millisecond): + } + + operationRelease() + select { + case release := <-reservationDone: + release() + case <-time.After(time.Second): + t.Fatal("rollback cleanup did not resume after anonymous operation release") + } +} + +func TestDuckLakeRollbackCleanupMatchesPhysicalOperationOwner(t *testing.T) { + provider := newDuckLakeOperationTestProvider() + target := new(stdsql.Tx) + ctx := mycontext.WithFrontendQuery(context.Background()) + operationRelease := provider.BeginDuckLakeOperationForOwner(ctx, target) + defer operationRelease() + + reservationDone := make(chan func(), 1) + go func() { + reservationDone <- provider.BeginDuckLakeRollbackCleanupForOwner(target, nil) + }() + waitForDuckLakeCleanupActive(t, provider) + select { + case release := <-reservationDone: + release() + case <-time.After(time.Second): + t.Fatal("rollback cleanup did not match the physical operation owner") + } +} + +func TestDuckLakeRollbackCleanupTreatsNonComparableOwnerAsAnonymous(t *testing.T) { + provider := newDuckLakeOperationTestProvider() + ctx := mycontext.WithFrontendQuery(context.Background()) + operationOwner := []string{"owner"} + operationRelease := provider.BeginDuckLakeOperationForOwner(ctx, operationOwner) + defer operationRelease() + + reservationDone := make(chan func(), 1) + go func() { + reservationDone <- provider.BeginDuckLakeRollbackCleanupForOwner(nil, operationOwner) + }() + waitForDuckLakeCleanupActive(t, provider) + select { + case release := <-reservationDone: + release() + t.Fatal("rollback cleanup exempted a non-comparable owner") + case <-time.After(50 * time.Millisecond): + } + + operationRelease() + select { + case release := <-reservationDone: + release() + case <-time.After(time.Second): + t.Fatal("rollback cleanup did not resume after non-comparable lease release") + } +} + +func TestDuckLakeOperationAdmissionBlocksDuringCleanup(t *testing.T) { + provider := newDuckLakeOperationTestProvider() + cleanupRelease := provider.beginDuckLakeCleanup() + + operationDone := make(chan struct{}) + go func() { + release := provider.BeginDuckLakeOperation(mycontext.WithFrontendQuery(context.Background())) + release() + close(operationDone) + }() + + select { + case <-operationDone: + t.Fatal("logical operation entered during cleanup") + case <-time.After(25 * time.Millisecond): + } + + cleanupRelease() + select { + case <-operationDone: + case <-time.After(time.Second): + t.Fatal("logical operation did not resume after cleanup release") + } + require.Zero(t, duckLakeOperationCount(provider)) +} + +func TestResetDuckLakeTransactionsDrainsLogicalOperationsBeforeReset(t *testing.T) { + provider := newDuckLakeOperationTestProvider() + oldRelease := provider.BeginDuckLakeOperation(mycontext.WithFrontendQuery(context.Background())) + + resetDone := make(chan struct{}) + go func() { + provider.resetDuckLakeTransactions() + close(resetDone) + }() + + // Wait until reset has announced the generation transition. It must retain + // the old lease instead of clearing its count underneath the iterator. + deadline := time.After(time.Second) + for { + provider.duckLakeTxnMu.Lock() + resetting := provider.duckLakeCleanupActive + provider.duckLakeTxnMu.Unlock() + if resetting { + break + } + select { + case <-deadline: + t.Fatal("reset did not announce a cleanup barrier") + default: + time.Sleep(time.Millisecond) + } + } + + newDone := make(chan func(), 1) + go func() { + newDone <- provider.BeginDuckLakeOperation(mycontext.WithFrontendQuery(context.Background())) + }() + select { + case <-newDone: + t.Fatal("new logical operation entered while reset was draining") + case <-time.After(25 * time.Millisecond): + } + + oldRelease() + select { + case <-resetDone: + case <-time.After(time.Second): + t.Fatal("reset did not finish after the old logical operation released") + } + + var newRelease func() + select { + case newRelease = <-newDone: + case <-time.After(time.Second): + t.Fatal("new logical operation did not resume after reset") + } + require.Equal(t, 1, duckLakeOperationCount(provider)) + newRelease() + require.Zero(t, duckLakeOperationCount(provider)) +} + +func TestDuckLakeGenerationTransitionDoesNotDeadlockProviderOperation(t *testing.T) { + provider := newDuckLakeOperationTestProvider() + provider.mu = &sync.RWMutex{} + operationRelease := provider.BeginDuckLakeOperation(mycontext.WithFrontendQuery(context.Background())) + + transitionDone := make(chan struct{}) + go func() { + release := provider.beginDuckLakeGenerationTransition() + release() + close(transitionDone) + }() + + // Wait until the transition has announced cleanup. The admitted operation + // must still be able to acquire the provider mutex and finish; the transition + // itself deliberately does not take that mutex while draining operations. + deadline := time.After(time.Second) + for { + provider.duckLakeTxnMu.Lock() + active := provider.duckLakeCleanupActive + provider.duckLakeTxnMu.Unlock() + if active { + break + } + select { + case <-deadline: + t.Fatal("generation transition did not announce cleanup") + default: + time.Sleep(time.Millisecond) + } + } + + operationDone := make(chan struct{}) + go func() { + provider.mu.Lock() + provider.mu.Unlock() + operationRelease() + close(operationDone) + }() + select { + case <-operationDone: + case <-time.After(time.Second): + t.Fatal("admitted operation could not finish while transition drained") + } + select { + case <-transitionDone: + case <-time.After(time.Second): + t.Fatal("generation transition did not finish after operation release") + } +} + +func TestProviderCloseLetsAdmittedOperationReleaseBeforePoolTeardown(t *testing.T) { + connector, err := duckdb.NewConnector("", nil) + require.NoError(t, err) + db := stdsql.OpenDB(connector) + p := NewConnectionPool(connector, db, "memory") + provider := &DatabaseProvider{ + mu: &sync.RWMutex{}, + pool: p, + duckLake: &duckLakeRuntime{config: configuration.DuckLakeConfig{ + MetadataPath: "/tmp/task77/catalog.ducklake", + DataPath: "/tmp/task77/data", + }}, + connector: connector, + } + p.SetTransactionLifecycleHooks(provider.beginDuckLakeTransaction, provider.untrackDuckLakeTransaction) + + operationRelease := provider.BeginDuckLakeOperation(mycontext.WithFrontendQuery(context.Background())) + closeDone := make(chan error, 1) + go func() { closeDone <- provider.Close() }() + + deadline := time.After(time.Second) + for { + provider.duckLakeTxnMu.Lock() + active := provider.duckLakeCleanupActive + provider.duckLakeTxnMu.Unlock() + if active { + break + } + select { + case <-deadline: + t.Fatal("provider close did not announce the generation transition") + default: + time.Sleep(time.Millisecond) + } + } + + operationDone := make(chan struct{}) + go func() { + // This models an admitted frontend operation that needs the provider + // mutex before it can release its operation lease. + provider.mu.Lock() + provider.mu.Unlock() + operationRelease() + close(operationDone) + }() + select { + case <-operationDone: + case <-time.After(time.Second): + t.Fatal("provider close held the mutex ahead of an admitted operation") + } + select { + case closeErr := <-closeDone: + require.NoError(t, closeErr) + case <-time.After(time.Second): + t.Fatal("provider close did not finish after operation release") + } + _ = connector.Close() +} + +func TestDuckLakeProviderCloseAllowsPhysicalOwnerLeaseToReleaseDuringPoolTeardown(t *testing.T) { + provider, closeProvider := newDuckLakeOperationPoolProvider(t) + const sessionID uint32 = 907 + + tx, err := provider.Pool().GetTxn(context.Background(), sessionID, "", nil) + require.NoError(t, err) + require.NotNil(t, tx) + + // COPY keeps its provider operation lease until the surrounding physical + // transaction finishes. The pool completion callback below models that + // production registration and must be allowed to run from provider.Close. + operationRelease := provider.BeginDuckLakeOperationForOwner( + mycontext.WithFrontendQuery(context.Background()), tx, + ) + callbackDone := make(chan struct{}) + var callbackOnce sync.Once + require.True(t, provider.Pool().RegisterTransactionCompletion(tx, func(bool) { + operationRelease() + callbackOnce.Do(func() { close(callbackDone) }) + })) + + closeDone := make(chan error, 1) + go func() { closeDone <- closeProvider() }() + + select { + case closeErr := <-closeDone: + require.NoError(t, closeErr) + case <-time.After(time.Second): + // Keep the test goroutine recoverable on an implementation that regresses + // to waiting for every operation before pool teardown. Releasing here is + // only cleanup; the timeout remains the assertion failure. + operationRelease() + select { + case <-closeDone: + case <-time.After(time.Second): + } + t.Fatal("provider close waited for a physical-owner operation lease before pool teardown") + } + + select { + case <-callbackDone: + case <-time.After(time.Second): + t.Fatal("pool teardown did not invoke the physical transaction completion callback") + } + require.Zero(t, duckLakeOperationCount(provider)) +} + +// newDuckLakeOperationPoolProvider creates the same provider/pool pairing used +// by a live configured provider, while leaving extension setup out of this +// lifecycle-only test. The non-empty service paths enable the provider's +// admission barriers; no DuckLake SQL is issued. +func newDuckLakeOperationPoolProvider(t *testing.T) (*DatabaseProvider, func() error) { + t.Helper() + connector, err := duckdb.NewConnector("", nil) + require.NoError(t, err) + db := stdsql.OpenDB(connector) + pool := NewConnectionPool(connector, db, "memory") + provider := &DatabaseProvider{ + mu: &sync.RWMutex{}, + pool: pool, + storage: db, + connector: connector, + duckLake: &duckLakeRuntime{config: configuration.DuckLakeConfig{ + MetadataPath: "/tmp/task77/catalog.ducklake", + DataPath: "/tmp/task77/data", + }}, + } + pool.SetTransactionLifecycleHooks(provider.beginDuckLakeTransaction, provider.untrackDuckLakeTransaction) + + var closeOnce sync.Once + var closeErr error + closeProvider := func() error { + closeOnce.Do(func() { + closeErr = provider.Close() + }) + return closeErr + } + t.Cleanup(func() { _ = closeProvider() }) + return provider, closeProvider +} + +type providerOperationChild struct { + events chan<- string + nextErr error + closeStarted chan struct{} + allowClose <-chan struct{} + closeOnce sync.Once +} + +func (c *providerOperationChild) Next(*sql.Context) (sql.Row, error) { + c.events <- "child-next" + return nil, c.nextErr +} + +func (c *providerOperationChild) Close(*sql.Context) error { + c.closeOnce.Do(func() { + c.events <- "child-close-start" + close(c.closeStarted) + <-c.allowClose + c.events <- "child-close-end" + }) + return nil +} + +type providerTerminalMode string + +const ( + providerTerminalEOF providerTerminalMode = "eof" + providerTerminalError providerTerminalMode = "error" + providerEarlyClose providerTerminalMode = "early-close" +) + +// runProviderChildTerminal is the small integration harness for the +// production iterator contract: a terminal child result is closed before the +// provider operation lease is released. The concrete wrapper's EOF/error and +// early-close behavior is covered in backend/ducklake_operation_test.go; this +// harness supplies a real provider and pool so generation finalization is +// exercised in the same ordering. +func runProviderChildTerminal( + ctx *sql.Context, + child *providerOperationChild, + mode providerTerminalMode, + release func(), +) error { + var terminalErr error + if mode == providerEarlyClose { + terminalErr = errProviderEarlyClose + } else { + _, terminalErr = child.Next(ctx) + } + terminalErr = errors.Join(terminalErr, child.Close(ctx)) + release() + return terminalErr +} + +func waitForDuckLakeCleanupActive(t *testing.T, provider *DatabaseProvider) { + t.Helper() + deadline := time.After(time.Second) + for { + provider.duckLakeTxnMu.Lock() + active := provider.duckLakeCleanupActive + provider.duckLakeTxnMu.Unlock() + if active { + return + } + select { + case <-deadline: + t.Fatal("provider did not announce the cleanup barrier") + default: + time.Sleep(time.Millisecond) + } + } +} + +func TestDuckLakeProviderCloseWaitsForChildCloseBeforeLeaseRelease(t *testing.T) { + cases := []struct { + name string + mode providerTerminalMode + nextErr error + expectErr error + }{ + {name: "eof", mode: providerTerminalEOF, nextErr: io.EOF, expectErr: io.EOF}, + {name: "child-error", mode: providerTerminalError, nextErr: errOperationProviderChild, expectErr: errOperationProviderChild}, + {name: "early-close", mode: providerEarlyClose, expectErr: errProviderEarlyClose}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + provider, closeProvider := newDuckLakeOperationPoolProvider(t) + events := make(chan string, 32) + allowClose := make(chan struct{}) + var allowCloseOnce sync.Once + openChildClose := func() { + allowCloseOnce.Do(func() { close(allowClose) }) + } + child := &providerOperationChild{ + events: events, + nextErr: tc.nextErr, + closeStarted: make(chan struct{}), + allowClose: allowClose, + } + operationRelease := provider.BeginDuckLakeOperation(mycontext.WithFrontendQuery(context.Background())) + t.Cleanup(func() { + openChildClose() + operationRelease() + }) + release := func() { + events <- "operation-release" + operationRelease() + } + + ctx := sql.NewContext(mycontext.WithFrontendQuery(context.Background())) + terminalDone := make(chan error, 1) + go func() { + terminalDone <- runProviderChildTerminal(ctx, child, tc.mode, release) + }() + select { + case <-child.closeStarted: + case <-time.After(time.Second): + t.Fatal("child close did not start") + } + + closeDone := make(chan error, 1) + go func() { + closeErr := closeProvider() + events <- "provider-close" + closeDone <- closeErr + }() + waitForDuckLakeCleanupActive(t, provider) + + select { + case closeErr := <-closeDone: + t.Fatalf("provider closed while child close was blocked: %v", closeErr) + case <-time.After(50 * time.Millisecond): + } + select { + case terminalErr := <-terminalDone: + t.Fatalf("terminal path completed before child close was released: %v", terminalErr) + case <-time.After(25 * time.Millisecond): + } + + openChildClose() + var terminalErr error + select { + case terminalErr = <-terminalDone: + case <-time.After(time.Second): + t.Fatal("terminal path did not finish after child close release") + } + require.ErrorIs(t, terminalErr, tc.expectErr) + + select { + case closeErr := <-closeDone: + require.NoError(t, closeErr) + case <-time.After(time.Second): + t.Fatal("provider close did not finish after operation lease release") + } + + got := make([]string, 0, 5) + for { + select { + case event := <-events: + got = append(got, event) + default: + goto drained + } + } + drained: + closeIndex := indexProviderEvent(got, "child-close-end") + releaseIndex := indexProviderEvent(got, "operation-release") + providerIndex := indexProviderEvent(got, "provider-close") + require.GreaterOrEqual(t, closeIndex, 0) + require.Greater(t, releaseIndex, closeIndex) + require.Greater(t, providerIndex, releaseIndex) + }) + } +} + +var ( + errOperationProviderChild = errors.New("provider child failed") + errProviderEarlyClose = errors.New("provider child closed early") +) + +func indexProviderEvent(events []string, want string) int { + for i, event := range events { + if event == want { + return i + } + } + return -1 +} + +func TestDuckLakeRollbackCleanupWaitsForChildCloseAndLeaseRelease(t *testing.T) { + provider, closeProvider := newDuckLakeOperationPoolProvider(t) + const sessionID uint32 = 881 + tx, err := provider.Pool().GetTxn(context.Background(), sessionID, "", nil) + require.NoError(t, err) + require.NotNil(t, tx) + + events := make(chan string, 32) + allowClose := make(chan struct{}) + var allowCloseOnce sync.Once + openChildClose := func() { + allowCloseOnce.Do(func() { close(allowClose) }) + } + child := &providerOperationChild{ + events: events, + nextErr: io.EOF, + closeStarted: make(chan struct{}), + allowClose: allowClose, + } + operationRelease := provider.BeginDuckLakeOperation(mycontext.WithFrontendQuery(context.Background())) + reservationDone := make(chan func(), 1) + var reservationRelease func() + t.Cleanup(func() { + openChildClose() + operationRelease() + if reservationRelease != nil { + reservationRelease() + return + } + select { + case release := <-reservationDone: + release() + case <-time.After(time.Second): + } + }) + release := func() { + events <- "operation-release" + operationRelease() + } + + go func() { + reservationDone <- provider.BeginDuckLakeRollbackCleanup(tx) + }() + waitForDuckLakeCleanupActive(t, provider) + + ctx := sql.NewContext(mycontext.WithFrontendQuery(context.Background())) + terminalDone := make(chan error, 1) + go func() { + terminalDone <- runProviderChildTerminal(ctx, child, providerTerminalEOF, release) + }() + select { + case <-child.closeStarted: + case <-time.After(time.Second): + t.Fatal("child close did not start") + } + select { + case <-reservationDone: + t.Fatal("rollback cleanup reservation entered while child close was blocked") + case <-time.After(50 * time.Millisecond): + } + + openChildClose() + select { + case terminalErr := <-terminalDone: + require.ErrorIs(t, terminalErr, io.EOF) + case <-time.After(time.Second): + t.Fatal("terminal path did not finish") + } + + select { + case reservationRelease = <-reservationDone: + events <- "rollback-reservation" + case <-time.After(time.Second): + t.Fatal("rollback cleanup reservation did not follow operation release") + } + + inactive, rollbackErr := provider.Pool().RollbackTxnWithCleanup(sessionID, tx, func(*stdsql.Conn) error { + events <- "rollback-cleanup" + return nil + }) + require.True(t, inactive) + require.NoError(t, rollbackErr) + events <- "rollback-finished" + reservationRelease() + require.NoError(t, closeProvider()) + + got := make([]string, 0, 8) + for { + select { + case event := <-events: + got = append(got, event) + default: + goto drainedRollback + } + } + +drainedRollback: + closeIndex := indexProviderEvent(got, "child-close-end") + releaseIndex := indexProviderEvent(got, "operation-release") + reservationIndex := indexProviderEvent(got, "rollback-reservation") + cleanupIndex := indexProviderEvent(got, "rollback-cleanup") + require.GreaterOrEqual(t, closeIndex, 0) + require.Greater(t, releaseIndex, closeIndex) + require.Greater(t, reservationIndex, releaseIndex) + require.Greater(t, cleanupIndex, reservationIndex) +} + +func TestDuckLakeGenerationTransitionKeepsPoolWriterOutOfAdmissionWait(t *testing.T) { + connector, err := duckdb.NewConnector("", nil) + require.NoError(t, err) + db := stdsql.OpenDB(connector) + p := NewConnectionPool(connector, db, "memory") + provider := &DatabaseProvider{duckLake: &duckLakeRuntime{config: configuration.DuckLakeConfig{ + MetadataPath: "/tmp/task77/catalog.ducklake", + DataPath: "/tmp/task77/data", + }}} + p.SetTransactionLifecycleHooks(provider.beginDuckLakeTransaction, provider.untrackDuckLakeTransaction) + + releaseTransition := provider.beginDuckLakeGenerationTransition() + t.Cleanup(func() { + releaseTransition() + _ = p.Close() + _ = connector.Close() + }) + + txnDone := make(chan error, 1) + go func() { + _, getErr := p.GetTxn(context.Background(), 771, "", nil) + txnDone <- getErr + }() + + // Pool shutdown must be able to acquire its writer even though GetTxn is + // waiting for cleanupActive. The admission wait happens before RLock. + closeDone := make(chan error, 1) + go func() { closeDone <- p.Close() }() + select { + case closeErr := <-closeDone: + require.NoError(t, closeErr) + case <-time.After(time.Second): + t.Fatal("pool close blocked behind a transaction admission wait") + } + + releaseTransition() + select { + case getErr := <-txnDone: + require.Error(t, getErr) + case <-time.After(time.Second): + t.Fatal("transaction admission did not resume after transition release") + } +} + +func TestDuckLakeOperationAdmissionRequiresFrontendAndCompleteConfiguration(t *testing.T) { + provider := newDuckLakeOperationTestProvider() + for _, ctx := range []context.Context{ + context.Background(), + mycontext.WithMaintenanceQuery(context.Background()), + mycontext.WithQueryOrigin(context.Background(), mycontext.MySQLReplicationQueryOrigin), + } { + release := provider.BeginDuckLakeOperation(ctx) + release() + } + require.Zero(t, duckLakeOperationCount(provider)) + + extensionOnly := &DatabaseProvider{duckLake: &duckLakeRuntime{}} + release := extensionOnly.BeginDuckLakeOperation(mycontext.WithFrontendQuery(context.Background())) + release() + require.Zero(t, duckLakeOperationCount(extensionOnly)) +} diff --git a/catalog/provider.go b/catalog/provider.go index 9a35c781..b494861f 100644 --- a/catalog/provider.go +++ b/catalog/provider.go @@ -3,9 +3,11 @@ package catalog import ( "context" "database/sql/driver" + "errors" "fmt" "os" "path/filepath" + "reflect" "sort" "strings" "sync" @@ -24,18 +26,118 @@ import ( ) type DatabaseProvider struct { - mu *sync.RWMutex - defaultTimeZone string - connector *duckdb.Connector - storage *stdsql.DB - pool *ConnectionPool - defaultCatalogName string // default database name in postgres - dataDir string - dbFile string - dsn string - externalProcedureRegistry sql.ExternalStoredProcedureRegistry - duckLake *duckLakeRuntime - ready bool + mu *sync.RWMutex + duckLakeCleanupMu sync.Mutex + duckLakeTxnMu sync.Mutex + duckLakeTxnCond *sync.Cond + duckLakeActiveTx map[*stdsql.Tx]struct{} + duckLakePendingTx int + duckLakeActiveOperations int + duckLakeOperationOwners map[any]int + duckLakeOperationLeases map[uint64]duckLakeOperationLease + duckLakeAnonymousOps int + duckLakeNextOperationID uint64 + duckLakeCleanupActive bool + // Lake DDL is committed on a dedicated connection because DuckDB refuses + // writes to a second attached database in the same transaction. Relations + // created that way stay invisible until local catalog publication commits; + // this list is the rollback/publication-failure compensation set. Guarded + // by duckLakeTxnMu. + duckLakeUncommitted []uncommittedLakeRelation + // A failed pool teardown can leave an active physical owner behind when its + // transaction marker is malformed. Preserve that evidence, but do not let a + // generation transition wait forever or reopen admission against the retired + // storage. The fields are guarded by duckLakeTxnMu. + duckLakeGenerationPoisoned bool + duckLakeGenerationErr error + defaultTimeZone string + connector *duckdb.Connector + storage *stdsql.DB + pool *ConnectionPool + defaultCatalogName string // default database name in postgres + dataDir string + dbFile string + dsn string + externalProcedureRegistry sql.ExternalStoredProcedureRegistry + duckLake *duckLakeRuntime + ready bool +} + +// duckLakeOperationLease records the owner identity captured at admission. +// Comparable identities are indexed in duckLakeOperationOwners; anonymous or +// non-comparable owners are kept as blocking leases in their own count. +type duckLakeOperationLease struct { + owner any + owned bool +} + +// uncommittedLakeRelation is a DuckLake table created outside the session +// transaction. Rollback cleanup must DROP it; a successful commit forgets it. +type uncommittedLakeRelation struct { + tx *stdsql.Tx + conn *stdsql.Conn + physical string +} + +// rollbackCleanupOwnerKey carries the physical connection reserved by a +// rollback finalizer. It is intentionally private; callers can only construct +// it through WithRollbackCleanupOwner. +type rollbackCleanupOwnerKey struct{} + +type rollbackCleanupOwner struct { + conn *stdsql.Conn +} + +// duckLakeCleanupLeaseKey marks a maintenance call that already owns the +// provider cleanup barrier. It is private so a client SQL context cannot +// manufacture the exemption; only the rollback finalizer creates it. +type duckLakeCleanupLeaseKey struct{} + +// WithRollbackCleanupOwner marks a maintenance context whose connection is +// reserved by the session rollback finalizer. Provider setup and cleanup can +// then avoid re-reading the session transaction map while that map's lock is +// held by the finalizer. +func WithRollbackCleanupOwner(ctx context.Context, conn *stdsql.Conn) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, rollbackCleanupOwnerKey{}, rollbackCleanupOwner{conn: conn}) +} + +func rollbackCleanupOwnerConn(ctx context.Context) *stdsql.Conn { + if ctx == nil { + return nil + } + owner, ok := ctx.Value(rollbackCleanupOwnerKey{}).(rollbackCleanupOwner) + if !ok { + return nil + } + return owner.conn +} + +// RollbackCleanupOwner returns the physical connection reserved by a rollback +// finalizer, if the context was marked with WithRollbackCleanupOwner. +func RollbackCleanupOwner(ctx context.Context) *stdsql.Conn { + return rollbackCleanupOwnerConn(ctx) +} + +// WithDuckLakeCleanupLease marks a context used by a rollback finalizer that +// already acquired the provider cleanup barrier before entering the pool +// lifecycle lock. CleanupDuckLakeOrphansOnConn must not acquire that barrier a +// second time in this scope. +func WithDuckLakeCleanupLease(ctx context.Context) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, duckLakeCleanupLeaseKey{}, struct{}{}) +} + +func duckLakeCleanupLeaseHeld(ctx context.Context) bool { + if ctx == nil { + return false + } + _, ok := ctx.Value(duckLakeCleanupLeaseKey{}).(struct{}) + return ok } var _ sql.DatabaseProvider = (*DatabaseProvider)(nil) @@ -121,6 +223,12 @@ func (prov *DatabaseProvider) openStorage(readOnly bool) error { } func (prov *DatabaseProvider) openStorageWithContext(readOnly bool, openCtx context.Context) error { + // A failed teardown leaves the generation permanently unusable. Do not let + // Restart (or another internal open path) silently create a fresh connector + // while the provider still retains the malformed/unknown owner state. + if err := prov.duckLakeGenerationPoisonError(); err != nil { + return err + } connectorDSN := prov.dsn if readOnly && prov.dsn == "" { connectorDSN = "?access_mode=read_only" @@ -165,6 +273,10 @@ func (prov *DatabaseProvider) openStorageWithContext(readOnly bool, openCtx cont if prov.pool == nil { prov.pool = NewConnectionPool(connector, storage, prov.defaultCatalogName) } else if err := prov.pool.Reset(connector, storage); err != nil { + // Reset may fail while preserving a malformed transaction/owner mapping. + // Poison before returning so a subsequent Restart cannot repeatedly retry + // the same unsafe generation or admit work through a stale pool. + prov.poisonDuckLakeGeneration(err) _ = storage.Close() _ = connector.Close() return err @@ -174,6 +286,16 @@ func (prov *DatabaseProvider) openStorageWithContext(readOnly bool, openCtx cont // address for a fresh connection. prov.duckLake.resetInitialized() } + // Keep transaction admission/finalization coupled to this provider even + // when callers reach the pool directly (for example connection-close and + // replication teardown paths). Standalone pools leave these hooks unset. + prov.pool.SetTransactionLifecycleHooksWithError(prov.beginDuckLakeTransactionWithError, prov.untrackDuckLakeTransaction) + // A malformed teardown poisons this storage generation. Keep the pool's + // physical transaction admission fail-closed as well as the provider's + // logical operation barrier; otherwise a caller that reaches GetTxn directly + // could open an untracked transaction after the generation was retired. + prov.pool.SetTransactionAdmissionGuard(prov.duckLakeGenerationPoisonError) + prov.pool.SetConnectionAdmissionGuard(prov.duckLakeConnectionAdmissionError) prov.pool.SetConnectionInitializer(prov.initializeConnection) prov.connector = connector prov.storage = storage @@ -250,9 +372,19 @@ func (prov *DatabaseProvider) openStorageWithContext(readOnly bool, openCtx cont } func (prov *DatabaseProvider) initCatalog() error { + return prov.initCatalogWithExecutor(prov.storage) +} + +// initCatalogWithExecutor creates provider metadata through one selected +// database/sql executor. CREATE DATABASE supplies its session connection here +// so ATTACH, USE, and initialization cannot split across pooled connections. +func (prov *DatabaseProvider) initCatalogWithExecutor(execer adapter.SQLExecutor) error { + if prov == nil || execer == nil { + return fmt.Errorf("catalog initialization executor is unavailable") + } for _, t := range internalSchemas { - if _, err := prov.storage.ExecContext( + if _, err := execer.ExecContext( context.Background(), "CREATE SCHEMA IF NOT EXISTS "+t.Schema, ); err != nil { @@ -261,20 +393,20 @@ func (prov *DatabaseProvider) initCatalog() error { } for _, t := range internalTables { - if _, err := prov.storage.ExecContext( + if _, err := execer.ExecContext( context.Background(), "CREATE SCHEMA IF NOT EXISTS "+t.Schema, ); err != nil { return fmt.Errorf("failed to create internal schema %q: %w", t.Schema, err) } - if _, err := prov.storage.ExecContext( + if _, err := execer.ExecContext( context.Background(), "CREATE TABLE IF NOT EXISTS "+t.QualifiedName()+"("+t.DDL+")", ); err != nil { return fmt.Errorf("failed to create internal table %q: %w", t.Name, err) } for _, row := range t.InitialData { - if _, err := prov.storage.ExecContext( + if _, err := execer.ExecContext( context.Background(), t.UpsertStmt(), row..., @@ -287,7 +419,7 @@ func (prov *DatabaseProvider) initCatalog() error { if initialFileContent != "" { var count int // Count rows in the internal table - if err := prov.storage.QueryRow(t.CountAllStmt()).Scan(&count); err != nil { + if err := execer.QueryRowContext(context.Background(), t.CountAllStmt()).Scan(&count); err != nil { return fmt.Errorf("failed to count rows in internal table %q: %w", t.Name, err) } @@ -311,7 +443,7 @@ func (prov *DatabaseProvider) initCatalog() error { } // Execute the COPY command to insert data into the table - if _, err := prov.storage.ExecContext( + if _, err := execer.ExecContext( context.Background(), fmt.Sprintf("COPY %s FROM '%s' (DELIMITER ',', HEADER, ESCAPE '\"')", t.QualifiedName(), tmpFile.Name()), ); err != nil { @@ -322,13 +454,13 @@ func (prov *DatabaseProvider) initCatalog() error { } for _, v := range InternalViews { - if _, err := prov.storage.ExecContext( + if _, err := execer.ExecContext( context.Background(), "CREATE SCHEMA IF NOT EXISTS "+v.Schema, ); err != nil { return fmt.Errorf("failed to create internal schema %q: %w", v.Schema, err) } - if _, err := prov.storage.ExecContext( + if _, err := execer.ExecContext( context.Background(), "CREATE VIEW IF NOT EXISTS "+v.QualifiedName()+" AS "+v.DDL, ); err != nil { @@ -337,7 +469,7 @@ func (prov *DatabaseProvider) initCatalog() error { } for _, m := range InternalMacros { - if _, err := prov.storage.ExecContext( + if _, err := execer.ExecContext( context.Background(), "CREATE SCHEMA IF NOT EXISTS "+m.Schema, ); err != nil { @@ -354,7 +486,7 @@ func (prov *DatabaseProvider) initCatalog() error { } definitions = append(definitions, fmt.Sprintf("\n(%s) AS %s%s", macroParams, asType, d.DDL)) } - if _, err := prov.storage.ExecContext( + if _, err := execer.ExecContext( context.Background(), "CREATE OR REPLACE MACRO "+m.QualifiedName()+strings.Join(definitions, ",")+";", ); err != nil { @@ -362,13 +494,13 @@ func (prov *DatabaseProvider) initCatalog() error { } } - if _, err := prov.pool.ExecContext(context.Background(), "PRAGMA enable_checkpoint_on_shutdown"); err != nil { + if _, err := execer.ExecContext(context.Background(), "PRAGMA enable_checkpoint_on_shutdown"); err != nil { logrus.WithError(err).Fatalln("Failed to enable checkpoint on shutdown") } // Postgres tables are created in the `public` schema by default. // Create the `public` schema if it doesn't exist. - _, err := prov.pool.ExecContext(context.Background(), "CREATE SCHEMA IF NOT EXISTS public") + _, err := execer.ExecContext(context.Background(), "CREATE SCHEMA IF NOT EXISTS public") if err != nil { logrus.WithError(err).Fatalln("Failed to create the `public` schema") } @@ -432,8 +564,49 @@ func (prov *DatabaseProvider) AttachCatalog(file interface { } func (prov *DatabaseProvider) CreateCatalog(name string, ifNotExists bool) error { + if prov == nil { + return fmt.Errorf("database provider is unavailable") + } + if prov.mu != nil { + prov.mu.Lock() + defer prov.mu.Unlock() + } + if prov.storage == nil { + return fmt.Errorf("database storage is unavailable") + } + ctx := context.Background() + conn, err := prov.storage.Conn(ctx) + if err != nil { + return err + } + defer conn.Close() + return prov.createCatalogOnConnLocked(ctx, conn, name, ifNotExists) +} + +// CreateCatalogOnConn creates and attaches a database using the caller's +// already-owned physical connection. The provider lock serializes the whole +// ATTACH/initialization sequence with Restart, while the connection preserves +// the PostgreSQL client's session catalog state. +func (prov *DatabaseProvider) CreateCatalogOnConn(ctx context.Context, conn *stdsql.Conn, name string, ifNotExists bool) error { + if prov == nil { + return fmt.Errorf("database provider is unavailable") + } + if prov.mu != nil { + prov.mu.Lock() + defer prov.mu.Unlock() + } + return prov.createCatalogOnConnLocked(ctx, conn, name, ifNotExists) +} + +func (prov *DatabaseProvider) createCatalogOnConnLocked(ctx context.Context, conn *stdsql.Conn, name string, ifNotExists bool) error { + if ctx == nil { + ctx = context.Background() + } + if conn == nil { + return fmt.Errorf("database connection is unavailable") + } name = strings.TrimSpace(name) - // in memory database does not need to be created + // The in-memory database is always present and does not need ATTACH. if name == "" || name == "memory" { return nil } @@ -442,80 +615,155 @@ func (prov *DatabaseProvider) CreateCatalog(name string, ifNotExists bool) error _, err := os.Stat(dsn) shouldInit := os.IsNotExist(err) - // attach attachSQL := "ATTACH" if ifNotExists { attachSQL += " IF NOT EXISTS" } quoted := QuoteIdentifierANSI(name) attachSQL += " '" + dsn + "' AS " + quoted - _, err = prov.storage.ExecContext(context.Background(), attachSQL) - if err != nil { + if _, err = conn.ExecContext(ctx, attachSQL); err != nil { return err } - if shouldInit { - res, err := prov.storage.QueryContext(context.Background(), "SELECT current_catalog") - if err != nil { + if !shouldInit { + return nil + } + res, err := conn.QueryContext(ctx, "SELECT current_catalog") + if err != nil { + return fmt.Errorf("failed to init catalog: %w", err) + } + defer res.Close() + lastCatalog := "" + for res.Next() { + if err := res.Scan(&lastCatalog); err != nil { return fmt.Errorf("failed to init catalog: %w", err) } - lastCatalog := "" - for res.Next() { - if err := res.Scan(&lastCatalog); err != nil { - return fmt.Errorf("failed to init catalog: %w", err) - } - } + } + if err := res.Err(); err != nil { + return fmt.Errorf("failed to init catalog: %w", err) + } - if _, err := prov.storage.ExecContext(context.Background(), "USE "+quoted); err != nil { - return fmt.Errorf("failed to switch to the new catalog: %w", err) - } + if _, err := conn.ExecContext(ctx, "USE "+quoted); err != nil { + return fmt.Errorf("failed to switch to the new catalog: %w", err) + } - defer func() { - if lastCatalog == "" { - return - } - if _, err := prov.storage.ExecContext(context.Background(), "USE "+QuoteIdentifierANSI(lastCatalog)); err != nil { - logrus.WithError(err).Errorln("Failed to switch back to the old catalog") - } - }() - err = prov.initCatalog() - if err != nil { - return err + defer func() { + if lastCatalog == "" { + return } - } - return nil + if _, err := conn.ExecContext(context.WithoutCancel(ctx), "USE "+QuoteIdentifierANSI(lastCatalog)); err != nil { + logrus.WithError(err).Errorln("Failed to switch back to the old catalog") + } + }() + return prov.initCatalogWithExecutor(conn) } func (prov *DatabaseProvider) DropCatalog(name string, ifExists bool) error { + if prov == nil { + return fmt.Errorf("database provider is unavailable") + } + if prov.mu != nil { + prov.mu.Lock() + defer prov.mu.Unlock() + } + if prov.storage == nil { + return fmt.Errorf("database storage is unavailable") + } + ctx := context.Background() + conn, err := prov.storage.Conn(ctx) + if err != nil { + return err + } + defer conn.Close() + return prov.dropCatalogOnConnLocked(ctx, conn, name, ifExists) +} + +// DropCatalogOnConn detaches and removes a database using the caller's +// physical connection. It is serialized with Restart for the full detach/file +// removal sequence. +func (prov *DatabaseProvider) DropCatalogOnConn(ctx context.Context, conn *stdsql.Conn, name string, ifExists bool) error { + if prov == nil { + return fmt.Errorf("database provider is unavailable") + } + if prov.mu != nil { + prov.mu.Lock() + defer prov.mu.Unlock() + } + return prov.dropCatalogOnConnLocked(ctx, conn, name, ifExists) +} + +func (prov *DatabaseProvider) dropCatalogOnConnLocked(ctx context.Context, conn *stdsql.Conn, name string, ifExists bool) error { + if ctx == nil { + ctx = context.Background() + } + if conn == nil { + return fmt.Errorf("database connection is unavailable") + } name = strings.TrimSpace(name) - // in memory database does not need to be created if name == "" || name == "memory" { return fmt.Errorf("cannot drop the in-memory catalog") } dsn := filepath.Join(prov.dataDir, name+".db") - // if file does not exist, return error - _, err := os.Stat(dsn) - if os.IsNotExist(err) { + if _, err := os.Stat(dsn); os.IsNotExist(err) { if ifExists { return nil } return fmt.Errorf("database file %s does not exist", dsn) } - // detach - if _, err := prov.storage.ExecContext(context.Background(), "DETACH "+QuoteIdentifierANSI(name)); err != nil { + if _, err := conn.ExecContext(ctx, "DETACH "+QuoteIdentifierANSI(name)); err != nil { return fmt.Errorf("failed to detach catalog %w", err) } - // delete the file - err = os.Remove(dsn) - if err != nil { + if err := os.Remove(dsn); err != nil { return fmt.Errorf("failed to delete database file %s: %w", dsn, err) } return nil } -func (prov *DatabaseProvider) Close() error { - defer prov.connector.Close() - return prov.pool.Close() +func (prov *DatabaseProvider) Close() (err error) { + if prov == nil { + return nil + } + // Announce the generation transition before taking the provider mutex. An + // already-admitted frontend operation may still need that mutex (for + // catalog DDL/metadata) before it can release its DuckLake operation lease. + // Holding prov.mu first would make pool teardown wait on the operation while + // the operation waits on prov.mu. + releaseTransition, transitionErr := prov.beginDuckLakeGenerationTransitionWithError() + if transitionErr != nil { + return transitionErr + } + defer func() { + err = errors.Join(err, releaseTransition()) + }() + if prov.mu != nil { + prov.mu.Lock() + defer prov.mu.Unlock() + } + err = prov.closeLocked() + return err +} + +// closeLocked retires the current storage generation. Callers that already +// hold prov.mu (notably Restart) use this helper to avoid self-deadlocking. +func (prov *DatabaseProvider) closeLocked() error { + if prov == nil { + return nil + } + var err error + if prov.pool != nil { + err = prov.pool.Close() + if err != nil { + // Pool teardown deliberately preserves malformed mappings and their + // unknown physical owners. Mark this generation unusable so the + // transition release can return instead of waiting on state that no + // type-safe path can finish. + prov.poisonDuckLakeGeneration(err) + } + } + if prov.connector != nil { + _ = prov.connector.Close() + } + return err } func (prov *DatabaseProvider) Connector() *duckdb.Connector { @@ -600,10 +848,1130 @@ func (prov *DatabaseProvider) DuckLakeEnabled() bool { return prov != nil && prov.duckLake != nil } +// DuckLakeObjectStorageEnabled reports whether the service has both sides of +// the DuckLake catalog configured. An enabled extension layer may intentionally +// omit these paths, but it cannot produce or clean lake-table files. +func (prov *DatabaseProvider) DuckLakeObjectStorageEnabled() bool { + return prov != nil && prov.duckLakeObjectStorageEnabled() +} + +// withoutCancelContext removes request cancellation while retaining the +// concrete *sql.Context type. Provider boundaries use the concrete type to +// detect a session transaction and avoid issuing *sql.Conn work alongside an +// active *sql.Tx. +func withoutCancelContext(ctx context.Context) context.Context { + if ctx == nil { + return context.Background() + } + if sqlCtx, ok := ctx.(*sql.Context); ok && sqlCtx != nil { + return sqlCtx.WithContext(context.WithoutCancel(sqlCtx)) + } + return context.WithoutCancel(ctx) +} + +// EnsureDuckLakeConnection makes the service-owned lake available on an +// already acquired physical connection. Object-table paths call this after +// resolving persisted metadata; ordinary local and catalog operations do not +// need to invoke it explicitly. +func (prov *DatabaseProvider) EnsureDuckLakeConnection(ctx context.Context, conn *stdsql.Conn) error { + return prov.ensureDuckLakeConnection(ctx, conn, nil, false) +} + +// EnsureDuckLakeConnectionWithBinding initializes DuckLake on a connection +// selected by an already-captured execution snapshot. bindingKnown is kept +// private in the implementation so this boundary cannot accidentally perform +// a second pool lookup while the caller retains a lifecycle lease. +// +// A non-nil tx proves that the pool initializer ran before BeginTx and that +// this exact physical owner is already attached. It does not bypass provider +// admission: frontend snapshots, including active transactions, are rejected +// while the storage generation is transitioning or poisoned. A nil tx means +// an admitted connection-only snapshot may perform the driver-level attach. +func (prov *DatabaseProvider) EnsureDuckLakeConnectionWithBinding( + ctx context.Context, + conn *stdsql.Conn, + tx *stdsql.Tx, +) error { + return prov.ensureDuckLakeConnection(ctx, conn, tx, true) +} + +// EnsureDuckLakeConnectionWithExecutor is the convenient form used by +// statement paths that already carry both the SQL executor and its physical +// owner. *sql.Tx is recognized without consulting the session/pool maps. +func (prov *DatabaseProvider) EnsureDuckLakeConnectionWithExecutor( + ctx context.Context, + execer adapter.SQLExecutor, + conn *stdsql.Conn, +) error { + var tx *stdsql.Tx + if typed, ok := execer.(*stdsql.Tx); ok { + tx = typed + } + return prov.ensureDuckLakeConnection(ctx, conn, tx, true) +} + +func (prov *DatabaseProvider) ensureDuckLakeConnection( + ctx context.Context, + conn *stdsql.Conn, + boundTx *stdsql.Tx, + bindingKnown bool, +) error { + if prov == nil || prov.duckLake == nil { + return fmt.Errorf("%w: DuckLake service configuration is disabled", ErrInvalidTableStorage) + } + if err := prov.duckLakeConnectionAdmissionError(ctx); err != nil { + return err + } + if conn == nil { + return fmt.Errorf("ducklake storage connection is unavailable") + } + // For an admitted request, a session transaction owns this physical + // connection. The pool initializer runs before BeginTx, so an active + // transaction is already initialized and attached; calling Conn.Raw here + // would race the transaction's driver calls (and can deadlock when the driver + // serializes access). Object callers use their SQLExecutor (*sql.Tx) for all + // statement work while the transaction is active. + if ownerConn := rollbackCleanupOwnerConn(ctx); ownerConn != nil { + if ownerConn != conn { + return fmt.Errorf("ducklake cleanup connection is not the rollback owner") + } + // The rollback finalizer holds the session lifecycle lock. The explicit + // owner marker is the atomic snapshot for this maintenance operation; + // calling back into the holder here would deadlock. + } else if bindingKnown { + // The caller supplied the atomic snapshot that selected conn. Do not + // re-enter ConnectionPool.GetTxnBinding here: a retained execution lease + // already holds lifecycleMu.RLock, and a pending pool writer would make + // that nested ordinary read wait on itself. A physical transaction has + // already been initialized by the pool before it was returned. + if boundTx != nil { + return nil + } + } else if sqlCtx, ok := ctx.(*sql.Context); ok && sqlCtx != nil { + if _, holder := sqlCtx.Session.(adapter.ConnectionHolder); holder { + boundConn, boundTx := adapter.TryGetTxnBinding(sqlCtx) + if boundConn != nil && boundConn != conn { + return fmt.Errorf("ducklake connection is not owned by the session") + } + if boundTx != nil { + // The transaction was initialized before BeginTx and owns this exact + // physical connection. Do not call Conn.Raw beside its executor. + return nil + } + } + } + return conn.Raw(func(driverConn any) error { + execer, ok := driverConn.(driver.ExecerContext) + if !ok { + return fmt.Errorf("duckdb connection does not support context execution") + } + physicalConn, _ := driverConn.(driver.Conn) + return prov.duckLake.EnsureAttached(ctx, physicalConn, execer) + }) +} + +func (prov *DatabaseProvider) ensureDuckLakeTxnStateLocked() { + if prov.duckLakeTxnCond == nil { + prov.duckLakeTxnCond = sync.NewCond(&prov.duckLakeTxnMu) + } + if prov.duckLakeActiveTx == nil { + prov.duckLakeActiveTx = make(map[*stdsql.Tx]struct{}) + } + if prov.duckLakeOperationOwners == nil { + prov.duckLakeOperationOwners = make(map[any]int) + } + if prov.duckLakeOperationLeases == nil { + prov.duckLakeOperationLeases = make(map[uint64]duckLakeOperationLease) + } +} + +var errDuckLakeGenerationPoisoned = errors.New("DuckLake storage generation is poisoned") +var errDuckLakeGenerationTransition = errors.New("DuckLake storage generation is transitioning") + +// duckLakeGenerationPoisonErrorLocked returns the stable error for a retired +// generation. Callers must hold duckLakeTxnMu; keeping this small locked helper +// avoids trying to reacquire the mutex from an admission path that is already +// checking the poison bit under the condition variable. +func (prov *DatabaseProvider) duckLakeGenerationPoisonErrorLocked() error { + if prov == nil || !prov.duckLakeGenerationPoisoned { + return nil + } + if prov.duckLakeGenerationErr == nil { + return errDuckLakeGenerationPoisoned + } + return errors.Join(errDuckLakeGenerationPoisoned, prov.duckLakeGenerationErr) +} + +// poisonDuckLakeGeneration records a teardown failure without discarding the +// mappings that caused it. A poisoned provider cannot safely reopen or admit +// new DuckLake work, but late physical finalizers may still untrack the exact +// transaction they own. +func (prov *DatabaseProvider) poisonDuckLakeGeneration(err error) { + if prov == nil || err == nil || !prov.duckLakeObjectStorageEnabled() { + return + } + prov.duckLakeTxnMu.Lock() + defer prov.duckLakeTxnMu.Unlock() + prov.ensureDuckLakeTxnStateLocked() + if !prov.duckLakeGenerationPoisoned { + prov.duckLakeGenerationPoisoned = true + prov.duckLakeGenerationErr = err + } else if prov.duckLakeGenerationErr == nil { + prov.duckLakeGenerationErr = err + } + prov.duckLakeTxnCond.Broadcast() +} + +// duckLakeGenerationPoisonError returns the stable teardown error, if this +// provider can no longer safely admit or reopen a DuckLake generation. +func (prov *DatabaseProvider) duckLakeGenerationPoisonError() error { + if prov == nil { + return nil + } + prov.duckLakeTxnMu.Lock() + defer prov.duckLakeTxnMu.Unlock() + return prov.duckLakeGenerationPoisonErrorLocked() +} + +// duckLakeConnectionAdmissionError protects every ordinary pool connection and +// snapshot entry point, not only transaction admission. A provider transition +// keeps cleanupActive set across pool Close/Reset and recovery reopen; reject +// frontend work during that interval so it cannot slip into the retired pool +// after the lifecycle writer unlocks but before poison/reopen is published. +// Provider-owned rollback cleanup and recovery carry explicit context markers +// and may continue on their already-reserved generation. +func (prov *DatabaseProvider) duckLakeConnectionAdmissionError(ctx context.Context) error { + if prov == nil || !prov.duckLakeObjectStorageEnabled() { + return nil + } + prov.duckLakeTxnMu.Lock() + defer prov.duckLakeTxnMu.Unlock() + if poisonErr := prov.duckLakeGenerationPoisonErrorLocked(); poisonErr != nil { + return poisonErr + } + origin := mycontext.QueryOrigin(ctx) + cleanupOwner := duckLakeCleanupLeaseHeld(ctx) && origin == mycontext.MaintenanceQueryOrigin + if prov.duckLakeCleanupActive && !cleanupOwner && origin != mycontext.RecoveryQueryOrigin { + return errDuckLakeGenerationTransition + } + return nil +} + +// comparableDuckLakeOwner returns an owner only when it can be used safely as +// a map key and compared through an interface. A non-comparable (or typed-nil) +// owner is deliberately treated as anonymous, which keeps rollback cleanup +// conservative instead of risking a comparison panic or an accidental +// exemption. +func comparableDuckLakeOwner(owner any) (any, bool) { + if owner == nil { + return nil, false + } + v := reflect.ValueOf(owner) + if !v.IsValid() || !v.Comparable() { + return nil, false + } + switch v.Kind() { + case reflect.Chan, reflect.Func, reflect.Map, reflect.Ptr, reflect.Slice, reflect.UnsafePointer: + if v.IsNil() { + return nil, false + } + } + return owner, true +} + +// duckLakeOperationOwnerFromContext derives the logical GMS transaction when +// the legacy API is called with a concrete SQL context. Contexts without a +// session/transaction remain anonymous and therefore continue to block a +// rollback cleanup reservation. +func duckLakeOperationOwnerFromContext(ctx context.Context) any { + sqlCtx, ok := ctx.(*sql.Context) + if !ok || sqlCtx == nil || sqlCtx.Session == nil { + return nil + } + // Prefer the physical transaction identity when the session exposes an + // atomic binding snapshot. PostgreSQL and explicit MySQL transactions can + // then be released by pool generation teardown without waiting on a logical + // lease whose callback lives only in the GMS transaction object. Implicit + // autocommit scopes still have no physical owner here and remain conservative + // logical/anonymous leases. + if holder, ok := sqlCtx.Session.(adapter.TransactionBindingHolder); ok { + if _, tx := holder.GetTxnBinding(); tx != nil { + return tx + } + } + return sqlCtx.GetTransaction() +} + +// duckLakeOperationBlocksCleanupLocked reports whether any admitted operation +// remains that the rollback caller cannot prove it owns. The caller must hold +// duckLakeTxnMu. The aggregate count is retained as a conservative fallback if +// bookkeeping was populated by an older generation or a legacy test double. +func (prov *DatabaseProvider) duckLakeOperationBlocksCleanupLocked(target *stdsql.Tx, owner any) bool { + ownerKey, ownerOK := comparableDuckLakeOwner(owner) + targetKey, targetOK := comparableDuckLakeOwner(target) + + if prov.duckLakeAnonymousOps > 0 { + return true + } + tracked := 0 + for operationOwner, count := range prov.duckLakeOperationOwners { + if count <= 0 { + continue + } + tracked += count + ownedByRollback := (ownerOK && operationOwner == ownerKey) || + (targetOK && operationOwner == targetKey) + if !ownedByRollback { + return true + } + } + // If an aggregate count exceeds the owner-indexed records, retain the old + // fail-closed behavior: the unclassified leases must still block cleanup. + return prov.duckLakeActiveOperations > tracked +} + +// beginDuckLakeTransaction reserves admission for a transaction before the +// pool starts the driver transaction. The reservation closes the small window +// in which cleanup could otherwise begin between BeginTx and registration. The +// returned completion function is idempotent and must be called with the +// transaction on success or nil on failure. +func (prov *DatabaseProvider) beginDuckLakeTransaction() func(*stdsql.Tx) { + release, _ := prov.beginDuckLakeTransactionWithError() + return release +} + +// beginDuckLakeTransactionWithError reserves admission for a physical +// transaction and preserves a poisoned-generation error for callers that can +// report it. The legacy closure-only wrapper above intentionally remains for +// compatibility with existing pool/test hooks; production pools install this +// error-bearing form so a failed reservation cannot be mistaken for success. +func (prov *DatabaseProvider) beginDuckLakeTransactionWithError() (func(*stdsql.Tx), error) { + if prov == nil || !prov.duckLakeObjectStorageEnabled() { + return func(*stdsql.Tx) {}, nil + } + prov.duckLakeTxnMu.Lock() + prov.ensureDuckLakeTxnStateLocked() + if poisonErr := prov.duckLakeGenerationPoisonErrorLocked(); poisonErr != nil { + prov.duckLakeTxnMu.Unlock() + return func(*stdsql.Tx) {}, poisonErr + } + for prov.duckLakeCleanupActive && !prov.duckLakeGenerationPoisoned { + prov.duckLakeTxnCond.Wait() + } + if poisonErr := prov.duckLakeGenerationPoisonErrorLocked(); poisonErr != nil { + prov.duckLakeTxnMu.Unlock() + return func(*stdsql.Tx) {}, poisonErr + } + prov.duckLakePendingTx++ + prov.duckLakeTxnMu.Unlock() + + var once sync.Once + return func(tx *stdsql.Tx) { + once.Do(func() { + prov.duckLakeTxnMu.Lock() + if prov.duckLakePendingTx > 0 { + prov.duckLakePendingTx-- + } + if tx != nil { + prov.duckLakeActiveTx[tx] = struct{}{} + } + prov.duckLakeTxnCond.Broadcast() + prov.duckLakeTxnMu.Unlock() + }) + }, nil +} + +// BeginDuckLakeOperation reserves admission for one top-level frontend +// operation that may read or write the service-managed DuckLake relation. +// +// GMS creates a logical autocommit transaction without a physical *sql.Tx +// when @@autocommit=1. The transaction barrier therefore cannot observe that +// work through duckLakeActiveTx alone. This lease covers the interval from +// execution-build through iterator consumption, allowing rollback cleanup to +// wait for the operation even when no physical transaction exists. +// +// The returned release function is idempotent. Callers must release it on +// build failure and from every terminal iterator path (EOF, error, or Close). +func (prov *DatabaseProvider) BeginDuckLakeOperation(ctx context.Context) func() { + release, _ := prov.BeginDuckLakeOperationWithError(ctx) + return release +} + +// BeginDuckLakeOperationWithError is the error-preserving admission surface. +// Legacy callers may continue to use BeginDuckLakeOperation, but new protocol +// paths should use this form so a poisoned generation cannot be mistaken for a +// successful no-op admission. +func (prov *DatabaseProvider) BeginDuckLakeOperationWithError(ctx context.Context) (func(), error) { + return prov.BeginDuckLakeOperationForOwnerWithError(ctx, duckLakeOperationOwnerFromContext(ctx)) +} + +// BeginDuckLakeOperationForOwner is the owner-aware form of +// BeginDuckLakeOperation. owner should be the logical transaction identity +// used by the caller's rollback path (for example, the concrete GMS +// transaction). A physical *sql.Tx is also accepted and is matched against a +// rollback target. Owners that are not safely comparable are retained as +// anonymous blocking leases. +func (prov *DatabaseProvider) BeginDuckLakeOperationForOwner(ctx context.Context, owner any) func() { + release, _ := prov.BeginDuckLakeOperationForOwnerWithError(ctx, owner) + return release +} + +// BeginDuckLakeOperationForOwnerWithError is the owner-aware, error-preserving +// admission surface. A non-nil error means no lease was admitted and the +// returned release function is a safe no-op. +func (prov *DatabaseProvider) BeginDuckLakeOperationForOwnerWithError(ctx context.Context, owner any) (func(), error) { + if prov == nil || !prov.duckLakeObjectStorageEnabled() || + mycontext.QueryOrigin(ctx) != mycontext.FrontendQueryOrigin { + return func() {}, nil + } + ownerKey, owned := comparableDuckLakeOwner(owner) + + prov.duckLakeTxnMu.Lock() + prov.ensureDuckLakeTxnStateLocked() + if poisonErr := prov.duckLakeGenerationPoisonErrorLocked(); poisonErr != nil { + prov.duckLakeTxnMu.Unlock() + return func() {}, poisonErr + } + for prov.duckLakeCleanupActive && !prov.duckLakeGenerationPoisoned { + prov.duckLakeTxnCond.Wait() + } + if poisonErr := prov.duckLakeGenerationPoisonErrorLocked(); poisonErr != nil { + prov.duckLakeTxnMu.Unlock() + return func() {}, poisonErr + } + prov.duckLakeNextOperationID++ + leaseID := prov.duckLakeNextOperationID + prov.duckLakeOperationLeases[leaseID] = duckLakeOperationLease{ + owner: ownerKey, + owned: owned, + } + if owned { + prov.duckLakeOperationOwners[ownerKey]++ + } else { + prov.duckLakeAnonymousOps++ + } + prov.duckLakeActiveOperations++ + prov.duckLakeTxnMu.Unlock() + + var once sync.Once + return func() { + once.Do(func() { + prov.duckLakeTxnMu.Lock() + prov.ensureDuckLakeTxnStateLocked() + lease, present := prov.duckLakeOperationLeases[leaseID] + if !present { + prov.duckLakeTxnMu.Unlock() + return + } + delete(prov.duckLakeOperationLeases, leaseID) + if lease.owned { + if count := prov.duckLakeOperationOwners[lease.owner]; count <= 1 { + delete(prov.duckLakeOperationOwners, lease.owner) + } else { + prov.duckLakeOperationOwners[lease.owner] = count - 1 + } + } else if prov.duckLakeAnonymousOps > 0 { + prov.duckLakeAnonymousOps-- + } + if prov.duckLakeActiveOperations > 0 { + prov.duckLakeActiveOperations-- + } + prov.duckLakeTxnCond.Broadcast() + prov.duckLakeTxnMu.Unlock() + }) + }, nil +} + +// BeginDuckLakeRollbackCleanup reserves the provider-wide DuckLake cleanup +// barrier for a rollback finalizer. Unlike beginDuckLakeCleanup, the target +// transaction is allowed to remain active while the reservation is acquired: +// the caller must still run the driver's Rollback before maintenance can use +// the owner's physical connection. Every other physical transaction, pending +// admission, and logical frontend operation is drained first, and no new work +// can enter until the returned release function is called. +// +// The reservation must be obtained before entering a pool/session finalizer. +// That ordering prevents a finalizer holding the pool release lock from +// waiting on an operation that itself needs the pool lock to finish. The +// returned callback is idempotent and safe to defer on all paths. +func (prov *DatabaseProvider) BeginDuckLakeRollbackCleanup(target *stdsql.Tx) func() { + release, _ := prov.BeginDuckLakeRollbackCleanupWithError(target) + return release +} + +// BeginDuckLakeRollbackCleanupWithError is the error-preserving counterpart to +// BeginDuckLakeRollbackCleanup. A non-nil error means the provider did not +// reserve its cleanup barrier; callers must still finalize the driver +// transaction, but must skip DuckLake maintenance for this poisoned generation. +func (prov *DatabaseProvider) BeginDuckLakeRollbackCleanupWithError(target *stdsql.Tx) (func(), error) { + return prov.BeginDuckLakeRollbackCleanupForOwnerWithError(target, nil) +} + +// BeginDuckLakeRollbackCleanupForOwner is the owner-aware rollback barrier. +// The physical target and logical owner are both considered when deciding +// which already-admitted operation leases belong to this rollback caller. Only +// those leases are exempt; every other owner, plus anonymous leases, keeps the +// reservation blocked. +func (prov *DatabaseProvider) BeginDuckLakeRollbackCleanupForOwner(target *stdsql.Tx, owner any) func() { + release, _ := prov.BeginDuckLakeRollbackCleanupForOwnerWithError(target, owner) + return release +} + +// BeginDuckLakeRollbackCleanupForOwnerWithError is the owner-aware, +// error-preserving rollback barrier. The returned release is always safe to +// call, including when admission fails before the provider mutexes are held. +func (prov *DatabaseProvider) BeginDuckLakeRollbackCleanupForOwnerWithError(target *stdsql.Tx, owner any) (func(), error) { + if prov == nil || !prov.duckLakeObjectStorageEnabled() { + return func() {}, nil + } + + prov.duckLakeCleanupMu.Lock() + prov.duckLakeTxnMu.Lock() + prov.ensureDuckLakeTxnStateLocked() + if poisonErr := prov.duckLakeGenerationPoisonErrorLocked(); poisonErr != nil { + prov.duckLakeTxnMu.Unlock() + prov.duckLakeCleanupMu.Unlock() + return func() {}, poisonErr + } + prov.duckLakeCleanupActive = true + for { + if poisonErr := prov.duckLakeGenerationPoisonErrorLocked(); poisonErr != nil { + prov.duckLakeCleanupActive = false + prov.duckLakeTxnCond.Broadcast() + prov.duckLakeTxnMu.Unlock() + prov.duckLakeCleanupMu.Unlock() + return func() {}, poisonErr + } + otherTransactions := len(prov.duckLakeActiveTx) + if target != nil { + if _, present := prov.duckLakeActiveTx[target]; present { + otherTransactions-- + } + } + if otherTransactions <= 0 && + prov.duckLakePendingTx == 0 && + !prov.duckLakeOperationBlocksCleanupLocked(target, owner) { + break + } + prov.duckLakeTxnCond.Wait() + } + prov.duckLakeTxnMu.Unlock() + + var once sync.Once + return func() { + once.Do(func() { + prov.duckLakeTxnMu.Lock() + prov.ensureDuckLakeTxnStateLocked() + prov.duckLakeCleanupActive = false + prov.duckLakeTxnCond.Broadcast() + prov.duckLakeTxnMu.Unlock() + prov.duckLakeCleanupMu.Unlock() + }) + }, nil +} + +// trackDuckLakeTransaction records a transaction that was admitted by the +// provider barrier. It is kept as a small idempotent primitive for lifecycle +// callers that already hold a transaction and need to repair bookkeeping. +func (prov *DatabaseProvider) trackDuckLakeTransaction(tx *stdsql.Tx) { + if prov == nil || tx == nil || !prov.duckLakeObjectStorageEnabled() { + return + } + prov.duckLakeTxnMu.Lock() + defer prov.duckLakeTxnMu.Unlock() + prov.ensureDuckLakeTxnStateLocked() + prov.duckLakeActiveTx[tx] = struct{}{} +} + +// untrackDuckLakeTransaction is idempotent so both normal finalization and a +// connection-close fallback may safely release the same transaction. +func (prov *DatabaseProvider) untrackDuckLakeTransaction(tx *stdsql.Tx) { + if prov == nil || tx == nil || !prov.duckLakeObjectStorageEnabled() { + return + } + prov.duckLakeTxnMu.Lock() + defer prov.duckLakeTxnMu.Unlock() + prov.ensureDuckLakeTxnStateLocked() + if _, exists := prov.duckLakeActiveTx[tx]; exists { + delete(prov.duckLakeActiveTx, tx) + prov.duckLakeTxnCond.Broadcast() + } +} + +func (prov *DatabaseProvider) registerUncommittedLakeRelation(tx *stdsql.Tx, conn *stdsql.Conn, physical string) { + if prov == nil || tx == nil || strings.TrimSpace(physical) == "" { + return + } + prov.duckLakeTxnMu.Lock() + prov.ensureDuckLakeTxnStateLocked() + prov.duckLakeUncommitted = append(prov.duckLakeUncommitted, uncommittedLakeRelation{ + tx: tx, + conn: conn, + physical: physical, + }) + prov.duckLakeTxnMu.Unlock() + if prov.pool != nil { + prov.pool.RegisterTransactionCompletion(tx, func(success bool) { + if success { + prov.forgetUncommittedLakeRelation(tx, physical) + } + }) + } +} + +func (prov *DatabaseProvider) forgetUncommittedLakeRelation(tx *stdsql.Tx, physical string) { + if prov == nil || tx == nil { + return + } + prov.duckLakeTxnMu.Lock() + defer prov.duckLakeTxnMu.Unlock() + prov.ensureDuckLakeTxnStateLocked() + kept := prov.duckLakeUncommitted[:0] + for _, rel := range prov.duckLakeUncommitted { + if rel.tx == tx && (physical == "" || rel.physical == physical) { + continue + } + kept = append(kept, rel) + } + prov.duckLakeUncommitted = kept +} + +func (prov *DatabaseProvider) takeUncommittedLakeRelationsForConn(conn *stdsql.Conn) []string { + if prov == nil || conn == nil { + return nil + } + prov.duckLakeTxnMu.Lock() + defer prov.duckLakeTxnMu.Unlock() + prov.ensureDuckLakeTxnStateLocked() + var taken []string + kept := prov.duckLakeUncommitted[:0] + for _, rel := range prov.duckLakeUncommitted { + if rel.conn == conn { + taken = append(taken, rel.physical) + continue + } + kept = append(kept, rel) + } + prov.duckLakeUncommitted = kept + return taken +} + +// beginDuckLakeCleanup obtains an exclusive provider barrier. It waits for +// existing physical transactions and logical frontend operations to finish, +// and blocks admission of new ones until the returned release function is +// called. +func (prov *DatabaseProvider) beginDuckLakeCleanup() func() { + release, _ := prov.beginDuckLakeCleanupWithError() + return release +} + +// beginDuckLakeCleanupWithError is the error-preserving maintenance barrier. +// It is intentionally private: public cleanup callers use +// CleanupDuckLakeOrphansOnConn, which turns a poisoned generation into an +// actionable error before issuing any SQL. +func (prov *DatabaseProvider) beginDuckLakeCleanupWithError() (func(), error) { + if prov == nil || !prov.duckLakeObjectStorageEnabled() { + return func() {}, nil + } + prov.duckLakeCleanupMu.Lock() + prov.duckLakeTxnMu.Lock() + prov.ensureDuckLakeTxnStateLocked() + if poisonErr := prov.duckLakeGenerationPoisonErrorLocked(); poisonErr != nil { + prov.duckLakeTxnMu.Unlock() + prov.duckLakeCleanupMu.Unlock() + return func() {}, poisonErr + } + prov.duckLakeCleanupActive = true + for !prov.duckLakeGenerationPoisoned && + (len(prov.duckLakeActiveTx) > 0 || prov.duckLakePendingTx > 0 || prov.duckLakeActiveOperations > 0) { + prov.duckLakeTxnCond.Wait() + } + if poisonErr := prov.duckLakeGenerationPoisonErrorLocked(); poisonErr != nil { + prov.duckLakeCleanupActive = false + prov.duckLakeTxnCond.Broadcast() + prov.duckLakeTxnMu.Unlock() + prov.duckLakeCleanupMu.Unlock() + return func() {}, poisonErr + } + prov.duckLakeTxnMu.Unlock() + return func() { + prov.duckLakeTxnMu.Lock() + prov.duckLakeCleanupActive = false + prov.duckLakeTxnCond.Broadcast() + prov.duckLakeTxnMu.Unlock() + prov.duckLakeCleanupMu.Unlock() + }, nil +} + +func (prov *DatabaseProvider) resetDuckLakeTransactions() { + if prov == nil { + return + } + // Keep this compatibility helper on the same generation-transition path as + // provider Close/Restart. The transition drains anonymous/non-physical + // operations before pool teardown, then lets physical transaction callbacks + // drain before it resets bookkeeping. + releaseTransition := prov.beginDuckLakeGenerationTransition() + releaseTransition() +} + +// duckLakeGenerationOperationsBlockLocked reports whether an admitted +// operation must finish before a storage-generation transition can enter pool +// teardown. A physical *sql.Tx owner is intentionally allowed through: pool +// teardown retires that transaction and the pool completion bridge releases +// the operation lease. Anonymous and logical/non-physical owners have no such +// physical completion boundary, so waiting for them here avoids entering a +// transition that could otherwise strand their lease. +// +// The caller must hold duckLakeTxnMu. Missing or inconsistent per-lease +// bookkeeping is treated conservatively as blocking rather than allowing a +// transition to reset state underneath an untracked operation. +func (prov *DatabaseProvider) duckLakeGenerationOperationsBlockLocked() bool { + if prov == nil { + return false + } + if prov.duckLakeAnonymousOps > 0 { + return true + } + tracked := 0 + for _, lease := range prov.duckLakeOperationLeases { + tracked++ + if !lease.owned { + return true + } + owner, physical := lease.owner.(*stdsql.Tx) + if !physical || owner == nil { + return true + } + } + // An aggregate count without a corresponding lease record can only come + // from legacy/partially initialized state. Fail closed until that count is + // released rather than silently dropping an operation during reset. + return prov.duckLakeActiveOperations != tracked +} + +// beginDuckLakeGenerationTransition announces an exclusive provider storage +// generation transition. It blocks admission of new DuckLake transactions and +// logical frontend operations, then waits for operations already admitted to +// finish. The returned release function must remain held across pool +// close/reset and reopen. Once those lifecycle calls have retired physical +// transactions, release waits for their finished hooks, resets the old +// generation bookkeeping, and reopens admission. +func (prov *DatabaseProvider) beginDuckLakeGenerationTransition() func() { + release, _ := prov.beginDuckLakeGenerationTransitionWithError() + return func() { + if release != nil { + _ = release() + } + } +} + +// beginDuckLakeGenerationTransitionWithError is the error-preserving form used +// by Close and Restart. A poisoned generation is never reopened or reused; the +// release callback also reports a poison discovered during pool teardown. +func (prov *DatabaseProvider) beginDuckLakeGenerationTransitionWithError() (func() error, error) { + if prov == nil || !prov.duckLakeObjectStorageEnabled() { + return func() error { return nil }, nil + } + + prov.duckLakeCleanupMu.Lock() + prov.duckLakeTxnMu.Lock() + prov.ensureDuckLakeTxnStateLocked() + if poisonErr := prov.duckLakeGenerationPoisonErrorLocked(); poisonErr != nil { + prov.duckLakeTxnMu.Unlock() + prov.duckLakeCleanupMu.Unlock() + return func() error { return nil }, poisonErr + } + prov.duckLakeCleanupActive = true + // Do not wait for physical-owner operation leases here. Pool close/reset is + // the operation that rolls those transactions back and invokes their + // completion callbacks (COPY relies on this deferred boundary). Anonymous + // and logical/non-physical leases have no pool callback that can release + // them, so they must drain before the provider enters pool teardown. + for !prov.duckLakeGenerationPoisoned && + (prov.duckLakePendingTx > 0 || prov.duckLakeGenerationOperationsBlockLocked()) { + prov.duckLakeTxnCond.Wait() + } + if poisonErr := prov.duckLakeGenerationPoisonErrorLocked(); poisonErr != nil { + prov.duckLakeCleanupActive = false + prov.duckLakeTxnCond.Broadcast() + prov.duckLakeTxnMu.Unlock() + prov.duckLakeCleanupMu.Unlock() + return func() error { return nil }, poisonErr + } + prov.duckLakeTxnMu.Unlock() + + var once sync.Once + var releaseErr error + release := func() error { + once.Do(func() { + prov.duckLakeTxnMu.Lock() + prov.ensureDuckLakeTxnStateLocked() + // Pool teardown should have rolled back and untracked every physical + // transaction and completed every operation callback. Keep the wait here + // so a delayed finished hook cannot be mistaken for a transaction or + // operation in the replacement generation. + for !prov.duckLakeGenerationPoisoned && + (len(prov.duckLakeActiveTx) > 0 || prov.duckLakeActiveOperations > 0 || prov.duckLakePendingTx > 0) { + prov.duckLakeTxnCond.Wait() + } + if poisonErr := prov.duckLakeGenerationPoisonErrorLocked(); poisonErr != nil { + // Keep malformed mappings and active-owner bookkeeping for recovery. + // This generation is permanently closed, so wake admission waiters + // without resetting state underneath a late finalizer. + releaseErr = poisonErr + prov.duckLakeCleanupActive = false + prov.duckLakeTxnCond.Broadcast() + prov.duckLakeTxnMu.Unlock() + prov.duckLakeCleanupMu.Unlock() + return + } + prov.duckLakeActiveTx = make(map[*stdsql.Tx]struct{}) + prov.duckLakePendingTx = 0 + prov.duckLakeActiveOperations = 0 + prov.duckLakeOperationOwners = make(map[any]int) + prov.duckLakeOperationLeases = make(map[uint64]duckLakeOperationLease) + prov.duckLakeAnonymousOps = 0 + // Keep operation IDs monotonic across generations. A late idempotent + // release closure from the retired generation must never collide with a + // newly admitted lease after admission reopens. + prov.duckLakeCleanupActive = false + prov.duckLakeTxnCond.Broadcast() + prov.duckLakeTxnMu.Unlock() + prov.duckLakeCleanupMu.Unlock() + }) + return releaseErr + } + return release, nil +} + +// duckLakeOrphanCleanupSQL is the product-owned cleanup form. The table +// function returns one path column; selecting that column keeps this call +// stable across DuckLake extension revisions and, importantly, performs the +// cleanup rather than the dry-run/reporting variant. +const duckLakeOrphanCleanupSQL = `SELECT path +FROM ducklake_delete_orphaned_files('__myduck_ducklake', cleanup_all => true)` + +// CleanupDuckLakeOrphans runs the product-owned orphan cleanup function on a +// newly acquired service connection. It is a compatibility entry point for +// maintenance callers that do not already hold a session connection. +func (prov *DatabaseProvider) CleanupDuckLakeOrphans(ctx context.Context) error { + if prov == nil || prov.duckLake == nil || !prov.duckLakeObjectStorageEnabled() { + return nil + } + if poisonErr := prov.duckLakeGenerationPoisonError(); poisonErr != nil { + return poisonErr + } + if !mycontext.IsDuckLakeEligibleQuery(ctx) { + return nil + } + if prov.storage == nil { + return fmt.Errorf("ducklake storage is unavailable") + } + cleanupCtx := withoutCancelContext(ctx) + conn, err := prov.storage.Conn(cleanupCtx) + if err != nil { + return fmt.Errorf("ducklake cleanup connection is unavailable") + } + cleanupErr := prov.CleanupDuckLakeOrphansOnConn(cleanupCtx, conn) + if closeErr := conn.Close(); closeErr != nil && !errors.Is(closeErr, stdsql.ErrConnDone) { + cleanupErr = errors.Join(cleanupErr, fmt.Errorf("ducklake cleanup connection close failed")) + } + return cleanupErr +} + +// CleanupDuckLakeOrphansOnConn performs orphan cleanup on the supplied +// physical connection. Callers that own a session transaction must invoke it +// only after that transaction has become inactive, and should pass the same +// connection so cleanup cannot acquire a second pooled connection. +func (prov *DatabaseProvider) CleanupDuckLakeOrphansOnConn(ctx context.Context, conn *stdsql.Conn) error { + if prov == nil || prov.duckLake == nil || !prov.duckLakeObjectStorageEnabled() { + return nil + } + if poisonErr := prov.duckLakeGenerationPoisonError(); poisonErr != nil { + return poisonErr + } + if !mycontext.IsDuckLakeEligibleQuery(ctx) { + return nil + } + if conn == nil { + return fmt.Errorf("ducklake cleanup connection is unavailable") + } + cleanupCtx := withoutCancelContext(ctx) + // Check the caller before entering the provider-wide barrier. The rollback + // finalizer may invoke this method while its transaction is still represented + // by the session, and beginDuckLakeCleanup would otherwise wait for that + // transaction forever. A finalizer supplies an explicit owner marker after + // it has removed the transaction; ordinary callers with an active session + // transaction are rejected immediately. + if ownerConn := rollbackCleanupOwnerConn(cleanupCtx); ownerConn != nil { + if ownerConn != conn { + return fmt.Errorf("ducklake cleanup connection is not the rollback owner") + } + } else if sqlCtx, ok := cleanupCtx.(*sql.Context); ok && sqlCtx != nil { + if _, holderOK := sqlCtx.Session.(adapter.ConnectionHolder); holderOK { + if _, activeTx := adapter.TryGetTxnBindingForRelease(sqlCtx); activeTx != nil { + return fmt.Errorf("ducklake orphan cleanup requires an inactive transaction") + } + } + } + // ducklake_delete_orphaned_files(cleanup_all => true) mutates shared + // catalog/data state. Exclude all active/new DuckLake transactions while it + // discovers and deletes files, so another session's uncommitted files cannot + // be classified as orphans. The admission check above must happen first so + // this call cannot wait on its own active transaction. + if !duckLakeCleanupLeaseHeld(cleanupCtx) { + releaseCleanup, admissionErr := prov.beginDuckLakeCleanupWithError() + if admissionErr != nil { + return admissionErr + } + defer releaseCleanup() + // The barrier above made cleanupActive visible to the connection guard. + // Mark this already-reserved maintenance scope before Ensure enters Raw so + // it does not reject its own cleanup transition. + cleanupCtx = WithDuckLakeCleanupLease(cleanupCtx) + } + if err := prov.EnsureDuckLakeConnection(cleanupCtx, conn); err != nil { + return fmt.Errorf("ducklake orphan cleanup setup failed: %w", err) + } + var cleanupErr error + for _, physical := range prov.takeUncommittedLakeRelationsForConn(conn) { + if _, err := conn.ExecContext(cleanupCtx, "DROP TABLE IF EXISTS "+physical); err != nil { + cleanupErr = errors.Join(cleanupErr, fmt.Errorf("ducklake uncommitted relation cleanup failed")) + } + } + rows, err := conn.QueryContext(cleanupCtx, duckLakeOrphanCleanupSQL) + if err != nil { + return errors.Join(cleanupErr, fmt.Errorf("ducklake orphan cleanup query failed")) + } + columns, columnsErr := rows.Columns() + if columnsErr != nil { + cleanupErr = fmt.Errorf("ducklake orphan cleanup result inspection failed") + } else { + values := make([]any, len(columns)) + dest := make([]any, len(columns)) + for i := range values { + dest[i] = &values[i] + } + for rows.Next() { + if err := rows.Scan(dest...); err != nil { + cleanupErr = fmt.Errorf("ducklake orphan cleanup result read failed") + break + } + } + } + if err := rows.Err(); err != nil { + cleanupErr = errors.Join(cleanupErr, fmt.Errorf("ducklake orphan cleanup execution failed")) + } + if err := rows.Close(); err != nil { + cleanupErr = errors.Join(cleanupErr, fmt.Errorf("ducklake orphan cleanup result close failed")) + } + return cleanupErr +} + +// duckLakeObjectStorageEnabled reports whether the service has a complete +// metadata/data pair. Extension-only configuration is valid for callers that +// need the DuckDB extensions but has no lake catalog for orphan cleanup. +func (prov *DatabaseProvider) duckLakeObjectStorageEnabled() bool { + if prov == nil || prov.duckLake == nil { + return false + } + return strings.TrimSpace(prov.duckLake.config.MetadataPath) != "" && + strings.TrimSpace(prov.duckLake.config.DataPath) != "" +} + +// LakeSchemaName returns the deterministic physical schema used for a logical +// MyDuck catalog/schema pair. The default database retains its schema names so +// the attached catalog remains inspectable with the normal DuckLake helpers; +// additional database files are prefixed to avoid collisions in the single +// service lake catalog. +func (prov *DatabaseProvider) LakeSchemaName(catalogName, schemaName string) string { + if prov == nil || catalogName == "" || catalogName == prov.defaultCatalogName { + return schemaName + } + return "__myduck_" + lakeIdentifierPart(catalogName) + "__" + lakeIdentifierPart(schemaName) +} + +func lakeIdentifierPart(value string) string { + value = strings.TrimSpace(value) + if value == "" { + return "default" + } + var b strings.Builder + for _, r := range value { + if (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') || (r >= '0' && r <= '9') || r == '_' { + b.WriteRune(r) + } else { + b.WriteByte('_') + } + } + if b.Len() == 0 { + return "default" + } + return b.String() +} + +// ObjectTableName resolves a persisted object-table comment to its physical +// DuckLake relation. It deliberately consults the durable local shadow table +// on every call, so a provider restart cannot retain stale in-memory routing. +func (prov *DatabaseProvider) ObjectTableName(ctx *sql.Context, catalogName, schemaName, tableName string) (string, bool, error) { + if prov == nil || prov.duckLake == nil { + return "", false, nil + } + if strings.TrimSpace(catalogName) == "" || strings.TrimSpace(schemaName) == "" || strings.TrimSpace(tableName) == "" { + return "", false, nil + } + execer, conn, _, release, err := adapter.GetCatalogExecutionSnapshotWithLease(ctx) + if err != nil { + return "", false, err + } + defer release() + if err := prov.EnsureDuckLakeConnectionWithExecutor(ctx, execer, conn); err != nil { + return "", false, err + } + return prov.objectTableNameWithExecutor(ctx, execer, catalogName, schemaName, tableName) +} + +// ObjectTableNameWithExecutor resolves a persisted object-table comment using +// the caller's execution snapshot. The executor and connection must come from +// the same adapter.GetExecutionSnapshot call; this is important while a +// session transaction is active because metadata written in that transaction +// is only visible through its *sql.Tx. +func (prov *DatabaseProvider) ObjectTableNameWithExecutor( + ctx *sql.Context, + execer adapter.SQLExecutor, + catalogName, schemaName, tableName string, +) (string, bool, error) { + if prov == nil || prov.duckLake == nil { + return "", false, nil + } + if strings.TrimSpace(catalogName) == "" || strings.TrimSpace(schemaName) == "" || strings.TrimSpace(tableName) == "" { + return "", false, nil + } + if execer == nil { + return "", false, fmt.Errorf("object-table metadata executor is unavailable") + } + return prov.objectTableNameWithExecutor(ctx, execer, catalogName, schemaName, tableName) +} + +func (prov *DatabaseProvider) objectTableNameWithExecutor( + ctx *sql.Context, + execer adapter.SQLExecutor, + catalogName, schemaName, tableName string, +) (string, bool, error) { + rows, err := execer.QueryContext(ctx, ` + SELECT comment + FROM duckdb_tables() + WHERE database_name = ? AND schema_name = ? AND table_name = ? + `, catalogName, schemaName, tableName) + if err != nil { + return "", false, err + } + defer rows.Close() + if !rows.Next() { + if err := rows.Err(); err != nil { + return "", false, err + } + return "", false, nil + } + var comment stdsql.NullString + if err := rows.Scan(&comment); err != nil { + return "", false, err + } + info := DecodeComment[ExtraTableInfo](comment.String).Meta + if info.StorageKind() != TableStorageObject { + return "", false, nil + } + return FullTableName(DuckLakeCatalogName, prov.LakeSchemaName(catalogName, schemaName), tableName), true, nil +} + +// ObjectTables returns all durable object-table mappings for a logical +// catalog. The result is used by the protocol SQL rewriter and is intentionally +// sourced from duckdb_tables() rather than process-local state. +type ObjectTableMapping struct { + Catalog string + Schema string + Table string + PhysicalName string + PhysicalSchema string +} + +func (prov *DatabaseProvider) ObjectTables(ctx *sql.Context, catalogName string) ([]ObjectTableMapping, error) { + if prov == nil || prov.duckLake == nil { + return nil, nil + } + execer, conn, _, release, err := adapter.GetCatalogExecutionSnapshotWithLease(ctx) + if err != nil { + return nil, err + } + defer release() + if err := prov.EnsureDuckLakeConnectionWithExecutor(ctx, execer, conn); err != nil { + return nil, err + } + return prov.objectTablesWithExecutor(ctx, execer, catalogName) +} + +// ObjectTablesWithExecutor returns durable object-table mappings through an +// already selected execution snapshot. It never acquires another connection. +func (prov *DatabaseProvider) ObjectTablesWithExecutor( + ctx *sql.Context, + execer adapter.SQLExecutor, + catalogName string, +) ([]ObjectTableMapping, error) { + if prov == nil || prov.duckLake == nil { + return nil, nil + } + if execer == nil { + return nil, fmt.Errorf("object-table metadata executor is unavailable") + } + return prov.objectTablesWithExecutor(ctx, execer, catalogName) +} + +func (prov *DatabaseProvider) objectTablesWithExecutor( + ctx *sql.Context, + execer adapter.SQLExecutor, + catalogName string, +) ([]ObjectTableMapping, error) { + rows, err := execer.QueryContext(ctx, ` + SELECT schema_name, table_name, comment + FROM duckdb_tables() + WHERE database_name = ? + `, catalogName) + if err != nil { + return nil, err + } + defer rows.Close() + var mappings []ObjectTableMapping + for rows.Next() { + var schemaName, tableName string + var comment stdsql.NullString + if err := rows.Scan(&schemaName, &tableName, &comment); err != nil { + return nil, err + } + info := DecodeComment[ExtraTableInfo](comment.String).Meta + if info.StorageKind() != TableStorageObject { + continue + } + physicalSchema := prov.LakeSchemaName(catalogName, schemaName) + mappings = append(mappings, ObjectTableMapping{ + Catalog: catalogName, Schema: schemaName, Table: tableName, + PhysicalSchema: physicalSchema, + PhysicalName: FullTableName(DuckLakeCatalogName, physicalSchema, tableName), + }) + } + if err := rows.Err(); err != nil { + return nil, err + } + return mappings, nil +} + func (prov *DatabaseProvider) initializeConnection(ctx context.Context, conn *stdsql.Conn) error { if prov.duckLake == nil || conn == nil { return nil } + if err := prov.duckLakeConnectionAdmissionError(ctx); err != nil { + return err + } return conn.Raw(func(driverConn any) error { execer, ok := driverConn.(driver.ExecerContext) if !ok { @@ -632,7 +2000,7 @@ func (prov *DatabaseProvider) AllDatabases(ctx *sql.Context) []sql.Database { catalogName := adapter.GetCurrentCatalog(ctx) rows, err := adapter.QueryCatalog(ctx, "SELECT DISTINCT schema_name FROM information_schema.schemata WHERE catalog_name = ?", catalogName) if err != nil { - panic(ErrDuckDB.New(err)) + return []sql.Database{} } defer rows.Close() @@ -640,7 +2008,7 @@ func (prov *DatabaseProvider) AllDatabases(ctx *sql.Context) []sql.Database { for rows.Next() { var schemaName string if err := rows.Scan(&schemaName); err != nil { - panic(ErrDuckDB.New(err)) + return []sql.Database{} } switch schemaName { @@ -648,7 +2016,7 @@ func (prov *DatabaseProvider) AllDatabases(ctx *sql.Context) []sql.Database { continue } - all = append(all, NewDatabase(schemaName, catalogName)) + all = append(all, newDatabase(schemaName, catalogName, prov)) } sort.Slice(all, func(i, j int) bool { @@ -670,7 +2038,7 @@ func (prov *DatabaseProvider) Database(ctx *sql.Context, name string) (sql.Datab } if ok { - return NewDatabase(name, catalogName), nil + return newDatabase(name, catalogName, prov), nil } return nil, sql.ErrDatabaseNotFound.New(name) } @@ -725,11 +2093,26 @@ func (prov *DatabaseProvider) DropDatabase(ctx *sql.Context, name string) error return nil } -func (prov *DatabaseProvider) Restart(readOnly bool) error { - prov.mu.Lock() - defer prov.mu.Unlock() +func (prov *DatabaseProvider) Restart(readOnly bool) (err error) { + if prov == nil { + return fmt.Errorf("database provider is unavailable") + } + // Keep the generation transition ahead of prov.mu for the same reason as + // Close: admitted frontend work must be able to finish provider catalog + // operations before pool teardown waits on its lifecycle. + releaseTransition, transitionErr := prov.beginDuckLakeGenerationTransitionWithError() + if transitionErr != nil { + return transitionErr + } + defer func() { + err = errors.Join(err, releaseTransition()) + }() + if prov.mu != nil { + prov.mu.Lock() + defer prov.mu.Unlock() + } - err := prov.Close() + err = prov.closeLocked() if err != nil { return err } diff --git a/catalog/provider_ducklake_test.go b/catalog/provider_ducklake_test.go index 20c8be14..a8f3aad8 100644 --- a/catalog/provider_ducklake_test.go +++ b/catalog/provider_ducklake_test.go @@ -3,13 +3,76 @@ package catalog import ( "context" stdsql "database/sql" + "database/sql/driver" "testing" + "time" + "github.com/apecloud/myduckserver/adapter" + "github.com/apecloud/myduckserver/configuration" "github.com/apecloud/myduckserver/mycontext" + "github.com/dolthub/go-mysql-server/memory" + "github.com/dolthub/go-mysql-server/sql" "github.com/duckdb/duckdb-go/v2" "github.com/stretchr/testify/require" ) +// poolBindingTestSession exposes the production pool binding through a GMS +// session so the regression can exercise the exact adapter lookup that used +// to deadlock when a lifecycle lease was held during shutdown. +type poolBindingTestSession struct { + *memory.Session + pool *ConnectionPool + id uint32 +} + +var _ adapter.ConnectionHolder = (*poolBindingTestSession)(nil) +var _ adapter.TransactionBindingHolder = (*poolBindingTestSession)(nil) + +func (s *poolBindingTestSession) GetConn(ctx context.Context) (*stdsql.Conn, error) { + return s.pool.GetConn(ctx, s.id) +} + +func (s *poolBindingTestSession) GetTxn(ctx context.Context, options *stdsql.TxOptions) (*stdsql.Tx, error) { + return s.pool.GetTxn(ctx, s.id, "", options) +} + +func (s *poolBindingTestSession) GetCatalogConn(ctx context.Context) (*stdsql.Conn, error) { + return s.pool.GetConn(ctx, s.id) +} + +func (s *poolBindingTestSession) GetCatalogTxn(ctx context.Context, options *stdsql.TxOptions) (*stdsql.Tx, error) { + return s.pool.GetTxn(ctx, s.id, "", options) +} + +func (s *poolBindingTestSession) TryGetTxn() *stdsql.Tx { + return s.pool.TryGetTxn(s.id) +} + +func (s *poolBindingTestSession) GetTxnBinding() (*stdsql.Conn, *stdsql.Tx) { + return s.pool.GetTxnBinding(s.id) +} + +func (s *poolBindingTestSession) GetCurrentCatalog() string { return "memory" } + +func (s *poolBindingTestSession) GetCurrentSchema() string { return "main" } + +func (s *poolBindingTestSession) CloseTxn() { s.pool.CloseTxn(s.id) } + +func (s *poolBindingTestSession) CloseConn() { _ = s.pool.CloseConn(s.id) } + +type recordingDuckLakeExecer struct { + queries []string + err error +} + +func (e *recordingDuckLakeExecer) ExecContext(_ context.Context, query string, _ []driver.NamedValue) (driver.Result, error) { + e.queries = append(e.queries, query) + if e.err != nil { + return nil, e.err + } + return driver.RowsAffected(0), nil +} + func TestStorageIsPassiveAndDoesNotInventOrigin(t *testing.T) { connector, err := duckdb.NewConnector("", nil) require.NoError(t, err) @@ -55,3 +118,141 @@ func TestDuckDBStringLiteralRoundTripsQuotesAndBackslashes(t *testing.T) { require.NoError(t, db.QueryRow("SELECT "+duckDBStringLiteral(want)).Scan(&got)) require.Equal(t, want, got) } + +func TestDuckLakeAttachUsesServicePaths(t *testing.T) { + runtime := &duckLakeRuntime{config: configuration.DuckLakeConfig{ + MetadataPath: "/var/lib/myduck/catalog.ducklake", + DataPath: "s3://test-bucket/data", + }} + execer := &recordingDuckLakeExecer{} + + require.NoError(t, runtime.attachLocked(context.Background(), nil, execer)) + require.Equal(t, []string{ + "ATTACH IF NOT EXISTS 'ducklake:/var/lib/myduck/catalog.ducklake' AS \"__myduck_ducklake\" (DATA_PATH 's3://test-bucket/data', DATA_INLINING_ROW_LIMIT 0, CREATE_IF_NOT_EXISTS true)", + }, execer.queries) +} + +func TestDuckLakeAttachRejectsRemoteCatalogURI(t *testing.T) { + runtime := &duckLakeRuntime{config: configuration.DuckLakeConfig{ + MetadataPath: "s3://test-bucket/catalog.ducklake", + DataPath: "s3://test-bucket/data", + }} + execer := &recordingDuckLakeExecer{} + + err := runtime.attachLocked(context.Background(), nil, execer) + require.Error(t, err) + require.Contains(t, err.Error(), "reason=invalid_configuration") + require.Empty(t, execer.queries) +} + +func TestDuckLakeObjectBoundaryRejectsIncompletePaths(t *testing.T) { + runtime := &duckLakeRuntime{config: configuration.DuckLakeConfig{ + MetadataPath: "/var/lib/myduck/catalog.ducklake", + }} + execer := &recordingDuckLakeExecer{} + + err := runtime.EnsureAttached(mycontext.WithFrontendQuery(context.Background()), nil, execer) + require.ErrorContains(t, err, "metadata and data paths are required") + require.Empty(t, execer.queries) +} + +func TestDisabledDuckLakeConnectionUsesTableStorageError(t *testing.T) { + var provider *DatabaseProvider + + err := provider.EnsureDuckLakeConnectionWithExecutor(context.Background(), nil, nil) + require.ErrorIs(t, err, ErrInvalidTableStorage) + require.ErrorContains(t, err, "DuckLake service configuration is disabled") +} + +func TestDisabledDuckLakeRuntimeUsesTableStorageError(t *testing.T) { + var runtime *duckLakeRuntime + + err := runtime.EnsureAttached( + mycontext.WithFrontendQuery(context.Background()), + nil, + nil, + ) + require.ErrorIs(t, err, ErrInvalidTableStorage) + require.ErrorContains(t, err, "DuckLake service configuration is disabled") +} + +func TestRecordObjectStorageSelectionRejectsDisabledDuckLake(t *testing.T) { + database := NewDatabaseWithProvider("main", "memory", &DatabaseProvider{}) + + err := database.RecordTableStorageSelectionWithExecutor( + nil, + "disabled_object", + TableStorageSelection{Kind: TableStorageObject, Explicit: true, Source: "test"}, + nil, + nil, + ) + require.ErrorIs(t, err, ErrInvalidTableStorage) + require.ErrorContains(t, err, "DuckLake service configuration is disabled") +} + +func TestEnsureDuckLakeConnectionWithBindingDoesNotReenterLifecycleLock(t *testing.T) { + connector, err := duckdb.NewConnector("", nil) + require.NoError(t, err) + db := stdsql.OpenDB(connector) + pool := NewConnectionPool(connector, db, "memory") + t.Cleanup(func() { + require.NoError(t, pool.Close()) + require.NoError(t, connector.Close()) + }) + + const sessionID = uint32(174) + conn, tx, release, err := pool.GetExecutionSnapshotLease( + context.Background(), sessionID, "", true, + ) + require.NoError(t, err) + require.NotNil(t, conn) + require.Nil(t, tx) + t.Cleanup(release) + + // Mark setup complete so the test isolates lifecycle admission rather than + // extension loading or ATTACH SQL. + runtime := &duckLakeRuntime{config: configuration.DuckLakeConfig{ + MetadataPath: "/tmp/task77/catalog.ducklake", + DataPath: "/tmp/task77/data", + }} + require.NoError(t, conn.Raw(func(raw any) error { + physical, ok := raw.(driver.Conn) + require.True(t, ok) + runtime.initialized.Store(physical, struct{}{}) + runtime.attached.Store(physical, struct{}{}) + return nil + })) + provider := &DatabaseProvider{duckLake: runtime} + session := &poolBindingTestSession{ + Session: memory.NewSession(sql.NewBaseSession(), nil), + pool: pool, + id: sessionID, + } + ctx := sql.NewContext(mycontext.WithFrontendQuery(context.Background()), sql.WithSession(session)) + + closeDone := make(chan error, 1) + go func() { closeDone <- pool.Close() }() + waitForPoolShutdown(t, pool) + + // The pool close has announced its writer and is waiting for the retained + // lease. The binding-aware setup must complete without calling + // pool.GetTxnBinding (which would queue behind that writer). + setupDone := make(chan error, 1) + go func() { + setupDone <- provider.EnsureDuckLakeConnectionWithBinding(ctx, conn, nil) + }() + select { + case setupErr := <-setupDone: + require.NoError(t, setupErr) + case <-time.After(time.Second): + t.Fatal("binding-aware DuckLake setup re-entered the pending lifecycle writer") + } + + release() + select { + case closeErr := <-closeDone: + require.NoError(t, closeErr) + case <-time.After(time.Second): + t.Fatal("pool close did not finish after execution lease release") + } +} diff --git a/catalog/provider_generation_test.go b/catalog/provider_generation_test.go new file mode 100644 index 00000000..5e62e2ed --- /dev/null +++ b/catalog/provider_generation_test.go @@ -0,0 +1,487 @@ +package catalog + +import ( + "context" + stdsql "database/sql" + "errors" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/DATA-DOG/go-sqlmock" + "github.com/apecloud/myduckserver/configuration" + "github.com/apecloud/myduckserver/mycontext" + "github.com/stretchr/testify/require" +) + +// newMalformedGenerationProvider creates a real provider/pool pairing and then +// corrupts only the logical transaction marker. The sql.Tx and its physical +// owner remain live, which is the state pool teardown must preserve. +func newMalformedGenerationProvider(t *testing.T, sessionID uint32) (*DatabaseProvider, *ConnectionPool, *stdsql.Tx, *stdsql.Conn, sqlmock.Sqlmock) { + t.Helper() + db, mock, err := sqlmock.New() + require.NoError(t, err) + p := NewConnectionPool(nil, db, "memory") + p.registerMySQLUDFsOnce.Do(func() {}) + provider := &DatabaseProvider{ + mu: &sync.RWMutex{}, + pool: p, + storage: db, + duckLake: &duckLakeRuntime{config: configuration.DuckLakeConfig{ + MetadataPath: "/tmp/task77/catalog.ducklake", + DataPath: "/tmp/task77/data", + }}, + } + p.SetTransactionLifecycleHooksWithError(provider.beginDuckLakeTransactionWithError, provider.untrackDuckLakeTransaction) + p.SetTransactionAdmissionGuard(provider.duckLakeGenerationPoisonError) + p.SetConnectionAdmissionGuard(provider.duckLakeConnectionAdmissionError) + + mock.ExpectBegin() + tx, err := p.GetTxn(context.Background(), sessionID, "", nil) + require.NoError(t, err) + ownerEntry, ok := p.txnConns.Load(sessionID) + require.True(t, ok) + owner, ok := ownerEntry.(*stdsql.Conn) + require.True(t, ok) + require.NotNil(t, owner) + + p.txns.Store(sessionID, "malformed transaction marker") + t.Cleanup(func() { + // A failed teardown intentionally leaves the owner available for an + // explicit recovery. Keep cleanup bounded even when an assertion above + // aborts the test before the normal repair below. + p.txns.Store(sessionID, tx) + _ = tx.Rollback() + p.CloseTxn(sessionID) + _ = owner.Close() + _ = db.Close() + }) + return provider, p, tx, owner, mock +} + +func newPoisonedAdmissionProvider(t *testing.T) (*DatabaseProvider, *ConnectionPool, sqlmock.Sqlmock) { + t.Helper() + db, mock, err := sqlmock.New() + require.NoError(t, err) + p := NewConnectionPool(nil, db, "memory") + provider := &DatabaseProvider{ + mu: &sync.RWMutex{}, + pool: p, + storage: db, + duckLake: &duckLakeRuntime{config: configuration.DuckLakeConfig{ + MetadataPath: "/tmp/task77/catalog.ducklake", + DataPath: "/tmp/task77/data", + }}, + } + p.SetTransactionLifecycleHooksWithError(provider.beginDuckLakeTransactionWithError, provider.untrackDuckLakeTransaction) + p.SetTransactionAdmissionGuard(provider.duckLakeGenerationPoisonError) + p.SetConnectionAdmissionGuard(provider.duckLakeConnectionAdmissionError) + provider.poisonDuckLakeGeneration(errors.New("malformed generation for admission test")) + t.Cleanup(func() { _ = db.Close() }) + return provider, p, mock +} + +func TestDuckLakeProviderCloseMalformedTransactionReturnsBoundedError(t *testing.T) { + const sessionID = uint32(184) + provider, pool, tx, owner, mock := newMalformedGenerationProvider(t, sessionID) + operationRelease := provider.BeginDuckLakeOperationForOwner( + mycontext.WithFrontendQuery(context.Background()), tx, + ) + t.Cleanup(operationRelease) + + closeDone := make(chan error, 1) + go func() { closeDone <- provider.Close() }() + select { + case closeErr := <-closeDone: + require.ErrorContains(t, closeErr, "active transaction mapping is invalid") + case <-time.After(time.Second): + t.Fatal("provider Close hung on a malformed transaction mapping") + } + + // The unknown physical owner and malformed marker are deliberately retained; + // closing either would risk blocking database/sql on an untyped live owner. + currentOwner, ok := pool.txnConns.Load(sessionID) + require.True(t, ok) + require.Same(t, owner, currentOwner) + currentMarker, ok := pool.txns.Load(sessionID) + require.True(t, ok) + require.Equal(t, "malformed transaction marker", currentMarker) + require.Equal(t, 1, duckLakeOperationCount(provider)) + require.ErrorIs(t, provider.duckLakeGenerationPoisonError(), errDuckLakeGenerationPoisoned) + + // Repair the marker and finish the owner explicitly so sqlmock can verify + // that no implicit/unknown-owner rollback was attempted by provider.Close. + pool.txns.Store(sessionID, tx) + mock.ExpectRollback() + require.NoError(t, tx.Rollback()) + pool.CloseTxn(sessionID) + operationRelease() + require.Zero(t, duckLakeOperationCount(provider)) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestDuckLakeProviderRestartMalformedTransactionReturnsBoundedError(t *testing.T) { + const sessionID = uint32(185) + provider, pool, tx, owner, mock := newMalformedGenerationProvider(t, sessionID) + operationRelease := provider.BeginDuckLakeOperationForOwner( + mycontext.WithFrontendQuery(context.Background()), tx, + ) + t.Cleanup(operationRelease) + + restartDone := make(chan error, 1) + go func() { restartDone <- provider.Restart(true) }() + select { + case restartErr := <-restartDone: + require.ErrorContains(t, restartErr, "active transaction mapping is invalid") + case <-time.After(time.Second): + t.Fatal("provider Restart hung on a malformed transaction mapping") + } + + currentOwner, ok := pool.txnConns.Load(sessionID) + require.True(t, ok) + require.Same(t, owner, currentOwner) + currentMarker, ok := pool.txns.Load(sessionID) + require.True(t, ok) + require.Equal(t, "malformed transaction marker", currentMarker) + require.Equal(t, 1, duckLakeOperationCount(provider)) + require.ErrorIs(t, provider.duckLakeGenerationPoisonError(), errDuckLakeGenerationPoisoned) + + pool.txns.Store(sessionID, tx) + mock.ExpectRollback() + require.NoError(t, tx.Rollback()) + pool.CloseTxn(sessionID) + operationRelease() + require.Zero(t, duckLakeOperationCount(provider)) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestDuckLakePoisonedGenerationRejectsNewPhysicalTransaction(t *testing.T) { + provider, pool, mock := newPoisonedAdmissionProvider(t) + + _, err := pool.GetTxn(context.Background(), 771, "", nil) + require.ErrorIs(t, err, errDuckLakeGenerationPoisoned) + _, mapped := pool.txns.Load(uint32(771)) + require.False(t, mapped, "poisoned generation must not create a transaction mapping") + require.NoError(t, mock.ExpectationsWereMet()) + _ = provider +} + +func TestDuckLakePoisonedGenerationRejectsConnectionAndSnapshotAdmission(t *testing.T) { + provider, pool, mock := newPoisonedAdmissionProvider(t) + var initializerCalls atomic.Int32 + pool.SetConnectionInitializer(func(context.Context, *stdsql.Conn) error { + initializerCalls.Add(1) + return nil + }) + ctx := mycontext.WithFrontendQuery(context.Background()) + + conn, err := pool.GetConn(ctx, 772) + require.Nil(t, conn) + require.ErrorIs(t, err, errDuckLakeGenerationPoisoned) + _, mapped := pool.conns.Load(uint32(772)) + require.False(t, mapped, "poisoned connection admission must not install a mapping") + + snapshotConn, snapshotTx, release, err := pool.GetExecutionSnapshotLease(ctx, 773, "", false) + require.Nil(t, snapshotConn) + require.Nil(t, snapshotTx) + require.NotNil(t, release) + release() + require.ErrorIs(t, err, errDuckLakeGenerationPoisoned) + _, mapped = pool.conns.Load(uint32(773)) + require.False(t, mapped, "poisoned snapshot admission must not install a mapping") + + require.Empty(t, pool.CurrentSchema(772)) + require.Empty(t, pool.CurrentCatalog(772)) + require.Zero(t, initializerCalls.Load(), "poison must be reported before connection initialization") + require.NoError(t, mock.ExpectationsWereMet()) + _ = provider +} + +func TestDuckLakePoisonedGenerationRejectsDirectConnectionInitialization(t *testing.T) { + provider, pool, mock := newPoisonedAdmissionProvider(t) + conn, err := pool.DB.Conn(context.Background()) + require.NoError(t, err) + t.Cleanup(func() { _ = conn.Close() }) + + err = provider.InitializeConnection(mycontext.WithFrontendQuery(context.Background()), conn) + require.ErrorIs(t, err, errDuckLakeGenerationPoisoned) + require.ErrorIs( + t, + provider.duckLakeConnectionAdmissionError(mycontext.WithRecoveryQuery(context.Background())), + errDuckLakeGenerationPoisoned, + ) + require.ErrorIs( + t, + provider.duckLakeConnectionAdmissionError( + WithDuckLakeCleanupLease(mycontext.WithMaintenanceQuery(context.Background())), + ), + errDuckLakeGenerationPoisoned, + ) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestDuckLakeGenerationTransitionConnectionAdmission(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + pool := NewConnectionPool(nil, db, "memory") + pool.registerMySQLUDFsOnce.Do(func() {}) + provider := &DatabaseProvider{ + mu: &sync.RWMutex{}, + pool: pool, + storage: db, + duckLake: &duckLakeRuntime{config: configuration.DuckLakeConfig{ + MetadataPath: "/tmp/task77/catalog.ducklake", + DataPath: "/tmp/task77/data", + }}, + } + pool.SetConnectionAdmissionGuard(provider.duckLakeConnectionAdmissionError) + var initializerCalls atomic.Int32 + pool.SetConnectionInitializer(func(context.Context, *stdsql.Conn) error { + initializerCalls.Add(1) + return nil + }) + + provider.duckLakeTxnMu.Lock() + provider.ensureDuckLakeTxnStateLocked() + provider.duckLakeCleanupActive = true + provider.duckLakeTxnMu.Unlock() + defer func() { + provider.duckLakeTxnMu.Lock() + provider.duckLakeCleanupActive = false + provider.duckLakeTxnCond.Broadcast() + provider.duckLakeTxnMu.Unlock() + _ = db.Close() + }() + + frontendCtx := mycontext.WithFrontendQuery(context.Background()) + maintenanceCtx := mycontext.WithMaintenanceQuery(context.Background()) + recoveryCtx := mycontext.WithRecoveryQuery(context.Background()) + cleanupCtx := WithDuckLakeCleanupLease(maintenanceCtx) + + require.ErrorIs(t, provider.duckLakeConnectionAdmissionError(frontendCtx), errDuckLakeGenerationTransition) + require.ErrorIs(t, provider.duckLakeConnectionAdmissionError(maintenanceCtx), errDuckLakeGenerationTransition) + require.ErrorIs( + t, + provider.duckLakeConnectionAdmissionError(WithDuckLakeCleanupLease(frontendCtx)), + errDuckLakeGenerationTransition, + ) + for _, origin := range []mycontext.QueryOriginKind{ + mycontext.UnknownQueryOrigin, + mycontext.InternalQueryOrigin, + mycontext.MySQLReplicationQueryOrigin, + mycontext.PostgresReplicationQueryOrigin, + } { + require.ErrorIs( + t, + provider.duckLakeConnectionAdmissionError(mycontext.WithQueryOrigin(context.Background(), origin)), + errDuckLakeGenerationTransition, + ) + } + require.NoError(t, provider.duckLakeConnectionAdmissionError(recoveryCtx)) + require.NoError(t, provider.duckLakeConnectionAdmissionError(cleanupCtx)) + + conn, err := pool.GetConn(frontendCtx, 774) + require.Nil(t, conn) + require.ErrorIs(t, err, errDuckLakeGenerationTransition) + _, mapped := pool.conns.Load(uint32(774)) + require.False(t, mapped) + + snapshotConn, snapshotTx, release, err := pool.GetExecutionSnapshotLease(frontendCtx, 775, "", false) + require.Nil(t, snapshotConn) + require.Nil(t, snapshotTx) + require.NotNil(t, release) + release() + require.ErrorIs(t, err, errDuckLakeGenerationTransition) + _, mapped = pool.conns.Load(uint32(775)) + require.False(t, mapped) + + recoveryConn, err := pool.GetConn(recoveryCtx, 776) + require.NoError(t, err) + require.NotNil(t, recoveryConn) + + cleanupConn, err := pool.GetConn(cleanupCtx, 776) + require.NoError(t, err) + require.Same(t, recoveryConn, cleanupConn) + + require.ErrorIs( + t, + provider.InitializeConnection(frontendCtx, recoveryConn), + errDuckLakeGenerationTransition, + ) + require.ErrorIs( + t, + provider.EnsureDuckLakeConnectionWithBinding(frontendCtx, recoveryConn, nil), + errDuckLakeGenerationTransition, + ) + require.ErrorIs( + t, + provider.EnsureDuckLakeConnectionWithBinding(frontendCtx, recoveryConn, new(stdsql.Tx)), + errDuckLakeGenerationTransition, + "a transaction pointer is not an admission token during transition", + ) + require.Empty(t, pool.CurrentSchema(774)) + require.Empty(t, pool.CurrentCatalog(774)) + require.Equal(t, int32(2), initializerCalls.Load()) + require.NoError(t, pool.CloseConn(776)) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestConnectionAdmissionGuardRechecksAfterLifecycleAdmission(t *testing.T) { + testAdmission := func(t *testing.T, acquire func(*ConnectionPool) error) { + t.Helper() + db, mock, err := sqlmock.New() + require.NoError(t, err) + pool := NewConnectionPool(nil, db, "memory") + pool.registerMySQLUDFsOnce.Do(func() {}) + var guardCalls atomic.Int32 + var initializerCalls atomic.Int32 + pool.SetConnectionAdmissionGuard(func(context.Context) error { + if guardCalls.Add(1) == 2 { + return errDuckLakeGenerationTransition + } + return nil + }) + pool.SetConnectionInitializer(func(context.Context, *stdsql.Conn) error { + initializerCalls.Add(1) + return nil + }) + + require.ErrorIs(t, acquire(pool), errDuckLakeGenerationTransition) + require.Equal(t, int32(2), guardCalls.Load()) + require.Zero(t, initializerCalls.Load()) + _, mapped := pool.conns.Load(uint32(778)) + require.False(t, mapped) + require.NoError(t, mock.ExpectationsWereMet()) + _ = db.Close() + } + + t.Run("connection", func(t *testing.T) { + testAdmission(t, func(pool *ConnectionPool) error { + conn, err := pool.GetConn(mycontext.WithFrontendQuery(context.Background()), 778) + require.Nil(t, conn) + return err + }) + }) + t.Run("execution snapshot", func(t *testing.T) { + testAdmission(t, func(pool *ConnectionPool) error { + conn, tx, release, err := pool.GetExecutionSnapshotLease( + mycontext.WithFrontendQuery(context.Background()), 778, "", false, + ) + require.Nil(t, conn) + require.Nil(t, tx) + require.NotNil(t, release) + release() + return err + }) + }) +} + +func TestDuckLakeGenerationTransitionRejectsActiveTransactionExecutionButAllowsReleaseLookup(t *testing.T) { + pool := NewConnectionPool(nil, nil, "memory") + provider := &DatabaseProvider{duckLake: &duckLakeRuntime{config: configuration.DuckLakeConfig{ + MetadataPath: "/tmp/task77/catalog.ducklake", + DataPath: "/tmp/task77/data", + }}} + pool.SetConnectionAdmissionGuard(provider.duckLakeConnectionAdmissionError) + + const sessionID uint32 = 779 + conn := new(stdsql.Conn) + tx := new(stdsql.Tx) + pool.conns.Store(sessionID, conn) + pool.txns.Store(sessionID, tx) + pool.txnConns.Store(sessionID, conn) + + provider.duckLakeTxnMu.Lock() + provider.ensureDuckLakeTxnStateLocked() + provider.duckLakeCleanupActive = true + provider.duckLakeTxnMu.Unlock() + defer func() { + provider.duckLakeTxnMu.Lock() + provider.duckLakeCleanupActive = false + provider.duckLakeTxnCond.Broadcast() + provider.duckLakeTxnMu.Unlock() + }() + + gotConn, gotTx, err := pool.GetTxnWithBinding( + mycontext.WithFrontendQuery(context.Background()), + sessionID, + "", + nil, + ) + require.Nil(t, gotConn) + require.Nil(t, gotTx) + require.ErrorIs(t, err, errDuckLakeGenerationTransition) + + releaseConn, releaseTx := pool.GetTxnBindingForRelease(sessionID) + require.Same(t, conn, releaseConn) + require.Same(t, tx, releaseTx) +} + +func TestActiveTransactionAdmissionGuardRechecksAfterLifecycleAdmission(t *testing.T) { + pool := NewConnectionPool(nil, nil, "memory") + const sessionID uint32 = 780 + conn := new(stdsql.Conn) + tx := new(stdsql.Tx) + pool.conns.Store(sessionID, conn) + pool.txns.Store(sessionID, tx) + pool.txnConns.Store(sessionID, conn) + + var guardCalls atomic.Int32 + pool.SetConnectionAdmissionGuard(func(context.Context) error { + if guardCalls.Add(1) == 2 { + return errDuckLakeGenerationTransition + } + return nil + }) + + gotConn, gotTx, err := pool.GetTxnWithBinding(context.Background(), sessionID, "", nil) + require.Nil(t, gotConn) + require.Nil(t, gotTx) + require.ErrorIs(t, err, errDuckLakeGenerationTransition) + require.Equal(t, int32(2), guardCalls.Load()) + + releaseConn, releaseTx := pool.GetTxnBindingForRelease(sessionID) + require.Same(t, conn, releaseConn) + require.Same(t, tx, releaseTx) +} + +func TestDuckLakePoisonedGenerationRejectsOperationAndCleanup(t *testing.T) { + provider, pool, mock := newPoisonedAdmissionProvider(t) + + release, err := provider.BeginDuckLakeOperationForOwnerWithError( + mycontext.WithFrontendQuery(context.Background()), "owner", + ) + require.ErrorIs(t, err, errDuckLakeGenerationPoisoned) + release() + require.Zero(t, duckLakeOperationCount(provider)) + + conn, err := pool.DB.Conn(context.Background()) + require.NoError(t, err) + cleanupErr := provider.CleanupDuckLakeOrphansOnConn( + mycontext.WithMaintenanceQuery(context.Background()), conn, + ) + require.ErrorIs(t, cleanupErr, errDuckLakeGenerationPoisoned) + require.NoError(t, conn.Close()) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestDuckLakePoisonedGenerationRestartFailsClosed(t *testing.T) { + provider, pool, mock := newPoisonedAdmissionProvider(t) + + restartDone := make(chan error, 1) + go func() { restartDone <- provider.Restart(true) }() + select { + case restartErr := <-restartDone: + require.ErrorIs(t, restartErr, errDuckLakeGenerationPoisoned) + case <-time.After(time.Second): + t.Fatal("Restart hung while generation was poisoned") + } + var transactionMappings int + pool.txns.Range(func(_, _ any) bool { + transactionMappings++ + return true + }) + require.Zero(t, transactionMappings) + require.NoError(t, mock.ExpectationsWereMet()) +} diff --git a/catalog/rollback_cleanup_test.go b/catalog/rollback_cleanup_test.go new file mode 100644 index 00000000..5ea40a63 --- /dev/null +++ b/catalog/rollback_cleanup_test.go @@ -0,0 +1,349 @@ +package catalog + +import ( + "context" + stdsql "database/sql" + "database/sql/driver" + "strings" + "sync" + "testing" + "time" + + "github.com/DATA-DOG/go-sqlmock" + "github.com/apecloud/myduckserver/adapter" + "github.com/apecloud/myduckserver/configuration" + "github.com/apecloud/myduckserver/mycontext" + "github.com/dolthub/go-mysql-server/memory" + "github.com/dolthub/go-mysql-server/sql" + "github.com/dolthub/go-mysql-server/sql/types" + "github.com/duckdb/duckdb-go/v2" + "github.com/stretchr/testify/require" +) + +// transactionTestSession supplies only the connection-holder surface needed by +// rowInserter while retaining a real GMS session for sql.Context. +type transactionTestSession struct { + *memory.Session + conn *stdsql.Conn + tx *stdsql.Tx +} + +var _ adapter.ConnectionHolder = (*transactionTestSession)(nil) +var _ adapter.TransactionBindingHolder = (*transactionTestSession)(nil) + +func (s *transactionTestSession) GetConn(context.Context) (*stdsql.Conn, error) { + return s.conn, nil +} + +func (s *transactionTestSession) GetTxn(context.Context, *stdsql.TxOptions) (*stdsql.Tx, error) { + return s.tx, nil +} + +func (s *transactionTestSession) GetCatalogConn(context.Context) (*stdsql.Conn, error) { + return s.conn, nil +} + +func (s *transactionTestSession) GetCatalogTxn(context.Context, *stdsql.TxOptions) (*stdsql.Tx, error) { + return s.tx, nil +} + +func (s *transactionTestSession) TryGetTxn() *stdsql.Tx { + return s.tx +} + +func (s *transactionTestSession) GetTxnBinding() (*stdsql.Conn, *stdsql.Tx) { + return s.conn, s.tx +} + +func (s *transactionTestSession) GetCurrentCatalog() string { + return "memory" +} + +func (s *transactionTestSession) GetCurrentSchema() string { + return "main" +} + +func (s *transactionTestSession) CloseTxn() { + s.tx = nil +} + +func (s *transactionTestSession) CloseConn() {} + +func TestRowInserterUsesActiveTransactionForObjectStyleDML(t *testing.T) { + ctx := context.Background() + connector, err := duckdb.NewConnector("", nil) + require.NoError(t, err) + db := stdsql.OpenDB(connector) + t.Cleanup(func() { + require.NoError(t, db.Close()) + require.NoError(t, connector.Close()) + }) + + conn, err := db.Conn(ctx) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, conn.Close()) }) + _, err = conn.ExecContext(ctx, `CREATE TABLE "object_like" (id INTEGER, payload VARCHAR)`) + require.NoError(t, err) + + tx, err := conn.BeginTx(ctx, nil) + require.NoError(t, err) + session := &transactionTestSession{ + Session: memory.NewSession(sql.NewBaseSession(), nil), + conn: conn, + tx: tx, + } + session.SetCurrentDatabase("memory") + sqlCtx := sql.NewContext(ctx, sql.WithSession(session)) + + ri := &rowInserter{ + catalog: "memory", + db: "main", + table: "object_like", + schema: sql.Schema{ + &sql.Column{Name: "id", Type: types.Int32, Nullable: false}, + &sql.Column{Name: "payload", Type: types.Text, Nullable: true}, + }, + } + require.NoError(t, ri.Insert(sqlCtx, sql.Row{int32(7), "uncommitted"})) + require.NoError(t, ri.Close(sqlCtx)) + + var inTx int + require.NoError(t, tx.QueryRowContext(ctx, `SELECT count(*) FROM "object_like"`).Scan(&inTx)) + require.Equal(t, 1, inTx) + require.NoError(t, tx.Rollback()) + + var afterRollback int + require.NoError(t, conn.QueryRowContext(ctx, `SELECT count(*) FROM "object_like"`).Scan(&afterRollback)) + require.Zero(t, afterRollback) +} + +func TestCleanupDuckLakeOrphansUsesCanonicalQueryAndClosesRows(t *testing.T) { + var actualQuery string + matcher := sqlmock.QueryMatcherFunc(func(_ string, actual string) error { + actualQuery = strings.Join(strings.Fields(actual), " ") + return nil + }) + db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(matcher)) + require.NoError(t, err) + t.Cleanup(func() { _ = db.Close() }) + conn, err := db.Conn(context.Background()) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, conn.Close()) }) + + runtime := &duckLakeRuntime{config: configuration.DuckLakeConfig{ + MetadataPath: "/tmp/task77/catalog.ducklake", + DataPath: "/tmp/task77/data", + }} + provider := &DatabaseProvider{duckLake: runtime, storage: db} + // Mark the mock physical connection as already initialized and attached; + // this isolates the cleanup SQL from extension loading while retaining the + // real *sql.Conn/Raw path used in production. + require.NoError(t, conn.Raw(func(raw any) error { + physical, ok := raw.(driver.Conn) + require.True(t, ok) + runtime.initialized.Store(physical, struct{}{}) + runtime.attached.Store(physical, struct{}{}) + return nil + })) + + mock.ExpectQuery("cleanup").WillReturnRows( + sqlmock.NewRows([]string{"path"}).AddRow("orphan-a.parquet").AddRow("orphan-b.parquet"), + ).RowsWillBeClosed() + + cleanupCtx := mycontext.WithMaintenanceQuery(context.Background()) + require.NoError(t, provider.CleanupDuckLakeOrphansOnConn(cleanupCtx, conn)) + require.Equal(t, strings.Join(strings.Fields(duckLakeOrphanCleanupSQL), " "), actualQuery) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestCleanupDuckLakeOrphansNoopsForExtensionOnlyConfiguration(t *testing.T) { + provider := &DatabaseProvider{duckLake: &duckLakeRuntime{}} + require.NoError(t, provider.CleanupDuckLakeOrphansOnConn( + mycontext.WithMaintenanceQuery(context.Background()), nil, + )) +} + +func TestCleanupDuckLakeOrphansRejectsActiveSessionTransaction(t *testing.T) { + db, _, err := sqlmock.New() + require.NoError(t, err) + t.Cleanup(func() { _ = db.Close() }) + conn, err := db.Conn(context.Background()) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, conn.Close()) }) + + runtime := &duckLakeRuntime{config: configuration.DuckLakeConfig{ + MetadataPath: "/tmp/task77/catalog.ducklake", + DataPath: "/tmp/task77/data", + }} + session := &transactionTestSession{ + Session: memory.NewSession(sql.NewBaseSession(), nil), + conn: conn, + tx: new(stdsql.Tx), + } + ctx := sql.NewContext(mycontext.WithMaintenanceQuery(context.Background()), sql.WithSession(session)) + provider := &DatabaseProvider{duckLake: runtime, storage: db} + + err = provider.CleanupDuckLakeOrphansOnConn(ctx, conn) + require.Error(t, err) + require.Contains(t, err.Error(), "inactive transaction") +} + +func TestCleanupDuckLakeOrphansPrechecksActiveTransactionBeforeBarrier(t *testing.T) { + connector, err := duckdb.NewConnector("", nil) + require.NoError(t, err) + db := stdsql.OpenDB(connector) + p := NewConnectionPool(connector, db, "memory") + provider := &DatabaseProvider{duckLake: &duckLakeRuntime{config: configuration.DuckLakeConfig{ + MetadataPath: "/tmp/task77/catalog.ducklake", + DataPath: "/tmp/task77/data", + }}} + // Use the same lifecycle hooks installed by a live provider. The active + // transaction must be visible to the provider barrier for this regression to + // catch the self-wait that occurred when the barrier ran first. + p.SetTransactionLifecycleHooks(provider.beginDuckLakeTransaction, provider.untrackDuckLakeTransaction) + t.Cleanup(func() { + require.NoError(t, p.Close()) + require.NoError(t, connector.Close()) + }) + + tx, err := p.GetTxn(context.Background(), 77, "", nil) + require.NoError(t, err) + conn, gotTx := p.GetTxnBinding(77) + require.Same(t, tx, gotTx) + require.NotNil(t, conn) + + session := &transactionTestSession{ + Session: memory.NewSession(sql.NewBaseSession(), nil), + conn: conn, + tx: tx, + } + ctx := sql.NewContext(mycontext.WithFrontendQuery(context.Background()), sql.WithSession(session)) + + result := make(chan error, 1) + go func() { + result <- provider.CleanupDuckLakeOrphansOnConn(ctx, conn) + }() + select { + case err := <-result: + require.Error(t, err) + require.Contains(t, err.Error(), "inactive transaction") + case <-time.After(time.Second): + t.Fatal("cleanup waited on the caller's active transaction") + } + + _, inactive, rollbackErr := p.RollbackTxn(77, tx) + require.True(t, inactive) + require.NoError(t, rollbackErr) +} + +func TestWithoutCancelContextRetainsSQLContext(t *testing.T) { + ctx := sql.NewContext(mycontext.WithFrontendQuery(context.Background())) + derived := withoutCancelContext(ctx) + require.IsType(t, (*sql.Context)(nil), derived) + require.Equal(t, mycontext.FrontendQueryOrigin, mycontext.QueryOrigin(derived)) +} + +func TestDuckLakeTransactionBarrierTracksPoolLifecycle(t *testing.T) { + connector, err := duckdb.NewConnector("", nil) + require.NoError(t, err) + db := stdsql.OpenDB(connector) + p := NewConnectionPool(connector, db, "memory") + provider := &DatabaseProvider{duckLake: &duckLakeRuntime{config: configuration.DuckLakeConfig{ + MetadataPath: "/tmp/task77/catalog.ducklake", + DataPath: "/tmp/task77/data", + }}} + p.SetTransactionLifecycleHooks(provider.beginDuckLakeTransaction, provider.untrackDuckLakeTransaction) + t.Cleanup(func() { + require.NoError(t, p.Close()) + require.NoError(t, connector.Close()) + }) + + tx, err := p.GetTxn(context.Background(), 1, "", nil) + require.NoError(t, err) + provider.duckLakeTxnMu.Lock() + _, tracked := provider.duckLakeActiveTx[tx] + provider.duckLakeTxnMu.Unlock() + require.True(t, tracked) + + // Cleanup waits for the active writer, while a repeated GetTxn call on that + // same session remains immediately usable and does not join the wait. + cleanupDone := make(chan struct{}) + go func() { + release := provider.beginDuckLakeCleanup() + release() + close(cleanupDone) + }() + select { + case <-cleanupDone: + t.Fatal("cleanup ran while a transaction was still active") + case <-time.After(20 * time.Millisecond): + } + got, err := p.GetTxn(context.Background(), 1, "", nil) + require.NoError(t, err) + require.Same(t, tx, got) + require.NoError(t, tx.Rollback()) + p.CloseTxn(1) + select { + case <-cleanupDone: + case <-time.After(time.Second): + t.Fatal("cleanup did not resume after transaction finalization") + } +} + +func TestDuckLakeTransactionAdmissionBlocksNewWriterDuringCleanup(t *testing.T) { + provider := &DatabaseProvider{duckLake: &duckLakeRuntime{config: configuration.DuckLakeConfig{ + MetadataPath: "/tmp/task77/catalog.ducklake", + DataPath: "/tmp/task77/data", + }}} + release := provider.beginDuckLakeCleanup() + + started := make(chan struct{}) + finished := make(chan struct{}) + var tx *stdsql.Tx + go func() { + close(started) + complete := provider.beginDuckLakeTransaction() + complete(tx) + close(finished) + }() + <-started + select { + case <-finished: + t.Fatal("new transaction was admitted during cleanup") + case <-time.After(20 * time.Millisecond): + } + release() + select { + case <-finished: + case <-time.After(time.Second): + t.Fatal("transaction admission did not resume after cleanup") + } +} + +func TestDuckLakeTransactionBarrierSupportsConcurrentCleanupCallers(t *testing.T) { + provider := &DatabaseProvider{duckLake: &duckLakeRuntime{config: configuration.DuckLakeConfig{ + MetadataPath: "/tmp/task77/catalog.ducklake", + DataPath: "/tmp/task77/data", + }}} + firstRelease := provider.beginDuckLakeCleanup() + + var wg sync.WaitGroup + started := make(chan struct{}, 1) + finished := make(chan struct{}, 1) + wg.Add(1) + go func() { + defer wg.Done() + close(started) + release := provider.beginDuckLakeCleanup() + release() + close(finished) + }() + <-started + select { + case <-finished: + t.Fatal("second cleanup entered before first released") + case <-time.After(20 * time.Millisecond): + } + firstRelease() + wg.Wait() +} diff --git a/catalog/table.go b/catalog/table.go index a5fadab1..8e9e3f12 100644 --- a/catalog/table.go +++ b/catalog/table.go @@ -32,6 +32,26 @@ type ExtraTableInfo struct { Replicated bool Sequence string Checks []sql.CheckDefinition + // Storage records the table's durable storage class. An empty value is + // treated as local for metadata written before table-level storage + // selection existed; newly-created tables always write the explicit value. + Storage TableStorageKind `json:"storage,omitempty"` +} + +// StorageKind returns the effective storage class for this metadata. Missing +// storage metadata is the backwards-compatible local-table behavior. +func (info ExtraTableInfo) StorageKind() TableStorageKind { + if info.Storage == TableStorageObject { + return TableStorageObject + } + return TableStorageLocal +} + +func (info *ExtraTableInfo) normalizeStorage() { + if info == nil { + return + } + info.Storage = info.StorageKind() } type ColumnInfo struct { @@ -75,6 +95,10 @@ func NewTable(db *Database, name string, hasPrimaryKey bool) *Table { } func (t *Table) withComment(comment *Comment[ExtraTableInfo]) *Table { + if comment == nil { + comment = NewComment[ExtraTableInfo]("") + } + comment.Meta.normalizeStorage() t.comment = comment return t } @@ -84,7 +108,18 @@ func (t *Table) withSchema(ctx *sql.Context) error { if err != nil { return err } + return t.applySchema(schema) +} +func (t *Table) withSchemaWithExecutor(ctx *sql.Context, execer adapter.SQLExecutor) error { + schema, err := getPKSchemaWithExecutor(ctx, execer, t.db.catalog, t.db.name, t.name) + if err != nil { + return err + } + return t.applySchema(schema) +} + +func (t *Table) applySchema(schema sql.PrimaryKeySchema) error { t.schema = schema // https://github.com/apecloud/myduckserver/issues/272 @@ -100,7 +135,12 @@ func (t *Table) withSchema(ctx *sql.Context) error { } func (t *Table) ExtraTableInfo() ExtraTableInfo { - return t.comment.Meta + if t.comment == nil { + return ExtraTableInfo{Storage: TableStorageLocal} + } + info := t.comment.Meta + info.normalizeStorage() + return info } func (t *Table) HasPrimaryKey() bool { @@ -135,8 +175,17 @@ func (t *Table) Schema(_ *sql.Context) sql.Schema { // RowCount implements sql.StatisticsTable. func (t *Table) RowCount(ctx *sql.Context) (uint64, bool, error) { var n uint64 - q := `SELECT COUNT(*) FROM ` + FullTableName(t.db.catalog, t.db.name, t.name) - if err := adapter.QueryRowCatalog(ctx, q).Scan(&n); err != nil { + execer, conn, _, release, err := adapter.GetExecutionSnapshotWithLease(ctx) + if err != nil { + return 0, false, err + } + defer release() + physical, err := t.physicalTableNameWithExecutor(ctx, execer, conn) + if err != nil { + return 0, false, err + } + q := `SELECT COUNT(*) FROM ` + physical + if err := execer.QueryRowContext(ctx, q).Scan(&n); err != nil { return 0, false, ErrDuckDB.New(err) } return n, true, nil @@ -177,9 +226,18 @@ func (t *Table) DataLength(ctx *sql.Context) (uint64, error) { } func getPKSchema(ctx *sql.Context, catalogName, dbName, tableName string) (sql.PrimaryKeySchema, error) { + execer, _, _, release, err := adapter.GetCatalogExecutionSnapshotWithLease(ctx) + if err != nil { + return sql.PrimaryKeySchema{}, err + } + defer release() + return getPKSchemaWithExecutor(ctx, execer, catalogName, dbName, tableName) +} + +func getPKSchemaWithExecutor(ctx *sql.Context, execer adapter.SQLExecutor, catalogName, dbName, tableName string) (sql.PrimaryKeySchema, error) { var schema sql.Schema - columns, err := queryColumns(ctx, catalogName, dbName, tableName) + columns, err := queryColumnsWithExecutor(ctx, execer, catalogName, dbName, tableName) if err != nil { return sql.PrimaryKeySchema{}, ErrDuckDB.New(err) } @@ -230,7 +288,7 @@ func getPKSchema(ctx *sql.Context, catalogName, dbName, tableName string) (sql.P } // Add primary key columns to the schema - primaryKeyOrdinals := getPrimaryKeyOrdinals(ctx, catalogName, dbName, tableName) + primaryKeyOrdinals := getPrimaryKeyOrdinalsWithExecutor(ctx, execer, catalogName, dbName, tableName) setPrimaryKeyColumns(schema, primaryKeyOrdinals) return sql.NewPrimaryKeySchema(schema, primaryKeyOrdinals...), nil @@ -253,7 +311,16 @@ func (t *Table) PrimaryKeySchema(_ *sql.Context) sql.PrimaryKeySchema { } func getPrimaryKeyOrdinals(ctx *sql.Context, catalogName, dbName, tableName string) []int { - rows, err := adapter.QueryCatalog(ctx, ` + execer, _, _, release, err := adapter.GetCatalogExecutionSnapshotWithLease(ctx) + if err != nil { + panic(ErrDuckDB.New(err)) + } + defer release() + return getPrimaryKeyOrdinalsWithExecutor(ctx, execer, catalogName, dbName, tableName) +} + +func getPrimaryKeyOrdinalsWithExecutor(ctx *sql.Context, execer adapter.SQLExecutor, catalogName, dbName, tableName string) []int { + rows, err := execer.QueryContext(ctx, ` SELECT constraint_column_indexes FROM duckdb_constraints() WHERE ((database_name = ? AND schema_name = ? AND table_name = ?) OR (database_name = 'temp' AND schema_name = 'main' AND table_name = ?)) AND constraint_type = 'PRIMARY KEY' LIMIT 1 `, catalogName, dbName, tableName, tableName) if err != nil { @@ -287,6 +354,9 @@ func getCreateSequence(temporary bool, sequenceName string) (createStmt, fullNam func (t *Table) AddColumn(ctx *sql.Context, column *sql.Column, order *sql.ColumnOrder) error { t.mu.Lock() defer t.mu.Unlock() + if t.objectStorage() { + return t.addObjectColumn(ctx, column) + } // TODO: Column order is ignored as DuckDB does not support it. @@ -375,6 +445,9 @@ func (t *Table) AddColumn(ctx *sql.Context, column *sql.Column, order *sql.Colum func (t *Table) DropColumn(ctx *sql.Context, columnName string) error { t.mu.Lock() defer t.mu.Unlock() + if t.objectStorage() { + return t.dropObjectColumn(ctx, columnName) + } // Check if the column is AUTO_INCREMENT autoIncrement := false @@ -413,6 +486,9 @@ func (t *Table) DropColumn(ctx *sql.Context, columnName string) error { func (t *Table) ModifyColumn(ctx *sql.Context, columnName string, column *sql.Column, order *sql.ColumnOrder) error { t.mu.Lock() defer t.mu.Unlock() + if t.objectStorage() { + return t.modifyObjectColumn(ctx, columnName, column) + } typ, err := DuckdbDataType(column.Type) if err != nil { @@ -580,10 +656,12 @@ func (t *Table) Updater(ctx *sql.Context) sql.RowUpdater { // Inserter implements sql.InsertableTable. func (t *Table) Inserter(*sql.Context) sql.RowInserter { return &rowInserter{ - db: t.db.Name(), - table: t.name, - schema: t.schema.Schema, - hasPK: t.hasPrimaryKey, + catalog: t.db.catalog, + db: t.db.Name(), + table: t.name, + schema: t.schema.Schema, + hasPK: t.hasPrimaryKey, + provider: t.db.provider, } } @@ -594,7 +672,16 @@ func (t *Table) Deleter(*sql.Context) sql.RowDeleter { // Truncate implements sql.TruncateableTable. func (t *Table) Truncate(ctx *sql.Context) (int, error) { - result, err := adapter.ExecCatalog(ctx, `TRUNCATE TABLE `+FullTableName(t.db.catalog, t.db.name, t.name)) + execer, conn, _, release, err := adapter.GetExecutionSnapshotWithLease(ctx) + if err != nil { + return 0, err + } + defer release() + physical, err := t.physicalTableNameWithExecutor(ctx, execer, conn) + if err != nil { + return 0, err + } + result, err := execer.ExecContext(ctx, `TRUNCATE TABLE `+physical) if err != nil { return 0, err } @@ -606,11 +693,13 @@ func (t *Table) Truncate(ctx *sql.Context) (int, error) { func (t *Table) Replacer(*sql.Context) sql.RowReplacer { hasKey := len(t.schema.PkOrdinals) > 0 || !sql.IsKeyless(t.schema.Schema) return &rowInserter{ - db: t.db.Name(), - table: t.name, - schema: t.schema.Schema, - hasPK: t.hasPrimaryKey, - replace: hasKey, + catalog: t.db.catalog, + db: t.db.Name(), + table: t.name, + schema: t.schema.Schema, + hasPK: t.hasPrimaryKey, + provider: t.db.provider, + replace: hasKey, } } @@ -619,6 +708,9 @@ func (t *Table) CreateIndex(ctx *sql.Context, indexDef sql.IndexDef) error { // Lock the table to ensure thread-safety during index creation t.mu.Lock() defer t.mu.Unlock() + if t.objectStorage() { + return fmt.Errorf("%w: object tables do not support indexes", ErrInvalidTableStorage) + } // https://github.com/apecloud/myduckserver/issues/272 if isIndexCreationDisabled(ctx) { @@ -694,6 +786,9 @@ func (t *Table) CreateIndex(ctx *sql.Context, indexDef sql.IndexDef) error { func (t *Table) DropIndex(ctx *sql.Context, indexName string) error { t.mu.Lock() defer t.mu.Unlock() + if t.objectStorage() { + return fmt.Errorf("%w: object tables do not support indexes", ErrInvalidTableStorage) + } // Construct the SQL statement for dropping the index // DuckDB requires switching context to the schema by USE statement @@ -720,9 +815,27 @@ func (t *Table) RenameIndex(ctx *sql.Context, fromIndexName string, toIndexName func (t *Table) GetIndexes(ctx *sql.Context) ([]sql.Index, error) { t.mu.RLock() defer t.mu.RUnlock() + if t.objectStorage() { + return nil, nil + } + + execer, _, _, release, err := adapter.GetCatalogExecutionSnapshotWithLease(ctx) + if err != nil { + return nil, ErrDuckDB.New(err) + } + defer release() + + columnsInfo, err := queryColumnsWithExecutor(ctx, execer, t.db.catalog, t.db.name, t.name) + if err != nil { + return nil, ErrDuckDB.New(err) + } + columnsInfoMap := make(map[string]*ColumnInfo) + for _, columnInfo := range columnsInfo { + columnsInfoMap[columnInfo.ColumnName] = columnInfo + } // Query to get the indexes for the table - rows, err := adapter.QueryCatalog(ctx, `SELECT index_name, is_unique, comment, sql FROM duckdb_indexes() WHERE (database_name = ? AND schema_name = ? AND table_name = ?) or (database_name = 'temp' AND schema_name = 'main' AND table_name = ?)`, + rows, err := execer.QueryContext(ctx, `SELECT index_name, is_unique, comment, sql FROM duckdb_indexes() WHERE (database_name = ? AND schema_name = ? AND table_name = ?) or (database_name = 'temp' AND schema_name = 'main' AND table_name = ?)`, t.db.catalog, t.db.name, t.name, t.name) if err != nil { return nil, ErrDuckDB.New(err) @@ -741,16 +854,6 @@ func (t *Table) GetIndexes(ctx *sql.Context) ([]sql.Index, error) { indexes = append(indexes, NewIndex(t.db.name, t.name, "PRIMARY", true, NewComment[IndexMeta](""), pkExprs)) } - columnsInfo, err := queryColumns(ctx, t.db.catalog, t.db.name, t.name) - columnsInfoMap := make(map[string]*ColumnInfo) - for _, columnInfo := range columnsInfo { - columnsInfoMap[columnInfo.ColumnName] = columnInfo - } - - if err != nil { - return nil, ErrDuckDB.New(err) - } - for rows.Next() { var encodedIndexName string var comment stdsql.NullString @@ -817,7 +920,16 @@ func (t *Table) ModifyComment(ctx *sql.Context, text string) error { } func queryColumns(ctx *sql.Context, catalogName, schemaName, tableName string) ([]*ColumnInfo, error) { - rows, err := adapter.QueryCatalog(ctx, ` + execer, _, _, release, err := adapter.GetCatalogExecutionSnapshotWithLease(ctx) + if err != nil { + return nil, err + } + defer release() + return queryColumnsWithExecutor(ctx, execer, catalogName, schemaName, tableName) +} + +func queryColumnsWithExecutor(ctx *sql.Context, execer adapter.SQLExecutor, catalogName, schemaName, tableName string) ([]*ColumnInfo, error) { + rows, err := execer.QueryContext(ctx, ` SELECT column_name, column_index, data_type, is_nullable, column_default, comment, numeric_precision, numeric_scale FROM duckdb_columns() WHERE (database_name = ? AND schema_name = ? AND table_name = ?) OR (database_name = 'temp' AND schema_name = 'main' AND table_name = ?) @@ -1059,6 +1171,9 @@ func (t *Table) GetChecks(ctx *sql.Context) ([]sql.CheckDefinition, error) { func (t *Table) CreateCheck(ctx *sql.Context, check *sql.CheckDefinition) error { t.mu.Lock() defer t.mu.Unlock() + if t.objectStorage() { + return fmt.Errorf("%w: object tables do not support CHECK constraints", ErrInvalidTableStorage) + } // TODO(fan): Implement this once DuckDB supports modifying check constraints. // https://duckdb.org/docs/sql/statements/alter_table.html#add--drop-constraint @@ -1073,6 +1188,9 @@ func (t *Table) CreateCheck(ctx *sql.Context, check *sql.CheckDefinition) error func (t *Table) DropCheck(ctx *sql.Context, checkName string) error { t.mu.Lock() defer t.mu.Unlock() + if t.objectStorage() { + return fmt.Errorf("%w: object tables do not support CHECK constraints", ErrInvalidTableStorage) + } checks := make([]sql.CheckDefinition, 0, max(len(t.comment.Meta.Checks)-1, 0)) found := false diff --git a/catalog/table_storage.go b/catalog/table_storage.go new file mode 100644 index 00000000..42d07c2a --- /dev/null +++ b/catalog/table_storage.go @@ -0,0 +1,423 @@ +package catalog + +import ( + "context" + "errors" + "fmt" + "sort" + "strings" + "unicode" + + "github.com/dolthub/go-mysql-server/sql" +) + +// TableStorageKind is the durable storage class for a user table. Resolver +// functions use local when no selector is present; object is the DuckLake- +// backed class selected by the 0.3 table-storage slice. +type TableStorageKind string + +const ( + TableStorageLocal TableStorageKind = "local" + TableStorageObject TableStorageKind = "object" + + // TableStorageDuckLake is an explicit spelling for callers that think in + // terms of the underlying implementation. It has the same durable value as + // TableStorageObject. + TableStorageDuckLake = TableStorageObject + // TableStorageObjectStorage is retained as a descriptive alias for API + // callers; object is the only persisted value. + TableStorageObjectStorage = TableStorageObject +) + +const ( + // TableStorageOptionName is the protocol-neutral PostgreSQL storage + // parameter. MySQL clients use ENGINE=DUCKLAKE because Vitess accepts only + // its built-in table-option grammar. + TableStorageOptionName = "myduck_storage" + TableStorageObjectValue = "object" + TableStorageLocalValue = "local" + TableStorageMySQLEngine = "ducklake" +) + +var ( + ErrInvalidTableStorage = errors.New("invalid table storage selection") + ErrTableStorageConflict = errors.New("conflicting table storage selections") + ErrTableStorageDuplicate = errors.New("duplicate table storage selection") + ErrTableStorageCredentialSQL = errors.New("table storage credentials must come from service configuration") + ErrTableStoragePathSQL = errors.New("table storage paths must come from service configuration") +) + +// TableStorageSelection is the normalized result of a CREATE TABLE storage +// selector. Explicit is false only when no selector was supplied, in which case +// Kind is always TableStorageLocal. +type TableStorageSelection struct { + Kind TableStorageKind + Explicit bool + Source string +} + +func DefaultTableStorageSelection() TableStorageSelection { + return TableStorageSelection{Kind: TableStorageLocal, Source: "default"} +} + +func (s TableStorageSelection) Validate() error { + if s.Kind != TableStorageLocal && s.Kind != TableStorageObject { + return fmt.Errorf("%w: unsupported kind %q", ErrInvalidTableStorage, s.Kind) + } + return nil +} + +func (s TableStorageSelection) IsObjectStorage() bool { + return s.Kind == TableStorageObject +} + +// EffectiveKind makes the local default explicit for callers that receive a +// zero-valued selection from an optional context or metadata field. +func (s TableStorageSelection) EffectiveKind() TableStorageKind { + if s.Kind == "" { + return TableStorageLocal + } + return s.Kind +} + +// TableStorageOption is a parser-neutral name/value pair. Keeping options as a +// slice lets protocol adapters preserve duplicates and reject them before a +// map-based planner silently overwrites one declaration with another. +type TableStorageOption struct { + Name string + Value any +} + +// NormalizeTableStorageOptions resolves the selectors that protocol parsers +// expose. Ordinary MySQL/PostgreSQL table options are ignored and retain their +// historical local behavior. Only the explicit selector and recognized +// ENGINE alias affect the result. +func NormalizeTableStorageOptions(options []TableStorageOption) (TableStorageSelection, error) { + selection := DefaultTableStorageSelection() + canonicalSeen := false + canonicalKind := TableStorageLocal + canonicalName := "" + engineSeen := false + engineKind := TableStorageLocal + engineSelected := false + + for _, option := range options { + name := normalizeTableStorageOptionName(option.Name) + if name == "" { + return TableStorageSelection{}, fmt.Errorf("%w: empty option name", ErrInvalidTableStorage) + } + + if isTableStorageCredentialOption(name) { + return TableStorageSelection{}, fmt.Errorf("%w: option %q", ErrTableStorageCredentialSQL, name) + } + if isTableStoragePathOption(name) { + return TableStorageSelection{}, fmt.Errorf("%w: option %q", ErrTableStoragePathSQL, name) + } + + switch name { + case TableStorageOptionName: + if canonicalSeen { + return TableStorageSelection{}, fmt.Errorf("%w: %q and %q", ErrTableStorageDuplicate, canonicalName, name) + } + kind, err := parseTableStorageKind(option.Value) + if err != nil { + return TableStorageSelection{}, fmt.Errorf("%w: option %q: %v", ErrInvalidTableStorage, name, err) + } + canonicalSeen = true + canonicalKind = kind + canonicalName = name + case "engine": + kind, selected, err := parseMySQLEngine(option.Value) + if err != nil { + return TableStorageSelection{}, err + } + if engineSeen { + if engineSelected && selected && engineKind != kind { + return TableStorageSelection{}, fmt.Errorf("%w: ENGINE %s versus %s", ErrTableStorageConflict, engineKind, kind) + } + return TableStorageSelection{}, fmt.Errorf("%w: repeated ENGINE selector", ErrTableStorageDuplicate) + } + engineSeen = true + engineKind = kind + engineSelected = selected + if !selected { + // A normal MySQL engine (for example InnoDB) has always been + // accepted by MyDuck and maps to the local DuckDB table. + continue + } + if canonicalSeen { + if canonicalKind != engineKind { + return TableStorageSelection{}, fmt.Errorf("%w: %s versus ENGINE", ErrTableStorageConflict, canonicalKind) + } + return TableStorageSelection{}, fmt.Errorf("%w: %q and ENGINE", ErrTableStorageDuplicate, canonicalName) + } + selection.Kind = kind + selection.Explicit = true + selection.Source = "mysql-engine" + default: + // Do not silently accept a future MyDuck/DuckLake control option. + // Standard PostgreSQL WITH parameters and standard MySQL options + // remain untouched for compatibility. + if strings.HasPrefix(name, "myduck_") || strings.HasPrefix(name, "ducklake_") || isTableStorageAliasName(name) { + return TableStorageSelection{}, fmt.Errorf("%w: unsupported option %q", ErrInvalidTableStorage, name) + } + } + } + + if canonicalSeen { + if engineSeen { + if engineSelected && engineKind == canonicalKind { + return TableStorageSelection{}, fmt.Errorf("%w: %q and ENGINE", ErrTableStorageDuplicate, canonicalName) + } + if engineKind != canonicalKind { + return TableStorageSelection{}, fmt.Errorf("%w: %s versus ENGINE", ErrTableStorageConflict, canonicalKind) + } + } + selection.Kind = canonicalKind + selection.Explicit = true + selection.Source = "storage-option" + } + if err := selection.Validate(); err != nil { + return TableStorageSelection{}, err + } + return selection, nil +} + +// ResolveMySQLTableStorage consumes the map produced by go-mysql-server's +// planner. The planner has already removed duplicate keys, so protocol adapters +// that need duplicate detection should call NormalizeTableStorageOptions on a +// token-preserving slice first. +func ResolveMySQLTableStorage(options map[string]interface{}) (TableStorageSelection, error) { + if len(options) == 0 { + return DefaultTableStorageSelection(), nil + } + keys := make([]string, 0, len(options)) + for key := range options { + keys = append(keys, key) + } + sort.Strings(keys) + normalized := make([]TableStorageOption, 0, len(keys)) + for _, key := range keys { + normalized = append(normalized, TableStorageOption{Name: key, Value: options[key]}) + } + return NormalizeTableStorageOptions(normalized) +} + +// ResolvePostgresStorageParams consumes literal values extracted from +// Cockroach's tree.CreateTable.StorageParams. Values may retain SQL quoting; +// the strict value parser removes one matching quote pair. +func ResolvePostgresStorageParams(params map[string]string) (TableStorageSelection, error) { + if len(params) == 0 { + return DefaultTableStorageSelection(), nil + } + keys := make([]string, 0, len(params)) + for key := range params { + keys = append(keys, key) + } + sort.Strings(keys) + options := make([]TableStorageOption, 0, len(keys)) + for _, key := range keys { + options = append(options, TableStorageOption{Name: key, Value: params[key]}) + } + return NormalizeTableStorageOptions(options) +} + +// ResolvePostgresStorageOptions is the duplicate-preserving form for callers +// that walk tree.StorageParams directly. +func ResolvePostgresStorageOptions(options []TableStorageOption) (TableStorageSelection, error) { + return NormalizeTableStorageOptions(options) +} + +// tableStorageSelectionContextKey is deliberately private so a caller cannot +// accidentally collide with another context value. The selector is request +// scoped: callers set it immediately before invoking the generic table-create +// executor, and the database consumes it while creating that table. +type tableStorageSelectionContextKey struct{} + +// tableStorageSelectionParentKey lets the request-scoped value be consumed +// without discarding cancellation, query-origin, or other context values. +// SetTableStorageSelection mutates sql.Context for the GMS hook API, so the +// original context is retained explicitly and restored by Clear... +type tableStorageSelectionParentKey struct{} + +// SetTableStorageSelection attaches a validated selector to a GMS SQL +// context. It mutates only the context wrapper passed by the caller; the +// underlying request/session remains otherwise unchanged. +func SetTableStorageSelection(ctx *sql.Context, selection TableStorageSelection) error { + if ctx == nil { + return fmt.Errorf("%w: nil SQL context", ErrInvalidTableStorage) + } + if err := selection.Validate(); err != nil { + return err + } + base := ctx.Context + if base == nil { + base = context.Background() + } + parent := base + if original, ok := base.Value(tableStorageSelectionParentKey{}).(context.Context); ok && original != nil { + parent = original + } + withParent := context.WithValue(base, tableStorageSelectionParentKey{}, parent) + ctx.Context = context.WithValue(withParent, tableStorageSelectionContextKey{}, selection) + return nil +} + +// ClearTableStorageSelection consumes the request-scoped selector installed by +// SetTableStorageSelection. This prevents a reusable sql.Context (for example +// a session transaction context) from carrying one CREATE TABLE's choice into +// a later statement. +func ClearTableStorageSelection(ctx *sql.Context) { + if ctx == nil || ctx.Context == nil { + return + } + if parent, ok := ctx.Context.Value(tableStorageSelectionParentKey{}).(context.Context); ok && parent != nil { + ctx.Context = parent + } +} + +// WithTableStorageSelection returns a copy of ctx carrying a validated table +// selector. The copy form is useful in hooks that must not mutate a context +// owned by another execution path. +func WithTableStorageSelection(ctx *sql.Context, selection TableStorageSelection) (*sql.Context, error) { + if ctx == nil { + return nil, fmt.Errorf("%w: nil SQL context", ErrInvalidTableStorage) + } + copy := ctx.WithContext(ctx.Context) + if err := SetTableStorageSelection(copy, selection); err != nil { + return nil, err + } + return copy, nil +} + +// TableStorageSelectionFromContext retrieves a selector attached by +// SetTableStorageSelection or WithTableStorageSelection. A missing or invalid +// value is reported as absent so callers can safely fall back to local. +func TableStorageSelectionFromContext(ctx *sql.Context) (TableStorageSelection, bool) { + if ctx == nil || ctx.Context == nil { + return TableStorageSelection{}, false + } + selection, ok := ctx.Value(tableStorageSelectionContextKey{}).(TableStorageSelection) + if !ok || selection.Validate() != nil { + return TableStorageSelection{}, false + } + return selection, true +} + +func normalizeTableStorageOptionName(raw string) string { + name := strings.ToLower(strings.TrimSpace(raw)) + if len(name) >= 2 { + first, last := name[0], name[len(name)-1] + if (first == '`' && last == '`') || (first == '"' && last == '"') || (first == '\'' && last == '\'') { + name = strings.TrimSpace(name[1 : len(name)-1]) + } + } + return name +} + +func normalizeTableStorageValue(value any) (string, error) { + var raw string + switch v := value.(type) { + case string: + raw = v + case []byte: + raw = string(v) + case fmt.Stringer: + raw = v.String() + default: + return "", fmt.Errorf("value must be a literal string") + } + raw = strings.TrimSpace(raw) + if len(raw) >= 2 { + first, last := raw[0], raw[len(raw)-1] + if (first == '\'' && last == '\'') || (first == '"' && last == '"') || (first == '`' && last == '`') { + raw = strings.TrimSpace(raw[1 : len(raw)-1]) + } + } + if raw == "" { + return "", fmt.Errorf("value is empty") + } + if strings.ContainsAny(raw, "\r\n\t;:/\\?#'\"") { + return "", fmt.Errorf("value must be a simple storage kind") + } + for _, r := range raw { + if unicode.IsSpace(r) || !(unicode.IsLetter(r) || unicode.IsDigit(r) || r == '_' || r == '-') { + return "", fmt.Errorf("value must be a simple storage kind") + } + } + return strings.ToLower(raw), nil +} + +func parseTableStorageKind(value any) (TableStorageKind, error) { + normalized, err := normalizeTableStorageValue(value) + if err != nil { + return "", err + } + switch normalized { + case TableStorageLocalValue: + return TableStorageLocal, nil + case TableStorageObjectValue: + return TableStorageObject, nil + default: + return "", errors.New("unknown storage kind") + } +} + +func parseMySQLEngine(value any) (TableStorageKind, bool, error) { + normalized, err := normalizeTableStorageValue(value) + if err != nil { + return "", false, fmt.Errorf("%w: ENGINE: %v", ErrInvalidTableStorage, err) + } + switch normalized { + case TableStorageMySQLEngine: + return TableStorageObject, true, nil + case "local", "duckdb": + return TableStorageLocal, true, nil + } + // These are common MySQL/Dolt engine names. MyDuck historically accepts + // them as compatibility decorations while storing the table in DuckDB. + switch normalized { + case "innodb", "myisam", "memory", "heap", "merge", "archive", "csv", "blackhole", "ndbcluster", "federated", "rocksdb", "aria", "columnstore": + return TableStorageLocal, false, nil + default: + // ENGINE is a long-standing MySQL compatibility decoration. GMS and + // MyDuck have historically accepted engines they do not implement and + // still materialized those tables locally. Preserve that behavior for + // arbitrary simple engine identifiers; DUCKLAKE remains the sole object + // storage spelling. + return TableStorageLocal, false, nil + } +} + +func isTableStorageAliasName(name string) bool { + switch name { + case "storage", "storage_kind", "storage_type", "myduck_storage_kind", "ducklake_storage": + return true + default: + return false + } +} + +func isTableStorageCredentialOption(name string) bool { + for _, part := range []string{ + "access_key", "access_key_id", "secret", "secret_key", "secret_access_key", + "credential", "credentials", "password", "token", "key_id", "session_token", + } { + if name == part || strings.HasPrefix(name, part+"_") || strings.HasSuffix(name, "_"+part) { + return true + } + } + return false +} + +func isTableStoragePathOption(name string) bool { + for _, part := range []string{ + "endpoint", "s3_endpoint", "data_path", "metadata_path", "object_path", + "lake_path", "bucket", "path", "region", "url_style", "use_ssl", "tls", + } { + if name == part || strings.HasPrefix(name, part+"_") || strings.HasSuffix(name, "_"+part) { + return true + } + } + return false +} diff --git a/catalog/table_storage_test.go b/catalog/table_storage_test.go new file mode 100644 index 00000000..9d1fdba0 --- /dev/null +++ b/catalog/table_storage_test.go @@ -0,0 +1,215 @@ +package catalog + +import ( + "errors" + "testing" + + "github.com/dolthub/go-mysql-server/sql" + "github.com/stretchr/testify/require" +) + +func TestNormalizeTableStorageOptionsDefaultsAndSelectors(t *testing.T) { + tests := []struct { + name string + options []TableStorageOption + want TableStorageSelection + }{ + { + name: "default", + want: DefaultTableStorageSelection(), + }, + { + name: "postgres object", + options: []TableStorageOption{{ + Name: TableStorageOptionName, + Value: "'object'", + }}, + want: TableStorageSelection{Kind: TableStorageObject, Explicit: true, Source: "storage-option"}, + }, + { + name: "postgres local", + options: []TableStorageOption{{ + Name: TableStorageOptionName, + Value: "local", + }}, + want: TableStorageSelection{Kind: TableStorageLocal, Explicit: true, Source: "storage-option"}, + }, + { + name: "mysql object", + options: []TableStorageOption{{Name: "ENGINE", Value: "DUCKLAKE"}}, + want: TableStorageSelection{Kind: TableStorageObject, Explicit: true, Source: "mysql-engine"}, + }, + { + name: "mysql explicit local", + options: []TableStorageOption{{Name: "engine", Value: "LOCAL"}}, + want: TableStorageSelection{Kind: TableStorageLocal, Explicit: true, Source: "mysql-engine"}, + }, + { + name: "ordinary mysql engine remains local", + options: []TableStorageOption{{Name: "engine", Value: "InnoDB"}}, + want: DefaultTableStorageSelection(), + }, + { + name: "ordinary options remain local", + options: []TableStorageOption{{Name: "comment", Value: "kept by the normal planner"}}, + want: DefaultTableStorageSelection(), + }, + { + name: "ordinary key option remains local", + options: []TableStorageOption{{Name: "key_block_size", Value: "8"}}, + want: DefaultTableStorageSelection(), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := NormalizeTableStorageOptions(tt.options) + require.NoError(t, err) + require.Equal(t, tt.want, got) + }) + } +} + +func TestResolveTableStorageProtocolMaps(t *testing.T) { + mysql, err := ResolveMySQLTableStorage(map[string]interface{}{"ENGINE": "DUCKLAKE"}) + require.NoError(t, err) + require.Equal(t, TableStorageSelection{Kind: TableStorageObject, Explicit: true, Source: "mysql-engine"}, mysql) + + postgres, err := ResolvePostgresStorageParams(map[string]string{TableStorageOptionName: "'object'"}) + require.NoError(t, err) + require.Equal(t, TableStorageSelection{Kind: TableStorageObject, Explicit: true, Source: "storage-option"}, postgres) + + local, err := ResolvePostgresStorageParams(nil) + require.NoError(t, err) + require.Equal(t, DefaultTableStorageSelection(), local) +} + +func TestNormalizeTableStorageOptionsRejectsInvalidSelectors(t *testing.T) { + tests := []struct { + name string + options []TableStorageOption + want error + }{ + { + name: "duplicate postgres selector", + options: []TableStorageOption{ + {Name: TableStorageOptionName, Value: "object"}, + {Name: TableStorageOptionName, Value: "object"}, + }, + want: ErrTableStorageDuplicate, + }, + { + name: "duplicate mysql selector", + options: []TableStorageOption{ + {Name: "engine", Value: "DUCKLAKE"}, + {Name: "engine", Value: "DUCKLAKE"}, + }, + want: ErrTableStorageDuplicate, + }, + { + name: "same value through both selectors", + options: []TableStorageOption{ + {Name: TableStorageOptionName, Value: "object"}, + {Name: "engine", Value: "DUCKLAKE"}, + }, + want: ErrTableStorageDuplicate, + }, + { + name: "same value through both selectors reversed", + options: []TableStorageOption{ + {Name: "engine", Value: "DUCKLAKE"}, + {Name: TableStorageOptionName, Value: "object"}, + }, + want: ErrTableStorageDuplicate, + }, + { + name: "conflicting selectors", + options: []TableStorageOption{ + {Name: TableStorageOptionName, Value: "object"}, + {Name: "engine", Value: "LOCAL"}, + }, + want: ErrTableStorageConflict, + }, + { + name: "unknown postgres value", + options: []TableStorageOption{{Name: TableStorageOptionName, Value: "remote"}}, + want: ErrInvalidTableStorage, + }, + { + name: "unknown selector alias", + options: []TableStorageOption{{Name: "storage_kind", Value: "object"}}, + want: ErrInvalidTableStorage, + }, + { + name: "unknown mysql engine remains local", + options: []TableStorageOption{{Name: "engine", Value: "remote"}}, + want: nil, + }, + { + name: "credential option", + options: []TableStorageOption{{Name: "myduck_access_key_id", Value: "not persisted"}}, + want: ErrTableStorageCredentialSQL, + }, + { + name: "path option", + options: []TableStorageOption{{Name: "ducklake_bucket", Value: "test-bucket"}}, + want: ErrTableStoragePathSQL, + }, + { + name: "nonliteral value", + options: []TableStorageOption{{Name: TableStorageOptionName, Value: 42}}, + want: ErrInvalidTableStorage, + }, + { + name: "injected value", + options: []TableStorageOption{{Name: TableStorageOptionName, Value: "object;DROP TABLE t"}}, + want: ErrInvalidTableStorage, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := NormalizeTableStorageOptions(tt.options) + if tt.want == nil { + require.NoError(t, err) + require.Equal(t, DefaultTableStorageSelection(), got) + return + } + require.Error(t, err) + require.Truef(t, errors.Is(err, tt.want), "error %q does not wrap %v", err, tt.want) + }) + } +} + +func TestTableStorageSelectionContext(t *testing.T) { + ctx := sql.NewEmptyContext() + object := TableStorageSelection{Kind: TableStorageObject, Explicit: true, Source: "test"} + require.NoError(t, SetTableStorageSelection(ctx, object)) + + got, ok := TableStorageSelectionFromContext(ctx) + require.True(t, ok) + require.Equal(t, object, got) + + local := TableStorageSelection{Kind: TableStorageLocal, Explicit: true, Source: "copy"} + copyCtx, err := WithTableStorageSelection(ctx, local) + require.NoError(t, err) + copyGot, copyOK := TableStorageSelectionFromContext(copyCtx) + require.True(t, copyOK) + require.Equal(t, local, copyGot) + + originalGot, originalOK := TableStorageSelectionFromContext(ctx) + require.True(t, originalOK) + require.Equal(t, object, originalGot) + + _, ok = TableStorageSelectionFromContext(sql.NewEmptyContext()) + require.False(t, ok) + require.Error(t, SetTableStorageSelection(nil, object)) + require.Error(t, SetTableStorageSelection(ctx, TableStorageSelection{})) + _, err = WithTableStorageSelection(nil, object) + require.Error(t, err) +} + +func TestTableStorageSelectionEffectiveKind(t *testing.T) { + require.Equal(t, TableStorageLocal, (TableStorageSelection{}).EffectiveKind()) + require.Equal(t, TableStorageObject, (TableStorageSelection{Kind: TableStorageObject}).EffectiveKind()) +} diff --git a/configuration/ducklake.go b/configuration/ducklake.go index 40725f14..7e6c66fd 100644 --- a/configuration/ducklake.go +++ b/configuration/ducklake.go @@ -15,6 +15,8 @@ const DefaultDuckLakeExtensionDir = "/usr/local/lib/myduckserver/duckdb-extensio const ( duckLakeEnabledEnv = "MYDUCK_DUCKLAKE_ENABLED" duckLakeExtensionDirEnv = "MYDUCK_DUCKLAKE_EXTENSION_DIR" + duckLakeMetadataPathEnv = "MYDUCK_DUCKLAKE_METADATA_PATH" + duckLakeDataPathEnv = "MYDUCK_DUCKLAKE_DATA_PATH" duckLakeEndpointEnv = "MYDUCK_DUCKLAKE_S3_ENDPOINT" duckLakeRegionEnv = "MYDUCK_DUCKLAKE_S3_REGION" duckLakeUseSSLEnv = "MYDUCK_DUCKLAKE_S3_USE_SSL" @@ -78,7 +80,15 @@ type S3Config = DuckLakeS3Config type DuckLakeConfig struct { Enabled bool ExtensionDir string - S3 DuckLakeS3Config + // MetadataPath is the direct DuckLake metadata URI/path. The runtime + // prepends the DuckLake scheme when constructing ATTACH, so callers must + // not provide a named DuckLake secret here. + MetadataPath string + // DataPath is the service-owned DATA_PATH passed to DuckLake ATTACH. It may + // be a local absolute path or an object-storage URI, but it is never read + // from table SQL. + DataPath string + S3 DuckLakeS3Config } // DuckLakeServiceConfig is an explicit alias for service configuration APIs. @@ -89,6 +99,8 @@ type DuckLakeServiceConfig = DuckLakeConfig const ( DuckLakeEnabledEnv = duckLakeEnabledEnv DuckLakeExtensionDirEnv = duckLakeExtensionDirEnv + DuckLakeMetadataPathEnv = duckLakeMetadataPathEnv + DuckLakeDataPathEnv = duckLakeDataPathEnv DuckLakeS3EndpointEnv = duckLakeEndpointEnv DuckLakeS3RegionEnv = duckLakeRegionEnv DuckLakeS3UseSSLEnv = duckLakeUseSSLEnv @@ -113,6 +125,21 @@ func (c DuckLakeConfig) Normalize() (DuckLakeConfig, error) { if !isAbsolutePath(out.ExtensionDir) { return DuckLakeConfig{}, fmt.Errorf("ducklake extension directory must be absolute") } + // Paths are optional for the first (extension-only) service slice. An + // enabled object-table request checks both values at the ATTACH boundary; + // keeping them optional here preserves existing #71 configurations. + if out.MetadataPath != "" { + if err := validateDuckLakeMetadataPath(out.MetadataPath); err != nil { + return DuckLakeConfig{}, err + } + out.MetadataPath = strings.TrimSpace(out.MetadataPath) + } + if out.DataPath != "" { + if err := validateDuckLakeDataPath(out.DataPath); err != nil { + return DuckLakeConfig{}, err + } + out.DataPath = strings.TrimSpace(out.DataPath) + } var err error out.S3.Endpoint, out.S3.UseSSL, err = normalizeEndpointAndTLS(out.S3.Endpoint, out.S3.UseSSL, out.S3.TLS, out.S3.UseSSLSet) @@ -273,7 +300,12 @@ func LoadDuckLakeConfig() (DuckLakeConfig, error) { if err != nil { return DuckLakeConfig{}, err } - c := DuckLakeConfig{Enabled: enabled, ExtensionDir: os.Getenv(duckLakeExtensionDirEnv)} + c := DuckLakeConfig{ + Enabled: enabled, + ExtensionDir: os.Getenv(duckLakeExtensionDirEnv), + MetadataPath: os.Getenv(duckLakeMetadataPathEnv), + DataPath: os.Getenv(duckLakeDataPathEnv), + } if !enabled { return c, nil } @@ -326,6 +358,79 @@ func parseBool(raw string) (bool, error) { } } +func validateDuckLakeMetadataPath(raw string) error { + path := strings.TrimSpace(raw) + if path == "" { + return nil + } + if strings.ContainsRune(path, '\x00') || strings.ContainsAny(path, "\r\n\t") { + return fmt.Errorf("ducklake metadata path is invalid") + } + if strings.HasPrefix(strings.ToLower(path), "ducklake:") { + return fmt.Errorf("ducklake metadata path must be a direct URI or file path") + } + // DuckLake treats an identifier-only value after the `ducklake:` prefix as + // the name of a DuckLake secret. MetadataPath is deliberately a direct + // service-owned path, so reject values that could be parsed as a secret + // name instead of silently changing the ATTACH meaning. + if isIdentifierOnly(path) { + return fmt.Errorf("ducklake metadata path must be a direct URI or file path") + } + if strings.Contains(path, "#") { + return fmt.Errorf("ducklake metadata path is invalid") + } + if strings.Contains(path, "://") { + u, err := url.Parse(path) + if err != nil || u.Scheme == "" || u.Host == "" || u.User != nil || u.RawQuery != "" || u.Fragment != "" { + return fmt.Errorf("ducklake metadata path is invalid") + } + if u.Scheme != "file" && u.Scheme != "http" && u.Scheme != "https" && u.Scheme != "s3" { + return fmt.Errorf("ducklake metadata path uses an unsupported URI scheme") + } + } + return nil +} + +func isIdentifierOnly(value string) bool { + if value == "" { + return false + } + for i, r := range value { + if (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') || r == '_' || (i > 0 && r >= '0' && r <= '9') { + continue + } + return false + } + return true +} + +func validateDuckLakeDataPath(raw string) error { + path := strings.TrimSpace(raw) + if path == "" { + return nil + } + if strings.ContainsRune(path, '\x00') || strings.ContainsAny(path, "\r\n\t") { + return fmt.Errorf("ducklake data path is invalid") + } + if strings.Contains(path, "#") { + return fmt.Errorf("ducklake data path is invalid") + } + if strings.Contains(path, "://") { + u, err := url.Parse(path) + if err != nil || u.Scheme == "" || u.Host == "" || u.User != nil || u.RawQuery != "" || u.Fragment != "" { + return fmt.Errorf("ducklake data path is invalid") + } + if u.Scheme != "file" && u.Scheme != "http" && u.Scheme != "https" && u.Scheme != "s3" { + return fmt.Errorf("ducklake data path uses an unsupported URI scheme") + } + return nil + } + if !isAbsolutePath(path) { + return fmt.Errorf("ducklake data path must be absolute or a URI") + } + return nil +} + // RedactedString intentionally omits all credential values and references. func (c DuckLakeConfig) RedactedString() string { if !c.Enabled { diff --git a/configuration/ducklake_test.go b/configuration/ducklake_test.go index 72caeeaf..4b51b59e 100644 --- a/configuration/ducklake_test.go +++ b/configuration/ducklake_test.go @@ -54,10 +54,11 @@ func TestDuckLakeConfigRejectsInvalidEnabledValues(t *testing.T) { AccessKeyID: "id", SecretAccessKey: "secret", }} for name, mutate := range map[string]func(*DuckLakeConfig){ - "relative extension dir": func(c *DuckLakeConfig) { c.ExtensionDir = "relative" }, - "missing endpoint": func(c *DuckLakeConfig) { c.S3.Endpoint = "" }, - "missing region": func(c *DuckLakeConfig) { c.S3.Region = "" }, - "bad url style": func(c *DuckLakeConfig) { c.S3.URLStyle = "query" }, + "relative extension dir": func(c *DuckLakeConfig) { c.ExtensionDir = "relative" }, + "missing endpoint": func(c *DuckLakeConfig) { c.S3.Endpoint = "" }, + "missing region": func(c *DuckLakeConfig) { c.S3.Region = "" }, + "bad url style": func(c *DuckLakeConfig) { c.S3.URLStyle = "query" }, + "metadata secret identifier": func(c *DuckLakeConfig) { c.MetadataPath = "lake_secret" }, } { t.Run(name, func(t *testing.T) { candidate := base diff --git a/docker/Dockerfile b/docker/Dockerfile index 913998f1..ebd8b888 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -112,6 +112,7 @@ RUN useradd --create-home --user-group --shell /bin/bash admin \ # verifies their hashes before loading them, and never downloads extensions at # runtime. COPY --chown=admin:admin --chmod=644 extensions/duckdb/v1.5.5/linux_${TARGETARCH}/ /usr/local/lib/myduckserver/duckdb-extensions/ +RUN chmod 755 /usr/local/lib/myduckserver /usr/local/lib/myduckserver/duckdb-extensions USER admin WORKDIR /home/admin diff --git a/docs/design-task77-rollback-repair.md b/docs/design-task77-rollback-repair.md new file mode 100644 index 00000000..9abb1232 --- /dev/null +++ b/docs/design-task77-rollback-repair.md @@ -0,0 +1,57 @@ +# task #77 回滚与孤儿清理修复设计 + +## 目标 + +在冻结提交 `5cfbe01d79b601c353fbecd9d0270a6489de6ef8` 之上，闭合对象表回滚的产品行为： + +1. MySQL 和 PostgreSQL 的显式 `BEGIN` 在默认 autocommit 下都创建实际的 DuckDB `*sql.Tx`。 +2. 对象表的 INSERT、UPDATE、DELETE、TRUNCATE、RETURNING 以及必要的 catalog 元数据操作复用当前 session transaction；回滚后未提交行不可见。 +3. 前端协议 `ROLLBACK` 成功后，先释放/确认底层事务已 inactive，再在同一物理连接上调用产品拥有的 DuckLake orphan cleanup。不能通过第二条池连接，也不能与活动 `*sql.Tx` 并发使用同一连接。 + +## 非目标与边界 + +- 不修改 `591bdd2b`、`5cfbe01d` 或其父提交，不 amend 旧 OID，不合并 PR #489，不构建或发布镜像/标签。 +- 不迁移旧表，不改变复制路径，不宣称 S3/MinIO 或全量对象存储 readiness；本轮产品验收限定为 fresh local bind-mount lake。 +- extension-only 或缺少 metadata/data 完整配置时 cleanup 是 no-op；不直接删除文件，删除只能由 `ducklake_delete_orphaned_files('__myduck_ducklake', cleanup_all => true)` 完成。 +- replication、recovery、maintenance 等非前端 origin 不触发用户回滚清理。 + +## 实现策略 + +- 在事务计划构建阶段先设置 `IgnoreAutoCommit`，让 GMS 的 `StartTransaction` 回调看到显式 BEGIN，而不把它误判为空的 autocommit wrapper。 +- 由 session/connection-pool 维护一条逻辑 session 到物理连接及底层 transaction 的绑定；所有对象表执行器在取 executor 时优先使用当前 transaction。 +- 回滚顺序固定为：回滚底层 DuckDB transaction -> 使 session transaction 映射失效 -> 回滚 GMS 内存 transaction -> 复用原物理连接执行 DuckLake cleanup。每一步错误都保留，cleanup 错误不能伪装成回滚成功。 +- PostgreSQL 的事务控制走 handler 直接执行，因此在 `ROLLBACK` 返回成功后调用同一连接 cleanup；MySQL 走 `backend.Session.Rollback` 的统一 hook。cleanup 上下文保留 `*sql.Context`，但切换到 service-owned maintenance origin 并去除请求取消。 +- 连接/事务收尾使用 identity-safe 的映射删除和串行化，避免旧 transaction 的 defer 清除后来建立的新 transaction；cleanup/临时表清理不得在活动 transaction 上执行。 +- 事务完成回调注册和完成使用同一把锁。完成先于注册时先写入带结果的完成标记，迟到的注册立即收到 `success=false`（或真实成功结果），不能把 operation lease 留在一个永远不会再被触发的状态。池关闭、连接淘汰和意外回滚也都算物理事务的终端事件。 +- pool finalizer 只在锁内更新事务/连接映射并收集回调，然后按“session 锁 -> lifecycle 锁 -> callback”的顺序退出。因此 callback 可以同步重入 pool 的释放和只读接口，并且依然在原 public finalizer 返回前完成。provider 代际切换的外层锁还可能持有，callback 不得重入 provider DDL、cleanup、`Close` 或 `Restart`。 +- callback 属于外部回调，可能 panic。一个 callback panic 时，pool 仍要执行同一事务的其余 callback、provider finished hook 和同一批的其余完成事件，最后把第一个 panic 原样抛回。这样既不吞异常，也不会因计数未释放而把代际切换永久卡住。 +- 完成标记只保留一个有限窗口：按完成时间懒清理过期 tombstone，并设置固定数量上限；只有首次终态派发的 owner 在 callback 和 provider finished hook 都结束后才能记录 tombstone，迟到注册只能消费结果，不能提前把仍在派发的状态变成可清理。清理只允许删除已记录的 tombstone，活动事务和仍等待 callback/finished hook 的条目不能被驱逐。窗口外的旧物理事务注册必须 fail-closed，由调用方立即走 fallback 释放路径，不能创建新的永久 pending 状态。`Close`/`Reset` 会清掉已完成缓存，但保留待处理或 malformed 状态供收尾。 +- 普通连接、execution snapshot、当前 schema/catalog 和连接初始化统一经过 provider 的 connection-admission guard。guard 在进入 pool lifecycle 前检查一次，拿到 lifecycle/session 锁后再检查一次，关闭切换窗口。poison 时所有新连接工作都拒绝；正常代际切换期间只允许已预约的 rollback cleanup 和显式 recovery 上下文使用已知物理连接，普通前台和 maintenance 请求都失败关闭。 +- `*sql.Tx` 只表示物理事务身份，不是执行许可。创建或查询事务的方法返回后，调用者如果还要执行 SQL、准备语句、访问 `Conn.Raw` 或消费结果，必须重新取得并一直保留 execution snapshot lease；lease 要覆盖到同步调用结束、`Rows`/`Row` 扫描结束或 iterator 关闭。代际切换可以发生在两次协议请求之间，但不能在已准入的 driver 调用中间回滚或关闭它的物理 owner。 +- adapter 的直接 `Exec`/`Query`/事务执行辅助方法也遵守同一规则：同步结果在返回前释放 lease；行结果把 lease 交给包装器，在 `Close`、EOF 或 `Scan` 后释放。复制路径如果已经捕获 `(conn, tx)`，必须把同一 lease 传过整段 flush/write/LSN 更新，不能再用裸事务指针开始下一次 driver 调用。 +- FlightSQL 不再直接把 `provider.Storage()` 当作无生命周期的连接池。每个一次性请求先取得 provider operation lease 和 pool direct-connection lease；查询流消费完成后才关闭连接并释放。独立 prepared statement 和 FlightSQL transaction 在 handle 整个存活期保留同一 lease，transaction 内 prepared statement 复用其 owner，不单独关闭连接。服务关闭时先拒绝新 handle，清理遗留 prepared/transaction handle 并释放 lease，再允许 provider 关闭或重启旧代际；不能由 FlightSQL 直接关闭 provider 的共享 `*sql.DB`。 +- rollback cleanup 的 provider 屏障必须在物理 finalizer 前预约。预约失败不能跳过已有物理事务的 rollback，但必须跳过 cleanup SQL 并把预约错误返回。对没有物理 `*sql.Tx` 的 GMS-only 回滚，预约失败后不再额外获取连接。 +- COPY 的 `discardToSync` 不把网络 EOF/解码错误当成普通返回；它先中止并等待 loader、释放 snapshot/operation lease、清除 pending completion，再把原始接收错误返回给上层。这样断开连接和半截 extended exchange 不会留下 COPY 状态。 + +## 错误与并发约束 + +- rollback 返回 driver 的真实错误；只有明确的 `ErrTxDone` 或 DuckDB “no transaction is active” 才视为底层已 inactive。 +- 非预期 rollback 错误不得继续复用失效 transaction；应清除绑定并让后续请求重新获取连接/事务。 +- 同一 session 的事务建立、获取、提交、回滚和关闭串行；不同 session 的 cleanup 也不得在同一物理连接上交叉执行。 +- 代际关闭不能只等待 operation lease 后才进入 pool teardown；已完成但等待 Sync 的 COPY lease 必须能由物理事务 teardown 触发终端回调。活动 loader 仍由 snapshot lease 保持到 child 结束，随后才允许关闭其物理 owner。 +- row inserter 必须原子地捕获 `(conn, tx, executor)` 快照；旧 transaction 的临时表清理只能在其 owner 仍匹配时使用 transaction executor，否则在事务已 inactive 后使用同一连接。 +- 如果 pool teardown 发现无法解释的 transaction/owner 映射，代际切换进入不可复用的 poisoned 终态：保留原始映射和未知物理 owner，不尝试类型不安全的 rollback/close；释放切换屏障并唤醒等待者，让 `Close`/`Restart` 在有界时间内返回原始错误。后续 DuckLake 事务和操作不再等待或重新进入这个代际，避免把失效状态当作新代际使用。 +- 已 poison 的代际没有 cleanup/recovery 豁免；所有新 SQL 和重开都返回同一个可观测错误。cleanup/recovery 豁免只适用于状态仍完整的正常代际切换，不能用来绕过 poison。 +- retained handle 的关闭和正在执行的 handle 操作必须串行。关闭只在该 handle 的同步调用或查询流退出后释放 generation lease；并发的 commit/rollback、prepared close 或服务 shutdown 不能让旧连接在仍有 driver 调用时重新进入池 teardown。 + +## 验收 + +先运行 focused unit tests、`gofmt`、`git diff --check` 和带 `duckdb_arrow`/ICU 的相关包测试；然后由本负责人启动唯一 fresh local runtime，逐条保存命令、exit、raw SHA 和 before/after inventory： + +- MySQL simple wire：`BEGIN; INSERT; SELECT; ROLLBACK; SELECT`，行数应 `1 -> 2 -> 1`，未登记 parquet 应被清理且登记文件保留。 +- PostgreSQL simple wire：同一序列，行数应 `1 -> 2 -> 1`，并验证 cleanup。 +- PostgreSQL extended wire：Parse/Bind/Execute 的 INSERT/ROLLBACK 路径至少覆盖一次，确认不绕过 active transaction。 +- 错误/边界：cleanup SQL 错误可见、非前端 origin no-op、extension-only no-op、rollback 后连接可再次使用。 +- 生命周期：正常 generation transition 等待已准入的事务 driver 调用、行消费和 FlightSQL 流/handle；transition 公布后发起的新调用在任何 driver SQL 或 `Conn.Raw` 前失败。FlightSQL 服务关闭后 prepared/transaction handle 清单为空，共享 provider storage 仍只由 provider 关闭。 + +所有结果只绑定本 child 的 exact OID、parent/tree、binary/runtime 和本地 lake；未覆盖的 restart/recreate、全树 orphan、S3/MinIO 必须明确列为证据缺口。 diff --git a/docs/object-storage.md b/docs/object-storage.md new file mode 100644 index 00000000..0a82493c --- /dev/null +++ b/docs/object-storage.md @@ -0,0 +1,23 @@ +# Object storage (DuckLake) first-cut notes + +Accepted candidate `450fe6f0` replayed onto current `main`. Delivery image +must not include diagnostic inject switches. + +## Usage boundaries + +1. **MySQL CREATE is an implicit commit.** `CREATE TABLE` is not undone by + `ROLLBACK`, for local tables and object tables alike. +2. **One transaction cannot write both storages.** Mixing a local-table write + and an object-table write in the same transaction returns 1105. This is + unsupported, not a bug to paper over. +3. **The local catalog must persist.** DuckLake metadata lives in the local + catalog file (`MYDUCK_DUCKLAKE_METADATA_PATH`). The object bucket stores + table data only. Losing the catalog cannot be recovered from the bucket + alone. + +## Scope + +New tables may choose object storage. Existing tables are not migrated. +Replication is unchanged. Merge and image publication are separate steps. + +FlightSQL prepared-statement close is idempotent after execute. diff --git a/flightsqlserver/managed_lifecycle.go b/flightsqlserver/managed_lifecycle.go new file mode 100644 index 00000000..ab1b4342 --- /dev/null +++ b/flightsqlserver/managed_lifecycle.go @@ -0,0 +1,809 @@ +package flightsqlserver + +import ( + "context" + "database/sql" + "database/sql/driver" + "errors" + "fmt" + "sync" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/flight" + "github.com/apache/arrow-go/v18/arrow/flight/flightsql" + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/apecloud/myduckserver/catalog" + "github.com/apecloud/myduckserver/mycontext" + "github.com/duckdb/duckdb-go/v2" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +type arrowQuerySetup struct { + schema *arrow.Schema + err error +} + +func (s *SQLiteFlightSQLServer) beginManagedRequest(ctx context.Context) (context.Context, *managedFlightRequest, error) { + if s == nil || s.provider == nil { + return nil, nil, fmt.Errorf("FlightSQL provider lifecycle is unavailable") + } + ctx = mycontext.WithFrontendQuery(ctx) + requestCtx, cancel := context.WithCancel(ctx) + request := &managedFlightRequest{cancel: cancel, done: make(chan struct{})} + s.requestMu.Lock() + if s.closed { + s.requestMu.Unlock() + cancel() + return nil, nil, catalog.ErrExternalSessionClosed + } + if s.requests == nil { + s.requests = make(map[*managedFlightRequest]struct{}) + } + s.requests[request] = struct{}{} + s.requestMu.Unlock() + return requestCtx, request, nil +} + +func (s *SQLiteFlightSQLServer) finishManagedRequest(request *managedFlightRequest) { + if request == nil { + return + } + request.once.Do(func() { + request.cancel() + s.requestMu.Lock() + delete(s.requests, request) + s.requestMu.Unlock() + close(request.done) + }) +} + +func (s *SQLiteFlightSQLServer) openExternalSession(ctx context.Context) (*catalog.ExternalSession, error) { + if s == nil || s.provider == nil || s.provider.Pool() == nil { + return nil, fmt.Errorf("FlightSQL provider lifecycle is unavailable") + } + return s.provider.Pool().OpenExternalSession(ctx) +} + +// acquireExternalOperation enters the provider operation barrier before the +// pool lifecycle reader. This lock order lets a provider transition announce +// cleanup without deadlocking behind an operation waiting for the pool. +func (s *SQLiteFlightSQLServer) acquireExternalOperation( + ctx context.Context, + session *catalog.ExternalSession, + owner any, +) (*sql.Conn, *sql.Tx, func(), error) { + if s == nil || s.provider == nil || session == nil { + return nil, nil, func() {}, fmt.Errorf("FlightSQL external session is unavailable") + } + releaseOperation, err := s.provider.BeginDuckLakeOperationForOwnerWithError(ctx, owner) + if err != nil { + return nil, nil, func() {}, err + } + conn, tx, releaseSession, err := session.ExecutionLease(ctx) + if err != nil { + releaseOperation() + return nil, nil, func() {}, err + } + var once sync.Once + release := func() { + once.Do(func() { + releaseSession() + releaseOperation() + }) + } + return conn, tx, release, nil +} + +func requireExternalTransaction(actual, expected *sql.Tx) error { + if actual != expected { + return catalog.ErrExternalSessionTransactionMismatch + } + return nil +} + +func (s *SQLiteFlightSQLServer) managedStatelessQuery( + ctx context.Context, + query string, + args ...interface{}, +) (*arrow.Schema, <-chan flight.StreamChunk, error) { + if err := catalog.RejectSensitiveSQL(query); err != nil { + return nil, nil, err + } + ctx, request, err := s.beginManagedRequest(ctx) + if err != nil { + return nil, nil, err + } + session, err := s.openExternalSession(ctx) + if err != nil { + s.finishManagedRequest(request) + return nil, nil, err + } + conn, tx, release, err := s.acquireExternalOperation(ctx, session, session) + if err != nil { + _ = session.Close() + s.finishManagedRequest(request) + return nil, nil, err + } + if tx != nil { + release() + _ = session.Close() + s.finishManagedRequest(request) + return nil, nil, catalog.ErrExternalSessionTransactionMismatch + } + terminal := func() { + release() + _ = session.Close() + s.finishManagedRequest(request) + } + return startArrowQuery(ctx, conn, query, [][]interface{}{args}, terminal) +} + +func (s *SQLiteFlightSQLServer) managedTablesQuery( + ctx context.Context, + query string, + includeSchema bool, +) (*arrow.Schema, <-chan flight.StreamChunk, error) { + ctx, request, err := s.beginManagedRequest(ctx) + if err != nil { + return nil, nil, err + } + session, err := s.openExternalSession(ctx) + if err != nil { + s.finishManagedRequest(request) + return nil, nil, err + } + conn, tx, release, err := s.acquireExternalOperation(ctx, session, session) + if err != nil { + _ = session.Close() + s.finishManagedRequest(request) + return nil, nil, err + } + terminal := func() { + release() + _ = session.Close() + s.finishManagedRequest(request) + } + if tx != nil { + terminal() + return nil, nil, catalog.ErrExternalSessionTransactionMismatch + } + if !includeSchema { + return startArrowQuery(ctx, conn, query, [][]interface{}{nil}, terminal) + } + return startMaterializedTablesQuery(ctx, s.Alloc, conn, query, terminal) +} + +func (s *SQLiteFlightSQLServer) managedStatementUpdate(ctx context.Context, cmd flightsql.StatementUpdate) (int64, error) { + ctx, request, err := s.beginManagedRequest(ctx) + if err != nil { + return 0, err + } + defer s.finishManagedRequest(request) + + if len(cmd.GetTransactionId()) > 0 { + value, loaded := s.openTransactions.Load(string(cmd.GetTransactionId())) + state, ok := value.(*transactionState) + if !loaded || !ok || state == nil || state.external == nil || state.tx == nil { + return -1, status.Error(codes.InvalidArgument, "invalid transaction handle provided") + } + conn, tx, release, err := s.acquireExternalOperation(ctx, state.external, state.tx) + _ = conn + if err != nil { + return 0, err + } + defer release() + if state.closed.Load() { + return 0, catalog.ErrExternalSessionClosed + } + if err := requireExternalTransaction(tx, state.tx); err != nil { + return 0, err + } + result, err := tx.ExecContext(ctx, cmd.GetQuery()) + if err != nil { + return 0, err + } + return result.RowsAffected() + } + + session, err := s.openExternalSession(ctx) + if err != nil { + return 0, err + } + defer session.Close() + conn, tx, release, err := s.acquireExternalOperation(ctx, session, session) + if err != nil { + return 0, err + } + defer release() + if tx != nil { + return 0, catalog.ErrExternalSessionTransactionMismatch + } + result, err := conn.ExecContext(ctx, cmd.GetQuery()) + if err != nil { + return 0, err + } + return result.RowsAffected() +} + +func (s *SQLiteFlightSQLServer) createManagedPreparedStatement( + ctx context.Context, + req flightsql.ActionCreatePreparedStatementRequest, +) (result flightsql.ActionCreatePreparedStatementResult, err error) { + ctx, request, err := s.beginManagedRequest(ctx) + if err != nil { + return result, err + } + defer s.finishManagedRequest(request) + + handle := genRandomString() + handleKey := string(handle) + state := &Statement{query: req.GetQuery()} + if len(req.GetTransactionId()) > 0 { + value, loaded := s.openTransactions.Load(string(req.GetTransactionId())) + txnState, ok := value.(*transactionState) + if !loaded || !ok || txnState == nil || txnState.external == nil || txnState.tx == nil { + return result, status.Error(codes.InvalidArgument, "invalid transaction handle provided") + } + conn, tx, release, leaseErr := s.acquireExternalOperation(ctx, txnState.external, txnState.tx) + if leaseErr != nil { + return result, leaseErr + } + defer release() + if txnState.closed.Load() { + return result, catalog.ErrExternalSessionClosed + } + if leaseErr = requireExternalTransaction(tx, txnState.tx); leaseErr != nil { + return result, leaseErr + } + prepared, prepareErr := tx.PrepareContext(ctx, req.GetQuery()) + if prepareErr != nil { + return result, prepareErr + } + state.stmt = prepared + state.conn = conn + state.external = txnState.external + state.transaction = txnState + s.prepared.Store(handleKey, state) + } else { + session, openErr := s.openExternalSession(context.WithoutCancel(ctx)) + if openErr != nil { + return result, openErr + } + state.external = session + state.ownsConn = true + if hookErr := session.SetCleanupHook(func() { + s.cleanupManagedStatement(handleKey, state) + }); hookErr != nil { + _ = session.Close() + return result, hookErr + } + conn, tx, release, leaseErr := s.acquireExternalOperation(ctx, session, session) + if leaseErr != nil { + _ = session.Close() + return result, leaseErr + } + if state.closed.Load() || tx != nil { + release() + _ = session.Close() + if state.closed.Load() { + return result, catalog.ErrExternalSessionClosed + } + return result, catalog.ErrExternalSessionTransactionMismatch + } + prepared, prepareErr := conn.PrepareContext(ctx, req.GetQuery()) + if prepareErr != nil { + release() + _ = session.Close() + return result, prepareErr + } + state.stmt = prepared + state.conn = conn + s.prepared.Store(handleKey, state) + release() + } + + result.Handle = handle + return result, nil +} + +func (s *SQLiteFlightSQLServer) managedPreparedQuery( + ctx context.Context, + cmd flightsql.PreparedStatementQuery, +) (*arrow.Schema, <-chan flight.StreamChunk, error) { + ctx, request, err := s.beginManagedRequest(ctx) + if err != nil { + return nil, nil, err + } + value, loaded := s.prepared.Load(string(cmd.GetPreparedStatementHandle())) + stmt, ok := value.(*Statement) + if !loaded || !ok || stmt == nil || stmt.external == nil { + s.finishManagedRequest(request) + return nil, nil, status.Error(codes.InvalidArgument, "prepared statement not found") + } + if err := catalog.RejectSensitiveSQL(stmt.query); err != nil { + s.finishManagedRequest(request) + return nil, nil, err + } + owner := any(stmt.external) + var expected *sql.Tx + if stmt.transaction != nil { + expected = stmt.transaction.tx + owner = expected + } + conn, tx, release, err := s.acquireExternalOperation(ctx, stmt.external, owner) + if err != nil { + s.finishManagedRequest(request) + return nil, nil, err + } + stmt.mu.Lock() + if stmt.closed.Load() || stmt.stmt == nil || (stmt.transaction != nil && stmt.transaction.closed.Load()) { + stmt.mu.Unlock() + release() + s.finishManagedRequest(request) + return nil, nil, catalog.ErrExternalSessionClosed + } + if err := requireExternalTransaction(tx, expected); err != nil { + stmt.mu.Unlock() + release() + s.finishManagedRequest(request) + return nil, nil, err + } + query := stmt.query + argSets := make([][]interface{}, len(stmt.params)) + for i := range stmt.params { + argSets[i] = append([]interface{}(nil), stmt.params[i]...) + } + stmt.mu.Unlock() + if len(argSets) == 0 { + argSets = [][]interface{}{nil} + } + terminal := func() { + release() + s.finishManagedRequest(request) + } + return startArrowQuery(ctx, conn, query, argSets, terminal) +} + +func (s *SQLiteFlightSQLServer) managedPreparedUpdate( + ctx context.Context, + cmd flightsql.PreparedStatementUpdate, + args [][]interface{}, +) (int64, error) { + ctx, request, err := s.beginManagedRequest(ctx) + if err != nil { + return 0, err + } + defer s.finishManagedRequest(request) + value, loaded := s.prepared.Load(string(cmd.GetPreparedStatementHandle())) + stmt, ok := value.(*Statement) + if !loaded || !ok || stmt == nil || stmt.external == nil { + return 0, status.Error(codes.InvalidArgument, "prepared statement not found") + } + owner := any(stmt.external) + var expected *sql.Tx + if stmt.transaction != nil { + expected = stmt.transaction.tx + owner = expected + } + _, tx, release, err := s.acquireExternalOperation(ctx, stmt.external, owner) + if err != nil { + return 0, err + } + defer release() + stmt.mu.Lock() + defer stmt.mu.Unlock() + if stmt.closed.Load() || stmt.stmt == nil || (stmt.transaction != nil && stmt.transaction.closed.Load()) { + return 0, catalog.ErrExternalSessionClosed + } + if err := requireExternalTransaction(tx, expected); err != nil { + return 0, err + } + if len(args) == 0 { + result, err := stmt.stmt.ExecContext(ctx) + if err != nil { + return 0, err + } + return result.RowsAffected() + } + var totalAffected int64 + for _, params := range args { + result, err := stmt.stmt.ExecContext(ctx, params...) + if err != nil { + return totalAffected, err + } + affected, err := result.RowsAffected() + if err != nil { + return totalAffected, err + } + totalAffected += affected + } + return totalAffected, nil +} + +func (s *SQLiteFlightSQLServer) beginManagedTransaction(ctx context.Context) ([]byte, error) { + ctx, request, err := s.beginManagedRequest(ctx) + if err != nil { + return nil, err + } + defer s.finishManagedRequest(request) + session, err := s.openExternalSession(ctx) + if err != nil { + return nil, err + } + handle := genRandomString() + state := &transactionState{external: session, handle: string(handle)} + if err := session.SetCleanupHook(func() { + s.cleanupManagedTransaction(state) + }); err != nil { + _ = session.Close() + return nil, err + } + releaseOperation, err := s.provider.BeginDuckLakeOperationForOwnerWithError(ctx, session) + if err != nil { + _ = session.Close() + return nil, err + } + tx, beginErr := session.BeginTx(ctx, nil) + releaseOperation() + if beginErr != nil { + _ = session.Close() + return nil, beginErr + } + state.tx = tx + conn, current, release, err := s.acquireExternalOperation(ctx, session, tx) + if err != nil { + _ = s.rollbackManagedTransaction(ctx, state) + _ = session.Close() + return nil, err + } + if state.closed.Load() || current != tx { + release() + _ = s.rollbackManagedTransaction(ctx, state) + _ = session.Close() + if state.closed.Load() { + return nil, catalog.ErrExternalSessionClosed + } + return nil, catalog.ErrExternalSessionTransactionMismatch + } + state.conn = conn + s.openTransactions.Store(state.handle, state) + release() + return handle, nil +} + +func (s *SQLiteFlightSQLServer) endManagedTransaction(ctx context.Context, req flightsql.ActionEndTransactionRequest) error { + ctx, request, err := s.beginManagedRequest(ctx) + if err != nil { + return err + } + defer s.finishManagedRequest(request) + if action := req.GetAction(); action != flightsql.EndTransactionCommit && action != flightsql.EndTransactionRollback { + return status.Error(codes.InvalidArgument, "must specify Commit or Rollback to end transaction") + } + handle := string(req.GetTransactionId()) + value, loaded := s.openTransactions.LoadAndDelete(handle) + state, ok := value.(*transactionState) + if !loaded || !ok || state == nil || state.external == nil || state.tx == nil { + return status.Error(codes.InvalidArgument, "transaction id not found") + } + state.closed.Store(true) + var finalErr error + switch req.GetAction() { + case flightsql.EndTransactionCommit: + finalErr = state.external.Commit(state.tx) + case flightsql.EndTransactionRollback: + finalErr = s.rollbackManagedTransaction(ctx, state) + } + return errors.Join(finalErr, state.external.Close()) +} + +// startArrowQuery keeps database/sql.Conn.Raw active for the complete Arrow +// reader lifetime. A driver.Conn obtained through Raw is valid only inside the +// callback; returning it and consuming a reader later can cross connection and +// storage-generation teardown. +func startArrowQuery( + ctx context.Context, + conn *sql.Conn, + query string, + argSets [][]interface{}, + terminal func(), +) (*arrow.Schema, <-chan flight.StreamChunk, error) { + if terminal == nil { + terminal = func() {} + } + setup := make(chan arrowQuerySetup, 1) + chunks := make(chan flight.StreamChunk, 2) + go func() { + defer terminal() + setupSent := false + streamOwnsChannel := false + defer func() { + if recovered := recover(); recovered != nil { + err := fmt.Errorf("panic while reading FlightSQL Arrow result: %v", recovered) + if !setupSent { + setup <- arrowQuerySetup{err: err} + } else if !streamOwnsChannel { + select { + case chunks <- flight.StreamChunk{Err: err}: + case <-ctx.Done(): + } + } + } + if !streamOwnsChannel { + close(chunks) + } + }() + + err := conn.Raw(func(driverConn any) error { + raw, ok := driverConn.(driver.Conn) + if !ok { + return fmt.Errorf("unexpected DuckDB driver connection type %T", driverConn) + } + arrowDB, err := duckdb.NewArrowFromConn(raw) + if err != nil { + return err + } + if len(argSets) == 0 { + argSets = [][]interface{}{nil} + } + readers := make([]array.RecordReader, 0, len(argSets)) + for _, args := range argSets { + rdr, queryErr := arrowDB.QueryContext(ctx, query, args...) + if queryErr != nil { + for _, opened := range readers { + opened.Release() + } + return queryErr + } + readers = append(readers, rdr) + } + setup <- arrowQuerySetup{schema: readers[0].Schema()} + setupSent = true + streamOwnsChannel = true + streamArrowReaders(ctx, readers, chunks) + return nil + }) + if !setupSent { + if err == nil { + err = fmt.Errorf("FlightSQL Arrow query returned without a result") + } + setup <- arrowQuerySetup{err: err} + } + }() + result := <-setup + if result.err != nil { + return nil, nil, result.err + } + return result.schema, chunks, nil +} + +func streamArrowReaders(ctx context.Context, readers []array.RecordReader, chunks chan<- flight.StreamChunk) { + defer close(chunks) + defer func() { + for _, rdr := range readers { + rdr.Release() + } + }() + defer func() { + if recovered := recover(); recovered != nil { + select { + case chunks <- flight.StreamChunk{Err: fmt.Errorf("panic while reading FlightSQL Arrow result: %v", recovered)}: + case <-ctx.Done(): + } + } + }() + for _, rdr := range readers { + for rdr.Next() { + record := rdr.RecordBatch() + record.Retain() + select { + case chunks <- flight.StreamChunk{Data: record}: + case <-ctx.Done(): + record.Release() + return + } + } + if err := rdr.Err(); err != nil { + select { + case chunks <- flight.StreamChunk{Err: err}: + case <-ctx.Done(): + } + return + } + } +} + +// DoGetTables needs additional queries while iterating its first result. Fully +// materialize that small metadata result inside Raw, then prepare the schema +// query after Raw returns while retaining the same external operation lease. +func startMaterializedTablesQuery( + ctx context.Context, + alloc memory.Allocator, + conn *sql.Conn, + query string, + terminal func(), +) (*arrow.Schema, <-chan flight.StreamChunk, error) { + var ( + schema *arrow.Schema + records []arrow.RecordBatch + ) + err := conn.Raw(func(driverConn any) error { + raw, ok := driverConn.(driver.Conn) + if !ok { + return fmt.Errorf("unexpected DuckDB driver connection type %T", driverConn) + } + arrowDB, err := duckdb.NewArrowFromConn(raw) + if err != nil { + return err + } + rdr, err := arrowDB.QueryContext(ctx, query) + if err != nil { + return err + } + defer rdr.Release() + schema = rdr.Schema() + for rdr.Next() { + record := rdr.RecordBatch() + record.Retain() + records = append(records, record) + } + return rdr.Err() + }) + if err != nil { + for _, record := range records { + record.Release() + } + terminal() + return nil, nil, err + } + rdr, err := array.NewRecordReader(schema, records) + for _, record := range records { + record.Release() + } + if err != nil { + terminal() + return nil, nil, err + } + tables, err := newSqliteTablesSchemaBatchReader(ctx, alloc, rdr, conn, query) + if err != nil { + terminal() + return nil, nil, err + } + chunks := make(chan flight.StreamChunk, 2) + go func() { + defer terminal() + flight.StreamChunksFromReader(ctx, tables, chunks) + }() + return tables.Schema(), chunks, nil +} + +func (s *SQLiteFlightSQLServer) cleanupManagedStatement(handle string, stmt *Statement) { + if stmt == nil { + return + } + s.prepared.CompareAndDelete(handle, stmt) + _ = stmt.closePrepared() +} + +func (stmt *Statement) closePrepared() error { + if stmt == nil { + return nil + } + stmt.closeOnce.Do(func() { + stmt.closed.Store(true) + if stmt.stmt != nil { + stmt.closeErr = stmt.stmt.Close() + } + }) + return stmt.closeErr +} + +func (s *SQLiteFlightSQLServer) cleanupManagedTransaction(state *transactionState) { + if state == nil { + return + } + state.closed.Store(true) + if state.handle != "" { + s.openTransactions.CompareAndDelete(state.handle, state) + } + s.prepared.Range(func(key, value any) bool { + stmt, ok := value.(*Statement) + if !ok || stmt == nil || stmt.transaction != state { + return true + } + if s.prepared.CompareAndDelete(key, stmt) { + _ = stmt.closePrepared() + } + return true + }) +} + +func (s *SQLiteFlightSQLServer) rollbackManagedTransaction(ctx context.Context, state *transactionState) error { + if state == nil || state.external == nil || state.tx == nil { + return catalog.ErrExternalSessionTransactionMismatch + } + releaseCleanup, reservationErr := s.provider.BeginDuckLakeRollbackCleanupForOwnerWithError(state.tx, state.tx) + if releaseCleanup != nil { + defer releaseCleanup() + } + var cleanup func(*sql.Conn) error + if reservationErr == nil { + cleanup = func(conn *sql.Conn) error { + cleanupCtx := mycontext.WithMaintenanceQuery(context.WithoutCancel(ctx)) + cleanupCtx = catalog.WithDuckLakeCleanupLease(cleanupCtx) + cleanupCtx = catalog.WithRollbackCleanupOwner(cleanupCtx, conn) + return s.provider.CleanupDuckLakeOrphansOnConn(cleanupCtx, conn) + } + } + return errors.Join(reservationErr, state.external.Rollback(state.tx, cleanup)) +} + +// Close rejects new provider-owned requests, cancels active streams, and +// releases every retained prepared/transaction handle. It deliberately does +// not close provider.Storage(); the provider alone owns that shared database. +func (s *SQLiteFlightSQLServer) Close() error { + if s == nil { + return nil + } + s.requestMu.Lock() + if s.closed { + s.requestMu.Unlock() + return nil + } + s.closed = true + requests := make([]*managedFlightRequest, 0, len(s.requests)) + for request := range s.requests { + requests = append(requests, request) + } + s.requestMu.Unlock() + for _, request := range requests { + request.cancel() + } + + var closeErr error + s.prepared.Range(func(key, value any) bool { + if !s.prepared.CompareAndDelete(key, value) { + return true + } + stmt, ok := value.(*Statement) + if !ok || stmt == nil { + return true + } + if stmt.external != nil && stmt.transaction == nil { + closeErr = errors.Join(closeErr, stmt.external.Close()) + } + closeErr = errors.Join(closeErr, stmt.closePrepared()) + if stmt.external == nil && stmt.ownsConn && stmt.conn != nil { + closeErr = errors.Join(closeErr, stmt.conn.Close()) + } + return true + }) + s.openTransactions.Range(func(key, value any) bool { + if !s.openTransactions.CompareAndDelete(key, value) { + return true + } + state, ok := value.(*transactionState) + if !ok || state == nil { + return true + } + state.closed.Store(true) + if state.external != nil { + closeErr = errors.Join(closeErr, s.rollbackManagedTransaction(context.Background(), state)) + closeErr = errors.Join(closeErr, state.external.Close()) + } else { + if state.tx != nil { + closeErr = errors.Join(closeErr, state.tx.Rollback()) + } + if state.conn != nil { + closeErr = errors.Join(closeErr, state.conn.Close()) + } + } + return true + }) + for _, request := range requests { + <-request.done + } + return closeErr +} diff --git a/flightsqlserver/managed_lifecycle_test.go b/flightsqlserver/managed_lifecycle_test.go new file mode 100644 index 00000000..32cb8276 --- /dev/null +++ b/flightsqlserver/managed_lifecycle_test.go @@ -0,0 +1,314 @@ +package flightsqlserver + +import ( + "context" + "database/sql" + "errors" + "sync/atomic" + "testing" + "time" + + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/flight" + "github.com/apache/arrow-go/v18/arrow/flight/flightsql" + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/apecloud/myduckserver/catalog" + "github.com/stretchr/testify/require" +) + +type testStatementQueryTicket struct{ handle []byte } + +func (t testStatementQueryTicket) GetStatementHandle() []byte { return t.handle } + +type testGetTables struct { + tableName *string + includeSchema bool +} + +func (t testGetTables) GetCatalog() *string { return nil } +func (t testGetTables) GetDBSchemaFilterPattern() *string { return nil } +func (t testGetTables) GetTableNameFilterPattern() *string { return t.tableName } +func (t testGetTables) GetTableTypes() []string { return nil } +func (t testGetTables) GetIncludeSchema() bool { return t.includeSchema } + +func newManagedFlightSQLServer(t *testing.T) (*catalog.DatabaseProvider, *SQLiteFlightSQLServer) { + t.Helper() + provider, err := catalog.NewDBProvider("", t.TempDir(), "managed_flight") + require.NoError(t, err) + t.Cleanup(func() { + require.NoError(t, provider.Close()) + }) + server, err := NewSQLiteFlightSQLServerWithProvider(provider) + require.NoError(t, err) + t.Cleanup(func() { + require.NoError(t, server.Close()) + }) + return provider, server +} + +func managedQueryTicket(t *testing.T, query string) testStatementQueryTicket { + t.Helper() + return testStatementQueryTicket{handle: append([]byte{':'}, []byte(query)...)} +} + +func managedQueryInt64(t *testing.T, server *SQLiteFlightSQLServer, query string) int64 { + t.Helper() + _, chunks, err := server.DoGetStatement(context.Background(), managedQueryTicket(t, query)) + require.NoError(t, err) + var ( + value int64 + seen bool + ) + for chunk := range chunks { + require.NoError(t, chunk.Err) + if chunk.Data == nil { + continue + } + column, ok := chunk.Data.Column(0).(*array.Int64) + require.True(t, ok) + if column.Len() > 0 { + value = column.Value(0) + seen = true + } + chunk.Data.Release() + } + require.True(t, seen) + return value +} + +func countSyncMap(values interface{ Range(func(any, any) bool) }) int { + count := 0 + values.Range(func(any, any) bool { + count++ + return true + }) + return count +} + +func TestManagedFlightSQLUsesReplacementProviderGeneration(t *testing.T) { + provider, server := newManagedFlightSQLServer(t) + ctx := context.Background() + + _, err := server.DoPutCommandStatementUpdate(ctx, testStatementUpdate{query: "CREATE TABLE managed_restart (id BIGINT)"}) + require.NoError(t, err) + _, err = server.DoPutCommandStatementUpdate(ctx, testStatementUpdate{query: "INSERT INTO managed_restart VALUES (1)"}) + require.NoError(t, err) + require.EqualValues(t, 1, managedQueryInt64(t, server, "SELECT COUNT(*) FROM managed_restart")) + + require.NoError(t, provider.Restart(false)) + require.EqualValues(t, 1, managedQueryInt64(t, server, "SELECT COUNT(*) FROM managed_restart")) +} + +func TestManagedFlightSQLGetTablesIncludesSchemaAcrossRestart(t *testing.T) { + provider, server := newManagedFlightSQLServer(t) + ctx := context.Background() + const tableName = "managed_tables_schema" + _, err := server.DoPutCommandStatementUpdate(ctx, testStatementUpdate{ + query: "CREATE TABLE " + tableName + " (id BIGINT NOT NULL, name VARCHAR)", + }) + require.NoError(t, err) + + check := func() { + t.Helper() + resultSchema, chunks, err := server.DoGetTables(ctx, testGetTables{ + tableName: func() *string { value := tableName; return &value }(), + includeSchema: true, + }) + require.NoError(t, err) + require.Len(t, resultSchema.FieldIndices("table_schema"), 1) + + matched := 0 + for chunk := range chunks { + require.NoError(t, chunk.Err) + if chunk.Data == nil { + continue + } + tableNames := chunk.Data.Column(chunk.Data.Schema().FieldIndices("table_name")[0]).(*array.String) + tableSchemas := chunk.Data.Column(chunk.Data.Schema().FieldIndices("table_schema")[0]).(*array.Binary) + for row := 0; row < tableNames.Len(); row++ { + if tableNames.Value(row) != tableName { + continue + } + matched++ + included, decodeErr := flight.DeserializeSchema(tableSchemas.Value(row), memory.DefaultAllocator) + require.NoError(t, decodeErr) + fields := included.Fields() + require.Len(t, fields, 2) + require.Equal(t, "id", fields[0].Name) + require.Equal(t, "name", fields[1].Name) + require.False(t, fields[0].Nullable) + require.True(t, fields[1].Nullable) + } + chunk.Data.Release() + } + require.Equal(t, 1, matched) + } + + check() + require.NoError(t, provider.Restart(false)) + check() +} + +func TestManagedFlightSQLStreamKeepsRestartBehindExecutionLease(t *testing.T) { + provider, server := newManagedFlightSQLServer(t) + _, chunks, err := server.DoGetStatement( + context.Background(), + managedQueryTicket(t, "SELECT range AS id FROM range(10000000)"), + ) + require.NoError(t, err) + + restartDone := make(chan error, 1) + go func() { + restartDone <- provider.Restart(false) + }() + select { + case err := <-restartDone: + t.Fatalf("provider restart completed before the FlightSQL stream drained: %v", err) + case <-time.After(100 * time.Millisecond): + } + require.NoError(t, drainFlightChunks(chunks)) + select { + case err := <-restartDone: + require.NoError(t, err) + case <-time.After(10 * time.Second): + t.Fatal("provider restart did not complete after the FlightSQL stream drained") + } + require.EqualValues(t, 1, managedQueryInt64(t, server, "SELECT COUNT(*) FROM range(1)")) +} + +func TestManagedFlightSQLRestartInvalidatesIdlePreparedAndTransactionHandles(t *testing.T) { + provider, server := newManagedFlightSQLServer(t) + ctx := context.Background() + _, err := server.DoPutCommandStatementUpdate(ctx, testStatementUpdate{query: "CREATE TABLE managed_txn_restart (id BIGINT)"}) + require.NoError(t, err) + + standalone, err := server.CreatePreparedStatement(ctx, testCreatePreparedRequest{query: "SELECT 1"}) + require.NoError(t, err) + txnID, err := server.BeginTransaction(ctx, struct{}{}) + require.NoError(t, err) + _, err = server.DoPutCommandStatementUpdate(ctx, testStatementUpdate{ + query: "INSERT INTO managed_txn_restart VALUES (1)", + txn: txnID, + }) + require.NoError(t, err) + _, err = server.CreatePreparedStatement(ctx, testCreatePreparedRequest{ + query: "INSERT INTO managed_txn_restart VALUES (2)", + txn: txnID, + }) + require.NoError(t, err) + + require.NoError(t, provider.Restart(false)) + require.Zero(t, countSyncMap(&server.prepared)) + require.Zero(t, countSyncMap(&server.openTransactions)) + _, err = server.GetFlightInfoPreparedStatement(ctx, testPreparedCommand{handle: standalone.Handle}, &flight.FlightDescriptor{}) + require.Error(t, err) + err = server.EndTransaction(ctx, testEndTransactionRequest{txn: txnID, action: flightsql.EndTransactionCommit}) + require.Error(t, err) + require.EqualValues(t, 0, managedQueryInt64(t, server, "SELECT COUNT(*) FROM managed_txn_restart")) + + fresh, err := server.CreatePreparedStatement(ctx, testCreatePreparedRequest{query: "SELECT 2"}) + require.NoError(t, err) + require.NoError(t, server.ClosePreparedStatement(ctx, testPreparedCommand{handle: fresh.Handle})) +} + +func TestManagedFlightSQLClosePreparedStatementAfterPythonStyleDDL(t *testing.T) { + _, server := newManagedFlightSQLServer(t) + ctx := context.Background() + + drop, err := server.CreatePreparedStatement(ctx, testCreatePreparedRequest{query: "DROP TABLE IF EXISTS intTable"}) + require.NoError(t, err) + _, err = server.DoPutCommandStatementUpdate(ctx, testStatementUpdate{query: "DROP TABLE IF EXISTS intTable"}) + require.NoError(t, err) + require.NoError(t, server.ClosePreparedStatement(ctx, testPreparedCommand{handle: drop.Handle})) + require.NoError(t, server.ClosePreparedStatement(ctx, testPreparedCommand{handle: drop.Handle})) + + create, err := server.CreatePreparedStatement(ctx, testCreatePreparedRequest{ + query: "CREATE TABLE IF NOT EXISTS intTable (id INTEGER PRIMARY KEY, name VARCHAR(50), value INT)", + }) + require.NoError(t, err) + _, err = server.DoPutCommandStatementUpdate(ctx, testStatementUpdate{ + query: "CREATE TABLE IF NOT EXISTS intTable (id INTEGER PRIMARY KEY, name VARCHAR(50), value INT)", + }) + require.NoError(t, err) + require.NoError(t, server.ClosePreparedStatement(ctx, testPreparedCommand{handle: create.Handle})) + require.NoError(t, server.ClosePreparedStatement(ctx, testPreparedCommand{handle: create.Handle})) +} + +func TestManagedFlightSQLPreparedStreamBlocksRestartButIdleHandleDoesNot(t *testing.T) { + provider, server := newManagedFlightSQLServer(t) + ctx := context.Background() + prepared, err := server.CreatePreparedStatement(ctx, testCreatePreparedRequest{ + query: "SELECT range AS id FROM range(10000000)", + }) + require.NoError(t, err) + command := testPreparedCommand{handle: prepared.Handle} + _, chunks, err := server.DoGetPreparedStatement(ctx, command) + require.NoError(t, err) + + restartDone := make(chan error, 1) + go func() { + restartDone <- provider.Restart(false) + }() + select { + case err := <-restartDone: + t.Fatalf("provider restart completed before the prepared stream drained: %v", err) + case <-time.After(100 * time.Millisecond): + } + require.NoError(t, drainFlightChunks(chunks)) + select { + case err := <-restartDone: + require.NoError(t, err) + case <-time.After(10 * time.Second): + t.Fatal("provider restart did not complete after the prepared stream drained") + } + require.Zero(t, countSyncMap(&server.prepared)) + + idle, err := server.CreatePreparedStatement(ctx, testCreatePreparedRequest{query: "SELECT 1"}) + require.NoError(t, err) + restartDone = make(chan error, 1) + go func() { + restartDone <- provider.Restart(false) + }() + select { + case err := <-restartDone: + require.NoError(t, err) + case <-time.After(10 * time.Second): + t.Fatal("idle prepared handle blocked provider restart") + } + _, present := server.prepared.Load(string(idle.Handle)) + require.False(t, present) +} + +func TestManagedFlightSQLAdmissionFailurePrecedesConnectionInitialization(t *testing.T) { + provider, server := newManagedFlightSQLServer(t) + sentinel := errors.New("managed admission rejected") + var initialized atomic.Int64 + provider.Pool().SetConnectionInitializer(func(context.Context, *sql.Conn) error { + initialized.Add(1) + return nil + }) + provider.Pool().SetConnectionAdmissionGuard(func(context.Context) error { + return sentinel + }) + + _, err := server.DoPutCommandStatementUpdate(context.Background(), testStatementUpdate{query: "SELECT 1"}) + require.ErrorIs(t, err, sentinel) + require.Zero(t, initialized.Load()) +} + +func TestManagedFlightSQLCloseCleansHandlesWithoutClosingProvider(t *testing.T) { + provider, server := newManagedFlightSQLServer(t) + ctx := context.Background() + _, err := server.CreatePreparedStatement(ctx, testCreatePreparedRequest{query: "SELECT 1"}) + require.NoError(t, err) + _, err = server.BeginTransaction(ctx, struct{}{}) + require.NoError(t, err) + + require.NoError(t, server.Close()) + require.NoError(t, server.Close()) + require.Zero(t, countSyncMap(&server.prepared)) + require.Zero(t, countSyncMap(&server.openTransactions)) + require.NoError(t, provider.Storage().PingContext(ctx)) + _, err = server.DoPutCommandStatementUpdate(ctx, testStatementUpdate{query: "SELECT 1"}) + require.ErrorIs(t, err, catalog.ErrExternalSessionClosed) +} diff --git a/flightsqlserver/sqlite_server.go b/flightsqlserver/sqlite_server.go index ee79022d..9cc81426 100644 --- a/flightsqlserver/sqlite_server.go +++ b/flightsqlserver/sqlite_server.go @@ -46,6 +46,7 @@ import ( "path/filepath" "strings" "sync" + "sync/atomic" "github.com/apache/arrow-go/v18/arrow" "github.com/apache/arrow-go/v18/arrow/array" @@ -208,6 +209,7 @@ func decodeTransactionQuery(ticket []byte) (txnID, query string, err error) { } type Statement struct { + mu sync.Mutex stmt *sql.Stmt query string params [][]interface{} @@ -219,21 +221,42 @@ type Statement struct { // must remain checked out until ClosePreparedStatement so the statement and // every execution stay on the initialized physical connection. ownsConn bool + // external is set for provider-owned FlightSQL handles. A standalone + // prepared statement owns the external session; a transaction-owned + // statement shares the transaction's session and leaves its lifetime to the + // transaction finalizer. + external *catalog.ExternalSession + transaction *transactionState + closed atomic.Bool + closeOnce sync.Once + closeErr error } type transactionState struct { - tx *sql.Tx - conn *sql.Conn + tx *sql.Tx + conn *sql.Conn + external *catalog.ExternalSession + handle string + closed atomic.Bool +} + +type managedFlightRequest struct { + cancel context.CancelFunc + done chan struct{} + once sync.Once } type SQLiteFlightSQLServer struct { flightsql.BaseServer db *sql.DB - conn *duckdb.Conn initializeStorage func(context.Context, *sql.Conn) error + provider *catalog.DatabaseProvider prepared sync.Map openTransactions sync.Map + requestMu sync.Mutex + requests map[*managedFlightRequest]struct{} + closed bool } // NewSQLiteFlightSQLServer accepts an optional service initializer. Keeping @@ -250,14 +273,34 @@ func NewSQLiteFlightSQLServer(db *sql.DB, initializer ...func(context.Context, * if len(initializer) == 1 { initializeConnection = initializer[0] } - ret := &SQLiteFlightSQLServer{db: db, initializeStorage: initializeConnection} - ret.Alloc = memory.DefaultAllocator - for k, v := range SqlInfoResultMap() { - ret.RegisterSqlInfo(flightsql.SqlInfo(k), v) + ret := &SQLiteFlightSQLServer{db: db, initializeStorage: initializeConnection, requests: make(map[*managedFlightRequest]struct{})} + ret.initialize() + return ret, nil +} + +// NewSQLiteFlightSQLServerWithProvider binds FlightSQL to the provider-owned +// connection pool instead of caching the startup *sql.DB. Pool-owned external +// sessions are invalidated on Close/Reset, while each in-flight operation keeps +// its storage generation alive until synchronous work or Arrow streaming ends. +func NewSQLiteFlightSQLServerWithProvider(provider *catalog.DatabaseProvider) (*SQLiteFlightSQLServer, error) { + if provider == nil || provider.Pool() == nil { + return nil, fmt.Errorf("FlightSQL database provider is unavailable") } + ret := &SQLiteFlightSQLServer{provider: provider, requests: make(map[*managedFlightRequest]struct{})} + ret.initialize() return ret, nil } +func (s *SQLiteFlightSQLServer) initialize() { + if s == nil { + return + } + s.Alloc = memory.DefaultAllocator + for k, v := range SqlInfoResultMap() { + s.RegisterSqlInfo(flightsql.SqlInfo(k), v) + } +} + // frontendConnection marks the request before selecting a physical // connection, then initializes that same connection. A failed initializer // never leaves the connection available to another request. @@ -280,26 +323,6 @@ func (s *SQLiteFlightSQLServer) frontendConnection(ctx context.Context) (context return frontendConnection(ctx, s.db, s.initializeStorage) } -func duckDBConn(conn *sql.Conn) (*duckdb.Conn, error) { - var duckConn *duckdb.Conn - if err := conn.Raw(func(driverConn any) error { - var ok bool - duckConn, ok = driverConn.(*duckdb.Conn) - if !ok { - return fmt.Errorf("unexpected DuckDB driver connection type %T", driverConn) - } - return nil - }); err != nil { - return nil, err - } - return duckConn, nil -} - -func streamReaderWithConn(ctx context.Context, rdr array.RecordReader, conn *sql.Conn, ch chan<- flight.StreamChunk) { - defer conn.Close() - flight.StreamChunksFromReader(ctx, rdr, ch) -} - func (s *SQLiteFlightSQLServer) flightInfoForCommand(desc *flight.FlightDescriptor, schema *arrow.Schema) *flight.FlightInfo { return &flight.FlightInfo{ Endpoint: []*flight.FlightEndpoint{{Ticket: &flight.Ticket{Ticket: desc.Cmd}}}, @@ -336,6 +359,9 @@ func (s *SQLiteFlightSQLServer) DoGetStatement(ctx context.Context, cmd flightsq if txnid != "" { return nil, nil, fmt.Errorf("transactions not yet supported with DuckDB") } + if s.provider != nil { + return s.managedStatelessQuery(ctx, query) + } // var db dbQueryCtx = s.db // if txnid != "" { @@ -423,38 +449,18 @@ func (s *SQLiteFlightSQLServer) DoGetTables(ctx context.Context, cmd flightsql.G if err := catalog.RejectSensitiveSQL(query); err != nil { return nil, nil, err } - ctx, conn, err := s.frontendConnection(ctx) - if err != nil { - return nil, nil, err - } - duckConn, err := duckDBConn(conn) - if err != nil { - _ = conn.Close() - return nil, nil, err - } - arrow, err := duckdb.NewArrowFromConn(duckConn) - if err != nil { - _ = conn.Close() - return nil, nil, err + if s.provider != nil { + return s.managedTablesQuery(ctx, query, cmd.GetIncludeSchema()) } - rdr, err := arrow.QueryContext(ctx, query) + ctx, conn, err := s.frontendConnection(ctx) if err != nil { - _ = conn.Close() return nil, nil, err } - - ch := make(chan flight.StreamChunk, 2) + terminal := func() { _ = conn.Close() } if cmd.GetIncludeSchema() { - rdr, err = newSqliteTablesSchemaBatchReader(ctx, s.Alloc, rdr, conn, query) - if err != nil { - _ = conn.Close() - return nil, nil, err - } + return startMaterializedTablesQuery(ctx, s.Alloc, conn, query, terminal) } - - schema := rdr.Schema() - go streamReaderWithConn(ctx, rdr, conn, ch) - return schema, ch, nil + return startArrowQuery(ctx, conn, query, [][]interface{}{nil}, terminal) } func (s *SQLiteFlightSQLServer) GetFlightInfoXdbcTypeInfo(_ context.Context, _ flightsql.GetXdbcTypeInfo, desc *flight.FlightDescriptor) (*flight.FlightInfo, error) { @@ -481,6 +487,9 @@ func (s *SQLiteFlightSQLServer) GetFlightInfoTableTypes(_ context.Context, desc func (s *SQLiteFlightSQLServer) DoGetTableTypes(ctx context.Context) (*arrow.Schema, <-chan flight.StreamChunk, error) { query := "SELECT DISTINCT type AS table_type FROM sqlite_master" + if s.provider != nil { + return s.managedStatelessQuery(ctx, query) + } return doGetQueryWithInitializer(ctx, s.initializeStorage, s.db, query, schema_ref.TableTypes) } @@ -488,6 +497,9 @@ func (s *SQLiteFlightSQLServer) DoPutCommandStatementUpdate(ctx context.Context, if err := catalog.RejectSensitiveSQL(cmd.GetQuery()); err != nil { return 0, err } + if s.provider != nil { + return s.managedStatementUpdate(ctx, cmd) + } ctx = mycontext.WithFrontendQuery(ctx) var ( res sql.Result @@ -500,7 +512,11 @@ func (s *SQLiteFlightSQLServer) DoPutCommandStatementUpdate(ctx context.Context, return -1, status.Error(codes.InvalidArgument, "invalid transaction handle provided") } - res, err = tx.(transactionState).tx.ExecContext(ctx, cmd.GetQuery()) + state, ok := tx.(*transactionState) + if !ok || state == nil || state.tx == nil { + return -1, status.Error(codes.InvalidArgument, "invalid transaction handle provided") + } + res, err = state.tx.ExecContext(ctx, cmd.GetQuery()) } else { var conn *sql.Conn ctx, conn, err = s.frontendConnection(ctx) @@ -523,6 +539,9 @@ func (s *SQLiteFlightSQLServer) CreatePreparedStatement(ctx context.Context, req if err := catalog.RejectSensitiveSQL(query); err != nil { return result, err } + if s.provider != nil { + return s.createManagedPreparedStatement(ctx, req) + } ctx = mycontext.WithFrontendQuery(ctx) var stmtConn *sql.Conn @@ -531,7 +550,10 @@ func (s *SQLiteFlightSQLServer) CreatePreparedStatement(ctx context.Context, req if !loaded { return result, status.Error(codes.InvalidArgument, "invalid transaction handle provided") } - state := tx.(transactionState) + state, ok := tx.(*transactionState) + if !ok || state == nil || state.tx == nil { + return result, status.Error(codes.InvalidArgument, "invalid transaction handle provided") + } stmt, err = state.tx.PrepareContext(ctx, req.GetQuery()) stmtConn = state.conn } else { @@ -552,7 +574,7 @@ func (s *SQLiteFlightSQLServer) CreatePreparedStatement(ctx context.Context, req } handle := genRandomString() - s.prepared.Store(string(handle), Statement{ + s.prepared.Store(string(handle), &Statement{ stmt: stmt, query: query, conn: stmtConn, @@ -567,8 +589,14 @@ func (s *SQLiteFlightSQLServer) CreatePreparedStatement(ctx context.Context, req func (s *SQLiteFlightSQLServer) ClosePreparedStatement(ctx context.Context, request flightsql.ActionClosePreparedStatementRequest) error { handle := request.GetPreparedStatementHandle() if val, loaded := s.prepared.LoadAndDelete(string(handle)); loaded { - stmt := val.(Statement) - stmtErr := stmt.stmt.Close() + stmt, ok := val.(*Statement) + if !ok || stmt == nil { + return status.Error(codes.InvalidArgument, "prepared statement not found") + } + if stmt.external != nil && stmt.transaction == nil { + return errors.Join(stmt.external.Close(), stmt.closePrepared()) + } + stmtErr := stmt.closePrepared() if stmt.conn == nil || !stmt.ownsConn { return stmtErr } @@ -576,12 +604,16 @@ func (s *SQLiteFlightSQLServer) ClosePreparedStatement(ctx context.Context, requ return errors.Join(stmtErr, connErr) } - return status.Error(codes.InvalidArgument, "prepared statement not found") + // ADBC/Python cursors close prepared statements from __exit__/__del__ + // after execute (and sometimes after commit already dropped the handle). + // Close is idempotent. + return nil } func (s *SQLiteFlightSQLServer) GetFlightInfoPreparedStatement(_ context.Context, cmd flightsql.PreparedStatementQuery, desc *flight.FlightDescriptor) (*flight.FlightInfo, error) { - _, ok := s.prepared.Load(string(cmd.GetPreparedStatementHandle())) - if !ok { + value, ok := s.prepared.Load(string(cmd.GetPreparedStatementHandle())) + stmt, valid := value.(*Statement) + if !ok || !valid || stmt == nil || stmt.closed.Load() { return nil, status.Error(codes.InvalidArgument, "prepared statement not found") } @@ -609,34 +641,22 @@ func doGetQueryWithInitializer(ctx context.Context, initializer func(context.Con if err != nil { return nil, nil, err } - duckConn, err := duckDBConn(conn) - if err != nil { - _ = conn.Close() - return nil, nil, err - } - arrow, err := duckdb.NewArrowFromConn(duckConn) - if err != nil { - _ = conn.Close() - return nil, nil, err - } - rdr, err := arrow.QueryContext(ctx, query, args...) - if err != nil { - _ = conn.Close() - return nil, nil, err - } - schema = rdr.Schema() - ch := make(chan flight.StreamChunk) - go streamReaderWithConn(ctx, rdr, conn, ch) - return schema, ch, nil + return startArrowQuery(ctx, conn, query, [][]interface{}{args}, func() { _ = conn.Close() }) } func (s *SQLiteFlightSQLServer) DoGetPreparedStatement(ctx context.Context, cmd flightsql.PreparedStatementQuery) (schema *arrow.Schema, out <-chan flight.StreamChunk, err error) { + if s.provider != nil { + return s.managedPreparedQuery(ctx, cmd) + } val, ok := s.prepared.Load(string(cmd.GetPreparedStatementHandle())) if !ok { return nil, nil, status.Error(codes.InvalidArgument, "prepared statement not found") } - stmt := val.(Statement) + stmt, ok := val.(*Statement) + if !ok || stmt == nil || stmt.closed.Load() { + return nil, nil, status.Error(codes.InvalidArgument, "prepared statement not found") + } if err := catalog.RejectSensitiveSQL(stmt.query); err != nil { return nil, nil, err } @@ -650,50 +670,17 @@ func (s *SQLiteFlightSQLServer) DoGetPreparedStatement(ctx context.Context, cmd return nil, nil, fmt.Errorf("prepared statement has no owning connection") } - duckConn, err := duckDBConn(conn) - if err != nil { - return nil, nil, err + stmt.mu.Lock() + query := stmt.query + argSets := make([][]interface{}, len(stmt.params)) + for i := range stmt.params { + argSets[i] = append([]interface{}(nil), stmt.params[i]...) } - - arrow, err := duckdb.NewArrowFromConn(duckConn) - if err != nil { - return nil, nil, err - } - - readers := make([]array.RecordReader, 0, len(stmt.params)) - if len(stmt.params) == 0 { - rdr, err := arrow.QueryContext(ctx, stmt.query) - if err != nil { - return nil, nil, err - } - schema = rdr.Schema() - readers = append(readers, rdr) - } else { - defer func() { - if err != nil { - for _, r := range readers { - r.Release() - } - } - }() - // if we have multiple rows of bound params, execute the query - // multiple times and concatenate the result sets. - for _, p := range stmt.params { - rdr, err := arrow.QueryContext(ctx, stmt.query, p...) - if err != nil { - return nil, nil, err - } - schema = rdr.Schema() - readers = append(readers, rdr) - } + stmt.mu.Unlock() + if len(argSets) == 0 { + argSets = [][]interface{}{nil} } - ch := make(chan flight.StreamChunk) - // Keep the owning connection checked out. Standalone statements release it - // from ClosePreparedStatement; transaction statements release it when the - // transaction ends. - go flight.ConcatenateReaders(readers, ch) - out = ch - return + return startArrowQuery(ctx, conn, query, argSets, func() {}) } func scalarToIFace(s scalar.Scalar) (interface{}, error) { @@ -767,7 +754,10 @@ func (s *SQLiteFlightSQLServer) DoPutPreparedStatementQuery(_ context.Context, c return nil, status.Error(codes.InvalidArgument, "prepared statement not found") } - stmt := val.(Statement) + stmt, ok := val.(*Statement) + if !ok || stmt == nil || stmt.closed.Load() { + return nil, status.Error(codes.InvalidArgument, "prepared statement not found") + } if err := catalog.RejectSensitiveSQL(stmt.query); err != nil { return nil, err } @@ -776,8 +766,12 @@ func (s *SQLiteFlightSQLServer) DoPutPreparedStatementQuery(_ context.Context, c return nil, status.Errorf(codes.Internal, "error gathering parameters for prepared statement query: %s", err.Error()) } + stmt.mu.Lock() + defer stmt.mu.Unlock() + if stmt.closed.Load() || (stmt.transaction != nil && stmt.transaction.closed.Load()) { + return nil, status.Error(codes.InvalidArgument, "prepared statement not found") + } stmt.params = args - s.prepared.Store(string(cmd.GetPreparedStatementHandle()), stmt) return cmd.GetPreparedStatementHandle(), nil } @@ -787,7 +781,10 @@ func (s *SQLiteFlightSQLServer) DoPutPreparedStatementUpdate(ctx context.Context return 0, status.Error(codes.InvalidArgument, "prepared statement not found") } - stmt := val.(Statement) + stmt, ok := val.(*Statement) + if !ok || stmt == nil || stmt.closed.Load() { + return 0, status.Error(codes.InvalidArgument, "prepared statement not found") + } if err := catalog.RejectSensitiveSQL(stmt.query); err != nil { return 0, err } @@ -796,7 +793,19 @@ func (s *SQLiteFlightSQLServer) DoPutPreparedStatementUpdate(ctx context.Context if err != nil { return 0, status.Errorf(codes.Internal, "error gathering parameters for prepared statement: %s", err.Error()) } + if s.provider != nil { + affected, err := s.managedPreparedUpdate(ctx, cmd, args) + if err != nil && strings.Contains(err.Error(), "no such table") { + return affected, status.Error(codes.NotFound, err.Error()) + } + return affected, err + } + stmt.mu.Lock() + defer stmt.mu.Unlock() + if stmt.closed.Load() { + return 0, status.Error(codes.InvalidArgument, "prepared statement not found") + } if len(args) == 0 { result, err := stmt.stmt.ExecContext(ctx) if err != nil { @@ -853,6 +862,9 @@ func (s *SQLiteFlightSQLServer) DoGetPrimaryKeys(ctx context.Context, cmd flight fmt.Fprintf(&b, " and table_name LIKE '%s'", cmd.Table) + if s.provider != nil { + return s.managedStatelessQuery(ctx, b.String()) + } return doGetQueryWithInitializer(ctx, s.initializeStorage, s.db, b.String(), schema_ref.PrimaryKeys) } @@ -869,6 +881,9 @@ func (s *SQLiteFlightSQLServer) DoGetImportedKeys(ctx context.Context, ref fligh filter += " AND fk_schema_name = '" + *ref.DBSchema + "'" } query := prepareQueryForGetKeys(filter) + if s.provider != nil { + return s.managedStatelessQuery(ctx, query) + } return doGetQueryWithInitializer(ctx, s.initializeStorage, s.db, query, schema_ref.ImportedKeys) } @@ -885,6 +900,9 @@ func (s *SQLiteFlightSQLServer) DoGetExportedKeys(ctx context.Context, ref fligh filter += " AND pk_schema_name = '" + *ref.DBSchema + "'" } query := prepareQueryForGetKeys(filter) + if s.provider != nil { + return s.managedStatelessQuery(ctx, query) + } return doGetQueryWithInitializer(ctx, s.initializeStorage, s.db, query, schema_ref.ExportedKeys) } @@ -911,10 +929,16 @@ func (s *SQLiteFlightSQLServer) DoGetCrossReference(ctx context.Context, cmd fli filter += " AND fk_schema_name = '" + *fkref.DBSchema + "'" } query := prepareQueryForGetKeys(filter) + if s.provider != nil { + return s.managedStatelessQuery(ctx, query) + } return doGetQueryWithInitializer(ctx, s.initializeStorage, s.db, query, schema_ref.ExportedKeys) } func (s *SQLiteFlightSQLServer) BeginTransaction(ctx context.Context, req flightsql.ActionBeginTransactionRequest) (id []byte, err error) { + if s.provider != nil { + return s.beginManagedTransaction(ctx) + } ctx, conn, err := s.frontendConnection(ctx) if err != nil { return nil, err @@ -926,18 +950,27 @@ func (s *SQLiteFlightSQLServer) BeginTransaction(ctx context.Context, req flight } handle := genRandomString() - s.openTransactions.Store(string(handle), transactionState{tx: tx, conn: conn}) + s.openTransactions.Store(string(handle), &transactionState{tx: tx, conn: conn}) return handle, nil } -func (s *SQLiteFlightSQLServer) EndTransaction(_ context.Context, req flightsql.ActionEndTransactionRequest) error { +func (s *SQLiteFlightSQLServer) EndTransaction(ctx context.Context, req flightsql.ActionEndTransactionRequest) error { if req.GetAction() == flightsql.EndTransactionUnspecified { return status.Error(codes.InvalidArgument, "must specify Commit or Rollback to end transaction") } + if s.provider != nil { + if err := s.endManagedTransaction(ctx, req); err != nil { + return status.Error(codes.Internal, err.Error()) + } + return nil + } handle := string(req.GetTransactionId()) if tx, loaded := s.openTransactions.LoadAndDelete(handle); loaded { - state := tx.(transactionState) + state, ok := tx.(*transactionState) + if !ok || state == nil || state.tx == nil { + return status.Error(codes.InvalidArgument, "transaction id not found") + } txn := state.tx var txErr error switch req.GetAction() { diff --git a/flightsqlserver/sqlite_server_ducklake_test.go b/flightsqlserver/sqlite_server_ducklake_test.go index dfa500e3..088c90ea 100644 --- a/flightsqlserver/sqlite_server_ducklake_test.go +++ b/flightsqlserver/sqlite_server_ducklake_test.go @@ -246,7 +246,7 @@ func TestFlightSQLPreparedAndTransactionOperationsStayOnInitializedConnection(t queryHandle := testPreparedCommand{handle: queryResult.Handle} queryState, ok := srv.prepared.Load(string(queryResult.Handle)) require.True(t, ok) - queryStatement := queryState.(Statement) + queryStatement := queryState.(*Statement) queryConnID, err := physicalConnectionID(queryStatement.conn) require.NoError(t, err) @@ -267,7 +267,7 @@ func TestFlightSQLPreparedAndTransactionOperationsStayOnInitializedConnection(t updateHandle := testPreparedCommand{handle: updateResult.Handle} updateState, ok := srv.prepared.Load(string(updateResult.Handle)) require.True(t, ok) - updateConnID, err := physicalConnectionID(updateState.(Statement).conn) + updateConnID, err := physicalConnectionID(updateState.(*Statement).conn) require.NoError(t, err) updateParams := oneIntMessageReader(t, 7) affected, err := srv.DoPutPreparedStatementUpdate(ctx, updateHandle, updateParams) @@ -283,7 +283,7 @@ func TestFlightSQLPreparedAndTransactionOperationsStayOnInitializedConnection(t require.NoError(t, err) txnValue, ok := srv.openTransactions.Load(string(txnID)) require.True(t, ok) - txnState := txnValue.(transactionState) + txnState := txnValue.(*transactionState) txnConnID, err := physicalConnectionID(txnState.conn) require.NoError(t, err) ids, origins := recorder.snapshot() diff --git a/flightsqlserver/sqlite_tables_schema_batch_reader.go b/flightsqlserver/sqlite_tables_schema_batch_reader.go index f066799d..8c82e369 100644 --- a/flightsqlserver/sqlite_tables_schema_batch_reader.go +++ b/flightsqlserver/sqlite_tables_schema_batch_reader.go @@ -60,10 +60,9 @@ func NewSqliteTablesSchemaBatchReader(ctx context.Context, mem memory.Allocator, // the historical *sql.DB API for external callers; FlightSQL uses this // connection-bound variant so a second uninitialized pool connection cannot // be opened while a request is streaming. -func newSqliteTablesSchemaBatchReader(ctx context.Context, mem memory.Allocator, rdr array.RecordReader, preparer statementPreparer, mainQuery string) (*SqliteTablesSchemaBatchReader, error) { - schemaQuery := `SELECT table_name, name, type, [notnull] - FROM pragma_table_info(table_name) - JOIN (` + mainQuery + `) WHERE table_name = ?` +func newSqliteTablesSchemaBatchReader(ctx context.Context, mem memory.Allocator, rdr array.RecordReader, preparer statementPreparer, _ string) (*SqliteTablesSchemaBatchReader, error) { + schemaQuery := `SELECT ? AS table_name, name, type, "notnull" + FROM pragma_table_info(?)` stmt, err := preparer.PrepareContext(ctx, schemaQuery) if err != nil { @@ -177,16 +176,16 @@ func (s *SqliteTablesSchemaBatchReader) Next() bool { columnFields := make([]arrow.Field, 0) for i := 0; i < tableNameArr.Len(); i++ { table := tableNameArr.Value(i) - rows, err := s.stmt.QueryContext(s.ctx, table) + rows, err := s.stmt.QueryContext(s.ctx, table, table) if err != nil { s.err = err return false } var tableName, name, typ string - var nn int + var notNull bool for rows.Next() { - if err := rows.Scan(&tableName, &name, &typ, &nn); err != nil { + if err := rows.Scan(&tableName, &name, &typ, ¬Null); err != nil { rows.Close() s.err = err return false @@ -195,7 +194,7 @@ func (s *SqliteTablesSchemaBatchReader) Next() bool { columnFields = append(columnFields, arrow.Field{ Name: name, Type: getArrowTypeFromString(typ), - Nullable: nn == 0, + Nullable: !notNull, Metadata: getColumnMetadata(bldr, getSqlTypeFromTypeName(typ), tableName), }) } diff --git a/main.go b/main.go index d87d7076..0d75a933 100644 --- a/main.go +++ b/main.go @@ -318,12 +318,15 @@ func main() { var flightServer flight.Server if flightsqlPort > 0 { - db := provider.Storage() - - srv, err := flightsqlserver.NewSQLiteFlightSQLServer(db, provider.InitializeConnection) + srv, err := flightsqlserver.NewSQLiteFlightSQLServerWithProvider(provider) if err != nil { log.Fatal(err) } + defer func() { + if err := srv.Close(); err != nil { + logrus.WithError(err).Warn("Failed to close FlightSQL lifecycle handles") + } + }() flightServer = flight.NewServerWithMiddleware(nil) flightServer.RegisterFlightService(flightsql.NewFlightServer(srv)) diff --git a/pgserver/arrow_lifecycle_test.go b/pgserver/arrow_lifecycle_test.go new file mode 100644 index 00000000..0e3f7675 --- /dev/null +++ b/pgserver/arrow_lifecycle_test.go @@ -0,0 +1,143 @@ +package pgserver + +import ( + "context" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/apecloud/myduckserver/backend" + "github.com/apecloud/myduckserver/catalog" + "github.com/apecloud/myduckserver/mycontext" + "github.com/cockroachdb/cockroachdb-parser/pkg/sql/sem/tree" + "github.com/dolthub/go-mysql-server/memory" + "github.com/dolthub/go-mysql-server/sql" + "github.com/stretchr/testify/require" +) + +func TestRejectArrowCopyInTransaction(t *testing.T) { + provider := catalog.NewInMemoryDBProvider() + t.Cleanup(func() { require.NoError(t, provider.Close()) }) + + session := backend.NewSession(memory.NewSession(sql.NewBaseSession(), provider), provider) + ctx := sql.NewContext(mycontext.WithFrontendQuery(context.Background()), sql.WithSession(session)) + require.NoError(t, beginPostgresTransaction(ctx, &tree.BeginTransaction{})) + + require.ErrorIs(t, rejectArrowCopyInTransaction(ctx), ErrArrowCopyInTransaction) + _, err := NewArrowDataLoader(ctx, nil, "", nil, nil, "") + require.ErrorIs(t, err, ErrArrowCopyInTransaction) + _, err = NewArrowWriter(ctx, nil, "", nil, nil, "", "") + require.ErrorIs(t, err, ErrArrowCopyInTransaction) + require.NoError(t, rollbackPostgresTransaction(ctx)) + require.NoError(t, rejectArrowCopyInTransaction(ctx)) +} + +type waitingCopyLoader struct { + abortStarted chan struct{} + release chan struct{} + waitCalled atomic.Bool + + mu sync.Mutex + abortCtx *sql.Context +} + +func (l *waitingCopyLoader) Start() <-chan error { return closedErrorChannel() } + +func (l *waitingCopyLoader) LoadChunk(*sql.Context, []byte) error { return nil } + +func (l *waitingCopyLoader) Abort(ctx *sql.Context) error { + l.mu.Lock() + l.abortCtx = ctx + l.mu.Unlock() + close(l.abortStarted) + return nil +} + +func (l *waitingCopyLoader) Finish(*sql.Context) (*LoadDataResults, error) { + return &LoadDataResults{}, nil +} + +func (l *waitingCopyLoader) Wait() { + l.waitCalled.Store(true) + <-l.release +} + +func closedErrorChannel() <-chan error { + ch := make(chan error) + close(ch) + return ch +} + +func TestHandleCopyFailAbortsAndWaitsForLoader(t *testing.T) { + loader := &waitingCopyLoader{ + abortStarted: make(chan struct{}), + release: make(chan struct{}), + } + ctx := sql.NewEmptyContext() + h := &ConnectionHandler{ + copyFromStdinState: ©FromStdinState{ + ctx: ctx, + dataLoader: loader, + }, + } + + type result struct { + stop, end bool + err error + } + resultCh := make(chan result, 1) + go func() { + stop, end, err := h.handleCopyFail(nil) + resultCh <- result{stop: stop, end: end, err: err} + }() + + select { + case <-loader.abortStarted: + case <-time.After(time.Second): + t.Fatal("COPY FAIL did not invoke DataLoader.Abort") + } + + select { + case <-resultCh: + t.Fatal("COPY FAIL returned before the loader finished") + case <-time.After(20 * time.Millisecond): + } + require.True(t, loader.waitCalled.Load()) + close(loader.release) + + select { + case got := <-resultCh: + require.False(t, got.stop) + require.True(t, got.end) + require.ErrorIs(t, got.err, ErrCopyAborted) + case <-time.After(time.Second): + t.Fatal("COPY FAIL did not return after loader completion") + } + require.Nil(t, h.copyFromStdinState) + loader.mu.Lock() + require.Same(t, ctx, loader.abortCtx) + loader.mu.Unlock() +} + +func TestArrowWriterCloseCancelsAndWaits(t *testing.T) { + ctx := sql.NewEmptyContext() + copyCtx, cancel := newCopyContext(ctx) + done := make(chan struct{}) + writer := &ArrowWriter{ + ctx: copyCtx, + cancel: cancel, + pipePath: t.TempDir() + "/arrow-copy", + done: done, + } + writer.started.Store(true) + + go func() { + <-copyCtx.Done() + close(done) + }() + + writer.Close() + require.ErrorIs(t, copyCtx.Err(), context.Canceled) + writer.Close() // Close is safe after the producer has already exited. +} diff --git a/pgserver/arrowloader.go b/pgserver/arrowloader.go index 1996a262..b8d21fb7 100644 --- a/pgserver/arrowloader.go +++ b/pgserver/arrowloader.go @@ -1,9 +1,11 @@ package pgserver import ( + "errors" "os" "strconv" "strings" + "syscall" "github.com/apache/arrow-go/v18/arrow/ipc" "github.com/apecloud/myduckserver/adapter" @@ -21,7 +23,130 @@ type ArrowDataLoader struct { var _ DataLoader = (*ArrowDataLoader)(nil) +// ErrArrowCopyInTransaction is returned before an Arrow COPY operation is +// started when the session already owns a transaction. Arrow's registration +// API requires a *sql.Conn, while normal statement execution must use the +// transaction's executor; rejecting this combination keeps the operation from +// silently splitting a statement across different physical owners. +var ErrArrowCopyInTransaction = errors.New("COPY FORMAT ARROW is not supported inside an active transaction") + +// rejectArrowCopyInTransaction enforces the current Arrow COPY lifecycle +// boundary. The Arrow implementation still uses Conn.Raw and therefore may +// only run when no session transaction is bound. Check both the GMS session +// transaction and the adapter's physical transaction binding because the +// PostgreSQL frontend binds explicit BEGIN directly in the adapter pool. +func rejectArrowCopyInTransaction(ctx *sql.Context) error { + if err := rejectArrowCopyGMSInTransaction(ctx); err != nil { + return err + } + if ctx == nil || ctx.Session == nil { + return nil + } + if _, ok := ctx.Session.(adapter.ConnectionHolder); ok { + if _, tx := adapter.TryGetTxnBinding(ctx); tx != nil { + return ErrArrowCopyInTransaction + } + } + return nil +} + +// rejectArrowCopyGMSInTransaction is the pool-free portion of the Arrow +// transaction gate. Protocol setup uses it before acquiring a retained +// snapshot so an already-visible GMS transaction fails without waiting on a +// pool lifecycle writer. +func rejectArrowCopyGMSInTransaction(ctx *sql.Context) error { + if ctx != nil && ctx.Session != nil && ctx.Session.GetTransaction() != nil { + return ErrArrowCopyInTransaction + } + return nil +} + +// rejectArrowCopyInTransactionForLease performs the same Arrow transaction +// gate while allowing callers that already hold a retained execution lease to +// skip a second pool lifecycle read. The lease's captured owner/transaction is +// authoritative for the asynchronous worker; a nested TryGetTxnBinding would +// otherwise block behind a pending pool shutdown writer. +func rejectArrowCopyInTransactionForLease(ctx *sql.Context, leaseHeld bool) error { + if ctx == nil || ctx.Session == nil { + return nil + } + if ctx.Session.GetTransaction() != nil { + return ErrArrowCopyInTransaction + } + if leaseHeld { + return nil + } + if _, ok := ctx.Session.(adapter.ConnectionHolder); ok { + if _, tx := adapter.TryGetTxnBinding(ctx); tx != nil { + return ErrArrowCopyInTransaction + } + } + return nil +} + +// rejectArrowCopyBinding applies the Arrow transaction boundary to a binding +// that was already captured under a retained lease. The lease lets us avoid a +// second pool lookup, but a captured physical transaction is still conclusive +// evidence that Arrow's Conn.Raw-based path cannot run safely. +func rejectArrowCopyBinding(ctx *sql.Context, binding postgresCopyBinding) error { + if binding.ownerTx != nil { + return ErrArrowCopyInTransaction + } + return rejectArrowCopyInTransactionForLease(ctx, binding.leaseHeld()) +} + func NewArrowDataLoader(ctx *sql.Context, handler *DuckHandler, schema string, table sql.InsertableTable, columns tree.NameList, options string) (DataLoader, error) { + if err := rejectArrowCopyGMSInTransaction(ctx); err != nil { + return nil, err + } + binding, err := capturePostgresCopyBindingWithHandler(ctx, handler, nil) + if err != nil { + return nil, err + } + loader, err := newArrowDataLoaderWithBinding(ctx, handler, schema, table, columns, options, binding) + if err != nil { + binding.releaseLease() + return nil, err + } + if arrowLoader, ok := loader.(*ArrowDataLoader); ok { + arrowLoader.ownsLease = true + } + return loader, nil +} + +func newArrowDataLoaderWithBinding( + ctx *sql.Context, handler *DuckHandler, schema string, table sql.InsertableTable, + columns tree.NameList, options string, binding postgresCopyBinding, +) (DataLoader, error) { + if err := rejectArrowCopyBinding(ctx, binding); err != nil { + return nil, err + } + ownedBinding := false + if !binding.snapshotCaptured { + var err error + binding, err = capturePostgresCopyBindingWithHandler(ctx, handler, nil) + if err != nil { + return nil, err + } + ownedBinding = true + } + if err := rejectArrowCopyBinding(ctx, binding); err != nil { + if ownedBinding { + binding.releaseLease() + } + return nil, err + } + cleanup := true + defer func() { + if cleanup && ownedBinding { + binding.releaseLease() + } + }() + target, err := resolvePostgresCopyTarget(ctx, handler, schema, table, binding.execer, binding.ownerConn) + if err != nil { + return nil, err + } + // Create the FIFO pipe duckBuilder := handler.e.Analyzer.ExecBuilder.PriorityBuilder.(*backend.DuckBuilder) pipePath, err := duckBuilder.CreatePipe(ctx, "pg-from-arrow") @@ -35,14 +160,21 @@ func NewArrowDataLoader(ctx *sql.Context, handler *DuckHandler, schema string, t loader := &ArrowDataLoader{ PipeDataLoader: PipeDataLoader{ - ctx: ctx, - cancel: cancel, - schema: schema, - table: table, - columns: columns, - pipePath: pipePath, - rowCount: make(chan int64, 1), - logger: ctx.GetLogger(), + ctx: ctx, + cancel: cancel, + execer: binding.execer, + ownerConn: binding.ownerConn, + ownerTx: binding.ownerTx, + snapshotCaptured: binding.snapshotCaptured, + lease: binding.lease, + schema: schema, + table: table, + target: target, + columns: columns, + pipePath: pipePath, + rowCount: make(chan int64, 1), + done: make(chan struct{}), + logger: ctx.GetLogger(), }, arrowName: arrowName, options: options, @@ -51,6 +183,7 @@ func NewArrowDataLoader(ctx *sql.Context, handler *DuckHandler, schema string, t loader.executeInsert(loader.buildSQL(), pipePath) } + cleanup = false return loader, nil } @@ -60,11 +193,15 @@ func (loader *ArrowDataLoader) buildSQL() string { b.Grow(256) b.WriteString("INSERT INTO ") - if loader.schema != "" { + if loader.target != "" { + b.WriteString(loader.target) + } else if loader.schema != "" { b.WriteString(loader.schema) b.WriteString(".") + b.WriteString(loader.table.Name()) + } else { + b.WriteString(loader.table.Name()) } - b.WriteString(loader.table.Name()) if len(loader.columns) > 0 { b.WriteString(" (") @@ -79,18 +216,52 @@ func (loader *ArrowDataLoader) buildSQL() string { } func (loader *ArrowDataLoader) executeInsert(sql string, pipePath string) { - defer close(loader.rowCount) + defer func() { + if loader.rowCount != nil { + close(loader.rowCount) + } + }() + + if err := loader.validateBinding(loader.ctx); err != nil { + loader.err.Store(&err) + if loader.blocked.Load() { + loader.unblockWriter() + } + return + } + + // The constructor and protocol handler reject active transactions. Repeat + // the check at execution time to cover a transaction that was bound between + // COPY setup and the first CopyData message. + if err := rejectArrowCopyBinding(loader.ctx, postgresCopyBinding{ + execer: loader.execer, + ownerConn: loader.ownerConn, + ownerTx: loader.ownerTx, + snapshotCaptured: loader.snapshotCaptured, + lease: loader.lease, + }); err != nil { + loader.err.Store(&err) + if loader.blocked.Load() { + loader.unblockWriter() + } + return + } // Open the pipe for reading. loader.logger.Debugf("Opening pipe for reading: %s", pipePath) pipe, err := os.OpenFile(pipePath, os.O_RDONLY, os.ModeNamedPipe) if err != nil { loader.err.Store(&err) - // Open the pipe once to unblock the writer - pipe, _ = os.OpenFile(pipePath, os.O_RDONLY, os.ModeNamedPipe) - loader.errPipe.Store(pipe) + // Open the pipe once to unblock a writer that may still be waiting in + // OpenFile. Use O_NONBLOCK: when the path disappeared or the reader + // failed before FIFO creation, a second blocking open would leak the + // loader goroutine forever. + if unblock, openErr := os.OpenFile(pipePath, os.O_RDONLY|syscall.O_NONBLOCK, os.ModeNamedPipe); openErr == nil { + loader.errPipe.Store(unblock) + } return } + defer pipe.Close() // Create an Arrow IPC reader from the pipe. loader.logger.Debugf("Creating Arrow IPC reader from pipe: %s", pipePath) @@ -101,8 +272,21 @@ func (loader *ArrowDataLoader) executeInsert(sql string, pipePath string) { } defer arrowReader.Release() - conn, err := adapter.GetConn(loader.ctx) - if err != nil { + conn := loader.ownerConn + if conn == nil { + err := errors.New("COPY connection owner is unavailable") + loader.err.Store(&err) + return + } + if loader.execer == nil { + loader.execer = conn + } + if conn == nil { + err := errors.New("COPY connection is unavailable") + loader.err.Store(&err) + return + } + if err := loader.validateBinding(loader.ctx); err != nil { loader.err.Store(&err) return } @@ -127,8 +311,12 @@ func (loader *ArrowDataLoader) executeInsert(sql string, pipePath string) { // Execute the INSERT statement. // This will block until the reader has finished reading the data. + if err := loader.validateBinding(loader.ctx); err != nil { + loader.err.Store(&err) + return + } loader.logger.Debugln("Executing SQL:", sql) - result, err := conn.ExecContext(loader.ctx, sql) + result, err := loader.execer.ExecContext(loader.ctx, sql) if err != nil { loader.err.Store(&err) return diff --git a/pgserver/arrowwriter.go b/pgserver/arrowwriter.go index 83a84ba7..0f89139b 100644 --- a/pgserver/arrowwriter.go +++ b/pgserver/arrowwriter.go @@ -1,14 +1,18 @@ package pgserver import ( + "context" + stdsql "database/sql" + "errors" "os" "strings" + "sync" "sync/atomic" + "syscall" "github.com/apache/arrow-go/v18/arrow/ipc" "github.com/apecloud/myduckserver/adapter" "github.com/apecloud/myduckserver/backend" - "github.com/apecloud/myduckserver/catalog" "github.com/cockroachdb/cockroachdb-parser/pkg/sql/sem/tree" "github.com/dolthub/go-mysql-server/sql" "github.com/duckdb/duckdb-go/v2" @@ -16,9 +20,22 @@ import ( type ArrowWriter struct { ctx *sql.Context + cancel context.CancelFunc + execer adapter.SQLExecutor + ownerConn *stdsql.Conn + ownerTx *stdsql.Tx duckSQL string pipePath string rawOptions string + + started atomic.Bool + blocked atomic.Bool + done chan struct{} + pipeMu sync.Mutex + pipe *os.File + lease *postgresCopyLease + close sync.Once + doneOnce sync.Once } func NewArrowWriter( @@ -28,6 +45,46 @@ func NewArrowWriter( query string, rawOptions string, ) (*ArrowWriter, error) { + if err := rejectArrowCopyGMSInTransaction(ctx); err != nil { + return nil, err + } + execer, ownerConn, ownerTx, lease, err := postgresCopySnapshotLease(ctx, handler, nil) + if err != nil { + return nil, err + } + if err := rejectArrowCopyBinding(ctx, postgresCopyBinding{ + execer: execer, + ownerConn: ownerConn, + ownerTx: ownerTx, + snapshotCaptured: ownerConn != nil, + lease: lease, + }); err != nil { + lease.Release() + return nil, err + } + if ownerTx != nil { + lease.Release() + return nil, ErrArrowCopyInTransaction + } + cleanup := true + defer func() { + if cleanup { + lease.Release() + } + }() + copyQuery := query + if table == nil && handler != nil { + if rewritten, rewriteErr := handler.rewritePostgresObjectRelationsWithSnapshot(ctx, query, execer, ownerConn); rewriteErr != nil { + return nil, rewriteErr + } else { + copyQuery = rewritten + } + } + target, err := resolvePostgresCopyTarget(ctx, handler, schema, table, execer, ownerConn) + if err != nil { + return nil, err + } + // Create the FIFO pipe db := handler.e.Analyzer.ExecBuilder.PriorityBuilder.(*backend.DuckBuilder) pipePath, err := db.CreatePipe(ctx, "pg-to-arrow") @@ -42,36 +99,61 @@ func NewArrowWriter( // https://duckdb.org/docs/sql/query_syntax/from.html#from-first-syntax // FROM table_name [ SELECT column_list ] builder.WriteString("FROM ") - if schema != "" { - builder.WriteString(catalog.QuoteIdentifierANSI(schema)) - builder.WriteString(".") - } - builder.WriteString(catalog.QuoteIdentifierANSI(table.Name())) + builder.WriteString(target) if columns != nil { builder.WriteString(" SELECT ") builder.WriteString(columns.String()) } } else { - builder.WriteString(query) + builder.WriteString(copyQuery) } - return &ArrowWriter{ - ctx: ctx, + // Keep the producer on a child context so Close can cancel a query that is + // still writing to the FIFO. The parent protocol context remains the source + // of cancellation and frontend origin metadata. + copyCtx, cancel := newCopyContext(ctx) + + writer := &ArrowWriter{ + ctx: copyCtx, + cancel: cancel, + execer: execer, + ownerConn: ownerConn, + ownerTx: ownerTx, duckSQL: builder.String(), pipePath: pipePath, rawOptions: rawOptions, // TODO(fan): parse rawOptions - }, nil + done: make(chan struct{}), + lease: lease, + } + cleanup = false + return writer, nil } func (dw *ArrowWriter) Start(globalErr *atomic.Pointer[error]) (string, chan CopyToResult, error) { // Execute the statement in a separate goroutine. + dw.ensureDone() ch := make(chan CopyToResult, 1) + dw.started.Store(true) + dw.blocked.Store(true) go func() { + defer close(dw.done) defer close(ch) dw.ctx.GetLogger().Tracef("Executing statement via Arrow interface: %s", dw.duckSQL) - conn, err := adapter.GetConn(dw.ctx) - if err != nil { + if err := rejectArrowCopyBinding(dw.ctx, postgresCopyBinding{ + execer: dw.execer, + ownerConn: dw.ownerConn, + ownerTx: dw.ownerTx, + snapshotCaptured: dw.ownerConn != nil, + lease: dw.lease, + }); err != nil { + globalErr.Store(&err) + ch <- CopyToResult{Err: err} + return + } + conn := dw.ownerConn + if conn == nil { + err := errors.New("COPY connection owner is unavailable") globalErr.Store(&err) ch <- CopyToResult{Err: err} return @@ -86,11 +168,14 @@ func (dw *ArrowWriter) Start(globalErr *atomic.Pointer[error]) (string, chan Cop // Open the pipe for writing. // This operation will block until the reader opens the pipe for reading. pipe, err := os.OpenFile(dw.pipePath, os.O_WRONLY, os.ModeNamedPipe) + dw.blocked.Store(false) if err != nil { globalErr.Store(&err) ch <- CopyToResult{Err: err} return } + dw.setPipe(pipe) + defer dw.clearPipe(pipe) defer pipe.Close() rowCount := int64(0) @@ -134,5 +219,72 @@ func (dw *ArrowWriter) Start(globalErr *atomic.Pointer[error]) (string, chan Cop } func (dw *ArrowWriter) Close() { - os.Remove(dw.pipePath) + dw.ensureDone() + dw.close.Do(func() { + if dw.cancel != nil { + dw.cancel() + } + + // Closing the producer side unblocks an in-flight IPC write. The protocol + // reader normally closes it first, but Close must also be safe on error + // paths where the reader has already returned. + dw.closePipe() + if dw.blocked.Load() { + // A producer blocked in OpenFile has no *os.File to close yet. A + // short-lived non-blocking reader lets that open return so cancellation + // can be observed by the producer goroutine. + if unblock, err := os.OpenFile(dw.pipePath, os.O_RDONLY|syscall.O_NONBLOCK, os.ModeNamedPipe); err == nil { + _ = unblock.Close() + } + } + }) + + if dw.started.Load() && dw.done != nil { + <-dw.done + } + _ = os.Remove(dw.pipePath) + if dw.lease != nil { + // Keep provider operation admission through a physical transaction's + // COMMIT/ROLLBACK callback; the pool snapshot itself is no longer needed + // after the producer and pipe reader have terminated. + dw.lease.ReleaseSnapshot() + if !dw.lease.OperationDeferred() { + dw.lease.ReleaseOperation() + } + } +} + +func (dw *ArrowWriter) ensureDone() { + if dw == nil { + return + } + dw.doneOnce.Do(func() { + if dw.done == nil { + dw.done = make(chan struct{}) + } + }) +} + +func (dw *ArrowWriter) setPipe(pipe *os.File) { + dw.pipeMu.Lock() + dw.pipe = pipe + dw.pipeMu.Unlock() +} + +func (dw *ArrowWriter) clearPipe(pipe *os.File) { + dw.pipeMu.Lock() + if dw.pipe == pipe { + dw.pipe = nil + } + dw.pipeMu.Unlock() +} + +func (dw *ArrowWriter) closePipe() { + dw.pipeMu.Lock() + pipe := dw.pipe + dw.pipe = nil + dw.pipeMu.Unlock() + if pipe != nil { + _ = pipe.Close() + } } diff --git a/pgserver/connection_data.go b/pgserver/connection_data.go index d623da66..77b4afc1 100644 --- a/pgserver/connection_data.go +++ b/pgserver/connection_data.go @@ -92,6 +92,11 @@ func (cs ConvertedStatement) QueryForAudit() string { // this statement is processed, the server accepts COPY DATA messages from the client with chunks of data to load // into a table. type copyFromStdinState struct { + // ctx is the frontend context that opened COPY. Keeping it on the state + // lets COPY FAIL abort the loader without manufacturing a plain/background + // context after the protocol has entered COPY mode. + ctx *sql.Context + // copyFromStdinNode stores the original CopyFrom statement that initiated the CopyData message sequence. This // node is used to look at what parameters were specified, such as which table to load data into, file format, // delimiters, etc. @@ -105,16 +110,36 @@ type copyFromStdinState struct { // dataLoader is the implementation of DataLoader that is used to load each individual CopyData chunk into the // target table. dataLoader DataLoader + // binding is captured before CopyInResponse is sent. Every later frontend + // message and asynchronous loader worker must continue to use this exact + // executor/connection/transaction identity. + binding postgresCopyBinding + // deferCommandComplete records that COPY was opened by the final statement + // of a simple Query. The handler clears its transient marker when the Query + // returns, but CopyDone arrives later and still needs the same protocol mode + // to order completion after the implicit commit. + deferCommandComplete bool // copyErr stores any error that was returned while processing a CopyData message and loading a chunk of data // to the target table. The server needs to keep track of any errors that were encountered while processing chunks // so that it can avoid sending a CommandComplete message if an error was encountered after the client already // sent a CopyDone message to the server. copyErr error + // copyAborted records that the loader has already been cancelled and + // drained. Keep the surrounding state as a terminal sentinel until the + // client sends CopyDone/CopyFail (or the protocol reaches Sync), because + // clients such as pgx send CopyDone after observing a server-side error. + // This makes repeated terminal cleanup idempotent while preventing a stale + // COPY message from being mistaken for a new operation. + copyAborted bool } type PortalData struct { - Statement ConvertedStatement - IsEmptyQuery bool + Statement ConvertedStatement + IsEmptyQuery bool + // Prepared indicates that this portal belongs to an engine-prepared + // statement. The underlying DuckDB statement is intentionally prepared at + // execution time and is never retained across protocol messages. + Prepared bool Fields []pgproto3.FieldDescription ResultFormatCodes []int16 Stmt *duckdb.Stmt @@ -123,7 +148,10 @@ type PortalData struct { } type PreparedStatementData struct { - Statement ConvertedStatement + Statement ConvertedStatement + // Prepared distinguishes an engine statement from an in-place protocol + // command. Stmt is transient and is normally nil after Parse. + Prepared bool ReturnFields []pgproto3.FieldDescription BindVarTypes []uint32 Stmt *duckdb.Stmt diff --git a/pgserver/connection_handler.go b/pgserver/connection_handler.go index dd5effc7..dffdef7a 100644 --- a/pgserver/connection_handler.go +++ b/pgserver/connection_handler.go @@ -19,6 +19,7 @@ import ( "bytes" "context" "crypto/tls" + stdsql "database/sql" "errors" "fmt" "io" @@ -39,6 +40,7 @@ import ( "github.com/dolthub/go-mysql-server/server" "github.com/dolthub/go-mysql-server/sql" "github.com/dolthub/vitess/go/mysql" + "github.com/jackc/pgx/v5/pgconn" "github.com/jackc/pgx/v5/pgproto3" "github.com/jackc/pgx/v5/pgtype" "github.com/sirupsen/logrus" @@ -54,9 +56,43 @@ type ConnectionHandler struct { backend *pgproto3.Backend pgTypeMap *pgtype.Map waitForSync bool + // implicitTx is the physical transaction opened for a protocol message + // group when the client has not issued an explicit BEGIN. It is kept + // separate from txStatus: PostgreSQL advertises T while the group is open, + // but returns to I as soon as the group is committed or rolled back. + implicitTx *stdsql.Tx + // implicitTxCtx and implicitTxConn are the exact lifecycle identity captured + // when implicitTx was opened. They must be reused for finalization; creating + // a fresh context can observe a replacement pool binding. + implicitTxCtx *sql.Context + implicitTxConn *stdsql.Conn + implicitTxPromoted bool + // deferredCommandComplete is used only by the final statement of a simple + // query. Its command tag is held until the implicit transaction commit has + // succeeded, matching PostgreSQL's completion ordering on commit failure. + deferredCommandComplete bool + // pendingCommandCompletes are response tags whose statement has finished, + // but whose message-group transaction still needs to commit. COPY uses this + // queue in the extended protocol because the client may send Sync only after + // receiving all data frames. Keep a queue rather than a single slot so a + // pipelined group cannot silently lose an earlier COPY completion. + pendingCommandCompletes []*pgproto3.CommandComplete + // pendingProtocolMessages is the ordered form of the deferred response + // queue. Most statements only need a CommandComplete, but in-place handlers + // such as SET also emit ParameterStatus; keeping both messages in one queue + // preserves their wire order until an implicit simple-query commit succeeds. + pendingProtocolMessages []pgproto3.BackendMessage + // txStatus is the transaction state advertised by ReadyForQuery. It is + // deliberately kept at the protocol layer because PostgreSQL transaction + // control is handled outside GMS's transaction iterator. + txStatus ReadyForQueryTransactionIndicator // copyFromStdinState is set when this connection is in the COPY FROM STDIN mode, meaning it is waiting on // COPY DATA messages from the client to import data into tables. copyFromStdinState *copyFromStdinState + // pendingCopyRelease keeps a successful extended-protocol COPY-IN lease + // alive from CopyDone through the following Sync, where the implicit + // transaction is committed. It is idempotent and nil outside that window. + pendingCopyRelease func() server *Server readOnly bool @@ -132,6 +168,7 @@ func NewConnectionHandler(conn net.Conn, handler mysql.Handler, engine *gms.Engi duckHandler: duckHandler, backend: pgproto3.NewBackend(conn, conn), pgTypeMap: pgtype.NewMap(), + txStatus: ReadyForQueryTransactionIndicator_Idle, server: server, readOnly: readOnly, @@ -144,13 +181,220 @@ func NewConnectionHandler(conn net.Conn, handler mysql.Handler, engine *gms.Engi return &connectionHandler } +// readyForQueryStatus returns a valid PostgreSQL transaction indicator. Keep +// the zero value usable for tests and for any handler constructed without the +// normal constructor. +func (h *ConnectionHandler) readyForQueryStatus() ReadyForQueryTransactionIndicator { + switch h.txStatus { + case ReadyForQueryTransactionIndicator_Idle, + ReadyForQueryTransactionIndicator_TransactionBlock, + ReadyForQueryTransactionIndicator_FailedTransactionBlock: + return h.txStatus + default: + h.txStatus = ReadyForQueryTransactionIndicator_Idle + return h.txStatus + } +} + +// markTransactionError records an error in an explicit transaction block. +// Errors while idle do not create a transaction, and a failed block remains +// failed until COMMIT or ROLLBACK finalizes it. +func (h *ConnectionHandler) markTransactionError() { + // An implicit protocol scope is rolled back immediately by its owning + // message handler. Do not expose that transient scope as PostgreSQL's sticky + // failed explicit block. + if h.implicitTx != nil && !h.implicitTxPromoted { + return + } + if h.readyForQueryStatus() == ReadyForQueryTransactionIndicator_TransactionBlock { + h.txStatus = ReadyForQueryTransactionIndicator_FailedTransactionBlock + } +} + +// observeTransactionBinding reports whether the session still has a physical +// transaction bound to this PostgreSQL connection. A real handler can inspect +// the pool through a fresh lifecycle context; minimal unit-test handlers do not +// have a DuckHandler/session and are treated as having no observable binding. +// The second return value is false when the production binding could not be +// inspected, in which case callers must fail closed and retain a non-idle +// protocol state. +func (h *ConnectionHandler) observeTransactionBinding() (*stdsql.Tx, bool) { + if h == nil { + return nil, true + } + if h.implicitTx != nil { + if h.implicitTxCtx == nil { + return h.implicitTx, true + } + _, tx := adapter.TryGetTxnBindingForRelease(h.implicitTxCtx) + return tx, true + } + if h.duckHandler == nil || h.mysqlConn == nil { + return nil, true + } + ctx, err := h.newProtocolSQLContext("") + if err != nil || ctx == nil { + return nil, false + } + _, tx := adapter.TryGetTxnBindingForRelease(ctx) + return tx, true +} + +// recordTransactionResult applies the protocol-visible state transition for a +// statement. COMMIT and ROLLBACK both leave the session idle even when their +// finalizer reports an error: the lifecycle code removes/evicts the physical +// transaction on such errors, and a rollback cleanup error must not advertise +// a transaction that has already been finalized. +func (h *ConnectionHandler) recordTransactionResult(stmt tree.Statement, err error) { + switch stmt.(type) { + case *tree.RollbackTransaction, *tree.CommitTransaction: + bound, known := h.observeTransactionBinding() + if known && bound == nil { + // The transaction-control executor has finalized (or evicted) the + // physical binding. Clear the protocol-owned marker as well so a + // later boundary cannot attempt to finalize a stale transaction. + h.clearImplicitScope() + h.clearPortals() + h.txStatus = ReadyForQueryTransactionIndicator_Idle + return + } + // A failed explicit control, or a control that claimed success while + // leaving a transaction bound, must not advertise Idle. Keep an + // existing failed state; otherwise transition an active block to E. + if !known || bound != nil { + h.txStatus = ReadyForQueryTransactionIndicator_FailedTransactionBlock + } + case *tree.BeginTransaction: + if err == nil && h.readyForQueryStatus() != ReadyForQueryTransactionIndicator_FailedTransactionBlock { + h.txStatus = ReadyForQueryTransactionIndicator_TransactionBlock + return + } + // BEGIN issued while T is already active returns 25001 but keeps the + // original transaction block alive. Other BEGIN failures while idle also + // leave the connection idle; neither case creates PostgreSQL's failed E + // state. + return + default: + if err != nil { + h.markTransactionError() + } + } +} + func (h *ConnectionHandler) closeBackendConn() { - ctx, err := h.duckHandler.NewContext(context.Background(), h.mysqlConn, "") + // A client can disconnect while COPY FROM STDIN is still being fed by a + // loader goroutine. Abort and wait for that loader before rolling back or + // closing its physical owner; otherwise it can continue inserting after the + // rollback (or race a reused connection). + if h.copyFromStdinState != nil { + if err := h.abortCopyFromStdinState(ErrCopyAborted); err != nil && h.logger != nil { + h.logger.WithError(err).Error("Failed to abort COPY while closing backend connection") + } + } + // An extended COPY may have completed its loader but is still waiting for + // Sync to commit the implicit transaction. Disconnecting at that boundary + // must release the deferred operation/snapshot lease as well. + h.releasePendingCopyLease() + if h.duckHandler == nil || h.mysqlConn == nil { + return + } + + // Resolve the binding once and retain its complete identity. A delayed close + // from an old protocol handler must never roll back or close a replacement + // transaction that reused the same logical session (or even the same + // physical connection). + ctx, err := h.newProtocolSQLContext("") if err != nil { - h.logger.WithError(err).Error("Failed to create context for closing backend connection") + if h.logger != nil { + h.logger.WithError(err).Error("Failed to inspect PostgreSQL transaction on close") + } return } - adapter.CloseConn(ctx) + if h.implicitTxCtx != nil { + ctx = h.implicitTxCtx + } + expectedConn, expectedTx := adapter.TryGetTxnBindingForRelease(ctx) + if h.implicitTx != nil { + expectedConn = h.implicitTxConn + expectedTx = h.implicitTx + } + if expectedConn == nil { + return + } + currentConn, currentTx := adapter.TryGetTxnBindingForRelease(ctx) + if currentConn != expectedConn || currentTx != expectedTx { + if h.logger != nil { + h.logger.Warn("Skipping PostgreSQL close after transaction binding replacement") + } + return + } + + var provider postgresRollbackOrphanCleaner + if catalogProvider := h.duckHandler.GetCatalogProvider(); catalogProvider != nil { + provider = catalogProvider + } + if expectedTx != nil { + var rollbackErr error + if h.implicitTx != nil && h.implicitScopeActive() { + // Keep the existing implicit finalizer semantics (including protocol + // state transitions); the post-finalization identity check below + // prevents this teardown from closing a replacement binding. + rollbackErr = h.finalizeImplicitPostgresTransaction(false) + } else { + rollbackErr = rollbackPostgresBinding(ctx, expectedConn, expectedTx, provider) + } + if rollbackErr != nil && h.logger != nil { + h.logger.WithError(rollbackErr).Error("Failed to rollback PostgreSQL transaction on close") + } + } + + // A successful rollback removes the old transaction mapping. Recheck both + // identities before closing; if a replacement appeared, leave it untouched. + currentConn, currentTx = adapter.TryGetTxnBindingForRelease(ctx) + if currentConn != expectedConn || currentTx != nil { + return + } + if closeErr := adapter.CloseConnIfBinding(ctx, expectedConn, nil); closeErr != nil && h.logger != nil { + h.logger.WithError(closeErr).Error("Failed to close backend connection") + } +} + +// rollbackPostgresBinding is the close-path variant of the protocol rollback +// helper. It carries the expected transaction through the adapter finalizer so +// a replacement cannot be selected between inspection and rollback. +func rollbackPostgresBinding( + ctx *sql.Context, + expectedConn *stdsql.Conn, + expectedTx *stdsql.Tx, + provider postgresRollbackOrphanCleaner, +) error { + if expectedTx == nil { + return nil + } + currentConn, currentTx := adapter.TryGetTxnBindingForRelease(ctx) + if currentConn != expectedConn || currentTx != expectedTx { + return fmt.Errorf("postgres transaction binding changed before close rollback") + } + var cleanup func(*stdsql.Conn) error + releaseCleanup, reservationErr := beginPostgresRollbackCleanupReservationWithError(ctx, provider, expectedTx) + if reservationErr == nil { + defer releaseCleanup() + } else { + // Keep the physical rollback path usable after a poisoned-generation + // admission failure; maintenance is skipped and the error is joined below. + releaseCleanup = nil + } + if provider != nil && reservationErr == nil { + cleanup = func(conn *stdsql.Conn) error { + base := context.WithoutCancel(ctx) + base = catalog.WithRollbackCleanupOwner(base, conn) + base = catalog.WithDuckLakeCleanupLease(base) + cleanupCtx := ctx.WithContext(base) + return cleanupPostgresRollbackOrphans(cleanupCtx, &tree.RollbackTransaction{}, provider, conn) + } + } + _, err := adapter.FinalizeRollbackWithCleanup(ctx, expectedTx, cleanup) + return errors.Join(reservationErr, err) } // HandleConnection handles a connection's session, reading messages, executing queries, and sending responses. @@ -184,8 +428,12 @@ func (h *ConnectionHandler) HandleConnection() { fmt.Println(returnErr.Error()) } - h.duckHandler.ConnectionClosed(h.mysqlConn) + h.closePreparedObjects() + // Release any transaction/connection-owned state before removing the + // GMS session; rollback cleanup needs the session binding to remain + // addressable. h.closeBackendConn() + h.duckHandler.ConnectionClosed(h.mysqlConn) if err := h.Conn().Close(); err != nil { fmt.Printf("Failed to properly close connection:\n%v\n", err) } @@ -393,12 +641,7 @@ func (h *ConnectionHandler) receiveMessage() (bool, error) { eomErr = fmt.Errorf("panic: %v", r) } - if !endOfMessages && h.waitForSync { - if syncErr := h.discardToSync(); syncErr != nil { - h.logger.Error(syncErr.Error()) - } - } - h.endOfMessages(eomErr) + h.handleProtocolError(eomErr, !endOfMessages && h.waitForSync) } }() } @@ -416,12 +659,9 @@ func (h *ConnectionHandler) receiveMessage() (bool, error) { var stop bool stop, endOfMessages, err = h.handleMessage(msg) if err != nil { - if !endOfMessages && h.waitForSync { - if syncErr := h.discardToSync(); syncErr != nil { - fmt.Println(syncErr.Error()) - } + if syncErr := h.handleProtocolError(err, !endOfMessages && h.waitForSync); syncErr != nil { + return false, syncErr } - h.endOfMessages(err) } else if endOfMessages { h.endOfMessages(nil) } @@ -435,11 +675,131 @@ func (h *ConnectionHandler) receiveMessage() (bool, error) { // and a READY FOR QUERY message should be sent back to the client, so it can send the next query. func (h *ConnectionHandler) handleMessage(msg pgproto3.Message) (stop, endOfMessages bool, err error) { logrus.Tracef("Handling message: %T", msg) + // COPY FROM STDIN switches the frontend into a dedicated sub-protocol. + // PostgreSQL explicitly ignores Flush and Sync while the copy is still + // active, accepts only CopyData/CopyDone/CopyFail, and treats every other + // message as a protocol error. A state carrying copyErr is already a + // terminal server-error sentinel; its messages are consumed by the normal + // recovery path instead of generating a second error. + if state := h.copyFromStdinState; state != nil { + switch msg.(type) { + case *pgproto3.Terminate, *pgproto3.CopyData, *pgproto3.CopyDone, *pgproto3.CopyFail: + // Handled by the switch below (Terminate may close the connection). + case *pgproto3.Flush: + if state.copyErr != nil && !h.waitForSync { + // Simple-query COPY already sent its ErrorResponse/ReadyForQuery; + // release the terminal sentinel before accepting a new exchange. + h.releaseCopyLease(state) + h.copyFromStdinState = nil + } + return false, false, nil + case *pgproto3.Sync: + if state.copyErr == nil { + // Sync is ignored during healthy COPY-in. In particular, do not + // clear waitForSync or finalize the implicit transaction yet. + return false, false, nil + } + if !h.waitForSync { + // A simple-query error has no discard-until-Sync phase. Clear the + // sentinel, then let the ordinary Sync branch provide its boundary. + h.releaseCopyLease(state) + h.copyFromStdinState = nil + } + default: + if state.copyErr != nil { + // Once COPY has reported an error, the extended protocol skips + // all messages until Sync. discardToSync normally owns this loop; + // retain the same behavior for direct handler calls as well. + if h.waitForSync { + return false, false, nil + } + // Simple-query COPY already returned to the normal command loop; + // release the sentinel and process this message as a new exchange. + h.releaseCopyLease(state) + h.copyFromStdinState = nil + break + } + protocolErr := &pgconn.PgError{ + Severity: string(ErrorResponseSeverity_Error), + Code: "08P01", + Message: fmt.Sprintf("unexpected message type %T during COPY FROM STDIN", msg), + } + if h.waitForSync { + // Keep the terminal sentinel until discardToSync consumes the + // client's Sync, just as a server-side CopyData failure does. + err = h.abortCopyLoaderState(state, protocolErr) + h.releaseCopyLease(state) + return false, false, err + } + // Simple-query COPY has no Sync recovery boundary. Abort and clear + // the state before returning the ErrorResponse/Ready pair. + err = h.abortCopyFromStdinState(protocolErr) + h.releaseCopyLease(state) + return false, true, err + } + } switch message := msg.(type) { case *pgproto3.Terminate: + if state := h.copyFromStdinState; state != nil { + // Terminate closes the connection without another frontend boundary; + // drain the loader and release its retained lease before teardown. + _ = h.abortCopyFromStdinState(ErrCopyAborted) + h.releaseCopyLease(state) + } return true, false, nil case *pgproto3.Sync: + // A healthy COPY-in returns early above, because PostgreSQL ignores Sync + // until CopyDone/CopyFail terminates the sub-protocol. This branch only + // handles an ordinary Sync or a terminal server-error sentinel. h.waitForSync = false + copyAborted := false + var syncErr error + if h.copyFromStdinState != nil { + // A server-side COPY error is reported before the client's Sync. In + // that case the state is only a terminal sentinel and must be cleared + // without emitting a second ErrorResponse/ReadyForQuery pair. An + // interrupted but otherwise successful COPY still needs normal abort + // reporting here. + copyAborted = true + state := h.copyFromStdinState + if state.copyErr != nil { + h.releaseCopyLease(state) + h.copyFromStdinState = nil + } else { + // Sync is the protocol boundary for an interrupted COPY, not a + // client COPY FAIL. Pass no synthetic cause so a clean loader abort + // remains a normal boundary; only an actual cleanup failure is + // surfaced to the client. + syncErr = h.abortCopyFromStdinState(nil) + h.releaseCopyLease(state) + } + // COPY cancellation is a statement failure for an explicit + // PostgreSQL block. An implicit message-group scope is rolled back + // below and remains protocol-idle, so markTransactionError leaves it + // alone in that case. + h.markTransactionError() + // Any earlier completions belong to the same message group. The + // rollback below invalidates them; do not report success for work that + // did not commit. + h.clearPendingCommandCompletes() + } + if h.implicitScopeActive() { + // A healthy COPY returns before reaching this branch because Sync is + // ignored until CopyDone. For an ordinary boundary (or a terminal + // server-error sentinel), finalize the implicit scope normally; the + // latter is already marked for rollback by the COPY error path. + finalizeErr := h.finalizeImplicitPostgresTransaction(!copyAborted) + if finalizeErr != nil { + syncErr = errors.Join(syncErr, finalizeErr) + } + } + h.releasePendingCopyLease() + if syncErr != nil { + return false, true, syncErr + } + if err := h.sendPendingCommandCompletes(); err != nil { + return false, true, err + } return false, true, nil case *pgproto3.Query: endOfMessages, err = h.handleQuery(message) @@ -452,6 +812,11 @@ func (h *ConnectionHandler) handleMessage(msg pgproto3.Message) (stop, endOfMess return false, false, h.handleBind(message) case *pgproto3.Execute: return false, false, h.handleExecute(message) + case *pgproto3.Flush: + if h.backend == nil { + return false, false, nil + } + return false, false, h.backend.Flush() case *pgproto3.Close: if message.ObjectType == 'S' { h.deletePreparedStatement(message.Name) @@ -466,7 +831,10 @@ func (h *ConnectionHandler) handleMessage(msg pgproto3.Message) (stop, endOfMess case *pgproto3.CopyFail: return h.handleCopyFail(message) default: - return false, true, fmt.Errorf(`unhandled message "%t"`, message) + // Unknown messages in an extended exchange are errors until the + // client's Sync. Returning endOfMessages=false makes receiveMessage + // send ErrorResponse, consume that Sync, and emit exactly one Ready. + return false, false, fmt.Errorf(`unhandled message "%t"`, message) } } @@ -474,12 +842,64 @@ func (h *ConnectionHandler) handleMessage(msg pgproto3.Message) (stop, endOfMess // expected as part of this query, in which case the server will send a READY FOR QUERY message back to the client so // that it can send its next query. func (h *ConnectionHandler) handleQuery(message *pgproto3.Query) (endOfMessages bool, err error) { + h.pendingCommandCompletes = nil + h.pendingProtocolMessages = nil + h.deferredCommandComplete = false + defer func() { + pending := h.takePendingProtocolMessages() + h.deferredCommandComplete = false + if err != nil { + err = h.finishImplicitPostgresError(err) + return + } + // A simple Query message is one implicit transaction scope. COPY FROM + // STDIN is the exception: its data arrives in later protocol messages, + // so handleCopyDone owns the eventual commit. + if endOfMessages && h.implicitScopeActive() { + err = h.finalizeImplicitPostgresTransaction(true) + if err != nil { + return + } + } + for _, response := range pending { + if err = h.send(response); err != nil { + return + } + } + }() + + // Sensitive SQL is normally rejected before any parser or shortcut. In a + // failed transaction, preserve PostgreSQL's 25P02 contract even for a + // sensitive payload, while still keeping the lexical security gate first. + if h.readyForQueryStatus() == ReadyForQueryTransactionIndicator_FailedTransactionBlock && catalog.IsSensitiveSQL(message.String) { + return true, postgresFailedTransactionError() + } if err := catalog.RejectSensitiveSQL(message.String); err != nil { return true, err } + statements, err := h.convertQuery(message.String) + if err != nil { + // Once a transaction is failed, even a statement that cannot be + // converted must be ignored rather than reaching a shortcut or the + // executor. Preserve the protocol-level SQLSTATE in that state. + if h.readyForQueryStatus() == ReadyForQueryTransactionIndicator_FailedTransactionBlock { + return true, postgresFailedTransactionError() + } + return true, err + } + if err := h.rejectFailedQueryStatements(statements); err != nil { + return true, err + } + // COPY FROM STDIN switches the connection into a message-driven mode: the + // client must send CopyData/CopyDone before any later SQL can run. Reject a + // trailing batch up front so no statement executes ahead of the COPY data. + if err := rejectCopyFromStdinTrailingStatements(statements, message.String); err != nil { + return true, err + } // usql use ";" to test if the connection is alive. If we don't handle it, this will return an error. So we need to - // manually handle it here. + // manually handle it here. Keep this shortcut after the failed-transaction + // gate and sensitive-SQL check so it cannot bypass protocol state handling. if message.String == ";" { err := h.send(makeCommandComplete("", 0)) if err != nil { @@ -487,47 +907,91 @@ func (h *ConnectionHandler) handleQuery(message *pgproto3.Query) (endOfMessages } return true, nil } + if len(statements) == 0 { + return true, h.send(&pgproto3.EmptyQueryResponse{}) + } + + // PSQL/workbench replacements are only safe after the transaction gate has + // classified the original SQL. In a failed block they must never execute a + // replacement SELECT behind the client's rejected command. + if h.readyForQueryStatus() != ReadyForQueryTransactionIndicator_FailedTransactionBlock { + if len(statements) == 1 { + // These replacements call run directly, before the normal statement + // loop sets its final-statement marker. Keep their completion behind + // the implicit message-group commit as well. + h.deferredCommandComplete = true + handled, shortcutErr := h.handledPSQLCommands(message.String) + h.deferredCommandComplete = false + if handled || shortcutErr != nil { + return true, shortcutErr + } - handled, err := h.handledPSQLCommands(message.String) - if handled || err != nil { - return true, err - } - - // TODO: Remove this once we support `SELECT * FROM function()` syntax - // Github issue: https://github.com/dolthub/doltgresql/issues/464 - handled, err = h.handledWorkbenchCommands(message.String) - if handled || err != nil { - return true, err - } - - statements, err := h.convertQuery(message.String) - if err != nil { - return true, err + // TODO: Remove this once we support `SELECT * FROM function()` syntax + // Github issue: https://github.com/dolthub/doltgresql/issues/464 + h.deferredCommandComplete = true + handled, shortcutErr = h.handledWorkbenchCommands(message.String) + h.deferredCommandComplete = false + if handled || shortcutErr != nil { + return true, shortcutErr + } + } } // A query message destroys the unnamed statement and the unnamed portal h.deletePreparedStatement("") h.deletePortal("") - for _, statement := range statements { + var handled bool + lastStatement := -1 + for i, statement := range statements { + if !isEmptyConvertedStatement(statement) { + lastStatement = i + } + } + for i, statement := range statements { statement.IsExtendedQuery = false - if err := h.rejectReadOnly(statement); err != nil { - return true, err + if gateErr := h.rejectStatementIfTransactionFailed(statement); gateErr != nil { + return true, gateErr + } + if statement.AST == nil && strings.TrimSpace(statement.String) == "" { + if err := h.send(&pgproto3.EmptyQueryResponse{}); err != nil { + return true, err + } + endOfMessages = true + continue + } + if ensureErr := h.ensureImplicitPostgresTransaction(statement); ensureErr != nil { + return true, ensureErr + } + if statementErr := h.rejectReadOnly(statement); statementErr != nil { + h.recordTransactionResult(statement.AST, statementErr) + return true, statementErr } // Certain statement types get handled directly by the handler instead of being passed to the engine - handled, endOfMessages, err = h.handleStatementOutsideEngine(statement) + var statementErr error + // In-place handlers such as PostgreSQL compatibility SELECT rewrites may + // call run themselves. Carry the same final-command deferral into those + // paths so an implicit commit failure cannot follow an already-sent + // CommandComplete. + h.deferredCommandComplete = i == lastStatement + handled, endOfMessages, statementErr = h.handleStatementOutsideEngine(statement) + h.deferredCommandComplete = false if handled { - if err != nil { - h.logger.Warnf("Failed to handle statement %s outside engine: %v", statementLogText(statement), err) - return true, err + h.recordTransactionResult(statement.AST, statementErr) + if statementErr != nil { + h.logger.Warnf("Failed to handle statement %s outside engine: %v", statementLogText(statement), statementErr) + return true, statementErr } } else { - if err != nil { - h.logger.Warnf("Failed to handle statement %s outside engine: %v", statementLogText(statement), err) + if statementErr != nil { + h.logger.Warnf("Failed to handle statement %s outside engine: %v", statementLogText(statement), statementErr) } - endOfMessages, err = true, h.run(statement) - if err != nil { - return true, err + h.deferredCommandComplete = i == lastStatement + endOfMessages, statementErr = true, h.run(statement) + h.deferredCommandComplete = false + h.recordTransactionResult(statement.AST, statementErr) + if statementErr != nil { + return true, statementErr } } } @@ -540,6 +1004,12 @@ func (h *ConnectionHandler) handleQuery(message *pgproto3.Query) (endOfMessages // if no more messages are expected for this query and server should send the client a READY FOR QUERY message, // and any error that occurred while handling the query. func (h *ConnectionHandler) handleStatementOutsideEngine(statement ConvertedStatement) (handled bool, endOfMessages bool, err error) { + if err := h.rejectStatementIfTransactionFailed(statement); err != nil { + return true, true, err + } + if err := h.ensureImplicitPostgresTransaction(statement); err != nil { + return true, true, err + } switch stmt := statement.AST.(type) { case *tree.Deallocate: // TODO: handle ALL keyword @@ -592,8 +1062,18 @@ func (h *ConnectionHandler) handleStatementOutsideEngine(statement ConvertedStat } // handleParse handles a parse message, returning any error that occurs -func (h *ConnectionHandler) handleParse(message *pgproto3.Parse) error { +func (h *ConnectionHandler) handleParse(message *pgproto3.Parse) (err error) { h.waitForSync = true + defer func() { + if err != nil { + err = h.finishImplicitPostgresError(err) + } + }() + // Keep the lexical sensitive-SQL gate ahead of parser/prepared-statement + // handling. A failed block still reports 25P02 for such a payload. + if h.readyForQueryStatus() == ReadyForQueryTransactionIndicator_FailedTransactionBlock && catalog.IsSensitiveSQL(message.Query) { + return postgresFailedTransactionError() + } if err := catalog.RejectSensitiveSQL(message.Query); err != nil { return err } @@ -601,12 +1081,24 @@ func (h *ConnectionHandler) handleParse(message *pgproto3.Parse) error { // TODO: "Named prepared statements must be explicitly closed before they can be redefined by another Parse message, but this is not required for the unnamed statement" statements, err := h.convertQuery(message.Query) if err != nil { + if h.readyForQueryStatus() == ReadyForQueryTransactionIndicator_FailedTransactionBlock { + return postgresFailedTransactionError() + } + return err + } + if err := h.rejectFailedQueryStatements(statements); err != nil { return err } // TODO(Noy): handle multiple statements + if len(statements) == 0 { + return fmt.Errorf("cannot prepare an empty statement list") + } statement := statements[0] statement.IsExtendedQuery = true + if err := h.rejectStatementIfTransactionFailed(statement); err != nil { + return err + } if err := h.rejectReadOnly(statement); err != nil { return err } @@ -617,6 +1109,9 @@ func (h *ConnectionHandler) handleParse(message *pgproto3.Parse) error { } return h.send(&pgproto3.ParseComplete{}) } + if err := h.ensureImplicitPostgresTransaction(statement); err != nil { + return err + } handledOutsideEngine, err := shouldQueryBeHandledInPlace(h, &statement) if err != nil { @@ -625,6 +1120,7 @@ func (h *ConnectionHandler) handleParse(message *pgproto3.Parse) error { if handledOutsideEngine { h.preparedStatements[message.Name] = PreparedStatementData{ Statement: statement, + Prepared: false, ReturnFields: nil, BindVarTypes: nil, Stmt: nil, @@ -639,7 +1135,11 @@ func (h *ConnectionHandler) handleParse(message *pgproto3.Parse) error { } if !statement.PgParsable { - statement.Tag = GetStatementTag(stmt) + if stmt != nil { + statement.Tag = GetStatementTag(stmt) + } else { + statement.Tag = GuessStatementTag(statement.String) + } } // https://www.postgresql.org/docs/current/protocol-flow.html#PROTOCOL-FLOW-EXT-QUERY @@ -661,6 +1161,7 @@ func (h *ConnectionHandler) handleParse(message *pgproto3.Parse) error { } h.preparedStatements[message.Name] = PreparedStatementData{ Statement: statement, + Prepared: true, ReturnFields: fields, BindVarTypes: bindVarTypes, Stmt: stmt, @@ -671,23 +1172,36 @@ func (h *ConnectionHandler) handleParse(message *pgproto3.Parse) error { } // handleDescribe handles a Describe message, returning any error that occurs -func (h *ConnectionHandler) handleDescribe(message *pgproto3.Describe) error { +func (h *ConnectionHandler) handleDescribe(message *pgproto3.Describe) (err error) { + defer func() { + if err != nil { + err = h.finishImplicitPostgresError(err) + } + }() var fields []pgproto3.FieldDescription var bindvarTypes []uint32 var tag string var returnsRows bool + var statement ConvertedStatement h.waitForSync = true if message.ObjectType == 'S' { preparedStatementData, ok := h.preparedStatements[message.Name] if !ok { + if h.readyForQueryStatus() == ReadyForQueryTransactionIndicator_FailedTransactionBlock { + return postgresFailedTransactionError() + } return fmt.Errorf("prepared statement %s does not exist", message.Name) } + if err := h.rejectStatementIfTransactionFailed(preparedStatementData.Statement); err != nil { + return err + } + statement = preparedStatementData.Statement // https://www.postgresql.org/docs/current/protocol-flow.html // > Note that since Bind has not yet been issued, the formats to be used for returned columns are not yet known to the backend; // > the format code fields in the RowDescription message will be zeroes in this case. - if preparedStatementData.Stmt != nil { + if preparedStatementData.Prepared { fields = slices.Clone(preparedStatementData.ReturnFields) for i := range fields { fields[i].Format = 0 @@ -704,10 +1218,20 @@ func (h *ConnectionHandler) handleDescribe(message *pgproto3.Describe) error { } else { portalData, ok := h.portals[message.Name] if !ok { + if h.readyForQueryStatus() == ReadyForQueryTransactionIndicator_FailedTransactionBlock { + return postgresFailedTransactionError() + } return fmt.Errorf("portal %s does not exist", message.Name) } + if err := h.rejectStatementIfTransactionFailed(portalData.Statement); err != nil { + return err + } + statement = portalData.Statement - if portalData.Stmt != nil { + // Prepared PostgreSQL portals retain SQL/metadata but no raw DuckDB + // statement. Use the explicit lifecycle marker rather than testing Stmt, + // so a portal Describe still returns fields for exec/cache_describe modes. + if portalData.Prepared { fields = portalData.Fields tag = portalData.Statement.Tag returnsRows = statementReturnsRows(portalData.Statement) @@ -717,6 +1241,9 @@ func (h *ConnectionHandler) handleDescribe(message *pgproto3.Describe) error { return nil } } + if err := h.ensureImplicitPostgresTransaction(statement); err != nil { + return err + } if !returnsRows { returnsRows = returnsRow(tag) @@ -725,15 +1252,26 @@ func (h *ConnectionHandler) handleDescribe(message *pgproto3.Describe) error { } // handleBind handles a bind message, returning any error that occurs -func (h *ConnectionHandler) handleBind(message *pgproto3.Bind) error { +func (h *ConnectionHandler) handleBind(message *pgproto3.Bind) (err error) { h.waitForSync = true + defer func() { + if err != nil { + err = h.finishImplicitPostgresError(err) + } + }() // TODO: a named portal object lasts till the end of the current transaction, unless explicitly destroyed // we need to destroy the named portal as a side effect of the transaction ending preparedData, ok := h.preparedStatements[message.PreparedStatement] if !ok { + if h.readyForQueryStatus() == ReadyForQueryTransactionIndicator_FailedTransactionBlock { + return postgresFailedTransactionError() + } return fmt.Errorf("prepared statement %s does not exist", message.PreparedStatement) } + if err := h.rejectStatementIfTransactionFailed(preparedData.Statement); err != nil { + return err + } if err := catalog.RejectSensitiveSQL(preparedData.Statement.String); err != nil { return err } @@ -742,9 +1280,12 @@ func (h *ConnectionHandler) handleBind(message *pgproto3.Bind) error { return err } } + if err := h.ensureImplicitPostgresTransaction(preparedData.Statement); err != nil { + return err + } logrus.Tracef("binding portal %q to prepared statement %s", message.DestinationPortal, message.PreparedStatement) - if preparedData.Stmt == nil { + if !preparedData.Prepared { h.portals[message.DestinationPortal] = PortalData{ Statement: preparedData.Statement, IsEmptyQuery: strings.TrimSpace(preparedData.Statement.String) == "", @@ -776,6 +1317,7 @@ func (h *ConnectionHandler) handleBind(message *pgproto3.Bind) error { h.portals[message.DestinationPortal] = PortalData{ Statement: preparedData.Statement, + Prepared: preparedData.Prepared, Fields: fields, ResultFormatCodes: message.ResultFormatCodes, Stmt: preparedData.Stmt, @@ -786,16 +1328,30 @@ func (h *ConnectionHandler) handleBind(message *pgproto3.Bind) error { } // handleExecute handles an execute message, returning any error that occurs -func (h *ConnectionHandler) handleExecute(message *pgproto3.Execute) error { +func (h *ConnectionHandler) handleExecute(message *pgproto3.Execute) (err error) { h.waitForSync = true + var statement tree.Statement + defer func() { + if err != nil { + err = h.finishImplicitPostgresError(err) + } + h.recordTransactionResult(statement, err) + }() // TODO: implement the RowMax portalData, ok := h.portals[message.Portal] if !ok { + if h.readyForQueryStatus() == ReadyForQueryTransactionIndicator_FailedTransactionBlock { + return postgresFailedTransactionError() + } return fmt.Errorf("portal %s does not exist", message.Portal) } query := portalData.Statement + statement = query.AST + if err := h.rejectStatementIfTransactionFailed(query); err != nil { + return err + } if err := catalog.RejectSensitiveSQL(query.String); err != nil { return err } @@ -804,10 +1360,16 @@ func (h *ConnectionHandler) handleExecute(message *pgproto3.Execute) error { return err } } + if err := h.ensureImplicitPostgresTransaction(query); err != nil { + return err + } + if err := h.sendNestedBeginWarning(query.AST); err != nil { + return err + } logrus.Tracef("executing portal %s with statement %s and %d bound values", message.Portal, statementLogText(query), len(portalData.Vars)) if portalData.IsEmptyQuery { - err := h.send(&pgproto3.NoData{}) + err = h.send(&pgproto3.NoData{}) if err != nil { return fmt.Errorf("error sending NoData message: %w", err) } @@ -815,9 +1377,12 @@ func (h *ConnectionHandler) handleExecute(message *pgproto3.Execute) error { } // Certain statement types get handled directly by the handler instead of being passed to the engine - if strings.ToUpper(query.Tag) != "SELECT" || portalData.Stmt == nil { - handled, _, err := h.handleStatementOutsideEngine(query) - if handled { + if strings.ToUpper(query.Tag) != "SELECT" || !portalData.Prepared { + handled, _, outsideErr := h.handleStatementOutsideEngine(query) + if outsideErr != nil { + err = outsideErr + } + if handled || outsideErr != nil { return err } } @@ -826,12 +1391,17 @@ func (h *ConnectionHandler) handleExecute(message *pgproto3.Execute) error { rowsAffected := int32(0) callback := h.spoolRowsCallback(query, &rowsAffected, true) - err := h.duckHandler.ComExecuteBound(frontendContext(context.Background()), h.mysqlConn, portalData, callback) + err = h.duckHandler.ComExecuteBound(h.protocolContext(), h.mysqlConn, portalData, callback) if err != nil { return err } - return h.send(makeCommandComplete(query.Tag, rowsAffected)) + commandComplete := makeCommandComplete(h.protocolCommandTag(query), rowsAffected) + if h.shouldDeferCommandComplete() { + h.queueCommandComplete(commandComplete) + return nil + } + return h.send(commandComplete) } func makeCommandComplete(tag string, rows int32) *pgproto3.CommandComplete { @@ -853,44 +1423,235 @@ func makeCommandComplete(tag string, rows int32) *pgproto3.CommandComplete { // messages are expected, and the server should tell the client that it is ready for the next query, and |err| contains // any error that occurred while processing the COPY DATA message. func (h *ConnectionHandler) handleCopyData(message *pgproto3.CopyData) (stop bool, endOfMessages bool, err error) { - helper, messages, err := h.handleCopyDataHelper(message) - if err != nil { - h.copyFromStdinState.copyErr = err + stop, endOfMessages, err = h.handleCopyDataHelper(message) + if err == nil { + return stop, endOfMessages, nil + } + + // A COPY DATA failure aborts the loader immediately, but retains a terminal + // state sentinel until CopyDone/CopyFail (or Sync) arrives. pgx sends + // CopyDone after it observes a server-side ErrorResponse; clearing the state + // here would make that protocol message look like a second, invalid COPY. + // The loader is still fully drained before the error is reported, so no + // reader goroutine can continue inserting after the transaction rollback. + state := h.copyFromStdinState + err = h.abortCopyLoaderState(state, err) + // The loader is fully drained before releasing its lease. Release before + // implicit rollback cleanup: the provider cleanup barrier waits for active + // logical operations, so retaining this operation while rolling back would + // deadlock. + h.releaseCopyLease(state) + err = h.finishImplicitPostgresError(err) + return stop, h.copyErrorEndOfMessages(), err +} + +// copyOperationContext returns the context associated with a COPY operation. +// Production COPY states retain their setup context; the empty-context +// fallback keeps lifecycle cleanup safe for minimal handlers used in tests. +func (h *ConnectionHandler) copyOperationContext(state *copyFromStdinState) (*sql.Context, error) { + if state != nil && state.ctx != nil { + return state.ctx, nil + } + if h != nil && h.duckHandler != nil { + return h.duckHandler.NewContext(frontendContext(context.Background()), h.mysqlConn, "") + } + return sql.NewEmptyContext(), nil +} + +// copyErrorEndOfMessages preserves the extended-query protocol's Sync +// discard path. A COPY error in an extended exchange must leave +// waitForSync=true so receiveMessage consumes the pending Sync before sending +// ReadyForQuery; simple-query and direct lifecycle calls can finish now. +func (h *ConnectionHandler) copyErrorEndOfMessages() bool { + return h == nil || !h.waitForSync +} + +func (h *ConnectionHandler) copySuccessEndOfMessages() bool { + return h == nil || !h.waitForSync +} + +// releaseCopyLease closes the state-owned COPY snapshot and releases provider +// operation admission when no physical transaction finalizer has registered a +// callback. A registered operation lease intentionally remains held through +// COMMIT/ROLLBACK; the backend Session invokes its callback from the exact +// physical finalizer. +func (h *ConnectionHandler) releaseCopyLease(state *copyFromStdinState) { + if state != nil { + state.binding.releaseTerminal() + } +} + +// releaseCopyOperationFallback releases only operation admission for a COPY +// whose physical transaction could not accept a finalization callback. A +// registered operation lease belongs to the transaction's COMMIT/ROLLBACK +// callback and must remain held until that callback runs. +func (h *ConnectionHandler) releaseCopyOperationFallback(state *copyFromStdinState) { + if state == nil { + return + } + lease := state.binding.lease + if lease == nil || !lease.OperationDeferred() { + state.binding.releaseOperation() + } +} + +func (h *ConnectionHandler) releasePendingCopyLease() { + if h == nil || h.pendingCopyRelease == nil { + return + } + release := h.pendingCopyRelease + h.pendingCopyRelease = nil + release() +} + +func (h *ConnectionHandler) deferCopyLeaseUntilSync(state *copyFromStdinState) { + if h == nil || state == nil { + return + } + // This fallback is used only by legacy/test sessions that cannot register a + // physical post-finalization callback. Keep the operation half until Sync + // so a successful extended COPY commits before it is released; registered + // leases are released by the finalizer and do not need this slot. + if state.binding.lease == nil || state.binding.lease.OperationDeferred() { + return + } + release := state.binding.releaseOperation + if release == nil { + return + } + if h.pendingCopyRelease == nil { + h.pendingCopyRelease = release + return + } + // There should normally be only one active COPY exchange per handler, but + // compose releases defensively if a test or pipelined path supplies both. + previous := h.pendingCopyRelease + h.pendingCopyRelease = func() { + previous() + release() + } +} + +// abortCopyLoaderState aborts and waits for an active COPY loader without +// changing the handler's state pointer. Keeping the state is important after a +// server-side CopyData failure: clients commonly send CopyDone after reading +// ErrorResponse, and that terminal message must be consumed exactly once. +// Repeated calls are idempotent, including for custom loaders whose Abort and +// Wait methods are not themselves idempotent. +func (h *ConnectionHandler) abortCopyLoaderState(state *copyFromStdinState, cause error) error { + if state == nil { + return cause } - return helper, messages, err + if state.copyAborted { + if state.copyErr != nil { + return state.copyErr + } + return cause + } + if cause == nil { + cause = state.copyErr + } else if state.copyErr != nil && !errors.Is(cause, state.copyErr) { + cause = errors.Join(state.copyErr, cause) + } else if state.copyErr != nil { + cause = state.copyErr + } + + var abortErr error + if state.dataLoader != nil { + ctx := state.ctx + if ctx == nil { + ctx = sql.NewEmptyContext() + } + abortErr = state.dataLoader.Abort(ctx) + if waiter, ok := state.dataLoader.(interface{ Wait() }); ok { + waiter.Wait() + } + } + state.copyAborted = true + if abortErr != nil { + if cause == nil { + cause = abortErr + } else if !errors.Is(cause, abortErr) { + cause = errors.Join(cause, abortErr) + } + } + state.copyErr = cause + return cause +} + +// abortCopyFromStdinState aborts and waits for the active COPY loader, then +// clears the connection's COPY state on every return path. The original cause +// is retained alongside any cleanup error so callers can report both. +func (h *ConnectionHandler) abortCopyFromStdinState(cause error) error { + if h == nil { + return cause + } + state := h.copyFromStdinState + if state == nil { + return cause + } + err := h.abortCopyLoaderState(state, cause) + // The state is cleared immediately after the loader has been drained. This + // helper is used by protocol error/close paths that may proceed directly to + // rollback cleanup, so release before returning to avoid blocking that + // cleanup on our own operation admission lease. + h.releaseCopyLease(state) + h.copyFromStdinState = nil + return err } -// handleCopyDataHelper is a helper function that should only be invoked by handleCopyData. handleCopyData wraps this -// function so that it can capture any returned error message and store it in the saved state. +// handleCopyDataHelper is a helper function that should only be invoked by +// handleCopyData. The wrapper owns aborting the loader; a failed COPY retains a +// terminal sentinel until the protocol consumes CopyDone/CopyFail or Sync. func (h *ConnectionHandler) handleCopyDataHelper(message *pgproto3.CopyData) (stop bool, endOfMessages bool, err error) { - if h.copyFromStdinState == nil { + state := h.copyFromStdinState + if state == nil { + if h.readyForQueryStatus() == ReadyForQueryTransactionIndicator_FailedTransactionBlock { + return false, h.copyErrorEndOfMessages(), postgresFailedTransactionError() + } return false, true, fmt.Errorf("COPY DATA message received without a COPY FROM STDIN operation in progress") } + // Once a prior chunk failed, ignore any additional data until the client + // sends its terminal CopyDone/CopyFail. Returning nil here avoids emitting a + // second error/Ready pair while preserving the sentinel for that terminal + // message. + if state.copyErr != nil { + return false, false, nil + } + if h.readyForQueryStatus() == ReadyForQueryTransactionIndicator_FailedTransactionBlock { + return false, h.copyErrorEndOfMessages(), postgresFailedTransactionError() + } // Grab a sql.Context. - sqlCtx, err := h.duckHandler.NewContext(frontendContext(context.Background()), h.mysqlConn, "") + sqlCtx, err := h.copyOperationContext(state) if err != nil { return false, false, err } + if err := state.binding.validate(sqlCtx); err != nil { + return false, false, err + } - dataLoader := h.copyFromStdinState.dataLoader + dataLoader := state.dataLoader if dataLoader == nil { - copyFrom := h.copyFromStdinState.copyFromStdinNode + copyFrom := state.copyFromStdinNode if copyFrom == nil { return false, false, fmt.Errorf("no COPY FROM STDIN node found") } - table := h.copyFromStdinState.targetTable + table := state.targetTable if table == nil { return false, true, fmt.Errorf("no target table found") } - rawOptions := h.copyFromStdinState.rawOptions + rawOptions := state.rawOptions switch copyFrom.Options.CopyFormat { case CopyFormatArrow: - dataLoader, err = NewArrowDataLoader( + if err := rejectArrowCopyBinding(sqlCtx, state.binding); err != nil { + return false, false, err + } + dataLoader, err = newArrowDataLoaderWithBinding( sqlCtx, h.duckHandler, copyFrom.Table.Schema(), table, copyFrom.Columns, - rawOptions, + rawOptions, state.binding, ) case tree.CopyFormatText: // Remove `\.` from the end of the message data, if it exists @@ -902,11 +1663,11 @@ func (h *ConnectionHandler) handleCopyDataHelper(message *pgproto3.CopyData) (st } fallthrough case tree.CopyFormatCSV: - dataLoader, err = NewCsvDataLoader( + dataLoader, err = newCsvDataLoaderWithBinding( sqlCtx, h.duckHandler, copyFrom.Table.Schema(), table, copyFrom.Columns, ©From.Options, - rawOptions, + rawOptions, state.binding, ) case tree.CopyFormatBinary: err = fmt.Errorf("BINARY format is not supported for COPY FROM") @@ -918,14 +1679,21 @@ func (h *ConnectionHandler) handleCopyDataHelper(message *pgproto3.CopyData) (st return false, false, err } + // Publish the loader before Start so a setup failure can still be + // aborted by handleCopyData's common cleanup path. + state.dataLoader = dataLoader + if err := state.binding.validate(sqlCtx); err != nil { + return false, false, err + } ready := dataLoader.Start() if err, hasErr := <-ready; hasErr { return false, false, err } - - h.copyFromStdinState.dataLoader = dataLoader } + if err := state.binding.validate(sqlCtx); err != nil { + return false, false, err + } if err = dataLoader.LoadChunk(sqlCtx, message.Data); err != nil { return false, false, err } @@ -940,63 +1708,155 @@ func (h *ConnectionHandler) handleCopyDataHelper(message *pgproto3.CopyData) (st // |endOfMessages| is true if no more COPY DATA messages are expected, and the server should tell the client that it is // ready for the next query, and |err| contains any error that occurred while processing the COPY DATA message. func (h *ConnectionHandler) handleCopyDone(_ *pgproto3.CopyDone) (stop bool, endOfMessages bool, err error) { - if h.copyFromStdinState == nil { + defer func() { + if err != nil { + err = h.finishImplicitPostgresError(err) + } + }() + + state := h.copyFromStdinState + if state == nil { + if h.readyForQueryStatus() == ReadyForQueryTransactionIndicator_FailedTransactionBlock { + err = h.abortCopyFromStdinState(postgresFailedTransactionError()) + return false, h.copyErrorEndOfMessages(), err + } return false, true, fmt.Errorf("COPY DONE message received without a COPY FROM STDIN operation in progress") } - // If there was a previous error returned from processing a CopyData message, then don't return an error here - // and don't send endOfMessage=true, since the CopyData error already sent endOfMessage=true. If we do send - // endOfMessage=true here, then the client gets confused about the unexpected/extra Idle message since the - // server has already reported it was idle in the last message after the returned error. - if h.copyFromStdinState.copyErr != nil { + // A previous COPY DATA failure may have been observed by the protocol + // layer before COPY DONE arrives. The loader was already aborted by + // handleCopyData; consume this terminal message and clear the sentinel + // without sending another ErrorResponse or ReadyForQuery. This is required + // for clients such as pgx, which send CopyDone after seeing the server error. + if state.copyErr != nil { + h.releaseCopyLease(state) + h.copyFromStdinState = nil return false, false, nil } + if h.readyForQueryStatus() == ReadyForQueryTransactionIndicator_FailedTransactionBlock { + err = h.abortCopyFromStdinState(postgresFailedTransactionError()) + return false, h.copyErrorEndOfMessages(), err + } - dataLoader := h.copyFromStdinState.dataLoader + dataLoader := state.dataLoader if dataLoader == nil { - return false, true, - fmt.Errorf("no data loader found for COPY FROM STDIN operation") + err = h.abortCopyFromStdinState(fmt.Errorf("no data loader found for COPY FROM STDIN operation")) + return false, h.copyErrorEndOfMessages(), err } - sqlCtx, err := h.duckHandler.NewContext(frontendContext(context.Background()), h.mysqlConn, "") + sqlCtx, err := h.copyOperationContext(state) if err != nil { - return false, false, err + err = h.abortCopyFromStdinState(err) + return false, h.copyErrorEndOfMessages(), err + } + if bindingErr := state.binding.validate(sqlCtx); bindingErr != nil { + err = h.abortCopyFromStdinState(bindingErr) + return false, h.copyErrorEndOfMessages(), err } loadDataResults, err := dataLoader.Finish(sqlCtx) if err != nil { - return false, false, err + err = h.abortCopyLoaderState(state, err) + h.releaseCopyLease(state) + h.copyFromStdinState = nil + return false, h.copyErrorEndOfMessages(), err + } + // Finish has drained the loader and its child iterator. The pool snapshot is + // therefore no longer needed, but a registered provider operation lease may + // still be protecting the surrounding physical transaction until its + // finalizer. Keep those two lifetimes independent. + state.binding.releaseSnapshot() + if loadDataResults == nil { + err = h.abortCopyLoaderState(state, fmt.Errorf("COPY FROM STDIN returned no load results")) + h.releaseCopyOperationFallback(state) + h.copyFromStdinState = nil + return false, h.copyErrorEndOfMessages(), err } h.copyFromStdinState = nil + if h.implicitScopeActive() { + if h.waitForSync { + // Extended COPY sends its completion before Sync. Keep the physical + // transaction's operation admission until that Sync commits the + // implicit transaction. The snapshot was released immediately after + // Finish above. + h.deferCopyLeaseUntilSync(state) + } else { + finalizeErr := h.finalizeImplicitPostgresTransaction(true) + // A finalizer callback normally releases registered operation leases. + // Release again unconditionally so a finalization error that evicts the + // transaction cannot strand a lease; postgresCopyLease is idempotent. + state.binding.releaseOperation() + if finalizeErr != nil { + return false, true, finalizeErr + } + } + } else { + // Explicit transactions remain open after COPY DONE, but the COPY + // producer/loader has completed. Only an unregistered fallback operation + // lease can end now; a registered lease belongs to the later physical + // COMMIT/ROLLBACK callback. + h.releaseCopyOperationFallback(state) + } + commandComplete := &pgproto3.CommandComplete{ + CommandTag: []byte(fmt.Sprintf("COPY %d", loadDataResults.RowsLoaded)), + } + if h.shouldDeferCopyCommandComplete() { + // A simple Query's final COPY is held until its implicit transaction + // commits. Extended COPY leaves this branch false so its tag preserves + // the protocol ordering relative to later pipelined messages. + h.queueCommandComplete(commandComplete) + return false, false, nil + } // We send back endOfMessage=true, since the COPY DONE message ends the COPY DATA flow and the server is ready // to accept the next query now. - return false, true, h.send(&pgproto3.CommandComplete{ - CommandTag: []byte(fmt.Sprintf("COPY %d", loadDataResults.RowsLoaded)), - }) + return false, h.copySuccessEndOfMessages(), h.send(commandComplete) } // handleCopyFail handles a COPY FAIL message by aborting the in-progress COPY DATA operation. The |stop| response // parameter is true if the connection handler should shut down the connection, |endOfMessages| is true if no more // COPY DATA messages are expected, and the server should tell the client that it is ready for the next query, and // |err| contains any error that occurred while processing the COPY DATA message. -func (h *ConnectionHandler) handleCopyFail(_ *pgproto3.CopyFail) (stop bool, endOfMessages bool, err error) { - if h.copyFromStdinState == nil { +func (h *ConnectionHandler) handleCopyFail(message *pgproto3.CopyFail) (stop bool, endOfMessages bool, err error) { + defer func() { + if err != nil { + err = h.finishImplicitPostgresError(err) + } + }() + + state := h.copyFromStdinState + if state == nil { + if h.readyForQueryStatus() == ReadyForQueryTransactionIndicator_FailedTransactionBlock { + return false, h.copyErrorEndOfMessages(), postgresFailedTransactionError() + } return false, true, fmt.Errorf("COPY FAIL message received without a COPY FROM STDIN operation in progress") } - - dataLoader := h.copyFromStdinState.dataLoader - if dataLoader == nil { - return false, true, - fmt.Errorf("no data loader found for COPY FROM STDIN operation") + // A server-side COPY error has already been reported and the loader has + // already been drained. CopyFail is then only the client's terminal + // acknowledgement; consume it and clear the sentinel without producing a + // duplicate response. + if state.copyErr != nil { + h.releaseCopyLease(state) + h.copyFromStdinState = nil + return false, false, nil } - h.copyFromStdinState = nil - // We send back endOfMessage=true, since the COPY FAIL message ends the COPY DATA flow and the server is ready - // to accept the next query now. - return false, true, nil + // COPY FAIL is itself a failed command. Preserve the client's reason when + // present, while retaining ErrCopyAborted for callers that need to classify + // an intentional cancellation. + cause := state.copyErr + if cause == nil { + cause = ErrCopyAborted + } + if message != nil && message.Message != "" { + cause = fmt.Errorf("%w: %s", cause, message.Message) + } + err = h.abortCopyFromStdinState(cause) + // We send back endOfMessage=true for a simple exchange; an extended query + // keeps waitForSync set so receiveMessage can discard the pending Sync. + return false, h.copyErrorEndOfMessages(), err } func (h *ConnectionHandler) deallocatePreparedStatement(name string, preparedStatements map[string]PreparedStatementData, query ConvertedStatement, conn net.Conn) error { @@ -1006,7 +1866,7 @@ func (h *ConnectionHandler) deallocatePreparedStatement(name string, preparedSta } h.deletePreparedStatement(name) - return h.send(&pgproto3.CommandComplete{ + return h.sendOrQueueProtocolMessage(&pgproto3.CommandComplete{ CommandTag: []byte(query.Tag), }) } @@ -1015,7 +1875,13 @@ func (h *ConnectionHandler) deletePreparedStatement(name string) { ps, ok := h.preparedStatements[name] if ok { delete(h.preparedStatements, name) - if ps.Closed.CompareAndSwap(false, true) { + if ps.Closed == nil { + if ps.Stmt != nil { + _ = ps.Stmt.Close() + } + return + } + if ps.Closed.CompareAndSwap(false, true) && ps.Stmt != nil { ps.Stmt.Close() } } @@ -1025,12 +1891,40 @@ func (h *ConnectionHandler) deletePortal(name string) { p, ok := h.portals[name] if ok { delete(h.portals, name) - if p.Closed.CompareAndSwap(false, true) { + if p.Closed == nil { + if p.Stmt != nil { + _ = p.Stmt.Close() + } + return + } + if p.Closed.CompareAndSwap(false, true) && p.Stmt != nil { p.Stmt.Close() } } } +// clearPortals releases portals whose lifetime is bounded by the transaction +// that created them. Prepared statement metadata is retained separately and +// can be rebound in a later protocol cycle. +func (h *ConnectionHandler) clearPortals() { + for name := range h.portals { + h.deletePortal(name) + } +} + +// closePreparedObjects releases any transient raw statements left by legacy +// or test handlers before the protocol connection is torn down. Current +// PostgreSQL prepared plans retain SQL/metadata only, but keeping teardown +// identity-safe makes connection replacement harmless for older entries. +func (h *ConnectionHandler) closePreparedObjects() { + for name := range h.preparedStatements { + h.deletePreparedStatement(name) + } + for name := range h.portals { + h.deletePortal(name) + } +} + // convertBindParameters handles the conversion from bind parameters to variable values. func (h *ConnectionHandler) convertBindParameters(types []uint32, formatCodes []int16, values [][]byte) ([]any, error) { if len(types) != len(values) { @@ -1054,6 +1948,12 @@ func (h *ConnectionHandler) convertBindParameters(types []uint32, formatCodes [] // run runs the given statement and sends a CommandComplete message to the client func (h *ConnectionHandler) run(statement ConvertedStatement) error { + if err := h.rejectStatementIfTransactionFailed(statement); err != nil { + return err + } + if err := h.ensureImplicitPostgresTransaction(statement); err != nil { + return err + } if err := catalog.RejectSensitiveSQL(statement.String); err != nil { return err } @@ -1062,6 +1962,9 @@ func (h *ConnectionHandler) run(statement ConvertedStatement) error { return err } } + if err := h.sendNestedBeginWarning(statement.AST); err != nil { + return err + } h.logger.Tracef("running statement %s", statementLogText(statement)) // |rowsAffected| gets altered by the callback below @@ -1099,7 +2002,7 @@ func (h *ConnectionHandler) run(statement ConvertedStatement) error { callback := h.spoolRowsCallback(statement, &rowsAffected, false) if err := h.duckHandler.ComQuery( - frontendContext(context.Background()), + h.protocolContext(), h.mysqlConn, statement.String, statement.QueryForAudit(), @@ -1109,7 +2012,12 @@ func (h *ConnectionHandler) run(statement ConvertedStatement) error { return fmt.Errorf("fallback statement execution failed: %w", err) } - return h.send(makeCommandComplete(statement.Tag, rowsAffected)) + commandComplete := makeCommandComplete(h.protocolCommandTag(statement), rowsAffected) + if h.shouldDeferCommandComplete() { + h.queueCommandComplete(commandComplete) + return nil + } + return h.send(commandComplete) } func (h *ConnectionHandler) rejectReadOnly(statement ConvertedStatement) error { @@ -1172,6 +2080,23 @@ func (h *ConnectionHandler) spoolRowsCallback(statement ConvertedStatement, rows } } +// sendNestedBeginWarning mirrors PostgreSQL's behavior for BEGIN issued while +// an explicit transaction block is already active: it is a successful no-op, +// but emits a WARNING with SQLSTATE 25001 before the command completion. +func (h *ConnectionHandler) sendNestedBeginWarning(statement tree.Statement) error { + if h == nil || h.readyForQueryStatus() != ReadyForQueryTransactionIndicator_TransactionBlock { + return nil + } + if _, ok := statement.(*tree.BeginTransaction); !ok { + return nil + } + return h.send(&pgproto3.NoticeResponse{ + Severity: string(ErrorResponseSeverity_Warning), + Code: "25001", + Message: "there is already a transaction in progress", + }) +} + // sendDescribeResponse sends a response message for a Describe message func (h *ConnectionHandler) sendDescribeResponse(fields []pgproto3.FieldDescription, types []uint32, returnsRows bool) error { // The prepared statement variant of the describe command returns the OIDs of the parameters. @@ -1309,17 +2234,156 @@ func (h *ConnectionHandler) handledWorkbenchCommands(statement string) (bool, er return false, nil } +// sendReadyForQuery emits the protocol boundary after a message group. Keep +// this separate from endOfMessages so extended-query errors can send their +// ErrorResponse first, consume the client's Sync, and only then advertise the +// final transaction state. +func (h *ConnectionHandler) sendReadyForQuery() error { + return h.send(&pgproto3.ReadyForQuery{ + TxStatus: byte(h.readyForQueryStatus()), + }) +} + +func (h *ConnectionHandler) queueCommandComplete(commandComplete *pgproto3.CommandComplete) { + if h == nil || commandComplete == nil { + return + } + h.pendingCommandCompletes = append(h.pendingCommandCompletes, commandComplete) + h.pendingProtocolMessages = append(h.pendingProtocolMessages, commandComplete) +} + +// queueProtocolMessage appends an in-place response to the same ordered queue +// used by deferred CommandComplete messages. It is intentionally private to +// the protocol handler: callers must still decide whether a message is safe to +// emit immediately or should be held behind an implicit commit. +func (h *ConnectionHandler) queueProtocolMessage(message pgproto3.BackendMessage) { + if h == nil || message == nil { + return + } + h.pendingProtocolMessages = append(h.pendingProtocolMessages, message) +} + +// sendOrQueueProtocolMessage keeps direct protocol handlers from advertising +// success before the final implicit simple-query transaction has committed. +// Extended Execute paths leave deferredCommandComplete false and retain their +// normal response timing. +func (h *ConnectionHandler) sendOrQueueProtocolMessage(message pgproto3.BackendMessage) error { + if h == nil || message == nil { + return nil + } + if h.shouldDeferCommandComplete() { + h.queueProtocolMessage(message) + return nil + } + return h.send(message) +} + +// takePendingProtocolMessages returns the ordered queue and clears both the +// unified queue and the legacy command-complete mirror. The fallback conversion +// keeps direct unit tests that seed pendingCommandCompletes working. +func (h *ConnectionHandler) takePendingProtocolMessages() []pgproto3.BackendMessage { + if h == nil { + return nil + } + if len(h.pendingProtocolMessages) == 0 && len(h.pendingCommandCompletes) > 0 { + pending := make([]pgproto3.BackendMessage, 0, len(h.pendingCommandCompletes)) + for _, commandComplete := range h.pendingCommandCompletes { + if commandComplete != nil { + pending = append(pending, commandComplete) + } + } + h.pendingCommandCompletes = nil + return pending + } + pending := h.pendingProtocolMessages + h.pendingProtocolMessages = nil + h.pendingCommandCompletes = nil + return pending +} + +func (h *ConnectionHandler) takePendingCommandCompletes() []*pgproto3.CommandComplete { + if h == nil || len(h.pendingCommandCompletes) == 0 { + return nil + } + pending := h.pendingCommandCompletes + h.pendingCommandCompletes = nil + return pending +} + +func (h *ConnectionHandler) clearPendingCommandCompletes() { + if h != nil { + h.pendingCommandCompletes = nil + h.pendingProtocolMessages = nil + } +} + +func (h *ConnectionHandler) shouldDeferCommandComplete() bool { + // Ordinary extended Execute messages must expose their completion as soon + // as the statement finishes. The client may be waiting on that tag before it + // sends or consumes the next pipelined Execute. Only a simple-query's final + // statement marks deferredCommandComplete; COPY uses the specialized helper + // below because its completion is produced by a later protocol message. + return h != nil && h.deferredCommandComplete +} + +// shouldDeferCopyCommandComplete preserves the simple-query final-statement +// deferral used to keep a completion tag behind an implicit commit. Extended +// COPY must not be deferred: PostgreSQL sends CommandComplete immediately +// after CopyDone, before any later pipelined Parse/Bind/Execute or Sync. COPY +// FROM/TO can finish in a later handler, but the simple-query marker remains +// authoritative for the protocol mode that opened it. +func (h *ConnectionHandler) shouldDeferCopyCommandComplete() bool { + return h != nil && h.deferredCommandComplete +} + +// sendPendingCommandCompletes emits delayed tags only after the surrounding +// extended-query transaction has finalized successfully. Callers clear the +// queue on any finalization error, so a failed group cannot report success. +func (h *ConnectionHandler) sendPendingCommandCompletes() error { + for _, response := range h.takePendingProtocolMessages() { + if err := h.send(response); err != nil { + return err + } + } + return nil +} + +// handleProtocolError applies the transaction error policy and completes an +// extended protocol exchange. PostgreSQL sends ErrorResponse immediately and +// ignores messages until Sync; waiting for Sync before sending the error can +// deadlock clients that send Sync only after observing the error. +func (h *ConnectionHandler) handleProtocolError(err error, waitForSync bool) error { + if err == nil { + return nil + } + err = h.finishImplicitPostgresError(err) + // Any delayed completion belongs to the failed message group and must not + // be emitted after its ErrorResponse/ReadyForQuery boundary. + h.clearPendingCommandCompletes() + if waitForSync && h.waitForSync { + h.sendError(err) + if syncErr := h.discardToSync(); syncErr != nil { + return syncErr + } + return h.sendReadyForQuery() + } + h.endOfMessages(err) + return nil +} + // endOfMessages should be called from HandleConnection or a function within HandleConnection. This represents the end // of the message slice, which may occur naturally (all relevant response messages have been sent) or on error. Once // endOfMessages has been called, no further messages should be sent, and the connection loop should wait for the next // query. A nil error should be provided if this is being called naturally. func (h *ConnectionHandler) endOfMessages(err error) { if err != nil { + // Keep the transaction indicator in sync before sending the error and + // trailing ReadyForQuery. This also covers parse/bind/protocol errors + // that do not have a ConvertedStatement available at this layer. + h.markTransactionError() h.sendError(err) } - if sendErr := h.send(&pgproto3.ReadyForQuery{ - TxStatus: byte(ReadyForQueryTransactionIndicator_Idle), - }); sendErr != nil { + if sendErr := h.sendReadyForQuery(); sendErr != nil { // We panic here for the same reason as above. panic(sendErr) } @@ -1328,10 +2392,22 @@ func (h *ConnectionHandler) endOfMessages(err error) { // sendError sends the given error to the client. This should generally never be called directly. func (h *ConnectionHandler) sendError(err error) { fmt.Println(err.Error()) + severity := string(ErrorResponseSeverity_Error) + code := "XX000" // internal_error for now + message := err.Error() + var pgErr *pgconn.PgError + if errors.As(err, &pgErr) { + severity = pgErr.Severity + if severity == "" { + severity = string(ErrorResponseSeverity_Error) + } + code = pgErr.Code + message = pgErr.Message + } if sendErr := h.send(&pgproto3.ErrorResponse{ - Severity: string(ErrorResponseSeverity_Error), - Code: "XX000", // internal_error for now - Message: err.Error(), + Severity: severity, + Code: code, + Message: message, }); sendErr != nil { // If we're unable to send anything to the connection, then there's something wrong with the connection and // we should terminate it. This will be caught in HandleConnection's defer block. @@ -1417,9 +2493,24 @@ func (h *ConnectionHandler) convertQuery(query string, modifiers ...QueryModifie // discardAll handles the DISCARD ALL command func (h *ConnectionHandler) discardAll(query ConvertedStatement) error { + // DISCARD ALL resets the physical session. PostgreSQL forbids it while a + // transaction block is active; closing the pooled connection here would + // otherwise roll back (or replace) the transaction while leaving the wire + // status at T. The same guard covers an implicit simple-query scope and a + // promoted BEGIN scope. + if h != nil { + status := h.readyForQueryStatus() + if h.implicitTx != nil || status == ReadyForQueryTransactionIndicator_TransactionBlock { + return &pgconn.PgError{ + Severity: "ERROR", + Code: "25001", + Message: "DISCARD ALL cannot run inside a transaction block", + } + } + } h.closeBackendConn() - return h.send(&pgproto3.CommandComplete{ + return h.sendOrQueueProtocolMessage(&pgproto3.CommandComplete{ CommandTag: []byte(query.Tag), }) } @@ -1431,6 +2522,9 @@ func (h *ConnectionHandler) handleCopyFromStdinQuery( query ConvertedStatement, copyFrom *tree.CopyFrom, rawOptions string, // For non-PG-parseable COPY FROM ) error { + if err := h.rejectStatementIfTransactionFailed(query); err != nil { + return err + } if err := catalog.RejectSensitiveSQL(query.String); err != nil { return err } @@ -1444,16 +2538,34 @@ func (h *ConnectionHandler) handleCopyFromStdinQuery( return err } sqlCtx.SetLogger(sqlCtx.GetLogger().WithField("query", catalog.RedactSensitiveSQL(query.QueryForAudit()))) + if copyFrom.Options.CopyFormat == CopyFormatArrow { + if err := rejectArrowCopyGMSInTransaction(sqlCtx); err != nil { + return err + } + } table, err := ValidateCopyFrom(copyFrom, sqlCtx) if err != nil { return err } + binding, err := capturePostgresCopyBindingWithHandler(sqlCtx, h.duckHandler, copyFrom) + if err != nil { + return err + } + if copyFrom.Options.CopyFormat == CopyFormatArrow { + if err := rejectArrowCopyBinding(sqlCtx, binding); err != nil { + binding.releaseLease() + return err + } + } h.copyFromStdinState = ©FromStdinState{ - copyFromStdinNode: copyFrom, - targetTable: table, - rawOptions: rawOptions, + ctx: sqlCtx, + copyFromStdinNode: copyFrom, + targetTable: table, + rawOptions: rawOptions, + binding: binding, + deferCommandComplete: h.deferredCommandComplete, } var format byte @@ -1464,21 +2576,113 @@ func (h *ConnectionHandler) handleCopyFromStdinQuery( format = 1 // binary format } - return h.send(&pgproto3.CopyInResponse{ + if err := h.send(&pgproto3.CopyInResponse{ OverallFormat: format, - }) + }); err != nil { + // The client never entered COPY mode, so no later terminal message can + // release this retained snapshot. Release it on the setup failure and + // clear the state before returning the wire error. + binding.releaseLease() + h.copyFromStdinState = nil + return err + } + return nil } // DiscardToSync discards all messages in the buffer until a Sync has been reached. If a Sync was never sent, then this // may cause the connection to lock until the client send a Sync, as their request structure was malformed. +// +// A receive failure is terminal for the current exchange. It cannot be treated +// like a normal loop return: an active COPY loader may still own a worker and a +// pool snapshot, while an implicit protocol transaction may still need its +// rollback callback to release the operation lease. Keep the receive error as +// the primary cause, but finish all local lifecycle work before returning it. +func (h *ConnectionHandler) cleanupDiscardToSyncReceiveError(receiveErr error) error { + if h == nil { + return receiveErr + } + + // There will be no later Sync after EOF/decode failure. Clear this marker + // up front so a teardown/recovery caller cannot leave the handler in a stale + // extended-exchange state. + h.waitForSync = false + state := h.copyFromStdinState + var cleanupErr error + if state != nil { + // Passing a nil cause keeps a successful abort from manufacturing a + // second error; Abort/Wait errors are still returned by the helper. + cleanupErr = h.abortCopyFromStdinState(nil) + } + + // Completions from the interrupted group must never be emitted after its + // receive boundary. A completed extended COPY can also have a lease held + // until Sync; release it before entering rollback so cleanup cannot wait on + // its own operation admission. + h.clearPendingCommandCompletes() + h.releasePendingCopyLease() + + var rollbackErr error + if h.implicitScopeActive() { + rollbackErr = h.finalizeImplicitPostgresTransaction(false) + // A registered physical callback normally releases this half during the + // finalizer. Keep the defensive idempotent release used by CopyDone so a + // failed/evicted finalizer cannot strand the lease. + if state != nil { + state.binding.releaseOperation() + } + } else { + // Explicit transaction blocks remain owned by their physical finalizer; + // record the protocol failure without claiming that the transaction was + // rolled back here. + h.markTransactionError() + } + + return errors.Join(receiveErr, cleanupErr, rollbackErr) +} + func (h *ConnectionHandler) discardToSync() error { for { message, err := h.backend.Receive() if err != nil { - return err + return h.cleanupDiscardToSyncReceiveError(err) } - if _, ok := message.(*pgproto3.Sync); ok { + switch message.(type) { + case *pgproto3.CopyDone, *pgproto3.CopyFail: + // COPY errors are sent before the client terminal message. Consume + // that message while discarding the extended exchange, then let Sync + // provide the single ReadyForQuery boundary. + if state := h.copyFromStdinState; state != nil && state.copyErr != nil { + h.releaseCopyLease(state) + h.copyFromStdinState = nil + } + case *pgproto3.Sync: + h.waitForSync = false + var syncErr error + if state := h.copyFromStdinState; state != nil { + if state.copyErr != nil { + // A terminal COPY sentinel is already aborted and only needs + // to be released at the protocol boundary. + h.releaseCopyLease(state) + h.copyFromStdinState = nil + } else { + syncErr = h.abortCopyFromStdinState(nil) + } + h.markTransactionError() + h.clearPendingCommandCompletes() + } + // Error paths normally roll back before entering this loop. Keep the + // boundary self-contained for callers that invoke discardToSync + // directly or for a panic recovered outside a statement handler. + if h.implicitScopeActive() { + if rollbackErr := h.finalizeImplicitPostgresTransaction(false); rollbackErr != nil { + syncErr = errors.Join(syncErr, rollbackErr) + } + } + h.releasePendingCopyLease() + if syncErr != nil { + return syncErr + } return nil } } @@ -1500,7 +2704,21 @@ func returnsRow(tag string) bool { } } +// copyToStdoutCancellationError preserves a concrete error recorded by the +// asynchronous COPY reader/producer when it cancels the protocol context. +// Keeping this selection in one place prevents the stale setup error from +// masking the failure that actually reached the client-facing path. +func copyToStdoutCancellationError(ctxErr error, globalErr *atomic.Pointer[error]) error { + if errPtr := globalErr.Load(); errPtr != nil { + return errors.Join(ctxErr, *errPtr) + } + return ctxErr +} + func (h *ConnectionHandler) handleCopyToStdout(query ConvertedStatement, copyTo *tree.CopyTo, subquery string, format tree.CopyFormat, rawOptions string) error { + if err := h.rejectStatementIfTransactionFailed(query); err != nil { + return err + } if err := catalog.RejectSensitiveSQL(query.String); err != nil { return err } @@ -1514,6 +2732,11 @@ func (h *ConnectionHandler) handleCopyToStdout(query ConvertedStatement, copyTo return err } ctx.SetLogger(ctx.GetLogger().WithField("query", catalog.RedactSensitiveSQL(query.QueryForAudit()))) + if format == CopyFormatArrow { + if err := rejectArrowCopyGMSInTransaction(ctx); err != nil { + return err + } + } // Create cancelable context childCtx, cancel := context.WithCancel(ctx) @@ -1687,10 +2910,10 @@ func (h *ConnectionHandler) handleCopyToStdout(query ConvertedStatement, copyTo select { case <-ctx.Done(): // Context is canceled <-done - if errPtr := globalErr.Load(); errPtr != nil { - return errors.Join(ctx.Err(), err) - } - return ctx.Err() + // The reader/producer goroutine may have canceled the context after + // recording a concrete COPY or wire error. Preserve that error rather + // than returning the stale setup error from writer.Start. + return copyToStdoutCancellationError(ctx.Err(), &globalErr) case result := <-ch: if blocked.Load() { // If the pipe is still opened for reading but the writer has exited, @@ -1720,8 +2943,13 @@ func (h *ConnectionHandler) handleCopyToStdout(query ConvertedStatement, copyTo // Send CommandComplete with the number of rows copied ctx.GetLogger().Debugf("sending CommandComplete to the client") - return h.send(&pgproto3.CommandComplete{ + commandComplete := &pgproto3.CommandComplete{ CommandTag: []byte(fmt.Sprintf("COPY %d", result.RowCount)), - }) + } + if h.shouldDeferCopyCommandComplete() { + h.queueCommandComplete(commandComplete) + return nil + } + return h.send(commandComplete) } } diff --git a/pgserver/connection_handler_copy_test.go b/pgserver/connection_handler_copy_test.go new file mode 100644 index 00000000..dfe59426 --- /dev/null +++ b/pgserver/connection_handler_copy_test.go @@ -0,0 +1,536 @@ +package pgserver + +import ( + "bytes" + "context" + stdsql "database/sql" + "errors" + "io" + "testing" + + "github.com/apecloud/myduckserver/adapter" + "github.com/apecloud/myduckserver/backend" + "github.com/apecloud/myduckserver/catalog" + "github.com/apecloud/myduckserver/mycontext" + sqle "github.com/dolthub/go-mysql-server" + "github.com/dolthub/go-mysql-server/memory" + "github.com/dolthub/go-mysql-server/sql" + "github.com/jackc/pgx/v5/pgconn" + "github.com/jackc/pgx/v5/pgproto3" + "github.com/stretchr/testify/require" +) + +type handlerCopyLifecycleLoader struct { + loadErr error + finishErr error + abortErr error + + loadCalls int + finishCalls int + abortCalls int + waitCalls int + abortCtx *sql.Context +} + +func (l *handlerCopyLifecycleLoader) Start() <-chan error { + ready := make(chan error) + close(ready) + return ready +} + +func (l *handlerCopyLifecycleLoader) LoadChunk(*sql.Context, []byte) error { + l.loadCalls++ + return l.loadErr +} + +func (l *handlerCopyLifecycleLoader) Abort(ctx *sql.Context) error { + l.abortCalls++ + l.abortCtx = ctx + return l.abortErr +} + +func (l *handlerCopyLifecycleLoader) Finish(*sql.Context) (*LoadDataResults, error) { + l.finishCalls++ + if l.finishErr != nil { + return nil, l.finishErr + } + return &LoadDataResults{}, nil +} + +func (l *handlerCopyLifecycleLoader) Wait() { + l.waitCalls++ +} + +func requireFailedTransactionSQLState(t *testing.T, err error) { + t.Helper() + var pgErr *pgconn.PgError + require.ErrorAs(t, err, &pgErr) + require.Equal(t, "25P02", pgErr.Code) +} + +func TestExtendedLookupUsesFailedTransactionErrorForMissingObject(t *testing.T) { + tests := []struct { + name string + run func(*ConnectionHandler) error + }{ + { + name: "describe statement", + run: func(h *ConnectionHandler) error { + return h.handleDescribe(&pgproto3.Describe{ObjectType: 'S', Name: "missing"}) + }, + }, + { + name: "describe portal", + run: func(h *ConnectionHandler) error { + return h.handleDescribe(&pgproto3.Describe{ObjectType: 'P', Name: "missing"}) + }, + }, + { + name: "bind", + run: func(h *ConnectionHandler) error { + return h.handleBind(&pgproto3.Bind{PreparedStatement: "missing"}) + }, + }, + { + name: "execute", + run: func(h *ConnectionHandler) error { + return h.handleExecute(&pgproto3.Execute{Portal: "missing"}) + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + h := &ConnectionHandler{txStatus: ReadyForQueryTransactionIndicator_FailedTransactionBlock} + requireFailedTransactionSQLState(t, tt.run(h)) + require.Equal(t, ReadyForQueryTransactionIndicator_FailedTransactionBlock, h.readyForQueryStatus()) + }) + } +} + +func TestCopyDataErrorAbortsAndRetainsTerminalState(t *testing.T) { + loadErr := errors.New("load chunk failed") + abortErr := errors.New("abort failed") + ctx := sql.NewEmptyContext() + loader := &handlerCopyLifecycleLoader{loadErr: loadErr, abortErr: abortErr} + h := &ConnectionHandler{ + txStatus: ReadyForQueryTransactionIndicator_TransactionBlock, + copyFromStdinState: ©FromStdinState{ + ctx: ctx, + dataLoader: loader, + }, + } + + stop, end, err := h.handleCopyData(&pgproto3.CopyData{Data: []byte("row")}) + require.False(t, stop) + require.True(t, end) + require.ErrorIs(t, err, loadErr) + require.ErrorIs(t, err, abortErr) + require.NotNil(t, h.copyFromStdinState) + require.ErrorIs(t, h.copyFromStdinState.copyErr, loadErr) + require.ErrorIs(t, h.copyFromStdinState.copyErr, abortErr) + require.True(t, h.copyFromStdinState.copyAborted) + require.Equal(t, ReadyForQueryTransactionIndicator_FailedTransactionBlock, h.readyForQueryStatus()) + require.Equal(t, 1, loader.loadCalls) + require.Equal(t, 1, loader.abortCalls) + require.Equal(t, 1, loader.waitCalls) + require.Same(t, ctx, loader.abortCtx) + + // pgx sends CopyDone after observing the server-side error. It must be + // consumed as a terminal acknowledgement without a second error/Ready and + // without invoking the loader lifecycle twice. + stop, end, err = h.handleCopyDone(&pgproto3.CopyDone{}) + require.False(t, stop) + require.False(t, end) + require.NoError(t, err) + require.Nil(t, h.copyFromStdinState) + require.Equal(t, 1, loader.abortCalls) + require.Equal(t, 1, loader.waitCalls) +} + +func TestCopyDoneErrorAbortsClearsStateAndFailsTransaction(t *testing.T) { + finishErr := errors.New("finish failed") + ctx := sql.NewEmptyContext() + loader := &handlerCopyLifecycleLoader{finishErr: finishErr} + h := &ConnectionHandler{ + txStatus: ReadyForQueryTransactionIndicator_TransactionBlock, + copyFromStdinState: ©FromStdinState{ + ctx: ctx, + dataLoader: loader, + }, + } + + stop, end, err := h.handleCopyDone(&pgproto3.CopyDone{}) + require.False(t, stop) + require.True(t, end) + require.ErrorIs(t, err, finishErr) + require.Nil(t, h.copyFromStdinState) + require.Equal(t, ReadyForQueryTransactionIndicator_FailedTransactionBlock, h.readyForQueryStatus()) + require.Equal(t, 1, loader.finishCalls) + require.Equal(t, 1, loader.abortCalls) + require.Equal(t, 1, loader.waitCalls) + require.Same(t, ctx, loader.abortCtx) +} + +func TestCopyFailAbortsClearsStateAndFailsTransaction(t *testing.T) { + ctx := sql.NewEmptyContext() + loader := &handlerCopyLifecycleLoader{} + h := &ConnectionHandler{ + txStatus: ReadyForQueryTransactionIndicator_TransactionBlock, + copyFromStdinState: ©FromStdinState{ + ctx: ctx, + dataLoader: loader, + }, + } + + stop, end, err := h.handleCopyFail(&pgproto3.CopyFail{Message: "client input failed"}) + require.False(t, stop) + require.True(t, end) + require.ErrorIs(t, err, ErrCopyAborted) + require.ErrorContains(t, err, "client input failed") + require.Nil(t, h.copyFromStdinState) + require.Equal(t, ReadyForQueryTransactionIndicator_FailedTransactionBlock, h.readyForQueryStatus()) + require.Equal(t, 1, loader.abortCalls) + require.Equal(t, 1, loader.waitCalls) + require.Same(t, ctx, loader.abortCtx) +} + +func TestCopyFailAfterServerErrorOnlyClearsTerminalState(t *testing.T) { + ctx := sql.NewEmptyContext() + loader := &handlerCopyLifecycleLoader{} + h := &ConnectionHandler{ + copyFromStdinState: ©FromStdinState{ + ctx: ctx, + dataLoader: loader, + copyErr: errors.New("server-side copy error"), + copyAborted: true, + }, + } + + stop, end, err := h.handleCopyFail(&pgproto3.CopyFail{Message: "client acknowledged"}) + require.False(t, stop) + require.False(t, end) + require.NoError(t, err) + require.Nil(t, h.copyFromStdinState) + require.Zero(t, loader.abortCalls) + require.Zero(t, loader.waitCalls) +} + +func TestCopyDataAfterServerErrorIsIgnoredUntilTerminalMessage(t *testing.T) { + copyErr := errors.New("server-side copy error") + loader := &handlerCopyLifecycleLoader{} + h := &ConnectionHandler{ + copyFromStdinState: ©FromStdinState{ + dataLoader: loader, + copyErr: copyErr, + copyAborted: true, + }, + } + + stop, end, err := h.handleCopyData(&pgproto3.CopyData{Data: []byte("ignored")}) + require.False(t, stop) + require.False(t, end) + require.NoError(t, err) + require.Same(t, loader, h.copyFromStdinState.dataLoader) + require.Zero(t, loader.loadCalls) +} + +func TestCloseBackendConnAbortsCopyLoaderBeforeConnectionRelease(t *testing.T) { + ctx := sql.NewEmptyContext() + loader := &handlerCopyLifecycleLoader{} + h := &ConnectionHandler{ + copyFromStdinState: ©FromStdinState{ + ctx: ctx, + dataLoader: loader, + }, + } + + // A minimal handler has no backend connection, but closeBackendConn still + // must drain an active COPY loader before it returns. Production handlers then + // proceed to rollback/close the physical owner. + h.closeBackendConn() + + require.Nil(t, h.copyFromStdinState) + require.Equal(t, 1, loader.abortCalls) + require.Equal(t, 1, loader.waitCalls) + require.Same(t, ctx, loader.abortCtx) +} + +func TestDiscardAllRejectsActiveTransaction(t *testing.T) { + h := &ConnectionHandler{ + txStatus: ReadyForQueryTransactionIndicator_TransactionBlock, + } + err := h.discardAll(ConvertedStatement{Tag: "DISCARD"}) + var pgErr *pgconn.PgError + require.ErrorAs(t, err, &pgErr) + require.Equal(t, "25001", pgErr.Code) + require.Equal(t, "DISCARD ALL cannot run inside a transaction block", pgErr.Message) + require.Equal(t, ReadyForQueryTransactionIndicator_TransactionBlock, h.readyForQueryStatus()) +} + +func TestDiscardAllRejectsImplicitTransactionScope(t *testing.T) { + h := &ConnectionHandler{ + implicitTx: new(stdsql.Tx), + } + err := h.discardAll(ConvertedStatement{Tag: "DISCARD"}) + var pgErr *pgconn.PgError + require.ErrorAs(t, err, &pgErr) + require.Equal(t, "25001", pgErr.Code) +} + +func TestExtendedCopyDoneSendsCommandCompleteBeforeSync(t *testing.T) { + ctx := sql.NewEmptyContext() + loader := &handlerCopyLifecycleLoader{} + var output bytes.Buffer + h := &ConnectionHandler{ + backend: pgproto3.NewBackend(bytes.NewReader(nil), &output), + waitForSync: true, + implicitTx: new(stdsql.Tx), + copyFromStdinState: ©FromStdinState{ + ctx: ctx, + dataLoader: loader, + }, + } + + stop, end, err := h.handleCopyDone(&pgproto3.CopyDone{}) + require.False(t, stop) + require.False(t, end) + require.NoError(t, err) + require.Nil(t, h.copyFromStdinState) + require.Empty(t, h.pendingCommandCompletes) + + // Extended COPY follows the same response ordering as an ordinary Execute: + // its completion is visible before a later pipelined response or Sync. + require.NoError(t, h.send(&pgproto3.ParseComplete{})) + frontend := pgproto3.NewFrontend(bytes.NewReader(output.Bytes()), io.Discard) + message, err := frontend.Receive() + require.NoError(t, err) + commandComplete, ok := message.(*pgproto3.CommandComplete) + require.True(t, ok) + require.Equal(t, "COPY 0", string(commandComplete.CommandTag)) + + message, err = frontend.Receive() + require.NoError(t, err) + _, ok = message.(*pgproto3.ParseComplete) + require.True(t, ok) +} + +func TestSyncDuringCopyIsIgnoredUntilCopyDone(t *testing.T) { + provider := catalog.NewInMemoryDBProvider() + t.Cleanup(func() { require.NoError(t, provider.Close()) }) + _, err := provider.Storage().ExecContext(context.Background(), "CREATE TABLE pg_sync_copy (id INTEGER)") + require.NoError(t, err) + + session := backend.NewSession(memory.NewSession(sql.NewBaseSession(), provider), provider) + ctx := sql.NewContext(mycontext.WithFrontendQuery(context.Background()), sql.WithSession(session)) + tx, err := adapter.GetTxn(ctx, &stdsql.TxOptions{}) + require.NoError(t, err) + _, err = tx.ExecContext(ctx, "INSERT INTO pg_sync_copy VALUES (1)") + require.NoError(t, err) + owner, bound := adapter.TryGetTxnBinding(ctx) + require.Same(t, tx, bound) + require.NotNil(t, owner) + + engine := sqle.NewDefault(provider) + t.Cleanup(func() { require.NoError(t, engine.Close()) }) + var output bytes.Buffer + loader := &handlerCopyLifecycleLoader{} + h := &ConnectionHandler{ + backend: pgproto3.NewBackend(bytes.NewReader(nil), &output), + duckHandler: &DuckHandler{e: engine}, + waitForSync: true, + implicitTx: tx, + implicitTxCtx: ctx, + implicitTxConn: owner, + copyFromStdinState: ©FromStdinState{ + ctx: ctx, + dataLoader: loader, + }, + } + + stop, end, err := h.handleMessage(&pgproto3.Sync{}) + require.False(t, stop) + require.False(t, end) + require.NoError(t, err) + require.NotNil(t, h.copyFromStdinState) + require.True(t, h.waitForSync) + require.Zero(t, loader.abortCalls) + require.Zero(t, loader.waitCalls) + _, bound = adapter.TryGetTxnBinding(ctx) + require.Same(t, tx, bound) + + // Once COPY terminates, its completion is sent immediately. A subsequent + // Sync is the transaction boundary that commits the implicit scope. + stop, end, err = h.handleCopyDone(&pgproto3.CopyDone{}) + require.False(t, stop) + require.False(t, end) + require.NoError(t, err) + require.Nil(t, h.copyFromStdinState) + require.Equal(t, 1, loader.finishCalls) + + stop, end, err = h.handleMessage(&pgproto3.Sync{}) + require.False(t, stop) + require.True(t, end) + require.NoError(t, err) + require.Nil(t, session.TryGetTxn()) + + var count int + require.NoError(t, provider.Storage().QueryRowContext(context.Background(), "SELECT count(*) FROM pg_sync_copy").Scan(&count)) + require.Equal(t, 1, count) +} + +func TestSyncDuringExplicitCopyIsIgnored(t *testing.T) { + ctx := sql.NewEmptyContext() + loader := &handlerCopyLifecycleLoader{} + h := &ConnectionHandler{ + waitForSync: true, + txStatus: ReadyForQueryTransactionIndicator_TransactionBlock, + copyFromStdinState: ©FromStdinState{ + ctx: ctx, + dataLoader: loader, + }, + } + + stop, end, err := h.handleMessage(&pgproto3.Sync{}) + require.False(t, stop) + require.False(t, end) + require.NoError(t, err) + require.NotNil(t, h.copyFromStdinState) + require.Equal(t, ReadyForQueryTransactionIndicator_TransactionBlock, h.readyForQueryStatus()) + require.Zero(t, loader.abortCalls) + require.Zero(t, loader.waitCalls) +} + +func TestCopyModeIgnoresFlushAndSync(t *testing.T) { + ctx := sql.NewEmptyContext() + loader := &handlerCopyLifecycleLoader{} + var output bytes.Buffer + h := &ConnectionHandler{ + backend: pgproto3.NewBackend(bytes.NewReader(nil), &output), + waitForSync: true, + copyFromStdinState: ©FromStdinState{ + ctx: ctx, + dataLoader: loader, + }, + } + + stop, end, err := h.handleMessage(&pgproto3.Flush{}) + require.False(t, stop) + require.False(t, end) + require.NoError(t, err) + require.NotNil(t, h.copyFromStdinState) + require.True(t, h.waitForSync) + + stop, end, err = h.handleMessage(&pgproto3.Sync{}) + require.False(t, stop) + require.False(t, end) + require.NoError(t, err) + require.NotNil(t, h.copyFromStdinState) + require.True(t, h.waitForSync) + require.Zero(t, loader.abortCalls) + require.Zero(t, loader.waitCalls) + require.Empty(t, output.Bytes()) +} + +func TestFlushIsHandledOutsideCopy(t *testing.T) { + var output bytes.Buffer + h := &ConnectionHandler{ + backend: pgproto3.NewBackend(bytes.NewReader(nil), &output), + } + + stop, end, err := h.handleMessage(&pgproto3.Flush{}) + require.False(t, stop) + require.False(t, end) + require.NoError(t, err) + require.Empty(t, output.Bytes()) +} + +func TestCopyModeRejectsNonCopyMessageAndAborts(t *testing.T) { + ctx := sql.NewEmptyContext() + loader := &handlerCopyLifecycleLoader{} + h := &ConnectionHandler{ + waitForSync: true, + copyFromStdinState: ©FromStdinState{ + ctx: ctx, + dataLoader: loader, + }, + } + + stop, end, err := h.handleMessage(&pgproto3.Parse{Query: "SELECT 1"}) + require.False(t, stop) + require.False(t, end) + var pgErr *pgconn.PgError + require.ErrorAs(t, err, &pgErr) + require.Equal(t, "08P01", pgErr.Code) + require.Contains(t, pgErr.Message, "COPY FROM STDIN") + require.NotNil(t, h.copyFromStdinState) + require.True(t, h.copyFromStdinState.copyAborted) + require.ErrorIs(t, h.copyFromStdinState.copyErr, err) + require.Equal(t, 1, loader.abortCalls) + require.Equal(t, 1, loader.waitCalls) +} + +func TestSimpleCopyModeRejectsNonCopyMessageAndClearsState(t *testing.T) { + ctx := sql.NewEmptyContext() + loader := &handlerCopyLifecycleLoader{} + h := &ConnectionHandler{ + copyFromStdinState: ©FromStdinState{ + ctx: ctx, + dataLoader: loader, + }, + } + + stop, end, err := h.handleMessage(&pgproto3.Query{String: "SELECT 1"}) + require.False(t, stop) + require.True(t, end) + var pgErr *pgconn.PgError + require.ErrorAs(t, err, &pgErr) + require.Equal(t, "08P01", pgErr.Code) + require.Nil(t, h.copyFromStdinState) + require.Equal(t, 1, loader.abortCalls) + require.Equal(t, 1, loader.waitCalls) +} + +func TestSimpleCopyErrorSentinelDoesNotSwallowNextQuery(t *testing.T) { + var output bytes.Buffer + h := &ConnectionHandler{ + backend: pgproto3.NewBackend(bytes.NewReader(nil), &output), + copyFromStdinState: ©FromStdinState{ + copyErr: errors.New("server-side copy error"), + copyAborted: true, + }, + } + + stop, end, err := h.handleMessage(&pgproto3.Query{String: ";"}) + require.False(t, stop) + require.True(t, end) + require.NoError(t, err) + require.Nil(t, h.copyFromStdinState) + + frontend := pgproto3.NewFrontend(bytes.NewReader(output.Bytes()), io.Discard) + message, err := frontend.Receive() + require.NoError(t, err) + commandComplete, ok := message.(*pgproto3.CommandComplete) + require.True(t, ok) + require.Empty(t, commandComplete.CommandTag) +} + +func TestProtocolErrorDropsDeferredCommandComplete(t *testing.T) { + var output bytes.Buffer + h := &ConnectionHandler{ + backend: pgproto3.NewBackend(bytes.NewReader(nil), &output), + pendingCommandCompletes: []*pgproto3.CommandComplete{{CommandTag: []byte("COPY 1")}}, + } + + require.NoError(t, h.handleProtocolError(errors.New("commit failed"), false)) + require.Empty(t, h.pendingCommandCompletes) + frontend := pgproto3.NewFrontend(bytes.NewReader(output.Bytes()), io.Discard) + message, err := frontend.Receive() + require.NoError(t, err) + _, ok := message.(*pgproto3.ErrorResponse) + require.True(t, ok) + message, err = frontend.Receive() + require.NoError(t, err) + _, ok = message.(*pgproto3.ReadyForQuery) + require.True(t, ok) +} diff --git a/pgserver/copy_binding.go b/pgserver/copy_binding.go new file mode 100644 index 00000000..58ebf89e --- /dev/null +++ b/pgserver/copy_binding.go @@ -0,0 +1,280 @@ +package pgserver + +import ( + stdsql "database/sql" + "fmt" + "reflect" + "sync" + "sync/atomic" + + "github.com/apecloud/myduckserver/adapter" + "github.com/cockroachdb/cockroachdb-parser/pkg/sql/sem/tree" + "github.com/dolthub/go-mysql-server/sql" +) + +// postgresCopyLease owns the two independent lifecycle resources captured for +// an asynchronous COPY exchange. The pool snapshot protects the physical +// executor only until the producer/loader has closed its child. Provider +// operation admission may outlive that child until the surrounding physical +// transaction has committed or rolled back. The pointer is shared by state, +// loader and writer copies so every terminal path can safely release either +// half exactly once. +type postgresCopyLease struct { + snapshotRelease func() + operationRelease func() + snapshotOnce sync.Once + operationOnce sync.Once + snapshotReleased atomic.Bool + operationDeferred atomic.Bool +} + +// newPostgresCopyLease keeps the one-argument form source-compatible with the +// older COPY helpers: that callback is treated as the snapshot release and +// Release still invokes it. New callers pass the provider-operation release as +// the optional second argument. +func newPostgresCopyLease(snapshotRelease func(), operationRelease ...func()) *postgresCopyLease { + var operation func() + if len(operationRelease) > 0 { + operation = operationRelease[0] + } + if snapshotRelease == nil && operation == nil { + return nil + } + return &postgresCopyLease{ + snapshotRelease: snapshotRelease, + operationRelease: operation, + } +} + +// ReleaseSnapshot ends the pool lifecycle lease. It deliberately marks the +// snapshot released only after the callback returns: a concurrent binding +// validation must continue to trust the captured identity while the pool +// reader is still held by the callback. +func (lease *postgresCopyLease) ReleaseSnapshot() { + if lease == nil || lease.snapshotRelease == nil { + return + } + lease.snapshotOnce.Do(func() { + lease.snapshotRelease() + lease.snapshotReleased.Store(true) + }) +} + +// ReleaseOperation ends provider operation admission. A physical transaction +// registrar may invoke this callback after commit/rollback; the once guard +// keeps that callback compatible with explicit teardown and fallback paths. +func (lease *postgresCopyLease) ReleaseOperation() { + if lease == nil || lease.operationRelease == nil { + return + } + lease.operationOnce.Do(lease.operationRelease) +} + +// Release releases both resources in dependency order. This is used only for +// construction/teardown paths that cannot retain a transaction-finalization +// callback; normal COPY terminal paths use ReleaseSnapshot plus a deferred +// operation release where appropriate. +func (lease *postgresCopyLease) Release() { + if lease == nil { + return + } + lease.ReleaseSnapshot() + lease.ReleaseOperation() +} + +// SnapshotHeld reports whether a pool snapshot callback is still active. It +// is used by binding validation to avoid a second lifecycle read while a +// shutdown writer is waiting on the retained snapshot. +func (lease *postgresCopyLease) SnapshotHeld() bool { + return lease != nil && lease.snapshotRelease != nil && !lease.snapshotReleased.Load() +} + +// OperationDeferred reports whether provider operation release is owned by a +// physical transaction finalizer. If registration is unavailable, callers +// must release the operation at their terminal boundary instead. +func (lease *postgresCopyLease) OperationDeferred() bool { + return lease != nil && lease.operationDeferred.Load() +} + +// registerPhysicalFinalization hands operation admission to the session's +// raw-transaction completion registry. The registry callback receives the +// commit/rollback outcome but the lease has the same release behavior for both +// outcomes. A false result leaves the caller responsible for a fallback +// terminal release. +func (lease *postgresCopyLease) registerPhysicalFinalization(ctx *sql.Context, tx *stdsql.Tx) bool { + if lease == nil || tx == nil || lease.operationRelease == nil || ctx == nil || ctx.Session == nil { + return false + } + registrar, ok := ctx.Session.(interface { + RegisterPostPhysicalTransactionCleanup(*sql.Context, *stdsql.Tx, func(bool)) bool + }) + if !ok { + return false + } + lease.operationDeferred.Store(true) + if !registrar.RegisterPostPhysicalTransactionCleanup(ctx, tx, func(bool) { + lease.ReleaseOperation() + }) { + lease.operationDeferred.Store(false) + return false + } + return true +} + +// postgresCopyBinding is the lifecycle identity captured when a COPY FROM +// STDIN exchange is accepted. COPY data arrives in later frontend messages, +// so reacquiring a connection at the first chunk can silently move execution +// outside the transaction that was active when COPY was opened. +type postgresCopyBinding struct { + execer adapter.SQLExecutor + ownerConn *stdsql.Conn + ownerTx *stdsql.Tx + snapshotCaptured bool + lease *postgresCopyLease +} + +// capturePostgresCopyBinding obtains the executor and its physical owner as a +// single snapshot. Minimal unit-test handlers may not have a session at all; +// those retain the legacy lifecycle and intentionally skip identity checks. +func capturePostgresCopyBinding(ctx *sql.Context) (postgresCopyBinding, error) { + if ctx == nil || ctx.Session == nil { + return postgresCopyBinding{}, nil + } + if _, ok := ctx.Session.(adapter.ConnectionHolder); !ok { + return postgresCopyBinding{}, nil + } + execer, ownerConn, ownerTx, err := adapter.GetExecutionSnapshot(ctx) + if err != nil { + return postgresCopyBinding{}, err + } + if execer == nil || ownerConn == nil { + return postgresCopyBinding{}, fmt.Errorf("COPY execution snapshot has no physical owner") + } + return postgresCopyBinding{ + execer: execer, + ownerConn: ownerConn, + ownerTx: ownerTx, + snapshotCaptured: true, + }, nil +} + +// capturePostgresCopyBindingWithHandler captures a retained snapshot for a +// frontend COPY exchange. Operation admission is acquired before the pool +// snapshot so rollback cleanup cannot overtake an accepted operation. The +// returned lease remains owned by the caller until the protocol has drained +// the loader and finalized its surrounding transaction. +func capturePostgresCopyBindingWithHandler( + ctx *sql.Context, + handler *DuckHandler, + statement tree.Statement, +) (postgresCopyBinding, error) { + if ctx == nil || ctx.Session == nil { + return postgresCopyBinding{}, nil + } + if _, ok := ctx.Session.(adapter.ConnectionHolder); !ok { + return postgresCopyBinding{}, nil + } + execer, ownerConn, ownerTx, lease, err := postgresCopySnapshotLease(ctx, handler, statement) + if err != nil { + return postgresCopyBinding{}, err + } + if execer == nil || ownerConn == nil { + if lease != nil { + lease.Release() + } + return postgresCopyBinding{}, fmt.Errorf("COPY execution snapshot has no physical owner") + } + return postgresCopyBinding{ + execer: execer, + ownerConn: ownerConn, + ownerTx: ownerTx, + snapshotCaptured: true, + lease: lease, + }, nil +} + +func (binding postgresCopyBinding) releaseLease() { + if binding.lease != nil { + binding.lease.Release() + } +} + +func (binding postgresCopyBinding) releaseSnapshot() { + if binding.lease != nil { + binding.lease.ReleaseSnapshot() + } +} + +func (binding postgresCopyBinding) releaseOperation() { + if binding.lease != nil { + binding.lease.ReleaseOperation() + } +} + +// releaseTerminal closes the physical snapshot and releases operation +// admission only when no transaction finalizer owns it. This is the common +// COPY error/abort boundary; a registered physical transaction callback keeps +// operation admission through COMMIT/ROLLBACK. +func (binding postgresCopyBinding) releaseTerminal() { + binding.releaseSnapshot() + if binding.lease == nil || !binding.lease.OperationDeferred() { + binding.releaseOperation() + } +} + +func (binding postgresCopyBinding) leaseHeld() bool { + return binding.lease != nil && binding.lease.SnapshotHeld() +} + +func copyExecutorForBinding(conn *stdsql.Conn, tx *stdsql.Tx) adapter.SQLExecutor { + if tx != nil { + return tx + } + return conn +} + +// sameCopyExecutor compares the concrete database/sql pointer behind the +// executor interface without assuming that every possible test double is +// directly comparable. +func sameCopyExecutor(left, right adapter.SQLExecutor) bool { + if left == nil || right == nil { + return left == nil && right == nil + } + leftType, rightType := reflect.TypeOf(left), reflect.TypeOf(right) + if leftType != rightType || !leftType.Comparable() { + return false + } + return reflect.ValueOf(left).Interface() == reflect.ValueOf(right).Interface() +} + +// validate checks both halves of the physical binding before a COPY operation +// touches the database. A missing binding and a replacement binding are both +// failures; neither is permission to acquire a fresh connection. +func (binding postgresCopyBinding) validate(ctx *sql.Context) error { + if !binding.snapshotCaptured { + return nil + } + // A retained execution lease pins the selected pool generation and the + // protocol serializes this session's transaction binding. Re-reading the + // binding here would re-enter the pool lifecycle lock; when shutdown is + // waiting, that second read can block behind the writer while the current + // lease is still held. The captured identity is therefore authoritative for + // leased COPY state. + if binding.leaseHeld() { + return nil + } + if ctx == nil || ctx.Session == nil { + return fmt.Errorf("COPY execution binding is unavailable") + } + if _, ok := ctx.Session.(adapter.ConnectionHolder); !ok { + return fmt.Errorf("COPY execution session has no connection holder") + } + currentConn, currentTx := adapter.TryGetTxnBinding(ctx) + if currentConn != binding.ownerConn || currentTx != binding.ownerTx { + return fmt.Errorf("COPY execution binding changed") + } + if expected := copyExecutorForBinding(currentConn, currentTx); !sameCopyExecutor(binding.execer, expected) { + return fmt.Errorf("COPY execution executor changed") + } + return nil +} diff --git a/pgserver/copy_binding_test.go b/pgserver/copy_binding_test.go new file mode 100644 index 00000000..46ad3d6c --- /dev/null +++ b/pgserver/copy_binding_test.go @@ -0,0 +1,105 @@ +package pgserver + +import ( + "context" + stdsql "database/sql" + "testing" + + "github.com/apecloud/myduckserver/adapter" + "github.com/dolthub/go-mysql-server/memory" + "github.com/dolthub/go-mysql-server/sql" + "github.com/stretchr/testify/require" +) + +type copyBindingTestSession struct { + *memory.Session + conn *stdsql.Conn + tx *stdsql.Tx +} + +var _ adapter.ConnectionHolder = (*copyBindingTestSession)(nil) +var _ adapter.TransactionBindingHolder = (*copyBindingTestSession)(nil) +var _ adapter.ExecutionSnapshotHolder = (*copyBindingTestSession)(nil) + +func (s *copyBindingTestSession) GetConn(context.Context) (*stdsql.Conn, error) { + return s.conn, nil +} + +func (s *copyBindingTestSession) GetTxn(context.Context, *stdsql.TxOptions) (*stdsql.Tx, error) { + return s.tx, nil +} + +func (s *copyBindingTestSession) GetCatalogConn(context.Context) (*stdsql.Conn, error) { + return s.conn, nil +} + +func (s *copyBindingTestSession) GetCatalogTxn(context.Context, *stdsql.TxOptions) (*stdsql.Tx, error) { + return s.tx, nil +} + +func (s *copyBindingTestSession) TryGetTxn() *stdsql.Tx { return s.tx } + +func (s *copyBindingTestSession) GetTxnBinding() (*stdsql.Conn, *stdsql.Tx) { + return s.conn, s.tx +} + +func (s *copyBindingTestSession) GetExecutionSnapshot(context.Context, bool) (*stdsql.Conn, *stdsql.Tx, error) { + return s.conn, s.tx, nil +} + +func (s *copyBindingTestSession) GetCurrentCatalog() string { return "memory" } + +func (s *copyBindingTestSession) GetCurrentSchema() string { return "main" } + +func (s *copyBindingTestSession) CloseTxn() {} + +func (s *copyBindingTestSession) CloseConn() {} + +func TestPostgresCopyBindingRejectsLifecycleReplacement(t *testing.T) { + ownerConn := new(stdsql.Conn) + ownerTx := new(stdsql.Tx) + session := ©BindingTestSession{ + Session: memory.NewSession(sql.NewBaseSession(), nil), + conn: ownerConn, + tx: ownerTx, + } + ctx := sql.NewContext(context.Background(), sql.WithSession(session)) + + binding, err := capturePostgresCopyBinding(ctx) + require.NoError(t, err) + require.True(t, binding.snapshotCaptured) + require.Same(t, ownerConn, binding.ownerConn) + require.Same(t, ownerTx, binding.ownerTx) + require.Same(t, ownerTx, binding.execer) + require.NoError(t, binding.validate(ctx)) + + session.tx = new(stdsql.Tx) + require.ErrorContains(t, binding.validate(ctx), "binding changed") + + session.tx = ownerTx + session.conn = new(stdsql.Conn) + require.ErrorContains(t, binding.validate(ctx), "binding changed") + + session.conn = nil + session.tx = nil + require.ErrorContains(t, binding.validate(ctx), "binding changed") +} + +func TestPostgresCopyBindingRejectsExecutorReplacement(t *testing.T) { + ownerConn := new(stdsql.Conn) + ownerTx := new(stdsql.Tx) + session := ©BindingTestSession{ + Session: memory.NewSession(sql.NewBaseSession(), nil), + conn: ownerConn, + tx: ownerTx, + } + ctx := sql.NewContext(context.Background(), sql.WithSession(session)) + binding := postgresCopyBinding{ + execer: ownerConn, + ownerConn: ownerConn, + ownerTx: ownerTx, + snapshotCaptured: true, + } + + require.ErrorContains(t, binding.validate(ctx), "executor changed") +} diff --git a/pgserver/copy_operation_lifecycle_test.go b/pgserver/copy_operation_lifecycle_test.go new file mode 100644 index 00000000..1aa5c87f --- /dev/null +++ b/pgserver/copy_operation_lifecycle_test.go @@ -0,0 +1,235 @@ +package pgserver + +import ( + "bytes" + "context" + stdsql "database/sql" + "errors" + "testing" + + "github.com/dolthub/go-mysql-server/memory" + "github.com/dolthub/go-mysql-server/sql" + "github.com/jackc/pgx/v5/pgproto3" + "github.com/stretchr/testify/require" +) + +// copyOperationLifecycleLoader is deliberately synchronous. The production +// loader may finish on a worker, but the handler must observe the same +// ownership boundary after Finish returns. +type copyOperationLifecycleLoader struct { + finishErr error + onFinish func() + onAbort func() + onWait func() +} + +func (l *copyOperationLifecycleLoader) Start() <-chan error { + ready := make(chan error) + close(ready) + return ready +} + +func (*copyOperationLifecycleLoader) LoadChunk(*sql.Context, []byte) error { return nil } + +func (l *copyOperationLifecycleLoader) Abort(*sql.Context) error { + if l.onAbort != nil { + l.onAbort() + } + return nil +} + +func (l *copyOperationLifecycleLoader) Finish(*sql.Context) (*LoadDataResults, error) { + if l.onFinish != nil { + l.onFinish() + } + if l.finishErr != nil { + return nil, l.finishErr + } + return &LoadDataResults{RowsLoaded: 1}, nil +} + +func (l *copyOperationLifecycleLoader) Wait() { + if l.onWait != nil { + l.onWait() + } +} + +// copyPhysicalCleanupRegistrar captures the callback used by a raw physical +// transaction finalizer without requiring a live database/sql driver. +type copyPhysicalCleanupRegistrar struct { + *memory.Session + callback func(bool) +} + +func (s *copyPhysicalCleanupRegistrar) RegisterPostPhysicalTransactionCleanup( + _ *sql.Context, + _ *stdsql.Tx, + callback func(bool), +) bool { + s.callback = callback + return true +} + +func TestCopyDoneReleasesSnapshotImmediatelyAfterFinish(t *testing.T) { + var events []string + lease := newPostgresCopyLease( + func() { events = append(events, "snapshot-release") }, + func() { events = append(events, "operation-release") }, + ) + loader := ©OperationLifecycleLoader{ + onFinish: func() { events = append(events, "finish") }, + } + h := &ConnectionHandler{ + backend: pgproto3.NewBackend(bytes.NewReader(nil), &bytes.Buffer{}), + txStatus: ReadyForQueryTransactionIndicator_TransactionBlock, + copyFromStdinState: ©FromStdinState{ + dataLoader: loader, + binding: postgresCopyBinding{ + lease: lease, + }, + }, + } + + stop, end, err := h.handleCopyDone(&pgproto3.CopyDone{}) + require.False(t, stop) + require.True(t, end) + require.NoError(t, err) + require.Equal(t, + []string{"finish", "snapshot-release", "operation-release"}, + events, + "COPY must release its snapshot as soon as Finish drains the loader", + ) + require.Nil(t, h.copyFromStdinState) +} + +func TestCopyDoneRegisteredOperationWaitsForPhysicalFinalization(t *testing.T) { + session := ©PhysicalCleanupRegistrar{ + Session: memory.NewSession(sql.NewBaseSession(), nil), + } + ctx := sql.NewContext(context.Background(), sql.WithSession(session)) + var snapshotCalls, operationCalls int + lease := newPostgresCopyLease( + func() { snapshotCalls++ }, + func() { operationCalls++ }, + ) + transaction := new(stdsql.Tx) + require.True(t, lease.registerPhysicalFinalization(ctx, transaction)) + require.True(t, lease.OperationDeferred()) + + h := &ConnectionHandler{ + deferredCommandComplete: true, + txStatus: ReadyForQueryTransactionIndicator_TransactionBlock, + copyFromStdinState: ©FromStdinState{ + ctx: ctx, + dataLoader: ©OperationLifecycleLoader{}, + binding: postgresCopyBinding{ + lease: lease, + }, + }, + } + + _, _, err := h.handleCopyDone(&pgproto3.CopyDone{}) + require.NoError(t, err) + require.Equal(t, 1, snapshotCalls) + require.Zero(t, operationCalls, + "COPY completion must not release an operation owned by the physical finalizer") + require.NotNil(t, session.callback) + + session.callback(false) + require.Equal(t, 1, operationCalls) + // A defensive duplicate completion must remain idempotent. + session.callback(false) + require.Equal(t, 1, operationCalls) +} + +type copyCommitLifecycleSession struct { + *replacingTransactionSession + onCommit func() +} + +func (s *copyCommitLifecycleSession) CommitTxn(tx *stdsql.Tx) error { + if s.onCommit != nil { + s.onCommit() + } + return s.replacingTransactionSession.CommitTxn(tx) +} + +func TestSimpleCopyDoneReleasesFallbackOperationAfterImplicitCommit(t *testing.T) { + owner := new(stdsql.Conn) + transaction := new(stdsql.Tx) + var events []string + base := &replacingTransactionSession{ + Session: memory.NewSession(sql.NewBaseSession(), nil), + conn: owner, + tx: transaction, + } + session := ©CommitLifecycleSession{ + replacingTransactionSession: base, + onCommit: func() { events = append(events, "commit") }, + } + h := &ConnectionHandler{ + deferredCommandComplete: true, + implicitTx: transaction, + implicitTxConn: owner, + copyFromStdinState: ©FromStdinState{ + dataLoader: ©OperationLifecycleLoader{ + onFinish: func() { events = append(events, "finish") }, + }, + binding: postgresCopyBinding{ + lease: newPostgresCopyLease( + func() { events = append(events, "snapshot-release") }, + func() { events = append(events, "operation-release") }, + ), + }, + }, + } + ctx := sql.NewContext( + withPostgresConnectionHandler(context.Background(), h), + sql.WithSession(session), + ) + h.implicitTxCtx = ctx + + _, _, err := h.handleCopyDone(&pgproto3.CopyDone{}) + require.NoError(t, err) + require.Equal(t, + []string{"finish", "snapshot-release", "commit", "operation-release"}, + events, + "simple-query COPY must finalize its implicit transaction before fallback release", + ) + require.Nil(t, h.implicitTx) + require.Equal(t, ReadyForQueryTransactionIndicator_Idle, h.readyForQueryStatus()) +} + +func TestCopyDoneErrorReleasesFallbackOperation(t *testing.T) { + finishErr := errors.New("finish failed") + var events []string + loader := ©OperationLifecycleLoader{ + finishErr: finishErr, + onFinish: func() { events = append(events, "finish") }, + onAbort: func() { events = append(events, "abort") }, + onWait: func() { events = append(events, "wait") }, + } + h := &ConnectionHandler{ + deferredCommandComplete: true, + txStatus: ReadyForQueryTransactionIndicator_TransactionBlock, + copyFromStdinState: ©FromStdinState{ + dataLoader: loader, + binding: postgresCopyBinding{ + lease: newPostgresCopyLease( + func() { events = append(events, "snapshot-release") }, + func() { events = append(events, "operation-release") }, + ), + }, + }, + } + + _, _, err := h.handleCopyDone(&pgproto3.CopyDone{}) + require.ErrorIs(t, err, finishErr) + require.Equal(t, + []string{"finish", "abort", "wait", "snapshot-release", "operation-release"}, + events, + "an unsuccessful COPY must release both fallback leases after draining the loader", + ) + require.Nil(t, h.copyFromStdinState) + require.Equal(t, ReadyForQueryTransactionIndicator_FailedTransactionBlock, h.readyForQueryStatus()) +} diff --git a/pgserver/copy_query.go b/pgserver/copy_query.go new file mode 100644 index 00000000..97af99d8 --- /dev/null +++ b/pgserver/copy_query.go @@ -0,0 +1,93 @@ +package pgserver + +import ( + "strings" + + "github.com/cockroachdb/cockroachdb-parser/pkg/sql/parser" + "github.com/cockroachdb/cockroachdb-parser/pkg/sql/sem/tree" + "github.com/jackc/pgx/v5/pgconn" +) + +// rejectCopyFromStdinTrailingStatements prevents a simple-query batch from +// entering COPY FROM STDIN mode while later SQL is still queued in the same +// Query message. The wire protocol delivers CopyData/CopyDone as separate +// messages, so executing the later statements here would run them before the +// COPY has completed. Custom COPY formats are not understood by the SQL +// parser; the lexical pass below covers those forms as well. +func rejectCopyFromStdinTrailingStatements(statements []ConvertedStatement, query string) error { + copySeen := false + for _, statement := range statements { + // Parser adapters may preserve an empty statement for a trailing + // separator. Empty queries are protocol no-ops and do not count as SQL + // queued after COPY FROM STDIN. + if isEmptyConvertedStatement(statement) { + continue + } + if copySeen { + return copyFromStdinTrailingStatementError() + } + copySeen = isCopyFromStdinStatement(statement) + } + + if rawQueryHasTrailingCopyFromStdin(query) { + return copyFromStdinTrailingStatementError() + } + return nil +} + +func isCopyFromStdinStatement(statement ConvertedStatement) bool { + if copyFrom, ok := statement.AST.(*tree.CopyFrom); ok { + return copyFrom.Stdin + } + // COPY FORMAT PARQUET/JSON/ARROW is handled by the custom parser path and + // therefore arrives here with a synthetic AST and a COPY tag. + if !strings.EqualFold(statement.Tag, "COPY") { + return false + } + _, _, _, ok := ParseCopyFrom(statement.String) + return ok +} + +// rawQueryHasTrailingCopyFromStdin handles custom COPY syntax for which +// convertQuery intentionally falls back to a synthetic statement after the +// PostgreSQL parser reports an unsupported format. SplitFirstStatement uses +// the parser's lexer, so semicolons inside quoted values/comments are not +// mistaken for statement boundaries. +func rawQueryHasTrailingCopyFromStdin(query string) bool { + copySeen := false + for len(query) > 0 { + position, ok := parser.SplitFirstStatement(query) + segment := query + if ok { + segment = query[:position] + query = query[position:] + } else { + query = "" + } + + // SplitFirstStatement includes the separator in each segment. A + // trailing `;;` therefore produces a segment containing only `;`, + // which is another separator rather than a queued statement. Strip + // separators only at the segment boundary; quoted semicolons remain + // untouched. + segment = strings.TrimSpace(RemoveComments(segment)) + if strings.Trim(segment, "; \t\r\n") == "" { + continue + } + if copySeen { + return true + } + if _, _, _, ok := ParseCopyFrom(segment); ok { + copySeen = true + } + } + return false +} + +func copyFromStdinTrailingStatementError() error { + return &pgconn.PgError{ + Severity: "ERROR", + Code: "0A000", + Message: "COPY FROM STDIN must be the last statement in a simple query", + } +} diff --git a/pgserver/copy_query_test.go b/pgserver/copy_query_test.go new file mode 100644 index 00000000..9fb7283e --- /dev/null +++ b/pgserver/copy_query_test.go @@ -0,0 +1,78 @@ +package pgserver + +import ( + "testing" + + "github.com/cockroachdb/cockroachdb-parser/pkg/sql/sem/tree" + "github.com/jackc/pgx/v5/pgconn" + "github.com/jackc/pgx/v5/pgproto3" + "github.com/stretchr/testify/require" +) + +func TestRejectCopyFromStdinTrailingStatements(t *testing.T) { + tests := []struct { + name string + query string + }{ + {name: "copy before select", query: "COPY rows FROM STDIN; SELECT 1"}, + {name: "copy between statements", query: "SELECT 1; COPY rows FROM STDIN; SELECT 2"}, + {name: "custom format", query: "COPY rows FROM STDIN (FORMAT PARQUET); SELECT 1"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + h := &ConnectionHandler{} + statements, err := h.convertQuery(tt.query) + require.NoError(t, err) + + err = rejectCopyFromStdinTrailingStatements(statements, tt.query) + var pgErr *pgconn.PgError + require.ErrorAs(t, err, &pgErr) + require.Equal(t, "0A000", pgErr.Code) + + // The guard runs before statement execution and therefore cannot leave + // the handler in COPY mode. + end, err := h.handleQuery(&pgproto3.Query{String: tt.query}) + require.True(t, end) + require.ErrorAs(t, err, &pgErr) + require.Nil(t, h.copyFromStdinState) + }) + } +} + +func TestRawQueryCopyFromStdinTrailingStatementLexing(t *testing.T) { + for _, query := range []string{ + "COPY rows FROM STDIN (FORMAT PARQUET, DELIMITER ';'); SELECT 1", + "COPY rows FROM STDIN (FORMAT PARQUET, DELIMITER ';'); SELECT ';'", + } { + require.True(t, rawQueryHasTrailingCopyFromStdin(query), query) + } + + for _, query := range []string{ + "COPY rows FROM STDIN;;", + "SELECT 'COPY rows FROM STDIN; SELECT 1'", + "COPY rows FROM 'input.csv'; SELECT 1", + } { + require.False(t, rawQueryHasTrailingCopyFromStdin(query), query) + } +} + +func TestRejectCopyFromStdinTrailingStatementsIgnoresEmptyConvertedStatements(t *testing.T) { + copyStmt := ConvertedStatement{ + AST: &tree.CopyFrom{Stdin: true}, + Tag: "COPY", + String: "COPY rows FROM STDIN", + } + for _, empty := range []ConvertedStatement{ + {}, + {String: ""}, + {String: ";"}, + } { + require.NoError(t, rejectCopyFromStdinTrailingStatements([]ConvertedStatement{copyStmt, empty}, copyStmt.String)) + } + + require.Error(t, rejectCopyFromStdinTrailingStatements([]ConvertedStatement{ + copyStmt, + {AST: &tree.Select{}, Tag: "SELECT", String: "SELECT 1"}, + }, "COPY rows FROM STDIN; SELECT 1")) +} diff --git a/pgserver/copy_stdout_error_test.go b/pgserver/copy_stdout_error_test.go new file mode 100644 index 00000000..9ec99f20 --- /dev/null +++ b/pgserver/copy_stdout_error_test.go @@ -0,0 +1,30 @@ +package pgserver + +import ( + "context" + "errors" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/require" +) + +// TestCopyToStdoutCancellationPreservesProducerError guards the error path +// used when the COPY reader cancels the context after recording its failure. +// The setup error must not mask the concrete producer/transport error. +func TestCopyToStdoutCancellationPreservesProducerError(t *testing.T) { + ctxErr := context.Canceled + setupErr := errors.New("writer setup failed") + producerErr := errors.New("copy transport failed") + + var globalErr atomic.Pointer[error] + globalErr.Store(&producerErr) + + // Keep the setup error in scope to make sure the stored error, not the + // stale local, is returned by the production cancellation helper. + got := copyToStdoutCancellationError(ctxErr, &globalErr) + + require.ErrorIs(t, got, ctxErr) + require.ErrorIs(t, got, producerErr) + require.NotErrorIs(t, got, setupErr) +} diff --git a/pgserver/dataloader.go b/pgserver/dataloader.go index 5ab3611a..5068d8dd 100644 --- a/pgserver/dataloader.go +++ b/pgserver/dataloader.go @@ -2,12 +2,15 @@ package pgserver import ( "context" + stdsql "database/sql" "errors" "fmt" "os" "strconv" "strings" + "sync" "sync/atomic" + "syscall" "github.com/apecloud/myduckserver/adapter" "github.com/apecloud/myduckserver/backend" @@ -49,98 +52,355 @@ type LoadDataResults struct { var ErrCopyAborted = fmt.Errorf("COPY operation aborted") type PipeDataLoader struct { - ctx *sql.Context - cancel context.CancelFunc - schema string - table sql.InsertableTable - columns tree.NameList - pipePath string - read func() - pipe atomic.Pointer[os.File] // for writing - blocked atomic.Bool // for writing - errPipe atomic.Pointer[os.File] // for error handling - rowCount chan int64 - err atomic.Pointer[error] - logger *logrus.Entry + ctx *sql.Context + cancel context.CancelFunc + execer adapter.SQLExecutor + ownerConn *stdsql.Conn + ownerTx *stdsql.Tx + snapshotCaptured bool + lease *postgresCopyLease + ownsLease bool + schema string + table sql.InsertableTable + target string + columns tree.NameList + pipePath string + read func() + pipe atomic.Pointer[os.File] // for writing + blocked atomic.Bool // for writing + errPipe atomic.Pointer[os.File] // for error handling + rowCount chan int64 + done chan struct{} + writerDone chan struct{} + ready chan error + startOnce sync.Once + started atomic.Bool + aborted atomic.Bool + abortOnce sync.Once + abortErr error + err atomic.Pointer[error] + logger *logrus.Entry } func (loader *PipeDataLoader) Start() <-chan error { - loader.blocked.Store(true) - - // Open the reader. - go loader.read() - - ready := make(chan error, 1) - go func() { - defer close(ready) - - // TODO(fan): If the reader fails to open the pipe, the writer will block forever. - // Open the pipe for writing. - // This operation will block until the reader opens the pipe for reading. - loader.logger.Debugf("Opening pipe for writing: %s", loader.pipePath) - pipe, err := os.OpenFile(loader.pipePath, os.O_WRONLY, os.ModeNamedPipe) - loader.blocked.Store(false) - if err != nil { - ready <- err + loader.startOnce.Do(func() { + loader.ready = make(chan error, 1) + if loader.done == nil { + loader.done = make(chan struct{}) + } + loader.writerDone = make(chan struct{}) + if loader.aborted.Load() { + close(loader.done) + close(loader.writerDone) + loader.ready <- ErrCopyAborted + close(loader.ready) + loader.releaseOwnedLease() return } - - // If the COPY operation failed to start, close the pipe and return the error. - if loader.errPipe.Load() != nil { - ready <- errors.Join(*loader.err.Load(), pipe.Close(), loader.errPipe.Load().Close()) + if err := loader.validateBinding(loader.ctx); err != nil { + loader.err.Store(&err) + close(loader.done) + close(loader.writerDone) + loader.ready <- err + close(loader.ready) + loader.releaseOwnedLease() + return + } + if loader.read == nil { + err := errors.New("COPY data loader has no reader") + loader.err.Store(&err) + close(loader.done) + close(loader.writerDone) + loader.ready <- err + close(loader.ready) + loader.releaseOwnedLease() return } - loader.pipe.Store(pipe) - }() - return ready + loader.started.Store(true) + loader.blocked.Store(true) + + // Open the reader. A panic in a loader implementation must still wake + // the protocol side and release any writer blocked on the FIFO. + go func() { + defer func() { + if recovered := recover(); recovered != nil { + err := fmt.Errorf("COPY data loader panicked: %v", recovered) + loader.err.Store(&err) + if loader.cancel != nil { + loader.cancel() + } + _ = loader.closePipe() + } + if loader.blocked.Load() { + // A reader that exits before opening the FIFO must release a + // writer blocked in OpenFile. This also covers non-panic setup + // failures in custom DataLoader implementations. + loader.unblockWriter() + } + close(loader.done) + if loader.currentError() != nil { + go func() { + loader.Wait() + loader.releaseOwnedLease() + }() + } + }() + loader.read() + }() + + go func() { + defer close(loader.writerDone) + defer close(loader.ready) + + // Open the pipe for writing. This operation blocks until the reader + // opens the FIFO, so Abort has an explicit unblock path below. + if loader.logger != nil { + loader.logger.Debugf("Opening pipe for writing: %s", loader.pipePath) + } + pipe, err := os.OpenFile(loader.pipePath, os.O_WRONLY, os.ModeNamedPipe) + loader.blocked.Store(false) + if err != nil { + loader.err.Store(&err) + if loader.cancel != nil { + loader.cancel() + } + loader.ready <- err + go func() { + loader.Wait() + loader.releaseOwnedLease() + }() + return + } + // If the COPY operation failed to start, close both ends and surface + // the original reader error. Keep this path nil-safe: a reader may + // have published an unblock descriptor before its primary error. + if errPipe := loader.errPipe.Swap(nil); errPipe != nil { + copyErr := loader.currentError() + if copyErr == nil { + copyErr = errors.New("COPY data loader failed before startup") + } + loader.ready <- errors.Join(copyErr, pipe.Close(), errPipe.Close()) + go func() { + loader.Wait() + loader.releaseOwnedLease() + }() + return + } + if copyErr := loader.currentError(); copyErr != nil { + _ = pipe.Close() + loader.ready <- copyErr + go func() { + loader.Wait() + loader.releaseOwnedLease() + }() + return + } + // Abort may race the writer opening the FIFO. Do not publish a pipe + // to a loader that has already been canceled. + if loader.aborted.Load() { + _ = pipe.Close() + loader.ready <- ErrCopyAborted + go func() { + loader.Wait() + loader.releaseOwnedLease() + }() + return + } + + loader.pipe.Store(pipe) + if loader.aborted.Load() { + if loader.pipe.CompareAndSwap(pipe, nil) { + _ = pipe.Close() + } + loader.ready <- ErrCopyAborted + go func() { + loader.Wait() + loader.releaseOwnedLease() + }() + } + }() + }) + return loader.ready } func (loader *PipeDataLoader) LoadChunk(ctx *sql.Context, data []byte) error { - if errp := loader.err.Load(); errp != nil { - return fmt.Errorf("COPY operation has been aborted: %w", *errp) + if ctx == nil { + ctx = loader.ctx + } + if err := loader.validateBinding(ctx); err != nil { + return err + } + if copyErr := loader.currentError(); copyErr != nil { + return fmt.Errorf("COPY operation has been aborted: %w", copyErr) + } + pipe := loader.pipe.Load() + if pipe == nil { + return errors.New("COPY data loader is not ready") + } + if loader.logger != nil { + loader.logger.Tracef("Copying %d bytes to pipe %s", len(data), loader.pipePath) } - loader.logger.Tracef("Copying %d bytes to pipe %s", len(data), loader.pipePath) // Write the data to the FIFO pipe. - _, err := loader.pipe.Load().Write(data) + _, err := pipe.Write(data) if err != nil { - loader.logger.Error("Copying data to pipe failed:", err) + if loader.logger != nil { + loader.logger.Error("Copying data to pipe failed:", err) + } loader.Abort(ctx) return err } - loader.logger.Tracef("Copied %d bytes to pipe %s", len(data), loader.pipePath) + if loader.logger != nil { + loader.logger.Tracef("Copied %d bytes to pipe %s", len(data), loader.pipePath) + } return nil } func (loader *PipeDataLoader) Abort(ctx *sql.Context) error { - defer os.Remove(loader.pipePath) - loader.err.Store(&ErrCopyAborted) - loader.cancel() - <-loader.rowCount // Ensure the reader has exited - return loader.pipe.Load().Close() + loader.abortOnce.Do(func() { + loader.aborted.Store(true) + loader.err.Store(&ErrCopyAborted) + if loader.cancel != nil { + loader.cancel() + } + + // Close the writer before waiting. Arrow IPC readers can be blocked in an + // operating-system read from the FIFO, which cancellation alone cannot + // interrupt; EOF is what lets the reader goroutine finish. + loader.abortErr = loader.closePipe() + if loader.blocked.Load() { + // If the writer side is still blocked in OpenFile, briefly opening a + // non-blocking reader gives it an EOF/error path so it can observe the + // aborted flag and exit. + loader.unblockWriter() + } + loader.Wait() + _ = os.Remove(loader.pipePath) + loader.releaseOwnedLease() + }) + return loader.abortErr +} + +func (loader *PipeDataLoader) currentError() error { + if errp := loader.err.Load(); errp != nil && *errp != nil { + return *errp + } + return nil +} + +func (loader *PipeDataLoader) releaseOwnedLease() { + if loader != nil && loader.ownsLease && loader.lease != nil { + loader.lease.Release() + } +} + +func (loader *PipeDataLoader) validateBinding(ctx *sql.Context) error { + if loader == nil || !loader.snapshotCaptured { + return nil + } + return (postgresCopyBinding{ + execer: loader.execer, + ownerConn: loader.ownerConn, + ownerTx: loader.ownerTx, + snapshotCaptured: true, + lease: loader.lease, + }).validate(ctx) +} + +func (loader *PipeDataLoader) closePipe() error { + pipe := loader.pipe.Swap(nil) + if pipe == nil { + return nil + } + return pipe.Close() +} + +func (loader *PipeDataLoader) unblockWriter() { + if loader.pipePath == "" { + return + } + if unblock, err := os.OpenFile(loader.pipePath, os.O_RDONLY|syscall.O_NONBLOCK, os.ModeNamedPipe); err == nil { + _ = unblock.Close() + } +} + +// Wait blocks until the reader goroutine has exited. Abort already waits for +// the row-count signal, but keeping a separate completion signal makes the +// lifecycle guarantee explicit for protocol handlers and custom loaders. +func (loader *PipeDataLoader) Wait() { + if loader.started.Load() && loader.done != nil { + <-loader.done + } + if loader.started.Load() && loader.writerDone != nil { + <-loader.writerDone + } + // The reader opens and then publishes errPipe, so the writer can wake and + // race past its initial Swap before that Store becomes visible. Once both + // goroutines are done, no later publisher exists and this final Swap closes + // any descriptor left by that race. Repeated Wait calls remain safe. + if errPipe := loader.errPipe.Swap(nil); errPipe != nil { + _ = errPipe.Close() + } } func (loader *PipeDataLoader) Finish(ctx *sql.Context) (*LoadDataResults, error) { - defer os.Remove(loader.pipePath) + defer loader.releaseOwnedLease() + defer func() { + _ = os.Remove(loader.pipePath) + }() + if ctx == nil { + ctx = loader.ctx + } + if err := loader.validateBinding(ctx); err != nil { + if loader.started.Load() { + abortErr := loader.Abort(ctx) + return nil, errors.Join(err, abortErr) + } + return nil, err + } - if errp := loader.err.Load(); errp != nil { - loader.logger.Errorln("COPY operation failed:", *errp) - return nil, *errp + if !loader.started.Load() { + if copyErr := loader.currentError(); copyErr != nil { + return nil, copyErr + } + return nil, errors.New("COPY data loader has not been started") } // Close the pipe to signal the reader to exit - if err := loader.pipe.Load().Close(); err != nil { - return nil, err + closeErr := loader.closePipe() + loader.Wait() + if loader.cancel != nil { + loader.cancel() } - rows := <-loader.rowCount + copyErr := loader.currentError() + if copyErr != nil { + if loader.logger != nil { + loader.logger.Errorln("COPY operation failed:", copyErr) + } + return nil, errors.Join(copyErr, closeErr) + } + if closeErr != nil { + return nil, closeErr + } + if loader.rowCount == nil { + return nil, errors.New("COPY data loader has no row-count channel") + } + rows, ok := <-loader.rowCount + if !ok { + if copyErr := loader.currentError(); copyErr != nil { + return nil, copyErr + } + return nil, errors.New("COPY data loader stopped without a row count") + } // Now the reader has exited, check the error again - if errp := loader.err.Load(); errp != nil { - loader.logger.Errorln("COPY operation failed:", *errp) - return nil, *errp + if copyErr := loader.currentError(); copyErr != nil { + if loader.logger != nil { + loader.logger.Errorln("COPY operation failed:", copyErr) + } + return nil, copyErr } - return &LoadDataResults{ RowsLoaded: int32(rows), }, nil @@ -158,6 +418,46 @@ func NewCsvDataLoader( schema string, table sql.InsertableTable, columns tree.NameList, options *tree.CopyOptions, rawOptions string, // For non-PG-parsable COPY FROM, unused for now ) (DataLoader, error) { + binding, err := capturePostgresCopyBindingWithHandler(ctx, handler, nil) + if err != nil { + return nil, err + } + loader, err := newCsvDataLoaderWithBinding(ctx, handler, schema, table, columns, options, rawOptions, binding) + if err != nil { + binding.releaseLease() + return nil, err + } + if pipeLoader, ok := loader.(*CsvDataLoader); ok { + pipeLoader.ownsLease = true + } + return loader, nil +} + +func newCsvDataLoaderWithBinding( + ctx *sql.Context, handler *DuckHandler, + schema string, table sql.InsertableTable, columns tree.NameList, options *tree.CopyOptions, + rawOptions string, binding postgresCopyBinding, +) (DataLoader, error) { + ownedBinding := false + if !binding.snapshotCaptured { + var err error + binding, err = capturePostgresCopyBindingWithHandler(ctx, handler, nil) + if err != nil { + return nil, err + } + ownedBinding = true + } + cleanup := true + defer func() { + if cleanup && ownedBinding { + binding.releaseLease() + } + }() + target, err := resolvePostgresCopyTarget(ctx, handler, schema, table, binding.execer, binding.ownerConn) + if err != nil { + return nil, err + } + // Create the FIFO pipe duckBuilder := handler.e.Analyzer.ExecBuilder.PriorityBuilder.(*backend.DuckBuilder) pipePath, err := duckBuilder.CreatePipe(ctx, "pg-copy-from") @@ -170,14 +470,21 @@ func NewCsvDataLoader( loader := &CsvDataLoader{ PipeDataLoader: PipeDataLoader{ - ctx: ctx, - cancel: cancel, - schema: schema, - table: table, - columns: columns, - pipePath: pipePath, - rowCount: make(chan int64, 1), - logger: ctx.GetLogger(), + ctx: ctx, + cancel: cancel, + execer: binding.execer, + ownerConn: binding.ownerConn, + ownerTx: binding.ownerTx, + snapshotCaptured: binding.snapshotCaptured, + lease: binding.lease, + schema: schema, + table: table, + target: target, + columns: columns, + pipePath: pipePath, + rowCount: make(chan int64, 1), + done: make(chan struct{}), + logger: ctx.GetLogger(), }, options: options, } @@ -185,6 +492,7 @@ func NewCsvDataLoader( loader.executeCopy(loader.buildSQL(), pipePath) } + cleanup = false return loader, nil } @@ -198,11 +506,15 @@ func (loader *CsvDataLoader) buildSQL() string { b.Grow(256) b.WriteString("COPY ") - if loader.schema != "" { + if loader.target != "" { + b.WriteString(loader.target) + } else if loader.schema != "" { b.WriteString(loader.schema) b.WriteString(".") + b.WriteString(loader.table.Name()) + } else { + b.WriteString(loader.table.Name()) } - b.WriteString(loader.table.Name()) if len(loader.columns) > 0 { b.WriteString(" (") @@ -257,23 +569,56 @@ func (loader *CsvDataLoader) buildSQL() string { } func (loader *CsvDataLoader) executeCopy(sql string, pipePath string) { - defer close(loader.rowCount) - loader.logger.Debugf("Executing COPY statement: %s", sql) - result, err := adapter.Exec(loader.ctx, sql) + defer func() { + if loader.rowCount != nil { + close(loader.rowCount) + } + }() + if err := loader.validateBinding(loader.ctx); err != nil { + loader.err.Store(&err) + if loader.blocked.Load() { + loader.unblockWriter() + } + return + } + if loader.logger != nil { + loader.logger.Debugf("Executing COPY statement: %s", sql) + } + if err := loader.validateBinding(loader.ctx); err != nil { + loader.err.Store(&err) + if loader.blocked.Load() { + loader.unblockWriter() + } + return + } + var result stdsql.Result + var err error + if loader.execer != nil { + result, err = loader.execer.ExecContext(loader.ctx, sql) + } else { + result, err = adapter.Exec(loader.ctx, sql) + } if err != nil { - loader.ctx.GetLogger().Error(err) + if loader.ctx != nil && loader.ctx.GetLogger() != nil { + loader.ctx.GetLogger().Error(err) + } loader.err.Store(&err) if loader.blocked.Load() { - // Open the pipe once to unblock the writer - pipe, _ := os.OpenFile(pipePath, os.O_RDONLY, os.ModeNamedPipe) - loader.errPipe.Store(pipe) + // Open the pipe once to unblock a writer that may still be waiting in + // OpenFile. O_NONBLOCK avoids turning this error path into a second + // permanently blocked goroutine when the FIFO has disappeared. + if unblock, openErr := os.OpenFile(pipePath, os.O_RDONLY|syscall.O_NONBLOCK, os.ModeNamedPipe); openErr == nil { + loader.errPipe.Store(unblock) + } } return } rows, err := result.RowsAffected() if err != nil { - loader.ctx.GetLogger().Error(err) + if loader.ctx != nil && loader.ctx.GetLogger() != nil { + loader.ctx.GetLogger().Error(err) + } loader.err.Store(&err) return } diff --git a/pgserver/dataloader_lifecycle_test.go b/pgserver/dataloader_lifecycle_test.go new file mode 100644 index 00000000..bc3ac00c --- /dev/null +++ b/pgserver/dataloader_lifecycle_test.go @@ -0,0 +1,206 @@ +package pgserver + +import ( + "errors" + "io" + "os" + "path/filepath" + "syscall" + "testing" + "time" + + "github.com/dolthub/go-mysql-server/sql" + "github.com/sirupsen/logrus" + "github.com/stretchr/testify/require" +) + +func newTestPipeDataLoader(t *testing.T, rows int64) *PipeDataLoader { + t.Helper() + dir := t.TempDir() + pipePath := filepath.Join(dir, "copy.pipe") + require.NoError(t, syscall.Mkfifo(pipePath, 0o600)) + + loader := &PipeDataLoader{ + pipePath: pipePath, + rowCount: make(chan int64, 1), + logger: logrus.New().WithField("test", t.Name()), + } + loader.read = func() { + reader, err := os.Open(pipePath) + if err != nil { + loader.err.Store(&err) + return + } + defer reader.Close() + _, _ = io.Copy(io.Discard, reader) + loader.rowCount <- rows + } + return loader +} + +func TestPipeDataLoaderFinishClosesAndWaits(t *testing.T) { + loader := newTestPipeDataLoader(t, 4) + ready := loader.Start() + require.NoError(t, <-ready) + require.NoError(t, loader.LoadChunk(sql.NewEmptyContext(), []byte("1\n"))) + + result, err := loader.Finish(sql.NewEmptyContext()) + require.NoError(t, err) + require.Equal(t, int32(4), result.RowsLoaded) + require.Eventually(t, func() bool { + select { + case <-loader.done: + return true + default: + return false + } + }, time.Second, time.Millisecond) + _, statErr := os.Stat(loader.pipePath) + require.ErrorIs(t, statErr, os.ErrNotExist) +} + +func TestPipeDataLoaderFinishCleansUpAfterReaderError(t *testing.T) { + loader := newTestPipeDataLoader(t, 0) + ready := loader.Start() + require.NoError(t, <-ready) + + readerErr := errors.New("reader failed") + loader.err.Store(&readerErr) + result, err := loader.Finish(sql.NewEmptyContext()) + require.Nil(t, result) + require.ErrorIs(t, err, readerErr) + select { + case <-loader.done: + case <-time.After(time.Second): + t.Fatal("Finish returned before the reader exited") + } + _, statErr := os.Stat(loader.pipePath) + require.ErrorIs(t, statErr, os.ErrNotExist) +} + +func TestPipeDataLoaderAbortIsIdempotentAndPreStartSafe(t *testing.T) { + loader := newTestPipeDataLoader(t, 0) + require.NoError(t, loader.Abort(sql.NewEmptyContext())) + require.NoError(t, loader.Abort(sql.NewEmptyContext())) + _, err := loader.Finish(sql.NewEmptyContext()) + require.Error(t, err) + + ready := loader.Start() + require.ErrorIs(t, <-ready, ErrCopyAborted) + select { + case <-loader.done: + default: + // Abort happened before Start, so no reader goroutine is expected. + } + _, statErr := os.Stat(loader.pipePath) + require.ErrorIs(t, statErr, os.ErrNotExist) +} + +func TestPipeDataLoaderStartIsIdempotentAndLoadBeforeReadyIsSafe(t *testing.T) { + loader := newTestPipeDataLoader(t, 1) + first := loader.Start() + second := loader.Start() + require.True(t, first == second) + require.NoError(t, <-first) + require.NoError(t, loader.LoadChunk(sql.NewEmptyContext(), []byte("x\n"))) + _, err := loader.Finish(sql.NewEmptyContext()) + require.NoError(t, err) + + var zero PipeDataLoader + require.NotPanics(t, func() { + require.Error(t, zero.LoadChunk(sql.NewEmptyContext(), []byte("x"))) + }) + _, finishErr := zero.Finish(sql.NewEmptyContext()) + require.Error(t, finishErr) +} + +func TestPipeDataLoaderStartWithoutReaderReturnsError(t *testing.T) { + loader := &PipeDataLoader{} + ready := loader.Start() + require.EqualError(t, <-ready, "COPY data loader has no reader") + require.EqualError(t, loader.currentError(), "COPY data loader has no reader") +} + +func TestPipeDataLoaderStartupErrorClosesUnblockPipe(t *testing.T) { + loader := newTestPipeDataLoader(t, 0) + readerErr := errors.New("reader startup failed") + loader.read = func() { + loader.err.Store(&readerErr) + unblock, err := os.OpenFile(loader.pipePath, os.O_RDONLY|syscall.O_NONBLOCK, os.ModeNamedPipe) + if err != nil { + loader.err.Store(&err) + return + } + loader.errPipe.Store(unblock) + } + t.Cleanup(func() { + if unblock := loader.errPipe.Swap(nil); unblock != nil { + _ = unblock.Close() + } + }) + + require.ErrorIs(t, <-loader.Start(), readerErr) + loader.Wait() + require.Nil(t, loader.errPipe.Load()) +} + +func TestPipeDataLoaderWaitClosesLateUnblockPipe(t *testing.T) { + loader := newTestPipeDataLoader(t, 0) + readerErr := errors.New("reader startup failed") + loader.err.Store(&readerErr) + + opened := make(chan *os.File, 1) + publish := make(chan struct{}) + loader.read = func() { + unblock, err := os.OpenFile(loader.pipePath, os.O_RDONLY|syscall.O_NONBLOCK, os.ModeNamedPipe) + if err != nil { + loader.err.Store(&err) + return + } + opened <- unblock + <-publish + loader.errPipe.Store(unblock) + } + + ready := loader.Start() + var unblock *os.File + select { + case unblock = <-opened: + case <-time.After(time.Second): + t.Fatal("reader did not open the unblock pipe") + } + require.ErrorIs(t, <-ready, readerErr) + close(publish) + loader.Wait() + require.Nil(t, loader.errPipe.Load()) + require.ErrorIs(t, unblock.Close(), os.ErrClosed) + loader.Wait() // Finalization is idempotent. +} + +func TestArrowDataLoaderMissingPipeDoesNotBlock(t *testing.T) { + loader := &ArrowDataLoader{ + PipeDataLoader: PipeDataLoader{ + ctx: sql.NewEmptyContext(), + pipePath: filepath.Join(t.TempDir(), "missing.pipe"), + rowCount: make(chan int64, 1), + logger: logrus.New().WithField("test", t.Name()), + }, + } + done := make(chan struct{}) + go func() { + loader.executeInsert("INSERT INTO missing FROM __arrow", loader.pipePath) + close(done) + }() + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("ArrowDataLoader blocked while reporting a missing FIFO") + } + require.Error(t, loader.currentError()) + select { + case _, ok := <-loader.rowCount: + require.False(t, ok) + default: + t.Fatal("ArrowDataLoader did not close its row-count channel") + } +} diff --git a/pgserver/datawriter.go b/pgserver/datawriter.go index 6797e7eb..ae7ad8cb 100644 --- a/pgserver/datawriter.go +++ b/pgserver/datawriter.go @@ -1,14 +1,17 @@ package pgserver import ( + "context" + stdsql "database/sql" "fmt" "os" "strings" + "sync" "sync/atomic" + "syscall" "github.com/apecloud/myduckserver/adapter" "github.com/apecloud/myduckserver/backend" - "github.com/apecloud/myduckserver/catalog" "github.com/cockroachdb/cockroachdb-parser/pkg/sql/sem/tree" "github.com/dolthub/go-mysql-server/sql" ) @@ -24,10 +27,19 @@ type CopyToResult struct { } type DuckDataWriter struct { - ctx *sql.Context - duckSQL string - options *tree.CopyOptions - pipePath string + ctx *sql.Context + execer adapter.SQLExecutor + ownerConn *stdsql.Conn + duckSQL string + options *tree.CopyOptions + pipePath string + cancel context.CancelFunc + done chan struct{} + started atomic.Bool + blocked atomic.Bool + close sync.Once + doneOnce sync.Once + lease *postgresCopyLease } func NewDuckDataWriter( @@ -37,6 +49,28 @@ func NewDuckDataWriter( query string, options *tree.CopyOptions, rawOptions string, ) (*DuckDataWriter, error) { + execer, ownerConn, _, lease, err := postgresCopySnapshotLease(ctx, handler, nil) + if err != nil { + return nil, err + } + cleanup := true + defer func() { + if cleanup { + lease.Release() + } + }() + copyQuery := query + if table == nil && handler != nil { + if rewritten, rewriteErr := handler.rewritePostgresObjectRelationsWithSnapshot(ctx, query, execer, ownerConn); rewriteErr != nil { + return nil, rewriteErr + } else { + copyQuery = rewritten + } + } + target, err := resolvePostgresCopyTarget(ctx, handler, schema, table, execer, ownerConn) + if err != nil { + return nil, err + } // Create the FIFO pipe db := handler.e.Analyzer.ExecBuilder.PriorityBuilder.(*backend.DuckBuilder) pipePath, err := db.CreatePipe(ctx, "pg-copy-to") @@ -51,11 +85,7 @@ func NewDuckDataWriter( builder.WriteString("COPY ") if table != nil { - if schema != "" { - builder.WriteString(catalog.QuoteIdentifierANSI(schema)) - builder.WriteString(".") - } - builder.WriteString(catalog.QuoteIdentifierANSI(table.Name())) + builder.WriteString(target) if columns != nil { builder.WriteString("(") builder.WriteString(columns.String()) @@ -63,7 +93,7 @@ func NewDuckDataWriter( } } else { // the parentheses have already been added - builder.WriteString(query) + builder.WriteString(copyQuery) } builder.WriteString(" TO '") @@ -138,24 +168,44 @@ func NewDuckDataWriter( return nil, fmt.Errorf("BINARY format is not supported for COPY TO") } - return &DuckDataWriter{ - ctx: ctx, - duckSQL: builder.String(), - options: options, - pipePath: pipePath, - }, nil + copyCtx, cancel := newCopyContext(ctx) + writer := &DuckDataWriter{ + ctx: ctx, + execer: execer, + ownerConn: ownerConn, + duckSQL: builder.String(), + options: options, + pipePath: pipePath, + cancel: cancel, + done: make(chan struct{}), + lease: lease, + } + writer.ctx = copyCtx + cleanup = false + return writer, nil } func (dw *DuckDataWriter) Start(globalErr *atomic.Pointer[error]) (string, chan CopyToResult, error) { // Execute the COPY TO statement in a separate goroutine. + dw.ensureDone() ch := make(chan CopyToResult, 1) + dw.started.Store(true) + dw.blocked.Store(true) go func() { + defer close(dw.done) defer close(ch) dw.ctx.GetLogger().Tracef("Executing COPY TO statement: %s", dw.duckSQL) // This operation will block until the reader opens the pipe for reading. - result, err := adapter.ExecCatalog(dw.ctx, dw.duckSQL) + var result stdsql.Result + var err error + if dw.execer != nil { + result, err = dw.execer.ExecContext(dw.ctx, dw.duckSQL) + } else { + result, err = adapter.ExecCatalog(dw.ctx, dw.duckSQL) + } + dw.blocked.Store(false) if err != nil { globalErr.Store(&err) ch <- CopyToResult{Err: err} @@ -169,5 +219,39 @@ func (dw *DuckDataWriter) Start(globalErr *atomic.Pointer[error]) (string, chan } func (dw *DuckDataWriter) Close() { - os.Remove(dw.pipePath) + dw.ensureDone() + dw.close.Do(func() { + if dw.cancel != nil { + dw.cancel() + } + if dw.blocked.Load() && dw.pipePath != "" { + if unblock, err := os.OpenFile(dw.pipePath, os.O_RDONLY|syscall.O_NONBLOCK, os.ModeNamedPipe); err == nil { + _ = unblock.Close() + } + } + }) + if dw.started.Load() && dw.done != nil { + <-dw.done + } + _ = os.Remove(dw.pipePath) + if dw.lease != nil { + // The physical snapshot is no longer used once the producer and pipe + // reader have both terminated. A registered operation lease remains + // admitted until the surrounding PostgreSQL transaction finalizer. + dw.lease.ReleaseSnapshot() + if !dw.lease.OperationDeferred() { + dw.lease.ReleaseOperation() + } + } +} + +func (dw *DuckDataWriter) ensureDone() { + if dw == nil { + return + } + dw.doneOnce.Do(func() { + if dw.done == nil { + dw.done = make(chan struct{}) + } + }) } diff --git a/pgserver/discard_to_sync_test.go b/pgserver/discard_to_sync_test.go new file mode 100644 index 00000000..3a937a07 --- /dev/null +++ b/pgserver/discard_to_sync_test.go @@ -0,0 +1,258 @@ +package pgserver + +import ( + "bytes" + "context" + stdsql "database/sql" + "encoding/binary" + "errors" + "io" + "sync/atomic" + "testing" + + "github.com/apecloud/myduckserver/mycontext" + "github.com/dolthub/go-mysql-server/memory" + "github.com/dolthub/go-mysql-server/sql" + "github.com/jackc/pgx/v5/pgproto3" + "github.com/stretchr/testify/require" +) + +// lifecycleAuditErrorReader lets pgproto3 decode a complete prefix and then +// surface a deterministic receive-side failure on the next message boundary. +// This models both a client EOF and a transport/decode error after an +// interrupted extended exchange. +type lifecycleAuditErrorReader struct { + remaining []byte + err error +} + +func malformedParseFrame() []byte { + // Parse contains two NUL-terminated strings followed by a two-byte + // parameter-count field. Keep the frame length valid but omit that field so + // pgproto3 returns its decoder error rather than a transport short-read. + frame := make([]byte, 7) + frame[0] = 'P' + binary.BigEndian.PutUint32(frame[1:5], 6) + copy(frame[5:], []byte{0, 0}) + return frame +} + +func (r *lifecycleAuditErrorReader) Read(p []byte) (int, error) { + if len(r.remaining) == 0 { + return 0, r.err + } + n := copy(p, r.remaining) + r.remaining = r.remaining[n:] + return n, nil +} + +// discardImplicitSession is a minimal raw-transaction session double kept in +// this file so the receive-error regression does not depend on another test +// file's uncommitted helper type. +type discardImplicitSession struct { + *memory.Session + conn *stdsql.Conn + tx *stdsql.Tx + rollbackCalled bool + callback func(bool) + events *[]string +} + +func (s *discardImplicitSession) GetConn(context.Context) (*stdsql.Conn, error) { + return s.conn, nil +} + +func (s *discardImplicitSession) GetTxn(context.Context, *stdsql.TxOptions) (*stdsql.Tx, error) { + return s.tx, nil +} + +func (s *discardImplicitSession) GetCatalogConn(context.Context) (*stdsql.Conn, error) { + return s.conn, nil +} + +func (s *discardImplicitSession) GetCatalogTxn(context.Context, *stdsql.TxOptions) (*stdsql.Tx, error) { + return s.tx, nil +} + +func (s *discardImplicitSession) TryGetTxn() *stdsql.Tx { + return s.tx +} + +func (s *discardImplicitSession) GetTxnBinding() (*stdsql.Conn, *stdsql.Tx) { + return s.conn, s.tx +} + +func (s *discardImplicitSession) GetCurrentCatalog() string { return "memory" } + +func (s *discardImplicitSession) GetCurrentSchema() string { return "main" } + +func (s *discardImplicitSession) CloseTxn() { s.tx = nil } + +func (s *discardImplicitSession) CloseConn() { s.conn = nil } + +func (s *discardImplicitSession) CommitTxn(*stdsql.Tx) error { return nil } + +func (s *discardImplicitSession) RollbackTxn(*stdsql.Tx) (*stdsql.Conn, bool, error) { + if s.events != nil { + *s.events = append(*s.events, "rollback") + } + s.rollbackCalled = true + s.tx = nil + if s.callback != nil { + callback := s.callback + s.callback = nil + callback(false) + } + return s.conn, true, nil +} + +func (s *discardImplicitSession) RegisterPostPhysicalTransactionCleanup( + _ *sql.Context, + _ *stdsql.Tx, + callback func(bool), +) bool { + s.callback = callback + return true +} + +func TestDiscardToSyncReceiveErrorCleansCopyLifecycle(t *testing.T) { + receiveErr := errors.New("copy receive failed") + loader := &handlerCopyLifecycleLoader{} + var snapshotCalls, operationCalls, pendingCalls atomic.Int32 + lease := newPostgresCopyLease( + func() { snapshotCalls.Add(1) }, + func() { operationCalls.Add(1) }, + ) + state := ©FromStdinState{ + ctx: sql.NewEmptyContext(), + dataLoader: loader, + binding: postgresCopyBinding{ + lease: lease, + }, + } + h := &ConnectionHandler{ + backend: pgproto3.NewBackend( + &lifecycleAuditErrorReader{err: receiveErr}, + io.Discard, + ), + waitForSync: true, + txStatus: ReadyForQueryTransactionIndicator_TransactionBlock, + copyFromStdinState: state, + pendingCopyRelease: func() { pendingCalls.Add(1) }, + pendingCommandCompletes: []*pgproto3.CommandComplete{ + {CommandTag: []byte("COPY 1")}, + }, + pendingProtocolMessages: []pgproto3.BackendMessage{ + &pgproto3.CommandComplete{CommandTag: []byte("COPY 1")}, + }, + } + + err := h.discardToSync() + require.ErrorIs(t, err, receiveErr) + require.Nil(t, h.copyFromStdinState) + require.False(t, h.waitForSync) + require.Equal(t, ReadyForQueryTransactionIndicator_FailedTransactionBlock, h.readyForQueryStatus()) + require.Equal(t, 1, loader.abortCalls) + require.Equal(t, 1, loader.waitCalls) + require.Equal(t, int32(1), snapshotCalls.Load()) + require.Equal(t, int32(1), operationCalls.Load()) + require.Equal(t, int32(1), pendingCalls.Load()) + require.Empty(t, h.pendingCommandCompletes) + require.Empty(t, h.pendingProtocolMessages) + + // A second teardown attempt sees no state or pending release and must not + // invoke any lifecycle callback again. + err = h.discardToSync() + require.ErrorIs(t, err, receiveErr) + require.Equal(t, 1, loader.abortCalls) + require.Equal(t, 1, loader.waitCalls) + require.Equal(t, int32(1), snapshotCalls.Load()) + require.Equal(t, int32(1), operationCalls.Load()) + require.Equal(t, int32(1), pendingCalls.Load()) +} + +func TestDiscardToSyncDecodeErrorRollsBackImplicitCopy(t *testing.T) { + malformed := malformedParseFrame() + probe := pgproto3.NewBackend(bytes.NewReader(malformed), io.Discard) + _, expectedDecodeErr := probe.Receive() + require.Error(t, expectedDecodeErr) + owner := new(stdsql.Conn) + tx := new(stdsql.Tx) + var events []string + session := &discardImplicitSession{ + Session: memory.NewSession(sql.NewBaseSession(), nil), + conn: owner, + tx: tx, + events: &events, + } + ctx := sql.NewContext( + mycontext.WithFrontendQuery(context.Background()), + sql.WithSession(session), + ) + loader := &handlerCopyLifecycleLoader{} + var snapshotCalls, operationCalls, pendingCalls atomic.Int32 + lease := newPostgresCopyLease( + func() { + snapshotCalls.Add(1) + events = append(events, "snapshot-release") + }, + func() { + operationCalls.Add(1) + events = append(events, "operation-release") + }, + ) + require.True(t, lease.registerPhysicalFinalization(ctx, tx)) + h := &ConnectionHandler{ + backend: pgproto3.NewBackend(bytes.NewReader(malformed), io.Discard), + waitForSync: true, + implicitTx: tx, + implicitTxCtx: ctx, + implicitTxConn: owner, + copyFromStdinState: ©FromStdinState{ctx: ctx, dataLoader: loader, binding: postgresCopyBinding{lease: lease}}, + pendingCopyRelease: func() { + pendingCalls.Add(1) + events = append(events, "pending-release") + }, + pendingCommandCompletes: []*pgproto3.CommandComplete{ + {CommandTag: []byte("COPY 1")}, + }, + txStatus: ReadyForQueryTransactionIndicator_Idle, + } + + err := h.discardToSync() + require.EqualError(t, err, expectedDecodeErr.Error()) + require.Nil(t, h.copyFromStdinState) + require.Nil(t, h.implicitTx) + require.Nil(t, session.tx) + require.True(t, session.rollbackCalled) + require.False(t, h.waitForSync) + require.Equal(t, ReadyForQueryTransactionIndicator_Idle, h.readyForQueryStatus()) + require.Equal(t, 1, loader.abortCalls) + require.Equal(t, 1, loader.waitCalls) + require.Equal(t, int32(1), snapshotCalls.Load()) + require.Equal(t, int32(1), operationCalls.Load()) + require.Equal(t, int32(1), pendingCalls.Load()) + require.Empty(t, h.pendingCommandCompletes) + require.Empty(t, h.pendingProtocolMessages) + require.Equal(t, []string{"snapshot-release", "pending-release", "rollback", "operation-release"}, events) +} + +func TestDiscardToSyncEOFWithoutCopyStillReleasesPendingLease(t *testing.T) { + leaseErr := io.ErrUnexpectedEOF + var pendingCalls atomic.Int32 + h := &ConnectionHandler{ + backend: pgproto3.NewBackend(&lifecycleAuditErrorReader{err: io.EOF}, io.Discard), + waitForSync: true, + txStatus: ReadyForQueryTransactionIndicator_TransactionBlock, + pendingCopyRelease: func() { + pendingCalls.Add(1) + }, + } + + err := h.discardToSync() + require.ErrorIs(t, err, leaseErr) + require.False(t, h.waitForSync) + require.Equal(t, ReadyForQueryTransactionIndicator_FailedTransactionBlock, h.readyForQueryStatus()) + require.Equal(t, int32(1), pendingCalls.Load()) + require.Nil(t, h.pendingCopyRelease) +} diff --git a/pgserver/duck_handler.go b/pgserver/duck_handler.go index 3b79c821..3c8d4216 100644 --- a/pgserver/duck_handler.go +++ b/pgserver/duck_handler.go @@ -19,11 +19,13 @@ import ( stdsql "database/sql" "database/sql/driver" "encoding/base64" + "errors" "fmt" "io" "os" "regexp" "runtime/trace" + "strings" "sync" "time" @@ -32,6 +34,7 @@ import ( "github.com/apecloud/myduckserver/catalog" "github.com/apecloud/myduckserver/mycontext" "github.com/apecloud/myduckserver/pgtypes" + "github.com/cockroachdb/cockroachdb-parser/pkg/sql/parser" "github.com/cockroachdb/cockroachdb-parser/pkg/sql/sem/tree" sqle "github.com/dolthub/go-mysql-server" "github.com/dolthub/go-mysql-server/server" @@ -120,6 +123,12 @@ func (h *DuckHandler) ComBind(ctx context.Context, c *mysql.Conn, prepared Prepa } } + // PostgreSQL prepared statements retain only SQL and metadata between Parse + // and Execute. The raw DuckDB statement used during Parse is closed before + // returning, so binding is validated at Execute on the active executor. + if prepared.Stmt == nil { + return prepared.ReturnFields, nil + } err := prepared.Stmt.Bind(vars) if err != nil { return nil, err @@ -168,47 +177,74 @@ func (h *DuckHandler) ComPrepareParsed(ctx context.Context, c *mysql.Conn, query return nil, nil, nil, err } - conn, err := adapter.GetConn(sqlCtx) + // Transaction control is executed by the protocol handler rather than + // DuckDB's prepared-statement machinery. Keeping it unprepared also avoids + // probing a second connection while a session transaction is active. + switch parsed.(type) { + case *tree.BeginTransaction, *tree.CommitTransaction, *tree.RollbackTransaction: + return nil, nil, nil, nil + } + + // Admit the operation before retaining the executor/physical owner snapshot. + // The release is deferred because all metadata probes and result-row schema + // reads complete before Parse returns. + execer, conn, tx, release, err := h.getPostgresExecutionSnapshotLease(sqlCtx, parsed) + if err != nil { + return nil, nil, nil, err + } + defer release() + + prepareQuery, err := postgresCreateTableQueryForPrepare(query, parsed) + if err != nil { + return nil, nil, nil, err + } + // Resolve object-table relations against the same lifecycle snapshot used + // for raw preparation and metadata probing. The physical relation can be + // visible only after the session's DuckLake connection is initialized; doing + // this before the probe keeps prepared INSERT ... RETURNING metadata aligned + // with the relation executed later by the bound plan. + prepareQuery, err = h.rewritePostgresObjectRelationsWithSnapshot(sqlCtx, prepareQuery, execer, conn) if err != nil { return nil, nil, nil, err } var ( - stmt *duckdb.Stmt stmtType duckdb.StmtType paramTypes []duckdb.Type ) - // This is a bit of a hack to get DuckDB's underlying prepared statement. - // But we know that the connection is a DuckDB connection and it is kept alive. - err = conn.Raw(func(driverConn interface{}) error { - dc := driverConn.(*duckdb.Conn) - s, err := dc.PrepareContext(sqlCtx, query) - if err != nil { - return err - } - n := s.NumInput() - stmt = s.(*duckdb.Stmt) - stmtType, err = stmt.StatementType() - if err != nil { - return err - } - paramTypes = make([]duckdb.Type, n) - for i := 0; i < n; i++ { - paramTypes[i], err = stmt.ParamType(i + 1) // 1-based index - if err != nil { - return err - } - } - return nil - }) + if tx != nil { + // Parse runs after the protocol has opened the implicit transaction. A + // Conn.Raw probe here would issue driver calls beside that *sql.Tx and + // can observe a different transaction state (or deadlock the driver). + // Use the parsed AST/tag and a parser/lexical placeholder count instead; + // unknown parameter OIDs are represented by TYPE_INVALID below. + stmtType = postgresDuckDBStatementType(parsed, prepareQuery) + paramTypes = make([]duckdb.Type, postgresPlaceholderCount(query)) + } else { + // There is no transaction owner, so a short-lived raw probe is safe and + // preserves DuckDB's exact statement/parameter type information. The + // retained snapshot itself prevents a transaction or replacement + // connection from appearing while the probe is in flight. + stmtType, paramTypes, err = prepareDuckDBStatementMetadata(sqlCtx, conn, prepareQuery) + } if err != nil { logrus.WithField("query", catalog.RedactSensitiveSQL(query)).Errorf("unable to prepare query: %s", err.Error()) return nil, nil, nil, err } - paramOIDs := make([]uint32, len(paramTypes)) for i, t := range paramTypes { paramOIDs[i] = pgtypes.DuckdbTypeToPostgresOID[t] + // A transaction-safe Parse path cannot call DuckDB's raw + // PrepareContext beside the active *sql.Tx, so an untyped + // placeholder has no driver-derived OID. PostgreSQL's protocol uses + // OID zero to mean "unspecified"; pgx then selects an encoder from + // the concrete Go value (rather than rejecting it as UnknownOID or + // assuming the value is a string). The server-side bind decoder + // intentionally converts that wire value to text before DuckDB's + // statement-context coercion. + if paramOIDs[i] == pgtype.UnknownOID { + paramOIDs[i] = 0 + } } var ( @@ -216,13 +252,17 @@ func (h *DuckHandler) ComPrepareParsed(ctx context.Context, c *mysql.Conn, query rows *stdsql.Rows ) if insert, ok := postgresInsertReturningRows(parsed); ok { - schema, schemaErr := inferPostgresInsertReturningSchema(sqlCtx, insert) + metadataQuery := postgresInsertReturningSchemaQuery(insert) + metadataQuery, routeErr := h.rewritePostgresObjectRelationsWithSnapshot(sqlCtx, metadataQuery, execer, conn) + if routeErr != nil { + return nil, nil, nil, routeErr + } + schema, schemaErr := inferPostgresInsertReturningSchema(sqlCtx, metadataQuery, execer) if schemaErr != nil { - defer stmt.Close() return nil, nil, nil, schemaErr } fields = schemaToFieldDescriptions(sqlCtx, schema, nil, ExtendedQueryMode) - return stmt, paramOIDs, fields, nil + return nil, paramOIDs, fields, nil } switch stmtType { case duckdb.STATEMENT_TYPE_SELECT, @@ -232,14 +272,18 @@ func (h *DuckHandler) ComPrepareParsed(ctx context.Context, c *mysql.Conn, query duckdb.STATEMENT_TYPE_EXPLAIN: // Execute the query with all NULL values as parameters to get the result types. - query := query + // Use the rewritten preparation query for both the LIMIT 0 probe and + // the later Execute path. Object-table relations only exist under + // their durable DuckLake name; probing the original logical relation + // would fail during Parse even though Execute correctly routes it. + query := prepareQuery if stmtType == duckdb.STATEMENT_TYPE_SELECT || stmtType == duckdb.STATEMENT_TYPE_RELATION { // Add LIMIT 0 to avoid executing the actual query. query = "SELECT * FROM (" + sql.RemoveSpaceAndDelimiter(query, ';') + ") LIMIT 0" } params := make([]any, len(paramTypes)) // all nil - rows, err = conn.QueryContext(sqlCtx, query, params...) + rows, err = execer.QueryContext(sqlCtx, query, params...) if err != nil { break } @@ -260,11 +304,10 @@ func (h *DuckHandler) ComPrepareParsed(ctx context.Context, c *mysql.Conn, query } } if err != nil { - defer stmt.Close() return nil, nil, nil, err } - return stmt, paramOIDs, fields, nil + return nil, paramOIDs, fields, nil } // ComQuery implements the Handler interface. @@ -294,6 +337,178 @@ func (h *DuckHandler) rejectReadOnly(statement ConvertedStatement) error { return h.connectionHandler.rejectReadOnly(statement) } +// isPostgresTransactionControl identifies statements whose lifecycle is owned +// by the protocol transaction handler. They must not acquire a DuckLake +// operation or execution snapshot: COMMIT/ROLLBACK in particular must remain +// able to recover a failed transaction while the ordinary execution path is +// unavailable. +func isPostgresTransactionControl(statement tree.Statement) bool { + switch statement.(type) { + case *tree.BeginTransaction, *tree.CommitTransaction, *tree.RollbackTransaction: + return true + default: + return false + } +} + +// postgresDuckLakeOperationProvider is the legacy admission surface retained +// for older provider/test doubles. Production providers also implement the +// owner-aware extension below. +type postgresDuckLakeOperationProvider interface { + BeginDuckLakeOperation(context.Context) func() +} + +type postgresDuckLakeOperationOwnerProvider interface { + BeginDuckLakeOperationForOwner(context.Context, any) func() +} + +// Error-bearing admission surfaces are implemented by the production +// provider. Legacy provider/test doubles continue to use the closure-only +// interfaces above; those adapters cannot report a poisoned generation. +type postgresDuckLakeOperationProviderWithError interface { + BeginDuckLakeOperationWithError(context.Context) (func(), error) +} + +type postgresDuckLakeOperationOwnerProviderWithError interface { + BeginDuckLakeOperationForOwnerWithError(context.Context, any) (func(), error) +} + +// beginPostgresDuckLakeOperationForProvider admits one frontend statement +// before it captures a pool execution snapshot. Passing the physical owner is +// important for PostgreSQL: its protocol transaction controls keep a raw +// *database/sql.Tx in the pool/handler rather than a GMS sql.Transaction, so +// the legacy context-derived owner would be anonymous and rollback cleanup +// would wait on its own operation lease. +func beginPostgresDuckLakeOperationForProvider( + ctx *sql.Context, + statement tree.Statement, + provider postgresDuckLakeOperationProvider, + owner any, +) func() { + release, _ := beginPostgresDuckLakeOperationForProviderWithError(ctx, statement, provider, owner) + return release +} + +func beginPostgresDuckLakeOperationForProviderWithError( + ctx *sql.Context, + statement tree.Statement, + provider postgresDuckLakeOperationProvider, + owner any, +) (func(), error) { + if ctx == nil || isPostgresTransactionControl(statement) || + mycontext.QueryOrigin(ctx) != mycontext.FrontendQueryOrigin || provider == nil { + return func() {}, nil + } + if ownerProvider, ok := provider.(postgresDuckLakeOperationOwnerProviderWithError); ok { + return ownerProvider.BeginDuckLakeOperationForOwnerWithError(ctx, owner) + } + if providerWithError, ok := provider.(postgresDuckLakeOperationProviderWithError); ok { + return providerWithError.BeginDuckLakeOperationWithError(ctx) + } + if ownerProvider, ok := provider.(postgresDuckLakeOperationOwnerProvider); ok { + return ownerProvider.BeginDuckLakeOperationForOwner(ctx, owner), nil + } + return provider.BeginDuckLakeOperation(ctx), nil +} + +// postgresDuckLakeOperationOwner resolves the same physical transaction +// identity that PostgreSQL rollback finalization will pass to the provider +// barrier. The handler's implicit marker is preferred because it is already +// captured for the current protocol scope; the atomic session binding covers +// explicit transactions and minimal handler contexts. A GMS logical +// transaction remains a final fallback for test/session doubles that do not +// expose a physical binding. +func postgresDuckLakeOperationOwner(ctx *sql.Context, handler *DuckHandler) any { + if handler != nil && handler.connectionHandler != nil && handler.connectionHandler.implicitTx != nil { + return handler.connectionHandler.implicitTx + } + if ctx == nil || ctx.Session == nil { + return nil + } + if _, ok := ctx.Session.(adapter.ConnectionHolder); ok { + if _, tx := adapter.TryGetTxnBinding(ctx); tx != nil { + return tx + } + } + return ctx.GetTransaction() +} + +// beginPostgresDuckLakeOperationForOwner is the explicit-owner entry point for +// asynchronous protocol paths that already captured a transaction identity. +// Keeping it separate lets COPY/extended callers pass their retained physical +// owner without re-reading a mutable session binding. +func (h *DuckHandler) beginPostgresDuckLakeOperationForOwner( + ctx *sql.Context, + statement tree.Statement, + owner any, +) func() { + release, _ := h.beginPostgresDuckLakeOperationForOwnerWithError(ctx, statement, owner) + return release +} + +func (h *DuckHandler) beginPostgresDuckLakeOperationForOwnerWithError( + ctx *sql.Context, + statement tree.Statement, + owner any, +) (func(), error) { + if h == nil || h.e == nil { + return func() {}, nil + } + provider := h.GetCatalogProvider() + return beginPostgresDuckLakeOperationForProviderWithError(ctx, statement, provider, owner) +} + +// beginPostgresDuckLakeOperation admits one frontend statement before it +// captures a pool execution snapshot. The provider itself fail-closes for +// replication/maintenance origins and disabled storage; the explicit origin +// check here keeps protocol code from ever registering those paths as logical +// frontend work. +func (h *DuckHandler) beginPostgresDuckLakeOperation(ctx *sql.Context, statement tree.Statement) func() { + return h.beginPostgresDuckLakeOperationForOwner(ctx, statement, postgresDuckLakeOperationOwner(ctx, h)) +} + +func (h *DuckHandler) beginPostgresDuckLakeOperationWithError(ctx *sql.Context, statement tree.Statement) (func(), error) { + return h.beginPostgresDuckLakeOperationForOwnerWithError(ctx, statement, postgresDuckLakeOperationOwner(ctx, h)) +} + +// combinePostgresExecutionReleases closes a retained execution snapshot before +// releasing the logical-operation admission. Keeping the pool generation alive +// until the child iterator is closed is the important ordering guarantee; the +// once guard also covers legacy holders whose callbacks are not idempotent. +func combinePostgresExecutionReleases(snapshotRelease, operationRelease func()) func() { + var once sync.Once + return func() { + once.Do(func() { + if snapshotRelease != nil { + snapshotRelease() + } + if operationRelease != nil { + operationRelease() + } + }) + } +} + +// getPostgresExecutionSnapshotLease admits a frontend operation before taking +// the retained pool snapshot. Callers that return an iterator transfer the +// returned release callback to backend.WrapRowIterWithRelease; metadata-only +// callers defer it locally. +func (h *DuckHandler) getPostgresExecutionSnapshotLease( + ctx *sql.Context, + statement tree.Statement, +) (adapter.SQLExecutor, *stdsql.Conn, *stdsql.Tx, func(), error) { + operationRelease, admissionErr := h.beginPostgresDuckLakeOperationWithError(ctx, statement) + if admissionErr != nil { + return nil, nil, nil, func() {}, admissionErr + } + execer, conn, tx, snapshotRelease, err := adapter.GetExecutionSnapshotWithLease(ctx) + if err != nil { + operationRelease() + return nil, nil, tx, func() {}, err + } + return execer, conn, tx, combinePostgresExecutionReleases(snapshotRelease, operationRelease), nil +} + // ComResetConnection implements the Handler interface. func (h *DuckHandler) ComResetConnection(c *mysql.Conn) error { logrus.WithField("connectionId", c.ConnectionID).Debug("COM_RESET_CONNECTION command received") @@ -347,23 +562,280 @@ func (h *DuckHandler) getStatementTag(mysqlConn *mysql.Conn, query string) (stri if err != nil { return "", err } - conn, err := adapter.GetConn(sqlCtx) + _, conn, tx, release, err := h.getPostgresExecutionSnapshotLease(sqlCtx, nil) if err != nil { return "", err } + defer release() + if tx != nil { + // Statement-tag lookup is called after the protocol has opened its + // implicit transaction. Do not prepare through Conn.Raw beside that + // transaction; the lexical tag is sufficient for CommandComplete. + return GuessStatementTag(query), nil + } var tag string - err = conn.Raw(func(driverConn any) error { - c := driverConn.(*duckdb.Conn) - s, err := c.PrepareContext(sqlCtx, query) - if err != nil { - return err + stmtType, _, prepareErr := prepareDuckDBStatementMetadata(sqlCtx, conn, query) + if prepareErr == nil { + tag = duckDBStatementTag(stmtType) + } + err = prepareErr + return tag, err +} + +// assertPostgresExecutionSnapshot verifies that a metadata probe did not race +// a transaction or connection replacement. The check is intentionally +// fail-closed: reacquiring a different connection would make the caller lose +// the transaction/physical-owner affinity it captured at the start. +func assertPostgresExecutionSnapshot(ctx *sql.Context, expectedConn *stdsql.Conn, expectedTx *stdsql.Tx) error { + if ctx == nil || expectedConn == nil { + return fmt.Errorf("PostgreSQL execution snapshot is unavailable") + } + _, actualConn, actualTx, err := adapter.GetExecutionSnapshot(ctx) + if err != nil { + return err + } + if actualConn != expectedConn || actualTx != expectedTx { + return fmt.Errorf("PostgreSQL execution snapshot changed during metadata probe") + } + return nil +} + +// prepareDuckDBStatementMetadata performs the driver-specific probe used only +// when no session transaction is active. The returned statement is always +// closed before the function returns; callers retain no raw driver state +// across protocol messages. +func prepareDuckDBStatementMetadata(ctx context.Context, conn *stdsql.Conn, query string) (stmtType duckdb.StmtType, paramTypes []duckdb.Type, err error) { + if conn == nil { + return duckdb.STATEMENT_TYPE_INVALID, nil, fmt.Errorf("DuckDB execution connection is unavailable") + } + err = conn.Raw(func(driverConn any) (callbackErr error) { + dc, ok := driverConn.(*duckdb.Conn) + if !ok { + return fmt.Errorf("prepared PostgreSQL statement connection has unexpected driver type") + } + driverStmt, prepareErr := dc.PrepareContext(ctx, query) + if prepareErr != nil { + return prepareErr + } + defer func() { + callbackErr = errors.Join(callbackErr, driverStmt.Close()) + }() + stmt, ok := driverStmt.(*duckdb.Stmt) + if !ok { + return fmt.Errorf("prepared PostgreSQL statement has unexpected driver statement type") + } + stmtType, prepareErr = stmt.StatementType() + if prepareErr != nil { + return prepareErr + } + paramTypes = make([]duckdb.Type, driverStmt.NumInput()) + for i := range paramTypes { + paramTypes[i], prepareErr = stmt.ParamType(i + 1) // 1-based index + if prepareErr != nil { + return prepareErr + } } - defer s.Close() - stmt := s.(*duckdb.Stmt) - tag = GetStatementTag(stmt) return nil }) - return tag, err + return stmtType, paramTypes, err +} + +// postgresDuckDBStatementType classifies a statement without touching a raw +// driver connection. Parsing the SQL again, when possible, distinguishes a +// real WITH/INSERT (or other CTE statement) from the synthetic SELECT AST used +// for DuckDB-only syntax. The supplied AST remains the fallback for rewritten +// SQL that the PostgreSQL parser cannot consume a second time. +func postgresDuckDBStatementType(parsed tree.Statement, query string) duckdb.StmtType { + parseQuery := sql.RemoveSpaceAndDelimiter(query, ';') + if parseQuery != "" { + if reparsed, parseErr := parser.ParseOne(parseQuery); parseErr == nil && reparsed.AST != nil { + parsed = reparsed.AST + } + } + if parsed != nil { + if stmtType := duckDBStatementTypeForTag(parsed.StatementTag()); stmtType != duckdb.STATEMENT_TYPE_INVALID { + return stmtType + } + } + return duckDBStatementTypeForTag(GuessStatementTag(query)) +} + +func duckDBStatementTypeForTag(tag string) duckdb.StmtType { + tag = strings.ToUpper(strings.TrimSpace(tag)) + switch { + case tag == "SELECT", tag == "SHOW", tag == "VALUES", tag == "WITH", tag == "FROM", strings.HasPrefix(tag, "SHOW "): + return duckdb.STATEMENT_TYPE_SELECT + case tag == "CALL", strings.HasPrefix(tag, "CALL "): + return duckdb.STATEMENT_TYPE_CALL + case tag == "PRAGMA", strings.HasPrefix(tag, "PRAGMA "): + return duckdb.STATEMENT_TYPE_PRAGMA + case tag == "EXPLAIN", strings.HasPrefix(tag, "EXPLAIN "): + return duckdb.STATEMENT_TYPE_EXPLAIN + case tag == "INSERT", strings.HasPrefix(tag, "INSERT "): + return duckdb.STATEMENT_TYPE_INSERT + case tag == "UPDATE", strings.HasPrefix(tag, "UPDATE "): + return duckdb.STATEMENT_TYPE_UPDATE + case tag == "DELETE", strings.HasPrefix(tag, "DELETE "): + return duckdb.STATEMENT_TYPE_DELETE + case tag == "COPY", strings.HasPrefix(tag, "COPY "): + return duckdb.STATEMENT_TYPE_COPY + case tag == "ALTER", strings.HasPrefix(tag, "ALTER "): + return duckdb.STATEMENT_TYPE_ALTER + case tag == "CREATE", strings.HasPrefix(tag, "CREATE "): + return duckdb.STATEMENT_TYPE_CREATE + case tag == "DROP", strings.HasPrefix(tag, "DROP "): + return duckdb.STATEMENT_TYPE_DROP + case tag == "PREPARE", strings.HasPrefix(tag, "PREPARE "): + return duckdb.STATEMENT_TYPE_PREPARE + case tag == "EXECUTE", strings.HasPrefix(tag, "EXECUTE "): + return duckdb.STATEMENT_TYPE_EXECUTE + case tag == "ATTACH", strings.HasPrefix(tag, "ATTACH "): + return duckdb.STATEMENT_TYPE_ATTACH + case tag == "DETACH", strings.HasPrefix(tag, "DETACH "): + return duckdb.STATEMENT_TYPE_DETACH + case tag == "TRANSACTION", tag == "BEGIN", tag == "COMMIT", tag == "ROLLBACK": + return duckdb.STATEMENT_TYPE_TRANSACTION + case tag == "ANALYZE", strings.HasPrefix(tag, "ANALYZE "): + return duckdb.STATEMENT_TYPE_ANALYZE + case tag == "SET", strings.HasPrefix(tag, "SET "): + return duckdb.STATEMENT_TYPE_SET + case tag == "LOAD", strings.HasPrefix(tag, "LOAD "): + return duckdb.STATEMENT_TYPE_LOAD + case tag == "EXPORT", strings.HasPrefix(tag, "EXPORT "): + return duckdb.STATEMENT_TYPE_EXPORT + default: + return duckdb.STATEMENT_TYPE_INVALID + } +} + +func duckDBStatementTag(stmtType duckdb.StmtType) string { + switch stmtType { + case duckdb.STATEMENT_TYPE_SELECT, duckdb.STATEMENT_TYPE_RELATION: + return "SELECT" + case duckdb.STATEMENT_TYPE_INSERT: + return "INSERT" + case duckdb.STATEMENT_TYPE_UPDATE: + return "UPDATE" + case duckdb.STATEMENT_TYPE_DELETE: + return "DELETE" + case duckdb.STATEMENT_TYPE_CALL: + return "CALL" + case duckdb.STATEMENT_TYPE_PRAGMA: + return "PRAGMA" + case duckdb.STATEMENT_TYPE_COPY: + return "COPY" + case duckdb.STATEMENT_TYPE_ALTER: + return "ALTER" + case duckdb.STATEMENT_TYPE_CREATE, duckdb.STATEMENT_TYPE_CREATE_FUNC: + return "CREATE" + case duckdb.STATEMENT_TYPE_DROP: + return "DROP" + case duckdb.STATEMENT_TYPE_PREPARE: + return "PREPARE" + case duckdb.STATEMENT_TYPE_EXECUTE: + return "EXECUTE" + case duckdb.STATEMENT_TYPE_ATTACH: + return "ATTACH" + case duckdb.STATEMENT_TYPE_DETACH: + return "DETACH" + case duckdb.STATEMENT_TYPE_TRANSACTION: + return "TRANSACTION" + case duckdb.STATEMENT_TYPE_ANALYZE: + return "ANALYZE" + case duckdb.STATEMENT_TYPE_EXPLAIN: + return "EXPLAIN" + case duckdb.STATEMENT_TYPE_SET: + return "SET" + case duckdb.STATEMENT_TYPE_VARIABLE_SET: + return "SET VARIABLE" + case duckdb.STATEMENT_TYPE_EXPORT: + return "EXPORT" + case duckdb.STATEMENT_TYPE_LOAD: + return "LOAD" + default: + return "UNKNOWN" + } +} + +// postgresPlaceholderCount returns PostgreSQL's highest placeholder index. A +// parser result is authoritative for valid PostgreSQL syntax; the lexical +// fallback covers DuckDB-only statements that are intentionally represented by +// a synthetic AST during protocol parsing. +func postgresPlaceholderCount(query string) int { + parseQuery := sql.RemoveSpaceAndDelimiter(query, ';') + if parseQuery != "" { + if parsed, parseErr := parser.ParseOne(parseQuery); parseErr == nil { + return parsed.NumPlaceholders + } + } + return lexicalPostgresPlaceholderCount(query) +} + +func lexicalPostgresPlaceholderCount(query string) int { + maxIndex, questionCount := 0, 0 + for i := 0; i < len(query); { + switch query[i] { + case '\'', '"', '`': + quote := query[i] + i++ + for i < len(query) { + if query[i] == '\\' && quote != '`' && i+1 < len(query) { + i += 2 + continue + } + if query[i] == quote { + if i+1 < len(query) && query[i+1] == quote { + i += 2 + continue + } + i++ + break + } + i++ + } + continue + case '-': + if i+1 < len(query) && query[i+1] == '-' { + i += 2 + for i < len(query) && query[i] != '\n' { + i++ + } + continue + } + case '/': + if i+1 < len(query) && query[i+1] == '*' { + if end := strings.Index(query[i+2:], "*/"); end >= 0 { + i += end + 4 + } else { + return maxIndex + } + continue + } + case '$': + if i+1 < len(query) && query[i+1] >= '0' && query[i+1] <= '9' { + j := i + 1 + for j < len(query) && query[j] >= '0' && query[j] <= '9' { + j++ + } + var index int + for k := i + 1; k < j; k++ { + index = index*10 + int(query[k]-'0') + } + if index > maxIndex { + maxIndex = index + } + i = j + continue + } + case '?': + questionCount++ + } + i++ + } + if questionCount > maxIndex { + return questionCount + } + return maxIndex } var queryLoggingRegex = regexp.MustCompile(`[\r\n\t ]+`) @@ -437,13 +909,32 @@ func (h *DuckHandler) doQuery(ctx context.Context, c *mysql.Conn, query string, sqlCtx.GetLogger().WithError(err).Warn("error running query") return err } + // Normalize the iterator boundary even for transaction-control and + // replication paths that intentionally carry no execution lease. The + // wrapper makes Close idempotent when a result helper already closed the + // child, and guarantees any retained snapshot lease is released before the + // session schema is inspected below. + rowIter = backend.WrapRowIterWithRelease(rowIter, nil) + defer func() { + if rowIter != nil { + closeErr := rowIter.Close(sqlCtx) + err = errors.Join(err, closeErr) + returnErr = errors.Join(returnErr, closeErr) + } + // CurrentSchemaOfUnderlyingConn re-enters the session/pool lifecycle + // lock. It must run only after the iterator wrapper has closed its child + // and released the retained execution snapshot. + if isPostgresTransactionControl(parsed) { + return + } + if session, ok := sqlCtx.Session.(*backend.Session); ok && session != nil { + if currentSchema := session.CurrentSchemaOfUnderlyingConn(); len(currentSchema) > 0 { + sqlCtx.SetCurrentDatabase(currentSchema) + } + } + }() rowIter = backend.ApplyQueryRowLimit(sqlCtx, schema, rowIter) - // If the query is "USE ", we need to update the current database in the session. - if currentSchema := sqlCtx.Session.(*backend.Session).CurrentSchemaOfUnderlyingConn(); len(currentSchema) > 0 { - sqlCtx.SetCurrentDatabase(currentSchema) - } - // create result before goroutines to avoid |ctx| racing var r *Result var processedAtLeastOneBatch bool @@ -486,7 +977,7 @@ type QueryExecutor func(ctx *sql.Context, query string, parsed tree.Statement, s // executeQuery is a QueryExecutor that calls QueryWithBindings on the given engine using the given query and parsed // statement, which may be nil. -func (h *DuckHandler) executeQuery(ctx *sql.Context, query string, parsed tree.Statement, _ *duckdb.Stmt, _ []any) (sql.Schema, sql.RowIter, *sql.QueryFlags, error) { +func (h *DuckHandler) executeQuery(ctx *sql.Context, query string, parsed tree.Statement, _ *duckdb.Stmt, _ []any) (schema sql.Schema, iter sql.RowIter, qFlags *sql.QueryFlags, returnErr error) { // return h.e.QueryWithBindings(ctx, query, parsed, nil, nil) sql.IncrementStatusVariable(ctx, "Questions", 1) @@ -495,20 +986,110 @@ func (h *DuckHandler) executeQuery(ctx *sql.Context, query string, parsed tree.S } var ( - schema sql.Schema - iter sql.RowIter rows *stdsql.Rows result stdsql.Result err error ) + // Transaction controls are protocol lifecycle operations. Resolve them + // before routing or acquiring an execution snapshot: a failed-block + // COMMIT/ROLLBACK must remain able to finalize its owner even when ordinary + // statement routing or connection acquisition is no longer usable. + switch transaction := parsed.(type) { + case *tree.BeginTransaction, *tree.CommitTransaction, *tree.RollbackTransaction: + handled, controlErr := postgresTransactionControl(ctx, transaction, h.GetCatalogProvider()) + if !handled { + return nil, nil, nil, fmt.Errorf("unsupported PostgreSQL transaction statement %T", parsed) + } + if controlErr != nil { + return nil, nil, nil, controlErr + } + return types.OkResultSchema, sql.RowsToRowIter(sql.NewRow(types.OkResult{})), nil, nil + } + if err := rejectPostgresDatabaseDDLInTransaction(ctx, parsed); err != nil { + return nil, nil, nil, err + } + execer, ownerConn, _, leaseRelease, snapshotErr := h.getPostgresExecutionSnapshotLease(ctx, parsed) + if snapshotErr != nil { + return nil, nil, nil, snapshotErr + } + // Transfer the retained snapshot/operation lease to the returned iterator. + // Wrapping even OK-result iterators gives doQuery one idempotent close point; + // on every error or nil-iterator return the wrapper releases immediately. + defer func() { + if leaseRelease != nil { + iter = backend.WrapRowIterWithRelease(iter, leaseRelease) + leaseRelease = nil + } + }() + routedQuery, routeErr := h.rewritePostgresObjectRelationsWithSnapshot(ctx, query, execer, ownerConn) + if routeErr != nil { + return nil, nil, nil, routeErr + } // NOTE: The query is parsed using Postgres parser, which does not support all DuckDB syntax. // Consequently, the following classification is not perfect. switch parsed := parsed.(type) { - case *tree.BeginTransaction, *tree.CommitTransaction, *tree.RollbackTransaction, - *tree.CreateTable, *tree.DropTable, *tree.AlterTable, *tree.CreateIndex, *tree.DropIndex, + case *tree.CreateTable: + selection, createErr := setPostgresCreateTableStorage(ctx, parsed) + if createErr != nil { + err = createErr + break + } + if parsed.Persistence == tree.PersistenceTemporary && selection.IsObjectStorage() { + err = fmt.Errorf("%w: temporary tables cannot use object storage", catalog.ErrInvalidTableStorage) + break + } + executionQuery, _, createErr := postgresCreateTableExecutionQuery(query, parsed) + if createErr != nil { + err = createErr + break + } + result, err = execer.ExecContext(ctx, executionQuery) + if err == nil { + err = persistPostgresCreateTableStorageWithExecutor(ctx, parsed, selection, h.GetCatalogProvider(), execer, ownerConn) + } + if err != nil { + break + } + affected, _ := result.RowsAffected() + insertId, _ := result.LastInsertId() + schema = types.OkResultSchema + iter = sql.RowsToRowIter(sql.NewRow(types.OkResult{ + RowsAffected: uint64(affected), + InsertID: uint64(insertId), + })) + case *tree.DropTable, *tree.AlterTable: + if ddlResult, handled, ddlErr := h.executePostgresObjectDDL(ctx, query, routedQuery, parsed, execer, ownerConn); handled { + result, err = ddlResult, ddlErr + if err != nil { + break + } + affected, _ := result.RowsAffected() + insertId, _ := result.LastInsertId() + schema = types.OkResultSchema + iter = sql.RowsToRowIter(sql.NewRow(types.OkResult{ + RowsAffected: uint64(affected), + InsertID: uint64(insertId), + })) + break + } else if ddlErr != nil { + err = ddlErr + break + } + result, err = execer.ExecContext(ctx, routedQuery) + if err != nil { + break + } + affected, _ := result.RowsAffected() + insertId, _ := result.LastInsertId() + schema = types.OkResultSchema + iter = sql.RowsToRowIter(sql.NewRow(types.OkResult{ + RowsAffected: uint64(affected), + InsertID: uint64(insertId), + })) + case *tree.CreateIndex, *tree.DropIndex, *tree.Update, *tree.Delete, *tree.Truncate, *tree.CopyFrom, *tree.CopyTo, *tree.SetVar: - result, err = adapter.Exec(ctx, query) + result, err = execer.ExecContext(ctx, routedQuery) if err != nil { break } @@ -521,7 +1102,7 @@ func (h *DuckHandler) executeQuery(ctx *sql.Context, query string, parsed tree.S })) case *tree.Insert: if _, ok := postgresInsertReturningRows(parsed); ok { - rows, err = adapter.QueryCatalog(ctx, query) + rows, err = execer.QueryContext(ctx, routedQuery) if err != nil { break } @@ -536,7 +1117,7 @@ func (h *DuckHandler) executeQuery(ctx *sql.Context, query string, parsed tree.S } break } - result, err = adapter.Exec(ctx, query) + result, err = execer.ExecContext(ctx, routedQuery) if err != nil { break } @@ -554,7 +1135,7 @@ func (h *DuckHandler) executeQuery(ctx *sql.Context, query string, parsed tree.S break } dbName := parsed.Name.String() - err = provider.CreateCatalog(dbName, parsed.IfNotExists) + err = provider.CreateCatalogOnConn(ctx, ownerConn, dbName, parsed.IfNotExists) if err != nil { break } @@ -567,14 +1148,14 @@ func (h *DuckHandler) executeQuery(ctx *sql.Context, query string, parsed tree.S break } dbName := parsed.Name.String() - err = provider.DropCatalog(dbName, parsed.IfExists) + err = provider.DropCatalogOnConn(ctx, ownerConn, dbName, parsed.IfExists) if err != nil { break } schema = types.OkResultSchema iter = sql.RowsToRowIter(sql.NewRow(types.OkResult{})) case *tree.Select, tree.SelectStatement: - rows, schema, err = queryCatalogWithJSONScan(ctx, query) + rows, schema, err = queryCatalogWithJSONScan(ctx, routedQuery, execer) if err != nil { break } @@ -584,7 +1165,7 @@ func (h *DuckHandler) executeQuery(ctx *sql.Context, query string, parsed tree.S break } default: - rows, err = adapter.QueryCatalog(ctx, query) + rows, err = execer.QueryContext(ctx, routedQuery) if err != nil { break } @@ -606,15 +1187,22 @@ func (h *DuckHandler) executeQuery(ctx *sql.Context, query string, parsed tree.S return schema, iter, nil, nil } -func queryCatalogWithJSONScan(ctx *sql.Context, query string, vars ...any) (*stdsql.Rows, sql.Schema, error) { +// queryCatalogWithJSONScan performs both the schema probe and the final query +// through the caller's captured executor. The caller must retain the matching +// execution lease until the returned rows have been consumed and closed; this +// helper deliberately never reacquires a session snapshot. +func queryCatalogWithJSONScan(ctx *sql.Context, query string, execer adapter.SQLExecutor, vars ...any) (*stdsql.Rows, sql.Schema, error) { + if execer == nil { + return nil, nil, fmt.Errorf("query executor is unavailable") + } probeQuery := "SELECT * FROM (" + sql.RemoveSpaceAndDelimiter(query, ';') + ") LIMIT 0" - probeRows, err := adapter.QueryCatalog(ctx, probeQuery, vars...) + probeRows, err := execer.QueryContext(ctx, probeQuery, vars...) if err != nil { // The PostgreSQL parser can classify DuckDB-only statements such as // CREATE OR REPLACE TABLE as a SelectStatement. The probe wrapper is // invalid for those statements, so preserve the pre-projection direct // execution path instead of returning the wrapper's parser error. - rows, directErr := adapter.QueryCatalog(ctx, query, vars...) + rows, directErr := execer.QueryContext(ctx, query, vars...) if directErr != nil { return nil, nil, directErr } @@ -634,7 +1222,7 @@ func queryCatalogWithJSONScan(ctx *sql.Context, query string, vars ...any) (*std return nil, nil, closeErr } - rows, err := adapter.QueryCatalog(ctx, backend.QueryForJSONScan(query, schema), vars...) + rows, err := execer.QueryContext(ctx, backend.QueryForJSONScan(query, schema), vars...) if err != nil { return nil, nil, err } @@ -643,7 +1231,7 @@ func queryCatalogWithJSONScan(ctx *sql.Context, query string, vars ...any) (*std // executeBoundPlan is a QueryExecutor that calls QueryWithBindings on the given engine using the given query and parsed // statement, which may be nil. -func (h *DuckHandler) executeBoundPlan(ctx *sql.Context, query string, parsed tree.Statement, stmt *duckdb.Stmt, vars []any) (sql.Schema, sql.RowIter, *sql.QueryFlags, error) { +func (h *DuckHandler) executeBoundPlan(ctx *sql.Context, query string, parsed tree.Statement, stmt *duckdb.Stmt, vars []any) (schema sql.Schema, iter sql.RowIter, qFlags *sql.QueryFlags, returnErr error) { // return h.e.PrepQueryPlanForExecution(ctx, query, plan, nil) // TODO(fan): Currently, the result of executing the bound query is occasionally incorrect. @@ -697,19 +1285,72 @@ func (h *DuckHandler) executeBoundPlan(ctx *sql.Context, query string, parsed tr // return nil, nil, nil, err // } - stmtType, err := stmt.StatementType() - if err != nil { - return nil, nil, nil, err + // As in simple-query execution, transaction controls must not depend on an + // executor snapshot. In particular, failed-state COMMIT is the recovery + // operation that clears a potentially unusable transaction binding. + switch transaction := parsed.(type) { + case *tree.BeginTransaction, *tree.CommitTransaction, *tree.RollbackTransaction: + handled, controlErr := postgresTransactionControl(ctx, transaction, h.GetCatalogProvider()) + if !handled { + return nil, nil, nil, fmt.Errorf("unsupported PostgreSQL transaction statement %T", parsed) + } + if controlErr != nil { + return nil, nil, nil, controlErr + } + return types.OkResultSchema, sql.RowsToRowIter(sql.NewRow(types.OkResult{})), nil, nil } var ( - schema sql.Schema - iter sql.RowIter rows *stdsql.Rows result stdsql.Result + err error ) + execer, ownerConn, tx, leaseRelease, snapshotErr := h.getPostgresExecutionSnapshotLease(ctx, parsed) + if snapshotErr != nil { + return nil, nil, nil, snapshotErr + } + // Keep the selected connection/transaction and provider admission alive + // through iterator consumption. The defer handles every setup error and + // transfers ownership to the returned iterator on success. + defer func() { + if leaseRelease != nil { + iter = backend.WrapRowIterWithRelease(iter, leaseRelease) + leaseRelease = nil + } + }() + routedQuery, routeErr := h.rewritePostgresObjectRelationsWithSnapshot(ctx, query, execer, ownerConn) + if routeErr != nil { + return nil, nil, nil, routeErr + } + var stmtType duckdb.StmtType + if tx != nil { + // A parsed PostgreSQL portal is executed through the active *sql.Tx. + // Never inspect or prepare a raw driver statement on ownerConn here; + // doing so bypasses the transaction executor. + stmtType = postgresDuckDBStatementType(parsed, routedQuery) + } else if stmt != nil { + stmtType, err = stmt.StatementType() + if err != nil { + return nil, nil, nil, err + } + } else { + // Parse/Bind deliberately retain no raw DuckDB statement. In the + // no-transaction case only, use a short-lived raw probe for exact + // DuckDB classification. The retained execution snapshot keeps the + // physical owner stable while the probe runs. + prepareQuery := routedQuery + if create, ok := parsed.(*tree.CreateTable); ok { + if normalized, normalizeErr := postgresCreateTableQueryForPrepare(query, create); normalizeErr == nil { + prepareQuery = normalized + } + } + stmtType, _, err = prepareDuckDBStatementMetadata(ctx, ownerConn, prepareQuery) + if err != nil { + return nil, nil, nil, err + } + } if _, ok := postgresInsertReturningRows(parsed); ok { - rows, err = adapter.QueryCatalog(ctx, query, vars...) + rows, err = execer.QueryContext(ctx, routedQuery, vars...) if err == nil { schema, err = pgtypes.InferSchema(rows) } @@ -722,10 +1363,69 @@ func (h *DuckHandler) executeBoundPlan(ctx *sql.Context, query string, parsed tr return schema, iter, nil, err } + // Prepared PostgreSQL DDL is executed through this bound-plan path rather + // than executeQuery. Keep CREATE TABLE's selector and metadata behavior + // identical to the simple-query path, and remove MyDuck-only WITH options + // before handing the statement to DuckDB. + if create, ok := parsed.(*tree.CreateTable); ok { + selection, createErr := setPostgresCreateTableStorage(ctx, create) + if createErr != nil { + return nil, nil, nil, createErr + } + if create.Persistence == tree.PersistenceTemporary && selection.IsObjectStorage() { + return nil, nil, nil, fmt.Errorf("%w: temporary tables cannot use object storage", catalog.ErrInvalidTableStorage) + } + executionQuery, _, createErr := postgresCreateTableExecutionQuery(query, create) + if createErr != nil { + return nil, nil, nil, createErr + } + result, createErr := execer.ExecContext(ctx, executionQuery, vars...) + if createErr != nil { + return nil, nil, nil, createErr + } + if createErr = persistPostgresCreateTableStorageWithExecutor(ctx, create, selection, h.GetCatalogProvider(), execer, ownerConn); createErr != nil { + return nil, nil, nil, createErr + } + affected, _ := result.RowsAffected() + insertID, _ := result.LastInsertId() + return types.OkResultSchema, sql.RowsToRowIter(sql.NewRow(types.OkResult{ + RowsAffected: uint64(affected), + InsertID: uint64(insertID), + })), nil, nil + } + if _, ddl := parsed.(*tree.DropTable); ddl { + result, handled, ddlErr := h.executePostgresObjectDDL(ctx, query, routedQuery, parsed, execer, ownerConn, vars...) + if ddlErr != nil { + return nil, nil, nil, ddlErr + } + if handled { + affected, _ := result.RowsAffected() + insertID, _ := result.LastInsertId() + return types.OkResultSchema, sql.RowsToRowIter(sql.NewRow(types.OkResult{ + RowsAffected: uint64(affected), + InsertID: uint64(insertID), + })), nil, nil + } + } + if _, ddl := parsed.(*tree.AlterTable); ddl { + result, handled, ddlErr := h.executePostgresObjectDDL(ctx, query, routedQuery, parsed, execer, ownerConn, vars...) + if ddlErr != nil { + return nil, nil, nil, ddlErr + } + if handled { + affected, _ := result.RowsAffected() + insertID, _ := result.LastInsertId() + return types.OkResultSchema, sql.RowsToRowIter(sql.NewRow(types.OkResult{ + RowsAffected: uint64(affected), + InsertID: uint64(insertID), + })), nil, nil + } + } + switch stmtType { case duckdb.STATEMENT_TYPE_SELECT, duckdb.STATEMENT_TYPE_RELATION: - rows, schema, err = queryCatalogWithJSONScan(ctx, query, vars...) + rows, schema, err = queryCatalogWithJSONScan(ctx, routedQuery, execer, vars...) if err != nil { break } @@ -737,7 +1437,7 @@ func (h *DuckHandler) executeBoundPlan(ctx *sql.Context, query string, parsed tr case duckdb.STATEMENT_TYPE_CALL, duckdb.STATEMENT_TYPE_PRAGMA, duckdb.STATEMENT_TYPE_EXPLAIN: - rows, err = adapter.QueryCatalog(ctx, query, vars...) + rows, err = execer.QueryContext(ctx, routedQuery, vars...) if err != nil { break } @@ -752,7 +1452,7 @@ func (h *DuckHandler) executeBoundPlan(ctx *sql.Context, query string, parsed tr break } default: - result, err = adapter.ExecCatalog(ctx, query, vars...) + result, err = execer.ExecContext(ctx, routedQuery, vars...) if err != nil { break } @@ -767,7 +1467,6 @@ func (h *DuckHandler) executeBoundPlan(ctx *sql.Context, query string, parsed tr if err != nil { return nil, nil, nil, err } - return schema, iter, nil, nil } @@ -780,11 +1479,26 @@ func postgresInsertReturningRows(stmt tree.Statement) (*tree.Insert, bool) { return insert, ok } -func inferPostgresInsertReturningSchema(ctx *sql.Context, insert *tree.Insert) (sql.Schema, error) { - returning := insert.Returning.(*tree.ReturningExprs) - query := "SELECT " + tree.AsString((*tree.SelectExprs)(returning)) + +func postgresInsertReturningSchemaQuery(insert *tree.Insert) string { + if insert == nil { + return "" + } + returning, ok := insert.Returning.(*tree.ReturningExprs) + if !ok { + return "" + } + return "SELECT " + tree.AsString((*tree.SelectExprs)(returning)) + " FROM " + tree.AsString(insert.Table) + " LIMIT 0" - rows, err := adapter.QueryCatalog(ctx, query) +} + +func inferPostgresInsertReturningSchema(ctx *sql.Context, query string, execer adapter.SQLExecutor) (sql.Schema, error) { + if execer == nil { + return nil, fmt.Errorf("INSERT ... RETURNING metadata executor is unavailable") + } + if strings.TrimSpace(query) == "" { + return nil, fmt.Errorf("INSERT ... RETURNING metadata query is empty") + } + rows, err := execer.QueryContext(ctx, query) if err != nil { return nil, err } diff --git a/pgserver/ducklake_operation_owner_test.go b/pgserver/ducklake_operation_owner_test.go new file mode 100644 index 00000000..19e1c66a --- /dev/null +++ b/pgserver/ducklake_operation_owner_test.go @@ -0,0 +1,77 @@ +package pgserver + +import ( + "context" + stdsql "database/sql" + "testing" + + "github.com/apecloud/myduckserver/mycontext" + "github.com/cockroachdb/cockroachdb-parser/pkg/sql/sem/tree" + "github.com/dolthub/go-mysql-server/memory" + "github.com/dolthub/go-mysql-server/sql" + "github.com/stretchr/testify/require" +) + +type postgresOperationOwnerProviderProbe struct { + legacyCalls int + ownerCalls int + owner any +} + +func (p *postgresOperationOwnerProviderProbe) BeginDuckLakeOperation(context.Context) func() { + p.legacyCalls++ + return func() {} +} + +func (p *postgresOperationOwnerProviderProbe) BeginDuckLakeOperationForOwner(_ context.Context, owner any) func() { + p.ownerCalls++ + p.owner = owner + return func() {} +} + +func TestBeginPostgresDuckLakeOperationPassesPhysicalOwner(t *testing.T) { + physicalTx := new(stdsql.Tx) + session := &rollbackCleanupSessionProbe{ + Session: memory.NewSession(sql.NewBaseSession(), nil), + tx: physicalTx, + } + ctx := sql.NewContext( + mycontext.WithFrontendQuery(context.Background()), + sql.WithSession(session), + ) + provider := &postgresOperationOwnerProviderProbe{} + + owner := postgresDuckLakeOperationOwner(ctx, nil) + require.Same(t, physicalTx, owner) + release := beginPostgresDuckLakeOperationForProvider(ctx, &tree.Select{}, provider, owner) + release() + + require.Equal(t, 1, provider.ownerCalls) + require.Equal(t, 0, provider.legacyCalls) + require.Same(t, physicalTx, provider.owner) +} + +func TestBeginPostgresDuckLakeOperationFallsBackForLegacyProvider(t *testing.T) { + physicalTx := new(stdsql.Tx) + session := &rollbackCleanupSessionProbe{ + Session: memory.NewSession(sql.NewBaseSession(), nil), + tx: physicalTx, + } + ctx := sql.NewContext( + mycontext.WithFrontendQuery(context.Background()), + sql.WithSession(session), + ) + provider := &legacyPostgresOperationProviderProbe{} + + release := beginPostgresDuckLakeOperationForProvider(ctx, &tree.Select{}, provider, physicalTx) + release() + + require.Equal(t, 1, provider.calls) +} + +type legacyPostgresOperationProviderProbe struct{ calls int } + +func (p *legacyPostgresOperationProviderProbe) BeginDuckLakeOperation(context.Context) func() { + p.calls++ + return func() {} +} diff --git a/pgserver/implicit_transaction_test.go b/pgserver/implicit_transaction_test.go new file mode 100644 index 00000000..77a0dd0c --- /dev/null +++ b/pgserver/implicit_transaction_test.go @@ -0,0 +1,115 @@ +package pgserver + +import ( + "context" + "os" + "strconv" + "testing" + + "github.com/apecloud/myduckserver/testutil" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" + "github.com/stretchr/testify/require" +) + +// A simple Query message containing multiple statements is one PostgreSQL +// transaction scope. The first write must not survive a later statement error. +func TestPostgresSimpleQueryImplicitTransactionRollsBackBatch(t *testing.T) { + ctx, conn := newImplicitTransactionTestConnection(t, pgx.QueryExecModeSimpleProtocol) + + _, err := conn.Exec(ctx, "INSERT INTO implicit_batch.implicit_batch_rows VALUES (2); INSERT INTO implicit_batch.implicit_batch_rows (missing) VALUES (3)") + require.Error(t, err) + + var count int + require.NoError(t, conn.QueryRow(ctx, "SELECT count(*) FROM implicit_batch.implicit_batch_rows").Scan(&count)) + require.Equal(t, 1, count) +} + +// pgx.Batch uses the extended protocol and emits one Sync after all queued +// statements. The server must retain the same implicit transaction across the +// whole exchange so an error in a later statement rolls back earlier writes. +func TestPostgresExtendedBatchImplicitTransactionRollsBack(t *testing.T) { + ctx, conn := newImplicitTransactionTestConnection(t, pgx.QueryExecModeExec) + + var batch pgx.Batch + batch.Queue("INSERT INTO implicit_batch.implicit_batch_rows VALUES (2)") + batch.Queue("INSERT INTO implicit_batch.implicit_batch_rows (missing) VALUES (3)") + results := conn.SendBatch(ctx, &batch) + _, err := results.Exec() + require.NoError(t, err) + _, err = results.Exec() + require.Error(t, err) + // The second Exec has consumed the server error and closed the result + // reader. Close is still required by pgx before reusing the connection. + _ = results.Close() + + var count int + require.NoError(t, conn.QueryRow(ctx, "SELECT count(*) FROM implicit_batch.implicit_batch_rows").Scan(&count)) + require.Equal(t, 1, count) + + // A subsequent clean extended exchange must start a fresh implicit scope + // and commit it when its Sync arrives. + var clean pgx.Batch + clean.Queue("INSERT INTO implicit_batch.implicit_batch_rows VALUES (4)") + cleanResults := conn.SendBatch(ctx, &clean) + _, err = cleanResults.Exec() + require.NoError(t, err) + require.NoError(t, cleanResults.Close()) + require.NoError(t, conn.QueryRow(ctx, "SELECT count(*) FROM implicit_batch.implicit_batch_rows").Scan(&count)) + require.Equal(t, 2, count) +} + +// An explicit BEGIN must promote the connection to PostgreSQL's sticky +// transaction-block state. A later statement error is not an implicit-scope +// error: the block remains failed until the client issues ROLLBACK. +func TestPostgresExplicitTransactionRemainsFailedUntilRollback(t *testing.T) { + ctx, conn := newImplicitTransactionTestConnection(t, pgx.QueryExecModeSimpleProtocol) + + _, err := conn.Exec(ctx, "BEGIN; INSERT INTO implicit_batch.implicit_batch_rows VALUES (2); INSERT INTO implicit_batch.implicit_batch_rows (missing) VALUES (3)") + require.Error(t, err) + + _, err = conn.Exec(ctx, "SELECT count(*) FROM implicit_batch.implicit_batch_rows") + require.Error(t, err) + var pgErr *pgconn.PgError + require.ErrorAs(t, err, &pgErr) + require.Equal(t, "25P02", pgErr.Code) + + _, err = conn.Exec(ctx, "ROLLBACK") + require.NoError(t, err) + + var count int + require.NoError(t, conn.QueryRow(ctx, "SELECT count(*) FROM implicit_batch.implicit_batch_rows").Scan(&count)) + require.Equal(t, 1, count) +} + +func newImplicitTransactionTestConnection(t *testing.T, mode pgx.QueryExecMode) (context.Context, *pgx.Conn) { + t.Helper() + + originalWorkingDir, err := os.Getwd() + require.NoError(t, err) + defer func() { require.NoError(t, os.Chdir(originalWorkingDir)) }() + + testDir := testutil.CreateTestDir(t) + testEnv := testutil.NewTestEnv() + require.NoError(t, testutil.StartDuckSqlServer(t, testDir, nil, testEnv)) + t.Cleanup(func() { + require.NoError(t, testEnv.MyDuckServer.Close()) + testutil.StopDuckSqlServer(t, testEnv.DuckProcess) + }) + + _, err = testEnv.MyDuckServer.Exec("CREATE DATABASE implicit_batch") + require.NoError(t, err) + _, err = testEnv.MyDuckServer.Exec("CREATE TABLE implicit_batch.implicit_batch_rows (id INTEGER)") + require.NoError(t, err) + _, err = testEnv.MyDuckServer.Exec("INSERT INTO implicit_batch.implicit_batch_rows VALUES (1)") + require.NoError(t, err) + + ctx := context.Background() + config, err := pgx.ParseConfig("postgresql://postgres@127.0.0.1:" + strconv.Itoa(testEnv.DuckPgPort) + "/postgres") + require.NoError(t, err) + config.DefaultQueryExecMode = mode + conn, err := pgx.ConnectConfig(ctx, config) + require.NoError(t, err) + t.Cleanup(func() { _ = conn.Close(ctx) }) + return ctx, conn +} diff --git a/pgserver/in_place_handler.go b/pgserver/in_place_handler.go index 5d3389ae..a29330b8 100644 --- a/pgserver/in_place_handler.go +++ b/pgserver/in_place_handler.go @@ -14,6 +14,7 @@ import ( "github.com/apecloud/myduckserver/pgserver/pgconfig" "github.com/cockroachdb/cockroachdb-parser/pkg/sql/sem/tree" "github.com/dolthub/go-mysql-server/sql" + "github.com/jackc/pgx/v5/pgconn" "github.com/jackc/pgx/v5/pgproto3" ) @@ -107,13 +108,14 @@ func (h *ConnectionHandler) setPgSessionVar(name string, value any, useDefault b if err != nil { return false, fmt.Errorf("error: %s variable was not found, err: %w", name, err) } - // Sent CommandComplete message - err = h.send(makeCommandComplete(tag, 0)) + // Keep both responses behind a final implicit simple-query commit. Sending + // either one directly could expose a successful SET/RESET before a prior + // statement's transaction fails during the protocol boundary. + err = h.sendOrQueueProtocolMessage(makeCommandComplete(tag, 0)) if err != nil { return true, err } - // Sent ParameterStatus message - if err := h.send(&pgproto3.ParameterStatus{ + if err := h.sendOrQueueProtocolMessage(&pgproto3.ParameterStatus{ Name: name, Value: fmt.Sprintf("%v", v), }); err != nil { @@ -329,13 +331,14 @@ var inPlaceHandlers = map[string]InPlaceHandler{ Tag: "SELECT", }) } - // TODO(sean): Implement SHOW ALL - _ = h.send(&pgproto3.ErrorResponse{ - Severity: string(ErrorResponseSeverity_Error), + // TODO(sean): Implement SHOW ALL. Return the error to the protocol + // layer so an active transaction enters PostgreSQL's failed state; + // sending it here would make the handler report a false success. + return true, &pgconn.PgError{ + Severity: "ERROR", Code: "0A000", Message: "Statement 'SHOW ALL' is not supported yet.", - }) - return true, nil + } }, }, "SET": { @@ -434,13 +437,13 @@ var inPlaceHandlers = map[string]InPlaceHandler{ if !resetVar.ResetAll { return h.setPgSessionVar(key, nil, true, "RESET") } - // TODO(sean): Implement RESET ALL - _ = h.send(&pgproto3.ErrorResponse{ - Severity: string(ErrorResponseSeverity_Error), + // TODO(sean): Implement RESET ALL. Return the error to the protocol + // layer so the transaction state transition is not lost. + return true, &pgconn.PgError{ + Severity: "ERROR", Code: "0A000", Message: "Statement 'RESET ALL' is not supported yet.", - }) - return true, nil + } }, }, } diff --git a/pgserver/logrepl/replication.go b/pgserver/logrepl/replication.go index 893f8b86..ef514ea1 100644 --- a/pgserver/logrepl/replication.go +++ b/pgserver/logrepl/replication.go @@ -28,6 +28,7 @@ import ( "github.com/apecloud/myduckserver/binlog" "github.com/apecloud/myduckserver/catalog" "github.com/apecloud/myduckserver/delta" + "github.com/apecloud/myduckserver/mycontext" "github.com/apecloud/myduckserver/pgtypes" "github.com/dolthub/go-mysql-server/sql" "github.com/jackc/pglogrepl" @@ -220,6 +221,15 @@ func (state *replicationState) reset(ctx *sql.Context, slotName string, lsn pglo // StartReplication starts the replication process for the given slot name. This function blocks until replication is // stopped via the Stop method, or an error occurs. func (r *LogicalReplicator) StartReplication(sqlCtx *sql.Context, slotName string) error { + if sqlCtx == nil { + return fmt.Errorf("replication context is nil") + } + // Replication is a long-lived service operation. Keep it independent from + // the request that created or enabled the subscription, and explicitly mark + // its origin so rollback/finalization paths can never be mistaken for a + // frontend transaction that is eligible for DuckLake orphan cleanup. + replicationBase := mycontext.WithQueryOrigin(context.WithoutCancel(sqlCtx), mycontext.PostgresReplicationQueryOrigin) + sqlCtx = sqlCtx.WithContext(replicationBase) sqlCtx.SetLogger(r.logger) standbyMessageTimeout := 10 * time.Second nextStandbyMessageDeadline := time.Now().Add(standbyMessageTimeout) @@ -450,13 +460,12 @@ func (r *LogicalReplicator) StartReplication(sqlCtx *sql.Context, slotName strin } func (r *LogicalReplicator) rollback(ctx *sql.Context) error { - defer adapter.CloseTxn(ctx) - txn := adapter.TryGetTxn(ctx) + _, txn := adapter.TryGetTxnBindingForRelease(ctx) if txn == nil { return nil } - err := txn.Rollback() - if err != nil && !strings.Contains(err.Error(), "no transaction is active") { + _, _, err := adapter.FinalizeRollback(ctx, txn) + if err != nil && !adapter.IsTransactionInactiveError(err) { r.logger.Debugf("Failed to roll back transaction: %v", err) return err } @@ -700,14 +709,30 @@ func (r *LogicalReplicator) processMessage( switch logicalMsg := logicalMsg.(type) { case *pglogrepl.RelationMessageV2: - _, exists := state.relations[logicalMsg.RelationID] - if exists { + previous, exists := state.relations[logicalMsg.RelationID] + // PostgreSQL repeats a relation message at the start of every + // transaction, even when the schema did not change. Only flush an + // existing batch when the relation metadata actually changed. + // Committing an empty transaction here would close the transaction + // opened by the preceding BeginMessage; the following row messages + // would then be buffered without a physical transaction and could + // never advance the replication LSN. + if exists && relationMessageChanged(previous, logicalMsg) && state.ongoingBatchTxn && state.dirtyTxn { // This means schema changes have occurred, so we need to // commit any buffered ongoing batch transactions. err := r.commitOngoingTxn(state, delta.DDLStmtFlushReason) if err != nil { return false, err } + // The Relation message is emitted inside the source transaction. + // Re-open the local transaction after flushing a previous batch so + // row messages that follow this relation remain transactional. + if state.processMessages { + if _, err := adapter.GetCatalogTxn(state.replicaCtx, nil); err != nil { + return false, err + } + state.ongoingBatchTxn = true + } } state.relations[logicalMsg.RelationID] = logicalMsg @@ -866,12 +891,22 @@ func (r *LogicalReplicator) processMessage( r.logger.Debugf("Truncate message: xid %d\n", logicalMsg.Xid) - // Flush the delta buffer first - r.flushDeltaBuffer(state, nil, nil, delta.DMLStmtFlushReason) + // Flush the delta buffer first. A truncate can arrive after row changes + // have been buffered in the current transaction, so use one lifecycle + // snapshot for the physical connection/transaction and surface any flush + // failure instead of continuing with a partially applied stream. + conn, tx, release, err := adapter.GetCatalogTxnExecutionSnapshotWithLease(state.replicaCtx, nil) + if err != nil { + return false, err + } + defer release() + if err := r.flushDeltaBuffer(state, conn, tx, delta.DMLStmtFlushReason); err != nil { + return false, err + } // Truncate the tables for _, relationID := range logicalMsg.RelationIDs { - if err := r.truncate(state, relationID); err != nil { + if err := r.truncate(state, tx, relationID); err != nil { return false, err } } @@ -903,6 +938,40 @@ func (r *LogicalReplicator) processMessage( return false, nil } +// relationMessageChanged reports whether a Relation message describes a +// different relation schema. PostgreSQL sends an otherwise identical Relation +// message at the start of each transaction, so treating every repeated message +// as a DDL boundary would prematurely commit the local transaction. +func relationMessageChanged(previous, current *pglogrepl.RelationMessageV2) bool { + if previous == nil || current == nil { + return true + } + if previous.RelationID != current.RelationID || + previous.Namespace != current.Namespace || + previous.RelationName != current.RelationName || + previous.ReplicaIdentity != current.ReplicaIdentity || + previous.ColumnNum != current.ColumnNum || + len(previous.Columns) != len(current.Columns) { + return true + } + for i, oldColumn := range previous.Columns { + newColumn := current.Columns[i] + if oldColumn == nil || newColumn == nil { + if oldColumn != newColumn { + return true + } + continue + } + if oldColumn.Flags != newColumn.Flags || + oldColumn.Name != newColumn.Name || + oldColumn.DataType != newColumn.DataType || + oldColumn.TypeModifier != newColumn.TypeModifier { + return true + } + } + return false +} + // whereClause returns a WHERE clause string with the contents of the builder if it's non-empty, or the empty // string otherwise func whereClause(str strings.Builder) string { @@ -981,19 +1050,36 @@ func (r *LogicalReplicator) commitOngoingTxnIfClean(state *replicationState, rea return nil } -// commitOngoingTxn commits the current transaction -func (r *LogicalReplicator) commitOngoingTxn(state *replicationState, flushReason delta.FlushReason) error { - conn, err := adapter.GetCatalogConn(state.replicaCtx) +// commitOngoingTxn commits the current transaction. +func (r *LogicalReplicator) commitOngoingTxn(state *replicationState, flushReason delta.FlushReason) (err error) { + // Acquire the physical connection and transaction from one lifecycle + // snapshot. A separate GetCatalogConn followed by TryGetTxn can pair a + // transaction with a replacement connection while a session is being + // finalized or recreated. + _, conn, tx, release, err := adapter.GetCatalogExecutionSnapshotWithLease(state.replicaCtx) if err != nil { return err } - tx := adapter.TryGetTxn(state.replicaCtx) + defer release() if tx == nil { return nil } - defer tx.Rollback() - defer adapter.CloseTxn(state.replicaCtx) + committed := false + defer func() { + if committed { + return + } + _, _, rollbackErr := adapter.FinalizeRollback(state.replicaCtx, tx) + if rollbackErr != nil && !adapter.IsTransactionInactiveError(rollbackErr) { + rollbackErr = fmt.Errorf("failed to roll back replication transaction: %w", rollbackErr) + if err == nil { + err = rollbackErr + } else { + err = errors.Join(err, rollbackErr) + } + } + }() // Flush the delta buffer if too large err = r.flushDeltaBuffer(state, conn, tx, flushReason) @@ -1002,14 +1088,15 @@ func (r *LogicalReplicator) commitOngoingTxn(state *replicationState, flushReaso } r.logger.Debugf("Writing LSN %s\n", state.lastCommitLSN) - if err = UpdateSubscriptionLsn(state.replicaCtx, state.lastCommitLSN.String(), r.subscription); err != nil { + if err = UpdateSubscriptionLsnInTxn(state.replicaCtx, tx, state.lastCommitLSN.String(), r.subscription); err != nil { return err } // Commit the transaction - if err := tx.Commit(); err != nil { + if err = adapter.FinalizeCommit(state.replicaCtx, tx); err != nil { return err } + committed = true // Reset transaction state state.ongoingBatchTxn = false @@ -1105,14 +1192,14 @@ func (r *LogicalReplicator) append(state *replicationState, relationID uint32, t return nil } -func (r *LogicalReplicator) truncate(state *replicationState, relationID uint32) error { +func (r *LogicalReplicator) truncate(state *replicationState, tx *stdsql.Tx, relationID uint32) error { rel, ok := state.relations[relationID] if !ok { return fmt.Errorf("unknown relation ID %d", relationID) } r.logger.Debugf("Truncating table %s.%s\n", rel.Namespace, rel.RelationName) - _, err := adapter.ExecInTxn(state.replicaCtx, `TRUNCATE `+catalog.ConnectIdentifiersANSI(rel.Namespace, rel.RelationName)) + _, err := tx.ExecContext(state.replicaCtx, `TRUNCATE `+catalog.ConnectIdentifiersANSI(rel.Namespace, rel.RelationName)) return err } diff --git a/pgserver/logrepl/subscription.go b/pgserver/logrepl/subscription.go index 7bef15b6..a50379a5 100644 --- a/pgserver/logrepl/subscription.go +++ b/pgserver/logrepl/subscription.go @@ -124,7 +124,24 @@ func DeleteSubscription(ctx *sql.Context, name string) error { } func UpdateSubscriptionLsn(ctx *sql.Context, lsn, name string) error { - _, err := adapter.ExecCatalogInTxn(ctx, catalog.InternalTables.PgSubscription.UpdateStmt(keyColumns, lsnValueColumns), lsn, name) + _, tx, release, err := adapter.GetCatalogTxnExecutionSnapshotWithLease(ctx, nil) + if err != nil { + return err + } + defer release() + return UpdateSubscriptionLsnInTxn(ctx, tx, lsn, name) +} + +// UpdateSubscriptionLsnInTxn updates the durable replication position through +// the transaction supplied by the caller. Replication finalization snapshots +// its physical connection and transaction together; reacquiring a catalog +// transaction here could otherwise bind the LSN update to a replacement +// transaction while the caller commits the old one. +func UpdateSubscriptionLsnInTxn(ctx *sql.Context, tx *stdsql.Tx, lsn, name string) error { + if tx == nil { + return fmt.Errorf("subscription LSN update requires an active transaction") + } + _, err := tx.ExecContext(ctx, catalog.InternalTables.PgSubscription.UpdateStmt(keyColumns, lsnValueColumns), lsn, name) return err } diff --git a/pgserver/object_routing.go b/pgserver/object_routing.go new file mode 100644 index 00000000..72118d22 --- /dev/null +++ b/pgserver/object_routing.go @@ -0,0 +1,438 @@ +package pgserver + +import ( + stdsql "database/sql" + "fmt" + "strings" + + "github.com/apecloud/myduckserver/adapter" + "github.com/apecloud/myduckserver/backend" + "github.com/apecloud/myduckserver/catalog" + "github.com/cockroachdb/cockroachdb-parser/pkg/sql/sem/tree" + "github.com/dolthub/go-mysql-server/sql" +) + +// rewritePostgresObjectRelations maps logical PostgreSQL table references to +// their service-owned DuckLake relations. The mapping is reread from durable +// shadow comments for every statement, so a restarted provider does not rely +// on process-local table state. CREATE/COMMENT statements intentionally remain +// on the logical shadow path because that is where GMS metadata is stored. +func (h *DuckHandler) rewritePostgresObjectRelations(ctx *sql.Context, query string) (string, error) { + if postgresShadowOnlyStatement(query) { + return query, nil + } + execer, conn, _, release, err := adapter.GetExecutionSnapshotWithLease(ctx) + if err != nil { + return "", err + } + defer release() + return h.rewritePostgresObjectRelationsWithSnapshot(ctx, query, execer, conn) +} + +// rewritePostgresObjectRelationsWithSnapshot resolves object metadata through +// the exact executor that will execute the returned SQL. It is used by the +// simple and extended PostgreSQL paths to close the metadata/execution TOCTOU +// window. +func (h *DuckHandler) rewritePostgresObjectRelationsWithSnapshot( + ctx *sql.Context, + query string, + execer adapter.SQLExecutor, + conn *stdsql.Conn, +) (string, error) { + if h == nil || ctx == nil || query == "" { + return query, nil + } + provider := h.GetCatalogProvider() + if provider == nil || !provider.DuckLakeEnabled() || postgresShadowOnlyStatement(query) { + return query, nil + } + // Resolve the namespace through the executor captured by the caller. + // Calling adapter.GetCurrentCatalog here would re-enter ConnectionPool's + // lifecycle lock; a pending shutdown writer can therefore deadlock a caller + // that is already holding the retained execution lease. The executor query + // observes the same transaction/connection snapshot as the eventual + // statement. + catalogName, _ := postgresCurrentNamespaceWithExecutor(ctx, execer, provider) + if conn != nil { + if err := h.ensurePostgresObjectConnectionOnSnapshot(ctx, execer, conn); err != nil { + return "", err + } + } + mappings, err := provider.ObjectTablesWithExecutor(ctx, execer, catalogName) + if err != nil { + return "", err + } + if len(mappings) == 0 { + return query, nil + } + routes := make(map[string]string, len(mappings)*3) + bare := make(map[string]string, len(mappings)) + ambiguous := make(map[string]bool) + for _, mapping := range mappings { + table := strings.ToLower(mapping.Table) + schema := strings.ToLower(mapping.Schema) + if schema != "" { + routes[schema+"."+table] = mapping.PhysicalName + } + if previous, ok := bare[table]; ok && previous != mapping.PhysicalName { + delete(bare, table) + ambiguous[table] = true + } else if !ambiguous[table] { + bare[table] = mapping.PhysicalName + } + } + for table, physical := range bare { + routes[table] = physical + } + rewritten, changed := backend.RewriteSQLRelations(query, routes) + _ = changed // retained for readability at call sites and future diagnostics + return rewritten, nil +} + +// postgresCurrentNamespaceWithExecutor returns the current catalog/schema from +// a previously captured SQL executor. It intentionally never consults the +// session holder, because callers may be retaining the pool lifecycle read +// lease while resolving object metadata. +func postgresCurrentNamespaceWithExecutor( + ctx *sql.Context, + execer adapter.SQLExecutor, + provider *catalog.DatabaseProvider, +) (catalogName, schemaName string) { + if provider != nil { + catalogName = provider.DefaultCatalogName() + } + if ctx != nil { + schemaName = ctx.GetCurrentDatabase() + } + if execer == nil { + return catalogName, schemaName + } + var currentCatalog, currentSchema string + if err := execer.QueryRowContext(ctx, "SELECT CURRENT_CATALOG, CURRENT_SCHEMA").Scan(¤tCatalog, ¤tSchema); err == nil { + if currentCatalog != "" { + catalogName = currentCatalog + } + if currentSchema != "" { + schemaName = currentSchema + } + } + return catalogName, schemaName +} + +// ensurePostgresObjectConnection initializes DuckLake on the session's +// execution connection. Object-table protocol helpers call this after they +// resolve a durable mapping; local-table paths remain untouched. +func (h *DuckHandler) ensurePostgresObjectConnection(ctx *sql.Context) error { + if h == nil || ctx == nil { + return nil + } + provider := h.GetCatalogProvider() + if provider == nil || !provider.DuckLakeEnabled() { + return nil + } + execer, conn, _, release, err := adapter.GetExecutionSnapshotWithLease(ctx) + if err != nil { + return err + } + defer release() + return h.ensurePostgresObjectConnectionOnSnapshot(ctx, execer, conn) +} + +// ensurePostgresObjectConnectionOnSnapshot initializes DuckLake using the +// executor/connection pair that selected the relation metadata. When the +// caller retains an execution lease, passing that executor is essential: +// falling back to a session binding lookup would try to acquire the pool's +// lifecycle read lock a second time while a shutdown writer is pending. +func (h *DuckHandler) ensurePostgresObjectConnectionOnSnapshot( + ctx *sql.Context, + execer adapter.SQLExecutor, + conn *stdsql.Conn, +) error { + if h == nil || ctx == nil || conn == nil { + return nil + } + provider := h.GetCatalogProvider() + if provider == nil || !provider.DuckLakeEnabled() { + return nil + } + if execer != nil { + return provider.EnsureDuckLakeConnectionWithExecutor(ctx, execer, conn) + } + return provider.EnsureDuckLakeConnection(ctx, conn) +} + +// postgresCopySnapshotLease selects and retains the executor/physical owner +// for an asynchronous COPY exchange. Frontend DuckLake admission is acquired +// before the pool snapshot so rollback cleanup cannot overtake the accepted +// operation. The returned lease keeps those resources independent: callers +// release its snapshot half when the producer/loader closes, while the +// operation half is handed to the physical transaction finalizer when the +// session supports that lifecycle surface. +func postgresCopySnapshotLease( + ctx *sql.Context, + handler *DuckHandler, + statement tree.Statement, +) (adapter.SQLExecutor, *stdsql.Conn, *stdsql.Tx, *postgresCopyLease, error) { + if ctx == nil || ctx.Session == nil { + return nil, nil, nil, nil, fmt.Errorf("COPY execution session is unavailable") + } + if _, ok := ctx.Session.(adapter.ConnectionHolder); !ok { + return nil, nil, nil, nil, fmt.Errorf("COPY execution session has no connection holder") + } + + // Capture the transaction owner before admitting the operation. PostgreSQL + // opens its implicit/explicit physical transaction before COPY setup, and + // passing that identity lets rollback cleanup exempt this operation while it + // is finalizing. If a transaction appears in the tiny snapshot gap below, + // the operation is reacquired with the newly observed owner. + var initialTx *stdsql.Tx + _, initialTx = adapter.TryGetTxnBinding(ctx) + operationRelease := func() {} + if handler != nil { + var admissionErr error + operationRelease, admissionErr = handler.beginPostgresDuckLakeOperationForOwnerWithError(ctx, statement, initialTx) + if admissionErr != nil { + return nil, nil, nil, nil, admissionErr + } + } + execer, conn, tx, snapshotRelease, err := adapter.GetExecutionSnapshotWithLease(ctx) + if err != nil { + operationRelease() + return nil, nil, tx, nil, err + } + if initialTx != tx { + if initialTx != nil { + // A transaction that disappeared or was replaced after admission is + // not a safe target for an asynchronous COPY. Do not execute on the + // replacement connection and do not leave either lease held. + snapshotRelease() + operationRelease() + return nil, nil, tx, nil, fmt.Errorf("COPY transaction binding changed while opening") + } + // The operation was admitted anonymously before a concurrent transaction + // became visible. Rebind it to the physical owner captured by the lease; + // the active transaction itself prevents rollback cleanup from overtaking + // this short reacquisition window. + operationRelease() + if handler != nil { + var admissionErr error + operationRelease, admissionErr = handler.beginPostgresDuckLakeOperationForOwnerWithError(ctx, statement, tx) + if admissionErr != nil { + snapshotRelease() + return nil, nil, tx, nil, admissionErr + } + } else { + operationRelease = func() {} + } + } + lease := newPostgresCopyLease(snapshotRelease, operationRelease) + if tx != nil { + lease.registerPhysicalFinalization(ctx, tx) + } + return execer, conn, tx, lease, nil +} + +// postgresCopySnapshot retains the historical snapshot-only helper for +// callers that do not span an asynchronous protocol exchange. New COPY +// producers/consumers use postgresCopySnapshotLease instead. +func postgresCopySnapshot(ctx *sql.Context) (adapter.SQLExecutor, *stdsql.Conn, error) { + execer, conn, _, lease, err := postgresCopySnapshotLease(ctx, nil, nil) + if lease != nil { + lease.Release() + } + if err != nil { + return nil, nil, err + } + return execer, conn, nil +} + +// resolvePostgresCopyTarget resolves a validated GMS table to the identifier +// that DuckDB should receive. Object tables are routed to their durable +// DuckLake relation; local tables retain their catalog-qualified name. The +// caller must pass the executor/connection pair returned by +// postgresCopySnapshot so metadata and COPY execution share one snapshot. +func resolvePostgresCopyTarget( + ctx *sql.Context, + handler *DuckHandler, + schema string, + table sql.Table, + execer adapter.SQLExecutor, + conn *stdsql.Conn, +) (string, error) { + if table == nil { + return "", nil + } + if handler != nil { + if provider := handler.GetCatalogProvider(); provider != nil { + if physical, _, err := provider.PhysicalTableNameForTableWithExecutor(ctx, table, execer, conn); err != nil { + return "", err + } else if physical != "" { + return physical, nil + } + } + } + if schema != "" { + return catalog.ConnectIdentifiersANSI(schema, table.Name()), nil + } + return catalog.QuoteIdentifierANSI(table.Name()), nil +} + +func postgresShadowOnlyStatement(query string) bool { + trimmed := strings.TrimSpace(query) + for len(trimmed) > 0 { + switch { + case strings.HasPrefix(trimmed, "--"): + if idx := strings.IndexByte(trimmed, '\n'); idx >= 0 { + trimmed = strings.TrimSpace(trimmed[idx+1:]) + continue + } + return true + case strings.HasPrefix(trimmed, "/*"): + if idx := strings.Index(trimmed[2:], "*/"); idx >= 0 { + trimmed = strings.TrimSpace(trimmed[idx+4:]) + continue + } + return true + } + break + } + if trimmed == "" { + return true + } + word := strings.ToLower(trimmed) + for i, r := range word { + if (r < 'a' || r > 'z') && (r < '0' || r > '9') && r != '_' { + word = word[:i] + break + } + } + switch word { + case "create", "comment", "show", "set", "begin", "commit", "rollback", "vacuum", "attach", "detach": + return true + default: + return false + } +} + +// executePostgresObjectDDL keeps the durable local shadow table in lockstep +// with its DuckLake relation. The normal PostgreSQL handler executes SQL +// directly through database/sql and therefore bypasses GMS's catalog drop and +// alter callbacks. For an object table, execute the routed statement against +// the physical relation and then apply the same operation to the shadow using +// the exact executor snapshot supplied by the caller. Local-table statements +// return handled=false and retain their existing one-statement behavior. +// +// The helper deliberately limits the shadow statement for DROP TABLE to object +// names only. A mixed DROP containing local tables must not attempt to drop the +// local names a second time. +func (h *DuckHandler) executePostgresObjectDDL( + ctx *sql.Context, + query, routedQuery string, + stmt tree.Statement, + execer adapter.SQLExecutor, + conn *stdsql.Conn, + args ...any, +) (result stdsql.Result, handled bool, err error) { + provider := h.GetCatalogProvider() + if provider == nil || !provider.DuckLakeEnabled() || execer == nil { + return nil, false, nil + } + currentCatalog, currentSchema := postgresCurrentNamespaceWithExecutor(ctx, execer, provider) + // Resolve durable object mappings through the same executor as the DDL. A + // GMS table lookup would call DatabaseProvider.Database and reacquire the + // session lifecycle lock while this caller retains its execution lease. + mappings, mappingErr := provider.ObjectTablesWithExecutor(ctx, execer, currentCatalog) + if mappingErr != nil { + return nil, false, mappingErr + } + + resolve := func(name *tree.TableName) (sql.Table, bool, string, error) { + if name == nil || name.Table() == "" { + return nil, false, "", nil + } + schemaName := name.Schema() + if schemaName == "" { + schemaName = currentSchema + } + catalogName := name.Catalog() + if catalogName == "" { + catalogName = currentCatalog + } + for _, mapping := range mappings { + if strings.EqualFold(mapping.Catalog, catalogName) && + strings.EqualFold(mapping.Schema, schemaName) && + strings.EqualFold(mapping.Table, name.Table()) { + logical := catalog.FullTableName(catalogName, schemaName, mapping.Table) + return nil, true, logical, nil + } + } + return nil, false, "", nil + } + + switch node := stmt.(type) { + case *tree.DropTable: + var shadows []string + for i := range node.Names { + name := &node.Names[i] + _, object, logical, lookupErr := resolve(name) + if lookupErr != nil { + if node.IfExists && sql.ErrTableNotFound.Is(lookupErr) { + continue + } + return nil, false, lookupErr + } + if object { + shadows = append(shadows, logical) + } + } + if len(shadows) == 0 { + return nil, false, nil + } + result, err = execer.ExecContext(ctx, routedQuery, args...) + if err != nil { + return nil, true, err + } + var shadow strings.Builder + shadow.WriteString("DROP TABLE ") + if node.IfExists { + shadow.WriteString("IF EXISTS ") + } + shadow.WriteString(strings.Join(shadows, ", ")) + if node.DropBehavior != tree.DropDefault { + shadow.WriteByte(' ') + shadow.WriteString(node.DropBehavior.String()) + } + if _, err = execer.ExecContext(ctx, shadow.String()); err != nil && !node.IfExists { + return result, true, err + } + return result, true, nil + + case *tree.AlterTable: + if node.Table == nil { + return nil, false, nil + } + name := node.Table.ToTableName() + _, object, _, lookupErr := resolve(&name) + if lookupErr != nil { + if node.IfExists && sql.ErrTableNotFound.Is(lookupErr) { + return nil, false, nil + } + return nil, false, lookupErr + } + if !object { + return nil, false, nil + } + result, err = execer.ExecContext(ctx, routedQuery, args...) + if err != nil { + return nil, true, err + } + // The raw query still names the logical shadow relation. It is executed + // only after the physical statement succeeds, and on the same executor. + if _, err = execer.ExecContext(ctx, query, args...); err != nil { + return result, true, err + } + return result, true, nil + default: + return nil, false, nil + } +} diff --git a/pgserver/pending_copy_release_test.go b/pgserver/pending_copy_release_test.go new file mode 100644 index 00000000..a6d1c863 --- /dev/null +++ b/pgserver/pending_copy_release_test.go @@ -0,0 +1,81 @@ +package pgserver + +import ( + "context" + stdsql "database/sql" + "errors" + "sync/atomic" + "testing" + + "github.com/apecloud/myduckserver/mycontext" + "github.com/dolthub/go-mysql-server/memory" + "github.com/dolthub/go-mysql-server/sql" + "github.com/stretchr/testify/require" +) + +// orderedRollbackSession lets this regression observe the exact boundary +// between releasing a completed COPY lease and starting implicit rollback +// without changing the production session implementation. +type orderedRollbackSession struct { + *replacingTransactionSession + onRollback func() +} + +func (s *orderedRollbackSession) RollbackTxn(tx *stdsql.Tx) (*stdsql.Conn, bool, error) { + if s.onRollback != nil { + s.onRollback() + } + return s.replacingTransactionSession.RollbackTxn(tx) +} + +// A completed extended COPY keeps its lease until Sync so a successful group +// can commit before releasing the provider operation. An error in a later +// pipelined message must release that lease before rollback, or the rollback +// reservation can wait on its own operation forever. +func TestImplicitErrorReleasesPendingCopyLeaseBeforeRollback(t *testing.T) { + owner := new(stdsql.Conn) + transaction := new(stdsql.Tx) + baseSession := &replacingTransactionSession{ + Session: memory.NewSession(sql.NewBaseSession(), nil), + conn: owner, + tx: transaction, + } + var events []string + session := &orderedRollbackSession{ + replacingTransactionSession: baseSession, + onRollback: func() { + events = append(events, "rollback-start") + }, + } + h := &ConnectionHandler{ + implicitTx: transaction, + implicitTxConn: owner, + txStatus: ReadyForQueryTransactionIndicator_Idle, + } + ctx := sql.NewContext( + withPostgresConnectionHandler(mycontext.WithFrontendQuery(context.Background()), h), + sql.WithSession(session), + ) + h.implicitTxCtx = ctx + + var releaseCalls atomic.Int32 + h.pendingCopyRelease = func() { + events = append(events, "lease-release") + releaseCalls.Add(1) + } + + err := h.finishImplicitPostgresError(errors.New("pipelined statement failed")) + require.ErrorContains(t, err, "pipelined statement failed") + require.Equal(t, int32(1), releaseCalls.Load()) + require.True(t, session.rollbackCalled) + require.Equal(t, []string{"lease-release", "rollback-start"}, events) + require.Nil(t, session.tx) + require.Nil(t, h.pendingCopyRelease) + require.Nil(t, h.implicitTx) + require.Equal(t, ReadyForQueryTransactionIndicator_Idle, h.readyForQueryStatus()) + + // The release slot is cleared at the error boundary, so a repeated cleanup + // call cannot invoke the callback again. + _ = h.finishImplicitPostgresError(nil) + require.Equal(t, int32(1), releaseCalls.Load()) +} diff --git a/pgserver/prepared_returning_test.go b/pgserver/prepared_returning_test.go new file mode 100644 index 00000000..4c68ad0f --- /dev/null +++ b/pgserver/prepared_returning_test.go @@ -0,0 +1,68 @@ +package pgserver + +import ( + "context" + "regexp" + "testing" + + "github.com/DATA-DOG/go-sqlmock" + "github.com/apecloud/myduckserver/backend" + "github.com/apecloud/myduckserver/pgtypes" + "github.com/cockroachdb/cockroachdb-parser/pkg/sql/parser" + "github.com/cockroachdb/cockroachdb-parser/pkg/sql/sem/tree" + "github.com/dolthub/go-mysql-server/sql" + "github.com/jackc/pgx/v5/pgtype" + "github.com/stretchr/testify/require" +) + +func TestPreparedInsertReturningMetadataUsesActiveExecutionSnapshot(t *testing.T) { + parsed, err := parser.ParseOne("INSERT INTO object_rows VALUES ($1, $2) RETURNING id, name") + require.NoError(t, err) + insert, ok := parsed.AST.(*tree.Insert) + require.True(t, ok) + + metadataQuery := postgresInsertReturningSchemaQuery(insert) + require.Equal(t, + "SELECT id, name FROM object_rows LIMIT 0", + metadataQuery, + ) + routedQuery, changed := backend.RewriteSQLRelations(metadataQuery, map[string]string{ + "object_rows": `"__myduck_ducklake"."app"."object_rows"`, + }) + require.True(t, changed) + require.Equal(t, + `SELECT id, name FROM "__myduck_ducklake"."app"."object_rows" LIMIT 0`, + routedQuery, + ) + + db, mock, err := sqlmock.New() + require.NoError(t, err) + t.Cleanup(func() { _ = db.Close() }) + + mock.ExpectBegin() + tx, err := db.BeginTx(context.Background(), nil) + require.NoError(t, err) + + mock.ExpectQuery(regexp.QuoteMeta(routedQuery)).WillReturnRows( + sqlmock.NewRowsWithColumnDefinition( + sqlmock.NewColumn("id").OfType("INTEGER", int64(0)).Nullable(false), + sqlmock.NewColumn("name").OfType("VARCHAR", "").Nullable(false), + ), + ) + + // Passing a context without a session makes an accidental fallback to + // adapter.QueryCatalog fail immediately. The explicit *sql.Tx is the only + // legal metadata executor for this active snapshot. + schema, err := inferPostgresInsertReturningSchema(sql.NewEmptyContext(), routedQuery, tx) + require.NoError(t, err) + require.Len(t, schema, 2) + require.Equal(t, "id", schema[0].Name) + require.Equal(t, "name", schema[1].Name) + idType, ok := schema[0].Type.(pgtypes.PostgresType) + require.True(t, ok) + require.Equal(t, uint32(pgtype.Int4OID), idType.PG.OID) + + mock.ExpectRollback() + require.NoError(t, tx.Rollback()) + require.NoError(t, mock.ExpectationsWereMet()) +} diff --git a/pgserver/ready_for_query_test.go b/pgserver/ready_for_query_test.go new file mode 100644 index 00000000..718d4bd8 --- /dev/null +++ b/pgserver/ready_for_query_test.go @@ -0,0 +1,202 @@ +package pgserver + +import ( + "bytes" + "errors" + "io" + "testing" + + "github.com/cockroachdb/cockroachdb-parser/pkg/sql/sem/tree" + "github.com/jackc/pgx/v5/pgconn" + "github.com/jackc/pgx/v5/pgproto3" + "github.com/stretchr/testify/require" +) + +func TestConnectionHandlerTransactionStatusTransitions(t *testing.T) { + h := &ConnectionHandler{} + + // A zero-value handler is equivalent to a newly opened connection. + require.Equal(t, ReadyForQueryTransactionIndicator_Idle, h.readyForQueryStatus()) + + h.recordTransactionResult(&tree.BeginTransaction{}, nil) + require.Equal(t, ReadyForQueryTransactionIndicator_TransactionBlock, h.readyForQueryStatus()) + + // Successful work inside the block preserves T. + h.recordTransactionResult(nil, nil) + require.Equal(t, ReadyForQueryTransactionIndicator_TransactionBlock, h.readyForQueryStatus()) + + // The first statement error moves an explicit block to E, and later + // non-control results cannot clear it. + h.recordTransactionResult(nil, errors.New("statement failed")) + require.Equal(t, ReadyForQueryTransactionIndicator_FailedTransactionBlock, h.readyForQueryStatus()) + h.recordTransactionResult(nil, nil) + require.Equal(t, ReadyForQueryTransactionIndicator_FailedTransactionBlock, h.readyForQueryStatus()) + + // Cleanup errors after a completed rollback must still advertise idle. + h.recordTransactionResult(&tree.RollbackTransaction{}, errors.New("cleanup failed")) + require.Equal(t, ReadyForQueryTransactionIndicator_Idle, h.readyForQueryStatus()) + + h.recordTransactionResult(&tree.BeginTransaction{}, nil) + require.Equal(t, ReadyForQueryTransactionIndicator_TransactionBlock, h.readyForQueryStatus()) + h.recordTransactionResult(&tree.CommitTransaction{}, errors.New("commit failed")) + require.Equal(t, ReadyForQueryTransactionIndicator_Idle, h.readyForQueryStatus()) + + // A nested BEGIN is a warning/no-op, so PostgreSQL keeps the original block + // active rather than entering the failed E state. + h.txStatus = ReadyForQueryTransactionIndicator_TransactionBlock + h.recordTransactionResult(&tree.BeginTransaction{}, nil) + require.Equal(t, ReadyForQueryTransactionIndicator_TransactionBlock, h.readyForQueryStatus()) +} + +func TestConnectionHandlerEndOfMessagesUsesTransactionStatus(t *testing.T) { + var output bytes.Buffer + h := &ConnectionHandler{ + backend: pgproto3.NewBackend(bytes.NewReader(nil), &output), + txStatus: ReadyForQueryTransactionIndicator_TransactionBlock, + } + + h.endOfMessages(errors.New("statement failed")) + + frontend := pgproto3.NewFrontend(bytes.NewReader(output.Bytes()), io.Discard) + message, err := frontend.Receive() + require.NoError(t, err) + pgErr, ok := message.(*pgproto3.ErrorResponse) + require.True(t, ok) + require.Equal(t, "XX000", pgErr.Code) + + message, err = frontend.Receive() + require.NoError(t, err) + ready, ok := message.(*pgproto3.ReadyForQuery) + require.True(t, ok) + require.Equal(t, byte(ReadyForQueryTransactionIndicator_FailedTransactionBlock), ready.TxStatus) +} + +func TestConnectionHandlerRejectsStatementsInFailedTransaction(t *testing.T) { + h := &ConnectionHandler{txStatus: ReadyForQueryTransactionIndicator_FailedTransactionBlock} + + for _, statement := range []ConvertedStatement{ + {AST: &tree.CommitTransaction{}, String: "COMMIT"}, + {AST: &tree.RollbackTransaction{}, String: "ROLLBACK"}, + {String: ""}, + } { + require.NoError(t, h.rejectStatementIfTransactionFailed(statement)) + } + + for _, statement := range []ConvertedStatement{ + {AST: &tree.BeginTransaction{}, String: "BEGIN"}, + {AST: &tree.Select{}, String: "SELECT 1"}, + {String: "not valid transaction syntax"}, + } { + err := h.rejectStatementIfTransactionFailed(statement) + var pgErr *pgconn.PgError + require.ErrorAs(t, err, &pgErr) + require.Equal(t, "25P02", pgErr.Code) + require.Equal(t, postgresFailedTransactionMessage, pgErr.Message) + } +} + +func TestConnectionHandlerFailedQueryGateModelsControlsInSimpleBatch(t *testing.T) { + h := &ConnectionHandler{txStatus: ReadyForQueryTransactionIndicator_FailedTransactionBlock} + + // A recovery control permits later statements in the same simple-query + // message to be evaluated from the newly idle state. + require.NoError(t, h.rejectFailedQueryStatements([]ConvertedStatement{ + {AST: &tree.RollbackTransaction{}, String: "ROLLBACK"}, + {AST: &tree.Select{}, String: "SELECT 1"}, + })) + + // Without a recovery control, the first ordinary statement is rejected. + var pgErr *pgconn.PgError + err := h.rejectFailedQueryStatements([]ConvertedStatement{ + {AST: &tree.Select{}, String: "SELECT 1"}, + {AST: &tree.RollbackTransaction{}, String: "ROLLBACK"}, + }) + require.ErrorAs(t, err, &pgErr) + require.Equal(t, "25P02", pgErr.Code) +} + +func TestConnectionHandlerFailedParseUsesTransactionStateBeforeSQLGates(t *testing.T) { + h := &ConnectionHandler{txStatus: ReadyForQueryTransactionIndicator_FailedTransactionBlock} + + // Parse is a protocol message, but its SQL still must be rejected while the + // transaction is failed. No engine/session is needed to exercise this gate. + err := h.handleParse(&pgproto3.Parse{Query: "SELECT 1", Name: "failed"}) + var pgErr *pgconn.PgError + require.ErrorAs(t, err, &pgErr) + require.Equal(t, "25P02", pgErr.Code) +} + +func TestConnectionHandlerSemicolonShortcutPreservesFailedState(t *testing.T) { + var output bytes.Buffer + h := &ConnectionHandler{ + backend: pgproto3.NewBackend(bytes.NewReader(nil), &output), + txStatus: ReadyForQueryTransactionIndicator_FailedTransactionBlock, + } + + end, err := h.handleQuery(&pgproto3.Query{String: ";"}) + require.NoError(t, err) + require.True(t, end) + require.Equal(t, ReadyForQueryTransactionIndicator_FailedTransactionBlock, h.readyForQueryStatus()) +} + +func TestConnectionHandlerCopyDataWithoutCopyStateReturnsError(t *testing.T) { + h := &ConnectionHandler{} + require.NotPanics(t, func() { + _, _, err := h.handleCopyData(&pgproto3.CopyData{Data: []byte("x")}) + require.Error(t, err) + }) +} + +func TestConnectionHandlerSerializesFailedTransactionSQLState(t *testing.T) { + var output bytes.Buffer + h := &ConnectionHandler{ + backend: pgproto3.NewBackend(bytes.NewReader(nil), &output), + txStatus: ReadyForQueryTransactionIndicator_FailedTransactionBlock, + } + + h.endOfMessages(postgresFailedTransactionError()) + frontend := pgproto3.NewFrontend(bytes.NewReader(output.Bytes()), io.Discard) + message, err := frontend.Receive() + require.NoError(t, err) + pgErr, ok := message.(*pgproto3.ErrorResponse) + require.True(t, ok) + require.Equal(t, "25P02", pgErr.Code) + require.Equal(t, postgresFailedTransactionMessage, pgErr.Message) + + message, err = frontend.Receive() + require.NoError(t, err) + ready, ok := message.(*pgproto3.ReadyForQuery) + require.True(t, ok) + require.Equal(t, byte(ReadyForQueryTransactionIndicator_FailedTransactionBlock), ready.TxStatus) +} + +func TestConnectionHandlerCommitInFailedTransactionReportsRollback(t *testing.T) { + h := &ConnectionHandler{txStatus: ReadyForQueryTransactionIndicator_FailedTransactionBlock} + statement := ConvertedStatement{AST: &tree.CommitTransaction{}, Tag: "COMMIT"} + require.Equal(t, "ROLLBACK", h.protocolCommandTag(statement)) + h.txStatus = ReadyForQueryTransactionIndicator_TransactionBlock + require.Equal(t, "COMMIT", h.protocolCommandTag(statement)) +} + +func TestConnectionHandlerDescribePreparedPortalWithoutRawStatement(t *testing.T) { + var output bytes.Buffer + h := &ConnectionHandler{ + backend: pgproto3.NewBackend(bytes.NewReader(nil), &output), + portals: map[string]PortalData{ + "portal": { + Prepared: true, + Statement: ConvertedStatement{Tag: "SELECT"}, + Fields: []pgproto3.FieldDescription{{Name: []byte("value"), DataTypeOID: 25}}, + }, + }, + } + + require.NoError(t, h.handleDescribe(&pgproto3.Describe{ObjectType: 'P', Name: "portal"})) + frontend := pgproto3.NewFrontend(bytes.NewReader(output.Bytes()), io.Discard) + message, err := frontend.Receive() + require.NoError(t, err) + description, ok := message.(*pgproto3.RowDescription) + require.True(t, ok) + require.Len(t, description.Fields, 1) + require.Equal(t, []byte("value"), description.Fields[0].Name) +} diff --git a/pgserver/returning_test.go b/pgserver/returning_test.go index aa8a9e9f..ae3e1650 100644 --- a/pgserver/returning_test.go +++ b/pgserver/returning_test.go @@ -72,6 +72,25 @@ func TestInsertReturningWireProtocols(t *testing.T) { 4, "postgres prepared") }) + t.Run("postgres prepared rollback", func(t *testing.T) { + conn, err := pgx.Connect(ctx, pgURL) + require.NoError(t, err) + defer conn.Close(ctx) + + tx, err := conn.Begin(ctx) + require.NoError(t, err) + _, err = tx.Prepare(ctx, "returning_wire_tx_prepared", "INSERT INTO returning_wire.returning_rows VALUES ($1, $2) RETURNING id, name") + require.NoError(t, err) + requireReturningWireRow(t, + tx.QueryRow(ctx, "returning_wire_tx_prepared", int32(5), "postgres prepared rollback"), + 5, "postgres prepared rollback") + require.NoError(t, tx.Rollback(ctx)) + + var count int + require.NoError(t, conn.QueryRow(ctx, "SELECT count(*) FROM returning_wire.returning_rows").Scan(&count)) + require.Equal(t, 3, count) + }) + var count int err = testEnv.MyDuckServer.QueryRow("SELECT count(*) FROM returning_rows").Scan(&count) require.NoError(t, err) diff --git a/pgserver/rollback_cleanup.go b/pgserver/rollback_cleanup.go new file mode 100644 index 00000000..5b5a3028 --- /dev/null +++ b/pgserver/rollback_cleanup.go @@ -0,0 +1,191 @@ +package pgserver + +import ( + "context" + stdsql "database/sql" + "fmt" + + "github.com/apecloud/myduckserver/adapter" + "github.com/apecloud/myduckserver/catalog" + "github.com/apecloud/myduckserver/mycontext" + "github.com/cockroachdb/cockroachdb-parser/pkg/sql/sem/tree" + "github.com/dolthub/go-mysql-server/sql" +) + +// postgresRollbackOrphanCleaner is the small provider surface needed by the +// PostgreSQL protocol path. Keeping it narrow makes the rollback hook easy to +// exercise without constructing a full server in unit tests. +type postgresRollbackOrphanCleaner interface { + DuckLakeEnabled() bool + CleanupDuckLakeOrphansOnConn(context.Context, *stdsql.Conn) error +} + +type postgresRollbackObjectStorageGate interface { + DuckLakeObjectStorageEnabled() bool +} + +// postgresRollbackCleanupReservation is implemented by the production +// provider. It is intentionally optional so the narrow provider probes used by +// unit tests do not need to model the provider's transaction barrier. +type postgresRollbackCleanupReservation interface { + BeginDuckLakeRollbackCleanup(*stdsql.Tx) func() +} + +type postgresRollbackCleanupReservationWithError interface { + BeginDuckLakeRollbackCleanupWithError(*stdsql.Tx) (func(), error) +} + +type postgresRollbackCleanupOwnerReservation interface { + BeginDuckLakeRollbackCleanupForOwner(*stdsql.Tx, any) func() +} + +type postgresRollbackCleanupOwnerReservationWithError interface { + BeginDuckLakeRollbackCleanupForOwnerWithError(*stdsql.Tx, any) (func(), error) +} + +func beginPostgresRollbackCleanupReservation( + ctx *sql.Context, + provider postgresRollbackOrphanCleaner, + target *stdsql.Tx, +) func() { + release, _ := beginPostgresRollbackCleanupReservationWithError(ctx, provider, target) + return release +} + +func beginPostgresRollbackCleanupReservationWithError( + ctx *sql.Context, + provider postgresRollbackOrphanCleaner, + target *stdsql.Tx, +) (func(), error) { + if ctx == nil || provider == nil || mycontext.QueryOrigin(ctx) != mycontext.FrontendQueryOrigin { + return func() {}, nil + } + if ownerGate, ownerOK := provider.(postgresRollbackCleanupOwnerReservationWithError); ownerOK { + owner := any(target) + if owner == nil { + if handler := postgresConnectionHandler(ctx); handler != nil && handler.implicitTx != nil { + owner = handler.implicitTx + } else if ctx.GetTransaction() != nil { + owner = ctx.GetTransaction() + } + } + return ownerGate.BeginDuckLakeRollbackCleanupForOwnerWithError(target, owner) + } + if gate, ok := provider.(postgresRollbackCleanupReservationWithError); ok { + return gate.BeginDuckLakeRollbackCleanupWithError(target) + } + gate, ok := provider.(postgresRollbackCleanupReservation) + if ownerGate, ownerOK := provider.(postgresRollbackCleanupOwnerReservation); ownerOK { + // PostgreSQL transaction control owns the raw physical *sql.Tx. Pass it as + // both identities when present so operation leases admitted before the + // snapshot can be exempted without weakening the anonymous/other-owner + // barrier. A GMS-only scope has no physical target; use its logical marker + // when the handler/context exposes one. + owner := any(target) + if owner == nil { + if handler := postgresConnectionHandler(ctx); handler != nil && handler.implicitTx != nil { + owner = handler.implicitTx + } else if ctx.GetTransaction() != nil { + owner = ctx.GetTransaction() + } + } + return ownerGate.BeginDuckLakeRollbackCleanupForOwner(target, owner), nil + } + if !ok { + return func() {}, nil + } + release := gate.BeginDuckLakeRollbackCleanup(target) + if release == nil { + return func() {}, nil + } + return release, nil +} + +// cleanupPostgresRollbackOrphans runs the product-owned DuckLake cleanup after +// a PostgreSQL ROLLBACK has completed. PostgreSQL transaction control is +// executed directly on the session connection, so it does not pass through +// backend.Session.Rollback (the MySQL cleanup hook). The same connection must +// be reused here to preserve its DuckLake attachment and session state. +func cleanupPostgresRollbackOrphans( + ctx *sql.Context, + statement tree.Statement, + provider postgresRollbackOrphanCleaner, + connections ...*stdsql.Conn, +) error { + if ctx == nil || provider == nil { + return nil + } + // The public entry point is a frontend rollback, while finalizer callbacks + // derive a service-owned maintenance context before invoking this helper. + // Accept both explicit origins, but continue to reject unknown and + // replication contexts so a replication rollback can never trigger lake + // cleanup. + origin := mycontext.QueryOrigin(ctx) + if origin != mycontext.FrontendQueryOrigin && origin != mycontext.MaintenanceQueryOrigin { + return nil + } + if _, ok := statement.(*tree.RollbackTransaction); !ok { + return nil + } + if !provider.DuckLakeEnabled() { + return nil + } + gate, ok := provider.(postgresRollbackObjectStorageGate) + if !ok || !gate.DuckLakeObjectStorageEnabled() { + return nil + } + var conn *stdsql.Conn + if len(connections) > 0 { + conn = connections[0] + } + ownerConn := catalog.RollbackCleanupOwner(ctx) + if ownerConn != nil { + if conn != nil && ownerConn != conn { + return fmt.Errorf("postgres rollback cleanup connection is not the rollback owner") + } + conn = ownerConn + } + var holder adapter.ConnectionHolder + if ownerConn == nil { + var ok bool + holder, ok = ctx.Session.(adapter.ConnectionHolder) + if !ok { + return fmt.Errorf("postgres rollback cleanup connection unavailable: session has no connection holder") + } + if _, activeTx := adapter.TryGetTxnBindingForRelease(ctx); activeTx != nil { + return fmt.Errorf("postgres rollback orphan cleanup requires an inactive transaction") + } + } + + // Keep the concrete sql.Context so provider-side transaction guards remain + // effective, while replacing the request origin with an explicit + // service-owned maintenance origin. + cleanupBase := context.WithoutCancel(ctx) + if origin == mycontext.FrontendQueryOrigin { + cleanupBase = mycontext.WithQueryOrigin(cleanupBase, mycontext.MaintenanceQueryOrigin) + } + if conn != nil { + cleanupBase = catalog.WithRollbackCleanupOwner(cleanupBase, conn) + } + cleanupCtx := ctx.WithContext(cleanupBase) + if conn == nil { + var err error + // Acquire only after cancellation has been removed. A client disconnect + // must not prevent the post-rollback maintenance operation from running. + conn, err = holder.GetConn(cleanupCtx) + if err != nil { + return fmt.Errorf("postgres rollback cleanup connection unavailable: %w", err) + } + } + // A replacement transaction may have been opened while the owner was being + // finalized. Never issue maintenance SQL through its connection. + if catalog.RollbackCleanupOwner(cleanupCtx) == nil { + if _, activeTx := adapter.TryGetTxnBindingForRelease(cleanupCtx); activeTx != nil { + return fmt.Errorf("postgres rollback orphan cleanup requires an inactive transaction") + } + } + if err := provider.CleanupDuckLakeOrphansOnConn(cleanupCtx, conn); err != nil { + return fmt.Errorf("postgres rollback orphan cleanup failed: %w", err) + } + return nil +} diff --git a/pgserver/rollback_cleanup_test.go b/pgserver/rollback_cleanup_test.go new file mode 100644 index 00000000..15cbd186 --- /dev/null +++ b/pgserver/rollback_cleanup_test.go @@ -0,0 +1,213 @@ +package pgserver + +import ( + "context" + stdsql "database/sql" + "errors" + "testing" + + "github.com/DATA-DOG/go-sqlmock" + "github.com/apecloud/myduckserver/mycontext" + "github.com/cockroachdb/cockroachdb-parser/pkg/sql/sem/tree" + "github.com/dolthub/go-mysql-server/memory" + "github.com/dolthub/go-mysql-server/sql" + "github.com/stretchr/testify/require" +) + +type rollbackCleanupProviderProbe struct { + enabled bool + objectStorage bool + calls int + gotCtx context.Context + gotConn *stdsql.Conn + cleanupErr error + reservationErr error +} + +func (p *rollbackCleanupProviderProbe) DuckLakeEnabled() bool { return p.enabled } + +func (p *rollbackCleanupProviderProbe) DuckLakeObjectStorageEnabled() bool { + return p.objectStorage +} + +func (p *rollbackCleanupProviderProbe) CleanupDuckLakeOrphansOnConn(ctx context.Context, conn *stdsql.Conn) error { + p.calls++ + p.gotCtx = ctx + p.gotConn = conn + return p.cleanupErr +} + +func (p *rollbackCleanupProviderProbe) BeginDuckLakeRollbackCleanupWithError(*stdsql.Tx) (func(), error) { + return func() {}, p.reservationErr +} + +type rollbackCleanupSessionProbe struct { + *memory.Session + conn *stdsql.Conn + tx *stdsql.Tx + getCalls int + commitCalled bool + rollbackCalled bool +} + +func (s *rollbackCleanupSessionProbe) GetConn(context.Context) (*stdsql.Conn, error) { + s.getCalls++ + return s.conn, nil +} + +func (s *rollbackCleanupSessionProbe) GetTxn(context.Context, *stdsql.TxOptions) (*stdsql.Tx, error) { + return s.tx, nil +} + +func (s *rollbackCleanupSessionProbe) GetCatalogConn(context.Context) (*stdsql.Conn, error) { + return s.conn, nil +} + +func (s *rollbackCleanupSessionProbe) GetCatalogTxn(context.Context, *stdsql.TxOptions) (*stdsql.Tx, error) { + return nil, nil +} + +func (s *rollbackCleanupSessionProbe) TryGetTxn() *stdsql.Tx { return s.tx } + +func (s *rollbackCleanupSessionProbe) GetTxnBinding() (*stdsql.Conn, *stdsql.Tx) { + return s.conn, s.tx +} + +func (s *rollbackCleanupSessionProbe) CommitTxn(*stdsql.Tx) error { + s.commitCalled = true + s.tx = nil + return nil +} + +func (s *rollbackCleanupSessionProbe) RollbackTxn(*stdsql.Tx) (*stdsql.Conn, bool, error) { + s.rollbackCalled = true + s.tx = nil + return s.conn, true, nil +} + +func (s *rollbackCleanupSessionProbe) GetCurrentCatalog() string { return "memory" } + +func (s *rollbackCleanupSessionProbe) GetCurrentSchema() string { return "main" } + +func (s *rollbackCleanupSessionProbe) CloseTxn() {} + +func (s *rollbackCleanupSessionProbe) CloseConn() {} + +func newRollbackCleanupContext(t *testing.T, conn *stdsql.Conn) *sql.Context { + t.Helper() + session := &rollbackCleanupSessionProbe{ + Session: memory.NewSession(sql.NewBaseSession(), nil), + conn: conn, + } + return sql.NewContext(mycontext.WithFrontendQuery(context.Background()), sql.WithSession(session)) +} + +func TestCleanupPostgresRollbackOrphansUsesSameConnectionAndMaintenanceOrigin(t *testing.T) { + db, _, err := sqlmock.New() + require.NoError(t, err) + t.Cleanup(func() { _ = db.Close() }) + conn, err := db.Conn(context.Background()) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, conn.Close()) }) + + provider := &rollbackCleanupProviderProbe{enabled: true, objectStorage: true} + ctx := newRollbackCleanupContext(t, conn) + require.NoError(t, cleanupPostgresRollbackOrphans(ctx, &tree.RollbackTransaction{}, provider)) + require.Equal(t, 1, provider.calls) + require.Same(t, conn, provider.gotConn) + require.IsType(t, (*sql.Context)(nil), provider.gotCtx) + require.Equal(t, mycontext.MaintenanceQueryOrigin, mycontext.QueryOrigin(provider.gotCtx)) +} + +func TestCleanupPostgresRollbackOrphansSkipsIneligibleOriginsAndConfiguration(t *testing.T) { + provider := &rollbackCleanupProviderProbe{enabled: true, objectStorage: true} + ctx := newRollbackCleanupContext(t, nil) + require.NoError(t, cleanupPostgresRollbackOrphans(ctx, &tree.CommitTransaction{}, provider)) + require.Equal(t, 0, provider.calls) + + provider.enabled = false + require.NoError(t, cleanupPostgresRollbackOrphans(ctx, &tree.RollbackTransaction{}, provider)) + require.Equal(t, 0, provider.calls) + + provider.enabled = true + provider.objectStorage = false + require.NoError(t, cleanupPostgresRollbackOrphans(ctx, &tree.RollbackTransaction{}, provider)) + require.Equal(t, 0, provider.calls) + + unknownCtx := sql.NewContext(context.Background(), sql.WithSession(ctx.Session)) + require.NoError(t, cleanupPostgresRollbackOrphans(unknownCtx, &tree.RollbackTransaction{}, provider)) + require.Equal(t, 0, provider.calls) + + replicationCtx := sql.NewContext( + mycontext.WithQueryOrigin(context.Background(), mycontext.PostgresReplicationQueryOrigin), + sql.WithSession(ctx.Session), + ) + require.NoError(t, cleanupPostgresRollbackOrphans(replicationCtx, &tree.RollbackTransaction{}, provider)) + require.Equal(t, 0, provider.calls) +} + +func TestCleanupPostgresRollbackOrphansRejectsActiveTransaction(t *testing.T) { + db, _, err := sqlmock.New() + require.NoError(t, err) + t.Cleanup(func() { _ = db.Close() }) + conn, err := db.Conn(context.Background()) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, conn.Close()) }) + tx := new(stdsql.Tx) + + session := &rollbackCleanupSessionProbe{ + Session: memory.NewSession(sql.NewBaseSession(), nil), + conn: conn, + tx: tx, + } + ctx := sql.NewContext(mycontext.WithFrontendQuery(context.Background()), sql.WithSession(session)) + provider := &rollbackCleanupProviderProbe{enabled: true, objectStorage: true} + err = cleanupPostgresRollbackOrphans(ctx, &tree.RollbackTransaction{}, provider) + require.Error(t, err) + require.Contains(t, err.Error(), "inactive transaction") + require.Zero(t, provider.calls) + require.Zero(t, session.getCalls) +} + +func TestPostgresRollbackReservationErrorStillFinalizesPhysicalTransaction(t *testing.T) { + reservationErr := errors.New("poisoned cleanup reservation") + conn := new(stdsql.Conn) + tx := new(stdsql.Tx) + session := &rollbackCleanupSessionProbe{ + Session: memory.NewSession(sql.NewBaseSession(), nil), + conn: conn, + tx: tx, + } + ctx := sql.NewContext(mycontext.WithFrontendQuery(context.Background()), sql.WithSession(session)) + provider := &rollbackCleanupProviderProbe{ + enabled: true, + objectStorage: true, + reservationErr: reservationErr, + } + + err := rollbackPostgresTransactionControl(ctx, &tree.RollbackTransaction{}, provider) + require.ErrorIs(t, err, reservationErr) + require.True(t, session.rollbackCalled, "cleanup admission failure must not skip physical rollback") + require.Nil(t, session.tx) + require.Zero(t, provider.calls, "poisoned cleanup reservation must not issue cleanup SQL") +} + +func TestDeletePreparedObjectsToleratesProtocolOnlyTransactionStatements(t *testing.T) { + h := &ConnectionHandler{ + preparedStatements: map[string]PreparedStatementData{ + "begin": {Statement: ConvertedStatement{String: "BEGIN"}}, + }, + portals: map[string]PortalData{ + "begin": {Statement: ConvertedStatement{String: "BEGIN"}}, + }, + } + + // Transaction-control statements deliberately have no DuckDB prepared + // statement. Closing either protocol object must remain nil-safe. + h.deletePreparedStatement("begin") + h.deletePortal("begin") + _, preparedExists := h.preparedStatements["begin"] + _, portalExists := h.portals["begin"] + require.False(t, preparedExists) + require.False(t, portalExists) +} diff --git a/pgserver/subscription_handler.go b/pgserver/subscription_handler.go index a07288c2..b873fd41 100644 --- a/pgserver/subscription_handler.go +++ b/pgserver/subscription_handler.go @@ -2,14 +2,16 @@ package pgserver import ( "context" + "errors" "fmt" + "regexp" + "strings" + "github.com/apecloud/myduckserver/adapter" "github.com/apecloud/myduckserver/catalog" "github.com/apecloud/myduckserver/pgserver/logrepl" "github.com/dolthub/go-mysql-server/sql" "github.com/jackc/pglogrepl" - "regexp" - "strings" ) // This file handles SQL statements for managing PostgreSQL subscriptions. It supports: @@ -176,7 +178,7 @@ func (h *ConnectionHandler) executeSubscriptionSQL(subscriptionConfig *Subscript } func (h *ConnectionHandler) executeEnableSubscription(subscriptionConfig *SubscriptionConfig) error { - sqlCtx, err := h.duckHandler.sm.NewContextWithQuery(context.Background(), h.mysqlConn, "") + sqlCtx, err := h.duckHandler.sm.NewContextWithQuery(frontendContext(context.Background()), h.mysqlConn, "") if err != nil { return fmt.Errorf("failed to create context for query: %w", err) } @@ -197,7 +199,7 @@ func (h *ConnectionHandler) executeEnableSubscription(subscriptionConfig *Subscr } func (h *ConnectionHandler) executeDisableSubscription(subscriptionConfig *SubscriptionConfig) error { - sqlCtx, err := h.duckHandler.sm.NewContextWithQuery(context.Background(), h.mysqlConn, "") + sqlCtx, err := h.duckHandler.sm.NewContextWithQuery(frontendContext(context.Background()), h.mysqlConn, "") if err != nil { return fmt.Errorf("failed to create context for query: %w", err) } @@ -218,7 +220,7 @@ func (h *ConnectionHandler) executeDisableSubscription(subscriptionConfig *Subsc } func (h *ConnectionHandler) executeDrop(subscriptionConfig *SubscriptionConfig) error { - sqlCtx, err := h.duckHandler.sm.NewContextWithQuery(context.Background(), h.mysqlConn, "") + sqlCtx, err := h.duckHandler.sm.NewContextWithQuery(frontendContext(context.Background()), h.mysqlConn, "") if err != nil { return fmt.Errorf("failed to create context for query: %w", err) } @@ -239,7 +241,7 @@ func (h *ConnectionHandler) executeDrop(subscriptionConfig *SubscriptionConfig) } func (h *ConnectionHandler) executeCreate(subscriptionConfig *SubscriptionConfig) error { - sqlCtx, err := h.duckHandler.sm.NewContextWithQuery(context.Background(), h.mysqlConn, "") + sqlCtx, err := h.duckHandler.sm.NewContextWithQuery(frontendContext(context.Background()), h.mysqlConn, "") if err != nil { return fmt.Errorf("failed to create context for query: %w", err) } @@ -257,14 +259,10 @@ func (h *ConnectionHandler) executeCreate(subscriptionConfig *SubscriptionConfig return nil } -func (h *ConnectionHandler) doSnapshot(sqlCtx *sql.Context, subscriptionConfig *SubscriptionConfig) (pglogrepl.LSN, error) { +func (h *ConnectionHandler) doSnapshot(sqlCtx *sql.Context, subscriptionConfig *SubscriptionConfig) (lsn pglogrepl.LSN, err error) { // If there is ongoing transcation, commit it - if txn := adapter.TryGetTxn(sqlCtx); txn != nil { - if err := func() error { - defer txn.Rollback() - defer adapter.CloseTxn(sqlCtx) - return txn.Commit() - }(); err != nil { + if _, txn := adapter.TryGetTxnBindingForRelease(sqlCtx); txn != nil { + if err := adapter.FinalizeCommit(sqlCtx, txn); err != nil { return 0, fmt.Errorf("failed to commit current transaction: %w", err) } } @@ -282,7 +280,7 @@ func (h *ConnectionHandler) doSnapshot(sqlCtx *sql.Context, subscriptionConfig * }() var currentLSN string - err := adapter.QueryRowCatalog( + err = adapter.QueryRowCatalog( sqlCtx, fmt.Sprintf("SELECT * FROM postgres_query('%s', 'SELECT pg_current_wal_lsn()')", attachName), ).Scan(¤tLSN) @@ -290,7 +288,7 @@ func (h *ConnectionHandler) doSnapshot(sqlCtx *sql.Context, subscriptionConfig * return 0, fmt.Errorf("failed to query WAL LSN: %w", err) } - lsn, err := pglogrepl.ParseLSN(currentLSN) + lsn, err = pglogrepl.ParseLSN(currentLSN) if err != nil { return 0, fmt.Errorf("failed to parse LSN: %w", err) } @@ -328,6 +326,26 @@ func (h *ConnectionHandler) doSnapshot(sqlCtx *sql.Context, subscriptionConfig * return 0, err } + txn, err := adapter.GetCatalogTxn(sqlCtx, nil) + if err != nil { + return 0, fmt.Errorf("failed to get transaction: %w", err) + } + committed := false + defer func() { + if committed { + return + } + _, _, rollbackErr := adapter.FinalizeRollback(sqlCtx, txn) + if rollbackErr != nil && !adapter.IsTransactionInactiveError(rollbackErr) { + rollbackErr = fmt.Errorf("failed to roll back snapshot transaction: %w", rollbackErr) + if err == nil { + err = rollbackErr + } else { + err = errors.Join(err, rollbackErr) + } + } + }() + // Create all schemas in the target database for _, t := range tables { if _, err := adapter.ExecCatalogInTxn(sqlCtx, `CREATE SCHEMA IF NOT EXISTS `+catalog.QuoteIdentifierANSI(t.schema)); err != nil { @@ -335,13 +353,6 @@ func (h *ConnectionHandler) doSnapshot(sqlCtx *sql.Context, subscriptionConfig * } } - txn, err := adapter.GetCatalogTxn(sqlCtx, nil) - if err != nil { - return 0, fmt.Errorf("failed to get transaction: %w", err) - } - defer txn.Rollback() - defer adapter.CloseTxn(sqlCtx) - for _, t := range tables { if _, err := adapter.ExecCatalogInTxn( sqlCtx, @@ -351,11 +362,15 @@ func (h *ConnectionHandler) doSnapshot(sqlCtx *sql.Context, subscriptionConfig * } } - return lsn, txn.Commit() + if err = adapter.FinalizeCommit(sqlCtx, txn); err != nil { + return 0, fmt.Errorf("failed to commit snapshot transaction: %w", err) + } + committed = true + return lsn, nil } -func (h *ConnectionHandler) doCreateSubscription(sqlCtx *sql.Context, subscriptionConfig *SubscriptionConfig, lsn pglogrepl.LSN) error { - err := logrepl.CreatePublicationIfNotExists(subscriptionConfig.ToDNS(), subscriptionConfig.PublicationName) +func (h *ConnectionHandler) doCreateSubscription(sqlCtx *sql.Context, subscriptionConfig *SubscriptionConfig, lsn pglogrepl.LSN) (err error) { + err = logrepl.CreatePublicationIfNotExists(subscriptionConfig.ToDNS(), subscriptionConfig.PublicationName) if err != nil { return fmt.Errorf("failed to create publication: %w", err) } @@ -364,16 +379,30 @@ func (h *ConnectionHandler) doCreateSubscription(sqlCtx *sql.Context, subscripti if err != nil { return fmt.Errorf("failed to get transaction: %w", err) } - defer tx.Rollback() - defer adapter.CloseTxn(sqlCtx) + committed := false + defer func() { + if committed { + return + } + _, _, rollbackErr := adapter.FinalizeRollback(sqlCtx, tx) + if rollbackErr != nil && !adapter.IsTransactionInactiveError(rollbackErr) { + rollbackErr = fmt.Errorf("failed to roll back subscription transaction: %w", rollbackErr) + if err == nil { + err = rollbackErr + } else { + err = errors.Join(err, rollbackErr) + } + } + }() if err = logrepl.CreateSubscription(sqlCtx, subscriptionConfig.SubscriptionName, subscriptionConfig.ToDNS(), subscriptionConfig.PublicationName, lsn.String(), true); err != nil { return fmt.Errorf("failed to write subscription: %w", err) } - if err = adapter.CommitAndCloseTxn(sqlCtx); err != nil { - return err + if err = adapter.FinalizeCommit(sqlCtx, tx); err != nil { + return fmt.Errorf("failed to commit subscription transaction: %w", err) } + committed = true if err = logrepl.UpdateSubscriptions(sqlCtx); err != nil { return fmt.Errorf("failed to update subscriptions: %w", err) diff --git a/pgserver/table_storage.go b/pgserver/table_storage.go new file mode 100644 index 00000000..fb1c9a17 --- /dev/null +++ b/pgserver/table_storage.go @@ -0,0 +1,148 @@ +package pgserver + +import ( + stdsql "database/sql" + "fmt" + "strings" + + "github.com/apecloud/myduckserver/adapter" + "github.com/apecloud/myduckserver/catalog" + "github.com/apecloud/myduckserver/mycontext" + "github.com/cockroachdb/cockroachdb-parser/pkg/sql/sem/tree" + "github.com/dolthub/go-mysql-server/sql" +) + +// postgresCreateTableStorage normalizes the table-level storage parameters +// while preserving their order (and therefore duplicate declarations). The +// Cockroach parser exposes expressions as String values; the catalog resolver +// validates that selector values are simple literals and contain no secrets or +// paths. +func postgresCreateTableStorage(stmt *tree.CreateTable) (catalog.TableStorageSelection, error) { + if stmt == nil { + return catalog.DefaultTableStorageSelection(), nil + } + options := make([]catalog.TableStorageOption, 0, len(stmt.StorageParams)) + for _, param := range stmt.StorageParams { + var value any + if param.Value != nil { + value = param.Value.String() + } + options = append(options, catalog.TableStorageOption{ + Name: param.Key.String(), + Value: value, + }) + } + return catalog.ResolvePostgresStorageOptions(options) +} + +// postgresCreateTableExecutionQuery returns a query that is safe to hand to +// DuckDB. PostgreSQL's custom WITH (myduck_storage=...) parameter is consumed +// by MyDuck and must not be forwarded as an unknown DuckDB storage option. +// Other PostgreSQL storage parameters remain in the formatted statement and +// retain the existing execution behavior. +func postgresCreateTableExecutionQuery(query string, stmt *tree.CreateTable) (string, catalog.TableStorageSelection, error) { + selection, err := postgresCreateTableStorage(stmt) + if err != nil { + return "", catalog.TableStorageSelection{}, err + } + if stmt == nil || len(stmt.StorageParams) == 0 { + return query, selection, nil + } + + custom := false + params := make(tree.StorageParams, 0, len(stmt.StorageParams)) + for _, param := range stmt.StorageParams { + if strings.EqualFold(strings.TrimSpace(param.Key.String()), catalog.TableStorageOptionName) { + custom = true + continue + } + params = append(params, param) + } + if !custom { + return query, selection, nil + } + + copyStmt := *stmt + if len(params) == 0 { + copyStmt.StorageParams = nil + } else { + copyStmt.StorageParams = params + } + return tree.AsString(©Stmt), selection, nil +} + +// setPostgresCreateTableStorage validates and attaches the request-scoped +// selection used by the catalog boundary. Replication-origin DDL deliberately +// keeps its historical path and metadata behavior unchanged. +func setPostgresCreateTableStorage(ctx *sql.Context, stmt *tree.CreateTable) (catalog.TableStorageSelection, error) { + selection, err := postgresCreateTableStorage(stmt) + if err != nil { + return catalog.TableStorageSelection{}, err + } + if err := catalog.SetTableStorageSelection(ctx, selection); err != nil { + return catalog.TableStorageSelection{}, err + } + return selection, nil +} + +func persistPostgresCreateTableStorage(ctx *sql.Context, stmt *tree.CreateTable, selection catalog.TableStorageSelection, provider *catalog.DatabaseProvider) error { + if stmt == nil || stmt.Persistence == tree.PersistenceTemporary || mycontext.IsReplicationQuery(ctx) { + return nil + } + execer, conn, _, release, err := adapter.GetCatalogExecutionSnapshotWithLease(ctx) + if err != nil { + return err + } + defer release() + return persistPostgresCreateTableStorageWithExecutor(ctx, stmt, selection, provider, execer, conn) +} + +// persistPostgresCreateTableStorageWithExecutor records PostgreSQL table +// storage metadata through the execution snapshot already held by the caller. +// Re-entering adapter.GetCurrentSchema/GetCurrentCatalog or the catalog pool +// from a retained PG lease can block forever behind a pending shutdown writer. +func persistPostgresCreateTableStorageWithExecutor( + ctx *sql.Context, + stmt *tree.CreateTable, + selection catalog.TableStorageSelection, + provider *catalog.DatabaseProvider, + execer adapter.SQLExecutor, + conn *stdsql.Conn, +) error { + if stmt == nil || stmt.Persistence == tree.PersistenceTemporary || mycontext.IsReplicationQuery(ctx) { + return nil + } + + schemaName := stmt.Table.Schema() + if schemaName == "" { + _, schemaName = postgresCurrentNamespaceWithExecutor(ctx, execer, provider) + } + if schemaName == "" { + schemaName = ctx.GetCurrentDatabase() + } + if schemaName == "" { + return fmt.Errorf("cannot determine PostgreSQL schema for table %q", stmt.Table.Table()) + } + catalogName := stmt.Table.Catalog() + if catalogName == "" { + catalogName, _ = postgresCurrentNamespaceWithExecutor(ctx, execer, provider) + } + if catalogName == "" { + return fmt.Errorf("cannot determine PostgreSQL catalog for table %q", stmt.Table.Table()) + } + + db := catalog.NewDatabaseWithProvider(schemaName, catalogName, provider) + return db.RecordTableStorageSelectionWithExecutor(ctx, stmt.Table.Table(), selection, execer, conn) +} + +// postgresCreateTableQueryForPrepare sanitizes only CREATE TABLE statements; +// it is shared by simple and extended PostgreSQL protocol paths so a +// prepared DDL sees the same selector semantics as a simple query. +func postgresCreateTableQueryForPrepare(query string, parsed tree.Statement) (string, error) { + stmt, ok := parsed.(*tree.CreateTable) + if !ok { + return query, nil + } + executionQuery, _, err := postgresCreateTableExecutionQuery(query, stmt) + return executionQuery, err +} diff --git a/pgserver/table_storage_wire_test.go b/pgserver/table_storage_wire_test.go new file mode 100644 index 00000000..dce21859 --- /dev/null +++ b/pgserver/table_storage_wire_test.go @@ -0,0 +1,64 @@ +package pgserver + +import ( + "context" + "fmt" + "os" + "strconv" + "testing" + + "github.com/apecloud/myduckserver/catalog" + "github.com/apecloud/myduckserver/testutil" + "github.com/jackc/pgx/v5" + "github.com/stretchr/testify/require" +) + +func TestTableStorageSelectionWireDefaultOff(t *testing.T) { + originalWorkingDir, err := os.Getwd() + require.NoError(t, err) + defer func() { require.NoError(t, os.Chdir(originalWorkingDir)) }() + + testDir := testutil.CreateTestDir(t) + testEnv := testutil.NewTestEnv() + require.NoError(t, testutil.StartDuckSqlServer(t, testDir, nil, testEnv)) + t.Cleanup(func() { + require.NoError(t, testEnv.MyDuckServer.Close()) + testutil.StopDuckSqlServer(t, testEnv.DuckProcess) + }) + + db := testEnv.MyDuckServer + _, err = db.Exec("CREATE DATABASE storage_wire") + require.NoError(t, err) + _, err = db.Exec("USE storage_wire") + require.NoError(t, err) + + _, err = db.Exec("CREATE TABLE mysql_object (id INT, payload VARCHAR(32)) ENGINE=DUCKLAKE") + require.ErrorContains(t, err, "DuckLake service configuration is disabled") + _, err = db.Exec("CREATE TABLE mysql_default (id INT)") + require.NoError(t, err) + _, err = db.Exec("CREATE TABLE mysql_local (id INT) ENGINE=LOCAL") + require.NoError(t, err) + + ctx := context.Background() + pg, err := pgx.Connect(ctx, "postgresql://postgres@127.0.0.1:"+strconv.Itoa(testEnv.DuckPgPort)+"/postgres") + require.NoError(t, err) + t.Cleanup(func() { _ = pg.Close(ctx) }) + requireStorageKindPG(t, ctx, pg, "storage_wire", "mysql_default", catalog.TableStorageLocal) + requireStorageKindPG(t, ctx, pg, "storage_wire", "mysql_local", catalog.TableStorageLocal) + + _, err = pg.Exec(ctx, "CREATE TABLE storage_wire.pg_object (id INTEGER) WITH (myduck_storage = 'object')") + require.ErrorContains(t, err, "DuckLake service configuration is disabled") + _, err = pg.Exec(ctx, "CREATE TABLE storage_wire.pg_default (id INTEGER)") + require.NoError(t, err) + requireStorageKindPG(t, ctx, pg, "storage_wire", "pg_default", catalog.TableStorageLocal) +} + +func requireStorageKindPG(t *testing.T, ctx context.Context, db *pgx.Conn, catalogName, tableName string, want catalog.TableStorageKind) { + t.Helper() + var raw string + // PostgreSQL exposes MyDuck's physical catalog as `myduck` and maps the + // user database to the schema name. + query := fmt.Sprintf("SELECT comment FROM duckdb_tables() WHERE database_name = 'myduck' AND schema_name = '%s' AND table_name = '%s'", catalogName, tableName) + require.NoError(t, db.QueryRow(ctx, query).Scan(&raw)) + require.Equal(t, want, catalog.DecodeComment[catalog.ExtraTableInfo](raw).Meta.StorageKind()) +} diff --git a/pgserver/transaction.go b/pgserver/transaction.go new file mode 100644 index 00000000..1cfd4474 --- /dev/null +++ b/pgserver/transaction.go @@ -0,0 +1,837 @@ +package pgserver + +import ( + "context" + stdsql "database/sql" + "errors" + "fmt" + "strings" + + "github.com/apecloud/myduckserver/adapter" + "github.com/apecloud/myduckserver/catalog" + "github.com/cockroachdb/cockroachdb-parser/pkg/sql/sem/tree" + "github.com/dolthub/go-mysql-server/sql" + "github.com/jackc/pgx/v5/pgconn" +) + +const postgresFailedTransactionMessage = "current transaction is aborted, commands ignored until end of transaction block" + +// postgresFailedTransactionControlKey marks a protocol context whose COMMIT +// must use the failed-transaction rollback semantics. The marker lives on the +// context rather than in the transaction map so direct handler calls and GMS +// query execution can share the same transaction-control implementation. +type postgresFailedTransactionControlKey struct{} + +// postgresConnectionHandlerKey carries the wire handler through the context +// used by the GMS executor. Transaction-control statements are executed below +// the protocol layer, but BEGIN still needs to promote a protocol-owned +// implicit scope when it appears after ordinary work in the same exchange. +type postgresConnectionHandlerKey struct{} + +func withPostgresConnectionHandler(ctx context.Context, handler *ConnectionHandler) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, postgresConnectionHandlerKey{}, handler) +} + +func postgresConnectionHandler(ctx context.Context) *ConnectionHandler { + if ctx == nil { + return nil + } + handler, _ := ctx.Value(postgresConnectionHandlerKey{}).(*ConnectionHandler) + return handler +} + +func withPostgresFailedTransactionControl(ctx context.Context) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, postgresFailedTransactionControlKey{}, true) +} + +func isPostgresFailedTransactionControl(ctx context.Context) bool { + if ctx == nil { + return false + } + failed, _ := ctx.Value(postgresFailedTransactionControlKey{}).(bool) + return failed +} + +func postgresFailedTransactionError() error { + return &pgconn.PgError{ + Severity: "ERROR", + Code: "25P02", + Message: postgresFailedTransactionMessage, + } +} + +// rejectStatementIfTransactionFailed implements PostgreSQL's sticky failed +// transaction state. COMMIT and ROLLBACK are the only controls supported by +// this handler that can end the block; all other SQL must be rejected before +// any handler, metadata probe, or executor is touched. +func (h *ConnectionHandler) rejectStatementIfTransactionFailed(statement ConvertedStatement) error { + if h == nil || h.readyForQueryStatus() != ReadyForQueryTransactionIndicator_FailedTransactionBlock { + return nil + } + if isEmptyConvertedStatement(statement) { + return nil + } + switch statement.AST.(type) { + case *tree.CommitTransaction, *tree.RollbackTransaction: + return nil + default: + return postgresFailedTransactionError() + } +} + +func isEmptyConvertedStatement(statement ConvertedStatement) bool { + if statement.AST != nil { + return false + } + switch strings.TrimSpace(statement.String) { + case "", ";": + // Empty queries are protocol no-ops and PostgreSQL keeps the failed + // state. A standalone semicolon follows the same rule. + return true + default: + return false + } +} + +// postgresDatabaseDDLName identifies database-file DDL handled by the +// PostgreSQL adapter rather than DuckDB's transactional executor. DuckDB's +// ATTACH/DETACH plus the backing-file mutation cannot be rolled back as part of +// a user transaction, so these statements must be rejected while a block is +// active and must not open an implicit transaction when run alone. +func postgresDatabaseDDLName(statement tree.Statement) (string, bool) { + switch statement.(type) { + case *tree.CreateDatabase: + return "CREATE DATABASE", true + case *tree.DropDatabase: + return "DROP DATABASE", true + default: + return "", false + } +} + +func postgresDatabaseDDLTransactionError(statement tree.Statement) error { + name, ok := postgresDatabaseDDLName(statement) + if !ok { + return nil + } + return &pgconn.PgError{ + Severity: "ERROR", + Code: "25001", + Message: name + " cannot run inside a transaction block", + } +} + +// rejectPostgresDatabaseDDLInTransaction is shared by protocol lifecycle +// gates and the direct PostgreSQL query handler. The handler state catches an +// explicit or message-group transaction even when a test/minimal session does +// not expose a transaction map; the binding check covers direct callers that +// have an active session transaction but no protocol handler marker. +func rejectPostgresDatabaseDDLInTransaction(ctx *sql.Context, statement tree.Statement) error { + if _, ok := postgresDatabaseDDLName(statement); !ok { + return nil + } + if handler := postgresConnectionHandler(ctx); handler != nil { + switch handler.readyForQueryStatus() { + case ReadyForQueryTransactionIndicator_FailedTransactionBlock: + return postgresFailedTransactionError() + case ReadyForQueryTransactionIndicator_TransactionBlock: + return postgresDatabaseDDLTransactionError(statement) + } + if handler.implicitTx != nil { + return postgresDatabaseDDLTransactionError(statement) + } + } + if ctx != nil { + if _, tx := adapter.TryGetTxnBinding(ctx); tx != nil { + return postgresDatabaseDDLTransactionError(statement) + } + } + return nil +} + +// rejectFailedQueryStatements checks a complete simple-query/Parse payload +// before sensitive-SQL gates or protocol shortcuts run. A simple query may +// contain multiple statements, so model the local status as transaction +// controls are encountered: COMMIT/ROLLBACK clear E and statements after the +// control are evaluated as idle. The handler's real state is changed only by +// successful execution. +func (h *ConnectionHandler) rejectFailedQueryStatements(statements []ConvertedStatement) error { + if h == nil { + return nil + } + status := h.readyForQueryStatus() + for _, statement := range statements { + if status != ReadyForQueryTransactionIndicator_FailedTransactionBlock { + continue + } + if isEmptyConvertedStatement(statement) { + continue + } + switch statement.AST.(type) { + case *tree.CommitTransaction, *tree.RollbackTransaction: + status = ReadyForQueryTransactionIndicator_Idle + default: + return postgresFailedTransactionError() + } + } + return nil +} + +// protocolContext returns the ordinary frontend context and carries the +// failed-state marker needed to turn COMMIT into rollback when appropriate. +func (h *ConnectionHandler) protocolContext() context.Context { + ctx := withPostgresConnectionHandler(frontendContext(context.Background()), h) + if h != nil && h.readyForQueryStatus() == ReadyForQueryTransactionIndicator_FailedTransactionBlock { + ctx = withPostgresFailedTransactionControl(ctx) + } + return ctx +} + +// protocolCommandTag reports PostgreSQL's command tag for a protocol +// transaction control statement. COMMIT in a failed block discards the block, +// so the server reports ROLLBACK just as PostgreSQL does. +func (h *ConnectionHandler) protocolCommandTag(statement ConvertedStatement) string { + if h != nil && h.readyForQueryStatus() == ReadyForQueryTransactionIndicator_FailedTransactionBlock { + if _, ok := statement.AST.(*tree.CommitTransaction); ok { + return "ROLLBACK" + } + } + return statement.Tag +} + +// isInactiveTransactionError recognizes only the two driver outcomes that +// prove a transaction is already inactive. Other rollback failures leave the +// physical connection untrusted and must be surfaced to the client. +func isInactiveTransactionError(err error) bool { + if err == nil { + return true + } + return errors.Is(err, stdsql.ErrTxDone) || + strings.Contains(strings.ToLower(err.Error()), "no transaction is active") +} + +func postgresBeginOptions(stmt *tree.BeginTransaction) *stdsql.TxOptions { + if stmt != nil && stmt.Modes.ReadWriteMode == tree.ReadOnly { + return &stdsql.TxOptions{ReadOnly: true} + } + return &stdsql.TxOptions{} +} + +// implicitScopeActive reports whether the handler owns a transaction that is +// still bounded by the current protocol message group. Once BEGIN promotes the +// scope, the same physical binding remains in the pool but Sync must no longer +// commit it automatically. +func (h *ConnectionHandler) implicitScopeActive() bool { + return h != nil && h.implicitTx != nil && !h.implicitTxPromoted +} + +func (h *ConnectionHandler) clearImplicitScope() { + if h == nil { + return + } + h.implicitTx = nil + h.implicitTxCtx = nil + h.implicitTxConn = nil + h.implicitTxPromoted = false +} + +func (h *ConnectionHandler) implicitScopeContext() (*sql.Context, error) { + if h == nil { + return nil, fmt.Errorf("postgres transaction context is unavailable") + } + if h.implicitTxCtx != nil { + return h.implicitTxCtx, nil + } + ctx, err := h.newProtocolSQLContext("") + if err != nil { + return nil, err + } + h.implicitTxCtx = ctx + return ctx, nil +} + +func samePostgresBinding(expectedConn, expectedTx, currentConn, currentTx any) bool { + return expectedConn == currentConn && expectedTx == currentTx +} + +// closeKnownPostgresTransaction removes only the supplied transaction mapping. +// It is deliberately best-effort: a malformed or replacement mapping must not +// be guessed at, and the caller still reports the lifecycle error that caused +// cleanup to be attempted. +func closeKnownPostgresTransaction(ctx *sql.Context, tx *stdsql.Tx) { + if ctx == nil || tx == nil { + return + } + adapter.CloseTxnIf(ctx, tx) +} + +// postgresTransactionBinding returns the complete physical identity for an +// active transaction. A transaction pointer without its owner is not a usable +// binding: remove that exact mapping when the session supports it and fail +// closed so callers cannot fall back to a replacement connection. +func postgresTransactionBinding(ctx *sql.Context, operation string) (*stdsql.Conn, *stdsql.Tx, error) { + if ctx == nil { + return nil, nil, fmt.Errorf("%s: postgres transaction context is unavailable", operation) + } + conn, tx := adapter.TryGetTxnBindingForRelease(ctx) + if tx == nil { + return conn, nil, nil + } + if conn == nil { + closeKnownPostgresTransaction(ctx, tx) + return nil, nil, fmt.Errorf("%s: active postgres transaction has no physical owner", operation) + } + return conn, tx, nil +} + +// invalidateImplicitScope closes the exact physical owner when possible and +// drops the local marker. It never closes the currently observed replacement +// connection: callers pass only the owner captured when the scope was opened. +func (h *ConnectionHandler) invalidateImplicitScope(cause error) error { + if h == nil { + return cause + } + owner := h.implicitTxConn + ctx := h.implicitTxCtx + tx := h.implicitTx + promoted := h.implicitTxPromoted + // Only evict the physical owner if the exact transaction is still bound. + // A replacement transaction may reuse the same *sql.Conn; closing by + // connection identity alone would take that replacement down with it. + ownerStillBound := owner != nil + if ownerStillBound && ctx != nil { + currentConn, currentTx := adapter.TryGetTxnBindingForRelease(ctx) + ownerStillBound = samePostgresBinding(owner, tx, currentConn, currentTx) + } + if ownerStillBound { + var closeErr error + if ctx != nil { + // Keep the transaction identity in the close operation. A replacement + // transaction may reuse the same physical connection; a connection-only + // close would evict that replacement after the binding changed. + closeErr = adapter.CloseConnIfBinding(ctx, owner, tx) + } else { + closeErr = owner.Close() + } + cause = errors.Join(cause, closeErr) + } + // If the owner was already missing, CloseConnIfBinding cannot remove the + // transaction entry. Remove only the captured transaction instead; an + // identity-aware session leaves a replacement transaction untouched. + closeKnownPostgresTransaction(ctx, tx) + remainingTx := false + if ctx != nil { + _, currentTx := adapter.TryGetTxnBindingForRelease(ctx) + remainingTx = currentTx != nil + } + h.clearImplicitScope() + // Any identity violation is unrecoverable at this layer, even when the map + // happens to be empty after cleanup: the driver state of the captured + // transaction was not proven inactive. Keep a sticky failed state rather + // than advertising ReadyForQuery=I over an uncertain lifecycle. + _ = promoted + _ = remainingTx + h.txStatus = ReadyForQueryTransactionIndicator_FailedTransactionBlock + return cause +} + +func (h *ConnectionHandler) validateImplicitBinding() error { + if h == nil || h.implicitTx == nil { + return nil + } + if h.implicitTxConn == nil { + return h.invalidateImplicitScope(fmt.Errorf("postgres implicit transaction owner is missing")) + } + ctx, err := h.implicitScopeContext() + if err != nil { + // The scope cannot be finalized without its saved context. Drop the + // marker and retain a sticky failed state so a later statement cannot + // silently open a different transaction. + h.clearImplicitScope() + h.txStatus = ReadyForQueryTransactionIndicator_FailedTransactionBlock + return err + } + currentConn, currentTx := adapter.TryGetTxnBinding(ctx) + if currentConn != nil && currentTx != nil && + samePostgresBinding(h.implicitTxConn, h.implicitTx, currentConn, currentTx) { + return nil + } + return h.invalidateImplicitScope(fmt.Errorf("postgres implicit transaction binding changed")) +} + +// promoteImplicitPostgresTransaction converts an implicit message-group +// transaction into PostgreSQL's explicit transaction block without opening a +// second physical transaction. The saved context/owner remains the identity +// authority for subsequent controls and cleanup. +func (h *ConnectionHandler) promoteImplicitPostgresTransaction() error { + if h == nil || h.implicitTx == nil { + return nil + } + if h.implicitTxPromoted { + return h.validateImplicitBinding() + } + if err := h.validateImplicitBinding(); err != nil { + return err + } + h.implicitTxPromoted = true + h.txStatus = ReadyForQueryTransactionIndicator_TransactionBlock + return nil +} + +// beginPostgresTransaction binds an explicit PostgreSQL BEGIN to the same +// session transaction map used by MySQL. The direct PostgreSQL handler does +// not run GMS's transaction iterator, so this binding must be established here. +func beginPostgresTransaction(ctx *sql.Context, stmt *tree.BeginTransaction) error { + h := postgresConnectionHandler(ctx) + if h != nil { + status := h.readyForQueryStatus() + // PostgreSQL's failed block is sticky. BEGIN is an ordinary command in + // that state and must not open a fresh physical transaction. + if status == ReadyForQueryTransactionIndicator_FailedTransactionBlock { + return postgresFailedTransactionError() + } + if h.implicitTx != nil { + // BEGIN after ordinary work in this protocol exchange promotes the + // existing physical transaction instead of becoming a nested BEGIN. + if !h.implicitTxPromoted { + return h.promoteImplicitPostgresTransaction() + } + return h.validateImplicitBinding() + } + currentConn, currentTx := adapter.TryGetTxnBinding(ctx) + switch status { + case ReadyForQueryTransactionIndicator_TransactionBlock: + if currentTx == nil { + h.txStatus = ReadyForQueryTransactionIndicator_FailedTransactionBlock + return fmt.Errorf("postgres explicit transaction binding is missing") + } + if currentConn == nil { + closeKnownPostgresTransaction(ctx, currentTx) + h.txStatus = ReadyForQueryTransactionIndicator_FailedTransactionBlock + return fmt.Errorf("postgres explicit transaction owner is missing") + } + case ReadyForQueryTransactionIndicator_Idle: + if currentTx == nil { + break + } + // An unmarked binding while the wire state is idle is stale. Evict + // only its observed owner and refuse to execute on a replacement. + staleErr := fmt.Errorf("postgres transaction binding exists while protocol is idle") + if currentConn != nil { + staleErr = errors.Join(staleErr, adapter.CloseConnIfBinding(ctx, currentConn, currentTx)) + } + closeKnownPostgresTransaction(ctx, currentTx) + if currentConn == nil { + h.txStatus = ReadyForQueryTransactionIndicator_FailedTransactionBlock + staleErr = errors.Join(staleErr, fmt.Errorf("postgres stale transaction owner is missing")) + } + return staleErr + } + } + // PostgreSQL treats BEGIN inside an active transaction block as a successful + // no-op. The wire handler emits the documented WARNING separately; retaining + // the existing binding here keeps the physical transaction untouched. + currentConn, currentTx := adapter.TryGetTxnBinding(ctx) + if currentTx != nil { + if currentConn == nil { + closeKnownPostgresTransaction(ctx, currentTx) + if h != nil { + h.txStatus = ReadyForQueryTransactionIndicator_FailedTransactionBlock + } + return fmt.Errorf("postgres transaction owner is missing") + } + return nil + } + tx, err := adapter.GetTxn(ctx, postgresBeginOptions(stmt)) + if err != nil { + return err + } + if tx == nil { + return fmt.Errorf("postgres BEGIN returned no transaction") + } + owner, boundTx := adapter.TryGetTxnBinding(ctx) + if owner == nil || boundTx != tx { + closeKnownPostgresTransaction(ctx, tx) + if owner != nil { + // The pair check prevents a replacement transaction from being + // evicted when GetTxn raced with another lifecycle operation. + _ = adapter.CloseConnIfBinding(ctx, owner, tx) + } + if h != nil { + h.txStatus = ReadyForQueryTransactionIndicator_FailedTransactionBlock + } + if boundTx != tx { + return fmt.Errorf("postgres explicit transaction binding changed while opening") + } + return fmt.Errorf("postgres explicit transaction owner missing while opening") + } + return nil +} + +// implicitPostgresStatement reports whether a statement participates in the +// transaction that PostgreSQL creates around a simple-query message or an +// extended-query exchange. Protocol controls are handled separately and a few +// session-object commands deliberately stay outside a DuckDB transaction +// because they operate on the connection itself (for example DISCARD ALL). +func implicitPostgresStatement(statement ConvertedStatement) bool { + if statement.AST == nil || isEmptyConvertedStatement(statement) { + return false + } + switch statement.AST.(type) { + case *tree.BeginTransaction, *tree.CommitTransaction, *tree.RollbackTransaction, + *tree.Deallocate, *tree.Discard, *tree.SetVar, *tree.SetSessionCharacteristics, + *tree.ShowVar, *tree.CreateDatabase, *tree.DropDatabase: + return false + case *tree.CopyFrom: + // Arrow COPY currently registers a view through *sql.Conn.Raw and is + // deliberately rejected when a transaction is already bound. Do not + // create an implicit binding that would make every ordinary Arrow COPY + // fail before its loader can start. + return statement.AST.(*tree.CopyFrom).Options.CopyFormat != CopyFormatArrow + case *tree.CopyTo: + return statement.AST.(*tree.CopyTo).Options.CopyFormat != CopyFormatArrow + default: + if statement.Tag == "COPY" { + if _, format, _, ok := ParseCopyFrom(statement.String); ok && format == CopyFormatArrow { + return false + } + if _, format, _, ok := ParseCopyTo(statement.String); ok && format == CopyFormatArrow { + return false + } + } + return true + } +} + +// newProtocolSQLContext creates a context for lifecycle work that is not tied +// to a single client request. In particular, implicit transaction finalizers +// must remain runnable after a request context has been canceled. +func (h *ConnectionHandler) newProtocolSQLContext(query string) (*sql.Context, error) { + if h == nil || h.duckHandler == nil { + return nil, fmt.Errorf("postgres transaction context is unavailable") + } + ctx := withPostgresConnectionHandler(frontendContext(context.Background()), h) + return h.duckHandler.NewContext(ctx, h.mysqlConn, query) +} + +// ensureImplicitPostgresTransaction starts the physical DuckDB transaction +// used by a PostgreSQL implicit transaction scope. The handler keeps the +// pointer so the scope can be committed at the end of a successful simple +// query or Sync, or rolled back after a protocol error. Explicit BEGIN owns its +// binding directly and therefore never sets implicitTx. +func (h *ConnectionHandler) ensureImplicitPostgresTransaction(statement ConvertedStatement) error { + if h == nil { + return nil + } + if _, databaseDDL := postgresDatabaseDDLName(statement.AST); databaseDDL { + // Database-file DDL is deliberately outside the implicit transaction + // scope. Check the protocol marker first so minimal/direct handlers do + // not need to manufacture a SQL context merely to reject an active block. + switch h.readyForQueryStatus() { + case ReadyForQueryTransactionIndicator_FailedTransactionBlock: + return postgresFailedTransactionError() + case ReadyForQueryTransactionIndicator_TransactionBlock: + return postgresDatabaseDDLTransactionError(statement.AST) + } + if h.implicitTx != nil { + return postgresDatabaseDDLTransactionError(statement.AST) + } + if h.duckHandler == nil || h.mysqlConn == nil { + return nil + } + ctx, err := h.newProtocolSQLContext(statement.String) + if err != nil { + return err + } + return rejectPostgresDatabaseDDLInTransaction(ctx, statement.AST) + } + if !implicitPostgresStatement(statement) { + return nil + } + if h.implicitTx != nil { + return h.validateImplicitBinding() + } + if h.readyForQueryStatus() == ReadyForQueryTransactionIndicator_FailedTransactionBlock { + return postgresFailedTransactionError() + } + ctx, err := h.newProtocolSQLContext(statement.String) + if err != nil { + return err + } + if err := rejectPostgresDatabaseDDLInTransaction(ctx, statement.AST); err != nil { + return err + } + status := h.readyForQueryStatus() + currentConn, existing := adapter.TryGetTxnBinding(ctx) + if status == ReadyForQueryTransactionIndicator_TransactionBlock { + if existing == nil { + h.txStatus = ReadyForQueryTransactionIndicator_FailedTransactionBlock + return fmt.Errorf("postgres explicit transaction binding is missing") + } + if currentConn == nil { + closeKnownPostgresTransaction(ctx, existing) + h.txStatus = ReadyForQueryTransactionIndicator_FailedTransactionBlock + return fmt.Errorf("postgres explicit transaction owner is missing") + } + // The explicit transaction is owned by the pool; do not claim it as + // implicit or let a later Sync commit it. + return nil + } + if existing != nil { + staleErr := fmt.Errorf("postgres transaction binding exists while protocol is idle") + if currentConn != nil { + staleErr = errors.Join(staleErr, adapter.CloseConnIfBinding(ctx, currentConn, existing)) + } + closeKnownPostgresTransaction(ctx, existing) + if currentConn == nil { + h.txStatus = ReadyForQueryTransactionIndicator_FailedTransactionBlock + staleErr = errors.Join(staleErr, fmt.Errorf("postgres stale transaction owner is missing")) + } + return staleErr + } + tx, err := adapter.GetTxn(ctx, &stdsql.TxOptions{}) + if err != nil { + return err + } + owner, boundTx := adapter.TryGetTxnBinding(ctx) + if boundTx != tx || owner == nil { + closeKnownPostgresTransaction(ctx, tx) + if owner != nil { + // The strict pair check prevents this cleanup from evicting a + // replacement transaction that won the race after GetTxn returned. + _ = adapter.CloseConnIfBinding(ctx, owner, tx) + } + if boundTx != tx { + h.txStatus = ReadyForQueryTransactionIndicator_FailedTransactionBlock + return fmt.Errorf("postgres implicit transaction binding changed while opening") + } + h.txStatus = ReadyForQueryTransactionIndicator_FailedTransactionBlock + return fmt.Errorf("postgres implicit transaction owner missing while opening") + } + h.implicitTx = tx + h.implicitTxCtx = ctx + h.implicitTxConn = owner + h.implicitTxPromoted = false + return nil +} + +// finalizeImplicitPostgresTransaction closes the handler-owned implicit scope +// and always returns the protocol to Idle. Rollback uses the same provider +// cleanup hook as an explicit PostgreSQL ROLLBACK, preserving the physical +// connection and making cleanup errors visible to the caller. +func (h *ConnectionHandler) finalizeImplicitPostgresTransaction(commit bool) error { + if h == nil || !h.implicitScopeActive() { + return nil + } + tx := h.implicitTx + ctx, err := h.implicitScopeContext() + if err != nil { + closeKnownPostgresTransaction(h.implicitTxCtx, tx) + h.clearImplicitScope() + h.txStatus = ReadyForQueryTransactionIndicator_FailedTransactionBlock + return err + } + if h.implicitTxConn == nil { + return h.invalidateImplicitScope(fmt.Errorf("postgres implicit transaction owner is missing before finalization")) + } + currentConn, current := adapter.TryGetTxnBindingForRelease(ctx) + if !samePostgresBinding(h.implicitTxConn, tx, currentConn, current) { + return h.invalidateImplicitScope(fmt.Errorf("postgres implicit transaction binding changed before finalization")) + } + if commit { + err = adapter.FinalizeCommit(ctx, tx) + } else { + var provider postgresRollbackOrphanCleaner + if h.duckHandler != nil { + provider = h.duckHandler.GetCatalogProvider() + } + if provider != nil { + err = rollbackPostgresTransactionControl(ctx, &tree.RollbackTransaction{}, provider) + } else { + err = rollbackPostgresTransactionControl(ctx, &tree.RollbackTransaction{}) + } + } + // A successful finalizer removes the exact mapping. The transaction pointer + // is the decisive state signal: afterTx == nil means the expected block is + // gone (even if the generic connection was replaced), while a non-nil afterTx + // means a transaction is still active and the protocol must not advertise I. + afterConn, afterTx := adapter.TryGetTxnBindingForRelease(ctx) + if afterTx == nil { + h.clearImplicitScope() + h.clearPortals() + h.txStatus = ReadyForQueryTransactionIndicator_Idle + return err + } + if err == nil { + // A finalizer claiming success while leaving the exact transaction + // mapped is inconsistent with the lifecycle contract. Evict only the + // still-matching owner and surface an explicit failure instead of + // advertising ReadyForQuery=I over an active transaction. + err = fmt.Errorf("postgres implicit transaction remained bound after finalization") + } + if h.implicitTxConn != nil { + err = errors.Join(err, adapter.CloseConnIfBinding(ctx, h.implicitTxConn, tx)) + afterConn, afterTx = adapter.TryGetTxnBindingForRelease(ctx) + } + closeKnownPostgresTransaction(ctx, tx) + afterConn, afterTx = adapter.TryGetTxnBindingForRelease(ctx) + if afterTx == nil { + h.clearImplicitScope() + h.clearPortals() + h.txStatus = ReadyForQueryTransactionIndicator_Idle + return err + } + // A different non-nil transaction is a replacement, not proof that the old + // scope finalized. Drop the stale local marker but retain a sticky failed + // state so later statements cannot run on that replacement silently. + if !samePostgresBinding(h.implicitTxConn, tx, afterConn, afterTx) { + err = errors.Join(err, fmt.Errorf("postgres implicit transaction binding replaced before finalization")) + } + h.clearImplicitScope() + h.clearPortals() + h.txStatus = ReadyForQueryTransactionIndicator_FailedTransactionBlock + return err +} + +// finishImplicitPostgresError rolls back a protocol-owned implicit scope and +// returns the original error together with any rollback/cleanup error. An +// explicit transaction has no implicit marker and therefore retains the +// normal sticky E state. +func (h *ConnectionHandler) finishImplicitPostgresError(err error) error { + if err == nil { + return nil + } + // Extended COPY FROM can finish its loader before the protocol receives the + // following Sync. Its snapshot/operation lease is retained in + // pendingCopyRelease so a successful group commits before releasing it. If a + // later pipelined message fails before Sync, however, rollback must not wait + // on that same lease: release the completed COPY first, then enter the + // provider rollback reservation. The callback is idempotent and is also safe + // on paths that already released the lease. + if h != nil { + h.releasePendingCopyLease() + } + if h != nil && h.implicitScopeActive() { + return errors.Join(err, h.finalizeImplicitPostgresTransaction(false)) + } + if h != nil { + h.markTransactionError() + } + return err +} + +func commitPostgresTransaction(ctx *sql.Context) error { + _, tx, err := postgresTransactionBinding(ctx, "postgres COMMIT") + if err != nil { + return err + } + if tx == nil { + return nil + } + return adapter.FinalizeCommit(ctx, tx) +} + +func rollbackPostgresTransactionWithOwner(ctx *sql.Context) (*stdsql.Conn, bool, error) { + conn, tx, err := postgresTransactionBinding(ctx, "postgres ROLLBACK") + if err != nil { + return nil, false, err + } + if tx == nil { + return conn, true, nil + } + return adapter.FinalizeRollback(ctx, tx) +} + +func rollbackPostgresTransaction(ctx *sql.Context) error { + _, _, err := rollbackPostgresTransactionWithOwner(ctx) + return err +} + +// rollbackPostgresTransactionControl performs the protocol rollback path for +// both ROLLBACK and COMMIT issued while a block is failed. The latter is +// deliberately represented as a synthetic RollbackTransaction for cleanup: +// PostgreSQL discards the failed block and reports a rollback-equivalent +// command tag rather than attempting to commit an already-aborted driver tx. +func rollbackPostgresTransactionControl( + ctx *sql.Context, + transaction *tree.RollbackTransaction, + providers ...postgresRollbackOrphanCleaner, +) error { + _, tx, bindingErr := postgresTransactionBinding(ctx, "postgres rollback cleanup") + if bindingErr != nil { + return bindingErr + } + var releaseCleanup func() + var reservationErr error + if len(providers) > 0 { + releaseCleanup, reservationErr = beginPostgresRollbackCleanupReservationWithError(ctx, providers[0], tx) + if reservationErr == nil { + defer releaseCleanup() + } else { + // A poisoned generation must not block the driver's rollback. Skip + // maintenance, but preserve the admission error for the protocol caller. + releaseCleanup = nil + } + } + if tx == nil { + if reservationErr != nil { + return reservationErr + } + if len(providers) == 0 { + return nil + } + // A GMS-only rollback has no physical *sql.Tx, but the reservation above + // still owns the provider cleanup barrier. Mark that ownership explicitly + // before entering the helper so it does not try to acquire the same + // provider mutex a second time. Reuse the captured idle connection when one + // exists; this keeps DuckLake attachment/session state stable. + cleanupBase := catalog.WithDuckLakeCleanupLease(context.WithoutCancel(ctx)) + cleanupCtx := ctx.WithContext(cleanupBase) + if conn, _ := adapter.TryGetTxnBindingForRelease(ctx); conn != nil { + return cleanupPostgresRollbackOrphans(cleanupCtx, transaction, providers[0], conn) + } + return cleanupPostgresRollbackOrphans(cleanupCtx, transaction, providers[0]) + } + var cleanup func(*stdsql.Conn) error + if len(providers) > 0 && reservationErr == nil { + provider := providers[0] + cleanup = func(conn *stdsql.Conn) error { + // The pool finalizer holds the session lock while this callback + // runs. Mark the exact owner so provider checks do not re-enter + // that lock or observe a replacement transaction. + base := context.WithoutCancel(ctx) + base = catalog.WithRollbackCleanupOwner(base, conn) + base = catalog.WithDuckLakeCleanupLease(base) + cleanupCtx := ctx.WithContext(base) + return cleanupPostgresRollbackOrphans(cleanupCtx, transaction, provider, conn) + } + } + _, rollbackErr := adapter.FinalizeRollbackWithCleanup(ctx, tx, cleanup) + return errors.Join(reservationErr, rollbackErr) +} + +// postgresTransactionControl handles transaction statements that bypass the +// GMS transaction iterator. A PostgreSQL rollback owns its connection through +// finalization, then performs any eligible DuckLake cleanup on that same +// connection before returning to the wire layer. +func postgresTransactionControl(ctx *sql.Context, stmt tree.Statement, providers ...postgresRollbackOrphanCleaner) (handled bool, err error) { + switch transaction := stmt.(type) { + case *tree.BeginTransaction: + return true, beginPostgresTransaction(ctx, transaction) + case *tree.CommitTransaction: + if isPostgresFailedTransactionControl(ctx) { + return true, rollbackPostgresTransactionControl(ctx, &tree.RollbackTransaction{}, providers...) + } + return true, commitPostgresTransaction(ctx) + case *tree.RollbackTransaction: + return true, rollbackPostgresTransactionControl(ctx, transaction, providers...) + default: + return false, nil + } +} diff --git a/pgserver/transaction_test.go b/pgserver/transaction_test.go new file mode 100644 index 00000000..94a215fe --- /dev/null +++ b/pgserver/transaction_test.go @@ -0,0 +1,388 @@ +package pgserver + +import ( + "context" + stdsql "database/sql" + "testing" + + "github.com/apecloud/myduckserver/adapter" + "github.com/apecloud/myduckserver/backend" + "github.com/apecloud/myduckserver/catalog" + "github.com/apecloud/myduckserver/mycontext" + "github.com/cockroachdb/cockroachdb-parser/pkg/sql/parser" + "github.com/cockroachdb/cockroachdb-parser/pkg/sql/sem/tree" + "github.com/dolthub/go-mysql-server/memory" + "github.com/dolthub/go-mysql-server/sql" + "github.com/jackc/pgx/v5/pgconn" + "github.com/stretchr/testify/require" +) + +// replacingTransactionSession models a finalizer that reports success while a +// different transaction has become current. The protocol finalizer must treat +// that as a replacement/corruption state, not as proof that its old implicit +// scope finalized. +type replacingTransactionSession struct { + *memory.Session + conn *stdsql.Conn + tx *stdsql.Tx + replacement *stdsql.Tx + replaceConn *stdsql.Conn + commitCalled bool + rollbackCalled bool + getTxnCalls int +} + +var _ adapter.ConnectionHolder = (*replacingTransactionSession)(nil) +var _ adapter.TransactionBindingHolder = (*replacingTransactionSession)(nil) +var _ adapter.TransactionFinalizer = (*replacingTransactionSession)(nil) + +func (s *replacingTransactionSession) GetConn(context.Context) (*stdsql.Conn, error) { + return s.conn, nil +} + +func (s *replacingTransactionSession) GetTxn(context.Context, *stdsql.TxOptions) (*stdsql.Tx, error) { + s.getTxnCalls++ + return s.tx, nil +} + +func (s *replacingTransactionSession) GetCatalogConn(context.Context) (*stdsql.Conn, error) { + return s.conn, nil +} + +func (s *replacingTransactionSession) GetCatalogTxn(context.Context, *stdsql.TxOptions) (*stdsql.Tx, error) { + return s.tx, nil +} + +func (s *replacingTransactionSession) TryGetTxn() *stdsql.Tx { return s.tx } + +func (s *replacingTransactionSession) GetTxnBinding() (*stdsql.Conn, *stdsql.Tx) { + return s.conn, s.tx +} + +func (s *replacingTransactionSession) GetCurrentCatalog() string { return "memory" } + +func (s *replacingTransactionSession) GetCurrentSchema() string { return "main" } + +func (s *replacingTransactionSession) CloseTxn() { s.tx = nil } + +func (s *replacingTransactionSession) CloseConn() { s.conn = nil } + +func (s *replacingTransactionSession) CommitTxn(*stdsql.Tx) error { + s.commitCalled = true + s.conn = s.replaceConn + s.tx = s.replacement + return nil +} + +func (s *replacingTransactionSession) RollbackTxn(*stdsql.Tx) (*stdsql.Conn, bool, error) { + s.rollbackCalled = true + s.tx = nil + return s.conn, true, nil +} + +func TestPostgresTransactionControlBindsAndFinalizesDuckDBTransaction(t *testing.T) { + provider := catalog.NewInMemoryDBProvider() + t.Cleanup(func() { require.NoError(t, provider.Close()) }) + _, err := provider.Storage().ExecContext(context.Background(), "CREATE TABLE pg_tx_control (id INTEGER)") + require.NoError(t, err) + + session := backend.NewSession(memory.NewSession(sql.NewBaseSession(), provider), provider) + ctx := sql.NewContext(mycontext.WithFrontendQuery(context.Background()), sql.WithSession(session)) + + handled, err := postgresTransactionControl(ctx, &tree.BeginTransaction{}, provider) + require.True(t, handled) + require.NoError(t, err) + _, tx := adapter.TryGetTxnBinding(ctx) + require.NotNil(t, tx) + _, err = tx.ExecContext(ctx, "INSERT INTO pg_tx_control VALUES (1)") + require.NoError(t, err) + + handled, err = postgresTransactionControl(ctx, &tree.RollbackTransaction{}, provider) + require.True(t, handled) + require.NoError(t, err) + require.Nil(t, session.TryGetTxn()) + + var count int + require.NoError(t, provider.Storage().QueryRowContext(context.Background(), "SELECT count(*) FROM pg_tx_control").Scan(&count)) + require.Zero(t, count) +} + +func TestPostgresTransactionControlCommitLeavesCommittedRows(t *testing.T) { + provider := catalog.NewInMemoryDBProvider() + t.Cleanup(func() { require.NoError(t, provider.Close()) }) + _, err := provider.Storage().ExecContext(context.Background(), "CREATE TABLE pg_tx_commit (id INTEGER)") + require.NoError(t, err) + + session := backend.NewSession(memory.NewSession(sql.NewBaseSession(), provider), provider) + ctx := sql.NewContext(mycontext.WithFrontendQuery(context.Background()), sql.WithSession(session)) + + err = beginPostgresTransaction(ctx, &tree.BeginTransaction{}) + require.NoError(t, err) + tx := session.TryGetTxn() + require.NotNil(t, tx) + _, err = tx.ExecContext(ctx, "INSERT INTO pg_tx_commit VALUES (1)") + require.NoError(t, err) + require.NoError(t, commitPostgresTransaction(ctx)) + require.Nil(t, session.TryGetTxn()) + + var count int + require.NoError(t, provider.Storage().QueryRowContext(context.Background(), "SELECT count(*) FROM pg_tx_commit").Scan(&count)) + require.Equal(t, 1, count) +} + +func TestPostgresExplicitBeginRetainsOwnerAcrossStatementContexts(t *testing.T) { + provider := catalog.NewInMemoryDBProvider() + t.Cleanup(func() { require.NoError(t, provider.Close()) }) + _, err := provider.Storage().ExecContext(context.Background(), "CREATE TABLE pg_tx_owner (id INTEGER)") + require.NoError(t, err) + + session := backend.NewSession(memory.NewSession(sql.NewBaseSession(), provider), provider) + beginCtx := sql.NewContext( + withPostgresConnectionHandler(mycontext.WithFrontendQuery(context.Background()), &ConnectionHandler{}), + sql.WithSession(session), + ) + require.NoError(t, beginPostgresTransaction(beginCtx, &tree.BeginTransaction{})) + owner, tx := adapter.TryGetTxnBinding(beginCtx) + require.NotNil(t, owner) + require.NotNil(t, tx) + + // Each context represents a separate frontend statement/request. The + // explicit BEGIN must keep both the transaction and its physical owner + // stable instead of reacquiring the session connection for the next query. + statementCtx := sql.NewContext(mycontext.WithFrontendQuery(context.Background()), sql.WithSession(session)) + gotOwner, gotTx := adapter.TryGetTxnBinding(statementCtx) + require.Same(t, owner, gotOwner) + require.Same(t, tx, gotTx) + execer, executionOwner, executionTx, err := adapter.GetExecutionSnapshot(statementCtx) + require.NoError(t, err) + require.Same(t, owner, executionOwner) + require.Same(t, tx, executionTx) + _, err = execer.ExecContext(statementCtx, "INSERT INTO pg_tx_owner VALUES (1)") + require.NoError(t, err) + + selectCtx := sql.NewContext(mycontext.WithFrontendQuery(context.Background()), sql.WithSession(session)) + execer, executionOwner, executionTx, err = adapter.GetExecutionSnapshot(selectCtx) + require.NoError(t, err) + require.Same(t, owner, executionOwner) + require.Same(t, tx, executionTx) + var count int + require.NoError(t, execer.QueryRowContext(selectCtx, "SELECT count(*) FROM pg_tx_owner").Scan(&count)) + require.Equal(t, 1, count) + + rollbackCtx := sql.NewContext(mycontext.WithFrontendQuery(context.Background()), sql.WithSession(session)) + handled, err := postgresTransactionControl(rollbackCtx, &tree.RollbackTransaction{}, provider) + require.True(t, handled) + require.NoError(t, err) + require.Nil(t, session.TryGetTxn()) + remainingOwner, remainingTx := adapter.TryGetTxnBinding(selectCtx) + require.Same(t, owner, remainingOwner) + require.Nil(t, remainingTx) + + require.NoError(t, provider.Storage().QueryRowContext(context.Background(), "SELECT count(*) FROM pg_tx_owner").Scan(&count)) + require.Zero(t, count) +} + +func TestFinalizeImplicitTransactionRejectsReplacementBinding(t *testing.T) { + oldConn := new(stdsql.Conn) + oldTx := new(stdsql.Tx) + replacementConn := new(stdsql.Conn) + replacementTx := new(stdsql.Tx) + session := &replacingTransactionSession{ + Session: memory.NewSession(sql.NewBaseSession(), nil), + conn: oldConn, + tx: oldTx, + replaceConn: replacementConn, + replacement: replacementTx, + } + ctx := sql.NewContext( + withPostgresConnectionHandler(mycontext.WithFrontendQuery(context.Background()), &ConnectionHandler{}), + sql.WithSession(session), + ) + h := &ConnectionHandler{ + implicitTx: oldTx, + implicitTxCtx: ctx, + implicitTxConn: oldConn, + txStatus: ReadyForQueryTransactionIndicator_Idle, + } + + err := h.finalizeImplicitPostgresTransaction(true) + require.ErrorContains(t, err, "binding replaced") + require.True(t, session.commitCalled) + require.Same(t, replacementTx, session.tx) + require.Same(t, replacementConn, session.conn) + require.Nil(t, h.implicitTx) + require.Equal(t, ReadyForQueryTransactionIndicator_FailedTransactionBlock, h.readyForQueryStatus()) +} + +func TestPostgresTransactionControlNestedBeginIsNoOp(t *testing.T) { + provider := catalog.NewInMemoryDBProvider() + t.Cleanup(func() { require.NoError(t, provider.Close()) }) + session := backend.NewSession(memory.NewSession(sql.NewBaseSession(), provider), provider) + ctx := sql.NewContext(mycontext.WithFrontendQuery(context.Background()), sql.WithSession(session)) + + require.NoError(t, beginPostgresTransaction(ctx, &tree.BeginTransaction{})) + first := session.TryGetTxn() + require.NotNil(t, first) + + require.NoError(t, beginPostgresTransaction(ctx, &tree.BeginTransaction{})) + require.Same(t, first, session.TryGetTxn()) + require.NoError(t, rollbackPostgresTransaction(ctx)) +} + +func TestPostgresFailedCommitRollsBackAndCleansBinding(t *testing.T) { + provider := catalog.NewInMemoryDBProvider() + t.Cleanup(func() { require.NoError(t, provider.Close()) }) + _, err := provider.Storage().ExecContext(context.Background(), "CREATE TABLE pg_failed_commit (id INTEGER)") + require.NoError(t, err) + + session := backend.NewSession(memory.NewSession(sql.NewBaseSession(), provider), provider) + ctx := sql.NewContext(mycontext.WithFrontendQuery(context.Background()), sql.WithSession(session)) + require.NoError(t, beginPostgresTransaction(ctx, &tree.BeginTransaction{})) + tx := session.TryGetTxn() + require.NotNil(t, tx) + _, err = tx.ExecContext(ctx, "INSERT INTO pg_failed_commit VALUES (1)") + require.NoError(t, err) + // A statement error puts the PostgreSQL protocol in E. COMMIT must then + // discard the underlying transaction instead of attempting tx.Commit(). + _, err = tx.ExecContext(ctx, "INSERT INTO pg_failed_commit (missing) VALUES (2)") + require.Error(t, err) + + failedCtx := ctx.WithContext(withPostgresFailedTransactionControl(ctx)) + handled, err := postgresTransactionControl(failedCtx, &tree.CommitTransaction{}, provider) + require.True(t, handled) + require.NoError(t, err) + require.Nil(t, session.TryGetTxn()) + + var count int + require.NoError(t, provider.Storage().QueryRowContext(context.Background(), "SELECT count(*) FROM pg_failed_commit").Scan(&count)) + require.Zero(t, count) +} + +func TestPostgresBeginInFailedBlockDoesNotOpenTransaction(t *testing.T) { + failedTx := new(stdsql.Tx) + session := &replacingTransactionSession{ + Session: memory.NewSession(sql.NewBaseSession(), nil), + conn: new(stdsql.Conn), + tx: failedTx, + } + h := &ConnectionHandler{txStatus: ReadyForQueryTransactionIndicator_FailedTransactionBlock} + ctx := sql.NewContext( + withPostgresConnectionHandler(context.Background(), h), + sql.WithSession(session), + ) + + err := beginPostgresTransaction(ctx, &tree.BeginTransaction{}) + var pgErr *pgconn.PgError + require.ErrorAs(t, err, &pgErr) + require.Equal(t, "25P02", pgErr.Code) + require.Zero(t, session.getTxnCalls, "failed BEGIN must not call GetTxn") + require.Same(t, failedTx, session.tx, "failed BEGIN must leave the existing binding untouched") + require.Equal(t, ReadyForQueryTransactionIndicator_FailedTransactionBlock, h.readyForQueryStatus()) +} + +func TestPostgresBeginWithOwnerlessExplicitBindingFailsClosed(t *testing.T) { + activeTx := new(stdsql.Tx) + session := &replacingTransactionSession{ + Session: memory.NewSession(sql.NewBaseSession(), nil), + tx: activeTx, + } + h := &ConnectionHandler{txStatus: ReadyForQueryTransactionIndicator_TransactionBlock} + ctx := sql.NewContext( + withPostgresConnectionHandler(context.Background(), h), + sql.WithSession(session), + ) + + err := beginPostgresTransaction(ctx, &tree.BeginTransaction{}) + require.ErrorContains(t, err, "explicit transaction owner is missing") + require.Zero(t, session.getTxnCalls) + require.Nil(t, session.tx, "the ownerless mapping must be removed by identity") + require.Equal(t, ReadyForQueryTransactionIndicator_FailedTransactionBlock, h.readyForQueryStatus()) +} + +func TestPostgresDatabaseDDLIsOutsideImplicitTransactions(t *testing.T) { + for _, query := range []string{ + "CREATE DATABASE ddl_scope_test", + "DROP DATABASE ddl_scope_test", + } { + statements, err := parser.Parse(query) + require.NoError(t, err) + require.Len(t, statements, 1) + converted := ConvertedStatement{AST: statements[0].AST, String: query} + require.False(t, implicitPostgresStatement(converted), query) + + idle := &ConnectionHandler{txStatus: ReadyForQueryTransactionIndicator_Idle} + require.NoError(t, idle.ensureImplicitPostgresTransaction(converted), query) + + active := &ConnectionHandler{txStatus: ReadyForQueryTransactionIndicator_TransactionBlock} + err = active.ensureImplicitPostgresTransaction(converted) + var pgErr *pgconn.PgError + require.ErrorAs(t, err, &pgErr, query) + require.Equal(t, "25001", pgErr.Code, query) + name, ok := postgresDatabaseDDLName(statements[0].AST) + require.True(t, ok) + require.Equal(t, name+" cannot run inside a transaction block", pgErr.Message, query) + } +} + +func TestPostgresDatabaseDDLRejectsFailedBlockWithStickyError(t *testing.T) { + statements, err := parser.Parse("CREATE DATABASE ddl_failed_test") + require.NoError(t, err) + ctx := sql.NewContext( + withPostgresConnectionHandler(context.Background(), &ConnectionHandler{ + txStatus: ReadyForQueryTransactionIndicator_FailedTransactionBlock, + }), + ) + err = rejectPostgresDatabaseDDLInTransaction(ctx, statements[0].AST) + var pgErr *pgconn.PgError + require.ErrorAs(t, err, &pgErr) + require.Equal(t, "25P02", pgErr.Code) +} + +func TestPostgresImplicitBindingWithoutOwnerRemainsFailed(t *testing.T) { + activeTx := new(stdsql.Tx) + session := &replacingTransactionSession{ + Session: memory.NewSession(sql.NewBaseSession(), nil), + tx: activeTx, + } + h := &ConnectionHandler{txStatus: ReadyForQueryTransactionIndicator_Idle, implicitTx: activeTx} + ctx := sql.NewContext( + withPostgresConnectionHandler(context.Background(), h), + sql.WithSession(session), + ) + h.implicitTxCtx = ctx + + err := h.validateImplicitBinding() + require.ErrorContains(t, err, "implicit transaction owner is missing") + require.Nil(t, h.implicitTx) + require.Nil(t, session.tx) + require.Equal(t, ReadyForQueryTransactionIndicator_FailedTransactionBlock, h.readyForQueryStatus()) +} + +func TestFinalizeImplicitRollbackWithoutDuckHandlerDoesNotPanic(t *testing.T) { + owner := new(stdsql.Conn) + activeTx := new(stdsql.Tx) + session := &replacingTransactionSession{ + Session: memory.NewSession(sql.NewBaseSession(), nil), + conn: owner, + tx: activeTx, + } + h := &ConnectionHandler{ + implicitTx: activeTx, + implicitTxCtx: nil, + implicitTxConn: owner, + duckHandler: nil, + } + ctx := sql.NewContext( + withPostgresConnectionHandler(context.Background(), h), + sql.WithSession(session), + ) + h.implicitTxCtx = ctx + + require.NotPanics(t, func() { + require.NoError(t, h.finalizeImplicitPostgresTransaction(false)) + }) + require.True(t, session.rollbackCalled) + require.Nil(t, h.implicitTx) + require.Nil(t, session.tx) + require.Equal(t, ReadyForQueryTransactionIndicator_Idle, h.readyForQueryStatus()) +}