intersect-sdk 0.6.1a1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (68) hide show
  1. intersect_sdk-0.6.1a1/PKG-INFO +39 -0
  2. intersect_sdk-0.6.1a1/README.md +16 -0
  3. intersect_sdk-0.6.1a1/pyproject.toml +222 -0
  4. intersect_sdk-0.6.1a1/src/intersect_sdk/__init__.py +64 -0
  5. intersect_sdk-0.6.1a1/src/intersect_sdk/_internal/__init__.py +4 -0
  6. intersect_sdk-0.6.1a1/src/intersect_sdk/_internal/compression.py.tmp +38 -0
  7. intersect_sdk-0.6.1a1/src/intersect_sdk/_internal/constants.py +10 -0
  8. intersect_sdk-0.6.1a1/src/intersect_sdk/_internal/control_plane/__init__.py +0 -0
  9. intersect_sdk-0.6.1a1/src/intersect_sdk/_internal/control_plane/brokers/__init__.py +0 -0
  10. intersect_sdk-0.6.1a1/src/intersect_sdk/_internal/control_plane/brokers/amqp_client.py +280 -0
  11. intersect_sdk-0.6.1a1/src/intersect_sdk/_internal/control_plane/brokers/broker_client.py +58 -0
  12. intersect_sdk-0.6.1a1/src/intersect_sdk/_internal/control_plane/brokers/mqtt_client.py +154 -0
  13. intersect_sdk-0.6.1a1/src/intersect_sdk/_internal/control_plane/control_plane_manager.py +147 -0
  14. intersect_sdk-0.6.1a1/src/intersect_sdk/_internal/control_plane/discovery_service.py +40 -0
  15. intersect_sdk-0.6.1a1/src/intersect_sdk/_internal/data_plane/__init__.py +0 -0
  16. intersect_sdk-0.6.1a1/src/intersect_sdk/_internal/data_plane/data_plane_manager.py +102 -0
  17. intersect_sdk-0.6.1a1/src/intersect_sdk/_internal/data_plane/minio_utils.py +149 -0
  18. intersect_sdk-0.6.1a1/src/intersect_sdk/_internal/event_metadata.py +60 -0
  19. intersect_sdk-0.6.1a1/src/intersect_sdk/_internal/exceptions.py +17 -0
  20. intersect_sdk-0.6.1a1/src/intersect_sdk/_internal/function_metadata.py +27 -0
  21. intersect_sdk-0.6.1a1/src/intersect_sdk/_internal/interfaces.py +20 -0
  22. intersect_sdk-0.6.1a1/src/intersect_sdk/_internal/logger.py +3 -0
  23. intersect_sdk-0.6.1a1/src/intersect_sdk/_internal/messages/__init__.py +0 -0
  24. intersect_sdk-0.6.1a1/src/intersect_sdk/_internal/messages/event.py +157 -0
  25. intersect_sdk-0.6.1a1/src/intersect_sdk/_internal/messages/lifecycle.py +174 -0
  26. intersect_sdk-0.6.1a1/src/intersect_sdk/_internal/messages/userspace.py +171 -0
  27. intersect_sdk-0.6.1a1/src/intersect_sdk/_internal/pydantic_schema_generator.py +504 -0
  28. intersect_sdk-0.6.1a1/src/intersect_sdk/_internal/schema.py +558 -0
  29. intersect_sdk-0.6.1a1/src/intersect_sdk/_internal/stoppable_thread.py +19 -0
  30. intersect_sdk-0.6.1a1/src/intersect_sdk/_internal/utils.py +32 -0
  31. intersect_sdk-0.6.1a1/src/intersect_sdk/_internal/version_resolver.py +58 -0
  32. intersect_sdk-0.6.1a1/src/intersect_sdk/app_lifecycle.py +147 -0
  33. intersect_sdk-0.6.1a1/src/intersect_sdk/capability/__init__.py +14 -0
  34. intersect_sdk-0.6.1a1/src/intersect_sdk/capability/base.py +101 -0
  35. intersect_sdk-0.6.1a1/src/intersect_sdk/client.py +443 -0
  36. intersect_sdk-0.6.1a1/src/intersect_sdk/client_callback_definitions.py +154 -0
  37. intersect_sdk-0.6.1a1/src/intersect_sdk/config/__init__.py +20 -0
  38. intersect_sdk-0.6.1a1/src/intersect_sdk/config/client.py +52 -0
  39. intersect_sdk-0.6.1a1/src/intersect_sdk/config/service.py +42 -0
  40. intersect_sdk-0.6.1a1/src/intersect_sdk/config/shared.py +168 -0
  41. intersect_sdk-0.6.1a1/src/intersect_sdk/constants.py +9 -0
  42. intersect_sdk-0.6.1a1/src/intersect_sdk/core_definitions.py +37 -0
  43. intersect_sdk-0.6.1a1/src/intersect_sdk/py.typed +0 -0
  44. intersect_sdk-0.6.1a1/src/intersect_sdk/schema.py +77 -0
  45. intersect_sdk-0.6.1a1/src/intersect_sdk/service.py +607 -0
  46. intersect_sdk-0.6.1a1/src/intersect_sdk/service_definitions.py +262 -0
  47. intersect_sdk-0.6.1a1/src/intersect_sdk/version.py +15 -0
  48. intersect_sdk-0.6.1a1/tests/__init__.py +0 -0
  49. intersect_sdk-0.6.1a1/tests/conftest.py +0 -0
  50. intersect_sdk-0.6.1a1/tests/e2e/__init__.py +0 -0
  51. intersect_sdk-0.6.1a1/tests/e2e/test_examples.py +130 -0
  52. intersect_sdk-0.6.1a1/tests/fixtures/__init__.py +0 -0
  53. intersect_sdk-0.6.1a1/tests/fixtures/example_schema.json +1062 -0
  54. intersect_sdk-0.6.1a1/tests/fixtures/example_schema.py +590 -0
  55. intersect_sdk-0.6.1a1/tests/fixtures/return_type_mismatch.py +7 -0
  56. intersect_sdk-0.6.1a1/tests/integration/__init__.py +0 -0
  57. intersect_sdk-0.6.1a1/tests/integration/test_return_type_mismatch.py +122 -0
  58. intersect_sdk-0.6.1a1/tests/integration/test_service.py +378 -0
  59. intersect_sdk-0.6.1a1/tests/unit/__init__.py +0 -0
  60. intersect_sdk-0.6.1a1/tests/unit/test_annotations.py +163 -0
  61. intersect_sdk-0.6.1a1/tests/unit/test_base_capability_implementation.py +117 -0
  62. intersect_sdk-0.6.1a1/tests/unit/test_config.py +193 -0
  63. intersect_sdk-0.6.1a1/tests/unit/test_invalid_schema_runtime.py +42 -0
  64. intersect_sdk-0.6.1a1/tests/unit/test_lifecycle_message.py +100 -0
  65. intersect_sdk-0.6.1a1/tests/unit/test_schema_invalids.py +790 -0
  66. intersect_sdk-0.6.1a1/tests/unit/test_schema_valid.py +79 -0
  67. intersect_sdk-0.6.1a1/tests/unit/test_userspace_message.py +105 -0
  68. intersect_sdk-0.6.1a1/tests/unit/test_version_resolver.py +134 -0
@@ -0,0 +1,39 @@
1
+ Metadata-Version: 2.1
2
+ Name: intersect-sdk
3
+ Version: 0.6.1a1
4
+ Summary: Python SDK to interact with INTERSECT
5
+ Keywords: intersect
6
+ 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>
7
+ Classifier: Programming Language :: Python :: 3
8
+ Classifier: Programming Language :: Python :: 3.8
9
+ Classifier: Programming Language :: Python :: 3.9
10
+ Classifier: Programming Language :: Python :: 3.10
11
+ Classifier: Programming Language :: Python :: 3.11
12
+ Requires-Python: <4.0,>=3.8.10
13
+ Requires-Dist: pydantic>=2.7.0
14
+ Requires-Dist: retrying<2.0.0,>=1.3.4
15
+ Requires-Dist: paho-mqtt<2.0.0,>=1.6.1
16
+ Requires-Dist: minio>=7.2.3
17
+ Requires-Dist: jsonschema[format-nongpl]>=4.21.1
18
+ Requires-Dist: pika<2.0.0,>=1.3.2; extra == "amqp"
19
+ Requires-Dist: eval-type-backport>=0.1.3; extra == "py38"
20
+ Provides-Extra: amqp
21
+ Provides-Extra: py38
22
+ Description-Content-Type: text/markdown
23
+
24
+ # INTERSECT-SDK
25
+
26
+ The INTERSECT-SDK is a framework for microservices to integrate themselves into the wider INTERSECT ecosystem.
27
+
28
+ Please note that this README is currently a work in progress.
29
+
30
+ ## What is INTERSECT?
31
+
32
+ INTERSECT was designed as a specific usecase - as an open federated hardware/software architecture for the laboratory of the future, which connects scientific instruments, robot-controlled laboratories and edge/center computing/data resources to enable autonomous experiments, self-driving laboratories, smart manufacturing, and AI-driven design, discovery and evaluation.
33
+
34
+ ## What are the core design philosophies of the SDK?
35
+
36
+ - Event-driven architecture
37
+ - Support core interaction types: request/response, events, commands, statuses
38
+ - Borrows several concepts from [AsyncAPI](https://www.asyncapi.com/docs/reference/specification/latest), and intends to support multiple different protocols. Currently, we support MQTT 3.1.1 and AMQP 0.9.1, but other protocols will be supported as well.
39
+ - Users automatically generate schema from code; schemas are part of the core contract of an INTERSECT microservice, and both external inputs and microservice outputs are required to uphold this contract.
@@ -0,0 +1,16 @@
1
+ # INTERSECT-SDK
2
+
3
+ The INTERSECT-SDK is a framework for microservices to integrate themselves into the wider INTERSECT ecosystem.
4
+
5
+ Please note that this README is currently a work in progress.
6
+
7
+ ## What is INTERSECT?
8
+
9
+ INTERSECT was designed as a specific usecase - as an open federated hardware/software architecture for the laboratory of the future, which connects scientific instruments, robot-controlled laboratories and edge/center computing/data resources to enable autonomous experiments, self-driving laboratories, smart manufacturing, and AI-driven design, discovery and evaluation.
10
+
11
+ ## What are the core design philosophies of the SDK?
12
+
13
+ - Event-driven architecture
14
+ - Support core interaction types: request/response, events, commands, statuses
15
+ - Borrows several concepts from [AsyncAPI](https://www.asyncapi.com/docs/reference/specification/latest), and intends to support multiple different protocols. Currently, we support MQTT 3.1.1 and AMQP 0.9.1, but other protocols will be supported as well.
16
+ - Users automatically generate schema from code; schemas are part of the core contract of an INTERSECT microservice, and both external inputs and microservice outputs are required to uphold this contract.
@@ -0,0 +1,222 @@
1
+ [project]
2
+ name = "intersect-sdk"
3
+ description = "Python SDK to interact with INTERSECT"
4
+ authors = [
5
+ { name = "Lance Drane", email = "dranelt@ornl.gov" },
6
+ { name = "Marshall McDonnell", email = "mcdonnellmt@ornl.gov" },
7
+ { name = "Seth Hitefield", email = "hitefieldsd@ornl.gov" },
8
+ { name = "Andrew Ayres", email = "ayresaf@ornl.gov" },
9
+ { name = "Gregory Cage", email = "cagege@ornl.gov" },
10
+ { name = "Jesse McGaha", email = "mcgahajr@ornl.gov" },
11
+ { name = "Robert Smith", email = "smithrw@ornl.gov" },
12
+ { name = "Gavin Wiggins", email = "wigginsg@ornl.gov" },
13
+ ]
14
+ readme = "README.md"
15
+ requires-python = ">=3.8.10,<4.0"
16
+ keywords = [
17
+ "intersect",
18
+ ]
19
+ version = "0.6.1a1"
20
+ classifiers = [
21
+ "Programming Language :: Python :: 3",
22
+ "Programming Language :: Python :: 3.8",
23
+ "Programming Language :: Python :: 3.9",
24
+ "Programming Language :: Python :: 3.10",
25
+ "Programming Language :: Python :: 3.11",
26
+ ]
27
+ dependencies = [
28
+ "pydantic>=2.7.0",
29
+ "retrying>=1.3.4,<2.0.0",
30
+ "paho-mqtt>=1.6.1,<2.0.0",
31
+ "minio>=7.2.3",
32
+ "jsonschema[format-nongpl]>=4.21.1",
33
+ ]
34
+
35
+ [project.optional-dependencies]
36
+ amqp = [
37
+ "pika>=1.3.2,<2.0.0",
38
+ ]
39
+ py38 = [
40
+ "eval-type-backport>=0.1.3",
41
+ ]
42
+
43
+ [tool.pdm.dev-dependencies]
44
+ lint = [
45
+ "pre-commit>=3.3.1",
46
+ "ruff>=0.4.2",
47
+ "mypy>=1.10.0",
48
+ "types-paho-mqtt>=1.6.0.20240106",
49
+ ]
50
+ test = [
51
+ "pytest>=7.3.2",
52
+ "pytest-cov>=4.1.0",
53
+ "httpretty>=1.1.4",
54
+ ]
55
+ doc = [
56
+ "sphinx>=5.3.0",
57
+ "furo>=2023.3.27",
58
+ ]
59
+ ":lint" = []
60
+
61
+ [tool.pdm.scripts]
62
+ test-all = "pytest tests/ --cov=src/intersect_sdk/ --cov-fail-under=80 --cov-report=html:reports/htmlcov/ --cov-report=xml:reports/coverage_report.xml --junitxml=reports/junit.xml"
63
+ test-all-debug = "pytest tests/ --cov=src/intersect_sdk/ --cov-fail-under=80 --cov-report=html:reports/htmlcov/ --cov-report=xml:reports/coverage_report.xml --junitxml=reports/junit.xml -s"
64
+ test-unit = "pytest tests/unit --cov=src/intersect_sdk/"
65
+ lint-format = "ruff format"
66
+ lint-ruff = "ruff check --fix"
67
+ lint-mypy = "mypy src/intersect_sdk/"
68
+
69
+ [tool.pdm.scripts.lint]
70
+ composite = [
71
+ "lint-format",
72
+ "lint-ruff",
73
+ "lint-mypy",
74
+ ]
75
+
76
+ [tool.pdm.build]
77
+ package-dir = "src"
78
+
79
+ [tool.pdm.version]
80
+ source = "file"
81
+ path = "src/intersect_sdk/version.py"
82
+
83
+ [tool.ruff]
84
+ line-length = 100
85
+
86
+ [tool.ruff.format]
87
+ quote-style = "single"
88
+
89
+ [tool.ruff.lint]
90
+ extend-select = [
91
+ "C90",
92
+ "I",
93
+ "N",
94
+ "D",
95
+ "UP",
96
+ "YTT",
97
+ "ANN",
98
+ "ASYNC",
99
+ "S",
100
+ "BLE",
101
+ "B",
102
+ "A",
103
+ "COM",
104
+ "C4",
105
+ "DTZ",
106
+ "T10",
107
+ "EM",
108
+ "FA",
109
+ "ISC",
110
+ "ICN",
111
+ "G",
112
+ "INP",
113
+ "PIE",
114
+ "T20",
115
+ "PYI",
116
+ "PT",
117
+ "Q",
118
+ "RSE",
119
+ "RET",
120
+ "SLF",
121
+ "SLOT",
122
+ "SIM",
123
+ "TCH",
124
+ "ARG",
125
+ "PTH",
126
+ "PGH",
127
+ "PL",
128
+ "TRY",
129
+ "FLY",
130
+ "RUF",
131
+ ]
132
+ ignore = [
133
+ "COM812",
134
+ "ISC001",
135
+ "SIM105",
136
+ "ANN101",
137
+ "ANN102",
138
+ "ANN401",
139
+ "PLR2004",
140
+ ]
141
+
142
+ [tool.ruff.lint.isort]
143
+ known-first-party = [
144
+ "src",
145
+ ]
146
+
147
+ [tool.ruff.lint.pydocstyle]
148
+ convention = "google"
149
+
150
+ [tool.ruff.lint.flake8-quotes]
151
+ inline-quotes = "single"
152
+ multiline-quotes = "double"
153
+
154
+ [tool.ruff.lint.mccabe]
155
+ max-complexity = 20
156
+
157
+ [tool.ruff.lint.pylint]
158
+ max-args = 10
159
+ max-branches = 20
160
+ max-returns = 10
161
+
162
+ [tool.ruff.lint.extend-per-file-ignores]
163
+ "__init__.py" = [
164
+ "F401",
165
+ ]
166
+ "docs/*" = [
167
+ "D",
168
+ "INP001",
169
+ ]
170
+ "examples/*" = [
171
+ "N999",
172
+ "T20",
173
+ "S106",
174
+ "D100",
175
+ "D104",
176
+ "TRY002",
177
+ "FA100",
178
+ ]
179
+ "tests/*" = [
180
+ "S101",
181
+ "S106",
182
+ "S311",
183
+ "SLF001",
184
+ "ANN",
185
+ "ARG",
186
+ "D",
187
+ "FA100",
188
+ ]
189
+
190
+ [tool.mypy]
191
+ strict = true
192
+ ignore_missing_imports = true
193
+ disallow_untyped_decorators = false
194
+ plugins = [
195
+ "pydantic.mypy",
196
+ ]
197
+
198
+ [tool.pydantic-mypy]
199
+ init_forbid_extra = true
200
+ init_typed = true
201
+ warn_required_dynamic_aliases = true
202
+ warn_untyped_fields = true
203
+
204
+ [tool.pytest.ini_options]
205
+ log_cli = true
206
+ addopts = "-ra"
207
+
208
+ [tool.coverage.report]
209
+ omit = [
210
+ "*__init__*",
211
+ ]
212
+ exclude_also = [
213
+ "pragma: no-cover",
214
+ "if (typing\\\\.)?TYPE_CHECKING:",
215
+ "@(abc\\\\.)?abstractmethod",
216
+ ]
217
+
218
+ [build-system]
219
+ requires = [
220
+ "pdm-backend",
221
+ ]
222
+ build-backend = "pdm.backend"
@@ -0,0 +1,64 @@
1
+ """The root module contains the intended public API for users of the INTERSECT-SDK.
2
+
3
+ Users should not need to import anything outside of the root.
4
+
5
+ In general, most breaking changes on version updates will relate to:
6
+ - Configuration classes (both adding and removing new config models). These configuration classes are relevant to the next point.
7
+ - When a new data service is integrated into INTERSECT, ALL adapters will need to update to support this data service, which will include new dependencies.
8
+ """
9
+
10
+ from .app_lifecycle import default_intersect_lifecycle_loop
11
+ from .capability.base import IntersectBaseCapabilityImplementation
12
+ from .client import IntersectClient
13
+ from .client_callback_definitions import (
14
+ INTERSECT_CLIENT_EVENT_CALLBACK_TYPE,
15
+ INTERSECT_CLIENT_RESPONSE_CALLBACK_TYPE,
16
+ INTERSECT_JSON_VALUE,
17
+ IntersectClientCallback,
18
+ IntersectClientMessageParams,
19
+ )
20
+ from .config.client import IntersectClientConfig
21
+ from .config.service import IntersectServiceConfig
22
+ from .config.shared import (
23
+ ControlPlaneConfig,
24
+ DataStoreConfig,
25
+ DataStoreConfigMap,
26
+ HierarchyConfig,
27
+ )
28
+ from .core_definitions import IntersectDataHandler, IntersectMimeType
29
+ from .schema import get_schema_from_capability_implementation
30
+ from .service import IntersectService
31
+ from .service_definitions import (
32
+ IntersectEventDefinition,
33
+ intersect_event,
34
+ intersect_message,
35
+ intersect_status,
36
+ )
37
+ from .version import __version__, version_info
38
+
39
+ __all__ = [
40
+ 'IntersectDataHandler',
41
+ 'IntersectEventDefinition',
42
+ 'IntersectMimeType',
43
+ 'intersect_event',
44
+ 'intersect_message',
45
+ 'intersect_status',
46
+ 'get_schema_from_capability_implementation',
47
+ 'IntersectService',
48
+ 'IntersectClient',
49
+ 'IntersectClientCallback',
50
+ 'IntersectClientMessageParams',
51
+ 'INTERSECT_CLIENT_RESPONSE_CALLBACK_TYPE',
52
+ 'INTERSECT_CLIENT_EVENT_CALLBACK_TYPE',
53
+ 'INTERSECT_JSON_VALUE',
54
+ 'IntersectBaseCapabilityImplementation',
55
+ 'default_intersect_lifecycle_loop',
56
+ 'IntersectClientConfig',
57
+ 'IntersectServiceConfig',
58
+ 'HierarchyConfig',
59
+ 'ControlPlaneConfig',
60
+ 'DataStoreConfig',
61
+ 'DataStoreConfigMap',
62
+ '__version__',
63
+ 'version_info',
64
+ ]
@@ -0,0 +1,4 @@
1
+ """NOTE: It's not advised for users to import anything in this module. These are internal definitions essential for INTERSECT's functions.
2
+
3
+ THERE IS NO GUARANTEE ANY API DEFINED HERE REMAINS STABLE.
4
+ """
@@ -0,0 +1,38 @@
1
+ """Compression logic used inside INTERSECT.
2
+
3
+ Our current compression algorithm is Brotli - see https://stackoverflow.com/a/59255343 for a fairly comprehensive overview of serialization formats compatible with generic JSON.
4
+ Zstd is worth considering for instances where data isn't being persisted for long, as it is generally faster than Brotli (especially for decompression) - however, Zstd requires multithreading to maximize its potential.
5
+
6
+ Currently, we will ALWAYS compress the following:
7
+ - The schemas sent as part of the lifecycle message
8
+ - The PAYLOADS of UserspaceMessages *if the data itself is being sent in UserspaceMessages*
9
+
10
+ We do NOT want to compress message headers, though.
11
+
12
+ MINIO should always handle compression itself, see https://min.io/docs/minio/linux/administration/object-management/data-compression.html for details.
13
+
14
+ TODO - how should the data API handle compression?
15
+
16
+ TODO - We should consider NOT compressing audio, video, image, or any data which is already compressed.
17
+ """
18
+
19
+ import brotli
20
+
21
+
22
+ def intersect_compress(message: bytes) -> bytes:
23
+ """Compress MESSAGE using the compression algorithm, and return compression result.
24
+
25
+ Current compression algorithm is Brotli with the highest level quality.
26
+ """
27
+ return brotli.compress(message, quality=11) # type: ignore[no-any-return]
28
+
29
+ def intersect_decompress(compressed: bytes) -> bytes:
30
+ """Decompress COMPRESSED using the decompression algorithm, and return decompressed message.
31
+
32
+ Current decompression algorithm is Brotli.
33
+ """
34
+ try:
35
+ return brotli.decompress(compressed) # type: ignore[no-any-return]
36
+ except brotli.error:
37
+ # if the parameters weren't compressed with Brotli, just return the input
38
+ return compressed
@@ -0,0 +1,10 @@
1
+ BASE_RESPONSE_ATTR = '__is_intersect_response__'
2
+ BASE_STATUS_ATTR = '__is_intersect_status__'
3
+ BASE_EVENT_ATTR = '__is_intersect_event__'
4
+ # in theory, as long as the next attributes are unique, they can be any string
5
+ REQUEST_CONTENT = '__request_content_type__'
6
+ RESPONSE_CONTENT = '__response_content_type__'
7
+ RESPONSE_DATA = '__response_data_transfer_handler__'
8
+ STRICT_VALIDATION = '__strict_validation__'
9
+ SHUTDOWN_KEYS = '__ignore_message__'
10
+ EVENT_ATTR_KEY = '__intersect_sdk_events__'
@@ -0,0 +1,280 @@
1
+ from __future__ import annotations
2
+
3
+ import functools
4
+ import threading
5
+ import uuid
6
+ from typing import TYPE_CHECKING, Callable
7
+
8
+ import pika
9
+ from retrying import retry
10
+
11
+ from .broker_client import BrokerClient
12
+
13
+ if TYPE_CHECKING:
14
+ from collections import defaultdict
15
+
16
+ from pika.channel import Channel
17
+ from pika.frame import Frame
18
+ from pika.spec import Basic, BasicProperties
19
+
20
+
21
+ class AMQPClient(BrokerClient):
22
+ """Client for performing broker actions backed by a AMQP broker.
23
+
24
+ NOTE: Currently, thread safety has been attempted, but may not be guaranteed
25
+
26
+ Attributes:
27
+ id: A string representation of the client's UUID.
28
+ _connection_params: connection information to the AMQP broker (includes credentials)
29
+ _publish_connection: AMQP connection dedicated to publishing messages
30
+ _consume_connection: AMQP connection dedicated to consuming messages
31
+ _topics_to_handlers: Dictionary of string topic names to lists of
32
+ Callables to invoke for messages on that topic.
33
+ """
34
+
35
+ def __init__(
36
+ self,
37
+ host: str,
38
+ port: int,
39
+ username: str,
40
+ password: str,
41
+ topics_to_handlers: Callable[[], defaultdict[str, set[Callable[[bytes], None]]]],
42
+ uid: str | None = None,
43
+ ) -> None:
44
+ """The default constructor.
45
+
46
+ Args:
47
+ host: String for hostname of AMQP broker
48
+ port: port number of AMQP broker
49
+ username: username credentials for AMQP broker
50
+ password: password credentials for AMQP broker
51
+ topics_to_handlers: callback function which gets the topic to handler map from the channel manager
52
+ uid: String for the client's UUID.
53
+ """
54
+ self.uid = uid if uid else str(uuid.uuid4())
55
+
56
+ self._connection_params = pika.ConnectionParameters(
57
+ host=host,
58
+ port=port,
59
+ virtual_host='/',
60
+ credentials=pika.PlainCredentials(username, password),
61
+ blocked_connection_timeout=1,
62
+ )
63
+
64
+ # The pika connection to the broker
65
+ self._publish_connection: pika.adapters.BlockingConnection = None
66
+ self._consume_connection: pika.adapters.SelectConnection = None
67
+ self._consume_connection_ready_event = threading.Event()
68
+ self._consume_subscription_ready_event = threading.Event()
69
+
70
+ # Callback to the topics_to_handler list inside of
71
+ self._topics_to_handlers = topics_to_handlers
72
+ # mapping of topics to callables which can unsubscribe from the topic
73
+ self._topics_to_channel_cancel_callbacks: dict[str, Callable[[], None]] = {}
74
+ self._consumer_thread = None
75
+
76
+ @retry(
77
+ stop_max_attempt_number=5,
78
+ wait_exponential_multiplier=1000,
79
+ wait_exponential_max=60000,
80
+ )
81
+ def connect(self) -> None:
82
+ """Connect to the defined broker.
83
+
84
+ Try to connect to the broker, performing exponential backoff if connection fails.
85
+ """
86
+ # need deamon=True otherwise if tests fails it hangs trying to acquire lock
87
+ self.thread = threading.Thread(target=self._start_consuming, daemon=True)
88
+ self.thread.start()
89
+ self._publish_connection = pika.adapters.BlockingConnection(self._connection_params)
90
+ self._consume_connection_ready_event.wait(timeout=5)
91
+
92
+ def disconnect(self) -> None:
93
+ """Close all connections.
94
+
95
+ Close both the public and consume connections and stop the consuming thread.
96
+ """
97
+ self._publish_connection.close()
98
+ self._publish_connection = None
99
+ self._consume_connection.ioloop.add_callback_threadsafe(self._close_consume_connection)
100
+ # as soon as connection is closed, the ioloop.stop will be
101
+ # called which in turn will terminate the consuming thread
102
+ self.thread.join(5)
103
+ self._consume_connection = None
104
+
105
+ def is_connected(self) -> bool:
106
+ """Check if there is an active connection to the broker.
107
+
108
+ Returns:
109
+ A boolean. True if there is a connection, False if not.
110
+ """
111
+ # We are connected to the broker if either the publish or consume connections is open
112
+ return (self._publish_connection is not None and self._publish_connection.is_open) or (
113
+ self._consume_connection is not None and self._consume_connection.is_open
114
+ )
115
+
116
+ def publish(self, topic: str, payload: bytes) -> None:
117
+ """Publish the given message.
118
+
119
+ Publish payload with the pre-existing connection (via connect()) on topic.
120
+
121
+ Args:
122
+ topic: The topic on which to publish the message as a string
123
+ payload: The message to publish, as raw bytes.
124
+ """
125
+ channel = self._publish_connection.channel()
126
+ channel.exchange_declare(topic, exchange_type='fanout', durable=True)
127
+ # this will send the message to topic exchange and distribute to all
128
+ # queues that subscribed to it
129
+ channel.basic_publish(
130
+ exchange=topic,
131
+ routing_key=topic,
132
+ body=payload,
133
+ properties=pika.BasicProperties(content_type='text/plain'),
134
+ )
135
+ channel.close()
136
+
137
+ def subscribe(self, topic: str) -> None:
138
+ self._consume_subscription_ready_event.clear()
139
+ self._subscribe_to_queue(topic)
140
+ self._consume_subscription_ready_event.wait()
141
+
142
+ def unsubscribe(self, topic: str) -> None:
143
+ old_channel = self._topics_to_channel_cancel_callbacks.get(topic, None)
144
+ if old_channel:
145
+ old_channel()
146
+ del self._topics_to_channel_cancel_callbacks[topic]
147
+
148
+ def _start_consuming(self) -> None:
149
+ """Start consuming messages from broker.
150
+
151
+ Open the consuming connection and start its io loop.
152
+ """
153
+ self._consume_connection = pika.adapters.SelectConnection(
154
+ parameters=self._connection_params,
155
+ on_close_callback=(
156
+ lambda _connection, _exception: self._consume_connection.ioloop.stop()
157
+ ),
158
+ on_open_callback=lambda _connection: self._consume_connection_ready_event.set(),
159
+ )
160
+
161
+ self._consume_connection.ioloop.start()
162
+
163
+ def _subscribe_to_queue(self, topic: str) -> None:
164
+ """Start consuming from the given topic.
165
+
166
+ Declares the correct exchange and queue on the broker if needed, then starts
167
+ consuming messages on that queue.
168
+
169
+ Args:
170
+ topic: Name of the topic on the broker to subscribe to as a string.
171
+ """
172
+ # in consumer thread: open channel-> declare queue (need that if we starte
173
+ # consuming before message is published (no queue yet)) -> start consuming
174
+ cb = functools.partial(self._open_channel, topic=topic)
175
+ self._consume_connection.ioloop.add_callback_threadsafe(cb)
176
+
177
+ def _open_channel(self, topic: str) -> None:
178
+ """Open a channel for the given topic.
179
+
180
+ Args:
181
+ topic: The topic to open a channel for as a string.
182
+ """
183
+ cb = functools.partial(self._on_channel_open, topic=topic)
184
+ self._consume_connection.channel(on_open_callback=cb)
185
+
186
+ def _on_channel_open(self, channel: Channel, topic: str) -> None:
187
+ """Create an exchange on the broker.
188
+
189
+ Used as a listener on channel open.
190
+
191
+ Args:
192
+ channel: The Channel being instantiated.
193
+ topic: The string name for the Channel on the broker.
194
+ """
195
+ cb = functools.partial(self._on_exchange_declareok, channel=channel, topic=topic)
196
+ channel.exchange_declare(exchange=topic, exchange_type='fanout', durable=True, callback=cb)
197
+
198
+ def _on_exchange_declareok(self, _unused_frame: Frame, channel: Channel, topic: str) -> None:
199
+ """Create a queue on the broker.
200
+
201
+ Used as a listener on exchange declaration.
202
+
203
+ Args:
204
+ _unused_frame: Object from pika. Ignored.
205
+ channel: The Channel being instantiated.
206
+ topic: The string name for the Channel on the broker.
207
+ """
208
+ cb = functools.partial(self._on_queue_declareok, channel=channel, topic=topic)
209
+ channel.queue_declare(queue=topic + '.' + self.uid, durable=True, callback=cb)
210
+
211
+ def _on_queue_declareok(self, _unused_frame: Frame, channel: Channel, topic: str) -> None:
212
+ """Begins listening on the given queue.
213
+
214
+ Used as a listener on queue declaration.
215
+
216
+ Args:
217
+ _unused_frame: Object from pika. Ignored.
218
+ channel: The Channel being instantiated.
219
+ topic: The string name for the Channel on the broker.
220
+ """
221
+ cb = functools.partial(self._on_queue_bindok, channel=channel, topic=topic)
222
+ channel.queue_bind(topic + '.' + self.uid, topic, routing_key=topic, callback=cb)
223
+
224
+ def _on_queue_bindok(self, _unused_frame: Frame, channel: Channel, topic: str) -> None:
225
+ """Consumes a message from the given channel.
226
+
227
+ Used as a listener on queue binding.
228
+
229
+ Args:
230
+ _unused_frame: Object from pika. Ignored.
231
+ channel: The Channel being instantiated.
232
+ topic: The string name for the Channel on the broker.
233
+ """
234
+ cb = functools.partial(self._on_consume_ok)
235
+ consumer_tag = channel.basic_consume(
236
+ queue=topic + '.' + self.uid,
237
+ auto_ack=True,
238
+ on_message_callback=self._consume_message,
239
+ callback=cb,
240
+ )
241
+ self._topics_to_channel_cancel_callbacks[topic] = lambda: channel.basic_cancel(consumer_tag)
242
+
243
+ def _on_consume_ok(self, _unused_frame: Frame) -> None:
244
+ """Sets the consume subscription ready even.
245
+
246
+ Used as a listener on consuming an initial message on a channel.
247
+
248
+ Args:
249
+ _unused_frame: Object from pika. Ignored.
250
+ topic: The string name for the Channel on the broker.
251
+ """
252
+ self._consume_subscription_ready_event.set()
253
+
254
+ def _consume_message(
255
+ self,
256
+ _unused_channel: Channel,
257
+ basic_deliver: Basic.Deliver,
258
+ _properties: BasicProperties,
259
+ body: bytes,
260
+ ) -> None:
261
+ """Handles incoming messages.
262
+
263
+ Looks up all handlers for the topic and delegates message handling to them.
264
+
265
+ Args:
266
+ _unused_channel: The Pika channel the message was received on. Ignored
267
+ basic_deliver: Contains internal Pika delivery information - i.e. the routing key.
268
+ _properties: Object from the Pika call. Ignored.
269
+ body: the pika message to be handled.
270
+ """
271
+ for handler in self._topics_to_handlers().get(basic_deliver.routing_key, []):
272
+ handler(body)
273
+
274
+ def _close_consume_connection(self) -> None:
275
+ """Closes the consume connection.
276
+
277
+ Used as a listener on the connection loop to safely shutdown.
278
+ """
279
+ self._consume_connection.close()
280
+ self._topics_to_channel_cancel_callbacks.clear()