From 2632fa36f4644aa546474fe989166d9bfeb06f86 Mon Sep 17 00:00:00 2001 From: Hussein Awala 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 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