DefStream-Redis 0.1.3__9-py3-none-any.whl

DefStream/Broker/Implements/Redis/Configuration.py

@@ -0,0 +1,59 @@
+from enum import Enum
+from typing import Optional
+
+from DefStream.Broker import Configuration as BaseConfiguration
+from .Serializer import (
+	Decoder,
+	Encoder,
+	JavaScriptObjectNotationDecoder,
+	JavaScriptObjectNotationEncoder,
+)
+
+__all__ = (
+	'Configuration',
+	'ConnectionType',
+)
+
+
+class ConnectionType(Enum):
+	Pool = 'pool'
+	Cluster = 'cluster'
+
+
+class Configuration(BaseConfiguration):
+	def __init__(
+		self,
+		host: str = '127.0.0.1',
+		port: int = 6379,
+		username: Optional[str] = None,
+		password: Optional[str] = None,
+		secure: bool = False,
+		type: ConnectionType = ConnectionType.Pool,
+		max: int = 10,
+		database: int = 0,
+		timeout: float = 5,
+		encode: Optional[Encoder] = None,
+		decode: Optional[Decoder] = None,
+	):
+		if not isinstance(
+			type,
+			ConnectionType,
+		):
+			raise ValueError('Redis connection type must be a ConnectionType')
+		if max < 1:
+			raise ValueError('Redis max connections must be positive')
+		if database < 0 or (type == ConnectionType.Cluster and database != 0):
+			raise ValueError('Redis database must be non-negative; Cluster requires database 0')
+		if timeout <= 0:
+			raise ValueError('Redis connection timeout must be positive')
+		self.host = host
+		self.port = port
+		self.username = username
+		self.password = password
+		self.secure = secure
+		self.type = type
+		self.max = max
+		self.database = database
+		self.timeout = timeout
+		self.encode = encode if encode is not None else JavaScriptObjectNotationEncoder()
+		self.decode = decode if decode is not None else JavaScriptObjectNotationDecoder()

DefStream/Broker/Implements/Redis/Connection.py

@@ -0,0 +1,174 @@
+from typing import (
+	Any,
+	Optional,
+	Union,
+)
+
+from redis import Redis
+from redis.cluster import RedisCluster
+
+from DefStream.Broker import (
+	Connection as BaseConnection,
+	GetQueue,
+	GetStream,
+	GetConsumer,
+	Declarable,
+	TaskSender,
+	EventHandler,
+)
+from .Configuration import (
+	Configuration,
+	ConnectionType,
+)
+from .Consumer import Consumer
+from .Queue import Queue
+from .Subscription import Subscription
+
+__all__ = (
+	'Connection',
+)
+
+
+class Connection(
+	BaseConnection,
+	GetQueue,
+	GetStream,
+	GetConsumer,
+	Declarable,
+	TaskSender,
+):
+	def __init__(
+		self,
+		conf: Configuration,
+	):
+		self.conf = conf
+		self.connection: Optional[Union[Redis, RedisCluster]] = None
+
+	def connect(self) -> None:
+		if self.connection is not None:
+			return
+		parameters = {
+			'host': self.conf.host,
+			'port': self.conf.port,
+			'username': self.conf.username,
+			'password': self.conf.password,
+			'ssl': self.conf.secure,
+			'max_connections': self.conf.max,
+			'socket_connect_timeout': self.conf.timeout,
+			'socket_timeout': self.conf.timeout,
+			'decode_responses': False,
+		}
+		if self.conf.type == ConnectionType.Cluster:
+			connection = RedisCluster(**parameters)
+		else:
+			connection = Redis(
+				db=self.conf.database,
+				**parameters,
+			)
+		try:
+			connection.ping()
+		except Exception:
+			connection.close()
+			raise
+		self.connection = connection
+
+	def connected(self) -> Union[Redis, RedisCluster]:
+		if self.connection is None:
+			raise RuntimeError('Connect the Redis broker before use')
+		return self.connection
+
+	def declare(
+		self,
+		parameters: Any,
+	) -> None:
+		subscription = Subscription.parse(parameters)
+		self.create_queue(
+			subscription.queue,
+			group=subscription.group,
+		)
+
+	def send(
+		self,
+		body: Any,
+		parameters: Any,
+	) -> str:
+		subscription = Subscription.parse(parameters)
+		return self.queue(
+			subscription.queue,
+			group=subscription.group,
+		).send(body)
+
+	def queue(
+		self,
+		name: str,
+		group: str = 'DefStream',
+	) -> Queue:
+		return Queue(
+			self.connected(),
+			self.conf.encode,
+			self.conf.decode,
+			name,
+			group,
+		)
+
+	def stream(
+		self,
+		name: str,
+		group: str = 'DefStream',
+	) -> Queue:
+		return self.queue(
+			name,
+			group,
+		)
+
+	def consumer(
+		self,
+		handler: Optional[EventHandler] = None,
+		qos: int = 1,
+		reclaim_idle: Optional[int] = 60000,
+	) -> Consumer:
+		return Consumer(
+			self.connected(),
+			self.conf.encode,
+			self.conf.decode,
+			handler,
+			qos,
+			reclaim_idle,
+		)
+
+	def create_queue(
+		self,
+		name: str,
+		group: str = 'DefStream',
+	) -> None:
+		self.queue(
+			name,
+			group,
+		).create_group()
+
+	def delete_queue(
+		self,
+		name: Union[str, Queue],
+	) -> int:
+		return self.connected().delete(str(name))
+
+	def create_stream(
+		self,
+		name: str,
+		group: str = 'DefStream',
+	) -> None:
+		self.create_queue(
+			name,
+			group,
+		)
+
+	def delete_stream(
+		self,
+		name: Union[str, Queue],
+	) -> int:
+		return self.delete_queue(name)
+
+	def close(self) -> None:
+		if self.connection is not None:
+			self.connection.close()
+			self.connection = None

DefStream/Broker/Implements/Redis/Consumer.py

@@ -0,0 +1,209 @@
+from collections import deque
+from queue import (
+	Empty,
+	SimpleQueue,
+)
+from threading import Event as Wakeup
+from typing import (
+	Any,
+	Callable,
+	Dict,
+	Optional,
+	Tuple,
+	Union,
+)
+from uuid import uuid4
+
+from redis import Redis
+from redis.cluster import RedisCluster
+
+from DefStream.Broker import (
+	Consumer as BaseConsumer,
+	Event as BaseEvent,
+	EventError,
+	EventHandler,
+	Requeue,
+)
+from .Event import Event
+from .Queue import Queue
+from .Serializer import (
+	Decoder,
+	Encoder,
+)
+from .Subscription import Subscription
+
+__all__ = (
+	'Consumer',
+)
+
+
+class Consumer(BaseConsumer):
+	def __init__(
+		self,
+		connection: Union[Redis, RedisCluster],
+		encode: Encoder,
+		decode: Decoder,
+		handler: Optional[EventHandler] = None,
+		qos: int = 1,
+		reclaim_idle: Optional[int] = 60000,
+	):
+		if isinstance(
+			qos,
+			bool,
+		) or not isinstance(
+			qos,
+			int,
+		) or qos < 1:
+			raise ValueError('Redis qos must be a positive integer')
+		if reclaim_idle is not None and reclaim_idle < 1:
+			raise ValueError('Redis reclaim_idle must be positive milliseconds or None')
+		self.connection = connection
+		self.encode = encode
+		self.decode = decode
+		self.handler = handler
+		self.qos = qos
+		self.reclaim_idle = reclaim_idle
+		self.name = uuid4().hex
+		self.subscriptions: Dict[str, Tuple[Queue, EventHandler]] = {}
+		self.rotation = deque()
+		self.pending: Dict[Tuple[str, str], Event] = {}
+		self.cursors: Dict[str, str] = {}
+		self.callbacks: SimpleQueue[Callable[[], None]] = SimpleQueue()
+		self.wakeup = Wakeup()
+		self.stopped = False
+		self.closed = False
+
+	def subs(
+		self,
+		queue: Any,
+		handler: Optional[EventHandler] = None,
+	) -> None:
+		if self.closed:
+			raise RuntimeError('Redis consumer is closed')
+		subscription = Subscription.parse(queue)
+		handler = handler if handler is not None else self.handler
+		if handler is None:
+			raise ValueError('subs requires a handler or a default handler')
+		if subscription.queue in self.subscriptions:
+			raise ValueError('A Run is already subscribed to queue {}'.format(subscription.queue))
+		stream = Queue(
+			self.connection,
+			self.encode,
+			self.decode,
+			subscription.queue,
+			subscription.group,
+		)
+		stream.create_group()
+		self.subscriptions[subscription.queue] = (
+			stream,
+			handler,
+		)
+		self.cursors[subscription.queue] = '0-0'
+		self.rotation.append(subscription.queue)
+
+	def call(
+		self,
+		callback: Callable[[], None],
+	) -> None:
+		if self.closed:
+			raise RuntimeError('Redis consumer is closed')
+		self.callbacks.put(callback)
+		self.wakeup.set()
+
+	def fail(
+		self,
+		event: BaseEvent,
+		policy: Optional[EventError] = None,
+	) -> None:
+		(policy if policy is not None else Requeue())(event)
+
+	def drain(self) -> None:
+		while True:
+			try:
+				callback = self.callbacks.get_nowait()
+			except Empty:
+				break
+			callback()
+		self.pending = {
+			key: event
+			for key, event in self.pending.items()
+			if not event.settled
+		}
+
+	def deliver(
+		self,
+		stream: Queue,
+		handler: EventHandler,
+		entries: list,
+	) -> None:
+		for id, fields in entries:
+			key = (
+				stream.queue,
+				id.decode('ascii'),
+			)
+			if key in self.pending:
+				continue
+			event = stream.event(
+				id,
+				fields,
+				stream.group,
+				self.name,
+			)
+			self.pending[key] = event
+			handler(event)
+			self.drain()
+
+	def run(self) -> None:
+		if self.closed:
+			raise RuntimeError('Redis consumer is closed')
+		if not self.subscriptions:
+			raise RuntimeError('Subscribe before running the Redis consumer')
+		while not self.stopped:
+			self.wakeup.clear()
+			self.drain()
+			for _ in range(len(self.rotation)):
+				available = self.qos - len(self.pending)
+				if self.stopped or available <= 0:
+					break
+				name = self.rotation.popleft()
+				self.rotation.append(name)
+				stream, handler = self.subscriptions[name]
+				if self.reclaim_idle is not None:
+					claimed = self.connection.xautoclaim(
+						stream.queue,
+						stream.group,
+						self.name,
+						self.reclaim_idle,
+						start_id=self.cursors[stream.queue],
+						count=available,
+					)
+					self.cursors[stream.queue] = claimed[0].decode('ascii')
+					self.deliver(
+						stream,
+						handler,
+						claimed[1],
+					)
+				available = self.qos - len(self.pending)
+				if self.stopped or available <= 0:
+					continue
+				messages = self.connection.xreadgroup(
+					stream.group,
+					self.name,
+					{stream.queue: '>'},
+					count=available,
+				)
+				for _, entries in messages:
+					self.deliver(
+						stream,
+						handler,
+						entries,
+					)
+			self.wakeup.wait(0.05)
+
+	def stop(self) -> None:
+		self.stopped = True
+		self.wakeup.set()
+
+	def close(self) -> None:
+		self.stop()
+		self.closed = True

DefStream/Broker/Implements/Redis/Event.py

@@ -0,0 +1,83 @@
+from typing import (
+	Any,
+	Optional,
+	Union,
+)
+
+from redis import Redis
+from redis.cluster import RedisCluster
+
+from DefStream.Broker import Event as BaseEvent
+
+__all__ = (
+	'Event',
+)
+
+
+class Event(BaseEvent):
+	def __init__(
+		self,
+		connection: Union[Redis, RedisCluster],
+		queue: str,
+		id: str,
+		body: Any,
+		group: Optional[str] = None,
+		consumer: Optional[str] = None,
+	):
+		self.connection = connection
+		self.queue = queue
+		self.id = id
+		self.payload = body
+		self.group = group
+		self.consumer = consumer
+		self.settled = False
+
+	@property
+	def body(self) -> Any:
+		return self.payload
+
+	@body.setter
+	def body(
+		self,
+		body: Any,
+	) -> None:
+		self.payload = body
+
+	@property
+	def src(self) -> str:
+		return self.queue
+
+	def ack(self) -> None:
+		if self.group is None:
+			raise RuntimeError('Events from read() have no consumer group to acknowledge')
+		if not self.settled:
+			self.connection.xack(
+				self.queue,
+				self.group,
+				self.id,
+			)
+			self.settled = True
+
+	def nack(self) -> None:
+		if self.group is None:
+			raise RuntimeError('Events from read() have no consumer group to requeue')
+		if not self.settled:
+			# Requeue and acknowledge atomically on one stream (also valid in Cluster).
+			self.connection.eval(
+				"""
+				local pending = redis.call('XPENDING', KEYS[1], ARGV[1], ARGV[2], ARGV[2], 1)
+				if #pending == 0 then return 0 end
+				local entries = redis.call('XRANGE', KEYS[1], ARGV[2], ARGV[2])
+				if #entries == 0 then return redis.error_reply('Pending entry was deleted') end
+				redis.call('XADD', KEYS[1], '*', unpack(entries[1][2]))
+				return redis.call('XACK', KEYS[1], ARGV[1], ARGV[2])
+				""",
+				1,
+				self.queue,
+				self.group,
+				self.id,
+			)
+			self.settled = True
+
+	def reject(self) -> None:
+		self.ack()

DefStream/Broker/Implements/Redis/Queue.py

@@ -0,0 +1,151 @@
+from time import (
+	monotonic,
+	sleep,
+)
+from typing import (
+	Any,
+	Iterable,
+	Optional,
+	Union,
+)
+from uuid import uuid4
+
+from redis import Redis
+from redis.cluster import RedisCluster
+from redis.exceptions import ResponseError
+
+from DefStream.Broker import (
+	Poppable,
+	Readable,
+	Sendable,
+)
+from .Event import Event
+from .Serializer import (
+	Decoder,
+	Encoder,
+)
+from .Subscription import Subscription
+
+__all__ = (
+	'Queue',
+)
+
+
+class Queue(
+	Sendable,
+	Poppable,
+	Readable,
+):
+	def __init__(
+		self,
+		connection: Union[Redis, RedisCluster],
+		encode: Encoder,
+		decode: Decoder,
+		name: str,
+		group: str = 'DefStream',
+	):
+		self.subscription = Subscription(
+			name,
+			group,
+		)
+		self.connection = connection
+		self.encode = encode
+		self.decode = decode
+		self.queue = name
+		self.group = group
+		self.consumer = uuid4().hex
+
+	def __str__(self) -> str:
+		return self.queue
+
+	def create_group(self) -> None:
+		try:
+			self.connection.xgroup_create(
+				self.queue,
+				self.group,
+				id='0-0',
+				mkstream=True,
+			)
+		except ResponseError as error:
+			if not str(error).startswith('BUSYGROUP '):
+				raise
+
+	def send(
+		self,
+		body: Any,
+	) -> str:
+		id = self.connection.xadd(
+			self.queue,
+			{'body': self.encode(body)},
+		)
+		return id.decode('ascii')
+
+	def event(
+		self,
+		id: bytes,
+		fields: dict,
+		group: Optional[str] = None,
+		consumer: Optional[str] = None,
+	) -> Event:
+		return Event(
+			self.connection,
+			self.queue,
+			id.decode('ascii'),
+			self.decode(fields[b'body']),
+			group,
+			consumer,
+		)
+
+	def pop(
+		self,
+		timeout: Optional[int] = None,
+	) -> Optional[Event]:
+		if timeout is not None and timeout < 0:
+			raise ValueError('Redis pop timeout must be non-negative milliseconds')
+		self.create_group()
+		deadline = monotonic() + timeout / 1000 if timeout else None
+		while True:
+			messages = self.connection.xreadgroup(
+				self.group,
+				self.consumer,
+				{self.queue: '>'},
+				count=1,
+			)
+			for _, entries in messages:
+				for id, fields in entries:
+					return self.event(
+						id,
+						fields,
+						self.group,
+						self.consumer,
+					)
+			if deadline is not None and monotonic() >= deadline:
+				return None
+			sleep(0.01)
+
+	def read(
+		self,
+		offset: Union[int, str] = 0,
+		timeout: Optional[int] = None,
+	) -> Iterable[Event]:
+		if timeout is not None and timeout < 0:
+			raise ValueError('Redis read timeout must be non-negative milliseconds')
+		cursor = str(offset)
+		deadline = monotonic() + timeout / 1000 if timeout else None
+		while True:
+			if deadline is not None and monotonic() >= deadline:
+				return
+			messages = self.connection.xread(
+				{self.queue: cursor},
+				count=1,
+			)
+			if not messages:
+				sleep(0.01)
+				continue
+			for _, entries in messages:
+				for id, fields in entries:
+					cursor = id.decode('ascii')
+					yield self.event(
+						id,
+						fields,
+					)

DefStream/Broker/Implements/Redis/Serializer.py

@@ -0,0 +1,48 @@
+from json import (
+	dumps,
+	loads,
+)
+from typing import (
+	Any,
+	Protocol,
+)
+
+__all__ = (
+	'Encoder',
+	'Decoder',
+	'JavaScriptObjectNotationEncoder',
+	'JavaScriptObjectNotationDecoder',
+)
+
+
+class Encoder(Protocol):
+	def __call__(
+		self,
+		body: Any,
+	) -> bytes: ...
+
+
+class Decoder(Protocol):
+	def __call__(
+		self,
+		body: bytes,
+	) -> Any: ...
+
+
+class JavaScriptObjectNotationEncoder(Encoder):
+	def __call__(
+		self,
+		body: Any,
+	) -> bytes:
+		return dumps(
+			body,
+			ensure_ascii=False,
+		).encode('utf-8')
+
+
+class JavaScriptObjectNotationDecoder(Decoder):
+	def __call__(
+		self,
+		body: bytes,
+	) -> Any:
+		return loads(body)