Merge pull request #601 from insano10/postgres-enums

Add support to output enum definitions in the db doc README
k1LoW · Oct 6, 2024 · 76d30f8 · 76d30f8
2 parents b85fa2f + 20b1e62
commit 76d30f8
Show file tree

Hide file tree

Showing 18 changed files with 184 additions and 0 deletions.
diff --git a/config/templates.go b/config/templates.go
@@ -15,6 +15,7 @@ type MD struct {
 	Index     string `yaml:"index,omitempty"`
 	Table     string `yaml:"table,omitempty"`
 	Viewpoint string `yaml:"viewpoint,omitempty"`
+	Enum      string `yaml:"enum,omitempty"`
 }
 
 // Dot holds the paths to the dot template files.

diff --git a/drivers/postgres/postgres.go b/drivers/postgres/postgres.go
@@ -329,6 +329,13 @@ ORDER BY tgrelid
 	}
 	s.Functions = functions
 
+	// Enums
+	enums, err := p.getEnums()
+	if err != nil {
+		return err
+	}
+	s.Enums = enums
+
 	s.Tables = tables
 
 	// Relations
@@ -490,6 +497,41 @@ func (p *Postgres) getFunctionsByQuery(query string) ([]*schema.Function, error)
 	return functions, nil
 }
 
+func (p *Postgres) getEnums() ([]*schema.Enum, error) {
+	enums := []*schema.Enum{}
+
+	enumsResult, err := p.db.Query(`SELECT n.nspname, t.typname AS enum_name, ARRAY_AGG(e.enumlabel) AS enum_values
+											FROM pg_type t, pg_enum e, pg_catalog.pg_namespace n
+											WHERE t.typcategory = 'E'
+											  AND t.oid = e.enumtypid
+											  AND n.oid = t.typnamespace
+											GROUP BY n.nspname, t.typname `)
+
+	if err != nil {
+		return nil, errors.WithStack(err)
+	}
+	defer enumsResult.Close()
+
+	for enumsResult.Next() {
+		var (
+			schemaName string
+			enumName   string
+			enumValues []string
+		)
+		err := enumsResult.Scan(&schemaName, &enumName, pq.Array(&enumValues))
+		if err != nil {
+			return enums, errors.WithStack(err)
+		}
+
+		enum := &schema.Enum{
+			Name:   fmt.Sprintf("%s.%s", schemaName, enumName),
+			Values: enumValues,
+		}
+		enums = append(enums, enum)
+	}
+	return enums, nil
+}
+
 func fullTableName(owner string, tableName string) string {
 	return fmt.Sprintf("%s.%s", owner, tableName)
 }

diff --git a/output/md/md.go b/output/md/md.go
@@ -7,6 +7,7 @@ import (
 	"io"
 	"os"
 	"path/filepath"
+	"sort"
 	"strconv"
 	"strings"
 	"text/template"
@@ -521,11 +522,15 @@ func (m *Md) makeSchemaTemplateData(s *schema.Schema) map[string]interface{} {
 	// Viewpoints
 	viewpointsData := m.viewpointsData(s.Viewpoints, number, adjust, showOnlyFirstParagraph)
 
+	// Enums
+	enumData := m.enumData(s.Enums)
+
 	return map[string]interface{}{
 		"Schema":     s,
 		"Tables":     tablesData,
 		"Functions":  functionsData,
 		"Viewpoints": viewpointsData,
+		"Enums":      enumData,
 	}
 }
 
@@ -873,6 +878,35 @@ func (m *Md) functionsData(functions []*schema.Function, number, adjust, showOnl
 	return data
 }
 
+func (m *Md) enumData(enums []*schema.Enum) [][]string {
+	data := [][]string{}
+
+	if len(enums) == 0 {
+		return data
+	}
+
+	header := []string{
+		m.config.MergedDict.Lookup("Name"),
+		m.config.MergedDict.Lookup("Values"),
+	}
+	headerLine := []string{"----", "-------"}
+	data = append(data,
+		header,
+		headerLine,
+	)
+
+	for _, e := range enums {
+		sort.Strings(e.Values)
+		d := []string{
+			e.Name,
+			strings.Join(e.Values, ", "),
+		}
+		data = append(data, d)
+	}
+
+	return data
+}
+
 func (m *Md) viewpointsData(viewpoints []*schema.Viewpoint, number, adjust, showOnlyFirstParagraph bool) [][]string {
 	data := [][]string{}
 	header := []string{

diff --git a/output/md/templates/index.md.tmpl b/output/md/templates/index.md.tmpl
@@ -31,6 +31,15 @@
 |{{ range $d := $t }} {{ $d | nl2br }} |{{ end }}
 {{- end -}}
 {{- end -}}
+
+{{ if ne (len .Enums) 0 }}
+
+## {{ "Enums" | lookup }}
+{{ range $t := .Enums }}
+|{{ range $d := $t }} {{ $d | nl2br }} |{{ end }}
+{{- end -}}
+{{- end -}}
+
 {{- if .er }}
 
 ## {{ "Relations" | lookup }}

diff --git a/sample/adjust/README.md b/sample/adjust/README.md
@@ -42,6 +42,12 @@ Sample PostgreSQL database document.
 | public.update_updated     | trigger    |                           | FUNCTION  |
 | public.reset_comment      | void       | IN comment_id integer     | PROCEDURE |
 
+## Enums
+
+| Name | Values |
+| ---- | ------- |
+| public.post_types | draft, private, public |
+
 ## Relations
 
 ![er](schema.svg)

diff --git a/sample/postgres/README.md b/sample/postgres/README.md
@@ -42,6 +42,12 @@ Sample PostgreSQL database document.
 | public.update_updated | trigger |  | FUNCTION |
 | public.reset_comment | void | IN comment_id integer | PROCEDURE |
 
+## Enums
+
+| Name | Values |
+| ---- | ------- |
+| public.post_types | draft, private, public |
+
 ## Relations
 
 ![er](schema.svg)

diff --git a/sample/postgres95/README.md b/sample/postgres95/README.md
@@ -41,6 +41,12 @@ Sample PostgreSQL database document.
 | public.uuid_generate_v5 | uuid | namespace uuid, name text | FUNCTION |
 | public.update_updated | trigger |  | FUNCTION |
 
+## Enums
+
+| Name | Values |
+| ---- | ------- |
+| public.post_types | draft, private, public |
+
 ## Relations
 
 ![er](schema.svg)

diff --git a/schema/json.go b/schema/json.go
@@ -19,6 +19,7 @@ func (s Schema) MarshalJSON() ([]byte, error) {
 		Tables     []*Table     `json:"tables"`
 		Relations  []*Relation  `json:"relations"`
 		Functions  []*Function  `json:"functions"`
+		Enums      []*Enum      `json:"enums,omitempty"`
 		Driver     *Driver      `json:"driver"`
 		Labels     Labels       `json:"labels,omitempty"`
 		Viewpoints []*Viewpoint `json:"viewpoints,omitempty"`
@@ -29,6 +30,7 @@ func (s Schema) MarshalJSON() ([]byte, error) {
 		Relations:  s.Relations,
 		Driver:     s.Driver,
 		Functions:  s.Functions,
+		Enums:      s.Enums,
 		Labels:     s.Labels,
 		Viewpoints: s.Viewpoints,
 	})
@@ -49,6 +51,17 @@ func (d Function) MarshalJSON() ([]byte, error) {
 	})
 }
 
+// MarshalJSON return custom JSON byte
+func (e Enum) MarshalJSON() ([]byte, error) {
+	return json.Marshal(&struct {
+		Name   string   `json:"name"`
+		Values []string `json:"values"`
+	}{
+		Name:   e.Name,
+		Values: e.Values,
+	})
+}
+
 // MarshalJSON return custom JSON byte
 func (d Driver) MarshalJSON() ([]byte, error) {
 	if d.Meta == nil {

diff --git a/schema/schema.go b/schema/schema.go
@@ -180,6 +180,11 @@ type Function struct {
 	Type       string `json:"type"`
 }
 
+type Enum struct {
+	Name   string   `json:"name"`
+	Values []string `json:"values"`
+}
+
 // Driver is the struct for tbls driver information
 type Driver struct {
 	Name            string      `json:"name"`
@@ -194,6 +199,7 @@ type Schema struct {
 	Tables     []*Table    `json:"tables"`
 	Relations  []*Relation `json:"relations"`
 	Functions  []*Function `json:"functions"`
+	Enums      []*Enum     `json:"enums,omitempty"`
 	Driver     *Driver     `json:"driver"`
 	Labels     Labels      `json:"labels,omitempty"`
 	Viewpoints Viewpoints  `json:"viewpoints,omitempty"`

diff --git a/testdata/json_output_schema.golden b/testdata/json_output_schema.golden
@@ -139,6 +139,16 @@
     }
   ],
   "functions": null,
+  "enums": [
+    {
+      "name": "enum",
+      "values": [
+        "one",
+        "two",
+        "three"
+      ]
+    }
+  ],
   "driver": {
     "name": "testdriver",
     "database_version": "1.0.0",

diff --git a/testdata/md_test_README.md.adjust.golden b/testdata/md_test_README.md.adjust.golden
@@ -17,6 +17,12 @@
 | [b](b.md)       | 2       | table b |      | `red` `green`  |
 | [view](view.md) | 1       | view    | VIEW |                |
 
+## Enums
+
+| Name | Values |
+| ---- | ------- |
+| enum | one, three, two |
+
 ---
 
 > Generated by [tbls](https://github.com/k1LoW/tbls)
diff --git a/testdata/md_test_README.md.first_para.golden b/testdata/md_test_README.md.first_para.golden
@@ -17,6 +17,12 @@
 | [b](b.md) | 2 | table b |  | `red` `green` |
 | [view](view.md) | 1 | view | VIEW |  |
 
+## Enums
+
+| Name | Values |
+| ---- | ------- |
+| enum | one, three, two |
+
 ## Relations
 
 ![er](schema.png)

diff --git a/testdata/md_test_README.md.golden b/testdata/md_test_README.md.golden
@@ -17,6 +17,12 @@
 | [b](b.md) | 2 | table b |  | `red` `green` |
 | [view](view.md) | 1 | view | VIEW |  |
 
+## Enums
+
+| Name | Values |
+| ---- | ------- |
+| enum | one, three, two |
+
 ## Relations
 
 ![er](schema.png)

diff --git a/testdata/md_test_README.md.mermaid.golden b/testdata/md_test_README.md.mermaid.golden
@@ -17,6 +17,12 @@
 | [b](b.md) | 2 | table b |  | `red` `green` |
 | [view](view.md) | 1 | view | VIEW |  |
 
+## Enums
+
+| Name | Values |
+| ---- | ------- |
+| enum | one, three, two |
+
 ## Relations
 
 ```mermaid

diff --git a/testdata/md_test_README.md.number.golden b/testdata/md_test_README.md.number.golden
@@ -17,6 +17,12 @@
 | 2 | [b](b.md) | 2 | table b |  | `red` `green` |
 | 3 | [view](view.md) | 1 | view | VIEW |  |
 
+## Enums
+
+| Name | Values |
+| ---- | ------- |
+| enum | one, three, two |
+
 ---
 
 > Generated by [tbls](https://github.com/k1LoW/tbls)
diff --git a/testdata/md_test_README.md.space_in_table_name.golden b/testdata/md_test_README.md.space_in_table_name.golden
@@ -17,6 +17,12 @@
 | [a b](a%20b.md) | 2 | table b |  | `red` `green` |
 | [view](view.md) | 1 | view | VIEW |  |
 
+## Enums
+
+| Name | Values |
+| ---- | ------- |
+| enum | one, three, two |
+
 ---
 
 > Generated by [tbls](https://github.com/k1LoW/tbls)
diff --git a/testdata/yaml_output_schema.golden b/testdata/yaml_output_schema.golden
@@ -92,6 +92,12 @@ relations:
   def: FOREIGN KEY (b) REFERENCES a(a)
   virtual: false
 functions: []
+enums:
+- name: enum
+  values:
+  - one
+  - two
+  - three
 driver:
   name: testdriver
   databaseVersion: 1.0.0

diff --git a/testutil/schema.go b/testutil/schema.go
@@ -14,6 +14,7 @@ func NewSchema(t *testing.T) *schema.Schema {
 		labelBlueName  = "blue"
 		labelRedName   = "red"
 		labelGreenName = "green"
+		enumName       = "enum"
 	)
 
 	labelBlue := &schema.Label{
@@ -108,6 +109,11 @@ func NewSchema(t *testing.T) *schema.Schema {
 		},
 	}
 
+	enum := &schema.Enum{
+		Name:   enumName,
+		Values: []string{"one", "two", "three"},
+	}
+
 	r := &schema.Relation{
 		Table:             tb,
 		Columns:           []*schema.Column{cb},
@@ -128,6 +134,9 @@ func NewSchema(t *testing.T) *schema.Schema {
 			tb,
 			tView,
 		},
+		Enums: []*schema.Enum{
+			enum,
+		},
 		Relations: []*schema.Relation{
 			r,
 		},