Skip to content

Commit e6ead7d

Browse files
presleyppresleyp
authored
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
1 parent 560d3c9 commit e6ead7d

File tree

12 files changed

+91
-41
lines changed

12 files changed

+91
-41
lines changed

‎coderd/autobuild/executor/lifecycle_executor.go

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -99,13 +99,14 @@ func (e *Executor) runOnce(t time.Time) Stats {
9999
// NOTE: If a workspace build is created with a given TTL and then the user either
100100
// changes or unsets the TTL, the deadline for the workspace build will not
101101
// have changed. This behavior is as expected per #2229.
102-
workspaces, err := e.db.GetWorkspaces(e.ctx, database.GetWorkspacesParams{
102+
workspaceRows, err := e.db.GetWorkspaces(e.ctx, database.GetWorkspacesParams{
103103
Deleted: false,
104104
})
105105
if err != nil {
106106
e.log.Error(e.ctx, "get workspaces for autostart or autostop", slog.Error(err))
107107
return stats
108108
}
109+
workspaces := database.ConvertWorkspaceRows(workspaceRows)
109110

110111
var eligibleWorkspaceIDs []uuid.UUID
111112
for _, ws := range workspaces {

‎coderd/database/databasefake/databasefake.go

Lines changed: 30 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -718,14 +718,14 @@ func (q *fakeQuerier) GetAuthorizationUserRoles(_ context.Context, userID uuid.U
718718
}, nil
719719
}
720720

721-
func (q *fakeQuerier) GetWorkspaces(ctx context.Context, arg database.GetWorkspacesParams) ([]database.Workspace, error) {
721+
func (q *fakeQuerier) GetWorkspaces(ctx context.Context, arg database.GetWorkspacesParams) ([]database.GetWorkspacesRow, error) {
722722
// A nil auth filter means no auth filter.
723-
workspaces, err := q.GetAuthorizedWorkspaces(ctx, arg, nil)
724-
return workspaces, err
723+
workspaceRows, err := q.GetAuthorizedWorkspaces(ctx, arg, nil)
724+
return workspaceRows, err
725725
}
726726

727727
//nolint:gocyclo
728-
func (q *fakeQuerier) GetAuthorizedWorkspaces(ctx context.Context, arg database.GetWorkspacesParams, authorizedFilter rbac.AuthorizeFilter) ([]database.Workspace, error) {
728+
func (q *fakeQuerier) GetAuthorizedWorkspaces(ctx context.Context, arg database.GetWorkspacesParams, authorizedFilter rbac.AuthorizeFilter) ([]database.GetWorkspacesRow, error) {
729729
q.mutex.RLock()
730730
defer q.mutex.RUnlock()
731731

@@ -866,20 +866,43 @@ func (q *fakeQuerier) GetAuthorizedWorkspaces(ctx context.Context, arg database.
866866
workspaces = append(workspaces, workspace)
867867
}
868868

869+
beforePageCount := len(workspaces)
870+
869871
if arg.Offset > 0 {
870872
if int(arg.Offset) > len(workspaces) {
871-
return []database.Workspace{}, nil
873+
return []database.GetWorkspacesRow{}, nil
872874
}
873875
workspaces = workspaces[arg.Offset:]
874876
}
875877
if arg.Limit > 0 {
876878
if int(arg.Limit) > len(workspaces) {
877-
return workspaces, nil
879+
return convertToWorkspaceRows(workspaces, int64(beforePageCount)), nil
878880
}
879881
workspaces = workspaces[:arg.Limit]
880882
}
881883

882-
return workspaces, nil
884+
return convertToWorkspaceRows(workspaces, int64(beforePageCount)), nil
885+
}
886+
887+
func convertToWorkspaceRows(workspaces []database.Workspace, count int64) []database.GetWorkspacesRow {
888+
rows := make([]database.GetWorkspacesRow, len(workspaces))
889+
for i, w := range workspaces {
890+
rows[i] = database.GetWorkspacesRow{
891+
ID: w.ID,
892+
CreatedAt: w.CreatedAt,
893+
UpdatedAt: w.UpdatedAt,
894+
OwnerID: w.OwnerID,
895+
OrganizationID: w.OrganizationID,
896+
TemplateID: w.TemplateID,
897+
Deleted: w.Deleted,
898+
Name: w.Name,
899+
AutostartSchedule: w.AutostartSchedule,
900+
Ttl: w.Ttl,
901+
LastUsedAt: w.LastUsedAt,
902+
Count: count,
903+
}
904+
}
905+
return rows
883906
}
884907

885908
func (q *fakeQuerier) GetWorkspaceByID(_ context.Context, id uuid.UUID) (database.Workspace, error) {

‎coderd/database/modelmethods.go

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -93,3 +93,24 @@ func ConvertUserRows(rows []GetUsersRow) []User {
9393

9494
return users
9595
}
96+
97+
func ConvertWorkspaceRows(rows []GetWorkspacesRow) []Workspace {
98+
workspaces := make([]Workspace, len(rows))
99+
for i, r := range rows {
100+
workspaces[i] = Workspace{
101+
ID: r.ID,
102+
CreatedAt: r.CreatedAt,
103+
UpdatedAt: r.UpdatedAt,
104+
OwnerID: r.OwnerID,
105+
OrganizationID: r.OrganizationID,
106+
TemplateID: r.TemplateID,
107+
Deleted: r.Deleted,
108+
Name: r.Name,
109+
AutostartSchedule: r.AutostartSchedule,
110+
Ttl: r.Ttl,
111+
LastUsedAt: r.LastUsedAt,
112+
}
113+
}
114+
115+
return workspaces
116+
}

‎coderd/database/modelqueries.go

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -112,13 +112,13 @@ func (q *sqlQuerier) GetTemplateGroupRoles(ctx context.Context, id uuid.UUID) ([
112112
}
113113

114114
type workspaceQuerier interface {
115-
GetAuthorizedWorkspaces(ctx context.Context, arg GetWorkspacesParams, authorizedFilter rbac.AuthorizeFilter) ([]Workspace, error)
115+
GetAuthorizedWorkspaces(ctx context.Context, arg GetWorkspacesParams, authorizedFilter rbac.AuthorizeFilter) ([]GetWorkspacesRow, error)
116116
}
117117

118118
// GetAuthorizedWorkspaces returns all workspaces that the user is authorized to access.
119119
// This code is copied from `GetWorkspaces` and adds the authorized filter WHERE
120120
// clause.
121-
func (q *sqlQuerier) GetAuthorizedWorkspaces(ctx context.Context, arg GetWorkspacesParams, authorizedFilter rbac.AuthorizeFilter) ([]Workspace, error) {
121+
func (q *sqlQuerier) GetAuthorizedWorkspaces(ctx context.Context, arg GetWorkspacesParams, authorizedFilter rbac.AuthorizeFilter) ([]GetWorkspacesRow, error) {
122122
// In order to properly use ORDER BY, OFFSET, and LIMIT, we need to inject the
123123
// authorizedFilter between the end of the where clause and those statements.
124124
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
139139
return nil, xerrors.Errorf("get authorized workspaces: %w", err)
140140
}
141141
defer rows.Close()
142-
var items []Workspace
142+
var items []GetWorkspacesRow
143143
for rows.Next() {
144-
var i Workspace
144+
var i GetWorkspacesRow
145145
if err := rows.Scan(
146146
&i.ID,
147147
&i.CreatedAt,
@@ -154,6 +154,7 @@ func (q *sqlQuerier) GetAuthorizedWorkspaces(ctx context.Context, arg GetWorkspa
154154
&i.AutostartSchedule,
155155
&i.Ttl,
156156
&i.LastUsedAt,
157+
&i.Count,
157158
); err != nil {
158159
return nil, err
159160
}

‎coderd/database/querier.go

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries.sql.go

Lines changed: 20 additions & 4 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries/workspaces.sql

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ LIMIT
1010

1111
-- name: GetWorkspaces :many
1212
SELECT
13-
workspaces.*
13+
workspaces.*, COUNT(*) OVER () as count
1414
FROM
1515
workspaces
1616
LEFT JOIN LATERAL (

‎coderd/telemetry/telemetry.go

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -380,10 +380,11 @@ func (r *remoteReporter) createSnapshot() (*Snapshot, error) {
380380
return nil
381381
})
382382
eg.Go(func() error {
383-
workspaces, err := r.options.Database.GetWorkspaces(ctx, database.GetWorkspacesParams{})
383+
workspaceRows, err := r.options.Database.GetWorkspaces(ctx, database.GetWorkspacesParams{})
384384
if err != nil {
385385
return xerrors.Errorf("get workspaces: %w", err)
386386
}
387+
workspaces := database.ConvertWorkspaceRows(workspaceRows)
387388
snapshot.Workspaces = make([]Workspace, 0, len(workspaces))
388389
for _, dbWorkspace := range workspaces {
389390
snapshot.Workspaces = append(snapshot.Workspaces, ConvertWorkspace(dbWorkspace))

‎coderd/workspaces.go

Lines changed: 8 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -127,27 +127,23 @@ func (api *API) workspaces(rw http.ResponseWriter, r *http.Request) {
127127
return
128128
}
129129

130-
workspaces, err := api.Database.GetAuthorizedWorkspaces(ctx, filter, sqlFilter)
130+
workspaceRows, err := api.Database.GetAuthorizedWorkspaces(ctx, filter, sqlFilter)
131131
if err != nil {
132132
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
133133
Message: "Internal error fetching workspaces.",
134134
Detail: err.Error(),
135135
})
136136
return
137137
}
138-
139-
// run the query again to get the total count for frontend pagination
140-
filter.Offset = 0
141-
filter.Limit = 0
142-
all, err := api.Database.GetAuthorizedWorkspaces(ctx, filter, sqlFilter)
143-
if err != nil {
144-
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
145-
Message: "Internal error fetching workspaces.",
146-
Detail: err.Error(),
138+
if len(workspaceRows) == 0 {
139+
httpapi.Write(ctx, rw, http.StatusOK, codersdk.WorkspacesResponse{
140+
Workspaces: []codersdk.Workspace{},
141+
Count: 0,
147142
})
148143
return
149144
}
150-
count := len(all)
145+
146+
workspaces := database.ConvertWorkspaceRows(workspaceRows)
151147

152148
data, err := api.workspaceData(ctx, workspaces)
153149
if err != nil {
@@ -169,7 +165,7 @@ func (api *API) workspaces(rw http.ResponseWriter, r *http.Request) {
169165

170166
httpapi.Write(ctx, rw, http.StatusOK, codersdk.WorkspacesResponse{
171167
Workspaces: wss,
172-
Count: count,
168+
Count: int(workspaceRows[0].Count),
173169
})
174170
}
175171

‎site/src/pages/WorkspacesPage/WorkspacesPage.tsx

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -27,8 +27,7 @@ const WorkspacesPage: FC = () => {
2727
},
2828
})
2929

30-
const { workspaceRefs, count, getWorkspacesError, getCountError } =
31-
workspacesState.context
30+
const { workspaceRefs, count, getWorkspacesError } = workspacesState.context
3231
const paginationRef = workspacesState.context
3332
.paginationRef as PaginationMachineRef
3433

@@ -44,7 +43,6 @@ const WorkspacesPage: FC = () => {
4443
workspaceRefs={workspaceRefs}
4544
count={count}
4645
getWorkspacesError={getWorkspacesError}
47-
getCountError={getCountError}
4846
onFilter={(query) => {
4947
send({
5048
type: "UPDATE_FILTER",

‎site/src/pages/WorkspacesPage/WorkspacesPageView.tsx

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,6 @@ export interface WorkspacesPageViewProps {
3232
workspaceRefs?: WorkspaceItemMachineRef[]
3333
count?: number
3434
getWorkspacesError: Error | unknown
35-
getCountError: Error | unknown
3635
filter?: string
3736
onFilter: (query: string) => void
3837
paginationRef: PaginationMachineRef
@@ -46,7 +45,6 @@ export const WorkspacesPageView: FC<
4645
workspaceRefs,
4746
count,
4847
getWorkspacesError,
49-
getCountError,
5048
filter,
5149
onFilter,
5250
paginationRef,
@@ -92,10 +90,6 @@ export const WorkspacesPageView: FC<
9290
/>
9391
</Maybe>
9492

95-
<Maybe condition={getCountError !== undefined}>
96-
<AlertBanner error={getCountError} severity="warning" />
97-
</Maybe>
98-
9993
<SearchBarWithFilter
10094
filter={filter}
10195
onFilter={onFilter}

‎site/src/xServices/workspaces/workspacesXService.ts

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -216,7 +216,6 @@ interface WorkspacesContext {
216216
filter: string
217217
count?: number
218218
getWorkspacesError?: Error | unknown
219-
getCountError?: Error | unknown
220219
paginationContext: PaginationContext
221220
}
222221

0 commit comments

Comments
 (0)