Impl ErdGenerator

sue445 · Dec 3, 2019 · fa3377d · fa3377d
1 parent d7aceee
commit fa3377d
Show file tree

Hide file tree

Showing 4 changed files with 241 additions and 0 deletions.
diff --git a/adapter/adapter.go b/adapter/adapter.go
@@ -0,0 +1,9 @@
+package adapter
+
+import "github.com/sue445/plant_erd/db"
+
+// Adapter represents database adapter
+type Adapter interface {
+	GetAllTableNames() ([]string, error)
+	GetTable(tableName string) (*db.Table, error)
+}
diff --git a/db/schema.go b/db/schema.go
@@ -11,6 +11,11 @@ type Schema struct {
 	Tables []*Table
 }
 
+// NewSchema returns a new Schema instance
+func NewSchema(tables []*Table) *Schema {
+	return &Schema{Tables: tables}
+}
+
 // ToErd returns ERD formatted schema
 func (s *Schema) ToErd() string {
 	var lines []string
@@ -37,3 +42,28 @@ func (s *Schema) SurroundingTablesWithin(tableName string, distance int) []strin
 	explorer := NewSchemaExplorer(s)
 	return explorer.Explore(tableName, distance)
 }
+
+// Subset returns subset of a schema
+func (s *Schema) Subset(tableName string, distance int) *Schema {
+	tableNames := s.SurroundingTablesWithin(tableName, distance)
+
+	var tables []*Table
+	for _, tableName := range tableNames {
+		table := s.findTable(tableName)
+
+		if table != nil {
+			tables = append(tables, table)
+		}
+	}
+
+	return NewSchema(tables)
+}
+
+func (s *Schema) findTable(tableName string) *Table {
+	for _, table := range s.Tables {
+		if table.Name == tableName {
+			return table
+		}
+	}
+	return nil
+}
diff --git a/erd_generator.go b/erd_generator.go
@@ -0,0 +1,41 @@
+package main
+
+import (
+	"fmt"
+	"github.com/sue445/plant_erd/db"
+	"io/ioutil"
+	"os"
+)
+
+// ErdGenerator represents ERD generator
+type ErdGenerator struct {
+	Filepath string
+	Table    string
+	Distance int
+}
+
+// Run performs generator
+func (g *ErdGenerator) Run(schema *db.Schema) error {
+	erd := g.generate(schema)
+	return g.outputErd(erd)
+}
+
+func (g *ErdGenerator) generate(schema *db.Schema) string {
+	if g.Table == "" || g.Distance <= 0 {
+		return schema.ToErd()
+	}
+
+	subset := schema.Subset(g.Table, g.Distance)
+	return subset.ToErd()
+}
+
+func (g *ErdGenerator) outputErd(content string) error {
+	if g.Filepath == "" {
+		// Print to stdout
+		fmt.Fprint(os.Stdout, content)
+		return nil
+	}
+
+	// Output to file
+	return ioutil.WriteFile(g.Filepath, []byte(content), 0644)
+}
diff --git a/erd_generator_test.go b/erd_generator_test.go
@@ -0,0 +1,161 @@
+package main
+
+import (
+	"bytes"
+	"github.com/stretchr/testify/assert"
+	"github.com/sue445/plant_erd/db"
+	"io"
+	"io/ioutil"
+	"os"
+	"path/filepath"
+	"testing"
+)
+
+func TestErdGenerator_generate(t *testing.T) {
+	tables := []*db.Table{
+		{
+			Name: "articles",
+			Columns: []*db.Column{
+				{
+					Name:       "id",
+					Type:       "integer",
+					NotNull:    true,
+					PrimaryKey: true,
+				},
+				{
+					Name:    "user_id",
+					Type:    "integer",
+					NotNull: true,
+				},
+			},
+			ForeignKeys: []*db.ForeignKey{
+				{
+					Sequence:   0,
+					FromColumn: "user_id",
+					ToTable:    "users",
+					ToColumn:   "id",
+				},
+			},
+		},
+		{
+			Name: "users",
+			Columns: []*db.Column{
+				{
+					Name:       "id",
+					Type:       "integer",
+					NotNull:    true,
+					PrimaryKey: true,
+				},
+				{
+					Name: "name",
+					Type: "text",
+				},
+			},
+		},
+	}
+	schema := db.NewSchema(tables)
+
+	type fields struct {
+		Filepath string
+		Table    string
+		Distance int
+	}
+	type args struct {
+		schema *db.Schema
+	}
+	tests := []struct {
+		name   string
+		fields fields
+		args   args
+	}{
+		{
+			name: "no table",
+			fields: fields{
+				Table:    "",
+				Distance: 0,
+			},
+			args: args{
+				schema: schema,
+			},
+		},
+		{
+			name: "with table and distance",
+			fields: fields{
+				Table:    "users",
+				Distance: 1,
+			},
+			args: args{
+				schema: schema,
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			g := &ErdGenerator{
+				Filepath: tt.fields.Filepath,
+				Table:    tt.fields.Table,
+				Distance: tt.fields.Distance,
+			}
+			got := g.generate(tt.args.schema)
+			assert.Greater(t, len(got), 0)
+		})
+	}
+}
+
+func TestErdGenerator_outputErd_ToFile(t *testing.T) {
+	dir, err := ioutil.TempDir("", "example")
+
+	if err != nil {
+		panic(err)
+	}
+
+	defer os.RemoveAll(dir)
+
+	filePath := filepath.Join(dir, "erd.txt")
+	g := &ErdGenerator{
+		Filepath: filePath,
+	}
+
+	g.outputErd("aaa")
+
+	data, err := ioutil.ReadFile(filePath)
+
+	if assert.NoError(t, err) {
+		str := string(data)
+		assert.Equal(t, "aaa", str)
+	}
+}
+
+// c.f. https://qiita.com/kami_zh/items/ff636f15da87dabebe6c
+func captureStdout(f func()) string {
+	r, w, err := os.Pipe()
+	if err != nil {
+		panic(err)
+	}
+
+	stdout := os.Stdout
+	os.Stdout = w
+
+	f()
+
+	os.Stdout = stdout
+	w.Close()
+
+	var buf bytes.Buffer
+	io.Copy(&buf, r)
+
+	return buf.String()
+}
+
+func TestErdGenerator_outputErd_ToStdout(t *testing.T) {
+	g := &ErdGenerator{
+		Filepath: "",
+	}
+
+	str := captureStdout(func() {
+		err := g.outputErd("aaa")
+		assert.NoError(t, err)
+	})
+
+	assert.Equal(t, "aaa", str)
+}