CDRM-Project/custom_functions/database/unified_db_ops.py

"""Unified database operations module that automatically uses the correct backend."""

import os
from typing import Optional, List, Dict, Any
import yaml

# Import both backend modules
try:
    import custom_functions.database.cache_to_db_sqlite as sqlite_db
except ImportError:
    sqlite_db = None

try:
    import custom_functions.database.cache_to_db_mariadb as mariadb_db
except ImportError:
    mariadb_db = None


class DatabaseOperations:
    """Unified database operations class that automatically selects the correct backend."""

    def __init__(self):
        self.backend = self._get_database_backend()
        self.db_module = self._get_db_module()

    def _get_database_backend(self) -> str:
        """Get the database backend from config, default to sqlite."""
        try:
            config_path = os.path.join(os.getcwd(), "configs", "config.yaml")
            with open(config_path, "r", encoding="utf-8") as file:
                config = yaml.safe_load(file)
            return config.get("database_type", "sqlite").lower()
        except (FileNotFoundError, KeyError, yaml.YAMLError):
            return "sqlite"

    def _get_db_module(self):
        """Get the appropriate database module based on backend."""
        if self.backend == "mariadb" and mariadb_db:
            return mariadb_db
        if sqlite_db:
            return sqlite_db
        raise ImportError(f"Database module for {self.backend} not available")

    def get_backend_info(self) -> Dict[str, str]:
        """Get information about the current database backend being used."""
        return {
            "backend": self.backend,
            "module": self.db_module.__name__ if self.db_module else "None",
        }

    def create_database(self) -> None:
        """Create the database using the configured backend."""
        return self.db_module.create_database()

    def cache_to_db(
        self,
        service: str = "",
        pssh: str = "",
        kid: str = "",
        key: str = "",
        license_url: str = "",
        headers: str = "",
        cookies: str = "",
        data: str = "",
    ) -> bool:
        """Cache data to the database using the configured backend."""
        return self.db_module.cache_to_db(
            service=service,
            pssh=pssh,
            kid=kid,
            key=key,
            license_url=license_url,
            headers=headers,
            cookies=cookies,
            data=data,
        )

    def search_by_pssh_or_kid(self, search_filter: str) -> List[Dict[str, str]]:
        """Search the database by PSSH or KID using the configured backend."""
        return self.db_module.search_by_pssh_or_kid(search_filter)

    def get_key_by_kid_and_service(self, kid: str, service: str) -> Optional[str]:
        """Get the key by KID and service using the configured backend."""
        return self.db_module.get_key_by_kid_and_service(kid, service)

    def get_kid_key_dict(self, service_name: str) -> Dict[str, str]:
        """Get the KID and key dictionary using the configured backend."""
        return self.db_module.get_kid_key_dict(service_name)

    def get_unique_services(self) -> List[str]:
        """Get the unique services using the configured backend."""
        return self.db_module.get_unique_services()

    def key_count(self) -> int:
        """Get the key count using the configured backend."""
        return self.db_module.key_count()


# Create a singleton instance for easy import and use
db_ops = DatabaseOperations()


# Convenience functions that use the singleton instance
def get_backend_info() -> Dict[str, str]:
    """Get information about the current database backend being used."""
    return db_ops.get_backend_info()


def create_database() -> None:
    """Create the database using the configured backend."""
    return db_ops.create_database()


def cache_to_db(
    service: str = "",
    pssh: str = "",
    kid: str = "",
    key: str = "",
    license_url: str = "",
    headers: str = "",
    cookies: str = "",
    data: str = "",
) -> bool:
    """Cache data to the database using the configured backend."""
    return db_ops.cache_to_db(
        service=service,
        pssh=pssh,
        kid=kid,
        key=key,
        license_url=license_url,
        headers=headers,
        cookies=cookies,
        data=data,
    )


def search_by_pssh_or_kid(search_filter: str) -> List[Dict[str, str]]:
    """Search the database by PSSH or KID using the configured backend."""
    return db_ops.search_by_pssh_or_kid(search_filter)


def get_key_by_kid_and_service(kid: str, service: str) -> Optional[str]:
    """Get the key by KID and service using the configured backend."""
    return db_ops.get_key_by_kid_and_service(kid, service)


def get_kid_key_dict(service_name: str) -> Dict[str, str]:
    """Get the KID and key dictionary using the configured backend."""
    return db_ops.get_kid_key_dict(service_name)


def get_unique_services() -> List[str]:
    """Get the unique services using the configured backend."""
    return db_ops.get_unique_services()


def key_count() -> int:
    """Get the key count using the configured backend."""
    return db_ops.key_count()