Skip to content

fix(migrate): close SQLite connection and clean temp palace on exception#1216

Open
arnoldwender wants to merge 1 commit intoMemPalace:developfrom
arnoldwender:fix/migrate-resource-cleanup
Open

fix(migrate): close SQLite connection and clean temp palace on exception#1216
arnoldwender wants to merge 1 commit intoMemPalace:developfrom
arnoldwender:fix/migrate-resource-cleanup

Conversation

@arnoldwender
Copy link
Copy Markdown
Contributor

What and Why

Two resource leaks in mempalace.migrate that bite hardest on Windows:

1. SQLite connection leak in extract_drawers_from_sqlite()

mempalace/migrate.py:55 opened a connection at the top of the function and only closed it on the happy path. Any exception during query/iteration leaked the handle. On Windows that holds a file lock on chroma.sqlite3 — the rest of the migration cannot shutil.copytree or os.replace the palace directory until the process exits. Users see a confusing PermissionError after a partial extraction.

2. Orphan temp dir in migrate()

mempalace/migrate.py:232 (now ~234) called tempfile.mkdtemp(prefix="mempalace_migrate_") with no try/finally. If anything between mkdtemp and the final os.replace(temp_palace, palace_path) raises (chromadb open failure, batch import failure, count mismatch, EXDEV fallback failure), the partial palace directory under the system temp root leaks forever. For a 50K-drawer palace this can be hundreds of MB the user has to find and clean by hand.

Fix

  1. extract_drawers_from_sqlite — wrap the connection in contextlib.closing(sqlite3.connect(db_path)) so close() runs on every exit path. detect_chromadb_version already had a try/finally and is left alone.
  2. migrate — wrap the temp-palace import-and-swap dance in try/finally. The cleanup checks os.path.exists(temp_palace) first; on the happy path os.replace consumes the temp dir so the guard makes cleanup a no-op.

Test plan

  • test_extract_drawers_returns_drawers — happy path still works after the closing refactor (real SQLite fixture with the schema extract_drawers_from_sqlite reads)
  • test_migrate_cleans_temp_palace_on_chromadb_failuremkdtemp's directory is removed when ChromaBackend.get_or_create_collection raises, asserted via a tracking_mkdtemp that captures the path
  • All 7 test_migrate.py tests pass
  • Full suite: 1318 passed locally (2 pre-existing test_backends.py::test_pin_hnsw_threads* failures unrelated — chromadb configuration_json["hnsw"] schema drift, also fails on clean develop)

Two resource leaks in mempalace.migrate that bite hardest on Windows:

1. extract_drawers_from_sqlite() opened a sqlite3 connection at the
   top and only closed it in the happy path. Any exception during
   query/iteration leaked the handle, leaving a file lock on
   chroma.sqlite3 that prevented the rest of the migration from
   touching the palace directory. Wrap with contextlib.closing().

2. migrate() called tempfile.mkdtemp() and never cleaned it up if a
   later step (chromadb open, batch import, count, swap) raised.
   The orphaned palace under the system temp root could be 100s of
   MB and never gets reclaimed. Wrap the import-and-swap dance in
   try/finally and rmtree the temp dir if it still exists at the
   end (os.replace consumes it on the happy path so the existence
   guard makes that case a no-op).

Tests cover the happy-path extraction and verify mkdtemp's directory
is removed when ChromaBackend.get_or_create_collection() raises.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant