chore: add default CQ columns (#66)

refs: #49

Author: mnorbury

lib/src/main/java/io/cloudquery/schema/Column.java

@@ -1,15 +1,55 @@
 package io.cloudquery.schema;
 
+import io.cloudquery.types.UUIDType;
 import lombok.Builder;
 import lombok.Getter;
 import org.apache.arrow.vector.types.pojo.ArrowType;
 
-@Builder
+@Builder(toBuilder = true)
 @Getter
 public class Column {
     private String name;
+    private String description;
     private ArrowType type;
     private ColumnResolver resolver;
     private boolean primaryKey;
+    private boolean notNull;
+    private boolean unique;
+    private boolean incrementalKey;
     private boolean ignoreInTests;
+
+    public static final Column CQ_ID_COLUMN = Column.builder().
+            name("_cq_id").
+            type(new UUIDType()).
+            description("Internal CQ ID of the row").
+            notNull(true).
+            unique(true).
+            build();
+    public static final Column CQ_PARENT_ID_COLUMN = Column.builder().
+            name("_cq_parent_id").
+            type(new UUIDType()).
+            description("Internal CQ ID of the parent row").
+            ignoreInTests(true).
+            build();
+
+    @Override
+    public String toString() {
+        StringBuilder sb = new StringBuilder();
+        sb.append(name);
+        sb.append(":");
+        sb.append(type.toString());
+        if (primaryKey) {
+            sb.append(":PK");
+        }
+        if (notNull) {
+            sb.append(":NotNull");
+        }
+        if (unique) {
+            sb.append(":Unique");
+        }
+        if (incrementalKey) {
+            sb.append(":IncrementalKey");
+        }
+        return sb.toString();
+    }
 }

lib/src/main/java/io/cloudquery/schema/Table.java

@@ -1,6 +1,7 @@
 package io.cloudquery.schema;
 
 import io.cloudquery.glob.Glob;
+import io.cloudquery.schema.Column.ColumnBuilder;
 import io.cloudquery.transformers.TransformerException;
 import lombok.Builder;
 import lombok.Getter;
@@ -17,7 +18,6 @@
 @Builder(toBuilder = true)
 @Getter
 public class Table {
-
     public interface Transform {
         void transformTable(Table table) throws TransformerException;
     }
@@ -113,7 +113,7 @@ public static int maxDepth(List<Table> tables) {
     private List<Column> columns = new ArrayList<>();
 
     @Builder.Default
-    private List<Table> relations = Collections.emptyList();
+    private List<Table> relations = new ArrayList<>();
 
     private Transform transform;
 
@@ -123,6 +123,30 @@ public void transform() throws TransformerException {
         }
     }
 
+    public void addCQIDs() {
+        boolean havePrimaryKeys = !primaryKeys().isEmpty();
+        ColumnBuilder cqIdColumnBuilder = Column.CQ_ID_COLUMN.toBuilder();
+        if (!havePrimaryKeys) {
+            cqIdColumnBuilder.primaryKey(true);
+        }
+
+        List<Column> newColumns = List.of(cqIdColumnBuilder.build(), Column.CQ_PARENT_ID_COLUMN);
+        for (int i = 0; i < newColumns.size(); i++) {
+            columns.add(i, newColumns.get(i));
+        }
+
+        for (Table relation : relations) {
+            relation.addCQIDs();
+        }
+    }
+
+    public List<String> primaryKeys() {
+        return columns.stream().
+                filter(Column::isPrimaryKey).
+                map(Column::getName).
+                toList();
+    }
+
     private Optional<Table> filterDfs(boolean parentMatched, Predicate<Table> include, Predicate<Table> exclude, boolean skipDependentTables) {
         if (exclude.test(this)) {
             return Optional.empty();

lib/src/test/java/io/cloudquery/schema/TableTest.java

@@ -0,0 +1,68 @@
+package io.cloudquery.schema;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Optional;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+class TableTest {
+    @Test
+    public void shouldAddDefaultCQColumns() {
+        Table table = Table.builder().name("test").build();
+
+        assertEquals(0, table.getColumns().size());
+        table.addCQIDs();
+
+        List<Column> expectedColumns = List.of(Column.CQ_ID_COLUMN, Column.CQ_PARENT_ID_COLUMN);
+        assertEquals(expectedColumns.size(), table.getColumns().size());
+        for (int i = 0; i < expectedColumns.size(); i++) {
+            assertEquals(expectedColumns.get(i).getName(), table.getColumns().get(i).getName());
+        }
+    }
+
+    @Test
+    public void shouldAddDefaultCQColumnsToRelations() {
+        Table table = Table.builder().name("test").build();
+        Table relation = Table.builder().name("relation").build();
+        table.getRelations().add(relation);
+
+        assertEquals(0, relation.getColumns().size());
+        table.addCQIDs();
+
+        List<Column> expectedColumns = List.of(Column.CQ_ID_COLUMN, Column.CQ_PARENT_ID_COLUMN);
+        assertEquals(expectedColumns.size(), relation.getColumns().size());
+        for (int i = 0; i < expectedColumns.size(); i++) {
+            assertEquals(expectedColumns.get(i).getName(), relation.getColumns().get(i).getName());
+        }
+    }
+
+    @Test
+    public void shouldSetDefaultColumnAsPrimaryKeyIfNoOtherPrimaryKey() {
+        Table table = Table.builder().name("test").build();
+
+        table.addCQIDs();
+
+        Optional<Column> column = table.getColumn(Column.CQ_ID_COLUMN.getName());
+        assertTrue(column.isPresent());
+        assertTrue(column.get().isPrimaryKey(), "CQ_ID_COLUMN should be primary key");
+    }
+
+    @Test
+    public void shouldUseExistingPrimaryKeyWhenPossible() {
+        Column pkColumn = Column.builder().name("mypk").primaryKey(true).build();
+        Table table = Table.builder().name("test").
+                columns(new ArrayList<>(List.of(pkColumn))).
+                build();
+
+        table.addCQIDs();
+
+        Optional<Column> column = table.getColumn(Column.CQ_ID_COLUMN.getName());
+        assertTrue(column.isPresent());
+        assertFalse(column.get().isPrimaryKey(), "CQ_ID_COLUMN should not be primary key");
+    }
+}