// 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" "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 } type entity interface { GetID() (int64, error) } // 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 := time.Now() defer func() { saveAudit(ctx, m.Type(), m.Op(), v, err, time.Since(start)) }() v, err = next.Mutate(ctx, m) return }) } func saveAudit(ctx context.Context, t string, op ent.Op, v ent.Value, 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" } big. reqID := ctx.Value("RequestID") ip := ctx.Value("RequestIP") ua := ctx.Value("RequestUA") if en, ok := v.(entity); ok { id, _ := en.GetID() logger.Info("%s %s %s-%s(%d) ip=%s ua=%s t=%v error=%e", reqID, status, op.String(), t, id, ip, ua, d, err) } else { logger.Info("%s %s %s-%s ip=%s ua=%s t=%v error=%e", reqID, status, op.String(), t, ip, ua, d, err) } }