apache-airflow-providers-amazon 8.21.0rc1__tar.gz → 8.22.0rc1__tar.gz

{apache_airflow_providers_amazon-8.21.0rc1 → apache_airflow_providers_amazon-8.22.0rc1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: apache-airflow-providers-amazon
-Version: 8.21.0rc1
+Version: 8.22.0rc1
 Summary: Provider package apache-airflow-providers-amazon for Apache Airflow
 Keywords: airflow-provider,amazon,airflow,integration
 Author-email: Apache Software Foundation <dev@airflow.apache.org>
@@ -29,11 +29,11 @@ Requires-Dist: asgiref
 Requires-Dist: boto3>=1.33.0
 Requires-Dist: botocore>=1.33.0
 Requires-Dist: inflection>=0.5.1
+Requires-Dist: jmespath
 Requires-Dist: jsonpath_ng>=1.5.3
 Requires-Dist: redshift_connector>=2.0.918
 Requires-Dist: sqlalchemy_redshift>=0.8.6
 Requires-Dist: watchtower>=2.0.1,<4
-Requires-Dist: xmlsec<1.3.14
 Requires-Dist: aiobotocore[boto3]>=2.5.3 ; extra == "aiobotocore"
 Requires-Dist: apache-airflow-providers-apache-hive ; extra == "apache.hive"
 Requires-Dist: apache-airflow-providers-cncf-kubernetes>=7.2.0rc0 ; extra == "cncf.kubernetes"
@@ -52,8 +52,8 @@ Requires-Dist: s3fs>=2023.10.0 ; extra == "s3fs"
 Requires-Dist: apache-airflow-providers-salesforce ; extra == "salesforce"
 Requires-Dist: apache-airflow-providers-ssh ; extra == "ssh"
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
-Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.21.0/changelog.html
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.21.0
+Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.22.0/changelog.html
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.22.0
 Project-URL: Slack Chat, https://s.apache.org/airflow-slack
 Project-URL: Source Code, https://github.com/apache/airflow
 Project-URL: Twitter, https://twitter.com/ApacheAirflow
@@ -120,7 +120,7 @@ Provides-Extra: ssh
 
 Package ``apache-airflow-providers-amazon``
 
-Release: ``8.21.0.rc1``
+Release: ``8.22.0.rc1``
 
 
 Amazon integration (including `Amazon Web Services (AWS) <https://aws.amazon.com/>`__).
@@ -133,7 +133,7 @@ This is a provider package for ``amazon`` provider. All classes for this provide
 are in ``airflow.providers.amazon`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.21.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.22.0/>`_.
 
 Installation
 ------------
@@ -162,7 +162,7 @@ PIP package                              Version required
 ``sqlalchemy_redshift``                  ``>=0.8.6``
 ``asgiref``
 ``PyAthena``                             ``>=3.0.10``
-``xmlsec``                               ``<1.3.14``
+``jmespath``
 =======================================  ==================
 
 Cross provider package dependencies
@@ -197,4 +197,4 @@ Dependent package
 ======================================================================================================================  ===================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.21.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.22.0/changelog.html>`_.

{apache_airflow_providers_amazon-8.21.0rc1 → apache_airflow_providers_amazon-8.22.0rc1}/README.rst

@@ -42,7 +42,7 @@
 
 Package ``apache-airflow-providers-amazon``
 
-Release: ``8.21.0.rc1``
+Release: ``8.22.0.rc1``
 
 
 Amazon integration (including `Amazon Web Services (AWS) <https://aws.amazon.com/>`__).
@@ -55,7 +55,7 @@ This is a provider package for ``amazon`` provider. All classes for this provide
 are in ``airflow.providers.amazon`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.21.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.22.0/>`_.
 
 Installation
 ------------
@@ -84,7 +84,7 @@ PIP package                              Version required
 ``sqlalchemy_redshift``                  ``>=0.8.6``
 ``asgiref``
 ``PyAthena``                             ``>=3.0.10``
-``xmlsec``                               ``<1.3.14``
+``jmespath``
 =======================================  ==================
 
 Cross provider package dependencies
@@ -119,4 +119,4 @@ Dependent package
 ======================================================================================================================  ===================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.21.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.22.0/changelog.html>`_.

{apache_airflow_providers_amazon-8.21.0rc1 → apache_airflow_providers_amazon-8.22.0rc1}/airflow/providers/amazon/__init__.py

@@ -25,14 +25,11 @@ from __future__ import annotations
 
 import packaging.version
 
-__all__ = ["__version__"]
+from airflow import __version__ as airflow_version
 
-__version__ = "8.21.0"
+__all__ = ["__version__"]
 
-try:
-    from airflow import __version__ as airflow_version
-except ImportError:
-    from airflow.version import version as airflow_version
+__version__ = "8.22.0"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
     "2.7.0"

{apache_airflow_providers_amazon-8.21.0rc1 → apache_airflow_providers_amazon-8.22.0rc1}/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py

@@ -400,7 +400,12 @@ class AwsEcsExecutor(BaseExecutor):
             else:
                 task = run_task_response["tasks"][0]
                 self.active_workers.add_task(task, task_key, queue, cmd, exec_config, attempt_number)
-                self.queued(task_key, task.task_arn)
+                try:
+                    self.running_state(task_key, task.task_arn)
+                except AttributeError:
+                    # running_state is newly added, and only needed to support task adoption (an optional
+                    # executor feature).
+                    pass
         if failure_reasons:
             self.log.error(
                 "Pending ECS tasks failed to launch for the following reasons: %s. Retrying later.",
@@ -522,7 +527,7 @@ class AwsEcsExecutor(BaseExecutor):
                         ti.queue,
                         ti.command_as_list(),
                         ti.executor_config,
-                        ti.prev_attempted_tries,
+                        ti.try_number,
                     )
                     adopted_tis.append(ti)
 

{apache_airflow_providers_amazon-8.21.0rc1 → apache_airflow_providers_amazon-8.22.0rc1}/airflow/providers/amazon/aws/hooks/bedrock.py

@@ -57,3 +57,23 @@ class BedrockRuntimeHook(AwsBaseHook):
     def __init__(self, *args, **kwargs) -> None:
         kwargs["client_type"] = self.client_type
         super().__init__(*args, **kwargs)
+
+
+class BedrockAgentHook(AwsBaseHook):
+    """
+    Interact with the Amazon Agents for Bedrock API.
+
+    Provide thin wrapper around :external+boto3:py:class:`boto3.client("bedrock-agent") <AgentsforBedrock.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`
+    """
+
+    client_type = "bedrock-agent"
+
+    def __init__(self, *args, **kwargs) -> None:
+        kwargs["client_type"] = self.client_type
+        super().__init__(*args, **kwargs)

apache_airflow_providers_amazon-8.22.0rc1/airflow/providers/amazon/aws/hooks/opensearch_serverless.py

@@ -0,0 +1,39 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+
+
+class OpenSearchServerlessHook(AwsBaseHook):
+    """
+    Interact with the Amazon OpenSearch Serverless API.
+
+    Provide thin wrapper around :external+boto3:py:class:`boto3.client("opensearchserverless") <OpenSearchServiceServerless.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`
+    """
+
+    client_type = "opensearchserverless"
+
+    def __init__(self, *args, **kwargs) -> None:
+        kwargs["client_type"] = self.client_type
+        super().__init__(*args, **kwargs)

{apache_airflow_providers_amazon-8.21.0rc1 → apache_airflow_providers_amazon-8.22.0rc1}/airflow/providers/amazon/aws/hooks/s3.py

@@ -52,7 +52,7 @@ from asgiref.sync import sync_to_async
 from boto3.s3.transfer import S3Transfer, TransferConfig
 from botocore.exceptions import ClientError
 
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowException, AirflowNotFoundException, AirflowProviderDeprecationWarning
 from airflow.providers.amazon.aws.exceptions import S3HookUriParseFailure
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 from airflow.providers.amazon.aws.utils.tags import format_tags
@@ -1390,7 +1390,7 @@ class S3Hook(AwsBaseHook):
             s3_obj = self.get_key(key, bucket_name)
         except ClientError as e:
             if e.response.get("Error", {}).get("Code") == 404:
-                raise AirflowException(
+                raise AirflowNotFoundException(
                     f"The source file in Bucket {bucket_name} with path {key} does not exist"
                 )
             else:

{apache_airflow_providers_amazon-8.21.0rc1 → apache_airflow_providers_amazon-8.22.0rc1}/airflow/providers/amazon/aws/operators/bedrock.py

@@ -17,16 +17,19 @@
 from __future__ import annotations
 
 import json
+from time import sleep
 from typing import TYPE_CHECKING, Any, Sequence
 
 from botocore.exceptions import ClientError
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
-from airflow.providers.amazon.aws.hooks.bedrock import BedrockHook, BedrockRuntimeHook
+from airflow.providers.amazon.aws.hooks.bedrock import BedrockAgentHook, BedrockHook, BedrockRuntimeHook
 from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
 from airflow.providers.amazon.aws.triggers.bedrock import (
     BedrockCustomizeModelCompletedTrigger,
+    BedrockIngestionJobTrigger,
+    BedrockKnowledgeBaseActiveTrigger,
     BedrockProvisionModelThroughputCompletedTrigger,
 )
 from airflow.providers.amazon.aws.utils import validate_execute_complete_event
@@ -351,3 +354,313 @@ class BedrockCreateProvisionedModelThroughputOperator(AwsBaseOperator[BedrockHoo
 
         self.log.info("Bedrock provisioned throughput job `%s` complete.", event["provisioned_model_id"])
         return event["provisioned_model_id"]
+
+
+class BedrockCreateKnowledgeBaseOperator(AwsBaseOperator[BedrockAgentHook]):
+    """
+    Create a knowledge base that contains data sources used by Amazon Bedrock LLMs and Agents.
+
+    To create a knowledge base, you must first set up your data sources and configure a supported vector store.
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:BedrockCreateKnowledgeBaseOperator`
+
+    :param name: The name of the knowledge base. (templated)
+    :param embedding_model_arn: ARN of the model used to create vector embeddings for the knowledge base. (templated)
+    :param role_arn: The ARN of the IAM role with permissions to create the knowledge base. (templated)
+    :param storage_config: Configuration details of the vector database used for the knowledge base. (templated)
+    :param wait_for_indexing: Vector indexing can take some time and there is no apparent way to check the state
+        before trying to create the Knowledge Base.  If this is True, and creation fails due to the index not
+        being available, the operator will wait and retry.  (default: True) (templated)
+    :param indexing_error_retry_delay: Seconds between retries if an index error is encountered. (default 5) (templated)
+    :param indexing_error_max_attempts: Maximum number of times to retry when encountering an index error. (default 20) (templated)
+    :param create_knowledge_base_kwargs: Any additional optional parameters to pass to the API call. (templated)
+
+    :param wait_for_completion: Whether to wait for cluster to stop. (default: True)
+    :param waiter_delay: Time in seconds to wait between status checks. (default: 60)
+    :param waiter_max_attempts: Maximum number of attempts to check for job completion. (default: 20)
+    :param deferrable: If True, the operator will wait asynchronously for the cluster to stop.
+        This implies waiting for completion. This mode requires aiobotocore module to be installed.
+        (default: False)
+    :param aws_conn_id: The Airflow connection used for AWS credentials.
+        If this is ``None`` or empty then the default boto3 behaviour is used. If
+        running Airflow in a distributed manner and aws_conn_id is None or
+        empty, then default boto3 configuration would be used (and must be
+        maintained on each worker node).
+    :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+    :param verify: Whether or not to verify SSL certificates. See:
+        https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
+    :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+        https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+    """
+
+    aws_hook_class = BedrockAgentHook
+    template_fields: Sequence[str] = aws_template_fields(
+        "name",
+        "embedding_model_arn",
+        "role_arn",
+        "storage_config",
+        "wait_for_indexing",
+        "indexing_error_retry_delay",
+        "indexing_error_max_attempts",
+        "create_knowledge_base_kwargs",
+    )
+
+    def __init__(
+        self,
+        name: str,
+        embedding_model_arn: str,
+        role_arn: str,
+        storage_config: dict[str, Any],
+        create_knowledge_base_kwargs: dict[str, Any] | None = None,
+        wait_for_indexing: bool = True,
+        indexing_error_retry_delay: int = 5,  # seconds
+        indexing_error_max_attempts: int = 20,
+        wait_for_completion: bool = True,
+        waiter_delay: int = 60,
+        waiter_max_attempts: int = 20,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.name = name
+        self.role_arn = role_arn
+        self.storage_config = storage_config
+        self.create_knowledge_base_kwargs = create_knowledge_base_kwargs or {}
+        self.embedding_model_arn = embedding_model_arn
+        self.knowledge_base_config = {
+            "type": "VECTOR",
+            "vectorKnowledgeBaseConfiguration": {"embeddingModelArn": self.embedding_model_arn},
+        }
+        self.wait_for_indexing = wait_for_indexing
+        self.indexing_error_retry_delay = indexing_error_retry_delay
+        self.indexing_error_max_attempts = indexing_error_max_attempts
+
+        self.wait_for_completion = wait_for_completion
+        self.waiter_delay = waiter_delay
+        self.waiter_max_attempts = waiter_max_attempts
+        self.deferrable = deferrable
+
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:
+        event = validate_execute_complete_event(event)
+
+        if event["status"] != "success":
+            raise AirflowException(f"Error while running job: {event}")
+
+        self.log.info("Bedrock knowledge base creation job `%s` complete.", self.name)
+        return event["knowledge_base_id"]
+
+    def execute(self, context: Context) -> str:
+        def _create_kb():
+            # This API call will return the following if the index has not completed, but there is no apparent
+            # way to check the state of the index beforehand, so retry on index failure if set to do so.
+            #       botocore.errorfactory.ValidationException: An error occurred (ValidationException)
+            #       when calling the CreateKnowledgeBase operation: The knowledge base storage configuration
+            #       provided is invalid... no such index [bedrock-sample-rag-index-abc108]
+            try:
+                return self.hook.conn.create_knowledge_base(
+                    name=self.name,
+                    roleArn=self.role_arn,
+                    knowledgeBaseConfiguration=self.knowledge_base_config,
+                    storageConfiguration=self.storage_config,
+                    **self.create_knowledge_base_kwargs,
+                )["knowledgeBase"]["knowledgeBaseId"]
+            except ClientError as error:
+                if all(
+                    [
+                        error.response["Error"]["Code"] == "ValidationException",
+                        "no such index" in error.response["Error"]["Message"],
+                        self.wait_for_indexing,
+                        self.indexing_error_max_attempts > 0,
+                    ]
+                ):
+                    self.indexing_error_max_attempts -= 1
+                    self.log.warning(
+                        "Vector index not ready, retrying in %s seconds.", self.indexing_error_retry_delay
+                    )
+                    self.log.debug("%s retries remaining.", self.indexing_error_max_attempts)
+                    sleep(self.indexing_error_retry_delay)
+                    return _create_kb()
+                raise
+
+        self.log.info("Creating Amazon Bedrock Knowledge Base %s", self.name)
+        knowledge_base_id = _create_kb()
+
+        if self.deferrable:
+            self.log.info("Deferring for Knowledge base creation.")
+            self.defer(
+                trigger=BedrockKnowledgeBaseActiveTrigger(
+                    knowledge_base_id=knowledge_base_id,
+                    waiter_delay=self.waiter_delay,
+                    waiter_max_attempts=self.waiter_max_attempts,
+                    aws_conn_id=self.aws_conn_id,
+                ),
+                method_name="execute_complete",
+            )
+        if self.wait_for_completion:
+            self.log.info("Waiting for Knowledge Base creation.")
+            self.hook.get_waiter("knowledge_base_active").wait(
+                knowledgeBaseId=knowledge_base_id,
+                WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+            )
+
+        return knowledge_base_id
+
+
+class BedrockCreateDataSourceOperator(AwsBaseOperator[BedrockAgentHook]):
+    """
+    Set up an Amazon Bedrock Data Source to be added to an Amazon Bedrock Knowledge Base.
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:BedrockCreateDataSourceOperator`
+
+    :param name: name for the Amazon Bedrock Data Source being created. (templated).
+    :param bucket_name: The name of the Amazon S3 bucket to use for data source storage. (templated)
+    :param knowledge_base_id: The unique identifier of the knowledge base to which to add the data source. (templated)
+    :param create_data_source_kwargs: Any additional optional parameters to pass to the API call. (templated)
+
+    :param aws_conn_id: The Airflow connection used for AWS credentials.
+        If this is ``None`` or empty then the default boto3 behaviour is used. If
+        running Airflow in a distributed manner and aws_conn_id is None or
+        empty, then default boto3 configuration would be used (and must be
+        maintained on each worker node).
+    :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+    :param verify: Whether or not to verify SSL certificates. See:
+        https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
+    :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+        https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+    """
+
+    aws_hook_class = BedrockAgentHook
+    template_fields: Sequence[str] = aws_template_fields(
+        "name",
+        "bucket_name",
+        "knowledge_base_id",
+        "create_data_source_kwargs",
+    )
+
+    def __init__(
+        self,
+        name: str,
+        knowledge_base_id: str,
+        bucket_name: str | None = None,
+        create_data_source_kwargs: dict[str, Any] | None = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.name = name
+        self.knowledge_base_id = knowledge_base_id
+        self.bucket_name = bucket_name
+        self.create_data_source_kwargs = create_data_source_kwargs or {}
+
+    def execute(self, context: Context) -> str:
+        create_ds_response = self.hook.conn.create_data_source(
+            name=self.name,
+            knowledgeBaseId=self.knowledge_base_id,
+            dataSourceConfiguration={
+                "type": "S3",
+                "s3Configuration": {"bucketArn": f"arn:aws:s3:::{self.bucket_name}"},
+            },
+            **self.create_data_source_kwargs,
+        )
+
+        return create_ds_response["dataSource"]["dataSourceId"]
+
+
+class BedrockIngestDataOperator(AwsBaseOperator[BedrockAgentHook]):
+    """
+    Begin an ingestion job, in which an Amazon Bedrock data source is added to an Amazon Bedrock knowledge base.
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:BedrockIngestDataOperator`
+
+    :param knowledge_base_id: The unique identifier of the knowledge base to which to add the data source. (templated)
+    :param data_source_id: The unique identifier of the data source to ingest. (templated)
+    :param ingest_data_kwargs: Any additional optional parameters to pass to the API call. (templated)
+
+    :param wait_for_completion: Whether to wait for cluster to stop. (default: True)
+    :param waiter_delay: Time in seconds to wait between status checks. (default: 60)
+    :param waiter_max_attempts: Maximum number of attempts to check for job completion. (default: 10)
+    :param deferrable: If True, the operator will wait asynchronously for the cluster to stop.
+        This implies waiting for completion. This mode requires aiobotocore module to be installed.
+        (default: False)
+    :param aws_conn_id: The Airflow connection used for AWS credentials.
+        If this is ``None`` or empty then the default boto3 behaviour is used. If
+        running Airflow in a distributed manner and aws_conn_id is None or
+        empty, then default boto3 configuration would be used (and must be
+        maintained on each worker node).
+    :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+    :param verify: Whether or not to verify SSL certificates. See:
+        https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
+    :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+        https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+    """
+
+    aws_hook_class = BedrockAgentHook
+    template_fields: Sequence[str] = aws_template_fields(
+        "knowledge_base_id",
+        "data_source_id",
+        "ingest_data_kwargs",
+    )
+
+    def __init__(
+        self,
+        knowledge_base_id: str,
+        data_source_id: str,
+        ingest_data_kwargs: dict[str, Any] | None = None,
+        wait_for_completion: bool = True,
+        waiter_delay: int = 60,
+        waiter_max_attempts: int = 10,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.knowledge_base_id = knowledge_base_id
+        self.data_source_id = data_source_id
+        self.ingest_data_kwargs = ingest_data_kwargs or {}
+
+        self.wait_for_completion = wait_for_completion
+        self.waiter_delay = waiter_delay
+        self.waiter_max_attempts = waiter_max_attempts
+        self.deferrable = deferrable
+
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:
+        event = validate_execute_complete_event(event)
+
+        if event["status"] != "success":
+            raise AirflowException(f"Error while running ingestion job: {event}")
+
+        self.log.info("Bedrock ingestion job `%s` complete.", event["ingestion_job_id"])
+
+        return event["ingestion_job_id"]
+
+    def execute(self, context: Context) -> str:
+        ingestion_job_id = self.hook.conn.start_ingestion_job(
+            knowledgeBaseId=self.knowledge_base_id, dataSourceId=self.data_source_id
+        )["ingestionJob"]["ingestionJobId"]
+
+        if self.deferrable:
+            self.log.info("Deferring for ingestion job.")
+            self.defer(
+                trigger=BedrockIngestionJobTrigger(
+                    knowledge_base_id=self.knowledge_base_id,
+                    data_source_id=self.data_source_id,
+                    ingestion_job_id=ingestion_job_id,
+                    waiter_delay=self.waiter_delay,
+                    waiter_max_attempts=self.waiter_max_attempts,
+                    aws_conn_id=self.aws_conn_id,
+                ),
+                method_name="execute_complete",
+            )
+        if self.wait_for_completion:
+            self.log.info("Waiting for ingestion job %s", ingestion_job_id)
+            self.hook.get_waiter(waiter_name="ingestion_job_complete").wait(
+                knowledgeBaseId=self.knowledge_base_id,
+                dataSourceId=self.data_source_id,
+                ingestionJobId=ingestion_job_id,
+            )
+
+        return ingestion_job_id

{apache_airflow_providers_amazon-8.21.0rc1 → apache_airflow_providers_amazon-8.22.0rc1}/airflow/providers/amazon/aws/operators/s3.py

@@ -30,8 +30,9 @@ from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 from airflow.utils.helpers import exactly_one
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from datetime import datetime
 
+    from airflow.utils.context import Context
 
 BUCKET_DOES_NOT_EXIST_MSG = "Bucket with name: %s doesn't exist"
 
@@ -473,6 +474,10 @@ class S3DeleteObjectsOperator(BaseOperator):
 
     :param prefix: Prefix of objects to delete. (templated)
         All objects matching this prefix in the bucket will be deleted.
+    :param from_datetime: Greater LastModified Date of objects to delete. (templated)
+        All objects which LastModified Date is greater than this datetime in the bucket will be deleted.
+    :param to_datetime: less LastModified Date of objects to delete. (templated)
+        All objects which LastModified Date is less than this datetime in the bucket will be deleted.
     :param aws_conn_id: Connection id of the S3 connection to use
     :param verify: Whether or not to verify SSL certificates for S3 connection.
         By default SSL certificates are verified.
@@ -487,7 +492,7 @@ class S3DeleteObjectsOperator(BaseOperator):
                  CA cert bundle than the one used by botocore.
     """
 
-    template_fields: Sequence[str] = ("keys", "bucket", "prefix")
+    template_fields: Sequence[str] = ("keys", "bucket", "prefix", "from_datetime", "to_datetime")
 
     def __init__(
         self,
@@ -495,6 +500,8 @@ class S3DeleteObjectsOperator(BaseOperator):
         bucket: str,
         keys: str | list | None = None,
         prefix: str | None = None,
+        from_datetime: datetime | None = None,
+        to_datetime: datetime | None = None,
         aws_conn_id: str | None = "aws_default",
         verify: str | bool | None = None,
         **kwargs,
@@ -503,23 +510,36 @@ class S3DeleteObjectsOperator(BaseOperator):
         self.bucket = bucket
         self.keys = keys
         self.prefix = prefix
+        self.from_datetime = from_datetime
+        self.to_datetime = to_datetime
         self.aws_conn_id = aws_conn_id
         self.verify = verify
 
         self._keys: str | list[str] = ""
 
-        if not exactly_one(prefix is None, keys is None):
-            raise AirflowException("Either keys or prefix should be set.")
+        if not exactly_one(keys is None, all(var is None for var in [prefix, from_datetime, to_datetime])):
+            raise AirflowException(
+                "Either keys or at least one of prefix, from_datetime, to_datetime should be set."
+            )
 
     def execute(self, context: Context):
-        if not exactly_one(self.keys is None, self.prefix is None):
-            raise AirflowException("Either keys or prefix should be set.")
+        if not exactly_one(
+            self.keys is None, all(var is None for var in [self.prefix, self.from_datetime, self.to_datetime])
+        ):
+            raise AirflowException(
+                "Either keys or at least one of prefix, from_datetime, to_datetime should be set."
+            )
 
         if isinstance(self.keys, (list, str)) and not self.keys:
             return
         s3_hook = S3Hook(aws_conn_id=self.aws_conn_id, verify=self.verify)
 
-        keys = self.keys or s3_hook.list_keys(bucket_name=self.bucket, prefix=self.prefix)
+        keys = self.keys or s3_hook.list_keys(
+            bucket_name=self.bucket,
+            prefix=self.prefix,
+            from_datetime=self.from_datetime,
+            to_datetime=self.to_datetime,
+        )
         if keys:
             s3_hook.delete_objects(bucket=self.bucket, keys=keys)
             self._keys = keys