diff --git a/libervia/backend/memory/sqla.py b/libervia/backend/memory/sqla.py index 27a7efc7..2b265cb8 100644 --- a/libervia/backend/memory/sqla.py +++ b/libervia/backend/memory/sqla.py @@ -1254,25 +1254,23 @@ class Storage: if create_kwargs is None: create_kwargs = {} try: - return await as_future( - self.set_pubsub_node(client, service, name, **create_kwargs) + return await self.set_pubsub_node.aio_func( + client, service, name, **create_kwargs ) except IntegrityError as e: if "unique" in str(e.orig).lower(): # the node may already exist, if it has been created just after # get_pubsub_node above log.debug("ignoring UNIQUE constraint error") - return await as_future( - self.get_pubsub_node( - client, - service, - name=name, - with_items=with_items, - with_subscriptions=with_subscriptions, - with_affiliations=with_affiliations, - with_parent=with_parent, - with_link=with_link, - ) + return await self.get_pubsub_node.aio_func( + client, + service, + name=name, + with_items=with_items, + with_subscriptions=with_subscriptions, + with_affiliations=with_affiliations, + with_parent=with_parent, + with_link=with_link, ) else: raise e @@ -2059,7 +2057,7 @@ class Storage: If affiliation is None, the entity's affiliation is removed. @raise exceptions.NotFound: Node doesn't exist. """ - node = await as_future(self.get_pubsub_node(client, service, node_name)) + node = await self.get_pubsub_node.aio_func(client, service, node_name) if node is None: raise exceptions.NotFound(f"Node {node_name!r} not found") diff --git a/libervia/backend/tools/utils.py b/libervia/backend/tools/utils.py index 0eaf663d..506dd87e 100644 --- a/libervia/backend/tools/utils.py +++ b/libervia/backend/tools/utils.py @@ -18,7 +18,7 @@ """various useful methods""" -from typing import Any, Optional, Union +from typing import Any, Callable, Optional, Union import unicodedata import os.path import datetime @@ -134,15 +134,20 @@ async def maybe_async(ret: Any) -> Any: return ret -def aio(func): - """Decorator to return a Deferred from asyncio coroutine +def aio(func: Callable[..., Any]) -> Callable[..., Any]: + """Decorator to return a Deferred from asyncio coroutine. - Functions with this decorator are run in asyncio context + Functions with this decorator are run in asyncio context. + + The original coroutine is exposed as func.aio_func for direct awaiting (e.g. when + called from another @aio function to avoid spawning an extra task via ensure_future). """ - def wrapper(*args, **kwargs): + @functools.wraps(func) + def wrapper(*args: Any, **kwargs: Any) -> defer.Deferred: return defer.Deferred.fromFuture(asyncio.ensure_future(func(*args, **kwargs))) + wrapper.aio_func = func # type: ignore return wrapper