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

[Bartok9]

Summary

Strip trailing semicolons on DataFusion unlimited queries.

Motivation

Limited/dry_run already stripped; unlimited left SELECT 1; intact and DataFusion rejects it.

Verification

.venv/bin/python -m pytest tests/unit/test_datafusion_semicolon.py -q
7 passed

Summary by CodeRabbit

• Bug Fixes
  • Improved SQL query handling by removing trailing semicolons consistently.
  • Fixed queries without a row limit that could fail due to semicolon-related parsing issues.

[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: e861bed5-ed49-49a6-acee-d7b59d936f56

📥 Commits

Reviewing files that changed from the base of the PR and between 06e4c42 and 7c205b5.

📒 Files selected for processing (2)

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

---

Walkthrough

DataFusionConnector now strips trailing semicolons before handling query limits, including the unlimited path. The corresponding unit test expects semicolon-free SQL for unlimited queries.

Changes

DataFusion query normalization

Layer / File(s)	Summary
Normalize SQL before execution core/wren/src/wren/connector/datafusion.py, core/wren/tests/unit/test_datafusion_semicolon.py	Trailing semicolons are removed before limit handling, and the unlimited-query test expects the normalized SQL.

Estimated code review effort: 1 (Trivial) | ~5 minutes

Possibly related PRs

• Canner/WrenAI#2585: Directly updates the same DataFusion query path and semicolon-focused test.
• Canner/WrenAI#2555: Applies the same unconditional preprocessing before limit wrapping.
• Canner/WrenAI#2430: Changes the same connector and related semicolon handling.

Suggested reviewers: goldmedal

Poem

I nibbled the semicolon away,
So DataFusion parses cleanly today.
Limited or free, the query runs bright,
Tests hop happily into sight.
Squeak! 🐇

🚥 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 for unlimited DataFusion queries.
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

The comment's claim is false. It states that stripping is required because:

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

Lines 29 to 40 in 7c205b5

	self._register_tables()
	def query(self, sql: str, limit: int | None = None) -> pa.Table:
	# Unlimited path must strip trailing terminators too — DataFusion's
	# SQL parser rejects ``SELECT 1;`` the same way subquery-wrap does.
	sql = strip_trailing_semicolon(sql)
	if limit is not None:
	sql = f"SELECT * FROM ({sql}) AS _q LIMIT {int(limit)}"
	ipc_bytes = self.ctx.query(sql)
	reader = ipc.open_stream(io.BytesIO(bytes(ipc_bytes)))
	return reader.read_all()

The connector executes via 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, the double terminator included. DataFusion's parser does not reject SELECT 1; here, so the sentence "the same way subquery-wrap does" is not accurate — the subquery wrap genuinely needs the strip, this path does not.

The change can still land as a consistency alignment with bigquery, duckdb, redshift and mysql (#2595), but it needs:

1. a refactor(...) title rather than fix(...), since no broken path is being repaired;
2. a comment that states the real reason instead of an engine behaviour that does not hold;
3. a decision with fix(datafusion): strip trailing semicolon on unlimited query path #2585, which is the same change to the same file with the same test filename.

Also worth a rebase before merge — this branch is 3 commits behind main.

[Bartok9]

Confirmed — SessionContext().query() handles trailing ; fine, so this strip isn't needed on the datafusion path. Closing. Thanks for verifying.