🔍 fix: Creation of custom_id Index (#135)

* Fixed creation of index on startup issue

* updated log statement

* updated log message

* fixed test

Author: shirsho-dutt

app/services/database.py

@@ -25,29 +25,10 @@ async def ensure_custom_id_index_on_embedding():
 
     pool = await PSQLDatabase.get_pool()
     async with pool.acquire() as conn:
-        # Check if the index exists
-        index_exists = await check_index_exists(conn, index_name)
-
-        if not index_exists:
-            # If the index does not exist, create it
-            await conn.execute(f"""
+        await conn.execute(f"""
                 CREATE INDEX IF NOT EXISTS {index_name} ON {table_name} ({column_name});
             """)
-            logger.debug(f"Created index '{index_name}' on '{table_name}({column_name})'")
-        else:
-            logger.debug(f"Index '{index_name}' already exists on '{table_name}({column_name})'")
-
-
-async def check_index_exists(conn, index_name: str) -> bool:
-    # Adjust the SQL query if necessary
-    result = await conn.fetchval("""
-        SELECT EXISTS (
-            SELECT FROM pg_class c
-            JOIN pg_namespace n ON n.oid = c.relnamespace
-            WHERE  c.relname = $1 AND n.nspname = 'public' -- Adjust schema if necessary
-        );
-    """, index_name)
-    return result
+        logger.debug(f"Checking if index '{index_name}' on '{table_name}({column_name}) exists, if not found then the index is created.'")
 
 
 async def pg_health_check() -> bool:

main.py

@@ -7,7 +7,7 @@
 
 from starlette.responses import JSONResponse
 
-from app.config import debug_mode, RAG_HOST, RAG_PORT, CHUNK_SIZE, CHUNK_OVERLAP, PDF_EXTRACT_IMAGES, VECTOR_DB_TYPE, \
+from app.config import VectorDBType, debug_mode, RAG_HOST, RAG_PORT, CHUNK_SIZE, CHUNK_OVERLAP, PDF_EXTRACT_IMAGES, VECTOR_DB_TYPE, \
     LogMiddleware, logger
 from app.middleware import security_middleware
 from app.routes import document_routes, pgvector_routes
@@ -16,7 +16,7 @@
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     # Startup logic goes here
-    if VECTOR_DB_TYPE == "pgvector":
+    if VECTOR_DB_TYPE == VectorDBType.PGVECTOR:
         await PSQLDatabase.get_pool()  # Initialize the pool
         await ensure_custom_id_index_on_embedding()
 

tests/services/test_database.py

@@ -34,9 +34,6 @@ def dummy_pool(monkeypatch):
 import asyncio
 @pytest.mark.asyncio
 async def test_ensure_custom_id_index_on_embedding(monkeypatch, dummy_pool):
-    async def dummy_check_index_exists(conn, index_name: str) -> bool:
-        return False
-    monkeypatch.setattr("app.services.database.check_index_exists", dummy_check_index_exists)
     result = await ensure_custom_id_index_on_embedding()
     # If no exceptions are raised, the function worked as expected.
     assert result is None