Repository: julioz/FloorPlan
Pull request #50 (2020-07-26 09:48:25)

Parse composite primary keys indexes
Indexes allow users to quickly locate and access the data. Users can define single or multi-column indexes.

One can define a composite primary key (in DBML) via the `indexes` block:

```
Table bookings {
  id integer
  country varchar
  booking_date date
  created_at timestamp

  indexes {
      (id, country) [pk] // composite primary key
  }
}
```

Here we handle that use-case by parsing the individual composite keys for each table.

<img width="192" alt="Screenshot 2020-07-26 at 11 49 39" src="https://user-images.githubusercontent.com/1541701/88476116-196d5280-cf36-11ea-9108-67907ef579e2.png">


Patch:
diff --git a/dbml-parser/src/main/kotlin/com/zynger/floorplan/dbmlparser/CompositePrimaryKey.kt b/dbml-parser/src/main/kotlin/com/zynger/floorplan/dbmlparser/CompositePrimaryKey.kt
new file mode 100644
index 0000000..110228f
--- /dev/null
+++ b/dbml-parser/src/main/kotlin/com/zynger/floorplan/dbmlparser/CompositePrimaryKey.kt
@@ -0,0 +1,3 @@
+package com.zynger.floorplan.dbmlparser
+
+data class CompositePrimaryKey(val columnNames: List<String>)
diff --git a/dbml-parser/src/main/kotlin/com/zynger/floorplan/lex/ColumnParser.kt b/dbml-parser/src/main/kotlin/com/zynger/floorplan/lex/ColumnParser.kt
index d4ef50d..bc3676b 100644
--- a/dbml-parser/src/main/kotlin/com/zynger/floorplan/lex/ColumnParser.kt
+++ b/dbml-parser/src/main/kotlin/com/zynger/floorplan/lex/ColumnParser.kt
@@ -2,6 +2,7 @@ package com.zynger.floorplan.lex
 
 import com.zynger.floorplan.dbml.Column
 import com.zynger.floorplan.dbml.Reference
+import com.zynger.floorplan.dbmlparser.CompositePrimaryKey
 import org.intellij.lang.annotations.Language
 
 object ColumnParser {
@@ -13,14 +14,19 @@ object ColumnParser {
 
     // TODO parse column default values
 
-    fun parseColumns(tableName: String, columnsInput: String): List<Column> {
+    fun parseColumns(
+        tableName: String,
+        columnsInput: String,
+        compositePrimaryKeys: List<CompositePrimaryKey> = emptyList()
+    ): List<Column> {
+        val columnNamesFromCompositePrimaryKeys = compositePrimaryKeys.map { compositePrimaryKey -> compositePrimaryKey.columnNames }.flatten()
         return COLUMN_REGEX.findAll(columnsInput).map {
             val rawValue = it.groups[0]!!.value
             val name = it.groups[1]!!.value.removeSurroundQuotes()
             val type = it.groups[2]!!.value
             val columnProperties = it.groups[3]!!.value.trim()
             val notNull = columnProperties.contains("not null")
-            val pk = columnProperties.contains("pk")
+            val pk = columnProperties.contains("pk") || columnNamesFromCompositePrimaryKeys.contains(name)
             val increment = columnProperties.contains("increment")
             val reference: Reference? = ColumnReferenceParser.parse(tableName, name, columnProperties)
 
diff --git a/dbml-parser/src/main/kotlin/com/zynger/floorplan/lex/IndexParser.kt b/dbml-parser/src/main/kotlin/com/zynger/floorplan/lex/IndexParser.kt
index 98513da..cb1fb33 100644
--- a/dbml-parser/src/main/kotlin/com/zynger/floorplan/lex/IndexParser.kt
+++ b/dbml-parser/src/main/kotlin/com/zynger/floorplan/lex/IndexParser.kt
@@ -1,6 +1,7 @@
 package com.zynger.floorplan.lex
 
 import com.zynger.floorplan.dbml.Index
+import com.zynger.floorplan.dbmlparser.CompositePrimaryKey
 import org.intellij.lang.annotations.Language
 
 object IndexParser {
@@ -16,7 +17,7 @@ object IndexParser {
 
     fun parseIndexes(indexContent: String): List<Index> {
         return INDEXES_BLOCK_REGEX.find(indexContent)?.let {
-            INDEX_REGEX.findAll(it.groups[1]!!.value).toList().map { indexMatch ->
+            INDEX_REGEX.findAll(it.groups[1]!!.value).toList().mapNotNull { indexMatch ->
                 val indexColumns = indexMatch.groups[1]!!.value.split(",").map { columnName -> columnName.trim() }
                 val indexProperties = indexMatch.groups[2]!!.value
 
@@ -26,10 +27,24 @@ object IndexParser {
                 } else {
                     UNNAMED_INDEX
                 }
-                Index(
-                    name = indexName,
-                    columnNames = indexColumns,
-                    unique = indexProperties.contains("unique", ignoreCase = true)
+                if (indexProperties.contains("pk")) {
+                    null
+                } else {
+                    Index(
+                        name = indexName,
+                        columnNames = indexColumns,
+                        unique = indexProperties.contains("unique", ignoreCase = true)
+                    )
+                }
+            }
+        } ?: emptyList()
+    }
+
+    fun parseCompositePrimaryKeys(indexContent: String): List<CompositePrimaryKey> {
+        return INDEXES_BLOCK_REGEX.find(indexContent)?.let {
+            it.groups[1]!!.value.lines().filter { indexLine -> indexLine.contains("pk") }.map { compositePrimaryKeyLine ->
+                CompositePrimaryKey(
+                    Regex(INDEX_NAME).find(compositePrimaryKeyLine)!!.groups[1]!!.value.split(",").map { columnName -> columnName.trim() }
                 )
             }
         } ?: emptyList()
diff --git a/dbml-parser/src/main/kotlin/com/zynger/floorplan/lex/TableParser.kt b/dbml-parser/src/main/kotlin/com/zynger/floorplan/lex/TableParser.kt
index e73ccdd..660d567 100644
--- a/dbml-parser/src/main/kotlin/com/zynger/floorplan/lex/TableParser.kt
+++ b/dbml-parser/src/main/kotlin/com/zynger/floorplan/lex/TableParser.kt
@@ -16,12 +16,14 @@ object TableParser {
         return TABLE_REGEX.findAll(dbmlInput).map {
             val tableName = it.groups[1]!!.value.trim().removeSurroundQuotes()
             val tableContent = it.groups[4]!!.value
+            val indexes = IndexParser.parseIndexes(tableContent)
+            val primaryKeysFromIndexes = IndexParser.parseCompositePrimaryKeys(tableContent)
 
             Table(
                 rawValue = it.groups[0]!!.value,
                 name = tableName,
-                columns = ColumnParser.parseColumns(tableName, removeIndexes(tableContent)),
-                indexes = IndexParser.parseIndexes(tableContent)
+                columns = ColumnParser.parseColumns(tableName, removeIndexes(tableContent), primaryKeysFromIndexes),
+                indexes = indexes
             )
         }.toList()
     }


Test patch:
diff --git a/dbml-parser/src/test/kotlin/com/zynger/floorplan/lex/ColumnParserTest.kt b/dbml-parser/src/test/kotlin/com/zynger/floorplan/lex/ColumnParserTest.kt
index 54218f6..ad2899f 100644
--- a/dbml-parser/src/test/kotlin/com/zynger/floorplan/lex/ColumnParserTest.kt
+++ b/dbml-parser/src/test/kotlin/com/zynger/floorplan/lex/ColumnParserTest.kt
@@ -1,6 +1,7 @@
 package com.zynger.floorplan.lex
 
 import com.zynger.floorplan.dbml.Column
+import com.zynger.floorplan.dbmlparser.CompositePrimaryKey
 import org.junit.Assert.*
 import org.junit.Test
 
@@ -88,6 +89,28 @@ class ColumnParserTest {
         assertEquals(true, columns.first().primaryKey)
     }
 
+    @Test
+    fun `considers primary key if it is part of a composite primary key`() {
+        val columnsInput = """
+            id int
+            user_id int
+            another_col blob
+
+        """.trimIndent()
+
+        val columns = ColumnParser.parseColumns(
+            TABLE_NAME,
+            columnsInput,
+            compositePrimaryKeys = listOf(
+                CompositePrimaryKey(
+                    listOf("id", "user_id")))
+        )
+
+        assertEquals(true, columns[0].primaryKey)
+        assertEquals(true, columns[1].primaryKey)
+        assertEquals(false, columns[2].primaryKey)
+    }
+
     @Test
     fun `parses not null`() {
         val columnsInput = """
diff --git a/dbml-parser/src/test/kotlin/com/zynger/floorplan/lex/IndexParserTest.kt b/dbml-parser/src/test/kotlin/com/zynger/floorplan/lex/IndexParserTest.kt
index 3a56153..0f8c0b4 100644
--- a/dbml-parser/src/test/kotlin/com/zynger/floorplan/lex/IndexParserTest.kt
+++ b/dbml-parser/src/test/kotlin/com/zynger/floorplan/lex/IndexParserTest.kt
@@ -72,6 +72,27 @@ class IndexParserTest {
         assertEquals(true, indexes[0].unique)
     }
 
+    @Test
+    fun `composite primary key in indexes block gets removed`() {
+        val input = """
+          "user_id" int(11) [not null, default: "0"]
+          "track_id" int(11) [not null, default: "0"]
+          "created_at" datetime(6) [default: NULL]
+
+          Indexes {
+            (user_id, created_at) [name: "index_user_id_created_at"]
+            (track_id, created_at) [name: "index_track_id_created_at"]
+            (user_id, track_id) [pk]
+          }
+        """.trimIndent()
+
+        val indexes = IndexParser.parseIndexes(input)
+
+        assertEquals(2, indexes.size)
+        assertEquals("index_user_id_created_at", indexes[0].name)
+        assertEquals("index_track_id_created_at", indexes[1].name)
+    }
+
     @Ignore(value = "Unsupported: advanced usage of index.")
     @Test
     fun `index as expression`() {
@@ -185,4 +206,48 @@ class IndexParserTest {
         assertEquals(listOf("post_id"), indexes[0].columnNames)
         assertEquals(listOf("tag_id", "post_id"), indexes[1].columnNames)
     }
+
+    @Test
+    fun `parses composite primary key`() {
+        val input = """
+          "user_id" int(11) [not null, default: "0"]
+          "track_id" int(11) [not null, default: "0"]
+          "created_at" datetime(6) [default: NULL]
+
+          Indexes {
+            (user_id, created_at) [name: "index_user_id_created_at"]
+            (track_id, created_at) [name: "index_track_id_created_at"]
+            (user_id, track_id) [pk]
+          }
+        """.trimIndent()
+
+        val compositePrimaryKeys = IndexParser.parseCompositePrimaryKeys(input)
+
+        assertEquals(1, compositePrimaryKeys.size)
+        assertEquals(2, compositePrimaryKeys[0].columnNames.size)
+        assertEquals("user_id", compositePrimaryKeys[0].columnNames[0])
+        assertEquals("track_id", compositePrimaryKeys[0].columnNames[1])
+    }
+
+    @Test
+    fun `parses composite primary key even with line comment`() {
+        val input = """
+          "user_id" int(11) [not null, default: "0"]
+          "track_id" int(11) [not null, default: "0"]
+          "created_at" datetime(6) [default: NULL]
+
+          Indexes {
+            (user_id, created_at) [name: "index_user_id_created_at"]
+            (track_id, created_at) [name: "index_track_id_created_at"]
+            (user_id, track_id) [pk] // composite primary key
+          }
+        """.trimIndent()
+
+        val compositePrimaryKeys = IndexParser.parseCompositePrimaryKeys(input)
+
+        assertEquals(1, compositePrimaryKeys.size)
+        assertEquals(2, compositePrimaryKeys[0].columnNames.size)
+        assertEquals("user_id", compositePrimaryKeys[0].columnNames[0])
+        assertEquals("track_id", compositePrimaryKeys[0].columnNames[1])
+    }
 }
\ No newline at end of file
diff --git a/dbml-parser/src/test/kotlin/com/zynger/floorplan/lex/TableParserTest.kt b/dbml-parser/src/test/kotlin/com/zynger/floorplan/lex/TableParserTest.kt
index 5c03372..d5d93eb 100644
--- a/dbml-parser/src/test/kotlin/com/zynger/floorplan/lex/TableParserTest.kt
+++ b/dbml-parser/src/test/kotlin/com/zynger/floorplan/lex/TableParserTest.kt
@@ -108,4 +108,27 @@ class TableParserTest {
 
         assertEquals(2, tables.first().indexes.size)
     }
+
+    @Test
+    fun `passes index content with composite primary keys to be parsed by delegation`() {
+        val input = """
+            table post_tags [note: 'hey table note']{
+              id int [pk]
+              post_id int
+              tag_id int
+
+              Indexes  {
+                (post_id) [name:'index_post_tags_post_id', unique]
+                (tag_id) [name:'index_post_tags_tag_id']
+                (id, post_id) [pk]
+              }
+            }
+        """.trimIndent()
+
+        val tables = TableParser.parseTables(input)
+
+        assertEquals(2, tables.first().indexes.size)
+        assertEquals(true, tables.first().columns[0].primaryKey)
+        assertEquals(true, tables.first().columns[1].primaryKey)
+    }
 }