feat:func SetViewpointsToTables

k1LoW · Oct 26, 2023 · 4235f6d · 4235f6d
1 parent 5059fd9
commit 4235f6d
Show file tree

Hide file tree

Showing 2 changed files with 86 additions and 13 deletions.
diff --git a/schema/schema.go b/schema/schema.go
@@ -131,19 +131,26 @@ type Column struct {
 	HideForER       bool            `json:"-"`
 }
 
+type TableViewpoint struct {
+	Index int    `json:"index"`
+	Name  string `json:"name"`
+	Desc  string `json:"desc"`
+}
+
 // Table is the struct for database table
 type Table struct {
-	Name             string        `json:"name"`
-	Type             string        `json:"type"`
-	Comment          string        `json:"comment"`
-	Columns          []*Column     `json:"columns"`
-	Indexes          []*Index      `json:"indexes"`
-	Constraints      []*Constraint `json:"constraints"`
-	Triggers         []*Trigger    `json:"triggers"`
-	Def              string        `json:"def"`
-	Labels           Labels        `json:"labels,omitempty"`
-	ReferencedTables []*Table      `json:"referenced_tables,omitempty" yaml:"referencedTables,omitempty"`
-	External         bool          `json:"-"` // Table external to the schema
+	Name             string            `json:"name"`
+	Type             string            `json:"type"`
+	Comment          string            `json:"comment"`
+	Columns          []*Column         `json:"columns"`
+	Viewpoints       []*TableViewpoint `json:"viewpoints"`
+	Indexes          []*Index          `json:"indexes"`
+	Constraints      []*Constraint     `json:"constraints"`
+	Triggers         []*Trigger        `json:"triggers"`
+	Def              string            `json:"def"`
+	Labels           Labels            `json:"labels,omitempty"`
+	ReferencedTables []*Table          `json:"referenced_tables,omitempty" yaml:"referencedTables,omitempty"`
+	External         bool              `json:"-"` // Table external to the schema
 }
 
 // Relation is the struct for table relation
@@ -192,6 +199,24 @@ type Schema struct {
 	Viewpoints Viewpoints  `json:"viewpoints,omitempty"`
 }
 
+func (s *Schema) SetViewpointsToTables() (*Schema, error) {
+	for vi, v := range s.Viewpoints {
+		// Add viewpoints to table
+		for _, t := range v.Tables {
+			table, err := s.FindTableByName(t)
+			if err != nil {
+				return s, err
+			}
+			table.Viewpoints = append(table.Viewpoints, &TableViewpoint{
+				Index: vi,
+				Name:  v.Name,
+				Desc:  v.Desc,
+			})
+		}
+	}
+	return s, nil
+}
+
 func (s *Schema) NormalizeTableName(name string) string {
 	if s.Driver != nil && s.Driver.Meta != nil && s.Driver.Meta.CurrentSchema != "" && (s.Driver.Name == "postgres" || s.Driver.Name == "redshift") && !strings.Contains(name, ".") {
 		return fmt.Sprintf("%s.%s", s.Driver.Meta.CurrentSchema, name)

diff --git a/schema/schema_test.go b/schema/schema_test.go
@@ -11,6 +11,54 @@ import (
 	"github.com/google/go-cmp/cmp"
 )
 
+func TestSetViewpointsToTables(t *testing.T) {
+	viewpointAName := "va"
+	viewpointBName := "vb"
+
+	tests := []struct {
+		viewpointATables     []string
+		viewpointBTables     []string
+		wantTableAViewpoints []*TableViewpoint
+	}{
+		{[]string{"a"}, []string{"b"}, []*TableViewpoint{{Name: viewpointAName}}},
+		{[]string{"a", "b"}, []string{"a"}, []*TableViewpoint{{
+			Index: 0,
+			Name:  viewpointAName,
+		}, {
+			Index: 1,
+			Name:  viewpointBName,
+		}}},
+	}
+
+	for _, tt := range tests {
+		t.Run(fmt.Sprintf("%v", tt.viewpointATables), func(t *testing.T) {
+			fmt.Println(tt.viewpointATables)
+			s := newTestSchema(t)
+			s.Viewpoints = []*Viewpoint{
+				{
+					Name:   viewpointAName,
+					Tables: tt.viewpointATables,
+				},
+				{
+					Name:   viewpointBName,
+					Tables: tt.viewpointBTables,
+				},
+			}
+			result, err := s.SetViewpointsToTables()
+			if err != nil {
+				t.Error(err)
+			}
+			gotTable, _ := result.FindTableByName("a")
+			got := gotTable.Viewpoints
+			want := tt.wantTableAViewpoints
+
+			if diff := cmp.Diff(got, want, nil); diff != "" {
+				t.Errorf("%s", diff)
+			}
+		})
+	}
+}
+
 func TestNormalizeTableName(t *testing.T) {
 	tests := []struct {
 		s    *Schema
@@ -212,11 +260,11 @@ func TestSchema_Sort(t *testing.T) {
 		},
 		Functions: []*Function{
 			&Function{
-				Name:   "b",
+				Name:      "b",
 				Arguments: "arg b",
 			},
 			&Function{
-				Name:   "b",
+				Name:      "b",
 				Arguments: "arg a",
 			},
 		},