mirror of
https://github.com/coder/coder.git
synced 2025-07-15 22:20:27 +00:00
fix: remove unnecessary user lookup in agent API calls (#17934)
# Use workspace.OwnerUsername instead of fetching the owner This PR optimizes the agent API by using the `workspace.OwnerUsername` field directly instead of making an additional database query to fetch the owner's username. The change removes the need to call `GetUserByID` in the manifest API and workspace agent RPC endpoints. An issue arose when the agent token was scoped without access to user data (`api_key_scope = "no_user_data"`), causing the agent to fail to fetch the manifest due to an RBAC issue. Change-Id: I3b6e7581134e2374b364ee059e3b18ece3d98b41 Signed-off-by: Thomas Kosiewski <tk@coder.com>
This commit is contained in:
@ -47,7 +47,6 @@ func (a *ManifestAPI) GetManifest(ctx context.Context, _ *agentproto.GetManifest
|
|||||||
scripts []database.WorkspaceAgentScript
|
scripts []database.WorkspaceAgentScript
|
||||||
metadata []database.WorkspaceAgentMetadatum
|
metadata []database.WorkspaceAgentMetadatum
|
||||||
workspace database.Workspace
|
workspace database.Workspace
|
||||||
owner database.User
|
|
||||||
devcontainers []database.WorkspaceAgentDevcontainer
|
devcontainers []database.WorkspaceAgentDevcontainer
|
||||||
)
|
)
|
||||||
|
|
||||||
@ -76,10 +75,6 @@ func (a *ManifestAPI) GetManifest(ctx context.Context, _ *agentproto.GetManifest
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
return xerrors.Errorf("getting workspace by id: %w", err)
|
return xerrors.Errorf("getting workspace by id: %w", err)
|
||||||
}
|
}
|
||||||
owner, err = a.Database.GetUserByID(ctx, workspace.OwnerID)
|
|
||||||
if err != nil {
|
|
||||||
return xerrors.Errorf("getting workspace owner by id: %w", err)
|
|
||||||
}
|
|
||||||
return err
|
return err
|
||||||
})
|
})
|
||||||
eg.Go(func() (err error) {
|
eg.Go(func() (err error) {
|
||||||
@ -98,7 +93,7 @@ func (a *ManifestAPI) GetManifest(ctx context.Context, _ *agentproto.GetManifest
|
|||||||
AppSlugOrPort: "{{port}}",
|
AppSlugOrPort: "{{port}}",
|
||||||
AgentName: workspaceAgent.Name,
|
AgentName: workspaceAgent.Name,
|
||||||
WorkspaceName: workspace.Name,
|
WorkspaceName: workspace.Name,
|
||||||
Username: owner.Username,
|
Username: workspace.OwnerUsername,
|
||||||
}
|
}
|
||||||
|
|
||||||
vscodeProxyURI := vscodeProxyURI(appSlug, a.AccessURL, a.AppHostname)
|
vscodeProxyURI := vscodeProxyURI(appSlug, a.AccessURL, a.AppHostname)
|
||||||
@ -115,7 +110,7 @@ func (a *ManifestAPI) GetManifest(ctx context.Context, _ *agentproto.GetManifest
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
apps, err := dbAppsToProto(dbApps, workspaceAgent, owner.Username, workspace)
|
apps, err := dbAppsToProto(dbApps, workspaceAgent, workspace.OwnerUsername, workspace)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, xerrors.Errorf("converting workspace apps: %w", err)
|
return nil, xerrors.Errorf("converting workspace apps: %w", err)
|
||||||
}
|
}
|
||||||
@ -128,7 +123,7 @@ func (a *ManifestAPI) GetManifest(ctx context.Context, _ *agentproto.GetManifest
|
|||||||
return &agentproto.Manifest{
|
return &agentproto.Manifest{
|
||||||
AgentId: workspaceAgent.ID[:],
|
AgentId: workspaceAgent.ID[:],
|
||||||
AgentName: workspaceAgent.Name,
|
AgentName: workspaceAgent.Name,
|
||||||
OwnerUsername: owner.Username,
|
OwnerUsername: workspace.OwnerUsername,
|
||||||
WorkspaceId: workspace.ID[:],
|
WorkspaceId: workspace.ID[:],
|
||||||
WorkspaceName: workspace.Name,
|
WorkspaceName: workspace.Name,
|
||||||
GitAuthConfigs: gitAuthConfigs,
|
GitAuthConfigs: gitAuthConfigs,
|
||||||
|
@ -48,6 +48,7 @@ func TestGetManifest(t *testing.T) {
|
|||||||
workspace = database.Workspace{
|
workspace = database.Workspace{
|
||||||
ID: uuid.New(),
|
ID: uuid.New(),
|
||||||
OwnerID: owner.ID,
|
OwnerID: owner.ID,
|
||||||
|
OwnerUsername: owner.Username,
|
||||||
Name: "cool-workspace",
|
Name: "cool-workspace",
|
||||||
}
|
}
|
||||||
agent = database.WorkspaceAgent{
|
agent = database.WorkspaceAgent{
|
||||||
@ -336,7 +337,6 @@ func TestGetManifest(t *testing.T) {
|
|||||||
}).Return(metadata, nil)
|
}).Return(metadata, nil)
|
||||||
mDB.EXPECT().GetWorkspaceAgentDevcontainersByAgentID(gomock.Any(), agent.ID).Return(devcontainers, nil)
|
mDB.EXPECT().GetWorkspaceAgentDevcontainersByAgentID(gomock.Any(), agent.ID).Return(devcontainers, nil)
|
||||||
mDB.EXPECT().GetWorkspaceByID(gomock.Any(), workspace.ID).Return(workspace, nil)
|
mDB.EXPECT().GetWorkspaceByID(gomock.Any(), workspace.ID).Return(workspace, nil)
|
||||||
mDB.EXPECT().GetUserByID(gomock.Any(), workspace.OwnerID).Return(owner, nil)
|
|
||||||
|
|
||||||
got, err := api.GetManifest(context.Background(), &agentproto.GetManifestRequest{})
|
got, err := api.GetManifest(context.Background(), &agentproto.GetManifestRequest{})
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
@ -404,7 +404,6 @@ func TestGetManifest(t *testing.T) {
|
|||||||
}).Return([]database.WorkspaceAgentMetadatum{}, nil)
|
}).Return([]database.WorkspaceAgentMetadatum{}, nil)
|
||||||
mDB.EXPECT().GetWorkspaceAgentDevcontainersByAgentID(gomock.Any(), childAgent.ID).Return([]database.WorkspaceAgentDevcontainer{}, nil)
|
mDB.EXPECT().GetWorkspaceAgentDevcontainersByAgentID(gomock.Any(), childAgent.ID).Return([]database.WorkspaceAgentDevcontainer{}, nil)
|
||||||
mDB.EXPECT().GetWorkspaceByID(gomock.Any(), workspace.ID).Return(workspace, nil)
|
mDB.EXPECT().GetWorkspaceByID(gomock.Any(), workspace.ID).Return(workspace, nil)
|
||||||
mDB.EXPECT().GetUserByID(gomock.Any(), workspace.OwnerID).Return(owner, nil)
|
|
||||||
|
|
||||||
got, err := api.GetManifest(context.Background(), &agentproto.GetManifestRequest{})
|
got, err := api.GetManifest(context.Background(), &agentproto.GetManifestRequest{})
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
@ -468,7 +467,6 @@ func TestGetManifest(t *testing.T) {
|
|||||||
}).Return(metadata, nil)
|
}).Return(metadata, nil)
|
||||||
mDB.EXPECT().GetWorkspaceAgentDevcontainersByAgentID(gomock.Any(), agent.ID).Return(devcontainers, nil)
|
mDB.EXPECT().GetWorkspaceAgentDevcontainersByAgentID(gomock.Any(), agent.ID).Return(devcontainers, nil)
|
||||||
mDB.EXPECT().GetWorkspaceByID(gomock.Any(), workspace.ID).Return(workspace, nil)
|
mDB.EXPECT().GetWorkspaceByID(gomock.Any(), workspace.ID).Return(workspace, nil)
|
||||||
mDB.EXPECT().GetUserByID(gomock.Any(), workspace.OwnerID).Return(owner, nil)
|
|
||||||
|
|
||||||
got, err := api.GetManifest(context.Background(), &agentproto.GetManifestRequest{})
|
got, err := api.GetManifest(context.Background(), &agentproto.GetManifestRequest{})
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
|
@ -439,14 +439,42 @@ func TestWorkspaceAgentConnectRPC(t *testing.T) {
|
|||||||
t.Run("Connect", func(t *testing.T) {
|
t.Run("Connect", func(t *testing.T) {
|
||||||
t.Parallel()
|
t.Parallel()
|
||||||
|
|
||||||
|
for _, tc := range []struct {
|
||||||
|
name string
|
||||||
|
apiKeyScope rbac.ScopeName
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "empty (backwards compat)",
|
||||||
|
apiKeyScope: "",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "all",
|
||||||
|
apiKeyScope: rbac.ScopeAll,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "no_user_data",
|
||||||
|
apiKeyScope: rbac.ScopeNoUserData,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "application_connect",
|
||||||
|
apiKeyScope: rbac.ScopeApplicationConnect,
|
||||||
|
},
|
||||||
|
} {
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
client, db := coderdtest.NewWithDatabase(t, nil)
|
client, db := coderdtest.NewWithDatabase(t, nil)
|
||||||
user := coderdtest.CreateFirstUser(t, client)
|
user := coderdtest.CreateFirstUser(t, client)
|
||||||
r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
|
r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
|
||||||
OrganizationID: user.OrganizationID,
|
OrganizationID: user.OrganizationID,
|
||||||
OwnerID: user.UserID,
|
OwnerID: user.UserID,
|
||||||
}).WithAgent().Do()
|
}).WithAgent(func(agents []*proto.Agent) []*proto.Agent {
|
||||||
|
for _, agent := range agents {
|
||||||
|
agent.ApiKeyScope = string(tc.apiKeyScope)
|
||||||
|
}
|
||||||
|
|
||||||
|
return agents
|
||||||
|
}).Do()
|
||||||
_ = agenttest.New(t, client.URL, r.AgentToken)
|
_ = agenttest.New(t, client.URL, r.AgentToken)
|
||||||
resources := coderdtest.AwaitWorkspaceAgents(t, client, r.Workspace.ID)
|
resources := coderdtest.NewWorkspaceAgentWaiter(t, client, r.Workspace.ID).AgentNames([]string{}).Wait()
|
||||||
|
|
||||||
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
|
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
|
||||||
defer cancel()
|
defer cancel()
|
||||||
@ -459,6 +487,8 @@ func TestWorkspaceAgentConnectRPC(t *testing.T) {
|
|||||||
}()
|
}()
|
||||||
conn.AwaitReachable(ctx)
|
conn.AwaitReachable(ctx)
|
||||||
})
|
})
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
t.Run("FailNonLatestBuild", func(t *testing.T) {
|
t.Run("FailNonLatestBuild", func(t *testing.T) {
|
||||||
t.Parallel()
|
t.Parallel()
|
||||||
|
@ -76,17 +76,8 @@ func (api *API) workspaceAgentRPC(rw http.ResponseWriter, r *http.Request) {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
owner, err := api.Database.GetUserByID(ctx, workspace.OwnerID)
|
|
||||||
if err != nil {
|
|
||||||
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
|
|
||||||
Message: "Internal error fetching user.",
|
|
||||||
Detail: err.Error(),
|
|
||||||
})
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
logger = logger.With(
|
logger = logger.With(
|
||||||
slog.F("owner", owner.Username),
|
slog.F("owner", workspace.OwnerUsername),
|
||||||
slog.F("workspace_name", workspace.Name),
|
slog.F("workspace_name", workspace.Name),
|
||||||
slog.F("agent_name", workspaceAgent.Name),
|
slog.F("agent_name", workspaceAgent.Name),
|
||||||
)
|
)
|
||||||
@ -170,7 +161,7 @@ func (api *API) workspaceAgentRPC(rw http.ResponseWriter, r *http.Request) {
|
|||||||
})
|
})
|
||||||
|
|
||||||
streamID := tailnet.StreamID{
|
streamID := tailnet.StreamID{
|
||||||
Name: fmt.Sprintf("%s-%s-%s", owner.Username, workspace.Name, workspaceAgent.Name),
|
Name: fmt.Sprintf("%s-%s-%s", workspace.OwnerUsername, workspace.Name, workspaceAgent.Name),
|
||||||
ID: workspaceAgent.ID,
|
ID: workspaceAgent.ID,
|
||||||
Auth: tailnet.AgentCoordinateeAuth{ID: workspaceAgent.ID},
|
Auth: tailnet.AgentCoordinateeAuth{ID: workspaceAgent.ID},
|
||||||
}
|
}
|
||||||
|
@ -13,6 +13,7 @@ import (
|
|||||||
"github.com/coder/coder/v2/coderd/database"
|
"github.com/coder/coder/v2/coderd/database"
|
||||||
"github.com/coder/coder/v2/coderd/database/dbfake"
|
"github.com/coder/coder/v2/coderd/database/dbfake"
|
||||||
"github.com/coder/coder/v2/coderd/database/dbtime"
|
"github.com/coder/coder/v2/coderd/database/dbtime"
|
||||||
|
"github.com/coder/coder/v2/coderd/rbac"
|
||||||
"github.com/coder/coder/v2/codersdk/agentsdk"
|
"github.com/coder/coder/v2/codersdk/agentsdk"
|
||||||
"github.com/coder/coder/v2/provisionersdk/proto"
|
"github.com/coder/coder/v2/provisionersdk/proto"
|
||||||
"github.com/coder/coder/v2/testutil"
|
"github.com/coder/coder/v2/testutil"
|
||||||
@ -22,6 +23,30 @@ import (
|
|||||||
func TestWorkspaceAgentReportStats(t *testing.T) {
|
func TestWorkspaceAgentReportStats(t *testing.T) {
|
||||||
t.Parallel()
|
t.Parallel()
|
||||||
|
|
||||||
|
for _, tc := range []struct {
|
||||||
|
name string
|
||||||
|
apiKeyScope rbac.ScopeName
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "empty (backwards compat)",
|
||||||
|
apiKeyScope: "",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "all",
|
||||||
|
apiKeyScope: rbac.ScopeAll,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "no_user_data",
|
||||||
|
apiKeyScope: rbac.ScopeNoUserData,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "application_connect",
|
||||||
|
apiKeyScope: rbac.ScopeApplicationConnect,
|
||||||
|
},
|
||||||
|
} {
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
tickCh := make(chan time.Time)
|
tickCh := make(chan time.Time)
|
||||||
flushCh := make(chan int, 1)
|
flushCh := make(chan int, 1)
|
||||||
client, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{
|
client, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{
|
||||||
@ -33,7 +58,15 @@ func TestWorkspaceAgentReportStats(t *testing.T) {
|
|||||||
OrganizationID: user.OrganizationID,
|
OrganizationID: user.OrganizationID,
|
||||||
OwnerID: user.UserID,
|
OwnerID: user.UserID,
|
||||||
LastUsedAt: dbtime.Now().Add(-time.Minute),
|
LastUsedAt: dbtime.Now().Add(-time.Minute),
|
||||||
}).WithAgent().Do()
|
}).WithAgent(
|
||||||
|
func(agent []*proto.Agent) []*proto.Agent {
|
||||||
|
for _, a := range agent {
|
||||||
|
a.ApiKeyScope = string(tc.apiKeyScope)
|
||||||
|
}
|
||||||
|
|
||||||
|
return agent
|
||||||
|
},
|
||||||
|
).Do()
|
||||||
|
|
||||||
ac := agentsdk.New(client.URL)
|
ac := agentsdk.New(client.URL)
|
||||||
ac.SetSessionToken(r.AgentToken)
|
ac.SetSessionToken(r.AgentToken)
|
||||||
@ -72,10 +105,36 @@ func TestWorkspaceAgentReportStats(t *testing.T) {
|
|||||||
newWorkspace.LastUsedAt.After(r.Workspace.LastUsedAt),
|
newWorkspace.LastUsedAt.After(r.Workspace.LastUsedAt),
|
||||||
"%s is not after %s", newWorkspace.LastUsedAt, r.Workspace.LastUsedAt,
|
"%s is not after %s", newWorkspace.LastUsedAt, r.Workspace.LastUsedAt,
|
||||||
)
|
)
|
||||||
|
})
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestAgentAPI_LargeManifest(t *testing.T) {
|
func TestAgentAPI_LargeManifest(t *testing.T) {
|
||||||
t.Parallel()
|
t.Parallel()
|
||||||
|
|
||||||
|
for _, tc := range []struct {
|
||||||
|
name string
|
||||||
|
apiKeyScope rbac.ScopeName
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "empty (backwards compat)",
|
||||||
|
apiKeyScope: "",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "all",
|
||||||
|
apiKeyScope: rbac.ScopeAll,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "no_user_data",
|
||||||
|
apiKeyScope: rbac.ScopeNoUserData,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "application_connect",
|
||||||
|
apiKeyScope: rbac.ScopeApplicationConnect,
|
||||||
|
},
|
||||||
|
} {
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
ctx := testutil.Context(t, testutil.WaitLong)
|
ctx := testutil.Context(t, testutil.WaitLong)
|
||||||
client, store := coderdtest.NewWithDatabase(t, nil)
|
client, store := coderdtest.NewWithDatabase(t, nil)
|
||||||
adminUser := coderdtest.CreateFirstUser(t, client)
|
adminUser := coderdtest.CreateFirstUser(t, client)
|
||||||
@ -93,6 +152,7 @@ func TestAgentAPI_LargeManifest(t *testing.T) {
|
|||||||
Script: string(longScript),
|
Script: string(longScript),
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
|
agents[0].ApiKeyScope = string(tc.apiKeyScope)
|
||||||
return agents
|
return agents
|
||||||
}).Do()
|
}).Do()
|
||||||
ac := agentsdk.New(client.URL)
|
ac := agentsdk.New(client.URL)
|
||||||
@ -107,4 +167,6 @@ func TestAgentAPI_LargeManifest(t *testing.T) {
|
|||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
require.Len(t, manifest.Scripts, 1)
|
require.Len(t, manifest.Scripts, 1)
|
||||||
require.Len(t, manifest.Scripts[0].Script, n)
|
require.Len(t, manifest.Scripts[0].Script, n)
|
||||||
|
})
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
Reference in New Issue
Block a user