apache-airflow-providers-amazon 9.4.0__py3-none-any.whl → 9.5.0rc1__py3-none-any.whl

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

@@ -53,6 +53,9 @@ class S3KeyTrigger(BaseTrigger):
         poke_interval: float = 5.0,
         should_check_fn: bool = False,
         use_regex: bool = False,
+        region_name: str | None = None,
+        verify: bool | str | None = None,
+        botocore_config: dict | None = None,
         **hook_params: Any,
     ):
         super().__init__()
@@ -64,6 +67,9 @@ class S3KeyTrigger(BaseTrigger):
         self.poke_interval = poke_interval
         self.should_check_fn = should_check_fn
         self.use_regex = use_regex
+        self.region_name = region_name
+        self.verify = verify
+        self.botocore_config = botocore_config
 
     def serialize(self) -> tuple[str, dict[str, Any]]:
         """Serialize S3KeyTrigger arguments and classpath."""
@@ -78,17 +84,25 @@ class S3KeyTrigger(BaseTrigger):
                 "poke_interval": self.poke_interval,
                 "should_check_fn": self.should_check_fn,
                 "use_regex": self.use_regex,
+                "region_name": self.region_name,
+                "verify": self.verify,
+                "botocore_config": self.botocore_config,
             },
         )
 
     @cached_property
     def hook(self) -> S3Hook:
-        return S3Hook(aws_conn_id=self.aws_conn_id, verify=self.hook_params.get("verify"))
+        return S3Hook(
+            aws_conn_id=self.aws_conn_id,
+            region_name=self.region_name,
+            verify=self.verify,
+            config=self.botocore_config,
+        )
 
     async def run(self) -> AsyncIterator[TriggerEvent]:
         """Make an asynchronous connection using S3HookAsync."""
         try:
-            async with self.hook.async_conn as client:
+            async with await self.hook.get_async_conn() as client:
                 while True:
                     if await self.hook.check_key_async(
                         client, self.bucket_name, self.bucket_key, self.wildcard_match, self.use_regex
@@ -143,7 +157,9 @@ class S3KeysUnchangedTrigger(BaseTrigger):
         allow_delete: bool = True,
         aws_conn_id: str | None = "aws_default",
         last_activity_time: datetime | None = None,
+        region_name: str | None = None,
         verify: bool | str | None = None,
+        botocore_config: dict | None = None,
         **hook_params: Any,
     ):
         super().__init__()
@@ -160,8 +176,10 @@ class S3KeysUnchangedTrigger(BaseTrigger):
         self.allow_delete = allow_delete
         self.aws_conn_id = aws_conn_id
         self.last_activity_time = last_activity_time
-        self.verify = verify
         self.polling_period_seconds = 0
+        self.region_name = region_name
+        self.verify = verify
+        self.botocore_config = botocore_config
         self.hook_params = hook_params
 
     def serialize(self) -> tuple[str, dict[str, Any]]:
@@ -179,19 +197,26 @@ class S3KeysUnchangedTrigger(BaseTrigger):
                 "aws_conn_id": self.aws_conn_id,
                 "last_activity_time": self.last_activity_time,
                 "hook_params": self.hook_params,
-                "verify": self.verify,
                 "polling_period_seconds": self.polling_period_seconds,
+                "region_name": self.region_name,
+                "verify": self.verify,
+                "botocore_config": self.botocore_config,
             },
         )
 
     @cached_property
     def hook(self) -> S3Hook:
-        return S3Hook(aws_conn_id=self.aws_conn_id, verify=self.hook_params.get("verify"))
+        return S3Hook(
+            aws_conn_id=self.aws_conn_id,
+            region_name=self.region_name,
+            verify=self.verify,
+            config=self.botocore_config,
+        )
 
     async def run(self) -> AsyncIterator[TriggerEvent]:
         """Make an asynchronous connection using S3Hook."""
         try:
-            async with self.hook.async_conn as client:
+            async with await self.hook.get_async_conn() as client:
                 while True:
                     result = await self.hook.is_keys_unchanged_async(
                         client=client,

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

@@ -108,7 +108,7 @@ class SageMakerTrigger(BaseTrigger):
 
     async def run(self):
         self.log.info("job name is %s and job type is %s", self.job_name, self.job_type)
-        async with self.hook.async_conn as client:
+        async with await self.hook.get_async_conn() as client:
             waiter = self.hook.get_waiter(
                 self._get_job_type_waiter(self.job_type), deferrable=True, client=client
             )
@@ -166,7 +166,7 @@ class SageMakerPipelineTrigger(BaseTrigger):
 
     async def run(self) -> AsyncIterator[TriggerEvent]:
         hook = SageMakerHook(aws_conn_id=self.aws_conn_id)
-        async with hook.async_conn as conn:
+        async with await hook.get_async_conn() as conn:
             waiter = hook.get_waiter(self._waiter_name[self.waiter_type], deferrable=True, client=conn)
             for _ in range(self.waiter_max_attempts):
                 try:

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

@@ -0,0 +1,66 @@
+# 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.
+
+"""This module contains the Amazon SageMaker Unified Studio Notebook job trigger."""
+
+from __future__ import annotations
+
+from airflow.triggers.base import BaseTrigger
+
+
+class SageMakerNotebookJobTrigger(BaseTrigger):
+    """
+    Watches for a notebook job, triggers when it finishes.
+
+    Examples:
+     .. code-block:: python
+
+        from airflow.providers.amazon.aws.triggers.sagemaker_unified_studio import SageMakerNotebookJobTrigger
+
+        notebook_trigger = SageMakerNotebookJobTrigger(
+            execution_id="notebook_job_1234",
+            execution_name="notebook_task",
+            waiter_delay=10,
+            waiter_max_attempts=1440,
+        )
+
+    :param execution_id: A unique, meaningful id for the task.
+    :param execution_name: A unique, meaningful name for the task.
+    :param waiter_delay: Interval in seconds to check the notebook execution status.
+    :param waiter_max_attempts: Number of attempts to wait before returning FAILED.
+    """
+
+    def __init__(self, execution_id, execution_name, waiter_delay, waiter_max_attempts, **kwargs):
+        super().__init__(**kwargs)
+        self.execution_id = execution_id
+        self.execution_name = execution_name
+        self.waiter_delay = waiter_delay
+        self.waiter_max_attempts = waiter_max_attempts
+
+    def serialize(self):
+        return (
+            # dynamically generate the fully qualified name of the class
+            self.__class__.__module__ + "." + self.__class__.__qualname__,
+            {
+                "execution_id": self.execution_id,
+                "execution_name": self.execution_name,
+                "poll_interval": self.poll_interval,
+            },
+        )
+
+    async def run(self):
+        pass

airflow/providers/amazon/aws/triggers/sqs.py

@@ -23,14 +23,22 @@ from typing import TYPE_CHECKING, Any
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.sqs import SqsHook
 from airflow.providers.amazon.aws.utils.sqs import process_response
-from airflow.triggers.base import BaseTrigger, TriggerEvent
+from airflow.providers.amazon.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.triggers.base import BaseEventTrigger, TriggerEvent
+else:
+    from airflow.triggers.base import (  # type: ignore
+        BaseTrigger as BaseEventTrigger,
+        TriggerEvent,
+    )
 
 if TYPE_CHECKING:
     from airflow.providers.amazon.aws.hooks.base_aws import BaseAwsConnection
     from airflow.providers.amazon.aws.utils.sqs import MessageFilteringType
 
 
-class SqsSensorTrigger(BaseTrigger):
+class SqsSensorTrigger(BaseEventTrigger):
     """
     Asynchronously get messages from an Amazon SQS queue and then delete the messages from the queue.
 
@@ -176,7 +184,7 @@ class SqsSensorTrigger(BaseTrigger):
         while True:
             # This loop will run indefinitely until the timeout, which is set in the self.defer
             # method, is reached.
-            async with self.hook.async_conn as client:
+            async with await self.hook.get_async_conn() as client:
                 result = await self.poke(client=client)
                 if result:
                     yield TriggerEvent({"status": "success", "message_batch": result})

airflow/providers/amazon/aws/{auth_manager/security_manager/__init__.py → utils/sagemaker_unified_studio.py}

@@ -14,3 +14,15 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+
+"""This module contains utils for the Amazon SageMaker Unified Studio Notebook plugin."""
+
+from __future__ import annotations
+
+import os
+
+workflows_env_key = "WORKFLOWS_ENV"
+
+
+def is_local_runner():
+    return os.getenv(workflows_env_key, "") == "Local"

airflow/providers/amazon/get_provider_info.py

@@ -27,8 +27,9 @@ def get_provider_info():
         "name": "Amazon",
         "description": "Amazon integration (including `Amazon Web Services (AWS) <https://aws.amazon.com/>`__).\n",
         "state": "ready",
-        "source-date-epoch": 1739958314,
+        "source-date-epoch": 1741507721,
         "versions": [
+            "9.5.0",
             "9.4.0",
             "9.2.0",
             "9.1.0",
@@ -275,6 +276,15 @@ def get_provider_info():
                 "how-to-guide": ["/docs/apache-airflow-providers-amazon/operators/sagemaker.rst"],
                 "tags": ["aws"],
             },
+            {
+                "integration-name": "Amazon SageMaker Unified Studio",
+                "external-doc-url": "https://aws.amazon.com/sagemaker/unified-studio/",
+                "logo": "/docs/integration-logos/Amazon-SageMaker_light-bg@4x.png",
+                "how-to-guide": [
+                    "/docs/apache-airflow-providers-amazon/operators/sagemakerunifiedstudio.rst"
+                ],
+                "tags": ["aws"],
+            },
             {
                 "integration-name": "Amazon SecretsManager",
                 "external-doc-url": "https://aws.amazon.com/secrets-manager/",
@@ -491,6 +501,10 @@ def get_provider_info():
                 "integration-name": "Amazon SageMaker",
                 "python-modules": ["airflow.providers.amazon.aws.operators.sagemaker"],
             },
+            {
+                "integration-name": "Amazon SageMaker Unified Studio",
+                "python-modules": ["airflow.providers.amazon.aws.operators.sagemaker_unified_studio"],
+            },
             {
                 "integration-name": "Amazon Simple Notification Service (SNS)",
                 "python-modules": ["airflow.providers.amazon.aws.operators.sns"],
@@ -604,6 +618,10 @@ def get_provider_info():
                 "integration-name": "Amazon Managed Service for Apache Flink",
                 "python-modules": ["airflow.providers.amazon.aws.sensors.kinesis_analytics"],
             },
+            {
+                "integration-name": "Amazon Managed Workflows for Apache Airflow (MWAA)",
+                "python-modules": ["airflow.providers.amazon.aws.sensors.mwaa"],
+            },
             {
                 "integration-name": "Amazon OpenSearch Serverless",
                 "python-modules": ["airflow.providers.amazon.aws.sensors.opensearch_serverless"],
@@ -624,6 +642,10 @@ def get_provider_info():
                 "integration-name": "Amazon SageMaker",
                 "python-modules": ["airflow.providers.amazon.aws.sensors.sagemaker"],
             },
+            {
+                "integration-name": "Amazon SageMaker Unified Studio",
+                "python-modules": ["airflow.providers.amazon.aws.sensors.sagemaker_unified_studio"],
+            },
             {
                 "integration-name": "Amazon Simple Queue Service (SQS)",
                 "python-modules": ["airflow.providers.amazon.aws.sensors.sqs"],
@@ -777,6 +799,10 @@ def get_provider_info():
                 "integration-name": "Amazon SageMaker",
                 "python-modules": ["airflow.providers.amazon.aws.hooks.sagemaker"],
             },
+            {
+                "integration-name": "Amazon SageMaker Unified Studio",
+                "python-modules": ["airflow.providers.amazon.aws.hooks.sagemaker_unified_studio"],
+            },
             {
                 "integration-name": "Amazon Simple Email Service (SES)",
                 "python-modules": ["airflow.providers.amazon.aws.hooks.ses"],
@@ -874,6 +900,10 @@ def get_provider_info():
                 "integration-name": "Amazon SageMaker",
                 "python-modules": ["airflow.providers.amazon.aws.triggers.sagemaker"],
             },
+            {
+                "integration-name": "Amazon SageMaker Unified Studio",
+                "python-modules": ["airflow.providers.amazon.aws.triggers.sagemaker_unified_studio"],
+            },
             {
                 "integration-name": "AWS Glue",
                 "python-modules": [
@@ -1068,6 +1098,7 @@ def get_provider_info():
             "airflow.providers.amazon.aws.links.glue.GlueJobRunDetailsLink",
             "airflow.providers.amazon.aws.links.logs.CloudWatchEventsLink",
             "airflow.providers.amazon.aws.links.sagemaker.SageMakerTransformJobLink",
+            "airflow.providers.amazon.aws.links.sagemaker_unified_studio.SageMakerUnifiedStudioLink",
             "airflow.providers.amazon.aws.links.step_function.StateMachineDetailsLink",
             "airflow.providers.amazon.aws.links.step_function.StateMachineExecutionsDetailsLink",
             "airflow.providers.amazon.aws.links.comprehend.ComprehendPiiEntitiesDetectionLink",
@@ -1350,6 +1381,7 @@ def get_provider_info():
             "PyAthena>=3.0.10",
             "jmespath>=0.7.0",
             "python3-saml>=1.16.0",
+            "sagemaker-studio>=1.0.9",
         ],
         "optional-dependencies": {
             "pandas": ["pandas>=2.1.2,<2.2"],
@@ -1359,6 +1391,7 @@ def get_provider_info():
             "python3-saml": ["python3-saml>=1.16.0"],
             "apache.hive": ["apache-airflow-providers-apache-hive"],
             "exasol": ["apache-airflow-providers-exasol"],
+            "fab": ["apache-airflow-providers-fab"],
             "ftp": ["apache-airflow-providers-ftp"],
             "google": ["apache-airflow-providers-google"],
             "imap": ["apache-airflow-providers-imap"],
@@ -1367,6 +1400,7 @@ def get_provider_info():
             "openlineage": ["apache-airflow-providers-openlineage"],
             "salesforce": ["apache-airflow-providers-salesforce"],
             "ssh": ["apache-airflow-providers-ssh"],
+            "standard": ["apache-airflow-providers-standard"],
         },
         "devel-dependencies": [
             "aiobotocore>=2.13.0",
@@ -1379,5 +1413,6 @@ def get_provider_info():
             "s3fs>=2023.10.0",
             "openapi-schema-validator>=0.6.2",
             "openapi-spec-validator>=0.7.1",
+            "opensearch-py>=2.2.0",
         ],
     }

{apache_airflow_providers_amazon-9.4.0.dist-info → apache_airflow_providers_amazon-9.5.0rc1.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: apache-airflow-providers-amazon
-Version: 9.4.0
+Version: 9.5.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>
@@ -20,9 +20,9 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Topic :: System :: Monitoring
-Requires-Dist: apache-airflow>=2.9.0
-Requires-Dist: apache-airflow-providers-common-compat>=1.3.0
-Requires-Dist: apache-airflow-providers-common-sql>=1.20.0
+Requires-Dist: apache-airflow>=2.9.0rc0
+Requires-Dist: apache-airflow-providers-common-compat>=1.3.0rc0
+Requires-Dist: apache-airflow-providers-common-sql>=1.20.0rc0
 Requires-Dist: apache-airflow-providers-http
 Requires-Dist: boto3>=1.34.90
 Requires-Dist: botocore>=1.34.90
@@ -34,10 +34,12 @@ Requires-Dist: asgiref>=2.3.0
 Requires-Dist: PyAthena>=3.0.10
 Requires-Dist: jmespath>=0.7.0
 Requires-Dist: python3-saml>=1.16.0
+Requires-Dist: sagemaker-studio>=1.0.9
 Requires-Dist: aiobotocore[boto3]>=2.13.0 ; extra == "aiobotocore"
 Requires-Dist: apache-airflow-providers-apache-hive ; extra == "apache-hive"
 Requires-Dist: apache-airflow-providers-cncf-kubernetes>=7.2.0 ; extra == "cncf-kubernetes"
 Requires-Dist: apache-airflow-providers-exasol ; extra == "exasol"
+Requires-Dist: apache-airflow-providers-fab ; extra == "fab"
 Requires-Dist: apache-airflow-providers-ftp ; extra == "ftp"
 Requires-Dist: apache-airflow-providers-google ; extra == "google"
 Requires-Dist: apache-airflow-providers-imap ; extra == "imap"
@@ -49,9 +51,10 @@ Requires-Dist: python3-saml>=1.16.0 ; extra == "python3-saml"
 Requires-Dist: s3fs>=2023.10.0 ; extra == "s3fs"
 Requires-Dist: apache-airflow-providers-salesforce ; extra == "salesforce"
 Requires-Dist: apache-airflow-providers-ssh ; extra == "ssh"
+Requires-Dist: apache-airflow-providers-standard ; extra == "standard"
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
-Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.4.0/changelog.html
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.4.0
+Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.5.0/changelog.html
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.5.0
 Project-URL: Slack Chat, https://s.apache.org/airflow-slack
 Project-URL: Source Code, https://github.com/apache/airflow
 Project-URL: Twitter, https://x.com/ApacheAirflow
@@ -60,6 +63,7 @@ Provides-Extra: aiobotocore
 Provides-Extra: apache-hive
 Provides-Extra: cncf-kubernetes
 Provides-Extra: exasol
+Provides-Extra: fab
 Provides-Extra: ftp
 Provides-Extra: google
 Provides-Extra: imap
@@ -71,34 +75,34 @@ Provides-Extra: python3-saml
 Provides-Extra: s3fs
 Provides-Extra: salesforce
 Provides-Extra: ssh
+Provides-Extra: standard
 
 
- .. 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
+.. 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
+..   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.
+.. 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.
 
- .. NOTE! THIS FILE IS AUTOMATICALLY GENERATED AND WILL BE OVERWRITTEN!
-
- .. IF YOU WANT TO MODIFY TEMPLATE FOR THIS FILE, YOU SHOULD MODIFY THE TEMPLATE
-    `PROVIDER_README_TEMPLATE.rst.jinja2` IN the `dev/breeze/src/airflow_breeze/templates` DIRECTORY
+.. NOTE! THIS FILE IS AUTOMATICALLY GENERATED AND WILL BE OVERWRITTEN!
 
+.. IF YOU WANT TO MODIFY TEMPLATE FOR THIS FILE, YOU SHOULD MODIFY THE TEMPLATE
+   ``PROVIDER_README_TEMPLATE.rst.jinja2`` IN the ``dev/breeze/src/airflow_breeze/templates`` DIRECTORY
 
 Package ``apache-airflow-providers-amazon``
 
-Release: ``9.4.0``
+Release: ``9.5.0``
 
 
 Amazon integration (including `Amazon Web Services (AWS) <https://aws.amazon.com/>`__).
@@ -111,7 +115,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/9.4.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.5.0/>`_.
 
 Installation
 ------------
@@ -142,6 +146,7 @@ PIP package                                 Version required
 ``PyAthena``                                ``>=3.0.10``
 ``jmespath``                                ``>=0.7.0``
 ``python3-saml``                            ``>=1.16.0``
+``sagemaker-studio``                        ``>=1.0.9``
 ==========================================  ======================
 
 Cross provider package dependencies
@@ -177,5 +182,5 @@ Dependent package
 ======================================================================================================================  ===================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.4.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.5.0/changelog.html>`_.