airflow-clickhouse-plugin から ClickHouse プロバイダへの移行

pip install apache-airflow-providers-clickhousedb

pip uninstall airflow-clickhouse-plugin clickhouse-driver

{
    "conn_type": "sqlite",
    "host": "ch.example.com",
    "port": 9440,
    "login": "airflow",
    "password": "secret",
    "schema": "analytics",
    "extra": {
        "secure": true,
        "settings": {"max_execution_time": 300},
        "compression": true
    }
}

{
    "conn_type": "clickhouse",
    "host": "ch.example.com",
    "port": 8443,
    "login": "airflow",
    "password": "secret",
    "schema": "analytics",
    "extra": {
        "secure": true,
        "session_settings": {"max_execution_time": 300},
        "compress": true
    }
}

airflow connections test clickhouse_default

from airflow_clickhouse_plugin.operators.clickhouse import ClickHouseOperator

update_income_aggregate = ClickHouseOperator(
    task_id="update_income_aggregate",
    clickhouse_conn_id="clickhouse_test",
    database="default",
    sql=(
        """
        INSERT INTO aggregate
        SELECT eventDt, sum(price * qty) AS income FROM sales
        WHERE eventDt = '{{ ds }}' GROUP BY eventDt
        """,
        """
        SELECT sum(income) FROM aggregate
        WHERE eventDt BETWEEN
            '{{ data_interval_start | ds }}' AND '{{ data_interval_end | ds }}'
        """,
    ),
    settings={"max_execution_time": 600},
)

from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator

update_income_aggregate = SQLExecuteQueryOperator(
    task_id="update_income_aggregate",
    conn_id="clickhouse_test",
    database="default",
    sql=[
        """
        INSERT INTO aggregate
        SELECT eventDt, sum(price * qty) AS income FROM sales
        WHERE eventDt = '{{ ds }}' GROUP BY eventDt
        """,
        """
        SELECT sum(income) FROM aggregate
        WHERE eventDt BETWEEN
            '{{ data_interval_start | ds }}' AND '{{ data_interval_end | ds }}'
        """,
    ],
    hook_params={"session_settings": {"max_execution_time": 600}},
)

SQLExecuteQueryOperator(
    task_id="update_income_aggregate",
    conn_id="clickhouse_test",
    sql="""
        INSERT INTO aggregate
        SELECT eventDt, sum(price * qty) AS income FROM sales
        WHERE eventDt = '{{ ds }}' GROUP BY eventDt;
        SELECT sum(income) FROM aggregate WHERE eventDt = '{{ ds }}'
    """,
    split_statements=True,
)

from airflow.providers.common.sql.operators.sql import SQLInsertRowsOperator

load_rows = SQLInsertRowsOperator(
    task_id="load_rows",
    conn_id="clickhouse_default",
    table_name="some_ch_table",
    columns=["id", "name"],
    rows=extract_task.output,
    insert_args={"executemany": True},
)

def fetch_with_column_types(cursor):
    return cursor.fetchall(), [(col[0], col[1]) for col in cursor.description]


SQLExecuteQueryOperator(
    task_id="typed_query",
    conn_id="clickhouse_default",
    sql="SELECT id, name FROM users LIMIT 10",
    handler=fetch_with_column_types,
)

from airflow_clickhouse_plugin.hooks.clickhouse import ClickHouseHook


def sqlite_to_clickhouse():
    records = SqliteHook().get_records("SELECT id, name FROM some_sqlite_table")
    ClickHouseHook().execute("INSERT INTO some_ch_table VALUES", records)

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


def sqlite_to_clickhouse():
    records = SqliteHook().get_records("SELECT id, name FROM some_sqlite_table")
    ClickHouseHook().bulk_insert_rows(
        "some_ch_table", records, column_names=["id", "name"], batch_size=100_000
    )

from clickhouse_connect.driver.external import ExternalData

hook = ClickHouseHook()
with hook.get_client() as client:
    ext = ExternalData(file_name="ids", structure="id UInt64", data=b"1\n2\n3\n")
    result = client.query(
        "SELECT * FROM events WHERE id IN ids",
        external_data=ext,
        settings={"query_id": "my-traceable-id"},
    )
    columns = result.result_columns

from airflow_clickhouse_plugin.sensors.clickhouse import ClickHouseSensor

ClickHouseSensor(
    task_id="poke_events_count",
    database="monitor",
    sql="SELECT count() FROM warnings WHERE eventDate = '{{ ds }}'",
    is_success=lambda result: result[0][0] > 10000,
)

from airflow.providers.common.sql.sensors.sql import SqlSensor

SqlSensor(
    task_id="poke_events_count",
    conn_id="clickhouse_default",
    hook_params={"database": "monitor"},
    sql="SELECT count() FROM warnings WHERE eventDate = '{{ ds }}'",
    success=lambda cnt: cnt > 10000,
)

Migrate this repository from airflow-clickhouse-plugin to
apache-airflow-providers-clickhousedb following
https://clickhouse.com/docs/integrations/airflow/migrating-from-airflow-clickhouse-plugin

- Replace every airflow_clickhouse_plugin import per the class mapping table.
- Rename clickhouse_conn_id to conn_id on operators and sensors, not on hooks; add
  conn_id="clickhouse_default" wherever a task relied on the plugin's default.
- Move settings= into hook_params={"session_settings": ...}.
- Replace hook.execute(...) with get_records, get_first, run or bulk_insert_rows
  according to the hook table; never pass a list of rows as parameters. Replace
  ClickHouseOperator(parameters=<rows>) with SQLInsertRowsOperator.
- Rewrite sensor callables to receive the first cell instead of the full result.
- Flag, but do not silently rewrite, any use of with_column_types, external_tables,
  columnar, query_id or types_check, and any XCom consumer of an INSERT task.
- List every Airflow connection that must change port and extras; do not edit
  connections yourself.
Show the diff and a summary of items flagged for human review.