Cancel current query on cursor.close()

Author: mdesmet

tests/integration/test_dbapi_integration.py

@@ -775,6 +775,18 @@ def test_cancel_query(trino_connection):
     assert "Cancel query failed; no running query" in str(cancel_error.value)
 
 
+def test_close_cursor(trino_connection):
+    cur = trino_connection.cursor()
+    cur.execute("SELECT * FROM tpch.sf1.customer")
+    cur.fetchone()
+    cur.close()  # would raise an exception if cancel fails
+
+    cur = trino_connection.cursor()
+    with pytest.raises(Exception) as cancel_error:
+        cur.close()
+    assert "Cancel query failed; no running query" in str(cancel_error.value)
+
+
 def test_session_properties(run_trino):
     _, host, port = run_trino
 

trino/dbapi.py

@@ -567,7 +567,9 @@ def cancel(self):
         self._query.cancel()
 
     def close(self):
-        self._connection.close()
+        self.cancel()
+        # TODO: Cancel not only the last query executed on this cursor
+        #  but also any other outstanding queries executed through this cursor.
 
 
 Date = datetime.date