Add Producer interface and allow ClickhouseProducer.dlq_producer to write to CH

.github/workflow-templates/test.template.yml

@@ -1,4 +1,4 @@
-name: Package {{ package.name }}
+name: Test {{ package.name }}
 
 on:
   push:

.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:

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

packages/tkati-core/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "tkati-core"
-version = "0.2.0.post3"
+version = "0.3.0-alpha.1"
 description = "Core abstractions for Kafka/ClickHouse streaming pipeline nodes"
 readme = "README.md"
 requires-python = ">=3.13"

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()

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()

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.
 

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: ...