refactor: introduce ContextInject helper to reduce code duplication

Callum Styan · claude · Callum Styan · commit 8c4a0f5ccdbc · 2025-12-10T00:06:27.000Z
Created CachedWorkspaceFields.ContextInject() to consolidate the repeated pattern of injecting workspace RBAC into context across connectionlog.go, metadata.go, and stats.go. This reduces code duplication and makes the codebase more maintainable by centralizing the RBAC injection logic with consistent error handling. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
diff --git a/coderd/agentapi/cached_workspace.go b/coderd/agentapi/cached_workspace.go
@@ -1,9 +1,12 @@
 package agentapi
 
 import (
+	"context"
 	"sync"
 
+	"cdr.dev/slog"
 	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/database/dbauthz"
 )
 
 // CachedWorkspaceFields contains workspace data that is safe to cache for the
@@ -50,3 +53,26 @@ func (cws *CachedWorkspaceFields) AsWorkspaceIdentity() (database.WorkspaceIdent
 	}
 	return cws.identity, true
 }
+
+// ContextInject attempts to inject the cached workspace RBAC object into the context.
+// This enables the dbauthz fast path and avoids expensive database calls like GetWorkspaceByAgentID.
+// If the cache is empty or RBAC injection fails, the original context is returned.
+//
+// Parameters:
+//   - ctx: the original context
+//   - log: logger for debug messages (optional, can be nil to skip logging)
+//
+// Returns the context with RBAC injected if successful, otherwise the original context.
+func (cws *CachedWorkspaceFields) ContextInject(ctx context.Context, log slog.Logger) context.Context {
+	if dbws, ok := cws.AsWorkspaceIdentity(); ok {
+		rbacCtx, err := dbauthz.WithWorkspaceRBAC(ctx, dbws.RBACObject())
+		if err != nil {
+			// Don't error level log here, will exit the function. We want to fall back to GetWorkspaceByAgentID.
+			//nolint:gocritic
+			log.Debug(ctx, "Cached workspace was present but RBAC object was invalid", slog.F("err", err))
+			return ctx
+		}
+		return rbacCtx
+	}
+	return ctx
+}
diff --git a/coderd/agentapi/connectionlog.go b/coderd/agentapi/connectionlog.go
@@ -14,7 +14,6 @@ import (
 	"github.com/coder/coder/v2/coderd/connectionlog"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/db2sdk"
-	"github.com/coder/coder/v2/coderd/database/dbauthz"
 )
 
 type ConnLogAPI struct {
@@ -55,16 +54,10 @@ func (a *ConnLogAPI) ReportConnection(ctx context.Context, req *agentproto.Repor
 
 	// Inject RBAC object into context for dbauthz fast path, avoid having to
 	// call GetWorkspaceByAgentID on every metadata update.
-	rbacCtx := ctx
+	rbacCtx := a.Workspace.ContextInject(ctx, a.Log)
 	var ws database.WorkspaceIdentity
 	if dbws, ok := a.Workspace.AsWorkspaceIdentity(); ok {
 		ws = dbws
-		rbacCtx, err = dbauthz.WithWorkspaceRBAC(ctx, dbws.RBACObject())
-		if err != nil {
-			// Don't error level log here, will exit the function. We want to fall back to GetWorkspaceByAgentID.
-			//nolint:gocritic
-			a.Log.Debug(ctx, "Cached workspace was present but RBAC object was invalid", slog.F("err", err))
-		}
 	}
 
 	// Fetch contextual data for this connection log event.
diff --git a/coderd/agentapi/metadata.go b/coderd/agentapi/metadata.go
@@ -12,7 +12,6 @@ import (
 	"cdr.dev/slog"
 	agentproto "github.com/coder/coder/v2/agent/proto"
 	"github.com/coder/coder/v2/coderd/database"
-	"github.com/coder/coder/v2/coderd/database/dbauthz"
 	"github.com/coder/coder/v2/coderd/database/dbtime"
 	"github.com/coder/coder/v2/coderd/database/pubsub"
 )
@@ -49,16 +48,7 @@ func (a *MetadataAPI) BatchUpdateMetadata(ctx context.Context, req *agentproto.B
 
 	// Inject RBAC object into context for dbauthz fast path, avoid having to
 	// call GetWorkspaceByAgentID on every metadata update.
-	var err error
-	rbacCtx := ctx
-	if dbws, ok := a.Workspace.AsWorkspaceIdentity(); ok {
-		rbacCtx, err = dbauthz.WithWorkspaceRBAC(ctx, dbws.RBACObject())
-		if err != nil {
-			// Don't error level log here, will exit the function. We want to fall back to GetWorkspaceByAgentID.
-			//nolint:gocritic
-			a.Log.Debug(ctx, "Cached workspace was present but RBAC object was invalid", slog.F("err", err))
-		}
-	}
+	rbacCtx := a.Workspace.ContextInject(ctx, a.Log)
 
 	workspaceAgent, err := a.AgentFn(rbacCtx)
 	if err != nil {
diff --git a/coderd/agentapi/stats.go b/coderd/agentapi/stats.go
@@ -10,7 +10,6 @@ import (
 	"cdr.dev/slog"
 	agentproto "github.com/coder/coder/v2/agent/proto"
 	"github.com/coder/coder/v2/coderd/database"
-	"github.com/coder/coder/v2/coderd/database/dbauthz"
 	"github.com/coder/coder/v2/coderd/database/dbtime"
 	"github.com/coder/coder/v2/coderd/workspacestats"
 	"github.com/coder/coder/v2/codersdk"
@@ -46,17 +45,7 @@ func (a *StatsAPI) UpdateStats(ctx context.Context, req *agentproto.UpdateStatsR
 
 	// Inject RBAC object into context for dbauthz fast path, avoid having to
 	// call GetWorkspaceAgentByID on every stats update.
-
-	rbacCtx := ctx
-	if dbws, ok := a.Workspace.AsWorkspaceIdentity(); ok {
-		var err error
-		rbacCtx, err = dbauthz.WithWorkspaceRBAC(ctx, dbws.RBACObject())
-		if err != nil {
-			// Don't error level log here, will exit the function. We want to fall back to GetWorkspaceByAgentID.
-			//nolint:gocritic
-			a.Log.Debug(ctx, "Cached workspace was present but RBAC object was invalid", slog.F("err", err))
-		}
-	}
+	rbacCtx := a.Workspace.ContextInject(ctx, a.Log)
 
 	workspaceAgent, err := a.AgentFn(rbacCtx)
 	if err != nil {
