feat(bot): модель данных, миграции Alembic, FSM в БД, TTL-чистка
- app/db/models.py: 13 таблиц (ARCHITECTURE §2.10); перечисления — VARCHAR + CHECK, JSON → JSONB на Postgres, BIGINT/INTEGER-ключи для SQLite - миграция 0001 (автогенерация + ручная чистка: импорт sa.Text, дубли CHECK) - миграции применяются при старте бота программно (app.db.migrate), script_location из пакета - SqlStorage: FSM aiogram поверх fsm_state — состояние переживает рестарт (перенесено из #4) - cleanup_loop: истёкшие черновики, снимки контейнеров и попытки входа старше суток - тесты: upgrade/downgrade, дрейф моделей и миграций, CHECK, CASCADE, FSM, чистка Refs #5 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,46 @@
|
||||
"""Periodic cleanup of TTL data: expired drafts, old container snapshots and access attempts."""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Any, cast
|
||||
|
||||
from sqlalchemy import CursorResult, Delete, delete
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.db.models import AccessAttempt, ContainerSnapshot, Draft, utcnow
|
||||
from app.db.session import SessionMaker
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
SNAPSHOT_TTL = timedelta(days=1)
|
||||
ATTEMPT_TTL = timedelta(days=1)
|
||||
|
||||
|
||||
async def _delete(session: AsyncSession, stmt: Delete) -> int:
|
||||
# DML executes into a CursorResult; the generic annotation of execute() just doesn't say so
|
||||
return cast(CursorResult[Any], await session.execute(stmt)).rowcount
|
||||
|
||||
|
||||
async def cleanup_expired(sessions: SessionMaker, now: datetime) -> dict[str, int]:
|
||||
async with sessions.begin() as session:
|
||||
return {
|
||||
"drafts": await _delete(session, delete(Draft).where(Draft.expires_at < now)),
|
||||
"snapshots": await _delete(
|
||||
session, delete(ContainerSnapshot).where(ContainerSnapshot.at < now - SNAPSHOT_TTL)
|
||||
),
|
||||
"access_attempts": await _delete(
|
||||
session, delete(AccessAttempt).where(AccessAttempt.at < now - ATTEMPT_TTL)
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
async def cleanup_loop(sessions: SessionMaker, interval: float = 3600.0) -> None:
|
||||
while True:
|
||||
try:
|
||||
removed = await cleanup_expired(sessions, utcnow())
|
||||
if any(removed.values()):
|
||||
log.info("cleanup: %s", removed)
|
||||
except Exception: # a failed cleanup must not kill the bot; retry next round
|
||||
log.exception("cleanup failed")
|
||||
await asyncio.sleep(interval)
|
||||
@@ -0,0 +1,52 @@
|
||||
"""Apply Alembic migrations programmatically at startup — no separate step in the container.
|
||||
|
||||
`python -m app.db.migrate [revision]` upgrades the DB from DATABASE_URL (default: head).
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import sys
|
||||
|
||||
from alembic import command
|
||||
from alembic.config import Config
|
||||
from sqlalchemy.engine import Connection
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine
|
||||
|
||||
from app.config import DatabaseSettings
|
||||
from app.db.session import make_engine
|
||||
|
||||
SCRIPT_LOCATION = "app.db:migrations"
|
||||
|
||||
|
||||
def alembic_config(connection: Connection | None = None) -> Config:
|
||||
cfg = Config()
|
||||
cfg.set_main_option("script_location", SCRIPT_LOCATION)
|
||||
# env.py reuses this connection instead of opening its own (async engine, one transaction)
|
||||
cfg.attributes["connection"] = connection
|
||||
return cfg
|
||||
|
||||
|
||||
async def upgrade(engine: AsyncEngine, revision: str = "head") -> None:
|
||||
async with engine.begin() as conn:
|
||||
await conn.run_sync(lambda sync_conn: command.upgrade(alembic_config(sync_conn), revision))
|
||||
|
||||
|
||||
async def downgrade(engine: AsyncEngine, revision: str) -> None:
|
||||
async with engine.begin() as conn:
|
||||
await conn.run_sync(lambda sync_conn: command.downgrade(alembic_config(sync_conn), revision))
|
||||
|
||||
|
||||
def main() -> None:
|
||||
revision = sys.argv[1] if len(sys.argv) > 1 else "head"
|
||||
|
||||
async def _run() -> None:
|
||||
engine = make_engine(DatabaseSettings().url)
|
||||
try:
|
||||
await upgrade(engine, revision)
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,39 @@
|
||||
"""Alembic environment. Runs on a connection handed over by app.db.migrate (or the CLI via alembic.ini)."""
|
||||
|
||||
import asyncio
|
||||
|
||||
from alembic import context
|
||||
from sqlalchemy.engine import Connection
|
||||
|
||||
from app.config import DatabaseSettings
|
||||
from app.db.models import Base
|
||||
from app.db.session import make_engine
|
||||
|
||||
target_metadata = Base.metadata
|
||||
|
||||
|
||||
def _run(connection: Connection) -> None:
|
||||
context.configure(
|
||||
connection=connection,
|
||||
target_metadata=target_metadata,
|
||||
render_as_batch=connection.dialect.name == "sqlite", # SQLite can't ALTER most things in place
|
||||
compare_type=True,
|
||||
)
|
||||
with context.begin_transaction():
|
||||
context.run_migrations()
|
||||
|
||||
|
||||
async def _run_cli() -> None:
|
||||
engine = make_engine(DatabaseSettings().url)
|
||||
try:
|
||||
async with engine.begin() as conn:
|
||||
await conn.run_sync(_run)
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
connection = context.config.attributes.get("connection")
|
||||
if connection is not None:
|
||||
_run(connection)
|
||||
else:
|
||||
asyncio.run(_run_cli())
|
||||
@@ -0,0 +1,25 @@
|
||||
"""${message}
|
||||
|
||||
Revision ID: ${up_revision}
|
||||
Revises: ${down_revision | comma,n}
|
||||
Create Date: ${create_date}
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
${imports if imports else ""}
|
||||
|
||||
revision: str = ${repr(up_revision)}
|
||||
down_revision: str | None = ${repr(down_revision)}
|
||||
branch_labels: str | Sequence[str] | None = ${repr(branch_labels)}
|
||||
depends_on: str | Sequence[str] | None = ${repr(depends_on)}
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
${upgrades if upgrades else "pass"}
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
${downgrades if downgrades else "pass"}
|
||||
@@ -0,0 +1,311 @@
|
||||
"""initial schema
|
||||
|
||||
Revision ID: 0001
|
||||
Revises:
|
||||
Create Date: 2026-10-14 17:41:05.672912
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
from sqlalchemy.dialects import postgresql
|
||||
|
||||
revision: str = "0001"
|
||||
down_revision: str | None = None
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.create_table(
|
||||
"access_attempts",
|
||||
sa.Column("id", sa.BigInteger().with_variant(sa.Integer(), "sqlite"), nullable=False),
|
||||
sa.Column("tg_id", sa.BigInteger(), nullable=False),
|
||||
sa.Column("at", sa.DateTime(timezone=True), nullable=False),
|
||||
sa.Column("ok", sa.Boolean(), nullable=False),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_access_attempts")),
|
||||
)
|
||||
with op.batch_alter_table("access_attempts", schema=None) as batch_op:
|
||||
batch_op.create_index("ix_access_attempts_tg_id_at", ["tg_id", "at"], unique=False)
|
||||
|
||||
op.create_table(
|
||||
"alerts",
|
||||
sa.Column("id", sa.BigInteger().with_variant(sa.Integer(), "sqlite"), nullable=False),
|
||||
sa.Column(
|
||||
"kind",
|
||||
sa.Enum(
|
||||
"crash",
|
||||
"oom",
|
||||
"restart_loop",
|
||||
"unhealthy",
|
||||
"overload",
|
||||
"deploy_refused",
|
||||
"cd_failed",
|
||||
"disk_low",
|
||||
"dependency_down",
|
||||
"agent_attack",
|
||||
name="alert_kind",
|
||||
native_enum=False,
|
||||
create_constraint=True,
|
||||
length=20,
|
||||
),
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column("service", sa.String(length=100), nullable=True),
|
||||
sa.Column(
|
||||
"payload", sa.JSON().with_variant(postgresql.JSONB(astext_type=sa.Text()), "postgresql"), nullable=False
|
||||
),
|
||||
sa.Column("count", sa.Integer(), nullable=False),
|
||||
sa.Column("opened_at", sa.DateTime(timezone=True), nullable=False),
|
||||
sa.Column("last_at", sa.DateTime(timezone=True), nullable=False),
|
||||
sa.Column("resolved_at", sa.DateTime(timezone=True), nullable=True),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_alerts")),
|
||||
)
|
||||
with op.batch_alter_table("alerts", schema=None) as batch_op:
|
||||
batch_op.create_index("ix_alerts_kind_service_resolved_at", ["kind", "service", "resolved_at"], unique=False)
|
||||
|
||||
op.create_table(
|
||||
"container_snapshots",
|
||||
sa.Column("id", sa.BigInteger().with_variant(sa.Integer(), "sqlite"), nullable=False),
|
||||
sa.Column("container", sa.String(length=100), nullable=False),
|
||||
sa.Column("at", sa.DateTime(timezone=True), nullable=False),
|
||||
sa.Column("restart_count", sa.Integer(), nullable=False),
|
||||
sa.Column("cpu_pct", sa.Float(), nullable=False),
|
||||
sa.Column("mem_pct", sa.Float(), nullable=False),
|
||||
sa.Column("health", sa.String(length=20), nullable=True),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_container_snapshots")),
|
||||
)
|
||||
with op.batch_alter_table("container_snapshots", schema=None) as batch_op:
|
||||
batch_op.create_index("ix_container_snapshots_container_at", ["container", "at"], unique=False)
|
||||
|
||||
op.create_table(
|
||||
"fsm_state",
|
||||
sa.Column("key", sa.String(length=255), nullable=False),
|
||||
sa.Column("state", sa.String(length=255), nullable=True),
|
||||
sa.Column(
|
||||
"data", sa.JSON().with_variant(postgresql.JSONB(astext_type=sa.Text()), "postgresql"), nullable=False
|
||||
),
|
||||
sa.PrimaryKeyConstraint("key", name=op.f("pk_fsm_state")),
|
||||
)
|
||||
op.create_table(
|
||||
"users",
|
||||
sa.Column("tg_id", sa.BigInteger(), autoincrement=False, nullable=False),
|
||||
sa.Column("username", sa.String(length=64), nullable=True),
|
||||
sa.Column("full_name", sa.String(length=256), nullable=False),
|
||||
sa.Column(
|
||||
"role",
|
||||
sa.Enum("member", "admin", name="role", native_enum=False, create_constraint=True, length=20),
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column("is_active", sa.Boolean(), nullable=False),
|
||||
sa.Column("invited_by", sa.BigInteger(), nullable=True),
|
||||
sa.Column("joined_at", sa.DateTime(timezone=True), nullable=False),
|
||||
sa.Column("last_seen_at", sa.DateTime(timezone=True), nullable=True),
|
||||
sa.ForeignKeyConstraint(["invited_by"], ["users.tg_id"], name=op.f("fk_users_invited_by_users")),
|
||||
sa.PrimaryKeyConstraint("tg_id", name=op.f("pk_users")),
|
||||
)
|
||||
op.create_table(
|
||||
"agent_requests",
|
||||
sa.Column("id", sa.BigInteger().with_variant(sa.Integer(), "sqlite"), nullable=False),
|
||||
sa.Column("user_id", sa.BigInteger(), nullable=False),
|
||||
sa.Column("prompt", sa.Text(), nullable=False),
|
||||
sa.Column(
|
||||
"verdict",
|
||||
sa.Enum("safe", "attack", "error", name="verdict", native_enum=False, create_constraint=True, length=20),
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column("reason", sa.Text(), nullable=True),
|
||||
sa.Column(
|
||||
"log_query", sa.JSON().with_variant(postgresql.JSONB(astext_type=sa.Text()), "postgresql"), nullable=True
|
||||
),
|
||||
sa.Column("tokens", sa.Integer(), nullable=False),
|
||||
sa.Column("at", sa.DateTime(timezone=True), nullable=False),
|
||||
sa.ForeignKeyConstraint(["user_id"], ["users.tg_id"], name=op.f("fk_agent_requests_user_id_users")),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_agent_requests")),
|
||||
)
|
||||
op.create_table(
|
||||
"alert_deliveries",
|
||||
sa.Column("id", sa.BigInteger().with_variant(sa.Integer(), "sqlite"), nullable=False),
|
||||
sa.Column("alert_id", sa.BigInteger().with_variant(sa.Integer(), "sqlite"), nullable=False),
|
||||
sa.Column("chat_id", sa.BigInteger(), nullable=False),
|
||||
sa.Column("message_id", sa.BigInteger(), nullable=False),
|
||||
sa.ForeignKeyConstraint(
|
||||
["alert_id"], ["alerts.id"], name=op.f("fk_alert_deliveries_alert_id_alerts"), ondelete="CASCADE"
|
||||
),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_alert_deliveries")),
|
||||
sa.UniqueConstraint("alert_id", "chat_id", name=op.f("uq_alert_deliveries_alert_id")),
|
||||
)
|
||||
op.create_table(
|
||||
"drafts",
|
||||
sa.Column("id", sa.BigInteger().with_variant(sa.Integer(), "sqlite"), nullable=False),
|
||||
sa.Column("user_id", sa.BigInteger(), nullable=False),
|
||||
sa.Column("base_sha", sa.String(length=40), nullable=False),
|
||||
sa.Column("message", sa.Text(), nullable=True),
|
||||
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
|
||||
sa.Column("expires_at", sa.DateTime(timezone=True), nullable=False),
|
||||
sa.ForeignKeyConstraint(["user_id"], ["users.tg_id"], name=op.f("fk_drafts_user_id_users"), ondelete="CASCADE"),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_drafts")),
|
||||
sa.UniqueConstraint("user_id", name=op.f("uq_drafts_user_id")),
|
||||
)
|
||||
op.create_table(
|
||||
"invites",
|
||||
sa.Column("id", sa.BigInteger().with_variant(sa.Integer(), "sqlite"), nullable=False),
|
||||
sa.Column("token_hash", sa.String(length=64), nullable=False),
|
||||
sa.Column("created_by", sa.BigInteger(), nullable=False),
|
||||
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
|
||||
sa.Column("expires_at", sa.DateTime(timezone=True), nullable=False),
|
||||
sa.Column("used_by", sa.BigInteger(), nullable=True),
|
||||
sa.Column("used_at", sa.DateTime(timezone=True), nullable=True),
|
||||
sa.ForeignKeyConstraint(["created_by"], ["users.tg_id"], name=op.f("fk_invites_created_by_users")),
|
||||
sa.ForeignKeyConstraint(["used_by"], ["users.tg_id"], name=op.f("fk_invites_used_by_users")),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_invites")),
|
||||
sa.UniqueConstraint("token_hash", name=op.f("uq_invites_token_hash")),
|
||||
)
|
||||
op.create_table(
|
||||
"log_queries",
|
||||
sa.Column("id", sa.BigInteger().with_variant(sa.Integer(), "sqlite"), nullable=False),
|
||||
sa.Column("user_id", sa.BigInteger(), nullable=False),
|
||||
sa.Column("container", sa.String(length=100), nullable=False),
|
||||
sa.Column(
|
||||
"mode",
|
||||
sa.Enum("tail", "grep", "agent", name="log_mode", native_enum=False, create_constraint=True, length=20),
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column(
|
||||
"query", sa.JSON().with_variant(postgresql.JSONB(astext_type=sa.Text()), "postgresql"), nullable=False
|
||||
),
|
||||
sa.Column("lines", sa.Integer(), nullable=False),
|
||||
sa.Column("at", sa.DateTime(timezone=True), nullable=False),
|
||||
sa.ForeignKeyConstraint(["user_id"], ["users.tg_id"], name=op.f("fk_log_queries_user_id_users")),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_log_queries")),
|
||||
)
|
||||
op.create_table(
|
||||
"pushes",
|
||||
sa.Column("id", sa.BigInteger().with_variant(sa.Integer(), "sqlite"), nullable=False),
|
||||
sa.Column("user_id", sa.BigInteger(), nullable=False),
|
||||
sa.Column("sha", sa.String(length=40), nullable=False),
|
||||
sa.Column("message", sa.Text(), nullable=False),
|
||||
sa.Column(
|
||||
"paths", sa.JSON().with_variant(postgresql.JSONB(astext_type=sa.Text()), "postgresql"), nullable=False
|
||||
),
|
||||
sa.Column("chat_id", sa.BigInteger(), nullable=False),
|
||||
sa.Column("message_id", sa.BigInteger(), nullable=True),
|
||||
sa.Column(
|
||||
"ci_status",
|
||||
sa.Enum(
|
||||
"pending",
|
||||
"running",
|
||||
"success",
|
||||
"failed",
|
||||
"timeout",
|
||||
"skipped",
|
||||
name="ci_status",
|
||||
native_enum=False,
|
||||
create_constraint=True,
|
||||
length=20,
|
||||
),
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column(
|
||||
"cd_status",
|
||||
sa.Enum(
|
||||
"pending",
|
||||
"running",
|
||||
"success",
|
||||
"failed",
|
||||
"timeout",
|
||||
"skipped",
|
||||
name="cd_status",
|
||||
native_enum=False,
|
||||
create_constraint=True,
|
||||
length=20,
|
||||
),
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
|
||||
sa.ForeignKeyConstraint(["user_id"], ["users.tg_id"], name=op.f("fk_pushes_user_id_users")),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_pushes")),
|
||||
sa.UniqueConstraint("sha", name=op.f("uq_pushes_sha")),
|
||||
)
|
||||
op.create_table(
|
||||
"draft_changes",
|
||||
sa.Column("id", sa.BigInteger().with_variant(sa.Integer(), "sqlite"), nullable=False),
|
||||
sa.Column("draft_id", sa.BigInteger().with_variant(sa.Integer(), "sqlite"), nullable=False),
|
||||
sa.Column("path", sa.String(length=500), nullable=False),
|
||||
sa.Column(
|
||||
"op",
|
||||
sa.Enum("add", "modify", "delete", name="change_op", native_enum=False, create_constraint=True, length=20),
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column("content", sa.LargeBinary(), nullable=True),
|
||||
sa.ForeignKeyConstraint(
|
||||
["draft_id"], ["drafts.id"], name=op.f("fk_draft_changes_draft_id_drafts"), ondelete="CASCADE"
|
||||
),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_draft_changes")),
|
||||
sa.UniqueConstraint("draft_id", "path", name=op.f("uq_draft_changes_draft_id")),
|
||||
)
|
||||
op.create_table(
|
||||
"pipeline_runs",
|
||||
sa.Column("id", sa.BigInteger().with_variant(sa.Integer(), "sqlite"), nullable=False),
|
||||
sa.Column("push_id", sa.BigInteger().with_variant(sa.Integer(), "sqlite"), nullable=False),
|
||||
sa.Column(
|
||||
"stage",
|
||||
sa.Enum("ci", "cd", name="stage", native_enum=False, create_constraint=True, length=20),
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column("run_slug", sa.String(length=50), nullable=False),
|
||||
sa.Column(
|
||||
"status",
|
||||
sa.Enum(
|
||||
"pending",
|
||||
"running",
|
||||
"success",
|
||||
"failed",
|
||||
"timeout",
|
||||
"skipped",
|
||||
name="run_status",
|
||||
native_enum=False,
|
||||
create_constraint=True,
|
||||
length=20,
|
||||
),
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column("failed_cube", sa.String(length=100), nullable=True),
|
||||
sa.Column("url", sa.String(length=500), nullable=True),
|
||||
sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False),
|
||||
sa.ForeignKeyConstraint(
|
||||
["push_id"], ["pushes.id"], name=op.f("fk_pipeline_runs_push_id_pushes"), ondelete="CASCADE"
|
||||
),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_pipeline_runs")),
|
||||
sa.UniqueConstraint("push_id", "stage", name=op.f("uq_pipeline_runs_push_id")),
|
||||
)
|
||||
# ### end Alembic commands ###
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.drop_table("pipeline_runs")
|
||||
op.drop_table("draft_changes")
|
||||
op.drop_table("pushes")
|
||||
op.drop_table("log_queries")
|
||||
op.drop_table("invites")
|
||||
op.drop_table("drafts")
|
||||
op.drop_table("alert_deliveries")
|
||||
op.drop_table("agent_requests")
|
||||
op.drop_table("users")
|
||||
op.drop_table("fsm_state")
|
||||
with op.batch_alter_table("container_snapshots", schema=None) as batch_op:
|
||||
batch_op.drop_index("ix_container_snapshots_container_at")
|
||||
|
||||
op.drop_table("container_snapshots")
|
||||
with op.batch_alter_table("alerts", schema=None) as batch_op:
|
||||
batch_op.drop_index("ix_alerts_kind_service_resolved_at")
|
||||
|
||||
op.drop_table("alerts")
|
||||
with op.batch_alter_table("access_attempts", schema=None) as batch_op:
|
||||
batch_op.drop_index("ix_access_attempts_tg_id_at")
|
||||
|
||||
op.drop_table("access_attempts")
|
||||
# ### end Alembic commands ###
|
||||
@@ -0,0 +1,283 @@
|
||||
"""Data model (ARCHITECTURE §2.10). One schema for PostgreSQL (prod) and SQLite (tests).
|
||||
|
||||
Enums are stored as VARCHAR + CHECK rather than native PG enums: portable and painless to migrate.
|
||||
"""
|
||||
|
||||
from datetime import UTC, datetime
|
||||
from enum import StrEnum
|
||||
from typing import Any, ClassVar
|
||||
|
||||
from sqlalchemy import (
|
||||
JSON,
|
||||
BigInteger,
|
||||
Boolean,
|
||||
DateTime,
|
||||
Enum,
|
||||
Float,
|
||||
ForeignKey,
|
||||
Index,
|
||||
Integer,
|
||||
LargeBinary,
|
||||
MetaData,
|
||||
String,
|
||||
Text,
|
||||
UniqueConstraint,
|
||||
)
|
||||
from sqlalchemy.dialects.postgresql import JSONB
|
||||
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship
|
||||
|
||||
# Stable constraint names so Alembic migrations stay deterministic across databases.
|
||||
NAMING = {
|
||||
"ix": "ix_%(column_0_label)s",
|
||||
"uq": "uq_%(table_name)s_%(column_0_name)s",
|
||||
"ck": "ck_%(table_name)s_%(constraint_name)s",
|
||||
"fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
|
||||
"pk": "pk_%(table_name)s",
|
||||
}
|
||||
|
||||
# Surrogate keys: BIGINT on Postgres, INTEGER on SQLite (only INTEGER PRIMARY KEY autoincrements there).
|
||||
Id = BigInteger().with_variant(Integer(), "sqlite")
|
||||
Json = JSON().with_variant(JSONB(), "postgresql")
|
||||
TgId = BigInteger
|
||||
|
||||
|
||||
def utcnow() -> datetime:
|
||||
return datetime.now(UTC)
|
||||
|
||||
|
||||
def enum(cls: type[StrEnum], name: str) -> Enum:
|
||||
return Enum(
|
||||
cls,
|
||||
name=name,
|
||||
native_enum=False,
|
||||
create_constraint=True, # SQLAlchemy 2 skips the CHECK for non-native enums unless asked
|
||||
length=20,
|
||||
values_callable=lambda e: [m.value for m in e],
|
||||
)
|
||||
|
||||
|
||||
class Role(StrEnum):
|
||||
MEMBER = "member"
|
||||
ADMIN = "admin"
|
||||
|
||||
|
||||
class ChangeOp(StrEnum):
|
||||
ADD = "add"
|
||||
MODIFY = "modify"
|
||||
DELETE = "delete"
|
||||
|
||||
|
||||
class Stage(StrEnum):
|
||||
CI = "ci"
|
||||
CD = "cd"
|
||||
|
||||
|
||||
class PipelineStatus(StrEnum):
|
||||
PENDING = "pending"
|
||||
RUNNING = "running"
|
||||
SUCCESS = "success"
|
||||
FAILED = "failed"
|
||||
TIMEOUT = "timeout"
|
||||
SKIPPED = "skipped"
|
||||
|
||||
|
||||
class AlertKind(StrEnum):
|
||||
CRASH = "crash"
|
||||
OOM = "oom"
|
||||
RESTART_LOOP = "restart_loop"
|
||||
UNHEALTHY = "unhealthy"
|
||||
OVERLOAD = "overload"
|
||||
DEPLOY_REFUSED = "deploy_refused"
|
||||
CD_FAILED = "cd_failed"
|
||||
DISK_LOW = "disk_low"
|
||||
DEPENDENCY_DOWN = "dependency_down"
|
||||
AGENT_ATTACK = "agent_attack"
|
||||
|
||||
|
||||
class LogMode(StrEnum):
|
||||
TAIL = "tail"
|
||||
GREP = "grep"
|
||||
AGENT = "agent"
|
||||
|
||||
|
||||
class Verdict(StrEnum):
|
||||
SAFE = "safe"
|
||||
ATTACK = "attack"
|
||||
ERROR = "error"
|
||||
|
||||
|
||||
class Base(DeclarativeBase):
|
||||
metadata = MetaData(naming_convention=NAMING)
|
||||
type_annotation_map: ClassVar[dict[Any, Any]] = {
|
||||
datetime: DateTime(timezone=True),
|
||||
dict[str, Any]: Json,
|
||||
list[str]: Json,
|
||||
}
|
||||
|
||||
|
||||
class User(Base):
|
||||
__tablename__ = "users"
|
||||
|
||||
tg_id: Mapped[int] = mapped_column(TgId, primary_key=True, autoincrement=False)
|
||||
username: Mapped[str | None] = mapped_column(String(64))
|
||||
full_name: Mapped[str] = mapped_column(String(256))
|
||||
role: Mapped[Role] = mapped_column(enum(Role, "role"), default=Role.MEMBER)
|
||||
is_active: Mapped[bool] = mapped_column(Boolean, default=True)
|
||||
invited_by: Mapped[int | None] = mapped_column(TgId, ForeignKey("users.tg_id"))
|
||||
joined_at: Mapped[datetime] = mapped_column(default=utcnow)
|
||||
last_seen_at: Mapped[datetime | None]
|
||||
|
||||
|
||||
class Invite(Base):
|
||||
__tablename__ = "invites"
|
||||
|
||||
id: Mapped[int] = mapped_column(Id, primary_key=True)
|
||||
# sha256 hex of the token; the token itself is never stored (SPEC A-3)
|
||||
token_hash: Mapped[str] = mapped_column(String(64), unique=True)
|
||||
created_by: Mapped[int] = mapped_column(TgId, ForeignKey("users.tg_id"))
|
||||
created_at: Mapped[datetime] = mapped_column(default=utcnow)
|
||||
expires_at: Mapped[datetime]
|
||||
used_by: Mapped[int | None] = mapped_column(TgId, ForeignKey("users.tg_id"))
|
||||
used_at: Mapped[datetime | None]
|
||||
|
||||
|
||||
class AccessAttempt(Base):
|
||||
__tablename__ = "access_attempts"
|
||||
__table_args__ = (Index("ix_access_attempts_tg_id_at", "tg_id", "at"),)
|
||||
|
||||
id: Mapped[int] = mapped_column(Id, primary_key=True)
|
||||
tg_id: Mapped[int] = mapped_column(TgId)
|
||||
at: Mapped[datetime] = mapped_column(default=utcnow)
|
||||
ok: Mapped[bool] = mapped_column(Boolean)
|
||||
|
||||
|
||||
class FsmState(Base):
|
||||
__tablename__ = "fsm_state"
|
||||
|
||||
key: Mapped[str] = mapped_column(String(255), primary_key=True)
|
||||
state: Mapped[str | None] = mapped_column(String(255))
|
||||
data: Mapped[dict[str, Any]] = mapped_column(default=dict)
|
||||
|
||||
|
||||
class Draft(Base):
|
||||
__tablename__ = "drafts"
|
||||
|
||||
id: Mapped[int] = mapped_column(Id, primary_key=True)
|
||||
user_id: Mapped[int] = mapped_column(TgId, ForeignKey("users.tg_id", ondelete="CASCADE"), unique=True)
|
||||
base_sha: Mapped[str] = mapped_column(String(40))
|
||||
message: Mapped[str | None] = mapped_column(Text)
|
||||
created_at: Mapped[datetime] = mapped_column(default=utcnow)
|
||||
expires_at: Mapped[datetime]
|
||||
|
||||
changes: Mapped[list["DraftChange"]] = relationship(
|
||||
back_populates="draft", cascade="all, delete-orphan", order_by="DraftChange.id"
|
||||
)
|
||||
|
||||
|
||||
class DraftChange(Base):
|
||||
__tablename__ = "draft_changes"
|
||||
__table_args__ = (UniqueConstraint("draft_id", "path"),)
|
||||
|
||||
id: Mapped[int] = mapped_column(Id, primary_key=True)
|
||||
draft_id: Mapped[int] = mapped_column(Id, ForeignKey("drafts.id", ondelete="CASCADE"))
|
||||
path: Mapped[str] = mapped_column(String(500))
|
||||
op: Mapped[ChangeOp] = mapped_column(enum(ChangeOp, "change_op"))
|
||||
content: Mapped[bytes | None] = mapped_column(LargeBinary)
|
||||
|
||||
draft: Mapped[Draft] = relationship(back_populates="changes")
|
||||
|
||||
|
||||
class Push(Base):
|
||||
__tablename__ = "pushes"
|
||||
|
||||
id: Mapped[int] = mapped_column(Id, primary_key=True)
|
||||
user_id: Mapped[int] = mapped_column(TgId, ForeignKey("users.tg_id"))
|
||||
sha: Mapped[str] = mapped_column(String(40), unique=True)
|
||||
message: Mapped[str] = mapped_column(Text)
|
||||
paths: Mapped[list[str]] = mapped_column(default=list)
|
||||
chat_id: Mapped[int] = mapped_column(TgId)
|
||||
message_id: Mapped[int | None] = mapped_column(BigInteger)
|
||||
ci_status: Mapped[PipelineStatus] = mapped_column(enum(PipelineStatus, "ci_status"), default=PipelineStatus.PENDING)
|
||||
cd_status: Mapped[PipelineStatus] = mapped_column(enum(PipelineStatus, "cd_status"), default=PipelineStatus.PENDING)
|
||||
created_at: Mapped[datetime] = mapped_column(default=utcnow)
|
||||
|
||||
|
||||
class PipelineRun(Base):
|
||||
__tablename__ = "pipeline_runs"
|
||||
__table_args__ = (UniqueConstraint("push_id", "stage"),)
|
||||
|
||||
id: Mapped[int] = mapped_column(Id, primary_key=True)
|
||||
push_id: Mapped[int] = mapped_column(Id, ForeignKey("pushes.id", ondelete="CASCADE"))
|
||||
# CI and CD are tasks of one SourceCraft workflow (check-and-deploy), hence "stage", not "workflow"
|
||||
stage: Mapped[Stage] = mapped_column(enum(Stage, "stage"))
|
||||
run_slug: Mapped[str] = mapped_column(String(50))
|
||||
status: Mapped[PipelineStatus] = mapped_column(enum(PipelineStatus, "run_status"))
|
||||
failed_cube: Mapped[str | None] = mapped_column(String(100))
|
||||
url: Mapped[str | None] = mapped_column(String(500))
|
||||
updated_at: Mapped[datetime] = mapped_column(default=utcnow, onupdate=utcnow)
|
||||
|
||||
|
||||
class ContainerSnapshot(Base):
|
||||
__tablename__ = "container_snapshots"
|
||||
__table_args__ = (Index("ix_container_snapshots_container_at", "container", "at"),)
|
||||
|
||||
id: Mapped[int] = mapped_column(Id, primary_key=True)
|
||||
container: Mapped[str] = mapped_column(String(100))
|
||||
at: Mapped[datetime] = mapped_column(default=utcnow)
|
||||
restart_count: Mapped[int] = mapped_column(Integer)
|
||||
cpu_pct: Mapped[float] = mapped_column(Float)
|
||||
mem_pct: Mapped[float] = mapped_column(Float)
|
||||
health: Mapped[str | None] = mapped_column(String(20))
|
||||
|
||||
|
||||
class Alert(Base):
|
||||
__tablename__ = "alerts"
|
||||
__table_args__ = (Index("ix_alerts_kind_service_resolved_at", "kind", "service", "resolved_at"),)
|
||||
|
||||
id: Mapped[int] = mapped_column(Id, primary_key=True)
|
||||
kind: Mapped[AlertKind] = mapped_column(enum(AlertKind, "alert_kind"))
|
||||
service: Mapped[str | None] = mapped_column(String(100))
|
||||
payload: Mapped[dict[str, Any]] = mapped_column(default=dict)
|
||||
count: Mapped[int] = mapped_column(Integer, default=1)
|
||||
opened_at: Mapped[datetime] = mapped_column(default=utcnow)
|
||||
last_at: Mapped[datetime] = mapped_column(default=utcnow)
|
||||
resolved_at: Mapped[datetime | None]
|
||||
|
||||
deliveries: Mapped[list["AlertDelivery"]] = relationship(back_populates="alert", cascade="all, delete-orphan")
|
||||
|
||||
|
||||
class AlertDelivery(Base):
|
||||
__tablename__ = "alert_deliveries"
|
||||
__table_args__ = (UniqueConstraint("alert_id", "chat_id"),)
|
||||
|
||||
id: Mapped[int] = mapped_column(Id, primary_key=True)
|
||||
alert_id: Mapped[int] = mapped_column(Id, ForeignKey("alerts.id", ondelete="CASCADE"))
|
||||
chat_id: Mapped[int] = mapped_column(TgId)
|
||||
message_id: Mapped[int] = mapped_column(BigInteger)
|
||||
|
||||
alert: Mapped[Alert] = relationship(back_populates="deliveries")
|
||||
|
||||
|
||||
class LogQuery(Base):
|
||||
__tablename__ = "log_queries"
|
||||
|
||||
id: Mapped[int] = mapped_column(Id, primary_key=True)
|
||||
user_id: Mapped[int] = mapped_column(TgId, ForeignKey("users.tg_id"))
|
||||
container: Mapped[str] = mapped_column(String(100))
|
||||
mode: Mapped[LogMode] = mapped_column(enum(LogMode, "log_mode"))
|
||||
query: Mapped[dict[str, Any]] = mapped_column(default=dict)
|
||||
lines: Mapped[int] = mapped_column(Integer, default=0)
|
||||
at: Mapped[datetime] = mapped_column(default=utcnow)
|
||||
|
||||
|
||||
class AgentRequest(Base):
|
||||
__tablename__ = "agent_requests"
|
||||
|
||||
id: Mapped[int] = mapped_column(Id, primary_key=True)
|
||||
user_id: Mapped[int] = mapped_column(TgId, ForeignKey("users.tg_id"))
|
||||
prompt: Mapped[str] = mapped_column(Text)
|
||||
verdict: Mapped[Verdict] = mapped_column(enum(Verdict, "verdict"))
|
||||
reason: Mapped[str | None] = mapped_column(Text)
|
||||
log_query: Mapped[dict[str, Any] | None]
|
||||
tokens: Mapped[int] = mapped_column(Integer, default=0)
|
||||
at: Mapped[datetime] = mapped_column(default=utcnow)
|
||||
@@ -0,0 +1,35 @@
|
||||
"""Async engine and session factory."""
|
||||
|
||||
from sqlalchemy import event
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine
|
||||
from sqlalchemy.pool import StaticPool
|
||||
|
||||
SessionMaker = async_sessionmaker[AsyncSession]
|
||||
|
||||
|
||||
def make_engine(url: str) -> AsyncEngine:
|
||||
if url.startswith("sqlite"):
|
||||
# In-memory SQLite lives per connection: share one connection so every session sees the same DB.
|
||||
kwargs = (
|
||||
{"poolclass": StaticPool, "connect_args": {"check_same_thread": False}}
|
||||
if ":memory:" in url or url.endswith("://")
|
||||
else {}
|
||||
)
|
||||
engine = create_async_engine(url, **kwargs)
|
||||
event.listen(engine.sync_engine, "connect", _sqlite_foreign_keys)
|
||||
return engine
|
||||
return create_async_engine(url, pool_pre_ping=True)
|
||||
|
||||
|
||||
def _sqlite_foreign_keys(dbapi_conn: object, _record: object) -> None:
|
||||
# SQLite ignores FOREIGN KEY / ON DELETE CASCADE unless asked per connection
|
||||
cursor = dbapi_conn.cursor() # type: ignore[attr-defined]
|
||||
cursor.execute("PRAGMA foreign_keys=ON")
|
||||
cursor.close()
|
||||
|
||||
|
||||
def make_sessionmaker(engine: AsyncEngine) -> SessionMaker:
|
||||
return async_sessionmaker(engine, expire_on_commit=False)
|
||||
|
||||
|
||||
__all__ = ["SessionMaker", "make_engine", "make_sessionmaker"]
|
||||
Reference in New Issue
Block a user