You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by ep...@apache.org on 2022/06/29 15:19:49 UTC
[airflow] 04/45: Fix closing connection dbapi.get_pandas_df (#23452)
This is an automated email from the ASF dual-hosted git repository.
ephraimanierobi pushed a commit to branch v2-3-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit bab5e1cc6fbc1d9333c3c4aacb18d4083c92d310
Author: Hubert PietroĊ <94...@users.noreply.github.com>
AuthorDate: Tue May 31 12:39:16 2022 +0200
Fix closing connection dbapi.get_pandas_df (#23452)
(cherry picked from commit ab1f637e463011a34d950c306583400b7a2fceb3)
---
airflow/hooks/dbapi.py | 18 ++++++++++++++++++
1 file changed, 18 insertions(+)
diff --git a/airflow/hooks/dbapi.py b/airflow/hooks/dbapi.py
index da33bacca8..0b9ce4377b 100644
--- a/airflow/hooks/dbapi.py
+++ b/airflow/hooks/dbapi.py
@@ -128,6 +128,24 @@ class DbApiHook(BaseHook):
with closing(self.get_conn()) as conn:
return psql.read_sql(sql, con=conn, params=parameters, **kwargs)
+ def get_pandas_df_by_chunks(self, sql, parameters=None, *, chunksize, **kwargs):
+ """
+ Executes the sql and returns a generator
+
+ :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 chunksize: number of rows to include in each chunk
+ :param kwargs: (optional) passed into pandas.io.sql.read_sql method
+ """
+ try:
+ from pandas.io import sql as psql
+ except ImportError:
+ raise Exception("pandas library not installed, run: pip install 'apache-airflow[pandas]'.")
+
+ with closing(self.get_conn()) as conn:
+ yield from psql.read_sql(sql, con=conn, params=parameters, chunksize=chunksize, **kwargs)
+
def get_records(self, sql, parameters=None):
"""
Executes the sql and returns a set of records.