Fix Redshift unload query when select query contains single quotes

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 ('{select_query}')
+                    UNLOAD ({escape}\n{select_query}{escape}\n)
                     TO 's3://{self.s3_bucket}/{s3_key}'
                     credentials
                     '{credentials_block}'

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
 
@@ -395,3 +397,39 @@ 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, warning",
+        [
+            ("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, warning
+    ):
+        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)
+        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