Repository: incubator-airflow Updated Branches: refs/heads/master fb6229cac -> 9b42fc9a1
[AIRFLOW-2247] Fix RedshiftToS3Transfer not to fail with ValueError RedshiftToS3Transfer calls S3Hook.get_credentials() and tries to take the return value as a tuple with two elements but fails. This is because the return value is ReadOnlyCredentials, which should be handled as a single object or a tuple with tree elements. This PR fixes the above problem. Closes #3158 from sekikn/AIRFLOW-2247 Project: http://git-wip-us.apache.org/repos/asf/incubator-airflow/repo Commit: http://git-wip-us.apache.org/repos/asf/incubator-airflow/commit/9b42fc9a Tree: http://git-wip-us.apache.org/repos/asf/incubator-airflow/tree/9b42fc9a Diff: http://git-wip-us.apache.org/repos/asf/incubator-airflow/diff/9b42fc9a Branch: refs/heads/master Commit: 9b42fc9a1a036a63aeec22a7c7ab310447b376f2 Parents: fb6229c Author: Kengo Seki <[email protected]> Authored: Mon Mar 26 21:26:30 2018 +0200 Committer: Fokko Driesprong <[email protected]> Committed: Mon Mar 26 21:26:30 2018 +0200 ---------------------------------------------------------------------- airflow/operators/redshift_to_s3_operator.py | 13 +-- tests/operators/test_redshift_to_s3_operator.py | 88 ++++++++++++++++++++ 2 files changed, 95 insertions(+), 6 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/9b42fc9a/airflow/operators/redshift_to_s3_operator.py ---------------------------------------------------------------------- diff --git a/airflow/operators/redshift_to_s3_operator.py b/airflow/operators/redshift_to_s3_operator.py index 971ea4a..3e127c1 100644 --- a/airflow/operators/redshift_to_s3_operator.py +++ b/airflow/operators/redshift_to_s3_operator.py @@ -68,7 +68,7 @@ class RedshiftToS3Transfer(BaseOperator): def execute(self, context): self.hook = PostgresHook(postgres_conn_id=self.redshift_conn_id) self.s3 = S3Hook(aws_conn_id=self.aws_conn_id) - a_key, s_key = self.s3.get_credentials() + credentials = self.s3.get_credentials() unload_options = '\n\t\t\t'.join(self.unload_options) self.log.info("Retrieving headers from %s.%s...", self.schema, self.table) @@ -83,10 +83,10 @@ class RedshiftToS3Transfer(BaseOperator): cursor = self.hook.get_conn().cursor() cursor.execute(columns_query) rows = cursor.fetchall() - columns = map(lambda row: row[0], rows) - column_names = ', '.join(map(lambda c: "\\'{0}\\'".format(c), columns)) - column_castings = ', '.join(map(lambda c: "CAST({0} AS text) AS {0}".format(c), - columns)) + columns = [row[0] for row in rows] + column_names = ', '.join("\\'{0}\\'".format(c) for c in columns) + column_castings = ', '.join("CAST({0} AS text) AS {0}".format(c) + for c in columns) unload_query = """ UNLOAD ('SELECT {0} @@ -98,7 +98,8 @@ class RedshiftToS3Transfer(BaseOperator): credentials 'aws_access_key_id={6};aws_secret_access_key={7}' {8}; """.format(column_names, column_castings, self.schema, self.table, - self.s3_bucket, self.s3_key, a_key, s_key, unload_options) + self.s3_bucket, self.s3_key, credentials.access_key, + credentials.secret_key, unload_options) self.log.info('Executing UNLOAD command...') self.hook.run(unload_query, self.autocommit) http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/9b42fc9a/tests/operators/test_redshift_to_s3_operator.py ---------------------------------------------------------------------- diff --git a/tests/operators/test_redshift_to_s3_operator.py b/tests/operators/test_redshift_to_s3_operator.py new file mode 100644 index 0000000..a2aeef0 --- /dev/null +++ b/tests/operators/test_redshift_to_s3_operator.py @@ -0,0 +1,88 @@ +# -*- coding: utf-8 -*- +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import mock +import re +import unittest + +from boto3.session import Session +from airflow.operators.redshift_to_s3_operator import RedshiftToS3Transfer + + +class TestRedshiftToS3Transfer(unittest.TestCase): + + @mock.patch("boto3.session.Session") + @mock.patch("airflow.hooks.postgres_hook.PostgresHook.run") + @mock.patch("airflow.hooks.postgres_hook.PostgresHook.get_conn") + def test_execute(self, mock_get_conn, mock_run, mock_Session): + column_name = "col" + cur = mock.MagicMock() + cur.fetchall.return_value = [(column_name, )] + mock_get_conn.return_value.cursor.return_value = cur + + access_key = "aws_access_key_id" + secret_key = "aws_secret_access_key" + mock_Session.return_value = Session(access_key, secret_key) + + schema = "schema" + table = "table" + s3_bucket = "bucket" + s3_key = "key" + unload_options = "" + + t = RedshiftToS3Transfer( + schema=schema, + table=table, + s3_bucket=s3_bucket, + s3_key=s3_key, + unload_options=unload_options, + redshift_conn_id="redshift_conn_id", + aws_conn_id="aws_conn_id", + task_id="task_id", + dag=None) + t.execute(None) + + columns_query = """ + SELECT column_name + FROM information_schema.columns + WHERE table_schema = '{0}' + AND table_name = '{1}' + ORDER BY ordinal_position + """.format(schema, table) + + unload_query = """ + UNLOAD ('SELECT \\'{0}\\' + UNION ALL + SELECT CAST({0} AS text) AS {0} + FROM {1}.{2} + ORDER BY 1 DESC') + TO 's3://{3}/{4}/{2}_' + with credentials + 'aws_access_key_id={5};aws_secret_access_key={6}' + {7}; + """.format(column_name, schema, table, + s3_bucket, s3_key, access_key, + secret_key, unload_options) + + def _trim(s): + return re.sub("\s+", " ", s.strip()) + + self.assertEqual(_trim(cur.execute.call_args[0][0]), + _trim(columns_query)) + cur.execute.assert_called_once() + + self.assertEqual(_trim(mock_run.call_args[0][0]), + _trim(unload_query)) + mock_run.assert_called_once()
