mirror of https://gitlab.com/ngerakines/tavern.git
118 lines
3.6 KiB
Go
118 lines
3.6 KiB
Go
package storage
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
|
|
"github.com/prometheus/client_golang/prometheus"
|
|
"github.com/prometheus/client_golang/prometheus/promauto"
|
|
"go.uber.org/zap"
|
|
)
|
|
|
|
type QueryExecute interface {
|
|
ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error)
|
|
QueryContext(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error)
|
|
QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row
|
|
}
|
|
|
|
type MetricDriver struct {
|
|
Driver QueryExecute
|
|
Gauge prometheus.Gauge
|
|
Counter prometheus.Counter
|
|
}
|
|
|
|
type TransactionSQLDriver struct {
|
|
Driver *sql.Tx
|
|
}
|
|
|
|
type LoggingSQLDriver struct {
|
|
Driver QueryExecute
|
|
Logger *zap.Logger
|
|
}
|
|
|
|
type SQLDriver struct {
|
|
Driver *sql.DB
|
|
}
|
|
|
|
func (d LoggingSQLDriver) ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error) {
|
|
d.Logger.Debug("ExecContext", zap.String("query", query), zap.Reflect("args", args))
|
|
return d.Driver.ExecContext(ctx, query, args...)
|
|
}
|
|
|
|
func (d LoggingSQLDriver) QueryContext(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error) {
|
|
d.Logger.Debug("QueryContext", zap.String("query", query), zap.Reflect("args", args))
|
|
return d.Driver.QueryContext(ctx, query, args...)
|
|
}
|
|
|
|
func (d LoggingSQLDriver) QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row {
|
|
d.Logger.Debug("QueryRowContext", zap.String("query", query), zap.Reflect("args", args))
|
|
return d.Driver.QueryRowContext(ctx, query, args...)
|
|
}
|
|
|
|
func (d SQLDriver) ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error) {
|
|
return d.Driver.ExecContext(ctx, query, args...)
|
|
}
|
|
|
|
func (d SQLDriver) QueryContext(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error) {
|
|
return d.Driver.QueryContext(ctx, query, args...)
|
|
}
|
|
|
|
func (d SQLDriver) QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row {
|
|
return d.Driver.QueryRowContext(ctx, query, args...)
|
|
}
|
|
|
|
func (d TransactionSQLDriver) ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error) {
|
|
return d.Driver.ExecContext(ctx, query, args...)
|
|
}
|
|
|
|
func (d TransactionSQLDriver) QueryContext(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error) {
|
|
return d.Driver.QueryContext(ctx, query, args...)
|
|
}
|
|
|
|
func (d TransactionSQLDriver) QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row {
|
|
return d.Driver.QueryRowContext(ctx, query, args...)
|
|
}
|
|
|
|
func (d MetricDriver) ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error) {
|
|
d.Counter.Inc()
|
|
d.Gauge.Inc()
|
|
defer d.Gauge.Dec()
|
|
|
|
return d.Driver.ExecContext(ctx, query, args...)
|
|
}
|
|
|
|
func (d MetricDriver) QueryContext(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error) {
|
|
d.Counter.Inc()
|
|
d.Gauge.Inc()
|
|
defer d.Gauge.Dec()
|
|
|
|
return d.Driver.QueryContext(ctx, query, args...)
|
|
}
|
|
|
|
func (d MetricDriver) QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row {
|
|
d.Counter.Inc()
|
|
d.Gauge.Inc()
|
|
defer d.Gauge.Dec()
|
|
|
|
return d.Driver.QueryRowContext(ctx, query, args...)
|
|
}
|
|
|
|
func WrapMetricDriver(metricFactory promauto.Factory, namespace, subsystem string, driver QueryExecute) QueryExecute {
|
|
return MetricDriver{
|
|
Driver: driver,
|
|
Gauge: metricFactory.NewGauge(prometheus.GaugeOpts{
|
|
Namespace: namespace,
|
|
Subsystem: subsystem,
|
|
Name: "queries_in_flight",
|
|
Help: "A gauge of in-flight DB queries.",
|
|
}),
|
|
Counter: metricFactory.NewCounter(
|
|
prometheus.CounterOpts{
|
|
Namespace: namespace,
|
|
Subsystem: subsystem,
|
|
Name: "queries_total",
|
|
Help: "A counter for queries.",
|
|
}),
|
|
}
|
|
}
|