Skip to content

Commit 40b62d8

Browse files
feat: SQL query-plan diagrams, and a reports/ folder for generated artifacts (#50)
## What? Adds `make sql-diagram` / `/sql-diagram`: a `sqlglot`-based tool that draws a SQL query either as its execution steps (default) or as column-level lineage, writing the analysed `.sql`, a Mermaid `.mmd` and a standalone `.svg`. Also collects generated artifacts under a new `reports/` folder — `coverage_reports/` becomes `reports/coverage/`, `cost_report/` becomes `reports/cost/`, diagrams land in `reports/sql-diagram/` — and folds in the pending doc consolidation in `specs/tooling.md` / `specs/workflow.md`. ## Why? Reading a non-trivial query in this repo means reconstructing its shape by hand, and the queries that matter most (the `system.billing` cost queries) live interpolated inside f-strings, so the SQL that actually runs exists nowhere in the tree. Existing Claude skills for this either analyse an `EXPLAIN` you paste in or draw generic Mermaid from model inference; none derives the diagram from the query. Parsing the AST gives edges that are what the query says rather than what a model guessed. ## How? `sqlglot` models a multi-table join as one n-ary step, so `build_plan` splits it back into `JOIN 1`, `JOIN 2`, … in written order — that is what makes an under-constrained join predicate visible, which is exactly the class of bug that caused the 3.5× fan-out fixed in #47. Each join carries its side and full `ON` condition; `--comments` annotates tables with their Unity Catalog comment via the `dev` profile, since the MCP service principal lacks `USE SCHEMA` on `system.billing`. It is opt-in as the only network call, and a table without a comment is left unannotated rather than described from guesswork. SVG is written by hand rather than shelled out to `mmdc`, which would pull a Node toolchain and headless Chromium into a Python repo — the same trade `scripts/star_history.py` already makes. One CSS class per box, because a multi-class cascade is not resolved by every SVG renderer. ## Validation? `uv run pytest` passes with coverage writing to the new path; ruff clean. Both modes exercised on the two cost queries, and the SVG rasterised and inspected at each step — which caught black-on-black boxes twice and mid-sentence truncation. Round-tripping the emitted `.sql` exposed non-determinism: `planner.Step.dependencies` is a `set`, so node numbering followed the hash seed and three runs over one input produced three different files, making the `.mmd` churn in git. Dependencies are now walked in a stable order and the round trip reproduces byte-for-byte across five runs in both modes. ## Impact in prod - [x] No table schema/data change — no production impact. No medallion table, job, or pipeline is touched. The one operational change is the CI coverage artifact path in `onpush.yml`, updated to `reports/coverage/` in the same commit as the `pyproject.toml` setting that produces it.
1 parent bd57d30 commit 40b62d8

17 files changed

Lines changed: 1106 additions & 199 deletions

File tree

.claude/commands/project-costs.md

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,10 @@ analysis into the generated markdown report.
44
## Steps
55

66
1. Run `make project-costs` via Bash and capture the full output. It prints to stdout **and**
7-
writes `cost_report/YYYY-MM-DD.md` (gitignored) containing every table in markdown, ending with
7+
writes `reports/cost/YYYY-MM-DD.md` (gitignored) containing every table in markdown, ending with
88
an empty `## Analysis` section. The script writes only numbers, never prose.
99
2. Analyze the output (see below).
10-
3. **Edit `cost_report/YYYY-MM-DD.md` and replace the `## Analysis` placeholder with your written
10+
3. **Edit `reports/cost/YYYY-MM-DD.md` and replace the `## Analysis` placeholder with your written
1111
analysis.** Do not retype the tables — they are already in the file, and re-transcribing numbers
1212
risks introducing errors. Reference them instead.
1313
4. Report the same analysis back in chat, and link the report path.

.claude/commands/sql-diagram.md

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
Diagram a SQL query and explain what it shows — either its execution steps or its column lineage.
2+
3+
## Steps
4+
5+
1. Get the SQL. If the user named a file, use it. If the query is embedded in Python (most of this
6+
repo's SQL lives in f-strings under `scripts/`), extract it to a scratch `.sql` file first and
7+
**replace the interpolated placeholders with literals**`sqlglot` parses SQL, not f-strings.
8+
2. Pick the mode. `mode=plan` (the default) answers "what does this query *do*, step by step";
9+
`mode=lineage` answers "where does this output column come from". When the user asks about
10+
joins, stages, filters or ordering, they want `plan`.
11+
3. Run `make sql-diagram sql=<path> name=<basename> comments=1` via Bash. It writes three files to
12+
`reports/sql-diagram/`: `<basename>.sql` (the query as analysed), `.mmd` and `.svg` — all
13+
gitignored, so `git add -f` them only if they are meant to be a committed example. Pass
14+
`--stdout` to `scripts/sql_diagram.py` for a throwaway look with no files written.
15+
4. Read the `.mmd`, show it in a ```mermaid fence, and explain it (see below). The `.svg` is the
16+
same graph for linking from prose where no Mermaid renderer is available.
17+
18+
The emitted `.sql` is what makes the diagram auditable: it is the query *after* any f-string
19+
placeholders were filled in, so `make sql-diagram sql=reports/sql-diagram/<basename>.sql` reproduces
20+
the diagram exactly. When you commit a diagram as an example, commit its `.sql` with it.
21+
22+
Both diagrams come from the parsed AST, so they are exactly what the query says — do not "improve"
23+
one by adding a node or edge you believe should be there. If it looks wrong, the query is the thing
24+
to question.
25+
26+
## Reading `mode=plan`
27+
28+
Nodes are the query's steps, bottom-up: `SCAN` per table, one `JOIN n` per individual join, then
29+
`WHERE`, `AGGREGATE`, `SORT`, `OUTPUT`. A CTE appears as its own sub-pipeline feeding the `SCAN`
30+
that reads it.
31+
32+
- **Each join is numbered in the order the query writes it** and carries its side and keys.
33+
`sqlglot` models a multi-table join as one n-ary step; the script splits it back apart. An
34+
`INNER JOIN` silently drops rows where a `LEFT JOIN` keeps them — always say which, because it
35+
changes what a blank in the output means.
36+
- **Extra `ON` predicates beyond the equality keys** are listed under the keys as `and …`. On a
37+
slowly-changing dimension those range predicates are what stop the join fanning out; call them
38+
out rather than treating them as noise.
39+
- **This is the logical plan, not the physical one.** Databricks reorders joins, chooses broadcast
40+
versus shuffle, and prunes columns. Say "as written" — and if the real execution matters, point
41+
at the query profile in the UI or `EXPLAIN FORMATTED`, which is the only authority on what ran.
42+
- `AGGREGATE` may show synthetic operand names (`_a_0`) for `DISTINCT`/expression arguments that
43+
`sqlglot` lifted out. Read the intent off the original SQL rather than repeating the placeholder.
44+
45+
## Reading `mode=lineage`
46+
47+
- **Subgraphs are source tables**, one node per source column actually read. A column the query
48+
never touches does not appear — that is the point.
49+
- **The `output` subgraph** is the projected column list, in select order.
50+
- **`(unqualified)`** collects columns referenced without a table prefix in a multi-table join.
51+
`sqlglot` will not guess which side they came from without the table schemas, and neither should
52+
you. Call it out: it is usually a readability defect in the query worth fixing at the source.
53+
- **Struct columns collapse to their root.** `u.usage_metadata.job_id` traces back to
54+
`usage_metadata`, not to the leaf field. Say so rather than implying field-level precision.
55+
- Columns in `WHERE`/`GROUP BY` but not in the output do **not** appear. Use `mode=plan` when
56+
filtering is the point.
57+
58+
## In either mode
59+
60+
**The grey line under a table name** is its Unity Catalog comment, present only when the run passed
61+
`comments=1` and the profile could read the table. It is fetched, never written by you — if a table
62+
has no comment the space is blank, and that absence is itself worth reporting.
63+
64+
## What to say about it
65+
66+
- The **shape** of the query first: how many tables, how many joins, what it groups by. A reader
67+
who cannot restate the query after your first paragraph has learned nothing.
68+
- Any source column or table feeding **many** outputs — the query's hub, where a schema change has
69+
the widest blast radius.
70+
- Any table contributing **only one or two** columns, especially through a `LEFT JOIN`. That is
71+
often a lookup that could be a smaller subquery, and a join whose only job is one column is a
72+
cheap thing to get wrong.
73+
- Join predicates that look under-constrained. A join on a slowly-changing dimension without a
74+
time-range predicate fans rows out and silently multiplies aggregates — this repo has been bitten
75+
by exactly that (see the `#47` entry in `specs/CHANGELOG.md`).
76+
77+
## Limits worth stating rather than hiding
78+
79+
- `SELECT *` errors out in lineage mode by design — tracing it needs the table schemas, which the
80+
script does not have. Plan mode draws it fine.
81+
- Lineage mode also refuses a query whose output projects the same column name twice
82+
(`SELECT a.id, b.id`): `sqlglot` resolves lineage by name and would trace both to the first
83+
match, drawing a confident wrong graph. Alias them, or use plan mode.
84+
- `CREATE TABLE … AS SELECT` and `INSERT … SELECT` are unwrapped to their SELECT and diagrammed.
85+
Anything with no SELECT at all (a `DELETE`, a DDL statement) exits with a one-line message.
86+
- Dialect defaults to `databricks`; pass `--dialect` to `scripts/sql_diagram.py` directly for others.
87+
- CTEs resolve through to their base tables, but a query reading a **view** stops at the view name;
88+
the view's own definition is not expanded.
89+
- `--comments` is the only part that touches the network, and it uses the `dev` profile: the MCP
90+
service principal lacks `USE SCHEMA` on `system.billing`.

.github/workflows/onpush.yml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ jobs:
5757
uses: actions/upload-artifact@v4
5858
with:
5959
name: coverage-report
60-
path: coverage_reports/
60+
path: reports/coverage/
6161
retention-days: 14
6262

6363
# Pin the CLI to a tagged release so an upstream change can't silently break CI.

.gitignore

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ notes/
88
.pytest_cache/
99
dist/
1010
build/
11-
coverage_reports/
11+
reports/coverage/
1212
src/template.egg-info/
1313
*.pyc
1414
*.lock
@@ -23,4 +23,7 @@ resources/orders_dashboard_deploy.lvdash.json
2323
# Generated spend reports — local artifacts containing account cost figures.
2424
# A report that has been committed stays tracked (.gitignore only governs untracked files), so
2525
# new dated reports are ignored by default; `git add -f` a specific one to keep it.
26-
cost_report/
26+
reports/cost/
27+
# Generated query diagrams — same rule: ignored by default, `git add -f` one to keep it as an
28+
# example. They are derived artifacts, regenerable from the query at any time.
29+
reports/sql-diagram/

CLAUDE.md

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,13 +8,14 @@ A production-ready PySpark/Databricks ETL pipeline template using medallion arch
88

99
## Tooling: MCP servers, CLI, skills → see [`specs/tooling.md`](specs/tooling.md)
1010

11-
Developed with the [Databricks AI Dev Kit](https://github.com/databricks-solutions/ai-dev-kit) (user-level, `~/.ai-dev-kit/`). MCP config (`.mcp.json`) and `.claude/` are gitignored — user-level tooling, not committed; the exception is `.claude/commands/`, which is un-ignored and committed. Quick decision list (full reference in `specs/tooling.md`):
11+
Developed with the [Databricks AI Dev Kit](https://github.com/databricks-solutions/ai-dev-kit) **user-level tooling** (`~/.ai-dev-kit/`), never installed into or committed to this repo ([why that matters](specs/tooling.md#install-layout)). Quick decision list (full reference in [`specs/tooling.md`](specs/tooling.md)):
1212

13-
- **Workspace / UC / Jobs / Pipelines / Apps / Serving / SQL** → prefer `mcp__databricks__*` tools over `databricks` CLI shell-outs or hand-rolled SDK scripts.
14-
- **Bundle / job changes**`databricks-bundles` / `databricks-jobs` skills, and route job edits through `scripts/sdk_generate_template_job.py` + `make deploy`.
13+
- **Workspace / UC / Jobs / Pipelines / Apps / Serving / SQL** → prefer `mcp__databricks__*` tools over `databricks` CLI shell-outs or hand-rolled SDK scripts ([servers](specs/tooling.md#mcp-servers)).
14+
- **Bundle / job changes**`databricks-bundles` / `databricks-jobs` skills, and route job edits through `scripts/sdk_generate_template_job.py` + `make deploy` ([skills](specs/tooling.md#skills)).
1515
- **Library/SDK docs** (PySpark, Databricks SDK, uv, ruff) → `context7` MCP, not memory or web search.
16-
- **Cloud spend / cost analysis**`aws-billing-cost` MCP (`AWS_PROFILE=costs`) + `/project-costs` skill. **AWS docs**`aws-documentation` MCP.
16+
- **Cloud spend / cost analysis**`aws-billing-cost` MCP (`AWS_PROFILE=costs`) + `/project-costs`. **AWS docs**`aws-documentation` MCP.
1717
- Use the `dev` profile unless told otherwise (`prod` for prod ops). If MCP tools are unavailable, fall back to CLI/SDK and flag it.
18+
- **MCP calls run as the prod SP, not as you**`dev` is your user account, but the `databricks` MCP server is pinned to `DEFAULT`, which resolves to the same `template-sp` that `prod` uses. It can read/write `prod` tables; the catalog is the guardrail ([why](specs/tooling.md#mcp-runs-as-the-production-service-principal)).
1819

1920
## Commands
2021

@@ -29,6 +30,7 @@ make run env=dev # Run integration test job on a target env (dev or stagin
2930
make drop env=dev # Drop all medallion tables in a target env (schema migrations; staging/prod need yes=--yes)
3031
make whoami # Print the identity the env's profile authenticates as (runs implicitly before deploy/run/drop)
3132
make project-costs # AWS + Databricks spend report (--aws-profile costs); backs the /project-costs skill
33+
make sql-diagram sql=q.sql # Query plan (or mode=lineage) .mmd + .svg into reports/sql-diagram/; backs /sql-diagram
3234
make star-history # Regenerate the README star-history SVGs (assets/star_history*.svg) from the GitHub API
3335
```
3436

@@ -71,8 +73,8 @@ The detailed specs live in [`specs/`](specs/) — read the relevant one **before
7173

7274
- **Ask "should I open a new branch?" before executing a plan**, and **never commit directly to `main`** — cut a feature branch and land via PR (a hook blocks direct commits and pushes to `main`).
7375
- **Hold commits until asked.** Before merging, update the PR description (a hook uses it as the merge commit message body) following the What / Why / How / Validation / **Impact in prod** template in [`.github/PULL_REQUEST_TEMPLATE.md`](.github/PULL_REQUEST_TEMPLATE.md); any table schema/data change needs the production-table impact check.
74-
- **Keep docs in sync in the same commit.** Don't ship changes to the CLI surface (`main.py:arg_parser`), runtime env vars, catalog/schema model, or production guardrails without updating `README.md`, the relevant doc under `specs/`, and this file (`CLAUDE.md`) together.
75-
- **Add a `specs/CHANGELOG.md` entry immediately before merging a PR** (not while the work is in progress — scope grows, and an early entry just gets rewritten). Append-only (never edit old ones). Each entry is **exactly 3 sentences** and **at most ~5 rendered lines** (~475 chars); keep it one unwrapped paragraph — the line cap is a length budget, not a wrap width.
76+
- **Keep docs in sync in the same commit** the CLI surface (`main.py:arg_parser`), runtime env vars, catalog/schema model, and production guardrails each need `README.md` + the relevant `specs/` doc + this file updated together ([full rule](specs/workflow.md#keep-docs-in-sync)).
77+
- **Add a `specs/CHANGELOG.md` entry immediately before merging** — never earlier; append-only; one unwrapped paragraph, ~1000 characters ([full rule](specs/workflow.md#changelog-discipline)).
7678

7779
## Keep It Simple
7880

Makefile

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,13 @@ project-costs: aws-profile ?= costs
3737
project-costs:
3838
uv run python ./scripts/project_costs.py $(if $(aws-profile),--aws-profile $(aws-profile),)
3939

40+
# Diagram a SQL query with sqlglot, writing .mmd + .svg into reports/sql-diagram/. Pass the query
41+
# with sql=path/to/query.sql (stdin if omitted); optionally name=basename, mode=plan|lineage
42+
# (default plan: the query's steps) and comments=1 (looks up each table's Unity Catalog comment).
43+
sql-diagram:
44+
uv run python ./scripts/sql_diagram.py $(if $(sql),--file $(sql),) $(if $(name),--name $(name),) \
45+
$(if $(mode),--mode $(mode),) $(if $(comments),--comments,)
46+
4047
# Regenerate the README star-history chart from the GitHub API. The SVGs are committed,
4148
# so the README renders from this repo rather than a third-party chart service.
4249
star-history:

README.md

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,8 @@ This project template demonstrates how to:
5858
- utilize the [Databricks SDK for Python](https://docs.databricks.com/en/dev-tools/sdk-python.html) to manage catalogs, schemas, workspaces, and accounts. Refer to the `scripts` folder for examples.
5959
- utilize [Databricks Unity Catalog](https://www.databricks.com/product/unity-catalog) to manage permissions and get data lineage.
6060
- enforce production guardrails out of the box — identity-locked CI deploys, a health-check task, wheel version pinning, per-task timeouts, schema-drift guards, queued runs, and on-call alerting.
61-
- track project cloud spend in USD across AWS (Cost Explorer) and Databricks ([`system.billing`](https://docs.databricks.com/aws/en/admin/system-tables/pricing)) with `make project-costs` — see an [example report](cost_report/2026-07-16.md).
61+
- track project cloud spend in USD across AWS (Cost Explorer) and Databricks ([`system.billing`](https://docs.databricks.com/aws/en/admin/system-tables/pricing)) with `make project-costs` — see an [example report](reports/cost/2026-07-22.md).
62+
- diagram any SQL query with `make sql-diagram sql=<file>`[`sqlglot`](https://github.com/tobymao/sqlglot) parses the AST and writes a Mermaid flowchart plus a standalone SVG to `reports/sql-diagram/`, either as the query's execution steps (each scan, each join with its keys, filter, aggregate, sort — [example](reports/sql-diagram/job_spend_plan.svg)) or as column-level lineage (`mode=lineage`), so what is drawn is what the query says rather than what a model guessed.
6263
- utilize serverless job clusters on [Databricks Free Edition](https://docs.databricks.com/aws/en/getting-started/free-edition) to deploy your pipelines.
6364

6465

pyproject.toml

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ dev = [
2424
"pytest-cov==5.0.0",
2525
"pyspark==4.1.0",
2626
"databricks-bundles>=0.298.0",
27+
"sqlglot==30.13.0",
2728
]
2829

2930
[project.scripts]
@@ -46,7 +47,7 @@ testpaths = ["tests"]
4647
python_files = ["test_*.py", "*_test.py", "unit_test_*.py"]
4748
python_classes = ["Test*"]
4849
python_functions = ["test_*"]
49-
addopts = "--cov=. --cov-report=term --cov-report=xml:coverage_reports/coverage.xml --cov-report=html:coverage_reports/html"
50+
addopts = "--cov=. --cov-report=term --cov-report=xml:reports/coverage/coverage.xml --cov-report=html:reports/coverage/html"
5051

5152
[tool.ruff]
5253
line-length = 120

0 commit comments

Comments
 (0)