Conectar Apache Airflow con ClickHouse

Apache Airflow es una plataforma de código abierto para crear, programar y supervisar flujos de trabajo como código. Los flujos de trabajo se definen como grafos acíclicos dirigidos (DAG) de tareas escritas en Python.

El proveedor apache-airflow-providers-clickhousedb conecta Airflow con ClickHouse, lo que le permite ejecutar consultas, crear tablas y cargar datos como parte de un DAG. Se conecta a través de la interfaz HTTP mediante el Client clickhouse-connect, y expone ClickHouse a través del marco SQL común de Airflow, por lo que el SQLExecuteQueryOperator estándar gestiona consultas DDL, DML y analíticas sin necesidad de un operador específico de ClickHouse.

Instala el proveedor en el entorno donde se ejecutan el scheduler y los workers de Airflow:

pip install apache-airflow-providers-clickhousedb

El proveedor depende de apache-airflow-providers-common-sql y clickhouse-connect, que se instalan junto con él. Para pasar los resultados de la consulta a DataFrames de pandas o polars, instala los extras opcionales:

pip install 'apache-airflow-providers-common-sql[pandas,polars]'

El proveedor registra un tipo de conexión clickhouse. Cree una conexión desde la UI de Airflow en Admin > Connections o defínala mediante la CLI o una variable de entorno.

En la UI, seleccione ClickHouse como tipo de conexión y complete los siguientes campos:

Para ClickHouse Cloud o cualquier clúster autohospedado con TLS habilitado, establezca secure en true en el campo Extra y use el puerto TLS (8443).

El proveedor expone opciones adicionales como campos específicos en el formulario de conexión. Si, en cambio, defines la conexión mediante URI, JSON o una variable de entorno, debes proporcionarlas como claves en el objeto JSON extra. Todas son opcionales:

Configure la conexión mediante una variable de entorno. El formato URI incluye el host, las credenciales y la base de datos:

export AIRFLOW_CONN_CLICKHOUSE_DEFAULT='clickhouse://default:password@localhost:8123/my_database'

Todos los componentes del URI deben estar codificados para URL. Para TLS, los tiempos de espera o la configuración de la sesión, use el formato JSON, que expone los campos de Extra:

export AIRFLOW_CONN_CLICKHOUSE_DEFAULT='{
    "conn_type": "clickhouse",
    "host": "abc123.clickhouse.cloud",
    "port": 8443,
    "login": "default",
    "password": "secret",
    "schema": "my_database",
    "extra": {
        "secure": true,
        "session_settings": {
            "max_execution_time": 300,
            "max_memory_usage": 10000000000
        }
    }
}'

Todos los hooks y operadores usan el ID de conexión clickhouse_default, a menos que especifiques otro.

Configura el conn_id del operador para que use tu conexión de ClickHouse. El siguiente DAG crea una tabla, inserta filas, las lee de nuevo y elimina la tabla:

from datetime import datetime

from airflow import DAG
from airflow.providers.common.sql.hooks.sql import fetch_all_handler
from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator

CLICKHOUSE_CONN_ID = "clickhouse_default"
CLICKHOUSE_TABLE = "airflow_example"

with DAG(
    dag_id="example_clickhouse",
    start_date=datetime(2021, 1, 1),
    default_args={"conn_id": CLICKHOUSE_CONN_ID},
    schedule="@once",
    catchup=False,
) as dag:
    create_table = SQLExecuteQueryOperator(
        task_id="create_table",
        sql=f"""
            CREATE TABLE IF NOT EXISTS {CLICKHOUSE_TABLE} (
                id   UInt32,
                name String,
                ts   DateTime DEFAULT now()
            ) ENGINE = MergeTree()
            ORDER BY id
        """,
    )

    insert_rows = SQLExecuteQueryOperator(
        task_id="insert_rows",
        sql=f"""
            INSERT INTO {CLICKHOUSE_TABLE} (id, name) VALUES
                (1, 'Alice'),
                (2, 'Bob'),
                (3, 'Charlie')
        """,
    )

    read_rows = SQLExecuteQueryOperator(
        task_id="read_rows",
        sql=f"SELECT id, name FROM {CLICKHOUSE_TABLE} ORDER BY id",
        handler=fetch_all_handler,
    )

    drop_table = SQLExecuteQueryOperator(
        task_id="drop_table",
        sql=f"DROP TABLE IF EXISTS {CLICKHOUSE_TABLE}",
    )

    create_table >> insert_rows >> read_rows >> drop_table

Los resultados de la consulta se obtienen con el handler predeterminado (fetch_all_handler). Para devolver algo distinto del conjunto completo de resultados, pase un handler diferente, como fetch_one_handler para devolver solo la primera fila.

Cuando una conexión apunta a un clúster y las tareas individuales consultan distintas bases de datos, sobrescriba la base de datos mediante hook_params en lugar de crear una conexión independiente:

read_rows = SQLExecuteQueryOperator(
    task_id="read_rows",
    conn_id=CLICKHOUSE_CONN_ID,
    sql="SELECT count() FROM events",
    hook_params={"database": "analytics"},
)

Para tareas que no encajan en un operador SQL — bulk inserts, streaming o llamadas del client específicas de ClickHouse — use ClickHouseHook dentro de una tarea de Python.

El método bulk_insert_rows del hook usa la ruta de inserción columnar nativa de clickhouse-connect, que es mucho más rápida que las inserciones fila por fila para grandes volúmenes de datos. Establezca batch_size para limitar el uso máximo de memoria con entradas muy grandes:

from airflow.providers.clickhousedb.hooks.clickhouse import ClickHouseHook

hook = ClickHouseHook(clickhouse_conn_id="clickhouse_default")

hook.bulk_insert_rows(
    table="events",
    rows=[("user1", "click"), ("user2", "view")],
    column_names=["user_id", "action"],
    batch_size=1000,
)

Llama a get_client() para acceder al client subyacente de clickhouse-connect para cualquier función que el hook no exponga directamente:

client = hook.get_client()
total = client.query("SELECT count() FROM events").result_rows[0][0]

Pasa la configuración de sesión al crear el hook, ya sea directamente o mediante el hook_params del operador. La configuración que se pasa al constructor se superpone a cualquier session_settings definido en el campo Extra de la conexión, y los valores del constructor prevalecen cuando hay conflicto entre claves:

hook = ClickHouseHook(
    clickhouse_conn_id="clickhouse_default",
    session_settings={"max_execution_time": 60, "max_threads": 4},
)

• clickhouse-connect Client para Python
• interfaz HTTP de ClickHouse
• referencia de la configuración de sesión de ClickHouse
• documentación de referencia de apache-airflow-providers-clickhousedb
• paquete del proveedor en PyPI