Fixing the behaviours of SQL Hooks and Operators finally #27912

potiuk · 2022-11-25T11:12:58Z

This is the more bold fix to common.sql and related providers. It implements more comprehensive tests and describes the intended behaviours in much more explicit way, also it simplifies the notion of "scalar" behaviour by introducing a new method to check on when scalar return is expected.

Comprehensive suite of tests have been added (the original tests were converted to Pytest test and parameterized was used to test all the combinations of parameters and returned values.

^ Add meaningful description above

Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in newsfragments.

potiuk · 2022-11-25T11:13:27Z

cc: @alexott

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

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

airflow/providers/databricks/hooks/databricks_sql.py

potiuk · 2022-11-25T11:37:32Z

cc: @alexott @kazanzhy - I have two important doubts about processing of ; (and sql stripping in general).

(Less important if it works) Stripping ";"

I followed original implementations I think and I left I think intact original behviour of stripping in both implementationations.

Look at the parameterized tests - and the "DBApiHook" behaves now differently than Databricks Hook in this regard - the Databricks hook always splits the ;

https://github.com/apache/airflow/pull/27912/files#diff-db2c1495ba766a1bad71474f36caeb3e8b4228522571e3229c1c64cede967027R99

Where the standard DBApiHook does not:

https://github.com/apache/airflow/pull/27912/files#diff-718d368660d07715b383dd4d94344bc41b835189da39bb95fcecaa3a423e6e47R110

The behaviour of cursor and connection opening/closing is different in both cases:

DBAPiHook opens connection, cursors - and iterates over multiple queries with the same cursor (runs execute command for each) https://github.com/apache/airflow/pull/27912/files#diff-6e1b2f961cb951d05d66d2d814ef5f6d8f8bf8f43c40fb5d40e27a031fed8dd7R288
Databricks Hook opens connection, but then creates a cursor for each sql statement separately:
https://github.com/apache/airflow/pull/27912/files#diff-f1c71c8b869a687043b306f047f592e3638b3d74d9ec3c972360cf6c9a7956a9R182

The comiit strategy is different between the two (and does not work in Databricks IMHO) after last statement in Databricks Hook. The DBApi hook performs commit after all statements when autocommit is off https://github.com/apache/airflow/pull/27912/files#diff-6e1b2f961cb951d05d66d2d814ef5f6d8f8bf8f43c40fb5d40e27a031fed8dd7R307 but the DAtabricks one has very weird behaviour - when autocommit is off, I think each statement result will be rolled back automatically - so it basically makes no sense to set autocommit to off, because there is no way to manually commit (COMMIT is a separate statement and previous statement is rolled back when it is executed).

https://peps.python.org/pep-0249/#Connection.close

Note that closing a connection without committing the changes first will cause an implicit rollback to be performed.

I am no sure what was the reasoning behind those differences - but from what I see those differences now between those two "run" methods. Question is - whether we have good reasons to behave differently for those cases?

airflow/providers/databricks/hooks/databricks_sql.py

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

airflow/providers/databricks/hooks/databricks_sql.py

tests/providers/databricks/operators/test_databricks_sql.py

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

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

airflow/providers/databricks/hooks/databricks_sql.py

potiuk · 2022-11-25T16:01:44Z

OK. There was some permission issue - but I think this one shoudl be ready to go. I added some descriptions and explanations about the behaviours and 🤞 it will pass all the docs and tests

alexott · 2022-11-25T16:05:56Z

Let me do a checkout in my environment and test.

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

vincbeck

So much work on that PR! Nothing to say on my side, very solid! Thanks for adding all these tests (on top of the other changes), it adds a lot of robustness!

potiuk · 2022-11-25T19:47:03Z

So much work on that PR! Nothing to say on my side, very solid! Thanks for adding all these tests (on top of the other changes), it adds a lot of robustness!

Indeed. I added already few small changes/clarifications after the initial push and adding all tests and it felt soooooooo reassuring to see green checkmark for all tests after my changes.

potiuk · 2022-11-25T19:53:15Z

BTW. I REALLY ❤️ Pytest style of tests vs. the "classic" unit.TestCase.

Partucularrly fixtures - when you get a grasp of them, parameterization and the integration of those running tests within IDEs are just great.

potiuk · 2022-11-25T22:18:17Z

All looks good :)

This is the more bold fix to common.sql and related providers. It implements more comprehensive tests and describes the intended behaviours in much more explicit way, also it simplifies the notion of "scalar" behaviour by introducing a new method to check on when scalar return is expected. Comprehensive suite of tests have been added (the original tests were converted to Pytest test and parameterized was used to test all the combinations of parameters and returned values.

dstandish · 2022-11-26T06:55:38Z

It's holiday in us and I'm traveling, won't be able to review until next week

alexott

thank you again!

alexott · 2022-11-26T09:11:27Z

airflow/providers/common/sql/CHANGELOG.rst

@@ -32,7 +32,9 @@ This release fixes a few errors that were introduced in common.sql operator whil
 * ``_process_output`` method in ``SQLExecuteQueryOperator`` has now consistent semantics and typing, it
  can also modify the returned (and stored in XCom) values in the operators that derive from the
  ``SQLExecuteQueryOperator``).
-* last description of the cursor whether to return scalar values are now stored in DBApiHook
+* descriptions of all returned results are stored as descriptions property in the DBApiHook


(could be done later) maybe it would be a bit easier to read if descriptions and last_description are formatted as code?

Yep. We can update it later :)

alexott · 2022-11-26T09:29:31Z

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

@@ -190,11 +189,17 @@ class SQLExecuteQueryOperator(BaseSQLOperator):
    Executes SQL code in a specific database
    :param sql: the SQL code or string pointing to a template file to be executed (templated).
    File must have a '.sql' extensions.
+
+    When implementing a specific Operator, you can also implement `_process_output` method in the


should we use double backquotes in docstrings?

Likely :). I think it will work anyway in docstring (we'll see in the API docs) but yeah - we can fix it as follow-up

potiuk · 2022-11-26T09:53:59Z

BTW. Release of rc3 is coming .

BillCM · 2022-11-28T17:44:10Z

Interesting note: the 3.3.0 DatabricksSQLOperator had a parameter for control writing to xcom. It appears to be removed in 4.0.0. This is causing the dump of a lot of PII into XCOM for my use case.

alexott · 2022-11-28T17:49:52Z

@BillCM do_xcom_push should be still supported as it's argument to the BaseOperator - can you check it?

The change apache#27912 fixed and unified behaviour of DBApiHooks across the board, but it missed two places where sql was mis-used and overridden in exasol and snowflake hooks. The check for "sql" type did not use the original sql parameter value but the one that was overridden later in the run method implementation. The fix is the same as applied in Databricks Hook and DBAPI generic run methods - using consistent typing and separate variable to convert the sql string into sql list. Related: astronomer/astro-sdk#1324

…27997) The change #27912 fixed and unified behaviour of DBApiHooks across the board, but it missed two places where sql was mis-used and overridden in exasol and snowflake hooks. The check for "sql" type did not use the original sql parameter value but the one that was overridden later in the run method implementation. The fix is the same as applied in Databricks Hook and DBAPI generic run methods - using consistent typing and separate variable to convert the sql string into sql list. Related: astronomer/astro-sdk#1324

…pache#27997) The change apache#27912 fixed and unified behaviour of DBApiHooks across the board, but it missed two places where sql was mis-used and overridden in exasol and snowflake hooks. The check for "sql" type did not use the original sql parameter value but the one that was overridden later in the run method implementation. The fix is the same as applied in Databricks Hook and DBAPI generic run methods - using consistent typing and separate variable to convert the sql string into sql list. Related: astronomer/astro-sdk#1324

potiuk requested review from eladkal, turbaszek and mik-laj as code owners November 25, 2022 11:12

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

potiuk requested review from uranusjr, dstandish and ephraimbuddy November 25, 2022 11:13