Closing all cursors of a connection on a connection close

[Jrmyy]

👋🏻 Hello community,

First, thanks for this package, we work with it extensively! 🥳

I was a previous maintainer of dbt-athena package. Under the hood, dbt-athena uses PyAthena to connect and execute queries on Athena.

We currently encounter an issue when we need to close a connection: dbt-labs/dbt-adapters#406. Basically, when we cancel a dbt invocation, we cannot properly cancel the running queries because a connection has no information about the created cursors, as for instance the big query client does.

Would you be open to add such feature in PyAthena so we can, on a cancel, stops all the cursors?

[laughingman7743]

I do not think it would be that difficult to implement the same. I will address this when I have time.

[laughingman7743]

https://github.com/googleapis/python-bigquery/blob/main/google/cloud/bigquery/dbapi/cursor.py#L104-L106
It seems that BigQuery clients do not cancel query execution when Cursor is closed, have you checked the behavior?

Cursor polls until the end of query execution.
https://github.com/laughingman7743/PyAthena/blob/master/pyathena/common.py#L486-L496
The query can be canceled by calling Cursor's cancel method from another thread.
https://github.com/laughingman7743/PyAthena/blob/master/pyathena/cursor.py#L124-L127

PyAthena/tests/pyathena/test_cursor.py

Lines 548 to 565 in 2787190

	def test_cancel(self, cursor):
	def cancel(c):
	time.sleep(randint(5, 10))
	c.cancel()
	with ThreadPoolExecutor(max_workers=1) as executor:
	executor.submit(cancel, cursor)
	pytest.raises(
	DatabaseError,
	lambda: cursor.execute(
	"""
	SELECT a.a * rand(), b.a * rand()
	FROM many_rows a
	CROSS JOIN many_rows b
	"""
	),
	)

https://github.com/PyMySQL/PyMySQL/blob/main/pymysql/connections.py#L414-L434
https://github.com/tlocke/pg8000/blob/main/src/pg8000/legacy.py#L754-L756
I checked other Python DB-API implementations and could not find one that manages the Cursor class created by Connection. shouldn't the Connection class manage Cursor? I don't think it is a good implementation.

[Jrmyy]

Indeed, on bigquery client, it does not seem the query is being stopped when the cursor is closed 🤔

Thanks for the research you did. The issue might need to be handled in dbt.