Skip to content

Commit d3a93d7

Browse files
authored
feat(databases-on-aws): add query-plan-explainability workflow to dsql skill (#141)
* feat(databases-on-aws): add query-plan-explainability workflow to dsql skill Adds Workflow 8 to the dsql skill: a phased diagnostic pipeline that captures an EXPLAIN ANALYZE VERBOSE plan, interrogates pg_class / pg_stats / pg_indexes for cardinality and index evidence, optionally runs GUC experiments (gated on execution time ≤30s), and produces a structured Markdown diagnostic report. When the aurora-dsql MCP server is unavailable, the workflow falls back to raw psql with a generated IAM token — multi-statement heredoc, exit-code checked — supporting GUC experiments and BEGIN/ROLLBACK wrappers. Four supporting references live under `references/query-plan/` (plan-interpretation, catalog-queries, guc-experiments, report-format) so the entry file stays under the 300-line guidance. DML safety. Plan capture is exclusively `readonly_query`; write-mode `transact` bypasses all MCP safety checks and is never used in Phase 1. INSERT, pl/pgsql, DO blocks, and functions are rejected outright (no meaningful plan). UPDATE and DELETE are rewritten to the equivalent SELECT (same join chain + WHERE) before EXPLAIN — the optimizer picks the same plan shape for the SELECT. The Experiment 2 GUC path still needs `transact` (readonly_query blocks multi-statement), but caller-side guard rails require the statement be verified-SELECT, passed as a single list element (no string concat that could smuggle a second statement), and halt-on-error without recovery chaining. SET LOCAL scopes GUC changes to the transaction. Adds a functional eval harness under tools/evals/databases-on-aws/ with 5 prompts / 42 assertions, a Python runner driving claude -p via stream-json, and an MCP warmup that retries once and aborts on repeated failure so evals never run against a cold MCP. Tool-required assertions accept both MCP tool calls and the psql fallback (Bash invoking psql against a DSQL endpoint). Other assertions match against agent output text only to avoid false positives from prompt echo. Per-eval diagnostic.json persists reference loads, MCP calls, and malformed-JSON counts for post-hoc triage. Eval 4's DML-recognition grader accepts DML → SELECT rewrite, INSERT rejection, or the legacy ROLLBACK-wrapping form. Also extends trigger_evals.json with three should-trigger explainability queries and three should-not-trigger PostgreSQL/RDS/Redshift EXPLAIN variants to confirm the skill stays DSQL-specific, and scopes the Bandit security scan to skip ./.tmp and ./node_modules so per-developer scratch code under .tmp/ doesn't break the build. Co-authored-by: anwesham-lab <64298192+anwesham-lab@users.noreply.github.com> * feat(databases-on-aws): harden Phase 1, add Phase 5 reassessment loop Three related skill hardenings for Workflow 8 query-plan-explainability: 1. Phase 1 hardening. ALWAYS run readonly_query(EXPLAIN ANALYZE VERBOSE …) on the user's query verbatim and capture a fresh plan from the cluster — even when the user describes the plan or reports an anomaly, because findings and the support template need real cluster data. Schema pre-checks via get_schema / information_schema are permitted (MAY). When EXPLAIN errors (relation does not exist, column does not exist), MUST report the error verbatim; MUST NOT invent DSQL-specific semantics (e.g., case sensitivity, identifier quoting) as the root cause. PostgreSQL/DSQL auto-lowercase unquoted identifiers so "case sensitivity" is not a valid DSQL-specific diagnosis in the first place. This preempts the class of failure where the agent substitutes a plausible-sounding theory for evidence when a diagnostic tool returns ambiguous empty results. 2. Workflow-opening reframe: "The structured Markdown diagnostic report is the deliverable, not a conversational answer — run the workflow end-to-end before answering." Fixes the pattern where the agent narrates causes in prose instead of running the workflow. 3. Phase 5 — Reassess after change. Workflow now ends with a Next Steps block inviting the user to signal "reassess" after applying any recommendation. When they do, the agent re-runs Phase 1–2 and appends an "Addendum: After-Change Performance" section (before/after comparison table with match-vs-expected-impact verdict) to the original report rather than producing a new one. Template for both lives in report-format.md so the structure is enforced. All normative statements use RFC 2119 keywords (MUST, MUST NOT, MAY) plus the repo's ALWAYS convention. Co-authored-by: anwesham-lab <64298192+anwesham-lab@users.noreply.github.com> * feat(databases-on-aws): add outcome-oriented evals 6-9 + grader enhancements Four new functional evals target outcome quality rather than workflow execution, covering failure modes the existing 1-5 miss: - Eval 6 (Phase 5 reassessment): agent re-runs EXPLAIN after the user applies a recommendation, appends an Addendum to the prior report rather than producing a fresh report, and compares observed improvement against the originally stated Expected Impact. - Eval 7 (mixed-case identifiers): SELECT FROM UserAccount resolves via PG auto-lowercasing. Agent MUST NOT invent "DSQL is case-sensitive" as a root cause. Guards against the real-user hallucination where get_schema returning empty got fabricated into a DSQL-specific rule. - Eval 8 (unknown table): agent surfaces relation does not exist verbatim, does NOT fabricate a diagnostic report for a query it could not EXPLAIN. - Eval 9 (stale pg_class.reltuples): real DSQL pattern where plan looks fine but stats lag reality. Agent MUST query both pg_class and COUNT(*), detect divergence, and recommend ANALYZE / note the auto-analyze schedule. Three grader enhancements support the new assertions (ordered before the fallback keyword matcher): - Grader (b) — Expected Impact contains concrete numbers. Extracts the Summary section, requires at least one x-factor, percentage, or ms/s value. Rejects hedging prose ("should improve", "significant reduction"). - Grader (c) — Addendum appended, not fresh report. Counts top-level H1s and requires Addendum header present. Fails when agent restarts with a new Diagnostic Report title on reassessment. - Grader (d) — No hallucinated DSQL-specific semantics. Regex blacklist for "DSQL is case-sensitive", "you must lowercase", etc. Passes if paired with an auto-lowercase correction. README updated: 9 evals / 70 assertions, with per-eval summary table. Co-authored-by: anwesham-lab <64298192+anwesham-lab@users.noreply.github.com> * feat(databases-on-aws): grader enhancements + anti-hallucination framing on Expected Impact Grader improvements for evals 6, 8, 9 — dedicated branches that were previously hitting the 0.8-keyword fallback matcher: - Before/after comparison table with numeric duration delta (eval 6) - Comments explicitly on match vs Expected Impact (eval 6) - Proposes next hypothesis on shortfall rather than declaring success (eval 6) - Surfaces 'relation ... does not exist' verbatim (eval 8) - Does NOT produce a fresh Diagnostic Report H1 for an un-EXPLAINable query (eval 8) - Does NOT fabricate plan metrics / findings / row counts (eval 8) - Asks user to confirm table / schema / cluster (eval 8) - Names stale statistics as root cause (eval 9) - Recommends ANALYZE or notes auto-analyze schedule (eval 9) Grader (b) reframed from "must include numbers" to "must be evidence-grounded". The old framing rewarded any number, which is directly hallucination-inducing (same pathology as the case-sensitivity fabrication): forcing a concrete prediction without evidence invites the agent to invent one. New framing accepts either a numeric prediction OR an honest "cannot predict without X evidence" admission. Plain hedging ("should improve performance") still fails. Matching skill update in report-format.md: Expected impact is now instructed to ground predictions in gathered evidence; when evidence is insufficient, the agent MUST name the missing piece rather than fabricate a number. Co-authored-by: anwesham-lab <64298192+anwesham-lab@users.noreply.github.com>
1 parent bd83f9b commit d3a93d7

11 files changed

Lines changed: 2175 additions & 14 deletions

File tree

mise.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -106,7 +106,7 @@ run = "uv run tools/init-skill.py"
106106
[tasks."security:bandit"]
107107
description = "Run Bandit"
108108
run = [
109-
"bandit -r ."
109+
"bandit -r . -x ./.tmp,./node_modules",
110110
]
111111

112112
[tasks."security:semgrep"]

plugins/databases-on-aws/skills/dsql/SKILL.md

Lines changed: 36 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
---
22
name: dsql
3-
description: "Build with Aurora DSQL — manage schemas, execute queries, handle migrations, and develop applications with a serverless, distributed SQL database. Covers IAM auth, multi-tenant patterns, MySQL-to-DSQL migration, and DDL operations. Triggers on phrases like: DSQL, Aurora DSQL, create DSQL table, DSQL schema, migrate to DSQL, distributed SQL database, serverless PostgreSQL-compatible database."
3+
description: "Build with Aurora DSQL — manage schemas, execute queries, handle migrations, diagnose query plans, and develop applications with a serverless, distributed SQL database. Covers IAM auth, multi-tenant patterns, MySQL-to-DSQL migration, DDL operations, and query plan explainability. Triggers on phrases like: DSQL, Aurora DSQL, create DSQL table, DSQL schema, migrate to DSQL, distributed SQL database, serverless PostgreSQL-compatible database, DSQL query plan, DSQL EXPLAIN ANALYZE, why is my DSQL query slow."
44
license: Apache-2.0
55
metadata:
66
tags: aws, aurora, dsql, distributed-sql, distributed, distributed-database, database, serverless, serverless-database, postgresql, postgres, sql, schema, migration, multi-tenant, iam-auth, aurora-dsql, mcp
@@ -109,6 +109,11 @@ sampled in [.mcp.json](../../.mcp.json)
109109
**When:** Load when migrating a complete MySQL table to DSQL
110110
**Contains:** End-to-end MySQL CREATE TABLE migration example with decision summary
111111

112+
### Query Plan Explainability (modular):
113+
114+
**When:** MUST load all four at Workflow 8 Phase 0 — [query-plan/plan-interpretation.md](references/query-plan/plan-interpretation.md), [query-plan/catalog-queries.md](references/query-plan/catalog-queries.md), [query-plan/guc-experiments.md](references/query-plan/guc-experiments.md), [query-plan/report-format.md](references/query-plan/report-format.md)
115+
**Contains:** DSQL node types + Node Duration math + estimation-error bands, pg_class/pg_stats/pg_indexes SQL + correlated-predicate verification, GUC experiment procedures + 30-second skip protocol, required report structure + element checklist + support request template
116+
112117
---
113118

114119
## MCP Tools Available
@@ -156,8 +161,6 @@ the exact number doesn't affect the user's decision.
156161
**Fallback:** If `awsknowledge` is unavailable, use the defaults above and note to the user
157162
that limits should be verified against [DSQL documentation](https://docs.aws.amazon.com/aurora-dsql/latest/userguide/).
158163

159-
---
160-
161164
## CLI Scripts Available
162165

163166
Bash scripts in [scripts/](../../scripts/) for cluster management (create, delete, list, cluster info), psql connection, and bulk data loading from local/s3 csv/tsv/parquet files.
@@ -188,8 +191,7 @@ Validate inputs carefully (no parameterized queries available)
188191
Use transact tool with list of SQL statements
189192
Follow one-DDL-per-transaction rule
190193
Always use CREATE INDEX ASYNC in separate transaction
191-
ALTER COLUMN TYPE, DROP COLUMN, DROP CONSTRAINT are NOT supported directly
192-
→ These require the Table Recreation Pattern (see Workflow 6)
194+
ALTER COLUMN TYPE, DROP COLUMN, DROP CONSTRAINT → Table Recreation Pattern (Workflow 6)
193195
```
194196

195197
---
@@ -250,16 +252,42 @@ MUST load [ddl-migrations/overview.md](references/ddl-migrations/overview.md) be
250252

251253
MUST load [mysql-migrations/type-mapping.md](references/mysql-migrations/type-mapping.md) for type mappings, feature alternatives, and migration steps.
252254

255+
### Workflow 8: Query Plan Explainability
256+
257+
Explains why the DSQL optimizer chose a particular plan. Triggered by slow queries, high DPU, unexpected Full Scans, or plans the user doesn't understand. **REQUIRES a structured Markdown diagnostic report is the deliverable** beyond conversation — run the workflow end-to-end before answering. Use the `aurora-dsql` MCP when connected; fall back to raw `psql` with a generated IAM token (see the fallback block below) otherwise.
258+
259+
**Phase 0 — Load reference material.** Read all four before starting — each has content later phases need verbatim (node-type math, exact catalog SQL, the `>30s` skip protocol, required report elements):
260+
261+
1. [query-plan/plan-interpretation.md](references/query-plan/plan-interpretation.md) — node types, duration math, anomalous values
262+
2. [query-plan/catalog-queries.md](references/query-plan/catalog-queries.md) — pg_class / pg_stats / pg_indexes SQL
263+
3. [query-plan/guc-experiments.md](references/query-plan/guc-experiments.md) — GUC procedures and `>30s` skip protocol
264+
4. [query-plan/report-format.md](references/query-plan/report-format.md) — required report structure
265+
266+
**Phase 1 — Capture the plan.** **ALWAYS** run `readonly_query("EXPLAIN ANALYZE VERBOSE …")` on the user's query verbatim (SELECT form) — **ALWAYS** capture a fresh plan from the cluster, even when the user describes the plan or reports an anomaly. **MAY** leverage `get_schema` or `information_schema` for schema sanity checks. When EXPLAIN errors (`relation does not exist`, `column does not exist`), **MUST** report the error verbatim — **MUST NOT** invent DSQL-specific semantics (e.g., case sensitivity, identifier quoting) as the root cause. Extract Query ID, Planning Time, Execution Time, DPU Estimate. **SELECT** runs as-is. **UPDATE/DELETE** rewrite to the equivalent SELECT (same join chain + WHERE) — the optimizer picks the same plan shape. **INSERT**, pl/pgsql, DO blocks, and functions **MUST** be rejected. **MUST NOT** use `transact --allow-writes` for plan capture; it bypasses MCP safety.
267+
268+
**Phase 2 — Gather evidence.** Using SQL from `catalog-queries.md`, query `pg_class`, `pg_stats`, `pg_indexes`, `COUNT(*)`, `COUNT(DISTINCT)`. Classify estimation errors per `plan-interpretation.md` (2x–5x minor, 5x–50x significant, 50x+ severe). Detect correlated predicates and data skew.
269+
270+
**Phase 3 — Experiment (conditional).** ≤30s: run GUC experiments per `guc-experiments.md` (default + merge-join-only) plus optional redundant-predicate test. >30s: skip experiments, include the manual GUC testing SQL verbatim in the report, and do not re-run for redundant-predicate testing. Anomalous values (impossible row counts): confirm query results are correct despite the anomalous EXPLAIN, flag as a potential DSQL bug, and produce the Support Request Template from `report-format.md`.
271+
272+
**Phase 4 — Produce the report, invite reassessment.** Produce the full diagnostic report per the "Required Elements Checklist" in [query-plan/report-format.md](references/query-plan/report-format.md) — structure is non-negotiable. End with the "Next Steps" block from that reference so the user can ask for a reassessment after applying a recommendation. When the user says "reassess" (or equivalent), re-run Phase 1–2 and **append an "Addendum: After-Change Performance"** to the original report (before/after table, match against expected impact) rather than producing a new report.
273+
274+
**psql fallback (MCP unavailable).** Pipe statements into `psql` via heredoc and check `$?`; report failures without proceeding on partial evidence:
275+
276+
```bash
277+
TOKEN=$(aws dsql generate-db-connect-admin-auth-token --hostname "$HOST" --region "$REGION")
278+
PGPASSWORD="$TOKEN" psql "host=$HOST port=5432 user=admin dbname=postgres sslmode=require" <<<"EXPLAIN ANALYZE VERBOSE <sql>;"
279+
```
280+
281+
**Safety.** Plan capture uses `readonly_query` exclusively — it rejects INSERT/UPDATE/DELETE/DDL at the MCP layer. Rewrite DML to SELECT (Phase 1) rather than asking `transact --allow-writes` to run it; write-mode `transact` bypasses all MCP safety checks. **MUST NOT** run arbitrary DDL/DML or pl/pgsql.
282+
253283
---
254284

255285
## Error Scenarios
256286

257-
- **MCP server unavailable:** Fall back to CLI scripts ([scripts/](../../scripts/)) or direct psql. Note the limitation to the user.
258287
- **`awsknowledge` returns no results:** Use the default limits in the table above and note that limits should be verified against [DSQL documentation](https://docs.aws.amazon.com/aurora-dsql/latest/userguide/).
259288
- **OCC serialization error:** Retry the transaction. If persistent, check for hot-key contention — see [troubleshooting.md](references/troubleshooting.md).
260289
- **Transaction exceeds limits:** Split into batches under 3,000 rows — see [batched-migration.md](references/ddl-migrations/batched-migration.md).
261-
- **Token expiration mid-operation:** Generate a fresh IAM token and reconnect — see [authentication-guide.md](references/auth/authentication-guide.md).
262-
- See [troubleshooting.md](references/troubleshooting.md) for other issues.
290+
- **Token expiration mid-operation:** Generate a fresh IAM token — see [authentication-guide.md](references/auth/authentication-guide.md). See [troubleshooting.md](references/troubleshooting.md) for other issues.
263291

264292
---
265293

@@ -268,5 +296,4 @@ MUST load [mysql-migrations/type-mapping.md](references/mysql-migrations/type-ma
268296
- [Aurora DSQL Documentation](https://docs.aws.amazon.com/aurora-dsql/latest/userguide/)
269297
- [Code Samples Repository](https://github.com/aws-samples/aurora-dsql-samples)
270298
- [PostgreSQL Compatibility](https://docs.aws.amazon.com/aurora-dsql/latest/userguide/working-with-postgresql-compatibility.html)
271-
- [IAM Authentication Guide](https://docs.aws.amazon.com/aurora-dsql/latest/userguide/using-database-and-iam-roles.html)
272299
- [CloudFormation Resource](https://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/aws-resource-dsql-cluster.html)

plugins/databases-on-aws/skills/dsql/references/auth/authentication-guide.md

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -147,3 +147,7 @@ For production applications:
147147

148148
For complete role setup instructions, schema separation patterns, and IAM configuration,
149149
see [access-control.md](../access-control.md).
150+
151+
## Additional Resources
152+
153+
- [IAM Authentication Guide (AWS documentation)](https://docs.aws.amazon.com/aurora-dsql/latest/userguide/using-database-and-iam-roles.html)
Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
1+
# Catalog Queries Reference
2+
3+
Exact SQL for interrogating optimizer statistics and actual cardinalities against the DSQL cluster.
4+
5+
## Table of Contents
6+
7+
1. [Table-Level Statistics (pg_class)](#table-level-statistics)
8+
2. [Column Statistics (pg_stats)](#column-statistics)
9+
3. [Index Definitions](#index-definitions)
10+
4. [Actual Row Counts](#actual-row-counts)
11+
5. [Actual Distinct Counts](#actual-distinct-counts)
12+
6. [Value Distribution Analysis](#value-distribution-analysis)
13+
14+
---
15+
16+
## Table-Level Statistics
17+
18+
Retrieve optimizer's view of table size for all referenced tables:
19+
20+
```sql
21+
SELECT
22+
schemaname,
23+
relname,
24+
reltuples::bigint AS estimated_rows,
25+
relpages
26+
FROM pg_class c
27+
JOIN pg_namespace n ON n.oid = c.relnamespace
28+
WHERE n.nspname = '{schema}'
29+
AND c.relname IN ('{table1}', '{table2}', '{table3}');
30+
```
31+
32+
Compare `reltuples` against actual `COUNT(*)`. A divergence >20% on the table-stats snapshot indicates stale `reltuples` requiring `ANALYZE`. This is distinct from the row-estimate-vs-actual error thresholds used for plan findings (see plan-interpretation.md: 2x–5x minor, 5x–50x significant, 50x+ severe).
33+
34+
## Column Statistics
35+
36+
Retrieve statistics for columns involved in joins, WHERE clauses, and estimation errors:
37+
38+
```sql
39+
SELECT
40+
tablename,
41+
attname,
42+
null_frac,
43+
n_distinct,
44+
most_common_vals,
45+
most_common_freqs,
46+
histogram_bounds,
47+
correlation
48+
FROM pg_stats
49+
WHERE schemaname = '{schema}'
50+
AND tablename = '{table}'
51+
AND attname IN ('{col1}', '{col2}');
52+
```
53+
54+
**Key fields:**
55+
56+
| Field | Use |
57+
| ------------------- | ------------------------------------------------------ |
58+
| `n_distinct` | Negative = fraction of rows; Positive = absolute count |
59+
| `most_common_vals` | Values the optimizer considers frequent |
60+
| `most_common_freqs` | Corresponding frequencies (sum < 1.0) |
61+
| `histogram_bounds` | Equal-frequency bucket boundaries for non-MCV values |
62+
| `correlation` | Physical row order correlation (-1 to 1) |
63+
64+
## Index Definitions
65+
66+
Retrieve existing indexes on referenced tables. DSQL does not populate the cumulative `pg_stat_user_indexes` counters (`idx_scan`, `idx_tup_read`, `idx_tup_fetch`) that standard PostgreSQL exposes — infer index usage from the EXPLAIN plan instead.
67+
68+
```sql
69+
SELECT
70+
tablename,
71+
indexname,
72+
indexdef
73+
FROM pg_indexes
74+
WHERE schemaname = '{schema}'
75+
AND tablename IN ('{table1}', '{table2}', '{table3}')
76+
ORDER BY tablename, indexname;
77+
```
78+
79+
## Actual Row Counts
80+
81+
Retrieve ground-truth row counts for comparison against `pg_class.reltuples`:
82+
83+
```sql
84+
SELECT COUNT(*) AS actual_rows FROM {schema}.{table};
85+
```
86+
87+
Run for each referenced table. Present results as:
88+
89+
| Table | pg_class.reltuples | Actual COUNT(*) | Difference |
90+
| ------ | ------------------ | --------------- | ------------------ |
91+
| table1 | N | M | X% over/undercount |
92+
93+
## Actual Distinct Counts
94+
95+
Retrieve actual distinct values for columns in joins and WHERE predicates:
96+
97+
```sql
98+
SELECT COUNT(DISTINCT {column}) AS distinct_count FROM {schema}.{table};
99+
```
100+
101+
Compare against `pg_stats.n_distinct`:
102+
103+
- If `n_distinct` is positive: compare directly
104+
- If `n_distinct` is negative: multiply absolute value by actual row count to get estimated distinct count
105+
106+
## Value Distribution Analysis
107+
108+
For columns with suspected data skew, retrieve the actual top-N value frequencies:
109+
110+
```sql
111+
SELECT
112+
{column},
113+
COUNT(*) AS freq,
114+
ROUND(COUNT(*)::numeric / (SELECT COUNT(*) FROM {schema}.{table}), 5) AS fraction
115+
FROM {schema}.{table}
116+
GROUP BY {column}
117+
ORDER BY freq DESC
118+
LIMIT 20;
119+
```
120+
121+
Compare results against `most_common_vals` and `most_common_freqs` from pg_stats. Flag:
122+
123+
- Values present in data but missing from `most_common_vals`
124+
- Values whose actual frequency differs >2x from `most_common_freqs`
125+
- Skewed distributions where top values account for >50% of rows
126+
127+
### Correlated Predicate Verification
128+
129+
To verify predicate correlation, measure the actual combined selectivity:
130+
131+
```sql
132+
SELECT COUNT(*) AS combined_count
133+
FROM {schema}.{table}
134+
WHERE {predicate1} AND {predicate2};
135+
```
136+
137+
Then compare against the independence assumption:
138+
139+
```
140+
Expected (independent) = (count_pred1 / total_rows) × (count_pred2 / total_rows) × total_rows
141+
Actual = combined_count
142+
Error = actual / expected
143+
```
144+
145+
An error >3x indicates significant predicate correlation.

0 commit comments

Comments
 (0)