Implement on_kill() for SQLExecuteQueryOperator #27514

kazanzhy · 2022-11-05T00:15:40Z

Discussion: https://apache-airflow.slack.com/archives/CCPRP7943/p1666718396900909

airflow/providers/common/sql/operators/sql.py

airflow/providers/common/sql/hooks/sql.py

kazanzhy · 2022-11-06T21:35:38Z

.kill() method is implemented only for Trino/Preso and Snowflake.
Cursors of these connectors return query_id which is used to stop running query.

For the other databases we could use something like:

select pg_terminate_backend(pid) 
from pg_stat_activity 
where query = '';

if we will save the latest queries as well as QueryId

potiuk

Cool. I like the way it is implemented now.

airflow/providers/snowflake/hooks/snowflake.py

airflow/providers/common/sql/operators/sql.py

airflow/providers/common/sql/hooks/sql.py

airflow/providers/common/sql/operators/sql.py

airflow/providers/common/sql/hooks/sql.py

dstandish · 2022-11-10T06:44:49Z

airflow/providers/common/sql/hooks/sql.py

@@ -264,6 +266,7 @@ def run(
                results = []
                for sql_statement in sql:
                    self._run_command(cur, sql_statement, parameters)
+                    self._update_query_ids(cur)


At this point, the query is done, so what is the utility of updating the query id list here?

You're right. I checked it for Snowflake.

Do you have access to any Trino and Presto instances to check it also for these DBs?

eladkal · 2022-11-11T12:48:04Z

@kazanzhy do we have something else to resolve or this PR is ready for final review?

dstandish · 2022-11-11T20:00:47Z

@kazanzhy do we have something else to resolve or this PR is ready for final review?

@eladkal why do you think we should merge this, given that it does not do what it tries to do?

The queries it says are running are not actually running. The queries it will kill in on_kill are not running. Before we add a new interface, shouldn't we start with code that does what it's supposed to do?

dstandish · 2022-11-11T20:23:48Z

airflow/providers/common/sql/hooks/sql.py

@@ -264,6 +266,7 @@ def run(
                results = []
                for sql_statement in sql:
                    self._run_command(cur, sql_statement, parameters)
+                    self._update_query_ids(cur)


i think that we need to address this issue, where the queries aren't actually running and therefore will not be killed, before adding this interface.

eladkal · 2022-11-11T20:24:13Z

@eladkal why do you think we should merge this, given that it does not do what it tries to do?

I didnt say we should merge this. I asked if its ready to review. I wasnt sure if the running queries problem was adressed.

ashb

Does this actually need changes in the hook to capture query IDs and kill them, couldn't that be solely on the operators?

From an architectural PoV it feels odd/out-of-place with the rest of the code base they were track what is effectively operator state (query ids) on the hook object rather than in the Operator object

dstandish · 2022-11-11T20:31:56Z

@eladkal why do you think we should merge this, given that it does not do what it tries to do?

I didnt say we should merge this. I asked if its ready to review. I wasnt sure if the running queries problem was adressed.

OK cool i thought it meant like... you were preparing to merge it :)

Cus I had commented after the recent updates like saying "this is still an issue" but then no response and the next thing i saw was "ok anything else"?

Maybe it was the word "final" which I interpreted to mean you thought it was basically done. But yeah definitely more work to do.

Anyway, thank you. Looking forward to getting this right... just ... want to get it right :)

kazanzhy · 2022-11-12T23:13:19Z

@dstandish I understand your concerns about this PR.
I must say that these changes are just returning deleted code in TrinoHook (df00436#diff-9a702240b5127432b36cea027e9846f74943e8c20c292f4301c182c6063db880) and refactor of SnowflakeHook to use DbApiHook.run() instead of SnowflakeHook.run().
So as we saw now there's no sense in using query_ids for on_kill(). But we can save them for other reasons.

I'm proposing the following steps when we will be sure that Trino cursor.execute for Trino is also synchronous.

remove on_kill() implementations for TrinoOperator and SQLExecuteQueryOperator
remove kill_query() implementations for TrinoHook, PrestoHook, SnowflakeHook, and DbApiHook.
rename this PR to Merge SnowflakeHook.run and TrinoHook.run to DbApiHook.run which will be the second part of Unify DbApiHook.run() method with the methods which override it #23971
close issue Generalize on_kill() for SQLExecuteQueryOperator #27314 because of impossibility to implement on_kill()

dstandish · 2022-11-16T21:54:26Z

I'm a little skeptical about Merge SnowflakeHook.run and TrinoHook.run to DbApiHook.run simply because there isn't clear utility apart from maybe consolidating shared code between two hooks. But we don't add to the base hook every time two subclasses share any code. And it's unclear to me anyway that this particular feature is actually useful / a good idea. But if you think there's a valuable interface to add, nothing wrong with opening the PR. Might make sense to do new PR for it though and just close this one. Then we can examine the interface from the simple perspective of... is this a feature we want to add to DbApiHook.

kazanzhy · 2022-11-18T00:33:29Z

This PR is really overloaded.
Let's discuss #27763 and #27762.
After that, we could come back here or just close this PR.
@dstandish

dstandish · 2022-11-18T07:14:11Z

This PR is really overloaded. Let's discuss #27763 and #27762. After that, we could come back here or just close this PR. @dstandish

great, thanks

kazanzhy · 2022-11-28T00:18:53Z

Can I conclude that this PR should be closed as well #27314 because of the absent possibility of canceling queries after their execution using a synchronous cursor?
@eladkal @dstandish

boring-cyborg bot added area:providers provider:common-sql provider:snowflake Issues related to Snowflake provider labels Nov 5, 2022

eladkal reviewed Nov 5, 2022

View reviewed changes

kazanzhy force-pushed the on_kill_actions_for_sql_provider branch 2 times, most recently from 23073ef to 622c37c Compare November 5, 2022 23:16

kazanzhy force-pushed the on_kill_actions_for_sql_provider branch from 622c37c to 9c0ad7c Compare November 6, 2022 21:37

kazanzhy marked this pull request as ready for review November 6, 2022 21:38

kazanzhy requested review from turbaszek, potiuk and mik-laj as code owners November 6, 2022 21:38

potiuk approved these changes Nov 6, 2022

View reviewed changes

potiuk requested a review from dstandish November 6, 2022 23:47

eladkal reviewed Nov 7, 2022

View reviewed changes

airflow/providers/snowflake/hooks/snowflake.py Show resolved Hide resolved

dstandish reviewed Nov 7, 2022

View reviewed changes

kazanzhy force-pushed the on_kill_actions_for_sql_provider branch 3 times, most recently from 1bbeb06 to 186036b Compare November 10, 2022 00:30

dstandish reviewed Nov 10, 2022

View reviewed changes

kazanzhy force-pushed the on_kill_actions_for_sql_provider branch from 186036b to 928a2cc Compare November 10, 2022 23:09

dstandish requested a review from potiuk November 11, 2022 20:01

dstandish requested changes Nov 11, 2022

View reviewed changes

ashb reviewed Nov 11, 2022

View reviewed changes

kazanzhy force-pushed the on_kill_actions_for_sql_provider branch from 928a2cc to 20538cf Compare November 13, 2022 21:52

potiuk requested a review from dstandish November 16, 2022 20:52

kazanzhy force-pushed the on_kill_actions_for_sql_provider branch from 20538cf to 68602b9 Compare November 17, 2022 23:45

Add on_kill logic for Presto, Trino, Snowflake

ec760ac

kazanzhy force-pushed the on_kill_actions_for_sql_provider branch from 68602b9 to ec760ac Compare November 18, 2022 00:22

kazanzhy marked this pull request as draft November 18, 2022 00:33

kazanzhy closed this Nov 30, 2022

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implement on_kill() for SQLExecuteQueryOperator #27514

Implement on_kill() for SQLExecuteQueryOperator #27514

kazanzhy commented Nov 5, 2022

kazanzhy commented Nov 6, 2022

potiuk left a comment

dstandish Nov 10, 2022

kazanzhy Nov 12, 2022

kazanzhy Nov 12, 2022

eladkal commented Nov 11, 2022

dstandish commented Nov 11, 2022

dstandish Nov 11, 2022

eladkal commented Nov 11, 2022

ashb left a comment

dstandish commented Nov 11, 2022 •

edited

Loading

kazanzhy commented Nov 12, 2022

dstandish commented Nov 16, 2022

kazanzhy commented Nov 18, 2022

dstandish commented Nov 18, 2022

kazanzhy commented Nov 28, 2022

Implement on_kill() for SQLExecuteQueryOperator #27514

Implement on_kill() for SQLExecuteQueryOperator #27514

Conversation

kazanzhy commented Nov 5, 2022

kazanzhy commented Nov 6, 2022

potiuk left a comment

Choose a reason for hiding this comment

dstandish Nov 10, 2022

Choose a reason for hiding this comment

kazanzhy Nov 12, 2022

Choose a reason for hiding this comment

kazanzhy Nov 12, 2022

Choose a reason for hiding this comment

eladkal commented Nov 11, 2022

dstandish commented Nov 11, 2022

dstandish Nov 11, 2022

Choose a reason for hiding this comment

eladkal commented Nov 11, 2022

ashb left a comment

Choose a reason for hiding this comment

dstandish commented Nov 11, 2022 • edited Loading

kazanzhy commented Nov 12, 2022

dstandish commented Nov 16, 2022

kazanzhy commented Nov 18, 2022

dstandish commented Nov 18, 2022

kazanzhy commented Nov 28, 2022

dstandish commented Nov 11, 2022 •

edited

Loading