fix(spark): strip trailing semicolon before sql/dry_run by Bartok9 · Pull Request #2464 · Canner/WrenAI

Bartok9 · 2026-07-10T06:02:28Z

Summary

Strip trailing ; / whitespace before Spark Connect sql() in both query and dry_run.
Client SQL ending in a semicolon no longer fails or mis-parses on Spark Connect.

Motivation

Postgres/redshift/duckdb/databricks/datafusion already sanitize trailing semicolons before plan/execute paths. Spark still passed raw SQL into SparkSession.sql(...), so statements like SELECT 1; hit Connect/parser multi-statement roughness and pouch provenance protection. Reuse the terminating-run strip helper so literals such as SELECT ';' stay intact while terminal ; is removed.

Verification

core/wren/.venv/bin/python -m pytest tests/unit/test_spark_semicolon.py -q — 4 passed
Mocked session: query("SELECT 1;", limit=2) calls sql("SELECT 1"); dry_run strips trailing whitespace+semicolon
Did NOT change: client-side Arrow slice limit behavior

License

core/wren/** only (Apache-2.0).

Summary by CodeRabbit

Bug Fixes
- Improved Spark SQL execution by automatically removing trailing semicolons (and surrounding whitespace/newlines) before running queries.
- Ensured semicolons inside string literals remain unchanged.
- Preserved dry-run behavior, including keeping the limit at zero.
Tests
- Added unit coverage for trailing-semicolon normalization for both query execution and dry-run.

coderabbitai · 2026-07-10T06:06:46Z

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

@coderabbitai resume to resume automatic reviews.
@coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

▶️ Resume reviews
🔍 Trigger review

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: e98483b1-7305-4bf7-bb55-d42fc10937b5

📥 Commits

Reviewing files that changed from the base of the PR and between 6a9b5f2 and fda0362.

📒 Files selected for processing (2)

core/wren/src/wren/connector/spark.py
core/wren/tests/unit/test_spark_semicolon.py

Walkthrough

Spark query and dry-run execution now strip trailing semicolons and whitespace before calling the Spark session. Unit tests cover normalization, string literals, unchanged SQL, and the dry-run zero-row limit.

Changes

Spark SQL normalization

Layer / File(s)	Summary
Normalize trailing semicolons `core/wren/src/wren/connector/spark.py`	Uses `strip_trailing_semicolon()` for `query()` and `dry_run()` before Spark SQL execution.
Validate semicolon behavior `core/wren/tests/unit/test_spark_semicolon.py`	Tests query and dry-run normalization, string-literal preservation, unchanged SQL, and the dry-run zero-row limit.

Estimated code review effort: 2 (Simple) | ~10 minutes

Possibly related issues

Strip trailing semicolon on the unlimited query() path across connectors #2497 — Directly covers Spark query and dry-run semicolon stripping.
Strip trailing semicolon on dry_run / EXPLAIN / describe across connectors #2498 — Covers applying strip_trailing_semicolon to Spark validation and execution paths.
Hoist trailing-semicolon strip to the engine.query/dry_run boundary #2499 — Covers trailing-semicolon normalization at connector query and dry-run boundaries.

Possibly related PRs

Canner/WrenAI#2422 — Applies related SQL preprocessing to the Databricks dry-run path.
Canner/WrenAI#2423 — Applies equivalent normalization to Oracle query and dry-run paths with tests.
Canner/WrenAI#2430 — Applies equivalent normalization to DataFusion query handling.

Suggested reviewers: goldmedal

Poem

A bunny trims the SQL tail,
While quoted marks remain without fail.
Queries hop, dry-runs count,
Zero rows are the amount. 🐇

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 9.09% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly matches the main change: Spark now strips trailing semicolons before sql and dry_run.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands.}

coderabbitai

🧹 Nitpick comments (1)

core/wren/tests/unit/test_spark_semicolon.py (1)
16-26: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Remove unused mock variables df and dfto.

Lines 18-20 create df and dfto as MagicMock objects that are never used. The actual mock chain is set up on line 24 via session.sql.return_value.toPandas.return_value, making these variables dead code that adds confusion about the test's mock setup.
🧹 Proposed cleanup
 def test_query_strips_trailing_semicolon_before_sql() -> None:
     connector, session = _make_mock_connector()
-    df = MagicMock()
-    df.attrs = {}
-    dfto = MagicMock()
     # pandas DF mock for pa.Table.from_pandas
     import pandas as pd
🤖 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 `@core/wren/tests/unit/test_spark_semicolon.py` around lines 16 - 26, Remove
the unused MagicMock variables df and dfto from
test_query_strips_trailing_semicolon_before_sql; retain only the
session.sql.return_value.toPandas.return_value setup needed by the test.

🤖 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.

Nitpick comments:
In `@core/wren/tests/unit/test_spark_semicolon.py`:
- Around line 16-26: Remove the unused MagicMock variables df and dfto from
test_query_strips_trailing_semicolon_before_sql; retain only the
session.sql.return_value.toPandas.return_value setup needed by the test.

ℹ️ Review info