wilsonehusin · jcarnu · Apr 30, 2024 · Apr 30, 2024 · wilsonehusin · May 1, 2024
diff --git a/config/config.go b/config/config.go
@@ -29,6 +29,7 @@ type Documentation struct {
 	Directory string `json:"directory"`
 	Filename  string `json:"filename"`
 	Stdout    bool   `json:"stdout"`
+	Mermaid   bool   `json:"mermaid"`
 }
 
 func Default() *Config {

diff --git a/config/example.yaml b/config/example.yaml
@@ -20,3 +20,5 @@ documentation:
   filename: "schema.md"
   # Write the output to STDOUT as well.
   stdout: true
+  # Write mermaid output
+  mermaid: true
diff --git a/db/db.go b/db/db.go
@@ -11,13 +11,43 @@ type Database interface {
 
 	ListTables(ctx context.Context, schema string) ([]string, error)
 	ListColumns(ctx context.Context, schema, table string) ([]TableColumn, error)
+	ListConstraints(ctx context.Context, schema, table string) ([]TableConstraint, error)
 }
 
 type TableColumn struct {
 	Name     string
 	Type     string
 	Nullable bool
 	Default  string
+	PK       bool
+	FK       string
+	FK_ref   string
+	UK       bool
+}
+
+/*
+	tc.constraint_schema
+	, tc.constraint_name
+	, tc.constraint_type
+	, kcu1.constraint_schema
+	, kcu1.constraint_name
+	, kcu1.table_name from_table
+	, kcu1.column_name from_column
+	, kcu1.ordinal_position from_position
+	, kcu2.constraint_schema
+	, kcu2.constraint_name
+	, kcu2.table_name to_table
+	, kcu2.column_name to_column
+	, kcu2.ordinal_position to_position
+*/
+type TableConstraint struct {
+	Schema       string
+	Name         string
+	Type         string
+	TableColumn  string // referring column
+	TargetSchema string // referred schema
+	TargetTable  string // referred table
+	TargetColumn string // referred column
 }
 
 func New(url string) (Database, error) {

diff --git a/db/postgres.go b/db/postgres.go
@@ -56,8 +56,35 @@ func (p *Postgres) ListTables(ctx context.Context, schema string) ([]string, err
 	return tables, nil
 }
 
-const pgListColumns = `SELECT column_name, data_type, is_nullable, column_default FROM information_schema.columns WHERE table_schema = $1 AND table_name = $2`
+const pgListColumns = `SELECT
+	c.column_name
+	, c.data_type
+	, c.is_nullable
+	, c.column_default
 
+FROM information_schema.columns c
+WHERE c.table_schema = $1 AND c.table_name = $2
+ORDER BY ordinal_position`
+
+/*
+const pgListColumns = `
+SELECT
+	c.column_name
+	, c.data_type
+	, c.is_nullable
+	, c.column_default
+	, tc.constraint_type,
+	, *
+FROM information_schema.columns c
+LEFT OUTER JOIN information_schema.key_column_usage k
+	ON c.column_name = k.column_name
+	AND c.table_name = k.table_name
+LEFT OUTER JOIN information_schema.table_constraints tc
+	ON tc.constraint_name = k.constraint_name
+WHERE c.table_name = $1
+	AND c.table_schema = $2
+`
+*/
 func (p *Postgres) ListColumns(ctx context.Context, schema, table string) ([]TableColumn, error) {
 	rows, err := p.db.QueryContext(ctx, pgListColumns, schema, table)
 	if err != nil {
@@ -86,3 +113,75 @@ func (p *Postgres) ListColumns(ctx context.Context, schema, table string) ([]Tab
 	}
 	return columns, nil
 }
+
+const pgListConstraints = `
+SELECT
+	  tc.constraint_schema
+	, tc.constraint_name
+	, tc.constraint_type
+--	, kcu1.constraint_schema
+--	, kcu1.constraint_name
+--	, kcu1.table_name from_table
+	, kcu1.column_name from_column
+--	, kcu1.ordinal_position from_position
+--	, kcu2.constraint_schema
+--	, kcu2.constraint_name
+	, kcu2.table_schema to_schema
+	, kcu2.table_name to_table
+	, kcu2.column_name to_column
+--  , kcu2.ordinal_position to_position
+FROM information_schema.table_constraints tc
+LEFT OUTER JOIN information_schema.referential_constraints rc
+  	USING(constraint_schema, constraint_name)
+JOIN information_schema.key_column_usage as kcu1
+    ON  kcu1.constraint_catalog = tc.constraint_catalog
+    AND kcu1.constraint_schema  = tc.constraint_schema
+    AND kcu1.constraint_name    = tc.constraint_name
+LEFT OUTER JOIN information_schema.key_column_usage as kcu2
+    ON kcu2.constraint_catalog = rc.unique_constraint_catalog
+    AND kcu2.constraint_schema = rc.unique_constraint_schema
+    AND kcu2.constraint_name   = rc.unique_constraint_name
+    AND kcu2.ordinal_position  = kcu1.ordinal_position
+WHERE   tc.table_name = $2
+	AND tc.table_schema = $1
+`
+
+func (p *Postgres) ListConstraints(ctx context.Context, schema, table string) ([]TableConstraint, error) {
+	rows, err := p.db.QueryContext(ctx, pgListConstraints, schema, table)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	var constraints []TableConstraint
+	for rows.Next() {
+		var c TableConstraint
+		var tblColumn sql.NullString
+		var tgtSchema sql.NullString
+		var tgtTable sql.NullString
+		var tgtColumn sql.NullString
+		if err := rows.Scan(&c.Schema, &c.Name, &c.Type, &tblColumn, &tgtSchema, &tgtTable, &tgtColumn); err != nil {
+			return nil, err
+		}
+		if tblColumn.Valid {
+			c.TableColumn = tblColumn.String
+		}
+		if tgtSchema.Valid {
+			c.TargetSchema = tgtSchema.String
+		}
+		if tgtTable.Valid {
+			c.TargetTable = tgtTable.String
+		}
+		if tgtColumn.Valid {
+			c.TargetColumn = tgtColumn.String
+		}
+
+		constraints = append(constraints, c)
+	}
+	if err := rows.Close(); err != nil {
+		return nil, err
+	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+	return constraints, nil
+}
diff --git a/md/md.go b/md/md.go
@@ -7,6 +7,7 @@ import (
 	"os"
 	"path/filepath"
 	"slices"
+	"strings"
 
 	"github.com/charmbracelet/glamour"
 	"github.com/nao1215/markdown"
@@ -100,30 +101,138 @@ func (m *MD) renderTable(ctx context.Context, w io.Writer, schema, tableName str
 	canvas.H1(markdown.Code(tableName))
 	canvas.PlainText("")
 	mdTableRows := [][]string{}
+	constraints, err := m.db.ListConstraints(ctx, schema, tableName)
+
+	if err != nil {
+		return fmt.Errorf("fetching constraints for '%s.%s': %w", schema, tableName, err)
+	}
+
 	columns, err := m.db.ListColumns(ctx, schema, tableName)
 	if err != nil {
 		return fmt.Errorf("fetching columns for '%s.%s': %w", schema, tableName, err)
 	}
-	for _, column := range columns {
+	for idx := range columns {
 		nullable := ""
-		if !column.Nullable {
+		if !columns[idx].Nullable {
 			nullable = "NOT NULL"
 		}
-		d := column.Default
+		d := columns[idx].Default
 		if d != "" {
 			d = markdown.Code(d)
 		}
+		c := []string{}
+		t := ""
+		for _, constraint := range constraints {
+			if constraint.TableColumn == columns[idx].Name {
+				if constraint.Type == "PRIMARY KEY" {
+					c = append(c, "PK")
+					columns[idx].PK = true
+				}
+				if constraint.Type == "UNIQUE" {
+					c = append(c, "UK")
+					columns[idx].UK = true
+				}
+				if constraint.Type == "FOREIGN KEY" {
+					c = append(c, "FK")
+					t = fmt.Sprintf("%s.%s", constraint.TargetSchema, constraint.TargetTable)
+					columns[idx].FK = t
+					t = fmt.Sprintf("%s.%s", t, constraint.TargetColumn)
+					columns[idx].FK_ref = constraint.TargetColumn
+
+				}
+			}
+
+		}
 		mdTableRows = append(mdTableRows, []string{
-			markdown.Code(column.Name),
-			column.Type,
+			markdown.Code(columns[idx].Name),
+			columns[idx].Type,
 			nullable,
 			d,
+			strings.Join(c[:], ", "),
+			t,
 		})
 	}
 	canvas.Table(markdown.TableSet{
-		Header: []string{"Name", "Type", "Nullable", "Default"},
+		Header: []string{"Name", "Type", "Nullable", "Default", "Key", "Target"},
 		Rows:   mdTableRows,
 	})
+
+	if m.conf.Documentation.Mermaid {
+		mermaid_dep := ""
+
+		mermaid := "```mermaid\nerDiagram\n"
+		mermaid = fmt.Sprintf("%s\"%s.%s\"{\n", mermaid, schema, tableName)
+		for _, column := range columns {
+			// nullable := ""
+			// if !column.Nullable {
+			// 	nullable = "NN"
+			// }
+			k := []string{}
+			if column.PK {
+				k = append(k, "PK")
+			}
+			if column.FK != "" {
+				k = append(k, "FK")
+				mermaid_dep = fmt.Sprintf("%s\n\"%s.%s\" o|--o| \"%s\":%s",
+					mermaid_dep,
+					schema, tableName,
+					column.FK,
+					column.FK_ref)
+			}
+
+			if column.UK {
+				k = append(k, "UK")
+			}
+
+			mermaid = fmt.Sprintf("%s%s %s %s\n", mermaid,
+				strings.Replace(column.Type, " ", "_", -1),
+				column.Name, strings.Join(k[:], ","))
+		}
+
+		mermaid = fmt.Sprintf("%s}\n%s\n```", mermaid, mermaid_dep)
+
+		canvas.PlainText(mermaid)
+	}
+
+	if len(constraints) > 0 {
+		canvas.PlainText("")
+		canvas.H2("Constraints")
+		canvas.PlainText("")
+		mdConstraintsRows := [][]string{}
+		for _, constraint := range constraints {
+
+			col := ""
+			if constraint.TableColumn != "" {
+				col = markdown.Code(constraint.TableColumn)
+			}
+			tgtsch := ""
+			if constraint.TargetSchema != "" {
+				tgtsch = markdown.Code(constraint.TargetSchema)
+			}
+			tgttbl := ""
+			if constraint.TargetTable != "" {
+				tgttbl = markdown.Code(constraint.TargetTable)
+			}
+			tgtcol := ""
+			if constraint.TargetColumn != "" {
+				tgtcol = markdown.Code(constraint.TargetColumn)
+			}
+			mdConstraintsRows = append(mdConstraintsRows, []string{
+				markdown.Code(constraint.Schema),
+				markdown.Code(constraint.Name),
+				markdown.Code(constraint.Type),
+				col,
+				tgtsch,
+				tgttbl,
+				tgtcol,
+			})
+		}
+		canvas.Table(markdown.TableSet{
+			Header: []string{"Schema", "Name", "Type", "Referring Column", "Target Schema", "Target Table", "Target Column"},
+			Rows:   mdConstraintsRows,
+		})
+	}
+
 	if m.conf.Documentation.Stdout {
 		out, err := glamour.Render(canvas.String(), "dark")
 		if err != nil {