routstr-core/routstr/core/db.py

import asyncio
import json
import os
import pathlib
import sqlite3
import time
import uuid
from contextlib import asynccontextmanager
from typing import AsyncGenerator

from alembic import command
from alembic.config import Config
from alembic.util.exc import CommandError
from sqlalchemy import UniqueConstraint, delete
from sqlalchemy.exc import IntegrityError, OperationalError
from sqlalchemy.ext.asyncio.engine import create_async_engine
from sqlalchemy.orm import aliased
from sqlmodel import Field, Relationship, SQLModel, col, func, select, update
from sqlmodel.ext.asyncio.session import AsyncSession

from .logging import get_logger

logger = get_logger(__name__)

DATABASE_URL = os.environ.get("DATABASE_URL", "sqlite+aiosqlite:///keys.db")


engine = create_async_engine(DATABASE_URL, echo=False)  # echo=True for debugging SQL


# Durable JSONL outbox for cashu transactions when the DB is down.
_cashu_outbox_lock = asyncio.Lock()


def _cashu_outbox_path() -> pathlib.Path:
    """Resolve the per-process cashu outbox file path."""
    override = os.environ.get("CASHU_OUTBOX_PATH")
    if override:
        return pathlib.Path(override)

    filename = f"cashu_outbox.{os.getpid()}.jsonl"

    if DATABASE_URL.startswith("sqlite"):
        raw = DATABASE_URL.split("///", 1)[-1]
        if not raw or raw == ":memory:":
            return pathlib.Path("/tmp").resolve() / filename
        db_file = pathlib.Path(raw)
        return db_file.parent.resolve() / filename

    return pathlib.Path(filename).resolve()


async def append_to_cashu_outbox(
    *,
    token: str,
    amount: int,
    unit: str,
    mint_url: str | None,
    typ: str,
    request_id: str | None,
    collected: bool,
    created_at: int | None,
    source: str,
    api_key_hashed_key: str | None,
) -> None:
    """Append a transaction payload to the durable outbox. Never raises."""
    entry = {
        "outbox_id": uuid.uuid4().hex,
        "queued_at": int(time.time()),
        "token": token,
        "amount": amount,
        "unit": unit,
        "mint_url": mint_url,
        "type": typ,
        "request_id": request_id,
        "collected": collected,
        "created_at": created_at,
        "source": source,
        "api_key_hashed_key": api_key_hashed_key,
    }
    line = json.dumps(entry, separators=(",", ":")) + "\n"
    async with _cashu_outbox_lock:
        try:
            path = _cashu_outbox_path()
            path.parent.mkdir(parents=True, exist_ok=True)
            with open(path, "a", encoding="utf-8") as fh:
                fh.write(line)
                fh.flush()
                os.fsync(fh.fileno())
            logger.warning(
                "Cashu transaction spooled to durable outbox",
                extra={"outbox_path": str(path), "type": typ, "request_id": request_id},
            )
        except Exception as outbox_exc:
            logger.critical(
                "cashu outbox append failed; token may be unrecoverable",
                extra={
                    "error": str(outbox_exc),
                    "type": typ,
                    "request_id": request_id,
                    "token": token,
                },
            )


class ApiKey(SQLModel, table=True):  # type: ignore
    __tablename__ = "api_keys"

    hashed_key: str = Field(primary_key=True)
    balance: int = Field(default=0, description="Balance in millisatoshis (msats)")
    reserved_balance: int = Field(
        default=0, description="Reserved balance in millisatoshis (msats)"
    )
    reserved_at: int | None = Field(
        default=None,
        description=(
            "Unix timestamp of the most recent balance reservation. Used to "
            "detect and release stale reservations (e.g. after client "
            "disconnects). NULL when no reservation has been made yet."
        ),
    )
    refund_address: str | None = Field(
        default=None,
        description="Lightning address to refund remaining balance after key expires",
    )
    key_expiry_time: int | None = Field(
        default=None,
        description="Unix-timestamp after which the cashu-token's balance gets refunded to the refund_address",
    )
    total_spent: int = Field(
        default=0, description="Total spent in millisatoshis (msats)"
    )
    total_requests: int = Field(default=0)
    created_at: int | None = Field(
        default_factory=lambda: int(time.time()),
        nullable=True,
        description=(
            "Unix timestamp when the key was created. Nullable: keys created "
            "before this column existed have no value and sort last."
        ),
    )
    refund_mint_url: str | None = Field(
        default=None,
        description="URL of the mint used to create the cashu-token",
    )
    refund_currency: str | None = Field(
        default=None,
        description="Currency of the cashu-token",
    )
    parent_key_hash: str | None = Field(
        default=None, foreign_key="api_keys.hashed_key", index=True
    )
    balance_limit: int | None = Field(
        default=None,
        description="Max spendable balance in msats for this key (mostly for child keys)",
    )
    balance_limit_reset: str | None = Field(
        default=None,
        description="Reset policy for balance limit (manual, daily, monthly, etc.)",
    )
    balance_limit_reset_date: int | None = Field(
        default=None,
        description="Unix timestamp of the last time the balance limit was reset",
    )
    validity_date: int | None = Field(
        default=None,
        description="Unix timestamp after which the key is no longer valid",
    )

    @property
    def total_balance(self) -> int:
        return self.balance - self.reserved_balance


async def reset_all_reserved_balances(session: AsyncSession) -> None:
    stmt = update(ApiKey).values(reserved_balance=0, reserved_at=None)
    await session.exec(stmt)  # type: ignore[call-overload]
    await session.commit()
    logger.info("Reset reserved balances on startup")


async def release_stale_reservations(
    session: AsyncSession, max_age_seconds: int
) -> int:
    """Release reservations whose last reserve is older than max_age_seconds.
    """
    cutoff = int(time.time()) - max_age_seconds
    stmt = (
        update(ApiKey)
        .where(col(ApiKey.reserved_balance) > 0)
        .where(col(ApiKey.reserved_at).is_not(None))
        .where(col(ApiKey.reserved_at) < cutoff)
        .values(reserved_balance=0, reserved_at=None)
    )
    result = await session.exec(stmt)  # type: ignore[call-overload]
    await session.commit()
    released = int(result.rowcount or 0)
    if released:
        logger.warning(
            "Released stale balance reservations",
            extra={"released_keys": released, "max_age_seconds": max_age_seconds},
        )
    return released


async def prune_dead_api_keys(session: AsyncSession, min_age_seconds: int) -> int:
    """Delete dead parentless API keys; return the count removed.

    Dead = 0 balance/reservation/spend/requests, older than the grace period,
    no parent, no children, no pending invoice. Cashu rows are unlinked (not
    deleted) first to keep the audit trail.
    """
    cutoff = int(time.time()) - min_age_seconds

    child = aliased(ApiKey)
    has_children = (
        select(child.hashed_key).where(
            col(child.parent_key_hash) == col(ApiKey.hashed_key)
        )
    ).exists()
    pending_invoice = (
        select(LightningInvoice.id)
        .where(col(LightningInvoice.api_key_hash) == col(ApiKey.hashed_key))
        .where(col(LightningInvoice.status) == "pending")
    ).exists()

    eligible_hashes = (
        select(ApiKey.hashed_key)
        .where(col(ApiKey.balance) == 0)
        .where(col(ApiKey.reserved_balance) == 0)
        .where(col(ApiKey.total_spent) == 0)
        .where(col(ApiKey.total_requests) == 0)
        .where(col(ApiKey.parent_key_hash).is_(None))
        .where(
            (col(ApiKey.created_at).is_(None)) | (col(ApiKey.created_at) < cutoff)
        )
        .where(~pending_invoice)
        .where(~has_children)
    )

    # Unlink transactions rather than cascade-deleting them, so the financial
    # audit trail survives. The eligibility predicate is re-evaluated inside both
    # statements so a key that gained balance mid-run is left untouched.
    await session.exec(  # type: ignore[call-overload]
        update(CashuTransaction)
        .where(col(CashuTransaction.api_key_hashed_key).in_(eligible_hashes))
        .values(api_key_hashed_key=None)
    )

    result = await session.exec(  # type: ignore[call-overload]
        delete(ApiKey).where(col(ApiKey.hashed_key).in_(eligible_hashes))
    )
    await session.commit()

    pruned = int(result.rowcount or 0)
    logger.info(
        "Pruned dead API keys",
        extra={"pruned_keys": pruned, "min_age_seconds": min_age_seconds},
    )
    return pruned


class ModelRow(SQLModel, table=True):  # type: ignore
    __tablename__ = "models"
    id: str = Field(primary_key=True)
    upstream_provider_id: int = Field(
        primary_key=True, foreign_key="upstream_providers.id", ondelete="CASCADE"
    )
    name: str = Field()
    created: int = Field()
    description: str = Field()
    context_length: int = Field()
    architecture: str = Field()
    pricing: str = Field()
    sats_pricing: str | None = Field(default=None)
    per_request_limits: str | None = Field(default=None)
    top_provider: str | None = Field(default=None)
    canonical_slug: str | None = Field(default=None, description="Canonical model slug")
    alias_ids: str | None = Field(
        default=None, description="JSON array of model alias IDs"
    )
    enabled: bool = Field(default=True, description="Whether this model is enabled")
    forwarded_model_id: str | None = Field(
        default=None,
        description="Model ID to use when forwarding requests to upstream provider. Defaults to id if not set.",
    )
    upstream_provider: "UpstreamProviderRow" = Relationship(back_populates="models")


class LightningInvoice(SQLModel, table=True):  # type: ignore
    __tablename__ = "lightning_invoices"

    id: str = Field(primary_key=True, description="Unique invoice identifier")
    bolt11: str = Field(description="BOLT11 invoice string", unique=True)
    amount_sats: int = Field(description="Amount in satoshis")
    description: str = Field(description="Invoice description")
    payment_hash: str = Field(description="Payment hash for tracking", unique=True)
    status: str = Field(
        default="pending", description="pending, paid, expired, cancelled"
    )
    api_key_hash: str | None = Field(
        default=None, description="Associated API key hash for topup operations"
    )
    purpose: str = Field(description="create or topup")
    mint_url: str | None = Field(
        default=None, description="Mint URL where the quote was created (fallback tracking)"
    )
    created_at: int = Field(
        default_factory=lambda: int(time.time()), description="Unix timestamp"
    )
    expires_at: int = Field(description="Unix timestamp when invoice expires")
    paid_at: int | None = Field(default=None, description="Unix timestamp when paid")
    balance_limit: int | None = Field(
        default=None,
        description="Max spendable msats for the created key",
    )
    balance_limit_reset: str | None = Field(
        default=None,
        description="Reset policy for balance limit (daily, weekly, monthly)",
    )
    validity_date: int | None = Field(
        default=None,
        description="Unix timestamp after which the created key expires",
    )


class CashuTransaction(SQLModel, table=True):  # type: ignore
    __tablename__ = "cashu_transactions"
    __table_args__ = (
        UniqueConstraint(
            "token", "type", name="uq_cashu_transactions_token_type"
        ),
    )

    id: str = Field(
        primary_key=True,
        default_factory=lambda: uuid.uuid4().hex,
        description="Unique transaction identifier",
    )
    token: str = Field(description="Serialized Cashu token")
    amount: int = Field(description="Amount in the token's unit")
    unit: str = Field(description="Token unit (sat or msat)")
    mint_url: str | None = Field(default=None, description="Mint URL for the token")
    type: str = Field(default="out", description="Transaction type: in or out")
    request_id: str | None = Field(default=None, description="Associated request ID")
    created_at: int = Field(
        default_factory=lambda: int(time.time()),
        description="Unix timestamp",
    )
    collected: bool = Field(default=False)
    swept: bool = Field(default=False)
    source: str = Field(
        default="x-cashu",
        description="Payment source: x-cashu or apikey",
    )
    api_key_hashed_key: str | None = Field(
        default=None,
        foreign_key="api_keys.hashed_key",
        index=True,
        description="Associated API key hash for wallet history",
    )


async def store_cashu_transaction(
    token: str,
    amount: int,
    unit: str,
    mint_url: str | None = None,
    typ: str = "out",
    request_id: str | None = None,
    collected: bool = False,
    created_at: int | None = None,
    source: str = "x-cashu",
    api_key_hashed_key: str | None = None,
) -> bool:
    """Persist a cashu transaction; idempotent on (token, type). Returns True on success."""
    try:
        async with create_session() as session:
            existing = await session.exec(
                select(CashuTransaction).where(
                    CashuTransaction.token == token,
                    CashuTransaction.type == typ,
                )
            )
            if existing.first() is not None:
                return True

            tx = CashuTransaction(
                token=token,
                amount=amount,
                unit=unit,
                mint_url=mint_url,
                type=typ,
                request_id=request_id,
                collected=collected,
                created_at=created_at or int(time.time()),
                source=source,
                api_key_hashed_key=api_key_hashed_key,
            )
            session.add(tx)
            await session.commit()
        return True
    except IntegrityError:
        # A concurrent insert of the same (token, type) won the race.
        async with create_session() as session:
            existing = await session.exec(
                select(CashuTransaction).where(
                    CashuTransaction.token == token,
                    CashuTransaction.type == typ,
                )
            )
            if existing.first() is not None:
                return True
        logger.error(
            f"Integrity error storing cashu transaction: non-duplicate violation (type={typ})",
            extra={
                "error": "non-duplicate IntegrityError",
                "type": typ,
                "request_id": request_id,
                "amount": amount,
                "unit": unit,
                "mint_url": mint_url,
            },
        )
        return False
    except Exception as e:
        logger.error(
            f"Failed to store cashu transaction: {e} (type={typ})",
            extra={
                "error": str(e),
                "type": typ,
                "request_id": request_id,
                "amount": amount,
                "unit": unit,
                "mint_url": mint_url,
                "token_preview": (token[:30] + "...") if len(token) > 30 else token,
            },
        )
        return False


async def store_cashu_transaction_with_retry(
    token: str,
    amount: int,
    unit: str,
    mint_url: str | None = None,
    typ: str = "out",
    request_id: str | None = None,
    collected: bool = False,
    created_at: int | None = None,
    source: str = "x-cashu",
    api_key_hashed_key: str | None = None,
    max_retries: int = 3,
) -> bool:
    """Retry ``store_cashu_transaction`` with backoff; spool to the outbox on exhaustion."""
    for attempt in range(max_retries):
        ok = await store_cashu_transaction(
            token=token,
            amount=amount,
            unit=unit,
            mint_url=mint_url,
            typ=typ,
            request_id=request_id,
            collected=collected,
            created_at=created_at,
            source=source,
            api_key_hashed_key=api_key_hashed_key,
        )
        if ok:
            return True
        if attempt < max_retries - 1:
            backoff = 0.5 * (2 ** attempt)
            logger.warning(
                "Retrying store_cashu_transaction",
                extra={
                    "attempt": attempt + 1,
                    "backoff_seconds": backoff,
                    "type": typ,
                    "request_id": request_id,
                },
            )
            await asyncio.sleep(backoff)

    logger.critical(
        "Cashu transaction spooled to outbox after retries exhausted",
        extra={
            "type": typ,
            "request_id": request_id,
            "amount": amount,
            "unit": unit,
            "mint_url": mint_url,
            "token": token,
        },
    )
    await append_to_cashu_outbox(
        token=token,
        amount=amount,
        unit=unit,
        mint_url=mint_url,
        typ=typ,
        request_id=request_id,
        collected=collected,
        created_at=created_at or int(time.time()),
        source=source,
        api_key_hashed_key=api_key_hashed_key,
    )
    return False


async def replay_cashu_outbox(path: pathlib.Path | None = None) -> int:
    """Replay spooled outbox entries into the DB. Returns the count persisted."""
    if path is None:
        path = _cashu_outbox_path()
    if not path.exists():
        return 0

    # Hold the lock for the full read-replay-rewrite cycle so a concurrent
    # appender cannot slip a new line between the read and the rewrite.
    async with _cashu_outbox_lock:
        try:
            raw_lines = path.read_text(encoding="utf-8").splitlines()
        except Exception as e:
            logger.error("Failed to read cashu outbox", extra={"error": str(e)})
            return 0

        entries: list[dict] = []
        for line in raw_lines:
            line = line.strip()
            if not line:
                continue
            try:
                entries.append(json.loads(line))
            except json.JSONDecodeError as e:
                logger.warning(
                    "Skipping malformed outbox line",
                    extra={"error": str(e), "line_preview": line[:80]},
                )

        if not entries:
            return 0

        persisted = 0
        remaining: list[dict] = []
        for entry in entries:
            ok = await store_cashu_transaction(
                token=entry["token"],
                amount=entry["amount"],
                unit=entry["unit"],
                mint_url=entry.get("mint_url"),
                typ=entry.get("type", "out"),
                request_id=entry.get("request_id"),
                collected=entry.get("collected", False),
                created_at=entry.get("created_at"),
                source=entry.get("source", "x-cashu"),
                api_key_hashed_key=entry.get("api_key_hashed_key"),
            )
            if ok:
                persisted += 1
                logger.info(
                    "Outbox entry replayed into DB",
                    extra={
                        "outbox_id": entry.get("outbox_id"),
                        "type": entry.get("type"),
                        "request_id": entry.get("request_id"),
                    },
                )
            else:
                remaining.append(entry)

        try:
            if remaining:
                tmp = path.with_suffix(path.suffix + ".tmp")
                with open(tmp, "w", encoding="utf-8") as fh:
                    for entry in remaining:
                        fh.write(json.dumps(entry, separators=(",", ":")) + "\n")
                    fh.flush()
                    os.fsync(fh.fileno())
                os.replace(tmp, path)
            else:
                tmp = path.with_suffix(path.suffix + ".tmp")
                tmp.write_text("")
                os.replace(tmp, path)
        except Exception as e:
            logger.error(
                "Failed to rewrite cashu outbox after replay",
                extra={"error": str(e), "remaining": len(remaining)},
            )

    return persisted


async def replay_all_outbox_files() -> int:
    """Replay every cashu outbox file in the outbox directory."""
    override = os.environ.get("CASHU_OUTBOX_PATH")
    if override:
        return await replay_cashu_outbox(pathlib.Path(override))

    outbox_dir = _cashu_outbox_path().parent

    total = 0
    for file in sorted(outbox_dir.glob("cashu_outbox.*.jsonl")):
        total += await replay_cashu_outbox(file)
    return total


async def periodic_cashu_outbox_replay() -> None:
    """Background loop that drains the cashu outbox into the database."""
    interval = float(os.environ.get("CASHU_OUTBOX_REPLAY_INTERVAL", "30"))
    logger.info("Starting cashu outbox replay loop", extra={"interval": interval})
    while True:
        try:
            await asyncio.sleep(interval)
            await replay_all_outbox_files()
        except asyncio.CancelledError:
            raise
        except Exception as e:
            logger.error(
                "Outbox replay iteration failed",
                extra={"error": str(e), "error_type": type(e).__name__},
            )


class UpstreamProviderRow(SQLModel, table=True):  # type: ignore
    __tablename__ = "upstream_providers"
    __table_args__ = (
        UniqueConstraint(
            "base_url", "api_key", name="uq_upstream_providers_base_url_api_key"
        ),
    )
    id: int | None = Field(default=None, primary_key=True)
    slug: str | None = Field(
        default=None,
        unique=True,
        index=True,
        description="Stable external slug used for updates via API key.",
    )
    provider_type: str = Field(
        description="Provider type: custom, openai, anthropic, azure, openrouter, etc."
    )
    base_url: str = Field(description="Base URL of the upstream API")
    api_key: str = Field(description="API key for the upstream provider")
    api_version: str | None = Field(
        default=None, description="API version for Azure OpenAI"
    )
    enabled: bool = Field(default=True, description="Whether this provider is enabled")
    provider_fee: float = Field(
        default=1.01, description="Provider fee multiplier (default 1%)"
    )
    provider_settings: str | None = Field(
        default=None, description="JSON string for provider-specific settings"
    )
    models: list["ModelRow"] = Relationship(
        back_populates="upstream_provider",
        sa_relationship_kwargs={"cascade": "all, delete-orphan"},
    )


class RoutstrFee(SQLModel, table=True):  # type: ignore
    __tablename__ = "routstr_fees"
    id: int = Field(default=1, primary_key=True)
    accumulated_msats: int = Field(default=0)
    total_paid_msats: int = Field(default=0)
    last_paid_at: int | None = Field(default=None)


class CliToken(SQLModel, table=True):  # type: ignore
    """Long-lived authorization token for CLI/agent use against admin endpoints."""

    __tablename__ = "cli_tokens"
    id: str = Field(
        primary_key=True, default_factory=lambda: uuid.uuid4().hex
    )
    token: str = Field(unique=True, index=True, description="Bearer token value")
    name: str = Field(description="Human-readable label for this token")
    created_at: int = Field(default_factory=lambda: int(time.time()))
    last_used_at: int | None = Field(default=None)
    expires_at: int | None = Field(
        default=None, description="Optional expiry unix timestamp; null = never expires"
    )


async def accumulate_routstr_fee(session: AsyncSession, amount_msats: int) -> None:
    stmt = (
        update(RoutstrFee)
        .where(col(RoutstrFee.id) == 1)
        .values(accumulated_msats=RoutstrFee.accumulated_msats + amount_msats)
    )
    result = await session.exec(stmt)  # type: ignore[call-overload]
    if result.rowcount == 0:
        session.add(RoutstrFee(id=1, accumulated_msats=amount_msats))
    await session.commit()


async def get_routstr_fee(session: AsyncSession) -> RoutstrFee:
    fee = await session.get(RoutstrFee, 1)
    if fee is None:
        fee = RoutstrFee(id=1, accumulated_msats=0, total_paid_msats=0)
        session.add(fee)
        await session.commit()
        await session.refresh(fee)
    return fee


async def reset_routstr_fee(session: AsyncSession, paid_msats: int) -> None:
    stmt = (
        update(RoutstrFee)
        .where(col(RoutstrFee.id) == 1)
        .values(
            accumulated_msats=RoutstrFee.accumulated_msats - paid_msats,
            total_paid_msats=RoutstrFee.total_paid_msats + paid_msats,
            last_paid_at=int(time.time()),
        )
    )
    await session.exec(stmt)  # type: ignore[call-overload]
    await session.commit()


async def balances_for_mint_and_unit(
    db_session: AsyncSession, mint_url: str, unit: str
) -> int:
    query = select(func.sum(ApiKey.balance)).where(
        ApiKey.refund_mint_url == mint_url, ApiKey.refund_currency == unit
    )
    result = await db_session.exec(query)
    return result.one() or 0


async def init_db() -> None:
    """Initializes the database and creates tables if they don't exist."""
    async with engine.begin() as conn:
        if DATABASE_URL.startswith("sqlite"):
            await conn.exec_driver_sql("PRAGMA journal_mode=WAL")
        await conn.run_sync(SQLModel.metadata.create_all)


async def get_session() -> AsyncGenerator[AsyncSession, None]:
    async with AsyncSession(engine, expire_on_commit=False) as session:
        yield session


@asynccontextmanager
async def create_session() -> AsyncGenerator[AsyncSession, None]:
    async with AsyncSession(engine, expire_on_commit=False) as session:
        yield session


def fix_cashu_migrations() -> None:
    """
    Fixes Cashu wallet migrations that are not idempotent.
    This specifically addresses the 'duplicate column name: public_keys' error
    in the keysets table of Cashu's internal SQLite databases.
    """
    project_root = pathlib.Path(__file__).resolve().parents[2]
    wallet_dir = project_root / ".wallet"

    if not wallet_dir.exists() or not wallet_dir.is_dir():
        return

    logger.info("Checking Cashu wallet databases for migration idempotency")

    for db_file in wallet_dir.glob("*.sqlite3"):
        try:
            conn = sqlite3.connect(db_file)
            cursor = conn.cursor()

            # Check if keysets table exists
            cursor.execute(
                "SELECT name FROM sqlite_master WHERE type='table' AND name='keysets'"
            )
            if not cursor.fetchone():
                conn.close()
                continue

            # Check if public_keys column exists
            cursor.execute("PRAGMA table_info(keysets)")
            columns = [info[1] for info in cursor.fetchall()]

            if "public_keys" not in columns:
                logger.info(f"Adding missing public_keys column to {db_file.name}")
                cursor.execute("ALTER TABLE keysets ADD COLUMN public_keys TEXT")
                conn.commit()

            conn.close()
        except Exception as e:
            logger.warning(f"Could not check/fix Cashu database {db_file}: {e}")


def _clear_alembic_version() -> None:
    """Clear the alembic_version table so stamp/upgrade can proceed."""
    sync_url = DATABASE_URL.replace("+aiosqlite", "")
    from sqlalchemy import create_engine, text

    eng = create_engine(sync_url)
    with eng.begin() as conn:
        conn.execute(text("DELETE FROM alembic_version"))
    eng.dispose()


def run_migrations() -> None:
    """Run Alembic migrations programmatically."""
    try:
        # Run Cashu migration fix first
        fix_cashu_migrations()

        # Get the path to the alembic.ini file
        project_root = pathlib.Path(__file__).resolve().parents[2]
        alembic_ini_path = project_root / "alembic.ini"

        if not alembic_ini_path.exists():
            raise FileNotFoundError(
                f"Alembic configuration file not found at {alembic_ini_path}"
            )

        # Create Alembic config object
        alembic_cfg = Config(str(alembic_ini_path))

        # Set the database URL in the config
        alembic_cfg.set_main_option("sqlalchemy.url", DATABASE_URL)

        try:
            command.upgrade(alembic_cfg, "head")
        except CommandError as e:
            if "Can't locate revision" in str(e):
                logger.warning(
                    "Database stamped with unknown revision (likely from another branch). "
                    "Re-stamping to current head.",
                    extra={"error": str(e)},
                )
                _clear_alembic_version()
                command.stamp(alembic_cfg, "head")
            else:
                raise
        except OperationalError as e:
            if "duplicate column name" in str(e).lower():
                logger.warning(
                    "Migration hit a column that already exists (likely added via "
                    "create_all on another branch). Stamping to current head.",
                    extra={"error": str(e)},
                )
                _clear_alembic_version()
                command.stamp(alembic_cfg, "head")
            else:
                raise

        logger.info("Database migrations completed successfully")

    except Exception as e:
        logger.error(
            "Database migration failed",
            extra={"error": str(e), "error_type": type(e).__name__},
        )
        raise