i2e_query.py

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.local import I2EToLocalFilesystemOperator, LocalFilesystemToI2EOperator


@dag(dag_id="example_dags.i2e_query", schedule=None, catchup=False)
def i2e_query():
    """
    DAG showing how to use I2E operators to upload source data from a local file to I2E,
    index the source data, run a query against the index, write the query results
    to a local file, and delete all transient resources from I2E.

    Requirements to use this DAG:
        - An I2E connection named `i2e_default`.
        - A source file at '/work_dir/source.txt' on the Airflow worker's filesystem.
        - A directory named '/work_dir/results/' on the Airflow worker's filesystem.
        - 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.
    """
    local_filesystem_to_i2e_task = i2e_source_data_uri = LocalFilesystemToI2EOperator(
        i2e_conn_id="i2e_default",
        local_file_path="/work_dir/source.txt",
    )

    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_local_filesystem_task = I2EToLocalFilesystemOperator(
        i2e_conn_id="i2e_default",
        i2e_query_results_uri=i2e_query_results_uri.output,
        local_folder="/work_dir/results/",
    )

    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(
        local_filesystem_to_i2e_task,
        i2e_make_index_task,
        i2e_run_query_task,
        i2e_to_local_filesystem_task,
        i2e_delete_resources_task,
    )


i2e_query()