// 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()) }