diff --git a/providers/amazon/src/airflow/providers/amazon/aws/transfers/s3_to_redshift.py b/providers/amazon/src/airflow/providers/amazon/aws/transfers/s3_to_redshift.py index a4c371a4df460..18636cf92e840 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/transfers/s3_to_redshift.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/transfers/s3_to_redshift.py @@ -117,21 +117,6 @@ def __init__( self.method = method self.upsert_keys = upsert_keys self.redshift_data_api_kwargs = redshift_data_api_kwargs or {} - # In execute() we attempt to fetch this aws connection to check for extras. If the user didn't - # actually provide a connection note that, because we don't want to let the exception bubble up in - # that case (since we're silently injecting a connection on their behalf). - self._aws_conn_id: str | None - if is_arg_set(aws_conn_id): - self.conn_set = True - self._aws_conn_id = aws_conn_id - else: - self.conn_set = False - self._aws_conn_id = "aws_default" - - if self.redshift_data_api_kwargs: - for arg in ["sql", "parameters"]: - if arg in self.redshift_data_api_kwargs: - raise AirflowException(f"Cannot include param '{arg}' in Redshift Data API kwargs") @property def use_redshift_data(self): @@ -154,15 +139,31 @@ def execute(self, context: Context) -> None: if self.method not in AVAILABLE_METHODS: raise AirflowException(f"Method not found! Available methods: {AVAILABLE_METHODS}") + if self.redshift_data_api_kwargs: + for arg in ["sql", "parameters"]: + if arg in self.redshift_data_api_kwargs: + raise AirflowException(f"Cannot include param '{arg}' in Redshift Data API kwargs") + + # We attempt to fetch this aws connection below to check for extras. If the user didn't + # actually provide a connection note that, because we don't want to let the exception bubble up + # in that case (since we're silently injecting a connection on their behalf). + aws_conn_id: str | None + if is_arg_set(self.aws_conn_id): + conn_set = True + aws_conn_id = self.aws_conn_id + else: + conn_set = False + aws_conn_id = "aws_default" + if self.use_redshift_data: redshift_data_hook = RedshiftDataHook(aws_conn_id=self.redshift_conn_id) else: redshift_sql_hook = RedshiftSQLHook(redshift_conn_id=self.redshift_conn_id) conn = ( - S3Hook.get_connection(conn_id=self._aws_conn_id) + S3Hook.get_connection(conn_id=aws_conn_id) # Only fetch the connection if it was set by the user and it is not None - if self.conn_set and self._aws_conn_id + if conn_set and aws_conn_id else None ) region_info = "" @@ -171,7 +172,7 @@ def execute(self, context: Context) -> None: if conn and conn.extra_dejson.get("role_arn", False): credentials_block = f"aws_iam_role={conn.extra_dejson['role_arn']}" else: - s3_hook = S3Hook(aws_conn_id=self._aws_conn_id, verify=self.verify) + s3_hook = S3Hook(aws_conn_id=aws_conn_id, verify=self.verify) credentials = s3_hook.get_credentials() credentials_block = build_credentials_block(credentials) diff --git a/providers/amazon/tests/unit/amazon/aws/transfers/test_s3_to_redshift.py b/providers/amazon/tests/unit/amazon/aws/transfers/test_s3_to_redshift.py index 8a73e88320ac6..cc88aa3dec37e 100644 --- a/providers/amazon/tests/unit/amazon/aws/transfers/test_s3_to_redshift.py +++ b/providers/amazon/tests/unit/amazon/aws/transfers/test_s3_to_redshift.py @@ -407,18 +407,19 @@ def test_execute_unavailable_method(self): @pytest.mark.parametrize("param", ["sql", "parameters"]) def test_invalid_param_in_redshift_data_api_kwargs(self, param): """ - Test passing invalid param in RS Data API kwargs raises an error + Test passing invalid param in RS Data API kwargs raises an error at execute time """ - with pytest.raises(AirflowException): - S3ToRedshiftOperator( - schema="schema", - table="table", - s3_bucket="bucket", - s3_key="key", - task_id="task_id", - dag=None, - redshift_data_api_kwargs={param: "param"}, - ) + op = S3ToRedshiftOperator( + schema="schema", + table="table", + s3_bucket="bucket", + s3_key="key", + task_id="task_id", + dag=None, + redshift_data_api_kwargs={param: "param"}, + ) + with pytest.raises(AirflowException, match=f"Cannot include param '{param}'"): + op.execute({}) @mock.patch("airflow.providers.amazon.aws.hooks.s3.S3Hook.get_connection") @mock.patch("airflow.models.connection.Connection") diff --git a/scripts/ci/prek/validate_operators_init_exemptions.txt b/scripts/ci/prek/validate_operators_init_exemptions.txt index e192ef70bafb8..6b22398dc3678 100644 --- a/scripts/ci/prek/validate_operators_init_exemptions.txt +++ b/scripts/ci/prek/validate_operators_init_exemptions.txt @@ -16,7 +16,6 @@ providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py::SageMa providers/amazon/src/airflow/providers/amazon/aws/operators/step_function.py::StepFunctionStartExecutionOperator providers/amazon/src/airflow/providers/amazon/aws/transfers/base.py::AwsToAwsBaseOperator providers/amazon/src/airflow/providers/amazon/aws/transfers/gcs_to_s3.py::GCSToS3Operator -providers/amazon/src/airflow/providers/amazon/aws/transfers/s3_to_redshift.py::S3ToRedshiftOperator providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py::KubernetesPodOperator providers/google/src/airflow/providers/google/cloud/operators/bigquery.py::BigQueryInsertJobOperator providers/google/src/airflow/providers/google/cloud/operators/cloud_batch.py::CloudBatchSubmitJobOperator