diff --git a/sdk/eventhub/azure-eventhub/CHANGELOG.md b/sdk/eventhub/azure-eventhub/CHANGELOG.md index a56924903cec..d0ca49ff82eb 100644 --- a/sdk/eventhub/azure-eventhub/CHANGELOG.md +++ b/sdk/eventhub/azure-eventhub/CHANGELOG.md @@ -1,6 +1,6 @@ # Release History -## 5.10.2 (Unreleased) +## 5.11.0 (Unreleased) ### Features Added @@ -88,9 +88,66 @@ This version and all future versions will require Python 3.7+, Python 3.6 is no ## 5.9.0b1 (2022-02-09) +- The following features have been temporarily pulled out of async `EventHubProducerClient` and `EventHubConsumerClient` which will be added back in future previews as we work towards a stable release: + - Passing the following keyword arguments to the constructors and `from_connection_string` methods of the `EventHubProducerClient` and `EventHubConsumerClient` is not supported: `transport_type`, `http_proxy`, `custom_endpoint_address`, and `connection_verify`. + +## 5.8.0b2 (2022-10-11) + +### Features Added + +- Updated the optional dependency for async transport using AMQP over WebSocket from `websocket-client` to `aiohttp` (Issue #24315, thanks @hansmbakker for the suggestion). + +## 5.8.0b1 (2022-09-22) + +This version and all future versions will require Python 3.7+. Python 3.6 is no longer supported. + +### Other Changes + +- Added the `uamqp_transport` optional parameter to the clients, to allow switching to the `uamqp` library as the transport. + +## 5.8.0a5 (2022-07-19) + +### Bugs Fixed + +- Fixed bug that prevented token refresh at regular intervals. +- Fixed bug that was improperly passing the debug keyword argument, so that network trace debug logs are output when requested. + +### Other Changes + +- Added logging added in to track proper token refreshes & fetches, output exception reason for producer init failure. + +## 5.8.0a4 (2022-06-07) + +### Features Added + +- Added support for connection using websocket and http proxy. +- Added support for custom endpoint connection over websocket. + +## 5.8.0a3 (2022-03-08) + +### Other Changes + +- Improved the performance of async sending and receiving. + +## 5.8.0a2 (2022-02-09) + ### Features Added -- The classmethod `from_message_data` has been added to `EventData` for interoperability with the Schema Registry Avro Encoder library, and takes `data` and `content_type` as positional parameters. +- Added support for async `EventHubProducerClient` and `EventHubConsumerClient`. + +## 5.8.0a1 (2022-01-13) + +Version 5.8.0a1 is our first efforts to build an Azure Event Hubs client library based on pure python implemented AMQP stack. + +### Breaking changes + +- The following features have been temporarily pulled out which will be added back in future previews as we work towards a stable release: + - Async is not supported. + - Passing the following keyword arguments to the constructors and `from_connection_string` methods of the `EventHubProducerClient` and `EventHubConsumerClient` is not supported: `transport_type`, `http_proxy`, `custom_endpoint_address`, and `connection_verify`. + +### Other Changes + +- uAMQP dependency is removed. ## 5.7.0 (2022-01-12) @@ -598,4 +655,6 @@ Version 5.0.0b1 is a preview of our efforts to create a client library that is u - Further testing and minor bug fixes. -![Impressions](https://azure-sdk-impressions.azurewebsites.net/api/impressions/azure-sdk-for-python/sdk/eventhub/azure-eventhub/HISTORY.png) +## 0.2.0a2 (2018-04-02) + +- Updated uAQMP dependency. diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_buffered_producer/_buffered_producer.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_buffered_producer/_buffered_producer.py index cb2725c22a4d..9babfcb7d2bb 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_buffered_producer/_buffered_producer.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_buffered_producer/_buffered_producer.py @@ -32,6 +32,7 @@ def __init__( max_message_size_on_link: int, executor: ThreadPoolExecutor, *, + amqp_transport: AmqpTransport, max_buffer_length: int, max_wait_time: float = 1 ): @@ -50,10 +51,11 @@ def __init__( self._max_message_size_on_link = max_message_size_on_link self._check_max_wait_time_future = None self.partition_id = partition_id + self._amqp_transport = amqp_transport def start(self): with self._lock: - self._cur_batch = EventDataBatch(self._max_message_size_on_link) + self._cur_batch = EventDataBatch(self._max_message_size_on_link, amqp_transport=self._amqp_transport) self._running = True if self._max_wait_time: self._last_send_time = time.time() @@ -113,12 +115,12 @@ def put_events(self, events, timeout_time=None): self._buffered_queue.put(self._cur_batch) self._buffered_queue.put(events) # create a new batch for incoming events - self._cur_batch = EventDataBatch(self._max_message_size_on_link) + self._cur_batch = EventDataBatch(self._max_message_size_on_link, amqp_transport=self._amqp_transport) except ValueError: # add single event exceeds the cur batch size, create new batch with self._lock: self._buffered_queue.put(self._cur_batch) - self._cur_batch = EventDataBatch(self._max_message_size_on_link) + self._cur_batch = EventDataBatch(self._max_message_size_on_link, amqp_transport=self._amqp_transport) self._cur_batch.add(events) with self._lock: self._cur_buffered_len += new_events_len @@ -197,7 +199,7 @@ def flush(self, timeout_time=None, raise_error=True): self._last_send_time = time.time() #reset buffered count self._cur_buffered_len = 0 - self._cur_batch = EventDataBatch(self._max_message_size_on_link) + self._cur_batch = EventDataBatch(self._max_message_size_on_link, amqp_transport=self._amqp_transport) _LOGGER.info("Partition %r finished flushing.", self.partition_id) def check_max_wait_time_worker(self): diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_buffered_producer/_buffered_producer_dispatcher.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_buffered_producer/_buffered_producer_dispatcher.py index 71f97f15fecd..7bad79bead64 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_buffered_producer/_buffered_producer_dispatcher.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_buffered_producer/_buffered_producer_dispatcher.py @@ -14,8 +14,8 @@ from ..exceptions import EventDataSendError, ConnectError, EventHubError if TYPE_CHECKING: - from .._producer_client import SendEventTypes from .._transport._base import AmqpTransport + from .._producer_client import SendEventTypes _LOGGER = logging.getLogger(__name__) @@ -31,6 +31,7 @@ def __init__( eventhub_name: str, max_message_size_on_link: int, *, + amqp_transport: AmqpTransport, max_buffer_length: int = 1500, max_wait_time: float = 1, executor: Optional[Union[ThreadPoolExecutor, int]] = None @@ -47,6 +48,7 @@ def __init__( self._max_wait_time = max_wait_time self._max_buffer_length = max_buffer_length self._existing_executor = False + self._amqp_transport = amqp_transport if not executor: self._executor = ThreadPoolExecutor() @@ -88,6 +90,7 @@ def enqueue_events( executor=self._executor, max_wait_time=self._max_wait_time, max_buffer_length=self._max_buffer_length, + amqp_transport = self._amqp_transport, ) buffered_producer.start() self._buffered_producers[pid] = buffered_producer diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_client_base.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_client_base.py index 61e59688cbc4..7a239d1dbd13 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_client_base.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_client_base.py @@ -11,7 +11,7 @@ import collections from typing import Any, Dict, Tuple, List, Optional, TYPE_CHECKING, cast, Union try: - from typing import TypeAlias + from typing import TypeAlias # type: ignore except ImportError: from typing_extensions import TypeAlias from datetime import timedelta @@ -25,11 +25,15 @@ from azure.core.utils import parse_connection_string as core_parse_connection_string from azure.core.pipeline.policies import RetryMode - -from ._transport._uamqp_transport import UamqpTransport +try: + from ._transport._uamqp_transport import UamqpTransport +except ImportError: + UamqpTransport = None # type: ignore +from ._transport._pyamqp_transport import PyamqpTransport from .exceptions import ClientClosedError from ._configuration import Configuration -from ._utils import utc_from_timestamp, parse_sas_credential, generate_sas_token +from ._utils import utc_from_timestamp, parse_sas_credential +from ._pyamqp.utils import generate_sas_token from ._connection_manager import get_connection_manager from ._constants import ( CONTAINER_PREFIX, @@ -43,8 +47,14 @@ if TYPE_CHECKING: from azure.core.credentials import TokenCredential - from uamqp import Message as uamqp_Message - from uamqp.authentication import JWTTokenAuth as uamqp_JWTTokenAuth + try: + from uamqp import Message as uamqp_Message + from uamqp.authentication import JWTTokenAuth as uamqp_JWTTokenAuth + except ImportError: + uamqp_Message = None + uamqp_JWTTokenAuth = None + from ._pyamqp.message import Message + from ._pyamqp.authentication import JWTTokenAuth _LOGGER = logging.getLogger(__name__) _Address = collections.namedtuple("_Address", "hostname path") @@ -165,7 +175,7 @@ def _get_backoff_time(retry_mode, backoff_factor, backoff_max, retried_times): if retry_mode == RetryMode.Fixed: backoff_value = backoff_factor else: - backoff_value = backoff_factor * (2**retried_times) + backoff_value = backoff_factor * (2 ** retried_times) return min(backoff_max, backoff_value) @@ -262,6 +272,7 @@ def get_token(self, *scopes, **kwargs): # pylint:disable=unused-argument return AccessToken(signature, expiry) +# separate TYPE_CHECKING block here for EventHubSharedKeyCredential, o/w mypy raised error even with forward referencing if TYPE_CHECKING: from azure.core.credentials import TokenCredential @@ -281,8 +292,10 @@ def __init__( credential: CredentialTypes, **kwargs: Any, ) -> None: - uamqp_transport = kwargs.pop("uamqp_transport", True) - self._amqp_transport = kwargs.pop("amqp_transport", UamqpTransport) + uamqp_transport = kwargs.pop("uamqp_transport", False) + if uamqp_transport and not UamqpTransport: + raise ValueError("To use the uAMQP transport, please install `uamqp>=1.6.0,<2.0.0`.") + self._amqp_transport = kwargs.pop("amqp_transport", UamqpTransport if uamqp_transport else PyamqpTransport) self.eventhub_name = eventhub_name if not eventhub_name: @@ -305,7 +318,10 @@ def __init__( **kwargs, ) self._debug = self._config.network_tracing - self._conn_manager = get_connection_manager(**kwargs) + kwargs["custom_endpoint_address"] = self._config.custom_endpoint_address + self._conn_manager = get_connection_manager( + amqp_transport=self._amqp_transport, + **kwargs) self._idle_timeout = kwargs.get("idle_timeout", None) @staticmethod @@ -322,7 +338,7 @@ def _from_connection_string(conn_str, **kwargs): kwargs["credential"] = EventHubSharedKeyCredential(policy, key) return kwargs - def _create_auth(self) -> uamqp_JWTTokenAuth: + def _create_auth(self) -> Union[uamqp_JWTTokenAuth, JWTTokenAuth]: """ Create an ~uamqp.authentication.SASTokenAuth instance to authenticate the session. @@ -381,7 +397,7 @@ def _backoff( raise last_exception def _management_request( - self, mgmt_msg: uamqp_Message, op_type: bytes + self, mgmt_msg: Union[uamqp_Message, Message], op_type: bytes ) -> Any: # pylint:disable=assignment-from-none retried_times = 0 @@ -401,7 +417,7 @@ def _management_request( mgmt_msg.application_properties[ "security_token" ] = self._amqp_transport.get_updated_token(mgmt_auth) - response = self._amqp_transport.mgmt_client_request( + status_code, description, response = self._amqp_transport.mgmt_client_request( mgmt_client, mgmt_msg, operation=READ_OPERATION, @@ -409,18 +425,21 @@ def _management_request( status_code_field=MGMT_STATUS_CODE, description_fields=MGMT_STATUS_DESC, ) - status_code = int(response.application_properties[MGMT_STATUS_CODE]) - description = response.application_properties.get( - MGMT_STATUS_DESC - ) # type: Optional[Union[str, bytes]] + status_code = int(status_code) if description and isinstance(description, bytes): description = description.decode("utf-8") if status_code < 400: return response raise self._amqp_transport.get_error(status_code, description) except Exception as exception: # pylint: disable=broad-except + # is_consumer=True passed in here, ALTHOUGH this method is shared by the producer and consumer. + # is_consumer will only be checked if FileNotFoundError is raised by self.mgmt_client.open() due to + # invalid/non-existent connection_verify filepath. The producer will encounter the FileNotFoundError + # when opening the SendClient, so is_consumer=True will not be passed to amqp_transport.handle_exception + # there. This is for uamqp exception parity, which raises FileNotFoundError in the consumer and + # EventHubError in the producer. TODO: Remove `is_consumer` kwarg when resolving issue #27128. last_exception = self._amqp_transport._handle_exception( # pylint: disable=protected-access - exception, self + exception, self, is_consumer=True ) self._backoff( retried_times=retried_times, last_exception=last_exception @@ -540,10 +559,10 @@ def _close_connection(self): self._close_handler() self._client._conn_manager.reset_connection_if_broken() # pylint: disable=protected-access - def _handle_exception(self, exception): + def _handle_exception(self, exception, *, is_consumer=False): exception = self._amqp_transport.check_timeout_exception(self, exception) return self._amqp_transport._handle_exception( # pylint: disable=protected-access - exception, self + exception, self, is_consumer=is_consumer ) def _do_retryable_operation(self, operation, timeout=None, **kwargs): diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_common.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_common.py index 969d9c5bfb5f..633a90f8d0a9 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_common.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_common.py @@ -5,6 +5,7 @@ from __future__ import unicode_literals, annotations import json +import warnings import datetime import logging import uuid @@ -52,10 +53,19 @@ AmqpMessageHeader, AmqpMessageProperties, ) -from ._transport._uamqp_transport import UamqpTransport +from ._pyamqp._message_backcompat import LegacyMessage, LegacyBatchMessage +from ._pyamqp.message import Message +from ._transport._pyamqp_transport import PyamqpTransport if TYPE_CHECKING: - from uamqp import Message as uamqp_Message, BatchMessage as uamqp_BatchMessage + try: + from uamqp import ( # pylint: disable=unused-import + Message as uamqp_Message, + BatchMessage, + ) + except ImportError: + uamqp_Message = None + BatchMessage = None from ._transport._base import AmqpTransport MessageContent = TypedDict("MessageContent", {"content": bytes, "content_type": str}) @@ -124,9 +134,8 @@ def __init__( self._raw_amqp_message = AmqpAnnotatedMessage( # type: ignore data_body=body, annotations={}, application_properties={} ) - # amqp message to be reset right before sending - self._message = UamqpTransport.to_outgoing_amqp_message(self._raw_amqp_message) - self.message = self._message + self._uamqp_message: Optional[Union[LegacyMessage, uamqp_Message]] = None + self._message: Message = None # type: ignore self._raw_amqp_message.header = AmqpMessageHeader() self._raw_amqp_message.properties = AmqpMessageProperties() self.message_id = None @@ -137,35 +146,42 @@ def __repr__(self) -> str: # pylint: disable=bare-except try: body_str = self.body_as_str() - except: + except Exception as e: # pylint: disable=broad-except + _LOGGER.debug("Message body read error: %r", e) body_str = "" event_repr = f"body='{body_str}'" try: event_repr += f", properties={self.properties}" - except: + except Exception as e: # pylint: disable=broad-except + _LOGGER.debug("Message properties read error: %r", e) event_repr += ", properties=" try: event_repr += f", offset={self.offset}" - except: + except Exception as e: # pylint: disable=broad-except + _LOGGER.debug("Message offset read error: %r", e) event_repr += ", offset=" try: event_repr += f", sequence_number={self.sequence_number}" - except: + except Exception as e: # pylint: disable=broad-except + _LOGGER.debug("Message sequence number read error: %r", e) event_repr += ", sequence_number=" try: event_repr += f", partition_key={self.partition_key!r}" - except: + except Exception as e: # pylint: disable=broad-except + _LOGGER.debug("Message partition key read error: %r", e) event_repr += ", partition_key=" try: event_repr += f", enqueued_time={self.enqueued_time!r}" - except: + except Exception as e: # pylint: disable=broad-except + _LOGGER.debug("Message enqueued time read error: %r", e) event_repr += ", enqueued_time=" return f"EventData({event_repr})" def __str__(self) -> str: try: body_str = self.body_as_str() - except: # pylint: disable=bare-except + except Exception as e: # pylint: disable=broad-except + _LOGGER.debug("Message body read error: %r", e) body_str = "" event_str = f"{{ body: '{body_str}'" try: @@ -178,8 +194,8 @@ def __str__(self) -> str: event_str += f", partition_key={self.partition_key!r}" if self.enqueued_time: event_str += f", enqueued_time={self.enqueued_time!r}" - except: # pylint: disable=bare-except - pass + except Exception as e: # pylint: disable=broad-except + _LOGGER.debug("Message metadata read error: %r", e) event_str += " }" return event_str @@ -210,7 +226,7 @@ def from_message_content( # pylint: disable=unused-argument @classmethod def _from_message( cls, - message: uamqp_Message, + message: Union[uamqp_Message, Message], raw_amqp_message: Optional[AmqpAnnotatedMessage] = None, ) -> EventData: # pylint:disable=protected-access @@ -225,7 +241,6 @@ def _from_message( event_data = cls(body="") # pylint: disable=protected-access event_data._message = message - event_data.message = message event_data._raw_amqp_message = ( raw_amqp_message if raw_amqp_message @@ -244,6 +259,35 @@ def _decode_non_data_body_as_str(self, encoding: str = "UTF-8") -> str: seq_list = [d for seq_section in body for d in seq_section] return str(decode_with_recurse(seq_list, encoding)) + @property + def message(self) -> LegacyMessage: + """DEPRECATED: Get the underlying LegacyMessage. + This is deprecated and will be removed in a later release. + + :rtype: LegacyMessage + """ + warnings.warn( + "The `message` property is deprecated and will be removed in future versions.", + DeprecationWarning, + ) + if not self._uamqp_message: + self._uamqp_message = LegacyMessage( + self._raw_amqp_message, + to_outgoing_amqp_message=PyamqpTransport().to_outgoing_amqp_message, + ) + return self._uamqp_message + + @message.setter + def message(self, value: "uamqp_Message") -> None: + """DEPRECATED: Set the underlying Message. + This is deprecated and will be removed in a later release. + """ + warnings.warn( + "The `message` property is deprecated and will be removed in future versions.", + DeprecationWarning, + ) + self._uamqp_message = value + @property def raw_amqp_message(self) -> AmqpAnnotatedMessage: """Advanced usage only. The internal AMQP message payload that is sent or received.""" @@ -379,9 +423,11 @@ def body_as_str(self, encoding: str = "UTF-8") -> str: if self.body_type != AmqpMessageBodyType.DATA: return self._decode_non_data_body_as_str(encoding=encoding) return "".join(b.decode(encoding) for b in cast(Iterable[bytes], data)) - except TypeError: + except UnicodeDecodeError as e: + raise TypeError(f"Message data is not compatible with string type: {e}") + except TypeError as e: return str(data) - except: # pylint: disable=bare-except + except Exception: # pylint: disable=broad-except pass try: return cast(bytes, data).decode(encoding) @@ -493,13 +539,12 @@ def __init__( self, max_size_in_bytes: Optional[int] = None, partition_id: Optional[str] = None, - partition_key: Optional[Union[str, bytes]] = None + partition_key: Optional[Union[str, bytes]] = None, + **kwargs, ) -> None: - self._amqp_transport = UamqpTransport + self._amqp_transport = kwargs.pop("amqp_transport", PyamqpTransport) - if partition_key and not isinstance( - partition_key, (str, bytes) - ): + if partition_key and not isinstance(partition_key, (str, bytes)): _LOGGER.info( "WARNING: Setting partition_key of non-string value on the events to be sent is discouraged " "as the partition_key will be ignored by the Event Hub service and events will be assigned " @@ -507,22 +552,25 @@ def __init__( "partition_key to only be string type, they might fail to parse the non-string value." ) - self.max_size_in_bytes = ( - max_size_in_bytes or self._amqp_transport.MAX_MESSAGE_LENGTH_BYTES - ) - self._message = self._amqp_transport.build_batch_message(data=[]) self._partition_id = partition_id self._partition_key = partition_key + self._message = self._amqp_transport.build_batch_message(data=[]) self._message = self._amqp_transport.set_message_partition_key( self._message, self._partition_key ) - self.message: uamqp_BatchMessage = self._message self._size = self._amqp_transport.get_batch_message_encoded_size(self._message) + self.max_size_in_bytes = ( + max_size_in_bytes or self._amqp_transport.MAX_MESSAGE_LENGTH_BYTES + ) + self._count = 0 - self._internal_events: List[ - Union[EventData, AmqpAnnotatedMessage] - ] = [] + self._internal_events: List[Union[EventData, AmqpAnnotatedMessage]] = [] + self._uamqp_message = ( + None + if PyamqpTransport.TIMEOUT_FACTOR == self._amqp_transport.TIMEOUT_FACTOR + else self._message + ) def __repr__(self) -> str: batch_repr = ( @@ -537,9 +585,12 @@ def __len__(self) -> int: @classmethod def _from_batch( cls, - batch_data: Iterable[EventData], + batch_data: Iterable[Union[AmqpAnnotatedMessage, EventData]], amqp_transport: AmqpTransport, partition_key: Optional[AnyStr] = None, + *, + max_size_in_bytes: Optional[int] = None, + partition_id: Optional[str] = None, ) -> EventDataBatch: outgoing_batch_data = [ transform_outbound_single_message( @@ -547,7 +598,12 @@ def _from_batch( ) for m in batch_data ] - batch_data_instance = cls(partition_key=partition_key) + batch_data_instance = cls( + partition_key=partition_key, + amqp_transport=amqp_transport, + max_size_in_bytes=max_size_in_bytes, + partition_id=partition_id, + ) for event_data in outgoing_batch_data: batch_data_instance.add(event_data) @@ -564,6 +620,36 @@ def _load_events(self, events): "or use EventDataBatch, which is guaranteed to be under the frame size limit" ) + @property + def message(self) -> Union["BatchMessage", LegacyBatchMessage]: + """DEPRECATED: Get the underlying uamqp.BatchMessage or LegacyBatchMessage. + This is deprecated and will be removed in a later release. + + :rtype: uamqp.BatchMessage or LegacyBatchMessage + """ + warnings.warn( + "The `message` property is deprecated and will be removed in future versions.", + DeprecationWarning, + ) + if not self._uamqp_message: + message = AmqpAnnotatedMessage(message=Message(*self._message)) + self._uamqp_message = LegacyBatchMessage( + message, + to_outgoing_amqp_message=PyamqpTransport().to_outgoing_amqp_message, + ) + return self._uamqp_message + + @message.setter + def message(self, value: "BatchMessage") -> None: + """DEPRECATED: Set the underlying BatchMessage. + This is deprecated and will be removed in a later release. + """ + warnings.warn( + "The `message` property is deprecated and will be removed in future versions.", + DeprecationWarning, + ) + self._uamqp_message = value + @property def size_in_bytes(self) -> int: """The combined size of the events in the batch, in bytes. @@ -598,7 +684,7 @@ def add(self, event_data: Union[EventData, AmqpAnnotatedMessage]) -> None: "The partition key of event_data does not match the partition key of this batch." ) if not outgoing_event_data.partition_key: - self._amqp_transport.set_message_partition_key( + outgoing_event_data._message = self._amqp_transport.set_message_partition_key( # pylint: disable=protected-access outgoing_event_data._message, # pylint: disable=protected-access self._partition_key, ) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_configuration.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_configuration.py index 00c03ca4197b..fe87d65b60cd 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_configuration.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_configuration.py @@ -9,6 +9,7 @@ from ._constants import TransportType, DEFAULT_AMQPS_PORT, DEFAULT_AMQP_WSS_PORT + class Configuration(object): # pylint:disable=too-many-instance-attributes def __init__(self, **kwargs): self.user_agent = kwargs.get("user_agent") # type: Optional[str] @@ -29,9 +30,7 @@ def __init__(self, **kwargs): self.max_batch_size = kwargs.get("max_batch_size", self.prefetch) # type: int self.receive_timeout = kwargs.get("receive_timeout", 0) # type: int self.send_timeout = kwargs.get("send_timeout", 60) # type: int - self.custom_endpoint_address = kwargs.get( - "custom_endpoint_address" - ) # type: Optional[str] + self.custom_endpoint_address = kwargs.get("custom_endpoint_address") # type: Optional[str] self.connection_verify = kwargs.get("connection_verify") # type: Optional[str] self.connection_port = DEFAULT_AMQPS_PORT self.custom_endpoint_hostname = None diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_connection_manager.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_connection_manager.py index f8e109a224cc..e51991a6677c 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_connection_manager.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_connection_manager.py @@ -4,16 +4,22 @@ # -------------------------------------------------------------------------------------------- from __future__ import annotations -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Optional, Union from threading import Lock from enum import Enum -from ._transport._uamqp_transport import UamqpTransport from ._constants import TransportType if TYPE_CHECKING: - from uamqp.authentication import JWTTokenAuth - from uamqp import Connection + from ._pyamqp.authentication import JWTTokenAuth + from ._pyamqp._connection import Connection + try: + from uamqp.authentication import JWTTokenAuth as uamqp_JWTTokenAuth + from uamqp import Connection as uamqp_Connection + except ImportError: + uamqp_JWTTokenAuth = None + uamqp_Connection = None + from ._transport._base import AmqpTransport try: from typing_extensions import Protocol @@ -22,8 +28,12 @@ class ConnectionManager(Protocol): def get_connection( - self, *, host: Optional[str] = None, auth: Optional[JWTTokenAuth] = None, endpoint: Optional[str] = None - ) -> Connection: + self, + *, + host: Optional[str] = None, + auth: Optional[Union[JWTTokenAuth, uamqp_JWTTokenAuth]] = None, + endpoint: Optional[str] = None, + ) -> Union[Connection, uamqp_Connection]: pass def close_connection(self): @@ -41,9 +51,10 @@ class _ConnectionMode(Enum): class _SharedConnectionManager(object): # pylint:disable=too-many-instance-attributes def __init__(self, **kwargs): self._lock = Lock() - self._conn: Connection = None + self._conn: Union[Connection, uamqp_Connection] = None self._container_id = kwargs.get("container_id") + self._custom_endpoint_address = kwargs.get("custom_endpoint_address") self._debug = kwargs.get("debug") self._error_policy = kwargs.get("error_policy") self._properties = kwargs.get("properties") @@ -53,20 +64,23 @@ def __init__(self, **kwargs): self._max_frame_size = kwargs.get("max_frame_size") self._channel_max = kwargs.get("channel_max") self._idle_timeout = kwargs.get("idle_timeout") - self._remote_idle_timeout_empty_frame_send_ratio = kwargs.get( - "remote_idle_timeout_empty_frame_send_ratio" - ) - self._amqp_transport = kwargs.get("amqp_transport", UamqpTransport) + self._remote_idle_timeout_empty_frame_send_ratio = kwargs.get("remote_idle_timeout_empty_frame_send_ratio") + self._amqp_transport: AmqpTransport = kwargs.pop("amqp_transport") def get_connection( - self, *, host: Optional[str] = None, auth: Optional[JWTTokenAuth] = None, endpoint: Optional[str] = None - ) -> Connection: + self, + *, + host: Optional[str] = None, + auth: Optional[Union[JWTTokenAuth, uamqp_JWTTokenAuth]] = None, + endpoint: Optional[str] = None, + ) -> Union[Connection, uamqp_Connection]: with self._lock: if self._conn is None: self._conn = self._amqp_transport.create_connection( host=host, auth=auth, endpoint=endpoint, + custom_endpoint_address=self._custom_endpoint_address, container_id=self._container_id, max_frame_size=self._max_frame_size, channel_max=self._channel_max, @@ -99,7 +113,11 @@ def __init__(self, **kwargs): pass def get_connection( # pylint:disable=unused-argument, no-self-use - self, *, host: Optional[str] = None, auth: Optional[JWTTokenAuth] = None, endpoint: Optional[str] = None + self, + *, + host: Optional[str] = None, + auth: Optional[Union[JWTTokenAuth, uamqp_JWTTokenAuth]] = None, + endpoint: Optional[str] = None, ) -> None: return None @@ -114,7 +132,7 @@ def reset_connection_if_broken(self): def get_connection_manager(**kwargs): # type: (...) -> 'ConnectionManager' - connection_mode = kwargs.get("connection_mode", _ConnectionMode.SeparateConnection) # type: ignore + connection_mode = kwargs.get("connection_mode", _ConnectionMode.SeparateConnection) # type: ignore if connection_mode == _ConnectionMode.ShareConnection: return _SharedConnectionManager(**kwargs) return _SeparateConnectionManager(**kwargs) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_constants.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_constants.py index eb8fd4f6198f..de5659411a84 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_constants.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_constants.py @@ -34,6 +34,7 @@ TIMEOUT_SYMBOL = b"com.microsoft:timeout" RECEIVER_RUNTIME_METRIC_SYMBOL = b"com.microsoft:enable-receiver-runtime-metric" +MAX_MESSAGE_LENGTH_BYTES = 1024 * 1024 MAX_USER_AGENT_LENGTH = 512 ALL_PARTITIONS = "all-partitions" CONTAINER_PREFIX = "eventhub.pysdk-" diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_consumer.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_consumer.py index 8f647a60c6e7..1f51155497b7 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_consumer.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_consumer.py @@ -8,7 +8,7 @@ import uuid import logging from collections import deque -from typing import TYPE_CHECKING, Callable, Dict, Optional, Any, Deque +from typing import TYPE_CHECKING, Callable, Dict, Optional, Any, Deque, Union, cast from ._common import EventData from ._client_base import ConsumerProducerMixin @@ -20,9 +20,20 @@ ) if TYPE_CHECKING: - from uamqp import ReceiveClient as uamqp_ReceiveClient, Message as uamqp_Message, types as uamqp_types - from uamqp.authentication import JWTTokenAuth as uamqp_JWTTokenAuth + from ._pyamqp import types + from ._pyamqp.message import Message + from ._pyamqp.authentication import JWTTokenAuth + from ._pyamqp.client import ReceiveClient + try: + from uamqp import ReceiveClient as uamqp_ReceiveClient, Message as uamqp_Message + from uamqp.types import AMQPType as uamqp_AMQPType + from uamqp.authentication import JWTTokenAuth as uamqp_JWTTokenAuth + except ImportError: + uamqp_ReceiveClient = None + uamqp_Message = None + uamqp_AMQPType = None + uamqp_JWTTokenAuth = None from ._consumer_client import EventHubConsumerClient @@ -49,7 +60,7 @@ class EventHubConsumer( :param client: The parent EventHubConsumerClient. :type client: ~azure.eventhub.EventHubConsumerClient :param source: The source EventHub from which to receive events. - :type source: ~uamqp.address.Source + :type source: ~azure.eventhub._pyamqp.endpoints.Source or ~uamqp.address.Source :keyword event_position: The position from which to start receiving. :paramtype event_position: int, str, datetime.datetime :keyword int prefetch: The number of events to prefetch from the service @@ -65,7 +76,9 @@ class EventHubConsumer( It is set to `False` by default. """ - def __init__(self, client: "EventHubConsumerClient", source: str, **kwargs: Any) -> None: + def __init__( + self, client: "EventHubConsumerClient", source: str, **kwargs: Any + ) -> None: event_position = kwargs.get("event_position", None) prefetch = kwargs.get("prefetch", 300) owner_level = kwargs.get("owner_level", None) @@ -93,39 +106,52 @@ def __init__(self, client: "EventHubConsumerClient", source: str, **kwargs: Any) self._owner_level = owner_level self._keep_alive = keep_alive self._auto_reconnect = auto_reconnect - self._retry_policy = self._amqp_transport.create_retry_policy(self._client._config) + self._retry_policy = self._amqp_transport.create_retry_policy( + self._client._config + ) self._reconnect_backoff = 1 - link_properties: Dict[uamqp_types.AMQPType, uamqp_types.AMQPType] = {} + link_properties: Dict[bytes, int] = {} self._error = None self._timeout = 0 - self._idle_timeout = (idle_timeout * self._amqp_transport.TIMEOUT_FACTOR) if idle_timeout else None + self._idle_timeout = ( + (idle_timeout * self._amqp_transport.TIMEOUT_FACTOR) + if idle_timeout + else None + ) self._partition = self._source.split("/")[-1] self._name = f"EHConsumer-{uuid.uuid4()}-partition{self._partition}" if owner_level is not None: link_properties[EPOCH_SYMBOL] = int(owner_level) link_property_timeout_ms = ( - self._client._config.receive_timeout or self._timeout # pylint:disable=protected-access + self._client._config.receive_timeout + or self._timeout # pylint:disable=protected-access ) * self._amqp_transport.TIMEOUT_FACTOR link_properties[TIMEOUT_SYMBOL] = int(link_property_timeout_ms) - self._link_properties = self._amqp_transport.create_link_properties(link_properties) - self._handler: Optional[uamqp_ReceiveClient] = None + self._link_properties: Union[ + Dict[uamqp_AMQPType, uamqp_AMQPType], Dict[types.AMQPTypes, types.AMQPTypes] + ] = self._amqp_transport.create_link_properties(link_properties) + self._handler: Optional[Union[uamqp_ReceiveClient, ReceiveClient]] = None self._track_last_enqueued_event_properties = ( track_last_enqueued_event_properties ) self._message_buffer: Deque[uamqp_Message] = deque() self._last_received_event: Optional[EventData] = None - self._receive_start_time: Optional[float]= None + self._receive_start_time: Optional[float] = None - def _create_handler(self, auth: uamqp_JWTTokenAuth) -> None: + def _create_handler(self, auth: Union[uamqp_JWTTokenAuth, JWTTokenAuth]) -> None: source = self._amqp_transport.create_source( self._source, self._offset, - event_position_selector(self._offset, self._offset_inclusive) + event_position_selector(self._offset, self._offset_inclusive), + ) + desired_capabilities = ( + [RECEIVER_RUNTIME_METRIC_SYMBOL] + if self._track_last_enqueued_event_properties + else None ) - desired_capabilities = [RECEIVER_RUNTIME_METRIC_SYMBOL] if self._track_last_enqueued_event_properties else None self._handler = self._amqp_transport.create_receive_client( - config=self._client._config, # pylint:disable=protected-access + config=self._client._config, # pylint:disable=protected-access source=source, auth=auth, network_trace=self._client._config.network_tracing, # pylint:disable=protected-access @@ -137,7 +163,8 @@ def _create_handler(self, auth: uamqp_JWTTokenAuth) -> None: keep_alive_interval=self._keep_alive, client_name=self._name, properties=create_properties( - self._client._config.user_agent, amqp_transport=self._amqp_transport # pylint:disable=protected-access + self._client._config.user_agent, # pylint:disable=protected-access + amqp_transport=self._amqp_transport, ), desired_capabilities=desired_capabilities, streaming_receive=True, @@ -147,7 +174,7 @@ def _create_handler(self, auth: uamqp_JWTTokenAuth) -> None: def _open_with_retry(self) -> None: self._do_retryable_operation(self._open, operation_need_param=False) - def _message_received(self, message: uamqp_Message) -> None: + def _message_received(self, message: Union[uamqp_Message, Message]) -> None: # pylint:disable=protected-access self._message_buffer.append(message) @@ -159,8 +186,7 @@ def _next_message_in_buffer(self): return event_data def _open(self) -> bool: - """Open the EventHubConsumer/EventHubProducer using the supplied connection. - """ + """Open the EventHubConsumer/EventHubProducer using the supplied connection.""" # pylint: disable=protected-access if not self.running: if self._handler: @@ -170,12 +196,15 @@ def _open(self) -> bool: conn = self._client._conn_manager.get_connection( # pylint: disable=protected-access host=self._client._address.hostname, auth=auth ) + self._handler = cast("ReceiveClient", self._handler) self._handler.open(connection=conn) - while not self._handler.client_ready(): - time.sleep(0.05) - self.handler_ready = True + self.handler_ready = False self.running = True + if not self.handler_ready: + if self._handler.client_ready(): # type: ignore + self.handler_ready = True + return self.handler_ready def receive(self, batch=False, max_batch_size=300, max_wait_time=None): @@ -184,10 +213,11 @@ def receive(self, batch=False, max_batch_size=300, max_wait_time=None): self._client._config.max_retries # pylint:disable=protected-access ) self._receive_start_time = self._receive_start_time or time.time() - deadline = self._receive_start_time + ( - max_wait_time or 0 - ) + deadline = self._receive_start_time + (max_wait_time or 0) if len(self._message_buffer) < max_batch_size: + # TODO: the retry here is a bit tricky as we are using low-level api from the amqp client. + # Currently we create a new client with the latest received event's offset per retry. + # Ideally we should reuse the same client reestablishing the connection/link with the latest offset. while retried_times <= max_retries: try: if self._open(): @@ -195,11 +225,13 @@ def receive(self, batch=False, max_batch_size=300, max_wait_time=None): break except Exception as exception: # pylint: disable=broad-except self._amqp_transport.check_link_stolen(self, exception) + # TODO: below block hangs when retry_total > 0 + # need to remove/refactor, issue #27137 if not self.running: # exit by close return if self._last_received_event: self._offset = self._last_received_event.offset - last_exception = self._handle_exception(exception) + last_exception = self._handle_exception(exception, is_consumer=True) retried_times += 1 if retried_times > max_retries: _LOGGER.info( diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_consumer_client.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_consumer_client.py index 5b48324bbe05..07309007369d 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_consumer_client.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_consumer_client.py @@ -126,6 +126,8 @@ class EventHubConsumerClient( :keyword str connection_verify: Path to the custom CA_BUNDLE file of the SSL certificate which is used to authenticate the identity of the connection endpoint. Default is None in which case `certifi.where()` will be used. + :keyword bool uamqp_transport: Whether to use the `uamqp` library as the underlying transport. The default value is + False and the Pure Python AMQP library will be used as the underlying transport. .. admonition:: Example: diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_eventprocessor/_eventprocessor_mixin.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_eventprocessor/_eventprocessor_mixin.py index 37a9b69d41c9..77945e502032 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_eventprocessor/_eventprocessor_mixin.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_eventprocessor/_eventprocessor_mixin.py @@ -3,6 +3,7 @@ # Licensed under the MIT License. See License.txt in the project root for license information. # -------------------------------------------------------------------------------------------- +from __future__ import annotations from datetime import datetime from contextlib import contextmanager from typing import ( @@ -27,8 +28,10 @@ from azure.core.tracing import AbstractSpan from .._common import EventData from .._consumer import EventHubConsumer + from ..aio._consumer_async import ( + EventHubConsumer as EventHubConsumerAsync + ) from .._consumer_client import EventHubConsumerClient - from ..aio._consumer_async import EventHubConsumer as EventHubConsumerAsync from ..aio._consumer_client_async import ( EventHubConsumerClient as EventHubConsumerClientAsync, ) @@ -36,9 +39,7 @@ class EventProcessorMixin(object): - _eventhub_client = ( - None - ) # type: Optional[Union[EventHubConsumerClient, EventHubConsumerClientAsync]] + _eventhub_client: Optional[Union[EventHubConsumerClient, EventHubConsumerClientAsync]] = None _consumer_group = "" # type: str _owner_level = None # type: Optional[int] _prefetch = None # type: Optional[int] @@ -78,7 +79,7 @@ def create_consumer( initial_event_position, # type: Union[str, int, datetime] initial_event_position_inclusive, # type: bool on_event_received, # type: Callable[[Union[Optional[EventData], List[EventData]]], None] - **kwargs # type: Any + **kwargs, # type: Any ): # type: (...) -> Union[EventHubConsumer, EventHubConsumerAsync] consumer = self._eventhub_client._create_consumer( # type: ignore # pylint: disable=protected-access @@ -90,7 +91,7 @@ def create_consumer( owner_level=self._owner_level, track_last_enqueued_event_properties=self._track_last_enqueued_event_properties, prefetch=self._prefetch, - **kwargs + **kwargs, ) return consumer diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_eventprocessor/event_processor.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_eventprocessor/event_processor.py index 029fb2de4684..2e167f4032e8 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_eventprocessor/event_processor.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_eventprocessor/event_processor.py @@ -34,6 +34,7 @@ from .._consumer import EventHubConsumer from .._consumer_client import EventHubConsumerClient + _LOGGER = logging.getLogger(__name__) @@ -347,6 +348,8 @@ def _do_receive(self, partition_id, consumer): error, ) self._process_error(self._partition_contexts[partition_id], error) + # TODO: close consumer if non-retryable. issue #27137 + # Does OWNERSHIP_LOST make sense for all errors? self._close_consumer(partition_id, consumer, CloseReason.OWNERSHIP_LOST) def start(self): diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_mixin.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_mixin.py index fa00d0ff8874..ff25ea4083d4 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_mixin.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_mixin.py @@ -8,7 +8,6 @@ Optional, ) - class DictMixin(object): def __setitem__(self, key, item): # type: (Any, Any) -> None diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_producer.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_producer.py index 5b42d964400e..66826a0fccf9 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_producer.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_producer.py @@ -15,6 +15,7 @@ AnyStr, List, TYPE_CHECKING, + cast ) # pylint: disable=unused-import from ._common import EventData, EventDataBatch @@ -33,8 +34,16 @@ if TYPE_CHECKING: from azure.core.tracing import AbstractSpan - from uamqp import constants as uamqp_constants, SendClient as uamqp_SendClient - from uamqp.authentication import JWTTokenAuth as uamqp_JWTTokenAuth + try: + from uamqp import SendClient as uamqp_SendClient + from uamqp.constants import MessageSendResult as uamqp_MessageSendResult + from uamqp.authentication import JWTTokenAuth as uamqp_JWTTokenAuth + except ImportError: + uamqp_MessageSendResult = None + uamqp_SendClient = None + uamqp_JWTTokenAuth = None + from ._pyamqp.client import SendClient + from ._pyamqp.authentication import JWTTokenAuth from ._transport._base import AmqpTransport from ._producer_client import EventHubProducerClient @@ -120,8 +129,8 @@ def __init__( if partition: self._target += "/Partitions/" + partition self._name += f"-partition{partition}" - self._handler: Optional[uamqp_SendClient] = None - self._outcome: Optional[uamqp_constants.MessageSendResult] = None + self._handler: Optional[Union[uamqp_SendClient, SendClient]] = None + self._outcome: Optional[uamqp_MessageSendResult] = None self._condition: Optional[Exception] = None self._lock = threading.Lock() self._link_properties = self._amqp_transport.create_link_properties( @@ -129,7 +138,7 @@ def __init__( ) def _create_handler( - self, auth: uamqp_JWTTokenAuth + self, auth: Union[uamqp_JWTTokenAuth, JWTTokenAuth] ) -> None: self._handler = self._amqp_transport.create_send_client( config=self._client._config, # pylint:disable=protected-access @@ -145,7 +154,7 @@ def _create_handler( self._client._config.user_agent, # pylint: disable=protected-access amqp_transport=self._amqp_transport, ), - msg_timeout=self._timeout * 1000, + msg_timeout=self._timeout * self._amqp_transport.TIMEOUT_FACTOR, ) def _open_with_retry(self) -> None: @@ -153,11 +162,11 @@ def _open_with_retry(self) -> None: def _on_outcome( self, - outcome: "uamqp_constants.MessageSendResult", + outcome: uamqp_MessageSendResult, condition: Optional[Exception], ) -> None: """ - Called when the outcome is received for a delivery. + ONLY USED FOR uamqp_transport=True. Called when the outcome is received for a delivery. :param outcome: The outcome of the message delivery - success or failure. :type outcome: ~uamqp.constants.MessageSendResult @@ -202,6 +211,16 @@ def _wrap_eventdata( ): # The partition_key in the param will be omitted. if not event_data: return event_data + # If AmqpTransports are not the same, create batch with correct BatchMessage. + if self._amqp_transport.TIMEOUT_FACTOR != event_data._amqp_transport.TIMEOUT_FACTOR: # pylint: disable=protected-access + # pylint: disable=protected-access + event_data = EventDataBatch._from_batch( + event_data._internal_events, + amqp_transport=self._amqp_transport, + partition_key=cast(AnyStr, event_data._partition_key), + partition_id=event_data._partition_id, + max_size_in_bytes=event_data.max_size_in_bytes + ) if ( partition_key and partition_key diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_producer_client.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_producer_client.py index efebf97f5fe0..656b68042957 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_producer_client.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_producer_client.py @@ -121,6 +121,8 @@ class EventHubProducerClient( :keyword str connection_verify: Path to the custom CA_BUNDLE file of the SSL certificate which is used to authenticate the identity of the connection endpoint. Default is None in which case `certifi.where()` will be used. + :keyword bool uamqp_transport: Whether to use the `uamqp` library as the underlying transport. The default value is + False and the Pure Python AMQP library will be used as the underlying transport. .. admonition:: Example: @@ -247,7 +249,8 @@ def _buffered_send(self, events, **kwargs): self._max_message_size_on_link, max_wait_time=self._max_wait_time, max_buffer_length=self._max_buffer_length, - executor=self._executor + executor=self._executor, + amqp_transport=self._amqp_transport, ) self._buffered_producer_dispatcher.enqueue_events(events, **kwargs) @@ -266,13 +269,13 @@ def _batch_preparer(self, event_data_batch, **kwargs): to_send_batch = self.create_batch( partition_id=partition_id, partition_key=partition_key ) - to_send_batch._load_events( # pylint:disable=protected-access + to_send_batch._load_events( # pylint:disable=protected-access event_data_batch ) return ( to_send_batch, - to_send_batch._partition_id, # pylint:disable=protected-access + to_send_batch._partition_id, # pylint:disable=protected-access partition_key, ) @@ -307,6 +310,7 @@ def _buffered_send_event(self, event, **kwargs): def _get_partitions(self): # type: () -> None if not self._partition_ids: + _LOGGER.debug("Populating partition IDs so producers can be started.") self._partition_ids = self.get_partition_ids() # type: ignore for p_id in cast(List[str], self._partition_ids): self._producers[p_id] = None @@ -668,7 +672,12 @@ def send_batch(self, event_data_batch, **kwargs): ) if self._on_success: self._on_success(batch._internal_events, pid) - except (KeyError, AttributeError, EventHubError): + except (KeyError, AttributeError, EventHubError) as e: + _LOGGER.debug( + "Producer for partition ID %s not available: %s. Rebuilding new producer.", + partition_id, + e, + ) self._start_producer(partition_id, send_timeout) cast(EventHubProducer, self._producers[partition_id]).send( batch, partition_key=pkey, timeout=send_timeout @@ -727,7 +736,8 @@ def create_batch(self, **kwargs): event_data_batch = EventDataBatch( max_size_in_bytes=(max_size_in_bytes or self._max_message_size_on_link), partition_id=partition_id, - partition_key=partition_key + partition_key=partition_key, + amqp_transport=self._amqp_transport, ) return event_data_batch diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/__init__.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/__init__.py new file mode 100644 index 000000000000..fc9544449266 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/__init__.py @@ -0,0 +1,21 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# ------------------------------------------------------------------------- + +__version__ = "2.0.0a1" + + +from ._connection import Connection +from ._transport import SSLTransport + +from .client import AMQPClient, ReceiveClient, SendClient + +__all__ = [ + "Connection", + "SSLTransport", + "AMQPClient", + "ReceiveClient", + "SendClient", +] diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_connection.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_connection.py new file mode 100644 index 000000000000..4412e9f58733 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_connection.py @@ -0,0 +1,854 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + +import uuid +import logging +import time +from urllib.parse import urlparse +import socket +from ssl import SSLError +from typing import Any, Tuple, Optional, NamedTuple, Union, cast + +from ._transport import Transport +from .sasl import SASLTransport, SASLWithWebSocket +from .session import Session +from .performatives import OpenFrame, CloseFrame +from .constants import ( + PORT, + SECURE_PORT, + WEBSOCKET_PORT, + MAX_CHANNELS, + MAX_FRAME_SIZE_BYTES, + HEADER_FRAME, + ConnectionState, + EMPTY_FRAME, + TransportType, +) + +from .error import ErrorCondition, AMQPConnectionError, AMQPError + +_LOGGER = logging.getLogger(__name__) +_CLOSING_STATES = ( + ConnectionState.OC_PIPE, + ConnectionState.CLOSE_PIPE, + ConnectionState.DISCARDING, + ConnectionState.CLOSE_SENT, + ConnectionState.END, +) + + +def get_local_timeout(now, idle_timeout, last_frame_received_time): + # type: (float, float, float) -> bool + """Check whether the local timeout has been reached since a new incoming frame was received. + + :param float now: The current time to check against. + :rtype: bool + :returns: Whether to shutdown the connection due to timeout. + """ + if idle_timeout and last_frame_received_time: + time_since_last_received = now - last_frame_received_time + return time_since_last_received > idle_timeout + return False + + +class Connection(object): # pylint:disable=too-many-instance-attributes + """An AMQP Connection. + + :ivar str state: The connection state. + :param str endpoint: The endpoint to connect to. Must be fully qualified with scheme and port number. + :keyword str container_id: The ID of the source container. If not set a GUID will be generated. + :keyword int max_frame_size: Proposed maximum frame size in bytes. Default value is 64kb. + :keyword int channel_max: The maximum channel number that may be used on the Connection. Default value is 65535. + :keyword int idle_timeout: Connection idle time-out in seconds. + :keyword list(str) outgoing_locales: Locales available for outgoing text. + :keyword list(str) incoming_locales: Desired locales for incoming text in decreasing level of preference. + :keyword list(str) offered_capabilities: The extension capabilities the sender supports. + :keyword list(str) desired_capabilities: The extension capabilities the sender may use if the receiver supports + :keyword dict properties: Connection properties. + :keyword bool allow_pipelined_open: Allow frames to be sent on the connection before a response Open frame + has been received. Default value is `True`. + :keyword float idle_timeout_empty_frame_send_ratio: Portion of the idle timeout time to wait before sending an + empty frame. The default portion is 50% of the idle timeout value (i.e. `0.5`). + :keyword float idle_wait_time: The time in seconds to sleep while waiting for a response from the endpoint. + Default value is `0.1`. + :keyword bool network_trace: Whether to log the network traffic. Default value is `False`. If enabled, frames + will be logged at the logging.INFO level. + :keyword str transport_type: Determines if the transport type is Amqp or AmqpOverWebSocket. + Defaults to TransportType.Amqp. It will be AmqpOverWebSocket if using http_proxy. + :keyword Dict http_proxy: HTTP proxy settings. This must be a dictionary with the following + keys: `'proxy_hostname'` (str value) and `'proxy_port'` (int value). When using these settings, + the transport_type would be AmqpOverWebSocket. + Additionally the following keys may also be present: `'username', 'password'`. + """ + + def __init__(self, endpoint, **kwargs): # pylint:disable=too-many-statements + # type(str, Any) -> None + parsed_url = urlparse(endpoint) + self._hostname = parsed_url.hostname + endpoint = self._hostname + if parsed_url.port: + self._port = parsed_url.port + elif parsed_url.scheme == "amqps": + self._port = SECURE_PORT + else: + self._port = PORT + self.state = None # type: Optional[ConnectionState] + + # Custom Endpoint + custom_endpoint_address = kwargs.get("custom_endpoint_address") + custom_endpoint = None + if custom_endpoint_address: + custom_parsed_url = urlparse(custom_endpoint_address) + custom_port = custom_parsed_url.port or WEBSOCKET_PORT + custom_endpoint = f"{custom_parsed_url.hostname}:{custom_port}{custom_parsed_url.path}" + self._container_id = kwargs.pop("container_id", None) or str(uuid.uuid4()) # type: str + self._network_trace = kwargs.get("network_trace", False) + self._network_trace_params = {"amqpConnection": self._container_id, "amqpSession": None, "amqpLink": None} + + transport = kwargs.get("transport") + self._transport_type = kwargs.pop("transport_type", TransportType.Amqp) + if transport: + self._transport = transport + elif "sasl_credential" in kwargs: + sasl_transport = SASLTransport + if self._transport_type.name == "AmqpOverWebsocket" or kwargs.get("http_proxy"): + sasl_transport = SASLWithWebSocket + endpoint = parsed_url.hostname + parsed_url.path + self._transport = sasl_transport( + host=endpoint, + credential=kwargs["sasl_credential"], + custom_endpoint=custom_endpoint, + network_trace_params=self._network_trace_params, + **kwargs + ) + else: + self._transport = Transport( + parsed_url.netloc, + transport_type=self._transport_type, + network_trace_params=self._network_trace_params, + **kwargs) + self._max_frame_size = kwargs.pop("max_frame_size", MAX_FRAME_SIZE_BYTES) # type: int + self._remote_max_frame_size = None # type: Optional[int] + self._channel_max = kwargs.pop("channel_max", MAX_CHANNELS) # type: int + self._idle_timeout = kwargs.pop("idle_timeout", None) # type: Optional[int] + self._outgoing_locales = kwargs.pop("outgoing_locales", None) # type: Optional[List[str]] + self._incoming_locales = kwargs.pop("incoming_locales", None) # type: Optional[List[str]] + self._offered_capabilities = None # type: Optional[str] + self._desired_capabilities = kwargs.pop("desired_capabilities", None) # type: Optional[str] + self._properties = kwargs.pop("properties", None) # type: Optional[Dict[str, str]] + + self._allow_pipelined_open = kwargs.pop("allow_pipelined_open", True) # type: bool + self._remote_idle_timeout = None # type: Optional[int] + self._remote_idle_timeout_send_frame = None # type: Optional[int] + self._idle_timeout_empty_frame_send_ratio = kwargs.get("idle_timeout_empty_frame_send_ratio", 0.5) + self._last_frame_received_time = None # type: Optional[float] + self._last_frame_sent_time = None # type: Optional[float] + self._idle_wait_time = kwargs.get("idle_wait_time", 0.1) # type: float + self._error = None + self._outgoing_endpoints = {} # type: Dict[int, Session] + self._incoming_endpoints = {} # type: Dict[int, Session] + + def __enter__(self): + self.open() + return self + + def __exit__(self, *args): + self.close() + + def _set_state(self, new_state): + # type: (ConnectionState) -> None + """Update the connection state.""" + if new_state is None: + return + previous_state = self.state + self.state = new_state + _LOGGER.info( + "Connection state changed: %r -> %r", + previous_state, + new_state, + extra=self._network_trace_params + ) + for session in self._outgoing_endpoints.values(): + session._on_connection_state_change() # pylint:disable=protected-access + + def _connect(self): + # type: () -> None + """Initiate the connection. + + If `allow_pipelined_open` is enabled, the incoming response header will be processed immediately + and the state on exiting will be HDR_EXCH. Otherwise, the function will return before waiting for + the response header and the final state will be HDR_SENT. + + :raises ValueError: If a reciprocating protocol header is not received during negotiation. + """ + try: + if not self.state: + self._transport.connect() + self._set_state(ConnectionState.START) + self._transport.negotiate() + self._outgoing_header() + self._set_state(ConnectionState.HDR_SENT) + if not self._allow_pipelined_open: + # TODO: List/tuple expected as variable args + self._read_frame(wait=True) + if self.state != ConnectionState.HDR_EXCH: + self._disconnect() + raise ValueError("Did not receive reciprocal protocol header. Disconnecting.") + else: + self._set_state(ConnectionState.HDR_SENT) + except (OSError, IOError, SSLError, socket.error) as exc: + # FileNotFoundError is being raised for exception parity with uamqp when invalid + # `connection_verify` file path is passed in. Remove later when resolving issue #27128. + if isinstance(exc, FileNotFoundError) and exc.filename and "ca_certs" in exc.filename: + raise + raise AMQPConnectionError( + ErrorCondition.SocketError, + description="Failed to initiate the connection due to exception: " + str(exc), + error=exc, + ) + + def _disconnect(self): + # type: () -> None + """Disconnect the transport and set state to END.""" + if self.state == ConnectionState.END: + return + self._set_state(ConnectionState.END) + self._transport.close() + + def _can_read(self): + # type: () -> bool + """Whether the connection is in a state where it is legal to read for incoming frames.""" + return self.state not in (ConnectionState.CLOSE_RCVD, ConnectionState.END) + + def _read_frame( + self, wait: Union[bool, float] = True, **kwargs: Any + ) -> bool: + """Read an incoming frame from the transport. + + :param Union[bool, float] wait: Whether to block on the socket while waiting for an incoming frame. + The default value is `False`, where the frame will block for the configured timeout only (0.1 seconds). + If set to `True`, socket will block indefinitely. If set to a timeout value in seconds, the socket will + block for at most that value. + :rtype: Tuple[int, Optional[Tuple[int, NamedTuple]]] + :returns: A tuple with the incoming channel number, and the frame in the form or a tuple of performative + descriptor and field values. + """ + if wait is False: + new_frame = self._transport.receive_frame(**kwargs) + elif wait is True: + with self._transport.block(): + new_frame = self._transport.receive_frame(**kwargs) + else: + with self._transport.block_with_timeout(timeout=wait): + new_frame = self._transport.receive_frame(**kwargs) + return self._process_incoming_frame(*new_frame) + + def _can_write(self): + # type: () -> bool + """Whether the connection is in a state where it is legal to write outgoing frames.""" + return self.state not in _CLOSING_STATES + + def _send_frame(self, channel, frame, timeout=None, **kwargs): + # type: (int, NamedTuple, Optional[int], Any) -> None + """Send a frame over the connection. + + :param int channel: The outgoing channel number. + :param NamedTuple: The outgoing frame. + :param int timeout: An optional timeout value to wait until the socket is ready to send the frame. + :rtype: None + """ + try: + raise self._error + except TypeError: + pass + + if self._can_write(): + try: + self._last_frame_sent_time = time.time() + if timeout: + with self._transport.block_with_timeout(timeout): + self._transport.send_frame(channel, frame, **kwargs) + else: + self._transport.send_frame(channel, frame, **kwargs) + except (OSError, IOError, SSLError, socket.error) as exc: + self._error = AMQPConnectionError( + ErrorCondition.SocketError, + description="Can not send frame out due to exception: " + str(exc), + error=exc, + ) + except Exception: # pylint:disable=try-except-raise + raise + else: + _LOGGER.info("Cannot write frame in current state: %r", self.state, extra=self._network_trace_params) + + def _get_next_outgoing_channel(self): + # type: () -> int + """Get the next available outgoing channel number within the max channel limit. + + :raises ValueError: If maximum channels has been reached. + :returns: The next available outgoing channel number. + :rtype: int + """ + if (len(self._incoming_endpoints) + len(self._outgoing_endpoints)) >= self._channel_max: + raise ValueError("Maximum number of channels ({}) has been reached.".format(self._channel_max)) + next_channel = next(i for i in range(1, self._channel_max) if i not in self._outgoing_endpoints) + return next_channel + + def _outgoing_empty(self): + # type: () -> None + """Send an empty frame to prevent the connection from reaching an idle timeout.""" + if self._network_trace: + _LOGGER.debug("-> EmptyFrame()", extra=self._network_trace_params) + try: + raise self._error + except TypeError: + pass + try: + if self._can_write(): + self._transport.write(EMPTY_FRAME) + self._last_frame_sent_time = time.time() + except (OSError, IOError, SSLError, socket.error) as exc: + self._error = AMQPConnectionError( + ErrorCondition.SocketError, + description="Can not send empty frame due to exception: " + str(exc), + error=exc, + ) + except Exception: # pylint:disable=try-except-raise + raise + + def _outgoing_header(self): + # type: () -> None + """Send the AMQP protocol header to initiate the connection.""" + self._last_frame_sent_time = time.time() + if self._network_trace: + _LOGGER.debug("-> Header(%r)", HEADER_FRAME, extra=self._network_trace_params) + self._transport.write(HEADER_FRAME) + + def _incoming_header(self, _, frame): + # type: (int, bytes) -> None + """Process an incoming AMQP protocol header and update the connection state.""" + if self._network_trace: + _LOGGER.debug("<- Header(%r)", frame, extra=self._network_trace_params) + if self.state == ConnectionState.START: + self._set_state(ConnectionState.HDR_RCVD) + elif self.state == ConnectionState.HDR_SENT: + self._set_state(ConnectionState.HDR_EXCH) + elif self.state == ConnectionState.OPEN_PIPE: + self._set_state(ConnectionState.OPEN_SENT) + + def _outgoing_open(self): + # type: () -> None + """Send an Open frame to negotiate the AMQP connection functionality.""" + open_frame = OpenFrame( + container_id=self._container_id, + hostname=self._hostname, + max_frame_size=self._max_frame_size, + channel_max=self._channel_max, + idle_timeout=self._idle_timeout * 1000 if self._idle_timeout else None, # Convert to milliseconds + outgoing_locales=self._outgoing_locales, + incoming_locales=self._incoming_locales, + offered_capabilities=self._offered_capabilities if self.state == ConnectionState.OPEN_RCVD else None, + desired_capabilities=self._desired_capabilities if self.state == ConnectionState.HDR_EXCH else None, + properties=self._properties, + ) + if self._network_trace: + _LOGGER.debug("-> %r", open_frame, extra=self._network_trace_params) + self._send_frame(0, open_frame) + + def _incoming_open(self, channel, frame): + # type: (int, Tuple[Any, ...]) -> None + """Process incoming Open frame to finish the connection negotiation. + + The incoming frame format is:: + + - frame[0]: container_id (str) + - frame[1]: hostname (str) + - frame[2]: max_frame_size (int) + - frame[3]: channel_max (int) + - frame[4]: idle_timeout (Optional[int]) + - frame[5]: outgoing_locales (Optional[List[bytes]]) + - frame[6]: incoming_locales (Optional[List[bytes]]) + - frame[7]: offered_capabilities (Optional[List[bytes]]) + - frame[8]: desired_capabilities (Optional[List[bytes]]) + - frame[9]: properties (Optional[Dict[bytes, bytes]]) + + :param int channel: The incoming channel number. + :param frame: The incoming Open frame. + :type frame: Tuple[Any, ...] + :rtype: None + """ + # TODO: Add type hints for full frame tuple contents. + if self._network_trace: + _LOGGER.debug("<- %r", OpenFrame(*frame), extra=self._network_trace_params) + if channel != 0: + _LOGGER.error("OPEN frame received on a channel that is not 0.", extra=self._network_trace_params) + self.close( + error=AMQPError( + condition=ErrorCondition.NotAllowed, description="OPEN frame received on a channel that is not 0." + ) + ) + self._set_state(ConnectionState.END) + if self.state == ConnectionState.OPENED: + _LOGGER.error("OPEN frame received in the OPENED state.", extra=self._network_trace_params) + self.close() + if frame[4]: + self._remote_idle_timeout = frame[4] / 1000 # Convert to seconds + self._remote_idle_timeout_send_frame = ( + self._idle_timeout_empty_frame_send_ratio * self._remote_idle_timeout + ) + + if frame[2] < 512: + # Max frame size is less than supported minimum. + # If any of the values in the received open frame are invalid then the connection shall be closed. + # The error amqp:invalid-field shall be set in the error.condition field of the CLOSE frame. + self.close( + error=AMQPError( + condition=ErrorCondition.InvalidField, + description="Failed parsing OPEN frame: Max frame size is less than supported minimum.", + ) + ) + _LOGGER.error( + "Failed parsing OPEN frame: Max frame size is less than supported minimum.", + extra=self._network_trace_params + ) + return + self._remote_max_frame_size = frame[2] + if self.state == ConnectionState.OPEN_SENT: + self._set_state(ConnectionState.OPENED) + elif self.state == ConnectionState.HDR_EXCH: + self._set_state(ConnectionState.OPEN_RCVD) + self._outgoing_open() + self._set_state(ConnectionState.OPENED) + else: + self.close( + error=AMQPError( + condition=ErrorCondition.IllegalState, + description=f"connection is an illegal state: {self.state}", + ) + ) + _LOGGER.error("Connection is an illegal state: %r", self.state, extra=self._network_trace_params) + + def _outgoing_close(self, error=None): + # type: (Optional[AMQPError]) -> None + """Send a Close frame to shutdown connection with optional error information.""" + close_frame = CloseFrame(error=error) + if self._network_trace: + _LOGGER.debug("-> %r", close_frame, extra=self._network_trace_params) + self._send_frame(0, close_frame) + + def _incoming_close(self, channel, frame): + # type: (int, Tuple[Any, ...]) -> None + """Process incoming Open frame to finish the connection negotiation. + + The incoming frame format is:: + + - frame[0]: error (Optional[AMQPError]) + + """ + if self._network_trace: + _LOGGER.debug("<- %r", CloseFrame(*frame), extra=self._network_trace_params) + disconnect_states = [ + ConnectionState.HDR_RCVD, + ConnectionState.HDR_EXCH, + ConnectionState.OPEN_RCVD, + ConnectionState.CLOSE_SENT, + ConnectionState.DISCARDING, + ] + if self.state in disconnect_states: + self._disconnect() + return + + close_error = None + if channel > self._channel_max: + _LOGGER.error( + "CLOSE frame received on a channel greated than support max.", + extra=self._network_trace_params + ) + close_error = AMQPError(condition=ErrorCondition.InvalidField, description="Invalid channel", info=None) + + self._set_state(ConnectionState.CLOSE_RCVD) + self._outgoing_close(error=close_error) + self._disconnect() + + if frame[0]: + self._error = AMQPConnectionError( + condition=frame[0][0], description=frame[0][1], info=frame[0][2] + ) + _LOGGER.error( + "Connection closed with error: %r", frame[0], + extra=self._network_trace_params + ) + + + def _incoming_begin(self, channel, frame): + # type: (int, Tuple[Any, ...]) -> None + """Process incoming Begin frame to finish negotiating a new session. + + The incoming frame format is:: + + - frame[0]: remote_channel (int) + - frame[1]: next_outgoing_id (int) + - frame[2]: incoming_window (int) + - frame[3]: outgoing_window (int) + - frame[4]: handle_max (int) + - frame[5]: offered_capabilities (Optional[List[bytes]]) + - frame[6]: desired_capabilities (Optional[List[bytes]]) + - frame[7]: properties (Optional[Dict[bytes, bytes]]) + + :param int channel: The incoming channel number. + :param frame: The incoming Begin frame. + :type frame: Tuple[Any, ...] + :rtype: None + """ + try: + existing_session = self._outgoing_endpoints[frame[0]] + self._incoming_endpoints[channel] = existing_session + self._incoming_endpoints[channel]._incoming_begin( # pylint:disable=protected-access + frame + ) + except KeyError: + new_session = Session.from_incoming_frame(self, channel) + self._incoming_endpoints[channel] = new_session + new_session._incoming_begin(frame) # pylint:disable=protected-access + + def _incoming_end(self, channel, frame): + # type: (int, Tuple[Any, ...]) -> None + """Process incoming End frame to close a session. + + The incoming frame format is:: + + - frame[0]: error (Optional[AMQPError]) + + :param int channel: The incoming channel number. + :param frame: The incoming End frame. + :type frame: Tuple[Any, ...] + :rtype: None + """ + try: + self._incoming_endpoints[channel]._incoming_end(frame) # pylint:disable=protected-access + self._incoming_endpoints.pop(channel) + self._outgoing_endpoints.pop(channel) + except KeyError: + #close the connection + self.close( + error=AMQPError( + condition=ErrorCondition.ConnectionCloseForced, + description="Invalid channel number received" + )) + _LOGGER.error( + "END frame received on invalid channel. Closing connection.", + extra=self._network_trace_params + ) + return + + def _process_incoming_frame(self, channel, frame): # pylint:disable=too-many-return-statements + # type: (int, Optional[Union[bytes, Tuple[int, Tuple[Any, ...]]]]) -> bool + """Process an incoming frame, either directly or by passing to the necessary Session. + + :param int channel: The channel the frame arrived on. + :param frame: A tuple containing the performative descriptor and the field values of the frame. + This parameter can be None in the case of an empty frame or a socket timeout. + :type frame: Optional[Tuple[int, NamedTuple]] + :rtype: bool + :returns: A boolean to indicate whether more frames in a batch can be processed or whether the + incoming frame has altered the state. If `True` is returned, the state has changed and the batch + should be interrupted. + """ + try: + performative, fields = cast(Union[bytes, Tuple], frame) + except TypeError: + return True # Empty Frame or socket timeout + fields = cast(Tuple[Any, ...], fields) + try: + self._last_frame_received_time = time.time() + if performative == 20: + self._incoming_endpoints[channel]._incoming_transfer( # pylint:disable=protected-access + fields + ) + return False + if performative == 21: + self._incoming_endpoints[channel]._incoming_disposition( # pylint:disable=protected-access + fields + ) + return False + if performative == 19: + self._incoming_endpoints[channel]._incoming_flow( # pylint:disable=protected-access + fields + ) + return False + if performative == 18: + self._incoming_endpoints[channel]._incoming_attach( # pylint:disable=protected-access + fields + ) + return False + if performative == 22: + self._incoming_endpoints[channel]._incoming_detach( # pylint:disable=protected-access + fields + ) + return True + if performative == 17: + self._incoming_begin(channel, fields) + return True + if performative == 23: + self._incoming_end(channel, fields) + return True + if performative == 16: + self._incoming_open(channel, fields) + return True + if performative == 24: + self._incoming_close(channel, fields) + return True + if performative == 0: + self._incoming_header(channel, cast(bytes, fields)) + return True + if performative == 1: + return False + _LOGGER.error("Unrecognized incoming frame: %r", frame, extra=self._network_trace_params) + return True + except KeyError: + return True # TODO: channel error + + def _process_outgoing_frame(self, channel, frame): + # type: (int, NamedTuple) -> None + """Send an outgoing frame if the connection is in a legal state. + + :raises ValueError: If the connection is not open or not in a valid state. + """ + if not self._allow_pipelined_open and self.state in [ + ConnectionState.OPEN_PIPE, + ConnectionState.OPEN_SENT, + ]: + raise ValueError("Connection not configured to allow pipeline send.") + if self.state not in [ + ConnectionState.OPEN_PIPE, + ConnectionState.OPEN_SENT, + ConnectionState.OPENED, + ]: + raise ValueError("Connection not open.") + now = time.time() + if get_local_timeout( + now, + cast(float, self._idle_timeout), + cast(float, self._last_frame_received_time), + ) or self._get_remote_timeout(now): + _LOGGER.info( + "No frame received for the idle timeout. Closing connection.", + extra=self._network_trace_params + ) + self.close( + error=AMQPError( + condition=ErrorCondition.ConnectionCloseForced, + description="No frame received for the idle timeout.", + ), + wait=False, + ) + return + self._send_frame(channel, frame) + + def _get_remote_timeout(self, now): + # type: (float) -> bool + """Check whether the local connection has reached the remote endpoints idle timeout since + the last outgoing frame was sent. + + If the time since the last since frame is greater than the allowed idle interval, an Empty + frame will be sent to maintain the connection. + + :param float now: The current time to check against. + :rtype: bool + :returns: Whether the local connection should be shutdown due to timeout. + """ + if self._remote_idle_timeout and self._last_frame_sent_time: + time_since_last_sent = now - self._last_frame_sent_time + if time_since_last_sent > cast(int, self._remote_idle_timeout_send_frame): + self._outgoing_empty() + return False + + def _wait_for_response(self, wait, end_state): + # type: (Union[bool, float], ConnectionState) -> None + """Wait for an incoming frame to be processed that will result in a desired state change. + + :param wait: Whether to wait for an incoming frame to be processed. Can be set to `True` to wait + indefinitely, or an int to wait for a specified amount of time (in seconds). To not wait, set to `False`. + :type wait: bool or float + :param ConnectionState end_state: The desired end state to wait until. + :rtype: None + """ + if wait is True: + self.listen(wait=False) + while self.state != end_state: + time.sleep(self._idle_wait_time) + self.listen(wait=False) + elif wait: + self.listen(wait=False) + timeout = time.time() + wait + while self.state != end_state: + if time.time() >= timeout: + break + time.sleep(self._idle_wait_time) + self.listen(wait=False) + + def listen(self, wait=False, batch=1, **kwargs): + # type: (Union[float, int, bool], int, Any) -> None + """Listen on the socket for incoming frames and process them. + + :param wait: Whether to block on the socket until a frame arrives. If set to `True`, socket will + block indefinitely. Alternatively, if set to a time in seconds, the socket will block for at most + the specified timeout. Default value is `False`, where the socket will block for its configured read + timeout (by default 0.1 seconds). + :type wait: int or float or bool + :param int batch: The number of frames to attempt to read and process before returning. The default value + is 1, i.e. process frames one-at-a-time. A higher value should only be used when a receiver is established + and is processing incoming Transfer frames. + :rtype: None + """ + try: + raise self._error + except TypeError: + pass + try: + if self.state not in _CLOSING_STATES: + now = time.time() + if get_local_timeout( + now, + cast(float, self._idle_timeout), + cast(float, self._last_frame_received_time), + ) or self._get_remote_timeout( + now + ): + _LOGGER.info( + "No frame received for the idle timeout. Closing connection.", + extra=self._network_trace_params + ) + self.close( + error=AMQPError( + condition=ErrorCondition.ConnectionCloseForced, + description="No frame received for the idle timeout.", + ), + wait=False, + ) + return + if self.state == ConnectionState.END: + self._error = AMQPConnectionError( + condition=ErrorCondition.ConnectionCloseForced, description="Connection was already closed." + ) + return + for _ in range(batch): + if self._can_read(): + if self._read_frame(wait=wait, **kwargs): + break + else: + _LOGGER.info( + "Connection cannot read frames in this state: %r", + self.state, + extra=self._network_trace_params + ) + break + except (OSError, IOError, SSLError, socket.error) as exc: + self._error = AMQPConnectionError( + ErrorCondition.SocketError, + description="Can not read frame due to exception: " + str(exc), + error=exc, + ) + except Exception: # pylint:disable=try-except-raise + raise + + def create_session(self, **kwargs): + # type: (Any) -> Session + """Create a new session within this connection. + + :keyword str name: The name of the connection. If not set a GUID will be generated. + :keyword int next_outgoing_id: The transfer-id of the first transfer id the sender will send. + Default value is 0. + :keyword int incoming_window: The initial incoming-window of the Session. Default value is 1. + :keyword int outgoing_window: The initial outgoing-window of the Session. Default value is 1. + :keyword int handle_max: The maximum handle value that may be used on the session. Default value is 4294967295. + :keyword list(str) offered_capabilities: The extension capabilities the session supports. + :keyword list(str) desired_capabilities: The extension capabilities the session may use if + the endpoint supports it. + :keyword dict properties: Session properties. + :keyword bool allow_pipelined_open: Allow frames to be sent on the connection before a response Open frame + has been received. Default value is that configured for the connection. + :keyword float idle_wait_time: The time in seconds to sleep while waiting for a response from the endpoint. + Default value is that configured for the connection. + :keyword bool network_trace: Whether to log the network traffic of this session. If enabled, frames + will be logged at the logging.INFO level. Default value is that configured for the connection. + """ + assigned_channel = self._get_next_outgoing_channel() + kwargs["allow_pipelined_open"] = self._allow_pipelined_open + kwargs["idle_wait_time"] = self._idle_wait_time + session = Session( + self, + assigned_channel, + network_trace=kwargs.pop("network_trace", self._network_trace), + network_trace_params=dict(self._network_trace_params), + **kwargs, + ) + self._outgoing_endpoints[assigned_channel] = session + return session + + def open(self, wait=False): + # type: (bool) -> None + """Send an Open frame to start the connection. + + Alternatively, this will be called on entering a Connection context manager. + + :param bool wait: Whether to wait to receive an Open response from the endpoint. Default is `False`. + :raises ValueError: If `wait` is set to `False` and `allow_pipelined_open` is disabled. + :rtype: None + """ + self._connect() + self._outgoing_open() + if self.state == ConnectionState.HDR_EXCH: + self._set_state(ConnectionState.OPEN_SENT) + elif self.state == ConnectionState.HDR_SENT: + self._set_state(ConnectionState.OPEN_PIPE) + if wait: + self._wait_for_response(wait, ConnectionState.OPENED) + elif not self._allow_pipelined_open: + raise ValueError( + "Connection has been configured to not allow piplined-open. Please set 'wait' parameter." + ) + + + def close(self, error=None, wait=False): + # type: (Optional[AMQPError], bool) -> None + """Close the connection and disconnect the transport. + + Alternatively this method will be called on exiting a Connection context manager. + + :param ~uamqp.AMQPError error: Optional error information to include in the close request. + :param bool wait: Whether to wait for a service Close response. Default is `False`. + :rtype: None + """ + try: + if self.state in [ + ConnectionState.END, + ConnectionState.CLOSE_SENT, + ConnectionState.DISCARDING, + ]: + return + self._outgoing_close(error=error) + if error: + self._error = AMQPConnectionError( + condition=error.condition, + description=error.description, + info=error.info, + ) + if self.state == ConnectionState.OPEN_PIPE: + self._set_state(ConnectionState.OC_PIPE) + elif self.state == ConnectionState.OPEN_SENT: + self._set_state(ConnectionState.CLOSE_PIPE) + elif error: + self._set_state(ConnectionState.DISCARDING) + else: + self._set_state(ConnectionState.CLOSE_SENT) + self._wait_for_response(wait, ConnectionState.END) + except Exception as exc: # pylint:disable=broad-except + # If error happened during closing, ignore the error and set state to END + _LOGGER.info("An error occurred when closing the connection: %r", exc, extra=self._network_trace_params) + self._set_state(ConnectionState.END) + finally: + self._disconnect() diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_decode.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_decode.py new file mode 100644 index 000000000000..099069712865 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_decode.py @@ -0,0 +1,349 @@ +#------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +#-------------------------------------------------------------------------- +# pylint: disable=redefined-builtin, import-error + +import struct +import uuid +import logging +from typing import List, Optional, Tuple, Dict, Callable, Any, cast, Union # pylint: disable=unused-import + + +from .message import Message, Header, Properties + +_LOGGER = logging.getLogger(__name__) +_HEADER_PREFIX = memoryview(b'AMQP') +_COMPOSITES = { + 35: 'received', + 36: 'accepted', + 37: 'rejected', + 38: 'released', + 39: 'modified', +} + +c_unsigned_char = struct.Struct('>B') +c_signed_char = struct.Struct('>b') +c_unsigned_short = struct.Struct('>H') +c_signed_short = struct.Struct('>h') +c_unsigned_int = struct.Struct('>I') +c_signed_int = struct.Struct('>i') +c_unsigned_long = struct.Struct('>L') +c_unsigned_long_long = struct.Struct('>Q') +c_signed_long_long = struct.Struct('>q') +c_float = struct.Struct('>f') +c_double = struct.Struct('>d') + + +def _decode_null(buffer): + # type: (memoryview) -> Tuple[memoryview, None] + return buffer, None + + +def _decode_true(buffer): + # type: (memoryview) -> Tuple[memoryview, bool] + return buffer, True + + +def _decode_false(buffer): + # type: (memoryview) -> Tuple[memoryview, bool] + return buffer, False + + +def _decode_zero(buffer): + # type: (memoryview) -> Tuple[memoryview, int] + return buffer, 0 + + +def _decode_empty(buffer): + # type: (memoryview) -> Tuple[memoryview, List[None]] + return buffer, [] + + +def _decode_boolean(buffer): + # type: (memoryview) -> Tuple[memoryview, bool] + return buffer[1:], buffer[:1] == b'\x01' + + +def _decode_ubyte(buffer): + # type: (memoryview) -> Tuple[memoryview, int] + return buffer[1:], buffer[0] + + +def _decode_ushort(buffer): + # type: (memoryview) -> Tuple[memoryview, int] + return buffer[2:], c_unsigned_short.unpack(buffer[:2])[0] + + +def _decode_uint_small(buffer): + # type: (memoryview) -> Tuple[memoryview, int] + return buffer[1:], buffer[0] + + +def _decode_uint_large(buffer): + # type: (memoryview) -> Tuple[memoryview, int] + return buffer[4:], c_unsigned_int.unpack(buffer[:4])[0] + + +def _decode_ulong_small(buffer): + # type: (memoryview) -> Tuple[memoryview, int] + return buffer[1:], buffer[0] + + +def _decode_ulong_large(buffer): + # type: (memoryview) -> Tuple[memoryview, int] + return buffer[8:], c_unsigned_long_long.unpack(buffer[:8])[0] + + +def _decode_byte(buffer): + # type: (memoryview) -> Tuple[memoryview, int] + return buffer[1:], c_signed_char.unpack(buffer[:1])[0] + + +def _decode_short(buffer): + # type: (memoryview) -> Tuple[memoryview, int] + return buffer[2:], c_signed_short.unpack(buffer[:2])[0] + + +def _decode_int_small(buffer): + # type: (memoryview) -> Tuple[memoryview, int] + return buffer[1:], c_signed_char.unpack(buffer[:1])[0] + + +def _decode_int_large(buffer): + # type: (memoryview) -> Tuple[memoryview, int] + return buffer[4:], c_signed_int.unpack(buffer[:4])[0] + + +def _decode_long_small(buffer): + # type: (memoryview) -> Tuple[memoryview, int] + return buffer[1:], c_signed_char.unpack(buffer[:1])[0] + + +def _decode_long_large(buffer): + # type: (memoryview) -> Tuple[memoryview, int] + return buffer[8:], c_signed_long_long.unpack(buffer[:8])[0] + + +def _decode_float(buffer): + # type: (memoryview) -> Tuple[memoryview, float] + return buffer[4:], c_float.unpack(buffer[:4])[0] + + +def _decode_double(buffer): + # type: (memoryview) -> Tuple[memoryview, float] + return buffer[8:], c_double.unpack(buffer[:8])[0] + + +def _decode_timestamp(buffer): + # type: (memoryview) -> Tuple[memoryview, int] + return buffer[8:], c_signed_long_long.unpack(buffer[:8])[0] + + +def _decode_uuid(buffer): + # type: (memoryview) -> Tuple[memoryview, uuid.UUID] + return buffer[16:], uuid.UUID(bytes=buffer[:16].tobytes()) + + +def _decode_binary_small(buffer): + # type: (memoryview) -> Tuple[memoryview, bytes] + length_index = buffer[0] + 1 + return buffer[length_index:], buffer[1:length_index].tobytes() + + +def _decode_binary_large(buffer): + # type: (memoryview) -> Tuple[memoryview, bytes] + length_index = c_unsigned_long.unpack(buffer[:4])[0] + 4 + return buffer[length_index:], buffer[4:length_index].tobytes() + + +def _decode_list_small(buffer): + # type: (memoryview) -> Tuple[memoryview, List[Any]] + count = buffer[1] + buffer = buffer[2:] + values = [None] * count + for i in range(count): + buffer, values[i] = _DECODE_BY_CONSTRUCTOR[buffer[0]](buffer[1:]) + return buffer, values + + +def _decode_list_large(buffer): + # type: (memoryview) -> Tuple[memoryview, List[Any]] + count = c_unsigned_long.unpack(buffer[4:8])[0] + buffer = buffer[8:] + values = [None] * count + for i in range(count): + buffer, values[i] = _DECODE_BY_CONSTRUCTOR[buffer[0]](buffer[1:]) + return buffer, values + + +def _decode_map_small(buffer): + # type: (memoryview) -> Tuple[memoryview, Dict[Any, Any]] + count = int(buffer[1]/2) + buffer = buffer[2:] + values = {} + for _ in range(count): + buffer, key = _DECODE_BY_CONSTRUCTOR[buffer[0]](buffer[1:]) + buffer, value = _DECODE_BY_CONSTRUCTOR[buffer[0]](buffer[1:]) + values[key] = value + return buffer, values + + +def _decode_map_large(buffer): + # type: (memoryview) -> Tuple[memoryview, Dict[Any, Any]] + count = int(c_unsigned_long.unpack(buffer[4:8])[0]/2) + buffer = buffer[8:] + values = {} + for _ in range(count): + buffer, key = _DECODE_BY_CONSTRUCTOR[buffer[0]](buffer[1:]) + buffer, value = _DECODE_BY_CONSTRUCTOR[buffer[0]](buffer[1:]) + values[key] = value + return buffer, values + + +def _decode_array_small(buffer): + # type: (memoryview) -> Tuple[memoryview, List[Any]] + count = buffer[1] # Ignore first byte (size) and just rely on count + if count: + subconstructor = buffer[2] + buffer = buffer[3:] + values = [None] * count + for i in range(count): + buffer, values[i] = _DECODE_BY_CONSTRUCTOR[subconstructor](buffer) + return buffer, values + return buffer[2:], [] + + +def _decode_array_large(buffer): + # type: (memoryview) -> Tuple[memoryview, List[Any]] + count = c_unsigned_long.unpack(buffer[4:8])[0] + if count: + subconstructor = buffer[8] + buffer = buffer[9:] + values = [None] * count + for i in range(count): + buffer, values[i] = _DECODE_BY_CONSTRUCTOR[subconstructor](buffer) + return buffer, values + return buffer[8:], [] + + +def _decode_described(buffer): + # type: (memoryview) -> Tuple[memoryview, Any] + # TODO: to move the cursor of the buffer to the described value based on size of the + # descriptor without decoding descriptor value + composite_type = buffer[0] + buffer, descriptor = _DECODE_BY_CONSTRUCTOR[composite_type](buffer[1:]) + buffer, value = _DECODE_BY_CONSTRUCTOR[buffer[0]](buffer[1:]) + try: + composite_type = cast(int, _COMPOSITES[descriptor]) + return buffer, {composite_type: value} + except KeyError: + return buffer, value + + +def decode_payload(buffer): + # type: (memoryview) -> Message + message: Dict[str, Union[Properties, Header, Dict, bytes, List]] = {} + while buffer: + # Ignore the first two bytes, they will always be the constructors for + # described type then ulong. + descriptor = buffer[2] + buffer, value = _DECODE_BY_CONSTRUCTOR[buffer[3]](buffer[4:]) + if descriptor == 112: + message["header"] = Header(*value) + elif descriptor == 113: + message["delivery_annotations"] = value + elif descriptor == 114: + message["message_annotations"] = value + elif descriptor == 115: + message["properties"] = Properties(*value) + elif descriptor == 116: + message["application_properties"] = value + elif descriptor == 117: + try: + cast(List, message["data"]).append(value) + except KeyError: + message["data"] = [value] + elif descriptor == 118: + try: + cast(List, message["sequence"]).append(value) + except KeyError: + message["sequence"] = [value] + elif descriptor == 119: + message["value"] = value + elif descriptor == 120: + message["footer"] = value + # TODO: we can possibly swap out the Message construct with a TypedDict + # for both input and output so we get the best of both. + return Message(**message) + + +def decode_frame(data): + # type: (memoryview) -> Tuple[int, List[Any]] + # Ignore the first two bytes, they will always be the constructors for + # described type then ulong. + frame_type = data[2] + compound_list_type = data[3] + if compound_list_type == 0xd0: + # list32 0xd0: data[4:8] is size, data[8:12] is count + count = c_signed_int.unpack(data[8:12])[0] + buffer = data[12:] + else: + # list8 0xc0: data[4] is size, data[5] is count + count = data[5] + buffer = data[6:] + fields: List[Optional[memoryview]] = [None] * count + for i in range(count): + buffer, fields[i] = _DECODE_BY_CONSTRUCTOR[buffer[0]](buffer[1:]) + if frame_type == 20: + fields.append(buffer) + return frame_type, fields + + +def decode_empty_frame(header): + # type: (memoryview) -> Tuple[int, bytes] + if header[0:4] == _HEADER_PREFIX: + return 0, header.tobytes() + if header[5] == 0: + return 1, b"EMPTY" + raise ValueError("Received unrecognized empty frame") + + +_DECODE_BY_CONSTRUCTOR: List[Callable] = cast(List[Callable], [None] * 256) +_DECODE_BY_CONSTRUCTOR[0] = _decode_described +_DECODE_BY_CONSTRUCTOR[64] = _decode_null +_DECODE_BY_CONSTRUCTOR[65] = _decode_true +_DECODE_BY_CONSTRUCTOR[66] = _decode_false +_DECODE_BY_CONSTRUCTOR[67] = _decode_zero +_DECODE_BY_CONSTRUCTOR[68] = _decode_zero +_DECODE_BY_CONSTRUCTOR[69] = _decode_empty +_DECODE_BY_CONSTRUCTOR[80] = _decode_ubyte +_DECODE_BY_CONSTRUCTOR[81] = _decode_byte +_DECODE_BY_CONSTRUCTOR[82] = _decode_uint_small +_DECODE_BY_CONSTRUCTOR[83] = _decode_ulong_small +_DECODE_BY_CONSTRUCTOR[84] = _decode_int_small +_DECODE_BY_CONSTRUCTOR[85] = _decode_long_small +_DECODE_BY_CONSTRUCTOR[86] = _decode_boolean +_DECODE_BY_CONSTRUCTOR[96] = _decode_ushort +_DECODE_BY_CONSTRUCTOR[97] = _decode_short +_DECODE_BY_CONSTRUCTOR[112] = _decode_uint_large +_DECODE_BY_CONSTRUCTOR[113] = _decode_int_large +_DECODE_BY_CONSTRUCTOR[114] = _decode_float +_DECODE_BY_CONSTRUCTOR[128] = _decode_ulong_large +_DECODE_BY_CONSTRUCTOR[129] = _decode_long_large +_DECODE_BY_CONSTRUCTOR[130] = _decode_double +_DECODE_BY_CONSTRUCTOR[131] = _decode_timestamp +_DECODE_BY_CONSTRUCTOR[152] = _decode_uuid +_DECODE_BY_CONSTRUCTOR[160] = _decode_binary_small +_DECODE_BY_CONSTRUCTOR[161] = _decode_binary_small +_DECODE_BY_CONSTRUCTOR[163] = _decode_binary_small +_DECODE_BY_CONSTRUCTOR[176] = _decode_binary_large +_DECODE_BY_CONSTRUCTOR[177] = _decode_binary_large +_DECODE_BY_CONSTRUCTOR[179] = _decode_binary_large +_DECODE_BY_CONSTRUCTOR[192] = _decode_list_small +_DECODE_BY_CONSTRUCTOR[193] = _decode_map_small +_DECODE_BY_CONSTRUCTOR[208] = _decode_list_large +_DECODE_BY_CONSTRUCTOR[209] = _decode_map_large +_DECODE_BY_CONSTRUCTOR[224] = _decode_array_small +_DECODE_BY_CONSTRUCTOR[240] = _decode_array_large diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_encode.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_encode.py new file mode 100644 index 000000000000..24267004c8b1 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_encode.py @@ -0,0 +1,921 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + +# TODO: fix mypy errors for _code/_definition/__defaults__ (issue #26500) +import calendar +import struct +import uuid +from datetime import datetime +from typing import ( + Iterable, + Union, + Tuple, + Dict, + Any, + cast, + Sized, + Optional, + List, + Callable, + TYPE_CHECKING, + Sequence, + Collection, +) + +try: + from typing import TypeAlias # type: ignore +except ImportError: + from typing_extensions import TypeAlias + +import six + +from .types import ( + TYPE, + VALUE, + AMQPTypes, + FieldDefinition, + ObjDefinition, + ConstructorBytes, +) +from .message import Message +from . import performatives + +if TYPE_CHECKING: + from .message import Header, Properties + + Performative: TypeAlias = Union[ + performatives.OpenFrame, + performatives.BeginFrame, + performatives.AttachFrame, + performatives.FlowFrame, + performatives.TransferFrame, + performatives.DispositionFrame, + performatives.DetachFrame, + performatives.EndFrame, + performatives.CloseFrame, + performatives.SASLMechanism, + performatives.SASLInit, + performatives.SASLChallenge, + performatives.SASLResponse, + performatives.SASLOutcome, + Message, + Header, + Properties, + ] + +_FRAME_OFFSET = b"\x02" +_FRAME_TYPE = b"\x00" + + +def _construct(byte, construct): + # type: (bytes, bool) -> bytes + return byte if construct else b"" + + +def encode_null(output, *args, **kwargs): # pylint: disable=unused-argument + # type: (bytearray, Any, Any) -> None + """ + encoding code="0x40" category="fixed" width="0" label="the null value" + """ + output.extend(ConstructorBytes.null) + + +def encode_boolean( + output, value, with_constructor=True, **kwargs # pylint: disable=unused-argument +): + # type: (bytearray, bool, bool, Any) -> None + """ + + + + """ + value = bool(value) + if with_constructor: + output.extend(_construct(ConstructorBytes.bool, with_constructor)) + output.extend(b"\x01" if value else b"\x00") + return + + output.extend(ConstructorBytes.bool_true if value else ConstructorBytes.bool_false) + + +def encode_ubyte( + output, value, with_constructor=True, **kwargs # pylint: disable=unused-argument +): + # type: (bytearray, Union[int, bytes], bool, Any) -> None + """ + + """ + try: + value = int(value) + except ValueError: + value = cast(bytes, value) + value = ord(value) + try: + output.extend(_construct(ConstructorBytes.ubyte, with_constructor)) + output.extend(struct.pack(">B", abs(value))) + except struct.error: + raise ValueError("Unsigned byte value must be 0-255") + + +def encode_ushort( + output, value, with_constructor=True, **kwargs # pylint: disable=unused-argument +): + # type: (bytearray, int, bool, Any) -> None + """ + + """ + value = int(value) + try: + output.extend(_construct(ConstructorBytes.ushort, with_constructor)) + output.extend(struct.pack(">H", abs(value))) + except struct.error: + raise ValueError("Unsigned byte value must be 0-65535") + + +def encode_uint(output, value, with_constructor=True, use_smallest=True): + # type: (bytearray, int, bool, bool) -> None + """ + + + + """ + value = int(value) + if value == 0: + output.extend(ConstructorBytes.uint_0) + return + try: + if use_smallest and value <= 255: + output.extend(_construct(ConstructorBytes.uint_small, with_constructor)) + output.extend(struct.pack(">B", abs(value))) + return + output.extend(_construct(ConstructorBytes.uint_large, with_constructor)) + output.extend(struct.pack(">I", abs(value))) + except struct.error: + raise ValueError("Value supplied for unsigned int invalid: {}".format(value)) + + +def encode_ulong(output, value, with_constructor=True, use_smallest=True): + # type: (bytearray, int, bool, bool) -> None + """ + + + + """ + value = int(value) + if value == 0: + output.extend(ConstructorBytes.ulong_0) + return + try: + if use_smallest and value <= 255: + output.extend(_construct(ConstructorBytes.ulong_small, with_constructor)) + output.extend(struct.pack(">B", abs(value))) + return + output.extend(_construct(ConstructorBytes.ulong_large, with_constructor)) + output.extend(struct.pack(">Q", abs(value))) + except struct.error: + raise ValueError("Value supplied for unsigned long invalid: {}".format(value)) + + +def encode_byte( + output, value, with_constructor=True, **kwargs # pylint: disable=unused-argument +): + # type: (bytearray, int, bool, Any) -> None + """ + + """ + value = int(value) + try: + output.extend(_construct(ConstructorBytes.byte, with_constructor)) + output.extend(struct.pack(">b", value)) + except struct.error: + raise ValueError("Byte value must be -128-127") + + +def encode_short( + output, value, with_constructor=True, **kwargs # pylint: disable=unused-argument +): + # type: (bytearray, int, bool, Any) -> None + """ + + """ + value = int(value) + try: + output.extend(_construct(ConstructorBytes.short, with_constructor)) + output.extend(struct.pack(">h", value)) + except struct.error: + raise ValueError("Short value must be -32768-32767") + + +def encode_int(output, value, with_constructor=True, use_smallest=True): + # type: (bytearray, int, bool, bool) -> None + """ + + + """ + value = int(value) + try: + if use_smallest and (-128 <= value <= 127): + output.extend(_construct(ConstructorBytes.int_small, with_constructor)) + output.extend(struct.pack(">b", value)) + return + output.extend(_construct(ConstructorBytes.int_large, with_constructor)) + output.extend(struct.pack(">i", value)) + except struct.error: + raise ValueError("Value supplied for int invalid: {}".format(value)) + + +def encode_long(output, value, with_constructor=True, use_smallest=True): + # type: (bytearray, int, bool, bool) -> None + """ + + + """ + if isinstance(value, datetime): + value = (calendar.timegm(value.utctimetuple()) * 1000) + ( + value.microsecond / 1000 + ) + value = int(value) + try: + if use_smallest and (-128 <= value <= 127): + output.extend(_construct(ConstructorBytes.long_small, with_constructor)) + output.extend(struct.pack(">b", value)) + return + output.extend(_construct(ConstructorBytes.long_large, with_constructor)) + output.extend(struct.pack(">q", value)) + except struct.error: + raise ValueError("Value supplied for long invalid: {}".format(value)) + + +def encode_float( + output, value, with_constructor=True, **kwargs # pylint: disable=unused-argument +): + # type: (bytearray, float, bool, Any) -> None + """ + + """ + value = float(value) + output.extend(_construct(ConstructorBytes.float, with_constructor)) + output.extend(struct.pack(">f", value)) + + +def encode_double( + output, value, with_constructor=True, **kwargs # pylint: disable=unused-argument +): + # type: (bytearray, float, bool, Any) -> None + """ + + """ + value = float(value) + output.extend(_construct(ConstructorBytes.double, with_constructor)) + output.extend(struct.pack(">d", value)) + + +def encode_timestamp( + output, value, with_constructor=True, **kwargs # pylint: disable=unused-argument +): + # type: (bytearray, Union[int, datetime], bool, Any) -> None + """ + + """ + value = cast(datetime, value) + if isinstance(value, datetime): + value = cast( + int, + (calendar.timegm(value.utctimetuple()) * 1000) + (value.microsecond / 1000), + ) + value = int(cast(int, value)) + output.extend(_construct(ConstructorBytes.timestamp, with_constructor)) + output.extend(struct.pack(">q", value)) + + +def encode_uuid( + output, value, with_constructor=True, **kwargs # pylint: disable=unused-argument +): + # type: (bytearray, Union[uuid.UUID, str, bytes], bool, Any) -> None + """ + + """ + if isinstance(value, six.text_type): + value = uuid.UUID(value).bytes + elif isinstance(value, uuid.UUID): + value = value.bytes + elif isinstance(value, six.binary_type): + value = uuid.UUID(bytes=value).bytes + else: + raise TypeError("Invalid UUID type: {}".format(type(value))) + output.extend(_construct(ConstructorBytes.uuid, with_constructor)) + output.extend(value) + + +def encode_binary(output, value, with_constructor=True, use_smallest=True): + # type: (bytearray, Union[bytes, bytearray], bool, bool) -> None + """ + + + """ + length = len(value) + if use_smallest and length <= 255: + output.extend(_construct(ConstructorBytes.binary_small, with_constructor)) + output.extend(struct.pack(">B", length)) + output.extend(value) + return + try: + output.extend(_construct(ConstructorBytes.binary_large, with_constructor)) + output.extend(struct.pack(">L", length)) + output.extend(value) + except struct.error: + raise ValueError("Binary data to long to encode") + + +def encode_string(output, value, with_constructor=True, use_smallest=True): + # type: (bytearray, Union[bytes, str], bool, bool) -> None + """ + + + """ + if isinstance(value, six.text_type): + value = value.encode("utf-8") + length = len(value) + if use_smallest and length <= 255: + output.extend(_construct(ConstructorBytes.string_small, with_constructor)) + output.extend(struct.pack(">B", length)) + output.extend(value) + return + try: + output.extend(_construct(ConstructorBytes.string_large, with_constructor)) + output.extend(struct.pack(">L", length)) + output.extend(value) + except struct.error: + raise ValueError("String value too long to encode.") + + +def encode_symbol(output, value, with_constructor=True, use_smallest=True): + # type: (bytearray, Union[bytes, str], bool, bool) -> None + """ + + + """ + if isinstance(value, six.text_type): + value = value.encode("utf-8") + length = len(value) + if use_smallest and length <= 255: + output.extend(_construct(ConstructorBytes.symbol_small, with_constructor)) + output.extend(struct.pack(">B", length)) + output.extend(value) + return + try: + output.extend(_construct(ConstructorBytes.symbol_large, with_constructor)) + output.extend(struct.pack(">L", length)) + output.extend(value) + except struct.error: + raise ValueError("Symbol value too long to encode.") + + +def encode_list(output, value, with_constructor=True, use_smallest=True): + # type: (bytearray, Iterable[Any], bool, bool) -> None + """ + + + + """ + count = len(cast(Sized, value)) + if use_smallest and count == 0: + output.extend(ConstructorBytes.list_0) + return + encoded_size = 0 + encoded_values = bytearray() + for item in value: + encode_value(encoded_values, item, with_constructor=True) + encoded_size += len(encoded_values) + if use_smallest and count <= 255 and encoded_size < 255: + output.extend(_construct(ConstructorBytes.list_small, with_constructor)) + output.extend(struct.pack(">B", encoded_size + 1)) + output.extend(struct.pack(">B", count)) + else: + try: + output.extend(_construct(ConstructorBytes.list_large, with_constructor)) + output.extend(struct.pack(">L", encoded_size + 4)) + output.extend(struct.pack(">L", count)) + except struct.error: + raise ValueError("List is too large or too long to be encoded.") + output.extend(encoded_values) + +def encode_map(output, value, with_constructor=True, use_smallest=True): + # type: (bytearray, Union[Dict[Any, Any], Iterable[Tuple[Any, Any]]], bool, bool) -> None + """ + + + """ + count = len(cast(Sized, value)) * 2 + encoded_size = 0 + encoded_values = bytearray() + try: + value = cast(Dict, value) + items = cast(Iterable, value.items()) + except AttributeError: + items = cast(Iterable, value) + for key, data in items: + encode_value(encoded_values, key, with_constructor=True) + encode_value(encoded_values, data, with_constructor=True) + encoded_size = len(encoded_values) + if use_smallest and count <= 255 and encoded_size < 255: + output.extend(_construct(ConstructorBytes.map_small, with_constructor)) + output.extend(struct.pack(">B", encoded_size + 1)) + output.extend(struct.pack(">B", count)) + else: + try: + output.extend(_construct(ConstructorBytes.map_large, with_constructor)) + output.extend(struct.pack(">L", encoded_size + 4)) + output.extend(struct.pack(">L", count)) + except struct.error: + raise ValueError("Map is too large or too long to be encoded.") + output.extend(encoded_values) + + +def _check_element_type(item, element_type): + if not element_type: + try: + return item["TYPE"] + except (KeyError, TypeError): + return type(item) + try: + if item["TYPE"] != element_type: + raise TypeError("All elements in an array must be the same type.") + except (KeyError, TypeError): + if not isinstance(item, element_type): + raise TypeError("All elements in an array must be the same type.") + return element_type + + +def encode_array(output, value, with_constructor=True, use_smallest=True): + # type: (bytearray, Iterable[Any], bool, bool) -> None + """ + + + """ + count = len(cast(Sized, value)) + encoded_size = 0 + encoded_values = bytearray() + first_item = True + element_type = None + for item in value: + element_type = _check_element_type(item, element_type) + encode_value( + encoded_values, item, with_constructor=first_item, use_smallest=False + ) + first_item = False + if item is None: + encoded_size -= 1 + break + encoded_size += len(encoded_values) + if use_smallest and count <= 255 and encoded_size < 255: + output.extend(_construct(ConstructorBytes.array_small, with_constructor)) + output.extend(struct.pack(">B", encoded_size + 1)) + output.extend(struct.pack(">B", count)) + else: + try: + output.extend(_construct(ConstructorBytes.array_large, with_constructor)) + output.extend(struct.pack(">L", encoded_size + 4)) + output.extend(struct.pack(">L", count)) + except struct.error: + raise ValueError("Array is too large or too long to be encoded.") + output.extend(encoded_values) + + +def encode_described(output: bytearray, value: Tuple[Any, Any], _: bool = None, **kwargs: Any) -> None: # type: ignore + output.extend(ConstructorBytes.descriptor) + encode_value(output, value[0], **kwargs) + encode_value(output, value[1], **kwargs) + + +def encode_fields(value): + # type: (Optional[Dict[str, Any]]) -> Dict[str, Any] + """A mapping from field name to value. + + The fields type is a map where the keys are restricted to be of type symbol (this excludes the possibility + of a null key). There is no further restriction implied by the fields type on the allowed values for the + entries or the set of allowed keys. + + + """ + if not value: + return {TYPE: AMQPTypes.null, VALUE: None} + fields = {TYPE: AMQPTypes.map, VALUE: []} + for key, data in value.items(): + if isinstance(key, str): + key = key.encode("utf-8") # type: ignore + cast(List, fields[VALUE]).append(({TYPE: AMQPTypes.symbol, VALUE: key}, data)) + return fields + + +def encode_annotations(value): + # type: (Optional[Dict[str, Any]]) -> Dict[str, Any] + """The annotations type is a map where the keys are restricted to be of type symbol or of type ulong. + + All ulong keys, and all symbolic keys except those beginning with "x-" are reserved. + On receiving an annotations map containing keys or values which it does not recognize, and for which the + key does not begin with the string 'x-opt-' an AMQP container MUST detach the link with the not-implemented + amqp-error. + + + """ + if not value: + return {TYPE: AMQPTypes.null, VALUE: None} + fields = {TYPE: AMQPTypes.map, VALUE: []} + for key, data in value.items(): + if isinstance(key, int): + field_key = {TYPE: AMQPTypes.ulong, VALUE: key} + else: + field_key = {TYPE: AMQPTypes.symbol, VALUE: key} + try: + cast(List, fields[VALUE]).append( + (field_key, {TYPE: data[TYPE], VALUE: data[VALUE]}) + ) + except (KeyError, TypeError): + cast(List, fields[VALUE]).append((field_key, {TYPE: None, VALUE: data})) + return fields + + +def encode_application_properties(value): + # type: (Optional[Dict[str, Any]]) -> Dict[str, Any] + """The application-properties section is a part of the bare message used for structured application data. + + + + + + Intermediaries may use the data within this structure for the purposes of filtering or routing. + The keys of this map are restricted to be of type string (which excludes the possibility of a null key) + and the values are restricted to be of simple types only, that is (excluding map, list, and array types). + """ + if not value: + return {TYPE: AMQPTypes.null, VALUE: None} + fields = {TYPE: AMQPTypes.map, VALUE: cast(List, [])} + for key, data in value.items(): + cast(List, fields[VALUE]).append(({TYPE: AMQPTypes.string, VALUE: key}, data)) + return fields + + +def encode_message_id(value): + # type: (Any) -> Dict[str, Union[int, uuid.UUID, bytes, str]] + """ + + + + + """ + if isinstance(value, int): + return {TYPE: AMQPTypes.ulong, VALUE: value} + if isinstance(value, uuid.UUID): + return {TYPE: AMQPTypes.uuid, VALUE: value} + if isinstance(value, six.binary_type): + return {TYPE: AMQPTypes.binary, VALUE: value} + if isinstance(value, six.text_type): + return {TYPE: AMQPTypes.string, VALUE: value} + raise TypeError("Unsupported Message ID type.") + + +def encode_node_properties(value): + # type: (Optional[Dict[str, Any]]) -> Dict[str, Any] + """Properties of a node. + + + + A symbol-keyed map containing properties of a node used when requesting creation or reporting + the creation of a dynamic node. The following common properties are defined:: + + - `lifetime-policy`: The lifetime of a dynamically generated node. Definitionally, the lifetime will + never be less than the lifetime of the link which caused its creation, however it is possible to extend + the lifetime of dynamically created node using a lifetime policy. The value of this entry MUST be of a type + which provides the lifetime-policy archetype. The following standard lifetime-policies are defined below: + delete-on-close, delete-on-no-links, delete-on-no-messages or delete-on-no-links-or-messages. + + - `supported-dist-modes`: The distribution modes that the node supports. The value of this entry MUST be one or + more symbols which are valid distribution-modes. That is, the value MUST be of the same type as would be valid + in a field defined with the following attributes: + type="symbol" multiple="true" requires="distribution-mode" + """ + if not value: + return {TYPE: AMQPTypes.null, VALUE: None} + # TODO + fields = {TYPE: AMQPTypes.map, VALUE: []} + # fields[{TYPE: AMQPTypes.symbol, VALUE: b'lifetime-policy'}] = { + # TYPE: AMQPTypes.described, + # VALUE: ( + # {TYPE: AMQPTypes.ulong, VALUE: value['lifetime_policy']}, + # {TYPE: AMQPTypes.list, VALUE: []} + # ) + # } + # fields[{TYPE: AMQPTypes.symbol, VALUE: b'supported-dist-modes'}] = {} + return fields + + +def encode_filter_set(value): + # type: (Optional[Dict[str, Any]]) -> Dict[str, Any] + """A set of predicates to filter the Messages admitted onto the Link. + + + + A set of named filters. Every key in the map MUST be of type symbol, every value MUST be either null or of a + described type which provides the archetype filter. A filter acts as a function on a message which returns a + boolean result indicating whether the message can pass through that filter or not. A message will pass + through a filter-set if and only if it passes through each of the named filters. If the value for a given key is + null, this acts as if there were no such key present (i.e., all messages pass through the null filter). + + Filter types are a defined extension point. The filter types that a given source supports will be indicated + by the capabilities of the source. + """ + if not value: + return {TYPE: AMQPTypes.null, VALUE: None} + fields = {TYPE: AMQPTypes.map, VALUE: cast(List, [])} + for name, data in value.items(): + described_filter: Dict[str, Union[Tuple[Dict[str, Any], Any], Optional[str]]] + if data is None: + described_filter = {TYPE: AMQPTypes.null, VALUE: None} + else: + if isinstance(name, str): + name = name.encode("utf-8") # type: ignore + try: + descriptor, filter_value = data + described_filter = { + TYPE: AMQPTypes.described, + VALUE: ({TYPE: AMQPTypes.symbol, VALUE: descriptor}, filter_value), + } + except ValueError: + described_filter = data + + cast(List, fields[VALUE]).append( + ({TYPE: AMQPTypes.symbol, VALUE: name}, described_filter) + ) + return fields + + +def encode_unknown(output, value, **kwargs): + # type: (bytearray, Optional[Any], Any) -> None + """ + Dynamic encoding according to the type of `value`. + """ + if value is None: + encode_null(output, **kwargs) + elif isinstance(value, bool): + encode_boolean(output, value, **kwargs) + elif isinstance(value, six.string_types): + encode_string(output, value, **kwargs) + elif isinstance(value, uuid.UUID): + encode_uuid(output, value, **kwargs) + elif isinstance(value, (bytearray, six.binary_type)): + encode_binary(output, value, **kwargs) + elif isinstance(value, float): + encode_double(output, value, **kwargs) + elif isinstance(value, six.integer_types): + encode_int(output, value, **kwargs) + elif isinstance(value, datetime): + encode_timestamp(output, value, **kwargs) + elif isinstance(value, list): + encode_list(output, value, **kwargs) + elif isinstance(value, tuple): + encode_described(output, cast(Tuple[Any, Any], value), **kwargs) + elif isinstance(value, dict): + encode_map(output, value, **kwargs) + else: + raise TypeError("Unable to encode unknown value: {}".format(value)) + + +_FIELD_DEFINITIONS = { + FieldDefinition.fields: encode_fields, + FieldDefinition.annotations: encode_annotations, + FieldDefinition.message_id: encode_message_id, + FieldDefinition.app_properties: encode_application_properties, + FieldDefinition.node_properties: encode_node_properties, + FieldDefinition.filter_set: encode_filter_set, +} + +_ENCODE_MAP = { + None: encode_unknown, + AMQPTypes.null: encode_null, + AMQPTypes.boolean: encode_boolean, + AMQPTypes.ubyte: encode_ubyte, + AMQPTypes.byte: encode_byte, + AMQPTypes.ushort: encode_ushort, + AMQPTypes.short: encode_short, + AMQPTypes.uint: encode_uint, + AMQPTypes.int: encode_int, + AMQPTypes.ulong: encode_ulong, + AMQPTypes.long: encode_long, + AMQPTypes.float: encode_float, + AMQPTypes.double: encode_double, + AMQPTypes.timestamp: encode_timestamp, + AMQPTypes.uuid: encode_uuid, + AMQPTypes.binary: encode_binary, + AMQPTypes.string: encode_string, + AMQPTypes.symbol: encode_symbol, + AMQPTypes.list: encode_list, + AMQPTypes.map: encode_map, + AMQPTypes.array: encode_array, + AMQPTypes.described: encode_described, +} + + +def encode_value(output, value, **kwargs): + # type: (bytearray, Any, Any) -> None + try: + cast(Callable, _ENCODE_MAP[value[TYPE]])(output, value[VALUE], **kwargs) + except (KeyError, TypeError): + encode_unknown(output, value, **kwargs) + + +def describe_performative(performative): + # type: (Performative) -> Dict[str, Sequence[Collection[str]]] + body: List[Dict[str, Any]] = [] + for index, value in enumerate(performative): + # TODO: fix mypy + field = performative._definition[index] # type: ignore # pylint: disable=protected-access + if value is None: + body.append({TYPE: AMQPTypes.null, VALUE: None}) + elif field is None: + continue + elif isinstance(field.type, FieldDefinition): + if field.multiple: + body.append( + { + TYPE: AMQPTypes.array, + VALUE: [_FIELD_DEFINITIONS[field.type](v) for v in value], + } + ) + else: + body.append(_FIELD_DEFINITIONS[field.type](value)) + elif isinstance(field.type, ObjDefinition): + body.append(describe_performative(value)) + else: + if field.multiple: + body.append( + { + TYPE: AMQPTypes.array, + VALUE: [{TYPE: field.type, VALUE: v} for v in value], + } + ) + else: + body.append({TYPE: field.type, VALUE: value}) + + return { + TYPE: AMQPTypes.described, + VALUE: ( + {TYPE: AMQPTypes.ulong, VALUE: performative._code}, # type: ignore # pylint: disable=protected-access + {TYPE: AMQPTypes.list, VALUE: body}, + ), + } + + +def encode_payload(output, payload): + # type: (bytearray, Message) -> bytes + + if payload[0]: # header + # TODO: Header and Properties encoding can be optimized to + # 1. not encoding trailing None fields + # Possible fix 1: + # header = payload[0] + # header = header[0:max(i for i, v in enumerate(header) if v is not None) + 1] + # Possible fix 2: + # itertools.dropwhile(lambda x: x is None, header[::-1]))[::-1] + # 2. encoding bool without constructor + # Possible fix 3: + # header = list(payload[0]) + # while header[-1] is None: + # del header[-1] + encode_value(output, describe_performative(payload[0])) + + if payload[2]: # message annotations + encode_value( + output, + { + TYPE: AMQPTypes.described, + VALUE: ( + {TYPE: AMQPTypes.ulong, VALUE: 0x00000072}, + encode_annotations(payload[2]), + ), + }, + ) + + if payload[3]: # properties + # TODO: Header and Properties encoding can be optimized to + # 1. not encoding trailing None fields + # 2. encoding bool without constructor + encode_value(output, describe_performative(payload[3])) + + if payload[4]: # application properties + encode_value( + output, + { + TYPE: AMQPTypes.described, + VALUE: ( + {TYPE: AMQPTypes.ulong, VALUE: 0x00000074}, + encode_application_properties(payload[4]), + ), + }, + ) + + if payload[5]: # data + for item_value in payload[5]: + encode_value( + output, + { + TYPE: AMQPTypes.described, + VALUE: ( + {TYPE: AMQPTypes.ulong, VALUE: 0x00000075}, + {TYPE: AMQPTypes.binary, VALUE: item_value}, + ), + }, + ) + + if payload[6]: # sequence + for item_value in payload[6]: + encode_value( + output, + { + TYPE: AMQPTypes.described, + VALUE: ( + {TYPE: AMQPTypes.ulong, VALUE: 0x00000076}, + {TYPE: None, VALUE: item_value}, + ), + }, + ) + + if payload[7]: # value + encode_value( + output, + { + TYPE: AMQPTypes.described, + VALUE: ( + {TYPE: AMQPTypes.ulong, VALUE: 0x00000077}, + {TYPE: None, VALUE: payload[7]}, + ), + }, + ) + + if payload[8]: # footer + encode_value( + output, + { + TYPE: AMQPTypes.described, + VALUE: ( + {TYPE: AMQPTypes.ulong, VALUE: 0x00000078}, + encode_annotations(payload[8]), + ), + }, + ) + + # TODO: + # currently the delivery annotations must be finally encoded instead of being encoded at the 2nd position + # otherwise the event hubs service would ignore the delivery annotations + # -- received message doesn't have it populated + # check with service team? + if payload[1]: # delivery annotations + encode_value( + output, + { + TYPE: AMQPTypes.described, + VALUE: ( + {TYPE: AMQPTypes.ulong, VALUE: 0x00000071}, + encode_annotations(payload[1]), + ), + }, + ) + + return output + + +def encode_frame(frame, frame_type=_FRAME_TYPE): + # type: (Optional[Performative], bytes) -> Tuple[bytes, Optional[bytes]] + # TODO: allow passing type specific bytes manually, e.g. Empty Frame needs padding + if frame is None: + size = 8 + header = size.to_bytes(4, "big") + _FRAME_OFFSET + frame_type + return header, None + + frame_description = describe_performative(frame) + frame_data = bytearray() + encode_value(frame_data, frame_description) + if isinstance(frame, performatives.TransferFrame): + frame_data += frame.payload + + size = len(frame_data) + 8 + header = size.to_bytes(4, "big") + _FRAME_OFFSET + frame_type + return header, frame_data diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_message_backcompat.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_message_backcompat.py new file mode 100644 index 000000000000..0e3c22213eda --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_message_backcompat.py @@ -0,0 +1,250 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + +# pylint: disable=too-many-lines +from typing import Callable, cast +from enum import Enum + +from ._encode import encode_payload +from .utils import get_message_encoded_size +from .error import AMQPError +from .message import Header, Properties + + +def _encode_property(value): + try: + return value.encode("UTF-8") + except AttributeError: + return value + + +class MessageState(Enum): + WaitingToBeSent = 0 + WaitingForSendAck = 1 + SendComplete = 2 + SendFailed = 3 + ReceivedUnsettled = 4 + ReceivedSettled = 5 + + def __eq__(self, __o: object) -> bool: + try: + return self.value == cast(Enum, __o).value + except AttributeError: + return super().__eq__(__o) + + +class MessageAlreadySettled(Exception): + pass + + +DONE_STATES = (MessageState.SendComplete, MessageState.SendFailed) +RECEIVE_STATES = (MessageState.ReceivedSettled, MessageState.ReceivedUnsettled) +PENDING_STATES = (MessageState.WaitingForSendAck, MessageState.WaitingToBeSent) + + +class LegacyMessage(object): # pylint: disable=too-many-instance-attributes + def __init__(self, message, **kwargs): + self._message = message + self.state = MessageState.SendComplete + self.idle_time = 0 + self.retries = 0 + self._settler = kwargs.get("settler") + self._encoding = kwargs.get("encoding") + self.delivery_no = kwargs.get("delivery_no") + self.delivery_tag = kwargs.get("delivery_tag") or None + self.on_send_complete = None + self.properties = ( + LegacyMessageProperties(self._message.properties) + if self._message.properties + else None + ) + self.application_properties = ( + self._message.application_properties + if any(self._message.application_properties) + else None + ) + self.annotations = ( + self._message.annotations if any(self._message.annotations) else None + ) + self.header = ( + LegacyMessageHeader(self._message.header) if self._message.header else None + ) + self.footer = self._message.footer + self.delivery_annotations = self._message.delivery_annotations + if self._settler: + self.state = MessageState.ReceivedUnsettled + elif self.delivery_no: + self.state = MessageState.ReceivedSettled + self._to_outgoing_amqp_message: Callable = kwargs.get( + "to_outgoing_amqp_message" + ) + + def __str__(self): + return str(self._message) + + def _can_settle_message(self): + if self.state not in RECEIVE_STATES: + raise TypeError("Only received messages can be settled.") + if self.settled: + return False + return True + + @property + def settled(self): + if self.state == MessageState.ReceivedUnsettled: + return False + return True + + def get_message_encoded_size(self): + return get_message_encoded_size(self._to_outgoing_amqp_message(self._message)) + + def encode_message(self): + output = bytearray() + encode_payload(output, self._to_outgoing_amqp_message(self._message)) + return bytes(output) + + def get_data(self): + return self._message.body + + def gather(self): + if self.state in RECEIVE_STATES: + raise TypeError("Only new messages can be gathered.") + if not self._message: + raise ValueError("Message data already consumed.") + if self.state in DONE_STATES: + raise MessageAlreadySettled() + return [self] + + def get_message(self): + return self._to_outgoing_amqp_message(self._message) + + def accept(self): + if self._can_settle_message(): + self._settler.settle_messages(self.delivery_no, "accepted") + self.state = MessageState.ReceivedSettled + return True + return False + + def reject(self, condition=None, description=None, info=None): + if self._can_settle_message(): + self._settler.settle_messages( + self.delivery_no, + "rejected", + error=AMQPError( + condition=condition, description=description, info=info + ), + ) + self.state = MessageState.ReceivedSettled + return True + return False + + def release(self): + if self._can_settle_message(): + self._settler.settle_messages(self.delivery_no, "released") + self.state = MessageState.ReceivedSettled + return True + return False + + def modify(self, failed, deliverable, annotations=None): + if self._can_settle_message(): + self._settler.settle_messages( + self.delivery_no, + "modified", + delivery_failed=failed, + undeliverable_here=deliverable, + message_annotations=annotations, + ) + self.state = MessageState.ReceivedSettled + return True + return False + + +class LegacyBatchMessage(LegacyMessage): + batch_format = 0x80013700 + max_message_length = 1024 * 1024 + size_offset = 0 + + +class LegacyMessageProperties(object): # pylint: disable=too-many-instance-attributes + def __init__(self, properties): + self.message_id = _encode_property(properties.message_id) + self.user_id = _encode_property(properties.user_id) + self.to = _encode_property(properties.to) + self.subject = _encode_property(properties.subject) + self.reply_to = _encode_property(properties.reply_to) + self.correlation_id = _encode_property(properties.correlation_id) + self.content_type = _encode_property(properties.content_type) + self.content_encoding = _encode_property(properties.content_encoding) + self.absolute_expiry_time = properties.absolute_expiry_time + self.creation_time = properties.creation_time + self.group_id = _encode_property(properties.group_id) + self.group_sequence = properties.group_sequence + self.reply_to_group_id = _encode_property(properties.reply_to_group_id) + + def __str__(self): + return str( + { + "message_id": self.message_id, + "user_id": self.user_id, + "to": self.to, + "subject": self.subject, + "reply_to": self.reply_to, + "correlation_id": self.correlation_id, + "content_type": self.content_type, + "content_encoding": self.content_encoding, + "absolute_expiry_time": self.absolute_expiry_time, + "creation_time": self.creation_time, + "group_id": self.group_id, + "group_sequence": self.group_sequence, + "reply_to_group_id": self.reply_to_group_id, + } + ) + + def get_properties_obj(self): + return Properties( + self.message_id, + self.user_id, + self.to, + self.subject, + self.reply_to, + self.correlation_id, + self.content_type, + self.content_encoding, + self.absolute_expiry_time, + self.creation_time, + self.group_id, + self.group_sequence, + self.reply_to_group_id, + ) + + +class LegacyMessageHeader(object): + def __init__(self, header): + self.delivery_count = header.delivery_count or 0 + self.time_to_live = header.time_to_live + self.first_acquirer = header.first_acquirer + self.durable = header.durable + self.priority = header.priority + + def __str__(self): + return str( + { + "delivery_count": self.delivery_count, + "time_to_live": self.time_to_live, + "first_acquirer": self.first_acquirer, + "durable": self.durable, + "priority": self.priority, + } + ) + + def get_header_obj(self): + return Header( + self.durable, + self.priority, + self.time_to_live, + self.first_acquirer, + self.delivery_count, + ) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_platform.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_platform.py new file mode 100644 index 000000000000..18d91f710041 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_platform.py @@ -0,0 +1,107 @@ +"""Platform compatibility.""" +# pylint: skip-file + +from __future__ import absolute_import, unicode_literals + +from typing import Tuple, cast +import platform +import re +import struct +import sys + +# Jython does not have this attribute +try: + from socket import SOL_TCP +except ImportError: # pragma: no cover + from socket import IPPROTO_TCP as SOL_TCP # noqa + + +RE_NUM = re.compile(r'(\d+).+') + + +def _linux_version_to_tuple(s): + # type: (str) -> Tuple[int, int, int] + return cast(Tuple[int, int, int], tuple(map(_versionatom, s.split('.')[:3]))) + + +def _versionatom(s): + # type: (str) -> int + if s.isdigit(): + return int(s) + match = RE_NUM.match(s) + return int(match.groups()[0]) if match else 0 + + +# available socket options for TCP level +KNOWN_TCP_OPTS = { + 'TCP_CORK', 'TCP_DEFER_ACCEPT', 'TCP_KEEPCNT', + 'TCP_KEEPIDLE', 'TCP_KEEPINTVL', 'TCP_LINGER2', + 'TCP_MAXSEG', 'TCP_NODELAY', 'TCP_QUICKACK', + 'TCP_SYNCNT', 'TCP_USER_TIMEOUT', 'TCP_WINDOW_CLAMP', +} + +LINUX_VERSION = None +if sys.platform.startswith('linux'): + LINUX_VERSION = _linux_version_to_tuple(platform.release()) + if LINUX_VERSION < (2, 6, 37): + KNOWN_TCP_OPTS.remove('TCP_USER_TIMEOUT') + + # Windows Subsystem for Linux is an edge-case: the Python socket library + # returns most TCP_* enums, but they aren't actually supported + if platform.release().endswith("Microsoft"): + KNOWN_TCP_OPTS = {'TCP_NODELAY', 'TCP_KEEPIDLE', 'TCP_KEEPINTVL', + 'TCP_KEEPCNT'} + +elif sys.platform.startswith('darwin'): + KNOWN_TCP_OPTS.remove('TCP_USER_TIMEOUT') + +elif 'bsd' in sys.platform: + KNOWN_TCP_OPTS.remove('TCP_USER_TIMEOUT') + +# According to MSDN Windows platforms support getsockopt(TCP_MAXSSEG) but not +# setsockopt(TCP_MAXSEG) on IPPROTO_TCP sockets. +elif sys.platform.startswith('win'): + KNOWN_TCP_OPTS = {'TCP_NODELAY'} + +elif sys.platform.startswith('cygwin'): + KNOWN_TCP_OPTS = {'TCP_NODELAY'} + +# illumos does not allow to set the TCP_MAXSEG socket option, +# even if the Oracle documentation says otherwise. +elif sys.platform.startswith('sunos'): + KNOWN_TCP_OPTS.remove('TCP_MAXSEG') + +# aix does not allow to set the TCP_MAXSEG +# or the TCP_USER_TIMEOUT socket options. +elif sys.platform.startswith('aix'): + KNOWN_TCP_OPTS.remove('TCP_MAXSEG') + KNOWN_TCP_OPTS.remove('TCP_USER_TIMEOUT') + +if sys.version_info < (2, 7, 7): # pragma: no cover + import functools + + def _to_bytes_arg(fun): + @functools.wraps(fun) + def _inner(s, *args, **kwargs): + return fun(s.encode(), *args, **kwargs) + return _inner + + pack = _to_bytes_arg(struct.pack) + pack_into = _to_bytes_arg(struct.pack_into) + unpack = _to_bytes_arg(struct.unpack) + unpack_from = _to_bytes_arg(struct.unpack_from) +else: + pack = struct.pack + pack_into = struct.pack_into + unpack = struct.unpack + unpack_from = struct.unpack_from + +__all__ = [ + 'LINUX_VERSION', + 'SOL_TCP', + 'KNOWN_TCP_OPTS', + 'pack', + 'pack_into', + 'unpack', + 'unpack_from', +] diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_transport.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_transport.py new file mode 100644 index 000000000000..570c8b5c0110 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_transport.py @@ -0,0 +1,805 @@ +# ------------------------------------------------------------------------- # pylint: disable=file-needs-copyright-header +# This is a fork of the transport.py which was originally written by Barry Pederson and +# maintained by the Celery project: https://github.com/celery/py-amqp. +# +# Copyright (C) 2009 Barry Pederson +# +# The license text can also be found here: +# http://www.opensource.org/licenses/BSD-3-Clause +# +# License +# ======= +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, +# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS +# BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF +# THE POSSIBILITY OF SUCH DAMAGE. +# ------------------------------------------------------------------------- + + +from __future__ import absolute_import, unicode_literals + +import errno +import re +import socket +import ssl +import struct +from ssl import SSLError +from contextlib import contextmanager +from io import BytesIO +import logging +from threading import Lock + +import certifi + +from ._platform import KNOWN_TCP_OPTS, SOL_TCP +from ._encode import encode_frame +from ._decode import decode_frame, decode_empty_frame +from .constants import ( + TLS_HEADER_FRAME, + WEBSOCKET_PORT, + TransportType, + AMQP_WS_SUBPROTOCOL, +) +from .error import AuthenticationException, ErrorCondition + + +try: + import fcntl +except ImportError: # pragma: no cover + fcntl = None # type: ignore # noqa + +def set_cloexec(fd, cloexec): # noqa + """Set flag to close fd after exec.""" + if fcntl is None: + return + try: + FD_CLOEXEC = fcntl.FD_CLOEXEC + except AttributeError: + raise NotImplementedError( + "close-on-exec flag not supported on this platform", + ) + flags = fcntl.fcntl(fd, fcntl.F_GETFD) + if cloexec: + flags |= FD_CLOEXEC + else: + flags &= ~FD_CLOEXEC + return fcntl.fcntl(fd, fcntl.F_SETFD, flags) + + +_LOGGER = logging.getLogger(__name__) +_UNAVAIL = {errno.EAGAIN, errno.EINTR, errno.ENOENT, errno.EWOULDBLOCK} + +AMQP_PORT = 5672 +AMQPS_PORT = 5671 +AMQP_FRAME = memoryview(b"AMQP") +EMPTY_BUFFER = bytes() +SIGNED_INT_MAX = 0x7FFFFFFF +TIMEOUT_INTERVAL = 1 +WS_TIMEOUT_INTERVAL = 1 +READ_TIMEOUT_INTERVAL = 0.2 + +# Match things like: [fe80::1]:5432, from RFC 2732 +IPV6_LITERAL = re.compile(r"\[([\.0-9a-f:]+)\](?::(\d+))?") + +DEFAULT_SOCKET_SETTINGS = { + "TCP_NODELAY": 1, + "TCP_USER_TIMEOUT": 1000, + "TCP_KEEPIDLE": 60, + "TCP_KEEPINTVL": 10, + "TCP_KEEPCNT": 9, +} + + +def get_errno(exc): + """Get exception errno (if set). + + Notes: + :exc:`socket.error` and :exc:`IOError` first got + the ``.errno`` attribute in Py2.7. + """ + try: + return exc.errno + except AttributeError: + try: + # e.args = (errno, reason) + if isinstance(exc.args, tuple) and len(exc.args) == 2: + return exc.args[0] + except AttributeError: + pass + return 0 + + +# TODO: fails when host = hostname:port/path. fix +def to_host_port(host, port=AMQP_PORT): + """Convert hostname:port string to host, port tuple.""" + m = IPV6_LITERAL.match(host) + if m: + host = m.group(1) + if m.group(2): + port = int(m.group(2)) + else: + if ":" in host: + host, port = host.rsplit(":", 1) + port = int(port) + return host, port + + +class UnexpectedFrame(Exception): + pass + + +class _AbstractTransport(object): # pylint: disable=too-many-instance-attributes + """Common superclass for TCP and SSL transports.""" + + def __init__( + self, + host, + *, + port=AMQP_PORT, + connect_timeout=None, + read_timeout=None, + socket_settings=None, + raise_on_initial_eintr=True, + **kwargs + ): + self._quick_recv = None + self.connected = False + self.sock = None + self.raise_on_initial_eintr = raise_on_initial_eintr + self._read_buffer = BytesIO() + self.host, self.port = to_host_port(host, port) + self.network_trace_params = kwargs.get('network_trace_params') + + self.connect_timeout = connect_timeout or TIMEOUT_INTERVAL + self.read_timeout = read_timeout or READ_TIMEOUT_INTERVAL + self.socket_settings = socket_settings + self.socket_lock = Lock() + + def connect(self): + try: + # are we already connected? + if self.connected: + return + self._connect(self.host, self.port, self.connect_timeout) + self._init_socket( + self.socket_settings, + self.read_timeout, + ) + # we've sent the banner; signal connect + # EINTR, EAGAIN, EWOULDBLOCK would signal that the banner + # has _not_ been sent + self.connected = True + except (OSError, IOError, SSLError) as e: + _LOGGER.info("Transport connection failed: %r", e, extra=self.network_trace_params) + # if not fully connected, close socket, and reraise error + if self.sock and not self.connected: + self.sock.close() + self.sock = None + raise + + @contextmanager + def block_with_timeout(self, timeout): + if timeout is None: + yield self.sock + else: + sock = self.sock + prev = sock.gettimeout() + if prev != timeout: + sock.settimeout(timeout) + try: + yield self.sock + except SSLError as exc: + if "timed out" in str(exc): + # http://bugs.python.org/issue10272 + raise socket.timeout() + if "The operation did not complete" in str(exc): + # Non-blocking SSL sockets can throw SSLError + raise socket.timeout() + raise + except socket.error as exc: + if get_errno(exc) == errno.EWOULDBLOCK: + raise socket.timeout() + raise + finally: + if timeout != prev: + sock.settimeout(prev) + + @contextmanager + def block(self): + bocking_timeout = None + sock = self.sock + prev = sock.gettimeout() + if prev != bocking_timeout: + sock.settimeout(bocking_timeout) + try: + yield self.sock + except SSLError as exc: + if "timed out" in str(exc): + # http://bugs.python.org/issue10272 + raise socket.timeout() + if "The operation did not complete" in str(exc): + # Non-blocking SSL sockets can throw SSLError + raise socket.timeout() + raise + except socket.error as exc: + if get_errno(exc) == errno.EWOULDBLOCK: + raise socket.timeout() + raise + finally: + if bocking_timeout != prev: + sock.settimeout(prev) + + @contextmanager + def non_blocking(self): + non_bocking_timeout = 0.0 + sock = self.sock + prev = sock.gettimeout() + if prev != non_bocking_timeout: + sock.settimeout(non_bocking_timeout) + try: + yield self.sock + except SSLError as exc: + if "timed out" in str(exc): + # http://bugs.python.org/issue10272 + raise socket.timeout() + if "The operation did not complete" in str(exc): + # Non-blocking SSL sockets can throw SSLError + raise socket.timeout() + raise + except socket.error as exc: + if get_errno(exc) == errno.EWOULDBLOCK: + raise socket.timeout() + raise + finally: + if non_bocking_timeout != prev: + sock.settimeout(prev) + + def _connect(self, host, port, timeout): + e = None + + # Below we are trying to avoid additional DNS requests for AAAA if A + # succeeds. This helps a lot in case when a hostname has an IPv4 entry + # in /etc/hosts but not IPv6. Without the (arguably somewhat twisted) + # logic below, getaddrinfo would attempt to resolve the hostname for + # both IP versions, which would make the resolver talk to configured + # DNS servers. If those servers are for some reason not available + # during resolution attempt (either because of system misconfiguration, + # or network connectivity problem), resolution process locks the + # _connect call for extended time. + addr_types = (socket.AF_INET, socket.AF_INET6) + addr_types_num = len(addr_types) + for n, family in enumerate(addr_types): + # first, resolve the address for a single address family + try: + entries = socket.getaddrinfo( + host, port, family, socket.SOCK_STREAM, SOL_TCP + ) + entries_num = len(entries) + except socket.gaierror: + # we may have depleted all our options + if n + 1 >= addr_types_num: + # if getaddrinfo succeeded before for another address + # family, reraise the previous socket.error since it's more + # relevant to users + raise e if e is not None else socket.error( + "failed to resolve broker hostname" + ) + continue # pragma: no cover + + # now that we have address(es) for the hostname, connect to broker + for i, res in enumerate(entries): + af, socktype, proto, _, sa = res + try: + self.sock = socket.socket(af, socktype, proto) + try: + set_cloexec(self.sock, True) + except NotImplementedError: + pass + self.sock.settimeout(timeout) + self.sock.connect(sa) + except socket.error as ex: + e = ex + if self.sock is not None: + self.sock.close() + self.sock = None + # we may have depleted all our options + if i + 1 >= entries_num and n + 1 >= addr_types_num: + raise + else: + # hurray, we established connection + return + + def _init_socket(self, socket_settings, read_timeout): + self.sock.settimeout(None) # set socket back to blocking mode + self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) + self._set_socket_options(socket_settings) + + # set socket timeouts + # for timeout, interval in ((socket.SO_SNDTIMEO, write_timeout), + # (socket.SO_RCVTIMEO, read_timeout)): + # if interval is not None: + # sec = int(interval) + # usec = int((interval - sec) * 1000000) + # self.sock.setsockopt( + # socket.SOL_SOCKET, timeout, + # pack('ll', sec, usec), + # ) + self._setup_transport() + # TODO: a greater timeout value is needed in long distance communication + # we should either figure out a reasonable value error/dynamically adjust the timeout + # 0.2 second is enough for perf analysis + self.sock.settimeout(read_timeout) # set socket back to non-blocking mode + + def _get_tcp_socket_defaults(self, sock): # pylint: disable=no-self-use + tcp_opts = {} + for opt in KNOWN_TCP_OPTS: + enum = None + if opt == "TCP_USER_TIMEOUT": + try: + from socket import TCP_USER_TIMEOUT as enum + except ImportError: + # should be in Python 3.6+ on Linux. + enum = 18 + elif hasattr(socket, opt): + enum = getattr(socket, opt) + + if enum: + if opt in DEFAULT_SOCKET_SETTINGS: + tcp_opts[enum] = DEFAULT_SOCKET_SETTINGS[opt] + elif hasattr(socket, opt): + tcp_opts[enum] = sock.getsockopt(SOL_TCP, getattr(socket, opt)) + return tcp_opts + + def _set_socket_options(self, socket_settings): + tcp_opts = self._get_tcp_socket_defaults(self.sock) + if socket_settings: + tcp_opts.update(socket_settings) + for opt, val in tcp_opts.items(): + self.sock.setsockopt(SOL_TCP, opt, val) + + def _read(self, n, initial=False, buffer=None, _errnos=None): + """Read exactly n bytes from the peer.""" + raise NotImplementedError("Must be overriden in subclass") + + def _setup_transport(self): + """Do any additional initialization of the class.""" + + def _shutdown_transport(self): + """Do any preliminary work in shutting down the connection.""" + + def _write(self, s): + """Completely write a string to the peer.""" + raise NotImplementedError("Must be overriden in subclass") + + def close(self): + with self.socket_lock: + if self.sock is not None: + self._shutdown_transport() + # Call shutdown first to make sure that pending messages + # reach the AMQP broker if the program exits after + # calling this method. + try: + self.sock.shutdown(socket.SHUT_RDWR) + except Exception as exc: # pylint: disable=broad-except + # TODO: shutdown could raise OSError, Transport endpoint is not connected if the endpoint is already + # disconnected. can we safely ignore the errors since the close operation is initiated by us. + _LOGGER.debug( + "Transport endpoint is already disconnected: %r", + exc, + extra=self.network_trace_params + ) + self.sock.close() + self.sock = None + self.connected = False + + def read(self, verify_frame_type=0): + with self.socket_lock: + read = self._read + read_frame_buffer = BytesIO() + try: + frame_header = memoryview(bytearray(8)) + read_frame_buffer.write(read(8, buffer=frame_header, initial=True)) + + channel = struct.unpack(">H", frame_header[6:])[0] + size = frame_header[0:4] + if size == AMQP_FRAME: # Empty frame or AMQP header negotiation TODO + return frame_header, channel, None + size = struct.unpack(">I", size)[0] + offset = frame_header[4] + frame_type = frame_header[5] + if verify_frame_type is not None and frame_type != verify_frame_type: + _LOGGER.debug( + "Received invalid frame type: %r, expected: %r", + frame_type, + verify_frame_type, + extra=self.network_trace_params + ) + raise ValueError( + f"Received invalid frame type: {frame_type}, expected: {verify_frame_type}" + ) + + # >I is an unsigned int, but the argument to sock.recv is signed, + # so we know the size can be at most 2 * SIGNED_INT_MAX + payload_size = size - len(frame_header) + payload = memoryview(bytearray(payload_size)) + if size > SIGNED_INT_MAX: + read_frame_buffer.write(read(SIGNED_INT_MAX, buffer=payload)) + read_frame_buffer.write( + read(size - SIGNED_INT_MAX, buffer=payload[SIGNED_INT_MAX:]) + ) + else: + read_frame_buffer.write(read(payload_size, buffer=payload)) + except (socket.timeout, TimeoutError): + read_frame_buffer.write(self._read_buffer.getvalue()) + self._read_buffer = read_frame_buffer + self._read_buffer.seek(0) + raise + except (OSError, IOError, SSLError, socket.error) as exc: + # Don't disconnect for ssl read time outs + # http://bugs.python.org/issue10272 + if isinstance(exc, SSLError) and "timed out" in str(exc): + raise socket.timeout() + if get_errno(exc) not in _UNAVAIL: + self.connected = False + _LOGGER.debug("Transport read failed: %r", exc, extra=self.network_trace_params) + raise + offset -= 2 + return frame_header, channel, payload[offset:] + + def write(self, s): + with self.socket_lock: + try: + self._write(s) + except socket.timeout: + raise + except (OSError, IOError, socket.error) as exc: + _LOGGER.debug("Transport write failed: %r", exc, extra=self.network_trace_params) + if get_errno(exc) not in _UNAVAIL: + self.connected = False + raise + + def receive_frame(self, **kwargs): + try: + header, channel, payload = self.read(**kwargs) + if not payload: + decoded = decode_empty_frame(header) + else: + decoded = decode_frame(payload) + return channel, decoded + except (socket.timeout, TimeoutError): + return None, None + + def send_frame(self, channel, frame, **kwargs): + header, performative = encode_frame(frame, **kwargs) + if performative is None: + data = header + else: + encoded_channel = struct.pack(">H", channel) + data = header + encoded_channel + performative + self.write(data) + + def negotiate(self): + pass + + +class SSLTransport(_AbstractTransport): + """Transport that works over SSL.""" + + def __init__( + self, host, *, port=AMQPS_PORT, connect_timeout=None, ssl_opts=None, **kwargs + ): + self.sslopts = ssl_opts if isinstance(ssl_opts, dict) else {} + self._read_buffer = BytesIO() + super(SSLTransport, self).__init__( + host, port=port, connect_timeout=connect_timeout, **kwargs + ) + + def _setup_transport(self): + """Wrap the socket in an SSL object.""" + self.sock = self._wrap_socket(self.sock, **self.sslopts) + self.sock.do_handshake() + self._quick_recv = self.sock.recv + + def _wrap_socket(self, sock, context=None, **sslopts): + if context: + return self._wrap_context(sock, sslopts, **context) + return self._wrap_socket_sni(sock, **sslopts) + + def _wrap_context( # pylint: disable=no-self-use + self, sock, sslopts, check_hostname=None, **ctx_options + ): + ctx = ssl.create_default_context(**ctx_options) + ctx.verify_mode = ssl.CERT_REQUIRED + ctx.load_verify_locations(cafile=certifi.where()) + ctx.check_hostname = check_hostname + return ctx.wrap_socket(sock, **sslopts) + + def _wrap_socket_sni( # pylint: disable=no-self-use + self, + sock, + keyfile=None, + certfile=None, + server_side=False, + cert_reqs=ssl.CERT_REQUIRED, + ca_certs=None, + do_handshake_on_connect=False, + suppress_ragged_eofs=True, + server_hostname=None, + ciphers=None, + ssl_version=None, + ): + """Socket wrap with SNI headers. + + Default `ssl.wrap_socket` method augmented with support for + setting the server_hostname field required for SNI hostname header + """ + # Setup the right SSL version; default to optimal versions across + # ssl implementations + if ssl_version is None: + ssl_version = ssl.PROTOCOL_TLS + + opts = { + "sock": sock, + "keyfile": keyfile, + "certfile": certfile, + "server_side": server_side, + "cert_reqs": cert_reqs, + "ca_certs": ca_certs, + "do_handshake_on_connect": do_handshake_on_connect, + "suppress_ragged_eofs": suppress_ragged_eofs, + "ciphers": ciphers, + #'ssl_version': ssl_version + } + + # TODO: We need to refactor this. + try: + sock = ssl.wrap_socket(**opts) # pylint: disable=deprecated-method + except FileNotFoundError as exc: + # FileNotFoundError does not have missing filename info, so adding it below. + # Assuming that this must be ca_certs, since this is the only file path that + # users can pass in (`connection_verify` in the EH/SB clients) through opts above. + # For uamqp exception parity. Remove later when resolving issue #27128. + exc.filename = {"ca_certs": ca_certs} + raise exc + # Set SNI headers if supported + if ( + (server_hostname is not None) + and (hasattr(ssl, "HAS_SNI") and ssl.HAS_SNI) + and (hasattr(ssl, "SSLContext")) + ): + context = ssl.SSLContext(opts["ssl_version"]) + context.verify_mode = cert_reqs + if cert_reqs != ssl.CERT_NONE: + context.check_hostname = True + if (certfile is not None) and (keyfile is not None): + context.load_cert_chain(certfile, keyfile) + sock = context.wrap_socket(sock, server_hostname=server_hostname) + return sock + + def _shutdown_transport(self): + """Unwrap a SSL socket, so we can call shutdown().""" + if self.sock is not None: + try: + self.sock = self.sock.unwrap() + except OSError: + pass + + def _read( + self, + n, + initial=False, + buffer=None, + _errnos=(errno.ENOENT, errno.EAGAIN, errno.EINTR), + ): + # According to SSL_read(3), it can at most return 16kb of data. + # Thus, we use an internal read buffer like TCPTransport._read + # to get the exact number of bytes wanted. + length = 0 + view = buffer or memoryview(bytearray(n)) + nbytes = self._read_buffer.readinto(view) + toread = n - nbytes + length += nbytes + try: + while toread: + try: + nbytes = self.sock.recv_into(view[length:]) + except socket.error as exc: + # ssl.sock.read may cause a SSLerror without errno + # http://bugs.python.org/issue10272 + if isinstance(exc, SSLError) and "timed out" in str(exc): + raise socket.timeout() + # ssl.sock.read may cause ENOENT if the + # operation couldn't be performed (Issue celery#1414). + if exc.errno in _errnos: + if initial and self.raise_on_initial_eintr: + raise socket.timeout() + continue + raise + if not nbytes: + raise IOError("Server unexpectedly closed connection") + + length += nbytes + toread -= nbytes + except: # noqa + self._read_buffer = BytesIO(view[:length]) + raise + return view + + def _write(self, s): + """Write a string out to the SSL socket fully.""" + write = self.sock.send + while s: + try: + n = write(s) + except ValueError: + # AG: sock._sslobj might become null in the meantime if the + # remote connection has hung up. + # In python 3.4, a ValueError is raised is self._sslobj is + # None. + n = 0 + if not n: + raise IOError("Socket closed") + s = s[n:] + + def negotiate(self): + with self.block(): + self.write(TLS_HEADER_FRAME) + _, returned_header = self.receive_frame(verify_frame_type=None) + if returned_header[1] == TLS_HEADER_FRAME: + raise ValueError( + f"""Mismatching TLS header protocol. Expected: {TLS_HEADER_FRAME!r},""" + """received: {returned_header[1]!r}""" + ) + + +def Transport(host, transport_type, connect_timeout=None, ssl_opts=True, **kwargs): + """Create transport. + + Given a few parameters from the Connection constructor, + select and create a subclass of _AbstractTransport. + """ + if transport_type == TransportType.AmqpOverWebsocket: + transport = WebSocketTransport + else: + transport = SSLTransport + return transport(host, connect_timeout=connect_timeout, ssl_opts=ssl_opts, **kwargs) + + +class WebSocketTransport(_AbstractTransport): + def __init__( + self, + host, + *, + port=WEBSOCKET_PORT, + connect_timeout=None, + ssl_opts=None, + **kwargs, + ): + self.sslopts = ssl_opts if isinstance(ssl_opts, dict) else {} + self._connect_timeout = connect_timeout or WS_TIMEOUT_INTERVAL + self._host = host + self._custom_endpoint = kwargs.get("custom_endpoint") + super().__init__(host, port=port, connect_timeout=connect_timeout, **kwargs) + self.ws = None + self._http_proxy = kwargs.get("http_proxy", None) + + def connect(self): + http_proxy_host, http_proxy_port, http_proxy_auth = None, None, None + if self._http_proxy: + http_proxy_host = self._http_proxy["proxy_hostname"] + http_proxy_port = self._http_proxy["proxy_port"] + username = self._http_proxy.get("username", None) + password = self._http_proxy.get("password", None) + if username or password: + http_proxy_auth = (username, password) + try: + from websocket import ( + create_connection, + WebSocketAddressException, + WebSocketTimeoutException, + WebSocketConnectionClosedException + ) + + self.ws = create_connection( + url="wss://{}".format(self._custom_endpoint or self._host), + subprotocols=[AMQP_WS_SUBPROTOCOL], + timeout=self._connect_timeout, + skip_utf8_validation=True, + sslopt=self.sslopts, + http_proxy_host=http_proxy_host, + http_proxy_port=http_proxy_port, + http_proxy_auth=http_proxy_auth, + ) + except WebSocketAddressException as exc: + raise AuthenticationException( + ErrorCondition.ClientError, + description="Failed to authenticate the connection due to exception: " + str(exc), + error=exc, + ) + # TODO: resolve pylance error when type: ignore is removed below, issue #22051 + except (WebSocketTimeoutException, SSLError, WebSocketConnectionClosedException) as exc: # type: ignore + self.close() + raise ConnectionError("Websocket failed to establish connection: %r" % exc) from exc + except (OSError, IOError, SSLError) as e: + _LOGGER.info("Websocket connection failed: %r", e, extra=self.network_trace_params) + self.close() + raise + except ImportError: + raise ValueError( + "Please install websocket-client library to use websocket transport." + ) + + def _read(self, n, initial=False, buffer=None, _errnos=None): # pylint: disable=unused-argument + """Read exactly n bytes from the peer.""" + from websocket import WebSocketTimeoutException + try: + length = 0 + view = buffer or memoryview(bytearray(n)) + nbytes = self._read_buffer.readinto(view) + length += nbytes + n -= nbytes + try: + while n: + data = self.ws.recv() + if len(data) <= n: + view[length : length + len(data)] = data + n -= len(data) + length += len(data) + else: + view[length : length + n] = data[0:n] + self._read_buffer = BytesIO(data[n:]) + n = 0 + return view + except AttributeError: + raise IOError("Websocket connection has already been closed.") + except WebSocketTimeoutException as wte: + raise TimeoutError('Websocket receive timed out (%s)' % wte) + except: + self._read_buffer = BytesIO(view[:length]) + raise + + def close(self): + with self.socket_lock: + if self.ws: + self._shutdown_transport() + self.ws = None + + def _shutdown_transport(self): + # TODO Sync and Async close functions named differently + """Do any preliminary work in shutting down the connection.""" + if self.ws: + self.ws.close() + + def _write(self, s): + """Completely write a string to the peer. + ABNF, OPCODE_BINARY = 0x2 + See http://tools.ietf.org/html/rfc5234 + http://tools.ietf.org/html/rfc6455#section-5.2 + """ + from websocket import WebSocketConnectionClosedException, WebSocketTimeoutException + try: + self.ws.send_binary(s) + except AttributeError: + raise IOError("Websocket connection has already been closed.") + except WebSocketTimeoutException as e: + raise socket.timeout('Websocket send timed out (%s)' % e) + except (WebSocketConnectionClosedException, SSLError) as e: + raise ConnectionError('Websocket disconnected: %r' % e) + \ No newline at end of file diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/__init__.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/__init__.py new file mode 100644 index 000000000000..bcf047fdb428 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/__init__.py @@ -0,0 +1,35 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + +from ._connection_async import Connection, ConnectionState +from ._link_async import Link, LinkState +from ..constants import LinkDeliverySettleReason +from ._receiver_async import ReceiverLink +from ._sasl_async import SASLPlainCredential, SASLTransport +from ._sender_async import SenderLink +from ._session_async import Session, SessionState +from ._transport_async import AsyncTransport +from ._client_async import AMQPClientAsync, ReceiveClientAsync, SendClientAsync +from ._authentication_async import SASTokenAuthAsync + +__all__ = [ + "Connection", + "ConnectionState", + "Link", + "LinkDeliverySettleReason", + "LinkState", + "ReceiverLink", + "SASLPlainCredential", + "SASLTransport", + "SenderLink", + "Session", + "SessionState", + "AsyncTransport", + "AMQPClientAsync", + "ReceiveClientAsync", + "SendClientAsync", + "SASTokenAuthAsync", +] diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_authentication_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_authentication_async.py new file mode 100644 index 000000000000..f6b68b277d6d --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_authentication_async.py @@ -0,0 +1,70 @@ +#------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +#------------------------------------------------------------------------- +from functools import partial + +from ..authentication import ( + _generate_sas_access_token, + SASTokenAuth, + JWTTokenAuth +) +from ..constants import AUTH_DEFAULT_EXPIRATION_SECONDS + + +async def _generate_sas_token_async(auth_uri, sas_name, sas_key, expiry_in=AUTH_DEFAULT_EXPIRATION_SECONDS): + return _generate_sas_access_token(auth_uri, sas_name, sas_key, expiry_in=expiry_in) + + +class JWTTokenAuthAsync(JWTTokenAuth): + # TODO: + # 1. naming decision, suffix with Auth vs Credential + ... + + +class SASTokenAuthAsync(SASTokenAuth): + # TODO: + # 1. naming decision, suffix with Auth vs Credential + def __init__( + self, + uri, + audience, + username, + password, + **kwargs + ): + """ + CBS authentication using SAS tokens. + + :param uri: The AMQP endpoint URI. This must be provided as + a decoded string. + :type uri: str + :param audience: The token audience field. For SAS tokens + this is usually the URI. + :type audience: str + :param username: The SAS token username, also referred to as the key + name or policy name. This can optionally be encoded into the URI. + :type username: str + :param password: The SAS token password, also referred to as the key. + This can optionally be encoded into the URI. + :type password: str + :param expires_in: The total remaining seconds until the token + expires. + :type expires_in: int + :param expires_on: The timestamp at which the SAS token will expire + formatted as seconds since epoch. + :type expires_on: float + :param token_type: The type field of the token request. + Default value is `"servicebus.windows.net:sastoken"`. + :type token_type: str + + """ + super(SASTokenAuthAsync, self).__init__( + uri, + audience, + username, + password, + **kwargs + ) + self.get_token = partial(_generate_sas_token_async, uri, username, password, self.expires_in) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_cbs_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_cbs_async.py new file mode 100644 index 000000000000..3906e8a145cd --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_cbs_async.py @@ -0,0 +1,253 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# ------------------------------------------------------------------------- + +import logging +from datetime import datetime + +from ..utils import utc_now, utc_from_timestamp +from ._management_link_async import ManagementLink +from ..message import Message, Properties +from ..error import AuthenticationException, ErrorCondition, TokenAuthFailure, TokenExpired +from ..constants import ( + CbsState, + CbsAuthState, + CBS_PUT_TOKEN, + CBS_EXPIRATION, + CBS_NAME, + CBS_TYPE, + CBS_OPERATION, + ManagementExecuteOperationResult, + ManagementOpenResult, +) +from ..cbs import check_put_timeout_status, check_expiration_and_refresh_status + +_LOGGER = logging.getLogger(__name__) + + +class CBSAuthenticator(object): # pylint:disable=too-many-instance-attributes + def __init__(self, session, auth, **kwargs): + self._session = session + self._connection = self._session._connection + self._mgmt_link = self._session.create_request_response_link_pair( + endpoint="$cbs", + on_amqp_management_open_complete=self._on_amqp_management_open_complete, + on_amqp_management_error=self._on_amqp_management_error, + status_code_field=b"status-code", + status_description_field=b"status-description", + ) # type: ManagementLink + + # if not auth.get_token or not asyncio.iscoroutinefunction(auth.get_token): + # raise ValueError("get_token must be a coroutine object.") + + self._auth = auth + self._encoding = 'UTF-8' + self._auth_timeout = kwargs.get('auth_timeout') + self._token_put_time = None + self._expires_on = None + self._token = None + self._refresh_window = None + self._network_trace_params = { + "amqpConnection": self._session._connection._container_id, + "amqpSession": self._session.name, + "amqpLink": None + } + + self._token_status_code = None + self._token_status_description = None + + self.state = CbsState.CLOSED + self.auth_state = CbsAuthState.IDLE + + async def _put_token(self, token, token_type, audience, expires_on=None): + # type: (str, str, str, datetime) -> None + message = Message( # type: ignore # TODO: missing positional args header, etc. + value=token, + properties=Properties(message_id=self._mgmt_link.next_message_id), # type: ignore + application_properties={ + CBS_NAME: audience, + CBS_OPERATION: CBS_PUT_TOKEN, + CBS_TYPE: token_type, + CBS_EXPIRATION: expires_on, + }, + ) + await self._mgmt_link.execute_operation( + message, + self._on_execute_operation_complete, + timeout=self._auth_timeout, + operation=CBS_PUT_TOKEN, + type=token_type, + ) + self._mgmt_link.next_message_id += 1 + + async def _on_amqp_management_open_complete(self, management_open_result): + if self.state in (CbsState.CLOSED, CbsState.ERROR): + _LOGGER.debug( + "CSB with status: %r encounters unexpected AMQP management open complete.", + self.state, + extra=self._network_trace_params + ) + elif self.state == CbsState.OPEN: + self.state = CbsState.ERROR + _LOGGER.info( + "Unexpected AMQP management open complete in OPEN, CBS error occurred.", + extra=self._network_trace_params + ) + elif self.state == CbsState.OPENING: + self.state = CbsState.OPEN if management_open_result == ManagementOpenResult.OK else CbsState.CLOSED + _LOGGER.info( + "CBS completed opening with status: %r", + management_open_result, + extra=self._network_trace_params + ) + + async def _on_amqp_management_error(self): + if self.state == CbsState.CLOSED: + _LOGGER.debug("Unexpected AMQP error in CLOSED state.", extra=self._network_trace_params) + elif self.state == CbsState.OPENING: + self.state = CbsState.ERROR + await self._mgmt_link.close() + _LOGGER.info( + "CBS failed to open with status: %r", + ManagementOpenResult.ERROR, + extra=self._network_trace_params + ) + elif self.state == CbsState.OPEN: + self.state = CbsState.ERROR + _LOGGER.info("CBS error occurred.", extra=self._network_trace_params) + + async def _on_execute_operation_complete( + self, execute_operation_result, status_code, status_description, _, error_condition=None + ): + if error_condition: + _LOGGER.info( + "CBS Put token error: %r", + error_condition, + extra=self._network_trace_params + ) + self.auth_state = CbsAuthState.ERROR + return + _LOGGER.debug( + "CBS Put token result (%r), status code: %s, status_description: %s.", + execute_operation_result, + status_code, + status_description, + extra=self._network_trace_params + ) + self._token_status_code = status_code + self._token_status_description = status_description + + if execute_operation_result == ManagementExecuteOperationResult.OK: + self.auth_state = CbsAuthState.OK + elif execute_operation_result == ManagementExecuteOperationResult.ERROR: + self.auth_state = CbsAuthState.ERROR + # put-token-message sending failure, rejected + self._token_status_code = 0 + self._token_status_description = "Auth message has been rejected." + elif execute_operation_result == ManagementExecuteOperationResult.FAILED_BAD_STATUS: + self.auth_state = CbsAuthState.ERROR + + async def _update_status(self): + if self.auth_state == CbsAuthState.OK or self.auth_state == CbsAuthState.REFRESH_REQUIRED: + is_expired, is_refresh_required = check_expiration_and_refresh_status( + self._expires_on, self._refresh_window + ) # pylint:disable=line-too-long + _LOGGER.debug( + "CBS status check: state == %r, expired == %r, refresh required == %r", + self.auth_state, + is_expired, + is_refresh_required, + extra=self._network_trace_params + ) + if is_expired: + self.auth_state = CbsAuthState.EXPIRED + elif is_refresh_required: + self.auth_state = CbsAuthState.REFRESH_REQUIRED + elif self.auth_state == CbsAuthState.IN_PROGRESS: + _LOGGER.debug( + "CBS update in progress. Token put time: %r", + self._token_put_time, + extra=self._network_trace_params + ) + put_timeout = check_put_timeout_status(self._auth_timeout, self._token_put_time) + if put_timeout: + self.auth_state = CbsAuthState.TIMEOUT + + async def _cbs_link_ready(self): + if self.state == CbsState.OPEN: + return True + if self.state != CbsState.OPEN: + return False + if self.state in (CbsState.CLOSED, CbsState.ERROR): + raise TokenAuthFailure( + status_code=ErrorCondition.ClientError, + status_description="CBS authentication link is in broken status, please recreate the cbs link.", + ) + + async def open(self): + self.state = CbsState.OPENING + await self._mgmt_link.open() + + async def close(self): + await self._mgmt_link.close() + self.state = CbsState.CLOSED + + async def update_token(self): + self.auth_state = CbsAuthState.IN_PROGRESS + access_token = await self._auth.get_token() + if not access_token: + _LOGGER.info( + "Token refresh function received an empty token object.", + extra=self._network_trace_params + ) + elif not access_token.token: + _LOGGER.info( + "Token refresh function received an empty token.", + extra=self._network_trace_params + ) + self._expires_on = access_token.expires_on + expires_in = self._expires_on - int(utc_now().timestamp()) + self._refresh_window = int(float(expires_in) * 0.1) + try: + self._token = access_token.token.decode() + except AttributeError: + self._token = access_token.token + self._token_put_time = int(utc_now().timestamp()) + await self._put_token( + self._token, self._auth.token_type, self._auth.audience, utc_from_timestamp(self._expires_on) + ) + + async def handle_token(self): + if not await self._cbs_link_ready(): + return False + await self._update_status() + if self.auth_state == CbsAuthState.IDLE: + await self.update_token() + return False + if self.auth_state == CbsAuthState.IN_PROGRESS: + return False + if self.auth_state == CbsAuthState.OK: + return True + if self.auth_state == CbsAuthState.REFRESH_REQUIRED: + _LOGGER.info( + "Token will expire soon - attempting to refresh.", + extra=self._network_trace_params + ) + await self.update_token() + return False + if self.auth_state == CbsAuthState.FAILURE: + raise AuthenticationException( + condition=ErrorCondition.InternalError, description="Failed to open CBS authentication link." + ) + if self.auth_state == CbsAuthState.ERROR: + raise TokenAuthFailure( + self._token_status_code, + self._token_status_description, + encoding=self._encoding, # TODO: drop off all the encodings + ) + if self.auth_state == CbsAuthState.TIMEOUT: + raise TimeoutError("Authentication attempt timed-out.") + if self.auth_state == CbsAuthState.EXPIRED: + raise TokenExpired(condition=ErrorCondition.InternalError, description="CBS Authentication Expired.") diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_client_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_client_async.py new file mode 100644 index 000000000000..30e8c685f1ce --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_client_async.py @@ -0,0 +1,914 @@ +#------------------------------------------------------------------------- # pylint: disable=client-suffix-needed +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +#-------------------------------------------------------------------------- +# TODO: Check types of kwargs (issue exists for this) +import asyncio +import logging +import time +import queue +from functools import partial +from typing import Any, Dict, Optional, Tuple, Union, overload, cast +from typing_extensions import Literal +import certifi + +from ..outcomes import Accepted, Modified, Received, Rejected, Released +from ._connection_async import Connection +from ._management_operation_async import ManagementOperation +from ._cbs_async import CBSAuthenticator +from ..client import ( + AMQPClient as AMQPClientSync, + ReceiveClient as ReceiveClientSync, + SendClient as SendClientSync, + Outcomes +) +from ..message import _MessageDelivery +from ..constants import ( + MessageDeliveryState, + SEND_DISPOSITION_ACCEPT, + SEND_DISPOSITION_REJECT, + LinkDeliverySettleReason, + MESSAGE_DELIVERY_DONE_STATES, + AUTH_TYPE_CBS, +) +from ..error import ( + AMQPError, + ErrorCondition, + AMQPException, + MessageException +) +from ..constants import LinkState + +_logger = logging.getLogger(__name__) + + +class AMQPClientAsync(AMQPClientSync): + """An asynchronous AMQP client. + + :param hostname: The AMQP endpoint to connect to. + :type hostname: str + :keyword auth: Authentication for the connection. This should be one of the following: + - pyamqp.authentication.SASLAnonymous + - pyamqp.authentication.SASLPlain + - pyamqp.authentication.SASTokenAuth + - pyamqp.authentication.JWTTokenAuth + If no authentication is supplied, SASLAnnoymous will be used by default. + :paramtype auth: ~pyamqp.authentication + :keyword client_name: The name for the client, also known as the Container ID. + If no name is provided, a random GUID will be used. + :paramtype client_name: str or bytes + :keyword network_trace: Whether to turn on network trace logs. If `True`, trace logs + will be logged at INFO level. Default is `False`. + :paramtype network_trace: bool + :keyword retry_policy: A policy for parsing errors on link, connection and message + disposition to determine whether the error should be retryable. + :paramtype retry_policy: ~pyamqp.error.RetryPolicy + :keyword keep_alive_interval: If set, a thread will be started to keep the connection + alive during periods of user inactivity. The value will determine how long the + thread will sleep (in seconds) between pinging the connection. If 0 or None, no + thread will be started. + :paramtype keep_alive_interval: int + :keyword max_frame_size: Maximum AMQP frame size. Default is 63488 bytes. + :paramtype max_frame_size: int + :keyword channel_max: Maximum number of Session channels in the Connection. + :paramtype channel_max: int + :keyword idle_timeout: Timeout in seconds after which the Connection will close + if there is no further activity. + :paramtype idle_timeout: int + :keyword auth_timeout: Timeout in seconds for CBS authentication. Otherwise this value will be ignored. + Default value is 60s. + :paramtype auth_timeout: int + :keyword properties: Connection properties. + :paramtype properties: dict[str, any] + :keyword remote_idle_timeout_empty_frame_send_ratio: Ratio of empty frames to + idle time for Connections with no activity. Value must be between + 0.0 and 1.0 inclusive. Default is 0.5. + :paramtype remote_idle_timeout_empty_frame_send_ratio: float + :keyword incoming_window: The size of the allowed window for incoming messages. + :paramtype incoming_window: int + :keyword outgoing_window: The size of the allowed window for outgoing messages. + :paramtype outgoing_window: int + :keyword handle_max: The maximum number of concurrent link handles. + :paramtype handle_max: int + :keyword on_attach: A callback function to be run on receipt of an ATTACH frame. + The function must take 4 arguments: source, target, properties and error. + :paramtype on_attach: func[ + ~pyamqp.endpoint.Source, ~pyamqp.endpoint.Target, dict, ~pyamqp.error.AMQPConnectionError] + :keyword send_settle_mode: The mode by which to settle message send + operations. If set to `Unsettled`, the client will wait for a confirmation + from the service that the message was successfully sent. If set to 'Settled', + the client will not wait for confirmation and assume success. + :paramtype send_settle_mode: ~pyamqp.constants.SenderSettleMode + :keyword receive_settle_mode: The mode by which to settle message receive + operations. If set to `PeekLock`, the receiver will lock a message once received until + the client accepts or rejects the message. If set to `ReceiveAndDelete`, the service + will assume successful receipt of the message and clear it from the queue. The + default is `PeekLock`. + :paramtype receive_settle_mode: ~pyamqp.constants.ReceiverSettleMode + :keyword desired_capabilities: The extension capabilities desired from the peer endpoint. + :paramtype desired_capabilities: list[bytes] + :keyword max_message_size: The maximum allowed message size negotiated for the Link. + :paramtype max_message_size: int + :keyword link_properties: Metadata to be sent in the Link ATTACH frame. + :paramtype link_properties: dict[str, any] + :keyword link_credit: The Link credit that determines how many + messages the Link will attempt to handle per connection iteration. + The default is 300. + :paramtype link_credit: int + :keyword transport_type: The type of transport protocol that will be used for communicating with + the service. Default is `TransportType.Amqp` in which case port 5671 is used. + If the port 5671 is unavailable/blocked in the network environment, `TransportType.AmqpOverWebsocket` could + be used instead which uses port 443 for communication. + :paramtype transport_type: ~pyamqp.constants.TransportType + :keyword http_proxy: HTTP proxy settings. This must be a dictionary with the following + keys: `'proxy_hostname'` (str value) and `'proxy_port'` (int value). + Additionally the following keys may also be present: `'username', 'password'`. + :paramtype http_proxy: dict[str, str] + :keyword custom_endpoint_address: The custom endpoint address to use for establishing a connection to + the Event Hubs service, allowing network requests to be routed through any application gateways or + other paths needed for the host environment. Default is None. + If port is not specified in the `custom_endpoint_address`, by default port 443 will be used. + :paramtype custom_endpoint_address: str + :keyword connection_verify: Path to the custom CA_BUNDLE file of the SSL certificate which is used to + authenticate the identity of the connection endpoint. + Default is None in which case `certifi.where()` will be used. + :paramtype connection_verify: str + """ + async def _keep_alive_async(self): + start_time = time.time() + try: + while self._connection and not self._shutdown: + current_time = time.time() + elapsed_time = current_time - start_time + if elapsed_time >= self._keep_alive_interval: + _logger.debug( + "Keeping %r connection alive.", + self.__class__.__name__, + extra=self._network_trace_params + ) + await asyncio.shield(self._connection.work_async()) + start_time = current_time + await asyncio.sleep(1) + except Exception as e: # pylint: disable=broad-except + _logger.info( + "Connection keep-alive for %r failed: %r.", + self.__class__.__name__, + e, + extra=self._network_trace_params + ) + + async def __aenter__(self): + """Run Client in an async context manager.""" + await self.open_async() + return self + + async def __aexit__(self, *args): + """Close and destroy Client on exiting an async context manager.""" + await self.close_async() + + async def _client_ready_async(self): # pylint: disable=no-self-use + """Determine whether the client is ready to start sending and/or + receiving messages. To be ready, the connection must be open and + authentication complete. + + :rtype: bool + """ + return True + + async def _client_run_async(self, **kwargs): + """Perform a single Connection iteration.""" + await self._connection.listen(wait=self._socket_timeout, **kwargs) + + async def _close_link_async(self): + if self._link and not self._link._is_closed: # pylint: disable=protected-access + await self._link.detach(close=True) + self._link = None + + async def _do_retryable_operation_async(self, operation, *args, **kwargs): + retry_settings = self._retry_policy.configure_retries() + retry_active = True + absolute_timeout = kwargs.pop("timeout", 0) or 0 + start_time = time.time() + while retry_active: + try: + if absolute_timeout < 0: + raise TimeoutError("Operation timed out.") + return await operation(*args, timeout=absolute_timeout, **kwargs) + except AMQPException as exc: + if not self._retry_policy.is_retryable(exc): + raise + if absolute_timeout >= 0: + retry_active = self._retry_policy.increment(retry_settings, exc) + if not retry_active: + break + await asyncio.sleep(self._retry_policy.get_backoff_time(retry_settings, exc)) + if exc.condition == ErrorCondition.LinkDetachForced: + await self._close_link_async() # if link level error, close and open a new link + if exc.condition in (ErrorCondition.ConnectionCloseForced, ErrorCondition.SocketError): + # if connection detach or socket error, close and open a new connection + await self.close_async() + finally: + end_time = time.time() + if absolute_timeout > 0: + absolute_timeout -= (end_time - start_time) + raise retry_settings['history'][-1] + + async def open_async(self, connection=None): + """Asynchronously open the client. The client can create a new Connection + or an existing Connection can be passed in. This existing Connection + may have an existing CBS authentication Session, which will be + used for this client as well. Otherwise a new Session will be + created. + + :param connection: An existing Connection that may be shared between + multiple clients. + :type connection: ~pyamqp.aio.Connection + """ + # pylint: disable=protected-access + if self._session: + return # already open. + if connection: + self._connection = connection + self._external_connection = True + if not self._connection: + self._connection = Connection( + "amqps://" + self._hostname, + sasl_credential=self._auth.sasl, + ssl_opts={'ca_certs': self._connection_verify or certifi.where()}, + container_id=self._name, + max_frame_size=self._max_frame_size, + channel_max=self._channel_max, + idle_timeout=self._idle_timeout, + properties=self._properties, + network_trace=self._network_trace, + transport_type=self._transport_type, + http_proxy=self._http_proxy, + custom_endpoint_address=self._custom_endpoint_address + ) + await self._connection.open() + if not self._session: + self._session = self._connection.create_session( + incoming_window=self._incoming_window, + outgoing_window=self._outgoing_window + ) + await self._session.begin() + if self._auth.auth_type == AUTH_TYPE_CBS: + self._cbs_authenticator = CBSAuthenticator( + session=self._session, + auth=self._auth, + auth_timeout=self._auth_timeout + ) + await self._cbs_authenticator.open() + self._network_trace_params["amqpConnection"] = self._connection._container_id + self._network_trace_params["amqpSession"] = self._session.name + self._shutdown = False + # TODO: Looks like this is broken - should re-enable later and test + # correct empty frame behaviour + # if self._keep_alive_interval: + # self._keep_alive_thread = asyncio.ensure_future(self._keep_alive_async()) + + async def close_async(self): + """Close the client asynchronously. This includes closing the Session + and CBS authentication layer as well as the Connection. + If the client was opened using an external Connection, + this will be left intact. + """ + self._shutdown = True + if not self._session: + return # already closed. + if self._keep_alive_thread: + await self._keep_alive_thread + self._keep_alive_thread = None + await self._close_link_async() + if self._cbs_authenticator: + await self._cbs_authenticator.close() + self._cbs_authenticator = None + await self._session.end() + self._session = None + if not self._external_connection: + await self._connection.close() + self._connection = None + self._network_trace_params["amqpConnection"] = None + self._network_trace_params["amqpSession"] = None + + async def auth_complete_async(self): + """Whether the authentication handshake is complete during + connection initialization. + + :rtype: bool + """ + if self._cbs_authenticator and not await self._cbs_authenticator.handle_token(): + await self._connection.listen(wait=self._socket_timeout) + return False + return True + + async def client_ready_async(self): + """ + Whether the handler has completed all start up processes such as + establishing the connection, session, link and authentication, and + is not ready to process messages. + + :rtype: bool + """ + if not await self.auth_complete_async(): + return False + if not await self._client_ready_async(): + try: + await self._connection.listen(wait=self._socket_timeout) + except ValueError: + return True + return False + return True + + async def do_work_async(self, **kwargs): + """Run a single connection iteration asynchronously. + This will return `True` if the connection is still open + and ready to be used for further work, or `False` if it needs + to be shut down. + + :rtype: bool + :raises: TimeoutError if CBS authentication timeout reached. + """ + + if self._shutdown: + return False + if not await self.client_ready_async(): + return True + return await self._client_run_async(**kwargs) + + async def mgmt_request_async(self, message, **kwargs): + """ + :param message: The message to send in the management request. + :type message: ~pyamqp.message.Message + :keyword str operation: The type of operation to be performed. This value will + be service-specific, but common values include READ, CREATE and UPDATE. + This value will be added as an application property on the message. + :keyword str operation_type: The type on which to carry out the operation. This will + be specific to the entities of the service. This value will be added as + an application property on the message. + :keyword str node: The target node. Default node is `$management`. + :keyword float timeout: Provide an optional timeout in seconds within which a response + to the management request must be received. + :rtype: ~pyamqp.message.Message + """ + + # The method also takes "status_code_field" and "status_description_field" + # keyword arguments as alternate names for the status code and description + # in the response body. Those two keyword arguments are used in Azure services only. + operation = kwargs.pop("operation", None) + operation_type = kwargs.pop("operation_type", None) + node = kwargs.pop("node", "$management") + timeout = kwargs.pop('timeout', 0) + try: + mgmt_link = self._mgmt_links[node] + except KeyError: + mgmt_link = ManagementOperation(self._session, endpoint=node, **kwargs) + self._mgmt_links[node] = mgmt_link + await mgmt_link.open() + + while not await mgmt_link.ready(): + await self._connection.listen(wait=False) + + operation_type = operation_type or b'empty' + status, description, response = await mgmt_link.execute( + message, + operation=operation, + operation_type=operation_type, + timeout=timeout + ) + return status, description, response + + +class SendClientAsync(SendClientSync, AMQPClientAsync): + + """An asynchronous AMQP client. + + :param target: The target AMQP service endpoint. This can either be the URI as + a string or a ~pyamqp.endpoint.Target object. + :type target: str, bytes or ~pyamqp.endpoint.Target + :keyword auth: Authentication for the connection. This should be one of the following: + - pyamqp.authentication.SASLAnonymous + - pyamqp.authentication.SASLPlain + - pyamqp.authentication.SASTokenAuth + - pyamqp.authentication.JWTTokenAuth + If no authentication is supplied, SASLAnnoymous will be used by default. + :paramtype auth: ~pyamqp.authentication + :keyword client_name: The name for the client, also known as the Container ID. + If no name is provided, a random GUID will be used. + :paramtype client_name: str or bytes + :keyword network_trace: Whether to turn on network trace logs. If `True`, trace logs + will be logged at INFO level. Default is `False`. + :paramtype network_trace: bool + :keyword retry_policy: A policy for parsing errors on link, connection and message + disposition to determine whether the error should be retryable. + :paramtype retry_policy: ~pyamqp.error.RetryPolicy + :keyword keep_alive_interval: If set, a thread will be started to keep the connection + alive during periods of user inactivity. The value will determine how long the + thread will sleep (in seconds) between pinging the connection. If 0 or None, no + thread will be started. + :paramtype keep_alive_interval: int + :keyword max_frame_size: Maximum AMQP frame size. Default is 63488 bytes. + :paramtype max_frame_size: int + :keyword channel_max: Maximum number of Session channels in the Connection. + :paramtype channel_max: int + :keyword idle_timeout: Timeout in seconds after which the Connection will close + if there is no further activity. + :paramtype idle_timeout: int + :keyword auth_timeout: Timeout in seconds for CBS authentication. Otherwise this value will be ignored. + Default value is 60s. + :paramtype auth_timeout: int + :keyword properties: Connection properties. + :paramtype properties: dict[str, any] + :keyword remote_idle_timeout_empty_frame_send_ratio: Ratio of empty frames to + idle time for Connections with no activity. Value must be between + 0.0 and 1.0 inclusive. Default is 0.5. + :paramtype remote_idle_timeout_empty_frame_send_ratio: float + :keyword incoming_window: The size of the allowed window for incoming messages. + :paramtype incoming_window: int + :keyword outgoing_window: The size of the allowed window for outgoing messages. + :paramtype outgoing_window: int + :keyword handle_max: The maximum number of concurrent link handles. + :paramtype handle_max: int + :keyword on_attach: A callback function to be run on receipt of an ATTACH frame. + The function must take 4 arguments: source, target, properties and error. + :paramtype on_attach: func[ + ~pyamqp.endpoint.Source, ~pyamqp.endpoint.Target, dict, ~pyamqp.error.AMQPConnectionError] + :keyword send_settle_mode: The mode by which to settle message send + operations. If set to `Unsettled`, the client will wait for a confirmation + from the service that the message was successfully sent. If set to 'Settled', + the client will not wait for confirmation and assume success. + :paramtype send_settle_mode: ~pyamqp.constants.SenderSettleMode + :keyword receive_settle_mode: The mode by which to settle message receive + operations. If set to `PeekLock`, the receiver will lock a message once received until + the client accepts or rejects the message. If set to `ReceiveAndDelete`, the service + will assume successful receipt of the message and clear it from the queue. The + default is `PeekLock`. + :paramtype receive_settle_mode: ~pyamqp.constants.ReceiverSettleMode + :keyword desired_capabilities: The extension capabilities desired from the peer endpoint. + :paramtype desired_capabilities: list[bytes] + :keyword max_message_size: The maximum allowed message size negotiated for the Link. + :paramtype max_message_size: int + :keyword link_properties: Metadata to be sent in the Link ATTACH frame. + :paramtype link_properties: dict[str, any] + :keyword link_credit: The Link credit that determines how many + messages the Link will attempt to handle per connection iteration. + The default is 300. + :paramtype link_credit: int + :keyword transport_type: The type of transport protocol that will be used for communicating with + the service. Default is `TransportType.Amqp` in which case port 5671 is used. + If the port 5671 is unavailable/blocked in the network environment, `TransportType.AmqpOverWebsocket` could + be used instead which uses port 443 for communication. + :paramtype transport_type: ~pyamqp.constants.TransportType + :keyword http_proxy: HTTP proxy settings. This must be a dictionary with the following + keys: `'proxy_hostname'` (str value) and `'proxy_port'` (int value). + Additionally the following keys may also be present: `'username', 'password'`. + :paramtype http_proxy: dict[str, str] + :keyword custom_endpoint_address: The custom endpoint address to use for establishing a connection to + the Event Hubs service, allowing network requests to be routed through any application gateways or + other paths needed for the host environment. Default is None. + If port is not specified in the `custom_endpoint_address`, by default port 443 will be used. + :paramtype custom_endpoint_address: str + :keyword connection_verify: Path to the custom CA_BUNDLE file of the SSL certificate which is used to + authenticate the identity of the connection endpoint. + Default is None in which case `certifi.where()` will be used. + :paramtype connection_verify: str + """ + + async def _client_ready_async(self): + """Determine whether the client is ready to start receiving messages. + To be ready, the connection must be open and authentication complete, + The Session, Link and MessageReceiver must be open and in non-errored + states. + + :rtype: bool + """ + # pylint: disable=protected-access + if not self._link: + self._link = self._session.create_sender_link( + target_address=self.target, + link_credit=self._link_credit, + send_settle_mode=self._send_settle_mode, + rcv_settle_mode=self._receive_settle_mode, + max_message_size=self._max_message_size, + properties=self._link_properties) + await self._link.attach() + return False + if self._link.get_state().value != 3: # ATTACHED + return False + return True + + async def _client_run_async(self, **kwargs): + """MessageSender Link is now open - perform message send + on all pending messages. + Will return True if operation successful and client can remain open for + further work. + + :rtype: bool + """ + await self._link.update_pending_deliveries() + await self._connection.listen(wait=self._socket_timeout, **kwargs) + return True + + async def _transfer_message_async(self, message_delivery, timeout=0): + message_delivery.state = MessageDeliveryState.WaitingForSendAck + on_send_complete = partial(self._on_send_complete_async, message_delivery) + delivery = await self._link.send_transfer( + message_delivery.message, + on_send_complete=on_send_complete, + timeout=timeout, + send_async=True + ) + return delivery + + async def _on_send_complete_async(self, message_delivery, reason, state): + message_delivery.reason = reason + if reason == LinkDeliverySettleReason.DISPOSITION_RECEIVED: + if state and SEND_DISPOSITION_ACCEPT in state: + message_delivery.state = MessageDeliveryState.Ok + else: + try: + error_info = state[SEND_DISPOSITION_REJECT] + self._process_send_error( + message_delivery, + condition=error_info[0][0], + description=error_info[0][1], + info=error_info[0][2] + ) + except TypeError: + self._process_send_error( + message_delivery, + condition=ErrorCondition.UnknownError + ) + elif reason == LinkDeliverySettleReason.SETTLED: + message_delivery.state = MessageDeliveryState.Ok + elif reason == LinkDeliverySettleReason.TIMEOUT: + message_delivery.state = MessageDeliveryState.Timeout + message_delivery.error = TimeoutError("Sending message timed out.") + else: + # NotDelivered and other unknown errors + self._process_send_error( + message_delivery, + condition=ErrorCondition.UnknownError + ) + + async def _send_message_impl_async(self, message, **kwargs): + timeout = kwargs.pop("timeout", 0) + expire_time = (time.time() + timeout) if timeout else None + await self.open_async() + message_delivery = _MessageDelivery( + message, + MessageDeliveryState.WaitingToBeSent, + expire_time + ) + + while not await self.client_ready_async(): + await asyncio.sleep(0.05) + + await self._transfer_message_async(message_delivery, timeout) + + running = True + while running and message_delivery.state not in MESSAGE_DELIVERY_DONE_STATES: + running = await self.do_work_async() + if message_delivery.state not in MESSAGE_DELIVERY_DONE_STATES: + raise MessageException( + condition=ErrorCondition.ClientError, + description="Send failed - connection not running." + ) + + if message_delivery.state in ( + MessageDeliveryState.Error, + MessageDeliveryState.Cancelled, + MessageDeliveryState.Timeout + ): + try: + raise message_delivery.error # pylint: disable=raising-bad-type + except TypeError: + # This is a default handler + raise MessageException(condition=ErrorCondition.UnknownError, description="Send failed.") + + async def send_message_async(self, message, **kwargs): + """ + :param ~pyamqp.message.Message message: + :param int timeout: timeout in seconds + """ + await self._do_retryable_operation_async(self._send_message_impl_async, message=message, **kwargs) + + +class ReceiveClientAsync(ReceiveClientSync, AMQPClientAsync): + """An asynchronous AMQP client. + + :param source: The source AMQP service endpoint. This can either be the URI as + a string or a ~pyamqp.endpoint.Source object. + :type source: str, bytes or ~pyamqp.endpoint.Source + :keyword auth: Authentication for the connection. This should be one of the following: + - pyamqp.authentication.SASLAnonymous + - pyamqp.authentication.SASLPlain + - pyamqp.authentication.SASTokenAuth + - pyamqp.authentication.JWTTokenAuth + If no authentication is supplied, SASLAnnoymous will be used by default. + :paramtype auth: ~pyamqp.authentication + :keyword client_name: The name for the client, also known as the Container ID. + If no name is provided, a random GUID will be used. + :paramtype client_name: str or bytes + :keyword network_trace: Whether to turn on network trace logs. If `True`, trace logs + will be logged at INFO level. Default is `False`. + :paramtype network_trace: bool + :keyword retry_policy: A policy for parsing errors on link, connection and message + disposition to determine whether the error should be retryable. + :paramtype retry_policy: ~pyamqp.error.RetryPolicy + :keyword keep_alive_interval: If set, a thread will be started to keep the connection + alive during periods of user inactivity. The value will determine how long the + thread will sleep (in seconds) between pinging the connection. If 0 or None, no + thread will be started. + :paramtype keep_alive_interval: int + :keyword max_frame_size: Maximum AMQP frame size. Default is 63488 bytes. + :paramtype max_frame_size: int + :keyword channel_max: Maximum number of Session channels in the Connection. + :paramtype channel_max: int + :keyword idle_timeout: Timeout in seconds after which the Connection will close + if there is no further activity. + :paramtype idle_timeout: int + :keyword auth_timeout: Timeout in seconds for CBS authentication. Otherwise this value will be ignored. + Default value is 60s. + :paramtype auth_timeout: int + :keyword properties: Connection properties. + :paramtype properties: dict[str, any] + :keyword remote_idle_timeout_empty_frame_send_ratio: Ratio of empty frames to + idle time for Connections with no activity. Value must be between + 0.0 and 1.0 inclusive. Default is 0.5. + :paramtype remote_idle_timeout_empty_frame_send_ratio: float + :keyword incoming_window: The size of the allowed window for incoming messages. + :paramtype incoming_window: int + :keyword outgoing_window: The size of the allowed window for outgoing messages. + :paramtype outgoing_window: int + :keyword handle_max: The maximum number of concurrent link handles. + :paramtype handle_max: int + :keyword on_attach: A callback function to be run on receipt of an ATTACH frame. + The function must take 4 arguments: source, target, properties and error. + :paramtype on_attach: func[ + ~pyamqp.endpoint.Source, ~pyamqp.endpoint.Target, dict, ~pyamqp.error.AMQPConnectionError] + :keyword send_settle_mode: The mode by which to settle message send + operations. If set to `Unsettled`, the client will wait for a confirmation + from the service that the message was successfully sent. If set to 'Settled', + the client will not wait for confirmation and assume success. + :paramtype send_settle_mode: ~pyamqp.constants.SenderSettleMode + :keyword receive_settle_mode: The mode by which to settle message receive + operations. If set to `PeekLock`, the receiver will lock a message once received until + the client accepts or rejects the message. If set to `ReceiveAndDelete`, the service + will assume successful receipt of the message and clear it from the queue. The + default is `PeekLock`. + :paramtype receive_settle_mode: ~pyamqp.constants.ReceiverSettleMode + :keyword desired_capabilities: The extension capabilities desired from the peer endpoint. + :paramtype desired_capabilities: list[bytes] + :keyword max_message_size: The maximum allowed message size negotiated for the Link. + :paramtype max_message_size: int + :keyword link_properties: Metadata to be sent in the Link ATTACH frame. + :paramtype link_properties: dict[str, any] + :keyword link_credit: The Link credit that determines how many + messages the Link will attempt to handle per connection iteration. + The default is 300. + :paramtype link_credit: int + :keyword transport_type: The type of transport protocol that will be used for communicating with + the service. Default is `TransportType.Amqp` in which case port 5671 is used. + If the port 5671 is unavailable/blocked in the network environment, `TransportType.AmqpOverWebsocket` could + be used instead which uses port 443 for communication. + :paramtype transport_type: ~pyamqp.constants.TransportType + :keyword http_proxy: HTTP proxy settings. This must be a dictionary with the following + keys: `'proxy_hostname'` (str value) and `'proxy_port'` (int value). + Additionally the following keys may also be present: `'username', 'password'`. + :paramtype http_proxy: dict[str, str] + :keyword custom_endpoint_address: The custom endpoint address to use for establishing a connection to + the Event Hubs service, allowing network requests to be routed through any application gateways or + other paths needed for the host environment. Default is None. + If port is not specified in the `custom_endpoint_address`, by default port 443 will be used. + :paramtype custom_endpoint_address: str + :keyword connection_verify: Path to the custom CA_BUNDLE file of the SSL certificate which is used to + authenticate the identity of the connection endpoint. + Default is None in which case `certifi.where()` will be used. + :paramtype connection_verify: str + """ + + async def _client_ready_async(self): + """Determine whether the client is ready to start receiving messages. + To be ready, the connection must be open and authentication complete, + The Session, Link and MessageReceiver must be open and in non-errored + states. + + :rtype: bool + """ + # pylint: disable=protected-access + if not self._link: + self._link = self._session.create_receiver_link( + source_address=self.source, + link_credit=self._link_credit, + send_settle_mode=self._send_settle_mode, + rcv_settle_mode=self._receive_settle_mode, + max_message_size=self._max_message_size, + on_transfer=self._message_received_async, + properties=self._link_properties, + desired_capabilities=self._desired_capabilities, + on_attach=self._on_attach + ) + await self._link.attach() + return False + if self._link.get_state().value != 3: # ATTACHED + return False + return True + + async def _client_run_async(self, **kwargs): + """MessageReceiver Link is now open - start receiving messages. + Will return True if operation successful and client can remain open for + further work. + + :rtype: bool + """ + try: + await self._link.flow() + await self._connection.listen(wait=self._socket_timeout, **kwargs) + except ValueError: + _logger.info("Timeout reached, closing receiver.", extra=self._network_trace_params) + self._shutdown = True + return False + return True + + async def _message_received_async(self, frame, message): + """Callback run on receipt of every message. If there is + a user-defined callback, this will be called. + Additionally if the client is retrieving messages for a batch + or iterator, the message will be added to an internal queue. + + :param message: Received message. + :type message: ~pyamqp.message.Message + """ + if self._message_received_callback: + await self._message_received_callback(message) + if not self._streaming_receive: + self._received_messages.put((frame, message)) + + async def _receive_message_batch_impl_async(self, max_batch_size=None, on_message_received=None, timeout=0): + self._message_received_callback = on_message_received + max_batch_size = max_batch_size or self._link_credit + timeout_time = time.time() + timeout if timeout else 0 + receiving = True + batch = [] + await self.open_async() + while len(batch) < max_batch_size: + try: + # TODO: This drops the transfer frame data + _, message = self._received_messages.get_nowait() + batch.append(message) + self._received_messages.task_done() + except queue.Empty: + break + else: + return batch + + to_receive_size = max_batch_size - len(batch) + before_queue_size = self._received_messages.qsize() + + while receiving and to_receive_size > 0: + now_time = time.time() + if timeout_time and now_time > timeout_time: + break + + try: + receiving = await asyncio.wait_for( + self.do_work_async(batch=to_receive_size), + timeout=timeout_time - now_time if timeout else None + ) + except asyncio.TimeoutError: + break + + cur_queue_size = self._received_messages.qsize() + # after do_work, check how many new messages have been received since previous iteration + received = cur_queue_size - before_queue_size + if to_receive_size < max_batch_size and received == 0: + # there are already messages in the batch, and no message is received in the current cycle + # return what we have + break + + to_receive_size -= received + before_queue_size = cur_queue_size + + while len(batch) < max_batch_size: + try: + _, message = self._received_messages.get_nowait() + batch.append(message) + self._received_messages.task_done() + except queue.Empty: + break + return batch + + async def close_async(self): + self._received_messages = queue.Queue() + await super(ReceiveClientAsync, self).close_async() + + async def receive_message_batch_async(self, **kwargs): + """Receive a batch of messages. Messages returned in the batch have already been + accepted - if you wish to add logic to accept or reject messages based on custom + criteria, pass in a callback. This method will return as soon as some messages are + available rather than waiting to achieve a specific batch size, and therefore the + number of messages returned per call will vary up to the maximum allowed. + + :keyword max_batch_size: The maximum number of messages that can be returned in + one call. This value cannot be larger than the prefetch value, and if not specified, + the prefetch value will be used. + :paramtype max_batch_size: int + :keyword on_message_received: A callback to process messages as they arrive from the + service. It takes a single argument, a ~pyamqp.message.Message object. + :paramtype on_message_received: callable[~pyamqp.message.Message] + :keyword timeout: Timeout in seconds for which to wait to receive any messages. + If no messages are received in this time, an empty list will be returned. If set to + 0, the client will continue to wait until at least one message is received. The + default is 0. + :paramtype timeout: float + """ + return await self._do_retryable_operation_async( + self._receive_message_batch_impl_async, + **kwargs + ) + + @overload + async def settle_messages_async( + self, + delivery_id: Union[int, Tuple[int, int]], + outcome: Literal["accepted"], + *, + batchable: Optional[bool] = None + ): + ... + + @overload + async def settle_messages_async( + self, + delivery_id: Union[int, Tuple[int, int]], + outcome: Literal["released"], + *, + batchable: Optional[bool] = None + ): + ... + + @overload + async def settle_messages_async( + self, + delivery_id: Union[int, Tuple[int, int]], + outcome: Literal["rejected"], + *, + error: Optional[AMQPError] = None, + batchable: Optional[bool] = None + ): + ... + + @overload + async def settle_messages_async( + self, + delivery_id: Union[int, Tuple[int, int]], + outcome: Literal["modified"], + *, + delivery_failed: Optional[bool] = None, + undeliverable_here: Optional[bool] = None, + message_annotations: Optional[Dict[Union[str, bytes], Any]] = None, + batchable: Optional[bool] = None + ): + ... + + @overload + async def settle_messages_async( + self, + delivery_id: Union[int, Tuple[int, int]], + outcome: Literal["received"], + *, + section_number: int, + section_offset: int, + batchable: Optional[bool] = None + ): + ... + + async def settle_messages_async(self, delivery_id: Union[int, Tuple[int, int]], outcome: str, **kwargs): + batchable = kwargs.pop('batchable', None) + if outcome.lower() == 'accepted': + state: Outcomes = Accepted() + elif outcome.lower() == 'released': + state = Released() + elif outcome.lower() == 'rejected': + state = Rejected(**kwargs) + elif outcome.lower() == 'modified': + state = Modified(**kwargs) + elif outcome.lower() == 'received': + state = Received(**kwargs) + else: + raise ValueError("Unrecognized message output: {}".format(outcome)) + try: + first, last = cast(Tuple, delivery_id) + except TypeError: + first = delivery_id + last = None + await self._link.send_disposition( + first_delivery_id=first, + last_delivery_id=last, + settled=True, + delivery_state=state, + batchable=batchable, + wait=True + ) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_connection_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_connection_async.py new file mode 100644 index 000000000000..aaaf43bfe420 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_connection_async.py @@ -0,0 +1,870 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + +import uuid +import logging +import time +from urllib.parse import urlparse +import socket +from ssl import SSLError +import asyncio +from typing import Any, Tuple, Optional, NamedTuple, Union, cast + +from ._transport_async import AsyncTransport +from ._sasl_async import SASLTransport, SASLWithWebSocket +from ._session_async import Session +from ..performatives import OpenFrame, CloseFrame +from .._connection import get_local_timeout, _CLOSING_STATES +from ..constants import ( + PORT, + SECURE_PORT, + WEBSOCKET_PORT, + MAX_CHANNELS, + MAX_FRAME_SIZE_BYTES, + HEADER_FRAME, + ConnectionState, + EMPTY_FRAME, + TransportType, +) + +from ..error import ErrorCondition, AMQPConnectionError, AMQPError + +_LOGGER = logging.getLogger(__name__) + + +class Connection(object): # pylint:disable=too-many-instance-attributes + """An AMQP Connection. + + :ivar str state: The connection state. + :param str endpoint: The endpoint to connect to. Must be fully qualified with scheme and port number. + :keyword str container_id: The ID of the source container. If not set a GUID will be generated. + :keyword int max_frame_size: Proposed maximum frame size in bytes. Default value is 64kb. + :keyword int channel_max: The maximum channel number that may be used on the Connection. Default value is 65535. + :keyword int idle_timeout: Connection idle time-out in seconds. + :keyword list(str) outgoing_locales: Locales available for outgoing text. + :keyword list(str) incoming_locales: Desired locales for incoming text in decreasing level of preference. + :keyword list(str) offered_capabilities: The extension capabilities the sender supports. + :keyword list(str) desired_capabilities: The extension capabilities the sender may use if the receiver supports + :keyword dict properties: Connection properties. + :keyword bool allow_pipelined_open: Allow frames to be sent on the connection before a response Open frame + has been received. Default value is `True`. + :keyword float idle_timeout_empty_frame_send_ratio: Portion of the idle timeout time to wait before sending an + empty frame. The default portion is 50% of the idle timeout value (i.e. `0.5`). + :keyword float idle_wait_time: The time in seconds to sleep while waiting for a response from the endpoint. + Default value is `0.1`. + :keyword bool network_trace: Whether to log the network traffic. Default value is `False`. If enabled, frames + will be logged at the logging.INFO level. + :keyword str transport_type: Determines if the transport type is Amqp or AmqpOverWebSocket. + Defaults to TransportType.Amqp. It will be AmqpOverWebSocket if using http_proxy. + :keyword Dict http_proxy: HTTP proxy settings. This must be a dictionary with the following + keys: `'proxy_hostname'` (str value) and `'proxy_port'` (int value). When using these settings, + the transport_type would be AmqpOverWebSocket. + Additionally the following keys may also be present: `'username', 'password'`. + """ + + def __init__(self, endpoint, **kwargs): # pylint:disable=too-many-statements + # type(str, Any) -> None + parsed_url = urlparse(endpoint) + self._hostname = parsed_url.hostname + endpoint = self._hostname + if parsed_url.port: + self._port = parsed_url.port + elif parsed_url.scheme == "amqps": + self._port = SECURE_PORT + else: + self._port = PORT + self.state = None # type: Optional[ConnectionState] + + # Custom Endpoint + custom_endpoint_address = kwargs.get("custom_endpoint_address") + custom_endpoint = None + if custom_endpoint_address: + custom_parsed_url = urlparse(custom_endpoint_address) + custom_port = custom_parsed_url.port or WEBSOCKET_PORT + custom_endpoint = f"{custom_parsed_url.hostname}:{custom_port}{custom_parsed_url.path}" + self._container_id = kwargs.pop("container_id", None) or str( + uuid.uuid4() + ) # type: str + self._network_trace = kwargs.get("network_trace", False) + self._network_trace_params = {"amqpConnection": self._container_id, "amqpSession": None, "amqpLink": None} + + transport = kwargs.get("transport") + self._transport_type = kwargs.pop("transport_type", TransportType.Amqp) + if transport: + self._transport = transport + elif "sasl_credential" in kwargs: + sasl_transport = SASLTransport + if self._transport_type.name == "AmqpOverWebsocket" or kwargs.get( + "http_proxy" + ): + sasl_transport = SASLWithWebSocket + endpoint = parsed_url.hostname + parsed_url.path + self._transport = sasl_transport( + host=endpoint, + credential=kwargs["sasl_credential"], + custom_endpoint=custom_endpoint, + network_trace_params=self._network_trace_params, + **kwargs, + ) + else: + self._transport = AsyncTransport( + parsed_url.netloc, + network_trace_params=self._network_trace_params, + **kwargs) + + self._max_frame_size = kwargs.pop( + "max_frame_size", MAX_FRAME_SIZE_BYTES + ) # type: int + self._remote_max_frame_size = None # type: Optional[int] + self._channel_max = kwargs.pop("channel_max", MAX_CHANNELS) # type: int + self._idle_timeout = kwargs.pop("idle_timeout", None) # type: Optional[int] + self._outgoing_locales = kwargs.pop( + "outgoing_locales", None + ) # type: Optional[List[str]] + self._incoming_locales = kwargs.pop( + "incoming_locales", None + ) # type: Optional[List[str]] + self._offered_capabilities = None # type: Optional[str] + self._desired_capabilities = kwargs.pop( + "desired_capabilities", None + ) # type: Optional[str] + self._properties = kwargs.pop( + "properties", None + ) # type: Optional[Dict[str, str]] + + self._allow_pipelined_open = kwargs.pop( + "allow_pipelined_open", True + ) # type: bool + self._remote_idle_timeout = None # type: Optional[int] + self._remote_idle_timeout_send_frame = None # type: Optional[int] + self._idle_timeout_empty_frame_send_ratio = kwargs.get( + "idle_timeout_empty_frame_send_ratio", 0.5 + ) + self._last_frame_received_time = None # type: Optional[float] + self._last_frame_sent_time = None # type: Optional[float] + self._idle_wait_time = kwargs.get("idle_wait_time", 0.1) # type: float + self._error = None + self._outgoing_endpoints = {} # type: Dict[int, Session] + self._incoming_endpoints = {} # type: Dict[int, Session] + + async def __aenter__(self): + await self.open() + return self + + async def __aexit__(self, *args): + await self.close() + + async def _set_state(self, new_state): + # type: (ConnectionState) -> None + """Update the connection state.""" + if new_state is None: + return + previous_state = self.state + self.state = new_state + _LOGGER.info( + "Connection state changed: %r -> %r", + previous_state, + new_state, + extra=self._network_trace_params + ) + for session in self._outgoing_endpoints.values(): + await session._on_connection_state_change() # pylint:disable=protected-access + + async def _connect(self): + # type: () -> None + """Initiate the connection. + + If `allow_pipelined_open` is enabled, the incoming response header will be processed immediately + and the state on exiting will be HDR_EXCH. Otherwise, the function will return before waiting for + the response header and the final state will be HDR_SENT. + + :raises ValueError: If a reciprocating protocol header is not received during negotiation. + """ + try: + if not self.state: + await self._transport.connect() + await self._set_state(ConnectionState.START) + await self._transport.negotiate() + await self._outgoing_header() + await self._set_state(ConnectionState.HDR_SENT) + if not self._allow_pipelined_open: + await self._read_frame(wait=True) + if self.state != ConnectionState.HDR_EXCH: + await self._disconnect() + raise ValueError( + "Did not receive reciprocal protocol header. Disconnecting." + ) + else: + await self._set_state(ConnectionState.HDR_SENT) + except (OSError, IOError, SSLError, socket.error, asyncio.TimeoutError) as exc: + # FileNotFoundError is being raised for exception parity with uamqp when invalid + # `connection_verify` file path is passed in. Remove later when resolving issue #27128. + if isinstance(exc, FileNotFoundError) and exc.filename and "ca_certs" in exc.filename: + raise + raise AMQPConnectionError( + ErrorCondition.SocketError, + description="Failed to initiate the connection due to exception: " + + str(exc), + error=exc, + ) + + async def _disconnect(self) -> None: + """Disconnect the transport and set state to END.""" + if self.state == ConnectionState.END: + return + await self._set_state(ConnectionState.END) + await self._transport.close() + + def _can_read(self): + # type: () -> bool + """Whether the connection is in a state where it is legal to read for incoming frames.""" + return self.state not in (ConnectionState.CLOSE_RCVD, ConnectionState.END) + + async def _read_frame(self, wait: Union[bool, int, float] = True, **kwargs) -> bool: + """Read an incoming frame from the transport. + + :param Union[bool, float] wait: Whether to block on the socket while waiting for an incoming frame. + The default value is `False`, where the frame will block for the configured timeout only (0.1 seconds). + If set to `True`, socket will block indefinitely. If set to a timeout value in seconds, the socket will + block for at most that value. + :rtype: Tuple[int, Optional[Tuple[int, NamedTuple]]] + :returns: A tuple with the incoming channel number, and the frame in the form or a tuple of performative + descriptor and field values. + """ + timeout: Optional[Union[int, float]] = None + if wait is False: + timeout = 1 # TODO: What should this default be? + elif wait is True: + timeout = None + else: + timeout = wait + new_frame = await self._transport.receive_frame(timeout=timeout, **kwargs) + return await self._process_incoming_frame(*new_frame) + + def _can_write(self): + # type: () -> bool + """Whether the connection is in a state where it is legal to write outgoing frames.""" + return self.state not in _CLOSING_STATES + + async def _send_frame(self, channel, frame, timeout=None, **kwargs): + # type: (int, NamedTuple, Optional[int], Any) -> None + """Send a frame over the connection. + + :param int channel: The outgoing channel number. + :param NamedTuple: The outgoing frame. + :param int timeout: An optional timeout value to wait until the socket is ready to send the frame. + :rtype: None + """ + try: + raise self._error + except TypeError: + pass + + if self._can_write(): + try: + self._last_frame_sent_time = time.time() + await asyncio.wait_for( + self._transport.send_frame(channel, frame, **kwargs), + timeout=timeout, + ) + except ( + OSError, + IOError, + SSLError, + socket.error, + asyncio.TimeoutError, + ) as exc: + self._error = AMQPConnectionError( + ErrorCondition.SocketError, + description="Can not send frame out due to exception: " + str(exc), + error=exc, + ) + else: + _LOGGER.info("Cannot write frame in current state: %r", self.state, extra=self._network_trace_params) + + def _get_next_outgoing_channel(self): + # type: () -> int + """Get the next available outgoing channel number within the max channel limit. + + :raises ValueError: If maximum channels has been reached. + :returns: The next available outgoing channel number. + :rtype: int + """ + if ( + len(self._incoming_endpoints) + len(self._outgoing_endpoints) + ) >= self._channel_max: + raise ValueError( + "Maximum number of channels ({}) has been reached.".format( + self._channel_max + ) + ) + next_channel = next( + i for i in range(1, self._channel_max) if i not in self._outgoing_endpoints + ) + return next_channel + + async def _outgoing_empty(self): + # type: () -> None + """Send an empty frame to prevent the connection from reaching an idle timeout.""" + if self._network_trace: + _LOGGER.debug("-> EmptyFrame()", extra=self._network_trace_params) + try: + raise self._error + except TypeError: + pass + try: + if self._can_write(): + await self._transport.write(EMPTY_FRAME) + self._last_frame_sent_time = time.time() + except (OSError, IOError, SSLError, socket.error) as exc: + self._error = AMQPConnectionError( + ErrorCondition.SocketError, + description="Can not send empty frame due to exception: " + str(exc), + error=exc, + ) + + async def _outgoing_header(self): + # type: () -> None + """Send the AMQP protocol header to initiate the connection.""" + self._last_frame_sent_time = time.time() + if self._network_trace: + _LOGGER.debug("-> Header(%r)", HEADER_FRAME, extra=self._network_trace_params) + await self._transport.write(HEADER_FRAME) + + async def _incoming_header(self, _, frame): + # type: (int, bytes) -> None + """Process an incoming AMQP protocol header and update the connection state.""" + if self._network_trace: + _LOGGER.debug("<- Header(%r)", frame, extra=self._network_trace_params) + if self.state == ConnectionState.START: + await self._set_state(ConnectionState.HDR_RCVD) + elif self.state == ConnectionState.HDR_SENT: + await self._set_state(ConnectionState.HDR_EXCH) + elif self.state == ConnectionState.OPEN_PIPE: + await self._set_state(ConnectionState.OPEN_SENT) + + async def _outgoing_open(self): + # type: () -> None + """Send an Open frame to negotiate the AMQP connection functionality.""" + open_frame = OpenFrame( + container_id=self._container_id, + hostname=self._hostname, + max_frame_size=self._max_frame_size, + channel_max=self._channel_max, + idle_timeout=self._idle_timeout * 1000 + if self._idle_timeout + else None, # Convert to milliseconds + outgoing_locales=self._outgoing_locales, + incoming_locales=self._incoming_locales, + offered_capabilities=self._offered_capabilities + if self.state == ConnectionState.OPEN_RCVD + else None, + desired_capabilities=self._desired_capabilities + if self.state == ConnectionState.HDR_EXCH + else None, + properties=self._properties, + ) + if self._network_trace: + _LOGGER.debug("-> %r", open_frame, extra=self._network_trace_params) + await self._send_frame(0, open_frame) + + async def _incoming_open(self, channel, frame): + # type: (int, Tuple[Any, ...]) -> None + """Process incoming Open frame to finish the connection negotiation. + + The incoming frame format is:: + + - frame[0]: container_id (str) + - frame[1]: hostname (str) + - frame[2]: max_frame_size (int) + - frame[3]: channel_max (int) + - frame[4]: idle_timeout (Optional[int]) + - frame[5]: outgoing_locales (Optional[List[bytes]]) + - frame[6]: incoming_locales (Optional[List[bytes]]) + - frame[7]: offered_capabilities (Optional[List[bytes]]) + - frame[8]: desired_capabilities (Optional[List[bytes]]) + - frame[9]: properties (Optional[Dict[bytes, bytes]]) + + :param int channel: The incoming channel number. + :param frame: The incoming Open frame. + :type frame: Tuple[Any, ...] + :rtype: None + """ + # TODO: Add type hints for full frame tuple contents. + if self._network_trace: + _LOGGER.debug("<- %r", OpenFrame(*frame), extra=self._network_trace_params) + if channel != 0: + _LOGGER.error("OPEN frame received on a channel that is not 0.", extra=self._network_trace_params) + await self.close( + error=AMQPError( + condition=ErrorCondition.NotAllowed, + description="OPEN frame received on a channel that is not 0.", + ) + ) + await self._set_state(ConnectionState.END) + if self.state == ConnectionState.OPENED: + _LOGGER.error("OPEN frame received in the OPENED state.", extra=self._network_trace_params) + await self.close() + if frame[4]: + self._remote_idle_timeout = frame[4] / 1000 # Convert to seconds + self._remote_idle_timeout_send_frame = ( + self._idle_timeout_empty_frame_send_ratio * self._remote_idle_timeout + ) + + if frame[2] < 512: + # Max frame size is less than supported minimum + # If any of the values in the received open frame are invalid then the connection shall be closed. + # The error amqp:invalid-field shall be set in the error.condition field of the CLOSE frame. + await self.close( + error=AMQPError( + condition=ErrorCondition.InvalidField, + description="Failed parsing OPEN frame: Max frame size is less than supported minimum.", + ) + ) + _LOGGER.error( + "Failed parsing OPEN frame: Max frame size is less than supported minimum.", + extra=self._network_trace_params + ) + return + self._remote_max_frame_size = frame[2] + if self.state == ConnectionState.OPEN_SENT: + await self._set_state(ConnectionState.OPENED) + elif self.state == ConnectionState.HDR_EXCH: + await self._set_state(ConnectionState.OPEN_RCVD) + await self._outgoing_open() + await self._set_state(ConnectionState.OPENED) + else: + await self.close( + error=AMQPError( + condition=ErrorCondition.IllegalState, + description=f"Connection is an illegal state: {self.state}", + ) + ) + _LOGGER.error("Connection is an illegal state: %r", self.state, extra=self._network_trace_params) + + async def _outgoing_close(self, error=None): + # type: (Optional[AMQPError]) -> None + """Send a Close frame to shutdown connection with optional error information.""" + close_frame = CloseFrame(error=error) + if self._network_trace: + _LOGGER.debug("-> %r", close_frame, extra=self._network_trace_params) + await self._send_frame(0, close_frame) + + async def _incoming_close(self, channel, frame): + # type: (int, Tuple[Any, ...]) -> None + """Process incoming Open frame to finish the connection negotiation. + + The incoming frame format is:: + + - frame[0]: error (Optional[AMQPError]) + + """ + if self._network_trace: + _LOGGER.debug("<- %r", CloseFrame(*frame), extra=self._network_trace_params) + disconnect_states = [ + ConnectionState.HDR_RCVD, + ConnectionState.HDR_EXCH, + ConnectionState.OPEN_RCVD, + ConnectionState.CLOSE_SENT, + ConnectionState.DISCARDING, + ] + if self.state in disconnect_states: + await self._disconnect() + return + + close_error = None + if channel > self._channel_max: + _LOGGER.error( + "CLOSE frame received on a channel greated than support max.", + extra=self._network_trace_params + ) + close_error = AMQPError( + condition=ErrorCondition.InvalidField, + description="Invalid channel", + info=None, + ) + + await self._set_state(ConnectionState.CLOSE_RCVD) + await self._outgoing_close(error=close_error) + await self._disconnect() + + if frame[0]: + self._error = AMQPConnectionError( + condition=frame[0][0], description=frame[0][1], info=frame[0][2] + ) + _LOGGER.error( + "Connection closed with error: %r", frame[0], + extra=self._network_trace_params + ) + + async def _incoming_begin(self, channel, frame): + # type: (int, Tuple[Any, ...]) -> None + """Process incoming Begin frame to finish negotiating a new session. + + The incoming frame format is:: + + - frame[0]: remote_channel (int) + - frame[1]: next_outgoing_id (int) + - frame[2]: incoming_window (int) + - frame[3]: outgoing_window (int) + - frame[4]: handle_max (int) + - frame[5]: offered_capabilities (Optional[List[bytes]]) + - frame[6]: desired_capabilities (Optional[List[bytes]]) + - frame[7]: properties (Optional[Dict[bytes, bytes]]) + + :param int channel: The incoming channel number. + :param frame: The incoming Begin frame. + :type frame: Tuple[Any, ...] + :rtype: None + """ + try: + existing_session = self._outgoing_endpoints[frame[0]] + self._incoming_endpoints[channel] = existing_session + await self._incoming_endpoints[channel]._incoming_begin( # pylint:disable=protected-access + frame + ) + except KeyError: + new_session = Session.from_incoming_frame(self, channel) + self._incoming_endpoints[channel] = new_session + await new_session._incoming_begin(frame) # pylint:disable=protected-access + + async def _incoming_end(self, channel, frame): + # type: (int, Tuple[Any, ...]) -> None + """Process incoming End frame to close a session. + + The incoming frame format is:: + + - frame[0]: error (Optional[AMQPError]) + + :param int channel: The incoming channel number. + :param frame: The incoming End frame. + :type frame: Tuple[Any, ...] + :rtype: None + """ + try: + await self._incoming_endpoints[channel]._incoming_end(frame) # pylint:disable=protected-access + self._incoming_endpoints.pop(channel) + self._outgoing_endpoints.pop(channel) + except KeyError: + #close the connection + await self.close( + error=AMQPError( + condition=ErrorCondition.ConnectionCloseForced, + description="Invalid channel number received" + )) + _LOGGER.error( + "END frame received on invalid channel. Closing connection.", + extra=self._network_trace_params + ) + return + + async def _process_incoming_frame( + self, channel, frame + ): # pylint:disable=too-many-return-statements + # type: (int, Optional[Union[bytes, Tuple[int, Tuple[Any, ...]]]]) -> bool + """Process an incoming frame, either directly or by passing to the necessary Session. + + :param int channel: The channel the frame arrived on. + :param frame: A tuple containing the performative descriptor and the field values of the frame. + This parameter can be None in the case of an empty frame or a socket timeout. + :type frame: Optional[Tuple[int, NamedTuple]] + :rtype: bool + :returns: A boolean to indicate whether more frames in a batch can be processed or whether the + incoming frame has altered the state. If `True` is returned, the state has changed and the batch + should be interrupted. + """ + try: + performative, fields = cast(Union[bytes, Tuple], frame) + except TypeError: + return True # Empty Frame or socket timeout + fields = cast(Tuple[Any, ...], fields) + try: + self._last_frame_received_time = time.time() + if performative == 20: + await self._incoming_endpoints[channel]._incoming_transfer( # pylint:disable=protected-access + fields + ) + return False + if performative == 21: + await self._incoming_endpoints[channel]._incoming_disposition( # pylint:disable=protected-access + fields + ) + return False + if performative == 19: + await self._incoming_endpoints[channel]._incoming_flow( # pylint:disable=protected-access + fields + ) + return False + if performative == 18: + await self._incoming_endpoints[channel]._incoming_attach( # pylint:disable=protected-access + fields + ) + return False + if performative == 22: + await self._incoming_endpoints[channel]._incoming_detach( # pylint:disable=protected-access + fields + ) + return True + if performative == 17: + await self._incoming_begin(channel, fields) + return True + if performative == 23: + await self._incoming_end(channel, fields) + return True + if performative == 16: + await self._incoming_open(channel, fields) + return True + if performative == 24: + await self._incoming_close(channel, fields) + return True + if performative == 0: + await self._incoming_header(channel, cast(bytes, fields)) + return True + if performative == 1: + return False # TODO: incoming EMPTY + _LOGGER.error("Unrecognized incoming frame: %r", frame, extra=self._network_trace_params) + return True + except KeyError: + return True # TODO: channel error + + async def _process_outgoing_frame(self, channel, frame): + # type: (int, NamedTuple) -> None + """Send an outgoing frame if the connection is in a legal state. + + :raises ValueError: If the connection is not open or not in a valid state. + """ + if not self._allow_pipelined_open and self.state in [ + ConnectionState.OPEN_PIPE, + ConnectionState.OPEN_SENT, + ]: + raise ValueError("Connection not configured to allow pipeline send.") + if self.state not in [ + ConnectionState.OPEN_PIPE, + ConnectionState.OPEN_SENT, + ConnectionState.OPENED, + ]: + raise ValueError("Connection not open.") + now = time.time() + if get_local_timeout( + now, + cast(float, self._idle_timeout), + cast(float, self._last_frame_received_time), + ) or (await self._get_remote_timeout(now)): + _LOGGER.info( + "No frame received for the idle timeout. Closing connection.", + extra=self._network_trace_params + ) + await self.close( + error=AMQPError( + condition=ErrorCondition.ConnectionCloseForced, + description="No frame received for the idle timeout.", + ), + wait=False, + ) + return + await self._send_frame(channel, frame) + + async def _get_remote_timeout(self, now): + # type: (float) -> bool + """Check whether the local connection has reached the remote endpoints idle timeout since + the last outgoing frame was sent. + + If the time since the last since frame is greater than the allowed idle interval, an Empty + frame will be sent to maintain the connection. + + :param float now: The current time to check against. + :rtype: bool + :returns: Whether the local connection should be shutdown due to timeout. + """ + if self._remote_idle_timeout and self._last_frame_sent_time: + time_since_last_sent = now - self._last_frame_sent_time + if time_since_last_sent > cast(int, self._remote_idle_timeout_send_frame): + await self._outgoing_empty() + return False + + async def _wait_for_response(self, wait, end_state): + # type: (Union[bool, float], ConnectionState) -> None + """Wait for an incoming frame to be processed that will result in a desired state change. + + :param wait: Whether to wait for an incoming frame to be processed. Can be set to `True` to wait + indefinitely, or an int to wait for a specified amount of time (in seconds). To not wait, set to `False`. + :type wait: bool or float + :param ConnectionState end_state: The desired end state to wait until. + :rtype: None + """ + if wait is True: + await self.listen(wait=False) + while self.state != end_state: + await asyncio.sleep(self._idle_wait_time) + await self.listen(wait=False) + elif wait: + await self.listen(wait=False) + timeout = time.time() + wait + while self.state != end_state: + if time.time() >= timeout: + break + await asyncio.sleep(self._idle_wait_time) + await self.listen(wait=False) + + async def listen(self, wait=False, batch=1, **kwargs): + # type: (Union[float, int, bool], int, Any) -> None + """Listen on the socket for incoming frames and process them. + + :param wait: Whether to block on the socket until a frame arrives. If set to `True`, socket will + block indefinitely. Alternatively, if set to a time in seconds, the socket will block for at most + the specified timeout. Default value is `False`, where the socket will block for its configured read + timeout (by default 0.1 seconds). + :type wait: int or float or bool + :param int batch: The number of frames to attempt to read and process before returning. The default value + is 1, i.e. process frames one-at-a-time. A higher value should only be used when a receiver is established + and is processing incoming Transfer frames. + :rtype: None + """ + try: + raise self._error + except TypeError: + pass + try: + if self.state not in _CLOSING_STATES: + now = time.time() + if get_local_timeout( + now, + cast(float, self._idle_timeout), + cast(float, self._last_frame_received_time), + ) or (await self._get_remote_timeout(now)): + _LOGGER.info( + "No frame received for the idle timeout. Closing connection.", + extra=self._network_trace_params + ) + await self.close( + error=AMQPError( + condition=ErrorCondition.ConnectionCloseForced, + description="No frame received for the idle timeout.", + ), + wait=False, + ) + return + if self.state == ConnectionState.END: + # TODO: check error condition + self._error = AMQPConnectionError( + condition=ErrorCondition.ConnectionCloseForced, + description="Connection was already closed.", + ) + return + for _ in range(batch): + if self._can_read(): + if await self._read_frame(wait=wait, **kwargs): + break + else: + _LOGGER.info( + "Connection cannot read frames in this state: %r", + self.state, + extra=self._network_trace_params + ) + break + except (OSError, IOError, SSLError, socket.error) as exc: + self._error = AMQPConnectionError( + ErrorCondition.SocketError, + description="Can not read frame due to exception: " + str(exc), + error=exc, + ) + + def create_session(self, **kwargs): + # type: (Any) -> Session + """Create a new session within this connection. + + :keyword str name: The name of the connection. If not set a GUID will be generated. + :keyword int next_outgoing_id: The transfer-id of the first transfer id the sender will send. + Default value is 0. + :keyword int incoming_window: The initial incoming-window of the Session. Default value is 1. + :keyword int outgoing_window: The initial outgoing-window of the Session. Default value is 1. + :keyword int handle_max: The maximum handle value that may be used on the session. Default value is 4294967295. + :keyword list(str) offered_capabilities: The extension capabilities the session supports. + :keyword list(str) desired_capabilities: The extension capabilities the session may use if + the endpoint supports it. + :keyword dict properties: Session properties. + :keyword bool allow_pipelined_open: Allow frames to be sent on the connection before a response Open frame + has been received. Default value is that configured for the connection. + :keyword float idle_wait_time: The time in seconds to sleep while waiting for a response from the endpoint. + Default value is that configured for the connection. + :keyword bool network_trace: Whether to log the network traffic of this session. If enabled, frames + will be logged at the logging.INFO level. Default value is that configured for the connection. + """ + assigned_channel = self._get_next_outgoing_channel() + kwargs["allow_pipelined_open"] = self._allow_pipelined_open + kwargs["idle_wait_time"] = self._idle_wait_time + session = Session( + self, + assigned_channel, + network_trace=kwargs.pop("network_trace", self._network_trace), + network_trace_params=dict(self._network_trace_params), + **kwargs, + ) + self._outgoing_endpoints[assigned_channel] = session + return session + + async def open(self, wait=False): + # type: (bool) -> None + """Send an Open frame to start the connection. + + Alternatively, this will be called on entering a Connection context manager. + + :param bool wait: Whether to wait to receive an Open response from the endpoint. Default is `False`. + :raises ValueError: If `wait` is set to `False` and `allow_pipelined_open` is disabled. + :rtype: None + """ + await self._connect() + await self._outgoing_open() + if self.state == ConnectionState.HDR_EXCH: + await self._set_state(ConnectionState.OPEN_SENT) + elif self.state == ConnectionState.HDR_SENT: + await self._set_state(ConnectionState.OPEN_PIPE) + if wait: + await self._wait_for_response(wait, ConnectionState.OPENED) + elif not self._allow_pipelined_open: + raise ValueError( + "Connection has been configured to not allow piplined-open. Please set 'wait' parameter." + ) + + async def close(self, error=None, wait=False): + # type: (Optional[AMQPError], bool) -> None + """Close the connection and disconnect the transport. + + Alternatively this method will be called on exiting a Connection context manager. + + :param ~uamqp.AMQPError error: Optional error information to include in the close request. + :param bool wait: Whether to wait for a service Close response. Default is `False`. + :rtype: None + """ + try: + if self.state in [ + ConnectionState.END, + ConnectionState.CLOSE_SENT, + ConnectionState.DISCARDING, + ]: + return + await self._outgoing_close(error=error) + if error: + self._error = AMQPConnectionError( + condition=error.condition, + description=error.description, + info=error.info, + ) + if self.state == ConnectionState.OPEN_PIPE: + await self._set_state(ConnectionState.OC_PIPE) + elif self.state == ConnectionState.OPEN_SENT: + await self._set_state(ConnectionState.CLOSE_PIPE) + elif error: + await self._set_state(ConnectionState.DISCARDING) + else: + await self._set_state(ConnectionState.CLOSE_SENT) + await self._wait_for_response(wait, ConnectionState.END) + except Exception as exc: # pylint:disable=broad-except + # If error happened during closing, ignore the error and set state to END + _LOGGER.info("An error occurred when closing the connection: %r", exc, extra=self._network_trace_params) + await self._set_state(ConnectionState.END) + finally: + await self._disconnect() diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_link_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_link_async.py new file mode 100644 index 000000000000..8b8b015d294e --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_link_async.py @@ -0,0 +1,262 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + +from typing import Optional +import uuid +import logging + +from ..endpoints import Source, Target +from ..constants import DEFAULT_LINK_CREDIT, SessionState, LinkState, Role, SenderSettleMode, ReceiverSettleMode +from ..performatives import ( + AttachFrame, + DetachFrame, +) + +from ..error import ErrorCondition, AMQPLinkError, AMQPLinkRedirect, AMQPConnectionError + +_LOGGER = logging.getLogger(__name__) + + +class Link(object): # pylint: disable=too-many-instance-attributes + """An AMQP Link. + + This object should not be used directly - instead use one of directional + derivatives: Sender or Receiver. + """ + + def __init__(self, session, handle, name, role, **kwargs): + self.state = LinkState.DETACHED + self.name = name or str(uuid.uuid4()) + self.handle = handle + self.remote_handle = None + self.role = role + source_address = kwargs["source_address"] + target_address = kwargs["target_address"] + self.source = ( + source_address + if isinstance(source_address, Source) + else Source( + address=kwargs["source_address"], + durable=kwargs.get("source_durable"), + expiry_policy=kwargs.get("source_expiry_policy"), + timeout=kwargs.get("source_timeout"), + dynamic=kwargs.get("source_dynamic"), + dynamic_node_properties=kwargs.get("source_dynamic_node_properties"), + distribution_mode=kwargs.get("source_distribution_mode"), + filters=kwargs.get("source_filters"), + default_outcome=kwargs.get("source_default_outcome"), + outcomes=kwargs.get("source_outcomes"), + capabilities=kwargs.get("source_capabilities"), + ) + ) + self.target = ( + target_address + if isinstance(target_address, Target) + else Target( + address=kwargs["target_address"], + durable=kwargs.get("target_durable"), + expiry_policy=kwargs.get("target_expiry_policy"), + timeout=kwargs.get("target_timeout"), + dynamic=kwargs.get("target_dynamic"), + dynamic_node_properties=kwargs.get("target_dynamic_node_properties"), + capabilities=kwargs.get("target_capabilities"), + ) + ) + self.link_credit = kwargs.pop("link_credit", None) or DEFAULT_LINK_CREDIT + self.current_link_credit = self.link_credit + self.send_settle_mode = kwargs.pop("send_settle_mode", SenderSettleMode.Mixed) + self.rcv_settle_mode = kwargs.pop("rcv_settle_mode", ReceiverSettleMode.First) + self.unsettled = kwargs.pop("unsettled", None) + self.incomplete_unsettled = kwargs.pop("incomplete_unsettled", None) + self.initial_delivery_count = kwargs.pop("initial_delivery_count", 0) + self.delivery_count = self.initial_delivery_count + self.received_delivery_id = None + self.max_message_size = kwargs.pop("max_message_size", None) + self.remote_max_message_size = None + self.available = kwargs.pop("available", None) + self.properties = kwargs.pop("properties", None) + self.offered_capabilities = None + self.desired_capabilities = kwargs.pop("desired_capabilities", None) + + self.network_trace = kwargs["network_trace"] + self.network_trace_params = kwargs["network_trace_params"] + self.network_trace_params["amqpLink"] = self.name + self._session = session + self._is_closed = False + self._on_link_state_change = kwargs.get("on_link_state_change") + self._on_attach = kwargs.get("on_attach") + self._error = None + + async def __aenter__(self): + await self.attach() + return self + + async def __aexit__(self, *args): + await self.detach(close=True) + + @classmethod + def from_incoming_frame(cls, session, handle, frame): + # check link_create_from_endpoint in C lib + raise NotImplementedError("Pending") # TODO: Assuming we establish all links for now... + + def get_state(self): + try: + raise self._error + except TypeError: + pass + return self.state + + def _check_if_closed(self): + if self._is_closed: + try: + raise self._error + except TypeError: + raise AMQPConnectionError(condition=ErrorCondition.InternalError, description="Link already closed.") + + async def _set_state(self, new_state): + # type: (LinkState) -> None + """Update the session state.""" + if new_state is None: + return + previous_state = self.state + self.state = new_state + _LOGGER.info("Link state changed: %r -> %r", previous_state, new_state, extra=self.network_trace_params) + try: + await self._on_link_state_change(previous_state, new_state) + except TypeError: + pass + except Exception as e: # pylint: disable=broad-except + _LOGGER.error("Link state change callback failed: '%r'", e, extra=self.network_trace_params) + + async def _on_session_state_change(self): + if self._session.state == SessionState.MAPPED: + if not self._is_closed and self.state == LinkState.DETACHED: + await self._outgoing_attach() + await self._set_state(LinkState.ATTACH_SENT) + elif self._session.state == SessionState.DISCARDING: + await self._set_state(LinkState.DETACHED) + + async def _outgoing_attach(self): + self.delivery_count = self.initial_delivery_count + attach_frame = AttachFrame( + name=self.name, + handle=self.handle, + role=self.role, + send_settle_mode=self.send_settle_mode, + rcv_settle_mode=self.rcv_settle_mode, + source=self.source, + target=self.target, + unsettled=self.unsettled, + incomplete_unsettled=self.incomplete_unsettled, + initial_delivery_count=self.initial_delivery_count if self.role == Role.Sender else None, + max_message_size=self.max_message_size, + offered_capabilities=self.offered_capabilities if self.state == LinkState.ATTACH_RCVD else None, + desired_capabilities=self.desired_capabilities if self.state == LinkState.DETACHED else None, + properties=self.properties, + ) + if self.network_trace: + _LOGGER.debug("-> %r", attach_frame, extra=self.network_trace_params) + await self._session._outgoing_attach(attach_frame) # pylint: disable=protected-access + + async def _incoming_attach(self, frame): + if self.network_trace: + _LOGGER.debug("<- %r", AttachFrame(*frame), extra=self.network_trace_params) + if self._is_closed: + raise ValueError("Invalid link") + if not frame[5] or not frame[6]: + _LOGGER.info("Cannot get source or target. Detaching link", extra=self.network_trace_params) + await self._set_state(LinkState.DETACHED) + raise ValueError("Invalid link") + self.remote_handle = frame[1] # handle + self.remote_max_message_size = frame[10] # max_message_size + self.offered_capabilities = frame[11] # offered_capabilities + if self.properties: + self.properties.update(frame[13]) # properties + else: + self.properties = frame[13] + if self.state == LinkState.DETACHED: + await self._set_state(LinkState.ATTACH_RCVD) + elif self.state == LinkState.ATTACH_SENT: + await self._set_state(LinkState.ATTACHED) + if self._on_attach: + try: + if frame[5]: + frame[5] = Source(*frame[5]) + if frame[6]: + frame[6] = Target(*frame[6]) + await self._on_attach(AttachFrame(*frame)) + except Exception as e: # pylint: disable=broad-except + _LOGGER.warning("Callback for link attach raised error: %s", e, extra=self.network_trace_params) + + async def _outgoing_flow(self, **kwargs): + flow_frame = { + "handle": self.handle, + "delivery_count": self.delivery_count, + "link_credit": self.current_link_credit, + "available": kwargs.get("available"), + "drain": kwargs.get("drain"), + "echo": kwargs.get("echo"), + "properties": kwargs.get("properties"), + } + await self._session._outgoing_flow(flow_frame) # pylint: disable=protected-access + + async def _incoming_flow(self, frame): + pass + + async def _incoming_disposition(self, frame): + pass + + async def _outgoing_detach(self, close=False, error=None): + detach_frame = DetachFrame(handle=self.handle, closed=close, error=error) + if self.network_trace: + _LOGGER.debug("-> %r", detach_frame, extra=self.network_trace_params) + await self._session._outgoing_detach(detach_frame) # pylint: disable=protected-access + if close: + self._is_closed = True + + async def _incoming_detach(self, frame): + if self.network_trace: + _LOGGER.debug("<- %r", DetachFrame(*frame), extra=self.network_trace_params) + if self.state == LinkState.ATTACHED: + await self._outgoing_detach(close=frame[1]) # closed + elif frame[1] and not self._is_closed and self.state in [LinkState.ATTACH_SENT, LinkState.ATTACH_RCVD]: + # Received a closing detach after we sent a non-closing detach. + # In this case, we MUST signal that we closed by reattaching and then sending a closing detach. + await self._outgoing_attach() + await self._outgoing_detach(close=True) + # TODO: on_detach_hook + if frame[2]: # error + # frame[2][0] is condition, frame[2][1] is description, frame[2][2] is info + error_cls = AMQPLinkRedirect if frame[2][0] == ErrorCondition.LinkRedirect else AMQPLinkError + self._error = error_cls(condition=frame[2][0], description=frame[2][1], info=frame[2][2]) + await self._set_state(LinkState.ERROR) + else: + await self._set_state(LinkState.DETACHED) + + async def attach(self): + if self._is_closed: + raise ValueError("Link already closed.") + await self._outgoing_attach() + await self._set_state(LinkState.ATTACH_SENT) + + async def detach(self, close=False, error=None): + if self.state in (LinkState.DETACHED, LinkState.ERROR): + return + try: + self._check_if_closed() + if self.state in [LinkState.ATTACH_SENT, LinkState.ATTACH_RCVD]: + await self._outgoing_detach(close=close, error=error) + await self._set_state(LinkState.DETACHED) + elif self.state == LinkState.ATTACHED: + await self._outgoing_detach(close=close, error=error) + await self._set_state(LinkState.DETACH_SENT) + except Exception as exc: # pylint: disable=broad-except + _LOGGER.info("An error occurred when detaching the link: %r", exc, extra=self.network_trace_params) + await self._set_state(LinkState.DETACHED) + + async def flow(self, *, link_credit: Optional[int] = None, **kwargs) -> None: + self.current_link_credit = link_credit if link_credit is not None else self.link_credit + await self._outgoing_flow(**kwargs) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_management_link_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_management_link_async.py new file mode 100644 index 000000000000..94f3163accfd --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_management_link_async.py @@ -0,0 +1,249 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + +import time +import logging +from functools import partial + +from ..management_link import PendingManagementOperation +from ._sender_async import SenderLink +from ._receiver_async import ReceiverLink +from ..constants import ( + ManagementLinkState, + LinkState, + SenderSettleMode, + ReceiverSettleMode, + ManagementExecuteOperationResult, + ManagementOpenResult, + SEND_DISPOSITION_REJECT, + MessageDeliveryState, + LinkDeliverySettleReason +) +from ..error import AMQPException, ErrorCondition +from ..message import Properties, _MessageDelivery + +_LOGGER = logging.getLogger(__name__) + + +class ManagementLink(object): # pylint:disable=too-many-instance-attributes + """ + # TODO: Fill in docstring + """ + + def __init__(self, session, endpoint, **kwargs): + self.next_message_id = 0 + self.state = ManagementLinkState.IDLE + self._pending_operations = [] + self._session = session + self._network_trace_params = kwargs.get('network_trace_params') + self._request_link: SenderLink = session.create_sender_link( + endpoint, + source_address=endpoint, + on_link_state_change=self._on_sender_state_change, + send_settle_mode=SenderSettleMode.Unsettled, + rcv_settle_mode=ReceiverSettleMode.First, + network_trace=kwargs.get("network_trace", False) + ) + self._response_link: ReceiverLink = session.create_receiver_link( + endpoint, + target_address=endpoint, + on_link_state_change=self._on_receiver_state_change, + on_transfer=self._on_message_received, + send_settle_mode=SenderSettleMode.Unsettled, + rcv_settle_mode=ReceiverSettleMode.First, + network_trace=kwargs.get("network_trace", False) + ) + self._on_amqp_management_error = kwargs.get("on_amqp_management_error") + self._on_amqp_management_open_complete = kwargs.get("on_amqp_management_open_complete") + + self._status_code_field = kwargs.get("status_code_field", b"statusCode") + self._status_description_field = kwargs.get("status_description_field", b"statusDescription") + + self._sender_connected = False + self._receiver_connected = False + + async def __aenter__(self): + await self.open() + return self + + async def __aexit__(self, *args): + await self.close() + + async def _on_sender_state_change(self, previous_state, new_state): + _LOGGER.info( + "Management link sender state changed: %r -> %r", + previous_state, + new_state, + extra=self._network_trace_params + ) + if new_state == previous_state: + return + if self.state == ManagementLinkState.OPENING: + if new_state == LinkState.ATTACHED: + self._sender_connected = True + if self._receiver_connected: + self.state = ManagementLinkState.OPEN + await self._on_amqp_management_open_complete(ManagementOpenResult.OK) + elif new_state in [LinkState.DETACHED, LinkState.DETACH_SENT, LinkState.DETACH_RCVD, LinkState.ERROR]: + self.state = ManagementLinkState.IDLE + await self._on_amqp_management_open_complete(ManagementOpenResult.ERROR) + elif self.state == ManagementLinkState.OPEN: + if new_state is not LinkState.ATTACHED: + self.state = ManagementLinkState.ERROR + await self._on_amqp_management_error() + elif self.state == ManagementLinkState.CLOSING: + if new_state not in [LinkState.DETACHED, LinkState.DETACH_SENT, LinkState.DETACH_RCVD]: + self.state = ManagementLinkState.ERROR + await self._on_amqp_management_error() + elif self.state == ManagementLinkState.ERROR: + # All state transitions shall be ignored. + return + + async def _on_receiver_state_change(self, previous_state, new_state): + _LOGGER.info( + "Management link receiver state changed: %r -> %r", + previous_state, + new_state, + extra=self._network_trace_params + ) + if new_state == previous_state: + return + if self.state == ManagementLinkState.OPENING: + if new_state == LinkState.ATTACHED: + self._receiver_connected = True + if self._sender_connected: + self.state = ManagementLinkState.OPEN + await self._on_amqp_management_open_complete(ManagementOpenResult.OK) + elif new_state in [LinkState.DETACHED, LinkState.DETACH_SENT, LinkState.DETACH_RCVD, LinkState.ERROR]: + self.state = ManagementLinkState.IDLE + await self._on_amqp_management_open_complete(ManagementOpenResult.ERROR) + elif self.state == ManagementLinkState.OPEN: + if new_state is not LinkState.ATTACHED: + self.state = ManagementLinkState.ERROR + await self._on_amqp_management_error() + elif self.state == ManagementLinkState.CLOSING: + if new_state not in [LinkState.DETACHED, LinkState.DETACH_SENT, LinkState.DETACH_RCVD]: + self.state = ManagementLinkState.ERROR + await self._on_amqp_management_error() + elif self.state == ManagementLinkState.ERROR: + # All state transitions shall be ignored. + return + + async def _on_message_received(self, _, message): + message_properties = message.properties + correlation_id = message_properties[5] + response_detail = message.application_properties + + status_code = response_detail.get(self._status_code_field) + status_description = response_detail.get(self._status_description_field) + + to_remove_operation = None + for operation in self._pending_operations: + if operation.message.properties.message_id == correlation_id: + to_remove_operation = operation + break + if to_remove_operation: + mgmt_result = ( + ManagementExecuteOperationResult.OK + if 200 <= status_code <= 299 + else ManagementExecuteOperationResult.FAILED_BAD_STATUS + ) + await to_remove_operation.on_execute_operation_complete( + mgmt_result, status_code, status_description, message, response_detail.get(b"error-condition") + ) + self._pending_operations.remove(to_remove_operation) + + async def _on_send_complete(self, message_delivery, reason, state): + if reason == LinkDeliverySettleReason.DISPOSITION_RECEIVED and SEND_DISPOSITION_REJECT in state: + # sample reject state: {'rejected': [[b'amqp:not-allowed', b"Invalid command 'RE1AD'.", None]]} + to_remove_operation = None + for operation in self._pending_operations: + if message_delivery.message == operation.message: + to_remove_operation = operation + break + self._pending_operations.remove(to_remove_operation) + # TODO: better error handling + # AMQPException is too general? to be more specific: MessageReject(Error) or AMQPManagementError? + # or should there an error mapping which maps the condition to the error type + + # The callback is defined in management_operation.py + await to_remove_operation.on_execute_operation_complete( + ManagementExecuteOperationResult.ERROR, + None, + None, + message_delivery.message, + error=AMQPException( + condition=state[SEND_DISPOSITION_REJECT][0][0], # 0 is error condition + description=state[SEND_DISPOSITION_REJECT][0][1], # 1 is error description + info=state[SEND_DISPOSITION_REJECT][0][2], # 2 is error info + ), + ) + + async def open(self): + if self.state != ManagementLinkState.IDLE: + raise ValueError("Management links are already open or opening.") + self.state = ManagementLinkState.OPENING + await self._response_link.attach() + await self._request_link.attach() + + async def execute_operation(self, message, on_execute_operation_complete, **kwargs): + """Execute a request and wait on a response. + + :param message: The message to send in the management request. + :type message: ~uamqp.message.Message + :param on_execute_operation_complete: Callback to be called when the operation is complete. + The following value will be passed to the callback: operation_id, operation_result, status_code, + status_description, raw_message and error. + :type on_execute_operation_complete: Callable[[str, str, int, str, ~uamqp.message.Message, Exception], None] + :keyword operation: The type of operation to be performed. This value will + be service-specific, but common values include READ, CREATE and UPDATE. + This value will be added as an application property on the message. + :paramtype operation: bytes or str + :keyword type: The type on which to carry out the operation. This will + be specific to the entities of the service. This value will be added as + an application property on the message. + :paramtype type: bytes or str + :keyword str locales: A list of locales that the sending peer permits for incoming + informational text in response messages. + :keyword float timeout: Provide an optional timeout in seconds within which a response + to the management request must be received. + :rtype: None + """ + timeout = kwargs.get("timeout") + message.application_properties["operation"] = kwargs.get("operation") + message.application_properties["type"] = kwargs.get("type") + if "locales" in kwargs: + message.application_properties["locales"] = kwargs.get("locales") + try: + # TODO: namedtuple is immutable, which may push us to re-think about the namedtuple approach for Message + new_properties = message.properties._replace(message_id=self.next_message_id) + except AttributeError: + new_properties = Properties(message_id=self.next_message_id) + message = message._replace(properties=new_properties) + expire_time = (time.time() + timeout) if timeout else None + message_delivery = _MessageDelivery(message, MessageDeliveryState.WaitingToBeSent, expire_time) + + on_send_complete = partial(self._on_send_complete, message_delivery) + + await self._request_link.send_transfer(message, on_send_complete=on_send_complete, timeout=timeout) + self.next_message_id += 1 + self._pending_operations.append(PendingManagementOperation(message, on_execute_operation_complete)) + + async def close(self): + if self.state != ManagementLinkState.IDLE: + self.state = ManagementLinkState.CLOSING + await self._response_link.detach(close=True) + await self._request_link.detach(close=True) + for pending_operation in self._pending_operations: + await pending_operation.on_execute_operation_complete( + ManagementExecuteOperationResult.LINK_CLOSED, + None, + None, + pending_operation.message, + AMQPException(condition=ErrorCondition.ClientError, description="Management link already closed."), + ) + self._pending_operations = [] + self.state = ManagementLinkState.IDLE diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_management_operation_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_management_operation_async.py new file mode 100644 index 000000000000..e5830d7d0ff8 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_management_operation_async.py @@ -0,0 +1,140 @@ +#------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +#-------------------------------------------------------------------------- +import logging +import uuid +import time +from functools import partial + +from ._management_link_async import ManagementLink +from ..error import ( + AMQPLinkError, + ErrorCondition +) + +from ..constants import ( + ManagementOpenResult, + ManagementExecuteOperationResult +) + +_LOGGER = logging.getLogger(__name__) + + +class ManagementOperation(object): + def __init__(self, session, endpoint='$management', **kwargs): + self._mgmt_link_open_status = None + + self._session = session + self._connection = self._session._connection + self._network_trace_params = { + "amqpConnection": self._session._connection._container_id, + "amqpSession": self._session.name, + "amqpLink": None + } + self._mgmt_link = self._session.create_request_response_link_pair( + endpoint=endpoint, + on_amqp_management_open_complete=self._on_amqp_management_open_complete, + on_amqp_management_error=self._on_amqp_management_error, + **kwargs + ) # type: ManagementLink + self._responses = {} + self._mgmt_error = None + + async def _on_amqp_management_open_complete(self, result): + """Callback run when the send/receive links are open and ready + to process messages. + + :param result: Whether the link opening was successful. + :type result: int + """ + self._mgmt_link_open_status = result + + async def _on_amqp_management_error(self): + """Callback run if an error occurs in the send/receive links.""" + # TODO: This probably shouldn't be ValueError + self._mgmt_error = ValueError("Management Operation error occurred.") + + async def _on_execute_operation_complete( + self, + operation_id, + operation_result, + status_code, + status_description, + raw_message, + error=None + ): + _LOGGER.debug( + "Management operation completed, id: %r; result: %r; code: %r; description: %r, error: %r", + operation_id, + operation_result, + status_code, + status_description, + error, + extra=self._network_trace_params + ) + + if operation_result in\ + (ManagementExecuteOperationResult.ERROR, ManagementExecuteOperationResult.LINK_CLOSED): + self._mgmt_error = error + _LOGGER.error( + "Failed to complete management operation due to error: %r.", + error, + extra=self._network_trace_params + ) + else: + self._responses[operation_id] = (status_code, status_description, raw_message) + + async def execute(self, message, operation=None, operation_type=None, timeout=0): + start_time = time.time() + operation_id = str(uuid.uuid4()) + self._responses[operation_id] = None + self._mgmt_error = None + + await self._mgmt_link.execute_operation( + message, + partial(self._on_execute_operation_complete, operation_id), + timeout=timeout, + operation=operation, + type=operation_type + ) + + while not self._responses[operation_id] and not self._mgmt_error: + if timeout and timeout > 0: + now = time.time() + if (now - start_time) >= timeout: + raise TimeoutError("Failed to receive mgmt response in {}ms".format(timeout)) + await self._connection.listen() + + if self._mgmt_error: + self._responses.pop(operation_id) + raise self._mgmt_error # pylint: disable=raising-bad-type + + response = self._responses.pop(operation_id) + return response + + async def open(self): + self._mgmt_link_open_status = ManagementOpenResult.OPENING + await self._mgmt_link.open() + + async def ready(self): + try: + raise self._mgmt_error # pylint: disable=raising-bad-type + except TypeError: + pass + + if self._mgmt_link_open_status == ManagementOpenResult.OPENING: + return False + if self._mgmt_link_open_status == ManagementOpenResult.OK: + return True + # ManagementOpenResult.ERROR or CANCELLED + # TODO: update below with correct status code + info + raise AMQPLinkError( + condition=ErrorCondition.ClientError, + description="Failed to open mgmt link, management link status: {}".format(self._mgmt_link_open_status), + info=None + ) + + async def close(self): + await self._mgmt_link.close() diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_receiver_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_receiver_async.py new file mode 100644 index 000000000000..7d3c6c540160 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_receiver_async.py @@ -0,0 +1,124 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + +import uuid +import logging +from typing import Optional, Union + +from .._decode import decode_payload +from ._link_async import Link +from ..constants import LinkState, Role +from ..performatives import ( + TransferFrame, + DispositionFrame, +) +from ..outcomes import Received, Accepted, Rejected, Released, Modified + + +_LOGGER = logging.getLogger(__name__) + + +class ReceiverLink(Link): + def __init__(self, session, handle, source_address, **kwargs): + name = kwargs.pop("name", None) or str(uuid.uuid4()) + role = Role.Receiver + if "target_address" not in kwargs: + kwargs["target_address"] = "receiver-link-{}".format(name) + super(ReceiverLink, self).__init__(session, handle, name, role, source_address=source_address, **kwargs) + self._on_transfer = kwargs.pop("on_transfer") + self._received_payload = bytearray() + + @classmethod + def from_incoming_frame(cls, session, handle, frame): + # TODO: Assuming we establish all links for now... + # check link_create_from_endpoint in C lib + raise NotImplementedError("Pending") + + async def _process_incoming_message(self, frame, message): + try: + return await self._on_transfer(frame, message) + except Exception as e: # pylint: disable=broad-except + _LOGGER.error("Transfer callback function failed with error: %r", e, extra=self.network_trace_params) + return None + + async def _incoming_attach(self, frame): + await super(ReceiverLink, self)._incoming_attach(frame) + if frame[9] is None: # initial_delivery_count + _LOGGER.info("Cannot get initial-delivery-count. Detaching link", extra=self.network_trace_params) + await self._set_state(LinkState.DETACHED) # TODO: Send detach now? + self.delivery_count = frame[9] + self.current_link_credit = self.link_credit + await self._outgoing_flow() + + async def _incoming_transfer(self, frame): + if self.network_trace: + _LOGGER.debug("<- %r", TransferFrame(payload=b"***", *frame[:-1]), extra=self.network_trace_params) + self.current_link_credit -= 1 + self.delivery_count += 1 + self.received_delivery_id = frame[1] # delivery_id + if not self.received_delivery_id and not self._received_payload: + pass # TODO: delivery error + if self._received_payload or frame[5]: # more + self._received_payload.extend(frame[11]) + if not frame[5]: + if self._received_payload: + message = decode_payload(memoryview(self._received_payload)) + self._received_payload = bytearray() + else: + message = decode_payload(frame[11]) + delivery_state = await self._process_incoming_message(frame, message) + if not frame[4] and delivery_state: # settled + await self._outgoing_disposition( + first=frame[1], + last=frame[1], + settled=True, + state=delivery_state, + batchable=None + ) + + async def _wait_for_response(self, wait: Union[bool, float]) -> None: + if wait is True: + await self._session._connection.listen(wait=False) # pylint: disable=protected-access + if self.state == LinkState.ERROR: + raise self._error + elif wait: + await self._session._connection.listen(wait=wait) # pylint: disable=protected-access + if self.state == LinkState.ERROR: + raise self._error + + async def _outgoing_disposition( + self, + first: int, + last: Optional[int], + settled: Optional[bool], + state: Optional[Union[Received, Accepted, Rejected, Released, Modified]], + batchable: Optional[bool], + ): + disposition_frame = DispositionFrame( + role=self.role, first=first, last=last, settled=settled, state=state, batchable=batchable + ) + if self.network_trace: + _LOGGER.debug("-> %r", DispositionFrame(*disposition_frame), extra=self.network_trace_params) + await self._session._outgoing_disposition(disposition_frame) # pylint: disable=protected-access + + async def attach(self): + await super().attach() + self._received_payload = bytearray() + + async def send_disposition( + self, + *, + wait: Union[bool, float] = False, + first_delivery_id: int, + last_delivery_id: Optional[int] = None, + settled: Optional[bool] = None, + delivery_state: Optional[Union[Received, Accepted, Rejected, Released, Modified]] = None, + batchable: Optional[bool] = None + ): + if self._is_closed: + raise ValueError("Link already closed.") + await self._outgoing_disposition(first_delivery_id, last_delivery_id, settled, delivery_state, batchable) + await self._wait_for_response(wait) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_sasl_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_sasl_async.py new file mode 100644 index 000000000000..441eb40ec874 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_sasl_async.py @@ -0,0 +1,149 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + +from ._transport_async import AsyncTransport, WebSocketTransportAsync +from ..constants import SASLCode, SASL_HEADER_FRAME, WEBSOCKET_PORT +from .._transport import AMQPS_PORT +from ..performatives import SASLInit + + +_SASL_FRAME_TYPE = b"\x01" + + +# TODO: do we need it here? it's a duplicate of the sync version +class SASLPlainCredential(object): + """PLAIN SASL authentication mechanism. + See https://tools.ietf.org/html/rfc4616 for details + """ + + mechanism = b"PLAIN" + + def __init__(self, authcid, passwd, authzid=None): + self.authcid = authcid + self.passwd = passwd + self.authzid = authzid + + def start(self): + if self.authzid: + login_response = self.authzid.encode("utf-8") + else: + login_response = b"" + login_response += b"\0" + login_response += self.authcid.encode("utf-8") + login_response += b"\0" + login_response += self.passwd.encode("utf-8") + return login_response + + +# TODO: do we need it here? it's a duplicate of the sync version +class SASLAnonymousCredential(object): + """ANONYMOUS SASL authentication mechanism. + See https://tools.ietf.org/html/rfc4505 for details + """ + + mechanism = b"ANONYMOUS" + + def start(self): # pylint: disable=no-self-use + return b"" + + +# TODO: do we need it here? it's a duplicate of the sync version +class SASLExternalCredential(object): + """EXTERNAL SASL mechanism. + Enables external authentication, i.e. not handled through this protocol. + Only passes 'EXTERNAL' as authentication mechanism, but no further + authentication data. + """ + + mechanism = b"EXTERNAL" + + def start(self): # pylint: disable=no-self-use + return b"" + + +class SASLTransportMixinAsync: # pylint: disable=no-member + async def _negotiate(self): + await self.write(SASL_HEADER_FRAME) + _, returned_header = await self.receive_frame() + if returned_header[1] != SASL_HEADER_FRAME: + raise ValueError( + f"""Mismatching AMQP header protocol. Expected: {SASL_HEADER_FRAME!r},""" + """received: {returned_header[1]!r}""" + ) + + _, supported_mechanisms = await self.receive_frame(verify_frame_type=1) + if ( + self.credential.mechanism not in supported_mechanisms[1][0] + ): # sasl_server_mechanisms + raise ValueError( + "Unsupported SASL credential type: {}".format(self.credential.mechanism) + ) + sasl_init = SASLInit( + mechanism=self.credential.mechanism, + initial_response=self.credential.start(), + hostname=self.host, + ) + await self.send_frame(0, sasl_init, frame_type=_SASL_FRAME_TYPE) + + _, next_frame = await self.receive_frame(verify_frame_type=1) + frame_type, fields = next_frame + if frame_type != 0x00000044: # SASLOutcome + raise NotImplementedError("Unsupported SASL challenge") + if fields[0] == SASLCode.Ok: # code + return + raise ValueError( + "SASL negotiation failed.\nOutcome: {}\nDetails: {}".format(*fields) + ) + + +class SASLTransport(AsyncTransport, SASLTransportMixinAsync): + def __init__( + self, + host, + credential, + *, + port=AMQPS_PORT, + connect_timeout=None, + ssl_opts=None, + **kwargs, + ): + self.credential = credential + ssl_opts = ssl_opts or True + super(SASLTransport, self).__init__( + host, + port=port, + connect_timeout=connect_timeout, + ssl_opts=ssl_opts, + **kwargs, + ) + + async def negotiate(self): + await self._negotiate() + + +class SASLWithWebSocket(WebSocketTransportAsync, SASLTransportMixinAsync): + def __init__( + self, + host, + credential, + *, + port=WEBSOCKET_PORT, + connect_timeout=None, + ssl_opts=None, + **kwargs, + ): + self.credential = credential + ssl_opts = ssl_opts or True + super().__init__( + host, + port=port, + connect_timeout=connect_timeout, + ssl_opts=ssl_opts, + **kwargs, + ) + + async def negotiate(self): + await self._negotiate() diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_sender_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_sender_async.py new file mode 100644 index 000000000000..29a4c052baa3 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_sender_async.py @@ -0,0 +1,203 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- +import struct +import uuid +import logging +import time +import asyncio + +from .._encode import encode_payload +from ._link_async import Link +from ..constants import SessionTransferState, LinkDeliverySettleReason, LinkState, Role, SenderSettleMode, SessionState +from ..error import AMQPLinkError, ErrorCondition, MessageException + +_LOGGER = logging.getLogger(__name__) + + +class PendingDelivery(object): + def __init__(self, **kwargs): + self.message = kwargs.get("message") + self.sent = False + self.frame = None + self.on_delivery_settled = kwargs.get("on_delivery_settled") + self.start = time.time() + self.transfer_state = None + self.timeout = kwargs.get("timeout") + self.settled = kwargs.get("settled", False) + self._network_trace_params = kwargs.get('network_trace_params') + + async def on_settled(self, reason, state): + if self.on_delivery_settled and not self.settled: + try: + await self.on_delivery_settled(reason, state) + except Exception as e: # pylint:disable=broad-except + _LOGGER.warning( + "Message 'on_send_complete' callback failed: %r", + e, + extra=self._network_trace_params + ) + self.settled = True + + +class SenderLink(Link): + def __init__(self, session, handle, target_address, **kwargs): + name = kwargs.pop("name", None) or str(uuid.uuid4()) + role = Role.Sender + if "source_address" not in kwargs: + kwargs["source_address"] = "sender-link-{}".format(name) + super(SenderLink, self).__init__(session, handle, name, role, target_address=target_address, **kwargs) + self._pending_deliveries = [] + + @classmethod + def from_incoming_frame(cls, session, handle, frame): + # TODO: Assuming we establish all links for now... + # check link_create_from_endpoint in C lib + raise NotImplementedError("Pending") + + # In theory we should not need to purge pending deliveries on attach/dettach - as a link should + # be resume-able, however this is not yet supported. + async def _incoming_attach(self, frame): + try: + await super(SenderLink, self)._incoming_attach(frame) + except AMQPLinkError: + await self._remove_pending_deliveries() + raise + self.current_link_credit = self.link_credit + await self._outgoing_flow() + await self.update_pending_deliveries() + + async def _incoming_detach(self, frame): + await super(SenderLink, self)._incoming_detach(frame) + await self._remove_pending_deliveries() + + async def _incoming_flow(self, frame): + rcv_link_credit = frame[6] # link_credit + rcv_delivery_count = frame[5] # delivery_count + if frame[4] is not None: # handle + if rcv_link_credit is None or rcv_delivery_count is None: + _LOGGER.info( + "Unable to get link-credit or delivery-count from incoming ATTACH. Detaching link.", + extra=self.network_trace_params + ) + await self._remove_pending_deliveries() + await self._set_state(LinkState.DETACHED) # TODO: Send detach now? + else: + self.current_link_credit = rcv_delivery_count + rcv_link_credit - self.delivery_count + await self.update_pending_deliveries() + + async def _outgoing_transfer(self, delivery): + output = bytearray() + encode_payload(output, delivery.message) + delivery_count = self.delivery_count + 1 + delivery.frame = { + "handle": self.handle, + "delivery_tag": struct.pack(">I", abs(delivery_count)), + "message_format": delivery.message._code, # pylint:disable=protected-access + "settled": delivery.settled, + "more": False, + "rcv_settle_mode": None, + "state": None, + "resume": None, + "aborted": None, + "batchable": None, + "payload": output, + } + await self._session._outgoing_transfer( # pylint:disable=protected-access + delivery, + self.network_trace_params if self.network_trace else None + ) + sent_and_settled = False + if delivery.transfer_state == SessionTransferState.OKAY: + self.delivery_count = delivery_count + self.current_link_credit -= 1 + delivery.sent = True + if delivery.settled: + await delivery.on_settled(LinkDeliverySettleReason.SETTLED, None) + sent_and_settled = True + # elif delivery.transfer_state == SessionTransferState.ERROR: + # TODO: Session wasn't mapped yet - re-adding to the outgoing delivery queue? + return sent_and_settled + + async def _incoming_disposition(self, frame): + if not frame[3]: # settled + return + range_end = (frame[2] or frame[1]) + 1 # first or last + settled_ids = list(range(frame[1], range_end)) + unsettled = [] + for delivery in self._pending_deliveries: + if delivery.sent and delivery.frame["delivery_id"] in settled_ids: + await delivery.on_settled(LinkDeliverySettleReason.DISPOSITION_RECEIVED, frame[4]) # state + continue + unsettled.append(delivery) + self._pending_deliveries = unsettled + + async def _remove_pending_deliveries(self): + futures = [] + for delivery in self._pending_deliveries: + futures.append(asyncio.ensure_future(delivery.on_settled(LinkDeliverySettleReason.NOT_DELIVERED, None))) + await asyncio.gather(*futures) + self._pending_deliveries = [] + + async def _on_session_state_change(self): + if self._session.state == SessionState.DISCARDING: + await self._remove_pending_deliveries() + await super()._on_session_state_change() + + async def update_pending_deliveries(self): + if self.current_link_credit <= 0: + self.current_link_credit = self.link_credit + await self._outgoing_flow() + now = time.time() + pending = [] + for delivery in self._pending_deliveries: + if delivery.timeout and (now - delivery.start) >= delivery.timeout: + delivery.on_settled(LinkDeliverySettleReason.TIMEOUT, None) + continue + if not delivery.sent: + sent_and_settled = await self._outgoing_transfer(delivery) + if sent_and_settled: + continue + pending.append(delivery) + self._pending_deliveries = pending + + async def send_transfer(self, message, *, send_async=False, **kwargs): + self._check_if_closed() + if self.state != LinkState.ATTACHED: + raise AMQPLinkError( + condition=ErrorCondition.ClientError, + description="Link is not attached." + ) + settled = self.send_settle_mode == SenderSettleMode.Settled + if self.send_settle_mode == SenderSettleMode.Mixed: + settled = kwargs.pop("settled", True) + delivery = PendingDelivery( + on_delivery_settled=kwargs.get("on_send_complete"), + timeout=kwargs.get("timeout"), + message=message, + settled=settled, + network_trace_params=self.network_trace_params + ) + if self.current_link_credit == 0 or send_async: + self._pending_deliveries.append(delivery) + else: + sent_and_settled = await self._outgoing_transfer(delivery) + if not sent_and_settled: + self._pending_deliveries.append(delivery) + return delivery + + async def cancel_transfer(self, delivery): + try: + index = self._pending_deliveries.index(delivery) + except ValueError: + raise ValueError("Found no matching pending transfer.") + delivery = self._pending_deliveries[index] + if delivery.sent: + raise MessageException( + ErrorCondition.ClientError, + message="Transfer cannot be cancelled. Message has already been sent and awaiting disposition.", + ) + await delivery.on_settled(LinkDeliverySettleReason.CANCELLED, None) + self._pending_deliveries.pop(index) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_session_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_session_async.py new file mode 100644 index 000000000000..fd1cb14218cf --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_session_async.py @@ -0,0 +1,458 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + +from __future__ import annotations +import uuid +import logging +import time +import asyncio +from typing import Optional, Union, List + +from ..constants import ConnectionState, SessionState, SessionTransferState, Role +from ._sender_async import SenderLink +from ._receiver_async import ReceiverLink +from ._management_link_async import ManagementLink +from ..performatives import ( + BeginFrame, + EndFrame, + FlowFrame, + TransferFrame, + DispositionFrame, +) +from .._encode import encode_frame +from ..error import AMQPError, ErrorCondition + +_LOGGER = logging.getLogger(__name__) + + +class Session(object): # pylint: disable=too-many-instance-attributes + """ + :param int remote_channel: The remote channel for this Session. + :param int next_outgoing_id: The transfer-id of the first transfer id the sender will send. + :param int incoming_window: The initial incoming-window of the sender. + :param int outgoing_window: The initial outgoing-window of the sender. + :param int handle_max: The maximum handle value that may be used on the Session. + :param list(str) offered_capabilities: The extension capabilities the sender supports. + :param list(str) desired_capabilities: The extension capabilities the sender may use if the receiver supports + :param dict properties: Session properties. + """ + + def __init__(self, connection, channel, **kwargs): + self.name = kwargs.pop("name", None) or str(uuid.uuid4()) + self.state = SessionState.UNMAPPED + self.handle_max = kwargs.get("handle_max", 4294967295) + self.properties = kwargs.pop("properties", None) + self.channel = channel + self.remote_channel = None + self.next_outgoing_id = kwargs.pop("next_outgoing_id", 0) + self.next_incoming_id = None + self.incoming_window = kwargs.pop("incoming_window", 1) + self.outgoing_window = kwargs.pop("outgoing_window", 1) + self.target_incoming_window = self.incoming_window + self.remote_incoming_window = 0 + self.remote_outgoing_window = 0 + self.offered_capabilities = None + self.desired_capabilities = kwargs.pop("desired_capabilities", None) + + self.allow_pipelined_open = kwargs.pop("allow_pipelined_open", True) + self.idle_wait_time = kwargs.get("idle_wait_time", 0.1) + self.network_trace = kwargs["network_trace"] + self.network_trace_params = kwargs["network_trace_params"] + self.network_trace_params["amqpSession"] = self.name + + self.links = {} + self._connection = connection + self._output_handles = {} + self._input_handles = {} + + async def __aenter__(self): + await self.begin() + return self + + async def __aexit__(self, *args): + await self.end() + + @classmethod + def from_incoming_frame(cls, connection, channel): + # check session_create_from_endpoint in C lib + new_session = cls(connection, channel) + return new_session + + async def _set_state(self, new_state): + # type: (SessionState) -> None + """Update the session state.""" + if new_state is None: + return + previous_state = self.state + self.state = new_state + _LOGGER.info( + "Session state changed: %r -> %r", + previous_state, + new_state, + extra=self.network_trace_params, + ) + for link in self.links.values(): + await link._on_session_state_change() # pylint: disable=protected-access + + async def _on_connection_state_change(self): + if self._connection.state in [ConnectionState.CLOSE_RCVD, ConnectionState.END]: + if self.state not in [SessionState.DISCARDING, SessionState.UNMAPPED]: + await self._set_state(SessionState.DISCARDING) + + def _get_next_output_handle(self): + # type: () -> int + """Get the next available outgoing handle number within the max handle limit. + + :raises ValueError: If maximum handle has been reached. + :returns: The next available outgoing handle number. + :rtype: int + """ + if len(self._output_handles) >= self.handle_max: + raise ValueError("Maximum number of handles ({}) has been reached.".format(self.handle_max)) + next_handle = next(i for i in range(1, self.handle_max) if i not in self._output_handles) + return next_handle + + async def _outgoing_begin(self): + begin_frame = BeginFrame( + remote_channel=self.remote_channel if self.state == SessionState.BEGIN_RCVD else None, + next_outgoing_id=self.next_outgoing_id, + outgoing_window=self.outgoing_window, + incoming_window=self.incoming_window, + handle_max=self.handle_max, + offered_capabilities=self.offered_capabilities if self.state == SessionState.BEGIN_RCVD else None, + desired_capabilities=self.desired_capabilities if self.state == SessionState.UNMAPPED else None, + properties=self.properties, + ) + if self.network_trace: + _LOGGER.debug("-> %r", begin_frame, extra=self.network_trace_params) + await self._connection._process_outgoing_frame(self.channel, begin_frame) # pylint: disable=protected-access + + async def _incoming_begin(self, frame): + if self.network_trace: + _LOGGER.debug("<- %r", BeginFrame(*frame), extra=self.network_trace_params) + self.handle_max = frame[4] # handle_max + self.next_incoming_id = frame[1] # next_outgoing_id + self.remote_incoming_window = frame[2] # incoming_window + self.remote_outgoing_window = frame[3] # outgoing_window + if self.state == SessionState.BEGIN_SENT: + self.remote_channel = frame[0] # remote_channel + await self._set_state(SessionState.MAPPED) + elif self.state == SessionState.UNMAPPED: + await self._set_state(SessionState.BEGIN_RCVD) + await self._outgoing_begin() + await self._set_state(SessionState.MAPPED) + + async def _outgoing_end(self, error=None): + end_frame = EndFrame(error=error) + if self.network_trace: + _LOGGER.debug("-> %r", end_frame, extra=self.network_trace_params) + await self._connection._process_outgoing_frame(self.channel, end_frame) # pylint: disable=protected-access + + async def _incoming_end(self, frame): + if self.network_trace: + _LOGGER.debug("<- %r", EndFrame(*frame), extra=self.network_trace_params) + if self.state not in [ + SessionState.END_RCVD, + SessionState.END_SENT, + SessionState.DISCARDING, + ]: + await self._set_state(SessionState.END_RCVD) + for _, link in self.links.items(): + await link.detach() + # TODO: handling error + await self._outgoing_end() + await self._set_state(SessionState.UNMAPPED) + + async def _outgoing_attach(self, frame): + await self._connection._process_outgoing_frame(self.channel, frame) # pylint: disable=protected-access + + async def _incoming_attach(self, frame): + try: + self._input_handles[frame[1]] = self.links[frame[0].decode("utf-8")] # name and handle + await self._input_handles[frame[1]]._incoming_attach(frame) # pylint: disable=protected-access + except KeyError: + try: + outgoing_handle = self._get_next_output_handle() + except ValueError: + _LOGGER.error( + "Unable to attach new link - cannot allocate more handles.", + extra=self.network_trace_params + ) + # detach the link that would have been set. + await self.links[frame[0].decode("utf-8")].detach( + error=AMQPError( + condition=ErrorCondition.LinkDetachForced, + description=f"Cannot allocate more handles, the max number of handles is {self.handle_max}. Detaching link", # pylint: disable=line-too-long + info=None, + ) + ) + return + if frame[2] == Role.Sender: + new_link = ReceiverLink.from_incoming_frame(self, outgoing_handle, frame) + else: + new_link = SenderLink.from_incoming_frame(self, outgoing_handle, frame) + await new_link._incoming_attach(frame) # pylint: disable=protected-access + self.links[frame[0]] = new_link + self._output_handles[outgoing_handle] = new_link + self._input_handles[frame[1]] = new_link + except ValueError as e: + # Reject Link + _LOGGER.error( + "Unable to attach new link: %r", + e, + extra=self.network_trace_params + ) + await self._input_handles[frame[1]].detach() + + async def _outgoing_flow(self, frame=None): + link_flow = frame or {} + link_flow.update( + { + "next_incoming_id": self.next_incoming_id, + "incoming_window": self.incoming_window, + "next_outgoing_id": self.next_outgoing_id, + "outgoing_window": self.outgoing_window, + } + ) + flow_frame = FlowFrame(**link_flow) + if self.network_trace: + _LOGGER.debug("-> %r", flow_frame, extra=self.network_trace_params) + await self._connection._process_outgoing_frame(self.channel, flow_frame) # pylint: disable=protected-access + + async def _incoming_flow(self, frame): + if self.network_trace: + _LOGGER.debug("<- %r", FlowFrame(*frame), extra=self.network_trace_params) + self.next_incoming_id = frame[2] # next_outgoing_id + remote_incoming_id = frame[0] or self.next_outgoing_id # next_incoming_id TODO "initial-outgoing-id" + self.remote_incoming_window = remote_incoming_id + frame[1] - self.next_outgoing_id # incoming_window + self.remote_outgoing_window = frame[3] # outgoing_window + if frame[4] is not None: # handle + await self._input_handles[frame[4]]._incoming_flow(frame) # pylint: disable=protected-access + else: + for link in self._output_handles.values(): + if self.remote_incoming_window > 0 and not link._is_closed: # pylint: disable=protected-access + await link._incoming_flow(frame) # pylint: disable=protected-access + + async def _outgoing_transfer(self, delivery, network_trace_params): + if self.state != SessionState.MAPPED: + delivery.transfer_state = SessionTransferState.ERROR + if self.remote_incoming_window <= 0: + delivery.transfer_state = SessionTransferState.BUSY + else: + payload = delivery.frame["payload"] + payload_size = len(payload) + + delivery.frame["delivery_id"] = self.next_outgoing_id + # calculate the transfer frame encoding size excluding the payload + delivery.frame["payload"] = b"" + # TODO: encoding a frame would be expensive, we might want to improve depending on the perf test results + encoded_frame = encode_frame(TransferFrame(**delivery.frame))[1] + transfer_overhead_size = len(encoded_frame) + + # available size for payload per frame is calculated as following: + # remote max frame size - transfer overhead (calculated) - header (8 bytes) + available_frame_size = ( + self._connection._remote_max_frame_size - transfer_overhead_size - 8 # pylint: disable=protected-access + ) + + start_idx = 0 + remaining_payload_cnt = payload_size + # encode n-1 frames if payload_size > available_frame_size + while remaining_payload_cnt > available_frame_size: + tmp_delivery_frame = { + "handle": delivery.frame["handle"], + "delivery_tag": delivery.frame["delivery_tag"], + "message_format": delivery.frame["message_format"], + "settled": delivery.frame["settled"], + "more": True, + "rcv_settle_mode": delivery.frame["rcv_settle_mode"], + "state": delivery.frame["state"], + "resume": delivery.frame["resume"], + "aborted": delivery.frame["aborted"], + "batchable": delivery.frame["batchable"], + "delivery_id": self.next_outgoing_id, + } + if network_trace_params: + # We determine the logging for the outgoing Transfer frames based on the source + # Link configuration rather than the Session, because it's only at the Session + # level that we can determine how many outgoing frames are needed and their + # delivery IDs. + # TODO: Obscuring the payload for now to investigate the potential for leaks. + _LOGGER.debug( + "-> %r", TransferFrame(payload=b"***", **tmp_delivery_frame), + extra=network_trace_params + ) + await self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, + TransferFrame( + payload=payload[start_idx : start_idx + available_frame_size], + **tmp_delivery_frame + ) + ) + start_idx += available_frame_size + remaining_payload_cnt -= available_frame_size + + # encode the last frame + tmp_delivery_frame = { + "handle": delivery.frame["handle"], + "delivery_tag": delivery.frame["delivery_tag"], + "message_format": delivery.frame["message_format"], + "settled": delivery.frame["settled"], + "more": False, + "rcv_settle_mode": delivery.frame["rcv_settle_mode"], + "state": delivery.frame["state"], + "resume": delivery.frame["resume"], + "aborted": delivery.frame["aborted"], + "batchable": delivery.frame["batchable"], + "delivery_id": self.next_outgoing_id, + } + if network_trace_params: + # We determine the logging for the outgoing Transfer frames based on the source + # Link configuration rather than the Session, because it's only at the Session + # level that we can determine how many outgoing frames are needed and their + # delivery IDs. + # TODO: Obscuring the payload for now to investigate the potential for leaks. + _LOGGER.debug( + "-> %r", TransferFrame(payload=b"***", **tmp_delivery_frame), + extra=network_trace_params + ) + await self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, + TransferFrame(payload=payload[start_idx:], **tmp_delivery_frame) + ) + self.next_outgoing_id += 1 + self.remote_incoming_window -= 1 + self.outgoing_window -= 1 + # TODO: We should probably handle an error at the connection and update state accordingly + delivery.transfer_state = SessionTransferState.OKAY + + async def _incoming_transfer(self, frame): + self.next_incoming_id += 1 + self.remote_outgoing_window -= 1 + self.incoming_window -= 1 + try: + await self._input_handles[frame[0]]._incoming_transfer(frame) # pylint: disable=protected-access + except KeyError: + _LOGGER.error( + "Received Transfer frame on unattached link. Ending session.", + extra=self.network_trace_params + ) + await self._set_state(SessionState.DISCARDING) + await self.end( + error=AMQPError( + condition=ErrorCondition.SessionUnattachedHandle, + description="""Invalid handle reference in received frame: """ + """Handle is not currently associated with an attached link""", + ) + ) + if self.incoming_window == 0: + self.incoming_window = self.target_incoming_window + await self._outgoing_flow() + + async def _outgoing_disposition(self, frame): + await self._connection._process_outgoing_frame(self.channel, frame) # pylint: disable=protected-access + + async def _incoming_disposition(self, frame): + if self.network_trace: + _LOGGER.debug("<- %r", DispositionFrame(*frame), extra=self.network_trace_params) + for link in self._input_handles.values(): + await link._incoming_disposition(frame) # pylint: disable=protected-access + + async def _outgoing_detach(self, frame): + await self._connection._process_outgoing_frame(self.channel, frame) # pylint: disable=protected-access + + async def _incoming_detach(self, frame): + try: + link = self._input_handles[frame[0]] # handle + await link._incoming_detach(frame) # pylint: disable=protected-access + # if link._is_closed: TODO + # self.links.pop(link.name, None) + # self._input_handles.pop(link.remote_handle, None) + # self._output_handles.pop(link.handle, None) + except KeyError: + await self._set_state(SessionState.DISCARDING) + await self._connection.close( + error=AMQPError( + condition=ErrorCondition.SessionUnattachedHandle, + description="""Invalid handle reference in received frame: """ + """Handle is not currently associated with an attached link""", + ) + ) + + async def _wait_for_response(self, wait, end_state): + # type: (Union[bool, float], SessionState) -> None + if wait is True: + await self._connection.listen(wait=False) + while self.state != end_state: + await asyncio.sleep(self.idle_wait_time) + await self._connection.listen(wait=False) + elif wait: + await self._connection.listen(wait=False) + timeout = time.time() + wait + while self.state != end_state: + if time.time() >= timeout: + break + await asyncio.sleep(self.idle_wait_time) + await self._connection.listen(wait=False) + + async def begin(self, wait=False): + await self._outgoing_begin() + await self._set_state(SessionState.BEGIN_SENT) + if wait: + await self._wait_for_response(wait, SessionState.BEGIN_SENT) + elif not self.allow_pipelined_open: + raise ValueError("Connection has been configured to not allow piplined-open. Please set 'wait' parameter.") + + async def end(self, error=None, wait=False): + # type: (Optional[AMQPError], bool) -> None + try: + if self.state not in [SessionState.UNMAPPED, SessionState.DISCARDING]: + await self._outgoing_end(error=error) + for _, link in self.links.items(): + await link.detach() + new_state = SessionState.DISCARDING if error else SessionState.END_SENT + await self._set_state(new_state) + await self._wait_for_response(wait, SessionState.UNMAPPED) + except Exception as exc: # pylint: disable=broad-except + _LOGGER.info("An error occurred when ending the session: %r", exc, extra=self.network_trace_params) + await self._set_state(SessionState.UNMAPPED) + + def create_receiver_link(self, source_address, **kwargs): + assigned_handle = self._get_next_output_handle() + link = ReceiverLink( + self, + handle=assigned_handle, + source_address=source_address, + network_trace=kwargs.pop("network_trace", self.network_trace), + network_trace_params=dict(self.network_trace_params), + **kwargs, + ) + self.links[link.name] = link + self._output_handles[assigned_handle] = link + return link + + def create_sender_link(self, target_address, **kwargs): + assigned_handle = self._get_next_output_handle() + link = SenderLink( + self, + handle=assigned_handle, + target_address=target_address, + network_trace=kwargs.pop("network_trace", self.network_trace), + network_trace_params=dict(self.network_trace_params), + **kwargs, + ) + self._output_handles[assigned_handle] = link + self.links[link.name] = link + return link + + def create_request_response_link_pair(self, endpoint, **kwargs): + return ManagementLink( + self, + endpoint, + network_trace=kwargs.pop("network_trace", self.network_trace), + network_trace_params=dict(self.network_trace_params), + **kwargs, + ) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_transport_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_transport_async.py new file mode 100644 index 000000000000..a5caf7418fc1 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_transport_async.py @@ -0,0 +1,589 @@ +# ------------------------------------------------------------------------- # pylint: disable=file-needs-copyright-header +# This is a fork of the transport.py which was originally written by Barry Pederson and +# maintained by the Celery project: https://github.com/celery/py-amqp. +# +# Copyright (C) 2009 Barry Pederson +# +# The license text can also be found here: +# http://www.opensource.org/licenses/BSD-3-Clause +# +# License +# ======= +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, +# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS +# BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF +# THE POSSIBILITY OF SUCH DAMAGE. +# ------------------------------------------------------------------------- + +import asyncio +import errno +import socket +import ssl +import struct +from ssl import SSLError +from io import BytesIO +import logging + + + +import certifi + +from .._platform import KNOWN_TCP_OPTS, SOL_TCP +from .._encode import encode_frame +from .._decode import decode_frame, decode_empty_frame +from ..constants import DEFAULT_WEBSOCKET_HEARTBEAT_SECONDS, TLS_HEADER_FRAME, WEBSOCKET_PORT, AMQP_WS_SUBPROTOCOL +from .._transport import ( + AMQP_FRAME, + get_errno, + to_host_port, + DEFAULT_SOCKET_SETTINGS, + SIGNED_INT_MAX, + _UNAVAIL, + set_cloexec, + AMQP_PORT, + TIMEOUT_INTERVAL, +) +from ..error import AuthenticationException, ErrorCondition + + +_LOGGER = logging.getLogger(__name__) + + +class AsyncTransportMixin: + async def receive_frame(self, timeout=None, **kwargs): + try: + header, channel, payload = await asyncio.wait_for( + self.read(**kwargs), timeout=timeout + ) + if not payload: + decoded = decode_empty_frame(header) + else: + decoded = decode_frame(payload) + return channel, decoded + except ( + TimeoutError, + socket.timeout, + asyncio.IncompleteReadError, + asyncio.TimeoutError, + ): + return None, None + + async def read(self, verify_frame_type=0): + async with self.socket_lock: + read_frame_buffer = BytesIO() + try: + frame_header = memoryview(bytearray(8)) + read_frame_buffer.write( + await self._read(8, buffer=frame_header, initial=True) + ) + + channel = struct.unpack(">H", frame_header[6:])[0] + size = frame_header[0:4] + if size == AMQP_FRAME: # Empty frame or AMQP header negotiation + return frame_header, channel, None + size = struct.unpack(">I", size)[0] + offset = frame_header[4] + frame_type = frame_header[5] + if verify_frame_type is not None and frame_type != verify_frame_type: + _LOGGER.debug( + "Received invalid frame type: %r, expected: %r", + frame_type, + verify_frame_type, + extra=self.network_trace_params + ) + raise ValueError( + f"Received invalid frame type: {frame_type}, expected: {verify_frame_type}" + ) + # >I is an unsigned int, but the argument to sock.recv is signed, + # so we know the size can be at most 2 * SIGNED_INT_MAX + payload_size = size - len(frame_header) + payload = memoryview(bytearray(payload_size)) + if size > SIGNED_INT_MAX: + read_frame_buffer.write( + await self._read(SIGNED_INT_MAX, buffer=payload) + ) + read_frame_buffer.write( + await self._read( + size - SIGNED_INT_MAX, buffer=payload[SIGNED_INT_MAX:] + ) + ) + else: + read_frame_buffer.write( + await self._read(payload_size, buffer=payload) + ) + except ( + asyncio.CancelledError, + asyncio.TimeoutError, + TimeoutError, + socket.timeout, + asyncio.IncompleteReadError + ): + read_frame_buffer.write(self._read_buffer.getvalue()) + self._read_buffer = read_frame_buffer + self._read_buffer.seek(0) + raise + except (OSError, IOError, SSLError, socket.error) as exc: + # Don't disconnect for ssl read time outs + # http://bugs.python.org/issue10272 + if isinstance(exc, SSLError) and "timed out" in str(exc): + raise socket.timeout() + if get_errno(exc) not in _UNAVAIL: + self.connected = False + _LOGGER.debug("Transport read failed: %r", exc, extra=self.network_trace_params) + raise + offset -= 2 + return frame_header, channel, payload[offset:] + + async def write(self, s): + async with self.socket_lock: + try: + await self._write(s) + except socket.timeout: + raise + except (OSError, IOError, socket.error) as exc: + _LOGGER.debug("Transport write failed: %r", exc, extra=self.network_trace_params) + if get_errno(exc) not in _UNAVAIL: + self.connected = False + raise + + async def send_frame(self, channel, frame, **kwargs): + header, performative = encode_frame(frame, **kwargs) + if performative is None: + data = header + else: + encoded_channel = struct.pack(">H", channel) + data = header + encoded_channel + performative + + await self.write(data) + + def _build_ssl_opts(self, sslopts): + if sslopts in [True, False, None, {}]: + return sslopts + try: + if "context" in sslopts: + return self._build_ssl_context(**sslopts.pop("context")) + ssl_version = sslopts.get("ssl_version") + if ssl_version is None: + ssl_version = ssl.PROTOCOL_TLS + + # Set SNI headers if supported + server_hostname = sslopts.get("server_hostname") + if ( + (server_hostname is not None) + and (hasattr(ssl, "HAS_SNI") and ssl.HAS_SNI) + and (hasattr(ssl, "SSLContext")) + ): + context = ssl.SSLContext(ssl_version) + cert_reqs = sslopts.get("cert_reqs", ssl.CERT_REQUIRED) + certfile = sslopts.get("certfile") + keyfile = sslopts.get("keyfile") + context.verify_mode = cert_reqs + if cert_reqs != ssl.CERT_NONE: + context.check_hostname = True + if (certfile is not None) and (keyfile is not None): + context.load_cert_chain(certfile, keyfile) + return context + ca_certs = sslopts.get("ca_certs") + if ca_certs: + context = ssl.SSLContext(ssl_version) + context.load_verify_locations(ca_certs) + return context + return True + except TypeError: + raise TypeError( + "SSL configuration must be a dictionary, or the value True." + ) + + def _build_ssl_context( + self, check_hostname=None, **ctx_options + ): # pylint: disable=no-self-use + ctx = ssl.create_default_context(**ctx_options) + ctx.verify_mode = ssl.CERT_REQUIRED + ctx.load_verify_locations(cafile=certifi.where()) + ctx.check_hostname = check_hostname + return ctx + + +class AsyncTransport( + AsyncTransportMixin +): # pylint: disable=too-many-instance-attributes + """Common superclass for TCP and SSL transports.""" + + def __init__( + self, + host, + *, + port=AMQP_PORT, + connect_timeout=None, + ssl_opts=False, + socket_settings=None, + raise_on_initial_eintr=True, + **kwargs, # pylint: disable=unused-argument + ): + self.connected = False + self.sock = None + self.reader = None + self.writer = None + self.raise_on_initial_eintr = raise_on_initial_eintr + self._read_buffer = BytesIO() + self.host, self.port = to_host_port(host, port) + + self.connect_timeout = connect_timeout + self.socket_settings = socket_settings + self.socket_lock = asyncio.Lock() + self.sslopts = ssl_opts + self.network_trace_params = kwargs.get('network_trace_params') + + async def connect(self): + try: + # are we already connected? + if self.connected: + return + await self._connect(self.host, self.port, self.connect_timeout) + self._init_socket(self.socket_settings) + self.reader, self.writer = await asyncio.open_connection( + sock=self.sock, + ssl=self.sslopts, + server_hostname=self.host if self.sslopts else None, + ) + # we've sent the banner; signal connect + # EINTR, EAGAIN, EWOULDBLOCK would signal that the banner + # has _not_ been sent + self.connected = True + except (OSError, IOError, SSLError) as e: + _LOGGER.info("Transport connect failed: %r", e, extra=self.network_trace_params) + # if not fully connected, close socket, and reraise error + if self.sock and not self.connected: + self.sock.close() + self.sock = None + raise + + async def _connect(self, host, port, timeout): + # Below we are trying to avoid additional DNS requests for AAAA if A + # succeeds. This helps a lot in case when a hostname has an IPv4 entry + # in /etc/hosts but not IPv6. Without the (arguably somewhat twisted) + # logic below, getaddrinfo would attempt to resolve the hostname for + # both IP versions, which would make the resolver talk to configured + # DNS servers. If those servers are for some reason not available + # during resolution attempt (either because of system misconfiguration, + # or network connectivity problem), resolution process locks the + # _connect call for extended time. + e = None + addr_types = (socket.AF_INET, socket.AF_INET6) + addr_types_num = len(addr_types) + for n, family in enumerate(addr_types): + # first, resolve the address for a single address family + try: + entries = await asyncio.get_event_loop().getaddrinfo( + host, port, family=family, type=socket.SOCK_STREAM, proto=SOL_TCP + ) + entries_num = len(entries) + except socket.gaierror: + # we may have depleted all our options + if n + 1 >= addr_types_num: + # if getaddrinfo succeeded before for another address + # family, reraise the previous socket.error since it's more + # relevant to users + raise e if e is not None else socket.error("failed to resolve broker hostname") + continue # pragma: no cover + # now that we have address(es) for the hostname, connect to broker + for i, res in enumerate(entries): + af, socktype, proto, _, sa = res + try: + self.sock = socket.socket(af, socktype, proto) + try: + set_cloexec(self.sock, True) + except NotImplementedError: + pass + self.sock.settimeout(timeout) + await asyncio.get_event_loop().sock_connect(self.sock, sa) + except socket.error as ex: + e = ex + if self.sock is not None: + self.sock.close() + self.sock = None + # we may have depleted all our options + if i + 1 >= entries_num and n + 1 >= addr_types_num: + raise + else: + # hurray, we established connection + return + + def _init_socket(self, socket_settings): + self.sock.settimeout(None) # set socket back to blocking mode + self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) + self._set_socket_options(socket_settings) + try: + # Building ssl opts here instead of constructor, so that invalid cert error is raised + # when client is connecting, rather then during creation. For uamqp exception parity. + self.sslopts = self._build_ssl_opts(self.sslopts) + except FileNotFoundError as exc: + # FileNotFoundError does not have missing filename info, so adding it below. + # Assuming that this must be ca_certs, since this is the only file path that + # users can pass in (`connection_verify` in the EH/SB clients) through sslopts above. + # For uamqp exception parity. Remove later when resolving issue #27128. + exc.filename = self.sslopts + raise exc + self.sock.settimeout(1) # set socket back to non-blocking mode + + def _get_tcp_socket_defaults(self, sock): # pylint: disable=no-self-use + tcp_opts = {} + for opt in KNOWN_TCP_OPTS: + enum = None + if opt == "TCP_USER_TIMEOUT": + try: + from socket import TCP_USER_TIMEOUT as enum + except ImportError: + # should be in Python 3.6+ on Linux. + enum = 18 + elif hasattr(socket, opt): + enum = getattr(socket, opt) + + if enum: + if opt in DEFAULT_SOCKET_SETTINGS: + tcp_opts[enum] = DEFAULT_SOCKET_SETTINGS[opt] + elif hasattr(socket, opt): + tcp_opts[enum] = sock.getsockopt(SOL_TCP, getattr(socket, opt)) + return tcp_opts + + def _set_socket_options(self, socket_settings): + tcp_opts = self._get_tcp_socket_defaults(self.sock) + if socket_settings: + tcp_opts.update(socket_settings) + for opt, val in tcp_opts.items(): + self.sock.setsockopt(SOL_TCP, opt, val) + + async def _read( + self, + toread, + initial=False, + buffer=None, + _errnos=(errno.ENOENT, errno.EAGAIN, errno.EINTR), + ): + # According to SSL_read(3), it can at most return 16kb of data. + # Thus, we use an internal read buffer like TCPTransport._read + # to get the exact number of bytes wanted. + length = 0 + view = buffer or memoryview(bytearray(toread)) + nbytes = self._read_buffer.readinto(view) + toread -= nbytes + length += nbytes + try: + while toread: + try: + view[nbytes : nbytes + toread] = await self.reader.readexactly( + toread + ) + nbytes = toread + except AttributeError: + # This means that close() was called concurrently + # self.reader has been set to None. + raise IOError("Connection has already been closed") + except asyncio.IncompleteReadError as exc: + pbytes = len(exc.partial) + view[nbytes : nbytes + pbytes] = exc.partial + nbytes = pbytes + except socket.error as exc: + # ssl.sock.read may cause a SSLerror without errno + # http://bugs.python.org/issue10272 + if isinstance(exc, SSLError) and "timed out" in str(exc): + raise socket.timeout() + # ssl.sock.read may cause ENOENT if the + # operation couldn't be performed (Issue celery#1414). + if exc.errno in _errnos: + if initial and self.raise_on_initial_eintr: + raise socket.timeout() + continue + raise + if not nbytes: + raise IOError("Server unexpectedly closed connection") + + length += nbytes + toread -= nbytes + except: # noqa + self._read_buffer = BytesIO(view[:length]) + raise + return view + + async def _write(self, s): + """Write a string out to the SSL socket fully.""" + try: + self.writer.write(s) + await self.writer.drain() + except AttributeError: + raise IOError("Connection has already been closed") + + async def close(self): + async with self.socket_lock: + try: + if self.writer is not None: + # Closing the writer closes the underlying socket. + self.writer.close() + if self.sslopts: + # see issue: https://github.com/encode/httpx/issues/914 + await asyncio.sleep(0) + self.writer.transport.abort() + await self.writer.wait_closed() + except Exception as e: # pylint: disable=broad-except + # Sometimes SSL raises APPLICATION_DATA_AFTER_CLOSE_NOTIFY here on close. + _LOGGER.debug("Error shutting down socket: %r", e, extra=self.network_trace_params) + self.writer, self.reader = None, None + self.sock = None + self.connected = False + + async def negotiate(self): + if not self.sslopts: + return + await self.write(TLS_HEADER_FRAME) + _, returned_header = await self.receive_frame(verify_frame_type=None) + if returned_header[1] == TLS_HEADER_FRAME: + raise ValueError( + f"""Mismatching TLS header protocol. Expected: {TLS_HEADER_FRAME!r},""" + """received: {returned_header[1]!r}""" + ) + + +class WebSocketTransportAsync( + AsyncTransportMixin +): # pylint: disable=too-many-instance-attributes + def __init__( + self, + host, + *, + port=WEBSOCKET_PORT, + connect_timeout=None, + ssl_opts=None, + **kwargs + ): + self._read_buffer = BytesIO() + self.socket_lock = asyncio.Lock() + self.sslopts = ssl_opts if isinstance(ssl_opts, dict) else None + self._connect_timeout = connect_timeout or TIMEOUT_INTERVAL + self._custom_endpoint = kwargs.get("custom_endpoint") + self.host, self.port = to_host_port(host, port) + self.ws = None + self.session = None + self._http_proxy = kwargs.get("http_proxy", None) + self.connected = False + self.network_trace_params = kwargs.get('network_trace_params') + + async def connect(self): + self.sslopts = self._build_ssl_opts(self.sslopts) + username, password = None, None + http_proxy_host, http_proxy_port = None, None + http_proxy_auth = None + + if self._http_proxy: + http_proxy_host = self._http_proxy["proxy_hostname"] + http_proxy_port = self._http_proxy["proxy_port"] + if http_proxy_host and http_proxy_port: + http_proxy_host = f"{http_proxy_host}:{http_proxy_port}" + username = self._http_proxy.get("username", None) + password = self._http_proxy.get("password", None) + + try: + from aiohttp import ClientSession, ClientConnectorError + from urllib.parse import urlsplit + + if username or password: + from aiohttp import BasicAuth + + http_proxy_auth = BasicAuth(login=username, password=password) + + self.session = ClientSession() + if self._custom_endpoint: + url = f"wss://{self._custom_endpoint}" + else: + url = f"wss://{self.host}" + parsed_url = urlsplit(url) + url = f"{parsed_url.scheme}://{parsed_url.netloc}:{self.port}{parsed_url.path}" + + try: + # Enabling heartbeat that sends a ping message every n seconds and waits for pong response. + # if pong response is not received then close connection. This raises an error when trying + # to communicate with the websocket which is no longer active. + # We are waiting a bug fix in aiohttp for these 2 bugs where aiohttp ws might hang on network disconnect + # and the heartbeat mechanism helps mitigate these two. + # https://github.com/aio-libs/aiohttp/pull/5860 + # https://github.com/aio-libs/aiohttp/issues/2309 + + self.ws = await self.session.ws_connect( + url=url, + timeout=self._connect_timeout, + protocols=[AMQP_WS_SUBPROTOCOL], + autoclose=False, + proxy=http_proxy_host, + proxy_auth=http_proxy_auth, + ssl=self.sslopts, + heartbeat=DEFAULT_WEBSOCKET_HEARTBEAT_SECONDS, + ) + except ClientConnectorError as exc: + _LOGGER.info("Websocket connect failed: %r", exc, extra=self.network_trace_params) + if self._custom_endpoint: + raise AuthenticationException( + ErrorCondition.ClientError, + description="Failed to authenticate the connection due to exception: " + str(exc), + error=exc, + ) + raise ConnectionError("Failed to establish websocket connection: " + str(exc)) + self.connected = True + except ImportError: + raise ValueError( + "Please install aiohttp library to use websocket transport." + ) + + async def _read(self, toread, buffer=None, **kwargs): # pylint: disable=unused-argument + """Read exactly n bytes from the peer.""" + length = 0 + view = buffer or memoryview(bytearray(toread)) + nbytes = self._read_buffer.readinto(view) + length += nbytes + toread -= nbytes + try: + while toread: + data = await self.ws.receive_bytes() + read_length = len(data) + if read_length <= toread: + view[length : length + read_length] = data + toread -= read_length + length += read_length + else: + view[length : length + toread] = data[0:toread] + self._read_buffer = BytesIO(data[toread:]) + toread = 0 + return view + except: + self._read_buffer = BytesIO(view[:length]) + raise + + async def close(self): + """Do any preliminary work in shutting down the connection.""" + async with self.socket_lock: + await self.ws.close() + await self.session.close() + self.connected = False + + async def _write(self, s): + """Completely write a string (byte array) to the peer. + ABNF, OPCODE_BINARY = 0x2 + See http://tools.ietf.org/html/rfc5234 + http://tools.ietf.org/html/rfc6455#section-5.2 + """ + await self.ws.send_bytes(s) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/authentication.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/authentication.py new file mode 100644 index 000000000000..43d7803c87d6 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/authentication.py @@ -0,0 +1,175 @@ +#------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +#------------------------------------------------------------------------- + +import time +from collections import namedtuple +from functools import partial + +from .sasl import SASLAnonymousCredential, SASLPlainCredential +from .utils import generate_sas_token + +from .constants import ( + AUTH_DEFAULT_EXPIRATION_SECONDS, + TOKEN_TYPE_JWT, + TOKEN_TYPE_SASTOKEN, + AUTH_TYPE_CBS, + AUTH_TYPE_SASL_PLAIN +) + +AccessToken = namedtuple("AccessToken", ["token", "expires_on"]) + + +def _generate_sas_access_token(auth_uri, sas_name, sas_key, expiry_in=AUTH_DEFAULT_EXPIRATION_SECONDS): + expires_on = int(time.time() + expiry_in) + token = generate_sas_token(auth_uri, sas_name, sas_key, expires_on) + return AccessToken( + token, + expires_on + ) + + +class SASLPlainAuth(object): + # TODO: + # 1. naming decision, suffix with Auth vs Credential + auth_type = AUTH_TYPE_SASL_PLAIN + + def __init__(self, authcid, passwd, authzid=None): + self.sasl = SASLPlainCredential(authcid, passwd, authzid) + + +class _CBSAuth(object): + # TODO: + # 1. naming decision, suffix with Auth vs Credential + auth_type = AUTH_TYPE_CBS + + def __init__( + self, + uri, + audience, + token_type, + get_token, + **kwargs + ): + """ + CBS authentication using JWT tokens. + + :param uri: The AMQP endpoint URI. This must be provided as + a decoded string. + :type uri: str + :param audience: The token audience field. For SAS tokens + this is usually the URI. + :type audience: str + :param get_token: The callback function used for getting and refreshing + tokens. It should return a valid jwt token each time it is called. + :type get_token: callable object + :param token_type: The type field of the token request. + Default value is `"jwt"`. + :type token_type: str + + """ + self.sasl = SASLAnonymousCredential() + self.uri = uri + self.audience = audience + self.token_type = token_type + self.get_token = get_token + self.expires_in = kwargs.pop("expires_in", AUTH_DEFAULT_EXPIRATION_SECONDS) + self.expires_on = kwargs.pop("expires_on", None) + + @staticmethod + def _set_expiry(expires_in, expires_on): + if not expires_on and not expires_in: + raise ValueError("Must specify either 'expires_on' or 'expires_in'.") + if not expires_on: + expires_on = time.time() + expires_in + else: + expires_in = expires_on - time.time() + if expires_in < 1: + raise ValueError("Token has already expired.") + return expires_in, expires_on + + +class JWTTokenAuth(_CBSAuth): + # TODO: + # 1. naming decision, suffix with Auth vs Credential + def __init__( + self, + uri, + audience, + get_token, + **kwargs + ): + """ + CBS authentication using JWT tokens. + + :param uri: The AMQP endpoint URI. This must be provided as + a decoded string. + :type uri: str + :param audience: The token audience field. For SAS tokens + this is usually the URI. + :type audience: str + :param get_token: The callback function used for getting and refreshing + tokens. It should return a valid jwt token each time it is called. + :type get_token: callable object + :param token_type: The type field of the token request. + Default value is `"jwt"`. + :type token_type: str + + """ + super(JWTTokenAuth, self).__init__(uri, audience, kwargs.pop("kwargs", TOKEN_TYPE_JWT), get_token) + self.get_token = get_token + + +class SASTokenAuth(_CBSAuth): + # TODO: + # 1. naming decision, suffix with Auth vs Credential + def __init__( + self, + uri, + audience, + username, + password, + **kwargs + ): + """ + CBS authentication using SAS tokens. + + :param uri: The AMQP endpoint URI. This must be provided as + a decoded string. + :type uri: str + :param audience: The token audience field. For SAS tokens + this is usually the URI. + :type audience: str + :param username: The SAS token username, also referred to as the key + name or policy name. This can optionally be encoded into the URI. + :type username: str + :param password: The SAS token password, also referred to as the key. + This can optionally be encoded into the URI. + :type password: str + :param expires_in: The total remaining seconds until the token + expires. + :type expires_in: int + :param expires_on: The timestamp at which the SAS token will expire + formatted as seconds since epoch. + :type expires_on: float + :param token_type: The type field of the token request. + Default value is `"servicebus.windows.net:sastoken"`. + :type token_type: str + + """ + self.username = username + self.password = password + expires_in = kwargs.pop("expires_in", AUTH_DEFAULT_EXPIRATION_SECONDS) + expires_on = kwargs.pop("expires_on", None) + expires_in, expires_on = self._set_expiry(expires_in, expires_on) + self.get_token = partial(_generate_sas_access_token, uri, username, password, expires_in) + super(SASTokenAuth, self).__init__( + uri, + audience, + kwargs.pop("token_type", TOKEN_TYPE_SASTOKEN), + self.get_token, + expires_in=expires_in, + expires_on=expires_on + ) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/cbs.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/cbs.py new file mode 100644 index 000000000000..f2eb796b587c --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/cbs.py @@ -0,0 +1,294 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# ------------------------------------------------------------------------- + +import logging +from datetime import datetime + +from .utils import utc_now, utc_from_timestamp +from .management_link import ManagementLink +from .message import Message, Properties +from .error import ( + AuthenticationException, + ErrorCondition, + TokenAuthFailure, + TokenExpired, +) +from .constants import ( + CbsState, + CbsAuthState, + CBS_PUT_TOKEN, + CBS_EXPIRATION, + CBS_NAME, + CBS_TYPE, + CBS_OPERATION, + ManagementExecuteOperationResult, + ManagementOpenResult, +) + +_LOGGER = logging.getLogger(__name__) + + +def check_expiration_and_refresh_status(expires_on, refresh_window): + seconds_since_epoc = int(utc_now().timestamp()) + is_expired = seconds_since_epoc >= expires_on + is_refresh_required = (expires_on - seconds_since_epoc) <= refresh_window + return is_expired, is_refresh_required + + +def check_put_timeout_status(auth_timeout, token_put_time): + if auth_timeout > 0: + return (int(utc_now().timestamp()) - token_put_time) >= auth_timeout + return False + + +class CBSAuthenticator(object): # pylint:disable=too-many-instance-attributes + def __init__(self, session, auth, **kwargs): + self._session = session + self._connection = self._session._connection + self._mgmt_link = self._session.create_request_response_link_pair( + endpoint="$cbs", + on_amqp_management_open_complete=self._on_amqp_management_open_complete, + on_amqp_management_error=self._on_amqp_management_error, + status_code_field=b"status-code", + status_description_field=b"status-description", + ) # type: ManagementLink + + if not auth.get_token or not callable(auth.get_token): + raise ValueError("get_token must be a callable object.") + + self._auth = auth + self._encoding = "UTF-8" + self._auth_timeout = kwargs.get("auth_timeout") + self._token_put_time = None + self._expires_on = None + self._token = None + self._refresh_window = None + self._network_trace_params = { + "amqpConnection": self._session._connection._container_id, + "amqpSession": self._session.name, + "amqpLink": None + } + + self._token_status_code = None + self._token_status_description = None + + self.state = CbsState.CLOSED + self.auth_state = CbsAuthState.IDLE + + def _put_token(self, token, token_type, audience, expires_on=None): + # type: (str, str, str, datetime) -> None + message = Message( # type: ignore # TODO: missing positional args header, etc. + value=token, + properties=Properties(message_id=self._mgmt_link.next_message_id), # type: ignore + application_properties={ + CBS_NAME: audience, + CBS_OPERATION: CBS_PUT_TOKEN, + CBS_TYPE: token_type, + CBS_EXPIRATION: expires_on, + }, + ) + self._mgmt_link.execute_operation( + message, + self._on_execute_operation_complete, + timeout=self._auth_timeout, + operation=CBS_PUT_TOKEN, + type=token_type, + ) + self._mgmt_link.next_message_id += 1 + + def _on_amqp_management_open_complete(self, management_open_result): + if self.state in (CbsState.CLOSED, CbsState.ERROR): + _LOGGER.debug( + "CSB with status: %r encounters unexpected AMQP management open complete.", + self.state, + extra=self._network_trace_params + ) + elif self.state == CbsState.OPEN: + self.state = CbsState.ERROR + _LOGGER.info( + "Unexpected AMQP management open complete in OPEN, CBS error occurred.", + extra=self._network_trace_params + ) + elif self.state == CbsState.OPENING: + self.state = ( + CbsState.OPEN + if management_open_result == ManagementOpenResult.OK + else CbsState.CLOSED + ) + _LOGGER.debug( + "CBS completed opening with status: %r", + management_open_result, + extra=self._network_trace_params + ) + + def _on_amqp_management_error(self): + if self.state == CbsState.CLOSED: + _LOGGER.info("Unexpected AMQP error in CLOSED state.", extra=self._network_trace_params) + elif self.state == CbsState.OPENING: + self.state = CbsState.ERROR + self._mgmt_link.close() + _LOGGER.info( + "CBS failed to open with status: %r", + ManagementOpenResult.ERROR, + extra=self._network_trace_params + ) + elif self.state == CbsState.OPEN: + self.state = CbsState.ERROR + _LOGGER.info("CBS error occurred.", extra=self._network_trace_params) + + def _on_execute_operation_complete( + self, + execute_operation_result, + status_code, + status_description, + _, + error_condition=None, + ): + if error_condition: + _LOGGER.info( + "CBS Put token error: %r", + error_condition, + extra=self._network_trace_params + ) + self.auth_state = CbsAuthState.ERROR + return + _LOGGER.debug( + "CBS Put token result (%r), status code: %s, status_description: %s.", + execute_operation_result, + status_code, + status_description, + extra=self._network_trace_params + ) + self._token_status_code = status_code + self._token_status_description = status_description + + if execute_operation_result == ManagementExecuteOperationResult.OK: + self.auth_state = CbsAuthState.OK + elif execute_operation_result == ManagementExecuteOperationResult.ERROR: + self.auth_state = CbsAuthState.ERROR + # put-token-message sending failure, rejected + self._token_status_code = 0 + self._token_status_description = "Auth message has been rejected." + elif ( + execute_operation_result + == ManagementExecuteOperationResult.FAILED_BAD_STATUS + ): + self.auth_state = CbsAuthState.ERROR + + def _update_status(self): + if ( + self.auth_state == CbsAuthState.OK + or self.auth_state == CbsAuthState.REFRESH_REQUIRED + ): + is_expired, is_refresh_required = check_expiration_and_refresh_status( + self._expires_on, self._refresh_window + ) + _LOGGER.debug( + "CBS status check: state == %r, expired == %r, refresh required == %r", + self.auth_state, + is_expired, + is_refresh_required, + extra=self._network_trace_params + ) + if is_expired: + self.auth_state = CbsAuthState.EXPIRED + elif is_refresh_required: + self.auth_state = CbsAuthState.REFRESH_REQUIRED + elif self.auth_state == CbsAuthState.IN_PROGRESS: + _LOGGER.debug( + "CBS update in progress. Token put time: %r", + self._token_put_time, + extra=self._network_trace_params + ) + put_timeout = check_put_timeout_status( + self._auth_timeout, self._token_put_time + ) + if put_timeout: + self.auth_state = CbsAuthState.TIMEOUT + + def _cbs_link_ready(self): + if self.state == CbsState.OPEN: + return True + if self.state != CbsState.OPEN: + return False + if self.state in (CbsState.CLOSED, CbsState.ERROR): + raise TokenAuthFailure( + status_code=ErrorCondition.ClientError, + status_description="CBS authentication link is in broken status, please recreate the cbs link.", + ) + + def open(self): + self.state = CbsState.OPENING + self._mgmt_link.open() + + def close(self): + self._mgmt_link.close() + self.state = CbsState.CLOSED + + def update_token(self): + self.auth_state = CbsAuthState.IN_PROGRESS + access_token = self._auth.get_token() + if not access_token: + _LOGGER.info( + "Token refresh function received an empty token object.", + extra=self._network_trace_params + ) + elif not access_token.token: + _LOGGER.info( + "Token refresh function received an empty token.", + extra=self._network_trace_params + ) + self._expires_on = access_token.expires_on + expires_in = self._expires_on - int(utc_now().timestamp()) + self._refresh_window = int(float(expires_in) * 0.1) + try: + self._token = access_token.token.decode() + except AttributeError: + self._token = access_token.token + self._token_put_time = int(utc_now().timestamp()) + self._put_token( + self._token, + self._auth.token_type, + self._auth.audience, + utc_from_timestamp(self._expires_on), + ) + + def handle_token(self): + if not self._cbs_link_ready(): + return False + self._update_status() + if self.auth_state == CbsAuthState.IDLE: + self.update_token() + return False + if self.auth_state == CbsAuthState.IN_PROGRESS: + return False + if self.auth_state == CbsAuthState.OK: + return True + if self.auth_state == CbsAuthState.REFRESH_REQUIRED: + _LOGGER.info( + "Token will expire soon - attempting to refresh.", + extra=self._network_trace_params + ) + self.update_token() + return False + if self.auth_state == CbsAuthState.FAILURE: + raise AuthenticationException( + condition=ErrorCondition.InternalError, + description="Failed to open CBS authentication link.", + ) + if self.auth_state == CbsAuthState.ERROR: + raise TokenAuthFailure( + self._token_status_code, + self._token_status_description, + encoding=self._encoding, # TODO: drop off all the encodings + ) + if self.auth_state == CbsAuthState.TIMEOUT: + raise TimeoutError("Authentication attempt timed-out.") + if self.auth_state == CbsAuthState.EXPIRED: + raise TokenExpired( + condition=ErrorCondition.InternalError, + description="CBS Authentication Expired.", + ) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/client.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/client.py new file mode 100644 index 000000000000..2f3bf5c668d0 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/client.py @@ -0,0 +1,982 @@ +# ------------------------------------------------------------------------- # pylint: disable=client-suffix-needed +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + +# pylint: disable=client-accepts-api-version-keyword +# pylint: disable=missing-client-constructor-parameter-credential +# pylint: disable=client-method-missing-type-annotations +# pylint: disable=too-many-lines +# TODO: Check types of kwargs (issue exists for this) +import logging +import queue +import time +import uuid +from functools import partial +from typing import Any, Dict, Optional, Tuple, Union, overload, cast +import certifi +from typing_extensions import Literal + +from ._connection import Connection +from .message import _MessageDelivery +from .error import ( + AMQPException, + ErrorCondition, + MessageException, + MessageSendFailed, + RetryPolicy, + AMQPError, +) +from .outcomes import Received, Rejected, Released, Accepted, Modified + +from .constants import ( + MAX_CHANNELS, + MessageDeliveryState, + SenderSettleMode, + ReceiverSettleMode, + LinkDeliverySettleReason, + TransportType, + SEND_DISPOSITION_ACCEPT, + SEND_DISPOSITION_REJECT, + AUTH_TYPE_CBS, + MAX_FRAME_SIZE_BYTES, + INCOMING_WINDOW, + OUTGOING_WINDOW, + DEFAULT_AUTH_TIMEOUT, + MESSAGE_DELIVERY_DONE_STATES, +) + +from .management_operation import ManagementOperation +from .cbs import CBSAuthenticator + +Outcomes = Union[Received, Rejected, Released, Accepted, Modified] + + +_logger = logging.getLogger(__name__) + + +class AMQPClient( + object +): # pylint: disable=too-many-instance-attributes + """An AMQP client. + :param hostname: The AMQP endpoint to connect to. + :type hostname: str + :keyword auth: Authentication for the connection. This should be one of the following: + - pyamqp.authentication.SASLAnonymous + - pyamqp.authentication.SASLPlain + - pyamqp.authentication.SASTokenAuth + - pyamqp.authentication.JWTTokenAuth + If no authentication is supplied, SASLAnnoymous will be used by default. + :paramtype auth: ~pyamqp.authentication + :keyword client_name: The name for the client, also known as the Container ID. + If no name is provided, a random GUID will be used. + :paramtype client_name: str or bytes + :keyword network_trace: Whether to turn on network trace logs. If `True`, trace logs + will be logged at INFO level. Default is `False`. + :paramtype network_trace: bool + :keyword retry_policy: A policy for parsing errors on link, connection and message + disposition to determine whether the error should be retryable. + :paramtype retry_policy: ~pyamqp.error.RetryPolicy + :keyword keep_alive_interval: If set, a thread will be started to keep the connection + alive during periods of user inactivity. The value will determine how long the + thread will sleep (in seconds) between pinging the connection. If 0 or None, no + thread will be started. + :paramtype keep_alive_interval: int + :keyword max_frame_size: Maximum AMQP frame size. Default is 63488 bytes. + :paramtype max_frame_size: int + :keyword channel_max: Maximum number of Session channels in the Connection. + :paramtype channel_max: int + :keyword idle_timeout: Timeout in seconds after which the Connection will close + if there is no further activity. + :paramtype idle_timeout: int + :keyword auth_timeout: Timeout in seconds for CBS authentication. Otherwise this value will be ignored. + Default value is 60s. + :paramtype auth_timeout: int + :keyword properties: Connection properties. + :paramtype properties: dict[str, any] + :keyword remote_idle_timeout_empty_frame_send_ratio: Portion of the idle timeout time to wait before sending an + empty frame. The default portion is 50% of the idle timeout value (i.e. `0.5`). + :paramtype remote_idle_timeout_empty_frame_send_ratio: float + :keyword incoming_window: The size of the allowed window for incoming messages. + :paramtype incoming_window: int + :keyword outgoing_window: The size of the allowed window for outgoing messages. + :paramtype outgoing_window: int + :keyword handle_max: The maximum number of concurrent link handles. + :paramtype handle_max: int + :keyword on_attach: A callback function to be run on receipt of an ATTACH frame. + The function must take 4 arguments: source, target, properties and error. + :paramtype on_attach: func[ + ~pyamqp.endpoint.Source, ~pyamqp.endpoint.Target, dict, ~pyamqp.error.AMQPConnectionError] + :keyword send_settle_mode: The mode by which to settle message send + operations. If set to `Unsettled`, the client will wait for a confirmation + from the service that the message was successfully sent. If set to 'Settled', + the client will not wait for confirmation and assume success. + :paramtype send_settle_mode: ~pyamqp.constants.SenderSettleMode + :keyword receive_settle_mode: The mode by which to settle message receive + operations. If set to `PeekLock`, the receiver will lock a message once received until + the client accepts or rejects the message. If set to `ReceiveAndDelete`, the service + will assume successful receipt of the message and clear it from the queue. The + default is `PeekLock`. + :paramtype receive_settle_mode: ~pyamqp.constants.ReceiverSettleMode + :keyword desired_capabilities: The extension capabilities desired from the peer endpoint. + :paramtype desired_capabilities: list[bytes] + :keyword max_message_size: The maximum allowed message size negotiated for the Link. + :paramtype max_message_size: int + :keyword link_properties: Metadata to be sent in the Link ATTACH frame. + :paramtype link_properties: dict[str, any] + :keyword link_credit: The Link credit that determines how many + messages the Link will attempt to handle per connection iteration. + The default is 300. + :paramtype link_credit: int + :keyword transport_type: The type of transport protocol that will be used for communicating with + the service. Default is `TransportType.Amqp` in which case port 5671 is used. + If the port 5671 is unavailable/blocked in the network environment, `TransportType.AmqpOverWebsocket` could + be used instead which uses port 443 for communication. + :paramtype transport_type: ~pyamqp.constants.TransportType + :keyword http_proxy: HTTP proxy settings. This must be a dictionary with the following + keys: `'proxy_hostname'` (str value) and `'proxy_port'` (int value). + Additionally the following keys may also be present: `'username', 'password'`. + :paramtype http_proxy: dict[str, str] + :keyword custom_endpoint_address: The custom endpoint address to use for establishing a connection to + the service, allowing network requests to be routed through any application gateways or + other paths needed for the host environment. Default is None. + If port is not specified in the `custom_endpoint_address`, by default port 443 will be used. + :paramtype custom_endpoint_address: str + :keyword connection_verify: Path to the custom CA_BUNDLE file of the SSL certificate which is used to + authenticate the identity of the connection endpoint. + Default is None in which case `certifi.where()` will be used. + :paramtype connection_verify: str + """ + + def __init__(self, hostname, **kwargs): + # I think these are just strings not instances of target or source + self._hostname = hostname + self._auth = kwargs.pop("auth", None) + self._name = kwargs.pop("client_name", str(uuid.uuid4())) + self._shutdown = False + self._connection = None + self._session = None + self._link = None + self._socket_timeout = False + self._external_connection = False + self._cbs_authenticator = None + self._auth_timeout = kwargs.pop("auth_timeout", DEFAULT_AUTH_TIMEOUT) + self._mgmt_links = {} + self._retry_policy = kwargs.pop("retry_policy", RetryPolicy()) + self._keep_alive_interval = int(kwargs.get("keep_alive_interval") or 0) + self._keep_alive_thread = None + + # Connection settings + self._max_frame_size = kwargs.pop("max_frame_size", MAX_FRAME_SIZE_BYTES) + self._channel_max = kwargs.pop("channel_max", MAX_CHANNELS) + self._idle_timeout = kwargs.pop("idle_timeout", None) + self._properties = kwargs.pop("properties", None) + self._remote_idle_timeout_empty_frame_send_ratio = kwargs.pop( + "remote_idle_timeout_empty_frame_send_ratio", None + ) + self._network_trace = kwargs.pop("network_trace", False) + self._network_trace_params = {"amqpConnection": None, "amqpSession": None, "amqpLink": None} + + # Session settings + self._outgoing_window = kwargs.pop("outgoing_window", OUTGOING_WINDOW) + self._incoming_window = kwargs.pop("incoming_window", INCOMING_WINDOW) + self._handle_max = kwargs.pop("handle_max", None) + + # Link settings + self._send_settle_mode = kwargs.pop( + "send_settle_mode", SenderSettleMode.Unsettled + ) + self._receive_settle_mode = kwargs.pop( + "receive_settle_mode", ReceiverSettleMode.Second + ) + self._desired_capabilities = kwargs.pop("desired_capabilities", None) + self._on_attach = kwargs.pop("on_attach", None) + + # transport + if ( + kwargs.get("transport_type") is TransportType.Amqp + and kwargs.get("http_proxy") is not None + ): + raise ValueError( + "Http proxy settings can't be passed if transport_type is explicitly set to Amqp" + ) + self._transport_type = kwargs.pop("transport_type", TransportType.Amqp) + self._http_proxy = kwargs.pop("http_proxy", None) + + # Custom Endpoint + self._custom_endpoint_address = kwargs.get("custom_endpoint_address") + self._connection_verify = kwargs.get("connection_verify") + + def __enter__(self): + """Run Client in a context manager.""" + self.open() + return self + + def __exit__(self, *args): + """Close and destroy Client on exiting a context manager.""" + self.close() + + def _client_ready(self): # pylint: disable=no-self-use + """Determine whether the client is ready to start sending and/or + receiving messages. To be ready, the connection must be open and + authentication complete. + + :rtype: bool + """ + return True + + def _client_run(self, **kwargs): + """Perform a single Connection iteration.""" + self._connection.listen(wait=self._socket_timeout, **kwargs) + + def _close_link(self): + if self._link and not self._link._is_closed: # pylint: disable=protected-access + self._link.detach(close=True) + self._link = None + + def _do_retryable_operation(self, operation, *args, **kwargs): + retry_settings = self._retry_policy.configure_retries() + retry_active = True + absolute_timeout = kwargs.pop("timeout", 0) or 0 + start_time = time.time() + while retry_active: + try: + if absolute_timeout < 0: + raise TimeoutError("Operation timed out.") + return operation(*args, timeout=absolute_timeout, **kwargs) + except AMQPException as exc: + if not self._retry_policy.is_retryable(exc): + raise + if absolute_timeout >= 0: + retry_active = self._retry_policy.increment(retry_settings, exc) + if not retry_active: + break + time.sleep(self._retry_policy.get_backoff_time(retry_settings, exc)) + if exc.condition == ErrorCondition.LinkDetachForced: + self._close_link() # if link level error, close and open a new link + if exc.condition in ( + ErrorCondition.ConnectionCloseForced, + ErrorCondition.SocketError, + ): + # if connection detach or socket error, close and open a new connection + self.close() + finally: + end_time = time.time() + if absolute_timeout > 0: + absolute_timeout -= end_time - start_time + raise retry_settings["history"][-1] + + def open(self, connection=None): + """Open the client. The client can create a new Connection + or an existing Connection can be passed in. This existing Connection + may have an existing CBS authentication Session, which will be + used for this client as well. Otherwise a new Session will be + created. + + :param connection: An existing Connection that may be shared between + multiple clients. + :type connection: ~pyamqp.Connection + """ + # pylint: disable=protected-access + if self._session: + return # already open. + if connection: + self._connection = connection + self._external_connection = True + elif not self._connection: + self._connection = Connection( + "amqps://" + self._hostname, + sasl_credential=self._auth.sasl, + ssl_opts={"ca_certs": self._connection_verify or certifi.where()}, + container_id=self._name, + max_frame_size=self._max_frame_size, + channel_max=self._channel_max, + idle_timeout=self._idle_timeout, + properties=self._properties, + network_trace=self._network_trace, + transport_type=self._transport_type, + http_proxy=self._http_proxy, + custom_endpoint_address=self._custom_endpoint_address, + ) + self._connection.open() + if not self._session: + self._session = self._connection.create_session( + incoming_window=self._incoming_window, + outgoing_window=self._outgoing_window, + ) + self._session.begin() + if self._auth.auth_type == AUTH_TYPE_CBS: + self._cbs_authenticator = CBSAuthenticator( + session=self._session, auth=self._auth, auth_timeout=self._auth_timeout + ) + self._cbs_authenticator.open() + self._network_trace_params["amqpConnection"] = self._connection._container_id + self._network_trace_params["amqpSession"] = self._session.name + self._shutdown = False + + def close(self): + """Close the client. This includes closing the Session + and CBS authentication layer as well as the Connection. + If the client was opened using an external Connection, + this will be left intact. + + No further messages can be sent or received and the client + cannot be re-opened. + + All pending, unsent messages will remain uncleared to allow + them to be inspected and queued to a new client. + """ + self._shutdown = True + if not self._session: + return # already closed. + self._close_link() + if self._cbs_authenticator: + self._cbs_authenticator.close() + self._cbs_authenticator = None + self._session.end() + self._session = None + if not self._external_connection: + self._connection.close() + self._connection = None + self._network_trace_params["amqpConnection"] = None + self._network_trace_params["amqpSession"] = None + + def auth_complete(self): + """Whether the authentication handshake is complete during + connection initialization. + + :rtype: bool + """ + if self._cbs_authenticator and not self._cbs_authenticator.handle_token(): + self._connection.listen(wait=self._socket_timeout) + return False + return True + + def client_ready(self): + """ + Whether the handler has completed all start up processes such as + establishing the connection, session, link and authentication, and + is not ready to process messages. + + :rtype: bool + """ + if not self.auth_complete(): + return False + if not self._client_ready(): + try: + self._connection.listen(wait=self._socket_timeout) + except ValueError: + return True + return False + return True + + def do_work(self, **kwargs): + """Run a single connection iteration. + This will return `True` if the connection is still open + and ready to be used for further work, or `False` if it needs + to be shut down. + + :rtype: bool + :raises: TimeoutError if CBS authentication timeout reached. + """ + if self._shutdown: + return False + if not self.client_ready(): + return True + return self._client_run(**kwargs) + + def mgmt_request(self, message, **kwargs): + """ + :param message: The message to send in the management request. + :type message: ~pyamqp.message.Message + :keyword str operation: The type of operation to be performed. This value will + be service-specific, but common values include READ, CREATE and UPDATE. + This value will be added as an application property on the message. + :keyword str operation_type: The type on which to carry out the operation. This will + be specific to the entities of the service. This value will be added as + an application property on the message. + :keyword str node: The target node. Default node is `$management`. + :keyword float timeout: Provide an optional timeout in seconds within which a response + to the management request must be received. + :rtype: ~pyamqp.message.Message + """ + + # The method also takes "status_code_field" and "status_description_field" + # keyword arguments as alternate names for the status code and description + # in the response body. Those two keyword arguments are used in Azure services only. + operation = kwargs.pop("operation", None) + operation_type = kwargs.pop("operation_type", None) + node = kwargs.pop("node", "$management") + timeout = kwargs.pop("timeout", 0) + try: + mgmt_link = self._mgmt_links[node] + except KeyError: + mgmt_link = ManagementOperation(self._session, endpoint=node, **kwargs) + self._mgmt_links[node] = mgmt_link + mgmt_link.open() + + while not mgmt_link.ready(): + self._connection.listen(wait=False) + + operation_type = operation_type or b"empty" + status, description, response = mgmt_link.execute( + message, operation=operation, operation_type=operation_type, timeout=timeout + ) + return status, description, response + + +class SendClient(AMQPClient): + """ + An AMQP client for sending messages. + :param target: The target AMQP service endpoint. This can either be the URI as + a string or a ~pyamqp.endpoint.Target object. + :type target: str, bytes or ~pyamqp.endpoint.Target + :keyword auth: Authentication for the connection. This should be one of the following: + - pyamqp.authentication.SASLAnonymous + - pyamqp.authentication.SASLPlain + - pyamqp.authentication.SASTokenAuth + - pyamqp.authentication.JWTTokenAuth + If no authentication is supplied, SASLAnnoymous will be used by default. + :paramtype auth: ~pyamqp.authentication + :keyword client_name: The name for the client, also known as the Container ID. + If no name is provided, a random GUID will be used. + :paramtype client_name: str or bytes + :keyword network_trace: Whether to turn on network trace logs. If `True`, trace logs + will be logged at INFO level. Default is `False`. + :paramtype network_trace: bool + :keyword retry_policy: A policy for parsing errors on link, connection and message + disposition to determine whether the error should be retryable. + :paramtype retry_policy: ~pyamqp.error.RetryPolicy + :keyword keep_alive_interval: If set, a thread will be started to keep the connection + alive during periods of user inactivity. The value will determine how long the + thread will sleep (in seconds) between pinging the connection. If 0 or None, no + thread will be started. + :paramtype keep_alive_interval: int + :keyword max_frame_size: Maximum AMQP frame size. Default is 63488 bytes. + :paramtype max_frame_size: int + :keyword channel_max: Maximum number of Session channels in the Connection. + :paramtype channel_max: int + :keyword idle_timeout: Timeout in seconds after which the Connection will close + if there is no further activity. + :paramtype idle_timeout: int + :keyword auth_timeout: Timeout in seconds for CBS authentication. Otherwise this value will be ignored. + Default value is 60s. + :paramtype auth_timeout: int + :keyword properties: Connection properties. + :paramtype properties: dict[str, any] + :keyword remote_idle_timeout_empty_frame_send_ratio: Portion of the idle timeout time to wait before sending an + empty frame. The default portion is 50% of the idle timeout value (i.e. `0.5`). + :paramtype remote_idle_timeout_empty_frame_send_ratio: float + :keyword incoming_window: The size of the allowed window for incoming messages. + :paramtype incoming_window: int + :keyword outgoing_window: The size of the allowed window for outgoing messages. + :paramtype outgoing_window: int + :keyword handle_max: The maximum number of concurrent link handles. + :paramtype handle_max: int + :keyword on_attach: A callback function to be run on receipt of an ATTACH frame. + The function must take 4 arguments: source, target, properties and error. + :paramtype on_attach: func[ + ~pyamqp.endpoint.Source, ~pyamqp.endpoint.Target, dict, ~pyamqp.error.AMQPConnectionError] + :keyword send_settle_mode: The mode by which to settle message send + operations. If set to `Unsettled`, the client will wait for a confirmation + from the service that the message was successfully sent. If set to 'Settled', + the client will not wait for confirmation and assume success. + :paramtype send_settle_mode: ~pyamqp.constants.SenderSettleMode + :keyword receive_settle_mode: The mode by which to settle message receive + operations. If set to `PeekLock`, the receiver will lock a message once received until + the client accepts or rejects the message. If set to `ReceiveAndDelete`, the service + will assume successful receipt of the message and clear it from the queue. The + default is `PeekLock`. + :paramtype receive_settle_mode: ~pyamqp.constants.ReceiverSettleMode + :keyword desired_capabilities: The extension capabilities desired from the peer endpoint. + :paramtype desired_capabilities: list[bytes] + :keyword max_message_size: The maximum allowed message size negotiated for the Link. + :paramtype max_message_size: int + :keyword link_properties: Metadata to be sent in the Link ATTACH frame. + :paramtype link_properties: dict[str, any] + :keyword link_credit: The Link credit that determines how many + messages the Link will attempt to handle per connection iteration. + The default is 300. + :paramtype link_credit: int + :keyword transport_type: The type of transport protocol that will be used for communicating with + the service. Default is `TransportType.Amqp` in which case port 5671 is used. + If the port 5671 is unavailable/blocked in the network environment, `TransportType.AmqpOverWebsocket` could + be used instead which uses port 443 for communication. + :paramtype transport_type: ~pyamqp.constants.TransportType + :keyword http_proxy: HTTP proxy settings. This must be a dictionary with the following + keys: `'proxy_hostname'` (str value) and `'proxy_port'` (int value). + Additionally the following keys may also be present: `'username', 'password'`. + :paramtype http_proxy: dict[str, str] + :keyword custom_endpoint_address: The custom endpoint address to use for establishing a connection to + the service, allowing network requests to be routed through any application gateways or + other paths needed for the host environment. Default is None. + If port is not specified in the `custom_endpoint_address`, by default port 443 will be used. + :paramtype custom_endpoint_address: str + :keyword connection_verify: Path to the custom CA_BUNDLE file of the SSL certificate which is used to + authenticate the identity of the connection endpoint. + Default is None in which case `certifi.where()` will be used. + :paramtype connection_verify: str + """ + + def __init__(self, hostname, target, **kwargs): + self.target = target + # Sender and Link settings + self._max_message_size = kwargs.pop("max_message_size", MAX_FRAME_SIZE_BYTES) + self._link_properties = kwargs.pop("link_properties", None) + self._link_credit = kwargs.pop("link_credit", None) + super(SendClient, self).__init__(hostname, **kwargs) + + def _client_ready(self): + """Determine whether the client is ready to start receiving messages. + To be ready, the connection must be open and authentication complete, + The Session, Link and MessageReceiver must be open and in non-errored + states. + + :rtype: bool + """ + # pylint: disable=protected-access + if not self._link: + self._link = self._session.create_sender_link( + target_address=self.target, + link_credit=self._link_credit, + send_settle_mode=self._send_settle_mode, + rcv_settle_mode=self._receive_settle_mode, + max_message_size=self._max_message_size, + properties=self._link_properties, + ) + self._link.attach() + return False + if self._link.get_state().value != 3: # ATTACHED + return False + return True + + def _client_run(self, **kwargs): + """MessageSender Link is now open - perform message send + on all pending messages. + Will return True if operation successful and client can remain open for + further work. + + :rtype: bool + """ + self._link.update_pending_deliveries() + self._connection.listen(wait=self._socket_timeout, **kwargs) + return True + + def _transfer_message(self, message_delivery, timeout=0): + message_delivery.state = MessageDeliveryState.WaitingForSendAck + on_send_complete = partial(self._on_send_complete, message_delivery) + delivery = self._link.send_transfer( + message_delivery.message, + on_send_complete=on_send_complete, + timeout=timeout, + send_async=True, + ) + return delivery + + @staticmethod + def _process_send_error(message_delivery, condition, description=None, info=None): + try: + amqp_condition = ErrorCondition(condition) + except ValueError: + error = MessageException(condition, description=description, info=info) + else: + error = MessageSendFailed( + amqp_condition, description=description, info=info + ) + message_delivery.state = MessageDeliveryState.Error + message_delivery.error = error + + def _on_send_complete(self, message_delivery, reason, state): + message_delivery.reason = reason + if reason == LinkDeliverySettleReason.DISPOSITION_RECEIVED: + if state and SEND_DISPOSITION_ACCEPT in state: + message_delivery.state = MessageDeliveryState.Ok + else: + try: + error_info = state[SEND_DISPOSITION_REJECT] + self._process_send_error( + message_delivery, + condition=error_info[0][0], + description=error_info[0][1], + info=error_info[0][2], + ) + except TypeError: + self._process_send_error( + message_delivery, condition=ErrorCondition.UnknownError + ) + elif reason == LinkDeliverySettleReason.SETTLED: + message_delivery.state = MessageDeliveryState.Ok + elif reason == LinkDeliverySettleReason.TIMEOUT: + message_delivery.state = MessageDeliveryState.Timeout + message_delivery.error = TimeoutError("Sending message timed out.") + else: + # NotDelivered and other unknown errors + self._process_send_error( + message_delivery, condition=ErrorCondition.UnknownError + ) + + def _send_message_impl(self, message, **kwargs): + timeout = kwargs.pop("timeout", 0) + expire_time = (time.time() + timeout) if timeout else None + self.open() + message_delivery = _MessageDelivery( + message, MessageDeliveryState.WaitingToBeSent, expire_time + ) + while not self.client_ready(): + time.sleep(0.05) + + self._transfer_message(message_delivery, timeout) + running = True + while running and message_delivery.state not in MESSAGE_DELIVERY_DONE_STATES: + running = self.do_work() + if message_delivery.state not in MESSAGE_DELIVERY_DONE_STATES: + raise MessageException( + condition=ErrorCondition.ClientError, + description="Send failed - connection not running." + ) + + if message_delivery.state in ( + MessageDeliveryState.Error, + MessageDeliveryState.Cancelled, + MessageDeliveryState.Timeout, + ): + try: + raise message_delivery.error # pylint: disable=raising-bad-type + except TypeError: + # This is a default handler + raise MessageException( + condition=ErrorCondition.UnknownError, description="Send failed." + ) + + def send_message(self, message, **kwargs): + """ + :param ~pyamqp.message.Message message: + :keyword float timeout: timeout in seconds. If set to + 0, the client will continue to wait until the message is sent or error happens. The + default is 0. + """ + self._do_retryable_operation(self._send_message_impl, message=message, **kwargs) + + +class ReceiveClient(AMQPClient): + """ + An AMQP client for receiving messages. + :param source: The source AMQP service endpoint. This can either be the URI as + a string or a ~pyamqp.endpoint.Source object. + :type source: str, bytes or ~pyamqp.endpoint.Source + :keyword auth: Authentication for the connection. This should be one of the following: + - pyamqp.authentication.SASLAnonymous + - pyamqp.authentication.SASLPlain + - pyamqp.authentication.SASTokenAuth + - pyamqp.authentication.JWTTokenAuth + If no authentication is supplied, SASLAnnoymous will be used by default. + :paramtype auth: ~pyamqp.authentication + :keyword client_name: The name for the client, also known as the Container ID. + If no name is provided, a random GUID will be used. + :paramtype client_name: str or bytes + :keyword network_trace: Whether to turn on network trace logs. If `True`, trace logs + will be logged at INFO level. Default is `False`. + :paramtype network_trace: bool + :keyword retry_policy: A policy for parsing errors on link, connection and message + disposition to determine whether the error should be retryable. + :paramtype retry_policy: ~pyamqp.error.RetryPolicy + :keyword keep_alive_interval: If set, a thread will be started to keep the connection + alive during periods of user inactivity. The value will determine how long the + thread will sleep (in seconds) between pinging the connection. If 0 or None, no + thread will be started. + :paramtype keep_alive_interval: int + :keyword max_frame_size: Maximum AMQP frame size. Default is 63488 bytes. + :paramtype max_frame_size: int + :keyword channel_max: Maximum number of Session channels in the Connection. + :paramtype channel_max: int + :keyword idle_timeout: Timeout in seconds after which the Connection will close + if there is no further activity. + :paramtype idle_timeout: int + :keyword auth_timeout: Timeout in seconds for CBS authentication. Otherwise this value will be ignored. + Default value is 60s. + :paramtype auth_timeout: int + :keyword properties: Connection properties. + :paramtype properties: dict[str, any] + :keyword remote_idle_timeout_empty_frame_send_ratio: Portion of the idle timeout time to wait before sending an + empty frame. The default portion is 50% of the idle timeout value (i.e. `0.5`). + :paramtype remote_idle_timeout_empty_frame_send_ratio: float + :keyword incoming_window: The size of the allowed window for incoming messages. + :paramtype incoming_window: int + :keyword outgoing_window: The size of the allowed window for outgoing messages. + :paramtype outgoing_window: int + :keyword handle_max: The maximum number of concurrent link handles. + :paramtype handle_max: int + :keyword on_attach: A callback function to be run on receipt of an ATTACH frame. + The function must take 4 arguments: source, target, properties and error. + :paramtype on_attach: func[ + ~pyamqp.endpoint.Source, ~pyamqp.endpoint.Target, dict, ~pyamqp.error.AMQPConnectionError] + :keyword send_settle_mode: The mode by which to settle message send + operations. If set to `Unsettled`, the client will wait for a confirmation + from the service that the message was successfully sent. If set to 'Settled', + the client will not wait for confirmation and assume success. + :paramtype send_settle_mode: ~pyamqp.constants.SenderSettleMode + :keyword receive_settle_mode: The mode by which to settle message receive + operations. If set to `PeekLock`, the receiver will lock a message once received until + the client accepts or rejects the message. If set to `ReceiveAndDelete`, the service + will assume successful receipt of the message and clear it from the queue. The + default is `PeekLock`. + :paramtype receive_settle_mode: ~pyamqp.constants.ReceiverSettleMode + :keyword desired_capabilities: The extension capabilities desired from the peer endpoint. + :paramtype desired_capabilities: list[bytes] + :keyword max_message_size: The maximum allowed message size negotiated for the Link. + :paramtype max_message_size: int + :keyword link_properties: Metadata to be sent in the Link ATTACH frame. + :paramtype link_properties: dict[str, any] + :keyword link_credit: The Link credit that determines how many + messages the Link will attempt to handle per connection iteration. + The default is 300. + :paramtype link_credit: int + :keyword transport_type: The type of transport protocol that will be used for communicating with + the service. Default is `TransportType.Amqp` in which case port 5671 is used. + If the port 5671 is unavailable/blocked in the network environment, `TransportType.AmqpOverWebsocket` could + be used instead which uses port 443 for communication. + :paramtype transport_type: ~pyamqp.constants.TransportType + :keyword http_proxy: HTTP proxy settings. This must be a dictionary with the following + keys: `'proxy_hostname'` (str value) and `'proxy_port'` (int value). + Additionally the following keys may also be present: `'username', 'password'`. + :paramtype http_proxy: dict[str, str] + :keyword custom_endpoint_address: The custom endpoint address to use for establishing a connection to + the service, allowing network requests to be routed through any application gateways or + other paths needed for the host environment. Default is None. + If port is not specified in the `custom_endpoint_address`, by default port 443 will be used. + :paramtype custom_endpoint_address: str + :keyword connection_verify: Path to the custom CA_BUNDLE file of the SSL certificate which is used to + authenticate the identity of the connection endpoint. + Default is None in which case `certifi.where()` will be used. + :paramtype connection_verify: str + """ + + def __init__(self, hostname, source, **kwargs): + self.source = source + self._streaming_receive = kwargs.pop("streaming_receive", False) + self._received_messages = queue.Queue() + self._message_received_callback = kwargs.pop("message_received_callback", None) + + # Sender and Link settings + self._max_message_size = kwargs.pop("max_message_size", MAX_FRAME_SIZE_BYTES) + self._link_properties = kwargs.pop("link_properties", None) + self._link_credit = kwargs.pop("link_credit", 300) + super(ReceiveClient, self).__init__(hostname, **kwargs) + + def _client_ready(self): + """Determine whether the client is ready to start receiving messages. + To be ready, the connection must be open and authentication complete, + The Session, Link and MessageReceiver must be open and in non-errored + states. + + :rtype: bool + """ + # pylint: disable=protected-access + if not self._link: + self._link = self._session.create_receiver_link( + source_address=self.source, + link_credit=self._link_credit, + send_settle_mode=self._send_settle_mode, + rcv_settle_mode=self._receive_settle_mode, + max_message_size=self._max_message_size, + on_transfer=self._message_received, + properties=self._link_properties, + desired_capabilities=self._desired_capabilities, + on_attach=self._on_attach, + ) + self._link.attach() + return False + if self._link.get_state().value != 3: # ATTACHED + return False + return True + + def _client_run(self, **kwargs): + """MessageReceiver Link is now open - start receiving messages. + Will return True if operation successful and client can remain open for + further work. + + :rtype: bool + """ + try: + self._link.flow() + self._connection.listen(wait=self._socket_timeout, **kwargs) + except ValueError: + _logger.info("Timeout reached, closing receiver.", extra=self._network_trace_params) + self._shutdown = True + return False + return True + + def _message_received(self, frame, message): + """Callback run on receipt of every message. If there is + a user-defined callback, this will be called. + Additionally if the client is retrieving messages for a batch + or iterator, the message will be added to an internal queue. + + :param message: Received message. + :type message: ~pyamqp.message.Message + """ + if self._message_received_callback: + self._message_received_callback(message) + if not self._streaming_receive: + self._received_messages.put((frame, message)) + + def _receive_message_batch_impl( + self, max_batch_size=None, on_message_received=None, timeout=0 + ): + self._message_received_callback = on_message_received + max_batch_size = max_batch_size or self._link_credit + timeout = time.time() + timeout if timeout else 0 + receiving = True + batch = [] + self.open() + while len(batch) < max_batch_size: + try: + # TODO: This drops the transfer frame data + _, message = self._received_messages.get_nowait() + batch.append(message) + self._received_messages.task_done() + except queue.Empty: + break + else: + return batch + + to_receive_size = max_batch_size - len(batch) + before_queue_size = self._received_messages.qsize() + + while receiving and to_receive_size > 0: + if timeout and time.time() > timeout: + break + + receiving = self.do_work(batch=to_receive_size) + cur_queue_size = self._received_messages.qsize() + # after do_work, check how many new messages have been received since previous iteration + received = cur_queue_size - before_queue_size + if to_receive_size < max_batch_size and received == 0: + # there are already messages in the batch, and no message is received in the current cycle + # return what we have + break + + to_receive_size -= received + before_queue_size = cur_queue_size + + while len(batch) < max_batch_size: + try: + _, message = self._received_messages.get_nowait() + batch.append(message) + self._received_messages.task_done() + except queue.Empty: + break + return batch + + def close(self): + self._received_messages = queue.Queue() + super(ReceiveClient, self).close() + + def receive_message_batch(self, **kwargs): + """Receive a batch of messages. Messages returned in the batch have already been + accepted - if you wish to add logic to accept or reject messages based on custom + criteria, pass in a callback. This method will return as soon as some messages are + available rather than waiting to achieve a specific batch size, and therefore the + number of messages returned per call will vary up to the maximum allowed. + + :param max_batch_size: The maximum number of messages that can be returned in + one call. This value cannot be larger than the prefetch value, and if not specified, + the prefetch value will be used. + :type max_batch_size: int + :param on_message_received: A callback to process messages as they arrive from the + service. It takes a single argument, a ~pyamqp.message.Message object. + :type on_message_received: callable[~pyamqp.message.Message] + :param timeout: The timeout in milliseconds for which to wait to receive any messages. + If no messages are received in this time, an empty list will be returned. If set to + 0, the client will continue to wait until at least one message is received. The + default is 0. + :type timeout: float + """ + return self._do_retryable_operation(self._receive_message_batch_impl, **kwargs) + + @overload + def settle_messages( + self, + delivery_id: Union[int, Tuple[int, int]], + outcome: Literal["accepted"], + *, + batchable: Optional[bool] = None + ): + ... + + @overload + def settle_messages( + self, + delivery_id: Union[int, Tuple[int, int]], + outcome: Literal["released"], + *, + batchable: Optional[bool] = None + ): + ... + + @overload + def settle_messages( + self, + delivery_id: Union[int, Tuple[int, int]], + outcome: Literal["rejected"], + *, + error: Optional[AMQPError] = None, + batchable: Optional[bool] = None + ): + ... + + @overload + def settle_messages( + self, + delivery_id: Union[int, Tuple[int, int]], + outcome: Literal["modified"], + *, + delivery_failed: Optional[bool] = None, + undeliverable_here: Optional[bool] = None, + message_annotations: Optional[Dict[Union[str, bytes], Any]] = None, + batchable: Optional[bool] = None + ): + ... + + @overload + def settle_messages( + self, + delivery_id: Union[int, Tuple[int, int]], + outcome: Literal["received"], + *, + section_number: int, + section_offset: int, + batchable: Optional[bool] = None + ): + ... + + def settle_messages( + self, delivery_id: Union[int, Tuple[int, int]], outcome: str, **kwargs + ): + batchable = kwargs.pop("batchable", None) + if outcome.lower() == "accepted": + state: Outcomes = Accepted() + elif outcome.lower() == "released": + state = Released() + elif outcome.lower() == "rejected": + state = Rejected(**kwargs) + elif outcome.lower() == "modified": + state = Modified(**kwargs) + elif outcome.lower() == "received": + state = Received(**kwargs) + else: + raise ValueError("Unrecognized message output: {}".format(outcome)) + try: + first, last = cast(Tuple, delivery_id) + except TypeError: + first = delivery_id + last = None + self._link.send_disposition( + first_delivery_id=first, + last_delivery_id=last, + settled=True, + delivery_state=state, + batchable=batchable, + wait=True, + ) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/constants.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/constants.py new file mode 100644 index 000000000000..efe4ebd24ccc --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/constants.py @@ -0,0 +1,338 @@ +#------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +#-------------------------------------------------------------------------- +from typing import cast +from collections import namedtuple +from enum import Enum +import struct + +_AS_BYTES = struct.Struct('>B') + +#: The IANA assigned port number for AMQP.The standard AMQP port number that has been assigned by IANA +#: for TCP, UDP, and SCTP.There are currently no UDP or SCTP mappings defined for AMQP. +#: The port number is reserved for future transport mappings to these protocols. +PORT = 5672 + +# default port for AMQP over Websocket +WEBSOCKET_PORT = 443 + +# subprotocol for AMQP over Websocket +AMQP_WS_SUBPROTOCOL = 'AMQPWSB10' + +#: The IANA assigned port number for secure AMQP (amqps).The standard AMQP port number that has been assigned +#: by IANA for secure TCP using TLS. Implementations listening on this port should NOT expect a protocol +#: handshake before TLS is negotiated. +SECURE_PORT = 5671 + + +# default port for AMQP over Websocket +WEBSOCKET_PORT = 443 + + +# subprotocol for AMQP over Websocket +AMQP_WS_SUBPROTOCOL = 'AMQPWSB10' + + +MAJOR = 1 #: Major protocol version. +MINOR = 0 #: Minor protocol version. +REV = 0 #: Protocol revision. +HEADER_FRAME = b"AMQP\x00" + _AS_BYTES.pack(MAJOR) + _AS_BYTES.pack(MINOR) + _AS_BYTES.pack(REV) + + +TLS_MAJOR = 1 #: Major protocol version. +TLS_MINOR = 0 #: Minor protocol version. +TLS_REV = 0 #: Protocol revision. +TLS_HEADER_FRAME = b"AMQP\x02" + _AS_BYTES.pack(TLS_MAJOR) + _AS_BYTES.pack(TLS_MINOR) + _AS_BYTES.pack(TLS_REV) + +SASL_MAJOR = 1 #: Major protocol version. +SASL_MINOR = 0 #: Minor protocol version. +SASL_REV = 0 #: Protocol revision. +SASL_HEADER_FRAME = b"AMQP\x03" + _AS_BYTES.pack(SASL_MAJOR) + _AS_BYTES.pack(SASL_MINOR) + _AS_BYTES.pack(SASL_REV) + +EMPTY_FRAME = b'\x00\x00\x00\x08\x02\x00\x00\x00' + +#: The lower bound for the agreed maximum frame size (in bytes). During the initial Connection negotiation, the +#: two peers must agree upon a maximum frame size. This constant defines the minimum value to which the maximum +#: frame size can be set. By defining this value, the peers can guarantee that they can send frames of up to this +#: size until they have agreed a definitive maximum frame size for that Connection. +MIN_MAX_FRAME_SIZE = 512 +MAX_FRAME_SIZE_BYTES = 1024 * 1024 +MAX_CHANNELS = 65535 +INCOMING_WINDOW = 64 * 1024 +OUTGOING_WINDOW = 64 * 1024 + +DEFAULT_LINK_CREDIT = 10000 + +FIELD = namedtuple('FIELD', 'name, type, mandatory, default, multiple') + +STRING_FILTER = b"apache.org:selector-filter:string" + +DEFAULT_AUTH_TIMEOUT = 60 +AUTH_DEFAULT_EXPIRATION_SECONDS = 3600 +TOKEN_TYPE_JWT = "jwt" +TOKEN_TYPE_SASTOKEN = "servicebus.windows.net:sastoken" +CBS_PUT_TOKEN = "put-token" +CBS_NAME = "name" +CBS_OPERATION = "operation" +CBS_TYPE = "type" +CBS_EXPIRATION = "expiration" + +SEND_DISPOSITION_ACCEPT = "accepted" +SEND_DISPOSITION_REJECT = "rejected" + +AUTH_TYPE_SASL_PLAIN = "AUTH_SASL_PLAIN" +AUTH_TYPE_CBS = "AUTH_CBS" + +DEFAULT_WEBSOCKET_HEARTBEAT_SECONDS = 10 + + +class ConnectionState(Enum): + #: In this state a Connection exists, but nothing has been sent or received. This is the state an + #: implementation would be in immediately after performing a socket connect or socket accept. + START = 0 + #: In this state the Connection header has been received from our peer, but we have not yet sent anything. + HDR_RCVD = 1 + #: In this state the Connection header has been sent to our peer, but we have not yet received anything. + HDR_SENT = 2 + #: In this state we have sent and received the Connection header, but we have not yet sent or + #: received an open frame. + HDR_EXCH = 3 + #: In this state we have sent both the Connection header and the open frame, but + #: we have not yet received anything. + OPEN_PIPE = 4 + #: In this state we have sent the Connection header, the open frame, any pipelined Connection traffic, + #: and the close frame, but we have not yet received anything. + OC_PIPE = 5 + #: In this state we have sent and received the Connection header, and received an open frame from + #: our peer, but have not yet sent an open frame. + OPEN_RCVD = 6 + #: In this state we have sent and received the Connection header, and sent an open frame to our peer, + #: but have not yet received an open frame. + OPEN_SENT = 7 + #: In this state we have send and received the Connection header, sent an open frame, any pipelined + #: Connection traffic, and the close frame, but we have not yet received an open frame. + CLOSE_PIPE = 8 + #: In this state the Connection header and the open frame have both been sent and received. + OPENED = 9 + #: In this state we have received a close frame indicating that our partner has initiated a close. + #: This means we will never have to read anything more from this Connection, however we can + #: continue to write frames onto the Connection. If desired, an implementation could do a TCP half-close + #: at this point to shutdown the read side of the Connection. + CLOSE_RCVD = 10 + #: In this state we have sent a close frame to our partner. It is illegal to write anything more onto + #: the Connection, however there may still be incoming frames. If desired, an implementation could do + #: a TCP half-close at this point to shutdown the write side of the Connection. + CLOSE_SENT = 11 + #: The DISCARDING state is a variant of the CLOSE_SENT state where the close is triggered by an error. + #: In this case any incoming frames on the connection MUST be silently discarded until the peer's close + #: frame is received. + DISCARDING = 12 + #: In this state it is illegal for either endpoint to write anything more onto the Connection. The + #: Connection may be safely closed and discarded. + END = 13 + + +class SessionState(Enum): + #: In the UNMAPPED state, the Session endpoint is not mapped to any incoming or outgoing channels on the + #: Connection endpoint. In this state an endpoint cannot send or receive frames. + UNMAPPED = 0 + #: In the BEGIN_SENT state, the Session endpoint is assigned an outgoing channel number, but there is no entry + #: in the incoming channel map. In this state the endpoint may send frames but cannot receive them. + BEGIN_SENT = 1 + #: In the BEGIN_RCVD state, the Session endpoint has an entry in the incoming channel map, but has not yet + #: been assigned an outgoing channel number. The endpoint may receive frames, but cannot send them. + BEGIN_RCVD = 2 + #: In the MAPPED state, the Session endpoint has both an outgoing channel number and an entry in the incoming + #: channel map. The endpoint may both send and receive frames. + MAPPED = 3 + #: In the END_SENT state, the Session endpoint has an entry in the incoming channel map, but is no longer + #: assigned an outgoing channel number. The endpoint may receive frames, but cannot send them. + END_SENT = 4 + #: In the END_RCVD state, the Session endpoint is assigned an outgoing channel number, but there is no entry in + #: the incoming channel map. The endpoint may send frames, but cannot receive them. + END_RCVD = 5 + #: The DISCARDING state is a variant of the END_SENT state where the end is triggered by an error. In this + #: case any incoming frames on the session MUST be silently discarded until the peer's end frame is received. + DISCARDING = 6 + + +class SessionTransferState(Enum): + + OKAY = 0 + ERROR = 1 + BUSY = 2 + + +class LinkDeliverySettleReason(Enum): + + DISPOSITION_RECEIVED = 0 + SETTLED = 1 + NOT_DELIVERED = 2 + TIMEOUT = 3 + CANCELLED = 4 + + +class LinkState(Enum): + + DETACHED = 0 + ATTACH_SENT = 1 + ATTACH_RCVD = 2 + ATTACHED = 3 + DETACH_SENT = 4 + DETACH_RCVD = 5 + ERROR = 6 + + +class ManagementLinkState(Enum): + + IDLE = 0 + OPENING = 1 + CLOSING = 2 + OPEN = 3 + ERROR = 4 + + +class ManagementOpenResult(Enum): + + OPENING = 0 + OK = 1 + ERROR = 2 + CANCELLED = 3 + + +class ManagementExecuteOperationResult(Enum): + + OK = 0 + ERROR = 1 + FAILED_BAD_STATUS = 2 + LINK_CLOSED = 3 + + +class CbsState(Enum): + CLOSED = 0 + OPENING = 1 + OPEN = 2 + ERROR = 3 + + +class CbsAuthState(Enum): + OK = 0 + IDLE = 1 + IN_PROGRESS = 2 + TIMEOUT = 3 + REFRESH_REQUIRED = 4 + EXPIRED = 5 + ERROR = 6 # Put token rejected or complete but fail authentication + FAILURE = 7 # Fail to open cbs links + + +class Role(object): + """Link endpoint role. + + Valid Values: + - False: Sender + - True: Receiver + + + + + + """ + Sender = False + Receiver = True + + +class SenderSettleMode(object): + """Settlement policy for a Sender. + + Valid Values: + - 0: The Sender will send all deliveries initially unsettled to the Receiver. + - 1: The Sender will send all deliveries settled to the Receiver. + - 2: The Sender may send a mixture of settled and unsettled deliveries to the Receiver. + + + + + + + """ + Unsettled = 0 + Settled = 1 + Mixed = 2 + + +class ReceiverSettleMode(object): + """Settlement policy for a Receiver. + + Valid Values: + - 0: The Receiver will spontaneously settle all incoming transfers. + - 1: The Receiver will only settle after sending the disposition to the Sender and + receiving a disposition indicating settlement of the delivery from the sender. + + + + + + """ + First = 0 + Second = 1 + + +class SASLCode(object): + """Codes to indicate the outcome of the sasl dialog. + + + + + + + + + """ + #: Connection authentication succeeded. + Ok = 0 + #: Connection authentication failed due to an unspecified problem with the supplied credentials. + Auth = 1 + #: Connection authentication failed due to a system error. + Sys = 2 + #: Connection authentication failed due to a system error that is unlikely to be corrected without intervention. + SysPerm = 3 + #: Connection authentication failed due to a transient system error. + SysTemp = 4 + + +class MessageDeliveryState(object): + + WaitingToBeSent = 0 + WaitingForSendAck = 1 + Ok = 2 + Error = 3 + Timeout = 4 + Cancelled = 5 + + +MESSAGE_DELIVERY_DONE_STATES = ( + MessageDeliveryState.Ok, + MessageDeliveryState.Error, + MessageDeliveryState.Timeout, + MessageDeliveryState.Cancelled +) + +class TransportType(Enum): + """Transport type + The underlying transport protocol type: + Amqp: AMQP over the default TCP transport protocol, it uses port 5671. + AmqpOverWebsocket: Amqp over the Web Sockets transport protocol, it uses + port 443. + """ + Amqp = 1 + AmqpOverWebsocket = 2 + + def __eq__(self, __o: object) -> bool: + try: + __o = cast(Enum, __o) + return self.value == __o.value + except AttributeError: + return super().__eq__(__o) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/endpoints.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/endpoints.py new file mode 100644 index 000000000000..2d2de0a2868e --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/endpoints.py @@ -0,0 +1,278 @@ +#------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +#-------------------------------------------------------------------------- + +# The messaging layer defines two concrete types (source and target) to be used as the source and target of a +# link. These types are supplied in the source and target fields of the attach frame when establishing or +# resuming link. The source is comprised of an address (which the container of the outgoing Link Endpoint will +# resolve to a Node within that container) coupled with properties which determine: +# +# - which messages from the sending Node will be sent on the Link +# - how sending the message affects the state of that message at the sending Node +# - the behavior of Messages which have been transferred on the Link, but have not yet reached a +# terminal state at the receiver, when the source is destroyed. + +# TODO: fix mypy errors for _code/_definition/__defaults__ (issue #26500) +from collections import namedtuple + +from .types import AMQPTypes, FieldDefinition, ObjDefinition +from .constants import FIELD +from .performatives import _CAN_ADD_DOCSTRING + + +class TerminusDurability(object): + """Durability policy for a terminus. + + + + + + + + Determines which state of the terminus is held durably. + """ + #: No Terminus state is retained durably + NoDurability = 0 + #: Only the existence and configuration of the Terminus is retained durably. + Configuration = 1 + #: In addition to the existence and configuration of the Terminus, the unsettled state for durable + #: messages is retained durably. + UnsettledState = 2 + + +class ExpiryPolicy(object): + """Expiry policy for a terminus. + + + + + + + + + Determines when the expiry timer of a terminus starts counting down from the timeout + value. If the link is subsequently re-attached before the terminus is expired, then the + count down is aborted. If the conditions for the terminus-expiry-policy are subsequently + re-met, the expiry timer restarts from its originally configured timeout value. + """ + #: The expiry timer starts when Terminus is detached. + LinkDetach = b"link-detach" + #: The expiry timer starts when the most recently associated session is ended. + SessionEnd = b"session-end" + #: The expiry timer starts when most recently associated connection is closed. + ConnectionClose = b"connection-close" + #: The Terminus never expires. + Never = b"never" + + +class DistributionMode(object): + """Link distribution policy. + + + + + + + Policies for distributing messages when multiple links are connected to the same node. + """ + #: Once successfully transferred over the link, the message will no longer be available + #: to other links from the same node. + Move = b'move' + #: Once successfully transferred over the link, the message is still available for other + #: links from the same node. + Copy = b'copy' + + +class LifeTimePolicy(object): + #: Lifetime of dynamic node scoped to lifetime of link which caused creation. + #: A node dynamically created with this lifetime policy will be deleted at the point that the link + #: which caused its creation ceases to exist. + DeleteOnClose = 0x0000002b + #: Lifetime of dynamic node scoped to existence of links to the node. + #: A node dynamically created with this lifetime policy will be deleted at the point that there remain + #: no links for which the node is either the source or target. + DeleteOnNoLinks = 0x0000002c + #: Lifetime of dynamic node scoped to existence of messages on the node. + #: A node dynamically created with this lifetime policy will be deleted at the point that the link which + #: caused its creation no longer exists and there remain no messages at the node. + DeleteOnNoMessages = 0x0000002d + #: Lifetime of node scoped to existence of messages on or links to the node. + #: A node dynamically created with this lifetime policy will be deleted at the point that the there are no + #: links which have this node as their source or target, and there remain no messages at the node. + DeleteOnNoLinksOrMessages = 0x0000002e + + +class SupportedOutcomes(object): + #: Indicates successful processing at the receiver. + accepted = b"amqp:accepted:list" + #: Indicates an invalid and unprocessable message. + rejected = b"amqp:rejected:list" + #: Indicates that the message was not (and will not be) processed. + released = b"amqp:released:list" + #: Indicates that the message was modified, but not processed. + modified = b"amqp:modified:list" + + +class ApacheFilters(object): + #: Exact match on subject - analogous to legacy AMQP direct exchange bindings. + legacy_amqp_direct_binding = b"apache.org:legacy-amqp-direct-binding:string" + #: Pattern match on subject - analogous to legacy AMQP topic exchange bindings. + legacy_amqp_topic_binding = b"apache.org:legacy-amqp-topic-binding:string" + #: Matching on message headers - analogous to legacy AMQP headers exchange bindings. + legacy_amqp_headers_binding = b"apache.org:legacy-amqp-headers-binding:map" + #: Filter out messages sent from the same connection as the link is currently associated with. + no_local_filter = b"apache.org:no-local-filter:list" + #: SQL-based filtering syntax. + selector_filter = b"apache.org:selector-filter:string" + + +Source = namedtuple( + 'Source', + [ + 'address', + 'durable', + 'expiry_policy', + 'timeout', + 'dynamic', + 'dynamic_node_properties', + 'distribution_mode', + 'filters', + 'default_outcome', + 'outcomes', + 'capabilities' + ]) +Source.__new__.__defaults__ = (None,) * len(Source._fields) # type: ignore +Source._code = 0x00000028 # type: ignore # pylint: disable=protected-access +Source._definition = ( # type: ignore # pylint: disable=protected-access + FIELD("address", AMQPTypes.string, False, None, False), + FIELD("durable", AMQPTypes.uint, False, "none", False), + FIELD("expiry_policy", AMQPTypes.symbol, False, ExpiryPolicy.SessionEnd, False), + FIELD("timeout", AMQPTypes.uint, False, 0, False), + FIELD("dynamic", AMQPTypes.boolean, False, False, False), + FIELD("dynamic_node_properties", FieldDefinition.node_properties, False, None, False), + FIELD("distribution_mode", AMQPTypes.symbol, False, None, False), + FIELD("filters", FieldDefinition.filter_set, False, None, False), + FIELD("default_outcome", ObjDefinition.delivery_state, False, None, False), + FIELD("outcomes", AMQPTypes.symbol, False, None, True), + FIELD("capabilities", AMQPTypes.symbol, False, None, True)) +if _CAN_ADD_DOCSTRING: + Source.__doc__ = """ + For containers which do not implement address resolution (and do not admit spontaneous link + attachment from their partners) but are instead only used as producers of messages, it is unnecessary to provide + spurious detail on the source. For this purpose it is possible to use a "minimal" source in which all the + fields are left unset. + + :param str address: The address of the source. + The address of the source MUST NOT be set when sent on a attach frame sent by the receiving Link Endpoint + where the dynamic fiag is set to true (that is where the receiver is requesting the sender to create an + addressable node). The address of the source MUST be set when sent on a attach frame sent by the sending + Link Endpoint where the dynamic fiag is set to true (that is where the sender has created an addressable + node at the request of the receiver and is now communicating the address of that created node). + The generated name of the address SHOULD include the link name and the container-id of the remote container + to allow for ease of identification. + :param ~uamqp.endpoints.TerminusDurability durable: Indicates the durability of the terminus. + Indicates what state of the terminus will be retained durably: the state of durable messages, only + existence and configuration of the terminus, or no state at all. + :param ~uamqp.endpoints.ExpiryPolicy expiry_policy: The expiry policy of the Source. + Determines when the expiry timer of a Terminus starts counting down from the timeout value. If the link + is subsequently re-attached before the Terminus is expired, then the count down is aborted. If the + conditions for the terminus-expiry-policy are subsequently re-met, the expiry timer restarts from its + originally configured timeout value. + :param int timeout: Duration that an expiring Source will be retained in seconds. + The Source starts expiring as indicated by the expiry-policy. + :param bool dynamic: Request dynamic creation of a remote Node. + When set to true by the receiving Link endpoint, this field constitutes a request for the sending peer + to dynamically create a Node at the source. In this case the address field MUST NOT be set. When set to + true by the sending Link Endpoint this field indicates creation of a dynamically created Node. In this case + the address field will contain the address of the created Node. The generated address SHOULD include the + Link name and Session-name or client-id in some recognizable form for ease of traceability. + :param dict dynamic_node_properties: Properties of the dynamically created Node. + If the dynamic field is not set to true this field must be left unset. When set by the receiving Link + endpoint, this field contains the desired properties of the Node the receiver wishes to be created. When + set by the sending Link endpoint this field contains the actual properties of the dynamically created node. + :param uamqp.endpoints.DistributionMode distribution_mode: The distribution mode of the Link. + This field MUST be set by the sending end of the Link if the endpoint supports more than one + distribution-mode. This field MAY be set by the receiving end of the Link to indicate a preference when a + Node supports multiple distribution modes. + :param dict filters: A set of predicates to filter the Messages admitted onto the Link. + The receiving endpoint sets its desired filter, the sending endpoint sets the filter actually in place + (including any filters defaulted at the node). The receiving endpoint MUST check that the filter in place + meets its needs and take responsibility for detaching if it does not. + Common filter types, along with the capabilities they are associated with are registered + here: http://www.amqp.org/specification/1.0/filters. + :param ~uamqp.outcomes.DeliveryState default_outcome: Default outcome for unsettled transfers. + Indicates the outcome to be used for transfers that have not reached a terminal state at the receiver + when the transfer is settled, including when the Source is destroyed. The value MUST be a valid + outcome (e.g. Released or Rejected). + :param list(bytes) outcomes: Descriptors for the outcomes that can be chosen on this link. + The values in this field are the symbolic descriptors of the outcomes that can be chosen on this link. + This field MAY be empty, indicating that the default-outcome will be assumed for all message transfers + (if the default-outcome is not set, and no outcomes are provided, then the accepted outcome must be + supported by the source). When present, the values MUST be a symbolic descriptor of a valid outcome, + e.g. "amqp:accepted:list". + :param list(bytes) capabilities: The extension capabilities the sender supports/desires. + See http://www.amqp.org/specification/1.0/source-capabilities. + """ + + +Target = namedtuple( + 'Target', + [ + 'address', + 'durable', + 'expiry_policy', + 'timeout', + 'dynamic', + 'dynamic_node_properties', + 'capabilities' + ]) +Target._code = 0x00000029 # type: ignore # pylint: disable=protected-access +Target.__new__.__defaults__ = (None,) * len(Target._fields) # type: ignore # type: ignore # pylint: disable=protected-access +Target._definition = ( # type: ignore # pylint: disable=protected-access + FIELD("address", AMQPTypes.string, False, None, False), + FIELD("durable", AMQPTypes.uint, False, "none", False), + FIELD("expiry_policy", AMQPTypes.symbol, False, ExpiryPolicy.SessionEnd, False), + FIELD("timeout", AMQPTypes.uint, False, 0, False), + FIELD("dynamic", AMQPTypes.boolean, False, False, False), + FIELD("dynamic_node_properties", FieldDefinition.node_properties, False, None, False), + FIELD("capabilities", AMQPTypes.symbol, False, None, True)) +if _CAN_ADD_DOCSTRING: + Target.__doc__ = """ + For containers which do not implement address resolution (and do not admit spontaneous link attachment + from their partners) but are instead only used as consumers of messages, it is unnecessary to provide spurious + detail on the source. For this purpose it is possible to use a 'minimal' target in which all the + fields are left unset. + + :param str address: The address of the source. + The address of the source MUST NOT be set when sent on a attach frame sent by the receiving Link Endpoint + where the dynamic fiag is set to true (that is where the receiver is requesting the sender to create an + addressable node). The address of the source MUST be set when sent on a attach frame sent by the sending + Link Endpoint where the dynamic fiag is set to true (that is where the sender has created an addressable + node at the request of the receiver and is now communicating the address of that created node). + The generated name of the address SHOULD include the link name and the container-id of the remote container + to allow for ease of identification. + :param ~uamqp.endpoints.TerminusDurability durable: Indicates the durability of the terminus. + Indicates what state of the terminus will be retained durably: the state of durable messages, only + existence and configuration of the terminus, or no state at all. + :param ~uamqp.endpoints.ExpiryPolicy expiry_policy: The expiry policy of the Source. + Determines when the expiry timer of a Terminus starts counting down from the timeout value. If the link + is subsequently re-attached before the Terminus is expired, then the count down is aborted. If the + conditions for the terminus-expiry-policy are subsequently re-met, the expiry timer restarts from its + originally configured timeout value. + :param int timeout: Duration that an expiring Source will be retained in seconds. + The Source starts expiring as indicated by the expiry-policy. + :param bool dynamic: Request dynamic creation of a remote Node. + When set to true by the receiving Link endpoint, this field constitutes a request for the sending peer + to dynamically create a Node at the source. In this case the address field MUST NOT be set. When set to + true by the sending Link Endpoint this field indicates creation of a dynamically created Node. In this case + the address field will contain the address of the created Node. The generated address SHOULD include the + Link name and Session-name or client-id in some recognizable form for ease of traceability. + :param dict dynamic_node_properties: Properties of the dynamically created Node. + If the dynamic field is not set to true this field must be left unset. When set by the receiving Link + endpoint, this field contains the desired properties of the Node the receiver wishes to be created. When + set by the sending Link endpoint this field contains the actual properties of the dynamically created node. + :param list(bytes) capabilities: The extension capabilities the sender supports/desires. + See http://www.amqp.org/specification/1.0/source-capabilities. + """ diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/error.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/error.py new file mode 100644 index 000000000000..91f3393eb8bf --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/error.py @@ -0,0 +1,356 @@ +#------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +#-------------------------------------------------------------------------- + +# TODO: fix mypy errors for _code/_definition/__defaults__ (issue #26500) +from enum import Enum +from collections import namedtuple + +from .constants import SECURE_PORT, FIELD +from .types import AMQPTypes, FieldDefinition + + +class ErrorCondition(bytes, Enum): + # Shared error conditions: + + #: An internal error occurred. Operator intervention may be required to resume normaloperation. + InternalError = b"amqp:internal-error" + #: A peer attempted to work with a remote entity that does not exist. + NotFound = b"amqp:not-found" + #: A peer attempted to work with a remote entity to which it has no access due tosecurity settings. + UnauthorizedAccess = b"amqp:unauthorized-access" + #: Data could not be decoded. + DecodeError = b"amqp:decode-error" + #: A peer exceeded its resource allocation. + ResourceLimitExceeded = b"amqp:resource-limit-exceeded" + #: The peer tried to use a frame in a manner that is inconsistent with the semantics defined in the specification. + NotAllowed = b"amqp:not-allowed" + #: An invalid field was passed in a frame body, and the operation could not proceed. + InvalidField = b"amqp:invalid-field" + #: The peer tried to use functionality that is not implemented in its partner. + NotImplemented = b"amqp:not-implemented" + #: The client attempted to work with a server entity to which it has no access + #: because another client is working with it. + ResourceLocked = b"amqp:resource-locked" + #: The client made a request that was not allowed because some precondition failed. + PreconditionFailed = b"amqp:precondition-failed" + #: A server entity the client is working with has been deleted. + ResourceDeleted = b"amqp:resource-deleted" + #: The peer sent a frame that is not permitted in the current state of the Session. + IllegalState = b"amqp:illegal-state" + #: The peer cannot send a frame because the smallest encoding of the performative with the currently + #: valid values would be too large to fit within a frame of the agreed maximum frame size. + FrameSizeTooSmall = b"amqp:frame-size-too-small" + + # Symbols used to indicate connection error conditions: + + #: An operator intervened to close the Connection for some reason. The client may retry at some later date. + ConnectionCloseForced = b"amqp:connection:forced" + #: A valid frame header cannot be formed from the incoming byte stream. + ConnectionFramingError = b"amqp:connection:framing-error" + #: The container is no longer available on the current connection. The peer should attempt reconnection + #: to the container using the details provided in the info map. + ConnectionRedirect = b"amqp:connection:redirect" + + # Symbols used to indicate session error conditions: + + #: The peer violated incoming window for the session. + SessionWindowViolation = b"amqp:session:window-violation" + #: Input was received for a link that was detached with an error. + SessionErrantLink = b"amqp:session:errant-link" + #: An attach was received using a handle that is already in use for an attached Link. + SessionHandleInUse = b"amqp:session:handle-in-use" + #: A frame (other than attach) was received referencing a handle which + #: is not currently in use of an attached Link. + SessionUnattachedHandle = b"amqp:session:unattached-handle" + + # Symbols used to indicate link error conditions: + + #: An operator intervened to detach for some reason. + LinkDetachForced = b"amqp:link:detach-forced" + #: The peer sent more Message transfers than currently allowed on the link. + LinkTransferLimitExceeded = b"amqp:link:transfer-limit-exceeded" + #: The peer sent a larger message than is supported on the link. + LinkMessageSizeExceeded = b"amqp:link:message-size-exceeded" + #: The address provided cannot be resolved to a terminus at the current container. + LinkRedirect = b"amqp:link:redirect" + #: The link has been attached elsewhere, causing the existing attachment to be forcibly closed. + LinkStolen = b"amqp:link:stolen" + + # Customized symbols used to indicate client error conditions. + # TODO: check whether Client/Unknown/Vendor Error are exposed in EH/SB as users might be depending + # on the code for error handling + ClientError = b"amqp:client-error" + UnknownError = b"amqp:unknown-error" + VendorError = b"amqp:vendor-error" + SocketError = b"amqp:socket-error" + + +class RetryMode(str, Enum): # pylint: disable=enum-must-inherit-case-insensitive-enum-meta + EXPONENTIAL = 'exponential' + FIXED = 'fixed' + + +class RetryPolicy: + + no_retry = [ + ErrorCondition.DecodeError, + ErrorCondition.LinkMessageSizeExceeded, + ErrorCondition.NotFound, + ErrorCondition.NotImplemented, + ErrorCondition.LinkRedirect, + ErrorCondition.NotAllowed, + ErrorCondition.UnauthorizedAccess, + ErrorCondition.LinkStolen, + ErrorCondition.ResourceLimitExceeded, + ErrorCondition.ConnectionRedirect, + ErrorCondition.PreconditionFailed, + ErrorCondition.InvalidField, + ErrorCondition.ResourceDeleted, + ErrorCondition.IllegalState, + ErrorCondition.FrameSizeTooSmall, + ErrorCondition.ConnectionFramingError, + ErrorCondition.SessionUnattachedHandle, + ErrorCondition.SessionHandleInUse, + ErrorCondition.SessionErrantLink, + ErrorCondition.SessionWindowViolation + ] + + def __init__( + self, + **kwargs + ): + """ + keyword int retry_total: + keyword float retry_backoff_factor: + keyword float retry_backoff_max: + keyword RetryMode retry_mode: + keyword list no_retry: + keyword dict custom_retry_policy: + """ + self.total_retries = kwargs.pop('retry_total', 3) + # TODO: A. consider letting retry_backoff_factor be either a float or a callback obj which returns a float + # to give more extensibility on customization of retry backoff time, the callback could take the exception + # as input. + self.backoff_factor = kwargs.pop('retry_backoff_factor', 0.8) + self.backoff_max = kwargs.pop('retry_backoff_max', 120) + self.retry_mode = kwargs.pop('retry_mode', RetryMode.EXPONENTIAL) + self.no_retry.extend(kwargs.get('no_retry', [])) + self.custom_condition_backoff = kwargs.pop("custom_condition_backoff", None) + # TODO: B. As an alternative of option A, we could have a new kwarg serve the goal + + def configure_retries(self, **kwargs): + return { + 'total': kwargs.pop("retry_total", self.total_retries), + 'backoff': kwargs.pop("retry_backoff_factor", self.backoff_factor), + 'max_backoff': kwargs.pop("retry_backoff_max", self.backoff_max), + 'retry_mode': kwargs.pop("retry_mode", self.retry_mode), + 'history': [] + } + + def increment(self, settings, error): # pylint: disable=no-self-use + settings['total'] -= 1 + settings['history'].append(error) + if settings['total'] < 0: + return False + return True + + def is_retryable(self, error): + try: + if error.condition in self.no_retry: + return False + except TypeError: + pass + return True + + def get_backoff_time(self, settings, error): + try: + return self.custom_condition_backoff[error.condition] + except (KeyError, TypeError): + pass + + consecutive_errors_len = len(settings['history']) + if consecutive_errors_len <= 1: + return 0 + + if self.retry_mode == RetryMode.FIXED: + backoff_value = settings['backoff'] + else: + backoff_value = settings['backoff'] * (2 ** (consecutive_errors_len - 1)) + return min(settings['max_backoff'], backoff_value) + + +AMQPError = namedtuple('AMQPError', ['condition', 'description', 'info'], defaults=[None, None]) +AMQPError.__new__.__defaults__ = (None,) * len(AMQPError._fields) # type: ignore +AMQPError._code = 0x0000001d # type: ignore # pylint: disable=protected-access +AMQPError._definition = ( # type: ignore # pylint: disable=protected-access + FIELD('condition', AMQPTypes.symbol, True, None, False), + FIELD('description', AMQPTypes.string, False, None, False), + FIELD('info', FieldDefinition.fields, False, None, False), +) + + +class AMQPException(Exception): + """Base exception for all errors. + + :param bytes condition: The error code. + :keyword str description: A description of the error. + :keyword dict info: A dictionary of additional data associated with the error. + """ + def __init__(self, condition, **kwargs): + self.condition = condition or ErrorCondition.UnknownError + self.description = kwargs.get("description", None) + self.info = kwargs.get("info", None) + self.message = kwargs.get("message", None) + self.inner_error = kwargs.get("error", None) + message = self.message or "Error condition: {}".format( + str(condition) if isinstance(condition, ErrorCondition) else condition.decode() + ) + if self.description: + try: + message += "\n Error Description: {}".format(self.description.decode()) + except (TypeError, AttributeError): + message += "\n Error Description: {}".format(self.description) + super(AMQPException, self).__init__(message) + + +class AMQPDecodeError(AMQPException): + """An error occurred while decoding an incoming frame. + + """ + + +class AMQPConnectionError(AMQPException): + """Details of a Connection-level error. + + """ + + +class AMQPConnectionRedirect(AMQPConnectionError): + """Details of a Connection-level redirect response. + + The container is no longer available on the current connection. + The peer should attempt reconnection to the container using the details provided. + + :param bytes condition: The error code. + :keyword str description: A description of the error. + :keyword dict info: A dictionary of additional data associated with the error. + """ + def __init__(self, condition, description=None, info=None): + self.hostname = info.get(b'hostname', b'').decode('utf-8') + self.network_host = info.get(b'network-host', b'').decode('utf-8') + self.port = int(info.get(b'port', SECURE_PORT)) + super(AMQPConnectionRedirect, self).__init__(condition, description=description, info=info) + + +class AMQPSessionError(AMQPException): + """Details of a Session-level error. + + :param bytes condition: The error code. + :keyword str description: A description of the error. + :keyword dict info: A dictionary of additional data associated with the error. + """ + + +class AMQPLinkError(AMQPException): + """Details of a Link-level error. + + :param bytes condition: The error code. + :keyword str description: A description of the error. + :keyword dict info: A dictionary of additional data associated with the error. + """ + + +class AMQPLinkRedirect(AMQPLinkError): + """Details of a Link-level redirect response. + + The address provided cannot be resolved to a terminus at the current container. + The supplied information may allow the client to locate and attach to the terminus. + + :param bytes condition: The error code. + :keyword str description: A description of the error. + :keyword dict info: A dictionary of additional data associated with the error. + """ + + def __init__(self, condition, description=None, info=None): + self.hostname = info.get(b'hostname', b'').decode('utf-8') + self.network_host = info.get(b'network-host', b'').decode('utf-8') + self.port = int(info.get(b'port', SECURE_PORT)) + self.address = info.get(b'address', b'').decode('utf-8') + super().__init__(condition, description=description, info=info) + + +class AuthenticationException(AMQPException): + """Details of a Authentication error. + + :param bytes condition: The error code. + :keyword str description: A description of the error. + :keyword dict info: A dictionary of additional data associated with the error. + """ + + +class TokenExpired(AuthenticationException): + """Details of a Token expiration error. + + :param bytes condition: The error code. + :keyword str description: A description of the error. + :keyword dict info: A dictionary of additional data associated with the error. + """ + + +class TokenAuthFailure(AuthenticationException): + """Failure to authenticate with token.""" + + def __init__(self, status_code, status_description, **kwargs): + encoding = kwargs.get("encoding", 'utf-8') + self.status_code = status_code + self.status_description = status_description + message = "CBS Token authentication failed.\nStatus code: {}".format(self.status_code) + if self.status_description: + try: + message += "\nDescription: {}".format(self.status_description.decode(encoding)) + except (TypeError, AttributeError): + message += "\nDescription: {}".format(self.status_description) + super(TokenAuthFailure, self).__init__(condition=ErrorCondition.ClientError, message=message) + + +class MessageException(AMQPException): + """Details of a Message error. + + :param bytes condition: The error code. + :keyword str description: A description of the error. + :keyword dict info: A dictionary of additional data associated with the error. + + """ + + +class MessageSendFailed(MessageException): + """Details of a Message send failed error. + + :param bytes condition: The error code. + :keyword str description: A description of the error. + :keyword dict info: A dictionary of additional data associated with the error. + """ + + +class ErrorResponse(object): + """AMQP error object.""" + + def __init__(self, **kwargs): + self.condition = kwargs.get("condition") + self.description = kwargs.get("description") + + info = kwargs.get("info") + error_info = kwargs.get("error_info") + if isinstance(error_info, list) and len(error_info) >= 1: + if isinstance(error_info[0], list) and len(error_info[0]) >= 1: + self.condition = error_info[0][0] + if len(error_info[0]) >= 2: + self.description = error_info[0][1] + if len(error_info[0]) >= 3: + info = error_info[0][2] + + self.info = info + self.error = error_info diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/link.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/link.py new file mode 100644 index 000000000000..ab3523566cb3 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/link.py @@ -0,0 +1,261 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + + +from typing import Optional +import uuid +import logging + +from .endpoints import Source, Target +from .constants import DEFAULT_LINK_CREDIT, SessionState, LinkState, Role, SenderSettleMode, ReceiverSettleMode +from .performatives import AttachFrame, DetachFrame + +from .error import ErrorCondition, AMQPLinkError, AMQPLinkRedirect, AMQPConnectionError + +_LOGGER = logging.getLogger(__name__) + + +class Link(object): # pylint: disable=too-many-instance-attributes + """An AMQP Link. + + This object should not be used directly - instead use one of directional + derivatives: Sender or Receiver. + """ + + def __init__(self, session, handle, name, role, **kwargs): + self.state = LinkState.DETACHED + self.name = name or str(uuid.uuid4()) + self.handle = handle + self.remote_handle = None + self.role = role + source_address = kwargs["source_address"] + target_address = kwargs["target_address"] + self.source = ( + source_address + if isinstance(source_address, Source) + else Source( + address=kwargs["source_address"], + durable=kwargs.get("source_durable"), + expiry_policy=kwargs.get("source_expiry_policy"), + timeout=kwargs.get("source_timeout"), + dynamic=kwargs.get("source_dynamic"), + dynamic_node_properties=kwargs.get("source_dynamic_node_properties"), + distribution_mode=kwargs.get("source_distribution_mode"), + filters=kwargs.get("source_filters"), + default_outcome=kwargs.get("source_default_outcome"), + outcomes=kwargs.get("source_outcomes"), + capabilities=kwargs.get("source_capabilities"), + ) + ) + self.target = ( + target_address + if isinstance(target_address, Target) + else Target( + address=kwargs["target_address"], + durable=kwargs.get("target_durable"), + expiry_policy=kwargs.get("target_expiry_policy"), + timeout=kwargs.get("target_timeout"), + dynamic=kwargs.get("target_dynamic"), + dynamic_node_properties=kwargs.get("target_dynamic_node_properties"), + capabilities=kwargs.get("target_capabilities"), + ) + ) + self.link_credit = kwargs.pop("link_credit", None) or DEFAULT_LINK_CREDIT + self.current_link_credit = self.link_credit + self.send_settle_mode = kwargs.pop("send_settle_mode", SenderSettleMode.Mixed) + self.rcv_settle_mode = kwargs.pop("rcv_settle_mode", ReceiverSettleMode.First) + self.unsettled = kwargs.pop("unsettled", None) + self.incomplete_unsettled = kwargs.pop("incomplete_unsettled", None) + self.initial_delivery_count = kwargs.pop("initial_delivery_count", 0) + self.delivery_count = self.initial_delivery_count + self.received_delivery_id = None + self.max_message_size = kwargs.pop("max_message_size", None) + self.remote_max_message_size = None + self.available = kwargs.pop("available", None) + self.properties = kwargs.pop("properties", None) + self.offered_capabilities = None + self.desired_capabilities = kwargs.pop("desired_capabilities", None) + + self.network_trace = kwargs["network_trace"] + self.network_trace_params = kwargs["network_trace_params"] + self.network_trace_params["amqpLink"] = self.name + self._session = session + self._is_closed = False + self._on_link_state_change = kwargs.get("on_link_state_change") + self._on_attach = kwargs.get("on_attach") + self._error = None + + def __enter__(self): + self.attach() + return self + + def __exit__(self, *args): + self.detach(close=True) + + @classmethod + def from_incoming_frame(cls, session, handle, frame): + # TODO: Assuming we establish all links for now... + # check link_create_from_endpoint in C lib + raise NotImplementedError("Pending") + + def get_state(self): + try: + raise self._error + except TypeError: + pass + return self.state + + def _check_if_closed(self): + if self._is_closed: + try: + raise self._error + except TypeError: + raise AMQPConnectionError(condition=ErrorCondition.InternalError, description="Link already closed.") + + def _set_state(self, new_state): + # type: (LinkState) -> None + """Update the session state.""" + if new_state is None: + return + previous_state = self.state + self.state = new_state + _LOGGER.info("Link state changed: %r -> %r", previous_state, new_state, extra=self.network_trace_params) + try: + self._on_link_state_change(previous_state, new_state) + except TypeError: + pass + except Exception as e: # pylint: disable=broad-except + _LOGGER.error("Link state change callback failed: '%r'", e, extra=self.network_trace_params) + + def _on_session_state_change(self): + if self._session.state == SessionState.MAPPED: + if not self._is_closed and self.state == LinkState.DETACHED: + self._outgoing_attach() + self._set_state(LinkState.ATTACH_SENT) + elif self._session.state == SessionState.DISCARDING: + self._set_state(LinkState.DETACHED) + + def _outgoing_attach(self): + self.delivery_count = self.initial_delivery_count + attach_frame = AttachFrame( + name=self.name, + handle=self.handle, + role=self.role, + send_settle_mode=self.send_settle_mode, + rcv_settle_mode=self.rcv_settle_mode, + source=self.source, + target=self.target, + unsettled=self.unsettled, + incomplete_unsettled=self.incomplete_unsettled, + initial_delivery_count=self.initial_delivery_count if self.role == Role.Sender else None, + max_message_size=self.max_message_size, + offered_capabilities=self.offered_capabilities if self.state == LinkState.ATTACH_RCVD else None, + desired_capabilities=self.desired_capabilities if self.state == LinkState.DETACHED else None, + properties=self.properties, + ) + if self.network_trace: + _LOGGER.debug("-> %r", attach_frame, extra=self.network_trace_params) + self._session._outgoing_attach(attach_frame) # pylint: disable=protected-access + + def _incoming_attach(self, frame): + if self.network_trace: + _LOGGER.debug("<- %r", AttachFrame(*frame), extra=self.network_trace_params) + if self._is_closed: + raise ValueError("Invalid link") + if not frame[5] or not frame[6]: + _LOGGER.info("Cannot get source or target. Detaching link", extra=self.network_trace_params) + self._set_state(LinkState.DETACHED) + raise ValueError("Invalid link") + self.remote_handle = frame[1] # handle + self.remote_max_message_size = frame[10] # max_message_size + self.offered_capabilities = frame[11] # offered_capabilities + if self.properties: + self.properties.update(frame[13]) # properties + else: + self.properties = frame[13] + if self.state == LinkState.DETACHED: + self._set_state(LinkState.ATTACH_RCVD) + elif self.state == LinkState.ATTACH_SENT: + self._set_state(LinkState.ATTACHED) + if self._on_attach: + try: + if frame[5]: + frame[5] = Source(*frame[5]) + if frame[6]: + frame[6] = Target(*frame[6]) + self._on_attach(AttachFrame(*frame)) + except Exception as e: # pylint: disable=broad-except + _LOGGER.warning("Callback for link attach raised error: %r", e, extra=self.network_trace_params) + + def _outgoing_flow(self, **kwargs): + flow_frame = { + "handle": self.handle, + "delivery_count": self.delivery_count, + "link_credit": self.current_link_credit, + "available": kwargs.get("available"), + "drain": kwargs.get("drain"), + "echo": kwargs.get("echo"), + "properties": kwargs.get("properties"), + } + self._session._outgoing_flow(flow_frame) # pylint: disable=protected-access + + def _incoming_flow(self, frame): + pass + + def _incoming_disposition(self, frame): + pass + + def _outgoing_detach(self, close=False, error=None): + detach_frame = DetachFrame(handle=self.handle, closed=close, error=error) + if self.network_trace: + _LOGGER.debug("-> %r", detach_frame, extra=self.network_trace_params) + self._session._outgoing_detach(detach_frame) # pylint: disable=protected-access + if close: + self._is_closed = True + + def _incoming_detach(self, frame): + if self.network_trace: + _LOGGER.debug("<- %r", DetachFrame(*frame), extra=self.network_trace_params) + if self.state == LinkState.ATTACHED: + self._outgoing_detach(close=frame[1]) # closed + elif frame[1] and not self._is_closed and self.state in [LinkState.ATTACH_SENT, LinkState.ATTACH_RCVD]: + # Received a closing detach after we sent a non-closing detach. + # In this case, we MUST signal that we closed by reattaching and then sending a closing detach. + self._outgoing_attach() + self._outgoing_detach(close=True) + # TODO: on_detach_hook + if frame[2]: # error + # frame[2][0] is condition, frame[2][1] is description, frame[2][2] is info + error_cls = AMQPLinkRedirect if frame[2][0] == ErrorCondition.LinkRedirect else AMQPLinkError + self._error = error_cls(condition=frame[2][0], description=frame[2][1], info=frame[2][2]) + self._set_state(LinkState.ERROR) + else: + self._set_state(LinkState.DETACHED) + + def attach(self): + if self._is_closed: + raise ValueError("Link already closed.") + self._outgoing_attach() + self._set_state(LinkState.ATTACH_SENT) + + def detach(self, close=False, error=None): + if self.state in (LinkState.DETACHED, LinkState.ERROR): + return + try: + self._check_if_closed() + if self.state in [LinkState.ATTACH_SENT, LinkState.ATTACH_RCVD]: + self._outgoing_detach(close=close, error=error) + self._set_state(LinkState.DETACHED) + elif self.state == LinkState.ATTACHED: + self._outgoing_detach(close=close, error=error) + self._set_state(LinkState.DETACH_SENT) + except Exception as exc: # pylint: disable=broad-except + _LOGGER.info("An error occurred when detaching the link: %r", exc, extra=self.network_trace_params) + self._set_state(LinkState.DETACHED) + + def flow(self, *, link_credit: Optional[int] = None, **kwargs) -> None: + self.current_link_credit = link_credit if link_credit is not None else self.link_credit + self._outgoing_flow(**kwargs) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/management_link.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/management_link.py new file mode 100644 index 000000000000..c5b1e6c0aa19 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/management_link.py @@ -0,0 +1,262 @@ +#------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +#-------------------------------------------------------------------------- + +import time +import logging +from functools import partial +from collections import namedtuple + +from .sender import SenderLink +from .receiver import ReceiverLink +from .constants import ( + ManagementLinkState, + LinkState, + SenderSettleMode, + ReceiverSettleMode, + ManagementExecuteOperationResult, + ManagementOpenResult, + SEND_DISPOSITION_REJECT, + MessageDeliveryState, + LinkDeliverySettleReason +) +from .error import AMQPException, ErrorCondition +from .message import Properties, _MessageDelivery + +_LOGGER = logging.getLogger(__name__) + +PendingManagementOperation = namedtuple('PendingManagementOperation', ['message', 'on_execute_operation_complete']) + + +class ManagementLink(object): # pylint:disable=too-many-instance-attributes + """ + # TODO: Fill in docstring + """ + def __init__(self, session, endpoint, **kwargs): + self.next_message_id = 0 + self.state = ManagementLinkState.IDLE + self._pending_operations = [] + self._session = session + self._network_trace_params = kwargs.get('network_trace_params') + self._request_link: SenderLink = session.create_sender_link( + endpoint, + source_address=endpoint, + on_link_state_change=self._on_sender_state_change, + send_settle_mode=SenderSettleMode.Unsettled, + rcv_settle_mode=ReceiverSettleMode.First, + network_trace=kwargs.get("network_trace", False) + ) + self._response_link: ReceiverLink = session.create_receiver_link( + endpoint, + target_address=endpoint, + on_link_state_change=self._on_receiver_state_change, + on_transfer=self._on_message_received, + send_settle_mode=SenderSettleMode.Unsettled, + rcv_settle_mode=ReceiverSettleMode.First, + network_trace=kwargs.get("network_trace", False) + ) + self._on_amqp_management_error = kwargs.get('on_amqp_management_error') + self._on_amqp_management_open_complete = kwargs.get('on_amqp_management_open_complete') + + self._status_code_field = kwargs.get('status_code_field', b'statusCode') + self._status_description_field = kwargs.get('status_description_field', b'statusDescription') + + self._sender_connected = False + self._receiver_connected = False + + def __enter__(self): + self.open() + return self + + def __exit__(self, *args): + self.close() + + def _on_sender_state_change(self, previous_state, new_state): + _LOGGER.info( + "Management link sender state changed: %r -> %r", + previous_state, + new_state, + extra=self._network_trace_params + ) + if new_state == previous_state: + return + if self.state == ManagementLinkState.OPENING: + if new_state == LinkState.ATTACHED: + self._sender_connected = True + if self._receiver_connected: + self.state = ManagementLinkState.OPEN + self._on_amqp_management_open_complete(ManagementOpenResult.OK) + elif new_state in [LinkState.DETACHED, LinkState.DETACH_SENT, LinkState.DETACH_RCVD, LinkState.ERROR]: + self.state = ManagementLinkState.IDLE + self._on_amqp_management_open_complete(ManagementOpenResult.ERROR) + elif self.state == ManagementLinkState.OPEN: + if new_state is not LinkState.ATTACHED: + self.state = ManagementLinkState.ERROR + self._on_amqp_management_error() + elif self.state == ManagementLinkState.CLOSING: + if new_state not in [LinkState.DETACHED, LinkState.DETACH_SENT, LinkState.DETACH_RCVD]: + self.state = ManagementLinkState.ERROR + self._on_amqp_management_error() + elif self.state == ManagementLinkState.ERROR: + # All state transitions shall be ignored. + return + + def _on_receiver_state_change(self, previous_state, new_state): + _LOGGER.info( + "Management link receiver state changed: %r -> %r", + previous_state, + new_state, + extra=self._network_trace_params + ) + if new_state == previous_state: + return + if self.state == ManagementLinkState.OPENING: + if new_state == LinkState.ATTACHED: + self._receiver_connected = True + if self._sender_connected: + self.state = ManagementLinkState.OPEN + self._on_amqp_management_open_complete(ManagementOpenResult.OK) + elif new_state in [LinkState.DETACHED, LinkState.DETACH_SENT, LinkState.DETACH_RCVD, LinkState.ERROR]: + self.state = ManagementLinkState.IDLE + self._on_amqp_management_open_complete(ManagementOpenResult.ERROR) + elif self.state == ManagementLinkState.OPEN: + if new_state is not LinkState.ATTACHED: + self.state = ManagementLinkState.ERROR + self._on_amqp_management_error() + elif self.state == ManagementLinkState.CLOSING: + if new_state not in [LinkState.DETACHED, LinkState.DETACH_SENT, LinkState.DETACH_RCVD]: + self.state = ManagementLinkState.ERROR + self._on_amqp_management_error() + elif self.state == ManagementLinkState.ERROR: + # All state transitions shall be ignored. + return + + def _on_message_received(self, _, message): + message_properties = message.properties + correlation_id = message_properties[5] + response_detail = message.application_properties + + status_code = response_detail.get(self._status_code_field) + status_description = response_detail.get(self._status_description_field) + + to_remove_operation = None + for operation in self._pending_operations: + if operation.message.properties.message_id == correlation_id: + to_remove_operation = operation + break + if to_remove_operation: + mgmt_result = ManagementExecuteOperationResult.OK \ + if 200 <= status_code <= 299 else ManagementExecuteOperationResult.FAILED_BAD_STATUS + to_remove_operation.on_execute_operation_complete( + mgmt_result, + status_code, + status_description, + message, + response_detail.get(b'error-condition') + ) + self._pending_operations.remove(to_remove_operation) + + def _on_send_complete(self, message_delivery, reason, state): # todo: reason is never used, should check spec + if reason == LinkDeliverySettleReason.DISPOSITION_RECEIVED and SEND_DISPOSITION_REJECT in state: + # sample reject state: {'rejected': [[b'amqp:not-allowed', b"Invalid command 'RE1AD'.", None]]} + to_remove_operation = None + for operation in self._pending_operations: + if message_delivery.message == operation.message: + to_remove_operation = operation + break + self._pending_operations.remove(to_remove_operation) + # TODO: better error handling + # AMQPException is too general? to be more specific: MessageReject(Error) or AMQPManagementError? + # or should there an error mapping which maps the condition to the error type + to_remove_operation.on_execute_operation_complete( # The callback is defined in management_operation.py + ManagementExecuteOperationResult.ERROR, + None, + None, + message_delivery.message, + error=AMQPException( + condition=state[SEND_DISPOSITION_REJECT][0][0], # 0 is error condition + description=state[SEND_DISPOSITION_REJECT][0][1], # 1 is error description + info=state[SEND_DISPOSITION_REJECT][0][2], # 2 is error info + ) + ) + + def open(self): + if self.state != ManagementLinkState.IDLE: + raise ValueError("Management links are already open or opening.") + self.state = ManagementLinkState.OPENING + self._response_link.attach() + self._request_link.attach() + + def execute_operation( + self, + message, + on_execute_operation_complete, + **kwargs + ): + """Execute a request and wait on a response. + + :param message: The message to send in the management request. + :type message: ~uamqp.message.Message + :param on_execute_operation_complete: Callback to be called when the operation is complete. + The following value will be passed to the callback: operation_id, operation_result, status_code, + status_description, raw_message and error. + :type on_execute_operation_complete: Callable[[str, str, int, str, ~uamqp.message.Message, Exception], None] + :keyword operation: The type of operation to be performed. This value will + be service-specific, but common values include READ, CREATE and UPDATE. + This value will be added as an application property on the message. + :paramtype operation: bytes or str + :keyword type: The type on which to carry out the operation. This will + be specific to the entities of the service. This value will be added as + an application property on the message. + :paramtype type: bytes or str + :keyword str locales: A list of locales that the sending peer permits for incoming + informational text in response messages. + :keyword float timeout: Provide an optional timeout in seconds within which a response + to the management request must be received. + :rtype: None + """ + timeout = kwargs.get("timeout") + message.application_properties["operation"] = kwargs.get("operation") + message.application_properties["type"] = kwargs.get("type") + if "locales" in kwargs: + message.application_properties["locales"] = kwargs.get("locales") + try: + # TODO: namedtuple is immutable, which may push us to re-think about the namedtuple approach for Message + new_properties = message.properties._replace(message_id=self.next_message_id) + except AttributeError: + new_properties = Properties(message_id=self.next_message_id) + message = message._replace(properties=new_properties) + expire_time = (time.time() + timeout) if timeout else None + message_delivery = _MessageDelivery( + message, + MessageDeliveryState.WaitingToBeSent, + expire_time + ) + + on_send_complete = partial(self._on_send_complete, message_delivery) + + self._request_link.send_transfer( + message, + on_send_complete=on_send_complete, + timeout=timeout + ) + self.next_message_id += 1 + self._pending_operations.append(PendingManagementOperation(message, on_execute_operation_complete)) + + def close(self): + if self.state != ManagementLinkState.IDLE: + self.state = ManagementLinkState.CLOSING + self._response_link.detach(close=True) + self._request_link.detach(close=True) + for pending_operation in self._pending_operations: + pending_operation.on_execute_operation_complete( + ManagementExecuteOperationResult.LINK_CLOSED, + None, + None, + pending_operation.message, + AMQPException(condition=ErrorCondition.ClientError, description="Management link already closed.") + ) + self._pending_operations = [] + self.state = ManagementLinkState.IDLE diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/management_operation.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/management_operation.py new file mode 100644 index 000000000000..475c3424a897 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/management_operation.py @@ -0,0 +1,140 @@ +#------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +#-------------------------------------------------------------------------- +import logging +import uuid +import time +from functools import partial + +from .management_link import ManagementLink +from .error import ( + AMQPLinkError, + ErrorCondition +) + +from .constants import ( + ManagementOpenResult, + ManagementExecuteOperationResult +) + +_LOGGER = logging.getLogger(__name__) + + +class ManagementOperation(object): + def __init__(self, session, endpoint='$management', **kwargs): + self._mgmt_link_open_status = None + + self._session = session + self._connection = self._session._connection + self._network_trace_params = { + "amqpConnection": self._session._connection._container_id, + "amqpSession": self._session.name, + "amqpLink": None + } + self._mgmt_link = self._session.create_request_response_link_pair( + endpoint=endpoint, + on_amqp_management_open_complete=self._on_amqp_management_open_complete, + on_amqp_management_error=self._on_amqp_management_error, + **kwargs + ) # type: ManagementLink + self._responses = {} + self._mgmt_error = None + + def _on_amqp_management_open_complete(self, result): + """Callback run when the send/receive links are open and ready + to process messages. + + :param result: Whether the link opening was successful. + :type result: int + """ + self._mgmt_link_open_status = result + + def _on_amqp_management_error(self): + """Callback run if an error occurs in the send/receive links.""" + # TODO: This probably shouldn't be ValueError + self._mgmt_error = ValueError("Management Operation error occurred.") + + def _on_execute_operation_complete( + self, + operation_id, + operation_result, + status_code, + status_description, + raw_message, + error=None + ): + _LOGGER.debug( + "Management operation completed, id: %r; result: %r; code: %r; description: %r, error: %r", + operation_id, + operation_result, + status_code, + status_description, + error, + extra=self._network_trace_params + ) + + if operation_result in\ + (ManagementExecuteOperationResult.ERROR, ManagementExecuteOperationResult.LINK_CLOSED): + self._mgmt_error = error + _LOGGER.error( + "Failed to complete management operation due to error: %r.", + error, + extra=self._network_trace_params + ) + else: + self._responses[operation_id] = (status_code, status_description, raw_message) + + def execute(self, message, operation=None, operation_type=None, timeout=0): + start_time = time.time() + operation_id = str(uuid.uuid4()) + self._responses[operation_id] = None + self._mgmt_error = None + + self._mgmt_link.execute_operation( + message, + partial(self._on_execute_operation_complete, operation_id), + timeout=timeout, + operation=operation, + type=operation_type + ) + + while not self._responses[operation_id] and not self._mgmt_error: + if timeout and timeout > 0: + now = time.time() + if (now - start_time) >= timeout: + raise TimeoutError("Failed to receive mgmt response in {}ms".format(timeout)) + self._connection.listen() + + if self._mgmt_error: + self._responses.pop(operation_id) + raise self._mgmt_error # pylint: disable=raising-bad-type + + response = self._responses.pop(operation_id) + return response + + def open(self): + self._mgmt_link_open_status = ManagementOpenResult.OPENING + self._mgmt_link.open() + + def ready(self): + try: + raise self._mgmt_error # pylint: disable=raising-bad-type + except TypeError: + pass + + if self._mgmt_link_open_status == ManagementOpenResult.OPENING: + return False + if self._mgmt_link_open_status == ManagementOpenResult.OK: + return True + # ManagementOpenResult.ERROR or CANCELLED + # TODO: update below with correct status code + info + raise AMQPLinkError( + condition=ErrorCondition.ClientError, + description="Failed to open mgmt link, management link status: {}".format(self._mgmt_link_open_status), + info=None + ) + + def close(self): + self._mgmt_link.close() diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/message.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/message.py new file mode 100644 index 000000000000..c4bc6b0e1d19 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/message.py @@ -0,0 +1,268 @@ +#------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +#-------------------------------------------------------------------------- + +# TODO: fix mypy errors for _code/_definition/__defaults__ (issue #26500) +from collections import namedtuple + +from .types import AMQPTypes, FieldDefinition +from .constants import FIELD, MessageDeliveryState +from .performatives import _CAN_ADD_DOCSTRING + + +Header = namedtuple( + 'Header', + [ + 'durable', + 'priority', + 'ttl', + 'first_acquirer', + 'delivery_count' + ]) +Header._code = 0x00000070 # type: ignore # pylint:disable=protected-access +Header.__new__.__defaults__ = (None,) * len(Header._fields) # type: ignore +Header._definition = ( # type: ignore # pylint:disable=protected-access + FIELD("durable", AMQPTypes.boolean, False, None, False), + FIELD("priority", AMQPTypes.ubyte, False, None, False), + FIELD("ttl", AMQPTypes.uint, False, None, False), + FIELD("first_acquirer", AMQPTypes.boolean, False, None, False), + FIELD("delivery_count", AMQPTypes.uint, False, None, False)) +if _CAN_ADD_DOCSTRING: + Header.__doc__ = """ + Transport headers for a Message. + + The header section carries standard delivery details about the transfer of a Message through the AMQP + network. If the header section is omitted the receiver MUST assume the appropriate default values for + the fields within the header unless other target or node specific defaults have otherwise been set. + + :param bool durable: Specify durability requirements. + Durable Messages MUST NOT be lost even if an intermediary is unexpectedly terminated and restarted. + A target which is not capable of fulfilling this guarantee MUST NOT accept messages where the durable + header is set to true: if the source allows the rejected outcome then the message should be rejected + with the precondition-failed error, otherwise the link must be detached by the receiver with the same error. + :param int priority: Relative Message priority. + This field contains the relative Message priority. Higher numbers indicate higher priority Messages. + Messages with higher priorities MAY be delivered before those with lower priorities. An AMQP intermediary + implementing distinct priority levels MUST do so in the following manner: + + - If n distince priorities are implemented and n is less than 10 - priorities 0 to (5 - ceiling(n/2)) + MUST be treated equivalently and MUST be the lowest effective priority. The priorities (4 + fioor(n/2)) + and above MUST be treated equivalently and MUST be the highest effective priority. The priorities + (5 ceiling(n/2)) to (4 + fioor(n/2)) inclusive MUST be treated as distinct priorities. + - If n distinct priorities are implemented and n is 10 or greater - priorities 0 to (n - 1) MUST be + distinct, and priorities n and above MUST be equivalent to priority (n - 1). Thus, for example, if 2 + distinct priorities are implemented, then levels 0 to 4 are equivalent, and levels 5 to 9 are equivalent + and levels 4 and 5 are distinct. If 3 distinct priorities are implements the 0 to 3 are equivalent, + 5 to 9 are equivalent and 3, 4 and 5 are distinct. This scheme ensures that if two priorities are distinct + for a server which implements m separate priority levels they are also distinct for a server which + implements n different priority levels where n > m. + + :param int ttl: Time to live in ms. + Duration in milliseconds for which the Message should be considered 'live'. If this is set then a message + expiration time will be computed based on the time of arrival at an intermediary. Messages that live longer + than their expiration time will be discarded (or dead lettered). When a message is transmitted by an + intermediary that was received with a ttl, the transmitted message's header should contain a ttl that is + computed as the difference between the current time and the formerly computed message expiration + time, i.e. the reduced ttl, so that messages will eventually die if they end up in a delivery loop. + :param bool first_acquirer: If this value is true, then this message has not been acquired by any other Link. + If this value is false, then this message may have previously been acquired by another Link or Links. + :param int delivery_count: The number of prior unsuccessful delivery attempts. + The number of unsuccessful previous attempts to deliver this message. If this value is non-zero it may + be taken as an indication that the delivery may be a duplicate. On first delivery, the value is zero. + It is incremented upon an outcome being settled at the sender, according to rules defined for each outcome. + """ + + +Properties = namedtuple( + 'Properties', + [ + 'message_id', + 'user_id', + 'to', + 'subject', + 'reply_to', + 'correlation_id', + 'content_type', + 'content_encoding', + 'absolute_expiry_time', + 'creation_time', + 'group_id', + 'group_sequence', + 'reply_to_group_id' + ]) +Properties._code = 0x00000073 # type: ignore # pylint:disable=protected-access +Properties.__new__.__defaults__ = (None,) * len(Properties._fields) # type: ignore +Properties._definition = ( # type: ignore # pylint:disable=protected-access + FIELD("message_id", FieldDefinition.message_id, False, None, False), + FIELD("user_id", AMQPTypes.binary, False, None, False), + FIELD("to", AMQPTypes.string, False, None, False), + FIELD("subject", AMQPTypes.string, False, None, False), + FIELD("reply_to", AMQPTypes.string, False, None, False), + FIELD("correlation_id", FieldDefinition.message_id, False, None, False), + FIELD("content_type", AMQPTypes.symbol, False, None, False), + FIELD("content_encoding", AMQPTypes.symbol, False, None, False), + FIELD("absolute_expiry_time", AMQPTypes.timestamp, False, None, False), + FIELD("creation_time", AMQPTypes.timestamp, False, None, False), + FIELD("group_id", AMQPTypes.string, False, None, False), + FIELD("group_sequence", AMQPTypes.uint, False, None, False), + FIELD("reply_to_group_id", AMQPTypes.string, False, None, False)) +if _CAN_ADD_DOCSTRING: + Properties.__doc__ = """ + Immutable properties of the Message. + + The properties section is used for a defined set of standard properties of the message. The properties + section is part of the bare message and thus must, if retransmitted by an intermediary, remain completely + unaltered. + + :param message_id: Application Message identifier. + Message-id is an optional property which uniquely identifies a Message within the Message system. + The Message producer is usually responsible for setting the message-id in such a way that it is assured + to be globally unique. A broker MAY discard a Message as a duplicate if the value of the message-id + matches that of a previously received Message sent to the same Node. + :param bytes user_id: Creating user id. + The identity of the user responsible for producing the Message. The client sets this value, and it MAY + be authenticated by intermediaries. + :param to: The address of the Node the Message is destined for. + The to field identifies the Node that is the intended destination of the Message. On any given transfer + this may not be the Node at the receiving end of the Link. + :param str subject: The subject of the message. + A common field for summary information about the Message content and purpose. + :param reply_to: The Node to send replies to. + The address of the Node to send replies to. + :param correlation_id: Application correlation identifier. + This is a client-specific id that may be used to mark or identify Messages between clients. + :param bytes content_type: MIME content type. + The RFC-2046 MIME type for the Message's application-data section (body). As per RFC-2046 this may contain + a charset parameter defining the character encoding used: e.g. 'text/plain; charset="utf-8"'. + For clarity, the correct MIME type for a truly opaque binary section is application/octet-stream. + When using an application-data section with a section code other than data, contenttype, if set, SHOULD + be set to a MIME type of message/x-amqp+?, where '?' is either data, map or list. + :param bytes content_encoding: MIME content type. + The Content-Encoding property is used as a modifier to the content-type. When present, its value indicates + what additional content encodings have been applied to the application-data, and thus what decoding + mechanisms must be applied in order to obtain the media-type referenced by the content-type header field. + Content-Encoding is primarily used to allow a document to be compressed without losing the identity of + its underlying content type. Content Encodings are to be interpreted as per Section 3.5 of RFC 2616. + Valid Content Encodings are registered at IANA as "Hypertext Transfer Protocol (HTTP) Parameters" + (http://www.iana.org/assignments/http-parameters/httpparameters.xml). Content-Encoding MUST not be set when + the application-data section is other than data. Implementations MUST NOT use the identity encoding. + Instead, implementations should not set this property. Implementations SHOULD NOT use the compress + encoding, except as to remain compatible with messages originally sent with other protocols, + e.g. HTTP or SMTP. Implementations SHOULD NOT specify multiple content encoding values except as to be + compatible with messages originally sent with other protocols, e.g. HTTP or SMTP. + :param datetime absolute_expiry_time: The time when this message is considered expired. + An absolute time when this message is considered to be expired. + :param datetime creation_time: The time when this message was created. + An absolute time when this message was created. + :param str group_id: The group this message belongs to. + Identifies the group the message belongs to. + :param int group_sequence: The sequence-no of this message within its group. + The relative position of this message within its group. + :param str reply_to_group_id: The group the reply message belongs to. + This is a client-specific id that is used so that client can send replies to this message to a specific group. + """ + +# TODO: should be a class, namedtuple or dataclass, immutability vs performance, need to collect performance data +Message = namedtuple( + 'Message', + [ + 'header', + 'delivery_annotations', + 'message_annotations', + 'properties', + 'application_properties', + 'data', + 'sequence', + 'value', + 'footer', + ]) +Message.__new__.__defaults__ = (None,) * len(Message._fields) # type: ignore +Message._code = 0 # type: ignore # pylint:disable=protected-access +Message._definition = ( # type: ignore # pylint:disable=protected-access + (0x00000070, FIELD("header", Header, False, None, False)), + (0x00000071, FIELD("delivery_annotations", FieldDefinition.annotations, False, None, False)), + (0x00000072, FIELD("message_annotations", FieldDefinition.annotations, False, None, False)), + (0x00000073, FIELD("properties", Properties, False, None, False)), + (0x00000074, FIELD("application_properties", AMQPTypes.map, False, None, False)), + (0x00000075, FIELD("data", AMQPTypes.binary, False, None, True)), + (0x00000076, FIELD("sequence", AMQPTypes.list, False, None, False)), + (0x00000077, FIELD("value", None, False, None, False)), + (0x00000078, FIELD("footer", FieldDefinition.annotations, False, None, False))) +if _CAN_ADD_DOCSTRING: + Message.__doc__ = """ + An annotated message consists of the bare message plus sections for annotation at the head and tail + of the bare message. + + There are two classes of annotations: annotations that travel with the message indefinitely, and + annotations that are consumed by the next node. + The exact structure of a message, together with its encoding, is defined by the message format. This document + defines the structure and semantics of message format 0 (MESSAGE-FORMAT). Altogether a message consists of the + following sections: + + - Zero or one header. + - Zero or one delivery-annotations. + - Zero or one message-annotations. + - Zero or one properties. + - Zero or one application-properties. + - The body consists of either: one or more data sections, one or more amqp-sequence sections, + or a single amqp-value section. + - Zero or one footer. + + :param ~uamqp.message.Header header: Transport headers for a Message. + The header section carries standard delivery details about the transfer of a Message through the AMQP + network. If the header section is omitted the receiver MUST assume the appropriate default values for + the fields within the header unless other target or node specific defaults have otherwise been set. + :param dict delivery_annotations: The delivery-annotations section is used for delivery-specific non-standard + properties at the head of the message. Delivery annotations convey information from the sending peer to + the receiving peer. If the recipient does not understand the annotation it cannot be acted upon and its + effects (such as any implied propagation) cannot be acted upon. Annotations may be specific to one + implementation, or common to multiple implementations. The capabilities negotiated on link attach and on + the source and target should be used to establish which annotations a peer supports. A registry of defined + annotations and their meanings can be found here: http://www.amqp.org/specification/1.0/delivery-annotations. + If the delivery-annotations section is omitted, it is equivalent to a delivery-annotations section + containing an empty map of annotations. + :param dict message_annotations: The message-annotations section is used for properties of the message which + are aimed at the infrastructure and should be propagated across every delivery step. Message annotations + convey information about the message. Intermediaries MUST propagate the annotations unless the annotations + are explicitly augmented or modified (e.g. by the use of the modified outcome). + The capabilities negotiated on link attach and on the source and target may be used to establish which + annotations a peer understands, however it a network of AMQP intermediaries it may not be possible to know + if every intermediary will understand the annotation. Note that for some annotation it may not be necessary + for the intermediary to understand their purpose - they may be being used purely as an attribute which can be + filtered on. A registry of defined annotations and their meanings can be found here: + http://www.amqp.org/specification/1.0/message-annotations. If the message-annotations section is omitted, + it is equivalent to a message-annotations section containing an empty map of annotations. + :param ~uamqp.message.Properties: Immutable properties of the Message. + The properties section is used for a defined set of standard properties of the message. The properties + section is part of the bare message and thus must, if retransmitted by an intermediary, remain completely + unaltered. + :param dict application_properties: The application-properties section is a part of the bare message used + for structured application data. Intermediaries may use the data within this structure for the purposes + of filtering or routing. The keys of this map are restricted to be of type string (which excludes the + possibility of a null key) and the values are restricted to be of simple types only (that is excluding + map, list, and array types). + :param list(bytes) data_body: A data section contains opaque binary data. + :param list sequence_body: A sequence section contains an arbitrary number of structured data elements. + :param value_body: An amqp-value section contains a single AMQP value. + :param dict footer: Transport footers for a Message. + The footer section is used for details about the message or delivery which can only be calculated or + evaluated once the whole bare message has been constructed or seen (for example message hashes, HMACs, + signatures and encryption details). A registry of defined footers and their meanings can be found + here: http://www.amqp.org/specification/1.0/footer. + """ + + +class BatchMessage(Message): + _code = 0x80013700 + + +class _MessageDelivery: + def __init__(self, message, state=MessageDeliveryState.WaitingToBeSent, expiry=None): + self.message = message + self.state = state + self.expiry = expiry + self.reason = None + self.delivery = None + self.error = None diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/outcomes.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/outcomes.py new file mode 100644 index 000000000000..64c5d09c7f66 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/outcomes.py @@ -0,0 +1,160 @@ +#------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +#-------------------------------------------------------------------------- + +# The Messaging layer defines a concrete set of delivery states which can be used (via the disposition frame) +# to indicate the state of the message at the receiver. + +# Delivery states may be either terminal or non-terminal. Once a delivery reaches a terminal delivery-state, +# the state for that delivery will no longer change. A terminal delivery-state is referred to as an outcome. + +# The following outcomes are formally defined by the messaging layer to indicate the result of processing at the +# receiver: + +# - accepted: indicates successful processing at the receiver +# - rejected: indicates an invalid and unprocessable message +# - released: indicates that the message was not (and will not be) processed +# - modified: indicates that the message was modified, but not processed + +# The following non-terminal delivery-state is formally defined by the messaging layer for use during link +# recovery to allow the sender to resume the transfer of a large message without retransmitting all the +# message data: + +# - received: indicates partial message data seen by the receiver as well as the starting point for a +# resumed transfer + +# TODO: fix mypy errors for _code/_definition/__defaults__ (issue #26500) +from collections import namedtuple + +from .types import AMQPTypes, FieldDefinition, ObjDefinition +from .constants import FIELD +from .performatives import _CAN_ADD_DOCSTRING + + +Received = namedtuple('Received', ['section_number', 'section_offset']) +Received._code = 0x00000023 # type: ignore # pylint:disable=protected-access +Received._definition = ( # type: ignore # pylint:disable=protected-access + FIELD("section_number", AMQPTypes.uint, True, None, False), + FIELD("section_offset", AMQPTypes.ulong, True, None, False)) +if _CAN_ADD_DOCSTRING: + Received.__doc__ = """ + At the target the received state indicates the furthest point in the payload of the message + which the target will not need to have resent if the link is resumed. At the source the received state represents + the earliest point in the payload which the Sender is able to resume transferring at in the case of link + resumption. When resuming a delivery, if this state is set on the first transfer performative it indicates + the offset in the payload at which the first resumed delivery is starting. The Sender MUST NOT send the + received state on transfer or disposition performatives except on the first transfer performative on a + resumed delivery. + + :param int section_number: + When sent by the Sender this indicates the first section of the message (with sectionnumber 0 being the + first section) for which data can be resent. Data from sections prior to the given section cannot be + retransmitted for this delivery. When sent by the Receiver this indicates the first section of the message + for which all data may not yet have been received. + :param int section_offset: + When sent by the Sender this indicates the first byte of the encoded section data of the section given by + section-number for which data can be resent (with section-offset 0 being the first byte). Bytes from the + same section prior to the given offset section cannot be retransmitted for this delivery. When sent by the + Receiver this indicates the first byte of the given section which has not yet been received. Note that if + a receiver has received all of section number X (which contains N bytes of data), but none of section + number X + 1, then it may indicate this by sending either Received(section-number=X, section-offset=N) or + Received(section-number=X+1, section-offset=0). The state Received(sectionnumber=0, section-offset=0) + indicates that no message data at all has been transferred. + """ + + +Accepted = namedtuple('Accepted', []) +Accepted._code = 0x00000024 # type: ignore # pylint:disable=protected-access +Accepted._definition = () # type: ignore # pylint:disable=protected-access +if _CAN_ADD_DOCSTRING: + Accepted.__doc__ = """ + The accepted outcome. + + At the source the accepted state means that the message has been retired from the node, and transfer of + payload data will not be able to be resumed if the link becomes suspended. A delivery may become accepted at + the source even before all transfer frames have been sent, this does not imply that the remaining transfers + for the delivery will not be sent - only the aborted fiag on the transfer performative can be used to indicate + a premature termination of the transfer. At the target, the accepted outcome is used to indicate that an + incoming Message has been successfully processed, and that the receiver of the Message is expecting the sender + to transition the delivery to the accepted state at the source. The accepted outcome does not increment the + delivery-count in the header of the accepted Message. + """ + + +Rejected = namedtuple('Rejected', ['error']) +Rejected.__new__.__defaults__ = (None,) * len(Rejected._fields) # type: ignore +Rejected._code = 0x00000025 # type: ignore # pylint:disable=protected-access +Rejected._definition = (FIELD("error", ObjDefinition.error, False, None, False),) # type: ignore # pylint:disable=protected-access +if _CAN_ADD_DOCSTRING: + Rejected.__doc__ = """ + The rejected outcome. + + At the target, the rejected outcome is used to indicate that an incoming Message is invalid and therefore + unprocessable. The rejected outcome when applied to a Message will cause the delivery-count to be incremented + in the header of the rejected Message. At the source, the rejected outcome means that the target has informed + the source that the message was rejected, and the source has taken the required action. The delivery SHOULD + NOT ever spontaneously attain the rejected state at the source. + + :param ~uamqp.error.AMQPError error: The error that caused the message to be rejected. + The value supplied in this field will be placed in the delivery-annotations of the rejected Message + associated with the symbolic key "rejected". + """ + + +Released = namedtuple('Released', []) +Released._code = 0x00000026 # type: ignore # pylint:disable=protected-access +Released._definition = () # type: ignore # pylint:disable=protected-access +if _CAN_ADD_DOCSTRING: + Released.__doc__ = """ + The released outcome. + + At the source the released outcome means that the message is no longer acquired by the receiver, and has been + made available for (re-)delivery to the same or other targets receiving from the node. The message is unchanged + at the node (i.e. the delivery-count of the header of the released Message MUST NOT be incremented). + As released is a terminal outcome, transfer of payload data will not be able to be resumed if the link becomes + suspended. A delivery may become released at the source even before all transfer frames have been sent, this + does not imply that the remaining transfers for the delivery will not be sent. The source MAY spontaneously + attain the released outcome for a Message (for example the source may implement some sort of time bound + acquisition lock, after which the acquisition of a message at a node is revoked to allow for delivery to an + alternative consumer). + + At the target, the released outcome is used to indicate that a given transfer was not and will not be acted upon. + """ + + +Modified = namedtuple('Modified', ['delivery_failed', 'undeliverable_here', 'message_annotations']) +Modified.__new__.__defaults__ = (None,) * len(Modified._fields) # type: ignore +Modified._code = 0x00000027 # type: ignore # pylint:disable=protected-access +Modified._definition = ( # type: ignore # pylint:disable=protected-access + FIELD('delivery_failed', AMQPTypes.boolean, False, None, False), + FIELD('undeliverable_here', AMQPTypes.boolean, False, None, False), + FIELD('message_annotations', FieldDefinition.fields, False, None, False)) +if _CAN_ADD_DOCSTRING: + Modified.__doc__ = """ + The modified outcome. + + At the source the modified outcome means that the message is no longer acquired by the receiver, and has been + made available for (re-)delivery to the same or other targets receiving from the node. The message has been + changed at the node in the ways indicated by the fields of the outcome. As modified is a terminal outcome, + transfer of payload data will not be able to be resumed if the link becomes suspended. A delivery may become + modified at the source even before all transfer frames have been sent, this does not imply that the remaining + transfers for the delivery will not be sent. The source MAY spontaneously attain the modified outcome for a + Message (for example the source may implement some sort of time bound acquisition lock, after which the + acquisition of a message at a node is revoked to allow for delivery to an alternative consumer with the + message modified in some way to denote the previous failed, e.g. with delivery-failed set to true). + At the target, the modified outcome is used to indicate that a given transfer was not and will not be acted + upon, and that the message should be modified in the specified ways at the node. + + :param bool delivery_failed: Count the transfer as an unsuccessful delivery attempt. + If the delivery-failed fiag is set, any Messages modified MUST have their deliverycount incremented. + :param bool undeliverable_here: Prevent redelivery. + If the undeliverable-here is set, then any Messages released MUST NOT be redelivered to the modifying + Link Endpoint. + :param dict message_annotations: Message attributes. + Map containing attributes to combine with the existing message-annotations held in the Message's header + section. Where the existing message-annotations of the Message contain an entry with the same key as an + entry in this field, the value in this field associated with that key replaces the one in the existing + headers; where the existing message-annotations has no such value, the value in this map is added. + """ diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/performatives.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/performatives.py new file mode 100644 index 000000000000..efcfc444ccd7 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/performatives.py @@ -0,0 +1,634 @@ +#------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +#-------------------------------------------------------------------------- + +# TODO: fix mypy errors for _code/_definition/__defaults__ (issue #26500) +from collections import namedtuple +import sys + +from .types import AMQPTypes, FieldDefinition, ObjDefinition +from .constants import FIELD + +_CAN_ADD_DOCSTRING = sys.version_info.major >= 3 + + +OpenFrame = namedtuple( + 'OpenFrame', + [ + 'container_id', + 'hostname', + 'max_frame_size', + 'channel_max', + 'idle_timeout', + 'outgoing_locales', + 'incoming_locales', + 'offered_capabilities', + 'desired_capabilities', + 'properties' + ]) +OpenFrame._code = 0x00000010 # type: ignore # pylint:disable=protected-access +OpenFrame._definition = ( # type: ignore # pylint:disable=protected-access + FIELD("container_id", AMQPTypes.string, True, None, False), + FIELD("hostname", AMQPTypes.string, False, None, False), + FIELD("max_frame_size", AMQPTypes.uint, False, 4294967295, False), + FIELD("channel_max", AMQPTypes.ushort, False, 65535, False), + FIELD("idle_timeout", AMQPTypes.uint, False, None, False), + FIELD("outgoing_locales", AMQPTypes.symbol, False, None, True), + FIELD("incoming_locales", AMQPTypes.symbol, False, None, True), + FIELD("offered_capabilities", AMQPTypes.symbol, False, None, True), + FIELD("desired_capabilities", AMQPTypes.symbol, False, None, True), + FIELD("properties", FieldDefinition.fields, False, None, False)) +if _CAN_ADD_DOCSTRING: + OpenFrame.__doc__ = """ + OPEN performative. Negotiate Connection parameters. + + The first frame sent on a connection in either direction MUST contain an Open body. + (Note that theConnection header which is sent first on the Connection is *not* a frame.) + The fields indicate thecapabilities and limitations of the sending peer. + + :param str container_id: The ID of the source container. + :param str hostname: The name of the target host. + The dns name of the host (either fully qualified or relative) to which the sendingpeer is connecting. + It is not mandatory to provide the hostname. If no hostname isprovided the receiving peer should select + a default based on its own configuration.This field can be used by AMQP proxies to determine the correct + back-end service toconnect the client to.This field may already have been specified by the sasl-init frame, + if a SASL layer is used, or, the server name indication extension as described in RFC-4366, if a TLSlayer + is used, in which case this field SHOULD be null or contain the same value. It is undefined what a different + value to those already specific means. + :param int max_frame_size: Proposed maximum frame size in bytes. + The largest frame size that the sending peer is able to accept on this Connection. + If this field is not set it means that the peer does not impose any specific limit. A peer MUST NOT send + frames larger than its partner can handle. A peer that receives an oversized frame MUST close the Connection + with the framing-error error-code. Both peers MUST accept frames of up to 512 (MIN-MAX-FRAME-SIZE) + octets large. + :param int channel_max: The maximum channel number that may be used on the Connection. + The channel-max value is the highest channel number that may be used on the Connection. This value plus one + is the maximum number of Sessions that can be simultaneously active on the Connection. A peer MUST not use + channel numbers outside the range that its partner can handle. A peer that receives a channel number + outside the supported range MUST close the Connection with the framing-error error-code. + :param int idle_timeout: Idle time-out in milliseconds. + The idle time-out required by the sender. A value of zero is the same as if it was not set (null). If the + receiver is unable or unwilling to support the idle time-out then it should close the connection with + an error explaining why (eg, because it is too small). If the value is not set, then the sender does not + have an idle time-out. However, senders doing this should be aware that implementations MAY choose to use + an internal default to efficiently manage a peer's resources. + :param list(str) outgoing_locales: Locales available for outgoing text. + A list of the locales that the peer supports for sending informational text. This includes Connection, + Session and Link error descriptions. A peer MUST support at least the en-US locale. Since this value + is always supported, it need not be supplied in the outgoing-locales. A null value or an empty list implies + that only en-US is supported. + :param list(str) incoming_locales: Desired locales for incoming text in decreasing level of preference. + A list of locales that the sending peer permits for incoming informational text. This list is ordered in + decreasing level of preference. The receiving partner will chose the first (most preferred) incoming locale + from those which it supports. If none of the requested locales are supported, en-US will be chosen. Note + that en-US need not be supplied in this list as it is always the fallback. A peer may determine which of the + permitted incoming locales is chosen by examining the partner's supported locales asspecified in the + outgoing_locales field. A null value or an empty list implies that only en-US is supported. + :param list(str) offered_capabilities: The extension capabilities the sender supports. + If the receiver of the offered-capabilities requires an extension capability which is not present in the + offered-capability list then it MUST close the connection. A list of commonly defined connection capabilities + and their meanings can be found here: http://www.amqp.org/specification/1.0/connection-capabilities. + :param list(str) required_capabilities: The extension capabilities the sender may use if the receiver supports + them. The desired-capability list defines which extension capabilities the sender MAY use if the receiver + offers them (i.e. they are in the offered-capabilities list received by the sender of the + desired-capabilities). If the receiver of the desired-capabilities offers extension capabilities which are + not present in the desired-capability list it received, then it can be sure those (undesired) capabilities + will not be used on the Connection. + :param dict properties: Connection properties. + The properties map contains a set of fields intended to indicate information about the connection and its + container. A list of commonly defined connection properties and their meanings can be found + here: http://www.amqp.org/specification/1.0/connection-properties. + """ + + +BeginFrame = namedtuple( + 'BeginFrame', + [ + 'remote_channel', + 'next_outgoing_id', + 'incoming_window', + 'outgoing_window', + 'handle_max', + 'offered_capabilities', + 'desired_capabilities', + 'properties' + ]) +BeginFrame._code = 0x00000011 # type: ignore # pylint:disable=protected-access +BeginFrame._definition = ( # type: ignore # pylint:disable=protected-access + FIELD("remote_channel", AMQPTypes.ushort, False, None, False), + FIELD("next_outgoing_id", AMQPTypes.uint, True, None, False), + FIELD("incoming_window", AMQPTypes.uint, True, None, False), + FIELD("outgoing_window", AMQPTypes.uint, True, None, False), + FIELD("handle_max", AMQPTypes.uint, False, 4294967295, False), + FIELD("offered_capabilities", AMQPTypes.symbol, False, None, True), + FIELD("desired_capabilities", AMQPTypes.symbol, False, None, True), + FIELD("properties", FieldDefinition.fields, False, None, False)) +if _CAN_ADD_DOCSTRING: + BeginFrame.__doc__ = """ + BEGIN performative. Begin a Session on a channel. + + Indicate that a Session has begun on the channel. + + :param int remote_channel: The remote channel for this Session. + If a Session is locally initiated, the remote-channel MUST NOT be set. When an endpoint responds to a + remotely initiated Session, the remote-channel MUST be set to the channel on which the remote Session + sent the begin. + :param int next_outgoing_id: The transfer-id of the first transfer id the sender will send. + The next-outgoing-id is used to assign a unique transfer-id to all outgoing transfer frames on a given + session. The next-outgoing-id may be initialized to an arbitrary value and is incremented after each + successive transfer according to RFC-1982 serial number arithmetic. + :param int incoming_window: The initial incoming-window of the sender. + The incoming-window defines the maximum number of incoming transfer frames that the endpoint can currently + receive. This identifies a current maximum incoming transfer-id that can be computed by subtracting one + from the sum of incoming-window and next-incoming-id. + :param int outgoing_window: The initial outgoing-window of the sender. + The outgoing-window defines the maximum number of outgoing transfer frames that the endpoint can currently + send. This identifies a current maximum outgoing transfer-id that can be computed by subtracting one from + the sum of outgoing-window and next-outgoing-id. + :param int handle_max: The maximum handle value that may be used on the Session. + The handle-max value is the highest handle value that may be used on the Session. A peer MUST NOT attempt + to attach a Link using a handle value outside the range that its partner can handle. A peer that receives + a handle outside the supported range MUST close the Connection with the framing-error error-code. + :param list(str) offered_capabilities: The extension capabilities the sender supports. + A list of commonly defined session capabilities and their meanings can be found + here: http://www.amqp.org/specification/1.0/session-capabilities. + :param list(str) desired_capabilities: The extension capabilities the sender may use if the receiver + supports them. + :param dict properties: Session properties. + The properties map contains a set of fields intended to indicate information about the session and its + container. A list of commonly defined session properties and their meanings can be found + here: http://www.amqp.org/specification/1.0/session-properties. + """ + + +AttachFrame = namedtuple( + 'AttachFrame', + [ + 'name', + 'handle', + 'role', + 'send_settle_mode', + 'rcv_settle_mode', + 'source', + 'target', + 'unsettled', + 'incomplete_unsettled', + 'initial_delivery_count', + 'max_message_size', + 'offered_capabilities', + 'desired_capabilities', + 'properties' + ]) +AttachFrame._code = 0x00000012 # type: ignore # pylint:disable=protected-access +AttachFrame._definition = ( # type: ignore # pylint:disable=protected-access + FIELD("name", AMQPTypes.string, True, None, False), + FIELD("handle", AMQPTypes.uint, True, None, False), + FIELD("role", AMQPTypes.boolean, True, None, False), + FIELD("send_settle_mode", AMQPTypes.ubyte, False, 2, False), + FIELD("rcv_settle_mode", AMQPTypes.ubyte, False, 0, False), + FIELD("source", ObjDefinition.source, False, None, False), + FIELD("target", ObjDefinition.target, False, None, False), + FIELD("unsettled", AMQPTypes.map, False, None, False), + FIELD("incomplete_unsettled", AMQPTypes.boolean, False, False, False), + FIELD("initial_delivery_count", AMQPTypes.uint, False, None, False), + FIELD("max_message_size", AMQPTypes.ulong, False, None, False), + FIELD("offered_capabilities", AMQPTypes.symbol, False, None, True), + FIELD("desired_capabilities", AMQPTypes.symbol, False, None, True), + FIELD("properties", FieldDefinition.fields, False, None, False)) +if _CAN_ADD_DOCSTRING: + AttachFrame.__doc__ = """ + ATTACH performative. Attach a Link to a Session. + + The attach frame indicates that a Link Endpoint has been attached to the Session. The opening flag + is used to indicate that the Link Endpoint is newly created. + + :param str name: The name of the link. + This name uniquely identifies the link from the container of the source to the container of the target + node, e.g. if the container of the source node is A, and the container of the target node is B, the link + may be globally identified by the (ordered) tuple(A,B,). + :param int handle: The handle of the link. + The handle MUST NOT be used for other open Links. An attempt to attach using a handle which is already + associated with a Link MUST be responded to with an immediate close carrying a Handle-in-usesession-error. + To make it easier to monitor AMQP link attach frames, it is recommended that implementations always assign + the lowest available handle to this field. + :param bool role: The role of the link endpoint. Either Role.Sender (False) or Role.Receiver (True). + :param str send_settle_mode: The settlement mode for the Sender. + Determines the settlement policy for deliveries sent at the Sender. When set at the Receiver this indicates + the desired value for the settlement mode at the Sender. When set at the Sender this indicates the actual + settlement mode in use. + :param str rcv_settle_mode: The settlement mode of the Receiver. + Determines the settlement policy for unsettled deliveries received at the Receiver. When set at the Sender + this indicates the desired value for the settlement mode at the Receiver. When set at the Receiver this + indicates the actual settlement mode in use. + :param ~uamqp.messaging.Source source: The source for Messages. + If no source is specified on an outgoing Link, then there is no source currently attached to the Link. + A Link with no source will never produce outgoing Messages. + :param ~uamqp.messaging.Target target: The target for Messages. + If no target is specified on an incoming Link, then there is no target currently attached to the Link. + A Link with no target will never permit incoming Messages. + :param dict unsettled: Unsettled delivery state. + This is used to indicate any unsettled delivery states when a suspended link is resumed. The map is keyed + by delivery-tag with values indicating the delivery state. The local and remote delivery states for a given + delivery-tag MUST be compared to resolve any in-doubt deliveries. If necessary, deliveries MAY be resent, + or resumed based on the outcome of this comparison. If the local unsettled map is too large to be encoded + within a frame of the agreed maximum frame size then the session may be ended with the + frame-size-too-smallerror. The endpoint SHOULD make use of the ability to send an incomplete unsettled map + to avoid sending an error. The unsettled map MUST NOT contain null valued keys. When reattaching + (as opposed to resuming), the unsettled map MUST be null. + :param bool incomplete_unsettled: + If set to true this field indicates that the unsettled map provided is not complete. When the map is + incomplete the recipient of the map cannot take the absence of a delivery tag from the map as evidence of + settlement. On receipt of an incomplete unsettled map a sending endpoint MUST NOT send any new deliveries + (i.e. deliveries where resume is not set to true) to its partner (and a receiving endpoint which sent an + incomplete unsettled map MUST detach with an error on receiving a transfer which does not have the resume + flag set to true). + :param int initial_delivery_count: This MUST NOT be null if role is sender, + and it is ignored if the role is receiver. + :param int max_message_size: The maximum message size supported by the link endpoint. + This field indicates the maximum message size supported by the link endpoint. Any attempt to deliver a + message larger than this results in a message-size-exceeded link-error. If this field is zero or unset, + there is no maximum size imposed by the link endpoint. + :param list(str) offered_capabilities: The extension capabilities the sender supports. + A list of commonly defined session capabilities and their meanings can be found + here: http://www.amqp.org/specification/1.0/link-capabilities. + :param list(str) desired_capabilities: The extension capabilities the sender may use if the receiver + supports them. + :param dict properties: Link properties. + The properties map contains a set of fields intended to indicate information about the link and its + container. A list of commonly defined link properties and their meanings can be found + here: http://www.amqp.org/specification/1.0/link-properties. + """ + + +FlowFrame = namedtuple( + 'FlowFrame', + [ + 'next_incoming_id', + 'incoming_window', + 'next_outgoing_id', + 'outgoing_window', + 'handle', + 'delivery_count', + 'link_credit', + 'available', + 'drain', + 'echo', + 'properties' + ]) +FlowFrame.__new__.__defaults__ = (None, None, None, None, None, None, None) # type: ignore +FlowFrame._code = 0x00000013 # type: ignore # pylint:disable=protected-access +FlowFrame._definition = ( # type: ignore # pylint:disable=protected-access + FIELD("next_incoming_id", AMQPTypes.uint, False, None, False), + FIELD("incoming_window", AMQPTypes.uint, True, None, False), + FIELD("next_outgoing_id", AMQPTypes.uint, True, None, False), + FIELD("outgoing_window", AMQPTypes.uint, True, None, False), + FIELD("handle", AMQPTypes.uint, False, None, False), + FIELD("delivery_count", AMQPTypes.uint, False, None, False), + FIELD("link_credit", AMQPTypes.uint, False, None, False), + FIELD("available", AMQPTypes.uint, False, None, False), + FIELD("drain", AMQPTypes.boolean, False, False, False), + FIELD("echo", AMQPTypes.boolean, False, False, False), + FIELD("properties", FieldDefinition.fields, False, None, False)) +if _CAN_ADD_DOCSTRING: + FlowFrame.__doc__ = """ + FLOW performative. Update link state. + + Updates the flow state for the specified Link. + + :param int next_incoming_id: Identifies the expected transfer-id of the next incoming transfer frame. + This value is not set if and only if the sender has not yet received the begin frame for the session. + :param int incoming_window: Defines the maximum number of incoming transfer frames that the endpoint + concurrently receive. + :param int next_outgoing_id: The transfer-id that will be assigned to the next outgoing transfer frame. + :param int outgoing_window: Defines the maximum number of outgoing transfer frames that the endpoint could + potentially currently send, if it was not constrained by restrictions imposed by its peer's incoming-window. + :param int handle: If set, indicates that the flow frame carries flow state information for the local Link + Endpoint associated with the given handle. If not set, the flow frame is carrying only information + pertaining to the Session Endpoint. If set to a handle that is not currently associated with an attached + Link, the recipient MUST respond by ending the session with an unattached-handle session error. + :param int delivery_count: The endpoint's delivery-count. + When the handle field is not set, this field MUST NOT be set. When the handle identifies that the flow + state is being sent from the Sender Link Endpoint to Receiver Link Endpoint this field MUST be set to the + current delivery-count of the Link Endpoint. When the flow state is being sent from the Receiver Endpoint + to the Sender Endpoint this field MUST be set to the last known value of the corresponding Sending Endpoint. + In the event that the Receiving Link Endpoint has not yet seen the initial attach frame from the Sender + this field MUST NOT be set. + :param int link_credit: The current maximum number of Messages that can be received. + The current maximum number of Messages that can be handled at the Receiver Endpoint of the Link. Only the + receiver endpoint can independently set this value. The sender endpoint sets this to the last known + value seen from the receiver. When the handle field is not set, this field MUST NOT be set. + :param int available: The number of available Messages. + The number of Messages awaiting credit at the link sender endpoint. Only the sender can independently set + this value. The receiver sets this to the last known value seen from the sender. When the handle field is + not set, this field MUST NOT be set. + :param bool drain: Indicates drain mode. + When flow state is sent from the sender to the receiver, this field contains the actual drain mode of the + sender. When flow state is sent from the receiver to the sender, this field contains the desired drain + mode of the receiver. When the handle field is not set, this field MUST NOT be set. + :param bool echo: Request link state from other endpoint. + :param dict properties: Link state properties. + A list of commonly defined link state properties and their meanings can be found + here: http://www.amqp.org/specification/1.0/link-state-properties. + """ + + +TransferFrame = namedtuple( + 'TransferFrame', + [ + 'handle', + 'delivery_id', + 'delivery_tag', + 'message_format', + 'settled', + 'more', + 'rcv_settle_mode', + 'state', + 'resume', + 'aborted', + 'batchable', + 'payload' + ]) +TransferFrame._code = 0x00000014 # type: ignore # pylint:disable=protected-access +TransferFrame._definition = ( # type: ignore # pylint:disable=protected-access + FIELD("handle", AMQPTypes.uint, True, None, False), + FIELD("delivery_id", AMQPTypes.uint, False, None, False), + FIELD("delivery_tag", AMQPTypes.binary, False, None, False), + FIELD("message_format", AMQPTypes.uint, False, 0, False), + FIELD("settled", AMQPTypes.boolean, False, None, False), + FIELD("more", AMQPTypes.boolean, False, False, False), + FIELD("rcv_settle_mode", AMQPTypes.ubyte, False, None, False), + FIELD("state", ObjDefinition.delivery_state, False, None, False), + FIELD("resume", AMQPTypes.boolean, False, False, False), + FIELD("aborted", AMQPTypes.boolean, False, False, False), + FIELD("batchable", AMQPTypes.boolean, False, False, False), + None) +if _CAN_ADD_DOCSTRING: + TransferFrame.__doc__ = """ + TRANSFER performative. Transfer a Message. + + The transfer frame is used to send Messages across a Link. Messages may be carried by a single transfer up + to the maximum negotiated frame size for the Connection. Larger Messages may be split across several + transfer frames. + + :param int handle: Specifies the Link on which the Message is transferred. + :param int delivery_id: Alias for delivery-tag. + The delivery-id MUST be supplied on the first transfer of a multi-transfer delivery. On continuation + transfers the delivery-id MAY be omitted. It is an error if the delivery-id on a continuation transfer + differs from the delivery-id on the first transfer of a delivery. + :param bytes delivery_tag: Uniquely identifies the delivery attempt for a given Message on this Link. + This field MUST be specified for the first transfer of a multi transfer message and may only be + omitted for continuation transfers. + :param int message_format: Indicates the message format. + This field MUST be specified for the first transfer of a multi transfer message and may only be omitted + for continuation transfers. + :param bool settled: If not set on the first (or only) transfer for a delivery, then the settled flag MUST + be interpreted as being false. For subsequent transfers if the settled flag is left unset then it MUST be + interpreted as true if and only if the value of the settled flag on any of the preceding transfers was + true; if no preceding transfer was sent with settled being true then the value when unset MUST be taken + as false. If the negotiated value for snd-settle-mode at attachment is settled, then this field MUST be + true on at least one transfer frame for a delivery (i.e. the delivery must be settled at the Sender at + the point the delivery has been completely transferred). If the negotiated value for snd-settle-mode at + attachment is unsettled, then this field MUST be false (or unset) on every transfer frame for a delivery + (unless the delivery is aborted). + :param bool more: Indicates that the Message has more content. + Note that if both the more and aborted fields are set to true, the aborted flag takes precedence. That is + a receiver should ignore the value of the more field if the transfer is marked as aborted. A sender + SHOULD NOT set the more flag to true if it also sets the aborted flag to true. + :param str rcv_settle_mode: If first, this indicates that the Receiver MUST settle the delivery once it has + arrived without waiting for the Sender to settle first. If second, this indicates that the Receiver MUST + NOT settle until sending its disposition to the Sender and receiving a settled disposition from the sender. + If not set, this value is defaulted to the value negotiated on link attach. If the negotiated link value is + first, then it is illegal to set this field to second. If the message is being sent settled by the Sender, + the value of this field is ignored. The (implicit or explicit) value of this field does not form part of the + transfer state, and is not retained if a link is suspended and subsequently resumed. + :param bytes state: The state of the delivery at the sender. + When set this informs the receiver of the state of the delivery at the sender. This is particularly useful + when transfers of unsettled deliveries are resumed after a resuming a link. Setting the state on the + transfer can be thought of as being equivalent to sending a disposition immediately before the transfer + performative, i.e. it is the state of the delivery (not the transfer) that existed at the point the frame + was sent. Note that if the transfer performative (or an earlier disposition performative referring to the + delivery) indicates that the delivery has attained a terminal state, then no future transfer or disposition + sent by the sender can alter that terminal state. + :param bool resume: Indicates a resumed delivery. + If true, the resume flag indicates that the transfer is being used to reassociate an unsettled delivery + from a dissociated link endpoint. The receiver MUST ignore resumed deliveries that are not in its local + unsettled map. The sender MUST NOT send resumed transfers for deliveries not in its local unsettledmap. + If a resumed delivery spans more than one transfer performative, then the resume flag MUST be set to true + on the first transfer of the resumed delivery. For subsequent transfers for the same delivery the resume + flag may be set to true, or may be omitted. In the case where the exchange of unsettled maps makes clear + that all message data has been successfully transferred to the receiver, and that only the final state + (andpotentially settlement) at the sender needs to be conveyed, then a resumed delivery may carry no + payload and instead act solely as a vehicle for carrying the terminal state of the delivery at the sender. + :param bool aborted: Indicates that the Message is aborted. + Aborted Messages should be discarded by the recipient (any payload within the frame carrying the performative + MUST be ignored). An aborted Message is implicitly settled. + :param bool batchable: Batchable hint. + If true, then the issuer is hinting that there is no need for the peer to urgently communicate updated + delivery state. This hint may be used to artificially increase the amount of batching an implementation + uses when communicating delivery states, and thereby save bandwidth. If the message being delivered is too + large to fit within a single frame, then the setting of batchable to true on any of the transfer + performatives for the delivery is equivalent to setting batchable to true for all the transfer performatives + for the delivery. The batchable value does not form part of the transfer state, and is not retained if a + link is suspended and subsequently resumed. + """ + + +DispositionFrame = namedtuple( + 'DispositionFrame', + [ + 'role', + 'first', + 'last', + 'settled', + 'state', + 'batchable' + ]) +DispositionFrame._code = 0x00000015 # type: ignore # pylint:disable=protected-access +DispositionFrame._definition = ( # type: ignore # pylint:disable=protected-access + FIELD("role", AMQPTypes.boolean, True, None, False), + FIELD("first", AMQPTypes.uint, True, None, False), + FIELD("last", AMQPTypes.uint, False, None, False), + FIELD("settled", AMQPTypes.boolean, False, False, False), + FIELD("state", ObjDefinition.delivery_state, False, None, False), + FIELD("batchable", AMQPTypes.boolean, False, False, False)) +if _CAN_ADD_DOCSTRING: + DispositionFrame.__doc__ = """ + DISPOSITION performative. Inform remote peer of delivery state changes. + + The disposition frame is used to inform the remote peer of local changes in the state of deliveries. + The disposition frame may reference deliveries from many different links associated with a session, + although all links MUST have the directionality indicated by the specified role. Note that it is possible + for a disposition sent from sender to receiver to refer to a delivery which has not yet completed + (i.e. a delivery which is spread over multiple frames and not all frames have yet been sent). The use of such + interleaving is discouraged in favor of carrying the modified state on the next transfer performative for + the delivery. The disposition performative may refer to deliveries on links that are no longer attached. + As long as the links have not been closed or detached with an error then the deliveries are still "live" and + the updated state MUST be applied. + + :param str role: Directionality of disposition. + The role identifies whether the disposition frame contains information about sending link endpoints + or receiving link endpoints. + :param int first: Lower bound of deliveries. + Identifies the lower bound of delivery-ids for the deliveries in this set. + :param int last: Upper bound of deliveries. + Identifies the upper bound of delivery-ids for the deliveries in this set. If not set, + this is taken to be the same as first. + :param bool settled: Indicates deliveries are settled. + If true, indicates that the referenced deliveries are considered settled by the issuing endpoint. + :param bytes state: Indicates state of deliveries. + Communicates the state of all the deliveries referenced by this disposition. + :param bool batchable: Batchable hint. + If true, then the issuer is hinting that there is no need for the peer to urgently communicate the impact + of the updated delivery states. This hint may be used to artificially increase the amount of batching an + implementation uses when communicating delivery states, and thereby save bandwidth. + """ + +DetachFrame = namedtuple('DetachFrame', ['handle', 'closed', 'error']) +DetachFrame._code = 0x00000016 # type: ignore # pylint:disable=protected-access +DetachFrame._definition = ( # type: ignore # pylint:disable=protected-access + FIELD("handle", AMQPTypes.uint, True, None, False), + FIELD("closed", AMQPTypes.boolean, False, False, False), + FIELD("error", ObjDefinition.error, False, None, False)) +if _CAN_ADD_DOCSTRING: + DetachFrame.__doc__ = """ + DETACH performative. Detach the Link Endpoint from the Session. + + Detach the Link Endpoint from the Session. This un-maps the handle and makes it available for + use by other Links + + :param int handle: The local handle of the link to be detached. + :param bool handle: If true then the sender has closed the link. + :param ~uamqp.error.AMQPError error: Error causing the detach. + If set, this field indicates that the Link is being detached due to an error condition. + The value of the field should contain details on the cause of the error. + """ + + +EndFrame = namedtuple('EndFrame', ['error']) +EndFrame._code = 0x00000017 # type: ignore # pylint:disable=protected-access +EndFrame._definition = (FIELD("error", ObjDefinition.error, False, None, False),) # type: ignore # pylint:disable=protected-access +if _CAN_ADD_DOCSTRING: + EndFrame.__doc__ = """ + END performative. End the Session. + + Indicates that the Session has ended. + + :param ~uamqp.error.AMQPError error: Error causing the end. + If set, this field indicates that the Session is being ended due to an error condition. + The value of the field should contain details on the cause of the error. + """ + + +CloseFrame = namedtuple('CloseFrame', ['error']) +CloseFrame._code = 0x00000018 # type: ignore # pylint:disable=protected-access +CloseFrame._definition = (FIELD("error", ObjDefinition.error, False, None, False),) # type: ignore # pylint:disable=protected-access +if _CAN_ADD_DOCSTRING: + CloseFrame.__doc__ = """ + CLOSE performative. Signal a Connection close. + + Sending a close signals that the sender will not be sending any more frames (or bytes of any other kind) on + the Connection. Orderly shutdown requires that this frame MUST be written by the sender. It is illegal to + send any more frames (or bytes of any other kind) after sending a close frame. + + :param ~uamqp.error.AMQPError error: Error causing the close. + If set, this field indicates that the Connection is being closed due to an error condition. + The value of the field should contain details on the cause of the error. + """ + + +SASLMechanism = namedtuple('SASLMechanism', ['sasl_server_mechanisms']) +SASLMechanism._code = 0x00000040 # type: ignore # pylint:disable=protected-access +SASLMechanism._definition = (FIELD('sasl_server_mechanisms', AMQPTypes.symbol, True, None, True),) # type: ignore # pylint:disable=protected-access +if _CAN_ADD_DOCSTRING: + SASLMechanism.__doc__ = """ + Advertise available sasl mechanisms. + + dvertises the available SASL mechanisms that may be used for authentication. + + :param list(bytes) sasl_server_mechanisms: Supported sasl mechanisms. + A list of the sasl security mechanisms supported by the sending peer. + It is invalid for this list to be null or empty. If the sending peer does not require its partner to + authenticate with it, then it should send a list of one element with its value as the SASL mechanism + ANONYMOUS. The server mechanisms are ordered in decreasing level of preference. + """ + + +SASLInit = namedtuple('SASLInit', ['mechanism', 'initial_response', 'hostname']) +SASLInit._code = 0x00000041 # type: ignore # pylint:disable=protected-access +SASLInit._definition = ( # type: ignore # pylint:disable=protected-access + FIELD('mechanism', AMQPTypes.symbol, True, None, False), + FIELD('initial_response', AMQPTypes.binary, False, None, False), + FIELD('hostname', AMQPTypes.string, False, None, False)) +if _CAN_ADD_DOCSTRING: + SASLInit.__doc__ = """ + Initiate sasl exchange. + + Selects the sasl mechanism and provides the initial response if needed. + + :param bytes mechanism: Selected security mechanism. + The name of the SASL mechanism used for the SASL exchange. If the selected mechanism is not supported by + the receiving peer, it MUST close the Connection with the authentication-failure close-code. Each peer + MUST authenticate using the highest-level security profile it can handle from the list provided by the + partner. + :param bytes initial_response: Security response data. + A block of opaque data passed to the security mechanism. The contents of this data are defined by the + SASL security mechanism. + :param str hostname: The name of the target host. + The DNS name of the host (either fully qualified or relative) to which the sending peer is connecting. It + is not mandatory to provide the hostname. If no hostname is provided the receiving peer should select a + default based on its own configuration. This field can be used by AMQP proxies to determine the correct + back-end service to connect the client to, and to determine the domain to validate the client's credentials + against. This field may already have been specified by the server name indication extension as described + in RFC-4366, if a TLS layer is used, in which case this field SHOULD benull or contain the same value. + It is undefined what a different value to those already specific means. + """ + + +SASLChallenge = namedtuple('SASLChallenge', ['challenge']) +SASLChallenge._code = 0x00000042 # type: ignore # pylint:disable=protected-access +SASLChallenge._definition = (FIELD('challenge', AMQPTypes.binary, True, None, False),) # type: ignore # pylint:disable=protected-access +if _CAN_ADD_DOCSTRING: + SASLChallenge.__doc__ = """ + Security mechanism challenge. + + Send the SASL challenge data as defined by the SASL specification. + + :param bytes challenge: Security challenge data. + Challenge information, a block of opaque binary data passed to the security mechanism. + """ + + +SASLResponse = namedtuple('SASLResponse', ['response']) +SASLResponse._code = 0x00000043 # type: ignore # pylint:disable=protected-access +SASLResponse._definition = (FIELD('response', AMQPTypes.binary, True, None, False),) # type: ignore # pylint:disable=protected-access +if _CAN_ADD_DOCSTRING: + SASLResponse.__doc__ = """ + Security mechanism response. + + Send the SASL response data as defined by the SASL specification. + + :param bytes response: Security response data. + """ + + +SASLOutcome = namedtuple('SASLOutcome', ['code', 'additional_data']) +SASLOutcome._code = 0x00000044 # type: ignore # pylint:disable=protected-access +SASLOutcome._definition = ( # type: ignore # pylint:disable=protected-access + FIELD('code', AMQPTypes.ubyte, True, None, False), + FIELD('additional_data', AMQPTypes.binary, False, None, False)) +if _CAN_ADD_DOCSTRING: + SASLOutcome.__doc__ = """ + Indicates the outcome of the sasl dialog. + + This frame indicates the outcome of the SASL dialog. Upon successful completion of the SASL dialog the + Security Layer has been established, and the peers must exchange protocol headers to either starta nested + Security Layer, or to establish the AMQP Connection. + + :param int code: Indicates the outcome of the sasl dialog. + A reply-code indicating the outcome of the SASL dialog. + :param bytes additional_data: Additional data as specified in RFC-4422. + The additional-data field carries additional data on successful authentication outcomeas specified by + the SASL specification (RFC-4422). If the authentication is unsuccessful, this field is not set. + """ diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/receiver.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/receiver.py new file mode 100644 index 000000000000..5713f51b4b8c --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/receiver.py @@ -0,0 +1,121 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + +import uuid +import logging +from typing import Optional, Union + +from ._decode import decode_payload +from .link import Link +from .constants import LinkState, Role +from .performatives import TransferFrame, DispositionFrame +from .outcomes import Received, Accepted, Rejected, Released, Modified + + +_LOGGER = logging.getLogger(__name__) + + +class ReceiverLink(Link): + def __init__(self, session, handle, source_address, **kwargs): + name = kwargs.pop("name", None) or str(uuid.uuid4()) + role = Role.Receiver + if "target_address" not in kwargs: + kwargs["target_address"] = "receiver-link-{}".format(name) + super(ReceiverLink, self).__init__(session, handle, name, role, source_address=source_address, **kwargs) + self._on_transfer = kwargs.pop("on_transfer") + self._received_payload = bytearray() + + @classmethod + def from_incoming_frame(cls, session, handle, frame): + # TODO: Assuming we establish all links for now... + # check link_create_from_endpoint in C lib + raise NotImplementedError("Pending") + + def _process_incoming_message(self, frame, message): + try: + return self._on_transfer(frame, message) + except Exception as e: # pylint: disable=broad-except + _LOGGER.error("Transfer callback function failed with error: %r", e, extra=self.network_trace_params) + return None + + def _incoming_attach(self, frame): + super(ReceiverLink, self)._incoming_attach(frame) + if frame[9] is None: # initial_delivery_count + _LOGGER.info("Cannot get initial-delivery-count. Detaching link", extra=self.network_trace_params) + self._set_state(LinkState.DETACHED) # TODO: Send detach now? + self.delivery_count = frame[9] + self.current_link_credit = self.link_credit + self._outgoing_flow() + + def _incoming_transfer(self, frame): + if self.network_trace: + _LOGGER.debug("<- %r", TransferFrame(payload=b"***", *frame[:-1]), extra=self.network_trace_params) + self.current_link_credit -= 1 + self.delivery_count += 1 + self.received_delivery_id = frame[1] # delivery_id + if not self.received_delivery_id and not self._received_payload: + pass # TODO: delivery error + if self._received_payload or frame[5]: # more + self._received_payload.extend(frame[11]) + if not frame[5]: + if self._received_payload: + message = decode_payload(memoryview(self._received_payload)) + self._received_payload = bytearray() + else: + message = decode_payload(frame[11]) + delivery_state = self._process_incoming_message(frame, message) + if not frame[4] and delivery_state: # settled + self._outgoing_disposition( + first=frame[1], + last=frame[1], + settled=True, + state=delivery_state, + batchable=None + ) + + def _wait_for_response(self, wait: Union[bool, float]) -> None: + if wait is True: + self._session._connection.listen(wait=False) # pylint: disable=protected-access + if self.state == LinkState.ERROR: + raise self._error + elif wait: + self._session._connection.listen(wait=wait) # pylint: disable=protected-access + if self.state == LinkState.ERROR: + raise self._error + + def _outgoing_disposition( + self, + first: int, + last: Optional[int], + settled: Optional[bool], + state: Optional[Union[Received, Accepted, Rejected, Released, Modified]], + batchable: Optional[bool], + ): + disposition_frame = DispositionFrame( + role=self.role, first=first, last=last, settled=settled, state=state, batchable=batchable + ) + if self.network_trace: + _LOGGER.debug("-> %r", DispositionFrame(*disposition_frame), extra=self.network_trace_params) + self._session._outgoing_disposition(disposition_frame) # pylint: disable=protected-access + + def attach(self): + super().attach() + self._received_payload = bytearray() + + def send_disposition( + self, + *, + wait: Union[bool, float] = False, + first_delivery_id: int, + last_delivery_id: Optional[int] = None, + settled: Optional[bool] = None, + delivery_state: Optional[Union[Received, Accepted, Rejected, Released, Modified]] = None, + batchable: Optional[bool] = None + ): + if self._is_closed: + raise ValueError("Link already closed.") + self._outgoing_disposition(first_delivery_id, last_delivery_id, settled, delivery_state, batchable) + self._wait_for_response(wait) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/sasl.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/sasl.py new file mode 100644 index 000000000000..c4ff9d265540 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/sasl.py @@ -0,0 +1,146 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + +from ._transport import SSLTransport, WebSocketTransport, AMQPS_PORT +from .constants import SASLCode, SASL_HEADER_FRAME, WEBSOCKET_PORT +from .performatives import SASLInit + + +_SASL_FRAME_TYPE = b"\x01" + + +class SASLPlainCredential(object): + """PLAIN SASL authentication mechanism. + See https://tools.ietf.org/html/rfc4616 for details + """ + + mechanism = b"PLAIN" + + def __init__(self, authcid, passwd, authzid=None): + self.authcid = authcid + self.passwd = passwd + self.authzid = authzid + + def start(self): + if self.authzid: + login_response = self.authzid.encode("utf-8") + else: + login_response = b"" + login_response += b"\0" + login_response += self.authcid.encode("utf-8") + login_response += b"\0" + login_response += self.passwd.encode("utf-8") + return login_response + + +class SASLAnonymousCredential(object): + """ANONYMOUS SASL authentication mechanism. + See https://tools.ietf.org/html/rfc4505 for details + """ + + mechanism = b"ANONYMOUS" + + def start(self): # pylint: disable=no-self-use + return b"" + + +class SASLExternalCredential(object): + """EXTERNAL SASL mechanism. + Enables external authentication, i.e. not handled through this protocol. + Only passes 'EXTERNAL' as authentication mechanism, but no further + authentication data. + """ + + mechanism = b"EXTERNAL" + + def start(self): # pylint: disable=no-self-use + return b"" + + +class SASLTransportMixin: + def _negotiate(self): + self.write(SASL_HEADER_FRAME) + _, returned_header = self.receive_frame() + if returned_header[1] != SASL_HEADER_FRAME: + raise ValueError( + f"""Mismatching AMQP header protocol. Expected: {SASL_HEADER_FRAME!r},""" + """received: {returned_header[1]!r}""" + ) + + _, supported_mechanisms = self.receive_frame(verify_frame_type=1) + if ( + self.credential.mechanism not in supported_mechanisms[1][0] + ): # sasl_server_mechanisms + raise ValueError( + "Unsupported SASL credential type: {}".format(self.credential.mechanism) + ) + sasl_init = SASLInit( + mechanism=self.credential.mechanism, + initial_response=self.credential.start(), + hostname=self.host, + ) + self.send_frame(0, sasl_init, frame_type=_SASL_FRAME_TYPE) + + _, next_frame = self.receive_frame(verify_frame_type=1) + frame_type, fields = next_frame + if frame_type != 0x00000044: # SASLOutcome + raise NotImplementedError("Unsupported SASL challenge") + if fields[0] == SASLCode.Ok: # code + return + raise ValueError( + "SASL negotiation failed.\nOutcome: {}\nDetails: {}".format(*fields) + ) + + +class SASLTransport(SSLTransport, SASLTransportMixin): + def __init__( + self, + host, + credential, + *, + port=AMQPS_PORT, + connect_timeout=None, + ssl_opts=None, + **kwargs, + ): + self.credential = credential + ssl_opts = ssl_opts or True + super(SASLTransport, self).__init__( + host, + port=port, + connect_timeout=connect_timeout, + ssl_opts=ssl_opts, + **kwargs, + ) + + def negotiate(self): + with self.block(): + self._negotiate() + + +class SASLWithWebSocket(WebSocketTransport, SASLTransportMixin): + def __init__( + self, + host, + credential, + *, + port=WEBSOCKET_PORT, + connect_timeout=None, + ssl_opts=None, + **kwargs, + ): + self.credential = credential + ssl_opts = ssl_opts or True + super().__init__( + host, + port=port, + connect_timeout=connect_timeout, + ssl_opts=ssl_opts, + **kwargs, + ) + + def negotiate(self): + self._negotiate() diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/sender.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/sender.py new file mode 100644 index 000000000000..26c78f5f9c17 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/sender.py @@ -0,0 +1,200 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- +import struct +import uuid +import logging +import time + +from ._encode import encode_payload +from .link import Link +from .constants import SessionTransferState, LinkDeliverySettleReason, LinkState, Role, SenderSettleMode, SessionState +from .error import AMQPLinkError, ErrorCondition, MessageException + +_LOGGER = logging.getLogger(__name__) + + +class PendingDelivery(object): + def __init__(self, **kwargs): + self.message = kwargs.get("message") + self.sent = False + self.frame = None + self.on_delivery_settled = kwargs.get("on_delivery_settled") + self.start = time.time() + self.transfer_state = None + self.timeout = kwargs.get("timeout") + self.settled = kwargs.get("settled", False) + self._network_trace_params = kwargs.get('network_trace_params') + + def on_settled(self, reason, state): + if self.on_delivery_settled and not self.settled: + try: + self.on_delivery_settled(reason, state) + except Exception as e: # pylint:disable=broad-except + _LOGGER.warning( + "Message 'on_send_complete' callback failed: %r", + e, + extra=self._network_trace_params + ) + self.settled = True + + +class SenderLink(Link): + def __init__(self, session, handle, target_address, **kwargs): + name = kwargs.pop("name", None) or str(uuid.uuid4()) + role = Role.Sender + if "source_address" not in kwargs: + kwargs["source_address"] = "sender-link-{}".format(name) + super(SenderLink, self).__init__(session, handle, name, role, target_address=target_address, **kwargs) + self._pending_deliveries = [] + + @classmethod + def from_incoming_frame(cls, session, handle, frame): + # TODO: Assuming we establish all links for now... + # check link_create_from_endpoint in C lib + raise NotImplementedError("Pending") + + # In theory we should not need to purge pending deliveries on attach/dettach - as a link should + # be resume-able, however this is not yet supported. + def _incoming_attach(self, frame): + try: + super(SenderLink, self)._incoming_attach(frame) + except AMQPLinkError: + self._remove_pending_deliveries() + raise + self.current_link_credit = self.link_credit + self._outgoing_flow() + self.update_pending_deliveries() + + def _incoming_detach(self, frame): + super(SenderLink, self)._incoming_detach(frame) + self._remove_pending_deliveries() + + def _incoming_flow(self, frame): + rcv_link_credit = frame[6] # link_credit + rcv_delivery_count = frame[5] # delivery_count + if frame[4] is not None: # handle + if rcv_link_credit is None or rcv_delivery_count is None: + _LOGGER.info( + "Unable to get link-credit or delivery-count from incoming ATTACH. Detaching link.", + extra=self.network_trace_params + ) + self._remove_pending_deliveries() + self._set_state(LinkState.DETACHED) # TODO: Send detach now? + else: + self.current_link_credit = rcv_delivery_count + rcv_link_credit - self.delivery_count + self.update_pending_deliveries() + + def _outgoing_transfer(self, delivery): + output = bytearray() + encode_payload(output, delivery.message) + delivery_count = self.delivery_count + 1 + delivery.frame = { + "handle": self.handle, + "delivery_tag": struct.pack(">I", abs(delivery_count)), + "message_format": delivery.message._code, # pylint:disable=protected-access + "settled": delivery.settled, + "more": False, + "rcv_settle_mode": None, + "state": None, + "resume": None, + "aborted": None, + "batchable": None, + "payload": output, + } + self._session._outgoing_transfer( # pylint:disable=protected-access + delivery, + self.network_trace_params if self.network_trace else None + ) + sent_and_settled = False + if delivery.transfer_state == SessionTransferState.OKAY: + self.delivery_count = delivery_count + self.current_link_credit -= 1 + delivery.sent = True + if delivery.settled: + delivery.on_settled(LinkDeliverySettleReason.SETTLED, None) + sent_and_settled = True + # elif delivery.transfer_state == SessionTransferState.ERROR: + # TODO: Session wasn't mapped yet - re-adding to the outgoing delivery queue? + return sent_and_settled + + def _incoming_disposition(self, frame): + if not frame[3]: # settled + return + range_end = (frame[2] or frame[1]) + 1 # first or last + settled_ids = list(range(frame[1], range_end)) + unsettled = [] + for delivery in self._pending_deliveries: + if delivery.sent and delivery.frame["delivery_id"] in settled_ids: + delivery.on_settled(LinkDeliverySettleReason.DISPOSITION_RECEIVED, frame[4]) # state + continue + unsettled.append(delivery) + self._pending_deliveries = unsettled + + def _remove_pending_deliveries(self): + for delivery in self._pending_deliveries: + delivery.on_settled(LinkDeliverySettleReason.NOT_DELIVERED, None) + self._pending_deliveries = [] + + def _on_session_state_change(self): + if self._session.state == SessionState.DISCARDING: + self._remove_pending_deliveries() + super()._on_session_state_change() + + def update_pending_deliveries(self): + if self.current_link_credit <= 0: + self.current_link_credit = self.link_credit + self._outgoing_flow() + now = time.time() + pending = [] + for delivery in self._pending_deliveries: + if delivery.timeout and (now - delivery.start) >= delivery.timeout: + delivery.on_settled(LinkDeliverySettleReason.TIMEOUT, None) + continue + if not delivery.sent: + sent_and_settled = self._outgoing_transfer(delivery) + if sent_and_settled: + continue + pending.append(delivery) + self._pending_deliveries = pending + + def send_transfer(self, message, *, send_async=False, **kwargs): + self._check_if_closed() + if self.state != LinkState.ATTACHED: + raise AMQPLinkError( + condition=ErrorCondition.ClientError, + description="Link is not attached." + ) + settled = self.send_settle_mode == SenderSettleMode.Settled + if self.send_settle_mode == SenderSettleMode.Mixed: + settled = kwargs.pop("settled", True) + delivery = PendingDelivery( + on_delivery_settled=kwargs.get("on_send_complete"), + timeout=kwargs.get("timeout"), + message=message, + settled=settled, + network_trace_params = self.network_trace_params + ) + if self.current_link_credit == 0 or send_async: + self._pending_deliveries.append(delivery) + else: + sent_and_settled = self._outgoing_transfer(delivery) + if not sent_and_settled: + self._pending_deliveries.append(delivery) + return delivery + + def cancel_transfer(self, delivery): + try: + index = self._pending_deliveries.index(delivery) + except ValueError: + raise ValueError("Found no matching pending transfer.") + delivery = self._pending_deliveries[index] + if delivery.sent: + raise MessageException( + ErrorCondition.ClientError, + message="Transfer cannot be cancelled. Message has already been sent and awaiting disposition.", + ) + delivery.on_settled(LinkDeliverySettleReason.CANCELLED, None) + self._pending_deliveries.pop(index) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/session.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/session.py new file mode 100644 index 000000000000..3582b2e64e48 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/session.py @@ -0,0 +1,505 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + +from __future__ import annotations +import uuid +import logging +import time +from typing import Union, Optional + +from .constants import ConnectionState, SessionState, SessionTransferState, Role +from .sender import SenderLink +from .receiver import ReceiverLink +from .management_link import ManagementLink +from .performatives import ( + BeginFrame, + EndFrame, + FlowFrame, + TransferFrame, + DispositionFrame, +) +from .error import AMQPError, ErrorCondition +from ._encode import encode_frame + +_LOGGER = logging.getLogger(__name__) + + +class Session(object): # pylint: disable=too-many-instance-attributes + """ + :param int remote_channel: The remote channel for this Session. + :param int next_outgoing_id: The transfer-id of the first transfer id the sender will send. + :param int incoming_window: The initial incoming-window of the sender. + :param int outgoing_window: The initial outgoing-window of the sender. + :param int handle_max: The maximum handle value that may be used on the Session. + :param list(str) offered_capabilities: The extension capabilities the sender supports. + :param list(str) desired_capabilities: The extension capabilities the sender may use if the receiver supports + :param dict properties: Session properties. + """ + + def __init__(self, connection, channel, **kwargs): + self.name = kwargs.pop("name", None) or str(uuid.uuid4()) + self.state = SessionState.UNMAPPED + self.handle_max = kwargs.get("handle_max", 4294967295) + self.properties = kwargs.pop("properties", None) + self.channel = channel + self.remote_channel = None + self.next_outgoing_id = kwargs.pop("next_outgoing_id", 0) + self.next_incoming_id = None + self.incoming_window = kwargs.pop("incoming_window", 1) + self.outgoing_window = kwargs.pop("outgoing_window", 1) + self.target_incoming_window = self.incoming_window + self.remote_incoming_window = 0 + self.remote_outgoing_window = 0 + self.offered_capabilities = None + self.desired_capabilities = kwargs.pop("desired_capabilities", None) + + self.allow_pipelined_open = kwargs.pop("allow_pipelined_open", True) + self.idle_wait_time = kwargs.get("idle_wait_time", 0.1) + self.network_trace = kwargs["network_trace"] + self.network_trace_params = kwargs["network_trace_params"] + self.network_trace_params["amqpSession"] = self.name + + self.links = {} + self._connection = connection + self._output_handles = {} + self._input_handles = {} + + def __enter__(self): + self.begin() + return self + + def __exit__(self, *args): + self.end() + + @classmethod + def from_incoming_frame(cls, connection, channel): + # TODO: check session_create_from_endpoint in C lib + new_session = cls(connection, channel) + return new_session + + def _set_state(self, new_state): + # type: (SessionState) -> None + """Update the session state.""" + if new_state is None: + return + previous_state = self.state + self.state = new_state + _LOGGER.info( + "Session state changed: %r -> %r", + previous_state, + new_state, + extra=self.network_trace_params, + ) + for link in self.links.values(): + link._on_session_state_change() # pylint: disable=protected-access + + def _on_connection_state_change(self): + if self._connection.state in [ConnectionState.CLOSE_RCVD, ConnectionState.END]: + if self.state not in [SessionState.DISCARDING, SessionState.UNMAPPED]: + self._set_state(SessionState.DISCARDING) + + def _get_next_output_handle(self): + # type: () -> int + """Get the next available outgoing handle number within the max handle limit. + + :raises ValueError: If maximum handle has been reached. + :returns: The next available outgoing handle number. + :rtype: int + """ + if len(self._output_handles) >= self.handle_max: + raise ValueError( + "Maximum number of handles ({}) has been reached.".format( + self.handle_max + ) + ) + next_handle = next( + i for i in range(1, self.handle_max) if i not in self._output_handles + ) + return next_handle + + def _outgoing_begin(self): + begin_frame = BeginFrame( + remote_channel=self.remote_channel + if self.state == SessionState.BEGIN_RCVD + else None, + next_outgoing_id=self.next_outgoing_id, + outgoing_window=self.outgoing_window, + incoming_window=self.incoming_window, + handle_max=self.handle_max, + offered_capabilities=self.offered_capabilities + if self.state == SessionState.BEGIN_RCVD + else None, + desired_capabilities=self.desired_capabilities + if self.state == SessionState.UNMAPPED + else None, + properties=self.properties, + ) + if self.network_trace: + _LOGGER.debug("-> %r", begin_frame, extra=self.network_trace_params) + self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, begin_frame + ) + + def _incoming_begin(self, frame): + if self.network_trace: + _LOGGER.debug("<- %r", BeginFrame(*frame), extra=self.network_trace_params) + self.handle_max = frame[4] # handle_max + self.next_incoming_id = frame[1] # next_outgoing_id + self.remote_incoming_window = frame[2] # incoming_window + self.remote_outgoing_window = frame[3] # outgoing_window + if self.state == SessionState.BEGIN_SENT: + self.remote_channel = frame[0] # remote_channel + self._set_state(SessionState.MAPPED) + elif self.state == SessionState.UNMAPPED: + self._set_state(SessionState.BEGIN_RCVD) + self._outgoing_begin() + self._set_state(SessionState.MAPPED) + + def _outgoing_end(self, error=None): + end_frame = EndFrame(error=error) + if self.network_trace: + _LOGGER.debug("-> %r", end_frame, extra=self.network_trace_params) + self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, end_frame + ) + + def _incoming_end(self, frame): + if self.network_trace: + _LOGGER.debug("<- %r", EndFrame(*frame), extra=self.network_trace_params) + if self.state not in [ + SessionState.END_RCVD, + SessionState.END_SENT, + SessionState.DISCARDING, + ]: + self._set_state(SessionState.END_RCVD) + for _, link in self.links.items(): + link.detach() + # TODO: handling error + self._outgoing_end() + self._set_state(SessionState.UNMAPPED) + + def _outgoing_attach(self, frame): + self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, frame + ) + + def _incoming_attach(self, frame): + try: + self._input_handles[frame[1]] = self.links[ + frame[0].decode("utf-8") + ] # name and handle + self._input_handles[frame[1]]._incoming_attach( # pylint: disable=protected-access + frame + ) + except KeyError: + try: + outgoing_handle = self._get_next_output_handle() + except ValueError: + _LOGGER.error( + "Unable to attach new link - cannot allocate more handles.", + extra=self.network_trace_params + ) + # detach the link that would have been set. + self.links[frame[0].decode("utf-8")].detach( + error=AMQPError( + condition=ErrorCondition.LinkDetachForced, + description="""Cannot allocate more handles, """ + """the max number of handles is {}. Detaching link""".format( + self.handle_max + ), + info=None, + ) + ) + return + if frame[2] == Role.Sender: # role + new_link = ReceiverLink.from_incoming_frame( + self, outgoing_handle, frame + ) + else: + new_link = SenderLink.from_incoming_frame(self, outgoing_handle, frame) + new_link._incoming_attach(frame) # pylint: disable=protected-access + self.links[frame[0]] = new_link + self._output_handles[outgoing_handle] = new_link + self._input_handles[frame[1]] = new_link + except ValueError as e: + # Reject Link + _LOGGER.error( + "Unable to attach new link: %r", + e, + extra=self.network_trace_params + ) + self._input_handles[frame[1]].detach() + + def _outgoing_flow(self, frame=None): + link_flow = frame or {} + link_flow.update( + { + "next_incoming_id": self.next_incoming_id, + "incoming_window": self.incoming_window, + "next_outgoing_id": self.next_outgoing_id, + "outgoing_window": self.outgoing_window, + } + ) + flow_frame = FlowFrame(**link_flow) + if self.network_trace: + _LOGGER.debug("-> %r", flow_frame, extra=self.network_trace_params) + self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, flow_frame + ) + + def _incoming_flow(self, frame): + if self.network_trace: + _LOGGER.debug("<- %r", FlowFrame(*frame), extra=self.network_trace_params) + self.next_incoming_id = frame[2] # next_outgoing_id + remote_incoming_id = ( + frame[0] or self.next_outgoing_id + ) # next_incoming_id TODO "initial-outgoing-id" + self.remote_incoming_window = ( + remote_incoming_id + frame[1] - self.next_outgoing_id + ) # incoming_window + self.remote_outgoing_window = frame[3] # outgoing_window + if frame[4] is not None: # handle + self._input_handles[frame[4]]._incoming_flow( # pylint: disable=protected-access + frame + ) + else: + for link in self._output_handles.values(): + if ( + self.remote_incoming_window > 0 and not link._is_closed # pylint: disable=protected-access + ): + link._incoming_flow(frame) # pylint: disable=protected-access + + def _outgoing_transfer(self, delivery, network_trace_params): + if self.state != SessionState.MAPPED: + delivery.transfer_state = SessionTransferState.ERROR + if self.remote_incoming_window <= 0: + delivery.transfer_state = SessionTransferState.BUSY + else: + payload = delivery.frame["payload"] + payload_size = len(payload) + + delivery.frame["delivery_id"] = self.next_outgoing_id + # calculate the transfer frame encoding size excluding the payload + delivery.frame["payload"] = b"" + # TODO: encoding a frame would be expensive, we might want to improve depending on the perf test results + encoded_frame = encode_frame(TransferFrame(**delivery.frame))[1] + transfer_overhead_size = len(encoded_frame) + + # available size for payload per frame is calculated as following: + # remote max frame size - transfer overhead (calculated) - header (8 bytes) + available_frame_size = ( + self._connection._remote_max_frame_size - transfer_overhead_size - 8 # pylint: disable=protected-access + ) + + start_idx = 0 + remaining_payload_cnt = payload_size + # encode n-1 frames if payload_size > available_frame_size + while remaining_payload_cnt > available_frame_size: + tmp_delivery_frame = { + "handle": delivery.frame["handle"], + "delivery_tag": delivery.frame["delivery_tag"], + "message_format": delivery.frame["message_format"], + "settled": delivery.frame["settled"], + "more": True, + "rcv_settle_mode": delivery.frame["rcv_settle_mode"], + "state": delivery.frame["state"], + "resume": delivery.frame["resume"], + "aborted": delivery.frame["aborted"], + "batchable": delivery.frame["batchable"], + "delivery_id": self.next_outgoing_id, + } + if network_trace_params: + # We determine the logging for the outgoing Transfer frames based on the source + # Link configuration rather than the Session, because it's only at the Session + # level that we can determine how many outgoing frames are needed and their + # delivery IDs. + # TODO: Obscuring the payload for now to investigate the potential for leaks. + _LOGGER.debug( + "-> %r", TransferFrame(payload=b"***", **tmp_delivery_frame), + extra=network_trace_params + ) + self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, + TransferFrame( + payload=payload[start_idx : start_idx + available_frame_size], + **tmp_delivery_frame + ) + ) + start_idx += available_frame_size + remaining_payload_cnt -= available_frame_size + + # encode the last frame + tmp_delivery_frame = { + "handle": delivery.frame["handle"], + "delivery_tag": delivery.frame["delivery_tag"], + "message_format": delivery.frame["message_format"], + "settled": delivery.frame["settled"], + "more": False, + "rcv_settle_mode": delivery.frame["rcv_settle_mode"], + "state": delivery.frame["state"], + "resume": delivery.frame["resume"], + "aborted": delivery.frame["aborted"], + "batchable": delivery.frame["batchable"], + "delivery_id": self.next_outgoing_id, + } + if network_trace_params: + # We determine the logging for the outgoing Transfer frames based on the source + # Link configuration rather than the Session, because it's only at the Session + # level that we can determine how many outgoing frames are needed and their + # delivery IDs. + # TODO: Obscuring the payload for now to investigate the potential for leaks. + _LOGGER.debug( + "-> %r", TransferFrame(payload=b"***", **tmp_delivery_frame), + extra=network_trace_params + ) + self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, + TransferFrame(payload=payload[start_idx:], **tmp_delivery_frame) + ) + self.next_outgoing_id += 1 + self.remote_incoming_window -= 1 + self.outgoing_window -= 1 + # TODO: We should probably handle an error at the connection and update state accordingly + delivery.transfer_state = SessionTransferState.OKAY + + def _incoming_transfer(self, frame): + self.next_incoming_id += 1 + self.remote_outgoing_window -= 1 + self.incoming_window -= 1 + try: + self._input_handles[frame[0]]._incoming_transfer( # pylint: disable=protected-access + frame + ) + except KeyError: + _LOGGER.error( + "Received Transfer frame on unattached link. Ending session.", + extra=self.network_trace_params + ) + self._set_state(SessionState.DISCARDING) + self.end( + error=AMQPError( + condition=ErrorCondition.SessionUnattachedHandle, + description="""Invalid handle reference in received frame: """ + """Handle is not currently associated with an attached link""", + ) + ) + return + if self.incoming_window == 0: + self.incoming_window = self.target_incoming_window + self._outgoing_flow() + + def _outgoing_disposition(self, frame): + self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, frame + ) + + def _incoming_disposition(self, frame): + if self.network_trace: + _LOGGER.debug( + "<- %r", DispositionFrame(*frame), extra=self.network_trace_params + ) + for link in self._input_handles.values(): + link._incoming_disposition(frame) # pylint: disable=protected-access + + def _outgoing_detach(self, frame): + self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, frame + ) + + def _incoming_detach(self, frame): + try: + link = self._input_handles[frame[0]] # handle + link._incoming_detach(frame) # pylint: disable=protected-access + # if link._is_closed: TODO + # self.links.pop(link.name, None) + # self._input_handles.pop(link.remote_handle, None) + # self._output_handles.pop(link.handle, None) + except KeyError: + self._set_state(SessionState.DISCARDING) + self._connection.close( + error=AMQPError( + condition=ErrorCondition.SessionUnattachedHandle, + description="""Invalid handle reference in received frame: """ + """Handle is not currently associated with an attached link""", + ) + ) + + def _wait_for_response(self, wait, end_state): + # type: (Union[bool, float], SessionState) -> None + if wait is True: + self._connection.listen(wait=False) + while self.state != end_state: + time.sleep(self.idle_wait_time) + self._connection.listen(wait=False) + elif wait: + self._connection.listen(wait=False) + timeout = time.time() + wait + while self.state != end_state: + if time.time() >= timeout: + break + time.sleep(self.idle_wait_time) + self._connection.listen(wait=False) + + def begin(self, wait=False): + self._outgoing_begin() + self._set_state(SessionState.BEGIN_SENT) + if wait: + self._wait_for_response(wait, SessionState.BEGIN_SENT) + elif not self.allow_pipelined_open: + raise ValueError( + "Connection has been configured to not allow piplined-open. Please set 'wait' parameter." + ) + + def end(self, error=None, wait=False): + # type: (Optional[AMQPError], bool) -> None + try: + if self.state not in [SessionState.UNMAPPED, SessionState.DISCARDING]: + self._outgoing_end(error=error) + for _, link in self.links.items(): + link.detach() + new_state = SessionState.DISCARDING if error else SessionState.END_SENT + self._set_state(new_state) + self._wait_for_response(wait, SessionState.UNMAPPED) + except Exception as exc: # pylint: disable=broad-except + _LOGGER.info("An error occurred when ending the session: %r", exc, extra=self.network_trace_params) + self._set_state(SessionState.UNMAPPED) + + def create_receiver_link(self, source_address, **kwargs): + assigned_handle = self._get_next_output_handle() + link = ReceiverLink( + self, + handle=assigned_handle, + source_address=source_address, + network_trace=kwargs.pop("network_trace", self.network_trace), + network_trace_params=dict(self.network_trace_params), + **kwargs, + ) + self.links[link.name] = link + self._output_handles[assigned_handle] = link + return link + + def create_sender_link(self, target_address, **kwargs): + assigned_handle = self._get_next_output_handle() + link = SenderLink( + self, + handle=assigned_handle, + target_address=target_address, + network_trace=kwargs.pop("network_trace", self.network_trace), + network_trace_params=dict(self.network_trace_params), + **kwargs, + ) + self._output_handles[assigned_handle] = link + self.links[link.name] = link + return link + + def create_request_response_link_pair(self, endpoint, **kwargs): + return ManagementLink( + self, + endpoint, + network_trace=kwargs.pop("network_trace", self.network_trace), + network_trace_params=dict(self.network_trace_params), + **kwargs, + ) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/types.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/types.py new file mode 100644 index 000000000000..db478af591c8 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/types.py @@ -0,0 +1,90 @@ +#------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +#-------------------------------------------------------------------------- + +from enum import Enum + + +TYPE = 'TYPE' +VALUE = 'VALUE' + + +class AMQPTypes(object): # pylint: disable=no-init + null = 'NULL' + boolean = 'BOOL' + ubyte = 'UBYTE' + byte = 'BYTE' + ushort = 'USHORT' + short = 'SHORT' + uint = 'UINT' + int = 'INT' + ulong = 'ULONG' + long = 'LONG' + float = 'FLOAT' + double = 'DOUBLE' + timestamp = 'TIMESTAMP' + uuid = 'UUID' + binary = 'BINARY' + string = 'STRING' + symbol = 'SYMBOL' + list = 'LIST' + map = 'MAP' + array = 'ARRAY' + described = 'DESCRIBED' + + +class FieldDefinition(Enum): + fields = "fields" + annotations = "annotations" + message_id = "message-id" + app_properties = "application-properties" + node_properties = "node-properties" + filter_set = "filter-set" + + +class ObjDefinition(Enum): + source = "source" + target = "target" + delivery_state = "delivery-state" + error = "error" + + +class ConstructorBytes(object): # pylint: disable=no-init + null = b'\x40' + bool = b'\x56' + bool_true = b'\x41' + bool_false = b'\x42' + ubyte = b'\x50' + byte = b'\x51' + ushort = b'\x60' + short = b'\x61' + uint_0 = b'\x43' + uint_small = b'\x52' + int_small = b'\x54' + uint_large = b'\x70' + int_large = b'\x71' + ulong_0 = b'\x44' + ulong_small = b'\x53' + long_small = b'\x55' + ulong_large = b'\x80' + long_large = b'\x81' + float = b'\x72' + double = b'\x82' + timestamp = b'\x83' + uuid = b'\x98' + binary_small = b'\xA0' + binary_large = b'\xB0' + string_small = b'\xA1' + string_large = b'\xB1' + symbol_small = b'\xA3' + symbol_large = b'\xB3' + list_0 = b'\x45' + list_small = b'\xC0' + list_large = b'\xD0' + map_small = b'\xC1' + map_large = b'\xD1' + array_small = b'\xE0' + array_large = b'\xF0' + descriptor = b'\x00' diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/utils.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/utils.py new file mode 100644 index 000000000000..5baf13992f44 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/utils.py @@ -0,0 +1,139 @@ +#------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +#-------------------------------------------------------------------------- +import datetime +from base64 import b64encode +from hashlib import sha256 +from hmac import HMAC +from urllib.parse import urlencode, quote_plus +import time +import six + +from .types import TYPE, VALUE, AMQPTypes +from ._encode import encode_payload + + +class UTC(datetime.tzinfo): + """Time Zone info for handling UTC""" + + def utcoffset(self, dt): + """UTF offset for UTC is 0.""" + return datetime.timedelta(0) + + def tzname(self, dt): + """Timestamp representation.""" + return "Z" + + def dst(self, dt): + """No daylight saving for UTC.""" + return datetime.timedelta(hours=1) + + +try: + from datetime import timezone # pylint: disable=ungrouped-imports + + TZ_UTC = timezone.utc # type: ignore +except ImportError: + TZ_UTC = UTC() # type: ignore + + +def utc_from_timestamp(timestamp): + return datetime.datetime.fromtimestamp(timestamp, tz=TZ_UTC) + + +def utc_now(): + return datetime.datetime.now(tz=TZ_UTC) + + +def encode(value, encoding='UTF-8'): + return value.encode(encoding) if isinstance(value, six.text_type) else value + + +def generate_sas_token(audience, policy, key, expiry=None): + """ + Generate a sas token according to the given audience, policy, key and expiry + + :param str audience: + :param str policy: + :param str key: + :param int expiry: abs expiry time + :rtype: str + """ + if not expiry: + expiry = int(time.time()) + 3600 # Default to 1 hour. + + encoded_uri = quote_plus(audience) + encoded_policy = quote_plus(policy).encode("utf-8") + encoded_key = key.encode("utf-8") + + ttl = int(expiry) + sign_key = '%s\n%d' % (encoded_uri, ttl) + signature = b64encode(HMAC(encoded_key, sign_key.encode('utf-8'), sha256).digest()) + result = { + 'sr': audience, + 'sig': signature, + 'se': str(ttl) + } + if policy: + result['skn'] = encoded_policy + return 'SharedAccessSignature ' + urlencode(result) + + +def add_batch(batch, message): + # Add a message to a batch + output = bytearray() + encode_payload(output, message) + batch[5].append(output) + + +def encode_str(data, encoding='utf-8'): + try: + return data.encode(encoding) + except AttributeError: + return data + + +def normalized_data_body(data, **kwargs): + # A helper method to normalize input into AMQP Data Body format + encoding = kwargs.get("encoding", "utf-8") + if isinstance(data, list): + return [encode_str(item, encoding) for item in data] + return [encode_str(data, encoding)] + + +def normalized_sequence_body(sequence): + # A helper method to normalize input into AMQP Sequence Body format + if isinstance(sequence, list) and all([isinstance(b, list) for b in sequence]): + return sequence + if isinstance(sequence, list): + return [sequence] + + +def get_message_encoded_size(message): + output = bytearray() + encode_payload(output, message) + return len(output) + + +def amqp_long_value(value): + # A helper method to wrap a Python int as AMQP long + # TODO: wrapping one line in a function is expensive, find if there's a better way to do it + return {TYPE: AMQPTypes.long, VALUE: value} + + +def amqp_uint_value(value): + # A helper method to wrap a Python int as AMQP uint + return {TYPE: AMQPTypes.uint, VALUE: value} + + +def amqp_string_value(value): + return {TYPE: AMQPTypes.string, VALUE: value} + + +def amqp_symbol_value(value): + return {TYPE: AMQPTypes.symbol, VALUE: value} + +def amqp_array_value(value): + return {TYPE: AMQPTypes.array, VALUE: value} diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_transport/_base.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_transport/_base.py index d67cceedcd40..84db0d7c13e1 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_transport/_base.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_transport/_base.py @@ -7,7 +7,10 @@ from abc import ABC, abstractmethod if TYPE_CHECKING: - from uamqp import types as uamqp_types + try: + from uamqp import types as uamqp_types + except ImportError: + uamqp_types = None class AmqpTransport(ABC): # pylint: disable=too-many-public-methods """ @@ -163,10 +166,10 @@ def set_message_partition_key(message, partition_key, **kwargs): @staticmethod @abstractmethod - def add_batch(batch_message, outgoing_event_data, event_data): + def add_batch(event_data_batch, outgoing_event_data, event_data): """ Add EventData to the data body of the BatchMessage. - :param batch_message: BatchMessage to add data to. + :param event_data_batch: BatchMessage to add data to. :param outgoing_event_data: Transformed EventData for sending. :param event_data: EventData to add to internal batch events. uamqp use only. :rtype: None diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_transport/_pyamqp_transport.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_transport/_pyamqp_transport.py new file mode 100644 index 000000000000..2983ec920fca --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_transport/_pyamqp_transport.py @@ -0,0 +1,616 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +import logging +import time +from typing import Optional, Union, Any, Tuple, cast + +from .._pyamqp import ( + error as errors, + utils, + SendClient, + constants, + AMQPClient, + ReceiveClient, +) +from .._pyamqp.message import Message, BatchMessage, Header, Properties +from .._pyamqp.authentication import JWTTokenAuth +from .._pyamqp.endpoints import Source, ApacheFilters +from .._pyamqp._connection import Connection, _CLOSING_STATES + +from ._base import AmqpTransport +from .._constants import ( + NO_RETRY_ERRORS, + PROP_PARTITION_KEY, + CUSTOM_CONDITION_BACKOFF, +) + +from ..exceptions import ( + ConnectError, + EventHubError, + AuthenticationError, + ConnectionLostError, + EventDataSendError, + OperationTimeoutError +) + +_LOGGER = logging.getLogger(__name__) + + +class PyamqpTransport(AmqpTransport): # pylint: disable=too-many-public-methods + """ + Class which defines uamqp-based methods used by the producer and consumer. + """ + + # define constants + MAX_FRAME_SIZE_BYTES = constants.MAX_FRAME_SIZE_BYTES + MAX_MESSAGE_LENGTH_BYTES = ( + constants.MAX_FRAME_SIZE_BYTES + ) # TODO: define actual value in pyamqp + TIMEOUT_FACTOR = 1 + CONNECTION_CLOSING_STATES: Tuple = _CLOSING_STATES + + # define symbols + PRODUCT_SYMBOL = "product" + VERSION_SYMBOL = "version" + FRAMEWORK_SYMBOL = "framework" + PLATFORM_SYMBOL = "platform" + USER_AGENT_SYMBOL = "user-agent" + PROP_PARTITION_KEY_AMQP_SYMBOL = PROP_PARTITION_KEY + + ERROR_CONDITIONS = [condition.value for condition in errors.ErrorCondition] + + @staticmethod + def build_message(**kwargs): + """ + Creates a pyamqp.Message with given arguments. + :rtype: pyamqp.Message + """ + return Message(**kwargs) + + @staticmethod + def build_batch_message(**kwargs): + """ + Creates a pyamqp.BatchMessage with given arguments. + :rtype: pyamqp.BatchMessage + """ + return BatchMessage(**kwargs) + + @staticmethod + def to_outgoing_amqp_message(annotated_message): + """ + Converts an AmqpAnnotatedMessage into an Amqp Message. + :param AmqpAnnotatedMessage annotated_message: AmqpAnnotatedMessage to convert. + :rtype: pyamqp.Message + """ + message_header = None + header_vals = annotated_message.header.values() if annotated_message.header else None + # If header and non-None header values, create outgoing header. + if annotated_message.header and header_vals.count(None) != len(header_vals): + message_header = Header( + delivery_count=annotated_message.header.delivery_count, + ttl=annotated_message.header.time_to_live, + first_acquirer=annotated_message.header.first_acquirer, + durable=annotated_message.header.durable, + priority=annotated_message.header.priority, + ) + + message_properties = None + properties_vals = annotated_message.properties.values() if annotated_message.properties else None + # If properties and non-None properties values, create outgoing properties. + if annotated_message.properties and properties_vals.count(None) != len(properties_vals): + message_properties = Properties( + message_id=annotated_message.properties.message_id, + user_id=annotated_message.properties.user_id, + to=annotated_message.properties.to, + subject=annotated_message.properties.subject, + reply_to=annotated_message.properties.reply_to, + correlation_id=annotated_message.properties.correlation_id, + content_type=annotated_message.properties.content_type, + content_encoding=annotated_message.properties.content_encoding, + creation_time=int(annotated_message.properties.creation_time) + if annotated_message.properties.creation_time + else None, + absolute_expiry_time=int( + annotated_message.properties.absolute_expiry_time + ) + if annotated_message.properties.absolute_expiry_time + else None, + group_id=annotated_message.properties.group_id, + group_sequence=annotated_message.properties.group_sequence, + reply_to_group_id=annotated_message.properties.reply_to_group_id, + ) + + message_dict = { + "header": message_header, + "properties": message_properties, + "application_properties": annotated_message.application_properties, + "message_annotations": annotated_message.annotations, + "delivery_annotations": annotated_message.delivery_annotations, + "data": annotated_message._data_body, # pylint: disable=protected-access + "sequence": annotated_message._sequence_body, # pylint: disable=protected-access + "value": annotated_message._value_body, # pylint: disable=protected-access + "footer": annotated_message.footer, + } + + return Message(**message_dict) + + @staticmethod + def get_batch_message_encoded_size(message): + """ + Gets the batch message encoded size given an underlying Message. + :param pyamqp.BatchMessage message: Message to get encoded size of. + :rtype: int + """ + return utils.get_message_encoded_size(message) + + @staticmethod + def get_message_encoded_size(message): + """ + Gets the message encoded size given an underlying Message. + :param pyamqp.Message: Message to get encoded size of. + :rtype: int + """ + return utils.get_message_encoded_size(message) + + @staticmethod + def get_remote_max_message_size(handler): + """ + Returns max peer message size. + :param AMQPClient handler: Client to get remote max message size on link from. + :rtype: int + """ + return handler._link.remote_max_message_size # pylint: disable=protected-access + + @staticmethod + def create_retry_policy(config): + """ + Creates the error retry policy. + :param ~azure.eventhub._configuration.Configuration config: Configuration. + """ + return errors.RetryPolicy( + retry_total=config.max_retries, # pylint:disable=protected-access + retry_backoff_factor=config.backoff_factor, # pylint:disable=protected-access + retry_backoff_max=config.backoff_max, # pylint:disable=protected-access + retry_mode=config.retry_mode, # pylint:disable=protected-access + no_retry_condition=NO_RETRY_ERRORS, + custom_condition_backoff=CUSTOM_CONDITION_BACKOFF, + ) + + @staticmethod + def create_link_properties(link_properties): + """ + Creates and returns the link properties. + :param dict[bytes, int] link_properties: The dict of symbols and corresponding values. + :rtype: dict + """ + return { + symbol: utils.amqp_long_value(value) + for (symbol, value) in link_properties.items() + } + + @staticmethod + def create_connection(**kwargs): + """ + Creates and returns the uamqp Connection object. + :keyword str host: The hostname, used by uamqp. + :keyword JWTTokenAuth auth: The auth, used by uamqp. + :keyword str endpoint: The endpoint, used by pyamqp. + :keyword str container_id: Required. + :keyword int max_frame_size: Required. + :keyword int channel_max: Required. + :keyword int idle_timeout: Required. + :keyword Dict properties: Required. + :keyword int remote_idle_timeout_empty_frame_send_ratio: Required. + :keyword error_policy: Required. + :keyword bool debug: Required. + :keyword str encoding: Required. + """ + endpoint = kwargs.pop("endpoint") + host = kwargs.pop("host") # pylint:disable=unused-variable + auth = kwargs.pop("auth") # pylint:disable=unused-variable + network_trace = kwargs.pop("debug") + return Connection(endpoint, network_trace=network_trace, **kwargs) + + @staticmethod + def close_connection(connection): + """ + Closes existing connection. + :param connection: uamqp or pyamqp Connection. + """ + connection.close() + + @staticmethod + def get_connection_state(connection): + """ + Gets connection state. + :param connection: uamqp or pyamqp Connection. + """ + return connection.state + + @staticmethod + def create_send_client(*, config, **kwargs): # pylint:disable=unused-argument + """ + Creates and returns the uamqp SendClient. + :param ~azure.eventhub._configuration.Configuration config: The configuration. + + :keyword str target: Required. The target. + :keyword JWTTokenAuth auth: Required. + :keyword int idle_timeout: Required. + :keyword network_trace: Required. + :keyword retry_policy: Required. + :keyword keep_alive_interval: Required. + :keyword str client_name: Required. + :keyword dict link_properties: Required. + :keyword properties: Required. + """ + + target = kwargs.pop("target") + # TODO: not used by pyamqp? + msg_timeout = kwargs.pop( # pylint: disable=unused-variable + "msg_timeout" + ) + + return SendClient( + config.hostname, + target, + custom_endpoint_address=config.custom_endpoint_address, + connection_verify=config.connection_verify, + transport_type=config.transport_type, + http_proxy=config.http_proxy, + **kwargs, + ) + + @staticmethod + def send_messages(producer, timeout_time, last_exception, logger): + """ + Handles sending of event data messages. + :param ~azure.eventhub._producer.EventHubProducer producer: The producer with handler to send messages. + :param int timeout_time: Timeout time. + :param last_exception: Exception to raise if message timed out. Only used by uamqp transport. + :param logger: Logger. + """ + # pylint: disable=protected-access + try: + producer._open() + timeout = timeout_time - time.time() if timeout_time else 0 + producer._handler.send_message( + producer._unsent_events[0], timeout=timeout + ) + # TODO: The unsent_events list will always be <= 1. Even for a batch, + # it gets the underlying singular BatchMessage. + # May want to refactor in the future so that this isn't a list. + producer._unsent_events = None + except TimeoutError as exc: + raise OperationTimeoutError(message=str(exc), details=exc) + + @staticmethod + def set_message_partition_key(message, partition_key, **kwargs): + # type: (Message, Optional[Union[bytes, str]], Any) -> Message + """Set the partition key as an annotation on a uamqp message. + :param Message message: The message to update. + :param str partition_key: The partition key value. + :rtype: Message + """ + encoding = kwargs.pop("encoding", "utf-8") + if partition_key: + annotations = message.message_annotations + if annotations is None: + annotations = {} + try: + partition_key = cast(bytes, partition_key).decode(encoding) + except AttributeError: + pass + annotations[ + PROP_PARTITION_KEY + ] = partition_key # pylint:disable=protected-access + header = Header(durable=True) # type: ignore + return message._replace(message_annotations=annotations, header=header) + return message + + @staticmethod + def add_batch( + event_data_batch, outgoing_event_data, event_data + ): # pylint: disable=unused-argument + """ + Add EventData to the data body of the BatchMessage. + :param event_data_batch: EventDataBatch to add data to. + :param outgoing_event_data: Transformed EventData for sending. + :param event_data: EventData to add to internal batch events. uamqp use only. + :rtype: None + """ + event_data_batch._internal_events.append( # pylint: disable=protected-access + event_data + ) + # pylint: disable=protected-access + utils.add_batch( + event_data_batch._message, outgoing_event_data._message + ) + + @staticmethod + def create_source(source, offset, selector): + """ + Creates and returns the Source. + + :param str source: Required. + :param int offset: Required. + :param bytes selector: Required. + """ + source = Source(address=source, filters={}) + if offset is not None: + filter_key = ApacheFilters.selector_filter + source.filters[filter_key] = (filter_key, utils.amqp_string_value(selector)) + return source + + @staticmethod + def create_receive_client(*, config, **kwargs): + """ + Creates and returns the receive client. + :param ~azure.eventhub._configuration.Configuration config: The configuration. + + :keyword str source: Required. The source. + :keyword str offset: Required. + :keyword str offset_inclusive: Required. + :keyword JWTTokenAuth auth: Required. + :keyword int idle_timeout: Required. + :keyword network_trace: Required. + :keyword retry_policy: Required. + :keyword str client_name: Required. + :keyword dict link_properties: Required. + :keyword properties: Required. + :keyword link_credit: Required. The prefetch. + :keyword keep_alive_interval: Required. Missing in pyamqp. + :keyword desired_capabilities: Required. + :keyword streaming_receive: Required. + :keyword message_received_callback: Required. + :keyword timeout: Required. + """ + + source = kwargs.pop("source") + return ReceiveClient( + config.hostname, + source, + receive_settle_mode=constants.ReceiverSettleMode.First, + http_proxy=config.http_proxy, + transport_type=config.transport_type, + custom_endpoint_address=config.custom_endpoint_address, + connection_verify=config.connection_verify, + **kwargs, + ) + + @staticmethod + def open_receive_client(*, handler, client, auth): + """ + Opens the receive client and returns ready status. + :param ReceiveClient handler: The receive client. + :param ~azure.eventhub.EventHubConsumerClient client: The consumer client. + :param auth: Auth. + :rtype: bool + """ + # pylint:disable=protected-access + handler.open( + connection=client._conn_manager.get_connection( + client._address.hostname, auth + ) + ) + + @staticmethod + def check_link_stolen(consumer, exception): + """ + Checks if link stolen and handles exception. + :param consumer: The EventHubConsumer. + :param exception: Exception to check. + """ + + if ( + isinstance(exception, errors.AMQPLinkError) + and exception.condition == errors.ErrorCondition.LinkStolen + ): + raise consumer._handle_exception( # pylint: disable=protected-access + exception + ) + + @staticmethod + def create_token_auth(auth_uri, get_token, token_type, config, **kwargs): + """ + Creates the JWTTokenAuth. + :param str auth_uri: The auth uri to pass to JWTTokenAuth. + :param get_token: The callback function used for getting and refreshing + tokens. It should return a valid jwt token each time it is called. + :param bytes token_type: Token type. + :param ~azure.eventhub._configuration.Configuration config: EH config. + + :keyword bool update_token: Whether to update token. If not updating token, then pass 300 to refresh_window. + """ + # TODO: figure out why we're passing all these args to pyamqp JWTTokenAuth, which aren't being used + update_token = kwargs.pop("update_token") # pylint: disable=unused-variable + if update_token: + # update_token not actually needed by pyamqp + # just using to detect wh + return JWTTokenAuth(auth_uri, auth_uri, get_token) + return JWTTokenAuth( + auth_uri, + auth_uri, + get_token, + token_type=token_type, + timeout=config.auth_timeout, + custom_endpoint_hostname=config.custom_endpoint_hostname, + port=config.connection_port, + verify=config.connection_verify, + ) + # if update_token: + # token_auth.update_token() # TODO: why don't we need to update in pyamqp? + + @staticmethod + def create_mgmt_client( + address, mgmt_auth, config + ): # pylint: disable=unused-argument + """ + Creates and returns the mgmt AMQP client. + :param _Address address: Required. The Address. + :param JWTTokenAuth mgmt_auth: Auth for client. + :param ~azure.eventhub._configuration.Configuration config: The configuration. + """ + + return AMQPClient( + config.hostname, + auth=mgmt_auth, + network_trace=config.network_tracing, + transport_type=config.transport_type, + http_proxy=config.http_proxy, + custom_endpoint_address=config.custom_endpoint_address, + connection_verify=config.connection_verify, + ) + + @staticmethod + def get_updated_token(mgmt_auth): + """ + Return updated auth token. + :param mgmt_auth: Auth. + """ + return mgmt_auth.get_token() + + @staticmethod + def mgmt_client_request(mgmt_client, mgmt_msg, **kwargs): + """ + Send mgmt request. + :param AMQPClient mgmt_client: Client to send request with. + :param str mgmt_msg: Message. + :keyword bytes operation: Operation. + :keyword operation_type: Op type. + :keyword status_code_field: mgmt status code. + :keyword description_fields: mgmt status desc. + """ + operation_type = kwargs.pop("operation_type") + operation = kwargs.pop("operation") + return mgmt_client.mgmt_request( + mgmt_msg, + operation=operation.decode(), + operation_type=operation_type.decode(), + **kwargs, + ) + + @staticmethod + def get_error(status_code, description): + """ + Gets error and passes in error message, and, if applicable, condition. + :param error: The error to raise. + :param str message: Error message. + :param condition: Optional error condition. Will not be used by uamqp. + """ + if status_code in [401]: + return errors.AuthenticationException( + errors.ErrorCondition.UnauthorizedAccess, + description=f"""Management authentication failed. Status code: {status_code}, """ + """Description: {description!r}""", + ) + if status_code in [404]: + return errors.AMQPConnectionError( + errors.ErrorCondition.NotFound, + description=f"Management connection failed. Status code: {status_code}, Description: {description!r}", + ) + return errors.AMQPConnectionError( + errors.ErrorCondition.UnknownError, + description=f"Management request error. Status code: {status_code}, Description: {description!r}", + ) + + @staticmethod + def check_timeout_exception(base, exception): + """ + Checks if timeout exception. + :param base: ClientBase. + :param exception: Exception to check. + """ + if not base.running and isinstance(exception, TimeoutError): + exception = errors.AuthenticationException( + errors.ErrorCondition.InternalError, + description="Authorization timeout.", + ) + return exception + + @staticmethod + def _create_eventhub_exception(exception, *, is_consumer=False): + if isinstance(exception, errors.AuthenticationException): + error = AuthenticationError(str(exception), exception) + elif isinstance(exception, errors.AMQPLinkError): + # For uamqp exception parity, raising ConnectionLostError for LinkDetaches. + # Else, vendor error condition that starts with "com.", so raise ConnectError. + if exception.condition in PyamqpTransport.ERROR_CONDITIONS: + error = ConnectionLostError(str(exception), exception) + else: + error = ConnectError(str(exception), exception) + # TODO: do we need MessageHandlerError in amqp any more + # if connection/session/link error are enough? + # elif isinstance(exception, errors.MessageHandlerError): + # error = ConnectionLostError(str(exception), exception) + elif isinstance(exception, errors.AMQPConnectionError): + error = ConnectError(str(exception), exception) + elif isinstance(exception, TimeoutError): + error = ConnectionLostError(str(exception), exception) + else: + if ( + isinstance(exception, FileNotFoundError) + and is_consumer + and exception.filename + and "ca_certs" in exception.filename + ): + + error = exception + else: + error = EventHubError(str(exception), exception) + return error + + @staticmethod + def _handle_exception( + exception, closable, *, is_consumer=False + ): # pylint:disable=too-many-branches, too-many-statements + try: # closable is a producer/consumer object + name = closable._name # pylint: disable=protected-access + except AttributeError: # closable is an client object + name = closable._container_id # pylint: disable=protected-access + if isinstance(exception, KeyboardInterrupt): # pylint:disable=no-else-raise + _LOGGER.info("%r stops due to keyboard interrupt", name) + closable._close_connection() # pylint:disable=protected-access + raise exception + elif isinstance(exception, EventHubError): + closable._close_handler() # pylint:disable=protected-access + raise exception + # TODO: The following errors seem to be useless in EH + # elif isinstance( + # exception, + # ( + # errors.MessageAccepted, + # errors.MessageAlreadySettled, + # errors.MessageModified, + # errors.MessageRejected, + # errors.MessageReleased, + # errors.MessageContentTooLarge, + # ), + # ): + # _LOGGER.info("%r Event data error (%r)", name, exception) + # error = EventDataError(str(exception), exception) + # raise error + elif isinstance(exception, errors.MessageException): + _LOGGER.info("%r Event data send error (%r)", name, exception) + error = EventDataSendError(str(exception), exception) + raise error + else: + if isinstance(exception, errors.AuthenticationException): + if hasattr(closable, "_close_connection"): + closable._close_connection() # pylint:disable=protected-access + elif isinstance(exception, errors.AMQPLinkError): + if hasattr(closable, "_close_handler"): + closable._close_handler() # pylint:disable=protected-access + elif isinstance(exception, errors.AMQPConnectionError): + if hasattr(closable, "_close_connection"): + closable._close_connection() # pylint:disable=protected-access + # TODO: add MessageHandlerError in amqp? + # elif isinstance(exception, errors.MessageHandlerError): + # if hasattr(closable, "_close_handler"): + # closable._close_handler() # pylint:disable=protected-access + else: # errors.AMQPConnectionError, compat.TimeoutException + if hasattr(closable, "_close_connection"): + closable._close_connection() # pylint:disable=protected-access + return PyamqpTransport._create_eventhub_exception(exception, is_consumer=is_consumer) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_transport/_uamqp_transport.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_transport/_uamqp_transport.py index 018d3611aa72..417320ef4d24 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_transport/_uamqp_transport.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_transport/_uamqp_transport.py @@ -103,16 +103,16 @@ class UamqpTransport(AmqpTransport): # pylint: disable=too-many-public-method @staticmethod def build_message(**kwargs): """ - Creates a uamqp.Message or pyamqp.Message with given arguments. - :rtype: uamqp.Message or pyamqp.Message + Creates a uamqp.Message with given arguments. + :rtype: uamqp.Message """ return Message(**kwargs) @staticmethod def build_batch_message(**kwargs): """ - Creates a uamqp.BatchMessage or pyamqp.BatchMessage with given arguments. - :rtype: uamqp.BatchMessage or pyamqp.BatchMessage + Creates a uamqp.BatchMessage with given arguments. + :rtype: uamqp.BatchMessage """ return BatchMessage(**kwargs) @@ -124,7 +124,9 @@ def to_outgoing_amqp_message(annotated_message): :rtype: uamqp.Message """ message_header = None - if annotated_message.header: + header_vals = annotated_message.header.values() if annotated_message.header else None + # If header and non-None header values, create outgoing header. + if annotated_message.header and header_vals.count(None) != len(header_vals): message_header = MessageHeader() message_header.delivery_count = annotated_message.header.delivery_count message_header.time_to_live = annotated_message.header.time_to_live @@ -133,7 +135,9 @@ def to_outgoing_amqp_message(annotated_message): message_header.priority = annotated_message.header.priority message_properties = None - if annotated_message.properties: + properties_vals = annotated_message.properties.values() if annotated_message.properties else None + # If properties and non-None properties values, create outgoing properties. + if annotated_message.properties and properties_vals.count(None) != len(properties_vals): message_properties = MessageProperties( message_id=annotated_message.properties.message_id, user_id=annotated_message.properties.user_id, @@ -238,6 +242,7 @@ def create_connection(**kwargs): :keyword str encoding: Required. """ endpoint = kwargs.pop("endpoint") # pylint:disable=unused-variable + custom_endpoint_address = kwargs.pop("custom_endpoint_address") # pylint:disable=unused-variable host = kwargs.pop("host") auth = kwargs.pop("auth") return Connection( @@ -348,17 +353,17 @@ def set_message_partition_key(message, partition_key, **kwargs): # pylint:disab return message @staticmethod - def add_batch(batch_message, outgoing_event_data, event_data): + def add_batch(event_data_batch, outgoing_event_data, event_data): """ Add EventData to the data body of the BatchMessage. - :param batch_message: BatchMessage to add data to. + :param event_data_batch: BatchMessage to add data to. :param outgoing_event_data: Transformed EventData for sending. :param event_data: EventData to add to internal batch events. uamqp use only. :rtype: None """ # pylint: disable=protected-access - batch_message._internal_events.append(event_data) - batch_message._message._body_gen.append( + event_data_batch._internal_events.append(event_data) + event_data_batch._message._body_gen.append( outgoing_event_data._message ) @@ -505,6 +510,15 @@ def create_mgmt_client(address, mgmt_auth, config): debug=config.network_tracing ) + @staticmethod + def open_mgmt_client(mgmt_client, conn): + """ + Opens the mgmt AMQP client. + :param AMQPClient mgmt_client: uamqp AMQPClient. + :param conn: Connection. + """ + mgmt_client.open(connection=conn) + @staticmethod def get_updated_token(mgmt_auth): """ @@ -526,12 +540,17 @@ def mgmt_client_request(mgmt_client, mgmt_msg, **kwargs): """ operation_type = kwargs.pop("operation_type") operation = kwargs.pop("operation") - return mgmt_client.mgmt_request( + response = mgmt_client.mgmt_request( mgmt_msg, operation, op_type=operation_type, **kwargs ) + status_code = response.application_properties[kwargs.get("status_code_field")] + description = response.application_properties.get( + kwargs.get("description_fields") + ) # type: Optional[Union[str, bytes]] + return status_code, description, response @staticmethod def get_error(status_code, description): @@ -562,8 +581,7 @@ def check_timeout_exception(base, exception): if not base.running and isinstance( exception, compat.TimeoutException ): - exception = UamqpTransport.get_error( - errors.AuthenticationException, + exception = errors.AuthenticationException( "Authorization timeout." ) return exception @@ -595,7 +613,7 @@ def _create_eventhub_exception(exception): @staticmethod def _handle_exception( - exception, closable + exception, closable, *, is_consumer=False # pylint:disable=unused-argument ): # pylint:disable=too-many-branches, too-many-statements try: # closable is a producer/consumer object name = closable._name # pylint: disable=protected-access diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_utils.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_utils.py index fc9b1bc8c9c9..e15f0b884732 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_utils.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_utils.py @@ -10,11 +10,6 @@ import datetime import calendar import logging -from base64 import b64encode -from hashlib import sha256 -from hmac import HMAC -from urllib.parse import urlencode, quote_plus -import time from typing import ( TYPE_CHECKING, cast, @@ -30,7 +25,6 @@ ) import six -from uamqp import types as uamqp_types from azure.core.settings import settings from azure.core.tracing import SpanKind, Link @@ -51,6 +45,11 @@ if TYPE_CHECKING: # pylint: disable=ungrouped-imports from ._transport._base import AmqpTransport + try: + from uamqp import types as uamqp_types + except ImportError: + uamqp_types = None + from ._pyamqp import types from azure.core.tracing import AbstractSpan from azure.core.credentials import AzureSasCredential from ._common import EventData @@ -94,7 +93,7 @@ def utc_from_timestamp(timestamp): def create_properties( user_agent: Optional[str] = None, *, amqp_transport: AmqpTransport -) -> Dict[uamqp_types.AMQPSymbol, str]: +) -> Union[Dict[uamqp_types.AMQPSymbol, str], Dict[str, str]]: """ Format the properties with which to instantiate the connection. This acts like a user agent over HTTP. @@ -345,32 +344,3 @@ def decode_with_recurse(data, encoding="UTF-8"): return decoded_list return data - - -def generate_sas_token(audience, policy, key, expiry=None): - """ - Generate a sas token according to the given audience, policy, key and expiry - :param str audience: - :param str policy: - :param str key: - :param int expiry: abs expiry time - :rtype: str - """ - if not expiry: - expiry = int(time.time()) + 3600 # Default to 1 hour. - - encoded_uri = quote_plus(audience) - encoded_policy = quote_plus(policy).encode("utf-8") - encoded_key = key.encode("utf-8") - - ttl = int(expiry) - sign_key = '%s\n%d' % (encoded_uri, ttl) - signature = b64encode(HMAC(encoded_key, sign_key.encode('utf-8'), sha256).digest()) - result = { - 'sr': audience, - 'sig': signature, - 'se': str(ttl) - } - if policy: - result['skn'] = encoded_policy - return 'SharedAccessSignature ' + urlencode(result) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_version.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_version.py index d3da8af8ca9a..7d861beb714e 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_version.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_version.py @@ -3,4 +3,4 @@ # Licensed under the MIT License. # ------------------------------------ -VERSION = "5.10.2" +VERSION = "5.11.0" diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_buffered_producer/_buffered_producer_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_buffered_producer/_buffered_producer_async.py index dbbee8f09879..0888aca876e5 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_buffered_producer/_buffered_producer_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_buffered_producer/_buffered_producer_async.py @@ -33,8 +33,9 @@ def __init__( ], max_message_size_on_link: int, *, - max_wait_time: float = 1, - max_buffer_length: int + amqp_transport: AmqpTransportAsync, + max_buffer_length: int, + max_wait_time: float = 1 ): self._buffered_queue: queue.Queue = queue.Queue() self._max_buffer_len = max_buffer_length @@ -50,10 +51,11 @@ def __init__( self._max_message_size_on_link = max_message_size_on_link self._check_max_wait_time_future = None self.partition_id = partition_id + self._amqp_transport = amqp_transport async def start(self): async with self._lock: - self._cur_batch = EventDataBatch(self._max_message_size_on_link) + self._cur_batch = EventDataBatch(self._max_message_size_on_link, amqp_transport=self._amqp_transport) self._running = True if self._max_wait_time: self._last_send_time = time.time() @@ -113,12 +115,12 @@ async def put_events(self, events, timeout_time=None): self._buffered_queue.put(self._cur_batch) self._buffered_queue.put(events) # create a new batch for incoming events - self._cur_batch = EventDataBatch(self._max_message_size_on_link) + self._cur_batch = EventDataBatch(self._max_message_size_on_link, amqp_transport=self._amqp_transport) except ValueError: # add single event exceeds the cur batch size, create new batch async with self._lock: self._buffered_queue.put(self._cur_batch) - self._cur_batch = EventDataBatch(self._max_message_size_on_link) + self._cur_batch = EventDataBatch(self._max_message_size_on_link, amqp_transport=self._amqp_transport) self._cur_batch.add(events) async with self._lock: self._cur_buffered_len += new_events_len @@ -147,7 +149,7 @@ async def _flush(self, timeout_time=None, raise_error=True): _LOGGER.info("Partition: %r started flushing.", self.partition_id) if self._cur_batch: # if there is batch, enqueue it to the buffer first self._buffered_queue.put(self._cur_batch) - self._cur_batch = EventDataBatch(self._max_message_size_on_link) + self._cur_batch = EventDataBatch(self._max_message_size_on_link, amqp_transport=self._amqp_transport) while self._buffered_queue.qsize() > 0: remaining_time = timeout_time - time.time() if timeout_time else None if (remaining_time and remaining_time > 0) or remaining_time is None: @@ -200,7 +202,7 @@ async def _flush(self, timeout_time=None, raise_error=True): self._last_send_time = time.time() #reset curr_buffered self._cur_buffered_len = 0 - self._cur_batch = EventDataBatch(self._max_message_size_on_link) + self._cur_batch = EventDataBatch(self._max_message_size_on_link, amqp_transport=self._amqp_transport) _LOGGER.info("Partition %r finished flushing.", self.partition_id) async def check_max_wait_time_worker(self): diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_buffered_producer/_buffered_producer_dispatcher_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_buffered_producer/_buffered_producer_dispatcher_async.py index d3f2135ff170..64e565944aaf 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_buffered_producer/_buffered_producer_dispatcher_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_buffered_producer/_buffered_producer_dispatcher_async.py @@ -33,6 +33,7 @@ def __init__( eventhub_name: str, max_message_size_on_link: int, *, + amqp_transport: AmqpTransportAsync, max_buffer_length: int = 1500, max_wait_time: float = 1 ): @@ -47,6 +48,7 @@ def __init__( self._partition_resolver = PartitionResolver(self._partition_ids) self._max_wait_time = max_wait_time self._max_buffer_length = max_buffer_length + self._amqp_transport = amqp_transport async def _get_partition_id(self, partition_id, partition_key): if partition_id: @@ -79,6 +81,7 @@ async def enqueue_events( self._max_message_size_on_link, max_wait_time=self._max_wait_time, max_buffer_length=self._max_buffer_length, + amqp_transport=self._amqp_transport, ) await buffered_producer.start() self._buffered_producers[pid] = buffered_producer diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_client_base_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_client_base_async.py index e9becf77a884..9aea94814252 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_client_base_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_client_base_async.py @@ -34,14 +34,26 @@ ) from ._async_utils import get_dict_with_loop_if_needed from ._connection_manager_async import get_connection_manager -from ._transport._uamqp_transport_async import UamqpTransportAsync +try: + from ._transport._uamqp_transport_async import UamqpTransportAsync +except ImportError: + UamqpTransportAsync = None # type: ignore +from ._transport._pyamqp_transport_async import PyamqpTransportAsync if TYPE_CHECKING: - from uamqp import ( - authentication, - Message, - AMQPClientAsync, - ) + from .._pyamqp.message import Message + from .._pyamqp.aio import AMQPClientAsync + from .._pyamqp.aio._authentication_async import JWTTokenAuthAsync + try: + from uamqp import ( + authentication as uamqp_authentication, + Message as uamqp_Message, + AMQPClientAsync as uamqp_AMQPClientAsync, + ) + except ImportError: + uamqp_authentication = None + uamqp_Message = None + uamqp_AMQPClientAsync = None from azure.core.credentials_async import AsyncTokenCredential CredentialTypes = Union[ @@ -58,8 +70,7 @@ class AbstractConsumerProducer(Protocol): @property - def _name(self): - # type: () -> str + def _name(self) -> str: """Name of the consumer or producer""" @_name.setter @@ -67,8 +78,7 @@ def _name(self, value): pass @property - def _client(self): - # type: () -> ClientBaseAsync + def _client(self) -> ClientBaseAsync: """The instance of EventHubComsumerClient or EventHubProducerClient""" @_client.setter @@ -76,13 +86,11 @@ def _client(self, value): pass @property - def _handler(self): - # type: () -> AMQPClientAsync + def _handler(self) -> Union[uamqp_AMQPClientAsync, AMQPClientAsync]: """The instance of SendClientAsync or ReceiveClientAsync""" @property - def _internal_kwargs(self): - # type: () -> dict + def _internal_kwargs(self) -> dict: """The dict with an event loop that users may pass in to wrap sync calls to async API. It's furthur passed to uamqp APIs """ @@ -100,7 +108,7 @@ def running(self): def running(self, value): pass - def _create_handler(self, auth: authentication.JWTTokenAsync) -> None: + def _create_handler(self, auth: Union[uamqp_authentication.JWTTokenAsync, JWTTokenAuthAsync]) -> None: pass _MIXIN_BASE = AbstractConsumerProducer @@ -163,8 +171,7 @@ class EventhubAzureNamedKeyTokenCredentialAsync(object): :type credential: ~azure.core.credentials.AzureNamedKeyCredential """ - def __init__(self, azure_named_key_credential): - # type: (AzureNamedKeyCredential) -> None + def __init__(self, azure_named_key_credential: AzureNamedKeyCredential) -> None: self._credential = azure_named_key_credential self.token_type = b"servicebus.windows.net:sastoken" @@ -208,8 +215,10 @@ def __init__( **kwargs: Any ) -> None: self._internal_kwargs = get_dict_with_loop_if_needed(kwargs.get("loop", None)) - uamqp_transport = kwargs.pop("uamqp_transport", True) - self._amqp_transport = UamqpTransportAsync + uamqp_transport = kwargs.get("uamqp_transport", False) + if uamqp_transport and not UamqpTransportAsync: + raise ValueError("To use the uAMQP transport, please install `uamqp>=1.6.0,<2.0.0`.") + self._amqp_transport = UamqpTransportAsync if uamqp_transport else PyamqpTransportAsync if isinstance(credential, AzureSasCredential): self._credential = EventhubAzureSasTokenCredentialAsync(credential) # type: ignore elif isinstance(credential, AzureNamedKeyCredential): @@ -220,11 +229,14 @@ def __init__( fully_qualified_namespace=fully_qualified_namespace, eventhub_name=eventhub_name, credential=self._credential, - uamqp_transport=uamqp_transport, amqp_transport=self._amqp_transport, **kwargs ) - self._conn_manager_async = get_connection_manager(**kwargs) + kwargs["custom_endpoint_address"] = self._config.custom_endpoint_address + self._conn_manager_async = get_connection_manager( + amqp_transport=self._amqp_transport, + **kwargs + ) def __enter__(self): raise TypeError( @@ -244,7 +256,7 @@ def _from_connection_string(conn_str: str, **kwargs) -> Dict[str, Any]: kwargs["credential"] = EventHubSharedKeyCredential(policy, key) return kwargs - async def _create_auth_async(self) -> authentication.JWTTokenAsync: + async def _create_auth_async(self) -> Union[uamqp_authentication.JWTTokenAsync, JWTTokenAuthAsync]: """ Create an ~uamqp.authentication.SASTokenAuthAsync instance to authenticate the session. @@ -305,7 +317,7 @@ async def _backoff_async( ) raise last_exception - async def _management_request_async(self, mgmt_msg: Message, op_type: bytes) -> Any: + async def _management_request_async(self, mgmt_msg: Union[Message, uamqp_Message], op_type: bytes) -> Any: retried_times = 0 last_exception = None while retried_times <= self._config.max_retries: @@ -323,7 +335,7 @@ async def _management_request_async(self, mgmt_msg: Message, op_type: bytes) -> mgmt_msg.application_properties[ "security_token" ] = await self._amqp_transport.get_updated_token_async(mgmt_auth) - response = await self._amqp_transport.mgmt_client_request_async( + status_code, description, response = await self._amqp_transport.mgmt_client_request_async( mgmt_client, mgmt_msg, operation=READ_OPERATION, @@ -331,10 +343,7 @@ async def _management_request_async(self, mgmt_msg: Message, op_type: bytes) -> status_code_field=MGMT_STATUS_CODE, description_fields=MGMT_STATUS_DESC, ) - status_code = int(response.application_properties[MGMT_STATUS_CODE]) - description = response.application_properties.get( - MGMT_STATUS_DESC - ) # type: Optional[Union[str, bytes]] + status_code = int(status_code) if description and isinstance(description, bytes): description = description.decode("utf-8") if status_code < 400: @@ -343,7 +352,15 @@ async def _management_request_async(self, mgmt_msg: Message, op_type: bytes) -> except asyncio.CancelledError: # pylint: disable=try-except-raise raise except Exception as exception: # pylint:disable=broad-except - last_exception = await self._amqp_transport._handle_exception_async(exception, self) # pylint: disable=protected-access + # is_consumer=True passed in here, ALTHOUGH this method is shared by the producer and consumer. + # is_consumer will only be checked if FileNotFoundError is raised by self.mgmt_client.open() due to + # invalid/non-existent connection_verify filepath. The producer will encounter the FileNotFoundError + # when opening the SendClient, so is_consumer=True will not be passed to amqp_transport.handle_exception + # there. This is for uamqp exception parity, which raises FileNotFoundError in the consumer and + # EventHubError in the producer. TODO: Remove `is_consumer` kwarg when resolving issue #27128. + last_exception = await self._amqp_transport._handle_exception_async( # pylint: disable=protected-access + exception, self, is_consumer=True + ) await self._backoff_async( retried_times=retried_times, last_exception=last_exception ) @@ -357,7 +374,7 @@ async def _management_request_async(self, mgmt_msg: Message, op_type: bytes) -> await mgmt_client.close_async() async def _get_eventhub_properties_async(self) -> Dict[str, Any]: - mgmt_msg = mgmt_msg = self._amqp_transport.build_message( + mgmt_msg = self._amqp_transport.build_message( application_properties={"name": self.eventhub_name} ) response = await self._management_request_async( @@ -465,11 +482,11 @@ async def _close_connection_async(self) -> None: await self._close_handler_async() await self._client._conn_manager_async.reset_connection_if_broken() # pylint:disable=protected-access - async def _handle_exception(self, exception: Exception) -> Exception: + async def _handle_exception(self, exception: Exception, *, is_consumer: bool = False) -> Exception: # pylint: disable=protected-access exception = self._client._amqp_transport.check_timeout_exception(self, exception) return await self._client._amqp_transport._handle_exception_async( - exception, self + exception, self, is_consumer=is_consumer ) async def _do_retryable_operation( diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_connection_manager_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_connection_manager_async.py index ec3c0fffaebf..5c81df19140f 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_connection_manager_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_connection_manager_async.py @@ -4,16 +4,22 @@ # -------------------------------------------------------------------------------------------- from __future__ import annotations -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Optional, Union from asyncio import Lock -from ._transport._uamqp_transport_async import UamqpTransportAsync from .._connection_manager import _ConnectionMode from .._constants import TransportType if TYPE_CHECKING: - from uamqp.authentication import JWTTokenAsync - from uamqp.async_ops import ConnectionAsync + from .._pyamqp.aio._authentication_async import JWTTokenAuthAsync + from .._pyamqp.aio._connection_async import Connection as ConnectionAsync + from ._transport._base_async import AmqpTransportAsync + try: + from uamqp.authentication import JWTTokenAsync as uamqp_JWTTokenAuthAsync + from uamqp.async_ops import ConnectionAsync as uamqp_ConnectionAsync + except ImportError: + uamqp_JWTTokenAuthAsync = None + uamqp_ConnectionAsync = None try: from typing_extensions import Protocol @@ -22,8 +28,12 @@ class ConnectionManager(Protocol): async def get_connection( - self, *, host: Optional[str] = None, auth: Optional[JWTTokenAsync] = None, endpoint: Optional[str] = None - ) -> ConnectionAsync: + self, + *, + host: Optional[str] = None, + auth: Optional[Union[uamqp_JWTTokenAuthAsync, JWTTokenAuthAsync]] = None, + endpoint: Optional[str] = None, + ) -> Union[ConnectionAsync, uamqp_ConnectionAsync]: pass async def close_connection(self) -> None: @@ -40,6 +50,7 @@ def __init__(self, **kwargs) -> None: self._conn = None self._container_id = kwargs.get("container_id") + self._custom_endpoint_address = kwargs.get("custom_endpoint_address") self._debug = kwargs.get("debug") self._error_policy = kwargs.get("error_policy") self._properties = kwargs.get("properties") @@ -49,20 +60,23 @@ def __init__(self, **kwargs) -> None: self._max_frame_size = kwargs.get("max_frame_size") self._channel_max = kwargs.get("channel_max") self._idle_timeout = kwargs.get("idle_timeout") - self._remote_idle_timeout_empty_frame_send_ratio = kwargs.get( - "remote_idle_timeout_empty_frame_send_ratio" - ) - self._amqp_transport = kwargs.get("amqp_transport", UamqpTransportAsync) + self._remote_idle_timeout_empty_frame_send_ratio = kwargs.get("remote_idle_timeout_empty_frame_send_ratio") + self._amqp_transport: AmqpTransportAsync = kwargs.pop("amqp_transport") async def get_connection( - self, *, host: Optional[str] = None, auth: Optional[JWTTokenAsync] = None, endpoint: Optional[str] = None - ) -> ConnectionAsync: + self, + *, + host: Optional[str] = None, + auth: Optional[Union[JWTTokenAuthAsync, uamqp_JWTTokenAuthAsync]] = None, + endpoint: Optional[str] = None, + ) -> Union[ConnectionAsync, uamqp_ConnectionAsync]: async with self._lock: if self._conn is None: self._conn = self._amqp_transport.create_connection_async( host=host, auth=auth, endpoint=endpoint, + custom_endpoint_address=self._custom_endpoint_address, container_id=self._container_id, max_frame_size=self._max_frame_size, channel_max=self._channel_max, @@ -94,7 +108,11 @@ def __init__(self, **kwargs) -> None: pass async def get_connection( - self, *, host: Optional[str] = None, auth: Optional[JWTTokenAsync] = None, endpoint: Optional[str] = None + self, + *, + host: Optional[str] = None, + auth: Optional[Union[JWTTokenAuthAsync, uamqp_JWTTokenAuthAsync]] = None, + endpoint: Optional[str] = None, ) -> None: pass # return None @@ -106,7 +124,7 @@ async def reset_connection_if_broken(self) -> None: def get_connection_manager(**kwargs) -> "ConnectionManager": - connection_mode = kwargs.get("connection_mode", _ConnectionMode.SeparateConnection) # type: ignore + connection_mode = kwargs.get("connection_mode", _ConnectionMode.SeparateConnection) # type: ignore if connection_mode == _ConnectionMode.ShareConnection: return _SharedConnectionManager(**kwargs) return _SeparateConnectionManager(**kwargs) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_consumer_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_consumer_async.py index c918bca559c8..c4e311082e68 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_consumer_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_consumer_async.py @@ -4,9 +4,11 @@ # -------------------------------------------------------------------------------------------- from __future__ import annotations import uuid +import asyncio import logging from collections import deque from typing import TYPE_CHECKING, Callable, Awaitable, Dict, Optional, Union, List +from functools import partial from ._client_base_async import ConsumerProducerMixin from ._async_utils import get_dict_with_loop_if_needed @@ -16,9 +18,24 @@ if TYPE_CHECKING: from typing import Deque - import uamqp - from uamqp import ReceiveClientAsync, Source, types - from uamqp.authentication import JWTTokenAsync + + try: + from uamqp import ( # pylint: disable=unused-import + ReceiveClientAsync as uamqp_ReceiveClientAsync, + Message as uamqp_Message, + ) + from uamqp.types import AMQPType as uamqp_AMQPType + from uamqp.authentication import JWTTokenAsync as uamqp_JWTTokenAsync + except ImportError: + uamqp_Message = None + uamqp_ReceiveClientAsync = None + uamqp_AMQPType = None + uamqp_JWTTokenAsync = None + + from .._pyamqp.aio._authentication_async import JWTTokenAuthAsync + from .._pyamqp.aio._client_async import ReceiveClientAsync + from .._pyamqp import types + from ._consumer_client_async import EventHubConsumerClient _LOGGER = logging.getLogger(__name__) @@ -76,9 +93,9 @@ def __init__(self, client: "EventHubConsumerClient", source: str, **kwargs) -> N self.closed = False self._amqp_transport = kwargs.pop("amqp_transport") - self._on_event_received: Callable[[Union[Optional[EventData], List[EventData]]], Awaitable[None]] = kwargs[ - "on_event_received" - ] + self._on_event_received: Callable[ + [Union[Optional[EventData], List[EventData]]], Awaitable[None] + ] = kwargs["on_event_received"] self._internal_kwargs = get_dict_with_loop_if_needed(kwargs.get("loop", None)) self._client = client self._source = source @@ -88,11 +105,17 @@ def __init__(self, client: "EventHubConsumerClient", source: str, **kwargs) -> N self._owner_level = owner_level self._keep_alive = keep_alive self._auto_reconnect = auto_reconnect - self._retry_policy = self._amqp_transport.create_retry_policy(self._client._config) + self._retry_policy = self._amqp_transport.create_retry_policy( + self._client._config + ) self._reconnect_backoff = 1 self._timeout = 0 - self._idle_timeout = (idle_timeout * self._amqp_transport.TIMEOUT_FACTOR) if idle_timeout else None - link_properties: Dict[types.AMQPType, types.AMQPType] = {} + self._idle_timeout = ( + (idle_timeout * self._amqp_transport.TIMEOUT_FACTOR) + if idle_timeout + else None + ) + link_properties: Dict[bytes, int] = {} self._partition = self._source.split("/")[-1] self._name = f"EHReceiver-{uuid.uuid4()}-partition{self._partition}" if owner_level is not None: @@ -102,24 +125,35 @@ def __init__(self, client: "EventHubConsumerClient", source: str, **kwargs) -> N or self._timeout # pylint:disable=protected-access ) * self._amqp_transport.TIMEOUT_FACTOR link_properties[TIMEOUT_SYMBOL] = int(link_property_timeout_ms) - self._link_properties = self._amqp_transport.create_link_properties(link_properties) + self._link_properties: Union[ + Dict[uamqp_AMQPType, uamqp_AMQPType], Dict[types.AMQPTypes, types.AMQPTypes] + ] = self._amqp_transport.create_link_properties(link_properties) self._handler: Optional[ReceiveClientAsync] = None self._track_last_enqueued_event_properties = ( track_last_enqueued_event_properties ) - self._message_buffer: Deque[uamqp.Message] = deque() + self._message_buffer: Deque[uamqp_Message] = deque() self._last_received_event: Optional[EventData] = None + self._message_buffer_lock = asyncio.Lock() + self._last_callback_called_time = None + self._callback_task_run = None - def _create_handler(self, auth: JWTTokenAsync) -> None: + def _create_handler( + self, auth: Union[uamqp_JWTTokenAsync, JWTTokenAuthAsync] + ) -> None: source = self._amqp_transport.create_source( self._source, self._offset, - event_position_selector(self._offset, self._offset_inclusive) + event_position_selector(self._offset, self._offset_inclusive), + ) + desired_capabilities = ( + [RECEIVER_RUNTIME_METRIC_SYMBOL] + if self._track_last_enqueued_event_properties + else None ) - desired_capabilities = [RECEIVER_RUNTIME_METRIC_SYMBOL] if self._track_last_enqueued_event_properties else None self._handler = self._amqp_transport.create_receive_client( - config=self._client._config, # pylint:disable=protected-access + config=self._client._config, # pylint:disable=protected-access source=source, auth=auth, network_trace=self._client._config.network_tracing, # pylint:disable=protected-access @@ -130,19 +164,19 @@ def _create_handler(self, auth: JWTTokenAsync) -> None: keep_alive_interval=self._keep_alive, client_name=self._name, properties=create_properties( - self._client._config.user_agent, amqp_transport=self._amqp_transport # pylint:disable=protected-access + self._client._config.user_agent, # pylint:disable=protected-access + amqp_transport=self._amqp_transport, ), desired_capabilities=desired_capabilities, streaming_receive=True, - message_received_callback=self._message_received, + message_received_callback=partial( + self._amqp_transport.message_received_async, self + ), ) async def _open_with_retry(self) -> None: await self._do_retryable_operation(self._open, operation_need_param=False) - def _message_received(self, message: uamqp.Message) -> None: - self._message_buffer.append(message) - def _next_message_in_buffer(self): # pylint:disable=protected-access message = self._message_buffer.popleft() @@ -153,4 +187,6 @@ def _next_message_in_buffer(self): async def receive( self, batch=False, max_batch_size=300, max_wait_time=None ) -> None: - await self._amqp_transport.receive_messages(self, batch, max_batch_size, max_wait_time) + await self._amqp_transport.receive_messages_async( + self, batch, max_batch_size, max_wait_time + ) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_consumer_client_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_consumer_client_async.py index 73ef31f79ac4..8e33b705dacb 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_consumer_client_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_consumer_client_async.py @@ -131,6 +131,8 @@ class EventHubConsumerClient( :keyword str connection_verify: Path to the custom CA_BUNDLE file of the SSL certificate which is used to authenticate the identity of the connection endpoint. Default is None in which case `certifi.where()` will be used. + :keyword bool uamqp_transport: Whether to use the `uamqp` library as the underlying transport. The default value is + False and the Pure Python AMQP library will be used as the underlying transport. .. admonition:: Example: @@ -205,7 +207,7 @@ def _create_consumer( source_url = "amqps://{}{}/ConsumerGroups/{}/Partitions/{}".format( self._address.hostname, self._address.path, consumer_group, partition_id ) - handler = EventHubConsumer( + handler = EventHubConsumer( # type: ignore self, source_url, on_event_received=on_event_received, diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_eventprocessor/event_processor.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_eventprocessor/event_processor.py index 69e8567b7a86..5a3408912afa 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_eventprocessor/event_processor.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_eventprocessor/event_processor.py @@ -2,6 +2,7 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for license information. # -------------------------------------------------------------------------------------------- +from __future__ import annotations import random from typing import ( Dict, @@ -28,7 +29,6 @@ from .in_memory_checkpoint_store import InMemoryCheckpointStore from .checkpoint_store import CheckpointStore from ._ownership_manager import OwnershipManager -from .utils import get_running_loop from .._async_utils import get_dict_with_loop_if_needed if TYPE_CHECKING: @@ -170,7 +170,7 @@ def _create_tasks_for_claimed_ownership( if partition_id not in self._tasks or self._tasks[partition_id].done(): checkpoint = checkpoints.get(partition_id) if checkpoints else None if self._running: - self._tasks[partition_id] = get_running_loop().create_task( + self._tasks[partition_id] = asyncio.create_task( self._receive(partition_id, checkpoint) ) _LOGGER.info( diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_producer_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_producer_async.py index 8f0fe7d6004f..8972e4c1efd4 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_producer_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_producer_async.py @@ -6,7 +6,7 @@ import uuid import asyncio import logging -from typing import Iterable, Union, Optional, Any, AnyStr, List, TYPE_CHECKING +from typing import Iterable, Union, Optional, Any, AnyStr, List, TYPE_CHECKING, cast from azure.core.tracing import AbstractSpan @@ -24,10 +24,20 @@ from ._async_utils import get_dict_with_loop_if_needed if TYPE_CHECKING: - from uamqp import types, constants, errors - from uamqp import SendClientAsync + try: + from uamqp import ( # pylint: disable=unused-import + constants, + SendClientAsync as uamqp_SendClientAsync, + ) + from uamqp.constants import MessageSendResult as uamqp_MessageSendResult + from uamqp.authentication import JWTTokenAsync as uamqp_JWTTokenAsync + except ImportError: + uamqp_MessageSendResult = None + uamqp_SendClientAsync = None + uamqp_JWTTokenAsync = None - from uamqp.authentication import JWTTokenAsync # pylint: disable=ungrouped-imports + from .._pyamqp.aio._client_async import SendClientAsync + from .._pyamqp.aio._authentication_async import JWTTokenAuthAsync from ._producer_client_async import EventHubProducerClient _LOGGER = logging.getLogger(__name__) @@ -94,16 +104,17 @@ def __init__(self, client: EventHubProducerClient, target: str, **kwargs) -> Non if partition: self._target += "/Partitions/" + partition self._name += "-partition{}".format(partition) - self._handler: Optional[SendClientAsync] = None - self._outcome: Optional[constants.MessageSendResult] = None + self._handler: Optional[Union[uamqp_SendClientAsync, SendClientAsync]] = None + self._outcome: Optional[uamqp_MessageSendResult] = None self._condition: Optional[Exception] = None self._lock = asyncio.Lock(**self._internal_kwargs) self._link_properties = self._amqp_transport.create_link_properties( {TIMEOUT_SYMBOL: int(self._timeout * self._amqp_transport.TIMEOUT_FACTOR)} ) - - def _create_handler(self, auth: "JWTTokenAsync") -> None: + def _create_handler( + self, auth: Union[uamqp_JWTTokenAsync, JWTTokenAuthAsync] + ) -> None: self._handler = self._amqp_transport.create_send_client( config=self._client._config, # pylint:disable=protected-access target=self._target, @@ -118,7 +129,7 @@ def _create_handler(self, auth: "JWTTokenAsync") -> None: self._client._config.user_agent, # pylint: disable=protected-access amqp_transport=self._amqp_transport, ), - msg_timeout=self._timeout * 1000, + msg_timeout=self._timeout * self._amqp_transport.TIMEOUT_FACTOR, ) async def _open_with_retry(self) -> Any: @@ -142,10 +153,10 @@ async def _send_event_data_with_retry( await self._do_retryable_operation(self._send_event_data, timeout=timeout) def _on_outcome( - self, outcome: constants.MessageSendResult, condition: Optional[Exception] + self, outcome: uamqp_MessageSendResult, condition: Optional[Exception] ) -> None: """ - Called when the outcome is received for a delivery. + ONLY USED FOR uamqp_transport=True. Called when the outcome is received for a delivery. :param outcome: The outcome of the message delivery - success or failure. :type outcome: ~uamqp.constants.MessageSendResult @@ -169,7 +180,8 @@ def _wrap_eventdata( ) if partition_key: self._amqp_transport.set_message_partition_key( - outgoing_event_data._message, partition_key # pylint: disable=protected-access + outgoing_event_data._message, # pylint: disable=protected-access + partition_key, ) wrapper_event_data = outgoing_event_data trace_message(wrapper_event_data, span) @@ -179,10 +191,21 @@ def _wrap_eventdata( ): # The partition_key in the param will be omitted. if not event_data: return event_data + # If AmqpTransports are not the same, create batch with correct BatchMessage. + if ( + self._amqp_transport.TIMEOUT_FACTOR + != event_data._amqp_transport.TIMEOUT_FACTOR # pylint: disable=protected-access + ): + # pylint: disable=protected-access + event_data = EventDataBatch._from_batch( + event_data._internal_events, + amqp_transport=self._amqp_transport, + partition_key=cast(AnyStr, event_data._partition_key), + partition_id=event_data._partition_id, + max_size_in_bytes=event_data.max_size_in_bytes, + ) if ( - partition_key - and partition_key - != event_data._partition_key # pylint: disable=protected-access + partition_key and partition_key != event_data._partition_key # pylint: disable=protected-access ): raise ValueError( "The partition_key does not match the one of the EventDataBatch" @@ -198,7 +221,7 @@ def _wrap_eventdata( event_data, partition_key, self._amqp_transport ) event_data = _set_trace_message(event_data, span) - wrapper_event_data = EventDataBatch._from_batch( # type: ignore # pylint: disable=protected-access + wrapper_event_data = EventDataBatch._from_batch( # type: ignore # pylint: disable=protected-access event_data, self._amqp_transport, partition_key ) return wrapper_event_data @@ -210,7 +233,7 @@ async def send( ], *, partition_key: Optional[AnyStr] = None, - timeout: Optional[float] = None + timeout: Optional[float] = None, ) -> None: """ Sends an event data and blocks until acknowledgement is @@ -246,7 +269,9 @@ async def send( if not wrapper_event_data: return - self._unsent_events = [wrapper_event_data._message] # pylint: disable=protected-access + self._unsent_events = [ + wrapper_event_data._message # pylint: disable=protected-access + ] if child: self._client._add_span_request_attributes( # pylint: disable=protected-access diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_producer_client_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_producer_client_async.py index af59da1efc5a..a47dd81dc5ff 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_producer_client_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_producer_client_async.py @@ -108,6 +108,8 @@ class EventHubProducerClient( :keyword str connection_verify: Path to the custom CA_BUNDLE file of the SSL certificate which is used to authenticate the identity of the connection endpoint. Default is None in which case `certifi.where()` will be used. + :keyword bool uamqp_transport: Whether to use the `uamqp` library as the underlying transport. The default value is + False and the Pure Python AMQP library will be used as the underlying transport. .. admonition:: Example: @@ -230,6 +232,7 @@ async def _buffered_send(self, events, **kwargs): self._max_message_size_on_link, max_wait_time=self._max_wait_time, max_buffer_length=self._max_buffer_length, + amqp_transport=self._amqp_transport ) await self._buffered_producer_dispatcher.enqueue_events(events, **kwargs) @@ -344,7 +347,7 @@ def _create_producer( self._config.send_timeout if send_timeout is None else send_timeout ) - handler = EventHubProducer( + handler = EventHubProducer( # type: ignore self, target, partition=partition_id, @@ -719,7 +722,8 @@ async def create_batch( event_data_batch = EventDataBatch( max_size_in_bytes=(max_size_in_bytes or self._max_message_size_on_link), partition_id=partition_id, - partition_key=partition_key + partition_key=partition_key, + amqp_transport=self._amqp_transport ) return event_data_batch diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_transport/_base_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_transport/_base_async.py index ce9342c607ed..259ea1358cf5 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_transport/_base_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_transport/_base_async.py @@ -3,11 +3,15 @@ # Licensed under the MIT License. See License.txt in the project root for license information. # -------------------------------------------------------------------------------------------- from __future__ import annotations -from typing import Tuple, Union, TYPE_CHECKING from abc import ABC, abstractmethod +from typing import Tuple, Union, TYPE_CHECKING +from typing_extensions import Literal if TYPE_CHECKING: - from uamqp import types as uamqp_types + try: + from uamqp import types as uamqp_types + except ImportError: + uamqp_types = None class AmqpTransportAsync(ABC): # pylint: disable=too-many-public-methods """ @@ -20,12 +24,12 @@ class AmqpTransportAsync(ABC): # pylint: disable=too-many-public-methods CONNECTION_CLOSING_STATES: Tuple # define symbols - PRODUCT_SYMBOL: Union[uamqp_types.AMQPSymbol, str, bytes] - VERSION_SYMBOL: Union[uamqp_types.AMQPSymbol, str, bytes] - FRAMEWORK_SYMBOL: Union[uamqp_types.AMQPSymbol, str, bytes] - PLATFORM_SYMBOL: Union[uamqp_types.AMQPSymbol, str, bytes] - USER_AGENT_SYMBOL: Union[uamqp_types.AMQPSymbol, str, bytes] - PROP_PARTITION_KEY_AMQP_SYMBOL: Union[uamqp_types.AMQPSymbol, str, bytes] + PRODUCT_SYMBOL: Union[uamqp_types.AMQPSymbol, Literal["product"]] + VERSION_SYMBOL: Union[uamqp_types.AMQPSymbol, Literal["version"]] + FRAMEWORK_SYMBOL: Union[uamqp_types.AMQPSymbol, Literal["framework"]] + PLATFORM_SYMBOL: Union[uamqp_types.AMQPSymbol, Literal["platform"]] + USER_AGENT_SYMBOL: Union[uamqp_types.AMQPSymbol, Literal["user-agent"]] + PROP_PARTITION_KEY_AMQP_SYMBOL: Union[uamqp_types.AMQPSymbol, Literal[b'x-opt-partition-key']] @staticmethod @@ -198,7 +202,7 @@ def create_receive_client(*, config, **kwargs): @staticmethod @abstractmethod - async def receive_messages(consumer, batch, max_batch_size, max_wait_time): + async def receive_messages_async(consumer, batch, max_batch_size, max_wait_time): """ Receives messages, creates events, and returns them by calling the on received callback. :param ~azure.eventhub.aio.EventHubConsumer consumer: The EventHubConsumer. diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_transport/_pyamqp_transport_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_transport/_pyamqp_transport_async.py new file mode 100644 index 000000000000..8a2432f543ea --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_transport/_pyamqp_transport_async.py @@ -0,0 +1,371 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +from __future__ import annotations +import asyncio +import time +import logging +from typing import Union, cast, TYPE_CHECKING, List + +from ..._pyamqp import constants, error as errors +from ..._pyamqp.aio import AMQPClientAsync, SendClientAsync, ReceiveClientAsync +from ..._pyamqp.aio._authentication_async import JWTTokenAuthAsync +from ..._pyamqp.aio._connection_async import Connection as ConnectionAsync + +from ._base_async import AmqpTransportAsync +from ..._transport._pyamqp_transport import PyamqpTransport +from ...exceptions import ( + EventHubError, + EventDataSendError, + OperationTimeoutError +) +from ..._common import EventData + +if TYPE_CHECKING: + from .._client_base_async import ClientBaseAsync, ConsumerProducerMixin + from ..._pyamqp.message import Message + +_LOGGER = logging.getLogger(__name__) + + +class PyamqpTransportAsync(PyamqpTransport, AmqpTransportAsync): + """ + Class which defines pyamqp-based methods used by the producer and consumer. + """ + + @staticmethod + async def create_connection_async(**kwargs): + """ + Creates and returns the pyamqp Connection object. + :keyword str host: The hostname, used by pyamqp. + :keyword JWTTokenAuthAsync auth: The auth, used by pyamqp. + :keyword str endpoint: The endpoint, used by pyamqp. + :keyword str container_id: Required. + :keyword int max_frame_size: Required. + :keyword int channel_max: Required. + :keyword int idle_timeout: Required. + :keyword Dict properties: Required. + :keyword int remote_idle_timeout_empty_frame_send_ratio: Required. + :keyword error_policy: Required. + :keyword bool debug: Required. + :keyword str encoding: Required. + """ + endpoint = kwargs.pop("endpoint") + host = kwargs.pop("host") # pylint:disable=unused-variable + auth = kwargs.pop("auth") # pylint:disable=unused-variable + network_trace = kwargs.pop("debug") + return ConnectionAsync(endpoint, network_trace=network_trace, **kwargs) + + @staticmethod + async def close_connection_async(connection): + """ + Closes existing connection. + :param connection: pyamqp Connection. + """ + await connection.close() + + @staticmethod + def create_send_client(*, config, **kwargs): # pylint:disable=unused-argument + """ + Creates and returns the pyamqp SendClient. + :param ~azure.eventhub._configuration.Configuration config: The configuration. + + :keyword str target: Required. The target. + :keyword JWTTokenAuth auth: Required. + :keyword int idle_timeout: Required. + :keyword network_trace: Required. + :keyword retry_policy: Required. + :keyword keep_alive_interval: Required. + :keyword str client_name: Required. + :keyword dict link_properties: Required. + :keyword properties: Required. + """ + target = kwargs.pop("target") + # TODO: extra passed in to pyamqp, but not used. should be used? + msg_timeout = kwargs.pop("msg_timeout") # pylint: disable=unused-variable # TODO: not used by pyamqp? + + return SendClientAsync( + config.hostname, + target, + custom_endpoint_address=config.custom_endpoint_address, + connection_verify=config.connection_verify, + transport_type=config.transport_type, + http_proxy=config.http_proxy, + **kwargs, + ) + + @staticmethod + async def send_messages_async(producer, timeout_time, last_exception, logger): + """ + Handles sending of event data messages. + :param ~azure.eventhub._producer.EventHubProducer producer: The producer with handler to send messages. + :param int timeout_time: Timeout time. + :param last_exception: Exception to raise if message timed out. Only used by pyamqp transport. + :param logger: Logger. + """ + # pylint: disable=protected-access + try: + await producer._open() + timeout = timeout_time - time.time() if timeout_time else 0 + await producer._handler.send_message_async(producer._unsent_events[0], timeout=timeout) + producer._unsent_events = None + except TimeoutError as exc: + raise OperationTimeoutError(message=str(exc), details=exc) + + @staticmethod + def create_receive_client(*, config, **kwargs): # pylint:disable=unused-argument + """ + Creates and returns the receive client. + :param ~azure.eventhub._configuration.Configuration config: The configuration. + + :keyword str source: Required. The source. + :keyword str offset: Required. + :keyword str offset_inclusive: Required. + :keyword JWTTokenAuth auth: Required. + :keyword int idle_timeout: Required. + :keyword network_trace: Required. + :keyword retry_policy: Required. + :keyword str client_name: Required. + :keyword dict link_properties: Required. + :keyword properties: Required. + :keyword link_credit: Required. The prefetch. + :keyword keep_alive_interval: Required. + :keyword desired_capabilities: Required. + :keyword streaming_receive: Required. + :keyword message_received_callback: Required. + :keyword timeout: Required. + """ + + source = kwargs.pop("source") + return ReceiveClientAsync( + config.hostname, + source, + receive_settle_mode=constants.ReceiverSettleMode.First, # TODO: make more descriptive in pyamqp? + http_proxy=config.http_proxy, + transport_type=config.transport_type, + custom_endpoint_address=config.custom_endpoint_address, + connection_verify=config.connection_verify, + **kwargs, + ) + + @staticmethod + async def _callback_task(consumer, batch, max_batch_size, max_wait_time): + while consumer._callback_task_run: # pylint: disable=protected-access + async with consumer._message_buffer_lock: # pylint: disable=protected-access + messages = [ + consumer._message_buffer.popleft() # pylint: disable=protected-access + for _ in range(min(max_batch_size, len(consumer._message_buffer))) # pylint: disable=protected-access + ] + events = [EventData._from_message(message) for message in messages] # pylint: disable=protected-access + now_time = time.time() + if len(events) > 0: + await consumer._on_event_received(events if batch else events[0]) # pylint: disable=protected-access + consumer._last_callback_called_time = now_time # pylint: disable=protected-access + else: + if max_wait_time and (now_time - consumer._last_callback_called_time) > max_wait_time: # pylint: disable=protected-access + # no events received, and need to callback + await consumer._on_event_received([] if batch else None) # pylint: disable=protected-access + consumer._last_callback_called_time = now_time # pylint: disable=protected-access + # backoff a bit to avoid throttling CPU when no events are coming + await asyncio.sleep(0.05) + + @staticmethod + async def _receive_task(consumer): + # pylint:disable=protected-access + max_retries = consumer._client._config.max_retries + retried_times = 0 + running = True + try: + while retried_times <= max_retries and running and consumer._callback_task_run: + try: + await consumer._open() # pylint: disable=protected-access + running = await cast(ReceiveClientAsync, consumer._handler).do_work_async(batch=consumer._prefetch) + except asyncio.CancelledError: # pylint: disable=try-except-raise + raise + except Exception as exception: # pylint: disable=broad-except + if ( + isinstance(exception, errors.AMQPLinkError) + and exception.condition == errors.ErrorCondition.LinkStolen # pylint: disable=no-member + ): + raise await consumer._handle_exception(exception) + if not consumer.running: # exit by close + return + if consumer._last_received_event: + consumer._offset = consumer._last_received_event.offset + last_exception = await consumer._handle_exception(exception) + retried_times += 1 + if retried_times > max_retries: + _LOGGER.info( + "%r operation has exhausted retry. Last exception: %r.", + consumer._name, + last_exception, + ) + raise last_exception + finally: + consumer._callback_task_run = False + + @staticmethod + async def message_received_async(consumer, message: Message) -> None: + async with consumer._message_buffer_lock: # pylint: disable=protected-access + consumer._message_buffer.append(message) # pylint: disable=protected-access + + @staticmethod + async def receive_messages_async(consumer, batch, max_batch_size, max_wait_time): + """ + Receives messages, creates events, and returns them by calling the on received callback. + :param ~azure.eventhub.aio.EventHubConsumer consumer: The EventHubConsumer. + :param bool batch: If receive batch or single event. + :param int max_batch_size: Max batch size. + :param int or None max_wait_time: Max wait time. + """ + # pylint:disable=protected-access + consumer._callback_task_run = True + consumer._last_callback_called_time = time.time() + callback_task = asyncio.create_task( + PyamqpTransportAsync._callback_task(consumer, batch, max_batch_size, max_wait_time) + ) + receive_task = asyncio.create_task(PyamqpTransportAsync._receive_task(consumer)) + + tasks = [callback_task, receive_task] + try: + await asyncio.gather(*tasks) + finally: + consumer._callback_task_run = False + for t in tasks: + if not t.done(): + await asyncio.wait([t], timeout=1) + + @staticmethod + async def create_token_auth_async(auth_uri, get_token, token_type, config, **kwargs): + """ + Creates the JWTTokenAuth. + :param str auth_uri: The auth uri to pass to JWTTokenAuth. + :param get_token: The callback function used for getting and refreshing + tokens. It should return a valid jwt token each time it is called. + :param bytes token_type: Token type. + :param ~azure.eventhub._configuration.Configuration config: EH config. + + :keyword bool update_token: Required. Whether to update token. If not updating token, + then pass 300 to refresh_window. + """ + # TODO: figure out why we're passing all these args to pyamqp JWTTokenAuth, which aren't being used + update_token = kwargs.pop("update_token") # pylint: disable=unused-variable + if update_token: + # update_token not actually needed by pyamqp + # just using to detect wh + return JWTTokenAuthAsync(auth_uri, auth_uri, get_token) + return JWTTokenAuthAsync( + auth_uri, + auth_uri, + get_token, + token_type=token_type, + timeout=config.auth_timeout, + custom_endpoint_hostname=config.custom_endpoint_hostname, + port=config.connection_port, + verify=config.connection_verify, + ) + # if update_token: + # token_auth.update_token() # TODO: why don't we need to update in pyamqp? + + @staticmethod + def create_mgmt_client(address, mgmt_auth, config): # pylint: disable=unused-argument + """ + Creates and returns the mgmt AMQP client. + :param _Address address: Required. The Address. + :param JWTTokenAuth mgmt_auth: Auth for client. + :param ~azure.eventhub._configuration.Configuration config: The configuration. + """ + + return AMQPClientAsync( + config.hostname, + auth=mgmt_auth, + network_trace=config.network_tracing, + transport_type=config.transport_type, + http_proxy=config.http_proxy, + custom_endpoint_address=config.custom_endpoint_address, + connection_verify=config.connection_verify, + ) + + @staticmethod + async def get_updated_token_async(mgmt_auth): + """ + Return updated auth token. + :param mgmt_auth: Auth. + """ + return await mgmt_auth.get_token() + + @staticmethod + async def mgmt_client_request_async(mgmt_client, mgmt_msg, **kwargs): + """ + Send mgmt request. + :param AMQPClientAsync mgmt_client: Client to send request with. + :param str mgmt_msg: Message. + :keyword bytes operation: Operation. + :keyword operation_type: Op type. + :keyword status_code_field: mgmt status code. + :keyword description_fields: mgmt status desc. + """ + operation_type = kwargs.pop("operation_type") + operation = kwargs.pop("operation") + return await mgmt_client.mgmt_request_async( + mgmt_msg, operation=operation.decode(), operation_type=operation_type.decode(), **kwargs + ) + + @staticmethod + async def _handle_exception_async( # pylint:disable=too-many-branches, too-many-statements + exception: Exception, closable: Union["ClientBaseAsync", "ConsumerProducerMixin"], *, is_consumer=False + ) -> Exception: + # pylint: disable=protected-access + if isinstance(exception, asyncio.CancelledError): + raise exception + error = exception + try: + name = cast("ConsumerProducerMixin", closable)._name + except AttributeError: + name = cast("ClientBaseAsync", closable)._container_id + if isinstance(exception, KeyboardInterrupt): # pylint:disable=no-else-raise + _LOGGER.info("%r stops due to keyboard interrupt", name) + await cast("ConsumerProducerMixin", closable)._close_connection_async() + raise error + elif isinstance(exception, EventHubError): + await cast("ConsumerProducerMixin", closable)._close_handler_async() + raise error + # TODO: The following errors seem to be useless in EH + # elif isinstance( + # exception, + # ( + # errors.MessageAccepted, + # errors.MessageAlreadySettled, + # errors.MessageModified, + # errors.MessageRejected, + # errors.MessageReleased, + # errors.MessageContentTooLarge, + # ), + # ): + # _LOGGER.info("%r Event data error (%r)", name, exception) + # error = EventDataError(str(exception), exception) + # raise error + elif isinstance(exception, errors.MessageException): + _LOGGER.info("%r Event data send error (%r)", name, exception) + error = EventDataSendError(str(exception), exception) + raise error + else: + try: + if isinstance(exception, errors.AuthenticationException): + await closable._close_connection_async() # pylint:disable=protected-access + elif isinstance(exception, errors.AMQPLinkError): + await cast("ConsumerProducerMixin", closable)._close_handler_async() # pylint:disable=protected-access + elif isinstance(exception, errors.AMQPConnectionError): + await closable._close_connection_async() # pylint:disable=protected-access + # TODO: add MessageHandlerError in amqp? + # elif isinstance(exception, errors.MessageHandlerError): + # if hasattr(closable, "_close_handler"): + # closable._close_handler() # pylint:disable=protected-access + else: # errors.AMQPConnectionError, compat.TimeoutException + await closable._close_connection_async() # pylint:disable=protected-access + return PyamqpTransportAsync._create_eventhub_exception(exception, is_consumer=is_consumer) + except AttributeError: + pass + return PyamqpTransportAsync._create_eventhub_exception(exception, is_consumer=is_consumer) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_transport/_uamqp_transport_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_transport/_uamqp_transport_async.py index 72b91aee7766..1a1178f51638 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_transport/_uamqp_transport_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_transport/_uamqp_transport_async.py @@ -7,22 +7,26 @@ import asyncio import time import logging -from typing import Union, cast, TYPE_CHECKING, List +from typing import Union, cast, TYPE_CHECKING, List, Optional -from uamqp import ( - constants, - types, - SendClientAsync, - ReceiveClientAsync, - utils, - authentication, - AMQPClientAsync, - errors, -) -from uamqp.async_ops import ConnectionAsync +try: + from uamqp import ( + constants, + types, + SendClientAsync, + ReceiveClientAsync, + utils, + authentication, + AMQPClientAsync, + errors, + ) + from uamqp.async_ops import ConnectionAsync + from ..._transport._uamqp_transport import UamqpTransport + uamqp_installed = True +except ImportError: + uamqp_installed = False from ._base_async import AmqpTransportAsync -from ..._transport._uamqp_transport import UamqpTransport from ...exceptions import ( OperationTimeoutError, EventHubError, @@ -33,340 +37,366 @@ if TYPE_CHECKING: from .._client_base_async import ClientBaseAsync, ConsumerProducerMixin from ..._common import EventData + try: + from uamqp import Message + except ImportError: + Message = None _LOGGER = logging.getLogger(__name__) -class UamqpTransportAsync(UamqpTransport, AmqpTransportAsync): - """ - Class which defines uamqp-based methods used by the producer and consumer. - """ - - @staticmethod - async def create_connection_async(**kwargs): +if uamqp_installed: + class UamqpTransportAsync(UamqpTransport, AmqpTransportAsync): """ - Creates and returns the uamqp async Connection object. - :keyword str host: The hostname, used by uamqp. - :keyword JWTTokenAuth auth: The auth, used by uamqp. - :keyword str endpoint: The endpoint, used by pyamqp. - :keyword str container_id: Required. - :keyword int max_frame_size: Required. - :keyword int channel_max: Required. - :keyword int idle_timeout: Required. - :keyword Dict properties: Required. - :keyword int remote_idle_timeout_empty_frame_send_ratio: Required. - :keyword error_policy: Required. - :keyword bool debug: Required. - :keyword str encoding: Required. + Class which defines uamqp-based methods used by the producer and consumer. """ - endpoint = kwargs.pop("endpoint") # pylint:disable=unused-variable - host = kwargs.pop("host") - auth = kwargs.pop("auth") - return ConnectionAsync( - host, - auth, - **kwargs - ) - @staticmethod - async def close_connection_async(connection): - """ - Closes existing connection. - :param connection: uamqp or pyamqp Connection. - """ - await connection.destroy_async() + @staticmethod + async def create_connection_async(**kwargs): + """ + Creates and returns the uamqp async Connection object. + :keyword str host: The hostname, used by uamqp. + :keyword JWTTokenAuth auth: The auth, used by uamqp. + :keyword str endpoint: The endpoint, used by pyamqp. + :keyword str container_id: Required. + :keyword int max_frame_size: Required. + :keyword int channel_max: Required. + :keyword int idle_timeout: Required. + :keyword Dict properties: Required. + :keyword int remote_idle_timeout_empty_frame_send_ratio: Required. + :keyword error_policy: Required. + :keyword bool debug: Required. + :keyword str encoding: Required. + """ + endpoint = kwargs.pop("endpoint") # pylint:disable=unused-variable + host = kwargs.pop("host") + auth = kwargs.pop("auth") + return ConnectionAsync( + host, + auth, + **kwargs + ) - @staticmethod - def create_send_client(*, config, **kwargs): # pylint:disable=unused-argument - """ - Creates and returns the uamqp SendClient. - :param ~azure.eventhub._configuration.Configuration config: The configuration. + @staticmethod + async def close_connection_async(connection): + """ + Closes existing connection. + :param connection: uamqp or pyamqp Connection. + """ + await connection.destroy_async() - :keyword str target: Required. The target. - :keyword JWTTokenAuth auth: Required. - :keyword int idle_timeout: Required. - :keyword network_trace: Required. - :keyword retry_policy: Required. - :keyword keep_alive_interval: Required. - :keyword str client_name: Required. - :keyword dict link_properties: Required. - :keyword properties: Required. - """ - target = kwargs.pop("target") - retry_policy = kwargs.pop("retry_policy") - network_trace = kwargs.pop("network_trace") + @staticmethod + def create_send_client(*, config, **kwargs): # pylint:disable=unused-argument + """ + Creates and returns the uamqp SendClient. + :param ~azure.eventhub._configuration.Configuration config: The configuration. - return SendClientAsync( - target, - debug=network_trace, # pylint:disable=protected-access - error_policy=retry_policy, - **kwargs - ) + :keyword str target: Required. The target. + :keyword JWTTokenAuth auth: Required. + :keyword int idle_timeout: Required. + :keyword network_trace: Required. + :keyword retry_policy: Required. + :keyword keep_alive_interval: Required. + :keyword str client_name: Required. + :keyword dict link_properties: Required. + :keyword properties: Required. + """ + target = kwargs.pop("target") + retry_policy = kwargs.pop("retry_policy") + network_trace = kwargs.pop("network_trace") - @staticmethod - async def send_messages_async(producer, timeout_time, last_exception, logger): - """ - Handles sending of event data messages. - :param ~azure.eventhub._producer.EventHubProducer producer: The producer with handler to send messages. - :param int timeout_time: Timeout time. - :param last_exception: Exception to raise if message timed out. Only used by uamqp transport. - :param logger: Logger. - """ - # pylint: disable=protected-access - await producer._open() - producer._unsent_events[0].on_send_complete = producer._on_outcome - UamqpTransportAsync._set_msg_timeout(producer, timeout_time, last_exception, logger) - producer._handler.queue_message(*producer._unsent_events) # type: ignore - await producer._handler.wait_async() # type: ignore - producer._unsent_events = producer._handler.pending_messages # type: ignore - if producer._outcome != constants.MessageSendResult.Ok: - if producer._outcome == constants.MessageSendResult.Timeout: - producer._condition = OperationTimeoutError("Send operation timed out") - if producer._condition: - raise producer._condition + return SendClientAsync( + target, + debug=network_trace, # pylint:disable=protected-access + error_policy=retry_policy, + **kwargs + ) - @staticmethod - def create_receive_client(*, config, **kwargs): # pylint:disable=unused-argument - """ - Creates and returns the receive client. - :param ~azure.eventhub._configuration.Configuration config: The configuration. + @staticmethod + async def send_messages_async(producer, timeout_time, last_exception, logger): + """ + Handles sending of event data messages. + :param ~azure.eventhub._producer.EventHubProducer producer: The producer with handler to send messages. + :param int timeout_time: Timeout time. + :param last_exception: Exception to raise if message timed out. Only used by uamqp transport. + :param logger: Logger. + """ + # pylint: disable=protected-access + await producer._open() + producer._unsent_events[0].on_send_complete = producer._on_outcome + UamqpTransportAsync._set_msg_timeout(producer, timeout_time, last_exception, logger) + producer._handler.queue_message(*producer._unsent_events) # type: ignore + await producer._handler.wait_async() # type: ignore + producer._unsent_events = producer._handler.pending_messages # type: ignore + if producer._outcome != constants.MessageSendResult.Ok: + if producer._outcome == constants.MessageSendResult.Timeout: + producer._condition = OperationTimeoutError("Send operation timed out") + if producer._condition: + raise producer._condition - :keyword str source: Required. The source. - :keyword str offset: Required. - :keyword str offset_inclusive: Required. - :keyword JWTTokenAuth auth: Required. - :keyword int idle_timeout: Required. - :keyword network_trace: Required. - :keyword retry_policy: Required. - :keyword str client_name: Required. - :keyword dict link_properties: Required. - :keyword properties: Required. - :keyword link_credit: Required. The prefetch. - :keyword keep_alive_interval: Required. - :keyword desired_capabilities: Required. - :keyword streaming_receive: Required. - :keyword message_received_callback: Required. - :keyword timeout: Required. - """ + @staticmethod + def create_receive_client(*, config, **kwargs): # pylint:disable=unused-argument + """ + Creates and returns the receive client. + :param ~azure.eventhub._configuration.Configuration config: The configuration. - source = kwargs.pop("source") - symbol_array = kwargs.pop("desired_capabilities") - desired_capabilities = None - if symbol_array: - symbol_array = [types.AMQPSymbol(symbol) for symbol in symbol_array] - desired_capabilities = utils.data_factory(types.AMQPArray(symbol_array)) - retry_policy = kwargs.pop("retry_policy") - network_trace = kwargs.pop("network_trace") - link_credit = kwargs.pop("link_credit") - streaming_receive = kwargs.pop("streaming_receive") - message_received_callback = kwargs.pop("message_received_callback") + :keyword str source: Required. The source. + :keyword str offset: Required. + :keyword str offset_inclusive: Required. + :keyword JWTTokenAuth auth: Required. + :keyword int idle_timeout: Required. + :keyword network_trace: Required. + :keyword retry_policy: Required. + :keyword str client_name: Required. + :keyword dict link_properties: Required. + :keyword properties: Required. + :keyword link_credit: Required. The prefetch. + :keyword keep_alive_interval: Required. + :keyword desired_capabilities: Required. + :keyword streaming_receive: Required. + :keyword message_received_callback: Required. + :keyword timeout: Required. + """ - client = ReceiveClientAsync( - source, - debug=network_trace, # pylint:disable=protected-access - error_policy=retry_policy, - desired_capabilities=desired_capabilities, - prefetch=link_credit, - receive_settle_mode=constants.ReceiverSettleMode.ReceiveAndDelete, - auto_complete=False, - **kwargs - ) - # pylint:disable=protected-access - client._streaming_receive = streaming_receive - client._message_received_callback = (message_received_callback) - return client + source = kwargs.pop("source") + symbol_array = kwargs.pop("desired_capabilities") + desired_capabilities = None + if symbol_array: + symbol_array = [types.AMQPSymbol(symbol) for symbol in symbol_array] + desired_capabilities = utils.data_factory(types.AMQPArray(symbol_array)) + retry_policy = kwargs.pop("retry_policy") + network_trace = kwargs.pop("network_trace") + link_credit = kwargs.pop("link_credit") + streaming_receive = kwargs.pop("streaming_receive") + message_received_callback = kwargs.pop("message_received_callback") - @staticmethod - async def receive_messages(consumer, batch, max_batch_size, max_wait_time): - """ - Receives messages, creates events, and returns them by calling the on received callback. - :param ~azure.eventhub.aio.EventHubConsumer consumer: The EventHubConsumer. - :param bool batch: If receive batch or single event. - :param int max_batch_size: Max batch size. - :param int or None max_wait_time: Max wait time. - """ - # pylint:disable=protected-access - max_retries = ( - consumer._client._config.max_retries # pylint:disable=protected-access - ) - has_not_fetched_once = True # ensure one trip when max_wait_time is very small - deadline = time.time() + (max_wait_time or 0) # max_wait_time can be None - while len(consumer._message_buffer) < max_batch_size and ( - time.time() < deadline or has_not_fetched_once - ): - retried_times = 0 - has_not_fetched_once = False - while retried_times <= max_retries: - try: - await consumer._open() - await cast( - ReceiveClientAsync, consumer._handler - ).do_work_async() # uamqp sleeps 0.05 if none received - break - except asyncio.CancelledError: # pylint: disable=try-except-raise - raise - except Exception as exception: # pylint: disable=broad-except - if ( - isinstance(exception, errors.LinkDetach) - and exception.condition == constants.ErrorCodes.LinkStolen # pylint: disable=no-member - ): - raise await consumer._handle_exception(exception) - if not consumer.running: # exit by close - return - if consumer._last_received_event: - consumer._offset = consumer._last_received_event.offset - last_exception = await consumer._handle_exception(exception) - retried_times += 1 - if retried_times > max_retries: - _LOGGER.info( - "%r operation has exhausted retry. Last exception: %r.", - consumer._name, - last_exception, - ) - raise last_exception + client = ReceiveClientAsync( + source, + debug=network_trace, # pylint:disable=protected-access + error_policy=retry_policy, + desired_capabilities=desired_capabilities, + prefetch=link_credit, + receive_settle_mode=constants.ReceiverSettleMode.ReceiveAndDelete, + auto_complete=False, + **kwargs + ) + # pylint:disable=protected-access + client._streaming_receive = streaming_receive + client._message_received_callback = (message_received_callback) + return client + + @staticmethod + def message_received_async(consumer, message: Message) -> None: + consumer._message_buffer.append(message) # pylint: disable=protected-access - if consumer._message_buffer: - while consumer._message_buffer: + @staticmethod + async def receive_messages_async(consumer, batch, max_batch_size, max_wait_time): + """ + Receives messages, creates events, and returns them by calling the on received callback. + :param ~azure.eventhub.aio.EventHubConsumer consumer: The EventHubConsumer. + :param bool batch: If receive batch or single event. + :param int max_batch_size: Max batch size. + :param int or None max_wait_time: Max wait time. + """ + # pylint:disable=protected-access + max_retries = ( + consumer._client._config.max_retries # pylint:disable=protected-access + ) + has_not_fetched_once = True # ensure one trip when max_wait_time is very small + deadline = time.time() + (max_wait_time or 0) # max_wait_time can be None + while len(consumer._message_buffer) < max_batch_size and ( + time.time() < deadline or has_not_fetched_once + ): + retried_times = 0 + has_not_fetched_once = False + while retried_times <= max_retries: + try: + await consumer._open() + await cast( + ReceiveClientAsync, consumer._handler + ).do_work_async() # uamqp sleeps 0.05 if none received + break + except asyncio.CancelledError: # pylint: disable=try-except-raise + raise + except Exception as exception: # pylint: disable=broad-except + if ( + isinstance(exception, errors.LinkDetach) + and exception.condition == constants.ErrorCodes.LinkStolen # pylint: disable=no-member + ): + raise await consumer._handle_exception(exception) + if not consumer.running: # exit by close + return + if consumer._last_received_event: + consumer._offset = consumer._last_received_event.offset + last_exception = await consumer._handle_exception(exception) + retried_times += 1 + if retried_times > max_retries: + _LOGGER.info( + "%r operation has exhausted retry. Last exception: %r.", + consumer._name, + last_exception, + ) + raise last_exception + + if consumer._message_buffer: + while consumer._message_buffer: + if batch: + events_for_callback: List[EventData] = [] + for _ in range(min(max_batch_size, len(consumer._message_buffer))): + events_for_callback.append(consumer._next_message_in_buffer()) + await consumer._on_event_received(events_for_callback) + else: + await consumer._on_event_received(consumer._next_message_in_buffer()) + elif max_wait_time: if batch: - events_for_callback: List[EventData] = [] - for _ in range(min(max_batch_size, len(consumer._message_buffer))): - events_for_callback.append(consumer._next_message_in_buffer()) - await consumer._on_event_received(events_for_callback) + await consumer._on_event_received([]) else: - await consumer._on_event_received(consumer._next_message_in_buffer()) - elif max_wait_time: - if batch: - await consumer._on_event_received([]) - else: - await consumer._on_event_received(None) + await consumer._on_event_received(None) - @staticmethod - async def create_token_auth_async(auth_uri, get_token, token_type, config, **kwargs): - """ - Creates the JWTTokenAuth. - :param str auth_uri: The auth uri to pass to JWTTokenAuth. - :param get_token: The callback function used for getting and refreshing - tokens. It should return a valid jwt token each time it is called. - :param bytes token_type: Token type. - :param ~azure.eventhub._configuration.Configuration config: EH config. + @staticmethod + async def create_token_auth_async(auth_uri, get_token, token_type, config, **kwargs): + """ + Creates the JWTTokenAuth. + :param str auth_uri: The auth uri to pass to JWTTokenAuth. + :param get_token: The callback function used for getting and refreshing + tokens. It should return a valid jwt token each time it is called. + :param bytes token_type: Token type. + :param ~azure.eventhub._configuration.Configuration config: EH config. - :keyword bool update_token: Required. Whether to update token. If not updating token, - then pass 300 to refresh_window. - """ - update_token = kwargs.pop("update_token") - refresh_window = 300 - if update_token: - refresh_window = 0 + :keyword bool update_token: Required. Whether to update token. If not updating token, + then pass 300 to refresh_window. + """ + update_token = kwargs.pop("update_token") + refresh_window = 300 + if update_token: + refresh_window = 0 - token_auth = authentication.JWTTokenAsync( - auth_uri, - auth_uri, - get_token, - token_type=token_type, - timeout=config.auth_timeout, - http_proxy=config.http_proxy, - transport_type=config.transport_type, - custom_endpoint_hostname=config.custom_endpoint_hostname, - port=config.connection_port, - verify=config.connection_verify, - refresh_window=refresh_window - ) - if update_token: - await token_auth.update_token() - return token_auth + token_auth = authentication.JWTTokenAsync( + auth_uri, + auth_uri, + get_token, + token_type=token_type, + timeout=config.auth_timeout, + http_proxy=config.http_proxy, + transport_type=config.transport_type, + custom_endpoint_hostname=config.custom_endpoint_hostname, + port=config.connection_port, + verify=config.connection_verify, + refresh_window=refresh_window + ) + if update_token: + await token_auth.update_token() + return token_auth - @staticmethod - def create_mgmt_client(address, mgmt_auth, config): - """ - Creates and returns the mgmt AMQP client. - :param _Address address: Required. The Address. - :param JWTTokenAuth mgmt_auth: Auth for client. - :param ~azure.eventhub._configuration.Configuration config: The configuration. - """ + @staticmethod + def create_mgmt_client(address, mgmt_auth, config): + """ + Creates and returns the mgmt AMQP client. + :param _Address address: Required. The Address. + :param JWTTokenAuth mgmt_auth: Auth for client. + :param ~azure.eventhub._configuration.Configuration config: The configuration. + """ - mgmt_target = f"amqps://{address.hostname}{address.path}" - return AMQPClientAsync( - mgmt_target, - auth=mgmt_auth, - debug=config.network_tracing - ) + mgmt_target = f"amqps://{address.hostname}{address.path}" + return AMQPClientAsync( + mgmt_target, + auth=mgmt_auth, + debug=config.network_tracing + ) - @staticmethod - async def get_updated_token_async(mgmt_auth): - """ - Return updated auth token. - :param mgmt_auth: Auth. - """ - return mgmt_auth.token + @staticmethod + async def get_updated_token_async(mgmt_auth): + """ + Return updated auth token. + :param mgmt_auth: Auth. + """ + return mgmt_auth.token - @staticmethod - async def mgmt_client_request_async(mgmt_client, mgmt_msg, **kwargs): - """ - Send mgmt request. - :param AMQP Client mgmt_client: Client to send request with. - :param str mgmt_msg: Message. - :keyword bytes operation: Operation. - :keyword operation_type: Op type. - :keyword status_code_field: mgmt status code. - :keyword description_fields: mgmt status desc. - """ - operation_type = kwargs.pop("operation_type") - operation = kwargs.pop("operation") - return await mgmt_client.mgmt_request_async( - mgmt_msg, - operation, - op_type=operation_type, - **kwargs - ) + @staticmethod + async def open_mgmt_client_async(mgmt_client, conn): + """ + Opens the mgmt AMQP client. + :param AMQPClient mgmt_client: uamqp AMQPClient. + :param conn: Connection. + """ + await mgmt_client.open_async(connection=conn) - @staticmethod - async def _handle_exception_async( # pylint:disable=too-many-branches, too-many-statements - exception: Exception, closable: Union["ClientBaseAsync", "ConsumerProducerMixin"] - ) -> Exception: - # pylint: disable=protected-access - if isinstance(exception, asyncio.CancelledError): - raise exception - error = exception - try: - name = cast("ConsumerProducerMixin", closable)._name - except AttributeError: - name = cast("ClientBaseAsync", closable)._container_id - if isinstance(exception, KeyboardInterrupt): # pylint:disable=no-else-raise - _LOGGER.info("%r stops due to keyboard interrupt", name) - await cast("ConsumerProducerMixin", closable)._close_connection_async() - raise error - elif isinstance(exception, EventHubError): - await cast("ConsumerProducerMixin", closable)._close_handler_async() - raise error - elif isinstance( - exception, - ( - errors.MessageAccepted, - errors.MessageAlreadySettled, - errors.MessageModified, - errors.MessageRejected, - errors.MessageReleased, - errors.MessageContentTooLarge, - ), - ): - _LOGGER.info("%r Event data error (%r)", name, exception) - error = EventDataError(str(exception), exception) - raise error - elif isinstance(exception, errors.MessageException): - _LOGGER.info("%r Event data send error (%r)", name, exception) - error = EventDataSendError(str(exception), exception) - raise error - else: + @staticmethod + async def mgmt_client_request_async(mgmt_client, mgmt_msg, **kwargs): + """ + Send mgmt request. + :param AMQP Client mgmt_client: Client to send request with. + :param str mgmt_msg: Message. + :keyword bytes operation: Operation. + :keyword operation_type: Op type. + :keyword status_code_field: mgmt status code. + :keyword description_fields: mgmt status desc. + """ + operation_type = kwargs.pop("operation_type") + operation = kwargs.pop("operation") + response = await mgmt_client.mgmt_request_async( + mgmt_msg, + operation, + op_type=operation_type, + **kwargs + ) + status_code = response.application_properties[kwargs.get("status_code_field")] + description = response.application_properties.get( + kwargs.get("description_fields") + ) # type: Optional[Union[str, bytes]] + return status_code, description, response + + @staticmethod + async def _handle_exception_async( # pylint:disable=too-many-branches, too-many-statements + exception: Exception, + closable: Union["ClientBaseAsync", "ConsumerProducerMixin"], + *, + is_consumer=False # pylint:disable=unused-argument + ) -> Exception: + # pylint: disable=protected-access + if isinstance(exception, asyncio.CancelledError): + raise exception + error = exception try: - if isinstance(exception, errors.AuthenticationException): - await closable._close_connection_async() - elif isinstance(exception, errors.LinkDetach): - await cast("ConsumerProducerMixin", closable)._close_handler_async() - elif isinstance(exception, errors.ConnectionClose): - await closable._close_connection_async() - elif isinstance(exception, errors.MessageHandlerError): - await cast("ConsumerProducerMixin", closable)._close_handler_async() - else: # errors.AMQPConnectionError, compat.TimeoutException, and any other errors - await closable._close_connection_async() + name = cast("ConsumerProducerMixin", closable)._name except AttributeError: - pass - return UamqpTransport._create_eventhub_exception(exception) + name = cast("ClientBaseAsync", closable)._container_id + if isinstance(exception, KeyboardInterrupt): # pylint:disable=no-else-raise + _LOGGER.info("%r stops due to keyboard interrupt", name) + await cast("ConsumerProducerMixin", closable)._close_connection_async() + raise error + elif isinstance(exception, EventHubError): + await cast("ConsumerProducerMixin", closable)._close_handler_async() + raise error + elif isinstance( + exception, + ( + errors.MessageAccepted, + errors.MessageAlreadySettled, + errors.MessageModified, + errors.MessageRejected, + errors.MessageReleased, + errors.MessageContentTooLarge, + ), + ): + _LOGGER.info("%r Event data error (%r)", name, exception) + error = EventDataError(str(exception), exception) + raise error + elif isinstance(exception, errors.MessageException): + _LOGGER.info("%r Event data send error (%r)", name, exception) + error = EventDataSendError(str(exception), exception) + raise error + else: + try: + if isinstance(exception, errors.AuthenticationException): + await closable._close_connection_async() + elif isinstance(exception, errors.LinkDetach): + await cast("ConsumerProducerMixin", closable)._close_handler_async() + elif isinstance(exception, errors.ConnectionClose): + await closable._close_connection_async() + elif isinstance(exception, errors.MessageHandlerError): + await cast("ConsumerProducerMixin", closable)._close_handler_async() + else: # errors.AMQPConnectionError, compat.TimeoutException, and any other errors + await closable._close_connection_async() + except AttributeError: + pass + return UamqpTransport._create_eventhub_exception(exception) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/amqp/_amqp_message.py b/sdk/eventhub/azure-eventhub/azure/eventhub/amqp/_amqp_message.py index d17816e765f7..88ea610c90c3 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/amqp/_amqp_message.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/amqp/_amqp_message.py @@ -22,6 +22,7 @@ class AmqpAnnotatedMessage(object): Please refer to the AMQP spec: http://docs.oasis-open.org/amqp/core/v1.0/os/amqp-core-messaging-v1.0-os.html#section-message-format for more information on the message format. + :keyword data_body: The body consists of one or more data sections and each section contains opaque binary data. :paramtype data_body: Union[str, bytes, List[Union[str, bytes]]] :keyword sequence_body: The body consists of one or more sequence sections and @@ -281,6 +282,7 @@ class AmqpMessageHeader(DictMixin): Please refer to the AMQP spec: http://docs.oasis-open.org/amqp/core/v1.0/os/amqp-core-messaging-v1.0-os.html#type-header for more information on the message header. + :keyword delivery_count: The number of unsuccessful previous attempts to deliver this message. If this value is non-zero it can be taken as an indication that the delivery might be a duplicate. On first delivery, the value is zero. It is @@ -354,6 +356,7 @@ class AmqpMessageProperties(DictMixin): Please refer to the AMQP spec: http://docs.oasis-open.org/amqp/core/v1.0/os/amqp-core-messaging-v1.0-os.html#type-properties for more information on the message properties. + :keyword message_id: Message-id, if set, uniquely identifies a message within the message system. The message producer is usually responsible for setting the message-id in such a way that it is assured to be globally unique. A broker MAY discard a message as a duplicate if the value diff --git a/sdk/eventhub/azure-eventhub/conftest.py b/sdk/eventhub/azure-eventhub/conftest.py index 981fcf68f53a..c2059f6644ff 100644 --- a/sdk/eventhub/azure-eventhub/conftest.py +++ b/sdk/eventhub/azure-eventhub/conftest.py @@ -16,8 +16,15 @@ from azure.mgmt.resource import ResourceManagementClient from azure.mgmt.eventhub import EventHubManagementClient from azure.eventhub import EventHubProducerClient -from uamqp import ReceiveClient -from uamqp.authentication import SASTokenAuth +from azure.eventhub._pyamqp import ReceiveClient +from azure.eventhub._pyamqp.authentication import SASTokenAuth +try: + import uamqp + uamqp_transport_params = [True, False] + uamqp_transport_ids = ["uamqp", "pyamqp"] +except (ModuleNotFoundError, ImportError): + uamqp_transport_params = [False] + uamqp_transport_ids = ["pyamqp"] from devtools_testutils import get_region_override @@ -42,7 +49,7 @@ def sleep(request): sleep = request.config.getoption("--sleep") return sleep.lower() in ('true', 'yes', '1', 'y') -@pytest.fixture(scope="session", params=[True]) +@pytest.fixture(scope="session", params=uamqp_transport_params, ids=uamqp_transport_ids) def uamqp_transport(request): return request.param @@ -198,21 +205,26 @@ def invalid_policy(live_eventhub): @pytest.fixture() -def connstr_receivers(live_eventhub): +def connstr_receivers(live_eventhub, uamqp_transport): connection_str = live_eventhub["connection_str"] partitions = [str(i) for i in range(PARTITION_COUNT)] receivers = [] for p in partitions: uri = "sb://{}/{}".format(live_eventhub['hostname'], live_eventhub['event_hub']) - sas_auth = SASTokenAuth.from_shared_access_key( - uri, live_eventhub['key_name'], live_eventhub['access_key']) - source = "amqps://{}/{}/ConsumerGroups/{}/Partitions/{}".format( live_eventhub['hostname'], live_eventhub['event_hub'], live_eventhub['consumer_group'], p) - receiver = ReceiveClient(source, auth=sas_auth, debug=False, timeout=0, prefetch=500) + if uamqp_transport: + sas_auth = uamqp.authentication.SASTokenAuth.from_shared_access_key( + uri, live_eventhub['key_name'], live_eventhub['access_key']) + receiver = uamqp.ReceiveClient(source, auth=sas_auth, debug=False, timeout=0, prefetch=500) + else: + sas_auth = SASTokenAuth( + uri, uri, live_eventhub['key_name'], live_eventhub['access_key'] + ) + receiver = ReceiveClient(live_eventhub['hostname'], source, auth=sas_auth, network_trace=False, timeout=0, link_credit=500) receiver.open() receivers.append(receiver) yield connection_str, receivers @@ -221,9 +233,9 @@ def connstr_receivers(live_eventhub): @pytest.fixture() -def connstr_senders(live_eventhub): +def connstr_senders(live_eventhub, uamqp_transport): connection_str = live_eventhub["connection_str"] - client = EventHubProducerClient.from_connection_string(connection_str) + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) partitions = client.get_partition_ids() senders = [] diff --git a/sdk/eventhub/azure-eventhub/dev_requirements.txt b/sdk/eventhub/azure-eventhub/dev_requirements.txt index e0e39c64561d..aeb67fc88af1 100644 --- a/sdk/eventhub/azure-eventhub/dev_requirements.txt +++ b/sdk/eventhub/azure-eventhub/dev_requirements.txt @@ -4,4 +4,6 @@ -e ../azure-mgmt-eventhub azure-mgmt-resource==20.0.0 aiohttp>=3.0 +websocket-client -e ../../../tools/azure-devtools +uamqp>=1.6.3,<2.0.0 diff --git a/sdk/eventhub/azure-eventhub/setup.py b/sdk/eventhub/azure-eventhub/setup.py index 30f3dc52f98b..630edc3caea7 100644 --- a/sdk/eventhub/azure-eventhub/setup.py +++ b/sdk/eventhub/azure-eventhub/setup.py @@ -70,7 +70,6 @@ packages=find_packages(exclude=exclude_packages), install_requires=[ "azure-core<2.0.0,>=1.14.0", - "uamqp>=1.6.3,<2.0.0", "typing-extensions>=4.0.1", ] ) diff --git a/sdk/eventhub/azure-eventhub/stress/.gitignore b/sdk/eventhub/azure-eventhub/stress/.gitignore new file mode 100644 index 000000000000..8157b28c778b --- /dev/null +++ b/sdk/eventhub/azure-eventhub/stress/.gitignore @@ -0,0 +1 @@ +generatedValues.yaml \ No newline at end of file diff --git a/sdk/eventhub/azure-eventhub/stress/.helmignore b/sdk/eventhub/azure-eventhub/stress/.helmignore new file mode 100644 index 000000000000..135bd2ba21d8 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/stress/.helmignore @@ -0,0 +1,8 @@ +stress +stress.exe +.env +Dockerfile +*.py +__pycache__ +scripts +*.txt diff --git a/sdk/eventhub/azure-eventhub/stress/Chart.lock b/sdk/eventhub/azure-eventhub/stress/Chart.lock index 0c49b45c0e7a..bd8b41142ec4 100644 --- a/sdk/eventhub/azure-eventhub/stress/Chart.lock +++ b/sdk/eventhub/azure-eventhub/stress/Chart.lock @@ -1,6 +1,6 @@ dependencies: - name: stress-test-addons repository: https://stresstestcharts.blob.core.windows.net/helm/ - version: 0.1.20 -digest: sha256:174a2f4b768cb47718d4b3d5a506330aa781abb31803fbeaeba3b7eef87a9f38 -generated: "2022-07-29T10:58:25.2774398-07:00" + version: 0.2.0 +digest: sha256:53cbe4c0fed047f6c611523bd34181b21a310e7a3a21cb14f649bb09e4a77648 +generated: "2022-11-16T10:34:09.3008939-08:00" diff --git a/sdk/eventhub/azure-eventhub/stress/Chart.yaml b/sdk/eventhub/azure-eventhub/stress/Chart.yaml index ba2319390f2c..a2d7a6ac27a0 100644 --- a/sdk/eventhub/azure-eventhub/stress/Chart.yaml +++ b/sdk/eventhub/azure-eventhub/stress/Chart.yaml @@ -1,5 +1,5 @@ apiVersion: v2 -name: python-eventhubs-stress-test +name: stress-py-eventhubs description: python event hubs stress test. version: 0.1.2 appVersion: v0.2 @@ -9,5 +9,5 @@ annotations: dependencies: - name: stress-test-addons - version: 0.1.20 + version: 0.2.0 repository: https://stresstestcharts.blob.core.windows.net/helm/ diff --git a/sdk/eventhub/azure-eventhub/stress/Dockerfile b/sdk/eventhub/azure-eventhub/stress/Dockerfile38 similarity index 89% rename from sdk/eventhub/azure-eventhub/stress/Dockerfile rename to sdk/eventhub/azure-eventhub/stress/Dockerfile38 index 033b01c3fa0c..6ac4055d743c 100644 --- a/sdk/eventhub/azure-eventhub/stress/Dockerfile +++ b/sdk/eventhub/azure-eventhub/stress/Dockerfile38 @@ -2,6 +2,7 @@ # public OSS users should simply leave this argument blank or ignore its presence entirely ARG REGISTRY="mcr.microsoft.com/mirror/docker/library/" FROM ${REGISTRY}python:3.8-slim-buster +# RUN apt-get -y update && apt-get -y install git WORKDIR /app diff --git a/sdk/eventhub/azure-eventhub/stress/scenarios-matrix.yaml b/sdk/eventhub/azure-eventhub/stress/scenarios-matrix.yaml new file mode 100644 index 000000000000..69b7af99d4b0 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/stress/scenarios-matrix.yaml @@ -0,0 +1,28 @@ +matrix: + image: + - Dockerfile38 + scenarios: + event-async: + testTarget: event-async + event-sync: + testTarget: event-sync + batch-async: + testTarget: batch-async + batch-sync: + testTarget: batch-sync + bplistsync: + testTarget: bplistsync + bpasync: + testTarget: bpasync + bplistasync: + testTarget: bplistasync + bpsync: + testTarget: bpsync + syncwebsockets: + testTarget: syncwebsockets + asyncwebsockets: + testTarget: asyncwebsockets + sync-batch-web: + testTarget: sync-batch-web + async-batch-web: + testTarget: async-batch-web diff --git a/sdk/eventhub/azure-eventhub/stress/scripts/app_insights_metric.py b/sdk/eventhub/azure-eventhub/stress/scripts/app_insights_metric.py index 9f814fdf2845..245f2bfd4735 100644 --- a/sdk/eventhub/azure-eventhub/stress/scripts/app_insights_metric.py +++ b/sdk/eventhub/azure-eventhub/stress/scripts/app_insights_metric.py @@ -22,13 +22,13 @@ def __init__(self, test_name, test_description=None): self.name = test_name self.desc = test_description - events_measure_name = "The number of events handled by " + self.name + events_measure_name = "NumEvents" + self.name events_measure_desc = "The number of events handled by" + self.desc if self.desc else None - memory_measure_name = "memory usage percentage for " + self.name - memory_measure_desc = "memory usage percentage for " + self.desc if self.desc else None - cpu_measure_name = "cpu usage percentage for " + self.name - cpu_measure_desc = "cpu usage percentage for " + self.desc if self.desc else None - error_measure_name = "error count for " + self.name + memory_measure_name = "Memory " + self.name + memory_measure_desc = "Memory usage percentage for " + self.desc if self.desc else None + cpu_measure_name = "Cpu " + self.name + cpu_measure_desc = "Cpu usage percentage for " + self.desc if self.desc else None + error_measure_name = "Errors " + self.name error_measure_desc = "The number of errors happened while running the test for " + self.desc if self.desc else None self.events_measure = measure_module.MeasureInt(events_measure_name, events_measure_desc, "events") diff --git a/sdk/eventhub/azure-eventhub/stress/scripts/azure_eventhub_consumer_stress_async.py b/sdk/eventhub/azure-eventhub/stress/scripts/azure_eventhub_consumer_stress_async.py index 5a0af259d765..a8f8fa3ddca6 100644 --- a/sdk/eventhub/azure-eventhub/stress/scripts/azure_eventhub_consumer_stress_async.py +++ b/sdk/eventhub/azure-eventhub/stress/scripts/azure_eventhub_consumer_stress_async.py @@ -85,9 +85,10 @@ def parse_starting_position(args): parser.add_argument("--aad_tenant_id", help="AAD tenant id") parser.add_argument("--storage_conn_str", help="conn str of storage blob to store ownership and checkpoint data") parser.add_argument("--storage_container_name", help="storage container name to store ownership and checkpoint data") -parser.add_argument("--uamqp_logging_enable", help="uamqp logging enable", action="store_true") +parser.add_argument("--pyamqp_logging_enable", help="pyamqp logging enable", action="store_true") parser.add_argument("--print_console", help="print to console", action="store_true") parser.add_argument("--log_filename", help="log file name", type=str) +parser.add_argument("--uamqp_mode", help="Flag for uamqp or pyamqp", action="store_true") args = parser.parse_args() starting_position = parse_starting_position(args) @@ -190,7 +191,8 @@ def create_client(args): auth_timeout=args.auth_timeout, http_proxy=http_proxy, transport_type=transport_type, - logging_enable=args.uamqp_logging_enable + logging_enable=args.pyamqp_logging_enable, + uamqp_transport=args.uamqp_mode, ) elif args.conn_str: client = EventHubConsumerClientTest.from_connection_string( @@ -202,7 +204,8 @@ def create_client(args): auth_timeout=args.auth_timeout, http_proxy=http_proxy, transport_type=transport_type, - logging_enable=args.uamqp_logging_enable + logging_enable=args.pyamqp_logging_enable, + uamqp_transport=args.uamqp_mode, ) elif args.hostname: client = EventHubConsumerClientTest( @@ -215,7 +218,8 @@ def create_client(args): auth_timeout=args.auth_timeout, http_proxy=http_proxy, transport_type=transport_type, - logging_enable=args.uamqp_logging_enable + logging_enable=args.pyamqp_logging_enable, + uamqp_transport=args.uamqp_mode, ) elif args.aad_client_id: credential = ClientSecretCredential(args.tenant_id, args.aad_client_id, args.aad_secret) @@ -229,7 +233,8 @@ def create_client(args): auth_timeout=args.auth_timeout, http_proxy=http_proxy, transport_type=transport_type, - logging_enable=args.uamqp_logging_enable + logging_enable=args.pyamqp_logging_enable, + uamqp_transport=args.uamqp_mode, ) return client diff --git a/sdk/eventhub/azure-eventhub/stress/scripts/azure_eventhub_consumer_stress_sync.py b/sdk/eventhub/azure-eventhub/stress/scripts/azure_eventhub_consumer_stress_sync.py index bd5770bbe493..f77d56b6b689 100644 --- a/sdk/eventhub/azure-eventhub/stress/scripts/azure_eventhub_consumer_stress_sync.py +++ b/sdk/eventhub/azure-eventhub/stress/scripts/azure_eventhub_consumer_stress_sync.py @@ -83,9 +83,10 @@ def parse_starting_position(args): parser.add_argument("--aad_tenant_id", help="AAD tenant id") parser.add_argument("--storage_conn_str", help="conn str of storage blob to store ownership and checkpoint data") parser.add_argument("--storage_container_name", help="storage container name to store ownership and checkpoint data") -parser.add_argument("--uamqp_logging_enable", help="uamqp logging enable", action="store_true") +parser.add_argument("--pyamqp_logging_enable", help="pyamqp logging enable", action="store_true") parser.add_argument("--print_console", help="print to console", action="store_true") parser.add_argument("--log_filename", help="log file name", type=str) +parser.add_argument("--uamqp_mode", help="Flag for uamqp or pyamqp", action="store_true") args = parser.parse_args() starting_position = parse_starting_position(args) @@ -191,7 +192,8 @@ def create_client(args): auth_timeout=args.auth_timeout, http_proxy=http_proxy, transport_type=transport_type, - logging_enable=args.uamqp_logging_enable + logging_enable=args.pyamqp_logging_enable, + uamqp_transport=args.uamqp_mode, ) elif args.conn_str: client = EventHubConsumerClientTest.from_connection_string( @@ -203,7 +205,8 @@ def create_client(args): auth_timeout=args.auth_timeout, http_proxy=http_proxy, transport_type=transport_type, - logging_enable=args.uamqp_logging_enable + logging_enable=args.pyamqp_logging_enable, + uamqp_transport=args.uamqp_mode, ) elif args.hostname: client = EventHubConsumerClientTest( @@ -216,7 +219,8 @@ def create_client(args): auth_timeout=args.auth_timeout, http_proxy=http_proxy, transport_type=transport_type, - logging_enable=args.uamqp_logging_enable + logging_enable=args.pyamqp_logging_enable, + uamqp_transport=args.uamqp_mode, ) elif args.aad_client_id: credential = ClientSecretCredential(args.tenant_id, args.aad_client_id, args.aad_secret) @@ -230,7 +234,8 @@ def create_client(args): auth_timeout=args.auth_timeout, http_proxy=http_proxy, transport_type=transport_type, - logging_enable=args.uamqp_logging_enable + logging_enable=args.pyamqp_logging_enable, + uamqp_transport=args.uamqp_mode, ) return client diff --git a/sdk/eventhub/azure-eventhub/stress/scripts/azure_eventhub_producer_stress.py b/sdk/eventhub/azure-eventhub/stress/scripts/azure_eventhub_producer_stress.py index 87311ca6b23f..a89c9d8291f3 100644 --- a/sdk/eventhub/azure-eventhub/stress/scripts/azure_eventhub_producer_stress.py +++ b/sdk/eventhub/azure-eventhub/stress/scripts/azure_eventhub_producer_stress.py @@ -32,6 +32,24 @@ def handle_exception(error, ignore_send_failure, stress_logger, azure_monitor_me return 0 raise error +def on_success(events, pid): + # sending succeeded + pass + + +def on_error(events, pid, error): + # sending failed + pass + +async def on_success_async(events, pid): + # sending succeeded + pass + + +async def on_error_async(events, pid, error): + # sending failed + pass + def stress_send_sync(producer: EventHubProducerClient, args, stress_logger, azure_monitor_metric): try: @@ -129,6 +147,7 @@ def __init__(self, argument_parser): action="store_true", help="Whether create new client for each sending", ) + self.argument_parser.add_argument("--buffered_mode", help="buffer producer", action="store_true") self.argument_parser.add_argument("--proxy_hostname", type=str) self.argument_parser.add_argument("--proxy_port", type=str) self.argument_parser.add_argument("--proxy_username", type=str) @@ -140,13 +159,14 @@ def __init__(self, argument_parser): self.argument_parser.add_argument("--aad_secret", help="AAD secret") self.argument_parser.add_argument("--aad_tenant_id", help="AAD tenant id") self.argument_parser.add_argument("--payload", help="payload size", type=int, default=1024) - self.argument_parser.add_argument("--uamqp_logging_enable", help="uamqp logging enable", action="store_true") + self.argument_parser.add_argument("--pyamqp_logging_enable", help="pyamqp logging enable", action="store_true") self.argument_parser.add_argument("--print_console", action="store_true") self.argument_parser.add_argument("--log_filename", help="log file name", type=str) self.argument_parser.add_argument("--retry_total", type=int, default=3) self.argument_parser.add_argument("--retry_backoff_factor", type=float, default=0.8) self.argument_parser.add_argument("--retry_backoff_max", type=float, default=120) self.argument_parser.add_argument("--ignore_send_failure", help="ignore sending failures", action="store_true") + self.argument_parser.add_argument("--uamqp_mode", help="Flag for uamqp or pyamqp", action="store_true") self.args, _ = parser.parse_known_args() if self.args.send_partition_key and self.args.send_partition_id: @@ -163,6 +183,7 @@ def create_client(self, client_class, is_async=False): "retry_backoff_factor": self.args.retry_backoff_factor, "retry_backoff_max": self.args.retry_backoff_max } + if self.args.proxy_hostname: http_proxy = { "proxy_hostname": self.args.proxy_hostname, @@ -170,8 +191,36 @@ def create_client(self, client_class, is_async=False): "username": self.args.proxy_username, "password": self.args.proxy_password, } - - if self.args.azure_identity: + if self.args.buffered_mode: + if is_async: + client = client_class.from_connection_string( + self.args.conn_str, + eventhub_name=self.args.eventhub, + auth_timeout=self.args.auth_timeout, + http_proxy=http_proxy, + transport_type=transport_type, + logging_enable=self.args.pyamqp_logging_enable, + buffered_mode=self.args.buffered_mode, + on_success=on_success_async, + on_error=on_error_async, + uamqp_transport=self.args.uamqp_mode, + **retry_options + ) + else: + client = client_class.from_connection_string( + self.args.conn_str, + eventhub_name=self.args.eventhub, + auth_timeout=self.args.auth_timeout, + http_proxy=http_proxy, + transport_type=transport_type, + logging_enable=self.args.pyamqp_logging_enable, + buffered_mode=self.args.buffered_mode, + on_success=on_success, + on_error=on_error, + uamqp_transport=self.args.uamqp_mode, + **retry_options + ) + elif self.args.azure_identity: print("Using Azure Identity") client = client_class( fully_qualified_namespace=self.args.hostname, @@ -180,7 +229,8 @@ def create_client(self, client_class, is_async=False): auth_timeout=self.args.auth_timeout, http_proxy=http_proxy, transport_type=transport_type, - logging_enable=self.args.uamqp_logging_enable, + logging_enable=self.args.pyamqp_logging_enable, + uamqp_transport=self.args.uamqp_mode, **retry_options ) elif self.args.conn_str: @@ -190,7 +240,8 @@ def create_client(self, client_class, is_async=False): auth_timeout=self.args.auth_timeout, http_proxy=http_proxy, transport_type=transport_type, - logging_enable=self.args.uamqp_logging_enable, + logging_enable=self.args.pyamqp_logging_enable, + uamqp_transport=self.args.uamqp_mode, **retry_options ) elif self.args.hostname: @@ -201,7 +252,8 @@ def create_client(self, client_class, is_async=False): auth_timeout=self.args.auth_timeout, http_proxy=http_proxy, transport_type=transport_type, - logging_enable=self.args.uamqp_logging_enable, + logging_enable=self.args.pyamqp_logging_enable, + uamqp_transport=self.args.uamqp_mode, **retry_options ) elif self.args.aad_client_id: @@ -216,7 +268,8 @@ def create_client(self, client_class, is_async=False): credential=credential, http_proxy=http_proxy, transport_type=transport_type, - logging_enable=self.args.uamqp_logging_enable, + logging_enable=self.args.pyamqp_logging_enable, + uamqp_transport=self.args.uamqp_mode, **retry_options ) else: diff --git a/sdk/eventhub/azure-eventhub/stress/scripts/dev_requirement.txt b/sdk/eventhub/azure-eventhub/stress/scripts/dev_requirement.txt index d3ffc7a78295..b7ff5a1bb209 100644 --- a/sdk/eventhub/azure-eventhub/stress/scripts/dev_requirement.txt +++ b/sdk/eventhub/azure-eventhub/stress/scripts/dev_requirement.txt @@ -2,8 +2,9 @@ psutil azure-eventhub azure-eventhub-checkpointstoreblob azure-eventhub-checkpointstoreblob-aio -azure-servicebus==0.50.3 +azure-servicebus==7.8.1 azure-storage-blob azure-identity opencensus-ext-azure python-dotenv +websocket-client diff --git a/sdk/eventhub/azure-eventhub/stress/stress-test-resources.bicep b/sdk/eventhub/azure-eventhub/stress/stress-test-resources.bicep index 856f6352e2d2..3ff69f8cb492 100644 --- a/sdk/eventhub/azure-eventhub/stress/stress-test-resources.bicep +++ b/sdk/eventhub/azure-eventhub/stress/stress-test-resources.bicep @@ -36,7 +36,7 @@ resource eventHubsNamespace_eventHubName 'Microsoft.EventHub/namespaces/eventhub name: '${eventHubsNamespace_var}/${eventHubName}' location: location properties: { - messageRetentionInDays: 1 + messageRetentionInDays: 5 partitionCount: 32 } dependsOn: [ @@ -104,4 +104,4 @@ output EVENT_HUB_SAS_POLICY string = eventHubAuthRuleName output EVENT_HUB_SAS_KEY string = listkeys(eventHubAuthRuleName, ehVersion).primaryKey output AZURE_STORAGE_CONN_STR string = 'DefaultEndpointsProtocol=https;AccountName=${storageAccount_var};AccountKey=${listKeys(storageAccountId, providers('Microsoft.Storage', 'storageAccounts').apiVersions[0]).keys[0].value};EndpointSuffix=${storageEndpointSuffix}' output AZURE_STORAGE_ACCOUNT string = storageAccount_var -output AZURE_STORAGE_ACCESS_KEY string = listKeys(storageAccountId, providers('Microsoft.Storage', 'storageAccounts').apiVersions[0]).keys[0].value \ No newline at end of file +output AZURE_STORAGE_ACCESS_KEY string = listKeys(storageAccountId, providers('Microsoft.Storage', 'storageAccounts').apiVersions[0]).keys[0].value diff --git a/sdk/eventhub/azure-eventhub/stress/templates/network_loss.yaml b/sdk/eventhub/azure-eventhub/stress/templates/network_loss.yaml index 2a7f5ee0aeb9..cfebba81384c 100644 --- a/sdk/eventhub/azure-eventhub/stress/templates/network_loss.yaml +++ b/sdk/eventhub/azure-eventhub/stress/templates/network_loss.yaml @@ -1,21 +1,31 @@ -{{- include "stress-test-addons.chaos-wrapper.tpl" (list . "stress.python-eh-network-chaos") -}} -{{- define "stress.python-eh-network-chaos" -}} -apiVersion: chaos-mesh.org/v1alpha1 -kind: NetworkChaos - -spec: - action: loss - direction: to - externalTargets: - - '{{ .Stress.ResourceGroupName }}.servicebus.windows.net' - mode: one - selector: - labelSelectors: - testInstance: "eventhub-{{ .Release.Name }}-{{ .Release.Revision }}" - chaos: "true" - namespaces: - - {{ .Release.Namespace }} - loss: - loss: "100" - correlation: "100" -{{- end -}} +# apiVersion: chaos-mesh.org/v1alpha1 +# kind: NetworkChaos +# metadata: +# name: '{{ .Release.Name }}-{{ .Release.Revision }}' +# namespace: {{ .Release.Namespace }} +# annotations: +# experiment.chaos-mesh.org/pause: 'true' +# labels: +# scenario: 'stress' +# release: '{{ .Release.Name }}' +# revision: '{{ .Release.Revision }}' +# spec: +# scheduler: +# cron: '@every 30s' +# duration: '10s' +# action: loss +# direction: to +# externalTargets: +# - 'servicebus.windows.net' +# mode: one +# selector: +# labelSelectors: +# testInstance: "eventhub-{{ .Release.Name }}-{{ .Release.Revision }}" +# chaos: 'true' +# namespaces: +# - {{ .Release.Namespace }} +# podPhaseSelectors: +# - 'Running' +# loss: +# loss: '100' +# correlation: '100' diff --git a/sdk/eventhub/azure-eventhub/stress/templates/testjob.yaml b/sdk/eventhub/azure-eventhub/stress/templates/testjob.yaml index e73b4f3bd699..d53032efd66c 100644 --- a/sdk/eventhub/azure-eventhub/stress/templates/testjob.yaml +++ b/sdk/eventhub/azure-eventhub/stress/templates/testjob.yaml @@ -4,55 +4,64 @@ metadata: labels: testName: "deploy-python-eh-stress" testInstance: "eventhub-{{ .Release.Name }}-{{ .Release.Revision }}" - chaos: "true" spec: + nodeSelector: + sku: 'd4v4' containers: - name: python-eh-stress - image: {{ .Values.image }} + image: {{ .Stress.imageTag }} imagePullPolicy: Always + resources: + limits: + memory: "2000Mi" + cpu: "1" - {{ if eq .Stress.Scenario "identity" }} - command: ['bash', '-c', 'python azure_eventhub_producer_stress.py -m stress_send_sync --azure_identity True --uamqp_logging_enable --print_console --duration 7200'] + {{ if eq .Stress.testTarget "event-async" }} + command: ['bash', '-c', 'python azure_eventhub_producer_stress.py -m stress_send_async --duration 259200 & python azure_eventhub_consumer_stress_async.py --duration 259200 '] {{- end -}} - {{ if eq .Stress.Scenario "sendsync" }} - command: ['bash', '-c', 'python azure_eventhub_producer_stress.py -m stress_send_sync --duration 7200'] + {{ if eq .Stress.testTarget "event-sync" }} + command: ['bash', '-c', 'python azure_eventhub_producer_stress.py -m stress_send_sync --duration 259200 & python azure_eventhub_consumer_stress_sync.py --duration 259200 '] {{- end -}} - {{ if eq .Stress.Scenario "sendlistsync" }} - command: ['bash', '-c', 'python azure_eventhub_producer_stress.py -m stress_send_list_sync --duration 7200'] + {{ if eq .Stress.testTarget "batch-async" }} + command: ['bash', '-c', 'python azure_eventhub_producer_stress.py -m stress_send_list_async --duration 259200 & python azure_eventhub_consumer_stress_async.py --duration 259200 '] {{- end -}} - {{ if eq .Stress.Scenario "sendasync" }} - command: ['bash', '-c', 'python azure_eventhub_producer_stress.py -m stress_send_sync --duration 7200'] + {{ if eq .Stress.testTarget "batch-sync" }} + command: ['bash', '-c', 'python azure_eventhub_producer_stress.py -m stress_send_list_sync --duration 259200 & python azure_eventhub_consumer_stress_sync.py --duration 259200 '] {{- end -}} - {{ if eq .Stress.Scenario "sendlistasync" }} - command: ['bash', '-c', 'python azure_eventhub_producer_stress.py -m stress_send_list_async --duration 7200'] + {{ if eq .Stress.testTarget "bplistsync" }} + command: ['bash', '-c', 'python azure_eventhub_producer_stress.py -m stress_send_list_sync --duration 259200 --buffered_mode & python azure_eventhub_consumer_stress_sync.py --duration 259200 '] {{- end -}} - {{ if eq .Stress.Scenario "sendconsumeasync" }} - command: ['bash', '-c', 'python azure_eventhub_producer_stress.py -m stress_send_async --duration 7200 & python azure_eventhub_consumer_stress_async.py --duration 7200'] + {{ if eq .Stress.testTarget "bpasync" }} + command: ['bash', '-c', 'python azure_eventhub_producer_stress.py -m stress_send_async --duration 259200 --buffered_mode & python azure_eventhub_consumer_stress_async.py --duration 259200'] {{- end -}} - {{ if eq .Stress.Scenario "sendconsumesync" }} - command: ['bash', '-c', 'python azure_eventhub_producer_stress.py -m stress_send_sync --duration 7200 & python azure_eventhub_consumer_stress_sync.py --duration 7200'] + {{ if eq .Stress.testTarget "bplistasync" }} + command: ['bash', '-c', 'python azure_eventhub_producer_stress.py -m stress_send_list_async --duration 259200 --buffered_mode & python azure_eventhub_consumer_stress_async.py --duration 259200 '] {{- end -}} - {{ if eq .Stress.Scenario "sendlistconsumeasync" }} - command: ['bash', '-c', 'python azure_eventhub_producer_stress.py -m stress_send_list_async --duration 7200 & python azure_eventhub_consumer_stress_async.py --duration 7200'] + {{ if eq .Stress.testTarget "bpsync" }} + command: ['bash', '-c', 'python azure_eventhub_producer_stress.py -m stress_send_sync --duration 259200 --buffered_mode & python azure_eventhub_consumer_stress_sync.py --duration 259200'] {{- end -}} - {{ if eq .Stress.Scenario "sendlistconsumesync" }} - command: ['bash', '-c', 'python azure_eventhub_producer_stress.py -m stress_send_list_sync --duration 7200 & python azure_eventhub_consumer_stress_sync.py --duration 7200'] + {{ if eq .Stress.testTarget "syncwebsockets" }} + command: ['bash', '-c', 'python azure_eventhub_producer_stress.py -m stress_send_sync --duration 259200 --transport_type 1 & python azure_eventhub_consumer_stress_sync.py --duration 259200 --transport_type 1'] {{- end -}} - {{ if eq .Stress.Scenario "consumeasyncidentity" }} - command: ['bash', '-c', 'python azure_eventhub_consumer_stress_async.py --azure_identity True --duration 7200'] + {{ if eq .Stress.testTarget "asyncwebsockets" }} + command: ['bash', '-c', 'python azure_eventhub_producer_stress.py -m stress_send_async --duration 259200 --transport_type 1 & python azure_eventhub_consumer_stress_async.py --duration 259200 --transport_type 1'] {{- end -}} - {{ if eq .Stress.Scenario "consumesyncidentity" }} - command: ['bash', '-c', 'python azure_eventhub_consumer_stress_sync.py --azure_identity True --duration 7200'] + {{ if eq .Stress.testTarget "sync-batch-web" }} + command: ['bash', '-c', 'python azure_eventhub_producer_stress.py -m stress_send_list_sync --duration 259200 --transport_type 1 & python azure_eventhub_consumer_stress_sync.py --duration 259200 --transport_type 1'] + {{- end -}} + + {{ if eq .Stress.testTarget "async-batch-web" }} + command: ['bash', '-c', 'python azure_eventhub_producer_stress.py -m stress_send_list_async --duration 259200 --transport_type 1 & python azure_eventhub_consumer_stress_async.py --duration 259200 --transport_type 1'] {{- end -}} {{- include "stress-test-addons.container-env" . | nindent 6 }} diff --git a/sdk/eventhub/azure-eventhub/stress/values.yaml b/sdk/eventhub/azure-eventhub/stress/values.yaml index e106bba45539..1cefb66bcd98 100644 --- a/sdk/eventhub/azure-eventhub/stress/values.yaml +++ b/sdk/eventhub/azure-eventhub/stress/values.yaml @@ -1,15 +1,16 @@ # Optional list of scenarios. If specified multiple stress test jobs will be generated, # one for each scenario in the list. The pod spec can then be configured to pass the # scenario name down to the test command, e.g. `command: ["node", "{{ .Scenario }}.js"]` -scenarios: -- "identity" -- "sendsync" -- "sendlistsync" -- "sendasync" -- "sendlistasync" -- "sendconsumeasync" -- "sendconsumesync" -- "sendlistconsumeasync" -- "sendlistconsumesync" -- "consumeasyncidentity" -- "consumesyncidentity" \ No newline at end of file +# scenarios: +# - "event-async" +# - "event-sync" +# - "batch-async" +# - "batch-sync" +# - "bufferedproducerlistsync" +# - "bufferedproducerasync" +# - "bufferedproducerlistasync" +# - "bufferedproducersync" +# - "syncwebsockets" +# - "asyncwebsockets" +# - "sync-batch-web" +# - "async-batch-web" diff --git a/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_auth_async.py b/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_auth_async.py index 9dc24bfd514e..3193db8f7f96 100644 --- a/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_auth_async.py +++ b/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_auth_async.py @@ -16,24 +16,30 @@ @pytest.mark.liveTest @pytest.mark.asyncio -async def test_client_secret_credential_async(live_eventhub): +async def test_client_secret_credential_async(live_eventhub, uamqp_transport): credential = EnvironmentCredential() producer_client = EventHubProducerClient(fully_qualified_namespace=live_eventhub['hostname'], eventhub_name=live_eventhub['event_hub'], credential=credential, - user_agent='customized information') + user_agent='customized information', + auth_timeout=3, + uamqp_transport=uamqp_transport + ) consumer_client = EventHubConsumerClient(fully_qualified_namespace=live_eventhub['hostname'], eventhub_name=live_eventhub['event_hub'], consumer_group='$default', credential=credential, - user_agent='customized information') + user_agent='customized information', + auth_timeout=3, + uamqp_transport=uamqp_transport + ) async with producer_client: batch = await producer_client.create_batch(partition_id='0') batch.add(EventData(body='A single message')) await producer_client.send_batch(batch) - def on_event(partition_context, event): + async def on_event(partition_context, event): on_event.called = True on_event.partition_id = partition_context.partition_id on_event.event = event @@ -49,11 +55,11 @@ def on_event(partition_context, event): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_client_sas_credential_async(live_eventhub): +async def test_client_sas_credential_async(live_eventhub, uamqp_transport): # This should "just work" to validate known-good. hostname = live_eventhub['hostname'] producer_client = EventHubProducerClient.from_connection_string(live_eventhub['connection_str'], - eventhub_name=live_eventhub['event_hub']) + eventhub_name=live_eventhub['event_hub'], uamqp_transport=uamqp_transport) async with producer_client: batch = await producer_client.create_batch(partition_id='0') @@ -66,7 +72,8 @@ async def test_client_sas_credential_async(live_eventhub): token = (await credential.get_token(auth_uri)).token producer_client = EventHubProducerClient(fully_qualified_namespace=hostname, eventhub_name=live_eventhub['event_hub'], - credential=EventHubSASTokenCredential(token, time.time() + 3000)) + credential=EventHubSASTokenCredential(token, time.time() + 3000), + uamqp_transport=uamqp_transport) async with producer_client: batch = await producer_client.create_batch(partition_id='0') @@ -76,7 +83,7 @@ async def test_client_sas_credential_async(live_eventhub): # Finally let's do it with SAS token + conn str token_conn_str = "Endpoint=sb://{}/;SharedAccessSignature={};".format(hostname, token.decode()) conn_str_producer_client = EventHubProducerClient.from_connection_string(token_conn_str, - eventhub_name=live_eventhub['event_hub']) + eventhub_name=live_eventhub['event_hub'], uamqp_transport=uamqp_transport) async with conn_str_producer_client: batch = await conn_str_producer_client.create_batch(partition_id='0') @@ -86,10 +93,10 @@ async def test_client_sas_credential_async(live_eventhub): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_client_azure_sas_credential_async(live_eventhub): +async def test_client_azure_sas_credential_async(live_eventhub, uamqp_transport): # This should "just work" to validate known-good. hostname = live_eventhub['hostname'] - producer_client = EventHubProducerClient.from_connection_string(live_eventhub['connection_str'], eventhub_name = live_eventhub['event_hub']) + producer_client = EventHubProducerClient.from_connection_string(live_eventhub['connection_str'], eventhub_name = live_eventhub['event_hub'], uamqp_transport=uamqp_transport) async with producer_client: batch = await producer_client.create_batch(partition_id='0') @@ -101,7 +108,8 @@ async def test_client_azure_sas_credential_async(live_eventhub): token = (await credential.get_token(auth_uri)).token.decode() producer_client = EventHubProducerClient(fully_qualified_namespace=hostname, eventhub_name=live_eventhub['event_hub'], - credential=AzureSasCredential(token)) + auth_timeout=3, + credential=AzureSasCredential(token), uamqp_transport=uamqp_transport) async with producer_client: batch = await producer_client.create_batch(partition_id='0') @@ -111,14 +119,15 @@ async def test_client_azure_sas_credential_async(live_eventhub): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_client_azure_named_key_credential_async(live_eventhub): +async def test_client_azure_named_key_credential_async(live_eventhub, uamqp_transport): credential = AzureNamedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']) consumer_client = EventHubConsumerClient(fully_qualified_namespace=live_eventhub['hostname'], eventhub_name=live_eventhub['event_hub'], consumer_group='$default', credential=credential, - user_agent='customized information') + auth_timeout=3, + user_agent='customized information', uamqp_transport=uamqp_transport) assert (await consumer_client.get_eventhub_properties()) is not None diff --git a/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_buffered_producer_async.py b/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_buffered_producer_async.py index cf88b9dd54e6..81983c71a356 100644 --- a/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_buffered_producer_async.py +++ b/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_buffered_producer_async.py @@ -38,25 +38,26 @@ async def random_pkey_generation(partitions): @pytest.mark.liveTest() @pytest.mark.asyncio -async def test_producer_client_constructor(connection_str): +async def test_producer_client_constructor(connection_str, uamqp_transport): async def on_success(events, pid): pass async def on_error(events, error, pid): pass with pytest.raises(TypeError): - EventHubProducerClient.from_connection_string(connection_str, buffered_mode=True) + EventHubProducerClient.from_connection_string(connection_str, buffered_mode=True, uamqp_transport=uamqp_transport) with pytest.raises(TypeError): - EventHubProducerClient.from_connection_string(connection_str, buffered_mode=True, on_success=on_success) + EventHubProducerClient.from_connection_string(connection_str, buffered_mode=True, on_success=on_success, uamqp_transport=uamqp_transport) with pytest.raises(TypeError): - EventHubProducerClient.from_connection_string(connection_str, buffered_mode=True, on_error=on_error) + EventHubProducerClient.from_connection_string(connection_str, buffered_mode=True, on_error=on_error, uamqp_transport=uamqp_transport) with pytest.raises(ValueError): EventHubProducerClient.from_connection_string( connection_str, buffered_mode=True, on_success=on_success, on_error=on_error, - max_wait_time=0 + max_wait_time=0, + uamqp_transport=uamqp_transport ) with pytest.raises(ValueError): EventHubProducerClient.from_connection_string( @@ -64,9 +65,34 @@ async def on_error(events, error, pid): buffered_mode=True, on_success=on_success, on_error=on_error, - max_buffer_length=0 + max_buffer_length=0, + uamqp_transport=uamqp_transport ) + def on_success_missing_params(events): + on_success_missing_params.events = events + + def on_error_missing_params(events, pid): + on_error_missing_params.events = events + + producer = EventHubProducerClient.from_connection_string( + connection_str, + buffered_mode=True, + buffer_concurrency=2, + on_success=on_success_missing_params, + on_error=on_error_missing_params, + uamqp_transport=uamqp_transport, + ) + + on_success_missing_params.events = None + on_error_missing_params.events = None + + # successfully send, but don't enter invalid callback + async with producer: + await producer.send_event(EventData('Single data')) + + assert not on_success_missing_params.events + assert not on_error_missing_params.events @pytest.mark.liveTest @pytest.mark.asyncio @@ -78,15 +104,16 @@ async def on_error(events, error, pid): (False, True) ] ) -async def test_basic_send_single_events_round_robin(connection_str, flush_after_sending, close_after_sending): +async def test_basic_send_single_events_round_robin(connection_str, flush_after_sending, close_after_sending, uamqp_transport): received_events = defaultdict(list) async def on_event(partition_context, event): received_events[partition_context.partition_id].append(event) - consumer = EventHubConsumerClient.from_connection_string(connection_str, consumer_group="$default") + consumer = EventHubConsumerClient.from_connection_string(connection_str, consumer_group="$default", uamqp_transport=uamqp_transport) receive_thread = asyncio.ensure_future(consumer.receive(on_event=on_event)) + await asyncio.sleep(5) sent_events = defaultdict(list) async def on_success(events, pid): @@ -103,7 +130,8 @@ async def on_error(events, pid, err): connection_str, buffered_mode=True, on_success=on_success, - on_error=on_error + on_error=on_error, + uamqp_transport=uamqp_transport ) async with producer: @@ -182,15 +210,16 @@ async def on_error(events, pid, err): (False, False) ] ) -async def test_basic_send_batch_events_round_robin(connection_str, flush_after_sending, close_after_sending): +async def test_basic_send_batch_events_round_robin(connection_str, flush_after_sending, close_after_sending, uamqp_transport): received_events = defaultdict(list) async def on_event(partition_context, event): received_events[partition_context.partition_id].append(event) - consumer = EventHubConsumerClient.from_connection_string(connection_str, consumer_group="$default") + consumer = EventHubConsumerClient.from_connection_string(connection_str, consumer_group="$default", uamqp_transport=uamqp_transport) receive_thread = asyncio.ensure_future(consumer.receive(on_event=on_event)) + await asyncio.sleep(5) sent_events = defaultdict(list) async def on_success(events, pid): @@ -204,7 +233,8 @@ async def on_error(events, pid, err): connection_str, buffered_mode=True, on_success=on_success, - on_error=on_error + on_error=on_error, + uamqp_transport=uamqp_transport ) async with producer: @@ -288,13 +318,13 @@ async def on_error(events, pid, err): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_with_hybrid_partition_assignment(connection_str): +async def test_send_with_hybrid_partition_assignment(connection_str, uamqp_transport): received_events = defaultdict(list) async def on_event(partition_context, event): received_events[partition_context.partition_id].append(event) - consumer = EventHubConsumerClient.from_connection_string(connection_str, consumer_group="$default") + consumer = EventHubConsumerClient.from_connection_string(connection_str, consumer_group="$default", uamqp_transport=uamqp_transport) receive_thread = asyncio.ensure_future(consumer.receive(on_event=on_event)) sent_events = defaultdict(list) @@ -310,7 +340,8 @@ async def on_error(events, pid, err): connection_str, buffered_mode=True, on_success=on_success, - on_error=on_error + on_error=on_error, + uamqp_transport=uamqp_transport ) async with producer: @@ -377,13 +408,13 @@ async def on_error(events, pid, err): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_with_timing_configuration(connection_str): +async def test_send_with_timing_configuration(connection_str, uamqp_transport): received_events = defaultdict(list) async def on_event(partition_context, event): received_events[partition_context.partition_id].append(event) - consumer = EventHubConsumerClient.from_connection_string(connection_str, consumer_group="$default") + consumer = EventHubConsumerClient.from_connection_string(connection_str, consumer_group="$default", uamqp_transport=uamqp_transport) receive_thread = asyncio.ensure_future(consumer.receive(on_event=on_event)) sent_events = defaultdict(list) @@ -402,7 +433,8 @@ async def on_error(events, pid, err): buffered_mode=True, max_wait_time=10, on_success=on_success, - on_error=on_error + on_error=on_error, + uamqp_transport=uamqp_transport ) async with producer: @@ -422,7 +454,8 @@ async def on_error(events, pid, err): max_wait_time=1000, max_buffer_length=10, on_success=on_success, - on_error=on_error + on_error=on_error, + uamqp_transport=uamqp_transport ) sent_events.clear() @@ -452,13 +485,13 @@ async def on_error(events, pid, err): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_long_sleep(connection_str): +async def test_long_sleep(connection_str, uamqp_transport): received_events = defaultdict(list) async def on_event(partition_context, event): received_events[partition_context.partition_id].append(event) - consumer = EventHubConsumerClient.from_connection_string(connection_str, consumer_group="$default") + consumer = EventHubConsumerClient.from_connection_string(connection_str, consumer_group="$default", uamqp_transport=uamqp_transport) receive_thread = asyncio.ensure_future(consumer.receive(on_event=on_event)) @@ -475,7 +508,8 @@ async def on_error(events, pid, err): connection_str, buffered_mode=True, on_success=on_success, - on_error=on_error + on_error=on_error, + uamqp_transport=uamqp_transport ) async with producer: @@ -494,13 +528,13 @@ async def on_error(events, pid, err): @pytest.mark.skip('not testing correctly + flaky, fix during MQ') @pytest.mark.liveTest @pytest.mark.asyncio -async def test_long_wait_small_buffer(connection_str): +async def test_long_wait_small_buffer(connection_str, uamqp_transport): received_events = defaultdict(list) async def on_event(partition_context, event): received_events[partition_context.partition_id].append(event) - consumer = EventHubConsumerClient.from_connection_string(connection_str, consumer_group="$default") + consumer = EventHubConsumerClient.from_connection_string(connection_str, consumer_group="$default", uamqp_transport=uamqp_transport) receive_thread = asyncio.ensure_future(consumer.receive(on_event=on_event)) @@ -523,7 +557,8 @@ async def on_error(events, pid, err): retry_mode='fixed', retry_backoff_factor=0.01, max_wait_time=10, - max_buffer_length=100 + max_buffer_length=100, + uamqp_transport=uamqp_transport ) async with producer: diff --git a/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_consumer_client_async.py b/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_consumer_client_async.py index 0ff6ae711cbf..9831cb9db439 100644 --- a/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_consumer_client_async.py +++ b/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_consumer_client_async.py @@ -8,11 +8,11 @@ @pytest.mark.liveTest @pytest.mark.asyncio -async def test_receive_no_partition_async(connstr_senders): +async def test_receive_no_partition_async(connstr_senders, uamqp_transport): connection_str, senders = connstr_senders senders[0].send(EventData("Test EventData")) senders[1].send(EventData("Test EventData")) - client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default') + client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default', uamqp_transport=uamqp_transport) async def on_event(partition_context, event): on_event.received += 1 @@ -49,10 +49,10 @@ async def on_event(partition_context, event): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_receive_partition_async(connstr_senders): +async def test_receive_partition_async(connstr_senders, uamqp_transport): connection_str, senders = connstr_senders senders[0].send(EventData("Test EventData")) - client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default') + client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default', uamqp_transport=uamqp_transport) async def on_event(partition_context, event): assert partition_context.partition_id == "0" @@ -72,13 +72,13 @@ async def on_event(partition_context, event): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_receive_load_balancing_async(connstr_senders): +async def test_receive_load_balancing_async(connstr_senders, uamqp_transport): connection_str, senders = connstr_senders cs = InMemoryCheckpointStore() client1 = EventHubConsumerClient.from_connection_string( - connection_str, consumer_group='$default', checkpoint_store=cs, load_balancing_interval=1) + connection_str, consumer_group='$default', checkpoint_store=cs, load_balancing_interval=1, uamqp_transport=uamqp_transport) client2 = EventHubConsumerClient.from_connection_string( - connection_str, consumer_group='$default', checkpoint_store=cs, load_balancing_interval=1) + connection_str, consumer_group='$default', checkpoint_store=cs, load_balancing_interval=1, uamqp_transport=uamqp_transport) async def on_event(partition_context, event): pass @@ -98,13 +98,13 @@ async def on_event(partition_context, event): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_receive_batch_no_max_wait_time_async(connstr_senders): +async def test_receive_batch_no_max_wait_time_async(connstr_senders, uamqp_transport): '''Test whether callback is called when max_wait_time is None and max_batch_size has reached ''' connection_str, senders = connstr_senders senders[0].send(EventData("Test EventData")) senders[1].send(EventData("Test EventData")) - client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default') + client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default', uamqp_transport=uamqp_transport) async def on_event_batch(partition_context, event_batch): on_event_batch.received += len(event_batch) @@ -144,10 +144,10 @@ async def on_event_batch(partition_context, event_batch): ]) @pytest.mark.liveTest @pytest.mark.asyncio -async def test_receive_batch_empty_with_max_wait_time_async(connection_str, max_wait_time, sleep_time, expected_result): +async def test_receive_batch_empty_with_max_wait_time_async(connection_str, max_wait_time, sleep_time, expected_result, uamqp_transport): '''Test whether event handler is called when max_wait_time > 0 and no event is received ''' - client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default') + client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default', uamqp_transport=uamqp_transport) async def on_event_batch(partition_context, event_batch): on_event_batch.event_batch = event_batch @@ -163,13 +163,13 @@ async def on_event_batch(partition_context, event_batch): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_receive_batch_early_callback_async(connstr_senders): +async def test_receive_batch_early_callback_async(connstr_senders, uamqp_transport): ''' Test whether the callback is called once max_batch_size reaches and before max_wait_time reaches. ''' connection_str, senders = connstr_senders for _ in range(10): senders[0].send(EventData("Test EventData")) - client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default') + client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default', uamqp_transport=uamqp_transport) async def on_event_batch(partition_context, event_batch): on_event_batch.received += len(event_batch) diff --git a/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_negative_async.py b/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_negative_async.py index 4d5880d9d2dc..b278f528e9c4 100644 --- a/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_negative_async.py +++ b/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_negative_async.py @@ -7,28 +7,31 @@ import asyncio import pytest import sys +import time from azure.eventhub import ( EventData, EventDataBatch, ) +from azure.identity.aio import EnvironmentCredential from azure.eventhub.exceptions import ( EventHubError, ConnectError, AuthenticationError, - EventDataSendError + OperationTimeoutError ) -from azure.eventhub.aio import EventHubConsumerClient, EventHubProducerClient +from azure.eventhub.aio import EventHubConsumerClient, EventHubProducerClient, EventHubSharedKeyCredential +from azure.eventhub.aio._client_base_async import EventHubSASTokenCredential @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_with_invalid_hostname_async(invalid_hostname, connstr_receivers): +async def test_send_with_invalid_hostname_async(invalid_hostname, connstr_receivers, uamqp_transport): if sys.platform.startswith('darwin'): pytest.skip("Skipping on OSX - it keeps reporting 'Unable to set external certificates' " "and blocking other tests") _, receivers = connstr_receivers - client = EventHubProducerClient.from_connection_string(invalid_hostname) + client = EventHubProducerClient.from_connection_string(invalid_hostname, uamqp_transport=uamqp_transport) async with client: with pytest.raises(ConnectError): batch = EventDataBatch() @@ -42,7 +45,7 @@ async def on_error(events, pid, err): on_error.err = err on_error.err = None - client = EventHubProducerClient.from_connection_string(invalid_hostname, on_error=on_error) + client = EventHubProducerClient.from_connection_string(invalid_hostname, on_error=on_error, uamqp_transport=uamqp_transport) async with client: batch = EventDataBatch() batch.add(EventData("test data")) @@ -50,7 +53,7 @@ async def on_error(events, pid, err): assert isinstance(on_error.err, ConnectError) on_error.err = None - client = EventHubProducerClient.from_connection_string(invalid_hostname, on_error=on_error) + client = EventHubProducerClient.from_connection_string(invalid_hostname, on_error=on_error, uamqp_transport=uamqp_transport) async with client: await client.send_event(EventData("test data")) assert isinstance(on_error.err, ConnectError) @@ -60,7 +63,7 @@ async def on_error(events, pid, err): ["hostname", "key_name", "access_key", "event_hub", "partition"]) @pytest.mark.liveTest @pytest.mark.asyncio -async def test_receive_with_invalid_param_async(live_eventhub, invalid_place): +async def test_receive_with_invalid_param_async(live_eventhub, invalid_place, uamqp_transport): eventhub_config = live_eventhub.copy() if invalid_place != "partition": eventhub_config[invalid_place] = "invalid " + invalid_place @@ -70,7 +73,7 @@ async def test_receive_with_invalid_param_async(live_eventhub, invalid_place): eventhub_config['access_key'], eventhub_config['event_hub']) - client = EventHubConsumerClient.from_connection_string(conn_str, consumer_group='$default', retry_total=0) + client = EventHubConsumerClient.from_connection_string(conn_str, consumer_group='$default', retry_total=0, uamqp_transport=uamqp_transport) async def on_event(partition_context, event): pass @@ -89,8 +92,8 @@ async def on_event(partition_context, event): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_with_invalid_key_async(invalid_key): - client = EventHubProducerClient.from_connection_string(invalid_key) +async def test_send_with_invalid_key_async(invalid_key, uamqp_transport): + client = EventHubProducerClient.from_connection_string(invalid_key, uamqp_transport=uamqp_transport) async with client: with pytest.raises(ConnectError): batch = EventDataBatch() @@ -100,8 +103,8 @@ async def test_send_with_invalid_key_async(invalid_key): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_with_invalid_policy_async(invalid_policy): - client = EventHubProducerClient.from_connection_string(invalid_policy) +async def test_send_with_invalid_policy_async(invalid_policy, uamqp_transport): + client = EventHubProducerClient.from_connection_string(invalid_policy, uamqp_transport=uamqp_transport) async with client: with pytest.raises(ConnectError): batch = EventDataBatch() @@ -111,8 +114,8 @@ async def test_send_with_invalid_policy_async(invalid_policy): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_non_existing_entity_sender_async(connection_str): - client = EventHubProducerClient.from_connection_string(connection_str, eventhub_name="nemo") +async def test_non_existing_entity_sender_async(connection_str, uamqp_transport): + client = EventHubProducerClient.from_connection_string(connection_str, eventhub_name="nemo", uamqp_transport=uamqp_transport) async with client: with pytest.raises(ConnectError): batch = EventDataBatch() @@ -122,10 +125,10 @@ async def test_non_existing_entity_sender_async(connection_str): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_to_invalid_partitions_async(connection_str): +async def test_send_to_invalid_partitions_async(connection_str, uamqp_transport): partitions = ["XYZ", "-1", "1000", "-"] for p in partitions: - client = EventHubProducerClient.from_connection_string(connection_str) + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) try: with pytest.raises(ConnectError): batch = await client.create_batch(partition_id=p) @@ -137,10 +140,10 @@ async def test_send_to_invalid_partitions_async(connection_str): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_too_large_message_async(connection_str): +async def test_send_too_large_message_async(connection_str, uamqp_transport): if sys.platform.startswith('darwin'): pytest.skip("Skipping on OSX - open issue regarding message size") - client = EventHubProducerClient.from_connection_string(connection_str) + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) try: data = EventData(b"A" * 1100000) with pytest.raises(ValueError): @@ -152,8 +155,8 @@ async def test_send_too_large_message_async(connection_str): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_null_body_async(connection_str): - client = EventHubProducerClient.from_connection_string(connection_str) +async def test_send_null_body_async(connection_str, uamqp_transport): + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) try: with pytest.raises(ValueError): data = EventData(None) @@ -166,11 +169,11 @@ async def test_send_null_body_async(connection_str): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_create_batch_with_invalid_hostname_async(invalid_hostname): +async def test_create_batch_with_invalid_hostname_async(invalid_hostname, uamqp_transport): if sys.platform.startswith('darwin'): pytest.skip("Skipping on OSX - it keeps reporting 'Unable to set external certificates' " "and blocking other tests") - client = EventHubProducerClient.from_connection_string(invalid_hostname) + client = EventHubProducerClient.from_connection_string(invalid_hostname, uamqp_transport=uamqp_transport) async with client: with pytest.raises(ConnectError): await client.create_batch(max_size_in_bytes=300) @@ -178,8 +181,259 @@ async def test_create_batch_with_invalid_hostname_async(invalid_hostname): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_create_batch_with_too_large_size_async(connection_str): - client = EventHubProducerClient.from_connection_string(connection_str) +async def test_create_batch_with_too_large_size_async(connection_str, uamqp_transport): + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) async with client: with pytest.raises(ValueError): await client.create_batch(max_size_in_bytes=5 * 1024 * 1024) + +@pytest.mark.liveTest +@pytest.mark.asyncio +async def test_invalid_proxy_server(connection_str, uamqp_transport): + if sys.platform.startswith('darwin') and uamqp_transport: + pytest.skip("Skipping on OSX - running forever and blocking other tests") + HTTP_PROXY = { + 'proxy_hostname': 'fakeproxy', # proxy hostname. + 'proxy_port': 3128, # proxy port. + } + + client = EventHubProducerClient.from_connection_string(connection_str, http_proxy=HTTP_PROXY, uamqp_transport=uamqp_transport) + async with client: + with pytest.raises(EventHubError): + batch = await client.create_batch() + +@pytest.mark.liveTest +@pytest.mark.asyncio +async def test_client_send_timeout(connstr_receivers, uamqp_transport): + connection_str, receivers = connstr_receivers + + async def on_success(events, pid): + pass + + async def on_error(events, pid, err): + pass + + producer = EventHubProducerClient.from_connection_string( + connection_str, uamqp_transport=uamqp_transport + ) + + async with producer: + with pytest.raises(OperationTimeoutError): + await producer.send_batch([EventData(b"Data")], timeout=-1) + + with pytest.raises(OperationTimeoutError): + await producer.send_event(EventData(b"Data"), timeout=-1) + + producer = EventHubProducerClient.from_connection_string( + connection_str, + buffered_mode=True, + on_success=on_success, + on_error=on_error, + uamqp_transport=uamqp_transport + ) + + async with producer: + with pytest.raises(OperationTimeoutError): + await producer.send_batch([EventData(b"Data")], timeout=-1) + + with pytest.raises(OperationTimeoutError): + await producer.send_event(EventData(b"Data"), timeout=-1) + +@pytest.mark.liveTest +@pytest.mark.asyncio +async def test_client_invalid_credential_async(live_eventhub, uamqp_transport): + + async def on_event(partition_context, event): + pass + + async def on_error(partition_context, error): + on_error.err = error + + env_credential = EnvironmentCredential() + # Skipping on OSX - it's raising a ConnectionLostError and blocking other tests + if not sys.platform.startswith('darwin'): + producer_client = EventHubProducerClient(fully_qualified_namespace="fakeeventhub.servicebus.windows.net", + eventhub_name=live_eventhub['event_hub'], + credential=env_credential, + user_agent='customized information', + retry_total=1, + retry_mode='exponential', + retry_backoff=0.02, + uamqp_transport=uamqp_transport) + consumer_client = EventHubConsumerClient(fully_qualified_namespace="fakeeventhub.servicebus.windows.net", + eventhub_name=live_eventhub['event_hub'], + credential=env_credential, + user_agent='customized information', + consumer_group='$Default', + retry_total=1, + retry_mode='exponential', + retry_backoff=0.02, + uamqp_transport=uamqp_transport) + async with producer_client: + with pytest.raises(ConnectError): + await producer_client.create_batch(partition_id='0') + + on_error.err = None + async with consumer_client: + task = asyncio.ensure_future(consumer_client.receive(on_event, + starting_position= "-1", on_error=on_error)) + await asyncio.sleep(15) + await task + assert isinstance(on_error.err, ConnectError) + + producer_client = EventHubProducerClient(fully_qualified_namespace=live_eventhub['hostname'], + eventhub_name='fakehub', + credential=env_credential, + uamqp_transport=uamqp_transport) + + consumer_client = EventHubConsumerClient(fully_qualified_namespace=live_eventhub['hostname'], + eventhub_name='fakehub', + credential=env_credential, + consumer_group='$Default', + retry_total=0, + uamqp_transport=uamqp_transport) + + async with producer_client: + with pytest.raises(ConnectError): + await producer_client.create_batch(partition_id='0') + + on_error.err = None + async with consumer_client: + task = asyncio.ensure_future(consumer_client.receive(on_event, + starting_position= "-1", on_error=on_error)) + await asyncio.sleep(15) + await task + assert isinstance(on_error.err, AuthenticationError) + + credential = EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']) + auth_uri = "sb://{}/{}".format(live_eventhub['hostname'], live_eventhub['event_hub']) + token = (await credential.get_token(auth_uri)).token + producer_client = EventHubProducerClient(fully_qualified_namespace=live_eventhub['hostname'], + eventhub_name=live_eventhub['event_hub'], + credential=EventHubSASTokenCredential(token[:-1], time.time() + 5), + uamqp_transport=uamqp_transport) + await asyncio.sleep(10) + # expired credential + async with producer_client: + with pytest.raises(AuthenticationError): + await producer_client.create_batch(partition_id='0') + + consumer_client = EventHubConsumerClient(fully_qualified_namespace=live_eventhub['hostname'], + eventhub_name=live_eventhub['event_hub'], + credential=EventHubSASTokenCredential(token, time.time() + 7), + consumer_group='$Default', + retry_total=0, + uamqp_transport=uamqp_transport) + on_error.err = None + async with consumer_client: + task = asyncio.ensure_future(consumer_client.receive(on_event, + starting_position= "-1", on_error=on_error)) + await asyncio.sleep(15) + await task + + # expired credential + assert isinstance(on_error.err, AuthenticationError) + + credential = EventHubSharedKeyCredential('fakekey', live_eventhub['access_key']) + producer_client = EventHubProducerClient(fully_qualified_namespace=live_eventhub['hostname'], + eventhub_name=live_eventhub['event_hub'], + credential=credential, + uamqp_transport=uamqp_transport) + + async with producer_client: + with pytest.raises(AuthenticationError): + await producer_client.create_batch(partition_id='0') + + consumer_client = EventHubConsumerClient(fully_qualified_namespace=live_eventhub['hostname'], + eventhub_name=live_eventhub['event_hub'], + credential=credential, + consumer_group='$Default', + retry_total=0, + uamqp_transport=uamqp_transport) + on_error.err = None + async with consumer_client: + task = asyncio.ensure_future(consumer_client.receive(on_event, + starting_position= "-1", on_error=on_error)) + await asyncio.sleep(15) + await task + + assert isinstance(on_error.err, AuthenticationError) + + credential = EventHubSharedKeyCredential(live_eventhub['key_name'], 'fakekey') + producer_client = EventHubProducerClient(fully_qualified_namespace=live_eventhub['hostname'], + eventhub_name=live_eventhub['event_hub'], + credential=credential, + uamqp_transport=uamqp_transport) + + async with producer_client: + with pytest.raises(AuthenticationError): + await producer_client.create_batch(partition_id='0') + + consumer_client = EventHubConsumerClient(fully_qualified_namespace=live_eventhub['hostname'], + eventhub_name=live_eventhub['event_hub'], + credential=credential, + consumer_group='$Default', + retry_total=0, + uamqp_transport=uamqp_transport) + on_error.err = None + async with consumer_client: + task = asyncio.ensure_future(consumer_client.receive(on_event, + starting_position= "-1", on_error=on_error)) + await asyncio.sleep(15) + await task + + assert isinstance(on_error.err, AuthenticationError) + + producer_client = EventHubProducerClient(fully_qualified_namespace=live_eventhub['hostname'], + eventhub_name=live_eventhub['event_hub'], + credential=env_credential, + connection_verify="cacert.pem", + uamqp_transport=uamqp_transport) + + # TODO: this seems like a bug from uamqp, should be ConnectError? + async with producer_client: + with pytest.raises(EventHubError): + await producer_client.create_batch(partition_id='0') + + consumer_client = EventHubConsumerClient(fully_qualified_namespace=live_eventhub['hostname'], + eventhub_name=live_eventhub['event_hub'], + consumer_group='$Default', + credential=env_credential, + retry_total=0, + connection_verify="cacert.pem", + uamqp_transport=uamqp_transport) + async with consumer_client: + task = asyncio.ensure_future(consumer_client.receive(on_event, + starting_position= "-1", on_error=on_error)) + await asyncio.sleep(15) + await task + + # TODO: this seems like a bug from uamqp, should be ConnectError? + assert isinstance(on_error.err, FileNotFoundError) + + # Skipping on OSX - it's raising a ConnectionLostError + if not sys.platform.startswith('darwin'): + producer_client = EventHubProducerClient(fully_qualified_namespace=live_eventhub['hostname'], + eventhub_name=live_eventhub['event_hub'], + credential=env_credential, + custom_endpoint_address="fakeaddr", + uamqp_transport=uamqp_transport) + + async with producer_client: + with pytest.raises(AuthenticationError): + await producer_client.create_batch(partition_id='0') + + consumer_client = EventHubConsumerClient(fully_qualified_namespace=live_eventhub['hostname'], + eventhub_name=live_eventhub['event_hub'], + consumer_group='$Default', + credential=env_credential, + retry_total=0, + custom_endpoint_address="fakeaddr", + uamqp_transport=uamqp_transport) + async with consumer_client: + task = asyncio.ensure_future(consumer_client.receive(on_event, + starting_position= "-1", on_error=on_error)) + await asyncio.sleep(15) + await task + + assert isinstance(on_error.err, AuthenticationError) diff --git a/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_properties_async.py b/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_properties_async.py index fe53764dd8dc..de618d8dab38 100644 --- a/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_properties_async.py +++ b/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_properties_async.py @@ -12,9 +12,10 @@ @pytest.mark.liveTest @pytest.mark.asyncio -async def test_get_properties(live_eventhub): +async def test_get_properties(live_eventhub, uamqp_transport): client = EventHubConsumerClient(live_eventhub['hostname'], live_eventhub['event_hub'], '$default', - EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']) + EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']), + uamqp_transport=uamqp_transport ) async with client: properties = await client.get_eventhub_properties() @@ -22,16 +23,18 @@ async def test_get_properties(live_eventhub): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_get_properties_with_auth_error_async(live_eventhub): +async def test_get_properties_with_auth_error_async(live_eventhub, uamqp_transport): client = EventHubConsumerClient(live_eventhub['hostname'], live_eventhub['event_hub'], '$default', - EventHubSharedKeyCredential(live_eventhub['key_name'], "AaBbCcDdEeFf=") + EventHubSharedKeyCredential(live_eventhub['key_name'], "AaBbCcDdEeFf="), + uamqp_transport=uamqp_transport ) async with client: with pytest.raises(AuthenticationError) as e: await client.get_eventhub_properties() client = EventHubConsumerClient(live_eventhub['hostname'], live_eventhub['event_hub'], '$default', - EventHubSharedKeyCredential("invalid", live_eventhub['access_key']) + EventHubSharedKeyCredential("invalid", live_eventhub['access_key']), + uamqp_transport=uamqp_transport ) async with client: with pytest.raises(AuthenticationError) as e: @@ -39,26 +42,29 @@ async def test_get_properties_with_auth_error_async(live_eventhub): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_get_properties_with_connect_error(live_eventhub): +async def test_get_properties_with_connect_error(live_eventhub, uamqp_transport): client = EventHubConsumerClient(live_eventhub['hostname'], "invalid", '$default', - EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']) + EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']), + uamqp_transport=uamqp_transport ) async with client: with pytest.raises(ConnectError) as e: await client.get_eventhub_properties() client = EventHubConsumerClient("invalid.servicebus.windows.net", live_eventhub['event_hub'], '$default', - EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']) + EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']), + uamqp_transport=uamqp_transport ) async with client: - with pytest.raises(EventHubError) as e: # This can be either ConnectError or ConnectionLostError + with pytest.raises(ConnectError) as e: await client.get_eventhub_properties() @pytest.mark.liveTest @pytest.mark.asyncio -async def test_get_partition_ids(live_eventhub): +async def test_get_partition_ids(live_eventhub, uamqp_transport): client = EventHubConsumerClient(live_eventhub['hostname'], live_eventhub['event_hub'], '$default', - EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']) + EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']), + uamqp_transport=uamqp_transport ) async with client: partition_ids = await client.get_partition_ids() @@ -67,9 +73,10 @@ async def test_get_partition_ids(live_eventhub): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_get_partition_properties(live_eventhub): +async def test_get_partition_properties(live_eventhub, uamqp_transport): client = EventHubProducerClient(live_eventhub['hostname'], live_eventhub['event_hub'], - EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']) + EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']), + uamqp_transport=uamqp_transport ) async with client: properties = await client.get_partition_properties('0') diff --git a/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_receive_async.py b/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_receive_async.py index 6fb153052023..18279b81cc2d 100644 --- a/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_receive_async.py +++ b/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_receive_async.py @@ -15,7 +15,7 @@ @pytest.mark.liveTest @pytest.mark.asyncio -async def test_receive_end_of_stream_async(connstr_senders): +async def test_receive_end_of_stream_async(connstr_senders, uamqp_transport): async def on_event(partition_context, event): if partition_context.partition_id == "0": assert event.body_as_str() == "Receiving only a single event" @@ -31,7 +31,7 @@ async def on_event(partition_context, event): on_event.called = False connection_str, senders = connstr_senders # test async producer client - producer_client = EventHubProducerClient.from_connection_string(connection_str) + producer_client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) partitions = await producer_client.get_partition_ids() senders = [] for p in partitions: @@ -61,7 +61,7 @@ async def on_event(partition_context, event): ("enqueued_time", False, "Exclusive")]) @pytest.mark.liveTest @pytest.mark.asyncio -async def test_receive_with_event_position_async(connstr_senders, position, inclusive, expected_result): +async def test_receive_with_event_position_async(connstr_senders, position, inclusive, expected_result, uamqp_transport): async def on_event(partition_context, event): assert partition_context.last_enqueued_event_properties.get('sequence_number') == event.sequence_number assert partition_context.last_enqueued_event_properties.get('offset') == event.offset @@ -79,7 +79,7 @@ async def on_event(partition_context, event): on_event.event_position = None connection_str, senders = connstr_senders senders[0].send(EventData(b"Inclusive")) - client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default') + client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default', uamqp_transport=uamqp_transport) async with client: task = asyncio.ensure_future(client.receive(on_event, starting_position="-1", @@ -89,7 +89,7 @@ async def on_event(partition_context, event): assert on_event.event_position is not None await task senders[0].send(EventData(expected_result)) - client2 = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default') + client2 = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default', uamqp_transport=uamqp_transport) async with client2: task = asyncio.ensure_future( client2.receive(on_event, @@ -102,7 +102,7 @@ async def on_event(partition_context, event): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_receive_owner_level_async(connstr_senders): +async def test_receive_owner_level_async(connstr_senders, uamqp_transport): app_prop = {"raw_prop": "raw_value"} async def on_event(partition_context, event): @@ -112,8 +112,8 @@ async def on_error(partition_context, error): on_error.error = None connection_str, senders = connstr_senders - client1 = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default') - client2 = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default') + client1 = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default', uamqp_transport=uamqp_transport) + client2 = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default', uamqp_transport=uamqp_transport) async with client1, client2: task1 = asyncio.ensure_future(client1.receive(on_event, partition_id="0", starting_position="-1", @@ -136,7 +136,7 @@ async def on_error(partition_context, error): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_receive_over_websocket_async(connstr_senders): +async def test_receive_over_websocket_async(connstr_senders, uamqp_transport): app_prop = {"raw_prop": "raw_value"} content_type = "text/plain" message_id_base = "mess_id_sample_" @@ -149,7 +149,8 @@ async def on_event(partition_context, event): on_event.app_prop = None connection_str, senders = connstr_senders client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default', - transport_type=TransportType.AmqpOverWebsocket) + transport_type=TransportType.AmqpOverWebsocket, + uamqp_transport=uamqp_transport) event_list = [] for i in range(5): diff --git a/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_reconnect_async.py b/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_reconnect_async.py index 1bea5b131507..aeaebe4cf938 100644 --- a/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_reconnect_async.py +++ b/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_reconnect_async.py @@ -9,20 +9,31 @@ import pytest import time +from azure.eventhub._pyamqp.aio._authentication_async import SASTokenAuthAsync +from azure.eventhub._pyamqp.aio import ReceiveClientAsync +from azure.eventhub._pyamqp import error, constants +from azure.eventhub._utils import transform_outbound_single_message +try: + import uamqp + from uamqp import compat + from azure.eventhub._transport._uamqp_transport import UamqpTransport +except (ModuleNotFoundError, ImportError): + uamqp = None + UamqpTransport = None + +from azure.eventhub._transport._pyamqp_transport import PyamqpTransport from azure.eventhub import EventData from azure.eventhub.aio import EventHubProducerClient, EventHubConsumerClient, EventHubSharedKeyCredential from azure.eventhub.exceptions import OperationTimeoutError -import uamqp -from uamqp import authentication, errors, c_uamqp, compat - @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_with_long_interval_async(live_eventhub, sleep): +async def test_send_with_long_interval_async(live_eventhub, sleep, uamqp_transport, timeout_factor): test_partition = "0" sender = EventHubProducerClient(live_eventhub['hostname'], live_eventhub['event_hub'], - EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key'])) + EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']), + uamqp_transport=uamqp_transport) async with sender: batch = await sender.create_batch(partition_id=test_partition) batch.add(EventData(b"A single event")) @@ -31,90 +42,139 @@ async def test_send_with_long_interval_async(live_eventhub, sleep): if sleep: await asyncio.sleep(250) # EH server side idle timeout is 240 second else: - await sender._producers[test_partition]._handler._connection._conn.destroy() + if uamqp_transport: + await sender._producers[test_partition]._handler._connection._conn.destroy() + else: + await sender._producers[test_partition]._handler._connection.close() batch = await sender.create_batch(partition_id=test_partition) batch.add(EventData(b"A single event")) await sender.send_batch(batch) received = [] uri = "sb://{}/{}".format(live_eventhub['hostname'], live_eventhub['event_hub']) - sas_auth = authentication.SASTokenAuth.from_shared_access_key( - uri, live_eventhub['key_name'], live_eventhub['access_key']) - source = "amqps://{}/{}/ConsumerGroups/{}/Partitions/{}".format( live_eventhub['hostname'], live_eventhub['event_hub'], live_eventhub['consumer_group'], test_partition) - receiver = uamqp.ReceiveClient(source, auth=sas_auth, debug=False, timeout=10000, prefetch=10) + if uamqp_transport: + sas_auth = uamqp.authentication.SASTokenAsync.from_shared_access_key( + uri, live_eventhub['key_name'], live_eventhub['access_key']) + receiver = uamqp.async_ops.client_async.ReceiveClientAsync(source, auth=sas_auth, debug=False, timeout=5000, prefetch=500) + else: + sas_auth = SASTokenAuthAsync( + uri, uri, live_eventhub['key_name'], live_eventhub['access_key'] + ) + receiver = ReceiveClientAsync(live_eventhub['hostname'], source, auth=sas_auth, debug=False, link_credit=500) try: - receiver.open() + await receiver.open_async() # receive_message_batch() returns immediately once it receives any messages before the max_batch_size # and timeout reach. Could be 1, 2, or any number between 1 and max_batch_size. # So call it twice to ensure the two events are received. - received.extend([EventData._from_message(x) for x in receiver.receive_message_batch(max_batch_size=1, timeout=5000)]) - received.extend([EventData._from_message(x) for x in receiver.receive_message_batch(max_batch_size=1, timeout=5000)]) + received.extend([EventData._from_message(x) for x in (await receiver.receive_message_batch_async(max_batch_size=1, timeout=5 * timeout_factor))]) + received.extend([EventData._from_message(x) for x in (await receiver.receive_message_batch_async(max_batch_size=1, timeout=5 * timeout_factor))]) finally: - receiver.close() - + await receiver.close_async() assert len(received) == 2 assert list(received[0].body)[0] == b"A single event" @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_connection_idle_timeout_and_reconnect_async(connstr_receivers): +async def test_send_connection_idle_timeout_and_reconnect_async(connstr_receivers, uamqp_transport, timeout_factor): connection_str, receivers = connstr_receivers - client = EventHubProducerClient.from_connection_string(conn_str=connection_str, idle_timeout=10) + if uamqp_transport: + amqp_transport = UamqpTransport + retry_total = 3 + timeout_exc = compat.TimeoutException + else: + amqp_transport = PyamqpTransport + retry_total = 0 + timeout_exc = TimeoutError + + client = EventHubProducerClient.from_connection_string(conn_str=connection_str, idle_timeout=10, retry_total=retry_total, uamqp_transport=uamqp_transport) async with client: ed = EventData('data') sender = client._create_producer(partition_id='0') async with sender: await sender._open_with_retry() - time.sleep(11) - sender._unsent_events = [ed.message] - ed.message.on_send_complete = sender._on_outcome - with pytest.raises((uamqp.errors.ConnectionClose, - uamqp.errors.MessageHandlerError, OperationTimeoutError)): - # Mac may raise OperationTimeoutError or MessageHandlerError - await sender._send_event_data() - await sender._send_event_data_with_retry() + await asyncio.sleep(11) + ed = transform_outbound_single_message(ed, EventData, amqp_transport.to_outgoing_amqp_message) + sender._unsent_events = [ed._message] + if uamqp_transport: + sender._unsent_events[0].on_send_complete = sender._on_outcome + with pytest.raises((uamqp.errors.ConnectionClose, + uamqp.errors.MessageHandlerError, OperationTimeoutError)): + await sender._send_event_data() + else: + with pytest.raises(error.AMQPConnectionError): + await sender._send_event_data() + if uamqp_transport: + await sender._send_event_data_with_retry() + + # with retry, should work + if not uamqp_transport: + client = EventHubProducerClient.from_connection_string(conn_str=connection_str, idle_timeout=10) + async with client: + ed = EventData('data') + sender = client._create_producer(partition_id='0') + async with sender: + await sender._open_with_retry() + await asyncio.sleep(11) + ed = transform_outbound_single_message(ed, EventData, amqp_transport.to_outgoing_amqp_message) + sender._unsent_events = [ed._message] + await sender._send_event_data() + retry = 0 while retry < 3: try: - messages = receivers[0].receive_message_batch(max_batch_size=10, timeout=10000) + messages = receivers[0].receive_message_batch(max_batch_size=10, timeout=10 * timeout_factor) if messages: received_ed1 = EventData._from_message(messages[0]) assert received_ed1.body_as_str() == 'data' break - except compat.TimeoutException: + except timeout_exc: retry += 1 @pytest.mark.liveTest @pytest.mark.asyncio -async def test_receive_connection_idle_timeout_and_reconnect_async(connstr_senders): +async def test_receive_connection_idle_timeout_and_reconnect_async(connstr_senders, uamqp_transport): connection_str, senders = connstr_senders client = EventHubConsumerClient.from_connection_string( conn_str=connection_str, consumer_group='$default', - idle_timeout=10 + idle_timeout=10, + uamqp_transport=uamqp_transport ) + async def on_event_received(event): on_event_received.event = event async with client: consumer = client._create_consumer("$default", "0", "-1", on_event_received) async with consumer: - await consumer._open_with_retry() - - time.sleep(11) + if uamqp_transport: + await consumer._open_with_retry() + else: + await consumer._open() + await asyncio.sleep(11) ed = EventData("Event") senders[0].send(ed) - await consumer._handler.do_work_async() - assert consumer._handler._connection._state == c_uamqp.ConnectionState.DISCARDING - await consumer.receive(batch=False, max_batch_size=1, max_wait_time=10) + if uamqp_transport: + await consumer._handler.do_work_async() + assert consumer._handler._connection._state == uamqp.c_uamqp.ConnectionState.DISCARDING + await consumer.receive(batch=False, max_batch_size=1, max_wait_time=10) + else: + with pytest.raises(error.AMQPConnectionError): + await consumer._handler.do_work_async() + assert consumer._handler._connection.state == constants.ConnectionState.END + try: + await asyncio.wait_for(consumer.receive(), timeout=10) + except asyncio.TimeoutError: + pass + assert on_event_received.event.body_as_str() == "Event" diff --git a/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_send_async.py b/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_send_async.py index 141d2b4861d6..743b2961b93d 100644 --- a/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_send_async.py +++ b/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_send_async.py @@ -10,7 +10,6 @@ import pytest import time import json -import uamqp from azure.eventhub import EventData, TransportType, EventDataBatch from azure.eventhub.aio import EventHubProducerClient, EventHubConsumerClient @@ -21,6 +20,18 @@ AmqpAnnotatedMessage, AmqpMessageProperties, ) +try: + import uamqp + from uamqp.constants import TransportType as uamqp_TransportType, MessageState + from uamqp.message import MessageProperties +except (ModuleNotFoundError, ImportError): + uamqp = None + uamqp_TransportType = TransportType + MessageProperties = None +from azure.eventhub._pyamqp.message import Properties +from azure.eventhub._pyamqp.authentication import SASTokenAuth +from azure.eventhub._pyamqp.client import ReceiveClient +from azure.eventhub._pyamqp.error import AMQPConnectionError @pytest.mark.liveTest @@ -145,9 +156,9 @@ async def on_event(partition_context, event): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_with_partition_key_async(connstr_receivers, live_eventhub): +async def test_send_with_partition_key_async(connstr_receivers, live_eventhub, uamqp_transport, timeout_factor): connection_str, receivers = connstr_receivers - client = EventHubProducerClient.from_connection_string(connection_str) + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) async with client: data_val = 0 for partition in [b"a", b"b", b"c", b"d", b"e", b"f"]: @@ -175,7 +186,7 @@ async def test_send_with_partition_key_async(connstr_receivers, live_eventhub): for index, partition in enumerate(receivers): retry_total = 0 while retry_total < 3: - timeout = 5000 + retry_total * 1000 + timeout = (5 * retry_total) * timeout_factor try: received = partition.receive_message_batch(timeout=timeout) for message in received: @@ -192,19 +203,24 @@ async def test_send_with_partition_key_async(connstr_receivers, live_eventhub): if received: break retry_total += 1 - except uamqp.errors.ConnectionClose: + except AMQPConnectionError: for r in reconnect_receivers: r.close() uri = "sb://{}/{}".format(live_eventhub['hostname'], live_eventhub['event_hub']) - sas_auth = uamqp.authentication.SASTokenAuth.from_shared_access_key( - uri, live_eventhub['key_name'], live_eventhub['access_key']) - source = "amqps://{}/{}/ConsumerGroups/{}/Partitions/{}".format( live_eventhub['hostname'], live_eventhub['event_hub'], live_eventhub['consumer_group'], index) - partition = uamqp.ReceiveClient(source, auth=sas_auth, debug=True, timeout=0, prefetch=500) + if uamqp_transport: + sas_auth = uamqp.authentication.SASTokenAuth.from_shared_access_key( + uri, live_eventhub['key_name'], live_eventhub['access_key']) + partition = uamqp.ReceiveClient(source, auth=sas_auth, debug=False, timeout=0, prefetch=500) + else: + sas_auth = SASTokenAuth( + uri, uri, live_eventhub['key_name'], live_eventhub['access_key']) + partition = ReceiveClient(live_eventhub['hostname'], source, auth=sas_auth, network_trace=False, timeout=0, link_credit=500) + partition.open() reconnect_receivers.append(partition) retry_total += 1 if retry_total == 3: @@ -220,9 +236,9 @@ async def test_send_with_partition_key_async(connstr_receivers, live_eventhub): @pytest.mark.parametrize("payload", [b"", b"A single event"]) @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_and_receive_small_body_async(connstr_receivers, payload): +async def test_send_and_receive_small_body_async(connstr_receivers, payload, uamqp_transport, timeout_factor): connection_str, receivers = connstr_receivers - client = EventHubProducerClient.from_connection_string(connection_str) + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) async with client: batch = await client.create_batch() batch.add(EventData(payload)) @@ -230,7 +246,7 @@ async def test_send_and_receive_small_body_async(connstr_receivers, payload): await client.send_event(EventData(payload)) received = [] for r in receivers: - received.extend([EventData._from_message(x) for x in r.receive_message_batch(timeout=5000)]) + received.extend([EventData._from_message(x) for x in r.receive_message_batch(timeout=5 * timeout_factor)]) assert len(received) == 2 assert list(received[0].body)[0] == payload @@ -239,9 +255,9 @@ async def test_send_and_receive_small_body_async(connstr_receivers, payload): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_partition_async(connstr_receivers): +async def test_send_partition_async(connstr_receivers, uamqp_transport, timeout_factor): connection_str, receivers = connstr_receivers - client = EventHubProducerClient.from_connection_string(connection_str) + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) async with client: batch = await client.create_batch() @@ -249,14 +265,15 @@ async def test_send_partition_async(connstr_receivers): await client.send_batch(batch) await client.send_event(EventData(b"Data")) + async with client: batch = await client.create_batch(partition_id="1") batch.add(EventData(b"Data")) await client.send_batch(batch) await client.send_event(EventData(b"Data"), partition_id="1") - partition_0 = receivers[0].receive_message_batch(timeout=5000) - partition_1 = receivers[1].receive_message_batch(timeout=5000) + partition_0 = receivers[0].receive_message_batch(timeout=10 * timeout_factor) + partition_1 = receivers[1].receive_message_batch(timeout=10 * timeout_factor) assert len(partition_1) >= 2 assert len(partition_0) + len(partition_1) == 4 @@ -265,25 +282,28 @@ async def test_send_partition_async(connstr_receivers): batch.add(EventData(b"Data")) await client.send_batch(batch) await client.send_event(EventData(b"Data")) - async with client: batch = await client.create_batch(partition_id="0") batch.add(EventData(b"Data")) await client.send_batch(batch) await client.send_event(EventData(b"Data"), partition_id="0") + async with client: + batch = EventDataBatch(partition_id="0") + batch.add(EventData(b"Data")) + await client.send_batch(batch) time.sleep(5) - partition_0 = receivers[0].receive_message_batch(timeout=5000) - partition_1 = receivers[1].receive_message_batch(timeout=5000) - assert len(partition_0) >= 2 - assert len(partition_0) + len(partition_1) == 4 + partition_0 = receivers[0].receive_message_batch(timeout=10 * timeout_factor) + partition_1 = receivers[1].receive_message_batch(timeout=10 * timeout_factor) + assert len(partition_0) >= 3 + assert len(partition_0) + len(partition_1) == 5 @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_non_ascii_async(connstr_receivers): +async def test_send_non_ascii_async(connstr_receivers, uamqp_transport, timeout_factor): connection_str, receivers = connstr_receivers - client = EventHubProducerClient.from_connection_string(connection_str) + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) async with client: batch = await client.create_batch(partition_id="0") batch.add(EventData(u"é,è,à,ù,â,ê,î,ô,û")) @@ -295,8 +315,8 @@ async def test_send_non_ascii_async(connstr_receivers): # receive_message_batch() returns immediately once it receives any messages before the max_batch_size # and timeout reach. Could be 1, 2, or any number between 1 and max_batch_size. # So call it twice to ensure the two events are received. - partition_0 = [EventData._from_message(x) for x in receivers[0].receive_message_batch(timeout=5000)] + \ - [EventData._from_message(x) for x in receivers[0].receive_message_batch(timeout=5000)] + partition_0 = [EventData._from_message(x) for x in receivers[0].receive_message_batch(timeout=5 * timeout_factor)] + \ + [EventData._from_message(x) for x in receivers[0].receive_message_batch(timeout=5 * timeout_factor)] assert len(partition_0) == 4 assert partition_0[0].body_as_str() == u"é,è,à,ù,â,ê,î,ô,û" @@ -307,12 +327,16 @@ async def test_send_non_ascii_async(connstr_receivers): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_multiple_partition_with_app_prop_async(connstr_receivers): +async def test_send_multiple_partition_with_app_prop_async(connstr_receivers, uamqp_transport, timeout_factor): connection_str, receivers = connstr_receivers app_prop_key = "raw_prop" app_prop_value = "raw_value" app_prop = {app_prop_key: app_prop_value} - client = EventHubProducerClient.from_connection_string(connection_str) + client = EventHubProducerClient.from_connection_string( + connection_str, + uamqp_transport=uamqp_transport, + transport_type=TransportType.Amqp + ) async with client: ed0 = EventData(b"Message 0") ed0.properties = app_prop @@ -327,23 +351,23 @@ async def test_send_multiple_partition_with_app_prop_async(connstr_receivers): batch.add(ed1) await client.send_batch(batch) await client.send_event(ed1, partition_id="1") - - partition_0 = [EventData._from_message(x) for x in receivers[0].receive_message_batch(timeout=5000)] + partition_0 = [EventData._from_message(x) for x in receivers[0].receive_message_batch(timeout=5 * timeout_factor)] assert len(partition_0) == 2 assert partition_0[0].properties[b"raw_prop"] == b"raw_value" assert partition_0[1].properties[b"raw_prop"] == b"raw_value" - partition_1 = [EventData._from_message(x) for x in receivers[1].receive_message_batch(timeout=5000)] - assert len(partition_1) == 2 + partition_1 = [EventData._from_message(x) for x in receivers[1].receive_message_batch(timeout=5 * timeout_factor)] + assert len(partition_0) == 2 assert partition_1[0].properties[b"raw_prop"] == b"raw_value" assert partition_0[1].properties[b"raw_prop"] == b"raw_value" @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_over_websocket_async(connstr_receivers): +async def test_send_over_websocket_async(connstr_receivers, uamqp_transport, timeout_factor): connection_str, receivers = connstr_receivers client = EventHubProducerClient.from_connection_string(connection_str, - transport_type=uamqp.constants.TransportType.AmqpOverWebsocket) + transport_type=uamqp_TransportType.AmqpOverWebsocket, + uamqp_transport=uamqp_transport) async with client: batch = await client.create_batch(partition_id="0") @@ -353,19 +377,20 @@ async def test_send_over_websocket_async(connstr_receivers): time.sleep(1) received = [] - received.extend(receivers[0].receive_message_batch(max_batch_size=5, timeout=10000)) + received.extend(receivers[0].receive_message_batch(max_batch_size=5, timeout=10 * timeout_factor)) assert len(received) == 2 @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_with_create_event_batch_async(connstr_receivers): +async def test_send_with_create_event_batch_async(connstr_receivers, uamqp_transport, timeout_factor): connection_str, receivers = connstr_receivers app_prop_key = "raw_prop" app_prop_value = "raw_value" app_prop = {app_prop_key: app_prop_value} client = EventHubProducerClient.from_connection_string(connection_str, - transport_type=TransportType.AmqpOverWebsocket) + transport_type=TransportType.AmqpOverWebsocket, + uamqp_transport=uamqp_transport) async with client: event_data_batch = await client.create_batch(max_size_in_bytes=100000) while True: @@ -378,22 +403,26 @@ async def test_send_with_create_event_batch_async(connstr_receivers): await client.send_batch(event_data_batch) received = [] for r in receivers: - received.extend(r.receive_message_batch(timeout=10000)) + received.extend(r.receive_message_batch(timeout=10 * timeout_factor)) assert len(received) >= 1 assert EventData._from_message(received[0]).properties[b"raw_prop"] == b"raw_value" @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_list_async(connstr_receivers): +async def test_send_list_async(connstr_receivers, uamqp_transport, timeout_factor): connection_str, receivers = connstr_receivers - client = EventHubProducerClient.from_connection_string(connection_str) + client = EventHubProducerClient.from_connection_string( + connection_str, + uamqp_transport=uamqp_transport, + transport_type=uamqp_TransportType.Amqp + ) payload = "A1" async with client: await client.send_batch([EventData(payload)]) received = [] for r in receivers: - received.extend([EventData._from_message(x) for x in r.receive_message_batch(timeout=10000)]) + received.extend([EventData._from_message(x) for x in r.receive_message_batch(timeout=10 * timeout_factor)]) assert len(received) == 1 assert received[0].body_as_str() == payload @@ -401,13 +430,13 @@ async def test_send_list_async(connstr_receivers): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_list_partition_async(connstr_receivers): +async def test_send_list_partition_async(connstr_receivers, uamqp_transport, timeout_factor): connection_str, receivers = connstr_receivers - client = EventHubProducerClient.from_connection_string(connection_str) + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) payload = "A1" async with client: await client.send_batch([EventData(payload)], partition_id="0") - message = receivers[0].receive_message_batch(timeout=10000)[0] + message = receivers[0].receive_message_batch(timeout=10 * timeout_factor)[0] received = EventData._from_message(message) assert received.body_as_str() == payload @@ -418,8 +447,8 @@ async def test_send_list_partition_async(connstr_receivers): ]) @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_list_wrong_data_async(connection_str, to_send, exception_type): - client = EventHubProducerClient.from_connection_string(connection_str) +async def test_send_list_wrong_data_async(connection_str, to_send, exception_type, uamqp_transport): + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) async with client: with pytest.raises(exception_type): await client.send_batch(to_send) @@ -428,9 +457,9 @@ async def test_send_list_wrong_data_async(connection_str, to_send, exception_typ @pytest.mark.parametrize("partition_id, partition_key", [("0", None), (None, "pk")]) @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_batch_pid_pk_async(invalid_hostname, partition_id, partition_key): +async def test_send_batch_pid_pk_async(invalid_hostname, partition_id, partition_key, uamqp_transport): # Use invalid_hostname because this is not a live test. - client = EventHubProducerClient.from_connection_string(invalid_hostname) + client = EventHubProducerClient.from_connection_string(invalid_hostname, uamqp_transport=uamqp_transport) batch = EventDataBatch(partition_id=partition_id, partition_key=partition_key) async with client: with pytest.raises(TypeError): @@ -439,7 +468,7 @@ async def test_send_batch_pid_pk_async(invalid_hostname, partition_id, partition @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_with_callback_async(connstr_receivers): +async def test_send_with_callback_async(connstr_receivers, uamqp_transport): async def on_error(events, pid, err): on_error.err = err @@ -450,7 +479,7 @@ async def on_success(events, pid): sent_events = [] on_error.err = None connection_str, receivers = connstr_receivers - client = EventHubProducerClient.from_connection_string(connection_str, on_success=on_success, on_error=on_error) + client = EventHubProducerClient.from_connection_string(connection_str, on_success=on_success, on_error=on_error, uamqp_transport=uamqp_transport) async with client: batch = await client.create_batch() diff --git a/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_auth.py b/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_auth.py index 3fd512bbf141..a8bf32f23a1a 100644 --- a/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_auth.py +++ b/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_auth.py @@ -111,7 +111,9 @@ def test_client_azure_sas_credential(live_eventhub, uamqp_transport): token = credential.get_token(auth_uri).token.decode() producer_client = EventHubProducerClient(fully_qualified_namespace=hostname, eventhub_name=live_eventhub['event_hub'], - credential=AzureSasCredential(token)) + credential=AzureSasCredential(token), + auth_timeout=3, + uamqp_transport=uamqp_transport) with producer_client: batch = producer_client.create_batch(partition_id='0') @@ -127,6 +129,7 @@ def test_client_azure_named_key_credential(live_eventhub, uamqp_transport): consumer_group='$default', credential=credential, user_agent='customized information', + auth_timeout=3, uamqp_transport=uamqp_transport) assert consumer_client.get_eventhub_properties() is not None diff --git a/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_buffered_producer.py b/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_buffered_producer.py index 0da22a65dd77..d858e23b45b8 100644 --- a/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_buffered_producer.py +++ b/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_buffered_producer.py @@ -46,6 +46,7 @@ def on_success(events, pid): def on_error(events, error, pid): pass + with pytest.raises(TypeError): EventHubProducerClient.from_connection_string(connection_str, buffered_mode=True, uamqp_transport=uamqp_transport) with pytest.raises(TypeError): @@ -71,6 +72,30 @@ def on_error(events, error, pid): uamqp_transport=uamqp_transport ) + def on_success_missing_params(events): + on_success_missing_params.events = events + + def on_error_missing_params(events, pid): + on_error_missing_params.events = events + + producer = EventHubProducerClient.from_connection_string( + connection_str, + buffered_mode=True, + buffer_concurrency=2, + on_success=on_success_missing_params, + on_error=on_error_missing_params, + uamqp_transport=uamqp_transport, + ) + + on_success_missing_params.events = None + on_error_missing_params.events = None + + # successfully send, but don't enter invalid callback + with producer: + producer.send_event(EventData('Single data')) + + assert not on_success_missing_params.events + assert not on_error_missing_params.events @pytest.mark.liveTest @pytest.mark.parametrize( @@ -93,6 +118,7 @@ def on_event(partition_context, event): receive_thread.daemon = True receive_thread.start() + time.sleep(10) sent_events = defaultdict(list) def on_success(events, pid): @@ -105,6 +131,7 @@ def on_error(events, pid, err): on_error.err = None # ensure no error on_success.batching = False # ensure batching happened + producer = EventHubProducerClient.from_connection_string( connection_str, buffered_mode=True, @@ -183,9 +210,9 @@ def on_error(events, pid, err): @pytest.mark.parametrize( "flush_after_sending, close_after_sending", [ + (False, False), (True, False), - (False, True), - (False, False) + (False, True) ] ) def test_basic_send_batch_events_round_robin(connection_str, flush_after_sending, close_after_sending, uamqp_transport): @@ -199,6 +226,7 @@ def on_event(partition_context, event): receive_thread.daemon = True receive_thread.start() + time.sleep(10) sent_events = defaultdict(list) def on_success(events, pid): @@ -270,7 +298,7 @@ def on_error(events, pid, err): # ensure all events are sent assert sum([len(sent_events[pid]) for pid in partitions]) == total_events_cnt - time.sleep(10) + time.sleep(20) assert len(sent_events) == len(received_events) == partitions_cnt # ensure all events are received in the correct partition @@ -307,6 +335,7 @@ def on_event(partition_context, event): receive_thread.daemon = True receive_thread.start() + time.sleep(5) sent_events = defaultdict(list) def on_success(events, pid): @@ -397,6 +426,7 @@ def on_event(partition_context, event): receive_thread.daemon = True receive_thread.start() + time.sleep(5) sent_events = defaultdict(list) def on_success(events, pid): @@ -475,6 +505,7 @@ def on_event(partition_context, event): receive_thread.daemon = True receive_thread.start() + time.sleep(5) sent_events = defaultdict(list) def on_success(events, pid): @@ -507,16 +538,17 @@ def on_error(events, pid, err): @pytest.mark.skip('not testing correctly + flaky, fix during MQ') @pytest.mark.liveTest -def test_long_wait_small_buffer(connection_str): +def test_long_wait_small_buffer(connection_str, uamqp_transport): received_events = defaultdict(list) def on_event(partition_context, event): received_events[partition_context.partition_id].append(event) - consumer = EventHubConsumerClient.from_connection_string(connection_str, consumer_group="$default") + consumer = EventHubConsumerClient.from_connection_string(connection_str, consumer_group="$default", uamqp_transport=uamqp_transport) receive_thread = Thread(target=consumer.receive, args=(on_event,)) receive_thread.daemon = True receive_thread.start() + time.sleep(10) sent_events = defaultdict(list) @@ -537,7 +569,8 @@ def on_error(events, pid, err): retry_mode='fixed', retry_backoff_factor=0.01, max_wait_time=10, - max_buffer_length=100 + max_buffer_length=100, + uamqp_transport=uamqp_transport ) with producer: diff --git a/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_negative.py b/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_negative.py index ba58cf2a00c7..53e13cacd3ac 100644 --- a/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_negative.py +++ b/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_negative.py @@ -9,25 +9,26 @@ import sys import threading +from azure.identity import EnvironmentCredential from azure.eventhub import ( EventData, EventDataBatch) from azure.eventhub.exceptions import ( ConnectError, AuthenticationError, - EventDataSendError + EventHubError, + OperationTimeoutError ) -from azure.eventhub import EventHubConsumerClient -from azure.eventhub import EventHubProducerClient -try: - from azure.eventhub._transport._uamqp_transport import UamqpTransport -except (ImportError, ModuleNotFoundError): - UamqpTransport = None +from azure.eventhub import ( + EventHubProducerClient, + EventHubConsumerClient, + EventHubSharedKeyCredential +) +from azure.eventhub._client_base import EventHubSASTokenCredential @pytest.mark.liveTest def test_send_batch_with_invalid_hostname(invalid_hostname, uamqp_transport): - amqp_transport = UamqpTransport if uamqp_transport else None if sys.platform.startswith('darwin'): pytest.skip("Skipping on OSX - it keeps reporting 'Unable to set external certificates' " "and blocking other tests") @@ -79,7 +80,6 @@ def on_event(partition_context, event): @pytest.mark.liveTest def test_send_batch_with_invalid_key(invalid_key, uamqp_transport): client = EventHubProducerClient.from_connection_string(invalid_key, uamqp_transport=uamqp_transport) - amqp_transport = UamqpTransport if uamqp_transport else None try: with pytest.raises(ConnectError): batch = EventDataBatch() @@ -147,3 +147,264 @@ def test_create_batch_with_too_large_size_sync(connection_str, uamqp_transport): with client: with pytest.raises(ValueError): client.create_batch(max_size_in_bytes=5 * 1024 * 1024) + +@pytest.mark.liveTest +def test_invalid_proxy_server(connection_str, uamqp_transport): + if sys.platform.startswith('darwin') and uamqp_transport: + pytest.skip("Skipping on OSX - running forever and blocking other tests") + HTTP_PROXY = { + 'proxy_hostname': 'fakeproxy', # proxy hostname. + 'proxy_port': 3128, # proxy port. + } + + client = EventHubProducerClient.from_connection_string(connection_str, http_proxy=HTTP_PROXY, uamqp_transport=uamqp_transport) + with client: + with pytest.raises(EventHubError): + batch = client.create_batch() + +@pytest.mark.liveTest +@pytest.mark.asyncio +def test_client_send_timeout(connstr_receivers, uamqp_transport): + connection_str, receivers = connstr_receivers + + def on_success(events, pid): + pass + + def on_error(events, pid, err): + pass + + producer = EventHubProducerClient.from_connection_string( + connection_str, uamqp_transport=uamqp_transport + ) + + with producer: + with pytest.raises(OperationTimeoutError): + producer.send_batch([EventData(b"Data")], timeout=-1) + + with pytest.raises(OperationTimeoutError): + producer.send_event(EventData(b"Data"), timeout=-1) + + producer = EventHubProducerClient.from_connection_string( + connection_str, + buffered_mode=True, + on_success=on_success, + on_error=on_error, + uamqp_transport=uamqp_transport + ) + + with producer: + with pytest.raises(OperationTimeoutError): + producer.send_batch([EventData(b"Data")], timeout=-1) + + with pytest.raises(OperationTimeoutError): + producer.send_event(EventData(b"Data"), timeout=-1) + + +@pytest.mark.liveTest +@pytest.mark.asyncio +def test_client_invalid_credential(live_eventhub, uamqp_transport): + + def on_event(partition_context, event): + pass + + def on_error(partition_context, error): + on_error.err = error + + env_credential = EnvironmentCredential() + # Skipping on OSX - it's raising a ConnectionLostError + if not sys.platform.startswith('darwin'): + producer_client = EventHubProducerClient(fully_qualified_namespace="fakeeventhub.servicebus.windows.net", + eventhub_name=live_eventhub['event_hub'], + credential=env_credential, + user_agent='customized information', + retry_total=1, + retry_mode='exponential', + retry_backoff=0.02, + uamqp_transport=uamqp_transport) + consumer_client = EventHubConsumerClient(fully_qualified_namespace="fakeeventhub.servicebus.windows.net", + eventhub_name=live_eventhub['event_hub'], + credential=env_credential, + user_agent='customized information', + consumer_group='$Default', + retry_total=1, + retry_mode='exponential', + retry_backoff=0.02, + uamqp_transport=uamqp_transport) + with producer_client: + with pytest.raises(ConnectError): + producer_client.create_batch(partition_id='0') + + on_error.err = None + with consumer_client: + thread = threading.Thread(target=consumer_client.receive, args=(on_event,), + kwargs={"starting_position": "-1", "on_error": on_error}) + thread.daemon = True + thread.start() + time.sleep(15) + thread.join() + assert isinstance(on_error.err, ConnectError) + + producer_client = EventHubProducerClient(fully_qualified_namespace=live_eventhub['hostname'], + eventhub_name='fakehub', + credential=env_credential, + uamqp_transport=uamqp_transport) + + consumer_client = EventHubConsumerClient(fully_qualified_namespace=live_eventhub['hostname'], + eventhub_name='fakehub', + credential=env_credential, + consumer_group='$Default', + retry_total=0, + uamqp_transport=uamqp_transport) + + with producer_client: + with pytest.raises(ConnectError): + producer_client.create_batch(partition_id='0') + + on_error.err = None + with consumer_client: + thread = threading.Thread(target=consumer_client.receive, args=(on_event,), + kwargs={"starting_position": "-1", "on_error": on_error}) + thread.daemon = True + thread.start() + time.sleep(15) + thread.join() + assert isinstance(on_error.err, AuthenticationError) + + credential = EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']) + auth_uri = "sb://{}/{}".format(live_eventhub['hostname'], live_eventhub['event_hub']) + token = credential.get_token(auth_uri).token + producer_client = EventHubProducerClient(fully_qualified_namespace=live_eventhub['hostname'], + eventhub_name=live_eventhub['event_hub'], + credential=EventHubSASTokenCredential(token, time.time() + 5), + uamqp_transport=uamqp_transport) + time.sleep(10) + # expired credential + with producer_client: + with pytest.raises(AuthenticationError): + producer_client.create_batch(partition_id='0') + + consumer_client = EventHubConsumerClient(fully_qualified_namespace=live_eventhub['hostname'], + eventhub_name=live_eventhub['event_hub'], + credential=EventHubSASTokenCredential(token, time.time() + 7), + consumer_group='$Default', + retry_total=0, + uamqp_transport=uamqp_transport) + on_error.err = None + with consumer_client: + thread = threading.Thread(target=consumer_client.receive, args=(on_event,), + kwargs={"starting_position": "-1", "on_error": on_error}) + thread.daemon = True + thread.start() + time.sleep(15) + thread.join() + + assert isinstance(on_error.err, AuthenticationError) + + credential = EventHubSharedKeyCredential('fakekey', live_eventhub['access_key']) + producer_client = EventHubProducerClient(fully_qualified_namespace=live_eventhub['hostname'], + eventhub_name=live_eventhub['event_hub'], + credential=credential, + uamqp_transport=uamqp_transport) + with producer_client: + with pytest.raises(AuthenticationError): + producer_client.create_batch(partition_id='0') + + consumer_client = EventHubConsumerClient(fully_qualified_namespace=live_eventhub['hostname'], + eventhub_name=live_eventhub['event_hub'], + credential=credential, + consumer_group='$Default', + retry_total=0, + uamqp_transport=uamqp_transport) + on_error.err = None + with consumer_client: + thread = threading.Thread(target=consumer_client.receive, args=(on_event,), + kwargs={"starting_position": "-1", "on_error": on_error}) + thread.daemon = True + thread.start() + time.sleep(15) + thread.join() + assert isinstance(on_error.err, AuthenticationError) + + credential = EventHubSharedKeyCredential(live_eventhub['key_name'], 'fakekey') + producer_client = EventHubProducerClient(fully_qualified_namespace=live_eventhub['hostname'], + eventhub_name=live_eventhub['event_hub'], + credential=credential, + uamqp_transport=uamqp_transport) + + with producer_client: + with pytest.raises(AuthenticationError): + producer_client.create_batch(partition_id='0') + + consumer_client = EventHubConsumerClient(fully_qualified_namespace=live_eventhub['hostname'], + eventhub_name=live_eventhub['event_hub'], + credential=credential, + consumer_group='$Default', + retry_total=0, + uamqp_transport=uamqp_transport) + on_error.err = None + with consumer_client: + thread = threading.Thread(target=consumer_client.receive, args=(on_event,), + kwargs={"starting_position": "-1", "on_error": on_error}) + thread.daemon = True + thread.start() + time.sleep(15) + thread.join() + assert isinstance(on_error.err, AuthenticationError) + + producer_client = EventHubProducerClient(fully_qualified_namespace=live_eventhub['hostname'], + eventhub_name=live_eventhub['event_hub'], + credential=env_credential, + connection_verify="cacert.pem", + uamqp_transport=uamqp_transport) + + # TODO: this seems like a bug from uamqp, should be ConnectError? + with producer_client: + with pytest.raises(EventHubError): + producer_client.create_batch(partition_id='0') + + consumer_client = EventHubConsumerClient(fully_qualified_namespace=live_eventhub['hostname'], + eventhub_name=live_eventhub['event_hub'], + consumer_group='$Default', + credential=env_credential, + retry_total=0, + connection_verify="cacert.pem", + uamqp_transport=uamqp_transport) + with consumer_client: + thread = threading.Thread(target=consumer_client.receive, args=(on_event,), + kwargs={"starting_position": "-1", "on_error": on_error}) + thread.daemon = True + thread.start() + time.sleep(15) + thread.join() + + # TODO: this seems like a bug from uamqp, should be ConnectError? + assert isinstance(on_error.err, FileNotFoundError) + + # Skipping on OSX - it's raising a ConnectionLostError + if not sys.platform.startswith('darwin'): + producer_client = EventHubProducerClient(fully_qualified_namespace=live_eventhub['hostname'], + eventhub_name=live_eventhub['event_hub'], + credential=env_credential, + custom_endpoint_address="fakeaddr", + uamqp_transport=uamqp_transport) + + with producer_client: + with pytest.raises(AuthenticationError): + producer_client.create_batch(partition_id='0') + + consumer_client = EventHubConsumerClient(fully_qualified_namespace=live_eventhub['hostname'], + eventhub_name=live_eventhub['event_hub'], + consumer_group='$Default', + credential=env_credential, + retry_total=0, + custom_endpoint_address="fakeaddr", + uamqp_transport=uamqp_transport) + with consumer_client: + thread = threading.Thread(target=consumer_client.receive, args=(on_event,), + kwargs={"starting_position": "-1", "on_error": on_error}) + thread.daemon = True + thread.start() + time.sleep(15) + thread.join() + + assert isinstance(on_error.err, AuthenticationError) diff --git a/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_properties.py b/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_properties.py index 6a4cb8b6eccf..875fae6c715c 100644 --- a/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_properties.py +++ b/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_properties.py @@ -8,7 +8,7 @@ from azure.eventhub import EventHubSharedKeyCredential from azure.eventhub import EventHubConsumerClient -from azure.eventhub.exceptions import AuthenticationError, ConnectError, EventHubError +from azure.eventhub.exceptions import AuthenticationError, ConnectError, ConnectionLostError, EventHubError @pytest.mark.liveTest @@ -53,7 +53,7 @@ def test_get_properties_with_connect_error(live_eventhub, uamqp_transport): uamqp_transport=uamqp_transport ) with client: - with pytest.raises(EventHubError) as e: # This can be either ConnectError or ConnectionLostError + with pytest.raises(ConnectError) as e: client.get_eventhub_properties() @pytest.mark.liveTest diff --git a/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_receive.py b/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_receive.py index 22133f7983e3..7597531904af 100644 --- a/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_receive.py +++ b/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_receive.py @@ -4,12 +4,9 @@ # license information. #-------------------------------------------------------------------------- -import os import threading import pytest import time -import datetime -import uamqp from azure.eventhub import EventData, TransportType, EventHubConsumerClient from azure.eventhub.exceptions import EventHubError diff --git a/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_reconnect.py b/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_reconnect.py index c07489acface..87770adf80f3 100644 --- a/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_reconnect.py +++ b/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_reconnect.py @@ -15,9 +15,15 @@ ) from azure.eventhub.exceptions import OperationTimeoutError from azure.eventhub._utils import transform_outbound_single_message -import uamqp -from uamqp import compat -from azure.eventhub._transport._uamqp_transport import UamqpTransport +from azure.eventhub._pyamqp.authentication import SASTokenAuth +from azure.eventhub._pyamqp import ReceiveClient, error, constants +from azure.eventhub._transport._pyamqp_transport import PyamqpTransport +try: + import uamqp + from uamqp import compat + from azure.eventhub._transport._uamqp_transport import UamqpTransport +except (ModuleNotFoundError, ImportError): + UamqpTransport = None @pytest.mark.liveTest @@ -37,7 +43,7 @@ def test_send_with_long_interval_sync(live_eventhub, sleep, uamqp_transport, tim if uamqp_transport: sender._producers[test_partition]._handler._connection._conn.destroy() else: - pass + sender._producers[test_partition]._handler._connection.close() batch = sender.create_batch(partition_id=test_partition) batch.add(EventData(b"A single event")) sender.send_batch(batch) @@ -45,16 +51,21 @@ def test_send_with_long_interval_sync(live_eventhub, sleep, uamqp_transport, tim received = [] uri = "sb://{}/{}".format(live_eventhub['hostname'], live_eventhub['event_hub']) - if uamqp_transport: - sas_auth = uamqp.authentication.SASTokenAuth.from_shared_access_key( - uri, live_eventhub['key_name'], live_eventhub['access_key']) source = "amqps://{}/{}/ConsumerGroups/{}/Partitions/{}".format( live_eventhub['hostname'], live_eventhub['event_hub'], live_eventhub['consumer_group'], test_partition) if uamqp_transport: + sas_auth = uamqp.authentication.SASTokenAuth.from_shared_access_key( + uri, live_eventhub['key_name'], live_eventhub['access_key']) receiver = uamqp.ReceiveClient(source, auth=sas_auth, debug=False, timeout=5000, prefetch=500) + else: + sas_auth = SASTokenAuth( + uri, uri, live_eventhub['key_name'], live_eventhub['access_key'] + ) + receiver = ReceiveClient(live_eventhub['hostname'], source, auth=sas_auth, debug=False, link_credit=500) + try: receiver.open() # receive_message_batch() returns immediately once it receives any messages before the max_batch_size @@ -71,9 +82,16 @@ def test_send_with_long_interval_sync(live_eventhub, sleep, uamqp_transport, tim @pytest.mark.liveTest def test_send_connection_idle_timeout_and_reconnect_sync(connstr_receivers, uamqp_transport, timeout_factor): connection_str, receivers = connstr_receivers - amqp_transport = UamqpTransport + if uamqp_transport: + amqp_transport = UamqpTransport + retry_total = 3 + timeout_exc = compat.TimeoutException + else: + amqp_transport = PyamqpTransport + retry_total = 0 + timeout_exc = TimeoutError client = EventHubProducerClient.from_connection_string( - conn_str=connection_str, idle_timeout=10, uamqp_transport=uamqp_transport + conn_str=connection_str, idle_timeout=10, retry_total=retry_total, uamqp_transport=uamqp_transport ) with client: ed = EventData('data') @@ -89,8 +107,8 @@ def test_send_connection_idle_timeout_and_reconnect_sync(connstr_receivers, uamq uamqp.errors.MessageHandlerError, OperationTimeoutError)): sender._send_event_data() else: - # for pyamqp add later - pass + with pytest.raises(error.AMQPConnectionError): + sender._send_event_data() if uamqp_transport: sender._send_event_data_with_retry() @@ -116,7 +134,7 @@ def test_send_connection_idle_timeout_and_reconnect_sync(connstr_receivers, uamq received_ed1 = EventData._from_message(messages[0]) assert received_ed1.body_as_str() == 'data' break - except (compat.TimeoutException, TimeoutError): + except timeout_exc: retry += 1 @@ -135,14 +153,20 @@ def on_event_received(event): with client: consumer = client._create_consumer("$default", "0", "-1", on_event_received) with consumer: - consumer._open() + while not consumer.handler_ready: + consumer._open() time.sleep(11) ed = EventData("Event") senders[0].send(ed) - consumer._handler.do_work() - assert consumer._handler._connection._state == uamqp.c_uamqp.ConnectionState.DISCARDING + if uamqp_transport: + consumer._handler.do_work() + assert consumer._handler._connection._state == uamqp.c_uamqp.ConnectionState.DISCARDING + else: + with pytest.raises(error.AMQPConnectionError): + consumer._handler.do_work() + assert consumer._handler._connection.state == constants.ConnectionState.END duration = 10 now_time = time.time() diff --git a/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_send.py b/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_send.py index d75b4d013470..7728701b3655 100644 --- a/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_send.py +++ b/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_send.py @@ -11,8 +11,6 @@ import json import sys -import uamqp -from uamqp.message import MessageProperties from azure.eventhub import EventData, TransportType, EventDataBatch from azure.eventhub import EventHubProducerClient, EventHubConsumerClient from azure.eventhub.exceptions import EventDataSendError, OperationTimeoutError @@ -23,9 +21,17 @@ AmqpMessageProperties, ) try: - from azure.eventhub._transport._uamqp_transport import UamqpTransport -except (ImportError, ModuleNotFoundError): - UamqpTransport = None + import uamqp + from uamqp.constants import TransportType as uamqp_TransportType, MessageState + from uamqp.message import MessageProperties +except (ModuleNotFoundError, ImportError): + uamqp = None + uamqp_TransportType = TransportType + MessageProperties = None +from azure.eventhub._pyamqp.message import Properties +from azure.eventhub._pyamqp.authentication import SASTokenAuth +from azure.eventhub._pyamqp import ReceiveClient +from azure.eventhub._pyamqp.error import AMQPConnectionError @pytest.mark.liveTest @@ -76,19 +82,24 @@ def test_send_with_partition_key(connstr_receivers, live_eventhub, uamqp_transpo if received: break retry_total += 1 - except uamqp.errors.ConnectionClose: + except AMQPConnectionError: for r in reconnect_receivers: r.close() uri = "sb://{}/{}".format(live_eventhub['hostname'], live_eventhub['event_hub']) - sas_auth = uamqp.authentication.SASTokenAuth.from_shared_access_key( - uri, live_eventhub['key_name'], live_eventhub['access_key']) - source = "amqps://{}/{}/ConsumerGroups/{}/Partitions/{}".format( live_eventhub['hostname'], live_eventhub['event_hub'], live_eventhub['consumer_group'], index) - partition = uamqp.ReceiveClient(source, auth=sas_auth, debug=True, timeout=0, prefetch=500) + if uamqp_transport: + sas_auth = uamqp.authentication.SASTokenAuth.from_shared_access_key( + uri, live_eventhub['key_name'], live_eventhub['access_key']) + partition = uamqp.ReceiveClient(source, auth=sas_auth, debug=False, timeout=0, prefetch=500) + else: + sas_auth = SASTokenAuth( + uri, uri, live_eventhub['key_name'], live_eventhub['access_key']) + partition = ReceiveClient(live_eventhub['hostname'], source, auth=sas_auth, network_trace=False, timeout=0, link_credit=500) + partition.open() reconnect_receivers.append(partition) retry_total += 1 if retry_total == 3: @@ -124,22 +135,6 @@ def test_send_and_receive_large_body_size(connstr_receivers, uamqp_transport, ti assert len(list(received[0].body)[0]) == payload assert len(list(received[1].body)[0]) == payload - client = EventHubProducerClient.from_connection_string(connection_str) - with client: - payload = 250 * 1024 - batch = client.create_batch() - batch.add(EventData("A" * payload)) - client.send_batch(batch) - client.send_event(EventData("A" * payload)) - - received = [] - for r in receivers: - received.extend([EventData._from_message(x) for x in r.receive_message_batch(timeout=timeout)]) - - assert len(received) == 2 - assert len(list(received[0].body)[0]) == payload - assert len(list(received[1].body)[0]) == payload - @pytest.mark.liveTest def test_send_amqp_annotated_message(connstr_receivers, uamqp_transport): @@ -287,7 +282,7 @@ def test_send_and_receive_small_body(connstr_receivers, payload, uamqp_transport @pytest.mark.liveTest def test_send_partition(connstr_receivers, uamqp_transport, timeout_factor): connection_str, receivers = connstr_receivers - timeout = 5 * timeout_factor + timeout = 10 * timeout_factor client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) with client: @@ -319,11 +314,16 @@ def test_send_partition(connstr_receivers, uamqp_transport, timeout_factor): client.send_batch(batch) client.send_event(EventData(b"Data"), partition_id="0") + with client: + batch = EventDataBatch(partition_id="0") + batch.add(EventData(b"Data")) + client.send_batch(batch) + time.sleep(5) partition_0 = receivers[0].receive_message_batch(timeout=timeout) partition_1 = receivers[1].receive_message_batch(timeout=timeout) - assert len(partition_0) >= 2 - assert len(partition_0) + len(partition_1) == 4 + assert len(partition_0) >= 3 + assert len(partition_0) + len(partition_1) == 5 @pytest.mark.liveTest @@ -358,7 +358,11 @@ def test_send_multiple_partitions_with_app_prop(connstr_receivers, uamqp_transpo app_prop_key = "raw_prop" app_prop_value = "raw_value" app_prop = {app_prop_key: app_prop_value} - client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) + client = EventHubProducerClient.from_connection_string( + connection_str, + uamqp_transport=uamqp_transport, + transport_type=TransportType.Amqp + ) with client: ed0 = EventData(b"Message 0") ed0.properties = app_prop @@ -389,7 +393,7 @@ def test_send_over_websocket_sync(connstr_receivers, uamqp_transport, timeout_fa timeout = 10 * timeout_factor connection_str, receivers = connstr_receivers client = EventHubProducerClient.from_connection_string( - connection_str, transport_type=uamqp.constants.TransportType.AmqpOverWebsocket, uamqp_transport=uamqp_transport + connection_str, transport_type=uamqp_TransportType.AmqpOverWebsocket, uamqp_transport=uamqp_transport ) with client: @@ -435,7 +439,11 @@ def test_send_with_create_event_batch_with_app_prop_sync(connstr_receivers, uamq def test_send_list(connstr_receivers, uamqp_transport, timeout_factor): connection_str, receivers = connstr_receivers timeout = 10 * timeout_factor - client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) + client = EventHubProducerClient.from_connection_string( + connection_str, + uamqp_transport=uamqp_transport, + transport_type=uamqp_TransportType.Amqp + ) payload = "A1" with client: client.send_batch([EventData(payload)]) @@ -460,7 +468,6 @@ def test_send_list_partition(connstr_receivers, uamqp_transport, timeout_factor) assert received.body_as_str() == payload - @pytest.mark.parametrize("to_send, exception_type", [([EventData("A"*1024)]*1100, ValueError), ("any str", AttributeError)]) @@ -472,11 +479,9 @@ def test_send_list_wrong_data(connection_str, to_send, exception_type, uamqp_tra client.send_batch(to_send) - @pytest.mark.parametrize("partition_id, partition_key", [("0", None), (None, "pk")]) def test_send_batch_pid_pk(invalid_hostname, partition_id, partition_key, uamqp_transport): # Use invalid_hostname because this is not a live test. - amqp_transport = UamqpTransport if uamqp_transport else None client = EventHubProducerClient.from_connection_string(invalid_hostname, uamqp_transport=uamqp_transport) batch = EventDataBatch(partition_id=partition_id, partition_key=partition_key) with client: @@ -484,7 +489,7 @@ def test_send_batch_pid_pk(invalid_hostname, partition_id, partition_key, uamqp_ client.send_batch(batch, partition_id=partition_id, partition_key=partition_key) - +@pytest.mark.liveTest def test_send_with_callback(connstr_receivers, uamqp_transport): def on_error(events, pid, err): @@ -530,63 +535,3 @@ def on_success(events, pid): assert sent_events[-1][1] == "0" assert not on_error.err - -# TODO: add more checks after LegacyMessage has been added -@pytest.mark.liveTest -def test_send_message_modify_backcompat(connstr_receivers, uamqp_transport, timeout_factor): - connection_str, receivers = connstr_receivers - if uamqp_transport: - properties = MessageProperties - - timeout = 10 * timeout_factor - outgoing_event_data = EventData(body="hello") - message = outgoing_event_data.message - message.properties = properties(user_id='fake_user') - assert outgoing_event_data.message.properties.user_id == b'fake_user' - assert outgoing_event_data.message.state == uamqp.constants.MessageState.WaitingToBeSent - assert outgoing_event_data.message.delivery_annotations is None - assert outgoing_event_data.message.delivery_no is None - assert outgoing_event_data.message.delivery_tag is None - assert outgoing_event_data.message.on_send_complete is None - assert outgoing_event_data.message.footer is None - assert outgoing_event_data.message.retries == 0 - assert outgoing_event_data.message.idle_time == 0 - client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) - with client: - client.send_batch([outgoing_event_data]) - received = [] - for r in receivers: - received.extend([EventData._from_message(x) for x in r.receive_message_batch(timeout=timeout)]) - - assert len(received) == 1 - received_ed = received[0] - # check that setting properties directly on uamqp message doesn't update the outgoing message from the event data - assert received_ed.message.properties.user_id is None - assert received_ed.message.state == uamqp.constants.MessageState.ReceivedSettled - assert received_ed.message.delivery_annotations is None - assert received_ed.message.delivery_no >= 1 - assert received_ed.message.delivery_tag is None - assert received_ed.message.on_send_complete is None - assert received_ed.message.footer is None - assert received_ed.message.retries >= 0 - assert received_ed.message.idle_time >= 0 - - # setting message properties by calling event data properties SHOULD update the outgoing uamqp message - received_ed.properties = {'prop': 'test'} - received_ed.message_id = "id_message" - received_ed.content_type = "content type" - received_ed.correlation_id = "correlation" - - client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) - with client: - client.send_batch([received_ed]) - received = [] - for r in receivers: - received.extend([EventData._from_message(x) for x in r.receive_message_batch(timeout=timeout)]) - - assert len(received) == 1 - received_ed = received[0] - assert received_ed.message.application_properties == {b"prop": b"test"} - assert received_ed.message_id == "id_message" - assert received_ed.content_type == "content type" - assert received_ed.correlation_id == "correlation" diff --git a/sdk/eventhub/azure-eventhub/tests/perfstress_tests/_test_base.py b/sdk/eventhub/azure-eventhub/tests/perfstress_tests/_test_base.py index dbef4b8f73c2..03b107a8fc5c 100644 --- a/sdk/eventhub/azure-eventhub/tests/perfstress_tests/_test_base.py +++ b/sdk/eventhub/azure-eventhub/tests/perfstress_tests/_test_base.py @@ -5,9 +5,10 @@ import asyncio from uuid import uuid4 +from datetime import datetime from azure_devtools.perfstress_tests import BatchPerfTest, EventPerfTest, get_random_bytes -from azure.eventhub import EventHubProducerClient, EventHubConsumerClient, EventData +from azure.eventhub import EventHubProducerClient, EventHubConsumerClient, EventData, TransportType from azure.eventhub.aio import ( EventHubProducerClient as AsyncEventHubProducerClient, EventHubConsumerClient as AsyncEventHubConsumerClient @@ -34,25 +35,45 @@ def __init__(self, arguments): self.checkpoint_store = BlobCheckpointStore.from_connection_string(storage_connection_str, self.container_name) self.async_checkpoint_store = AsyncBlobCheckpointStore.from_connection_string(storage_connection_str, self.container_name) + transport_type = TransportType.AmqpOverWebsocket if arguments.transport_type==1 else TransportType.Amqp + self.consumer = EventHubConsumerClient.from_connection_string( connection_string, _EventHubProcessorTest.consumer_group, eventhub_name=eventhub_name, checkpoint_store=self.checkpoint_store, - load_balancing_strategy=arguments.load_balancing_strategy + load_balancing_strategy=arguments.load_balancing_strategy, + transport_type=transport_type, + uamqp_transport=arguments.uamqp_transport ) self.async_consumer = AsyncEventHubConsumerClient.from_connection_string( connection_string, _EventHubProcessorTest.consumer_group, eventhub_name=eventhub_name, checkpoint_store=self.async_checkpoint_store, - load_balancing_strategy=arguments.load_balancing_strategy + load_balancing_strategy=arguments.load_balancing_strategy, + transport_type=transport_type, + uamqp_transport=arguments.uamqp_transport ) if arguments.preload: - self.async_producer = AsyncEventHubProducerClient.from_connection_string(connection_string, eventhub_name=eventhub_name) + self.data = get_random_bytes(self.args.event_size) + self.async_producer = AsyncEventHubProducerClient.from_connection_string(connection_string, eventhub_name=eventhub_name, transport_type=transport_type, uamqp_transport=arguments.uamqp_transport) + + def _build_event(self): + event = EventData(self.data) + if self.args.event_extra: + event.raw_amqp_message.header.first_acquirer = True + event.raw_amqp_message.properties.subject = 'perf' + event.properties = { + "key1": b"data", + "key2": 42, + "key3": datetime.now(), + "key4": "foobar", + "key5": uuid4() + } + return event async def _preload_eventhub(self): - data = get_random_bytes(self.args.event_size) async with self.async_producer as producer: partitions = await producer.get_partition_ids() total_events = 0 @@ -65,13 +86,13 @@ async def _preload_eventhub(self): batch = await producer.create_batch() for i in range(events_to_add): try: - batch.add(EventData(data)) + batch.add(self._build_event()) except ValueError: # Batch full await producer.send_batch(batch) print(f"Loaded {i} of {events_to_add} events.") batch = await producer.create_batch() - batch.add(EventData(data)) + batch.add(self._build_event()) await producer.send_batch(batch) print(f"Finished loading {events_to_add} events.") @@ -120,6 +141,10 @@ def add_arguments(parser): parser.add_argument('--processing-delay-strategy', nargs='?', type=str, help="Whether to 'sleep' or 'spin' during processing delay. Default is 'sleep'.", default='sleep') parser.add_argument('--preload', nargs='?', type=int, help='Ensure the specified number of events are available across all partitions. Default is 0.', default=0) parser.add_argument('--use-storage-checkpoint', action="store_true", help="Use Blob storage for checkpointing. Default is False (in-memory checkpointing).", default=False) + parser.add_argument('--uamqp-transport', action="store_true", help="Switch to use uamqp transport. Default is False (pyamqp).", default=False) + parser.add_argument('--transport-type', nargs='?', type=int, help="Use Amqp (0) or Websocket (1) transport type. Default is Amqp.", default=0) + parser.add_argument('--event-extra', action="store_true", help="Add properties to the events to increase payload and serialization. Default is False.", default=False) + class _SendTest(BatchPerfTest): @@ -129,13 +154,20 @@ def __init__(self, arguments): super().__init__(arguments) connection_string = self.get_from_env("AZURE_EVENTHUB_CONNECTION_STRING") eventhub_name = self.get_from_env("AZURE_EVENTHUB_NAME") + + transport_type = TransportType.AmqpOverWebsocket if arguments.transport_type==1 else TransportType.Amqp + self.producer = EventHubProducerClient.from_connection_string( connection_string, - eventhub_name=eventhub_name + eventhub_name=eventhub_name, + transport_type=transport_type, + uamqp_transport=arguments.uamqp_transport ) self.async_producer = AsyncEventHubProducerClient.from_connection_string( connection_string, - eventhub_name=eventhub_name + eventhub_name=eventhub_name, + transport_type=transport_type, + uamqp_transport=arguments.uamqp_transport ) async def setup(self): @@ -156,3 +188,6 @@ def add_arguments(parser): super(_SendTest, _SendTest).add_arguments(parser) parser.add_argument('--event-size', nargs='?', type=int, help='Size of event body (in bytes). Defaults to 100 bytes', default=100) parser.add_argument('--batch-size', nargs='?', type=int, help='The number of events that should be included in each batch. Defaults to 100', default=100) + parser.add_argument('--uamqp-transport', action="store_true", help="Switch to use uamqp transport. Default is False (pyamqp).", default=False) + parser.add_argument('--transport-type', nargs='?', type=int, help="Use Amqp (0) or Websocket (1) transport type. Default is Amqp.", default=0) + parser.add_argument('--event-extra', action="store_true", help="Add properties to the events to increase payload and serialization. Default is False.", default=False) diff --git a/sdk/eventhub/azure-eventhub/tests/perfstress_tests/process_events_batch.py b/sdk/eventhub/azure-eventhub/tests/perfstress_tests/process_events_batch.py index 224652255a8e..15662ba70aed 100644 --- a/sdk/eventhub/azure-eventhub/tests/perfstress_tests/process_events_batch.py +++ b/sdk/eventhub/azure-eventhub/tests/perfstress_tests/process_events_batch.py @@ -28,7 +28,7 @@ def process_events_sync(self, partition_context, events): pass # Consume properties and body. - _ = [(e.properties, e.body) for e in events] + _ = [(list(e.body), str(e)) for e in events] if self.args.checkpoint_interval: self._partition_event_count[partition_context.partition_id] += len(events) @@ -51,9 +51,9 @@ async def process_events_async(self, partition_context, events): starttime = time.time() while (time.time() - starttime) < delay_in_seconds: pass - + # Consume properties and body. - _ = [(e.properties, e.body) for e in events] + _ = [(list(e.body), str(e)) for e in events] if self.args.checkpoint_interval: self._partition_event_count[partition_context.partition_id] += len(events) @@ -66,9 +66,11 @@ async def process_events_async(self, partition_context, events): await self.error_raised_async(e) def process_error_sync(self, _, error): + print(error) self.error_raised_sync(error) async def process_error_async(self, _, error): + print(error) await self.error_raised_async(error) def start_events_sync(self) -> None: @@ -100,4 +102,4 @@ async def start_events_async(self) -> None: @staticmethod def add_arguments(parser): super(ProcessEventsBatchTest, ProcessEventsBatchTest).add_arguments(parser) - parser.add_argument('--max-batch-size', nargs='?', type=int, help='Maximum number of events to process in a single batch. Defaults to 100.', default=100) + parser.add_argument('--max-batch-size', nargs='?', type=int, help='Maximum number of events to process in a single batch. Defaults to 100.', default=100) \ No newline at end of file diff --git a/sdk/eventhub/azure-eventhub/tests/perfstress_tests/send_event_batch.py b/sdk/eventhub/azure-eventhub/tests/perfstress_tests/send_event_batch.py index 538f2cc271d9..2a2bf0a506f7 100644 --- a/sdk/eventhub/azure-eventhub/tests/perfstress_tests/send_event_batch.py +++ b/sdk/eventhub/azure-eventhub/tests/perfstress_tests/send_event_batch.py @@ -3,6 +3,9 @@ # Licensed under the MIT License. See License.txt in the project root for license information. # -------------------------------------------------------------------------------------------- +from datetime import datetime +from uuid import uuid4 + from ._test_base import _SendTest from azure_devtools.perfstress_tests import get_random_bytes @@ -15,16 +18,30 @@ def __init__(self, arguments): super().__init__(arguments) self.data = get_random_bytes(self.args.event_size) + def _build_event(self): + event = EventData(self.data) + if self.args.event_extra: + event.raw_amqp_message.header.first_acquirer = True + event.raw_amqp_message.properties.subject = 'perf' + event.properties = { + "key1": b"data", + "key2": 42, + "key3": datetime.now(), + "key4": "foobar", + "key5": uuid4() + } + return event + def run_batch_sync(self): batch = self.producer.create_batch() for _ in range(self.args.batch_size): - batch.add(EventData(self.data)) + batch.add(self._build_event()) self.producer.send_batch(batch) return self.args.batch_size async def run_batch_async(self): batch = await self.async_producer.create_batch() for _ in range(self.args.batch_size): - batch.add(EventData(self.data)) + batch.add(self._build_event()) await self.async_producer.send_batch(batch) return self.args.batch_size diff --git a/sdk/eventhub/azure-eventhub/tests/perfstress_tests/send_events.py b/sdk/eventhub/azure-eventhub/tests/perfstress_tests/send_events.py index 8be9b87df685..851a84be1540 100644 --- a/sdk/eventhub/azure-eventhub/tests/perfstress_tests/send_events.py +++ b/sdk/eventhub/azure-eventhub/tests/perfstress_tests/send_events.py @@ -2,6 +2,8 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for license information. # -------------------------------------------------------------------------------------------- +from datetime import datetime +from uuid import uuid4 from ._test_base import _SendTest @@ -15,20 +17,34 @@ def __init__(self, arguments): super().__init__(arguments) self.data = get_random_bytes(self.args.event_size) + def _build_event(self): + event = EventData(self.data) + if self.args.event_extra: + event.raw_amqp_message.header.first_acquirer = True + event.raw_amqp_message.properties.subject = 'perf' + event.properties = { + "key1": b"data", + "key2": 42, + "key3": datetime.now(), + "key4": "foobar", + "key5": uuid4() + } + return event + def run_batch_sync(self): if self.args.batch_size > 1: self.producer.send_batch( - [EventData(self.data) for _ in range(self.args.batch_size)] + [self._build_event() for _ in range(self.args.batch_size)] ) else: - self.producer.send_event(EventData(self.data)) + self.producer.send_event(self._build_event()) return self.args.batch_size async def run_batch_async(self): if self.args.batch_size > 1: await self.async_producer.send_batch( - [EventData(self.data) for _ in range(self.args.batch_size)] + [self._build_event() for _ in range(self.args.batch_size)] ) else: - await self.async_producer.send_event(EventData(self.data)) + await self.async_producer.send_event(self._build_event()) return self.args.batch_size diff --git a/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/async/test_send_receive_pyamqp_async.py b/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/async/test_send_receive_pyamqp_async.py new file mode 100644 index 000000000000..3bda48ef1653 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/async/test_send_receive_pyamqp_async.py @@ -0,0 +1,56 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +import pytest +import asyncio +import logging + +from azure.eventhub._pyamqp.aio import _authentication_async +from azure.eventhub._pyamqp.aio import ReceiveClientAsync, SendClientAsync +from azure.eventhub._pyamqp.constants import TransportType +from azure.eventhub._pyamqp.message import Message + + +async def send_message(live_eventhub): + uri = "sb://{}/{}".format(live_eventhub['hostname'], live_eventhub['event_hub']) + sas_auth = _authentication_async.SASTokenAuthAsync( + uri=uri, + audience=uri, + username=live_eventhub['key_name'], + password=live_eventhub['access_key'] + ) + + target = "amqps://{}/{}/Partitions/{}".format( + live_eventhub['hostname'], + live_eventhub['event_hub'], + live_eventhub['partition']) + + message = Message(value="Single Message") + + async with SendClientAsync(live_eventhub['hostname'], target, auth=sas_auth, debug=True, transport_type=TransportType.Amqp) as send_client: + await send_client.send_message_async(message) + +@pytest.mark.asyncio +async def test_event_hubs_client_amqp_async(live_eventhub): + uri = "sb://{}/{}".format(live_eventhub['hostname'], live_eventhub['event_hub']) + sas_auth = _authentication_async.SASTokenAuthAsync( + uri=uri, + audience=uri, + username=live_eventhub['key_name'], + password=live_eventhub['access_key'] + ) + + source = "amqps://{}/{}/ConsumerGroups/{}/Partitions/{}".format( + live_eventhub['hostname'], + live_eventhub['event_hub'], + live_eventhub['consumer_group'], + live_eventhub['partition']) + + await send_message(live_eventhub=live_eventhub) + + async with ReceiveClientAsync(live_eventhub['hostname'], source, auth=sas_auth, debug=False, timeout=500, prefetch=1, transport_type=TransportType.Amqp) as receive_client: + messages = await receive_client.receive_message_batch_async(max_batch_size=1) + assert len(messages) > 0 + diff --git a/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/async/test_websocket_async.py b/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/async/test_websocket_async.py new file mode 100644 index 000000000000..48e9c0b00692 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/async/test_websocket_async.py @@ -0,0 +1,56 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +import pytest +import asyncio +import logging + +from azure.eventhub._pyamqp.aio import _authentication_async +from azure.eventhub._pyamqp.aio import ReceiveClientAsync, SendClientAsync +from azure.eventhub._pyamqp.constants import TransportType +from azure.eventhub._pyamqp.message import Message + + +async def send_message(live_eventhub): + uri = "sb://{}/{}".format(live_eventhub['hostname'], live_eventhub['event_hub']) + sas_auth = _authentication_async.SASTokenAuthAsync( + uri=uri, + audience=uri, + username=live_eventhub['key_name'], + password=live_eventhub['access_key'] + ) + + target = "amqps://{}/{}/Partitions/{}".format( + live_eventhub['hostname'], + live_eventhub['event_hub'], + live_eventhub['partition']) + + message = Message(value="Single Message") + + async with SendClientAsync(live_eventhub['hostname'], target, auth=sas_auth, debug=True, transport_type=TransportType.Amqp) as send_client: + await send_client.send_message_async(message) + +@pytest.mark.asyncio +async def test_event_hubs_client_web_socket_async(live_eventhub): + uri = "sb://{}/{}".format(live_eventhub['hostname'], live_eventhub['event_hub']) + sas_auth = _authentication_async.SASTokenAuthAsync( + uri=uri, + audience=uri, + username=live_eventhub['key_name'], + password=live_eventhub['access_key'] + ) + + source = "amqps://{}/{}/ConsumerGroups/{}/Partitions/{}".format( + live_eventhub['hostname'], + live_eventhub['event_hub'], + live_eventhub['consumer_group'], + live_eventhub['partition']) + + await send_message(live_eventhub=live_eventhub) + + async with ReceiveClientAsync(live_eventhub['hostname'] + '/$servicebus/websocket/', source, auth=sas_auth, debug=False, timeout=500, prefetch=1, transport_type=TransportType.AmqpOverWebsocket) as receive_client: + messages = await receive_client.receive_message_batch_async(max_batch_size=1) + assert len(messages) > 0 + diff --git a/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/synctests/test_mgmt_pyamqp.py b/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/synctests/test_mgmt_pyamqp.py new file mode 100644 index 000000000000..31edff23dc10 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/synctests/test_mgmt_pyamqp.py @@ -0,0 +1,69 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +import pytest +import time +from azure.identity import EnvironmentCredential, DefaultAzureCredential +from azure.eventhub import EventHubProducerClient, EventHubSharedKeyCredential +from azure.eventhub._client_base import EventHubSASTokenCredential +from azure.core.credentials import AzureSasCredential, AzureNamedKeyCredential + +@pytest.mark.livetest +def test_mgmt_call_conn_str(connstr_receivers): + connection_str, receivers = connstr_receivers + client = EventHubProducerClient.from_connection_string(connection_str) + client._start_producer("0",60) + +@pytest.mark.livetest +def test_mgmt_call_default_azure_credential(live_eventhub): + credential = DefaultAzureCredential() + client = EventHubProducerClient(fully_qualified_namespace=live_eventhub['hostname'], + eventhub_name=live_eventhub['event_hub'], + credential=credential, + user_agent='customized information') + client._start_producer("0",60) + +@pytest.mark.livetest +def test_mgmt_call_credential(live_eventhub): + credential = EnvironmentCredential() + client = EventHubProducerClient(fully_qualified_namespace=live_eventhub['hostname'], + eventhub_name=live_eventhub['event_hub'], + credential=credential, + user_agent='customized information') + client._start_producer("0",60) + +@pytest.mark.livetest +def test_mgmt_call_sas(live_eventhub): + hostname = live_eventhub["hostname"] + credential = EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']) + auth_uri = "sb://{}/{}".format(hostname, live_eventhub['event_hub']) + token = credential.get_token(auth_uri).token + client = EventHubProducerClient(fully_qualified_namespace=hostname, + eventhub_name=live_eventhub['event_hub'], + credential=EventHubSASTokenCredential(token, time.time() + 3000)) + client._start_producer("0",60) + assert True + +@pytest.mark.livetest +def test_mgmt_call_sas_credential(live_eventhub): + hostname = live_eventhub["hostname"] + credential = EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']) + auth_uri = "sb://{}/{}".format(hostname, live_eventhub['event_hub']) + token = credential.get_token(auth_uri).token.decode() + client = EventHubProducerClient(fully_qualified_namespace=hostname, + eventhub_name=live_eventhub['event_hub'], + credential=AzureSasCredential(token)) + client._start_producer("0",60) + assert True + +@pytest.mark.livetest +def test_mgmt_call_azure_named_key_credential(live_eventhub): + credential = AzureNamedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']) + client = EventHubProducerClient(fully_qualified_namespace=live_eventhub['hostname'], + eventhub_name=live_eventhub['event_hub'], + credential=credential) + + client._start_producer("0",60) + assert True diff --git a/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/synctests/test_send_receive_pyamqp.py b/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/synctests/test_send_receive_pyamqp.py new file mode 100644 index 000000000000..b6811cce0f81 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/synctests/test_send_receive_pyamqp.py @@ -0,0 +1,52 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +import pytest + +from azure.eventhub._pyamqp import authentication, ReceiveClient, SendClient +from azure.eventhub._pyamqp.constants import TransportType +from azure.eventhub._pyamqp.message import Message + + +def send_message(live_eventhub): + uri = "sb://{}/{}".format(live_eventhub['hostname'], live_eventhub['event_hub']) + sas_auth = authentication.SASTokenAuth( + uri=uri, + audience=uri, + username=live_eventhub['key_name'], + password=live_eventhub['access_key'] + ) + + target = "amqps://{}/{}/Partitions/{}".format( + live_eventhub['hostname'], + live_eventhub['event_hub'], + live_eventhub['partition']) + + message = Message(value="Single Message") + + with SendClient(live_eventhub['hostname'], target, auth=sas_auth, debug=True, transport_type=TransportType.Amqp) as send_client: + send_client.send_message(message) + + +def test_event_hubs_client_amqp(live_eventhub): + uri = "sb://{}/{}".format(live_eventhub['hostname'], live_eventhub['event_hub']) + sas_auth = authentication.SASTokenAuth( + uri=uri, + audience=uri, + username=live_eventhub['key_name'], + password=live_eventhub['access_key'] + ) + + source = "amqps://{}/{}/ConsumerGroups/{}/Partitions/{}".format( + live_eventhub['hostname'], + live_eventhub['event_hub'], + live_eventhub['consumer_group'], + live_eventhub['partition']) + + send_message(live_eventhub=live_eventhub) + + with ReceiveClient(live_eventhub['hostname'], source, auth=sas_auth, debug=False, timeout=500, prefetch=1, transport_type=TransportType.Amqp) as receive_client: + messages = receive_client.receive_message_batch(max_batch_size=1) + assert len(messages) > 0 diff --git a/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/synctests/test_websocket.py b/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/synctests/test_websocket.py new file mode 100644 index 000000000000..5d1a32c92c17 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/synctests/test_websocket.py @@ -0,0 +1,52 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +import pytest + +from azure.eventhub._pyamqp import authentication, ReceiveClient, SendClient +from azure.eventhub._pyamqp.constants import TransportType +from azure.eventhub._pyamqp.message import Message + + +def send_message(live_eventhub): + uri = "sb://{}/{}".format(live_eventhub['hostname'], live_eventhub['event_hub']) + sas_auth = authentication.SASTokenAuth( + uri=uri, + audience=uri, + username=live_eventhub['key_name'], + password=live_eventhub['access_key'] + ) + + target = "amqps://{}/{}/Partitions/{}".format( + live_eventhub['hostname'], + live_eventhub['event_hub'], + live_eventhub['partition']) + + message = Message(value="Single Message") + + with SendClient(live_eventhub['hostname'], target, auth=sas_auth, debug=True, transport_type=TransportType.Amqp) as send_client: + send_client.send_message(message) + + +def test_event_hubs_client_web_socket(live_eventhub): + uri = "sb://{}/{}".format(live_eventhub['hostname'], live_eventhub['event_hub']) + sas_auth = authentication.SASTokenAuth( + uri=uri, + audience=uri, + username=live_eventhub['key_name'], + password=live_eventhub['access_key'] + ) + + source = "amqps://{}/{}/ConsumerGroups/{}/Partitions/{}".format( + live_eventhub['hostname'], + live_eventhub['event_hub'], + live_eventhub['consumer_group'], + live_eventhub['partition']) + + send_message(live_eventhub=live_eventhub) + + with ReceiveClient(live_eventhub['hostname'] + '/$servicebus/websocket/', source, auth=sas_auth, debug=False, timeout=500, prefetch=1, transport_type=TransportType.AmqpOverWebsocket) as receive_client: + messages = receive_client.receive_message_batch(max_batch_size=1) + assert len(messages) > 0 diff --git a/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/unittest/test_amqp_value.py b/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/unittest/test_amqp_value.py new file mode 100644 index 000000000000..248370edd24e --- /dev/null +++ b/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/unittest/test_amqp_value.py @@ -0,0 +1,33 @@ +#------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +#-------------------------------------------------------------------------- +from unittest.mock import Mock +from base64 import encode +import os +import sys +import pytest +import uuid + +root_path = os.path.realpath('.') +sys.path.append(root_path) + +from azure.eventhub._pyamqp.types import AMQPTypes +from azure.eventhub._pyamqp.utils import amqp_uint_value, amqp_long_value, amqp_string_value + +def test_uint_value(): + value = amqp_uint_value(255) + assert value.get("VALUE") == 255 + assert value.get("TYPE") == AMQPTypes.uint + + +def test_long_value(): + value = amqp_long_value(255) + assert value.get("VALUE") == 255 + assert value.get("TYPE") == AMQPTypes.long + +def test_string_value(): + value = amqp_string_value("hello") + assert value.get("VALUE") == "hello" + assert value.get("TYPE") == AMQPTypes.string \ No newline at end of file diff --git a/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/unittest/test_authentication_pyamqp.py b/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/unittest/test_authentication_pyamqp.py new file mode 100644 index 000000000000..04e71380830c --- /dev/null +++ b/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/unittest/test_authentication_pyamqp.py @@ -0,0 +1,44 @@ + +import pytest +import functools + +from unittest.mock import Mock +from azure.eventhub._pyamqp.sasl import SASLAnonymousCredential, SASLPlainCredential +from azure.eventhub._pyamqp.authentication import SASLPlainAuth, JWTTokenAuth, SASTokenAuth + + +def test_sasl_plain_auth(): + auth = SASLPlainAuth( + authcid="authcid", + passwd="passwd", + authzid="Some Authzid" + ) + assert auth.auth_type=="AUTH_SASL_PLAIN" + assert auth.sasl.mechanism==b"PLAIN" + assert auth.sasl.start() == b'Some Authzid\x00authcid\x00passwd' + +def test_jwt_token_auth(): + credential = Mock() + attr = {"get_token.return_value": "my_token"} + credential.configure_mock(**attr) + auth = JWTTokenAuth( + uri="my_uri", + audience="my_audience_field", + get_token=functools.partial(credential.get_token, "my_auth_uri") + ) + + assert auth.uri == "my_uri" + assert auth.audience == "my_audience_field" + +def test_sas_token_auth(): + auth = SASTokenAuth( + uri="my_uri", + audience="my_audience", + username="username", + password="password" + ) + + assert auth.uri == "my_uri" + assert auth.audience == "my_audience" + assert auth.username == "username" + assert auth.password == "password" \ No newline at end of file diff --git a/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/unittest/test_client_creation_pyamqp.py b/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/unittest/test_client_creation_pyamqp.py new file mode 100644 index 000000000000..bb777b53c0cc --- /dev/null +++ b/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/unittest/test_client_creation_pyamqp.py @@ -0,0 +1,28 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- +import pytest + +from azure.eventhub._pyamqp import SendClient, ReceiveClient + +def test_send_client_creation(): + + sender = SendClient( + "fake.host.com", + "fake_eh", + auth="my_fake_auth" + ) + assert sender.target == "fake_eh" + assert sender._auth == "my_fake_auth" + + +def test_receive_client_creation(): + + receiver = ReceiveClient( + "fake.host.com", + "fake_eh", + auth="my_fake_auth" + ) + assert receiver.source == "fake_eh" + assert receiver._auth == "my_fake_auth" diff --git a/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/unittest/test_exceptions_pyamqp.py b/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/unittest/test_exceptions_pyamqp.py new file mode 100644 index 000000000000..d45f13147d22 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/unittest/test_exceptions_pyamqp.py @@ -0,0 +1,49 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- +import pytest + +from azure.eventhub._pyamqp import SendClient, Connection, authentication +from azure.eventhub._pyamqp.error import AMQPConnectionError + +def test_client_creation_exceptions(): + with pytest.raises(TypeError): + sender = SendClient( + "fake.host.com", + ) + assert sender._hostname == "fake.host.com" + +def test_connection_endpoint_exceptions(): + with pytest.raises(AMQPConnectionError): + endpoint = "fake.host.com" + connection = Connection(endpoint) + connection.open() + +def test_connection_sas_authentication_exception(): + uri = "sb://{}/{}".format("fake.host.come", "fake_eh") + + target = "amqps://{}/{}/Partitions/{}".format( + "fake.host.com", + "fake_eh", + "0") + sas_auth = authentication.SASTokenAuth( + uri=uri, + audience=uri, + username="key", + password="" + ) + with pytest.raises(AttributeError): + sender = SendClient("fake.host.com", target, auth=sas_auth) + sender.client_ready() + +def test_connection_sasl_annon_authentication_exception(): + target = "amqps://{}/{}/Partitions/{}".format( + "fake.host.com", + "fake_eh", + "0") + + sas_auth = authentication.SASLAnonymousCredential() + with pytest.raises(AttributeError): + sender = SendClient("fake.host.com", target, auth=sas_auth) + sender.client_ready() diff --git a/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/unittest/test_message.py b/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/unittest/test_message.py new file mode 100644 index 000000000000..bd75283463d0 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/unittest/test_message.py @@ -0,0 +1,65 @@ +#------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +#-------------------------------------------------------------------------- +from azure.eventhub._pyamqp.message import Message +from azure.eventhub._pyamqp.utils import AMQPTypes +from azure.eventhub._pyamqp._encode import encode_payload +from azure.eventhub._pyamqp._decode import decode_payload +from azure.eventhub.amqp._amqp_message import AmqpMessageProperties + + +def test_message(): + value = Message(value="test_message") + + assert value.data == None + assert value.value == "test_message" + + +def test_body_value(): + message = Message(value="test_message") + body_value = b"body" + + message = message._replace(data=body_value) + assert message.data == body_value + + output = bytearray() + encode_payload(output, message) + + message = decode_payload(memoryview(output)) + output = bytearray() + + body = message.data + assert body[0].get(b"TYPE").decode("utf-8") == AMQPTypes.binary + + +def test_delivery_tag(): + message = Message(value="test_message") + assert not message.delivery_annotations + + +def test_message_properties(): + + value = AmqpMessageProperties() + assert not value.user_id + + value = AmqpMessageProperties() + value.user_id = bytearray(b'testuseridlongstring') + assert value.user_id == b'testuseridlongstring' + + value = AmqpMessageProperties() + value.user_id = bytearray(b'') + assert value.user_id == b'' + + value = AmqpMessageProperties() + value.user_id =bytearray(b'short') + assert value.user_id == b'short' + + value = AmqpMessageProperties() + value.user_id = bytearray(b'!@#$%^&*()+_?') + assert value.user_id == b'!@#$%^&*()+_?' + + value = AmqpMessageProperties() + value.user_id = bytearray(b'\nweird\0user\1id\0\t') + assert value.user_id == b'\nweird\0user\1id\0\t' \ No newline at end of file diff --git a/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/unittest/test_message_components.py b/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/unittest/test_message_components.py new file mode 100644 index 000000000000..9cd4d036a876 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/unittest/test_message_components.py @@ -0,0 +1,136 @@ +#------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +#-------------------------------------------------------------------------- +import copy +import pickle +import pytest + +from azure.eventhub._pyamqp.message import ( + Properties, + Header, + BatchMessage, + Message +) + +def test_message_properties(): + + properties = Properties() + assert properties.user_id is None + + properties = Properties() + properties = properties._replace(user_id=b'') + assert properties.user_id == b'' + + properties = Properties() + properties = properties._replace(user_id=b'1') + assert properties.user_id == b'1' + + properties = Properties() + properties = properties._replace(user_id=b'short') + assert properties.user_id == b'short' + + properties = Properties() + properties = properties._replace(user_id=b'longuseridstring') + assert properties.user_id == b'longuseridstring' + + properties = Properties() + properties = properties._replace(user_id=b'!@#$%^&*()_+1234567890') + assert properties.user_id == b'!@#$%^&*()_+1234567890' + + properties = Properties() + properties = properties._replace(user_id=b'werid/0\0\1\t\n') + assert properties.user_id == b'werid/0\0\1\t\n' + +def test_deepcopy_batch_message(): + ## DEEPCOPY WITH MESSAGES IN BATCH THAT HAVE HEADER/PROPERTIES + properties = Properties() + properties = properties._replace(message_id = '2') + properties = properties._replace(user_id = '1') + properties = properties._replace(to = 'dkfj') + properties = properties._replace(subject = 'dsljv') + properties = properties._replace(reply_to = "kdjfk") + properties = properties._replace(correlation_id = 'ienag') + properties = properties._replace(content_type = 'b') + properties = properties._replace(content_encoding = '39ru') + properties = properties._replace(absolute_expiry_time = 24) + properties = properties._replace(creation_time = 10) + properties = properties._replace(group_id = '3irow') + properties = properties._replace(group_sequence = 39) + properties = properties._replace(reply_to_group_id = '39rud') + + header = Header() + header = header._replace(delivery_count = 3) + header = header._replace(ttl = 5) + header = header._replace(first_acquirer = 'dkfj') + header = header._replace(durable = True) + header = header._replace(priority = 4) + + message = Message(value="test", properties=properties, header=header) + message = message._replace(footer = {'a':2}) + # message = message._replace(state = constants.MessageState.ReceivedSettled) + + message_batch = BatchMessage(message) + message_batch_copy = copy.deepcopy(message_batch) + batch_message = list(message_batch)[0] + batch_copy_message = list(message_batch_copy)[0] + assert len(list(message_batch)) == len(list(message_batch_copy)) + + # check message attributes are equal to deepcopied message attributes + assert batch_message.footer == batch_copy_message.footer + assert batch_message.application_properties == batch_copy_message.application_properties + assert batch_message.delivery_annotations == batch_copy_message.delivery_annotations + # assert batch_message.settled == batch_copy_message.settled + assert batch_message.properties.message_id == batch_copy_message.properties.message_id + assert batch_message.properties.user_id == batch_copy_message.properties.user_id + assert batch_message.properties.to == batch_copy_message.properties.to + assert batch_message.properties.subject == batch_copy_message.properties.subject + assert batch_message.properties.reply_to == batch_copy_message.properties.reply_to + assert batch_message.properties.correlation_id == batch_copy_message.properties.correlation_id + assert batch_message.properties.content_type == batch_copy_message.properties.content_type + assert batch_message.properties.content_encoding == batch_copy_message.properties.content_encoding + assert batch_message.properties.absolute_expiry_time == batch_copy_message.properties.absolute_expiry_time + assert batch_message.properties.creation_time == batch_copy_message.properties.creation_time + assert batch_message.properties.group_id == batch_copy_message.properties.group_id + assert batch_message.properties.group_sequence == batch_copy_message.properties.group_sequence + assert batch_message.properties.reply_to_group_id == batch_copy_message.properties.reply_to_group_id + assert batch_message.header.delivery_count == batch_copy_message.header.delivery_count + assert batch_message.header.ttl == batch_copy_message.header.ttl + assert batch_message.header.first_acquirer == batch_copy_message.header.first_acquirer + assert batch_message.header.durable == batch_copy_message.header.durable + assert batch_message.header.priority == batch_copy_message.header.priority + +def test_message_auto_body_type(): + single_data = b'!@#$%^&*()_+1234567890' + single_data_message = Message(data=single_data) + check_list = [data for data in single_data_message.data] + assert len(check_list) == 22 + assert(str(single_data_message)) + + multiple_data = [b'!@#$%^&*()_+1234567890', 'abcdefg~123'] + multiple_data_message = Message(data=multiple_data) + check_list = [data for data in multiple_data_message.data] + assert len(check_list) == 2 + assert check_list[0] == multiple_data[0] + assert check_list[1] == multiple_data[1] + assert (str(multiple_data_message)) + + list_mixed_body = [b'!@#$%^&*()_+1234567890', 'abcdefg~123', False, 1.23] + list_mixed_body_message = Message(data=list_mixed_body) + check_data = list_mixed_body_message.data + assert isinstance(check_data, list) + assert len(check_data) == 4 + assert check_data[0] == list_mixed_body[0] + assert check_data[1] == list_mixed_body[1] + assert check_data[2] == list_mixed_body[2] + assert check_data[3] == list_mixed_body[3] + assert (str(list_mixed_body_message)) + + dic_mixed_body = {b'key1': b'value', b'key2': False, b'key3': -1.23} + dic_mixed_body_message = Message(data=dic_mixed_body) + check_data = dic_mixed_body_message.data + assert isinstance(check_data, dict) + assert len(check_data) == 3 + assert check_data == dic_mixed_body + assert (str(dic_mixed_body_message)) diff --git a/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/unittest/test_websocket_exception.py b/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/unittest/test_websocket_exception.py new file mode 100644 index 000000000000..cfbea0a4656b --- /dev/null +++ b/sdk/eventhub/azure-eventhub/tests/pyamqp_tests/unittest/test_websocket_exception.py @@ -0,0 +1,14 @@ +import pytest +import asyncio +from unittest.mock import patch + +import aiohttp +from azure.eventhub._pyamqp.aio._transport_async import WebSocketTransportAsync + + +# class WebsocketException(unittest.TestCase): +async def test_websocket_aiohttp_exception(): + with patch.object(aiohttp.ClientSession,'ws_connect', side_effect=aiohttp.ClientOSError): + transport = WebSocketTransportAsync(host="my_host") + with pytest.raises(ConnectionError): + await transport.connect() diff --git a/sdk/eventhub/azure-eventhub/tests/scripts/pyamqp/async/async_receive.py b/sdk/eventhub/azure-eventhub/tests/scripts/pyamqp/async/async_receive.py new file mode 100644 index 000000000000..e2857df8dd7c --- /dev/null +++ b/sdk/eventhub/azure-eventhub/tests/scripts/pyamqp/async/async_receive.py @@ -0,0 +1,209 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +import asyncio +import os +import dotenv +import logging +from logging.handlers import RotatingFileHandler +import time + +from azure.eventhub.aio import EventHubConsumerClient +from azure.eventhub import parse_connection_string + +logger = logging.getLogger('ASYNC_RECEIVE_PERF_TEST') +logger.setLevel(logging.INFO) +logger.addHandler(RotatingFileHandler("async_receive_perf_test.log")) + +dotenv.load_dotenv() +CONN_STRS = [ + os.environ["EVENT_HUB_CONN_STR_BASIC_NORTHEU"], + os.environ["EVENT_HUB_CONN_STR_STANDARD_NORTHEU"], + os.environ["EVENT_HUB_CONN_STR_BASIC_WESTUS2"], + os.environ["EVENT_HUB_CONN_STR_STANDARD_WESTUS2"] +] +EH_NAME_EVENT_SIZE_PAIR = [ + ('pyamqp_512', 512), +] + +PREFETCH_LIST = [300, 3000] +PARTITION_ID = "0" +RUN_DURATION = 30 +FIXED_AMOUNT = 100_000 + + +async def receive_fixed_time_interval( + conn_str, + eventhub_name, + single_event_size, + prefetch=300, + batch_receiving=False, + description=None, + run_duration=30, + partition_id="0" +): + consumer_client = EventHubConsumerClient.from_connection_string( + conn_str, + consumer_group="$Default", + eventhub_name=eventhub_name + ) + + last_received_count = [0] + received_count = [0] + run_flag = [True] + all_perf_records = [] + check_interval = 1 + + async def on_event(partition_context, event): + received_count[0] += 1 + + async def on_event_batch(partition_context, events): + received_count[0] += len(events) + + async def monitor(): + while run_flag[0]: + snap = received_count[0] + perf = (snap - last_received_count[0]) / check_interval + last_received_count[0] = snap + all_perf_records.append(perf) + await asyncio.sleep(check_interval) + + target = consumer_client.receive_batch if batch_receiving else consumer_client.receive + kwargs = { + "partition_id": partition_id, + "starting_position": "-1", # "-1" is from the beginning of the partition. + "prefetch": prefetch + } + if batch_receiving: + kwargs["max_batch_size"] = prefetch + kwargs["on_event_batch"] = on_event_batch + else: + kwargs["on_event"] = on_event + + recv_future = asyncio.create_task(target(**kwargs)) + monitor_future = asyncio.create_task(monitor()) + + await asyncio.sleep(run_duration) + await consumer_client.close() + run_flag[0] = False + await recv_future + await monitor_future + + valid_perf_records = all_perf_records[10:] # skip the first 10 records to let the receiving program be stable + avg_perf = sum(valid_perf_records) / len(valid_perf_records) + + logger.info( + "EH Namespace: {}.\nMethod: {}, The average performance is {} events/s, throughput: {} bytes/s.\n" + "Configs are: Single message size: {} bytes, Run duration: {} seconds, Batch: {}.\n" + "Prefetch: {}".format( + parse_connection_string(conn_str).fully_qualified_namespace, + description or "receive_fixed_time_interval", + avg_perf, + avg_perf * single_event_size, + single_event_size, + run_duration, + batch_receiving, + prefetch + ) + ) + + +async def receive_fixed_amount( + conn_str, + eventhub_name, + single_event_size, + prefetch=300, + batch_receiving=False, + description=None, + partition_id="0", + run_times=1, + fixed_amount=100_000 +): + consumer_client = EventHubConsumerClient.from_connection_string( + conn_str, + consumer_group="$Default", + eventhub_name=eventhub_name, + ) + perf_records = [] + received_count = [0] + + async def on_event(partition_context, event): + received_count[0] += 1 + if received_count[0] == fixed_amount: + await consumer_client.close() + + async def on_event_batch(partition_context, events): + received_count[0] += len(events) + if received_count[0] >= fixed_amount: + await consumer_client.close() + + for i in range(run_times): + start_time = time.time() + async with consumer_client: + if batch_receiving: + await consumer_client.receive_batch( + on_event_batch=on_event_batch, + partition_id=partition_id, + starting_position="-1", + max_batch_size=prefetch, + prefetch=prefetch + ) + else: + await consumer_client.receive( + on_event=on_event, + partition_id=partition_id, + starting_position="-1", + prefetch=prefetch + ) + end_time = time.time() + total_time = end_time - start_time + speed = fixed_amount/total_time + perf_records.append(speed) + received_count[0] = 0 + avg_perf = sum(perf_records) / len(perf_records) + + logger.info( + "EH Namespace: {}.\nMethod: {}, The average performance is {} events/s, throughput: {} bytes/s.\n" + "Configs are: Single message size: {} bytes, Total events to receive: {}, Batch: {}.\n" + "Prefetch: {}".format( + parse_connection_string(conn_str).fully_qualified_namespace, + description or "receive_fixed_amount", + avg_perf, + avg_perf * single_event_size, + single_event_size, + fixed_amount, + batch_receiving, + prefetch + ) + ) + + +if __name__ == "__main__": + for conn_str in CONN_STRS: + for eh_name, single_event_size in EH_NAME_EVENT_SIZE_PAIR: + for prefetch in PREFETCH_LIST: + for batch_receiving in [True, False]: + print('------------------- receiving fixed amount -------------------') + asyncio.run( + receive_fixed_amount( + conn_str=conn_str, + eventhub_name=eh_name, + single_event_size=single_event_size, + prefetch=prefetch, + batch_receiving=batch_receiving, + fixed_amount=FIXED_AMOUNT + ) + ) + print('------------------- receiving fixed interval -------------------') + asyncio.run( + receive_fixed_time_interval( + conn_str=conn_str, + eventhub_name=eh_name, + single_event_size=single_event_size, + prefetch=prefetch, + batch_receiving=batch_receiving, + run_duration=RUN_DURATION + ) + ) diff --git a/sdk/eventhub/azure-eventhub/tests/scripts/pyamqp/async/async_send.py b/sdk/eventhub/azure-eventhub/tests/scripts/pyamqp/async/async_send.py new file mode 100644 index 000000000000..0dc9ef7b3f85 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/tests/scripts/pyamqp/async/async_send.py @@ -0,0 +1,180 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- +import asyncio +import os +import dotenv +import time +import logging +from logging.handlers import RotatingFileHandler + +from azure.eventhub.aio import EventHubProducerClient +from azure.eventhub import EventData + +logger = logging.getLogger('ASYNC_SEND_PERF_TEST') +logger.setLevel(logging.INFO) +logger.addHandler(RotatingFileHandler("async_send_perf_test.log")) + +dotenv.load_dotenv() +CONN_STRS = [ + os.environ["EVENT_HUB_CONN_STR_BASIC_NORTHEU"], + os.environ["EVENT_HUB_CONN_STR_STANDARD_NORTHEU"], + os.environ["EVENT_HUB_CONN_STR_BASIC_WESTUS2"], + os.environ["EVENT_HUB_CONN_STR_STANDARD_WESTUS2"] +] +EVENTHUB_NAME = "pyamqp" + +SINGLE_EVENT_SIZE_LIST = [512] +PARALLEL_COROUTINE_COUNT_LIST = [1] +FIXED_AMOUNT_OF_EVENTS = 100_000 +RUN_DURATION = 30 + + +async def pre_prepare_client(client, data): + await client.create_batch() # precall to retrieve sender link settings + await client.send_batch([EventData(data)]) # precall to set up the sender link + + +async def send_batch_message(conn_str, eventhub_name, num_of_events, single_event_size, run_times=1, description=None): + + client = EventHubProducerClient.from_connection_string( + conn_str=conn_str, eventhub_name=eventhub_name + ) + + data = b'a' * single_event_size + perf_records = [] + await pre_prepare_client(client, data) + + for _ in range(run_times): # run run_times and calculate the avg performance + start_time = time.time() + batch = await client.create_batch() + for _ in range(num_of_events): + try: + batch.add(EventData(data)) + except ValueError: + # Batch full + await client.send_batch(batch) + batch = await client.create_batch() + batch.add(EventData(data)) + await client.send_batch(batch) + + end_time = time.time() + + total_time = end_time - start_time + speed = num_of_events / total_time + perf_records.append(speed) + + await client.close() + avg_perf = round(sum(perf_records) / len(perf_records), 2) + logger.info( + "Method: {}, The average performance is {} events/s, throughput: {} bytes/s, run times: {}.\n" + "Configs are: Num of events: {} events, Single message size: {} bytes.".format( + description or "send_batch_message", + avg_perf, + avg_perf * single_event_size, + run_times, + num_of_events, + single_event_size + ) + ) + return avg_perf + + +async def send_batch_message_worker_coroutine(client, data, run_flag): + total_cnt = 0 + while run_flag[0]: + batch = await client.create_batch() + try: + while True: + event_data = EventData(body=data) + batch.add(event_data) + except ValueError: + await client.send_batch(batch) + total_cnt += len(batch) + return total_cnt + + +async def send_batch_message_in_parallel(conn_str, eventhub_name, single_event_size, parallel_coroutine_count=4, run_times=1, run_duration=30, description=None): + + perf_records = [] + + for _ in range(run_times): + + futures = [] + clients = [ + EventHubProducerClient.from_connection_string( + conn_str=conn_str, eventhub_name=eventhub_name + ) for _ in range(parallel_coroutine_count) + ] + + data = b'a' * single_event_size + + for client in clients: + await pre_prepare_client(client, data) + + run_flag = [True] + for i in range(parallel_coroutine_count): + futures.append(asyncio.create_task( + send_batch_message_worker_coroutine( + clients[i], + data, + run_flag + ) + )) + + await asyncio.sleep(run_duration) + run_flag[0] = False + await asyncio.gather(*futures) + perf_records.append(sum([future.result() for future in futures]) / run_duration) + + for client in clients: + await client.close() + + avg_perf = round(sum(perf_records) / len(perf_records), 2) + + logger.info( + "Method: {}, The average performance is {} events/s, throughput: {} bytes/s, run times: {}.\n" + "Configs are: Single message size: {} bytes, Parallel thread count: {} threads, Run duration: {} seconds.".format( + description or "send_batch_message_in_parallel", + avg_perf, + avg_perf * single_event_size, + run_times, + single_event_size, + parallel_coroutine_count, + run_duration + ) + ) + + +if __name__ == '__main__': + logger.info('------------------- START OF TEST -------------------') + + for conn_str in CONN_STRS: + for single_event_size in SINGLE_EVENT_SIZE_LIST: + print('------------------- sending fixed amount of message -------------------') + asyncio.run( + send_batch_message( + conn_str=conn_str, + eventhub_name=EVENTHUB_NAME, + num_of_events=FIXED_AMOUNT_OF_EVENTS, + single_event_size=single_event_size, + description='sending fixed amount message' + ) + ) + + for parallel_coroutine_count in PARALLEL_COROUTINE_COUNT_LIST: + for single_event_size in SINGLE_EVENT_SIZE_LIST: + print('------------------- multiple coroutines sending messages for a fixed period -------------------') + asyncio.run( + send_batch_message_in_parallel( + conn_str=conn_str, + eventhub_name=EVENTHUB_NAME, + single_event_size=single_event_size, + parallel_coroutine_count=parallel_coroutine_count, + run_duration=RUN_DURATION, + description='multiple coroutine sending messages' + ) + ) + + logger.info('------------------- END OF TEST -------------------') diff --git a/sdk/eventhub/azure-eventhub/tests/scripts/pyamqp/dev_requirements.txt b/sdk/eventhub/azure-eventhub/tests/scripts/pyamqp/dev_requirements.txt new file mode 100644 index 000000000000..3e338bfa253e --- /dev/null +++ b/sdk/eventhub/azure-eventhub/tests/scripts/pyamqp/dev_requirements.txt @@ -0,0 +1 @@ +python-dotenv \ No newline at end of file diff --git a/sdk/eventhub/azure-eventhub/tests/scripts/pyamqp/preload_eventhub_script_for_receiving.py b/sdk/eventhub/azure-eventhub/tests/scripts/pyamqp/preload_eventhub_script_for_receiving.py new file mode 100644 index 000000000000..6e73bba69481 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/tests/scripts/pyamqp/preload_eventhub_script_for_receiving.py @@ -0,0 +1,94 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- +import os +import dotenv +import time +import logging +from concurrent.futures import ThreadPoolExecutor + +from azure.eventhub import EventHubProducerClient, EventData + +logger = logging.getLogger('PRELOAD_EVENTS') +logger.setLevel(logging.INFO) + +dotenv.load_dotenv() +connect_strs = [ + os.environ["EVENT_HUB_CONN_STR_BASIC_NORTHEU"], + os.environ["EVENT_HUB_CONN_STR_STANDARD_NORTHEU"], + os.environ["EVENT_HUB_CONN_STR_BASIC_WESTUS2"], + os.environ["EVENT_HUB_CONN_STR_STANDARD_WESTUS2"] +] + +eh_name_size_pairs = [ + ('pyamqp_512', 512), +] + +EVENT_DATA_COUNT = 2_000_000 +PARTITION_ID = "0" + + +def pre_prepare_client(client, data): + client.create_batch() # precall to retrieve sender link settings + client.send_batch([EventData(data)]) # precall to set up the sender link + + +def send_batch_message(conn_str, eventhub_name, num_of_events, single_event_size, run_times=1, description=None): + client = EventHubProducerClient.from_connection_string( + conn_str=conn_str, eventhub_name=eventhub_name + ) + + data = b'a' * single_event_size + pre_prepare_client(client, data) + + for _ in range(run_times): # run run_times and calculate the avg performance + start_time = time.time() + batch = client.create_batch(partition_id="0") + for _ in range(num_of_events): + try: + batch.add(EventData(data)) + except ValueError: + # Batch full + client.send_batch(batch) + logger.info( + 'Time{}: {} events of size {} sent to eh {}/{}'.format( + time.time(), len(batch), single_event_size, conn_str, eventhub_name + ) + ) + batch = client.create_batch(partition_id=PARTITION_ID) + batch.add(EventData(data)) + + client.send_batch(batch) + logger.info( + 'Finished! Time{}: {} events of size {} sent to eh {}/{}'.format( + time.time(), len(batch), single_event_size, conn_str, eventhub_name + ) + ) + + client.close() + return "Success for {}".format(eventhub_name) + + +if __name__ == '__main__': + logger.info('------------------- START PREPARATION -------------------') + + executor = ThreadPoolExecutor() + + futures = [] + for conn_str in connect_strs: + for eh_name, event_size in eh_name_size_pairs: + futures.append( + executor.submit( + send_batch_message, + conn_str, + eh_name, + EVENT_DATA_COUNT, + event_size + ) + ) + + for future in futures: + print(future.result()) + + logger.info('------------------- END PREPARATION -------------------') diff --git a/sdk/eventhub/azure-eventhub/tests/scripts/pyamqp/receive_perf_test.log b/sdk/eventhub/azure-eventhub/tests/scripts/pyamqp/receive_perf_test.log new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/sdk/eventhub/azure-eventhub/tests/scripts/pyamqp/sync_receive.py b/sdk/eventhub/azure-eventhub/tests/scripts/pyamqp/sync_receive.py new file mode 100644 index 000000000000..67de93c64e25 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/tests/scripts/pyamqp/sync_receive.py @@ -0,0 +1,211 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +import os +import dotenv +import logging +import threading +from logging.handlers import RotatingFileHandler +import time + +from azure.eventhub import EventHubConsumerClient +from azure.eventhub import parse_connection_string + +logger = logging.getLogger('RECEIVE_PERF_TEST') +logger.setLevel(logging.INFO) +logger.addHandler(RotatingFileHandler("receive_perf_test.log")) + +dotenv.load_dotenv() +CONN_STRS = [ + os.environ["EVENT_HUB_CONN_STR_BASIC_NORTHEU"], + os.environ["EVENT_HUB_CONN_STR_STANDARD_NORTHEU"], + os.environ["EVENT_HUB_CONN_STR_BASIC_WESTUS2"], + os.environ["EVENT_HUB_CONN_STR_STANDARD_WESTUS2"] +] +EH_NAME_EVENT_SIZE_PAIR = [ + ('pyamqp_512', 512), +] +PREFETCH_LIST = [300, 3000] +PARTITION_ID = "0" +RUN_DURATION = 30 +FIXED_AMOUNT = 100_000 + + +def receive_fixed_time_interval( + conn_str, + eventhub_name, + single_event_size, + prefetch=300, + batch_receiving=False, + description=None, + run_duration=30, + partition_id="0" +): + + consumer_client = EventHubConsumerClient.from_connection_string( + conn_str, + consumer_group="$Default", + eventhub_name=eventhub_name + ) + + last_received_count = [0] + received_count = [0] + run_flag = [True] + all_perf_records = [] + check_interval = 1 + + def on_event(partition_context, event): + received_count[0] += 1 + + def on_event_batch(partition_context, events): + received_count[0] += len(events) + + def monitor(): + while run_flag[0]: + snap = received_count[0] + perf = (snap - last_received_count[0]) / check_interval + last_received_count[0] = snap + all_perf_records.append(perf) + time.sleep(check_interval) + + target = consumer_client.receive_batch if batch_receiving else consumer_client.receive + kwargs = { + "partition_id": partition_id, + "starting_position": "-1", # "-1" is from the beginning of the partition. + "prefetch": prefetch + } + if batch_receiving: + kwargs["max_batch_size"] = prefetch + kwargs["on_event_batch"] = on_event_batch + else: + kwargs["on_event"] = on_event + + thread = threading.Thread( + target=target, + kwargs=kwargs + ) + + monitor_thread = threading.Thread( + target=monitor + ) + + thread.daemon = True + monitor_thread.daemon = True + + thread.start() + monitor_thread.start() + time.sleep(run_duration) + consumer_client.close() + run_flag[0] = False + + valid_perf_records = all_perf_records[10:] # skip the first 10 records to let the receiving program be stable + avg_perf = sum(valid_perf_records) / len(valid_perf_records) + + logger.info( + "EH Namespace: {}.\nMethod: {}, The average performance is {} events/s, throughput: {} bytes/s.\n" + "Configs are: Single message size: {} bytes, Run duration: {} seconds.\n" + "Prefetch: {}.".format( + parse_connection_string(conn_str).fully_qualified_namespace, + description or "receive_fixed_time_interval", + avg_perf, + avg_perf * single_event_size, + single_event_size, + run_duration, + prefetch + ) + ) + + +def receive_fixed_amount( + conn_str, + eventhub_name, + single_event_size, + prefetch=300, + batch_receiving=False, + description=None, + partition_id="0", + run_times=1, + fixed_amount=100_000 +): + consumer_client = EventHubConsumerClient.from_connection_string( + conn_str, + consumer_group="$Default", + eventhub_name=eventhub_name, + prefetch=prefetch + ) + perf_records = [] + received_count = [0] + + def on_event(partition_context, event): + received_count[0] += 1 + if received_count[0] == fixed_amount: + consumer_client.close() + + def on_event_batch(partition_context, events): + received_count[0] += len(events) + if received_count[0] >= fixed_amount: + consumer_client.close() + + for i in range(run_times): + start_time = time.time() + with consumer_client: + if batch_receiving: + consumer_client.receive_batch( + on_event_batch=on_event_batch, + partition_id=partition_id, + starting_position="-1", + max_batch_size=prefetch, + prefetch=prefetch + ) + else: + consumer_client.receive( + on_event=on_event, + partition_id=partition_id, + starting_position="-1", + prefetch=prefetch + ) + end_time = time.time() + total_time = end_time - start_time + speed = fixed_amount/total_time + perf_records.append(speed) + received_count[0] = 0 + avg_perf = sum(perf_records) / len(perf_records) + + logger.info( + "EH Namespace: {}.\nMethod: {}, The average performance is {} events/s, throughput: {} bytes/s.\n" + "Configs are: Single message size: {} bytes, Total events to receive: {}.\n" + "Prefetch:{}.".format( + parse_connection_string(conn_str).fully_qualified_namespace, + description or "receive_fixed_amount", + avg_perf, + avg_perf * single_event_size, + single_event_size, + fixed_amount, + prefetch + ) + ) + + +if __name__ == "__main__": + for conn_str in CONN_STRS: + for eh_name, single_event_size in EH_NAME_EVENT_SIZE_PAIR: + for prefetch in PREFETCH_LIST: + for batch_receiving in [True, False]: + print('------------------- receiving fixed amount -------------------') + receive_fixed_amount( + conn_str=conn_str, + eventhub_name=eh_name, + single_event_size=single_event_size, + prefetch=prefetch, + batch_receiving=batch_receiving + ) + print('------------------- receiving fixed interval -------------------') + receive_fixed_time_interval( + conn_str=conn_str, + eventhub_name=eh_name, + single_event_size=single_event_size, + prefetch=prefetch, + batch_receiving=batch_receiving + ) diff --git a/sdk/eventhub/azure-eventhub/tests/scripts/pyamqp/sync_send.py b/sdk/eventhub/azure-eventhub/tests/scripts/pyamqp/sync_send.py new file mode 100644 index 000000000000..55d72daeb14e --- /dev/null +++ b/sdk/eventhub/azure-eventhub/tests/scripts/pyamqp/sync_send.py @@ -0,0 +1,177 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +import os +import dotenv +import time +import logging +from logging.handlers import RotatingFileHandler +from concurrent.futures import ThreadPoolExecutor + +from azure.eventhub import EventHubProducerClient, EventData + +logger = logging.getLogger('SEND_PERF_TEST') +logger.setLevel(logging.INFO) +logger.addHandler(RotatingFileHandler("send_perf_test.log")) + +dotenv.load_dotenv() +CONN_STRS = [ + os.environ["EVENT_HUB_CONN_STR_BASIC_NORTHEU"], + os.environ["EVENT_HUB_CONN_STR_STANDARD_NORTHEU"], + os.environ["EVENT_HUB_CONN_STR_BASIC_WESTUS2"], + os.environ["EVENT_HUB_CONN_STR_STANDARD_WESTUS2"] +] +EVENTHUB_NAME = "pyamqp" + +SINGLE_EVENT_SIZE_LIST = [512] +PARALLEL_THREAD_COUNT_LIST = [1] +FIXED_AMOUNT_OF_EVENTS = 100_000 +RUN_DURATION = 30 + + +def pre_prepare_client(client, data): + client.create_batch() # precall to retrieve sender link settings + client.send_batch([EventData(data)]) # precall to set up the sender link + + +def send_batch_message(conn_str, eventhub_name, num_of_events, single_event_size, run_times=1, description=None): + + client = EventHubProducerClient.from_connection_string( + conn_str=conn_str, eventhub_name=eventhub_name + ) + + data = b'a' * single_event_size + perf_records = [] + pre_prepare_client(client, data) + + for _ in range(run_times): # run run_times and calculate the avg performance + start_time = time.time() + batch = client.create_batch() + for _ in range(num_of_events): + try: + batch.add(EventData(data)) + except ValueError: + # Batch full + client.send_batch(batch) + batch = client.create_batch() + batch.add(EventData(data)) + client.send_batch(batch) + + end_time = time.time() + + total_time = end_time - start_time + speed = num_of_events / total_time + perf_records.append(speed) + + client.close() + avg_perf = round(sum(perf_records) / len(perf_records), 2) + logger.info( + "Method: {}, The average performance is {} events/s, throughput: {} bytes/s, run times: {}.\n" + "Configs are: Num of events: {} events, Single message size: {} bytes.".format( + description or "send_batch_message", + avg_perf, + avg_perf * single_event_size, + run_times, + num_of_events, + single_event_size + ) + ) + return avg_perf + + +def send_batch_message_worker_thread(client, data, run_flag): + total_cnt = 0 + while run_flag[0]: + batch = client.create_batch() + try: + while True: + event_data = EventData(body=data) + batch.add(event_data) + except ValueError: + client.send_batch(batch) + total_cnt += len(batch) + return total_cnt + + +def send_batch_message_in_parallel(conn_str, eventhub_name, single_event_size, parallel_thread_count=4, run_times=1, run_duration=30, description=None): + + perf_records = [] + + for _ in range(run_times): + + futures = [] + clients = [ + EventHubProducerClient.from_connection_string( + conn_str=conn_str, eventhub_name=eventhub_name + ) for _ in range(parallel_thread_count) + ] + + data = b'a' * single_event_size + + for client in clients: + pre_prepare_client(client, data) + + with ThreadPoolExecutor(max_workers=parallel_thread_count) as executor: + run_flag = [True] + for i in range(parallel_thread_count): + futures.append( + executor.submit( + send_batch_message_worker_thread, + clients[i], + data, + run_flag + ) + ) + + time.sleep(run_duration) + run_flag[0] = False + perf_records.append(sum([future.result() for future in futures]) / run_duration) + + for client in clients: + client.close() + + avg_perf = round(sum(perf_records) / len(perf_records), 2) + + logger.info( + "Method: {}, The average performance is {} events/s, throughput: {} bytes/s, run times: {}.\n" + "Configs are: Single message size: {} bytes, Parallel thread count: {} threads, Run duration: {} seconds.".format( + description or "send_batch_message_in_parallel", + avg_perf, + avg_perf * single_event_size, + run_times, + single_event_size, + parallel_thread_count, + run_duration + ) + ) + + +if __name__ == '__main__': + logger.info('------------------- START OF TEST -------------------') + + for conn_str in CONN_STRS: + for single_event_size in SINGLE_EVENT_SIZE_LIST: + print('------------------- sending fixed amount of message -------------------') + send_batch_message( + conn_str=conn_str, + eventhub_name=EVENTHUB_NAME, + num_of_events=FIXED_AMOUNT_OF_EVENTS, + single_event_size=single_event_size, + description='sending fixed amount message' + ) + + for parallel_thread_count in PARALLEL_THREAD_COUNT_LIST: + for single_event_size in SINGLE_EVENT_SIZE_LIST: + print('------------------- multiple threads sending messages for a fixed period -------------------') + send_batch_message_in_parallel( + conn_str=conn_str, + eventhub_name=EVENTHUB_NAME, + single_event_size=single_event_size, + parallel_thread_count=parallel_thread_count, + run_duration=RUN_DURATION, + description='multiple threads sending messages' + ) + + logger.info('------------------- END OF TEST -------------------') diff --git a/sdk/eventhub/azure-eventhub/tests/unittest/asynctests/test_client_creation_async.py b/sdk/eventhub/azure-eventhub/tests/unittest/asynctests/test_client_creation_async.py index bb219fb51198..fb070d9a8a95 100644 --- a/sdk/eventhub/azure-eventhub/tests/unittest/asynctests/test_client_creation_async.py +++ b/sdk/eventhub/azure-eventhub/tests/unittest/asynctests/test_client_creation_async.py @@ -123,4 +123,4 @@ def test_custom_certificate_async(): None, connection_verify='D:/local/certfile' ) - assert consumer._config.connection_verify == 'D:/local/certfile' \ No newline at end of file + assert consumer._config.connection_verify == 'D:/local/certfile' diff --git a/sdk/eventhub/azure-eventhub/tests/unittest/asynctests/test_in_memory_checkpointstore.py b/sdk/eventhub/azure-eventhub/tests/unittest/asynctests/test_in_memory_checkpointstore.py index 294c131ef658..8e5503511e03 100644 --- a/sdk/eventhub/azure-eventhub/tests/unittest/asynctests/test_in_memory_checkpointstore.py +++ b/sdk/eventhub/azure-eventhub/tests/unittest/asynctests/test_in_memory_checkpointstore.py @@ -121,5 +121,3 @@ async def test_update_checkpoint(): assert listed_checkpoint[0] == checkpoint assert listed_checkpoint[0]["offset"] == "0" assert listed_checkpoint[0]["sequencenumber"] == 0 - - diff --git a/sdk/eventhub/azure-eventhub/tests/unittest/test_event_data.py b/sdk/eventhub/azure-eventhub/tests/unittest/test_event_data.py index 37a8b9056758..6eebc6f3050e 100644 --- a/sdk/eventhub/azure-eventhub/tests/unittest/test_event_data.py +++ b/sdk/eventhub/azure-eventhub/tests/unittest/test_event_data.py @@ -11,12 +11,14 @@ try: import uamqp from azure.eventhub._transport._uamqp_transport import UamqpTransport -except ImportError: +except (ModuleNotFoundError, ImportError): + uamqp = None UamqpTransport = None - pass +from azure.eventhub._transport._pyamqp_transport import PyamqpTransport +from azure.eventhub._pyamqp.message import Message, Properties, Header from azure.eventhub.amqp import AmqpAnnotatedMessage, AmqpMessageHeader, AmqpMessageProperties + from azure.eventhub import _common -from azure.eventhub._utils import transform_outbound_single_message pytestmark = pytest.mark.skipif(platform.python_implementation() == "PyPy", reason="This is ignored for PyPy") @@ -87,7 +89,23 @@ def test_sys_properties(uamqp_transport): message = uamqp.message.Message(properties=properties) message.annotations = {_common.PROP_OFFSET: "@latest"} else: - pass + properties = Properties( + message_id="message_id", + user_id="user_id", + to="to", + subject="subject", + reply_to="reply_to", + correlation_id="correlation_id", + content_type="content_type", + content_encoding="content_encoding", + absolute_expiry_time=1, + creation_time=1, + group_id="group_id", + group_sequence=1, + reply_to_group_id="reply_to_group_id" + ) + message_annotations = {_common.PROP_OFFSET: "@latest"} + message = Message(properties=properties, message_annotations=message_annotations) ed = EventData._from_message(message) # type: EventData assert ed.system_properties[_common.PROP_OFFSET] == "@latest" @@ -108,14 +126,16 @@ def test_sys_properties(uamqp_transport): def test_event_data_batch(uamqp_transport): if uamqp_transport: - amqp_transport = UamqpTransport() if version.parse(uamqp.__version__) >= version.parse("1.2.8"): - expected_result = 101 + expected_result = 97 else: expected_result = 93 + amqp_transport=UamqpTransport else: - pass - batch = EventDataBatch(max_size_in_bytes=110, partition_key="par") + expected_result = 99 + amqp_transport=PyamqpTransport + + batch = EventDataBatch(max_size_in_bytes=110, partition_key="par", amqp_transport=amqp_transport) batch.add(EventData("A")) assert str(batch) == "EventDataBatch(max_size_in_bytes=110, partition_id=None, partition_key='par', event_count=1)" assert repr(batch) == "EventDataBatch(max_size_in_bytes=110, partition_id=None, partition_key='par', event_count=1)" @@ -126,12 +146,11 @@ def test_event_data_batch(uamqp_transport): batch.add(EventData("A")) - def test_event_data_from_message(uamqp_transport): if uamqp_transport: - amqp_transport = UamqpTransport() + amqp_transport = UamqpTransport else: - pass + amqp_transport = PyamqpTransport annotated_message = AmqpAnnotatedMessage(data_body=b'A') message = amqp_transport.to_outgoing_amqp_message(annotated_message) event = EventData._from_message(message) @@ -154,6 +173,38 @@ def test_amqp_message_str_repr(): assert str(message) == 'A' assert 'AmqpAnnotatedMessage(body=A, body_type=data' in repr(message) +def test_outgoing_amqp_message_header_properties(uamqp_transport): + if uamqp_transport: + amqp_transport = UamqpTransport + else: + amqp_transport = PyamqpTransport + ann_message = AmqpAnnotatedMessage(data_body=b'A') + ann_message.header = AmqpMessageHeader() + ann_message.properties = AmqpMessageProperties() + amqp_message = amqp_transport.to_outgoing_amqp_message(ann_message) + + assert not amqp_message.header + assert not amqp_message.properties + + ann_message = AmqpAnnotatedMessage(data_body=b'A') + ann_message.header = AmqpMessageHeader() + ann_message.properties = AmqpMessageProperties() + ann_message.header.first_acquirer = False + ann_message.properties.creation_time = 0 + amqp_message = amqp_transport.to_outgoing_amqp_message(ann_message) + + assert amqp_message.header + assert amqp_message.properties + + ann_message = AmqpAnnotatedMessage(data_body=b'A') + ann_message.header = AmqpMessageHeader() + ann_message.properties = AmqpMessageProperties() + ann_message.properties.message_id = "" + amqp_message = amqp_transport.to_outgoing_amqp_message(ann_message) + + assert not amqp_message.header + assert amqp_message.properties + def test_amqp_message_from_message(uamqp_transport): if uamqp_transport: @@ -180,7 +231,30 @@ def test_amqp_message_from_message(uamqp_transport): message = uamqp.message.Message(header=header, properties=properties) message.annotations = {_common.PROP_OFFSET: "@latest"} else: - pass + header = Header( + delivery_count=1, + ttl=10000, + first_acquirer=True, + durable=True, + priority=1 + ) + properties = Properties( + message_id="message_id", + user_id="user_id", + to="to", + subject="subject", + reply_to="reply_to", + correlation_id="correlation_id", + content_type="content_type", + content_encoding="content_encoding", + absolute_expiry_time=1, + creation_time=1, + group_id="group_id", + group_sequence=1, + reply_to_group_id="reply_to_group_id" + ) + message_annotations = {_common.PROP_OFFSET: "@latest"} + message = Message(properties=properties, header=header, message_annotations=message_annotations) amqp_message = AmqpAnnotatedMessage(message=message) assert amqp_message.properties.message_id == message.properties.message_id @@ -201,3 +275,92 @@ def test_amqp_message_from_message(uamqp_transport): assert amqp_message.header.durable == message.header.durable assert amqp_message.header.priority == message.header.priority assert amqp_message.annotations == message.message_annotations + +def test_legacy_message(uamqp_transport): + if uamqp_transport: + header = uamqp.message.MessageHeader() + header.delivery_count = 1 + header.time_to_live = 10000 + header.first_acquirer = True + header.durable = True + header.priority = 1 + properties = uamqp.message.MessageProperties() + properties.message_id = "message_id" + properties.user_id = "user_id" + properties.to = "to" + properties.subject = "subject" + properties.reply_to = "reply_to" + properties.correlation_id = "correlation_id" + properties.content_type = "content_type" + properties.content_encoding = "content_encoding" + properties.absolute_expiry_time = 1 + properties.creation_time = 1 + properties.group_id = "group_id" + properties.group_sequence = 1 + properties.reply_to_group_id = "reply_to_group_id" + message = uamqp.message.Message(body=b'abc', header=header, properties=properties) + message.annotations = {_common.PROP_OFFSET: "@latest"} + amqp_transport = UamqpTransport + else: + header = Header( + delivery_count=1, + ttl=10000, + first_acquirer=True, + durable=True, + priority=1 + ) + properties = Properties( + message_id="message_id", + user_id="user_id", + to="to", + subject="subject", + reply_to="reply_to", + correlation_id="correlation_id", + content_type="content_type", + content_encoding="content_encoding", + absolute_expiry_time=1, + creation_time=1, + group_id="group_id", + group_sequence=1, + reply_to_group_id="reply_to_group_id" + ) + message_annotations = {_common.PROP_OFFSET: "@latest"} + message = Message(data=b'abc', properties=properties, header=header, message_annotations=message_annotations) + amqp_transport = PyamqpTransport + event_data = EventData._from_message(message=message) + assert event_data.message.properties.user_id == b'user_id' + assert event_data.message.properties.message_id == b'message_id' + assert event_data.message.properties.to == b'to' + assert event_data.message.properties.subject == b'subject' + assert event_data.message.properties.reply_to == b"reply_to" + assert event_data.message.properties.correlation_id == b"correlation_id" + assert event_data.message.properties.content_type == b"content_type" + assert event_data.message.properties.content_encoding == b"content_encoding" + assert event_data.message.properties.absolute_expiry_time == 1 + assert event_data.message.properties.creation_time == 1 + assert event_data.message.properties.group_id == b"group_id" + assert event_data.message.properties.group_sequence == 1 + assert event_data.message.properties.reply_to_group_id == b"reply_to_group_id" + assert event_data.message.state.value == 2 + assert event_data.message.delivery_annotations == {} + assert event_data.message.delivery_no is None + assert event_data.message.delivery_tag is None + assert event_data.message.on_send_complete is None + assert event_data.message.footer == {} + assert event_data.message.retries == 0 + assert event_data.message.idle_time == 0 + + event_data_batch = EventDataBatch(partition_key=b'par', partition_id='1', amqp_transport=amqp_transport) + event_data_batch.add(event_data) + assert event_data_batch.message.max_message_length == 1024 * 1024 + assert event_data_batch.message.size_offset == 0 + assert event_data_batch.message.batch_format == 0x80013700 + assert len(event_data_batch.message.annotations) == 1 + assert event_data_batch.message.application_properties is None + assert event_data_batch.message.header.delivery_count == 0 + assert event_data_batch.message.header.time_to_live is None + assert event_data_batch.message.header.first_acquirer is None + assert event_data_batch.message.header.durable is True + assert event_data_batch.message.header.priority is None + assert event_data_batch.message.on_send_complete is None + assert event_data_batch.message.properties is None diff --git a/sdk/eventhub/tests.yml b/sdk/eventhub/tests.yml index 90381cb9ce64..2373f37b7ff5 100644 --- a/sdk/eventhub/tests.yml +++ b/sdk/eventhub/tests.yml @@ -4,6 +4,7 @@ stages: - template: ../../eng/pipelines/templates/stages/archetype-sdk-tests.yml parameters: ServiceDirectory: eventhub + TestTimeoutInMinutes: 240 BuildTargetingString: azure-eventhub* MatrixReplace: - TestSamples=.*/true diff --git a/shared_requirements.txt b/shared_requirements.txt index ead962dd65fc..5c3d9be7a25f 100644 --- a/shared_requirements.txt +++ b/shared_requirements.txt @@ -233,7 +233,6 @@ opentelemetry-sdk<2.0.0,>=1.5.0,!=1.10a0 #override azure-eventhub-checkpointstoreblob-aio azure-core<2.0.0,>=1.20.1 #override azure-eventhub-checkpointstoreblob-aio aiohttp<4.0,>=3.8.3 #override azure-eventhub-checkpointstoretable azure-core<2.0.0,>=1.14.0 -#override azure-eventhub uamqp>=1.6.3,<2.0.0 #override azure-appconfiguration azure-core<2.0.0,>=1.24.0 #override azure-mgmt-maintenance msrest>=0.7.1 #override azure-appconfiguration-provider azure-appconfiguration<2.0.0,>=1.3.0