Skip to content

Commit 1fc2fc1

Browse files
committed
Kafka: Expand default_msg_processor into a miniature decoding unit
- Accept a bunch of decoding options per `KafkaDecodingOptions` - Provide a bunch of output formatting options per `KafkaEvent` - Tie both elements together using `KafkaEventProcessor` The machinery is effectively the same like before, but provides a few more options to allow type decoding for Kafka event's key/value slots, a selection mechanism to limit the output to specific fields only, and a small projection mechanism to optionally drill down into a specific field. In combination, those decoding options allow users to relay JSON-encoded Kafka event values directly into a destination table, without any metadata wrappings. The output formatter provides three different variants out of the box. More variants can be added in the future, as other users or use cases may have different requirements in the same area. Most importantly, the decoding unit is very compact, so relevant tasks don't need a corresponding transformation unit down the pipeline, to keep the whole ensemble lean, in the very spirit of ingestr.
1 parent 0c6e52f commit 1fc2fc1

9 files changed

Lines changed: 463 additions & 59 deletions

File tree

docs/supported-sources/kafka.md

Lines changed: 57 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,29 +10,81 @@ The URI format for Apache Kafka is as follows:
1010
kafka://?bootstrap_servers=localhost:9092&group_id=test_group&security_protocol=SASL_SSL&sasl_mechanisms=PLAIN&sasl_username=example_username&sasl_password=example_secret&batch_size=1000&batch_timeout=3
1111
```
1212

13-
URI parameters:
13+
### URI parameters
14+
15+
Connectivity options:
1416
- `bootstrap_servers`: Required, the Kafka server or servers to connect to, typically in the form of a host and port, e.g. `localhost:9092`
1517
- `group_id`: Required, the consumer group ID used for identifying the client when consuming messages.
1618
- `security_protocol`: The protocol used to communicate with brokers, e.g. `SASL_SSL` for secure communication.
1719
- `sasl_mechanisms`: The SASL mechanism to be used for authentication, e.g. `PLAIN`.
1820
- `sasl_username`: The username for SASL authentication.
1921
- `sasl_password`: The password for SASL authentication.
22+
23+
Transfer options:
2024
- `batch_size`: The number of messages to fetch in a single batch, defaults to 3000.
2125
- `batch_timeout`: The maximum time to wait for messages, defaults to 3 seconds.
2226

27+
Decoding options:
28+
- `key_type`: The data type of the Kafka event `key` field. Possible values: `json`.
29+
- `value_type`: The data type of the Kafka event `value_type` field. Possible values: `json`.
30+
- `include`: A list of event attributes to include in the output, comma-separated.
31+
- `select`: A single event attribute to select and drill down into.
32+
Use `select=value` to relay the Kafka event **payload data** only.
33+
- `format`: The output format/layout. Possible values: `standard_v1` (default),
34+
`standard_v2`, `flexible`. When using the `include` or `select` option, the
35+
decoder will automatically select the `flexible` output format.
36+
2337
The URI is used to connect to the Kafka brokers for ingesting messages.
2438

2539
### Group ID
2640
The group ID is used to identify the consumer group that reads messages from a topic. Kafka uses the group ID to manage consumer offsets and assign partitions to consumers, which means that the group ID is the key to reading messages from the correct partition and position in the topic.
2741

28-
Once you have your Kafka server, credentials, and group ID set up, here's a sample command to ingest messages from a Kafka topic into a DuckDB database:
42+
## Examples
43+
44+
### Kafka to DuckDB
45+
46+
Once you have your Kafka server, credentials, and group ID set up,
47+
here are a few sample commands to ingest messages from a Kafka topic into a destination database:
2948

49+
Transfer data using the traditional `standard_v1` output format into DuckDB.
50+
The result of this command will be a table in the `kafka.duckdb` database with JSON columns.
3051
```sh
3152
ingestr ingest \
32-
--source-uri 'kafka://?bootstrap_servers=localhost:9092&group_id=test_group' \
53+
--source-uri 'kafka://?bootstrap_servers=localhost:9092&group_id=test' \
3354
--source-table 'my-topic' \
34-
--dest-uri duckdb:///kafka.duckdb \
55+
--dest-uri 'duckdb:///kafka.duckdb' \
3556
--dest-table 'dest.my_topic'
3657
```
3758

38-
The result of this command will be a table in the `kafka.duckdb` database with JSON columns.
59+
### Kafka to PostgreSQL
60+
61+
Transfer data converging the Kafka event `value` into a PostgreSQL destination
62+
table, after decoding from JSON, using the `flexible` output format.
63+
```sh
64+
echo '{"sensor_id":1,"ts":"2025-06-01 10:00","reading":42.42}' | kcat -P -b localhost -t demo
65+
```
66+
```sh
67+
ingestr ingest \
68+
--source-uri 'kafka://?bootstrap_servers=localhost:9092&group_id=test&value_type=json&select=value' \
69+
--source-table 'demo' \
70+
--dest-uri 'postgres://postgres:postgres@localhost:5432/?sslmode=disable' \
71+
--dest-table 'public.kafka_demo'
72+
```
73+
The result of this command will be the `public.kafka_demo` table using
74+
the Kafka event `value`'s top-level JSON keys as table columns.
75+
```sh
76+
psql "postgresql://postgres:postgres@localhost:5432/" \
77+
-c '\d+ public.kafka_demo' \
78+
-c 'select * from public.kafka_demo;'
79+
```
80+
```text
81+
Table "public.kafka_demo"
82+
83+
Column | Type |
84+
--------------+--------------------------+
85+
sensor_id | bigint |
86+
ts | timestamp with time zone |
87+
reading | double precision |
88+
_dlt_load_id | character varying |
89+
_dlt_id | character varying |
90+
```

ingestr/main_test.py

Lines changed: 141 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
import requests
3232
import sqlalchemy
3333
from confluent_kafka import Producer # type: ignore
34+
from confluent_kafka.admin import AdminClient # type: ignore
3435
from dlt.sources.filesystem import glob_files
3536
from elasticsearch import Elasticsearch
3637
from fsspec.implementations.memory import MemoryFileSystem # type: ignore
@@ -1510,19 +1511,39 @@ def as_datetime2(date_str: str) -> datetime:
15101511
return datetime.strptime(date_str, "%Y-%m-%d")
15111512

15121513

1514+
@pytest.fixture(scope="session")
1515+
def kafka_service():
1516+
"""
1517+
Provide a Kafka service container for the whole test session.
1518+
"""
1519+
container = KafkaContainer("confluentinc/cp-kafka:7.6.0")
1520+
container.start()
1521+
yield container
1522+
container.stop()
1523+
1524+
1525+
@pytest.fixture(scope="function")
1526+
def kafka(kafka_service):
1527+
"""
1528+
Provide a Kafka service container using a clean canvas.
1529+
Before invoking the test case, delete all relevant topics completely.
1530+
"""
1531+
admin = AdminClient({"bootstrap.servers": kafka_service.get_bootstrap_server()})
1532+
admin.delete_topics(["test_topic"])
1533+
admin.poll(1)
1534+
return kafka_service
1535+
1536+
15131537
@pytest.mark.parametrize(
15141538
"dest", list(DESTINATIONS.values()), ids=list(DESTINATIONS.keys())
15151539
)
1516-
def test_kafka_to_db(dest):
1540+
def test_kafka_to_db_incremental(kafka, dest):
1541+
"""
1542+
Validate standard Kafka event decoding, focusing on both metadata and data payload.
1543+
"""
15171544
with ThreadPoolExecutor() as executor:
15181545
dest_future = executor.submit(dest.start)
1519-
source_future = executor.submit(
1520-
KafkaContainer("confluentinc/cp-kafka:7.6.0").start, timeout=120
1521-
)
15221546
dest_uri = dest_future.result()
1523-
kafka = source_future.result()
1524-
1525-
# kafka = KafkaContainer("confluentinc/cp-kafka:7.6.0").start(timeout=60)
15261547

15271548
# Create Kafka producer
15281549
producer = Producer({"bootstrap.servers": kafka.get_bootstrap_server()})
@@ -1583,7 +1604,119 @@ def get_output_table():
15831604
assert res[2] == ("message3",)
15841605
assert res[3] == ("message4",)
15851606

1586-
kafka.stop()
1607+
1608+
@pytest.mark.parametrize(
1609+
"dest", list(DESTINATIONS.values()), ids=list(DESTINATIONS.keys())
1610+
)
1611+
def test_kafka_to_db_decode_json(kafka, dest):
1612+
"""
1613+
Validate slightly more advanced Kafka event decoding, focusing on the payload value this time.
1614+
1615+
This exercise uses the `value_type=json` and `select=value` URL parameters.
1616+
"""
1617+
with ThreadPoolExecutor() as executor:
1618+
dest_future = executor.submit(dest.start)
1619+
dest_uri = dest_future.result()
1620+
1621+
# Create Kafka producer
1622+
producer = Producer({"bootstrap.servers": kafka.get_bootstrap_server()})
1623+
1624+
# Create topic and send messages
1625+
topic = "test_topic"
1626+
messages = [
1627+
{"id": 1, "temperature": 42.42, "humidity": 82},
1628+
{"id": 2, "temperature": 451.00, "humidity": 15},
1629+
]
1630+
1631+
for message in messages:
1632+
producer.produce(topic, json.dumps(message))
1633+
producer.flush()
1634+
1635+
def run():
1636+
res = invoke_ingest_command(
1637+
f"kafka://?bootstrap_servers={kafka.get_bootstrap_server()}&group_id=test_group&value_type=json&select=value",
1638+
"test_topic",
1639+
dest_uri,
1640+
"testschema.output",
1641+
)
1642+
assert res.exit_code == 0
1643+
1644+
def get_output_table():
1645+
dest_engine = sqlalchemy.create_engine(dest_uri)
1646+
with dest_engine.connect() as conn:
1647+
res = (
1648+
conn.execute(
1649+
"SELECT id, temperature, humidity FROM testschema.output WHERE temperature >= 38.00 ORDER BY id ASC"
1650+
)
1651+
.mappings()
1652+
.fetchall()
1653+
)
1654+
dest_engine.dispose()
1655+
return res
1656+
1657+
run()
1658+
1659+
res = get_output_table()
1660+
assert len(res) == 2
1661+
assert res[0] == messages[0]
1662+
assert res[1] == messages[1]
1663+
1664+
1665+
@pytest.mark.parametrize(
1666+
"dest", list(DESTINATIONS.values()), ids=list(DESTINATIONS.keys())
1667+
)
1668+
def test_kafka_to_db_include_metadata(kafka, dest):
1669+
"""
1670+
Validate slightly more advanced Kafka event decoding, focusing on metadata this time.
1671+
1672+
This exercise uses the `include=` URL parameter.
1673+
"""
1674+
with ThreadPoolExecutor() as executor:
1675+
dest_future = executor.submit(dest.start)
1676+
dest_uri = dest_future.result()
1677+
1678+
# Create Kafka producer
1679+
producer = Producer({"bootstrap.servers": kafka.get_bootstrap_server()})
1680+
1681+
# Create topic and send messages
1682+
topic = "test_topic"
1683+
messages = [
1684+
{"id": 1, "temperature": 42.42, "humidity": 82},
1685+
{"id": 2, "temperature": 451.00, "humidity": 15},
1686+
]
1687+
1688+
for message in messages:
1689+
producer.produce(topic=topic, value=json.dumps(message), key="test")
1690+
producer.flush()
1691+
1692+
def run():
1693+
res = invoke_ingest_command(
1694+
f"kafka://?bootstrap_servers={kafka.get_bootstrap_server()}&group_id=test_group&include=partition,topic,key,offset,ts",
1695+
"test_topic",
1696+
dest_uri,
1697+
"testschema.output",
1698+
)
1699+
assert res.exit_code == 0
1700+
1701+
def get_output_table():
1702+
dest_engine = sqlalchemy.create_engine(dest_uri)
1703+
with dest_engine.connect() as conn:
1704+
res = (
1705+
conn.execute(
1706+
'SELECT "partition", "topic", "key", "offset" FROM testschema.output ORDER BY "ts__value" ASC'
1707+
)
1708+
.mappings()
1709+
.fetchall()
1710+
)
1711+
dest_engine.dispose()
1712+
return res
1713+
1714+
run()
1715+
1716+
res = get_output_table()
1717+
assert len(res) == 2
1718+
assert res[0] == {"partition": 0, "topic": "test_topic", "key": "test", "offset": 0}
1719+
assert res[1] == {"partition": 0, "topic": "test_topic", "key": "test", "offset": 1}
15871720

15881721

15891722
@pytest.mark.parametrize(

ingestr/src/kafka/__init__.py

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414

1515
"""A source to extract Kafka messages.
1616
17-
When extraction starts, partitions length is checked -
17+
When extraction starts, partition length is checked -
1818
data is read only up to it, overriding the default Kafka's
1919
behavior of waiting for new messages in endless loop.
2020
"""
@@ -30,8 +30,8 @@
3030

3131
from .helpers import (
3232
KafkaCredentials,
33+
KafkaEventProcessor,
3334
OffsetTracker,
34-
default_msg_processor,
3535
)
3636

3737

@@ -43,9 +43,7 @@
4343
def kafka_consumer(
4444
topics: Union[str, List[str]],
4545
credentials: Union[KafkaCredentials, Consumer] = dlt.secrets.value,
46-
msg_processor: Optional[
47-
Callable[[Message], Dict[str, Any]]
48-
] = default_msg_processor,
46+
msg_processor: Optional[Callable[[Message], Dict[str, Any]]] = None,
4947
batch_size: Optional[int] = 3000,
5048
batch_timeout: Optional[int] = 3,
5149
start_from: Optional[TAnyDateTime] = None,
@@ -63,17 +61,19 @@ def kafka_consumer(
6361
Auth credentials or an initiated Kafka consumer. By default,
6462
is taken from secrets.
6563
msg_processor(Optional[Callable]): A function-converter,
66-
which'll process every Kafka message after it's read and
67-
before it's transfered to the destination.
64+
which will process every Kafka message after it is read and
65+
before it is transferred to the destination.
6866
batch_size (Optional[int]): Messages batch size to read at once.
6967
batch_timeout (Optional[int]): Maximum time to wait for a batch
70-
consume, in seconds.
68+
to be consumed in seconds.
7169
start_from (Optional[TAnyDateTime]): A timestamp, at which to start
7270
reading. Older messages are ignored.
7371
7472
Yields:
7573
Iterable[TDataItem]: Kafka messages.
7674
"""
75+
msg_processor = msg_processor or KafkaEventProcessor().process
76+
7777
if not isinstance(topics, list):
7878
topics = [topics]
7979

0 commit comments

Comments
 (0)