Add primary key constraint to SQL AST

4 files changed, 120 insertions, 13 deletions

diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/sql/SqlWriter.java b/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/sql/SqlWriter.java
index 20b28c3..b7d0caf 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/sql/SqlWriter.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/sql/SqlWriter.java
@@ -1,10 +1,7 @@
 package org.zwobble.hobgoblin.compiler.output.lang.sql;
 
 import org.zwobble.hobgoblin.compiler.output.CodeWriter;
-import org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlColumnDefinition;
-import org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlCreateTableStatement;
-import org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlIdentifier;
-import org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlStatement;
+import org.zwobble.hobgoblin.compiler.output.lang.sql.ast.*;
 
 import java.io.IOException;
 import java.nio.file.Path;
@@ -40,6 +37,8 @@ public class SqlWriter {
         this.writer = writer;
     }
 
+    // == Statements ==
+
     void writeStatement(SqlStatement statement) throws IOException {
         switch (statement) {
             case SqlCreateTableStatement createTableStatement -> {
@@ -64,6 +63,12 @@ public class SqlWriter {
             }
         );
 
+        for (var tableConstraint : createTableStatement.tableConstraints()) {
+            this.writer.write(",");
+            this.writer.newLine();
+            writeTableConstraint(tableConstraint);
+        }
+
         this.writer.dedent();
         this.writer.newLine();
         this.writer.write(");");
@@ -75,6 +80,28 @@ public class SqlWriter {
         this.writeIdentifier(column.typeName());
     }
 
+    // == Table constraints ==
+
+    void writeTableConstraint(SqlTableConstraint tableConstraint) throws IOException {
+        switch (tableConstraint) {
+            case SqlPrimaryKeyConstraint primaryKeyConstraint -> {
+                writePrimaryKeyConstraint(primaryKeyConstraint);
+            }
+        }
+    }
+
+    private void writePrimaryKeyConstraint(SqlPrimaryKeyConstraint primaryKeyConstraint) throws IOException {
+        this.writer.write("PRIMARY KEY (");
+        writeWithSeparator(
+            primaryKeyConstraint.columns(),
+            this::writeIdentifier,
+            () -> this.writer.write(", ")
+        );
+        this.writer.write(")");
+    }
+
+    // == Identifiers ==
+
     private void writeIdentifier(SqlIdentifier name) throws IOException {
         // TODO: handle escaping
         this.writer.write(name.value());
diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/sql/ast/SqlCreateTableStatement.java b/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/sql/ast/SqlCreateTableStatement.java
index a5fec23..85594a2 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/sql/ast/SqlCreateTableStatement.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/sql/ast/SqlCreateTableStatement.java
@@ -7,42 +7,60 @@ package org.zwobble.hobgoblin.compiler.output.lang.sql.ast;
 
 public record SqlCreateTableStatement(
     org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlIdentifier name,
-    java.util.List<org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlColumnDefinition> columns
+    java.util.List<org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlColumnDefinition> columns,
+    java.util.List<org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlTableConstraint> tableConstraints
 ) implements org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlStatement {
     public static org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlCreateTableStatement.Builder arbitrary() {
-        return new org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlCreateTableStatement.Builder(org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlIdentifier.arbitrary().build(), java.util.List.of());
+        return new org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlCreateTableStatement.Builder(org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlIdentifier.arbitrary().build(), java.util.List.of(), java.util.List.of());
     }
 
     public record Builder(
         org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlIdentifier name,
-        java.util.List<org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlColumnDefinition> columns
+        java.util.List<org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlColumnDefinition> columns,
+        java.util.List<org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlTableConstraint> tableConstraints
     ) implements org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlStatement.Builder {
         public org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlCreateTableStatement build() {
-            return new org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlCreateTableStatement(name, columns);
+            return new org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlCreateTableStatement(name, columns, tableConstraints);
         }
 
         public org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlCreateTableStatement.Builder withName(org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlIdentifier name) {
-            return new org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlCreateTableStatement.Builder(name, columns);
+            return new org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlCreateTableStatement.Builder(name, columns, tableConstraints);
         }
 
         public org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlCreateTableStatement.Builder withName(org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlIdentifier.Builder name) {
-            return new org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlCreateTableStatement.Builder(name.build(), columns);
+            return new org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlCreateTableStatement.Builder(name.build(), columns, tableConstraints);
         }
 
         public org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlCreateTableStatement.Builder withColumns(java.util.List<org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlColumnDefinition> columns) {
-            return new org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlCreateTableStatement.Builder(name, columns);
+            return new org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlCreateTableStatement.Builder(name, columns, tableConstraints);
         }
 
         public org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlCreateTableStatement.Builder addColumn(org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlColumnDefinition column) {
             var columns = new java.util.ArrayList<org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlColumnDefinition>(this.columns);
             columns.add(column);
-            return new org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlCreateTableStatement.Builder(name, columns);
+            return new org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlCreateTableStatement.Builder(name, columns, tableConstraints);
         }
 
         public org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlCreateTableStatement.Builder addColumn(org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlColumnDefinition.Builder column) {
             var columns = new java.util.ArrayList<org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlColumnDefinition>(this.columns);
             columns.add(column.build());
-            return new org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlCreateTableStatement.Builder(name, columns);
+            return new org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlCreateTableStatement.Builder(name, columns, tableConstraints);
+        }
+
+        public org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlCreateTableStatement.Builder withTableConstraints(java.util.List<org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlTableConstraint> tableConstraints) {
+            return new org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlCreateTableStatement.Builder(name, columns, tableConstraints);
+        }
+
+        public org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlCreateTableStatement.Builder addTableConstraint(org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlTableConstraint tableConstraint) {
+            var tableConstraints = new java.util.ArrayList<org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlTableConstraint>(this.tableConstraints);
+            tableConstraints.add(tableConstraint);
+            return new org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlCreateTableStatement.Builder(name, columns, tableConstraints);
+        }
+
+        public org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlCreateTableStatement.Builder addTableConstraint(org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlTableConstraint.Builder tableConstraint) {
+            var tableConstraints = new java.util.ArrayList<org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlTableConstraint>(this.tableConstraints);
+            tableConstraints.add(tableConstraint.build());
+            return new org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlCreateTableStatement.Builder(name, columns, tableConstraints);
         }
 
         // Custom area start: org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlCreateTableStatement.Builder body
diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/sql/ast/SqlPrimaryKeyConstraint.java b/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/sql/ast/SqlPrimaryKeyConstraint.java
new file mode 100644
index 0000000..1876219
--- /dev/null
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/sql/ast/SqlPrimaryKeyConstraint.java
@@ -0,0 +1,44 @@
+// Generated by hobgoblin.
+
+package org.zwobble.hobgoblin.compiler.output.lang.sql.ast;
+
+// Custom area start: org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlPrimaryKeyConstraint imports
+// Custom area end: org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlPrimaryKeyConstraint imports
+
+public record SqlPrimaryKeyConstraint(
+    java.util.List<org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlIdentifier> columns
+) implements org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlTableConstraint {
+    public static org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlPrimaryKeyConstraint.Builder arbitrary() {
+        return new org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlPrimaryKeyConstraint.Builder(java.util.List.of());
+    }
+
+    public record Builder(
+        java.util.List<org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlIdentifier> columns
+    ) implements org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlTableConstraint.Builder {
+        public org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlPrimaryKeyConstraint build() {
+            return new org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlPrimaryKeyConstraint(columns);
+        }
+
+        public org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlPrimaryKeyConstraint.Builder withColumns(java.util.List<org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlIdentifier> columns) {
+            return new org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlPrimaryKeyConstraint.Builder(columns);
+        }
+
+        public org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlPrimaryKeyConstraint.Builder addColumn(org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlIdentifier column) {
+            var columns = new java.util.ArrayList<org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlIdentifier>(this.columns);
+            columns.add(column);
+            return new org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlPrimaryKeyConstraint.Builder(columns);
+        }
+
+        public org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlPrimaryKeyConstraint.Builder addColumn(org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlIdentifier.Builder column) {
+            var columns = new java.util.ArrayList<org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlIdentifier>(this.columns);
+            columns.add(column.build());
+            return new org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlPrimaryKeyConstraint.Builder(columns);
+        }
+
+        // Custom area start: org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlPrimaryKeyConstraint.Builder body
+        // Custom area end: org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlPrimaryKeyConstraint.Builder body
+    }
+
+    // Custom area start: org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlPrimaryKeyConstraint body
+    // Custom area end: org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlPrimaryKeyConstraint body
+}
diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/sql/ast/SqlTableConstraint.java b/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/sql/ast/SqlTableConstraint.java
new file mode 100644
index 0000000..e35162f
--- /dev/null
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/sql/ast/SqlTableConstraint.java
@@ -0,0 +1,18 @@
+// Generated by hobgoblin.
+
+package org.zwobble.hobgoblin.compiler.output.lang.sql.ast;
+
+// Custom area start: org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlTableConstraint imports
+// Custom area end: org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlTableConstraint imports
+
+public sealed interface SqlTableConstraint permits org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlPrimaryKeyConstraint {
+    public interface Builder {
+        public org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlTableConstraint build();
+
+        // Custom area start: org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlTableConstraint.Builder body
+        // Custom area end: org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlTableConstraint.Builder body
+    }
+
+    // Custom area start: org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlTableConstraint body
+    // Custom area end: org.zwobble.hobgoblin.compiler.output.lang.sql.ast.SqlTableConstraint body
+}