From 8bae5e3d7ecec68f38401953f97a96682a1da94e Mon Sep 17 00:00:00 2001
From: Goffi <goffi@goffi.org>
Date: Wed, 20 Jul 2022 17:41:44 +0200
Subject: [PATCH] core (memory/sqla): add argument to `getPubsubNode` to
 auto-create a node:

if the node is not in cache, the new `create` and `create_kwargs` arguments can be used to
automatically call `setPubsubNode`

rel 370
---
 sat/memory/sqla.py | 40 +++++++++++++++++++++++++++++++++++++---
 1 file changed, 37 insertions(+), 3 deletions(-)

diff --git a/sat/memory/sqla.py b/sat/memory/sqla.py
index 29cc2c9e..38897e3d 100644
--- a/sat/memory/sqla.py
+++ b/sat/memory/sqla.py
@@ -47,7 +47,7 @@ from sat.core import exceptions
 from sat.core.log import getLogger
 from sat.core.constants import Const as C
 from sat.core.core_types import SatXMPPEntity
-from sat.tools.utils import aio
+from sat.tools.utils import aio, as_future
 from sat.tools.common import uri
 from sat.memory import migration
 from sat.memory import sqla_config
@@ -1059,8 +1059,18 @@ class Storage:
         name: str,
         with_items: bool = False,
         with_subscriptions: bool = False,
+        create: bool = False,
+        create_kwargs: Optional[dict] = None
     ) -> Optional[PubsubNode]:
-        """
+        """Retrieve a PubsubNode from DB
+
+        @param service: service hosting the node
+        @param name: node's name
+        @param with_items: retrieve items in the same query
+        @param with_subscriptions: retrieve subscriptions in the same query
+        @param create: if the node doesn't exist in DB, create it
+        @param create_kwargs: keyword arguments to use with ``setPubsubNode`` if the node
+            needs to be created.
         """
         async with self.session() as session:
             stmt = (
@@ -1080,7 +1090,31 @@ class Storage:
                     joinedload(PubsubNode.subscriptions)
                 )
             result = await session.execute(stmt)
-        return result.unique().scalar_one_or_none()
+        ret = result.unique().scalar_one_or_none()
+        if ret is None and create:
+            # we auto-create the node
+            if create_kwargs is None:
+                create_kwargs = {}
+            try:
+                return await as_future(self.setPubsubNode(
+                    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
+                    # getPubsubNode above
+                    log.debug("ignoring UNIQUE constraint error")
+                    cached_node = await as_future(self.getPubsubNode(
+                        client,
+                        service,
+                        name,
+                        with_items=with_items,
+                        with_subscriptions=with_subscriptions
+                    ))
+                else:
+                    raise e
+        else:
+            return ret
 
     @aio
     async def setPubsubNode(