From 2632fa36f4644aa546474fe989166d9bfeb06f86 Mon Sep 17 00:00:00 2001
From: Hussein Awala <houssein.awala.96@gmail.com>
Date: Sat, 25 Mar 2023 14:35:37 +0100
Subject: [PATCH 1/2] Fix redshift unload query when select query contains
 single quotes

---
 .../amazon/aws/transfers/redshift_to_s3.py    |  2 +-
 .../aws/transfers/test_redshift_to_s3.py      | 23 +++++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/airflow/providers/amazon/aws/transfers/redshift_to_s3.py b/airflow/providers/amazon/aws/transfers/redshift_to_s3.py
index 45f1e1783c49d..902ca1e562ff8 100644
--- a/airflow/providers/amazon/aws/transfers/redshift_to_s3.py
+++ b/airflow/providers/amazon/aws/transfers/redshift_to_s3.py
@@ -144,7 +144,7 @@ def _build_unload_query(
         self, credentials_block: str, select_query: str, s3_key: str, unload_options: str
     ) -> str:
         return f"""
-                    UNLOAD ('{select_query}')
+                    UNLOAD ($$\n{select_query}$$\n)
                     TO 's3://{self.s3_bucket}/{s3_key}'
                     credentials
                     '{credentials_block}'
diff --git a/tests/providers/amazon/aws/transfers/test_redshift_to_s3.py b/tests/providers/amazon/aws/transfers/test_redshift_to_s3.py
index 67982821e8e10..b27d38ea0e135 100644
--- a/tests/providers/amazon/aws/transfers/test_redshift_to_s3.py
+++ b/tests/providers/amazon/aws/transfers/test_redshift_to_s3.py
@@ -395,3 +395,26 @@ def test_table_unloading_using_redshift_data_api(
         )
         # test sql arg
         assert_equal_ignore_multiple_spaces(self, mock_rs.execute_statement.call_args[1]["Sql"], unload_query)
+
+    @pytest.mark.parametrize(
+        "select_query, expected_unload_query",
+        [
+            ("SELECT * FROM schema.table", "UNLOAD ($$\nSELECT * FROM schema.table$$\n)"),
+            (
+                "SELECT * FROM schema.table WHERE foo='boo'",
+                "UNLOAD ($$\nSELECT * FROM schema.table WHERE foo='boo'$$\n)",
+            ),
+        ],
+    )
+    @mock.patch("boto3.session.Session")
+    def test_build_unload_query_single_quotes(self, mock_session, select_query, expected_unload_query):
+        op = RedshiftToS3Operator(
+            s3_bucket="s3_bucket",
+            s3_key="s3_key",
+            select_query="select_query",
+            task_id="task_id",
+            dag=None,
+        )
+        credentials_block = build_credentials_block(mock_session.return_value)
+        unload_query = op._build_unload_query(credentials_block, select_query, "s3_key", "HEADER")
+        assert expected_unload_query in unload_query

From 9a17ac58fb819e3794aad5aed9770f2a1ad65894 Mon Sep 17 00:00:00 2001
From: Hussein Awala <houssein.awala.96@gmail.com>
Date: Fri, 14 Apr 2023 02:16:26 +0200
Subject: [PATCH 2/2] Check if the select query contains double quotes to skip
 adding escape char

---
 .../amazon/aws/transfers/redshift_to_s3.py    | 12 +++++++++-
 .../aws/transfers/test_redshift_to_s3.py      | 23 +++++++++++++++----
 2 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/airflow/providers/amazon/aws/transfers/redshift_to_s3.py b/airflow/providers/amazon/aws/transfers/redshift_to_s3.py
index 902ca1e562ff8..11fdb429e1d32 100644
--- a/airflow/providers/amazon/aws/transfers/redshift_to_s3.py
+++ b/airflow/providers/amazon/aws/transfers/redshift_to_s3.py
@@ -18,6 +18,7 @@
 """Transfers data from AWS Redshift into a S3 Bucket."""
 from __future__ import annotations
 
+import logging
 from typing import TYPE_CHECKING, Iterable, Mapping, Sequence
 
 from airflow.exceptions import AirflowException
@@ -143,8 +144,17 @@ def __init__(
     def _build_unload_query(
         self, credentials_block: str, select_query: str, s3_key: str, unload_options: str
     ) -> str:
+        # TODO: remove in provider 8.0.0
+        if "''" in select_query:
+            logging.warning(
+                "Skip adding escape characters for select query. This check will be removed soon, "
+                "consider deleting the '' and let Airflow add the escape characters"
+            )
+            escape = ""
+        else:
+            escape = "$$"
         return f"""
-                    UNLOAD ($$\n{select_query}$$\n)
+                    UNLOAD ({escape}\n{select_query}{escape}\n)
                     TO 's3://{self.s3_bucket}/{s3_key}'
                     credentials
                     '{credentials_block}'
diff --git a/tests/providers/amazon/aws/transfers/test_redshift_to_s3.py b/tests/providers/amazon/aws/transfers/test_redshift_to_s3.py
index b27d38ea0e135..9d04e3340aca1 100644
--- a/tests/providers/amazon/aws/transfers/test_redshift_to_s3.py
+++ b/tests/providers/amazon/aws/transfers/test_redshift_to_s3.py
@@ -17,6 +17,8 @@
 # under the License.
 from __future__ import annotations
 
+import io
+from contextlib import redirect_stdout
 from copy import deepcopy
 from unittest import mock
 
@@ -397,17 +399,25 @@ def test_table_unloading_using_redshift_data_api(
         assert_equal_ignore_multiple_spaces(self, mock_rs.execute_statement.call_args[1]["Sql"], unload_query)
 
     @pytest.mark.parametrize(
-        "select_query, expected_unload_query",
+        "select_query, expected_unload_query, warning",
         [
-            ("SELECT * FROM schema.table", "UNLOAD ($$\nSELECT * FROM schema.table$$\n)"),
+            ("SELECT * FROM schema.table", "UNLOAD ($$\nSELECT * FROM schema.table$$\n)", False),
             (
                 "SELECT * FROM schema.table WHERE foo='boo'",
                 "UNLOAD ($$\nSELECT * FROM schema.table WHERE foo='boo'$$\n)",
+                False,
+            ),
+            (
+                "SELECT * FROM schema.table WHERE foo=''boo''",
+                "UNLOAD (\nSELECT * FROM schema.table WHERE foo=''boo''\n)",
+                True,
             ),
         ],
     )
     @mock.patch("boto3.session.Session")
-    def test_build_unload_query_single_quotes(self, mock_session, select_query, expected_unload_query):
+    def test_build_unload_query_single_quotes(
+        self, mock_session, select_query, expected_unload_query, warning
+    ):
         op = RedshiftToS3Operator(
             s3_bucket="s3_bucket",
             s3_key="s3_key",
@@ -416,5 +426,10 @@ def test_build_unload_query_single_quotes(self, mock_session, select_query, expe
             dag=None,
         )
         credentials_block = build_credentials_block(mock_session.return_value)
-        unload_query = op._build_unload_query(credentials_block, select_query, "s3_key", "HEADER")
+        with redirect_stdout(io.StringIO()) as stdout:
+            unload_query = op._build_unload_query(credentials_block, select_query, "s3_key", "HEADER")
+            stdout = stdout.getvalue()
         assert expected_unload_query in unload_query
+        warning_msg = "Skip adding escape characters for select query. This check will be removed soon, "
+        "consider deleting the '' and let Airflow add the escape characters"
+        assert (warning_msg in stdout) == warning
