This is an automated email from the ASF dual-hosted git repository.
shahar1 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 4c4b9c2d330 Raise a clear error when SqlToS3Operator has no Parquet
engine (#74817)
4c4b9c2d330 is described below
commit 4c4b9c2d330278c84e99157e72b24625d56cc3d0
Author: Dev-iL <[email protected]>
AuthorDate: Sun Oct 18 17:25:49 2026 +0300
Raise a clear error when SqlToS3Operator has no Parquet engine (#74817)
Check pandas engine requirements before querying, requiring pyarrow for the
pyarrow dtype backend and either pyarrow or fastparquet for Parquet output.
Point missing dependencies to the Amazon provider's parquet extra.
---
.../providers/amazon/aws/transfers/sql_to_s3.py | 30 ++++++-
.../unit/amazon/aws/transfers/test_sql_to_s3.py | 96 +++++++++++++++++++++-
2 files changed, 122 insertions(+), 4 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 17ef7ef7c86..5313aa26c4e 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
@@ -19,6 +19,7 @@ from __future__ import annotations
import enum
import gzip
+import importlib.util
import io
import warnings
from collections import namedtuple
@@ -27,7 +28,12 @@ from typing import TYPE_CHECKING, Any, Literal, cast
from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
-from airflow.providers.common.compat.sdk import AirflowException, BaseHook,
BaseOperator
+from airflow.providers.common.compat.sdk import (
+ AirflowException,
+ AirflowOptionalProviderFeatureException,
+ BaseHook,
+ BaseOperator,
+)
if TYPE_CHECKING:
import pandas as pd
@@ -185,8 +191,6 @@ class SqlToS3Operator(BaseOperator):
import numpy as np
import pandas as pd
except ImportError as e:
- from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException
-
raise AirflowOptionalProviderFeatureException(e)
for col in df:
@@ -221,7 +225,27 @@ class SqlToS3Operator(BaseOperator):
return path[: -len(suffix)]
return path
+ def _ensure_parquet_engine(self) -> None:
+ if self.df_type != "pandas":
+ return
+ engines: tuple[str, ...]
+ if ("dtype_backend", "pyarrow") in self.read_kwargs.items():
+ # pandas builds the frame through pyarrow itself here, so
fastparquet cannot stand in
+ # for it even when the output format is Parquet.
+ trigger, engines = 'read_kwargs={"dtype_backend": "pyarrow"}',
("pyarrow",)
+ elif self.file_format == FILE_FORMAT.PARQUET:
+ trigger, engines = 'file_format="parquet"', ("pyarrow",
"fastparquet")
+ else:
+ return
+ if any(importlib.util.find_spec(engine) for engine in engines):
+ return
+
+ raise AirflowOptionalProviderFeatureException(
+ f"{trigger} needs {' or '.join(engines)}. Install
apache-airflow-providers-amazon[parquet]."
+ )
+
def execute(self, context: Context) -> None:
+ self._ensure_parquet_engine()
sql_hook = self._get_hook()
s3_conn = S3Hook(aws_conn_id=self.aws_conn_id, verify=self.verify)
data_df = sql_hook.get_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 c1100ea1b79..b97ee952b42 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
@@ -18,6 +18,7 @@
from __future__ import annotations
import gzip
+import importlib.util
import io
from unittest import mock
@@ -29,10 +30,30 @@ from packaging.version import Version
from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.models import Connection
from airflow.providers.amazon.aws.transfers.sql_to_s3 import SqlToS3Operator
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException,
AirflowOptionalProviderFeatureException
PANDAS_3_PLUS = Version(pd.__version__).major >= 3
+PARQUET_ENGINES = ("pyarrow", "fastparquet")
+REAL_FIND_SPEC = importlib.util.find_spec
+
+
+def build_find_spec_stub(available_engines):
+ """
+ Fake only the Parquet engines and resolve every other name for real.
+
+ The patch is installed process-wide and stays active while the operator is
constructed and
+ executed, so a stub answering None for every name would also answer
optional-dependency
+ probes these tests say nothing about, and a test could then pass for the
wrong reason.
+ """
+
+ def find_spec(name, package=None):
+ if name in PARQUET_ENGINES:
+ return object() if name in available_engines else None
+ return REAL_FIND_SPEC(name, package)
+
+ return find_spec
+
class TestSqlToS3Operator:
@pytest.mark.parametrize(
@@ -248,6 +269,79 @@ class TestSqlToS3Operator:
mock_fix_dtypes.assert_not_called()
+ @pytest.mark.parametrize(
+ ("df_type", "file_format", "read_kwargs", "available_engines",
"should_raise"),
+ [
+ pytest.param("pandas", "parquet", {}, set(), True,
id="pandas-parquet-no-engine"),
+ pytest.param("pandas", "parquet", {}, {"fastparquet"}, False,
id="pandas-parquet-fastparquet"),
+ pytest.param(
+ "pandas", "csv", {"dtype_backend": "pyarrow"}, set(), True,
id="pandas-pyarrow-backend"
+ ),
+ pytest.param(
+ "pandas",
+ "csv",
+ {"dtype_backend": "pyarrow"},
+ {"fastparquet"},
+ True,
+ id="pandas-pyarrow-backend-fastparquet-only",
+ ),
+ pytest.param(
+ "pandas",
+ "parquet",
+ {"dtype_backend": "pyarrow"},
+ {"fastparquet"},
+ True,
+ id="pandas-parquet-pyarrow-backend-fastparquet-only",
+ ),
+ pytest.param("pandas", "csv", {}, set(), False,
id="pandas-csv-no-engine"),
+ pytest.param("polars", "parquet", {}, set(), False,
id="polars-parquet-no-engine"),
+ ],
+ )
+ @mock.patch("importlib.util.find_spec", autospec=True)
+ def test_parquet_engine_required(
+ self, mock_find_spec, df_type, file_format, read_kwargs,
available_engines, should_raise
+ ):
+ mock_find_spec.side_effect = build_find_spec_stub(available_engines)
+ op = SqlToS3Operator(
+ query="query",
+ s3_bucket="bucket",
+ s3_key="key",
+ sql_conn_id="mysql_conn_id",
+ task_id="task_id",
+ df_type=df_type,
+ file_format=file_format,
+ read_kwargs=read_kwargs,
+ dag=None,
+ )
+
+ if should_raise:
+ with pytest.raises(AirflowOptionalProviderFeatureException,
match=r"\[parquet\]"):
+ op._ensure_parquet_engine()
+ else:
+ op._ensure_parquet_engine()
+
+ @mock.patch("airflow.providers.amazon.aws.transfers.sql_to_s3.S3Hook")
+ @mock.patch("importlib.util.find_spec", autospec=True)
+ def test_execute_checks_parquet_engine_before_querying(self,
mock_find_spec, mock_s3_hook):
+ mock_find_spec.side_effect = build_find_spec_stub(set())
+ mock_dbapi_hook = mock.Mock()
+ mock_dbapi_hook.return_value.get_df.return_value = pd.DataFrame({"a":
[1]})
+ op = SqlToS3Operator(
+ query="query",
+ s3_bucket="bucket",
+ s3_key="key",
+ sql_conn_id="mysql_conn_id",
+ task_id="task_id",
+ file_format="parquet",
+ dag=None,
+ )
+ op._get_hook = mock_dbapi_hook
+
+ with pytest.raises(AirflowOptionalProviderFeatureException,
match=r"\[parquet\]"):
+ op.execute(None)
+
+ mock_dbapi_hook.return_value.get_df.assert_not_called()
+
def test_invalid_file_format(self):
with pytest.raises(AirflowException):
SqlToS3Operator(
