sqllocks-shape-kafka 0.9.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (27) hide show
  1. sqllocks_shape_kafka-0.9.1/LICENSE +21 -0
  2. sqllocks_shape_kafka-0.9.1/PKG-INFO +125 -0
  3. sqllocks_shape_kafka-0.9.1/README.md +97 -0
  4. sqllocks_shape_kafka-0.9.1/pyproject.toml +43 -0
  5. sqllocks_shape_kafka-0.9.1/setup.cfg +4 -0
  6. sqllocks_shape_kafka-0.9.1/src/shape_kafka/__init__.py +17 -0
  7. sqllocks_shape_kafka-0.9.1/src/shape_kafka/emitter.py +361 -0
  8. sqllocks_shape_kafka-0.9.1/src/shape_kafka/formats.py +546 -0
  9. sqllocks_shape_kafka-0.9.1/src/shape_kafka/registry.py +130 -0
  10. sqllocks_shape_kafka-0.9.1/src/shape_kafka/source.py +264 -0
  11. sqllocks_shape_kafka-0.9.1/src/shape_kafka/testing.py +486 -0
  12. sqllocks_shape_kafka-0.9.1/src/sqllocks_shape_kafka.egg-info/PKG-INFO +125 -0
  13. sqllocks_shape_kafka-0.9.1/src/sqllocks_shape_kafka.egg-info/SOURCES.txt +25 -0
  14. sqllocks_shape_kafka-0.9.1/src/sqllocks_shape_kafka.egg-info/dependency_links.txt +1 -0
  15. sqllocks_shape_kafka-0.9.1/src/sqllocks_shape_kafka.egg-info/entry_points.txt +5 -0
  16. sqllocks_shape_kafka-0.9.1/src/sqllocks_shape_kafka.egg-info/requires.txt +8 -0
  17. sqllocks_shape_kafka-0.9.1/src/sqllocks_shape_kafka.egg-info/top_level.txt +1 -0
  18. sqllocks_shape_kafka-0.9.1/tests/test_cli.py +106 -0
  19. sqllocks_shape_kafka-0.9.1/tests/test_dead_letter.py +214 -0
  20. sqllocks_shape_kafka-0.9.1/tests/test_emitter.py +138 -0
  21. sqllocks_shape_kafka-0.9.1/tests/test_emitter_emulator.py +180 -0
  22. sqllocks_shape_kafka-0.9.1/tests/test_emulator.py +226 -0
  23. sqllocks_shape_kafka-0.9.1/tests/test_live.py +59 -0
  24. sqllocks_shape_kafka-0.9.1/tests/test_registry_formats.py +676 -0
  25. sqllocks_shape_kafka-0.9.1/tests/test_skeleton.py +9 -0
  26. sqllocks_shape_kafka-0.9.1/tests/test_source.py +319 -0
  27. sqllocks_shape_kafka-0.9.1/tests/test_source_bounds.py +19 -0
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2025-2026 SQLLocks (Jonathan Stewart)
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,125 @@
1
+ Metadata-Version: 2.4
2
+ Name: sqllocks-shape-kafka
3
+ Version: 0.9.1
4
+ Summary: Shape plugin: Kafka stream source and event emitter
5
+ License-Expression: MIT
6
+ Project-URL: Homepage, https://github.com/sqllocks/shape
7
+ Project-URL: Repository, https://github.com/sqllocks/shape
8
+ Project-URL: Issues, https://github.com/sqllocks/shape/issues
9
+ Project-URL: Documentation, https://github.com/sqllocks/shape/tree/main/plugins/shape-kafka
10
+ Classifier: Programming Language :: Python :: 3
11
+ Classifier: Programming Language :: Python :: 3 :: Only
12
+ Classifier: Programming Language :: Python :: 3.11
13
+ Classifier: Programming Language :: Python :: 3.12
14
+ Classifier: Programming Language :: Python :: 3.13
15
+ Classifier: Programming Language :: Python :: 3.14
16
+ Classifier: Operating System :: OS Independent
17
+ Classifier: Topic :: Database
18
+ Requires-Python: >=3.11
19
+ Description-Content-Type: text/markdown
20
+ License-File: LICENSE
21
+ Requires-Dist: sqllocks-shape==0.9.1
22
+ Requires-Dist: confluent-kafka<4,>=2.5
23
+ Provides-Extra: avro
24
+ Requires-Dist: fastavro>=1.9; extra == "avro"
25
+ Provides-Extra: protobuf
26
+ Requires-Dist: protobuf>=5; extra == "protobuf"
27
+ Dynamic: license-file
28
+
29
+ # sqllocks-shape-kafka
30
+
31
+ Shape plugin: a Kafka stream source (`shape.stream_sources`, scheme `kafka://`).
32
+
33
+ ```
34
+ pip install sqllocks-shape-kafka # or: pip install 'sqllocks-shape[kafka]'
35
+ shape stream-profile kafka://broker:9092/orders -o orders.json
36
+ ```
37
+
38
+ On Alpine/musllinux, `confluent-kafka` has no ARM64 or x86_64 wheels and builds from
39
+ source. Installing this plugin, `sqllocks-shape[kafka]` or `sqllocks-shape[all]` needs
40
+ librdkafka **2.16 or newer**, development headers and C build tools. Alpine's
41
+ `librdkafka-dev` is currently older. Build the library first using the
42
+ [Alpine installation steps](../../docs/INSTALL.md#alpine-and-musllinux), which match
43
+ `scripts/release_smoke_alpine.sh`, or
44
+ [Confluent's build instructions](https://github.com/confluentinc/librdkafka#building).
45
+
46
+ `kafka://host1:9092,host2:9092/TOPIC` reads one topic as Arrow micro-batches. Message bodies are
47
+ JSON objects; their fields become columns, and `_shape_event_time` carries the event time (the
48
+ body's `_shape_event_time`, or the message timestamp). Shape keeps its own position: the consumer
49
+ is *assigned* every partition at explicit offsets, nothing is committed to the broker, and a
50
+ `StreamOffset` is `{partition: next offset}`, so a checkpoint resumes exactly where it stopped.
51
+
52
+ Options, offsets, delivery guarantees and the checkpoint format are described in
53
+ `docs/plugins/streaming.md`. Extra `confluent-kafka` settings (SASL, TLS) go in through the
54
+ `config` source option, preferably from a file (`--options-file`).
55
+
56
+ ## Emitting events: `kafka://`
57
+
58
+ ```
59
+ shape emit retail --realtime --rate 5000 --sink kafka://broker:9092/retail-events
60
+ ```
61
+
62
+ `shape.emitters` `kafka`: one message per event, body = the event JSON (flat, or `--envelope
63
+ cloudevents`), **message key = the idempotency key `<table>/<seq>`**, header `shape-table`. The
64
+ producer uses `acks=all` and `enable.idempotence=true` (overridable through the `config` option, for
65
+ SASL/TLS and the like). `emit` returns after the broker acknowledged every message, so the
66
+ runtime's checkpoint never passes an undelivered event; a failed delivery is retried by the runtime
67
+ (a message may arrive twice: keep the first of each key). A full local queue is waited for.
68
+ See `docs/EMIT.md`.
69
+
70
+ ### Registry formats: Avro, Protobuf, JSON Schema
71
+
72
+ ```
73
+ pip install 'sqllocks-shape-kafka[avro]' # fastavro; or [protobuf]; json-schema needs nothing
74
+ shape emit retail --sink kafka://broker:9092/orders --event-format avro \
75
+ --sink-config kafka.schema_registry_url=https://registry:8081
76
+ ```
77
+
78
+ `--event-format json|avro|protobuf|json-schema` (default `json`, today's bytes). The emitter
79
+ derives the schema from the table's Arrow schema, registers it with a Confluent-compatible registry
80
+ (`kafka.schema_registry_url`; `kafka.schema_registry_username` and
81
+ `kafka.schema_registry_password=env://NAME`, a credential reference; `kafka.subject_strategy` =
82
+ `topic` (default, `<topic>-value`), `record` or `topic_record`) and writes the Confluent wire format:
83
+ byte 0, the 4-byte big-endian schema id, the payload (Protobuf: a zero byte for the message index
84
+ `[0]` first). The record (message) name is the table name in the namespace/package `shape.events`;
85
+ the table and its columns must be valid identifiers for Avro and Protobuf. The key stays
86
+ `<table>/<seq>`. The encoders are imported on first use; a missing one stops the run (exit 2) with the
87
+ `pip install` command above.
88
+
89
+ Type mapping (a nullable Arrow field is a nullable field: an Avro union with `null`, a Protobuf
90
+ `optional`, a JSON `null`; `_shape_table` and `_shape_seq` are never null; a float is always nullable
91
+ because NaN and infinities are sent as `null`, as in the flat JSON event):
92
+
93
+ | Arrow | Avro | Protobuf | JSON Schema |
94
+ |---|---|---|---|
95
+ | bool | `boolean` | `bool` | `boolean` |
96
+ | int8, int16, int32 | `int` | `int32` | `integer` |
97
+ | uint8, uint16 | `int` | `uint32` | `integer` |
98
+ | int64, uint32 | `long` | `int64`, `uint32` | `integer` |
99
+ | uint64 | `long` (a value above 2^63-1 cannot be encoded) | `uint64` | `integer` |
100
+ | float32, float64 | `float`, `double` | `float`, `double` | `number` |
101
+ | string | `string` | `string` | `string` |
102
+ | uuid (the Arrow extension type) | `string`, logical type `uuid` | `string` | `string`, format `uuid` |
103
+ | binary | `bytes` | `bytes` | `string`, `contentEncoding: base64` |
104
+ | decimal(p, s) | `bytes`, logical type `decimal(p, s)` | `string` (canonical text) | `string` (canonical text) |
105
+ | date | `int`, logical type `date` | `int32` (days since 1970-01-01) | `string`, format `date` |
106
+ | time | `time-millis` (s, ms) or `time-micros` | `int64` (microseconds) | `string` |
107
+ | timestamp, no time zone | `local-timestamp-millis` (s, ms) or `-micros` (us, ns) | `int64` (wall-clock microseconds since the epoch) | `string`, "timestamp without time zone" |
108
+ | timestamp with time zone | `timestamp-millis` (s, ms) or `-micros` (us, ns) | `int64` (UTC microseconds since the epoch) | `string`, format `date-time` |
109
+
110
+ A nanosecond value is cut to microseconds. A dictionary column is its value type. Lists, structs and
111
+ maps are not mapped (the table is refused with exit 2; use `json`). Protobuf field numbers follow the
112
+ column order. Tests decode every message back from the registry's copy of the schema and compare it
113
+ with the flat event (`shape_kafka.testing.FakeRegistry`, `decode_messages`).
114
+
115
+ Tests: contract tests run on every PR against in-memory fakes
116
+ (`shape_kafka.testing.FakeBroker`, `FakeProducer`; the emitter contract is
117
+ `shape.streaming.emit.contract`); `pytest -m emulator plugins/shape-kafka/tests` runs the
118
+ end-to-end tests against a real broker (`docker compose -f ci/emulators/docker-compose.yml up -d
119
+ --wait kafka`), nightly in CI.
120
+
121
+ Its version always equals core's (`sqllocks-shape`), and it is released together with core.
122
+ How plugins are written: `docs/plugins/authoring.md` in the repository.
123
+
124
+ Live tests (`pytest -m live plugins/shape-kafka/tests`) need `KAFKA_SERVERS`, `KAFKA_TOPIC` and,
125
+ optionally, `KAFKA_CONFIG` (JSON); the nightly job runs them only where those secrets exist.
@@ -0,0 +1,97 @@
1
+ # sqllocks-shape-kafka
2
+
3
+ Shape plugin: a Kafka stream source (`shape.stream_sources`, scheme `kafka://`).
4
+
5
+ ```
6
+ pip install sqllocks-shape-kafka # or: pip install 'sqllocks-shape[kafka]'
7
+ shape stream-profile kafka://broker:9092/orders -o orders.json
8
+ ```
9
+
10
+ On Alpine/musllinux, `confluent-kafka` has no ARM64 or x86_64 wheels and builds from
11
+ source. Installing this plugin, `sqllocks-shape[kafka]` or `sqllocks-shape[all]` needs
12
+ librdkafka **2.16 or newer**, development headers and C build tools. Alpine's
13
+ `librdkafka-dev` is currently older. Build the library first using the
14
+ [Alpine installation steps](../../docs/INSTALL.md#alpine-and-musllinux), which match
15
+ `scripts/release_smoke_alpine.sh`, or
16
+ [Confluent's build instructions](https://github.com/confluentinc/librdkafka#building).
17
+
18
+ `kafka://host1:9092,host2:9092/TOPIC` reads one topic as Arrow micro-batches. Message bodies are
19
+ JSON objects; their fields become columns, and `_shape_event_time` carries the event time (the
20
+ body's `_shape_event_time`, or the message timestamp). Shape keeps its own position: the consumer
21
+ is *assigned* every partition at explicit offsets, nothing is committed to the broker, and a
22
+ `StreamOffset` is `{partition: next offset}`, so a checkpoint resumes exactly where it stopped.
23
+
24
+ Options, offsets, delivery guarantees and the checkpoint format are described in
25
+ `docs/plugins/streaming.md`. Extra `confluent-kafka` settings (SASL, TLS) go in through the
26
+ `config` source option, preferably from a file (`--options-file`).
27
+
28
+ ## Emitting events: `kafka://`
29
+
30
+ ```
31
+ shape emit retail --realtime --rate 5000 --sink kafka://broker:9092/retail-events
32
+ ```
33
+
34
+ `shape.emitters` `kafka`: one message per event, body = the event JSON (flat, or `--envelope
35
+ cloudevents`), **message key = the idempotency key `<table>/<seq>`**, header `shape-table`. The
36
+ producer uses `acks=all` and `enable.idempotence=true` (overridable through the `config` option, for
37
+ SASL/TLS and the like). `emit` returns after the broker acknowledged every message, so the
38
+ runtime's checkpoint never passes an undelivered event; a failed delivery is retried by the runtime
39
+ (a message may arrive twice: keep the first of each key). A full local queue is waited for.
40
+ See `docs/EMIT.md`.
41
+
42
+ ### Registry formats: Avro, Protobuf, JSON Schema
43
+
44
+ ```
45
+ pip install 'sqllocks-shape-kafka[avro]' # fastavro; or [protobuf]; json-schema needs nothing
46
+ shape emit retail --sink kafka://broker:9092/orders --event-format avro \
47
+ --sink-config kafka.schema_registry_url=https://registry:8081
48
+ ```
49
+
50
+ `--event-format json|avro|protobuf|json-schema` (default `json`, today's bytes). The emitter
51
+ derives the schema from the table's Arrow schema, registers it with a Confluent-compatible registry
52
+ (`kafka.schema_registry_url`; `kafka.schema_registry_username` and
53
+ `kafka.schema_registry_password=env://NAME`, a credential reference; `kafka.subject_strategy` =
54
+ `topic` (default, `<topic>-value`), `record` or `topic_record`) and writes the Confluent wire format:
55
+ byte 0, the 4-byte big-endian schema id, the payload (Protobuf: a zero byte for the message index
56
+ `[0]` first). The record (message) name is the table name in the namespace/package `shape.events`;
57
+ the table and its columns must be valid identifiers for Avro and Protobuf. The key stays
58
+ `<table>/<seq>`. The encoders are imported on first use; a missing one stops the run (exit 2) with the
59
+ `pip install` command above.
60
+
61
+ Type mapping (a nullable Arrow field is a nullable field: an Avro union with `null`, a Protobuf
62
+ `optional`, a JSON `null`; `_shape_table` and `_shape_seq` are never null; a float is always nullable
63
+ because NaN and infinities are sent as `null`, as in the flat JSON event):
64
+
65
+ | Arrow | Avro | Protobuf | JSON Schema |
66
+ |---|---|---|---|
67
+ | bool | `boolean` | `bool` | `boolean` |
68
+ | int8, int16, int32 | `int` | `int32` | `integer` |
69
+ | uint8, uint16 | `int` | `uint32` | `integer` |
70
+ | int64, uint32 | `long` | `int64`, `uint32` | `integer` |
71
+ | uint64 | `long` (a value above 2^63-1 cannot be encoded) | `uint64` | `integer` |
72
+ | float32, float64 | `float`, `double` | `float`, `double` | `number` |
73
+ | string | `string` | `string` | `string` |
74
+ | uuid (the Arrow extension type) | `string`, logical type `uuid` | `string` | `string`, format `uuid` |
75
+ | binary | `bytes` | `bytes` | `string`, `contentEncoding: base64` |
76
+ | decimal(p, s) | `bytes`, logical type `decimal(p, s)` | `string` (canonical text) | `string` (canonical text) |
77
+ | date | `int`, logical type `date` | `int32` (days since 1970-01-01) | `string`, format `date` |
78
+ | time | `time-millis` (s, ms) or `time-micros` | `int64` (microseconds) | `string` |
79
+ | timestamp, no time zone | `local-timestamp-millis` (s, ms) or `-micros` (us, ns) | `int64` (wall-clock microseconds since the epoch) | `string`, "timestamp without time zone" |
80
+ | timestamp with time zone | `timestamp-millis` (s, ms) or `-micros` (us, ns) | `int64` (UTC microseconds since the epoch) | `string`, format `date-time` |
81
+
82
+ A nanosecond value is cut to microseconds. A dictionary column is its value type. Lists, structs and
83
+ maps are not mapped (the table is refused with exit 2; use `json`). Protobuf field numbers follow the
84
+ column order. Tests decode every message back from the registry's copy of the schema and compare it
85
+ with the flat event (`shape_kafka.testing.FakeRegistry`, `decode_messages`).
86
+
87
+ Tests: contract tests run on every PR against in-memory fakes
88
+ (`shape_kafka.testing.FakeBroker`, `FakeProducer`; the emitter contract is
89
+ `shape.streaming.emit.contract`); `pytest -m emulator plugins/shape-kafka/tests` runs the
90
+ end-to-end tests against a real broker (`docker compose -f ci/emulators/docker-compose.yml up -d
91
+ --wait kafka`), nightly in CI.
92
+
93
+ Its version always equals core's (`sqllocks-shape`), and it is released together with core.
94
+ How plugins are written: `docs/plugins/authoring.md` in the repository.
95
+
96
+ Live tests (`pytest -m live plugins/shape-kafka/tests`) need `KAFKA_SERVERS`, `KAFKA_TOPIC` and,
97
+ optionally, `KAFKA_CONFIG` (JSON); the nightly job runs them only where those secrets exist.
@@ -0,0 +1,43 @@
1
+ [build-system]
2
+ requires = ["setuptools>=77"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "sqllocks-shape-kafka"
7
+ # Lockstep with core (T-09): scripts/check_plugin_skeletons.py fails when this drifts.
8
+ version = "0.9.1"
9
+ description = "Shape plugin: Kafka stream source and event emitter"
10
+ readme = "README.md"
11
+ requires-python = ">=3.11"
12
+ license = "MIT"
13
+ license-files = ["LICENSE"]
14
+ classifiers = [
15
+ "Programming Language :: Python :: 3",
16
+ "Programming Language :: Python :: 3 :: Only",
17
+ "Programming Language :: Python :: 3.11",
18
+ "Programming Language :: Python :: 3.12",
19
+ "Programming Language :: Python :: 3.13",
20
+ "Programming Language :: Python :: 3.14",
21
+ "Operating System :: OS Independent",
22
+ "Topic :: Database",
23
+ ]
24
+ dependencies = ["sqllocks-shape==0.9.1", "confluent-kafka>=2.5,<4"]
25
+
26
+ [project.urls]
27
+ Homepage = "https://github.com/sqllocks/shape"
28
+ Repository = "https://github.com/sqllocks/shape"
29
+ Issues = "https://github.com/sqllocks/shape/issues"
30
+ Documentation = "https://github.com/sqllocks/shape/tree/main/plugins/shape-kafka"
31
+
32
+ [project.optional-dependencies]
33
+ avro = ["fastavro>=1.9"]
34
+ protobuf = ["protobuf>=5"]
35
+
36
+ [project.entry-points."shape.stream_sources"]
37
+ kafka = "shape_kafka:KafkaStreamSource"
38
+
39
+ [project.entry-points."shape.emitters"]
40
+ kafka = "shape_kafka:KafkaEmitter"
41
+
42
+ [tool.setuptools.packages.find]
43
+ where = ["src"]
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,17 @@
1
+ """Shape plugin: Kafka.
2
+
3
+ * ``shape.stream_sources``: ``kafka`` reads a topic (``kafka://host:9092/topic``) as Arrow
4
+ micro-batches with ``{partition: next offset}`` offsets, for ``shape stream-profile``.
5
+
6
+ * ``shape.emitters``: ``kafka`` sends events to a topic (``kafka://host:9092/topic``), one
7
+ message per event keyed by the idempotency key ``<table>/<seq>``, for ``shape emit``.
8
+
9
+ ``confluent-kafka`` loads only when a read or a send starts.
10
+ """
11
+
12
+ from .emitter import KafkaEmitter
13
+ from .source import KafkaStreamSource
14
+
15
+ SHAPE_API = "1.0"
16
+
17
+ __all__ = ["SHAPE_API", "KafkaEmitter", "KafkaStreamSource"]
@@ -0,0 +1,361 @@
1
+ """The ``kafka://`` emitter (P5-02).
2
+
3
+ ``kafka://broker1:9092,broker2:9092/topic`` sends each event as one message: the body is the
4
+ event's JSON (flat, or the CloudEvents envelope), the **message key is the D-12 idempotency key**
5
+ ``<table>/<seq>`` (so a replay of a row lands in the same partition with the same key, and a
6
+ log-compacted topic keeps one message per row), and the header ``shape-table`` names the table.
7
+
8
+ Delivery is at-least-once. ``emit`` returns after the broker has acknowledged every message of
9
+ its batches (``acks=all``, idempotent producer): a failed delivery raises ``ConnectionError``,
10
+ which the runtime retries, so a repeated message is possible and a lost one is not. A full local
11
+ queue is waited for (the producer is polled until it drains), which is the backpressure the
12
+ runtime passes back to the generator. A message the broker did not acknowledge within
13
+ ``flush_timeout`` seconds raises ``TimeoutError``.
14
+
15
+ Options of :meth:`KafkaEmitter.emit`:
16
+
17
+ ``envelope`` ``"flat"`` (default) or ``"cloudevents"`` (JSON only).
18
+ ``event_format`` ``"json"`` (default: the event's JSON, unchanged), ``"avro"``, ``"protobuf"`` or
19
+ ``"json-schema"``: the schema is derived from the table's Arrow schema,
20
+ registered with the registry at ``schema_registry_url``, and the message value
21
+ is the Confluent wire format (see :mod:`shape_kafka.formats`).
22
+ ``schema_registry_url``, ``subject_strategy`` (``topic``, ``record``, ``topic_record``),
23
+ ``schema_registry_username`` and ``schema_registry_password``
24
+ the registry (``--sink-config kafka.KEY=VALUE``; the password is a credential
25
+ reference resolved before it reaches the emitter).
26
+ ``config`` extra ``confluent-kafka`` producer settings (security, SASL, compression, ...);
27
+ ``bootstrap.servers`` comes from the URI. Defaults: ``acks=all``,
28
+ ``enable.idempotence=true``, ``linger.ms=5``.
29
+ ``flush_timeout`` seconds to wait for acknowledgements (default 60).
30
+ ``resuming`` accepted and ignored (a topic has no file to continue).
31
+
32
+ Rejections: a non-retryable error about one message (too large, an invalid record) or an event the
33
+ chosen ``event_format`` cannot hold does not fail the batch. Every other event is delivered and
34
+ acknowledged, then :class:`shape.streaming.emit.RejectedEvents` is raised with the keys and
35
+ reasons (``shape emit --dead-letter`` routes them). A dead-letter record (batch column
36
+ ``_shape_dead_letter_reason``) also carries the header ``shape-dead-letter-reason``.
37
+
38
+ ``confluent-kafka`` is imported when the first message is sent, never at plugin load.
39
+ """
40
+
41
+ from __future__ import annotations
42
+
43
+ import json
44
+ from collections.abc import Callable, Iterable, Mapping
45
+ from functools import partial
46
+ from typing import Any
47
+ from urllib.parse import unquote, urlsplit
48
+
49
+ import pyarrow as pa
50
+
51
+ from shape.errors import ShapeError
52
+ from shape.streaming.emit import RejectedEvents, Rejection
53
+ from shape.streaming.emit.formats import (
54
+ ENVELOPES,
55
+ FIELD_DEAD_REASON,
56
+ FIELD_POISON,
57
+ FIELD_SEQ,
58
+ FIELD_TABLE,
59
+ EncodedEvent,
60
+ encode_events,
61
+ poison_body,
62
+ )
63
+
64
+ from .formats import EVENT_FORMATS, EncodeError, TableCodec, make_codec, wire
65
+ from .registry import SUBJECT_STRATEGIES, RegistryClient, Transport, subject_name
66
+
67
+ DEFAULT_CONFIG: dict[str, Any] = {"acks": "all", "enable.idempotence": True, "linger.ms": 5}
68
+ HEADER_TABLE = "shape-table"
69
+ HEADER_SYNTHETIC = "shape-synthetic"
70
+ HEADER_DEAD_LETTER_REASON = "shape-dead-letter-reason"
71
+ # Errors about one message, not about the cluster or the topic: the broker or the client refuses
72
+ # this message for good. Such a message is a rejection (dead-lettered when there is a destination
73
+ # for them); every other failure is the destination's and stops the run after the retries.
74
+ PER_MESSAGE_ERRORS = frozenset(
75
+ {
76
+ "MSG_SIZE_TOO_LARGE",
77
+ "_MSG_SIZE_TOO_LARGE",
78
+ "INVALID_MSG",
79
+ "INVALID_MSG_SIZE",
80
+ "INVALID_RECORD",
81
+ "CORRUPT_MESSAGE",
82
+ "RECORD_LIST_TOO_LARGE",
83
+ }
84
+ )
85
+ _POLL = 0.05 # seconds one wait for a full queue to drain
86
+
87
+
88
+ def parse_uri(uri: str) -> tuple[str, str]:
89
+ """``(bootstrap servers, topic)`` of a ``kafka://`` URI."""
90
+ parts = urlsplit(uri)
91
+ if parts.scheme != "kafka" or not parts.netloc:
92
+ raise ShapeError(f"not a kafka URI: {uri!r} (kafka://host:9092[,host2:9092]/topic)")
93
+ topic = unquote(parts.path.lstrip("/"))
94
+ if not topic or "/" in topic:
95
+ raise ShapeError(f"the kafka URI needs one topic name: {uri!r}")
96
+ return parts.netloc, topic
97
+
98
+
99
+ def _per_message_error(err: Any) -> str | None:
100
+ """The name of ``err`` when it is a non-retryable error about one message, else ``None``."""
101
+ name = getattr(err, "name", None)
102
+ if not callable(name):
103
+ return None
104
+ retriable = getattr(err, "retriable", None)
105
+ if callable(retriable) and retriable():
106
+ return None
107
+ text = str(name())
108
+ return text if text in PER_MESSAGE_ERRORS else None
109
+
110
+
111
+ def _confluent_producer(config: dict[str, Any]) -> Any:
112
+ try:
113
+ from confluent_kafka import Producer
114
+ except ImportError as exc:
115
+ raise ShapeError(
116
+ "emitting to Kafka needs the confluent-kafka package (pip install sqllocks-shape-kafka)"
117
+ ) from exc
118
+ return Producer(config)
119
+
120
+
121
+ class KafkaEmitter:
122
+ """Events to a Kafka topic, keyed by the idempotency key.
123
+
124
+ ``producer_factory(config)`` is for tests: it stands in for ``confluent_kafka.Producer``.
125
+ """
126
+
127
+ name = "kafka"
128
+ accepts_poison = True # one JSON text per message: a cut-off body is a poison message
129
+ supports_synthetic = True # the `synthetic` option marks every message with a header
130
+ schemes = ("kafka",)
131
+ event_formats = EVENT_FORMATS # `--event-format`
132
+ sink_config_keys = ( # `--sink-config kafka.KEY=VALUE`
133
+ "schema_registry_url",
134
+ "subject_strategy",
135
+ "schema_registry_username",
136
+ "schema_registry_password",
137
+ )
138
+
139
+ def __init__(
140
+ self,
141
+ producer_factory: Callable[[dict[str, Any]], Any] | None = None,
142
+ registry_transport: Transport | None = None,
143
+ ) -> None:
144
+ self._factory = producer_factory or _confluent_producer
145
+ self._registry_transport = registry_transport
146
+ self._producers: dict[str, Any] = {}
147
+ self._registries: dict[tuple[Any, ...], RegistryClient] = {}
148
+ self._codecs: dict[tuple[Any, ...], tuple[TableCodec, int]] = {}
149
+
150
+ def _registry(self, url: str, user: str | None, password: str | None) -> RegistryClient:
151
+ key = (url, user, password)
152
+ client = self._registries.get(key)
153
+ if client is None:
154
+ client = self._registries[key] = RegistryClient(
155
+ url, username=user, password=password, transport=self._registry_transport
156
+ )
157
+ return client
158
+
159
+ def _codec(
160
+ self,
161
+ fmt: str,
162
+ topic: str,
163
+ table: str,
164
+ schema: pa.Schema,
165
+ strategy: str,
166
+ registry: RegistryClient,
167
+ ) -> tuple[TableCodec, int]:
168
+ """The codec of ``table`` and the id its schema got from the registry (registered on
169
+ the table's first batch, once per run). A table whose columns change during the run (a
170
+ drift plan that adds or drops one) gets a new schema, registered under the same subject:
171
+ whether that is a compatible new version is the registry's decision."""
172
+ shape = tuple((f.name, str(f.type), f.nullable) for f in schema if f.name != FIELD_POISON)
173
+ key = (fmt, topic, strategy, table, shape)
174
+ hit = self._codecs.get(key)
175
+ if hit is None:
176
+ codec = make_codec(fmt, table, schema)
177
+ subject = subject_name(strategy, topic, codec.record)
178
+ hit = self._codecs[key] = (
179
+ codec,
180
+ registry.register(subject, codec.schema_text, codec.schema_type),
181
+ )
182
+ return hit
183
+
184
+ def _registry_events(
185
+ self,
186
+ batch: pa.RecordBatch,
187
+ fmt: str,
188
+ topic: str,
189
+ strategy: str,
190
+ registry: RegistryClient,
191
+ ) -> tuple[list[EncodedEvent], list[Rejection]]:
192
+ """``batch`` as wire-format events, in order, for ``fmt``; an event the format cannot
193
+ hold is a :class:`Rejection` (the original flat JSON is its body) and is not in the
194
+ list."""
195
+ marks = (
196
+ [bool(v) for v in batch.column(FIELD_POISON).to_pylist()]
197
+ if FIELD_POISON in batch.schema.names
198
+ else None
199
+ )
200
+ if marks is not None:
201
+ batch = batch.select([n for n in batch.schema.names if n != FIELD_POISON])
202
+ tables = batch.column(FIELD_TABLE).to_pylist()
203
+ seqs = batch.column(FIELD_SEQ).to_pylist()
204
+ times = (
205
+ batch.column("_shape_event_time").cast(pa.string()).to_pylist()
206
+ if "_shape_event_time" in batch.schema.names
207
+ else [None] * len(tables)
208
+ )
209
+ bodies: list[bytes | EncodeError | None] = [None] * len(tables)
210
+ for table in dict.fromkeys(tables):
211
+ idx = [i for i, t in enumerate(tables) if t == table]
212
+ part = batch if len(idx) == len(tables) else batch.take(pa.array(idx))
213
+ codec, schema_id = self._codec(fmt, topic, table, part.schema, strategy, registry)
214
+ for i, payload in zip(idx, codec.encode_batch(part), strict=True):
215
+ bodies[i] = (
216
+ payload if isinstance(payload, EncodeError) else wire(schema_id, payload, fmt)
217
+ )
218
+ events: list[EncodedEvent] = []
219
+ rejected: list[Rejection] = []
220
+ for i, body in enumerate(bodies):
221
+ key = f"{tables[i]}/{seqs[i]}"
222
+ if isinstance(body, EncodeError):
223
+ rejected.append(Rejection(key, f"cannot encode as {fmt}: {body}"))
224
+ continue
225
+ assert body is not None
226
+ if marks is not None and marks[i]:
227
+ body = poison_body(body)
228
+ events.append(EncodedEvent(key, tables[i], seqs[i], times[i], body))
229
+ return events, rejected
230
+
231
+ def _producer(self, servers: str, extra: Mapping[str, Any] | None) -> Any:
232
+ config = {**DEFAULT_CONFIG, **(extra or {}), "bootstrap.servers": servers}
233
+ key = json.dumps(config, sort_keys=True, default=str)
234
+ producer = self._producers.get(key)
235
+ if producer is None:
236
+ producer = self._producers[key] = self._factory(config)
237
+ return producer
238
+
239
+ def emit(
240
+ self,
241
+ uri: str,
242
+ batches: Iterable[pa.RecordBatch],
243
+ *,
244
+ envelope: str = "flat",
245
+ resuming: bool = False,
246
+ config: Mapping[str, Any] | None = None,
247
+ flush_timeout: float = 60.0,
248
+ synthetic: bool = False,
249
+ event_format: str = "json",
250
+ schema_registry_url: str | None = None,
251
+ subject_strategy: str = "topic",
252
+ schema_registry_username: str | None = None,
253
+ schema_registry_password: str | None = None,
254
+ **options: Any,
255
+ ) -> int:
256
+ """Send every batch; return the number of events, after the broker acknowledged them."""
257
+ if options:
258
+ raise ShapeError(f"unknown kafka emitter options: {sorted(options)}")
259
+ if envelope not in ENVELOPES:
260
+ raise ShapeError(f"unknown envelope {envelope!r}; choose from {', '.join(ENVELOPES)}")
261
+ if event_format not in EVENT_FORMATS:
262
+ raise ShapeError(
263
+ f"unknown event format {event_format!r}; choose from {', '.join(EVENT_FORMATS)}"
264
+ )
265
+ registry: RegistryClient | None = None
266
+ if event_format != "json":
267
+ if envelope != "flat":
268
+ raise ShapeError(
269
+ f"--envelope {envelope} is JSON only: it cannot be combined with "
270
+ f"--event-format {event_format}"
271
+ )
272
+ if subject_strategy not in SUBJECT_STRATEGIES:
273
+ raise ShapeError(
274
+ f"unknown subject strategy {subject_strategy!r}; "
275
+ f"choose from {', '.join(SUBJECT_STRATEGIES)}"
276
+ )
277
+ if not schema_registry_url:
278
+ raise ShapeError(
279
+ f"--event-format {event_format} needs a schema registry: "
280
+ "--sink-config kafka.schema_registry_url=URL"
281
+ )
282
+ registry = self._registry(
283
+ schema_registry_url, schema_registry_username, schema_registry_password
284
+ )
285
+ servers, topic = parse_uri(uri)
286
+ producer = self._producer(servers, config)
287
+ failures: list[Any] = []
288
+ rejected: list[Rejection] = []
289
+
290
+ def reported(key: str, body: bytes, err: Any, _msg: Any = None) -> None:
291
+ if err is None:
292
+ return
293
+ name = _per_message_error(err)
294
+ if name is not None:
295
+ rejected.append(Rejection(key, name, body))
296
+ else:
297
+ failures.append(err)
298
+
299
+ sent = 0
300
+ marker = [(HEADER_SYNTHETIC, b"true")] if synthetic else []
301
+ for batch in batches:
302
+ reasons = (
303
+ batch.column(FIELD_DEAD_REASON).to_pylist()
304
+ if FIELD_DEAD_REASON in batch.schema.names
305
+ else None
306
+ )
307
+ if registry is None:
308
+ events = encode_events(batch, envelope)
309
+ else:
310
+ events, refused = self._registry_events(
311
+ batch, event_format, topic, subject_strategy, registry
312
+ )
313
+ rejected.extend(refused)
314
+ for i, event in enumerate(events):
315
+ headers = [(HEADER_TABLE, event.table.encode("utf-8")), *marker]
316
+ if reasons is not None:
317
+ headers.append((HEADER_DEAD_LETTER_REASON, str(reasons[i]).encode("utf-8")))
318
+ while True:
319
+ try:
320
+ producer.produce(
321
+ topic,
322
+ value=event.body,
323
+ key=event.key.encode("utf-8"),
324
+ headers=headers,
325
+ on_delivery=partial(reported, event.key, event.body),
326
+ )
327
+ break
328
+ except BufferError:
329
+ # The local queue is full: serve delivery reports until it drains.
330
+ producer.poll(_POLL)
331
+ except Exception as exc: # a KafkaException: the client refuses the message
332
+ name = _per_message_error(exc.args[0] if exc.args else None)
333
+ if name is None:
334
+ raise
335
+ rejected.append(Rejection(event.key, name, event.body))
336
+ break
337
+ sent += 1
338
+ producer.poll(0)
339
+ pending = producer.flush(flush_timeout)
340
+ if failures:
341
+ raise ConnectionError(
342
+ f"kafka: {len(failures)} of {sent} messages were not delivered to {topic!r}: "
343
+ f"{failures[0]}"
344
+ )
345
+ if pending:
346
+ raise TimeoutError(
347
+ f"kafka: {pending} messages to {topic!r} not acknowledged in {flush_timeout}s"
348
+ )
349
+ if rejected:
350
+ raise RejectedEvents(rejected)
351
+ return sent
352
+
353
+ def flush(self) -> None:
354
+ for producer in self._producers.values():
355
+ producer.flush(60.0)
356
+
357
+ def close(self) -> None:
358
+ self.flush()
359
+ self._producers.clear()
360
+ self._codecs.clear()
361
+ self._registries.clear()