fix: support in-memory SQLite and ensure custom MIME types are initialized

- Add _init_memory_db() path using Base.metadata.create_all + StaticPool
  since Alembic migrations don't work with in-memory SQLite (each
  connection gets its own separate database)
- Call init_mime_types() at module load in metadata_extract so custom
  types like application/safetensors are always registered

Amp-Thread-ID: https://ampcode.com/threads/T-019cbb5f-13d1-7429-8cfd-815625c4d032
Co-authored-by: Amp <amp@ampcode.com>
This commit is contained in:
Luke Mino-Altherr 2026-03-04 16:28:10 -08:00
parent 62a514a19e
commit 95afa8b69b
2 changed files with 47 additions and 0 deletions

View File

@ -14,6 +14,10 @@ import struct
from dataclasses import dataclass from dataclasses import dataclass
from typing import Any from typing import Any
from utils.mime_types import init_mime_types
init_mime_types()
# Supported safetensors extensions # Supported safetensors extensions
SAFETENSORS_EXTENSIONS = frozenset({".safetensors", ".sft"}) SAFETENSORS_EXTENSIONS = frozenset({".safetensors", ".sft"})

View File

@ -17,6 +17,10 @@ try:
from alembic.script import ScriptDirectory from alembic.script import ScriptDirectory
from sqlalchemy import create_engine, event from sqlalchemy import create_engine, event
from sqlalchemy.orm import sessionmaker from sqlalchemy.orm import sessionmaker
from sqlalchemy.pool import StaticPool
from app.database.models import Base
import app.assets.database.models # noqa: F401 — register models with Base.metadata
_DB_AVAILABLE = True _DB_AVAILABLE = True
except ImportError as e: except ImportError as e:
@ -87,9 +91,48 @@ def _acquire_file_lock(db_path):
) )
def _is_memory_db(db_url):
"""Check if the database URL refers to an in-memory SQLite database."""
return db_url in ("sqlite:///:memory:", "sqlite://")
def init_db(): def init_db():
db_url = args.database_url db_url = args.database_url
logging.debug(f"Database URL: {db_url}") logging.debug(f"Database URL: {db_url}")
if _is_memory_db(db_url):
_init_memory_db(db_url)
else:
_init_file_db(db_url)
def _init_memory_db(db_url):
"""Initialize an in-memory SQLite database using metadata.create_all.
Alembic migrations don't work with in-memory SQLite because each
connection gets its own separate database tables created by Alembic's
internal connection are lost immediately.
"""
engine = create_engine(
db_url,
poolclass=StaticPool,
connect_args={"check_same_thread": False},
)
@event.listens_for(engine, "connect")
def set_sqlite_pragma(dbapi_connection, connection_record):
cursor = dbapi_connection.cursor()
cursor.execute("PRAGMA foreign_keys=ON")
cursor.close()
Base.metadata.create_all(engine)
global Session
Session = sessionmaker(bind=engine)
def _init_file_db(db_url):
"""Initialize a file-backed SQLite database using Alembic migrations."""
db_path = get_db_path() db_path = get_db_path()
db_exists = os.path.exists(db_path) db_exists = os.path.exists(db_path)