From 1dc6553d7b863868bccd932903592992d84757aa Mon Sep 17 00:00:00 2001 From: Joffrey Bienvenu Date: Wed, 13 Dec 2023 12:43:38 +0100 Subject: [PATCH 01/22] feat: Databricks sql hook returns json-serializable namedtuple NamedTuple can be used as a tuple or a dict. databricks.Row object is a dict-like object. This should break fewer use-cases than switching directly to a tuple/list object. --- .../databricks/hooks/databricks_sql.py | 11 +++++-- .../databricks/hooks/test_databricks_sql.py | 33 ++++++++++++------- 2 files changed, 30 insertions(+), 14 deletions(-) diff --git a/airflow/providers/databricks/hooks/databricks_sql.py b/airflow/providers/databricks/hooks/databricks_sql.py index dc728c5ed70cf..66f64ae4ca8ed 100644 --- a/airflow/providers/databricks/hooks/databricks_sql.py +++ b/airflow/providers/databricks/hooks/databricks_sql.py @@ -16,6 +16,7 @@ # under the License. from __future__ import annotations +from collections import namedtuple from contextlib import closing from copy import copy from typing import TYPE_CHECKING, Any, Callable, Iterable, Mapping, TypeVar, overload @@ -243,11 +244,15 @@ def run( @staticmethod def _make_serializable(result): - """Transform the databricks Row objects into JSON-serializable lists.""" + """Transform the databricks Row objects into JSON-serializable namedtuple.""" + columns: list[str] | None = None if isinstance(result, list): - return [list(row) for row in result] + columns = result[0].__fields__ + row_object = namedtuple("Row", columns) + return [row_object(*row) for row in result] elif isinstance(result, Row): - return list(result) + columns = result.__fields__ + return namedtuple("Row", columns)(*result) return result def bulk_dump(self, table, tmp_file): diff --git a/tests/providers/databricks/hooks/test_databricks_sql.py b/tests/providers/databricks/hooks/test_databricks_sql.py index 64cd0b9c06e2f..a1e3f69a40057 100644 --- a/tests/providers/databricks/hooks/test_databricks_sql.py +++ b/tests/providers/databricks/hooks/test_databricks_sql.py @@ -18,6 +18,7 @@ # from __future__ import annotations +from collections import namedtuple from unittest import mock from unittest.mock import patch @@ -58,6 +59,10 @@ def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: return [(field,) for field in fields] +# Serializable Row object similar to the one returned by the Hook +SerializableRow = namedtuple("Row", ["id", "value"]) + + @pytest.mark.parametrize( "return_last, split_statements, sql, cursor_calls," "cursor_descriptions, cursor_results, hook_descriptions, hook_results, ", @@ -70,7 +75,7 @@ def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: [["id", "value"]], ([Row(id=1, value=2), Row(id=11, value=12)],), [[("id",), ("value",)]], - [[1, 2], [11, 12]], + [SerializableRow(1, 2), SerializableRow(11, 12)], id="The return_last set and no split statements set on single query in string", ), pytest.param( @@ -81,7 +86,7 @@ def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: [["id", "value"]], ([Row(id=1, value=2), Row(id=11, value=12)],), [[("id",), ("value",)]], - [[1, 2], [11, 12]], + [SerializableRow(1, 2), SerializableRow(11, 12)], id="The return_last not set and no split statements set on single query in string", ), pytest.param( @@ -92,7 +97,7 @@ def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: [["id", "value"]], ([Row(id=1, value=2), Row(id=11, value=12)],), [[("id",), ("value",)]], - [[1, 2], [11, 12]], + [SerializableRow(1, 2), SerializableRow(11, 12)], id="The return_last set and split statements set on single query in string", ), pytest.param( @@ -103,7 +108,7 @@ def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: [["id", "value"]], ([Row(id=1, value=2), Row(id=11, value=12)],), [[("id",), ("value",)]], - [[[1, 2], [11, 12]]], + [[SerializableRow(1, 2), SerializableRow(11, 12)]], id="The return_last not set and split statements set on single query in string", ), pytest.param( @@ -114,7 +119,7 @@ def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: [["id", "value"], ["id2", "value2"]], ([Row(id=1, value=2), Row(id=11, value=12)], [Row(id=3, value=4), Row(id=13, value=14)]), [[("id2",), ("value2",)]], - [[3, 4], [13, 14]], + [SerializableRow(3, 4), SerializableRow(13, 14)], id="The return_last set and split statements set on multiple queries in string", ), pytest.param( @@ -125,7 +130,10 @@ def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: [["id", "value"], ["id2", "value2"]], ([Row(id=1, value=2), Row(id=11, value=12)], [Row(id=3, value=4), Row(id=13, value=14)]), [[("id",), ("value",)], [("id2",), ("value2",)]], - [[[1, 2], [11, 12]], [[3, 4], [13, 14]]], + [ + [SerializableRow(1, 2), SerializableRow(11, 12)], + [SerializableRow(3, 4), SerializableRow(13, 14)], + ], id="The return_last not set and split statements set on multiple queries in string", ), pytest.param( @@ -136,7 +144,7 @@ def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: [["id", "value"]], ([Row(id=1, value=2), Row(id=11, value=12)],), [[("id",), ("value",)]], - [[[1, 2], [11, 12]]], + [[SerializableRow(1, 2), SerializableRow(11, 12)]], id="The return_last set on single query in list", ), pytest.param( @@ -147,7 +155,7 @@ def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: [["id", "value"]], ([Row(id=1, value=2), Row(id=11, value=12)],), [[("id",), ("value",)]], - [[[1, 2], [11, 12]]], + [[SerializableRow(1, 2), SerializableRow(11, 12)]], id="The return_last not set on single query in list", ), pytest.param( @@ -158,7 +166,7 @@ def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: [["id", "value"], ["id2", "value2"]], ([Row(id=1, value=2), Row(id=11, value=12)], [Row(id=3, value=4), Row(id=13, value=14)]), [[("id2",), ("value2",)]], - [[3, 4], [13, 14]], + [SerializableRow(3, 4), SerializableRow(13, 14)], id="The return_last set on multiple queries in list", ), pytest.param( @@ -169,7 +177,10 @@ def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: [["id", "value"], ["id2", "value2"]], ([Row(id=1, value=2), Row(id=11, value=12)], [Row(id=3, value=4), Row(id=13, value=14)]), [[("id",), ("value",)], [("id2",), ("value2",)]], - [[[1, 2], [11, 12]], [[3, 4], [13, 14]]], + [ + [SerializableRow(1, 2), SerializableRow(11, 12)], + [SerializableRow(3, 4), SerializableRow(13, 14)], + ], id="The return_last not set on multiple queries not set", ), pytest.param( @@ -180,7 +191,7 @@ def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: [["id", "value"]], (Row(id=1, value=2),), [[("id",), ("value",)]], - [1, 2], + SerializableRow(1, 2), id="The return_last set and no split statements set on single query in string", ), ], From 6a35e30aaecdbd1e0e837ed1d1bac20f286c12f1 Mon Sep 17 00:00:00 2001 From: Joffrey Bienvenu Date: Wed, 13 Dec 2023 13:02:55 +0100 Subject: [PATCH 02/22] fix: Ignore mypy namedtuple has to be named "Row", to be similar to the one returned by the Hook. But its variable name should not be Row in this context, to differentiate the real "databricks.Row" object from the namedtuple one. --- tests/providers/databricks/hooks/test_databricks_sql.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/providers/databricks/hooks/test_databricks_sql.py b/tests/providers/databricks/hooks/test_databricks_sql.py index a1e3f69a40057..4d4fe62e757b1 100644 --- a/tests/providers/databricks/hooks/test_databricks_sql.py +++ b/tests/providers/databricks/hooks/test_databricks_sql.py @@ -60,7 +60,7 @@ def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: # Serializable Row object similar to the one returned by the Hook -SerializableRow = namedtuple("Row", ["id", "value"]) +SerializableRow = namedtuple("Row", ["id", "value"]) # type: ignore[name-match] @pytest.mark.parametrize( From e09a34b2a101002996ee3292d6543c7d1ee8802e Mon Sep 17 00:00:00 2001 From: Joffrey Bienvenu Date: Wed, 13 Dec 2023 21:30:36 +0100 Subject: [PATCH 03/22] fix: correctly describe serialization --- airflow/providers/common/sql/hooks/sql.py | 2 +- airflow/providers/databricks/hooks/databricks_sql.py | 2 +- airflow/providers/odbc/hooks/odbc.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/airflow/providers/common/sql/hooks/sql.py b/airflow/providers/common/sql/hooks/sql.py index bb85dedc1cdbd..57c3e37e6cc9f 100644 --- a/airflow/providers/common/sql/hooks/sql.py +++ b/airflow/providers/common/sql/hooks/sql.py @@ -124,7 +124,7 @@ class DbApiHook(BaseHook): When subclassing, maintainers can override the `_make_serializable` method: This method transforms the result of the handler method (typically `cursor.fetchall()`) into - JSON-serializable objects. Most of the time, the underlying SQL library already returns tuples from + serializable objects. Most of the time, the underlying SQL library already returns tuples from its cursor, and the `_make_serializable` method can be ignored. :param schema: Optional DB schema that overrides the schema specified in the connection. Make sure that diff --git a/airflow/providers/databricks/hooks/databricks_sql.py b/airflow/providers/databricks/hooks/databricks_sql.py index 66f64ae4ca8ed..a67c6bb6a9526 100644 --- a/airflow/providers/databricks/hooks/databricks_sql.py +++ b/airflow/providers/databricks/hooks/databricks_sql.py @@ -244,7 +244,7 @@ def run( @staticmethod def _make_serializable(result): - """Transform the databricks Row objects into JSON-serializable namedtuple.""" + """Transform the databricks Row objects into serializable namedtuple.""" columns: list[str] | None = None if isinstance(result, list): columns = result[0].__fields__ diff --git a/airflow/providers/odbc/hooks/odbc.py b/airflow/providers/odbc/hooks/odbc.py index d84933dc29441..799e9f835c04f 100644 --- a/airflow/providers/odbc/hooks/odbc.py +++ b/airflow/providers/odbc/hooks/odbc.py @@ -214,7 +214,7 @@ def get_sqlalchemy_connection( @staticmethod def _make_serializable(result: list[pyodbc.Row] | pyodbc.Row | None) -> list[NamedTuple] | None: - """Transform the pyodbc.Row objects returned from an SQL command into JSON-serializable NamedTuple.""" + """Transform the pyodbc.Row objects returned from an SQL command into serializable NamedTuple.""" # Below ignored lines respect NamedTuple docstring, but mypy do not support dynamically # instantiated Namedtuple, and will never do: https://github.com/python/mypy/issues/848 columns: list[tuple[str, type]] | None = None From b11f88380eb810aef3e889427610baacf0fc3245 Mon Sep 17 00:00:00 2001 From: Joffrey Bienvenu Date: Wed, 13 Dec 2023 22:15:36 +0100 Subject: [PATCH 04/22] feat: disable the return of serializable namedtuple by default --- .../databricks/hooks/databricks_sql.py | 36 ++++++++++++----- .../databricks/hooks/test_databricks_sql.py | 40 ++++++++++++------- 2 files changed, 51 insertions(+), 25 deletions(-) diff --git a/airflow/providers/databricks/hooks/databricks_sql.py b/airflow/providers/databricks/hooks/databricks_sql.py index a67c6bb6a9526..74dd1df6e0473 100644 --- a/airflow/providers/databricks/hooks/databricks_sql.py +++ b/airflow/providers/databricks/hooks/databricks_sql.py @@ -16,6 +16,7 @@ # under the License. from __future__ import annotations +import warnings from collections import namedtuple from contextlib import closing from copy import copy @@ -24,7 +25,7 @@ from databricks import sql # type: ignore[attr-defined] from databricks.sql.types import Row -from airflow.exceptions import AirflowException +from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning from airflow.providers.common.sql.hooks.sql import DbApiHook, return_single_query_results from airflow.providers.databricks.hooks.databricks_base import BaseDatabricksHook @@ -53,6 +54,8 @@ class DatabricksSqlHook(BaseDatabricksHook, DbApiHook): on every request :param catalog: An optional initial catalog to use. Requires DBR version 9.0+ :param schema: An optional initial schema to use. Requires DBR version 9.0+ + :param return_serializable: Return a namedtuple "Row" object instead of a `databricks.sql.Row` object. + (default: False). In a future version of this provider, this will become True bu default. :param kwargs: Additional parameters internal to Databricks SQL Connector parameters """ @@ -69,6 +72,7 @@ def __init__( catalog: str | None = None, schema: str | None = None, caller: str = "DatabricksSqlHook", + return_serializable: bool = False, **kwargs, ) -> None: super().__init__(databricks_conn_id, caller=caller) @@ -81,8 +85,18 @@ def __init__( self.http_headers = http_headers self.catalog = catalog self.schema = schema + self.return_serializable = return_serializable self.additional_params = kwargs + if not self.return_serializable: + warnings.warn( + """Returning a raw `databricks.sql.Row` object is deprecated, and will be removed in a future + release of the databricks provider. Use `return_serializable=True` instead to receive a + serializable "Row" namedtuple.""", + AirflowProviderDeprecationWarning, + stacklevel=2, + ) + def _get_extra_config(self) -> dict[str, Any | None]: extra_params = copy(self.databricks_conn.extra_dejson) for arg in ["http_path", "session_configuration", *self.extra_parameters]: @@ -242,17 +256,17 @@ def run( else: return results - @staticmethod - def _make_serializable(result): + def _make_serializable(self, result): """Transform the databricks Row objects into serializable namedtuple.""" - columns: list[str] | None = None - if isinstance(result, list): - columns = result[0].__fields__ - row_object = namedtuple("Row", columns) - return [row_object(*row) for row in result] - elif isinstance(result, Row): - columns = result.__fields__ - return namedtuple("Row", columns)(*result) + if self.return_serializable: + columns: list[str] | None = None + if isinstance(result, list): + columns = result[0].__fields__ + row_object = namedtuple("Row", columns) + return [row_object(*row) for row in result] + elif isinstance(result, Row): + columns = result.__fields__ + return namedtuple("Row", columns)(*result) return result def bulk_dump(self, table, tmp_file): diff --git a/tests/providers/databricks/hooks/test_databricks_sql.py b/tests/providers/databricks/hooks/test_databricks_sql.py index 4d4fe62e757b1..3a1516f9ece83 100644 --- a/tests/providers/databricks/hooks/test_databricks_sql.py +++ b/tests/providers/databricks/hooks/test_databricks_sql.py @@ -64,7 +64,7 @@ def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: @pytest.mark.parametrize( - "return_last, split_statements, sql, cursor_calls," + "return_last, split_statements, sql, cursor_calls, return_serializable," "cursor_descriptions, cursor_results, hook_descriptions, hook_results, ", [ pytest.param( @@ -72,10 +72,11 @@ def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: False, "select * from test.test", ["select * from test.test"], + False, [["id", "value"]], ([Row(id=1, value=2), Row(id=11, value=12)],), [[("id",), ("value",)]], - [SerializableRow(1, 2), SerializableRow(11, 12)], + [Row(id=1, value=2), Row(id=11, value=12)], id="The return_last set and no split statements set on single query in string", ), pytest.param( @@ -83,10 +84,11 @@ def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: False, "select * from test.test;", ["select * from test.test"], + False, [["id", "value"]], ([Row(id=1, value=2), Row(id=11, value=12)],), [[("id",), ("value",)]], - [SerializableRow(1, 2), SerializableRow(11, 12)], + [Row(id=1, value=2), Row(id=11, value=12)], id="The return_last not set and no split statements set on single query in string", ), pytest.param( @@ -94,10 +96,11 @@ def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: True, "select * from test.test;", ["select * from test.test"], + False, [["id", "value"]], ([Row(id=1, value=2), Row(id=11, value=12)],), [[("id",), ("value",)]], - [SerializableRow(1, 2), SerializableRow(11, 12)], + [Row(id=1, value=2), Row(id=11, value=12)], id="The return_last set and split statements set on single query in string", ), pytest.param( @@ -105,10 +108,11 @@ def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: True, "select * from test.test;", ["select * from test.test"], + False, [["id", "value"]], ([Row(id=1, value=2), Row(id=11, value=12)],), [[("id",), ("value",)]], - [[SerializableRow(1, 2), SerializableRow(11, 12)]], + [[Row(id=1, value=2), Row(id=11, value=12)]], id="The return_last not set and split statements set on single query in string", ), pytest.param( @@ -116,10 +120,11 @@ def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: True, "select * from test.test;select * from test.test2;", ["select * from test.test", "select * from test.test2"], + False, [["id", "value"], ["id2", "value2"]], ([Row(id=1, value=2), Row(id=11, value=12)], [Row(id=3, value=4), Row(id=13, value=14)]), [[("id2",), ("value2",)]], - [SerializableRow(3, 4), SerializableRow(13, 14)], + [Row(id=3, value=4), Row(id=13, value=14)], id="The return_last set and split statements set on multiple queries in string", ), pytest.param( @@ -127,12 +132,13 @@ def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: True, "select * from test.test;select * from test.test2;", ["select * from test.test", "select * from test.test2"], + False, [["id", "value"], ["id2", "value2"]], ([Row(id=1, value=2), Row(id=11, value=12)], [Row(id=3, value=4), Row(id=13, value=14)]), [[("id",), ("value",)], [("id2",), ("value2",)]], [ - [SerializableRow(1, 2), SerializableRow(11, 12)], - [SerializableRow(3, 4), SerializableRow(13, 14)], + [Row(id=1, value=2), Row(id=11, value=12)], + [Row(id=3, value=4), Row(id=13, value=14)], ], id="The return_last not set and split statements set on multiple queries in string", ), @@ -141,10 +147,11 @@ def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: True, ["select * from test.test;"], ["select * from test.test"], + False, [["id", "value"]], ([Row(id=1, value=2), Row(id=11, value=12)],), [[("id",), ("value",)]], - [[SerializableRow(1, 2), SerializableRow(11, 12)]], + [[Row(id=1, value=2), Row(id=11, value=12)]], id="The return_last set on single query in list", ), pytest.param( @@ -152,10 +159,11 @@ def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: True, ["select * from test.test;"], ["select * from test.test"], + False, [["id", "value"]], ([Row(id=1, value=2), Row(id=11, value=12)],), [[("id",), ("value",)]], - [[SerializableRow(1, 2), SerializableRow(11, 12)]], + [[Row(id=1, value=2), Row(id=11, value=12)]], id="The return_last not set on single query in list", ), pytest.param( @@ -163,10 +171,11 @@ def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: True, "select * from test.test;select * from test.test2;", ["select * from test.test", "select * from test.test2"], + False, [["id", "value"], ["id2", "value2"]], ([Row(id=1, value=2), Row(id=11, value=12)], [Row(id=3, value=4), Row(id=13, value=14)]), [[("id2",), ("value2",)]], - [SerializableRow(3, 4), SerializableRow(13, 14)], + [Row(id=3, value=4), Row(id=13, value=14)], id="The return_last set on multiple queries in list", ), pytest.param( @@ -174,12 +183,13 @@ def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: True, "select * from test.test;select * from test.test2;", ["select * from test.test", "select * from test.test2"], + False, [["id", "value"], ["id2", "value2"]], ([Row(id=1, value=2), Row(id=11, value=12)], [Row(id=3, value=4), Row(id=13, value=14)]), [[("id",), ("value",)], [("id2",), ("value2",)]], [ - [SerializableRow(1, 2), SerializableRow(11, 12)], - [SerializableRow(3, 4), SerializableRow(13, 14)], + [Row(id=1, value=2), Row(id=11, value=12)], + [Row(id=3, value=4), Row(id=13, value=14)], ], id="The return_last not set on multiple queries not set", ), @@ -188,6 +198,7 @@ def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: False, "select * from test.test", ["select * from test.test"], + True, [["id", "value"]], (Row(id=1, value=2),), [[("id",), ("value",)]], @@ -197,11 +208,11 @@ def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: ], ) def test_query( - databricks_hook, return_last, split_statements, sql, cursor_calls, + return_serializable, cursor_descriptions, cursor_results, hook_descriptions, @@ -238,6 +249,7 @@ def test_query( cursors.append(cur) connections.append(conn) mock_conn.side_effect = connections + databricks_hook = DatabricksSqlHook(sql_endpoint_name="Test", return_serializable=return_serializable) results = databricks_hook.run( sql=sql, handler=fetch_all_handler, return_last=return_last, split_statements=split_statements ) From 837d4608b70e101d4c2741a9c3db2da0206c3282 Mon Sep 17 00:00:00 2001 From: Joffrey Bienvenu Date: Wed, 13 Dec 2023 22:22:07 +0100 Subject: [PATCH 05/22] fix: make DatabricksSqlOperator return serializable object by default This operator was already returning a serializable object since #31780. But the PR #31780 was reverted in favor of a serialization at the Hook level. To avoid regression, the DatabricksSqlOperator has to return a serializable --- airflow/providers/databricks/operators/databricks_sql.py | 1 + 1 file changed, 1 insertion(+) diff --git a/airflow/providers/databricks/operators/databricks_sql.py b/airflow/providers/databricks/operators/databricks_sql.py index a03cfa729c714..d5349d38dcbc9 100644 --- a/airflow/providers/databricks/operators/databricks_sql.py +++ b/airflow/providers/databricks/operators/databricks_sql.py @@ -113,6 +113,7 @@ def get_db_hook(self) -> DatabricksSqlHook: "catalog": self.catalog, "schema": self.schema, "caller": "DatabricksSqlOperator", + "return_serializable": True, **self.client_parameters, **self.hook_params, } From c653798ad3180bbcfada20b0b440c519593882ed Mon Sep 17 00:00:00 2001 From: Joffrey Bienvenu Date: Wed, 13 Dec 2023 22:38:04 +0100 Subject: [PATCH 06/22] fix: add `return_serializable` to databricks operator tests --- tests/providers/databricks/operators/test_databricks_sql.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/providers/databricks/operators/test_databricks_sql.py b/tests/providers/databricks/operators/test_databricks_sql.py index e7885740cfca9..d791985d54b8d 100644 --- a/tests/providers/databricks/operators/test_databricks_sql.py +++ b/tests/providers/databricks/operators/test_databricks_sql.py @@ -133,6 +133,7 @@ def test_exec_success(sql, return_last, split_statement, hook_results, hook_desc db_mock_class.assert_called_once_with( DEFAULT_CONN_ID, http_path=None, + return_serializable=True, session_configuration=None, sql_endpoint_name=None, http_headers=None, @@ -276,6 +277,7 @@ def test_exec_write_file( db_mock_class.assert_called_once_with( DEFAULT_CONN_ID, http_path=None, + return_serializable=True, session_configuration=None, sql_endpoint_name=None, http_headers=None, From 7c00732cc726453de77b7fc00e1d6c3964c6a79b Mon Sep 17 00:00:00 2001 From: Joffrey Bienvenu Date: Wed, 13 Dec 2023 23:01:33 +0100 Subject: [PATCH 07/22] fix: docstring spellchecks --- airflow/providers/databricks/hooks/databricks_sql.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/airflow/providers/databricks/hooks/databricks_sql.py b/airflow/providers/databricks/hooks/databricks_sql.py index 74dd1df6e0473..6fcdfc46b8ace 100644 --- a/airflow/providers/databricks/hooks/databricks_sql.py +++ b/airflow/providers/databricks/hooks/databricks_sql.py @@ -54,8 +54,8 @@ class DatabricksSqlHook(BaseDatabricksHook, DbApiHook): on every request :param catalog: An optional initial catalog to use. Requires DBR version 9.0+ :param schema: An optional initial schema to use. Requires DBR version 9.0+ - :param return_serializable: Return a namedtuple "Row" object instead of a `databricks.sql.Row` object. - (default: False). In a future version of this provider, this will become True bu default. + :param return_serializable: Return a ``namedtuple`` "Row" object instead of a ``databricks.sql.Row`` + object. (default: False). In a future version of the provider, this will become True by default. :param kwargs: Additional parameters internal to Databricks SQL Connector parameters """ From eceec0b15b74a83b1fb98186e20fd773b0faab7c Mon Sep 17 00:00:00 2001 From: Joffrey Bienvenu Date: Thu, 14 Dec 2023 21:43:50 +0100 Subject: [PATCH 08/22] fix: rewrite docstring and warning message --- airflow/providers/databricks/hooks/databricks_sql.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/airflow/providers/databricks/hooks/databricks_sql.py b/airflow/providers/databricks/hooks/databricks_sql.py index 6fcdfc46b8ace..7c4081c7c620d 100644 --- a/airflow/providers/databricks/hooks/databricks_sql.py +++ b/airflow/providers/databricks/hooks/databricks_sql.py @@ -54,8 +54,10 @@ class DatabricksSqlHook(BaseDatabricksHook, DbApiHook): on every request :param catalog: An optional initial catalog to use. Requires DBR version 9.0+ :param schema: An optional initial schema to use. Requires DBR version 9.0+ - :param return_serializable: Return a ``namedtuple`` "Row" object instead of a ``databricks.sql.Row`` - object. (default: False). In a future version of the provider, this will become True by default. + :param return_serializable: Return a serializable``namedtuple`` object instead of a ``databricks.sql.Row`` + object. Default to False. In a future release of the provider, this will become True by default. + This parameter ensures backward-compatibility during the transition phase, and will be removed too + in a future release. :param kwargs: Additional parameters internal to Databricks SQL Connector parameters """ @@ -90,9 +92,9 @@ def __init__( if not self.return_serializable: warnings.warn( - """Returning a raw `databricks.sql.Row` object is deprecated, and will be removed in a future - release of the databricks provider. Use `return_serializable=True` instead to receive a - serializable "Row" namedtuple.""", + """Returning a raw `databricks.sql.Row` object is deprecated. A namedtuple will be + returned instead in a future release of the databricks provider. Set + `return_serializable=True` to enable this behavior.""", AirflowProviderDeprecationWarning, stacklevel=2, ) From 67d1b903573ffe75f4f2a525d2fa5aed139252b2 Mon Sep 17 00:00:00 2001 From: Joffrey Bienvenu Date: Thu, 14 Dec 2023 22:22:45 +0100 Subject: [PATCH 09/22] fix: docstring spellchecks --- airflow/providers/databricks/hooks/databricks_sql.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/airflow/providers/databricks/hooks/databricks_sql.py b/airflow/providers/databricks/hooks/databricks_sql.py index 7c4081c7c620d..5ad4b0306bbd1 100644 --- a/airflow/providers/databricks/hooks/databricks_sql.py +++ b/airflow/providers/databricks/hooks/databricks_sql.py @@ -54,10 +54,10 @@ class DatabricksSqlHook(BaseDatabricksHook, DbApiHook): on every request :param catalog: An optional initial catalog to use. Requires DBR version 9.0+ :param schema: An optional initial schema to use. Requires DBR version 9.0+ - :param return_serializable: Return a serializable``namedtuple`` object instead of a ``databricks.sql.Row`` - object. Default to False. In a future release of the provider, this will become True by default. - This parameter ensures backward-compatibility during the transition phase, and will be removed too - in a future release. + :param return_serializable: Return a serializable ``namedtuple`` object instead of a + ``databricks.sql.Row`` object. Default to False. In a future release of the provider, this will + become True by default. This parameter ensures backward-compatibility during the transition phase, + and will be removed too in a future release. :param kwargs: Additional parameters internal to Databricks SQL Connector parameters """ From 1e84a757fe12cacf35d2d2a70ac1ba6fd1ad485f Mon Sep 17 00:00:00 2001 From: Joffrey Bienvenu Date: Sat, 16 Dec 2023 16:40:27 +0100 Subject: [PATCH 10/22] feat: rename `make_serializable` into `make_common_data_structure` for clarity Apply the ADR 0002 - https://github.com/apache/airflow/blob/main/airflow/providers/common/sql/doc/adr/0002-return-common-data-structure-from-dbapihook-derived-hooks.md --- airflow/providers/common/sql/hooks/sql.py | 17 ++++++++--------- .../databricks/hooks/databricks_sql.py | 6 +++--- airflow/providers/odbc/hooks/odbc.py | 5 ++--- 3 files changed, 13 insertions(+), 15 deletions(-) diff --git a/airflow/providers/common/sql/hooks/sql.py b/airflow/providers/common/sql/hooks/sql.py index 57c3e37e6cc9f..956342a14e691 100644 --- a/airflow/providers/common/sql/hooks/sql.py +++ b/airflow/providers/common/sql/hooks/sql.py @@ -392,7 +392,7 @@ def run( self._run_command(cur, sql_statement, parameters) if handler is not None: - result = self._make_serializable(handler(cur)) + result = self._make_common_data_structure(handler(cur)) if return_single_query_results(sql, return_last, split_statements): _last_result = result _last_description = cur.description @@ -412,17 +412,16 @@ def run( else: return results - @staticmethod - def _make_serializable(result: Any) -> Any: - """Ensure the data returned from an SQL command is JSON-serializable. + def _make_common_data_structure(self, result: Any) -> Any: + """Ensure the data returned from an SQL command is standard. This method is intended to be overridden by subclasses of the `DbApiHook`. Its purpose is to - transform the result of an SQL command (typically returned by cursor methods) into a - JSON-serializable format. + transform the result of an SQL command (typically returned by cursor methods) into a common + data structure (a tuple or list[tuple]) across all DBApiHook derived Hooks, as defined in the + ADR-0002 of the sql provider. - If this method is not overridden, the result data is returned as-is. - If the output of the cursor is already JSON-serializable, this method - should be ignored. + If this method is not overridden, the result data is returned as-is. If the output of the cursor + is already a common data structure, this method should be ignored. """ return result diff --git a/airflow/providers/databricks/hooks/databricks_sql.py b/airflow/providers/databricks/hooks/databricks_sql.py index 5ad4b0306bbd1..b7d5e9f5430a2 100644 --- a/airflow/providers/databricks/hooks/databricks_sql.py +++ b/airflow/providers/databricks/hooks/databricks_sql.py @@ -240,7 +240,7 @@ def run( with closing(conn.cursor()) as cur: self._run_command(cur, sql_statement, parameters) if handler is not None: - result = self._make_serializable(handler(cur)) + result = self._make_common_data_structure(handler(cur)) if return_single_query_results(sql, return_last, split_statements): results = [result] self.descriptions = [cur.description] @@ -258,8 +258,8 @@ def run( else: return results - def _make_serializable(self, result): - """Transform the databricks Row objects into serializable namedtuple.""" + def _make_common_data_structure(self, result: list[Row] | Row | None) -> list[tuple] | tuple | None: + """Transform the databricks Row objects into namedtuple.""" if self.return_serializable: columns: list[str] | None = None if isinstance(result, list): diff --git a/airflow/providers/odbc/hooks/odbc.py b/airflow/providers/odbc/hooks/odbc.py index 799e9f835c04f..97374cb05d91a 100644 --- a/airflow/providers/odbc/hooks/odbc.py +++ b/airflow/providers/odbc/hooks/odbc.py @@ -212,9 +212,8 @@ def get_sqlalchemy_connection( cnx = engine.connect(**(connect_kwargs or {})) return cnx - @staticmethod - def _make_serializable(result: list[pyodbc.Row] | pyodbc.Row | None) -> list[NamedTuple] | None: - """Transform the pyodbc.Row objects returned from an SQL command into serializable NamedTuple.""" + def _make_common_data_structure(self, result: list[pyodbc.Row] | pyodbc.Row | None) -> list[NamedTuple] | NamedTuple | None: + """Transform the pyodbc.Row objects returned from an SQL command into typed NamedTuples.""" # Below ignored lines respect NamedTuple docstring, but mypy do not support dynamically # instantiated Namedtuple, and will never do: https://github.com/python/mypy/issues/848 columns: list[tuple[str, type]] | None = None From e76d5b8cbc4a272c582db772f56c879b34bef74d Mon Sep 17 00:00:00 2001 From: Joffrey Bienvenu Date: Sat, 16 Dec 2023 18:18:18 +0100 Subject: [PATCH 11/22] feat: Add typing for `_make_common_data_structure` in databricks --- airflow/providers/common/sql/hooks/sql.py | 13 +++++---- .../databricks/hooks/databricks_sql.py | 28 +++++++++++-------- airflow/providers/odbc/hooks/odbc.py | 14 +++++----- airflow/providers/oracle/hooks/oracle.py | 2 +- 4 files changed, 32 insertions(+), 25 deletions(-) diff --git a/airflow/providers/common/sql/hooks/sql.py b/airflow/providers/common/sql/hooks/sql.py index 956342a14e691..e2e1449cac44a 100644 --- a/airflow/providers/common/sql/hooks/sql.py +++ b/airflow/providers/common/sql/hooks/sql.py @@ -29,7 +29,7 @@ Sequence, TypeVar, cast, - overload, + overload, Generic, Union, ) from urllib.parse import urlparse @@ -46,6 +46,7 @@ from airflow.providers.openlineage.sqlparser import DatabaseInfo +Common = TypeVar("Common", bound=tuple) T = TypeVar("T") @@ -122,10 +123,10 @@ class DbApiHook(BaseHook): """ Abstract base class for sql hooks. - When subclassing, maintainers can override the `_make_serializable` method: + When subclassing, maintainers can override the `_make_common_data_structure` method: This method transforms the result of the handler method (typically `cursor.fetchall()`) into - serializable objects. Most of the time, the underlying SQL library already returns tuples from - its cursor, and the `_make_serializable` method can be ignored. + objects common across all Hooks derived from this class (tuples). Most of the time, the underlying SQL + library already returns tuples from its cursor, and the `_make_common_data_structure` method can be ignored. :param schema: Optional DB schema that overrides the schema specified in the connection. Make sure that if you change the schema parameter value in the constructor of the derived Hook, such change @@ -305,7 +306,7 @@ def run( handler: Callable[[Any], T] = ..., split_statements: bool = ..., return_last: bool = ..., - ) -> T | list[T]: + ) -> T | tuple | list[T] | list[tuple] | list[Union[T, tuple, list[T], list[tuple], None]]: ... def run( @@ -316,7 +317,7 @@ def run( handler: Callable[[Any], T] | None = None, split_statements: bool = False, return_last: bool = True, - ) -> T | list[T] | None: + ) -> T | tuple | list[T] | list[tuple] | list[Union[T, tuple, list[T], list[tuple], None]] | None: """Run a command or a list of commands. Pass a list of SQL statements to the sql parameter to get them to diff --git a/airflow/providers/databricks/hooks/databricks_sql.py b/airflow/providers/databricks/hooks/databricks_sql.py index b7d5e9f5430a2..9f8e43bf7c0ff 100644 --- a/airflow/providers/databricks/hooks/databricks_sql.py +++ b/airflow/providers/databricks/hooks/databricks_sql.py @@ -20,7 +20,8 @@ from collections import namedtuple from contextlib import closing from copy import copy -from typing import TYPE_CHECKING, Any, Callable, Iterable, Mapping, TypeVar, overload +from typing import TYPE_CHECKING, Any, Callable, Iterable, Mapping, TypeVar, overload, cast, Type, Tuple, \ + Union from databricks import sql # type: ignore[attr-defined] from databricks.sql.types import Row @@ -35,6 +36,7 @@ LIST_SQL_ENDPOINTS_ENDPOINT = ("GET", "api/2.0/sql/endpoints") +# Common = Tuple[Any, ...] T = TypeVar("T") @@ -184,7 +186,7 @@ def run( handler: Callable[[Any], T] = ..., split_statements: bool = ..., return_last: bool = ..., - ) -> T | list[T]: + ) -> T | tuple | list[T] | list[tuple] | list[Union[T, tuple, list[T], list[tuple], None]]: ... def run( @@ -195,7 +197,7 @@ def run( handler: Callable[[Any], T] | None = None, split_statements: bool = True, return_last: bool = True, - ) -> T | list[T] | None: + ) -> T | tuple | list[T] | list[tuple] | list[Union[T, tuple, list[T], list[tuple], None]] | None: """ Run a command or a list of commands. @@ -258,17 +260,21 @@ def run( else: return results - def _make_common_data_structure(self, result: list[Row] | Row | None) -> list[tuple] | tuple | None: + def _make_common_data_structure(self, result: list[T] | T | None) -> list[T] | T | list[tuple] | tuple | None: """Transform the databricks Row objects into namedtuple.""" + # Below ignored lines respect namedtuple docstring, but mypy do not support dynamically + # instantiated namedtuple, and will never do: https://github.com/python/mypy/issues/848 if self.return_serializable: - columns: list[str] | None = None - if isinstance(result, list): - columns = result[0].__fields__ - row_object = namedtuple("Row", columns) - return [row_object(*row) for row in result] + if isinstance(result, list) and all(isinstance(item, Row) for item in result): + rows: list[Row] = result + rows_fields = rows[0].__fields__ + rows_object = namedtuple("Row", rows_fields) # type: ignore[misc] + return cast(list[tuple], [rows_object(*row) for row in rows]) elif isinstance(result, Row): - columns = result.__fields__ - return namedtuple("Row", columns)(*result) + row: Row = result + row_fields = row.__fields__ + row_object = namedtuple("Row", row_fields) # type: ignore[misc] + return cast(tuple, row_object(*row)) return result def bulk_dump(self, table, tmp_file): diff --git a/airflow/providers/odbc/hooks/odbc.py b/airflow/providers/odbc/hooks/odbc.py index 97374cb05d91a..49459782227f1 100644 --- a/airflow/providers/odbc/hooks/odbc.py +++ b/airflow/providers/odbc/hooks/odbc.py @@ -215,13 +215,13 @@ def get_sqlalchemy_connection( def _make_common_data_structure(self, result: list[pyodbc.Row] | pyodbc.Row | None) -> list[NamedTuple] | NamedTuple | None: """Transform the pyodbc.Row objects returned from an SQL command into typed NamedTuples.""" # Below ignored lines respect NamedTuple docstring, but mypy do not support dynamically - # instantiated Namedtuple, and will never do: https://github.com/python/mypy/issues/848 - columns: list[tuple[str, type]] | None = None - if isinstance(result, list): - columns = [col[:2] for col in result[0].cursor_description] - row_object = NamedTuple("Row", columns) # type: ignore[misc] + # instantiated typed Namedtuple, and will never do: https://github.com/python/mypy/issues/848 + field_names: list[tuple[str, type]] | None = None + if isinstance(result, list) and all(isinstance(item, pyodbc.Row) for item in result): + field_names = [col[:2] for col in result[0].cursor_description] + row_object = NamedTuple("Row", field_names) # type: ignore[misc] return [row_object(*row) for row in result] elif isinstance(result, pyodbc.Row): - columns = [col[:2] for col in result.cursor_description] - return NamedTuple("Row", columns)(*result) # type: ignore[misc, operator] + field_names = [col[:2] for col in result.cursor_description] + return NamedTuple("Row", field_names)(*result) # type: ignore[misc, operator] return result diff --git a/airflow/providers/oracle/hooks/oracle.py b/airflow/providers/oracle/hooks/oracle.py index e9c333a3477ec..de78ab726ff71 100644 --- a/airflow/providers/oracle/hooks/oracle.py +++ b/airflow/providers/oracle/hooks/oracle.py @@ -372,7 +372,7 @@ def callproc( identifier: str, autocommit: bool = False, parameters: list | dict | None = None, - ) -> list | dict | None: + ) -> list | dict | tuple | None: """ Call the stored procedure identified by the provided string. From a82553eb852b87c30f2258976b5ec847fe09c8c1 Mon Sep 17 00:00:00 2001 From: Joffrey Bienvenu Date: Sat, 16 Dec 2023 19:41:48 +0100 Subject: [PATCH 12/22] feat: Apply typing to ODBC hook --- airflow/providers/common/sql/hooks/sql.py | 7 +++---- .../databricks/hooks/databricks_sql.py | 21 +++++++++++++------ airflow/providers/odbc/hooks/odbc.py | 18 ++++++++++------ 3 files changed, 30 insertions(+), 16 deletions(-) diff --git a/airflow/providers/common/sql/hooks/sql.py b/airflow/providers/common/sql/hooks/sql.py index e2e1449cac44a..5da05f4f4b583 100644 --- a/airflow/providers/common/sql/hooks/sql.py +++ b/airflow/providers/common/sql/hooks/sql.py @@ -29,7 +29,7 @@ Sequence, TypeVar, cast, - overload, Generic, Union, + overload, ) from urllib.parse import urlparse @@ -46,7 +46,6 @@ from airflow.providers.openlineage.sqlparser import DatabaseInfo -Common = TypeVar("Common", bound=tuple) T = TypeVar("T") @@ -306,7 +305,7 @@ def run( handler: Callable[[Any], T] = ..., split_statements: bool = ..., return_last: bool = ..., - ) -> T | tuple | list[T] | list[tuple] | list[Union[T, tuple, list[T], list[tuple], None]]: + ) -> T | tuple | list[T] | list[tuple] | list[T | tuple | list[T] | list[tuple] | None]: ... def run( @@ -317,7 +316,7 @@ def run( handler: Callable[[Any], T] | None = None, split_statements: bool = False, return_last: bool = True, - ) -> T | tuple | list[T] | list[tuple] | list[Union[T, tuple, list[T], list[tuple], None]] | None: + ) -> T | tuple | list[T] | list[tuple] | list[T | tuple | list[T] | list[tuple] | None] | None: """Run a command or a list of commands. Pass a list of SQL statements to the sql parameter to get them to diff --git a/airflow/providers/databricks/hooks/databricks_sql.py b/airflow/providers/databricks/hooks/databricks_sql.py index 9f8e43bf7c0ff..a0f2c463a15d9 100644 --- a/airflow/providers/databricks/hooks/databricks_sql.py +++ b/airflow/providers/databricks/hooks/databricks_sql.py @@ -20,8 +20,16 @@ from collections import namedtuple from contextlib import closing from copy import copy -from typing import TYPE_CHECKING, Any, Callable, Iterable, Mapping, TypeVar, overload, cast, Type, Tuple, \ - Union +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Iterable, + Mapping, + TypeVar, + cast, + overload, +) from databricks import sql # type: ignore[attr-defined] from databricks.sql.types import Row @@ -36,7 +44,6 @@ LIST_SQL_ENDPOINTS_ENDPOINT = ("GET", "api/2.0/sql/endpoints") -# Common = Tuple[Any, ...] T = TypeVar("T") @@ -186,7 +193,7 @@ def run( handler: Callable[[Any], T] = ..., split_statements: bool = ..., return_last: bool = ..., - ) -> T | tuple | list[T] | list[tuple] | list[Union[T, tuple, list[T], list[tuple], None]]: + ) -> T | tuple | list[T] | list[tuple] | list[T | tuple | list[T] | list[tuple] | None]: ... def run( @@ -197,7 +204,7 @@ def run( handler: Callable[[Any], T] | None = None, split_statements: bool = True, return_last: bool = True, - ) -> T | tuple | list[T] | list[tuple] | list[Union[T, tuple, list[T], list[tuple], None]] | None: + ) -> T | tuple | list[T] | list[tuple] | list[T | tuple | list[T] | list[tuple] | None] | None: """ Run a command or a list of commands. @@ -260,7 +267,9 @@ def run( else: return results - def _make_common_data_structure(self, result: list[T] | T | None) -> list[T] | T | list[tuple] | tuple | None: + def _make_common_data_structure( + self, result: list[T] | T | None + ) -> list[T] | T | list[tuple] | tuple | None: """Transform the databricks Row objects into namedtuple.""" # Below ignored lines respect namedtuple docstring, but mypy do not support dynamically # instantiated namedtuple, and will never do: https://github.com/python/mypy/issues/848 diff --git a/airflow/providers/odbc/hooks/odbc.py b/airflow/providers/odbc/hooks/odbc.py index 49459782227f1..40bfe08b33439 100644 --- a/airflow/providers/odbc/hooks/odbc.py +++ b/airflow/providers/odbc/hooks/odbc.py @@ -17,7 +17,7 @@ """This module contains ODBC hook.""" from __future__ import annotations -from typing import Any, NamedTuple +from typing import Any, NamedTuple, TypeVar, cast from urllib.parse import quote_plus import pyodbc @@ -25,6 +25,8 @@ from airflow.providers.common.sql.hooks.sql import DbApiHook from airflow.utils.helpers import merge_dicts +T = TypeVar("T") + class OdbcHook(DbApiHook): """ @@ -212,16 +214,20 @@ def get_sqlalchemy_connection( cnx = engine.connect(**(connect_kwargs or {})) return cnx - def _make_common_data_structure(self, result: list[pyodbc.Row] | pyodbc.Row | None) -> list[NamedTuple] | NamedTuple | None: + def _make_common_data_structure( + self, result: list[T] | T | None + ) -> list[T] | T | list[NamedTuple] | NamedTuple | None: """Transform the pyodbc.Row objects returned from an SQL command into typed NamedTuples.""" # Below ignored lines respect NamedTuple docstring, but mypy do not support dynamically # instantiated typed Namedtuple, and will never do: https://github.com/python/mypy/issues/848 field_names: list[tuple[str, type]] | None = None if isinstance(result, list) and all(isinstance(item, pyodbc.Row) for item in result): - field_names = [col[:2] for col in result[0].cursor_description] + rows: list[pyodbc.Row] = result + field_names = [col[:2] for col in rows[0].cursor_description] row_object = NamedTuple("Row", field_names) # type: ignore[misc] - return [row_object(*row) for row in result] + return cast(list[NamedTuple], [row_object(*row) for row in rows]) elif isinstance(result, pyodbc.Row): - field_names = [col[:2] for col in result.cursor_description] - return NamedTuple("Row", field_names)(*result) # type: ignore[misc, operator] + row: pyodbc.Row = result + field_names = [col[:2] for col in row.cursor_description] + return cast(NamedTuple, NamedTuple("Row", field_names)(*row)) # type: ignore[misc, operator] return result From d3510efa97a980e5acd6f3ba0aa963896232bf70 Mon Sep 17 00:00:00 2001 From: Joffrey Bienvenu Date: Sun, 17 Dec 2023 15:45:49 +0100 Subject: [PATCH 13/22] fix: patch pyodbc.Row to make isinstance() checks pass with row_mock object --- tests/providers/odbc/hooks/test_odbc.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/providers/odbc/hooks/test_odbc.py b/tests/providers/odbc/hooks/test_odbc.py index 3740a5654c71c..b04c34dd09e08 100644 --- a/tests/providers/odbc/hooks/test_odbc.py +++ b/tests/providers/odbc/hooks/test_odbc.py @@ -305,7 +305,9 @@ def test_pyodbc_mock(self): """ assert hasattr(pyodbc.Row, "cursor_description") - def test_query_return_serializable_result_with_fetchall(self, pyodbc_row_mock): + def test_query_return_serializable_result_with_fetchall( + self, pyodbc_row_mock, monkeypatch, pyodbc_instancecheck + ): """ Simulate a cursor.fetchall which returns an iterable of pyodbc.Row object, and check if this iterable get converted into a list of tuples. @@ -317,7 +319,9 @@ def mock_handler(*_): return pyodbc_result hook = self.get_hook() - result = hook.run("SQL", handler=mock_handler) + with monkeypatch.context() as patcher: + patcher.setattr("pyodbc.Row", pyodbc_instancecheck) + result = hook.run("SQL", handler=mock_handler) assert hook_result == result def test_query_return_serializable_result_with_fetchone( From 0f20dfba36cb31eba7da0a2e39dc5b3502ab681c Mon Sep 17 00:00:00 2001 From: Joffrey Bienvenu Date: Sun, 17 Dec 2023 15:46:22 +0100 Subject: [PATCH 14/22] fix: Use List in type casting for python38 compatibility --- airflow/providers/odbc/hooks/odbc.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/airflow/providers/odbc/hooks/odbc.py b/airflow/providers/odbc/hooks/odbc.py index 40bfe08b33439..7d5be3ab62114 100644 --- a/airflow/providers/odbc/hooks/odbc.py +++ b/airflow/providers/odbc/hooks/odbc.py @@ -17,7 +17,7 @@ """This module contains ODBC hook.""" from __future__ import annotations -from typing import Any, NamedTuple, TypeVar, cast +from typing import Any, NamedTuple, TypeVar, cast, List from urllib.parse import quote_plus import pyodbc @@ -225,7 +225,7 @@ def _make_common_data_structure( rows: list[pyodbc.Row] = result field_names = [col[:2] for col in rows[0].cursor_description] row_object = NamedTuple("Row", field_names) # type: ignore[misc] - return cast(list[NamedTuple], [row_object(*row) for row in rows]) + return cast(List[NamedTuple], [row_object(*row) for row in rows]) elif isinstance(result, pyodbc.Row): row: pyodbc.Row = result field_names = [col[:2] for col in row.cursor_description] From ab871914600fb0420611ed63fc940aeef17ade98 Mon Sep 17 00:00:00 2001 From: Joffrey Bienvenu Date: Sun, 17 Dec 2023 15:56:01 +0100 Subject: [PATCH 15/22] feat: Remove all mentions of serialization in databricks hook --- .../databricks/hooks/databricks_sql.py | 20 +++++++++---------- airflow/providers/odbc/hooks/odbc.py | 2 +- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/airflow/providers/databricks/hooks/databricks_sql.py b/airflow/providers/databricks/hooks/databricks_sql.py index a0f2c463a15d9..89480b1459e52 100644 --- a/airflow/providers/databricks/hooks/databricks_sql.py +++ b/airflow/providers/databricks/hooks/databricks_sql.py @@ -63,10 +63,10 @@ class DatabricksSqlHook(BaseDatabricksHook, DbApiHook): on every request :param catalog: An optional initial catalog to use. Requires DBR version 9.0+ :param schema: An optional initial schema to use. Requires DBR version 9.0+ - :param return_serializable: Return a serializable ``namedtuple`` object instead of a - ``databricks.sql.Row`` object. Default to False. In a future release of the provider, this will - become True by default. This parameter ensures backward-compatibility during the transition phase, - and will be removed too in a future release. + :param return_tuple: Return a ``namedtuple`` object instead of a ``databricks.sql.Row`` object. Default + to False. In a future release of the provider, this will become True by default. This parameter + ensures backward-compatibility during the transition phase to common tuple objects for all hooks based + on DbApiHook. This flag will also be removed in a future release. :param kwargs: Additional parameters internal to Databricks SQL Connector parameters """ @@ -83,7 +83,7 @@ def __init__( catalog: str | None = None, schema: str | None = None, caller: str = "DatabricksSqlHook", - return_serializable: bool = False, + return_tuple: bool = False, **kwargs, ) -> None: super().__init__(databricks_conn_id, caller=caller) @@ -96,14 +96,14 @@ def __init__( self.http_headers = http_headers self.catalog = catalog self.schema = schema - self.return_serializable = return_serializable + self.return_tuple = return_tuple self.additional_params = kwargs - if not self.return_serializable: + if not self.return_tuple: warnings.warn( """Returning a raw `databricks.sql.Row` object is deprecated. A namedtuple will be - returned instead in a future release of the databricks provider. Set - `return_serializable=True` to enable this behavior.""", + returned instead in a future release of the databricks provider. Set `return_tuple=True` to + enable this behavior.""", AirflowProviderDeprecationWarning, stacklevel=2, ) @@ -273,7 +273,7 @@ def _make_common_data_structure( """Transform the databricks Row objects into namedtuple.""" # Below ignored lines respect namedtuple docstring, but mypy do not support dynamically # instantiated namedtuple, and will never do: https://github.com/python/mypy/issues/848 - if self.return_serializable: + if self.return_tuple: if isinstance(result, list) and all(isinstance(item, Row) for item in result): rows: list[Row] = result rows_fields = rows[0].__fields__ diff --git a/airflow/providers/odbc/hooks/odbc.py b/airflow/providers/odbc/hooks/odbc.py index 7d5be3ab62114..ef327bf847110 100644 --- a/airflow/providers/odbc/hooks/odbc.py +++ b/airflow/providers/odbc/hooks/odbc.py @@ -17,7 +17,7 @@ """This module contains ODBC hook.""" from __future__ import annotations -from typing import Any, NamedTuple, TypeVar, cast, List +from typing import Any, List, NamedTuple, TypeVar, cast from urllib.parse import quote_plus import pyodbc From 10e22c74938c2dbc9feb5cea49e22601286486ac Mon Sep 17 00:00:00 2001 From: Joffrey Bienvenu Date: Mon, 18 Dec 2023 09:37:48 +0100 Subject: [PATCH 16/22] feat: Rename `_make_serializable` into `_make_common_data_structure` in static checks --- .pre-commit-config.yaml | 4 +-- STATIC_CODE_CHECKS.rst | 4 +-- .../src/airflow_breeze/pre_commit_ids.py | 2 +- images/breeze/output_static-checks.svg | 36 +++++++++---------- images/breeze/output_static-checks.txt | 2 +- .../pre_commit_check_common_sql_dependency.py | 10 +++--- 6 files changed, 29 insertions(+), 29 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 18c1d7d64ca34..3cabc586182aa 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -385,8 +385,8 @@ repos: files: ^dev/breeze/src/airflow_breeze/utils/docker_command_utils\.py$|^scripts/ci/docker_compose/local\.yml$ pass_filenames: false additional_dependencies: ['rich>=12.4.4'] - - id: check-common-sql-dependency-make-serializable - name: Check dependency of SQL Providers with '_make_serializable' + - id: check-sql-dependency-common-data-structure + name: Check dependency of SQL Providers with common data structure entry: ./scripts/ci/pre_commit/pre_commit_check_common_sql_dependency.py language: python files: ^airflow/providers/.*/hooks/.*\.py$ diff --git a/STATIC_CODE_CHECKS.rst b/STATIC_CODE_CHECKS.rst index f16d119f6ea94..42aeb7d2327a8 100644 --- a/STATIC_CODE_CHECKS.rst +++ b/STATIC_CODE_CHECKS.rst @@ -170,8 +170,6 @@ require Breeze Docker image to be built locally. +-----------------------------------------------------------+--------------------------------------------------------------+---------+ | check-cncf-k8s-only-for-executors | Check cncf.kubernetes imports used for executors only | | +-----------------------------------------------------------+--------------------------------------------------------------+---------+ -| check-common-sql-dependency-make-serializable | Check dependency of SQL Providers with '_make_serializable' | | -+-----------------------------------------------------------+--------------------------------------------------------------+---------+ | check-core-deprecation-classes | Verify usage of Airflow deprecation classes in core | | +-----------------------------------------------------------+--------------------------------------------------------------+---------+ | check-daysago-import-from-utils | Make sure days_ago is imported from airflow.utils.dates | | @@ -240,6 +238,8 @@ require Breeze Docker image to be built locally. +-----------------------------------------------------------+--------------------------------------------------------------+---------+ | check-setup-order | Check order of dependencies in setup.cfg and setup.py | | +-----------------------------------------------------------+--------------------------------------------------------------+---------+ +| check-sql-dependency-common-data-structure | Check dependency of SQL Providers with common data structure | | ++-----------------------------------------------------------+--------------------------------------------------------------+---------+ | check-start-date-not-used-in-defaults | start_date not to be defined in default_args in example_dags | | +-----------------------------------------------------------+--------------------------------------------------------------+---------+ | check-system-tests-present | Check if system tests have required segments of code | | diff --git a/dev/breeze/src/airflow_breeze/pre_commit_ids.py b/dev/breeze/src/airflow_breeze/pre_commit_ids.py index 1297448cabf71..2806f0f1e923c 100644 --- a/dev/breeze/src/airflow_breeze/pre_commit_ids.py +++ b/dev/breeze/src/airflow_breeze/pre_commit_ids.py @@ -38,7 +38,6 @@ "check-builtin-literals", "check-changelog-has-no-duplicates", "check-cncf-k8s-only-for-executors", - "check-common-sql-dependency-make-serializable", "check-core-deprecation-classes", "check-daysago-import-from-utils", "check-decorated-operator-implements-custom-name", @@ -73,6 +72,7 @@ "check-revision-heads-map", "check-safe-filter-usage-in-html", "check-setup-order", + "check-sql-dependency-common-data-structure", "check-start-date-not-used-in-defaults", "check-system-tests-present", "check-system-tests-tocs", diff --git a/images/breeze/output_static-checks.svg b/images/breeze/output_static-checks.svg index 290f6a3264f64..6d0d3ea168ac4 100644 --- a/images/breeze/output_static-checks.svg +++ b/images/breeze/output_static-checks.svg @@ -313,24 +313,24 @@ │check-base-operator-partial-arguments | check-base-operator-usage |              │ │check-boring-cyborg-configuration | check-breeze-top-dependencies-limited |      │ │check-builtin-literals | check-changelog-has-no-duplicates |                     │ -│check-cncf-k8s-only-for-executors | check-common-sql-dependency-make-serializable│ -│| check-core-deprecation-classes | check-daysago-import-from-utils |             │ -│check-decorated-operator-implements-custom-name | check-deferrable-default-value │ -│| check-docstring-param-types | check-example-dags-urls |                        │ -│check-executables-have-shebangs | check-extra-packages-references |              │ -│check-extras-order | check-fab-migrations | check-for-inclusive-language |       │ -│check-google-re2-as-dependency | check-hooks-apply |                             │ -│check-incorrect-use-of-LoggingMixin | check-init-decorator-arguments |           │ -│check-lazy-logging | check-links-to-example-dags-do-not-use-hardcoded-versions | │ -│check-merge-conflict | check-newsfragments-are-valid |                           │ -│check-no-airflow-deprecation-in-providers | check-no-providers-in-core-examples |│ -│check-no-relative-imports | check-only-new-session-with-provide-session |        │ -│check-persist-credentials-disabled-in-github-workflows |                         │ -│check-pre-commit-information-consistent | check-provide-create-sessions-imports |│ -│check-provider-docs-valid | check-provider-yaml-valid |                          │ -│check-providers-init-file-missing | check-providers-subpackages-init-file-exist |│ -│check-pydevd-left-in-code | check-revision-heads-map |                           │ -│check-safe-filter-usage-in-html | check-setup-order |                            │ +│check-cncf-k8s-only-for-executors | check-core-deprecation-classes |             │ +│check-daysago-import-from-utils | check-decorated-operator-implements-custom-name│ +│| check-deferrable-default-value | check-docstring-param-types |                 │ +│check-example-dags-urls | check-executables-have-shebangs |                      │ +│check-extra-packages-references | check-extras-order | check-fab-migrations |    │ +│check-for-inclusive-language | check-google-re2-as-dependency | check-hooks-apply│ +│| check-incorrect-use-of-LoggingMixin | check-init-decorator-arguments |         │ +│check-lazy-logging | check-links-to-example-dags-do-not-use-hardcoded-versions | │ +│check-merge-conflict | check-newsfragments-are-valid |                           │ +│check-no-airflow-deprecation-in-providers | check-no-providers-in-core-examples |│ +│check-no-relative-imports | check-only-new-session-with-provide-session |        │ +│check-persist-credentials-disabled-in-github-workflows |                         │ +│check-pre-commit-information-consistent | check-provide-create-sessions-imports |│ +│check-provider-docs-valid | check-provider-yaml-valid |                          │ +│check-providers-init-file-missing | check-providers-subpackages-init-file-exist |│ +│check-pydevd-left-in-code | check-revision-heads-map |                           │ +│check-safe-filter-usage-in-html | check-setup-order |                            │ +│check-sql-dependency-common-data-structure |                                     │ │check-start-date-not-used-in-defaults | check-system-tests-present |             │ │check-system-tests-tocs | check-tests-unittest-testcase |                        │ │check-urlparse-usage-in-code | check-usage-of-re2-over-re | check-xml | codespell│ diff --git a/images/breeze/output_static-checks.txt b/images/breeze/output_static-checks.txt index 5b32905ea3617..9ffe4f833aec9 100644 --- a/images/breeze/output_static-checks.txt +++ b/images/breeze/output_static-checks.txt @@ -1 +1 @@ -1197108ac5d3038067e599375d5130dd +6fb4fd65fb7d3b1430a7de7a17c85e22 diff --git a/scripts/ci/pre_commit/pre_commit_check_common_sql_dependency.py b/scripts/ci/pre_commit/pre_commit_check_common_sql_dependency.py index 4b335d1cf639a..69bd1779c6802 100755 --- a/scripts/ci/pre_commit/pre_commit_check_common_sql_dependency.py +++ b/scripts/ci/pre_commit/pre_commit_check_common_sql_dependency.py @@ -33,7 +33,7 @@ COMMON_SQL_PROVIDER_NAME: str = "apache-airflow-providers-common-sql" COMMON_SQL_PROVIDER_MIN_COMPATIBLE_VERSIONS: str = "1.8.1" COMMON_SQL_PROVIDER_LATEST_INCOMPATIBLE_VERSION: str = "1.8.0" -MAKE_SERIALIZABLE_METHOD_NAME: str = "_make_serializable" +MAKE_COMMON_METHOD_NAME: str = "_make_common_data_structure" def get_classes(file_path: str) -> Iterable[ast.ClassDef]: @@ -54,9 +54,9 @@ def is_subclass_of_dbapihook(node: ast.ClassDef) -> bool: def has_make_serializable_method(node: ast.ClassDef) -> bool: - """Return True if the given class implements `_make_serializable` method.""" + """Return True if the given class implements `_make_common_data_structure` method.""" for body_element in node.body: - if isinstance(body_element, ast.FunctionDef) and (body_element.name == MAKE_SERIALIZABLE_METHOD_NAME): + if isinstance(body_element, ast.FunctionDef) and (body_element.name == MAKE_COMMON_METHOD_NAME): return True return False @@ -109,11 +109,11 @@ def check_sql_providers_dependency(): f"\n[yellow]Provider {provider_metadata['name']} must have " f"'{COMMON_SQL_PROVIDER_NAME}>={COMMON_SQL_PROVIDER_MIN_COMPATIBLE_VERSIONS}' as " f"dependency, because `{clazz.name}` overrides the " - f"`{MAKE_SERIALIZABLE_METHOD_NAME}` method." + f"`{MAKE_COMMON_METHOD_NAME}` method." ) if error_count: console.print( - f"The `{MAKE_SERIALIZABLE_METHOD_NAME}` method was introduced in {COMMON_SQL_PROVIDER_NAME} " + f"The `{MAKE_COMMON_METHOD_NAME}` method was introduced in {COMMON_SQL_PROVIDER_NAME} " f"{COMMON_SQL_PROVIDER_MIN_COMPATIBLE_VERSIONS}. You cannot rely on an older version of this " "provider to override this method." ) From 4796882f988d694f3ac89f83a65386c7f024e490 Mon Sep 17 00:00:00 2001 From: Joffrey Bienvenu Date: Mon, 18 Dec 2023 11:12:45 +0100 Subject: [PATCH 17/22] fix: use `return_tuple` in databricks tests --- airflow/providers/databricks/operators/databricks_sql.py | 2 +- tests/providers/databricks/hooks/test_databricks_sql.py | 6 +++--- tests/providers/databricks/operators/test_databricks_sql.py | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/airflow/providers/databricks/operators/databricks_sql.py b/airflow/providers/databricks/operators/databricks_sql.py index d5349d38dcbc9..96cadd827d320 100644 --- a/airflow/providers/databricks/operators/databricks_sql.py +++ b/airflow/providers/databricks/operators/databricks_sql.py @@ -113,7 +113,7 @@ def get_db_hook(self) -> DatabricksSqlHook: "catalog": self.catalog, "schema": self.schema, "caller": "DatabricksSqlOperator", - "return_serializable": True, + "return_tuple": True, **self.client_parameters, **self.hook_params, } diff --git a/tests/providers/databricks/hooks/test_databricks_sql.py b/tests/providers/databricks/hooks/test_databricks_sql.py index 3a1516f9ece83..a5d85880bec30 100644 --- a/tests/providers/databricks/hooks/test_databricks_sql.py +++ b/tests/providers/databricks/hooks/test_databricks_sql.py @@ -64,7 +64,7 @@ def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: @pytest.mark.parametrize( - "return_last, split_statements, sql, cursor_calls, return_serializable," + "return_last, split_statements, sql, cursor_calls, return_tuple," "cursor_descriptions, cursor_results, hook_descriptions, hook_results, ", [ pytest.param( @@ -212,7 +212,7 @@ def test_query( split_statements, sql, cursor_calls, - return_serializable, + return_tuple, cursor_descriptions, cursor_results, hook_descriptions, @@ -249,7 +249,7 @@ def test_query( cursors.append(cur) connections.append(conn) mock_conn.side_effect = connections - databricks_hook = DatabricksSqlHook(sql_endpoint_name="Test", return_serializable=return_serializable) + databricks_hook = DatabricksSqlHook(sql_endpoint_name="Test", return_tuple=return_tuple) results = databricks_hook.run( sql=sql, handler=fetch_all_handler, return_last=return_last, split_statements=split_statements ) diff --git a/tests/providers/databricks/operators/test_databricks_sql.py b/tests/providers/databricks/operators/test_databricks_sql.py index d791985d54b8d..d68dda41ab515 100644 --- a/tests/providers/databricks/operators/test_databricks_sql.py +++ b/tests/providers/databricks/operators/test_databricks_sql.py @@ -133,7 +133,7 @@ def test_exec_success(sql, return_last, split_statement, hook_results, hook_desc db_mock_class.assert_called_once_with( DEFAULT_CONN_ID, http_path=None, - return_serializable=True, + return_tuple=True, session_configuration=None, sql_endpoint_name=None, http_headers=None, @@ -277,7 +277,7 @@ def test_exec_write_file( db_mock_class.assert_called_once_with( DEFAULT_CONN_ID, http_path=None, - return_serializable=True, + return_tuple=True, session_configuration=None, sql_endpoint_name=None, http_headers=None, From 61255ede79ea54e64b7e9a752d1d13d2b36054d5 Mon Sep 17 00:00:00 2001 From: Joffrey Bienvenu Date: Mon, 18 Dec 2023 11:22:11 +0100 Subject: [PATCH 18/22] fix: add deprecated _make_serializable method in odbc.py --- airflow/providers/odbc/hooks/odbc.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/airflow/providers/odbc/hooks/odbc.py b/airflow/providers/odbc/hooks/odbc.py index ef327bf847110..388e191650972 100644 --- a/airflow/providers/odbc/hooks/odbc.py +++ b/airflow/providers/odbc/hooks/odbc.py @@ -17,11 +17,13 @@ """This module contains ODBC hook.""" from __future__ import annotations +import warnings from typing import Any, List, NamedTuple, TypeVar, cast from urllib.parse import quote_plus import pyodbc +from airflow.exceptions import AirflowProviderDeprecationWarning from airflow.providers.common.sql.hooks.sql import DbApiHook from airflow.utils.helpers import merge_dicts @@ -214,6 +216,13 @@ def get_sqlalchemy_connection( cnx = engine.connect(**(connect_kwargs or {})) return cnx + def _make_serializable(self, result: Any) -> Any: + """Use `airflow.providers.odbc.hooks.odbc.OdbcHook._make_common_data_structure`. + + This method is deprecated. + """ + return self._make_common_data_structure(result=result) + def _make_common_data_structure( self, result: list[T] | T | None ) -> list[T] | T | list[NamedTuple] | NamedTuple | None: From e54ca541c221ede3d089dead9963f84131eecf4a Mon Sep 17 00:00:00 2001 From: Joffrey Bienvenu Date: Mon, 18 Dec 2023 11:23:12 +0100 Subject: [PATCH 19/22] bump: min sql.common to 1.9.1 for odbc and databricks --- airflow/providers/databricks/provider.yaml | 2 +- airflow/providers/odbc/provider.yaml | 2 +- generated/provider_dependencies.json | 4 ++-- .../ci/pre_commit/pre_commit_check_common_sql_dependency.py | 4 ++-- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/airflow/providers/databricks/provider.yaml b/airflow/providers/databricks/provider.yaml index 84cc7b4593c46..53c52bd02f196 100644 --- a/airflow/providers/databricks/provider.yaml +++ b/airflow/providers/databricks/provider.yaml @@ -59,7 +59,7 @@ versions: dependencies: - apache-airflow>=2.6.0 - - apache-airflow-providers-common-sql>=1.8.1 + - apache-airflow-providers-common-sql>=1.9.1 - requests>=2.27,<3 # The connector 2.9.0 released on Aug 10, 2023 has a bug that it does not properly declare urllib3 and # it needs to be excluded. See https://github.com/databricks/databricks-sql-python/issues/190 diff --git a/airflow/providers/odbc/provider.yaml b/airflow/providers/odbc/provider.yaml index d83f4cf6891a6..30e7f4f0cae5e 100644 --- a/airflow/providers/odbc/provider.yaml +++ b/airflow/providers/odbc/provider.yaml @@ -45,7 +45,7 @@ versions: dependencies: - apache-airflow>=2.6.0 - - apache-airflow-providers-common-sql>=1.8.1 + - apache-airflow-providers-common-sql>=1.9.1 - pyodbc integrations: diff --git a/generated/provider_dependencies.json b/generated/provider_dependencies.json index b755f6c839c3c..cb10b8ec6e250 100644 --- a/generated/provider_dependencies.json +++ b/generated/provider_dependencies.json @@ -291,7 +291,7 @@ "databricks": { "deps": [ "aiohttp>=3.6.3, <4", - "apache-airflow-providers-common-sql>=1.8.1", + "apache-airflow-providers-common-sql>=1.9.1", "apache-airflow>=2.6.0", "databricks-sql-connector>=2.0.0, <3.0.0, !=2.9.0", "requests>=2.27,<3" @@ -653,7 +653,7 @@ }, "odbc": { "deps": [ - "apache-airflow-providers-common-sql>=1.8.1", + "apache-airflow-providers-common-sql>=1.9.1", "apache-airflow>=2.6.0", "pyodbc" ], diff --git a/scripts/ci/pre_commit/pre_commit_check_common_sql_dependency.py b/scripts/ci/pre_commit/pre_commit_check_common_sql_dependency.py index 69bd1779c6802..9719310a7174d 100755 --- a/scripts/ci/pre_commit/pre_commit_check_common_sql_dependency.py +++ b/scripts/ci/pre_commit/pre_commit_check_common_sql_dependency.py @@ -31,8 +31,8 @@ COMMON_SQL_PROVIDER_NAME: str = "apache-airflow-providers-common-sql" -COMMON_SQL_PROVIDER_MIN_COMPATIBLE_VERSIONS: str = "1.8.1" -COMMON_SQL_PROVIDER_LATEST_INCOMPATIBLE_VERSION: str = "1.8.0" +COMMON_SQL_PROVIDER_MIN_COMPATIBLE_VERSIONS: str = "1.9.1" +COMMON_SQL_PROVIDER_LATEST_INCOMPATIBLE_VERSION: str = "1.9.0" MAKE_COMMON_METHOD_NAME: str = "_make_common_data_structure" From e0c7ad4d22fdd5b9bbce169d0633100d0907e735 Mon Sep 17 00:00:00 2001 From: Joffrey Bienvenu Date: Mon, 18 Dec 2023 21:12:07 +0100 Subject: [PATCH 20/22] fix: add version 1.9.1 to provider.yaml of common.sql --- airflow/providers/common/sql/provider.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/airflow/providers/common/sql/provider.yaml b/airflow/providers/common/sql/provider.yaml index 92ce645375e75..1a64721e2cf51 100644 --- a/airflow/providers/common/sql/provider.yaml +++ b/airflow/providers/common/sql/provider.yaml @@ -24,6 +24,7 @@ description: | suspended: false source-date-epoch: 1701983370 versions: + - 1.9.1 - 1.9.0 - 1.8.1 - 1.8.0 From 74a763dd4bc460d08b1d638465f3bab81e4e189b Mon Sep 17 00:00:00 2001 From: Joffrey Bienvenu Date: Mon, 18 Dec 2023 21:13:08 +0100 Subject: [PATCH 21/22] fix: move back-compat in common.sql --- airflow/providers/common/sql/hooks/sql.py | 15 ++++++++++++++- airflow/providers/odbc/hooks/odbc.py | 9 --------- 2 files changed, 14 insertions(+), 10 deletions(-) diff --git a/airflow/providers/common/sql/hooks/sql.py b/airflow/providers/common/sql/hooks/sql.py index 5da05f4f4b583..7b9f128e44e44 100644 --- a/airflow/providers/common/sql/hooks/sql.py +++ b/airflow/providers/common/sql/hooks/sql.py @@ -16,6 +16,8 @@ # under the License. from __future__ import annotations +import contextlib +import warnings from contextlib import closing from datetime import datetime from typing import ( @@ -36,7 +38,7 @@ import sqlparse from sqlalchemy import create_engine -from airflow.exceptions import AirflowException +from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning from airflow.hooks.base import BaseHook if TYPE_CHECKING: @@ -423,6 +425,17 @@ def _make_common_data_structure(self, result: Any) -> Any: If this method is not overridden, the result data is returned as-is. If the output of the cursor is already a common data structure, this method should be ignored. """ + # Back-compatibility call for providers implementing old ´_make_serializable' method. + with contextlib.suppress(AttributeError): + result = self._make_serializable(result=result) # type: ignore[attr-defined] + warnings.warn( + "The `_make_serializable` method is deprecated and support will be removed in a future " + f"version of the common.sql provider. Please update the {self.__class__.__name__}'s provider " + "to a version based on common.sql >= 1.9.1.", + AirflowProviderDeprecationWarning, + stacklevel=2, + ) + return result return result def _run_command(self, cur, sql_statement, parameters): diff --git a/airflow/providers/odbc/hooks/odbc.py b/airflow/providers/odbc/hooks/odbc.py index 388e191650972..ef327bf847110 100644 --- a/airflow/providers/odbc/hooks/odbc.py +++ b/airflow/providers/odbc/hooks/odbc.py @@ -17,13 +17,11 @@ """This module contains ODBC hook.""" from __future__ import annotations -import warnings from typing import Any, List, NamedTuple, TypeVar, cast from urllib.parse import quote_plus import pyodbc -from airflow.exceptions import AirflowProviderDeprecationWarning from airflow.providers.common.sql.hooks.sql import DbApiHook from airflow.utils.helpers import merge_dicts @@ -216,13 +214,6 @@ def get_sqlalchemy_connection( cnx = engine.connect(**(connect_kwargs or {})) return cnx - def _make_serializable(self, result: Any) -> Any: - """Use `airflow.providers.odbc.hooks.odbc.OdbcHook._make_common_data_structure`. - - This method is deprecated. - """ - return self._make_common_data_structure(result=result) - def _make_common_data_structure( self, result: list[T] | T | None ) -> list[T] | T | list[NamedTuple] | NamedTuple | None: From affd2f45d25410c41b62699786f91004c54b514c Mon Sep 17 00:00:00 2001 From: Joffrey Bienvenu Date: Thu, 21 Dec 2023 18:59:51 +0100 Subject: [PATCH 22/22] feat: Implement strict `tuple`/`list[tuple]` return in ODBC and Databricks hooks --- airflow/providers/common/sql/hooks/sql.py | 15 ++++--- .../databricks/hooks/databricks_sql.py | 41 ++++++++++--------- airflow/providers/exasol/hooks/exasol.py | 6 +-- airflow/providers/exasol/provider.yaml | 2 +- airflow/providers/odbc/hooks/odbc.py | 29 +++++-------- .../providers/snowflake/hooks/snowflake.py | 6 +-- airflow/providers/snowflake/provider.yaml | 2 +- generated/provider_dependencies.json | 4 +- 8 files changed, 52 insertions(+), 53 deletions(-) diff --git a/airflow/providers/common/sql/hooks/sql.py b/airflow/providers/common/sql/hooks/sql.py index 7b9f128e44e44..707263cd50732 100644 --- a/airflow/providers/common/sql/hooks/sql.py +++ b/airflow/providers/common/sql/hooks/sql.py @@ -26,6 +26,7 @@ Callable, Generator, Iterable, + List, Mapping, Protocol, Sequence, @@ -307,7 +308,7 @@ def run( handler: Callable[[Any], T] = ..., split_statements: bool = ..., return_last: bool = ..., - ) -> T | tuple | list[T] | list[tuple] | list[T | tuple | list[T] | list[tuple] | None]: + ) -> tuple | list[tuple] | list[list[tuple] | tuple] | None: ... def run( @@ -318,7 +319,7 @@ def run( handler: Callable[[Any], T] | None = None, split_statements: bool = False, return_last: bool = True, - ) -> T | tuple | list[T] | list[tuple] | list[T | tuple | list[T] | list[tuple] | None] | None: + ) -> tuple | list[tuple] | list[list[tuple] | tuple] | None: """Run a command or a list of commands. Pass a list of SQL statements to the sql parameter to get them to @@ -414,8 +415,8 @@ def run( else: return results - def _make_common_data_structure(self, result: Any) -> Any: - """Ensure the data returned from an SQL command is standard. + def _make_common_data_structure(self, result: T | Sequence[T]) -> tuple | list[tuple]: + """Ensure the data returned from an SQL command is a standard tuple or list[tuple]. This method is intended to be overridden by subclasses of the `DbApiHook`. Its purpose is to transform the result of an SQL command (typically returned by cursor methods) into a common @@ -435,8 +436,10 @@ def _make_common_data_structure(self, result: Any) -> Any: AirflowProviderDeprecationWarning, stacklevel=2, ) - return result - return result + + if isinstance(result, Sequence): + return cast(List[tuple], result) + return cast(tuple, result) def _run_command(self, cur, sql_statement, parameters): """Run a statement using an already open cursor.""" diff --git a/airflow/providers/databricks/hooks/databricks_sql.py b/airflow/providers/databricks/hooks/databricks_sql.py index 89480b1459e52..6c31691c451bc 100644 --- a/airflow/providers/databricks/hooks/databricks_sql.py +++ b/airflow/providers/databricks/hooks/databricks_sql.py @@ -25,14 +25,15 @@ Any, Callable, Iterable, + List, Mapping, + Sequence, TypeVar, cast, overload, ) from databricks import sql # type: ignore[attr-defined] -from databricks.sql.types import Row from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning from airflow.providers.common.sql.hooks.sql import DbApiHook, return_single_query_results @@ -40,6 +41,7 @@ if TYPE_CHECKING: from databricks.sql.client import Connection + from databricks.sql.types import Row LIST_SQL_ENDPOINTS_ENDPOINT = ("GET", "api/2.0/sql/endpoints") @@ -193,7 +195,7 @@ def run( handler: Callable[[Any], T] = ..., split_statements: bool = ..., return_last: bool = ..., - ) -> T | tuple | list[T] | list[tuple] | list[T | tuple | list[T] | list[tuple] | None]: + ) -> tuple | list[tuple] | list[list[tuple] | tuple] | None: ... def run( @@ -204,7 +206,7 @@ def run( handler: Callable[[Any], T] | None = None, split_statements: bool = True, return_last: bool = True, - ) -> T | tuple | list[T] | list[tuple] | list[T | tuple | list[T] | list[tuple] | None] | None: + ) -> tuple | list[tuple] | list[list[tuple] | tuple] | None: """ Run a command or a list of commands. @@ -249,7 +251,12 @@ def run( with closing(conn.cursor()) as cur: self._run_command(cur, sql_statement, parameters) if handler is not None: - result = self._make_common_data_structure(handler(cur)) + raw_result = handler(cur) + if self.return_tuple: + result = self._make_common_data_structure(raw_result) + else: + # Returning raw result is deprecated, and do not comply with current common.sql interface + result = raw_result # type: ignore[assignment] if return_single_query_results(sql, return_last, split_statements): results = [result] self.descriptions = [cur.description] @@ -267,24 +274,20 @@ def run( else: return results - def _make_common_data_structure( - self, result: list[T] | T | None - ) -> list[T] | T | list[tuple] | tuple | None: + def _make_common_data_structure(self, result: Sequence[Row] | Row) -> list[tuple] | tuple: """Transform the databricks Row objects into namedtuple.""" # Below ignored lines respect namedtuple docstring, but mypy do not support dynamically # instantiated namedtuple, and will never do: https://github.com/python/mypy/issues/848 - if self.return_tuple: - if isinstance(result, list) and all(isinstance(item, Row) for item in result): - rows: list[Row] = result - rows_fields = rows[0].__fields__ - rows_object = namedtuple("Row", rows_fields) # type: ignore[misc] - return cast(list[tuple], [rows_object(*row) for row in rows]) - elif isinstance(result, Row): - row: Row = result - row_fields = row.__fields__ - row_object = namedtuple("Row", row_fields) # type: ignore[misc] - return cast(tuple, row_object(*row)) - return result + if isinstance(result, list): + rows: list[Row] = result + rows_fields = rows[0].__fields__ + rows_object = namedtuple("Row", rows_fields) # type: ignore[misc] + return cast(List[tuple], [rows_object(*row) for row in rows]) + else: + row: Row = result + row_fields = row.__fields__ + row_object = namedtuple("Row", row_fields) # type: ignore[misc] + return cast(tuple, row_object(*row)) def bulk_dump(self, table, tmp_file): raise NotImplementedError() diff --git a/airflow/providers/exasol/hooks/exasol.py b/airflow/providers/exasol/hooks/exasol.py index 98955a2579cfb..6d52b5122b633 100644 --- a/airflow/providers/exasol/hooks/exasol.py +++ b/airflow/providers/exasol/hooks/exasol.py @@ -183,7 +183,7 @@ def run( handler: Callable[[Any], T] = ..., split_statements: bool = ..., return_last: bool = ..., - ) -> T | list[T]: + ) -> tuple | list[tuple] | list[list[tuple] | tuple] | None: ... def run( @@ -194,7 +194,7 @@ def run( handler: Callable[[Any], T] | None = None, split_statements: bool = False, return_last: bool = True, - ) -> T | list[T] | None: + ) -> tuple | list[tuple] | list[list[tuple] | tuple] | None: """Run a command or a list of commands. Pass a list of SQL statements to the SQL parameter to get them to @@ -232,7 +232,7 @@ def run( with closing(conn.execute(sql_statement, parameters)) as exa_statement: self.log.info("Running statement: %s, parameters: %s", sql_statement, parameters) if handler is not None: - result = handler(exa_statement) + result = self._make_common_data_structure(handler(exa_statement)) if return_single_query_results(sql, return_last, split_statements): _last_result = result _last_columns = self.get_description(exa_statement) diff --git a/airflow/providers/exasol/provider.yaml b/airflow/providers/exasol/provider.yaml index d134431a98209..89b399e428b28 100644 --- a/airflow/providers/exasol/provider.yaml +++ b/airflow/providers/exasol/provider.yaml @@ -52,7 +52,7 @@ versions: dependencies: - apache-airflow>=2.6.0 - - apache-airflow-providers-common-sql>=1.3.1 + - apache-airflow-providers-common-sql>=1.9.1 - pyexasol>=0.5.1 - pandas>=0.17.1 diff --git a/airflow/providers/odbc/hooks/odbc.py b/airflow/providers/odbc/hooks/odbc.py index ef327bf847110..44d5aaceaea49 100644 --- a/airflow/providers/odbc/hooks/odbc.py +++ b/airflow/providers/odbc/hooks/odbc.py @@ -17,16 +17,14 @@ """This module contains ODBC hook.""" from __future__ import annotations -from typing import Any, List, NamedTuple, TypeVar, cast +from typing import Any, List, NamedTuple, Sequence, cast from urllib.parse import quote_plus -import pyodbc +from pyodbc import Connection, Row, connect from airflow.providers.common.sql.hooks.sql import DbApiHook from airflow.utils.helpers import merge_dicts -T = TypeVar("T") - class OdbcHook(DbApiHook): """ @@ -195,9 +193,9 @@ def connect_kwargs(self) -> dict: return merged_connect_kwargs - def get_conn(self) -> pyodbc.Connection: + def get_conn(self) -> Connection: """Returns a pyodbc connection object.""" - conn = pyodbc.connect(self.odbc_connection_string, **self.connect_kwargs) + conn = connect(self.odbc_connection_string, **self.connect_kwargs) return conn def get_uri(self) -> str: @@ -214,20 +212,15 @@ def get_sqlalchemy_connection( cnx = engine.connect(**(connect_kwargs or {})) return cnx - def _make_common_data_structure( - self, result: list[T] | T | None - ) -> list[T] | T | list[NamedTuple] | NamedTuple | None: + def _make_common_data_structure(self, result: Sequence[Row] | Row) -> list[tuple] | tuple: """Transform the pyodbc.Row objects returned from an SQL command into typed NamedTuples.""" # Below ignored lines respect NamedTuple docstring, but mypy do not support dynamically # instantiated typed Namedtuple, and will never do: https://github.com/python/mypy/issues/848 field_names: list[tuple[str, type]] | None = None - if isinstance(result, list) and all(isinstance(item, pyodbc.Row) for item in result): - rows: list[pyodbc.Row] = result - field_names = [col[:2] for col in rows[0].cursor_description] + if isinstance(result, Sequence): + field_names = [col[:2] for col in result[0].cursor_description] row_object = NamedTuple("Row", field_names) # type: ignore[misc] - return cast(List[NamedTuple], [row_object(*row) for row in rows]) - elif isinstance(result, pyodbc.Row): - row: pyodbc.Row = result - field_names = [col[:2] for col in row.cursor_description] - return cast(NamedTuple, NamedTuple("Row", field_names)(*row)) # type: ignore[misc, operator] - return result + return cast(List[tuple], [row_object(*row) for row in result]) + else: + field_names = [col[:2] for col in result.cursor_description] + return cast(tuple, NamedTuple("Row", field_names)(*result)) # type: ignore[misc, operator] diff --git a/airflow/providers/snowflake/hooks/snowflake.py b/airflow/providers/snowflake/hooks/snowflake.py index ead3e92274300..4b0d13b5e5027 100644 --- a/airflow/providers/snowflake/hooks/snowflake.py +++ b/airflow/providers/snowflake/hooks/snowflake.py @@ -323,7 +323,7 @@ def run( split_statements: bool = ..., return_last: bool = ..., return_dictionaries: bool = ..., - ) -> T | list[T]: + ) -> tuple | list[tuple] | list[list[tuple] | tuple] | None: ... def run( @@ -335,7 +335,7 @@ def run( split_statements: bool = True, return_last: bool = True, return_dictionaries: bool = False, - ) -> T | list[T] | None: + ) -> tuple | list[tuple] | list[list[tuple] | tuple] | None: """Runs a command or a list of commands. Pass a list of SQL statements to the SQL parameter to get them to @@ -388,7 +388,7 @@ def run( self._run_command(cur, sql_statement, parameters) if handler is not None: - result = handler(cur) + result = self._make_common_data_structure(handler(cur)) if return_single_query_results(sql, return_last, split_statements): _last_result = result _last_description = cur.description diff --git a/airflow/providers/snowflake/provider.yaml b/airflow/providers/snowflake/provider.yaml index b7355e69b4216..ee6fa214acc14 100644 --- a/airflow/providers/snowflake/provider.yaml +++ b/airflow/providers/snowflake/provider.yaml @@ -67,7 +67,7 @@ versions: dependencies: - apache-airflow>=2.6.0 - - apache-airflow-providers-common-sql>=1.3.1 + - apache-airflow-providers-common-sql>=1.9.1 - snowflake-connector-python>=2.7.8 - snowflake-sqlalchemy>=1.1.0 diff --git a/generated/provider_dependencies.json b/generated/provider_dependencies.json index cb10b8ec6e250..b20543ab35399 100644 --- a/generated/provider_dependencies.json +++ b/generated/provider_dependencies.json @@ -364,7 +364,7 @@ }, "exasol": { "deps": [ - "apache-airflow-providers-common-sql>=1.3.1", + "apache-airflow-providers-common-sql>=1.9.1", "apache-airflow>=2.6.0", "pandas>=0.17.1", "pyexasol>=0.5.1" @@ -864,7 +864,7 @@ }, "snowflake": { "deps": [ - "apache-airflow-providers-common-sql>=1.3.1", + "apache-airflow-providers-common-sql>=1.9.1", "apache-airflow>=2.6.0", "snowflake-connector-python>=2.7.8", "snowflake-sqlalchemy>=1.1.0"