Lazy match escaped quotes in RedshiftToS3Operator (#40206)

apache · Jun 12, 2024 · 835f28c · 835f28c
1 parent c2a93ea
commit 835f28c
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/airflow/providers/amazon/aws/transfers/redshift_to_s3.py b/airflow/providers/amazon/aws/transfers/redshift_to_s3.py
@@ -128,7 +128,7 @@ def _build_unload_query(
         self, credentials_block: str, select_query: str, s3_key: str, unload_options: str
     ) -> str:
         # Un-escape already escaped queries
-        select_query = re.sub(r"''(.+)''", r"'\1'", select_query)
+        select_query = re.sub(r"''(.+?)''", r"'\1'", select_query)
         return f"""
                     UNLOAD ($${select_query}$$)
                     TO 's3://{self.s3_bucket}/{s3_key}'

diff --git a/tests/providers/amazon/aws/transfers/test_redshift_to_s3.py b/tests/providers/amazon/aws/transfers/test_redshift_to_s3.py
@@ -245,6 +245,12 @@ def test_custom_select_query_unloading(
                 "SELECT 'Single Quotes Break this Operator'",
             ],
             [False, "key", "SELECT ''", "SELECT ''"],
+            [
+                False,
+                "key",
+                "SELECT ''Single Quotes '' || ''Break this Operator''",
+                "SELECT 'Single Quotes ' || 'Break this Operator'",
+            ],
         ],
     )
     @mock.patch("airflow.providers.amazon.aws.hooks.s3.S3Hook.get_connection")