fix(datafusion): strip trailing semicolon on unlimited query path

[Bartok9]

Summary

• Strip terminating semicolon/whitespace on DataFusion unlimited query() before ctx.query.
• Keep limited-path subquery wrap behavior; add regression tests.

Motivation

Limited queries already stripped before wrap. Unlimited still forwarded raw SELECT 1;, which LocalRuntime rejects for single-statement execution.

Verification

• core/wren/.venv/bin/python -m pytest tests/unit/test_datafusion_semicolon.py -q — 8 passed

Duplicate check

• Prior unlimited strip PRs fix(datafusion): strip trailing semicolon on unlimited query and dry_run #2483/fix(datafusion): strip trailing semicolon in dry_run #2506 closed/stale against older tips; main still lacked unlimited strip (test asserted pass-through).

Summary by CodeRabbit

• Bug Fixes
  • Improved SQL query handling by consistently removing trailing semicolons, including those followed by whitespace or newlines.
  • Ensured queries with row limits are executed reliably when wrapped as subqueries.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 724a72b7-4be2-4fe9-9993-de413fd48a27

📥 Commits

Reviewing files that changed from the base of the PR and between d472877 and 5ef0ed4.

📒 Files selected for processing (2)

• core/wren/src/wren/connector/datafusion.py
• core/wren/tests/unit/test_datafusion_semicolon.py

---

Walkthrough

The DataFusion connector now strips trailing semicolons before executing queries in both limited and unlimited paths. Unit tests cover semicolons with and without trailing whitespace or newlines.

Changes

DataFusion SQL cleanup

Layer / File(s)	Summary
Normalize and test query SQL core/wren/src/wren/connector/datafusion.py, core/wren/tests/unit/test_datafusion_semicolon.py	The connector reuses cleaned SQL for both query paths, with tests covering trailing semicolons and whitespace.

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

Possibly related PRs

• Canner/WrenAI#2480: Updates DataFusion query SQL cleaning and related unit tests.
• Canner/WrenAI#2506: Applies shared semicolon stripping in DataFusion execution paths.
• Canner/WrenAI#2555: Normalizes SQL before limit handling and updates unlimited-path tests.

Suggested reviewers: goldmedal

Poem

I’m a bunny with queries, neat and bright,
Trimming semicolons before flight.
Whitespace hops away too,
Tests check each path anew,
Clean SQL lands just right.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 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: stripping trailing semicolons on the unlimited DataFusion query path.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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

• X
• Mastodon
• Reddit
• LinkedIn

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

[goldmedal]

Code review

Two problems, one of which applies to #2594 as well.

1. Duplicate. #2594 makes the same change to the same file with the same test filename, opened a day later. One of the two should close; #2594's diff is the smaller of the pair.

2. The stated justification is not true. The comment claims the unlimited path needs stripping because:

WrenAI/core/wren/src/wren/connector/datafusion.py

Lines 29 to 42 in 5ef0ed4

	self._register_tables()
	def query(self, sql: str, limit: int | None = None) -> pa.Table:
	# Always strip terminating ``;`` / whitespace. Limited path needs it
	# for the subquery wrap; unlimited ``ctx.query`` also rejects a trailing
	# statement terminator on single-statement LocalRuntime execution.
	cleaned = strip_trailing_semicolon(sql)
	if limit is not None:
	sql = f"SELECT * FROM ({cleaned}) AS _q LIMIT {int(limit)}"
	else:
	sql = cleaned
	ipc_bytes = self.ctx.query(sql)
	reader = ipc.open_stream(io.BytesIO(bytes(ipc_bytes)))
	return reader.read_all()

The DataFusion connector executes through self.ctx.query(sql), where ctx is wren_core.SessionContext(). That call accepts trailing terminators:

>>> from wren_core import SessionContext
>>> ctx = SessionContext()
>>> for s in ("SELECT 1", "SELECT 1;", "SELECT 1;  \n", "SELECT 1;;"):
...     len(bytes(ctx.query(s)))
520
520
520
520

All four succeed, including the double terminator. There is no rejection to avoid on this path.

That does not necessarily kill the change — aligning DataFusion with bigquery, duckdb, redshift and now mysql (#2595) is a reasonable consistency argument on its own. But it should be labelled refactor(...) and carry a comment that says what it actually does, the way #2595 was adjusted. Please don't assert engine behaviour in a code comment without checking it; that comment will outlive the PR and mislead whoever reads it next.

[Bartok9]

Thanks for the careful check — you're right, ctx.query() accepts trailing terminators here so the datafusion unlimited path doesn't need the strip. Closing this one (and #2594 as the duplicate). Appreciate the correction.