i2e_query_with_i2e_hook.py

import os
from shutil import copyfileobj

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

from iqvia_nlp_provider.hooks.i2e import I2EHook


@dag(dag_id="example_dags.i2e_query_with_i2e_hook", schedule=None, catchup=False)
def i2e_query_with_i2e_hook():
    """
    DAG showing how to use the I2EHook directly 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.
    """

    @task
    def local_filesystem_to_i2e() -> str:
        i2e = I2EHook(i2e_conn_id="i2e_default")
        with open("/work_dir/source.txt", "rb") as fd:
            return i2e.upload_source_data(filename="source.txt", content=fd)

    @task
    def i2e_make_index(i2e_source_data_uri: str) -> str:
        i2e = I2EHook(i2e_conn_id="i2e_default")
        i2e_make_index_task_uri = i2e.post_make_index_task(
            source_data_uri=i2e_source_data_uri,
            index_template="data-factory-workflow-example",
            index_settings_overrides={"useDaemon": True, "checkDaemon": True, "ontologyIndexSuppression": "All"},
        )
        return i2e.wait_for_make_index_task(make_index_task_uri=i2e_make_index_task_uri)

    @task
    def i2e_run_query(i2e_index_uri: str) -> str:
        i2e = I2EHook(i2e_conn_id="i2e_default")
        i2e_query_task_uri = i2e.post_query_task(
            index_uri=i2e_index_uri,
            query_template="data-factory-workflow-example",
            query_settings_overrides={"checkClasses": False},
        )
        return i2e.wait_for_query_task(query_task_uri=i2e_query_task_uri)

    @task
    def i2e_to_local_filesystem(i2e_query_results_uri: str):
        i2e = I2EHook(i2e_conn_id="i2e_default")
        i2e_query_results = i2e.get_query_results(query_results_uri=i2e_query_results_uri)
        filename = os.path.split(i2e_query_results_uri)[1]
        with open(f"/work_dir/results/{filename}", "wb") as fd:
            copyfileobj(fsrc=i2e_query_results, fdst=fd)

    @task
    def i2e_delete_resources(i2e_source_data_uri: str, i2e_index_uri: str, i2e_query_results_uri: str):
        i2e = I2EHook(i2e_conn_id="i2e_default")
        i2e.delete_resource(i2e_source_data_uri)
        i2e.delete_resource(i2e_index_uri)
        i2e.delete_resource(i2e_query_results_uri)

    i2e_source_data_uri = local_filesystem_to_i2e()
    i2e_index_uri = i2e_make_index(i2e_source_data_uri)
    i2e_query_results_uri = i2e_run_query(i2e_index_uri)

    chain(
        i2e_source_data_uri,
        i2e_index_uri,
        i2e_query_results_uri,
        i2e_to_local_filesystem(i2e_query_results_uri),
        i2e_delete_resources(i2e_source_data_uri, i2e_index_uri, i2e_query_results_uri),
    )


i2e_query_with_i2e_hook()