Thanks to visit codestin.com
Credit goes to github.com

Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .github/workflows/integration.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ jobs:
GHSA-w596-4wvx-j9j6 # subversion related git pull, dependency for pytest. There is no impact here.
CVE-2026-26007 # dependency for entraid tests
CVE-2026-32597 # PyJWT does not validate the crit (Critical) Header Parameter defined in RFC 7515, this will be fixed in the next release
CVE-2026-4539 # pygments ReDoS in AdlLexer, local access only, no fix available yet

lint:
name: Code linters
Expand Down
65 changes: 64 additions & 1 deletion redis/asyncio/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -981,6 +981,8 @@ def __init__(
self.pending_unsubscribe_channels = set()
self.patterns = {}
self.pending_unsubscribe_patterns = set()
self.shard_channels = {}
self.pending_unsubscribe_shard_channels = set()
self._lock = asyncio.Lock()

async def __aenter__(self):
Expand Down Expand Up @@ -1009,6 +1011,8 @@ async def aclose(self):
self.pending_unsubscribe_channels = set()
self.patterns = {}
self.pending_unsubscribe_patterns = set()
self.shard_channels = {}
self.pending_unsubscribe_shard_channels = set()

@deprecated_function(version="5.0.1", reason="Use aclose() instead", name="close")
async def close(self) -> None:
Expand All @@ -1033,6 +1037,7 @@ async def on_connect(self, connection: Connection):
# that no decoding is required.
self.pending_unsubscribe_channels.clear()
self.pending_unsubscribe_patterns.clear()
self.pending_unsubscribe_shard_channels.clear()
if self.channels:
channels_with_handlers = {}
channels_without_handlers = []
Expand All @@ -1057,11 +1062,21 @@ async def on_connect(self, connection: Connection):
await self.psubscribe(
*patterns_without_handlers, **patterns_with_handlers
)
if self.shard_channels:
shard_with_handlers = {}
shard_without_handlers = []
for k, v in self.shard_channels.items():
if v is not None:
shard_with_handlers[self.encoder.decode(k, force=True)] = v
else:
shard_without_handlers.append(k)
if shard_with_handlers or shard_without_handlers:
await self.ssubscribe(*shard_without_handlers, **shard_with_handlers)

@property
def subscribed(self):
"""Indicates if there are subscriptions to any channels or patterns"""
return bool(self.channels or self.patterns)
return bool(self.channels or self.patterns or self.shard_channels)

async def execute_command(self, *args: EncodableT):
"""Execute a publish/subscribe command"""
Expand Down Expand Up @@ -1341,6 +1356,40 @@ def unsubscribe(self, *args) -> Awaitable:
self.pending_unsubscribe_channels.update(channels)
return self.execute_command("UNSUBSCRIBE", *parsed_args)

async def ssubscribe(self, *args, target_node=None, **kwargs):
"""
Subscribes the client to the specified shard channels.
Channels supplied as keyword arguments expect a channel name as the key
and a callable as the value. A channel's callable will be invoked automatically
when a message is received on that channel rather than producing a message via
``listen()`` or ``get_sharded_message()``.
"""
if args:
args = list_or_args(args[0], args[1:])
new_s_channels = dict.fromkeys(args)
new_s_channels.update(kwargs)
ret_val = await self.execute_command("SSUBSCRIBE", *new_s_channels.keys())
# update the s_channels dict AFTER we send the command. we don't want to
# subscribe twice to these channels, once for the command and again
# for the reconnection.
new_s_channels = self._normalize_keys(new_s_channels)
self.shard_channels.update(new_s_channels)
self.pending_unsubscribe_shard_channels.difference_update(new_s_channels)
return ret_val
Comment thread
petyaslavova marked this conversation as resolved.

def sunsubscribe(self, *args, target_node=None) -> Awaitable:
"""
Unsubscribe from the supplied shard_channels. If empty, unsubscribe from
all shard_channels
"""
if args:
args = list_or_args(args[0], args[1:])
s_channels = self._normalize_keys(dict.fromkeys(args))
else:
s_channels = self.shard_channels
self.pending_unsubscribe_shard_channels.update(s_channels)
return self.execute_command("SUNSUBSCRIBE", *args)

async def listen(self) -> AsyncIterator:
"""Listen for messages on channels this client has been subscribed to"""
while self.subscribed:
Expand Down Expand Up @@ -1412,6 +1461,13 @@ async def handle_message(self, response, ignore_subscribe_messages=False):
direction=PubSubDirection.RECEIVE,
channel=channel,
)
elif message_type == "smessage":
channel = str_if_bytes(message["channel"])
await record_pubsub_message(
direction=PubSubDirection.RECEIVE,
channel=channel,
sharded=True,
)

# if this is an unsubscribe message, remove it from memory
if message_type in self.UNSUBSCRIBE_MESSAGE_TYPES:
Expand All @@ -1420,6 +1476,11 @@ async def handle_message(self, response, ignore_subscribe_messages=False):
if pattern in self.pending_unsubscribe_patterns:
self.pending_unsubscribe_patterns.remove(pattern)
self.patterns.pop(pattern, None)
elif message_type == "sunsubscribe":
s_channel = response[1]
if s_channel in self.pending_unsubscribe_shard_channels:
self.pending_unsubscribe_shard_channels.remove(s_channel)
self.shard_channels.pop(s_channel, None)
else:
channel = response[1]
if channel in self.pending_unsubscribe_channels:
Expand All @@ -1430,6 +1491,8 @@ async def handle_message(self, response, ignore_subscribe_messages=False):
# if there's a message handler, invoke it
if message_type == "pmessage":
handler = self.patterns.get(message["pattern"], None)
elif message_type == "smessage":
handler = self.shard_channels.get(message["channel"], None)
else:
handler = self.channels.get(message["channel"], None)
if handler:
Expand Down
Loading
Loading