From 4813d690f29c33b86a3a3e227650e1f31d35f0dc Mon Sep 17 00:00:00 2001 From: viggy28 Date: Sun, 30 Aug 2026 23:17:52 -0700 Subject: [PATCH 1/2] feat: query DuckLake through attached catalog --- cmd/streambed/main.go | 15 ++- internal/ducklake/writer.go | 6 + internal/server/ducklake_test.go | 141 +++++++++++++++++------ internal/server/server.go | 184 +++++++++++++++++-------------- 4 files changed, 229 insertions(+), 117 deletions(-) diff --git a/cmd/streambed/main.go b/cmd/streambed/main.go index f2f286e..19cd937 100644 --- a/cmd/streambed/main.go +++ b/cmd/streambed/main.go @@ -3,6 +3,7 @@ package main import ( "bufio" "context" + "database/sql" "fmt" "log/slog" "os" @@ -326,8 +327,13 @@ func runSync(cmd *cobra.Command, args []string) error { } // Start query server if --query-addr is set + var querySrv *server.Server if cfg.QueryAddr != "" { - querySrv, err := server.NewServer(server.ServerConfig{ + var duckLakeDB *sql.DB + if duckWriter != nil { + duckLakeDB = duckWriter.DB() + } + querySrv, err = server.NewServer(server.ServerConfig{ ListenAddr: cfg.QueryAddr, S3Bucket: cfg.S3Bucket, S3Prefix: cfg.S3Prefix, @@ -337,6 +343,7 @@ func runSync(cmd *cobra.Command, args []string) error { DuckLakeCatalog: cfg.DuckLakeCatalog, DuckLakeCatalogStore: cfg.DuckLakeCatalogStore, DuckLakeDataPath: cfg.EffectiveDuckLakeDataPath(), + DuckLakeDB: duckLakeDB, }, s3Client, logger) if err != nil { return fmt.Errorf("create query server: %w", err) @@ -526,6 +533,12 @@ func runSync(cmd *cobra.Command, args []string) error { logger.Error("reconnect: writer setup failed", "error", err) continue } + if querySrv != nil && duckWriter != nil { + if err := querySrv.SetDuckLakeDB(duckWriter.DB()); err != nil { + logger.Error("reconnect: query server setup failed", "error", err) + continue + } + } p = pipeline.New(pgConn, cfg.SlotName, pubName, startLSN, cfg.ExcludeTables, logger, stateStore, tableFlushLSN, writer, cfg.FlushInterval, metaQuerier) diff --git a/internal/ducklake/writer.go b/internal/ducklake/writer.go index 165b376..3b3045f 100644 --- a/internal/ducklake/writer.go +++ b/internal/ducklake/writer.go @@ -176,6 +176,12 @@ func catalogName(cfg Config) string { return cfg.CatalogName } +// DB returns the configured DuckDB handle so the in-process query server can +// share the same DuckLake attachment and observe writer commits immediately. +func (w *Writer) DB() *sql.DB { + return w.db +} + func (w *Writer) Close() error { return w.db.Close() } diff --git a/internal/server/ducklake_test.go b/internal/server/ducklake_test.go index 36341b7..7310dc2 100644 --- a/internal/server/ducklake_test.go +++ b/internal/server/ducklake_test.go @@ -2,6 +2,7 @@ package server import ( "context" + "fmt" "log/slog" "os" "path/filepath" @@ -12,34 +13,88 @@ import ( "github.com/viggy28/streambed/internal/wal" ) -func TestDuckLakeServerRegistersViews(t *testing.T) { +func TestDuckLakeServerUsesAttachedCatalogDirectly(t *testing.T) { ctx := context.Background() dir := t.TempDir() - catalogPath := filepath.Join(dir, "catalog.sqlite") + catalogPath := filepath.Join(dir, "catalog.ducklake") dataPath := filepath.Join(dir, "data") + "/" logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError})) writer, err := ducklake.NewWriter(ctx, ducklake.Config{ - CatalogPath: catalogPath, - DataPath: dataPath, - }, nil, 10, time.Second, logger) + CatalogPath: catalogPath, + CatalogStore: "duckdb", + DataPath: dataPath, + }, nil, 100, time.Hour, logger) if err != nil { t.Fatalf("ducklake writer: %v", err) } - _, err = writer.HandleEvent(ctx, wal.RowEvent{ - Schema: "public", - Table: "orders", - Columns: []wal.Column{{Name: "id", OID: 23, IsKey: true}, {Name: "name", OID: 25}}, - KeyColumns: []int{0}, - Op: wal.OpInsert, - Values: []wal.ColumnValue{ - {Name: "id", OID: 23, Value: []byte("1")}, - {Name: "name", OID: 25, Value: []byte("alice")}, - }, - }) + defer writer.Close() + + writeDuckLakeRow(t, writer, "orders", 1, "alice") + if err := writer.FlushAll(ctx); err != nil { + t.Fatalf("initial flush: %v", err) + } + + srv, err := NewServer(ServerConfig{ + ListenAddr: ":0", + TargetFormat: "ducklake", + DuckLakeCatalog: catalogPath, + DuckLakeCatalogStore: "duckdb", + DuckLakeDataPath: dataPath, + DuckLakeDB: writer.DB(), + }, nil, logger) if err != nil { - t.Fatalf("handle event: %v", err) + t.Fatalf("NewServer: %v", err) } + defer srv.Close() + + assertDuckLakeCount(t, srv, `streambed.public.orders`, 1) // fully qualified + assertDuckLakeCount(t, srv, `public.orders`, 1) // current catalog + assertDuckLakeCount(t, srv, `orders`, 1) // current catalog and schema + + // Existing tables expose newly committed rows without recreating a view. + writeDuckLakeRow(t, writer, "orders", 2, "bob") + if err := writer.FlushAll(ctx); err != nil { + t.Fatalf("existing-table flush: %v", err) + } + assertDuckLakeCount(t, srv, `public.orders`, 2) + + // Newly committed tables are visible without waiting for catalog discovery. + writeDuckLakeRow(t, writer, "customers", 1, "carol") + if err := writer.FlushAll(ctx); err != nil { + t.Fatalf("new-table flush: %v", err) + } + assertDuckLakeCount(t, srv, `public.customers`, 1) + + var viewCount int + if err := srv.duckDB.QueryRow(` + SELECT count(*) + FROM information_schema.tables + WHERE table_type = 'VIEW' AND table_name IN ('orders', 'public_orders', 'customers', 'public_customers') + `).Scan(&viewCount); err != nil { + t.Fatalf("list compatibility views: %v", err) + } + if viewCount != 0 { + t.Fatalf("got %d DuckLake compatibility views, want none", viewCount) + } +} + +func TestDuckLakeServerAttachesCatalogInStandaloneMode(t *testing.T) { + ctx := context.Background() + dir := t.TempDir() + catalogPath := filepath.Join(dir, "catalog.ducklake") + dataPath := filepath.Join(dir, "data") + "/" + logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError})) + + writer, err := ducklake.NewWriter(ctx, ducklake.Config{ + CatalogPath: catalogPath, + CatalogStore: "duckdb", + DataPath: dataPath, + }, nil, 100, time.Hour, logger) + if err != nil { + t.Fatalf("ducklake writer: %v", err) + } + writeDuckLakeRow(t, writer, "orders", 1, "alice") if err := writer.FlushAll(ctx); err != nil { t.Fatalf("flush: %v", err) } @@ -48,29 +103,47 @@ func TestDuckLakeServerRegistersViews(t *testing.T) { } srv, err := NewServer(ServerConfig{ - ListenAddr: ":0", - TargetFormat: "ducklake", - DuckLakeCatalog: catalogPath, - DuckLakeDataPath: dataPath, + ListenAddr: ":0", + TargetFormat: "ducklake", + DuckLakeCatalog: catalogPath, + DuckLakeCatalogStore: "duckdb", + DuckLakeDataPath: dataPath, }, nil, logger) if err != nil { t.Fatalf("NewServer: %v", err) } defer srv.Close() - if err := srv.refreshAndRegister(ctx); err != nil { - t.Fatalf("refreshAndRegister: %v", err) - } - var count int - if err := srv.duckDB.QueryRow(`SELECT count(*) FROM orders`).Scan(&count); err != nil { - t.Fatalf("query unqualified view: %v", err) - } - if count != 1 { - t.Fatalf("got count=%d, want 1", count) + + assertDuckLakeCount(t, srv, `streambed.public.orders`, 1) + assertDuckLakeCount(t, srv, `public.orders`, 1) + assertDuckLakeCount(t, srv, `orders`, 1) +} + +func writeDuckLakeRow(t *testing.T, writer *ducklake.Writer, table string, id int, name string) { + t.Helper() + _, err := writer.HandleEvent(context.Background(), wal.RowEvent{ + Schema: "public", + Table: table, + Columns: []wal.Column{{Name: "id", OID: 23, IsKey: true}, {Name: "name", OID: 25}}, + KeyColumns: []int{0}, + Op: wal.OpInsert, + Values: []wal.ColumnValue{ + {Name: "id", OID: 23, Value: []byte(fmt.Sprint(id))}, + {Name: "name", OID: 25, Value: []byte(name)}, + }, + }) + if err != nil { + t.Fatalf("write %s row: %v", table, err) } - if err := srv.duckDB.QueryRow(`SELECT count(*) FROM public_orders`).Scan(&count); err != nil { - t.Fatalf("query qualified view: %v", err) +} + +func assertDuckLakeCount(t *testing.T, srv *Server, table string, want int) { + t.Helper() + var got int + if err := srv.duckDB.QueryRow(`SELECT count(*) FROM ` + table).Scan(&got); err != nil { + t.Fatalf("query %s: %v", table, err) } - if count != 1 { - t.Fatalf("got qualified count=%d, want 1", count) + if got != want { + t.Fatalf("query %s: got count=%d, want %d", table, got, want) } } diff --git a/internal/server/server.go b/internal/server/server.go index 97cbb0d..1c1fe7d 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -7,6 +7,7 @@ import ( "log/slog" "os" "strings" + "sync" "time" duckdb "github.com/duckdb/duckdb-go/v2" @@ -27,58 +28,116 @@ type ServerConfig struct { DuckLakeCatalog string DuckLakeCatalogStore string DuckLakeDataPath string + DuckLakeDB *sql.DB // optional shared writer connection for in-process sync + query } // Server implements a Postgres-wire-compatible query interface backed by DuckDB. -// It reads Iceberg tables from S3 and serves them to any Postgres client. +// It serves Iceberg views or tables from an attached DuckLake catalog. type Server struct { - cfg ServerConfig - catalog *TableCatalog - duckDB *sql.DB - logger *slog.Logger + cfg ServerConfig + catalog *TableCatalog + duckDB *sql.DB + duckDBMu sync.RWMutex + ownsDuckDB bool + logger *slog.Logger } -// NewServer creates a query server. Initializes an embedded DuckDB instance -// configured with S3 credentials and loads the Iceberg + httpfs extensions. +// NewServer creates a query server and initializes DuckDB for the selected +// lakehouse target. func NewServer(cfg ServerConfig, s3Client storage.ObjectStorage, logger *slog.Logger) (*Server, error) { if cfg.TargetFormat == "" { cfg.TargetFormat = "iceberg" } - db, err := sql.Open("duckdb", "") - if err != nil { - return nil, fmt.Errorf("open duckdb: %w", err) + if cfg.DuckLakeDB != nil && cfg.TargetFormat != "ducklake" { + return nil, fmt.Errorf("shared DuckLake DB requires ducklake target") + } + db := cfg.DuckLakeDB + ownsDuckDB := db == nil + if db == nil { + var err error + db, err = sql.Open("duckdb", "") + if err != nil { + return nil, fmt.Errorf("open duckdb: %w", err) + } + } + closeOnError := func() { + if ownsDuckDB { + db.Close() + } } var catalog *TableCatalog if cfg.TargetFormat == "ducklake" { - if err := ducklake.Configure(context.Background(), db, ducklake.Config{ - CatalogPath: cfg.DuckLakeCatalog, - CatalogStore: cfg.DuckLakeCatalogStore, - DataPath: cfg.DuckLakeDataPath, - S3Endpoint: cfg.S3Endpoint, - S3Region: cfg.S3Region, - }); err != nil { - db.Close() - return nil, fmt.Errorf("configure ducklake: %w", err) + // DuckLake attachments and USE are session-local. Keep all query-server + // work on the configured DuckDB session. Sync mode shares the writer's + // session so DuckDB-backed catalogs expose commits without a stale cache. + db.SetMaxOpenConns(1) + if cfg.DuckLakeDB == nil { + if err := ducklake.Configure(context.Background(), db, ducklake.Config{ + CatalogPath: cfg.DuckLakeCatalog, + CatalogStore: cfg.DuckLakeCatalogStore, + DataPath: cfg.DuckLakeDataPath, + S3Endpoint: cfg.S3Endpoint, + S3Region: cfg.S3Region, + }); err != nil { + closeOnError() + return nil, fmt.Errorf("configure ducklake: %w", err) + } + } + if err := configureDuckLakeQuerySession(db); err != nil { + closeOnError() + return nil, err } } else { conn, _ := db.Conn(context.Background()) // configure DuckDB on a specific connection configureDuckDBPerConn(context.Background(), conn, cfg) if err := configureDuckDB(db, cfg); err != nil { - db.Close() + closeOnError() return nil, fmt.Errorf("configure duckdb: %w", err) } catalog = NewTableCatalog(s3Client, cfg.S3Bucket, cfg.S3Prefix, logger) } return &Server{ - cfg: cfg, - catalog: catalog, - duckDB: db, - logger: logger, + cfg: cfg, + catalog: catalog, + duckDB: db, + ownsDuckDB: ownsDuckDB, + logger: logger, }, nil } +func configureDuckLakeQuerySession(db *sql.DB) error { + if _, err := db.Exec("USE streambed"); err != nil { + return fmt.Errorf("use ducklake catalog: %w", err) + } + // Match PostgreSQL's usual default schema while retaining DuckLake's main + // schema as a fallback. Missing schemas in the search path are allowed. + if _, err := db.Exec("SET search_path = 'streambed.public,streambed.main'"); err != nil { + return fmt.Errorf("configure ducklake search path: %w", err) + } + return nil +} + +// SetDuckLakeDB moves an in-process query server to a replacement writer +// connection after pipeline reconnection. +func (s *Server) SetDuckLakeDB(db *sql.DB) error { + if s.cfg.TargetFormat != "ducklake" { + return fmt.Errorf("cannot replace DuckDB connection for %s target", s.cfg.TargetFormat) + } + if s.ownsDuckDB { + return fmt.Errorf("cannot replace query server-owned DuckDB connection") + } + db.SetMaxOpenConns(1) + if err := configureDuckLakeQuerySession(db); err != nil { + return err + } + s.duckDBMu.Lock() + s.duckDB = db + s.duckDBMu.Unlock() + return nil +} + // configureDuckDBPerConn configures DuckDB on a specific connection -- delete it func configureDuckDBPerConn(ctx context.Context, con *sql.Conn, cfg ServerConfig) error { stmts := []string{ @@ -189,17 +248,18 @@ func configureDuckDB(db *sql.DB, cfg ServerConfig) error { } // Start begins listening for Postgres client connections and serving queries. -// It performs an initial catalog refresh, starts a background refresh goroutine, -// and blocks until ctx is cancelled. +// Iceberg mode refreshes discovered views in the background; DuckLake mode uses +// its attached catalog directly. Start blocks until ctx is cancelled. func (s *Server) Start(ctx context.Context) error { - // Initial catalog refresh and view registration - if err := s.refreshAndRegister(ctx); err != nil { - s.logger.Warn("initial catalog refresh failed (will retry)", "error", err) + // Iceberg tables must be discovered in S3 and exposed as views. DuckLake + // tables are already visible through the directly attached catalog. + if s.cfg.TargetFormat != "ducklake" { + if err := s.refreshAndRegister(ctx); err != nil { + s.logger.Warn("initial catalog refresh failed (will retry)", "error", err) + } + go s.refreshLoop(ctx) } - // Periodic catalog refresh in background - go s.refreshLoop(ctx) - // Create psql-wire server srv, err := wire.NewServer(s.handleParse, wire.Logger(s.logger), @@ -241,6 +301,9 @@ func (s *Server) handleParse(ctx context.Context, query string) (wire.PreparedSt s.logger.Debug("query received", "query", query) + s.duckDBMu.RLock() + defer s.duckDBMu.RUnlock() + // Execute query against DuckDB rows, err := s.duckDB.QueryContext(ctx, query) if err != nil { @@ -311,9 +374,8 @@ func (s *Server) handleParse(ctx context.Context, query string) (wire.PreparedSt // If DuckDB's engine was invalidated by a FATAL error (e.g., corrupt Iceberg // table), it re-opens the DuckDB instance and retries view registration. func (s *Server) refreshAndRegister(ctx context.Context) error { - if s.cfg.TargetFormat == "ducklake" { - return s.registerDuckLakeViews(ctx) - } + s.duckDBMu.Lock() + defer s.duckDBMu.Unlock() if err := s.catalog.Refresh(ctx); err != nil { return err } @@ -340,49 +402,6 @@ func (s *Server) refreshAndRegister(ctx context.Context) error { return s.catalog.RegisterViews(s.duckDB) } -func (s *Server) registerDuckLakeViews(ctx context.Context) error { - const catalogName = "streambed" - rows, err := s.duckDB.QueryContext(ctx, - "SELECT table_schema, table_name FROM information_schema.tables WHERE table_catalog = ? AND table_type = 'BASE TABLE'", - catalogName, - ) - if err != nil { - return fmt.Errorf("list ducklake tables: %w", err) - } - defer rows.Close() - type tableInfo struct { - schema string - table string - } - var tables []tableInfo - nameCount := map[string]int{} - for rows.Next() { - var info tableInfo - if err := rows.Scan(&info.schema, &info.table); err != nil { - return err - } - tables = append(tables, info) - nameCount[info.table]++ - } - if err := rows.Err(); err != nil { - return err - } - for _, info := range tables { - src := fmt.Sprintf("%s.%s.%s", quoteIdent(catalogName), quoteIdent(info.schema), quoteIdent(info.table)) - qualified := quoteIdent(info.schema + "_" + info.table) - if _, err := s.duckDB.ExecContext(ctx, fmt.Sprintf("CREATE OR REPLACE VIEW %s AS SELECT * FROM %s", qualified, src)); err != nil { - return fmt.Errorf("register ducklake view %s_%s: %w", info.schema, info.table, err) - } - if nameCount[info.table] == 1 { - if _, err := s.duckDB.ExecContext(ctx, fmt.Sprintf("CREATE OR REPLACE VIEW %s AS SELECT * FROM %s", quoteIdent(info.table), src)); err != nil { - return fmt.Errorf("register ducklake view %s: %w", info.table, err) - } - } - } - s.logger.Info("ducklake views registered", "count", len(tables)) - return nil -} - // refreshLoop periodically refreshes the catalog and re-registers views. func (s *Server) refreshLoop(ctx context.Context) { ticker := time.NewTicker(30 * time.Second) @@ -402,6 +421,11 @@ func (s *Server) refreshLoop(ctx context.Context) { // Close shuts down the DuckDB connection. func (s *Server) Close() error { + if !s.ownsDuckDB { + return nil + } + s.duckDBMu.Lock() + defer s.duckDBMu.Unlock() return s.duckDB.Close() } @@ -456,7 +480,3 @@ func normalizeValue(v any) any { return v } } - -func quoteIdent(s string) string { - return `"` + strings.ReplaceAll(s, `"`, `""`) + `"` -} From 606f9426bd40bf7872b4c240fafedb6341fb4c2c Mon Sep 17 00:00:00 2001 From: viggy28 Date: Sun, 30 Aug 2026 23:42:16 -0700 Subject: [PATCH 2/2] fix: isolate DuckLake query sessions --- cmd/streambed/main.go | 15 +--- internal/ducklake/writer.go | 17 ++-- internal/server/ducklake_test.go | 15 +++- internal/server/server.go | 130 +++++++++++++++---------------- 4 files changed, 86 insertions(+), 91 deletions(-) diff --git a/cmd/streambed/main.go b/cmd/streambed/main.go index 19cd937..f2f286e 100644 --- a/cmd/streambed/main.go +++ b/cmd/streambed/main.go @@ -3,7 +3,6 @@ package main import ( "bufio" "context" - "database/sql" "fmt" "log/slog" "os" @@ -327,13 +326,8 @@ func runSync(cmd *cobra.Command, args []string) error { } // Start query server if --query-addr is set - var querySrv *server.Server if cfg.QueryAddr != "" { - var duckLakeDB *sql.DB - if duckWriter != nil { - duckLakeDB = duckWriter.DB() - } - querySrv, err = server.NewServer(server.ServerConfig{ + querySrv, err := server.NewServer(server.ServerConfig{ ListenAddr: cfg.QueryAddr, S3Bucket: cfg.S3Bucket, S3Prefix: cfg.S3Prefix, @@ -343,7 +337,6 @@ func runSync(cmd *cobra.Command, args []string) error { DuckLakeCatalog: cfg.DuckLakeCatalog, DuckLakeCatalogStore: cfg.DuckLakeCatalogStore, DuckLakeDataPath: cfg.EffectiveDuckLakeDataPath(), - DuckLakeDB: duckLakeDB, }, s3Client, logger) if err != nil { return fmt.Errorf("create query server: %w", err) @@ -533,12 +526,6 @@ func runSync(cmd *cobra.Command, args []string) error { logger.Error("reconnect: writer setup failed", "error", err) continue } - if querySrv != nil && duckWriter != nil { - if err := querySrv.SetDuckLakeDB(duckWriter.DB()); err != nil { - logger.Error("reconnect: query server setup failed", "error", err) - continue - } - } p = pipeline.New(pgConn, cfg.SlotName, pubName, startLSN, cfg.ExcludeTables, logger, stateStore, tableFlushLSN, writer, cfg.FlushInterval, metaQuerier) diff --git a/internal/ducklake/writer.go b/internal/ducklake/writer.go index 3b3045f..28123da 100644 --- a/internal/ducklake/writer.go +++ b/internal/ducklake/writer.go @@ -27,6 +27,7 @@ type Config struct { S3Endpoint string S3Region string CatalogName string + ReadOnly bool } type Writer struct { @@ -143,10 +144,16 @@ func Configure(ctx context.Context, db *sql.DB, cfg Config) error { return fmt.Errorf("exec %q: %w", stmt, err) } } - attach := fmt.Sprintf("ATTACH '%s' AS %s (DATA_PATH '%s')", + attachOptions := []string{ + fmt.Sprintf("DATA_PATH '%s'", strings.ReplaceAll(cfg.DataPath, "'", "''")), + } + if cfg.ReadOnly { + attachOptions = append(attachOptions, "READ_ONLY") + } + attach := fmt.Sprintf("ATTACH '%s' AS %s (%s)", duckLakeAttachPath(cfg.CatalogPath, catalogStore), quoteIdent(catalogName(cfg)), - strings.ReplaceAll(cfg.DataPath, "'", "''"), + strings.Join(attachOptions, ", "), ) if _, err := db.ExecContext(ctx, attach); err != nil { return fmt.Errorf("attach ducklake catalog: %w", err) @@ -176,12 +183,6 @@ func catalogName(cfg Config) string { return cfg.CatalogName } -// DB returns the configured DuckDB handle so the in-process query server can -// share the same DuckLake attachment and observe writer commits immediately. -func (w *Writer) DB() *sql.DB { - return w.db -} - func (w *Writer) Close() error { return w.db.Close() } diff --git a/internal/server/ducklake_test.go b/internal/server/ducklake_test.go index 7310dc2..b0dd276 100644 --- a/internal/server/ducklake_test.go +++ b/internal/server/ducklake_test.go @@ -41,7 +41,6 @@ func TestDuckLakeServerUsesAttachedCatalogDirectly(t *testing.T) { DuckLakeCatalog: catalogPath, DuckLakeCatalogStore: "duckdb", DuckLakeDataPath: dataPath, - DuckLakeDB: writer.DB(), }, nil, logger) if err != nil { t.Fatalf("NewServer: %v", err) @@ -52,6 +51,15 @@ func TestDuckLakeServerUsesAttachedCatalogDirectly(t *testing.T) { assertDuckLakeCount(t, srv, `public.orders`, 1) // current catalog assertDuckLakeCount(t, srv, `orders`, 1) // current catalog and schema + if _, err := srv.duckDB.Exec(`INSERT INTO streambed.public.orders VALUES (99, 'mallory')`); err == nil { + t.Fatal("query-side DuckLake attachment accepted a write") + } + // Mutating the query session must not disturb the writer's attachment. The + // next query also gets a fresh session, so this DETACH is self-contained. + if _, err := srv.duckDB.Exec(`USE memory; DETACH streambed`); err != nil { + t.Fatalf("detach query-side catalog: %v", err) + } + // Existing tables expose newly committed rows without recreating a view. writeDuckLakeRow(t, writer, "orders", 2, "bob") if err := writer.FlushAll(ctx); err != nil { @@ -139,6 +147,11 @@ func writeDuckLakeRow(t *testing.T, writer *ducklake.Writer, table string, id in func assertDuckLakeCount(t *testing.T, srv *Server, table string, want int) { t.Helper() + srv.duckDBMu.Lock() + defer srv.duckDBMu.Unlock() + if err := srv.resetDuckLakeQueryDB(context.Background()); err != nil { + t.Fatalf("reset DuckLake query DB: %v", err) + } var got int if err := srv.duckDB.QueryRow(`SELECT count(*) FROM ` + table).Scan(&got); err != nil { t.Fatalf("query %s: %v", table, err) diff --git a/internal/server/server.go b/internal/server/server.go index 1c1fe7d..c195c67 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -28,18 +28,16 @@ type ServerConfig struct { DuckLakeCatalog string DuckLakeCatalogStore string DuckLakeDataPath string - DuckLakeDB *sql.DB // optional shared writer connection for in-process sync + query } // Server implements a Postgres-wire-compatible query interface backed by DuckDB. // It serves Iceberg views or tables from an attached DuckLake catalog. type Server struct { - cfg ServerConfig - catalog *TableCatalog - duckDB *sql.DB - duckDBMu sync.RWMutex - ownsDuckDB bool - logger *slog.Logger + cfg ServerConfig + catalog *TableCatalog + duckDB *sql.DB + duckDBMu sync.Mutex + logger *slog.Logger } // NewServer creates a query server and initializes DuckDB for the selected @@ -48,43 +46,28 @@ func NewServer(cfg ServerConfig, s3Client storage.ObjectStorage, logger *slog.Lo if cfg.TargetFormat == "" { cfg.TargetFormat = "iceberg" } - if cfg.DuckLakeDB != nil && cfg.TargetFormat != "ducklake" { - return nil, fmt.Errorf("shared DuckLake DB requires ducklake target") - } - db := cfg.DuckLakeDB - ownsDuckDB := db == nil - if db == nil { - var err error - db, err = sql.Open("duckdb", "") - if err != nil { - return nil, fmt.Errorf("open duckdb: %w", err) - } - } - closeOnError := func() { - if ownsDuckDB { - db.Close() - } + db, err := sql.Open("duckdb", "") + if err != nil { + return nil, fmt.Errorf("open duckdb: %w", err) } var catalog *TableCatalog if cfg.TargetFormat == "ducklake" { - // DuckLake attachments and USE are session-local. Keep all query-server - // work on the configured DuckDB session. Sync mode shares the writer's - // session so DuckDB-backed catalogs expose commits without a stale cache. + // The query server owns an isolated, read-only DuckDB session so client + // SQL cannot mutate the writer's DuckLake attachment. db.SetMaxOpenConns(1) - if cfg.DuckLakeDB == nil { - if err := ducklake.Configure(context.Background(), db, ducklake.Config{ - CatalogPath: cfg.DuckLakeCatalog, - CatalogStore: cfg.DuckLakeCatalogStore, - DataPath: cfg.DuckLakeDataPath, - S3Endpoint: cfg.S3Endpoint, - S3Region: cfg.S3Region, - }); err != nil { - closeOnError() - return nil, fmt.Errorf("configure ducklake: %w", err) - } + if err := ducklake.Configure(context.Background(), db, ducklake.Config{ + CatalogPath: cfg.DuckLakeCatalog, + CatalogStore: cfg.DuckLakeCatalogStore, + DataPath: cfg.DuckLakeDataPath, + S3Endpoint: cfg.S3Endpoint, + S3Region: cfg.S3Region, + ReadOnly: true, + }); err != nil { + db.Close() + return nil, fmt.Errorf("configure ducklake: %w", err) } if err := configureDuckLakeQuerySession(db); err != nil { - closeOnError() + db.Close() return nil, err } } else { @@ -92,21 +75,46 @@ func NewServer(cfg ServerConfig, s3Client storage.ObjectStorage, logger *slog.Lo // configure DuckDB on a specific connection configureDuckDBPerConn(context.Background(), conn, cfg) if err := configureDuckDB(db, cfg); err != nil { - closeOnError() + db.Close() return nil, fmt.Errorf("configure duckdb: %w", err) } catalog = NewTableCatalog(s3Client, cfg.S3Bucket, cfg.S3Prefix, logger) } return &Server{ - cfg: cfg, - catalog: catalog, - duckDB: db, - ownsDuckDB: ownsDuckDB, - logger: logger, + cfg: cfg, + catalog: catalog, + duckDB: db, + logger: logger, }, nil } +func (s *Server) resetDuckLakeQueryDB(ctx context.Context) error { + db, err := sql.Open("duckdb", "") + if err != nil { + return fmt.Errorf("open fresh duckdb query session: %w", err) + } + db.SetMaxOpenConns(1) + if err := ducklake.Configure(ctx, db, ducklake.Config{ + CatalogPath: s.cfg.DuckLakeCatalog, + CatalogStore: s.cfg.DuckLakeCatalogStore, + DataPath: s.cfg.DuckLakeDataPath, + S3Endpoint: s.cfg.S3Endpoint, + S3Region: s.cfg.S3Region, + ReadOnly: true, + }); err != nil { + db.Close() + return fmt.Errorf("configure fresh ducklake query session: %w", err) + } + if err := configureDuckLakeQuerySession(db); err != nil { + db.Close() + return err + } + oldDB := s.duckDB + s.duckDB = db + return oldDB.Close() +} + func configureDuckLakeQuerySession(db *sql.DB) error { if _, err := db.Exec("USE streambed"); err != nil { return fmt.Errorf("use ducklake catalog: %w", err) @@ -119,25 +127,6 @@ func configureDuckLakeQuerySession(db *sql.DB) error { return nil } -// SetDuckLakeDB moves an in-process query server to a replacement writer -// connection after pipeline reconnection. -func (s *Server) SetDuckLakeDB(db *sql.DB) error { - if s.cfg.TargetFormat != "ducklake" { - return fmt.Errorf("cannot replace DuckDB connection for %s target", s.cfg.TargetFormat) - } - if s.ownsDuckDB { - return fmt.Errorf("cannot replace query server-owned DuckDB connection") - } - db.SetMaxOpenConns(1) - if err := configureDuckLakeQuerySession(db); err != nil { - return err - } - s.duckDBMu.Lock() - s.duckDB = db - s.duckDBMu.Unlock() - return nil -} - // configureDuckDBPerConn configures DuckDB on a specific connection -- delete it func configureDuckDBPerConn(ctx context.Context, con *sql.Conn, cfg ServerConfig) error { stmts := []string{ @@ -301,8 +290,16 @@ func (s *Server) handleParse(ctx context.Context, query string) (wire.PreparedSt s.logger.Debug("query received", "query", query) - s.duckDBMu.RLock() - defer s.duckDBMu.RUnlock() + s.duckDBMu.Lock() + defer s.duckDBMu.Unlock() + if s.cfg.TargetFormat == "ducklake" { + // Give every client query an isolated, read-only attachment. Besides + // containing session mutations such as DETACH, reopening refreshes the + // snapshot cached by DuckDB-backed metadata catalogs. + if err := s.resetDuckLakeQueryDB(ctx); err != nil { + return nil, err + } + } // Execute query against DuckDB rows, err := s.duckDB.QueryContext(ctx, query) @@ -421,9 +418,6 @@ func (s *Server) refreshLoop(ctx context.Context) { // Close shuts down the DuckDB connection. func (s *Server) Close() error { - if !s.ownsDuckDB { - return nil - } s.duckDBMu.Lock() defer s.duckDBMu.Unlock() return s.duckDB.Close()