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

.go-version

@@ -0,0 +1 @@
+1.19.6

config/config.go

@@ -18,6 +18,7 @@ const (
 	UseAllSchemasKey               = "useAllSchemas"
 	ShowSchemaPrefix               = "showSchemaPrefix"
 	SchemaPrefixSeparator          = "schemaPrefixSeparator"
+	RelationshipLabelsKey          = "relationshipLabels"
 )
 
 type config struct{}
@@ -38,6 +39,7 @@ type MermerdConfig interface {
 	UseAllSchemas() bool
 	ShowSchemaPrefix() bool
 	SchemaPrefixSeparator() string
+	RelationshipLabels() []RelationshipLabel
 }
 
 func NewConfig() MermerdConfig {
@@ -72,6 +74,11 @@ func (c config) SelectedTables() []string {
 	return viper.GetStringSlice(SelectedTablesKey)
 }
 
+func (c config) RelationshipLabels() []RelationshipLabel {
+	labels := viper.GetStringSlice(RelationshipLabelsKey)
+	return parseLabels(labels)
+}
+
 func (c config) EncloseWithMermaidBackticks() bool {
 	return viper.GetBool(EncloseWithMermaidBackticksKey)
 }

config/config_test.go

@@ -2,9 +2,10 @@ package config
 
 import (
 	"bytes"
+	"testing"
+
 	"github.com/spf13/viper"
 	"github.com/stretchr/testify/assert"
-	"testing"
 )
 
 func TestYamlConfig(t *testing.T) {
@@ -32,6 +33,10 @@ showDescriptions:
   - enumValues
   - columnComments
   - notNull
+relationshipLabels:
+  - "schema.table1 schema.table2 : is_a"
+  - "table-name  another-table-name  : has_many"
+  - "incorrect format"
 useAllSchemas: true
 showSchemaPrefix: true
 schemaPrefixSeparator: "_"
@@ -63,4 +68,19 @@ connectionStringSuggestions:
 	assert.True(t, config.UseAllSchemas())
 	assert.True(t, config.ShowSchemaPrefix())
 	assert.Equal(t, "_", config.SchemaPrefixSeparator())
+	assert.ElementsMatch(t,
+		config.RelationshipLabels(),
+		[]RelationshipLabel{
+			RelationshipLabel{
+				PkName: "schema.table1",
+				FkName: "schema.table2",
+				Label:  "is_a",
+			},
+			RelationshipLabel{
+				PkName: "table-name",
+				FkName: "another-table-name",
+				Label:  "has_many",
+			},
+		},
+	)
 }

config/relationship_label.go

@@ -0,0 +1,60 @@
+package config
+
+import (
+	"errors"
+	"regexp"
+	"strings"
+
+	"github.com/sirupsen/logrus"
+)
+
+type RelationshipLabel struct {
+	PkName string
+	FkName string
+	Label  string
+}
+
+func parseLabels(labels []string) []RelationshipLabel {
+	var relationshipLabels []RelationshipLabel
+	for _, label := range labels {
+		parsed, err := parseLabel(label)
+		if err != nil {
+			logrus.Warnf("label '%s' is not in the correct format", label)
+			continue
+		}
+		relationshipLabels = append(relationshipLabels, parsed)
+	}
+	return relationshipLabels
+}
+
+func parseLabel(label string) (RelationshipLabel, error) {
+	label = strings.Trim(label, " \t")
+	matched, groups := match(label)
+	if !matched {
+		return RelationshipLabel{}, errors.New("invalid relationship label")
+	}
+
+	return RelationshipLabel{
+		PkName: string(groups[1]),
+		FkName: string(groups[2]),
+		Label:  string(groups[3]),
+	}, nil
+}
+
+// The regex works by creating three capture groups
+// Each group allows for all word characters, `.`, `_` and `-` any number of times
+// The first two groups (the table names) are separated by any amount of whitespace characters
+// The table names and label are are separated by
+// - any number of whitespace characters
+// - a `:`
+// - and then any other number of whitespace characters
+// The string must start with the first table name and it must end with the label
+var labelRegex = regexp.MustCompile(`^([\w\._-]+)[\s]+([\w\._-]+)[\s]+:[\s]+([\w._-]+)$`)
+
+func match(label string) (bool, [][]byte) {
+	groups := labelRegex.FindSubmatch([]byte(label))
+	if groups == nil {
+		return false, [][]byte{}
+	}
+	return true, groups
+}

diagram/diagram.go

@@ -59,12 +59,13 @@ func (d diagram) Create(result *database.Result) error {
 	}
 
 	var constraints []ErdConstraintData
+	relationshipLabelMap := BuildRelationshipLabelMapFromConfig(d.config)
 	for _, constraint := range allConstraints {
 		if shouldSkipConstraint(d.config, tableData, constraint) {
 			continue
 		}
 
-		constraints = append(constraints, getConstraintData(d.config, constraint))
+		constraints = append(constraints, getConstraintData(d.config, relationshipLabelMap, constraint))
 	}
 
 	diagramData := ErdDiagramData{

diagram/diagram_util.go

@@ -2,9 +2,10 @@ package diagram
 
 import (
 	"fmt"
-	"github.com/sirupsen/logrus"
 	"strings"
 
+	"github.com/sirupsen/logrus"
+
 	"github.com/KarnerTh/mermerd/config"
 	"github.com/KarnerTh/mermerd/database"
 )
@@ -90,15 +91,21 @@ func shouldSkipConstraint(config config.MermerdConfig, tables []ErdTableData, co
 	return !(tableNameInSlice(tables, constraint.PkTable) && tableNameInSlice(tables, constraint.FkTable))
 }
 
-func getConstraintData(config config.MermerdConfig, constraint database.ConstraintResult) ErdConstraintData {
+func getConstraintData(config config.MermerdConfig, labelMap RelationshipLabelMap, constraint database.ConstraintResult) ErdConstraintData {
+	pkTableName := getTableName(config, database.TableDetail{Schema: constraint.PkSchema, Name: constraint.PkTable})
+	fkTableName := getTableName(config, database.TableDetail{Schema: constraint.FkSchema, Name: constraint.FkTable})
+
 	constraintLabel := constraint.ColumnName
 	if config.OmitConstraintLabels() {
 		constraintLabel = ""
 	}
+	if relationshipLabel, found := labelMap.LookupRelationshipLabel(pkTableName, fkTableName); found {
+		constraintLabel = relationshipLabel.Label
+	}
 
 	return ErdConstraintData{
-		PkTableName:     getTableName(config, database.TableDetail{Schema: constraint.PkSchema, Name: constraint.PkTable}),
-		FkTableName:     getTableName(config, database.TableDetail{Schema: constraint.FkSchema, Name: constraint.FkTable}),
+		PkTableName:     pkTableName,
+		FkTableName:     fkTableName,
 		Relation:        getRelation(constraint),
 		ConstraintLabel: constraintLabel,
 	}

diagram/diagram_util_test.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/stretchr/testify/assert"
 
+	"github.com/KarnerTh/mermerd/config"
 	"github.com/KarnerTh/mermerd/database"
 	"github.com/KarnerTh/mermerd/mocks"
 )
@@ -284,20 +285,87 @@ func TestShouldSkipConstraint(t *testing.T) {
 }
 
 func TestGetConstraintData(t *testing.T) {
+	t.Run("The column name is used as the constraint label", func(t *testing.T) {
+		// Arrange
+		configMock := mocks.MermerdConfig{}
+		configMock.On("OmitConstraintLabels").Return(false).Once()
+		configMock.On("ShowSchemaPrefix").Return(false).Twice()
+		constraint := database.ConstraintResult{ColumnName: "Column1"}
+
+		// Act
+		result := getConstraintData(&configMock, &relationshipLabelMap{}, constraint)
+
+		// Assert
+		configMock.AssertExpectations(t)
+		assert.Equal(t, result.ConstraintLabel, "Column1")
+	})
 	t.Run("OmitConstraintLabels should remove the constraint label", func(t *testing.T) {
 		// Arrange
 		configMock := mocks.MermerdConfig{}
 		configMock.On("OmitConstraintLabels").Return(true).Once()
 		configMock.On("ShowSchemaPrefix").Return(false).Twice()
+
 		constraint := database.ConstraintResult{ColumnName: "Column1"}
 
 		// Act
-		result := getConstraintData(&configMock, constraint)
+		result := getConstraintData(&configMock, &relationshipLabelMap{}, constraint)
 
 		// Assert
 		configMock.AssertExpectations(t)
 		assert.Equal(t, result.ConstraintLabel, "")
 	})
+	t.Run("If a relationship label exists, it should be used", func(t *testing.T) {
+		// Arrange
+		configMock := mocks.MermerdConfig{}
+		configMock.On("OmitConstraintLabels").Return(true).Once()
+		configMock.On("ShowSchemaPrefix").Return(false).Twice()
+
+		labelsMap := &relationshipLabelMap{}
+		labelsMap.AddRelationshipLabel(config.RelationshipLabel{
+			PkName: "pk",
+			FkName: "fk",
+			Label:  "relationship-label",
+		})
+
+		constraint := database.ConstraintResult{
+			PkTable:    "pk",
+			FkTable:    "fk",
+			ColumnName: "Column1",
+		}
+
+		// Act
+		result := getConstraintData(&configMock, labelsMap, constraint)
+
+		// Assert
+		configMock.AssertExpectations(t)
+		assert.Equal(t, result.ConstraintLabel, "relationship-label")
+	})
+	t.Run("If a relationship label exists, it should be used even if we don't omit constraint labels", func(t *testing.T) {
+		// Arrange
+		configMock := mocks.MermerdConfig{}
+		configMock.On("OmitConstraintLabels").Return(false).Once()
+		configMock.On("ShowSchemaPrefix").Return(false).Twice()
+
+		labelsMap := &relationshipLabelMap{}
+		labelsMap.AddRelationshipLabel(config.RelationshipLabel{
+			PkName: "pk",
+			FkName: "fk",
+			Label:  "relationship-label",
+		})
+
+		constraint := database.ConstraintResult{
+			PkTable:    "pk",
+			FkTable:    "fk",
+			ColumnName: "Column1",
+		}
+
+		// Act
+		result := getConstraintData(&configMock, labelsMap, constraint)
+
+		// Assert
+		configMock.AssertExpectations(t)
+		assert.Equal(t, result.ConstraintLabel, "relationship-label")
+	})
 }
 
 func TestGetTableName(t *testing.T) {

diagram/relationship_label_map.go

@@ -0,0 +1,49 @@
+package diagram
+
+import (
+	"fmt"
+
+	"github.com/KarnerTh/mermerd/config"
+)
+
+type RelationshipLabelMap interface {
+	AddRelationshipLabel(label config.RelationshipLabel)
+	LookupRelationshipLabel(pkName, fkName string) (label config.RelationshipLabel, found bool)
+}
+
+type relationshipLabelMap struct {
+	mapping map[string]config.RelationshipLabel
+}
+
+func (r *relationshipLabelMap) AddRelationshipLabel(label config.RelationshipLabel) {
+	if r.mapping == nil {
+		r.mapping = make(map[string]config.RelationshipLabel)
+	}
+	key := r.buildMapKey(label.PkName, label.FkName)
+	r.mapping[key] = label
+}
+
+func (r *relationshipLabelMap) LookupRelationshipLabel(pkName, fkName string) (label config.RelationshipLabel, found bool) {
+	if r.mapping == nil {
+		return config.RelationshipLabel{}, false
+	}
+	key := r.buildMapKey(pkName, fkName)
+	label, found = r.mapping[key]
+	return
+}
+
+func (r *relationshipLabelMap) buildMapKey(pkName, fkName string) string {
+	return fmt.Sprintf("%s-%s", pkName, fkName)
+}
+
+func BuildRelationshipLabelMapFromConfig(c config.MermerdConfig) RelationshipLabelMap {
+	return BuildRelationshipLabelMap(c.RelationshipLabels())
+}
+
+func BuildRelationshipLabelMap(labels []config.RelationshipLabel) RelationshipLabelMap {
+	labelMap := &relationshipLabelMap{}
+	for _, label := range labels {
+		labelMap.AddRelationshipLabel(label)
+	}
+	return labelMap
+}

diagram/relationship_label_map_test.go

@@ -0,0 +1,33 @@
+package diagram_test
+
+import (
+	"testing"
+
+	"github.com/KarnerTh/mermerd/config"
+	"github.com/KarnerTh/mermerd/diagram"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestEmptyRelationshipMapDoesNotError(t *testing.T) {
+	relationshipMap := diagram.BuildRelationshipLabelMap([]config.RelationshipLabel{})
+
+	_, found := relationshipMap.LookupRelationshipLabel("pk", "fk")
+
+	assert.False(t, found)
+}
+
+func TestRelationshipMapCanAddAndLookupLabel(t *testing.T) {
+	relationshipMap := diagram.BuildRelationshipLabelMap([]config.RelationshipLabel{})
+
+	exampleLabel := config.RelationshipLabel{
+		PkName: "name",
+		FkName: "another-name",
+		Label:  "a-label",
+	}
+	relationshipMap.AddRelationshipLabel(exampleLabel)
+
+	actual, found := relationshipMap.LookupRelationshipLabel("name", "another-name")
+
+	assert.True(t, found)
+	assert.Equal(t, actual, exampleLabel)
+}

exampleRunConfig.yaml

@@ -24,5 +24,7 @@ showDescriptions:
   - enumValues
   - columnComments
   - notNull
+relationshipLabels:
+  - "public_article public_article_comment : has_many"
 showSchemaPrefix: true
 schemaPrefixSeparator: "_"