database.py

from huggingface_hub import hf_hub_download, upload_file, space_info
from huggingface_hub.utils import HfHubHTTPError
from pathlib import Path
from typing import Set

DATABASE_REPO_ID = "spaces-ci-bot/webhook"
DATABASE_FILE = "registered_spaces.txt"

MAX_NB_ATTEMPTS = 10


def is_space_existing(space_id: str) -> bool:
    try:
        space_info(repo_id=space_id)
        return True
    except HfHubHTTPError:
        return False


def is_space_registered(space_id: str) -> bool:
    return space_id in get_registered_spaces()


def get_registered_spaces() -> Set[str]:
    return _read_database(_get_latest_file())


def update_status(space_id: str, should_watch: bool) -> None:
    nb_attempts = 0
    while True:
        # Get registered spaces
        filepath = _get_latest_file()
        registered_spaces = _read_database(filepath)

        # Do nothing if:
        # - need to register and already registered
        # - need to unregister and already not registered
        if (should_watch and space_id in registered_spaces) or (
            not should_watch and space_id not in registered_spaces
        ):
            return

        # Else, (un)register new space
        latest_revision = filepath.parent.name
        if should_watch:
            registered_spaces.add(space_id)
        else:
            registered_spaces.remove(space_id)

        # Re-upload database and ensure no concurrent call happened
        try:
            nb_attempts += 1
            upload_file(
                path_in_repo=DATABASE_FILE,
                repo_id=DATABASE_REPO_ID,
                repo_type="dataset",
                path_or_fileobj="\n".join(sorted(registered_spaces)).encode(),
                commit_message=(
                    f"Register {space_id}" if should_watch else f"Unregister {space_id}"
                ),
                parent_commit=latest_revision,  # ensure consistency
            )
            return
        except HfHubHTTPError:
            # Retry X times before giving up (in case multiple registrations at the same time)
            if nb_attempts == MAX_NB_ATTEMPTS:
                raise


def _read_database(filepath: Path) -> Set[str]:
    return set(filepath.read_text().split())


def _get_latest_file() -> Path:
    return Path(hf_hub_download(DATABASE_REPO_ID, DATABASE_FILE, repo_type="dataset"))