diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 3184de6e..82938970 100644 --- a/tests/integration/conftest.py +++ b/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): diff --git a/tests/integration/test_dbapi_integration.py b/tests/integration/test_dbapi_integration.py index 35ef0b30..e6123a1f 100644 --- a/tests/integration/test_dbapi_integration.py +++ b/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", [ diff --git a/tests/integration/test_sqlalchemy_integration.py b/tests/integration/test_sqlalchemy_integration.py index e5160bb0..571f67d7 100644 --- a/tests/integration/test_sqlalchemy_integration.py +++ b/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" diff --git a/tests/unit/test_dbapi.py b/tests/unit/test_dbapi.py index 17e024df..25fde653 100644 --- a/tests/unit/test_dbapi.py +++ b/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" diff --git a/trino/client.py b/trino/client.py index abf1c2f2..a7e43568 100644 --- a/trino/client.py +++ b/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):