Skip to content

Commit 2bd2007

Browse files
guan404mingeladkal
andauthored
Update deprecated functions to use stable functions (#50275)
Co-authored-by: Elad Kalif <45845474+eladkal@users.noreply.github.com>
1 parent 8583fba commit 2bd2007

2 files changed

Lines changed: 9 additions & 4 deletions

File tree

  • providers/common/sql

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -399,7 +399,7 @@ def get_pandas_df(
399399
:param parameters: The parameters to render the SQL query with.
400400
:param kwargs: (optional) passed into pandas.io.sql.read_sql method
401401
"""
402-
return self._get_pandas_df(sql, parameters, **kwargs)
402+
return self.get_df(sql, parameters, df_type="pandas", **kwargs)
403403

404404
@deprecated(
405405
reason="Replaced by function `get_df_by_chunks`.",
@@ -414,7 +414,7 @@ def get_pandas_df_by_chunks(
414414
chunksize: int,
415415
**kwargs,
416416
) -> Generator[PandasDataFrame, None, None]:
417-
return self._get_pandas_df_by_chunks(sql, parameters, chunksize=chunksize, **kwargs)
417+
return self.get_df_by_chunks(sql, parameters, chunksize=chunksize, df_type="pandas", **kwargs)
418418

419419
@overload
420420
def get_df(

providers/common/sql/tests/unit/common/sql/hooks/test_sql.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -313,11 +313,16 @@ def test_uri_with_schema(self):
313313
@pytest.mark.parametrize(
314314
"df_type, expected_type",
315315
[
316+
("test_default_df_type", pd.DataFrame),
316317
("pandas", pd.DataFrame),
317318
("polars", pl.DataFrame),
318319
],
319320
)
320321
def test_get_df_with_df_type(db, df_type, expected_type):
321322
dbapi_hook = mock_db_hook(DbApiHook)
322-
df = dbapi_hook.get_df("SQL", df_type=df_type)
323-
assert isinstance(df, expected_type)
323+
if df_type == "test_default_df_type":
324+
df = dbapi_hook.get_df("SQL")
325+
assert isinstance(df, pd.DataFrame)
326+
else:
327+
df = dbapi_hook.get_df("SQL", df_type=df_type)
328+
assert isinstance(df, expected_type)

0 commit comments

Comments
 (0)