Allow composite keys in primary keys

src/main/java/me/cobeine/sqlava/connection/database/table/Table.java

@@ -4,9 +4,7 @@
 import me.cobeine.sqlava.connection.database.table.column.Column;
 import org.jetbrains.annotations.NotNull;
 
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.List;
+import java.util.*;
 
 /**
  * @author <a href="https://github.com/Cobeine">Cobeine</a>
@@ -16,9 +14,9 @@
 
     @Getter private final String name;
     private final List<Column> columns;
-    private String primaryKey;
+    private List<String> primaryKeys = new ArrayList<>();
     private final List<ForeignKey> foreignKeys;
-    private final HashMap<String,String[]> uniqueKeys;
+    private final HashMap<String, String[]> uniqueKeys;
 
     public Table(String name) {
         this.name = name;
@@ -36,6 +34,7 @@
         foreignKeys.add(entry);
         return entry;
     }
+
     public void uniqueKey(String key, String... columns) {
         uniqueKeys.put(key, columns);
     }
@@ -46,53 +45,78 @@
         }
     }
 
-    public void setPrimaryKey(@NotNull Column primaryKey) {
-        this.primaryKey = primaryKey.getName();
-    }
-    public void setPrimaryKey(@NotNull String primaryKey) {
-        this.primaryKey = primaryKey;
+    public void setPrimaryKey(String... keys) {
+        primaryKeys = Arrays.asList(keys);
     }
 
     @Override
     public String toString() {
-        StringBuilder builder = new StringBuilder();
-        builder.append("CREATE TABLE IF NOT EXISTS `").append(name).append("` (");
+        StringBuilder builder = new StringBuilder("CREATE TABLE IF NOT EXISTS `" + name + "` (");
+
+        this.addColumns(builder);
+        this.addPrimaryKeys(builder);
+        this.addForeignKeys(builder);
+        this.addUniqueKeys(builder);
+
+        builder.append(")");
+        return builder.toString();
+    }
+
+    private void addColumns(StringBuilder builder) {
         for (Column column : columns) {
             builder.append(column.toString()).append(", ");
         }
-        if (primaryKey != null) {
-            builder.append("PRIMARY KEY (`").append(primaryKey).append("`)");
-        } else {
+    }
+
+    private void addPrimaryKeys(StringBuilder builder) {
+        if (primaryKeys.isEmpty()) {
             builder.deleteCharAt(builder.length() - 1);
             builder.deleteCharAt(builder.length() - 1); //to remove the last ", "
+            return;
         }
-        if (!foreignKeys.isEmpty()) {
-            for (ForeignKey entry : foreignKeys) {
-                if (entry.referencedColumn != null) {
-                    builder.append(", FOREIGN KEY (`").append(entry.foreignKey).append("`)");
-                    builder.append(" REFERENCES `").append(entry.referencedTable).append("`(`").append(entry.referencedColumn).append("`)");
-                    if (entry.onDelete != null) {
-                        builder.append(" ON DELETE ").append(entry.onDelete.name().replace("_"," "));
-                    }
-                }
-            }
-        }
-        if (!uniqueKeys.isEmpty()) {
-            for (String key : uniqueKeys.keySet()) {
-                String[] columns = uniqueKeys.get(key);
-                builder.append(", UNIQUE KEY `").append(key).append("` (");
-                for (int i = 0; i < columns.length; i++) {
-                    builder.append("`").append(columns[i]).append("`");
-                    if (i < columns.length - 1) {
-                        builder.append(", ");
-                    }
-                }
-                builder.append(")");
-            }
+
+        builder.append("PRIMARY KEY (");
+
+        int size = primaryKeys.size();
+        for (int i = 0; i < size; i++) {
+            builder.append("`").append(primaryKeys.get(i)).append("`");
+            if (i >= size - 1) return;
+
+            builder.append(", ");
         }
+
         builder.append(")");
-        return builder.toString();
     }
 
+    private void addForeignKeys(StringBuilder builder) {
+        if (foreignKeys.isEmpty()) {
+            return;
+        }
+
+        for (ForeignKey entry : foreignKeys) {
+            if (entry.referencedColumn == null) continue;
+            builder.append(", FOREIGN KEY (`").append(entry.foreignKey).append("`)");
+            builder.append(" REFERENCES `").append(entry.referencedTable).append("`(`").append(entry.referencedColumn).append("`)");
+            if (entry.onDelete == null) return;
+
+            builder.append(" ON DELETE ").append(entry.onDelete.name().replace("_"," "));
+        }
+    }
+
+    private void addUniqueKeys(StringBuilder builder) {
+        if (uniqueKeys.isEmpty()) {
+            return;
+        }
 
-}
+        for (String key : uniqueKeys.keySet()) {
+            String[] columns = uniqueKeys.get(key);
+            builder.append(", UNIQUE KEY `").append(key).append("` (");
+            for (int i = 0; i < columns.length; i++) {
+                builder.append("`").append(columns[i]).append("`");
+                if (i >= columns.length - 1) continue;
+                builder.append(", ");
+            }
+            builder.append(")");
+        }
+    }
+}