Do not cancel completed DML statements on cursor close#617
Conversation
An INSERT/UPDATE/DELETE reports its affected row count as a single synthetic row but Trino still sends a final `nextUri` that must be consumed for the query to reach a terminal state. Before this change `TrinoQuery.execute()` stopped as soon as that count row arrived leaving `nextUri` unconsumed. For consumers that read the row count without fetching results and then closed the cursor observed that Trino reported the query as USER_CANCELLED even though the query had functionally completed (DML was comitted). This change drains DML statements to a terminal state in `execute()` so `nextUri` is fully consumed, and make `Cursor.close()` only cancel queries that are still running.
`TRINO_RUNNING_PORT` is read from the environment as a string and passed straight into `socket.connect_ex`, which requires an integer port and raises `TypeError`. This made it impossible to point the integration suite at an already-running server via the env var. Cast the value to an int at read time, matching the int `DEFAULT_PORT` fallback.
|
@coderabbitai full review |
✅ Action performedFull review finished. |
📝 WalkthroughWalkthroughThe client now drains remaining result pages for update statements before considering execution complete. Unit tests verify that finished update queries are not canceled while unfinished queries are canceled on cursor closure. New DBAPI and SQLAlchemy integration tests verify successful insert completion without Assessment against linked issues
Assessment against linked issues: Out-of-scope changes
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. Comment |
![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)
tests/integration/test_dbapi_integration.py (1)
236-258: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse a unique table name to prevent test flakiness.
Using a hardcoded table name (
test_issue_601) can cause flaky tests if the integration test suite is executed concurrently against the same Trino cluster. Consider generating a unique table name usinguuid.uuid4().hex, similar to the approach used intest_sqlalchemy_integration.py.♻️ Proposed refactor
Ensure you add
import uuidat the top of the file, then update the test as follows:cur = trino_connection.cursor() - cur.execute("CREATE TABLE IF NOT EXISTS memory.default.test_issue_601 (key int)") + table_name = f"test_issue_601_{uuid.uuid4().hex[:12]}" + cur.execute(f"CREATE TABLE IF NOT EXISTS memory.default.{table_name} (key int)") cur.fetchall() try: cur = trino_connection.cursor() - cur.execute("INSERT INTO memory.default.test_issue_601 VALUES (1), (2), (3)") + cur.execute(f"INSERT INTO memory.default.{table_name} VALUES (1), (2), (3)") # Emulate a consumer that reads the row count but never fetches the rows. assert cur.rowcount == 3 insert_query_id = cur.query_id cur.close() state_cur = trino_connection.cursor() state_cur.execute( "SELECT state, error_code FROM system.runtime.queries " f"WHERE query_id = '{insert_query_id}'" ) state, error_code = state_cur.fetchall()[0] assert state == "FINISHED" assert error_code is None finally: cur = trino_connection.cursor() - cur.execute("DROP TABLE IF EXISTS memory.default.test_issue_601") + cur.execute(f"DROP TABLE IF EXISTS memory.default.{table_name}") cur.fetchall()🤖 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 `@tests/integration/test_dbapi_integration.py` around lines 236 - 258, Update the integration test around the test_issue_601 table usage to generate a unique table name with uuid.uuid4().hex, and add the uuid import. Reuse that generated name consistently in the CREATE, INSERT, and DROP statements so concurrent test runs do not share the hardcoded table.
🤖 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 `@trino/client.py`:
- Around line 1016-1017: Update the loop that accumulates results in the query
update flow to handle both list and iterator values in self._result.rows: use
list.extend() when it is a list, and combine it with self.fetch() via
itertools.chain() when it is an iterator. Preserve the existing accumulation
behavior and avoid using += on iterator-backed rows.
---
Nitpick comments:
In `@tests/integration/test_dbapi_integration.py`:
- Around line 236-258: Update the integration test around the test_issue_601
table usage to generate a unique table name with uuid.uuid4().hex, and add the
uuid import. Reuse that generated name consistently in the CREATE, INSERT, and
DROP statements so concurrent test runs do not share the hardcoded table.
🪄 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: 925401f0-ca62-4d22-9c49-b627176e0200
📒 Files selected for processing (5)
tests/integration/conftest.pytests/integration/test_dbapi_integration.pytests/integration/test_sqlalchemy_integration.pytests/unit/test_dbapi.pytrino/client.py
| while self._update_type is not None and not self.finished and not self.cancelled: | ||
| self._result.rows += self.fetch() |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Prevent TypeError when self._result.rows is an iterator.
If the query uses the spooling protocol, the first loop may assign an itertools.chain iterator to self._result.rows. Using the += operator on an itertools.chain object will raise a TypeError: unsupported operand type(s) for +=.
To robustly handle both lists and iterators, use extend() if self._result.rows is a list, and itertools.chain() if it is an iterator.
🛠️ Proposed fix to handle iterators safely
while self._update_type is not None and not self.finished and not self.cancelled:
- self._result.rows += self.fetch()
+ new_rows = self.fetch()
+ if isinstance(self._result.rows, list):
+ self._result.rows.extend(new_rows)
+ else:
+ self._result.rows = itertools.chain(self._result.rows, new_rows)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| while self._update_type is not None and not self.finished and not self.cancelled: | |
| self._result.rows += self.fetch() | |
| while self._update_type is not None and not self.finished and not self.cancelled: | |
| new_rows = self.fetch() | |
| if isinstance(self._result.rows, list): | |
| self._result.rows.extend(new_rows) | |
| else: | |
| self._result.rows = itertools.chain(self._result.rows, new_rows) |
🤖 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 `@trino/client.py` around lines 1016 - 1017, Update the loop that accumulates
results in the query update flow to handle both list and iterator values in
self._result.rows: use list.extend() when it is a list, and combine it with
self.fetch() via itertools.chain() when it is an iterator. Preserve the existing
accumulation behavior and avoid using += on iterator-backed rows.
2 participants
Description
An INSERT/UPDATE/DELETE reports its affected row count as a single
synthetic row but Trino still sends a final
nextUrithat must beconsumed for the query to reach a terminal state. Before this change
TrinoQuery.execute()stopped as soon as that count row arrived leavingnextUriunconsumed.For consumers that read the row count without fetching results and then
closed the cursor observed that Trino reported the query as
USER_CANCELLED even though the query had functionally completed (DML was
comitted).
This change drains DML statements to a terminal state in
execute()sonextUriis fully consumed, and makeCursor.close()only cancelqueries that are still running.
Non-technical explanation
Do not cancel completed DML statements on cursor close. Fixes #601.
Release notes
( ) This is not user-visible or docs only and no release notes are required.
( ) Release notes are required, please propose a release note for me.
(x) Release notes are required, with the following suggested text: