diff --git a/eng/pipelines/templates/steps/build-artifacts.yml b/eng/pipelines/templates/steps/build-artifacts.yml index 2b8d5b625acb..3381e8534356 100644 --- a/eng/pipelines/templates/steps/build-artifacts.yml +++ b/eng/pipelines/templates/steps/build-artifacts.yml @@ -90,9 +90,9 @@ steps: Write-Host "##vso[task.setvariable variable=PIP_INDEX_URL]https://pypi.python.org/simple" displayName: Reset PIP Index For APIStubGen - #- template: /eng/pipelines/templates/steps/run_apistub.yml - # parameters: - # ServiceDirectory: ${{ parameters.ServiceDirectory }} + - template: /eng/pipelines/templates/steps/run_apistub.yml + parameters: + ServiceDirectory: ${{ parameters.ServiceDirectory }} - ${{ parameters.BeforePublishSteps }} diff --git a/eng/tox/allowed_pylint_failures.py b/eng/tox/allowed_pylint_failures.py index 8ff5355b2756..a67e0bfd0f31 100644 --- a/eng/tox/allowed_pylint_failures.py +++ b/eng/tox/allowed_pylint_failures.py @@ -60,7 +60,5 @@ "azure-agrifood-farming", "azure-ai-language-questionanswering", "azure-ai-language-conversations", - "azure-eventhub", - "azure-eventhub-checkpointstoreblob-aio", "azure-developer-loadtesting" ] diff --git a/eng/tox/mypy_hard_failure_packages.py b/eng/tox/mypy_hard_failure_packages.py index 19ba31d80094..5a8c9c54cb30 100644 --- a/eng/tox/mypy_hard_failure_packages.py +++ b/eng/tox/mypy_hard_failure_packages.py @@ -7,6 +7,7 @@ MYPY_HARD_FAILURE_OPTED = [ "azure-core", + "azure-eventhub", "azure-identity", "azure-keyvault-administration", "azure-keyvault-certificates", diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_client_base.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_client_base.py index 1a39889b52f1..540b80a95895 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,10 @@ from azure.core.utils import parse_connection_string as core_parse_connection_string from azure.core.pipeline.policies import RetryMode - try: from ._transport._uamqp_transport import UamqpTransport except ImportError: - UamqpTransport = None + UamqpTransport = None # type: ignore from ._transport._pyamqp_transport import PyamqpTransport from .exceptions import ClientClosedError from ._configuration import Configuration diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_common.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_common.py index 3b5c88f927e5..631a0b4e1aff 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_common.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_common.py @@ -59,7 +59,10 @@ if TYPE_CHECKING: try: - from uamqp import Message as uamqp_Message, BatchMessage + from uamqp import ( # pylint: disable=unused-import + Message as uamqp_Message, + BatchMessage, + ) except ImportError: uamqp_Message = None BatchMessage = None @@ -131,8 +134,8 @@ def __init__( self._raw_amqp_message = AmqpAnnotatedMessage( # type: ignore data_body=body, annotations={}, application_properties={} ) - self._uamqp_message = None - self._message = None + 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 @@ -231,7 +234,11 @@ def _from_message( event_data = cls(body="") # pylint: disable=protected-access event_data._message = message - event_data._raw_amqp_message = raw_amqp_message if raw_amqp_message else AmqpAnnotatedMessage(message=message) + event_data._raw_amqp_message = ( + raw_amqp_message + if raw_amqp_message + else AmqpAnnotatedMessage(message=message) + ) return event_data def _decode_non_data_body_as_str(self, encoding: str = "UTF-8") -> str: @@ -252,7 +259,10 @@ def message(self) -> LegacyMessage: :rtype: LegacyMessage """ - warnings.warn("The `message` property is deprecated and will be removed in future versions.", DeprecationWarning) + 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, @@ -263,9 +273,12 @@ def message(self) -> LegacyMessage: @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. + 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) + warnings.warn( + "The `message` property is deprecated and will be removed in future versions.", + DeprecationWarning, + ) self._uamqp_message = value @property @@ -534,14 +547,20 @@ def __init__( 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 = self._amqp_transport.set_message_partition_key( + self._message, self._partition_key + ) 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.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._uamqp_message = ( - None if PyamqpTransport.TIMEOUT_FACTOR == self._amqp_transport.TIMEOUT_FACTOR else self._message + None + if PyamqpTransport.TIMEOUT_FACTOR == self._amqp_transport.TIMEOUT_FACTOR + else self._message ) def __repr__(self) -> str: @@ -557,7 +576,7 @@ 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, *, @@ -565,7 +584,10 @@ def _from_batch( partition_id: Optional[str] = None, ) -> EventDataBatch: outgoing_batch_data = [ - transform_outbound_single_message(m, EventData, amqp_transport.to_outgoing_amqp_message) for m in batch_data + transform_outbound_single_message( + m, EventData, amqp_transport.to_outgoing_amqp_message + ) + for m in batch_data ] batch_data_instance = cls( partition_key=partition_key, @@ -596,20 +618,27 @@ def message(self) -> Union["BatchMessage", LegacyBatchMessage]: :rtype: uamqp.BatchMessage or LegacyBatchMessage """ - warnings.warn("The `message` property is deprecated and will be removed in future versions.", DeprecationWarning) + 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 + 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. + 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) + warnings.warn( + "The `message` property is deprecated and will be removed in future versions.", + DeprecationWarning, + ) self._uamqp_message = value @property @@ -638,10 +667,15 @@ def add(self, event_data: Union[EventData, AmqpAnnotatedMessage]) -> None: ) if self._partition_key: - if outgoing_event_data.partition_key and outgoing_event_data.partition_key != self._partition_key: - raise ValueError("The partition key of event_data does not match the partition key of this batch.") + if ( + outgoing_event_data.partition_key + and outgoing_event_data.partition_key != self._partition_key + ): + raise ValueError( + "The partition key of event_data does not match the partition key of this batch." + ) if not outgoing_event_data.partition_key: - outgoing_event_data._message = self._amqp_transport.set_message_partition_key( # pylint: disable=protected-access + 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, ) @@ -653,11 +687,15 @@ def add(self, event_data: Union[EventData, AmqpAnnotatedMessage]) -> None: # For a BatchMessage, if the encoded_message_size of event_data is < 256, then the overhead cost to encode that # message into the BatchMessage would be 5 bytes, if >= 256, it would be 8 bytes. size_after_add = ( - self._size + event_data_size + _BATCH_MESSAGE_OVERHEAD_COST[0 if (event_data_size < 256) else 1] + self._size + + event_data_size + + _BATCH_MESSAGE_OVERHEAD_COST[0 if (event_data_size < 256) else 1] ) if size_after_add > self.max_size_in_bytes: - raise ValueError(f"EventDataBatch has reached its size limit: {self.max_size_in_bytes}") + raise ValueError( + f"EventDataBatch has reached its size limit: {self.max_size_in_bytes}" + ) self._amqp_transport.add_batch(self, outgoing_event_data, event_data) self._size = size_after_add diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_connection_manager.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_connection_manager.py index a00c489568b3..d33aa59376d0 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_connection_manager.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_connection_manager.py @@ -15,6 +15,7 @@ from ._pyamqp._connection import Connection from uamqp.authentication import JWTTokenAuth as uamqp_JWTTokenAuth from uamqp import Connection as uamqp_Connection + from ._transport._base import AmqpTransport try: from typing_extensions import Protocol @@ -60,7 +61,7 @@ def __init__(self, **kwargs): 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") + self._amqp_transport: AmqpTransport = kwargs.pop("amqp_transport") def get_connection( self, diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_consumer.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_consumer.py index d3e88d21e592..4953105ef625 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, Union +from typing import TYPE_CHECKING, Callable, Dict, Optional, Any, Deque, Union, cast from ._common import EventData from ._client_base import ConsumerProducerMixin @@ -20,6 +20,11 @@ ) if TYPE_CHECKING: + 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 @@ -29,10 +34,6 @@ uamqp_Message = None uamqp_AMQPType = None uamqp_JWTTokenAuth = None - from ._pyamqp import types - from ._pyamqp.client import ReceiveClient - from ._pyamqp.message import Message - from ._pyamqp.authentication import JWTTokenAuth from ._consumer_client import EventHubConsumerClient @@ -75,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) @@ -103,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: Union[Dict[uamqp_AMQPType, uamqp_AMQPType], Dict[types.AMQPTypes, types.AMQPTypes]] = {} + 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._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: 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 @@ -147,7 +163,8 @@ def _create_handler(self, auth: Union[uamqp_JWTTokenAuth, 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, @@ -169,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: @@ -180,6 +196,7 @@ 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) @@ -194,9 +211,7 @@ 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. 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 01a8dda668e4..6d0d4cfbdcb7 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 ( @@ -28,14 +29,18 @@ 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_client_async import ( + EventHubConsumerClient as EventHubConsumerClientAsync, + ) class EventProcessorMixin(object): - _eventhub_client = ( - None - ) # type: Optional[EventHubConsumerClient] + _eventhub_client: Optional[Union[EventHubConsumerClient, EventHubConsumerClientAsync]] = None _consumer_group = "" # type: str _owner_level = None # type: Optional[int] _prefetch = None # type: Optional[int] @@ -75,9 +80,9 @@ 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: (...) -> EventHubConsumer + # type: (...) -> Union[EventHubConsumer, EventHubConsumerAsync] consumer = self._eventhub_client._create_consumer( # type: ignore # pylint: disable=protected-access self._consumer_group, partition_id, @@ -87,7 +92,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/_producer.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_producer.py index 990e2d46b622..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 @@ -216,7 +217,7 @@ def _wrap_eventdata( event_data = EventDataBatch._from_batch( event_data._internal_events, amqp_transport=self._amqp_transport, - partition_key=event_data._partition_key, + partition_key=cast(AnyStr, event_data._partition_key), partition_id=event_data._partition_id, max_size_in_bytes=event_data.max_size_in_bytes ) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_producer_client.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_producer_client.py index 2f4a461f3877..656b68042957 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_producer_client.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_producer_client.py @@ -250,7 +250,7 @@ def _buffered_send(self, events, **kwargs): max_wait_time=self._max_wait_time, max_buffer_length=self._max_buffer_length, executor=self._executor, - amqp_transport=self._amqp_transport + amqp_transport=self._amqp_transport, ) self._buffered_producer_dispatcher.enqueue_events(events, **kwargs) @@ -269,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, ) @@ -674,7 +674,10 @@ def send_batch(self, event_data_batch, **kwargs): self._on_success(batch._internal_events, pid) except (KeyError, AttributeError, EventHubError) as e: _LOGGER.debug( - "Producer for partition ID '{}' not available: {}. Rebuilding new producer.".format(partition_id, e)) + "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 @@ -734,7 +737,7 @@ def create_batch(self, **kwargs): max_size_in_bytes=(max_size_in_bytes or self._max_message_size_on_link), partition_id=partition_id, partition_key=partition_key, - amqp_transport=self._amqp_transport + amqp_transport=self._amqp_transport, ) return event_data_batch diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_connection.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_connection.py index 53d2fb64b37b..c494eb5c3d13 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_connection.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_connection.py @@ -10,6 +10,7 @@ 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 @@ -102,7 +103,7 @@ def __init__(self, endpoint, **kwargs): # pylint:disable=too-many-statements if custom_endpoint_address: custom_parsed_url = urlparse(custom_endpoint_address) custom_port = custom_parsed_url.port or WEBSOCKET_PORT - custom_endpoint = "{}:{}{}".format(custom_parsed_url.hostname, custom_port, custom_parsed_url.path) + custom_endpoint = f"{custom_parsed_url.hostname}:{custom_port}{custom_parsed_url.path}" transport = kwargs.get("transport") self._transport_type = kwargs.pop("transport_type", TransportType.Amqp) @@ -179,7 +180,8 @@ def _connect(self): self._outgoing_header() self._set_state(ConnectionState.HDR_SENT) if not self._allow_pipelined_open: - self._process_incoming_frame(*self._read_frame(wait=True)) + # TODO: List/tuple expected as variable args + self._process_incoming_frame(*self._read_frame(wait=True)) # type: ignore if self.state != ConnectionState.HDR_EXCH: self._disconnect() raise ValueError("Did not receive reciprocal protocol header. Disconnecting.") @@ -207,8 +209,9 @@ def _can_read(self): """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=True, **kwargs): - # type: (bool, Any) -> Tuple[int, Optional[Tuple[int, NamedTuple]]] + def _read_frame( # type: ignore # TODO: missing return + self, wait: Union[bool, float] = True, **kwargs: Any + ) -> Tuple[int, Optional[Tuple[int, NamedTuple]]]: """Read an incoming frame from the transport. :param Union[bool, float] wait: Whether to block on the socket while waiting for an incoming frame. @@ -380,19 +383,26 @@ def _incoming_open(self, channel, frame): 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 + self._remote_idle_timeout_send_frame = ( + self._idle_timeout_empty_frame_send_ratio * self._remote_idle_timeout + ) - if frame[2] < 512: # Ensure minimum max frame size. - # Close with error - # Codes_S_R_S_CONNECTION_01_143: [If any of the values in the received open frame are invalid then the connection shall be closed.] - # Codes_S_R_S_CONNECTION_01_220: [The error amqp:invalid-field shall be set in the error.condition field of the CLOSE frame.] + 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=AMQPConnectionError( - condition=ErrorCondition.InvalidField, - description="connection_endpoint_frame_received::failed parsing OPEN frame", + error=cast( + AMQPError, + AMQPConnectionError( + condition=ErrorCondition.InvalidField, + description="Failed parsing OPEN frame: Max frame size is less than supported minimum.", + ), ) ) - _LOGGER.error("connection_endpoint_frame_received::failed parsing OPEN frame") + _LOGGER.error( + "Failed parsing OPEN frame: Max frame size is less than supported minimum." + ) else: self._remote_max_frame_size = frame[2] if self.state == ConnectionState.OPEN_SENT: @@ -402,9 +412,12 @@ def _incoming_open(self, channel, frame): 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}")) + 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) def _outgoing_close(self, error=None): @@ -449,8 +462,13 @@ def _incoming_close(self, channel, frame): self._set_state(ConnectionState.END) if frame[0]: - self._error = AMQPConnectionError(condition=frame[0][0], description=frame[0][1], info=frame[0][2]) - _LOGGER.error("Connection error: {}".format(frame[0])) # pylint:disable=logging-format-interpolation + self._error = AMQPConnectionError( + condition=frame[0][0], description=frame[0][1], info=frame[0][2] + ) + _LOGGER.error( + "Connection error: {}".format(frame[0]) # pylint:disable=logging-format-interpolation + ) + def _incoming_begin(self, channel, frame): # type: (int, Tuple[Any, ...]) -> None @@ -475,9 +493,11 @@ def _incoming_begin(self, channel, frame): try: existing_session = self._outgoing_endpoints[frame[0]] self._incoming_endpoints[channel] = existing_session - self._incoming_endpoints[channel]._incoming_begin(frame) # pylint:disable=protected-access + self._incoming_endpoints[channel]._incoming_begin( # pylint:disable=protected-access + frame + ) except KeyError: - new_session = Session.from_incoming_frame(self, channel, frame) + new_session = Session.from_incoming_frame(self, channel) self._incoming_endpoints[channel] = new_session new_session._incoming_begin(frame) # pylint:disable=protected-access @@ -521,25 +541,36 @@ def _process_incoming_frame(self, channel, frame): # pylint:disable=too-many-re should be interrupted. """ try: - performative, fields = frame + 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(fields) # pylint:disable=protected-access + self._incoming_endpoints[channel]._incoming_transfer( # pylint:disable=protected-access + fields + ) return False if performative == 21: - self._incoming_endpoints[channel]._incoming_disposition(fields) # pylint:disable=protected-access + self._incoming_endpoints[channel]._incoming_disposition( # pylint:disable=protected-access + fields + ) return False if performative == 19: - self._incoming_endpoints[channel]._incoming_flow(fields) # pylint:disable=protected-access + self._incoming_endpoints[channel]._incoming_flow( # pylint:disable=protected-access + fields + ) return False if performative == 18: - self._incoming_endpoints[channel]._incoming_attach(fields) # pylint:disable=protected-access + self._incoming_endpoints[channel]._incoming_attach( # pylint:disable=protected-access + fields + ) return False if performative == 22: - self._incoming_endpoints[channel]._incoming_detach(fields) # pylint:disable=protected-access + self._incoming_endpoints[channel]._incoming_detach( # pylint:disable=protected-access + fields + ) return True if performative == 17: self._incoming_begin(channel, fields) @@ -554,15 +585,12 @@ def _process_incoming_frame(self, channel, frame): # pylint:disable=too-many-re self._incoming_close(channel, fields) return True if performative == 0: - self._incoming_header(channel, fields) - return True - if performative == 1: # pylint:disable=no-else-return - return False # TODO: incoming EMPTY - else: - _LOGGER.error( - "Unrecognized incoming frame: {}".format(frame) - ) # pylint:disable=logging-format-interpolation + self._incoming_header(channel, cast(bytes, fields)) return True + if performative == 1: + return False + _LOGGER.error("Unrecognized incoming frame: %s", frame) + return True except KeyError: return True # TODO: channel error @@ -572,12 +600,23 @@ def _process_outgoing_frame(self, channel, frame): :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]: + 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]: + 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, self._idle_timeout, self._last_frame_received_time) or self._get_remote_timeout(now): + if get_local_timeout( + now, + cast(float, self._idle_timeout), + cast(float, self._last_frame_received_time), + ) or self._get_remote_timeout(now): self.close( error=AMQPError( condition=ErrorCondition.ConnectionCloseForced, @@ -602,7 +641,7 @@ def _get_remote_timeout(self, now): """ 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 > self._remote_idle_timeout_send_frame: + if time_since_last_sent > cast(int, self._remote_idle_timeout_send_frame): self._outgoing_empty() return False @@ -651,7 +690,13 @@ def listen(self, wait=False, batch=1, **kwargs): try: if self.state not in _CLOSING_STATES: now = time.time() - if get_local_timeout(now, self._idle_timeout, self._last_frame_received_time) or self._get_remote_timeout(now): # pylint:disable=line-too-long + if get_local_timeout( + now, + cast(float, self._idle_timeout), + cast(float, self._last_frame_received_time), + ) or self._get_remote_timeout( + now + ): self.close( error=AMQPError( condition=ErrorCondition.ConnectionCloseForced, @@ -707,7 +752,7 @@ def create_session(self, **kwargs): assigned_channel, network_trace=kwargs.pop("network_trace", self._network_trace), network_trace_params=dict(self._network_trace_params), - **kwargs + **kwargs, ) self._outgoing_endpoints[assigned_channel] = session return session @@ -731,7 +776,10 @@ def open(self, wait=False): 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.") + 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 @@ -743,13 +791,19 @@ def close(self, error=None, wait=False): :param bool wait: Whether to wait for a service Close response. Default is `False`. :rtype: None """ - if self.state in [ConnectionState.END, ConnectionState.CLOSE_SENT, ConnectionState.DISCARDING]: + if self.state in [ + ConnectionState.END, + ConnectionState.CLOSE_SENT, + ConnectionState.DISCARDING, + ]: return try: self._outgoing_close(error=error) if error: self._error = AMQPConnectionError( - condition=error.condition, description=error.description, info=error.info + condition=error.condition, + description=error.description, + info=error.info, ) if self.state == ConnectionState.OPEN_PIPE: self._set_state(ConnectionState.OC_PIPE) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_decode.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_decode.py index 53915069be81..099069712865 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_decode.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_decode.py @@ -8,7 +8,7 @@ import struct import uuid import logging -from typing import List, Union, Tuple, Dict, Callable # pylint: disable=unused-import +from typing import List, Optional, Tuple, Dict, Callable, Any, cast, Union # pylint: disable=unused-import from .message import Message, Header, Properties @@ -236,7 +236,7 @@ def _decode_described(buffer): buffer, descriptor = _DECODE_BY_CONSTRUCTOR[composite_type](buffer[1:]) buffer, value = _DECODE_BY_CONSTRUCTOR[buffer[0]](buffer[1:]) try: - composite_type = _COMPOSITES[descriptor] + composite_type = cast(int, _COMPOSITES[descriptor]) return buffer, {composite_type: value} except KeyError: return buffer, value @@ -244,7 +244,7 @@ def _decode_described(buffer): def decode_payload(buffer): # type: (memoryview) -> Message - 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. @@ -262,12 +262,12 @@ def decode_payload(buffer): message["application_properties"] = value elif descriptor == 117: try: - message["data"].append(value) + cast(List, message["data"]).append(value) except KeyError: message["data"] = [value] elif descriptor == 118: try: - message["sequence"].append(value) + cast(List, message["sequence"]).append(value) except KeyError: message["sequence"] = [value] elif descriptor == 119: @@ -293,7 +293,7 @@ def decode_frame(data): # list8 0xc0: data[4] is size, data[5] is count count = data[5] buffer = data[6:] - fields = [None] * count + 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: @@ -302,7 +302,7 @@ def decode_frame(data): def decode_empty_frame(header): - # type: (memory) -> bytes + # type: (memoryview) -> Tuple[int, bytes] if header[0:4] == _HEADER_PREFIX: return 0, header.tobytes() if header[5] == 0: @@ -310,7 +310,7 @@ def decode_empty_frame(header): raise ValueError("Received unrecognized empty frame") -_DECODE_BY_CONSTRUCTOR = [None] * 256 # type: List[Callable[memoryview]] +_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 diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_encode.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_encode.py index 8ecca3adb2cb..24267004c8b1 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_encode.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_encode.py @@ -4,18 +4,68 @@ # 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 # pylint: disable=unused-import +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 .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" @@ -33,7 +83,9 @@ def encode_null(output, *args, **kwargs): # pylint: disable=unused-argument output.extend(ConstructorBytes.null) -def encode_boolean(output, value, with_constructor=True, **kwargs): # pylint: disable=unused-argument +def encode_boolean( + output, value, with_constructor=True, **kwargs # pylint: disable=unused-argument +): # type: (bytearray, bool, bool, Any) -> None """ @@ -50,7 +102,9 @@ def encode_boolean(output, value, with_constructor=True, **kwargs): # pylint: d output.extend(ConstructorBytes.bool_true if value else ConstructorBytes.bool_false) -def encode_ubyte(output, value, with_constructor=True, **kwargs): # pylint: disable=unused-argument +def encode_ubyte( + output, value, with_constructor=True, **kwargs # pylint: disable=unused-argument +): # type: (bytearray, Union[int, bytes], bool, Any) -> None """ @@ -58,6 +112,7 @@ def encode_ubyte(output, value, with_constructor=True, **kwargs): # pylint: dis try: value = int(value) except ValueError: + value = cast(bytes, value) value = ord(value) try: output.extend(_construct(ConstructorBytes.ubyte, with_constructor)) @@ -66,7 +121,9 @@ def encode_ubyte(output, value, with_constructor=True, **kwargs): # pylint: dis raise ValueError("Unsigned byte value must be 0-255") -def encode_ushort(output, value, with_constructor=True, **kwargs): # pylint: disable=unused-argument +def encode_ushort( + output, value, with_constructor=True, **kwargs # pylint: disable=unused-argument +): # type: (bytearray, int, bool, Any) -> None """ @@ -110,10 +167,7 @@ def encode_ulong(output, value, with_constructor=True, use_smallest=True): label="unsigned long value in the range 0 to 255 inclusive"/> """ - try: - value = long(value) - except NameError: - value = int(value) + value = int(value) if value == 0: output.extend(ConstructorBytes.ulong_0) return @@ -128,7 +182,9 @@ def encode_ulong(output, value, with_constructor=True, use_smallest=True): raise ValueError("Value supplied for unsigned long invalid: {}".format(value)) -def encode_byte(output, value, with_constructor=True, **kwargs): # pylint: disable=unused-argument +def encode_byte( + output, value, with_constructor=True, **kwargs # pylint: disable=unused-argument +): # type: (bytearray, int, bool, Any) -> None """ @@ -141,7 +197,9 @@ def encode_byte(output, value, with_constructor=True, **kwargs): # pylint: disa raise ValueError("Byte value must be -128-127") -def encode_short(output, value, with_constructor=True, **kwargs): # pylint: disable=unused-argument +def encode_short( + output, value, with_constructor=True, **kwargs # pylint: disable=unused-argument +): # type: (bytearray, int, bool, Any) -> None """ @@ -179,11 +237,10 @@ def encode_long(output, value, with_constructor=True, use_smallest=True): """ if isinstance(value, datetime): - value = (calendar.timegm(value.utctimetuple()) * 1000) + (value.microsecond / 1000) - try: - value = long(value) - except NameError: - value = int(value) + 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)) @@ -195,7 +252,9 @@ def encode_long(output, value, with_constructor=True, use_smallest=True): raise ValueError("Value supplied for long invalid: {}".format(value)) -def encode_float(output, value, with_constructor=True, **kwargs): # pylint: disable=unused-argument +def encode_float( + output, value, with_constructor=True, **kwargs # pylint: disable=unused-argument +): # type: (bytearray, float, bool, Any) -> None """ @@ -205,7 +264,9 @@ def encode_float(output, value, with_constructor=True, **kwargs): # pylint: dis output.extend(struct.pack(">f", value)) -def encode_double(output, value, with_constructor=True, **kwargs): # pylint: disable=unused-argument +def encode_double( + output, value, with_constructor=True, **kwargs # pylint: disable=unused-argument +): # type: (bytearray, float, bool, Any) -> None """ @@ -215,20 +276,28 @@ def encode_double(output, value, with_constructor=True, **kwargs): # pylint: di output.extend(struct.pack(">d", value)) -def encode_timestamp(output, value, with_constructor=True, **kwargs): # pylint: disable=unused-argument +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 = (calendar.timegm(value.utctimetuple()) * 1000) + (value.microsecond / 1000) - value = int(value) + 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 +def encode_uuid( + output, value, with_constructor=True, **kwargs # pylint: disable=unused-argument +): # type: (bytearray, Union[uuid.UUID, str, bytes], bool, Any) -> None """ @@ -323,7 +392,7 @@ def encode_list(output, value, with_constructor=True, use_smallest=True): """ - count = len(value) + count = len(cast(Sized, value)) if use_smallest and count == 0: output.extend(ConstructorBytes.list_0) return @@ -353,13 +422,14 @@ def encode_map(output, value, with_constructor=True, use_smallest=True): """ - count = len(value) * 2 + count = len(cast(Sized, value)) * 2 encoded_size = 0 encoded_values = bytearray() try: - items = value.items() + value = cast(Dict, value) + items = cast(Iterable, value.items()) except AttributeError: - items = value + 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) @@ -376,7 +446,6 @@ def encode_map(output, value, with_constructor=True, use_smallest=True): except struct.error: raise ValueError("Map is too large or too long to be encoded.") output.extend(encoded_values) - return def _check_element_type(item, element_type): @@ -402,14 +471,16 @@ def encode_array(output, value, with_constructor=True, use_smallest=True): """ - count = len(value) + 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) + encode_value( + encoded_values, item, with_constructor=first_item, use_smallest=False + ) first_item = False if item is None: encoded_size -= 1 @@ -429,8 +500,7 @@ def encode_array(output, value, with_constructor=True, use_smallest=True): output.extend(encoded_values) -def encode_described(output, value, _=None, **kwargs): - # type: (bytearray, Tuple(Any, Any), bool, Any) -> None +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) @@ -450,9 +520,9 @@ def encode_fields(value): return {TYPE: AMQPTypes.null, VALUE: None} fields = {TYPE: AMQPTypes.map, VALUE: []} for key, data in value.items(): - if isinstance(key, six.text_type): - key = key.encode("utf-8") - fields[VALUE].append(({TYPE: AMQPTypes.symbol, VALUE: key}, data)) + if isinstance(key, str): + key = key.encode("utf-8") # type: ignore + cast(List, fields[VALUE]).append(({TYPE: AMQPTypes.symbol, VALUE: key}, data)) return fields @@ -476,9 +546,11 @@ def encode_annotations(value): else: field_key = {TYPE: AMQPTypes.symbol, VALUE: key} try: - fields[VALUE].append((field_key, {TYPE: data[TYPE], VALUE: data[VALUE]})) + cast(List, fields[VALUE]).append( + (field_key, {TYPE: data[TYPE], VALUE: data[VALUE]}) + ) except (KeyError, TypeError): - fields[VALUE].append((field_key, {TYPE: None, VALUE: data})) + cast(List, fields[VALUE]).append((field_key, {TYPE: None, VALUE: data})) return fields @@ -496,9 +568,9 @@ def encode_application_properties(value): """ if not value: return {TYPE: AMQPTypes.null, VALUE: None} - fields = {TYPE: AMQPTypes.map, VALUE: []} + fields = {TYPE: AMQPTypes.map, VALUE: cast(List, [])} for key, data in value.items(): - fields[VALUE].append(({TYPE: AMQPTypes.string, VALUE: key}, data)) + cast(List, fields[VALUE]).append(({TYPE: AMQPTypes.string, VALUE: key}, data)) return fields @@ -573,13 +645,14 @@ def encode_filter_set(value): """ if not value: return {TYPE: AMQPTypes.null, VALUE: None} - fields = {TYPE: AMQPTypes.map, VALUE: []} + 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, six.text_type): - name = name.encode("utf-8") + if isinstance(name, str): + name = name.encode("utf-8") # type: ignore try: descriptor, filter_value = data described_filter = { @@ -589,7 +662,9 @@ def encode_filter_set(value): except ValueError: described_filter = data - fields[VALUE].append(({TYPE: AMQPTypes.symbol, VALUE: name}, described_filter)) + cast(List, fields[VALUE]).append( + ({TYPE: AMQPTypes.symbol, VALUE: name}, described_filter) + ) return fields @@ -617,7 +692,7 @@ def encode_unknown(output, value, **kwargs): elif isinstance(value, list): encode_list(output, value, **kwargs) elif isinstance(value, tuple): - encode_described(output, value, **kwargs) + encode_described(output, cast(Tuple[Any, Any], value), **kwargs) elif isinstance(value, dict): encode_map(output, value, **kwargs) else: @@ -662,36 +737,50 @@ def encode_unknown(output, value, **kwargs): def encode_value(output, value, **kwargs): # type: (bytearray, Any, Any) -> None try: - _ENCODE_MAP[value[TYPE]](output, value[VALUE], **kwargs) + cast(Callable, _ENCODE_MAP[value[TYPE]])(output, value[VALUE], **kwargs) except (KeyError, TypeError): encode_unknown(output, value, **kwargs) def describe_performative(performative): - # type: (Performative) -> Tuple(bytes, bytes) - body = [] + # type: (Performative) -> Dict[str, Sequence[Collection[str]]] + body: List[Dict[str, Any]] = [] for index, value in enumerate(performative): - field = performative._definition[index] # pylint: disable=protected-access + # 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]}) + 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]}) + 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: AMQPTypes.list, VALUE: body}), # pylint: disable=protected-access + VALUE: ( + {TYPE: AMQPTypes.ulong, VALUE: performative._code}, # type: ignore # pylint: disable=protected-access + {TYPE: AMQPTypes.list, VALUE: body}, + ), } @@ -736,7 +825,10 @@ def encode_payload(output, payload): output, { TYPE: AMQPTypes.described, - VALUE: ({TYPE: AMQPTypes.ulong, VALUE: 0x00000074}, encode_application_properties(payload[4])), + VALUE: ( + {TYPE: AMQPTypes.ulong, VALUE: 0x00000074}, + encode_application_properties(payload[4]), + ), }, ) @@ -746,7 +838,10 @@ def encode_payload(output, payload): output, { TYPE: AMQPTypes.described, - VALUE: ({TYPE: AMQPTypes.ulong, VALUE: 0x00000075}, {TYPE: AMQPTypes.binary, VALUE: item_value}), + VALUE: ( + {TYPE: AMQPTypes.ulong, VALUE: 0x00000075}, + {TYPE: AMQPTypes.binary, VALUE: item_value}, + ), }, ) @@ -756,7 +851,10 @@ def encode_payload(output, payload): output, { TYPE: AMQPTypes.described, - VALUE: ({TYPE: AMQPTypes.ulong, VALUE: 0x00000076}, {TYPE: None, VALUE: item_value}), + VALUE: ( + {TYPE: AMQPTypes.ulong, VALUE: 0x00000076}, + {TYPE: None, VALUE: item_value}, + ), }, ) @@ -765,7 +863,10 @@ def encode_payload(output, payload): output, { TYPE: AMQPTypes.described, - VALUE: ({TYPE: AMQPTypes.ulong, VALUE: 0x00000077}, {TYPE: None, VALUE: payload[7]}), + VALUE: ( + {TYPE: AMQPTypes.ulong, VALUE: 0x00000077}, + {TYPE: None, VALUE: payload[7]}, + ), }, ) @@ -802,7 +903,7 @@ def encode_payload(output, payload): def encode_frame(frame, frame_type=_FRAME_TYPE): - # type: (Performative) -> Tuple(bytes, bytes) + # 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 diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_message_backcompat.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_message_backcompat.py index 20f3c595d455..0e3c22213eda 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_message_backcompat.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_message_backcompat.py @@ -5,7 +5,7 @@ # -------------------------------------------------------------------------- # pylint: disable=too-many-lines -from typing import Callable +from typing import Callable, cast from enum import Enum from ._encode import encode_payload @@ -14,6 +14,13 @@ 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 @@ -24,7 +31,7 @@ class MessageState(Enum): def __eq__(self, __o: object) -> bool: try: - return self.value == __o.value + return self.value == cast(Enum, __o).value except AttributeError: return super().__eq__(__o) @@ -38,28 +45,42 @@ class MessageAlreadySettled(Exception): PENDING_STATES = (MessageState.WaitingForSendAck, MessageState.WaitingToBeSent) -class LegacyMessage(object): +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._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.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') + self._to_outgoing_amqp_message: Callable = kwargs.get( + "to_outgoing_amqp_message" + ) def __str__(self): return str(self._message) @@ -102,7 +123,7 @@ def get_message(self): def accept(self): if self._can_settle_message(): - self._settler.settle_messages(self.delivery_no, 'accepted') + self._settler.settle_messages(self.delivery_no, "accepted") self.state = MessageState.ReceivedSettled return True return False @@ -111,12 +132,10 @@ def reject(self, condition=None, description=None, info=None): if self._can_settle_message(): self._settler.settle_messages( self.delivery_no, - 'rejected', + "rejected", error=AMQPError( - condition=condition, - description=description, - info=info - ) + condition=condition, description=description, info=info + ), ) self.state = MessageState.ReceivedSettled return True @@ -124,7 +143,7 @@ def reject(self, condition=None, description=None, info=None): def release(self): if self._can_settle_message(): - self._settler.settle_messages(self.delivery_no, 'released') + self._settler.settle_messages(self.delivery_no, "released") self.state = MessageState.ReceivedSettled return True return False @@ -133,7 +152,7 @@ def modify(self, failed, deliverable, annotations=None): if self._can_settle_message(): self._settler.settle_messages( self.delivery_no, - 'modified', + "modified", delivery_failed=failed, undeliverable_here=deliverable, message_annotations=annotations, @@ -149,22 +168,21 @@ class LegacyBatchMessage(LegacyMessage): size_offset = 0 -class LegacyMessageProperties(object): - +class LegacyMessageProperties(object): # pylint: disable=too-many-instance-attributes def __init__(self, properties): - self.message_id = self._encode_property(properties.message_id) - self.user_id = self._encode_property(properties.user_id) - self.to = self._encode_property(properties.to) - self.subject = self._encode_property(properties.subject) - self.reply_to = self._encode_property(properties.reply_to) - self.correlation_id = self._encode_property(properties.correlation_id) - self.content_type = self._encode_property(properties.content_type) - self.content_encoding = self._encode_property(properties.content_encoding) + 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 = self._encode_property(properties.group_id) + self.group_id = _encode_property(properties.group_id) self.group_sequence = properties.group_sequence - self.reply_to_group_id = self._encode_property(properties.reply_to_group_id) + self.reply_to_group_id = _encode_property(properties.reply_to_group_id) def __str__(self): return str( @@ -185,12 +203,6 @@ def __str__(self): } ) - def _encode_property(self, value): - try: - return value.encode("UTF-8") - except AttributeError: - return value - def get_properties_obj(self): return Properties( self.message_id, @@ -205,12 +217,11 @@ def get_properties_obj(self): self.creation_time, self.group_id, self.group_sequence, - self.reply_to_group_id + 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 @@ -235,5 +246,5 @@ def get_header_obj(self): self.priority, self.time_to_live, self.first_acquirer, - self.delivery_count + 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 index e52153aa20a2..18d91f710041 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_platform.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_platform.py @@ -3,6 +3,7 @@ from __future__ import absolute_import, unicode_literals +from typing import Tuple, cast import platform import re import struct @@ -20,7 +21,7 @@ def _linux_version_to_tuple(s): # type: (str) -> Tuple[int, int, int] - return tuple(map(_versionatom, s.split('.')[:3])) + return cast(Tuple[int, int, int], tuple(map(_versionatom, s.split('.')[:3]))) def _versionatom(s): diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_transport.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_transport.py index f86cb192ddb9..c846a6667bae 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_transport.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_transport.py @@ -1,4 +1,4 @@ -# ------------------------------------------------------------------------- +# ------------------------------------------------------------------------- # 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. # @@ -51,33 +51,35 @@ 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 .constants import ( + TLS_HEADER_FRAME, + WEBSOCKET_PORT, + TransportType, + AMQP_WS_SUBPROTOCOL, +) try: import fcntl except ImportError: # pragma: no cover - fcntl = None # noqa -try: - from os import set_cloexec # Python 3.4? -except ImportError: # pragma: no cover - # TODO: Drop this once we drop Python 2.7 support - 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) + 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__) @@ -123,6 +125,7 @@ def get_errno(exc): 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) @@ -141,20 +144,21 @@ class UnexpectedFrame(Exception): pass -class _AbstractTransport(object): +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, - write_timeout=None, socket_settings=None, raise_on_initial_eintr=True, - **kwargs + **kwargs # pylint: disable=unused-argument ): + self._quick_recv = None self.connected = False self.sock = None self.raise_on_initial_eintr = raise_on_initial_eintr @@ -163,7 +167,6 @@ def __init__( self.connect_timeout = connect_timeout or TIMEOUT_INTERVAL self.read_timeout = read_timeout or READ_TIMEOUT_INTERVAL - self.write_timeout = write_timeout self.socket_settings = socket_settings self.socket_lock = Lock() @@ -176,7 +179,6 @@ def connect(self): self._init_socket( self.socket_settings, self.read_timeout, - self.write_timeout, ) # we've sent the banner; signal connect # EINTR, EAGAIN, EWOULDBLOCK would signal that the banner @@ -283,7 +285,9 @@ def _connect(self, host, port, timeout): 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 = socket.getaddrinfo( + host, port, family, socket.SOCK_STREAM, SOL_TCP + ) entries_num = len(entries) # now that we have address(es) for the hostname, connect to broker for i, res in enumerate(entries): @@ -313,12 +317,12 @@ def _connect(self, host, port, timeout): # 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") + raise e if e is not None else socket.error( + "failed to resolve broker hostname" + ) continue # pragma: no cover - - - def _init_socket(self, socket_settings, read_timeout, write_timeout): + 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) @@ -339,7 +343,7 @@ def _init_socket(self, socket_settings, read_timeout, write_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): + def _get_tcp_socket_defaults(self, sock): # pylint: disable=no-self-use tcp_opts = {} for opt in KNOWN_TCP_OPTS: enum = None @@ -366,7 +370,7 @@ def _set_socket_options(self, socket_settings): for opt, val in tcp_opts.items(): self.sock.setsockopt(SOL_TCP, opt, val) - def _read(self, n, initial=False): + def _read(self, n, initial=False, buffer=None, _errnos=None): """Read exactly n bytes from the peer.""" raise NotImplementedError("Must be overriden in subclass") @@ -388,7 +392,7 @@ def close(self): # calling this method. try: self.sock.shutdown(socket.SHUT_RDWR) - except Exception as exc: + 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.info("Transport endpoint is already disconnected: %r", exc) @@ -396,7 +400,7 @@ def close(self): self.sock = None self.connected = False - def read(self, verify_frame_type=0, **kwargs): + def read(self, verify_frame_type=0): read = self._read read_frame_buffer = BytesIO() try: @@ -410,6 +414,10 @@ def read(self, verify_frame_type=0, **kwargs): 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: + 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 @@ -417,7 +425,9 @@ def read(self, verify_frame_type=0, **kwargs): 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:])) + 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): @@ -446,7 +456,7 @@ def write(self, s): self.connected = False raise - def receive_frame(self, *args, **kwargs): + def receive_frame(self, **kwargs): try: header, channel, payload = self.read(**kwargs) if not payload: @@ -466,17 +476,21 @@ def send_frame(self, channel, frame, **kwargs): data = header + encoded_channel + performative self.write(data) - def negotiate(self, encode, decode): + def negotiate(self): pass class SSLTransport(_AbstractTransport): """Transport that works over SSL.""" - def __init__(self, host, port=AMQPS_PORT, connect_timeout=None, ssl=None, **kwargs): - self.sslopts = ssl if isinstance(ssl, dict) else {} + 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) + super(SSLTransport, self).__init__( + host, port=port, connect_timeout=connect_timeout, **kwargs + ) def _setup_transport(self): """Wrap the socket in an SSL object.""" @@ -489,14 +503,16 @@ def _wrap_socket(self, sock, context=None, **sslopts): return self._wrap_context(sock, sslopts, **context) return self._wrap_socket_sni(sock, **sslopts) - def _wrap_context(self, sock, sslopts, check_hostname=None, **ctx_options): + 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( + def _wrap_socket_sni( # pylint: disable=no-self-use self, sock, keyfile=None, @@ -533,9 +549,14 @@ def _wrap_socket_sni( #'ssl_version': ssl_version } - sock = ssl.wrap_socket(**opts) + # TODO: We need to refactor this. + sock = ssl.wrap_socket(**opts) # pylint: disable=deprecated-method # Set SNI headers if supported - if (server_hostname is not None) and (hasattr(ssl, "HAS_SNI") and ssl.HAS_SNI) and (hasattr(ssl, "SSLContext")): + 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: @@ -553,14 +574,20 @@ def _shutdown_transport(self): except OSError: pass - def _read(self, toread, initial=False, buffer=None, _errnos=(errno.ENOENT, errno.EAGAIN, errno.EINTR)): + 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(toread)) + view = buffer or memoryview(bytearray(n)) nbytes = self._read_buffer.readinto(view) - toread -= nbytes + toread = n - nbytes length += nbytes try: while toread: @@ -607,51 +634,15 @@ def _write(self, s): def negotiate(self): with self.block(): self.write(TLS_HEADER_FRAME) - channel, returned_header = self.receive_frame(verify_frame_type=None) + _, returned_header = self.receive_frame(verify_frame_type=None) if returned_header[1] == TLS_HEADER_FRAME: raise ValueError( - "Mismatching TLS header protocol. Excpected: {}, received: {}".format( - TLS_HEADER_FRAME, returned_header[1] - ) + f"""Mismatching TLS header protocol. Expected: {TLS_HEADER_FRAME!r},""" + """received: {returned_header[1]!r}""" ) -class TCPTransport(_AbstractTransport): - """Transport that deals directly with TCP socket.""" - - def _setup_transport(self): - # Setup to _write() directly to the socket, and - # do our own buffered reads. - self._write = self.sock.sendall - self._read_buffer = EMPTY_BUFFER - self._quick_recv = self.sock.recv - - def _read(self, n, initial=False, _errnos=(errno.EAGAIN, errno.EINTR)): - """Read exactly n bytes from the socket.""" - recv = self._quick_recv - rbuf = self._read_buffer - try: - while len(rbuf) < n: - try: - s = self.sock.read(n - len(rbuf)) - except socket.error as exc: - if exc.errno in _errnos: - if initial and self.raise_on_initial_eintr: - raise socket.timeout() - continue - raise - if not s: - raise IOError("Server unexpectedly closed connection") - rbuf += s - except: # noqa - self._read_buffer = rbuf - raise - - result, self._read_buffer = rbuf[:n], rbuf[n:] - return result - - -def Transport(host, transport_type, connect_timeout=None, ssl=False, **kwargs): +def Transport(host, transport_type, connect_timeout=None, ssl_opts=True, **kwargs): """Create transport. Given a few parameters from the Connection constructor, @@ -660,17 +651,25 @@ def Transport(host, transport_type, connect_timeout=None, ssl=False, **kwargs): if transport_type == TransportType.AmqpOverWebsocket: transport = WebSocketTransport else: - transport = SSLTransport if ssl else TCPTransport - return transport(host, connect_timeout=connect_timeout, ssl=ssl, **kwargs) + 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=None, **kwargs): - self.sslopts = ssl if isinstance(ssl, dict) else {} + 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, connect_timeout, **kwargs) + super().__init__(host, port=port, connect_timeout=connect_timeout, **kwargs) self.ws = None self._http_proxy = kwargs.get("http_proxy", None) @@ -697,9 +696,11 @@ def connect(self): http_proxy_auth=http_proxy_auth, ) except ImportError: - raise ValueError("Please install websocket-client library to use websocket transport.") + raise ValueError( + "Please install websocket-client library to use websocket transport." + ) - def _read(self, n, initial=False, buffer=None, **kwargs): # pylint: disable=unused-arguments + 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 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 index 6348008cb38f..f6b68b277d6d 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_authentication_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_authentication_async.py @@ -18,9 +18,9 @@ async def _generate_sas_token_async(auth_uri, sas_name, sas_key, expiry_in=AUTH_ class JWTTokenAuthAsync(JWTTokenAuth): - """""" # TODO: # 1. naming decision, suffix with Auth vs Credential + ... class SASTokenAuthAsync(SASTokenAuth): 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 index a38fb50ad80d..c4859df0e8ff 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_cbs_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_cbs_async.py @@ -6,7 +6,6 @@ import logging from datetime import datetime -import asyncio from ..utils import utc_now, utc_from_timestamp from ._management_link_async import ManagementLink @@ -22,7 +21,6 @@ CBS_OPERATION, ManagementExecuteOperationResult, ManagementOpenResult, - DEFAULT_AUTH_TIMEOUT, ) from ..cbs import check_put_timeout_status, check_expiration_and_refresh_status @@ -60,9 +58,9 @@ def __init__(self, session, auth, **kwargs): async def _put_token(self, token, token_type, audience, expires_on=None): # type: (str, str, str, datetime) -> None - message = Message( + message = Message( # type: ignore # TODO: missing positional args header, etc. value=token, - properties=Properties(message_id=self._mgmt_link.next_message_id), + properties=Properties(message_id=self._mgmt_link.next_message_id), # type: ignore application_properties={ CBS_NAME: audience, CBS_OPERATION: CBS_PUT_TOKEN, @@ -114,8 +112,12 @@ async def _on_amqp_management_error(self): ) # pylint:disable=protected-access async def _on_execute_operation_complete( - self, execute_operation_result, status_code, status_description, message, error_condition=None - ): # TODO: message and error_condition never used + self, execute_operation_result, status_code, status_description, _, error_condition=None + ): + if error_condition: + _LOGGER.info("CBS Put token error: %r", error_condition) + self.auth_state = CbsAuthState.ERROR + return _LOGGER.info( "CBS Put token result (%r), status code: %s, status_description: %s.", execute_operation_result, @@ -159,8 +161,8 @@ async def _cbs_link_ready(self): return False if self.state in (CbsState.CLOSED, CbsState.ERROR): raise TokenAuthFailure( - condition=ErrorCondition.ClientError, - description="CBS authentication link is in broken status, please recreate the cbs link.", + status_code=ErrorCondition.ClientError, + status_description="CBS authentication link is in broken status, please recreate the cbs link.", ) async def open(self): 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 index 143555c41714..ba6c24ad1125 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_client_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_client_async.py @@ -1,4 +1,4 @@ -#------------------------------------------------------------------------- +#------------------------------------------------------------------------- # 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. @@ -9,7 +9,7 @@ import time import queue from functools import partial -from typing import Any, Dict, Optional, Tuple, Union, overload +from typing import Any, Dict, Optional, Tuple, Union, overload, cast from typing_extensions import Literal import certifi @@ -17,9 +17,12 @@ from ._connection_async import Connection from ._management_operation_async import ManagementOperation from ._cbs_async import CBSAuthenticator -from ..client import AMQPClient as AMQPClientSync -from ..client import ReceiveClient as ReceiveClientSync -from ..client import SendClient as SendClientSync +from ..client import ( + AMQPClient as AMQPClientSync, + ReceiveClient as ReceiveClientSync, + SendClient as SendClientSync, + Outcomes +) from ..message import _MessageDelivery from ..constants import ( MessageDeliveryState, @@ -90,7 +93,8 @@ class AMQPClientAsync(AMQPClientSync): :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] + :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', @@ -104,13 +108,13 @@ class AMQPClientAsync(AMQPClientSync): :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. + :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. + :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. @@ -131,6 +135,21 @@ class AMQPClientAsync(AMQPClientSync): 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.info("Keeping %r connection alive. %r", + self.__class__.__name__, + self._connection.container_id) + 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) async def __aenter__(self): """Run Client in an async context manager.""" @@ -152,9 +171,9 @@ async def _client_ready_async(self): # pylint: disable=no-self-use async def _client_run_async(self, **kwargs): """Perform a single Connection iteration.""" - await self._connection.listen(wait=self._socket_timeout) + await self._connection.listen(wait=self._socket_timeout, **kwargs) - async def _close_link_async(self, **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 @@ -182,8 +201,6 @@ async def _do_retryable_operation_async(self, operation, *args, **kwargs): if exc.condition in (ErrorCondition.ConnectionCloseForced, ErrorCondition.SocketError): # if connection detach or socket error, close and open a new connection await self.close_async() - except Exception: - raise finally: end_time = time.time() if absolute_timeout > 0: @@ -212,7 +229,7 @@ async def open_async(self, connection=None): self._connection = Connection( "amqps://" + self._hostname, sasl_credential=self._auth.sasl, - ssl={'ca_certs': self._connection_verify or certifi.where()}, + 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, @@ -238,6 +255,8 @@ async def open_async(self, connection=None): ) await self._cbs_authenticator.open() self._shutdown = False + 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 @@ -248,7 +267,10 @@ async def close_async(self): self._shutdown = True if not self._session: return # already closed. - await self._close_link_async(close=True) + 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 @@ -349,8 +371,8 @@ 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. + :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 @@ -397,7 +419,8 @@ class SendClientAsync(SendClientSync, AMQPClientAsync): :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] + :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', @@ -411,13 +434,13 @@ class SendClientAsync(SendClientSync, AMQPClientAsync): :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. + :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. + :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. @@ -546,7 +569,7 @@ async def _send_message_impl_async(self, message, **kwargs): MessageDeliveryState.Timeout ): try: - raise message_delivery.error + raise message_delivery.error # pylint: disable=raising-bad-type except TypeError: # This is a default handler raise MessageException(condition=ErrorCondition.UnknownError, description="Send failed.") @@ -610,7 +633,8 @@ class ReceiveClientAsync(ReceiveClientSync, AMQPClientAsync): :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] + :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', @@ -624,13 +648,13 @@ class ReceiveClientAsync(ReceiveClientSync, AMQPClientAsync): :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. + :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. + :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. @@ -855,7 +879,7 @@ async def settle_messages_async( 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 = Accepted() + state: Outcomes = Accepted() elif outcome.lower() == 'released': state = Released() elif outcome.lower() == 'rejected': @@ -867,7 +891,7 @@ async def settle_messages_async(self, delivery_id: Union[int, Tuple[int, int]], else: raise ValueError("Unrecognized message output: {}".format(outcome)) try: - first, last = delivery_id + first, last = cast(Tuple, delivery_id) except TypeError: first = delivery_id last = None @@ -879,4 +903,3 @@ async def settle_messages_async(self, delivery_id: Union[int, Tuple[int, int]], 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 index f03648302333..3a28c21aa2c4 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_connection_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_connection_async.py @@ -11,6 +11,7 @@ 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 @@ -29,14 +30,7 @@ TransportType, ) - -from ..error import ( - AMQPException, - ErrorCondition, - AMQPConnectionError, - AMQPError -) - +from ..error import ErrorCondition, AMQPConnectionError, AMQPError _LOGGER = logging.getLogger(__name__) @@ -90,7 +84,7 @@ def __init__(self, endpoint, **kwargs): # pylint:disable=too-many-statements if custom_endpoint_address: custom_parsed_url = urlparse(custom_endpoint_address) custom_port = custom_parsed_url.port or WEBSOCKET_PORT - custom_endpoint = "{}:{}{}".format(custom_parsed_url.hostname, custom_port, custom_parsed_url.path) + custom_endpoint = f"{custom_parsed_url.hostname}:{custom_port}{custom_parsed_url.path}" transport = kwargs.get("transport") self._transport_type = kwargs.pop("transport_type", TransportType.Amqp) @@ -98,35 +92,60 @@ def __init__(self, endpoint, **kwargs): # pylint:disable=too-many-statements self._transport = transport elif "sasl_credential" in kwargs: sasl_transport = SASLTransport - if self._transport_type.name == "AmqpOverWebsocket" or kwargs.get("http_proxy"): + 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, **kwargs + host=endpoint, + credential=kwargs["sasl_credential"], + custom_endpoint=custom_endpoint, + **kwargs, ) else: - self._transport = AsyncTransport(parsed_url.netloc, transport_type=self._transport_type, **kwargs) - - self._container_id = kwargs.pop("container_id", None) or str(uuid.uuid4()) # type: str - self._max_frame_size = kwargs.pop("max_frame_size", MAX_FRAME_SIZE_BYTES) # type: int + self._transport = AsyncTransport(parsed_url.netloc, **kwargs) + + self._container_id = kwargs.pop("container_id", None) or str( + uuid.uuid4() + ) # type: str + 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._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._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._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._network_trace = kwargs.get("network_trace", False) - self._network_trace_params = {"connection": self._container_id, "session": None, "link": None} + self._network_trace_params = { + "connection": self._container_id, + "session": None, + "link": None, + } self._error = None self._outgoing_endpoints = {} # type: Dict[int, Session] self._incoming_endpoints = {} # type: Dict[int, Session] @@ -145,7 +164,12 @@ async def _set_state(self, new_state): return previous_state = self.state self.state = new_state - _LOGGER.info("Connection '%s' state changed: %r -> %r", self._container_id, previous_state, new_state) + _LOGGER.info( + "Connection '%s' state changed: %r -> %r", + self._container_id, + previous_state, + new_state, + ) for session in self._outgoing_endpoints.values(): await session._on_connection_state_change() # pylint:disable=protected-access @@ -170,17 +194,20 @@ async def _connect(self): await self._process_incoming_frame(*(await self._read_frame(wait=True))) if self.state != ConnectionState.HDR_EXCH: await self._disconnect() - raise ValueError("Did not receive reciprocal protocol header. Disconnecting.") + 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: raise AMQPConnectionError( ErrorCondition.SocketError, - description="Failed to initiate the connection due to exception: " + str(exc), + description="Failed to initiate the connection due to exception: " + + str(exc), error=exc, ) - async def _disconnect(self, *args) -> None: + async def _disconnect(self) -> None: """Disconnect the transport and set state to END.""" if self.state == ConnectionState.END: return @@ -192,7 +219,7 @@ def _can_read(self): """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=True, **kwargs): + async def _read_frame(self, wait=True, **kwargs): # type: ignore # TODO: missing return # type: (bool, Any) -> Tuple[int, Optional[Tuple[int, NamedTuple]]] """Read an incoming frame from the transport. @@ -236,8 +263,17 @@ async def _send_frame(self, channel, frame, timeout=None, **kwargs): 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: + 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), @@ -254,9 +290,17 @@ def _get_next_outgoing_channel(self): :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) + 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): @@ -284,7 +328,9 @@ async def _outgoing_header(self): """Send the AMQP protocol header to initiate the connection.""" self._last_frame_sent_time = time.time() if self._network_trace: - _LOGGER.info("-> header(%r)", HEADER_FRAME, extra=self._network_trace_params) + _LOGGER.info( + "-> header(%r)", HEADER_FRAME, extra=self._network_trace_params + ) await self._transport.write(HEADER_FRAME) async def _incoming_header(self, _, frame): @@ -307,11 +353,17 @@ async def _outgoing_open(self): 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 + 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, + 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: @@ -345,28 +397,38 @@ async def _incoming_open(self, channel, frame): _LOGGER.info("<- %r", OpenFrame(*frame), extra=self._network_trace_params) if channel != 0: _LOGGER.error("OPEN frame received on a channel that is not 0.") - await self.close(error=AMQPError( - condition=ErrorCondition.NotAllowed, - description="OPEN frame received on a channel that is not 0")) + 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.") 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 + self._remote_idle_timeout_send_frame = ( + self._idle_timeout_empty_frame_send_ratio * self._remote_idle_timeout + ) if frame[2] < 512: - # Close with error - # Codes_S_R_S_CONNECTION_01_143: [If any of the values in the received open frame are invalid then the connection shall be closed.] - # Codes_S_R_S_CONNECTION_01_220: [The error amqp:invalid-field shall be set in the error.condition field of the CLOSE frame.] + # 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=AMQPConnectionError( - condition=ErrorCondition.InvalidField, - description="connection_endpoint_frame_received::failed parsing OPEN frame", + error=cast( + AMQPError, + AMQPConnectionError( + condition=ErrorCondition.InvalidField, + description="Failed parsing OPEN frame: Max frame size is less than supported minimum.", + ), ) ) - _LOGGER.error("connection_endpoint_frame_received::failed parsing OPEN frame") + _LOGGER.error( + "Failed parsing OPEN frame: Max frame size is less than supported minimum." + ) else: self._remote_max_frame_size = frame[2] if self.state == ConnectionState.OPEN_SENT: @@ -376,9 +438,12 @@ async def _incoming_open(self, channel, frame): 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}")) + 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) async def _outgoing_close(self, error=None): @@ -415,7 +480,11 @@ async def _incoming_close(self, channel, frame): close_error = None if channel > self._channel_max: _LOGGER.error("Invalid channel") - close_error = AMQPError(condition=ErrorCondition.InvalidField, description="Invalid channel", info=None) + 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) @@ -423,8 +492,12 @@ async def _incoming_close(self, channel, frame): await self._set_state(ConnectionState.END) if frame[0]: - self._error = AMQPConnectionError(condition=frame[0][0], description=frame[0][1], info=frame[0][2]) - _LOGGER.error("Connection error: {}".format(frame[0])) # pylint:disable=logging-format-interpolation + self._error = AMQPConnectionError( + condition=frame[0][0], description=frame[0][1], info=frame[0][2] + ) + _LOGGER.error( + "Connection error: {}".format(frame[0]) # pylint:disable=logging-format-interpolation + ) async def _incoming_begin(self, channel, frame): # type: (int, Tuple[Any, ...]) -> None @@ -449,9 +522,11 @@ async def _incoming_begin(self, channel, frame): try: existing_session = self._outgoing_endpoints[frame[0]] self._incoming_endpoints[channel] = existing_session - await self._incoming_endpoints[channel]._incoming_begin(frame) # pylint:disable=protected-access + await self._incoming_endpoints[channel]._incoming_begin( # pylint:disable=protected-access + frame + ) except KeyError: - new_session = Session.from_incoming_frame(self, channel, frame) + new_session = Session.from_incoming_frame(self, channel) self._incoming_endpoints[channel] = new_session await new_session._incoming_begin(frame) # pylint:disable=protected-access @@ -481,7 +556,9 @@ async def _incoming_end(self, channel, frame): )) return - async def _process_incoming_frame(self, channel, frame): # pylint:disable=too-many-return-statements + 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. @@ -495,25 +572,36 @@ async def _process_incoming_frame(self, channel, frame): # pylint:disable=too-m should be interrupted. """ try: - performative, fields = frame + 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(fields) # pylint:disable=protected-access + await self._incoming_endpoints[channel]._incoming_transfer( # pylint:disable=protected-access + fields + ) return False if performative == 21: - await self._incoming_endpoints[channel]._incoming_disposition(fields) # pylint:disable=protected-access + await self._incoming_endpoints[channel]._incoming_disposition( # pylint:disable=protected-access + fields + ) return False if performative == 19: - await self._incoming_endpoints[channel]._incoming_flow(fields) # pylint:disable=protected-access + await self._incoming_endpoints[channel]._incoming_flow( # pylint:disable=protected-access + fields + ) return False if performative == 18: - await self._incoming_endpoints[channel]._incoming_attach(fields) # pylint:disable=protected-access + await self._incoming_endpoints[channel]._incoming_attach( # pylint:disable=protected-access + fields + ) return False if performative == 22: - await self._incoming_endpoints[channel]._incoming_detach(fields) # pylint:disable=protected-access + await self._incoming_endpoints[channel]._incoming_detach( # pylint:disable=protected-access + fields + ) return True if performative == 17: await self._incoming_begin(channel, fields) @@ -528,15 +616,12 @@ async def _process_incoming_frame(self, channel, frame): # pylint:disable=too-m await self._incoming_close(channel, fields) return True if performative == 0: - await self._incoming_header(channel, fields) + await self._incoming_header(channel, cast(bytes, fields)) return True - if performative == 1: # pylint:disable=no-else-return + if performative == 1: return False # TODO: incoming EMPTY - else: - _LOGGER.error( - "Unrecognized incoming frame: {}".format(frame) - ) # pylint:disable=logging-format-interpolation - return True + _LOGGER.error("Unrecognized incoming frame: %s", frame) + return True except KeyError: return True # TODO: channel error @@ -546,14 +631,23 @@ async def _process_outgoing_frame(self, channel, frame): :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]: + 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]: + 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, self._idle_timeout, self._last_frame_received_time) or ( - await self._get_remote_timeout(now) - ): + if get_local_timeout( + now, + cast(float, self._idle_timeout), + cast(float, self._last_frame_received_time), + ) or (await self._get_remote_timeout(now)): await self.close( error=AMQPError( condition=ErrorCondition.ConnectionCloseForced, @@ -578,7 +672,7 @@ async def _get_remote_timeout(self, now): """ 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 > self._remote_idle_timeout_send_frame: + if time_since_last_sent > cast(int, self._remote_idle_timeout_send_frame): await self._outgoing_empty() return False @@ -631,8 +725,11 @@ async def listen(self, wait=False, batch=1, **kwargs): try: if self.state not in _CLOSING_STATES: now = time.time() - if get_local_timeout(now, self._idle_timeout, self._last_frame_received_time) or ( - await self._get_remote_timeout(now)): + if get_local_timeout( + now, + cast(float, self._idle_timeout), + cast(float, self._last_frame_received_time), + ) or (await self._get_remote_timeout(now)): await self.close( error=AMQPError( condition=ErrorCondition.ConnectionCloseForced, @@ -642,12 +739,16 @@ async def listen(self, wait=False, batch=1, **kwargs): ) return if self.state == ConnectionState.END: + # TODO: check error condition self._error = AMQPConnectionError( - condition=ErrorCondition.ConnectionCloseForced, description="Connection was already closed." + condition=ErrorCondition.ConnectionCloseForced, + description="Connection was already closed.", ) return for _ in range(batch): - if await asyncio.ensure_future(self._listen_one_frame(wait=wait, **kwargs)): + if await asyncio.ensure_future( + self._listen_one_frame(wait=wait, **kwargs) + ): # TODO: compare the perf difference between ensure_future and direct await break except (OSError, IOError, SSLError, socket.error) as exc: @@ -686,7 +787,7 @@ def create_session(self, **kwargs): assigned_channel, network_trace=kwargs.pop("network_trace", self._network_trace), network_trace_params=dict(self._network_trace_params), - **kwargs + **kwargs, ) self._outgoing_endpoints[assigned_channel] = session return session @@ -710,7 +811,9 @@ async def open(self, wait=False): 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.") + 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 @@ -722,13 +825,19 @@ async def close(self, error=None, wait=False): :param bool wait: Whether to wait for a service Close response. Default is `False`. :rtype: None """ - if self.state in [ConnectionState.END, ConnectionState.CLOSE_SENT, ConnectionState.DISCARDING]: + if self.state in [ + ConnectionState.END, + ConnectionState.CLOSE_SENT, + ConnectionState.DISCARDING, + ]: return try: await self._outgoing_close(error=error) if error: self._error = AMQPConnectionError( - condition=error.condition, description=error.description, info=error.info + condition=error.condition, + description=error.description, + info=error.info, ) if self.state == ConnectionState.OPEN_PIPE: await self._set_state(ConnectionState.OC_PIPE) 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 index dc79c37c30a4..174fb61ee128 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_link_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_link_async.py @@ -8,8 +8,6 @@ import uuid import logging -import asyncio - from ..endpoints import Source, Target from ..constants import DEFAULT_LINK_CREDIT, SessionState, LinkState, Role, SenderSettleMode, ReceiverSettleMode from ..performatives import ( @@ -17,20 +15,17 @@ DetachFrame, ) - -from ..error import ( - AMQPError, - ErrorCondition, - AMQPLinkError, - AMQPLinkRedirect, - AMQPConnectionError -) +from ..error import ErrorCondition, AMQPLinkError, AMQPLinkRedirect, AMQPConnectionError _LOGGER = logging.getLogger(__name__) -class Link(object): - """ """ +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 @@ -170,20 +165,11 @@ async def _incoming_attach(self, frame): if self.network_trace: _LOGGER.info("<- %r", AttachFrame(*frame), extra=self.network_trace_params) if self._is_closed: - raise AMQPLinkError( - condition=ErrorCondition.ClientError, - description="Received attach frame on a link that is already closed", - info=None, - ) - elif not frame[5] or not frame[6]: + raise ValueError("Invalid link") + if not frame[5] or not frame[6]: _LOGGER.info("Cannot get source or target. Detaching link") - await self.detach( - error=AMQPError( - condition=ErrorCondition.LinkDetachForced, - description="Cannot get source or target from the frame. Detaching link", - info=None, - ) - ) + 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 @@ -202,8 +188,8 @@ async def _incoming_attach(self, frame): if frame[6]: frame[6] = Target(*frame[6]) await self._on_attach(AttachFrame(*frame)) - except Exception as e: - _LOGGER.warning("Callback for link attach raised error: {}".format(e)) + except Exception as e: # pylint: disable=broad-except + _LOGGER.warning("Callback for link attach raised error: %s", e) async def _outgoing_flow(self, **kwargs): flow_frame = { @@ -267,7 +253,7 @@ async def detach(self, close=False, error=None): elif self.state == LinkState.ATTACHED: await self._outgoing_detach(close=close, error=error) await self._set_state(LinkState.DETACH_SENT) - except Exception as exc: + except Exception as exc: # pylint: disable=broad-except _LOGGER.info("An error occurred when detaching the link: %r", exc) await self._set_state(LinkState.DETACHED) 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 index d6dfde96e7b1..3928f93d2ff7 100644 --- 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 @@ -20,6 +20,7 @@ ManagementOpenResult, SEND_DISPOSITION_REJECT, MessageDeliveryState, + LinkDeliverySettleReason ) from ..error import AMQPException, ErrorCondition from ..message import Properties, _MessageDelivery @@ -142,8 +143,8 @@ async def _on_message_received(self, _, message): ) self._pending_operations.remove(to_remove_operation) - async def _on_send_complete(self, message_delivery, reason, state): # todo: reason is never used, should check spec - if SEND_DISPOSITION_REJECT in state: + 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: @@ -154,7 +155,9 @@ async def _on_send_complete(self, message_delivery, reason, state): # todo: rea # 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 - await to_remove_operation.on_execute_operation_complete( # The callback is defined in management_operation.py + + # The callback is defined in management_operation.py + await to_remove_operation.on_execute_operation_complete( ManagementExecuteOperationResult.ERROR, None, None, 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 index d6f225a34427..f7ebb5f667bf 100644 --- 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 @@ -104,7 +104,7 @@ async def execute(self, message, operation=None, operation_type=None, timeout=0) if self._mgmt_error: self._responses.pop(operation_id) - raise self._mgmt_error + raise self._mgmt_error # pylint: disable=raising-bad-type response = self._responses.pop(operation_id) return response @@ -115,7 +115,7 @@ async def open(self): async def ready(self): try: - raise self._mgmt_error + raise self._mgmt_error # pylint: disable=raising-bad-type except TypeError: pass 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 index dea7ac8bda91..b5748909c747 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_receiver_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_receiver_async.py @@ -31,10 +31,16 @@ def __init__(self, session, handle, 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: + except Exception as e: # pylint: disable=broad-except _LOGGER.error("Handler function failed with error: %r", e) return None @@ -67,7 +73,13 @@ async def _incoming_transfer(self, frame): _LOGGER.info(" %r", message, extra=self.network_trace_params) delivery_state = await self._process_incoming_message(frame, message) if not frame[4] and delivery_state: # settled - await self._outgoing_disposition(first=frame[1], settled=True, state=delivery_state) + 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: 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 index acb1079eae44..441eb40ec874 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_sasl_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_sasl_async.py @@ -1,8 +1,8 @@ -#------------------------------------------------------------------------- +# ------------------------------------------------------------------------- # 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 @@ -10,7 +10,7 @@ from ..performatives import SASLInit -_SASL_FRAME_TYPE = b'\x01' +_SASL_FRAME_TYPE = b"\x01" # TODO: do we need it here? it's a duplicate of the sync version @@ -19,7 +19,7 @@ class SASLPlainCredential(object): See https://tools.ietf.org/html/rfc4616 for details """ - mechanism = b'PLAIN' + mechanism = b"PLAIN" def __init__(self, authcid, passwd, authzid=None): self.authcid = authcid @@ -28,13 +28,13 @@ def __init__(self, authcid, passwd, authzid=None): def start(self): if self.authzid: - login_response = self.authzid.encode('utf-8') + 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') + 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 @@ -44,10 +44,10 @@ class SASLAnonymousCredential(object): See https://tools.ietf.org/html/rfc4505 for details """ - mechanism = b'ANONYMOUS' + mechanism = b"ANONYMOUS" - def start(self): - return b'' + def start(self): # pylint: disable=no-self-use + return b"" # TODO: do we need it here? it's a duplicate of the sync version @@ -58,27 +58,34 @@ class SASLExternalCredential(object): authentication data. """ - mechanism = b'EXTERNAL' + mechanism = b"EXTERNAL" - def start(self): - return b'' + def start(self): # pylint: disable=no-self-use + return b"" -class SASLTransportMixinAsync(): +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("Mismatching AMQP header protocol. Expected: {}, received: {}".format( - SASL_HEADER_FRAME, returned_header[1])) + 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)) + 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) + 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) @@ -87,33 +94,56 @@ async def _negotiate(self): raise NotImplementedError("Unsupported SASL challenge") if fields[0] == SASLCode.Ok: # code return - raise ValueError("SASL negotiation failed.\nOutcome: {}\nDetails: {}".format(*fields)) + 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=None, **kwargs): + def __init__( + self, + host, + credential, + *, + port=AMQPS_PORT, + connect_timeout=None, + ssl_opts=None, + **kwargs, + ): self.credential = credential - ssl = ssl or True - super(SASLTransport, self).__init__(host, port=port, connect_timeout=connect_timeout, ssl=ssl, **kwargs) + 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=None, **kwargs): + def __init__( + self, + host, + credential, + *, + port=WEBSOCKET_PORT, + connect_timeout=None, + ssl_opts=None, + **kwargs, + ): self.credential = credential - ssl = ssl or True - self._transport = WebSocketTransportAsync( + ssl_opts = ssl_opts or True + super().__init__( host, port=port, connect_timeout=connect_timeout, - ssl=ssl, - **kwargs + ssl_opts=ssl_opts, + **kwargs, ) - super().__init__(host, port, connect_timeout, ssl, **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 index 47f6c5be41cf..37bb95a705f4 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_sender_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_sender_async.py @@ -49,6 +49,12 @@ def __init__(self, session, handle, target_address, **kwargs): 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): @@ -157,7 +163,7 @@ async def update_pending_deliveries(self): async def send_transfer(self, message, *, send_async=False, **kwargs): self._check_if_closed() if self.state != LinkState.ATTACHED: - raise AMQPLinkError( + raise AMQPLinkError( condition=ErrorCondition.ClientError, description="Link is not attached." ) 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 index 2e16ca20ddec..12ad6e3516d6 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_session_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_session_async.py @@ -4,33 +4,31 @@ # license information. # -------------------------------------------------------------------------- +from __future__ import annotations import uuid import logging import time import asyncio -from typing import Optional, Union +from typing import Optional, Union, List -from azure.eventhub._pyamqp.error import AMQPError, AMQPSessionError, ErrorCondition - -from ..constants import ( - INCOMING_WINDOW, - OUTGOING_WINDOW, - ConnectionState, - SessionState, - SessionTransferState, - Role -) -from ..endpoints import Source, Target +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 ..performatives import ( + BeginFrame, + EndFrame, + FlowFrame, + TransferFrame, + DispositionFrame, +) from .._encode import encode_frame +from ..error import AMQPError, ErrorCondition _LOGGER = logging.getLogger(__name__) -class Session(object): +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. @@ -78,7 +76,7 @@ async def __aexit__(self, *args): await self.end() @classmethod - def from_incoming_frame(cls, connection, channel, frame): + def from_incoming_frame(cls, connection, channel): # check session_create_from_endpoint in C lib new_session = cls(connection, channel) return new_session @@ -90,11 +88,20 @@ async def _set_state(self, new_state): 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) - futures = [] - for link in self.links.values(): - futures.append(asyncio.ensure_future(link._on_session_state_change())) # pylint: disable=protected-access - await asyncio.gather(*futures) + _LOGGER.info( + "Session state changed: %r -> %r", + previous_state, + new_state, + extra=self.network_trace_params, + ) + await asyncio.gather( + *[ + asyncio.ensure_future( + link._on_session_state_change() # pylint: disable=protected-access + ) + for link in self.links.values() + ] + ) async def _on_connection_state_change(self): if self._connection.state in [ConnectionState.CLOSE_RCVD, ConnectionState.END]: @@ -110,24 +117,38 @@ def _get_next_output_handle(self): :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) + 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, + 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, + 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.info("-> %r", begin_frame, extra=self.network_trace_params) - await self._connection._process_outgoing_frame(self.channel, begin_frame) # pylint: disable=protected-access + await self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, begin_frame + ) async def _incoming_begin(self, frame): if self.network_trace: @@ -148,12 +169,18 @@ async def _outgoing_end(self, error=None): end_frame = EndFrame(error=error) if self.network_trace: _LOGGER.info("-> %r", end_frame, extra=self.network_trace_params) - await self._connection._process_outgoing_frame(self.channel, end_frame) # pylint: disable=protected-access + await self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, end_frame + ) async def _incoming_end(self, frame): if self.network_trace: _LOGGER.info("<- %r", EndFrame(*frame), extra=self.network_trace_params) - if self.state not in [SessionState.END_RCVD, SessionState.END_SENT, SessionState.DISCARDING]: + 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() @@ -162,21 +189,30 @@ async def _incoming_end(self, frame): 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 + await self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, frame + ) 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 + self._input_handles[frame[1]] = self.links[ + frame[0].decode("utf-8") + ] # name and handle + await self._input_handles[ # pylint: disable=protected-access + frame[1] + ]._incoming_attach( + frame + ) except KeyError: try: outgoing_handle = self._get_next_output_handle() except ValueError: # detach the link that would have been set. - await self.links[frame[0].decode('utf-8')].detach( + await 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( + description="""Cannot allocate more handles, """ + """the max number of handles is {}. Detaching link""".format( self.handle_max ), info=None, @@ -184,10 +220,12 @@ async def _incoming_attach(self, frame): ) return if frame[2] == Role.Sender: - new_link = ReceiverLink.from_incoming_frame(self, outgoing_handle, frame) + 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 + 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 @@ -208,22 +246,37 @@ async def _outgoing_flow(self, frame=None): flow_frame = FlowFrame(**link_flow) if self.network_trace: _LOGGER.info("-> %r", flow_frame, extra=self.network_trace_params) - await self._connection._process_outgoing_frame(self.channel, flow_frame) # pylint: disable=protected-access + await self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, flow_frame + ) async def _incoming_flow(self, frame): if self.network_trace: _LOGGER.info("<- %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 + 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 + await self._input_handles[ # pylint: disable=protected-access + frame[4] + ]._incoming_flow( + frame + ) else: futures = [] for link in self._output_handles.values(): - if self.remote_incoming_window > 0 and not link._is_closed: # pylint: disable=protected-access - futures.append(link._incoming_flow(frame)) # pylint: disable=protected-access + if ( + self.remote_incoming_window > 0 + and not link._is_closed # pylint: disable=protected-access + ): + futures.append( + link._incoming_flow(frame) # pylint: disable=protected-access + ) await asyncio.gather(*futures) async def _outgoing_transfer(self, delivery): @@ -244,7 +297,11 @@ async def _outgoing_transfer(self, delivery): # 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 + available_frame_size = ( + self._connection._remote_max_frame_size # pylint: disable=protected-access + - transfer_overhead_size + - 8 + ) start_idx = 0 remaining_payload_cnt = payload_size @@ -264,7 +321,9 @@ async def _outgoing_transfer(self, delivery): "payload": payload[start_idx : start_idx + available_frame_size], "delivery_id": self.next_outgoing_id, } - await self._connection._process_outgoing_frame(self.channel, TransferFrame(**tmp_delivery_frame)) # pylint: disable=protected-access + await self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, TransferFrame(**tmp_delivery_frame) + ) start_idx += available_frame_size remaining_payload_cnt -= available_frame_size @@ -283,7 +342,9 @@ async def _outgoing_transfer(self, delivery): "payload": payload[start_idx:], "delivery_id": self.next_outgoing_id, } - await self._connection._process_outgoing_frame(self.channel, TransferFrame(**tmp_delivery_frame)) # pylint: disable=protected-access + await self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, TransferFrame(**tmp_delivery_frame) + ) self.next_outgoing_id += 1 self.remote_incoming_window -= 1 self.outgoing_window -= 1 @@ -295,43 +356,63 @@ async def _incoming_transfer(self, frame): self.remote_outgoing_window -= 1 self.incoming_window -= 1 try: - await self._input_handles[frame[0]]._incoming_transfer(frame) # pylint: disable=protected-access + await self._input_handles[ # pylint: disable=protected-access + frame[0] + ]._incoming_transfer( + frame + ) except KeyError: await self._set_state(SessionState.DISCARDING) - await self.end(error=AMQPError( + await self.end( + error=AMQPError( condition=ErrorCondition.SessionUnattachedHandle, - description="Invalid handle reference in received frame: Handle is not currently associated with an attached link")) + 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 + await self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, frame + ) async def _incoming_disposition(self, frame): if self.network_trace: - _LOGGER.info("<- %r", DispositionFrame(*frame), extra=self.network_trace_params) + _LOGGER.info( + "<- %r", DispositionFrame(*frame), extra=self.network_trace_params + ) futures = [] for link in self._input_handles.values(): - asyncio.ensure_future(link._incoming_disposition(frame)) # pylint: disable=protected-access + asyncio.ensure_future( + link._incoming_disposition(frame) # pylint: disable=protected-access + ) await asyncio.gather(*futures) async def _outgoing_detach(self, frame): - await self._connection._process_outgoing_frame(self.channel, frame) # pylint: disable=protected-access + await self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, frame + ) async def _incoming_detach(self, frame): try: link = self._input_handles[frame[0]] # handle - await link._incoming_detach(frame) # pylint: disable=protected-access + 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")) + 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 @@ -355,7 +436,9 @@ async def begin(self, wait=False): 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.") + 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 @@ -367,7 +450,7 @@ async def end(self, error=None, wait=False): 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: + except Exception as exc: # pylint: disable=broad-except _LOGGER.info("An error occurred when ending the session: %r", exc) await self._set_state(SessionState.UNMAPPED) @@ -379,7 +462,7 @@ def create_receiver_link(self, source_address, **kwargs): source_address=source_address, network_trace=kwargs.pop("network_trace", self.network_trace), network_trace_params=dict(self.network_trace_params), - **kwargs + **kwargs, ) self.links[link.name] = link self._output_handles[assigned_handle] = link @@ -393,11 +476,16 @@ def create_sender_link(self, target_address, **kwargs): target_address=target_address, network_trace=kwargs.pop("network_trace", self.network_trace), network_trace_params=dict(self.network_trace_params), - **kwargs + **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), **kwargs) + return ManagementLink( + self, + endpoint, + network_trace=kwargs.pop("network_trace", self.network_trace), + **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 index e3f3e6734d92..3380e215eb2c 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_transport_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_transport_async.py @@ -1,4 +1,4 @@ -# ------------------------------------------------------------------------- +# ------------------------------------------------------------------------- # 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. # @@ -37,12 +37,10 @@ import socket import ssl import struct -from ssl import SSLContext, SSLError -from contextlib import contextmanager +from ssl import SSLError from io import BytesIO import logging -from threading import Lock -from typing import Optional + import certifi @@ -60,49 +58,40 @@ set_cloexec, AMQP_PORT, TIMEOUT_INTERVAL, - WebSocketTransport, ) _LOGGER = logging.getLogger(__name__) -def get_running_loop(): - try: - import asyncio # pylint: disable=import-error - - return asyncio.get_running_loop() - except AttributeError: # 3.6 - loop = None - try: - loop = asyncio._get_running_loop() # pylint: disable=protected-access - except AttributeError: - _LOGGER.warning("This version of Python is deprecated, please upgrade to >= v3.6") - if loop is None: - _LOGGER.warning("No running event loop") - loop = asyncio.get_event_loop() - return loop - - class AsyncTransportMixin: - async def receive_frame(self, timeout=None, *args, **kwargs): + async def receive_frame(self, timeout=None, **kwargs): try: - header, channel, payload = await asyncio.wait_for(self.read(**kwargs), timeout=timeout) + 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) _LOGGER.info("ICH%d <- %r", channel, decoded) return channel, decoded - except (TimeoutError, socket.timeout, asyncio.IncompleteReadError, asyncio.TimeoutError): + except ( + TimeoutError, + socket.timeout, + asyncio.IncompleteReadError, + asyncio.TimeoutError, + ): return None, None - async def read(self, verify_frame_type=0, **kwargs): + 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)) + 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] @@ -111,16 +100,28 @@ async def read(self, verify_frame_type=0, **kwargs): 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: + 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:])) + 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)) + read_frame_buffer.write( + await self._read(payload_size, buffer=payload) + ) except (TimeoutError, socket.timeout, asyncio.IncompleteReadError): read_frame_buffer.write(self._read_buffer.getvalue()) self._read_buffer = read_frame_buffer @@ -148,13 +149,12 @@ async def send_frame(self, channel, frame, **kwargs): await self.write(data) # _LOGGER.info("OCH%d -> %r", channel, frame) - 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, **sslopts.pop("context")) + return self._build_ssl_context(**sslopts.pop("context")) ssl_version = sslopts.get("ssl_version") if ssl_version is None: ssl_version = ssl.PROTOCOL_TLS @@ -178,9 +178,13 @@ def _build_ssl_opts(self, sslopts): return context return True except TypeError: - raise TypeError("SSL configuration must be a dictionary, or the value True.") + raise TypeError( + "SSL configuration must be a dictionary, or the value True." + ) - def _build_ssl_context(self, sslopts, check_hostname=None, **ctx_options): + 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()) @@ -188,20 +192,21 @@ def _build_ssl_context(self, sslopts, check_hostname=None, **ctx_options): return ctx -class AsyncTransport(AsyncTransportMixin): +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, - read_timeout=None, - write_timeout=None, - ssl=False, + ssl_opts=False, socket_settings=None, raise_on_initial_eintr=True, - **kwargs, + **kwargs, # pylint: disable=unused-argument ): self.connected = False self.sock = None @@ -212,14 +217,10 @@ def __init__( self.host, self.port = to_host_port(host, port) self.connect_timeout = connect_timeout - self.read_timeout = read_timeout - self.write_timeout = write_timeout self.socket_settings = socket_settings - self.loop = get_running_loop() + self.loop = asyncio.get_running_loop() self.socket_lock = asyncio.Lock() - self.sslopts = self._build_ssl_opts(ssl) - - + self.sslopts = self._build_ssl_opts(ssl_opts) async def connect(self): try: @@ -227,13 +228,11 @@ async def connect(self): if self.connected: return await self._connect(self.host, self.port, self.connect_timeout) - self._init_socket( - self.socket_settings, - self.read_timeout, - self.write_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 + 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 @@ -262,7 +261,9 @@ async def _connect(self, host, port, timeout): for n, family in enumerate(addr_types): # first, resolve the address for a single address family try: - entries = await self.loop.getaddrinfo(host, port, family=family, type=socket.SOCK_STREAM, proto=SOL_TCP) + entries = await self.loop.getaddrinfo( + host, port, family=family, type=socket.SOCK_STREAM, proto=SOL_TCP + ) entries_num = len(entries) # now that we have address(es) for the hostname, connect to broker for i, res in enumerate(entries): @@ -292,29 +293,18 @@ async def _connect(self, host, port, timeout): # 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")) - - + raise e if e is not None else socket.error( + "failed to resolve broker hostname" + ) + continue # pragma: no cover - def _init_socket(self, socket_settings, read_timeout, write_timeout): + 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) - - # 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.sock.settimeout(1) # set socket back to non-blocking mode - def _get_tcp_socket_defaults(self, sock): + def _get_tcp_socket_defaults(self, sock): # pylint: disable=no-self-use tcp_opts = {} for opt in KNOWN_TCP_OPTS: enum = None @@ -341,7 +331,13 @@ def _set_socket_options(self, 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)): + 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. @@ -353,7 +349,9 @@ async def _read(self, toread, initial=False, buffer=None, _errnos=(errno.ENOENT, try: while toread: try: - view[nbytes : nbytes + toread] = await self.reader.readexactly(toread) + view[nbytes : nbytes + toread] = await self.reader.readexactly( + toread + ) nbytes = toread except asyncio.IncompleteReadError as exc: pbytes = len(exc.partial) @@ -405,7 +403,7 @@ async def write(self, s): self.connected = False raise - async def receive_frame_with_lock(self, *args, **kwargs): + async def receive_frame_with_lock(self, **kwargs): try: async with self.socket_lock: header, channel, payload = await self.read(**kwargs) @@ -421,26 +419,37 @@ async def negotiate(self): if not self.sslopts: return await self.write(TLS_HEADER_FRAME) - channel, returned_header = await self.receive_frame(verify_frame_type=None) + _, returned_header = await self.receive_frame(verify_frame_type=None) if returned_header[1] == TLS_HEADER_FRAME: raise ValueError( - "Mismatching TLS header protocol. Excpected: {}, received: {}".format( - TLS_HEADER_FRAME, returned_header[1] - ) + f"""Mismatching TLS header protocol. Expected: {TLS_HEADER_FRAME!r},""" + """received: {returned_header[1]!r}""" ) -class WebSocketTransportAsync(AsyncTransportMixin): - def __init__(self, host, port=WEBSOCKET_PORT, connect_timeout=None, ssl=None, **kwargs): +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 = self._build_ssl_opts(ssl) if isinstance(ssl, dict) else None + self.sslopts = self._build_ssl_opts(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 = host + 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 + async def connect(self): username, password = None, None http_proxy_host, http_proxy_port = None, None @@ -456,14 +465,23 @@ async def connect(self): try: from aiohttp import ClientSession + 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}" + self.ws = await self.session.ws_connect( - url="wss://{}".format(self._custom_endpoint or self.host), + url=url, timeout=self._connect_timeout, protocols=[AMQP_WS_SUBPROTOCOL], autoclose=False, @@ -471,11 +489,14 @@ async def connect(self): proxy_auth=http_proxy_auth, ssl=self.sslopts, ) - + self.connected = True + except ImportError: - raise ValueError("Please install aiohttp library to use websocket transport.") + raise ValueError( + "Please install aiohttp library to use websocket transport." + ) - async def _read(self, n, buffer=None, **kwargs): # pylint: disable=unused-arguments + async def _read(self, n, buffer=None, **kwargs): # pylint: disable=unused-argument """Read exactly n bytes from the peer.""" length = 0 @@ -483,7 +504,7 @@ async def _read(self, n, buffer=None, **kwargs): # pylint: disable=unused-argum nbytes = self._read_buffer.readinto(view) length += nbytes n -= nbytes - + try: while n: data = await self.ws.receive_bytes() @@ -493,9 +514,9 @@ async def _read(self, n, buffer=None, **kwargs): # pylint: disable=unused-argum else: view[length : length + n] = data[0:n] self._read_buffer = BytesIO(data[n:]) - n = 0 + n = 0 return view - except (asyncio.TimeoutError) as wex: + except asyncio.TimeoutError: raise TimeoutError() async def close(self): @@ -509,5 +530,5 @@ async def write(self, s): 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/cbs.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/cbs.py index 4d26b4ff2b1d..9270346faa6a 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/cbs.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/cbs.py @@ -10,7 +10,12 @@ 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 .error import ( + AuthenticationException, + ErrorCondition, + TokenAuthFailure, + TokenExpired, +) from .constants import ( CbsState, CbsAuthState, @@ -21,7 +26,6 @@ CBS_OPERATION, ManagementExecuteOperationResult, ManagementOpenResult, - DEFAULT_AUTH_TIMEOUT, ) _LOGGER = logging.getLogger(__name__) @@ -56,8 +60,8 @@ def __init__(self, session, auth, **kwargs): raise ValueError("get_token must be a callable object.") self._auth = auth - self._encoding = 'UTF-8' - self._auth_timeout = kwargs.get('auth_timeout') + self._encoding = "UTF-8" + self._auth_timeout = kwargs.get("auth_timeout") self._token_put_time = None self._expires_on = None self._token = None @@ -71,9 +75,9 @@ def __init__(self, session, auth, **kwargs): def _put_token(self, token, token_type, audience, expires_on=None): # type: (str, str, str, datetime) -> None - message = Message( + message = Message( # type: ignore # TODO: missing positional args header, etc. value=token, - properties=Properties(message_id=self._mgmt_link.next_message_id), + properties=Properties(message_id=self._mgmt_link.next_message_id), # type: ignore application_properties={ CBS_NAME: audience, CBS_OPERATION: CBS_PUT_TOKEN, @@ -92,7 +96,10 @@ def _put_token(self, token, token_type, audience, expires_on=None): 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) + _LOGGER.debug( + "CSB with status: %r encounters unexpected AMQP management open complete.", + self.state, + ) elif self.state == CbsState.OPEN: self.state = CbsState.ERROR _LOGGER.info( @@ -100,10 +107,14 @@ def _on_amqp_management_open_complete(self, management_open_result): self._connection._container_id, # pylint:disable=protected-access ) elif self.state == CbsState.OPENING: - self.state = CbsState.OPEN if management_open_result == ManagementOpenResult.OK else CbsState.CLOSED + self.state = ( + CbsState.OPEN + if management_open_result == ManagementOpenResult.OK + else CbsState.CLOSED + ) _LOGGER.info( "CBS for connection %r completed opening with status: %r", - self._connection._container_id, # pylint: disable=protected-access + self._connection._container_id, # pylint: disable=protected-access management_open_result, ) # pylint:disable=protected-access @@ -125,8 +136,17 @@ def _on_amqp_management_error(self): ) # pylint:disable=protected-access def _on_execute_operation_complete( - self, execute_operation_result, status_code, status_description, message, error_condition=None - ): # TODO: message and error_condition never used + self, + execute_operation_result, + status_code, + status_description, + _, + error_condition=None, + ): + if error_condition: + _LOGGER.info("CBS Put token error: %r", error_condition) + self.auth_state = CbsAuthState.ERROR + return _LOGGER.info( "CBS Put token result (%r), status code: %s, status_description: %s.", execute_operation_result, @@ -143,23 +163,38 @@ def _on_execute_operation_complete( # 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: + 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: + if ( + self.auth_state == CbsAuthState.OK + or self.auth_state == CbsAuthState.REFRESH_REQUIRED + ): _LOGGER.debug("update_status In refresh required or OK.") is_expired, is_refresh_required = check_expiration_and_refresh_status( self._expires_on, self._refresh_window ) - _LOGGER.debug("is expired == %r, is refresh required == %r", is_expired, is_refresh_required) + _LOGGER.debug( + "is expired == %r, is refresh required == %r", + is_expired, + is_refresh_required, + ) 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("In update status, in progress. token put time: %r", self._token_put_time) - put_timeout = check_put_timeout_status(self._auth_timeout, self._token_put_time) + _LOGGER.debug( + "In update status, in progress. token put time: %r", + self._token_put_time, + ) + put_timeout = check_put_timeout_status( + self._auth_timeout, self._token_put_time + ) if put_timeout: self.auth_state = CbsAuthState.TIMEOUT @@ -170,8 +205,8 @@ def _cbs_link_ready(self): return False if self.state in (CbsState.CLOSED, CbsState.ERROR): raise TokenAuthFailure( - condition=ErrorCondition.ClientError, - description="CBS authentication link is in broken status, please recreate the cbs link.", + status_code=ErrorCondition.ClientError, + status_description="CBS authentication link is in broken status, please recreate the cbs link.", ) def open(self): @@ -197,7 +232,12 @@ def update_token(self): 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)) + 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(): @@ -212,13 +252,15 @@ def handle_token(self): return True if self.auth_state == CbsAuthState.REFRESH_REQUIRED: _LOGGER.info( - "Token on connection %r will expire soon - attempting to refresh.", self._connection._container_id + "Token on connection %r will expire soon - attempting to refresh.", + self._connection._container_id, ) # pylint:disable=protected-access self.update_token() return False if self.auth_state == CbsAuthState.FAILURE: raise AuthenticationException( - condition=ErrorCondition.InternalError, description="Failed to open CBS authentication link." + condition=ErrorCondition.InternalError, + description="Failed to open CBS authentication link.", ) if self.auth_state == CbsAuthState.ERROR: raise TokenAuthFailure( @@ -229,4 +271,7 @@ def handle_token(self): 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.") + 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 index e1c088f41c8d..8a17b202ef4e 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/client.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/client.py @@ -1,9 +1,12 @@ -# ------------------------------------------------------------------------- +# ------------------------------------------------------------------------- # 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 @@ -11,34 +14,21 @@ import time import uuid from functools import partial -from typing import Any, Dict, Optional, Tuple, Union, overload +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 .session import Session -from .sender import SenderLink -from .receiver import ReceiverLink -from .sasl import SASLAnonymousCredential, SASLTransport -from .endpoints import Source, Target from .error import ( - AMQPConnectionError, AMQPException, - ErrorResponse, ErrorCondition, MessageException, MessageSendFailed, RetryPolicy, - AMQPError -) -from .outcomes import( - Received, - Rejected, - Released, - Accepted, - Modified + AMQPError, ) +from .outcomes import Received, Rejected, Released, Accepted, Modified from .constants import ( MAX_CHANNELS, @@ -60,11 +50,15 @@ from .management_operation import ManagementOperation from .cbs import CBSAuthenticator +Outcomes = Union[Received, Rejected, Released, Accepted, Modified] + _logger = logging.getLogger(__name__) -class AMQPClient(object): +class AMQPClient( + object +): # pylint: disable=too-many-instance-attributes """An AMQP client. :param hostname: The AMQP endpoint to connect to. :type hostname: str @@ -112,7 +106,8 @@ class AMQPClient(object): :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] + :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', @@ -126,13 +121,13 @@ class AMQPClient(object): :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. + :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. @@ -170,30 +165,41 @@ def __init__(self, hostname, **kwargs): 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._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) + "remote_idle_timeout_empty_frame_send_ratio", None + ) self._network_trace = kwargs.pop("network_trace", False) # 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) + 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._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") + 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) @@ -221,10 +227,10 @@ def _client_ready(self): # pylint: disable=no-self-use def _client_run(self, **kwargs): """Perform a single Connection iteration.""" - self._connection.listen(wait=self._socket_timeout) + self._connection.listen(wait=self._socket_timeout, **kwargs) - def _close_link(self, **kwargs): - if self._link and not self._link._is_closed: # pylint: disable=protected-access + 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 @@ -248,11 +254,12 @@ def _do_retryable_operation(self, operation, *args, **kwargs): 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 exc.condition in ( + ErrorCondition.ConnectionCloseForced, + ErrorCondition.SocketError, + ): # if connection detach or socket error, close and open a new connection self.close() - except Exception: - raise finally: end_time = time.time() if absolute_timeout > 0: @@ -270,7 +277,6 @@ def open(self, connection=None): multiple clients. :type connection: ~pyamqp.Connection """ - # pylint: disable=protected-access if self._session: return # already open. @@ -282,7 +288,7 @@ def open(self, connection=None): self._connection = Connection( "amqps://" + self._hostname, sasl_credential=self._auth.sasl, - ssl={"ca_certs": self._connection_verify or certifi.where()}, + 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, @@ -296,7 +302,8 @@ def open(self, connection=None): self._connection.open() if not self._session: self._session = self._connection.create_session( - incoming_window=self._incoming_window, outgoing_window=self._outgoing_window + incoming_window=self._incoming_window, + outgoing_window=self._outgoing_window, ) self._session.begin() if self._auth.auth_type == AUTH_TYPE_CBS: @@ -321,7 +328,7 @@ def close(self): self._shutdown = True if not self._session: return # already closed. - self._close_link(close=True) + self._close_link() if self._cbs_authenticator: self._cbs_authenticator.close() self._cbs_authenticator = None @@ -416,10 +423,10 @@ def mgmt_request(self, message, **kwargs): 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. + :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 @@ -465,7 +472,8 @@ class SendClient(AMQPClient): :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] + :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', @@ -479,13 +487,13 @@ class SendClient(AMQPClient): :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. + :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. @@ -510,9 +518,9 @@ class SendClient(AMQPClient): 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) + 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): @@ -560,7 +568,10 @@ 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 + message_delivery.message, + on_send_complete=on_send_complete, + timeout=timeout, + send_async=True, ) return delivery @@ -571,7 +582,9 @@ def _process_send_error(message_delivery, condition, description=None, info=None except ValueError: error = MessageException(condition, description=description, info=info) else: - error = MessageSendFailed(amqp_condition, description=description, info=info) + error = MessageSendFailed( + amqp_condition, description=description, info=info + ) message_delivery.state = MessageDeliveryState.Error message_delivery.error = error @@ -590,7 +603,9 @@ def _on_send_complete(self, message_delivery, reason, state): info=error_info[0][2], ) except TypeError: - self._process_send_error(message_delivery, condition=ErrorCondition.UnknownError) + self._process_send_error( + message_delivery, condition=ErrorCondition.UnknownError + ) elif reason == LinkDeliverySettleReason.SETTLED: message_delivery.state = MessageDeliveryState.Ok elif reason == LinkDeliverySettleReason.TIMEOUT: @@ -598,13 +613,17 @@ def _on_send_complete(self, message_delivery, reason, state): message_delivery.error = TimeoutError("Sending message timed out.") else: # NotDelivered and other unknown errors - self._process_send_error(message_delivery, condition=ErrorCondition.UnknownError) + 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) + message_delivery = _MessageDelivery( + message, MessageDeliveryState.WaitingToBeSent, expire_time + ) while not self.client_ready(): time.sleep(0.05) @@ -618,10 +637,12 @@ def _send_message_impl(self, message, **kwargs): MessageDeliveryState.Timeout, ): try: - raise message_delivery.error # pylint: disable=raising-bad-type + raise message_delivery.error # pylint: disable=raising-bad-type except TypeError: # This is a default handler - raise MessageException(condition=ErrorCondition.UnknownError, description="Send failed.") + raise MessageException( + condition=ErrorCondition.UnknownError, description="Send failed." + ) def send_message(self, message, **kwargs): """ @@ -683,7 +704,8 @@ class ReceiveClient(AMQPClient): :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] + :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', @@ -697,13 +719,13 @@ class ReceiveClient(AMQPClient): :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. + :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. @@ -727,14 +749,14 @@ class ReceiveClient(AMQPClient): def __init__(self, hostname, source, **kwargs): self.source = source - self._streaming_receive = kwargs.pop("streaming_receive", False) + self._streaming_receive = kwargs.pop("streaming_receive", False) self._received_messages = queue.Queue() - self._message_received_callback = kwargs.pop("message_received_callback", None) + 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) + 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): @@ -794,7 +816,9 @@ def _message_received(self, frame, 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): + 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 @@ -922,10 +946,12 @@ def settle_messages( ): ... - def settle_messages(self, delivery_id: Union[int, Tuple[int, int]], outcome: str, **kwargs): + def settle_messages( + self, delivery_id: Union[int, Tuple[int, int]], outcome: str, **kwargs + ): batchable = kwargs.pop("batchable", None) if outcome.lower() == "accepted": - state = Accepted() + state: Outcomes = Accepted() elif outcome.lower() == "released": state = Released() elif outcome.lower() == "rejected": @@ -937,7 +963,7 @@ def settle_messages(self, delivery_id: Union[int, Tuple[int, int]], outcome: str else: raise ValueError("Unrecognized message output: {}".format(outcome)) try: - first, last = delivery_id + first, last = cast(Tuple, delivery_id) except TypeError: first = delivery_id last = None diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/constants.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/constants.py index 2e26ea451667..e55474d33103 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/constants.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/constants.py @@ -3,6 +3,7 @@ # 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 @@ -64,7 +65,7 @@ DEFAULT_LINK_CREDIT = 10000 -FIELD = namedtuple('field', 'name, type, mandatory, default, multiple') +FIELD = namedtuple('FIELD', 'name, type, mandatory, default, multiple') STRING_FILTER = b"apache.org:selector-filter:string" @@ -329,6 +330,7 @@ class TransportType(Enum): 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 index a2d0b4a240e7..2d2de0a2868e 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/endpoints.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/endpoints.py @@ -14,6 +14,7 @@ # - 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 @@ -128,7 +129,7 @@ class ApacheFilters(object): Source = namedtuple( - 'source', + 'Source', [ 'address', 'durable', @@ -142,9 +143,9 @@ class ApacheFilters(object): 'outcomes', 'capabilities' ]) -Source.__new__.__defaults__ = (None,) * len(Source._fields) -Source._code = 0x00000028 # pylint: disable=protected-access -Source._definition = ( # pylint: disable=protected-access +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), @@ -217,7 +218,7 @@ class ApacheFilters(object): Target = namedtuple( - 'target', + 'Target', [ 'address', 'durable', @@ -227,9 +228,9 @@ class ApacheFilters(object): 'dynamic_node_properties', 'capabilities' ]) -Target._code = 0x00000029 # pylint: disable=protected-access -Target.__new__.__defaults__ = (None,) * len(Target._fields) # pylint: disable=protected-access -Target._definition = ( # pylint: disable=protected-access +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), diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/error.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/error.py index 96c7803fa0bb..91f3393eb8bf 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/error.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/error.py @@ -4,6 +4,7 @@ # license information. #-------------------------------------------------------------------------- +# TODO: fix mypy errors for _code/_definition/__defaults__ (issue #26500) from enum import Enum from collections import namedtuple @@ -87,7 +88,7 @@ class ErrorCondition(bytes, Enum): SocketError = b"amqp:socket-error" -class RetryMode(str, Enum): +class RetryMode(str, Enum): # pylint: disable=enum-must-inherit-case-insensitive-enum-meta EXPONENTIAL = 'exponential' FIXED = 'fixed' @@ -149,7 +150,7 @@ def configure_retries(self, **kwargs): 'history': [] } - def increment(self, settings, error): + def increment(self, settings, error): # pylint: disable=no-self-use settings['total'] -= 1 settings['history'].append(error) if settings['total'] < 0: @@ -181,10 +182,10 @@ def get_backoff_time(self, settings, error): return min(settings['max_backoff'], backoff_value) -AMQPError = namedtuple('error', ['condition', 'description', 'info']) -AMQPError.__new__.__defaults__ = (None,) * len(AMQPError._fields) -AMQPError._code = 0x0000001d # pylint: disable=protected-access -AMQPError._definition = ( # pylint: disable=protected-access +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), @@ -278,7 +279,7 @@ def __init__(self, condition, description=None, info=None): 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(AMQPLinkError, self).__init__(condition, description=description, info=info) + super().__init__(condition, description=description, info=info) class AuthenticationException(AMQPException): @@ -300,9 +301,8 @@ class TokenExpired(AuthenticationException): 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 @@ -332,13 +332,12 @@ class MessageSendFailed(MessageException): :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") diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/link.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/link.py index 05a71c46a3c9..54a81e8fc989 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/link.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/link.py @@ -13,20 +13,17 @@ from .constants import DEFAULT_LINK_CREDIT, SessionState, LinkState, Role, SenderSettleMode, ReceiverSettleMode from .performatives import AttachFrame, DetachFrame - -from .error import ( - AMQPError, - ErrorCondition, - AMQPLinkError, - AMQPLinkRedirect, - AMQPConnectionError -) +from .error import ErrorCondition, AMQPLinkError, AMQPLinkRedirect, AMQPConnectionError _LOGGER = logging.getLogger(__name__) -class Link(object): - """ """ +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 @@ -100,8 +97,9 @@ def __exit__(self, *args): @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") # TODO: Assuming we establish all links for now... + raise NotImplementedError("Pending") def get_state(self): try: @@ -166,20 +164,11 @@ def _incoming_attach(self, frame): if self.network_trace: _LOGGER.info("<- %r", AttachFrame(*frame), extra=self.network_trace_params) if self._is_closed: - raise AMQPLinkError( - condition=ErrorCondition.ClientError, - description="Received attach frame on a link that is already closed", - info=None, - ) - elif not frame[5] or not frame[6]: + raise ValueError("Invalid link") + if not frame[5] or not frame[6]: _LOGGER.info("Cannot get source or target. Detaching link") - self.detach( - error=AMQPError( - condition=ErrorCondition.LinkDetachForced, - description="Cannot get source or target from the frame. Detaching link", - info=None, - ) - ) + 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 @@ -198,8 +187,8 @@ def _incoming_attach(self, frame): if frame[6]: frame[6] = Target(*frame[6]) self._on_attach(AttachFrame(*frame)) - except Exception as e: - _LOGGER.warning("Callback for link attach raised error: {}".format(e)) + except Exception as e: # pylint: disable=broad-except + _LOGGER.warning("Callback for link attach raised error: %r", e) def _outgoing_flow(self, **kwargs): flow_frame = { @@ -263,7 +252,7 @@ def detach(self, close=False, error=None): elif self.state == LinkState.ATTACHED: self._outgoing_detach(close=close, error=error) self._set_state(LinkState.DETACH_SENT) - except Exception as exc: + except Exception as exc: # pylint: disable=broad-except _LOGGER.info("An error occurred when detaching the link: %r", exc) self._set_state(LinkState.DETACHED) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/management_link.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/management_link.py index ac1b7bc08029..87290435af9b 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/management_link.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/management_link.py @@ -19,7 +19,8 @@ ManagementExecuteOperationResult, ManagementOpenResult, SEND_DISPOSITION_REJECT, - MessageDeliveryState + MessageDeliveryState, + LinkDeliverySettleReason ) from .error import AMQPException, ErrorCondition from .message import Properties, _MessageDelivery @@ -145,7 +146,7 @@ def _on_message_received(self, _, message): 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 SEND_DISPOSITION_REJECT in 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: diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/message.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/message.py index 890929e27582..c4bc6b0e1d19 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/message.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/message.py @@ -4,6 +4,7 @@ # license information. #-------------------------------------------------------------------------- +# TODO: fix mypy errors for _code/_definition/__defaults__ (issue #26500) from collections import namedtuple from .types import AMQPTypes, FieldDefinition @@ -12,7 +13,7 @@ Header = namedtuple( - 'header', + 'Header', [ 'durable', 'priority', @@ -20,9 +21,9 @@ 'first_acquirer', 'delivery_count' ]) -Header._code = 0x00000070 # pylint:disable=protected-access -Header.__new__.__defaults__ = (None,) * len(Header._fields) -Header._definition = ( # pylint:disable=protected-access +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), @@ -75,7 +76,7 @@ Properties = namedtuple( - 'properties', + 'Properties', [ 'message_id', 'user_id', @@ -91,9 +92,9 @@ 'group_sequence', 'reply_to_group_id' ]) -Properties._code = 0x00000073 # pylint:disable=protected-access -Properties.__new__.__defaults__ = (None,) * len(Properties._fields) -Properties._definition = ( # pylint:disable=protected-access +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), @@ -165,7 +166,7 @@ # TODO: should be a class, namedtuple or dataclass, immutability vs performance, need to collect performance data Message = namedtuple( - 'message', + 'Message', [ 'header', 'delivery_annotations', @@ -177,9 +178,9 @@ 'value', 'footer', ]) -Message.__new__.__defaults__ = (None,) * len(Message._fields) -Message._code = 0 # pylint:disable=protected-access -Message._definition = ( # pylint:disable=protected-access +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)), diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/outcomes.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/outcomes.py index 2056db2f1a38..64c5d09c7f66 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/outcomes.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/outcomes.py @@ -25,6 +25,7 @@ # - 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 @@ -32,9 +33,9 @@ from .performatives import _CAN_ADD_DOCSTRING -Received = namedtuple('received', ['section_number', 'section_offset']) -Received._code = 0x00000023 # pylint:disable=protected-access -Received._definition = ( # pylint:disable=protected-access +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: @@ -64,9 +65,9 @@ """ -Accepted = namedtuple('accepted', []) -Accepted._code = 0x00000024 # pylint:disable=protected-access -Accepted._definition = () # pylint:disable=protected-access +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. @@ -82,10 +83,10 @@ """ -Rejected = namedtuple('rejected', ['error']) -Rejected.__new__.__defaults__ = (None,) * len(Rejected._fields) -Rejected._code = 0x00000025 # pylint:disable=protected-access -Rejected._definition = (FIELD("error", ObjDefinition.error, False, None, False),) # pylint:disable=protected-access +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. @@ -102,9 +103,9 @@ """ -Released = namedtuple('released', []) -Released._code = 0x00000026 # pylint:disable=protected-access -Released._definition = () # pylint:disable=protected-access +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. @@ -123,10 +124,10 @@ """ -Modified = namedtuple('modified', ['delivery_failed', 'undeliverable_here', 'message_annotations']) -Modified.__new__.__defaults__ = (None,) * len(Modified._fields) -Modified._code = 0x00000027 # pylint:disable=protected-access -Modified._definition = ( # pylint:disable=protected-access +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)) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/performatives.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/performatives.py index 3280cde01f08..efcfc444ccd7 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/performatives.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/performatives.py @@ -4,6 +4,7 @@ # license information. #-------------------------------------------------------------------------- +# TODO: fix mypy errors for _code/_definition/__defaults__ (issue #26500) from collections import namedtuple import sys @@ -14,7 +15,7 @@ OpenFrame = namedtuple( - 'open', + 'OpenFrame', [ 'container_id', 'hostname', @@ -27,8 +28,8 @@ 'desired_capabilities', 'properties' ]) -OpenFrame._code = 0x00000010 # pylint:disable=protected-access -OpenFrame._definition = ( # pylint:disable=protected-access +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), @@ -103,7 +104,7 @@ BeginFrame = namedtuple( - 'begin', + 'BeginFrame', [ 'remote_channel', 'next_outgoing_id', @@ -114,8 +115,8 @@ 'desired_capabilities', 'properties' ]) -BeginFrame._code = 0x00000011 # pylint:disable=protected-access -BeginFrame._definition = ( # pylint:disable=protected-access +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), @@ -163,7 +164,7 @@ AttachFrame = namedtuple( - 'attach', + 'AttachFrame', [ 'name', 'handle', @@ -180,8 +181,8 @@ 'desired_capabilities', 'properties' ]) -AttachFrame._code = 0x00000012 # pylint:disable=protected-access -AttachFrame._definition = ( # pylint:disable=protected-access +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), @@ -262,7 +263,7 @@ FlowFrame = namedtuple( - 'flow', + 'FlowFrame', [ 'next_incoming_id', 'incoming_window', @@ -276,9 +277,9 @@ 'echo', 'properties' ]) -FlowFrame.__new__.__defaults__ = (None, None, None, None, None, None, None) -FlowFrame._code = 0x00000013 # pylint:disable=protected-access -FlowFrame._definition = ( # pylint:disable=protected-access +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), @@ -334,7 +335,7 @@ TransferFrame = namedtuple( - 'transfer', + 'TransferFrame', [ 'handle', 'delivery_id', @@ -349,8 +350,8 @@ 'batchable', 'payload' ]) -TransferFrame._code = 0x00000014 # pylint:disable=protected-access -TransferFrame._definition = ( # pylint:disable=protected-access +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), @@ -435,7 +436,7 @@ DispositionFrame = namedtuple( - 'disposition', + 'DispositionFrame', [ 'role', 'first', @@ -444,8 +445,8 @@ 'state', 'batchable' ]) -DispositionFrame._code = 0x00000015 # pylint:disable=protected-access -DispositionFrame._definition = ( # pylint:disable=protected-access +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), @@ -484,9 +485,9 @@ implementation uses when communicating delivery states, and thereby save bandwidth. """ -DetachFrame = namedtuple('detach', ['handle', 'closed', 'error']) -DetachFrame._code = 0x00000016 # pylint:disable=protected-access -DetachFrame._definition = ( # pylint:disable=protected-access +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)) @@ -505,9 +506,9 @@ """ -EndFrame = namedtuple('end', ['error']) -EndFrame._code = 0x00000017 # pylint:disable=protected-access -EndFrame._definition = (FIELD("error", ObjDefinition.error, False, None, False),) # pylint:disable=protected-access +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. @@ -520,9 +521,9 @@ """ -CloseFrame = namedtuple('close', ['error']) -CloseFrame._code = 0x00000018 # pylint:disable=protected-access -CloseFrame._definition = (FIELD("error", ObjDefinition.error, False, None, False),) # pylint:disable=protected-access +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. @@ -537,9 +538,9 @@ """ -SASLMechanism = namedtuple('sasl_mechanism', ['sasl_server_mechanisms']) -SASLMechanism._code = 0x00000040 # pylint:disable=protected-access -SASLMechanism._definition = (FIELD('sasl_server_mechanisms', AMQPTypes.symbol, True, None, True),) # pylint:disable=protected-access +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. @@ -554,9 +555,9 @@ """ -SASLInit = namedtuple('sasl_init', ['mechanism', 'initial_response', 'hostname']) -SASLInit._code = 0x00000041 # pylint:disable=protected-access -SASLInit._definition = ( # pylint:disable=protected-access +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)) @@ -585,9 +586,9 @@ """ -SASLChallenge = namedtuple('sasl_challenge', ['challenge']) -SASLChallenge._code = 0x00000042 # pylint:disable=protected-access -SASLChallenge._definition = (FIELD('challenge', AMQPTypes.binary, True, None, False),) # pylint:disable=protected-access +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. @@ -599,9 +600,9 @@ """ -SASLResponse = namedtuple('sasl_response', ['response']) -SASLResponse._code = 0x00000043 # pylint:disable=protected-access -SASLResponse._definition = (FIELD('response', AMQPTypes.binary, True, None, False),) # pylint:disable=protected-access +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. @@ -612,9 +613,9 @@ """ -SASLOutcome = namedtuple('sasl_outcome', ['code', 'additional_data']) -SASLOutcome._code = 0x00000044 # pylint:disable=protected-access -SASLOutcome._definition = ( # pylint:disable=protected-access +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: diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/receiver.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/receiver.py index 2923ddaebc19..a7abe9c1536a 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/receiver.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/receiver.py @@ -28,10 +28,16 @@ def __init__(self, session, handle, 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: + except Exception as e: # pylint: disable=broad-except _LOGGER.error("Handler function failed with error: %r", e) return None @@ -64,7 +70,13 @@ def _incoming_transfer(self, frame): _LOGGER.info(" %r", message, extra=self.network_trace_params) delivery_state = self._process_incoming_message(frame, message) if not frame[4] and delivery_state: # settled - self._outgoing_disposition(first=frame[1], settled=True, state=delivery_state) + 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: diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/sasl.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/sasl.py index b927cf5b627a..c4ff9d265540 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/sasl.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/sasl.py @@ -1,15 +1,15 @@ -#------------------------------------------------------------------------- +# ------------------------------------------------------------------------- # 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' +_SASL_FRAME_TYPE = b"\x01" class SASLPlainCredential(object): @@ -17,7 +17,7 @@ class SASLPlainCredential(object): See https://tools.ietf.org/html/rfc4616 for details """ - mechanism = b'PLAIN' + mechanism = b"PLAIN" def __init__(self, authcid, passwd, authzid=None): self.authcid = authcid @@ -26,13 +26,13 @@ def __init__(self, authcid, passwd, authzid=None): def start(self): if self.authzid: - login_response = self.authzid.encode('utf-8') + 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') + 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 @@ -41,10 +41,10 @@ class SASLAnonymousCredential(object): See https://tools.ietf.org/html/rfc4505 for details """ - mechanism = b'ANONYMOUS' + mechanism = b"ANONYMOUS" - def start(self): - return b'' + def start(self): # pylint: disable=no-self-use + return b"" class SASLExternalCredential(object): @@ -54,27 +54,34 @@ class SASLExternalCredential(object): authentication data. """ - mechanism = b'EXTERNAL' + mechanism = b"EXTERNAL" - def start(self): - return b'' + def start(self): # pylint: disable=no-self-use + return b"" -class SASLTransportMixin(): +class SASLTransportMixin: def _negotiate(self): self.write(SASL_HEADER_FRAME) _, returned_header = self.receive_frame() if returned_header[1] != SASL_HEADER_FRAME: - raise ValueError("Mismatching AMQP header protocol. Expected: {}, received: {}".format( - SASL_HEADER_FRAME, returned_header[1])) + 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)) + 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) + hostname=self.host, + ) self.send_frame(0, sasl_init, frame_type=_SASL_FRAME_TYPE) _, next_frame = self.receive_frame(verify_frame_type=1) @@ -83,33 +90,57 @@ def _negotiate(self): raise NotImplementedError("Unsupported SASL challenge") if fields[0] == SASLCode.Ok: # code return - raise ValueError("SASL negotiation failed.\nOutcome: {}\nDetails: {}".format(*fields)) + 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=None, **kwargs): + def __init__( + self, + host, + credential, + *, + port=AMQPS_PORT, + connect_timeout=None, + ssl_opts=None, + **kwargs, + ): self.credential = credential - ssl = ssl or True - super(SASLTransport, self).__init__(host, port=port, connect_timeout=connect_timeout, ssl=ssl, **kwargs) + 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=None, **kwargs): +class SASLWithWebSocket(WebSocketTransport, SASLTransportMixin): + def __init__( + self, + host, + credential, + *, + port=WEBSOCKET_PORT, + connect_timeout=None, + ssl_opts=None, + **kwargs, + ): self.credential = credential - ssl = ssl or True - self._transport = WebSocketTransport( + ssl_opts = ssl_opts or True + super().__init__( host, port=port, connect_timeout=connect_timeout, - ssl=ssl, - **kwargs + ssl_opts=ssl_opts, + **kwargs, ) - super().__init__(host, port, connect_timeout, ssl, **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 index e91636d7cff9..c0e9e64cd6e2 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/sender.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/sender.py @@ -48,6 +48,12 @@ def __init__(self, session, handle, target_address, **kwargs): 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): diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/session.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/session.py index 0e5546443ae4..ea36c5b1be1c 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/session.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/session.py @@ -4,30 +4,30 @@ # license information. # -------------------------------------------------------------------------- +from __future__ import annotations import uuid import logging import time +from typing import Union, Optional -from .constants import ( - INCOMING_WINDOW, - OUTGOING_WINDOW, - ConnectionState, - SessionState, - SessionTransferState, - Role -) -from .endpoints import Source, Target +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 .performatives import ( + BeginFrame, + EndFrame, + FlowFrame, + TransferFrame, + DispositionFrame, +) +from .error import AMQPError, ErrorCondition from ._encode import encode_frame -from azure.eventhub._pyamqp.error import AMQPError, AMQPSessionError, ErrorCondition _LOGGER = logging.getLogger(__name__) -class Session(object): +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. @@ -75,8 +75,8 @@ def __exit__(self, *args): self.end() @classmethod - def from_incoming_frame(cls, connection, channel, frame): - # check session_create_from_endpoint in C lib + def from_incoming_frame(cls, connection, channel): + # TODO: check session_create_from_endpoint in C lib new_session = cls(connection, channel) return new_session @@ -87,9 +87,14 @@ def _set_state(self, new_state): 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) + _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 + 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]: @@ -105,24 +110,38 @@ def _get_next_output_handle(self): :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) + 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, + 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, + 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.info("-> %r", begin_frame, extra=self.network_trace_params) - self._connection._process_outgoing_frame(self.channel, begin_frame) # pylint: disable=protected-access + self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, begin_frame + ) def _incoming_begin(self, frame): if self.network_trace: @@ -143,12 +162,18 @@ def _outgoing_end(self, error=None): end_frame = EndFrame(error=error) if self.network_trace: _LOGGER.info("-> %r", end_frame, extra=self.network_trace_params) - self._connection._process_outgoing_frame(self.channel, end_frame) # pylint: disable=protected-access + self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, end_frame + ) def _incoming_end(self, frame): if self.network_trace: _LOGGER.info("<- %r", EndFrame(*frame), extra=self.network_trace_params) - if self.state not in [SessionState.END_RCVD, SessionState.END_SENT, SessionState.DISCARDING]: + 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() @@ -157,21 +182,28 @@ def _incoming_end(self, frame): self._set_state(SessionState.UNMAPPED) def _outgoing_attach(self, frame): - self._connection._process_outgoing_frame(self.channel, frame) # pylint: disable=protected-access + 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(frame) # pylint: disable=protected-access + 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: # detach the link that would have been set. - self.links[frame[0].decode('utf-8')].detach( + 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( + description="""Cannot allocate more handles, """ + """the max number of handles is {}. Detaching link""".format( self.handle_max ), info=None, @@ -179,10 +211,12 @@ def _incoming_attach(self, frame): ) return if frame[2] == Role.Sender: # role - new_link = ReceiverLink.from_incoming_frame(self, outgoing_handle, frame) + 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 + 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 @@ -203,21 +237,31 @@ def _outgoing_flow(self, frame=None): flow_frame = FlowFrame(**link_flow) if self.network_trace: _LOGGER.info("-> %r", flow_frame, extra=self.network_trace_params) - self._connection._process_outgoing_frame(self.channel, flow_frame) # pylint: disable=protected-access + self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, flow_frame + ) def _incoming_flow(self, frame): if self.network_trace: _LOGGER.info("<- %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 + 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(frame) # pylint: disable=protected-access + 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 + 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): if self.state != SessionState.MAPPED: @@ -237,7 +281,9 @@ def _outgoing_transfer(self, delivery): # 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 + 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 @@ -257,7 +303,9 @@ def _outgoing_transfer(self, delivery): "payload": payload[start_idx : start_idx + available_frame_size], "delivery_id": self.next_outgoing_id, } - self._connection._process_outgoing_frame(self.channel, TransferFrame(**tmp_delivery_frame)) # pylint: disable=protected-access + self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, TransferFrame(**tmp_delivery_frame) + ) start_idx += available_frame_size remaining_payload_cnt -= available_frame_size @@ -276,7 +324,9 @@ def _outgoing_transfer(self, delivery): "payload": payload[start_idx:], "delivery_id": self.next_outgoing_id, } - self._connection._process_outgoing_frame(self.channel, TransferFrame(**tmp_delivery_frame)) # pylint: disable=protected-access + self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, TransferFrame(**tmp_delivery_frame) + ) self.next_outgoing_id += 1 self.remote_incoming_window -= 1 self.outgoing_window -= 1 @@ -288,41 +338,57 @@ def _incoming_transfer(self, frame): self.remote_outgoing_window -= 1 self.incoming_window -= 1 try: - self._input_handles[frame[0]]._incoming_transfer(frame) # pylint: disable=protected-access + self._input_handles[frame[0]]._incoming_transfer( # pylint: disable=protected-access + frame + ) except KeyError: self._set_state(SessionState.DISCARDING) - self.end(error=AMQPError( + self.end( + error=AMQPError( condition=ErrorCondition.SessionUnattachedHandle, - description="Invalid handle reference in received frame: Handle is not currently associated with an attached link")) + 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 self._outgoing_flow() def _outgoing_disposition(self, frame): - self._connection._process_outgoing_frame(self.channel, frame) # pylint: disable=protected-access + self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, frame + ) def _incoming_disposition(self, frame): if self.network_trace: - _LOGGER.info("<- %r", DispositionFrame(*frame), extra=self.network_trace_params) + _LOGGER.info( + "<- %r", DispositionFrame(*frame), extra=self.network_trace_params + ) for link in self._input_handles.values(): - link._incoming_disposition(frame) # pylint: disable=protected-access + link._incoming_disposition(frame) # pylint: disable=protected-access def _outgoing_detach(self, frame): - self._connection._process_outgoing_frame(self.channel, frame) # pylint: disable=protected-access + 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 + 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")) + 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 @@ -346,7 +412,9 @@ def begin(self, wait=False): 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.") + 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 @@ -358,7 +426,7 @@ def end(self, error=None, wait=False): 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: + except Exception as exc: # pylint: disable=broad-except _LOGGER.info("An error occurred when ending the session: %r", exc) self._set_state(SessionState.UNMAPPED) @@ -370,7 +438,7 @@ def create_receiver_link(self, source_address, **kwargs): source_address=source_address, network_trace=kwargs.pop("network_trace", self.network_trace), network_trace_params=dict(self.network_trace_params), - **kwargs + **kwargs, ) self.links[link.name] = link self._output_handles[assigned_handle] = link @@ -384,11 +452,16 @@ def create_sender_link(self, target_address, **kwargs): target_address=target_address, network_trace=kwargs.pop("network_trace", self.network_trace), network_trace_params=dict(self.network_trace_params), - **kwargs + **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), **kwargs) + return ManagementLink( + self, + endpoint, + network_trace=kwargs.pop("network_trace", self.network_trace), + **kwargs, + ) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_transport/_base.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_transport/_base.py index d67cceedcd40..4981f6fce643 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_transport/_base.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_transport/_base.py @@ -163,10 +163,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 index 205a3b3419d7..074cb3e7d11f 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_transport/_pyamqp_transport.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_transport/_pyamqp_transport.py @@ -5,7 +5,7 @@ import logging import time -from typing import Optional, Union, Any, Tuple +from typing import Optional, Union, Any, Tuple, cast from .._pyamqp import ( error as errors, @@ -38,14 +38,16 @@ _LOGGER = logging.getLogger(__name__) -class PyamqpTransport(AmqpTransport): +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 + MAX_MESSAGE_LENGTH_BYTES = ( + constants.MAX_FRAME_SIZE_BYTES + ) # TODO: define actual value in pyamqp TIMEOUT_FACTOR = 1 CONNECTION_CLOSING_STATES: Tuple = _CLOSING_STATES @@ -104,7 +106,9 @@ def to_outgoing_amqp_message(annotated_message): 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) + 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, @@ -175,7 +179,10 @@ def create_link_properties(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()} + return { + symbol: utils.amqp_long_value(value) + for (symbol, value) in link_properties.items() + } @staticmethod def create_connection(**kwargs): @@ -234,8 +241,10 @@ def create_send_client(*, config, **kwargs): # pylint:disable=unused-argument """ 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? + # TODO: not used by pyamqp? + msg_timeout = kwargs.pop( # pylint: disable=unused-variable + "msg_timeout" + ) return SendClient( config.hostname, @@ -288,16 +297,20 @@ def set_message_partition_key(message, partition_key, **kwargs): if annotations is None: annotations = {} try: - partition_key = partition_key.decode(encoding) + 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) + 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 + 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. @@ -305,8 +318,13 @@ def add_batch(event_data_batch, outgoing_event_data, event_data): # pylint: dis :param event_data: EventData to add to internal batch events. uamqp use only. :rtype: None """ - event_data_batch._internal_events.append(event_data) # pylint: disable=protected-access - utils.add_batch(event_data_batch._message, outgoing_event_data._message) # pylint: disable=protected-access + 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): @@ -369,7 +387,11 @@ def open_receive_client(*, handler, client, auth): :rtype: bool """ # pylint:disable=protected-access - handler.open(connection=client._conn_manager.get_connection(client._address.hostname, auth)) + handler.open( + connection=client._conn_manager.get_connection( + client._address.hostname, auth + ) + ) @staticmethod def check_link_stolen(consumer, exception): @@ -379,8 +401,13 @@ def check_link_stolen(consumer, exception): :param exception: Exception to check. """ - if isinstance(exception, errors.AMQPLinkError) and exception.condition == errors.ErrorCondition.LinkStolen: - raise consumer._handle_exception(exception) # pylint: disable=protected-access + 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): @@ -414,7 +441,9 @@ def create_token_auth(auth_uri, get_token, token_type, config, **kwargs): # 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 + 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. @@ -454,7 +483,10 @@ def mgmt_client_request(mgmt_client, mgmt_msg, **kwargs): 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 + mgmt_msg, + operation=operation.decode(), + operation_type=operation_type.decode(), + **kwargs, ) @staticmethod @@ -468,7 +500,8 @@ def get_error(status_code, description): if status_code in [401]: return errors.AuthenticationException( errors.ErrorCondition.UnauthorizedAccess, - description=f"Management authentication failed. Status code: {status_code}, Description: {description!r}", + description=f"""Management authentication failed. Status code: {status_code}, """ + """Description: {description!r}""", ) if status_code in [404]: return errors.AMQPConnectionError( @@ -489,7 +522,8 @@ def check_timeout_exception(base, exception): """ if not base.running and isinstance(exception, TimeoutError): exception = errors.AuthenticationException( - errors.ErrorCondition.InternalError, description="Authorization timeout." + errors.ErrorCondition.InternalError, + description="Authorization timeout.", ) return exception @@ -512,7 +546,9 @@ def _create_eventhub_exception(exception): return error @staticmethod - def _handle_exception(exception, closable): # pylint:disable=too-many-branches, too-many-statements + def _handle_exception( + exception, closable + ): # 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 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 478906f27a38..fbc8f1d6af73 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_transport/_uamqp_transport.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_transport/_uamqp_transport.py @@ -349,17 +349,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 ) 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 6f8533a64adb..4cd1bb9f5b78 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 @@ -37,7 +37,7 @@ try: from ._transport._uamqp_transport_async import UamqpTransportAsync except ImportError: - UamqpTransportAsync = None + UamqpTransportAsync = None # type: ignore from ._transport._pyamqp_transport_async import PyamqpTransportAsync if TYPE_CHECKING: 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 e1be06565dca..57f544b2ac2f 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 @@ -13,6 +13,7 @@ if TYPE_CHECKING: from .._pyamqp.aio._authentication_async import JWTTokenAuthAsync from .._pyamqp.aio._connection_async import Connection as ConnectionAsync + from ._transport._base_async import AmqpTransportAsync from uamqp.authentication import JWTTokenAsync as uamqp_JWTTokenAuthAsync from uamqp.async_ops import ConnectionAsync as uamqp_ConnectionAsync @@ -56,7 +57,7 @@ def __init__(self, **kwargs) -> None: 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") + self._amqp_transport: AmqpTransportAsync = kwargs.pop("amqp_transport") async def get_connection( self, 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 c1e4fec6d275..c4e311082e68 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_consumer_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_consumer_async.py @@ -20,7 +20,10 @@ from typing import Deque try: - from uamqp import ReceiveClientAsync as uamqp_ReceiveClientAsync, Message as uamqp_Message + 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: @@ -38,7 +41,9 @@ _LOGGER = logging.getLogger(__name__) -class EventHubConsumer(ConsumerProducerMixin): # pylint:disable=too-many-instance-attributes +class EventHubConsumer( + ConsumerProducerMixin +): # pylint:disable=too-many-instance-attributes """ A consumer responsible for reading EventData from a specific Event Hub partition and as a member of a specific consumer group. @@ -79,16 +84,18 @@ def __init__(self, client: "EventHubConsumerClient", source: str, **kwargs) -> N owner_level = kwargs.get("owner_level", None) keep_alive = kwargs.get("keep_alive", None) auto_reconnect = kwargs.get("auto_reconnect", True) - track_last_enqueued_event_properties = kwargs.get("track_last_enqueued_event_properties", False) + track_last_enqueued_event_properties = kwargs.get( + "track_last_enqueued_event_properties", False + ) idle_timeout = kwargs.get("idle_timeout", None) self.running = False 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 @@ -98,33 +105,52 @@ 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: Union[Dict[uamqp_AMQPType, uamqp_AMQPType], Dict[types.AMQPTypes, types.AMQPTypes]] = {} + 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: 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._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._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._message_buffer_lock = asyncio.Lock() self._last_callback_called_time = None self._callback_task_run = None - def _create_handler(self, auth: Union[uamqp_JWTTokenAsync, JWTTokenAuthAsync]) -> 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) + self._source, + self._offset, + 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 @@ -138,11 +164,14 @@ def _create_handler(self, auth: Union[uamqp_JWTTokenAsync, JWTTokenAuthAsync]) - 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=partial(self._amqp_transport.message_received_async, self), + message_received_callback=partial( + self._amqp_transport.message_received_async, self + ), ) async def _open_with_retry(self) -> None: @@ -155,5 +184,9 @@ def _next_message_in_buffer(self): self._last_received_event = event_data return event_data - async def receive(self, batch=False, max_batch_size=300, max_wait_time=None) -> None: - await self._amqp_transport.receive_messages_async(self, batch, max_batch_size, max_wait_time) + async def receive( + self, batch=False, max_batch_size=300, max_wait_time=None + ) -> None: + 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 73a91416b1a4..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 @@ -207,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..6413b80bff48 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, 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 a5bb2ee2c68f..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 @@ -25,7 +25,10 @@ if TYPE_CHECKING: try: - from uamqp import constants, SendClientAsync as uamqp_SendClientAsync + 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: @@ -40,7 +43,9 @@ _LOGGER = logging.getLogger(__name__) -class EventHubProducer(ConsumerProducerMixin): # pylint: disable=too-many-instance-attributes +class EventHubProducer( + ConsumerProducerMixin +): # pylint: disable=too-many-instance-attributes """A producer responsible for transmitting batches of EventData to a specific Event Hub. Depending on the options specified at creation, the producer may @@ -83,9 +88,15 @@ def __init__(self, client: EventHubProducerClient, target: str, **kwargs) -> Non self._keep_alive = keep_alive self._auto_reconnect = auto_reconnect self._timeout = send_timeout - 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._retry_policy = self._amqp_transport.create_retry_policy(config=self._client._config) + self._retry_policy = self._amqp_transport.create_retry_policy( + config=self._client._config + ) self._reconnect_backoff = 1 self._name = "EHProducer-{}".format(uuid.uuid4()) self._unsent_events = [] # type: List[Any] @@ -101,7 +112,9 @@ def __init__(self, client: EventHubProducerClient, target: str, **kwargs) -> Non {TIMEOUT_SYMBOL: int(self._timeout * self._amqp_transport.TIMEOUT_FACTOR)} ) - def _create_handler(self, auth: Union[uamqp_JWTTokenAsync, JWTTokenAuthAsync]) -> 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, @@ -120,7 +133,9 @@ def _create_handler(self, auth: Union[uamqp_JWTTokenAsync, JWTTokenAuthAsync]) - ) async def _open_with_retry(self) -> Any: - return await self._do_retryable_operation(self._open, operation_need_param=False) + return await self._do_retryable_operation( + self._open, operation_need_param=False + ) async def _send_event_data( self, @@ -128,12 +143,18 @@ async def _send_event_data( last_exception: Optional[Exception] = None, ) -> None: if self._unsent_events: - await self._amqp_transport.send_messages_async(self, timeout_time, last_exception, _LOGGER) + await self._amqp_transport.send_messages_async( + self, timeout_time, last_exception, _LOGGER + ) - async def _send_event_data_with_retry(self, timeout: Optional[float] = None) -> None: + async def _send_event_data_with_retry( + self, timeout: Optional[float] = None + ) -> None: await self._do_retryable_operation(self._send_event_data, timeout=timeout) - def _on_outcome(self, outcome: uamqp_MessageSendResult, condition: Optional[Exception]) -> None: + def _on_outcome( + self, outcome: uamqp_MessageSendResult, condition: Optional[Exception] + ) -> None: """ ONLY USED FOR uamqp_transport=True. Called when the outcome is received for a delivery. @@ -147,7 +168,9 @@ def _on_outcome(self, outcome: uamqp_MessageSendResult, condition: Optional[Exce def _wrap_eventdata( self, - event_data: Union[EventData, AmqpAnnotatedMessage, EventDataBatch, Iterable[EventData]], + event_data: Union[ + EventData, AmqpAnnotatedMessage, EventDataBatch, Iterable[EventData] + ], span: Optional[AbstractSpan], partition_key: Optional[AnyStr], ) -> Union[EventData, EventDataBatch]: @@ -157,32 +180,46 @@ 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) else: - if isinstance(event_data, EventDataBatch): # The partition_key in the param will be omitted. + if isinstance( + event_data, EventDataBatch + ): # 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 + 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=event_data._partition_key, + 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 - raise ValueError("The partition_key does not match the one of the EventDataBatch") - for event in event_data._message.data: # pylint: disable=protected-access + if ( + 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" + ) + for ( + event + ) in event_data._message.data: # pylint: disable=protected-access trace_message(event, span) wrapper_event_data = event_data # type:ignore else: if partition_key: - event_data = _set_partition_key(event_data, partition_key, self._amqp_transport) + event_data = _set_partition_key( + 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 event_data, self._amqp_transport, partition_key @@ -191,7 +228,9 @@ def _wrap_eventdata( async def send( self, - event_data: Union[EventData, AmqpAnnotatedMessage, EventDataBatch, Iterable[EventData]], + event_data: Union[ + EventData, AmqpAnnotatedMessage, EventDataBatch, Iterable[EventData] + ], *, partition_key: Optional[AnyStr] = None, timeout: Optional[float] = None, @@ -223,17 +262,25 @@ async def send( async with self._lock: with send_context_manager() as child: self._check_closed() - wrapper_event_data = self._wrap_eventdata(event_data, child, partition_key) + wrapper_event_data = self._wrap_eventdata( + event_data, child, partition_key + ) 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(child) # pylint: disable=protected-access + self._client._add_span_request_attributes( # pylint: disable=protected-access + child + ) - await self._send_event_data_with_retry(timeout=timeout) # pylint:disable=unexpected-keyword-arg + await self._send_event_data_with_retry( + timeout=timeout + ) # pylint:disable=unexpected-keyword-arg async def close(self) -> None: """ 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 098833c43f7e..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 @@ -347,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, 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 81c654c88019..d481a9b71f44 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,8 +3,9 @@ # 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 @@ -20,12 +21,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 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 index adca8c81b736..eff4aa03d21a 100644 --- 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 @@ -58,7 +58,7 @@ async def create_connection_async(**kwargs): return ConnectionAsync(endpoint, network_trace=network_trace, **kwargs) @staticmethod - async def close_connection(connection): + async def close_connection_async(connection): """ Closes existing connection. :param connection: pyamqp Connection. @@ -344,7 +344,7 @@ async def _handle_exception_async( # pylint:disable=too-many-branches, too-many if isinstance(exception, errors.AuthenticationException): await closable._close_connection_async() # pylint:disable=protected-access elif isinstance(exception, errors.AMQPLinkError): - await closable._close_handler_async() # pylint:disable=protected-access + 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? diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_base_handler.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_base_handler.py index b97d2175a7f6..51fcf259ba75 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_base_handler.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_base_handler.py @@ -18,7 +18,7 @@ from ._pyamqp.utils import generate_sas_token, amqp_string_value from ._pyamqp.message import Message, Properties -from ._pyamqp.client import AMQPClientSync +from ._pyamqp.client import AMQPClient as AMQPClientSync from ._common._configuration import Configuration from .exceptions import ( diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/__init__.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/__init__.py index 4795dde0e65a..fc9544449266 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/__init__.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/__init__.py @@ -10,12 +10,12 @@ from ._connection import Connection from ._transport import SSLTransport -from .client import AMQPClientSync, ReceiveClientSync, SendClientSync +from .client import AMQPClient, ReceiveClient, SendClient __all__ = [ "Connection", "SSLTransport", - "AMQPClientSync", - "ReceiveClientSync", - "SendClientSync", + "AMQPClient", + "ReceiveClient", + "SendClient", ] diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/_connection.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/_connection.py index 207cca0cde39..a5ed80dea995 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/_connection.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/_connection.py @@ -103,9 +103,7 @@ def __init__(self, endpoint, **kwargs): # pylint:disable=too-many-statements if custom_endpoint_address: custom_parsed_url = urlparse(custom_endpoint_address) custom_port = custom_parsed_url.port or WEBSOCKET_PORT - custom_endpoint = "{}:{}{}".format( - custom_parsed_url.hostname, custom_port, custom_parsed_url.path - ) + custom_endpoint = "{}:{}{}".format(custom_parsed_url.hostname, custom_port, custom_parsed_url.path) transport = kwargs.get("transport") self._transport_type = kwargs.pop("transport_type", TransportType.Amqp) @@ -113,62 +111,35 @@ def __init__(self, endpoint, **kwargs): # pylint:disable=too-many-statements self._transport = transport elif "sasl_credential" in kwargs: sasl_transport = SASLTransport - if self._transport_type.name == "AmqpOverWebsocket" or kwargs.get( - "http_proxy" - ): + 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, - **kwargs, + host=endpoint, credential=kwargs["sasl_credential"], custom_endpoint=custom_endpoint, **kwargs ) else: - self._transport = Transport( - parsed_url.netloc, transport_type=self._transport_type, **kwargs - ) + self._transport = Transport(parsed_url.netloc, transport_type=self._transport_type, **kwargs) - self._container_id = kwargs.pop("container_id", None) or str( - uuid.uuid4() - ) # type: str - self._max_frame_size = kwargs.pop( - "max_frame_size", MAX_FRAME_SIZE_BYTES - ) # type: int + self._container_id = kwargs.pop("container_id", None) or str(uuid.uuid4()) # type: str + 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._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._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._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._network_trace = kwargs.get("network_trace", False) - self._network_trace_params = { - "connection": self._container_id, - "session": None, - "link": None, - } + self._network_trace_params = {"connection": self._container_id, "session": None, "link": None} self._error = None self._outgoing_endpoints = {} # type: Dict[int, Session] self._incoming_endpoints = {} # type: Dict[int, Session] @@ -187,12 +158,7 @@ def _set_state(self, new_state): return previous_state = self.state self.state = new_state - _LOGGER.info( - "Connection '%s' state changed: %r -> %r", - self._container_id, - previous_state, - new_state, - ) + _LOGGER.info("Connection '%s' state changed: %r -> %r", self._container_id, previous_state, new_state) for session in self._outgoing_endpoints.values(): session._on_connection_state_change() # pylint:disable=protected-access @@ -218,16 +184,13 @@ def _connect(self): self._process_incoming_frame(*self._read_frame(wait=True)) # type: ignore if self.state != ConnectionState.HDR_EXCH: self._disconnect() - raise ValueError( - "Did not receive reciprocal protocol header. Disconnecting." - ) + raise ValueError("Did not receive reciprocal protocol header. Disconnecting.") else: self._set_state(ConnectionState.HDR_SENT) except (OSError, IOError, SSLError, socket.error) as exc: raise AMQPConnectionError( ErrorCondition.SocketError, - description="Failed to initiate the connection due to exception: " - + str(exc), + description="Failed to initiate the connection due to exception: " + str(exc), error=exc, ) except Exception: # pylint:disable=try-except-raise @@ -316,17 +279,9 @@ def _get_next_outgoing_channel(self): :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 - ) + 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): @@ -356,9 +311,7 @@ def _outgoing_header(self): """Send the AMQP protocol header to initiate the connection.""" self._last_frame_sent_time = time.time() if self._network_trace: - _LOGGER.info( - "-> header(%r)", HEADER_FRAME, extra=self._network_trace_params - ) + _LOGGER.info("-> header(%r)", HEADER_FRAME, extra=self._network_trace_params) self._transport.write(HEADER_FRAME) def _incoming_header(self, _, frame): @@ -381,17 +334,11 @@ def _outgoing_open(self): 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 + 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, + 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: @@ -427,8 +374,7 @@ def _incoming_open(self, channel, frame): _LOGGER.error("OPEN frame received on a channel that is not 0.") self.close( error=AMQPError( - condition=ErrorCondition.NotAllowed, - description="OPEN frame received on a channel that is not 0.", + condition=ErrorCondition.NotAllowed, description="OPEN frame received on a channel that is not 0." ) ) self._set_state(ConnectionState.END) @@ -508,11 +454,7 @@ def _incoming_close(self, channel, frame): close_error = None if channel > self._channel_max: _LOGGER.error("Invalid channel") - close_error = AMQPError( - condition=ErrorCondition.InvalidField, - description="Invalid channel", - info=None, - ) + close_error = AMQPError(condition=ErrorCondition.InvalidField, description="Invalid channel", info=None) self._set_state(ConnectionState.CLOSE_RCVD) self._outgoing_close(error=close_error) @@ -527,6 +469,7 @@ def _incoming_close(self, channel, frame): "Connection error: {}".format(frame[0]) # pylint:disable=logging-format-interpolation ) + def _incoming_begin(self, channel, frame): # type: (int, Tuple[Any, ...]) -> None """Process incoming Begin frame to finish negotiating a new session. @@ -572,23 +515,21 @@ def _incoming_end(self, channel, frame): :rtype: None """ try: - self._incoming_endpoints[channel]._incoming_end( # pylint:disable=protected-access - frame - ) + self._incoming_endpoints[channel]._incoming_end(frame) # pylint:disable=protected-access self._incoming_endpoints.pop(channel) self._outgoing_endpoints.pop(channel) except KeyError: - end_error = AMQPError( - condition=ErrorCondition.InvalidField, - description=f"Invalid channel {channel}", - info=None, - ) - _LOGGER.error("Received END frame with invalid channel %s", channel) - self.close(error=end_error) + #close the connection + self.close( + error=AMQPError( + condition=ErrorCondition.ConnectionCloseForced, + description="Invalid channel number received" + )) + return + self._incoming_endpoints.pop(channel) + self._outgoing_endpoints.pop(channel) - def _process_incoming_frame( - self, channel, frame - ): # pylint:disable=too-many-return-statements + 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. @@ -649,7 +590,7 @@ def _process_incoming_frame( self._incoming_header(channel, cast(bytes, fields)) return True if performative == 1: - return False # TODO: incoming EMPTY + return False _LOGGER.error("Unrecognized incoming frame: %s", frame) return True except KeyError: @@ -679,7 +620,6 @@ def _process_outgoing_frame(self, channel, frame): cast(float, self._last_frame_received_time), ) or self._get_remote_timeout(now): self.close( - # TODO: check error condition error=AMQPError( condition=ErrorCondition.ConnectionCloseForced, description="No frame received for the idle timeout.", @@ -758,8 +698,7 @@ def listen(self, wait=False, batch=1, **kwargs): cast(float, self._last_frame_received_time), ) or self._get_remote_timeout( now - ): # pylint:disable=line-too-long - # TODO: check error condition + ): self.close( error=AMQPError( condition=ErrorCondition.ConnectionCloseForced, @@ -769,10 +708,8 @@ def listen(self, wait=False, batch=1, **kwargs): ) return if self.state == ConnectionState.END: - # TODO: check error condition self._error = AMQPConnectionError( - condition=ErrorCondition.ConnectionCloseForced, - description="Connection was already closed.", + condition=ErrorCondition.ConnectionCloseForced, description="Connection was already closed." ) return for _ in range(batch): @@ -845,6 +782,7 @@ def open(self, wait=False): "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. diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/_encode.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/_encode.py index e8c952c34f0e..4e6a86c6dd4b 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/_encode.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/_encode.py @@ -83,8 +83,8 @@ def encode_null(output, *args, **kwargs): # pylint: disable=unused-argument def encode_boolean( - output, value, with_constructor=True, **kwargs -): # pylint: disable=unused-argument + output, value, with_constructor=True, **kwargs # pylint: disable=unused-argument +): # type: (bytearray, bool, bool, Any) -> None """ @@ -102,8 +102,8 @@ def encode_boolean( def encode_ubyte( - output, value, with_constructor=True, **kwargs -): # pylint: disable=unused-argument + output, value, with_constructor=True, **kwargs # pylint: disable=unused-argument +): # type: (bytearray, Union[int, bytes], bool, Any) -> None """ @@ -121,8 +121,8 @@ def encode_ubyte( def encode_ushort( - output, value, with_constructor=True, **kwargs -): # pylint: disable=unused-argument + output, value, with_constructor=True, **kwargs # pylint: disable=unused-argument +): # type: (bytearray, int, bool, Any) -> None """ @@ -182,8 +182,8 @@ def encode_ulong(output, value, with_constructor=True, use_smallest=True): def encode_byte( - output, value, with_constructor=True, **kwargs -): # pylint: disable=unused-argument + output, value, with_constructor=True, **kwargs # pylint: disable=unused-argument +): # type: (bytearray, int, bool, Any) -> None """ @@ -197,8 +197,8 @@ def encode_byte( def encode_short( - output, value, with_constructor=True, **kwargs -): # pylint: disable=unused-argument + output, value, with_constructor=True, **kwargs # pylint: disable=unused-argument +): # type: (bytearray, int, bool, Any) -> None """ @@ -252,8 +252,8 @@ def encode_long(output, value, with_constructor=True, use_smallest=True): def encode_float( - output, value, with_constructor=True, **kwargs -): # pylint: disable=unused-argument + output, value, with_constructor=True, **kwargs # pylint: disable=unused-argument +): # type: (bytearray, float, bool, Any) -> None """ @@ -264,8 +264,8 @@ def encode_float( def encode_double( - output, value, with_constructor=True, **kwargs -): # pylint: disable=unused-argument + output, value, with_constructor=True, **kwargs # pylint: disable=unused-argument +): # type: (bytearray, float, bool, Any) -> None """ @@ -276,8 +276,8 @@ def encode_double( def encode_timestamp( - output, value, with_constructor=True, **kwargs -): # pylint: disable=unused-argument + output, value, with_constructor=True, **kwargs # pylint: disable=unused-argument +): # type: (bytearray, Union[int, datetime], bool, Any) -> None """ None """ @@ -413,7 +413,6 @@ def encode_list(output, value, with_constructor=True, use_smallest=True): 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 """ diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/_message_backcompat.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/_message_backcompat.py index b14bf24aad78..0e3c22213eda 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/_message_backcompat.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/_message_backcompat.py @@ -51,22 +51,36 @@ def __init__(self, message, **kwargs): 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._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 - self.annotations = self._message.annotations - self.header = LegacyMessageHeader(self._message.header) if self._message.header else 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') + self._to_outgoing_amqp_message: Callable = kwargs.get( + "to_outgoing_amqp_message" + ) def __str__(self): return str(self._message) @@ -109,7 +123,7 @@ def get_message(self): def accept(self): if self._can_settle_message(): - self._settler.settle_messages(self.delivery_no, 'accepted') + self._settler.settle_messages(self.delivery_no, "accepted") self.state = MessageState.ReceivedSettled return True return False @@ -118,12 +132,10 @@ def reject(self, condition=None, description=None, info=None): if self._can_settle_message(): self._settler.settle_messages( self.delivery_no, - 'rejected', + "rejected", error=AMQPError( - condition=condition, - description=description, - info=info - ) + condition=condition, description=description, info=info + ), ) self.state = MessageState.ReceivedSettled return True @@ -131,7 +143,7 @@ def reject(self, condition=None, description=None, info=None): def release(self): if self._can_settle_message(): - self._settler.settle_messages(self.delivery_no, 'released') + self._settler.settle_messages(self.delivery_no, "released") self.state = MessageState.ReceivedSettled return True return False @@ -140,7 +152,7 @@ def modify(self, failed, deliverable, annotations=None): if self._can_settle_message(): self._settler.settle_messages( self.delivery_no, - 'modified', + "modified", delivery_failed=failed, undeliverable_here=deliverable, message_annotations=annotations, @@ -157,7 +169,6 @@ class LegacyBatchMessage(LegacyMessage): 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) @@ -206,14 +217,13 @@ def get_properties_obj(self): self.creation_time, self.group_id, self.group_sequence, - self.reply_to_group_id + self.reply_to_group_id, ) class LegacyMessageHeader(object): - def __init__(self, header): - self.delivery_count = header.delivery_count # or 0 + 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 @@ -236,5 +246,5 @@ def get_header_obj(self): self.priority, self.time_to_live, self.first_acquirer, - self.delivery_count + self.delivery_count, ) diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/_transport.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/_transport.py index 63cc78f23cda..20935e6170c9 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/_transport.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/_transport.py @@ -91,6 +91,8 @@ def set_cloexec(fd, cloexec): # noqa 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+))?") @@ -150,6 +152,7 @@ def __init__( *, port=AMQP_PORT, connect_timeout=None, + read_timeout=None, socket_settings=None, raise_on_initial_eintr=True, **kwargs # pylint: disable=unused-argument @@ -162,6 +165,7 @@ def __init__( self.host, self.port = to_host_port(host, port) 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() @@ -171,8 +175,10 @@ def connect(self): if self.connected: return self._connect(self.host, self.port, self.connect_timeout) - self._init_socket(self.socket_settings) - self.sock.settimeout(0.2) + 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 @@ -316,17 +322,28 @@ def _connect(self, host, port, timeout): # hurray, we established connection return - def _init_socket(self, socket_settings): + 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 - # 1 second is enough for perf analysis - self.sock.settimeout(1) # set socket back to non-blocking mode + # 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 + def _get_tcp_socket_defaults(self, sock): # pylint: disable=no-self-use tcp_opts = {} for opt in KNOWN_TCP_OPTS: enum = None @@ -486,16 +503,16 @@ def _wrap_socket(self, sock, context=None, **sslopts): return self._wrap_context(sock, sslopts, **context) return self._wrap_socket_sni(sock, **sslopts) - def _wrap_context( + def _wrap_context( # pylint: disable=no-self-use self, sock, sslopts, 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.wrap_socket(sock, **sslopts) - def _wrap_socket_sni( + def _wrap_socket_sni( # pylint: disable=no-self-use self, sock, keyfile=None, @@ -508,7 +525,7 @@ def _wrap_socket_sni( server_hostname=None, ciphers=None, ssl_version=None, - ): # pylint: disable=no-self-use + ): """Socket wrap with SNI headers. Default `ssl.wrap_socket` method augmented with support for @@ -649,7 +666,7 @@ def __init__( **kwargs, ): self.sslopts = ssl_opts if isinstance(ssl_opts, dict) else {} - self._connect_timeout = connect_timeout or TIMEOUT_INTERVAL + 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) @@ -683,7 +700,7 @@ def connect(self): "Please install websocket-client library to use websocket transport." ) - def _read(self, n, initial=False, buffer=None, _errnos=None): + 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 diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_cbs_async.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_cbs_async.py index 7e6fcc91d2f2..7667f846dc3a 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_cbs_async.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_cbs_async.py @@ -160,9 +160,7 @@ async def _cbs_link_ready(self): if self.state != CbsState.OPEN: return False if self.state in (CbsState.CLOSED, CbsState.ERROR): - # TODO: raise proper error type also should this be a ClientError? - # Think how upper layer handle this exception + condition code - raise AuthenticationException( + raise TokenAuthFailure( # pylint: disable = no-value-for-parameter condition=ErrorCondition.ClientError, description="CBS authentication link is in broken status, please recreate the cbs link.", ) diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_client_async.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_client_async.py index 3bb8593c202b..ba6c24ad1125 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_client_async.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_client_async.py @@ -17,7 +17,12 @@ from ._connection_async import Connection from ._management_operation_async import ManagementOperation from ._cbs_async import CBSAuthenticator -from ..client import AMQPClientSync, ReceiveClientSync, SendClientSync, Outcomes +from ..client import ( + AMQPClient as AMQPClientSync, + ReceiveClient as ReceiveClientSync, + SendClient as SendClientSync, + Outcomes +) from ..message import _MessageDelivery from ..constants import ( MessageDeliveryState, diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_connection_async.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_connection_async.py index f55c8b59cc90..13e7267938de 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_connection_async.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_connection_async.py @@ -214,7 +214,7 @@ async def _disconnect(self) -> None: if self.state == ConnectionState.END: return await self._set_state(ConnectionState.END) - self._transport.close() + await self._transport.close() def _can_read(self): # type: () -> bool @@ -653,7 +653,6 @@ async def _process_outgoing_frame(self, channel, frame): cast(float, self._last_frame_received_time), ) or (await self._get_remote_timeout(now)): await self.close( - # TODO: check error condition error=AMQPError( condition=ErrorCondition.ConnectionCloseForced, description="No frame received for the idle timeout.", @@ -735,7 +734,6 @@ async def listen(self, wait=False, batch=1, **kwargs): cast(float, self._idle_timeout), cast(float, self._last_frame_received_time), ) or (await self._get_remote_timeout(now)): - # TODO: check error condition await self.close( error=AMQPError( condition=ErrorCondition.ConnectionCloseForced, diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_sender_async.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_sender_async.py index ce7ce7eb3ee0..37bb95a705f4 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_sender_async.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_sender_async.py @@ -60,7 +60,7 @@ def from_incoming_frame(cls, session, handle, frame): async def _incoming_attach(self, frame): try: await super(SenderLink, self)._incoming_attach(frame) - except ValueError: # TODO: This should NOT be a ValueError + except AMQPLinkError: await self._remove_pending_deliveries() raise self.current_link_credit = self.link_credit @@ -163,9 +163,9 @@ async def update_pending_deliveries(self): async def send_transfer(self, message, *, send_async=False, **kwargs): self._check_if_closed() if self.state != LinkState.ATTACHED: - raise AMQPLinkError( # TODO: should we introduce MessageHandler to indicate the handler is in wrong state - condition=ErrorCondition.ClientError, # TODO: should this be a ClientError? - description="Link is not 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: diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_session_async.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_session_async.py index 96c707bc18ab..02f5eab4caae 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_session_async.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_session_async.py @@ -9,21 +9,21 @@ import logging import time import asyncio -from typing import Optional, Union, TYPE_CHECKING +from typing import Optional, Union -from ..constants import ( - ConnectionState, - SessionState, - SessionTransferState, - Role -) +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 ..performatives import ( + BeginFrame, + EndFrame, + FlowFrame, + TransferFrame, + DispositionFrame, +) from .._encode import encode_frame -if TYPE_CHECKING: - from ..error import AMQPError +from ..error import AMQPError, ErrorCondition _LOGGER = logging.getLogger(__name__) @@ -88,10 +88,17 @@ async def _set_state(self, new_state): 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) + _LOGGER.info( + "Session state changed: %r -> %r", + previous_state, + new_state, + extra=self.network_trace_params, + ) futures = [] for link in self.links.values(): - futures.append(asyncio.ensure_future(link._on_session_state_change())) # pylint: disable=protected-access + futures.append( + asyncio.ensure_future(link._on_session_state_change()) # pylint: disable=protected-access + ) await asyncio.gather(*futures) async def _on_connection_state_change(self): @@ -108,24 +115,38 @@ def _get_next_output_handle(self): :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) + 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, + 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, + 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.info("-> %r", begin_frame, extra=self.network_trace_params) - await self._connection._process_outgoing_frame(self.channel, begin_frame) # pylint: disable=protected-access + await self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, begin_frame + ) async def _incoming_begin(self, frame): if self.network_trace: @@ -146,12 +167,18 @@ async def _outgoing_end(self, error=None): end_frame = EndFrame(error=error) if self.network_trace: _LOGGER.info("-> %r", end_frame, extra=self.network_trace_params) - await self._connection._process_outgoing_frame(self.channel, end_frame) # pylint: disable=protected-access + await self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, end_frame + ) async def _incoming_end(self, frame): if self.network_trace: _LOGGER.info("<- %r", EndFrame(*frame), extra=self.network_trace_params) - if self.state not in [SessionState.END_RCVD, SessionState.END_SENT, SessionState.DISCARDING]: + 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() @@ -160,19 +187,41 @@ async def _incoming_end(self, frame): 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 + await self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, frame + ) 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 + self._input_handles[frame[1]] = self.links[ + frame[0].decode("utf-8") + ] # name and handle + await self._input_handles[frame[1]]._incoming_attach( # pylint: disable=protected-access + frame + ) except KeyError: - outgoing_handle = self._get_next_output_handle() # TODO: catch max-handles error - if frame[2] == Role.Sender: # role - new_link = ReceiverLink.from_incoming_frame(self, outgoing_handle, frame) + try: + outgoing_handle = self._get_next_output_handle() + except ValueError: + # detach the link that would have been set. + await 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: + 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 + 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 @@ -193,22 +242,34 @@ async def _outgoing_flow(self, frame=None): flow_frame = FlowFrame(**link_flow) if self.network_trace: _LOGGER.info("-> %r", flow_frame, extra=self.network_trace_params) - await self._connection._process_outgoing_frame(self.channel, flow_frame) # pylint: disable=protected-access + await self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, flow_frame + ) async def _incoming_flow(self, frame): if self.network_trace: _LOGGER.info("<- %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 + 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 + await self._input_handles[frame[4]]._incoming_flow( # pylint: disable=protected-access + frame + ) else: futures = [] for link in self._output_handles.values(): - if self.remote_incoming_window > 0 and not link._is_closed: # pylint: disable=protected-access - futures.append(link._incoming_flow(frame)) # pylint: disable=protected-access + if ( + self.remote_incoming_window > 0 and not link._is_closed # pylint: disable=protected-access + ): + futures.append( + link._incoming_flow(frame) # pylint: disable=protected-access + ) await asyncio.gather(*futures) async def _outgoing_transfer(self, delivery): @@ -229,7 +290,9 @@ async def _outgoing_transfer(self, delivery): # 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 + 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 @@ -249,7 +312,9 @@ async def _outgoing_transfer(self, delivery): "payload": payload[start_idx : start_idx + available_frame_size], "delivery_id": self.next_outgoing_id, } - await self._connection._process_outgoing_frame(self.channel, TransferFrame(**tmp_delivery_frame)) # pylint: disable=protected-access + await self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, TransferFrame(**tmp_delivery_frame) + ) start_idx += available_frame_size remaining_payload_cnt -= available_frame_size @@ -268,7 +333,9 @@ async def _outgoing_transfer(self, delivery): "payload": payload[start_idx:], "delivery_id": self.next_outgoing_id, } - await self._connection._process_outgoing_frame(self.channel, TransferFrame(**tmp_delivery_frame)) # pylint: disable=protected-access + await self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, TransferFrame(**tmp_delivery_frame) + ) self.next_outgoing_id += 1 self.remote_incoming_window -= 1 self.outgoing_window -= 1 @@ -280,37 +347,61 @@ async def _incoming_transfer(self, frame): self.remote_outgoing_window -= 1 self.incoming_window -= 1 try: - await self._input_handles[frame[0]]._incoming_transfer(frame) # pylint: disable=protected-access + await self._input_handles[frame[0]]._incoming_transfer( # pylint: disable=protected-access + frame + ) except KeyError: - pass # TODO: "unattached handle" + 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 + await self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, frame + ) async def _incoming_disposition(self, frame): if self.network_trace: - _LOGGER.info("<- %r", DispositionFrame(*frame), extra=self.network_trace_params) + _LOGGER.info( + "<- %r", DispositionFrame(*frame), extra=self.network_trace_params + ) futures = [] for link in self._input_handles.values(): - asyncio.ensure_future(link._incoming_disposition(frame)) # pylint: disable=protected-access + asyncio.ensure_future( + link._incoming_disposition(frame) # pylint: disable=protected-access + ) await asyncio.gather(*futures) async def _outgoing_detach(self, frame): - await self._connection._process_outgoing_frame(self.channel, frame) # pylint: disable=protected-access + await self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, frame + ) async def _incoming_detach(self, frame): try: link = self._input_handles[frame[0]] # handle - await link._incoming_detach(frame) # pylint: disable=protected-access + 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: - pass # TODO: close session with unattached-handle + 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 @@ -334,7 +425,9 @@ async def begin(self, wait=False): 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.") + 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 @@ -358,7 +451,7 @@ def create_receiver_link(self, source_address, **kwargs): source_address=source_address, network_trace=kwargs.pop("network_trace", self.network_trace), network_trace_params=dict(self.network_trace_params), - **kwargs + **kwargs, ) self.links[link.name] = link self._output_handles[assigned_handle] = link @@ -372,11 +465,16 @@ def create_sender_link(self, target_address, **kwargs): target_address=target_address, network_trace=kwargs.pop("network_trace", self.network_trace), network_trace_params=dict(self.network_trace_params), - **kwargs + **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), **kwargs) + return ManagementLink( + self, + endpoint, + network_trace=kwargs.pop("network_trace", self.network_trace), + **kwargs, + ) diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_transport_async.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_transport_async.py index b26a1dc956c6..255db54c9c4a 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_transport_async.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_transport_async.py @@ -41,6 +41,7 @@ from io import BytesIO import logging + import certifi from .._platform import KNOWN_TCP_OPTS, SOL_TCP @@ -66,14 +67,21 @@ class AsyncTransportMixin: async def receive_frame(self, timeout=None, **kwargs): try: - header, channel, payload = await asyncio.wait_for(self.read(**kwargs), timeout=timeout) + 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) _LOGGER.info("ICH%d <- %r", channel, decoded) return channel, decoded - except (TimeoutError, socket.timeout, asyncio.IncompleteReadError, asyncio.TimeoutError): + except ( + TimeoutError, + socket.timeout, + asyncio.IncompleteReadError, + asyncio.TimeoutError, + ): return None, None async def read(self, verify_frame_type=0): @@ -81,7 +89,9 @@ async def read(self, verify_frame_type=0): read_frame_buffer = BytesIO() try: frame_header = memoryview(bytearray(8)) - read_frame_buffer.write(await self._read(8, buffer=frame_header, initial=True)) + 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] @@ -91,17 +101,27 @@ async def read(self, verify_frame_type=0): offset = frame_header[4] frame_type = frame_header[5] if verify_frame_type is not None and frame_type != verify_frame_type: - raise ValueError(f"Received invalid frame type: {frame_type}, expected: {verify_frame_type}") + 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:])) + 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)) + read_frame_buffer.write( + await self._read(payload_size, buffer=payload) + ) except (TimeoutError, socket.timeout, asyncio.IncompleteReadError): read_frame_buffer.write(self._read_buffer.getvalue()) self._read_buffer = read_frame_buffer @@ -129,35 +149,6 @@ async def send_frame(self, channel, frame, **kwargs): await self.write(data) # _LOGGER.info("OCH%d -> %r", channel, frame) - -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.loop = asyncio.get_running_loop() - self.socket_lock = asyncio.Lock() - self.sslopts = self._build_ssl_opts(ssl_opts) - def _build_ssl_opts(self, sslopts): if sslopts in [True, False, None, {}]: return sslopts @@ -187,15 +178,50 @@ def _build_ssl_opts(self, sslopts): return context return True except TypeError: - raise TypeError("SSL configuration must be a dictionary, or the value True.") + 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 + 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.loop = asyncio.get_running_loop() + self.socket_lock = asyncio.Lock() + self.sslopts = self._build_ssl_opts(ssl_opts) + async def connect(self): try: # are we already connected? @@ -204,7 +230,9 @@ async def connect(self): 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 + 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 @@ -233,7 +261,9 @@ async def _connect(self, host, port, timeout): for n, family in enumerate(addr_types): # first, resolve the address for a single address family try: - entries = await self.loop.getaddrinfo(host, port, family=family, type=socket.SOCK_STREAM, proto=SOL_TCP) + entries = await self.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 @@ -241,7 +271,9 @@ async def _connect(self, host, port, timeout): # 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") + 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 @@ -300,7 +332,13 @@ def _set_socket_options(self, 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)): + 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. @@ -312,7 +350,9 @@ async def _read(self, toread, initial=False, buffer=None, _errnos=(errno.ENOENT, try: while toread: try: - view[nbytes : nbytes + toread] = await self.reader.readexactly(toread) + view[nbytes : nbytes + toread] = await self.reader.readexactly( + toread + ) nbytes = toread except asyncio.IncompleteReadError as exc: pbytes = len(exc.partial) @@ -344,7 +384,7 @@ async def _write(self, s): """Write a string out to the SSL socket fully.""" self.writer.write(s) - def close(self): + async def close(self): if self.writer is not None: if self.sslopts: # see issue: https://github.com/encode/httpx/issues/914 @@ -388,57 +428,77 @@ async def negotiate(self): ) -class WebSocketTransportAsync(AsyncTransportMixin): # pylint: disable=too-many-instance-attributes - def __init__(self, host, *, port=WEBSOCKET_PORT, connect_timeout=None, ssl_opts=None, **kwargs): +class WebSocketTransportAsync(AsyncTransportMixin): + def __init__( + self, + host, + *, + port=WEBSOCKET_PORT, # pylint: disable=unused-argument + connect_timeout=None, + ssl_opts=None, + **kwargs + ): self._read_buffer = BytesIO() - self.loop = asyncio.get_running_loop() self.socket_lock = asyncio.Lock() - self.sslopts = ssl_opts if isinstance(ssl_opts, dict) else {} + self.sslopts = self._build_ssl_opts(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.host = host self.ws = None - self.connected = False + self.session = None self._http_proxy = kwargs.get("http_proxy", None) + self.connected = False async def connect(self): - http_proxy_host, http_proxy_port, http_proxy_auth = None, None, None + 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) - if username or password: - http_proxy_auth = (username, password) + try: - from websocket import create_connection + from aiohttp import ClientSession - self.ws = create_connection( + if username or password: + from aiohttp import BasicAuth + + http_proxy_auth = BasicAuth(login=username, password=password) + + self.session = ClientSession() + self.ws = await self.session.ws_connect( 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, + protocols=[AMQP_WS_SUBPROTOCOL], + autoclose=False, + proxy=http_proxy_host, + proxy_auth=http_proxy_auth, + ssl=self.sslopts, ) + self.connected = True + except ImportError: - raise ValueError("Please install websocket-client library to use websocket transport.") + raise ValueError( + "Please install aiohttp library to use websocket transport." + ) - async def _read(self, n, initial=False, buffer=None): # pylint: disable=unused-argument + async def _read(self, n, buffer=None, **kwargs): # pylint: disable=unused-argument """Read exactly n bytes from the peer.""" - from websocket import WebSocketTimeoutException length = 0 view = buffer or memoryview(bytearray(n)) nbytes = self._read_buffer.readinto(view) length += nbytes n -= nbytes + try: while n: - data = await self.loop.run_in_executor(None, self.ws.recv) - + data = await self.ws.receive_bytes() if len(data) <= n: view[length : length + len(data)] = data n -= len(data) @@ -446,20 +506,20 @@ async def _read(self, n, initial=False, buffer=None): # pylint: disable=unused-a view[length : length + n] = data[0:n] self._read_buffer = BytesIO(data[n:]) n = 0 - return view - except WebSocketTimeoutException: + except asyncio.TimeoutError: raise TimeoutError() - def close(self): + async def close(self): """Do any preliminary work in shutting down the connection.""" - self.ws.close() + await self.ws.close() + await self.session.close() self.connected = False async def write(self, s): - """Completely write a string to the peer. + """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.loop.run_in_executor(None, self.ws.send_binary, s) + await self.ws.send_bytes(s) diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/cbs.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/cbs.py index 6a5259eb9f95..2f80c7a839e1 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/cbs.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/cbs.py @@ -204,9 +204,7 @@ def _cbs_link_ready(self): if self.state != CbsState.OPEN: return False if self.state in (CbsState.CLOSED, CbsState.ERROR): - # TODO: raise proper error type also should this be a ClientError? - # Think how upper layer handle this exception + condition code - raise AuthenticationException( + raise TokenAuthFailure( # pylint: disable = no-value-for-parameter condition=ErrorCondition.ClientError, description="CBS authentication link is in broken status, please recreate the cbs link.", ) diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/client.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/client.py index 6dad85f4ebcf..8a17b202ef4e 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/client.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/client.py @@ -4,6 +4,9 @@ # 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 @@ -53,7 +56,9 @@ _logger = logging.getLogger(__name__) -class AMQPClientSync(object): # pylint: disable=too-many-instance-attributes +class AMQPClient( + object +): # pylint: disable=too-many-instance-attributes """An AMQP client. :param hostname: The AMQP endpoint to connect to. :type hostname: str @@ -272,7 +277,6 @@ def open(self, connection=None): multiple clients. :type connection: ~pyamqp.Connection """ - # pylint: disable=protected-access if self._session: return # already open. @@ -418,7 +422,7 @@ def mgmt_request(self, message, **kwargs): return status, description, response -class SendClientSync(AMQPClientSync): +class SendClient(AMQPClient): """ An AMQP client for sending messages. :param target: The target AMQP service endpoint. This can either be the URI as @@ -517,7 +521,7 @@ def __init__(self, hostname, target, **kwargs): 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(SendClientSync, self).__init__(hostname, **kwargs) + super(SendClient, self).__init__(hostname, **kwargs) def _client_ready(self): """Determine whether the client is ready to start receiving messages. @@ -650,7 +654,7 @@ def send_message(self, message, **kwargs): self._do_retryable_operation(self._send_message_impl, message=message, **kwargs) -class ReceiveClientSync(AMQPClientSync): +class ReceiveClient(AMQPClient): """ An AMQP client for receiving messages. :param source: The source AMQP service endpoint. This can either be the URI as @@ -753,7 +757,7 @@ def __init__(self, hostname, source, **kwargs): 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(ReceiveClientSync, self).__init__(hostname, **kwargs) + super(ReceiveClient, self).__init__(hostname, **kwargs) def _client_ready(self): """Determine whether the client is ready to start receiving messages. @@ -862,7 +866,7 @@ def _receive_message_batch_impl( def close(self): self._received_messages = queue.Queue() - super(ReceiveClientSync, self).close() + super(ReceiveClient, self).close() def receive_message_batch(self, **kwargs): """Receive a batch of messages. Messages returned in the batch have already been diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/sasl.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/sasl.py index 6c89343dd33a..c4ff9d265540 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/sasl.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/sasl.py @@ -127,7 +127,7 @@ def __init__( host, credential, *, - port=WEBSOCKET_PORT, # TODO: NOT KWARGS IN EH PYAMQP + port=WEBSOCKET_PORT, connect_timeout=None, ssl_opts=None, **kwargs, diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/sender.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/sender.py index 70e9bc62cfca..c0e9e64cd6e2 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/sender.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/sender.py @@ -59,7 +59,7 @@ def from_incoming_frame(cls, session, handle, frame): def _incoming_attach(self, frame): try: super(SenderLink, self)._incoming_attach(frame) - except ValueError: # TODO: This should NOT be a ValueError + except AMQPLinkError: self._remove_pending_deliveries() raise self.current_link_credit = self.link_credit @@ -160,9 +160,9 @@ def update_pending_deliveries(self): def send_transfer(self, message, *, send_async=False, **kwargs): self._check_if_closed() if self.state != LinkState.ATTACHED: - raise AMQPLinkError( # TODO: should we introduce MessageHandler to indicate the handler is in wrong state - condition=ErrorCondition.ClientError, # TODO: should this be a ClientError? - description="Link is not 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: diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/session.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/session.py index 0cdb2cdc7a8e..ea36c5b1be1c 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/session.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/session.py @@ -8,21 +8,21 @@ import uuid import logging import time -from typing import Union, Optional, TYPE_CHECKING +from typing import Union, Optional -from .constants import ( - ConnectionState, - SessionState, - SessionTransferState, - Role -) +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 .performatives import ( + BeginFrame, + EndFrame, + FlowFrame, + TransferFrame, + DispositionFrame, +) +from .error import AMQPError, ErrorCondition from ._encode import encode_frame -if TYPE_CHECKING: - from .error import AMQPError _LOGGER = logging.getLogger(__name__) @@ -87,9 +87,14 @@ def _set_state(self, new_state): 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) + _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 + 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]: @@ -105,24 +110,38 @@ def _get_next_output_handle(self): :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) + 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, + 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, + 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.info("-> %r", begin_frame, extra=self.network_trace_params) - self._connection._process_outgoing_frame(self.channel, begin_frame) # pylint: disable=protected-access + self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, begin_frame + ) def _incoming_begin(self, frame): if self.network_trace: @@ -143,12 +162,18 @@ def _outgoing_end(self, error=None): end_frame = EndFrame(error=error) if self.network_trace: _LOGGER.info("-> %r", end_frame, extra=self.network_trace_params) - self._connection._process_outgoing_frame(self.channel, end_frame) # pylint: disable=protected-access + self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, end_frame + ) def _incoming_end(self, frame): if self.network_trace: _LOGGER.info("<- %r", EndFrame(*frame), extra=self.network_trace_params) - if self.state not in [SessionState.END_RCVD, SessionState.END_SENT, SessionState.DISCARDING]: + 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() @@ -157,19 +182,41 @@ def _incoming_end(self, frame): self._set_state(SessionState.UNMAPPED) def _outgoing_attach(self, frame): - self._connection._process_outgoing_frame(self.channel, frame) # pylint: disable=protected-access + 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(frame) # pylint: disable=protected-access + 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: - outgoing_handle = self._get_next_output_handle() # TODO: catch max-handles error + try: + outgoing_handle = self._get_next_output_handle() + except ValueError: + # 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) + 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 + 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 @@ -190,21 +237,31 @@ def _outgoing_flow(self, frame=None): flow_frame = FlowFrame(**link_flow) if self.network_trace: _LOGGER.info("-> %r", flow_frame, extra=self.network_trace_params) - self._connection._process_outgoing_frame(self.channel, flow_frame) # pylint: disable=protected-access + self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, flow_frame + ) def _incoming_flow(self, frame): if self.network_trace: _LOGGER.info("<- %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 + 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(frame) # pylint: disable=protected-access + 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 + 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): if self.state != SessionState.MAPPED: @@ -224,7 +281,9 @@ def _outgoing_transfer(self, delivery): # 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 + 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 @@ -244,7 +303,9 @@ def _outgoing_transfer(self, delivery): "payload": payload[start_idx : start_idx + available_frame_size], "delivery_id": self.next_outgoing_id, } - self._connection._process_outgoing_frame(self.channel, TransferFrame(**tmp_delivery_frame)) # pylint: disable=protected-access + self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, TransferFrame(**tmp_delivery_frame) + ) start_idx += available_frame_size remaining_payload_cnt -= available_frame_size @@ -263,7 +324,9 @@ def _outgoing_transfer(self, delivery): "payload": payload[start_idx:], "delivery_id": self.next_outgoing_id, } - self._connection._process_outgoing_frame(self.channel, TransferFrame(**tmp_delivery_frame)) # pylint: disable=protected-access + self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, TransferFrame(**tmp_delivery_frame) + ) self.next_outgoing_id += 1 self.remote_incoming_window -= 1 self.outgoing_window -= 1 @@ -275,35 +338,57 @@ def _incoming_transfer(self, frame): self.remote_outgoing_window -= 1 self.incoming_window -= 1 try: - self._input_handles[frame[0]]._incoming_transfer(frame) # pylint: disable=protected-access + self._input_handles[frame[0]]._incoming_transfer( # pylint: disable=protected-access + frame + ) except KeyError: - pass # TODO: "unattached handle" + 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""", + ) + ) 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(self.channel, frame) # pylint: disable=protected-access + self._connection._process_outgoing_frame( # pylint: disable=protected-access + self.channel, frame + ) def _incoming_disposition(self, frame): if self.network_trace: - _LOGGER.info("<- %r", DispositionFrame(*frame), extra=self.network_trace_params) + _LOGGER.info( + "<- %r", DispositionFrame(*frame), extra=self.network_trace_params + ) for link in self._input_handles.values(): - link._incoming_disposition(frame) # pylint: disable=protected-access + link._incoming_disposition(frame) # pylint: disable=protected-access def _outgoing_detach(self, frame): - self._connection._process_outgoing_frame(self.channel, frame) # pylint: disable=protected-access + 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 + 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: - pass # TODO: close session with unattached-handle + 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 @@ -327,7 +412,9 @@ def begin(self, wait=False): 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.") + 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 @@ -351,7 +438,7 @@ def create_receiver_link(self, source_address, **kwargs): source_address=source_address, network_trace=kwargs.pop("network_trace", self.network_trace), network_trace_params=dict(self.network_trace_params), - **kwargs + **kwargs, ) self.links[link.name] = link self._output_handles[assigned_handle] = link @@ -365,11 +452,16 @@ def create_sender_link(self, target_address, **kwargs): target_address=target_address, network_trace=kwargs.pop("network_trace", self.network_trace), network_trace_params=dict(self.network_trace_params), - **kwargs + **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), **kwargs) + return ManagementLink( + self, + endpoint, + network_trace=kwargs.pop("network_trace", self.network_trace), + **kwargs, + ) diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_servicebus_receiver.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_servicebus_receiver.py index 2cdda7760222..575973fdbfd3 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_servicebus_receiver.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_servicebus_receiver.py @@ -16,7 +16,7 @@ #from uamqp.authentication.common import AMQPAuth from ._pyamqp.message import Message from ._pyamqp.constants import SenderSettleMode -from ._pyamqp.client import ReceiveClientSync +from ._pyamqp.client import ReceiveClient as ReceiveClientSync from ._pyamqp import utils from ._pyamqp.error import AMQPError diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_servicebus_sender.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_servicebus_sender.py index 55992dbeb7ca..d419770efe23 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_servicebus_sender.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_servicebus_sender.py @@ -10,7 +10,7 @@ from typing import Any, TYPE_CHECKING, Union, List, Optional, Mapping, cast #from uamqp.authentication.common import AMQPAuth -from ._pyamqp.client import SendClientSync +from ._pyamqp.client import SendClient as SendClientSync from ._pyamqp.utils import amqp_long_value, amqp_array_value from ._pyamqp.error import MessageException diff --git a/sdk/servicebus/azure-servicebus/tests/test_message.py b/sdk/servicebus/azure-servicebus/tests/test_message.py index 6c71bc823a9b..ce13c6488dc6 100644 --- a/sdk/servicebus/azure-servicebus/tests/test_message.py +++ b/sdk/servicebus/azure-servicebus/tests/test_message.py @@ -521,6 +521,7 @@ def test_message_backcompat_peek_lock_databody(self, servicebus_namespace_connec assert not incoming_message.message.reject() assert not incoming_message.message.modify(True, True) + @pytest.mark.skip("unskip after adding PyamqpTransport + pass in _to_outgoing_amqp_message to LegacyMessage") @pytest.mark.liveTest @pytest.mark.live_test_only @CachedResourceGroupPreparer(name_prefix='servicebustest') @@ -557,6 +558,7 @@ def test_message_backcompat_receive_and_delete_valuebody(self, servicebus_namesp assert not incoming_message.message.reject() assert not incoming_message.message.modify(True, True) + @pytest.mark.skip("unskip after adding PyamqpTransport + pass in _to_outgoing_amqp_message to LegacyMessage") @pytest.mark.liveTest @pytest.mark.live_test_only @CachedResourceGroupPreparer(name_prefix='servicebustest') @@ -598,6 +600,7 @@ def test_message_backcompat_peek_lock_valuebody(self, servicebus_namespace_conne assert not incoming_message.message.reject() assert not incoming_message.message.modify(True, True) + @pytest.mark.skip("unskip after adding PyamqpTransport + pass in _to_outgoing_amqp_message to LegacyMessage") @pytest.mark.liveTest @pytest.mark.live_test_only @CachedResourceGroupPreparer(name_prefix='servicebustest') @@ -634,6 +637,7 @@ def test_message_backcompat_receive_and_delete_sequencebody(self, servicebus_nam assert not incoming_message.message.reject() assert not incoming_message.message.modify(True, True) + @pytest.mark.skip("unskip after adding PyamqpTransport + pass in _to_outgoing_amqp_message to LegacyMessage") @pytest.mark.liveTest @pytest.mark.live_test_only @CachedResourceGroupPreparer(name_prefix='servicebustest') diff --git a/sdk/servicebus/azure-servicebus/tests/test_queues.py b/sdk/servicebus/azure-servicebus/tests/test_queues.py index 4a87fde884a7..ca10b5a884de 100644 --- a/sdk/servicebus/azure-servicebus/tests/test_queues.py +++ b/sdk/servicebus/azure-servicebus/tests/test_queues.py @@ -2344,22 +2344,22 @@ def _hack_sb_receiver_settle_message(self, message, settle_operation, dead_lette servicebus_namespace_connection_string, logging_enable=False) as sb_client: sender = sb_client.get_queue_sender(servicebus_queue.name) receiver = sb_client.get_queue_receiver(servicebus_queue.name, max_wait_time=5) - original_settlement = client.ReceiveClientSync.settle_messages + original_settlement = client.ReceiveClient.settle_messages try: with sender, receiver: # negative settlement via receiver link sender.send_messages(ServiceBusMessage("body"), timeout=10) message = receiver.receive_messages()[0] - client.ReceiveClientSync.settle_messages = types.MethodType(_hack_amqp_message_complete, receiver._handler) + client.ReceiveClient.settle_messages = types.MethodType(_hack_amqp_message_complete, receiver._handler) receiver.complete_message(message) # settle via mgmt link - origin_amqp_client_mgmt_request_method = client.AMQPClientSync.mgmt_request + origin_amqp_client_mgmt_request_method = client.AMQPClient.mgmt_request try: - client.AMQPClientSync.mgmt_request = _hack_amqp_mgmt_request + client.AMQPClient.mgmt_request = _hack_amqp_mgmt_request with pytest.raises(ServiceBusConnectionError): receiver.peek_messages() finally: - client.AMQPClientSync.mgmt_request = origin_amqp_client_mgmt_request_method + client.AMQPClient.mgmt_request = origin_amqp_client_mgmt_request_method sender.send_messages(ServiceBusMessage("body"), timeout=10) @@ -2374,7 +2374,7 @@ def _hack_sb_receiver_settle_message(self, message, settle_operation, dead_lette message = receiver.receive_messages(max_wait_time=6)[0] receiver.complete_message(message) finally: - client.ReceiveClientSync.settle_messages = original_settlement + client.ReceiveClient.settle_messages = original_settlement @pytest.mark.skip(reason="TODO: iterator support") @pytest.mark.liveTest