github.com/machinefi/w3bstream@v1.6.5-rc9.0.20240426031326-b8c7c4876e72/pkg/depends/kit/sqlx/driver/postgres/driver.go (about)

     1  package postgres
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"database/sql/driver"
     7  	"strconv"
     8  	"strings"
     9  
    10  	"github.com/lib/pq"
    11  	"github.com/pkg/errors"
    12  
    13  	"github.com/machinefi/w3bstream/pkg/depends/kit/logr"
    14  	"github.com/machinefi/w3bstream/pkg/depends/kit/sqlx"
    15  	"github.com/machinefi/w3bstream/pkg/depends/x/misc/timer"
    16  )
    17  
    18  type Driver struct {
    19  	drv pq.Driver
    20  }
    21  
    22  func (d *Driver) Open(dsn string) (driver.Conn, error) {
    23  	cfg, err := pq.ParseURL(dsn)
    24  	if err != nil {
    25  		return nil, err
    26  	}
    27  	opts := ParseOption(cfg)
    28  	if passwd, ok := opts["password"]; ok {
    29  		opts["password"] = strings.Repeat("*", len(passwd))
    30  	}
    31  	conn, err := d.drv.Open(cfg)
    32  	if err != nil {
    33  		return nil, errors.Wrapf(err, "Driver.Open")
    34  	}
    35  	return &LoggingConn{opts, conn}, nil
    36  }
    37  
    38  type LoggingConn struct {
    39  	opts Opts
    40  	driver.Conn
    41  }
    42  
    43  var _ interface {
    44  	driver.ConnBeginTx
    45  	driver.ExecerContext
    46  	driver.QueryerContext
    47  } = (*LoggingConn)(nil)
    48  
    49  func (c *LoggingConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
    50  	l := logr.FromContext(ctx)
    51  	l.Debug("=========== Beginning Transaction ===========")
    52  
    53  	tx, err := c.Conn.(driver.ConnBeginTx).BeginTx(ctx, opts)
    54  	if err != nil {
    55  		l.Error(errors.Wrap(err, "failed to begin transaction"))
    56  		return nil, err
    57  	}
    58  	return &LoggingTx{tx: tx, l: l}, nil
    59  }
    60  
    61  func (c *LoggingConn) Prepare(string) (driver.Stmt, error) { panic("dont use Prepare") }
    62  
    63  func (c *LoggingConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (rows driver.Rows, err error) {
    64  	cost := timer.Start()
    65  	_ctx, l := logr.Start(ctx, "postgres.Query")
    66  
    67  	defer func() {
    68  		qs := interpolate(query, args)
    69  		du := cost().Microseconds()
    70  
    71  		l = l.WithValues(
    72  			"@cst", du,
    73  			"@sql", qs.String(),
    74  		)
    75  
    76  		if err != nil {
    77  			if pgErr, ok := sqlx.UnwrapAll(err).(*pq.Error); !ok {
    78  				l.Error(err)
    79  			} else {
    80  				l.Warn(pgErr)
    81  			}
    82  		} else {
    83  			l.Debug("")
    84  		}
    85  
    86  		l.End()
    87  	}()
    88  
    89  	rows, err = c.Conn.(driver.QueryerContext).QueryContext(_ctx, replaceValueHolder(query), args)
    90  	return
    91  }
    92  
    93  func (c *LoggingConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (res driver.Result, err error) {
    94  	cost := timer.Start()
    95  
    96  	ctx, l := logr.Start(ctx, "postgres.Exec")
    97  
    98  	defer func() {
    99  		qs := interpolate(query, args)
   100  		du := strconv.FormatInt(cost().Microseconds(), 10) + "μs"
   101  
   102  		l = l.WithValues(
   103  			"@cst", du,
   104  			"@sql", qs.String(),
   105  		)
   106  
   107  		if err != nil {
   108  			if pgError, ok := sqlx.UnwrapAll(err).(*pq.Error); !ok {
   109  				l.Error(err)
   110  			} else if pgError.Code == "23505" {
   111  				l.Warn(pgError)
   112  			} else {
   113  				l.Error(pgError)
   114  			}
   115  			return
   116  		}
   117  
   118  		l.Debug("")
   119  		l.End()
   120  	}()
   121  
   122  	res, err = c.Conn.(driver.ExecerContext).ExecContext(ctx, replaceValueHolder(query), args)
   123  	return
   124  }
   125  
   126  type LoggingTx struct {
   127  	l  logr.Logger
   128  	tx driver.Tx
   129  }
   130  
   131  func (tx *LoggingTx) Commit() error {
   132  	if err := tx.tx.Commit(); err != nil {
   133  		tx.l.Debug("failed to commit transaction: %s", err)
   134  		return err
   135  	}
   136  	tx.l.Debug("=========== Committed Transaction ===========")
   137  	return nil
   138  
   139  }
   140  
   141  func (tx *LoggingTx) Rollback() error {
   142  	if err := tx.tx.Rollback(); err != nil {
   143  		tx.l.Debug("failed to rollback transaction: %s", err)
   144  		return err
   145  	}
   146  	tx.l.Debug("=========== Rollback Transaction ===========")
   147  	return nil
   148  }
   149  
   150  func replaceValueHolder(query string) string {
   151  	qc := 0
   152  	buf := bytes.NewBuffer(nil)
   153  
   154  	for _, c := range []byte(query) {
   155  		switch c {
   156  		case '?':
   157  			buf.WriteByte('$')
   158  			buf.WriteString(strconv.Itoa(qc + 1))
   159  			qc++
   160  		default:
   161  			buf.WriteByte(c)
   162  		}
   163  	}
   164  	return buf.String()
   165  }