fix(backends,storage,core): setuid in exec wrapper, missing-column migration, runner config dir for mirror, tagged prompt sources
This commit is contained in:
1 parent
e1f242a87a
commit
50b7057fa4
6 files changed
+127
-23
No files matched your search
@@ -18,6 +18,7 @@ import json
|
||||
from datetime import UTC, datetime
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from sqlalchemy import inspect, text
|
||||
from sqlalchemy.ext.asyncio import create_async_engine
|
||||
from sqlmodel import SQLModel, select
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
@@ -27,6 +28,7 @@ from beaver_gateway.storage.models import AuditLog, Token, Usage
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
from sqlalchemy import Connection
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine
|
||||
|
||||
|
||||
@@ -72,9 +74,14 @@ class Database:
|
||||
self._engine: AsyncEngine = create_async_engine(_to_async_url(url), echo=False)
|
||||
|
||||
async def create_all(self) -> None:
|
||||
"""Idempotent ``CREATE TABLE IF NOT EXISTS`` for every model."""
|
||||
"""Create missing tables, then add columns models gained since.
|
||||
|
||||
Nullable columns only - that is the whole migration story until
|
||||
Alembic is worth it.
|
||||
"""
|
||||
async with self._engine.begin() as conn:
|
||||
await conn.run_sync(SQLModel.metadata.create_all)
|
||||
await conn.run_sync(_add_missing_columns)
|
||||
|
||||
def session(self) -> AsyncSession:
|
||||
"""Open a fresh :class:`AsyncSession` (use as ``async with``).
|
||||
@@ -90,6 +97,19 @@ class Database:
|
||||
await self._engine.dispose()
|
||||
|
||||
|
||||
def _add_missing_columns(conn: Connection) -> None:
|
||||
inspector = inspect(conn)
|
||||
for table in SQLModel.metadata.sorted_tables:
|
||||
existing = {c["name"] for c in inspector.get_columns(table.name)}
|
||||
for column in table.columns:
|
||||
if column.name in existing:
|
||||
continue
|
||||
kind = column.type.compile(conn.dialect)
|
||||
conn.execute(
|
||||
text(f"ALTER TABLE {table.name} ADD COLUMN {column.name} {kind}")
|
||||
)
|
||||
|
||||
|
||||
# ---- Token CRUD ---------------------------------------------------------
|
||||
|
||||
|
||||
|
||||
Reference in new issue
Block a user