feat: add prometheus metrics to database.Store (#7713)

johnstcn · ammario · commit 49fcc365d895 · 2023-06-01T05:14:46.000Z
* Adds dbmetrics package and wraps database.Store with a Prometheus HistogramVec of timings.
* Adds Wrappers method to database.Store to avoid double-wrapping interfaces
* Fixes test flake in TestLicensesListFake
diff --git a/cli/server.go b/cli/server.go
@@ -65,6 +65,7 @@ import (
 	"github.com/coder/coder/coderd/autobuild/executor"
 	"github.com/coder/coder/coderd/database"
 	"github.com/coder/coder/coderd/database/dbfake"
+	"github.com/coder/coder/coderd/database/dbmetrics"
 	"github.com/coder/coder/coderd/database/dbpurge"
 	"github.com/coder/coder/coderd/database/migrations"
 	"github.com/coder/coder/coderd/devtunnel"
@@ -586,7 +587,8 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
 			}
 
 			if cfg.InMemoryDatabase {
-				options.Database = dbfake.New()
+				// This is only used for testing.
+				options.Database = dbmetrics.New(dbfake.New(), options.PrometheusRegistry)
 				options.Pubsub = database.NewPubsubInMemory()
 			} else {
 				sqlDB, err := connectToPostgres(ctx, logger, sqlDriver, cfg.PostgresURL.String())
@@ -597,7 +599,7 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
 					_ = sqlDB.Close()
 				}()
 
-				options.Database = database.New(sqlDB)
+				options.Database = dbmetrics.New(database.New(sqlDB), options.PrometheusRegistry)
 				options.Pubsub, err = database.NewPubsub(ctx, sqlDB, cfg.PostgresURL.String())
 				if err != nil {
 					return xerrors.Errorf("create pubsub: %w", err)
diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -47,6 +47,7 @@ import (
 	"github.com/coder/coder/coderd/awsidentity"
 	"github.com/coder/coder/coderd/database"
 	"github.com/coder/coder/coderd/database/dbauthz"
+	"github.com/coder/coder/coderd/database/dbmetrics"
 	"github.com/coder/coder/coderd/database/dbtype"
 	"github.com/coder/coder/coderd/gitauth"
 	"github.com/coder/coder/coderd/gitsshkey"
@@ -176,6 +177,11 @@ func New(options *Options) *API {
 		options = &Options{}
 	}
 
+	// Safety check: if we're not running a unit test, we *must* have a Prometheus registry.
+	if options.PrometheusRegistry == nil && flag.Lookup("test.v") == nil {
+		panic("developer error: options.PrometheusRegistry is nil and not running a unit test")
+	}
+
 	if options.DeploymentValues.DisableOwnerWorkspaceExec {
 		rbac.ReloadBuiltinRoles(&rbac.RoleOptions{
 			NoOwnerWorkspaceExec: true,
@@ -185,6 +191,10 @@ func New(options *Options) *API {
 	if options.Authorizer == nil {
 		options.Authorizer = rbac.NewCachingAuthorizer(options.PrometheusRegistry)
 	}
+	// The below are no-ops if already wrapped.
+	if options.PrometheusRegistry != nil {
+		options.Database = dbmetrics.New(options.Database, options.PrometheusRegistry)
+	}
 	options.Database = dbauthz.New(
 		options.Database,
 		options.Authorizer,
diff --git a/coderd/database/db.go b/coderd/database/db.go
@@ -24,11 +24,20 @@ type Store interface {
 	querier
 	// customQuerier contains custom queries that are not generated.
 	customQuerier
+	// wrapper allows us to detect if the interface has been wrapped.
+	wrapper
 
 	Ping(ctx context.Context) (time.Duration, error)
 	InTx(func(Store) error, *sql.TxOptions) error
 }
 
+type wrapper interface {
+	// Wrappers returns a list of wrappers that have been applied to the store.
+	// This is used to detect if the store has already wrapped, and avoid
+	// double-wrapping.
+	Wrappers() []string
+}
+
 // DBTX represents a database connection or transaction.
 type DBTX interface {
 	ExecContext(context.Context, string, ...interface{}) (sql.Result, error)
@@ -60,6 +69,10 @@ type sqlQuerier struct {
 	db  DBTX
 }
 
+func (*sqlQuerier) Wrappers() []string {
+	return []string{}
+}
+
 // Ping returns the time it takes to ping the database.
 func (q *sqlQuerier) Ping(ctx context.Context) (time.Duration, error) {
 	start := time.Now()
diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 
 	"github.com/google/uuid"
+	"golang.org/x/exp/slices"
 	"golang.org/x/xerrors"
 
 	"github.com/open-policy-agent/opa/topdown"
@@ -17,6 +18,8 @@ import (
 
 var _ database.Store = (*querier)(nil)
 
+const wrapname = "dbauthz.querier"
+
 // NoActorError wraps ErrNoRows for the api to return a 404. This is the correct
 // response when the user is not authorized.
 var NoActorError = xerrors.Errorf("no authorization actor in context: %w", sql.ErrNoRows)
@@ -89,7 +92,7 @@ type querier struct {
 func New(db database.Store, authorizer rbac.Authorizer, logger slog.Logger) database.Store {
 	// If the underlying db store is already a querier, return it.
 	// Do not double wrap.
-	if _, ok := db.(*querier); ok {
+	if slices.Contains(db.Wrappers(), wrapname) {
 		return db
 	}
 	return &querier{
@@ -99,6 +102,10 @@ func New(db database.Store, authorizer rbac.Authorizer, logger slog.Logger) data
 	}
 }
 
+func (q *querier) Wrappers() []string {
+	return append(q.db.Wrappers(), wrapname)
+}
+
 // authorizeContext is a helper function to authorize an action on an object.
 func (q *querier) authorizeContext(ctx context.Context, action rbac.Action, object rbac.Objecter) error {
 	act, ok := ActorFromContext(ctx)
diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go
@@ -138,7 +138,7 @@ func TestDBAuthzRecursive(t *testing.T) {
 		for i := 2; i < method.Type.NumIn(); i++ {
 			ins = append(ins, reflect.New(method.Type.In(i)).Elem())
 		}
-		if method.Name == "InTx" || method.Name == "Ping" {
+		if method.Name == "InTx" || method.Name == "Ping" || method.Name == "Wrappers" {
 			continue
 		}
 		// Log the name of the last method, so if there is a panic, it is
diff --git a/coderd/database/dbauthz/setup_test.go b/coderd/database/dbauthz/setup_test.go
@@ -8,6 +8,7 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/golang/mock/gomock"
 	"github.com/google/uuid"
 	"github.com/open-policy-agent/opa/topdown"
 	"github.com/stretchr/testify/require"
@@ -19,14 +20,16 @@ import (
 	"github.com/coder/coder/coderd/database"
 	"github.com/coder/coder/coderd/database/dbauthz"
 	"github.com/coder/coder/coderd/database/dbfake"
+	"github.com/coder/coder/coderd/database/dbmock"
 	"github.com/coder/coder/coderd/rbac"
 	"github.com/coder/coder/coderd/rbac/regosql"
 	"github.com/coder/coder/coderd/util/slice"
 )
 
 var skipMethods = map[string]string{
-	"InTx": "Not relevant",
-	"Ping": "Not relevant",
+	"InTx":     "Not relevant",
+	"Ping":     "Not relevant",
+	"Wrappers": "Not relevant",
 }
 
 // TestMethodTestSuite runs MethodTestSuite.
@@ -52,7 +55,11 @@ type MethodTestSuite struct {
 // SetupSuite sets up the suite by creating a map of all methods on querier
 // and setting their count to 0.
 func (s *MethodTestSuite) SetupSuite() {
-	az := dbauthz.New(nil, nil, slog.Make())
+	ctrl := gomock.NewController(s.T())
+	mockStore := dbmock.NewMockStore(ctrl)
+	// We intentionally set no expectations apart from this.
+	mockStore.EXPECT().Wrappers().Return([]string{}).AnyTimes()
+	az := dbauthz.New(mockStore, nil, slog.Make())
 	// Take the underlying type of the interface.
 	azt := reflect.TypeOf(az).Elem()
 	s.methodAccounting = make(map[string]int)
diff --git a/coderd/database/dbfake/databasefake.go b/coderd/database/dbfake/databasefake.go
@@ -97,6 +97,10 @@ type fakeQuerier struct {
 	*data
 }
 
+func (*fakeQuerier) Wrappers() []string {
+	return []string{}
+}
+
 type fakeTx struct {
 	*fakeQuerier
 	locks map[int64]struct{}
diff --git a/coderd/database/dbmetrics/dbmetrics.go b/coderd/database/dbmetrics/dbmetrics.go
diff --git a/coderd/database/dbmock/store.go b/coderd/database/dbmock/store.go
diff --git a/enterprise/cli/licenses_test.go b/enterprise/cli/licenses_test.go

Original file line number	Diff line number	Diff line change
`@@ -138,7 +138,7 @@ func TestDBAuthzRecursive(t *testing.T) {`
`138`	`138`	`for i := 2; i < method.Type.NumIn(); i++ {`
`139`	`139`	`ins = append(ins, reflect.New(method.Type.In(i)).Elem())`
`140`	`140`	`}`
`141`		`- if method.Name == "InTx" \|\| method.Name == "Ping" {`
	`141`	`+ if method.Name == "InTx" \|\| method.Name == "Ping" \|\| method.Name == "Wrappers" {`
`142`	`142`	`continue`
`143`	`143`	`}`
`144`	`144`	`// Log the name of the last method, so if there is a panic, it is`