Airflow ClickHouse Plugin

Since there is no dedicated ClickHouse connection type, create a new connection in Airflow using SQLite or any other SQL database type.

Connection Mapping

The following Airflow connection attributes are mapped to clickhouse_driver.Client arguments:

Important Notes

• Database Overrides: The database argument in ClickHouseOperator, ClickHouseHook, or ClickHouseSensor will override the schema attribute defined in the Airflow connection.
• Default Host: If host is not set, it defaults to localhost.
• Default Connection ID: The plugin uses clickhouse_default as the default connection ID if none is specified.

Unit Tests

Run unit tests using:

python3 -m unittest discover -t tests -s unit

Integration Tests

Integration tests require a running ClickHouse server. You can use Docker to set up a local environment:

1. Start ClickHouse:

   docker run -p 9000:9000 --ulimit nofile=262144:262144 -it clickhouse/clickhouse-server

2. Run Tests: Use the AIRFLOW_CONN_CLICKHOUSE_DEFAULT environment variable to point to your local container:

   PYTHONPATH=src AIRFLOW_CONN_CLICKHOUSE_DEFAULT=clickhouse://localhost python3 -m unittest discover -t tests -s integration

Run all tests (Unit & Integration)

PYTHONPATH=src AIRFLOW_CONN_CLICKHOUSE_DEFAULT=clickhouse://localhost python3 -m unittest discover -s tests

# Run all tests with ClickHouse connection defined
PYTHONPATH=src AIRFLOW_CONN_CLICKHOUSE_DEFAULT=clickhouse://localhost python3 -m unittest discover -s tests

For compatibility with the Python DB API 2.0 standard, use the clickhouse_dbapi module. These classes are useful when you need to work with standard DB-API patterns.

Key components:

• ClickHouseSqlSensor: A sensor that uses hook_params to configure the connection (e.g., schema).
• ClickHouseSQLExecuteQueryOperator: An operator for executing queries using the DB-API interface.

from airflow import DAG
from airflow_clickhouse_plugin.sensors.clickhouse_dbapi import ClickHouseSqlSensor
from airflow_clickhouse_plugin.operators.clickhouse_dbapi import ClickHouseSQLExecuteQueryOperator
from airflow.utils.dates import days_ago


with DAG(
        dag_id='listen_warnings',
        start_date=days_ago(2),
) as dag:
    dag >> ClickHouseSqlSensor(
        task_id='poke_events_count',
        hook_params=dict(schema='monitor'),
        sql="SELECT count() FROM warnings WHERE eventDate = '{{ ds }}'",
        success=lambda cnt: cnt > 10000,
        conn_id=None,  # required by common.sql SqlSensor; use None for default
    ) >> ClickHouseSQLExecuteQueryOperator(
        task_id='create_alert',
        database='alerts',
        sql='''
            INSERT INTO events SELECT eventDate, count()
            FROM monitor.warnings WHERE eventDate = '{{ ds }}'
        ''',
    )

The ClickHouseOperator allows you to execute one or multiple SQL statements. If multiple statements are provided in a list, the results of the last query are pushed to XCom.

Key parameters:

• sql: A string or a list of strings containing the SQL commands.
• database: The ClickHouse database to use (overrides connection schema).
• clickhouse_conn_id: The Airflow connection ID to use.
• query_id: A templated string used to identify the query in ClickHouse logs.

from airflow import DAG
from airflow_clickhouse_plugin.operators.clickhouse import ClickHouseOperator
from airflow.utils.dates import days_ago

with DAG(
        dag_id='update_income_aggregate',
        start_date=days_ago(2),
) as dag:
    ClickHouseOperator(
        task_id='update_income_aggregate',
        database='default',
        sql=(
            '''
                INSERT INTO aggregate
                SELECT eventDt, sum(price * qty) AS income FROM sales
                WHERE eventDt = '{{ ds }}' GROUP BY eventDt
            ''', '''
                OPTIMIZE TABLE aggregate ON CLUSTER {{ var.value.cluster_name }}
                PARTITION toDate('{{ execution_date.format('%Y-%m-01') }}')
            ''',
        ),
        query_id='{{ ti.dag_id }}-{{ ti.task_id }}-{{ ti.run_id }}-{{ ti.try_number }}',
        clickhouse_conn_id='clickhouse_test',
    )

The ClickHouseHook provides a programmatic interface to ClickHouse.

Warning on INSERTs: When using ch_hook.execute(), do not use literal string forms like INSERT INTO table VALUES (1). Because of the native ClickHouse protocol requirements, clickhouse-driver requires that values for INSERT queries be provided via the parameters argument.

from airflow import DAG
from airflow_clickhouse_plugin.hooks.clickhouse import ClickHouseHook
from airflow.providers.sqlite.hooks.sqlite import SqliteHook
from airflow.operators.python import PythonOperator
from airflow.utils.dates import days_ago


def sqlite_to_clickhouse():
    sqlite_hook = SqliteHook()
    ch_hook = ClickHouseHook()
    records = sqlite_hook.get_records('SELECT * FROM some_sqlite_table')
    # Note: records should be passed via parameters in a real scenario
    ch_hook.execute('INSERT INTO some_ch_table VALUES', records)


with DAG(
        dag_id='sqlite_to_clickhouse',
        start_date=days_ago(2),
) as dag:
    dag >> PythonOperator(
        task_id='sqlite_to_clickhouse',
        python_callable=sqlite_to_clickhouse,
    )

The ClickHouseSensor pokes a ClickHouse database to check if a condition is met based on a SQL query.

Key parameters:

• sql: The query to execute.
• is_success: A lambda or function that takes the query result as input and returns True if the condition is met.
• database: The ClickHouse database to query.

from airflow import DAG
from airflow_clickhouse_plugin.sensors.clickhouse import ClickHouseSensor
from airflow_clickhouse_plugin.operators.clickhouse import ClickHouseOperator
from airflow.utils.dates import days_ago


with DAG(
        dag_id='listen_warnings',
        start_date=days_ago(2),
) as dag:
    dag >> ClickHouseSensor(
        task_id='poke_events_count',
        database='monitor',
        sql="SELECT count() FROM warnings WHERE eventDate = '{{ ds }}'",
        is_success=lambda cnt: cnt > 10000,
    ) >> ClickHouseOperator(
        task_id='create_alert',
        database='alerts',
        sql='''
            INSERT INTO events SELECT eventDate, count()
            FROM monitor.warnings WHERE eventDate = '{{ ds }}'
        ''',
    )