github.com/authzed/spicedb@v1.32.1-0.20240520085336-ebda56537386/internal/datastore/mysql/connection.go (about)

     1  package mysql
     2  
     3  import (
     4  	"context"
     5  	"database/sql/driver"
     6  	"fmt"
     7  	"strconv"
     8  	"time"
     9  
    10  	log "github.com/authzed/spicedb/internal/logging"
    11  
    12  	"github.com/prometheus/client_golang/prometheus"
    13  )
    14  
    15  var (
    16  	connectHistogram = prometheus.NewHistogram(prometheus.HistogramOpts{
    17  		Namespace: "spicedb",
    18  		Subsystem: "datastore",
    19  		Name:      "mysql_connect_duration",
    20  		Help:      "distribution in seconds of time spent opening a new MySQL connection.",
    21  		Buckets:   []float64{0.01, 0.1, 0.5, 1, 5, 10, 25, 60, 120},
    22  	})
    23  	connectCount = prometheus.NewCounterVec(prometheus.CounterOpts{
    24  		Namespace: "spicedb",
    25  		Subsystem: "datastore",
    26  		Name:      "mysql_connect_count_total",
    27  		Help:      "number of mysql connections opened.",
    28  	}, []string{"success"})
    29  )
    30  
    31  // instrumentedConnector wraps the default MySQL driver connector
    32  // to get metrics and tracing when creating a new connection
    33  type instrumentedConnector struct {
    34  	conn driver.Connector
    35  	drv  driver.Driver
    36  }
    37  
    38  func (d *instrumentedConnector) Connect(ctx context.Context) (driver.Conn, error) {
    39  	ctx, span := tracer.Start(ctx, "openMySQLConnection")
    40  	defer span.End()
    41  
    42  	startTime := time.Now()
    43  	defer func() {
    44  		connectHistogram.Observe(time.Since(startTime).Seconds())
    45  	}()
    46  
    47  	conn, err := d.conn.Connect(ctx)
    48  	connectCount.WithLabelValues(strconv.FormatBool(err == nil)).Inc()
    49  	if err != nil {
    50  		span.RecordError(err)
    51  		log.Ctx(ctx).Error().Err(err).Msg("failed to open mysql connection")
    52  		return nil, fmt.Errorf("failed to open connection to mysql: %w", err)
    53  	}
    54  
    55  	return conn, nil
    56  }
    57  
    58  func (d *instrumentedConnector) Driver() driver.Driver {
    59  	return d.drv
    60  }
    61  
    62  func instrumentConnector(c driver.Connector) (driver.Connector, error) {
    63  	err := prometheus.Register(connectHistogram)
    64  	if err != nil {
    65  		return nil, fmt.Errorf("unable to register metric: %w", err)
    66  	}
    67  
    68  	err = prometheus.Register(connectCount)
    69  	if err != nil {
    70  		return nil, fmt.Errorf("unable to register metric: %w", err)
    71  	}
    72  
    73  	return &instrumentedConnector{
    74  		conn: c,
    75  		drv:  c.Driver(),
    76  	}, nil
    77  }
    78  
    79  type sessionVariableConnector struct {
    80  	conn driver.Connector
    81  	drv  driver.Driver
    82  
    83  	statements []string
    84  }
    85  
    86  func (s *sessionVariableConnector) Connect(ctx context.Context) (driver.Conn, error) {
    87  	ctx, span := tracer.Start(ctx, "setSessionVariables")
    88  	defer span.End()
    89  
    90  	conn, err := s.conn.Connect(ctx)
    91  	if err != nil {
    92  		span.RecordError(err)
    93  		log.Ctx(ctx).Error().Err(err).Msg("failed to open db connection")
    94  		return nil, fmt.Errorf("failed to open connection to db: %w", err)
    95  	}
    96  
    97  	// The go mysql driver implements the ExecerContext interface, assert that here.
    98  	execConn := conn.(driver.ExecerContext)
    99  
   100  	for _, stmt := range s.statements {
   101  		if _, err := execConn.ExecContext(ctx, stmt, nil); err != nil {
   102  			return nil, fmt.Errorf("unable to execute statement `%s`: %w", stmt, err)
   103  		}
   104  	}
   105  
   106  	return conn, nil
   107  }
   108  
   109  func (s *sessionVariableConnector) Driver() driver.Driver {
   110  	return s.drv
   111  }
   112  
   113  func addSessionVariables(c driver.Connector, variables map[string]string) (driver.Connector, error) {
   114  	statements := make([]string, 0, len(variables))
   115  	for sessionVar, value := range variables {
   116  		statements = append(statements, "SET SESSION "+sessionVar+"="+value)
   117  	}
   118  
   119  	return &sessionVariableConnector{
   120  		conn:       c,
   121  		drv:        c.Driver(),
   122  		statements: statements,
   123  	}, nil
   124  }