From 4c29116d918d8ed4fa3886c80630441a78c4f4a0 Mon Sep 17 00:00:00 2001 From: Jarek Potiuk Date: Tue, 16 Jan 2024 23:56:34 +0100 Subject: [PATCH 1/2] Fix databricks_sql hook query failing on empty result for return_tuple The hook was failing with list_index_out_of_range when returned results were empty list. Also it fixes the retunr_tuple test that was returning wrong cursor response and checked wrong hook response Fixes: #36822 --- .../providers/databricks/hooks/databricks_sql.py | 2 ++ .../databricks/hooks/test_databricks_sql.py | 16 ++++++++++++++-- 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/airflow/providers/databricks/hooks/databricks_sql.py b/airflow/providers/databricks/hooks/databricks_sql.py index d18d804b04d1f..7cb9f1b0e31bb 100644 --- a/airflow/providers/databricks/hooks/databricks_sql.py +++ b/airflow/providers/databricks/hooks/databricks_sql.py @@ -280,6 +280,8 @@ def _make_common_data_structure(self, result: Sequence[Row] | Row) -> list[tuple # instantiated namedtuple, and will never do: https://github.com/python/mypy/issues/848 if isinstance(result, list): rows: list[Row] = result + if len(rows) == 0: + return [] 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]) diff --git a/tests/providers/databricks/hooks/test_databricks_sql.py b/tests/providers/databricks/hooks/test_databricks_sql.py index a5d85880bec30..01392fb0a7dc3 100644 --- a/tests/providers/databricks/hooks/test_databricks_sql.py +++ b/tests/providers/databricks/hooks/test_databricks_sql.py @@ -200,11 +200,23 @@ def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: ["select * from test.test"], True, [["id", "value"]], - (Row(id=1, value=2),), + ([Row(id=1, value=2)],), [[("id",), ("value",)]], - SerializableRow(1, 2), + [SerializableRow(1, 2)], id="The return_last set and no split statements set on single query in string", ), + pytest.param( + True, + False, + "select * from test.test", + ["select * from test.test"], + True, + [["id", "value"]], + ([],), + [[("id",), ("value",)]], + [], + id="Empty list", + ), ], ) def test_query( From 37c755980067351e70bf3d28531874e371e63b5c Mon Sep 17 00:00:00 2001 From: Jarek Potiuk Date: Wed, 17 Jan 2024 00:06:57 +0100 Subject: [PATCH 2/2] Update airflow/providers/databricks/hooks/databricks_sql.py Co-authored-by: Jed Cunningham <66968678+jedcunningham@users.noreply.github.com> --- airflow/providers/databricks/hooks/databricks_sql.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/databricks/hooks/databricks_sql.py b/airflow/providers/databricks/hooks/databricks_sql.py index 7cb9f1b0e31bb..5ffe404e317e5 100644 --- a/airflow/providers/databricks/hooks/databricks_sql.py +++ b/airflow/providers/databricks/hooks/databricks_sql.py @@ -280,7 +280,7 @@ def _make_common_data_structure(self, result: Sequence[Row] | Row) -> list[tuple # instantiated namedtuple, and will never do: https://github.com/python/mypy/issues/848 if isinstance(result, list): rows: list[Row] = result - if len(rows) == 0: + if not rows: return [] rows_fields = rows[0].__fields__ rows_object = namedtuple("Row", rows_fields) # type: ignore[misc]