parent
3f89e7aea6
commit
a93961290e
@ -1,4 +1,4 @@
|
|||||||
# SPDX-FileCopyrightText: 2025-present NoahLan <6995syu@163.com>
|
# SPDX-FileCopyrightText: 2025-present NoahLan <6995syu@163.com>
|
||||||
#
|
#
|
||||||
# SPDX-License-Identifier: MIT
|
# SPDX-License-Identifier: MIT
|
||||||
__version__ = "0.4.0"
|
__version__ = "0.4.1"
|
||||||
|
|||||||
@ -0,0 +1,104 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import sqlite3
|
||||||
|
import threading
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Protocol
|
||||||
|
|
||||||
|
|
||||||
|
class MQOffsetStore(Protocol):
|
||||||
|
def get(self, consumer_name: str, topic: str, partition: int) -> int | None:
|
||||||
|
"""Return the next offset to consume."""
|
||||||
|
|
||||||
|
def set(self, consumer_name: str, topic: str, partition: int, offset: int) -> None:
|
||||||
|
"""Persist the next offset to consume."""
|
||||||
|
|
||||||
|
def close(self) -> None:
|
||||||
|
"""Release store resources."""
|
||||||
|
|
||||||
|
|
||||||
|
class SQLiteMQOffsetStore:
|
||||||
|
def __init__(self, path: str | Path) -> None:
|
||||||
|
self.path = Path(path)
|
||||||
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
self._lock = threading.Lock()
|
||||||
|
self._conn = sqlite3.connect(self.path, check_same_thread=False)
|
||||||
|
self._closed = False
|
||||||
|
self._conn.execute(
|
||||||
|
"""
|
||||||
|
CREATE TABLE IF NOT EXISTS mq_offsets (
|
||||||
|
consumer_name TEXT NOT NULL,
|
||||||
|
topic TEXT NOT NULL,
|
||||||
|
partition INTEGER NOT NULL,
|
||||||
|
offset INTEGER NOT NULL,
|
||||||
|
updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||||
|
PRIMARY KEY (consumer_name, topic, partition)
|
||||||
|
)
|
||||||
|
"""
|
||||||
|
)
|
||||||
|
self._conn.commit()
|
||||||
|
|
||||||
|
def get(self, consumer_name: str, topic: str, partition: int) -> int | None:
|
||||||
|
with self._lock:
|
||||||
|
if self._closed:
|
||||||
|
return None
|
||||||
|
row = self._conn.execute(
|
||||||
|
"""
|
||||||
|
SELECT offset
|
||||||
|
FROM mq_offsets
|
||||||
|
WHERE consumer_name = ? AND topic = ? AND partition = ?
|
||||||
|
""",
|
||||||
|
(consumer_name, topic, partition),
|
||||||
|
).fetchone()
|
||||||
|
if row is None:
|
||||||
|
return None
|
||||||
|
return int(row[0])
|
||||||
|
|
||||||
|
def set(self, consumer_name: str, topic: str, partition: int, offset: int) -> None:
|
||||||
|
with self._lock:
|
||||||
|
if self._closed:
|
||||||
|
return
|
||||||
|
self._conn.execute(
|
||||||
|
"""
|
||||||
|
INSERT INTO mq_offsets (consumer_name, topic, partition, offset, updated_at)
|
||||||
|
VALUES (?, ?, ?, ?, CURRENT_TIMESTAMP)
|
||||||
|
ON CONFLICT(consumer_name, topic, partition)
|
||||||
|
DO UPDATE SET offset = excluded.offset, updated_at = CURRENT_TIMESTAMP
|
||||||
|
""",
|
||||||
|
(consumer_name, topic, partition, offset),
|
||||||
|
)
|
||||||
|
self._conn.commit()
|
||||||
|
|
||||||
|
def close(self) -> None:
|
||||||
|
with self._lock:
|
||||||
|
if self._closed:
|
||||||
|
return
|
||||||
|
self._conn.close()
|
||||||
|
self._closed = True
|
||||||
|
|
||||||
|
|
||||||
|
class MemoryMQOffsetStore:
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self._offsets: dict[tuple[str, str, int], int] = {}
|
||||||
|
self._lock = threading.Lock()
|
||||||
|
|
||||||
|
def get(self, consumer_name: str, topic: str, partition: int) -> int | None:
|
||||||
|
with self._lock:
|
||||||
|
return self._offsets.get((consumer_name, topic, partition))
|
||||||
|
|
||||||
|
def set(self, consumer_name: str, topic: str, partition: int, offset: int) -> None:
|
||||||
|
with self._lock:
|
||||||
|
self._offsets[(consumer_name, topic, partition)] = offset
|
||||||
|
|
||||||
|
def close(self) -> None:
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
def create_offset_store(config: dict | None, *, default_path: str | Path) -> MQOffsetStore:
|
||||||
|
config = dict(config or {})
|
||||||
|
store_type = str(config.get("type", "sqlite")).strip().lower()
|
||||||
|
if store_type == "memory":
|
||||||
|
return MemoryMQOffsetStore()
|
||||||
|
if store_type == "sqlite":
|
||||||
|
return SQLiteMQOffsetStore(config.get("path") or default_path)
|
||||||
|
raise ValueError(f"unsupported mq offset store type: {store_type!r}")
|
||||||
Loading…
Reference in New Issue