Skip to content

Commit 3d14fe4

Browse files
mafredrimafredri
committed
fix: implement GetTemplateAppInsights in dbmem
1 parent 6fde859 commit 3d14fe4

File tree

3 files changed

+273
-90
lines changed

3 files changed

+273
-90
lines changed

‎coderd/database/dbmem/dbmem.go

+271-90
Original file line numberDiff line numberDiff line change
@@ -799,6 +799,92 @@ func least[T constraints.Ordered](a, b T) T {
799799
return b
800800
}
801801

802+
func (q *FakeQuerier) getLatestWorkspaceAppByTemplateIDUserIDSlugNoLock(ctx context.Context, templateID, userID uuid.UUID, slug string) (database.WorkspaceApp, error) {
803+
/*
804+
SELECT
805+
app.display_name,
806+
app.icon,
807+
app.slug
808+
FROM
809+
workspace_apps AS app
810+
JOIN
811+
workspace_agents AS agent
812+
ON
813+
agent.id = app.agent_id
814+
JOIN
815+
workspace_resources AS resource
816+
ON
817+
resource.id = agent.resource_id
818+
JOIN
819+
workspace_builds AS build
820+
ON
821+
build.job_id = resource.job_id
822+
JOIN
823+
workspaces AS workspace
824+
ON
825+
workspace.id = build.workspace_id
826+
WHERE
827+
-- Requires lateral join.
828+
app.slug = app_usage.key
829+
AND workspace.owner_id = tus.user_id
830+
AND workspace.template_id = tus.template_id
831+
ORDER BY
832+
app.created_at DESC
833+
LIMIT 1
834+
*/
835+
836+
var workspaces []database.Workspace
837+
for _, w := range q.workspaces {
838+
if w.TemplateID != templateID || w.OwnerID != userID {
839+
continue
840+
}
841+
workspaces = append(workspaces, w)
842+
}
843+
slices.SortFunc(workspaces, func(a, b database.Workspace) int {
844+
if a.CreatedAt.Before(b.CreatedAt) {
845+
return 1
846+
} else if a.CreatedAt.Equal(b.CreatedAt) {
847+
return 0
848+
}
849+
return -1
850+
})
851+
852+
for _, workspace := range workspaces {
853+
854+
build, err := q.getLatestWorkspaceBuildByWorkspaceIDNoLock(ctx, workspace.ID)
855+
if err != nil {
856+
continue
857+
}
858+
859+
resources, err := q.getWorkspaceResourcesByJobIDNoLock(ctx, build.JobID)
860+
if err != nil {
861+
continue
862+
}
863+
var resourceIDs []uuid.UUID
864+
for _, resource := range resources {
865+
resourceIDs = append(resourceIDs, resource.ID)
866+
}
867+
868+
agents, err := q.getWorkspaceAgentsByResourceIDsNoLock(ctx, resourceIDs)
869+
if err != nil {
870+
continue
871+
}
872+
873+
for _, agent := range agents {
874+
app, err := q.getWorkspaceAppByAgentIDAndSlugNoLock(ctx, database.GetWorkspaceAppByAgentIDAndSlugParams{
875+
AgentID: agent.ID,
876+
Slug: slug,
877+
})
878+
if err != nil {
879+
continue
880+
}
881+
return app, nil
882+
}
883+
}
884+
885+
return database.WorkspaceApp{}, sql.ErrNoRows
886+
}
887+
802888
func (*FakeQuerier) AcquireLock(_ context.Context, _ int64) error {
803889
return xerrors.New("AcquireLock must only be called within a transaction")
804890
}
@@ -2892,119 +2978,214 @@ func (q *FakeQuerier) GetTemplateAppInsights(ctx context.Context, arg database.G
28922978
q.mutex.RLock()
28932979
defer q.mutex.RUnlock()
28942980

2895-
type appKey struct {
2896-
AccessMethod string
2897-
SlugOrPort string
2898-
Slug string
2899-
DisplayName string
2900-
Icon string
2981+
/*
2982+
WITH
2983+
*/
2984+
2985+
/*
2986+
app_insights AS (
2987+
SELECT
2988+
tus.user_id,
2989+
array_agg(DISTINCT tus.template_id)::uuid[] AS template_ids,
2990+
app_usage.key::text AS app_name,
2991+
COALESCE(wa.display_name, '') AS display_name,
2992+
COALESCE(wa.icon, '') AS icon,
2993+
(wa.slug IS NOT NULL)::boolean AS is_app,
2994+
LEAST(SUM(app_usage.value::int), 30) AS app_usage_mins
2995+
FROM
2996+
template_usage_stats AS tus, jsonb_each(app_usage_mins) AS app_usage
2997+
LEFT JOIN LATERAL (
2998+
-- Fetch the latest app info for each app based on slug and template.
2999+
SELECT
3000+
app.display_name,
3001+
app.icon,
3002+
app.slug
3003+
FROM
3004+
workspace_apps AS app
3005+
JOIN
3006+
workspace_agents AS agent
3007+
ON
3008+
agent.id = app.agent_id
3009+
JOIN
3010+
workspace_resources AS resource
3011+
ON
3012+
resource.id = agent.resource_id
3013+
JOIN
3014+
workspace_builds AS build
3015+
ON
3016+
build.job_id = resource.job_id
3017+
JOIN
3018+
workspaces AS workspace
3019+
ON
3020+
workspace.id = build.workspace_id
3021+
WHERE
3022+
-- Requires lateral join.
3023+
app.slug = app_usage.key
3024+
AND workspace.owner_id = tus.user_id
3025+
AND workspace.template_id = tus.template_id
3026+
ORDER BY
3027+
app.created_at DESC
3028+
LIMIT 1
3029+
) AS wa
3030+
ON
3031+
true
3032+
WHERE
3033+
tus.start_time >= @start_time::timestamptz
3034+
AND tus.end_time <= @end_time::timestamptz
3035+
AND CASE WHEN COALESCE(array_length(@template_ids::uuid[], 1), 0) > 0 THEN tus.template_id = ANY(@template_ids::uuid[]) ELSE TRUE END
3036+
GROUP BY
3037+
tus.start_time, tus.user_id, app_usage.key::text, wa.display_name, wa.icon, wa.slug
3038+
),
3039+
*/
3040+
3041+
type appInsightsGroupBy struct {
3042+
StartTime time.Time
3043+
UserID uuid.UUID
3044+
AppName string
3045+
DisplayName string
3046+
Icon string
3047+
IsApp bool
29013048
}
2902-
type uniqueKey struct {
2903-
TemplateID uuid.UUID
2904-
UserID uuid.UUID
2905-
AgentID uuid.UUID
2906-
AppKey appKey
3049+
type appInsightsRow struct {
3050+
appInsightsGroupBy
3051+
TemplateIDs []uuid.UUID
3052+
AppUsageMins int64
29073053
}
2908-
2909-
appUsageIntervalsByUserAgentApp := make(map[uniqueKey]map[time.Time]int64)
2910-
for _, s := range q.workspaceAppStats {
2911-
// (was.session_started_at >= ts.from_ AND was.session_started_at < ts.to_)
2912-
// OR (was.session_ended_at > ts.from_ AND was.session_ended_at < ts.to_)
2913-
// OR (was.session_started_at < ts.from_ AND was.session_ended_at >= ts.to_)
2914-
if !(((s.SessionStartedAt.After(arg.StartTime) || s.SessionStartedAt.Equal(arg.StartTime)) && s.SessionStartedAt.Before(arg.EndTime)) ||
2915-
(s.SessionEndedAt.After(arg.StartTime) && s.SessionEndedAt.Before(arg.EndTime)) ||
2916-
(s.SessionStartedAt.Before(arg.StartTime) && (s.SessionEndedAt.After(arg.EndTime) || s.SessionEndedAt.Equal(arg.EndTime)))) {
3054+
appInsightRows := make(map[appInsightsGroupBy]appInsightsRow)
3055+
// FROM
3056+
for _, stat := range q.templateUsageStats {
3057+
// WHERE
3058+
if stat.StartTime.Before(arg.StartTime) || stat.StartTime.After(arg.EndTime) {
29173059
continue
29183060
}
2919-
2920-
w, err := q.getWorkspaceByIDNoLock(ctx, s.WorkspaceID)
2921-
if err != nil {
2922-
return nil, err
2923-
}
2924-
2925-
if len(arg.TemplateIDs) > 0 && !slices.Contains(arg.TemplateIDs, w.TemplateID) {
3061+
if len(arg.TemplateIDs) > 0 && !slices.Contains(arg.TemplateIDs, stat.TemplateID) {
29263062
continue
29273063
}
29283064

2929-
app, _ := q.getWorkspaceAppByAgentIDAndSlugNoLock(ctx, database.GetWorkspaceAppByAgentIDAndSlugParams{
2930-
AgentID: s.AgentID,
2931-
Slug: s.SlugOrPort,
2932-
})
3065+
// json_each
3066+
for appName, appUsage := range stat.AppUsageMins {
3067+
// LEFT JOIN LATERAL
3068+
app, _ := q.getLatestWorkspaceAppByTemplateIDUserIDSlugNoLock(ctx, stat.TemplateID, stat.UserID, appName)
29333069

2934-
key := uniqueKey{
2935-
TemplateID: w.TemplateID,
2936-
UserID: s.UserID,
2937-
AgentID: s.AgentID,
2938-
AppKey: appKey{
2939-
AccessMethod: s.AccessMethod,
2940-
SlugOrPort: s.SlugOrPort,
2941-
Slug: app.Slug,
2942-
DisplayName: app.DisplayName,
2943-
Icon: app.Icon,
2944-
},
2945-
}
2946-
if appUsageIntervalsByUserAgentApp[key] == nil {
2947-
appUsageIntervalsByUserAgentApp[key] = make(map[time.Time]int64)
3070+
// SELECT
3071+
key := appInsightsGroupBy{
3072+
StartTime: stat.StartTime,
3073+
UserID: stat.UserID,
3074+
AppName: appName,
3075+
DisplayName: app.DisplayName,
3076+
Icon: app.Icon,
3077+
IsApp: app.Slug != "",
3078+
}
3079+
row, ok := appInsightRows[key]
3080+
if !ok {
3081+
row = appInsightsRow{
3082+
appInsightsGroupBy: key,
3083+
}
3084+
}
3085+
row.TemplateIDs = append(row.TemplateIDs, stat.TemplateID)
3086+
row.AppUsageMins = least(row.AppUsageMins+appUsage, 30)
3087+
appInsightRows[key] = row
29483088
}
3089+
}
29493090

2950-
t := s.SessionStartedAt.Truncate(5 * time.Minute)
2951-
if t.Before(arg.StartTime) {
2952-
t = arg.StartTime
2953-
}
2954-
for t.Before(s.SessionEndedAt) && t.Before(arg.EndTime) {
2955-
appUsageIntervalsByUserAgentApp[key][t] = 60 // 1 minute.
2956-
t = t.Add(1 * time.Minute)
2957-
}
3091+
/*
3092+
templates AS (
3093+
SELECT
3094+
app_name,
3095+
display_name,
3096+
icon,
3097+
is_app,
3098+
array_agg(DISTINCT template_id)::uuid[] AS template_ids
3099+
FROM
3100+
app_insights, unnest(template_ids) AS template_id
3101+
GROUP BY
3102+
app_name, display_name, icon, is_app
3103+
)
3104+
*/
3105+
3106+
type appGroupBy struct {
3107+
AppName string
3108+
DisplayName string
3109+
Icon string
3110+
IsApp bool
3111+
}
3112+
type templateRow struct {
3113+
appGroupBy
3114+
TemplateIDs []uuid.UUID
29583115
}
29593116

2960-
appUsageTemplateIDs := make(map[appKey]map[uuid.UUID]struct{})
2961-
appUsageUserIDs := make(map[appKey]map[uuid.UUID]struct{})
2962-
appUsage := make(map[appKey]int64)
2963-
for uniqueKey, usage := range appUsageIntervalsByUserAgentApp {
2964-
for _, seconds := range usage {
2965-
if appUsageTemplateIDs[uniqueKey.AppKey] == nil {
2966-
appUsageTemplateIDs[uniqueKey.AppKey] = make(map[uuid.UUID]struct{})
2967-
}
2968-
appUsageTemplateIDs[uniqueKey.AppKey][uniqueKey.TemplateID] = struct{}{}
2969-
if appUsageUserIDs[uniqueKey.AppKey] == nil {
2970-
appUsageUserIDs[uniqueKey.AppKey] = make(map[uuid.UUID]struct{})
3117+
templateRows := make(map[appGroupBy]templateRow)
3118+
for _, row := range appInsightRows {
3119+
key := appGroupBy{
3120+
AppName: row.AppName,
3121+
DisplayName: row.DisplayName,
3122+
Icon: row.Icon,
3123+
IsApp: row.IsApp,
3124+
}
3125+
row, ok := templateRows[key]
3126+
if !ok {
3127+
row = templateRow{
3128+
appGroupBy: key,
29713129
}
2972-
appUsageUserIDs[uniqueKey.AppKey][uniqueKey.UserID] = struct{}{}
2973-
appUsage[uniqueKey.AppKey] += seconds
29743130
}
3131+
row.TemplateIDs = append(row.TemplateIDs, row.TemplateIDs...)
3132+
templateRows[key] = row
29753133
}
29763134

2977-
var rows []database.GetTemplateAppInsightsRow
2978-
for appKey, usage := range appUsage {
2979-
templateIDs := make([]uuid.UUID, 0, len(appUsageTemplateIDs[appKey]))
2980-
for templateID := range appUsageTemplateIDs[appKey] {
2981-
templateIDs = append(templateIDs, templateID)
2982-
}
2983-
slices.SortFunc(templateIDs, func(a, b uuid.UUID) int {
2984-
return slice.Ascending(a.String(), b.String())
2985-
})
2986-
activeUserIDs := make([]uuid.UUID, 0, len(appUsageUserIDs[appKey]))
2987-
for userID := range appUsageUserIDs[appKey] {
2988-
activeUserIDs = append(activeUserIDs, userID)
3135+
/*
3136+
SELECT
3137+
t.template_ids,
3138+
array_agg(DISTINCT ai.user_id)::uuid[] AS active_user_ids,
3139+
ai.app_name AS slug_or_port,
3140+
ai.display_name,
3141+
ai.icon,
3142+
ai.is_app,
3143+
(SUM(ai.app_usage_mins) * 60)::bigint AS usage_seconds
3144+
FROM
3145+
app_insights AS ai
3146+
JOIN
3147+
templates AS t
3148+
ON
3149+
ai.app_name = t.app_name
3150+
AND ai.display_name = t.display_name
3151+
AND ai.icon = t.icon
3152+
AND ai.is_app = t.is_app
3153+
GROUP BY
3154+
t.template_ids, ai.app_name, ai.display_name, ai.icon, ai.is_app;
3155+
*/
3156+
3157+
type templateAppInsightsRow struct {
3158+
TemplateIDs []uuid.UUID
3159+
ActiveUserIDs []uuid.UUID
3160+
UsageSeconds int64
3161+
}
3162+
groupedRows := make(map[appGroupBy]templateAppInsightsRow)
3163+
for _, aiRow := range appInsightRows {
3164+
key := appGroupBy{
3165+
AppName: aiRow.AppName,
3166+
DisplayName: aiRow.DisplayName,
3167+
Icon: aiRow.Icon,
3168+
IsApp: aiRow.IsApp,
29893169
}
2990-
slices.SortFunc(activeUserIDs, func(a, b uuid.UUID) int {
2991-
return slice.Ascending(a.String(), b.String())
2992-
})
3170+
row := groupedRows[key]
3171+
row.TemplateIDs = append(row.TemplateIDs, aiRow.TemplateIDs...)
3172+
row.ActiveUserIDs = append(row.ActiveUserIDs, aiRow.UserID)
3173+
row.UsageSeconds += aiRow.AppUsageMins * 60
3174+
}
29933175

3176+
var rows []database.GetTemplateAppInsightsRow
3177+
for k, gr := range groupedRows {
29943178
rows = append(rows, database.GetTemplateAppInsightsRow{
2995-
TemplateIDs: templateIDs,
2996-
ActiveUserIDs: activeUserIDs,
2997-
AccessMethod: appKey.AccessMethod,
2998-
SlugOrPort: appKey.SlugOrPort,
2999-
DisplayName: appKey.DisplayName,
3000-
Icon: appKey.Icon,
3001-
IsApp: appKey.Slug != "",
3002-
UsageSeconds: usage,
3179+
TemplateIDs: uniqueSortedUUIDs(gr.TemplateIDs),
3180+
ActiveUsers: int64(len(uniqueSortedUUIDs(gr.ActiveUserIDs))),
3181+
DisplayName: k.DisplayName,
3182+
SlugOrPort: k.AppName,
3183+
UsageSeconds: gr.UsageSeconds,
30033184
})
30043185
}
30053186

30063187
// NOTE(mafredri): Add sorting if we decide on how to handle PostgreSQL collations.
3007-
// ORDER BY access_method, slug_or_port, display_name, icon, is_app
3188+
// ORDER BY slug_or_port, display_name, icon, is_app
30083189
return rows, nil
30093190
}
30103191

0 commit comments

Comments
 (0)