Skip to content

Commit d29fcb3

Browse files
mafredrimafredri
committed
fix dbmem query logic
1 parent dc93229 commit d29fcb3

File tree

3 files changed

+168
-154
lines changed

3 files changed

+168
-154
lines changed

‎coderd/database/dbmem/dbmem.go

+116-106
Original file line numberDiff line numberDiff line change
@@ -4119,7 +4119,7 @@ func (q *FakeQuerier) GetUnexpiredLicenses(_ context.Context) ([]database.Licens
41194119
return results, nil
41204120
}
41214121

4122-
func (q *FakeQuerier) GetUserActivityInsights(ctx context.Context, arg database.GetUserActivityInsightsParams) ([]database.GetUserActivityInsightsRow, error) {
4122+
func (q *FakeQuerier) GetUserActivityInsights(_ context.Context, arg database.GetUserActivityInsightsParams) ([]database.GetUserActivityInsightsRow, error) {
41234123
err := validateDatabaseType(arg)
41244124
if err != nil {
41254125
return nil, err
@@ -4128,130 +4128,140 @@ func (q *FakeQuerier) GetUserActivityInsights(ctx context.Context, arg database.
41284128
q.mutex.RLock()
41294129
defer q.mutex.RUnlock()
41304130

4131-
type uniqueKey struct {
4132-
TemplateID uuid.UUID
4133-
UserID uuid.UUID
4134-
}
4135-
4136-
combinedStats := make(map[uniqueKey]map[time.Time]int64)
4137-
4138-
// Get application stats
4139-
for _, s := range q.workspaceAppStats {
4140-
if !(((s.SessionStartedAt.After(arg.StartTime) || s.SessionStartedAt.Equal(arg.StartTime)) && s.SessionStartedAt.Before(arg.EndTime)) ||
4141-
(s.SessionEndedAt.After(arg.StartTime) && s.SessionEndedAt.Before(arg.EndTime)) ||
4142-
(s.SessionStartedAt.Before(arg.StartTime) && (s.SessionEndedAt.After(arg.EndTime) || s.SessionEndedAt.Equal(arg.EndTime)))) {
4143-
continue
4144-
}
4145-
4146-
w, err := q.getWorkspaceByIDNoLock(ctx, s.WorkspaceID)
4147-
if err != nil {
4148-
return nil, err
4149-
}
4150-
4151-
if len(arg.TemplateIDs) > 0 && !slices.Contains(arg.TemplateIDs, w.TemplateID) {
4152-
continue
4153-
}
4154-
4155-
key := uniqueKey{
4156-
TemplateID: w.TemplateID,
4157-
UserID: s.UserID,
4158-
}
4159-
if combinedStats[key] == nil {
4160-
combinedStats[key] = make(map[time.Time]int64)
4161-
}
4131+
/*
4132+
WITH
4133+
*/
4134+
/*
4135+
deployment_stats AS (
4136+
SELECT
4137+
start_time,
4138+
user_id,
4139+
array_agg(template_id) AS template_ids,
4140+
-- See motivation in GetTemplateInsights for LEAST(SUM(n), 30).
4141+
LEAST(SUM(usage_mins), 30) AS usage_mins
4142+
FROM
4143+
template_usage_stats
4144+
WHERE
4145+
start_time >= @start_time::timestamptz
4146+
AND end_time <= @end_time::timestamptz
4147+
AND CASE WHEN COALESCE(array_length(@template_ids::uuid[], 1), 0) > 0 THEN template_id = ANY(@template_ids::uuid[]) ELSE TRUE END
4148+
GROUP BY
4149+
start_time, user_id
4150+
),
4151+
*/
41624152

4163-
t := s.SessionStartedAt.Truncate(time.Minute)
4164-
if t.Before(arg.StartTime) {
4165-
t = arg.StartTime
4166-
}
4167-
for t.Before(s.SessionEndedAt) && t.Before(arg.EndTime) {
4168-
combinedStats[key][t] = 60
4169-
t = t.Add(1 * time.Minute)
4170-
}
4153+
type deploymentStatsGroupBy struct {
4154+
StartTime time.Time
4155+
UserID uuid.UUID
41714156
}
4172-
4173-
// Get session stats
4174-
for _, s := range q.workspaceAgentStats {
4175-
if s.CreatedAt.Before(arg.StartTime) || s.CreatedAt.Equal(arg.EndTime) || s.CreatedAt.After(arg.EndTime) {
4176-
continue
4177-
}
4178-
if len(arg.TemplateIDs) > 0 && !slices.Contains(arg.TemplateIDs, s.TemplateID) {
4157+
type deploymentStatsRow struct {
4158+
deploymentStatsGroupBy
4159+
TemplateIDs []uuid.UUID
4160+
UsageMins int16
4161+
}
4162+
deploymentStatsRows := make(map[deploymentStatsGroupBy]deploymentStatsRow)
4163+
for _, stat := range q.templateUsageStats {
4164+
if stat.StartTime.Before(arg.StartTime) || stat.EndTime.After(arg.EndTime) {
41794165
continue
41804166
}
4181-
if s.ConnectionCount == 0 {
4167+
if len(arg.TemplateIDs) > 0 && !slices.Contains(arg.TemplateIDs, stat.TemplateID) {
41824168
continue
41834169
}
4184-
4185-
key := uniqueKey{
4186-
TemplateID: s.TemplateID,
4187-
UserID: s.UserID,
4188-
}
4189-
4190-
if combinedStats[key] == nil {
4191-
combinedStats[key] = make(map[time.Time]int64)
4170+
key := deploymentStatsGroupBy{
4171+
StartTime: stat.StartTime,
4172+
UserID: stat.UserID,
41924173
}
4193-
4194-
if s.SessionCountJetBrains > 0 || s.SessionCountVSCode > 0 || s.SessionCountReconnectingPTY > 0 || s.SessionCountSSH > 0 {
4195-
t := s.CreatedAt.Truncate(time.Minute)
4196-
combinedStats[key][t] = 60
4174+
row, ok := deploymentStatsRows[key]
4175+
if !ok {
4176+
row = deploymentStatsRow{
4177+
deploymentStatsGroupBy: key,
4178+
}
41974179
}
4180+
row.TemplateIDs = append(row.TemplateIDs, stat.TemplateID)
4181+
row.UsageMins = least(row.UsageMins+stat.UsageMins, 30)
4182+
deploymentStatsRows[key] = row
41984183
}
41994184

4200-
// Use temporary maps for aggregation purposes
4201-
mUserIDTemplateIDs := map[uuid.UUID]map[uuid.UUID]struct{}{}
4202-
mUserIDUsageSeconds := map[uuid.UUID]int64{}
4203-
4204-
for key, times := range combinedStats {
4205-
if mUserIDTemplateIDs[key.UserID] == nil {
4206-
mUserIDTemplateIDs[key.UserID] = make(map[uuid.UUID]struct{})
4207-
mUserIDUsageSeconds[key.UserID] = 0
4208-
}
4209-
4210-
if _, ok := mUserIDTemplateIDs[key.UserID][key.TemplateID]; !ok {
4211-
mUserIDTemplateIDs[key.UserID][key.TemplateID] = struct{}{}
4212-
}
4185+
/*
4186+
template_ids AS (
4187+
SELECT
4188+
user_id,
4189+
array_agg(DISTINCT template_id) AS ids
4190+
FROM
4191+
deployment_stats, unnest(template_ids) template_id
4192+
GROUP BY
4193+
user_id
4194+
)
4195+
*/
42134196

4214-
for _, t := range times {
4215-
mUserIDUsageSeconds[key.UserID] += t
4197+
type templateIDsRow struct {
4198+
UserID uuid.UUID
4199+
TemplateIDs []uuid.UUID
4200+
}
4201+
templateIDs := make(map[uuid.UUID]templateIDsRow)
4202+
for _, dsRow := range deploymentStatsRows {
4203+
row, ok := templateIDs[dsRow.UserID]
4204+
if !ok {
4205+
row = templateIDsRow{
4206+
UserID: row.UserID,
4207+
}
42164208
}
4209+
row.TemplateIDs = uniqueSortedUUIDs(append(row.TemplateIDs, dsRow.TemplateIDs...))
4210+
templateIDs[dsRow.UserID] = row
42174211
}
42184212

4219-
userIDs := make([]uuid.UUID, 0, len(mUserIDUsageSeconds))
4220-
for userID := range mUserIDUsageSeconds {
4221-
userIDs = append(userIDs, userID)
4222-
}
4223-
sort.Slice(userIDs, func(i, j int) bool {
4224-
return userIDs[i].String() < userIDs[j].String()
4225-
})
4213+
/*
4214+
SELECT
4215+
ds.user_id,
4216+
u.username,
4217+
u.avatar_url,
4218+
t.ids::uuid[] AS template_ids,
4219+
(SUM(ds.usage_mins) * 60)::bigint AS usage_seconds
4220+
FROM
4221+
deployment_stats ds
4222+
JOIN
4223+
users u
4224+
ON
4225+
u.id = ds.user_id
4226+
JOIN
4227+
template_ids t
4228+
ON
4229+
ds.user_id = t.user_id
4230+
GROUP BY
4231+
ds.user_id, u.username, u.avatar_url, t.ids
4232+
ORDER BY
4233+
ds.user_id ASC;
4234+
*/
42264235

4227-
// Finally, select stats
42284236
var rows []database.GetUserActivityInsightsRow
4229-
4230-
for _, userID := range userIDs {
4231-
user, err := q.getUserByIDNoLock(userID)
4232-
if err != nil {
4233-
return nil, err
4234-
}
4235-
4236-
tids := mUserIDTemplateIDs[userID]
4237-
templateIDs := make([]uuid.UUID, 0, len(tids))
4238-
for key := range tids {
4239-
templateIDs = append(templateIDs, key)
4240-
}
4241-
sort.Slice(templateIDs, func(i, j int) bool {
4242-
return templateIDs[i].String() < templateIDs[j].String()
4243-
})
4244-
4245-
row := database.GetUserActivityInsightsRow{
4246-
UserID: user.ID,
4247-
Username: user.Username,
4248-
AvatarURL: user.AvatarURL,
4249-
TemplateIDs: templateIDs,
4250-
UsageSeconds: mUserIDUsageSeconds[userID],
4237+
groupedRows := make(map[uuid.UUID]database.GetUserActivityInsightsRow)
4238+
for _, dsRow := range deploymentStatsRows {
4239+
row, ok := groupedRows[dsRow.UserID]
4240+
if !ok {
4241+
user, err := q.getUserByIDNoLock(dsRow.UserID)
4242+
if err != nil {
4243+
return nil, err
4244+
}
4245+
row = database.GetUserActivityInsightsRow{
4246+
UserID: user.ID,
4247+
Username: user.Username,
4248+
AvatarURL: user.AvatarURL,
4249+
TemplateIDs: templateIDs[user.ID].TemplateIDs,
4250+
}
42514251
}
4252-
4252+
row.UsageSeconds += int64(dsRow.UsageMins) * 60
4253+
groupedRows[dsRow.UserID] = row
4254+
}
4255+
for _, row := range groupedRows {
42534256
rows = append(rows, row)
42544257
}
4258+
if len(rows) == 0 {
4259+
return nil, sql.ErrNoRows
4260+
}
4261+
slices.SortFunc(rows, func(a, b database.GetUserActivityInsightsRow) int {
4262+
return slice.Ascending(a.UserID.String(), b.UserID.String())
4263+
})
4264+
42554265
return rows, nil
42564266
}
42574267

‎coderd/database/queries.sql.go

+26-24
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries/insights.sql

+26-24
Original file line numberDiff line numberDiff line change
@@ -33,30 +33,32 @@ ORDER BY
3333
-- requesting deployment-wide (or multiple template) data. Cumulative
3434
-- produces a bloated value if a user has used multiple templates
3535
-- simultaneously.
36-
WITH deployment_stats AS (
37-
SELECT
38-
start_time,
39-
user_id,
40-
array_agg(template_id) AS template_ids,
41-
-- See motivation in GetTemplateInsights for LEAST(SUM(n), 30).
42-
LEAST(SUM(usage_mins), 30) AS usage_mins
43-
FROM
44-
template_usage_stats
45-
WHERE
46-
start_time >= @start_time::timestamptz
47-
AND end_time <= @end_time::timestamptz
48-
AND CASE WHEN COALESCE(array_length(@template_ids::uuid[], 1), 0) > 0 THEN template_id = ANY(@template_ids::uuid[]) ELSE TRUE END
49-
GROUP BY
50-
start_time, user_id
51-
), template_ids AS (
52-
SELECT
53-
user_id,
54-
array_agg(DISTINCT template_id) AS ids
55-
FROM
56-
deployment_stats, unnest(template_ids) template_id
57-
GROUP BY
58-
user_id
59-
)
36+
WITH
37+
deployment_stats AS (
38+
SELECT
39+
start_time,
40+
user_id,
41+
array_agg(template_id) AS template_ids,
42+
-- See motivation in GetTemplateInsights for LEAST(SUM(n), 30).
43+
LEAST(SUM(usage_mins), 30) AS usage_mins
44+
FROM
45+
template_usage_stats
46+
WHERE
47+
start_time >= @start_time::timestamptz
48+
AND end_time <= @end_time::timestamptz
49+
AND CASE WHEN COALESCE(array_length(@template_ids::uuid[], 1), 0) > 0 THEN template_id = ANY(@template_ids::uuid[]) ELSE TRUE END
50+
GROUP BY
51+
start_time, user_id
52+
),
53+
template_ids AS (
54+
SELECT
55+
user_id,
56+
array_agg(DISTINCT template_id) AS ids
57+
FROM
58+
deployment_stats, unnest(template_ids) template_id
59+
GROUP BY
60+
user_id
61+
)
6062

6163
SELECT
6264
ds.user_id,

0 commit comments

Comments
 (0)