task_management_api/app/utils/idempotency.py - master - tinypot/integr

task_management_api/app/utils/idempotency.py
120 строк4 KB
Tim Polus
lab1
21 янв 2026, 23:15
21 янв 2026, 23:15dd1578c
"""
Idempotency utilities
"""

import hashlib
import json
from typing import Any, Optional
from fastapi import Request, HTTPException, status

from app.utils.redis_client import get_redis_client


class IdempotencyKey:
    """Handle idempotency keys for POST operations"""

    @staticmethod
    def generate_key(request: Request, user_id: int) -> str:
        """Generate idempotency key from request data"""
        # Create a hash of the request body and user ID
        body = b""
        if hasattr(request, "_body"):
            body = request._body

        # Create unique key combining user_id, method, url and body
        key_data = {
            "user_id": user_id,
            "method": request.method,
            "url": str(request.url),
            "body": body.decode() if body else ""
        }

        # Hash the key data
        key_string = json.dumps(key_data, sort_keys=True)
        return hashlib.sha256(key_string.encode()).hexdigest()

    @staticmethod
    async def check_and_store(key: str, response_data: Any, ttl: int = 3600) -> Optional[Any]:
        """
        Check if idempotency key exists, return stored response if found.
        Otherwise store the response data.

        Returns: stored response data if key exists, None otherwise
        """
        redis_client = await get_redis_client()

        # Check if key exists
        stored_data = await redis_client.get(f"idempotency:{key}")
        if stored_data:
            return json.loads(stored_data)

        # Store response data
        await redis_client.setex(
            f"idempotency:{key}",
            ttl,
            json.dumps(response_data)
        )

        return None

    @staticmethod
    async def get_stored_response(key: str) -> Optional[Any]:
        """Get stored response for idempotency key"""
        redis_client = await get_redis_client()
        stored_data = await redis_client.get(f"idempotency:{key}")
        if stored_data:
            return json.loads(stored_data)
        return None


async def idempotency_middleware(request: Request, call_next, user_id: Optional[int] = None):
    """Middleware for idempotency handling"""
    # Only apply to POST requests
    if request.method != "POST":
        return await call_next(request)

    # Check for idempotency key in headers
    idempotency_key = request.headers.get("Idempotency-Key")
    if not idempotency_key:
        return await call_next(request)

    if not user_id:
        # Try to get user_id from request state (set by auth middleware)
        user_id = getattr(request.state, "user_id", None)
        if not user_id:
            return await call_next(request)

    # Generate internal key
    internal_key = IdempotencyKey.generate_key(request, user_id)

    # Check for stored response
    stored_response = await IdempotencyKey.get_stored_response(internal_key)
    if stored_response:
        # Return stored response
        from fastapi.responses import JSONResponse
        return JSONResponse(
            content=stored_response["data"],
            status_code=stored_response["status_code"],
            headers=stored_response.get("headers", {})
        )

    # Process the request
    response = await call_next(request)

    # Store the response for future idempotent requests
    if hasattr(response, "body"):
        try:
            response_data = json.loads(response.body.decode())
            await IdempotencyKey.check_and_store(
                internal_key,
                {
                    "data": response_data,
                    "status_code": response.status_code,
                    "headers": dict(response.headers)
                }
            )
        except (json.JSONDecodeError, UnicodeDecodeError):
            # If response is not JSON, don't store it
            pass

    return response