Add Template Fields to RedshiftToS3Operator & S3ToRedshiftOperator (#…

…11844)
apache · Oct 29, 2020 · 8222851 · 8222851
1 parent 9320fe6
commit 8222851
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 5 deletions.
diff --git a/airflow/providers/amazon/aws/transfers/redshift_to_s3.py b/airflow/providers/amazon/aws/transfers/redshift_to_s3.py
@@ -63,7 +63,7 @@ class RedshiftToS3Operator(BaseOperator):
     :type table_as_file_name: bool
     """
 
-    template_fields = ('s3_key',)
+    template_fields = ('s3_bucket', 's3_key', 'schema', 'table', 'unload_options')
     template_ext = ()
     ui_color = '#ededed'
 

diff --git a/airflow/providers/amazon/aws/transfers/s3_to_redshift.py b/airflow/providers/amazon/aws/transfers/s3_to_redshift.py
@@ -60,10 +60,7 @@ class S3ToRedshiftOperator(BaseOperator):
     :type truncate_table: bool
     """
 
-    template_fields = (
-        's3_key',
-        'table',
-    )
+    template_fields = ('s3_bucket', 's3_key', 'schema', 'table', 'copy_options')
     template_ext = ()
     ui_color = '#99e699'
 

diff --git a/tests/providers/amazon/aws/transfers/test_redshift_to_s3.py b/tests/providers/amazon/aws/transfers/test_redshift_to_s3.py
@@ -87,3 +87,8 @@ def test_execute(
 
         assert mock_run.call_count == 1
         assert_equal_ignore_multiple_spaces(self, mock_run.call_args[0][0], unload_query)
+
+    def test_template_fields_overrides(self):
+        self.assertEqual(
+            RedshiftToS3Operator.template_fields, ('s3_bucket', 's3_key', 'schema', 'table', 'unload_options')
+        )
diff --git a/tests/providers/amazon/aws/transfers/test_s3_to_redshift.py b/tests/providers/amazon/aws/transfers/test_s3_to_redshift.py
@@ -109,3 +109,8 @@ def test_truncate(self, mock_run, mock_session):
         assert_equal_ignore_multiple_spaces(self, mock_run.call_args[0][0], transaction)
 
         assert mock_run.call_count == 1
+
+    def test_template_fields_overrides(self):
+        self.assertEqual(
+            S3ToRedshiftOperator.template_fields, ('s3_bucket', 's3_key', 'schema', 'table', 'copy_options')
+        )