diff --git a/api/auth/oauth_handlers.py b/api/auth/oauth_handlers.py index 4e58c5a9..b9c216cb 100644 --- a/api/auth/oauth_handlers.py +++ b/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, diff --git a/api/graph.py b/api/graph.py index d2611ab5..72f136a4 100644 --- a/api/graph.py +++ b/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 diff --git a/api/loaders/csv_loader.py b/api/loaders/csv_loader.py index 8beda5a5..54a057c4 100644 --- a/api/loaders/csv_loader.py +++ b/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: diff --git a/api/loaders/json_loader.py b/api/loaders/json_loader.py index 0b74ec60..3f7dd521 100644 --- a/api/loaders/json_loader.py +++ b/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" diff --git a/api/loaders/mysql_loader.py b/api/loaders/mysql_loader.py index 5a49193f..a174577c 100644 --- a/api/loaders/mysql_loader.py +++ b/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.") diff --git a/api/loaders/odata_loader.py b/api/loaders/odata_loader.py index 77125d4d..5558c878 100644 --- a/api/loaders/odata_loader.py +++ b/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" diff --git a/api/loaders/postgres_loader.py b/api/loaders/postgres_loader.py index 33dd6130..d3e60989 100644 --- a/api/loaders/postgres_loader.py +++ b/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.") diff --git a/api/routes/database.py b/api/routes/database.py index 4254d45f..9e8f61e2 100644 --- a/api/routes/database.py +++ b/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( diff --git a/api/routes/graphs.py b/api/routes/graphs.py index 4e59b7fd..bc670b55 100644 --- a/api/routes/graphs.py +++ b/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,7 +235,7 @@ 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") @@ -243,13 +243,13 @@ async def load_graph(request: Request, data: GraphData = None, file: UploadFile 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({ diff --git a/tests/test_mysql_loader.py b/tests/test_mysql_loader.py index 99bbaebf..c40522c8 100644 --- a/tests/test_mysql_loader.py +++ b/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 diff --git a/tests/test_postgres_loader.py b/tests/test_postgres_loader.py index 895220f5..9a4d5f3e 100644 --- a/tests/test_postgres_loader.py +++ b/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)