From 1855d2000e06bc1e015df102cb439e3862270158 Mon Sep 17 00:00:00 2001
From: Dominik Krenn <dominik.krenn@1tool.com>
Date: Mon, 15 Sep 2025 15:35:02 +0200
Subject: [PATCH] Enhance migration process by adding storage table creation
 and improving connection handling in Database and Model classes

---
 src/Migration.java             | 12 +++++++++--
 src/squirrel/Database.java     | 18 ++++++++---------
 src/squirrel/Model.java        | 37 +++++++++++++++++++++++++++++-----
 src/squirrel/ModelManager.java |  6 +++---
 4 files changed, 54 insertions(+), 19 deletions(-)

diff --git a/src/Migration.java b/src/Migration.java
index 6d39736..2feaca0 100644
--- a/src/Migration.java
+++ b/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();
         }
diff --git a/src/squirrel/Database.java b/src/squirrel/Database.java
index 5b13f66..1955d8b 100644
--- a/src/squirrel/Database.java
+++ b/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 {
-        if (conn == null || conn.isClosed()) {
-            String url = "jdbc:sqlite:instance/test.db";
-            conn = DriverManager.getConnection(url);
-            if (migrate != null) {
-                migrate.accept(conn);
+    public static Connection getConnection() {
+        try {
+            if (conn == null || conn.isClosed()) {
+                String url = "jdbc:sqlite:instance/test.db";
+                conn = DriverManager.getConnection(url);
             }
+        } catch (SQLException e) {
+            e.printStackTrace();
+            throw new RuntimeException("Failed to connect to database", e);
         }
         return conn;
     }
diff --git a/src/squirrel/Model.java b/src/squirrel/Model.java
index 81f291e..04879f1 100644
--- a/src/squirrel/Model.java
+++ b/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();
-            } catch (SQLException e) {
-                throw new RuntimeException("Failed to get DB connection for model '" + getTableName() + "'", e);
+                T instance = clazz.getDeclaredConstructor().newInstance();
+                instance.rowMode = false;
+                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)
     // -------------------------------
diff --git a/src/squirrel/ModelManager.java b/src/squirrel/ModelManager.java
index 05b2a76..d5f6248 100644
--- a/src/squirrel/ModelManager.java
+++ b/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();