This is an automated email from the ASF dual-hosted git repository.
eladkal 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 81dbef9f4bd Align hook run() annotations with None-able handler
results (#69230)
81dbef9f4bd is described below
commit 81dbef9f4bdd648e5ddf12709a50d460cad69b21
Author: Noritaka Sekiyama <[email protected]>
AuthorDate: Thu Jul 2 16:31:24 2026 +0900
Align hook run() annotations with None-able handler results (#69230)
* Fix DatabricksSqlHook TypeError when query returns no rows with fetch_one
handler
DBAPI fetchone() returns None on an empty result set, but the hook's
_make_common_data_structure only accepted Row objects or lists of them,
so run() with a scalar handler crashed on empty results instead of
returning None as the common.sql contract and other providers do.
* Widen DbApiHook._make_common_data_structure return type to include None
The base implementation passes the handler result through unchanged, and
scalar handlers such as fetch_one_handler legitimately return None for
empty result sets, so the annotation excluded a value the method already
returns. Mypy's override check rejected the Databricks subclass returning
None against the narrower supertype annotation.
* Align exasol and snowflake run() annotations with None-able handler
results
These hooks copy the base run() implementation, so their returned result
lists can also contain None once the base _make_common_data_structure
annotation admits it; mypy rejected the narrower list element type.
---
.../sql/src/airflow/providers/common/sql/hooks/sql.py | 2 +-
.../providers/databricks/hooks/databricks_sql.py | 10 +++++++---
.../unit/databricks/hooks/test_databricks_sql.py | 19 ++++++++++++++++++-
.../src/airflow/providers/exasol/hooks/exasol.py | 4 ++--
.../airflow/providers/snowflake/hooks/snowflake.py | 4 ++--
5 files changed, 30 insertions(+), 9 deletions(-)
diff --git a/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py
b/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py
index 7a02f42c62e..86fb4a466f9 100644
--- a/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py
+++ b/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py
@@ -839,7 +839,7 @@ class DbApiHook(BaseHook):
return _last_result
return results
- def _make_common_data_structure(self, result: T | Sequence[T]) -> tuple |
list[tuple]:
+ def _make_common_data_structure(self, result: T | Sequence[T]) -> tuple |
list[tuple] | None:
"""
Ensure the data returned from an SQL command is a standard tuple or
list[tuple].
diff --git
a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_sql.py
b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_sql.py
index ffc089607c3..ce2bc28b459 100644
---
a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_sql.py
+++
b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_sql.py
@@ -298,7 +298,7 @@ class DatabricksSqlHook(BaseDatabricksHook, DbApiHook):
split_statements: bool = ...,
return_last: bool = ...,
execution_timeout: timedelta | None = None,
- ) -> tuple | list[tuple] | list[list[tuple] | tuple] | None: ...
+ ) -> tuple | list[tuple] | list[list[tuple] | tuple | None] | None: ...
def run(
self,
@@ -309,7 +309,7 @@ class DatabricksSqlHook(BaseDatabricksHook, DbApiHook):
split_statements: bool = True,
return_last: bool = True,
execution_timeout: timedelta | None = None,
- ) -> tuple | list[tuple] | list[list[tuple] | tuple] | None:
+ ) -> tuple | list[tuple] | list[list[tuple] | tuple | None] | None:
"""
Run a command or a list of commands.
@@ -401,10 +401,14 @@ class DatabricksSqlHook(BaseDatabricksHook, DbApiHook):
return results[-1]
return results
- def _make_common_data_structure(self, result: T | Sequence[T]) ->
tuple[Any, ...] | list[tuple[Any, ...]]:
+ def _make_common_data_structure(
+ self, result: T | Sequence[T] | None
+ ) -> tuple[Any, ...] | list[tuple[Any, ...]] | 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 result is None:
+ return None
if isinstance(result, list):
rows: Sequence[Row] = result
if not rows:
diff --git
a/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py
b/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py
index 6fd835b7872..ec0718025ea 100644
--- a/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py
+++ b/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py
@@ -31,7 +31,7 @@ from databricks.sql.types import Row
from airflow.models import Connection
from airflow.providers.common.compat.sdk import AirflowException,
AirflowOptionalProviderFeatureException
-from airflow.providers.common.sql.hooks.handlers import fetch_all_handler
+from airflow.providers.common.sql.hooks.handlers import fetch_all_handler,
fetch_one_handler
from airflow.providers.databricks.hooks.databricks_sql import (
DatabricksSqlHook,
_format_query_tag_value,
@@ -377,6 +377,23 @@ def test_query(
cur.close.assert_called()
+def test_make_common_data_structure_none_result():
+ assert DatabricksSqlHook()._make_common_data_structure(None) is None
+
+
+def test_query_with_fetch_one_handler_on_empty_result(mock_get_conn,
mock_get_requests):
+ conn = mock.MagicMock()
+ cur = mock.MagicMock(rowcount=0,
description=get_cursor_descriptions(["id"]))
+ cur.fetchone.return_value = None
+ conn.cursor.return_value = cur
+ mock_get_conn.side_effect = [conn]
+
+ databricks_hook = DatabricksSqlHook(sql_endpoint_name="Test")
+ result = databricks_hook.run(sql="select * from test.test",
handler=fetch_one_handler)
+
+ assert result is None
+
+
@pytest.mark.parametrize(
"empty_statement",
[
diff --git a/providers/exasol/src/airflow/providers/exasol/hooks/exasol.py
b/providers/exasol/src/airflow/providers/exasol/hooks/exasol.py
index 0d418bc798e..476c5259ba5 100644
--- a/providers/exasol/src/airflow/providers/exasol/hooks/exasol.py
+++ b/providers/exasol/src/airflow/providers/exasol/hooks/exasol.py
@@ -287,7 +287,7 @@ class ExasolHook(DbApiHook):
handler: Callable[[Any], T] = ...,
split_statements: bool = ...,
return_last: bool = ...,
- ) -> tuple | list[tuple] | list[list[tuple] | tuple] | None: ...
+ ) -> tuple | list[tuple] | list[list[tuple] | tuple | None] | None: ...
def run(
self,
@@ -297,7 +297,7 @@ class ExasolHook(DbApiHook):
handler: Callable[[Any], T] | None = None,
split_statements: bool = False,
return_last: bool = True,
- ) -> tuple | list[tuple] | list[list[tuple] | tuple] | None:
+ ) -> tuple | list[tuple] | list[list[tuple] | tuple | None] | None:
"""
Run a command or a list of commands.
diff --git
a/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py
b/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py
index 15f3cf6f90a..fdd55893ca0 100644
--- a/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py
+++ b/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py
@@ -893,7 +893,7 @@ class SnowflakeHook(DbApiHook):
split_statements: bool = ...,
return_last: bool = ...,
return_dictionaries: bool = ...,
- ) -> tuple | list[tuple] | list[list[tuple] | tuple] | None: ...
+ ) -> tuple | list[tuple] | list[list[tuple] | tuple | None] | None: ...
def run(
self,
@@ -904,7 +904,7 @@ class SnowflakeHook(DbApiHook):
split_statements: bool = True,
return_last: bool = True,
return_dictionaries: bool = False,
- ) -> tuple | list[tuple] | list[list[tuple] | tuple] | None:
+ ) -> tuple | list[tuple] | list[list[tuple] | tuple | None] | None:
"""
Run a command or list of commands.
