sqlobjects 0.1.0__py3-none-any.whl

sqlobjects/database.py

@@ -0,0 +1,586 @@
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Any, Protocol
+
+from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
+
+from .session import SessionContextManager
+
+
+__all__ = [
+    "DatabaseConfig",
+    "Database",
+    "DatabaseManager",
+    "DatabaseObserver",
+    "init_db",
+    "init_dbs",
+    "create_tables",
+    "drop_tables",
+    "close_db",
+    "close_dbs",
+    "close_all_dbs",
+    "set_default_db",
+]
+
+
+class DatabaseObserver(Protocol):
+    """Database event observer protocol"""
+
+    def on_database_added(self, name: str, database: "Database", is_default: bool) -> None: ...
+    def on_database_closed(self, name: str) -> None: ...
+    def on_default_changed(self, old_default: str | None, new_default: str | None) -> None: ...
+
+
+@dataclass(init=False)
+class DatabaseConfig:
+    """Database configuration class
+
+    Uses dataclass to automatically generate initialization and other methods.
+
+    Attributes:
+        url: Database connection URL
+        echo: Whether to print SQL statements
+        pool_size: Connection pool size
+        max_overflow: Maximum pool overflow count
+        pool_timeout: Connection timeout in seconds
+        pool_recycle: Connection recycle time in seconds
+        engine_kwargs: Additional SQLAlchemy engine parameters
+
+    Examples:
+        >>> config = DatabaseConfig(
+        ...     url="postgresql+asyncpg://user:pass@localhost/mydb",
+        ...     pool_size=10,
+        ...     echo=True,
+        ...     isolation_level="READ_COMMITTED",
+        ... )
+    """
+
+    url: str
+    echo: bool
+    pool_size: int
+    max_overflow: int
+    pool_timeout: int
+    pool_recycle: int
+    engine_kwargs: dict[str, Any]
+
+    def __init__(
+        self,
+        url: str,
+        echo: bool = False,
+        pool_size: int = 5,
+        max_overflow: int = 10,
+        pool_timeout: int = 30,
+        pool_recycle: int = 3600,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize database configuration
+
+        Args:
+            url: Database connection URL
+            echo: Whether to print SQL statements
+            pool_size: Connection pool size
+            max_overflow: Maximum pool overflow count
+            pool_timeout: Connection timeout in seconds
+            pool_recycle: Connection recycle time in seconds
+            **kwargs: Additional SQLAlchemy engine parameters
+        """
+        self.url = url
+        self.echo = echo
+        self.pool_size = pool_size
+        self.max_overflow = max_overflow
+        self.pool_timeout = pool_timeout
+        self.pool_recycle = pool_recycle
+        self.engine_kwargs = kwargs
+
+
+class Database:
+    """Single database connection with event handling and table operations
+
+    Represents a single database connection, providing unified event registration
+    interface and table operation methods.
+
+    Attributes:
+        name: Unique name for the database connection
+        config: Database configuration
+        engine: SQLAlchemy async engine
+
+    Examples:
+        >>> config = DatabaseConfig(url="sqlite+aiosqlite:///test.db")
+        >>> db = Database("main", config)
+        >>> @db.on("connect")
+        ... def on_connect(conn, record):
+        ...     print("Database connected")
+    """
+
+    def __init__(self, name: str, config: DatabaseConfig) -> None:
+        """Initialize database instance
+
+        Args:
+            name: Unique name for the database connection
+            config: Database configuration
+        """
+        self.name = name
+        self.config = config
+
+        # Build engine parameters
+        engine_kwargs: dict[str, Any] = {
+            "echo": config.echo,
+            **config.engine_kwargs,
+        }
+
+        # Add connection pool parameters for non-SQLite databases
+        if not config.url.startswith("sqlite"):
+            engine_kwargs.update(
+                {
+                    "pool_size": config.pool_size,
+                    "max_overflow": config.max_overflow,
+                    "pool_timeout": config.pool_timeout,
+                    "pool_recycle": config.pool_recycle,
+                }
+            )
+
+        # Create async engine
+        self.engine: AsyncEngine = create_async_engine(config.url, **engine_kwargs)
+
+    def on(self, event_name: str, target=None):
+        """Unified event registration method
+
+        Args:
+            event_name: Event name (connect, close, before_commit, etc.)
+            target: Event target, automatically selected by default
+
+        Returns:
+            SQLAlchemy event listener decorator
+
+        Examples:
+            >>> @db.on("connect")
+            ... def on_connect(conn, record):
+            ...     print("Database connected")
+
+            >>> @db.on("before_commit")
+            ... def before_commit(conn):
+            ...     print("About to commit transaction")
+        """
+        from sqlalchemy import event
+
+        # Automatically select event target
+        if target is None:
+            target = self.engine.sync_engine
+
+        return event.listens_for(target, event_name)
+
+    async def create_tables(self, metadata) -> None:
+        """Create all tables defined in the metadata
+
+        Creates all tables, indexes, and constraints defined in the provided
+        SQLAlchemy metadata object using the database engine.
+
+        Args:
+            metadata: SQLAlchemy metadata object containing table definitions
+
+        Examples:
+            >>> from sqlobjects.base import ObjectModel
+            >>> await db.create_tables(ObjectModel.__registry__)
+        """
+        async with self.engine.begin() as conn:
+            await conn.run_sync(metadata.create_all)
+
+    async def drop_tables(self, metadata) -> None:
+        """Drop all tables defined in the metadata
+
+        Drops all tables, indexes, and constraints defined in the provided
+        SQLAlchemy metadata object from the database.
+
+        Args:
+            metadata: SQLAlchemy metadata object containing table definitions
+
+        Examples:
+            >>> from sqlobjects.base import ObjectModel
+            >>> await db.drop_tables(ObjectModel.__registry__)
+        """
+        async with self.engine.begin() as conn:
+            await conn.run_sync(metadata.drop_all)
+
+    async def disconnect(self) -> None:
+        """Disconnect database and clean up resources
+
+        Properly disposes the SQLAlchemy engine and closes all connections.
+        Should be called when the database is no longer needed.
+        """
+        await self.engine.dispose()
+
+
+class DatabaseManager:
+    """Multi-database connection manager
+
+    Manages multiple database connections, handles default database selection,
+    provides table operations and connection lifecycle management.
+    Uses observer pattern to decouple from ConnectionContextManager.
+
+    Examples:
+        >>> manager = DatabaseManager()
+        >>> await manager.add_database("main", main_config, is_default=True)
+        >>> await manager.add_database("analytics", analytics_config)
+        >>> await manager.create_tables(ObjectModel, "main")
+    """
+
+    def __init__(self) -> None:
+        """Initialize database manager"""
+        self._databases: dict[str, Database] = {}
+        self._default_db: str | None = None
+        self._observers: list[DatabaseObserver] = []
+
+    def _notify_observers(self, event: str, **kwargs) -> None:
+        """Notify all observers"""
+        for observer in self._observers:
+            getattr(observer, event)(**kwargs)
+
+    def add_observer(self, observer: DatabaseObserver) -> None:
+        """Add database event observer
+
+        Args:
+            observer: Observer instance implementing DatabaseObserver protocol
+        """
+        self._observers.append(observer)
+
+    def remove_observer(self, observer: DatabaseObserver) -> None:
+        """Remove database event observer
+
+        Args:
+            observer: Observer instance to remove
+
+        Raises:
+            ValueError: When observer is not found in the list
+        """
+        self._observers.remove(observer)
+
+    async def add_database(self, name: str, config: DatabaseConfig, is_default: bool = False) -> Database:
+        """Add database connection
+
+        Args:
+            name: Unique database name
+            config: Database configuration
+            is_default: Whether to set as default database
+
+        Returns:
+            Created database instance
+
+        Raises:
+            ValueError: When database connection fails
+        """
+        try:
+            database = Database(name, config)
+            self._databases[name] = database
+
+            old_default = self._default_db
+            if is_default:
+                self._default_db = name
+
+            # Notify observers
+            self._notify_observers("on_database_added", name=name, database=database, is_default=is_default)
+            if is_default and old_default != name:
+                self._notify_observers("on_default_changed", old_default=old_default, new_default=name)
+
+            return database
+        except Exception as e:
+            raise ValueError(f"Failed to connect to database '{name}': {e}") from e
+
+    def get_database(self, db_name: str | None = None) -> Database:
+        """Get database instance
+
+        Args:
+            db_name: Database name, uses default database when None
+
+        Returns:
+            Database instance
+
+        Raises:
+            ValueError: When database does not exist
+        """
+        name = db_name or self._default_db
+        if not name or name not in self._databases:
+            raise ValueError(f"Database '{name}' not found")
+        return self._databases[name]
+
+    async def create_tables(self, base_class, db_name: str | None = None) -> None:
+        """Create all tables defined in the base class registry
+
+        Creates all tables, indexes, and constraints for models registered
+        in the base class registry using the specified database connection.
+
+        Args:
+            base_class: SQLObjects base class containing model registry
+            db_name: Database name to use, uses default database if None
+
+        Raises:
+            ValueError: When specified database does not exist
+
+        Examples:
+            >>> from sqlobjects.base import ObjectModel
+            >>> await manager.create_tables(ObjectModel)
+            >>> await manager.create_tables(ObjectModel, "analytics")
+        """
+        database = self.get_database(db_name)
+        await database.create_tables(base_class.__registry__)
+
+    async def drop_tables(self, base_class, db_name: str | None = None) -> None:
+        """Drop all tables defined in the base class registry
+
+        Drops all tables, indexes, and constraints for models registered
+        in the base class registry from the specified database connection.
+
+        Args:
+            base_class: SQLObjects base class containing model registry
+            db_name: Database name to use, uses default database if None
+
+        Raises:
+            ValueError: When specified database does not exist
+
+        Examples:
+            >>> from sqlobjects.base import ObjectModel
+            >>> await manager.drop_tables(ObjectModel)
+            >>> await manager.drop_tables(ObjectModel, "analytics")
+        """
+        database = self.get_database(db_name)
+        await database.drop_tables(base_class.__registry__)
+
+    async def close(self, db_name: str | None = None, auto_default: bool = False) -> None:
+        """Close database connection and clean up resources
+
+        Closes the specified database connection and removes it from the manager.
+        Handles default database reassignment when closing the default database.
+
+        Args:
+            db_name: Database name to close, closes default when None
+            auto_default: Whether to automatically select new default when closing default database
+
+        Raises:
+            ValueError: When specified database does not exist
+        """
+        # Determine target database name
+        target_db = db_name or self._default_db
+        if not target_db or target_db not in self._databases:
+            raise ValueError(f"Database '{target_db}' not found")
+
+        # Close specified database
+        await self._databases[target_db].engine.dispose()
+        del self._databases[target_db]
+        self._notify_observers("on_database_closed", name=target_db)
+
+        # Handle default database change if closing default
+        if self._default_db == target_db:
+            old_default = self._default_db
+            if auto_default:
+                self._default_db = next(iter(self._databases), None)
+            else:
+                self._default_db = None
+            self._notify_observers("on_default_changed", old_default=old_default, new_default=self._default_db)
+
+    async def close_all(self) -> None:
+        """Close all database connections and clean up resources
+
+        Closes all managed database connections, disposes their engines,
+        and resets the manager to initial state.
+        """
+        for name, db in self._databases.items():
+            await db.engine.dispose()
+            self._notify_observers("on_database_closed", name=name)
+
+        old_default = self._default_db
+        self._databases.clear()
+        self._default_db = None
+        if old_default:
+            self._notify_observers("on_default_changed", old_default=old_default, new_default=None)
+
+    def set_default_db(self, db_name: str) -> None:
+        """Set default database for operations
+
+        Changes the default database used when no specific database
+        name is provided in operations.
+
+        Args:
+            db_name: Database name to set as default
+
+        Raises:
+            ValueError: When database does not exist
+        """
+        if db_name not in self._databases:
+            raise ValueError(f"Database '{db_name}' not found")
+
+        old_default = self._default_db
+        self._default_db = db_name
+        self._notify_observers("on_default_changed", old_default=old_default, new_default=db_name)
+
+
+# Global database manager instance
+_manager = DatabaseManager()
+
+# Register TransactionContextManager as observer
+_manager.add_observer(SessionContextManager)
+
+
+async def init_db(
+    url: str,
+    name: str | None = None,
+    echo: bool = False,
+    pool_size: int = 5,
+    max_overflow: int = 10,
+    pool_timeout: int = 30,
+    pool_recycle: int = 3600,
+    is_default: bool = True,
+    **engine_kwargs: Any,
+) -> Database:
+    """Initialize single database connection.
+
+    Args:
+        url: Database URL (e.g., 'sqlite+aiosqlite:///db.sqlite', 'postgresql+asyncpg://user:pass@host/db')
+        name: Name for the database connection, uses "default" if None
+        echo: Whether to log all SQL statements
+        pool_size: Number of connections to maintain in the pool
+        max_overflow: Maximum number of connections that can overflow the pool
+        pool_timeout: Timeout in seconds for getting connection from pool
+        pool_recycle: Time in seconds to recycle connections
+        is_default: Whether this database should be set as the default database
+        **engine_kwargs: Additional SQLAlchemy engine arguments
+
+    Returns:
+        Database instance with configured connection
+
+    Raises:
+        ValueError: If database URL format is invalid
+        DatabaseError: If connection to database fails
+        ImportError: If required database driver is not installed
+    """
+    config = DatabaseConfig(
+        url=url,
+        echo=echo,
+        pool_size=pool_size,
+        max_overflow=max_overflow,
+        pool_timeout=pool_timeout,
+        pool_recycle=pool_recycle,
+        **engine_kwargs,
+    )
+    db_name = name if name is not None else "default"
+    return await _manager.add_database(db_name, config, is_default=is_default)
+
+
+async def init_dbs(
+    databases: Mapping[str, dict[str, Any] | DatabaseConfig],
+    default: str | None = None,
+) -> tuple[Database, ...]:
+    """Initialize multiple database connections.
+
+    Args:
+        databases: Dictionary mapping database names to their configurations
+        default: Name of the default database to use when none is specified, or None for no default
+
+    Returns:
+        Tuple of Database instances in the order they appear in the databases dict
+
+    Raises:
+        ValueError: If default database name is not in databases dict or URL format is invalid
+        DatabaseError: If connection to any database fails
+        ImportError: If required database drivers are not installed
+    """
+    db_instances = []
+
+    for name, config_data in databases.items():
+        if isinstance(config_data, DatabaseConfig):
+            config = config_data
+        else:
+            config = DatabaseConfig(**config_data)
+
+        is_default = default is not None and name == default
+        database = await _manager.add_database(name, config, is_default)
+        db_instances.append(database)
+
+    return tuple(db_instances)
+
+
+async def create_tables(base_class, db_name: str | None = None) -> None:
+    """Create all tables defined in the base class registry
+
+    Creates all tables, indexes, and constraints for models registered
+    in the base class registry using the specified database connection.
+
+    Args:
+        base_class: SQLObjects base class containing model registry
+        db_name: Name of the database, uses default if None
+
+    Raises:
+        ValueError: When specified database does not exist
+    """
+    await _manager.create_tables(base_class, db_name)
+
+
+async def drop_tables(base_class, db_name: str | None = None) -> None:
+    """Drop all tables defined in the base class registry
+
+    Drops all tables, indexes, and constraints for models registered
+    in the base class registry from the specified database connection.
+
+    Args:
+        base_class: SQLObjects base class containing model registry
+        db_name: Name of the database, uses default if None
+
+    Raises:
+        ValueError: When specified database does not exist
+    """
+    await _manager.drop_tables(base_class, db_name)
+
+
+async def close_db(db_name: str | None = None, auto_default: bool = False) -> None:
+    """Close database connection and clean up resources
+
+    Closes the specified database connection and removes it from the manager.
+    Handles default database reassignment when closing the default database.
+
+    Args:
+        db_name: Name of specific database to close, closes default if None
+        auto_default: Whether to update default database when closing the default database
+
+    Raises:
+        ValueError: When specified database does not exist
+    """
+    await _manager.close(db_name, auto_default)
+
+
+async def close_dbs(db_names: list[str], auto_default: bool = False) -> None:
+    """Close multiple specific database connections
+
+    Closes all specified database connections in sequence.
+    Handles default database reassignment if the default database is closed.
+
+    Args:
+        db_names: List of database names to close
+        auto_default: Whether to update default database when closing the default database
+
+    Raises:
+        ValueError: When any specified database does not exist
+    """
+    for db_name in db_names:
+        await _manager.close(db_name, auto_default)
+
+
+async def close_all_dbs() -> None:
+    """Close all database connections and clean up resources
+
+    Closes all managed database connections, disposes their engines,
+    and resets the global manager to initial state.
+    """
+    await _manager.close_all()
+
+
+def set_default_db(db_name: str) -> None:
+    """Set the default database by name
+
+    Changes the default database used when no specific database
+    name is provided in operations.
+
+    Args:
+        db_name: Name of the database to set as default
+
+    Raises:
+        ValueError: If database is not found
+    """
+    _manager.set_default_db(db_name)