shahar1 commented on code in PR #74817:
URL: https://github.com/apache/airflow/pull/74817#discussion_r4306616490
##########
providers/amazon/src/airflow/providers/amazon/aws/transfers/sql_to_s3.py:
##########
@@ -220,7 +221,23 @@ def _strip_suffixes(
return path[: -len(suffix)]
return path
+ def _ensure_parquet_engine(self) -> None:
+ needs_parquet = self.file_format == FILE_FORMAT.PARQUET or (
+ ("dtype_backend", "pyarrow") in self.read_kwargs.items()
+ )
+ if self.df_type != "pandas" or not needs_parquet:
+ return
+ if any(importlib.util.find_spec(engine) for engine in ("pyarrow",
"fastparquet")):
+ return
+
+ from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException
Review Comment:
**nit** — `airflow.providers.common.compat.sdk` is already imported at
module level (`AirflowException, BaseHook, BaseOperator`), so this
function-level import gains nothing. Add
`AirflowOptionalProviderFeatureException` to the top-level import.
> Imports at top of file. Valid exceptions: circular imports, lazy loading
for worker isolation, `TYPE_CHECKING` blocks. (AGENTS.md → Coding Standards)
##########
providers/amazon/tests/unit/amazon/aws/transfers/test_sql_to_s3.py:
##########
@@ -199,6 +199,41 @@ def test_fix_dtypes_not_called(self, mock_s3_hook):
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", {}, set(), False,
id="pandas-csv-no-engine"),
+ pytest.param("polars", "parquet", {}, set(), False,
id="polars-parquet-no-engine"),
+ ],
+ )
+ @mock.patch("importlib.util.find_spec")
+ def test_parquet_engine_required(
+ self, mock_find_spec, df_type, file_format, read_kwargs,
available_engines, should_raise
+ ):
+ mock_find_spec.side_effect = lambda name: object() if name in
available_engines else None
+ 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()
Review Comment:
**minor** — this tests the helper but not the wiring. Removing
`self._ensure_parquet_engine()` from `execute()` keeps every test green, so the
"checks before querying" behaviour is untested. Drive the raising case through
`op.execute(None)` and assert the hook's `get_df` was never called.
> Every changed or added behaviour must have a test; every test must fail
without the PR's change. (AGENTS.md → Testing Standards)
##########
providers/amazon/tests/unit/amazon/aws/transfers/test_sql_to_s3.py:
##########
@@ -199,6 +199,41 @@ def test_fix_dtypes_not_called(self, mock_s3_hook):
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", {}, set(), False,
id="pandas-csv-no-engine"),
+ pytest.param("polars", "parquet", {}, set(), False,
id="polars-parquet-no-engine"),
+ ],
+ )
+ @mock.patch("importlib.util.find_spec")
Review Comment:
**nit** — add `autospec=True` here; the patch is also active process-wide
during operator construction, not just the helper call.
> Use `spec`/`autospec` when mocking. (AGENTS.md → Testing Standards)
##########
providers/amazon/src/airflow/providers/amazon/aws/transfers/sql_to_s3.py:
##########
@@ -220,7 +221,23 @@ def _strip_suffixes(
return path[: -len(suffix)]
return path
+ def _ensure_parquet_engine(self) -> None:
+ needs_parquet = self.file_format == FILE_FORMAT.PARQUET or (
+ ("dtype_backend", "pyarrow") in self.read_kwargs.items()
+ )
+ if self.df_type != "pandas" or not needs_parquet:
+ return
+ if any(importlib.util.find_spec(engine) for engine in ("pyarrow",
"fastparquet")):
Review Comment:
**minor** — `dtype_backend="pyarrow"` needs pyarrow itself, so fastparquet
must not satisfy this check. With fastparquet installed and pyarrow missing,
the check passes and `read_sql` then fails with pandas' generic ``ImportError:
`Import pyarrow` failed.`` (reproduced on pandas 3.0.6). Require `pyarrow` when
`dtype_backend="pyarrow"`, and accept either engine only for Parquet output. A
parametrize case with `dtype_backend="pyarrow"` and only `fastparquet`
available would cover it.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
