prole/knoe/ekosystem.py

"""knoe.ekosystem — Python counterpart to knoe.ekosystem_id() in PostgreSQL.

Bit layout (82 bits → 16 base36 chars):
  [49: timestamp_ms][12: tenant_id][10: shard_id][11: seq]

Epoch: 2026-01-01 00:00:00 UTC  (1767225600000 ms)

Produces identical output to knoe.ekosystem_id() for the same inputs and
timestamp. Use this when a PostgreSQL connection is unavailable (e.g. from
ClickHouse or SQLite application nodes).
"""

import threading
import time
from datetime import datetime, timezone
from typing import Optional

EPOCH_MS = 1767225600000  # 2026-01-01 00:00:00 UTC
ALPHABET = "0123456789abcdefghijklmnopqrstuvwxyz"

# Bit widths
_TS_BITS = 49
_TENANT_BITS = 12
_SHARD_BITS = 10
_SEQ_BITS = 11
_TOTAL_BITS = _TS_BITS + _TENANT_BITS + _SHARD_BITS + _SEQ_BITS  # 82

# Shift offsets
_SEQ_SHIFT = 0
_SHARD_SHIFT = _SEQ_BITS
_TENANT_SHIFT = _SEQ_BITS + _SHARD_BITS
_TS_SHIFT = _SEQ_BITS + _SHARD_BITS + _TENANT_BITS

# Masks
_SEQ_MASK = (1 << _SEQ_BITS) - 1
_SHARD_MASK = (1 << _SHARD_BITS) - 1
_TENANT_MASK = (1 << _TENANT_BITS) - 1
_TS_MASK = (1 << _TS_BITS) - 1

_BASE = len(ALPHABET)
_UUID_LEN = 16  # ceil(82 / log2(36)) = 16


def _to_base36(n: int) -> str:
    """Encode integer n as a zero-padded 16-character base36 string."""
    chars = []
    for _ in range(_UUID_LEN):
        chars.append(ALPHABET[n % _BASE])
        n //= _BASE
    return "".join(reversed(chars))


def _from_base36(s: str) -> int:
    """Decode a base36 string to an integer."""
    n = 0
    for ch in s:
        n = n * _BASE + ALPHABET.index(ch)
    return n


class EkosystemID:
    """Thread-safe ekosystem UUID generator (Snowflake pattern).

    Args:
        tenant_id: Integer tenant identifier (0 = knoe-db root, 1 = prole, …)
        shard_id:  Integer shard identifier within the tenant (usually 0)
    """

    def __init__(self, tenant_id: int, shard_id: int) -> None:
        if not (0 <= tenant_id < (1 << _TENANT_BITS)):
            raise ValueError(f"tenant_id must be 0–{(1 << _TENANT_BITS) - 1}")
        if not (0 <= shard_id < (1 << _SHARD_BITS)):
            raise ValueError(f"shard_id must be 0–{(1 << _SHARD_BITS) - 1}")
        self.tenant_id = tenant_id
        self.shard_id = shard_id
        self._lock = threading.Lock()
        self._last_ms: int = -1
        self._seq: int = 0

    def generate(self) -> str:
        """Generate a new ekosystem UUID string (thread-safe)."""
        with self._lock:
            now_ms = int(time.time() * 1000) - EPOCH_MS
            if now_ms < 0:
                raise RuntimeError("System clock is before the ekosystem epoch (2026-01-01)")
            if now_ms == self._last_ms:
                self._seq = (self._seq + 1) & _SEQ_MASK
                if self._seq == 0:
                    # Sequence exhausted — spin until next millisecond
                    while now_ms == self._last_ms:
                        now_ms = int(time.time() * 1000) - EPOCH_MS
            else:
                self._seq = 0
                self._last_ms = now_ms

            bits = (
                ((now_ms & _TS_MASK) << _TS_SHIFT)
                | ((self.tenant_id & _TENANT_MASK) << _TENANT_SHIFT)
                | ((self.shard_id & _SHARD_MASK) << _SHARD_SHIFT)
                | (self._seq & _SEQ_MASK)
            )
        return _to_base36(bits)

    def project_prefix(self) -> str:
        """Return a project namespace prefix: 'p_{uuid}_'."""
        return f"p_{self.generate()}_"

    def object_name(self, kind: str) -> str:
        """Return a namespaced object name: '{kind}_{uuid}'."""
        return f"{kind}_{self.generate()}"


def decode(uid: str) -> dict:
    """Decode an ekosystem UUID string into its component fields.

    Returns:
        dict with keys: tenant_id, shard_id, ts_ms, seq, issued_at (UTC datetime)
    """
    if len(uid) != _UUID_LEN:
        raise ValueError(f"ekosystem UUID must be {_UUID_LEN} characters, got {len(uid)}")
    bits = _from_base36(uid)
    ts_ms = (bits >> _TS_SHIFT) & _TS_MASK
    tenant_id = (bits >> _TENANT_SHIFT) & _TENANT_MASK
    shard_id = (bits >> _SHARD_SHIFT) & _SHARD_MASK
    seq = bits & _SEQ_MASK
    issued_at = datetime.fromtimestamp((ts_ms + EPOCH_MS) / 1000.0, tz=timezone.utc)
    return {
        "tenant_id": tenant_id,
        "shard_id": shard_id,
        "ts_ms": ts_ms,
        "seq": seq,
        "issued_at": issued_at,
    }


def can_access(uid: str, kind: str, tenant_id: int, grants: list) -> bool:
    """Check whether tenant_id can access object uid of the given kind.

    Args:
        uid:       ekosystem UUID of the object
        kind:      object kind string (e.g. 'table', 'embedding')
        tenant_id: the requesting tenant's integer ID
        grants:    list of grant dicts with keys 'object_uuid', 'kind', 'grantee_tenant_id'

    Returns:
        True if the object's own tenant matches tenant_id, or a matching grant exists.
    """
    info = decode(uid)
    if info["tenant_id"] == tenant_id:
        return True
    for g in grants:
        if (
            g.get("object_uuid") == uid
            and g.get("kind") == kind
            and g.get("grantee_tenant_id") == tenant_id
        ):
            return True
    return False