148 lines
4.5 KiB
Python
148 lines
4.5 KiB
Python
"""Database registry with singleton pattern for proper model management."""
|
|
|
|
from typing import Dict, Optional, Type, Any
|
|
from sqlalchemy import MetaData, inspect
|
|
from sqlalchemy.ext.declarative import declarative_base as _declarative_base
|
|
from sqlalchemy.orm import DeclarativeMeta
|
|
import threading
|
|
|
|
|
|
class DatabaseRegistry:
|
|
"""
|
|
Singleton registry for database models and metadata.
|
|
|
|
This ensures that:
|
|
1. Base is only created once
|
|
2. Models are registered only once
|
|
3. Tables can be safely re-imported without errors
|
|
4. Proper cleanup and reset for testing
|
|
"""
|
|
|
|
_instance: Optional['DatabaseRegistry'] = None
|
|
_lock = threading.Lock()
|
|
|
|
def __new__(cls) -> 'DatabaseRegistry':
|
|
if cls._instance is None:
|
|
with cls._lock:
|
|
if cls._instance is None:
|
|
cls._instance = super().__new__(cls)
|
|
cls._instance._initialized = False
|
|
return cls._instance
|
|
|
|
def __init__(self):
|
|
"""Initialize the registry only once."""
|
|
if self._initialized:
|
|
return
|
|
|
|
self._initialized = True
|
|
self._base: Optional[DeclarativeMeta] = None
|
|
self._metadata: Optional[MetaData] = None
|
|
self._models: Dict[str, Type[Any]] = {}
|
|
self._tables_created = False
|
|
|
|
@property
|
|
def Base(self) -> DeclarativeMeta:
|
|
"""Get or create the declarative base."""
|
|
if self._base is None:
|
|
self._metadata = MetaData()
|
|
self._base = _declarative_base(metadata=self._metadata)
|
|
return self._base
|
|
|
|
@property
|
|
def metadata(self) -> MetaData:
|
|
"""Get the metadata instance."""
|
|
if self._metadata is None:
|
|
_ = self.Base # Ensure Base is created
|
|
return self._metadata
|
|
|
|
def register_model(self, model_class: Type[Any]) -> Type[Any]:
|
|
"""
|
|
Register a model class with the registry.
|
|
|
|
This prevents duplicate registration and handles re-imports safely.
|
|
|
|
Args:
|
|
model_class: The SQLAlchemy model class to register
|
|
|
|
Returns:
|
|
The registered model class (may be the existing one if already registered)
|
|
"""
|
|
table_name = model_class.__tablename__
|
|
|
|
# If model already registered, return the existing one
|
|
if table_name in self._models:
|
|
existing_model = self._models[table_name]
|
|
# Update the class reference to the existing model
|
|
return existing_model
|
|
|
|
# Register new model
|
|
self._models[table_name] = model_class
|
|
return model_class
|
|
|
|
def get_model(self, table_name: str) -> Optional[Type[Any]]:
|
|
"""Get a registered model by table name."""
|
|
return self._models.get(table_name)
|
|
|
|
def create_all_tables(self, engine):
|
|
"""
|
|
Create all tables in the database.
|
|
|
|
Handles existing tables and indexes gracefully with checkfirst=True.
|
|
"""
|
|
# Create all tables with checkfirst=True to skip existing ones
|
|
# This also handles indexes properly
|
|
self.metadata.create_all(bind=engine, checkfirst=True)
|
|
self._tables_created = True
|
|
|
|
def drop_all_tables(self, engine):
|
|
"""Drop all tables from the database."""
|
|
self.metadata.drop_all(bind=engine)
|
|
self._tables_created = False
|
|
|
|
def clear_models(self):
|
|
"""
|
|
Clear all registered models.
|
|
|
|
Useful for testing to ensure clean state.
|
|
"""
|
|
self._models.clear()
|
|
self._tables_created = False
|
|
|
|
def reset(self):
|
|
"""
|
|
Complete reset of the registry.
|
|
|
|
WARNING: This should only be used in testing.
|
|
"""
|
|
self._base = None
|
|
self._metadata = None
|
|
self._models.clear()
|
|
self._tables_created = False
|
|
|
|
def table_exists(self, engine, table_name: str) -> bool:
|
|
"""Check if a table exists in the database."""
|
|
inspector = inspect(engine)
|
|
return table_name in inspector.get_table_names()
|
|
|
|
|
|
# Global registry instance
|
|
registry = DatabaseRegistry()
|
|
|
|
|
|
def get_base() -> DeclarativeMeta:
|
|
"""Get the declarative base from the registry."""
|
|
return registry.Base
|
|
|
|
|
|
def get_metadata() -> MetaData:
|
|
"""Get the metadata from the registry."""
|
|
return registry.metadata
|
|
|
|
|
|
def declarative_base(**kwargs) -> DeclarativeMeta:
|
|
"""
|
|
Replacement for SQLAlchemy's declarative_base that uses the registry.
|
|
|
|
This ensures only one Base is ever created.
|
|
"""
|
|
return registry.Base |