Upload 62 files

app/ai_assistant.py

@@ -451,6 +451,17 @@ class AIAssistant:
 
     def _create_sample_database(self, user_store, database: str) -> Dict:
         """Create multiple sample tables with data in a database"""
+        # Create database if it doesn't exist
+        from app.metadata import metadata
+        dbs = metadata.list_databases(user_store)
+        db_list = dbs if isinstance(dbs, list) else dbs.get('databases', [])
+        db_names_list = [d['name'] if isinstance(d, dict) else d for d in db_list]
+
+        if database not in db_names_list:
+            result = metadata.create_database(database, user_store)
+            if not result.get('ok'):
+                return {'type': 'error', 'message': f'Failed to create database: {result.get("error")}'}
+
         tables_to_create = [
             ('customers', [
                 (1, 'John Doe', 'john@example.com', '555-0101', '123 Main St'),
@@ -524,6 +535,17 @@ class AIAssistant:
     def _create_sample_table(self, user_store, database: str, table: str) -> Dict:
         """Create a sample table with realistic data based on table name"""
 
+        # Create database if it doesn't exist
+        from app.metadata import metadata
+        dbs = metadata.list_databases(user_store)
+        db_list = dbs if isinstance(dbs, list) else dbs.get('databases', [])
+        db_names_list = [d['name'] if isinstance(d, dict) else d for d in db_list]
+
+        if database not in db_names_list:
+            result = metadata.create_database(database, user_store)
+            if not result.get('ok'):
+                return {'type': 'error', 'message': f'Failed to create database: {result.get("error")}'}
+
         # Detect context from table name
         table_lower = table.lower()
 

app/query_parser.py

@@ -126,10 +126,23 @@ class QueryParser:
     
     def extract_tables(self, sql: str) -> List[str]:
         """Extract table names from SQL query"""
-        # Simple regex to find table names after FROM and JOIN
-        pattern = r'(?:FROM|JOIN)\s+["\']?(\w+\.\w+)["\']?'
-        matches = re.findall(pattern, sql, re.IGNORECASE)
-        return [self.sanitize_identifier(m) for m in matches]
+        tables = set()
+        # FROM and JOIN clauses
+        for m in re.findall(r'(?:FROM|JOIN)\s+["\']?(\w+\.\w+)["\']?', sql, re.IGNORECASE):
+            tables.add(m)
+        # CREATE TABLE
+        for m in re.findall(r'CREATE\s+(?:OR\s+REPLACE\s+)?TABLE\s+(?:IF\s+NOT\s+EXISTS\s+)?["\']?(\w+\.\w+)["\']?', sql, re.IGNORECASE):
+            tables.add(m)
+        # INSERT INTO
+        for m in re.findall(r'INSERT\s+INTO\s+["\']?(\w+\.\w+)["\']?', sql, re.IGNORECASE):
+            tables.add(m)
+        # UPDATE
+        for m in re.findall(r'UPDATE\s+["\']?(\w+\.\w+)["\']?', sql, re.IGNORECASE):
+            tables.add(m)
+        # DELETE FROM
+        for m in re.findall(r'DELETE\s+FROM\s+["\']?(\w+\.\w+)["\']?', sql, re.IGNORECASE):
+            tables.add(m)
+        return [self.sanitize_identifier(m) for m in tables]
     
     def build_safe_query(self, database: str, table: str, filters: Dict = None, 
                         limit: int = None, offset: int = None) -> str: