[Bug] Cancelling dbt invocation does not cancel the query on Athena

[Jrmyy]

👋🏻 Hello dbt-athena community,

Is this a new bug in dbt-athena?

• I believe this is a new bug in dbt-athena
• I have searched the existing issues, and I could not find an existing issue for this bug

Current Behavior

When I try to cancel a dbt invocation, the invocation can be cancelled but the query is not cancelled in Athena. Therefore it can lead to concurrent requests when retrying.

Expected Behavior

When the dbt invocation cancel process is done, the query should be cancelled in Athena

Steps To Reproduce

1. Run a model
2. Cancel the invocation (for instance with CTRL + c)
3. Check in Athena console if the query is still running or not

Relevant log output

No response

Environment

- OS: MacOS (the same problem in Docker image based on Debian)
- Python: 3.12.7
- dbt-core: 1.8.9
- dbt-athena: 1.8.4

Additional Context

No response

[Jrmyy]

The adapter cancel behavior is defined here:

dbt-adapters/dbt/adapters/base/impl.py

Line 1327 in b6152d6

def cancel_open_connections(self):

Then if we go to the SQLConnectionsManager, the cancel_open is defined here:

dbt-adapters/dbt/adapters/sql/connections.py

Line 53 in b6152d6

def cancel_open(self) -> List[str]:

But in the AthenaConnectionsManager, we don't implement the cancel method: https://github.com/dbt-labs/dbt-athena/blob/8e2aa424256e354103d619cc07f9a79d85fadf98/dbt-athena/src/dbt/adapters/athena/connections.py#L321-L322

The Connection object comes directly from pyathena: https://github.com/laughingman7743/PyAthena/blob/master/pyathena/connection.py#L42

[Jrmyy]

After deep-diving the code I think I found the issue.

The cancel query happens here: https://github.com/dbt-labs/dbt-core/blob/3de3b827bfffdc43845780f484d4d53011f20a37/core/dbt/task/runnable.py#L429-L448

As it should be, the raise of the exception is done after the cancel of the queries is done.
In dbt-athena, we wait for this exception to try to cancel the query. But it is never caught by our code. I think it is because the KeyboardInterrupt exception is sent to the main thread while the run of the dbt-model is sent to a thread run

https://github.com/dbt-labs/dbt-adapters/blob/main/dbt-athena/src/dbt/adapters/athena/connections.py#L213

Therefore, I think implemented the cancel method will fix the error. But I also think the _poll method is not working as expected (https://github.com/dbt-labs/dbt-adapters/blob/main/dbt-athena/src/dbt/adapters/athena/connections.py#L128-L138).