Support attribute keys (PK and FK) (#22)

* add attribute keys for all supported databases

* added tests for attribute key

* add flag and docs

* update mocks

.mermerd.example

@@ -1,6 +1,9 @@
 showAllConstraints: true
 encloseWithMermaidBackticks: false
 outputFileName: "my-db.mmd"
+debug: false
+omitConstraintLabels: false
+omitAttributeKeys: false
 
 # These connection strings are available as suggestions in the cli (use tab to access)
 connectionStringSuggestions:

cmd/root.go

@@ -67,6 +67,7 @@ func init() {
 	rootCmd.Flags().Bool(config.UseAllTablesKey, false, "use all available tables")
 	rootCmd.Flags().Bool(config.DebugKey, false, "show debug logs")
 	rootCmd.Flags().Bool(config.OmitConstraintLabelsKey, false, "omit the constraint labels")
+	rootCmd.Flags().Bool(config.OmitAttributeKeysKey, false, "omit the attribute keys (PK, FK)")
 	rootCmd.Flags().BoolP(config.EncloseWithMermaidBackticksKey, "e", false, "enclose output with mermaid backticks (needed for e.g. in markdown viewer)")
 	rootCmd.Flags().StringP(config.ConnectionStringKey, "c", "", "connection string that should be used")
 	rootCmd.Flags().StringP(config.SchemaKey, "s", "", "schema that should be used")
@@ -77,6 +78,7 @@ func init() {
 	bindFlagToViper(config.UseAllTablesKey)
 	bindFlagToViper(config.DebugKey)
 	bindFlagToViper(config.OmitConstraintLabelsKey)
+	bindFlagToViper(config.OmitAttributeKeysKey)
 	bindFlagToViper(config.EncloseWithMermaidBackticksKey)
 	bindFlagToViper(config.ConnectionStringKey)
 	bindFlagToViper(config.SchemaKey)

config/config.go

@@ -13,6 +13,7 @@ const (
 	EncloseWithMermaidBackticksKey = "encloseWithMermaidBackticks"
 	DebugKey                       = "debug"
 	OmitConstraintLabelsKey        = "omitConstraintLabels"
+	OmitAttributeKeysKey           = "omitAttributeKeys"
 )
 
 type config struct{}
@@ -28,6 +29,7 @@ type MermerdConfig interface {
 	EncloseWithMermaidBackticks() bool
 	Debug() bool
 	OmitConstraintLabels() bool
+	OmitAttributeKeys() bool
 }
 
 func NewConfig() MermerdConfig {
@@ -73,3 +75,7 @@ func (c config) Debug() bool {
 func (c config) OmitConstraintLabels() bool {
 	return viper.GetBool(OmitConstraintLabelsKey)
 }
+
+func (c config) OmitAttributeKeys() bool {
+	return viper.GetBool(OmitAttributeKeysKey)
+}

database/database_integration_test.go

@@ -8,6 +8,12 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
+type columnTestResult struct {
+	Name      string
+	isPrimary bool
+	isForeign bool
+}
+
 func TestDatabaseIntegrations(t *testing.T) {
 	if testing.Short() {
 		t.Skip("skipping integration test")
@@ -103,33 +109,59 @@ func TestDatabaseIntegrations(t *testing.T) {
 				connector := getConnectionAndConnect(t)
 				testCases := []struct {
 					tableName       string
-					expectedColumns []string
+					expectedColumns []columnTestResult
 				}{
-					{tableName: "article", expectedColumns: []string{"id", "title"}},
-					{tableName: "article_detail", expectedColumns: []string{"id", "created_at"}},
-					{tableName: "article_comment", expectedColumns: []string{"id", "article_id", "comment"}},
-					{tableName: "label", expectedColumns: []string{"id", "label"}},
-					{tableName: "article_label", expectedColumns: []string{"article_id", "label_id"}},
-					{tableName: "test_1_a", expectedColumns: []string{"id", "xid"}},
-					{tableName: "test_1_b", expectedColumns: []string{"aid", "bid"}},
+					{tableName: "article", expectedColumns: []columnTestResult{
+						{Name: "id", isPrimary: true, isForeign: false},
+						{Name: "title", isPrimary: false, isForeign: false},
+					}},
+					{tableName: "article_detail", expectedColumns: []columnTestResult{
+						{Name: "id", isPrimary: true, isForeign: true},
+						{Name: "created_at", isPrimary: false, isForeign: false},
+					}},
+					{tableName: "article_comment", expectedColumns: []columnTestResult{
+						{Name: "id", isPrimary: true, isForeign: false},
+						{Name: "article_id", isPrimary: false, isForeign: true},
+						{Name: "comment", isPrimary: false, isForeign: false},
+					}},
+					{tableName: "label", expectedColumns: []columnTestResult{
+						{Name: "id", isPrimary: true, isForeign: false},
+						{Name: "label", isPrimary: false, isForeign: false},
+					}},
+					{tableName: "article_label", expectedColumns: []columnTestResult{
+						{Name: "article_id", isPrimary: true, isForeign: true},
+						{Name: "label_id", isPrimary: true, isForeign: true},
+					}},
+					{tableName: "test_1_a", expectedColumns: []columnTestResult{
+						{Name: "id", isPrimary: true, isForeign: false},
+						{Name: "xid", isPrimary: true, isForeign: false},
+					}},
+					{tableName: "test_1_b", expectedColumns: []columnTestResult{
+						{Name: "aid", isPrimary: true, isForeign: true},
+						{Name: "bid", isPrimary: true, isForeign: true},
+					}},
 				}
 
 				for index, testCase := range testCases {
 					t.Run(fmt.Sprintf("run #%d", index), func(t *testing.T) {
 						// Arrange
 						tableName := testCase.tableName
-						var columnNames []string
+						var columnResult []columnTestResult
 
 						// Act
 						columns, err := connector.GetColumns(tableName)
 
 						// Assert
 						for _, column := range columns {
-							columnNames = append(columnNames, column.Name)
+							columnResult = append(columnResult, columnTestResult{
+								Name:      column.Name,
+								isPrimary: column.IsPrimary,
+								isForeign: column.IsForeign,
+							})
 						}
 
 						assert.Nil(t, err)
-						assert.ElementsMatch(t, testCase.expectedColumns, columnNames)
+						assert.ElementsMatch(t, testCase.expectedColumns, columnResult)
 					})
 				}
 			})

database/mssql.go

@@ -79,10 +79,23 @@ func (c *mssqlConnector) GetTables(schemaName string) ([]string, error) {
 
 func (c *mssqlConnector) GetColumns(tableName string) ([]ColumnResult, error) {
 	rows, err := c.db.Query(`
-		select column_name, data_type
-		from information_schema.columns
-		where table_name = @p1
-		order by ordinal_position
+		select c.column_name,
+			   c.data_type,
+			   (select count(*)
+				from information_schema.key_column_usage cu
+						 left join information_schema.table_constraints tc on tc.constraint_name = cu.constraint_name
+				where cu.column_name = c.column_name
+				  and cu.table_name = c.table_name
+				  and tc.constraint_type = 'PRIMARY KEY') as is_primary,
+			   (select count(*)
+				from information_schema.key_column_usage cu
+						 left join information_schema.table_constraints tc on tc.constraint_name = cu.constraint_name
+				where cu.column_name = c.column_name
+				  and cu.table_name = c.table_name
+				  and tc.constraint_type = 'FOREIGN KEY') as is_foreign
+		from information_schema.columns c
+		where c.table_name = @p1
+		order by c.ordinal_position;
 		`, tableName)
 	if err != nil {
 		return nil, err
@@ -91,7 +104,7 @@ func (c *mssqlConnector) GetColumns(tableName string) ([]ColumnResult, error) {
 	var columns []ColumnResult
 	for rows.Next() {
 		var column ColumnResult
-		if err = rows.Scan(&column.Name, &column.DataType); err != nil {
+		if err = rows.Scan(&column.Name, &column.DataType, &column.IsPrimary, &column.IsForeign); err != nil {
 			return nil, err
 		}
 

database/mysql.go

@@ -79,10 +79,22 @@ func (c *mySqlConnector) GetTables(schemaName string) ([]string, error) {
 
 func (c *mySqlConnector) GetColumns(tableName string) ([]ColumnResult, error) {
 	rows, err := c.db.Query(`
-		select column_name, data_type
-		from information_schema.columns
-		where table_name = ?
-		order by ordinal_position
+		select c.column_name,
+			   c.data_type,
+			   (select count(*)
+				from information_schema.KEY_COLUMN_USAGE
+				where table_name = c.table_name
+				  and column_name = c.column_name
+				  and constraint_name = 'PRIMARY')        as is_primary,
+			   (select count(*)
+				from information_schema.key_column_usage cu
+						 left join information_schema.table_constraints tc on tc.constraint_name = cu.constraint_name
+				where cu.column_name = c.column_name
+				  and cu.table_name = c.table_name
+				  and tc.constraint_type = 'FOREIGN KEY') as is_foreign
+		from information_schema.columns c
+		where c.table_name = ?
+		order by c.ordinal_position;
 		`, tableName)
 	if err != nil {
 		return nil, err
@@ -91,7 +103,7 @@ func (c *mySqlConnector) GetColumns(tableName string) ([]ColumnResult, error) {
 	var columns []ColumnResult
 	for rows.Next() {
 		var column ColumnResult
-		if err = rows.Scan(&column.Name, &column.DataType); err != nil {
+		if err = rows.Scan(&column.Name, &column.DataType, &column.IsPrimary, &column.IsForeign); err != nil {
 			return nil, err
 		}
 

database/postgres.go

@@ -79,10 +79,23 @@ func (c *postgresConnector) GetTables(schemaName string) ([]string, error) {
 
 func (c *postgresConnector) GetColumns(tableName string) ([]ColumnResult, error) {
 	rows, err := c.db.Query(`
-		select column_name, data_type
-		from information_schema.columns
-		where table_name = $1
-		order by ordinal_position
+		select c.column_name,
+			   c.data_type,
+			   (select count(*)
+				from information_schema.key_column_usage cu
+						 left join information_schema.table_constraints tc on tc.constraint_name = cu.constraint_name
+				where cu.column_name = c.column_name
+				  and cu.table_name = c.table_name
+				  and tc.constraint_type = 'PRIMARY KEY') as is_primary,
+			   (select count(*)
+				from information_schema.key_column_usage cu
+						 left join information_schema.table_constraints tc on tc.constraint_name = cu.constraint_name
+				where cu.column_name = c.column_name
+				  and cu.table_name = c.table_name
+				  and tc.constraint_type = 'FOREIGN KEY') as is_foreign
+		from information_schema.columns c
+		where c.table_name = $1
+		order by c.ordinal_position;
 		`, tableName)
 	if err != nil {
 		return nil, err
@@ -91,7 +104,7 @@ func (c *postgresConnector) GetColumns(tableName string) ([]ColumnResult, error)
 	var columns []ColumnResult
 	for rows.Next() {
 		var column ColumnResult
-		if err = rows.Scan(&column.Name, &column.DataType); err != nil {
+		if err = rows.Scan(&column.Name, &column.DataType, &column.IsPrimary, &column.IsForeign); err != nil {
 			return nil, err
 		}
 

database/result.go

@@ -11,8 +11,10 @@ type TableResult struct {
 }
 
 type ColumnResult struct {
-	Name     string
-	DataType string
+	Name      string
+	DataType  string
+	IsPrimary bool
+	IsForeign bool
 }
 
 type ConstraintResultList []ConstraintResult

diagram/diagram.go

@@ -49,9 +49,15 @@ func (d diagram) Create(result *database.Result) error {
 
 		columnData := make([]ErdColumnData, len(table.Columns))
 		for columnIndex, column := range table.Columns {
+			attributeKey := getAttributeKey(column)
+			if d.config.OmitAttributeKeys() {
+				attributeKey = none
+			}
+
 			columnData[columnIndex] = ErdColumnData{
-				Name:     column.Name,
-				DataType: column.DataType,
+				Name:         column.Name,
+				DataType:     column.DataType,
+				AttributeKey: attributeKey,
 			}
 		}
 
@@ -110,3 +116,15 @@ func tableNameInSlice(slice []ErdTableData, tableName string) bool {
 
 	return false
 }
+
+func getAttributeKey(column database.ColumnResult) ErdAttributeKey {
+	if column.IsPrimary {
+		return primaryKey
+	}
+
+	if column.IsForeign {
+		return foreignKey
+	}
+
+	return none
+}

diagram/diagram_data.go

@@ -7,6 +7,14 @@ const (
 	relationManyToOne ErdRelationType = "}o--||"
 )
 
+type ErdAttributeKey string
+
+const (
+	primaryKey ErdAttributeKey = "PK"
+	foreignKey ErdAttributeKey = "FK"
+	none       ErdAttributeKey = ""
+)
+
 type ErdDiagramData struct {
 	EncloseWithMermaidBackticks bool
 	Tables                      []ErdTableData
@@ -19,8 +27,9 @@ type ErdTableData struct {
 }
 
 type ErdColumnData struct {
-	Name     string
-	DataType string
+	Name         string
+	DataType     string
+	AttributeKey ErdAttributeKey
 }
 
 type ErdConstraintData struct {

diagram/diagram_test.go

@@ -40,3 +40,60 @@ func TestGetRelation(t *testing.T) {
 		})
 	}
 }
+
+func TestGetAttributeKey(t *testing.T) {
+	testCases := []struct {
+		column                  database.ColumnResult
+		expectedAttributeResult ErdAttributeKey
+	}{
+		{
+			column: database.ColumnResult{
+				Name:      "",
+				DataType:  "",
+				IsPrimary: true,
+				IsForeign: false,
+			},
+			expectedAttributeResult: primaryKey,
+		},
+		{
+			column: database.ColumnResult{
+				Name:      "",
+				DataType:  "",
+				IsPrimary: false,
+				IsForeign: true,
+			},
+			expectedAttributeResult: foreignKey,
+		},
+		{
+			column: database.ColumnResult{
+				Name:      "",
+				DataType:  "",
+				IsPrimary: true,
+				IsForeign: true,
+			},
+			expectedAttributeResult: primaryKey,
+		},
+		{
+			column: database.ColumnResult{
+				Name:      "",
+				DataType:  "",
+				IsPrimary: false,
+				IsForeign: false,
+			},
+			expectedAttributeResult: none,
+		},
+	}
+
+	for index, testCase := range testCases {
+		t.Run(fmt.Sprintf("run #%d", index), func(t *testing.T) {
+			// Arrange
+			column := testCase.column
+
+			// Act
+			result := getAttributeKey(column)
+
+			// Assert
+			assert.Equal(t, testCase.expectedAttributeResult, result)
+		})
+	}
+}

diagram/erd_template.gommd

@@ -3,7 +3,7 @@ erDiagram
 {{- range .Tables}}
     {{.Name}} {
     {{- range .Columns}}
-        {{.DataType}} {{.Name}}
+        {{.DataType}} {{.Name}} {{.AttributeKey}}
     {{- end}}
     }
 {{end -}}

exampleRunConfig.yaml

@@ -14,3 +14,4 @@ outputFileName: "my-db.mmd"
 encloseWithMermaidBackticks: false
 debug: false
 omitConstraintLabels: false
+omitAttributeKeys: false