kraken: defer state mutations until after image pull succeeds

Author: nicoschmdt

core/services/kraken/api/v1/routers/extension.py

@@ -21,7 +21,7 @@
 @extension_to_http_exception
 async def install_extension(body: ExtensionSource) -> StreamingResponse:
     extension = Extension(body)
-    return StreamingResponse(streamer(extension.install(atomic=True)))
+    return StreamingResponse(streamer(extension.install()))
 
 
 @extension_router_v1.post("/uninstall", status_code=status.HTTP_200_OK)

core/services/kraken/api/v2/routers/extension.py

@@ -83,7 +83,7 @@ async def install(body: ExtensionSource) -> StreamingResponse:
     can install incompatible extensions. Make sure to check the extension source before installing it.
     """
     extension = Extension(body)
-    return StreamingResponse(streamer(extension.install(atomic=True)))
+    return StreamingResponse(streamer(extension.install()))
 
 
 @extension_router_v2.post("/{identifier}/install", status_code=status.HTTP_201_CREATED)
@@ -149,13 +149,13 @@ async def update_to_latest(identifier: str, purge: bool = True, stable: bool = T
 
 @extension_router_v2.put("/{identifier}/{tag}", status_code=status.HTTP_200_OK)
 @extension_to_http_exception
-async def update_to_tag(identifier: str, tag: str, purge: bool = True) -> Response:
+async def update_to_tag(identifier: str, tag: str, purge: bool = True, should_enable: bool = True) -> Response:
     """
     Update a given extension by its identifier and tag to latest version on the higher priority manifest and by default
     purge all other tags, if purge is set to false it will keep all other versions disabled only.
     """
     extension = cast(Extension, await Extension.from_manifest(identifier, tag))
-    return StreamingResponse(streamer(extension.update(purge)))
+    return StreamingResponse(streamer(extension.update(purge, should_enable)))
 
 
 @extension_router_v2.delete("/{identifier}", status_code=status.HTTP_202_ACCEPTED)
@@ -254,4 +254,4 @@ async def finalize_extension(
     new_extension = await Extension.finalize_temporary_extension(temp_extension, body.identifier, body)
 
     # Install the extension
-    return StreamingResponse(streamer(new_extension.install(atomic=True)))
+    return StreamingResponse(streamer(new_extension.install()))

core/services/kraken/extension/extension.py

@@ -171,18 +171,17 @@ async def _disable_running_extension(self) -> Optional["Extension"]:
         except ExtensionNotRunning:
             return None
 
-    def _create_extension_settings(self) -> ExtensionSettings:
+    def _create_extension_settings(self, should_enable: bool = True) -> ExtensionSettings:
         """Create and save extension settings."""
         new_extension = ExtensionSettings(
             identifier=self.identifier,
             name=self.source.name,
             docker=self.source.docker,
             tag=self.tag,
             permissions=self.source.permissions,
-            enabled=True,
+            enabled=should_enable,
             user_permissions=self.source.user_permissions,
         )
-        # Save in settings first, if the image fails to install it will try to fetch after in main kraken check loop
         self._save_settings(new_extension)
         return new_extension
 
@@ -216,54 +215,50 @@ async def _pull_docker_image(self, docker_auth: Optional[str]) -> AsyncGenerator
             if self.digest:
                 await client.images.tag(tag, f"{self.source.docker}:{self.tag}")
 
+    @staticmethod
+    def _status_message(message: str) -> bytes:
+        return json.dumps({"status": message}).encode("utf-8")
+
     async def _clear_remaining_tags(self) -> None:
         """Uninstall all other tags for this extension."""
         logger.info(f"Clearing remaining tags for extension {self.identifier}")
         to_clear: List[Extension] = cast(List[Extension], await self.from_settings(self.identifier))
         to_clear = [version for version in to_clear if version.source.tag != self.tag]
         await asyncio.gather(*(version.uninstall() for version in to_clear))
 
-    async def install(self, clear_remaining_tags: bool = True, atomic: bool = False) -> AsyncGenerator[bytes, None]:
+    async def install(
+        self, clear_remaining_tags: bool = True, should_enable: bool = True
+    ) -> AsyncGenerator[bytes, None]:
         logger.info(f"Installing extension {self.identifier}:{self.tag}")
 
-        # First we should make sure no other tag is running
-        running_ext = await self._disable_running_extension()
-
-        self._create_extension_settings()
         try:
             self.lock(self.unique_entry)
 
             docker_auth = self._prepare_docker_auth()
             async for line in self._pull_docker_image(docker_auth):
                 yield line
         except Exception as error:
-            # In case of some external installs kraken shouldn't try to install it again so we remove from settings
-            if atomic:
-                should_raise = False
-                if await self._image_is_available_locally():
-                    logger.info(f"Pull failed but image {self.identifier}:{self.tag} is already available locally")
-                else:
-                    if not running_ext or self.unique_entry != running_ext.unique_entry:
-                        should_raise = True
-                        await self.uninstall()
-                    if running_ext:
-                        await running_ext.enable()
-
-                if should_raise:
-                    raise ExtensionPullFailed(f"Failed to pull extension {self.identifier}:{self.tag}") from error
-                # Reached only if the extensions are the same, the change is in permissions, not installation failure.
-                return
+            if await self._image_is_available_locally():
+                logger.info(f"Pull failed but image {self.identifier}:{self.tag} is already available locally")
+            else:
+                raise ExtensionPullFailed(f"Failed to pull extension {self.identifier}:{self.tag}") from error
         finally:
             self.unlock(self.unique_entry)
             self.reset_start_attempt(self.unique_entry)
 
-        logger.info(f"Extension {self.identifier}:{self.tag} installed")
+        await self._disable_running_extension()
+        self._create_extension_settings(should_enable)
+        yield self._status_message(f"Extension {self.identifier}:{self.tag} registered")
+
         # Uninstall all other tags in case user wants to clear them
         if clear_remaining_tags:
             await self._clear_remaining_tags()
+            yield self._status_message(f"Previous versions of {self.identifier} cleared")
+
+        logger.info(f"Extension {self.identifier}:{self.tag} installed")
 
-    async def update(self, clear_remaining_tags: bool) -> AsyncGenerator[bytes, None]:
-        async for data in self.install(clear_remaining_tags):
+    async def update(self, clear_remaining_tags: bool, should_enable: bool = True) -> AsyncGenerator[bytes, None]:
+        async for data in self.install(clear_remaining_tags, should_enable=should_enable):
             yield data
 
     async def uninstall(self) -> None: