fix the db not exist bug (#131)

Author: TechNomad-ds

dataflow/operators/filter/Text2SQL/sql_execution_filter.py

@@ -62,18 +62,34 @@ def run(self, storage: DataFlowStorage,
         self.check_column(dataframe)
         results = []
         self.logger.info(f"Start to filter {len(dataframe)} SQLs")
+        db_id_need_to_check = dataframe[input_db_id_key].unique()
+        for db_id in db_id_need_to_check:
+            if not self.database_manager.registry.database_exists(db_id):
+                self.logger.warning(f"Database {db_id} not found in registry, please check the database folder")
+                continue
+
         for _, row in tqdm(dataframe.iterrows(), total=len(dataframe), desc="Processing SQLs"):
             db_id = row[input_db_id_key]
             sql = row[input_sql_key]
 
             if not self.filter_select_sql(sql):
                 continue
-
-            ans = self.database_manager.analyze_sql_execution_plan(db_id, sql, 5)
-            if not ans['success']:
+            
+            try:
+                ans = self.database_manager.analyze_sql_execution_plan(db_id, sql, 5)
+            except Exception as e:
+                self.logger.error(f"Error analyzing SQL execution plan: {e}")
                 continue
 
-            ans = self.database_manager.execute_query(db_id, sql, 10)
+            if not ans['success']:
+                continue
+
+            try:
+                ans = self.database_manager.execute_query(db_id, sql, 10)
+            except Exception as e:
+                self.logger.error(f"Error executing SQL query: {e}")
+                continue
+                
             if not ans['success']:
                 continue
 

dataflow/utils/text2sql/database_manager.py

@@ -6,6 +6,7 @@
 import os
 import glob
 from dataclasses import dataclass
+from dataflow import get_logger
 
 
 @dataclass
@@ -267,7 +268,7 @@ class DatabaseManager:
     def __init__(self, db_type: str = "sqlite", config: Optional[Dict] = None, logger=None):
         self.db_type = db_type.lower()
         self.config = config or {}
-        self.logger = logger
+        self.logger = get_logger()
 
         self.registry = DatabaseRegistry(logger)
         self.cache = SchemaCache()
@@ -427,7 +428,11 @@ def execute_query(self, db_id: str, sql: str, timeout: float = 5.0) -> Dict[str,
     def analyze_sql_execution_plan(self, db_id: str, sql: str, timeout: float = 5.0) -> Dict[str, Any]:
         db_info = self.registry.get_database(db_id)
         if not db_info:
-            raise ValueError(f"Database {db_id} not found")
+            self.logger.error(f"Database {db_id} not found")
+            return {
+                'success': False,
+                'error': f'Database {db_id} not found'
+            }
 
         connector = self.connectors[db_info.db_type]
         result = None