Do not cancel completed DML statements on cursor close

tests/integration/conftest.py

@@ -43,7 +43,7 @@ def get_local_port():
 @pytest.fixture(scope="session")
 def run_trino():
     host = os.environ.get("TRINO_RUNNING_HOST", TRINO_HOST)
-    port = os.environ.get("TRINO_RUNNING_PORT", DEFAULT_PORT)
+    port = int(os.environ.get("TRINO_RUNNING_PORT", DEFAULT_PORT))
 
     # Is there any local Trino available
     if is_trino_available(host, port):

tests/integration/test_dbapi_integration.py

@@ -226,6 +226,38 @@ def test_execute_many(legacy_prepared_statements, run_trino):
         cur.execute("DROP TABLE IF EXISTS memory.default.test_execute_many")
 
 
+def test_update_query_close_without_fetch_is_not_cancelled(trino_connection):
+    """Regression test for https://github.com/trinodb/trino-python-client/issues/601
+
+    A consumer (e.g. the SQLAlchemy dialect) that executes an INSERT and closes
+    the cursor without fetching the result must not cause the already-completed
+    query to be reported as USER_CANCELED by the server.
+    """
+    cur = trino_connection.cursor()
+    cur.execute("CREATE TABLE IF NOT EXISTS memory.default.test_issue_601 (key int)")
+    cur.fetchall()
+    try:
+        cur = trino_connection.cursor()
+        cur.execute("INSERT INTO memory.default.test_issue_601 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.fetchall()
+
+
 @pytest.mark.parametrize(
     "legacy_prepared_statements",
     [

tests/integration/test_sqlalchemy_integration.py

@@ -237,6 +237,53 @@ def test_insert_multiple_statements(trino_connection):
     metadata.drop_all(engine)
 
 
+@pytest.mark.skipif(
+    sqlalchemy_version() < "1.4",
+    reason="columns argument to select() must be a Python list or other iterable"
+)
+@pytest.mark.parametrize('trino_connection', ['memory'], indirect=True)
+def test_insert_is_not_cancelled(trino_connection):
+    """Regression test for https://github.com/trinodb/trino-python-client/issues/601
+
+    The SQLAlchemy dialect closes the cursor after a DML statement without
+    fetching its result. Previously that issued a spurious cancel and the
+    server reported the already-completed INSERT as USER_CANCELED. The rows are
+    inserted either way, so this asserts on the server-side query state.
+    """
+    engine, conn = trino_connection
+    if not engine.dialect.has_schema(conn, "test"):
+        with engine.begin() as connection:
+            connection.execute(sqla.schema.CreateSchema("test"))
+    metadata = sqla.MetaData()
+    probes = sqla.Table(f"test_issue_601_{uuid.uuid4().hex[:12]}",
+                        metadata,
+                        sqla.Column('id', sqla.Integer),
+                        sqla.Column('name', sqla.String),
+                        schema="test")
+    metadata.create_all(engine)
+
+    insert_query_ids = []
+
+    def capture_insert_query_id(conn_, cursor, statement, parameters, context, executemany):
+        if "INSERT INTO" in statement.upper():
+            insert_query_ids.append(cursor.query_id)
+
+    sqla.event.listen(engine, "after_cursor_execute", capture_insert_query_id)
+    try:
+        conn.execute(probes.insert(), [{"id": i, "name": f"n{i}"} for i in range(100)])
+
+        assert insert_query_ids, "did not observe an INSERT statement being executed"
+        for query_id in insert_query_ids:
+            state, error_code = conn.execute(sqla.text(
+                "SELECT state, error_code FROM system.runtime.queries WHERE query_id = :query_id"
+            ), {"query_id": query_id}).one()
+            assert state == "FINISHED", f"query {query_id} ended in state {state} ({error_code})"
+            assert error_code is None
+    finally:
+        sqla.event.remove(engine, "after_cursor_execute", capture_insert_query_id)
+        metadata.drop_all(engine)
+
+
 @pytest.mark.skipif(
     sqlalchemy_version() < "1.4",
     reason="columns argument to select() must be a Python list or other iterable"

tests/unit/test_dbapi.py

@@ -9,6 +9,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import json
 import threading
 import uuid
 from unittest.mock import patch
@@ -373,3 +374,110 @@ def test_error_when_auth_over_http():
 
 def test_no_error_when_auth_over_https():
     Connection("mytrinoserver.domain", http_scheme=constants.HTTPS, auth=BasicAuthentication("u", "p"))
+
+
+def _statement_uri(query_id, token):
+    return f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}/{query_id}/{token}"
+
+
+@httprettified
+def test_cursor_close_does_not_cancel_finished_update_query():
+    """Regression test for https://github.com/trinodb/trino-python-client/issues/601
+
+    An update statement (INSERT/UPDATE/DELETE) reports its affected row count as
+    a single synthetic row while Trino still returns a final nextUri. Closing the
+    cursor without fetching must drain that nextUri instead of issuing a DELETE,
+    otherwise the already-completed statement is reported as USER_CANCELED.
+    """
+    query_id = "20210817_140827_00000_arvdv"
+    statement_path = f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}"
+
+    post_response = {
+        "id": query_id,
+        "nextUri": _statement_uri(query_id, 1),
+        "infoUri": f"{SERVER_ADDRESS}/query.html?{query_id}",
+        "stats": {"state": "QUEUED"},
+    }
+    # The update-count row arrives together with a still-present nextUri.
+    update_response = {
+        "id": query_id,
+        "nextUri": _statement_uri(query_id, 2),
+        "infoUri": f"{SERVER_ADDRESS}/query.html?{query_id}",
+        "updateType": "INSERT",
+        "updateCount": 1000,
+        "columns": [{
+            "name": "rows",
+            "type": "bigint",
+            "typeSignature": {"rawType": "bigint", "arguments": [], "typeArguments": []},
+        }],
+        "data": [[1000]],
+        "stats": {"state": "FINISHED"},
+    }
+    # Final response transitions the query to a terminal state (no nextUri).
+    # Trino keeps repeating updateType/updateCount on the trailing pages.
+    final_response = {
+        "id": query_id,
+        "infoUri": f"{SERVER_ADDRESS}/query.html?{query_id}",
+        "updateType": "INSERT",
+        "updateCount": 1000,
+        "columns": update_response["columns"],
+        "stats": {"state": "FINISHED"},
+    }
+
+    httpretty.register_uri(method=httpretty.POST, uri=statement_path, body=json.dumps(post_response))
+    httpretty.register_uri(method=httpretty.GET, uri=_statement_uri(query_id, 1), body=json.dumps(update_response))
+    httpretty.register_uri(method=httpretty.GET, uri=_statement_uri(query_id, 2), body=json.dumps(final_response))
+    httpretty.register_uri(method=httpretty.DELETE, uri=_statement_uri(query_id, 2), status=204)
+
+    with connect("coordinator", user="test", http_scheme=constants.HTTPS) as conn:
+        cur = conn.cursor()
+        cur.execute("INSERT INTO some_table VALUES (1), (2), (3)")
+        # execute() must have drained the query to a terminal state.
+        assert cur._query.finished is True
+        assert cur.rowcount == 1000
+        cur.close()
+
+    delete_requests = [r for r in httpretty.latest_requests() if r.method == "DELETE"]
+    assert delete_requests == [], "closing a finished update query must not issue a cancel"
+
+
+@httprettified
+def test_cursor_close_cancels_unfinished_query():
+    """Closing a cursor whose result set has not been fully consumed must still
+    cancel the running query so the server can free its resources.
+    """
+    query_id = "20210817_140827_00000_arvdv"
+    statement_path = f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}"
+
+    post_response = {
+        "id": query_id,
+        "nextUri": _statement_uri(query_id, 1),
+        "infoUri": f"{SERVER_ADDRESS}/query.html?{query_id}",
+        "stats": {"state": "QUEUED"},
+    }
+    # A SELECT that returns a first page of data with more still pending.
+    data_response = {
+        "id": query_id,
+        "nextUri": _statement_uri(query_id, 2),
+        "infoUri": f"{SERVER_ADDRESS}/query.html?{query_id}",
+        "columns": [{
+            "name": "x",
+            "type": "bigint",
+            "typeSignature": {"rawType": "bigint", "arguments": [], "typeArguments": []},
+        }],
+        "data": [[1]],
+        "stats": {"state": "RUNNING"},
+    }
+
+    httpretty.register_uri(method=httpretty.POST, uri=statement_path, body=json.dumps(post_response))
+    httpretty.register_uri(method=httpretty.GET, uri=_statement_uri(query_id, 1), body=json.dumps(data_response))
+    httpretty.register_uri(method=httpretty.DELETE, uri=_statement_uri(query_id, 2), status=204)
+
+    with connect("coordinator", user="test", http_scheme=constants.HTTPS) as conn:
+        cur = conn.cursor()
+        cur.execute("SELECT x FROM some_table")
+        assert cur._query.finished is False
+        cur.close()
+
+    delete_requests = [r for r in httpretty.latest_requests() if r.method == "DELETE"]
+    assert len(delete_requests) == 1, "closing an unfinished query must cancel it"

trino/client.py

@@ -1006,6 +1006,16 @@ def execute(self, additional_http_headers=None) -> TrinoResult:
                 except StopIteration:
                     self._result.rows = []
 
+        # Update statements (INSERT/UPDATE/DELETE/DDL/...) report their affected
+        # row count as a single synthetic row, but Trino still returns a final
+        # nextUri that must be consumed for the query to reach a terminal state.
+        # Unlike a SELECT there is no result set to stream, so drain the remaining
+        # pages now. Otherwise closing the cursor without fetching would issue a
+        # DELETE against an already-completed statement, which Trino reports as
+        # USER_CANCELED (see https://github.com/trinodb/trino-python-client/issues/601).
+        while self._update_type is not None and not self.finished and not self.cancelled:
+            self._result.rows += self.fetch()
+
         return self._result
 
     def _update_state(self, status):