From ed8cba20bb8c4297b790d64a8e50c9d37d3e3ed1 Mon Sep 17 00:00:00 2001 From: Mariusz Strzelecki Date: Thu, 4 Jun 2020 08:43:07 +0200 Subject: [PATCH 1/5] Fixed S3FileTransformOperator to support S3 Select transformation only Documentation for S3FileTransformOperator states that users can skip transformation script if S3 Select experession is specified, but in this case the created file is always zero bytes long. This fix changes the behaviour, so in case of no transformation given, the source file (a result of S3Select) is uploaded. --- .../amazon/aws/operators/s3_file_transform.py | 2 +- .../aws/operators/test_s3_file_transform.py | 24 ++++++++++--------- 2 files changed, 14 insertions(+), 12 deletions(-) diff --git a/airflow/providers/amazon/aws/operators/s3_file_transform.py b/airflow/providers/amazon/aws/operators/s3_file_transform.py index 39788918a8bdc..1141e9b25ed50 100644 --- a/airflow/providers/amazon/aws/operators/s3_file_transform.py +++ b/airflow/providers/amazon/aws/operators/s3_file_transform.py @@ -167,7 +167,7 @@ def execute(self, context): self.log.info("Uploading transformed file to S3") f_dest.flush() dest_s3.load_file( - filename=f_dest.name, + filename=f_dest.name if self.transform_script is not None else f_source.name, key=self.dest_s3_key, replace=self.replace ) diff --git a/tests/providers/amazon/aws/operators/test_s3_file_transform.py b/tests/providers/amazon/aws/operators/test_s3_file_transform.py index 602eef1ab3e57..76fd7df0a5266 100644 --- a/tests/providers/amazon/aws/operators/test_s3_file_transform.py +++ b/tests/providers/amazon/aws/operators/test_s3_file_transform.py @@ -36,6 +36,10 @@ class TestS3FileTransformOperator(unittest.TestCase): def setUp(self): + self.bucket = "bucket" + self.input_key = "foo" + self.output_key = "bar" + self.bio = io.BytesIO(b"input") self.tmp_dir = mkdtemp(prefix='test_tmpS3FileTransform_') self.transform_script = os.path.join(self.tmp_dir, "transform.py") os.mknod(self.transform_script) @@ -123,6 +127,10 @@ def test_execute_with_select_expression(self, mock_select_key): expression=select_expression ) + conn = boto3.client('s3') + result = conn.get_object(Bucket=self.bucket, Key=self.output_key) + self.assertEqual(b'input', result['Body'].read()) + @staticmethod def mock_process(mock_popen, return_code=0, process_output=None): process = mock_popen.return_value @@ -130,19 +138,13 @@ def mock_process(mock_popen, return_code=0, process_output=None): process.wait.return_value = None process.returncode = return_code - @staticmethod - def s3_paths(): - bucket = "bucket" - input_key = "foo" - output_key = "bar" - bio = io.BytesIO(b"input") - + def s3_paths(self): conn = boto3.client('s3') - conn.create_bucket(Bucket=bucket) - conn.upload_fileobj(Bucket=bucket, Key=input_key, Fileobj=bio) + conn.create_bucket(Bucket=self.bucket) + conn.upload_fileobj(Bucket=self.bucket, Key=self.input_key, Fileobj=self.bio) s3_url = "s3://{0}/{1}" - input_path = s3_url.format(bucket, input_key) - output_path = s3_url.format(bucket, output_key) + input_path = s3_url.format(self.bucket, self.input_key) + output_path = s3_url.format(self.bucket, self.output_key) return input_path, output_path From 00b2e4abf5d0fc5d39d396a7107cd94c6c3434c2 Mon Sep 17 00:00:00 2001 From: Mariusz Strzelecki Date: Wed, 15 Jul 2020 12:07:11 +0200 Subject: [PATCH 2/5] Update airflow/providers/amazon/aws/operators/s3_file_transform.py Co-authored-by: Felix Uellendall --- airflow/providers/amazon/aws/operators/s3_file_transform.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/amazon/aws/operators/s3_file_transform.py b/airflow/providers/amazon/aws/operators/s3_file_transform.py index 1141e9b25ed50..4f0e626fc37e8 100644 --- a/airflow/providers/amazon/aws/operators/s3_file_transform.py +++ b/airflow/providers/amazon/aws/operators/s3_file_transform.py @@ -167,7 +167,7 @@ def execute(self, context): self.log.info("Uploading transformed file to S3") f_dest.flush() dest_s3.load_file( - filename=f_dest.name if self.transform_script is not None else f_source.name, + filename=f_dest.name if self.transform_script else f_source.name, key=self.dest_s3_key, replace=self.replace ) From 8a0ead001bc58bbda95033f4405b5adf77f2da1d Mon Sep 17 00:00:00 2001 From: Mariusz Strzelecki Date: Wed, 15 Jul 2020 12:09:40 +0200 Subject: [PATCH 3/5] Assertion simplified --- tests/providers/amazon/aws/operators/test_s3_file_transform.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/providers/amazon/aws/operators/test_s3_file_transform.py b/tests/providers/amazon/aws/operators/test_s3_file_transform.py index 76fd7df0a5266..091ded02d4b55 100644 --- a/tests/providers/amazon/aws/operators/test_s3_file_transform.py +++ b/tests/providers/amazon/aws/operators/test_s3_file_transform.py @@ -129,7 +129,7 @@ def test_execute_with_select_expression(self, mock_select_key): conn = boto3.client('s3') result = conn.get_object(Bucket=self.bucket, Key=self.output_key) - self.assertEqual(b'input', result['Body'].read()) + assert self.bio.getvalue() == result['Body'] @staticmethod def mock_process(mock_popen, return_code=0, process_output=None): From 2b72850218d32ee8e95b2e2a2797a7de8835440e Mon Sep 17 00:00:00 2001 From: Mariusz Strzelecki Date: Wed, 15 Jul 2020 12:11:49 +0200 Subject: [PATCH 4/5] Assertion fix --- tests/providers/amazon/aws/operators/test_s3_file_transform.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/providers/amazon/aws/operators/test_s3_file_transform.py b/tests/providers/amazon/aws/operators/test_s3_file_transform.py index 091ded02d4b55..9c453f4953c2e 100644 --- a/tests/providers/amazon/aws/operators/test_s3_file_transform.py +++ b/tests/providers/amazon/aws/operators/test_s3_file_transform.py @@ -129,7 +129,7 @@ def test_execute_with_select_expression(self, mock_select_key): conn = boto3.client('s3') result = conn.get_object(Bucket=self.bucket, Key=self.output_key) - assert self.bio.getvalue() == result['Body'] + assert self.bio.getvalue() == result['Body'].read() @staticmethod def mock_process(mock_popen, return_code=0, process_output=None): From 5d0f04f45e272dd620a61895288abc7d79ffc742 Mon Sep 17 00:00:00 2001 From: Mariusz Strzelecki Date: Thu, 16 Jul 2020 08:24:44 +0200 Subject: [PATCH 5/5] Fix comparision --- .../providers/amazon/aws/operators/test_s3_file_transform.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/providers/amazon/aws/operators/test_s3_file_transform.py b/tests/providers/amazon/aws/operators/test_s3_file_transform.py index 9c453f4953c2e..a1bf7beeb5ed0 100644 --- a/tests/providers/amazon/aws/operators/test_s3_file_transform.py +++ b/tests/providers/amazon/aws/operators/test_s3_file_transform.py @@ -36,10 +36,11 @@ class TestS3FileTransformOperator(unittest.TestCase): def setUp(self): + self.content = b"input" self.bucket = "bucket" self.input_key = "foo" self.output_key = "bar" - self.bio = io.BytesIO(b"input") + self.bio = io.BytesIO(self.content) self.tmp_dir = mkdtemp(prefix='test_tmpS3FileTransform_') self.transform_script = os.path.join(self.tmp_dir, "transform.py") os.mknod(self.transform_script) @@ -129,7 +130,7 @@ def test_execute_with_select_expression(self, mock_select_key): conn = boto3.client('s3') result = conn.get_object(Bucket=self.bucket, Key=self.output_key) - assert self.bio.getvalue() == result['Body'].read() + self.assertEqual(self.content, result['Body'].read()) @staticmethod def mock_process(mock_popen, return_code=0, process_output=None):