From da16babca99b6eeaf52f2a8048ab1174fcb341ab Mon Sep 17 00:00:00 2001 From: Andrey Tatarinov Date: Wed, 8 Jul 2026 09:18:24 +0400 Subject: [PATCH 1/3] Add Producer interface and allow ClickhouseProducer.dlq_producer to write to CH --- packages/tkati-core/CHANGELOG.md | 6 +++ packages/tkati-core/pyproject.toml | 2 +- packages/tkati-core/tests/test_ch_producer.py | 46 +++++++++++++++++++ .../tkati_core/clickhouse/producer.py | 27 +++++++---- .../tkati-core/tkati_core/kafka/producer.py | 4 +- packages/tkati-core/tkati_core/producer.py | 21 +++++++++ 6 files changed, 95 insertions(+), 11 deletions(-) create mode 100644 packages/tkati-core/tkati_core/producer.py diff --git a/packages/tkati-core/CHANGELOG.md b/packages/tkati-core/CHANGELOG.md index 5ced7d4..eca9f5b 100644 --- a/packages/tkati-core/CHANGELOG.md +++ b/packages/tkati-core/CHANGELOG.md @@ -1,3 +1,9 @@ +# 0.3.0 + +* Add shared `Producer` base class implemented by `KafkaProducer` and `ClickhouseProducer` +* `ClickhouseProducer` now supports `produce_pylist`, `flush`, and `close`, and + can be used as `dlq_producer` for another `ClickhouseProducer` + # 0.2.0 * Initial implementation of ClickhouseProducer diff --git a/packages/tkati-core/pyproject.toml b/packages/tkati-core/pyproject.toml index 8109f4b..feb6299 100644 --- a/packages/tkati-core/pyproject.toml +++ b/packages/tkati-core/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "tkati-core" -version = "0.2.0.post3" +version = "0.3.0" description = "Core abstractions for Kafka/ClickHouse streaming pipeline nodes" readme = "README.md" requires-python = ">=3.13" diff --git a/packages/tkati-core/tests/test_ch_producer.py b/packages/tkati-core/tests/test_ch_producer.py index 0352093..ebaf6b4 100644 --- a/packages/tkati-core/tests/test_ch_producer.py +++ b/packages/tkati-core/tests/test_ch_producer.py @@ -164,3 +164,49 @@ def test_ch_producer_failure_without_dlq_raises() -> None: with patch("time.sleep"): with pytest.raises(Exception, match="CH down"): producer.produce_arrow(arrow_table) + + +def test_ch_producer_as_dlq_for_another_ch_producer() -> None: + """A ClickhouseProducer can itself be used as the dlq_producer for another ClickhouseProducer.""" + primary_ch_client = MagicMock() + dlq_ch_client = MagicMock() + arrow_table = _make_arrow_table(1) + + primary_ch_client.insert_arrow.side_effect = Exception("CH down") + + dlq_producer = ClickhouseProducer(ch_client=dlq_ch_client, table="traffic_event_dlq") + producer = ClickhouseProducer( + ch_client=primary_ch_client, table="traffic_event", dlq_producer=dlq_producer + ) + + with patch("time.sleep"): + producer.produce_arrow(arrow_table) + + dlq_ch_client.insert_arrow.assert_called_once() + + +def test_ch_producer_flush_is_noop() -> None: + """flush() on ClickhouseProducer is a no-op and never touches ch_client.""" + ch_client = MagicMock() + ClickhouseProducer(ch_client=ch_client, table="traffic_event").flush() + ch_client.assert_not_called() + + +def test_ch_producer_produce_pylist() -> None: + """produce_pylist converts rows to an Arrow table and inserts them.""" + ch_client = MagicMock() + rows = [{"uid": "uid-0", "traffic_in": 100}, {"uid": "uid-1", "traffic_in": 101}] + + producer = ClickhouseProducer(ch_client=ch_client, table="traffic_event") + producer.produce_pylist(rows) + + ch_client.insert_arrow.assert_called_once() + sent = ch_client.insert_arrow.call_args.kwargs["arrow_table"] + assert sent.to_pylist() == rows + + +def test_ch_producer_close() -> None: + """close() delegates to the underlying ch_client.""" + ch_client = MagicMock() + ClickhouseProducer(ch_client=ch_client, table="traffic_event").close() + ch_client.close.assert_called_once() diff --git a/packages/tkati-core/tkati_core/clickhouse/producer.py b/packages/tkati-core/tkati_core/clickhouse/producer.py index 03dff6d..0025121 100644 --- a/packages/tkati-core/tkati_core/clickhouse/producer.py +++ b/packages/tkati-core/tkati_core/clickhouse/producer.py @@ -3,9 +3,9 @@ import pyarrow as pa from loguru import logger from tenacity import RetryCallState, retry, stop_after_attempt, wait_fixed -from tkati_core.kafka.producer import KafkaProducer from tkati_core.clickhouse.settings import ClickHouseOutputSettings +from tkati_core.producer import Producer def log_retry_attempt(retry_state: RetryCallState) -> None: @@ -35,7 +35,7 @@ def _insert_with_dlq_fallback( table: pa.Table, ch_client: ch_driver.Client, ch_table: str, - dlq_producer: KafkaProducer, + dlq_producer: Producer, split_factor: int, ) -> None: try: @@ -55,12 +55,12 @@ def _insert_with_dlq_fallback( _insert_with_dlq_fallback(chunk, ch_client, ch_table, dlq_producer, split_factor) -class ClickhouseProducer: +class ClickhouseProducer(Producer): def __init__( self, ch_client: ch_driver.Client, table: str, - dlq_producer: KafkaProducer | None = None, + dlq_producer: Producer | None = None, split_factor: int = 10, ) -> None: self._ch_client = ch_client @@ -72,7 +72,7 @@ def __init__( def from_output_settings( cls, settings: ClickHouseOutputSettings, - dlq_producer: KafkaProducer | None = None, + dlq_producer: Producer | None = None, split_factor: int = 10, ) -> "ClickhouseProducer": ch_client = ch.get_client( @@ -85,21 +85,30 @@ def from_output_settings( ) return cls(ch_client=ch_client, table=settings.table, dlq_producer=dlq_producer, split_factor=split_factor) - def produce_arrow(self, arrow_table: pa.Table) -> None: + def produce_arrow(self, data: pa.Table) -> None: try: - _insert_with_retry(ch_client=self._ch_client, table=self._table, arrow_table=arrow_table) + _insert_with_retry(ch_client=self._ch_client, table=self._table, arrow_table=data) except Exception as err: if self._dlq_producer is None: raise logger.warning( - f"Batch of {len(arrow_table)} rows failed ({err}), " + f"Batch of {len(data)} rows failed ({err}), " f"switching to recursive fallback with split_factor={self._split_factor}" ) _insert_with_dlq_fallback( - table=arrow_table, + table=data, ch_client=self._ch_client, ch_table=self._table, dlq_producer=self._dlq_producer, split_factor=self._split_factor, ) self._dlq_producer.flush() + + def produce_pylist(self, rows: list[dict]) -> None: + self.produce_arrow(pa.Table.from_pylist(rows)) + + def flush(self) -> None: + """No-op: ClickHouse inserts are synchronous, nothing to flush.""" + + def close(self) -> None: + self._ch_client.close() diff --git a/packages/tkati-core/tkati_core/kafka/producer.py b/packages/tkati-core/tkati_core/kafka/producer.py index b9e6613..9d117fb 100644 --- a/packages/tkati-core/tkati_core/kafka/producer.py +++ b/packages/tkati-core/tkati_core/kafka/producer.py @@ -7,11 +7,13 @@ from confluent_kafka import Producer from loguru import logger +from tkati_core.producer import Producer as ProducerBase + if TYPE_CHECKING: from tkati_core.kafka.settings import KafkaOutputSettings, KafkaTopicSettings -class KafkaProducer: +class KafkaProducer(ProducerBase): """ A Kafka producer wrapper that writes data as messages. diff --git a/packages/tkati-core/tkati_core/producer.py b/packages/tkati-core/tkati_core/producer.py new file mode 100644 index 0000000..91891cc --- /dev/null +++ b/packages/tkati-core/tkati_core/producer.py @@ -0,0 +1,21 @@ +"""Shared producer interface, implemented by KafkaProducer and ClickhouseProducer.""" + +from abc import ABC, abstractmethod + +import pyarrow as pa + + +class Producer(ABC): + """Base class for anything that can act as a producer target.""" + + @abstractmethod + def produce_arrow(self, data: pa.Table) -> None: ... + + @abstractmethod + def produce_pylist(self, rows: list[dict]) -> None: ... + + @abstractmethod + def flush(self) -> None: ... + + @abstractmethod + def close(self) -> None: ... From edba05a94953825ae772bf2034376a7360f10199 Mon Sep 17 00:00:00 2001 From: Andrey Tatarinov Date: Wed, 8 Jul 2026 11:41:59 +0400 Subject: [PATCH 2/3] v0.3.0-alpha.1 --- packages/tkati-core/pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/tkati-core/pyproject.toml b/packages/tkati-core/pyproject.toml index feb6299..5aedc63 100644 --- a/packages/tkati-core/pyproject.toml +++ b/packages/tkati-core/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "tkati-core" -version = "0.3.0" +version = "0.3.0-alpha.1" description = "Core abstractions for Kafka/ClickHouse streaming pipeline nodes" readme = "README.md" requires-python = ">=3.13" From d08e9da48379564af3b25a073bb88a5ec2f894e3 Mon Sep 17 00:00:00 2001 From: Andrey Tatarinov Date: Wed, 8 Jul 2026 11:40:41 +0400 Subject: [PATCH 3/3] Update workflow name --- .github/workflow-templates/test.template.yml | 2 +- .github/workflows/test-tkati-core.yml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflow-templates/test.template.yml b/.github/workflow-templates/test.template.yml index 60fc777..4ee3071 100644 --- a/.github/workflow-templates/test.template.yml +++ b/.github/workflow-templates/test.template.yml @@ -1,4 +1,4 @@ -name: Package {{ package.name }} +name: Test {{ package.name }} on: push: diff --git a/.github/workflows/test-tkati-core.yml b/.github/workflows/test-tkati-core.yml index e59bac1..b817ec0 100644 --- a/.github/workflows/test-tkati-core.yml +++ b/.github/workflows/test-tkati-core.yml @@ -2,7 +2,7 @@ # 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: Package tkati-core +name: Test tkati-core on: push: