fix-async-queries

api/auth/oauth_handlers.py

@@ -34,7 +34,7 @@ async def handle_google_callback(_request: Request,
                 return False
 
             # Check if identity exists in Organizations graph, create if new
-            _, _ = ensure_user_in_organizations(
+            _, _ = await ensure_user_in_organizations(
                 user_id,
                 email,
                 name,
@@ -62,7 +62,7 @@ async def handle_github_callback(_request: Request,
                 return False
 
             # Check if identity exists in Organizations graph, create if new
-            _, _ = ensure_user_in_organizations(
+            _, _ = await ensure_user_in_organizations(
                 user_id,
                 email,
                 name,

api/graph.py

@@ -287,6 +287,9 @@ async def find(
     json_data = json.loads(completion_result.choices[0].message.content)
     descriptions = Descriptions(**json_data)
     descriptions_text = [desc.description for desc in descriptions.tables_descriptions] + [desc.description for desc in descriptions.columns_descriptions]
+    if not descriptions_text:
+        return []
+
     embedding_results = Config.EMBEDDING_MODEL.embed(descriptions_text)
 
     # Split embeddings back into table and column embeddings

api/loaders/csv_loader.py

@@ -14,7 +14,7 @@ class CSVLoader(BaseLoader):
     """CSV data loader for processing CSV files and loading them into graph database."""
 
     @staticmethod
-    def load(graph_id: str, data) -> Tuple[bool, str]:
+    async def load(graph_id: str, data) -> Tuple[bool, str]:
         """
         Load the data dictionary CSV file into the graph database.
 
@@ -192,12 +192,12 @@ def load(graph_id: str, data) -> Tuple[bool, str]:
                                             "source_column": key,
                                             "target_column": key,
                                             "note": "many-many",
-                                        }
-                                    )
+                                    }
+                                )
 
-            load_to_graph(graph_id, tables, relationships, db_name=db_name)
+            await load_to_graph(graph_id, tables, relationships, db_name=db_name)
             return True, "Data dictionary loaded successfully into graph"
-
+        
         except Exception as e:
             return False, f"Error loading CSV: {str(e)}"
             # else:

api/loaders/json_loader.py

@@ -24,7 +24,7 @@ class JSONLoader(BaseLoader):
     """JSON schema loader for loading database schemas from JSON files."""
 
     @staticmethod
-    def load(graph_id: str, data) -> Tuple[bool, str]:
+    async def load(graph_id: str, data) -> Tuple[bool, str]:
         """
         Load the graph data into the database.
         It gets the Graph name as an argument and expects
@@ -66,6 +66,6 @@ def load(graph_id: str, data) -> Tuple[bool, str]:
                         "note": fk_name,
                     }
                 )
-        load_to_graph(graph_id, data["tables"], relationships, db_name=data["database"])
+        await load_to_graph(graph_id, data["tables"], relationships, db_name=data["database"])
 
         return True, "Graph loaded successfully"

api/loaders/mysql_loader.py

@@ -125,7 +125,7 @@ def _parse_mysql_url(connection_url: str) -> Dict[str, str]:
         }
 
     @staticmethod
-    def load(prefix: str, connection_url: str) -> Tuple[bool, str]:
+    async def load(prefix: str, connection_url: str) -> Tuple[bool, str]:
         """
         Load the graph data from a MySQL database into the graph database.
 
@@ -158,7 +158,7 @@ def load(prefix: str, connection_url: str) -> Tuple[bool, str]:
             conn.close()
 
             # Load data into graph
-            load_to_graph(prefix + "_" + db_name, entities, relationships,
+            await load_to_graph(prefix + "_" + db_name, entities, relationships,
                          db_name=db_name, db_url=connection_url)
 
             return True, (f"MySQL schema loaded successfully. "
@@ -436,7 +436,7 @@ async def refresh_graph_schema(graph_id: str, db_url: str) -> Tuple[bool, str]:
                 prefix = graph_id
 
             # Reuse the existing load method to reload the schema
-            success, message = MySQLLoader.load(prefix, db_url)
+            success, message = await MySQLLoader.load(prefix, db_url)
 
             if success:
                 logging.info("Graph schema refreshed successfully.")

api/loaders/odata_loader.py

@@ -14,7 +14,7 @@ class ODataLoader(BaseLoader):
     """
 
     @staticmethod
-    def load(graph_id: str, data) -> Tuple[bool, str]:
+    async def load(graph_id: str, data) -> Tuple[bool, str]:
         """Load XML ODATA schema into a Graph."""
 
         try:
@@ -23,7 +23,7 @@ def load(graph_id: str, data) -> Tuple[bool, str]:
         except ET.ParseError:
             return False, "Invalid XML content"
 
-        load_to_graph(graph_id, entities, relationships, db_name="ERP system")
+        await load_to_graph(graph_id, entities, relationships, db_name="ERP system")
 
         return True, "Graph loaded successfully"
 

api/loaders/postgres_loader.py

@@ -64,7 +64,7 @@ def _serialize_value(value):
             return value
 
     @staticmethod
-    def load(prefix: str, connection_url: str) -> Tuple[bool, str]:
+    async def load(prefix: str, connection_url: str) -> Tuple[bool, str]:
         """
         Load the graph data from a PostgreSQL database into the graph database.
 
@@ -96,7 +96,7 @@ def load(prefix: str, connection_url: str) -> Tuple[bool, str]:
             conn.close()
 
             # Load data into graph
-            load_to_graph(prefix + "_" + db_name, entities, relationships,
+            await load_to_graph(prefix + "_" + db_name, entities, relationships,
                          db_name=db_name, db_url=connection_url)
 
             return True, (f"PostgreSQL schema loaded successfully. "
@@ -397,7 +397,7 @@ async def refresh_graph_schema(graph_id: str, db_url: str) -> Tuple[bool, str]:
                 prefix = graph_id
 
             # Reuse the existing load method to reload the schema
-            success, message = PostgresLoader.load(prefix, db_url)
+            success, message = await PostgresLoader.load(prefix, db_url)
 
             if success:
                 logging.info("Graph schema refreshed successfully.")

api/routes/database.py

@@ -45,7 +45,7 @@ async def connect_database(request: Request, db_request: DatabaseConnectionReque
         if url.startswith("postgres://") or url.startswith("postgresql://"):
             try:
                 # Attempt to connect/load using the PostgreSQL loader
-                success, result = PostgresLoader.load(request.state.user_id, url)
+                success, result = await PostgresLoader.load(request.state.user_id, url)
             except (ValueError, ConnectionError) as e:
                 logging.error("PostgreSQL connection error: %s", str(e))
                 raise HTTPException(
@@ -57,7 +57,7 @@ async def connect_database(request: Request, db_request: DatabaseConnectionReque
         elif url.startswith("mysql://"):
             try:
                 # Attempt to connect/load using the MySQL loader
-                success, result = MySQLLoader.load(request.state.user_id, url)
+                success, result = await MySQLLoader.load(request.state.user_id, url)
             except (ValueError, ConnectionError) as e:
                 logging.error("MySQL connection error: %s", str(e))
                 raise HTTPException(

api/routes/graphs.py

@@ -141,8 +141,8 @@ async def get_graph_data(request: Request, graph_id: str):
     """
 
     try:
-        tables_res = await graph.query(tables_query).result_set
-        links_res = await graph.query(links_query).result_set
+        tables_res = (await graph.query(tables_query)).result_set
+        links_res = (await graph.query(links_query)).result_set
     except Exception as e:
         logging.error("Error querying graph data for %s: %s", sanitize_log_input(namespaced), e)
         return JSONResponse(content={"error": "Failed to read graph data"}, status_code=500)
@@ -223,7 +223,7 @@ async def load_graph(request: Request, data: GraphData = None, file: UploadFile
             raise HTTPException(status_code=400, detail="Invalid JSON data")
 
         graph_id = request.state.user_id + "_" + data.database
-        success, result = JSONLoader.load(graph_id, data.dict())
+        success, result = await JSONLoader.load(graph_id, data.dict())
 
     # ✅ Handle File Upload
     elif file:
@@ -235,21 +235,21 @@ async def load_graph(request: Request, data: GraphData = None, file: UploadFile
             try:
                 data = json.loads(content.decode("utf-8"))
                 graph_id = request.state.user_id + "_" + data.get("database", "")
-                success, result = JSONLoader.load(graph_id, data)
+                success, result = await JSONLoader.load(graph_id, data)
             except json.JSONDecodeError:
                 raise HTTPException(status_code=400, detail="Invalid JSON file")
 
         # ✅ Check if file is XML
         elif filename.endswith(".xml"):
             xml_data = content.decode("utf-8")
             graph_id = request.state.user_id + "_" + filename.replace(".xml", "")
-            success, result = ODataLoader.load(graph_id, xml_data)
+            success, result = await ODataLoader.load(graph_id, xml_data)
 
         # ✅ Check if file is csv
         elif filename.endswith(".csv"):
             csv_data = content.decode("utf-8")
             graph_id = request.state.user_id + "_" + filename.replace(".csv", "")
-            success, result = CSVLoader.load(graph_id, csv_data)
+            success, result = await CSVLoader.load(graph_id, csv_data)
 
         else:
             raise HTTPException(status_code=415, detail="Unsupported file type")
@@ -457,7 +457,7 @@ async def generate():
                                          "refreshing graph...")}
                         yield json.dumps(step) + MESSAGE_DELIMITER
 
-                        refresh_result = loader_class.refresh_graph_schema(
+                        refresh_result = await loader_class.refresh_graph_schema(
                             graph_id, db_url)
                         refresh_success, refresh_message = refresh_result
 
@@ -556,7 +556,7 @@ async def confirm_destructive_operation(
     async def generate_confirmation():
         if confirmation == "CONFIRM":
             try:
-                db_description, db_url = get_db_description(graph_id)
+                db_description, db_url = await get_db_description(graph_id)
 
                 # Determine database type and get appropriate loader
                 db_type, loader_class = get_database_type_and_loader(db_url)
@@ -591,7 +591,7 @@ async def generate_confirmation():
                     yield json.dumps(step) + MESSAGE_DELIMITER
 
                     refresh_success, refresh_message = (
-                        loader_class.refresh_graph_schema(graph_id, db_url)
+                        await loader_class.refresh_graph_schema(graph_id, db_url)
                     )
 
                     if refresh_success:
@@ -664,7 +664,7 @@ async def refresh_graph_schema(request: Request, graph_id: str):
 
     try:
         # Get database connection details
-        _, db_url = get_db_description(graph_id)
+        _, db_url = await get_db_description(graph_id)
 
         if not db_url or db_url == "No URL available for this database.":
             return JSONResponse({
@@ -682,7 +682,7 @@ async def refresh_graph_schema(request: Request, graph_id: str):
             }, status_code=400)
 
         # Perform schema refresh using the appropriate loader
-        success, message = loader_class.refresh_graph_schema(graph_id, db_url)
+        success, message = await loader_class.refresh_graph_schema(graph_id, db_url)
 
         if success:
             return JSONResponse({

tests/test_mysql_loader.py

@@ -1,5 +1,6 @@
 """Tests for MySQL loader functionality."""
 
+import asyncio
 import datetime
 import decimal
 from unittest.mock import patch, MagicMock
@@ -110,7 +111,7 @@ def test_connection_error(self, mock_connect):
         # Mock connection failure
         mock_connect.side_effect = Exception("Connection failed")
 
-        success, message = MySQLLoader.load("test_prefix", "mysql://user:pass@host:3306/db")
+        success, message = asyncio.run(MySQLLoader.load("test_prefix", "mysql://user:pass@host:3306/db"))
 
         assert success is False
         assert "Error loading MySQL schema" in message
@@ -132,9 +133,9 @@ def test_successful_load(self, mock_load_to_graph, mock_connect):
         with patch.object(MySQLLoader, 'extract_tables_info',
                           return_value={'users': {'description': 'User table'}}):
             with patch.object(MySQLLoader, 'extract_relationships', return_value={}):
-                success, message = MySQLLoader.load(
+                success, message = asyncio.run(MySQLLoader.load(
                     "test_prefix", "mysql://user:pass@localhost:3306/testdb"
-                )
+                ))
 
         assert success is True
         assert "MySQL schema loaded successfully" in message

tests/test_postgres_loader.py

@@ -5,6 +5,7 @@
 This script provides basic tests for the PostgreSQL loader functionality.
 """
 
+import asyncio
 import unittest
 from unittest.mock import Mock, patch
 
@@ -40,7 +41,7 @@ def test_successful_load(self, mock_load_to_graph, mock_connect):
         mock_load_to_graph.return_value = None
 
         # Test the loader
-        success, message = PostgresLoader.load(self.test_graph_id, self.test_connection_url)
+        success, message = asyncio.run(PostgresLoader.load(self.test_graph_id, self.test_connection_url))
 
         # Assertions
         self.assertTrue(success)
@@ -55,7 +56,7 @@ def test_connection_error(self, mock_connect):
         mock_connect.side_effect = Exception("Connection failed")
 
         # Test the loader
-        success, message = PostgresLoader.load(self.test_graph_id, self.test_connection_url)
+        success, message = asyncio.run(PostgresLoader.load(self.test_graph_id, self.test_connection_url))
 
         # Assertions
         self.assertFalse(success)