Add Amazon Redshift-data to S3<>RS Transfer Operators

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

@@ -17,9 +17,11 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from time import sleep
+from typing import TYPE_CHECKING, Any, Iterable
 
 from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
+from airflow.providers.amazon.aws.utils import trim_none_values
 
 if TYPE_CHECKING:
     from mypy_boto3_redshift_data import RedshiftDataAPIServiceClient  # noqa
@@ -43,3 +45,144 @@ class RedshiftDataHook(AwsGenericHook["RedshiftDataAPIServiceClient"]):
     def __init__(self, *args, **kwargs) -> None:
         kwargs["client_type"] = "redshift-data"
         super().__init__(*args, **kwargs)
+
+    def execute_query(
+        self,
+        database: str,
+        sql: str | list[str],
+        cluster_identifier: str | None = None,
+        db_user: str | None = None,
+        parameters: Iterable | None = None,
+        secret_arn: str | None = None,
+        statement_name: str | None = None,
+        with_event: bool = False,
+        wait_for_completion: bool = True,
+        poll_interval: int = 10,
+    ) -> str:
+        """
+        Execute a statement against Amazon Redshift
+
+        :param database: the name of the database
+        :param sql: the SQL statement or list of  SQL statement to run
+        :param cluster_identifier: unique identifier of a cluster
+        :param db_user: the database username
+        :param parameters: the parameters for the SQL statement
+        :param secret_arn: the name or ARN of the secret that enables db access
+        :param statement_name: the name of the SQL statement
+        :param with_event: indicates whether to send an event to EventBridge
+        :param wait_for_completion: indicates whether to wait for a result, if True wait, if False don't wait
+        :param poll_interval: how often in seconds to check the query status
+
+        :returns statement_id: str, the UUID of the statement
+        """
+        kwargs: dict[str, Any] = {
+            "ClusterIdentifier": cluster_identifier,
+            "Database": database,
+            "DbUser": db_user,
+            "Parameters": parameters,
+            "WithEvent": with_event,
+            "SecretArn": secret_arn,
+            "StatementName": statement_name,
+        }
+        if isinstance(sql, list):
+            kwargs["Sqls"] = sql
+            resp = self.conn.batch_execute_statement(**trim_none_values(kwargs))
+        else:
+            kwargs["Sql"] = sql
+            resp = self.conn.execute_statement(**trim_none_values(kwargs))
+
+        statement_id = resp["Id"]
+
+        if wait_for_completion:
+            self.wait_for_results(statement_id, poll_interval=poll_interval)
+
+        return statement_id
+
+    def wait_for_results(self, statement_id, poll_interval):
+        while True:
+            self.log.info("Polling statement %s", statement_id)
+            resp = self.conn.describe_statement(
+                Id=statement_id,
+            )
+            status = resp["Status"]
+            if status == "FINISHED":
+                return status
+            elif status == "FAILED" or status == "ABORTED":
+                raise ValueError(
+                    f"Statement {statement_id!r} terminated with status {status}, "
+                    f"error msg: {resp.get('Error')}"
+                )
+            else:
+                self.log.info("Query %s", status)
+            sleep(poll_interval)
+
+    def get_table_primary_key(
+        self,
+        table: str,
+        database: str,
+        schema: str | None = "public",
+        cluster_identifier: str | None = None,
+        db_user: str | None = None,
+        secret_arn: str | None = None,
+        statement_name: str | None = None,
+        with_event: bool = False,
+        wait_for_completion: bool = True,
+        poll_interval: int = 10,
+    ) -> list[str] | None:
+        """
+        Helper method that returns the table primary key.
+
+        Copied from ``RedshiftSQLHook.get_table_primary_key()``
+
+        :param table: Name of the target table
+        :param database: the name of the database
+        :param schema: Name of the target schema, public by default
+        :param sql: the SQL statement or list of  SQL statement to run
+        :param cluster_identifier: unique identifier of a cluster
+        :param db_user: the database username
+        :param secret_arn: the name or ARN of the secret that enables db access
+        :param statement_name: the name of the SQL statement
+        :param with_event: indicates whether to send an event to EventBridge
+        :param wait_for_completion: indicates whether to wait for a result, if True wait, if False don't wait
+        :param poll_interval: how often in seconds to check the query status
+
+        :return: Primary key columns list
+        """
+        sql = f"""
+            select kcu.column_name
+            from information_schema.table_constraints tco
+                    join information_schema.key_column_usage kcu
+                        on kcu.constraint_name = tco.constraint_name
+                            and kcu.constraint_schema = tco.constraint_schema
+                            and kcu.constraint_name = tco.constraint_name
+            where tco.constraint_type = 'PRIMARY KEY'
+            and kcu.table_schema = {schema}
+            and kcu.table_name = {table}
+        """
+        stmt_id = self.execute_query(
+            sql=sql,
+            database=database,
+            cluster_identifier=cluster_identifier,
+            db_user=db_user,
+            secret_arn=secret_arn,
+            statement_name=statement_name,
+            with_event=with_event,
+            wait_for_completion=wait_for_completion,
+            poll_interval=poll_interval,
+        )
+        pk_columns = []
+        token = ""
+        while True:
+            kwargs = dict(Id=stmt_id)
+            if token:
+                kwargs["NextToken"] = token
+            response = self.conn.get_statement_result(**kwargs)
+            # we only select a single column (that is a string),
+            # so safe to assume that there is only a single col in the record
+            pk_columns += [y["stringValue"] for x in response["Records"] for y in x]
+            if "NextToken" not in response.keys():
+                break
+            else:
+                token = response["NextToken"]
+
+        return pk_columns or None

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

@@ -17,13 +17,12 @@
 # under the License.
 from __future__ import annotations
 
-from time import sleep
-from typing import TYPE_CHECKING, Any
+import warnings
+from typing import TYPE_CHECKING
 
 from airflow.compat.functools import cached_property
 from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.redshift_data import RedshiftDataHook
-from airflow.providers.amazon.aws.utils import trim_none_values
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -99,67 +98,80 @@ def __init__(
             )
         self.aws_conn_id = aws_conn_id
         self.region = region
-        self.statement_id = None
+        self.statement_id: str | None = None
 
     @cached_property
     def hook(self) -> RedshiftDataHook:
         """Create and return an RedshiftDataHook."""
         return RedshiftDataHook(aws_conn_id=self.aws_conn_id, region_name=self.region)
 
-    def execute_query(self):
-        kwargs: dict[str, Any] = {
-            "ClusterIdentifier": self.cluster_identifier,
-            "Database": self.database,
-            "Sql": self.sql,
-            "DbUser": self.db_user,
-            "Parameters": self.parameters,
-            "WithEvent": self.with_event,
-            "SecretArn": self.secret_arn,
-            "StatementName": self.statement_name,
-        }
-
-        resp = self.hook.conn.execute_statement(**trim_none_values(kwargs))
-        return resp["Id"]
-
-    def execute_batch_query(self):
-        kwargs: dict[str, Any] = {
-            "ClusterIdentifier": self.cluster_identifier,
-            "Database": self.database,
-            "Sqls": self.sql,
-            "DbUser": self.db_user,
-            "Parameters": self.parameters,
-            "WithEvent": self.with_event,
-            "SecretArn": self.secret_arn,
-            "StatementName": self.statement_name,
-        }
-        resp = self.hook.conn.batch_execute_statement(**trim_none_values(kwargs))
-        return resp["Id"]
-
-    def wait_for_results(self, statement_id):
-        while True:
-            self.log.info("Polling statement %s", statement_id)
-            resp = self.hook.conn.describe_statement(
-                Id=statement_id,
-            )
-            status = resp["Status"]
-            if status == "FINISHED":
-                return status
-            elif status == "FAILED" or status == "ABORTED":
-                raise ValueError(f"Statement {statement_id!r} terminated with status {status}.")
-            else:
-                self.log.info("Query %s", status)
-            sleep(self.poll_interval)
-
-    def execute(self, context: Context) -> None:
+    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.await_result,
+            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.await_result,
+            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) -> str:
         """Execute a statement against Amazon Redshift"""
         self.log.info("Executing statement: %s", self.sql)
-        if isinstance(self.sql, list):
-            self.statement_id = self.execute_batch_query()
-        else:
-            self.statement_id = self.execute_query()
 
-        if self.await_result:
-            self.wait_for_results(self.statement_id)
+        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.await_result,
+            poll_interval=self.poll_interval,
+        )
 
         return self.statement_id
 

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

@@ -20,7 +20,9 @@
 
 from typing import TYPE_CHECKING, Iterable, Mapping, Sequence
 
+from airflow.exceptions import AirflowException
 from airflow.models import BaseOperator
+from airflow.providers.amazon.aws.hooks.redshift_data import RedshiftDataHook
 from airflow.providers.amazon.aws.hooks.redshift_sql import RedshiftSQLHook
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 from airflow.providers.amazon.aws.utils.redshift import build_credentials_block
@@ -67,6 +69,9 @@ class RedshiftToS3Operator(BaseOperator):
     :param parameters: (optional) the parameters to render the SQL query with.
     :param table_as_file_name: If set to True, the s3 file will be named as the table.
         Applicable when ``table`` param provided.
+    :param redshift_data_api_kwargs: If using the Redshift Data API instead of the SQL-based connection,
+        dict of arguments for the hook's ``execute_query`` method.
+        Cannot include any of these kwargs: ``{'sql', 'parameters'}``
     """
 
     template_fields: Sequence[str] = (
@@ -98,6 +103,7 @@ def __init__(
         include_header: bool = False,
         parameters: Iterable | Mapping | None = None,
         table_as_file_name: bool = True,  # Set to True by default for not breaking current workflows
+        redshift_data_api_kwargs: dict = {},
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -113,6 +119,7 @@ def __init__(
         self.include_header = include_header
         self.parameters = parameters
         self.table_as_file_name = table_as_file_name
+        self.redshift_data_api_kwargs = redshift_data_api_kwargs
 
         if select_query:
             self.select_query = select_query
@@ -128,6 +135,11 @@ def __init__(
                 "HEADER",
             ]
 
+        if self.redshift_data_api_kwargs:
+            for arg in ["sql", "parameters"]:
+                if arg in self.redshift_data_api_kwargs.keys():
+                    raise AirflowException(f"Cannot include param '{arg}' in Redshift Data API kwargs")
+
     def _build_unload_query(
         self, credentials_block: str, select_query: str, s3_key: str, unload_options: str
     ) -> str:
@@ -140,7 +152,11 @@ def _build_unload_query(
         """
 
     def execute(self, context: Context) -> None:
-        redshift_hook = RedshiftSQLHook(redshift_conn_id=self.redshift_conn_id)
+        redshift_hook: RedshiftDataHook | RedshiftSQLHook
+        if self.redshift_data_api_kwargs:
+            redshift_hook = RedshiftDataHook(aws_conn_id=self.redshift_conn_id)
+        else:
+            redshift_hook = RedshiftSQLHook(redshift_conn_id=self.redshift_conn_id)
         conn = S3Hook.get_connection(conn_id=self.aws_conn_id)
         if conn.extra_dejson.get("role_arn", False):
             credentials_block = f"aws_iam_role={conn.extra_dejson['role_arn']}"
@@ -156,5 +172,10 @@ def execute(self, context: Context) -> None:
         )
 
         self.log.info("Executing UNLOAD command...")
-        redshift_hook.run(unload_query, self.autocommit, parameters=self.parameters)
+        if isinstance(redshift_hook, RedshiftDataHook):
+            redshift_hook.execute_query(
+                sql=unload_query, parameters=self.parameters, **self.redshift_data_api_kwargs
+            )
+        else:
+            redshift_hook.run(unload_query, self.autocommit, parameters=self.parameters)
         self.log.info("UNLOAD command complete...")