esmd-pysdk 0.1.0__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 (34) hide show
  1. esmd_pysdk-0.1.0/LICENSE +21 -0
  2. esmd_pysdk-0.1.0/PKG-INFO +143 -0
  3. esmd_pysdk-0.1.0/README.md +118 -0
  4. esmd_pysdk-0.1.0/esmd_pysdk/__init__.py +54 -0
  5. esmd_pysdk-0.1.0/esmd_pysdk/claim.py +67 -0
  6. esmd_pysdk-0.1.0/esmd_pysdk/client.py +102 -0
  7. esmd_pysdk-0.1.0/esmd_pysdk/cluster.py +111 -0
  8. esmd_pysdk-0.1.0/esmd_pysdk/compress.py +64 -0
  9. esmd_pysdk-0.1.0/esmd_pysdk/config.py +126 -0
  10. esmd_pysdk-0.1.0/esmd_pysdk/constants.py +53 -0
  11. esmd_pysdk-0.1.0/esmd_pysdk/consumer.py +87 -0
  12. esmd_pysdk-0.1.0/esmd_pysdk/driver.py +32 -0
  13. esmd_pysdk-0.1.0/esmd_pysdk/errors.py +63 -0
  14. esmd_pysdk-0.1.0/esmd_pysdk/limiter.py +34 -0
  15. esmd_pysdk-0.1.0/esmd_pysdk/message.py +242 -0
  16. esmd_pysdk-0.1.0/esmd_pysdk/packet.py +60 -0
  17. esmd_pysdk-0.1.0/esmd_pysdk/partition.py +208 -0
  18. esmd_pysdk-0.1.0/esmd_pysdk/producer.py +101 -0
  19. esmd_pysdk-0.1.0/esmd_pysdk/protocol.py +102 -0
  20. esmd_pysdk-0.1.0/esmd_pysdk/session.py +114 -0
  21. esmd_pysdk-0.1.0/esmd_pysdk/slot.py +100 -0
  22. esmd_pysdk-0.1.0/esmd_pysdk/utils.py +42 -0
  23. esmd_pysdk-0.1.0/examples/consumer.py +43 -0
  24. esmd_pysdk-0.1.0/examples/delete.py +31 -0
  25. esmd_pysdk-0.1.0/examples/matrix.py +32 -0
  26. esmd_pysdk-0.1.0/examples/producer.py +48 -0
  27. esmd_pysdk-0.1.0/examples/pull.py +28 -0
  28. esmd_pysdk-0.1.0/examples/query.py +35 -0
  29. esmd_pysdk-0.1.0/examples/schedule.py +45 -0
  30. esmd_pysdk-0.1.0/examples/subscribe.py +33 -0
  31. esmd_pysdk-0.1.0/pyproject.toml +43 -0
  32. esmd_pysdk-0.1.0/requirements.txt +1 -0
  33. esmd_pysdk-0.1.0/tests/test_hash.py +6 -0
  34. esmd_pysdk-0.1.0/tests/test_packet.py +9 -0
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Liu Zhenhao
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,143 @@
1
+ Metadata-Version: 2.4
2
+ Name: esmd-pysdk
3
+ Version: 0.1.0
4
+ Summary: Asyncio-based Python driver for esmdMQ
5
+ Project-URL: Homepage, https://pypi.org/project/esmd-pysdk/
6
+ Project-URL: Repository, https://pypi.org/project/esmd-pysdk/
7
+ Project-URL: Documentation, https://pypi.org/project/esmd-pysdk/
8
+ Author: Liu Zhenhao
9
+ License-Expression: MIT
10
+ License-File: LICENSE
11
+ Keywords: asyncio,esmd,message-queue,mq,sdk
12
+ Classifier: Development Status :: 3 - Alpha
13
+ Classifier: Intended Audience :: Developers
14
+ Classifier: License :: OSI Approved :: MIT License
15
+ Classifier: Programming Language :: Python :: 3
16
+ Classifier: Programming Language :: Python :: 3 :: Only
17
+ Classifier: Programming Language :: Python :: 3.10
18
+ Classifier: Programming Language :: Python :: 3.11
19
+ Classifier: Programming Language :: Python :: 3.12
20
+ Classifier: Topic :: Software Development :: Libraries
21
+ Classifier: Topic :: System :: Networking
22
+ Requires-Python: >=3.10
23
+ Requires-Dist: lz4
24
+ Description-Content-Type: text/markdown
25
+
26
+ # esmd-pysdk
27
+
28
+ Asyncio-based Python driver for esmdMQ. This project lives at the repo root and is parallel to `drivers/`.
29
+
30
+ ## Requirements
31
+
32
+ - Python 3.10+
33
+ - `lz4` (see `requirements.txt` or `pyproject.toml`)
34
+
35
+ ## Quick Start
36
+
37
+ ```bash
38
+ pip install -r requirements.txt
39
+ ```
40
+
41
+ Install from PyPI:
42
+
43
+ ```bash
44
+ pip install esmd-pysdk
45
+ ```
46
+
47
+ Check installed SDK version:
48
+
49
+ ```bash
50
+ python -c "import esmd_pysdk; print(esmd_pysdk.__version__)"
51
+ ```
52
+
53
+ Producer example:
54
+
55
+ ```bash
56
+ python examples/producer.py
57
+ ```
58
+
59
+ Consumer example:
60
+
61
+ ```bash
62
+ python examples/consumer.py
63
+ ```
64
+
65
+ Query example:
66
+
67
+ ```bash
68
+ python examples/query.py
69
+ ```
70
+
71
+ Matrix example:
72
+
73
+ ```bash
74
+ python examples/matrix.py
75
+ ```
76
+
77
+ Delete example:
78
+
79
+ ```bash
80
+ python examples/delete.py
81
+ ```
82
+
83
+ Schedule example (DelayTime + Punctual):
84
+
85
+ ```bash
86
+ python examples/schedule.py
87
+ ```
88
+
89
+ ## Notes
90
+
91
+ - JSON serialization is canonicalized (sorted keys, compact separators) to keep hash routing consistent with Go.
92
+ - Driver uses asyncio; call APIs inside an event loop.
93
+ - Producer schedule:
94
+ - `ScheduleMode.SERVER_TIME`: 服务端当前时间(默认)
95
+ - `ScheduleMode.DELAY_TIME`: 延迟投递,参数单位为毫秒(例如 `2000`)
96
+ - `ScheduleMode.PUNCTUAL`: 绝对时间投递,参数为 Unix 毫秒时间戳(例如 `int(time.time() * 1000) + 5000`)
97
+ - Consumer delay:
98
+ - `DelayMode.BASE_SERVER`: 基于服务端时间消费(推荐默认)
99
+ - `DelayMode.IGNORE_DELAY`: 忽略延迟立即消费
100
+ - `DelayMode.TIME_DELAY`: 指定延迟窗口(毫秒)
101
+
102
+ ## Recommended Demo Flow
103
+
104
+ ```bash
105
+ # 1) 先写入一条固定索引数据
106
+ python examples/producer.py
107
+
108
+ # 2) 查询和矩阵查询(使用与 producer 一致的 filter)
109
+ python examples/query.py
110
+ python examples/matrix.py
111
+
112
+ # 3) 删除同一条件下的数据
113
+ python examples/delete.py
114
+ ```
115
+
116
+ ## Release Notes
117
+
118
+ - Version is managed in `pyproject.toml` (`[project].version`).
119
+ - Release guide: `RELEASING.md` (TestPyPI -> PyPI).
120
+
121
+ Pre-release checklist:
122
+
123
+ ```bash
124
+ # 1) version check
125
+ python - <<'PY'
126
+ import re
127
+ print(re.search(r'^version\\s*=\\s*\"([^\"]+)\"', open('pyproject.toml').read(), re.M).group(1))
128
+ PY
129
+
130
+ # 2) tests
131
+ python -m pytest tests
132
+
133
+ # 3) build + metadata check
134
+ python -m build
135
+ python -m twine check dist/*
136
+
137
+ # 4) install/import smoke
138
+ python -m venv .venv-release
139
+ source .venv-release/bin/activate
140
+ python -m pip install --upgrade pip
141
+ python -m pip install dist/*.whl
142
+ python -c "import esmd_pysdk; print('import ok')"
143
+ ```
@@ -0,0 +1,118 @@
1
+ # esmd-pysdk
2
+
3
+ Asyncio-based Python driver for esmdMQ. This project lives at the repo root and is parallel to `drivers/`.
4
+
5
+ ## Requirements
6
+
7
+ - Python 3.10+
8
+ - `lz4` (see `requirements.txt` or `pyproject.toml`)
9
+
10
+ ## Quick Start
11
+
12
+ ```bash
13
+ pip install -r requirements.txt
14
+ ```
15
+
16
+ Install from PyPI:
17
+
18
+ ```bash
19
+ pip install esmd-pysdk
20
+ ```
21
+
22
+ Check installed SDK version:
23
+
24
+ ```bash
25
+ python -c "import esmd_pysdk; print(esmd_pysdk.__version__)"
26
+ ```
27
+
28
+ Producer example:
29
+
30
+ ```bash
31
+ python examples/producer.py
32
+ ```
33
+
34
+ Consumer example:
35
+
36
+ ```bash
37
+ python examples/consumer.py
38
+ ```
39
+
40
+ Query example:
41
+
42
+ ```bash
43
+ python examples/query.py
44
+ ```
45
+
46
+ Matrix example:
47
+
48
+ ```bash
49
+ python examples/matrix.py
50
+ ```
51
+
52
+ Delete example:
53
+
54
+ ```bash
55
+ python examples/delete.py
56
+ ```
57
+
58
+ Schedule example (DelayTime + Punctual):
59
+
60
+ ```bash
61
+ python examples/schedule.py
62
+ ```
63
+
64
+ ## Notes
65
+
66
+ - JSON serialization is canonicalized (sorted keys, compact separators) to keep hash routing consistent with Go.
67
+ - Driver uses asyncio; call APIs inside an event loop.
68
+ - Producer schedule:
69
+ - `ScheduleMode.SERVER_TIME`: 服务端当前时间(默认)
70
+ - `ScheduleMode.DELAY_TIME`: 延迟投递,参数单位为毫秒(例如 `2000`)
71
+ - `ScheduleMode.PUNCTUAL`: 绝对时间投递,参数为 Unix 毫秒时间戳(例如 `int(time.time() * 1000) + 5000`)
72
+ - Consumer delay:
73
+ - `DelayMode.BASE_SERVER`: 基于服务端时间消费(推荐默认)
74
+ - `DelayMode.IGNORE_DELAY`: 忽略延迟立即消费
75
+ - `DelayMode.TIME_DELAY`: 指定延迟窗口(毫秒)
76
+
77
+ ## Recommended Demo Flow
78
+
79
+ ```bash
80
+ # 1) 先写入一条固定索引数据
81
+ python examples/producer.py
82
+
83
+ # 2) 查询和矩阵查询(使用与 producer 一致的 filter)
84
+ python examples/query.py
85
+ python examples/matrix.py
86
+
87
+ # 3) 删除同一条件下的数据
88
+ python examples/delete.py
89
+ ```
90
+
91
+ ## Release Notes
92
+
93
+ - Version is managed in `pyproject.toml` (`[project].version`).
94
+ - Release guide: `RELEASING.md` (TestPyPI -> PyPI).
95
+
96
+ Pre-release checklist:
97
+
98
+ ```bash
99
+ # 1) version check
100
+ python - <<'PY'
101
+ import re
102
+ print(re.search(r'^version\\s*=\\s*\"([^\"]+)\"', open('pyproject.toml').read(), re.M).group(1))
103
+ PY
104
+
105
+ # 2) tests
106
+ python -m pytest tests
107
+
108
+ # 3) build + metadata check
109
+ python -m build
110
+ python -m twine check dist/*
111
+
112
+ # 4) install/import smoke
113
+ python -m venv .venv-release
114
+ source .venv-release/bin/activate
115
+ python -m pip install --upgrade pip
116
+ python -m pip install dist/*.whl
117
+ python -c "import esmd_pysdk; print('import ok')"
118
+ ```
@@ -0,0 +1,54 @@
1
+ from importlib.metadata import PackageNotFoundError, version
2
+
3
+ from .claim import Claim, new_claim, new_opinion
4
+ from .config import Config, Endpoint, new_config
5
+ from .constants import DelayMode, ScheduleMode
6
+ from .errors import (
7
+ Error,
8
+ ErrHashSlotIsFull,
9
+ ErrInvalidTopicName,
10
+ ErrMessageTooLarge,
11
+ ErrNoDataFound,
12
+ ErrSessionClosed,
13
+ ErrUnAckSizeExceeded,
14
+ )
15
+ from .message import ConsumerMessage, Messages, ProducerMessage, new_producer_message
16
+ from .driver import new_async_producer, new_consumer, new_producer
17
+
18
+ try:
19
+ __version__ = version("esmd-pysdk")
20
+ except PackageNotFoundError:
21
+ # Fallback for source checkout before package installation.
22
+ __version__ = "0.1.0"
23
+
24
+
25
+ def get_version() -> str:
26
+ return __version__
27
+
28
+
29
+ __all__ = [
30
+ "__version__",
31
+ "get_version",
32
+ "Claim",
33
+ "new_claim",
34
+ "new_opinion",
35
+ "Config",
36
+ "Endpoint",
37
+ "new_config",
38
+ "DelayMode",
39
+ "ScheduleMode",
40
+ "Error",
41
+ "ErrHashSlotIsFull",
42
+ "ErrInvalidTopicName",
43
+ "ErrMessageTooLarge",
44
+ "ErrNoDataFound",
45
+ "ErrSessionClosed",
46
+ "ErrUnAckSizeExceeded",
47
+ "ConsumerMessage",
48
+ "Messages",
49
+ "ProducerMessage",
50
+ "new_producer_message",
51
+ "new_producer",
52
+ "new_async_producer",
53
+ "new_consumer",
54
+ ]
@@ -0,0 +1,67 @@
1
+ from __future__ import annotations
2
+
3
+ import asyncio
4
+ import uuid
5
+ from dataclasses import dataclass, field
6
+ from typing import Any, Awaitable, Callable, Dict, List, Optional
7
+
8
+ from .constants import DelayMode
9
+ from .limiter import AsyncLimiter
10
+
11
+
12
+ ConsumeHandler = Callable[[asyncio.Queue], Awaitable[None]]
13
+
14
+
15
+ @dataclass
16
+ class Claim:
17
+ topics: List[str]
18
+ handle: Optional[ConsumeHandler] = None
19
+ delay: int = 0
20
+ filter: Optional[Dict[str, Any]] = None
21
+ consume_group: str = ""
22
+ max_load: int = 1
23
+ cid: str = field(default_factory=lambda: uuid.uuid4().hex)
24
+ msg: asyncio.Queue = field(default_factory=lambda: asyncio.Queue(maxsize=20000))
25
+ token_bucket: Optional[AsyncLimiter] = None
26
+ best_pos: int = 0
27
+
28
+ def set_max_load(self, max_load: int) -> None:
29
+ self.max_load = max_load
30
+
31
+ def ignore_data(self) -> None:
32
+ self.max_load = 0
33
+
34
+ def set_pop_limiter(self, size: int) -> None:
35
+ self.token_bucket = AsyncLimiter(1.0, size)
36
+
37
+ def set_custom_limiter(self, limiter: AsyncLimiter) -> None:
38
+ self.token_bucket = limiter
39
+
40
+ def set_delay(self, mode: DelayMode, delay_ms: Optional[int] = None) -> None:
41
+ if mode == DelayMode.IGNORE_DELAY:
42
+ self.delay = -1
43
+ elif mode == DelayMode.BASE_SERVER:
44
+ self.delay = 0
45
+ elif mode == DelayMode.TIME_DELAY:
46
+ if delay_ms is None or delay_ms <= 1:
47
+ raise ValueError("time value > 1ms required")
48
+ self.delay = int(delay_ms)
49
+
50
+ def try_filter(self, filter_data: Dict[str, Any]) -> None:
51
+ self.filter = filter_data
52
+
53
+ def set_consume_group(self, group: str) -> None:
54
+ self.consume_group = group
55
+
56
+ async def add_message(self, msg) -> None:
57
+ if self.token_bucket:
58
+ await self.token_bucket.wait()
59
+ await self.msg.put(msg)
60
+
61
+
62
+ def new_claim(topics: List[str], handle: Optional[ConsumeHandler]) -> Claim:
63
+ return Claim(topics=topics, handle=handle)
64
+
65
+
66
+ def new_opinion() -> Claim:
67
+ return Claim(topics=[], max_load=1)
@@ -0,0 +1,102 @@
1
+ from __future__ import annotations
2
+
3
+ import asyncio
4
+ import json
5
+ from typing import List, Optional, Tuple
6
+
7
+ from .cluster import partitions_from_cluster, watch_cluster
8
+ from .config import Config, Endpoint
9
+ from .constants import MessageType, VERSION_V1_2_0
10
+ from .errors import Error, new_error
11
+ from .partition import Partition
12
+ from .session import Session
13
+ from .slot import PeerSlot
14
+
15
+
16
+ class Client:
17
+ def __init__(self, conf: Config, slot: PeerSlot):
18
+ self.conf = conf
19
+ self.slot = slot
20
+ self.closed = False
21
+ self.stop_event = asyncio.Event()
22
+ self._watch_tasks: List[asyncio.Task] = []
23
+
24
+ async def close(self) -> None:
25
+ if self.closed:
26
+ return
27
+ self.closed = True
28
+ self.stop_event.set()
29
+ for task in self._watch_tasks:
30
+ task.cancel()
31
+ if self._watch_tasks:
32
+ await asyncio.gather(*self._watch_tasks, return_exceptions=True)
33
+ for part in self.slot.get_node_list():
34
+ while not part.sessions.empty():
35
+ sess = await part.sessions.get()
36
+ await sess.close()
37
+
38
+ def keep_retry(self, duration_sec: Optional[float] = None) -> None:
39
+ if duration_sec is None or duration_sec <= 0:
40
+ self.conf.retry_times = -1
41
+ return
42
+ interval = int(duration_sec / 0.5)
43
+ self.conf.retry_times = interval if duration_sec % 0.5 == 0 else interval + 1
44
+
45
+ def get_partition(self, slot: int) -> Optional[Partition]:
46
+ return self.slot.get_partition(slot)
47
+
48
+ def get_partitions(self) -> List[Partition]:
49
+ return self.slot.get_node_list()
50
+
51
+
52
+ async def new_client(raw_url: str, conf: Optional[Config] = None) -> Tuple[Optional[Client], Optional[Error]]:
53
+ conf = conf or Config()
54
+ conf.fix_config()
55
+ endpoints = []
56
+ for raw in raw_url.split(","):
57
+ endpoints.append(conf.parse_url(raw.strip()))
58
+ parts, err = await get_part_list(endpoints, conf)
59
+ if err:
60
+ return None, err
61
+ if not parts:
62
+ return None, new_error(Exception("no partitions available"))
63
+ slot = PeerSlot(conf)
64
+ err = await slot.initial(parts)
65
+ if err:
66
+ return None, err
67
+ client = Client(conf, slot)
68
+ for part in parts:
69
+ task = asyncio.create_task(
70
+ watch_cluster(part, slot.rebalance, stop_event=client.stop_event)
71
+ )
72
+ client._watch_tasks.append(task)
73
+ return client, None
74
+
75
+
76
+ async def get_part_list(
77
+ endpoints: List[Endpoint], conf: Config
78
+ ) -> Tuple[List[Partition], Optional[Error]]:
79
+ for ep in endpoints:
80
+ part = Partition(endpoint=ep, address=ep.address, no=0, conf=conf)
81
+ sess = Session(part, ep.version)
82
+ err = await sess.connect()
83
+ if err:
84
+ await sess.close()
85
+ continue
86
+ req = conf.new_protocol(ep.version)
87
+ resp, err = await sess.request(MessageType.PARTITION, req)
88
+ await sess.close()
89
+ if err or resp is None:
90
+ continue
91
+ try:
92
+ block = resp.block
93
+ if ep.version == VERSION_V1_2_0:
94
+ data = json.loads(block)
95
+ else:
96
+ data = json.loads(block)
97
+ parts = partitions_from_cluster(data, ep.version, conf, fallback_endpoint=ep)
98
+ if parts:
99
+ return parts, None
100
+ except Exception as exc:
101
+ return [], new_error(exc)
102
+ return [], new_error(Exception("no partitions available"))
@@ -0,0 +1,111 @@
1
+ from __future__ import annotations
2
+
3
+ import asyncio
4
+ from typing import Any, List, Optional
5
+
6
+ from .constants import MessageType, VERSION_V1_2_0
7
+ from .config import Endpoint
8
+ from .partition import Partition
9
+ from .session import Session
10
+
11
+
12
+ def _endpoint_from_node(node: dict, version: str, fallback: Optional[Endpoint]) -> Endpoint:
13
+ ep = node.get("endpoint") or {}
14
+ address = (ep.get("address") or node.get("address") or (fallback.address if fallback else "")).strip()
15
+ scheme = (ep.get("scheme") or (fallback.scheme if fallback else "esmd")).strip()
16
+ username = ep.get("username") or (fallback.username if fallback else "")
17
+ password = ep.get("password") or (fallback.password if fallback else "")
18
+ v = ep.get("version") or version or (fallback.version if fallback else "")
19
+ return Endpoint(address=address, scheme=scheme, username=username, password=password, version=v)
20
+
21
+
22
+ def partitions_from_cluster(
23
+ data: Any, version: str, conf, fallback_endpoint: Optional[Endpoint] = None
24
+ ) -> List[Partition]:
25
+ nodes = []
26
+ if version == VERSION_V1_2_0:
27
+ raw_nodes = data or []
28
+ else:
29
+ raw_nodes = (data or {}).get("nodes", [])
30
+ for node in raw_nodes:
31
+ endpoint = _endpoint_from_node(node, version, fallback_endpoint)
32
+ if not endpoint.address:
33
+ continue
34
+ part = Partition(
35
+ endpoint=endpoint,
36
+ address=node.get("address") or endpoint.address,
37
+ no=node.get("partition", 0),
38
+ master=node.get("master", False),
39
+ stat=node.get("stat", node.get("status", 0)),
40
+ conf=conf,
41
+ )
42
+ nodes.append(part)
43
+ return nodes
44
+
45
+
46
+ async def _sleep_or_stop(stop_event: Optional[asyncio.Event], seconds: float) -> None:
47
+ if stop_event is None:
48
+ await asyncio.sleep(seconds)
49
+ return
50
+ try:
51
+ await asyncio.wait_for(stop_event.wait(), timeout=seconds)
52
+ except asyncio.TimeoutError:
53
+ return
54
+
55
+
56
+ async def watch_cluster(
57
+ partition: Partition, rebalance_cb, stop_event: Optional[asyncio.Event] = None
58
+ ) -> None:
59
+ try:
60
+ while stop_event is None or not stop_event.is_set():
61
+ sess = Session(partition, partition.get_version())
62
+ err = await sess.connect()
63
+ if err:
64
+ await _sleep_or_stop(stop_event, 1)
65
+ continue
66
+
67
+ try:
68
+ req = partition.conf.new_protocol(partition.get_version())
69
+ _, req_err = await sess.request(MessageType.CLUSTER, req)
70
+ if req_err:
71
+ await _sleep_or_stop(stop_event, 1)
72
+ continue
73
+
74
+ while stop_event is None or not stop_event.is_set():
75
+ data, err = await sess.read_json()
76
+ if err:
77
+ break
78
+ if not data or data.get("code") != 0:
79
+ continue
80
+ block = data.get("block")
81
+ try:
82
+ import base64
83
+
84
+ if isinstance(block, str):
85
+ block = base64.b64decode(block)
86
+ except Exception:
87
+ pass
88
+ try:
89
+ import json
90
+
91
+ decoded = (
92
+ json.loads(block)
93
+ if isinstance(block, (bytes, bytearray))
94
+ else block
95
+ )
96
+ except Exception:
97
+ decoded = block
98
+ parts = partitions_from_cluster(
99
+ decoded,
100
+ partition.get_version(),
101
+ partition.conf,
102
+ fallback_endpoint=partition.endpoint,
103
+ )
104
+ if parts:
105
+ rebalance_cb(parts)
106
+ finally:
107
+ await sess.close()
108
+
109
+ await _sleep_or_stop(stop_event, 1)
110
+ except asyncio.CancelledError:
111
+ return
@@ -0,0 +1,64 @@
1
+ from __future__ import annotations
2
+
3
+ import gzip
4
+ import zlib
5
+ from dataclasses import dataclass
6
+ from typing import Optional
7
+
8
+ try:
9
+ import lz4.frame as _lz4_frame
10
+ except Exception: # pragma: no cover - optional dependency
11
+ _lz4_frame = None
12
+
13
+
14
+ class CompressType:
15
+ GZIP = 1
16
+ ZLIB = 2
17
+ LZ4 = 3
18
+
19
+
20
+ class CompressLevel:
21
+ NORMAL = 0
22
+ LEVEL1 = 1
23
+ LEVEL2 = 2
24
+ LEVEL3 = 3
25
+ LEVEL4 = 4
26
+ LEVEL5 = 5
27
+ LEVEL6 = 6
28
+ LEVEL7 = 7
29
+ LEVEL8 = 8
30
+ LEVEL9 = 9
31
+
32
+
33
+ @dataclass
34
+ class Compression:
35
+ zip_type: int
36
+ level: int
37
+
38
+ def compress(self, data: bytes) -> bytes:
39
+ if self.zip_type == CompressType.GZIP:
40
+ return gzip.compress(data, compresslevel=self.level or 1)
41
+ if self.zip_type == CompressType.ZLIB:
42
+ return zlib.compress(data, level=self.level or 1)
43
+ if self.zip_type == CompressType.LZ4:
44
+ if _lz4_frame is None:
45
+ raise RuntimeError("lz4 is not installed")
46
+ return _lz4_frame.compress(data)
47
+ return data
48
+
49
+ def decompress(self, data: bytes) -> bytes:
50
+ if self.zip_type == CompressType.GZIP:
51
+ return gzip.decompress(data)
52
+ if self.zip_type == CompressType.ZLIB:
53
+ return zlib.decompress(data)
54
+ if self.zip_type == CompressType.LZ4:
55
+ if _lz4_frame is None:
56
+ raise RuntimeError("lz4 is not installed")
57
+ return _lz4_frame.decompress(data)
58
+ return data
59
+
60
+
61
+ def new_compress(zip_type: int, level: Optional[int] = None) -> Optional[Compression]:
62
+ if zip_type not in (CompressType.GZIP, CompressType.ZLIB, CompressType.LZ4):
63
+ return None
64
+ return Compression(zip_type=zip_type, level=level or 0)