From c4396601a7f59b1e2ac00fbb648340bce465ce40 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Niklas=20Schr=C3=B6der?=
 <33390735+lnschroeder@users.noreply.github.com>
Date: Wed, 1 Nov 2023 17:28:39 +0100
Subject: [PATCH] Support multiple key constraints on a single attribute (#51)
 (#52)

---
 diagram/diagram_data.go      |  8 ++++----
 diagram/diagram_util.go      | 21 +++++++++++----------
 diagram/diagram_util_test.go | 24 ++++++++++++------------
 diagram/erd_template.gommd   |  7 ++++++-
 4 files changed, 33 insertions(+), 27 deletions(-)

diff --git a/diagram/diagram_data.go b/diagram/diagram_data.go
index 6810e74..bd54da3 100644
--- a/diagram/diagram_data.go
+++ b/diagram/diagram_data.go
@@ -27,10 +27,10 @@ type ErdTableData struct {
 }
 
 type ErdColumnData struct {
-	Name         string
-	DataType     string
-	Description  string
-	AttributeKey ErdAttributeKey
+	Name          string
+	DataType      string
+	Description   string
+	AttributeKeys []ErdAttributeKey
 }
 
 type ErdConstraintData struct {
diff --git a/diagram/diagram_util.go b/diagram/diagram_util.go
index fd0d1fd..2d795dd 100644
--- a/diagram/diagram_util.go
+++ b/diagram/diagram_util.go
@@ -28,29 +28,30 @@ func tableNameInSlice(slice []ErdTableData, tableName string) bool {
 	return false
 }
 
-func getAttributeKey(column database.ColumnResult) ErdAttributeKey {
+func getAttributeKeys(column database.ColumnResult) []ErdAttributeKey {
+	var attributeKeys []ErdAttributeKey
 	if column.IsPrimary {
-		return primaryKey
+		attributeKeys = append(attributeKeys, primaryKey)
 	}
 
 	if column.IsForeign {
-		return foreignKey
+		attributeKeys = append(attributeKeys, foreignKey)
 	}
 
-	return none
+	return attributeKeys
 }
 
 func getColumnData(config config.MermerdConfig, column database.ColumnResult) ErdColumnData {
-	attributeKey := getAttributeKey(column)
+	attributeKeys := getAttributeKeys(column)
 	if config.OmitAttributeKeys() {
-		attributeKey = none
+		attributeKeys = []ErdAttributeKey(nil)
 	}
 
 	return ErdColumnData{
-		Name:         column.Name,
-		DataType:     column.DataType,
-		Description:  getDescription(config.ShowDescriptions(), column),
-		AttributeKey: attributeKey,
+		Name:          column.Name,
+		DataType:      column.DataType,
+		Description:   getDescription(config.ShowDescriptions(), column),
+		AttributeKeys: attributeKeys,
 	}
 }
 
diff --git a/diagram/diagram_util_test.go b/diagram/diagram_util_test.go
index 27dc5d3..ad713aa 100644
--- a/diagram/diagram_util_test.go
+++ b/diagram/diagram_util_test.go
@@ -46,7 +46,7 @@ func TestGetRelation(t *testing.T) {
 func TestGetAttributeKey(t *testing.T) {
 	testCases := []struct {
 		column                  database.ColumnResult
-		expectedAttributeResult ErdAttributeKey
+		expectedAttributeResult []ErdAttributeKey
 	}{
 		{
 			column: database.ColumnResult{
@@ -55,7 +55,7 @@ func TestGetAttributeKey(t *testing.T) {
 				IsPrimary: true,
 				IsForeign: false,
 			},
-			expectedAttributeResult: primaryKey,
+			expectedAttributeResult: []ErdAttributeKey{primaryKey},
 		},
 		{
 			column: database.ColumnResult{
@@ -64,7 +64,7 @@ func TestGetAttributeKey(t *testing.T) {
 				IsPrimary: false,
 				IsForeign: true,
 			},
-			expectedAttributeResult: foreignKey,
+			expectedAttributeResult: []ErdAttributeKey{foreignKey},
 		},
 		{
 			column: database.ColumnResult{
@@ -73,7 +73,7 @@ func TestGetAttributeKey(t *testing.T) {
 				IsPrimary: true,
 				IsForeign: true,
 			},
-			expectedAttributeResult: primaryKey,
+			expectedAttributeResult: []ErdAttributeKey{primaryKey, foreignKey},
 		},
 		{
 			column: database.ColumnResult{
@@ -82,7 +82,7 @@ func TestGetAttributeKey(t *testing.T) {
 				IsPrimary: false,
 				IsForeign: false,
 			},
-			expectedAttributeResult: none,
+			expectedAttributeResult: []ErdAttributeKey(nil),
 		},
 	}
 
@@ -92,7 +92,7 @@ func TestGetAttributeKey(t *testing.T) {
 			column := testCase.column
 
 			// Act
-			result := getAttributeKey(column)
+			result := getAttributeKeys(column)
 
 			// Assert
 			assert.Equal(t, testCase.expectedAttributeResult, result)
@@ -152,7 +152,7 @@ func TestGetColumnData(t *testing.T) {
 		configMock.AssertExpectations(t)
 		assert.Equal(t, columnName, result.Name)
 		assert.Equal(t, "<"+enumValues+"> "+expectedComment, result.Description)
-		assert.Equal(t, primaryKey, result.AttributeKey)
+		assert.Equal(t, []ErdAttributeKey{primaryKey}, result.AttributeKeys)
 	})
 
 	t.Run("Get all fields with enum values", func(t *testing.T) {
@@ -168,7 +168,7 @@ func TestGetColumnData(t *testing.T) {
 		configMock.AssertExpectations(t)
 		assert.Equal(t, columnName, result.Name)
 		assert.Equal(t, "<"+enumValues+">", result.Description)
-		assert.Equal(t, primaryKey, result.AttributeKey)
+		assert.Equal(t, []ErdAttributeKey{primaryKey}, result.AttributeKeys)
 	})
 
 	t.Run("Get all fields with column comments", func(t *testing.T) {
@@ -184,7 +184,7 @@ func TestGetColumnData(t *testing.T) {
 		configMock.AssertExpectations(t)
 		assert.Equal(t, columnName, result.Name)
 		assert.Equal(t, expectedComment, result.Description)
-		assert.Equal(t, primaryKey, result.AttributeKey)
+		assert.Equal(t, []ErdAttributeKey{primaryKey}, result.AttributeKeys)
 	})
 
 	t.Run("Get all fields except description", func(t *testing.T) {
@@ -200,7 +200,7 @@ func TestGetColumnData(t *testing.T) {
 		configMock.AssertExpectations(t)
 		assert.Equal(t, columnName, result.Name)
 		assert.Equal(t, "", result.Description)
-		assert.Equal(t, primaryKey, result.AttributeKey)
+		assert.Equal(t, []ErdAttributeKey{primaryKey}, result.AttributeKeys)
 	})
 
 	t.Run("Get all fields except attribute key", func(t *testing.T) {
@@ -216,7 +216,7 @@ func TestGetColumnData(t *testing.T) {
 		configMock.AssertExpectations(t)
 		assert.Equal(t, columnName, result.Name)
 		assert.Equal(t, "<"+enumValues+"> "+expectedComment, result.Description)
-		assert.Equal(t, none, result.AttributeKey)
+		assert.Equal(t, []ErdAttributeKey(nil), result.AttributeKeys)
 	})
 
 	t.Run("Get only minimal fields", func(t *testing.T) {
@@ -232,7 +232,7 @@ func TestGetColumnData(t *testing.T) {
 		configMock.AssertExpectations(t)
 		assert.Equal(t, columnName, result.Name)
 		assert.Equal(t, "", result.Description)
-		assert.Equal(t, none, result.AttributeKey)
+		assert.Equal(t, []ErdAttributeKey(nil), result.AttributeKeys)
 	})
 }
 
diff --git a/diagram/erd_template.gommd b/diagram/erd_template.gommd
index 250d47b..0c5961b 100644
--- a/diagram/erd_template.gommd
+++ b/diagram/erd_template.gommd
@@ -3,7 +3,12 @@ erDiagram
 {{- range .Tables}}
     {{.Name}} {
     {{- range .Columns}}
-        {{.DataType}} {{.Name}} {{- if .AttributeKey}} {{.AttributeKey}}{{end -}} {{- if .Description}} "{{.Description}}"{{end -}}
+        {{.DataType}} {{.Name}} {{- if .AttributeKeys}} {{range $index, $attributeKey := .AttributeKeys}}
+            {{- if $index}},{{end -}}
+            {{$attributeKey}}
+        {{- end}}{{end}} {{if .Description -}}
+            "{{.Description}}"
+        {{- end}}
     {{- end}}
     }
 {{end -}}