chore: generate dbmetrics automatically when adding new queries (#7999)

* chore: rename dbgen package files and remove small file

* chore: automatically generate dbmetrics when new queries are added

coderd/database/dbmetrics/dbmetrics.go

@@ -1,3 +1,6 @@
+// Code generated by coderd/database/gen/metrics.
+// Any function can be edited and will not be overwritten.
+// New database functions are automatically generated!
 package dbmetrics
 
 import (
@@ -70,6 +73,41 @@ func (m metricsStore) InTx(f func(database.Store) error, options *sql.TxOptions)
 	return err
 }
 
+func (m metricsStore) GetAuthorizedTemplates(ctx context.Context, arg database.GetTemplatesWithFilterParams, prepared rbac.PreparedAuthorized) ([]database.Template, error) {
+	start := time.Now()
+	templates, err := m.s.GetAuthorizedTemplates(ctx, arg, prepared)
+	m.queryLatencies.WithLabelValues("GetAuthorizedTemplates").Observe(time.Since(start).Seconds())
+	return templates, err
+}
+
+func (m metricsStore) GetTemplateGroupRoles(ctx context.Context, id uuid.UUID) ([]database.TemplateGroup, error) {
+	start := time.Now()
+	roles, err := m.s.GetTemplateGroupRoles(ctx, id)
+	m.queryLatencies.WithLabelValues("GetTemplateGroupRoles").Observe(time.Since(start).Seconds())
+	return roles, err
+}
+
+func (m metricsStore) GetTemplateUserRoles(ctx context.Context, id uuid.UUID) ([]database.TemplateUser, error) {
+	start := time.Now()
+	roles, err := m.s.GetTemplateUserRoles(ctx, id)
+	m.queryLatencies.WithLabelValues("GetTemplateUserRoles").Observe(time.Since(start).Seconds())
+	return roles, err
+}
+
+func (m metricsStore) GetAuthorizedWorkspaces(ctx context.Context, arg database.GetWorkspacesParams, prepared rbac.PreparedAuthorized) ([]database.GetWorkspacesRow, error) {
+	start := time.Now()
+	workspaces, err := m.s.GetAuthorizedWorkspaces(ctx, arg, prepared)
+	m.queryLatencies.WithLabelValues("GetAuthorizedWorkspaces").Observe(time.Since(start).Seconds())
+	return workspaces, err
+}
+
+func (m metricsStore) GetAuthorizedUserCount(ctx context.Context, arg database.GetFilteredUserCountParams, prepared rbac.PreparedAuthorized) (int64, error) {
+	start := time.Now()
+	count, err := m.s.GetAuthorizedUserCount(ctx, arg, prepared)
+	m.queryLatencies.WithLabelValues("GetAuthorizedUserCount").Observe(time.Since(start).Seconds())
+	return count, err
+}
+
 func (m metricsStore) AcquireLock(ctx context.Context, pgAdvisoryXactLock int64) error {
 	start := time.Now()
 	err := m.s.AcquireLock(ctx, pgAdvisoryXactLock)
@@ -231,6 +269,13 @@ func (m metricsStore) GetDERPMeshKey(ctx context.Context) (string, error) {
 	return key, err
 }
 
+func (m metricsStore) GetDefaultProxyConfig(ctx context.Context) (database.GetDefaultProxyConfigRow, error) {
+	start := time.Now()
+	resp, err := m.s.GetDefaultProxyConfig(ctx)
+	m.queryLatencies.WithLabelValues("GetDefaultProxyConfig").Observe(time.Since(start).Seconds())
+	return resp, err
+}
+
 func (m metricsStore) GetDeploymentDAUs(ctx context.Context, tzOffset int32) ([]database.GetDeploymentDAUsRow, error) {
 	start := time.Now()
 	rows, err := m.s.GetDeploymentDAUs(ctx, tzOffset)
@@ -1360,9 +1405,9 @@ func (m metricsStore) UpdateWorkspaceAgentConnectionByID(ctx context.Context, ar
 
 func (m metricsStore) UpdateWorkspaceAgentLifecycleStateByID(ctx context.Context, arg database.UpdateWorkspaceAgentLifecycleStateByIDParams) error {
 	start := time.Now()
-	err := m.s.UpdateWorkspaceAgentLifecycleStateByID(ctx, arg)
+	r0 := m.s.UpdateWorkspaceAgentLifecycleStateByID(ctx, arg)
 	m.queryLatencies.WithLabelValues("UpdateWorkspaceAgentLifecycleStateByID").Observe(time.Since(start).Seconds())
-	return err
+	return r0
 }
 
 func (m metricsStore) UpdateWorkspaceAgentMetadata(ctx context.Context, arg database.UpdateWorkspaceAgentMetadataParams) error {
@@ -1437,98 +1482,56 @@ func (m metricsStore) UpdateWorkspaceProxy(ctx context.Context, arg database.Upd
 
 func (m metricsStore) UpdateWorkspaceProxyDeleted(ctx context.Context, arg database.UpdateWorkspaceProxyDeletedParams) error {
 	start := time.Now()
-	err := m.s.UpdateWorkspaceProxyDeleted(ctx, arg)
+	r0 := m.s.UpdateWorkspaceProxyDeleted(ctx, arg)
 	m.queryLatencies.WithLabelValues("UpdateWorkspaceProxyDeleted").Observe(time.Since(start).Seconds())
-	return err
+	return r0
 }
 
 func (m metricsStore) UpdateWorkspaceTTL(ctx context.Context, arg database.UpdateWorkspaceTTLParams) error {
 	start := time.Now()
-	err := m.s.UpdateWorkspaceTTL(ctx, arg)
+	r0 := m.s.UpdateWorkspaceTTL(ctx, arg)
 	m.queryLatencies.WithLabelValues("UpdateWorkspaceTTL").Observe(time.Since(start).Seconds())
-	return err
+	return r0
 }
 
 func (m metricsStore) UpdateWorkspaceTTLToBeWithinTemplateMax(ctx context.Context, arg database.UpdateWorkspaceTTLToBeWithinTemplateMaxParams) error {
 	start := time.Now()
-	err := m.s.UpdateWorkspaceTTLToBeWithinTemplateMax(ctx, arg)
+	r0 := m.s.UpdateWorkspaceTTLToBeWithinTemplateMax(ctx, arg)
 	m.queryLatencies.WithLabelValues("UpdateWorkspaceTTLToBeWithinTemplateMax").Observe(time.Since(start).Seconds())
-	return err
+	return r0
 }
 
 func (m metricsStore) UpsertAppSecurityKey(ctx context.Context, value string) error {
 	start := time.Now()
-	err := m.s.UpsertAppSecurityKey(ctx, value)
+	r0 := m.s.UpsertAppSecurityKey(ctx, value)
 	m.queryLatencies.WithLabelValues("UpsertAppSecurityKey").Observe(time.Since(start).Seconds())
-	return err
+	return r0
-}
-
-func (m metricsStore) UpsertLastUpdateCheck(ctx context.Context, value string) error {
-	start := time.Now()
-	err := m.s.UpsertLastUpdateCheck(ctx, value)
-	m.queryLatencies.WithLabelValues("UpsertLastUpdateCheck").Observe(time.Since(start).Seconds())
-	return err
-}
-
-func (m metricsStore) UpsertLogoURL(ctx context.Context, value string) error {
-	start := time.Now()
-	err := m.s.UpsertLogoURL(ctx, value)
-	m.queryLatencies.WithLabelValues("UpsertLogoURL").Observe(time.Since(start).Seconds())
-	return err
-}
-
-func (m metricsStore) UpsertServiceBanner(ctx context.Context, value string) error {
-	start := time.Now()
-	err := m.s.UpsertServiceBanner(ctx, value)
-	m.queryLatencies.WithLabelValues("UpsertServiceBanner").Observe(time.Since(start).Seconds())
-	return err
-}
-
-func (m metricsStore) GetAuthorizedTemplates(ctx context.Context, arg database.GetTemplatesWithFilterParams, prepared rbac.PreparedAuthorized) ([]database.Template, error) {
-	start := time.Now()
-	templates, err := m.s.GetAuthorizedTemplates(ctx, arg, prepared)
-	m.queryLatencies.WithLabelValues("GetAuthorizedTemplates").Observe(time.Since(start).Seconds())
-	return templates, err
-}
-
-func (m metricsStore) GetTemplateGroupRoles(ctx context.Context, id uuid.UUID) ([]database.TemplateGroup, error) {
-	start := time.Now()
-	roles, err := m.s.GetTemplateGroupRoles(ctx, id)
-	m.queryLatencies.WithLabelValues("GetTemplateGroupRoles").Observe(time.Since(start).Seconds())
-	return roles, err
-}
-
-func (m metricsStore) GetTemplateUserRoles(ctx context.Context, id uuid.UUID) ([]database.TemplateUser, error) {
-	start := time.Now()
-	roles, err := m.s.GetTemplateUserRoles(ctx, id)
-	m.queryLatencies.WithLabelValues("GetTemplateUserRoles").Observe(time.Since(start).Seconds())
-	return roles, err
-}
-
-func (m metricsStore) GetAuthorizedWorkspaces(ctx context.Context, arg database.GetWorkspacesParams, prepared rbac.PreparedAuthorized) ([]database.GetWorkspacesRow, error) {
-	start := time.Now()
-	workspaces, err := m.s.GetAuthorizedWorkspaces(ctx, arg, prepared)
-	m.queryLatencies.WithLabelValues("GetAuthorizedWorkspaces").Observe(time.Since(start).Seconds())
-	return workspaces, err
-}
-
-func (m metricsStore) GetAuthorizedUserCount(ctx context.Context, arg database.GetFilteredUserCountParams, prepared rbac.PreparedAuthorized) (int64, error) {
-	start := time.Now()
-	count, err := m.s.GetAuthorizedUserCount(ctx, arg, prepared)
-	m.queryLatencies.WithLabelValues("GetAuthorizedUserCount").Observe(time.Since(start).Seconds())
-	return count, err
 }
 
 func (m metricsStore) UpsertDefaultProxy(ctx context.Context, arg database.UpsertDefaultProxyParams) error {
 	start := time.Now()
-	err := m.s.UpsertDefaultProxy(ctx, arg)
+	r0 := m.s.UpsertDefaultProxy(ctx, arg)
 	m.queryLatencies.WithLabelValues("UpsertDefaultProxy").Observe(time.Since(start).Seconds())
-	return err
+	return r0
 }
 
-func (m metricsStore) GetDefaultProxyConfig(ctx context.Context) (database.GetDefaultProxyConfigRow, error) {
+func (m metricsStore) UpsertLastUpdateCheck(ctx context.Context, value string) error {
 	start := time.Now()
-	resp, err := m.s.GetDefaultProxyConfig(ctx)
+	r0 := m.s.UpsertLastUpdateCheck(ctx, value)
-	m.queryLatencies.WithLabelValues("GetDefaultProxyConfig").Observe(time.Since(start).Seconds())
+	m.queryLatencies.WithLabelValues("UpsertLastUpdateCheck").Observe(time.Since(start).Seconds())
-	return resp, err
+	return r0
+}
+
+func (m metricsStore) UpsertLogoURL(ctx context.Context, value string) error {
+	start := time.Now()
+	r0 := m.s.UpsertLogoURL(ctx, value)
+	m.queryLatencies.WithLabelValues("UpsertLogoURL").Observe(time.Since(start).Seconds())
+	return r0
+}
+
+func (m metricsStore) UpsertServiceBanner(ctx context.Context, value string) error {
+	start := time.Now()
+	r0 := m.s.UpsertServiceBanner(ctx, value)
+	m.queryLatencies.WithLabelValues("UpsertServiceBanner").Observe(time.Since(start).Seconds())
+	return r0
 }

coderd/database/gen/metrics/main.go

@@ -0,0 +1,210 @@
+package main
+
+import (
+	"fmt"
+	"go/format"
+	"go/token"
+	"log"
+	"os"
+	"strings"
+
+	"github.com/dave/dst"
+	"github.com/dave/dst/decorator"
+	"github.com/dave/dst/decorator/resolver/goast"
+	"github.com/dave/dst/decorator/resolver/guess"
+	"golang.org/x/xerrors"
+)
+
+func main() {
+	err := run()
+	if err != nil {
+		log.Fatal(err)
+	}
+}
+
+func run() error {
+	funcs, err := readStoreInterface()
+	if err != nil {
+		return err
+	}
+	funcByName := map[string]struct{}{}
+	for _, f := range funcs {
+		funcByName[f.Name] = struct{}{}
+	}
+	declByName := map[string]*dst.FuncDecl{}
+
+	dbmetrics, err := os.ReadFile("./dbmetrics/dbmetrics.go")
+	if err != nil {
+		return xerrors.Errorf("read dbfake: %w", err)
+	}
+
+	// Required to preserve imports!
+	f, err := decorator.NewDecoratorWithImports(token.NewFileSet(), "dbmetrics", goast.New()).Parse(dbmetrics)
+	if err != nil {
+		return xerrors.Errorf("parse dbfake: %w", err)
+	}
+
+	for i := 0; i < len(f.Decls); i++ {
+		funcDecl, ok := f.Decls[i].(*dst.FuncDecl)
+		if !ok || funcDecl.Recv == nil || len(funcDecl.Recv.List) == 0 {
+			continue
+		}
+		// Check if the receiver is the struct we're interested in
+		_, ok = funcDecl.Recv.List[0].Type.(*dst.Ident)
+		if !ok {
+			continue
+		}
+		if _, ok := funcByName[funcDecl.Name.Name]; !ok {
+			continue
+		}
+		declByName[funcDecl.Name.Name] = funcDecl
+		f.Decls = append(f.Decls[:i], f.Decls[i+1:]...)
+		i--
+	}
+
+	for _, fn := range funcs {
+		decl, ok := declByName[fn.Name]
+		if !ok {
+			params := make([]string, 0)
+			if fn.Func.Params != nil {
+				for _, p := range fn.Func.Params.List {
+					for _, name := range p.Names {
+						params = append(params, name.Name)
+					}
+				}
+			}
+			returns := make([]string, 0)
+			if fn.Func.Results != nil {
+				for i := range fn.Func.Results.List {
+					returns = append(returns, fmt.Sprintf("r%d", i))
+				}
+			}
+
+			code := fmt.Sprintf(`
+package stub
+
+func stub() {
+	start := time.Now()
+	%s := m.s.%s(%s)
+	m.queryLatencies.WithLabelValues("%s").Observe(time.Since(start).Seconds())
+	return %s
+}
+`, strings.Join(returns, ","), fn.Name, strings.Join(params, ","), fn.Name, strings.Join(returns, ","))
+			file, err := decorator.Parse(code)
+			if err != nil {
+				return xerrors.Errorf("parse code: %w", err)
+			}
+			stmt, ok := file.Decls[0].(*dst.FuncDecl)
+			if !ok {
+				return xerrors.Errorf("not ok %T", file.Decls[0])
+			}
+
+			// Not implemented!
+			// When a function isn't implemented, we automatically stub it!
+			decl = &dst.FuncDecl{
+				Name: dst.NewIdent(fn.Name),
+				Type: fn.Func,
+				Recv: &dst.FieldList{
+					List: []*dst.Field{{
+						Names: []*dst.Ident{dst.NewIdent("m")},
+						Type:  dst.NewIdent("metricsStore"),
+					}},
+				},
+				Decs: dst.FuncDeclDecorations{
+					NodeDecs: dst.NodeDecs{
+						Before: dst.EmptyLine,
+						After:  dst.EmptyLine,
+					},
+				},
+				Body: stmt.Body,
+			}
+		}
+		f.Decls = append(f.Decls, decl)
+	}
+
+	file, err := os.OpenFile("./dbmetrics/dbmetrics.go", os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0o755)
+	if err != nil {
+		return xerrors.Errorf("open dbfake: %w", err)
+	}
+	defer file.Close()
+
+	// Required to preserve imports!
+	restorer := decorator.NewRestorerWithImports("dbmetrics", guess.New())
+	restored, err := restorer.RestoreFile(f)
+	if err != nil {
+		return xerrors.Errorf("restore dbfake: %w", err)
+	}
+	err = format.Node(file, restorer.Fset, restored)
+	return err
+}
+
+type storeMethod struct {
+	Name string
+	Func *dst.FuncType
+}
+
+func readStoreInterface() ([]storeMethod, error) {
+	querier, err := os.ReadFile("./querier.go")
+	if err != nil {
+		return nil, xerrors.Errorf("read querier: %w", err)
+	}
+	f, err := decorator.Parse(querier)
+	if err != nil {
+		return nil, err
+	}
+
+	var sqlcQuerier *dst.InterfaceType
+	for _, decl := range f.Decls {
+		genDecl, ok := decl.(*dst.GenDecl)
+		if !ok {
+			continue
+		}
+
+		for _, spec := range genDecl.Specs {
+			typeSpec, ok := spec.(*dst.TypeSpec)
+			if !ok {
+				continue
+			}
+			if typeSpec.Name.Name != "sqlcQuerier" {
+				continue
+			}
+			sqlcQuerier, ok = typeSpec.Type.(*dst.InterfaceType)
+			if !ok {
+				return nil, xerrors.Errorf("unexpected sqlcQuerier type: %T", typeSpec.Type)
+			}
+			break
+		}
+	}
+	if sqlcQuerier == nil {
+		return nil, xerrors.Errorf("sqlcQuerier not found")
+	}
+	funcs := []storeMethod{}
+	for _, method := range sqlcQuerier.Methods.List {
+		funcType, ok := method.Type.(*dst.FuncType)
+		if !ok {
+			continue
+		}
+
+		for _, t := range []*dst.FieldList{funcType.Params, funcType.Results} {
+			if t == nil {
+				continue
+			}
+			for _, f := range t.List {
+				ident, ok := f.Type.(*dst.Ident)
+				if !ok {
+					continue
+				}
+				if !ident.IsExported() {
+					continue
+				}
+				ident.Path = "github.com/coder/coder/coderd/database"
+			}
+		}
+
+		funcs = append(funcs, storeMethod{
+			Name: method.Names[0].Name,
+			Func: funcType,
+		})
+	}
+	return funcs, nil
+}

coderd/database/generate.sh

@@ -62,4 +62,7 @@ SCRIPT_DIR=$(dirname "${BASH_SOURCE[0]}")
 	# Generate the database fake!
 	go run gen/fake/main.go
 	go run golang.org/x/tools/cmd/goimports@latest -w ./dbfake/dbfake.go
+
+	go run gen/metrics/main.go
+	go run golang.org/x/tools/cmd/goimports@latest -w ./dbmetrics/dbmetrics.go
 )