[ServiceBus&EventHubs] api fixes (Azure#22418)

* api updates

* fix mypy 3.7 errors

* Annas comments

* mypy literal variable error

* lint

* update type hint session_id

* docstring

* lint

* add SBSessionFilter to public API

* add typing_extensions to setup.py/add back kwargs with warnings

* add kwargs + warn to amqp message

* typing-extensions add to shared reqs

* lint

* add bool to body ED type

* anna + adam comments

* move to extras require

* add back to install reqs

Author: swathipil

sdk/eventhub/azure-eventhub/azure/eventhub/_common.py

@@ -6,6 +6,7 @@
 
 import json
 import logging
+import uuid
 from typing import (
     Union,
     Dict,
@@ -59,6 +60,17 @@
 if TYPE_CHECKING:
     import datetime
 
+PrimitiveTypes = Optional[Union[
+    int,
+    float,
+    bytes,
+    bool,
+    str,
+    Dict,
+    List,
+    uuid.UUID,
+]]
+
 _LOGGER = logging.getLogger(__name__)
 
 # event_data.encoded_size < 255, batch encode overhead is 5, >=256, overhead is 8 each
@@ -319,7 +331,7 @@ def system_properties(self):
 
     @property
     def body(self):
-        # type: () -> Any
+        # type: () -> PrimitiveTypes
         """The body of the Message. The format may vary depending on the body type:
         For :class:`azure.eventhub.amqp.AmqpMessageBodyType.DATA<azure.eventhub.amqp.AmqpMessageBodyType.DATA>`,
         the body could be bytes or Iterable[bytes].
@@ -328,7 +340,7 @@ def body(self):
         For :class:`azure.eventhub.amqp.AmqpMessageBodyType.VALUE<azure.eventhub.amqp.AmqpMessageBodyType.VALUE>`,
         the body could be any type.
 
-        :rtype: Any
+        :rtype: int or bool or float or bytes or str or dict or list or uuid.UUID
         """
         try:
             return self._raw_amqp_message.body

sdk/servicebus/azure-servicebus/CHANGELOG.md

@@ -8,6 +8,7 @@ This version and all future versions will require Python 3.7+. Python 2.7 and 3.
 
 - Added support for fixed (linear) retry backoff:
   - Sync/async `ServiceBusClient` constructors and `from_connection_string` take `retry_mode` as a keyword argument.
+- Added new enum class `ServiceBusSessionFilter`, which is the type of existing `NEXT_AVAILABLE_SESSION` value.
 
 ### Breaking Changes
 

sdk/servicebus/azure-servicebus/azure/servicebus/__init__.py

@@ -21,6 +21,7 @@
 from ._common.constants import (
     ServiceBusReceiveMode,
     ServiceBusSubQueue,
+    ServiceBusSessionFilter,
     NEXT_AVAILABLE_SESSION,
 )
 from ._common.auto_lock_renewer import AutoLockRenewer
@@ -37,6 +38,7 @@
     "ServiceBusReceivedMessage",
     "NEXT_AVAILABLE_SESSION",
     "ServiceBusSubQueue",
+    "ServiceBusSessionFilter",
     "ServiceBusReceiveMode",
     "ServiceBusClient",
     "ServiceBusReceiver",

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

@@ -5,6 +5,7 @@
 from typing import Any, Union, Optional, TYPE_CHECKING
 import logging
 from weakref import WeakSet
+from typing_extensions import Literal
 
 import uamqp
 
@@ -25,6 +26,7 @@
 from ._common.constants import (
     ServiceBusSubQueue,
     ServiceBusReceiveMode,
+    ServiceBusSessionFilter,
 )
 
 if TYPE_CHECKING:
@@ -34,6 +36,8 @@
         AzureNamedKeyCredential,
     )
 
+NextAvailableSessionType = Literal[ServiceBusSessionFilter.NEXT_AVAILABLE]
+
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -269,7 +273,7 @@ def get_queue_receiver(
         self,
         queue_name: str,
         *,
-        session_id: Optional[str] = None,
+        session_id: Optional[Union[str, NextAvailableSessionType]] = None,
         sub_queue: Optional[Union[ServiceBusSubQueue, str]] = None,
         receive_mode: Union[
             ServiceBusReceiveMode, str
@@ -285,13 +289,14 @@ def get_queue_receiver(
         :keyword session_id: A specific session from which to receive. This must be specified for a
          sessionful queue, otherwise it must be None. In order to receive messages from the next available
          session, set this to ~azure.servicebus.NEXT_AVAILABLE_SESSION.
-        :paramtype session_id: Union[str, ~azure.servicebus.NEXT_AVAILABLE_SESSION]
-        :keyword Optional[Union[ServiceBusSubQueue, str]] sub_queue: If specified, the subqueue this receiver will
+        :paramtype session_id: str or ~azure.servicebus.NEXT_AVAILABLE_SESSION
+        :keyword sub_queue: If specified, the subqueue this receiver will
          connect to.
          This includes the DEAD_LETTER and TRANSFER_DEAD_LETTER queues, holds messages that can't be delivered to any
          receiver or messages that can't be processed.
          The default is None, meaning connect to the primary queue.  Can be assigned values from `ServiceBusSubQueue`
          enum or equivalent string values "deadletter" and "transferdeadletter".
+        :paramtype sub_queue: str or ~azure.servicebus.ServiceBusSubQueue
         :keyword receive_mode: The receive_mode with which messages will be retrieved from the entity. The two options
          are PEEK_LOCK and RECEIVE_AND_DELETE. Messages received with PEEK_LOCK must be settled within a given
          lock period before they will be removed from the queue. Messages received with RECEIVE_AND_DELETE
@@ -421,7 +426,7 @@ def get_subscription_receiver(
         topic_name: str,
         subscription_name: str,
         *,
-        session_id: Optional[str] = None,
+        session_id: Optional[Union[str, NextAvailableSessionType]] = None,
         sub_queue: Optional[Union[ServiceBusSubQueue, str]] = None,
         receive_mode: Union[
             ServiceBusReceiveMode, str
@@ -439,13 +444,14 @@ def get_subscription_receiver(
         :keyword session_id: A specific session from which to receive. This must be specified for a
          sessionful subscription, otherwise it must be None. In order to receive messages from the next available
          session, set this to ~azure.servicebus.NEXT_AVAILABLE_SESSION.
-        :paramtype session_id: Union[str, ~azure.servicebus.NEXT_AVAILABLE_SESSION]
-        :keyword Optional[Union[ServiceBusSubQueue, str]] sub_queue: If specified, the subqueue this receiver will
+        :paramtype session_id: str or ~azure.servicebus.NEXT_AVAILABLE_SESSION
+        :keyword sub_queue: If specified, the subqueue this receiver will
          connect to.
          This includes the DEAD_LETTER and TRANSFER_DEAD_LETTER queues, holds messages that can't be delivered to any
          receiver or messages that can't be processed.
          The default is None, meaning connect to the primary queue.  Can be assigned values from `ServiceBusSubQueue`
          enum or equivalent string values "deadletter" and "transferdeadletter".
+        :paramtype sub_queue: str or ~azure.servicebus.ServiceBusSubQueue
         :keyword receive_mode: The receive_mode with which messages will be retrieved from the entity. The two options
          are PEEK_LOCK and RECEIVE_AND_DELETE. Messages received with PEEK_LOCK must be settled within a given
          lock period before they will be removed from the subscription. Messages received with RECEIVE_AND_DELETE

sdk/servicebus/azure-servicebus/azure/servicebus/_servicebus_receiver.py

@@ -7,6 +7,7 @@
 import functools
 import uuid
 import datetime
+import warnings
 from typing import Any, List, Optional, Dict, Iterator, Union, TYPE_CHECKING
 
 import six
@@ -660,7 +661,8 @@ def receive_deferred_messages(
         self,
         sequence_numbers: Union[int, List[int]],
         *,
-        timeout: Optional[float] = None
+        timeout: Optional[float] = None,
+        **kwargs: Any
     ) -> List[ServiceBusReceivedMessage]:
         """Receive messages that have previously been deferred.
 
@@ -683,6 +685,8 @@ def receive_deferred_messages(
                 :caption: Receive deferred messages from ServiceBus.
 
         """
+        if kwargs:
+            warnings.warn(f"Unsupported keyword args: {kwargs}")
         self._check_live()
         if timeout is not None and timeout <= 0:
             raise ValueError("The timeout must be greater than 0.")
@@ -734,7 +738,8 @@ def peek_messages(
         max_message_count: int = 1,
         *,
         sequence_number: int = 0,
-        timeout: Optional[float] = None
+        timeout: Optional[float] = None,
+        **kwargs: Any
     ) -> List[ServiceBusReceivedMessage]:
         """Browse messages currently pending in the queue.
 
@@ -759,6 +764,8 @@ def peek_messages(
                 :caption: Look at pending messages in the queue.
 
         """
+        if kwargs:
+            warnings.warn(f"Unsupported keyword args: {kwargs}")
         self._check_live()
         if timeout is not None and timeout <= 0:
             raise ValueError("The timeout must be greater than 0.")
@@ -899,7 +906,8 @@ def renew_message_lock(
         self,
         message: ServiceBusReceivedMessage,
         *,
-        timeout: Optional[float] = None
+        timeout: Optional[float] = None,
+        **kwargs: Any
     ) -> datetime.datetime:
         # pylint: disable=protected-access,no-member
         """Renew the message lock.
@@ -933,6 +941,8 @@ def renew_message_lock(
                 :caption: Renew the lock on a received message.
 
         """
+        if kwargs:
+            warnings.warn(f"Unsupported keyword args: {kwargs}")
         # type: ignore
         try:
             if self.session:

sdk/servicebus/azure-servicebus/azure/servicebus/_servicebus_sender.py

@@ -6,6 +6,7 @@
 import time
 import uuid
 import datetime
+import warnings
 from typing import Any, TYPE_CHECKING, Union, List, Optional, Mapping, cast
 
 import uamqp
@@ -278,7 +279,8 @@ def schedule_messages(
         messages: "MessageTypes",
         schedule_time_utc: datetime.datetime,
         *,
-        timeout: Optional[float] = None
+        timeout: Optional[float] = None,
+        **kwargs: Any
     ) -> List[int]:
         """Send Message or multiple Messages to be enqueued at a specific time.
         Returns a list of the sequence numbers of the enqueued messages.
@@ -301,6 +303,8 @@ def schedule_messages(
                 :dedent: 4
                 :caption: Schedule a message to be sent in future
         """
+        if kwargs:
+            warnings.warn(f"Unsupported keyword args: {kwargs}")
         # pylint: disable=protected-access
 
         self._check_live()
@@ -332,7 +336,8 @@ def cancel_scheduled_messages(
         self,
         sequence_numbers: Union[int, List[int]],
         *,
-        timeout: Optional[float] = None
+        timeout: Optional[float] = None,
+        **kwargs: Any
     ) -> None:
         """
         Cancel one or more messages that have previously been scheduled and are still pending.
@@ -354,6 +359,8 @@ def cancel_scheduled_messages(
                 :dedent: 4
                 :caption: Cancelling messages scheduled to be sent in future
         """
+        if kwargs:
+            warnings.warn(f"Unsupported keyword args: {kwargs}")
         self._check_live()
         if timeout is not None and timeout <= 0:
             raise ValueError("The timeout must be greater than 0.")
@@ -375,7 +382,8 @@ def send_messages(
         self,
         message: Union["MessageTypes", ServiceBusMessageBatch],
         *,
-        timeout: Optional[float] = None
+        timeout: Optional[float] = None,
+        **kwargs: Any
     ) -> None:
         """Sends message and blocks until acknowledgement is received or operation times out.
 
@@ -407,6 +415,8 @@ def send_messages(
                 :caption: Send message.
 
         """
+        if kwargs:
+            warnings.warn(f"Unsupported keyword args: {kwargs}")
         self._check_live()
         if timeout is not None and timeout <= 0:
             raise ValueError("The timeout must be greater than 0.")

sdk/servicebus/azure-servicebus/azure/servicebus/_servicebus_session.py

@@ -4,7 +4,8 @@
 # --------------------------------------------------------------------------------------------
 import logging
 import datetime
-from typing import TYPE_CHECKING, Union, Optional
+import warnings
+from typing import TYPE_CHECKING, Any, Union, Optional
 import six
 
 from ._common.utils import utc_from_timestamp, utc_now
@@ -88,15 +89,15 @@ class ServiceBusSession(BaseSession):
             :caption: Get session from a receiver
     """
 
-    def get_state(self, *, timeout: Optional[float] = None) -> bytes:
+    def get_state(self, *, timeout: Optional[float] = None, **kwargs: Any) -> bytes:
         # pylint: disable=protected-access
         """Get the session state.
 
         Returns None if no state has been set.
 
         :keyword float timeout: The total operation timeout in seconds including all the retries. The value must be
          greater than 0 if specified. The default value is None, meaning no timeout.
-        :rtype: str
+        :rtype: bytes
 
         .. admonition:: Example:
 
@@ -107,6 +108,8 @@ def get_state(self, *, timeout: Optional[float] = None) -> bytes:
                 :dedent: 4
                 :caption: Get the session state
         """
+        if kwargs:
+            warnings.warn(f"Unsupported keyword args: {kwargs}")
         self._receiver._check_live()  # pylint: disable=protected-access
         if timeout is not None and timeout <= 0:
             raise ValueError("The timeout must be greater than 0.")
@@ -119,7 +122,7 @@ def get_state(self, *, timeout: Optional[float] = None) -> bytes:
         session_state = response.get(MGMT_RESPONSE_SESSION_STATE)  # type: ignore
         return session_state
 
-    def set_state(self, state: Union[str, bytes, bytearray], *, timeout: Optional[float] = None) -> None:
+    def set_state(self, state: Union[str, bytes, bytearray], *, timeout: Optional[float] = None, **kwargs: Any) -> None:
         # pylint: disable=protected-access
         """Set the session state.
 
@@ -137,6 +140,8 @@ def set_state(self, state: Union[str, bytes, bytearray], *, timeout: Optional[fl
                 :dedent: 4
                 :caption: Set the session state
         """
+        if kwargs:
+            warnings.warn(f"Unsupported keyword args: {kwargs}")
         self._receiver._check_live()  # pylint: disable=protected-access
         if timeout is not None and timeout <= 0:
             raise ValueError("The timeout must be greater than 0.")
@@ -153,7 +158,7 @@ def set_state(self, state: Union[str, bytes, bytearray], *, timeout: Optional[fl
             timeout=timeout,
         )
 
-    def renew_lock(self, *, timeout: Optional[float] = None) -> datetime.datetime:
+    def renew_lock(self, *, timeout: Optional[float] = None, **kwargs: Any) -> datetime.datetime:
         # pylint: disable=protected-access
         """Renew the session lock.
 
@@ -179,6 +184,8 @@ def renew_lock(self, *, timeout: Optional[float] = None) -> datetime.datetime:
                 :dedent: 4
                 :caption: Renew the session lock before it expires
         """
+        if kwargs:
+            warnings.warn(f"Unsupported keyword args: {kwargs}")
         self._receiver._check_live()  # pylint: disable=protected-access
         if timeout is not None and timeout <= 0:
             raise ValueError("The timeout must be greater than 0.")