guan404ming commented on code in PR #48875:
URL: https://github.com/apache/airflow/pull/48875#discussion_r2040610356
##########
providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py:
##########
@@ -375,6 +379,56 @@ def get_reserved_words(self, dialect_name: str) ->
set[str]:
self.log.debug("reserved words for '%s': %s", dialect_name, result)
return result
+ def get_df(
+ self,
+ sql,
+ parameters: list | tuple | Mapping[str, Any] | None = None,
+ *,
+ df_type: Literal["pandas", "polars"] = "pandas",
+ **kwargs,
+ ) -> DataFrame | PolarsDataFrame:
+ """
+ Execute the sql and returns a dataframe.
+
+ :param sql: the sql statement to be executed (str) or a list of sql
statements to execute
+ :param parameters: The parameters to render the SQL query with.
+ :param df_type: Type of dataframe to return, either "pandas" or
"polars"
+ :param kwargs: (optional) passed into `pandas.io.sql.read_sql` or
`polars.read_database` method
+ """
+ if df_type == "pandas":
+ return self._get_pandas_df(sql, parameters, **kwargs)
+ elif df_type == "polars":
+ return self._get_polars_df(sql, parameters, **kwargs)
+
+ def _get_pandas_df(
+ self,
+ sql,
+ parameters: list | tuple | Mapping[str, Any] | None = None,
+ **kwargs,
+ ) -> DataFrame:
+ """
+ Execute the sql and returns a pandas dataframe.
+
+ :param sql: the sql statement to be executed (str) or a list of sql
statements to execute
+ :param parameters: The parameters to render the SQL query with.
+ :param kwargs: (optional) passed into pandas.io.sql.read_sql method
+ """
+ try:
+ from pandas.io import sql as psql
+ except ImportError:
+ raise AirflowOptionalProviderFeatureException(
+ "pandas library not installed, run: pip install "
+ "'apache-airflow-providers-common-sql[pandas]'."
+ )
+
+ with closing(self.get_conn()) as conn:
+ return psql.read_sql(sql, con=conn, params=parameters, **kwargs)
+
+ @deprecated(
+ reason="Replaced by function `get_df`.",
+ category=AirflowProviderDeprecationWarning,
+ action="ignore",
+ )
Review Comment:
I made this use the `_get_pandas_df` instead of `get_df` since the type
issue!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscr...@airflow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
