s3_timespan_transfer.py

from datetime import datetime, timedelta

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

from iqvia_nlp_provider.operators.i2e import (
    I2EDeleteResourcesOperator,
    I2EMakeIndexOperator,
    I2EPrintTaskLogOperator,
    I2ERunQueryOperator,
)
from iqvia_nlp_provider.transfers.s3 import I2EToS3Operator, S3ToI2EOperator


@dag(
    dag_id="example_dags.s3_timespan_transfer",
    start_date=datetime(year=2024, month=1, day=1),
    schedule=timedelta(hours=1),
    catchup=False,
    render_template_as_native_obj=True,
)
def s3_timespan_transfer():
    """
    DAG showing how to use I2E operators to periodically upload recent source data
    (from the last hour) from an S3 bucket to I2E, index the source data, run a query
    against the index, write the query results to an S3 bucket, and delete all
    transient resources from I2E.

    The DAG is scheduled to run at every hour, and will only process data created after the data
    interval start and before the data interval end (typically the previous hour). When the DAG run
    has no data interval (for example some manual triggers), a one-hour window ending at
    ``dag_run.logical_date``, ``dag_run.run_after``, or ``dag_run.start_date`` is used instead (in that order).

    Requirements to use this DAG:
        - An I2E connection named `i2e_default`.
        - An AWS connection named `aws_conn_id`.
        - An S3 bucket named `s3-bucket`.
        - Multiple source data files in the S3 bucket with the key prefix `source/`.
        - An index template named `data-factory-workflow-example` installed on the I2E server.
        - A query template named `data-factory-workflow-example` installed on the I2E server.
    """
    s3_to_i2e_task = i2e_source_data_uri = S3ToI2EOperator(
        aws_conn_id="aws_conn_id",
        i2e_conn_id="i2e_default",
        aws_bucket_name="s3-bucket",
        aws_prefix="source/",
        # Airflow 3: scheduled runs have `data_interval_*`; manual/API runs may omit them and set `logical_date` to
        # None — use `run_after` (or `start_date`) as the anchor for the one-hour window.
        select_data_added_between=(
            "{{ dag_run.data_interval_start or (("
            "dag_run.logical_date or dag_run.run_after or dag_run.start_date) - "
            "macros.timedelta(hours=1)) }}",
            "{{ dag_run.data_interval_end or dag_run.logical_date or dag_run.run_after or dag_run.start_date }}",
        ),
    )

    i2e_make_index_task = i2e_index_uri = I2EMakeIndexOperator(
        i2e_conn_id="i2e_default",
        i2e_source_data_uri=i2e_source_data_uri.output,
        i2e_index_template="data-factory-workflow-example",
        i2e_index_settings_overrides={
            "useDaemon": True,
            "checkDaemon": True,
            "ontologyIndexSuppression": "All",
        },
    )

    i2e_print_task_log_task = I2EPrintTaskLogOperator(
        i2e_conn_id="i2e_default",
    )

    i2e_run_query_task = i2e_query_results_uri = I2ERunQueryOperator(
        i2e_conn_id="i2e_default",
        i2e_index_uri=i2e_index_uri.output,
        i2e_query_template="data-factory-workflow-example",
        i2e_query_settings_overrides={
            "checkClasses": False,
        },
    )

    i2e_to_s3_task = I2EToS3Operator(
        i2e_conn_id="i2e_default",
        aws_conn_id="aws_conn_id",
        i2e_query_results_uri=i2e_query_results_uri.output,
        aws_bucket_name="s3-bucket",
    )

    i2e_delete_resources_task = I2EDeleteResourcesOperator(
        i2e_conn_id="i2e_default",
        i2e_resource_uris=[i2e_source_data_uri.output, i2e_index_uri.output, i2e_query_results_uri.output],
    )

    # Print task log runs only if make index or query fails
    i2e_make_index_task >> i2e_print_task_log_task
    i2e_run_query_task >> i2e_print_task_log_task

    chain(
        s3_to_i2e_task,
        i2e_make_index_task,
        i2e_run_query_task,
        i2e_to_s3_task,
        i2e_delete_resources_task,
    )


s3_timespan_transfer()