You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by bo...@apache.org on 2023/12/22 15:38:04 UTC
(airflow) branch main updated: Return common data structure in DBApi derived classes
This is an automated email from the ASF dual-hosted git repository.
bolke pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 5fe5d31a46 Return common data structure in DBApi derived classes
5fe5d31a46 is described below
commit 5fe5d31a46885fbb2fb6ba9c0bd551a6b57d129a
Author: Joffrey Bienvenu <jo...@infrabel.be>
AuthorDate: Fri Dec 22 16:37:56 2023 +0100
Return common data structure in DBApi derived classes
The ADR for Airflow' s DB API specifies it needs to return a
named tuple SerializableRow or a list of them.
---
.pre-commit-config.yaml | 4 +-
STATIC_CODE_CHECKS.rst | 4 +-
airflow/providers/common/sql/hooks/sql.py | 49 ++++++++++------
airflow/providers/common/sql/provider.yaml | 1 +
.../providers/databricks/hooks/databricks_sql.py | 65 +++++++++++++++++-----
.../databricks/operators/databricks_sql.py | 1 +
airflow/providers/databricks/provider.yaml | 2 +-
airflow/providers/exasol/hooks/exasol.py | 6 +-
airflow/providers/exasol/provider.yaml | 2 +-
airflow/providers/odbc/hooks/odbc.py | 32 +++++------
airflow/providers/odbc/provider.yaml | 2 +-
airflow/providers/oracle/hooks/oracle.py | 2 +-
airflow/providers/snowflake/hooks/snowflake.py | 6 +-
airflow/providers/snowflake/provider.yaml | 2 +-
dev/breeze/src/airflow_breeze/pre_commit_ids.py | 2 +-
generated/provider_dependencies.json | 8 +--
images/breeze/output_static-checks.svg | 36 ++++++------
images/breeze/output_static-checks.txt | 2 +-
.../pre_commit_check_common_sql_dependency.py | 14 ++---
.../databricks/hooks/test_databricks_sql.py | 49 +++++++++++-----
.../databricks/operators/test_databricks_sql.py | 2 +
tests/providers/odbc/hooks/test_odbc.py | 8 ++-
22 files changed, 191 insertions(+), 108 deletions(-)
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 18c1d7d64c..3cabc58618 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 f16d119f6e..42aeb7d232 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/airflow/providers/common/sql/hooks/sql.py b/airflow/providers/common/sql/hooks/sql.py
index 0d4b112b32..ee3387a76c 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 (
@@ -24,6 +26,7 @@ from typing import (
Callable,
Generator,
Iterable,
+ List,
Mapping,
Protocol,
Sequence,
@@ -36,7 +39,7 @@ from urllib.parse import urlparse
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:
@@ -122,10 +125,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
- JSON-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
@@ -308,7 +311,7 @@ class DbApiHook(BaseHook):
handler: Callable[[Any], T] = ...,
split_statements: bool = ...,
return_last: bool = ...,
- ) -> T | list[T]:
+ ) -> tuple | list[tuple] | list[list[tuple] | tuple] | None:
...
def run(
@@ -319,7 +322,7 @@ class DbApiHook(BaseHook):
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
@@ -395,7 +398,7 @@ class DbApiHook(BaseHook):
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
@@ -415,19 +418,31 @@ class DbApiHook(BaseHook):
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: 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
- 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 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,
+ )
- 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.
- """
- 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/common/sql/provider.yaml b/airflow/providers/common/sql/provider.yaml
index 92ce645375..1a64721e2c 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
diff --git a/airflow/providers/databricks/hooks/databricks_sql.py b/airflow/providers/databricks/hooks/databricks_sql.py
index dc728c5ed7..6c31691c45 100644
--- a/airflow/providers/databricks/hooks/databricks_sql.py
+++ b/airflow/providers/databricks/hooks/databricks_sql.py
@@ -16,19 +16,32 @@
# under the License.
from __future__ import annotations
+import warnings
+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,
+ 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
+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
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")
@@ -52,6 +65,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_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
"""
@@ -68,6 +85,7 @@ class DatabricksSqlHook(BaseDatabricksHook, DbApiHook):
catalog: str | None = None,
schema: str | None = None,
caller: str = "DatabricksSqlHook",
+ return_tuple: bool = False,
**kwargs,
) -> None:
super().__init__(databricks_conn_id, caller=caller)
@@ -80,8 +98,18 @@ class DatabricksSqlHook(BaseDatabricksHook, DbApiHook):
self.http_headers = http_headers
self.catalog = catalog
self.schema = schema
+ self.return_tuple = return_tuple
self.additional_params = kwargs
+ 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_tuple=True` to
+ enable this behavior.""",
+ 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]:
@@ -167,7 +195,7 @@ class DatabricksSqlHook(BaseDatabricksHook, DbApiHook):
handler: Callable[[Any], T] = ...,
split_statements: bool = ...,
return_last: bool = ...,
- ) -> T | list[T]:
+ ) -> tuple | list[tuple] | list[list[tuple] | tuple] | None:
...
def run(
@@ -178,7 +206,7 @@ class DatabricksSqlHook(BaseDatabricksHook, DbApiHook):
handler: Callable[[Any], T] | None = None,
split_statements: bool = True,
return_last: bool = True,
- ) -> T | list[T] | None:
+ ) -> tuple | list[tuple] | list[list[tuple] | tuple] | None:
"""
Run a command or a list of commands.
@@ -223,7 +251,12 @@ class DatabricksSqlHook(BaseDatabricksHook, DbApiHook):
with closing(conn.cursor()) as cur:
self._run_command(cur, sql_statement, parameters)
if handler is not None:
- result = self._make_serializable(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]
@@ -241,14 +274,20 @@ class DatabricksSqlHook(BaseDatabricksHook, DbApiHook):
else:
return results
- @staticmethod
- def _make_serializable(result):
- """Transform the databricks Row objects into JSON-serializable lists."""
+ 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 isinstance(result, list):
- return [list(row) for row in result]
- elif isinstance(result, Row):
- return list(result)
- return 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])
+ 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/databricks/operators/databricks_sql.py b/airflow/providers/databricks/operators/databricks_sql.py
index a03cfa729c..96cadd827d 100644
--- a/airflow/providers/databricks/operators/databricks_sql.py
+++ b/airflow/providers/databricks/operators/databricks_sql.py
@@ -113,6 +113,7 @@ class DatabricksSqlOperator(SQLExecuteQueryOperator):
"catalog": self.catalog,
"schema": self.schema,
"caller": "DatabricksSqlOperator",
+ "return_tuple": True,
**self.client_parameters,
**self.hook_params,
}
diff --git a/airflow/providers/databricks/provider.yaml b/airflow/providers/databricks/provider.yaml
index 84cc7b4593..53c52bd02f 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/exasol/hooks/exasol.py b/airflow/providers/exasol/hooks/exasol.py
index 98955a2579..6d52b5122b 100644
--- a/airflow/providers/exasol/hooks/exasol.py
+++ b/airflow/providers/exasol/hooks/exasol.py
@@ -183,7 +183,7 @@ class ExasolHook(DbApiHook):
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 @@ class ExasolHook(DbApiHook):
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 @@ class ExasolHook(DbApiHook):
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 d134431a98..89b399e428 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 a913626a00..80ea09c9e8 100644
--- a/airflow/providers/odbc/hooks/odbc.py
+++ b/airflow/providers/odbc/hooks/odbc.py
@@ -17,10 +17,10 @@
"""This module contains ODBC hook."""
from __future__ import annotations
-from typing import Any, NamedTuple
+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
@@ -195,9 +195,9 @@ class OdbcHook(DbApiHook):
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
@property
@@ -228,17 +228,15 @@ class OdbcHook(DbApiHook):
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 JSON-serializable NamedTuple."""
+ 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 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]
- 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]
- return result
+ # 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, Sequence):
+ field_names = [col[:2] for col in result[0].cursor_description]
+ row_object = NamedTuple("Row", field_names) # type: ignore[misc]
+ 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/odbc/provider.yaml b/airflow/providers/odbc/provider.yaml
index d83f4cf689..30e7f4f0ca 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/airflow/providers/oracle/hooks/oracle.py b/airflow/providers/oracle/hooks/oracle.py
index e9c333a347..de78ab726f 100644
--- a/airflow/providers/oracle/hooks/oracle.py
+++ b/airflow/providers/oracle/hooks/oracle.py
@@ -372,7 +372,7 @@ class OracleHook(DbApiHook):
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.
diff --git a/airflow/providers/snowflake/hooks/snowflake.py b/airflow/providers/snowflake/hooks/snowflake.py
index ead3e92274..4b0d13b5e5 100644
--- a/airflow/providers/snowflake/hooks/snowflake.py
+++ b/airflow/providers/snowflake/hooks/snowflake.py
@@ -323,7 +323,7 @@ class SnowflakeHook(DbApiHook):
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 @@ class SnowflakeHook(DbApiHook):
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 @@ class SnowflakeHook(DbApiHook):
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 b7355e69b4..ee6fa214ac 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/dev/breeze/src/airflow_breeze/pre_commit_ids.py b/dev/breeze/src/airflow_breeze/pre_commit_ids.py
index 1297448cab..2806f0f1e9 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 @@ PRE_COMMIT_LIST = [
"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 @@ PRE_COMMIT_LIST = [
"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/generated/provider_dependencies.json b/generated/provider_dependencies.json
index b755f6c839..b20543ab35 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"
@@ -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"
@@ -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"
],
@@ -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"
diff --git a/images/breeze/output_static-checks.svg b/images/breeze/output_static-checks.svg
index 290f6a3264..6d0d3ea168 100644
--- a/images/breeze/output_static-checks.svg
+++ b/images/breeze/output_static-checks.svg
@@ -313,24 +313,24 @@
</text><text class="breeze-static-checks-r5" x="0" y="264" textLength="12.2" clip-path="url(#breeze-static-checks-line-10)">│</text><text class="breeze-static-checks-r7" x="451.4" y="264" textLength="988.2" clip-path="url(#breeze-static-checks-line-10)">check-base-operator-partial-arguments | check-base-operator-usage |              </text><text class="breeze-static-checks-r5" x="1451.8" y="264" textLeng [...]
</text><text class="breeze-static-checks-r5" x="0" y="288.4" textLength="12.2" clip-path="url(#breeze-static-checks-line-11)">│</text><text class="breeze-static-checks-r7" x="451.4" y="288.4" textLength="988.2" clip-path="url(#breeze-static-checks-line-11)">check-boring-cyborg-configuration | check-breeze-top-dependencies-limited |      </text><text class="breeze-static-checks-r5" x="1451.8" y="288.4" textLength="12.2" clip-path="url(#breeze-s [...]
</text><text class="breeze-static-checks-r5" x="0" y="312.8" textLength="12.2" clip-path="url(#breeze-static-checks-line-12)">│</text><text class="breeze-static-checks-r7" x="451.4" y="312.8" textLength="988.2" clip-path="url(#breeze-static-checks-line-12)">check-builtin-literals | check-changelog-has-no-duplicates |                     </text><text class="breeze-static [...]
-</text><text class="breeze-static-checks-r5" x="0" y="337.2" textLength="12.2" clip-path="url(#breeze-static-checks-line-13)">│</text><text class="breeze-static-checks-r7" x="451.4" y="337.2" textLength="988.2" clip-path="url(#breeze-static-checks-line-13)">check-cncf-k8s-only-for-executors | check-common-sql-dependency-make-serializable</text><text class="breeze-static-checks-r5" x="1451.8" y="337.2" textLength="12.2" clip-path="url(#breeze-static-checks-line-13)">│</text><tex [...]
-</text><text class="breeze-static-checks-r5" x="0" y="361.6" textLength="12.2" clip-path="url(#breeze-static-checks-line-14)">│</text><text class="breeze-static-checks-r7" x="451.4" y="361.6" textLength="988.2" clip-path="url(#breeze-static-checks-line-14)">| check-core-deprecation-classes | check-daysago-import-from-utils |             </text><text class="breeze-static-checks-r5" x="1451.8" y="361.6" te [...]
-</text><text class="breeze-static-checks-r5" x="0" y="386" textLength="12.2" clip-path="url(#breeze-static-checks-line-15)">│</text><text class="breeze-static-checks-r7" x="451.4" y="386" textLength="988.2" clip-path="url(#breeze-static-checks-line-15)">check-decorated-operator-implements-custom-name | check-deferrable-default-value </text><text class="breeze-static-checks-r5" x="1451.8" y="386" textLength="12.2" clip-path="url(#breeze-static-checks-line-15)">│</text><text [...]
-</text><text class="breeze-static-checks-r5" x="0" y="410.4" textLength="12.2" clip-path="url(#breeze-static-checks-line-16)">│</text><text class="breeze-static-checks-r7" x="451.4" y="410.4" textLength="988.2" clip-path="url(#breeze-static-checks-line-16)">| check-docstring-param-types | check-example-dags-urls |                        </text><text [...]
-</text><text class="breeze-static-checks-r5" x="0" y="434.8" textLength="12.2" clip-path="url(#breeze-static-checks-line-17)">│</text><text class="breeze-static-checks-r7" x="451.4" y="434.8" textLength="988.2" clip-path="url(#breeze-static-checks-line-17)">check-executables-have-shebangs | check-extra-packages-references |              </text><text class="breeze-static-checks-r5" x="1451.8" y="434.8" te [...]
-</text><text class="breeze-static-checks-r5" x="0" y="459.2" textLength="12.2" clip-path="url(#breeze-static-checks-line-18)">│</text><text class="breeze-static-checks-r7" x="451.4" y="459.2" textLength="988.2" clip-path="url(#breeze-static-checks-line-18)">check-extras-order | check-fab-migrations | check-for-inclusive-language |       </text><text class="breeze-static-checks-r5" x="1451.8" y="459.2" textLength="12.2" clip-path [...]
-</text><text class="breeze-static-checks-r5" x="0" y="483.6" textLength="12.2" clip-path="url(#breeze-static-checks-line-19)">│</text><text class="breeze-static-checks-r7" x="451.4" y="483.6" textLength="988.2" clip-path="url(#breeze-static-checks-line-19)">check-google-re2-as-dependency | check-hooks-apply |                             [...]
-</text><text class="breeze-static-checks-r5" x="0" y="508" textLength="12.2" clip-path="url(#breeze-static-checks-line-20)">│</text><text class="breeze-static-checks-r7" x="451.4" y="508" textLength="988.2" clip-path="url(#breeze-static-checks-line-20)">check-incorrect-use-of-LoggingMixin | check-init-decorator-arguments |           </text><text class="breeze-static-checks-r5" x="1451.8" y="508" textLength="12.2" clip- [...]
-</text><text class="breeze-static-checks-r5" x="0" y="532.4" textLength="12.2" clip-path="url(#breeze-static-checks-line-21)">│</text><text class="breeze-static-checks-r7" x="451.4" y="532.4" textLength="988.2" clip-path="url(#breeze-static-checks-line-21)">check-lazy-logging | check-links-to-example-dags-do-not-use-hardcoded-versions | </text><text class="breeze-static-checks-r5" x="1451.8" y="532.4" textLength="12.2" clip-path="url(#breeze-static-checks-line-21)">│< [...]
-</text><text class="breeze-static-checks-r5" x="0" y="556.8" textLength="12.2" clip-path="url(#breeze-static-checks-line-22)">│</text><text class="breeze-static-checks-r7" x="451.4" y="556.8" textLength="988.2" clip-path="url(#breeze-static-checks-line-22)">check-merge-conflict | check-newsfragments-are-valid |                           </t [...]
-</text><text class="breeze-static-checks-r5" x="0" y="581.2" textLength="12.2" clip-path="url(#breeze-static-checks-line-23)">│</text><text class="breeze-static-checks-r7" x="451.4" y="581.2" textLength="988.2" clip-path="url(#breeze-static-checks-line-23)">check-no-airflow-deprecation-in-providers | check-no-providers-in-core-examples |</text><text class="breeze-static-checks-r5" x="1451.8" y="581.2" textLength="12.2" clip-path="url(#breeze-static-checks-line-23)">│</text [...]
-</text><text class="breeze-static-checks-r5" x="0" y="605.6" textLength="12.2" clip-path="url(#breeze-static-checks-line-24)">│</text><text class="breeze-static-checks-r7" x="451.4" y="605.6" textLength="988.2" clip-path="url(#breeze-static-checks-line-24)">check-no-relative-imports | check-only-new-session-with-provide-session |        </text><text class="breeze-static-checks-r5" x="1451.8" y="605.6" textLength="12.2" clip-path="url [...]
-</text><text class="breeze-static-checks-r5" x="0" y="630" textLength="12.2" clip-path="url(#breeze-static-checks-line-25)">│</text><text class="breeze-static-checks-r7" x="451.4" y="630" textLength="988.2" clip-path="url(#breeze-static-checks-line-25)">check-persist-credentials-disabled-in-github-workflows |                         </text><text class="breeze- [...]
-</text><text class="breeze-static-checks-r5" x="0" y="654.4" textLength="12.2" clip-path="url(#breeze-static-checks-line-26)">│</text><text class="breeze-static-checks-r7" x="451.4" y="654.4" textLength="988.2" clip-path="url(#breeze-static-checks-line-26)">check-pre-commit-information-consistent | check-provide-create-sessions-imports |</text><text class="breeze-static-checks-r5" x="1451.8" y="654.4" textLength="12.2" clip-path="url(#breeze-static-checks-line-26)">│</text [...]
-</text><text class="breeze-static-checks-r5" x="0" y="678.8" textLength="12.2" clip-path="url(#breeze-static-checks-line-27)">│</text><text class="breeze-static-checks-r7" x="451.4" y="678.8" textLength="988.2" clip-path="url(#breeze-static-checks-line-27)">check-provider-docs-valid | check-provider-yaml-valid |                          </text>< [...]
-</text><text class="breeze-static-checks-r5" x="0" y="703.2" textLength="12.2" clip-path="url(#breeze-static-checks-line-28)">│</text><text class="breeze-static-checks-r7" x="451.4" y="703.2" textLength="988.2" clip-path="url(#breeze-static-checks-line-28)">check-providers-init-file-missing | check-providers-subpackages-init-file-exist |</text><text class="breeze-static-checks-r5" x="1451.8" y="703.2" textLength="12.2" clip-path="url(#breeze-static-checks-line-28)">│</text [...]
-</text><text class="breeze-static-checks-r5" x="0" y="727.6" textLength="12.2" clip-path="url(#breeze-static-checks-line-29)">│</text><text class="breeze-static-checks-r7" x="451.4" y="727.6" textLength="988.2" clip-path="url(#breeze-static-checks-line-29)">check-pydevd-left-in-code | check-revision-heads-map |                           </t [...]
-</text><text class="breeze-static-checks-r5" x="0" y="752" textLength="12.2" clip-path="url(#breeze-static-checks-line-30)">│</text><text class="breeze-static-checks-r7" x="451.4" y="752" textLength="988.2" clip-path="url(#breeze-static-checks-line-30)">check-safe-filter-usage-in-html | check-setup-order |                            </ [...]
+</text><text class="breeze-static-checks-r5" x="0" y="337.2" textLength="12.2" clip-path="url(#breeze-static-checks-line-13)">│</text><text class="breeze-static-checks-r7" x="451.4" y="337.2" textLength="988.2" clip-path="url(#breeze-static-checks-line-13)">check-cncf-k8s-only-for-executors | check-core-deprecation-classes |             </text><text class="breeze-static-checks-r5" x="1451.8" y="337.2" textLen [...]
+</text><text class="breeze-static-checks-r5" x="0" y="361.6" textLength="12.2" clip-path="url(#breeze-static-checks-line-14)">│</text><text class="breeze-static-checks-r7" x="451.4" y="361.6" textLength="988.2" clip-path="url(#breeze-static-checks-line-14)">check-daysago-import-from-utils | check-decorated-operator-implements-custom-name</text><text class="breeze-static-checks-r5" x="1451.8" y="361.6" textLength="12.2" clip-path="url(#breeze-static-checks-line-14)">│</text><tex [...]
+</text><text class="breeze-static-checks-r5" x="0" y="386" textLength="12.2" clip-path="url(#breeze-static-checks-line-15)">│</text><text class="breeze-static-checks-r7" x="451.4" y="386" textLength="988.2" clip-path="url(#breeze-static-checks-line-15)">| check-deferrable-default-value | check-docstring-param-types |                 </text><text class="breeze-static-checks-r5" x="1451 [...]
+</text><text class="breeze-static-checks-r5" x="0" y="410.4" textLength="12.2" clip-path="url(#breeze-static-checks-line-16)">│</text><text class="breeze-static-checks-r7" x="451.4" y="410.4" textLength="988.2" clip-path="url(#breeze-static-checks-line-16)">check-example-dags-urls | check-executables-have-shebangs |                      </text><text class="breeze-s [...]
+</text><text class="breeze-static-checks-r5" x="0" y="434.8" textLength="12.2" clip-path="url(#breeze-static-checks-line-17)">│</text><text class="breeze-static-checks-r7" x="451.4" y="434.8" textLength="988.2" clip-path="url(#breeze-static-checks-line-17)">check-extra-packages-references | check-extras-order | check-fab-migrations |    </text><text class="breeze-static-checks-r5" x="1451.8" y="434.8" textLength="12.2" clip-path="url(#breeze-s [...]
+</text><text class="breeze-static-checks-r5" x="0" y="459.2" textLength="12.2" clip-path="url(#breeze-static-checks-line-18)">│</text><text class="breeze-static-checks-r7" x="451.4" y="459.2" textLength="988.2" clip-path="url(#breeze-static-checks-line-18)">check-for-inclusive-language | check-google-re2-as-dependency | check-hooks-apply</text><text class="breeze-static-checks-r5" x="1451.8" y="459.2" textLength="12.2" clip-path="url(#breeze-static-checks-line-18)">│< [...]
+</text><text class="breeze-static-checks-r5" x="0" y="483.6" textLength="12.2" clip-path="url(#breeze-static-checks-line-19)">│</text><text class="breeze-static-checks-r7" x="451.4" y="483.6" textLength="988.2" clip-path="url(#breeze-static-checks-line-19)">| check-incorrect-use-of-LoggingMixin | check-init-decorator-arguments |         </text><text class="breeze-static-checks-r5" x="1451.8" y="483.6" textLength="12.2" clip [...]
+</text><text class="breeze-static-checks-r5" x="0" y="508" textLength="12.2" clip-path="url(#breeze-static-checks-line-20)">│</text><text class="breeze-static-checks-r7" x="451.4" y="508" textLength="988.2" clip-path="url(#breeze-static-checks-line-20)">check-lazy-logging | check-links-to-example-dags-do-not-use-hardcoded-versions | </text><text class="breeze-static-checks-r5" x="1451.8" y="508" textLength="12.2" clip-path="url(#breeze-static-checks-line-20)">│</text> [...]
+</text><text class="breeze-static-checks-r5" x="0" y="532.4" textLength="12.2" clip-path="url(#breeze-static-checks-line-21)">│</text><text class="breeze-static-checks-r7" x="451.4" y="532.4" textLength="988.2" clip-path="url(#breeze-static-checks-line-21)">check-merge-conflict | check-newsfragments-are-valid |                           </t [...]
+</text><text class="breeze-static-checks-r5" x="0" y="556.8" textLength="12.2" clip-path="url(#breeze-static-checks-line-22)">│</text><text class="breeze-static-checks-r7" x="451.4" y="556.8" textLength="988.2" clip-path="url(#breeze-static-checks-line-22)">check-no-airflow-deprecation-in-providers | check-no-providers-in-core-examples |</text><text class="breeze-static-checks-r5" x="1451.8" y="556.8" textLength="12.2" clip-path="url(#breeze-static-checks-line-22)">│</text [...]
+</text><text class="breeze-static-checks-r5" x="0" y="581.2" textLength="12.2" clip-path="url(#breeze-static-checks-line-23)">│</text><text class="breeze-static-checks-r7" x="451.4" y="581.2" textLength="988.2" clip-path="url(#breeze-static-checks-line-23)">check-no-relative-imports | check-only-new-session-with-provide-session |        </text><text class="breeze-static-checks-r5" x="1451.8" y="581.2" textLength="12.2" clip-path="url [...]
+</text><text class="breeze-static-checks-r5" x="0" y="605.6" textLength="12.2" clip-path="url(#breeze-static-checks-line-24)">│</text><text class="breeze-static-checks-r7" x="451.4" y="605.6" textLength="988.2" clip-path="url(#breeze-static-checks-line-24)">check-persist-credentials-disabled-in-github-workflows |                         </text><text class="bre [...]
+</text><text class="breeze-static-checks-r5" x="0" y="630" textLength="12.2" clip-path="url(#breeze-static-checks-line-25)">│</text><text class="breeze-static-checks-r7" x="451.4" y="630" textLength="988.2" clip-path="url(#breeze-static-checks-line-25)">check-pre-commit-information-consistent | check-provide-create-sessions-imports |</text><text class="breeze-static-checks-r5" x="1451.8" y="630" textLength="12.2" clip-path="url(#breeze-static-checks-line-25)">│</text><text [...]
+</text><text class="breeze-static-checks-r5" x="0" y="654.4" textLength="12.2" clip-path="url(#breeze-static-checks-line-26)">│</text><text class="breeze-static-checks-r7" x="451.4" y="654.4" textLength="988.2" clip-path="url(#breeze-static-checks-line-26)">check-provider-docs-valid | check-provider-yaml-valid |                          </text>< [...]
+</text><text class="breeze-static-checks-r5" x="0" y="678.8" textLength="12.2" clip-path="url(#breeze-static-checks-line-27)">│</text><text class="breeze-static-checks-r7" x="451.4" y="678.8" textLength="988.2" clip-path="url(#breeze-static-checks-line-27)">check-providers-init-file-missing | check-providers-subpackages-init-file-exist |</text><text class="breeze-static-checks-r5" x="1451.8" y="678.8" textLength="12.2" clip-path="url(#breeze-static-checks-line-27)">│</text [...]
+</text><text class="breeze-static-checks-r5" x="0" y="703.2" textLength="12.2" clip-path="url(#breeze-static-checks-line-28)">│</text><text class="breeze-static-checks-r7" x="451.4" y="703.2" textLength="988.2" clip-path="url(#breeze-static-checks-line-28)">check-pydevd-left-in-code | check-revision-heads-map |                           </t [...]
+</text><text class="breeze-static-checks-r5" x="0" y="727.6" textLength="12.2" clip-path="url(#breeze-static-checks-line-29)">│</text><text class="breeze-static-checks-r7" x="451.4" y="727.6" textLength="988.2" clip-path="url(#breeze-static-checks-line-29)">check-safe-filter-usage-in-html | check-setup-order |                            [...]
+</text><text class="breeze-static-checks-r5" x="0" y="752" textLength="12.2" clip-path="url(#breeze-static-checks-line-30)">│</text><text class="breeze-static-checks-r7" x="451.4" y="752" textLength="988.2" clip-path="url(#breeze-static-checks-line-30)">check-sql-dependency-common-data-structure |                                [...]
</text><text class="breeze-static-checks-r5" x="0" y="776.4" textLength="12.2" clip-path="url(#breeze-static-checks-line-31)">│</text><text class="breeze-static-checks-r7" x="451.4" y="776.4" textLength="988.2" clip-path="url(#breeze-static-checks-line-31)">check-start-date-not-used-in-defaults | check-system-tests-present |             </text><text class="breeze-static-checks-r5" x="1451.8" y="776.4" textLen [...]
</text><text class="breeze-static-checks-r5" x="0" y="800.8" textLength="12.2" clip-path="url(#breeze-static-checks-line-32)">│</text><text class="breeze-static-checks-r7" x="451.4" y="800.8" textLength="988.2" clip-path="url(#breeze-static-checks-line-32)">check-system-tests-tocs | check-tests-unittest-testcase |                        </text><text class [...]
</text><text class="breeze-static-checks-r5" x="0" y="825.2" textLength="12.2" clip-path="url(#breeze-static-checks-line-33)">│</text><text class="breeze-static-checks-r7" x="451.4" y="825.2" textLength="988.2" clip-path="url(#breeze-static-checks-line-33)">check-urlparse-usage-in-code | check-usage-of-re2-over-re | check-xml | codespell</text><text class="breeze-static-checks-r5" x="1451.8" y="825.2" textLength="12.2" clip-path="url(#breeze-static-checks-li [...]
diff --git a/images/breeze/output_static-checks.txt b/images/breeze/output_static-checks.txt
index 5b32905ea3..9ffe4f833a 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 4b335d1cf6..9719310a71 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,9 +31,9 @@ console = Console(color_system="standard", width=200)
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"
+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"
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."
)
diff --git a/tests/providers/databricks/hooks/test_databricks_sql.py b/tests/providers/databricks/hooks/test_databricks_sql.py
index 64cd0b9c06..a5d85880be 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,8 +59,12 @@ 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"]) # type: ignore[name-match]
+
+
@pytest.mark.parametrize(
- "return_last, split_statements, sql, cursor_calls,"
+ "return_last, split_statements, sql, cursor_calls, return_tuple,"
"cursor_descriptions, cursor_results, hook_descriptions, hook_results, ",
[
pytest.param(
@@ -67,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",)]],
- [[1, 2], [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(
@@ -78,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",)]],
- [[1, 2], [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(
@@ -89,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",)]],
- [[1, 2], [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(
@@ -100,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",)]],
- [[[1, 2], [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(
@@ -111,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",)]],
- [[3, 4], [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(
@@ -122,10 +132,14 @@ 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",)]],
- [[[1, 2], [11, 12]], [[3, 4], [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",
),
pytest.param(
@@ -133,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",)]],
- [[[1, 2], [11, 12]]],
+ [[Row(id=1, value=2), Row(id=11, value=12)]],
id="The return_last set on single query in list",
),
pytest.param(
@@ -144,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",)]],
- [[[1, 2], [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(
@@ -155,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",)]],
- [[3, 4], [13, 14]],
+ [Row(id=3, value=4), Row(id=13, value=14)],
id="The return_last set on multiple queries in list",
),
pytest.param(
@@ -166,10 +183,14 @@ 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",)]],
- [[[1, 2], [11, 12]], [[3, 4], [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",
),
pytest.param(
@@ -177,20 +198,21 @@ 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",)]],
- [1, 2],
+ SerializableRow(1, 2),
id="The return_last set and no split statements set on single query in string",
),
],
)
def test_query(
- databricks_hook,
return_last,
split_statements,
sql,
cursor_calls,
+ return_tuple,
cursor_descriptions,
cursor_results,
hook_descriptions,
@@ -227,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_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 e7885740cf..d68dda41ab 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_tuple=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_tuple=True,
session_configuration=None,
sql_endpoint_name=None,
http_headers=None,
diff --git a/tests/providers/odbc/hooks/test_odbc.py b/tests/providers/odbc/hooks/test_odbc.py
index 8c0fbee40f..391a766df9 100644
--- a/tests/providers/odbc/hooks/test_odbc.py
+++ b/tests/providers/odbc/hooks/test_odbc.py
@@ -310,7 +310,9 @@ class TestOdbcHook:
"""
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.
@@ -322,7 +324,9 @@ class TestOdbcHook:
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(