sql_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.sql import I2EToSQLOperator, SQLTimestampToI2EOperator


@dag(
    dag_id="example_dags.sql_timespan_transfer",
    start_date=datetime(year=2024, month=1, day=1),
    schedule=timedelta(hours=1),
    catchup=False,
    render_template_as_native_obj=True,
)
def sql_timespan_transfer():
    """
    DAG showing how to use I2E operators to periodically upload recent source data (from the last hour)
    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.

    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 SQL connection named `postgres_default`.
        - An SQL table named `input` with columns `id`, `text`, `created_at`, where `created_at` is type TIMESTAMP
          (or equivalent) and others 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 = SQLTimestampToI2EOperator(
        sql_conn_id="postgres_default",
        i2e_conn_id="i2e_default",
        sql_table="input",
        sql_timestamp_column="created_at",
        # 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",
        },
        # Integration tests run many make-index tasks back-to-back; 300s is too tight on CI runners.
        i2e_wait_timeout=600,
    )

    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_timespan_transfer()