diff --git a/tests/mw-dev/scenario/perf/adapter_test.go b/tests/mw-dev/scenario/perf/adapter_test.go index 5079a0cd..33821b38 100644 --- a/tests/mw-dev/scenario/perf/adapter_test.go +++ b/tests/mw-dev/scenario/perf/adapter_test.go @@ -99,6 +99,48 @@ func TestExecutorRunsPerfStepAndWritesArtifacts(t *testing.T) { } } +func TestExecutorRecordsBothPGWireCacheVariantsWithEqualResources(t *testing.T) { + targets := []perfcore.Protocol{perfcore.ProtocolPGWireUncached, perfcore.ProtocolPGWireCached} + factory := &fakeDriverFactory{} + executor := NewExecutor(ExecutorConfig{ + Connection: scenariosql.ConnectionConfig{DialHost: "127.0.0.1", SNISuffix: ".example.test", SSLMode: "require"}, + OutputDir: t.TempDir(), DriverFactory: factory, + }) + err := executor.ExecuteStep(context.Background(), core.Step{ + ID: "cache-comparison", Type: StepTypePerfQueries, + With: map[string]any{ + "org_id": "scenario-org", "username": "root", "password": "test-password", + "catalog_file": writePerfCatalog(t, targets), "targets": []any{"pgwire_uncached", "pgwire_cached"}, + "run_id": "cache-run", "worker_cpu": "3", "worker_memory": "12Gi", + }, + }) + if err != nil { + t.Fatal(err) + } + result, ok := executor.State().Result("cache-comparison") + if !ok || result.Summary.RunID != "cache-run" || result.Summary.TotalQueries != 2 || result.Summary.WarmupQueries != 2 { + t.Fatalf("unexpected result: %+v", result) + } + if len(factory.pgwireConnections) != 2 { + t.Fatalf("PGWire factories called for %d targets, want 2", len(factory.pgwireConnections)) + } + for _, protocol := range targets { + connection, ok := factory.pgwireConnections[protocol] + if !ok || !strings.Contains(connection.DSN, "options='-c duckgres.worker_cpu=3 -c duckgres.worker_memory=12Gi'") { + t.Fatalf("%s connection missing standard resources: %+v", protocol, connection) + } + } + raw, err := os.ReadFile(filepath.Join(result.OutputDir, "query_results.csv")) + if err != nil { + t.Fatal(err) + } + for _, protocol := range targets { + if !strings.Contains(string(raw), "\nq1,i1,1,"+string(protocol)+",ok,") { + t.Fatalf("missing %s result for shared query/intent: %s", protocol, raw) + } + } +} + func TestExecutorRestrictsCatalogToStepTargets(t *testing.T) { catalogPath := writePerfCatalog(t, []perfcore.Protocol{perfcore.ProtocolPGWire}) provisionState := provision.NewState() @@ -487,20 +529,25 @@ func writePerfCatalog(t *testing.T, targets []perfcore.Protocol) string { } type fakeDriverFactory struct { - pgwireConnection scenariosql.PGWireConnection - pgwireErr error - pgwireDriver *fakeProtocolDriver - trinoConnection trinodriver.ConnectionConfig - trinoContext context.Context - trinoDriver *fakeProtocolDriver - athenaConnection athenadriver.ConnectionConfig - athenaContext context.Context - athenaDriver *fakeProtocolDriver + pgwireConnections map[perfcore.Protocol]scenariosql.PGWireConnection + pgwireConnection scenariosql.PGWireConnection + pgwireErr error + pgwireDriver *fakeProtocolDriver + trinoConnection trinodriver.ConnectionConfig + trinoContext context.Context + trinoDriver *fakeProtocolDriver + athenaConnection athenadriver.ConnectionConfig + athenaContext context.Context + athenaDriver *fakeProtocolDriver } -func (f *fakeDriverFactory) NewPGWire(connection scenariosql.PGWireConnection) (perfcore.ProtocolDriver, error) { +func (f *fakeDriverFactory) NewPGWire(connection scenariosql.PGWireConnection, protocol perfcore.Protocol) (perfcore.ProtocolDriver, error) { + if f.pgwireConnections == nil { + f.pgwireConnections = make(map[perfcore.Protocol]scenariosql.PGWireConnection) + } + f.pgwireConnections[protocol] = connection f.pgwireConnection = connection - f.pgwireDriver = &fakeProtocolDriver{protocol: perfcore.ProtocolPGWire, err: f.pgwireErr} + f.pgwireDriver = &fakeProtocolDriver{protocol: protocol, err: f.pgwireErr} return f.pgwireDriver, nil } diff --git a/tests/mw-dev/scenario/perf/steps.go b/tests/mw-dev/scenario/perf/steps.go index d262ef0f..d685ddf8 100644 --- a/tests/mw-dev/scenario/perf/steps.go +++ b/tests/mw-dev/scenario/perf/steps.go @@ -21,7 +21,7 @@ import ( const StepTypePerfQueries = "perf_queries" type DriverFactory interface { - NewPGWire(connection scenariosql.PGWireConnection) (perfcore.ProtocolDriver, error) + NewPGWire(connection scenariosql.PGWireConnection, protocol perfcore.Protocol) (perfcore.ProtocolDriver, error) NewTrino(ctx context.Context, connection trinodriver.ConnectionConfig) (perfcore.ProtocolDriver, error) NewAthena(ctx context.Context, connection athenadriver.ConnectionConfig) (perfcore.ProtocolDriver, error) } @@ -312,7 +312,7 @@ func targetsFromWith(step core.Step) ([]perfcore.Protocol, error) { } target := perfcore.Protocol(value) switch target { - case perfcore.ProtocolPGWire, perfcore.ProtocolTrino, perfcore.ProtocolAthena: + case perfcore.ProtocolPGWire, perfcore.ProtocolPGWireUncached, perfcore.ProtocolPGWireCached, perfcore.ProtocolTrino, perfcore.ProtocolAthena: default: return nil, classified(ErrorClassConfig, fmt.Errorf("step %s with.targets[%d] has unsupported perf protocol %q", step.ID, i, target)) } @@ -356,12 +356,12 @@ func (e *Executor) driversForCatalog(ctx context.Context, catalog perfcore.Catal continue } switch target { - case perfcore.ProtocolPGWire: + case perfcore.ProtocolPGWire, perfcore.ProtocolPGWireUncached, perfcore.ProtocolPGWireCached: connection, err := e.pgwireConnection(spec) if err != nil { return nil, err } - driver, err := e.driverFactory.NewPGWire(connection) + driver, err := e.driverFactory.NewPGWire(connection, target) if err != nil { return nil, classified(ErrorClassConfig, fmt.Errorf("create pgwire perf driver: %w", err)) } @@ -483,12 +483,17 @@ func closeDrivers(drivers map[perfcore.Protocol]perfcore.ProtocolDriver) { } } -func (defaultDriverFactory) NewPGWire(connection scenariosql.PGWireConnection) (perfcore.ProtocolDriver, error) { +func (defaultDriverFactory) NewPGWire(connection scenariosql.PGWireConnection, protocol perfcore.Protocol) (perfcore.ProtocolDriver, error) { db, err := connection.OpenDB() if err != nil { return nil, err } - return pgdriver.NewWithDB(db), nil + driver, err := pgdriver.NewWithDBAndProtocol(db, protocol) + if err != nil { + _ = db.Close() + return nil, err + } + return driver, nil } func (defaultDriverFactory) NewTrino(ctx context.Context, connection trinodriver.ConnectionConfig) (perfcore.ProtocolDriver, error) { diff --git a/tests/mw-dev/scenario/runner_test.go b/tests/mw-dev/scenario/runner_test.go index 9e3f42a3..5b395cdd 100644 --- a/tests/mw-dev/scenario/runner_test.go +++ b/tests/mw-dev/scenario/runner_test.go @@ -745,8 +745,8 @@ func assertPerfTargetsOnlyPGWire(t *testing.T, step core.Step) { func assertPerfTargetsPGWireTrinoAndAthena(t *testing.T, step core.Step) { t.Helper() targets, ok := step.With["targets"].([]any) - if !ok || len(targets) != 3 || targets[0] != "pgwire" || targets[1] != "trino" || targets[2] != "athena" { - t.Fatalf("perf step %s targets = %#v, want [pgwire trino athena]", step.ID, step.With["targets"]) + if !ok || len(targets) != 4 || targets[0] != "pgwire_uncached" || targets[1] != "pgwire_cached" || targets[2] != "trino" || targets[3] != "athena" { + t.Fatalf("perf step %s targets = %#v, want [pgwire_uncached pgwire_cached trino athena]", step.ID, step.With["targets"]) } } diff --git a/tests/mw-dev/scenario/scenarios/posthog_frozen_perf.yaml b/tests/mw-dev/scenario/scenarios/posthog_frozen_perf.yaml index 9d796185..ed582459 100644 --- a/tests/mw-dev/scenario/scenarios/posthog_frozen_perf.yaml +++ b/tests/mw-dev/scenario/scenarios/posthog_frozen_perf.yaml @@ -80,7 +80,7 @@ steps: org_id: ${env:DUCKGRES_SCENARIO_ORG_ID} catalog: ducklake catalog_file: ../../../perf/queries/ducklake_posthog_tables.yaml - targets: [pgwire, trino, athena] + targets: [pgwire_uncached, pgwire_cached, trino, athena] trino_ca_cert_file: ${env:DUCKGRES_SCENARIO_TRINO_CA_CERT} trino_startup_timeout: 2m trino_startup_poll_interval: 2s diff --git a/tests/perf/README.md b/tests/perf/README.md index 6414c755..aef39dd2 100644 --- a/tests/perf/README.md +++ b/tests/perf/README.md @@ -4,7 +4,8 @@ This package contains the golden-query performance harness. ## Protocol Drivers -Catalogs may target `pgwire`, `trino`, `athena`, or any supported combination. All drivers execute the same +Catalogs may target `pgwire`, `pgwire_uncached`, `pgwire_cached`, `trino`, +`athena`, or any supported combination. All drivers execute the same rendered statement stored in the existing `pgwire_sql` catalog field; the legacy field name is retained for catalog compatibility and must not be used to create a second, protocol-specific query definition. Keep shared benchmark @@ -32,9 +33,48 @@ readiness, Kubernetes Secret projection, and file-authenticator refresh. When a catalog targets multiple protocols, the runner completes all warmup and measured iterations for one protocol before starting the next protocol in the catalog's declared target order. This keeps each protocol's connection and -worker cache active throughout its measurements and prevents slow queries in +worker context active throughout its measurements and prevents slow queries in one protocol from changing another protocol's cache context. +### DuckDB cache comparison + +`posthog_frozen_perf` runs four separately labeled targets in one result set, +in order: **`pgwire_uncached` (baseline)**, `pgwire_cached`, `trino`, `athena`. +Both DuckDB variants use identical queries, worker resource requests, warmup +counts, and measured iterations. Each variant finishes before the next begins; +Trino and Athena run once, not once per cache mode. Legacy `pgwire` catalogs +keep their existing behavior and are not relabeled as uncached history. + +Only the perf driver changes cache settings. It pins its PGWire connection and +applies `SET GLOBAL` before that variant's first warmup and outside query timing: + +| Setting | Uncached baseline | Cached | +| --- | --- | --- | +| `enable_external_file_cache` | `false` | `true` | +| `parquet_metadata_cache` | `false` | `false` | +| `enable_http_metadata_cache` | `false` | `false` | + +The cached variant matches DuckDB's default remote-file caching policy rather +than enabling metadata caches that ordinary workers leave off. Setup is lazy, +so the cached driver's construction cannot enable caching during the uncached +phase. A lost pinned connection fails the query instead of silently connecting +to an unconfigured worker. Worker startup and production configuration are +unchanged. + +Use the isolated scenario workflow for this comparison. Local +`just scenario-frozen-perf` runs require a dedicated test warehouse: these +settings affect the underlying worker globally, and the two variants must not +run concurrently against a shared worker. The scenario's sequential phases +also work when connections reuse the same worker. If cache setup fails, inspect +the artifact errors and recreate the test stack; do not publish fallback runs +under either explicit cache-mode label. Teardown removes the test workers. + +Warmup does not imply that the entire dataset fits in memory. Query-local +buffering, prefetching, DuckLake catalog caching, and separate cache +proxies/extensions are unaffected, so "uncached" here is not a fully cold +end-to-end read path. The paired query and intent IDs use `balanced_v4` to +separate this methodology from `balanced_v3`; the dataset is unchanged. + ## Paired Query Catalogs Existing catalogs continue to use `queries:` unchanged. A catalog may contain diff --git a/tests/perf/core/catalog.go b/tests/perf/core/catalog.go index 78a83acd..1a035598 100644 --- a/tests/perf/core/catalog.go +++ b/tests/perf/core/catalog.go @@ -389,7 +389,7 @@ func validateCatalog(c Catalog) error { seenTargets := map[Protocol]struct{}{} for _, target := range c.Targets { switch target { - case ProtocolPGWire, ProtocolTrino, ProtocolAthena: + case ProtocolPGWire, ProtocolPGWireUncached, ProtocolPGWireCached, ProtocolTrino, ProtocolAthena: default: return fmt.Errorf("unsupported target protocol %q", target) } diff --git a/tests/perf/core/catalog_test.go b/tests/perf/core/catalog_test.go index c700297f..318547ab 100644 --- a/tests/perf/core/catalog_test.go +++ b/tests/perf/core/catalog_test.go @@ -26,7 +26,7 @@ func TestCheckedInCatalogsLoad(t *testing.T) { } wantTargets := []Protocol{ProtocolPGWire} if filepath.Base(path) == "ducklake_posthog_tables.yaml" { - wantTargets = []Protocol{ProtocolPGWire, ProtocolTrino, ProtocolAthena} + wantTargets = []Protocol{ProtocolPGWireUncached, ProtocolPGWireCached, ProtocolTrino, ProtocolAthena} } if !reflect.DeepEqual(catalog.Targets, wantTargets) { t.Fatalf("catalog targets = %v, want %v", catalog.Targets, wantTargets) @@ -49,27 +49,27 @@ func TestCheckedInPostHogCatalogPublishesCompleteStablePairs(t *testing.T) { t.Fatalf("LoadCatalog: %v", err) } want := []string{ - "q_events_total_balanced_v3__raw_view", - "q_events_total_balanced_v3__ducklake_table", - "q_events_total_balanced_v3__athena_external", - "q_events_count_one_day_balanced_v3__raw_view", - "q_events_count_one_day_balanced_v3__ducklake_table", - "q_events_count_one_day_balanced_v3__athena_external", - "q_events_by_name_march_2026_balanced_v3__raw_view", - "q_events_by_name_march_2026_balanced_v3__ducklake_table", - "q_events_by_name_march_2026_balanced_v3__athena_external", - "q_events_distinct_persons_balanced_v3__raw_view", - "q_events_distinct_persons_balanced_v3__ducklake_table", - "q_events_distinct_persons_balanced_v3__athena_external", - "q_persons_total_balanced_v3__raw_view", - "q_persons_total_balanced_v3__ducklake_table", - "q_persons_total_balanced_v3__athena_external", - "q_persons_daily_april_2026_balanced_v3__raw_view", - "q_persons_daily_april_2026_balanced_v3__ducklake_table", - "q_persons_daily_april_2026_balanced_v3__athena_external", - "q_events_daily_march_2026_balanced_v3__raw_view", - "q_events_daily_march_2026_balanced_v3__ducklake_table", - "q_events_daily_march_2026_balanced_v3__athena_external", + "q_events_total_balanced_v4__raw_view", + "q_events_total_balanced_v4__ducklake_table", + "q_events_total_balanced_v4__athena_external", + "q_events_count_one_day_balanced_v4__raw_view", + "q_events_count_one_day_balanced_v4__ducklake_table", + "q_events_count_one_day_balanced_v4__athena_external", + "q_events_by_name_march_2026_balanced_v4__raw_view", + "q_events_by_name_march_2026_balanced_v4__ducklake_table", + "q_events_by_name_march_2026_balanced_v4__athena_external", + "q_events_distinct_persons_balanced_v4__raw_view", + "q_events_distinct_persons_balanced_v4__ducklake_table", + "q_events_distinct_persons_balanced_v4__athena_external", + "q_persons_total_balanced_v4__raw_view", + "q_persons_total_balanced_v4__ducklake_table", + "q_persons_total_balanced_v4__athena_external", + "q_persons_daily_april_2026_balanced_v4__raw_view", + "q_persons_daily_april_2026_balanced_v4__ducklake_table", + "q_persons_daily_april_2026_balanced_v4__athena_external", + "q_events_daily_march_2026_balanced_v4__raw_view", + "q_events_daily_march_2026_balanced_v4__ducklake_table", + "q_events_daily_march_2026_balanced_v4__athena_external", } if got := queryIDs(catalog); !reflect.DeepEqual(got, want) { t.Fatalf("checked-in PostHog query IDs changed: got %v want %v", got, want) @@ -78,7 +78,7 @@ func TestCheckedInPostHogCatalogPublishesCompleteStablePairs(t *testing.T) { t.Fatalf("checked-in PostHog measure iterations = %d, want 4 for balanced target order", catalog.MeasureIterations) } for _, query := range catalog.Queries { - if !strings.HasSuffix(query.IntentID, "_balanced_v3") { + if !strings.HasSuffix(query.IntentID, "_balanced_v4") { t.Fatalf("checked-in PostHog query %s has unversioned methodology intent %q", query.QueryID, query.IntentID) } if strings.Contains(query.PGWireSQL, "TIMESTAMPTZ '") { diff --git a/tests/perf/core/intent_matcher.go b/tests/perf/core/intent_matcher.go index 6b3b5032..dcd60f82 100644 --- a/tests/perf/core/intent_matcher.go +++ b/tests/perf/core/intent_matcher.go @@ -13,7 +13,7 @@ func NewIntentMatcher() *IntentMatcher { func (m *IntentMatcher) SQLFor(query Query, protocol Protocol) (string, error) { switch protocol { - case ProtocolPGWire, ProtocolTrino: + case ProtocolPGWire, ProtocolPGWireUncached, ProtocolPGWireCached, ProtocolTrino: if strings.TrimSpace(query.CanonicalSQL()) == "" { return "", fmt.Errorf("query %s missing canonical SQL", query.QueryID) } diff --git a/tests/perf/core/intent_matcher_test.go b/tests/perf/core/intent_matcher_test.go index e6672493..01f5faad 100644 --- a/tests/perf/core/intent_matcher_test.go +++ b/tests/perf/core/intent_matcher_test.go @@ -9,7 +9,7 @@ func TestIntentMatcherReturnsCanonicalSQLForBothProtocols(t *testing.T) { IntentID: "i1", PGWireSQL: "SELECT 1", } - for _, protocol := range []Protocol{ProtocolPGWire, ProtocolTrino} { + for _, protocol := range []Protocol{ProtocolPGWire, ProtocolPGWireUncached, ProtocolPGWireCached, ProtocolTrino} { if got, err := m.SQLFor(q, protocol); err != nil || got != "SELECT 1" { t.Fatalf("unexpected %s SQL result: sql=%q err=%v", protocol, got, err) } diff --git a/tests/perf/core/runner.go b/tests/perf/core/runner.go index 3a669e20..bc4e6e2f 100644 --- a/tests/perf/core/runner.go +++ b/tests/perf/core/runner.go @@ -168,9 +168,9 @@ func querySupportsProtocol(query Query, protocol Protocol) bool { case "": return true case StorageTargetRawView: - return protocol == ProtocolPGWire + return protocol == ProtocolPGWire || protocol == ProtocolPGWireUncached || protocol == ProtocolPGWireCached case StorageTargetDuckLakeTable: - return protocol == ProtocolPGWire || protocol == ProtocolTrino + return protocol == ProtocolPGWire || protocol == ProtocolPGWireUncached || protocol == ProtocolPGWireCached || protocol == ProtocolTrino case StorageTargetAthenaExternal: return protocol == ProtocolAthena default: diff --git a/tests/perf/core/runner_test.go b/tests/perf/core/runner_test.go index 1f892eec..fb9c8e8d 100644 --- a/tests/perf/core/runner_test.go +++ b/tests/perf/core/runner_test.go @@ -305,6 +305,56 @@ func TestRunnerRoutesEachStorageVariantOnlyToItsComparableProtocol(t *testing.T) } } +func TestRunnerRunsUncachedAndCachedPGWireAsDistinctComparableResults(t *testing.T) { + uncached, cached := ProtocolPGWireUncached, ProtocolPGWireCached + var events []string + uncachedDriver := &testDriver{protocol: uncached, events: &events} + cachedDriver := &testDriver{protocol: cached, events: &events} + sink := &inMemorySink{} + runner := NewQueryRunner(RunnerConfig{ + RunID: "cache-comparison", + Catalog: Catalog{ + Targets: []Protocol{uncached, cached}, WarmupIterations: 1, MeasureIterations: 1, + Queries: []Query{ + {QueryID: "q__raw_view", IntentID: "intent", StorageTarget: StorageTargetRawView}, + {QueryID: "q__ducklake_table", IntentID: "intent", StorageTarget: StorageTargetDuckLakeTable}, + {QueryID: "q__athena_external", IntentID: "intent", StorageTarget: StorageTargetAthenaExternal}, + }, + }, + Drivers: map[Protocol]ProtocolDriver{uncached: uncachedDriver, cached: cachedDriver}, Sink: sink, + }) + summary, err := runner.Run(context.Background()) + if err != nil { + t.Fatal(err) + } + if summary.RunID != "cache-comparison" || summary.TotalQueries != 4 || summary.WarmupQueries != 4 { + t.Fatalf("unexpected summary: %+v", summary) + } + for _, driver := range []*testDriver{uncachedDriver, cachedDriver} { + if want := []string{"q__raw_view", "q__ducklake_table", "q__raw_view", "q__ducklake_table"}; !reflect.DeepEqual(driver.queryIDs, want) { + t.Fatalf("%s routed queries = %v, want %v", driver.protocol, driver.queryIDs, want) + } + } + wantEvents := []string{ + "pgwire_uncached/q__raw_view", "pgwire_uncached/q__ducklake_table", + "pgwire_uncached/q__raw_view", "pgwire_uncached/q__ducklake_table", + "pgwire_cached/q__raw_view", "pgwire_cached/q__ducklake_table", + "pgwire_cached/q__raw_view", "pgwire_cached/q__ducklake_table", + } + if !reflect.DeepEqual(events, wantEvents) { + t.Fatalf("execution order = %v, want %v", events, wantEvents) + } + for index, result := range sink.results { + wantProtocol := uncached + if index >= 2 { + wantProtocol = cached + } + if result.Protocol != wantProtocol || result.IntentID != "intent" || result.MeasureIteration != 1 { + t.Fatalf("result %d = %+v", index, result) + } + } +} + func (d *testDriver) Close() error { return nil } type inMemorySink struct { diff --git a/tests/perf/core/types.go b/tests/perf/core/types.go index 3aed9a62..9d76306c 100644 --- a/tests/perf/core/types.go +++ b/tests/perf/core/types.go @@ -5,9 +5,11 @@ import "time" type Protocol string const ( - ProtocolPGWire Protocol = "pgwire" - ProtocolTrino Protocol = "trino" - ProtocolAthena Protocol = "athena" + ProtocolPGWire Protocol = "pgwire" + ProtocolPGWireUncached Protocol = "pgwire_uncached" + ProtocolPGWireCached Protocol = "pgwire_cached" + ProtocolTrino Protocol = "trino" + ProtocolAthena Protocol = "athena" ) // StorageTarget identifies the physical relation family selected for a paired diff --git a/tests/perf/drivers/pgwire/driver.go b/tests/perf/drivers/pgwire/driver.go index 329d197c..78e0f598 100644 --- a/tests/perf/drivers/pgwire/driver.go +++ b/tests/perf/drivers/pgwire/driver.go @@ -3,6 +3,7 @@ package pgwire import ( "context" "database/sql" + "errors" "fmt" "time" @@ -16,17 +17,41 @@ type Executor interface { } type Driver struct { - exec Executor + exec Executor + protocol core.Protocol + prepare func(context.Context) error } func NewWithExecutor(exec Executor) *Driver { - return &Driver{exec: exec} + return &Driver{exec: exec, protocol: core.ProtocolPGWire} } func NewWithDB(db *sql.DB) *Driver { return NewWithExecutor(&sqlExecutor{db: db}) } +// NewWithDBAndProtocol takes ownership of db on success. Explicit cache variants +// pin their connection and apply settings lazily, before the first warmup query. +// Creating a later variant must not change the worker used by an earlier phase. +func NewWithDBAndProtocol(db *sql.DB, protocol core.Protocol) (*Driver, error) { + if protocol == core.ProtocolPGWire { + return NewWithDB(db), nil + } + if protocol != core.ProtocolPGWireUncached && protocol != core.ProtocolPGWireCached { + return nil, fmt.Errorf("unsupported pgwire protocol %q", protocol) + } + externalCache := protocol == core.ProtocolPGWireCached + exec := &sqlExecutor{ + db: db, + settings: []string{ + fmt.Sprintf("SET GLOBAL enable_external_file_cache = %t", externalCache), + "SET GLOBAL parquet_metadata_cache = false", + "SET GLOBAL enable_http_metadata_cache = false", + }, + } + return &Driver{exec: exec, protocol: protocol, prepare: exec.prepare}, nil +} + func NewFromDSN(dsn string) (*Driver, error) { db, err := sql.Open("postgres", dsn) if err != nil { @@ -36,7 +61,7 @@ func NewFromDSN(dsn string) (*Driver, error) { } func (d *Driver) Protocol() core.Protocol { - return core.ProtocolPGWire + return d.protocol } func (d *Driver) Execute(ctx context.Context, query core.Query, args []any) (core.ExecutionResult, error) { @@ -47,6 +72,11 @@ func (d *Driver) Execute(ctx context.Context, query core.Query, args []any) (cor if sqlText == "" { return core.ExecutionResult{}, fmt.Errorf("query %s missing pgwire_sql", query.QueryID) } + if d.prepare != nil { + if err := d.prepare(ctx); err != nil { + return core.ExecutionResult{}, err + } + } started := time.Now() rows, err := d.exec.Execute(ctx, sqlText, args) return core.ExecutionResult{ @@ -63,11 +93,41 @@ func (d *Driver) Close() error { } type sqlExecutor struct { - db *sql.DB + db *sql.DB + conn *sql.Conn + settings []string + prepared bool + prepareErr error +} + +func (e *sqlExecutor) prepare(ctx context.Context) error { + if e.prepared { + return e.prepareErr + } + e.prepared = true + conn, err := e.db.Conn(ctx) + if err != nil { + e.prepareErr = fmt.Errorf("pin pgwire cache-variant connection: %w", err) + return e.prepareErr + } + e.conn = conn + for _, setting := range e.settings { + if _, err := conn.ExecContext(ctx, setting); err != nil { + e.prepareErr = fmt.Errorf("configure pgwire cache variant (%s): %w", setting, err) + return e.prepareErr + } + } + return nil } func (e *sqlExecutor) Execute(ctx context.Context, query string, args []any) (int64, error) { - rows, err := e.db.QueryContext(ctx, query, args...) + queryContext := e.db.QueryContext + execContext := e.db.ExecContext + if e.conn != nil { + queryContext = e.conn.QueryContext + execContext = e.conn.ExecContext + } + rows, err := queryContext(ctx, query, args...) if err == nil { defer func() { _ = rows.Close() @@ -94,7 +154,7 @@ func (e *sqlExecutor) Execute(ctx context.Context, query string, args []any) (in return count, nil } - res, execErr := e.db.ExecContext(ctx, query, args...) + res, execErr := execContext(ctx, query, args...) if execErr != nil { return 0, execErr } @@ -106,5 +166,12 @@ func (e *sqlExecutor) Execute(ctx context.Context, query string, args []any) (in } func (e *sqlExecutor) Close() error { - return e.db.Close() + var connErr error + if e.conn != nil { + connErr = e.conn.Close() + if errors.Is(connErr, sql.ErrConnDone) { + connErr = nil + } + } + return errors.Join(connErr, e.db.Close()) } diff --git a/tests/perf/drivers/pgwire/driver_test.go b/tests/perf/drivers/pgwire/driver_test.go index 93811612..66cc68f2 100644 --- a/tests/perf/drivers/pgwire/driver_test.go +++ b/tests/perf/drivers/pgwire/driver_test.go @@ -2,7 +2,14 @@ package pgwire import ( "context" + "database/sql" + "database/sql/driver" + "errors" + "io" + "reflect" + "strings" "testing" + "time" "github.com/posthog/duckgres/tests/perf/core" ) @@ -33,3 +40,186 @@ func TestDriverUsesCanonicalRenderedSQL(t *testing.T) { t.Fatalf("expected canonical rendered SQL, got %q", exec.lastQuery) } } + +func TestCacheVariantsPrepareLazilyOnPinnedConnection(t *testing.T) { + for _, tc := range []struct { + protocol core.Protocol + external string + }{ + {core.ProtocolPGWireUncached, "false"}, + {core.ProtocolPGWireCached, "true"}, + } { + t.Run(string(tc.protocol), func(t *testing.T) { + connector := &cacheTestConnector{setupDelay: 10 * time.Millisecond} + db := sql.OpenDB(connector) + // Without pinning, each operation would close its connection. + db.SetMaxIdleConns(0) + d, err := NewWithDBAndProtocol(db, tc.protocol) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = d.Close() }) + if d.Protocol() != tc.protocol || connector.connections != 0 { + t.Fatalf("constructor must retain protocol without connecting: protocol=%s connections=%d", d.Protocol(), connector.connections) + } + + started := time.Now() + result, err := d.Execute(context.Background(), core.Query{QueryID: "warmup", PGWireSQL: "SELECT 1"}, nil) + elapsed := time.Since(started) + if err != nil || result.Rows != 1 { + t.Fatalf("warmup result=%+v error=%v", result, err) + } + if elapsed-result.Duration < 3*connector.setupDelay { + t.Fatalf("cache setup was included in query duration: elapsed=%s query=%s", elapsed, result.Duration) + } + if _, err := d.Execute(context.Background(), core.Query{QueryID: "measure", PGWireSQL: "SELECT 2"}, nil); err != nil { + t.Fatal(err) + } + want := []string{ + "SET GLOBAL enable_external_file_cache = " + tc.external, + "SET GLOBAL parquet_metadata_cache = false", + "SET GLOBAL enable_http_metadata_cache = false", + "SELECT 1", "SELECT 2", + } + if !reflect.DeepEqual(connector.statements, want) || connector.connections != 1 { + t.Fatalf("all setup and queries must use one connection: statements=%v connections=%d", connector.statements, connector.connections) + } + if err := d.Close(); err != nil { + t.Fatal(err) + } + if connector.closes != 1 { + t.Fatalf("Close did not release pinned connection: closes=%d", connector.closes) + } + if err := db.Ping(); err == nil { + t.Fatal("Close must also close the owned DB") + } + }) + } +} + +func TestCacheVariantFailsClosedOnSetupError(t *testing.T) { + connector := &cacheTestConnector{failStatement: "SET GLOBAL parquet_metadata_cache = false"} + d, err := NewWithDBAndProtocol(sql.OpenDB(connector), core.ProtocolPGWireUncached) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = d.Close() }) + for range 2 { + _, err := d.Execute(context.Background(), core.Query{QueryID: "q", PGWireSQL: "SELECT 1"}, nil) + if err == nil || !strings.Contains(err.Error(), "parquet_metadata_cache") { + t.Fatalf("expected contextual cache setup error, got %v", err) + } + } + want := []string{"SET GLOBAL enable_external_file_cache = false", "SET GLOBAL parquet_metadata_cache = false"} + if !reflect.DeepEqual(connector.statements, want) { + t.Fatalf("failed setup must never execute queries or retry: %v", connector.statements) + } +} + +func TestCacheVariantDoesNotReconnectAfterLostConnection(t *testing.T) { + connector := &cacheTestConnector{loseConnection: true} + d, err := NewWithDBAndProtocol(sql.OpenDB(connector), core.ProtocolPGWireUncached) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = d.Close() }) + for range 2 { + if _, err := d.Execute(context.Background(), core.Query{QueryID: "q", PGWireSQL: "SELECT 1"}, nil); err == nil { + t.Fatal("lost pinned connection must fail rather than run on an unconfigured replacement") + } + } + if connector.connections != 1 || len(connector.statements) != 4 { + t.Fatalf("unexpected replacement connection or query retry: connections=%d statements=%v", connector.connections, connector.statements) + } +} + +func TestDefaultDriverDoesNotChangeCacheSettings(t *testing.T) { + connector := &cacheTestConnector{} + d := NewWithDB(sql.OpenDB(connector)) + t.Cleanup(func() { _ = d.Close() }) + if _, err := d.Execute(context.Background(), core.Query{QueryID: "q", PGWireSQL: "SELECT 1"}, nil); err != nil { + t.Fatal(err) + } + if d.Protocol() != core.ProtocolPGWire || !reflect.DeepEqual(connector.statements, []string{"SELECT 1"}) { + t.Fatalf("default pgwire behavior changed: protocol=%s statements=%v", d.Protocol(), connector.statements) + } +} + +func TestCacheVariantRejectsUnsupportedProtocol(t *testing.T) { + connector := &cacheTestConnector{} + db := sql.OpenDB(connector) + t.Cleanup(func() { _ = db.Close() }) + if _, err := NewWithDBAndProtocol(db, core.ProtocolTrino); err == nil { + t.Fatal("expected unsupported protocol error") + } + if connector.connections != 0 { + t.Fatal("invalid protocol must not open a connection") + } +} + +type cacheTestConnector struct { + connections int + closes int + statements []string + setupDelay time.Duration + failStatement string + loseConnection bool +} + +func (c *cacheTestConnector) Connect(context.Context) (driver.Conn, error) { + c.connections++ + return &cacheTestConn{connector: c}, nil +} + +func (c *cacheTestConnector) Driver() driver.Driver { return cacheTestDriver{} } + +type cacheTestDriver struct{} + +func (cacheTestDriver) Open(string) (driver.Conn, error) { + return nil, errors.New("use connector") +} + +type cacheTestConn struct{ connector *cacheTestConnector } + +func (c *cacheTestConn) Prepare(string) (driver.Stmt, error) { + return nil, errors.New("unexpected prepare") +} + +func (c *cacheTestConn) Close() error { + c.connector.closes++ + return nil +} + +func (c *cacheTestConn) Begin() (driver.Tx, error) { + return nil, errors.New("unexpected transaction") +} + +func (c *cacheTestConn) ExecContext(_ context.Context, query string, _ []driver.NamedValue) (driver.Result, error) { + c.connector.statements = append(c.connector.statements, query) + time.Sleep(c.connector.setupDelay) + if query == c.connector.failStatement { + return nil, errors.New("setting rejected") + } + return driver.RowsAffected(0), nil +} + +func (c *cacheTestConn) QueryContext(_ context.Context, query string, _ []driver.NamedValue) (driver.Rows, error) { + c.connector.statements = append(c.connector.statements, query) + if c.connector.loseConnection { + return nil, driver.ErrBadConn + } + return &cacheTestRows{}, nil +} + +type cacheTestRows struct{ read bool } + +func (*cacheTestRows) Columns() []string { return []string{"value"} } +func (*cacheTestRows) Close() error { return nil } +func (r *cacheTestRows) Next(values []driver.Value) error { + if r.read { + return io.EOF + } + r.read = true + values[0] = int64(1) + return nil +} diff --git a/tests/perf/queries/ducklake_posthog_tables.yaml b/tests/perf/queries/ducklake_posthog_tables.yaml index 4a6e9877..d1be73a9 100644 --- a/tests/perf/queries/ducklake_posthog_tables.yaml +++ b/tests/perf/queries/ducklake_posthog_tables.yaml @@ -1,9 +1,10 @@ -name: posthog-frozen-ducklake-golden-v3 +name: posthog-frozen-ducklake-golden-v4 description: Identical query shapes with protocol-isolated, balanced execution order over frozen raw Parquet views and production-shaped DuckLake tables. seed: 42 dataset_scale: 1 targets: - - pgwire + - pgwire_uncached + - pgwire_cached - trino - athena warmup_iterations: 1 @@ -21,14 +22,14 @@ relation_variants: persons: persons paired_queries: - - query_id_base: q_events_total_balanced_v3 - intent_id: intent_events_total_balanced_v3 + - query_id_base: q_events_total_balanced_v4 + intent_id: intent_events_total_balanced_v4 tags: [nightly, frozen, posthog, events, aggregate, paired] params: {} sql_template: SELECT COUNT(*) AS events FROM {{ relation "events" }} - - query_id_base: q_events_count_one_day_balanced_v3 - intent_id: intent_events_count_one_day_balanced_v3 + - query_id_base: q_events_count_one_day_balanced_v4 + intent_id: intent_events_count_one_day_balanced_v4 tags: [nightly, frozen, posthog, events, aggregate, one-day, paired] params: {} sql_template: > @@ -37,32 +38,32 @@ paired_queries: WHERE "timestamp" >= CAST('2026-03-01 00:00:00+00:00' AS TIMESTAMP WITH TIME ZONE) AND "timestamp" < CAST('2026-03-02 00:00:00+00:00' AS TIMESTAMP WITH TIME ZONE) - - query_id_base: q_events_by_name_march_2026_balanced_v3 - intent_id: intent_events_by_name_march_2026_balanced_v3 + - query_id_base: q_events_by_name_march_2026_balanced_v4 + intent_id: intent_events_by_name_march_2026_balanced_v4 tags: [nightly, frozen, posthog, events, aggregate, analytics, paired] params: {} sql_template: SELECT event, COUNT(*) AS events FROM {{ relation "events" }} WHERE "timestamp" >= CAST('2026-03-01 00:00:00+00:00' AS TIMESTAMP WITH TIME ZONE) AND "timestamp" < CAST('2026-03-18 00:00:00+00:00' AS TIMESTAMP WITH TIME ZONE) GROUP BY event ORDER BY events DESC, event LIMIT 20 - - query_id_base: q_events_distinct_persons_balanced_v3 - intent_id: intent_events_distinct_persons_balanced_v3 + - query_id_base: q_events_distinct_persons_balanced_v4 + intent_id: intent_events_distinct_persons_balanced_v4 tags: [nightly, frozen, posthog, events, aggregate, distinct, paired] params: {} sql_template: SELECT COUNT(DISTINCT person_id) AS distinct_persons FROM {{ relation "events" }} WHERE person_id IS NOT NULL - - query_id_base: q_persons_total_balanced_v3 - intent_id: intent_persons_total_balanced_v3 + - query_id_base: q_persons_total_balanced_v4 + intent_id: intent_persons_total_balanced_v4 tags: [nightly, frozen, posthog, persons, aggregate, paired] params: {} sql_template: SELECT COUNT(*) AS persons FROM {{ relation "persons" }} - - query_id_base: q_persons_daily_april_2026_balanced_v3 - intent_id: intent_persons_daily_april_2026_balanced_v3 + - query_id_base: q_persons_daily_april_2026_balanced_v4 + intent_id: intent_persons_daily_april_2026_balanced_v4 tags: [nightly, frozen, posthog, persons, aggregate, time-series, paired] params: {} sql_template: SELECT date_trunc('day', _timestamp) AS day, COUNT(*) AS persons FROM {{ relation "persons" }} WHERE _timestamp >= CAST('2026-04-01 00:00:00+00:00' AS TIMESTAMP WITH TIME ZONE) AND _timestamp < CAST('2026-05-01 00:00:00+00:00' AS TIMESTAMP WITH TIME ZONE) GROUP BY 1 ORDER BY 1 - - query_id_base: q_events_daily_march_2026_balanced_v3 - intent_id: intent_events_daily_march_2026_balanced_v3 + - query_id_base: q_events_daily_march_2026_balanced_v4 + intent_id: intent_events_daily_march_2026_balanced_v4 tags: [nightly, frozen, posthog, events, aggregate, time-series, paired] params: {} sql_template: SELECT date_trunc('day', "timestamp") AS day, COUNT(*) AS events FROM {{ relation "events" }} WHERE "timestamp" >= CAST('2026-03-01 00:00:00+00:00' AS TIMESTAMP WITH TIME ZONE) AND "timestamp" < CAST('2026-03-18 00:00:00+00:00' AS TIMESTAMP WITH TIME ZONE) GROUP BY 1 ORDER BY 1