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:
Presley Pizzo
2022-11-16 10:16:37 -05:00
committed by GitHub
parent 560d3c9fd0
commit e6ead7d915
12 changed files with 93 additions and 43 deletions

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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