#224: Databricks deterministic sampling + materialise_sample + get_row_count - #257
Merged
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (5)
✅ Files skipped from review due to trivial changes (2)
🚧 Files skipped from review as they are similar to previous changes (3)
📝 WalkthroughWalkthroughThis change relaxes ChangesDatabricks adapter validation and sampling
Sequence Diagram(s)sequenceDiagram
participant DatabricksAdapter
participant FakeDatabricksConnection
participant map_databricks_exception
DatabricksAdapter->>FakeDatabricksConnection: execute COUNT(*), sample SQL, or CTAS SQL
FakeDatabricksConnection-->>DatabricksAdapter: rows or exception
DatabricksAdapter->>map_databricks_exception: classify connector failure
map_databricks_exception-->>DatabricksAdapter: WarehouseError subtype
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related issues
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
Comment |
…olumn_stats parity issue #258
…g catalog names Add validate_catalog_or_project to _sql_safety.py composing the existing validate_identifier and validate_project_id, and route TableRef.__post_init__'s project check through it (DEC-005 of #224). Short Unity Catalog catalogs (main), default catalogs (workspace), and short Snowflake databases now construct, while injection-shaped project values still raise InvalidIdentifierError. validate_project_id is unchanged; TableRef stays dialect-neutral (no new field).
…Unity Catalog catalogs
… cleanup + error mapping #224 US-002 — the Databricks analogue of the Snowflake #122/#124 connection / cleanup / error-mapping work. - Add tests/warehouse/_fake_databricks.py (FakeDatabricksConnection) mirroring the Snowflake fake: injectable session_id, close_raises, a cursor with execute/fetchall/description/close, expect_execute/assert_all_expectations_met, and close_call_count. - DatabricksAdapter._cleanup_active_session now reads the connection's opaque session id (defensive: session_id attr first, get_session_id_hex() fallback), emits the raw id ONLY in the swallow-and-warn failure WARNING (no manual cleanup command, no auto-expire countdown — temp objects are session-local), and the hashed id (_hash_session_id) in the success INFO. State resets in finally for idempotency; _connection is never nulled so a second __exit__ is a clean no-op. - Extend map_databricks_exception to the full taxonomy mirroring map_snowflake_exception: auth -> WarehouseAuthError; table-not-found -> TableNotFoundError; unresolved-column -> ColumnNotFoundError; residual Spark/SQL error -> QuerySyntaxError; else passthrough. The Table/Column/Syntax split is scoped to ServerOperationError/ProgrammingError so a transient OperationalError is not mis-mapped. The databricks.sql.exc import stays lazy + confined to the shim. No new WarehouseError subclass. - tests/warehouse/test_databricks_adapter.py covers the fake, the connection seam, fail-soft cleanup (raw-id WARNING / hashed-id INFO / idempotent second exit), repr redaction, and the error-mapper taxonomy (built from the connector's own databricks.sql.exc classes).
…+ fake + cleanup + error mapping
…ws (#224 US-003) Implement the first real DatabricksAdapter warehouse I/O, mirroring the Snowflake adapter (DEC-002 / DEC-003 / DEC-008 of #224): - _quote / _fold / _quote_identifier: fold-then-quote per component using DATABRICKS_DIALECT (identifier_case='lower' + backtick, quote_qualified_per_component=True), identical folding to the prune compiler so CREATE-vs-REFERENCE never diverge (the #124 lesson). - _execute / _execute_to_dicts / _rows_to_dicts: connection-bound cursor execution, error mapping via map_databricks_exception, dict-row shaping via cursor.description (no dict-cursor dependency). - get_row_count: SELECT COUNT(*) (DESCRIBE DETAIL has no reliable numRows; COUNT is metadata-cheap on Delta) -> int, WarehouseError -> None; overrides the ABC RowCountNotSupportedError degrade. - _resolve_sample_bucket: shared fail-loud sizing (UnknownTableSizeError / SamplingRequiresPartitionFilterError / bucket=1000 fallback / bucket math). - sample_rows: n<=0 ValueError; inline-predicate shape via render_sample_select(order_by_hash=True) (Databricks has no Snowflake-style HASH(*) predicate restriction; sample_hash_in_projection=False); hash expression read from the dialect, never hard-coded. - _render_partition_filter: datetime/date/str via dialect literal templates + escape, column fold-then-quoted. Tests: extend test_databricks_adapter.py (determinism, all three sizing branches, get_row_count happy/dict/error->None/empty/passthrough, dict-row shaping, partition-filter rendering, project=None two-part quoting, SDK error mapping); new ungated sqlglot databricks-dialect parse-guard (test_databricks_sql_parse.py) over every emitted statement plus a planted-violation self-check. Update the two now-stale stub tests in test_databricks_stub.py (sample_rows + get_row_count graduated).
…sizing + sample_rows
wjduenow
marked this pull request as ready for review
June 25, 2026 20:05
…224 US-004) Implements the materialise + test-execution surface on DatabricksAdapter, mirroring the Snowflake adapter and reusing the US-003 helpers. - materialise_sample: overrides the ABC MaterialisationNotSupportedError degrade with a qualified `CREATE TEMPORARY TABLE <cat>.<sch>._sf_sample_<run_id> AS <inline-predicate sample body>`. run_id via the shared _compute_run_id recipe (byte-identical to BigQuery/Snowflake); sample body via the shared render_sample_select (inline shape, sample_hash_in_projection=False). Returns the qualified temp TableRef; pins the connection so a follow-up run_test_sql reaches the session-scoped temp table. Cursor errors route through map_databricks_exception then wrap in MaterialisationFailedError(cause=...); one INFO log with the hashed session id + table + sample_rows + run_id. - run_test_sql: overrides the NotImplementedError stub. Count via `SELECT COUNT(*) AS failures FROM (<sql>) AS t`; capture via a second per-row `to_json(struct(*)) ... LIMIT k` query, json.loads-ing each JSON string individually (not a single outer decode like Snowflake's VARIANT). Case-insensitive column resolution via cursor.description. - Extends the ungated sqlglot databricks parse-guard over the emitted CTAS + COUNT wrap + to_json capture SQL; adds the #116 custom_sql {{ this }} substitution test (source name -> quoted temp ref, no full-scan leakage). - Retires the two now-obsolete stub assertions in test_databricks_stub.py. #226 live-cert caveats (shape-only here, noted in test docstrings): qualified temp-table-name acceptance, connector session persistence across queries, and the per-row to_json capture marshalling.
…le + run_test_sql
There was a problem hiding this comment.
Pull request overview
Adds a “super plan” document for issue #224, outlining the intended Databricks adapter implementation for deterministic sampling, sample materialisation, and row-count sizing, including decisions (DEC-001…DEC-012) and a story-by-story breakdown to drive beads task creation and later implementation PRs.
Changes:
- Introduces a detailed plan for implementing
sample_rows,get_row_count,materialise_sample,run_test_sql, andcolumn_statsfor the Databricks adapter. - Records key architectural decisions and testing strategy (fakes + sqlglot parse guards + gated live cert in #226).
- Documents a beads task breakdown and links the Snowflake
column_statsparity follow-up.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
Caution Failed to replace (edit) comment. This is likely due to insufficient permissions or the comment being deleted. Error details |
|
Caution Failed to replace (edit) comment. This is likely due to insufficient permissions or the comment being deleted. Error details |
…224 US-005) Implement DatabricksAdapter.column_stats, overriding the v0.x NotImplementedError stub. Databricks ships aggregate-only profiling AHEAD of Snowflake (which stubs it); the Snowflake-parity decision is tracked as GitHub issue #258. A single aggregate query over the fold-then-quoted table computes the ColumnStats contract for the fold-then-quoted column, mirroring BigQueryAdapter.column_stats semantics in Spark SQL: - count = COUNT(<col>) (non-null count) - distinct= COUNT(DISTINCT <col>) - nulls = COUNT_IF(<col> IS NULL) - min/max = MIN(<col>) / MAX(<col>) - data_type = MAX(typeof(<col>)) (NULL on empty table -> "") Reuses the US-003 helpers (_quote / _quote_identifier / _execute_to_dicts); no context-manager batching (single per-call round-trip, mirroring sample_rows / run_test_sql / get_row_count). Errors route through map_databricks_exception via _execute_to_dicts; the column is validated by validate_identifier before reaching SQL (DEC-013). Tests: ColumnStats shape against FakeDatabricksConnection (every field), identifier fold-to-lower, case-insensitive alias resolution, empty-table data_type coercion, two-part quoting, error mapping, identifier validation. Adds column_stats coverage to the ungated sqlglot databricks parse-guard. Removes the obsolete column_stats NotImplementedError stub test. Traces to DEC-001, DEC-011 (plans/super/224-databricks-sampling.md).
… fix stale skeleton docstrings, document column_stats complex-type MIN/MAX divergence (#226), add partition-filter run_id test
…on in warehouse-adapters.md, ops-doc sampling surface + #226 live-cert ledger
wjduenow
changed the title
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
src/signalforge/warehouse/models.py (1)
391-396: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueStale class docstring vs. relaxed contract.
The docstring still reads "Fully-qualified BigQuery table identity" and describes
projectpurely as a BigQuery default-project deferral, but the validation is now dialect-neutral (BigQuery project / Snowflake database / Unity Catalog catalog). Consider refreshing the class docstring so the public contract matches the relaxedvalidate_catalog_or_projectbehavior introduced below.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/signalforge/warehouse/models.py` around lines 391 - 396, Update the class docstring for the table identity model so it matches the relaxed, dialect-neutral contract enforced by validate_catalog_or_project. Replace the BigQuery-only wording with language that covers the supported catalog/project behavior across backends, and adjust the description of the optional top-level identifier so it no longer implies only BigQuery default-project deferral.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/signalforge/warehouse/adapters/databricks.py`:
- Around line 385-393: The cursor handling in Databricks query helpers is
leaking resources because `_execute`, `_execute_to_dicts`, and
`materialise_sample` open cursors via `self._get_connection().cursor()` without
guaranteeing cleanup. Update each of these methods to wrap cursor use in a
try/finally block so the cursor is always closed after `cursor.execute(...)`,
`fetchall()`, or exception mapping, and keep the existing exception behavior in
`map_databricks_exception` intact while ensuring release on both success and
failure.
---
Nitpick comments:
In `@src/signalforge/warehouse/models.py`:
- Around line 391-396: Update the class docstring for the table identity model
so it matches the relaxed, dialect-neutral contract enforced by
validate_catalog_or_project. Replace the BigQuery-only wording with language
that covers the supported catalog/project behavior across backends, and adjust
the description of the optional top-level identifier so it no longer implies
only BigQuery default-project deferral.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 7ff16ede-adf5-45ba-8408-8569597814ce
📒 Files selected for processing (13)
.claude/rules/warehouse-adapters.mddocs/warehouse-adapter-ops.mdplans/super/224-databricks-sampling.mdsrc/signalforge/warehouse/_sql_safety.pysrc/signalforge/warehouse/adapters/_databricks_client.pysrc/signalforge/warehouse/adapters/databricks.pysrc/signalforge/warehouse/models.pytests/warehouse/_fake_databricks.pytests/warehouse/test_databricks_adapter.pytests/warehouse/test_databricks_sql_parse.pytests/warehouse/test_databricks_stub.pytests/warehouse/test_models.pytests/warehouse/test_sql_safety.py
…CodeRabbit), reconcile plan phase to complete
Owner
Author
PR Review SummaryFixed (1 item)
Reconciled (2 items)
All 7 stories merged; full suite green (4387 passed); pyright 0 errors. No deferrals. |
…view compounding)
This was referenced Jun 29, 2026
wjduenow
added a commit
that referenced
this pull request
Jul 2, 2026
…_execute_to_dicts/run_test_sql) Wrap the execute/fetch in try/finally cursor.close() for the three cursor-opening methods that previously leaked server-side handles, mirroring _execute_scalar (the one already-correct path) and the Databricks PR #257 shape. _execute_to_dicts uses an OUTER try/finally (close) wrapping an INNER try/except (exception mapping) so the cursor closes AFTER _rows_to_dicts reads cursor.description. run_test_sql likewise closes only after description + rows are captured into locals. TDD: extend FakeSnowflakeConnection to track vended cursors (.cursors) and expose per-cursor .closed, then six new tests in tests/warehouse/ test_snowflake_adapter.py assert closed on both success and failure paths. Traces to DEC-007 of plans/super/258-snowflake-column-stats.md.
wjduenow
added a commit
that referenced
this pull request
Jul 3, 2026
* Add super plan for #258: Snowflake column_stats parity * Mark #258 plan published (PR #262) * Devolve #258 plan to beads (epic bd_1-scaffolding-om0) * bd_1-scaffolding-om0.1: fix Snowflake adapter cursor leaks (_execute/_execute_to_dicts/run_test_sql) Wrap the execute/fetch in try/finally cursor.close() for the three cursor-opening methods that previously leaked server-side handles, mirroring _execute_scalar (the one already-correct path) and the Databricks PR #257 shape. _execute_to_dicts uses an OUTER try/finally (close) wrapping an INNER try/except (exception mapping) so the cursor closes AFTER _rows_to_dicts reads cursor.description. run_test_sql likewise closes only after description + rows are captured into locals. TDD: extend FakeSnowflakeConnection to track vended cursors (.cursors) and expose per-cursor .closed, then six new tests in tests/warehouse/ test_snowflake_adapter.py assert closed on both success and failure paths. Traces to DEC-007 of plans/super/258-snowflake-column-stats.md. * bd_1-scaffolding-om0.2: implement Snowflake column_stats (full-batch, catalog pre-filter) Replace the NotImplementedError stub with a BigQuery-style context-manager batched column_stats (issue #258 US-002; DEC-001..DEC-006): - __enter__/__exit__ open/reset per-table pending + results caches, coexisting with the connection-session state. - column_stats: validate identifier -> DEC-025 with-block guard -> cache hit -> queue + flush. - _flush_column_stats_batch: one INFORMATION_SCHEMA.COLUMNS catalog lookup (mirroring _get_num_rows escaping) for the data_type field + MIN/MAX skip decision, then ONE aggregate for every queued column (index-suffixed aliases; COUNT(*)-COUNT null count per DEC-005; MIN/MAX gated per DEC-003). - _get_column_types keyed by lower(COLUMN_NAME); a requested column absent from the catalog raises ColumnNotFoundError before the aggregate. - _COMPLEX_SNOWFLAKE_TYPES {ARRAY,OBJECT,VARIANT,GEOGRAPHY,GEOMETRY} + _is_complex_snowflake_type (BINARY orderable); Decimal->float min/max (DEC-004). Full default-suite unit coverage via the hand-fake (no drift detector - ColumnStats is frozen/in-process). Drops the obsolete stub NotImplementedError test. * bd_1-scaffolding-om0.5: flip Snowflake column_stats parity docs (ops/README/CHANGELOG) * bd_1-scaffolding-om0.4: gated live cert for Snowflake column_stats (scalar + complex skip-set + aggregate-only smoke) * bd_1-scaffolding-om0.3: fakesnow offline execution tests for Snowflake column_stats Add five fakesnow-executed offline tests driving the real SnowflakeAdapter column_stats path end-to-end against DuckDB (US-003, DEC-008): - scalar NUMBER column: COUNT / COUNT(DISTINCT) / COUNT(*)-COUNT(col) null count / MIN / MAX execute; Decimal min/max coerced to float (DEC-004) - INFORMATION_SCHEMA.COLUMNS catalog lookup (_get_column_types) round-trip — fakesnow returns consumable Snowflake DATA_TYPE strings, so this is a full EXECUTION assertion (no parse-only degrade needed) - VARCHAR string MIN/MAX bounds pass through unchanged - VARIANT complex-type MIN/MAX skip decision (DEC-003) executes; live cert owns the analysis-time-rejection validation (#227 lesson) - all-NULL scalar column: null-count arithmetic + empty MIN/MAX None path Document in the module docstring that column_stats is fully executable under fakesnow (no parse-only degrade) and that ColumnStats needs NO drift detector (frozen, in-process, never read back from disk). * bd_1-scaffolding-om0.6: Quality gate — coerce out-of-union min/max (BINARY/TIME), doc GEOGRAPHY DISTINCT limit QG review (2 independent lenses) found BINARY (bytearray) and TIME (datetime.time) column MIN/MAX return types are outside ColumnStats.min/max's union, raising a non-WarehouseError ValidationError that fails the whole batch at the panic tier. Fix: _coerce_min_max nulls any out-of-union value (return-type analogue of the skip-set). Add BINARY/TIME coercion tests (offline + gated live columns). Document the pre-existing GEOGRAPHY/GEOMETRY COUNT(DISTINCT) limitation. * bd_1-scaffolding-om0.7: Patterns & Memory — flip Snowflake column_stats to shipped (#258) Update warehouse-adapters.md: flip the 'column_stats stays NotImplementedError' / 'parity is the open follow-up #258' statements to shipped; add the § 'Snowflake column_stats (issue #258)' subsection capturing the two-tier MIN/MAX defence (SQL-raise skip-set vs out-of-union return-type coercion net), the GEOGRAPHY DISTINCT limitation, and the deferred live cert. * #258: Address PR review feedback (CodeRabbit + Copilot) - CRITICAL (CodeRabbit): drain column_stats pending batch UP FRONT so a failed column (ColumnNotFoundError / GEOGRAPHY COUNT(DISTINCT)) no longer poisons every subsequent column of the same table within the with-block. + regression test. - Live test: move CREATE/INSERT inside the outer try/finally (no orphaned table on setup failure). - Docs/README/CHANGELOG: qualify aggregate-only with the GEOGRAPHY/GEOMETRY COUNT(DISTINCT) limitation; soften 'certified'/'confirmed' to maintainer-run gated cert. - MD018: reflow lines so #258 is not at a wrapped-line start (ops doc + README). - warehouse-adapters.md: update the stale Snowflake cursor-leak note (only materialise_sample leaks now). * #258: wire key-pair (JWT) auth through make_real_client (enables live cert on MFA accounts) The #120 profile model + from_profile + SnowflakeAdapter.__init__ already parsed and stored private_key_path / private_key_passphrase / authenticator, but make_real_client dropped them at the connect() call (they were never used). Thread them through: key-pair auth when private_key_path is set (private_key_file [+ _pwd]), else password; authenticator passed through when set. This is the non-interactive path MFA-enforced accounts require. The gated live cert (_make_adapter + _skip_reason) now accepts key-pair OR password auth. * #258: flip column_stats live cert to live-certified (passed via key-pair auth) Ran the gated column_stats live cert against a real Snowflake warehouse (key-pair / JWT auth — the restored account enforces MFA, which the headless password path can't satisfy). Result: scalar min/max populate; ARRAY/OBJECT/VARIANT return min=max=None without raising (DEC-003 skip-set confirmed complete); BINARY/TIME coerce to None. Since the cert profiles ARRAY/OBJECT/VARIANT (emitting COUNT(DISTINCT) on them) and passed, COUNT(DISTINCT) on semi-structured types does NOT raise — only GEOGRAPHY/GEOMETRY remain the documented limit. Docs/rules updated from 'pending a maintainer run' to 'live-certified'. * #258: fix Snowflake connection reuse across with-blocks + wire key-pair into e2e smoke Exposed by the aggregate-only e2e: generate with safety=aggregate-only uses one adapter across two 'with adapter:' blocks (safety-aggregate column_stats, then prune_tests). Snowflake's __exit__ closes the connection (the connection IS the session) but #122 didn't null self._connection, so the 2nd block reused the closed connection -> 250002 (08003): Connection is closed. Fix: track self._owns_connection (True only on lazy build); cleanup nulls _connection after close only when owned, so the next block rebuilds; injected fakes stay intact. + 2 lifecycle regression tests. Also wire key-pair (JWT) auth into the e2e smoke profiles.yml (SNOWFLAKE_PRIVATE_KEY_PATH), so the aggregate-only generate e2e is now live-certified against a real MFA-enforced Snowflake account.
Merged
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Super plan for #224 — Databricks adapter sampling surface (
sample_rows,get_row_count,run_test_sql,materialise_sample,column_stats). Epic #219; follows #221/#222/#223; models on Snowflake #122/#139/#140.Phase: complete (implemented + reviewed; live cert deferred to #226)
Stories: 5 implementation + Quality Gate + Patterns & Memory
Decisions: 12 (DEC-001 … DEC-012)
Key decisions
CREATE TEMPORARY TABLEmirroring Snowflake feat: Snowflake deterministic sampling (sample_rows HASH-mod) + materialise_sample (TEMP TABLE) #122 (reuses TableRef+compiler path, zero model churn). Qualified-temp validity + connector session persistence are Databricks: test harness + gated live e2e + ops docs #226 live-cert blockers.SELECT COUNT(*)(the ticket'sDESCRIBE DETAIL → numRowscolumn does not exist; numRows lives in thestatisticsmap only afterANALYZE).main; fixes latent short-Snowflake-DB case).column_stats(ahead of Snowflake); follow-up issue to decide Snowflake parity.HASH(*)predicate restriction); noDATABRICKS_DIALECTchange.Plan document
See
plans/super/224-databricks-sampling.md.Next steps
Implementation landed (Ralph epic
bd_1-scaffolding-v4o). All 5 stories + Quality Gate + Patterns & Memory merged; full suite green (4383 passed). Shipped:sample_rows/get_row_count(COUNT(*) sizing) /materialise_sample(qualified CREATE TEMPORARY TABLE) /run_test_sql/column_stats,validate_catalog_or_projectfor Unity Catalog catalogs, fakes + ungated sqlglotdatabricksparse-guard. Live validity (qualified-temp acceptance, connector session persistence, to_json capture marshalling, column_stats complex-type MIN/MAX) is the #226 live-cert ledger. Snowflakecolumn_statsparity → #258.Summary by CodeRabbit
New Features
Bug Fixes
Tests