im_server/server/db/mysql/metrics_driver.go

//go:build mysql
// +build mysql

package mysql

import (
	"context"
	"database/sql"
	"database/sql/driver"
	"time"

	ms "github.com/go-sql-driver/mysql"
	immetrics "github.com/tinode/chat/server/metrics"
)

const instrumentedMySQLDriverName = "meetingim_mysql_metrics"

func init() {
	sql.Register(instrumentedMySQLDriverName, instrumentedMySQLDriver{})
}

type instrumentedMySQLDriver struct{}

func (instrumentedMySQLDriver) Open(dsn string) (driver.Conn, error) {
	connector, err := instrumentedMySQLDriver{}.OpenConnector(dsn)
	if err != nil {
		return nil, err
	}
	return connector.Connect(context.Background())
}

func (instrumentedMySQLDriver) OpenConnector(dsn string) (driver.Connector, error) {
	cfg, err := ms.ParseDSN(dsn)
	if err != nil {
		return nil, err
	}
	connector, err := ms.NewConnector(cfg)
	if err != nil {
		return nil, err
	}
	return instrumentedConnector{connector: connector}, nil
}

type instrumentedConnector struct {
	connector driver.Connector
}

func (c instrumentedConnector) Connect(ctx context.Context) (driver.Conn, error) {
	conn, err := c.connector.Connect(ctx)
	if err != nil {
		return nil, err
	}
	return instrumentedConn{conn: conn}, nil
}

func (c instrumentedConnector) Driver() driver.Driver {
	return instrumentedMySQLDriver{}
}

type instrumentedConn struct {
	conn driver.Conn
}

func (c instrumentedConn) Prepare(query string) (driver.Stmt, error) {
	stmt, err := c.conn.Prepare(query)
	if err != nil {
		return nil, err
	}
	return instrumentedStmt{stmt: stmt}, nil
}

func (c instrumentedConn) Close() error {
	return c.conn.Close()
}

func (c instrumentedConn) Begin() (driver.Tx, error) {
	started := time.Now()
	tx, err := c.conn.Begin()
	immetrics.ObserveDB("mysql", "begin", started, err)
	if err != nil {
		return nil, err
	}
	return instrumentedTx{tx: tx}, nil
}

func (c instrumentedConn) Ping(ctx context.Context) error {
	pinger, ok := c.conn.(driver.Pinger)
	if !ok {
		return driver.ErrSkip
	}
	started := time.Now()
	err := pinger.Ping(ctx)
	immetrics.ObserveDB("mysql", "ping", started, err)
	return err
}

func (c instrumentedConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
	execer, ok := c.conn.(driver.ExecerContext)
	if !ok {
		return nil, driver.ErrSkip
	}
	started := time.Now()
	result, err := execer.ExecContext(ctx, query, args)
	immetrics.ObserveDB("mysql", "exec", started, err)
	return result, err
}

func (c instrumentedConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
	queryer, ok := c.conn.(driver.QueryerContext)
	if !ok {
		return nil, driver.ErrSkip
	}
	started := time.Now()
	rows, err := queryer.QueryContext(ctx, query, args)
	immetrics.ObserveDB("mysql", "query", started, err)
	return rows, err
}

func (c instrumentedConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
	beginner, ok := c.conn.(driver.ConnBeginTx)
	if !ok {
		return nil, driver.ErrSkip
	}
	started := time.Now()
	tx, err := beginner.BeginTx(ctx, opts)
	immetrics.ObserveDB("mysql", "begin", started, err)
	if err != nil {
		return nil, err
	}
	return instrumentedTx{tx: tx}, nil
}

func (c instrumentedConn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
	preparer, ok := c.conn.(driver.ConnPrepareContext)
	if !ok {
		return nil, driver.ErrSkip
	}
	stmt, err := preparer.PrepareContext(ctx, query)
	if err != nil {
		return nil, err
	}
	return instrumentedStmt{stmt: stmt}, nil
}

func (c instrumentedConn) CheckNamedValue(value *driver.NamedValue) error {
	checker, ok := c.conn.(driver.NamedValueChecker)
	if !ok {
		return driver.ErrSkip
	}
	return checker.CheckNamedValue(value)
}

func (c instrumentedConn) ResetSession(ctx context.Context) error {
	resetter, ok := c.conn.(driver.SessionResetter)
	if !ok {
		return nil
	}
	return resetter.ResetSession(ctx)
}

func (c instrumentedConn) IsValid() bool {
	validator, ok := c.conn.(driver.Validator)
	if !ok {
		return true
	}
	return validator.IsValid()
}

type instrumentedStmt struct {
	stmt driver.Stmt
}

func (s instrumentedStmt) Close() error {
	return s.stmt.Close()
}

func (s instrumentedStmt) NumInput() int {
	return s.stmt.NumInput()
}

func (s instrumentedStmt) Exec(args []driver.Value) (driver.Result, error) {
	started := time.Now()
	result, err := s.stmt.Exec(args)
	immetrics.ObserveDB("mysql", "exec", started, err)
	return result, err
}

func (s instrumentedStmt) Query(args []driver.Value) (driver.Rows, error) {
	started := time.Now()
	rows, err := s.stmt.Query(args)
	immetrics.ObserveDB("mysql", "query", started, err)
	return rows, err
}

func (s instrumentedStmt) ExecContext(ctx context.Context, args []driver.NamedValue) (driver.Result, error) {
	execer, ok := s.stmt.(driver.StmtExecContext)
	if !ok {
		return nil, driver.ErrSkip
	}
	started := time.Now()
	result, err := execer.ExecContext(ctx, args)
	immetrics.ObserveDB("mysql", "exec", started, err)
	return result, err
}

func (s instrumentedStmt) QueryContext(ctx context.Context, args []driver.NamedValue) (driver.Rows, error) {
	queryer, ok := s.stmt.(driver.StmtQueryContext)
	if !ok {
		return nil, driver.ErrSkip
	}
	started := time.Now()
	rows, err := queryer.QueryContext(ctx, args)
	immetrics.ObserveDB("mysql", "query", started, err)
	return rows, err
}

func (s instrumentedStmt) CheckNamedValue(value *driver.NamedValue) error {
	checker, ok := s.stmt.(driver.NamedValueChecker)
	if !ok {
		return driver.ErrSkip
	}
	return checker.CheckNamedValue(value)
}

type instrumentedTx struct {
	tx driver.Tx
}

func (tx instrumentedTx) Commit() error {
	started := time.Now()
	err := tx.tx.Commit()
	immetrics.ObserveDB("mysql", "commit", started, err)
	return err
}

func (tx instrumentedTx) Rollback() error {
	started := time.Now()
	err := tx.tx.Rollback()
	immetrics.ObserveDB("mysql", "rollback", started, err)
	return err
}

var (
	_ driver.Driver             = instrumentedMySQLDriver{}
	_ driver.DriverContext      = instrumentedMySQLDriver{}
	_ driver.Connector          = instrumentedConnector{}
	_ driver.Conn               = instrumentedConn{}
	_ driver.Pinger             = instrumentedConn{}
	_ driver.ExecerContext      = instrumentedConn{}
	_ driver.QueryerContext     = instrumentedConn{}
	_ driver.ConnBeginTx        = instrumentedConn{}
	_ driver.ConnPrepareContext = instrumentedConn{}
	_ driver.NamedValueChecker  = instrumentedConn{}
	_ driver.SessionResetter    = instrumentedConn{}
	_ driver.Validator          = instrumentedConn{}
	_ driver.Stmt               = instrumentedStmt{}
	_ driver.StmtExecContext    = instrumentedStmt{}
	_ driver.StmtQueryContext   = instrumentedStmt{}
	_ driver.NamedValueChecker  = instrumentedStmt{}
	_ driver.Tx                 = instrumentedTx{}
)