diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/appflow.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/appflow.py index 390a22679ead1..9a60273b25946 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/appflow.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/appflow.py @@ -84,8 +84,6 @@ def __init__( **kwargs, ) -> None: super().__init__(**kwargs) - if source is not None and source not in SUPPORTED_SOURCES: - raise ValueError(f"{source} is not a supported source (options: {SUPPORTED_SOURCES})!") self.filter_date = filter_date self.flow_name = flow_name self.source = source @@ -96,6 +94,8 @@ def __init__( self.wait_for_completion = wait_for_completion def execute(self, context: Context) -> None: + self._validate_source() + self._validate_filter_date() self.filter_date_parsed: datetime | None = ( datetime.fromisoformat(self.filter_date) if self.filter_date else None ) @@ -109,6 +109,13 @@ def execute(self, context: Context) -> None: self._run_flow(context) + def _validate_source(self) -> None: + if self.source is not None and self.source not in SUPPORTED_SOURCES: + raise ValueError(f"{self.source} is not a supported source (options: {SUPPORTED_SOURCES})!") + + def _validate_filter_date(self) -> None: + pass + def _get_connector_type(self) -> str: response = self.hook.conn.describe_flow(flowName=self.flow_name) connector_type = response["sourceFlowConfig"]["connectorType"] @@ -190,8 +197,6 @@ def __init__( wait_for_completion: bool = True, **kwargs, ) -> None: - if source not in {"salesforce", "zendesk"}: - raise ValueError(NOT_SUPPORTED_SOURCE_MSG.format(source=source, entity="AppflowRunFullOperator")) super().__init__( source=source, flow_name=flow_name, @@ -203,6 +208,12 @@ def __init__( **kwargs, ) + def _validate_source(self) -> None: + if self.source not in {"salesforce", "zendesk"}: + raise ValueError( + NOT_SUPPORTED_SOURCE_MSG.format(source=self.source, entity="AppflowRunFullOperator") + ) + class AppflowRunBeforeOperator(AppflowBaseOperator): """ @@ -236,12 +247,6 @@ def __init__( wait_for_completion: bool = True, **kwargs, ) -> None: - if not filter_date: - raise ValueError(MANDATORY_FILTER_DATE_MSG.format(entity="AppflowRunBeforeOperator")) - if source != "salesforce": - raise ValueError( - NOT_SUPPORTED_SOURCE_MSG.format(source=source, entity="AppflowRunBeforeOperator") - ) super().__init__( source=source, flow_name=flow_name, @@ -253,6 +258,16 @@ def __init__( **kwargs, ) + def _validate_source(self) -> None: + if self.source != "salesforce": + raise ValueError( + NOT_SUPPORTED_SOURCE_MSG.format(source=self.source, entity="AppflowRunBeforeOperator") + ) + + def _validate_filter_date(self) -> None: + if not self.filter_date: + raise ValueError(MANDATORY_FILTER_DATE_MSG.format(entity="AppflowRunBeforeOperator")) + def _update_flow(self) -> None: if not self.filter_date_parsed: raise ValueError(f"Invalid filter_date argument parser value: {self.filter_date_parsed}") @@ -298,10 +313,6 @@ def __init__( wait_for_completion: bool = True, **kwargs, ) -> None: - if not filter_date: - raise ValueError(MANDATORY_FILTER_DATE_MSG.format(entity="AppflowRunAfterOperator")) - if source not in {"salesforce", "zendesk"}: - raise ValueError(NOT_SUPPORTED_SOURCE_MSG.format(source=source, entity="AppflowRunAfterOperator")) super().__init__( source=source, flow_name=flow_name, @@ -313,6 +324,16 @@ def __init__( **kwargs, ) + def _validate_source(self) -> None: + if self.source not in {"salesforce", "zendesk"}: + raise ValueError( + NOT_SUPPORTED_SOURCE_MSG.format(source=self.source, entity="AppflowRunAfterOperator") + ) + + def _validate_filter_date(self) -> None: + if not self.filter_date: + raise ValueError(MANDATORY_FILTER_DATE_MSG.format(entity="AppflowRunAfterOperator")) + def _update_flow(self) -> None: if not self.filter_date_parsed: raise ValueError(f"Invalid filter_date argument parser value: {self.filter_date_parsed}") @@ -358,10 +379,6 @@ def __init__( wait_for_completion: bool = True, **kwargs, ) -> None: - if not filter_date: - raise ValueError(MANDATORY_FILTER_DATE_MSG.format(entity="AppflowRunDailyOperator")) - if source != "salesforce": - raise ValueError(NOT_SUPPORTED_SOURCE_MSG.format(source=source, entity="AppflowRunDailyOperator")) super().__init__( source=source, flow_name=flow_name, @@ -373,6 +390,16 @@ def __init__( **kwargs, ) + def _validate_source(self) -> None: + if self.source != "salesforce": + raise ValueError( + NOT_SUPPORTED_SOURCE_MSG.format(source=self.source, entity="AppflowRunDailyOperator") + ) + + def _validate_filter_date(self) -> None: + if not self.filter_date: + raise ValueError(MANDATORY_FILTER_DATE_MSG.format(entity="AppflowRunDailyOperator")) + def _update_flow(self) -> None: if not self.filter_date_parsed: raise ValueError(f"Invalid filter_date argument parser value: {self.filter_date_parsed}") diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_appflow.py b/providers/amazon/tests/unit/amazon/aws/operators/test_appflow.py index 86f765808d5a0..961a70a91162a 100644 --- a/providers/amazon/tests/unit/amazon/aws/operators/test_appflow.py +++ b/providers/amazon/tests/unit/amazon/aws/operators/test_appflow.py @@ -192,6 +192,45 @@ def test_run_daily(appflow_conn, ctx, waiter_mock): ) +@pytest.mark.db_test +def test_run_daily_with_templated_validation_fields(appflow_conn, ctx, waiter_mock): + operator = AppflowRunDailyOperator( + source="{{ params.source }}", + flow_name=FLOW_NAME, + source_field="{{ params.source_field }}", + filter_date="{{ params.filter_date }}", + poll_interval=0, + task_id=TASK_ID, + ) + context = { + **ctx, + "params": { + "source": SOURCE, + "source_field": "col0", + "filter_date": "2022-05-26T00:00+00:00", + }, + } + + operator.render_template_fields(context) + operator.execute(context) + + run_assertions_base( + appflow_conn, + [ + { + "taskType": "Filter", + "connectorOperator": {"Salesforce": "BETWEEN"}, + "sourceFields": ["col0"], + "taskProperties": { + "DATA_TYPE": "datetime", + "LOWER_BOUND": "1653523199999", + "UPPER_BOUND": "1653609600000", + }, + } + ], + ) + + @pytest.mark.db_test def test_short_circuit(appflow_conn, ctx): with mock.patch("airflow.models.TaskInstance.xcom_pull") as mock_xcom_pull: diff --git a/scripts/ci/prek/validate_operators_init_exemptions.txt b/scripts/ci/prek/validate_operators_init_exemptions.txt index e192ef70bafb8..47ca0504b848e 100644 --- a/scripts/ci/prek/validate_operators_init_exemptions.txt +++ b/scripts/ci/prek/validate_operators_init_exemptions.txt @@ -6,7 +6,6 @@ # Fixing a class (moving template-field validation/transformation out of __init__ into # execute()) MUST remove its entry in the same PR — the hook fails on stale entries. # Burn-down tracked at https://github.com/apache/airflow/issues/70296 -providers/amazon/src/airflow/providers/amazon/aws/operators/appflow.py::AppflowBaseOperator providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py::EmrAddStepsOperator providers/amazon/src/airflow/providers/amazon/aws/operators/neptune.py::NeptuneStartDbClusterOperator providers/amazon/src/airflow/providers/amazon/aws/operators/neptune.py::NeptuneStopDbClusterOperator