Added implementation of pgComment.

Author: vvuwei

database/pgComment.go

@@ -4,62 +4,79 @@ import (
 	"context"
 	"github.com/dipdup-net/go-lib/hasura"
 	"github.com/go-pg/pg/v10"
-	"github.com/pkg/errors"
 	"reflect"
 	"strings"
 )
 
 func makeComments(ctx context.Context, conn PgGoConnection, model interface{}) error {
-	typ := reflect.TypeOf(model)
-
-	// 1. go through fields
-	// 2. if tableName field -
-	//		2.1 read value from pg tab, if not exist from model type name with snake case convertion - remember as tableName
-	// 		2.2 read value from pg-comment tag, if not exist continue
-	//		2.3 set comment with SQL statement
-	// 3. other
-	//		3.1 read comment from pg-comment, it not exist continue
-	// 		3.2 read pg tag first value if it exist, if not - then snake case name of field and set as columnName
-	// 		3.3 set comment with SQL statement
-	for i := 0; i < typ.NumField(); i++ {
-		fieldType := typ.Field(i)
-		pgTag, ok := fieldType.Tag.Lookup("pg")
-		if !ok {
-			continue
-		}
+	modelType := reflect.TypeOf(model)
+	var tableName pg.Safe
 
-		tags := strings.Split(pgTag, ",")
+	for i := 0; i < modelType.NumField(); i++ {
+		fieldType := modelType.Field(i)
 
-		var name string
-		for i := range tags {
-			if i == 0 {
-				if name == "" {
-					name = hasura.ToSnakeCase(fieldType.Name)
-				} else {
-					name = tags[i]
-				}
-				continue
+		if fieldType.Name == "tableName" {
+
+			readFromModelName := true
+			tableName, readFromModelName = getPgName(fieldType)
+			if readFromModelName {
+				tableName = pg.Safe(hasura.ToSnakeCase(modelType.Name()))
 			}
 
-			parts := strings.Split(tags[i], ":")
-			if parts[0] == "comment" {
-				if len(parts) != 2 {
-					return errors.Errorf("invalid comments format: %s", pgTag)
-				}
-				if fieldType.Name == "tableName" {
-					// typ.Name() to
-					if _, err := conn.DB().ExecContext(ctx, `COMMENT ON TABLE ? IS ?`, pg.Safe(typ.Name()), parts[1]); err != nil {
-						return err
-					}
-				} else {
-					if _, err := conn.DB().ExecContext(ctx, `COMMENT ON COLUMN ?.? IS ?`, pg.Safe(typ.Name()), pg.Safe(name), parts[1]); err != nil {
-						return err
-					}
-				}
+			pgCommentTag, ok := getPgComment(fieldType)
+			if !ok {
 				continue
 			}
+
+			if _, err := conn.DB().ExecContext(ctx,
+				`COMMENT ON TABLE ? IS ?`,
+				tableName, pgCommentTag); err != nil {
+				return err
+			}
+
+			continue
+		}
+
+		pgCommentTag, ok := getPgComment(fieldType)
+		if !ok {
+			continue
+		}
+
+		columnName, readFromFieldName := getPgName(fieldType)
+		if readFromFieldName {
+			columnName = pg.Safe(hasura.ToSnakeCase(fieldType.Name))
+		}
+
+		if _, err := conn.DB().ExecContext(ctx,
+			`COMMENT ON COLUMN ?.? IS ?`,
+			tableName, columnName, pgCommentTag); err != nil {
+			return err
 		}
 	}
 
 	return nil
 }
+
+func getPgName(fieldType reflect.StructField) (name pg.Safe, ok bool) {
+	pgTag, ok := fieldType.Tag.Lookup("pg")
+	if ok {
+		tags := strings.Split(pgTag, ",")
+
+		if tags[0] != "" {
+			name = pg.Safe(tags[0])
+			ok = false
+		}
+	}
+
+	return name, ok
+}
+
+func getPgComment(fieldType reflect.StructField) (pg.Safe, bool) {
+	pgCommentTag, ok := fieldType.Tag.Lookup("pg-comment")
+
+	if ok {
+		return pg.Safe(pgCommentTag), ok
+	}
+
+	return "", false
+}

database/pgComment_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"github.com/go-pg/pg/v10"
 	"github.com/stretchr/testify/mock"
+	"sort"
 	"testing"
 )
 
@@ -25,16 +26,16 @@ type PgDBMock struct {
 	mock.Mock
 }
 
-func (db *PgDBMock) ExecContext(_ context.Context, query interface{}, params ...interface{}) (pg.Result, error) {
-	args := db.Called(query, params)
+func (db *PgDBMock) ExecContext(ctx context.Context, query interface{}, params ...interface{}) (pg.Result, error) {
+	args := db.Called(ctx, query, params)
 
 	return nil, args.Error(0)
 }
 
 func TestMakeCommentsWithTableName(t *testing.T) {
 	type Ballot struct {
 		//nolint
-		tableName struct{} `pg:"ballots"`
+		tableName struct{} `pg:"ballots" pg-comment:"Ballot table"`
 		Ballot    string   `json:"ballot"`
 	}
 
@@ -47,5 +48,26 @@ func TestMakeCommentsWithTableName(t *testing.T) {
 	makeComments(ctx, pgGo, model)
 
 	// assert params of ExecContext
-	pgGo.conn.AssertCalled(t, "ExecContext", ctx, `COMMENT ON TABLE ? IS ?`, "ballot", "")
+	pgGo.conn.AssertCalled(t, "ExecContext",
+		ctx, "COMMENT ON TABLE ? IS ?", []pg.Safe{"ballots", "Ballot table"})
+
+	//paramsMatcher := mock.MatchedBy(func(params []string) bool {
+	//	return IsEqual(params, []string{"ballots", "Ballot table"})
+	//})
+	//pgGo.conn.AssertCalled(t, "ExecContext", paramsMatcher)
+}
+
+func IsEqual(a1 []string, a2 []string) bool {
+	sort.Strings(a1)
+	sort.Strings(a2)
+	if len(a1) == len(a2) {
+		for i, v := range a1 {
+			if v != a2[i] {
+				return false
+			}
+		}
+	} else {
+		return false
+	}
+	return true
 }