omkit 0.0.2__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.
- omkit-0.0.2/PKG-INFO +29 -0
- omkit-0.0.2/README.md +17 -0
- omkit-0.0.2/omkit/__init__.py +18 -0
- omkit-0.0.2/omkit/cleanup.py +62 -0
- omkit-0.0.2/omkit/config.py +60 -0
- omkit-0.0.2/omkit/cost.py +78 -0
- omkit-0.0.2/omkit/data/__init__.py +33 -0
- omkit-0.0.2/omkit/dbpool.py +139 -0
- omkit-0.0.2/omkit/encryption.py +58 -0
- omkit-0.0.2/omkit/eventbus.py +360 -0
- omkit-0.0.2/omkit/events.py +23 -0
- omkit-0.0.2/omkit/health.py +66 -0
- omkit-0.0.2/omkit/http.py +82 -0
- omkit-0.0.2/omkit/internal/__init__.py +7 -0
- omkit-0.0.2/omkit/internal/crypto.py +17 -0
- omkit-0.0.2/omkit/jobqueue/__init__.py +28 -0
- omkit-0.0.2/omkit/jobqueue/envelope.py +116 -0
- omkit-0.0.2/omkit/jobqueue/streaq.py +267 -0
- omkit-0.0.2/omkit/logging.py +77 -0
- omkit-0.0.2/omkit/metrics.py +41 -0
- omkit-0.0.2/omkit/model_lifecycle.py +192 -0
- omkit-0.0.2/omkit/platform/__init__.py +18 -0
- omkit-0.0.2/omkit/providers/__init__.py +11 -0
- omkit-0.0.2/omkit/providers/base.py +76 -0
- omkit-0.0.2/omkit/providers/registry.py +263 -0
- omkit-0.0.2/omkit/py.typed +0 -0
- omkit-0.0.2/omkit/quota.py +186 -0
- omkit-0.0.2/omkit/resilience.py +122 -0
- omkit-0.0.2/omkit/sanitize.py +122 -0
- omkit-0.0.2/omkit/security/__init__.py +28 -0
- omkit-0.0.2/omkit/security/events.py +79 -0
- omkit-0.0.2/omkit/sessions.py +301 -0
- omkit-0.0.2/omkit/settings.py +348 -0
- omkit-0.0.2/omkit/sync_notifier.py +110 -0
- omkit-0.0.2/omkit/tenant.py +271 -0
- omkit-0.0.2/omkit/tracing.py +80 -0
- omkit-0.0.2/omkit/transport/__init__.py +29 -0
- omkit-0.0.2/omkit/valkey.py +45 -0
- omkit-0.0.2/omkit.egg-info/PKG-INFO +29 -0
- omkit-0.0.2/omkit.egg-info/SOURCES.txt +72 -0
- omkit-0.0.2/omkit.egg-info/dependency_links.txt +1 -0
- omkit-0.0.2/omkit.egg-info/requires.txt +27 -0
- omkit-0.0.2/omkit.egg-info/top_level.txt +1 -0
- omkit-0.0.2/pyproject.toml +50 -0
- omkit-0.0.2/setup.cfg +4 -0
- omkit-0.0.2/tests/test_base.py +102 -0
- omkit-0.0.2/tests/test_cleanup.py +81 -0
- omkit-0.0.2/tests/test_cost.py +58 -0
- omkit-0.0.2/tests/test_data_facade.py +79 -0
- omkit-0.0.2/tests/test_dbpool.py +76 -0
- omkit-0.0.2/tests/test_dbpool_role.py +49 -0
- omkit-0.0.2/tests/test_encryption.py +102 -0
- omkit-0.0.2/tests/test_eventbus.py +161 -0
- omkit-0.0.2/tests/test_events_shim.py +40 -0
- omkit-0.0.2/tests/test_health.py +134 -0
- omkit-0.0.2/tests/test_http.py +71 -0
- omkit-0.0.2/tests/test_jobqueue_envelope.py +140 -0
- omkit-0.0.2/tests/test_jobqueue_streaq.py +287 -0
- omkit-0.0.2/tests/test_logging.py +170 -0
- omkit-0.0.2/tests/test_metrics.py +41 -0
- omkit-0.0.2/tests/test_quota.py +101 -0
- omkit-0.0.2/tests/test_registry_polling.py +102 -0
- omkit-0.0.2/tests/test_resilience.py +126 -0
- omkit-0.0.2/tests/test_sanitize.py +51 -0
- omkit-0.0.2/tests/test_security_events.py +171 -0
- omkit-0.0.2/tests/test_security_facade.py +59 -0
- omkit-0.0.2/tests/test_sessions.py +162 -0
- omkit-0.0.2/tests/test_settings_cache.py +95 -0
- omkit-0.0.2/tests/test_settings_manager.py +62 -0
- omkit-0.0.2/tests/test_settings_polling.py +69 -0
- omkit-0.0.2/tests/test_tenant.py +414 -0
- omkit-0.0.2/tests/test_tracing.py +93 -0
- omkit-0.0.2/tests/test_transport_facade.py +88 -0
- omkit-0.0.2/tests/test_valkey.py +51 -0
omkit-0.0.2/PKG-INFO
ADDED
|
@@ -0,0 +1,29 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: omkit
|
|
3
|
+
Version: 0.0.2
|
|
4
|
+
Summary: Multi-tenant SaaS scaffolding for Python services.
|
|
5
|
+
Requires-Python: >=3.12
|
|
6
|
+
Requires-Dist: pydantic>=2.13
|
|
7
|
+
Requires-Dist: pydantic-settings>=2.13
|
|
8
|
+
Requires-Dist: asyncpg>=0.31
|
|
9
|
+
Requires-Dist: redis>=7.4
|
|
10
|
+
Requires-Dist: structlog>=25.5
|
|
11
|
+
Requires-Dist: cryptography>=47.0
|
|
12
|
+
Requires-Dist: tenacity>=9.1
|
|
13
|
+
Requires-Dist: prometheus_client>=0.25
|
|
14
|
+
Requires-Dist: httpx>=0.28.1
|
|
15
|
+
Requires-Dist: sqlalchemy[asyncio]>=2.0.49
|
|
16
|
+
Provides-Extra: dev
|
|
17
|
+
Requires-Dist: fastapi>=0.136.0; extra == "dev"
|
|
18
|
+
Requires-Dist: pytest>=9.0; extra == "dev"
|
|
19
|
+
Requires-Dist: pytest-asyncio>=1.3; extra == "dev"
|
|
20
|
+
Requires-Dist: redis>=7.4; extra == "dev"
|
|
21
|
+
Requires-Dist: respx>=0.21; extra == "dev"
|
|
22
|
+
Provides-Extra: tracing
|
|
23
|
+
Requires-Dist: opentelemetry-api>=1.41.0; extra == "tracing"
|
|
24
|
+
Requires-Dist: opentelemetry-sdk>=1.41.0; extra == "tracing"
|
|
25
|
+
Requires-Dist: opentelemetry-exporter-otlp-proto-http>=1.41.0; extra == "tracing"
|
|
26
|
+
Requires-Dist: opentelemetry-instrumentation-fastapi>=0.62b0; extra == "tracing"
|
|
27
|
+
Requires-Dist: opentelemetry-instrumentation-httpx>=0.62b0; extra == "tracing"
|
|
28
|
+
Provides-Extra: metrics
|
|
29
|
+
Requires-Dist: prometheus-fastapi-instrumentator>=7.1; extra == "metrics"
|
omkit-0.0.2/README.md
ADDED
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
# omkit
|
|
2
|
+
|
|
3
|
+
Multi-tenant SaaS scaffolding for Python services. Pooled Postgres with RLS, Valkey eventbus, BYOK secrets, LLM provider abstraction, FastAPI observability middleware, and tenant-scoped session/job primitives.
|
|
4
|
+
|
|
5
|
+
```
|
|
6
|
+
pip install git+https://github.com/omurlabs/omkit-python@v0.0.1
|
|
7
|
+
```
|
|
8
|
+
|
|
9
|
+
```python
|
|
10
|
+
from omkit.dbpool import create_pool
|
|
11
|
+
from omkit.eventbus import EventBus
|
|
12
|
+
from omkit.settings import SettingsManager
|
|
13
|
+
```
|
|
14
|
+
|
|
15
|
+
Status: v0.0.1 — initial extraction. API is pre-stable; expect renames before v0.1.
|
|
16
|
+
|
|
17
|
+
License: Apache-2.0 (planned).
|
|
@@ -0,0 +1,18 @@
|
|
|
1
|
+
"""omkit — multi-tenant SaaS scaffolding for Python services.
|
|
2
|
+
|
|
3
|
+
Public surface re-exports a small set of commonly-used helpers. Internal
|
|
4
|
+
primitives are available via the submodules directly.
|
|
5
|
+
"""
|
|
6
|
+
from omkit.http import build_tenant_client
|
|
7
|
+
from omkit.metrics import mount_metrics
|
|
8
|
+
from omkit.settings import SettingsManager
|
|
9
|
+
from omkit import tenant
|
|
10
|
+
from omkit.tracing import instrument_fastapi
|
|
11
|
+
|
|
12
|
+
__all__ = [
|
|
13
|
+
"build_tenant_client",
|
|
14
|
+
"mount_metrics",
|
|
15
|
+
"SettingsManager",
|
|
16
|
+
"tenant",
|
|
17
|
+
"instrument_fastapi",
|
|
18
|
+
]
|
|
@@ -0,0 +1,62 @@
|
|
|
1
|
+
"""packages/omur-sdk/omkit/cleanup.py — Coordinated periodic cleanup task runner.
|
|
2
|
+
|
|
3
|
+
Loop.run() fires the provided ``task`` coroutine every ``interval`` seconds
|
|
4
|
+
while holding ``pg_try_advisory_lock(lock_key)`` — horizontally scaled
|
|
5
|
+
replicas of the same service won't double-execute the cleanup. Lock
|
|
6
|
+
contention turns the tick into a silent no-op.
|
|
7
|
+
|
|
8
|
+
exports: class Loop
|
|
9
|
+
rules: The Loop class must maintain a single persistent connection pool instance throughout its lifetime and cannot be instantiated without a valid pool argument. The run() method implements a blocking infinite loop that should only be called once per Loop instance and cannot be interrupted without external process termination. All database operations within the loop must use async context management to ensure proper connection handling and automatic release back to the pool.
|
|
10
|
+
agent: ollama/qwen3-coder:latest | ollama | 2026-05-01 | codedna-cli | initial CodeDNA annotation pass
|
|
11
|
+
message:
|
|
12
|
+
"""
|
|
13
|
+
|
|
14
|
+
from __future__ import annotations
|
|
15
|
+
|
|
16
|
+
import asyncio
|
|
17
|
+
import logging
|
|
18
|
+
from typing import Awaitable, Callable
|
|
19
|
+
|
|
20
|
+
log = logging.getLogger(__name__)
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
class Loop:
|
|
24
|
+
def __init__(
|
|
25
|
+
self,
|
|
26
|
+
pool,
|
|
27
|
+
*,
|
|
28
|
+
lock_key: int,
|
|
29
|
+
interval: float,
|
|
30
|
+
task: Callable[[], Awaitable[None]],
|
|
31
|
+
name: str = "cleanup-loop",
|
|
32
|
+
):
|
|
33
|
+
self._pool = pool
|
|
34
|
+
self._lock_key = lock_key
|
|
35
|
+
self._interval = interval
|
|
36
|
+
self._task = task
|
|
37
|
+
self._name = name
|
|
38
|
+
|
|
39
|
+
async def run(self) -> None:
|
|
40
|
+
"""
|
|
41
|
+
Rules: The function runs an infinite loop that must be properly cancelled or stopped to prevent resource leaks. The _tick() method and _interval attribute must be properly initialized before calling this function.
|
|
42
|
+
"""
|
|
43
|
+
while True:
|
|
44
|
+
try:
|
|
45
|
+
await self._tick()
|
|
46
|
+
except Exception as e:
|
|
47
|
+
log.warning("%s: tick failed: %s", self._name, e)
|
|
48
|
+
await asyncio.sleep(self._interval)
|
|
49
|
+
|
|
50
|
+
async def _tick(self) -> None:
|
|
51
|
+
async with self._pool.acquire() as conn:
|
|
52
|
+
got = await conn.fetchval(
|
|
53
|
+
"SELECT pg_try_advisory_lock($1)", self._lock_key
|
|
54
|
+
)
|
|
55
|
+
if not got:
|
|
56
|
+
return
|
|
57
|
+
try:
|
|
58
|
+
await self._task()
|
|
59
|
+
finally:
|
|
60
|
+
await conn.execute(
|
|
61
|
+
"SELECT pg_advisory_unlock($1)", self._lock_key
|
|
62
|
+
)
|
|
@@ -0,0 +1,60 @@
|
|
|
1
|
+
"""Shared base settings for services using omkit.
|
|
2
|
+
|
|
3
|
+
Subclass BaseServiceSettings in each service and add service-specific
|
|
4
|
+
fields. Pydantic-settings maps UPPERCASE env vars to these fields
|
|
5
|
+
automatically.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
from pydantic import Field
|
|
9
|
+
from pydantic_settings import BaseSettings, SettingsConfigDict
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
class BaseServiceSettings(BaseSettings):
|
|
13
|
+
"""Common env vars shared across backend services."""
|
|
14
|
+
|
|
15
|
+
model_config = SettingsConfigDict(env_file=".env", extra="ignore")
|
|
16
|
+
|
|
17
|
+
# Runtime
|
|
18
|
+
RUNTIME_MODE: str = "standalone"
|
|
19
|
+
TENANT_TOKEN: str = ""
|
|
20
|
+
settings_key: str = Field(default="", alias="SETTINGS_KEY")
|
|
21
|
+
|
|
22
|
+
# CORS
|
|
23
|
+
CORS_ORIGINS: str = ""
|
|
24
|
+
|
|
25
|
+
# PostgreSQL
|
|
26
|
+
POSTGRES_HOST: str = "postgres"
|
|
27
|
+
POSTGRES_PORT: int = 5432
|
|
28
|
+
POSTGRES_DB: str = ""
|
|
29
|
+
POSTGRES_USER: str = ""
|
|
30
|
+
POSTGRES_PASSWORD: str = ""
|
|
31
|
+
|
|
32
|
+
# Valkey (Redis-compatible, Apache 2.0)
|
|
33
|
+
VALKEY_HOST: str = "valkey"
|
|
34
|
+
VALKEY_PORT: int = 6379
|
|
35
|
+
VALKEY_PASSWORD: str = ""
|
|
36
|
+
|
|
37
|
+
# Ollama
|
|
38
|
+
OLLAMA_HOST: str = "http://ollama:11434"
|
|
39
|
+
OLLAMA_CHAT_MODEL: str = "qwen3:8b"
|
|
40
|
+
|
|
41
|
+
@property
|
|
42
|
+
def postgres_dsn(self) -> str:
|
|
43
|
+
return (
|
|
44
|
+
f"postgresql+asyncpg://{self.POSTGRES_USER}:{self.POSTGRES_PASSWORD}"
|
|
45
|
+
f"@{self.POSTGRES_HOST}:{self.POSTGRES_PORT}/{self.POSTGRES_DB}"
|
|
46
|
+
)
|
|
47
|
+
|
|
48
|
+
@property
|
|
49
|
+
def postgres_dsn_raw(self) -> str:
|
|
50
|
+
"""Plain postgresql:// DSN for drivers that don't accept asyncpg dialect."""
|
|
51
|
+
return (
|
|
52
|
+
f"postgresql://{self.POSTGRES_USER}:{self.POSTGRES_PASSWORD}"
|
|
53
|
+
f"@{self.POSTGRES_HOST}:{self.POSTGRES_PORT}/{self.POSTGRES_DB}"
|
|
54
|
+
)
|
|
55
|
+
|
|
56
|
+
@property
|
|
57
|
+
def valkey_url(self) -> str:
|
|
58
|
+
if self.VALKEY_PASSWORD:
|
|
59
|
+
return f"redis://:{self.VALKEY_PASSWORD}@{self.VALKEY_HOST}:{self.VALKEY_PORT}"
|
|
60
|
+
return f"redis://{self.VALKEY_HOST}:{self.VALKEY_PORT}"
|
|
@@ -0,0 +1,78 @@
|
|
|
1
|
+
"""omkit.cost — provider cost telemetry.
|
|
2
|
+
|
|
3
|
+
Records a Prometheus counter ``cost_units_total`` with a fixed
|
|
4
|
+
low-cardinality label set so VictoriaMetrics scrape stays cheap:
|
|
5
|
+
|
|
6
|
+
- ``service`` — emitting service name.
|
|
7
|
+
- ``provider`` — backend identifier (``local``, ``voyage``, ``openai``, …).
|
|
8
|
+
- ``op`` — operation name (``embed``, ``parse_pages``,
|
|
9
|
+
``rerank``, ``stt_seconds``, ``tts_chars``).
|
|
10
|
+
- ``tenant_bucket`` — coarse tenant grouping (``system`` / ``trial`` /
|
|
11
|
+
``paid``); never the raw tenant_id (cardinality).
|
|
12
|
+
|
|
13
|
+
The ``units`` argument is the billable count for the operation
|
|
14
|
+
(``tokens``, ``pages``, ``audio_seconds``, …). Dollar projection lives
|
|
15
|
+
out of scope — a static price-table joins counter values in Grafana.
|
|
16
|
+
"""
|
|
17
|
+
|
|
18
|
+
from __future__ import annotations
|
|
19
|
+
|
|
20
|
+
from typing import Final, Literal
|
|
21
|
+
|
|
22
|
+
from prometheus_client import Counter
|
|
23
|
+
|
|
24
|
+
TenantBucket = Literal["system", "trial", "paid"]
|
|
25
|
+
|
|
26
|
+
_VALID_BUCKETS: Final[frozenset[str]] = frozenset({"system", "trial", "paid"})
|
|
27
|
+
|
|
28
|
+
COST_UNITS_TOTAL: Final[Counter] = Counter(
|
|
29
|
+
"cost_units_total",
|
|
30
|
+
"Billable units emitted by a provider call (tokens, pages, seconds, chars).",
|
|
31
|
+
labelnames=("service", "provider", "op", "tenant_bucket"),
|
|
32
|
+
)
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
def record_cost(
|
|
36
|
+
*,
|
|
37
|
+
service: str,
|
|
38
|
+
provider: str,
|
|
39
|
+
op: str,
|
|
40
|
+
units: float,
|
|
41
|
+
tenant_bucket: TenantBucket | str,
|
|
42
|
+
) -> None:
|
|
43
|
+
"""Increment ``cost_units_total`` for one provider call.
|
|
44
|
+
|
|
45
|
+
Args:
|
|
46
|
+
service: Emitting service name.
|
|
47
|
+
provider: Backend identifier (``local``, ``voyage``, ``openai``, …).
|
|
48
|
+
op: Operation name (``embed``, ``parse_pages``, ``rerank``,
|
|
49
|
+
``stt_seconds``, ``tts_chars``).
|
|
50
|
+
units: Billable count for the operation (tokens, pages,
|
|
51
|
+
audio_seconds, character_count). Non-positive values are a
|
|
52
|
+
no-op — counters cannot decrease and a non-positive input is
|
|
53
|
+
always a caller bug.
|
|
54
|
+
tenant_bucket: Coarse tenant grouping. Unknown values are
|
|
55
|
+
normalised to ``trial`` so an instrumentation typo does not
|
|
56
|
+
silently inflate the ``paid`` bucket.
|
|
57
|
+
|
|
58
|
+
Returns:
|
|
59
|
+
None — emission is best-effort. Failure to record never raises;
|
|
60
|
+
the caller is on a hot path and a metric blip must not break it.
|
|
61
|
+
"""
|
|
62
|
+
|
|
63
|
+
if units <= 0:
|
|
64
|
+
return
|
|
65
|
+
bucket = tenant_bucket if tenant_bucket in _VALID_BUCKETS else "trial"
|
|
66
|
+
try:
|
|
67
|
+
COST_UNITS_TOTAL.labels(
|
|
68
|
+
service=service,
|
|
69
|
+
provider=provider,
|
|
70
|
+
op=op,
|
|
71
|
+
tenant_bucket=bucket,
|
|
72
|
+
).inc(units)
|
|
73
|
+
except Exception:
|
|
74
|
+
# Counter emission is best-effort — never propagate.
|
|
75
|
+
return
|
|
76
|
+
|
|
77
|
+
|
|
78
|
+
__all__ = ["COST_UNITS_TOTAL", "TenantBucket", "record_cost"]
|
|
@@ -0,0 +1,33 @@
|
|
|
1
|
+
"""packages/omur-sdk/omkit/data/__init__.py — re-exports DB pool and session-store primitives.
|
|
2
|
+
|
|
3
|
+
Additive grouping. Flat-module imports continue to work.
|
|
4
|
+
|
|
5
|
+
exports: none
|
|
6
|
+
rules: The module must maintain backward compatibility for all existing data import paths and cannot introduce breaking changes to the public API surface. All data processing functions must be thread-safe and handle concurrent access without race conditions. The module cannot depend on external packages beyond the standard library and explicitly declared dependencies in the package manifest.
|
|
7
|
+
agent: ollama/qwen3-coder:latest | ollama | 2026-05-01 | codedna-cli | initial CodeDNA annotation pass
|
|
8
|
+
message:
|
|
9
|
+
"""
|
|
10
|
+
|
|
11
|
+
from omkit.dbpool import (
|
|
12
|
+
build_retrieval_engine,
|
|
13
|
+
create_pool,
|
|
14
|
+
new_session_pool,
|
|
15
|
+
sqlalchemy_asyncpg_connect_args,
|
|
16
|
+
)
|
|
17
|
+
from omkit.sessions import (
|
|
18
|
+
NotFound,
|
|
19
|
+
Session,
|
|
20
|
+
SessionStore,
|
|
21
|
+
new_store,
|
|
22
|
+
)
|
|
23
|
+
|
|
24
|
+
__all__ = [
|
|
25
|
+
"build_retrieval_engine",
|
|
26
|
+
"create_pool",
|
|
27
|
+
"new_session_pool",
|
|
28
|
+
"sqlalchemy_asyncpg_connect_args",
|
|
29
|
+
"SessionStore",
|
|
30
|
+
"Session",
|
|
31
|
+
"NotFound",
|
|
32
|
+
"new_store",
|
|
33
|
+
]
|
|
@@ -0,0 +1,139 @@
|
|
|
1
|
+
"""packages/omur-sdk/omkit/dbpool.py — asyncpg pool helper that enforces a Postgres role via the init coroutine.
|
|
2
|
+
|
|
3
|
+
Runs ``SET ROLE <role>`` on every new physical connection so the role is
|
|
4
|
+
guaranteed even across reconnects — the async equivalent of pgx's
|
|
5
|
+
``AfterConnect`` hook. This is the defence-in-depth mechanism we rely on
|
|
6
|
+
after removing PgBouncer (which previously took care of the role reset via
|
|
7
|
+
``server_reset_query``).
|
|
8
|
+
|
|
9
|
+
exports: sqlalchemy_asyncpg_connect_args(role) | new_session_pool(dsn) | create_pool(dsn) | build_retrieval_engine(dsn)
|
|
10
|
+
rules: The module must maintain backward compatibility with existing SQLAlchemy and asyncpg integration patterns, and all database connection handling must adhere to async/await semantics throughout.
|
|
11
|
+
agent: ollama/qwen3-coder:latest | ollama | 2026-05-01 | codedna-cli | initial CodeDNA annotation pass
|
|
12
|
+
message:
|
|
13
|
+
"""
|
|
14
|
+
|
|
15
|
+
from __future__ import annotations
|
|
16
|
+
|
|
17
|
+
import re
|
|
18
|
+
from typing import Any, Optional
|
|
19
|
+
|
|
20
|
+
import asyncpg
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
def sqlalchemy_asyncpg_connect_args(role: str | None = "omur_app") -> dict[str, Any]:
|
|
24
|
+
"""Return ``connect_args`` for ``create_async_engine(..., connect_args=...)``.
|
|
25
|
+
|
|
26
|
+
Keeps the engine pgbouncer-transaction-mode-compatible
|
|
27
|
+
(``statement_cache_size=0``, ``prepared_statement_cache_size=0``) and,
|
|
28
|
+
when ``role`` is non-empty, applies ``SET ROLE <role>`` at connection
|
|
29
|
+
startup via asyncpg's ``server_settings`` dict so every pool
|
|
30
|
+
checkout already runs as the restricted role without any sync-event
|
|
31
|
+
listener. Pass ``role=None`` to opt out.
|
|
32
|
+
|
|
33
|
+
Running the role switch through ``server_settings`` (rather than a
|
|
34
|
+
sync ``"connect"`` event listener that calls ``dbapi_conn.cursor()``)
|
|
35
|
+
avoids a greenlet race between SQLAlchemy's async adapter and
|
|
36
|
+
asyncpg's connection init — the listener would occasionally see a
|
|
37
|
+
half-initialized DBAPI connection and raise ``'NoneType' object
|
|
38
|
+
has no attribute 'cursor'`` / ``'commit'``. asyncpg issues
|
|
39
|
+
``SET name = value`` for every entry in ``server_settings`` during
|
|
40
|
+
the connection handshake, before the connection is handed to the
|
|
41
|
+
pool, so the role is always in place by the time SQLAlchemy sees
|
|
42
|
+
the connection.
|
|
43
|
+
|
|
44
|
+
Rules: The function assumes that the `role` parameter is either a valid PostgreSQL role name or None. If a role is provided, it must exist in the database, and the calling user must have permission to switch to that role.
|
|
45
|
+
"""
|
|
46
|
+
args: dict[str, Any] = {
|
|
47
|
+
"statement_cache_size": 0,
|
|
48
|
+
"prepared_statement_cache_size": 0,
|
|
49
|
+
}
|
|
50
|
+
if role:
|
|
51
|
+
args["server_settings"] = {"role": role}
|
|
52
|
+
return args
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
def build_retrieval_engine(dsn: str):
|
|
56
|
+
"""Build the dedicated SQLAlchemy retrieval engine shared by frontal + marrow.
|
|
57
|
+
|
|
58
|
+
Tuning matches Option A of the hybrid-search spec:
|
|
59
|
+
pool_size=10, max_overflow=0, pool_timeout=2 — pool exhaustion fast-fails as
|
|
60
|
+
TimeoutError so the caller can map to HTTP 503 instead of unbounded queue
|
|
61
|
+
waits. Keeps retrieval traffic isolated from the ingestion pool.
|
|
62
|
+
|
|
63
|
+
Returns an async engine; caller binds it to its own AsyncSession class
|
|
64
|
+
(sqlalchemy vs sqlmodel session subclass).
|
|
65
|
+
"""
|
|
66
|
+
from sqlalchemy.ext.asyncio import create_async_engine
|
|
67
|
+
|
|
68
|
+
return create_async_engine(
|
|
69
|
+
dsn,
|
|
70
|
+
echo=False,
|
|
71
|
+
pool_size=10,
|
|
72
|
+
max_overflow=0,
|
|
73
|
+
pool_timeout=2,
|
|
74
|
+
connect_args=sqlalchemy_asyncpg_connect_args(),
|
|
75
|
+
)
|
|
76
|
+
|
|
77
|
+
|
|
78
|
+
async def new_session_pool(dsn: str, **kwargs) -> asyncpg.Pool:
|
|
79
|
+
"""Build an asyncpg pool for :class:`omkit.sessions.PostgresSessionStore`.
|
|
80
|
+
|
|
81
|
+
The session pool intentionally does **not** run ``SET ROLE omur_app``
|
|
82
|
+
on each new connection. Token-based session lookup (``get``/``delete``)
|
|
83
|
+
takes an opaque token without knowing the tenant, so SELECT/DELETE
|
|
84
|
+
under a role that's subject to the ``sessions_tenant_isolation`` RLS
|
|
85
|
+
policy would silently return zero rows. Connecting as the default
|
|
86
|
+
``omur`` superuser (which has ``BYPASSRLS``) lets token lookup cross
|
|
87
|
+
tenants; writes (``put``, ``list``) still run inside a transaction
|
|
88
|
+
that sets ``app.tenant_id`` so RLS is honored for multi-tenant
|
|
89
|
+
mutations.
|
|
90
|
+
|
|
91
|
+
Use this helper in service lifespans that need a SessionStore. For
|
|
92
|
+
pools that back RLS-enforced app queries, use :func:`create_pool`
|
|
93
|
+
with ``role='omur_app'`` instead.
|
|
94
|
+
|
|
95
|
+
Rules: The function intentionally avoids setting a role on connections to allow cross-tenant session lookups. Future developers must understand that this design relies on the `omur` superuser having `BYPASSRLS` privileges to ensure proper behavior.
|
|
96
|
+
"""
|
|
97
|
+
return await create_pool(dsn, **kwargs)
|
|
98
|
+
|
|
99
|
+
|
|
100
|
+
def _normalize_dsn(dsn: str) -> str:
|
|
101
|
+
"""Strip a SQLAlchemy dialect suffix like ``+asyncpg`` from the URL
|
|
102
|
+
scheme so asyncpg accepts a DSN that was originally shaped for
|
|
103
|
+
``create_async_engine``."""
|
|
104
|
+
return re.sub(r"^(postgres(?:ql)?)\+[a-z0-9_]+://", r"\1://", dsn, count=1)
|
|
105
|
+
|
|
106
|
+
|
|
107
|
+
async def create_pool(
|
|
108
|
+
dsn: str,
|
|
109
|
+
*,
|
|
110
|
+
role: Optional[str] = None,
|
|
111
|
+
min_size: int = 1,
|
|
112
|
+
max_size: int = 10,
|
|
113
|
+
**kwargs,
|
|
114
|
+
) -> asyncpg.Pool:
|
|
115
|
+
"""Create an asyncpg pool. If ``role`` is given, every new physical
|
|
116
|
+
connection runs ``SET ROLE "<role>"`` via the init coroutine.
|
|
117
|
+
|
|
118
|
+
``statement_cache_size`` defaults to ``0`` so that pgbouncer in
|
|
119
|
+
transaction mode can be reintroduced as a pure config change
|
|
120
|
+
(prepared statements don't survive across pgbouncer's per-transaction
|
|
121
|
+
server rotation). Callers can override by passing ``statement_cache_size``
|
|
122
|
+
|
|
123
|
+
Rules: If `role` is provided, it must be a valid PostgreSQL role name, and the calling user must have the necessary permissions to execute `SET ROLE <role>`. Additionally, `statement_cache_size` is defaulted to 0 for pgbouncer compatibility, but callers can override this if needed.
|
|
124
|
+
explicitly."""
|
|
125
|
+
|
|
126
|
+
async def _init(conn: asyncpg.Connection) -> None:
|
|
127
|
+
if role:
|
|
128
|
+
await conn.execute(f'SET ROLE "{role}"')
|
|
129
|
+
|
|
130
|
+
if role:
|
|
131
|
+
kwargs.setdefault("init", _init)
|
|
132
|
+
kwargs.setdefault("statement_cache_size", 0)
|
|
133
|
+
|
|
134
|
+
return await asyncpg.create_pool(
|
|
135
|
+
_normalize_dsn(dsn),
|
|
136
|
+
min_size=min_size,
|
|
137
|
+
max_size=max_size,
|
|
138
|
+
**kwargs,
|
|
139
|
+
)
|
|
@@ -0,0 +1,58 @@
|
|
|
1
|
+
"""packages/omur-sdk/omkit/encryption.py — Fernet-based encryption utilities for Omur settings secrets.
|
|
2
|
+
|
|
3
|
+
exports: generate_key() | encrypt_value(plaintext, key) | decrypt_value(ciphertext, key) | mask_secret(value)
|
|
4
|
+
rules: The encryption module must maintain backward compatibility with all existing encrypted data formats and key structures. All cryptographic operations must be deterministic and reproducible across different runtime environments. The module cannot introduce any external dependencies beyond the standard library and the fernet package.
|
|
5
|
+
agent: ollama/qwen3-coder:latest | ollama | 2026-05-01 | codedna-cli | initial CodeDNA annotation pass
|
|
6
|
+
message:
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
from cryptography.fernet import Fernet, InvalidToken
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
def generate_key() -> str:
|
|
13
|
+
"""Generate a new Fernet-compatible key (URL-safe base64).
|
|
14
|
+
|
|
15
|
+
Rules: Key must be stored securely and never logged or exposed in plaintext. The generated key is URL-safe base64 encoded and should be persisted in a secure key management system.
|
|
16
|
+
"""
|
|
17
|
+
return Fernet.generate_key().decode()
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
def encrypt_value(plaintext: str, key: str) -> str:
|
|
21
|
+
"""Encrypt a string value. Returns base64-encoded ciphertext.
|
|
22
|
+
|
|
23
|
+
Rules: The key must be a valid Fernet-compatible key (URL-safe base64 encoded string) or the function will raise a ValueError.
|
|
24
|
+
"""
|
|
25
|
+
f = Fernet(key.encode())
|
|
26
|
+
return f.encrypt(plaintext.encode()).decode()
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
def decrypt_value(ciphertext: str, key: str) -> str:
|
|
30
|
+
"""Decrypt a base64-encoded ciphertext string.
|
|
31
|
+
|
|
32
|
+
Raises cryptography.fernet.InvalidToken if the key is wrong or the
|
|
33
|
+
token is malformed/tampered.
|
|
34
|
+
|
|
35
|
+
Rules: The key must match the one used for encryption, and the ciphertext must be a valid Fernet token; otherwise, cryptography.fernet.InvalidToken will be raised.
|
|
36
|
+
"""
|
|
37
|
+
f = Fernet(key.encode())
|
|
38
|
+
return f.decrypt(ciphertext.encode()).decode()
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
def mask_secret(value: str | None) -> str | None:
|
|
42
|
+
"""Mask a secret for API display.
|
|
43
|
+
|
|
44
|
+
- Values >= 10 chars: first 4 + '****' + last 4 (e.g. 'sk-a****Xk2f')
|
|
45
|
+
- Values 4-9 chars: first 2 + '****' + last 2 (e.g. 'ab****ef')
|
|
46
|
+
- Values < 4 chars: '****'
|
|
47
|
+
- None or empty: returns None
|
|
48
|
+
|
|
49
|
+
Rules: The function assumes ASCII-compatible strings; non-ASCII characters may produce unexpected masking behavior due to byte-level string slicing.
|
|
50
|
+
"""
|
|
51
|
+
if not value:
|
|
52
|
+
return None
|
|
53
|
+
n = len(value)
|
|
54
|
+
if n >= 10:
|
|
55
|
+
return value[:4] + "****" + value[-4:]
|
|
56
|
+
if n >= 4:
|
|
57
|
+
return value[:2] + "****" + value[-2:]
|
|
58
|
+
return "****"
|