import asyncio import os import sys from contextlib import asynccontextmanager, suppress from enum import Enum, auto from pathlib import Path from typing import AsyncGenerator, Optional from basic_memory.config import BasicMemoryConfig, ConfigManager, DatabaseBackend from alembic import command from alembic.config import Config from loguru import logger from sqlalchemy import text, event from sqlalchemy.ext.asyncio import ( create_async_engine, async_sessionmaker, AsyncSession, AsyncEngine, async_scoped_session, ) from sqlalchemy.pool import NullPool from basic_memory.repository.postgres_search_repository import PostgresSearchRepository from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository # ----------------------------------------------------------------------------- # Windows event loop policy # ----------------------------------------------------------------------------- # On Windows, the default ProactorEventLoop has known rough edges with aiosqlite # during shutdown/teardown (threads posting results to a loop that's closing), # which can manifest as: # - "RuntimeError: Event loop is closed" # - "IndexError: pop from an empty deque" # # The SelectorEventLoop doesn't support subprocess operations, so code that uses # asyncio.create_subprocess_shell() (like sync_service._quick_count_files) must # detect Windows and use fallback implementations. if sys.platform == "win32": # pragma: no cover asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy()) def maybe_install_uvloop(config: BasicMemoryConfig) -> bool: """Install the uvloop event-loop policy for the Postgres backend. Trigger: process entrypoint starting with database_backend == postgres, uvloop importable, and a non-Windows platform. Why: asyncpg engine teardown (engine.dispose()) races the stdlib asyncio loop shutdown and surfaces "IndexError: pop from an empty deque" from base_events._run_once (see #831/#877). uvloop's C scheduler has no self._ready.popleft() codepath, so that class of crash cannot fire under it. Outcome: Postgres deployments run on uvloop; SQLite users keep the default loop (no behavior change, smaller blast radius). Must run before the event loop is created, i.e. before asyncio.run(). Returns: True if the uvloop policy was installed, False otherwise. """ # uvloop is not available on Windows; the default loop already differs there. if sys.platform == "win32": # pragma: no cover return False # Limit the change to the backend that actually hits the asyncpg dispose race. if config.database_backend != DatabaseBackend.POSTGRES: return False # Deferred import: uvloop is an optional, platform-gated dependency and the # default (SQLite) path must not require it to be installed. try: import uvloop except ImportError: # pragma: no cover logger.warning("uvloop not available - using default event loop for Postgres backend") return False asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) logger.info("Installed uvloop event-loop policy for Postgres backend") return True # Module level state _engine: Optional[AsyncEngine] = None _session_maker: Optional[async_sessionmaker[AsyncSession]] = None class DatabaseType(Enum): """Types of supported databases.""" MEMORY = auto() FILESYSTEM = auto() POSTGRES = auto() @classmethod def get_db_url( cls, db_path: Path, db_type: "DatabaseType", config: Optional[BasicMemoryConfig] = None ) -> str: """Get SQLAlchemy URL for database path. Args: db_path: Path to SQLite database file (ignored for Postgres) db_type: Type of database (MEMORY, FILESYSTEM, or POSTGRES) config: Optional config to check for database backend and URL Returns: SQLAlchemy connection URL """ # Load config if not provided if config is None: config = ConfigManager().config # Handle explicit Postgres type if db_type == cls.POSTGRES: if not config.database_url: raise ValueError("DATABASE_URL must be set when using Postgres backend") logger.info(f"Using Postgres database: {config.database_url}") return config.database_url # Check if Postgres backend is configured (for backward compatibility) if config.database_backend == DatabaseBackend.POSTGRES: if not config.database_url: raise ValueError("DATABASE_URL must be set when using Postgres backend") logger.info(f"Using Postgres database: {config.database_url}") return config.database_url # SQLite databases if db_type == cls.MEMORY: logger.info("Using in-memory SQLite database") return "sqlite+aiosqlite://" return f"sqlite+aiosqlite:///{db_path}" # pragma: no cover def get_scoped_session_factory( session_maker: async_sessionmaker[AsyncSession], ) -> async_scoped_session: """Create a scoped session factory scoped to current task.""" return async_scoped_session(session_maker, scopefunc=asyncio.current_task) @asynccontextmanager async def scoped_session( session_maker: async_sessionmaker[AsyncSession], ) -> AsyncGenerator[AsyncSession, None]: """ Get a scoped session with proper lifecycle management. Args: session_maker: Session maker to create scoped sessions from """ factory = get_scoped_session_factory(session_maker) session = factory() try: # Only enable foreign keys for SQLite (Postgres has them enabled by default) # Detect database type from session's bind (engine) dialect engine = session.get_bind() dialect_name = engine.dialect.name if dialect_name == "sqlite": await session.execute(text("PRAGMA foreign_keys=ON")) yield session await session.commit() except Exception: await session.rollback() raise finally: await session.close() await factory.remove() def _configure_sqlite_connection(dbapi_conn, enable_wal: bool = True) -> None: """Configure SQLite connection with WAL mode and optimizations. Args: dbapi_conn: Database API connection object enable_wal: Whether to enable WAL mode (should be False for in-memory databases) """ cursor = dbapi_conn.cursor() try: # Enable WAL mode for better concurrency (not supported for in-memory databases) if enable_wal: cursor.execute("PRAGMA journal_mode=WAL") # Set busy timeout to handle locked databases cursor.execute("PRAGMA busy_timeout=10000") # 10 seconds # Optimize for performance cursor.execute("PRAGMA synchronous=NORMAL") cursor.execute("PRAGMA cache_size=-64000") # 64MB cache cursor.execute("PRAGMA temp_store=MEMORY") # Windows-specific optimizations if os.name == "nt": cursor.execute("PRAGMA locking_mode=NORMAL") # Ensure normal locking on Windows except Exception as e: # Log but don't fail - some PRAGMAs may not be supported logger.warning(f"Failed to configure SQLite connection: {e}") finally: cursor.close() def _create_sqlite_engine(db_url: str, db_type: DatabaseType) -> AsyncEngine: """Create SQLite async engine with appropriate configuration. Args: db_url: SQLite connection URL db_type: Database type (MEMORY or FILESYSTEM) Returns: Configured async engine for SQLite """ # Configure connection args with Windows-specific settings connect_args: dict[str, bool | float | None] = {"check_same_thread": False} # Add Windows-specific parameters to improve reliability if os.name == "nt": # Windows connect_args.update( { "timeout": 30.0, # Increase timeout to 30 seconds for Windows "isolation_level": None, # Use autocommit mode } ) # Use NullPool for Windows filesystem databases to avoid connection pooling issues # Important: Do NOT use NullPool for in-memory databases as it will destroy the database # between connections if db_type == DatabaseType.FILESYSTEM: engine = create_async_engine( db_url, connect_args=connect_args, poolclass=NullPool, # Disable connection pooling on Windows echo=False, ) else: # In-memory databases need connection pooling to maintain state engine = create_async_engine(db_url, connect_args=connect_args) else: engine = create_async_engine(db_url, connect_args=connect_args) # Enable WAL mode for better concurrency and reliability # Note: WAL mode is not supported for in-memory databases enable_wal = db_type != DatabaseType.MEMORY @event.listens_for(engine.sync_engine, "connect") def enable_wal_mode(dbapi_conn, connection_record): """Enable WAL mode on each connection.""" _configure_sqlite_connection(dbapi_conn, enable_wal=enable_wal) return engine def _create_postgres_engine(db_url: str, config: BasicMemoryConfig) -> AsyncEngine: """Create Postgres async engine with appropriate configuration. Args: db_url: Postgres connection URL (postgresql+asyncpg://...) config: BasicMemoryConfig with pool settings Returns: Configured async engine for Postgres """ # Use NullPool connection issues. # Assume connection pooler like PgBouncer handles connection pooling. engine = create_async_engine( db_url, echo=False, poolclass=NullPool, # No pooling - fresh connection per request connect_args={ # Disable statement cache to avoid issues with prepared statements on reconnect "statement_cache_size": 0, # Allow 30s for commands (Neon cold start can take 2-5s, sometimes longer) "command_timeout": 30, # Allow 30s for initial connection (Neon wake-up time) "timeout": 30, "server_settings": { "application_name": "basic-memory", # Statement timeout for queries (30s to allow for cold start) "statement_timeout": "30s", }, }, ) logger.debug("Created Postgres engine with NullPool (no connection pooling)") return engine def _create_engine_and_session( db_path: Path, db_type: DatabaseType = DatabaseType.FILESYSTEM, config: Optional[BasicMemoryConfig] = None, ) -> tuple[AsyncEngine, async_sessionmaker[AsyncSession]]: """Internal helper to create engine and session maker. Args: db_path: Path to database file (used for SQLite, ignored for Postgres) db_type: Type of database (MEMORY, FILESYSTEM, or POSTGRES) config: Optional explicit config. If not provided, reads from ConfigManager. Prefer passing explicitly from composition roots. Returns: Tuple of (engine, session_maker) """ # Prefer explicit parameter; fall back to ConfigManager for backwards compatibility if config is None: config = ConfigManager().config db_url = DatabaseType.get_db_url(db_path, db_type, config) logger.debug(f"Creating engine for db_url: {db_url}") # Delegate to backend-specific engine creation # Check explicit POSTGRES type first, then config setting if db_type == DatabaseType.POSTGRES or config.database_backend == DatabaseBackend.POSTGRES: engine = _create_postgres_engine(db_url, config) else: engine = _create_sqlite_engine(db_url, db_type) session_maker = async_sessionmaker(engine, expire_on_commit=False) return engine, session_maker async def get_or_create_db( db_path: Path, db_type: DatabaseType = DatabaseType.FILESYSTEM, ensure_migrations: bool = True, config: Optional[BasicMemoryConfig] = None, ) -> tuple[AsyncEngine, async_sessionmaker[AsyncSession]]: # pragma: no cover """Get or create database engine and session maker. Args: db_path: Path to database file db_type: Type of database ensure_migrations: Whether to run migrations config: Optional explicit config. If not provided, reads from ConfigManager. Prefer passing explicitly from composition roots. """ global _engine, _session_maker # Prefer explicit parameter; fall back to ConfigManager for backwards compatibility if config is None: config = ConfigManager().config if _engine is None: _engine, _session_maker = _create_engine_and_session(db_path, db_type, config) # Run migrations automatically unless explicitly disabled if ensure_migrations: await run_migrations(config, db_type) # These checks should never fail since we just created the engine and session maker # if they were None, but we'll check anyway for the type checker if _engine is None: logger.error("Failed to create database engine", db_path=str(db_path)) raise RuntimeError("Database engine initialization failed") if _session_maker is None: logger.error("Failed to create session maker", db_path=str(db_path)) raise RuntimeError("Session maker initialization failed") return _engine, _session_maker async def shutdown_db() -> None: # pragma: no cover """Clean up database connections.""" global _engine, _session_maker if _engine: # Trigger: teardown can run while the surrounding task is being cancelled # (e.g. lifespan shutdown, unshielded CLI cleanup). # Why: a cancellation landing mid-dispose surfaces the asyncpg # "IndexError: pop from an empty deque" race (#831/#877); shielding lets # dispose finish atomically, and suppressing CancelledError keeps a # cancelled shutdown from re-raising the underlying race. # Outcome: connections always close cleanly even under cancellation. with suppress(asyncio.CancelledError): await asyncio.shield(_engine.dispose()) _engine = None _session_maker = None @asynccontextmanager async def engine_session_factory( db_path: Path, db_type: DatabaseType = DatabaseType.MEMORY, config: Optional[BasicMemoryConfig] = None, ) -> AsyncGenerator[tuple[AsyncEngine, async_sessionmaker[AsyncSession]], None]: """Create engine and session factory. Note: This is primarily used for testing where we want a fresh database for each test. For production use, use get_or_create_db() instead. Args: db_path: Path to database file db_type: Type of database config: Optional explicit config. If not provided, reads from ConfigManager. """ global _engine, _session_maker # Use the same helper function as production code. # # Keep local references so teardown can deterministically dispose the # specific engine created by this context manager, even if other code calls # shutdown_db() and mutates module-level globals mid-test. created_engine, created_session_maker = _create_engine_and_session(db_path, db_type, config) _engine, _session_maker = created_engine, created_session_maker try: # Verify that engine and session maker are initialized if created_engine is None: # pragma: no cover logger.error("Database engine is None in engine_session_factory") raise RuntimeError("Database engine initialization failed") if created_session_maker is None: # pragma: no cover logger.error("Session maker is None in engine_session_factory") raise RuntimeError("Session maker initialization failed") yield created_engine, created_session_maker finally: # Trigger: context-manager teardown can run while the surrounding task is # being cancelled (e.g. a test aborting mid-fixture). # Why: on the asyncpg backend a cancellation landing mid-dispose surfaces # the "IndexError: pop from an empty deque" race (#831/#877); shield the # dispose and suppress CancelledError to match the other dispose seams. # Outcome: the per-context engine always disposes cleanly under cancellation. with suppress(asyncio.CancelledError): await asyncio.shield(created_engine.dispose()) # Only clear module-level globals if they still point to this context's # engine/session. This avoids clobbering newer globals from other callers. if _engine is created_engine: _engine = None if _session_maker is created_session_maker: _session_maker = None async def run_migrations( app_config: BasicMemoryConfig, database_type=DatabaseType.FILESYSTEM ): # pragma: no cover """Run any pending alembic migrations. Note: Alembic tracks which migrations have been applied via the alembic_version table, so it's safe to call this multiple times - it will only run pending migrations. """ logger.info("Running database migrations...") temp_engine: AsyncEngine | None = None try: # Get the absolute path to the alembic directory relative to this file alembic_dir = Path(__file__).parent / "alembic" config = Config() # Set required Alembic config options programmatically config.set_main_option("script_location", str(alembic_dir)) config.set_main_option( "file_template", "%%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s", ) config.set_main_option("timezone", "UTC") config.set_main_option("revision_environment", "false") # Get the correct database URL based on backend configuration # No URL conversion needed - env.py now handles both async and sync engines db_url = DatabaseType.get_db_url(app_config.database_path, database_type, app_config) config.set_main_option("sqlalchemy.url", db_url) command.upgrade(config, "head") logger.info("Migrations completed successfully") # Get session maker - ensure we don't trigger recursive migration calls if _session_maker is None: temp_engine, session_maker = _create_engine_and_session( app_config.database_path, database_type, app_config ) else: session_maker = _session_maker # Initialize the search index schema # For SQLite: Create FTS5 virtual table # For Postgres: No-op (tsvector column added by migrations) # The project_id is not used for init_search_index, so we pass a dummy value if ( database_type == DatabaseType.POSTGRES or app_config.database_backend == DatabaseBackend.POSTGRES ): await PostgresSearchRepository(session_maker, 1).init_search_index() else: await SQLiteSearchRepository(session_maker, 1).init_search_index() except Exception as e: # pragma: no cover logger.error(f"Error running migrations: {e}") raise finally: # Trigger: run_migrations() created a temporary engine while module-level # session maker was not initialized. # Why: temporary aiosqlite worker threads can outlive CLI command execution # and block process shutdown if the engine is not disposed. On the asyncpg # backend a cancellation landing mid-dispose surfaces the same "IndexError: # pop from an empty deque" race as the other dispose seams (#831/#877), so # shield the dispose and suppress CancelledError to match them. # Outcome: always dispose temporary engines cleanly, even under cancellation. if temp_engine is not None: with suppress(asyncio.CancelledError): await asyncio.shield(temp_engine.dispose())