数据库¶

`PasswordVault` ¶

Password vault manager that provides access to password vault implementations.

Caches vault instances for performance.

Source code in aloha/db/base.py

class PasswordVault:
    """
    Password vault manager that provides access to password vault implementations.

    Caches vault instances for performance.
    """

    _dict_cache_vault = {}

    @staticmethod
    def get_vault(vault_type: str | None = None, vault_config: dict | None = None, **kwargs) -> vault.BaseVault:
        """
        Get a password vault instance.

        Supports multiple vault types:
        - 'plain' or 'aes': AES-based vault (default fallback)
        - 'cyberark': CyberArk vault
        - Other/None: Dummy vault (plain text)

        :param vault_type: Type of vault to use (overrides config)
        :param vault_config: Vault configuration dictionary
        :param args: Additional arguments
        :param kwargs: Additional keyword arguments
        :return: Vault instance implementing BaseVault interface
        :raises RuntimeError: If CyberArk vault is requested but config is missing
        """
        encryption_method = vault_type or SETTINGS.config.get("PASSWORD_ENCRYPTION")
        LOG.debug("Using password vault: %s", encryption_method)  # nosemgrep

        cache_key = "%s:%s" % (encryption_method, str(vault_config))
        if cache_key not in PasswordVault._dict_cache_vault:
            if encryption_method in ("plain", "aes") or encryption_method is True:
                v = vault.AesVault(**(vault_config or {}))
            elif encryption_method == "cyberark":
                config_cyberark = vault_config or SETTINGS.config.get("CYBERARK_CONFIG")
                if config_cyberark is None:
                    raise RuntimeError("Missing [CYBERARK_CONFIG] in config!")
                v = vault.CyberArkVault(**config_cyberark)
            else:
                msg = "Using plain password vault as unknown value of PASSWORD_ENCRYPTION=%s in config." % encryption_method
                LOG.info(msg)  # nosemgrep
                v = vault.DummyVault(**(vault_config or {}))
            PasswordVault._dict_cache_vault[cache_key] = v

        return PasswordVault._dict_cache_vault[cache_key]

`get_vault(vault_type=None, vault_config=None, **kwargs)` `staticmethod` ¶

Get a password vault instance.

Supports multiple vault types: - 'plain' or 'aes': AES-based vault (default fallback) - 'cyberark': CyberArk vault - Other/None: Dummy vault (plain text)

Parameters:

Name	Type	Description	Default
`vault_type`	`str \| None`	Type of vault to use (overrides config)	`None`
`vault_config`	`dict \| None`	Vault configuration dictionary	`None`
`args`		Additional arguments	required
`kwargs`		Additional keyword arguments	`{}`

Returns:

Type	Description
`BaseVault`	Vault instance implementing BaseVault interface

Raises:

Type	Description
`RuntimeError`	If CyberArk vault is requested but config is missing

Source code in aloha/db/base.py

@staticmethod
def get_vault(vault_type: str | None = None, vault_config: dict | None = None, **kwargs) -> vault.BaseVault:
    """
    Get a password vault instance.

    Supports multiple vault types:
    - 'plain' or 'aes': AES-based vault (default fallback)
    - 'cyberark': CyberArk vault
    - Other/None: Dummy vault (plain text)

    :param vault_type: Type of vault to use (overrides config)
    :param vault_config: Vault configuration dictionary
    :param args: Additional arguments
    :param kwargs: Additional keyword arguments
    :return: Vault instance implementing BaseVault interface
    :raises RuntimeError: If CyberArk vault is requested but config is missing
    """
    encryption_method = vault_type or SETTINGS.config.get("PASSWORD_ENCRYPTION")
    LOG.debug("Using password vault: %s", encryption_method)  # nosemgrep

    cache_key = "%s:%s" % (encryption_method, str(vault_config))
    if cache_key not in PasswordVault._dict_cache_vault:
        if encryption_method in ("plain", "aes") or encryption_method is True:
            v = vault.AesVault(**(vault_config or {}))
        elif encryption_method == "cyberark":
            config_cyberark = vault_config or SETTINGS.config.get("CYBERARK_CONFIG")
            if config_cyberark is None:
                raise RuntimeError("Missing [CYBERARK_CONFIG] in config!")
            v = vault.CyberArkVault(**config_cyberark)
        else:
            msg = "Using plain password vault as unknown value of PASSWORD_ENCRYPTION=%s in config." % encryption_method
            LOG.info(msg)  # nosemgrep
            v = vault.DummyVault(**(vault_config or {}))
        PasswordVault._dict_cache_vault[cache_key] = v

    return PasswordVault._dict_cache_vault[cache_key]

`main()` ¶

Command-line tool to decrypt passwords from config.

Usage: python -m aloha.db.base

Source code in aloha/db/base.py

def main():
    """
    Command-line tool to decrypt passwords from config.

    Usage: python -m aloha.db.base <config_key>
    """
    import sys

    config_key = sys.argv[-1]
    LOG.debug("Getting pwd for deploy key [deploy.%s]" % config_key)
    try:
        db_config = SETTINGS.config["deploy"][config_key]
        password_vault = PasswordVault.get_vault()
        p = password_vault.get_password(db_config.get("password"))
        LOG.debug("Decrypted PWD: %s" % p)
    except KeyError:
        LOG.error("Please make sure config key [deploy.%s] exists!" % config_key)

DuckDB connection helpers.

`DuckOperator` ¶

Create and use a DuckDB connection through SQLAlchemy.

Source code in aloha/db/duckdb.py

class DuckOperator:
    """Create and use a DuckDB connection through SQLAlchemy."""

    def __init__(self, db_config, **kwargs):
        """Build a DuckDB engine, creating the database file if necessary."""
        """db_config example:
        {
            "path": "/path/to/db.duckdb",     # file path of duckdb, use ":memory:" for in-memory mode
            "schema": "sales",                # optional, 'main' by default
            "read_only": True,                # optional, False by default, (will set to False if in in-memory mode)
            "config": {"memory_limit": "500mb"}, # optional, duckdb connection configs
        }
        """
        self._config = {
            "path": db_config.get("path", ":memory:"),
            "schema": db_config.get("schema", "main"),
            "read_only": bool(db_config.get("read_only", False)),
            "config": db_config.get("config", {}),
            "auto_commit": db_config.get("auto_commit", True),
        }

        if not self._config["path"] or self._config["path"] == ":memory:":  # in-memory mode
            self._config["path"] = ":memory:"

            if self._config["read_only"]:  # in-memory mode cannot be read-only
                LOG.warning("In-memory database cannot be read-only. Setting read_only=False.")
                self._config["read_only"] = False

        else:
            self._prepare_database()

        try:
            str_connection = f"duckdb:///{self._config['path']}"
            self.engine = create_engine(
                str_connection,
                connect_args={"read_only": self._config["read_only"], "config": self._config["config"]},
                **kwargs,
            )

            self._initialize_schema()
            msg = f"DuckDB connected: {self._config['path']} [schema={self._config['schema']}, read_only={self._config['read_only']}]"
            LOG.debug(msg)
        except Exception as e:
            LOG.exception(e)
            raise RuntimeError("Failed to connect to DuckDB")

    def _prepare_database(self):
        """Prepare the database file and its parent directory."""
        path = self._config["path"]
        path_obj = Path(path)

        parent_dir = path_obj.parent
        if not parent_dir.exists():
            if self._config["read_only"]:
                raise RuntimeError(f"Directory '{parent_dir}' does not exist and read_only=True")
            try:
                parent_dir.mkdir(parents=True, exist_ok=True)
                LOG.debug(f"Created directory: {parent_dir}")
            except Exception as e:
                raise RuntimeError(f"Failed to create directory '{parent_dir}': {e}")

        if not path_obj.exists():
            if self._config["read_only"]:
                raise RuntimeError(f"DuckDB file '{path}' does not exist and read_only=True")
            try:
                LOG.debug(f"Database file not found, creating: {path}")
                duckdb.connect(path).close()
            except Exception as e:
                raise RuntimeError(f"Failed to create database file '{path}': {e}")

    def _initialize_schema(self):
        """Create or select the requested schema."""
        if self._config["schema"] == "main":
            return

        try:
            if self._config["read_only"]:
                result = self.engine.connext().execute(
                    text("SELECT schema_name FROM information_schema.schemata WHERE schema_name = :schema"),
                    {"schema": self._config["schema"]},
                )
                if not result.fetchone():
                    raise RuntimeError(f"Schema '{self._config['schema']}' does not exist and read_only=True")
            else:
                self.engine.connect().execute(text(f"CREATE SCHEMA IF NOT EXISTS {self._config['schema']}"))

            self.engine.connect().execute(text(f"SET schema '{self._config['schema']}'"))
        except Exception as e:
            raise RuntimeError(f"Failed to initialize schema: {e}")

    @property
    def connection(self):
        return self.engine

    conn = connection

    def execute_query(self, sql, *args, **kwargs):
        """Execute a SQL statement and return the cursor result."""
        with self.engine.connect() as conn:
            cur = conn.execute(text(sql), *args, *kwargs)
            if self._config.get("auto_commit", True):
                conn.commit()
            return cur

    @property
    def connection_str(self) -> str:
        """Return a human-readable connection string."""
        return f"duckdb:///{self._config['path']} [schema={self._config['schema']}, read_only={self._config['read_only']}]"

`connection_str` `property` ¶

Return a human-readable connection string.

`init(db_config, **kwargs)` ¶

Build a DuckDB engine, creating the database file if necessary.

Source code in aloha/db/duckdb.py

def __init__(self, db_config, **kwargs):
    """Build a DuckDB engine, creating the database file if necessary."""
    """db_config example:
    {
        "path": "/path/to/db.duckdb",     # file path of duckdb, use ":memory:" for in-memory mode
        "schema": "sales",                # optional, 'main' by default
        "read_only": True,                # optional, False by default, (will set to False if in in-memory mode)
        "config": {"memory_limit": "500mb"}, # optional, duckdb connection configs
    }
    """
    self._config = {
        "path": db_config.get("path", ":memory:"),
        "schema": db_config.get("schema", "main"),
        "read_only": bool(db_config.get("read_only", False)),
        "config": db_config.get("config", {}),
        "auto_commit": db_config.get("auto_commit", True),
    }

    if not self._config["path"] or self._config["path"] == ":memory:":  # in-memory mode
        self._config["path"] = ":memory:"

        if self._config["read_only"]:  # in-memory mode cannot be read-only
            LOG.warning("In-memory database cannot be read-only. Setting read_only=False.")
            self._config["read_only"] = False

    else:
        self._prepare_database()

    try:
        str_connection = f"duckdb:///{self._config['path']}"
        self.engine = create_engine(
            str_connection,
            connect_args={"read_only": self._config["read_only"], "config": self._config["config"]},
            **kwargs,
        )

        self._initialize_schema()
        msg = f"DuckDB connected: {self._config['path']} [schema={self._config['schema']}, read_only={self._config['read_only']}]"
        LOG.debug(msg)
    except Exception as e:
        LOG.exception(e)
        raise RuntimeError("Failed to connect to DuckDB")

`execute_query(sql, *args, **kwargs)` ¶

Execute a SQL statement and return the cursor result.

Source code in aloha/db/duckdb.py

def execute_query(self, sql, *args, **kwargs):
    """Execute a SQL statement and return the cursor result."""
    with self.engine.connect() as conn:
        cur = conn.execute(text(sql), *args, *kwargs)
        if self._config.get("auto_commit", True):
            conn.commit()
        return cur