Combining task groups, dynamic mapping and DockerOperator

[odhondt]

Hi

Consider the following DAG:

from airflow import DAG
from airflow.decorators import task, task_group
from datetime import datetime, timedelta
from time import sleep

with DAG(
    dag_id="queries_and_download",
    schedule=timedelta(minutes=10),
    start_date=datetime(2022, 12, 22),
    catchup=False,
) as dag:

    @task
    def query_all_sources():
        sleep(5)
        return [1, 2, 3, 4, 5]

    @task_group
    def calculations(x):
        @task(pool="download")
        def download_product(product):
            sleep(5)
            print(f"downloading product {product}")
            return product

        @task(pool="process")
        def process_product(product):
            sleep(10)
            print(f"processing product {product}")

        process_product(download_product(x))

    calculations.expand(x=query_all_sources())

I would like to do something similar replacing download_product and process_product with two DockerOperator which would fire up scripts (say download_product.py, process_product.py) that both would take the dynamically generated outputs from query_all_sources as an argument.

So far I have tried to use various combinations of xcoms and templates to automatically fill the command parameter of the operators but with no success.

Can this be achieved?

[odhondt]

To answer my own question: I have found a way to make it work by using the @task.docker decorator and the subprocess module to launch scripts inside a container.