kaptive.db.manager¶

@classmethod
def add(
    cls,
    owner: str,
    repo_name: str,
    db_name: str,
    branch: str = "main",
    local_meta: DatabaseMetadata | None = None,
) -> Database | None:
    r"""Add or update a database directly from a specified remote Git repository.

    This is the primary method for adding custom or official databases from GitHub. The procedure:

    1. Constructs raw GitHub URL endpoints for the repository's `.toml` metadata and `.gbk` GenBank files.
    2. Downloads and parses the remote TOML metadata to extract version information.
    3. Compares the remote version against local metadata (if available). If up-to-date, skips remaining steps and
       returns `None`.
    4. Downloads the raw GenBank file content over HTTP.
    5. Writes source files into a temporary directory and compiles them using
       [`from_genbank`][kaptive.db.core.Database.from_genbank].
    6. Serializes and caches the compiled [`Database`][kaptive.db.core.Database] object into the local storage
       directory.

    Args:
        owner (str): Owner or organization of the GitHub repository (e.g., `'klebgenomics'`).
        repo_name (str): Name of the GitHub repository (e.g., `'KpSC_surface_antigen_loci'`).
        db_name (str): Base name of the database files in the repository
            (e.g., `'Klebsiella_pneumoniae_Species_Complex_K'`).
        branch (str): Git branch name to fetch from. Defaults to `'main'`.
        local_meta (DatabaseMetadata | None): Pre-loaded metadata of local database installation.
            Defaults to `None`.

    Returns:
        Database | None: The newly compiled [`Database`][kaptive.db.core.Database] object if installed or updated,
            or `None` if the local version was already up-to-date.

    Raises:
        DatabaseError: If repository files are not found, network issues occur, or file compilation fails.

    See Also:
        [`DatabaseManager`][kaptive.db.manager.DatabaseManager],
        [`Database`][kaptive.db.core.Database],
        [`DatabaseMetadata`][kaptive.db.models.DatabaseMetadata],
        [`DatabaseError`][kaptive.db.models.DatabaseError]
    """
    fetched = cls._fetch_files(owner, repo_name, db_name, branch=branch, local_meta=local_meta)
    if fetched is None:
        return None
    return cls._compile_and_save(*fetched)

@classmethod
def get(cls, file_or_keyword: str | Path) -> Database:
    r"""Load a Database from a file path or resolve and load it by keyword.

    If `file_or_keyword` points to an existing file, it is loaded directly.
    Otherwise, it is treated as a keyword. If the keyword is not installed locally,
    it will be automatically downloaded and installed.

    Args:
        file_or_keyword (str | Path): File path or recognized database keyword.

    Returns:
        Database: The loaded [`Database`][kaptive.db.core.Database] instance.
    """
    from kaptive.db.core import Database

    try:
        file_path = Path(file_or_keyword)
        if file_path.is_file():
            return Database.load(file_path)
    except (TypeError, ValueError, OSError):
        pass

    try:
        return cls.load(str(file_or_keyword))
    except DatabaseError:
        result = cls.install(str(file_or_keyword))
        if isinstance(result, list):
            result = result[0]
        if result is None:
            return cls.load(str(file_or_keyword))
        return result

@classmethod
def install(cls, kwd: str | list[str]) -> Database | list[Database | None]:
    r"""Install known, officially supported databases by keyword.

    Looks up the repository details (owner, repo, database name) associated with the provided keyword(s)
    in the internal registry (`_KNOWN`) and delegates file retrieval
    and compilation to [`add`][kaptive.db.manager.DatabaseManager.add]. If `'all'` or a list of keywords is
    supplied, fetching is performed concurrently via a thread pool executor.

    Args:
        kwd (str | list[str]): The keyword(s) of the known database(s) to install (e.g., `'kpsc_k'`,
            `['kpsc_k', 'ab_k']`, or `'all'`).

    Returns:
        Database | list[Database | None]: For a single keyword, returns the compiled
            [`Database`][kaptive.db.core.Database] object (or `None` if already up-to-date). For a list of
            keywords or `'all'`, returns a list of compiled [`Database`][kaptive.db.core.Database] objects
            (or `None` for entries that were up-to-date).

    Raises:
        DatabaseError: If any keyword is not recognized in the list of known databases, or if network/parsing
            errors occur.

    See Also:
        [`known`][kaptive.db.manager.DatabaseManager.known],
        [`DatabaseManager`][kaptive.db.manager.DatabaseManager],
        [`add`][kaptive.db.manager.DatabaseManager.add],
        [`Database`][kaptive.db.core.Database],
        [`DatabaseError`][kaptive.db.models.DatabaseError]
    """
    if kwd == "all":
        kwd = list(cls._KNOWN.keys())

    if isinstance(kwd, list):

        def _fetch_one(k: str):
            if (known_info := cls._KNOWN.get(k, None)) is None:
                raise DatabaseError(f'"{k}" is not a known database, choose from {list(cls._KNOWN.keys())}')
            return cls._fetch_files(*known_info)

        with concurrent.futures.ThreadPoolExecutor() as executor:
            fetched_list = list(executor.map(_fetch_one, kwd))

        results = []
        for fetched in fetched_list:
            if fetched is None:
                results.append(None)
            else:
                results.append(cls._compile_and_save(*fetched))
        return results

    if (known_info := cls._KNOWN.get(kwd, None)) is None:
        raise DatabaseError(f'"{kwd}" is not a known database, choose from {list(cls._KNOWN.keys())}')
    return cls.add(*known_info)  # type: ignore

@classmethod
def installed(cls) -> list[str]:
    r"""Return a list of keywords for all currently installed databases.

    Scans the local storage directory for `.pkl` files and extracts their keywords from the file stems.

    Returns:
        list[str]: A list of database keywords corresponding to installed `.pkl` database files.
            Returns an empty list if no databases are installed or if the storage directory does not exist.

    See Also:
        [`DatabaseManager`][kaptive.db.manager.DatabaseManager],
        [`known`][kaptive.db.manager.DatabaseManager.known]
    """
    if not cls._DB_DIR.exists():
        return []
    return [p.stem for p in cls._DB_DIR.glob("*.pkl")]

@classmethod
def known(cls) -> list[str]:
    r"""Return a list of keywords for all currently known, officially supported databases.

    These databases can be installed directly by providing their keyword to
    [`install`][kaptive.db.manager.DatabaseManager.install].

    Returns:
        list[str]: A list of known database keywords (e.g., `['kpsc_k', 'kpsc_o', 'kosc_k', ...]`).

    See Also:
        [`DatabaseManager`][kaptive.db.manager.DatabaseManager],
        [`install`][kaptive.db.manager.DatabaseManager.install]
    """
    return list(cls._KNOWN.keys())

@classmethod
def load(cls, kwd: str) -> Database:
    r"""Load a locally installed, compiled database using its keyword.

    Reads and unpickles the serialized `.pkl` database file from the local cache directory.

    Args:
        kwd (str): The keyword identifier of the database to load (e.g., `'kpsc_k'`).

    Returns:
        Database: The deserialized [`Database`][kaptive.db.core.Database] instance.

    Raises:
        DatabaseError: If the specified database is not installed locally.

    See Also:
        [`DatabaseManager`][kaptive.db.manager.DatabaseManager],
        [`Database`][kaptive.db.core.Database],
        [`DatabaseError`][kaptive.db.models.DatabaseError]
    """
    return pickle.loads(cls._get_existing_db_path(kwd).read_bytes())

@classmethod
def reset(cls) -> None:
    r"""Remove all installed databases by deleting their compiled files from the local directory.

    This clears the user's `~/.kaptive` cache directory of any `.pkl` database files and `.json` metadata
    sidecar files, effectively uninstalling all downloaded and compiled databases.

    Returns:
        None

    See Also:
        [`DatabaseManager`][kaptive.db.manager.DatabaseManager],
        [`uninstall`][kaptive.db.manager.DatabaseManager.uninstall]
    """
    if cls._DB_DIR.exists():
        for file_path in cls._DB_DIR.glob("*.pkl"):
            file_path.unlink()
        for file_path in cls._DB_DIR.glob("*.json"):
            file_path.unlink()

@classmethod
def save(cls, db: Database) -> int:
    r"""Serialize and save a compiled Database object and its metadata to local storage.

    Saves the database as a `.pkl` file named `{keyword}.pkl` in the local cache directory (`_DB_DIR`).
    Also writes a companion `{keyword}.json` file containing the serialized metadata for fast version checking.

    Args:
        db (Database): The compiled [`Database`][kaptive.db.core.Database] object to save.

    Returns:
        int: The total number of bytes written to the `.pkl` database file.

    See Also:
        [`DatabaseManager`][kaptive.db.manager.DatabaseManager],
        [`Database`][kaptive.db.core.Database],
        [`DatabaseMetadata`][kaptive.db.models.DatabaseMetadata]
    """
    db_path = cls._get_db_path(db.metadata.keyword)
    db_path.with_suffix(".json").write_text(json.dumps(asdict(db.metadata)))
    return db_path.write_bytes(pickle.dumps(db, protocol=pickle.HIGHEST_PROTOCOL))

@classmethod
def uninstall(cls, kwd: str) -> None:
    r"""Uninstall a specific database by removing its compiled local `.pkl` and `.json` files.

    Args:
        kwd (str): The keyword of the database to uninstall (e.g., `'kpsc_k'`).

    Returns:
        None

    Raises:
        DatabaseError: If the specified database is not currently installed locally.

    See Also:
        [`DatabaseManager`][kaptive.db.manager.DatabaseManager],
        [`DatabaseError`][kaptive.db.models.DatabaseError]
    """
    db_path = cls._get_existing_db_path(kwd)
    db_path.unlink()
    if db_path.with_suffix(".json").exists():
        db_path.with_suffix(".json").unlink()

@classmethod
def update(cls, kwd: str | list[str] = "all") -> Generator[Database, None, None]:
    r"""Update installed databases by checking against their remote GitHub repositories.

    Extracts local metadata to determine the source repository and version, then checks the remote GitHub
    repository for a newer version. If a newer version is available, the source files (`.gbk` and `.toml`)
    are fetched, compiled into a new [`Database`][kaptive.db.core.Database] object, and saved to disk. When
    updating multiple databases or `"all"`, remote fetches are executed concurrently using a thread pool executor.

    Args:
        kwd (str | list[str]): The keyword(s) of the database to update (e.g., `'kpsc_k'`,
            `['kpsc_k', 'ab_k']`, or `'all'`). Defaults to `"all"`, which updates all currently installed databases.

    Yields:
        Database: The newly compiled [`Database`][kaptive.db.core.Database] object for each database that
            required an update. Databases that are already up-to-date yield nothing.

    Raises:
        DatabaseError: If a requested database is not installed locally, or if network/parsing failures occur
            during update.

    See Also:
        [`installed`][kaptive.db.manager.DatabaseManager.installed],
        [`add`][kaptive.db.manager.DatabaseManager.add],
        [`DatabaseManager`][kaptive.db.manager.DatabaseManager],
        [`Database`][kaptive.db.core.Database],
        [`DatabaseMetadata`][kaptive.db.models.DatabaseMetadata],
        [`DatabaseError`][kaptive.db.models.DatabaseError]
    """
    if kwd == "all":
        kwd = cls.installed()
        if not kwd:
            return

    if isinstance(kwd, list):

        def _fetch_update_one(k: str):
            db_path = cls._get_existing_db_path(k)
            json_path = db_path.with_suffix(".json")
            if json_path.is_file():
                meta = DatabaseMetadata.from_dict(json.loads(json_path.read_text()))
            else:
                meta = pickle.loads(db_path.read_bytes()).metadata
            db_name = Path(meta.genbank).with_suffix("").name
            return cls._fetch_files(meta.owner, meta.repo, db_name, branch=meta.branch, local_meta=meta)

        with concurrent.futures.ThreadPoolExecutor() as executor:
            fetched_list = list(executor.map(_fetch_update_one, kwd))

        for fetched in fetched_list:
            if fetched is not None:
                yield cls._compile_and_save(*fetched)
    else:
        db_path = cls._get_existing_db_path(kwd)
        json_path = db_path.with_suffix(".json")
        if json_path.is_file():
            meta = DatabaseMetadata.from_dict(json.loads(json_path.read_text()))
        else:
            meta = pickle.loads(db_path.read_bytes()).metadata
        db_name = Path(meta.genbank).with_suffix("").name
        if (res := cls.add(meta.owner, meta.repo, db_name, branch=meta.branch, local_meta=meta)) is not None:
            yield res