diff --git a/README.md b/README.md index cc43c7a..f47b269 100644 --- a/README.md +++ b/README.md @@ -129,8 +129,17 @@ MPU_LOGIN и MPU_PASSWORD в свой `.env`, дальше `scripts/check_lk_api ``` Поднимается на stdio-транспорте — так подключается локальный MCP-клиент -(например, Claude Desktop/Code). Пока один инструмент, `list_groups` -(ФОД-001) — список групп с кэшем в БД на 15 минут. +(например, Claude Desktop/Code). Инструменты: + +| Инструмент | Что делает | +|---|---| +| `list_groups` | список групп университета (ФОД-001), кэш в БД на 15 минут | +| `add_tracked_group(user, group)` | добавить группу в отслеживаемые пользователя; группа сверяется со списком `list_groups`, повторное добавление ничего не дублирует | +| `list_tracked_groups(user)` | отслеживаемые группы пользователя в порядке добавления | +| `remove_tracked_group(user, group)` | убрать группу из отслеживаемых | + +Отслеживаемые группы лежат в таблице `tracked_groups`. Авторизации на +сервере пока нет, так что `user` — просто логин, который передаёт клиент. ### HTTP-режим @@ -180,11 +189,12 @@ make docker-down `POSTGRES_USER`/`POSTGRES_PASSWORD`/`POSTGRES_DB` в `.env` — compose читает его сам, а `DATABASE_URL` оттуда для контейнера не используется. -Посмотреть, что кэш реально лёг в контейнерную БД (это же удобно показывать -после рестарта): +Посмотреть, что данные реально легли в контейнерную БД (это же удобно +показывать после рестарта): ```bash docker compose exec db psql -U mospolytech_mcp -c "select id, jsonb_array_length(groups), fetched_at from groups_cache" +docker compose exec db psql -U mospolytech_mcp -c "select * from tracked_groups" ``` ## Структура репозитория diff --git a/alembic/versions/5b1c7e2a9f40_create_tracked_groups.py b/alembic/versions/5b1c7e2a9f40_create_tracked_groups.py new file mode 100644 index 0000000..eb4d545 --- /dev/null +++ b/alembic/versions/5b1c7e2a9f40_create_tracked_groups.py @@ -0,0 +1,34 @@ +"""create tracked_groups + +Revision ID: 5b1c7e2a9f40 +Revises: d3ae49746cfb +Create Date: 2026-09-23 12:00:00.000000 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +revision: str = '5b1c7e2a9f40' +down_revision: Union[str, Sequence[str], None] = 'd3ae49746cfb' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.create_table('tracked_groups', + sa.Column('id', sa.Integer(), nullable=False), + sa.Column('user_login', sa.String(length=100), nullable=False), + sa.Column('group_name', sa.String(length=50), nullable=False), + sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.text('now()'), nullable=False), + sa.PrimaryKeyConstraint('id'), + sa.UniqueConstraint('user_login', 'group_name') + ) + op.create_index(op.f('ix_tracked_groups_user_login'), 'tracked_groups', ['user_login'], unique=False) + + +def downgrade() -> None: + op.drop_index(op.f('ix_tracked_groups_user_login'), table_name='tracked_groups') + op.drop_table('tracked_groups') diff --git a/src/mospolytech_mcp/db/__init__.py b/src/mospolytech_mcp/db/__init__.py index 94d40a1..b0b628d 100644 --- a/src/mospolytech_mcp/db/__init__.py +++ b/src/mospolytech_mcp/db/__init__.py @@ -1,7 +1,7 @@ # Пакет db: доступ к PostgreSQL через SQLAlchemy, схема в alembic/. from .base import Base -from .models import GroupsCache +from .models import GroupsCache, TrackedGroup from .session import make_engine, make_sessionmaker -__all__ = ["Base", "GroupsCache", "make_engine", "make_sessionmaker"] +__all__ = ["Base", "GroupsCache", "TrackedGroup", "make_engine", "make_sessionmaker"] diff --git a/src/mospolytech_mcp/db/models.py b/src/mospolytech_mcp/db/models.py index 5b066a6..291b928 100644 --- a/src/mospolytech_mcp/db/models.py +++ b/src/mospolytech_mcp/db/models.py @@ -4,7 +4,7 @@ from __future__ import annotations from datetime import datetime -from sqlalchemy import DateTime +from sqlalchemy import DateTime, String, UniqueConstraint, func from sqlalchemy.dialects.postgresql import JSONB from sqlalchemy.orm import Mapped, mapped_column @@ -17,3 +17,15 @@ class GroupsCache(Base): id: Mapped[int] = mapped_column(primary_key=True) groups: Mapped[list[str]] = mapped_column(JSONB) fetched_at: Mapped[datetime] = mapped_column(DateTime(timezone=True)) + + +class TrackedGroup(Base): + __tablename__ = "tracked_groups" + __table_args__ = (UniqueConstraint("user_login", "group_name"),) + + id: Mapped[int] = mapped_column(primary_key=True) + user_login: Mapped[str] = mapped_column(String(100), index=True) + group_name: Mapped[str] = mapped_column(String(50)) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), server_default=func.now() + ) diff --git a/src/mospolytech_mcp/server.py b/src/mospolytech_mcp/server.py index 5c8585f..3e064c1 100644 --- a/src/mospolytech_mcp/server.py +++ b/src/mospolytech_mcp/server.py @@ -8,11 +8,13 @@ import sys from pathlib import Path from mcp.server.mcpserver import MCPServer +from mcp.server.mcpserver.exceptions import ToolError from ._dotenv import load_dotenv from .api import UniversityAPI from .cache import GroupsCacheStore from .db import make_engine, make_sessionmaker +from .tracked import TrackedGroupsStore, UnknownGroupError if sys.platform == "win32": asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy()) @@ -23,7 +25,9 @@ mcp = MCPServer(name="mospolytech-mcp", version="0.1.0") _uni = UniversityAPI() _engine = make_engine() -_groups_cache = GroupsCacheStore(_uni.open, make_sessionmaker(_engine)) +_sessionmaker = make_sessionmaker(_engine) +_groups_cache = GroupsCacheStore(_uni.open, _sessionmaker) +_tracked = TrackedGroupsStore(_sessionmaker, _groups_cache.get_groups) @mcp.tool() @@ -32,6 +36,32 @@ async def list_groups() -> list[str]: return await _groups_cache.get_groups() +@mcp.tool() +async def add_tracked_group(user: str, group: str) -> str: + """Добавить группу в отслеживаемые для пользователя. group — номер вида 241-327.""" + try: + added = await _tracked.add(user, group) + except UnknownGroupError as exc: + raise ToolError(str(exc)) from exc + if added: + return f"группа {group} добавлена в отслеживаемые для {user}" + return f"группа {group} уже отслеживается для {user}" + + +@mcp.tool() +async def list_tracked_groups(user: str) -> list[str]: + """Группы, которые отслеживает пользователь, в порядке добавления.""" + return await _tracked.list(user) + + +@mcp.tool() +async def remove_tracked_group(user: str, group: str) -> str: + """Убрать группу из отслеживаемых для пользователя.""" + if await _tracked.remove(user, group): + return f"группа {group} больше не отслеживается для {user}" + return f"группа {group} не была в отслеживаемых для {user}" + + def main() -> None: parser = argparse.ArgumentParser(prog="mospolytech-mcp") parser.add_argument("--http", action="store_true", help="HTTP вместо stdio") diff --git a/src/mospolytech_mcp/tracked.py b/src/mospolytech_mcp/tracked.py new file mode 100644 index 0000000..2387957 --- /dev/null +++ b/src/mospolytech_mcp/tracked.py @@ -0,0 +1,64 @@ +# Отслеживаемые пользователем группы: хранятся в tracked_groups. + +from __future__ import annotations + +from collections.abc import Awaitable, Callable + +from sqlalchemy import delete, select +from sqlalchemy.dialects.postgresql import insert +from sqlalchemy.ext.asyncio import async_sessionmaker + +from .db import TrackedGroup + + +class UnknownGroupError(ValueError): + pass + + +class TrackedGroupsStore: + def __init__( + self, + sessionmaker: async_sessionmaker, + known_groups: Callable[[], Awaitable[list[str]]], + ) -> None: + self._sessionmaker = sessionmaker + self._known_groups = known_groups + + async def add(self, user_login: str, group_name: str) -> bool: + group_name = group_name.strip() + if group_name not in await self._known_groups(): + raise UnknownGroupError(f"группы {group_name!r} нет в списке групп университета") + + stmt = ( + insert(TrackedGroup) + .values(user_login=user_login, group_name=group_name) + .on_conflict_do_nothing(index_elements=["user_login", "group_name"]) + .returning(TrackedGroup.id) + ) + async with self._sessionmaker() as session: + new_id = await session.scalar(stmt) + await session.commit() + return new_id is not None + + async def list(self, user_login: str) -> list[str]: + stmt = ( + select(TrackedGroup.group_name) + .where(TrackedGroup.user_login == user_login) + .order_by(TrackedGroup.created_at, TrackedGroup.id) + ) + async with self._sessionmaker() as session: + return list(await session.scalars(stmt)) + + async def remove(self, user_login: str, group_name: str) -> bool: + stmt = ( + delete(TrackedGroup) + .where( + TrackedGroup.user_login == user_login, + TrackedGroup.group_name == group_name.strip(), + ) + .returning(TrackedGroup.id) + ) + async with self._sessionmaker() as session: + removed_id = await session.scalar(stmt) + await session.commit() + return removed_id is not None diff --git a/tests/test_tracked_offline.py b/tests/test_tracked_offline.py new file mode 100644 index 0000000..0c7391f --- /dev/null +++ b/tests/test_tracked_offline.py @@ -0,0 +1,27 @@ +# Офлайн-тесты TrackedGroupsStore: то, что проверяется до обращения к БД. + +from __future__ import annotations + +import pytest + +from mospolytech_mcp.tracked import TrackedGroupsStore, UnknownGroupError + + +def _no_db(): + raise AssertionError("до БД дойти не должны") + + +async def _known(): + return ["241-327", "221-111"] + + +async def test_add_rejects_unknown_group(): + store = TrackedGroupsStore(_no_db, _known) + with pytest.raises(UnknownGroupError, match="999-999"): + await store.add("r.v.starkov", "999-999") + + +async def test_add_rejects_unknown_group_after_strip(): + store = TrackedGroupsStore(_no_db, _known) + with pytest.raises(UnknownGroupError): + await store.add("r.v.starkov", " 241-328 ")