sql_query_transfer.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.sql import I2EToSQLOperator, SQLQueryToI2EOperator


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

    Requirements to use this DAG:
        - An I2E connection named `i2e_default`.
        - An SQL connection named `postgres_default`.
        - An SQL table named `input` with columns `id` and `text`, both with type TEXT (or equivalent).
        - An SQL table named `results` with columns `docs`, `date`, `age`, `number_hits`, all with type TEXT
          (or equivalent).
        - 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.
    """
    sql_to_i2e_task = i2e_source_data_uri = SQLQueryToI2EOperator(
        sql_conn_id="postgres_default",
        i2e_conn_id="i2e_default",
        sql_query="SELECT id, text FROM input",
    )

    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,
            "queryProperties": {"outputSettings": {"outputMode": "tsv"}},
        },
    )

    i2e_to_sql_task = I2EToSQLOperator(
        i2e_conn_id="i2e_default",
        sql_conn_id="postgres_default",
        i2e_query_results_uri=i2e_query_results_uri.output,
        sql_table="results",
        schema_mapping={
            "Doc": "docs",
            "Date": "date",
            "Age": "age",
            "#Hits": "number_hits",
        },
        use_mapped_columns_only=True,
    )

    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(
        sql_to_i2e_task,
        i2e_make_index_task,
        i2e_run_query_task,
        i2e_to_sql_task,
        i2e_delete_resources_task,
    )


sql_query_transfer()