vs-queue 0.1.0__py3-none-any.whl

vs_queue/__init__.py

@@ -0,0 +1,6 @@
+from vs_queue.registry.vs_queue_registry import VsQueueRegistry
+from vs_queue.provider.vs_redis_queue import VsRedisQueue
+from vs_queue.provider.vs_rabbitmq_queue import VsRabbitMQQueue
+
+VsQueueRegistry.register("redis", VsRedisQueue)
+VsQueueRegistry.register("rabbitmq", VsRabbitMQQueue)

vs_queue/base/vs_base_consumer.py

@@ -0,0 +1,25 @@
+from abc import ABC, abstractmethod
+from typing import Optional
+
+from vs_queue.schema.vs_message import VsMessage
+
+
+class VsBaseConsumer(ABC):
+
+    def __init__(self):
+        self.uuid: Optional[str] = None
+
+    @abstractmethod
+    async def handle(self, message: VsMessage) -> None:
+        pass
+
+    @abstractmethod
+    async def start(self) -> None:
+        pass
+
+    @abstractmethod
+    async def stop(self) -> None:
+        pass
+
+    async def on_error(self, message: VsMessage, error: Exception) -> None:
+        pass

vs_queue/base/vs_base_queue.py

@@ -0,0 +1,38 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Optional
+
+from vs_queue.schema.vs_message import VsMessage
+
+
+class VsBaseQueue(ABC):
+
+    def __init__(self, host: str, port: int, credentials: Optional[Dict[str, Any]] = None):
+        self._host = host
+        self._port = port
+        self._credentials = credentials or {}
+
+    @abstractmethod
+    async def connect(self) -> None:
+        pass
+
+    @abstractmethod
+    async def disconnect(self) -> None:
+        pass
+
+    @abstractmethod
+    async def health_check(self) -> bool:
+        pass
+
+    @abstractmethod
+    async def publish(self, queue: str, message: VsMessage) -> None:
+        pass
+
+    @abstractmethod
+    async def subscribe(self, queue: str, consumer: "VsBaseConsumer") -> None:
+        pass
+
+    async def blpop(self, key: str, timeout: int = 30) -> Optional[str]:
+        raise NotImplementedError(f"{type(self).__name__} does not support blpop")
+
+
+from vs_queue.base.vs_base_consumer import VsBaseConsumer

vs_queue/decorator/vs_queue_listener.py

@@ -0,0 +1,38 @@
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Type
+
+
+@dataclass
+class VsListenerMeta:
+    queue: str
+    concurrency: int
+    max_retries: int
+    retry_backoff_seconds: float
+    dead_letter_queue: Optional[str]
+
+
+_listener_registry: List[tuple] = []
+
+
+def queue_listener(
+    queue: str,
+    concurrency: int = 1,
+    max_retries: int = 3,
+    retry_backoff_seconds: float = 1.0,
+    dead_letter_queue: Optional[str] = None,
+):
+    def decorator(cls):
+        meta = VsListenerMeta(
+            queue=queue,
+            concurrency=concurrency,
+            max_retries=max_retries,
+            retry_backoff_seconds=retry_backoff_seconds,
+            dead_letter_queue=dead_letter_queue,
+        )
+        _listener_registry.append((cls, meta))
+        return cls
+    return decorator
+
+
+def get_registered_listeners() -> List[tuple]:
+    return list(_listener_registry)

vs_queue/factory/vs_queue_factory.py

@@ -0,0 +1,18 @@
+from typing import Any, Dict, Optional
+
+from vs_queue.base.vs_base_queue import VsBaseQueue
+from vs_queue.registry.vs_queue_registry import VsQueueRegistry
+
+
+class VsQueueFactory:
+
+    @classmethod
+    def create(
+        cls,
+        provider: str,
+        host: str,
+        port: int,
+        credentials: Optional[Dict[str, Any]] = None,
+    ) -> VsBaseQueue:
+        queue_class = VsQueueRegistry.get(provider)
+        return queue_class(host=host, port=port, credentials=credentials)

vs_queue/manager/vs_queue_manager.py

@@ -0,0 +1,116 @@
+import asyncio
+import logging
+from typing import Dict, List, Optional, Tuple
+from uuid import uuid4
+
+from vs_queue.base.vs_base_consumer import VsBaseConsumer
+from vs_queue.base.vs_base_queue import VsBaseQueue
+
+_logger = logging.getLogger(__name__)
+
+
+class VsQueueManager:
+
+    def __init__(self, queue: VsBaseQueue, max_consumer_restarts: int = 3):
+        self._queue = queue
+        self._max_consumer_restarts = max_consumer_restarts
+        self._consumers: Dict[str, Tuple[str, VsBaseConsumer]] = {}  # uuid -> (queue_name, consumer)
+        self._tasks: Dict[str, asyncio.Task] = {}  # uuid -> task
+        self._connected = False
+        self._listeners_registered = False
+
+    async def register(self, queue_name: str, consumer: VsBaseConsumer) -> str:
+        if not self._connected:
+            await self._queue.connect()
+            self._connected = True
+
+        uid = uuid4().hex
+        consumer.uuid = uid
+        self._consumers[uid] = (queue_name, consumer)
+
+        task = asyncio.create_task(
+            self._run(queue_name, consumer),
+            name=f"consumer:{queue_name}:{uid}",
+        )
+        self._tasks[uid] = task
+        _logger.info(f"Listener registered | queue={queue_name} uuid={uid}")
+        return uid
+
+    async def register_listeners(self) -> None:
+        if self._listeners_registered:
+            _logger.warning("register_listeners() already called, skipping")
+            return
+        self._listeners_registered = True
+        from vs_queue.decorator.vs_queue_listener import get_registered_listeners
+        for consumer_class, meta in get_registered_listeners():
+            for _ in range(meta.concurrency):
+                await self.register(meta.queue, consumer_class())
+            _logger.info(f"Listeners registered | queue={meta.queue} concurrency={meta.concurrency}")
+
+    def get_listeners(self) -> Dict[str, List[VsBaseConsumer]]:
+        result: Dict[str, List[VsBaseConsumer]] = {}
+        for uid, (queue_name, consumer) in self._consumers.items():
+            result.setdefault(queue_name, []).append(consumer)
+        return result
+
+    async def stop_listener(self, queue_name: str, uuid: Optional[str] = None) -> None:
+        if uuid and uuid in self._consumers:
+            actual_queue, _ = self._consumers[uuid]
+            if actual_queue != queue_name:
+                raise ValueError(f"UUID '{uuid}' belongs to queue '{actual_queue}', not '{queue_name}'")
+        targets = {
+            uid: consumer
+            for uid, (qname, consumer) in self._consumers.items()
+            if qname == queue_name and (uuid is None or uid == uuid)
+        }
+        if not targets:
+            _logger.warning(f"No running listener found | queue={queue_name} uuid={uuid}")
+            return
+        for uid in targets:
+            await self._cancel(uid)
+
+    async def stop_all(self) -> None:
+        for uid in list(self._tasks.keys()):
+            await self._cancel(uid)
+        await self._queue.disconnect()
+        self._connected = False
+        _logger.info("VsQueueManager stopped")
+
+    async def health_check(self) -> bool:
+        if not self._connected:
+            return False
+        return await self._queue.health_check()
+
+    async def _cancel(self, uid: str) -> None:
+        task = self._tasks.pop(uid, None)
+        if task:
+            task.cancel()
+            try:
+                await task
+            except asyncio.CancelledError:
+                pass
+        queue_name, _ = self._consumers.pop(uid, (None, None))
+        _logger.info(f"Listener stopped | queue={queue_name} uuid={uid}")
+
+    async def _run(self, queue_name: str, consumer: VsBaseConsumer) -> None:
+        attempts = 0
+        while True:
+            try:
+                await self._queue.subscribe(queue_name, consumer)
+                break
+            except asyncio.CancelledError:
+                await consumer.stop()
+                break
+            except Exception as e:
+                attempts += 1
+                if attempts > self._max_consumer_restarts:
+                    _logger.critical(
+                        f"Consumer permanently failed after {attempts} restarts | queue={queue_name} error={e}",
+                        exc_info=True,
+                    )
+                    break
+                backoff = 2 ** (attempts - 1)
+                _logger.error(
+                    f"Consumer crashed, restarting in {backoff}s | queue={queue_name} attempt={attempts} error={e}"
+                )
+                await asyncio.sleep(backoff)

vs_queue/provider/vs_rabbitmq_queue.py

@@ -0,0 +1,91 @@
+import json
+import logging
+from typing import Any, Dict, Optional
+
+from vs_queue.base.vs_base_consumer import VsBaseConsumer
+from vs_queue.base.vs_base_queue import VsBaseQueue
+from vs_queue.retry.vs_retry_policy import VsRetryPolicy
+from vs_queue.schema.vs_message import VsMessage
+
+_logger = logging.getLogger(__name__)
+
+
+class VsRabbitMQQueue(VsBaseQueue):
+
+    def __init__(self, host: str, port: int, credentials: Optional[Dict[str, Any]] = None):
+        super().__init__(host, port, credentials)
+        self._connection = None
+        self._channel = None
+
+    async def connect(self) -> None:
+        try:
+            import aio_pika
+            username = self._credentials.get("username", "guest")
+            password = self._credentials.get("password", "guest")
+            vhost = self._credentials.get("vhost", "/")
+            self._connection = await aio_pika.connect_robust(
+                host=self._host,
+                port=self._port,
+                login=username,
+                password=password,
+                virtualhost=vhost,
+            )
+            self._channel = await self._connection.channel()
+            _logger.info(f"RabbitMQ connected | host={self._host}:{self._port}")
+        except ImportError:
+            raise ImportError("Install aio-pika: pip install aio-pika")
+
+    async def disconnect(self) -> None:
+        if self._connection:
+            await self._connection.close()
+            self._connection = None
+            self._channel = None
+
+    async def health_check(self) -> bool:
+        return self._connection is not None and not self._connection.is_closed
+
+    async def publish(self, queue: str, message: VsMessage) -> None:
+        import aio_pika
+        await self._channel.declare_queue(queue, durable=True)
+        body = json.dumps({
+            "id": message.id,
+            "timestamp": message.timestamp.isoformat(),
+            "retry_count": message.retry_count,
+            "headers": message.headers,
+            "payload": message.payload or {},
+        }).encode()
+        await self._channel.default_exchange.publish(
+            aio_pika.Message(body=body, delivery_mode=aio_pika.DeliveryMode.PERSISTENT),
+            routing_key=queue,
+        )
+
+    async def subscribe(self, queue: str, consumer: VsBaseConsumer, retry_policy: Optional[VsRetryPolicy] = None) -> None:
+        q = await self._channel.declare_queue(queue, durable=True)
+        await consumer.start()
+        try:
+            async with q.iterator() as iterator:
+                async for amqp_message in iterator:
+                    data = json.loads(amqp_message.body)
+                    message = VsMessage(
+                        id=data["id"],
+                        timestamp=data["timestamp"],
+                        retry_count=data["retry_count"],
+                        headers=data["headers"],
+                        payload=data["payload"],
+                    )
+                    try:
+                        await consumer.handle(message)
+                        await amqp_message.ack()
+                    except Exception as e:
+                        await consumer.on_error(message, e)
+                        if retry_policy:
+                            retried = await retry_policy.handle_failure(message, e, self)
+                            if retried:
+                                await amqp_message.ack()
+                            else:
+                                await amqp_message.nack(requeue=False)
+                        else:
+                            await amqp_message.nack(requeue=False)
+                            _logger.error(f"Message failed, no retry policy | id={message.id} error={e}")
+        finally:
+            await consumer.stop()

vs_queue/provider/vs_redis_queue.py

@@ -0,0 +1,104 @@
+import asyncio
+import json
+import logging
+from typing import Any, Dict, Optional
+
+import redis.asyncio as aioredis
+
+from vs_queue.base.vs_base_consumer import VsBaseConsumer
+from vs_queue.base.vs_base_queue import VsBaseQueue
+from vs_queue.retry.vs_retry_policy import VsRetryPolicy
+from vs_queue.schema.vs_message import VsMessage
+
+_logger = logging.getLogger(__name__)
+
+
+class VsRedisQueue(VsBaseQueue):
+
+    def __init__(self, host: str, port: int, credentials: Optional[Dict[str, Any]] = None):
+        super().__init__(host, port, credentials)
+        self._client: Optional[aioredis.Redis] = None
+
+    async def connect(self) -> None:
+        password = self._credentials.get("password")
+        self._client = aioredis.Redis(
+            host=self._host,
+            port=self._port,
+            password=password,
+            decode_responses=True,
+        )
+        await self._client.ping()
+
+    async def disconnect(self) -> None:
+        if self._client:
+            await self._client.aclose()
+            self._client = None
+
+    async def health_check(self) -> bool:
+        try:
+            await self._client.ping()
+            return True
+        except Exception:
+            return False
+
+    async def publish(self, queue: str, message: VsMessage) -> None:
+        data = {
+            "id": message.id,
+            "timestamp": message.timestamp.isoformat(),
+            "retry_count": str(message.retry_count),
+            "headers": json.dumps(message.headers),
+            "payload": json.dumps(message.payload or {}),
+        }
+        await self._client.xadd(queue, data)
+
+    async def blpop(self, key: str, timeout: int = 30) -> Optional[str]:
+        result = await self._client.blpop(key, timeout=timeout)
+        if result is None:
+            return None
+        _, value = result
+        return value
+
+    async def subscribe(self, queue: str, consumer: VsBaseConsumer, retry_policy: Optional[VsRetryPolicy] = None) -> None:
+        group = f"{queue}:group"
+        consumer_name = f"{queue}:consumer"
+
+        try:
+            await self._client.xgroup_create(queue, group, id="0", mkstream=True)
+        except Exception:
+            pass
+
+        await consumer.start()
+        try:
+            while True:
+                results = await self._client.xreadgroup(
+                    groupname=group,
+                    consumername=consumer_name,
+                    streams={queue: ">"},
+                    count=1,
+                    block=5000,
+                )
+                if not results:
+                    continue
+
+                for _, messages in results:
+                    for msg_id, data in messages:
+                        message = VsMessage(
+                            id=data["id"],
+                            timestamp=data["timestamp"],
+                            retry_count=int(data["retry_count"]),
+                            headers=json.loads(data["headers"]),
+                            payload=json.loads(data["payload"]),
+                        )
+                        try:
+                            await consumer.handle(message)
+                            await self._client.xack(queue, group, msg_id)
+                        except Exception as e:
+                            await consumer.on_error(message, e)
+                            if retry_policy:
+                                retried = await retry_policy.handle_failure(message, e, self)
+                                if retried:
+                                    await self._client.xack(queue, group, msg_id)
+                            else:
+                                _logger.error(f"Message failed, no retry policy | id={message.id} error={e}")
+        finally:
+            await consumer.stop()

vs_queue/registry/vs_queue_registry.py

@@ -0,0 +1,24 @@
+from typing import Dict, Type
+
+from vs_queue.base.vs_base_queue import VsBaseQueue
+
+
+class VsQueueRegistry:
+
+    _registry: Dict[str, Type[VsBaseQueue]] = {}
+
+    @classmethod
+    def register(cls, name: str, queue_class: Type[VsBaseQueue]) -> None:
+        if not issubclass(queue_class, VsBaseQueue):
+            raise TypeError(f"'{queue_class.__name__}' must extend VsBaseQueue")
+        cls._registry[name] = queue_class
+
+    @classmethod
+    def get(cls, name: str) -> Type[VsBaseQueue]:
+        if name not in cls._registry:
+            raise KeyError(f"Queue provider '{name}' is not registered. Available: {list(cls._registry.keys())}")
+        return cls._registry[name]
+
+    @classmethod
+    def available(cls) -> list:
+        return list(cls._registry.keys())

vs_queue/retry/vs_retry_policy.py

@@ -0,0 +1,30 @@
+import asyncio
+import logging
+from typing import Optional
+
+from vs_queue.schema.vs_message import VsMessage
+
+_logger = logging.getLogger(__name__)
+
+
+class VsRetryPolicy:
+
+    def __init__(self, max_retries: int = 3, backoff_seconds: float = 1.0, dead_letter_queue: Optional[str] = None):
+        self.max_retries = max_retries
+        self.backoff_seconds = backoff_seconds
+        self.dead_letter_queue = dead_letter_queue
+
+    async def handle_failure(self, message: VsMessage, error: Exception, queue) -> bool:
+        if message.retry_count < self.max_retries:
+            message.retry_count += 1
+            delay = self.backoff_seconds * (2 ** (message.retry_count - 1))
+            _logger.warning(f"Retrying message | id={message.id} attempt={message.retry_count} delay={delay}s error={error}")
+            await asyncio.sleep(delay)
+            await queue.publish(message.headers.get("source_queue", ""), message)
+            return True
+
+        _logger.error(f"Message exceeded max retries | id={message.id} retries={message.retry_count}")
+        if self.dead_letter_queue:
+            await queue.publish(self.dead_letter_queue, message)
+            _logger.info(f"Message sent to DLQ | id={message.id} dlq={self.dead_letter_queue}")
+        return False

vs_queue/schema/vs_message.py

@@ -0,0 +1,13 @@
+import uuid
+from datetime import datetime, timezone
+from typing import Any, Dict, Optional
+
+from pydantic import BaseModel, Field
+
+
+class VsMessage(BaseModel):
+    id: str = Field(default_factory=lambda: uuid.uuid4().hex)
+    timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
+    retry_count: int = 0
+    headers: Dict[str, str] = Field(default_factory=dict)
+    payload: Optional[Dict[str, Any]] = None