From c51e7351e8ec2de212eca4b20c98279e43454c14 Mon Sep 17 00:00:00 2001
From: Vlad Stan <stan.v.vlad@gmail.com>
Date: Wed, 21 Feb 2024 12:08:37 +0200
Subject: [PATCH] fix: extension stop background work (#2281)

* feat: add helper methods

* fix: do not try to stop background work on first install

* fix: first stop via function call then try REST API

* fix: `make check`

* fix: prepare for `{ext_id}_stop`
---
 lnbits/core/helpers.py   | 43 +++++++++++++++++++++++++++++++++++++++-
 lnbits/core/views/api.py |  5 +++--
 lnbits/settings.py       | 10 ++++++++++
 3 files changed, 55 insertions(+), 3 deletions(-)

diff --git a/lnbits/core/helpers.py b/lnbits/core/helpers.py
index a89797989..839d58b71 100644
--- a/lnbits/core/helpers.py
+++ b/lnbits/core/helpers.py
@@ -53,8 +53,45 @@ async def stop_extension_background_work(
 ):
     """
     Stop background work for extension (like asyncio.Tasks, WebSockets, etc).
-    Extensions SHOULD expose a DELETE enpoint at the root level of their API.
+    Extensions SHOULD expose a `api_stop()` function and/or a DELETE enpoint
+    at the root level of their API.
     """
+    stopped = await _stop_extension_background_work(ext_id)
+
+    if not stopped:
+        # fallback to REST API call
+        await _stop_extension_background_work_via_api(ext_id, user, access_token)
+
+
+async def _stop_extension_background_work(ext_id) -> bool:
+    upgrade_hash = settings.extension_upgrade_hash(ext_id) or ""
+    ext = Extension(ext_id, True, False, upgrade_hash=upgrade_hash)
+
+    try:
+        logger.info(f"Stopping background work for extension '{ext.module_name}'.")
+        old_module = importlib.import_module(ext.module_name)
+
+        # Extensions must expose an `{ext_id}_stop()` function at the module level
+        # The `api_stop()` function is for backwards compatibility (will be deprecated)
+        stop_fns = [f"{ext_id}_stop", "api_stop"]
+        stop_fn_name = next((fn for fn in stop_fns if hasattr(old_module, fn)), None)
+        assert stop_fn_name, "No stop function found for '{ext.module_name}'"
+
+        await getattr(old_module, stop_fn_name)()
+
+        logger.info(f"Stopped background work for extension '{ext.module_name}'.")
+    except Exception as ex:
+        logger.warning(f"Failed to stop background work for '{ext.module_name}'.")
+        logger.warning(ex)
+        return False
+
+    return True
+
+
+async def _stop_extension_background_work_via_api(ext_id, user, access_token):
+    logger.info(
+        f"Stopping background work for extension '{ext_id}' using the REST API."
+    )
     async with httpx.AsyncClient() as client:
         try:
             url = f"http://{settings.host}:{settings.port}/{ext_id}/api/v1?usr={user}"
@@ -63,7 +100,11 @@ async def stop_extension_background_work(
             )
             resp = await client.delete(url=url, headers=headers)
             resp.raise_for_status()
+            logger.info(f"Stopped background work for extension '{ext_id}'.")
         except Exception as ex:
+            logger.warning(
+                f"Failed to stop background work for '{ext_id}' using the REST API."
+            )
             logger.warning(ex)
 
 
diff --git a/lnbits/core/views/api.py b/lnbits/core/views/api.py
index 82db728f7..74fd980e6 100644
--- a/lnbits/core/views/api.py
+++ b/lnbits/core/views/api.py
@@ -824,8 +824,9 @@ async def api_install_extension(
 
         await add_installed_extension(ext_info)
 
-        # call stop while the old routes are still active
-        await stop_extension_background_work(data.ext_id, user.id, access_token)
+        if extension.is_upgrade_extension:
+            # call stop while the old routes are still active
+            await stop_extension_background_work(data.ext_id, user.id, access_token)
 
         if data.ext_id not in settings.lnbits_deactivated_extensions:
             settings.lnbits_deactivated_extensions += [data.ext_id]
diff --git a/lnbits/settings.py b/lnbits/settings.py
index 534504177..1cbc667d7 100644
--- a/lnbits/settings.py
+++ b/lnbits/settings.py
@@ -68,6 +68,16 @@ class InstalledExtensionsSettings(LNbitsSettings):
     # list of redirects that extensions want to perform
     lnbits_extensions_redirects: List[Any] = Field(default=[])
 
+    def extension_upgrade_path(self, ext_id: str) -> Optional[str]:
+        return next(
+            (e for e in self.lnbits_upgraded_extensions if e.endswith(f"/{ext_id}")),
+            None,
+        )
+
+    def extension_upgrade_hash(self, ext_id: str) -> Optional[str]:
+        path = settings.extension_upgrade_path(ext_id)
+        return path.split("/")[0] if path else None
+
 
 class ThemesSettings(LNbitsSettings):
     lnbits_site_title: str = Field(default="LNbits")