apache-airflow-providers-amazon 7.4.1rc1__py3-none-any.whl → 8.0.0rc2__py3-none-any.whl

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

@@ -22,7 +22,10 @@ from typing import TYPE_CHECKING, Any, Sequence
 from airflow.exceptions import AirflowException
 from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.redshift_cluster import RedshiftHook
-from airflow.providers.amazon.aws.triggers.redshift_cluster import RedshiftClusterTrigger
+from airflow.providers.amazon.aws.triggers.redshift_cluster import (
+    RedshiftClusterTrigger,
+    RedshiftCreateClusterTrigger,
+)
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -88,6 +91,7 @@ class RedshiftCreateClusterOperator(BaseOperator):
     :param wait_for_completion: Whether wait for the cluster to be in ``available`` state
     :param max_attempt: The maximum number of attempts to be made. Default: 5
     :param poll_interval: The amount of time in seconds to wait between attempts. Default: 60
+    :param deferrable: If True, the operator will run in deferrable mode
     """
 
     template_fields: Sequence[str] = (
@@ -140,6 +144,7 @@ class RedshiftCreateClusterOperator(BaseOperator):
         wait_for_completion: bool = False,
         max_attempt: int = 5,
         poll_interval: int = 60,
+        deferrable: bool = False,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -180,6 +185,7 @@ class RedshiftCreateClusterOperator(BaseOperator):
         self.wait_for_completion = wait_for_completion
         self.max_attempt = max_attempt
         self.poll_interval = poll_interval
+        self.deferrable = deferrable
         self.kwargs = kwargs
 
     def execute(self, context: Context):
@@ -252,6 +258,16 @@ class RedshiftCreateClusterOperator(BaseOperator):
             self.master_user_password,
             params,
         )
+        if self.deferrable:
+            self.defer(
+                trigger=RedshiftCreateClusterTrigger(
+                    cluster_identifier=self.cluster_identifier,
+                    poll_interval=self.poll_interval,
+                    max_attempt=self.max_attempt,
+                    aws_conn_id=self.aws_conn_id,
+                ),
+                method_name="execute_complete",
+            )
         if self.wait_for_completion:
             redshift_hook.get_conn().get_waiter("cluster_available").wait(
                 ClusterIdentifier=self.cluster_identifier,
@@ -264,6 +280,11 @@ class RedshiftCreateClusterOperator(BaseOperator):
         self.log.info("Created Redshift cluster %s", self.cluster_identifier)
         self.log.info(cluster)
 
+    def execute_complete(self, context, event=None):
+        if event["status"] != "success":
+            raise AirflowException(f"Error creating cluster: {event}")
+        return
+
 
 class RedshiftCreateClusterSnapshotOperator(BaseOperator):
     """

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

@@ -17,7 +17,6 @@
 # under the License.
 from __future__ import annotations
 
-import warnings
 from typing import TYPE_CHECKING
 
 from airflow.compat.functools import cached_property
@@ -83,7 +82,6 @@ class RedshiftDataOperator(BaseOperator):
         return_sql_result: bool = False,
         aws_conn_id: str = "aws_default",
         region: str | None = None,
-        await_result: bool | None = None,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -95,16 +93,7 @@ class RedshiftDataOperator(BaseOperator):
         self.secret_arn = secret_arn
         self.statement_name = statement_name
         self.with_event = with_event
-        self.await_result = await_result
         self.wait_for_completion = wait_for_completion
-        if await_result:
-            warnings.warn(
-                f"Parameter `{self.__class__.__name__}.await_result` is deprecated and will be removed "
-                "in a future release. Please use method `wait_for_completion` instead.",
-                DeprecationWarning,
-                stacklevel=2,
-            )
-            self.wait_for_completion = await_result
         if poll_interval > 0:
             self.poll_interval = poll_interval
         else:
@@ -122,57 +111,6 @@ class RedshiftDataOperator(BaseOperator):
         """Create and return an RedshiftDataHook."""
         return RedshiftDataHook(aws_conn_id=self.aws_conn_id, region_name=self.region)
 
-    def execute_query(self) -> str:
-        warnings.warn(
-            "This method is deprecated and has been moved to the hook "
-            "`airflow.providers.amazon.aws.hooks.redshift_data.RedshiftDataHook`.",
-            DeprecationWarning,
-            stacklevel=2,
-        )
-        self.statement_id = self.hook.execute_query(
-            database=self.database,
-            sql=self.sql,
-            cluster_identifier=self.cluster_identifier,
-            db_user=self.db_user,
-            parameters=self.parameters,
-            secret_arn=self.secret_arn,
-            statement_name=self.statement_name,
-            with_event=self.with_event,
-            wait_for_completion=self.wait_for_completion,
-            poll_interval=self.poll_interval,
-        )
-        return self.statement_id
-
-    def execute_batch_query(self) -> str:
-        warnings.warn(
-            "This method is deprecated and has been moved to the hook "
-            "`airflow.providers.amazon.aws.hooks.redshift_data.RedshiftDataHook`.",
-            DeprecationWarning,
-            stacklevel=2,
-        )
-        self.statement_id = self.hook.execute_query(
-            database=self.database,
-            sql=self.sql,
-            cluster_identifier=self.cluster_identifier,
-            db_user=self.db_user,
-            parameters=self.parameters,
-            secret_arn=self.secret_arn,
-            statement_name=self.statement_name,
-            with_event=self.with_event,
-            wait_for_completion=self.wait_for_completion,
-            poll_interval=self.poll_interval,
-        )
-        return self.statement_id
-
-    def wait_for_results(self, statement_id: str):
-        warnings.warn(
-            "This method is deprecated and has been moved to the hook "
-            "`airflow.providers.amazon.aws.hooks.redshift_data.RedshiftDataHook`.",
-            DeprecationWarning,
-            stacklevel=2,
-        )
-        return self.hook.wait_for_results(statement_id=statement_id, poll_interval=self.poll_interval)
-
     def execute(self, context: Context) -> GetStatementResultResponseTypeDef | str:
         """Execute a statement against Amazon Redshift"""
         self.log.info("Executing statement: %s", self.sql)

airflow/providers/amazon/aws/secrets/secrets_manager.py

@@ -267,23 +267,6 @@ class SecretsManagerBackend(BaseSecretsBackend, LoggingMixin):
         else:
             return secret
 
-    def get_conn_uri(self, conn_id: str) -> str | None:
-        """
-        Return URI representation of Connection conn_id.
-
-        As of Airflow version 2.3.0 this method is deprecated.
-
-        :param conn_id: the connection id
-        :return: deserialized Connection
-        """
-        warnings.warn(
-            f"Method `{self.__class__.__name__}.get_conn_uri` is deprecated and will be removed "
-            "in a future release. Please use method `get_conn_value` instead.",
-            DeprecationWarning,
-            stacklevel=2,
-        )
-        return self.get_conn_value(conn_id)
-
     def get_variable(self, key: str) -> str | None:
         """
         Get Airflow Variable

airflow/providers/amazon/aws/secrets/systems_manager.py

@@ -19,7 +19,6 @@
 from __future__ import annotations
 
 import re
-import warnings
 
 from airflow.compat.functools import cached_property
 from airflow.providers.amazon.aws.utils import trim_none_values
@@ -143,26 +142,6 @@ class SystemsManagerParameterStoreBackend(BaseSecretsBackend, LoggingMixin):
 
         return self._get_secret(self.connections_prefix, conn_id, self.connections_lookup_pattern)
 
-    def get_conn_uri(self, conn_id: str) -> str | None:
-        """
-        Return URI representation of Connection conn_id.
-
-        As of Airflow version 2.3.0 this method is deprecated.
-
-        :param conn_id: the connection id
-        """
-        warnings.warn(
-            f"Method `{self.__class__.__name__}.get_conn_uri` is deprecated and will be removed "
-            "in a future release. Please use method `get_conn_value` instead.",
-            DeprecationWarning,
-            stacklevel=2,
-        )
-        value = self.get_conn_value(conn_id)
-        if value is None:
-            return None
-
-        return self.deserialize_connection(conn_id, value).get_uri()
-
     def get_variable(self, key: str) -> str | None:
         """
         Get Airflow Variable

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

@@ -0,0 +1,97 @@
+# 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 typing import TYPE_CHECKING, Any
+
+from airflow.compat.functools import cached_property
+from airflow.providers.amazon.aws.hooks.dynamodb import DynamoDBHook
+from airflow.sensors.base import BaseSensorOperator
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class DynamoDBValueSensor(BaseSensorOperator):
+    """
+    Waits for an attribute value to be present for an item in a DynamoDB table.
+
+    :param partition_key_name: DynamoDB partition key name
+    :param partition_key_value: DynamoDB partition key value
+    :param attribute_name: DynamoDB attribute name
+    :param attribute_value: DynamoDB attribute value
+    :param sort_key_name: (optional) DynamoDB sort key name
+    :param sort_key_value: (optional) DynamoDB sort key value
+    """
+
+    def __init__(
+        self,
+        table_name: str,
+        partition_key_name: str,
+        partition_key_value: str,
+        attribute_name: str,
+        attribute_value: str,
+        sort_key_name: str | None = None,
+        sort_key_value: str | None = None,
+        aws_conn_id: str | None = DynamoDBHook.default_conn_name,
+        region_name: str | None = None,
+        **kwargs: Any,
+    ):
+        super().__init__(**kwargs)
+        self.table_name = table_name
+        self.partition_key_name = partition_key_name
+        self.partition_key_value = partition_key_value
+        self.attribute_name = attribute_name
+        self.attribute_value = attribute_value
+        self.sort_key_name = sort_key_name
+        self.sort_key_value = sort_key_value
+        self.aws_conn_id = aws_conn_id
+        self.region_name = region_name
+
+    def poke(self, context: Context) -> bool:
+        """Test DynamoDB item for matching attribute value"""
+        key = {self.partition_key_name: self.partition_key_value}
+        msg = (
+            f"Checking table {self.table_name} for "
+            + f"item Partition Key: {self.partition_key_name}={self.partition_key_value}"
+        )
+
+        if self.sort_key_name and self.sort_key_value:
+            key = {self.partition_key_name: self.partition_key_value, self.sort_key_name: self.sort_key_value}
+            msg += f"\nSort Key: {self.sort_key_name}={self.sort_key_value}"
+
+        msg += f"\nattribute: {self.attribute_name}={self.attribute_value}"
+
+        self.log.info(msg)
+        table = self.hook.conn.Table(self.table_name)
+        self.log.info("Table: %s", table)
+        self.log.info("Key: %s", key)
+        response = table.get_item(Key=key)
+        try:
+            self.log.info("Response: %s", response)
+            self.log.info("Want: %s = %s", self.attribute_name, self.attribute_value)
+            self.log.info(
+                "Got: {response['Item'][self.attribute_name]} = %s", response["Item"][self.attribute_name]
+            )
+            return response["Item"][self.attribute_name] == self.attribute_value
+        except KeyError:
+            return False
+
+    @cached_property
+    def hook(self) -> DynamoDBHook:
+        """Create and return a DynamoDBHook"""
+        return DynamoDBHook(self.aws_conn_id, region_name=self.region_name)

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

@@ -25,7 +25,7 @@ from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.emr import EmrContainerHook, EmrHook, EmrServerlessHook
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 from airflow.providers.amazon.aws.links.emr import EmrLogsLink
-from airflow.sensors.base import BaseSensorOperator, poke_mode_only
+from airflow.sensors.base import BaseSensorOperator
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -451,7 +451,6 @@ class EmrJobFlowSensor(EmrBaseSensor):
         return None
 
 
-@poke_mode_only
 class EmrStepSensor(EmrBaseSensor):
     """
     Asks for the state of the step until it reaches any of the target states.

airflow/providers/amazon/aws/transfers/dynamodb_to_s3.py

@@ -87,7 +87,7 @@ class DynamoDBToS3Operator(AwsToAwsBaseOperator):
     :param dynamodb_scan_kwargs: kwargs pass to <https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/dynamodb.html#DynamoDB.Table.scan>
     :param s3_key_prefix: Prefix of s3 object key
     :param process_func: How we transforms a dynamodb item to bytes. By default we dump the json
-    """  # noqa: E501
+    """
 
     template_fields: Sequence[str] = (
         *AwsToAwsBaseOperator.template_fields,

airflow/providers/amazon/aws/transfers/gcs_to_s3.py

@@ -19,7 +19,6 @@
 from __future__ import annotations
 
 import os
-import warnings
 from typing import TYPE_CHECKING, Sequence
 
 from airflow.models import BaseOperator
@@ -45,11 +44,6 @@ class GCSToS3Operator(BaseOperator):
         For e.g to lists the CSV files from in a directory in GCS you would use
         delimiter='.csv'.
     :param gcp_conn_id: (Optional) The connection ID used to connect to Google Cloud.
-    :param google_cloud_storage_conn_id: (Deprecated) The connection ID used to connect to Google Cloud.
-        This parameter has been deprecated. You should pass the gcp_conn_id parameter instead.
-    :param delegate_to: Google account to impersonate using domain-wide delegation of authority,
-        if any. For this to work, the service account making the request must have
-        domain-wide delegation enabled.
     :param dest_aws_conn_id: The destination S3 connection
     :param dest_s3_key: The base S3 key to be used to store the files. (templated)
     :param dest_verify: Whether or not to verify SSL certificates for S3 connection.
@@ -100,8 +94,6 @@ class GCSToS3Operator(BaseOperator):
         prefix: str | None = None,
         delimiter: str | None = None,
         gcp_conn_id: str = "google_cloud_default",
-        google_cloud_storage_conn_id: str | None = None,
-        delegate_to: str | None = None,
         dest_aws_conn_id: str = "aws_default",
         dest_s3_key: str,
         dest_verify: str | bool | None = None,
@@ -114,20 +106,10 @@ class GCSToS3Operator(BaseOperator):
     ) -> None:
         super().__init__(**kwargs)
 
-        if google_cloud_storage_conn_id:
-            warnings.warn(
-                "The google_cloud_storage_conn_id parameter has been deprecated. You should pass "
-                "the gcp_conn_id parameter.",
-                DeprecationWarning,
-                stacklevel=3,
-            )
-            gcp_conn_id = google_cloud_storage_conn_id
-
         self.bucket = bucket
         self.prefix = prefix
         self.delimiter = delimiter
         self.gcp_conn_id = gcp_conn_id
-        self.delegate_to = delegate_to
         self.dest_aws_conn_id = dest_aws_conn_id
         self.dest_s3_key = dest_s3_key
         self.dest_verify = dest_verify
@@ -141,7 +123,6 @@ class GCSToS3Operator(BaseOperator):
         # list all files in an Google Cloud Storage bucket
         hook = GCSHook(
             gcp_conn_id=self.gcp_conn_id,
-            delegate_to=self.delegate_to,
             impersonation_chain=self.google_impersonation_chain,
         )
 

airflow/providers/amazon/aws/transfers/glacier_to_gcs.py

@@ -47,10 +47,7 @@ class GlacierToGCSOperator(BaseOperator):
     :param object_name: the name of the object to check in the Google cloud
         storage bucket.
     :param gzip: option to compress local file or file data for upload
-    :param chunk_size: size of chunk in bytes the that will downloaded from Glacier vault
-    :param delegate_to: The account to impersonate using domain-wide delegation of authority,
-        if any. For this to work, the service account making the request must have
-        domain-wide delegation enabled.
+    :param chunk_size: size of chunk in bytes the that will be downloaded from Glacier vault
     :param google_impersonation_chain: Optional Google service account to impersonate using
         short-term credentials, or chained list of accounts required to get the access_token
         of the last account in the list, which will be impersonated in the request.
@@ -73,7 +70,6 @@ class GlacierToGCSOperator(BaseOperator):
         object_name: str,
         gzip: bool,
         chunk_size: int = 1024,
-        delegate_to: str | None = None,
         google_impersonation_chain: str | Sequence[str] | None = None,
         **kwargs,
     ) -> None:
@@ -85,14 +81,12 @@ class GlacierToGCSOperator(BaseOperator):
         self.object_name = object_name
         self.gzip = gzip
         self.chunk_size = chunk_size
-        self.delegate_to = delegate_to
         self.impersonation_chain = google_impersonation_chain
 
     def execute(self, context: Context) -> str:
         glacier_hook = GlacierHook(aws_conn_id=self.aws_conn_id)
         gcs_hook = GCSHook(
             gcp_conn_id=self.gcp_conn_id,
-            delegate_to=self.delegate_to,
             impersonation_chain=self.impersonation_chain,
         )
         job_id = glacier_hook.retrieve_inventory(vault_name=self.vault_name)

airflow/providers/amazon/aws/transfers/google_api_to_s3.py

@@ -22,12 +22,14 @@ import json
 import sys
 from typing import TYPE_CHECKING, Sequence
 
-from airflow.models import BaseOperator, TaskInstance
+from airflow.models import BaseOperator
 from airflow.models.xcom import MAX_XCOM_SIZE, XCOM_RETURN_KEY
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 from airflow.providers.google.common.hooks.discovery_api import GoogleDiscoveryApiHook
 
 if TYPE_CHECKING:
+    from airflow.models import TaskInstance
+    from airflow.serialization.pydantic.taskinstance import TaskInstancePydantic
     from airflow.utils.context import Context
 
 
@@ -71,13 +73,10 @@ class GoogleApiToS3Operator(BaseOperator):
 
         .. note:: This means the response will be a list of responses.
 
-    :param google_api_num_retries: Define the number of retries for the google api requests being made
+    :param google_api_num_retries: Define the number of retries for the Google API requests being made
         if it fails.
     :param s3_overwrite: Specifies whether the s3 file will be overwritten if exists.
     :param gcp_conn_id: The connection ID to use when fetching connection info.
-    :param delegate_to: Google account to impersonate using domain-wide delegation of authority,
-        if any. For this to work, the service account making the request must have
-        domain-wide delegation enabled.
     :param aws_conn_id: The connection id specifying the authentication information for the S3 Bucket.
     :param google_impersonation_chain: Optional Google service account to impersonate using
         short-term credentials, or chained list of accounts required to get the access_token
@@ -113,7 +112,6 @@ class GoogleApiToS3Operator(BaseOperator):
         google_api_num_retries: int = 0,
         s3_overwrite: bool = False,
         gcp_conn_id: str = "google_cloud_default",
-        delegate_to: str | None = None,
         aws_conn_id: str = "aws_default",
         google_impersonation_chain: str | Sequence[str] | None = None,
         **kwargs,
@@ -131,7 +129,6 @@ class GoogleApiToS3Operator(BaseOperator):
         self.google_api_num_retries = google_api_num_retries
         self.s3_overwrite = s3_overwrite
         self.gcp_conn_id = gcp_conn_id
-        self.delegate_to = delegate_to
         self.aws_conn_id = aws_conn_id
         self.google_impersonation_chain = google_impersonation_chain
 
@@ -156,7 +153,6 @@ class GoogleApiToS3Operator(BaseOperator):
     def _retrieve_data_from_google_api(self) -> dict:
         google_discovery_api_hook = GoogleDiscoveryApiHook(
             gcp_conn_id=self.gcp_conn_id,
-            delegate_to=self.delegate_to,
             api_service_name=self.google_api_service_name,
             api_version=self.google_api_service_version,
             impersonation_chain=self.google_impersonation_chain,
@@ -177,7 +173,9 @@ class GoogleApiToS3Operator(BaseOperator):
             replace=self.s3_overwrite,
         )
 
-    def _update_google_api_endpoint_params_via_xcom(self, task_instance: TaskInstance) -> None:
+    def _update_google_api_endpoint_params_via_xcom(
+        self, task_instance: TaskInstance | TaskInstancePydantic
+    ) -> None:
 
         if self.google_api_endpoint_params_via_xcom:
             google_api_endpoint_params = task_instance.xcom_pull(
@@ -186,7 +184,9 @@ class GoogleApiToS3Operator(BaseOperator):
             )
             self.google_api_endpoint_params.update(google_api_endpoint_params)
 
-    def _expose_google_api_response_via_xcom(self, task_instance: TaskInstance, data: dict) -> None:
+    def _expose_google_api_response_via_xcom(
+        self, task_instance: TaskInstance | TaskInstancePydantic, data: dict
+    ) -> None:
         if sys.getsizeof(data) < MAX_XCOM_SIZE:
             task_instance.xcom_push(key=self.google_api_response_via_xcom or XCOM_RETURN_KEY, value=data)
         else:

airflow/providers/amazon/aws/transfers/imap_attachment_to_s3.py

@@ -18,7 +18,6 @@
 """This module allows you to transfer mail attachments from a mail server into s3 bucket."""
 from __future__ import annotations
 
-import warnings
 from typing import TYPE_CHECKING, Sequence
 
 from airflow.models import BaseOperator
@@ -28,10 +27,6 @@ from airflow.providers.imap.hooks.imap import ImapHook
 if TYPE_CHECKING:
     from airflow.utils.context import Context
 
-_DEPRECATION_MSG = (
-    "The s3_conn_id parameter has been deprecated. You should pass instead the aws_conn_id parameter."
-)
-
 
 class ImapAttachmentToS3Operator(BaseOperator):
     """
@@ -66,15 +61,10 @@ class ImapAttachmentToS3Operator(BaseOperator):
         imap_mail_filter: str = "All",
         s3_overwrite: bool = False,
         imap_conn_id: str = "imap_default",
-        s3_conn_id: str | None = None,
         aws_conn_id: str = "aws_default",
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
-        if s3_conn_id:
-            warnings.warn(_DEPRECATION_MSG, DeprecationWarning, stacklevel=3)
-            aws_conn_id = s3_conn_id
-
         self.imap_attachment_name = imap_attachment_name
         self.s3_bucket = s3_bucket
         self.s3_key = s3_key

airflow/providers/amazon/aws/transfers/mongo_to_s3.py

@@ -18,7 +18,6 @@
 from __future__ import annotations
 
 import json
-import warnings
 from typing import TYPE_CHECKING, Any, Iterable, Sequence, cast
 
 from bson import json_util
@@ -31,11 +30,6 @@ if TYPE_CHECKING:
     from airflow.utils.context import Context
 
 
-_DEPRECATION_MSG = (
-    "The s3_conn_id parameter has been deprecated. You should pass instead the aws_conn_id parameter."
-)
-
-
 class MongoToS3Operator(BaseOperator):
     """Operator meant to move data from mongo via pymongo to s3 via boto.
 
@@ -66,7 +60,6 @@ class MongoToS3Operator(BaseOperator):
     def __init__(
         self,
         *,
-        s3_conn_id: str | None = None,
         mongo_conn_id: str = "mongo_default",
         aws_conn_id: str = "aws_default",
         mongo_collection: str,
@@ -81,10 +74,6 @@ class MongoToS3Operator(BaseOperator):
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
-        if s3_conn_id:
-            warnings.warn(_DEPRECATION_MSG, DeprecationWarning, stacklevel=3)
-            aws_conn_id = s3_conn_id
-
         self.mongo_conn_id = mongo_conn_id
         self.aws_conn_id = aws_conn_id
         self.mongo_db = mongo_db

airflow/providers/amazon/aws/transfers/s3_to_sftp.py

@@ -17,7 +17,6 @@
 # under the License.
 from __future__ import annotations
 
-import warnings
 from tempfile import NamedTemporaryFile
 from typing import TYPE_CHECKING, Sequence
 from urllib.parse import urlsplit
@@ -29,10 +28,6 @@ from airflow.providers.ssh.hooks.ssh import SSHHook
 if TYPE_CHECKING:
     from airflow.utils.context import Context
 
-_DEPRECATION_MSG = (
-    "The s3_conn_id parameter has been deprecated. You should pass instead the aws_conn_id parameter."
-)
-
 
 class S3ToSFTPOperator(BaseOperator):
     """
@@ -62,15 +57,10 @@ class S3ToSFTPOperator(BaseOperator):
         s3_key: str,
         sftp_path: str,
         sftp_conn_id: str = "ssh_default",
-        s3_conn_id: str | None = None,
         aws_conn_id: str = "aws_default",
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
-        if s3_conn_id:
-            warnings.warn(_DEPRECATION_MSG, DeprecationWarning, stacklevel=3)
-            aws_conn_id = s3_conn_id
-
         self.sftp_conn_id = sftp_conn_id
         self.sftp_path = sftp_path
         self.s3_bucket = s3_bucket

airflow/providers/amazon/aws/transfers/sql_to_s3.py

@@ -80,6 +80,7 @@ class SqlToS3Operator(BaseOperator):
                 CA cert bundle than the one used by botocore.
     :param file_format: the destination file format, only string 'csv', 'json' or 'parquet' is accepted.
     :param pd_kwargs: arguments to include in DataFrame ``.to_parquet()``, ``.to_json()`` or ``.to_csv()``.
+    :param groupby_kwargs: argument to include in DataFrame ``groupby()``.
     """
 
     template_fields: Sequence[str] = (
@@ -107,6 +108,7 @@ class SqlToS3Operator(BaseOperator):
         verify: bool | str | None = None,
         file_format: Literal["csv", "json", "parquet"] = "csv",
         pd_kwargs: dict | None = None,
+        groupby_kwargs: dict | None = None,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -119,6 +121,7 @@ class SqlToS3Operator(BaseOperator):
         self.replace = replace
         self.pd_kwargs = pd_kwargs or {}
         self.parameters = parameters
+        self.groupby_kwargs = groupby_kwargs or {}
 
         if "path_or_buf" in self.pd_kwargs:
             raise AirflowException("The argument path_or_buf is not allowed, please remove it")
@@ -170,15 +173,26 @@ class SqlToS3Operator(BaseOperator):
         self._fix_dtypes(data_df, self.file_format)
         file_options = FILE_OPTIONS_MAP[self.file_format]
 
-        with NamedTemporaryFile(mode=file_options.mode, suffix=file_options.suffix) as tmp_file:
-
-            self.log.info("Writing data to temp file")
-            getattr(data_df, file_options.function)(tmp_file.name, **self.pd_kwargs)
-
-            self.log.info("Uploading data to S3")
-            s3_conn.load_file(
-                filename=tmp_file.name, key=self.s3_key, bucket_name=self.s3_bucket, replace=self.replace
-            )
+        for group_name, df in self._partition_dataframe(df=data_df):
+            with NamedTemporaryFile(mode=file_options.mode, suffix=file_options.suffix) as tmp_file:
+
+                self.log.info("Writing data to temp file")
+                getattr(df, file_options.function)(tmp_file.name, **self.pd_kwargs)
+
+                self.log.info("Uploading data to S3")
+                object_key = f"{self.s3_key}_{group_name}" if group_name else self.s3_key
+                s3_conn.load_file(
+                    filename=tmp_file.name, key=object_key, bucket_name=self.s3_bucket, replace=self.replace
+                )
+
+    def _partition_dataframe(self, df: DataFrame) -> Iterable[tuple[str, DataFrame]]:
+        """Partition dataframe using pandas groupby() method"""
+        if not self.groupby_kwargs:
+            yield "", df
+        else:
+            grouped_df = df.groupby(**self.groupby_kwargs)
+            for group_label in grouped_df.groups.keys():
+                yield group_label, grouped_df.get_group(group_label).reset_index(drop=True)
 
     def _get_hook(self) -> DbApiHook:
         self.log.debug("Get connection for %s", self.sql_conn_id)