135 lines
2.9 KiB
Go
135 lines
2.9 KiB
Go
// Copyright 2024 Patial Tech. All rights reserved.
|
|
// Use of this source code is governed by a BSD-style
|
|
// license that can be found in the LICENSE file.
|
|
|
|
package db
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"database/sql/driver"
|
|
"fmt"
|
|
"reflect"
|
|
"strings"
|
|
"time"
|
|
|
|
"contrib.go.opencensus.io/integrations/ocsql"
|
|
"entgo.io/ent/dialect"
|
|
entsql "entgo.io/ent/dialect/sql"
|
|
pgx "github.com/jackc/pgx/v5/stdlib"
|
|
"gitserver.in/patialtech/rano/config"
|
|
"gitserver.in/patialtech/rano/db/ent"
|
|
"gitserver.in/patialtech/rano/util/logger"
|
|
)
|
|
|
|
type connector struct {
|
|
dsn string
|
|
}
|
|
|
|
// New *sql.DB instance
|
|
func New() *sql.DB {
|
|
databaseUrl := config.Read().DbURL
|
|
db := sql.OpenDB(connector{dsn: databaseUrl})
|
|
db.SetMaxOpenConns(50)
|
|
db.SetMaxIdleConns(5)
|
|
db.SetConnMaxLifetime(0) // Maximum amount of time a connection can be reused (0 means no limit)
|
|
return db
|
|
}
|
|
|
|
func (c connector) Connect(context.Context) (driver.Conn, error) {
|
|
return c.Driver().Open(c.dsn)
|
|
}
|
|
|
|
func (connector) Driver() driver.Driver {
|
|
return ocsql.Wrap(
|
|
pgx.GetDefaultDriver(),
|
|
ocsql.WithAllTraceOptions(),
|
|
ocsql.WithRowsClose(false),
|
|
ocsql.WithRowsNext(false),
|
|
ocsql.WithDisableErrSkip(true),
|
|
)
|
|
}
|
|
|
|
// Client for pgx
|
|
//
|
|
// https://entgo.io/docs/sql-integration
|
|
func Client() *ent.Client {
|
|
// Create an ent.Driver from `db`.
|
|
drv := entsql.OpenDB(dialect.Postgres, New())
|
|
cl := ent.NewClient(ent.Driver(drv))
|
|
cl.Use(AuditHook)
|
|
return cl
|
|
}
|
|
|
|
// A AuditHook is an example for audit-log hook.
|
|
func AuditHook(next ent.Mutator) ent.Mutator {
|
|
return ent.MutateFunc(func(ctx context.Context, m ent.Mutation) (v ent.Value, err error) {
|
|
// start timer
|
|
start := time.Now()
|
|
// do the operation
|
|
v, err = next.Mutate(ctx, m)
|
|
|
|
// audit log
|
|
var id int64
|
|
if v != nil {
|
|
ev := reflect.Indirect(reflect.ValueOf(v))
|
|
f := ev.FieldByName("ID")
|
|
if f.IsValid() {
|
|
id = f.Int()
|
|
}
|
|
}
|
|
|
|
saveAudit(ctx, m.Type(), m.Op(), id, err, time.Since(start))
|
|
|
|
return
|
|
})
|
|
}
|
|
|
|
func saveAudit(ctx context.Context, t string, op ent.Op, id int64, err error, d time.Duration) {
|
|
if t == "Audit" {
|
|
// skip Audit table operations
|
|
return
|
|
}
|
|
|
|
var entOp string
|
|
switch {
|
|
case op.Is(ent.OpCreate):
|
|
entOp = "Create"
|
|
case op.Is(ent.OpDelete):
|
|
entOp = "Delete"
|
|
case op.Is(ent.OpDeleteOne):
|
|
entOp = "DeleteOne"
|
|
case op.Is(ent.OpUpdate):
|
|
entOp = "Update"
|
|
case op.Is(ent.OpUpdateOne):
|
|
entOp = "UpdateOne"
|
|
}
|
|
|
|
var sb strings.Builder
|
|
if reqID, ok := ctx.Value("RequestID").(string); ok {
|
|
sb.WriteString(reqID + " ")
|
|
}
|
|
|
|
if err != nil {
|
|
sb.WriteString("failed ")
|
|
}
|
|
|
|
sb.WriteString(fmt.Sprintf("%s:%s:%d ", entOp, t, id))
|
|
|
|
if ip, ok := ctx.Value("RequestIP").(string); ok {
|
|
sb.WriteString(fmt.Sprintf("ip=%s ", ip))
|
|
}
|
|
|
|
if ua, ok := ctx.Value("RequestUA").(string); ok {
|
|
sb.WriteString(fmt.Sprintf("ip=%s ", ua))
|
|
}
|
|
|
|
if err != nil {
|
|
sb.WriteString(fmt.Sprintf("error=%s ", err))
|
|
}
|
|
|
|
sb.WriteString(fmt.Sprintf("t=%s", d))
|
|
|
|
logger.Info(sb.String())
|
|
}
|