Skip to content

Commit 8ddb8c6

Browse files
vvuweivvuwei
authored
Merge pull request #13 from dipdup-io/GO-80-add-pg-comment-tag
GO 80 add pg comment tag
2 parents 53a6788 + 384c216 commit 8ddb8c6

File tree

6 files changed

+431
-2
lines changed

6 files changed

+431
-2
lines changed

‎database/pg.go‎

Lines changed: 65 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,75 @@ package database
33
import (
44
"context"
55
"fmt"
6-
76
"github.com/dipdup-net/go-lib/config"
87
pg "github.com/go-pg/pg/v10"
98
"github.com/pkg/errors"
109
)
1110

11+
type PgGoConnection interface {
12+
DB() PgDB
13+
}
14+
15+
type PgDB interface {
16+
/*
17+
Begin() (*pg.Tx, error)
18+
BeginContext(ctx context.Context) (*pg.Tx, error)
19+
RunInTransaction(ctx context.Context, fn func(*pg.Tx) error) error
20+
AddQueryHook(hook pg.QueryHook)
21+
beforeQuery(ctx context.Context, ormDB orm.DB, model interface{}, query interface{}, params []interface{}, fmtedQuery []byte) (context.Context, *QueryEvent, error)
22+
afterQuery(ctx context.Context, event *pg.QueryEvent, res pg.Result, err error) error
23+
afterQueryFromIndex(ctx context.Context, event *pg.QueryEvent, hookIndex int) error
24+
//startup(c context.Context, cn *pool.Conn, user string, password string, database string, appName string) error
25+
//enableSSL(c context.Context, cn *pool.Conn, tlsConf *tls.Config) error
26+
//auth(c context.Context, cn *pool.Conn, rd *pool.ReaderContext, user string, password string) error
27+
//logStartupNotice(rd *pool.ReaderContext) error
28+
//authCleartext(c context.Context, cn *pool.Conn, rd *pool.ReaderContext, password string) error
29+
//authMD5(c context.Context, cn *pool.Conn, rd *pool.ReaderContext, user string, password string) error
30+
//authSASL(c context.Context, cn *pool.Conn, rd *pool.ReaderContext, user string, password string) error
31+
PoolStats() *pg.PoolStats
32+
//clone() *baseDB
33+
//withPool(p pool.Pooler) *baseDB
34+
//WithTimeout(d time.Duration) *baseDB
35+
//WithParam(param string, value interface{}) *baseDB
36+
Param(param string) interface{}
37+
retryBackoff(retry int) time.Duration
38+
//getConn(ctx context.Context) (*pool.Conn, error)
39+
//initConn(ctx context.Context, cn *pool.Conn) error
40+
//releaseConn(ctx context.Context, cn *pool.Conn, err error)
41+
//withConn(ctx context.Context, fn func(context.Context, *pool.Conn) error) error
42+
shouldRetry(err error) bool
43+
Close() error
44+
Exec(query interface{}, params ...interface{}) (res pg.Result, err error)
45+
*/
46+
ExecContext(c context.Context, query interface{}, params ...interface{}) (pg.Result, error)
47+
/*
48+
//exec(ctx context.Context, query interface{}, params ...interface{}) (pg.Result, error)
49+
ExecOne(query interface{}, params ...interface{}) (pg.Result, error)
50+
ExecOneContext(ctx context.Context, query interface{}, params ...interface{}) (pg.Result, error)
51+
execOne(c context.Context, query interface{}, params ...interface{}) (pg.Result, error)
52+
Query(model interface{}, query interface{}, params ...interface{}) (res pg.Result, err error)
53+
QueryContext(c context.Context, model interface{}, query interface{}, params ...interface{}) (pg.Result, error)
54+
query(ctx context.Context, model interface{}, query interface{}, params ...interface{}) (pg.Result, error)
55+
QueryOne(model interface{}, query interface{}, params ...interface{}) (pg.Result, error)
56+
QueryOneContext(ctx context.Context, model interface{}, query interface{}, params ...interface{}) (pg.Result, error)
57+
queryOne(ctx context.Context, model interface{}, query interface{}, params ...interface{}) (pg.Result, error)
58+
CopyFrom(r io.Reader, query interface{}, params ...interface{}) (res pg.Result, err error)
59+
//copyFrom(ctx context.Context, cn *pool.Conn, r io.Reader, query interface{}, params ...interface{}) (res Result, err error)
60+
CopyTo(w io.Writer, query interface{}, params ...interface{}) (res pg.Result, err error)
61+
//copyTo(ctx context.Context, cn *pool.Conn, w io.Writer, query interface{}, params ...interface{}) (res Result, err error)
62+
Ping(ctx context.Context) error
63+
Model(model ...interface{}) *pg.Query
64+
ModelContext(c context.Context, model ...interface{}) *pg.Query
65+
Formatter() orm.QueryFormatter
66+
//cancelRequest(processID int32, secretKey int32) error
67+
//simpleQuery(c context.Context, cn *pool.Conn, wb *pool.WriteBuffer) (*result, error)
68+
//simpleQueryData(c context.Context, cn *pool.Conn, model interface{}, wb *pool.WriteBuffer) (*result, error)
69+
Prepare(q string) (*pg.Stmt, error)
70+
//prepare(c context.Context, cn *pool.Conn, q string) (string, []types.ColumnInfo, error)
71+
//closeStmt(c context.Context, cn *pool.Conn, name string) error
72+
*/
73+
}
74+
1275
// PgGo -
1376
type PgGo struct {
1477
conn *pg.DB
@@ -20,7 +83,7 @@ func NewPgGo() *PgGo {
2083
}
2184

2285
// DB -
23-
func (db *PgGo) DB() *pg.DB {
86+
func (db *PgGo) DB() PgDB {
2487
return db.conn
2588
}
2689

‎database/pgComment.go‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
package database
2+
3+
import (
4+
"context"
5+
"github.com/dipdup-net/go-lib/hasura"
6+
"github.com/go-pg/pg/v10"
7+
"reflect"
8+
"strings"
9+
)
10+
11+
func makeComments(ctx context.Context, conn PgGoConnection, models ...interface{}) error {
12+
for _, model := range models {
13+
modelType := reflect.TypeOf(model)
14+
var tableName pg.Safe
15+
16+
for i := 0; i < modelType.NumField(); i++ {
17+
fieldType := modelType.Field(i)
18+
19+
if fieldType.Name == "tableName" {
20+
var ok bool
21+
tableName, ok = getPgName(fieldType)
22+
if !ok {
23+
tableName = pg.Safe(hasura.ToSnakeCase(modelType.Name()))
24+
}
25+
26+
pgCommentTag, ok := getPgComment(fieldType)
27+
if !ok {
28+
continue
29+
}
30+
31+
if _, err := conn.DB().ExecContext(ctx,
32+
`COMMENT ON TABLE ? IS ?`,
33+
tableName, pgCommentTag); err != nil {
34+
return err
35+
}
36+
37+
continue
38+
}
39+
40+
pgCommentTag, ok := getPgComment(fieldType)
41+
if !ok {
42+
continue
43+
}
44+
45+
columnName, ok := getPgName(fieldType)
46+
if !ok {
47+
columnName = pg.Safe(hasura.ToSnakeCase(fieldType.Name))
48+
}
49+
50+
if _, err := conn.DB().ExecContext(ctx,
51+
`COMMENT ON COLUMN ?.? IS ?`,
52+
tableName, columnName, pgCommentTag); err != nil {
53+
return err
54+
}
55+
}
56+
}
57+
return nil
58+
}
59+
60+
func getPgName(fieldType reflect.StructField) (name pg.Safe, ok bool) {
61+
pgTag, ok := fieldType.Tag.Lookup("pg")
62+
if !ok {
63+
return "", false
64+
}
65+
66+
tags := strings.Split(pgTag, ",")
67+
68+
if tags[0] != "" {
69+
name = pg.Safe(tags[0])
70+
return name, ok
71+
}
72+
73+
return "", false
74+
}
75+
76+
func getPgComment(fieldType reflect.StructField) (pg.Safe, bool) {
77+
pgCommentTag, ok := fieldType.Tag.Lookup("pg-comment")
78+
79+
if ok {
80+
return pg.Safe(pgCommentTag), ok
81+
}
82+
83+
return "", false
84+
}

‎database/pgComment_test.go‎

Lines changed: 218 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,218 @@
1+
package database
2+
3+
import (
4+
"context"
5+
"github.com/dipdup-net/go-lib/mocks"
6+
"github.com/go-pg/pg/v10"
7+
"github.com/golang/mock/gomock"
8+
"github.com/stretchr/testify/assert"
9+
"testing"
10+
)
11+
12+
type PgGoMock struct {
13+
conn *mocks.MockPgDB
14+
}
15+
16+
func (p *PgGoMock) DB() PgDB {
17+
return p.conn
18+
}
19+
20+
func newPgGoMock(pgDB *mocks.MockPgDB) *PgGoMock {
21+
return &PgGoMock{
22+
conn: pgDB,
23+
}
24+
}
25+
26+
func TestMakeCommentsWithTableName(t *testing.T) {
27+
type Ballot struct {
28+
//nolint
29+
tableName struct{} `pg:"ballots" pg-comment:"Ballot table"`
30+
Ballot string `json:"ballot"`
31+
}
32+
33+
mockCtrl, mockPgDB, pgGo, ctx := createPgDbMock(t)
34+
defer mockCtrl.Finish()
35+
36+
model := Ballot{}
37+
38+
// Assert prepare
39+
expectedParams := toInterfaceSlice([]pg.Safe{"ballots", "Ballot table"})
40+
mockPgDB.
41+
EXPECT().
42+
ExecContext(ctx, "COMMENT ON TABLE ? IS ?",
43+
gomock.Eq(expectedParams)).
44+
Times(1).
45+
Return(nil, nil)
46+
47+
// Act
48+
err := makeComments(ctx, pgGo, model)
49+
50+
// Assert
51+
assert.Empty(t, err)
52+
}
53+
54+
func TestMakeCommentsWithoutPgComment(t *testing.T) {
55+
type Ballot struct {
56+
//nolint
57+
tableName struct{} `pg:"ballots"`
58+
Ballot string `json:"ballot"`
59+
}
60+
61+
mockCtrl, mockPgDB, pgGo, ctx := createPgDbMock(t)
62+
defer mockCtrl.Finish()
63+
64+
model := Ballot{}
65+
66+
// Assert prepare
67+
mockPgDB.
68+
EXPECT().
69+
ExecContext(ctx, "COMMENT ON TABLE ? IS ?", gomock.Any()).
70+
Times(0).
71+
Return(nil, nil)
72+
73+
mockPgDB.
74+
EXPECT().
75+
ExecContext(ctx, "COMMENT ON COLUMN ?.? IS ?", gomock.Any()).
76+
Times(0).
77+
Return(nil, nil)
78+
79+
// Act
80+
err := makeComments(ctx, pgGo, model)
81+
82+
// Assert
83+
assert.Empty(t, err)
84+
}
85+
86+
func TestMakeCommentsFieldWithPgComment(t *testing.T) {
87+
type Ballot struct {
88+
//nolint
89+
tableName struct{} `pg:"ballots"`
90+
Ballot string `json:"ballot" pg-comment:"This is field comment"`
91+
}
92+
93+
mockCtrl, mockPgDB, pgGo, ctx := createPgDbMock(t)
94+
defer mockCtrl.Finish()
95+
96+
model := Ballot{}
97+
98+
// Assert prepare
99+
expectedParams := toInterfaceSlice([]pg.Safe{"ballots", "ballot", "This is field comment"})
100+
mockPgDB.
101+
EXPECT().
102+
ExecContext(ctx, "COMMENT ON COLUMN ?.? IS ?",
103+
gomock.Eq(expectedParams)).
104+
Times(1).
105+
Return(nil, nil)
106+
107+
// Act
108+
err := makeComments(ctx, pgGo, model)
109+
110+
// Assert
111+
assert.Empty(t, err)
112+
}
113+
114+
func TestMakeCommentsWithTableNameAndFieldsWithPgComment(t *testing.T) {
115+
type Ballot struct {
116+
//nolint
117+
tableName struct{} `pg:"ballots" pg-comment:"Ballot table"`
118+
CreatedAt int64 `json:"-" pg-comment:"This is field comment"`
119+
UpdatedAt int64 `json:"-" pg-comment:"This is field comment"`
120+
Network string `json:"network" pg:",pk" pg-comment:"This is field comment"`
121+
Hash string `json:"hash" pg:",pk" pg-comment:"This is field comment"`
122+
Branch string `json:"branch" pg-comment:"This is field comment"`
123+
Status string `json:"status" pg-comment:"This is field comment"`
124+
Kind string `json:"kind" pg-comment:"This is field comment"`
125+
Signature string `json:"signature" pg-comment:"This is field comment"`
126+
Protocol string `json:"protocol" pg-comment:"This is field comment"`
127+
Level uint64 `json:"level" pg-comment:"This is field comment"`
128+
Errors interface{} `json:"errors,omitempty" pg:"type:jsonb" pg-comment:"This is field comment"`
129+
ExpirationLevel *uint64 `json:"expiration_level" pg-comment:"This is field comment"`
130+
Raw interface{} `json:"raw,omitempty" pg:"type:jsonb" pg-comment:"This is field comment"`
131+
Ballot string `json:"ballot" pg-comment:"This is field comment"`
132+
Period int64 `json:"period" pg-comment:"This is field comment"`
133+
}
134+
135+
mockCtrl, mockPgDB, pgGo, ctx := createPgDbMock(t)
136+
defer mockCtrl.Finish()
137+
138+
model := Ballot{}
139+
140+
// Assert prepare
141+
expectedParams := toInterfaceSlice([]pg.Safe{"ballots", "Ballot table"})
142+
commentOnTableCall := mockPgDB.
143+
EXPECT().
144+
ExecContext(ctx, "COMMENT ON TABLE ? IS ?",
145+
gomock.Eq(expectedParams)).
146+
Times(1).
147+
Return(nil, nil)
148+
149+
mockPgDB.
150+
EXPECT().
151+
ExecContext(ctx, "COMMENT ON COLUMN ?.? IS ?", gomock.Any()).
152+
Times(15).
153+
After(commentOnTableCall).
154+
Return(nil, nil)
155+
156+
// Act
157+
err := makeComments(ctx, pgGo, model)
158+
159+
// Assert
160+
assert.Empty(t, err)
161+
}
162+
163+
func TestMakeCommentsWithMultipleModels(t *testing.T) {
164+
type Ballot struct {
165+
//nolint
166+
tableName struct{} `pg:"ballots" pg-comment:"This multiple table name comment"`
167+
Ballot string `json:"ballot" pg-comment:"This is multiple field comment"`
168+
}
169+
170+
mockCtrl, mockPgDB, pgGo, ctx := createPgDbMock(t)
171+
defer mockCtrl.Finish()
172+
173+
models := []interface{}{Ballot{}, Ballot{}, Ballot{}}
174+
175+
// Assert prepare
176+
expectedTableParams := toInterfaceSlice([]pg.Safe{"ballots", "This multiple table name comment"})
177+
mockPgDB.
178+
EXPECT().
179+
ExecContext(ctx, "COMMENT ON TABLE ? IS ?",
180+
gomock.Eq(expectedTableParams)).
181+
Times(3).
182+
Return(nil, nil)
183+
184+
// Be aware: there is on issue with default order in checking
185+
// methods call: https://github.com/golang/mock/issues/653
186+
expectedParams := toInterfaceSlice([]pg.Safe{"ballots", "ballot", "This is multiple field comment"})
187+
mockPgDB.
188+
EXPECT().
189+
ExecContext(ctx, "COMMENT ON COLUMN ?.? IS ?",
190+
gomock.Eq(expectedParams)).
191+
Times(3).
192+
Return(nil, nil)
193+
194+
// Act
195+
err := makeComments(ctx, pgGo, models...)
196+
197+
// Assert
198+
assert.Empty(t, err)
199+
}
200+
201+
func createPgDbMock(t *testing.T) (*gomock.Controller, *mocks.MockPgDB, *PgGoMock, context.Context) {
202+
mockCtrl := gomock.NewController(t)
203+
mockPgDB := mocks.NewMockPgDB(mockCtrl)
204+
pgGo := newPgGoMock(mockPgDB)
205+
ctx := context.Background()
206+
207+
return mockCtrl, mockPgDB, pgGo, ctx
208+
}
209+
210+
func toInterfaceSlice(origin []pg.Safe) []interface{} {
211+
res := make([]interface{}, len(origin))
212+
213+
for i := range origin {
214+
res[i] = origin[i]
215+
}
216+
217+
return res
218+
}

‎go.mod‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ require (
77
github.com/ebellocchia/go-base58 v0.1.0
88
github.com/go-pg/pg/v10 v10.10.6
99
github.com/go-playground/validator/v10 v10.9.0
10+
github.com/golang/mock v1.6.0
1011
github.com/gorilla/websocket v1.4.2
1112
github.com/iancoleman/strcase v0.2.0
1213
github.com/json-iterator/go v1.1.12

0 commit comments

Comments
 (0)