From b117d6a6701d9e46eb1d03a3efecae5777590b4a Mon Sep 17 00:00:00 2001
From: James Greenhill <fuziontech@gmail.com>
Date: Thu, 30 Jul 2026 14:41:13 +0000
Subject: [PATCH] feat(duckdbservice): set late_materialization_max_rows=6000
 on all workers

Apply SET GLOBAL late_materialization_max_rows = 6000 when a worker opens
its DuckDB pair. The setting is session-scoped (LOCAL) in DuckDB, so SET
GLOBAL is required for the value to reach every session connection on the
shared connector; a plain SET would only affect the warmup connection.
---
 duckdbservice/duckdb_pair.go      | 10 +++++++++
 duckdbservice/duckdb_pair_test.go | 35 +++++++++++++++++++++++++++++++
 2 files changed, 45 insertions(+)
 create mode 100644 duckdbservice/duckdb_pair_test.go

diff --git a/duckdbservice/duckdb_pair.go b/duckdbservice/duckdb_pair.go
index 3a015a8e..c5bcd834 100644
--- a/duckdbservice/duckdb_pair.go
+++ b/duckdbservice/duckdb_pair.go
@@ -5,6 +5,7 @@ import (
 	"database/sql"
 	"database/sql/driver"
 	"fmt"
+	"log/slog"
 	"time"
 
 	duckdb "github.com/duckdb/duckdb-go/v2"
@@ -86,6 +87,12 @@ type workerRequiredExtensionExecer interface {
 	Exec(query string, args ...any) (sql.Result, error)
 }
 
+// workerLateMaterializationMaxRows caps how many rows DuckDB late-materializes
+// for LIMIT/SAMPLE queries. late_materialization_max_rows is session-scoped
+// (LOCAL) in DuckDB, so it must be applied with SET GLOBAL for the value to
+// reach every session connection on the shared connector.
+const workerLateMaterializationMaxRows = 6000
+
 func loadWorkerRequiredExtensions(db workerRequiredExtensionExecer) error {
 	if _, err := db.Exec("LOAD postgres_scanner"); err != nil {
 		return fmt.Errorf("load required worker extension postgres_scanner: %w", err)
@@ -147,6 +154,9 @@ func OpenDuckDBPair(cfg server.Config, username string) (*DuckDBPair, error) {
 		_ = connector.Close()
 		return nil, err
 	}
+	if _, err := mainDB.Exec(fmt.Sprintf("SET GLOBAL late_materialization_max_rows = %d", workerLateMaterializationMaxRows)); err != nil {
+		slog.Warn("Failed to set DuckDB late_materialization_max_rows.", "late_materialization_max_rows", workerLateMaterializationMaxRows, "error", err)
+	}
 	if err := loadWorkerRequiredExtensions(mainDB); err != nil {
 		_ = mainDB.Close()
 		_ = controlDB.Close()
diff --git a/duckdbservice/duckdb_pair_test.go b/duckdbservice/duckdb_pair_test.go
new file mode 100644
index 00000000..6f8d755f
--- /dev/null
+++ b/duckdbservice/duckdb_pair_test.go
@@ -0,0 +1,35 @@
+package duckdbservice
+
+import (
+	"database/sql"
+	"testing"
+
+	"github.com/posthog/duckgres/server"
+)
+
+// TestOpenDuckDBPairSetsLateMaterializationMaxRows is the regression guard for
+// the worker rollout that caps late materialization: every worker DB opened via
+// OpenDuckDBPair must have late_materialization_max_rows = 6000.
+//
+// The assertion runs on a *fresh* connection from the same connector — the
+// same shape as a session-pool connection serving user queries — so it also
+// proves the setting is visible beyond the warmup connection that applied it.
+func TestOpenDuckDBPairSetsLateMaterializationMaxRows(t *testing.T) {
+	cfg := server.Config{DataDir: t.TempDir()}
+	pair, err := OpenDuckDBPair(cfg, "worker")
+	if err != nil {
+		t.Fatalf("OpenDuckDBPair: %v", err)
+	}
+	defer func() { _ = pair.Close() }()
+
+	sessionDB := sql.OpenDB(pair.connector)
+	defer func() { _ = sessionDB.Close() }()
+
+	var got string
+	if err := sessionDB.QueryRow("SELECT current_setting('late_materialization_max_rows')").Scan(&got); err != nil {
+		t.Fatalf("read late_materialization_max_rows: %v", err)
+	}
+	if got != "6000" {
+		t.Errorf("late_materialization_max_rows = %q, want %q", got, "6000")
+	}
+}