File size: 4,626 Bytes
48d895c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | """SQL config backend (MySQL / PostgreSQL)."""
from typing import Any
import sqlalchemy as sa
from sqlalchemy.ext.asyncio import AsyncEngine
from .base import ConfigBackend
from ._serde import flatten, unflatten
_TABLE = "config_store"
_VERSION_KEY = "__version__"
_metadata = sa.MetaData()
config_store_table = sa.Table(
_TABLE,
_metadata,
sa.Column("key", sa.String(255), primary_key=True),
sa.Column("value", sa.Text, nullable=False, default=""),
)
class SqlConfigBackend(ConfigBackend):
"""Flat key-value storage in a ``config_store`` table.
Each dotted config key is one row. ``apply_patch`` only UPSERTs the
changed rows — the rest of the table is untouched.
Version token is stored as the integer value of the ``__version__`` row.
"""
def __init__(
self,
engine: AsyncEngine,
*,
dialect: str = "postgresql",
dispose_engine: bool = True,
) -> None:
self._engine = engine
self._dialect = dialect # "mysql" | "postgresql"
self._ready = False
self._dispose_engine = dispose_engine
async def _ensure_table(self) -> None:
if self._ready:
return
async with self._engine.begin() as conn:
await conn.run_sync(_metadata.create_all)
self._ready = True
async def load(self) -> dict[str, Any]:
await self._ensure_table()
async with self._engine.connect() as conn:
rows = await conn.execute(
sa.select(config_store_table)
.where(config_store_table.c.key != _VERSION_KEY)
)
flat = {row.key: row.value for row in rows}
return unflatten(flat)
async def apply_patch(self, patch: dict[str, Any]) -> None:
await self._ensure_table()
flat = flatten(patch)
if not flat:
return
async with self._engine.begin() as conn:
for k, v in flat.items():
await conn.execute(self._upsert(k, v))
# Atomically increment version counter (no read-modify-write race).
await conn.execute(self._upsert_incr_version())
async def version(self) -> object:
await self._ensure_table()
async with self._engine.connect() as conn:
row = await conn.execute(
sa.select(config_store_table.c.value)
.where(config_store_table.c.key == _VERSION_KEY)
)
val = row.scalar()
return int(val) if val else 0
def _upsert(self, key: str, value: str) -> sa.Insert:
if self._dialect == "postgresql":
from sqlalchemy.dialects.postgresql import insert
return (
insert(config_store_table)
.values(key=key, value=value)
.on_conflict_do_update(index_elements=["key"], set_={"value": value})
)
else: # mysql
from sqlalchemy.dialects.mysql import insert
return (
insert(config_store_table)
.values(key=key, value=value)
.on_duplicate_key_update(value=value)
)
def _upsert_incr_version(self) -> sa.Insert:
"""Atomically insert or increment the version counter row."""
if self._dialect == "postgresql":
from sqlalchemy.dialects.postgresql import insert
return (
insert(config_store_table)
.values(key=_VERSION_KEY, value="1")
.on_conflict_do_update(
index_elements=["key"],
set_={"value": sa.func.cast(
sa.func.cast(config_store_table.c.value, sa.Integer) + 1,
sa.Text,
)},
)
)
else: # mysql
from sqlalchemy.dialects.mysql import insert
return (
insert(config_store_table)
.values(key=_VERSION_KEY, value="1")
.on_duplicate_key_update(
value=sa.func.cast(
sa.func.cast(config_store_table.c.value, sa.Integer) + 1,
sa.Text,
)
)
)
async def close(self) -> None:
if self._dispose_engine:
from app.control.account.backends.sql import _evict_cached_engine
_evict_cached_engine(self._engine)
await self._engine.dispose()
|