rano/db/client.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())
}