perf: reduce number of queries made by /api/v2/workspaceagents/{id} (#21522)

Relates to https://github.com/coder/internal/issues/1214

The `ExtractWorkspaceAgentParam` middleware ends up making 4 database
queries to follow the chain of `WorkspaceAgent` -> `WorkspaceResource`
-> `ProvisionerJob` -> `WorkspaceBuild` -- but then dropping all that
hard work on the floor. The `api.workspaceAgent` handler that references
this middleware then has to do all of that work again, plus one more
query to get the related `User` so we can get the username. This pattern
is also mirrored in `getDatabaseTerminal` but without the middleware.

This PR:
* Adds a new query `GetWorkspaceAgentAndWorkspaceByID` to fetch all
this information at once to avoid the multiple round-trips,
* Updates the existing usage of `GetWorkspaceAgentByID` to this new
query instead,
* Updates `ExtractWorkspaceAgentParam` to also store the workspace in
the request context

Dalibo: [0.63ms](https://explain.dalibo.com/plan/40bb597f3539gc6c)
This commit is contained in:
Cian Johnston
2026-01-19 12:36:33 +00:00
committed by GitHub
parent d176714f90
commit 08343a7a9f
17 changed files with 259 additions and 230 deletions
+1 -3
View File
@@ -1436,9 +1436,7 @@ func New(options *Options) *API {
Optional: true, Optional: true,
}), }),
httpmw.RequireAPIKeyOrWorkspaceProxyAuth(), httpmw.RequireAPIKeyOrWorkspaceProxyAuth(),
httpmw.ExtractWorkspaceAgentAndWorkspaceParam(options.Database),
httpmw.ExtractWorkspaceAgentParam(options.Database),
httpmw.ExtractWorkspaceParam(options.Database),
) )
r.Get("/", api.workspaceAgent) r.Get("/", api.workspaceAgent)
r.Get("/watch-metadata", api.watchWorkspaceAgentMetadataSSE) r.Get("/watch-metadata", api.watchWorkspaceAgentMetadataSSE)
+1 -1
View File
@@ -569,7 +569,7 @@ func AppSubdomain(dbApp database.WorkspaceApp, agentName, workspaceName, ownerNa
}.String() }.String()
} }
func Apps(dbApps []database.WorkspaceApp, statuses []database.WorkspaceAppStatus, agent database.WorkspaceAgent, ownerName string, workspace database.Workspace) []codersdk.WorkspaceApp { func Apps(dbApps []database.WorkspaceApp, statuses []database.WorkspaceAppStatus, agent database.WorkspaceAgent, ownerName string, workspace database.WorkspaceTable) []codersdk.WorkspaceApp {
sort.Slice(dbApps, func(i, j int) bool { sort.Slice(dbApps, func(i, j int) bool {
if dbApps[i].DisplayOrder != dbApps[j].DisplayOrder { if dbApps[i].DisplayOrder != dbApps[j].DisplayOrder {
return dbApps[i].DisplayOrder < dbApps[j].DisplayOrder return dbApps[i].DisplayOrder < dbApps[j].DisplayOrder
+4
View File
@@ -3614,6 +3614,10 @@ func (q *querier) GetWorkspaceAgentAndLatestBuildByAuthToken(ctx context.Context
return q.db.GetWorkspaceAgentAndLatestBuildByAuthToken(ctx, authToken) return q.db.GetWorkspaceAgentAndLatestBuildByAuthToken(ctx, authToken)
} }
func (q *querier) GetWorkspaceAgentAndWorkspaceByID(ctx context.Context, id uuid.UUID) (database.GetWorkspaceAgentAndWorkspaceByIDRow, error) {
return fetch(q.log, q.auth, q.db.GetWorkspaceAgentAndWorkspaceByID)(ctx, id)
}
func (q *querier) GetWorkspaceAgentByID(ctx context.Context, id uuid.UUID) (database.WorkspaceAgent, error) { func (q *querier) GetWorkspaceAgentByID(ctx context.Context, id uuid.UUID) (database.WorkspaceAgent, error) {
// Fast path: Check if we have a workspace RBAC object in context. // Fast path: Check if we have a workspace RBAC object in context.
// In the agent API this is set at agent connection time to avoid the expensive // In the agent API this is set at agent connection time to avoid the expensive
+5
View File
@@ -1828,6 +1828,11 @@ func (s *MethodTestSuite) TestWorkspace() {
dbm.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agt.ID).Return(agt, nil).AnyTimes() dbm.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agt.ID).Return(agt, nil).AnyTimes()
check.Args(agt.ID).Asserts(w, policy.ActionRead).Returns(agt) check.Args(agt.ID).Asserts(w, policy.ActionRead).Returns(agt)
})) }))
s.Run("GetWorkspaceAgentAndWorkspaceByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
aww := testutil.Fake(s.T(), faker, database.GetWorkspaceAgentAndWorkspaceByIDRow{})
dbm.EXPECT().GetWorkspaceAgentAndWorkspaceByID(gomock.Any(), aww.WorkspaceAgent.ID).Return(aww, nil).AnyTimes()
check.Args(aww.WorkspaceAgent.ID).Asserts(aww.WorkspaceTable, policy.ActionRead).Returns(aww)
}))
s.Run("GetWorkspaceAgentsByWorkspaceAndBuildNumber", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { s.Run("GetWorkspaceAgentsByWorkspaceAndBuildNumber", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
w := testutil.Fake(s.T(), faker, database.Workspace{}) w := testutil.Fake(s.T(), faker, database.Workspace{})
agt := testutil.Fake(s.T(), faker, database.WorkspaceAgent{}) agt := testutil.Fake(s.T(), faker, database.WorkspaceAgent{})
@@ -2215,6 +2215,14 @@ func (m queryMetricsStore) GetWorkspaceAgentAndLatestBuildByAuthToken(ctx contex
return r0, r1 return r0, r1
} }
func (m queryMetricsStore) GetWorkspaceAgentAndWorkspaceByID(ctx context.Context, id uuid.UUID) (database.GetWorkspaceAgentAndWorkspaceByIDRow, error) {
start := time.Now()
r0, r1 := m.s.GetWorkspaceAgentAndWorkspaceByID(ctx, id)
m.queryLatencies.WithLabelValues("GetWorkspaceAgentByIDWithWorkspace").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetWorkspaceAgentByIDWithWorkspace").Inc()
return r0, r1
}
func (m queryMetricsStore) GetWorkspaceAgentByID(ctx context.Context, id uuid.UUID) (database.WorkspaceAgent, error) { func (m queryMetricsStore) GetWorkspaceAgentByID(ctx context.Context, id uuid.UUID) (database.WorkspaceAgent, error) {
start := time.Now() start := time.Now()
r0, r1 := m.s.GetWorkspaceAgentByID(ctx, id) r0, r1 := m.s.GetWorkspaceAgentByID(ctx, id)
+15
View File
@@ -4137,6 +4137,21 @@ func (mr *MockStoreMockRecorder) GetWorkspaceAgentAndLatestBuildByAuthToken(ctx,
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceAgentAndLatestBuildByAuthToken", reflect.TypeOf((*MockStore)(nil).GetWorkspaceAgentAndLatestBuildByAuthToken), ctx, authToken) return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceAgentAndLatestBuildByAuthToken", reflect.TypeOf((*MockStore)(nil).GetWorkspaceAgentAndLatestBuildByAuthToken), ctx, authToken)
} }
// GetWorkspaceAgentAndWorkspaceByID mocks base method.
func (m *MockStore) GetWorkspaceAgentAndWorkspaceByID(ctx context.Context, id uuid.UUID) (database.GetWorkspaceAgentAndWorkspaceByIDRow, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetWorkspaceAgentAndWorkspaceByID", ctx, id)
ret0, _ := ret[0].(database.GetWorkspaceAgentAndWorkspaceByIDRow)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetWorkspaceAgentAndWorkspaceByID indicates an expected call of GetWorkspaceAgentAndWorkspaceByID.
func (mr *MockStoreMockRecorder) GetWorkspaceAgentAndWorkspaceByID(ctx, id any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceAgentAndWorkspaceByID", reflect.TypeOf((*MockStore)(nil).GetWorkspaceAgentAndWorkspaceByID), ctx, id)
}
// GetWorkspaceAgentByID mocks base method. // GetWorkspaceAgentByID mocks base method.
func (m *MockStore) GetWorkspaceAgentByID(ctx context.Context, id uuid.UUID) (database.WorkspaceAgent, error) { func (m *MockStore) GetWorkspaceAgentByID(ctx context.Context, id uuid.UUID) (database.WorkspaceAgent, error) {
m.ctrl.T.Helper() m.ctrl.T.Helper()
+5
View File
@@ -884,3 +884,8 @@ func WorkspaceIdentityFromWorkspace(w Workspace) WorkspaceIdentity {
AutostartSchedule: w.AutostartSchedule, AutostartSchedule: w.AutostartSchedule,
} }
} }
// A workspace agent belongs to the owner of the associated workspace.
func (r GetWorkspaceAgentAndWorkspaceByIDRow) RBACObject() rbac.Object {
return r.WorkspaceTable.RBACObject()
}
+1
View File
@@ -467,6 +467,7 @@ type sqlcQuerier interface {
GetWebpushVAPIDKeys(ctx context.Context) (GetWebpushVAPIDKeysRow, error) GetWebpushVAPIDKeys(ctx context.Context) (GetWebpushVAPIDKeysRow, error)
GetWorkspaceACLByID(ctx context.Context, id uuid.UUID) (GetWorkspaceACLByIDRow, error) GetWorkspaceACLByID(ctx context.Context, id uuid.UUID) (GetWorkspaceACLByIDRow, error)
GetWorkspaceAgentAndLatestBuildByAuthToken(ctx context.Context, authToken uuid.UUID) (GetWorkspaceAgentAndLatestBuildByAuthTokenRow, error) GetWorkspaceAgentAndLatestBuildByAuthToken(ctx context.Context, authToken uuid.UUID) (GetWorkspaceAgentAndLatestBuildByAuthTokenRow, error)
GetWorkspaceAgentAndWorkspaceByID(ctx context.Context, id uuid.UUID) (GetWorkspaceAgentAndWorkspaceByIDRow, error)
GetWorkspaceAgentByID(ctx context.Context, id uuid.UUID) (WorkspaceAgent, error) GetWorkspaceAgentByID(ctx context.Context, id uuid.UUID) (WorkspaceAgent, error)
GetWorkspaceAgentByInstanceID(ctx context.Context, authInstanceID string) (WorkspaceAgent, error) GetWorkspaceAgentByInstanceID(ctx context.Context, authInstanceID string) (WorkspaceAgent, error)
GetWorkspaceAgentDevcontainersByAgentID(ctx context.Context, workspaceAgentID uuid.UUID) ([]WorkspaceAgentDevcontainer, error) GetWorkspaceAgentDevcontainersByAgentID(ctx context.Context, workspaceAgentID uuid.UUID) ([]WorkspaceAgentDevcontainer, error)
+93
View File
@@ -18155,6 +18155,99 @@ func (q *sqlQuerier) GetWorkspaceAgentAndLatestBuildByAuthToken(ctx context.Cont
return i, err return i, err
} }
const getWorkspaceAgentAndWorkspaceByID = `-- name: GetWorkspaceAgentAndWorkspaceByID :one
SELECT
workspace_agents.id, workspace_agents.created_at, workspace_agents.updated_at, workspace_agents.name, workspace_agents.first_connected_at, workspace_agents.last_connected_at, workspace_agents.disconnected_at, workspace_agents.resource_id, workspace_agents.auth_token, workspace_agents.auth_instance_id, workspace_agents.architecture, workspace_agents.environment_variables, workspace_agents.operating_system, workspace_agents.instance_metadata, workspace_agents.resource_metadata, workspace_agents.directory, workspace_agents.version, workspace_agents.last_connected_replica_id, workspace_agents.connection_timeout_seconds, workspace_agents.troubleshooting_url, workspace_agents.motd_file, workspace_agents.lifecycle_state, workspace_agents.expanded_directory, workspace_agents.logs_length, workspace_agents.logs_overflowed, workspace_agents.started_at, workspace_agents.ready_at, workspace_agents.subsystems, workspace_agents.display_apps, workspace_agents.api_version, workspace_agents.display_order, workspace_agents.parent_id, workspace_agents.api_key_scope, workspace_agents.deleted,
workspaces.id, workspaces.created_at, workspaces.updated_at, workspaces.owner_id, workspaces.organization_id, workspaces.template_id, workspaces.deleted, workspaces.name, workspaces.autostart_schedule, workspaces.ttl, workspaces.last_used_at, workspaces.dormant_at, workspaces.deleting_at, workspaces.automatic_updates, workspaces.favorite, workspaces.next_start_at, workspaces.group_acl, workspaces.user_acl,
users.username as owner_username
FROM
workspace_agents
JOIN
workspace_resources ON workspace_agents.resource_id = workspace_resources.id
JOIN
provisioner_jobs ON workspace_resources.job_id = provisioner_jobs.id
JOIN
workspace_builds ON provisioner_jobs.id = workspace_builds.job_id
JOIN
workspaces ON workspace_builds.workspace_id = workspaces.id
JOIN
users ON workspaces.owner_id = users.id
WHERE
workspace_agents.id = $1
AND workspace_agents.deleted = FALSE
AND provisioner_jobs.type = 'workspace_build'::provisioner_job_type
AND workspaces.deleted = FALSE
AND users.deleted = FALSE
LIMIT 1
`
type GetWorkspaceAgentAndWorkspaceByIDRow struct {
WorkspaceAgent WorkspaceAgent `db:"workspace_agent" json:"workspace_agent"`
WorkspaceTable WorkspaceTable `db:"workspace_table" json:"workspace_table"`
OwnerUsername string `db:"owner_username" json:"owner_username"`
}
func (q *sqlQuerier) GetWorkspaceAgentAndWorkspaceByID(ctx context.Context, id uuid.UUID) (GetWorkspaceAgentAndWorkspaceByIDRow, error) {
row := q.db.QueryRowContext(ctx, getWorkspaceAgentAndWorkspaceByID, id)
var i GetWorkspaceAgentAndWorkspaceByIDRow
err := row.Scan(
&i.WorkspaceAgent.ID,
&i.WorkspaceAgent.CreatedAt,
&i.WorkspaceAgent.UpdatedAt,
&i.WorkspaceAgent.Name,
&i.WorkspaceAgent.FirstConnectedAt,
&i.WorkspaceAgent.LastConnectedAt,
&i.WorkspaceAgent.DisconnectedAt,
&i.WorkspaceAgent.ResourceID,
&i.WorkspaceAgent.AuthToken,
&i.WorkspaceAgent.AuthInstanceID,
&i.WorkspaceAgent.Architecture,
&i.WorkspaceAgent.EnvironmentVariables,
&i.WorkspaceAgent.OperatingSystem,
&i.WorkspaceAgent.InstanceMetadata,
&i.WorkspaceAgent.ResourceMetadata,
&i.WorkspaceAgent.Directory,
&i.WorkspaceAgent.Version,
&i.WorkspaceAgent.LastConnectedReplicaID,
&i.WorkspaceAgent.ConnectionTimeoutSeconds,
&i.WorkspaceAgent.TroubleshootingURL,
&i.WorkspaceAgent.MOTDFile,
&i.WorkspaceAgent.LifecycleState,
&i.WorkspaceAgent.ExpandedDirectory,
&i.WorkspaceAgent.LogsLength,
&i.WorkspaceAgent.LogsOverflowed,
&i.WorkspaceAgent.StartedAt,
&i.WorkspaceAgent.ReadyAt,
pq.Array(&i.WorkspaceAgent.Subsystems),
pq.Array(&i.WorkspaceAgent.DisplayApps),
&i.WorkspaceAgent.APIVersion,
&i.WorkspaceAgent.DisplayOrder,
&i.WorkspaceAgent.ParentID,
&i.WorkspaceAgent.APIKeyScope,
&i.WorkspaceAgent.Deleted,
&i.WorkspaceTable.ID,
&i.WorkspaceTable.CreatedAt,
&i.WorkspaceTable.UpdatedAt,
&i.WorkspaceTable.OwnerID,
&i.WorkspaceTable.OrganizationID,
&i.WorkspaceTable.TemplateID,
&i.WorkspaceTable.Deleted,
&i.WorkspaceTable.Name,
&i.WorkspaceTable.AutostartSchedule,
&i.WorkspaceTable.Ttl,
&i.WorkspaceTable.LastUsedAt,
&i.WorkspaceTable.DormantAt,
&i.WorkspaceTable.DeletingAt,
&i.WorkspaceTable.AutomaticUpdates,
&i.WorkspaceTable.Favorite,
&i.WorkspaceTable.NextStartAt,
&i.WorkspaceTable.GroupACL,
&i.WorkspaceTable.UserACL,
&i.OwnerUsername,
)
return i, err
}
const getWorkspaceAgentByID = `-- name: GetWorkspaceAgentByID :one const getWorkspaceAgentByID = `-- name: GetWorkspaceAgentByID :one
SELECT SELECT
id, created_at, updated_at, name, first_connected_at, last_connected_at, disconnected_at, resource_id, auth_token, auth_instance_id, architecture, environment_variables, operating_system, instance_metadata, resource_metadata, directory, version, last_connected_replica_id, connection_timeout_seconds, troubleshooting_url, motd_file, lifecycle_state, expanded_directory, logs_length, logs_overflowed, started_at, ready_at, subsystems, display_apps, api_version, display_order, parent_id, api_key_scope, deleted id, created_at, updated_at, name, first_connected_at, last_connected_at, disconnected_at, resource_id, auth_token, auth_instance_id, architecture, environment_variables, operating_system, instance_metadata, resource_metadata, directory, version, last_connected_replica_id, connection_timeout_seconds, troubleshooting_url, motd_file, lifecycle_state, expanded_directory, logs_length, logs_overflowed, started_at, ready_at, subsystems, display_apps, api_version, display_order, parent_id, api_key_scope, deleted
@@ -393,3 +393,28 @@ AND wb.build_number = (
WHERE wb2.workspace_id = w.id WHERE wb2.workspace_id = w.id
) )
AND workspace_agents.deleted = FALSE; AND workspace_agents.deleted = FALSE;
-- name: GetWorkspaceAgentAndWorkspaceByID :one
SELECT
sqlc.embed(workspace_agents),
sqlc.embed(workspaces),
users.username as owner_username
FROM
workspace_agents
JOIN
workspace_resources ON workspace_agents.resource_id = workspace_resources.id
JOIN
provisioner_jobs ON workspace_resources.job_id = provisioner_jobs.id
JOIN
workspace_builds ON provisioner_jobs.id = workspace_builds.job_id
JOIN
workspaces ON workspace_builds.workspace_id = workspaces.id
JOIN
users ON workspaces.owner_id = users.id
WHERE
workspace_agents.id = @id
AND workspace_agents.deleted = FALSE
AND provisioner_jobs.type = 'workspace_build'::provisioner_job_type
AND workspaces.deleted = FALSE
AND users.deleted = FALSE
LIMIT 1;
+12 -51
View File
@@ -13,19 +13,19 @@ import (
"github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk"
) )
type workspaceAgentParamContextKey struct{} type workspaceAgentAndWorkspaceParamContextKey struct{}
// WorkspaceAgentParam returns the workspace agent from the ExtractWorkspaceAgentParam handler. // WorkspaceAgentAndWorkspaceParam returns the workspace agent and its associated workspace from the ExtractWorkspaceAgentParam handler.
func WorkspaceAgentParam(r *http.Request) database.WorkspaceAgent { func WorkspaceAgentAndWorkspaceParam(r *http.Request) database.GetWorkspaceAgentAndWorkspaceByIDRow {
user, ok := r.Context().Value(workspaceAgentParamContextKey{}).(database.WorkspaceAgent) aw, ok := r.Context().Value(workspaceAgentAndWorkspaceParamContextKey{}).(database.GetWorkspaceAgentAndWorkspaceByIDRow)
if !ok { if !ok {
panic("developer error: agent middleware not provided") panic("developer error: agent middleware not provided")
} }
return user return aw
} }
// ExtractWorkspaceAgentParam grabs a workspace agent from the "workspaceagent" URL parameter. // ExtractWorkspaceAgentAndWorkspaceParam grabs a workspace agent and its associated workspace from the "workspaceagent" URL parameter.
func ExtractWorkspaceAgentParam(db database.Store) func(http.Handler) http.Handler { func ExtractWorkspaceAgentAndWorkspaceParam(db database.Store) func(http.Handler) http.Handler {
return func(next http.Handler) http.Handler { return func(next http.Handler) http.Handler {
return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
ctx := r.Context() ctx := r.Context()
@@ -34,60 +34,21 @@ func ExtractWorkspaceAgentParam(db database.Store) func(http.Handler) http.Handl
return return
} }
agent, err := db.GetWorkspaceAgentByID(ctx, agentUUID) agentWithWorkspace, err := db.GetWorkspaceAgentAndWorkspaceByID(ctx, agentUUID)
if httpapi.Is404Error(err) { if httpapi.Is404Error(err) {
httpapi.Write(ctx, rw, http.StatusNotFound, codersdk.Response{ httpapi.Write(ctx, rw, http.StatusNotFound, codersdk.Response{
Message: "Agent doesn't exist with that id, or you do not have access to it.", Message: "Agent doesn't exist with that id, or you do not have access to it.",
}) })
return return
} }
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error fetching workspace agent.",
Detail: err.Error(),
})
return
}
resource, err := db.GetWorkspaceResourceByID(ctx, agent.ResourceID) ctx = context.WithValue(ctx, workspaceAgentAndWorkspaceParamContextKey{}, agentWithWorkspace)
if err != nil { chi.RouteContext(ctx).URLParams.Add("workspace", agentWithWorkspace.WorkspaceTable.ID.String())
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error fetching workspace resource.",
Detail: err.Error(),
})
return
}
job, err := db.GetProvisionerJobByID(ctx, resource.JobID)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error fetching provisioner job.",
Detail: err.Error(),
})
return
}
if job.Type != database.ProvisionerJobTypeWorkspaceBuild {
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: "Workspace agents can only be fetched for builds.",
})
return
}
build, err := db.GetWorkspaceBuildByJobID(ctx, job.ID)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error fetching workspace build.",
Detail: err.Error(),
})
return
}
ctx = context.WithValue(ctx, workspaceAgentParamContextKey{}, agent)
chi.RouteContext(ctx).URLParams.Add("workspace", build.WorkspaceID.String())
if rlogger := loggermw.RequestLoggerFromContext(ctx); rlogger != nil { if rlogger := loggermw.RequestLoggerFromContext(ctx); rlogger != nil {
rlogger.WithFields( rlogger.WithFields(
slog.F("workspace_name", resource.Name), slog.F("workspace_name", agentWithWorkspace.WorkspaceTable.Name),
slog.F("agent_name", agent.Name), slog.F("agent_name", agentWithWorkspace.WorkspaceAgent.Name),
) )
} }
+5 -5
View File
@@ -86,7 +86,7 @@ func TestWorkspaceAgentParam(t *testing.T) {
db, _ := dbtestutil.NewDB(t) db, _ := dbtestutil.NewDB(t)
dbtestutil.DisableForeignKeysAndTriggers(t, db) dbtestutil.DisableForeignKeysAndTriggers(t, db)
rtr := chi.NewRouter() rtr := chi.NewRouter()
rtr.Use(httpmw.ExtractWorkspaceAgentParam(db)) rtr.Use(httpmw.ExtractWorkspaceAgentAndWorkspaceParam(db))
rtr.Get("/", nil) rtr.Get("/", nil)
r, _ := setupAuthentication(db) r, _ := setupAuthentication(db)
@@ -113,10 +113,10 @@ func TestWorkspaceAgentParam(t *testing.T) {
RedirectToLogin: false, RedirectToLogin: false,
}), }),
// Only fail authz in this middleware // Only fail authz in this middleware
httpmw.ExtractWorkspaceAgentParam(dbFail), httpmw.ExtractWorkspaceAgentAndWorkspaceParam(dbFail),
) )
rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) { rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) {
_ = httpmw.WorkspaceAgentParam(r) _ = httpmw.WorkspaceAgentAndWorkspaceParam(r)
rw.WriteHeader(http.StatusOK) rw.WriteHeader(http.StatusOK)
}) })
@@ -140,10 +140,10 @@ func TestWorkspaceAgentParam(t *testing.T) {
DB: db, DB: db,
RedirectToLogin: false, RedirectToLogin: false,
}), }),
httpmw.ExtractWorkspaceAgentParam(db), httpmw.ExtractWorkspaceAgentAndWorkspaceParam(db),
) )
rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) { rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) {
_ = httpmw.WorkspaceAgentParam(r) _ = httpmw.WorkspaceAgentAndWorkspaceParam(r)
rw.WriteHeader(http.StatusOK) rw.WriteHeader(http.StatusOK)
}) })
+48 -92
View File
@@ -60,10 +60,9 @@ import (
// @Success 200 {object} codersdk.WorkspaceAgent // @Success 200 {object} codersdk.WorkspaceAgent
// @Router /workspaceagents/{workspaceagent} [get] // @Router /workspaceagents/{workspaceagent} [get]
func (api *API) workspaceAgent(rw http.ResponseWriter, r *http.Request) { func (api *API) workspaceAgent(rw http.ResponseWriter, r *http.Request) {
ctx := r.Context()
workspaceAgent := httpmw.WorkspaceAgentParam(r)
var ( var (
ctx = r.Context()
waws = httpmw.WorkspaceAgentAndWorkspaceParam(r)
dbApps []database.WorkspaceApp dbApps []database.WorkspaceApp
scripts []database.WorkspaceAgentScript scripts []database.WorkspaceAgentScript
logSources []database.WorkspaceAgentLogSource logSources []database.WorkspaceAgentLogSource
@@ -71,17 +70,17 @@ func (api *API) workspaceAgent(rw http.ResponseWriter, r *http.Request) {
var eg errgroup.Group var eg errgroup.Group
eg.Go(func() (err error) { eg.Go(func() (err error) {
dbApps, err = api.Database.GetWorkspaceAppsByAgentID(ctx, workspaceAgent.ID) dbApps, err = api.Database.GetWorkspaceAppsByAgentID(ctx, waws.WorkspaceAgent.ID)
return err return err
}) })
eg.Go(func() (err error) { eg.Go(func() (err error) {
//nolint:gocritic // TODO: can we make this not require system restricted? //nolint:gocritic // TODO: can we make this not require system restricted?
scripts, err = api.Database.GetWorkspaceAgentScriptsByAgentIDs(dbauthz.AsSystemRestricted(ctx), []uuid.UUID{workspaceAgent.ID}) scripts, err = api.Database.GetWorkspaceAgentScriptsByAgentIDs(dbauthz.AsSystemRestricted(ctx), []uuid.UUID{waws.WorkspaceAgent.ID})
return err return err
}) })
eg.Go(func() (err error) { eg.Go(func() (err error) {
//nolint:gocritic // TODO: can we make this not require system restricted? //nolint:gocritic // TODO: can we make this not require system restricted?
logSources, err = api.Database.GetWorkspaceAgentLogSourcesByAgentIDs(dbauthz.AsSystemRestricted(ctx), []uuid.UUID{workspaceAgent.ID}) logSources, err = api.Database.GetWorkspaceAgentLogSourcesByAgentIDs(dbauthz.AsSystemRestricted(ctx), []uuid.UUID{waws.WorkspaceAgent.ID})
return err return err
}) })
err := eg.Wait() err := eg.Wait()
@@ -111,41 +110,8 @@ func (api *API) workspaceAgent(rw http.ResponseWriter, r *http.Request) {
return return
} }
resource, err := api.Database.GetWorkspaceResourceByID(ctx, workspaceAgent.ResourceID)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error fetching workspace resource.",
Detail: err.Error(),
})
return
}
build, err := api.Database.GetWorkspaceBuildByJobID(ctx, resource.JobID)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error fetching workspace build.",
Detail: err.Error(),
})
return
}
workspace, err := api.Database.GetWorkspaceByID(ctx, build.WorkspaceID)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error fetching workspace.",
Detail: err.Error(),
})
return
}
owner, err := api.Database.GetUserByID(ctx, workspace.OwnerID)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error fetching workspace owner.",
Detail: err.Error(),
})
return
}
apiAgent, err := db2sdk.WorkspaceAgent( apiAgent, err := db2sdk.WorkspaceAgent(
api.DERPMap(), *api.TailnetCoordinator.Load(), workspaceAgent, db2sdk.Apps(dbApps, statuses, workspaceAgent, owner.Username, workspace), convertScripts(scripts), convertLogSources(logSources), api.AgentInactiveDisconnectTimeout, api.DERPMap(), *api.TailnetCoordinator.Load(), waws.WorkspaceAgent, db2sdk.Apps(dbApps, statuses, waws.WorkspaceAgent, waws.OwnerUsername, waws.WorkspaceTable), convertScripts(scripts), convertLogSources(logSources), api.AgentInactiveDisconnectTimeout,
api.DeploymentValues.AgentFallbackTroubleshootingURL.String(), api.DeploymentValues.AgentFallbackTroubleshootingURL.String(),
) )
if err != nil { if err != nil {
@@ -546,12 +512,12 @@ func (api *API) enqueueAITaskStateNotification(
func (api *API) workspaceAgentLogs(rw http.ResponseWriter, r *http.Request) { func (api *API) workspaceAgentLogs(rw http.ResponseWriter, r *http.Request) {
// This mostly copies how provisioner job logs are streamed! // This mostly copies how provisioner job logs are streamed!
var ( var (
ctx = r.Context() ctx = r.Context()
workspaceAgent = httpmw.WorkspaceAgentParam(r) waws = httpmw.WorkspaceAgentAndWorkspaceParam(r)
logger = api.Logger.With(slog.F("workspace_agent_id", workspaceAgent.ID)) logger = api.Logger.With(slog.F("workspace_agent_id", waws.WorkspaceAgent.ID))
follow = r.URL.Query().Has("follow") follow = r.URL.Query().Has("follow")
afterRaw = r.URL.Query().Get("after") afterRaw = r.URL.Query().Get("after")
noCompression = r.URL.Query().Has("no_compression") noCompression = r.URL.Query().Has("no_compression")
) )
var after int64 var after int64
@@ -571,7 +537,7 @@ func (api *API) workspaceAgentLogs(rw http.ResponseWriter, r *http.Request) {
} }
logs, err := api.Database.GetWorkspaceAgentLogsAfter(ctx, database.GetWorkspaceAgentLogsAfterParams{ logs, err := api.Database.GetWorkspaceAgentLogsAfter(ctx, database.GetWorkspaceAgentLogsAfterParams{
AgentID: workspaceAgent.ID, AgentID: waws.WorkspaceAgent.ID,
CreatedAfter: after, CreatedAfter: after,
}) })
if errors.Is(err, sql.ErrNoRows) { if errors.Is(err, sql.ErrNoRows) {
@@ -593,15 +559,6 @@ func (api *API) workspaceAgentLogs(rw http.ResponseWriter, r *http.Request) {
return return
} }
workspace, err := api.Database.GetWorkspaceByAgentID(ctx, workspaceAgent.ID)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error fetching workspace by agent id.",
Detail: err.Error(),
})
return
}
api.WebsocketWaitMutex.Lock() api.WebsocketWaitMutex.Lock()
api.WebsocketWaitGroup.Add(1) api.WebsocketWaitGroup.Add(1)
api.WebsocketWaitMutex.Unlock() api.WebsocketWaitMutex.Unlock()
@@ -650,13 +607,13 @@ func (api *API) workspaceAgentLogs(rw http.ResponseWriter, r *http.Request) {
notifyCh <- struct{}{} notifyCh <- struct{}{}
// Subscribe to workspace to detect new builds. // Subscribe to workspace to detect new builds.
closeSubscribeWorkspace, err := api.Pubsub.SubscribeWithErr(wspubsub.WorkspaceEventChannel(workspace.OwnerID), closeSubscribeWorkspace, err := api.Pubsub.SubscribeWithErr(wspubsub.WorkspaceEventChannel(waws.WorkspaceTable.OwnerID),
wspubsub.HandleWorkspaceEvent( wspubsub.HandleWorkspaceEvent(
func(_ context.Context, e wspubsub.WorkspaceEvent, err error) { func(_ context.Context, e wspubsub.WorkspaceEvent, err error) {
if err != nil { if err != nil {
return return
} }
if e.Kind == wspubsub.WorkspaceEventKindStateChange && e.WorkspaceID == workspace.ID { if e.Kind == wspubsub.WorkspaceEventKindStateChange && e.WorkspaceID == waws.WorkspaceTable.ID {
select { select {
case workspaceNotifyCh <- struct{}{}: case workspaceNotifyCh <- struct{}{}:
default: default:
@@ -672,7 +629,7 @@ func (api *API) workspaceAgentLogs(rw http.ResponseWriter, r *http.Request) {
} }
defer closeSubscribeWorkspace() defer closeSubscribeWorkspace()
// Subscribe early to prevent missing log events. // Subscribe early to prevent missing log events.
closeSubscribe, err := api.Pubsub.Subscribe(agentsdk.LogsNotifyChannel(workspaceAgent.ID), func(_ context.Context, _ []byte) { closeSubscribe, err := api.Pubsub.Subscribe(agentsdk.LogsNotifyChannel(waws.WorkspaceAgent.ID), func(_ context.Context, _ []byte) {
// The message is not important, we're tracking lastSentLogID manually. // The message is not important, we're tracking lastSentLogID manually.
select { select {
case notifyCh <- struct{}{}: case notifyCh <- struct{}{}:
@@ -726,7 +683,7 @@ func (api *API) workspaceAgentLogs(rw http.ResponseWriter, r *http.Request) {
t.Reset(recheckInterval) t.Reset(recheckInterval)
} }
agents, err := api.Database.GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx, workspace.ID) agents, err := api.Database.GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx, waws.WorkspaceTable.ID)
if err != nil && !xerrors.Is(err, sql.ErrNoRows) { if err != nil && !xerrors.Is(err, sql.ErrNoRows) {
if xerrors.Is(err, context.Canceled) { if xerrors.Is(err, context.Canceled) {
return return
@@ -736,7 +693,7 @@ func (api *API) workspaceAgentLogs(rw http.ResponseWriter, r *http.Request) {
} }
// If the agent is no longer in the latest build, we can stop after // If the agent is no longer in the latest build, we can stop after
// checking once. // checking once.
keepGoing = slices.ContainsFunc(agents, func(agent database.WorkspaceAgent) bool { return agent.ID == workspaceAgent.ID }) keepGoing = slices.ContainsFunc(agents, func(agent database.WorkspaceAgent) bool { return agent.ID == waws.WorkspaceAgent.ID })
logger.Debug( logger.Debug(
ctx, ctx,
@@ -753,7 +710,7 @@ func (api *API) workspaceAgentLogs(rw http.ResponseWriter, r *http.Request) {
} }
logs, err := api.Database.GetWorkspaceAgentLogsAfter(ctx, database.GetWorkspaceAgentLogsAfterParams{ logs, err := api.Database.GetWorkspaceAgentLogsAfter(ctx, database.GetWorkspaceAgentLogsAfterParams{
AgentID: workspaceAgent.ID, AgentID: waws.WorkspaceAgent.ID,
CreatedAfter: lastSentLogID, CreatedAfter: lastSentLogID,
}) })
if err != nil { if err != nil {
@@ -816,7 +773,7 @@ func (api *API) workspaceAgentLogs(rw http.ResponseWriter, r *http.Request) {
// @Router /workspaceagents/{workspaceagent}/listening-ports [get] // @Router /workspaceagents/{workspaceagent}/listening-ports [get]
func (api *API) workspaceAgentListeningPorts(rw http.ResponseWriter, r *http.Request) { func (api *API) workspaceAgentListeningPorts(rw http.ResponseWriter, r *http.Request) {
ctx := r.Context() ctx := r.Context()
workspaceAgent := httpmw.WorkspaceAgentParam(r) waws := httpmw.WorkspaceAgentAndWorkspaceParam(r)
// If the agent is unreachable, the request will hang. Assume that if we // If the agent is unreachable, the request will hang. Assume that if we
// don't get a response after 30s that the agent is unreachable. // don't get a response after 30s that the agent is unreachable.
@@ -824,7 +781,7 @@ func (api *API) workspaceAgentListeningPorts(rw http.ResponseWriter, r *http.Req
defer cancel() defer cancel()
apiAgent, err := db2sdk.WorkspaceAgent( apiAgent, err := db2sdk.WorkspaceAgent(
api.DERPMap(), *api.TailnetCoordinator.Load(), workspaceAgent, nil, nil, nil, api.AgentInactiveDisconnectTimeout, api.DERPMap(), *api.TailnetCoordinator.Load(), waws.WorkspaceAgent, nil, nil, nil, api.AgentInactiveDisconnectTimeout,
api.DeploymentValues.AgentFallbackTroubleshootingURL.String(), api.DeploymentValues.AgentFallbackTroubleshootingURL.String(),
) )
if err != nil { if err != nil {
@@ -841,7 +798,7 @@ func (api *API) workspaceAgentListeningPorts(rw http.ResponseWriter, r *http.Req
return return
} }
agentConn, release, err := api.agentProvider.AgentConn(ctx, workspaceAgent.ID) agentConn, release, err := api.agentProvider.AgentConn(ctx, waws.WorkspaceAgent.ID)
if err != nil { if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error dialing workspace agent.", Message: "Internal error dialing workspace agent.",
@@ -861,7 +818,7 @@ func (api *API) workspaceAgentListeningPorts(rw http.ResponseWriter, r *http.Req
} }
// Get a list of ports that are in-use by applications. // Get a list of ports that are in-use by applications.
apps, err := api.Database.GetWorkspaceAppsByAgentID(ctx, workspaceAgent.ID) apps, err := api.Database.GetWorkspaceAppsByAgentID(ctx, waws.WorkspaceAgent.ID)
if xerrors.Is(err, sql.ErrNoRows) { if xerrors.Is(err, sql.ErrNoRows) {
apps = []database.WorkspaceApp{} apps = []database.WorkspaceApp{}
err = nil err = nil
@@ -926,9 +883,9 @@ func (api *API) workspaceAgentListeningPorts(rw http.ResponseWriter, r *http.Req
// @Router /workspaceagents/{workspaceagent}/containers/watch [get] // @Router /workspaceagents/{workspaceagent}/containers/watch [get]
func (api *API) watchWorkspaceAgentContainers(rw http.ResponseWriter, r *http.Request) { func (api *API) watchWorkspaceAgentContainers(rw http.ResponseWriter, r *http.Request) {
var ( var (
ctx = r.Context() ctx = r.Context()
workspaceAgent = httpmw.WorkspaceAgentParam(r) waws = httpmw.WorkspaceAgentAndWorkspaceParam(r)
logger = api.Logger.Named("agent_container_watcher").With(slog.F("agent_id", workspaceAgent.ID)) logger = api.Logger.Named("agent_container_watcher").With(slog.F("agent_id", waws.WorkspaceAgent.ID))
) )
// If the agent is unreachable, the request will hang. Assume that if we // If the agent is unreachable, the request will hang. Assume that if we
@@ -938,7 +895,7 @@ func (api *API) watchWorkspaceAgentContainers(rw http.ResponseWriter, r *http.Re
apiAgent, err := db2sdk.WorkspaceAgent( apiAgent, err := db2sdk.WorkspaceAgent(
api.DERPMap(), api.DERPMap(),
*api.TailnetCoordinator.Load(), *api.TailnetCoordinator.Load(),
workspaceAgent, waws.WorkspaceAgent,
nil, nil,
nil, nil,
nil, nil,
@@ -959,7 +916,7 @@ func (api *API) watchWorkspaceAgentContainers(rw http.ResponseWriter, r *http.Re
return return
} }
agentConn, release, err := api.agentProvider.AgentConn(dialCtx, workspaceAgent.ID) agentConn, release, err := api.agentProvider.AgentConn(dialCtx, waws.WorkspaceAgent.ID)
if err != nil { if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error dialing workspace agent.", Message: "Internal error dialing workspace agent.",
@@ -1034,7 +991,7 @@ func (api *API) watchWorkspaceAgentContainers(rw http.ResponseWriter, r *http.Re
// @Router /workspaceagents/{workspaceagent}/containers [get] // @Router /workspaceagents/{workspaceagent}/containers [get]
func (api *API) workspaceAgentListContainers(rw http.ResponseWriter, r *http.Request) { func (api *API) workspaceAgentListContainers(rw http.ResponseWriter, r *http.Request) {
ctx := r.Context() ctx := r.Context()
workspaceAgent := httpmw.WorkspaceAgentParam(r) waws := httpmw.WorkspaceAgentAndWorkspaceParam(r)
labelParam, ok := r.URL.Query()["label"] labelParam, ok := r.URL.Query()["label"]
if !ok { if !ok {
@@ -1060,7 +1017,7 @@ func (api *API) workspaceAgentListContainers(rw http.ResponseWriter, r *http.Req
apiAgent, err := db2sdk.WorkspaceAgent( apiAgent, err := db2sdk.WorkspaceAgent(
api.DERPMap(), api.DERPMap(),
*api.TailnetCoordinator.Load(), *api.TailnetCoordinator.Load(),
workspaceAgent, waws.WorkspaceAgent,
nil, nil,
nil, nil,
nil, nil,
@@ -1081,7 +1038,7 @@ func (api *API) workspaceAgentListContainers(rw http.ResponseWriter, r *http.Req
return return
} }
agentConn, release, err := api.agentProvider.AgentConn(ctx, workspaceAgent.ID) agentConn, release, err := api.agentProvider.AgentConn(ctx, waws.WorkspaceAgent.ID)
if err != nil { if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error dialing workspace agent.", Message: "Internal error dialing workspace agent.",
@@ -1131,10 +1088,9 @@ func (api *API) workspaceAgentListContainers(rw http.ResponseWriter, r *http.Req
// @Router /workspaceagents/{workspaceagent}/containers/devcontainers/{devcontainer} [delete] // @Router /workspaceagents/{workspaceagent}/containers/devcontainers/{devcontainer} [delete]
func (api *API) workspaceAgentDeleteDevcontainer(rw http.ResponseWriter, r *http.Request) { func (api *API) workspaceAgentDeleteDevcontainer(rw http.ResponseWriter, r *http.Request) {
ctx := r.Context() ctx := r.Context()
workspaceAgent := httpmw.WorkspaceAgentParam(r) waws := httpmw.WorkspaceAgentAndWorkspaceParam(r)
workspace := httpmw.WorkspaceParam(r)
if !api.Authorize(r, policy.ActionUpdate, workspace) { if !api.Authorize(r, policy.ActionUpdate, waws.WorkspaceTable) {
httpapi.Forbidden(rw) httpapi.Forbidden(rw)
return return
} }
@@ -1153,7 +1109,7 @@ func (api *API) workspaceAgentDeleteDevcontainer(rw http.ResponseWriter, r *http
apiAgent, err := db2sdk.WorkspaceAgent( apiAgent, err := db2sdk.WorkspaceAgent(
api.DERPMap(), api.DERPMap(),
*api.TailnetCoordinator.Load(), *api.TailnetCoordinator.Load(),
workspaceAgent, waws.WorkspaceAgent,
nil, nil,
nil, nil,
nil, nil,
@@ -1178,7 +1134,7 @@ func (api *API) workspaceAgentDeleteDevcontainer(rw http.ResponseWriter, r *http
// don't get a response after 30s that the agent is unreachable. // don't get a response after 30s that the agent is unreachable.
dialCtx, dialCancel := context.WithTimeout(ctx, 30*time.Second) dialCtx, dialCancel := context.WithTimeout(ctx, 30*time.Second)
defer dialCancel() defer dialCancel()
agentConn, release, err := api.agentProvider.AgentConn(dialCtx, workspaceAgent.ID) agentConn, release, err := api.agentProvider.AgentConn(dialCtx, waws.WorkspaceAgent.ID)
if err != nil { if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error dialing workspace agent.", Message: "Internal error dialing workspace agent.",
@@ -1222,7 +1178,7 @@ func (api *API) workspaceAgentDeleteDevcontainer(rw http.ResponseWriter, r *http
// @Router /workspaceagents/{workspaceagent}/containers/devcontainers/{devcontainer}/recreate [post] // @Router /workspaceagents/{workspaceagent}/containers/devcontainers/{devcontainer}/recreate [post]
func (api *API) workspaceAgentRecreateDevcontainer(rw http.ResponseWriter, r *http.Request) { func (api *API) workspaceAgentRecreateDevcontainer(rw http.ResponseWriter, r *http.Request) {
ctx := r.Context() ctx := r.Context()
workspaceAgent := httpmw.WorkspaceAgentParam(r) waws := httpmw.WorkspaceAgentAndWorkspaceParam(r)
devcontainer := chi.URLParam(r, "devcontainer") devcontainer := chi.URLParam(r, "devcontainer")
if devcontainer == "" { if devcontainer == "" {
@@ -1238,7 +1194,7 @@ func (api *API) workspaceAgentRecreateDevcontainer(rw http.ResponseWriter, r *ht
apiAgent, err := db2sdk.WorkspaceAgent( apiAgent, err := db2sdk.WorkspaceAgent(
api.DERPMap(), api.DERPMap(),
*api.TailnetCoordinator.Load(), *api.TailnetCoordinator.Load(),
workspaceAgent, waws.WorkspaceAgent,
nil, nil,
nil, nil,
nil, nil,
@@ -1263,7 +1219,7 @@ func (api *API) workspaceAgentRecreateDevcontainer(rw http.ResponseWriter, r *ht
// don't get a response after 30s that the agent is unreachable. // don't get a response after 30s that the agent is unreachable.
dialCtx, dialCancel := context.WithTimeout(ctx, 30*time.Second) dialCtx, dialCancel := context.WithTimeout(ctx, 30*time.Second)
defer dialCancel() defer dialCancel()
agentConn, release, err := api.agentProvider.AgentConn(dialCtx, workspaceAgent.ID) agentConn, release, err := api.agentProvider.AgentConn(dialCtx, waws.WorkspaceAgent.ID)
if err != nil { if err != nil {
httpapi.Write(dialCtx, rw, http.StatusInternalServerError, codersdk.Response{ httpapi.Write(dialCtx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error dialing workspace agent.", Message: "Internal error dialing workspace agent.",
@@ -1441,8 +1397,8 @@ func (api *API) workspaceAgentClientCoordinate(rw http.ResponseWriter, r *http.R
// This route accepts user API key auth and workspace proxy auth. The moon actor has // This route accepts user API key auth and workspace proxy auth. The moon actor has
// full permissions so should be able to pass this authz check. // full permissions so should be able to pass this authz check.
workspace := httpmw.WorkspaceParam(r) waws := httpmw.WorkspaceAgentAndWorkspaceParam(r)
if !api.Authorize(r, policy.ActionSSH, workspace) { if !api.Authorize(r, policy.ActionSSH, waws) {
httpapi.ResourceNotFound(rw) httpapi.ResourceNotFound(rw)
return return
} }
@@ -1482,7 +1438,6 @@ func (api *API) workspaceAgentClientCoordinate(rw http.ResponseWriter, r *http.R
api.WebsocketWaitGroup.Add(1) api.WebsocketWaitGroup.Add(1)
api.WebsocketWaitMutex.Unlock() api.WebsocketWaitMutex.Unlock()
defer api.WebsocketWaitGroup.Done() defer api.WebsocketWaitGroup.Done()
workspaceAgent := httpmw.WorkspaceAgentParam(r)
conn, err := websocket.Accept(rw, r, nil) conn, err := websocket.Accept(rw, r, nil)
if err != nil { if err != nil {
@@ -1502,7 +1457,7 @@ func (api *API) workspaceAgentClientCoordinate(rw http.ResponseWriter, r *http.R
Name: "client", Name: "client",
ID: peerID, ID: peerID,
Auth: tailnet.ClientCoordinateeAuth{ Auth: tailnet.ClientCoordinateeAuth{
AgentID: workspaceAgent.ID, AgentID: waws.WorkspaceAgent.ID,
}, },
}) })
if err != nil && !xerrors.Is(err, io.EOF) && !xerrors.Is(err, context.Canceled) { if err != nil && !xerrors.Is(err, io.EOF) && !xerrors.Is(err, context.Canceled) {
@@ -1647,7 +1602,7 @@ func (api *API) workspaceAgentReinit(rw http.ResponseWriter, r *http.Request) {
// convertProvisionedApps converts applications that are in the middle of provisioning process. // convertProvisionedApps converts applications that are in the middle of provisioning process.
// It means that they may not have an agent or workspace assigned (dry-run job). // It means that they may not have an agent or workspace assigned (dry-run job).
func convertProvisionedApps(dbApps []database.WorkspaceApp) []codersdk.WorkspaceApp { func convertProvisionedApps(dbApps []database.WorkspaceApp) []codersdk.WorkspaceApp {
return db2sdk.Apps(dbApps, []database.WorkspaceAppStatus{}, database.WorkspaceAgent{}, "", database.Workspace{}) return db2sdk.Apps(dbApps, []database.WorkspaceAppStatus{}, database.WorkspaceAgent{}, "", database.WorkspaceTable{})
} }
func convertLogSources(dbLogSources []database.WorkspaceAgentLogSource) []codersdk.WorkspaceAgentLogSource { func convertLogSources(dbLogSources []database.WorkspaceAgentLogSource) []codersdk.WorkspaceAgentLogSource {
@@ -1719,15 +1674,16 @@ func (api *API) watchWorkspaceAgentMetadata(
defer cancel() defer cancel()
r = r.WithContext(ctx) // Rewire context for SSE cancellation. r = r.WithContext(ctx) // Rewire context for SSE cancellation.
workspaceAgent := httpmw.WorkspaceAgentParam(r) waws := httpmw.WorkspaceAgentAndWorkspaceParam(r)
log := api.Logger.Named("workspace_metadata_watcher").With( log := api.Logger.Named("workspace_metadata_watcher").With(
slog.F("workspace_agent_id", workspaceAgent.ID), slog.F("workspace_agent_id", waws.WorkspaceAgent.ID),
slog.F("workspace_id", waws.WorkspaceTable.ID),
) )
// Send metadata on updates, we must ensure subscription before sending // Send metadata on updates, we must ensure subscription before sending
// initial metadata to guarantee that events in-between are not missed. // initial metadata to guarantee that events in-between are not missed.
update := make(chan agentapi.WorkspaceAgentMetadataChannelPayload, 1) update := make(chan agentapi.WorkspaceAgentMetadataChannelPayload, 1)
cancelSub, err := api.Pubsub.Subscribe(agentapi.WatchWorkspaceAgentMetadataChannel(workspaceAgent.ID), func(_ context.Context, byt []byte) { cancelSub, err := api.Pubsub.Subscribe(agentapi.WatchWorkspaceAgentMetadataChannel(waws.WorkspaceAgent.ID), func(_ context.Context, byt []byte) {
if ctx.Err() != nil { if ctx.Err() != nil {
return return
} }
@@ -1758,7 +1714,7 @@ func (api *API) watchWorkspaceAgentMetadata(
// We always use the original Request context because it contains // We always use the original Request context because it contains
// the RBAC actor. // the RBAC actor.
initialMD, err := api.Database.GetWorkspaceAgentMetadata(ctx, database.GetWorkspaceAgentMetadataParams{ initialMD, err := api.Database.GetWorkspaceAgentMetadata(ctx, database.GetWorkspaceAgentMetadataParams{
WorkspaceAgentID: workspaceAgent.ID, WorkspaceAgentID: waws.WorkspaceAgent.ID,
Keys: nil, Keys: nil,
}) })
if err != nil { if err != nil {
@@ -1839,7 +1795,7 @@ func (api *API) watchWorkspaceAgentMetadata(
return return
case payload := <-update: case payload := <-update:
md, err := api.Database.GetWorkspaceAgentMetadata(ctx, database.GetWorkspaceAgentMetadataParams{ md, err := api.Database.GetWorkspaceAgentMetadata(ctx, database.GetWorkspaceAgentMetadataParams{
WorkspaceAgentID: workspaceAgent.ID, WorkspaceAgentID: waws.WorkspaceAgent.ID,
Keys: payload.Keys, Keys: payload.Keys,
}) })
if err != nil { if err != nil {
+24 -42
View File
@@ -96,8 +96,6 @@ func TestWatchAgentContainers(t *testing.T) {
workspaceID = uuid.New() workspaceID = uuid.New()
agentID = uuid.New() agentID = uuid.New()
resourceID = uuid.New() resourceID = uuid.New()
jobID = uuid.New()
buildID = uuid.New()
containersCh = make(chan codersdk.WorkspaceAgentListContainersResponse) containersCh = make(chan codersdk.WorkspaceAgentListContainersResponse)
@@ -120,24 +118,17 @@ func TestWatchAgentContainers(t *testing.T) {
api.agentProvider = fAgentProvider api.agentProvider = fAgentProvider
// Setup: Allow `ExtractWorkspaceAgentParams` to complete. // Setup: Allow `ExtractWorkspaceAgentParams` to complete.
mDB.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agentID).Return(database.WorkspaceAgent{ mDB.EXPECT().GetWorkspaceAgentAndWorkspaceByID(gomock.Any(), agentID).Return(database.GetWorkspaceAgentAndWorkspaceByIDRow{
ID: agentID, WorkspaceAgent: database.WorkspaceAgent{
ResourceID: resourceID, ID: agentID,
LifecycleState: database.WorkspaceAgentLifecycleStateReady, ResourceID: resourceID,
FirstConnectedAt: sql.NullTime{Valid: true, Time: dbtime.Now()}, LifecycleState: database.WorkspaceAgentLifecycleStateReady,
LastConnectedAt: sql.NullTime{Valid: true, Time: dbtime.Now()}, FirstConnectedAt: sql.NullTime{Valid: true, Time: dbtime.Now()},
}, nil) LastConnectedAt: sql.NullTime{Valid: true, Time: dbtime.Now()},
mDB.EXPECT().GetWorkspaceResourceByID(gomock.Any(), resourceID).Return(database.WorkspaceResource{ },
ID: resourceID, WorkspaceTable: database.WorkspaceTable{
JobID: jobID, ID: workspaceID,
}, nil) },
mDB.EXPECT().GetProvisionerJobByID(gomock.Any(), jobID).Return(database.ProvisionerJob{
ID: jobID,
Type: database.ProvisionerJobTypeWorkspaceBuild,
}, nil)
mDB.EXPECT().GetWorkspaceBuildByJobID(gomock.Any(), jobID).Return(database.WorkspaceBuild{
WorkspaceID: workspaceID,
ID: buildID,
}, nil) }, nil)
// And: Allow `db2dsk.WorkspaceAgent` to complete. // And: Allow `db2dsk.WorkspaceAgent` to complete.
@@ -152,7 +143,7 @@ func TestWatchAgentContainers(t *testing.T) {
}) })
// And: We mount the HTTP Handler // And: We mount the HTTP Handler
r.With(httpmw.ExtractWorkspaceAgentParam(mDB)). r.With(httpmw.ExtractWorkspaceAgentAndWorkspaceParam(mDB)).
Get("/workspaceagents/{workspaceagent}/containers/watch", api.watchWorkspaceAgentContainers) Get("/workspaceagents/{workspaceagent}/containers/watch", api.watchWorkspaceAgentContainers)
// Given: We create the HTTP server // Given: We create the HTTP server
@@ -222,8 +213,6 @@ func TestWatchAgentContainers(t *testing.T) {
workspaceID = uuid.New() workspaceID = uuid.New()
agentID = uuid.New() agentID = uuid.New()
resourceID = uuid.New() resourceID = uuid.New()
jobID = uuid.New()
buildID = uuid.New()
containersCh = make(chan codersdk.WorkspaceAgentListContainersResponse) containersCh = make(chan codersdk.WorkspaceAgentListContainersResponse)
@@ -246,24 +235,17 @@ func TestWatchAgentContainers(t *testing.T) {
api.agentProvider = fAgentProvider api.agentProvider = fAgentProvider
// Setup: Allow `ExtractWorkspaceAgentParams` to complete. // Setup: Allow `ExtractWorkspaceAgentParams` to complete.
mDB.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agentID).Return(database.WorkspaceAgent{ mDB.EXPECT().GetWorkspaceAgentAndWorkspaceByID(gomock.Any(), agentID).Return(database.GetWorkspaceAgentAndWorkspaceByIDRow{
ID: agentID, WorkspaceAgent: database.WorkspaceAgent{
ResourceID: resourceID, ID: agentID,
LifecycleState: database.WorkspaceAgentLifecycleStateReady, ResourceID: resourceID,
FirstConnectedAt: sql.NullTime{Valid: true, Time: dbtime.Now()}, LifecycleState: database.WorkspaceAgentLifecycleStateReady,
LastConnectedAt: sql.NullTime{Valid: true, Time: dbtime.Now()}, FirstConnectedAt: sql.NullTime{Valid: true, Time: dbtime.Now()},
}, nil) LastConnectedAt: sql.NullTime{Valid: true, Time: dbtime.Now()},
mDB.EXPECT().GetWorkspaceResourceByID(gomock.Any(), resourceID).Return(database.WorkspaceResource{ },
ID: resourceID, WorkspaceTable: database.WorkspaceTable{
JobID: jobID, ID: workspaceID,
}, nil) },
mDB.EXPECT().GetProvisionerJobByID(gomock.Any(), jobID).Return(database.ProvisionerJob{
ID: jobID,
Type: database.ProvisionerJobTypeWorkspaceBuild,
}, nil)
mDB.EXPECT().GetWorkspaceBuildByJobID(gomock.Any(), jobID).Return(database.WorkspaceBuild{
WorkspaceID: workspaceID,
ID: buildID,
}, nil) }, nil)
// And: Allow `db2dsk.WorkspaceAgent` to complete. // And: Allow `db2dsk.WorkspaceAgent` to complete.
@@ -274,7 +256,7 @@ func TestWatchAgentContainers(t *testing.T) {
Return(containersCh, io.NopCloser(&bytes.Buffer{}), nil) Return(containersCh, io.NopCloser(&bytes.Buffer{}), nil)
// And: We mount the HTTP Handler // And: We mount the HTTP Handler
r.With(httpmw.ExtractWorkspaceAgentParam(mDB)). r.With(httpmw.ExtractWorkspaceAgentAndWorkspaceParam(mDB)).
Get("/workspaceagents/{workspaceagent}/containers/watch", api.watchWorkspaceAgentContainers) Get("/workspaceagents/{workspaceagent}/containers/watch", api.watchWorkspaceAgentContainers)
// Given: We create the HTTP server // Given: We create the HTTP server
+1 -1
View File
@@ -148,7 +148,7 @@ func (p *DBTokenProvider) Issue(ctx context.Context, rw http.ResponseWriter, r *
aReq.dbReq = dbReq // Update audit request. aReq.dbReq = dbReq // Update audit request.
token.UserID = dbReq.User.ID token.UserID = dbReq.UserID
token.WorkspaceID = dbReq.Workspace.ID token.WorkspaceID = dbReq.Workspace.ID
token.AgentID = dbReq.Agent.ID token.AgentID = dbReq.Agent.ID
if dbReq.AppURL != nil { if dbReq.AppURL != nil {
+10 -34
View File
@@ -188,10 +188,10 @@ func (r Request) Check() error {
type databaseRequest struct { type databaseRequest struct {
Request Request
// User is the user that owns the app. // UserID is the ID of the user that owns the app.
User database.User UserID uuid.UUID
// Workspace is the workspace that the app is in. // Workspace is the workspace that the app is in.
Workspace database.Workspace Workspace database.WorkspaceTable
// Agent is the agent that the app is running on. // Agent is the agent that the app is running on.
Agent database.WorkspaceAgent Agent database.WorkspaceAgent
// App is the app that the user is trying to access. // App is the app that the user is trying to access.
@@ -420,8 +420,8 @@ func (r Request) getDatabase(ctx context.Context, db database.Store) (*databaseR
return &databaseRequest{ return &databaseRequest{
Request: r, Request: r,
User: user, UserID: user.ID,
Workspace: workspace, Workspace: workspace.WorkspaceTable(),
Agent: agent, Agent: agent,
App: app, App: app,
AppURL: appURLParsed, AppURL: appURLParsed,
@@ -443,40 +443,16 @@ func (r Request) getDatabaseTerminal(ctx context.Context, db database.Store) (*d
} }
var err error var err error
agent, err := db.GetWorkspaceAgentByID(ctx, agentID) aw, err := db.GetWorkspaceAgentAndWorkspaceByID(ctx, agentID)
if err != nil { if err != nil {
return nil, xerrors.Errorf("get workspace agent %q: %w", agentID, err) return nil, xerrors.Errorf("get workspace agent %q with workspace: %w", agentID, err)
}
// Get the corresponding resource.
res, err := db.GetWorkspaceResourceByID(ctx, agent.ResourceID)
if err != nil {
return nil, xerrors.Errorf("get workspace agent resource %q: %w", agent.ResourceID, err)
}
// Get the corresponding workspace build.
build, err := db.GetWorkspaceBuildByJobID(ctx, res.JobID)
if err != nil {
return nil, xerrors.Errorf("get workspace build by job ID %q: %w", res.JobID, err)
}
// Get the corresponding workspace.
workspace, err := db.GetWorkspaceByID(ctx, build.WorkspaceID)
if err != nil {
return nil, xerrors.Errorf("get workspace %q: %w", build.WorkspaceID, err)
}
// Get the workspace's owner.
user, err := db.GetUserByID(ctx, workspace.OwnerID)
if err != nil {
return nil, xerrors.Errorf("get user %q: %w", workspace.OwnerID, err)
} }
return &databaseRequest{ return &databaseRequest{
Request: r, Request: r,
User: user, UserID: aw.WorkspaceTable.OwnerID,
Workspace: workspace, Workspace: aw.WorkspaceTable,
Agent: agent, Agent: aw.WorkspaceAgent,
AppURL: nil, AppURL: nil,
AppSharingLevel: database.AppSharingLevelOwner, AppSharingLevel: database.AppSharingLevelOwner,
}, nil }, nil
+1 -1
View File
@@ -1252,7 +1252,7 @@ func (api *API) convertWorkspaceBuild(
statuses := statusesByAgentID[agent.ID] statuses := statusesByAgentID[agent.ID]
logSources := logSourcesByAgentID[agent.ID] logSources := logSourcesByAgentID[agent.ID]
apiAgent, err := db2sdk.WorkspaceAgent( apiAgent, err := db2sdk.WorkspaceAgent(
api.DERPMap(), *api.TailnetCoordinator.Load(), agent, db2sdk.Apps(apps, statuses, agent, workspace.OwnerUsername, workspace), convertScripts(scripts), convertLogSources(logSources), api.AgentInactiveDisconnectTimeout, api.DERPMap(), *api.TailnetCoordinator.Load(), agent, db2sdk.Apps(apps, statuses, agent, workspace.OwnerUsername, workspace.WorkspaceTable()), convertScripts(scripts), convertLogSources(logSources), api.AgentInactiveDisconnectTimeout,
api.DeploymentValues.AgentFallbackTroubleshootingURL.String(), api.DeploymentValues.AgentFallbackTroubleshootingURL.String(),
) )
if err != nil { if err != nil {