i2e_query_without_transfers.py

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

from iqvia_nlp_provider.operators.i2e import I2EMakeIndexOperator, I2EPrintTaskLogOperator, I2ERunQueryOperator


@dag(dag_id="example_dags.i2e_query_without_transfers", schedule=None, catchup=False)
def i2e_query_without_transfers():
    """
    DAG showing how to use I2E operators to index a source data file in I2E and
    run a query against the index. This DAG does not upload any source data to I2E,
    and does not copy the query results to anywhere outside I2E.

    Requirements to use this DAG:
        - An I2E connection named `i2e_default`.
        - Source data named 'source.txt' already present on the I2E server.
        - 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.
    """
    i2e_make_index_task = I2EMakeIndexOperator(
        i2e_conn_id="i2e_default",
        i2e_source_data_uri="/api;type=source_data/source.txt",
        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 = I2ERunQueryOperator(
        i2e_conn_id="i2e_default",
        i2e_index_uri=i2e_make_index_task.output,
        i2e_query_template="data-factory-workflow-example",
        i2e_query_settings_overrides={
            "checkClasses": False,
        },
    )

    # 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(i2e_make_index_task, i2e_run_query_task)


i2e_query_without_transfers()