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

coderd/autobuild/executor/lifecycle_executor.go

@@ -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 {

coderd/database/databasefake/databasefake.go

@@ -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) {

coderd/database/modelmethods.go

@@ -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
+}

coderd/database/modelqueries.go

@@ -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
 		}

coderd/database/queries/workspaces.sql

@@ -10,7 +10,7 @@ LIMIT
 
 -- name: GetWorkspaces :many
 SELECT
-	workspaces.*
+	workspaces.*, COUNT(*) OVER () as count
 FROM
 	workspaces
 LEFT JOIN LATERAL (

coderd/telemetry/telemetry.go

@@ -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))

coderd/workspaces.go

@@ -127,27 +127,23 @@ 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.",
 			Detail:  err.Error(),
 		})
 		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(),
+	if len(workspaceRows) == 0 {
+		httpapi.Write(ctx, rw, http.StatusOK, codersdk.WorkspacesResponse{
+			Workspaces: []codersdk.Workspace{},
+			Count:      0,
 		})
 		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),
 	})
 }
 

site/src/pages/WorkspacesPage/WorkspacesPage.tsx

@@ -27,8 +27,7 @@ const WorkspacesPage: FC = () => {
     },
   })
 
-  const { workspaceRefs, count, getWorkspacesError, getCountError } =
-    workspacesState.context
+  const { workspaceRefs, count, getWorkspacesError } = workspacesState.context
   const paginationRef = workspacesState.context
     .paginationRef as PaginationMachineRef
 
@@ -44,7 +43,6 @@ const WorkspacesPage: FC = () => {
         workspaceRefs={workspaceRefs}
         count={count}
         getWorkspacesError={getWorkspacesError}
-        getCountError={getCountError}
         onFilter={(query) => {
           send({
             type: "UPDATE_FILTER",

site/src/pages/WorkspacesPage/WorkspacesPageView.tsx

@@ -32,7 +32,6 @@ export interface WorkspacesPageViewProps {
   workspaceRefs?: WorkspaceItemMachineRef[]
   count?: number
   getWorkspacesError: Error | unknown
-  getCountError: Error | unknown
   filter?: string
   onFilter: (query: string) => void
   paginationRef: PaginationMachineRef
@@ -46,7 +45,6 @@ export const WorkspacesPageView: FC<
   workspaceRefs,
   count,
   getWorkspacesError,
-  getCountError,
   filter,
   onFilter,
   paginationRef,
@@ -92,10 +90,6 @@ export const WorkspacesPageView: FC<
           />
         </Maybe>
 
-        <Maybe condition={getCountError !== undefined}>
-          <AlertBanner error={getCountError} severity="warning" />
-        </Maybe>
-
         <SearchBarWithFilter
           filter={filter}
           onFilter={onFilter}

site/src/xServices/workspaces/workspacesXService.ts

@@ -216,7 +216,6 @@ interface WorkspacesContext {
   filter: string
   count?: number
   getWorkspacesError?: Error | unknown
-  getCountError?: Error | unknown
   paginationContext: PaginationContext
 }