fix(pgvector): use open=False to prevent ConnectionPool hang in Docker (#5155)

Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>

mem0/vector_stores/pgvector.py

@@ -181,6 +181,7 @@ class PGVector(VectorStoreBase):
         self.use_hnsw = hnsw
         self.embedding_model_dims = embedding_model_dims
         self.connection_pool = None
+        self._collection_ensured = False
 
         # Connection setup with priority: connection_pool > connection_string > individual parameters
         if connection_pool is not None:
@@ -196,15 +197,25 @@ class PGVector(VectorStoreBase):
         
         if self.connection_pool is None:
             if PSYCOPG_VERSION == 3:
-                # psycopg3 ConnectionPool
-                self.connection_pool = ConnectionPool(conninfo=connection_string, min_size=minconn, max_size=maxconn, open=True)
+                # open=False avoids blocking when DB DNS is not yet resolvable (e.g. Docker startup)
+                self.connection_pool = ConnectionPool(
+                    conninfo=connection_string,
+                    min_size=minconn,
+                    max_size=maxconn,
+                    open=False,
+                )
+                self.connection_pool.open(wait=False)
             else:
                 # psycopg2 ThreadedConnectionPool
                 self.connection_pool = ConnectionPool(minconn=minconn, maxconn=maxconn, dsn=connection_string)
 
+    def _ensure_collection(self):
+        if self._collection_ensured:
+            return
         collections = self.list_cols()
-        if collection_name not in collections:
+        if self.collection_name not in collections:
             self.create_col()
+        self._collection_ensured = True
 
     @contextmanager
     def _get_cursor(self, commit: bool = False):
@@ -294,6 +305,7 @@ class PGVector(VectorStoreBase):
             )
 
     def insert(self, vectors: list[list[float]], payloads=None, ids=None) -> None:
+        self._ensure_collection()
         logger.info(f"Inserting {len(vectors)} vectors into collection {self.collection_name}")
         json_payloads = [json.dumps(payload) for payload in payloads]
 
@@ -331,6 +343,7 @@ class PGVector(VectorStoreBase):
         Returns:
             list: Search results.
         """
+        self._ensure_collection()
         filter_conditions, filter_params = _build_filter_conditions(filters)
         filter_clause = sql.SQL("WHERE " + " AND ".join(filter_conditions)) if filter_conditions else sql.SQL("")
 
@@ -361,6 +374,7 @@ class PGVector(VectorStoreBase):
         Returns:
             List[OutputData]: Search results ranked by text relevance.
         """
+        self._ensure_collection()
         filter_conditions, filter_params = _build_filter_conditions(filters)
         filter_clause = sql.SQL("AND " + " AND ".join(filter_conditions)) if filter_conditions else sql.SQL("")
 
@@ -391,6 +405,7 @@ class PGVector(VectorStoreBase):
         Args:
             vector_id (str): ID of the vector to delete.
         """
+        self._ensure_collection()
         with self._get_cursor(commit=True) as cur:
             cur.execute(sql.SQL("DELETE FROM {} WHERE id = %s").format(self._col()), (vector_id,))
 
@@ -408,6 +423,7 @@ class PGVector(VectorStoreBase):
             vector (List[float], optional): Updated vector.
             payload (Dict, optional): Updated payload.
         """
+        self._ensure_collection()
         with self._get_cursor(commit=True) as cur:
             if vector:
                cur.execute(
@@ -440,6 +456,7 @@ class PGVector(VectorStoreBase):
         Returns:
             OutputData: Retrieved vector.
         """
+        self._ensure_collection()
         with self._get_cursor() as cur:
             cur.execute(
                 sql.SQL("SELECT id, vector, payload FROM {} WHERE id = %s").format(self._col()),
@@ -473,6 +490,7 @@ class PGVector(VectorStoreBase):
         Returns:
             Dict[str, Any]: Collection information.
         """
+        self._ensure_collection()
         with self._get_cursor() as cur:
             cur.execute(
                 sql.SQL("""
@@ -503,6 +521,7 @@ class PGVector(VectorStoreBase):
         Returns:
             List[OutputData]: List of vectors.
         """
+        self._ensure_collection()
         filter_conditions, filter_params = _build_filter_conditions(filters)
         filter_clause = sql.SQL("WHERE " + " AND ".join(filter_conditions)) if filter_conditions else sql.SQL("")
 
@@ -534,6 +553,7 @@ class PGVector(VectorStoreBase):
 
     def reset(self) -> None:
         """Reset the index by deleting and recreating it."""
+        self._ensure_collection()
         logger.warning(f"Resetting index {self.collection_name}...")
         self.delete_col()
         self.create_col()