Skip to content
Snippets Groups Projects
Unverified Commit dfdce6a9 authored by shirsho-dutt's avatar shirsho-dutt Committed by GitHub
Browse files

:mag: fix: Creation of `custom_id` Index (#135)

* Fixed creation of index on startup issue

* updated log statement

* updated log message

* fixed test
parent f7a32545
No related branches found
No related tags found
No related merge requests found
...@@ -25,29 +25,10 @@ async def ensure_custom_id_index_on_embedding(): ...@@ -25,29 +25,10 @@ async def ensure_custom_id_index_on_embedding():
pool = await PSQLDatabase.get_pool() pool = await PSQLDatabase.get_pool()
async with pool.acquire() as conn: async with pool.acquire() as conn:
# Check if the index exists await conn.execute(f"""
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"""
CREATE INDEX IF NOT EXISTS {index_name} ON {table_name} ({column_name}); CREATE INDEX IF NOT EXISTS {index_name} ON {table_name} ({column_name});
""") """)
logger.debug(f"Created index '{index_name}' on '{table_name}({column_name})'") logger.debug(f"Checking if index '{index_name}' on '{table_name}({column_name}) exists, if not found then the index is created.'")
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
async def pg_health_check() -> bool: async def pg_health_check() -> bool:
......
...@@ -7,7 +7,7 @@ from contextlib import asynccontextmanager ...@@ -7,7 +7,7 @@ from contextlib import asynccontextmanager
from starlette.responses import JSONResponse 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 LogMiddleware, logger
from app.middleware import security_middleware from app.middleware import security_middleware
from app.routes import document_routes, pgvector_routes from app.routes import document_routes, pgvector_routes
...@@ -16,7 +16,7 @@ from app.services.database import PSQLDatabase, ensure_custom_id_index_on_embedd ...@@ -16,7 +16,7 @@ from app.services.database import PSQLDatabase, ensure_custom_id_index_on_embedd
@asynccontextmanager @asynccontextmanager
async def lifespan(app: FastAPI): async def lifespan(app: FastAPI):
# Startup logic goes here # Startup logic goes here
if VECTOR_DB_TYPE == "pgvector": if VECTOR_DB_TYPE == VectorDBType.PGVECTOR:
await PSQLDatabase.get_pool() # Initialize the pool await PSQLDatabase.get_pool() # Initialize the pool
await ensure_custom_id_index_on_embedding() await ensure_custom_id_index_on_embedding()
......
...@@ -34,9 +34,6 @@ def dummy_pool(monkeypatch): ...@@ -34,9 +34,6 @@ def dummy_pool(monkeypatch):
import asyncio import asyncio
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_ensure_custom_id_index_on_embedding(monkeypatch, dummy_pool): 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() result = await ensure_custom_id_index_on_embedding()
# If no exceptions are raised, the function worked as expected. # If no exceptions are raised, the function worked as expected.
assert result is None assert result is None
\ No newline at end of file
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment