Added better connection handling on the database

Author: nullchimp

examples/8-url-loader.py

@@ -16,6 +16,7 @@
     username=os.environ.get("MEMGRAPH_USERNAME", "memgraph"),
     password=os.environ.get("MEMGRAPH_PASSWORD", "memgraph"),
 )
+db.connect()
 
 loader = WebLoader("http://localhost:4000/en/enterprise-cloud@latest/copilot")
 embedder = TextEmbedding3Small()
@@ -39,6 +40,8 @@ async def main():
         vectors = []
         await embedder.process_chunks(chunks, callback=lambda v: vectors.append(v))
         store(source, doc, chunks, vectors)
+
+    db.close()
 
 
 if __name__ == "__main__":

src/core/rag/dbhandler/__init__.py

@@ -43,12 +43,6 @@ def __init__(
         self._conn = None
         self._cur = None
 
-        try:
-            self.connect(**kwargs)
-        except Exception as e:
-            print(f"Connection error: {str(e)}")
-            raise ConnectionError(f"Failed to connect to Memgraph at {host}:{port}: {str(e)}") from e
-
     def connect(self, *args, **kwargs) -> None:
         raise ConnectionError("Not implemented")
 
@@ -57,7 +51,13 @@ def close(self) -> None:
         self._conn.close()
 
     # Context-manager support
-    def __enter__(self) -> GraphClient:
+    def __enter__(self, *args, **kwargs) -> GraphClient:
+        try:
+            self.connect(**kwargs)
+        except Exception as e:
+            print(f"Connection error: {str(e)}")
+            raise ConnectionError(f"Failed to connect to Memgraph at {self.host}:{self.port}: {str(e)}") from e
+
         return self
 
     def __exit__(self, exc_type, exc, tb) -> None: