diff --git a/.gitignore b/.gitignore index ce098f2..305f680 100644 --- a/.gitignore +++ b/.gitignore @@ -23,3 +23,6 @@ build/ .vscode/ .claude/worktrees/ .DS_Store + +# local SQLite DB (bot without Docker) +*.db diff --git a/alembic.ini b/alembic.ini new file mode 100644 index 0000000..3ee13ca --- /dev/null +++ b/alembic.ini @@ -0,0 +1,5 @@ +# Developer CLI: `uv run alembic revision --autogenerate -m "..."`. The DB URL comes from DATABASE_URL (.env), +# the bot itself applies migrations at startup (app/db/migrate.py). +[alembic] +script_location = app.db:migrations +file_template = %%(year)d%%(month).2d%%(day).2d_%%(rev)s_%%(slug)s diff --git a/app/bot/fsm.py b/app/bot/fsm.py index 4633dfc..6e5893e 100644 --- a/app/bot/fsm.py +++ b/app/bot/fsm.py @@ -1,8 +1,10 @@ -"""FSM storage factory. In-memory for now; Postgres-backed storage that survives restarts arrives with the DB (#5).""" +"""FSM storage factory: dialog state lives in the DB and survives restarts (SPEC M-4).""" from aiogram.fsm.storage.base import BaseStorage -from aiogram.fsm.storage.memory import MemoryStorage + +from app.bot.fsm_storage import SqlStorage +from app.db.session import SessionMaker -def make_storage() -> BaseStorage: - return MemoryStorage() +def make_storage(sessions: SessionMaker) -> BaseStorage: + return SqlStorage(sessions) diff --git a/app/bot/fsm_storage.py b/app/bot/fsm_storage.py new file mode 100644 index 0000000..6964c4a --- /dev/null +++ b/app/bot/fsm_storage.py @@ -0,0 +1,54 @@ +"""aiogram FSM storage on top of the `fsm_state` table — dialog state survives bot restarts (SPEC M-4).""" + +from collections.abc import Mapping +from typing import Any + +from aiogram.fsm.state import State +from aiogram.fsm.storage.base import BaseStorage, StateType, StorageKey + +from app.db.models import FsmState +from app.db.session import SessionMaker + + +def key_to_str(key: StorageKey) -> str: + parts = (key.bot_id, key.chat_id, key.user_id, key.thread_id, key.business_connection_id, key.destiny) + return ":".join("" if p is None else str(p) for p in parts) + + +class SqlStorage(BaseStorage): + def __init__(self, sessions: SessionMaker) -> None: + self._sessions = sessions + + async def set_state(self, key: StorageKey, state: StateType = None) -> None: + value = state.state if isinstance(state, State) else state + await self._write(key, state=value) + + async def get_state(self, key: StorageKey) -> str | None: + async with self._sessions() as session: + row = await session.get(FsmState, key_to_str(key)) + return row.state if row else None + + async def set_data(self, key: StorageKey, data: Mapping[str, Any]) -> None: + await self._write(key, data=dict(data)) + + async def get_data(self, key: StorageKey) -> dict[str, Any]: + async with self._sessions() as session: + row = await session.get(FsmState, key_to_str(key)) + return dict(row.data) if row else {} + + async def close(self) -> None: + pass # the engine is owned and disposed by the app + + async def _write(self, key: StorageKey, **fields: Any) -> None: + async with self._sessions.begin() as session: + row = await session.get(FsmState, key_to_str(key)) + state = fields.get("state", row.state if row else None) + data = fields.get("data", row.data if row else {}) + if state is None and not data: + if row is not None: + await session.delete(row) # empty state: keep the table small + return + if row is None: + session.add(FsmState(key=key_to_str(key), state=state, data=data)) + else: + row.state, row.data = state, data diff --git a/app/config.py b/app/config.py index 64ed71b..7d3ea49 100644 --- a/app/config.py +++ b/app/config.py @@ -29,6 +29,13 @@ class AccessSettings(BaseSettings): return value +class DatabaseSettings(BaseSettings): + model_config = SettingsConfigDict(**_ENV, env_prefix="DATABASE_") + + # Compose sets a PostgreSQL URL; the default lets the bot run locally without Docker. + url: str = "sqlite+aiosqlite:///./tgit.db" + + class RuntimeSettings(BaseSettings): model_config = _ENV @@ -46,4 +53,5 @@ class Settings(BaseSettings): telegram: TelegramSettings = Field(default_factory=TelegramSettings) access: AccessSettings = Field(default_factory=AccessSettings) + database: DatabaseSettings = Field(default_factory=DatabaseSettings) runtime: RuntimeSettings = Field(default_factory=RuntimeSettings) diff --git a/app/db/__init__.py b/app/db/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/app/db/maintenance.py b/app/db/maintenance.py new file mode 100644 index 0000000..8486dc5 --- /dev/null +++ b/app/db/maintenance.py @@ -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) diff --git a/app/db/migrate.py b/app/db/migrate.py new file mode 100644 index 0000000..8169f5f --- /dev/null +++ b/app/db/migrate.py @@ -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() diff --git a/app/db/migrations/__init__.py b/app/db/migrations/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/app/db/migrations/env.py b/app/db/migrations/env.py new file mode 100644 index 0000000..72314d1 --- /dev/null +++ b/app/db/migrations/env.py @@ -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()) diff --git a/app/db/migrations/script.py.mako b/app/db/migrations/script.py.mako new file mode 100644 index 0000000..e67e889 --- /dev/null +++ b/app/db/migrations/script.py.mako @@ -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"} diff --git a/app/db/migrations/versions/20261014_0001_initial_schema.py b/app/db/migrations/versions/20261014_0001_initial_schema.py new file mode 100644 index 0000000..acdf64e --- /dev/null +++ b/app/db/migrations/versions/20261014_0001_initial_schema.py @@ -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 ### diff --git a/app/db/migrations/versions/__init__.py b/app/db/migrations/versions/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/app/db/models.py b/app/db/models.py new file mode 100644 index 0000000..69c886f --- /dev/null +++ b/app/db/models.py @@ -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) diff --git a/app/db/session.py b/app/db/session.py new file mode 100644 index 0000000..c38ef53 --- /dev/null +++ b/app/db/session.py @@ -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"] diff --git a/app/main.py b/app/main.py index 4199ad0..097eefb 100644 --- a/app/main.py +++ b/app/main.py @@ -11,6 +11,9 @@ from app.bot.fsm import make_storage from app.bot.middlewares.access import AccessMiddleware from app.bot.routers.common import build_router from app.config import Settings +from app.db.maintenance import cleanup_loop +from app.db.migrate import upgrade +from app.db.session import SessionMaker, make_engine, make_sessionmaker from app.health import heartbeat from app.logs_setup import setup_logging from app.texts import Texts, load_texts @@ -18,13 +21,13 @@ from app.texts import Texts, load_texts log = logging.getLogger("tgit") -def build_dispatcher(settings: Settings, texts: Texts) -> Dispatcher: +def build_dispatcher(settings: Settings, texts: Texts, sessions: SessionMaker) -> Dispatcher: admin_ids = settings.access.admin_ids async def is_member(tg_id: int) -> bool: return tg_id in admin_ids - dp = Dispatcher(storage=make_storage()) + dp = Dispatcher(storage=make_storage(sessions)) access = AccessMiddleware(is_member, texts) # Outer middleware: runs before filters, so nothing — not even a stale button — reaches a handler unchecked. dp.message.outer_middleware(access) @@ -35,18 +38,27 @@ def build_dispatcher(settings: Settings, texts: Texts) -> Dispatcher: async def run(settings: Settings) -> None: texts = load_texts(settings.runtime.config_dir) - dp = build_dispatcher(settings, texts) + engine = make_engine(settings.database.url) + await upgrade(engine) + sessions = make_sessionmaker(engine) + dp = build_dispatcher(settings, texts, sessions) rt = settings.runtime - beat = asyncio.create_task(heartbeat(rt.heartbeat_file, rt.heartbeat_interval)) + background = [ + asyncio.create_task(heartbeat(rt.heartbeat_file, rt.heartbeat_interval)), + asyncio.create_task(cleanup_loop(sessions)), + ] bot = Bot(settings.telegram.bot_token.get_secret_value()) log.info("tgit %s starting, %d bootstrap admin(s)", __version__, len(settings.access.admin_ids)) try: await dp.start_polling(bot, handle_signals=True) finally: - beat.cancel() - with contextlib.suppress(asyncio.CancelledError): - await beat + for task in background: + task.cancel() + for task in background: + with contextlib.suppress(asyncio.CancelledError): + await task await bot.session.close() + await engine.dispose() log.info("tgit stopped") diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 7e319a7..cbcc2e6 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -138,7 +138,7 @@ SQLAlchemy 2 (async) + Alembic. PostgreSQL в Compose; в тестах — SQLit | `drafts` | `id`, `user_id` UNIQUE, `base_sha`, `message`, `created_at`, `expires_at` | | `draft_changes` | `id`, `draft_id`, `path`, `op` (`add`/`modify`/`delete`), `content` BYTEA | | `pushes` | `id`, `user_id`, `sha`, `message`, `paths` JSON, `chat_id`, `message_id`, `ci_status`, `cd_status`, `created_at` | -| `pipeline_runs` | `id`, `push_id`, `workflow` (`check`/`deploy`), `run_slug`, `status`, `failed_cube`, `url`, `updated_at` | +| `pipeline_runs` | `id`, `push_id`, `stage` (`ci`/`cd` — задачи одного workflow `check-and-deploy`), `run_slug`, `status`, `failed_cube`, `url`, `updated_at` | | `container_snapshots` | `container`, `at`, `restart_count`, `cpu_pct`, `mem_pct`, `health` — для проверок «за 5 мин» (TTL 1 сутки) | | `alerts` | `id`, `kind`, `service`, `payload` JSON, `count`, `opened_at`, `last_at`, `resolved_at` | | `alert_deliveries` | `alert_id`, `chat_id`, `message_id` | diff --git a/tests/bot/test_common.py b/tests/bot/test_common.py index cfc510f..0be65f0 100644 --- a/tests/bot/test_common.py +++ b/tests/bot/test_common.py @@ -4,6 +4,7 @@ from aiogram.types import InlineKeyboardMarkup, ReplyKeyboardMarkup from app.bot.menu import Nav, NavAction from app.config import Settings +from app.db.session import SessionMaker from app.main import build_dispatcher from app.texts import Texts from tests.conftest import ADMIN, OUTSIDER @@ -11,8 +12,8 @@ from tests.fakes import TgClient @pytest.fixture -def tg(settings: Settings, texts: Texts) -> TgClient: - return TgClient(build_dispatcher(settings, texts)) +def tg(settings: Settings, texts: Texts, sessions: SessionMaker) -> TgClient: + return TgClient(build_dispatcher(settings, texts, sessions)) def buttons(markup: ReplyKeyboardMarkup) -> list[str]: diff --git a/tests/bot/test_fsm_storage.py b/tests/bot/test_fsm_storage.py new file mode 100644 index 0000000..5a639f6 --- /dev/null +++ b/tests/bot/test_fsm_storage.py @@ -0,0 +1,48 @@ +from aiogram.fsm.state import State, StatesGroup +from aiogram.fsm.storage.base import StorageKey +from sqlalchemy import func, select + +from app.bot.fsm_storage import SqlStorage, key_to_str +from app.db.models import FsmState +from app.db.session import SessionMaker + +KEY = StorageKey(bot_id=42, chat_id=7, user_id=7) + + +class Flow(StatesGroup): + name = State() + + +async def test_state_and_data_survive_a_new_storage(sessions: SessionMaker) -> None: + first = SqlStorage(sessions) + await first.set_state(KEY, Flow.name) + await first.set_data(KEY, {"path": "services/a", "page": 2}) + + # A fresh storage on the same DB = the bot after a restart + second = SqlStorage(sessions) + assert await second.get_state(KEY) == Flow.name.state + assert await second.get_data(KEY) == {"path": "services/a", "page": 2} + + +async def test_unknown_key_is_empty(sessions: SessionMaker) -> None: + storage = SqlStorage(sessions) + assert await storage.get_state(KEY) is None + assert await storage.get_data(KEY) == {} + + +async def test_cleared_state_removes_row(sessions: SessionMaker) -> None: + storage = SqlStorage(sessions) + await storage.set_state(KEY, "x") + await storage.set_data(KEY, {"a": 1}) + await storage.set_state(KEY, None) + await storage.set_data(KEY, {}) + await storage.close() + + async with sessions() as s: + assert await s.scalar(select(func.count()).select_from(FsmState)) == 0 + + +def test_key_includes_every_part() -> None: + key = StorageKey(bot_id=1, chat_id=2, user_id=3, thread_id=4, business_connection_id="b", destiny="d") + assert key_to_str(key) == "1:2:3:4:b:d" + assert key_to_str(KEY) == "42:7:7:::default" diff --git a/tests/conftest.py b/tests/conftest.py index b743f75..1d71aa1 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,8 +1,12 @@ +from collections.abc import AsyncIterator from pathlib import Path import pytest +from sqlalchemy.ext.asyncio import AsyncEngine from app.config import AccessSettings, RuntimeSettings, Settings, TelegramSettings +from app.db.migrate import upgrade +from app.db.session import SessionMaker, make_engine, make_sessionmaker from app.texts import Texts, load_texts ROOT = Path(__file__).resolve().parent.parent @@ -22,3 +26,17 @@ def settings(tmp_path: Path) -> Settings: access=AccessSettings(admin_ids=frozenset({ADMIN})), runtime=RuntimeSettings(config_dir=ROOT / "config", heartbeat_file=tmp_path / "hb"), ) + + +@pytest.fixture +async def engine() -> AsyncIterator[AsyncEngine]: + """Fresh in-memory SQLite with all migrations applied — the same path the bot takes at startup.""" + eng = make_engine("sqlite+aiosqlite://") + await upgrade(eng) + yield eng + await eng.dispose() + + +@pytest.fixture +def sessions(engine: AsyncEngine) -> SessionMaker: + return make_sessionmaker(engine) diff --git a/tests/db/__init__.py b/tests/db/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/db/test_maintenance.py b/tests/db/test_maintenance.py new file mode 100644 index 0000000..e4c33d9 --- /dev/null +++ b/tests/db/test_maintenance.py @@ -0,0 +1,30 @@ +from datetime import timedelta + +from sqlalchemy import func, select + +from app.db.maintenance import cleanup_expired +from app.db.models import AccessAttempt, ContainerSnapshot, Draft, User, utcnow +from app.db.session import SessionMaker + + +async def test_cleanup_removes_only_expired(sessions: SessionMaker) -> None: + now = utcnow() + async with sessions.begin() as s: + s.add_all([User(tg_id=1, full_name="a"), User(tg_id=2, full_name="b")]) + await s.flush() # no relationship() to order inserts: users first for the FK + s.add_all( + [ + Draft(user_id=1, base_sha="a" * 40, expires_at=now - timedelta(minutes=1)), + Draft(user_id=2, base_sha="b" * 40, expires_at=now + timedelta(hours=1)), + ContainerSnapshot(container="c", at=now - timedelta(days=2), restart_count=0, cpu_pct=0, mem_pct=0), + ContainerSnapshot(container="c", at=now, restart_count=0, cpu_pct=0, mem_pct=0), + AccessAttempt(tg_id=9, at=now - timedelta(days=2), ok=False), + ] + ) + + removed = await cleanup_expired(sessions, now) + + assert removed == {"drafts": 1, "snapshots": 1, "access_attempts": 1} + async with sessions() as s: + assert await s.scalar(select(func.count()).select_from(Draft)) == 1 + assert await s.scalar(select(func.count()).select_from(ContainerSnapshot)) == 1 diff --git a/tests/db/test_migrations.py b/tests/db/test_migrations.py new file mode 100644 index 0000000..f81b126 --- /dev/null +++ b/tests/db/test_migrations.py @@ -0,0 +1,82 @@ +from datetime import timedelta + +import pytest +from alembic.autogenerate import compare_metadata +from alembic.migration import MigrationContext +from sqlalchemy import inspect, text +from sqlalchemy.exc import IntegrityError +from sqlalchemy.ext.asyncio import AsyncEngine + +from app.db.migrate import downgrade, upgrade +from app.db.models import Base, ChangeOp, Draft, DraftChange, User, utcnow +from app.db.session import SessionMaker, make_engine + +TABLES = { + "users", + "invites", + "access_attempts", + "fsm_state", + "drafts", + "draft_changes", + "pushes", + "pipeline_runs", + "container_snapshots", + "alerts", + "alert_deliveries", + "log_queries", + "agent_requests", +} + + +async def table_names(engine: AsyncEngine) -> set[str]: + async with engine.connect() as conn: + return set(await conn.run_sync(lambda c: inspect(c).get_table_names())) - {"alembic_version"} + + +async def test_upgrade_creates_all_tables(engine: AsyncEngine) -> None: + assert await table_names(engine) == TABLES + + +async def test_models_match_migrations(engine: AsyncEngine) -> None: + # A model change without a migration shows up here as a non-empty diff. + async with engine.connect() as conn: + diff = await conn.run_sync(lambda c: compare_metadata(MigrationContext.configure(c), Base.metadata)) + assert diff == [] + + +async def test_downgrade_to_base_and_back() -> None: + eng = make_engine("sqlite+aiosqlite://") + try: + await upgrade(eng) + await downgrade(eng, "base") + assert await table_names(eng) == set() + await upgrade(eng) + assert await table_names(eng) == TABLES + finally: + await eng.dispose() + + +INSERT_USER = text( + "INSERT INTO users (tg_id, full_name, role, is_active, joined_at) VALUES (:id, 'u', :role, 1, '2026-01-01')" +) + + +async def test_enum_check_rejects_garbage(engine: AsyncEngine) -> None: + async with engine.begin() as conn: + await conn.execute(INSERT_USER, {"id": 1, "role": "admin"}) + with pytest.raises(IntegrityError): + async with engine.begin() as conn: + await conn.execute(INSERT_USER, {"id": 2, "role": "root"}) + + +async def test_deleting_draft_cascades_to_changes(sessions: SessionMaker) -> None: + async with sessions.begin() as s: + s.add(User(tg_id=1, full_name="a")) + await s.flush() + draft = Draft(user_id=1, base_sha="0" * 40, expires_at=utcnow() + timedelta(hours=1)) + draft.changes.append(DraftChange(path="services/a/main.py", op=ChangeOp.ADD, content=b"print(1)")) + s.add(draft) + async with sessions.begin() as s: + await s.execute(text("DELETE FROM drafts")) + async with sessions() as s: + assert (await s.execute(text("SELECT count(*) FROM draft_changes"))).scalar_one() == 0 diff --git a/tests/test_health.py b/tests/test_health.py index 287400f..60b4929 100644 --- a/tests/test_health.py +++ b/tests/test_health.py @@ -23,8 +23,9 @@ def test_age_decides(tmp_path: Path) -> None: async def test_heartbeat_touches_file(tmp_path: Path) -> None: hb = tmp_path / "hb" task = asyncio.create_task(health.heartbeat(hb, interval=10)) - await asyncio.sleep(0) - assert hb.exists() + async with asyncio.timeout(2): + while not hb.exists(): # noqa: ASYNC110 - polling the filesystem; touch runs in a worker thread + await asyncio.sleep(0.01) task.cancel() with contextlib.suppress(asyncio.CancelledError): await task