Update src/db_connector.py

src/db_connector.py

@@ -8,24 +8,30 @@ class Database:
     def __init__(self):
         load_dotenv()
         self.db_uri = os.getenv("DB_URI")
+        self.local_db_path = "./chatbot.db"
         
+        # 1. Automatic Fallback to Internal DB
         if not self.db_uri:
-            print("⚠️ WARNING: DB_URI not found. Defaulting to 'sqlite:///./demo.db'")
-            self.db_uri = "sqlite:///./demo.db"
-
-        self.parsed = urlparse(self.db_uri)
-        
-        if "sqlite" in self.parsed.scheme:
+            print("⚠️ DB_URI not found. Using internal 'chatbot.db'...")
+            self.db_uri = f"sqlite:///{self.local_db_path}"
             self.type = "sqlite"
-            self.db_path = self.parsed.path if self.parsed.path else "./demo.db"
-            if self.db_path.startswith("/."): self.db_path = self.db_path[1:]
+            self.db_path = self.local_db_path
+            
+            # 🚀 AUTO-INIT: Create the data if it doesn't exist
+            self._init_sqlite_data()
         else:
-            self.type = "mysql"
-            self.host = self.parsed.hostname
-            self.port = self.parsed.port or 3306
-            self.user = self.parsed.username
-            self.password = unquote(self.parsed.password)
-            self.db_name = self.parsed.path[1:]
+            # Parse external DB_URI (if you eventually add a cloud DB)
+            parsed = urlparse(self.db_uri)
+            if "sqlite" in parsed.scheme:
+                self.type = "sqlite"
+                self.db_path = parsed.path
+            else:
+                self.type = "mysql"
+                self.host = parsed.hostname
+                self.port = parsed.port or 3306
+                self.user = parsed.username
+                self.password = unquote(parsed.password)
+                self.db_name = parsed.path[1:]
 
     def get_connection(self):
         if self.type == "sqlite":
@@ -34,14 +40,49 @@ class Database:
             return conn
         else:
             return pymysql.connect(
-                host=self.host,
-                user=self.user,
-                password=self.password,
-                database=self.db_name,
-                port=self.port,
+                host=self.host, user=self.user, password=self.password,
+                database=self.db_name, port=self.port,
                 cursorclass=pymysql.cursors.DictCursor
             )
 
+    def _init_sqlite_data(self):
+        """Creates the User's specific schema inside the Space."""
+        if os.path.exists(self.local_db_path):
+            return # Data already exists
+
+        print("⚡ Initializing internal database with User Data...")
+        conn = sqlite3.connect(self.local_db_path)
+        cursor = conn.cursor()
+        
+        # Your SQL (Adapted for SQLite)
+        commands = [
+            """CREATE TABLE IF NOT EXISTS employees (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                name TEXT,
+                department TEXT,
+                salary REAL,
+                hire_date TEXT
+            );""",
+            """CREATE TABLE IF NOT EXISTS sales (
+                sale_id INTEGER PRIMARY KEY AUTOINCREMENT,
+                employee_id INTEGER,
+                amount REAL,
+                sale_date TEXT,
+                FOREIGN KEY (employee_id) REFERENCES employees(id)
+            );""",
+            "INSERT INTO employees (name, department, salary, hire_date) VALUES ('Alice', 'Sales', 70000, '2023-01-15');",
+            "INSERT INTO employees (name, department, salary, hire_date) VALUES ('Bob', 'Engineering', 90000, '2022-05-20');",
+            "INSERT INTO sales (employee_id, amount, sale_date) VALUES (1, 500.00, '2023-06-01');",
+            "INSERT INTO sales (employee_id, amount, sale_date) VALUES (1, 1200.50, '2023-06-03');"
+        ]
+        
+        for cmd in commands:
+            cursor.execute(cmd)
+        
+        conn.commit()
+        conn.close()
+        print("✅ Internal Database Created Successfully!")
+
     def run_query(self, query):
         conn = self.get_connection()
         try:
@@ -52,8 +93,7 @@ class Database:
             else:
                 cursor = conn.cursor()
                 cursor.execute(query)
-                items = [dict(row) for row in cursor.fetchall()]
-                return items
+                return [dict(row) for row in cursor.fetchall()]
         except Exception as e:
             return [f"Error: {e}"]
         finally:
@@ -65,13 +105,13 @@ class Database:
             cursor = conn.cursor()
             if self.type == "sqlite":
                 cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
-                tables = [row[0] for row in cursor.fetchall()]
+                # Filter out system tables (sqlite_sequence, etc.)
+                tables = [row[0] for row in cursor.fetchall() if row[0] not in ['sqlite_sequence']]
             else:
                 cursor.execute("SHOW TABLES")
                 tables = [list(row.values())[0] for row in cursor.fetchall()]
             
-            # 🔍 DEBUG: PRINT TABLES TO LOGS
-            print(f"\n🔍 FOUND TABLES IN DATABASE: {tables}\n")
+            print(f"\n🔍 FOUND TABLES: {tables}\n")
             return tables
         except Exception as e:
             print(f"Error fetching tables: {e}")