Add configurable prefetch_count (#20)

* Add basic implementation of amqp_params

* Fix typing

* Update docs and tests

* Fix tests

Author: WiZeYAR

.devcontainer/devcontainer.json

@@ -25,6 +25,7 @@
 				"github.vscode-github-actions"
 			],
 			"settings": {
+				"python.defaultInterpreterPath": "/workspaces/${localWorkspaceFolderBasename}/.venv/bin/python",
 				"yaml.schemas": {
 					"https://asyncapi.com/schema-store/3.0.0-without-$id.json": [
 						"file:///workspaces/asyncapi-python/examples/*.yaml"

README.md

@@ -76,7 +76,7 @@ asyncapi_python_service(
 )
 ```
 
-This will be generating python module named `asyncapi_app` on `codegen-export` and `export` goals.
+This will be generating python module named `asyncapi_app` on `codegen-export`, and `export` goals.
 This target can later be used as a dependency of `python_sources`.
 
 ```python
@@ -100,51 +100,16 @@ Note that this plugin does not do dependency injection, so asyncapi-python must
 asyncapi-python[amqp]
 ```
 
-### Deploying this plugin into pants monorepo
-
-In order to deploy this plugin into your pants monorepo, create the following structure inside your plugins folder:
-
-```bash
-pants-plugins/
-└── asyncapi_python_plugin
-    ├── BUILD
-    ├── __init__.py
-    ├── register.py
-    └── requirements.txt
-```
-
-`requirements.txt` must contain:
-
-```text
-asyncapi-python
-```
-
-`register.py` should have:
-
-```python
-from asyncapi_python_pants.register import *
-```
-
-`BUILD` must include:
-
-```python
-python_sources(
-    dependencies=[":reqs"],
-)
-
-python_requirements(
-    name="reqs",
-)
-```
-
-`__init__.py` can be empty, but it has to exist.
-
-Finally, add `pants-plugins` to your `PYTHONPATH`, and add the created folder as a backend package:
+### Deploying this plugin into your pants monorepo
 
 ```toml
 # pants.toml
+plugins = [
+  "asyncapi_python[codegen]==0.2.5",  # Plugin version MUST match the version of your python clients
+  ...
+]
 backend_packages = [
-    "asyncapi_python_plugin",
+    "asyncapi_python_pants",
     ...
 ]
 pythonpath = ["%(buildroot)s/pants-plugins"]

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "asyncapi-python"
-version = "0.2.4"
+version = "0.2.5"
 license = "Apache-2.0"
 description = "Easily generate type-safe and async Python applications from AsyncAPI 3 specifications."
 authors = ["Yaroslav Petrov <yaroslav.v.petrov@gmail.com>"]
@@ -52,9 +52,10 @@ pex = "*"
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
 
-[tool.pytest.ini_options]
-asyncio_mode = "auto"
-asyncio_default_fixture_loop_scope = "function"
-
 [tool.poetry.requires-plugins]
 poetry-plugin-export = ">=1.8"
+
+[tool.pytest.ini_options]
+asyncio_mode = "auto"
+asyncio_default_fixture_loop_scope = "session"
+asyncio_default_test_loop_scope = "session"

src/asyncapi_python/amqp/__init__.py

@@ -19,9 +19,11 @@
 from .operation import Operation
 from .utils import union_model
 from .error import Rejection, RejectedError
+from .params import AmqpParams
 
 __all__ = [
     "channel_pool",
+    "AmqpParams",
     "AmqpPool",
     "BaseApplication",
     "Router",

src/asyncapi_python/amqp/base_application.py

@@ -23,6 +23,7 @@
 from .endpoint import EndpointParams
 from .connection import channel_pool
 from .utils import encode_message, decode_message
+from .params import AmqpParams
 from typing import Generic, Optional, TypeVar
 
 
@@ -61,6 +62,7 @@ def __init__(
         amqp_uri: str,
         producer_factory: type[P],
         consumer_factory: type[C],
+        amqp_params: AmqpParams,
     ):
         self.__params = EndpointParams(
             pool=channel_pool(amqp_uri),
@@ -69,6 +71,7 @@ def __init__(
             register_correlation_id=self.__register_correlation_id,
             stop_application=self.stop,
             app_id=str(uuid4()),
+            amqp_params=amqp_params,
         )
         self.__reply_futures: dict[
             str,

src/asyncapi_python/amqp/endpoint/base.py

@@ -31,6 +31,7 @@
 from ..error import Rejection, RejectedError
 from ..connection import AmqpPool
 from ..operation import Operation
+from ..params import AmqpParams
 from aio_pika.abc import (
     AbstractRobustChannel,
     AbstractRobustQueue,
@@ -64,6 +65,7 @@ class EndpointParams:
     ]
     app_id: str
     stop_application: Callable[[], Awaitable[None]]
+    amqp_params: AmqpParams
 
     @property
     def reply_queue_name(self) -> str:

src/asyncapi_python/amqp/endpoint/receiver.py

@@ -45,6 +45,8 @@ async def start(self) -> None:
         print("start", self._op)
         if self._fn:
             async with self._params.pool.acquire() as ch:
+                if prefetch_count := self._params.amqp_params.get("prefetch_count"):
+                    await ch.set_qos(prefetch_count=prefetch_count)
                 q = self._queue = await self._declare(ch)
                 self._consumer_tag = await q.consume(self._consumer)
             return

src/asyncapi_python/amqp/params.py

@@ -0,0 +1,5 @@
+from typing import TypedDict
+
+
+class AmqpParams(TypedDict, total=False):
+    prefetch_count: int

src/asyncapi_python_codegen/generators/amqp/templates/application.py.j2

@@ -13,9 +13,9 @@
    limitations under the License. #}
 from .consumer import _Router_0 as Consumer
 from .producer import _Router_0 as Producer
-from asyncapi_python.amqp.base_application import BaseApplication
+from asyncapi_python.amqp import BaseApplication, AmqpParams
 
 
 class Application(BaseApplication[Producer, Consumer]):
-    def __init__(self, amqp_uri: str):
-        super().__init__(amqp_uri, Producer, Consumer)
+    def __init__(self, amqp_uri: str, amqp_params: AmqpParams = {}):
+        super().__init__(amqp_uri, Producer, Consumer, amqp_params)

tests/core/amqp/conftest.py

@@ -20,5 +20,6 @@
 
 @pytest_asyncio.fixture(scope="function")
 async def amqp_pool(amqp_uri: str) -> AsyncGenerator[AmqpPool, None]:
+    channel_pool.cache_clear()
     pool = channel_pool(amqp_uri)
     yield pool