mirror of
https://github.com/coder/coder.git
synced 2025-07-08 11:39:50 +00:00
chore: refactor workspaces query to use window function (#5079)
* Use window function in query * Convert workspace rows and unpack count * Update types * Fix Scan bug * Remove getCountError
This commit is contained in:
@ -99,13 +99,14 @@ func (e *Executor) runOnce(t time.Time) Stats {
|
||||
// NOTE: If a workspace build is created with a given TTL and then the user either
|
||||
// changes or unsets the TTL, the deadline for the workspace build will not
|
||||
// have changed. This behavior is as expected per #2229.
|
||||
workspaces, err := e.db.GetWorkspaces(e.ctx, database.GetWorkspacesParams{
|
||||
workspaceRows, err := e.db.GetWorkspaces(e.ctx, database.GetWorkspacesParams{
|
||||
Deleted: false,
|
||||
})
|
||||
if err != nil {
|
||||
e.log.Error(e.ctx, "get workspaces for autostart or autostop", slog.Error(err))
|
||||
return stats
|
||||
}
|
||||
workspaces := database.ConvertWorkspaceRows(workspaceRows)
|
||||
|
||||
var eligibleWorkspaceIDs []uuid.UUID
|
||||
for _, ws := range workspaces {
|
||||
|
@ -718,14 +718,14 @@ func (q *fakeQuerier) GetAuthorizationUserRoles(_ context.Context, userID uuid.U
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (q *fakeQuerier) GetWorkspaces(ctx context.Context, arg database.GetWorkspacesParams) ([]database.Workspace, error) {
|
||||
func (q *fakeQuerier) GetWorkspaces(ctx context.Context, arg database.GetWorkspacesParams) ([]database.GetWorkspacesRow, error) {
|
||||
// A nil auth filter means no auth filter.
|
||||
workspaces, err := q.GetAuthorizedWorkspaces(ctx, arg, nil)
|
||||
return workspaces, err
|
||||
workspaceRows, err := q.GetAuthorizedWorkspaces(ctx, arg, nil)
|
||||
return workspaceRows, err
|
||||
}
|
||||
|
||||
//nolint:gocyclo
|
||||
func (q *fakeQuerier) GetAuthorizedWorkspaces(ctx context.Context, arg database.GetWorkspacesParams, authorizedFilter rbac.AuthorizeFilter) ([]database.Workspace, error) {
|
||||
func (q *fakeQuerier) GetAuthorizedWorkspaces(ctx context.Context, arg database.GetWorkspacesParams, authorizedFilter rbac.AuthorizeFilter) ([]database.GetWorkspacesRow, error) {
|
||||
q.mutex.RLock()
|
||||
defer q.mutex.RUnlock()
|
||||
|
||||
@ -866,20 +866,43 @@ func (q *fakeQuerier) GetAuthorizedWorkspaces(ctx context.Context, arg database.
|
||||
workspaces = append(workspaces, workspace)
|
||||
}
|
||||
|
||||
beforePageCount := len(workspaces)
|
||||
|
||||
if arg.Offset > 0 {
|
||||
if int(arg.Offset) > len(workspaces) {
|
||||
return []database.Workspace{}, nil
|
||||
return []database.GetWorkspacesRow{}, nil
|
||||
}
|
||||
workspaces = workspaces[arg.Offset:]
|
||||
}
|
||||
if arg.Limit > 0 {
|
||||
if int(arg.Limit) > len(workspaces) {
|
||||
return workspaces, nil
|
||||
return convertToWorkspaceRows(workspaces, int64(beforePageCount)), nil
|
||||
}
|
||||
workspaces = workspaces[:arg.Limit]
|
||||
}
|
||||
|
||||
return workspaces, nil
|
||||
return convertToWorkspaceRows(workspaces, int64(beforePageCount)), nil
|
||||
}
|
||||
|
||||
func convertToWorkspaceRows(workspaces []database.Workspace, count int64) []database.GetWorkspacesRow {
|
||||
rows := make([]database.GetWorkspacesRow, len(workspaces))
|
||||
for i, w := range workspaces {
|
||||
rows[i] = database.GetWorkspacesRow{
|
||||
ID: w.ID,
|
||||
CreatedAt: w.CreatedAt,
|
||||
UpdatedAt: w.UpdatedAt,
|
||||
OwnerID: w.OwnerID,
|
||||
OrganizationID: w.OrganizationID,
|
||||
TemplateID: w.TemplateID,
|
||||
Deleted: w.Deleted,
|
||||
Name: w.Name,
|
||||
AutostartSchedule: w.AutostartSchedule,
|
||||
Ttl: w.Ttl,
|
||||
LastUsedAt: w.LastUsedAt,
|
||||
Count: count,
|
||||
}
|
||||
}
|
||||
return rows
|
||||
}
|
||||
|
||||
func (q *fakeQuerier) GetWorkspaceByID(_ context.Context, id uuid.UUID) (database.Workspace, error) {
|
||||
|
@ -93,3 +93,24 @@ func ConvertUserRows(rows []GetUsersRow) []User {
|
||||
|
||||
return users
|
||||
}
|
||||
|
||||
func ConvertWorkspaceRows(rows []GetWorkspacesRow) []Workspace {
|
||||
workspaces := make([]Workspace, len(rows))
|
||||
for i, r := range rows {
|
||||
workspaces[i] = Workspace{
|
||||
ID: r.ID,
|
||||
CreatedAt: r.CreatedAt,
|
||||
UpdatedAt: r.UpdatedAt,
|
||||
OwnerID: r.OwnerID,
|
||||
OrganizationID: r.OrganizationID,
|
||||
TemplateID: r.TemplateID,
|
||||
Deleted: r.Deleted,
|
||||
Name: r.Name,
|
||||
AutostartSchedule: r.AutostartSchedule,
|
||||
Ttl: r.Ttl,
|
||||
LastUsedAt: r.LastUsedAt,
|
||||
}
|
||||
}
|
||||
|
||||
return workspaces
|
||||
}
|
||||
|
@ -112,13 +112,13 @@ func (q *sqlQuerier) GetTemplateGroupRoles(ctx context.Context, id uuid.UUID) ([
|
||||
}
|
||||
|
||||
type workspaceQuerier interface {
|
||||
GetAuthorizedWorkspaces(ctx context.Context, arg GetWorkspacesParams, authorizedFilter rbac.AuthorizeFilter) ([]Workspace, error)
|
||||
GetAuthorizedWorkspaces(ctx context.Context, arg GetWorkspacesParams, authorizedFilter rbac.AuthorizeFilter) ([]GetWorkspacesRow, error)
|
||||
}
|
||||
|
||||
// GetAuthorizedWorkspaces returns all workspaces that the user is authorized to access.
|
||||
// This code is copied from `GetWorkspaces` and adds the authorized filter WHERE
|
||||
// clause.
|
||||
func (q *sqlQuerier) GetAuthorizedWorkspaces(ctx context.Context, arg GetWorkspacesParams, authorizedFilter rbac.AuthorizeFilter) ([]Workspace, error) {
|
||||
func (q *sqlQuerier) GetAuthorizedWorkspaces(ctx context.Context, arg GetWorkspacesParams, authorizedFilter rbac.AuthorizeFilter) ([]GetWorkspacesRow, error) {
|
||||
// In order to properly use ORDER BY, OFFSET, and LIMIT, we need to inject the
|
||||
// authorizedFilter between the end of the where clause and those statements.
|
||||
filter := strings.Replace(getWorkspaces, "-- @authorize_filter", fmt.Sprintf(" AND %s", authorizedFilter.SQLString(rbac.NoACLConfig())), 1)
|
||||
@ -139,9 +139,9 @@ func (q *sqlQuerier) GetAuthorizedWorkspaces(ctx context.Context, arg GetWorkspa
|
||||
return nil, xerrors.Errorf("get authorized workspaces: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
var items []Workspace
|
||||
var items []GetWorkspacesRow
|
||||
for rows.Next() {
|
||||
var i Workspace
|
||||
var i GetWorkspacesRow
|
||||
if err := rows.Scan(
|
||||
&i.ID,
|
||||
&i.CreatedAt,
|
||||
@ -154,6 +154,7 @@ func (q *sqlQuerier) GetAuthorizedWorkspaces(ctx context.Context, arg GetWorkspa
|
||||
&i.AutostartSchedule,
|
||||
&i.Ttl,
|
||||
&i.LastUsedAt,
|
||||
&i.Count,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
@ -124,7 +124,7 @@ type sqlcQuerier interface {
|
||||
GetWorkspaceResourcesByJobID(ctx context.Context, jobID uuid.UUID) ([]WorkspaceResource, error)
|
||||
GetWorkspaceResourcesByJobIDs(ctx context.Context, ids []uuid.UUID) ([]WorkspaceResource, error)
|
||||
GetWorkspaceResourcesCreatedAfter(ctx context.Context, createdAt time.Time) ([]WorkspaceResource, error)
|
||||
GetWorkspaces(ctx context.Context, arg GetWorkspacesParams) ([]Workspace, error)
|
||||
GetWorkspaces(ctx context.Context, arg GetWorkspacesParams) ([]GetWorkspacesRow, error)
|
||||
InsertAPIKey(ctx context.Context, arg InsertAPIKeyParams) (APIKey, error)
|
||||
InsertAgentStat(ctx context.Context, arg InsertAgentStatParams) (AgentStat, error)
|
||||
// We use the organization_id as the id
|
||||
|
@ -6228,7 +6228,7 @@ func (q *sqlQuerier) GetWorkspaceOwnerCountsByTemplateIDs(ctx context.Context, i
|
||||
|
||||
const getWorkspaces = `-- name: GetWorkspaces :many
|
||||
SELECT
|
||||
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.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, COUNT(*) OVER () as count
|
||||
FROM
|
||||
workspaces
|
||||
LEFT JOIN LATERAL (
|
||||
@ -6375,7 +6375,22 @@ type GetWorkspacesParams struct {
|
||||
Limit int32 `db:"limit_" json:"limit_"`
|
||||
}
|
||||
|
||||
func (q *sqlQuerier) GetWorkspaces(ctx context.Context, arg GetWorkspacesParams) ([]Workspace, error) {
|
||||
type GetWorkspacesRow struct {
|
||||
ID uuid.UUID `db:"id" json:"id"`
|
||||
CreatedAt time.Time `db:"created_at" json:"created_at"`
|
||||
UpdatedAt time.Time `db:"updated_at" json:"updated_at"`
|
||||
OwnerID uuid.UUID `db:"owner_id" json:"owner_id"`
|
||||
OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"`
|
||||
TemplateID uuid.UUID `db:"template_id" json:"template_id"`
|
||||
Deleted bool `db:"deleted" json:"deleted"`
|
||||
Name string `db:"name" json:"name"`
|
||||
AutostartSchedule sql.NullString `db:"autostart_schedule" json:"autostart_schedule"`
|
||||
Ttl sql.NullInt64 `db:"ttl" json:"ttl"`
|
||||
LastUsedAt time.Time `db:"last_used_at" json:"last_used_at"`
|
||||
Count int64 `db:"count" json:"count"`
|
||||
}
|
||||
|
||||
func (q *sqlQuerier) GetWorkspaces(ctx context.Context, arg GetWorkspacesParams) ([]GetWorkspacesRow, error) {
|
||||
rows, err := q.db.QueryContext(ctx, getWorkspaces,
|
||||
arg.Deleted,
|
||||
arg.Status,
|
||||
@ -6391,9 +6406,9 @@ func (q *sqlQuerier) GetWorkspaces(ctx context.Context, arg GetWorkspacesParams)
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var items []Workspace
|
||||
var items []GetWorkspacesRow
|
||||
for rows.Next() {
|
||||
var i Workspace
|
||||
var i GetWorkspacesRow
|
||||
if err := rows.Scan(
|
||||
&i.ID,
|
||||
&i.CreatedAt,
|
||||
@ -6406,6 +6421,7 @@ func (q *sqlQuerier) GetWorkspaces(ctx context.Context, arg GetWorkspacesParams)
|
||||
&i.AutostartSchedule,
|
||||
&i.Ttl,
|
||||
&i.LastUsedAt,
|
||||
&i.Count,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
@ -10,7 +10,7 @@ LIMIT
|
||||
|
||||
-- name: GetWorkspaces :many
|
||||
SELECT
|
||||
workspaces.*
|
||||
workspaces.*, COUNT(*) OVER () as count
|
||||
FROM
|
||||
workspaces
|
||||
LEFT JOIN LATERAL (
|
||||
|
@ -380,10 +380,11 @@ func (r *remoteReporter) createSnapshot() (*Snapshot, error) {
|
||||
return nil
|
||||
})
|
||||
eg.Go(func() error {
|
||||
workspaces, err := r.options.Database.GetWorkspaces(ctx, database.GetWorkspacesParams{})
|
||||
workspaceRows, err := r.options.Database.GetWorkspaces(ctx, database.GetWorkspacesParams{})
|
||||
if err != nil {
|
||||
return xerrors.Errorf("get workspaces: %w", err)
|
||||
}
|
||||
workspaces := database.ConvertWorkspaceRows(workspaceRows)
|
||||
snapshot.Workspaces = make([]Workspace, 0, len(workspaces))
|
||||
for _, dbWorkspace := range workspaces {
|
||||
snapshot.Workspaces = append(snapshot.Workspaces, ConvertWorkspace(dbWorkspace))
|
||||
|
@ -127,7 +127,7 @@ func (api *API) workspaces(rw http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
|
||||
workspaces, err := api.Database.GetAuthorizedWorkspaces(ctx, filter, sqlFilter)
|
||||
workspaceRows, err := api.Database.GetAuthorizedWorkspaces(ctx, filter, sqlFilter)
|
||||
if err != nil {
|
||||
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
|
||||
Message: "Internal error fetching workspaces.",
|
||||
@ -135,19 +135,15 @@ func (api *API) workspaces(rw http.ResponseWriter, r *http.Request) {
|
||||
})
|
||||
return
|
||||
}
|
||||
if len(workspaceRows) == 0 {
|
||||
httpapi.Write(ctx, rw, http.StatusOK, codersdk.WorkspacesResponse{
|
||||
Workspaces: []codersdk.Workspace{},
|
||||
Count: 0,
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// run the query again to get the total count for frontend pagination
|
||||
filter.Offset = 0
|
||||
filter.Limit = 0
|
||||
all, err := api.Database.GetAuthorizedWorkspaces(ctx, filter, sqlFilter)
|
||||
if err != nil {
|
||||
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
|
||||
Message: "Internal error fetching workspaces.",
|
||||
Detail: err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
count := len(all)
|
||||
workspaces := database.ConvertWorkspaceRows(workspaceRows)
|
||||
|
||||
data, err := api.workspaceData(ctx, workspaces)
|
||||
if err != nil {
|
||||
@ -169,7 +165,7 @@ func (api *API) workspaces(rw http.ResponseWriter, r *http.Request) {
|
||||
|
||||
httpapi.Write(ctx, rw, http.StatusOK, codersdk.WorkspacesResponse{
|
||||
Workspaces: wss,
|
||||
Count: count,
|
||||
Count: int(workspaceRows[0].Count),
|
||||
})
|
||||
}
|
||||
|
||||
|
Reference in New Issue
Block a user