Rename the import package app to fluksio
A wheel whose top-level module is `app` collides with anything else in a user's venv, so the package that is about to be published takes the name it is published under. Only the Python package moves; the repo, the Docker WORKDIR and the compose project keep theirs. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Claude Opus 5
@@ -0,0 +1,605 @@
|
||||
"""
|
||||
State management module with pluggable backends.
|
||||
|
||||
This module provides state storage backends for pipeline execution,
|
||||
supporting both in-memory storage and Redis for distributed execution.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from abc import ABC, abstractmethod
|
||||
from collections import deque
|
||||
from collections.abc import Iterator
|
||||
from contextlib import contextmanager
|
||||
from threading import RLock
|
||||
from typing import Any, cast
|
||||
|
||||
import redis
|
||||
|
||||
# A sparkline only means something for numbers, so the history keeps the values
|
||||
# it can plot and nothing else. 120 points fill a panel-wide chart while leaving
|
||||
# Redis a cache rather than a time-series database.
|
||||
HISTORY_LIMIT = 120
|
||||
|
||||
|
||||
def as_number(value: Any) -> float | None:
|
||||
"""The plottable form of a value, or None if it is not a number."""
|
||||
# bool is an int subclass; a flag is not a measurement.
|
||||
if isinstance(value, bool) or not isinstance(value, (int, float)):
|
||||
return None
|
||||
return float(value)
|
||||
|
||||
|
||||
class StateBackend(ABC):
|
||||
"""
|
||||
Abstract base class for state storage backends.
|
||||
|
||||
Provides a dict-like interface for storing and retrieving state values
|
||||
with support for distributed locking.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def get(self, key: str, default: Any = None) -> Any:
|
||||
"""
|
||||
Get a value by key.
|
||||
|
||||
:param key: The key to retrieve.
|
||||
:type key: str
|
||||
:param default: Default value if key not found.
|
||||
:type default: Any
|
||||
:returns: The stored value or default.
|
||||
:rtype: Any
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def set(self, key: str, value: Any) -> None:
|
||||
"""
|
||||
Set a value by key.
|
||||
|
||||
:param key: The key to set.
|
||||
:type key: str
|
||||
:param value: The value to store.
|
||||
:type value: Any
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def exists(self, key: str) -> bool:
|
||||
"""
|
||||
Check if a key exists.
|
||||
|
||||
:param key: The key to check.
|
||||
:type key: str
|
||||
:returns: True if key exists.
|
||||
:rtype: bool
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def delete(self, key: str) -> None:
|
||||
"""Forget one key and any history kept for it."""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def clear(self) -> None:
|
||||
"""Clear all keys in the state."""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def keys(self) -> list[str]:
|
||||
"""
|
||||
Get all keys.
|
||||
|
||||
:returns: List of all keys.
|
||||
:rtype: list[str]
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def update(self, mapping: dict[str, Any]) -> None:
|
||||
"""
|
||||
Update multiple keys atomically.
|
||||
|
||||
:param mapping: Dict of key-value pairs to set.
|
||||
:type mapping: dict[str, Any]
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
@contextmanager
|
||||
def lock(self, timeout: float = 10.0) -> Iterator[None]:
|
||||
"""
|
||||
Acquire a lock for thread-safe operations.
|
||||
|
||||
:param timeout: Lock timeout in seconds.
|
||||
:type timeout: float
|
||||
:yields: Lock context.
|
||||
"""
|
||||
...
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# Atomic operations for synchronous node support
|
||||
# -------------------------------------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
def increment(self, key: str) -> int:
|
||||
"""
|
||||
Atomically increment a counter and return the new value.
|
||||
|
||||
:param key: The key to increment.
|
||||
:type key: str
|
||||
:returns: The new value after incrementing.
|
||||
:rtype: int
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def get_multi(self, keys: list[str]) -> dict[str, Any]:
|
||||
"""
|
||||
Get multiple values atomically.
|
||||
|
||||
:param keys: List of keys to retrieve.
|
||||
:type keys: list[str]
|
||||
:returns: Dict mapping keys to values (missing keys have None).
|
||||
:rtype: dict[str, Any]
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def compare_and_swap_multi(
|
||||
self,
|
||||
expected: dict[str, Any],
|
||||
updates: dict[str, Any],
|
||||
) -> bool:
|
||||
"""
|
||||
Atomically check expected values and apply updates if all match.
|
||||
|
||||
This is the key primitive for race-condition-free synchronous node execution.
|
||||
If any expected value doesn't match the current value, no updates are applied.
|
||||
|
||||
:param expected: Dict of keys to their expected values.
|
||||
:type expected: dict[str, Any]
|
||||
:param updates: Dict of keys to their new values (applied only if expected matches).
|
||||
:type updates: dict[str, Any]
|
||||
:returns: True if the swap succeeded, False if expected values didn't match.
|
||||
:rtype: bool
|
||||
"""
|
||||
...
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# Message history
|
||||
# -------------------------------------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
def append_history(
|
||||
self, values: dict[str, Any], ts: float, limits: dict[str, int] | None = None
|
||||
) -> None:
|
||||
"""
|
||||
Record freshly published values in each message's capped history.
|
||||
|
||||
Non-numeric payloads are skipped — the history exists to be plotted.
|
||||
|
||||
:param values: Message names mapped to the value just published.
|
||||
:type values: dict[str, Any]
|
||||
:param ts: When they were published.
|
||||
:param limits: How many points to keep per message, where a chart asks
|
||||
for more than the default.
|
||||
:type ts: float
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def history(self, key: str) -> list[tuple[float, float]]:
|
||||
"""
|
||||
The recorded ``(timestamp, value)`` pairs of one message, oldest first.
|
||||
|
||||
:param key: The message name.
|
||||
:type key: str
|
||||
:returns: At most ``HISTORY_LIMIT`` points; empty if nothing was recorded.
|
||||
:rtype: list[tuple[float, float]]
|
||||
"""
|
||||
...
|
||||
|
||||
def __contains__(self, key: str) -> bool:
|
||||
return self.exists(key)
|
||||
|
||||
def __getitem__(self, key: str) -> Any:
|
||||
if not self.exists(key):
|
||||
raise KeyError(key)
|
||||
return self.get(key)
|
||||
|
||||
def __setitem__(self, key: str, value: Any) -> None:
|
||||
self.set(key, value)
|
||||
|
||||
def __iter__(self) -> Iterator[str]:
|
||||
return iter(self.keys())
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.keys())
|
||||
|
||||
|
||||
class MemoryState(StateBackend):
|
||||
"""
|
||||
In-memory state backend using a thread-safe dict.
|
||||
|
||||
Suitable for single-process execution and testing.
|
||||
|
||||
:example:
|
||||
>>> state = MemoryState()
|
||||
>>> state["key"] = "value"
|
||||
>>> state["key"]
|
||||
'value'
|
||||
"""
|
||||
|
||||
__slots__ = ("_data", "_history", "_lock")
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._data: dict[str, Any] = {}
|
||||
self._history: dict[str, deque[tuple[float, float]]] = {}
|
||||
self._lock = RLock() # Reentrant lock for nested access
|
||||
|
||||
def get(self, key: str, default: Any = None) -> Any:
|
||||
with self._lock:
|
||||
return self._data.get(key, default)
|
||||
|
||||
def set(self, key: str, value: Any) -> None:
|
||||
with self._lock:
|
||||
self._data[key] = value
|
||||
|
||||
def exists(self, key: str) -> bool:
|
||||
with self._lock:
|
||||
return key in self._data
|
||||
|
||||
def delete(self, key: str) -> None:
|
||||
with self._lock:
|
||||
self._data.pop(key, None)
|
||||
self._history.pop(key, None)
|
||||
|
||||
def clear(self) -> None:
|
||||
with self._lock:
|
||||
self._data.clear()
|
||||
self._history.clear()
|
||||
|
||||
def keys(self) -> list[str]:
|
||||
with self._lock:
|
||||
return list(self._data.keys())
|
||||
|
||||
def update(self, mapping: dict[str, Any]) -> None:
|
||||
with self._lock:
|
||||
self._data.update(mapping)
|
||||
|
||||
@contextmanager
|
||||
def lock(self, timeout: float = 10.0) -> Iterator[None]:
|
||||
"""Acquire the internal lock."""
|
||||
acquired = self._lock.acquire(timeout=timeout)
|
||||
if not acquired:
|
||||
raise TimeoutError("Failed to acquire lock")
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
self._lock.release()
|
||||
|
||||
def increment(self, key: str) -> int:
|
||||
"""Atomically increment a counter and return the new value."""
|
||||
with self._lock:
|
||||
value: int = self._data.get(key, 0)
|
||||
new_value = value + 1
|
||||
self._data[key] = new_value
|
||||
return new_value
|
||||
|
||||
def get_multi(self, keys: list[str]) -> dict[str, Any]:
|
||||
"""Get multiple values atomically."""
|
||||
with self._lock:
|
||||
return {k: self._data.get(k) for k in keys}
|
||||
|
||||
def compare_and_swap_multi(
|
||||
self,
|
||||
expected: dict[str, Any],
|
||||
updates: dict[str, Any],
|
||||
) -> bool:
|
||||
"""Atomically check expected values and apply updates if all match."""
|
||||
with self._lock:
|
||||
# Check all expected values
|
||||
for key, exp_value in expected.items():
|
||||
actual = self._data.get(key)
|
||||
if actual != exp_value:
|
||||
return False
|
||||
|
||||
# All matched, apply updates
|
||||
self._data.update(updates)
|
||||
return True
|
||||
|
||||
def append_history(
|
||||
self, values: dict[str, Any], ts: float, limits: dict[str, int] | None = None
|
||||
) -> None:
|
||||
"""Append the numeric values to their message's bounded series."""
|
||||
with self._lock:
|
||||
for key, value in values.items():
|
||||
number = as_number(value)
|
||||
if number is None:
|
||||
continue
|
||||
cap = (limits or {}).get(key, HISTORY_LIMIT)
|
||||
series = self._history.get(key)
|
||||
if series is None or series.maxlen != cap:
|
||||
# A widget asking for a deeper series re-caps it in place.
|
||||
series = deque(series or (), maxlen=cap)
|
||||
self._history[key] = series
|
||||
series.append((ts, number))
|
||||
|
||||
def history(self, key: str) -> list[tuple[float, float]]:
|
||||
"""The recorded points of one message, oldest first."""
|
||||
with self._lock:
|
||||
return list(self._history.get(key, ()))
|
||||
|
||||
|
||||
class RedisState(StateBackend):
|
||||
"""
|
||||
Redis-based state backend for distributed execution.
|
||||
|
||||
Values are stored as JSON — anything a node passes through must be
|
||||
serializable. Uses a namespace prefix to isolate pipeline executions.
|
||||
|
||||
:param host: Redis host address.
|
||||
:type host: str
|
||||
:param port: Redis port number.
|
||||
:type port: int
|
||||
:param db: Redis database number.
|
||||
:type db: int
|
||||
:param namespace: Key prefix for isolation.
|
||||
:type namespace: str
|
||||
:param password: Redis password (optional).
|
||||
:type password: str | None
|
||||
:param ttl: Time-to-live for keys in seconds (optional).
|
||||
:type ttl: int | None
|
||||
|
||||
:example:
|
||||
>>> state = RedisState(host="localhost", namespace="my_pipeline")
|
||||
>>> state["result"] = {"value": 42}
|
||||
>>> state["result"]
|
||||
{'value': 42}
|
||||
"""
|
||||
|
||||
__slots__ = ("_client", "_namespace", "_ttl", "_lock_name")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
host: str = "localhost",
|
||||
port: int = 6379,
|
||||
db: int = 0,
|
||||
namespace: str = "pipeline",
|
||||
password: str | None = None,
|
||||
ttl: int | None = None,
|
||||
) -> None:
|
||||
self._client = redis.Redis(
|
||||
host=host,
|
||||
port=port,
|
||||
db=db,
|
||||
password=password,
|
||||
decode_responses=False,
|
||||
)
|
||||
self._namespace = namespace
|
||||
self._ttl = ttl
|
||||
self._lock_name = f"{namespace}:_lock"
|
||||
|
||||
def _key(self, key: str) -> str:
|
||||
"""Generate namespaced key."""
|
||||
return f"{self._namespace}:{key}"
|
||||
|
||||
def _serialize(self, value: Any) -> bytes:
|
||||
"""Serialize value for storage."""
|
||||
return json.dumps(value).encode()
|
||||
|
||||
def _deserialize(self, data: bytes | None) -> Any:
|
||||
"""Deserialize value from storage."""
|
||||
return json.loads(data) if data else None
|
||||
|
||||
# redis-py types every command as a sync/async union; this is the
|
||||
# synchronous client, so the results are narrowed where they are consumed.
|
||||
def get(self, key: str, default: Any = None) -> Any:
|
||||
data = cast(bytes | None, self._client.get(self._key(key)))
|
||||
return self._deserialize(data) if data is not None else default
|
||||
|
||||
def set(self, key: str, value: Any) -> None:
|
||||
data = self._serialize(value)
|
||||
if self._ttl:
|
||||
self._client.setex(self._key(key), self._ttl, data)
|
||||
else:
|
||||
self._client.set(self._key(key), data)
|
||||
|
||||
def exists(self, key: str) -> bool:
|
||||
return bool(self._client.exists(self._key(key)))
|
||||
|
||||
def delete(self, key: str) -> None:
|
||||
self._client.delete(self._key(key), self._history_key(key))
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear all keys in the namespace."""
|
||||
pattern = f"{self._namespace}:*"
|
||||
cursor = 0
|
||||
while True:
|
||||
cursor, keys = cast(
|
||||
tuple[int, list[bytes]],
|
||||
self._client.scan(cursor, match=pattern, count=100),
|
||||
)
|
||||
if keys:
|
||||
self._client.delete(*keys)
|
||||
if cursor == 0:
|
||||
break
|
||||
|
||||
def keys(self) -> list[str]:
|
||||
"""Get all keys in the namespace (without prefix)."""
|
||||
pattern = f"{self._namespace}:*"
|
||||
prefix_len = len(self._namespace) + 1
|
||||
result = []
|
||||
cursor = 0
|
||||
while True:
|
||||
cursor, keys = cast(
|
||||
tuple[int, list[bytes]],
|
||||
self._client.scan(cursor, match=pattern, count=100),
|
||||
)
|
||||
for k in keys:
|
||||
key_str = k.decode("utf-8")
|
||||
# Skip the lock key
|
||||
if not key_str.endswith(":_lock"):
|
||||
result.append(key_str[prefix_len:])
|
||||
if cursor == 0:
|
||||
break
|
||||
return result
|
||||
|
||||
def update(self, mapping: dict[str, Any]) -> None:
|
||||
"""Update multiple keys atomically using a pipeline."""
|
||||
if not mapping:
|
||||
return
|
||||
pipe = self._client.pipeline()
|
||||
for key, value in mapping.items():
|
||||
data = self._serialize(value)
|
||||
if self._ttl:
|
||||
pipe.setex(self._key(key), self._ttl, data)
|
||||
else:
|
||||
pipe.set(self._key(key), data)
|
||||
pipe.execute()
|
||||
|
||||
@contextmanager
|
||||
def lock(self, timeout: float = 10.0) -> Iterator[None]:
|
||||
"""
|
||||
Acquire a distributed lock using Redis.
|
||||
|
||||
:param timeout: Lock timeout in seconds.
|
||||
:type timeout: float
|
||||
:yields: Lock context.
|
||||
:raises TimeoutError: If lock cannot be acquired.
|
||||
"""
|
||||
redis_lock = self._client.lock(
|
||||
self._lock_name,
|
||||
timeout=timeout,
|
||||
blocking_timeout=timeout,
|
||||
)
|
||||
acquired = redis_lock.acquire()
|
||||
if not acquired:
|
||||
raise TimeoutError("Failed to acquire Redis lock")
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
try:
|
||||
redis_lock.release()
|
||||
except redis.exceptions.LockError:
|
||||
pass # Lock may have expired
|
||||
|
||||
def ping(self) -> bool:
|
||||
"""
|
||||
Check if Redis connection is available.
|
||||
|
||||
:returns: True if connected.
|
||||
:rtype: bool
|
||||
"""
|
||||
try:
|
||||
return bool(self._client.ping())
|
||||
except redis.ConnectionError:
|
||||
return False
|
||||
|
||||
def increment(self, key: str) -> int:
|
||||
"""Atomically increment a counter using Redis INCR."""
|
||||
return cast(int, self._client.incr(self._key(key)))
|
||||
|
||||
def get_multi(self, keys: list[str]) -> dict[str, Any]:
|
||||
"""Get multiple values atomically using Redis MGET."""
|
||||
if not keys:
|
||||
return {}
|
||||
|
||||
full_keys = [self._key(k) for k in keys]
|
||||
values = cast(list[bytes | None], self._client.mget(full_keys))
|
||||
|
||||
result = {}
|
||||
for key, value in zip(keys, values, strict=True):
|
||||
result[key] = self._deserialize(value) if value is not None else None
|
||||
|
||||
return result
|
||||
|
||||
def compare_and_swap_multi(
|
||||
self,
|
||||
expected: dict[str, Any],
|
||||
updates: dict[str, Any],
|
||||
) -> bool:
|
||||
"""
|
||||
Atomically check expected values and apply updates using Redis WATCH/MULTI.
|
||||
|
||||
Uses optimistic locking with WATCH to detect concurrent modifications.
|
||||
If any watched key changes before EXEC, the transaction is aborted.
|
||||
|
||||
:param expected: Dict of keys to their expected values.
|
||||
:type expected: dict[str, Any]
|
||||
:param updates: Dict of keys to their new values.
|
||||
:type updates: dict[str, Any]
|
||||
:returns: True if the swap succeeded, False otherwise.
|
||||
:rtype: bool
|
||||
"""
|
||||
if not expected:
|
||||
# No expectations, just set the updates
|
||||
self.update(updates)
|
||||
return True
|
||||
|
||||
watch_keys = [self._key(k) for k in expected.keys()]
|
||||
|
||||
try:
|
||||
# Start watching the keys for changes
|
||||
self._client.watch(*watch_keys)
|
||||
|
||||
# Check current values
|
||||
current = self.get_multi(list(expected.keys()))
|
||||
|
||||
for key, exp_value in expected.items():
|
||||
if current.get(key) != exp_value:
|
||||
self._client.unwatch()
|
||||
return False
|
||||
|
||||
# All values match - execute the transaction
|
||||
pipe = self._client.pipeline(True) # True = use MULTI/EXEC
|
||||
|
||||
for key, value in updates.items():
|
||||
data = self._serialize(value)
|
||||
if self._ttl:
|
||||
pipe.setex(self._key(key), self._ttl, data)
|
||||
else:
|
||||
pipe.set(self._key(key), data)
|
||||
|
||||
pipe.execute()
|
||||
return True
|
||||
|
||||
except redis.WatchError:
|
||||
# Another client modified one of the watched keys
|
||||
return False
|
||||
|
||||
def _history_key(self, key: str) -> str:
|
||||
return self._key(f"__history__:{key}")
|
||||
|
||||
def append_history(
|
||||
self, values: dict[str, Any], ts: float, limits: dict[str, int] | None = None
|
||||
) -> None:
|
||||
"""Push the numeric values onto their capped list, in one round-trip."""
|
||||
pipe = self._client.pipeline()
|
||||
queued = False
|
||||
for key, value in values.items():
|
||||
number = as_number(value)
|
||||
if number is None:
|
||||
continue
|
||||
history_key = self._history_key(key)
|
||||
cap = (limits or {}).get(key, HISTORY_LIMIT)
|
||||
pipe.lpush(history_key, json.dumps([ts, number]))
|
||||
pipe.ltrim(history_key, 0, cap - 1)
|
||||
if self._ttl:
|
||||
pipe.expire(history_key, self._ttl)
|
||||
queued = True
|
||||
if queued:
|
||||
pipe.execute()
|
||||
|
||||
def history(self, key: str) -> list[tuple[float, float]]:
|
||||
"""The recorded points of one message, oldest first."""
|
||||
entries = cast(list[bytes], self._client.lrange(self._history_key(key), 0, -1))
|
||||
# LPUSH puts the newest first, a chart reads the other way round.
|
||||
points: list[tuple[float, float]] = []
|
||||
for entry in reversed(entries):
|
||||
ts, value = json.loads(entry)
|
||||
points.append((float(ts), float(value)))
|
||||
return points
|
||||
Reference in New Issue
Block a user