Skip to content

Commit 49fcc36

Browse files
johnstcnammario
johnstcn
authored and
ammario
committed
feat: add prometheus metrics to database.Store (coder#7713)
* 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
1 parent cd65755 commit 49fcc36

File tree

10 files changed

+1641
-8
lines changed

10 files changed

+1641
-8
lines changed

‎cli/server.go

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,7 @@ import (
6565
"github.com/coder/coder/coderd/autobuild/executor"
6666
"github.com/coder/coder/coderd/database"
6767
"github.com/coder/coder/coderd/database/dbfake"
68+
"github.com/coder/coder/coderd/database/dbmetrics"
6869
"github.com/coder/coder/coderd/database/dbpurge"
6970
"github.com/coder/coder/coderd/database/migrations"
7071
"github.com/coder/coder/coderd/devtunnel"
@@ -586,7 +587,8 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
586587
}
587588

588589
if cfg.InMemoryDatabase {
589-
options.Database = dbfake.New()
590+
// This is only used for testing.
591+
options.Database = dbmetrics.New(dbfake.New(), options.PrometheusRegistry)
590592
options.Pubsub = database.NewPubsubInMemory()
591593
} else {
592594
sqlDB, err := connectToPostgres(ctx, logger, sqlDriver, cfg.PostgresURL.String())
@@ -597,7 +599,7 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
597599
_ = sqlDB.Close()
598600
}()
599601

600-
options.Database = database.New(sqlDB)
602+
options.Database = dbmetrics.New(database.New(sqlDB), options.PrometheusRegistry)
601603
options.Pubsub, err = database.NewPubsub(ctx, sqlDB, cfg.PostgresURL.String())
602604
if err != nil {
603605
return xerrors.Errorf("create pubsub: %w", err)

‎coderd/coderd.go

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ import (
4747
"github.com/coder/coder/coderd/awsidentity"
4848
"github.com/coder/coder/coderd/database"
4949
"github.com/coder/coder/coderd/database/dbauthz"
50+
"github.com/coder/coder/coderd/database/dbmetrics"
5051
"github.com/coder/coder/coderd/database/dbtype"
5152
"github.com/coder/coder/coderd/gitauth"
5253
"github.com/coder/coder/coderd/gitsshkey"
@@ -176,6 +177,11 @@ func New(options *Options) *API {
176177
options = &Options{}
177178
}
178179

180+
// Safety check: if we're not running a unit test, we *must* have a Prometheus registry.
181+
if options.PrometheusRegistry == nil && flag.Lookup("test.v") == nil {
182+
panic("developer error: options.PrometheusRegistry is nil and not running a unit test")
183+
}
184+
179185
if options.DeploymentValues.DisableOwnerWorkspaceExec {
180186
rbac.ReloadBuiltinRoles(&rbac.RoleOptions{
181187
NoOwnerWorkspaceExec: true,
@@ -185,6 +191,10 @@ func New(options *Options) *API {
185191
if options.Authorizer == nil {
186192
options.Authorizer = rbac.NewCachingAuthorizer(options.PrometheusRegistry)
187193
}
194+
// The below are no-ops if already wrapped.
195+
if options.PrometheusRegistry != nil {
196+
options.Database = dbmetrics.New(options.Database, options.PrometheusRegistry)
197+
}
188198
options.Database = dbauthz.New(
189199
options.Database,
190200
options.Authorizer,

‎coderd/database/db.go

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,11 +24,20 @@ type Store interface {
2424
querier
2525
// customQuerier contains custom queries that are not generated.
2626
customQuerier
27+
// wrapper allows us to detect if the interface has been wrapped.
28+
wrapper
2729

2830
Ping(ctx context.Context) (time.Duration, error)
2931
InTx(func(Store) error, *sql.TxOptions) error
3032
}
3133

34+
type wrapper interface {
35+
// Wrappers returns a list of wrappers that have been applied to the store.
36+
// This is used to detect if the store has already wrapped, and avoid
37+
// double-wrapping.
38+
Wrappers() []string
39+
}
40+
3241
// DBTX represents a database connection or transaction.
3342
type DBTX interface {
3443
ExecContext(context.Context, string, ...interface{}) (sql.Result, error)
@@ -60,6 +69,10 @@ type sqlQuerier struct {
6069
db DBTX
6170
}
6271

72+
func (*sqlQuerier) Wrappers() []string {
73+
return []string{}
74+
}
75+
6376
// Ping returns the time it takes to ping the database.
6477
func (q *sqlQuerier) Ping(ctx context.Context) (time.Duration, error) {
6578
start := time.Now()

‎coderd/database/dbauthz/dbauthz.go

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66
"fmt"
77

88
"github.com/google/uuid"
9+
"golang.org/x/exp/slices"
910
"golang.org/x/xerrors"
1011

1112
"github.com/open-policy-agent/opa/topdown"
@@ -17,6 +18,8 @@ import (
1718

1819
var _ database.Store = (*querier)(nil)
1920

21+
const wrapname = "dbauthz.querier"
22+
2023
// NoActorError wraps ErrNoRows for the api to return a 404. This is the correct
2124
// response when the user is not authorized.
2225
var NoActorError = xerrors.Errorf("no authorization actor in context: %w", sql.ErrNoRows)
@@ -89,7 +92,7 @@ type querier struct {
8992
func New(db database.Store, authorizer rbac.Authorizer, logger slog.Logger) database.Store {
9093
// If the underlying db store is already a querier, return it.
9194
// Do not double wrap.
92-
if _, ok := db.(*querier); ok {
95+
if slices.Contains(db.Wrappers(), wrapname) {
9396
return db
9497
}
9598
return &querier{
@@ -99,6 +102,10 @@ func New(db database.Store, authorizer rbac.Authorizer, logger slog.Logger) data
99102
}
100103
}
101104

105+
func (q *querier) Wrappers() []string {
106+
return append(q.db.Wrappers(), wrapname)
107+
}
108+
102109
// authorizeContext is a helper function to authorize an action on an object.
103110
func (q *querier) authorizeContext(ctx context.Context, action rbac.Action, object rbac.Objecter) error {
104111
act, ok := ActorFromContext(ctx)

‎coderd/database/dbauthz/dbauthz_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -138,7 +138,7 @@ func TestDBAuthzRecursive(t *testing.T) {
138138
for i := 2; i < method.Type.NumIn(); i++ {
139139
ins = append(ins, reflect.New(method.Type.In(i)).Elem())
140140
}
141-
if method.Name == "InTx" || method.Name == "Ping" {
141+
if method.Name == "InTx" || method.Name == "Ping" || method.Name == "Wrappers" {
142142
continue
143143
}
144144
// Log the name of the last method, so if there is a panic, it is

‎coderd/database/dbauthz/setup_test.go

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import (
88
"strings"
99
"testing"
1010

11+
"github.com/golang/mock/gomock"
1112
"github.com/google/uuid"
1213
"github.com/open-policy-agent/opa/topdown"
1314
"github.com/stretchr/testify/require"
@@ -19,14 +20,16 @@ import (
1920
"github.com/coder/coder/coderd/database"
2021
"github.com/coder/coder/coderd/database/dbauthz"
2122
"github.com/coder/coder/coderd/database/dbfake"
23+
"github.com/coder/coder/coderd/database/dbmock"
2224
"github.com/coder/coder/coderd/rbac"
2325
"github.com/coder/coder/coderd/rbac/regosql"
2426
"github.com/coder/coder/coderd/util/slice"
2527
)
2628

2729
var skipMethods = map[string]string{
28-
"InTx": "Not relevant",
29-
"Ping": "Not relevant",
30+
"InTx": "Not relevant",
31+
"Ping": "Not relevant",
32+
"Wrappers": "Not relevant",
3033
}
3134

3235
// TestMethodTestSuite runs MethodTestSuite.
@@ -52,7 +55,11 @@ type MethodTestSuite struct {
5255
// SetupSuite sets up the suite by creating a map of all methods on querier
5356
// and setting their count to 0.
5457
func (s *MethodTestSuite) SetupSuite() {
55-
az := dbauthz.New(nil, nil, slog.Make())
58+
ctrl := gomock.NewController(s.T())
59+
mockStore := dbmock.NewMockStore(ctrl)
60+
// We intentionally set no expectations apart from this.
61+
mockStore.EXPECT().Wrappers().Return([]string{}).AnyTimes()
62+
az := dbauthz.New(mockStore, nil, slog.Make())
5663
// Take the underlying type of the interface.
5764
azt := reflect.TypeOf(az).Elem()
5865
s.methodAccounting = make(map[string]int)

‎coderd/database/dbfake/databasefake.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,10 @@ type fakeQuerier struct {
9797
*data
9898
}
9999

100+
func (*fakeQuerier) Wrappers() []string {
101+
return []string{}
102+
}
103+
100104
type fakeTx struct {
101105
*fakeQuerier
102106
locks map[int64]struct{}

0 commit comments

Comments
 (0)