Add unique constraint to key column (#53)

* Support multiple key constraints on a single attribute (#51) * Add unique constraint to key column (#43) * Relationship labels (#50) * Adds relationship label types and parser * Lookup label based on pk and fk names; overrides omitting the label and the constraint label * First full working version * Use a map for faster lookup * Fix example labels * Adds comments for label regex * Adds basic tests for relationship label map * Support multiple key constraints on a single attribute (#51) (#52) * Add unique constraint to key column (#43) --------- Co-authored-by: Dan Goslen <dwgoslen@gmail.com>
KarnerTh · Nov 18, 2023 · bb145c8 · bb145c8
1 parent c439660
commit bb145c8
Show file tree

Hide file tree

Showing 11 changed files with 93 additions and 24 deletions.
diff --git a/cmd/root.go b/cmd/root.go
@@ -68,7 +68,7 @@ func init() {
 	rootCmd.Flags().Bool(config.UseAllSchemasKey, false, "use all available schemas")
 	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().Bool(config.OmitAttributeKeysKey, false, "omit the attribute keys (PK, FK, UK)")
 	rootCmd.Flags().Bool(config.ShowSchemaPrefix, false, "show schema prefix in table name")
 	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")

diff --git a/database/database_integration_test.go b/database/database_integration_test.go
@@ -12,6 +12,7 @@ type columnTestResult struct {
 	Name       string
 	isPrimary  bool
 	isForeign  bool
+	isUnique   bool
 	isNullable bool
 }
 
@@ -136,34 +137,34 @@ func TestDatabaseIntegrations(t *testing.T) {
 					expectedColumns []columnTestResult
 				}{
 					{tableName: "article", expectedColumns: []columnTestResult{
-						{Name: "id", isPrimary: true, isForeign: false, isNullable: false},
-						{Name: "title", isPrimary: false, isForeign: false, isNullable: false},
-						{Name: "subtitle", isPrimary: false, isForeign: false, isNullable: true},
+						{Name: "id", isPrimary: true, isForeign: false, isUnique: false, isNullable: false},
+						{Name: "title", isPrimary: false, isForeign: false, isUnique: false, isNullable: false},
+						{Name: "subtitle", isPrimary: false, isForeign: false, isUnique: false, isNullable: true},
 					}},
 					{tableName: "article_detail", expectedColumns: []columnTestResult{
-						{Name: "id", isPrimary: true, isForeign: true, isNullable: false},
-						{Name: "created_at", isPrimary: false, isForeign: false, isNullable: false},
+						{Name: "id", isPrimary: true, isForeign: true, isUnique: false, isNullable: false},
+						{Name: "created_at", isPrimary: false, isForeign: false, isUnique: false, isNullable: false},
 					}},
 					{tableName: "article_comment", expectedColumns: []columnTestResult{
-						{Name: "id", isPrimary: true, isForeign: false, isNullable: false},
-						{Name: "article_id", isPrimary: false, isForeign: true, isNullable: false},
-						{Name: "comment", isPrimary: false, isForeign: false, isNullable: false},
+						{Name: "id", isPrimary: true, isForeign: false, isUnique: false, isNullable: false},
+						{Name: "article_id", isPrimary: false, isForeign: true, isUnique: false, isNullable: false},
+						{Name: "comment", isPrimary: false, isForeign: false, isUnique: false, isNullable: false},
 					}},
 					{tableName: "label", expectedColumns: []columnTestResult{
-						{Name: "id", isPrimary: true, isForeign: false, isNullable: false},
-						{Name: "label", isPrimary: false, isForeign: false, isNullable: false},
+						{Name: "id", isPrimary: true, isForeign: false, isUnique: false, isNullable: false},
+						{Name: "label", isPrimary: false, isForeign: false, isUnique: true, isNullable: false},
 					}},
 					{tableName: "article_label", expectedColumns: []columnTestResult{
-						{Name: "article_id", isPrimary: true, isForeign: true, isNullable: false},
-						{Name: "label_id", isPrimary: true, isForeign: true, isNullable: false},
+						{Name: "article_id", isPrimary: true, isForeign: true, isUnique: false, isNullable: false},
+						{Name: "label_id", isPrimary: true, isForeign: true, isUnique: false, isNullable: false},
 					}},
 					{tableName: "test_1_a", expectedColumns: []columnTestResult{
-						{Name: "id", isPrimary: true, isForeign: false, isNullable: false},
-						{Name: "xid", isPrimary: true, isForeign: false, isNullable: false},
+						{Name: "id", isPrimary: true, isForeign: false, isUnique: false, isNullable: false},
+						{Name: "xid", isPrimary: true, isForeign: false, isUnique: false, isNullable: false},
 					}},
 					{tableName: "test_1_b", expectedColumns: []columnTestResult{
-						{Name: "aid", isPrimary: true, isForeign: true, isNullable: false},
-						{Name: "bid", isPrimary: true, isForeign: true, isNullable: false},
+						{Name: "aid", isPrimary: true, isForeign: true, isUnique: false, isNullable: false},
+						{Name: "bid", isPrimary: true, isForeign: true, isUnique: false, isNullable: false},
 					}},
 				}
 
@@ -182,6 +183,7 @@ func TestDatabaseIntegrations(t *testing.T) {
 								Name:       column.Name,
 								isPrimary:  column.IsPrimary,
 								isForeign:  column.IsForeign,
+								isUnique:   column.IsUnique,
 								isNullable: column.IsNullable,
 							})
 						}

diff --git a/database/mssql.go b/database/mssql.go
@@ -102,6 +102,12 @@ func (c *mssqlConnector) GetColumns(tableName TableDetail) ([]ColumnResult, erro
 				where cu.column_name = c.column_name
 				  and cu.table_name = c.table_name
 				  and tc.constraint_type = 'FOREIGN KEY') as is_foreign,
+			   (select IIF(count(*) > 0, 1, 0)
+				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 = 'UNIQUE') as is_unique,
 		    	case when c.is_nullable = 'YES' then 1 else 0 end as is_nullable,
 			   (select ISNULL(ep.value, '') from sys.tables t
 			      inner join sys.columns col on col.object_id = t.object_id and col.name = c.column_name
@@ -118,7 +124,7 @@ func (c *mssqlConnector) GetColumns(tableName TableDetail) ([]ColumnResult, erro
 	var columns []ColumnResult
 	for rows.Next() {
 		var column ColumnResult
-		if err = rows.Scan(&column.Name, &column.DataType, &column.IsPrimary, &column.IsForeign, &column.IsNullable, &column.Comment); err != nil {
+		if err = rows.Scan(&column.Name, &column.DataType, &column.IsPrimary, &column.IsForeign, &column.IsUnique, &column.IsNullable, &column.Comment); err != nil {
 			return nil, err
 		}
 

diff --git a/database/mysql.go b/database/mysql.go
@@ -99,6 +99,12 @@ func (c *mySqlConnector) GetColumns(tableName TableDetail) ([]ColumnResult, erro
 				where cu.column_name = c.column_name
 				  and cu.table_name = c.table_name
 				  and tc.constraint_type = 'FOREIGN KEY') as is_foreign,
+				(select count(*) > 0
+				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 = 'UNIQUE') as is_unique,
 		    	IF(c.is_nullable = 'YES', 1, 0) as is_nullable,
         case when c.data_type = 'enum' then REPLACE(REPLACE(REPLACE(REPLACE(c.column_type, 'enum', ''), '\'', ''), '(', ''), ')', '') else '' end as enum_values,
 		c.column_comment as comment
@@ -113,7 +119,7 @@ func (c *mySqlConnector) GetColumns(tableName TableDetail) ([]ColumnResult, erro
 	var columns []ColumnResult
 	for rows.Next() {
 		var column ColumnResult
-		if err = rows.Scan(&column.Name, &column.DataType, &column.IsPrimary, &column.IsForeign, &column.IsNullable, &column.EnumValues, &column.Comment); err != nil {
+		if err = rows.Scan(&column.Name, &column.DataType, &column.IsPrimary, &column.IsForeign, &column.IsUnique, &column.IsNullable, &column.EnumValues, &column.Comment); err != nil {
 			return nil, err
 		}
 

diff --git a/database/postgres.go b/database/postgres.go
@@ -101,6 +101,12 @@ func (c *postgresConnector) GetColumns(tableName TableDetail) ([]ColumnResult, e
                 where cu.column_name = c.column_name
                   and cu.table_name = c.table_name
                   and tc.constraint_type = 'FOREIGN KEY')                      as is_foreign,
+				(select count(*) > 0
+				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 = 'UNIQUE')                           as is_unique,
                bool_or(c.is_nullable = 'YES')                                  as is_not_null,
                coalesce(string_agg(enumlabel, ',' order by enumsortorder), '') as enum_values,
                coalesce(pd.description, '')                   				   as comment
@@ -126,7 +132,7 @@ func (c *postgresConnector) GetColumns(tableName TableDetail) ([]ColumnResult, e
 	var columns []ColumnResult
 	for rows.Next() {
 		var column ColumnResult
-		if err = rows.Scan(&column.Name, &column.DataType, &column.IsPrimary, &column.IsForeign, &column.IsNullable, &column.EnumValues, &column.Comment); err != nil {
+		if err = rows.Scan(&column.Name, &column.DataType, &column.IsPrimary, &column.IsForeign, &column.IsUnique, &column.IsNullable, &column.EnumValues, &column.Comment); err != nil {
 			return nil, err
 		}
 

diff --git a/database/result.go b/database/result.go
@@ -20,6 +20,7 @@ type ColumnResult struct {
 	DataType   string
 	IsPrimary  bool
 	IsForeign  bool
+	IsUnique   bool
 	IsNullable bool
 	EnumValues string
 	Comment    string

diff --git a/diagram/diagram_data.go b/diagram/diagram_data.go
@@ -12,7 +12,7 @@ type ErdAttributeKey string
 const (
 	primaryKey ErdAttributeKey = "PK"
 	foreignKey ErdAttributeKey = "FK"
-	none       ErdAttributeKey = ""
+	uniqueKey  ErdAttributeKey = "UK"
 )
 
 type ErdDiagramData struct {

diff --git a/diagram/diagram_util.go b/diagram/diagram_util.go
@@ -38,6 +38,10 @@ func getAttributeKeys(column database.ColumnResult) []ErdAttributeKey {
 		attributeKeys = append(attributeKeys, foreignKey)
 	}
 
+	if column.IsUnique {
+		attributeKeys = append(attributeKeys, uniqueKey)
+	}
+
 	return attributeKeys
 }
 

diff --git a/diagram/diagram_util_test.go b/diagram/diagram_util_test.go
@@ -54,6 +54,7 @@ func TestGetAttributeKey(t *testing.T) {
 				DataType:  "",
 				IsPrimary: true,
 				IsForeign: false,
+				IsUnique:  false,
 			},
 			expectedAttributeResult: []ErdAttributeKey{primaryKey},
 		},
@@ -63,6 +64,7 @@ func TestGetAttributeKey(t *testing.T) {
 				DataType:  "",
 				IsPrimary: false,
 				IsForeign: true,
+				IsUnique:  false,
 			},
 			expectedAttributeResult: []ErdAttributeKey{foreignKey},
 		},
@@ -72,6 +74,7 @@ func TestGetAttributeKey(t *testing.T) {
 				DataType:  "",
 				IsPrimary: true,
 				IsForeign: true,
+				IsUnique:  false,
 			},
 			expectedAttributeResult: []ErdAttributeKey{primaryKey, foreignKey},
 		},
@@ -81,6 +84,47 @@ func TestGetAttributeKey(t *testing.T) {
 				DataType:  "",
 				IsPrimary: false,
 				IsForeign: false,
+				IsUnique:  true,
+			},
+			expectedAttributeResult: []ErdAttributeKey{uniqueKey},
+		},
+		{
+			column: database.ColumnResult{
+				Name:      "",
+				DataType:  "",
+				IsPrimary: true,
+				IsForeign: false,
+				IsUnique:  true,
+			},
+			expectedAttributeResult: []ErdAttributeKey{primaryKey, uniqueKey},
+		},
+		{
+			column: database.ColumnResult{
+				Name:      "",
+				DataType:  "",
+				IsPrimary: false,
+				IsForeign: true,
+				IsUnique:  true,
+			},
+			expectedAttributeResult: []ErdAttributeKey{foreignKey, uniqueKey},
+		},
+		{
+			column: database.ColumnResult{
+				Name:      "",
+				DataType:  "",
+				IsPrimary: true,
+				IsForeign: true,
+				IsUnique:  true,
+			},
+			expectedAttributeResult: []ErdAttributeKey{primaryKey, foreignKey, uniqueKey},
+		},
+		{
+			column: database.ColumnResult{
+				Name:      "",
+				DataType:  "",
+				IsPrimary: false,
+				IsForeign: false,
+				IsUnique:  false,
 			},
 			expectedAttributeResult: []ErdAttributeKey(nil),
 		},

diff --git a/readme.md b/readme.md
@@ -42,7 +42,7 @@ for your operating system. To be able to use it globally on your system, add the
 * Interactive cli (multiselect, search for tables and schemas, etc.)
 * Use it in CI/CD pipeline via a run configuration
 * Either generate plain mermaid syntax or enclose it with mermaid backticks to use directly in e.g. GitHub markdown
-* Show primary and foreign keys
+* Show primary keys, foreign keys, and unique constraints
 * Show enum values of enum column
 * Show column comments
 * Show NOT NULL constraints
@@ -74,7 +74,7 @@ via `mermerd -h`
       --debug                         show debug logs        
   -e, --encloseWithMermaidBackticks   enclose output with mermaid backticks (needed for e.g. in markdown viewer)
   -h, --help                          help for mermerd
-      --omitAttributeKeys             omit the attribute keys (PK, FK)
+      --omitAttributeKeys             omit the attribute keys (PK, FK, UK)
       --omitConstraintLabels          omit the constraint labels
   -o, --outputFileName string         output file name (default "result.mmd")
       --runConfig string              run configuration (replaces global configuration)

diff --git a/test/db-table-setup.sql b/test/db-table-setup.sql
@@ -21,7 +21,7 @@ create table article_comment
 create table label
 (
     id    int          not null primary key,
-    label varchar(255) not null
+    label varchar(255) not null unique
 );
 
 create table article_label