server/server.py - master - tinypot/lab4

server/server.py
213 строк6 KB
Tim Polus
lab4: Task Management API через RabbitMQ
02 мар 2026, 06:17
02 мар 2026, 06:1752ab147
from __future__ import annotations

import json
import logging
import time

import pika

from database import (
    init_db, Session, User, ApiKey, IdempotencyRecord,
)
from handlers import dispatch

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)
log = logging.getLogger("server")

RABBITMQ_HOST = "localhost"
REQUEST_QUEUE = "api.requests"
DLX_EXCHANGE = "api.dlx"
DLQ_QUEUE = "api.dlq"
MAX_RETRIES = 3


def authenticate(api_key: str | None) -> int | None:
    if not api_key:
        return None
    session = Session()
    try:
        record = session.query(ApiKey).filter_by(key=api_key, is_active=True).first()
        return record.user_id if record else None
    finally:
        session.close()


def check_idempotency(request_id: str) -> dict | None:
    session = Session()
    try:
        rec = session.query(IdempotencyRecord).filter_by(request_id=request_id).first()
        if rec:
            return json.loads(rec.response_body)
        return None
    finally:
        session.close()


def save_idempotency(request_id: str, response: dict) -> None:
    session = Session()
    try:
        session.add(IdempotencyRecord(
            request_id=request_id,
            response_body=json.dumps(response, ensure_ascii=False),
        ))
        session.commit()
    finally:
        session.close()


def seed_admin():
    session = Session()
    try:
        if session.query(User).filter_by(username="admin").first():
            return
        from passlib.hash import bcrypt
        user = User(
            email="admin@example.com",
            username="admin",
            hashed_password=bcrypt.hash("admin"),
            full_name="Administrator",
            is_superuser=True,
        )
        session.add(user)
        session.flush()
        session.add(ApiKey(key="master-api-key", user_id=user.id))
        session.commit()
        log.info("Создан admin + API-ключ: master-api-key")
    finally:
        session.close()


def build_response(
    correlation_id: str,
    status: str = "ok",
    data: dict | None = None,
    error: str | None = None,
) -> dict:
    return {
        "correlation_id": correlation_id,
        "status": status,
        "data": data,
        "error": error,
    }


def process_message(body: bytes) -> dict:
    msg = json.loads(body)
    request_id = msg.get("id")
    version = msg.get("version", "v1")
    action = msg.get("action")
    data = msg.get("data", {})
    auth = msg.get("auth")

    if not request_id or not action:
        return build_response(
            request_id or "unknown",
            status="error",
            error="Поля 'id' и 'action' обязательны",
        )

    cached = check_idempotency(request_id)
    if cached is not None:
        log.info("Идемпотентный дубликат: %s", request_id)
        return cached

    user_id = authenticate(auth)
    if user_id is None:
        return build_response(request_id, "error", error="Неверный или отсутствующий API-ключ")

    try:
        result = dispatch(version, action, data, user_id)
        response = build_response(request_id, "ok", data=result)
    except ValueError as e:
        response = build_response(request_id, "error", error=str(e))
    except Exception as e:
        log.exception("Необработанная ошибка при action=%s", action)
        response = build_response(request_id, "error", error=f"Внутренняя ошибка: {e}")

    save_idempotency(request_id, response)
    return response


def on_message(channel, method, properties, body):
    log.info("Получено сообщение: %s", body[:200])

    headers = properties.headers or {}
    retry_count = headers.get("x-retry-count", 0)

    try:
        response = process_message(body)
    except Exception:
        if retry_count < MAX_RETRIES:
            log.warning("Retry %d/%d для сообщения", retry_count + 1, MAX_RETRIES)
            channel.basic_publish(
                exchange="",
                routing_key=REQUEST_QUEUE,
                body=body,
                properties=pika.BasicProperties(
                    headers={"x-retry-count": retry_count + 1},
                    reply_to=properties.reply_to,
                    correlation_id=properties.correlation_id,
                ),
            )
            channel.basic_ack(delivery_tag=method.delivery_tag)
            return
        else:
            log.error("Сообщение отправлено в DLQ после %d попыток", MAX_RETRIES)
            channel.basic_nack(delivery_tag=method.delivery_tag, requeue=False)
            return

    if properties.reply_to:
        channel.basic_publish(
            exchange="",
            routing_key=properties.reply_to,
            body=json.dumps(response, ensure_ascii=False),
            properties=pika.BasicProperties(
                correlation_id=properties.correlation_id,
            ),
        )

    channel.basic_ack(delivery_tag=method.delivery_tag)
    log.info("Ответ отправлен: correlation_id=%s, status=%s",
             response.get("correlation_id"), response.get("status"))


def main():
    init_db()
    seed_admin()

    log.info("Подключение к RabbitMQ...")
    connection = pika.BlockingConnection(
        pika.ConnectionParameters(host=RABBITMQ_HOST)
    )
    channel = connection.channel()

    channel.exchange_declare(exchange=DLX_EXCHANGE, exchange_type="fanout", durable=True)
    channel.queue_declare(queue=DLQ_QUEUE, durable=True)
    channel.queue_bind(queue=DLQ_QUEUE, exchange=DLX_EXCHANGE)

    channel.queue_declare(
        queue=REQUEST_QUEUE,
        durable=True,
        arguments={
            "x-dead-letter-exchange": DLX_EXCHANGE,
        },
    )

    channel.basic_qos(prefetch_count=1)
    channel.basic_consume(queue=REQUEST_QUEUE, on_message_callback=on_message)

    log.info("Сервер запущен. Ожидание сообщений в '%s'...", REQUEST_QUEUE)
    try:
        channel.start_consuming()
    except KeyboardInterrupt:
        log.info("Остановка сервера...")
        channel.stop_consuming()
    finally:
        connection.close()


if __name__ == "__main__":
    main()