panther 3.9.0__py3-none-any.whl → 4.0.0__py3-none-any.whl

panther/base_request.py

@@ -1,6 +1,9 @@
 from collections import namedtuple
 from collections.abc import Callable
 
+from panther.db import Model
+from panther.exceptions import InvalidPathVariableAPIError
+
 
 class Headers:
     accept: str
@@ -31,14 +34,12 @@ class Headers:
     def __getattr__(self, item: str):
         if result := self.__pythonic_headers.get(item):
             return result
-        else:
-            return self.__headers.get(item)
+        return self.__headers.get(item)
 
     def __getitem__(self, item: str):
         if result := self.__headers.get(item):
             return result
-        else:
-            return self.__pythonic_headers.get(item)
+        return self.__pythonic_headers.get(item)
 
     def __str__(self):
         items = ', '.join(f'{k}={v}' for k, v in self.__headers.items())
@@ -51,7 +52,7 @@ class Headers:
         return self.__headers
 
 
-Address = namedtuple('Client', ['ip', 'port'])
+Address = namedtuple('Address', ['ip', 'port'])
 
 
 class BaseRequest:
@@ -59,11 +60,10 @@ class BaseRequest:
         self.scope = scope
         self.asgi_send = send
         self.asgi_receive = receive
-        self._data = ...
-        self._validated_data = None
-        self._user = None
         self._headers: Headers | None = None
         self._params: dict | None = None
+        self.user: Model | None = None
+        self.path_variables: dict | None = None
 
     @property
     def headers(self) -> Headers:
@@ -103,9 +103,33 @@ class BaseRequest:
     def scheme(self) -> str:
         return self.scope['scheme']
 
-    @property
-    def user(self):
-        return self._user
-
-    def set_user(self, user) -> None:
-        self._user = user
+    def collect_path_variables(self, found_path: str):
+        self.path_variables = {
+            variable.strip('< >'): value
+            for variable, value in zip(
+                found_path.strip('/').split('/'),
+                self.path.strip('/').split('/')
+            )
+            if variable.startswith('<')
+        }
+
+    def clean_parameters(self, func: Callable) -> dict:
+        kwargs = self.path_variables.copy()
+
+        for variable_name, variable_type in func.__annotations__.items():
+            # Put Request/ Websocket In kwargs (If User Wants It)
+            if issubclass(variable_type, BaseRequest):
+                kwargs[variable_name] = self
+
+            elif variable_name in kwargs:
+                # Cast To Boolean
+                if variable_type is bool:
+                    kwargs[variable_name] = kwargs[variable_name].lower() not in ['false', '0']
+
+                # Cast To Int
+                elif variable_type is int:
+                    try:
+                        kwargs[variable_name] = int(kwargs[variable_name])
+                    except ValueError:
+                        raise InvalidPathVariableAPIError(value=kwargs[variable_name], variable_type=variable_type)
+        return kwargs

panther/base_websocket.py

@@ -1,28 +1,28 @@
 from __future__ import annotations
 
 import asyncio
-import contextlib
 import logging
-from multiprocessing import Manager
+from multiprocessing.managers import SyncManager
 from typing import TYPE_CHECKING, Literal
 
 import orjson as json
 
 from panther import status
-from panther._utils import generate_ws_connection_id
 from panther.base_request import BaseRequest
 from panther.configs import config
-from panther.db.connection import redis
-from panther.utils import Singleton
+from panther.db.connections import redis
+from panther.exceptions import AuthenticationAPIError, InvalidPathVariableAPIError
+from panther.monitoring import Monitoring
+from panther.utils import Singleton, ULID
 
 if TYPE_CHECKING:
-    from redis import Redis
+    from redis.asyncio import Redis
 
 logger = logging.getLogger('panther')
 
 
 class PubSub:
-    def __init__(self, manager):
+    def __init__(self, manager: SyncManager):
         self._manager = manager
         self._subscribers = self._manager.list()
 
@@ -37,17 +37,29 @@ class PubSub:
 
 
 class WebsocketConnections(Singleton):
-    def __init__(self, manager: Manager = None):
+    def __init__(self, pubsub_connection: Redis | SyncManager):
         self.connections = {}
         self.connections_count = 0
-        self.manager = manager
+        self.pubsub_connection = pubsub_connection
 
-    def __call__(self, r: Redis | None):
-        if r:
-            subscriber = r.pubsub()
-            subscriber.subscribe('websocket_connections')
+        if isinstance(self.pubsub_connection, SyncManager):
+            self.pubsub = PubSub(manager=self.pubsub_connection)
+
+    async def __call__(self):
+        if isinstance(self.pubsub_connection, SyncManager):
+            # We don't have redis connection, so use the `multiprocessing.Manager`
+            self.pubsub: PubSub
+            queue = self.pubsub.subscribe()
+            logger.info("Subscribed to 'websocket_connections' queue")
+            while True:
+                received_message = queue.get()
+                await self._handle_received_message(received_message=received_message)
+        else:
+            # We have a redis connection, so use it for pubsub
+            self.pubsub = self.pubsub_connection.pubsub()
+            await self.pubsub.subscribe('websocket_connections')
             logger.info("Subscribed to 'websocket_connections' channel")
-            for channel_data in subscriber.listen():
+            async for channel_data in self.pubsub.listen():
                 match channel_data['type']:
                     # Subscribed
                     case 'subscribe':
@@ -56,19 +68,12 @@ class WebsocketConnections(Singleton):
                     # Message Received
                     case 'message':
                         loaded_data = json.loads(channel_data['data'].decode())
-                        self._handle_received_message(received_message=loaded_data)
+                        await self._handle_received_message(received_message=loaded_data)
 
                     case unknown_type:
-                        logger.debug(f'Unknown Channel Type: {unknown_type}')
-        else:
-            self.pubsub = PubSub(manager=self.manager)
-            queue = self.pubsub.subscribe()
-            logger.info("Subscribed to 'websocket_connections' queue")
-            while True:
-                received_message = queue.get()
-                self._handle_received_message(received_message=received_message)
+                        logger.error(f'Unknown Channel Type: {unknown_type}')
 
-    def _handle_received_message(self, received_message):
+    async def _handle_received_message(self, received_message):
         if (
                 isinstance(received_message, dict)
                 and (connection_id := received_message.get('connection_id'))
@@ -79,71 +84,156 @@ class WebsocketConnections(Singleton):
             # Check Action of WS
             match received_message['action']:
                 case 'send':
-                    asyncio.run(self.connections[connection_id].send(data=received_message['data']))
+                    await self.connections[connection_id].send(data=received_message['data'])
                 case 'close':
-                    with contextlib.suppress(RuntimeError):
-                        asyncio.run(self.connections[connection_id].close(
-                            code=received_message['data']['code'],
-                            reason=received_message['data']['reason']
-                        ))
-                        # We are trying to disconnect the connection between a thread and a user
-                        # from another thread, it's working, but we have to find another solution for it
-                        #
-                        # Error:
-                        # Task <Task pending coro=<Websocket.close()>> got Future
-                        # <Task pending coro=<WebSocketCommonProtocol.transfer_data()>>
-                        # attached to a different loop
+                    await self.connections[connection_id].close(
+                        code=received_message['data']['code'],
+                        reason=received_message['data']['reason']
+                    )
                 case unknown_action:
-                    logger.debug(f'Unknown Message Action: {unknown_action}')
+                    logger.error(f'Unknown Message Action: {unknown_action}')
 
-    def publish(self, connection_id: str, action: Literal['send', 'close'], data: any):
+    async def publish(self, connection_id: str, action: Literal['send', 'close'], data: any):
         publish_data = {'connection_id': connection_id, 'action': action, 'data': data}
 
         if redis.is_connected:
-            redis.publish('websocket_connections', json.dumps(publish_data))
+            await redis.publish('websocket_connections', json.dumps(publish_data))
         else:
             self.pubsub.publish(publish_data)
 
-    async def new_connection(self, connection: Websocket) -> None:
-        await connection.connect(**connection.path_variables)
-        if not hasattr(connection, '_connection_id'):
-            # User didn't even call the `self.accept()` so close the connection
-            await connection.close()
+    async def listen(self, connection: Websocket) -> None:
+        # 1. Authentication
+        if not connection.is_rejected:
+            await self.handle_authentication(connection=connection)
 
-        if connection.is_connected:
-            self.connections_count += 1
+        # 2. Permissions
+        if not connection.is_rejected:
+            await self.handle_permissions(connection=connection)
 
-            # Save New ConnectionID
-            self.connections[connection.connection_id] = connection
+        if connection.is_rejected:
+            # Connection is rejected so don't continue the flow ...
+            return None
 
-    def remove_connection(self, connection: Websocket) -> None:
-        if connection.is_connected:
-            self.connections_count -= 1
-            del self.connections[connection.connection_id]
+        # 3. Put PathVariables and Request(If User Wants It) In kwargs
+        try:
+            kwargs = connection.clean_parameters(connection.connect)
+        except InvalidPathVariableAPIError as e:
+            connection.log(e.detail)
+            return await connection.close()
 
+        # 4. Connect To Endpoint
+        await connection.connect(**kwargs)
 
-class Websocket(BaseRequest):
-    is_connected: bool = False
+        # 5. Check Connection
+        if not connection.is_connected and not connection.is_rejected:
+            # User didn't call the `self.accept()` or `self.close()` so we `close()` the connection (reject)
+            return await connection.close()
 
-    async def connect(self, **kwargs) -> None:
-        """Check your conditions then self.accept() the connection"""
-        await self.accept()
+        # 6. Listen Connection
+        await self.listen_connection(connection=connection)
 
-    async def accept(self, subprotocol: str | None = None, headers: dict | None = None) -> None:
-        await self.asgi_send({'type': 'websocket.accept', 'subprotocol': subprotocol, 'headers': headers or {}})
-        self.is_connected = True
+    async def listen_connection(self, connection: Websocket):
+        while True:
+            response = await connection.asgi_receive()
+            if response['type'] == 'websocket.connect':
+                continue
+
+            if response['type'] == 'websocket.disconnect':
+                # Connect has to be closed by the client
+                await self.connection_closed(connection=connection)
+                break
 
+            if 'text' in response:
+                await connection.receive(data=response['text'])
+            else:
+                await connection.receive(data=response['bytes'])
+
+    async def connection_accepted(self, connection: Websocket) -> None:
         # Generate ConnectionID
-        connection_id = generate_ws_connection_id()
-        while connection_id in config['websocket_connections'].connections:
-            connection_id = generate_ws_connection_id()
+        connection._connection_id = ULID.new()
+
+        # Save Connection
+        self.connections[connection.connection_id] = connection
 
-        # Set ConnectionID
-        self.set_connection_id(connection_id)
+        # Logs
+        await connection.monitoring.after('Accepted')
+        connection.log(f'Accepted {connection.connection_id}')
+
+    async def connection_closed(self, connection: Websocket, from_server: bool = False) -> None:
+        if connection.is_connected:
+            del self.connections[connection.connection_id]
+            await connection.monitoring.after('Closed')
+            connection.log(f'Closed {connection.connection_id}')
+            connection._connection_id = ''
+
+        elif connection.is_rejected is False and from_server is True:
+            await connection.monitoring.after('Rejected')
+            connection.log('Rejected')
+            connection._is_rejected = True
+
+    async def start(self):
+        """
+        Start Websocket Listener (Redis/ Queue)
+
+        Cause of --preload in gunicorn we have to keep this function here,
+        and we can't move it to __init__ of Panther
+
+            * Each process should start this listener for itself,
+              but they have same Manager()
+        """
+
+        if config.HAS_WS:
+            # Schedule the async function to run in the background,
+            #   We don't need to await for this task
+            asyncio.create_task(self())
+
+    @classmethod
+    async def handle_authentication(cls, connection: Websocket):
+        """Return True if connection is closed, False otherwise."""
+        if connection.auth:
+            if not config.WS_AUTHENTICATION:
+                logger.critical('`WS_AUTHENTICATION` has not been set in configs')
+                await connection.close()
+            else:
+                try:
+                    connection.user = await config.WS_AUTHENTICATION.authentication(connection)
+                except AuthenticationAPIError as e:
+                    connection.log(e.detail)
+                    await connection.close()
+
+    @classmethod
+    async def handle_permissions(cls, connection: Websocket):
+        """Return True if connection is closed, False otherwise."""
+        for perm in connection.permissions:
+            if type(perm.authorization).__name__ != 'method':
+                logger.critical(f'{perm.__name__}.authorization should be "classmethod"')
+                await connection.close()
+            elif await perm.authorization(connection) is False:
+                connection.log('Permission Denied')
+                await connection.close()
+
+
+class Websocket(BaseRequest):
+    auth: bool = False
+    permissions: list = []
+    _connection_id: str = ''
+    _is_rejected: bool = False
+    _monitoring: Monitoring
+
+    def __init_subclass__(cls, **kwargs):
+        if cls.__module__ != 'panther.websocket':
+            config.HAS_WS = True
+
+    async def connect(self, **kwargs) -> None:
+        pass
 
     async def receive(self, data: str | bytes) -> None:
         pass
 
+    async def accept(self, subprotocol: str | None = None, headers: dict | None = None) -> None:
+        await self.asgi_send({'type': 'websocket.accept', 'subprotocol': subprotocol, 'headers': headers or {}})
+        await config.WEBSOCKET_CONNECTIONS.connection_accepted(connection=self)
+
     async def send(self, data: any = None) -> None:
         logger.debug(f'Sending WS Message to {self.connection_id}')
         if data:
@@ -161,38 +251,26 @@ class Websocket(BaseRequest):
         await self.asgi_send({'type': 'websocket.send', 'bytes': bytes_data})
 
     async def close(self, code: int = status.WS_1000_NORMAL_CLOSURE, reason: str = '') -> None:
-        logger.debug(f'Closing WS Connection {self.connection_id}')
-        self.is_connected = False
-        config['websocket_connections'].remove_connection(self)
         await self.asgi_send({'type': 'websocket.close', 'code': code, 'reason': reason})
+        await config.WEBSOCKET_CONNECTIONS.connection_closed(connection=self, from_server=True)
 
-    async def listen(self) -> None:
-        while self.is_connected:
-            response = await self.asgi_receive()
-            if response['type'] == 'websocket.connect':
-                continue
-
-            if response['type'] == 'websocket.disconnect':
-                break
-
-            if 'text' in response:
-                await self.receive(data=response['text'])
-            else:
-                await self.receive(data=response['bytes'])
-
-    def set_path_variables(self, path_variables: dict) -> None:
-        self._path_variables = path_variables
+    @property
+    def connection_id(self) -> str:
+        if self.is_connected:
+            return self._connection_id
+        logger.error('You should first `self.accept()` the connection then use the `self.connection_id`')
 
     @property
-    def path_variables(self) -> dict:
-        return getattr(self, '_path_variables', {})
+    def is_connected(self) -> bool:
+        return bool(self._connection_id)
 
-    def set_connection_id(self, connection_id: str) -> None:
-        self._connection_id = connection_id
+    @property
+    def is_rejected(self) -> bool:
+        return self._is_rejected
 
     @property
-    def connection_id(self) -> str:
-        connection_id = getattr(self, '_connection_id', None)
-        if connection_id is None:
-            logger.error('You should first `self.accept()` the connection then use the `self.connection_id`')
-        return connection_id
+    def monitoring(self) -> Monitoring:
+        return self._monitoring
+
+    def log(self, message: str):
+        logger.debug(f'WS {self.path} --> {message}')

panther/caching.py

@@ -6,54 +6,56 @@ from types import NoneType
 import orjson as json
 
 from panther.configs import config
-from panther.db.connection import redis
+from panther.db.connections import redis
 from panther.request import Request
 from panther.response import Response, ResponseDataTypes
+from panther.throttling import throttling_storage
 from panther.utils import generate_hash_value_from_string, round_datetime
 
 logger = logging.getLogger('panther')
 
 caches = {}
-CachedResponse = namedtuple('Cached', ['data', 'status_code'])
+CachedResponse = namedtuple('CachedResponse', ['data', 'status_code'])
 
 
-def cache_key(request: Request, /) -> str:
+def api_cache_key(request: Request, cache_exp_time: timedelta | None = None) -> str:
     client = request.user and request.user.id or request.client.ip
     query_params_hash = generate_hash_value_from_string(request.scope['query_string'].decode('utf-8'))
-    return f'{client}-{request.path}-{query_params_hash}-{request.validated_data}'
+    key = f'{client}-{request.path}-{query_params_hash}-{request.validated_data}'
 
-
-def local_cache_key(*, request: Request, cache_exp_time: timedelta | None = None) -> str:
-    key = cache_key(request)
     if cache_exp_time:
         time = round_datetime(datetime.now(), cache_exp_time)
         return f'{time}-{key}'
-    else:
-        return key
 
+    return key
 
-def get_cached_response_data(*, request: Request, cache_exp_time: timedelta) -> CachedResponse | None:
+
+def throttling_cache_key(request: Request, duration: timedelta) -> str:
+    client = request.user and request.user.id or request.client.ip
+    time = round_datetime(datetime.now(), duration)
+    return f'{time}-{client}-{request.path}'
+
+
+async def get_response_from_cache(*, request: Request, cache_exp_time: timedelta) -> CachedResponse | None:
     """
     If redis.is_connected:
         Get Cached Data From Redis
     else:
         Get Cached Data From Memory
     """
-    if redis.is_connected:  # noqa: Unresolved References
-        key = cache_key(request)
-        data = (redis.get(key) or b'{}').decode()
+    if redis.is_connected:
+        key = api_cache_key(request=request)
+        data = (await redis.get(key) or b'{}').decode()
         if cached_value := json.loads(data):
             return CachedResponse(*cached_value)
 
     else:
-        key = local_cache_key(request=request, cache_exp_time=cache_exp_time)
+        key = api_cache_key(request=request, cache_exp_time=cache_exp_time)
         if cached_value := caches.get(key):
             return CachedResponse(*cached_value)
 
-    return None
-
 
-def set_cache_response(*, request: Request, response: Response, cache_exp_time: timedelta | int) -> None:
+async def set_response_in_cache(*, request: Request, response: Response, cache_exp_time: timedelta | int) -> None:
     """
     If redis.is_connected:
         Cache The Data In Redis
@@ -63,10 +65,10 @@ def set_cache_response(*, request: Request, response: Response, cache_exp_time:
 
     cache_data: tuple[ResponseDataTypes, int] = (response.data, response.status_code)
 
-    if redis.is_connected:  # noqa: Unresolved References
-        key = cache_key(request)
+    if redis.is_connected:
+        key = api_cache_key(request=request)
 
-        cache_exp_time = cache_exp_time or config['default_cache_exp']
+        cache_exp_time = cache_exp_time or config.DEFAULT_CACHE_EXP
         cache_data: bytes = json.dumps(cache_data)
 
         if not isinstance(cache_exp_time, timedelta | int | NoneType):
@@ -76,15 +78,48 @@ def set_cache_response(*, request: Request, response: Response, cache_exp_time:
         if cache_exp_time is None:
             logger.warning(
                 'your response are going to cache in redis forever '
-                '** set DEFAULT_CACHE_EXP in configs or pass the cache_exp_time in @API.get() for prevent this **'
+                '** set DEFAULT_CACHE_EXP in `configs` or set the `cache_exp_time` in `@API.get()` to prevent this **'
             )
-            redis.set(key, cache_data)
+            await redis.set(key, cache_data)
         else:
-            redis.set(key, cache_data, ex=cache_exp_time)
+            await redis.set(key, cache_data, ex=cache_exp_time)
 
     else:
-        key = local_cache_key(request=request, cache_exp_time=cache_exp_time)
+        key = api_cache_key(request=request, cache_exp_time=cache_exp_time)
         caches[key] = cache_data
 
         if cache_exp_time:
-            logger.info('"cache_exp_time" is not very accurate when redis is not connected.')
+            logger.info('`cache_exp_time` is not very accurate when `redis` is not connected.')
+
+
+async def get_throttling_from_cache(request: Request, duration: timedelta) -> int:
+    """
+    If redis.is_connected:
+        Get Cached Data From Redis
+    else:
+        Get Cached Data From Memory
+    """
+    key = throttling_cache_key(request=request, duration=duration)
+
+    if redis.is_connected:
+        data = (await redis.get(key) or b'0').decode()
+        return json.loads(data)
+
+    else:
+        return throttling_storage[key]
+
+
+async def increment_throttling_in_cache(request: Request, duration: timedelta) -> None:
+    """
+    If redis.is_connected:
+        Increment The Data In Redis
+    else:
+        Increment The Data In Memory
+    """
+    key = throttling_cache_key(request=request, duration=duration)
+
+    if redis.is_connected:
+        await redis.incrby(key, amount=1)
+
+    else:
+        throttling_storage[key] += 1