🩺 fix: Concurrent Upload Isolation, MongoDB Connection Leaks, and Validation Crashes (#256)

* Fix race condition in embed_file_upload temp file path

Concurrent uploads with the same filename would write to the same temp
path because embed_file_upload used the filename directly under
RAG_UPLOAD_DIR, unlike embed_file which namespaces under user_id. Two
users uploading "report.pdf" simultaneously would corrupt each other's
data. Add user_id subdirectory to match embed_file's behavior.

* Fix connection leak and event loop blocking in mongo_health_check

Every health check call created a MongoClient, opened a connection pool,
and never closed it. Since health checks run frequently, this leaked
connections continuously. Also, MongoClient() and admin.command() are
blocking I/O calls that were running directly in an async function,
blocking the event loop.

Fix both: wrap blocking calls in asyncio.to_thread() and always close
the client in a finally block.

* Fix validation_exception_handler crash and body echo

request.body() re-reads the entire upload into memory inside the error
handler, and body.decode() raises UnicodeDecodeError on binary/non-ASCII
content, turning a 422 into a 500. The response also echoed potentially
large or sensitive upload content back to the caller. Remove the body
read entirely — the validation errors already contain all useful info.

* Remove dead exception branch in generate_digest

str.encode("utf-8") on a Python str never raises UnicodeEncodeError for
any valid Unicode string. The except branch handled surrogate characters
with an encode-decode-encode round-trip that produces the same result as
encode("utf-8", "ignore") directly. Replace with a single expression.

* Enable EMBEDDING_BATCH_SIZE by default (500)

The batch processing pipeline is fully built but disabled by default
(EMBEDDING_BATCH_SIZE=0). This means large documents embed all chunks in
one blocking call. Set default to 500 to enable streaming embeddings
into the DB while the rest are still being computed — the highest-
leverage memory improvement with zero new code.

* Stream CSV encoding conversion instead of reading entire file

Non-UTF-8 CSV files were read entirely into a single Python string
before writing to the temp file. For large files (e.g. 100MB UTF-16),
this causes unnecessary peak memory allocation. Stream in 64KB chunks
instead.

* Close asyncpg connection pool on shutdown

The pool was created during startup but never closed on shutdown.
Rolling deploys or OOM-killed pods leave connections in idle state until
TCP timeout, consuming PostgreSQL's finite connection limit. Close the
pool cleanly in the lifespan teardown.

* Address review findings: unique temp paths, guarded shutdown, tests

Finding 1: Wrap close_pool() in try/except so an exception doesn't
skip thread_pool.shutdown().

Finding 2+5: Use UUID-suffixed temp filenames via _make_unique_temp_path
in all three upload endpoints (embed, embed-upload, text). This prevents
same-user concurrent upload collisions (not just cross-user). Also moves
save_upload_file_async inside the try/finally block so temp files are
cleaned up even when the save itself fails.

Finding 3+6: Restore tuning guidance comment for EMBEDDING_BATCH_SIZE
explaining the default of 500 and when to increase.

Finding 4: Add tests for _make_unique_temp_path (uniqueness, user
isolation, extension preservation, traversal rejection) and
generate_digest (normal strings, unicode, surrogates, determinism).

* Clean up remaining nits from review

- Drop redundant second validate_file_path call in _make_unique_temp_path;
  the UUID suffix is [0-9a-f]{32} with no path separators, so it cannot
  escape the already-validated directory. Resolve the path directly.
- Fold os.makedirs into the try/finally block in all three upload
  endpoints so an OSError (permissions, disk full) is caught by the
  existing error handler instead of propagating as an unhandled 500.
- Fix hardcoded Unix separator in test assertion; use Path.parent.name
  instead of string containment check.

Author: danny-avila

app/config.py

@@ -77,10 +77,11 @@ def get_env_variable(
 # Trade-offs:
 # - Smaller batch size = lower memory, more DB round trips
 # - Larger batch size = higher memory, fewer DB round trips
-# - 0 = disable batching, process all at once (original behavior)
+# - 0 = disable batching, process all at once
 #
-# Recommended: 750 for text-embedding-3-small (good balance of speed and memory)
-EMBEDDING_BATCH_SIZE = int(get_env_variable("EMBEDDING_BATCH_SIZE", "0"))
+# Default of 500 is conservative and works well for most embedding providers.
+# Increase to 750 for higher throughput at the cost of higher peak memory.
+EMBEDDING_BATCH_SIZE = int(get_env_variable("EMBEDDING_BATCH_SIZE", "500"))
 
 # Maximum number of batches to buffer in memory during async processing.
 # Higher values allow more parallelism but use more memory.

app/routes/document_routes.py

@@ -1,5 +1,6 @@
 # app/routes/document_routes.py
 import os
+import uuid
 from pathlib import Path
 import hashlib
 import traceback
@@ -123,6 +124,20 @@ def validate_file_path(base_dir: str, file_path: str) -> Optional[str]:
         return None
 
 
+def _make_unique_temp_path(user_id: str, filename: str) -> Optional[str]:
+    """Build a unique temp file path under RAG_UPLOAD_DIR/{user_id}/ to prevent
+    concurrent upload collisions. Returns a validated absolute path, or None if
+    the raw filename would escape RAG_UPLOAD_DIR (path traversal rejection)."""
+    # Validate the raw filename to reject traversal attempts
+    if validate_file_path(RAG_UPLOAD_DIR, os.path.join(user_id, filename)) is None:
+        return None
+    # unique_name is stem + "_" + [0-9a-f]{32} + suffix — no path separators,
+    # so it cannot escape the directory validated above.
+    p = Path(filename)
+    unique_name = f"{p.stem}_{uuid.uuid4().hex}{p.suffix}"
+    return str(Path(RAG_UPLOAD_DIR, user_id, unique_name).resolve())
+
+
 async def load_file_content(
     filename: str, content_type: str, file_path: str, executor
 ) -> tuple:
@@ -612,14 +627,8 @@ async def _process_documents_batched_sync(
     return all_ids
 
 
-def generate_digest(page_content: str):
-    try:
-        hash_obj = hashlib.md5(page_content.encode("utf-8"))
-    except UnicodeEncodeError:
-        hash_obj = hashlib.md5(
-            page_content.encode("utf-8", "ignore").decode("utf-8").encode("utf-8")
-        )
-    return hash_obj.hexdigest()
+def generate_digest(page_content: str) -> str:
+    return hashlib.md5(page_content.encode("utf-8", "ignore")).hexdigest()
 
 
 def _prepare_documents_sync(
@@ -793,9 +802,7 @@ async def embed_file(
     known_type = None
 
     user_id = get_user_id(request, entity_id)
-    validated_file_path = validate_file_path(
-        RAG_UPLOAD_DIR, os.path.join(user_id, file.filename)
-    )
+    validated_file_path = _make_unique_temp_path(user_id, file.filename)
 
     if validated_file_path is None:
         logger.warning("Path validation failed for embed: %s", file.filename)
@@ -804,11 +811,9 @@ async def embed_file(
             detail=ERROR_MESSAGES.DEFAULT("Invalid request"),
         )
 
-    os.makedirs(os.path.dirname(validated_file_path), exist_ok=True)
-
-    await save_upload_file_async(file, validated_file_path)
-
     try:
+        os.makedirs(os.path.dirname(validated_file_path), exist_ok=True)
+        await save_upload_file_async(file, validated_file_path)
         data, known_type, file_ext = await load_file_content(
             file.filename,
             file.content_type,
@@ -931,9 +936,7 @@ async def embed_file_upload(
 ):
     user_id = get_user_id(request, entity_id)
 
-    validated_temp_file_path = validate_file_path(
-        RAG_UPLOAD_DIR, uploaded_file.filename
-    )
+    validated_temp_file_path = _make_unique_temp_path(user_id, uploaded_file.filename)
 
     if validated_temp_file_path is None:
         logger.warning(
@@ -944,9 +947,9 @@ async def embed_file_upload(
             detail=ERROR_MESSAGES.DEFAULT("Invalid request"),
         )
 
-    await save_upload_file_async(uploaded_file, validated_temp_file_path)
-
     try:
+        os.makedirs(os.path.dirname(validated_temp_file_path), exist_ok=True)
+        await save_upload_file_async(uploaded_file, validated_temp_file_path)
         data, known_type, file_ext = await load_file_content(
             uploaded_file.filename,
             uploaded_file.content_type,
@@ -1053,9 +1056,7 @@ async def extract_text_from_file(
     Returns the raw text content for text parsing purposes.
     """
     user_id = get_user_id(request, entity_id)
-    validated_temp_file_path = validate_file_path(
-        RAG_UPLOAD_DIR, os.path.join(user_id, file.filename)
-    )
+    validated_temp_file_path = _make_unique_temp_path(user_id, file.filename)
 
     if validated_temp_file_path is None:
         logger.warning("Path validation failed for text extraction: %s", file.filename)
@@ -1064,12 +1065,9 @@ async def extract_text_from_file(
             detail=ERROR_MESSAGES.DEFAULT("Invalid request"),
         )
 
-    # create base directory only if the file path is valid
-    os.makedirs(os.path.dirname(validated_temp_file_path), exist_ok=True)
-
-    await save_upload_file_async(file, validated_temp_file_path)
-
     try:
+        os.makedirs(os.path.dirname(validated_temp_file_path), exist_ok=True)
+        await save_upload_file_async(file, validated_temp_file_path)
         data, known_type, file_ext = await load_file_content(
             file.filename,
             file.content_type,

app/services/mongo_client.py

@@ -1,16 +1,25 @@
 # app/services/mongo_client.py
+import asyncio
 import logging
 from pymongo import MongoClient
 from pymongo.errors import PyMongoError
 from app.config import ATLAS_MONGO_DB_URI
 
 logger = logging.getLogger(__name__)
 
+
 async def mongo_health_check() -> bool:
+    client = None
     try:
-        client = MongoClient(ATLAS_MONGO_DB_URI)
-        client.admin.command("ping")
+        client = await asyncio.to_thread(MongoClient, ATLAS_MONGO_DB_URI)
+        await asyncio.to_thread(client.admin.command, "ping")
         return True
     except PyMongoError as e:
-        logger.error(f"MongoDB health check failed: {e}")
-        return False
+        logger.error("MongoDB health check failed: %s", e)
+        return False
+    finally:
+        if client is not None:
+            try:
+                await asyncio.to_thread(client.close)
+            except Exception as e:
+                logger.debug("Failed to close health check client: %s", e)

app/utils/document_loader.py

@@ -82,29 +82,27 @@ def get_loader(filename: str, file_content_type: str, filepath: str):
         encoding = detect_file_encoding(filepath)
 
         if encoding != "utf-8":
-            # For non-UTF-8 encodings, we need to convert the file first
-            # Create a temporary UTF-8 file
+            # For non-UTF-8 encodings, convert to UTF-8 using streaming
+            # to avoid holding the entire file in memory as a single string
             temp_file = None
             try:
                 with tempfile.NamedTemporaryFile(
                     mode="w", encoding="utf-8", suffix=".csv", delete=False
                 ) as temp_file:
-                    # Read the original file with detected encoding
                     with open(
                         filepath, "r", encoding=encoding, errors="replace"
                     ) as original_file:
-                        content = original_file.read()
-                        temp_file.write(content)
+                        while True:
+                            chunk = original_file.read(64 * 1024)
+                            if not chunk:
+                                break
+                            temp_file.write(chunk)
 
                     temp_filepath = temp_file.name
 
-                # Use the temporary UTF-8 file with CSVLoader
                 loader = CSVLoader(temp_filepath)
-
-                # Store the temp file path for cleanup
                 loader._temp_filepath = temp_filepath
             except Exception as e:
-                # If temp file was created but there was an error, clean it up
                 if temp_file and os.path.exists(temp_file.name):
                     os.unlink(temp_file.name)
                 raise e

docker-compose.yaml

@@ -15,6 +15,7 @@ services:
     environment:
       - DB_HOST=db
       - DB_PORT=5432
+      - EMBEDDING_BATCH_SIZE=${EMBEDDING_BATCH_SIZE:-500}
     ports:
       - "8000:8000"
     volumes:

main.py

@@ -47,6 +47,14 @@ async def lifespan(app: FastAPI):
     yield
 
     # Cleanup logic
+    if VECTOR_DB_TYPE == VectorDBType.PGVECTOR:
+        try:
+            logger.info("Closing asyncpg connection pool")
+            await PSQLDatabase.close_pool()
+            logger.info("asyncpg connection pool closed")
+        except Exception as e:
+            logger.warning("Failed to close asyncpg pool: %s", e)
+
     logger.info("Shutting down thread pool")
     app.state.thread_pool.shutdown(wait=True)
     logger.info("Thread pool shutdown complete")
@@ -79,17 +87,10 @@ async def lifespan(app: FastAPI):
 
 @app.exception_handler(RequestValidationError)
 async def validation_exception_handler(request: Request, exc: RequestValidationError):
-    body = await request.body()
-    logger.debug(f"Validation error occurred")
-    logger.debug(f"Raw request body: {body.decode()}")
-    logger.debug(f"Validation errors: {exc.errors()}")
+    logger.debug("Validation error: %s", exc.errors())
     return JSONResponse(
         status_code=422,
-        content={
-            "detail": exc.errors(),
-            "body": body.decode(),
-            "message": "Request validation failed",
-        },
+        content={"detail": exc.errors(), "message": "Request validation failed"},
     )
 
 

tests/test_upload_isolation.py

@@ -0,0 +1,88 @@
+"""
+Tests for upload temp file isolation and generate_digest correctness.
+
+Validates:
+- _make_unique_temp_path produces unique paths per call (no concurrent collisions)
+- _make_unique_temp_path isolates users into separate subdirectories
+- _make_unique_temp_path rejects path traversal filenames
+- generate_digest is consistent for all string inputs including surrogates
+"""
+
+import hashlib
+import os
+from pathlib import Path
+import pytest
+
+from app.routes.document_routes import _make_unique_temp_path, generate_digest
+
+
+class TestMakeUniqueTempPath:
+    """Ensure temp file paths are unique and user-isolated."""
+
+    def test_two_calls_same_filename_produce_different_paths(
+        self, monkeypatch, tmp_path
+    ):
+        monkeypatch.setattr("app.routes.document_routes.RAG_UPLOAD_DIR", str(tmp_path))
+        path_a = _make_unique_temp_path("user1", "report.pdf")
+        path_b = _make_unique_temp_path("user1", "report.pdf")
+        assert path_a != path_b, "Same user+filename must produce unique paths"
+
+    def test_different_users_produce_different_directories(self, monkeypatch, tmp_path):
+        monkeypatch.setattr("app.routes.document_routes.RAG_UPLOAD_DIR", str(tmp_path))
+        path_a = _make_unique_temp_path("user1", "report.pdf")
+        path_b = _make_unique_temp_path("user2", "report.pdf")
+        assert os.path.dirname(path_a) != os.path.dirname(path_b)
+        assert Path(path_a).parent.name == "user1"
+        assert Path(path_b).parent.name == "user2"
+
+    def test_preserves_file_extension(self, monkeypatch, tmp_path):
+        monkeypatch.setattr("app.routes.document_routes.RAG_UPLOAD_DIR", str(tmp_path))
+        path = _make_unique_temp_path("user1", "data.csv")
+        assert path.endswith(".csv")
+
+    def test_path_stays_within_upload_dir(self, monkeypatch, tmp_path):
+        monkeypatch.setattr("app.routes.document_routes.RAG_UPLOAD_DIR", str(tmp_path))
+        path = _make_unique_temp_path("user1", "file.txt")
+        assert path.startswith(str(tmp_path))
+
+    @pytest.mark.parametrize(
+        "malicious_filename",
+        [
+            "../../etc/passwd",
+            "../../../etc/shadow",
+            "/etc/passwd",
+        ],
+    )
+    def test_rejects_path_traversal(self, monkeypatch, tmp_path, malicious_filename):
+        monkeypatch.setattr("app.routes.document_routes.RAG_UPLOAD_DIR", str(tmp_path))
+        result = _make_unique_temp_path("user1", malicious_filename)
+        assert result is None
+
+
+class TestGenerateDigest:
+    """Ensure generate_digest is correct for all inputs."""
+
+    def test_normal_string(self):
+        content = "hello world"
+        expected = hashlib.md5(content.encode("utf-8")).hexdigest()
+        assert generate_digest(content) == expected
+
+    def test_empty_string(self):
+        expected = hashlib.md5(b"").hexdigest()
+        assert generate_digest("") == expected
+
+    def test_unicode_content(self):
+        content = "café résumé naïve"
+        expected = hashlib.md5(content.encode("utf-8")).hexdigest()
+        assert generate_digest(content) == expected
+
+    def test_surrogate_characters(self):
+        """Surrogate chars are stripped by encode('utf-8', 'ignore')."""
+        content = "hello\ud800world"
+        expected = hashlib.md5(content.encode("utf-8", "ignore")).hexdigest()
+        assert generate_digest(content) == expected
+        assert len(generate_digest(content)) == 32
+
+    def test_deterministic(self):
+        content = "same input"
+        assert generate_digest(content) == generate_digest(content)