From 97d12eac5a4e929afff4b764127307ac0adfad2b Mon Sep 17 00:00:00 2001
From: DShi <shibisoma@gmail.com>
Date: Mon, 28 Oct 2024 11:29:53 -0400
Subject: [PATCH] fix: fixes static checks and moves Row import

---
 .../src/airflow/providers/common/sql/sensors/sql.py   |  2 +-
 .../providers/databricks/hooks/databricks_sql.py      | 11 ++++-------
 2 files changed, 5 insertions(+), 8 deletions(-)

diff --git a/providers/src/airflow/providers/common/sql/sensors/sql.py b/providers/src/airflow/providers/common/sql/sensors/sql.py
index b213d2da0600d..ac7eb7b4a4005 100644
--- a/providers/src/airflow/providers/common/sql/sensors/sql.py
+++ b/providers/src/airflow/providers/common/sql/sensors/sql.py
@@ -80,7 +80,7 @@ def __init__(
         self.parameters = parameters
         self.success = success
         self.failure = failure
-        self.selector = selector or itemgetter(0)
+        self.selector = selector
         self.fail_on_empty = fail_on_empty
         self.hook_params = hook_params
         super().__init__(**kwargs)
diff --git a/providers/src/airflow/providers/databricks/hooks/databricks_sql.py b/providers/src/airflow/providers/databricks/hooks/databricks_sql.py
index dd485ec586655..44568dd569e4d 100644
--- a/providers/src/airflow/providers/databricks/hooks/databricks_sql.py
+++ b/providers/src/airflow/providers/databricks/hooks/databricks_sql.py
@@ -25,17 +25,15 @@
     Any,
     Callable,
     Iterable,
-    List,
-    Optional,
     Mapping,
     Sequence,
     TypeVar,
-    Union,
     cast,
     overload,
 )
 
 from databricks import sql  # type: ignore[attr-defined]
+from databricks.sql.types import Row
 
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.models.connection import Connection as AirflowConnection
@@ -44,7 +42,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")
 
@@ -52,7 +50,6 @@
 T = TypeVar("T")
 
 
-
 class DatabricksSqlHook(BaseDatabricksHook, DbApiHook):
     """
     Hook to interact with Databricks SQL.
@@ -94,7 +91,7 @@ def __init__(
         **kwargs,
     ) -> None:
         super().__init__(databricks_conn_id, caller=caller)
-        self._sql_conn: Optional[Connection] = None
+        self._sql_conn: Connection | None = None
         self._token: str | None = None
         self._http_path = http_path
         self._sql_endpoint_name = sql_endpoint_name
@@ -275,7 +272,7 @@ def run(
         else:
             return results
 
-    def _make_common_data_structure(self, result: Union[T, Sequence[T]]) -> Union[tuple[Any, ...], list[tuple[Any, ...]]]:
+    def _make_common_data_structure(self, result: T | Sequence[T]) -> tuple[Any, ...] | list[tuple[Any, ...]]:
         """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