This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new bdd4708aa08 Enhancement: AWS Provider sql to s3 operator pd.read_sql
kwargs (#53399)
bdd4708aa08 is described below
commit bdd4708aa08d47d3bc6d997b936a82f6dfe7b7cd
Author: Bryan Corder <[email protected]>
AuthorDate: Fri Jul 18 02:18:51 2025 -0400
Enhancement: AWS Provider sql to s3 operator pd.read_sql kwargs (#53399)
* initial commit for changes
* added unit test
---
.../providers/amazon/aws/transfers/sql_to_s3.py | 11 ++++-
.../unit/amazon/aws/transfers/test_sql_to_s3.py | 49 ++++++++++++++++++++--
2 files changed, 54 insertions(+), 6 deletions(-)
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/transfers/sql_to_s3.py
b/providers/amazon/src/airflow/providers/amazon/aws/transfers/sql_to_s3.py
index e31652af4d4..344935215b7 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/transfers/sql_to_s3.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/transfers/sql_to_s3.py
@@ -69,6 +69,7 @@ class SqlToS3Operator(BaseOperator):
:param sql_hook_params: Extra config params to be passed to the underlying
hook.
Should match the desired hook constructor params.
:param parameters: (optional) the parameters to render the SQL query with.
+ :param read_pd_kwargs: arguments to include in DataFrame when
``pd.read_sql()`` is called.
:param aws_conn_id: reference to a specific S3 connection
:param verify: Whether or not to verify SSL certificates for S3 connection.
By default SSL certificates are verified.
@@ -97,6 +98,7 @@ class SqlToS3Operator(BaseOperator):
template_fields_renderers = {
"query": "sql",
"pd_kwargs": "json",
+ "read_pd_kwargs": "json",
}
def __init__(
@@ -108,6 +110,7 @@ class SqlToS3Operator(BaseOperator):
sql_conn_id: str,
sql_hook_params: dict | None = None,
parameters: None | Mapping[str, Any] | list | tuple = None,
+ read_pd_kwargs: dict | None = None,
replace: bool = False,
aws_conn_id: str | None = "aws_default",
verify: bool | str | None = None,
@@ -127,6 +130,7 @@ class SqlToS3Operator(BaseOperator):
self.replace = replace
self.pd_kwargs = pd_kwargs or {}
self.parameters = parameters
+ self.read_pd_kwargs = read_pd_kwargs or {}
self.max_rows_per_file = max_rows_per_file
self.groupby_kwargs = groupby_kwargs or {}
self.sql_hook_params = sql_hook_params
@@ -185,9 +189,12 @@ class SqlToS3Operator(BaseOperator):
def execute(self, context: Context) -> None:
sql_hook = self._get_hook()
s3_conn = S3Hook(aws_conn_id=self.aws_conn_id, verify=self.verify)
- data_df = sql_hook.get_df(sql=self.query, parameters=self.parameters,
df_type="pandas")
+ data_df = sql_hook.get_df(
+ sql=self.query, parameters=self.parameters, df_type="pandas",
**self.read_pd_kwargs
+ )
self.log.info("Data from SQL obtained")
- self._fix_dtypes(data_df, self.file_format)
+ if ("dtype_backend", "pyarrow") not in self.read_pd_kwargs.items():
+ self._fix_dtypes(data_df, self.file_format)
file_options = FILE_OPTIONS_MAP[self.file_format]
for group_name, df in self._partition_dataframe(df=data_df):
diff --git a/providers/amazon/tests/unit/amazon/aws/transfers/test_sql_to_s3.py
b/providers/amazon/tests/unit/amazon/aws/transfers/test_sql_to_s3.py
index 61745829157..0d3791cc8d5 100644
--- a/providers/amazon/tests/unit/amazon/aws/transfers/test_sql_to_s3.py
+++ b/providers/amazon/tests/unit/amazon/aws/transfers/test_sql_to_s3.py
@@ -30,8 +30,9 @@ from airflow.providers.amazon.aws.transfers.sql_to_s3 import
SqlToS3Operator
class TestSqlToS3Operator:
+ @pytest.mark.parametrize("dtype_backend", ["numpy_nullable", "pyarrow"])
@mock.patch("airflow.providers.amazon.aws.transfers.sql_to_s3.S3Hook")
- def test_execute_csv(self, mock_s3_hook):
+ def test_execute_csv(self, mock_s3_hook, dtype_backend):
query = "query"
s3_bucket = "bucket"
s3_key = "key"
@@ -49,6 +50,7 @@ class TestSqlToS3Operator:
aws_conn_id="aws_conn_id",
task_id="task_id",
replace=True,
+ read_pd_kwargs={"dtype_backend": dtype_backend},
pd_kwargs={"index": False, "header": False},
dag=None,
)
@@ -56,15 +58,18 @@ class TestSqlToS3Operator:
op.execute(None)
mock_s3_hook.assert_called_once_with(aws_conn_id="aws_conn_id",
verify=None)
- get_df_mock.assert_called_once_with(sql=query, parameters=None,
df_type="pandas")
+ get_df_mock.assert_called_once_with(
+ sql=query, parameters=None, df_type="pandas",
dtype_backend=dtype_backend
+ )
file_obj =
mock_s3_hook.return_value.load_file_obj.call_args[1]["file_obj"]
assert isinstance(file_obj, io.BytesIO)
mock_s3_hook.return_value.load_file_obj.assert_called_once_with(
file_obj=file_obj, key=s3_key, bucket_name=s3_bucket, replace=True
)
+ @pytest.mark.parametrize("dtype_backend", ["numpy_nullable", "pyarrow"])
@mock.patch("airflow.providers.amazon.aws.transfers.sql_to_s3.S3Hook")
- def test_execute_parquet(self, mock_s3_hook):
+ def test_execute_parquet(self, mock_s3_hook, dtype_backend):
query = "query"
s3_bucket = "bucket"
s3_key = "key"
@@ -82,6 +87,7 @@ class TestSqlToS3Operator:
sql_conn_id="mysql_conn_id",
aws_conn_id="aws_conn_id",
task_id="task_id",
+ read_pd_kwargs={"dtype_backend": dtype_backend},
file_format="parquet",
replace=True,
dag=None,
@@ -90,7 +96,10 @@ class TestSqlToS3Operator:
op.execute(None)
mock_s3_hook.assert_called_once_with(aws_conn_id="aws_conn_id",
verify=None)
- get_df_mock.assert_called_once_with(sql=query, parameters=None,
df_type="pandas")
+ get_df_mock.assert_called_once_with(
+ sql=query, parameters=None, df_type="pandas",
dtype_backend=dtype_backend
+ )
+
file_obj =
mock_s3_hook.return_value.load_file_obj.call_args[1]["file_obj"]
assert isinstance(file_obj, io.BytesIO)
mock_s3_hook.return_value.load_file_obj.assert_called_once_with(
@@ -192,6 +201,38 @@ class TestSqlToS3Operator:
assert dirty_df["strings"].values[2] == params["null_string_result"]
assert dirty_df["ints"].dtype.kind == "i"
+ @mock.patch("airflow.providers.amazon.aws.transfers.sql_to_s3.S3Hook")
+ def test_fix_dtypes_not_called(self, mock_s3_hook):
+ query = "query"
+ s3_bucket = "bucket"
+ s3_key = "key"
+
+ mock_dbapi_hook = mock.Mock()
+ test_df = pd.DataFrame({"a": "1", "b": "2"}, index=[0, 1])
+
+ get_df_mock = mock_dbapi_hook.return_value.get_df
+ get_df_mock.return_value = test_df
+
+ op = SqlToS3Operator(
+ query=query,
+ s3_bucket=s3_bucket,
+ s3_key=s3_key,
+ sql_conn_id="mysql_conn_id",
+ aws_conn_id="aws_conn_id",
+ task_id="task_id",
+ read_pd_kwargs={"dtype_backend": "pyarrow"},
+ file_format="parquet",
+ replace=True,
+ dag=None,
+ )
+
+ op._get_hook = mock_dbapi_hook
+
+ with mock.patch.object(SqlToS3Operator, "_fix_dtypes") as
mock_fix_dtypes:
+ op.execute(None)
+
+ mock_fix_dtypes.assert_not_called()
+
def test_invalid_file_format(self):
with pytest.raises(AirflowException):
SqlToS3Operator(
