Refine type overload for common-sql

providers/common/sql/src/airflow/providers/common/sql/hooks/sql.pyi

@@ -37,8 +37,10 @@ from functools import cached_property as cached_property
 from typing import Any, Callable, Protocol, TypeVar, overload
 
 from _typeshed import Incomplete as Incomplete
-from pandas import DataFrame as DataFrame
+from pandas import DataFrame as PandasDataFrame
+from polars import DataFrame as PolarsDataFrame
 from sqlalchemy.engine import URL as URL, Engine as Engine, Inspector as Inspector
+from typing_extensions import Literal
 
 from airflow.hooks.base import BaseHook as BaseHook
 from airflow.models import Connection as Connection
@@ -105,10 +107,65 @@ class DbApiHook(BaseHook):
     def get_reserved_words(self, dialect_name: str) -> set[str]: ...
     def get_pandas_df(
         self, sql, parameters: list | tuple | Mapping[str, Any] | None = None, **kwargs
-    ) -> DataFrame: ...
+    ) -> PandasDataFrame: ...
     def get_pandas_df_by_chunks(
         self, sql, parameters: list | tuple | Mapping[str, Any] | None = None, *, chunksize: int, **kwargs
-    ) -> Generator[DataFrame, None, None]: ...
+    ) -> Generator[PandasDataFrame, None, None]: ...
+    @overload
+    def get_df(
+        self,
+        sql: str | list[str],
+        parameters: list | tuple | Mapping[str, Any] | None = None,
+        *,
+        df_type: Literal["pandas"] = "pandas",
+        **kwargs: Any,
+    ) -> PandasDataFrame: ...
+    @overload
+    def get_df(
+        self,
+        sql: str | list[str],
+        parameters: list | tuple | Mapping[str, Any] | None = None,
+        *,
+        df_type: Literal["polars"] = "polars",
+        **kwargs: Any,
+    ) -> PolarsDataFrame: ...
+    @overload
+    def get_df(  # fallback overload
+        self,
+        sql: str | list[str],
+        parameters: list | tuple | Mapping[str, Any] | None = None,
+        *,
+        df_type: Literal["pandas", "polars"] = "pandas",
+    ) -> PandasDataFrame | PolarsDataFrame: ...
+    @overload
+    def get_df_by_chunks(
+        self,
+        sql,
+        parameters: list | tuple | Mapping[str, Any] | None = None,
+        *,
+        chunksize: int,
+        df_type: Literal["pandas"] = "pandas",
+        **kwargs,
+    ) -> Generator[PandasDataFrame, None, None]: ...
+    @overload
+    def get_df_by_chunks(
+        self,
+        sql,
+        parameters: list | tuple | Mapping[str, Any] | None = None,
+        *,
+        chunksize: int,
+        df_type: Literal["polars"],
+        **kwargs,
+    ) -> Generator[PolarsDataFrame, None, None]: ...
+    @overload
+    def get_df_by_chunks(  # fallback overload
+        self,
+        sql,
+        parameters: list | tuple | Mapping[str, Any] | None = None,
+        *,
+        chunksize: int,
+        df_type: Literal["pandas", "polars"] = "pandas",
+    ) -> Generator[PandasDataFrame | PolarsDataFrame, None, None]: ...
     def get_records(
         self, sql: str | list[str], parameters: Iterable | Mapping[str, Any] | None = None
     ) -> Any: ...