Skip to content

Commit e1efce8

Browse files
feat(servicebus): add list_queue_sessions and list_subscription_sessions
- Add _SessionBrowser and _SessionBrowserAsync for management-only AMQP connections - Add create_mgmt_client to pyamqp and uamqp transports - Add create_mgmt_client_async to async pyamqp transport - Add list_queue_sessions/list_subscription_sessions to sync and async ServiceBusClient - Handle 404+SessionNotFound as empty result in list_sessions_op handler - Add live tests for queue sessions, subscription sessions, and updated_since mode
1 parent 3fa067a commit e1efce8

9 files changed

Lines changed: 654 additions & 0 deletions

File tree

sdk/servicebus/azure-servicebus/azure/servicebus/_common/mgmt_handlers.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,9 @@ def list_sessions_op( # pylint: disable=inconsistent-return-statements
7676
return parsed
7777
if status_code in [202, 204]:
7878
return []
79+
# 404 + SessionNotFound means no sessions exist for this entity.
80+
if status_code == 404:
81+
return []
7982

8083
amqp_transport.handle_amqp_mgmt_error(_LOGGER, "List sessions failed.", condition, description, status_code)
8184

sdk/servicebus/azure-servicebus/azure/servicebus/_servicebus_client.py

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
# --------------------------------------------------------------------------------------------
55
# pylint: disable=client-method-missing-tracing-decorator
66
from typing import Any, Union, Optional, TYPE_CHECKING, Type
7+
from datetime import datetime
78
import logging
89
import warnings
910
from weakref import WeakSet
@@ -17,6 +18,7 @@
1718
)
1819
from ._servicebus_sender import ServiceBusSender
1920
from ._servicebus_receiver import ServiceBusReceiver
21+
from ._session_browser import _SessionBrowser
2022
from ._common.auto_lock_renewer import AutoLockRenewer
2123
from ._common._configuration import Configuration
2224
from ._common.utils import (
@@ -699,3 +701,90 @@ def get_subscription_receiver(
699701
)
700702
self._handlers.add(handler)
701703
return handler
704+
705+
def _create_session_browser(self, entity_name, subscription_name=None, **kwargs):
706+
"""Create an internal _SessionBrowser for management-only operations."""
707+
browser = _SessionBrowser(
708+
fully_qualified_namespace=self.fully_qualified_namespace,
709+
entity_name=entity_name,
710+
credential=self._credential,
711+
logging_enable=self._config.logging_enable,
712+
transport_type=self._config.transport_type,
713+
http_proxy=self._config.http_proxy,
714+
connection=self._connection,
715+
user_agent=self._config.user_agent,
716+
retry_mode=self._config.retry_mode,
717+
retry_total=self._config.retry_total,
718+
retry_backoff_factor=self._config.retry_backoff_factor,
719+
retry_backoff_max=self._config.retry_backoff_max,
720+
custom_endpoint_address=self._custom_endpoint_address,
721+
connection_verify=self._connection_verify,
722+
ssl_context=self._ssl_context,
723+
amqp_transport=self._amqp_transport,
724+
use_tls=self._config.use_tls,
725+
subscription_name=subscription_name,
726+
**kwargs,
727+
)
728+
return browser
729+
730+
def list_queue_sessions(
731+
self,
732+
queue_name: str,
733+
*,
734+
updated_since: Optional["datetime"] = None,
735+
timeout: Optional[float] = None,
736+
**kwargs: Any,
737+
):
738+
# type: (...) -> list[str]
739+
"""List session IDs with active messages in a session-enabled queue.
740+
741+
If ``updated_since`` is specified, only sessions whose state was updated
742+
after that time are returned. If not specified, returns sessions with
743+
active messages in the queue.
744+
745+
:param str queue_name: The name of the session-enabled queue.
746+
:keyword ~datetime.datetime updated_since: If specified, only sessions whose state was
747+
updated after this time are returned.
748+
:keyword float timeout: The total operation timeout in seconds.
749+
:returns: A list of session ID strings.
750+
:rtype: list[str]
751+
"""
752+
if kwargs:
753+
warnings.warn(f"Unsupported keyword args: {kwargs}")
754+
browser = self._create_session_browser(queue_name)
755+
try:
756+
return browser.list_sessions(updated_since=updated_since, timeout=timeout)
757+
finally:
758+
browser.close()
759+
760+
def list_subscription_sessions(
761+
self,
762+
topic_name: str,
763+
subscription_name: str,
764+
*,
765+
updated_since: Optional["datetime"] = None,
766+
timeout: Optional[float] = None,
767+
**kwargs: Any,
768+
):
769+
# type: (...) -> list[str]
770+
"""List session IDs with active messages in a session-enabled subscription.
771+
772+
If ``updated_since`` is specified, only sessions whose state was updated
773+
after that time are returned. If not specified, returns sessions with
774+
active messages in the subscription.
775+
776+
:param str topic_name: The name of the topic.
777+
:param str subscription_name: The name of the subscription.
778+
:keyword ~datetime.datetime updated_since: If specified, only sessions whose state was
779+
updated after this time are returned.
780+
:keyword float timeout: The total operation timeout in seconds.
781+
:returns: A list of session ID strings.
782+
:rtype: list[str]
783+
"""
784+
if kwargs:
785+
warnings.warn(f"Unsupported keyword args: {kwargs}")
786+
browser = self._create_session_browser(topic_name, subscription_name=subscription_name)
787+
try:
788+
return browser.list_sessions(updated_since=updated_since, timeout=timeout)
789+
finally:
790+
browser.close()
Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
# --------------------------------------------------------------------------------------------
2+
# Copyright (c) Microsoft Corporation. All rights reserved.
3+
# Licensed under the MIT License. See License.txt in the project root for license information.
4+
# --------------------------------------------------------------------------------------------
5+
import time
6+
import logging
7+
import uuid
8+
from datetime import datetime, timezone
9+
from typing import Any, List, Optional, Union
10+
11+
from ._base_handler import BaseHandler
12+
from ._common.utils import create_authentication
13+
from ._common.constants import (
14+
REQUEST_RESPONSE_GET_MESSAGE_SESSIONS_OPERATION,
15+
)
16+
from ._common import mgmt_handlers
17+
from ._pyamqp.types import AMQPTypes, TYPE, VALUE
18+
19+
_LOGGER = logging.getLogger(__name__)
20+
21+
# The service checks for DateTime.MaxValue (C# 9999-12-31T23:59:59.9999999) to switch
22+
# between "active messages" mode and "updated since" mode. On the AMQP wire, timestamps
23+
# have millisecond precision, so DateTime.MaxValue becomes 253402300799999 ms from epoch.
24+
# Python's datetime.timestamp() float math rounds this up by 1ms, so we use the
25+
# pre-computed constant directly for the sentinel.
26+
_MAX_DATETIME_MS = 253402300799999
27+
_MAX_DATETIME = datetime(9999, 12, 31, 23, 59, 59, 999999, tzinfo=timezone.utc)
28+
_PAGE_SIZE = 100
29+
30+
31+
def _amqp_int_value(value):
32+
return {TYPE: AMQPTypes.int, VALUE: value}
33+
34+
35+
def _amqp_timestamp_value(value):
36+
"""Convert a datetime to an AMQP timestamp (milliseconds since epoch)."""
37+
return {TYPE: AMQPTypes.timestamp, VALUE: int(value.timestamp() * 1000)}
38+
39+
40+
class _SessionBrowser(BaseHandler):
41+
"""Internal handler that opens an AMQP connection for management-only operations.
42+
43+
Unlike ServiceBusSender/ServiceBusReceiver, this does NOT create a sender or
44+
receiver link. It only opens a connection and authenticates, then sends
45+
management requests to the $management endpoint.
46+
47+
Used for entity-level management operations like get-message-sessions where
48+
a sender or receiver link is not needed.
49+
"""
50+
51+
def __init__(self, fully_qualified_namespace, entity_name, credential, **kwargs):
52+
super().__init__(
53+
fully_qualified_namespace=fully_qualified_namespace,
54+
entity_name=entity_name,
55+
credential=credential,
56+
**kwargs,
57+
)
58+
self._auth_uri = f"sb://{self.fully_qualified_namespace}/{self._entity_name}"
59+
self._error_policy = self._amqp_transport.create_retry_policy(self._config)
60+
self._name = f"SBSessionBrowser-{uuid.uuid4()}"
61+
self._connection = kwargs.get("connection")
62+
63+
def _create_handler(self, auth):
64+
self._handler = self._amqp_transport.create_mgmt_client(
65+
config=self._config,
66+
auth=auth,
67+
properties=self._properties,
68+
retry_policy=self._error_policy,
69+
client_name=self._name,
70+
)
71+
72+
def _open(self):
73+
if self._running:
74+
return
75+
if self._handler:
76+
self._handler.close()
77+
78+
auth = None if self._connection else create_authentication(self)
79+
self._create_handler(auth)
80+
try:
81+
self._handler.open(connection=self._connection)
82+
while not self._handler.client_ready():
83+
time.sleep(0.05)
84+
self._running = True
85+
except:
86+
self._close_handler()
87+
raise
88+
89+
def list_sessions(
90+
self,
91+
*,
92+
updated_since: Optional[datetime] = None,
93+
timeout: Optional[float] = None,
94+
) -> List[str]:
95+
"""List session IDs for this entity.
96+
97+
:keyword ~datetime.datetime updated_since: If specified, only sessions whose state was
98+
updated after this time are returned. If not specified, returns sessions with
99+
active messages in the entity.
100+
:keyword float timeout: The total operation timeout in seconds.
101+
:returns: A list of session ID strings.
102+
:rtype: list[str]
103+
"""
104+
# DateTime.MaxValue triggers "active messages" mode on the service side.
105+
# A real timestamp triggers "sessions updated since" mode.
106+
use_sentinel = updated_since is None
107+
last_updated_time_ms = (
108+
_MAX_DATETIME_MS if use_sentinel
109+
else int(updated_since.timestamp() * 1000)
110+
)
111+
112+
all_session_ids: List[str] = []
113+
skip = 0
114+
115+
while True:
116+
message = {
117+
b"last-updated-time": {TYPE: AMQPTypes.timestamp, VALUE: last_updated_time_ms},
118+
b"skip": _amqp_int_value(skip),
119+
b"top": _amqp_int_value(_PAGE_SIZE),
120+
}
121+
result = self._mgmt_request_response_with_retry(
122+
REQUEST_RESPONSE_GET_MESSAGE_SESSIONS_OPERATION,
123+
message,
124+
mgmt_handlers.list_sessions_op,
125+
keep_alive_associated_link=False,
126+
timeout=timeout,
127+
)
128+
if not result:
129+
break
130+
all_session_ids.extend(result)
131+
skip += len(result)
132+
133+
return all_session_ids

sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_pyamqp_transport.py

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
from .._pyamqp import (
1616
utils,
17+
AMQPClient,
1718
SendClient,
1819
constants,
1920
ReceiveClient,
@@ -442,6 +443,36 @@ def close_connection(connection: "Connection") -> None:
442443
"""
443444
connection.close()
444445

446+
@staticmethod
447+
def create_mgmt_client(config: "Configuration", **kwargs: Any) -> "AMQPClient":
448+
"""Creates and returns a pyamqp AMQPClient for management-only operations.
449+
450+
Unlike SendClient/ReceiveClient, this client does not create a sender or
451+
receiver link. It only opens a connection and authenticates, suitable for
452+
management requests that don't need an associated link.
453+
454+
:param ~azure.servicebus._configuration.Configuration config: The configuration. Required.
455+
:keyword ~pyamqp.authentication.JWTTokenAuth auth: Required.
456+
:keyword retry_policy: Required.
457+
:keyword str client_name: Required.
458+
:keyword dict properties: Required.
459+
:return: AMQPClient
460+
:rtype: ~pyamqp.AMQPClient
461+
"""
462+
return AMQPClient(
463+
config.hostname,
464+
network_trace=config.logging_enable,
465+
keep_alive_interval=config.keep_alive,
466+
custom_endpoint_address=config.custom_endpoint_address,
467+
connection_verify=config.connection_verify,
468+
ssl_context=config.ssl_context,
469+
transport_type=config.transport_type,
470+
http_proxy=config.http_proxy,
471+
socket_timeout=config.socket_timeout,
472+
use_tls=config.use_tls,
473+
**kwargs,
474+
)
475+
445476
@staticmethod
446477
def create_send_client(config: "Configuration", **kwargs: Any) -> "SendClient": # pylint: disable=docstring-keyword-should-match-keyword-only
447478
"""

sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_uamqp_transport.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -507,6 +507,28 @@ def close_connection(connection: "Connection") -> None:
507507
"""
508508
connection.destroy()
509509

510+
@staticmethod
511+
def create_mgmt_client(config: "Configuration", **kwargs: Any) -> "AMQPClient":
512+
"""Creates and returns a uamqp AMQPClient for management-only operations.
513+
514+
:param ~azure.servicebus._common._configuration.Configuration config: The configuration.
515+
:keyword JWTTokenAuth auth: Required.
516+
:keyword retry_policy: Required.
517+
:keyword str client_name: Required.
518+
:keyword dict properties: Required.
519+
:return: AMQPClient
520+
:rtype: ~uamqp.AMQPClient
521+
"""
522+
retry_policy = kwargs.pop("retry_policy")
523+
return AMQPClient(
524+
"amqps://" + config.hostname,
525+
debug=config.logging_enable,
526+
error_policy=retry_policy,
527+
keep_alive_interval=config.keep_alive,
528+
encoding=config.encoding,
529+
**kwargs,
530+
)
531+
510532
@staticmethod
511533
def create_send_client(config: "Configuration", **kwargs: Any) -> "SendClient": # pylint:disable=docstring-keyword-should-match-keyword-only
512534
"""

0 commit comments

Comments
 (0)