i2e_resource_bundles.py

from airflow.sdk import dag, task
from airflow.sdk.bases.operator import chain

from iqvia_nlp_provider.transfers.local import LocalFilesystemToI2EResourceBundleOperator
from iqvia_nlp_provider.transfers.s3 import S3ToI2EResourceBundleOperator
from iqvia_nlp_provider.utils.i2e_connections import get_all_i2e_connection_ids


@dag(dag_id="example_dags.i2e_resource_bundles", schedule=None, catchup=False)
def i2e_resource_bundles():
    """
    DAG showing how to use the S3ToI2EResourceBundleOperator and LocalFilesystemToI2EResourceBundleOperator to
    retrieve an I2E resource bundle from S3 or a local file, and install the bundle on multiple I2E servers.

    The bundle from S3 will be installed on the I2E servers 'i2e_default' and 'i2e_default_2'.
    The bundle from the local file system will be installed on all I2E servers added to Airflow.

    Requirements to use this DAG:
        - I2E connections named `i2e_default` and `i2e_default_2`.
        - An AWS connection named `aws_conn_id`.
        - An S3 bucket named `s3-bucket`.
        - An I2E resource bundle in the S3 bucket with key `my_bundle_1.zip`
        - An I2E resource bundle at `/work_dir/my_bundle_2.zip` on the Airflow worker's local filesystem.
    """

    @task
    def get_all_i2e_connection_ids_task():
        return get_all_i2e_connection_ids()

    i2e_conn_ids = get_all_i2e_connection_ids_task()

    s3_to_i2e_task = S3ToI2EResourceBundleOperator(
        i2e_conn_ids=["i2e_default", "i2e_default_2"],
        aws_conn_id="aws_conn_id",
        aws_bucket_name="s3-bucket",
        aws_key="my_bundle_1.zip",
    )

    local_filesystem_to_i2e_task = LocalFilesystemToI2EResourceBundleOperator(
        i2e_conn_ids=i2e_conn_ids, local_file_path="/work_dir/my_bundle_2.zip"
    )

    chain(s3_to_i2e_task, i2e_conn_ids, local_filesystem_to_i2e_task)


i2e_resource_bundles()