From f67f3fd4ce1534a4d412ac89c2d78030d9e39c92 Mon Sep 17 00:00:00 2001 From: Andrey Tatarinov Date: Wed, 8 Jul 2026 14:41:41 +0400 Subject: [PATCH] Move k2ch to tkati --- .github/workflow-templates/test.template.yml | 7 +- .github/workflows/test-k2ch.yml | 76 ++++++++++++++ .github/workflows/test-tkati-core.yml | 7 +- packages/tkati-k2ch/README.md | 52 ++++++++++ packages/tkati-k2ch/pyproject.toml | 32 ++++++ packages/tkati-k2ch/settings.test.toml | 33 ++++++ packages/tkati-k2ch/src/k2ch/__init__.py | 0 packages/tkati-k2ch/src/k2ch/__main__.py | 4 + packages/tkati-k2ch/src/k2ch/main.py | 46 ++++++++ packages/tkati-k2ch/src/k2ch/py.typed | 0 packages/tkati-k2ch/src/k2ch/settings.py | 15 +++ packages/tkati-k2ch/tests/conftest.py | 96 +++++++++++++++++ packages/tkati-k2ch/tests/test_k2ch.py | 104 +++++++++++++++++++ pyproject.toml | 2 + 14 files changed, 472 insertions(+), 2 deletions(-) create mode 100644 .github/workflows/test-k2ch.yml create mode 100644 packages/tkati-k2ch/README.md create mode 100644 packages/tkati-k2ch/pyproject.toml create mode 100644 packages/tkati-k2ch/settings.test.toml create mode 100644 packages/tkati-k2ch/src/k2ch/__init__.py create mode 100644 packages/tkati-k2ch/src/k2ch/__main__.py create mode 100644 packages/tkati-k2ch/src/k2ch/main.py create mode 100644 packages/tkati-k2ch/src/k2ch/py.typed create mode 100644 packages/tkati-k2ch/src/k2ch/settings.py create mode 100644 packages/tkati-k2ch/tests/conftest.py create mode 100644 packages/tkati-k2ch/tests/test_k2ch.py diff --git a/.github/workflow-templates/test.template.yml b/.github/workflow-templates/test.template.yml index 4ee3071..8a4980d 100644 --- a/.github/workflow-templates/test.template.yml +++ b/.github/workflow-templates/test.template.yml @@ -63,7 +63,12 @@ jobs: run: uv sync --package {{ package.name }} - name: Check formatting - run: uv run ruff check {{ package.path }} + run: uvx ruff check {{ package.path }} + + - name: Check typing + shell: bash + run: | + uv run mypy -p {{ package.package_name }} - name: Run tests run: uv run pytest {{ package.path }} diff --git a/.github/workflows/test-k2ch.yml b/.github/workflows/test-k2ch.yml new file mode 100644 index 0000000..514742f --- /dev/null +++ b/.github/workflows/test-k2ch.yml @@ -0,0 +1,76 @@ +# This file was automatically generated by uv-workspace-codegen +# For more information, see: https://github.com/epoch8/uv-workspace-codegen/blob/master/README.md +# Do not edit this file manually - changes will be overwritten + +name: Test k2ch + +on: + push: + paths: + - "packages/tkati-k2ch/**" + - "packages/tkati-core/**" + +jobs: + build: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + with: + fetch-depth: 0 + + - name: Install uv + uses: astral-sh/setup-uv@v7 + + - name: Build wheel + run: uv build + working-directory: packages/tkati-k2ch + + test: + runs-on: ubuntu-latest + + services: + redpanda: + image: redpandadata/redpanda:v25.2.1 + ports: + - 9092:9092 + - 9644:9644 + options: >- + --health-cmd "curl -f http://localhost:9644/v1/status/ready" + --health-interval 5s + --health-timeout 5s + --health-retries 10 + # clickhouse: + # image: clickhouse/clickhouse-server:latest + # ports: + # - 8123:8123 + # - 9000:9000 + # env: + # CLICKHOUSE_DB: default + # CLICKHOUSE_USER: default + # CLICKHOUSE_PASSWORD: default + # options: >- + # --health-cmd "curl -f http://localhost:8123/ping" + # --health-interval 5s + # --health-timeout 5s + # --health-retries 10 + + steps: + - name: Checkout code + uses: actions/checkout@v6 + + - name: Install the latest version of uv + uses: astral-sh/setup-uv@v6 + + - name: Install deps + run: uv sync --package k2ch + + - name: Check formatting + run: uvx ruff check packages/tkati-k2ch + + - name: Check typing + shell: bash + run: | + uv run mypy -p k2ch + + - name: Run tests + run: uv run pytest packages/tkati-k2ch \ No newline at end of file diff --git a/.github/workflows/test-tkati-core.yml b/.github/workflows/test-tkati-core.yml index b817ec0..d50e208 100644 --- a/.github/workflows/test-tkati-core.yml +++ b/.github/workflows/test-tkati-core.yml @@ -64,7 +64,12 @@ jobs: run: uv sync --package tkati-core - name: Check formatting - run: uv run ruff check packages/tkati-core + run: uvx ruff check packages/tkati-core + + - name: Check typing + shell: bash + run: | + uv run mypy -p tkati_core - name: Run tests run: uv run pytest packages/tkati-core \ No newline at end of file diff --git a/packages/tkati-k2ch/README.md b/packages/tkati-k2ch/README.md new file mode 100644 index 0000000..4adf8a6 --- /dev/null +++ b/packages/tkati-k2ch/README.md @@ -0,0 +1,52 @@ +# k2ch — Kafka to ClickHouse + +Reads batches of JSON messages from a Kafka topic and inserts them into a ClickHouse table using the native Arrow protocol. Offsets are committed only after a successful insert (at-least-once delivery). + +## Configuration + +Settings are loaded from a TOML file. Set the `SETTINGS_FILE` environment variable to point to it (defaults to `settings.toml`). + +```toml +[input.topic] +broker = "redpanda:29092" +name = "traffic_event" + +[input.topic.schema] +uid = "string" +time = "timestamp[ms]" +traffic_in = "uint32" +# … other columns + +[input.consumer] +group_id = "k2ch-group" +batch_size = 1000 +batch_timeout_sec = 10 +auto_offset_reset = "latest" + +[output] +host = "clickhouse" +port = 9000 +user = "default" +password = "" +database = "default" +table = "traffic_event" +secure = false + +[dlq] +topic = "k2ch-dlq" +# broker defaults to input.topic.broker when omitted +split_factor = 10 +``` + +## DLQ semantics + +When a batch insert fails after all retries, the app switches to a recursive fallback to isolate the problematic rows: + +1. The failing batch is split into `split_factor` equal sub-batches and each is retried independently. +2. If a sub-batch also fails it is split again — this repeats until individual rows are reached. +3. A single row that ClickHouse still rejects is written to the DLQ Kafka topic in Arrow IPC (`arrow-batch`) format, preserving the full schema. +4. After all rows are handled (inserted or DLQ'd), the Kafka offset is committed and the app resumes normal large-batch processing. + +With `split_factor=10` and a 1 000-row batch this takes at most 3 recursive levels (1000 → 100 → 10 → 1). + +**Delivery guarantee: at-least-once.** If the process crashes mid-recursion the uncommitted batch is re-read on restart and re-processed from the beginning, which may produce duplicate rows in ClickHouse and duplicate messages in the DLQ topic. diff --git a/packages/tkati-k2ch/pyproject.toml b/packages/tkati-k2ch/pyproject.toml new file mode 100644 index 0000000..ef3bdbc --- /dev/null +++ b/packages/tkati-k2ch/pyproject.toml @@ -0,0 +1,32 @@ +[project] +name = "k2ch" +version = "0.1.0" +description = "Add your description here" +readme = "README.md" +requires-python = ">=3.13" +dependencies = [ + "tkati-core", + "loguru>=0.7.0", + "pydantic-settings>=2.11.0", +] + +[project.scripts] +k2ch = "k2ch.main:main" + +[dependency-groups] +dev = ["pytest>=9.0.1", "confluent-kafka>=2.11.0"] + +[build-system] +requires = ["setuptools", "wheel"] +build-backend = "setuptools.build_meta" + +[tool.setuptools.packages.find] +where = ["src"] +include = ["k2ch*"] + +[tool.uv] +package = true + +[tool.uv-workspace-codegen] +generate = true +template_type = "test" diff --git a/packages/tkati-k2ch/settings.test.toml b/packages/tkati-k2ch/settings.test.toml new file mode 100644 index 0000000..45e4326 --- /dev/null +++ b/packages/tkati-k2ch/settings.test.toml @@ -0,0 +1,33 @@ +[input.topic] +broker = "redpanda:29092" +name = "traffic_event" + +[input.topic.schema] +uid = "string" +time = "timestamp[ms]" +package_id = "int32" +user_hash = "string" +sdk_hash = "string" +conn_type = "string" +country = "string" +local_ip = "string" +frontend_ip = "string" +dest_addr = "string" +client_ip = "string" +traffic_in = "uint32" +traffic_out = "uint32" + +[input.consumer] +group_id = "k2ch-group" +auto_offset_reset = "earliest" +batch_size = 1000 +batch_timeout_sec = 10 + +[output] +host = "clickhouse" +port = 8123 +user = "default" +password = "default" +database = "default" +table = "traffic_event" +secure = false diff --git a/packages/tkati-k2ch/src/k2ch/__init__.py b/packages/tkati-k2ch/src/k2ch/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/packages/tkati-k2ch/src/k2ch/__main__.py b/packages/tkati-k2ch/src/k2ch/__main__.py new file mode 100644 index 0000000..40e2b01 --- /dev/null +++ b/packages/tkati-k2ch/src/k2ch/__main__.py @@ -0,0 +1,4 @@ +from .main import main + +if __name__ == "__main__": + main() diff --git a/packages/tkati-k2ch/src/k2ch/main.py b/packages/tkati-k2ch/src/k2ch/main.py new file mode 100644 index 0000000..6a5f560 --- /dev/null +++ b/packages/tkati-k2ch/src/k2ch/main.py @@ -0,0 +1,46 @@ +from loguru import logger +from tkati_core.kafka.consumer import KafkaConsumer +from tkati_core.kafka.producer import KafkaProducer + +from tkati_core.clickhouse.producer import ClickhouseProducer +from k2ch.settings import AppSettings + + +def run_one_iteration( + kafka_consumer: KafkaConsumer, + ch_producer: ClickhouseProducer, + settings: AppSettings, +) -> None: + batch = kafka_consumer.read_arrow( + max_events_to_aggregate=settings.input.consumer.batch_size, + aggregation_interval_seconds=settings.input.consumer.batch_timeout_sec, + ) + if batch is None: + return + ch_producer.produce_arrow(batch) + kafka_consumer.commit() + logger.info(f"Inserted {len(batch)} rows to {ch_producer._table}") + + +def main() -> None: + settings = AppSettings() # type: ignore + + kafka_consumer = KafkaConsumer.from_input_settings(settings.input) + + dlq_producer: KafkaProducer | None = None + if settings.dlq is not None: + dlq_producer = KafkaProducer.from_topic_settings(settings.dlq.topic) + + ch_producer = ClickhouseProducer.from_output_settings( + settings=settings.output, + dlq_producer=dlq_producer, + split_factor=settings.dlq.split_factor if settings.dlq else 10, + ) + + try: + while True: + run_one_iteration(kafka_consumer, ch_producer, settings) + finally: + kafka_consumer.close() + if dlq_producer is not None: + dlq_producer.close() diff --git a/packages/tkati-k2ch/src/k2ch/py.typed b/packages/tkati-k2ch/src/k2ch/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/packages/tkati-k2ch/src/k2ch/settings.py b/packages/tkati-k2ch/src/k2ch/settings.py new file mode 100644 index 0000000..d2a841e --- /dev/null +++ b/packages/tkati-k2ch/src/k2ch/settings.py @@ -0,0 +1,15 @@ +from pydantic import BaseModel +from tkati_core.clickhouse.settings import ClickHouseOutputSettings +from tkati_core.kafka.settings import KafkaInputSettings, KafkaTopicSettings +from tkati_core.settings import TomlBaseSettings + + +class DLQSettings(BaseModel): + topic: KafkaTopicSettings + split_factor: int = 10 + + +class AppSettings(TomlBaseSettings): + input: KafkaInputSettings + output: ClickHouseOutputSettings + dlq: DLQSettings | None = None diff --git a/packages/tkati-k2ch/tests/conftest.py b/packages/tkati-k2ch/tests/conftest.py new file mode 100644 index 0000000..0caebe2 --- /dev/null +++ b/packages/tkati-k2ch/tests/conftest.py @@ -0,0 +1,96 @@ +import uuid +from collections.abc import Generator +from unittest.mock import MagicMock + +import pytest +from confluent_kafka import Producer +from confluent_kafka.admin import AdminClient, NewTopic +from k2ch.settings import AppSettings, ClickHouseOutputSettings, DLQSettings +from tkati_core.kafka.settings import ( + KafkaConsumerSettings, + KafkaInputSettings, + KafkaTopicSettings, +) +from tkati_core.kafka.testing import kafka_admin_client # noqa: F401 + + +@pytest.fixture(scope="function") +def test_settings() -> AppSettings: + run_id = str(uuid.uuid4())[:8] + return AppSettings( + input=KafkaInputSettings( + topic=KafkaTopicSettings( + broker="localhost:9092", + name=f"e2e_k2ch_{run_id}", + schema={ + "uid": "string", + "time": "timestamp[ms]", + "package_id": "int32", + "user_hash": "string", + "sdk_hash": "string", + "conn_type": "string", + "country": "string", + "local_ip": "string", + "frontend_ip": "string", + "dest_addr": "string", + "client_ip": "string", + "traffic_in": "uint32", + "traffic_out": "uint32", + }, + ), + consumer=KafkaConsumerSettings( + group_id=f"test-k2ch-{run_id}", + auto_offset_reset="earliest", + batch_size=100, + batch_timeout_sec=5, + ), + ), + output=ClickHouseOutputSettings( + host="localhost", + port=8123, + user="default", + password="", + database="default", + table="traffic_event", + secure=False, + ), + dlq=DLQSettings( + topic=KafkaTopicSettings(broker="localhost:9092", name="k2ch-dlq"), + split_factor=2, + ), + ) + + +@pytest.fixture(scope="function") +def mock_ch_client() -> MagicMock: + """Mock for clickhouse_connect client. insert_arrow is the only called method.""" + client = MagicMock() + client.insert_arrow = MagicMock() + return client + + +@pytest.fixture(scope="function") +def mock_dlq_producer() -> MagicMock: + producer = MagicMock() + producer.produce_arrow = MagicMock() + producer.flush = MagicMock() + return producer + + +@pytest.fixture(scope="function") +def kafka_producer_and_topic( + test_settings: AppSettings, + kafka_admin_client: AdminClient, # noqa: F811 +) -> Generator[Producer, None, None]: + """Creates the input topic and yields a Producer for it.""" + topic = test_settings.input.topic.name + fs = kafka_admin_client.create_topics( + [NewTopic(topic, num_partitions=1, replication_factor=1)] + ) + for t, f in fs.items(): + try: + f.result() + except Exception as e: + print(f"Failed to create topic {t}: {e}") + + yield Producer({"bootstrap.servers": test_settings.input.topic.broker}) diff --git a/packages/tkati-k2ch/tests/test_k2ch.py b/packages/tkati-k2ch/tests/test_k2ch.py new file mode 100644 index 0000000..1b0352e --- /dev/null +++ b/packages/tkati-k2ch/tests/test_k2ch.py @@ -0,0 +1,104 @@ +import time +from unittest.mock import MagicMock + +import orjson +import pyarrow as pa +import pytest +from confluent_kafka import Producer +from k2ch.main import run_one_iteration +from k2ch.settings import AppSettings +from tkati_core.clickhouse.producer import ClickhouseProducer +from tkati_core.kafka.consumer import KafkaConsumer + + +def _make_consumer(test_settings: AppSettings) -> KafkaConsumer: + return KafkaConsumer( + kafka_config={ + "bootstrap.servers": test_settings.input.topic.broker, + "group.id": test_settings.input.consumer.group_id, + "auto.offset.reset": "earliest", + "enable.auto.commit": False, + }, + topic_name=test_settings.input.topic.name, + input_schema=test_settings.input.topic.schema, + ) + + +def _make_ch_producer(mock_ch_client: MagicMock, mock_dlq_producer: MagicMock, table: str = "traffic_event") -> ClickhouseProducer: + return ClickhouseProducer(ch_client=mock_ch_client, table=table, dlq_producer=mock_dlq_producer) + + +def test_k2ch_valid_flow( + kafka_producer_and_topic: Producer, + mock_ch_client: MagicMock, + mock_dlq_producer: MagicMock, + test_settings: AppSettings, +) -> None: + """Produce a valid event to Kafka, run one iteration, verify CH insert was called.""" + event = { + "uid": "abc123", + "time": int(time.time() * 1000), + "package_id": 1, + "user_hash": "uhash", + "sdk_hash": "shash", + "conn_type": "https", + "country": "US", + "local_ip": "10.0.0.1", + "frontend_ip": "1.2.3.4", + "dest_addr": "8.8.8.8", + "client_ip": "192.168.1.1", + "traffic_in": 100, + "traffic_out": 200, + } + + kafka_producer_and_topic.produce( + test_settings.input.topic.name, value=orjson.dumps(event) + ) + kafka_producer_and_topic.flush() + + consumer = _make_consumer(test_settings) + ch_producer = ClickhouseProducer( + ch_client=mock_ch_client, + table=test_settings.output.table, + dlq_producer=mock_dlq_producer, + ) + try: + run_one_iteration(consumer, ch_producer, test_settings) + finally: + consumer.close() + + mock_ch_client.insert_arrow.assert_called_once() + _, kwargs = mock_ch_client.insert_arrow.call_args + assert kwargs["table"] == test_settings.output.table + result_table = kwargs["arrow_table"] + assert isinstance(result_table, pa.Table) + assert len(result_table) == 1 + assert result_table.column("uid")[0].as_py() == "abc123" + mock_dlq_producer.produce_arrow.assert_not_called() + + +def test_k2ch_malformed_data( + kafka_producer_and_topic: Producer, + mock_ch_client: MagicMock, + mock_dlq_producer: MagicMock, + test_settings: AppSettings, +) -> None: + """Produce malformed JSON to Kafka; run_one_iteration must raise with 'JSON parse error'.""" + kafka_producer_and_topic.produce( + test_settings.input.topic.name, value=b"not a json object" + ) + kafka_producer_and_topic.flush() + + consumer = _make_consumer(test_settings) + ch_producer = ClickhouseProducer( + ch_client=mock_ch_client, + table=test_settings.output.table, + dlq_producer=mock_dlq_producer, + ) + try: + with pytest.raises(Exception, match="JSON parse error"): + run_one_iteration(consumer, ch_producer, test_settings) + finally: + consumer.close() + + mock_ch_client.insert_arrow.assert_not_called() diff --git a/pyproject.toml b/pyproject.toml index e1e731c..4d566fc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,12 +11,14 @@ members = ["packages/*"] [tool.uv.sources] tkati-core = { workspace = true } +tkati-k2ch = { workspace = true } [tool.pyright] reportPrivateImportUsage = "none" [dependency-groups] dev = [ + "mypy>=2.2.0", "pytest>=9.0.3", "ruff>=0.15.17", ]