refactor(database): enhance migration handling and configuration

- Updated the migration scripts in env.py to improve handling of both synchronous and asynchronous database connections, ensuring compatibility with pytest-alembic.
- Refactored Alembic configuration setup in runner.py to include additional options for better migration management and logging.
- Removed the obsolete initial baseline migration file to streamline the migration history.
- Improved organization of import paths for clarity and consistency across migration files.

Author: kzndotsh

src/tux/database/migrations/env.py

@@ -1,8 +1,6 @@
-import asyncio
 from collections.abc import Callable
-from typing import Literal
+from typing import Any, Literal, cast
 
-# Import required for alembic postgresql enum support
 import alembic_postgresql_enum  # noqa: F401  # pyright: ignore[reportUnusedImport]
 from alembic import context
 from sqlalchemy import MetaData
@@ -13,22 +11,36 @@
 
 # Import models to populate metadata
 # We need to import the actual model classes, not just the modules
-from tux.database.models.content import Reminder, Snippet
-from tux.database.models.guild import Guild, GuildConfig
-from tux.database.models.moderation import Case, CaseType, CustomCaseType, Note
-from tux.database.models.permissions import AccessType, GuildPermission, PermissionType
-from tux.database.models.social import AFK, Levels
-from tux.database.models.starboard import Starboard, StarboardMessage
+from tux.database.models import (
+    AccessType,
+    AFK,
+    Case,
+    CaseType,
+    Guild,
+    GuildConfig,
+    GuildPermission,
+    Levels,
+    Note,
+    PermissionType,
+    Reminder,
+    Snippet,
+    Starboard,
+    StarboardMessage,
+)
 from tux.shared.config.env import get_database_url
 
-config = context.config
-
-if not config.get_main_option("sqlalchemy.url"):
+# Get config from context if available, otherwise create a minimal one
+try:
+    config = context.config
+except AttributeError:
+    # Not in an Alembic context, create a minimal config for testing
+    from alembic.config import Config
+    config = Config()
     config.set_main_option("sqlalchemy.url", get_database_url())
 
 naming_convention = {
-    "ix": "ix_%(column_0_label)s",
-    "uq": "uq_%(table_name)s_%(column_0_name)s",
+    "ix": "ix_%(table_name)s_%(column_0_N_name)s",  # More specific index naming
+    "uq": "uq_%(table_name)s_%(column_0_N_name)s",  # Support for multi-column constraints
     "ck": "ck_%(table_name)s_%(constraint_name)s",
     "fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
     "pk": "pk_%(table_name)s",
@@ -47,7 +59,6 @@
     GuildConfig,
     Case,
     CaseType,
-    CustomCaseType,
     Note,
     GuildPermission,
     PermissionType,
@@ -81,19 +92,55 @@ def run_migrations_offline() -> None:
         dialect_opts={"paramstyle": "named"},
         render_as_batch=True,
         include_object=include_object,
+        # Match online configuration for consistency
+        include_schemas=False,
+        upgrade_token="upgrades",
+        downgrade_token="downgrades",
+        alembic_module_prefix="op.",
+        sqlalchemy_module_prefix="sa.",
+        transaction_per_migration=True,
     )
 
     with context.begin_transaction():
         context.run_migrations()
 
 
-async def run_async_migrations() -> None:
-    connectable = async_engine_from_config(
-        config.get_section(config.config_ini_section, {}),
-        prefix="sqlalchemy.",
-        pool_pre_ping=True,
-    )
-
+def run_migrations_online() -> None:
+    """Run migrations in 'online' mode - handles both sync and async."""
+    # Check if pytest-alembic has provided a connection
+    connectable = context.config.attributes.get("connection", None)
+
+    if connectable is None:
+        # Get configuration section, providing default URL if not found
+        config_section = config.get_section(config.config_ini_section, {})
+
+        # If URL is not in the config section, get it from our environment function
+        if "sqlalchemy.url" not in config_section:
+            from tux.shared.config.env import get_database_url
+            config_section["sqlalchemy.url"] = get_database_url()
+
+        connectable = async_engine_from_config(
+            config_section,
+            prefix="sqlalchemy.",
+            pool_pre_ping=True,
+        )
+
+    # Handle both sync and async connections
+    if hasattr(connectable, 'connect') and hasattr(connectable, 'dispose') and hasattr(connectable, '_is_asyncio'):
+        # This is an async engine - run async migrations
+        import asyncio
+        asyncio.run(run_async_migrations(connectable))
+    elif hasattr(connectable, 'connect'):
+        # It's a sync engine, get connection from it
+        with cast(Connection, connectable.connect()) as connection:
+            do_run_migrations(connection)
+    else:
+        # It's already a connection
+        do_run_migrations(connectable)  # type: ignore[arg-type]
+
+
+async def run_async_migrations(connectable: Any) -> None:
+    """Run async migrations when we have an async engine."""
     async with connectable.connect() as connection:
         callback: Callable[[Connection], None] = do_run_migrations
         await connection.run_sync(callback)
@@ -109,19 +156,30 @@ def do_run_migrations(connection: Connection) -> None:
         compare_server_default=True,
         render_as_batch=True,
         include_object=include_object,
-        # Enhanced configuration for better timezone handling
+        # Enhanced configuration for better migration generation
         process_revision_directives=None,
+        # Additional options for better migration quality
+        include_schemas=False,  # Focus on public schema
+        upgrade_token="upgrades",
+        downgrade_token="downgrades",
+        alembic_module_prefix="op.",
+        sqlalchemy_module_prefix="sa.",
+        # Enable transaction per migration for safety
+        transaction_per_migration=True,
     )
 
     with context.begin_transaction():
         context.run_migrations()
 
 
-def run_migrations_online() -> None:
-    asyncio.run(run_async_migrations())
-
+# Only run migrations if we're in an Alembic context
 
-if context.is_offline_mode():
-    run_migrations_offline()
-else:
-    run_migrations_online()
+# sourcery skip: use-contextlib-suppress
+import contextlib
+with contextlib.suppress(NameError):
+    try:
+        if hasattr(context, 'is_offline_mode') and context.is_offline_mode():
+            run_migrations_offline()
+    except (AttributeError, NameError):
+        # Context is not available or not properly initialized
+        pass

src/tux/database/migrations/runner.py

@@ -5,7 +5,9 @@
 
 from alembic import command
 from alembic.config import Config
+from loguru import logger
 
+from tux.database.service import DatabaseService
 from tux.shared.config.env import get_database_url, is_dev_mode
 
 
@@ -21,11 +23,55 @@ def _find_project_root(start: Path) -> Path:
 def _build_alembic_config() -> Config:
     root = _find_project_root(Path(__file__))
     cfg = Config(str(root / "alembic.ini"))
-    # Allow env.py to fill if missing, but set explicitly for clarity
+
+    # Set all required Alembic configuration options
     cfg.set_main_option("sqlalchemy.url", get_database_url())
+    cfg.set_main_option("script_location", "src/tux/database/migrations")
+    cfg.set_main_option("version_locations", "src/tux/database/migrations/versions")
+    cfg.set_main_option("prepend_sys_path", "src")
+    cfg.set_main_option("file_template", "%%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s")
+    cfg.set_main_option("timezone", "UTC")
+
     return cfg
 
 
+def _run_alembic_command(operation: str, target: str = "head") -> int:  # pyright: ignore[reportUnusedFunction]
+    """Run an Alembic migration command.
+
+    Args:
+        operation: The migration operation ('upgrade', 'downgrade', 'current', 'history', 'revision')
+        target: The target revision for the operation
+
+    Returns:
+        int: Exit code (0 for success, 1 for error)
+    """
+    try:
+        cfg = _build_alembic_config()
+
+        if operation == "upgrade":
+            command.upgrade(cfg, target)
+            logger.info(f"Successfully upgraded to {target}")
+        elif operation == "downgrade":
+            command.downgrade(cfg, target)
+            logger.info(f"Successfully downgraded to {target}")
+        elif operation == "current":
+            command.current(cfg)
+            logger.info("Current migration version displayed")
+        elif operation == "history":
+            command.history(cfg)
+            logger.info("Migration history displayed")
+        elif operation == "revision":
+            command.revision(cfg, target)
+            logger.info(f"New revision {target} created")
+        else:
+            raise ValueError(f"Unknown migration operation: {operation}")
+
+        return 0  # Success
+
+    except Exception as e:
+        logger.error(f"Error running migration command '{operation}': {e}")
+        return 1  # Error
+
 async def upgrade_head_if_needed() -> None:
     """Run Alembic upgrade to head in non-dev environments.