From b754c50e9287bb3c70e840880c7a09bf2cd27ee8 Mon Sep 17 00:00:00 2001
From: Luke Mino-Altherr <luke@comfy.org>
Date: Wed, 4 Mar 2026 16:28:10 -0800
Subject: [PATCH] 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>
---
 app/assets/services/metadata_extract.py |  4 +++
 app/database/db.py                      | 43 +++++++++++++++++++++++++
 2 files changed, 47 insertions(+)

diff --git a/app/assets/services/metadata_extract.py b/app/assets/services/metadata_extract.py
index 1cdd26a40..ab66ce9ac 100644
--- a/app/assets/services/metadata_extract.py
+++ b/app/assets/services/metadata_extract.py
@@ -14,6 +14,10 @@ import struct
 from dataclasses import dataclass
 from typing import Any
 
+from utils.mime_types import init_mime_types
+
+init_mime_types()
+
 # Supported safetensors extensions
 SAFETENSORS_EXTENSIONS = frozenset({".safetensors", ".sft"})
 
diff --git a/app/database/db.py b/app/database/db.py
index d397fb8c1..0aab09a49 100644
--- a/app/database/db.py
+++ b/app/database/db.py
@@ -17,6 +17,10 @@ try:
     from alembic.script import ScriptDirectory
     from sqlalchemy import create_engine, event
     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
 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():
     db_url = args.database_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_exists = os.path.exists(db_path)