From 04330c938336a50140a5c7419cac7df9621ec7d9 Mon Sep 17 00:00:00 2001 From: Joe Harris Date: Fri, 9 Oct 2026 18:20:41 -0400 Subject: [PATCH 1/5] fix(tuning): record tuning statements issued through .sql() and .query() The applied-tuning ledger only saw execute, cursor().execute and executescript, so tuned schema DDL on Spark, PySpark, LakeSail, Velox, DataFusion and BigQuery was never recorded and a tuned Spark run reported noop. Adapters now declare ledger_execute_verbs and the recording proxy wraps those methods. Readbacks stay unrecorded and capture failures never reach the run. BigQuery statements are recorded when the job reaches its final state, so a job that fails after submission is recorded as failed, and dry-run jobs are skipped. Multi-statement strings are split before recording, including several statements on one line, so a leading SET can no longer make a tuning statement look transient. The schema footprint now also matches PARTITIONED BY, CLUSTERED BY, DISTKEY, DISTSTYLE, DUPLICATE KEY, PRIMARY INDEX, distribution =, partitioning =, sorted_by and bucketed_by. Spark table definitions now place USING before a trailing PARTITIONED BY or CLUSTER BY clause; Spark rejects the previous order. --- CHANGELOG.md | 11 + benchbox/core/tuning/applied_ledger.py | 145 +++++-- benchbox/platforms/_spark_helpers.py | 26 +- benchbox/platforms/base/adapter.py | 8 +- benchbox/platforms/base/tuning.py | 7 +- benchbox/platforms/base/tuning_config.py | 7 +- benchbox/platforms/bigquery.py | 1 + benchbox/platforms/datafusion.py | 1 + .../test_spark_ledger_capture_live.py | 70 ++++ .../core/tuning/test_ledger_execute_verbs.py | 378 ++++++++++++++++++ tests/unit/platforms/test_spark_helpers.py | 8 + 11 files changed, 624 insertions(+), 38 deletions(-) create mode 100644 tests/integration/platforms/test_spark_ledger_capture_live.py create mode 100644 tests/unit/core/tuning/test_ledger_execute_verbs.py diff --git a/CHANGELOG.md b/CHANGELOG.md index a559c418a0..6956b1961b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -310,6 +310,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 there. Receipt entries for constraints carry `kind: "constraint"` and a `constraint_type`. Quoted non-ASCII table names, such as "É", are now found consistently with DuckDB's `lower(table_name)` catalog filters. +- **Tuning statements issued through `.sql()` and `.query()` are recorded.** + Spark, PySpark, LakeSail, Velox, DataFusion and BigQuery tuned schema DDL now + reaches the applied-tuning ledger, so a tuned Spark run no longer reports + `noop`. BigQuery statements are recorded with the final state of the job, not + at submission. Multi-statement strings are split before recording, so a + leading `SET` can no longer hide a tuning statement. Schema capture also + recognises `PARTITIONED BY`, `CLUSTERED BY`, `DISTKEY`, `DISTSTYLE`, + `DUPLICATE KEY`, `PRIMARY INDEX`, `distribution =`, `partitioning =`, + `sorted_by` and `bucketed_by`. + Spark table definitions now put `USING ` before a trailing + `PARTITIONED BY` or `CLUSTER BY`, which Spark requires. ### Changed diff --git a/benchbox/core/tuning/applied_ledger.py b/benchbox/core/tuning/applied_ledger.py index 5725242144..6b0fd532ea 100644 --- a/benchbox/core/tuning/applied_ledger.py +++ b/benchbox/core/tuning/applied_ledger.py @@ -67,8 +67,9 @@ re.IGNORECASE, ) _SCHEMA_TUNING_FOOTPRINT_RE = re.compile( - r"\b(?:cluster\s+by|distributed\s+by|foreign\s+key|order\s+by|partition\s+by|primary\s+key|" - r"sortkey|unique|check)\b", + r"\b(?:cluster\s+by|clustered\s+by|distributed\s+by|foreign\s+key|order\s+by|partition\s+by|" + r"partitioned\s+by|primary\s+key|primary\s+index|duplicate\s+key|sortkey|distkey|diststyle|" + r"sorted_by|bucketed_by|unique|check)\b|\b(?:distribution|partitioning)\s*=", re.IGNORECASE, ) @@ -97,16 +98,18 @@ def _split_sql_script(script: Any) -> list[str]: import sqlite3 statements: list[str] = [] - pending: list[str] = [] - for line in text.splitlines(keepends=True): - pending.append(line) - candidate = "".join(pending) + start = 0 + for index, char in enumerate(text): + if char != ";": + continue + candidate = text[start : index + 1] if sqlite3.complete_statement(candidate): if candidate.strip(): statements.append(candidate.strip()) - pending.clear() - if pending and "".join(pending).strip(): - statements.append("".join(pending).strip()) + start = index + 1 + tail = text[start:].strip() + if tail: + statements.append(tail) return statements except Exception: # pragma: no cover return [text] @@ -322,9 +325,94 @@ def _should_record(self, statement: Any) -> bool: logger.debug("applied-ledger statement filter degraded: %s", exc) return True + def _recordable_statements(self, statement: Any) -> list[str]: + try: + pieces = _split_sql_script(statement) + if len(pieces) <= 1: + pieces = [statement] + return [str(piece) for piece in pieces if self._should_record(piece)] + except Exception as exc: + logger.debug("applied-ledger statement capture degraded: %s", exc) + return [] + + def _record_all(self, statements: list[str], status: str, error: Any | None = None) -> None: + for statement in statements: + self._ledger.record(statement, self._phase, status=status, error=error) + + def _run(self, fn: Any, statement: Any, args: tuple, kwargs: dict) -> Any: + statements = self._recordable_statements(statement) + try: + result = fn(statement, *args, **kwargs) + except Exception as exc: + self._record_all(statements, STATEMENT_FAILED, exc) + raise + self._record_all(statements, EXECUTED) + return result + + def _run_verb(self, fn: Any, statement: Any, args: tuple, kwargs: dict) -> Any: + statements = [] if _is_dry_run(args, kwargs) else self._recordable_statements(statement) + try: + result = fn(statement, *args, **kwargs) + except Exception as exc: + self._record_all(statements, STATEMENT_FAILED, exc) + raise + if not statements: + return result + if callable(getattr(result, "result", None)) and callable(getattr(result, "done", None)): + return _RecordingJob(result, self._ledger, self._phase, statements) + self._record_all(statements, EXECUTED) + return result + + +def _is_dry_run(args: tuple, kwargs: dict) -> bool: + config = kwargs.get("job_config") or (args[0] if args else None) + return getattr(config, "dry_run", False) is True + + +class _RecordingJob: + __slots__ = ("_job", "_ledger", "_phase", "_statements", "_settled") + + def __init__(self, job: Any, ledger: AppliedTuningLedger, phase: str, statements: list[str]) -> None: + object.__setattr__(self, "_job", job) + object.__setattr__(self, "_ledger", ledger) + object.__setattr__(self, "_phase", phase) + object.__setattr__(self, "_statements", statements) + object.__setattr__(self, "_settled", False) + + def _settle(self, status: str, error: Any | None = None) -> None: + if self._settled: + return + object.__setattr__(self, "_settled", True) + for statement in self._statements: + self._ledger.record(statement, self._phase, status=status, error=error) + + def result(self, *args: Any, **kwargs: Any) -> Any: + from concurrent.futures import TimeoutError as FutureTimeoutError + + try: + outcome = self._job.result(*args, **kwargs) + except FutureTimeoutError: + raise + except Exception as exc: + self._settle(STATEMENT_FAILED, exc) + raise + self._settle(EXECUTED) + return outcome + + def exception(self, *args: Any, **kwargs: Any) -> Any: + error = self._job.exception(*args, **kwargs) + if error is not None: + self._settle(STATEMENT_FAILED, error) + elif self._job.done(): + self._settle(EXECUTED) + return error + + def __getattr__(self, name: str) -> Any: + return getattr(self._job, name) + class RecordingConnection(_RecordingProxy): - __slots__ = ("_conn",) + __slots__ = ("_conn", "_verbs") def __init__( self, @@ -332,9 +420,11 @@ def __init__( ledger: AppliedTuningLedger, phase: str, statement_filter: Callable[[Any], bool] | None = None, + execute_verbs: tuple[str, ...] = (), ) -> None: super().__init__(ledger, phase, statement_filter) object.__setattr__(self, "_conn", connection) + object.__setattr__(self, "_verbs", frozenset(execute_verbs)) def execute(self, statement: Any, *args: Any, **kwargs: Any) -> Any: return self._run(self._conn.execute, statement, args, kwargs) @@ -349,20 +439,15 @@ def executescript(self, script: Any, *args: Any, **kwargs: Any) -> Any: self._ledger.record(statement, self._phase, status=EXECUTED) return result - def _run(self, fn: Any, statement: Any, args: tuple, kwargs: dict) -> Any: - record = self._should_record(statement) - try: - result = fn(statement, *args, **kwargs) - except Exception as exc: - if record: - self._ledger.record(statement, self._phase, status=STATEMENT_FAILED, error=exc) - raise - if record: - self._ledger.record(statement, self._phase, status=EXECUTED) - return result - def __getattr__(self, name: str) -> Any: - return getattr(self._conn, name) + attribute = getattr(self._conn, name) + if name in self._verbs and callable(attribute): + + def call(statement: Any, *args: Any, **kwargs: Any) -> Any: + return self._run_verb(attribute, statement, args, kwargs) + + return call + return attribute def __setattr__(self, name: str, value: Any) -> None: setattr(self._conn, name, value) @@ -393,16 +478,7 @@ def __init__( object.__setattr__(self, "_cur", cursor) def execute(self, statement: Any, *args: Any, **kwargs: Any) -> Any: - record = self._should_record(statement) - try: - result = self._cur.execute(statement, *args, **kwargs) - except Exception as exc: - if record: - self._ledger.record(statement, self._phase, status=STATEMENT_FAILED, error=exc) - raise - if record: - self._ledger.record(statement, self._phase, status=EXECUTED) - return result + return self._run(self._cur.execute, statement, args, kwargs) def __getattr__(self, name: str) -> Any: return getattr(self._cur, name) @@ -426,11 +502,12 @@ def recording_connection( ledger: AppliedTuningLedger | None, phase: str, statement_filter: Callable[[Any], bool] | None = None, + execute_verbs: tuple[str, ...] = (), ) -> Any: if ledger is None: return connection try: - return RecordingConnection(connection, ledger, phase, statement_filter) + return RecordingConnection(connection, ledger, phase, statement_filter, execute_verbs) except Exception as exc: # pragma: no cover logger.debug("applied-ledger connection wrap degraded: %s", exc) return connection diff --git a/benchbox/platforms/_spark_helpers.py b/benchbox/platforms/_spark_helpers.py index 7868cd0df9..5f6e4c6f01 100644 --- a/benchbox/platforms/_spark_helpers.py +++ b/benchbox/platforms/_spark_helpers.py @@ -157,6 +157,26 @@ def get_spark_query_plan( _LEADING_COMMENTS_RE = re.compile(r"\A(?:\s*(?:--[^\n]*(?:\n|$)|/\*.*?\*/))+\s*", re.DOTALL) _FIXED_SIZE_ARRAY_TYPE_RE = re.compile(r"\b([A-Za-z]\w*)\s*\[\s*\d+\s*\]") _ARRAY_FIXED_SIZE_SUFFIX_RE = re.compile(r"(>)\s*\[\s*\d+\s*\]") +_CREATE_TABLE_COLUMN_LIST_RE = re.compile( + r"\A\s*CREATE\s+TABLE\s+(?:IF\s+NOT\s+EXISTS\s+)?[^\s(]+\s*\(", + re.IGNORECASE, +) + + +def _column_list_end(statement: str) -> int: + match = _CREATE_TABLE_COLUMN_LIST_RE.match(statement) + if not match: + return len(statement) + depth = 1 + cursor = match.end() + while cursor < len(statement) and depth: + ch = statement[cursor] + if ch == "(": + depth += 1 + elif ch == ")": + depth -= 1 + cursor += 1 + return cursor if depth == 0 else len(statement) def _strip_balanced_paren_constraints(statement: str) -> str: @@ -227,7 +247,9 @@ def optimize_spark_table_definition( if ")" in statement: fmt = (table_format or "parquet").upper() - statement = statement.rstrip(";").rstrip() + f" USING {fmt}" + statement = statement.rstrip(";").rstrip() + end = _column_list_end(statement) + statement = f"{statement[:end]} USING {fmt}{statement[end:]}" return statement @@ -255,6 +277,8 @@ def purge_orphaned_warehouse_directory(spark: Any, *, logger: logging.Logger) -> class SparkLikeAdapterMixin: + ledger_execute_verbs = ("sql",) + def apply_constraint_configuration( self, primary_key_config: Any, diff --git a/benchbox/platforms/base/adapter.py b/benchbox/platforms/base/adapter.py index 6b198f504f..9465461681 100644 --- a/benchbox/platforms/base/adapter.py +++ b/benchbox/platforms/base/adapter.py @@ -594,7 +594,12 @@ def run_enhanced_benchmark(self, benchmark, **run_config) -> EnhancedBenchmarkRe benchmark_type = run_config.get("benchmark_type", "olap") self.configure_for_benchmark( - recording_connection(connection, self._applied_tuning_ledger, PHASE_SESSION), + recording_connection( + connection, + self._applied_tuning_ledger, + PHASE_SESSION, + execute_verbs=self.ledger_execute_verbs, + ), benchmark_type, ) @@ -921,6 +926,7 @@ def _setup_fresh_database_phases(self, benchmark, connection: Any, effective_tun getattr(self, "_applied_tuning_ledger", None), PHASE_DDL, statement_filter=is_schema_tuning_statement, + execute_verbs=self.ledger_execute_verbs, ) if schema_records_ddl else connection diff --git a/benchbox/platforms/base/tuning.py b/benchbox/platforms/base/tuning.py index b116d248dd..23b616220d 100644 --- a/benchbox/platforms/base/tuning.py +++ b/benchbox/platforms/base/tuning.py @@ -62,6 +62,7 @@ class TuningHooksMixin: _supported_tuning_type_names: Iterable[str] | None = None physical_identifier_case: str = "preserve" post_load_connection_recording: bool = True + ledger_execute_verbs: tuple[str, ...] = () @staticmethod def table_tuning_for(effective_config: Any, table_name: str) -> Any: @@ -85,7 +86,11 @@ def run_post_load_tunings(self, table_name: str, effective_config: Any, connecti ledger = getattr(self, "_applied_tuning_ledger", None) record = ledger is not None and self.post_load_connection_recording - target = recording_connection(connection, ledger, PHASE_POST_LOAD) if record else connection + target = ( + recording_connection(connection, ledger, PHASE_POST_LOAD, execute_verbs=self.ledger_execute_verbs) + if record + else connection + ) performed = False start = mono_time() try: diff --git a/benchbox/platforms/base/tuning_config.py b/benchbox/platforms/base/tuning_config.py index 3b84d27265..84dce35299 100644 --- a/benchbox/platforms/base/tuning_config.py +++ b/benchbox/platforms/base/tuning_config.py @@ -13,7 +13,12 @@ def apply_standard_unified_tuning(adapter: Any, unified_config: UnifiedTuningCon from benchbox.core.tuning.applied_ledger import PHASE_DDL, recording_connection - recording = recording_connection(connection, getattr(adapter, "_applied_tuning_ledger", None), PHASE_DDL) + recording = recording_connection( + connection, + getattr(adapter, "_applied_tuning_ledger", None), + PHASE_DDL, + execute_verbs=getattr(adapter, "ledger_execute_verbs", ()), + ) adapter.apply_constraint_configuration(unified_config.primary_keys, unified_config.foreign_keys, recording) if unified_config.platform_optimizations: diff --git a/benchbox/platforms/bigquery.py b/benchbox/platforms/bigquery.py index 6daa2d9806..65ae7b4d94 100644 --- a/benchbox/platforms/bigquery.py +++ b/benchbox/platforms/bigquery.py @@ -86,6 +86,7 @@ class BigQueryAdapter(PlatformAdapter): physical_identifier_case = "upper" _DELIMITED_FORMATS = frozenset({"tbl", "csv"}) plan_capture_phase_eligible = False + ledger_execute_verbs = ("query",) def __init__(self, **config): super().__init__(**config) diff --git a/benchbox/platforms/datafusion.py b/benchbox/platforms/datafusion.py index 3c26edd651..cb5d8f1703 100644 --- a/benchbox/platforms/datafusion.py +++ b/benchbox/platforms/datafusion.py @@ -136,6 +136,7 @@ class DataFusionAdapter(NoConstraintEnforcementMixin, PlatformAdapter): driver_isolation_capability = DriverIsolationCapability.SUPPORTED supports_external_tables = True plan_capture_phase_eligible = True + ledger_execute_verbs = ("sql",) _process_working_dir_lock_depth: dict[str, int] = {} _process_working_dir_lock_guard = threading.Lock() diff --git a/tests/integration/platforms/test_spark_ledger_capture_live.py b/tests/integration/platforms/test_spark_ledger_capture_live.py new file mode 100644 index 0000000000..45abd90f23 --- /dev/null +++ b/tests/integration/platforms/test_spark_ledger_capture_live.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +import logging +from types import SimpleNamespace + +import pytest + +from benchbox.core.tuning.applied_ledger import ( + EXECUTED, + NOOP, + PHASE_DDL, + AppliedTuningLedger, + is_schema_tuning_statement, + recording_connection, +) + +from .delta_live_helpers import delta_live_skip_reason, make_delta_spark_session + +pytestmark = [ + pytest.mark.integration, + pytest.mark.live_integration, + pytest.mark.live_delta, + pytest.mark.skipif( + delta_live_skip_reason() is not None, + reason=delta_live_skip_reason() or "PySpark + Delta Lake runtime unavailable", + ), +] + + +@pytest.fixture(scope="module") +def spark(tmp_path_factory, pyspark_test_environment): + session = make_delta_spark_session(tmp_path_factory.mktemp("ledger_capture_warehouse"), app_name="ledger-capture") + session.sql("CREATE DATABASE IF NOT EXISTS ledger_capture") + session.sql("USE ledger_capture") + yield session + session.stop() + + +def test_partitioned_and_clustered_ddl_is_recorded_and_status_is_not_noop(spark) -> None: + from benchbox.platforms.spark import SparkAdapter + + ledger = AppliedTuningLedger() + assert ledger.overall_status(tuning_enabled=True, has_config=True) == NOOP + connection = recording_connection( + spark, + ledger, + PHASE_DDL, + statement_filter=is_schema_tuning_statement, + execute_verbs=SparkAdapter.ledger_execute_verbs, + ) + adapter = SimpleNamespace( + _create_schema_with_tuning=lambda *args, **kwargs: ( + "CREATE TABLE orders (o_id INT, o_date DATE)\nPARTITIONED BY (o_date);\n" + "CREATE TABLE lineitem (l_id INT, l_key INT)\nCLUSTER BY (l_key);\n" + "CREATE TABLE plain (p_id INT);" + ), + table_format="delta", + logger=logging.getLogger(__name__), + _remove_orphaned_table_location=lambda *args, **kwargs: None, + ) + + SparkAdapter.create_schema(adapter, SimpleNamespace(), connection) + + recorded = [(s.phase, s.status, s.statement) for s in ledger.statements] + assert [(phase, status) for phase, status, _ in recorded] == [(PHASE_DDL, EXECUTED)] * 2 + assert "PARTITIONED BY (o_date)" in recorded[0][2] + assert "CLUSTER BY (l_key)" in recorded[1][2] + assert ledger.overall_status(tuning_enabled=True, has_config=True) != NOOP + assert spark.sql("SELECT COUNT(*) FROM plain").collect()[0][0] == 0 + assert {table.name for table in spark.catalog.listTables()} >= {"orders", "lineitem", "plain"} diff --git a/tests/unit/core/tuning/test_ledger_execute_verbs.py b/tests/unit/core/tuning/test_ledger_execute_verbs.py new file mode 100644 index 0000000000..acb39e86ff --- /dev/null +++ b/tests/unit/core/tuning/test_ledger_execute_verbs.py @@ -0,0 +1,378 @@ +from __future__ import annotations + +import logging +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from benchbox.core.tuning.applied_ledger import ( + EXECUTED, + PHASE_DDL, + STATEMENT_FAILED, + AppliedTuningLedger, + is_schema_tuning_statement, + recording_connection, +) +from benchbox.core.tuning.introspection import corroborate + +pytestmark = [ + pytest.mark.unit, + pytest.mark.fast, +] + + +class _SqlDriver: + def __init__(self, fail_on: str | None = None) -> None: + self.calls: list[str] = [] + self.conf = SimpleNamespace(value="kept") + self._fail_on = fail_on + + def sql(self, statement, *args, **kwargs): + if self._fail_on is not None and self._fail_on in str(statement): + raise RuntimeError("driver rejected statement") + self.calls.append(str(statement)) + return f"frame::{statement}" + + def execute(self, statement, *args, **kwargs): + self.calls.append(str(statement)) + return "executed" + + +class _FakeJob: + def __init__(self, error: Exception | None = None) -> None: + self._error = error + self.state = "RUNNING" + self.job_id = "job-1" + + def result(self, *args, **kwargs): + if self._error is not None: + raise self._error + self.state = "DONE" + return "rows" + + def done(self) -> bool: + return self.state == "DONE" + + def exception(self, *args, **kwargs): + return self._error + + +class _QueryClient: + def __init__(self, error: Exception | None = None) -> None: + self.submitted: list[str] = [] + self.jobs: list[_FakeJob] = [] + self._error = error + + def query(self, statement, job_config=None, **kwargs): + self.submitted.append(str(statement)) + job = _FakeJob(self._error) + self.jobs.append(job) + return job + + def dataset(self, dataset_id): + return f"ref::{dataset_id}" + + def get_dataset(self, ref): + return ref + + +def _ddl_proxy(driver, verbs, ledger=None): + ledger = ledger if ledger is not None else AppliedTuningLedger() + return ( + recording_connection( + driver, + ledger, + PHASE_DDL, + statement_filter=is_schema_tuning_statement, + execute_verbs=verbs, + ), + ledger, + ) + + +class TestSqlVerb: + def test_partitioned_and_clustered_ddl_is_recorded_as_ddl(self) -> None: + driver = _SqlDriver() + proxy, ledger = _ddl_proxy(driver, ("sql",)) + + result = proxy.sql("CREATE TABLE t (a INT, b INT) USING delta PARTITIONED BY (a) CLUSTER BY (b)") + + assert result.startswith("frame::") + assert [(s.phase, s.status) for s in ledger.statements] == [(PHASE_DDL, EXECUTED)] + assert "PARTITIONED BY (a)" in ledger.statements[0].statement + + def test_readback_through_sql_executes_but_is_not_recorded(self) -> None: + driver = _SqlDriver() + proxy, ledger = _ddl_proxy(driver, ("sql",)) + + for readback in ("SELECT 1", "SHOW TABLES", "DESCRIBE t", "EXPLAIN SELECT 1"): + proxy.sql(readback) + + assert len(driver.calls) == 4 + assert ledger.is_empty() + + def test_verb_is_not_wrapped_unless_declared(self) -> None: + driver = _SqlDriver() + proxy, ledger = _ddl_proxy(driver, ()) + + proxy.sql("CREATE TABLE t (a INT) PARTITIONED BY (a)") + + assert driver.calls == ["CREATE TABLE t (a INT) PARTITIONED BY (a)"] + assert ledger.is_empty() + + def test_failed_statement_is_recorded_then_reraised(self) -> None: + driver = _SqlDriver(fail_on="PARTITIONED") + proxy, ledger = _ddl_proxy(driver, ("sql",)) + + with pytest.raises(RuntimeError, match="driver rejected"): + proxy.sql("CREATE TABLE t (a INT) PARTITIONED BY (a)") + + assert [s.status for s in ledger.statements] == [STATEMENT_FAILED] + assert "driver rejected" in (ledger.statements[0].error or "") + + def test_non_verb_attributes_delegate_unchanged(self) -> None: + driver = _SqlDriver() + proxy, _ledger = _ddl_proxy(driver, ("sql",)) + + assert proxy.conf.value == "kept" + + def test_capture_failure_does_not_break_the_statement(self) -> None: + class _Unprintable: + def __str__(self) -> str: + raise ValueError("cannot render") + + class _Driver: + def sql(self, statement): + return "ran" + + proxy, ledger = _ddl_proxy(_Driver(), ("sql",)) + + assert proxy.sql(_Unprintable()) == "ran" + assert ledger.is_empty() + + +class TestScriptSplitting: + def test_transient_prefix_does_not_launder_a_tuning_statement(self) -> None: + driver = _SqlDriver() + ledger = AppliedTuningLedger() + proxy = recording_connection(driver, ledger, PHASE_DDL, execute_verbs=("sql",)) + + proxy.sql("SET x=1; CREATE TABLE t (a INT, z INT) ORDER BY (z)") + + assert [s.statement for s in ledger.statements] == [ + "SET x=1;", + "CREATE TABLE t (a INT, z INT) ORDER BY (z)", + ] + receipt = corroborate(ledger, None) + verdicts = {entry.statement: entry.verdict for entry in receipt.entries} + assert verdicts["SET x=1;"] == "transient" + assert verdicts["CREATE TABLE t (a INT, z INT) ORDER BY (z)"] != "transient" + + def test_execute_splits_multi_statement_string_under_the_filter(self) -> None: + driver = _SqlDriver() + proxy, ledger = _ddl_proxy(driver, ()) + + proxy.execute("SET x=1; CREATE TABLE t (a INT, z INT) ORDER BY (z)") + + assert driver.calls == ["SET x=1; CREATE TABLE t (a INT, z INT) ORDER BY (z)"] + assert [s.statement for s in ledger.statements] == ["CREATE TABLE t (a INT, z INT) ORDER BY (z)"] + + def test_readback_prefix_does_not_hide_a_following_mutation(self) -> None: + driver = _SqlDriver() + ledger = AppliedTuningLedger() + proxy = recording_connection(driver, ledger, PHASE_DDL, execute_verbs=("sql",)) + + proxy.sql("SELECT 1; CREATE TABLE t (a INT) PARTITIONED BY (a)") + + assert [s.statement for s in ledger.statements] == ["CREATE TABLE t (a INT) PARTITIONED BY (a)"] + + def test_single_statement_is_recorded_unchanged(self) -> None: + driver = _SqlDriver() + ledger = AppliedTuningLedger() + proxy = recording_connection(driver, ledger, PHASE_DDL, execute_verbs=("sql",)) + statement = " CREATE TABLE t (a INT) \n" + + proxy.sql(statement) + + assert [s.statement for s in ledger.statements] == [statement] + + +class TestQueryJobs: + def test_submission_alone_is_not_recorded(self) -> None: + client = _QueryClient() + proxy, ledger = _ddl_proxy(client, ("query",)) + + proxy.query("CREATE TABLE d.t (a INT64) PARTITION BY DATE(ts) CLUSTER BY a") + + assert client.submitted + assert ledger.is_empty() + + def test_successful_job_is_recorded_when_it_finishes(self) -> None: + client = _QueryClient() + proxy, ledger = _ddl_proxy(client, ("query",)) + + job = proxy.query("CREATE TABLE d.t (a INT64) PARTITION BY DATE(ts) CLUSTER BY a") + assert job.job_id == "job-1" + assert job.result() == "rows" + job.result() + + assert [(s.phase, s.status) for s in ledger.statements] == [(PHASE_DDL, EXECUTED)] + + def test_job_failing_after_submission_is_recorded_failed(self) -> None: + client = _QueryClient(error=RuntimeError("Quota exceeded")) + proxy, ledger = _ddl_proxy(client, ("query",)) + + job = proxy.query("CREATE TABLE d.t (a INT64) PARTITION BY DATE(ts)") + with pytest.raises(RuntimeError, match="Quota exceeded"): + job.result() + + assert [s.status for s in ledger.statements] == [STATEMENT_FAILED] + assert "Quota exceeded" in (ledger.statements[0].error or "") + + def test_exception_accessor_settles_the_final_state(self) -> None: + client = _QueryClient(error=RuntimeError("Not found")) + proxy, ledger = _ddl_proxy(client, ("query",)) + + job = proxy.query("CREATE TABLE d.t (a INT64) CLUSTER BY a") + assert str(job.exception()) == "Not found" + + assert [s.status for s in ledger.statements] == [STATEMENT_FAILED] + + def test_timeout_while_waiting_is_not_a_final_state(self) -> None: + from concurrent.futures import TimeoutError as FutureTimeoutError + + client = _QueryClient(error=FutureTimeoutError("still running")) + proxy, ledger = _ddl_proxy(client, ("query",)) + + job = proxy.query("CREATE TABLE d.t (a INT64) CLUSTER BY a") + with pytest.raises(FutureTimeoutError): + job.result(timeout=1) + + assert ledger.is_empty() + + def test_dry_run_and_readback_jobs_are_not_recorded(self) -> None: + client = _QueryClient() + proxy, ledger = _ddl_proxy(client, ("query",)) + + dry = proxy.query( + "CREATE TABLE d.t (a INT64) CLUSTER BY a", + job_config=SimpleNamespace(dry_run=True), + ) + select = proxy.query("SELECT * FROM d.t") + dry.result() + select.result() + + assert ledger.is_empty() + + def test_submission_error_is_recorded_failed(self) -> None: + class _Rejecting(_QueryClient): + def query(self, statement, job_config=None, **kwargs): + raise RuntimeError("bad request") + + proxy, ledger = _ddl_proxy(_Rejecting(), ("query",)) + + with pytest.raises(RuntimeError, match="bad request"): + proxy.query("CREATE TABLE d.t (a INT64) CLUSTER BY a") + + assert [s.status for s in ledger.statements] == [STATEMENT_FAILED] + + +class TestFootprint: + @pytest.mark.parametrize( + "statement", + [ + "CREATE TABLE t (a INT) USING parquet PARTITIONED BY (a)", + "CREATE TABLE t (a INT) USING parquet CLUSTERED BY (a) INTO 4 BUCKETS", + "CREATE TABLE t (a INT) DISTKEY(a)", + "CREATE TABLE t (a INT) DISTSTYLE KEY", + "CREATE TABLE t (a INT) DUPLICATE KEY (a) DISTRIBUTED BY HASH (a)", + "CREATE TABLE t (a INT) PRIMARY INDEX (a)", + "CREATE TABLE t (a INT) WITH (distribution = 'hash')", + "CREATE TABLE t (a INT) WITH (partitioning = ARRAY['a'])", + "CREATE TABLE t (a INT) WITH (sorted_by = ARRAY['a'])", + "CREATE TABLE t (a INT) WITH (bucketed_by = ARRAY['a'], bucket_count = 4)", + ], + ) + def test_new_footprint_keywords_are_schema_tuning(self, statement: str) -> None: + assert is_schema_tuning_statement(statement) + + def test_plain_table_is_still_not_schema_tuning(self) -> None: + assert not is_schema_tuning_statement("CREATE TABLE t (a INT, b INT)") + + def test_keyword_inside_a_literal_does_not_count(self) -> None: + assert not is_schema_tuning_statement("CREATE TABLE t (a INT) COMMENT 'partitioned by day'") + + +class TestAdapterCreateSchema: + def test_spark_family_declares_sql_verb(self) -> None: + from benchbox.platforms._spark_helpers import SparkLikeAdapterMixin + + assert SparkLikeAdapterMixin.ledger_execute_verbs == ("sql",) + + def test_declared_verbs_on_non_dbapi_adapters(self) -> None: + from benchbox.platforms.bigquery import BigQueryAdapter + from benchbox.platforms.datafusion import DataFusionAdapter + + assert BigQueryAdapter.ledger_execute_verbs == ("query",) + assert DataFusionAdapter.ledger_execute_verbs == ("sql",) + + def test_spark_tuned_schema_ddl_reaches_the_ledger(self) -> None: + from benchbox.platforms.spark import SparkAdapter + + driver = _SqlDriver() + proxy, ledger = _ddl_proxy(driver, SparkAdapter.ledger_execute_verbs) + fake_self = SimpleNamespace( + _create_schema_with_tuning=lambda *a, **k: ( + "CREATE TABLE orders (o_id INT, o_date DATE) USING parquet PARTITIONED BY (o_date);\n" + "CREATE TABLE lineitem (l_id INT, l_key INT) USING delta CLUSTER BY (l_key);\n" + "CREATE TABLE plain (p_id INT);" + ), + table_format="parquet", + logger=logging.getLogger("test"), + _remove_orphaned_table_location=lambda *a, **k: None, + ) + + SparkAdapter.create_schema(fake_self, SimpleNamespace(), proxy) + + assert len(driver.calls) == 3 + recorded = [s.statement for s in ledger.statements] + assert len(recorded) == 2 + assert any("PARTITIONED BY" in s for s in recorded) + assert any("CLUSTER BY" in s for s in recorded) + assert {s.phase for s in ledger.statements} == {PHASE_DDL} + assert ledger.overall_status(tuning_enabled=True, has_config=True) == "applied_unverified" + + def test_bigquery_tuned_schema_ddl_records_final_job_state(self) -> None: + from benchbox.platforms.bigquery import BigQueryAdapter + + client = _QueryClient() + proxy, ledger = _ddl_proxy(client, BigQueryAdapter.ledger_execute_verbs) + fake_self = SimpleNamespace( + dataset_id="ds", + location="US", + _create_schema_with_tuning=lambda *a, **k: ( + "CREATE TABLE t1 (a INT64, ts TIMESTAMP) PARTITION BY DATE(ts) CLUSTER BY a;\nCREATE TABLE t2 (b INT64);" + ), + _convert_to_bigquery_table=lambda statement: statement, + logger=logging.getLogger("test"), + ) + + with patch("benchbox.platforms.bigquery.bigquery"): + BigQueryAdapter.create_schema(fake_self, SimpleNamespace(), proxy) + + assert len(client.submitted) == 2 + assert [s.status for s in ledger.statements] == [EXECUTED] + assert "CLUSTER BY a" in ledger.statements[0].statement + + def test_datafusion_tuned_ddl_through_sql_is_recorded(self) -> None: + from benchbox.platforms.datafusion import DataFusionAdapter + + driver = _SqlDriver() + proxy, ledger = _ddl_proxy(driver, DataFusionAdapter.ledger_execute_verbs) + + proxy.sql("CREATE TABLE t (a INT) PARTITIONED BY (a)") + proxy.sql("SELECT COUNT(*) FROM t") + + assert len(driver.calls) == 2 + assert [(s.phase, s.status) for s in ledger.statements] == [(PHASE_DDL, EXECUTED)] diff --git a/tests/unit/platforms/test_spark_helpers.py b/tests/unit/platforms/test_spark_helpers.py index 8301eaa047..22ec36d873 100644 --- a/tests/unit/platforms/test_spark_helpers.py +++ b/tests/unit/platforms/test_spark_helpers.py @@ -150,6 +150,14 @@ def test_strips_existing_using_clause(self) -> None: assert "DELTA" not in result.upper() assert result.endswith("USING PARQUET") + @pytest.mark.parametrize("clause", ["PARTITIONED BY (id)", "CLUSTER BY (id)"]) + def test_using_precedes_trailing_tuning_clause(self, clause: str) -> None: + result = optimize_spark_table_definition( + f"CREATE TABLE t (id INT, amount DECIMAL(10,2))\n{clause};", + table_format="delta", + ) + assert result == f"CREATE TABLE t (id INT, amount DECIMAL(10,2)) USING DELTA\n{clause}" + def test_non_create_table_unchanged(self) -> None: sql = "INSERT INTO t VALUES (1)" assert optimize_spark_table_definition(sql, table_format="parquet") == sql From 339c1ad6efd363dcc48c9dfd0199a21ad6d71b90 Mon Sep 17 00:00:00 2001 From: Joe Harris Date: Fri, 9 Oct 2026 18:42:06 -0400 Subject: [PATCH 2/5] fix(tuning): accept keyword SQL arguments and quoted parentheses in tuning capture The recording proxy now forwards the declared SQL verb with the caller's own positional or keyword form (sql(sqlQuery=...), query(query=...)) and extracts the statement from whichever was used. Spark table definitions place USING after the real closing delimiter of the column list, ignoring parentheses inside string literals and quoted identifiers, so defaults are no longer altered. Tests that only froze verb declarations and attribute delegation are removed, along with assertions on mock return values. --- benchbox/core/tuning/applied_ledger.py | 32 +++++++++++++--- benchbox/platforms/_spark_helpers.py | 20 +++++++--- .../core/tuning/test_ledger_execute_verbs.py | 38 ++++++------------- tests/unit/platforms/test_spark_helpers.py | 20 ++++++++++ 4 files changed, 74 insertions(+), 36 deletions(-) diff --git a/benchbox/core/tuning/applied_ledger.py b/benchbox/core/tuning/applied_ledger.py index 6b0fd532ea..a368879b01 100644 --- a/benchbox/core/tuning/applied_ledger.py +++ b/benchbox/core/tuning/applied_ledger.py @@ -349,10 +349,22 @@ def _run(self, fn: Any, statement: Any, args: tuple, kwargs: dict) -> Any: self._record_all(statements, EXECUTED) return result - def _run_verb(self, fn: Any, statement: Any, args: tuple, kwargs: dict) -> Any: - statements = [] if _is_dry_run(args, kwargs) else self._recordable_statements(statement) + def _run_verb(self, fn: Any, args: tuple, kwargs: dict) -> Any: + statements: list[str] = [] try: - result = fn(statement, *args, **kwargs) + name = _first_parameter_name(fn) + if args: + statement, rest, rest_kwargs = args[0], args[1:], kwargs + elif name in kwargs: + statement, rest, rest_kwargs = kwargs[name], args, {k: v for k, v in kwargs.items() if k != name} + else: + raise LookupError(name) + if not _is_dry_run(rest, rest_kwargs): + statements = self._recordable_statements(statement) + except Exception as exc: + logger.debug("applied-ledger statement capture degraded: %s", exc) + try: + result = fn(*args, **kwargs) except Exception as exc: self._record_all(statements, STATEMENT_FAILED, exc) raise @@ -364,6 +376,16 @@ def _run_verb(self, fn: Any, statement: Any, args: tuple, kwargs: dict) -> Any: return result +def _first_parameter_name(fn: Any) -> str | None: + import inspect + + for parameter in inspect.signature(fn).parameters.values(): + if parameter.kind in (parameter.POSITIONAL_OR_KEYWORD, parameter.KEYWORD_ONLY): + return parameter.name + break + return None + + def _is_dry_run(args: tuple, kwargs: dict) -> bool: config = kwargs.get("job_config") or (args[0] if args else None) return getattr(config, "dry_run", False) is True @@ -443,8 +465,8 @@ def __getattr__(self, name: str) -> Any: attribute = getattr(self._conn, name) if name in self._verbs and callable(attribute): - def call(statement: Any, *args: Any, **kwargs: Any) -> Any: - return self._run_verb(attribute, statement, args, kwargs) + def call(*args: Any, **kwargs: Any) -> Any: + return self._run_verb(attribute, args, kwargs) return call return attribute diff --git a/benchbox/platforms/_spark_helpers.py b/benchbox/platforms/_spark_helpers.py index 5f6e4c6f01..10e72597a0 100644 --- a/benchbox/platforms/_spark_helpers.py +++ b/benchbox/platforms/_spark_helpers.py @@ -158,7 +158,7 @@ def get_spark_query_plan( _FIXED_SIZE_ARRAY_TYPE_RE = re.compile(r"\b([A-Za-z]\w*)\s*\[\s*\d+\s*\]") _ARRAY_FIXED_SIZE_SUFFIX_RE = re.compile(r"(>)\s*\[\s*\d+\s*\]") _CREATE_TABLE_COLUMN_LIST_RE = re.compile( - r"\A\s*CREATE\s+TABLE\s+(?:IF\s+NOT\s+EXISTS\s+)?[^\s(]+\s*\(", + r"\A\s*CREATE\s+TABLE\s+(?:IF\s+NOT\s+EXISTS\s+)?(?:\"[^\"]*\"|`[^`]*`|[^\s(\"`])+\s*\(", re.IGNORECASE, ) @@ -168,15 +168,25 @@ def _column_list_end(statement: str) -> int: if not match: return len(statement) depth = 1 + quote = "" cursor = match.end() - while cursor < len(statement) and depth: + while cursor < len(statement): ch = statement[cursor] - if ch == "(": + cursor += 1 + if quote: + if ch == "\\" and quote == "'": + cursor += 1 + elif ch == quote: + quote = "" + elif ch in "'\"`": + quote = ch + elif ch == "(": depth += 1 elif ch == ")": depth -= 1 - cursor += 1 - return cursor if depth == 0 else len(statement) + if depth == 0: + return cursor + return len(statement) def _strip_balanced_paren_constraints(statement: str) -> str: diff --git a/tests/unit/core/tuning/test_ledger_execute_verbs.py b/tests/unit/core/tuning/test_ledger_execute_verbs.py index acb39e86ff..4277126c1e 100644 --- a/tests/unit/core/tuning/test_ledger_execute_verbs.py +++ b/tests/unit/core/tuning/test_ledger_execute_verbs.py @@ -25,14 +25,13 @@ class _SqlDriver: def __init__(self, fail_on: str | None = None) -> None: self.calls: list[str] = [] - self.conf = SimpleNamespace(value="kept") self._fail_on = fail_on def sql(self, statement, *args, **kwargs): if self._fail_on is not None and self._fail_on in str(statement): raise RuntimeError("driver rejected statement") self.calls.append(str(statement)) - return f"frame::{statement}" + return None def execute(self, statement, *args, **kwargs): self.calls.append(str(statement)) @@ -96,9 +95,8 @@ def test_partitioned_and_clustered_ddl_is_recorded_as_ddl(self) -> None: driver = _SqlDriver() proxy, ledger = _ddl_proxy(driver, ("sql",)) - result = proxy.sql("CREATE TABLE t (a INT, b INT) USING delta PARTITIONED BY (a) CLUSTER BY (b)") + proxy.sql("CREATE TABLE t (a INT, b INT) USING delta PARTITIONED BY (a) CLUSTER BY (b)") - assert result.startswith("frame::") assert [(s.phase, s.status) for s in ledger.statements] == [(PHASE_DDL, EXECUTED)] assert "PARTITIONED BY (a)" in ledger.statements[0].statement @@ -131,24 +129,24 @@ def test_failed_statement_is_recorded_then_reraised(self) -> None: assert [s.status for s in ledger.statements] == [STATEMENT_FAILED] assert "driver rejected" in (ledger.statements[0].error or "") - def test_non_verb_attributes_delegate_unchanged(self) -> None: - driver = _SqlDriver() - proxy, _ledger = _ddl_proxy(driver, ("sql",)) - - assert proxy.conf.value == "kept" - def test_capture_failure_does_not_break_the_statement(self) -> None: class _Unprintable: def __str__(self) -> str: raise ValueError("cannot render") class _Driver: + def __init__(self) -> None: + self.ran = False + def sql(self, statement): - return "ran" + self.ran = True - proxy, ledger = _ddl_proxy(_Driver(), ("sql",)) + driver = _Driver() + proxy, ledger = _ddl_proxy(driver, ("sql",)) + + proxy.sql(_Unprintable()) - assert proxy.sql(_Unprintable()) == "ran" + assert driver.ran assert ledger.is_empty() @@ -214,7 +212,7 @@ def test_successful_job_is_recorded_when_it_finishes(self) -> None: job = proxy.query("CREATE TABLE d.t (a INT64) PARTITION BY DATE(ts) CLUSTER BY a") assert job.job_id == "job-1" - assert job.result() == "rows" + job.result() job.result() assert [(s.phase, s.status) for s in ledger.statements] == [(PHASE_DDL, EXECUTED)] @@ -305,18 +303,6 @@ def test_keyword_inside_a_literal_does_not_count(self) -> None: class TestAdapterCreateSchema: - def test_spark_family_declares_sql_verb(self) -> None: - from benchbox.platforms._spark_helpers import SparkLikeAdapterMixin - - assert SparkLikeAdapterMixin.ledger_execute_verbs == ("sql",) - - def test_declared_verbs_on_non_dbapi_adapters(self) -> None: - from benchbox.platforms.bigquery import BigQueryAdapter - from benchbox.platforms.datafusion import DataFusionAdapter - - assert BigQueryAdapter.ledger_execute_verbs == ("query",) - assert DataFusionAdapter.ledger_execute_verbs == ("sql",) - def test_spark_tuned_schema_ddl_reaches_the_ledger(self) -> None: from benchbox.platforms.spark import SparkAdapter diff --git a/tests/unit/platforms/test_spark_helpers.py b/tests/unit/platforms/test_spark_helpers.py index 22ec36d873..f428e10cd6 100644 --- a/tests/unit/platforms/test_spark_helpers.py +++ b/tests/unit/platforms/test_spark_helpers.py @@ -158,6 +158,26 @@ def test_using_precedes_trailing_tuning_clause(self, clause: str) -> None: ) assert result == f"CREATE TABLE t (id INT, amount DECIMAL(10,2)) USING DELTA\n{clause}" + @pytest.mark.parametrize( + ("columns", "name"), + [ + ("id INT, note STRING DEFAULT ' ) '", "t"), + ("id INT, note STRING DEFAULT 'it''s ) here'", "t"), + ("id INT, note STRING DEFAULT 'a\\' ) b'", "t"), + ("id INT, `we ) ird` STRING", "t"), + ("id INT, amount DECIMAL(10,2), tags MAP>", "`my table`"), + ], + ) + def test_using_stays_outside_column_list_with_quoted_parentheses(self, columns: str, name: str) -> None: + result = optimize_spark_table_definition( + f"CREATE TABLE {name} ({columns}) PARTITIONED BY (id);", + table_format="delta", + ) + assert f"({columns})" in result + assert result.count("USING DELTA") == 1 + assert result.index(f"({columns})") + len(f"({columns})") < result.index("USING DELTA") + assert result.index("USING DELTA") < result.index("PARTITIONED BY") + def test_non_create_table_unchanged(self) -> None: sql = "INSERT INTO t VALUES (1)" assert optimize_spark_table_definition(sql, table_format="parquet") == sql From 2ab94155fb9f776e1f0e1db44716bbeb05b1bf78 Mon Sep 17 00:00:00 2001 From: Joe Harris Date: Fri, 9 Oct 2026 23:00:55 -0400 Subject: [PATCH 3/5] fix(tuning): close ledger capture gaps in footprint, CTAS sort, and script splitting Require bare schema-tuning keywords to sit in clause position so ordinary column names no longer record executed tuning, while every footprint keyword still records in its real clause form. Record CTAS sorts where they execute by wrapping the connection at the sort site, so the DataFusion load-time sort through execute reaches the ledger instead of staying silent. Skip line and block comments when locating the Spark column-list boundary, in addition to quoted text, so USING lands after the real closing delimiter. Exclude comment-only script fragments from recording on every proxy path, so trailing comments no longer corrupt tuning verdicts. --- benchbox/core/tuning/applied_ledger.py | 26 ++++-- benchbox/platforms/_spark_helpers.py | 6 ++ benchbox/platforms/base/sorted_ingestion.py | 29 ++++++- .../core/tuning/test_ledger_execute_verbs.py | 87 +++++++++++++++++++ .../test_local_platform_ctas_sort.py | 40 +++++++++ tests/unit/platforms/test_spark_helpers.py | 19 ++++ 6 files changed, 198 insertions(+), 9 deletions(-) diff --git a/benchbox/core/tuning/applied_ledger.py b/benchbox/core/tuning/applied_ledger.py index a368879b01..c9074e26ec 100644 --- a/benchbox/core/tuning/applied_ledger.py +++ b/benchbox/core/tuning/applied_ledger.py @@ -66,12 +66,22 @@ r"^\s*(?:create\s+index|create\s+(?:or\s+replace\s+)?table|alter\s+table)\b", re.IGNORECASE, ) -_SCHEMA_TUNING_FOOTPRINT_RE = re.compile( +_SCHEMA_TUNING_PHRASE_RE = re.compile( r"\b(?:cluster\s+by|clustered\s+by|distributed\s+by|foreign\s+key|order\s+by|partition\s+by|" - r"partitioned\s+by|primary\s+key|primary\s+index|duplicate\s+key|sortkey|distkey|diststyle|" - r"sorted_by|bucketed_by|unique|check)\b|\b(?:distribution|partitioning)\s*=", + r"partitioned\s+by|primary\s+key|primary\s+index|duplicate\s+key)\b" + r"|\b(?:distribution|partitioning)\s*=", re.IGNORECASE, ) +_SCHEMA_TUNING_KEYWORD_RE = re.compile( + r"\b(?:sortkey|distkey|diststyle|sorted_by|bucketed_by|unique|check)\b", + re.IGNORECASE, +) + + +def _bare_keyword_is_clause(shape: str, match: re.Match[str]) -> bool: + if shape[match.end() :].lstrip()[:1] in ("(", "="): + return True + return not shape[: match.start()].rstrip().endswith(("(", ",")) def _sql_shape(statement: Any) -> str: @@ -86,7 +96,9 @@ def is_schema_tuning_statement(statement: Any) -> bool: return False if re.match(r"^\s*create\s+index\b", shape, re.IGNORECASE): return True - return _SCHEMA_TUNING_FOOTPRINT_RE.search(shape) is not None + if _SCHEMA_TUNING_PHRASE_RE.search(shape) is not None: + return True + return any(_bare_keyword_is_clause(shape, match) for match in _SCHEMA_TUNING_KEYWORD_RE.finditer(shape)) def _split_sql_script(script: Any) -> list[str]: @@ -130,10 +142,12 @@ def normalize_ledger_phase(phase: Any) -> str: def _is_recordable_statement(statement: Any) -> bool: try: - text = str(statement).lstrip().lstrip("(").lstrip().lower() + text = _SQL_COMMENT_RE.sub(" ", str(statement)) except Exception: # pragma: no cover return True - return not text.startswith(_READBACK_PREFIXES) + if not text.strip(): + return False + return not text.lstrip().lstrip("(").lstrip().lower().startswith(_READBACK_PREFIXES) @dataclass diff --git a/benchbox/platforms/_spark_helpers.py b/benchbox/platforms/_spark_helpers.py index 10e72597a0..456136bcc9 100644 --- a/benchbox/platforms/_spark_helpers.py +++ b/benchbox/platforms/_spark_helpers.py @@ -180,6 +180,12 @@ def _column_list_end(statement: str) -> int: quote = "" elif ch in "'\"`": quote = ch + elif ch == "-" and statement[cursor : cursor + 1] == "-": + newline = statement.find("\n", cursor + 1) + cursor = len(statement) if newline == -1 else newline + 1 + elif ch == "/" and statement[cursor : cursor + 1] == "*": + close = statement.find("*/", cursor + 1) + cursor = len(statement) if close == -1 else close + 2 elif ch == "(": depth += 1 elif ch == ")": diff --git a/benchbox/platforms/base/sorted_ingestion.py b/benchbox/platforms/base/sorted_ingestion.py index 606768c8cf..d2eaaf9d56 100644 --- a/benchbox/platforms/base/sorted_ingestion.py +++ b/benchbox/platforms/base/sorted_ingestion.py @@ -184,6 +184,28 @@ def _resolve_ctas_sort_columns(self, table_name: str, tuning_config: Any) -> lis return sort_columns + def _recording_ctas_connection(self, connection: Any) -> Any: + from benchbox.core.tuning.applied_ledger import ( + PHASE_POST_LOAD, + RecordingConnection, + is_schema_tuning_statement, + recording_connection, + ) + + if getattr(self, "_applied_tuning_ledger", None) is None: + return connection + if isinstance(connection, RecordingConnection): + return connection + if not callable(getattr(connection, "execute", None)) and not callable(getattr(connection, "sql", None)): + return connection + return recording_connection( + connection, + self._applied_tuning_ledger, + PHASE_POST_LOAD, + statement_filter=is_schema_tuning_statement, + execute_verbs=("sql",), + ) + def _execute_ctas_sort( self, table_name: str, @@ -206,6 +228,7 @@ def _execute_ctas_sort( return True ledger_statement = "\n".join(statements) + target = self._recording_ctas_connection(connection) transaction_started = False try: if atomic: @@ -214,7 +237,7 @@ def _execute_ctas_sort( self._record_sorted_ingestion_skip(validated_table, sorted_columns, reason) self.logger.info("Skipping constraint-preserving CTAS sort for %s: %s", table_name, reason) return False - self._execute_sql_statement(connection, "BEGIN TRANSACTION;") + self._execute_sql_statement(target, "BEGIN TRANSACTION;") transaction_started = True if atomic and self._has_populated_fk_dependents(connection, validated_table): self._rollback_transaction(connection) @@ -228,13 +251,13 @@ def _execute_ctas_sort( ) return False for statement in statements: - self._execute_sql_statement(connection, statement) + self._execute_sql_statement(target, statement) if atomic: commit_method = getattr(connection, "commit", None) if callable(commit_method): commit_method() else: - self._execute_sql_statement(connection, "COMMIT;") + self._execute_sql_statement(target, "COMMIT;") transaction_started = False except Exception as exc: rollback_error: Exception | None = None diff --git a/tests/unit/core/tuning/test_ledger_execute_verbs.py b/tests/unit/core/tuning/test_ledger_execute_verbs.py index 4277126c1e..f00b5fd61e 100644 --- a/tests/unit/core/tuning/test_ledger_execute_verbs.py +++ b/tests/unit/core/tuning/test_ledger_execute_verbs.py @@ -196,6 +196,59 @@ def test_single_statement_is_recorded_unchanged(self) -> None: assert [s.statement for s in ledger.statements] == [statement] +_TRAILING_SQL_COMMENT = "-- trailing comment" +_READBACK_SQL_COMMENT = "-- readback with a trailing comment" + + +class TestCommentOnlyFragments: + def test_trailing_comment_after_execute_is_not_recorded(self) -> None: + driver = _SqlDriver() + ledger = AppliedTuningLedger() + proxy = recording_connection(driver, ledger, PHASE_DDL) + + proxy.execute(f"CREATE INDEX idx ON t(x); {_TRAILING_SQL_COMMENT}") + + assert driver.calls == [f"CREATE INDEX idx ON t(x); {_TRAILING_SQL_COMMENT}"] + assert [s.statement for s in ledger.statements] == ["CREATE INDEX idx ON t(x);"] + + def test_trailing_comment_after_verb_call_is_not_recorded(self) -> None: + driver = _SqlDriver() + ledger = AppliedTuningLedger() + proxy = recording_connection(driver, ledger, PHASE_DDL, execute_verbs=("sql",)) + + proxy.sql(f"CREATE INDEX idx ON t(x); {_TRAILING_SQL_COMMENT}") + + assert [s.statement for s in ledger.statements] == ["CREATE INDEX idx ON t(x);"] + receipt = corroborate(ledger, None) + assert [entry.statement for entry in receipt.entries] == ["CREATE INDEX idx ON t(x);"] + + def test_trailing_comment_after_executescript_is_not_recorded(self) -> None: + class _ScriptDriver: + def __init__(self) -> None: + self.scripts: list[str] = [] + + def executescript(self, script, *args, **kwargs): + self.scripts.append(str(script)) + + driver = _ScriptDriver() + ledger = AppliedTuningLedger() + proxy = recording_connection(driver, ledger, PHASE_DDL) + + proxy.executescript(f"CREATE INDEX idx ON t(x); {_TRAILING_SQL_COMMENT}") + + assert [s.statement for s in ledger.statements] == ["CREATE INDEX idx ON t(x);"] + + def test_readback_script_with_comment_tail_records_nothing(self) -> None: + driver = _SqlDriver() + ledger = AppliedTuningLedger() + proxy = recording_connection(driver, ledger, PHASE_DDL) + + proxy.execute(f"SELECT 1; {_READBACK_SQL_COMMENT}") + + assert driver.calls == [f"SELECT 1; {_READBACK_SQL_COMMENT}"] + assert ledger.is_empty() + + class TestQueryJobs: def test_submission_alone_is_not_recorded(self) -> None: client = _QueryClient() @@ -301,6 +354,40 @@ def test_plain_table_is_still_not_schema_tuning(self) -> None: def test_keyword_inside_a_literal_does_not_count(self) -> None: assert not is_schema_tuning_statement("CREATE TABLE t (a INT) COMMENT 'partitioned by day'") + @pytest.mark.parametrize( + "statement", + [ + "CREATE TABLE t (distkey INT)", + "CREATE TABLE t (a INT, distkey TEXT)", + "CREATE TABLE t (diststyle VARCHAR(16))", + "CREATE TABLE t (sortkey INT)", + "CREATE TABLE t (a INT, sorted_by DOUBLE)", + "CREATE TABLE t (bucketed_by INT)", + "CREATE TABLE t (unique INT)", + "CREATE TABLE t (check INT)", + ], + ) + def test_bare_keyword_as_a_column_name_is_not_schema_tuning(self, statement: str) -> None: + assert not is_schema_tuning_statement(statement) + + @pytest.mark.parametrize( + "statement", + [ + "CREATE TABLE t (a INT) PARTITIONED BY (a)", + "CREATE TABLE t (a INT) CLUSTERED BY (a) INTO 4 BUCKETS", + "CREATE TABLE t (a INT) DISTSTYLE EVEN DISTKEY (a)", + "CREATE TABLE t (a INT) DUPLICATE KEY (a)", + "CREATE TABLE t (a INT) PRIMARY INDEX (a)", + "CREATE TABLE t (a INT) WITH (distribution = 'hash')", + "CREATE TABLE t (a INT) WITH (partitioning = ARRAY['a'])", + "CREATE TABLE t (a INT) WITH (sorted_by = ARRAY['a'])", + "CREATE TABLE t (a INT) WITH (bucketed_by = ARRAY['a'])", + "CREATE TABLE t (a INT UNIQUE)", + ], + ) + def test_bare_keyword_in_clause_syntax_is_schema_tuning(self, statement: str) -> None: + assert is_schema_tuning_statement(statement) + class TestAdapterCreateSchema: def test_spark_tuned_schema_ddl_reaches_the_ledger(self) -> None: diff --git a/tests/unit/platforms/test_local_platform_ctas_sort.py b/tests/unit/platforms/test_local_platform_ctas_sort.py index 01f7c8e956..133fedde64 100644 --- a/tests/unit/platforms/test_local_platform_ctas_sort.py +++ b/tests/unit/platforms/test_local_platform_ctas_sort.py @@ -93,3 +93,43 @@ def test_datafusion_apply_ctas_sort_orders_rows(self, tmp_path): rows = connection.sql("SELECT l_shipdate, l_orderkey FROM lineitem").to_pydict() assert rows["l_shipdate"] == ["2024-01-01", "2024-02-01", "2024-03-01"] assert rows["l_orderkey"] == [1, 2, 3] + + def test_datafusion_real_ctas_sort_is_recorded_in_ledger(self, tmp_path): + from benchbox.core.tuning.applied_ledger import AppliedTuningLedger + from benchbox.core.tuning.interface import TuningColumn, TuningType + + adapter = DataFusionAdapter(working_dir=str(tmp_path / "df_workdir"), data_format="csv") + adapter._applied_tuning_ledger = AppliedTuningLedger() + connection = adapter.create_connection() + + connection.execute( + "CREATE TABLE lineitem AS " + "SELECT * FROM (VALUES (3, '2024-03-01'), (1, '2024-01-01'), (2, '2024-02-01')) " + "AS t(l_orderkey, l_shipdate)" + ) + + sort_columns = [ + TuningColumn(name="l_shipdate", type="VARCHAR", order=1), + TuningColumn(name="l_orderkey", type="INTEGER", order=2), + ] + table_tuning = Mock() + table_tuning.get_columns_by_type.side_effect = lambda tuning_type: ( + sort_columns if tuning_type == TuningType.SORTING else [] + ) + tuning_config = Mock() + tuning_config.table_tunings = {"lineitem": table_tuning} + + ledger = adapter._applied_tuning_ledger + assert ledger.overall_status(tuning_enabled=True, has_config=True) == "noop" + + assert adapter.apply_ctas_sort("lineitem", tuning_config, connection) is True + + executed = [s for s in ledger.statements if s.status == "executed"] + assert len(executed) == 1 + assert executed[0].phase == "post_load" + assert "ORDER BY l_shipdate, l_orderkey" in executed[0].statement + assert ledger.overall_status(tuning_enabled=True, has_config=True) == "applied_unverified" + + rows = connection.sql("SELECT l_shipdate, l_orderkey FROM lineitem").to_pydict() + assert rows["l_shipdate"] == ["2024-01-01", "2024-02-01", "2024-03-01"] + assert rows["l_orderkey"] == [1, 2, 3] diff --git a/tests/unit/platforms/test_spark_helpers.py b/tests/unit/platforms/test_spark_helpers.py index f428e10cd6..761f1f4b75 100644 --- a/tests/unit/platforms/test_spark_helpers.py +++ b/tests/unit/platforms/test_spark_helpers.py @@ -178,6 +178,25 @@ def test_using_stays_outside_column_list_with_quoted_parentheses(self, columns: assert result.index(f"({columns})") + len(f"({columns})") < result.index("USING DELTA") assert result.index("USING DELTA") < result.index("PARTITIONED BY") + @pytest.mark.parametrize( + "columns", + [ + "id INT /* ) */, amount INT", + "id INT /* ( unbalanced */, amount INT", + "id INT -- )\n, amount INT", + "id INT -- ( unbalanced\n, amount INT", + ], + ) + def test_using_stays_outside_column_list_with_comment_parentheses(self, columns: str) -> None: + result = optimize_spark_table_definition( + f"CREATE TABLE t ({columns}) PARTITIONED BY (id);", + table_format="delta", + ) + assert f"({columns})" in result + assert result.count("USING DELTA") == 1 + assert result.index(f"({columns})") + len(f"({columns})") < result.index("USING DELTA") + assert result.index("USING DELTA") < result.index("PARTITIONED BY") + def test_non_create_table_unchanged(self) -> None: sql = "INSERT INTO t VALUES (1)" assert optimize_spark_table_definition(sql, table_format="parquet") == sql From 593df7dfb90fd8f811cffab34d93b7c909629721 Mon Sep 17 00:00:00 2001 From: Joe Harris Date: Sat, 10 Oct 2026 03:15:59 -0400 Subject: [PATCH 4/5] fix(tuning): record UNIQUE NULLS DISTINCT constraints and drop delimiter-only fragments Recognize UNIQUE NULLS DISTINCT and NULLS NOT DISTINCT table constraints as clause syntax so executed uniqueness tuning reaches the ledger instead of being mistaken for a column name, while bare UNIQUE as a column name stays unrecorded. Treat fragments containing only statement delimiters, comments and whitespace as non-recordable on every proxy path, so empty split tails no longer corrupt tuning verdicts or the ledger hash. --- benchbox/core/tuning/applied_ledger.py | 11 +++-- .../core/tuning/test_ledger_execute_verbs.py | 46 +++++++++++++++++++ 2 files changed, 54 insertions(+), 3 deletions(-) diff --git a/benchbox/core/tuning/applied_ledger.py b/benchbox/core/tuning/applied_ledger.py index c9074e26ec..98275d8317 100644 --- a/benchbox/core/tuning/applied_ledger.py +++ b/benchbox/core/tuning/applied_ledger.py @@ -76,10 +76,14 @@ r"\b(?:sortkey|distkey|diststyle|sorted_by|bucketed_by|unique|check)\b", re.IGNORECASE, ) +_UNIQUE_NULLS_DISTINCT_RE = re.compile(r"\s*nulls\s+(?:not\s+distinct|distinct)\s*\(", re.IGNORECASE) def _bare_keyword_is_clause(shape: str, match: re.Match[str]) -> bool: - if shape[match.end() :].lstrip()[:1] in ("(", "="): + rest = shape[match.end() :] + if rest.lstrip()[:1] in ("(", "="): + return True + if match.group(0).lower() == "unique" and _UNIQUE_NULLS_DISTINCT_RE.match(rest) is not None: return True return not shape[: match.start()].rstrip().endswith(("(", ",")) @@ -145,9 +149,10 @@ def _is_recordable_statement(statement: Any) -> bool: text = _SQL_COMMENT_RE.sub(" ", str(statement)) except Exception: # pragma: no cover return True - if not text.strip(): + stripped = text.strip().strip(";").strip() + if not stripped: return False - return not text.lstrip().lstrip("(").lstrip().lower().startswith(_READBACK_PREFIXES) + return not stripped.lstrip("(").lstrip().lower().startswith(_READBACK_PREFIXES) @dataclass diff --git a/tests/unit/core/tuning/test_ledger_execute_verbs.py b/tests/unit/core/tuning/test_ledger_execute_verbs.py index f00b5fd61e..65e6369e8b 100644 --- a/tests/unit/core/tuning/test_ledger_execute_verbs.py +++ b/tests/unit/core/tuning/test_ledger_execute_verbs.py @@ -248,6 +248,50 @@ def test_readback_script_with_comment_tail_records_nothing(self) -> None: assert driver.calls == [f"SELECT 1; {_READBACK_SQL_COMMENT}"] assert ledger.is_empty() + def test_empty_split_fragment_after_execute_is_not_recorded(self) -> None: + driver = _SqlDriver() + ledger = AppliedTuningLedger() + proxy = recording_connection(driver, ledger, PHASE_DDL) + + proxy.execute("CREATE INDEX idx ON t(x);;") + + assert driver.calls == ["CREATE INDEX idx ON t(x);;"] + assert [s.statement for s in ledger.statements] == ["CREATE INDEX idx ON t(x);"] + + def test_empty_split_fragment_after_verb_call_is_not_recorded(self) -> None: + driver = _SqlDriver() + ledger = AppliedTuningLedger() + proxy = recording_connection(driver, ledger, PHASE_DDL, execute_verbs=("sql",)) + + proxy.sql("CREATE INDEX idx ON t(x);;") + + assert [s.statement for s in ledger.statements] == ["CREATE INDEX idx ON t(x);"] + + def test_empty_split_fragment_after_executescript_is_not_recorded(self) -> None: + class _ScriptDriver: + def __init__(self) -> None: + self.scripts: list[str] = [] + + def executescript(self, script, *args, **kwargs): + self.scripts.append(str(script)) + + driver = _ScriptDriver() + ledger = AppliedTuningLedger() + proxy = recording_connection(driver, ledger, PHASE_DDL) + + proxy.executescript("CREATE INDEX idx ON t(x);;") + + assert [s.statement for s in ledger.statements] == ["CREATE INDEX idx ON t(x);"] + + def test_delimiter_and_comment_tail_after_execute_is_not_recorded(self) -> None: + driver = _SqlDriver() + ledger = AppliedTuningLedger() + proxy = recording_connection(driver, ledger, PHASE_DDL) + + proxy.execute(f"CREATE INDEX idx ON t(x);; {_TRAILING_SQL_COMMENT}") + + assert [s.statement for s in ledger.statements] == ["CREATE INDEX idx ON t(x);"] + class TestQueryJobs: def test_submission_alone_is_not_recorded(self) -> None: @@ -383,6 +427,8 @@ def test_bare_keyword_as_a_column_name_is_not_schema_tuning(self, statement: str "CREATE TABLE t (a INT) WITH (sorted_by = ARRAY['a'])", "CREATE TABLE t (a INT) WITH (bucketed_by = ARRAY['a'])", "CREATE TABLE t (a INT UNIQUE)", + "CREATE TABLE t (a INT, UNIQUE NULLS NOT DISTINCT (a))", + "CREATE TABLE t (a INT, UNIQUE NULLS DISTINCT (a))", ], ) def test_bare_keyword_in_clause_syntax_is_schema_tuning(self, statement: str) -> None: From 53e28f27d6c984536283c6506dde45d1ed756326 Mon Sep 17 00:00:00 2001 From: Joe Harris Date: Sat, 10 Oct 2026 07:45:32 -0400 Subject: [PATCH 5/5] fix(tuning): defer ledger capture for lazy frames and record executescript failures --- benchbox/core/tuning/applied_ledger.py | 70 ++++++++- benchbox/platforms/_spark_helpers.py | 13 +- .../core/tuning/test_ledger_execute_verbs.py | 135 ++++++++++++++++++ 3 files changed, 214 insertions(+), 4 deletions(-) diff --git a/benchbox/core/tuning/applied_ledger.py b/benchbox/core/tuning/applied_ledger.py index 98275d8317..225d08a4ad 100644 --- a/benchbox/core/tuning/applied_ledger.py +++ b/benchbox/core/tuning/applied_ledger.py @@ -391,6 +391,8 @@ def _run_verb(self, fn: Any, args: tuple, kwargs: dict) -> Any: return result if callable(getattr(result, "result", None)) and callable(getattr(result, "done", None)): return _RecordingJob(result, self._ledger, self._phase, statements) + if _is_lazy_frame(result): + return _RecordingFrame(result, self._ledger, self._phase, statements) self._record_all(statements, EXECUTED) return result @@ -452,6 +454,66 @@ def __getattr__(self, name: str) -> Any: return getattr(self._job, name) +_LAZY_FRAME_SUBMISSION_METHODS = frozenset({"collect", "show", "fetchall", "fetch_all"}) +_LAZY_FRAME_ACTION_METHODS = frozenset( + { + "collect", + "show", + "fetchall", + "fetch_all", + "count", + "to_pandas", + "toPandas", + "to_polars", + "to_arrow", + } +) + + +def _is_lazy_frame(result: Any) -> bool: + try: + for name in _LAZY_FRAME_SUBMISSION_METHODS: + if callable(getattr(result, name, None)): + return True + except Exception: + return False + return False + + +class _RecordingFrame: + __slots__ = ("_frame", "_ledger", "_phase", "_statements", "_settled") + + def __init__(self, frame: Any, ledger: AppliedTuningLedger, phase: str, statements: list[str]) -> None: + object.__setattr__(self, "_frame", frame) + object.__setattr__(self, "_ledger", ledger) + object.__setattr__(self, "_phase", phase) + object.__setattr__(self, "_statements", statements) + object.__setattr__(self, "_settled", False) + + def _settle(self, status: str, error: Any | None = None) -> None: + if self._settled: + return + object.__setattr__(self, "_settled", True) + for statement in self._statements: + self._ledger.record(statement, self._phase, status=status, error=error) + + def __getattr__(self, name: str) -> Any: + attribute = getattr(self._frame, name) + if name in _LAZY_FRAME_ACTION_METHODS and callable(attribute): + + def call(*args: Any, **kwargs: Any) -> Any: + try: + outcome = attribute(*args, **kwargs) + except Exception as exc: + self._settle(STATEMENT_FAILED, exc) + raise + self._settle(EXECUTED) + return outcome + + return call + return attribute + + class RecordingConnection(_RecordingProxy): __slots__ = ("_conn", "_verbs") @@ -474,7 +536,13 @@ def cursor(self, *args: Any, **kwargs: Any) -> Any: return _RecordingCursor(self._conn.cursor(*args, **kwargs), self._ledger, self._phase, self._statement_filter) def executescript(self, script: Any, *args: Any, **kwargs: Any) -> Any: - result = self._conn.executescript(script, *args, **kwargs) + try: + result = self._conn.executescript(script, *args, **kwargs) + except Exception as exc: + for statement in _split_sql_script(script): + if self._should_record(statement): + self._ledger.record(statement, self._phase, status=STATEMENT_FAILED, error=exc) + raise for statement in _split_sql_script(script): if self._should_record(statement): self._ledger.record(statement, self._phase, status=EXECUTED) diff --git a/benchbox/platforms/_spark_helpers.py b/benchbox/platforms/_spark_helpers.py index 456136bcc9..dfd15de49a 100644 --- a/benchbox/platforms/_spark_helpers.py +++ b/benchbox/platforms/_spark_helpers.py @@ -360,6 +360,13 @@ def _record_spark_conf_ledger(ledger: Any, key: str, value: Any, *, applied: boo pass +def _run_schema_statement(spark: Any, statement: str) -> None: + result = spark.sql(statement) + collect = getattr(result, "collect", None) + if callable(collect): + collect() + + def run_spark_schema_creation_loop( spark: Any, statements: list[str], @@ -381,7 +388,7 @@ def run_spark_schema_creation_loop( if not statement.strip(): continue try: - spark.sql(statement) + _run_schema_statement(spark, statement) logger.debug(f"Executed schema statement: {statement[:100]}...") except Exception as exc: error_lower = str(exc).lower() @@ -399,7 +406,7 @@ def run_spark_schema_creation_loop( "drop the conflicting object manually." ) from exc - spark.sql(f"DROP TABLE IF EXISTS {table_name}") + _run_schema_statement(spark, f"DROP TABLE IF EXISTS {table_name}") if on_location_collision is not None and "location_already_exists" in error_lower: on_location_collision(spark, table_name) - spark.sql(statement) + _run_schema_statement(spark, statement) diff --git a/tests/unit/core/tuning/test_ledger_execute_verbs.py b/tests/unit/core/tuning/test_ledger_execute_verbs.py index 65e6369e8b..7e0b046c5a 100644 --- a/tests/unit/core/tuning/test_ledger_execute_verbs.py +++ b/tests/unit/core/tuning/test_ledger_execute_verbs.py @@ -38,6 +38,43 @@ def execute(self, statement, *args, **kwargs): return "executed" +class _LazyFrame: + def __init__(self, error: Exception | None = None) -> None: + self._error = error + self.actions: list[str] = [] + + def collect(self, *args, **kwargs): + self.actions.append("collect") + if self._error is not None: + raise self._error + return [("row",)] + + def show(self, *args, **kwargs): + self.actions.append("show") + if self._error is not None: + raise self._error + return None + + def count(self, *args, **kwargs): + self.actions.append("count") + if self._error is not None: + raise self._error + return 1 + + +class _LazySqlDriver: + def __init__(self, error: Exception | None = None) -> None: + self.calls: list[str] = [] + self.frames: list[_LazyFrame] = [] + self._error = error + + def sql(self, statement, *args, **kwargs): + self.calls.append(str(statement)) + frame = _LazyFrame(self._error) + self.frames.append(frame) + return frame + + class _FakeJob: def __init__(self, error: Exception | None = None) -> None: self._error = error @@ -373,6 +410,77 @@ def query(self, statement, job_config=None, **kwargs): assert [s.status for s in ledger.statements] == [STATEMENT_FAILED] +class TestLazyFrames: + def test_submission_alone_records_nothing(self) -> None: + driver = _LazySqlDriver() + proxy, ledger = _ddl_proxy(driver, ("sql",)) + + proxy.sql("CREATE TABLE t (a INT) PARTITIONED BY (a)") + + assert driver.calls == ["CREATE TABLE t (a INT) PARTITIONED BY (a)"] + assert ledger.is_empty() + assert ledger.overall_status(tuning_enabled=True, has_config=True) == "noop" + + def test_action_records_executed_once(self) -> None: + driver = _LazySqlDriver() + proxy, ledger = _ddl_proxy(driver, ("sql",)) + + frame = proxy.sql("CREATE TABLE t (a INT) PARTITIONED BY (a)") + assert ledger.is_empty() + frame.collect() + frame.show() + + assert [(s.phase, s.status) for s in ledger.statements] == [(PHASE_DDL, EXECUTED)] + assert "PARTITIONED BY (a)" in ledger.statements[0].statement + + def test_failing_action_records_failed(self) -> None: + driver = _LazySqlDriver(error=RuntimeError("action failed")) + proxy, ledger = _ddl_proxy(driver, ("sql",)) + + frame = proxy.sql("CREATE TABLE t (a INT) PARTITIONED BY (a)") + assert ledger.is_empty() + with pytest.raises(RuntimeError, match="action failed"): + frame.collect() + + assert [s.status for s in ledger.statements] == [STATEMENT_FAILED] + assert "action failed" in (ledger.statements[0].error or "") + assert ledger.overall_status(tuning_enabled=True, has_config=True) == "failed" + + def test_readback_frame_action_is_not_recorded(self) -> None: + driver = _LazySqlDriver() + proxy, ledger = _ddl_proxy(driver, ("sql",)) + + frame = proxy.sql("SELECT COUNT(*) FROM t") + frame.collect() + + assert len(driver.calls) == 1 + assert ledger.is_empty() + + def test_eager_none_result_still_records_at_submission(self) -> None: + driver = _SqlDriver() + proxy, ledger = _ddl_proxy(driver, ("sql",)) + + proxy.sql("CREATE TABLE t (a INT) PARTITIONED BY (a)") + + assert [(s.phase, s.status) for s in ledger.statements] == [(PHASE_DDL, EXECUTED)] + + +class TestExecutescriptFailures: + def test_failing_script_records_failed_per_statement(self) -> None: + class _FailingScriptDriver: + def executescript(self, script, *args, **kwargs): + raise RuntimeError("script boom") + + proxy, ledger = _ddl_proxy(_FailingScriptDriver(), ()) + + with pytest.raises(RuntimeError, match="script boom"): + proxy.executescript("CREATE INDEX i ON t(x); CREATE INDEX j ON t(y);") + + assert [s.status for s in ledger.statements] == [STATEMENT_FAILED, STATEMENT_FAILED] + assert all("script boom" in (s.error or "") for s in ledger.statements) + assert ledger.overall_status(tuning_enabled=True, has_config=True) == "failed" + + class TestFootprint: @pytest.mark.parametrize( "statement", @@ -462,6 +570,33 @@ def test_spark_tuned_schema_ddl_reaches_the_ledger(self) -> None: assert {s.phase for s in ledger.statements} == {PHASE_DDL} assert ledger.overall_status(tuning_enabled=True, has_config=True) == "applied_unverified" + def test_spark_tuned_schema_ddl_with_lazy_frames_reaches_the_ledger(self) -> None: + from benchbox.platforms.spark import SparkAdapter + + driver = _LazySqlDriver() + proxy, ledger = _ddl_proxy(driver, SparkAdapter.ledger_execute_verbs) + fake_self = SimpleNamespace( + _create_schema_with_tuning=lambda *a, **k: ( + "CREATE TABLE orders (o_id INT, o_date DATE) USING parquet PARTITIONED BY (o_date);\n" + "CREATE TABLE lineitem (l_id INT, l_key INT) USING delta CLUSTER BY (l_key);\n" + "CREATE TABLE plain (p_id INT);" + ), + table_format="parquet", + logger=logging.getLogger("test"), + _remove_orphaned_table_location=lambda *a, **k: None, + ) + + SparkAdapter.create_schema(fake_self, SimpleNamespace(), proxy) + + assert len(driver.calls) == 3 + assert all(frame.actions for frame in driver.frames) + recorded = [s.statement for s in ledger.statements] + assert len(recorded) == 2 + assert any("PARTITIONED BY" in s for s in recorded) + assert any("CLUSTER BY" in s for s in recorded) + assert {s.phase for s in ledger.statements} == {PHASE_DDL} + assert ledger.overall_status(tuning_enabled=True, has_config=True) == "applied_unverified" + def test_bigquery_tuned_schema_ddl_records_final_job_state(self) -> None: from benchbox.platforms.bigquery import BigQueryAdapter