apache-airflow-providers-amazon 9.6.0__py3-none-any.whl → 9.6.1rc1__py3-none-any.whl

airflow/providers/amazon/__init__.py

@@ -29,7 +29,7 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "9.6.0"
+__version__ = "9.6.1"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
     "2.9.0"

airflow/providers/amazon/aws/auth_manager/cli/avp_commands.py

@@ -107,26 +107,23 @@ def _create_policy_store(client: BaseClient, args) -> tuple[str | None, bool]:
             f"There is already a policy store with description '{args.policy_store_description}' in Amazon Verified Permissions: '{existing_policy_stores[0]['policyStoreId']}'."
         )
         return existing_policy_stores[0]["policyStoreId"], False
-    else:
-        print(f"No policy store with description '{args.policy_store_description}' found, creating one.")
-        if args.dry_run:
-            print(
-                f"Dry run, not creating the policy store with description '{args.policy_store_description}'."
-            )
-            return None, True
-
-        response = client.create_policy_store(
-            validationSettings={
-                "mode": "STRICT",
-            },
-            description=args.policy_store_description,
-        )
-        if args.verbose:
-            log.debug("Response from create_policy_store: %s", response)
+    print(f"No policy store with description '{args.policy_store_description}' found, creating one.")
+    if args.dry_run:
+        print(f"Dry run, not creating the policy store with description '{args.policy_store_description}'.")
+        return None, True
+
+    response = client.create_policy_store(
+        validationSettings={
+            "mode": "STRICT",
+        },
+        description=args.policy_store_description,
+    )
+    if args.verbose:
+        log.debug("Response from create_policy_store: %s", response)
 
-        print(f"Policy store created: '{response['policyStoreId']}'")
+    print(f"Policy store created: '{response['policyStoreId']}'")
 
-        return response["policyStoreId"], True
+    return response["policyStoreId"], True
 
 
 def _set_schema(client: BaseClient, policy_store_id: str, args) -> None:

airflow/providers/amazon/aws/auth_manager/router/login.py

@@ -91,7 +91,7 @@ def login_callback(request: Request):
 
 def _init_saml_auth(request: Request) -> OneLogin_Saml2_Auth:
     request_data = _prepare_request(request)
-    base_url = conf.get(section="api", key="base_url")
+    base_url = conf.get(section="api", key="base_url", fallback="/")
     settings = {
         # We want to keep this flag on in case of errors.
         # It provides an error reasons, if turned off, it does not

airflow/providers/amazon/aws/executors/ecs/ecs_executor.py

@@ -23,6 +23,7 @@ Each Airflow task gets delegated out to an Amazon ECS Task.
 
 from __future__ import annotations
 
+import contextlib
 import time
 from collections import defaultdict, deque
 from collections.abc import Sequence
@@ -449,13 +450,11 @@ class AwsEcsExecutor(BaseExecutor):
             else:
                 task = run_task_response["tasks"][0]
                 self.active_workers.add_task(task, task_key, queue, cmd, exec_config, attempt_number)
-                try:
-                    self.running_state(task_key, task.task_arn)
-                except AttributeError:
+                with contextlib.suppress(AttributeError):
                     # running_state is newly added, and only needed to support task adoption (an optional
                     # executor feature).
                     # TODO: remove when min airflow version >= 2.9.2
-                    pass
+                    self.running_state(task_key, task.task_arn)
 
     def _run_task(
         self, task_id: TaskInstanceKey, cmd: CommandType, queue: str, exec_config: ExecutorConfigType

airflow/providers/amazon/aws/executors/ecs/ecs_executor_config.py

@@ -72,7 +72,7 @@ def build_task_kwargs() -> dict:
         raise ValueError(
             "capacity_provider_strategy and launch_type are mutually exclusive, you can not provide both."
         )
-    elif "cluster" in task_kwargs and not (has_capacity_provider or has_launch_type):
+    if "cluster" in task_kwargs and not (has_capacity_provider or has_launch_type):
         # Default API behavior if neither is provided is to fall back on the default capacity
         # provider if it exists. Since it is not a required value, check if there is one
         # before using it, and if there is not then use the FARGATE launch_type as

airflow/providers/amazon/aws/executors/ecs/utils.py

@@ -138,7 +138,7 @@ class EcsExecutorTask:
         """
         if self.last_status == "RUNNING":
             return State.RUNNING
-        elif self.desired_status == "RUNNING":
+        if self.desired_status == "RUNNING":
             return State.QUEUED
         is_finished = self.desired_status == "STOPPED"
         has_exit_codes = all(["exit_code" in x for x in self.containers])

airflow/providers/amazon/aws/hooks/athena.py

@@ -208,7 +208,7 @@ class AthenaHook(AwsBaseHook):
         if query_state is None:
             self.log.error("Invalid Query state. Query execution id: %s", query_execution_id)
             return None
-        elif query_state in self.INTERMEDIATE_STATES or query_state in self.FAILURE_STATES:
+        if query_state in self.INTERMEDIATE_STATES or query_state in self.FAILURE_STATES:
             self.log.error(
                 'Query is in "%s" state. Cannot fetch results. Query execution id: %s',
                 query_state,

airflow/providers/amazon/aws/hooks/base_aws.py

@@ -191,15 +191,13 @@ class BaseSessionFactory(LoggingMixin):
                 session = self.get_async_session()
                 self._apply_session_kwargs(session)
                 return session
-            else:
-                return boto3.session.Session(region_name=self.region_name)
-        elif not self.role_arn:
+            return boto3.session.Session(region_name=self.region_name)
+        if not self.role_arn:
             if deferrable:
                 session = self.get_async_session()
                 self._apply_session_kwargs(session)
                 return session
-            else:
-                return self.basic_session
+            return self.basic_session
 
         # Values stored in ``AwsConnectionWrapper.session_kwargs`` are intended to be used only
         # to create the initial boto3 session.
@@ -624,7 +622,7 @@ class AwsGenericHook(BaseHook, Generic[BaseAwsConnection]):
                 if is_resource_type:
                     raise LookupError("Requested `resource_type`, but `client_type` was set instead.")
                 return self.client_type
-            elif self.resource_type:
+            if self.resource_type:
                 if not is_resource_type:
                     raise LookupError("Requested `client_type`, but `resource_type` was set instead.")
                 return self.resource_type
@@ -840,15 +838,14 @@ class AwsGenericHook(BaseHook, Generic[BaseAwsConnection]):
         """
         if "/" in role:
             return role
-        else:
-            session = self.get_session(region_name=region_name)
-            _client = session.client(
-                service_name="iam",
-                endpoint_url=self.conn_config.get_service_endpoint_url("iam"),
-                config=self.config,
-                verify=self.verify,
-            )
-            return _client.get_role(RoleName=role)["Role"]["Arn"]
+        session = self.get_session(region_name=region_name)
+        _client = session.client(
+            service_name="iam",
+            endpoint_url=self.conn_config.get_service_endpoint_url("iam"),
+            config=self.config,
+            verify=self.verify,
+        )
+        return _client.get_role(RoleName=role)["Role"]["Arn"]
 
     @staticmethod
     def retry(should_retry: Callable[[Exception], bool]):

airflow/providers/amazon/aws/hooks/batch_client.py

@@ -142,6 +142,17 @@ class BatchProtocol(Protocol):
         """
         ...
 
+    def create_compute_environment(self, **kwargs) -> dict:
+        """
+        Create an AWS Batch compute environment.
+
+        :param kwargs: Arguments for boto3 create_compute_environment
+
+        .. seealso::
+            - https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/batch/client/create_compute_environment.html
+        """
+        ...
+
 
 # Note that the use of invalid-name parameters should be restricted to the boto3 mappings only;
 # all the Airflow wrappers of boto3 clients should not adopt invalid-names to match boto3.

airflow/providers/amazon/aws/hooks/cloud_formation.py

@@ -62,8 +62,7 @@ class CloudFormationHook(AwsBaseHook):
         except ClientError as e:
             if "does not exist" in str(e):
                 return None
-            else:
-                raise e
+            raise e
 
     def create_stack(self, stack_name: str, cloudformation_parameters: dict) -> None:
         """

airflow/providers/amazon/aws/hooks/datasync.py

@@ -312,9 +312,9 @@ class DataSyncHook(AwsBaseHook):
             self.log.info("status=%s", status)
             if status in self.TASK_EXECUTION_SUCCESS_STATES:
                 return True
-            elif status in self.TASK_EXECUTION_FAILURE_STATES:
+            if status in self.TASK_EXECUTION_FAILURE_STATES:
                 return False
-            elif status is None or status in self.TASK_EXECUTION_INTERMEDIATE_STATES:
+            if status is None or status in self.TASK_EXECUTION_INTERMEDIATE_STATES:
                 time.sleep(self.wait_interval_seconds)
             else:
                 raise AirflowException(f"Unknown status: {status}")  # Should never happen

airflow/providers/amazon/aws/hooks/dms.py

@@ -108,9 +108,8 @@ class DmsHook(AwsBaseHook):
             status = replication_tasks[0]["Status"]
             self.log.info('Replication task with ARN(%s) has status "%s".', replication_task_arn, status)
             return status
-        else:
-            self.log.info("Replication task with ARN(%s) is not found.", replication_task_arn)
-            return None
+        self.log.info("Replication task with ARN(%s) is not found.", replication_task_arn)
+        return None
 
     def create_replication_task(
         self,

airflow/providers/amazon/aws/hooks/dynamodb.py

@@ -102,5 +102,4 @@ class DynamoDBHook(AwsBaseHook):
             error_code = e.response.get("Error", {}).get("Code")
             if error_code == "ImportNotFoundException":
                 raise AirflowException("S3 import into Dynamodb job not found.")
-            else:
-                raise e
+            raise e

airflow/providers/amazon/aws/hooks/emr.py

@@ -83,11 +83,10 @@ class EmrHook(AwsBaseHook):
             cluster_id = matching_clusters[0]["Id"]
             self.log.info("Found cluster name = %s id = %s", emr_cluster_name, cluster_id)
             return cluster_id
-        elif len(matching_clusters) > 1:
+        if len(matching_clusters) > 1:
             raise AirflowException(f"More than one cluster found for name {emr_cluster_name}")
-        else:
-            self.log.info("No cluster found for name %s", emr_cluster_name)
-            return None
+        self.log.info("No cluster found for name %s", emr_cluster_name)
+        return None
 
     def create_job_flow(self, job_flow_overrides: dict[str, Any]) -> dict[str, Any]:
         """
@@ -387,12 +386,11 @@ class EmrContainerHook(AwsBaseHook):
 
         if response["ResponseMetadata"]["HTTPStatusCode"] != 200:
             raise AirflowException(f"Create EMR EKS Cluster failed: {response}")
-        else:
-            self.log.info(
-                "Create EMR EKS Cluster success - virtual cluster id %s",
-                response["id"],
-            )
-            return response["id"]
+        self.log.info(
+            "Create EMR EKS Cluster success - virtual cluster id %s",
+            response["id"],
+        )
+        return response["id"]
 
     def submit_job(
         self,
@@ -446,13 +444,12 @@ class EmrContainerHook(AwsBaseHook):
 
         if response["ResponseMetadata"]["HTTPStatusCode"] != 200:
             raise AirflowException(f"Start Job Run failed: {response}")
-        else:
-            self.log.info(
-                "Start Job Run success - Job Id %s and virtual cluster id %s",
-                response["id"],
-                response["virtualClusterId"],
-            )
-            return response["id"]
+        self.log.info(
+            "Start Job Run success - Job Id %s and virtual cluster id %s",
+            response["id"],
+            response["virtualClusterId"],
+        )
+        return response["id"]
 
     def get_job_failure_reason(self, job_id: str) -> str | None:
         """

airflow/providers/amazon/aws/hooks/glue.py

@@ -320,8 +320,7 @@ class GlueJobHook(AwsBaseHook):
             if ret:
                 time.sleep(sleep_before_return)
                 return ret
-            else:
-                time.sleep(self.job_poll_interval)
+            time.sleep(self.job_poll_interval)
 
     async def async_job_completion(self, job_name: str, run_id: str, verbose: bool = False) -> dict[str, str]:
         """
@@ -338,8 +337,7 @@ class GlueJobHook(AwsBaseHook):
             ret = self._handle_state(job_run_state, job_name, run_id, verbose, next_log_tokens)
             if ret:
                 return ret
-            else:
-                await asyncio.sleep(self.job_poll_interval)
+            await asyncio.sleep(self.job_poll_interval)
 
     def _handle_state(
         self,
@@ -367,13 +365,12 @@ class GlueJobHook(AwsBaseHook):
             job_error_message = f"Exiting Job {run_id} Run State: {state}"
             self.log.info(job_error_message)
             raise AirflowException(job_error_message)
-        else:
-            self.log.info(
-                "Polling for AWS Glue Job %s current run state with status %s",
-                job_name,
-                state,
-            )
-            return None
+        self.log.info(
+            "Polling for AWS Glue Job %s current run state with status %s",
+            job_name,
+            state,
+        )
+        return None
 
     def has_job(self, job_name) -> bool:
         """
@@ -414,8 +411,7 @@ class GlueJobHook(AwsBaseHook):
             self.conn.update_job(JobName=job_name, JobUpdate=job_kwargs)
             self.log.info("Updated configurations: %s", update_config)
             return True
-        else:
-            return False
+        return False
 
     def get_or_create_glue_job(self) -> str | None:
         """

airflow/providers/amazon/aws/hooks/mwaa.py

@@ -101,13 +101,12 @@ class MwaaHook(AwsBaseHook):
                     "Access Denied due to missing airflow:InvokeRestApi in IAM policy. Trying again by generating local token..."
                 )
                 return self._invoke_rest_api_using_local_session_token(**api_kwargs)
-            else:
-                to_log = e.response
-                # ResponseMetadata is removed because it contains data that is either very unlikely to be
-                # useful in XComs and logs, or redundant given the data already included in the response
-                to_log.pop("ResponseMetadata", None)
-                self.log.error(to_log)
-                raise
+            to_log = e.response
+            # ResponseMetadata is removed because it contains data that is either very unlikely to be
+            # useful in XComs and logs, or redundant given the data already included in the response
+            to_log.pop("ResponseMetadata", None)
+            self.log.error(to_log)
+            raise
 
     def _invoke_rest_api_using_local_session_token(
         self,

airflow/providers/amazon/aws/hooks/redshift_data.py

@@ -181,7 +181,7 @@ class RedshiftDataHook(AwsGenericHook["RedshiftDataAPIServiceClient"]):
             if num_rows is not None:
                 self.log.info("Processed %s rows", num_rows)
             return True
-        elif status in FAILURE_STATES:
+        if status in FAILURE_STATES:
             exception_cls = (
                 RedshiftDataQueryFailedError if status == FAILED_STATE else RedshiftDataQueryAbortedError
             )

airflow/providers/amazon/aws/hooks/redshift_sql.py

@@ -245,15 +245,14 @@ class RedshiftSQLHook(DbApiHook):
         parts = hostname.split(".")
         if hostname.endswith("amazonaws.com") and len(parts) == 6:
             return f"{parts[0]}.{parts[2]}"
-        else:
-            self.log.debug(
-                """Could not parse identifier from hostname '%s'.
+        self.log.debug(
+            """Could not parse identifier from hostname '%s'.
             You are probably using IP to connect to Redshift cluster.
             Expected format: 'cluster_identifier.id.region_name.redshift.amazonaws.com'
             Falling back to whole hostname.""",
-                hostname,
-            )
-            return hostname
+            hostname,
+        )
+        return hostname
 
     def get_openlineage_database_dialect(self, connection: Connection) -> str:
         """Return redshift dialect."""

airflow/providers/amazon/aws/hooks/s3.py

@@ -442,8 +442,7 @@ class S3Hook(AwsBaseHook):
         except ClientError as e:
             if e.response["ResponseMetadata"]["HTTPStatusCode"] == 404:
                 return head_object_val
-            else:
-                raise e
+            raise e
 
     async def list_prefixes_async(
         self,
@@ -936,8 +935,7 @@ class S3Hook(AwsBaseHook):
         except ClientError as e:
             if e.response["ResponseMetadata"]["HTTPStatusCode"] == 404:
                 return None
-            else:
-                raise e
+            raise e
 
     @unify_bucket_name_and_key
     @provide_bucket_name
@@ -1469,8 +1467,7 @@ class S3Hook(AwsBaseHook):
                 raise AirflowNotFoundException(
                     f"The source file in Bucket {bucket_name} with path {key} does not exist"
                 )
-            else:
-                raise e
+            raise e
 
         if preserve_file_name:
             local_dir = local_path or gettempdir()

airflow/providers/amazon/aws/hooks/sagemaker.py

@@ -750,7 +750,7 @@ class SageMakerHook(AwsBaseHook):
 
             if status in self.failed_states:
                 raise AirflowException(f"SageMaker resource failed because {response['FailureReason']}")
-            elif status not in non_terminal_states:
+            if status not in non_terminal_states:
                 break
 
             if max_ingestion_time and sec > max_ingestion_time:
@@ -1010,8 +1010,7 @@ class SageMakerHook(AwsBaseHook):
             if "NextToken" not in response or (max_results is not None and len(results) == max_results):
                 # Return when there are no results left (no NextToken) or when we've reached max_results.
                 return results
-            else:
-                next_token = response["NextToken"]
+            next_token = response["NextToken"]
 
     @staticmethod
     def _name_matches_pattern(
@@ -1172,9 +1171,8 @@ class SageMakerHook(AwsBaseHook):
                     ):
                         self.log.warning("Cannot stop pipeline execution, as it was not running: %s", ce)
                         break
-                    else:
-                        self.log.error(ce)
-                        raise
+                    self.log.error(ce)
+                    raise
             else:
                 break
 
@@ -1214,9 +1212,8 @@ class SageMakerHook(AwsBaseHook):
                 # log msg only so it doesn't look like an error
                 self.log.info("%s", e.response["Error"]["Message"])
                 return False
-            else:
-                self.log.error("Error when trying to create Model Package Group: %s", e)
-                raise
+            self.log.error("Error when trying to create Model Package Group: %s", e)
+            raise
 
     def _describe_auto_ml_job(self, job_name: str):
         res = self.conn.describe_auto_ml_job(AutoMLJobName=job_name)

airflow/providers/amazon/aws/hooks/sagemaker_unified_studio.py

@@ -180,9 +180,8 @@ class SageMakerNotebookHook(BaseHook):
         if status in finished_states:
             self.log.info(execution_message)
             return {"Status": status, "ExecutionId": execution_id}
-        else:
-            log_error_message = f"Execution {execution_id} failed with error: {error_message}"
-            self.log.error(log_error_message)
-            if error_message == "":
-                error_message = execution_message
-            raise AirflowException(error_message)
+        log_error_message = f"Execution {execution_id} failed with error: {error_message}"
+        self.log.error(log_error_message)
+        if error_message == "":
+            error_message = execution_message
+        raise AirflowException(error_message)

airflow/providers/amazon/aws/links/base_aws.py

@@ -49,9 +49,9 @@ class BaseAwsLink(BaseOperatorLink):
     def get_aws_domain(aws_partition) -> str | None:
         if aws_partition == "aws":
             return "aws.amazon.com"
-        elif aws_partition == "aws-cn":
+        if aws_partition == "aws-cn":
             return "amazonaws.cn"
-        elif aws_partition == "aws-us-gov":
+        if aws_partition == "aws-us-gov":
             return "amazonaws-us-gov.com"
 
         return None

airflow/providers/amazon/aws/links/emr.py

@@ -127,8 +127,7 @@ class EmrServerlessLogsLink(BaseAwsLink):
         )
         if url:
             return url._replace(path="/logs/SPARK_DRIVER/stdout.gz").geturl()
-        else:
-            return ""
+        return ""
 
 
 class EmrServerlessDashboardLink(BaseAwsLink):
@@ -145,8 +144,7 @@ class EmrServerlessDashboardLink(BaseAwsLink):
         )
         if url:
             return url.geturl()
-        else:
-            return ""
+        return ""
 
 
 class EmrServerlessS3LogsLink(BaseAwsLink):

airflow/providers/amazon/aws/log/cloudwatch_task_handler.py

@@ -17,6 +17,7 @@
 # under the License.
 from __future__ import annotations
 
+import contextlib
 import copy
 import json
 import logging
@@ -56,8 +57,7 @@ def json_serialize_legacy(value: Any) -> str | None:
     """
     if isinstance(value, (date, datetime)):
         return value.isoformat()
-    else:
-        return None
+    return None
 
 
 def json_serialize(value: Any) -> str | None:
@@ -134,10 +134,8 @@ class CloudWatchRemoteLogIO(LoggingMixin):  # noqa: D101
             msg = copy.copy(event)
             created = None
             if ts := msg.pop("timestamp", None):
-                try:
+                with contextlib.suppress(Exception):
                     created = datetime.fromisoformat(ts)
-                except Exception:
-                    pass
             record = logRecordFactory(
                 name, level, pathname="", lineno=0, msg=msg, args=(), exc_info=None, func=None, sinfo=None
             )

airflow/providers/amazon/aws/log/s3_task_handler.py

@@ -162,8 +162,7 @@ class S3RemoteLogIO(LoggingMixin):  # noqa: D101
             for key in keys:
                 logs.append(self.s3_read(key, return_error=True))
             return messages, logs
-        else:
-            return messages, None
+        return messages, None
 
 
 class S3TaskHandler(FileTaskHandler, LoggingMixin):

airflow/providers/amazon/aws/operators/athena.py

@@ -168,7 +168,7 @@ class AthenaOperator(AwsBaseOperator[AthenaHook]):
                 f"Final state of Athena job is {query_status}, query_execution_id is "
                 f"{self.query_execution_id}. Error: {error_message}"
             )
-        elif not query_status or query_status in AthenaHook.INTERMEDIATE_STATES:
+        if not query_status or query_status in AthenaHook.INTERMEDIATE_STATES:
             raise AirflowException(
                 f"Final state of Athena job is {query_status}. Max tries of poll status exceeded, "
                 f"query_execution_id is {self.query_execution_id}."