Enhance migration process by adding storage table creation and improving connection handling in Database and Model classes

src/Migration.java

@@ -1,11 +1,10 @@
-package src;
-
 import java.sql.Connection;
 import java.sql.Statement;
 
 public class Migration {
     public static void run(Connection conn) {
         try (Statement stmt = conn.createStatement()) {
+            System.out.println("Running migrations...");
             stmt.execute("""
                 CREATE TABLE IF NOT EXISTS users (
                     id INTEGER PRIMARY KEY AUTOINCREMENT,
@@ -52,6 +51,15 @@ public class Migration {
                 )
             """);
 
+            stmt.execute("""
+                CREATE TABLE IF NOT EXISTS storage (
+                    key TEXT PRIMARY KEY,
+                    type TEXT,
+                    data TEXT NOT NULL
+                )
+            """);
+            System.out.println("Migrations completed.");
+
         } catch (Exception e) {
             e.printStackTrace();
         }

src/squirrel/Database.java

@@ -1,21 +1,21 @@
-package src.squirrel;
+package squirrel;
 
 import java.sql.Connection;
 import java.sql.DriverManager;
 import java.sql.SQLException;
-import java.util.function.Consumer;
 
 public class Database {
     private static Connection conn;
-    public static Consumer<Connection> migrate;
 
-    public static Connection getConnection() throws SQLException {
+    public static Connection getConnection() {
-        if (conn == null || conn.isClosed()) {
+        try {
-            String url = "jdbc:sqlite:instance/test.db";
+            if (conn == null || conn.isClosed()) {
-            conn = DriverManager.getConnection(url);
+                String url = "jdbc:sqlite:instance/test.db";
-            if (migrate != null) {
+                conn = DriverManager.getConnection(url);
-                migrate.accept(conn);
             }
+        } catch (SQLException e) {
+            e.printStackTrace();
+            throw new RuntimeException("Failed to connect to database", e);
         }
         return conn;
     }

src/squirrel/Model.java

@@ -1,4 +1,4 @@
-package src.squirrel;
+package squirrel;
 
 import java.sql.*;
 import java.util.*;
@@ -10,18 +10,45 @@ public abstract class Model {
     protected Set<String> columns = new HashSet<>();
     protected Map<String, Object> attributes = new LinkedHashMap<>();
 
-    protected boolean rowMode = false;
+    protected boolean rowMode = true;
+
+    // Only one classmode instance per subclass
+    private static final Map<Class<? extends Model>, Model> classModeInstances = new HashMap<>();
 
     public Model() {
         if (conn == null) {
+            conn = Database.getConnection();
+        }
+        // Always rowMode by default
+        this.rowMode = true;
+    }
+
+    // Call once per subclass to create the classmode instance
+    public static <T extends Model> T initializeClassMode(Class<T> clazz) {
+        synchronized (classModeInstances) {
+            if (classModeInstances.containsKey(clazz)) {
+                throw new IllegalStateException("Classmode instance already initialized for " + clazz.getSimpleName());
+            }
             try {
-                conn = Database.getConnection();
+                T instance = clazz.getDeclaredConstructor().newInstance();
-            } catch (SQLException e) {
+                instance.rowMode = false;
-                throw new RuntimeException("Failed to get DB connection for model '" + getTableName() + "'", e);
+                classModeInstances.put(clazz, instance);
+                return instance;
+            } catch (Exception e) {
+                throw new RuntimeException("Failed to initialize classmode instance for " + clazz.getSimpleName(), e);
             }
         }
     }
 
+    // Get the classmode instance for a subclass
+    public static <T extends Model> T getClassMode(Class<T> clazz) {
+        Model instance = classModeInstances.get(clazz);
+        if (instance == null) {
+            throw new IllegalStateException("Classmode instance not initialized for " + clazz.getSimpleName());
+        }
+        return clazz.cast(instance);
+    }
+
     // -------------------------------
     // Query (class/global mode only)
     // -------------------------------

src/squirrel/ModelManager.java

@@ -1,8 +1,8 @@
-package src.squirrel;
+package squirrel;
 
 import java.util.*;
 
-import src.models.UserModel;
+import models.UserModel;
 
 public class ModelManager {
     private static final Map<Class<? extends Model>, Model> models = new HashMap<>();
@@ -13,7 +13,7 @@ public class ModelManager {
         );
         for (Class<? extends Model> clazz : modelClasses) {
             try {
-                Model instance = clazz.getDeclaredConstructor().newInstance();
+                Model instance = Model.initializeClassMode(clazz);
                 models.put(clazz, instance);
             } catch (Exception e) {
                 e.printStackTrace();