Source code for q2_sdk.core.install_steps.vault_info

import base64
import sys
import inspect
import importlib
from typing import List, Optional, Unpack

from q2_sdk.core import vault
from q2_sdk.core.install_steps.base import (
    InstallStep,
    InstallStepAttribute,
    InstallStepArguments,
)


class BasenameTransform:
    def run(self, raw_value):
        raise NotImplementedError

    def __json__(self):
        return {
            "name": self.__class__.__name__,
            "location": inspect.getmodule(self.__class__).__name__,
            "parameters": vars(self),
        }


class LowercaseTransform(BasenameTransform):
    def run(self, raw_value):
        return raw_value.lower()


class UppercaseTransform(BasenameTransform):
    def run(self, raw_value):
        return raw_value.upper()


class EncodeAsBase64Transform(BasenameTransform):
    def run(self, raw_value):
        return base64.b64encode(raw_value.encode()).decode()


class RightStripTransform(BasenameTransform):
    def __init__(self, suffix_to_strip):
        self.suffix_to_strip = suffix_to_strip

    def run(self, raw_value):
        return raw_value.split(self.suffix_to_strip)[0]


class LeftStripTransform(BasenameTransform):
    def __init__(self, prefix_to_strip):
        self.prefix_to_strip = prefix_to_strip

    def run(self, raw_value):
        return raw_value.split(self.prefix_to_strip)[-1]




class VaultInfo(InstallStep):
    """Adds json blob to vault"""

    def __init__(
        self,
        key_prefix,
        key_basename,
        value: dict,
        basename_transforms: Optional[List[BasenameTransform]] = None,
        use_default_vault_prefix=False,
        **kwargs: Unpack[InstallStepArguments],
    ):
        super().__init__(**kwargs)
        assert isinstance(value, dict)
        if not basename_transforms:
            basename_transforms = []
        self.q2_vault_obj = None
        self.key_prefix = InstallStepAttribute(key_prefix)
        self.key_basename = InstallStepAttribute(key_basename, is_editable=True)
        self.value = InstallStepAttribute(value, is_editable=True)
        self.basename_transforms = basename_transforms
        self.use_default_vault_prefix = use_default_vault_prefix
        self.install_order = 5
        self.not_install_var_list = ["basename_transforms", "use_default_vault_prefix"]

    def __json__(self):
        return {
            **self.install_vars,
            **{"basename_transforms": self.basename_transforms},
        }



    async def install(self):
        await super().install()
        self.q2_vault_obj = vault.get_client()
        if not self.q2_vault_obj:
            self.logger.error("Cannot write to Vault without a vault token")
            sys.exit(1)

        basename = self.key_basename.value
        for transform in self.basename_transforms:
            if not isinstance(transform, BasenameTransform):
                module = importlib.import_module(transform["location"])
                class_ = getattr(module, transform["name"])
                instance = class_(**transform.get("parameters"))

                basename = instance.run(basename)
            else:
                basename = transform.run(basename)

        key_location = f"{self.key_prefix.value}/{basename}"
        prefix = ""
        if self.use_default_vault_prefix:
            prefix = vault.DEFAULT_PREFIX
        self.q2_vault_obj.write(key_location, self.value.value, prefix=prefix)