From e5125e780ecbf3a6fb7179f2e9a5c123187170db Mon Sep 17 00:00:00 2001 From: Bill Guowei Yang Date: Thu, 3 Sep 2026 17:18:34 -0400 Subject: [PATCH 1/4] feat(perf): compare frozen queries with Athena --- .github/workflows/scenario-dev.yml | 4 + docs/runbooks/scenario-dev.md | 39 ++- docs/runbooks/scenario-runner.md | 11 +- go.mod | 9 +- go.sum | 18 +- tests/mw-dev/README.md | 22 +- tests/mw-dev/manifests.tmpl.yaml | 6 + tests/mw-dev/run.sh | 18 ++ tests/mw-dev/run_sh_test.go | 26 ++ tests/mw-dev/scenario/perf/adapter_test.go | 62 ++++- tests/mw-dev/scenario/perf/steps.go | 98 +++++-- tests/mw-dev/scenario/runner_test.go | 22 +- .../scenarios/posthog_frozen_perf.yaml | 13 +- tests/mw-dev/scenario/script_test.go | 8 + tests/perf/README.md | 29 +- tests/perf/core/catalog.go | 45 +++- tests/perf/core/catalog_test.go | 87 +++++- tests/perf/core/runner.go | 23 +- tests/perf/core/runner_test.go | 43 +++ tests/perf/core/sink.go | 89 ++++++- tests/perf/core/sink_test.go | 64 +++++ tests/perf/core/types.go | 46 +++- tests/perf/drivers/athena/aws_client.go | 123 +++++++++ tests/perf/drivers/athena/driver.go | 247 ++++++++++++++++++ tests/perf/drivers/athena/driver_test.go | 215 +++++++++++++++ tests/perf/harness_test.go | 2 +- .../perf/queries/ducklake_posthog_tables.yaml | 4 + 27 files changed, 1268 insertions(+), 105 deletions(-) create mode 100644 tests/perf/drivers/athena/aws_client.go create mode 100644 tests/perf/drivers/athena/driver.go create mode 100644 tests/perf/drivers/athena/driver_test.go diff --git a/.github/workflows/scenario-dev.yml b/.github/workflows/scenario-dev.yml index 8e856e077..b9f1c96d6 100644 --- a/.github/workflows/scenario-dev.yml +++ b/.github/workflows/scenario-dev.yml @@ -74,6 +74,10 @@ jobs: # Frozen perf uses a namespace-local Trino cell. The dedicated role is # consumed only when E2E_SUITE selects that isolated deployment. TRINO_POD_IDENTITY_ROLE: ${{ secrets.MW_DEV_TRINO_POD_IDENTITY_ROLE }} + SCENARIO_POD_IDENTITY_ROLE: ${{ secrets.MW_DEV_ATHENA_POD_IDENTITY_ROLE }} + DUCKGRES_SCENARIO_ATHENA_WORKGROUP: ${{ vars.MW_DEV_ATHENA_WORKGROUP }} + DUCKGRES_SCENARIO_ATHENA_DATABASE: ${{ vars.MW_DEV_ATHENA_DATABASE }} + DUCKGRES_SCENARIO_ATHENA_RESULTS_S3_URI: ${{ vars.MW_DEV_ATHENA_RESULTS_S3_URI }} TRINO_IMAGE: ghcr.io/posthog/trino:4505364c570d6b51edecd299b603fca4b6693d86@sha256:ac80c275fd18a439d25da5652ab5cd3c80bcbdd2d88d64c9722dc3e8bb68ba07 E2E_SUITE: ${{ (github.event_name == 'schedule' || inputs.scenario == 'posthog_frozen_perf') && 'trino' || 'neutral' }} PR_NUMBER: ${{ github.run_id }} diff --git a/docs/runbooks/scenario-dev.md b/docs/runbooks/scenario-dev.md index 78448eaa5..a15b28d9f 100644 --- a/docs/runbooks/scenario-dev.md +++ b/docs/runbooks/scenario-dev.md @@ -8,17 +8,18 @@ identity, diagnostics, and teardown. ## Scheduled Runs The daily scheduled trigger runs `posthog_frozen_perf`, which creates -production-shaped DuckLake tables from the frozen dataset, executes paired -raw-view and DuckLake-table queries, and publishes the comparison history. +production-shaped DuckLake tables from the frozen dataset, executes equivalent +PGWire, Trino, and on-demand Athena queries, and publishes the comparison history. Scheduled runs no longer execute `full-suite`. Manual runs remain selectable and default to `full-suite`, which covers the frozen dataset metadata, legacy perf, and dbt workloads with one shared warehouse lifecycle. -Scenario jobs override the shared harness's default worker request to 2 CPU and -8Gi memory for the frozen pgwire perf workload. This adds process headroom for -repeated full-dataset aggregates without increasing DuckDB's thread count. The +Scenario jobs override the shared harness's default worker request to 3 CPU and +12Gi memory for the frozen PGWire perf workload. The isolated Trino workers +use the same aggregate execution resources. This adds process headroom for +repeated full-dataset aggregates without changing DuckDB's thread count. The shared e2e workflow keeps the harness defaults documented in `tests/mw-dev/README.md`. @@ -40,12 +41,24 @@ Configure these repository secrets: - `AWS_ECR_PUBLISH_IAM_ROLE` - `MW_DEV_ACCOUNT_ID` +- `MW_DEV_TRINO_POD_IDENTITY_ROLE` +- `MW_DEV_ATHENA_POD_IDENTITY_ROLE` Configure these repository variables: - `TS_WIF_CLIENT_ID_MW_DEV` - `TS_WIF_AUDIENCE_MW_DEV` - `MW_DEV_SCENARIO_PERF_SECRET_ID` +- `MW_DEV_ATHENA_WORKGROUP` +- `MW_DEV_ATHENA_DATABASE` +- `MW_DEV_ATHENA_RESULTS_S3_URI` + +The Athena values must identify an engine-v3 on-demand workgroup, a Glue +database with explicit `events` and `persons` external tables over the frozen +Parquet prefixes, and an encrypted S3 result root ending in `/`. Disable +workgroup result reuse and configure lifecycle expiry on the result root. The +scenario still disables result reuse on every query and validates that the +effective result location remains under the configured root. `MW_DEV_SCENARIO_PERF_SECRET_ID` names the AWS Secrets Manager JSON secret used only by historical perf publishing. The secret must contain `host`, `port`, @@ -61,6 +74,14 @@ The harness also injects `DUCKGRES_SCENARIO_ORG_ID=ci-pr--cnpg` for successful CNPG scenarios. That name matches the Crossplane composition's exact CI credential ownership boundary and the harness's cleanup convention. +The scenario Job uses the dedicated `duckgres-scenario` ServiceAccount. Its +Pod Identity role needs only Athena query APIs for the configured workgroup, +read-only Glue metadata access to the benchmark database, read/list access to +the frozen dataset, and result-prefix access (plus the configured KMS key when +SSE-KMS is used). The workflow deployer needs `iam:PassRole`/`iam:GetRole` for +that role. Teardown and the stale-run sweep already delete every Pod Identity +association in the disposable namespace. + For a direct local run with `just scenario`, set `DUCKGRES_SCENARIO_ORG_ID` explicitly to an authorized, disposable org id. The local runner has no default because reusing a fixed warehouse identity can @@ -83,6 +104,8 @@ uploadable in a visible `*.partial/` directory with an published in the GitHub Actions job summary. - `step_results.csv` contains per-step duration, status, and error class. - `events.jsonl` contains the execution timeline. +- `perf/query_service_metrics.csv` separates Athena queue/planning/engine time + and records bytes scanned for cost analysis. Perf query errors mark the `perf_queries` DAG step failed. They do not stop independent sibling branches: for example, `dbt_models` still runs because it @@ -104,3 +127,9 @@ If cleanup did not complete, the scenario-created org is control plane. Also inspect the workflow diagnostics from `tests/mw-dev/run.sh diagnostics`; the namespace is `duckgres-ci-pr-`. + +If an Athena query times out, the driver calls `StopQueryExecution`. A hard pod +termination can still leave result objects, so S3 lifecycle expiry is the +reliable cleanup backstop. Check the Athena query history for scan-limit or +catalog errors, confirm the workgroup did not override results outside the +configured root, and rerun after fixing the private Glue/IAM configuration. diff --git a/docs/runbooks/scenario-runner.md b/docs/runbooks/scenario-runner.md index 0b05f7330..a84807894 100644 --- a/docs/runbooks/scenario-runner.md +++ b/docs/runbooks/scenario-runner.md @@ -54,11 +54,16 @@ shape deployed with the scenario stack: ```bash export DUCKGRES_K8S_WORKER_CPU_REQUEST="3" export DUCKGRES_K8S_WORKER_MEMORY_REQUEST="12Gi" +export DUCKGRES_SCENARIO_ATHENA_REGION="" +export DUCKGRES_SCENARIO_ATHENA_WORKGROUP="" +export DUCKGRES_SCENARIO_ATHENA_DATABASE="" +export DUCKGRES_SCENARIO_ATHENA_RESULTS_S3_URI="s3:////" ``` The full and fast suites exercise PGWire only. The targeted frozen perf -scenario compares PGWire and Trino. It records per-query success and failure -rows in `query_results.csv`. +scenario compares PGWire, Trino, and on-demand Athena. It records per-query +success and failure rows in `query_results.csv` and Athena service details in +`query_service_metrics.csv`. Measured query errors fail the perf DAG step after its artifacts are written; independent sibling steps continue to run. @@ -150,7 +155,7 @@ The frozen perf scenario uses: - `tests/mw-dev/scenario/scenarios/posthog_frozen_perf.yaml` - `tests/perf/queries/ducklake_frozen.yaml` -Perf artifacts are written under `artifacts/scenario//perf/` using the existing `tests/perf/core` artifact schema, including `query_results.csv`, `summary.json`, and `server_metrics.prom`. +Perf artifacts are written under `artifacts/scenario//perf/` using the existing `tests/perf/core` artifact schema, including `query_results.csv`, `query_service_metrics.csv`, `summary.json`, and `server_metrics.prom`. The frozen dbt scenario uses: diff --git a/go.mod b/go.mod index 4a3df31e5..2e2ebe9c1 100644 --- a/go.mod +++ b/go.mod @@ -4,10 +4,11 @@ go 1.25.4 require ( github.com/apache/arrow-go/v18 v18.5.1 - github.com/aws/aws-sdk-go-v2 v1.42.1 + github.com/aws/aws-sdk-go-v2 v1.45.1 github.com/aws/aws-sdk-go-v2/config v1.32.29 github.com/aws/aws-sdk-go-v2/credentials v1.19.28 github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.32 + github.com/aws/aws-sdk-go-v2/service/athena v1.64.0 github.com/aws/aws-sdk-go-v2/service/route53 v1.62.3 github.com/aws/aws-sdk-go-v2/service/s3 v1.105.0 github.com/aws/aws-sdk-go-v2/service/sts v1.44.0 @@ -52,8 +53,8 @@ require ( github.com/agnivade/levenshtein v1.2.1 // indirect github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.14 // indirect github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.30 // indirect - github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.30 // indirect - github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.30 // indirect + github.com/aws/aws-sdk-go-v2/internal/configsources v1.5.1 // indirect + github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.8.1 // indirect github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.31 // indirect github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.13 // indirect github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.23 // indirect @@ -62,7 +63,7 @@ require ( github.com/aws/aws-sdk-go-v2/service/signin v1.4.0 // indirect github.com/aws/aws-sdk-go-v2/service/sso v1.32.0 // indirect github.com/aws/aws-sdk-go-v2/service/ssooidc v1.37.0 // indirect - github.com/aws/smithy-go v1.27.3 // indirect + github.com/aws/smithy-go v1.28.1 // indirect github.com/beorn7/perks v1.0.1 // indirect github.com/bytedance/gopkg v0.1.3 // indirect github.com/bytedance/sonic v1.15.0 // indirect diff --git a/go.sum b/go.sum index a3ecd2093..d0078308a 100644 --- a/go.sum +++ b/go.sum @@ -28,8 +28,8 @@ github.com/arbovm/levenshtein v0.0.0-20160628152529-48b4e1c0c4d0 h1:jfIu9sQUG6Ig github.com/arbovm/levenshtein v0.0.0-20160628152529-48b4e1c0c4d0/go.mod h1:t2tdKJDJF9BV14lnkjHmOQgcvEKgtqs5a1N3LNdJhGE= github.com/aws/aws-sdk-go v1.55.8 h1:JRmEUbU52aJQZ2AjX4q4Wu7t4uZjOu71uyNmaWlUkJQ= github.com/aws/aws-sdk-go v1.55.8/go.mod h1:ZkViS9AqA6otK+JBBNH2++sx1sgxrPKcSzPPvQkUtXk= -github.com/aws/aws-sdk-go-v2 v1.42.1 h1:9eOTgu1z/dVtYpNZ3/8/XbbaX0x/BqE3HUzAzs6K0ek= -github.com/aws/aws-sdk-go-v2 v1.42.1/go.mod h1:5pKeft2eJj+gElQ38Jqg4ibCqh+/AK33/0X3hip7IjM= +github.com/aws/aws-sdk-go-v2 v1.45.1 h1:iIoG3NaLhV6UZpPXyPXlDj2I9oS8tV/nMcMnITCC6Ks= +github.com/aws/aws-sdk-go-v2 v1.45.1/go.mod h1:bttEH6JqnUL8LepvDVfdrds/fZ5bCIxzpe3abyUrhDU= github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.14 h1:3IZY0XAJquT3aHzbkHfPzy4ACPcEjVG0x87KOwtpqGY= github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.14/go.mod h1:zwM6veDkhGgQFqkBy+uT28AAYpLu+uFMlPl+rCg/73E= github.com/aws/aws-sdk-go-v2/config v1.32.29 h1:BcMHHnpiWKogf+gGfpj3K1w+Sktz29XDo/cPSAPO3FU= @@ -40,12 +40,14 @@ github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.30 h1:/hi1JADLEW9YYryEz1w4GQ github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.30/go.mod h1:/3AOgy4K17Dm4ucMZVC/MJkzy5kmfKUcINRHZyo0koQ= github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.32 h1:kUb0wd0/NfYv2RDoDfogxBy/Hevby5yLIL12iCcs0hY= github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.32/go.mod h1:+rUx79uJZfEavbKROY+U1ez7bUMie7oOAdrjgL0C1FQ= -github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.30 h1:xM/Is9cKMHa8Jj8zkvWhvrFkZsXJV9E+BB4g0HW0duQ= -github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.30/go.mod h1:WueJeNDZvK1fMYEWJIkcivBfEzUkTpBhzlrUKKY8EuA= -github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.30 h1:jn46zC9LdsVR/ZpMIJqMqb8hHv31BlLx3ulVqNspUOk= -github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.30/go.mod h1:1hTMsAgbdS/AtUi4bw8+gUuh1pceo+eXRLfpSuSQj3M= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.5.1 h1:pc138gM1CW+XPc60rEwUlwwuwWFQK16CI1T7v1F9Oec= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.5.1/go.mod h1:1+koxpPIbfBdfzP6vojm5/zTpTQ/micYwlxIiNB3TxI= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.8.1 h1:K0JsbZQj+1h208Ro1zHeA4l7bMp0NvRffHQ91q8Ol1s= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.8.1/go.mod h1:W3/vL6EtCIatICGy9ab29QhMuae+cOKPWcMxv02CO+Q= github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.31 h1:3GUprIsfmGcC5SACIyB0e7E0BM1O1b3Erl5CePYIAeQ= github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.31/go.mod h1:7PuV1yl5e2xnUbm+RqvVg5i2iBM8EyijZNoI9wsOoOc= +github.com/aws/aws-sdk-go-v2/service/athena v1.64.0 h1:HnHMA98zWRuaRdIYa/fvPCHeEQbJykVsXfiK74x6v2o= +github.com/aws/aws-sdk-go-v2/service/athena v1.64.0/go.mod h1:PvOgT+mAEoXeoG7Nun4D1h8jiJ21dH3ms80DZFrw+hU= github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.13 h1:mbRIur/BiHK6SKPjoBIXSE/hJ6g6JGRLuxQy1jGjlN4= github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.13/go.mod h1:ITg9em2KbJx1s0y4aqRX5OYWG6HBZ5TVR//OdpEZ2CQ= github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.23 h1:9Fjh6fi/U5JEStVZijmaMpUwE/gvBJj7x2B/PjbO9To= @@ -66,8 +68,8 @@ github.com/aws/aws-sdk-go-v2/service/ssooidc v1.37.0 h1:fpOlDPI55HdszaxapEGk6HsG github.com/aws/aws-sdk-go-v2/service/ssooidc v1.37.0/go.mod h1:DMPWJBjYs6+3+f/qhBFEFPPlQ6NlhWjai3dJNvipJ84= github.com/aws/aws-sdk-go-v2/service/sts v1.44.0 h1:bLZ0PolJ8J+HkJHztcXORUpHXBye2U8298lCEMi6ZCU= github.com/aws/aws-sdk-go-v2/service/sts v1.44.0/go.mod h1:9gdl4RrflIdpDb2TlXshWgR1F9TeCkvqDx77Vpr4Z/Q= -github.com/aws/smithy-go v1.27.3 h1:F3Zb497UhhskkfpJmfkXswyo+t0sh9OTBnIHjogWbVY= -github.com/aws/smithy-go v1.27.3/go.mod h1:YE2RhdIuDbA5E5bTdciG9KrW3+TiEONeUWCqxX9i1Fc= +github.com/aws/smithy-go v1.28.1 h1:R/nXH00c8qcfCzQVELtRw+eLQWtzv+VAIEFJ1/xxXlQ= +github.com/aws/smithy-go v1.28.1/go.mod h1:YE2RhdIuDbA5E5bTdciG9KrW3+TiEONeUWCqxX9i1Fc= github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw= github.com/bytecodealliance/wasmtime-go/v43 v43.0.2 h1:EZJlEpDanv6j/Y5Mcl2ndMjgK5Tw2QVDTJCzmLWNozg= diff --git a/tests/mw-dev/README.md b/tests/mw-dev/README.md index 8fed41ed0..54e133be5 100644 --- a/tests/mw-dev/README.md +++ b/tests/mw-dev/README.md @@ -103,12 +103,20 @@ Optional perf-step settings are: - `trino_ca_cert_file` (default empty, using system roots) - `trino_startup_timeout` (default `2m`) - `trino_startup_poll_interval` (default `2s`) +- `athena_catalog` (default `AwsDataCatalog`) +- `athena_poll_interval` (default `500ms`) +- `athena_query_timeout` (default `30m`) + +An Athena target additionally requires explicit `athena_region`, +`athena_workgroup`, `athena_database`, and `athena_output_location` settings. +The output location must be the same `s3://` prefix enforced by the workgroup. +Athena uses unique query execution IDs for result object names. The startup window contains an authenticated `SELECT 1` retry and completes before warmup or measured statements run. For the isolated mw-dev cell, use `trino_ca_cert_file: /trino-ca/ca.crt`. -`posthog_frozen_perf` enables Trino and selects the isolated Trino suite. Its +`posthog_frozen_perf` enables Trino and Athena and selects the isolated Trino suite. Its scenario Job mounts the per-run CA from `duckgres-trino-tls` and passes that path through `DUCKGRES_SCENARIO_TRINO_CA_CERT`; the perf adapter verifies the coordinator certificate and retries its first authenticated query for the @@ -116,8 +124,12 @@ bounded Secret-projection window. It also passes the deployed `DUCKGRES_K8S_WORKER_CPU_REQUEST` and `DUCKGRES_K8S_WORKER_MEMORY_REQUEST` into the perf step, which requests that exact shape through PGWire startup options and therefore bypasses the exploratory worker tier. The paired catalog remains the single SQL -source: direct-Parquet `raw_view` members run only through PGWire, while each -production-shaped `ducklake_table` member runs through both PGWire and Trino. +source: direct-Parquet `raw_view` members run only through PGWire, +production-shaped `ducklake_table` members run through PGWire and Trino, and +`athena_external` members run through Athena against Glue tables over the same +immutable Parquet objects. Athena is on-demand, result reuse is disabled, and +the harness records service-side timing and scanned bytes in +`query_service_metrics.csv`. To reproduce the scheduled run, deploy and test with `E2E_SUITE=trino` and the same `TRINO_POD_IDENTITY_ROLE` required by the isolated Trino lane. Teardown and the scheduled cleanup sweep remove both namespace-local workloads and their @@ -504,6 +516,10 @@ finalizers are still running. | var | `TS_WIF_AUDIENCE_MW_DEV` | Tailscale WIF audience | | secret | `MW_DEV_ACCOUNT_ID` | mw-dev AWS account id (kept out of committed code; ARNs are built from it) | | secret | `MW_DEV_TRINO_POD_IDENTITY_ROLE` | full ARN of the dedicated mw-dev Trino Pod Identity role (consumed only by the Trino lane) | +| secret | `MW_DEV_ATHENA_POD_IDENTITY_ROLE` | full ARN of the least-privilege Athena scenario Pod Identity role | +| var | `MW_DEV_ATHENA_WORKGROUP` | on-demand Athena engine-v3 workgroup | +| var | `MW_DEV_ATHENA_DATABASE` | Glue database containing the frozen external tables | +| var | `MW_DEV_ATHENA_RESULTS_S3_URI` | encrypted result prefix ending in `/`; lifecycle expiry is required | | secret | `AWS_ECR_PUBLISH_IAM_ROLE` | ECR push (already exists; used by CD) | | (role) | `github-duckgres-e2e` | dedicated stripped role in the mw-dev account (posthog-cloud-infra) — `eks:DescribeCluster` + Pod Identity association calls + `iam:PassRole`/`iam:GetRole` on the CP and dedicated Trino roles + an EKS access entry for kubectl. The workflow assumes `arn:aws:iam:::role/github-duckgres-e2e`. | | repo setting | "Require approval for all outside collaborators" | the access gate (see below) | diff --git a/tests/mw-dev/manifests.tmpl.yaml b/tests/mw-dev/manifests.tmpl.yaml index 8305ca433..bef9ac382 100644 --- a/tests/mw-dev/manifests.tmpl.yaml +++ b/tests/mw-dev/manifests.tmpl.yaml @@ -128,6 +128,12 @@ metadata: namespace: ${NAMESPACE} automountServiceAccountToken: false --- +apiVersion: v1 +kind: ServiceAccount +metadata: + name: duckgres-scenario + namespace: ${NAMESPACE} +--- # In-namespace CP permissions: spawn/patch worker pods, manage per-worker # secrets, read the worker configmap, hold the janitor lease. Mirrors the # chart's duckgres-control-plane Role. diff --git a/tests/mw-dev/run.sh b/tests/mw-dev/run.sh index 02902cd7f..bfee09d10 100755 --- a/tests/mw-dev/run.sh +++ b/tests/mw-dev/run.sh @@ -20,6 +20,7 @@ KUBECTL=(kubectl --context "$CTX") EKS_CLUSTER_NAME="${EKS_CLUSTER_NAME:-posthog-mw-dev}" AWS_REGION="${AWS_REGION:-us-east-1}" SA_NAME="duckgres" +SCENARIO_SA_NAME="duckgres-scenario" FROZEN_S3_URI="${DUCKGRES_SCENARIO_FROZEN_S3_URI:-s3://posthog-duckgres-scenario-frozen-data-mw-dev/frozen_v1/}" SCENARIO_JOB_WATCH_TIMEOUT_SECONDS="${SCENARIO_JOB_WATCH_TIMEOUT_SECONDS:-16200}" SCENARIO_JOB_CLEANUP_TIMEOUT_SECONDS="${SCENARIO_JOB_CLEANUP_TIMEOUT_SECONDS:-180}" @@ -164,6 +165,14 @@ ensure_trino_pod_identity() { create_pod_identity trino "$TRINO_POD_IDENTITY_ROLE" } +ensure_scenario_pod_identity() { + : "${SCENARIO_POD_IDENTITY_ROLE:?SCENARIO_POD_IDENTITY_ROLE is required for the Athena perf scenario}" + create_pod_identity "$SCENARIO_SA_NAME" "$SCENARIO_POD_IDENTITY_ROLE" + # Pod Identity is injected only at pod admission. Let the association reach + # the node agent before test-scenario creates the runner Job. + sleep 15 +} + delete_pod_identity() { local ids id ids="$(aws eks list-pod-identity-associations --region "$AWS_REGION" \ @@ -328,6 +337,10 @@ cmd_deploy() { ensure_pod_identity restart_cp_with_identity + if [ "$SCENARIO_NAME" = "posthog_frozen_perf" ]; then + ensure_scenario_pod_identity + fi + if [ "$E2E_SUITE" = "trino" ]; then # Associate before admitting Trino pods: the Pod Identity agent injects # credentials only at admission and never retrofits an existing pod. @@ -520,6 +533,7 @@ spec: # holds the only artifact copy. Voluntary disruption would lose both. karpenter.sh/do-not-disrupt: "true" spec: + serviceAccountName: $SCENARIO_SA_NAME restartPolicy: Never nodeSelector: kubernetes.io/arch: arm64 @@ -535,6 +549,10 @@ spec: - { name: DUCKGRES_SCENARIO_SNI_SUFFIX, value: "$suffix" } - { name: DUCKGRES_SCENARIO_FROZEN_S3_URI, value: "$FROZEN_S3_URI" } - { name: DUCKGRES_SCENARIO_TRINO_CA_CERT, value: "/trino-ca/ca.crt" } + - { name: DUCKGRES_SCENARIO_ATHENA_REGION, value: "$AWS_REGION" } + - { name: DUCKGRES_SCENARIO_ATHENA_WORKGROUP, value: "${DUCKGRES_SCENARIO_ATHENA_WORKGROUP:-}" } + - { name: DUCKGRES_SCENARIO_ATHENA_DATABASE, value: "${DUCKGRES_SCENARIO_ATHENA_DATABASE:-}" } + - { name: DUCKGRES_SCENARIO_ATHENA_RESULTS_S3_URI, value: "${DUCKGRES_SCENARIO_ATHENA_RESULTS_S3_URI:-}" } - { name: DUCKGRES_SCENARIO_DBT_BIN, value: "dbt" } - { name: DUCKGRES_K8S_WORKER_CPU_REQUEST, value: "$DUCKGRES_K8S_WORKER_CPU_REQUEST" } - { name: DUCKGRES_K8S_WORKER_MEMORY_REQUEST, value: "$DUCKGRES_K8S_WORKER_MEMORY_REQUEST" } diff --git a/tests/mw-dev/run_sh_test.go b/tests/mw-dev/run_sh_test.go index b683ca0ec..a245eda78 100644 --- a/tests/mw-dev/run_sh_test.go +++ b/tests/mw-dev/run_sh_test.go @@ -135,6 +135,24 @@ func TestTrinoDeployStartsWorkloadsWithoutScaleSubresource(t *testing.T) { } } +func TestDeployCreatesDedicatedScenarioPodIdentityForAthenaPerf(t *testing.T) { + fakes := newRunSHFakes(t) + cmd := runSHCommand(t, fakes.binDir, "deploy", + "SCENARIO_DEV_ALLOW_DUCKLING_DELETE=1", + "SCENARIO_NAME=posthog_frozen_perf", + "SCENARIO_POD_IDENTITY_ROLE=arn:aws:iam::123456789012:role/athena-perf", + ) + out, err := cmd.CombinedOutput() + if err != nil { + t.Fatalf("Athena perf deploy failed: %v\n%s", err, out) + } + calls := fakes.calls(t) + want := "aws eks create-pod-identity-association --region us-east-1 --cluster-name test-cluster --namespace duckgres-ci-pr-123 --service-account duckgres-scenario --role-arn arn:aws:iam::123456789012:role/athena-perf" + if !strings.Contains(calls, want) { + t.Fatalf("deploy did not create dedicated scenario Pod Identity; calls:\n%s", calls) + } +} + func TestTrinoWorkersMatchDuckgresAggregateCompute(t *testing.T) { raw, err := os.ReadFile("manifests.trino.tmpl.yaml") if err != nil { @@ -405,6 +423,9 @@ func TestScenarioRunsSelectedScenarioAgainstIsolatedStack(t *testing.T) { cmd := runSHCommand(t, fakes.binDir, "test-scenario", "SCENARIO_RUNNER_IMAGE=example.invalid/duckgres:scenario", "SCENARIO_NAME=fast-suite", + "DUCKGRES_SCENARIO_ATHENA_WORKGROUP=benchmark", + "DUCKGRES_SCENARIO_ATHENA_DATABASE=benchmark_frozen", + "DUCKGRES_SCENARIO_ATHENA_RESULTS_S3_URI=s3://benchmark-results/root/", ) out, err := cmd.CombinedOutput() if err != nil { @@ -416,9 +437,14 @@ func TestScenarioRunsSelectedScenarioAgainstIsolatedStack(t *testing.T) { "kubectl --context test-context -n duckgres-ci-pr-123 get svc duckgres-control-plane -o jsonpath={.spec.clusterIP}", "kubectl --context test-context -n duckgres-ci-pr-123 apply -f -", "name: artifact-keeper", + "serviceAccountName: duckgres-scenario", "value: \"isolated-test-secret\"", "name: DUCKGRES_SCENARIO_ORG_ID, value: \"ci-pr-123-cnpg\"", "name: DUCKGRES_SCENARIO_TRINO_CA_CERT, value: \"/trino-ca/ca.crt\"", + "name: DUCKGRES_SCENARIO_ATHENA_REGION, value: \"us-east-1\"", + "name: DUCKGRES_SCENARIO_ATHENA_WORKGROUP, value: \"benchmark\"", + "name: DUCKGRES_SCENARIO_ATHENA_DATABASE, value: \"benchmark_frozen\"", + "name: DUCKGRES_SCENARIO_ATHENA_RESULTS_S3_URI, value: \"s3://benchmark-results/root/\"", "name: trino-ca, mountPath: /trino-ca, readOnly: true", "name: trino-ca, secret: { secretName: duckgres-trino-tls, optional: true", "kubectl --context test-context -n duckgres-ci-pr-123 logs -f pod/duckgres-scenario-pod", diff --git a/tests/mw-dev/scenario/perf/adapter_test.go b/tests/mw-dev/scenario/perf/adapter_test.go index 65e6acb82..5079a0cd2 100644 --- a/tests/mw-dev/scenario/perf/adapter_test.go +++ b/tests/mw-dev/scenario/perf/adapter_test.go @@ -13,6 +13,7 @@ import ( "github.com/posthog/duckgres/tests/mw-dev/scenario/provision" scenariosql "github.com/posthog/duckgres/tests/mw-dev/scenario/sql" perfcore "github.com/posthog/duckgres/tests/perf/core" + athenadriver "github.com/posthog/duckgres/tests/perf/drivers/athena" trinodriver "github.com/posthog/duckgres/tests/perf/drivers/trino" ) @@ -80,7 +81,7 @@ func TestExecutorRunsPerfStepAndWritesArtifacts(t *testing.T) { } perfDir := filepath.Join(executor.OutputDir(), "perf") - for _, name := range []string{"summary.json", "query_results.csv", "server_metrics.prom"} { + for _, name := range []string{"summary.json", "query_results.csv", "query_service_metrics.csv", "server_metrics.prom"} { if _, err := os.Stat(filepath.Join(perfDir, name)); err != nil { t.Fatalf("expected perf artifact %s: %v", name, err) } @@ -221,6 +222,55 @@ func TestExecutorBuildsTrinoDriverFromReadinessState(t *testing.T) { } } +func TestExecutorBuildsAthenaDriverFromExplicitOnDemandConfig(t *testing.T) { + catalogPath := writePerfCatalog(t, []perfcore.Protocol{perfcore.ProtocolAthena}) + provisionState := provision.NewState() + provisionState.StoreProvisionResponse("scenario-org", provision.ProvisionResponse{Password: "root-password"}) + factory := &fakeDriverFactory{} + executor := NewExecutor(ExecutorConfig{ + ProvisionState: provisionState, + OutputDir: t.TempDir(), + DriverFactory: factory, + }) + + err := executor.ExecuteStep(context.Background(), core.Step{ + ID: "perf_queries", + Type: StepTypePerfQueries, + With: map[string]any{ + "org_id": "scenario-org", + "catalog_file": catalogPath, + "run_id": "scenario-run-1", + "targets": []any{"athena"}, + "athena_region": "us-east-1", + "athena_workgroup": "benchmark", + "athena_catalog": "AwsDataCatalog", + "athena_database": "benchmark_frozen", + "athena_output_location": "s3://benchmark-results/run/", + "athena_poll_interval": "750ms", + "athena_query_timeout": "20m", + }, + }) + if err != nil { + t.Fatalf("ExecuteStep returned error: %v", err) + } + + want := athenadriver.ConnectionConfig{ + Region: "us-east-1", + WorkGroup: "benchmark", + Catalog: "AwsDataCatalog", + Database: "benchmark_frozen", + OutputLocation: "s3://benchmark-results/run/", + PollInterval: 750 * time.Millisecond, + QueryTimeout: 20 * time.Minute, + } + if factory.athenaConnection != want { + t.Fatalf("Athena connection = %+v, want %+v", factory.athenaConnection, want) + } + if factory.athenaContext == nil { + t.Fatal("Athena factory did not receive scenario context") + } +} + func TestExecutorRejectsTrinoWithoutReadinessState(t *testing.T) { catalogPath := writePerfCatalog(t, []perfcore.Protocol{perfcore.ProtocolTrino}) provisionState := provision.NewState() @@ -443,6 +493,9 @@ type fakeDriverFactory struct { 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) { @@ -458,6 +511,13 @@ func (f *fakeDriverFactory) NewTrino(ctx context.Context, connection trinodriver return f.trinoDriver, nil } +func (f *fakeDriverFactory) NewAthena(ctx context.Context, connection athenadriver.ConnectionConfig) (perfcore.ProtocolDriver, error) { + f.athenaContext = ctx + f.athenaConnection = connection + f.athenaDriver = &fakeProtocolDriver{protocol: perfcore.ProtocolAthena} + return f.athenaDriver, nil +} + type fakeProtocolDriver struct { protocol perfcore.Protocol err error diff --git a/tests/mw-dev/scenario/perf/steps.go b/tests/mw-dev/scenario/perf/steps.go index 4bfe679b2..d262ef0f0 100644 --- a/tests/mw-dev/scenario/perf/steps.go +++ b/tests/mw-dev/scenario/perf/steps.go @@ -13,6 +13,7 @@ import ( "github.com/posthog/duckgres/tests/mw-dev/scenario/provision" scenariosql "github.com/posthog/duckgres/tests/mw-dev/scenario/sql" perfcore "github.com/posthog/duckgres/tests/perf/core" + athenadriver "github.com/posthog/duckgres/tests/perf/drivers/athena" pgdriver "github.com/posthog/duckgres/tests/perf/drivers/pgwire" trinodriver "github.com/posthog/duckgres/tests/perf/drivers/trino" ) @@ -22,6 +23,7 @@ const StepTypePerfQueries = "perf_queries" type DriverFactory interface { NewPGWire(connection scenariosql.PGWireConnection) (perfcore.ProtocolDriver, error) NewTrino(ctx context.Context, connection trinodriver.ConnectionConfig) (perfcore.ProtocolDriver, error) + NewAthena(ctx context.Context, connection athenadriver.ConnectionConfig) (perfcore.ProtocolDriver, error) } type ExecutorConfig struct { @@ -54,22 +56,29 @@ type StepResult struct { } type stepSpec struct { - OrgID string - Username string - Password string - CatalogFile string - Targets []perfcore.Protocol - RunID string - DatasetVersion string - Database string - OutputSubdir string - ReadOnly bool - FailOnQueryErrors bool - WorkerCPU string - WorkerMemory string - TrinoSchema string - TrinoCACertFile string - TrinoStartup trinodriver.StartupOptions + OrgID string + Username string + Password string + CatalogFile string + Targets []perfcore.Protocol + RunID string + DatasetVersion string + Database string + OutputSubdir string + ReadOnly bool + FailOnQueryErrors bool + WorkerCPU string + WorkerMemory string + TrinoSchema string + TrinoCACertFile string + TrinoStartup trinodriver.StartupOptions + AthenaRegion string + AthenaWorkGroup string + AthenaCatalog string + AthenaDatabase string + AthenaOutputLocation string + AthenaPollInterval time.Duration + AthenaQueryTimeout time.Duration } type defaultDriverFactory struct{} @@ -229,6 +238,14 @@ func (e *Executor) parseStep(step core.Step) (stepSpec, error) { if err != nil { return stepSpec{}, err } + athenaPollInterval, err := durationFromWith(step, "athena_poll_interval") + if err != nil { + return stepSpec{}, err + } + athenaQueryTimeout, err := durationFromWith(step, "athena_query_timeout") + if err != nil { + return stepSpec{}, err + } username := stringFromWith(step, "username", "root") password := stringFromWith(step, "password", "") @@ -266,6 +283,13 @@ func (e *Executor) parseStep(step core.Step) (stepSpec, error) { Timeout: trinoStartupTimeout, PollInterval: trinoStartupPollInterval, }, + AthenaRegion: stringFromWith(step, "athena_region", ""), + AthenaWorkGroup: stringFromWith(step, "athena_workgroup", ""), + AthenaCatalog: stringFromWith(step, "athena_catalog", "AwsDataCatalog"), + AthenaDatabase: stringFromWith(step, "athena_database", ""), + AthenaOutputLocation: stringFromWith(step, "athena_output_location", ""), + AthenaPollInterval: athenaPollInterval, + AthenaQueryTimeout: athenaQueryTimeout, }, nil } @@ -288,7 +312,7 @@ func targetsFromWith(step core.Step) ([]perfcore.Protocol, error) { } target := perfcore.Protocol(value) switch target { - case perfcore.ProtocolPGWire, perfcore.ProtocolTrino: + case perfcore.ProtocolPGWire, 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)) } @@ -352,6 +376,16 @@ func (e *Executor) driversForCatalog(ctx context.Context, catalog perfcore.Catal return nil, classified(ErrorClassConfig, fmt.Errorf("create Trino perf driver: %w", err)) } drivers[target] = driver + case perfcore.ProtocolAthena: + connection, err := e.athenaConnection(spec) + if err != nil { + return nil, err + } + driver, err := e.driverFactory.NewAthena(ctx, connection) + if err != nil { + return nil, classified(ErrorClassConfig, fmt.Errorf("create Athena perf driver: %w", err)) + } + drivers[target] = driver default: return nil, classified(ErrorClassConfig, fmt.Errorf("unsupported perf target protocol %q", target)) } @@ -360,6 +394,32 @@ func (e *Executor) driversForCatalog(ctx context.Context, catalog perfcore.Catal return drivers, nil } +func (e *Executor) athenaConnection(spec stepSpec) (athenadriver.ConnectionConfig, error) { + required := []struct { + key string + value string + }{ + {key: "athena_region", value: spec.AthenaRegion}, + {key: "athena_workgroup", value: spec.AthenaWorkGroup}, + {key: "athena_database", value: spec.AthenaDatabase}, + {key: "athena_output_location", value: spec.AthenaOutputLocation}, + } + for _, setting := range required { + if strings.TrimSpace(setting.value) == "" { + return athenadriver.ConnectionConfig{}, classified(ErrorClassConfig, fmt.Errorf("athena perf target requires with.%s", setting.key)) + } + } + return athenadriver.ConnectionConfig{ + Region: spec.AthenaRegion, + WorkGroup: spec.AthenaWorkGroup, + Catalog: spec.AthenaCatalog, + Database: spec.AthenaDatabase, + OutputLocation: spec.AthenaOutputLocation, + PollInterval: spec.AthenaPollInterval, + QueryTimeout: spec.AthenaQueryTimeout, + }, nil +} + func (e *Executor) trinoConnection(spec stepSpec) (trinodriver.ConnectionConfig, error) { if e.provisionState == nil { return trinodriver.ConnectionConfig{}, classified(ErrorClassConfig, fmt.Errorf("provision state is required for Trino perf target")) @@ -435,6 +495,10 @@ func (defaultDriverFactory) NewTrino(ctx context.Context, connection trinodriver return trinodriver.New(ctx, connection) } +func (defaultDriverFactory) NewAthena(ctx context.Context, connection athenadriver.ConnectionConfig) (perfcore.ProtocolDriver, error) { + return athenadriver.New(ctx, connection) +} + func requiredString(step core.Step, key string) (string, error) { value, ok := step.With[key].(string) if !ok || value == "" { diff --git a/tests/mw-dev/scenario/runner_test.go b/tests/mw-dev/scenario/runner_test.go index b4e45ee2b..9e3f42a30 100644 --- a/tests/mw-dev/scenario/runner_test.go +++ b/tests/mw-dev/scenario/runner_test.go @@ -164,6 +164,7 @@ func TestProvisionSmokeScenarioUsesIsolatedStackWarehouseIdentityAndSupportedSte func TestFrozenSuccessScenariosUseIsolatedStackWarehouseIdentity(t *testing.T) { const scenarioOrgID = "ci-pr-123-cnpg" t.Setenv("DUCKGRES_SCENARIO_TRINO_CA_CERT", "/tmp/test-trino-ca.crt") + setAthenaPerfEnv(t) t.Setenv("DUCKGRES_K8S_WORKER_CPU_REQUEST", "3") t.Setenv("DUCKGRES_K8S_WORKER_MEMORY_REQUEST", "12Gi") @@ -452,6 +453,7 @@ func TestFrozenPerfScenarioUsesSupportedStepsAndRelativeCatalog(t *testing.T) { t.Setenv("DUCKGRES_SCENARIO_FROZEN_S3_URI", "s3://example-frozen/frozen_v1/") t.Setenv("DUCKGRES_SCENARIO_ORG_ID", "ci-pr-123-cnpg") t.Setenv("DUCKGRES_SCENARIO_TRINO_CA_CERT", "/tmp/test-trino-ca.crt") + setAthenaPerfEnv(t) t.Setenv("DUCKGRES_K8S_WORKER_CPU_REQUEST", "3") t.Setenv("DUCKGRES_K8S_WORKER_MEMORY_REQUEST", "12Gi") @@ -490,7 +492,7 @@ func TestFrozenPerfScenarioUsesSupportedStepsAndRelativeCatalog(t *testing.T) { t.Fatal("frozen perf scenario should not configure the deprecated Flight endpoint") } assertPerfQueryErrorsFailStep(t, step) - assertPerfTargetsPGWireAndTrino(t, step) + assertPerfTargetsPGWireTrinoAndAthena(t, step) if got, _ := step.With["trino_ca_cert_file"].(string); got != "/tmp/test-trino-ca.crt" { t.Fatalf("perf Trino CA file = %q, want resolved environment path", got) } @@ -500,6 +502,9 @@ func TestFrozenPerfScenarioUsesSupportedStepsAndRelativeCatalog(t *testing.T) { if got, _ := step.With["worker_memory"].(string); got != "12Gi" { t.Fatalf("perf worker_memory = %q, want standard worker memory", got) } + if got, _ := step.With["athena_output_location"].(string); got != "s3://benchmark-results/root/" { + t.Fatalf("Athena output location = %q, want workgroup-enforced results prefix", got) + } } if !foundPerf { t.Fatal("expected frozen perf scenario to include a perf_queries step") @@ -510,6 +515,7 @@ func TestFrozenPerfScenarioBuildsAndValidatesPostHogTablesBeforePerf(t *testing. t.Setenv("DUCKGRES_SCENARIO_FROZEN_S3_URI", "s3://example-frozen/frozen_v1/") t.Setenv("DUCKGRES_SCENARIO_ORG_ID", "ci-pr-123-cnpg") t.Setenv("DUCKGRES_SCENARIO_TRINO_CA_CERT", "/tmp/test-trino-ca.crt") + setAthenaPerfEnv(t) t.Setenv("DUCKGRES_K8S_WORKER_CPU_REQUEST", "3") t.Setenv("DUCKGRES_K8S_WORKER_MEMORY_REQUEST", "12Gi") @@ -736,14 +742,22 @@ func assertPerfTargetsOnlyPGWire(t *testing.T, step core.Step) { } } -func assertPerfTargetsPGWireAndTrino(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) != 2 || targets[0] != "pgwire" || targets[1] != "trino" { - t.Fatalf("perf step %s targets = %#v, want [pgwire trino]", step.ID, step.With["targets"]) + 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"]) } } +func setAthenaPerfEnv(t *testing.T) { + t.Helper() + t.Setenv("DUCKGRES_SCENARIO_ATHENA_REGION", "us-east-1") + t.Setenv("DUCKGRES_SCENARIO_ATHENA_WORKGROUP", "benchmark") + t.Setenv("DUCKGRES_SCENARIO_ATHENA_DATABASE", "benchmark_frozen") + t.Setenv("DUCKGRES_SCENARIO_ATHENA_RESULTS_S3_URI", "s3://benchmark-results/root/") +} + func TestFrozenDBTProjectModelsRealisticProductAnalyticsWorkload(t *testing.T) { projectDir := filepath.Join("dbt", "posthog_frozen_project") requiredModels := map[string][]string{ diff --git a/tests/mw-dev/scenario/scenarios/posthog_frozen_perf.yaml b/tests/mw-dev/scenario/scenarios/posthog_frozen_perf.yaml index 7f7add961..9d7961854 100644 --- a/tests/mw-dev/scenario/scenarios/posthog_frozen_perf.yaml +++ b/tests/mw-dev/scenario/scenarios/posthog_frozen_perf.yaml @@ -4,6 +4,10 @@ required_env: - DUCKGRES_SCENARIO_ORG_ID - DUCKGRES_SCENARIO_FROZEN_S3_URI - DUCKGRES_SCENARIO_TRINO_CA_CERT + - DUCKGRES_SCENARIO_ATHENA_REGION + - DUCKGRES_SCENARIO_ATHENA_WORKGROUP + - DUCKGRES_SCENARIO_ATHENA_DATABASE + - DUCKGRES_SCENARIO_ATHENA_RESULTS_S3_URI - DUCKGRES_K8S_WORKER_CPU_REQUEST - DUCKGRES_K8S_WORKER_MEMORY_REQUEST steps: @@ -76,10 +80,17 @@ steps: org_id: ${env:DUCKGRES_SCENARIO_ORG_ID} catalog: ducklake catalog_file: ../../../perf/queries/ducklake_posthog_tables.yaml - targets: [pgwire, trino] + targets: [pgwire, trino, athena] trino_ca_cert_file: ${env:DUCKGRES_SCENARIO_TRINO_CA_CERT} trino_startup_timeout: 2m trino_startup_poll_interval: 2s + athena_region: ${env:DUCKGRES_SCENARIO_ATHENA_REGION} + athena_workgroup: ${env:DUCKGRES_SCENARIO_ATHENA_WORKGROUP} + athena_catalog: AwsDataCatalog + athena_database: ${env:DUCKGRES_SCENARIO_ATHENA_DATABASE} + athena_output_location: ${env:DUCKGRES_SCENARIO_ATHENA_RESULTS_S3_URI} + athena_poll_interval: 500ms + athena_query_timeout: 30m run_id: ${run_id} dataset_version: posthog-file-views-v1 fail_on_query_errors: true diff --git a/tests/mw-dev/scenario/script_test.go b/tests/mw-dev/scenario/script_test.go index 72cfde35c..422cf0efc 100644 --- a/tests/mw-dev/scenario/script_test.go +++ b/tests/mw-dev/scenario/script_test.go @@ -51,6 +51,10 @@ func TestScenarioRunScriptCheckEnvIncludesScenarioRequiredEnv(t *testing.T) { "DUCKGRES_SCENARIO_ORG_ID", "DUCKGRES_SCENARIO_FROZEN_S3_URI", "DUCKGRES_SCENARIO_TRINO_CA_CERT", + "DUCKGRES_SCENARIO_ATHENA_REGION", + "DUCKGRES_SCENARIO_ATHENA_WORKGROUP", + "DUCKGRES_SCENARIO_ATHENA_DATABASE", + "DUCKGRES_SCENARIO_ATHENA_RESULTS_S3_URI", "DUCKGRES_K8S_WORKER_CPU_REQUEST", "DUCKGRES_K8S_WORKER_MEMORY_REQUEST", } { @@ -90,6 +94,10 @@ func TestDevScenarioWorkflowUsesUnifiedMwDevHarness(t *testing.T) { "EKS_CLUSTER_NAME: posthog-mw-dev", "CP_POD_IDENTITY_ROLE: arn:aws:iam::${{ secrets.MW_DEV_ACCOUNT_ID }}:role/duckgres-control-plane-dev", "TRINO_POD_IDENTITY_ROLE: ${{ secrets.MW_DEV_TRINO_POD_IDENTITY_ROLE }}", + "SCENARIO_POD_IDENTITY_ROLE: ${{ secrets.MW_DEV_ATHENA_POD_IDENTITY_ROLE }}", + "DUCKGRES_SCENARIO_ATHENA_WORKGROUP: ${{ vars.MW_DEV_ATHENA_WORKGROUP }}", + "DUCKGRES_SCENARIO_ATHENA_DATABASE: ${{ vars.MW_DEV_ATHENA_DATABASE }}", + "DUCKGRES_SCENARIO_ATHENA_RESULTS_S3_URI: ${{ vars.MW_DEV_ATHENA_RESULTS_S3_URI }}", "TRINO_IMAGE: ghcr.io/posthog/trino:", "E2E_SUITE: ${{ (github.event_name == 'schedule' || inputs.scenario == 'posthog_frozen_perf') && 'trino' || 'neutral' }}", "DUCKGRES_K8S_WORKER_CPU_REQUEST: \"3\"", diff --git a/tests/perf/README.md b/tests/perf/README.md index 96dc32440..fad9a249d 100644 --- a/tests/perf/README.md +++ b/tests/perf/README.md @@ -4,11 +4,19 @@ This package contains the golden-query performance harness. ## Protocol Drivers -Catalogs may target `pgwire`, `trino`, or both. Both drivers execute the same +Catalogs may target `pgwire`, `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 -SQL within the intersection supported by DuckDB and Trino. +SQL within the intersection supported by DuckDB, Trino, and Athena. + +The Athena driver uses on-demand capacity in an explicitly configured +workgroup, database, and S3 result prefix. It disables result reuse for every +execution and rejects a response which reports reuse. Timing is end to end: +it includes queueing, engine execution, and paginated result retrieval. The +driver defaults to catalog `AwsDataCatalog`, a 500ms status poll interval, and +a 30-minute query timeout. Workgroup, database, result prefix, and AWS region +are explicit scenario settings. Cancellation stops any unfinished query. The Trino driver requires an HTTPS coordinator and always verifies its TLS certificate. It uses system roots by default, or the explicitly configured CA @@ -54,10 +62,10 @@ paired_queries: ORDER BY 1 ``` -Paired catalogs must declare exactly the `raw_view` and `ducklake_table` -variants. A template expands in declaration order, with `raw_view` before -`ducklake_table`, into `q_events_daily__raw_view` and -`q_events_daily__ducklake_table`. Generated queries retain the same +Paired catalogs without Athena declare exactly the `raw_view` and +`ducklake_table` variants. Athena catalogs add `athena_external`, whose generic +table names are resolved in the configured Glue database. A template expands +in stable order: `raw_view`, `ducklake_table`, then `athena_external`. Generated queries retain the same `intent_id`, tags, parameters, and semantic template; only declared relation placeholders differ. They carry in-memory storage-target metadata, so later code does not need to infer the target from the generated ID. Legacy queries @@ -89,6 +97,8 @@ stored in the PGWire SQL field. This abstraction preserves the artifact contract while allowing downstream dashboards to compare paired targets by their generated query-ID suffixes. +PGWire executes `raw_view` and `ducklake_table`; Trino executes only +`ducklake_table`; Athena executes only `athena_external`. ## Local Smoke Run @@ -126,6 +136,7 @@ Artifacts are written to `artifacts/perf/`: - `summary.json` - `query_results.csv` +- `query_service_metrics.csv` - `server_metrics.prom` - `runner.log` - `dataset_manifest.json` (only when `DUCKGRES_PERF_DATASET_VERSION` is set) @@ -149,6 +160,12 @@ Artifacts are written to `artifacts/perf/`: `duration_ms` is emitted as milliseconds with fixed precision, and `started_at` is UTC RFC3339Nano. No CSV schema mutation is expected in this phase. +`query_service_metrics.csv` is an additive sidecar. Provider-backed rows record +queue, planning, engine, and service time; bytes scanned; DPU count when the +service returns it; result reuse; and engine version. `query_results.csv` +remains the canonical latency/status artifact and keeps its v1 header +unchanged. + ## Nightly Run ```bash diff --git a/tests/perf/core/catalog.go b/tests/perf/core/catalog.go index 7ad22c9d1..78a83acd6 100644 --- a/tests/perf/core/catalog.go +++ b/tests/perf/core/catalog.go @@ -70,7 +70,7 @@ func ParseCatalog(raw []byte) (Catalog, error) { return Catalog{}, err } if len(entries) > 0 { - if err := validateRelationVariants(file.RelationVariants, entries); err != nil { + if err := validateRelationVariants(file.Targets, file.RelationVariants, entries); err != nil { return Catalog{}, err } } @@ -129,7 +129,7 @@ func catalogEntries(raw []byte) ([]catalogEntry, error) { return entries, nil } -func validateRelationVariants(variants map[StorageTarget]map[string]string, entries []catalogEntry) error { +func validateRelationVariants(targets []Protocol, variants map[StorageTarget]map[string]string, entries []catalogEntry) error { hasPairedQueries := false for _, entry := range entries { if entry.paired != nil { @@ -140,17 +140,32 @@ func validateRelationVariants(variants map[StorageTarget]map[string]string, entr if !hasPairedQueries { return nil } - if len(variants) != 2 { - return fmt.Errorf("paired catalogs must declare exactly the raw_view and ducklake_table storage variants") + requiredTargets := []StorageTarget{StorageTargetRawView, StorageTargetDuckLakeTable} + for _, target := range targets { + if target == ProtocolAthena { + requiredTargets = append(requiredTargets, StorageTargetAthenaExternal) + break + } } - for _, target := range []StorageTarget{StorageTargetRawView, StorageTargetDuckLakeTable} { + if len(variants) != len(requiredTargets) { + return fmt.Errorf("paired catalogs must declare exactly the %s storage variants", storageTargetList(requiredTargets)) + } + for _, target := range requiredTargets { if _, ok := variants[target]; !ok { - return fmt.Errorf("paired catalogs must declare exactly the raw_view and ducklake_table storage variants") + return fmt.Errorf("paired catalogs must declare exactly the %s storage variants", storageTargetList(requiredTargets)) } } return nil } +func storageTargetList(targets []StorageTarget) string { + names := make([]string, 0, len(targets)) + for _, target := range targets { + names = append(names, string(target)) + } + return strings.Join(names, ", ") +} + func expandPairedQuery(def pairedQueryDefinition, variants map[StorageTarget]map[string]string) ([]Query, error) { if def.QueryIDBase == "" { return nil, fmt.Errorf("paired query missing query_id_base") @@ -166,8 +181,12 @@ func expandPairedQuery(def pairedQueryDefinition, variants map[StorageTarget]map return nil, fmt.Errorf("paired query %s must contain at least one relation placeholder", def.QueryIDBase) } - queries := make([]Query, 0, 2) - for _, target := range []StorageTarget{StorageTargetRawView, StorageTargetDuckLakeTable} { + targets := []StorageTarget{StorageTargetRawView, StorageTargetDuckLakeTable} + if _, ok := variants[StorageTargetAthenaExternal]; ok { + targets = append(targets, StorageTargetAthenaExternal) + } + queries := make([]Query, 0, len(targets)) + for _, target := range targets { rendered, err := renderRelationTemplate(def.QueryIDBase, def.SQLTemplate, placeholders, variants[target], target) if err != nil { return nil, err @@ -185,8 +204,12 @@ func expandPairedQuery(def pairedQueryDefinition, variants map[StorageTarget]map StorageTarget: target, }) } - if queries[0].PGWireSQL == queries[1].PGWireSQL { - return nil, fmt.Errorf("paired query %s relation bindings must differ between storage targets", def.QueryIDBase) + renderedSQL := make(map[string]struct{}, len(queries)) + for _, query := range queries { + if _, exists := renderedSQL[query.PGWireSQL]; exists { + return nil, fmt.Errorf("paired query %s relation bindings must differ between storage targets", def.QueryIDBase) + } + renderedSQL[query.PGWireSQL] = struct{}{} } return queries, nil } @@ -366,7 +389,7 @@ func validateCatalog(c Catalog) error { seenTargets := map[Protocol]struct{}{} for _, target := range c.Targets { switch target { - case ProtocolPGWire, ProtocolTrino: + case ProtocolPGWire, 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 4ba056064..c700297f8 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} + wantTargets = []Protocol{ProtocolPGWire, ProtocolTrino, ProtocolAthena} } if !reflect.DeepEqual(catalog.Targets, wantTargets) { t.Fatalf("catalog targets = %v, want %v", catalog.Targets, wantTargets) @@ -51,18 +51,25 @@ func TestCheckedInPostHogCatalogPublishesCompleteStablePairs(t *testing.T) { 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", } if got := queryIDs(catalog); !reflect.DeepEqual(got, want) { t.Fatalf("checked-in PostHog query IDs changed: got %v want %v", got, want) @@ -79,39 +86,76 @@ func TestCheckedInPostHogCatalogPublishesCompleteStablePairs(t *testing.T) { } } for index, query := range catalog.Queries { - wantTarget := StorageTargetRawView - if index%2 == 1 { - wantTarget = StorageTargetDuckLakeTable - } + wantTarget := []StorageTarget{StorageTargetRawView, StorageTargetDuckLakeTable, StorageTargetAthenaExternal}[index%3] if query.StorageTarget != wantTarget { t.Fatalf("query %s storage target = %q, want %q", query.QueryID, query.StorageTarget, wantTarget) } - if index%2 != 1 { + if index%3 != 2 { continue } - rawQuery := catalog.Queries[index-1] - if query.IntentID != rawQuery.IntentID { - t.Fatalf("query pair %s/%s has mismatched intents %q/%q", rawQuery.QueryID, query.QueryID, rawQuery.IntentID, query.IntentID) + rawQuery := catalog.Queries[index-2] + duckLakeQuery := catalog.Queries[index-1] + if query.IntentID != rawQuery.IntentID || duckLakeQuery.IntentID != rawQuery.IntentID { + t.Fatalf("query variants have mismatched intents: %q/%q/%q", rawQuery.IntentID, duckLakeQuery.IntentID, query.IntentID) } - if !reflect.DeepEqual(query.Tags, rawQuery.Tags) || !reflect.DeepEqual(query.Params, rawQuery.Params) { - t.Fatalf("query pair %s/%s must share tags and params", rawQuery.QueryID, query.QueryID) + if !reflect.DeepEqual(query.Tags, rawQuery.Tags) || !reflect.DeepEqual(query.Params, rawQuery.Params) || + !reflect.DeepEqual(duckLakeQuery.Tags, rawQuery.Tags) || !reflect.DeepEqual(duckLakeQuery.Params, rawQuery.Params) { + t.Fatalf("query variants for %s must share tags and params", rawQuery.IntentID) } rawRelation := `"frozen_v1"."events_file_view"` duckLakeRelation := `"posthog"."events"` + athenaRelation := `"events"` if strings.HasPrefix(query.IntentID, "intent_persons_") { rawRelation = `"frozen_v1"."persons_file_view"` duckLakeRelation = `"posthog"."persons"` + athenaRelation = `"persons"` } rawShape := strings.ReplaceAll(rawQuery.PGWireSQL, rawRelation, "") - duckLakeShape := strings.ReplaceAll(query.PGWireSQL, duckLakeRelation, "") - if rawShape != duckLakeShape { - t.Fatalf("query pair %s/%s differs beyond its relation:\nraw: %s\ntable: %s", rawQuery.QueryID, query.QueryID, rawQuery.PGWireSQL, query.PGWireSQL) + duckLakeShape := strings.ReplaceAll(duckLakeQuery.PGWireSQL, duckLakeRelation, "") + athenaShape := strings.ReplaceAll(query.PGWireSQL, athenaRelation, "") + if rawShape != duckLakeShape || rawShape != athenaShape { + t.Fatalf("query variants for %s differ beyond their relation", rawQuery.IntentID) } } } +func TestParseCatalogExpandsAthenaExternalVariantWhenAthenaIsTargeted(t *testing.T) { + catalog, err := ParseCatalog([]byte(athenaCatalogYAML(` +paired_queries: + - query_id_base: q_events + intent_id: ph.events.v1 + sql_template: SELECT COUNT(*) FROM {{ relation "events" }} +`))) + if err != nil { + t.Fatalf("ParseCatalog returned error: %v", err) + } + if got, want := queryIDs(catalog), []string{"q_events__raw_view", "q_events__ducklake_table", "q_events__athena_external"}; !reflect.DeepEqual(got, want) { + t.Fatalf("unexpected generated query order: got %v want %v", got, want) + } + athenaQuery := catalog.Queries[2] + if got, want := athenaQuery.StorageTarget, StorageTargetAthenaExternal; got != want { + t.Fatalf("Athena query target: got %q want %q", got, want) + } + if got, want := athenaQuery.PGWireSQL, `SELECT COUNT(*) FROM "events"`; got != want { + t.Fatalf("Athena query SQL: got %q want %q", got, want) + } +} + +func TestParseCatalogRejectsAthenaTargetWithoutExternalVariant(t *testing.T) { + raw := strings.Replace(pairedCatalogYAML(` +paired_queries: + - query_id_base: q_events + intent_id: ph.events.v1 + sql_template: SELECT COUNT(*) FROM {{ relation "events" }} +`), "targets: [pgwire]", "targets: [pgwire, athena]", 1) + _, err := ParseCatalog([]byte(raw)) + if err == nil || !strings.Contains(err.Error(), "athena_external") { + t.Fatalf("ParseCatalog error = %v, want missing athena_external variant", err) + } +} + func TestParseCatalogSuccess(t *testing.T) { raw := ` name: smoke @@ -425,6 +469,21 @@ relation_variants: `, "\t") + body) } +func athenaCatalogYAML(body string) string { + return strings.Replace(catalogYAML(strings.TrimSuffix(` +relation_variants: + raw_view: + events: frozen_v1.events_file_view + persons: frozen_v1.persons_file_view + ducklake_table: + events: posthog.events + persons: posthog.persons + athena_external: + events: events + persons: persons + `, "\t")+body), "targets: [pgwire]", "targets: [pgwire, trino, athena]", 1) +} + func catalogYAML(body string) string { return ` name: paired diff --git a/tests/perf/core/runner.go b/tests/perf/core/runner.go index cb57405ff..3a669e20b 100644 --- a/tests/perf/core/runner.go +++ b/tests/perf/core/runner.go @@ -132,6 +132,7 @@ func (r *QueryRunner) executeIteration(ctx context.Context, protocol Protocol, m } result.Duration = execResult.Duration result.Rows = execResult.Rows + result.ServiceMetrics = execResult.ServiceMetrics if err != nil { result.Status = "error" result.Error = err.Error() @@ -158,13 +159,23 @@ func (r *QueryRunner) executeIteration(ctx context.Context, protocol Protocol, m return nil } -// Raw-view members of a paired catalog are DuckDB views over read_parquet and -// therefore exist only behind PGWire. DuckLake-table members are backed by the -// shared catalog and intentionally run through every selected protocol. This -// keeps one canonical paired query definition while preventing a copied -// Trino-only catalog from drifting away from the PGWire workload. +// Each physical relation family is routed only to protocols which expose it. +// PGWire measures both the raw Parquet view and production-shaped DuckLake +// table, Trino measures the shared DuckLake table, and Athena measures its +// Glue external table over the same immutable Parquet files. func querySupportsProtocol(query Query, protocol Protocol) bool { - return query.StorageTarget != StorageTargetRawView || protocol == ProtocolPGWire + switch query.StorageTarget { + case "": + return true + case StorageTargetRawView: + return protocol == ProtocolPGWire + case StorageTargetDuckLakeTable: + return protocol == ProtocolPGWire || protocol == ProtocolTrino + case StorageTargetAthenaExternal: + return protocol == ProtocolAthena + default: + return false + } } // queriesForIteration alternates each generated raw-view/DuckLake-table pair diff --git a/tests/perf/core/runner_test.go b/tests/perf/core/runner_test.go index 9b8195b31..1f892eeca 100644 --- a/tests/perf/core/runner_test.go +++ b/tests/perf/core/runner_test.go @@ -262,6 +262,49 @@ func TestRunnerKeepsRawViewsOnPGWireAndRunsDuckLakeTablesOnEveryTarget(t *testin } } +func TestRunnerRoutesEachStorageVariantOnlyToItsComparableProtocol(t *testing.T) { + pg := &testDriver{protocol: ProtocolPGWire} + trinoDriver := &testDriver{protocol: ProtocolTrino} + athenaDriver := &testDriver{protocol: ProtocolAthena} + sink := &inMemorySink{} + runner := NewQueryRunner(RunnerConfig{ + Catalog: Catalog{ + Name: "three-engine-comparison", + MeasureIterations: 1, + Targets: []Protocol{ProtocolPGWire, ProtocolTrino, ProtocolAthena}, + 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{ + ProtocolPGWire: pg, + ProtocolTrino: trinoDriver, + ProtocolAthena: athenaDriver, + }, + Sink: sink, + Now: func() time.Time { return time.Unix(1700000000, 0) }, + }) + + summary, err := runner.Run(context.Background()) + if err != nil { + t.Fatalf("Run returned error: %v", err) + } + if got, want := pg.queryIDs, []string{"q__raw_view", "q__ducklake_table"}; !reflect.DeepEqual(got, want) { + t.Fatalf("PGWire query IDs: got %v want %v", got, want) + } + if got, want := trinoDriver.queryIDs, []string{"q__ducklake_table"}; !reflect.DeepEqual(got, want) { + t.Fatalf("Trino query IDs: got %v want %v", got, want) + } + if got, want := athenaDriver.queryIDs, []string{"q__athena_external"}; !reflect.DeepEqual(got, want) { + t.Fatalf("Athena query IDs: got %v want %v", got, want) + } + if summary.TotalQueries != 4 { + t.Fatalf("total measured queries = %d, want 4", summary.TotalQueries) + } +} + func (d *testDriver) Close() error { return nil } type inMemorySink struct { diff --git a/tests/perf/core/sink.go b/tests/perf/core/sink.go index 85d4e5b2a..3769b7661 100644 --- a/tests/perf/core/sink.go +++ b/tests/perf/core/sink.go @@ -11,10 +11,12 @@ import ( ) type ArtifactSink struct { - dir string - csvFile *os.File - csvWriter *csv.Writer - closed bool + dir string + csvFile *os.File + csvWriter *csv.Writer + serviceMetricsFile *os.File + serviceMetricsWriter *csv.Writer + closed bool } func NewArtifactSink(dir string) (*ArtifactSink, error) { @@ -48,10 +50,44 @@ func NewArtifactSink(dir string) (*ArtifactSink, error) { _ = f.Close() return nil, fmt.Errorf("flush csv header: %w", err) } + serviceMetricsPath := filepath.Join(dir, "query_service_metrics.csv") + serviceMetricsFile, err := os.Create(serviceMetricsPath) + if err != nil { + _ = f.Close() + return nil, fmt.Errorf("create service metrics artifact: %w", err) + } + serviceMetricsWriter := csv.NewWriter(serviceMetricsFile) + serviceMetricsHeader := []string{ + "query_id", + "intent_id", + "measure_iteration", + "protocol", + "queue_ms", + "planning_ms", + "engine_ms", + "service_ms", + "bytes_scanned", + "dpu_count", + "result_reused", + "engine_version", + } + if err := serviceMetricsWriter.Write(serviceMetricsHeader); err != nil { + _ = serviceMetricsFile.Close() + _ = f.Close() + return nil, fmt.Errorf("write service metrics header: %w", err) + } + serviceMetricsWriter.Flush() + if err := serviceMetricsWriter.Error(); err != nil { + _ = serviceMetricsFile.Close() + _ = f.Close() + return nil, fmt.Errorf("flush service metrics header: %w", err) + } return &ArtifactSink{ - dir: dir, - csvFile: f, - csvWriter: w, + dir: dir, + csvFile: f, + csvWriter: w, + serviceMetricsFile: serviceMetricsFile, + serviceMetricsWriter: serviceMetricsWriter, }, nil } @@ -78,6 +114,30 @@ func (s *ArtifactSink) Record(result QueryResult) error { if err := s.csvWriter.Error(); err != nil { return fmt.Errorf("flush csv row: %w", err) } + if result.ServiceMetrics != nil { + metrics := result.ServiceMetrics + serviceMetricsRow := []string{ + result.QueryID, + result.IntentID, + strconv.Itoa(result.MeasureIteration), + string(result.Protocol), + formatMilliseconds(metrics.QueueDuration), + formatMilliseconds(metrics.PlanningDuration), + formatMilliseconds(metrics.EngineDuration), + formatMilliseconds(metrics.ServiceDuration), + strconv.FormatInt(metrics.BytesScanned, 10), + strconv.FormatFloat(metrics.DPUCount, 'f', -1, 64), + strconv.FormatBool(metrics.ResultReused), + metrics.EngineVersion, + } + if err := s.serviceMetricsWriter.Write(serviceMetricsRow); err != nil { + return fmt.Errorf("write service metrics row: %w", err) + } + s.serviceMetricsWriter.Flush() + if err := s.serviceMetricsWriter.Error(); err != nil { + return fmt.Errorf("flush service metrics row: %w", err) + } + } return nil } @@ -97,6 +157,17 @@ func (s *ArtifactSink) Close(summary RunSummary, serverMetrics string) error { return fmt.Errorf("close csv file: %w", err) } } + if s.serviceMetricsWriter != nil { + s.serviceMetricsWriter.Flush() + if err := s.serviceMetricsWriter.Error(); err != nil { + return fmt.Errorf("flush service metrics close: %w", err) + } + } + if s.serviceMetricsFile != nil { + if err := s.serviceMetricsFile.Close(); err != nil { + return fmt.Errorf("close service metrics file: %w", err) + } + } summaryPath := filepath.Join(s.dir, "summary.json") summaryFile, err := os.Create(summaryPath) @@ -119,3 +190,7 @@ func (s *ArtifactSink) Close(summary RunSummary, serverMetrics string) error { } return nil } + +func formatMilliseconds(duration time.Duration) string { + return strconv.FormatFloat(float64(duration)/float64(time.Millisecond), 'f', 6, 64) +} diff --git a/tests/perf/core/sink_test.go b/tests/perf/core/sink_test.go index eceb9e8c5..62878bdfc 100644 --- a/tests/perf/core/sink_test.go +++ b/tests/perf/core/sink_test.go @@ -142,3 +142,67 @@ paired_queries: t.Fatalf("CSV query IDs: got %v want %v", got, want) } } + +func TestArtifactSinkWritesAthenaServiceMetricsWithoutChangingQueryResultsV1(t *testing.T) { + dir := t.TempDir() + sink, err := NewArtifactSink(dir) + if err != nil { + t.Fatalf("NewArtifactSink returned error: %v", err) + } + if err := sink.Record(QueryResult{ + QueryID: "q1__athena_external", + IntentID: "i1", + MeasureIteration: 1, + Protocol: ProtocolAthena, + Status: "ok", + Rows: 1, + Duration: 3 * time.Second, + StartedAt: time.Unix(1700000000, 0), + ServiceMetrics: &ServiceMetrics{ + QueueDuration: 100 * time.Millisecond, + PlanningDuration: 200 * time.Millisecond, + EngineDuration: 2 * time.Second, + ServiceDuration: 2500 * time.Millisecond, + BytesScanned: 4096, + DPUCount: 4, + ResultReused: false, + EngineVersion: "Athena engine version 3", + }, + }); err != nil { + t.Fatalf("Record returned error: %v", err) + } + if err := sink.Close(RunSummary{}, ""); err != nil { + t.Fatalf("Close returned error: %v", err) + } + + queryFile, err := os.Open(filepath.Join(dir, "query_results.csv")) + if err != nil { + t.Fatalf("open query_results.csv: %v", err) + } + queryRecords, err := csv.NewReader(queryFile).ReadAll() + _ = queryFile.Close() + if err != nil { + t.Fatalf("read query_results.csv: %v", err) + } + wantQueryHeader := []string{"query_id", "intent_id", "measure_iteration", "protocol", "status", "error", "error_class", "rows", "duration_ms", "started_at"} + if !reflect.DeepEqual(queryRecords[0], wantQueryHeader) { + t.Fatalf("query_results.csv v1 header changed: got %v want %v", queryRecords[0], wantQueryHeader) + } + + metricsFile, err := os.Open(filepath.Join(dir, "query_service_metrics.csv")) + if err != nil { + t.Fatalf("open query_service_metrics.csv: %v", err) + } + metricsRecords, err := csv.NewReader(metricsFile).ReadAll() + _ = metricsFile.Close() + if err != nil { + t.Fatalf("read query_service_metrics.csv: %v", err) + } + wantMetricsHeader := []string{"query_id", "intent_id", "measure_iteration", "protocol", "queue_ms", "planning_ms", "engine_ms", "service_ms", "bytes_scanned", "dpu_count", "result_reused", "engine_version"} + if !reflect.DeepEqual(metricsRecords[0], wantMetricsHeader) { + t.Fatalf("service metrics header: got %v want %v", metricsRecords[0], wantMetricsHeader) + } + if got, want := metricsRecords[1], []string{"q1__athena_external", "i1", "1", "athena", "100.000000", "200.000000", "2000.000000", "2500.000000", "4096", "4", "false", "Athena engine version 3"}; !reflect.DeepEqual(got, want) { + t.Fatalf("service metrics row: got %v want %v", got, want) + } +} diff --git a/tests/perf/core/types.go b/tests/perf/core/types.go index 4f2e7e854..3aed9a624 100644 --- a/tests/perf/core/types.go +++ b/tests/perf/core/types.go @@ -7,6 +7,7 @@ type Protocol string const ( ProtocolPGWire Protocol = "pgwire" ProtocolTrino Protocol = "trino" + ProtocolAthena Protocol = "athena" ) // StorageTarget identifies the physical relation family selected for a paired @@ -15,8 +16,9 @@ const ( type StorageTarget string const ( - StorageTargetRawView StorageTarget = "raw_view" - StorageTargetDuckLakeTable StorageTarget = "ducklake_table" + StorageTargetRawView StorageTarget = "raw_view" + StorageTargetDuckLakeTable StorageTarget = "ducklake_table" + StorageTargetAthenaExternal StorageTarget = "athena_external" ) type Catalog struct { @@ -48,21 +50,37 @@ func (q Query) CanonicalSQL() string { } type ExecutionResult struct { - Rows int64 - Duration time.Duration + Rows int64 + Duration time.Duration + ServiceMetrics *ServiceMetrics +} + +// ServiceMetrics captures provider-side execution details which are useful +// for separating queueing and planning from engine work. It is optional so +// PGWire and Trino keep their existing artifact contract. +type ServiceMetrics struct { + QueueDuration time.Duration `json:"queue_duration_ns"` + PlanningDuration time.Duration `json:"planning_duration_ns"` + EngineDuration time.Duration `json:"engine_duration_ns"` + ServiceDuration time.Duration `json:"service_duration_ns"` + BytesScanned int64 `json:"bytes_scanned"` + DPUCount float64 `json:"dpu_count"` + ResultReused bool `json:"result_reused"` + EngineVersion string `json:"engine_version"` } type QueryResult struct { - QueryID string `json:"query_id"` - IntentID string `json:"intent_id"` - MeasureIteration int `json:"measure_iteration"` - Protocol Protocol `json:"protocol"` - Status string `json:"status"` - Error string `json:"error,omitempty"` - ErrorClass string `json:"error_class,omitempty"` - Rows int64 `json:"rows"` - Duration time.Duration `json:"duration_ns"` - StartedAt time.Time `json:"started_at"` + QueryID string `json:"query_id"` + IntentID string `json:"intent_id"` + MeasureIteration int `json:"measure_iteration"` + Protocol Protocol `json:"protocol"` + Status string `json:"status"` + Error string `json:"error,omitempty"` + ErrorClass string `json:"error_class,omitempty"` + Rows int64 `json:"rows"` + Duration time.Duration `json:"duration_ns"` + StartedAt time.Time `json:"started_at"` + ServiceMetrics *ServiceMetrics `json:"service_metrics,omitempty"` } type RunSummary struct { diff --git a/tests/perf/drivers/athena/aws_client.go b/tests/perf/drivers/athena/aws_client.go new file mode 100644 index 000000000..c0d79b3f1 --- /dev/null +++ b/tests/perf/drivers/athena/aws_client.go @@ -0,0 +1,123 @@ +package athena + +import ( + "context" + "fmt" + "time" + + "github.com/aws/aws-sdk-go-v2/aws" + awsconfig "github.com/aws/aws-sdk-go-v2/config" + awsathena "github.com/aws/aws-sdk-go-v2/service/athena" + athenatypes "github.com/aws/aws-sdk-go-v2/service/athena/types" +) + +type athenaAPI interface { + StartQueryExecution(context.Context, *awsathena.StartQueryExecutionInput, ...func(*awsathena.Options)) (*awsathena.StartQueryExecutionOutput, error) + GetQueryExecution(context.Context, *awsathena.GetQueryExecutionInput, ...func(*awsathena.Options)) (*awsathena.GetQueryExecutionOutput, error) + GetQueryResults(context.Context, *awsathena.GetQueryResultsInput, ...func(*awsathena.Options)) (*awsathena.GetQueryResultsOutput, error) + StopQueryExecution(context.Context, *awsathena.StopQueryExecutionInput, ...func(*awsathena.Options)) (*awsathena.StopQueryExecutionOutput, error) +} + +type awsClient struct { + api athenaAPI +} + +func newAWSClient(ctx context.Context, region string) (*awsClient, error) { + options := []func(*awsconfig.LoadOptions) error{} + if region != "" { + options = append(options, awsconfig.WithRegion(region)) + } + cfg, err := awsconfig.LoadDefaultConfig(ctx, options...) + if err != nil { + return nil, fmt.Errorf("load AWS configuration for Athena: %w", err) + } + return &awsClient{api: awsathena.NewFromConfig(cfg)}, nil +} + +func (c *awsClient) StartQuery(ctx context.Context, input StartQueryInput) (string, error) { + output, err := c.api.StartQueryExecution(ctx, &awsathena.StartQueryExecutionInput{ + QueryString: aws.String(input.SQL), + WorkGroup: aws.String(input.WorkGroup), + QueryExecutionContext: &athenatypes.QueryExecutionContext{ + Catalog: aws.String(input.Catalog), + Database: aws.String(input.Database), + }, + ResultConfiguration: &athenatypes.ResultConfiguration{ + OutputLocation: aws.String(input.OutputLocation), + }, + ResultReuseConfiguration: &athenatypes.ResultReuseConfiguration{ + ResultReuseByAgeConfiguration: &athenatypes.ResultReuseByAgeConfiguration{ + Enabled: input.ResultReuseEnabled, + }, + }, + }) + if err != nil { + return "", err + } + if output.QueryExecutionId == nil || *output.QueryExecutionId == "" { + return "", fmt.Errorf("athena returned an empty query execution ID") + } + return *output.QueryExecutionId, nil +} + +func (c *awsClient) GetQuery(ctx context.Context, queryID string) (QueryExecution, error) { + output, err := c.api.GetQueryExecution(ctx, &awsathena.GetQueryExecutionInput{QueryExecutionId: aws.String(queryID)}) + if err != nil { + return QueryExecution{}, err + } + if output.QueryExecution == nil || output.QueryExecution.Status == nil { + return QueryExecution{}, fmt.Errorf("athena returned incomplete execution state") + } + execution := QueryExecution{ + State: QueryState(output.QueryExecution.Status.State), + StateChangeReason: aws.ToString(output.QueryExecution.Status.StateChangeReason), + } + if output.QueryExecution.ResultConfiguration != nil { + execution.OutputLocation = aws.ToString(output.QueryExecution.ResultConfiguration.OutputLocation) + } + if output.QueryExecution.EngineVersion != nil { + execution.EngineVersion = aws.ToString(output.QueryExecution.EngineVersion.EffectiveEngineVersion) + } + if statistics := output.QueryExecution.Statistics; statistics != nil { + execution.Statistics = QueryStatistics{ + QueueDuration: millis(statistics.QueryQueueTimeInMillis), + PlanningDuration: millis(statistics.QueryPlanningTimeInMillis), + EngineDuration: millis(statistics.EngineExecutionTimeInMillis), + ServiceDuration: millis(statistics.TotalExecutionTimeInMillis), + BytesScanned: aws.ToInt64(statistics.DataScannedInBytes), + DPUCount: aws.ToFloat64(statistics.DpuCount), + } + if statistics.ResultReuseInformation != nil { + execution.Statistics.ResultReused = statistics.ResultReuseInformation.ReusedPreviousResult + } + } + return execution, nil +} + +func (c *awsClient) GetResults(ctx context.Context, queryID, nextToken string) (ResultPage, error) { + input := &awsathena.GetQueryResultsInput{ + QueryExecutionId: aws.String(queryID), + MaxResults: aws.Int32(1000), + } + if nextToken != "" { + input.NextToken = aws.String(nextToken) + } + output, err := c.api.GetQueryResults(ctx, input) + if err != nil { + return ResultPage{}, err + } + page := ResultPage{NextToken: aws.ToString(output.NextToken)} + if output.ResultSet != nil { + page.RowCount = int64(len(output.ResultSet.Rows)) + } + return page, nil +} + +func (c *awsClient) StopQuery(ctx context.Context, queryID string) error { + _, err := c.api.StopQueryExecution(ctx, &awsathena.StopQueryExecutionInput{QueryExecutionId: aws.String(queryID)}) + return err +} + +func millis(value *int64) time.Duration { + return time.Duration(aws.ToInt64(value)) * time.Millisecond +} diff --git a/tests/perf/drivers/athena/driver.go b/tests/perf/drivers/athena/driver.go new file mode 100644 index 000000000..ad5893800 --- /dev/null +++ b/tests/perf/drivers/athena/driver.go @@ -0,0 +1,247 @@ +package athena + +import ( + "context" + "fmt" + "strings" + "time" + + perfcore "github.com/posthog/duckgres/tests/perf/core" +) + +const ( + defaultCatalog = "AwsDataCatalog" + defaultPollInterval = 500 * time.Millisecond + defaultQueryTimeout = 30 * time.Minute + stopTimeout = 5 * time.Second +) + +type QueryState string + +const ( + QueryStateQueued QueryState = "QUEUED" + QueryStateRunning QueryState = "RUNNING" + QueryStateSucceeded QueryState = "SUCCEEDED" + QueryStateFailed QueryState = "FAILED" + QueryStateCancelled QueryState = "CANCELLED" +) + +type ConnectionConfig struct { + Region string + WorkGroup string + Catalog string + Database string + OutputLocation string + PollInterval time.Duration + QueryTimeout time.Duration +} + +type StartQueryInput struct { + SQL string + WorkGroup string + Catalog string + Database string + OutputLocation string + ResultReuseEnabled bool +} + +type QueryStatistics struct { + QueueDuration time.Duration + PlanningDuration time.Duration + EngineDuration time.Duration + ServiceDuration time.Duration + BytesScanned int64 + DPUCount float64 + ResultReused bool +} + +type QueryExecution struct { + State QueryState + StateChangeReason string + OutputLocation string + EngineVersion string + Statistics QueryStatistics +} + +type ResultPage struct { + RowCount int64 + NextToken string +} + +type Client interface { + StartQuery(context.Context, StartQueryInput) (string, error) + GetQuery(context.Context, string) (QueryExecution, error) + GetResults(context.Context, string, string) (ResultPage, error) + StopQuery(context.Context, string) error +} + +type DriverOptions struct { + Now func() time.Time + Sleep func(context.Context, time.Duration) error +} + +type Driver struct { + client Client + cfg ConnectionConfig + now func() time.Time + sleep func(context.Context, time.Duration) error +} + +func New(ctx context.Context, cfg ConnectionConfig) (*Driver, error) { + client, err := newAWSClient(ctx, cfg.Region) + if err != nil { + return nil, err + } + return NewWithClient(client, cfg, DriverOptions{}) +} + +func NewWithClient(client Client, cfg ConnectionConfig, options DriverOptions) (*Driver, error) { + if client == nil { + return nil, fmt.Errorf("athena client is required") + } + if strings.TrimSpace(cfg.WorkGroup) == "" { + return nil, fmt.Errorf("athena workgroup is required") + } + if strings.TrimSpace(cfg.Database) == "" { + return nil, fmt.Errorf("athena database is required") + } + if !strings.HasPrefix(cfg.OutputLocation, "s3://") { + return nil, fmt.Errorf("athena output location must be an s3:// URI") + } + if cfg.Catalog == "" { + cfg.Catalog = defaultCatalog + } + if cfg.PollInterval <= 0 { + cfg.PollInterval = defaultPollInterval + } + if cfg.QueryTimeout <= 0 { + cfg.QueryTimeout = defaultQueryTimeout + } + if options.Now == nil { + options.Now = time.Now + } + if options.Sleep == nil { + options.Sleep = sleepWithContext + } + return &Driver{client: client, cfg: cfg, now: options.Now, sleep: options.Sleep}, nil +} + +func (d *Driver) Protocol() perfcore.Protocol { return perfcore.ProtocolAthena } + +func (d *Driver) Execute(ctx context.Context, query perfcore.Query, args []any) (result perfcore.ExecutionResult, err error) { + if len(args) > 0 { + return result, fmt.Errorf("athena perf queries do not support positional parameters") + } + queryCtx, cancel := context.WithTimeout(ctx, d.cfg.QueryTimeout) + defer cancel() + + startedAt := d.now() + queryID, err := d.client.StartQuery(queryCtx, StartQueryInput{ + SQL: query.CanonicalSQL(), + WorkGroup: d.cfg.WorkGroup, + Catalog: d.cfg.Catalog, + Database: d.cfg.Database, + OutputLocation: d.cfg.OutputLocation, + ResultReuseEnabled: false, + }) + if err != nil { + return result, fmt.Errorf("start Athena query: %w", err) + } + completed := false + defer func() { + if completed { + return + } + stopCtx, stopCancel := context.WithTimeout(context.Background(), stopTimeout) + defer stopCancel() + _ = d.client.StopQuery(stopCtx, queryID) + }() + + var execution QueryExecution + for { + execution, err = d.client.GetQuery(queryCtx, queryID) + if err != nil { + return result, fmt.Errorf("get Athena query %s: %w", queryID, err) + } + switch execution.State { + case QueryStateQueued, QueryStateRunning: + if err := d.sleep(queryCtx, d.cfg.PollInterval); err != nil { + return result, err + } + case QueryStateSucceeded: + completed = true + goto queryComplete + case QueryStateFailed, QueryStateCancelled: + completed = true + return result, fmt.Errorf("athena query %s ended in state %s: %s", queryID, execution.State, execution.StateChangeReason) + default: + return result, fmt.Errorf("athena query %s returned unknown state %q", queryID, execution.State) + } + } + +queryComplete: + if execution.Statistics.ResultReused { + return result, fmt.Errorf("athena query %s reused a previous result despite result reuse being disabled", queryID) + } + if !outputWithinRoot(execution.OutputLocation, d.cfg.OutputLocation) { + return result, fmt.Errorf("athena query output %q is outside configured output location %q", execution.OutputLocation, d.cfg.OutputLocation) + } + + rows, err := d.countRows(queryCtx, queryID) + if err != nil { + return result, err + } + result.Rows = rows + result.Duration = d.now().Sub(startedAt) + result.ServiceMetrics = &perfcore.ServiceMetrics{ + QueueDuration: execution.Statistics.QueueDuration, + PlanningDuration: execution.Statistics.PlanningDuration, + EngineDuration: execution.Statistics.EngineDuration, + ServiceDuration: execution.Statistics.ServiceDuration, + BytesScanned: execution.Statistics.BytesScanned, + DPUCount: execution.Statistics.DPUCount, + ResultReused: execution.Statistics.ResultReused, + EngineVersion: execution.EngineVersion, + } + return result, nil +} + +func (d *Driver) countRows(ctx context.Context, queryID string) (int64, error) { + var rows int64 + var nextToken string + firstPage := true + for { + page, err := d.client.GetResults(ctx, queryID, nextToken) + if err != nil { + return 0, fmt.Errorf("get Athena query results %s: %w", queryID, err) + } + pageRows := page.RowCount + if firstPage && pageRows > 0 { + pageRows-- // Athena returns the column header as the first result row. + } + rows += pageRows + firstPage = false + if page.NextToken == "" { + return rows, nil + } + nextToken = page.NextToken + } +} + +func (d *Driver) Close() error { return nil } + +func sleepWithContext(ctx context.Context, duration time.Duration) error { + timer := time.NewTimer(duration) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } +} + +func outputWithinRoot(outputLocation, configuredRoot string) bool { + root := strings.TrimSuffix(configuredRoot, "/") + "/" + return strings.HasPrefix(outputLocation, root) +} diff --git a/tests/perf/drivers/athena/driver_test.go b/tests/perf/drivers/athena/driver_test.go new file mode 100644 index 000000000..ccb27f8e1 --- /dev/null +++ b/tests/perf/drivers/athena/driver_test.go @@ -0,0 +1,215 @@ +package athena + +import ( + "context" + "errors" + "reflect" + "strings" + "testing" + "time" + + perfcore "github.com/posthog/duckgres/tests/perf/core" +) + +func TestDriverExecutesOnDemandQueryWithoutResultReuseAndCountsAllRows(t *testing.T) { + client := &fakeClient{ + executions: []QueryExecution{ + {State: QueryStateQueued}, + {State: QueryStateRunning}, + { + State: QueryStateSucceeded, + OutputLocation: "s3://benchmark-results/run/query.csv", + EngineVersion: "Athena engine version 3", + Statistics: QueryStatistics{ + QueueDuration: 120 * time.Millisecond, + PlanningDuration: 80 * time.Millisecond, + EngineDuration: 2 * time.Second, + ServiceDuration: 2*time.Second + 300*time.Millisecond, + BytesScanned: 4096, + }, + }, + }, + resultPages: []ResultPage{ + {RowCount: 3, NextToken: "page-2"}, // header + two data rows + {RowCount: 2}, + }, + } + now := time.Unix(1700000000, 0) + driver, err := NewWithClient(client, ConnectionConfig{ + WorkGroup: "benchmark", + Catalog: "AwsDataCatalog", + Database: "benchmark_frozen", + OutputLocation: "s3://benchmark-results/run/", + PollInterval: time.Millisecond, + QueryTimeout: time.Minute, + }, DriverOptions{ + Now: func() time.Time { + now = now.Add(time.Second) + return now + }, + Sleep: func(context.Context, time.Duration) error { return nil }, + }) + if err != nil { + t.Fatalf("NewWithClient returned error: %v", err) + } + + result, err := driver.Execute(context.Background(), perfcore.Query{PGWireSQL: "SELECT * FROM events"}, nil) + if err != nil { + t.Fatalf("Execute returned error: %v", err) + } + if result.Rows != 4 { + t.Fatalf("rows = %d, want 4 data rows excluding the Athena header", result.Rows) + } + if result.Duration != time.Second { + t.Fatalf("duration = %s, want end-to-end duration of 1s", result.Duration) + } + if got, want := client.startInput, (StartQueryInput{ + SQL: "SELECT * FROM events", + WorkGroup: "benchmark", + Catalog: "AwsDataCatalog", + Database: "benchmark_frozen", + OutputLocation: "s3://benchmark-results/run/", + ResultReuseEnabled: false, + }); !reflect.DeepEqual(got, want) { + t.Fatalf("start input = %+v, want %+v", got, want) + } + if got, want := client.resultTokens, []string{"", "page-2"}; !reflect.DeepEqual(got, want) { + t.Fatalf("result page tokens = %v, want %v", got, want) + } + if client.stopCalls != 0 { + t.Fatalf("StopQuery calls = %d, want 0 after success", client.stopCalls) + } + if result.ServiceMetrics == nil || result.ServiceMetrics.BytesScanned != 4096 || result.ServiceMetrics.EngineVersion != "Athena engine version 3" { + t.Fatalf("service metrics = %+v, want Athena execution statistics", result.ServiceMetrics) + } +} + +func TestDriverStopsAthenaQueryWhenContextIsCancelled(t *testing.T) { + client := &fakeClient{executions: []QueryExecution{{State: QueryStateRunning}}} + ctx, cancel := context.WithCancel(context.Background()) + driver, err := NewWithClient(client, ConnectionConfig{ + WorkGroup: "benchmark", + Database: "benchmark_frozen", + OutputLocation: "s3://benchmark-results/run/", + PollInterval: time.Millisecond, + QueryTimeout: time.Minute, + }, DriverOptions{ + Sleep: func(context.Context, time.Duration) error { + cancel() + return ctx.Err() + }, + }) + if err != nil { + t.Fatalf("NewWithClient returned error: %v", err) + } + + _, err = driver.Execute(ctx, perfcore.Query{PGWireSQL: "SELECT 1"}, nil) + if !errors.Is(err, context.Canceled) { + t.Fatalf("Execute error = %v, want context.Canceled", err) + } + if client.stopCalls != 1 || client.stoppedID != "query-1" { + t.Fatalf("StopQuery = %d calls for %q, want one call for query-1", client.stopCalls, client.stoppedID) + } +} + +func TestDriverReturnsAthenaFailureReason(t *testing.T) { + client := &fakeClient{executions: []QueryExecution{{ + State: QueryStateFailed, + StateChangeReason: "scan limit exceeded", + }}} + driver, err := NewWithClient(client, ConnectionConfig{ + WorkGroup: "benchmark", + Database: "benchmark_frozen", + OutputLocation: "s3://benchmark-results/run/", + }, DriverOptions{}) + if err != nil { + t.Fatalf("NewWithClient returned error: %v", err) + } + + _, err = driver.Execute(context.Background(), perfcore.Query{PGWireSQL: "SELECT 1"}, nil) + if err == nil || !strings.Contains(err.Error(), "scan limit exceeded") { + t.Fatalf("Execute error = %v, want Athena failure reason", err) + } +} + +func TestDriverRejectsResultOutsideConfiguredOutputRoot(t *testing.T) { + client := &fakeClient{executions: []QueryExecution{{ + State: QueryStateSucceeded, + OutputLocation: "s3://unexpected-bucket/query.csv", + }}} + driver, err := NewWithClient(client, ConnectionConfig{ + WorkGroup: "benchmark", + Database: "benchmark_frozen", + OutputLocation: "s3://benchmark-results/run/", + }, DriverOptions{}) + if err != nil { + t.Fatalf("NewWithClient returned error: %v", err) + } + + _, err = driver.Execute(context.Background(), perfcore.Query{PGWireSQL: "SELECT 1"}, nil) + if err == nil || !strings.Contains(err.Error(), "outside configured output location") { + t.Fatalf("Execute error = %v, want output-location validation error", err) + } +} + +func TestDriverRejectsUnexpectedlyReusedAthenaResult(t *testing.T) { + client := &fakeClient{executions: []QueryExecution{{ + State: QueryStateSucceeded, + OutputLocation: "s3://benchmark-results/run/query.csv", + Statistics: QueryStatistics{ResultReused: true}, + }}} + driver, err := NewWithClient(client, ConnectionConfig{ + WorkGroup: "benchmark", + Database: "benchmark_frozen", + OutputLocation: "s3://benchmark-results/run/", + }, DriverOptions{}) + if err != nil { + t.Fatalf("NewWithClient returned error: %v", err) + } + + _, err = driver.Execute(context.Background(), perfcore.Query{PGWireSQL: "SELECT 1"}, nil) + if err == nil || !strings.Contains(err.Error(), "reused a previous result") { + t.Fatalf("Execute error = %v, want invalid result-reuse error", err) + } +} + +type fakeClient struct { + startInput StartQueryInput + executions []QueryExecution + resultPages []ResultPage + resultTokens []string + stopCalls int + stoppedID string +} + +func (f *fakeClient) StartQuery(_ context.Context, input StartQueryInput) (string, error) { + f.startInput = input + return "query-1", nil +} + +func (f *fakeClient) GetQuery(_ context.Context, _ string) (QueryExecution, error) { + if len(f.executions) == 0 { + return QueryExecution{}, errors.New("unexpected GetQuery") + } + execution := f.executions[0] + if len(f.executions) > 1 { + f.executions = f.executions[1:] + } + return execution, nil +} + +func (f *fakeClient) GetResults(_ context.Context, _ string, nextToken string) (ResultPage, error) { + f.resultTokens = append(f.resultTokens, nextToken) + if len(f.resultPages) == 0 { + return ResultPage{}, errors.New("unexpected GetResults") + } + page := f.resultPages[0] + f.resultPages = f.resultPages[1:] + return page, nil +} + +func (f *fakeClient) StopQuery(_ context.Context, queryID string) error { + f.stopCalls++ + f.stoppedID = queryID + return nil +} diff --git a/tests/perf/harness_test.go b/tests/perf/harness_test.go index 90234ce63..042f4ed11 100644 --- a/tests/perf/harness_test.go +++ b/tests/perf/harness_test.go @@ -153,7 +153,7 @@ func TestGoldenQueryPerformanceHarness(t *testing.T) { t.Fatalf("write runner.log: %v", err) } - for _, artifact := range []string{"summary.json", "query_results.csv", "server_metrics.prom", "runner.log"} { + for _, artifact := range []string{"summary.json", "query_results.csv", "query_service_metrics.csv", "server_metrics.prom", "runner.log"} { if _, err := os.Stat(filepath.Join(outputDir, artifact)); err != nil { t.Fatalf("expected artifact %s: %v", artifact, err) } diff --git a/tests/perf/queries/ducklake_posthog_tables.yaml b/tests/perf/queries/ducklake_posthog_tables.yaml index edb6721a1..4a6e98772 100644 --- a/tests/perf/queries/ducklake_posthog_tables.yaml +++ b/tests/perf/queries/ducklake_posthog_tables.yaml @@ -5,6 +5,7 @@ dataset_scale: 1 targets: - pgwire - trino + - athena warmup_iterations: 1 measure_iterations: 4 @@ -15,6 +16,9 @@ relation_variants: ducklake_table: events: posthog.events persons: posthog.persons + athena_external: + events: events + persons: persons paired_queries: - query_id_base: q_events_total_balanced_v3 From 9e7df40cfcb287e04d8d11d5f73a7838056983e3 Mon Sep 17 00:00:00 2001 From: Bill Guowei Yang Date: Mon, 7 Sep 2026 09:16:10 -0400 Subject: [PATCH 2/4] fix(perf): preserve Athena error metrics and simplify SDK client --- tests/perf/README.md | 6 +- tests/perf/drivers/athena/aws_client.go | 123 ---------- tests/perf/drivers/athena/driver.go | 194 +++++++-------- tests/perf/drivers/athena/driver_test.go | 293 +++++++++++++---------- 4 files changed, 264 insertions(+), 352 deletions(-) delete mode 100644 tests/perf/drivers/athena/aws_client.go diff --git a/tests/perf/README.md b/tests/perf/README.md index fad9a249d..6414c7557 100644 --- a/tests/perf/README.md +++ b/tests/perf/README.md @@ -16,7 +16,11 @@ execution and rejects a response which reports reuse. Timing is end to end: it includes queueing, engine execution, and paginated result retrieval. The driver defaults to catalog `AwsDataCatalog`, a 500ms status poll interval, and a 30-minute query timeout. Workgroup, database, result prefix, and AWS region -are explicit scenario settings. Cancellation stops any unfinished query. +are explicit scenario settings. Cancellation attempts to stop any unfinished +query; Athena's independent server-side timeout still applies if that request +fails. Available terminal service statistics are retained even when execution, +result validation, or result retrieval fails. Executions without service +statistics do not emit a service-metrics row. The Trino driver requires an HTTPS coordinator and always verifies its TLS certificate. It uses system roots by default, or the explicitly configured CA diff --git a/tests/perf/drivers/athena/aws_client.go b/tests/perf/drivers/athena/aws_client.go deleted file mode 100644 index c0d79b3f1..000000000 --- a/tests/perf/drivers/athena/aws_client.go +++ /dev/null @@ -1,123 +0,0 @@ -package athena - -import ( - "context" - "fmt" - "time" - - "github.com/aws/aws-sdk-go-v2/aws" - awsconfig "github.com/aws/aws-sdk-go-v2/config" - awsathena "github.com/aws/aws-sdk-go-v2/service/athena" - athenatypes "github.com/aws/aws-sdk-go-v2/service/athena/types" -) - -type athenaAPI interface { - StartQueryExecution(context.Context, *awsathena.StartQueryExecutionInput, ...func(*awsathena.Options)) (*awsathena.StartQueryExecutionOutput, error) - GetQueryExecution(context.Context, *awsathena.GetQueryExecutionInput, ...func(*awsathena.Options)) (*awsathena.GetQueryExecutionOutput, error) - GetQueryResults(context.Context, *awsathena.GetQueryResultsInput, ...func(*awsathena.Options)) (*awsathena.GetQueryResultsOutput, error) - StopQueryExecution(context.Context, *awsathena.StopQueryExecutionInput, ...func(*awsathena.Options)) (*awsathena.StopQueryExecutionOutput, error) -} - -type awsClient struct { - api athenaAPI -} - -func newAWSClient(ctx context.Context, region string) (*awsClient, error) { - options := []func(*awsconfig.LoadOptions) error{} - if region != "" { - options = append(options, awsconfig.WithRegion(region)) - } - cfg, err := awsconfig.LoadDefaultConfig(ctx, options...) - if err != nil { - return nil, fmt.Errorf("load AWS configuration for Athena: %w", err) - } - return &awsClient{api: awsathena.NewFromConfig(cfg)}, nil -} - -func (c *awsClient) StartQuery(ctx context.Context, input StartQueryInput) (string, error) { - output, err := c.api.StartQueryExecution(ctx, &awsathena.StartQueryExecutionInput{ - QueryString: aws.String(input.SQL), - WorkGroup: aws.String(input.WorkGroup), - QueryExecutionContext: &athenatypes.QueryExecutionContext{ - Catalog: aws.String(input.Catalog), - Database: aws.String(input.Database), - }, - ResultConfiguration: &athenatypes.ResultConfiguration{ - OutputLocation: aws.String(input.OutputLocation), - }, - ResultReuseConfiguration: &athenatypes.ResultReuseConfiguration{ - ResultReuseByAgeConfiguration: &athenatypes.ResultReuseByAgeConfiguration{ - Enabled: input.ResultReuseEnabled, - }, - }, - }) - if err != nil { - return "", err - } - if output.QueryExecutionId == nil || *output.QueryExecutionId == "" { - return "", fmt.Errorf("athena returned an empty query execution ID") - } - return *output.QueryExecutionId, nil -} - -func (c *awsClient) GetQuery(ctx context.Context, queryID string) (QueryExecution, error) { - output, err := c.api.GetQueryExecution(ctx, &awsathena.GetQueryExecutionInput{QueryExecutionId: aws.String(queryID)}) - if err != nil { - return QueryExecution{}, err - } - if output.QueryExecution == nil || output.QueryExecution.Status == nil { - return QueryExecution{}, fmt.Errorf("athena returned incomplete execution state") - } - execution := QueryExecution{ - State: QueryState(output.QueryExecution.Status.State), - StateChangeReason: aws.ToString(output.QueryExecution.Status.StateChangeReason), - } - if output.QueryExecution.ResultConfiguration != nil { - execution.OutputLocation = aws.ToString(output.QueryExecution.ResultConfiguration.OutputLocation) - } - if output.QueryExecution.EngineVersion != nil { - execution.EngineVersion = aws.ToString(output.QueryExecution.EngineVersion.EffectiveEngineVersion) - } - if statistics := output.QueryExecution.Statistics; statistics != nil { - execution.Statistics = QueryStatistics{ - QueueDuration: millis(statistics.QueryQueueTimeInMillis), - PlanningDuration: millis(statistics.QueryPlanningTimeInMillis), - EngineDuration: millis(statistics.EngineExecutionTimeInMillis), - ServiceDuration: millis(statistics.TotalExecutionTimeInMillis), - BytesScanned: aws.ToInt64(statistics.DataScannedInBytes), - DPUCount: aws.ToFloat64(statistics.DpuCount), - } - if statistics.ResultReuseInformation != nil { - execution.Statistics.ResultReused = statistics.ResultReuseInformation.ReusedPreviousResult - } - } - return execution, nil -} - -func (c *awsClient) GetResults(ctx context.Context, queryID, nextToken string) (ResultPage, error) { - input := &awsathena.GetQueryResultsInput{ - QueryExecutionId: aws.String(queryID), - MaxResults: aws.Int32(1000), - } - if nextToken != "" { - input.NextToken = aws.String(nextToken) - } - output, err := c.api.GetQueryResults(ctx, input) - if err != nil { - return ResultPage{}, err - } - page := ResultPage{NextToken: aws.ToString(output.NextToken)} - if output.ResultSet != nil { - page.RowCount = int64(len(output.ResultSet.Rows)) - } - return page, nil -} - -func (c *awsClient) StopQuery(ctx context.Context, queryID string) error { - _, err := c.api.StopQueryExecution(ctx, &awsathena.StopQueryExecutionInput{QueryExecutionId: aws.String(queryID)}) - return err -} - -func millis(value *int64) time.Duration { - return time.Duration(aws.ToInt64(value)) * time.Millisecond -} diff --git a/tests/perf/drivers/athena/driver.go b/tests/perf/drivers/athena/driver.go index ad5893800..090c61824 100644 --- a/tests/perf/drivers/athena/driver.go +++ b/tests/perf/drivers/athena/driver.go @@ -6,6 +6,11 @@ import ( "strings" "time" + "github.com/aws/aws-sdk-go-v2/aws" + awsconfig "github.com/aws/aws-sdk-go-v2/config" + awsathena "github.com/aws/aws-sdk-go-v2/service/athena" + athenatypes "github.com/aws/aws-sdk-go-v2/service/athena/types" + perfcore "github.com/posthog/duckgres/tests/perf/core" ) @@ -16,16 +21,6 @@ const ( stopTimeout = 5 * time.Second ) -type QueryState string - -const ( - QueryStateQueued QueryState = "QUEUED" - QueryStateRunning QueryState = "RUNNING" - QueryStateSucceeded QueryState = "SUCCEEDED" - QueryStateFailed QueryState = "FAILED" - QueryStateCancelled QueryState = "CANCELLED" -) - type ConnectionConfig struct { Region string WorkGroup string @@ -36,66 +31,33 @@ type ConnectionConfig struct { QueryTimeout time.Duration } -type StartQueryInput struct { - SQL string - WorkGroup string - Catalog string - Database string - OutputLocation string - ResultReuseEnabled bool -} - -type QueryStatistics struct { - QueueDuration time.Duration - PlanningDuration time.Duration - EngineDuration time.Duration - ServiceDuration time.Duration - BytesScanned int64 - DPUCount float64 - ResultReused bool -} - -type QueryExecution struct { - State QueryState - StateChangeReason string - OutputLocation string - EngineVersion string - Statistics QueryStatistics -} - -type ResultPage struct { - RowCount int64 - NextToken string -} - -type Client interface { - StartQuery(context.Context, StartQueryInput) (string, error) - GetQuery(context.Context, string) (QueryExecution, error) - GetResults(context.Context, string, string) (ResultPage, error) - StopQuery(context.Context, string) error -} - -type DriverOptions struct { - Now func() time.Time - Sleep func(context.Context, time.Duration) error +type athenaAPI interface { + StartQueryExecution(context.Context, *awsathena.StartQueryExecutionInput, ...func(*awsathena.Options)) (*awsathena.StartQueryExecutionOutput, error) + GetQueryExecution(context.Context, *awsathena.GetQueryExecutionInput, ...func(*awsathena.Options)) (*awsathena.GetQueryExecutionOutput, error) + GetQueryResults(context.Context, *awsathena.GetQueryResultsInput, ...func(*awsathena.Options)) (*awsathena.GetQueryResultsOutput, error) + StopQueryExecution(context.Context, *awsathena.StopQueryExecutionInput, ...func(*awsathena.Options)) (*awsathena.StopQueryExecutionOutput, error) } type Driver struct { - client Client + client athenaAPI cfg ConnectionConfig now func() time.Time sleep func(context.Context, time.Duration) error } func New(ctx context.Context, cfg ConnectionConfig) (*Driver, error) { - client, err := newAWSClient(ctx, cfg.Region) + options := []func(*awsconfig.LoadOptions) error{} + if cfg.Region != "" { + options = append(options, awsconfig.WithRegion(cfg.Region)) + } + awsCfg, err := awsconfig.LoadDefaultConfig(ctx, options...) if err != nil { - return nil, err + return nil, fmt.Errorf("load AWS configuration for Athena: %w", err) } - return NewWithClient(client, cfg, DriverOptions{}) + return newWithClient(awsathena.NewFromConfig(awsCfg), cfg) } -func NewWithClient(client Client, cfg ConnectionConfig, options DriverOptions) (*Driver, error) { +func newWithClient(client athenaAPI, cfg ConnectionConfig) (*Driver, error) { if client == nil { return nil, fmt.Errorf("athena client is required") } @@ -117,13 +79,7 @@ func NewWithClient(client Client, cfg ConnectionConfig, options DriverOptions) ( if cfg.QueryTimeout <= 0 { cfg.QueryTimeout = defaultQueryTimeout } - if options.Now == nil { - options.Now = time.Now - } - if options.Sleep == nil { - options.Sleep = sleepWithContext - } - return &Driver{client: client, cfg: cfg, now: options.Now, sleep: options.Sleep}, nil + return &Driver{client: client, cfg: cfg, now: time.Now, sleep: sleepWithContext}, nil } func (d *Driver) Protocol() perfcore.Protocol { return perfcore.ProtocolAthena } @@ -136,17 +92,25 @@ func (d *Driver) Execute(ctx context.Context, query perfcore.Query, args []any) defer cancel() startedAt := d.now() - queryID, err := d.client.StartQuery(queryCtx, StartQueryInput{ - SQL: query.CanonicalSQL(), - WorkGroup: d.cfg.WorkGroup, - Catalog: d.cfg.Catalog, - Database: d.cfg.Database, - OutputLocation: d.cfg.OutputLocation, - ResultReuseEnabled: false, + defer func() { result.Duration = d.now().Sub(startedAt) }() + started, err := d.client.StartQueryExecution(queryCtx, &awsathena.StartQueryExecutionInput{ + QueryString: aws.String(query.CanonicalSQL()), + WorkGroup: aws.String(d.cfg.WorkGroup), + QueryExecutionContext: &athenatypes.QueryExecutionContext{ + Catalog: aws.String(d.cfg.Catalog), Database: aws.String(d.cfg.Database), + }, + ResultConfiguration: &athenatypes.ResultConfiguration{OutputLocation: aws.String(d.cfg.OutputLocation)}, + ResultReuseConfiguration: &athenatypes.ResultReuseConfiguration{ + ResultReuseByAgeConfiguration: &athenatypes.ResultReuseByAgeConfiguration{Enabled: false}, + }, }) if err != nil { return result, fmt.Errorf("start Athena query: %w", err) } + if started == nil || aws.ToString(started.QueryExecutionId) == "" { + return result, fmt.Errorf("athena returned an empty query execution ID") + } + queryID := started.QueryExecutionId completed := false defer func() { if completed { @@ -154,77 +118,101 @@ func (d *Driver) Execute(ctx context.Context, query perfcore.Query, args []any) } stopCtx, stopCancel := context.WithTimeout(context.Background(), stopTimeout) defer stopCancel() - _ = d.client.StopQuery(stopCtx, queryID) + _, _ = d.client.StopQueryExecution(stopCtx, &awsathena.StopQueryExecutionInput{QueryExecutionId: queryID}) }() - var execution QueryExecution + var execution *athenatypes.QueryExecution for { - execution, err = d.client.GetQuery(queryCtx, queryID) + output, err := d.client.GetQueryExecution(queryCtx, &awsathena.GetQueryExecutionInput{QueryExecutionId: queryID}) if err != nil { - return result, fmt.Errorf("get Athena query %s: %w", queryID, err) + return result, fmt.Errorf("get Athena query %s: %w", *queryID, err) + } + if output == nil || output.QueryExecution == nil || output.QueryExecution.Status == nil { + return result, fmt.Errorf("athena returned incomplete execution state for query %s", *queryID) } - switch execution.State { - case QueryStateQueued, QueryStateRunning: + execution = output.QueryExecution + switch execution.Status.State { + case athenatypes.QueryExecutionStateQueued, athenatypes.QueryExecutionStateRunning: if err := d.sleep(queryCtx, d.cfg.PollInterval); err != nil { return result, err } - case QueryStateSucceeded: + case athenatypes.QueryExecutionStateSucceeded, athenatypes.QueryExecutionStateFailed, athenatypes.QueryExecutionStateCancelled: completed = true goto queryComplete - case QueryStateFailed, QueryStateCancelled: - completed = true - return result, fmt.Errorf("athena query %s ended in state %s: %s", queryID, execution.State, execution.StateChangeReason) default: - return result, fmt.Errorf("athena query %s returned unknown state %q", queryID, execution.State) + return result, fmt.Errorf("athena query %s returned unknown state %q", *queryID, execution.Status.State) } } queryComplete: - if execution.Statistics.ResultReused { - return result, fmt.Errorf("athena query %s reused a previous result despite result reuse being disabled", queryID) + // Preserve final service statistics even if execution or result retrieval fails. + result.ServiceMetrics = serviceMetrics(execution) + if execution.Status.State != athenatypes.QueryExecutionStateSucceeded { + return result, fmt.Errorf("athena query %s ended in state %s: %s", *queryID, execution.Status.State, aws.ToString(execution.Status.StateChangeReason)) + } + if result.ServiceMetrics != nil && result.ServiceMetrics.ResultReused { + return result, fmt.Errorf("athena query %s reused a previous result despite result reuse being disabled", *queryID) } - if !outputWithinRoot(execution.OutputLocation, d.cfg.OutputLocation) { - return result, fmt.Errorf("athena query output %q is outside configured output location %q", execution.OutputLocation, d.cfg.OutputLocation) + var outputLocation string + if execution.ResultConfiguration != nil { + outputLocation = aws.ToString(execution.ResultConfiguration.OutputLocation) + } + if !outputWithinRoot(outputLocation, d.cfg.OutputLocation) { + return result, fmt.Errorf("athena query output %q is outside configured output location %q", outputLocation, d.cfg.OutputLocation) } - rows, err := d.countRows(queryCtx, queryID) + rows, err := d.countRows(queryCtx, *queryID) if err != nil { return result, err } result.Rows = rows - result.Duration = d.now().Sub(startedAt) - result.ServiceMetrics = &perfcore.ServiceMetrics{ - QueueDuration: execution.Statistics.QueueDuration, - PlanningDuration: execution.Statistics.PlanningDuration, - EngineDuration: execution.Statistics.EngineDuration, - ServiceDuration: execution.Statistics.ServiceDuration, - BytesScanned: execution.Statistics.BytesScanned, - DPUCount: execution.Statistics.DPUCount, - ResultReused: execution.Statistics.ResultReused, - EngineVersion: execution.EngineVersion, - } return result, nil } +func serviceMetrics(execution *athenatypes.QueryExecution) *perfcore.ServiceMetrics { + statistics := execution.Statistics + if statistics == nil { + return nil + } + metrics := &perfcore.ServiceMetrics{ + QueueDuration: time.Duration(aws.ToInt64(statistics.QueryQueueTimeInMillis)) * time.Millisecond, + PlanningDuration: time.Duration(aws.ToInt64(statistics.QueryPlanningTimeInMillis)) * time.Millisecond, + EngineDuration: time.Duration(aws.ToInt64(statistics.EngineExecutionTimeInMillis)) * time.Millisecond, + ServiceDuration: time.Duration(aws.ToInt64(statistics.TotalExecutionTimeInMillis)) * time.Millisecond, + BytesScanned: aws.ToInt64(statistics.DataScannedInBytes), + DPUCount: aws.ToFloat64(statistics.DpuCount), + } + if statistics.ResultReuseInformation != nil { + metrics.ResultReused = statistics.ResultReuseInformation.ReusedPreviousResult + } + if execution.EngineVersion != nil { + metrics.EngineVersion = aws.ToString(execution.EngineVersion.EffectiveEngineVersion) + } + return metrics +} + func (d *Driver) countRows(ctx context.Context, queryID string) (int64, error) { var rows int64 - var nextToken string + input := &awsathena.GetQueryResultsInput{QueryExecutionId: aws.String(queryID), MaxResults: aws.Int32(1000)} firstPage := true for { - page, err := d.client.GetResults(ctx, queryID, nextToken) + page, err := d.client.GetQueryResults(ctx, input) if err != nil { return 0, fmt.Errorf("get Athena query results %s: %w", queryID, err) } - pageRows := page.RowCount + var pageRows int64 + if page.ResultSet != nil { + pageRows = int64(len(page.ResultSet.Rows)) + } if firstPage && pageRows > 0 { pageRows-- // Athena returns the column header as the first result row. } rows += pageRows firstPage = false - if page.NextToken == "" { + if aws.ToString(page.NextToken) == "" { return rows, nil } - nextToken = page.NextToken + input.NextToken = page.NextToken } } diff --git a/tests/perf/drivers/athena/driver_test.go b/tests/perf/drivers/athena/driver_test.go index ccb27f8e1..5238f8e8a 100644 --- a/tests/perf/drivers/athena/driver_test.go +++ b/tests/perf/drivers/athena/driver_test.go @@ -8,50 +8,32 @@ import ( "testing" "time" + "github.com/aws/aws-sdk-go-v2/aws" + awsathena "github.com/aws/aws-sdk-go-v2/service/athena" + athenatypes "github.com/aws/aws-sdk-go-v2/service/athena/types" + perfcore "github.com/posthog/duckgres/tests/perf/core" ) func TestDriverExecutesOnDemandQueryWithoutResultReuseAndCountsAllRows(t *testing.T) { client := &fakeClient{ - executions: []QueryExecution{ - {State: QueryStateQueued}, - {State: QueryStateRunning}, - { - State: QueryStateSucceeded, - OutputLocation: "s3://benchmark-results/run/query.csv", - EngineVersion: "Athena engine version 3", - Statistics: QueryStatistics{ - QueueDuration: 120 * time.Millisecond, - PlanningDuration: 80 * time.Millisecond, - EngineDuration: 2 * time.Second, - ServiceDuration: 2*time.Second + 300*time.Millisecond, - BytesScanned: 4096, - }, - }, + executions: []*athenatypes.QueryExecution{ + {Status: &athenatypes.QueryExecutionStatus{State: athenatypes.QueryExecutionStateQueued}}, + {Status: &athenatypes.QueryExecutionStatus{State: athenatypes.QueryExecutionStateRunning}}, + terminalExecution(athenatypes.QueryExecutionStateSucceeded), }, - resultPages: []ResultPage{ - {RowCount: 3, NextToken: "page-2"}, // header + two data rows - {RowCount: 2}, + resultPages: []*awsathena.GetQueryResultsOutput{ + {ResultSet: &athenatypes.ResultSet{Rows: make([]athenatypes.Row, 3)}, NextToken: aws.String("page-2")}, + {ResultSet: &athenatypes.ResultSet{Rows: make([]athenatypes.Row, 2)}}, }, } + driver := testDriver(t, client) now := time.Unix(1700000000, 0) - driver, err := NewWithClient(client, ConnectionConfig{ - WorkGroup: "benchmark", - Catalog: "AwsDataCatalog", - Database: "benchmark_frozen", - OutputLocation: "s3://benchmark-results/run/", - PollInterval: time.Millisecond, - QueryTimeout: time.Minute, - }, DriverOptions{ - Now: func() time.Time { - now = now.Add(time.Second) - return now - }, - Sleep: func(context.Context, time.Duration) error { return nil }, - }) - if err != nil { - t.Fatalf("NewWithClient returned error: %v", err) + driver.now = func() time.Time { + now = now.Add(time.Second) + return now } + driver.sleep = func(context.Context, time.Duration) error { return nil } result, err := driver.Execute(context.Background(), perfcore.Query{PGWireSQL: "SELECT * FROM events"}, nil) if err != nil { @@ -63,153 +45,214 @@ func TestDriverExecutesOnDemandQueryWithoutResultReuseAndCountsAllRows(t *testin if result.Duration != time.Second { t.Fatalf("duration = %s, want end-to-end duration of 1s", result.Duration) } - if got, want := client.startInput, (StartQueryInput{ - SQL: "SELECT * FROM events", - WorkGroup: "benchmark", - Catalog: "AwsDataCatalog", - Database: "benchmark_frozen", - OutputLocation: "s3://benchmark-results/run/", - ResultReuseEnabled: false, - }); !reflect.DeepEqual(got, want) { - t.Fatalf("start input = %+v, want %+v", got, want) + wantInput := &awsathena.StartQueryExecutionInput{ + QueryString: aws.String("SELECT * FROM events"), + WorkGroup: aws.String("benchmark"), + QueryExecutionContext: &athenatypes.QueryExecutionContext{ + Catalog: aws.String("AwsDataCatalog"), Database: aws.String("benchmark_frozen"), + }, + ResultConfiguration: &athenatypes.ResultConfiguration{OutputLocation: aws.String("s3://benchmark-results/run/")}, + ResultReuseConfiguration: &athenatypes.ResultReuseConfiguration{ + ResultReuseByAgeConfiguration: &athenatypes.ResultReuseByAgeConfiguration{Enabled: false}, + }, + } + if !reflect.DeepEqual(client.startInput, wantInput) { + t.Fatalf("start input = %+v, want %+v", client.startInput, wantInput) } if got, want := client.resultTokens, []string{"", "page-2"}; !reflect.DeepEqual(got, want) { t.Fatalf("result page tokens = %v, want %v", got, want) } if client.stopCalls != 0 { - t.Fatalf("StopQuery calls = %d, want 0 after success", client.stopCalls) - } - if result.ServiceMetrics == nil || result.ServiceMetrics.BytesScanned != 4096 || result.ServiceMetrics.EngineVersion != "Athena engine version 3" { - t.Fatalf("service metrics = %+v, want Athena execution statistics", result.ServiceMetrics) + t.Fatalf("StopQueryExecution calls = %d, want 0 after success", client.stopCalls) } + assertServiceMetrics(t, result.ServiceMetrics) } func TestDriverStopsAthenaQueryWhenContextIsCancelled(t *testing.T) { - client := &fakeClient{executions: []QueryExecution{{State: QueryStateRunning}}} + client := &fakeClient{executions: []*athenatypes.QueryExecution{{ + Status: &athenatypes.QueryExecutionStatus{State: athenatypes.QueryExecutionStateRunning}, + }}} ctx, cancel := context.WithCancel(context.Background()) - driver, err := NewWithClient(client, ConnectionConfig{ - WorkGroup: "benchmark", - Database: "benchmark_frozen", - OutputLocation: "s3://benchmark-results/run/", - PollInterval: time.Millisecond, - QueryTimeout: time.Minute, - }, DriverOptions{ - Sleep: func(context.Context, time.Duration) error { - cancel() - return ctx.Err() - }, - }) - if err != nil { - t.Fatalf("NewWithClient returned error: %v", err) + defer cancel() + driver := testDriver(t, client) + driver.sleep = func(context.Context, time.Duration) error { + cancel() + return ctx.Err() } - _, err = driver.Execute(ctx, perfcore.Query{PGWireSQL: "SELECT 1"}, nil) + _, err := driver.Execute(ctx, perfcore.Query{PGWireSQL: "SELECT 1"}, nil) if !errors.Is(err, context.Canceled) { t.Fatalf("Execute error = %v, want context.Canceled", err) } if client.stopCalls != 1 || client.stoppedID != "query-1" { - t.Fatalf("StopQuery = %d calls for %q, want one call for query-1", client.stopCalls, client.stoppedID) + t.Fatalf("StopQueryExecution = %d calls for %q, want one call for query-1", client.stopCalls, client.stoppedID) } -} - -func TestDriverReturnsAthenaFailureReason(t *testing.T) { - client := &fakeClient{executions: []QueryExecution{{ - State: QueryStateFailed, - StateChangeReason: "scan limit exceeded", - }}} - driver, err := NewWithClient(client, ConnectionConfig{ - WorkGroup: "benchmark", - Database: "benchmark_frozen", - OutputLocation: "s3://benchmark-results/run/", - }, DriverOptions{}) - if err != nil { - t.Fatalf("NewWithClient returned error: %v", err) + if client.stopContextErr != nil { + t.Fatalf("stop context already cancelled: %v", client.stopContextErr) } +} - _, err = driver.Execute(context.Background(), perfcore.Query{PGWireSQL: "SELECT 1"}, nil) - if err == nil || !strings.Contains(err.Error(), "scan limit exceeded") { - t.Fatalf("Execute error = %v, want Athena failure reason", err) +func TestDriverPreservesMetricsWhenExecutionEndsWithError(t *testing.T) { + for _, state := range []athenatypes.QueryExecutionState{ + athenatypes.QueryExecutionStateFailed, athenatypes.QueryExecutionStateCancelled, athenatypes.QueryExecutionStateSucceeded, + } { + t.Run(string(state), func(t *testing.T) { + client := &fakeClient{executions: []*athenatypes.QueryExecution{terminalExecution(state)}} + // No result pages: the succeeded case fails while downloading results. + driver := testDriver(t, client) + result, err := driver.Execute(context.Background(), perfcore.Query{PGWireSQL: "SELECT 1"}, nil) + wantError := "scan limit exceeded" + if state == athenatypes.QueryExecutionStateSucceeded { + wantError = "get Athena query results query-1" + } + if err == nil || !strings.Contains(err.Error(), wantError) { + t.Fatalf("Execute error = %v, want %q", err, wantError) + } + assertServiceMetrics(t, result.ServiceMetrics) + if result.Duration <= 0 { + t.Fatalf("duration = %s, want elapsed time despite error", result.Duration) + } + if client.stopCalls != 0 { + t.Fatal("terminal execution should not be stopped again") + } + }) } } func TestDriverRejectsResultOutsideConfiguredOutputRoot(t *testing.T) { - client := &fakeClient{executions: []QueryExecution{{ - State: QueryStateSucceeded, - OutputLocation: "s3://unexpected-bucket/query.csv", - }}} - driver, err := NewWithClient(client, ConnectionConfig{ - WorkGroup: "benchmark", - Database: "benchmark_frozen", - OutputLocation: "s3://benchmark-results/run/", - }, DriverOptions{}) - if err != nil { - t.Fatalf("NewWithClient returned error: %v", err) - } - - _, err = driver.Execute(context.Background(), perfcore.Query{PGWireSQL: "SELECT 1"}, nil) + execution := terminalExecution(athenatypes.QueryExecutionStateSucceeded) + execution.ResultConfiguration.OutputLocation = aws.String("s3://unexpected-bucket/query.csv") + driver := testDriver(t, &fakeClient{executions: []*athenatypes.QueryExecution{execution}}) + result, err := driver.Execute(context.Background(), perfcore.Query{PGWireSQL: "SELECT 1"}, nil) if err == nil || !strings.Contains(err.Error(), "outside configured output location") { t.Fatalf("Execute error = %v, want output-location validation error", err) } + assertServiceMetrics(t, result.ServiceMetrics) } func TestDriverRejectsUnexpectedlyReusedAthenaResult(t *testing.T) { - client := &fakeClient{executions: []QueryExecution{{ - State: QueryStateSucceeded, - OutputLocation: "s3://benchmark-results/run/query.csv", - Statistics: QueryStatistics{ResultReused: true}, - }}} - driver, err := NewWithClient(client, ConnectionConfig{ - WorkGroup: "benchmark", - Database: "benchmark_frozen", + execution := terminalExecution(athenatypes.QueryExecutionStateSucceeded) + execution.Statistics.ResultReuseInformation = &athenatypes.ResultReuseInformation{ReusedPreviousResult: true} + driver := testDriver(t, &fakeClient{executions: []*athenatypes.QueryExecution{execution}}) + result, err := driver.Execute(context.Background(), perfcore.Query{PGWireSQL: "SELECT 1"}, nil) + if err == nil || !strings.Contains(err.Error(), "reused a previous result") { + t.Fatalf("Execute error = %v, want invalid result-reuse error", err) + } + if result.ServiceMetrics == nil || !result.ServiceMetrics.ResultReused { + t.Fatal("rejected reuse should still be recorded in service metrics") + } +} + +func TestDriverHandlesMissingExecutionFields(t *testing.T) { + for _, tc := range []struct { + name string + execution *athenatypes.QueryExecution + wantError string + wantStops int + }{ + {"execution", nil, "incomplete execution state", 1}, + {"status", &athenatypes.QueryExecution{}, "incomplete execution state", 1}, + {"output", &athenatypes.QueryExecution{Status: &athenatypes.QueryExecutionStatus{State: athenatypes.QueryExecutionStateSucceeded}}, "outside configured output location", 0}, + } { + t.Run(tc.name, func(t *testing.T) { + client := &fakeClient{executions: []*athenatypes.QueryExecution{tc.execution}} + driver := testDriver(t, client) + result, err := driver.Execute(context.Background(), perfcore.Query{PGWireSQL: "SELECT 1"}, nil) + if err == nil || !strings.Contains(err.Error(), tc.wantError) { + t.Fatalf("Execute error = %v, want %q", err, tc.wantError) + } + if result.ServiceMetrics != nil { + t.Fatal("missing statistics should not become zero-valued metrics") + } + if client.stopCalls != tc.wantStops { + t.Fatalf("stop calls = %d, want %d", client.stopCalls, tc.wantStops) + } + }) + } +} + +func testDriver(t *testing.T, client athenaAPI) *Driver { + t.Helper() + driver, err := newWithClient(client, ConnectionConfig{ + WorkGroup: "benchmark", Database: "benchmark_frozen", OutputLocation: "s3://benchmark-results/run/", - }, DriverOptions{}) + }) if err != nil { - t.Fatalf("NewWithClient returned error: %v", err) + t.Fatal(err) } + return driver +} - _, err = driver.Execute(context.Background(), perfcore.Query{PGWireSQL: "SELECT 1"}, nil) - if err == nil || !strings.Contains(err.Error(), "reused a previous result") { - t.Fatalf("Execute error = %v, want invalid result-reuse error", err) +func terminalExecution(state athenatypes.QueryExecutionState) *athenatypes.QueryExecution { + return &athenatypes.QueryExecution{ + Status: &athenatypes.QueryExecutionStatus{State: state, StateChangeReason: aws.String("scan limit exceeded")}, + ResultConfiguration: &athenatypes.ResultConfiguration{OutputLocation: aws.String("s3://benchmark-results/run/query.csv")}, + EngineVersion: &athenatypes.EngineVersion{EffectiveEngineVersion: aws.String("Athena engine version 3")}, + Statistics: &athenatypes.QueryExecutionStatistics{ + QueryQueueTimeInMillis: aws.Int64(120), QueryPlanningTimeInMillis: aws.Int64(80), + EngineExecutionTimeInMillis: aws.Int64(2000), TotalExecutionTimeInMillis: aws.Int64(2300), + DataScannedInBytes: aws.Int64(4096), DpuCount: aws.Float64(10), + }, + } +} + +func assertServiceMetrics(t *testing.T, got *perfcore.ServiceMetrics) { + t.Helper() + want := &perfcore.ServiceMetrics{ + QueueDuration: 120 * time.Millisecond, PlanningDuration: 80 * time.Millisecond, + EngineDuration: 2 * time.Second, ServiceDuration: 2300 * time.Millisecond, + BytesScanned: 4096, DPUCount: 10, EngineVersion: "Athena engine version 3", + } + if !reflect.DeepEqual(got, want) { + t.Fatalf("service metrics = %+v, want %+v", got, want) } } type fakeClient struct { - startInput StartQueryInput - executions []QueryExecution - resultPages []ResultPage - resultTokens []string - stopCalls int - stoppedID string + startInput *awsathena.StartQueryExecutionInput + executions []*athenatypes.QueryExecution + resultPages []*awsathena.GetQueryResultsOutput + resultTokens []string + stopCalls int + stoppedID string + stopContextErr error } -func (f *fakeClient) StartQuery(_ context.Context, input StartQueryInput) (string, error) { +func (f *fakeClient) StartQueryExecution(_ context.Context, input *awsathena.StartQueryExecutionInput, _ ...func(*awsathena.Options)) (*awsathena.StartQueryExecutionOutput, error) { f.startInput = input - return "query-1", nil + return &awsathena.StartQueryExecutionOutput{QueryExecutionId: aws.String("query-1")}, nil } -func (f *fakeClient) GetQuery(_ context.Context, _ string) (QueryExecution, error) { +func (f *fakeClient) GetQueryExecution(_ context.Context, input *awsathena.GetQueryExecutionInput, _ ...func(*awsathena.Options)) (*awsathena.GetQueryExecutionOutput, error) { + if aws.ToString(input.QueryExecutionId) != "query-1" { + return nil, errors.New("unexpected query ID") + } if len(f.executions) == 0 { - return QueryExecution{}, errors.New("unexpected GetQuery") + return nil, errors.New("unexpected GetQueryExecution") } execution := f.executions[0] if len(f.executions) > 1 { f.executions = f.executions[1:] } - return execution, nil + return &awsathena.GetQueryExecutionOutput{QueryExecution: execution}, nil } -func (f *fakeClient) GetResults(_ context.Context, _ string, nextToken string) (ResultPage, error) { - f.resultTokens = append(f.resultTokens, nextToken) +func (f *fakeClient) GetQueryResults(_ context.Context, input *awsathena.GetQueryResultsInput, _ ...func(*awsathena.Options)) (*awsathena.GetQueryResultsOutput, error) { + if aws.ToString(input.QueryExecutionId) != "query-1" || aws.ToInt32(input.MaxResults) != 1000 { + return nil, errors.New("unexpected GetQueryResults input") + } + f.resultTokens = append(f.resultTokens, aws.ToString(input.NextToken)) if len(f.resultPages) == 0 { - return ResultPage{}, errors.New("unexpected GetResults") + return nil, errors.New("unexpected GetQueryResults") } page := f.resultPages[0] f.resultPages = f.resultPages[1:] return page, nil } -func (f *fakeClient) StopQuery(_ context.Context, queryID string) error { +func (f *fakeClient) StopQueryExecution(ctx context.Context, input *awsathena.StopQueryExecutionInput, _ ...func(*awsathena.Options)) (*awsathena.StopQueryExecutionOutput, error) { f.stopCalls++ - f.stoppedID = queryID - return nil + f.stoppedID = aws.ToString(input.QueryExecutionId) + f.stopContextErr = ctx.Err() + return &awsathena.StopQueryExecutionOutput{}, nil } From 60703c3daca00f57b2d9dda23b986e06d200b458 Mon Sep 17 00:00:00 2001 From: Bill Guowei Yang Date: Mon, 7 Sep 2026 12:36:22 -0400 Subject: [PATCH 3/4] feat(perf): load Athena configuration from Terraform-managed SSM --- .github/workflows/scenario-dev.yml | 8 +- docs/runbooks/scenario-dev.md | 21 ++++- scripts/scenario_athena_config.sh | 34 +++++++ tests/mw-dev/README.md | 17 +++- .../scenario/athena_config_script_test.go | 94 +++++++++++++++++++ tests/mw-dev/scenario/script_test.go | 18 +++- 6 files changed, 176 insertions(+), 16 deletions(-) create mode 100644 scripts/scenario_athena_config.sh create mode 100644 tests/mw-dev/scenario/athena_config_script_test.go diff --git a/.github/workflows/scenario-dev.yml b/.github/workflows/scenario-dev.yml index b9f1c96d6..6db5b211c 100644 --- a/.github/workflows/scenario-dev.yml +++ b/.github/workflows/scenario-dev.yml @@ -74,10 +74,6 @@ jobs: # Frozen perf uses a namespace-local Trino cell. The dedicated role is # consumed only when E2E_SUITE selects that isolated deployment. TRINO_POD_IDENTITY_ROLE: ${{ secrets.MW_DEV_TRINO_POD_IDENTITY_ROLE }} - SCENARIO_POD_IDENTITY_ROLE: ${{ secrets.MW_DEV_ATHENA_POD_IDENTITY_ROLE }} - DUCKGRES_SCENARIO_ATHENA_WORKGROUP: ${{ vars.MW_DEV_ATHENA_WORKGROUP }} - DUCKGRES_SCENARIO_ATHENA_DATABASE: ${{ vars.MW_DEV_ATHENA_DATABASE }} - DUCKGRES_SCENARIO_ATHENA_RESULTS_S3_URI: ${{ vars.MW_DEV_ATHENA_RESULTS_S3_URI }} TRINO_IMAGE: ghcr.io/posthog/trino:4505364c570d6b51edecd299b603fca4b6693d86@sha256:ac80c275fd18a439d25da5652ab5cd3c80bcbdd2d88d64c9722dc3e8bb68ba07 E2E_SUITE: ${{ (github.event_name == 'schedule' || inputs.scenario == 'posthog_frozen_perf') && 'trino' || 'neutral' }} PR_NUMBER: ${{ github.run_id }} @@ -123,6 +119,10 @@ jobs: - name: Update kubeconfig run: aws eks update-kubeconfig --name "$CLUSTER_NAME" --region "$AWS_REGION" --alias "$KUBE_CONTEXT" + - name: Load Athena perf configuration + if: env.SCENARIO_NAME == 'posthog_frozen_perf' + run: bash scripts/scenario_athena_config.sh >> "$GITHUB_ENV" + - name: Deploy isolated Duckgres stack run: tests/mw-dev/run.sh deploy diff --git a/docs/runbooks/scenario-dev.md b/docs/runbooks/scenario-dev.md index a15b28d9f..fe61dcbd9 100644 --- a/docs/runbooks/scenario-dev.md +++ b/docs/runbooks/scenario-dev.md @@ -42,16 +42,22 @@ Configure these repository secrets: - `AWS_ECR_PUBLISH_IAM_ROLE` - `MW_DEV_ACCOUNT_ID` - `MW_DEV_TRINO_POD_IDENTITY_ROLE` -- `MW_DEV_ATHENA_POD_IDENTITY_ROLE` Configure these repository variables: - `TS_WIF_CLIENT_ID_MW_DEV` - `TS_WIF_AUDIENCE_MW_DEV` - `MW_DEV_SCENARIO_PERF_SECRET_ID` -- `MW_DEV_ATHENA_WORKGROUP` -- `MW_DEV_ATHENA_DATABASE` -- `MW_DEV_ATHENA_RESULTS_S3_URI` + +Athena requires no additional GitHub settings. Terraform publishes its +configuration to the SSM String parameter `/duckgres/perf/athena` in the dev +account and grants the workflow's existing OIDC role `ssm:GetParameter` on +that parameter only. Apply this infrastructure before enabling Athena runs. +For frozen perf, `scripts/scenario_athena_config.sh` loads and validates the +JSON after AWS authentication and before deploying the isolated stack. Other +scenarios do not load it. The parameter contains `pod_identity_role_arn`, +`workgroup_name`, `glue_database_name`, and `results_s3_uri`, derived from +Terraform resources; they are configuration identifiers, not credentials. The Athena values must identify an engine-v3 on-demand workgroup, a Glue database with explicit `events` and `persons` external tables over the frozen @@ -90,6 +96,13 @@ required as reported by `scripts/scenario_run.sh --check-env`. ## Failure Recovery +If loading `/duckgres/perf/athena` fails, confirm the Athena Terraform apply +completed, the workflow is using the expected AWS account/region, and its +OIDC role has the exact-parameter read grant. Invalid or incomplete JSON fails +before deployment and exports no partial settings. Fix the Terraform source +and apply it again rather than editing the parameter manually. Rerun the +workflow after the parameter and permission are available. + Every scenario sub-run provisions its own org or warehouse name and includes cleanup steps. The isolated stack also tears down the temporary namespace after the payload finishes. The harness keeps the artifact volume attached to a diff --git a/scripts/scenario_athena_config.sh b/scripts/scenario_athena_config.sh new file mode 100644 index 000000000..38f3e6ea2 --- /dev/null +++ b/scripts/scenario_athena_config.sh @@ -0,0 +1,34 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Emit validated GitHub environment-file assignments, never shell code to eval. +# Terraform owns this parameter and its exact-resource read permission. +parameter_name="/duckgres/perf/athena" +: "${AWS_REGION:?AWS_REGION is required to load Athena perf configuration}" + +if ! config=$(aws ssm get-parameter --name "$parameter_name" --region "$AWS_REGION" \ + --query Parameter.Value --output text); then + echo "Could not load Athena perf configuration from $parameter_name. Apply the Athena infrastructure and its CI read permission first." >&2 + exit 1 +fi + +# Validate the complete document before emitting anything: a bad value must not +# append a partial configuration or inject another entry into GITHUB_ENV. +if ! assignments=$(jq -ers ' + if length != 1 then error("expected one configuration object") else .[0] end + | { + SCENARIO_POD_IDENTITY_ROLE: .pod_identity_role_arn, + DUCKGRES_SCENARIO_ATHENA_WORKGROUP: .workgroup_name, + DUCKGRES_SCENARIO_ATHENA_DATABASE: .glue_database_name, + DUCKGRES_SCENARIO_ATHENA_RESULTS_S3_URI: .results_s3_uri + } + | if all(.[]; type == "string" and test("\\S") and (test("[[:cntrl:]]") | not)) + then to_entries[] | "\(.key)=\(.value)" + else error("expected nonempty single-line strings for all four Athena settings") + end +' <<< "$config"); then + echo "Invalid Athena perf configuration in $parameter_name. Check the Terraform-managed parameter; no settings were exported." >&2 + exit 1 +fi + +printf '%s\n' "$assignments" diff --git a/tests/mw-dev/README.md b/tests/mw-dev/README.md index 54e133be5..d5c770e0d 100644 --- a/tests/mw-dev/README.md +++ b/tests/mw-dev/README.md @@ -516,14 +516,23 @@ finalizers are still running. | var | `TS_WIF_AUDIENCE_MW_DEV` | Tailscale WIF audience | | secret | `MW_DEV_ACCOUNT_ID` | mw-dev AWS account id (kept out of committed code; ARNs are built from it) | | secret | `MW_DEV_TRINO_POD_IDENTITY_ROLE` | full ARN of the dedicated mw-dev Trino Pod Identity role (consumed only by the Trino lane) | -| secret | `MW_DEV_ATHENA_POD_IDENTITY_ROLE` | full ARN of the least-privilege Athena scenario Pod Identity role | -| var | `MW_DEV_ATHENA_WORKGROUP` | on-demand Athena engine-v3 workgroup | -| var | `MW_DEV_ATHENA_DATABASE` | Glue database containing the frozen external tables | -| var | `MW_DEV_ATHENA_RESULTS_S3_URI` | encrypted result prefix ending in `/`; lifecycle expiry is required | | secret | `AWS_ECR_PUBLISH_IAM_ROLE` | ECR push (already exists; used by CD) | | (role) | `github-duckgres-e2e` | dedicated stripped role in the mw-dev account (posthog-cloud-infra) — `eks:DescribeCluster` + Pod Identity association calls + `iam:PassRole`/`iam:GetRole` on the CP and dedicated Trino roles + an EKS access entry for kubectl. The workflow assumes `arn:aws:iam:::role/github-duckgres-e2e`. | | repo setting | "Require approval for all outside collaborators" | the access gate (see below) | +Athena adds no one-time GitHub configuration. Its Terraform unit publishes the +SSM String parameter `/duckgres/perf/athena` and grants the existing workflow +OIDC role `ssm:GetParameter` on that exact parameter. The JSON keys are +`pod_identity_role_arn`, `workgroup_name`, `glue_database_name`, and +`results_s3_uri`, all derived from the deployed resources. The frozen-perf +workflow loads them with `bash scripts/scenario_athena_config.sh` and exports +the existing scenario environment variables before deployment. The loader +requires AWS CLI, jq, and an explicit `AWS_REGION`; it fails without exporting +partial settings if fetching or validation fails. Apply the infrastructure +first, and fix configuration in Terraform rather than editing SSM manually. +Direct local scenario invocations still accept the documented explicit Athena +environment variables. + The `scenario-dev` workflow requests a 16,200-second session from `github-duckgres-e2e`, matching its 270-minute job timeout. The role's `max_session_duration` in posthog-cloud-infra must be at least 16,200 seconds diff --git a/tests/mw-dev/scenario/athena_config_script_test.go b/tests/mw-dev/scenario/athena_config_script_test.go new file mode 100644 index 000000000..06932abe9 --- /dev/null +++ b/tests/mw-dev/scenario/athena_config_script_test.go @@ -0,0 +1,94 @@ +package scenario + +import ( + "bytes" + "encoding/json" + "os" + "os/exec" + "path/filepath" + "strings" + "testing" +) + +func TestAthenaConfigScriptLoadsAndValidatesSSMConfiguration(t *testing.T) { + for _, tc := range []struct { + name string + change func(map[string]any) + raw string + awsError bool + wantError string + }{ + {name: "valid"}, + {name: "missing role", change: func(c map[string]any) { delete(c, "pod_identity_role_arn") }, wantError: "Invalid Athena perf configuration"}, + {name: "empty workgroup", change: func(c map[string]any) { c["workgroup_name"] = "" }, wantError: "Invalid Athena perf configuration"}, + {name: "wrong database type", change: func(c map[string]any) { c["glue_database_name"] = 42 }, wantError: "Invalid Athena perf configuration"}, + {name: "newline injection", change: func(c map[string]any) { c["results_s3_uri"] = "s3://example/results/\nUNEXPECTED=value" }, wantError: "Invalid Athena perf configuration"}, + {name: "carriage return", change: func(c map[string]any) { c["workgroup_name"] = "benchmark\rOTHER=value" }, wantError: "Invalid Athena perf configuration"}, + {name: "invalid JSON", raw: "not-json", wantError: "Invalid Athena perf configuration"}, + {name: "multiple documents", raw: "{}\n{}", wantError: "Invalid Athena perf configuration"}, + {name: "AWS failure", awsError: true, wantError: "Could not load Athena perf configuration"}, + } { + t.Run(tc.name, func(t *testing.T) { + config := map[string]any{ + "pod_identity_role_arn": "arn:aws:iam::123456789012:role/example-benchmark", + "workgroup_name": "benchmark", + "glue_database_name": "benchmark_frozen", + "results_s3_uri": "s3://example/results/", + } + if tc.change != nil { + tc.change(config) + } + raw, err := json.Marshal(config) + if err != nil { + t.Fatal(err) + } + if tc.raw != "" { + raw = []byte(tc.raw) + } + binDir := t.TempDir() + argsPath := filepath.Join(binDir, "args") + fakeAWS := "#!/bin/sh\nprintf '%s\\n' \"$*\" > \"$ATHENA_TEST_ARGS\"\nif [ \"$ATHENA_TEST_AWS_ERROR\" = 1 ]; then exit 1; fi\nprintf '%s\\n' \"$ATHENA_TEST_CONFIG\"\n" + if err := os.WriteFile(filepath.Join(binDir, "aws"), []byte(fakeAWS), 0o700); err != nil { + t.Fatal(err) + } + cmd := exec.Command("bash", filepath.Join("..", "..", "..", "scripts", "scenario_athena_config.sh")) + cmd.Env = []string{ + "PATH=" + binDir + string(os.PathListSeparator) + os.Getenv("PATH"), + "AWS_REGION=us-east-1", "ATHENA_TEST_ARGS=" + argsPath, + "ATHENA_TEST_CONFIG=" + string(raw), + } + if tc.awsError { + cmd.Env = append(cmd.Env, "ATHENA_TEST_AWS_ERROR=1") + } + var stdout, stderr bytes.Buffer + cmd.Stdout, cmd.Stderr = &stdout, &stderr + err = cmd.Run() + if tc.wantError != "" { + if err == nil || !strings.Contains(stderr.String(), tc.wantError) { + t.Fatalf("err=%v stderr=%q, want %q", err, stderr.String(), tc.wantError) + } + if stdout.Len() != 0 { + t.Fatalf("failed validation emitted partial environment: %q", stdout.String()) + } + return + } + if err != nil { + t.Fatalf("script failed: %v: %s", err, stderr.String()) + } + want := "SCENARIO_POD_IDENTITY_ROLE=arn:aws:iam::123456789012:role/example-benchmark\n" + + "DUCKGRES_SCENARIO_ATHENA_WORKGROUP=benchmark\n" + + "DUCKGRES_SCENARIO_ATHENA_DATABASE=benchmark_frozen\n" + + "DUCKGRES_SCENARIO_ATHENA_RESULTS_S3_URI=s3://example/results/\n" + if stdout.String() != want { + t.Fatalf("output=%q, want %q", stdout.String(), want) + } + args, err := os.ReadFile(argsPath) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(string(args), "ssm get-parameter --name /duckgres/perf/athena --region us-east-1 --query Parameter.Value --output text") { + t.Fatalf("unexpected AWS request: %s", args) + } + }) + } +} diff --git a/tests/mw-dev/scenario/script_test.go b/tests/mw-dev/scenario/script_test.go index 422cf0efc..f018125d9 100644 --- a/tests/mw-dev/scenario/script_test.go +++ b/tests/mw-dev/scenario/script_test.go @@ -94,10 +94,9 @@ func TestDevScenarioWorkflowUsesUnifiedMwDevHarness(t *testing.T) { "EKS_CLUSTER_NAME: posthog-mw-dev", "CP_POD_IDENTITY_ROLE: arn:aws:iam::${{ secrets.MW_DEV_ACCOUNT_ID }}:role/duckgres-control-plane-dev", "TRINO_POD_IDENTITY_ROLE: ${{ secrets.MW_DEV_TRINO_POD_IDENTITY_ROLE }}", - "SCENARIO_POD_IDENTITY_ROLE: ${{ secrets.MW_DEV_ATHENA_POD_IDENTITY_ROLE }}", - "DUCKGRES_SCENARIO_ATHENA_WORKGROUP: ${{ vars.MW_DEV_ATHENA_WORKGROUP }}", - "DUCKGRES_SCENARIO_ATHENA_DATABASE: ${{ vars.MW_DEV_ATHENA_DATABASE }}", - "DUCKGRES_SCENARIO_ATHENA_RESULTS_S3_URI: ${{ vars.MW_DEV_ATHENA_RESULTS_S3_URI }}", + "- name: Load Athena perf configuration", + "if: env.SCENARIO_NAME == 'posthog_frozen_perf'", + "bash scripts/scenario_athena_config.sh >> \"$GITHUB_ENV\"", "TRINO_IMAGE: ghcr.io/posthog/trino:", "E2E_SUITE: ${{ (github.event_name == 'schedule' || inputs.scenario == 'posthog_frozen_perf') && 'trino' || 'neutral' }}", "DUCKGRES_K8S_WORKER_CPU_REQUEST: \"3\"", @@ -126,6 +125,10 @@ func TestDevScenarioWorkflowUsesUnifiedMwDevHarness(t *testing.T) { } for _, forbidden := range []string{ + "MW_DEV_ATHENA_POD_IDENTITY_ROLE", + "MW_DEV_ATHENA_WORKGROUP", + "MW_DEV_ATHENA_DATABASE", + "MW_DEV_ATHENA_RESULTS_S3_URI", "skip_slow:", "inputs.skip_slow", "scenario-skipped:", @@ -163,6 +166,13 @@ func TestDevScenarioWorkflowUsesUnifiedMwDevHarness(t *testing.T) { } } + authIndex := strings.Index(workflow, "- name: Configure AWS credentials (OIDC)") + configIndex := strings.Index(workflow, "- name: Load Athena perf configuration") + deployIndex := strings.Index(workflow, "- name: Deploy isolated Duckgres stack") + if authIndex < 0 || configIndex < authIndex || deployIndex < configIndex { + t.Fatal("Athena config must load after AWS authentication and before deploying the stack") + } + teardownIndex := strings.Index(workflow, "- name: Teardown") publishPerfIndex := strings.Index(workflow, "- name: Publish scenario perf results") uploadIndex := strings.Index(workflow, "- name: Upload scenario artifacts") From 739e60106de692d8988d3b49046b38a3f50a58e0 Mon Sep 17 00:00:00 2001 From: Bill Guowei Yang Date: Mon, 7 Sep 2026 13:20:21 -0400 Subject: [PATCH 4/4] fix(tests): isolate scenario workflow environment --- tests/mw-dev/run_sh_test.go | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/tests/mw-dev/run_sh_test.go b/tests/mw-dev/run_sh_test.go index a245eda78..65297d97d 100644 --- a/tests/mw-dev/run_sh_test.go +++ b/tests/mw-dev/run_sh_test.go @@ -102,6 +102,9 @@ func TestDiagnosticsHarnessPodJSONPathParses(t *testing.T) { } func TestTrinoDeployStartsWorkloadsWithoutScaleSubresource(t *testing.T) { + // Scenario workflows export these before running the script tests. + t.Setenv("SCENARIO_NAME", "posthog_frozen_perf") + t.Setenv("SCENARIO_POD_IDENTITY_ROLE", "") fakes := newRunSHFakes(t) secretDir := filepath.Join(filepath.Dir(fakes.binDir), "secrets") for _, name := range []string{"duckgres-ci-trino-ca.crt", "duckgres-ci-trino-server.p12"} { @@ -702,10 +705,10 @@ func TestScenarioArtifactTokenCollisionCreatesANewResultDirectory(t *testing.T) } func TestScenarioDefaultsToFullSuite(t *testing.T) { - t.Setenv("SCENARIO_NAME", "") fakes := newRunSHFakes(t) cmd := runSHCommand(t, fakes.binDir, "test-scenario", + "SCENARIO_NAME=", "SCENARIO_RUNNER_IMAGE=example.invalid/duckgres:scenario", ) out, err := cmd.CombinedOutput() @@ -1122,6 +1125,9 @@ func TestE2EHarnessCoversRemoteBinaryCopy(t *testing.T) { } func TestDeployCreatesConfiguredSecretDirectoryPrivately(t *testing.T) { + t.Setenv("SCENARIO_NAME", "posthog_frozen_perf") + t.Setenv("SCENARIO_POD_IDENTITY_ROLE", "") + t.Setenv("E2E_SUITE", "trino") fakes := newRunSHFakes(t) secretDir := filepath.Join(filepath.Dir(fakes.binDir), "generated", "secrets") @@ -2039,6 +2045,9 @@ func runSHCommand(t *testing.T, binDir, subcommand string, extraEnv ...string) * "CP_POD_IDENTITY_ROLE=arn:aws:iam::123456789012:role/duckgres-control-plane-dev", "EKS_CLUSTER_NAME=test-cluster", "AWS_REGION=us-east-1", + "E2E_SUITE=neutral", + "SCENARIO_NAME=full-suite", + "SCENARIO_POD_IDENTITY_ROLE=", "SCENARIO_ARTIFACTS_DIR="+filepath.Join(filepath.Dir(binDir), "scenario-artifacts"), "DUCKGRES_CI_SECRET_DIR="+filepath.Join(filepath.Dir(binDir), "secrets"), )