Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
70 changes: 45 additions & 25 deletions core/services/kraken/extension/extension.py
Original file line number Diff line number Diff line change
Expand Up @@ -145,18 +145,18 @@ async def remove(cls, container_name: str, delete_image: bool = True) -> None:
finally:
cls.unlock(container_name)

async def install(self, clear_remaining_tags: bool = True, atomic: bool = False) -> AsyncGenerator[bytes, None]:
logger.info(f"Installing extension {self.identifier}:{self.tag}")

# First we should make sure no other tag is running
running_ext = None
async def _disable_running_extension(self) -> Optional["Extension"]:
"""Disable any currently running extension with the same identifier."""
try:
running_ext = await self.from_running(self.identifier)
if running_ext:
await running_ext.disable()
return running_ext
except ExtensionNotRunning:
pass
return None

def _create_extension_settings(self) -> ExtensionSettings:
"""Create and save extension settings."""
new_extension = ExtensionSettings(
identifier=self.identifier,
name=self.source.name,
Expand All @@ -168,25 +168,48 @@ async def install(self, clear_remaining_tags: bool = True, atomic: bool = False)
)
# 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

def _prepare_docker_auth(self) -> Optional[str]:
"""Prepare Docker authentication string from source auth credentials."""
if self.source.auth is None:
return None
docker_auth = f"{self.source.auth.username}:{self.source.auth.password}"
return base64.b64encode(docker_auth.encode("utf-8")).decode("utf-8")
Comment on lines +173 to +178
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

suggestion: Consider handling empty username or password in Docker auth preparation.

Validate that both username and password are non-empty before encoding, and raise an error if either is missing to prevent invalid Docker auth strings.

Suggested change
def _prepare_docker_auth(self) -> Optional[str]:
"""Prepare Docker authentication string from source auth credentials."""
if self.source.auth is None:
return None
docker_auth = f"{self.source.auth.username}:{self.source.auth.password}"
return base64.b64encode(docker_auth.encode("utf-8")).decode("utf-8")
def _prepare_docker_auth(self) -> Optional[str]:
"""Prepare Docker authentication string from source auth credentials."""
if self.source.auth is None:
return None
username = getattr(self.source.auth, "username", None)
password = getattr(self.source.auth, "password", None)
if not username or not password:
raise ValueError("Docker authentication requires both non-empty username and password.")
docker_auth = f"{username}:{password}"
return base64.b64encode(docker_auth.encode("utf-8")).decode("utf-8")


async def _pull_docker_image(self, docker_auth: Optional[str]) -> AsyncGenerator[bytes, None]:
"""Pull Docker image and yield progress updates."""
tag = f"{self.source.docker}:{self.tag}" + (f"@{self.digest}" if self.digest else "")
async with DockerCtx() as client:
async for line in client.images.pull(
tag, repo=self.source.docker, tag=self.tag, auth=docker_auth, stream=True
):
# TODO - Plug Error detection from docker image here
yield json.dumps(line).encode("utf-8")
# Make sure to add correct tag if a digest was used since docker messes up the tag
if self.digest:
await client.images.tag(tag, f"{self.source.docker}:{self.tag}")

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]:
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()
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

suggestion (bug_risk): Extension settings are created before acquiring the lock.

Moving the creation and saving of extension settings inside the locked section will prevent race conditions during concurrent installs.

try:
self.lock(self.unique_entry)

docker_auth: Optional[str] = None
if self.source.auth is not None:
docker_auth = f"{self.source.auth.username}:{self.source.auth.password}"
docker_auth = base64.b64encode(docker_auth.encode("utf-8")).decode("utf-8")

tag = f"{self.source.docker}:{self.tag}" + (f"@{self.digest}" if self.digest else "")
async with DockerCtx() as client:
async for line in client.images.pull(
tag, repo=self.source.docker, tag=self.tag, auth=docker_auth, stream=True
):
# TODO - Plug Error detection from docker image here
yield json.dumps(line).encode("utf-8")
# Make sure to add correct tag if a digest was used since docker messes up the tag
if self.digest:
await client.images.tag(tag, f"{self.source.docker}:{self.tag}")
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:
Expand All @@ -209,10 +232,7 @@ async def install(self, clear_remaining_tags: bool = True, atomic: bool = False)
logger.info(f"Extension {self.identifier}:{self.tag} installed")
# Uninstall all other tags in case user wants to clear them
if clear_remaining_tags:
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))
await self._clear_remaining_tags()

async def update(self, clear_remaining_tags: bool) -> AsyncGenerator[bytes, None]:
async for data in self.install(clear_remaining_tags):
Expand Down