token-storage-refactored

auth/tokens/oauth.py

@@ -0,0 +1,157 @@
+"""
+Управление OAuth токенов
+"""
+
+import json
+import time
+from typing import Optional
+
+from services.redis import redis as redis_adapter
+from utils.logger import root_logger as logger
+
+from .base import BaseTokenManager
+from .types import DEFAULT_TTL, TokenData, TokenType
+
+
+class OAuthTokenManager(BaseTokenManager):
+    """
+    Менеджер OAuth токенов
+    """
+
+    async def store_oauth_tokens(
+        self,
+        user_id: str,
+        provider: str,
+        access_token: str,
+        refresh_token: Optional[str] = None,
+        expires_in: Optional[int] = None,
+        additional_data: Optional[TokenData] = None,
+    ) -> bool:
+        """Сохраняет OAuth токены"""
+        try:
+            # Сохраняем access token
+            access_data = {
+                "token": access_token,
+                "provider": provider,
+                "expires_in": expires_in,
+                **(additional_data or {}),
+            }
+
+            access_ttl = expires_in if expires_in else DEFAULT_TTL["oauth_access"]
+            await self._create_oauth_token(user_id, access_data, access_ttl, provider, "oauth_access")
+
+            # Сохраняем refresh token если есть
+            if refresh_token:
+                refresh_data = {
+                    "token": refresh_token,
+                    "provider": provider,
+                }
+                await self._create_oauth_token(
+                    user_id, refresh_data, DEFAULT_TTL["oauth_refresh"], provider, "oauth_refresh"
+                )
+
+            return True
+
+        except Exception as e:
+            logger.error(f"Ошибка сохранения OAuth токенов: {e}")
+            return False
+
+    async def _create_oauth_token(
+        self, user_id: str, token_data: TokenData, ttl: int, provider: str, token_type: TokenType
+    ) -> str:
+        """Оптимизированное создание OAuth токена"""
+        if not provider:
+            error_msg = "OAuth токены требуют указания провайдера"
+            raise ValueError(error_msg)
+
+        identifier = f"{user_id}:{provider}"
+        token_key = self._make_token_key(token_type, identifier)
+
+        # Добавляем метаданные
+        token_data.update(
+            {"user_id": user_id, "token_type": token_type, "provider": provider, "created_at": int(time.time())}
+        )
+
+        # Используем SETEX для атомарной операции
+        serialized_data = json.dumps(token_data, ensure_ascii=False)
+        await redis_adapter.execute("setex", token_key, ttl, serialized_data)
+
+        logger.info(f"Создан {token_type} токен для пользователя {user_id}, провайдер {provider}")
+        return token_key
+
+    async def get_token(self, user_id: int, provider: str, token_type: TokenType) -> Optional[TokenData]:
+        """Получает токен"""
+        if isinstance(token_type, TokenType):
+            if token_type.startswith("oauth_"):
+                return await self._get_oauth_data_optimized(token_type, str(user_id), provider)  # type: ignore[arg-type]
+            return await self._get_token_data_optimized(token_type, str(user_id), provider)  # type: ignore[arg-type]
+        return None
+
+    async def _get_oauth_data_optimized(
+        self, token_type: TokenType, user_id: str, provider: str
+    ) -> Optional[TokenData]:
+        """Оптимизированное получение OAuth данных"""
+        if not user_id or not provider:
+            error_msg = "OAuth токены требуют user_id и provider"
+            raise ValueError(error_msg)
+
+        identifier = f"{user_id}:{provider}"
+        token_key = self._make_token_key(token_type, identifier)
+
+        # Получаем данные и TTL в одном pipeline
+        async with redis_adapter.pipeline() as pipe:
+            await pipe.get(token_key)
+            await pipe.ttl(token_key)
+            results = await pipe.execute()
+
+        if results[0]:
+            token_data = json.loads(results[0])
+            if results[1] > 0:
+                token_data["ttl_remaining"] = results[1]
+            return token_data
+        return None
+
+    async def revoke_oauth_tokens(self, user_id: str, provider: str) -> bool:
+        """Удаляет все OAuth токены для провайдера"""
+        try:
+            result1 = await self._revoke_oauth_token_optimized("oauth_access", user_id, provider)
+            result2 = await self._revoke_oauth_token_optimized("oauth_refresh", user_id, provider)
+            return result1 or result2
+        except Exception as e:
+            logger.error(f"Ошибка удаления OAuth токенов: {e}")
+            return False
+
+    async def _revoke_oauth_token_optimized(self, token_type: TokenType, user_id: str, provider: str) -> bool:
+        """Оптимизированный отзыв OAuth токена"""
+        if not user_id or not provider:
+            error_msg = "OAuth токены требуют user_id и provider"
+            raise ValueError(error_msg)
+
+        identifier = f"{user_id}:{provider}"
+        token_key = self._make_token_key(token_type, identifier)
+        result = await redis_adapter.delete(token_key)
+        return result > 0
+
+    async def revoke_user_oauth_tokens(self, user_id: str, token_type: TokenType) -> int:
+        """Оптимизированный отзыв OAuth токенов пользователя используя SCAN"""
+        count = 0
+        cursor = 0
+        delete_keys = []
+        pattern = f"{token_type}:{user_id}:*"
+
+        # Используем SCAN для безопасного поиска токенов
+        while True:
+            cursor, keys = await redis_adapter.execute("scan", cursor, pattern, 100)
+
+            if keys:
+                delete_keys.extend(keys)
+                count += len(keys)
+
+            if cursor == 0:
+                break
+
+        # Удаляем найденные токены пакетно
+        if delete_keys:
+            await redis_adapter.delete(*delete_keys)
+
+        return count