intersect-sdk 0.8.2__tar.gz → 0.8.4__tar.gz

{intersect_sdk-0.8.2 → intersect_sdk-0.8.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: intersect-sdk
-Version: 0.8.2
+Version: 0.8.4
 Summary: Python SDK to interact with INTERSECT
 Keywords: intersect
 Author-Email: Lance Drane <dranelt@ornl.gov>, Marshall McDonnell <mcdonnellmt@ornl.gov>, Seth Hitefield <hitefieldsd@ornl.gov>, Andrew Ayres <ayresaf@ornl.gov>, Gregory Cage <cagege@ornl.gov>, Jesse McGaha <mcgahajr@ornl.gov>, Robert Smith <smithrw@ornl.gov>, Gavin Wiggins <wigginsg@ornl.gov>, Michael Brim <brimmj@ornl.gov>, Rick Archibald <archibaldrk@ornl.gov>, Addi Malviya Thakur <malviyaa@ornl.gov>

{intersect_sdk-0.8.2 → intersect_sdk-0.8.4}/pyproject.toml

@@ -35,7 +35,7 @@ dependencies = [
     "jsonschema[format-nongpl]>=4.21.1",
     "eval-type-backport>=0.1.3;python_version<'3.10'",
 ]
-version = "0.8.2"
+version = "0.8.4"
 
 [project.license]
 text = "BSD-3-Clause"

{intersect_sdk-0.8.2 → intersect_sdk-0.8.4}/src/intersect_sdk/_internal/control_plane/brokers/amqp_client.py

@@ -33,11 +33,33 @@ if TYPE_CHECKING:
 
 _AMQP_MAX_RETRIES = 10
 
+_PREFETCH_COUNT = 1
+"""This determines the maximum amount of messages a single consumer will handle. The consumer must acknowledge the message before it gets another message."""
+
 # Note that we deliberately do NOT want this configurable at runtime. Any two INTERSECT services/clients could potentially exchange messages between one another.
 _INTERSECT_MESSAGE_EXCHANGE = 'intersect-messages'
 """All INTERSECT messages get published to one exchange on the broker."""
 
 
+class ConsumerTagInfo:
+    def __init__(self) -> None:
+        self.consumer_tag = ''
+        self.consumer_tag_event = threading.Event()
+
+    def obtain_consumer_tag(self, consumer_tag: str) -> None:
+        self.consumer_tag = consumer_tag
+        self.consumer_tag_event.set()
+
+    def consumer_tag_obtained(self) -> bool:
+        return self.consumer_tag_event.is_set()
+
+    def wait(self, time: float) -> None:
+        self.consumer_tag_event.wait(time)
+
+    def __repr__(self) -> str:
+        return f'{self.consumer_tag} -- OBTAINED: {self.consumer_tag_obtained()}'
+
+
 def _get_queue_name(routing_key: str) -> str:
     """Generate a valid queue name from the routing key.
 
@@ -101,6 +123,10 @@ class AMQPClient(BrokerClient):
             virtual_host='/',
             credentials=pika.PlainCredentials(username, password),
             connection_attempts=3,
+            # if not specified, this value is obtained by the broker. RabbitMQ sets it to 60s by default
+            heartbeat=10,
+            blocked_connection_timeout=5.0,
+            retry_delay=0.5,
         )
 
         # The pika connection to the broker
@@ -113,7 +139,8 @@ class AMQPClient(BrokerClient):
         # Callback to the topics_to_handler list inside of
         self._topics_to_handlers = topics_to_handlers
         # mapping of topics to callables which can unsubscribe from the topic
-        self._topics_to_consumer_tags: dict[str, str] = {}
+        self._topics_to_consumer_tags: dict[str, ConsumerTagInfo] = {}
+        self._consumer_tags_to_threads: dict[str, threading.Thread] = {}
 
         self._should_disconnect = False
         self._connection_retries = 0
@@ -129,7 +156,7 @@ class AMQPClient(BrokerClient):
         self._should_disconnect = False
         self._channel_flags.unset_all()
 
-        if not self._thread:
+        if not self._thread or not self._thread.is_alive():
             self._thread = threading.Thread(target=self._init_connection, daemon=True)
             self._thread.start()
 
@@ -139,8 +166,8 @@ class AMQPClient(BrokerClient):
     def disconnect(self) -> None:
         """Close all connections."""
         self._should_disconnect = True
-        for topic, tag in self._topics_to_consumer_tags.items():
-            self._cancel_consumer_tag(topic, tag)
+        for topic, tag_info in self._topics_to_consumer_tags.items():
+            self._cancel_consumer_tag(topic, tag_info.consumer_tag)
 
         # since _should_disconnect was set, _connection.ioloop.stop() will now execute after explicit connection close
         self._connection.close()
@@ -175,18 +202,23 @@ class AMQPClient(BrokerClient):
             persist: True if message should persist until consumers available, False if message should be removed immediately.
         """
         topic = _hierarchy_2_amqp(topic)
-        self._channel_out.basic_publish(
-            exchange=_INTERSECT_MESSAGE_EXCHANGE,
-            routing_key=topic,
-            body=payload,
-            properties=pika.BasicProperties(
-                content_type='text/plain',
-                delivery_mode=pika.delivery_mode.DeliveryMode.Persistent
-                if persist
-                else pika.delivery_mode.DeliveryMode.Transient,
-                # expiration=None if persist else '8640000',
-            ),
-        )
+        while not self._channel_flags.is_set():
+            self._channel_flags.wait(1.0)
+        if self._connection and self._connection.is_open:
+            self._channel_out.basic_publish(
+                exchange=_INTERSECT_MESSAGE_EXCHANGE,
+                routing_key=topic,
+                body=payload,
+                properties=pika.BasicProperties(
+                    content_type='text/plain',
+                    delivery_mode=pika.delivery_mode.DeliveryMode.Persistent
+                    if persist
+                    else pika.delivery_mode.DeliveryMode.Transient,
+                    # expiration=None if persist else '8640000',
+                ),
+            )
+        else:
+            logger.error('Unable to publish message %s on topic %s', payload, topic)
 
     def subscribe(self, topic: str, persist: bool) -> None:
         """Subscribe to a topic.
@@ -210,25 +242,41 @@ class AMQPClient(BrokerClient):
         Therefore, transient queues will be cleaned up.
         """
         amqp_topic = _hierarchy_2_amqp(topic)
-        consumer_tag = self._topics_to_consumer_tags.get(amqp_topic, None)
-        if consumer_tag:
-            self._cancel_consumer_tag(amqp_topic, consumer_tag)
+        consumer_tag_info = self._topics_to_consumer_tags.get(amqp_topic, None)
+        if consumer_tag_info:
+            self._cancel_consumer_tag(amqp_topic, consumer_tag_info.consumer_tag)
 
     def _cancel_consumer_tag(self, topic: str, consumer_tag: str) -> None:
         if self._channel_in and self._channel_in.is_open:
-            cb = functools.partial(self._cancel_consumer_tag_cb, topic=topic)
+            cb = functools.partial(
+                self._cancel_consumer_tag_cb, topic=topic, consumer_tag=consumer_tag
+            )
             self._channel_in.basic_cancel(
                 consumer_tag,
                 callback=cb,
             )
 
-    def _cancel_consumer_tag_cb(self, _frame: pika.frame.Frame, topic: str) -> None:
+    def _cancel_consumer_tag_cb(
+        self, _frame: pika.frame.Frame, topic: str, consumer_tag: str
+    ) -> None:
         try:
             del self._topics_to_consumer_tags[topic]
         except KeyError:
             # shouldn't happen because ControlPlaneManager gatekeeps consecutive remove_subscription_channel() calls
-            pass
+            logger.error(
+                'Unable to clean up consumer tag related to topic %s , please inform an INTERSECT-SDK developer if you somehow see this message.',
+                topic,
+            )
         logger.info('Unsubscribed from %s', topic)
+        try:
+            thread = self._consumer_tags_to_threads[consumer_tag]
+            # kill thread immediately if not recoverable, wait to send last message if we are shutting down gracefully
+            thread.join(0 if self.considered_unrecoverable() else None)
+            del self._consumer_tags_to_threads[consumer_tag]
+            logger.debug('Consumer cancelled')
+        except KeyError:
+            # This will commonly be encountered, as several topics will often share a single consumer
+            pass
 
     # BEGIN CALLBACKS + THREADSAFE FUNCTIONS #
 
@@ -238,13 +286,18 @@ class AMQPClient(BrokerClient):
         NOTE: ANY functions which are not eventually called from this function
         should be called via self._connection.ioloop.add_callback_threadsafe(cb)
         """
-        while not self._should_disconnect:
-            self._connection = pika.adapters.SelectConnection(
-                parameters=self._connection_params,
-                on_close_callback=self._on_connection_closed,
-                on_open_error_callback=self._on_connection_open_error,
-                on_open_callback=self._on_connection_open,
-            )
+        self._connection = None
+        self._channel_in = None
+        self._channel_out = None
+
+        while not self._should_disconnect and not self.considered_unrecoverable():
+            if not self._connection or self._connection.is_closed:
+                self._connection = pika.adapters.SelectConnection(
+                    parameters=self._connection_params,
+                    on_close_callback=self._on_connection_closed,
+                    on_open_error_callback=self._on_connection_open_error,
+                    on_open_callback=self._on_connection_open,
+                )
 
             # Loops forever until ioloop.stop is called WHEN self._should_disconnect is True
             self._connection.ioloop.start()
@@ -415,6 +468,32 @@ class AMQPClient(BrokerClient):
         topic: str,
         queue_name: str,
         persist: bool,
+    ) -> None:
+        """Sets up basic QOS on the current channel.
+
+        Args:
+            _unused_frame: AMQP response from binding to the queue. Ignored.
+            channel: The Channel being instantiated.
+            topic: Name of the topic on the broker.
+            queue_name: The name of the queue on the AMQP broker.
+            persist: Whether or not our queue should persist on either broker or application shutdown.
+        """
+        cb = functools.partial(
+            self._on_basic_qos_set,
+            channel=channel,
+            topic=topic,
+            queue_name=queue_name,
+            persist=persist,
+        )
+        channel.basic_qos(prefetch_count=_PREFETCH_COUNT, callback=cb)
+
+    def _on_basic_qos_set(
+        self,
+        _unused_frame: Frame,
+        channel: Channel,
+        topic: str,
+        queue_name: str,
+        persist: bool,
     ) -> None:
         """Consumes a message from the given channel.
 
@@ -429,13 +508,15 @@ class AMQPClient(BrokerClient):
         """
         cb = functools.partial(self._on_consume_ok, topic=topic)
         message_cb = functools.partial(self._consume_message, persist=persist)
+        info = ConsumerTagInfo()
+        self._topics_to_consumer_tags[topic] = info
         consumer_tag = channel.basic_consume(
             queue=queue_name,
             auto_ack=not persist,  # persistent messages should be manually acked and we have no reason to NACK a message for now
             on_message_callback=message_cb,
             callback=cb,
         )
-        self._topics_to_consumer_tags[topic] = consumer_tag
+        info.obtain_consumer_tag(consumer_tag)
 
     def _on_consume_ok(self, _unused_frame: Frame, topic: str) -> None:
         """Sets the consume subscription ready event.
@@ -468,14 +549,50 @@ class AMQPClient(BrokerClient):
             body: the AMQP message to be handled.
             persist: Whether or not our queue should persist on either broker or application shutdown.
         """
+        # if we got a message when we shouldn't, quickly try to requeue it
+        if self.considered_unrecoverable() or self._should_disconnect:
+            logger.warning(
+                "WARNING: A message for topic %s has been received when it shouldn't, attempting requeue"
+            )
+            channel.basic_reject(basic_deliver.delivery_tag)
+            return
+
         tth_key = _amqp_2_hierarchy(basic_deliver.routing_key)
         topic_handler = self._topics_to_handlers().get(tth_key)
         if topic_handler:
-            for cb in topic_handler.callbacks:
-                cb(body)
+            consumer_tag_info = self._topics_to_consumer_tags.get(basic_deliver.routing_key)
+            if not consumer_tag_info:
+                logger.error(
+                    'Could not fetch consumer tag for topic %s, please inform an INTERSECT-SDK developer that you saw this message',
+                    tth_key,
+                )
+                return
+            while not consumer_tag_info.consumer_tag_obtained():
+                consumer_tag_info.wait(1.0)
+            thrd = threading.Thread(
+                target=self._consume_message_subthread,
+                args=(channel, topic_handler.callbacks, body, basic_deliver.delivery_tag, persist),
+            )
+            self._consumer_tags_to_threads[consumer_tag_info.consumer_tag] = thrd
+            thrd.start()
+        elif persist:
+            # we somehow got a message that we don't want, discard it
+            channel.basic_ack(basic_deliver.delivery_tag)
+
+    def _consume_message_subthread(
+        self,
+        channel: Channel,
+        callbacks: set[Callable[[bytes], None]],
+        body: bytes,
+        delivery_tag: int,
+        persist: bool,
+    ) -> None:
+        """This is a subthread which executes the consumer code without blocking the IO loop. Without using a subthread, the AMQP heartbeat checker will be blocked."""
+        for cb in callbacks:
+            cb(body)
         # With persistent messages, we only acknowledge the message AFTER we are done processing
         # (this removes the message from the broker queue)
         # this allows us to retry a message if the broker OR this application goes down
         # We currently never NACK or reject a message because in INTERSECT, applications currently never "share" a queue.
         if persist:
-            channel.basic_ack(basic_deliver.delivery_tag)
+            channel.basic_ack(delivery_tag)

{intersect_sdk-0.8.2 → intersect_sdk-0.8.4}/src/intersect_sdk/_internal/control_plane/brokers/mqtt_client.py

@@ -81,6 +81,7 @@ class MQTTClient(BrokerClient):
         """Connect to the defined broker."""
         # Create a client to connect to RabbitMQ
         # TODO MQTT v5 implementations should set clean_start to NEVER here
+        self._should_disconnect = False
         self._connected_flag.clear()
         self._connection.connect(self.host, self.port, 60)
         self._connection.loop_start()
@@ -90,8 +91,9 @@ class MQTTClient(BrokerClient):
     def disconnect(self) -> None:
         """Disconnect from the broker."""
         self._should_disconnect = True
-        self._connection.disconnect()
-        self._connection.loop_stop()
+        if self._connection:
+            self._connection.disconnect()
+            self._connection.loop_stop()
 
     def is_connected(self) -> bool:
         """Check if there is an active connection to the broker.

{intersect_sdk-0.8.2 → intersect_sdk-0.8.4}/src/intersect_sdk/_internal/schema.py

@@ -594,7 +594,7 @@ def get_schema_and_functions_from_capability_implementations(
             'messageTraits': {
                 # this is where we can define our message headers
                 'commonHeaders': {
-                    'messageHeaders': TypeAdapter(UserspaceMessageHeader).json_schema(
+                    'userspaceHeaders': TypeAdapter(UserspaceMessageHeader).json_schema(
                         ref_template='#/components/messageTraits/commonHeaders/userspaceHeaders/$defs/{model}',
                         mode='serialization',
                     ),

{intersect_sdk-0.8.2 → intersect_sdk-0.8.4}/src/intersect_sdk/version.py

@@ -8,7 +8,7 @@ from __future__ import annotations
 from ._internal.version import strip_version_metadata
 
 # may include build metadata
-__version__ = '0.8.2'
+__version__ = '0.8.4'
 
 version_string = strip_version_metadata(__version__)
 """

{intersect_sdk-0.8.2 → intersect_sdk-0.8.4}/tests/fixtures/example_schema.json

@@ -1,6 +1,6 @@
 {
   "asyncapi": "2.6.0",
-  "x-intersect-version": "0.8.2",
+  "x-intersect-version": "0.8.4",
   "info": {
     "title": "test.test.test.test.test",
     "description": "INTERSECT schema",
@@ -943,7 +943,7 @@
     },
     "messageTraits": {
       "commonHeaders": {
-        "messageHeaders": {
+        "userspaceHeaders": {
           "$defs": {
             "IntersectDataHandler": {
               "description": "What data transfer type do you want to use for handling the request/response?\n\nDefault: MESSAGE",
@@ -982,11 +982,7 @@
               "type": "string"
             },
             "data_handler": {
-              "allOf": [
-                {
-                  "$ref": "#/components/messageTraits/commonHeaders/userspaceHeaders/$defs/IntersectDataHandler"
-                }
-              ],
+              "$ref": "#/components/messageTraits/commonHeaders/userspaceHeaders/$defs/IntersectDataHandler",
               "default": 0,
               "description": "Code signifying where data is stored."
             },
@@ -1039,11 +1035,7 @@
               "type": "string"
             },
             "data_handler": {
-              "allOf": [
-                {
-                  "$ref": "#/components/messageTraits/commonHeaders/eventHeaders/$defs/IntersectDataHandler"
-                }
-              ],
+              "$ref": "#/components/messageTraits/commonHeaders/eventHeaders/$defs/IntersectDataHandler",
               "default": 0,
               "description": "Code signifying where data is stored."
             },

{intersect_sdk-0.8.2 → intersect_sdk-0.8.4}/tests/unit/test_schema_invalids.py

@@ -197,7 +197,7 @@ def test_disallow_dynamic_list_subtyping(caplog: pytest.LogCaptureFixture):
     with pytest.raises(SystemExit):
         get_schema_helper([MockAnyList])
     assert "parameter 'param' type annotation" in caplog.text
-    assert 'list subtyping may not be dynamic in INTERSECT' in caplog.text
+    assert 'dynamic typing is not allowed for INTERSECT schemas' in caplog.text
 
 
 def test_disallow_dynamic_list_subtyping_complex(caplog: pytest.LogCaptureFixture):
@@ -227,7 +227,7 @@ def test_disallow_dynamic_set_subtyping(caplog: pytest.LogCaptureFixture):
     with pytest.raises(SystemExit):
         get_schema_helper([MockAnySet])
     assert "parameter 'param' type annotation" in caplog.text
-    assert 'set subtyping may not be dynamic in INTERSECT' in caplog.text
+    assert 'dynamic typing is not allowed for INTERSECT schemas' in caplog.text
 
 
 def test_disallow_dynamic_frozenset_subtyping(caplog: pytest.LogCaptureFixture):
@@ -241,7 +241,7 @@ def test_disallow_dynamic_frozenset_subtyping(caplog: pytest.LogCaptureFixture):
     with pytest.raises(SystemExit):
         get_schema_helper([MockAnyFrozenSet])
     assert "parameter 'param' type annotation" in caplog.text
-    assert 'frozenset subtyping may not be dynamic in INTERSECT' in caplog.text
+    assert 'dynamic typing is not allowed for INTERSECT schemas' in caplog.text
 
 
 def test_disallow_dynamic_generator_subtyping(caplog: pytest.LogCaptureFixture):
@@ -288,7 +288,7 @@ def test_disallow_dynamic_dict_value_type(caplog: pytest.LogCaptureFixture):
     with pytest.raises(SystemExit):
         get_schema_helper([MockAnyDictValue])
     assert "parameter 'param' type annotation" in caplog.text
-    assert 'dict or mapping: value type cannot be Any/object for INTERSECT' in caplog.text
+    assert 'dynamic typing is not allowed for INTERSECT schemas' in caplog.text
 
 
 def test_disallow_dynamic_tuple_subtyping(caplog: pytest.LogCaptureFixture):

intersect_sdk-0.8.4/tests/unit/test_schema_ref_resolution.py

@@ -0,0 +1,55 @@
+import json
+from pathlib import Path
+from typing import Iterable, List
+
+from referencing import Registry, Resource
+from referencing.exceptions import Unresolvable
+from referencing.jsonschema import DRAFT202012
+
+from intersect_sdk.schema import get_schema_from_capability_implementations
+from tests.fixtures.example_schema import FAKE_HIERARCHY_CONFIG, DummyCapabilityImplementation
+
+
+def get_fixture_path(fixture: str) -> Path:
+    return Path(__file__).absolute().parents[1] / 'fixtures' / fixture
+
+
+def iter_ref_values(schema: object) -> Iterable[str]:
+    if isinstance(schema, dict):
+        for key, value in schema.items():
+            if key == '$ref' and isinstance(value, str):
+                yield value
+            else:
+                yield from iter_ref_values(value)
+    elif isinstance(schema, list):
+        for item in schema:
+            yield from iter_ref_values(item)
+
+
+def assert_refs_resolve(schema: dict) -> None:
+    resource = Resource.from_contents(schema, default_specification=DRAFT202012)
+    base_uri = ''
+    resolver = Registry().with_resource(base_uri, resource).resolver(base_uri)
+    unresolved: List[str] = []
+    for ref in iter_ref_values(schema):
+        try:
+            resolver.lookup(ref)
+        except Unresolvable:
+            unresolved.append(ref)
+    assert not unresolved, f'Unresolved $ref entries: {sorted(set(unresolved))}'
+
+
+def test_fixture_schema_refs_resolve() -> None:
+    """Verifies that the generated JSON schema contains valid $ref entries."""
+    with Path.open(get_fixture_path('example_schema.json'), 'rb') as f:
+        schema = json.load(f)
+    assert_refs_resolve(schema)
+
+
+def test_generated_schema_refs_resolve() -> None:
+    """Assert that the INTERSECT-SDK can generate a schema with valid $ref entries."""
+    schema = get_schema_from_capability_implementations(
+        [DummyCapabilityImplementation],
+        FAKE_HIERARCHY_CONFIG,
+    )
+    assert_refs_resolve(schema)