app/worker.py - lab4 - DisVic/task-tracker

app/worker.py
100 строк4 KB
DisappearedVictory
Final version of Lab 4
28 дек 2025, 00:11
28 дек 2025, 00:11d1ba2d7
import asyncio
import json
import logging
from aio_pika import connect, Message, IncomingMessage
from sqlalchemy.future import select
from app.db import SessionLocal
from app.models import User, Project
from app.core.config import settings

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

RABBIT_URL = "amqp://guest:guest@localhost:5672/"

async def process_message(message: IncomingMessage):
    async with message.process():
        try:
            body = json.loads(message.body.decode())
            correlation_id = message.correlation_id
            reply_to = message.reply_to
            
            logger.info(f"Received request: {body['action']} (ID: {body.get('id')})")

            # --- 1. Аутентификация---
            token = body.get("auth")
            if not token:
                raise ValueError("No auth token")
            
            # --- 2. Роутинг по action ---
            response_data = {}
            status = "ok"
            error = None

            with SessionLocal() as db:
                action = body.get("action")
                data = body.get("data", {})

                if action == "create_project":
                    user = db.query(User).filter(User.email == "user1@example.com").first()
                    if not user:
                         user = User(email="user1@example.com", password_hash="hash")
                         db.add(user)
                         db.commit()

                    proj = Project(name=data["name"], owner_id=user.id)
                    db.add(proj)
                    db.commit()
                    response_data = {"id": proj.id, "name": proj.name}

                elif action == "get_projects":
                    projects = db.query(Project).limit(10).all()
                    response_data = [{"id": p.id, "name": p.name} for p in projects]
                
                else:
                    status = "error"
                    error = f"Unknown action: {action}"

            # --- 3. Формирование ответа ---
            response_body = {
                "correlation_id": correlation_id,
                "status": status,
                "data": response_data,
                "error": error
            }

            # --- 4. Отправка ответа ---
            if reply_to:
                connection = await connect(RABBIT_URL)
                async with connection:
                    channel = await connection.channel()
                    await channel.default_exchange.publish(
                        Message(
                            body=json.dumps(response_body).encode(),
                            correlation_id=correlation_id
                        ),
                        routing_key=reply_to
                    )
                    logger.info(f"Sent response to {reply_to}")

        except Exception as e:
            logger.error(f"Error processing message: {e}")

async def main():
    connection = await connect(RABBIT_URL)
    async with connection:
        channel = await connection.channel()
        
        # Объявляем очередь запросов
        queue = await channel.declare_queue("api.requests", auto_delete=False)
        
        logger.info("Worker started. Waiting for messages...")
        
        # Слушаем очередь
        await queue.consume(process_message)
        
        # Бесконечный цикл, чтобы скрипт не завершился
        await asyncio.Future()

if __name__ == "__main__":
    asyncio.run(main())