apache-airflow-providers-amazon 8.23.0__py3-none-any.whl → 8.24.0__py3-none-any.whl

airflow/providers/amazon/LICENSE

@@ -215,7 +215,7 @@ Third party Apache 2.0 licenses
 
 The following components are provided under the Apache 2.0 License.
 See project link for details. The text of each license is also included
-at licenses/LICENSE-[project].txt.
+at 3rd-party-licenses/LICENSE-[project].txt.
 
     (ALv2 License) hue v4.3.0 (https://github.com/cloudera/hue/)
     (ALv2 License) jqclock v2.3.0 (https://github.com/JohnRDOrazio/jQuery-Clock-Plugin)
@@ -227,7 +227,7 @@ MIT licenses
 ========================================================================
 
 The following components are provided under the MIT License. See project link for details.
-The text of each license is also included at licenses/LICENSE-[project].txt.
+The text of each license is also included at 3rd-party-licenses/LICENSE-[project].txt.
 
     (MIT License) jquery v3.5.1 (https://jquery.org/license/)
     (MIT License) dagre-d3 v0.6.4 (https://github.com/cpettitt/dagre-d3)
@@ -243,11 +243,11 @@ The text of each license is also included at licenses/LICENSE-[project].txt.
 BSD 3-Clause licenses
 ========================================================================
 The following components are provided under the BSD 3-Clause license. See project links for details.
-The text of each license is also included at licenses/LICENSE-[project].txt.
+The text of each license is also included at 3rd-party-licenses/LICENSE-[project].txt.
 
     (BSD 3 License) d3 v5.16.0 (https://d3js.org)
     (BSD 3 License) d3-shape v2.1.0 (https://github.com/d3/d3-shape)
     (BSD 3 License) cgroupspy 0.2.1 (https://github.com/cloudsigma/cgroupspy)
 
 ========================================================================
-See licenses/LICENSES-ui.txt for packages used in `/airflow/www`
+See 3rd-party-licenses/LICENSES-ui.txt for packages used in `/airflow/www`

airflow/providers/amazon/__init__.py

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

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

@@ -156,7 +156,9 @@ class BaseSessionFactory(LoggingMixin):
 
         return async_get_session()
 
-    def create_session(self, deferrable: bool = False) -> boto3.session.Session:
+    def create_session(
+        self, deferrable: bool = False
+    ) -> boto3.session.Session | aiobotocore.session.AioSession:
         """Create boto3 or aiobotocore Session from connection config."""
         if not self.conn:
             self.log.info(
@@ -198,7 +200,7 @@ class BaseSessionFactory(LoggingMixin):
 
     def _create_session_with_assume_role(
         self, session_kwargs: dict[str, Any], deferrable: bool = False
-    ) -> boto3.session.Session:
+    ) -> boto3.session.Session | aiobotocore.session.AioSession:
         if self.conn.assume_role_method == "assume_role_with_web_identity":
             # Deferred credentials have no initial credentials
             credential_fetcher = self._get_web_identity_credential_fetcher()
@@ -239,7 +241,10 @@ class BaseSessionFactory(LoggingMixin):
         session._credentials = credentials
         session.set_config_variable("region", self.basic_session.region_name)
 
-        return boto3.session.Session(botocore_session=session, **session_kwargs)
+        if not deferrable:
+            return boto3.session.Session(botocore_session=session, **session_kwargs)
+
+        return session
 
     def _refresh_credentials(self) -> dict[str, Any]:
         self.log.debug("Refreshing credentials")

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

@@ -102,6 +102,7 @@ class BatchProtocol(Protocol):
         arrayProperties: dict,
         parameters: dict,
         containerOverrides: dict,
+        ecsPropertiesOverride: dict,
         tags: dict,
     ) -> dict:
         """
@@ -119,6 +120,8 @@ class BatchProtocol(Protocol):
 
         :param containerOverrides: the same parameter that boto3 will receive
 
+        :param ecsPropertiesOverride: the same parameter that boto3 will receive
+
         :param tags: the same parameter that boto3 will receive
 
         :return: an API response

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

@@ -19,11 +19,17 @@
 
 from __future__ import annotations
 
-from typing import Iterable
+from functools import cached_property
+from typing import TYPE_CHECKING, Iterable
+
+from botocore.exceptions import ClientError
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 
+if TYPE_CHECKING:
+    from botocore.client import BaseClient
+
 
 class DynamoDBHook(AwsBaseHook):
     """
@@ -50,6 +56,11 @@ class DynamoDBHook(AwsBaseHook):
         kwargs["resource_type"] = "dynamodb"
         super().__init__(*args, **kwargs)
 
+    @cached_property
+    def client(self) -> BaseClient:
+        """Return boto3 client."""
+        return self.get_conn().meta.client
+
     def write_batch_data(self, items: Iterable) -> bool:
         """
         Write batch items to DynamoDB table with provisioned throughout capacity.
@@ -70,3 +81,25 @@ class DynamoDBHook(AwsBaseHook):
             return True
         except Exception as general_error:
             raise AirflowException(f"Failed to insert items in dynamodb, error: {general_error}")
+
+    def get_import_status(self, import_arn: str) -> tuple[str, str | None, str | None]:
+        """
+        Get import status from Dynamodb.
+
+        :param import_arn: The Amazon Resource Name (ARN) for the import.
+        :return: Import status, Error code and Error message
+        """
+        self.log.info("Poking for Dynamodb import %s", import_arn)
+
+        try:
+            describe_import = self.client.describe_import(ImportArn=import_arn)
+            status = describe_import["ImportTableDescription"]["ImportStatus"]
+            error_code = describe_import["ImportTableDescription"].get("FailureCode")
+            error_msg = describe_import["ImportTableDescription"].get("FailureMessage")
+            return status, error_code, error_msg
+        except ClientError as e:
+            error_code = e.response.get("Error", {}).get("Code")
+            if error_code == "ImportNotFoundException":
+                raise AirflowException("S3 import into Dynamodb job not found.")
+            else:
+                raise e

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

@@ -20,6 +20,7 @@ from __future__ import annotations
 import asyncio
 import time
 from functools import cached_property
+from typing import Any
 
 from botocore.exceptions import ClientError
 
@@ -430,3 +431,125 @@ class GlueJobHook(AwsBaseHook):
             self.conn.create_job(**config)
 
         return self.job_name
+
+
+class GlueDataQualityHook(AwsBaseHook):
+    """
+    Interact with AWS Glue Data Quality.
+
+    Provide thick wrapper around :external+boto3:py:class:`boto3.client("glue") <Glue.Client>`.
+
+    Additional arguments (such as ``aws_conn_id``) may be specified and
+    are passed down to the underlying AwsBaseHook.
+
+    .. seealso::
+        - :class:`airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook`
+    """
+
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ):
+        kwargs["client_type"] = "glue"
+        super().__init__(*args, **kwargs)
+
+    def has_data_quality_ruleset(self, name: str) -> bool:
+        try:
+            self.conn.get_data_quality_ruleset(Name=name)
+            return True
+        except self.conn.exceptions.EntityNotFoundException:
+            return False
+
+    def _log_results(self, result: dict[str, Any]) -> None:
+        """
+        Print the outcome of evaluation run, An evaluation run can involve multiple rulesets evaluated against a data source (Glue table).
+
+        Name    Description                                     Result        EvaluatedMetrics                                                                    EvaluationMessage
+        Rule_1    RowCount between 150000 and 600000             PASS        {'Dataset.*.RowCount': 300000.0}                                                       NaN
+        Rule_2    IsComplete "marketplace"                       PASS        {'Column.marketplace.Completeness': 1.0}                                               NaN
+        Rule_3    ColumnLength "marketplace" between 1 and 2     FAIL        {'Column.marketplace.MaximumLength': 9.0, 'Column.marketplace.MinimumLength': 3.0}     Value: 9.0 does not meet the constraint requirement!
+
+        """
+        import pandas as pd
+
+        pd.set_option("display.max_rows", None)
+        pd.set_option("display.max_columns", None)
+        pd.set_option("display.width", None)
+        pd.set_option("display.max_colwidth", None)
+
+        self.log.info(
+            "AWS Glue data quality ruleset evaluation result for RulesetName: %s RulesetEvaluationRunId: %s Score: %s",
+            result.get("RulesetName"),
+            result.get("RulesetEvaluationRunId"),
+            result.get("Score"),
+        )
+
+        rule_results = result["RuleResults"]
+        rule_results_df = pd.DataFrame(rule_results)
+        self.log.info(rule_results_df)
+
+    def get_evaluation_run_results(self, run_id: str) -> dict[str, Any]:
+        response = self.conn.get_data_quality_ruleset_evaluation_run(RunId=run_id)
+
+        return self.conn.batch_get_data_quality_result(ResultIds=response["ResultIds"])
+
+    def validate_evaluation_run_results(
+        self, evaluation_run_id: str, show_results: bool = True, verify_result_status: bool = True
+    ) -> None:
+        results = self.get_evaluation_run_results(evaluation_run_id)
+        total_failed_rules = 0
+
+        if results.get("ResultsNotFound"):
+            self.log.info(
+                "AWS Glue data quality ruleset evaluation run, results not found for %s",
+                results["ResultsNotFound"],
+            )
+
+        for result in results["Results"]:
+            rule_results = result["RuleResults"]
+
+            total_failed_rules += len(
+                list(
+                    filter(
+                        lambda result: result.get("Result") == "FAIL" or result.get("Result") == "ERROR",
+                        rule_results,
+                    )
+                )
+            )
+
+            if show_results:
+                self._log_results(result)
+
+        self.log.info(
+            "AWS Glue data quality ruleset evaluation run, total number of rules failed: %s",
+            total_failed_rules,
+        )
+
+        if verify_result_status and total_failed_rules > 0:
+            raise AirflowException(
+                "AWS Glue data quality ruleset evaluation run failed for one or more rules"
+            )
+
+    def log_recommendation_results(self, run_id: str) -> None:
+        """
+        Print the outcome of recommendation run, recommendation run generates multiple rules against a data source (Glue table) in Data Quality Definition Language (DQDL) format.
+
+        Rules = [
+        IsComplete "NAME",
+        ColumnLength "EMP_ID" between 1 and 12,
+        IsUnique "EMP_ID",
+        ColumnValues "INCOME" > 50000
+        ]
+        """
+        result = self.conn.get_data_quality_rule_recommendation_run(RunId=run_id)
+
+        if result.get("RecommendedRuleset"):
+            self.log.info(
+                "AWS Glue data quality recommended rules for DatabaseName: %s TableName: %s",
+                result["DataSource"]["GlueTable"]["DatabaseName"],
+                result["DataSource"]["GlueTable"]["TableName"],
+            )
+            self.log.info(result["RecommendedRuleset"])
+        else:
+            self.log.info("AWS Glue data quality, no recommended rules available for RunId: %s", run_id)

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

@@ -65,6 +65,7 @@ class BatchOperator(BaseOperator):
     :param job_queue: the queue name on AWS Batch
     :param overrides: DEPRECATED, use container_overrides instead with the same value.
     :param container_overrides: the `containerOverrides` parameter for boto3 (templated)
+    :param ecs_properties_override: the `ecsPropertiesOverride` parameter for boto3 (templated)
     :param node_overrides: the `nodeOverrides` parameter for boto3 (templated)
     :param share_identifier: The share identifier for the job. Don't specify this parameter if the job queue
         doesn't have a scheduling policy.
@@ -112,6 +113,7 @@ class BatchOperator(BaseOperator):
         "job_queue",
         "container_overrides",
         "array_properties",
+        "ecs_properties_override",
         "node_overrides",
         "parameters",
         "retry_strategy",
@@ -124,6 +126,7 @@ class BatchOperator(BaseOperator):
     template_fields_renderers = {
         "container_overrides": "json",
         "parameters": "json",
+        "ecs_properties_override": "json",
         "node_overrides": "json",
         "retry_strategy": "json",
     }
@@ -160,6 +163,7 @@ class BatchOperator(BaseOperator):
         overrides: dict | None = None,  # deprecated
         container_overrides: dict | None = None,
         array_properties: dict | None = None,
+        ecs_properties_override: dict | None = None,
         node_overrides: dict | None = None,
         share_identifier: str | None = None,
         scheduling_priority_override: int | None = None,
@@ -201,6 +205,7 @@ class BatchOperator(BaseOperator):
                 stacklevel=2,
             )
 
+        self.ecs_properties_override = ecs_properties_override
         self.node_overrides = node_overrides
         self.share_identifier = share_identifier
         self.scheduling_priority_override = scheduling_priority_override
@@ -296,6 +301,8 @@ class BatchOperator(BaseOperator):
             self.log.info("AWS Batch job - container overrides: %s", self.container_overrides)
         if self.array_properties:
             self.log.info("AWS Batch job - array properties: %s", self.array_properties)
+        if self.ecs_properties_override:
+            self.log.info("AWS Batch job - ECS properties: %s", self.ecs_properties_override)
         if self.node_overrides:
             self.log.info("AWS Batch job - node properties: %s", self.node_overrides)
 
@@ -307,6 +314,7 @@ class BatchOperator(BaseOperator):
             "parameters": self.parameters,
             "tags": self.tags,
             "containerOverrides": self.container_overrides,
+            "ecsPropertiesOverride": self.ecs_properties_override,
             "nodeOverrides": self.node_overrides,
             "retryStrategy": self.retry_strategy,
             "shareIdentifier": self.share_identifier,

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

@@ -20,7 +20,6 @@ import json
 from time import sleep
 from typing import TYPE_CHECKING, Any, Sequence
 
-import botocore
 from botocore.exceptions import ClientError
 
 from airflow.configuration import conf
@@ -38,7 +37,7 @@ from airflow.providers.amazon.aws.triggers.bedrock import (
     BedrockKnowledgeBaseActiveTrigger,
     BedrockProvisionModelThroughputCompletedTrigger,
 )
-from airflow.providers.amazon.aws.utils import get_botocore_version, validate_execute_complete_event
+from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
 from airflow.utils.helpers import prune_dict
 from airflow.utils.timezone import utcnow
@@ -799,24 +798,11 @@ class BedrockRaGOperator(AwsBaseOperator[BedrockAgentRuntimeHook]):
     def execute(self, context: Context) -> Any:
         self.validate_inputs()
 
-        try:
-            result = self.hook.conn.retrieve_and_generate(
-                input={"text": self.input},
-                retrieveAndGenerateConfiguration=self.build_rag_config(),
-                **self.rag_kwargs,
-            )
-        except botocore.exceptions.ParamValidationError as error:
-            if (
-                'Unknown parameter in retrieveAndGenerateConfiguration: "externalSourcesConfiguration"'
-                in str(error)
-            ) and (self.source_type == "EXTERNAL_SOURCES"):
-                self.log.error(
-                    "You are attempting to use External Sources and the BOTO API returned an "
-                    "error message which may indicate the need to update botocore to do this.  \n"
-                    "Support for External Sources was added in botocore 1.34.90 and you are using botocore %s",
-                    ".".join(map(str, get_botocore_version())),
-                )
-            raise
+        result = self.hook.conn.retrieve_and_generate(
+            input={"text": self.input},
+            retrieveAndGenerateConfiguration=self.build_rag_config(),
+            **self.rag_kwargs,
+        )
 
         self.log.info(
             "\nPrompt: %s\nResponse: %s\nCitations: %s",

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

@@ -141,7 +141,7 @@ class EcsCreateClusterOperator(EcsBaseOperator):
                     waiter_delay=self.waiter_delay,
                     waiter_max_attempts=self.waiter_max_attempts,
                     aws_conn_id=self.aws_conn_id,
-                    region_name=self.region,
+                    region_name=self.region_name,
                 ),
                 method_name="_complete_exec_with_cluster_desc",
                 # timeout is set to ensure that if a trigger dies, the timeout does not restart
@@ -218,7 +218,7 @@ class EcsDeleteClusterOperator(EcsBaseOperator):
                     waiter_delay=self.waiter_delay,
                     waiter_max_attempts=self.waiter_max_attempts,
                     aws_conn_id=self.aws_conn_id,
-                    region_name=self.region,
+                    region_name=self.region_name,
                 ),
                 method_name="_complete_exec_with_cluster_desc",
                 # timeout is set to ensure that if a trigger dies, the timeout does not restart
@@ -495,7 +495,7 @@ class EcsRunTaskOperator(EcsBaseOperator):
         self.number_logs_exception = number_logs_exception
 
         if self.awslogs_region is None:
-            self.awslogs_region = self.region
+            self.awslogs_region = self.region_name
 
         self.arn: str | None = None
         self._started_by: str | None = None
@@ -546,7 +546,7 @@ class EcsRunTaskOperator(EcsBaseOperator):
                     waiter_delay=self.waiter_delay,
                     waiter_max_attempts=self.waiter_max_attempts,
                     aws_conn_id=self.aws_conn_id,
-                    region=self.region,
+                    region=self.region_name,
                     log_group=self.awslogs_group,
                     log_stream=self._get_logs_stream_name(),
                 ),
@@ -589,7 +589,7 @@ class EcsRunTaskOperator(EcsBaseOperator):
         self._after_execution()
         if self._aws_logs_enabled():
             # same behavior as non-deferrable mode, return last line of logs of the task.
-            logs_client = AwsLogsHook(aws_conn_id=self.aws_conn_id, region_name=self.region).conn
+            logs_client = AwsLogsHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name).conn
             one_log = logs_client.get_log_events(
                 logGroupName=self.awslogs_group,
                 logStreamName=self._get_logs_stream_name(),

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

@@ -263,30 +263,34 @@ class EmrStartNotebookExecutionOperator(BaseOperator):
         wait_for_completion: bool = False,
         aws_conn_id: str | None = "aws_default",
         # TODO: waiter_max_attempts and waiter_delay should default to None when the other two are deprecated.
-        waiter_max_attempts: int | None | ArgNotSet = NOTSET,
-        waiter_delay: int | None | ArgNotSet = NOTSET,
-        waiter_countdown: int = 25 * 60,
-        waiter_check_interval_seconds: int = 60,
+        waiter_max_attempts: int | None = None,
+        waiter_delay: int | None = None,
+        waiter_countdown: int | None = None,
+        waiter_check_interval_seconds: int | None = None,
         **kwargs: Any,
     ):
-        if waiter_max_attempts is NOTSET:
+        if waiter_check_interval_seconds:
             warnings.warn(
-                "The parameter waiter_countdown has been deprecated to standardize "
-                "naming conventions.  Please use waiter_max_attempts instead.  In the "
+                "The parameter `waiter_check_interval_seconds` has been deprecated to "
+                "standardize naming conventions.  Please `use waiter_delay instead`.  In the "
                 "future this will default to None and defer to the waiter's default value.",
                 AirflowProviderDeprecationWarning,
                 stacklevel=2,
             )
-            waiter_max_attempts = waiter_countdown // waiter_check_interval_seconds
-        if waiter_delay is NOTSET:
+        else:
+            waiter_check_interval_seconds = 60
+        if waiter_countdown:
             warnings.warn(
-                "The parameter waiter_check_interval_seconds has been deprecated to "
-                "standardize naming conventions.  Please use waiter_delay instead.  In the "
+                "The parameter waiter_countdown has been deprecated to standardize "
+                "naming conventions.  Please use waiter_max_attempts instead.  In the "
                 "future this will default to None and defer to the waiter's default value.",
                 AirflowProviderDeprecationWarning,
                 stacklevel=2,
             )
-            waiter_delay = waiter_check_interval_seconds
+            # waiter_countdown defaults to never timing out, which is not supported
+            # by boto waiters, so we will set it here to "a very long time" for now.
+            waiter_max_attempts = (waiter_countdown or 999) // waiter_check_interval_seconds
+
         super().__init__(**kwargs)
         self.editor_id = editor_id
         self.relative_path = relative_path
@@ -298,8 +302,8 @@ class EmrStartNotebookExecutionOperator(BaseOperator):
         self.wait_for_completion = wait_for_completion
         self.cluster_id = cluster_id
         self.aws_conn_id = aws_conn_id
-        self.waiter_max_attempts = waiter_max_attempts
-        self.waiter_delay = waiter_delay
+        self.waiter_max_attempts = waiter_max_attempts or 25
+        self.waiter_delay = waiter_delay or waiter_check_interval_seconds or 60
         self.master_instance_security_group_id = master_instance_security_group_id
 
     def execute(self, context: Context):
@@ -387,36 +391,40 @@ class EmrStopNotebookExecutionOperator(BaseOperator):
         wait_for_completion: bool = False,
         aws_conn_id: str | None = "aws_default",
         # TODO: waiter_max_attempts and waiter_delay should default to None when the other two are deprecated.
-        waiter_max_attempts: int | None | ArgNotSet = NOTSET,
-        waiter_delay: int | None | ArgNotSet = NOTSET,
-        waiter_countdown: int = 25 * 60,
-        waiter_check_interval_seconds: int = 60,
+        waiter_max_attempts: int | None = None,
+        waiter_delay: int | None = None,
+        waiter_countdown: int | None = None,
+        waiter_check_interval_seconds: int | None = None,
         **kwargs: Any,
     ):
-        if waiter_max_attempts is NOTSET:
+        if waiter_check_interval_seconds:
             warnings.warn(
-                "The parameter waiter_countdown has been deprecated to standardize "
-                "naming conventions.  Please use waiter_max_attempts instead.  In the "
+                "The parameter `waiter_check_interval_seconds` has been deprecated to "
+                "standardize naming conventions.  Please `use waiter_delay instead`.  In the "
                 "future this will default to None and defer to the waiter's default value.",
                 AirflowProviderDeprecationWarning,
                 stacklevel=2,
             )
-            waiter_max_attempts = waiter_countdown // waiter_check_interval_seconds
-        if waiter_delay is NOTSET:
+        else:
+            waiter_check_interval_seconds = 60
+        if waiter_countdown:
             warnings.warn(
-                "The parameter waiter_check_interval_seconds has been deprecated to "
-                "standardize naming conventions.  Please use waiter_delay instead.  In the "
+                "The parameter waiter_countdown has been deprecated to standardize "
+                "naming conventions.  Please use waiter_max_attempts instead.  In the "
                 "future this will default to None and defer to the waiter's default value.",
                 AirflowProviderDeprecationWarning,
                 stacklevel=2,
             )
-            waiter_delay = waiter_check_interval_seconds
+            # waiter_countdown defaults to never timing out, which is not supported
+            # by boto waiters, so we will set it here to "a very long time" for now.
+            waiter_max_attempts = (waiter_countdown or 999) // waiter_check_interval_seconds
+
         super().__init__(**kwargs)
         self.notebook_execution_id = notebook_execution_id
         self.wait_for_completion = wait_for_completion
         self.aws_conn_id = aws_conn_id
-        self.waiter_max_attempts = waiter_max_attempts
-        self.waiter_delay = waiter_delay
+        self.waiter_max_attempts = waiter_max_attempts or 25
+        self.waiter_delay = waiter_delay or waiter_check_interval_seconds or 60
 
     def execute(self, context: Context) -> None:
         emr_hook = EmrHook(aws_conn_id=self.aws_conn_id)
@@ -822,8 +830,8 @@ class EmrCreateJobFlowOperator(BaseOperator):
                 trigger=EmrCreateJobFlowTrigger(
                     job_flow_id=self._job_flow_id,
                     aws_conn_id=self.aws_conn_id,
-                    poll_interval=self.waiter_delay,
-                    max_attempts=self.waiter_max_attempts,
+                    waiter_delay=self.waiter_delay,
+                    waiter_max_attempts=self.waiter_max_attempts,
                 ),
                 method_name="execute_complete",
                 # timeout is set to ensure that if a trigger dies, the timeout does not restart