From 8830a1daaf6b70b77cc50adaf4960961f6597375 Mon Sep 17 00:00:00 2001 From: Oleksandr Yarushevskyi Date: Mon, 19 Jul 2021 19:46:29 +0300 Subject: [PATCH] Fixed template_fields_renderers for Amazon provider --- airflow/providers/amazon/aws/operators/batch.py | 2 +- airflow/providers/amazon/aws/operators/datasync.py | 10 +++++----- airflow/providers/amazon/aws/operators/ecs.py | 7 ++++++- .../providers/amazon/aws/operators/emr_add_steps.py | 1 + airflow/providers/amazon/aws/operators/glue.py | 5 ++++- .../amazon/aws/operators/s3_bucket_tagging.py | 1 + airflow/providers/amazon/aws/operators/sns.py | 2 +- airflow/providers/amazon/aws/operators/sqs.py | 3 ++- airflow/providers/amazon/aws/transfers/mongo_to_s3.py | 2 +- airflow/providers/amazon/aws/transfers/mysql_to_s3.py | 2 +- .../providers/amazon/aws/transfers/redshift_to_s3.py | 3 ++- 11 files changed, 25 insertions(+), 13 deletions(-) diff --git a/airflow/providers/amazon/aws/operators/batch.py b/airflow/providers/amazon/aws/operators/batch.py index 46c551530b5c4..6fb5bedd2e2da 100644 --- a/airflow/providers/amazon/aws/operators/batch.py +++ b/airflow/providers/amazon/aws/operators/batch.py @@ -99,7 +99,7 @@ class AwsBatchOperator(BaseOperator): "overrides", "parameters", ) - template_fields_renderers = {"overrides": "py", "parameters": "py"} + template_fields_renderers = {"overrides": "json", "parameters": "json"} def __init__( self, diff --git a/airflow/providers/amazon/aws/operators/datasync.py b/airflow/providers/amazon/aws/operators/datasync.py index 2eaa9142fb3fe..8fe4cc1988d12 100644 --- a/airflow/providers/amazon/aws/operators/datasync.py +++ b/airflow/providers/amazon/aws/operators/datasync.py @@ -118,11 +118,11 @@ class AWSDataSyncOperator(BaseOperator): "task_execution_kwargs", ) template_fields_renderers = { - "create_task_kwargs": "py", - "create_source_location_kwargs": "py", - "create_destination_location_kwargs": "py", - "update_task_kwargs": "py", - "task_execution_kwargs": "py", + "create_task_kwargs": "json", + "create_source_location_kwargs": "json", + "create_destination_location_kwargs": "json", + "update_task_kwargs": "json", + "task_execution_kwargs": "json", } ui_color = "#44b5e2" diff --git a/airflow/providers/amazon/aws/operators/ecs.py b/airflow/providers/amazon/aws/operators/ecs.py index ea83166dcf222..02ba27728331d 100644 --- a/airflow/providers/amazon/aws/operators/ecs.py +++ b/airflow/providers/amazon/aws/operators/ecs.py @@ -147,7 +147,12 @@ class ECSOperator(BaseOperator): ui_color = '#f0ede4' template_fields = ('overrides',) - template_fields_renderers = {"overrides": "json"} + template_fields_renderers = { + "overrides": "json", + "network_configuration": "json", + "tags": "json", + "quota_retry": "json", + } REATTACH_XCOM_KEY = "ecs_task_arn" REATTACH_XCOM_TASK_ID_TEMPLATE = "{task_id}_task_arn" diff --git a/airflow/providers/amazon/aws/operators/emr_add_steps.py b/airflow/providers/amazon/aws/operators/emr_add_steps.py index a410aa36fa303..03a2b93847918 100644 --- a/airflow/providers/amazon/aws/operators/emr_add_steps.py +++ b/airflow/providers/amazon/aws/operators/emr_add_steps.py @@ -47,6 +47,7 @@ class EmrAddStepsOperator(BaseOperator): template_fields = ['job_flow_id', 'job_flow_name', 'cluster_states', 'steps'] template_ext = ('.json',) + template_fields_renderers = {"steps": "json"} ui_color = '#f9c915' def __init__( diff --git a/airflow/providers/amazon/aws/operators/glue.py b/airflow/providers/amazon/aws/operators/glue.py index 41ed467c4bcce..34dbc5b611ce9 100644 --- a/airflow/providers/amazon/aws/operators/glue.py +++ b/airflow/providers/amazon/aws/operators/glue.py @@ -56,7 +56,10 @@ class AwsGlueJobOperator(BaseOperator): template_fields = ('script_args',) template_ext = () - template_fields_renderers = {"script_args": "py"} + template_fields_renderers = { + "script_args": "json", + "create_job_kwargs": "json", + } ui_color = '#ededed' def __init__( diff --git a/airflow/providers/amazon/aws/operators/s3_bucket_tagging.py b/airflow/providers/amazon/aws/operators/s3_bucket_tagging.py index 9e2cc80f40710..6f73cc0406cf0 100644 --- a/airflow/providers/amazon/aws/operators/s3_bucket_tagging.py +++ b/airflow/providers/amazon/aws/operators/s3_bucket_tagging.py @@ -86,6 +86,7 @@ class S3PutBucketTaggingOperator(BaseOperator): """ template_fields = ("bucket_name",) + template_fields_renderers = {"tag_set": "json"} def __init__( self, diff --git a/airflow/providers/amazon/aws/operators/sns.py b/airflow/providers/amazon/aws/operators/sns.py index 0cd0a72db0e75..2f2c9a9218603 100644 --- a/airflow/providers/amazon/aws/operators/sns.py +++ b/airflow/providers/amazon/aws/operators/sns.py @@ -42,7 +42,7 @@ class SnsPublishOperator(BaseOperator): template_fields = ['message', 'subject', 'message_attributes'] template_ext = () - template_fields_renderers = {"message_attributes": "py"} + template_fields_renderers = {"message_attributes": "json"} def __init__( self, diff --git a/airflow/providers/amazon/aws/operators/sqs.py b/airflow/providers/amazon/aws/operators/sqs.py index c7ac46ce2e1fe..fe8e0c0b83421 100644 --- a/airflow/providers/amazon/aws/operators/sqs.py +++ b/airflow/providers/amazon/aws/operators/sqs.py @@ -39,7 +39,8 @@ class SQSPublishOperator(BaseOperator): :type aws_conn_id: str """ - template_fields = ('sqs_queue', 'message_content', 'delay_seconds') + template_fields = ('sqs_queue', 'message_content', 'delay_seconds', 'message_attributes') + template_fields_renderers = {'message_attributes': 'json'} ui_color = '#6ad3fa' def __init__( diff --git a/airflow/providers/amazon/aws/transfers/mongo_to_s3.py b/airflow/providers/amazon/aws/transfers/mongo_to_s3.py index 95b2bfc12d58e..f21c5385a31e9 100644 --- a/airflow/providers/amazon/aws/transfers/mongo_to_s3.py +++ b/airflow/providers/amazon/aws/transfers/mongo_to_s3.py @@ -58,7 +58,7 @@ class MongoToS3Operator(BaseOperator): template_fields = ('s3_bucket', 's3_key', 'mongo_query', 'mongo_collection') ui_color = '#589636' - template_fields_renderers = {"mongo_query": "py"} + template_fields_renderers = {"mongo_query": "json"} def __init__( self, diff --git a/airflow/providers/amazon/aws/transfers/mysql_to_s3.py b/airflow/providers/amazon/aws/transfers/mysql_to_s3.py index 223d6bc52270b..2b5fa8d2c2f9c 100644 --- a/airflow/providers/amazon/aws/transfers/mysql_to_s3.py +++ b/airflow/providers/amazon/aws/transfers/mysql_to_s3.py @@ -68,7 +68,7 @@ class MySQLToS3Operator(BaseOperator): 'query', ) template_ext = ('.sql',) - template_fields_renderers = {"query": "sql"} + template_fields_renderers = {"query": "sql", "pd_csv_kwargs": "json"} def __init__( self, diff --git a/airflow/providers/amazon/aws/transfers/redshift_to_s3.py b/airflow/providers/amazon/aws/transfers/redshift_to_s3.py index 98c784f90e6f5..9da6b9c245b5c 100644 --- a/airflow/providers/amazon/aws/transfers/redshift_to_s3.py +++ b/airflow/providers/amazon/aws/transfers/redshift_to_s3.py @@ -72,7 +72,8 @@ class RedshiftToS3Operator(BaseOperator): """ template_fields = ('s3_bucket', 's3_key', 'schema', 'table', 'unload_options', 'select_query') - template_ext = () + template_ext = ('.sql',) + template_fields_renderers = {'select_query': 'sql'} ui_color = '#ededed' def __init__(