airflow.contrib.operators.pubsub_operator.PubSubPublishOperator can't handle multiple messages in a template

[prodhype]

Apache Airflow version: 1.10.10

Kubernetes version (if you are using kubernetes) (use kubectl version): 1.19

Environment: composer

• Cloud provider or hardware configuration: GCP
• OS (e.g. from /etc/os-release): Debian Buster
• Kernel (e.g. uname -a): Linux cs-1078880547532-default-default-cnxpx 5.4.49+ #1 SMP Sun Oct 18 19:43:35 PDT 2020 x86_64 GNU/Linux
• Install tools: GCP
• Others: none

What happened: When a template is used to send multiple pubsub messages in a list, pubsub returns the error ERROR - ('Error publishing to topic projects/my-cloud-composer/topics/my-devtest', <HttpError 400 when requesting https://pubsub.googleapis.com/v1/projects/my-cloud-composer/topics/mytopic:publish?alt=json returned "Invalid value at 'messages' (type.googleapis.com/google.pubsub.v1.PubsubMessage), "[{'data': 'eyJwb2 (... trimmed ...) Details: "[{'@type': 'type.googleapis.com/google.rpc.BadRequest', 'fieldViolations': [{'field': 'messages', 'description': 'Invalid value at \'messages\' (type.googleapis.com/google.pubsub.v1.PubsubMessage), .... For example, if the following template contains a list of messages I get the error:

with dag:
    my_messages = "{{ ti.xcom_pull(task_ids='transform', key='messages') }}"

    enqueue_data = PubSubPublishOperator(
        task_id="enqueue",
        project=PUBSUB_PROJECT,
        topic=PUBSUB_TOPIC,
        messages=my_messages,
        retries=2,
        retry_delay=30,
        sla=datetime.timedelta(seconds=1200),
        dag=dag,
    )

When I test with a single message that is not templated it works as intended.

When I try it with a single message in a template I get the error as well

What you expected to happen: The messages in the list should get enqueued to pubsub correctly.

How to reproduce it: See the above example dag and run it.

Anything else we need to know: none

[boring-cyborg]

Thanks for opening your first issue here! Be sure to follow the issue template!

[mik-laj]

Jinja always returns text when the messages parameter accepts a list of dictionaries.
https://airflow.readthedocs.io/en/latest/howto/operator/google/cloud/pubsub.html

[prodhype]

The fact that Jinja is returning text means that a list of messages cannot be used in the templated argument for messages. json.loads() doesn't work because it tries to loads the template string. Using Jinjas NativeEnvironment doesn't work because there's no way to set ti= when calling render() on the NativeEnvironment instance.

[prodhype]

I was thinking this could be fixed with something like: https://github.com/prodhype/airflow/pull/1

[turbaszek]

I was thinking this could be fixed with something like: prodhype#1

Some time ago I did this to solve similar problem:

airflow/airflow/providers/google/cloud/operators/automl.py

Lines 1257 to 1263 in ba2b3b2

	def _parse_dataset_id(dataset_id: Union[str, List[str]]) -> List[str]:
	if not isinstance(dataset_id, str):
	return dataset_id
	try:
	return ast.literal_eval(dataset_id)
	except (SyntaxError, ValueError):
	return dataset_id.split(",")

But in general I agree that it would be good to have better support for lazy evaluation with python native types.

[hamsapriyak]

It is important that the render_templates(context) returns python native types. Without this, even if there are airflow operators available we end up creating additional custom operators to handle this .

Eg:. We need to delete files [‘file_1’, ‘file_2’, ‘file_3’] from s3 bucket. This file list is passed from a previous task (get_files) via Xcom. Although S3DeleteObjectsOperator accepts list of keys , due to templating it is changed as "[‘file_1’, ‘file_2’, ‘file_3’]". To handle this we need to create a custom operator which would receive this keys via context and call the s3_hook.

Sample Code using templating:

get_files = S3ListOperator(
    task_id='get_files',
    bucket='source_bucket',
    dag=dag
)

delete_files = S3DeleteObjectsOperator(
    task_id='delete_files',
    bucket=source_bucket,
    keys=" {{ task_instance.xcom_pull(task_ids='get_files') }}",
    dag=dag
)

get_files >> delete_files

Log output:

{s3_delete_objects_operator.py:83} INFO - Deleted: [“[‘file_1’, ‘file_2’, ‘file_3]”]

[turbaszek]

It seems that in 2.0 you can use output attribute to do what you want:

class CustomOp(BaseOperator):
    @apply_defaults
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def execute(self, context: Any):
        return [{"a": 1, "b": 2} for _ in range(5)]


class CustomOp2(BaseOperator):
    template_fields = ["xs"]

    @apply_defaults
    def __init__(self, xs, **kwargs):
        super().__init__(**kwargs)
        self.xs = xs

    def execute(self, context: Any):
        self.log.info(self.xs)
        for x in self.xs:
            self.log.info(x)


with DAG(
    'test_xcom_native',
    default_args=default_args
) as dag2:
    op = CustomOp(task_id="generate")
    CustomOp2(task_id="print", xs=op.output)

and you get:

[2021-01-13 07:44:12,266] {xcom.py:52} INFO - [{'a': 1, 'b': 2}, {'a': 1, 'b': 2}, {'a': 1, 'b': 2}, {'a': 1, 'b': 2}, {'a': 1, 'b': 2}]
[2021-01-13 07:44:12,267] {xcom.py:54} INFO - {'a': 1, 'b': 2}
[2021-01-13 07:44:12,269] {xcom.py:54} INFO - {'a': 1, 'b': 2}
[2021-01-13 07:44:12,270] {xcom.py:54} INFO - {'a': 1, 'b': 2}
[2021-01-13 07:44:12,271] {xcom.py:54} INFO - {'a': 1, 'b': 2}
[2021-01-13 07:44:12,272] {xcom.py:54} INFO - {'a': 1, 'b': 2}

To learn more see:
https://airflow.apache.org/docs/apache-airflow/stable/tutorial_taskflow_api.html
https://www.polidea.com/blog/airflow-2-0-dag-authoring-redesigned/

[hamsapriyak]

Thank you @turbaszek . This sure helps. Will try this out in 2.0.

[turbaszek]

@hamsapriyak so I'm closing this now as the issue seems to be solvable in 2.0