Skip to content

Commit ec47926

Browse files
committed
fix: implement GetTemplateInsightsByInterval in dbmem
1 parent 0b2b56c commit ec47926

File tree

1 file changed

+71
-64
lines changed

1 file changed

+71
-64
lines changed

coderd/database/dbmem/dbmem.go

+71-64
Original file line numberDiff line numberDiff line change
@@ -3403,7 +3403,7 @@ func (q *FakeQuerier) GetTemplateInsights(_ context.Context, arg database.GetTem
34033403
return row, nil
34043404
}
34053405

3406-
func (q *FakeQuerier) GetTemplateInsightsByInterval(ctx context.Context, arg database.GetTemplateInsightsByIntervalParams) ([]database.GetTemplateInsightsByIntervalRow, error) {
3406+
func (q *FakeQuerier) GetTemplateInsightsByInterval(_ context.Context, arg database.GetTemplateInsightsByIntervalParams) ([]database.GetTemplateInsightsByIntervalRow, error) {
34073407
err := validateDatabaseType(arg)
34083408
if err != nil {
34093409
return nil, err
@@ -3412,82 +3412,89 @@ func (q *FakeQuerier) GetTemplateInsightsByInterval(ctx context.Context, arg dat
34123412
q.mutex.RLock()
34133413
defer q.mutex.RUnlock()
34143414

3415-
type statByInterval struct {
3416-
startTime, endTime time.Time
3417-
userSet map[uuid.UUID]struct{}
3418-
templateIDSet map[uuid.UUID]struct{}
3419-
}
3415+
/*
3416+
WITH
3417+
ts AS (
3418+
SELECT
3419+
d::timestamptz AS from_,
3420+
CASE
3421+
WHEN (d::timestamptz + (@interval_days::int || ' day')::interval) <= @end_time::timestamptz
3422+
THEN (d::timestamptz + (@interval_days::int || ' day')::interval)
3423+
ELSE @end_time::timestamptz
3424+
END AS to_
3425+
FROM
3426+
-- Subtract 1 microsecond from end_time to avoid including the next interval in the results.
3427+
generate_series(@start_time::timestamptz, (@end_time::timestamptz) - '1 microsecond'::interval, (@interval_days::int || ' day')::interval) AS d
3428+
)
34203429
3421-
statsByInterval := []statByInterval{{arg.StartTime, arg.StartTime.AddDate(0, 0, int(arg.IntervalDays)), make(map[uuid.UUID]struct{}), make(map[uuid.UUID]struct{})}}
3422-
for statsByInterval[len(statsByInterval)-1].endTime.Before(arg.EndTime) {
3423-
statsByInterval = append(statsByInterval, statByInterval{statsByInterval[len(statsByInterval)-1].endTime, statsByInterval[len(statsByInterval)-1].endTime.AddDate(0, 0, int(arg.IntervalDays)), make(map[uuid.UUID]struct{}), make(map[uuid.UUID]struct{})})
3424-
}
3425-
if statsByInterval[len(statsByInterval)-1].endTime.After(arg.EndTime) {
3426-
statsByInterval[len(statsByInterval)-1].endTime = arg.EndTime
3427-
}
3430+
SELECT
3431+
ts.from_ AS start_time,
3432+
ts.to_ AS end_time,
3433+
array_remove(array_agg(DISTINCT tus.template_id), NULL)::uuid[] AS template_ids,
3434+
COUNT(DISTINCT tus.user_id) AS active_users
3435+
FROM
3436+
ts
3437+
LEFT JOIN
3438+
template_usage_stats AS tus
3439+
ON
3440+
tus.start_time >= ts.from_
3441+
AND tus.end_time <= ts.to_
3442+
AND CASE WHEN COALESCE(array_length(@template_ids::uuid[], 1), 0) > 0 THEN tus.template_id = ANY(@template_ids::uuid[]) ELSE TRUE END
3443+
GROUP BY
3444+
ts.from_, ts.to_;
3445+
*/
34283446

3429-
for _, s := range q.workspaceAgentStats {
3430-
if s.CreatedAt.Before(arg.StartTime) || s.CreatedAt.Equal(arg.EndTime) || s.CreatedAt.After(arg.EndTime) {
3431-
continue
3432-
}
3433-
if len(arg.TemplateIDs) > 0 && !slices.Contains(arg.TemplateIDs, s.TemplateID) {
3434-
continue
3435-
}
3436-
if s.ConnectionCount == 0 {
3437-
continue
3447+
type interval struct {
3448+
From time.Time
3449+
To time.Time
3450+
}
3451+
var ts []interval
3452+
for d := arg.StartTime; d.Before(arg.EndTime); d = d.AddDate(0, 0, int(arg.IntervalDays)) {
3453+
to := d.AddDate(0, 0, int(arg.IntervalDays))
3454+
if to.After(arg.EndTime) {
3455+
to = arg.EndTime
34383456
}
3457+
ts = append(ts, interval{From: d, To: to})
3458+
}
34393459

3440-
for _, ds := range statsByInterval {
3441-
if s.CreatedAt.Before(ds.startTime) || s.CreatedAt.Equal(ds.endTime) || s.CreatedAt.After(ds.endTime) {
3460+
type grouped struct {
3461+
TemplateIDs map[uuid.UUID]struct{}
3462+
UserIDs map[uuid.UUID]struct{}
3463+
}
3464+
groupedByInterval := make(map[interval]grouped)
3465+
for _, tus := range q.templateUsageStats {
3466+
for _, t := range ts {
3467+
if tus.StartTime.Before(t.From) || tus.EndTime.After(t.To) {
34423468
continue
34433469
}
3444-
ds.userSet[s.UserID] = struct{}{}
3445-
ds.templateIDSet[s.TemplateID] = struct{}{}
3446-
}
3447-
}
3448-
3449-
for _, s := range q.workspaceAppStats {
3450-
w, err := q.getWorkspaceByIDNoLock(ctx, s.WorkspaceID)
3451-
if err != nil {
3452-
return nil, err
3453-
}
3454-
3455-
if len(arg.TemplateIDs) > 0 && !slices.Contains(arg.TemplateIDs, w.TemplateID) {
3456-
continue
3457-
}
3458-
3459-
for _, ds := range statsByInterval {
3460-
// (was.session_started_at >= ts.from_ AND was.session_started_at < ts.to_)
3461-
// OR (was.session_ended_at > ts.from_ AND was.session_ended_at < ts.to_)
3462-
// OR (was.session_started_at < ts.from_ AND was.session_ended_at >= ts.to_)
3463-
if !(((s.SessionStartedAt.After(ds.startTime) || s.SessionStartedAt.Equal(ds.startTime)) && s.SessionStartedAt.Before(ds.endTime)) ||
3464-
(s.SessionEndedAt.After(ds.startTime) && s.SessionEndedAt.Before(ds.endTime)) ||
3465-
(s.SessionStartedAt.Before(ds.startTime) && (s.SessionEndedAt.After(ds.endTime) || s.SessionEndedAt.Equal(ds.endTime)))) {
3470+
if len(arg.TemplateIDs) > 0 && !slices.Contains(arg.TemplateIDs, tus.TemplateID) {
34663471
continue
34673472
}
3468-
3469-
ds.userSet[s.UserID] = struct{}{}
3470-
ds.templateIDSet[w.TemplateID] = struct{}{}
3473+
g, ok := groupedByInterval[t]
3474+
if !ok {
3475+
g = grouped{
3476+
TemplateIDs: make(map[uuid.UUID]struct{}),
3477+
UserIDs: make(map[uuid.UUID]struct{}),
3478+
}
3479+
}
3480+
g.TemplateIDs[tus.TemplateID] = struct{}{}
3481+
g.UserIDs[tus.UserID] = struct{}{}
3482+
groupedByInterval[t] = g
34713483
}
34723484
}
34733485

3474-
var result []database.GetTemplateInsightsByIntervalRow
3475-
for _, ds := range statsByInterval {
3476-
templateIDs := make([]uuid.UUID, 0, len(ds.templateIDSet))
3477-
for templateID := range ds.templateIDSet {
3478-
templateIDs = append(templateIDs, templateID)
3486+
var rows []database.GetTemplateInsightsByIntervalRow
3487+
for _, t := range ts { // Ordered by interval.
3488+
row := database.GetTemplateInsightsByIntervalRow{
3489+
StartTime: t.From,
3490+
EndTime: t.To,
34793491
}
3480-
slices.SortFunc(templateIDs, func(a, b uuid.UUID) int {
3481-
return slice.Ascending(a.String(), b.String())
3482-
})
3483-
result = append(result, database.GetTemplateInsightsByIntervalRow{
3484-
StartTime: ds.startTime,
3485-
EndTime: ds.endTime,
3486-
TemplateIDs: templateIDs,
3487-
ActiveUsers: int64(len(ds.userSet)),
3488-
})
3492+
row.TemplateIDs = uniqueSortedUUIDs(maps.Keys(groupedByInterval[t].TemplateIDs))
3493+
row.ActiveUsers = int64(len(groupedByInterval[t].UserIDs))
3494+
rows = append(rows, row)
34893495
}
3490-
return result, nil
3496+
3497+
return rows, nil
34913498
}
34923499

34933500
func (q *FakeQuerier) GetTemplateInsightsByTemplate(_ context.Context, arg database.GetTemplateInsightsByTemplateParams) ([]database.GetTemplateInsightsByTemplateRow, error) {

0 commit comments

Comments
 (0)