Fixed S3FileTransformOperator to support S3 Select transformation only

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 else f_source.name,
                 key=self.dest_s3_key,
                 replace=self.replace
             )

tests/providers/amazon/aws/operators/test_s3_file_transform.py

@@ -36,6 +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(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)
@@ -123,26 +128,24 @@ 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(self.content, result['Body'].read())
+
     @staticmethod
     def mock_process(mock_popen, return_code=0, process_output=None):
         process = mock_popen.return_value
         process.stdout.readline.side_effect = process_output or []
         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