Lancedb async lock (#1222)

<!-- .github/pull_request_template.md -->

## Description
<!-- Provide a clear description of the changes in this PR -->

## DCO Affirmation
I affirm that all code in every commit of this pull request conforms to
the terms of the Topoteretes Developer Certificate of Origin.

---------

Co-authored-by: Vasilije <8619304+Vasilije1990@users.noreply.github.com>

cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py

@@ -53,6 +53,7 @@ class SQLAlchemyAdapter:
             self.engine = create_async_engine(
                 connection_string,
                 poolclass=NullPool,
+                connect_args={"timeout": 30},
             )
         else:
             self.engine = create_async_engine(

cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py

@@ -51,6 +51,7 @@ class LanceDBAdapter(VectorDBInterface):
         self.url = url
         self.api_key = api_key
         self.embedding_engine = embedding_engine
+        self.VECTOR_DB_LOCK = asyncio.Lock()
 
     async def get_connection(self):
         """
@@ -126,6 +127,8 @@ class LanceDBAdapter(VectorDBInterface):
             vector: Vector(vector_size)
             payload: payload_schema
 
+        if not await self.has_collection(collection_name):
+            async with self.VECTOR_DB_LOCK:
                 if not await self.has_collection(collection_name):
                     connection = await self.get_connection()
                     return await connection.create_table(
@@ -144,6 +147,8 @@ class LanceDBAdapter(VectorDBInterface):
     async def create_data_points(self, collection_name: str, data_points: list[DataPoint]):
         payload_schema = type(data_points[0])
 
+        if not await self.has_collection(collection_name):
+            async with self.VECTOR_DB_LOCK:
                 if not await self.has_collection(collection_name):
                     await self.create_collection(
                         collection_name,
@@ -188,6 +193,7 @@ class LanceDBAdapter(VectorDBInterface):
             for (data_point_index, data_point) in enumerate(data_points)
         ]
 
+        async with self.VECTOR_DB_LOCK:
             await (
                 collection.merge_insert("id")
                 .when_matched_update_all()

cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py

@@ -54,6 +54,7 @@ class PGVectorAdapter(SQLAlchemyAdapter, VectorDBInterface):
         self.api_key = api_key
         self.embedding_engine = embedding_engine
         self.db_uri: str = connection_string
+        self.VECTOR_DB_LOCK = asyncio.Lock()
 
         relational_db = get_relational_engine()
 
@@ -124,6 +125,7 @@ class PGVectorAdapter(SQLAlchemyAdapter, VectorDBInterface):
         data_point_types = get_type_hints(DataPoint)
         vector_size = self.embedding_engine.get_vector_size()
 
+        async with self.VECTOR_DB_LOCK:
             if not await self.has_collection(collection_name):
 
                 class PGVectorDataPoint(Base):

cognee/infrastructure/databases/vector/pinecone/adapter.py

@@ -1,8 +0,0 @@
-# class PineconeVectorDB(VectorDB):
-#     def __init__(self, *args, **kwargs):
-#         super().__init__(*args, **kwargs)
-#         self.init_pinecone(self.index_name)
-#
-#     def init_pinecone(self, index_name):
-#         # Pinecone initialization logic
-#         pass

cognee/infrastructure/databases/vector/qdrant/QDrantAdapter.py

@@ -132,7 +132,6 @@ class QDrantAdapter(VectorDBInterface):
 
     def __init__(self, url, api_key, embedding_engine: EmbeddingEngine, qdrant_path=None):
         self.embedding_engine = embedding_engine
-
         if qdrant_path is not None:
             self.qdrant_path = qdrant_path
         else:

cognee/tasks/storage/index_data_points.py

@@ -1,5 +1,3 @@
-import asyncio
-
 from cognee.shared.logging_utils import get_logger
 
 from cognee.infrastructure.databases.exceptions.EmbeddingException import EmbeddingException
@@ -8,9 +6,6 @@ from cognee.infrastructure.engine import DataPoint
 
 logger = get_logger("index_data_points")
 
-# A single lock shared by all coroutines
-vector_index_lock = asyncio.Lock()
-
 
 async def index_data_points(data_points: list[DataPoint]):
     created_indexes = {}
@@ -27,8 +22,6 @@ async def index_data_points(data_points: list[DataPoint]):
 
             index_name = f"{data_point_type.__name__}_{field_name}"
 
-            # Add async lock to make sure two different coroutines won't create a table at the same time
-            async with vector_index_lock:
             if index_name not in created_indexes:
                 await vector_engine.create_vector_index(data_point_type.__name__, field_name)
                 created_indexes[index_name] = True