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 5c45e8b19fa19..e988c423e1f64 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 @@ -399,7 +399,7 @@ def get_pandas_df( :param parameters: The parameters to render the SQL query with. :param kwargs: (optional) passed into pandas.io.sql.read_sql method """ - return self._get_pandas_df(sql, parameters, **kwargs) + return self.get_df(sql, parameters, df_type="pandas", **kwargs) @deprecated( reason="Replaced by function `get_df_by_chunks`.", @@ -414,7 +414,7 @@ def get_pandas_df_by_chunks( chunksize: int, **kwargs, ) -> Generator[PandasDataFrame, None, None]: - return self._get_pandas_df_by_chunks(sql, parameters, chunksize=chunksize, **kwargs) + return self.get_df_by_chunks(sql, parameters, chunksize=chunksize, df_type="pandas", **kwargs) @overload def get_df( diff --git a/providers/common/sql/tests/unit/common/sql/hooks/test_sql.py b/providers/common/sql/tests/unit/common/sql/hooks/test_sql.py index d7a6a64b1c706..5d1aff2f1ed88 100644 --- a/providers/common/sql/tests/unit/common/sql/hooks/test_sql.py +++ b/providers/common/sql/tests/unit/common/sql/hooks/test_sql.py @@ -313,11 +313,16 @@ def test_uri_with_schema(self): @pytest.mark.parametrize( "df_type, expected_type", [ + ("test_default_df_type", pd.DataFrame), ("pandas", pd.DataFrame), ("polars", pl.DataFrame), ], ) def test_get_df_with_df_type(db, df_type, expected_type): dbapi_hook = mock_db_hook(DbApiHook) - df = dbapi_hook.get_df("SQL", df_type=df_type) - assert isinstance(df, expected_type) + if df_type == "test_default_df_type": + df = dbapi_hook.get_df("SQL") + assert isinstance(df, pd.DataFrame) + else: + df = dbapi_hook.get_df("SQL", df_type=df_type) + assert isinstance(df, expected_type)