"""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()