Skip to content

Commit 5c165f1

Browse files
mafredrimafredri
committed
fix: implement GetTemplateAppInsights in dbmem
1 parent 1c951d1 commit 5c165f1

File tree

3 files changed

+272
-90
lines changed

3 files changed

+272
-90
lines changed

‎coderd/database/dbmem/dbmem.go

+270-90
Original file line numberDiff line numberDiff line change
@@ -799,6 +799,91 @@ 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+
build, err := q.getLatestWorkspaceBuildByWorkspaceIDNoLock(ctx, workspace.ID)
854+
if err != nil {
855+
continue
856+
}
857+
858+
resources, err := q.getWorkspaceResourcesByJobIDNoLock(ctx, build.JobID)
859+
if err != nil {
860+
continue
861+
}
862+
var resourceIDs []uuid.UUID
863+
for _, resource := range resources {
864+
resourceIDs = append(resourceIDs, resource.ID)
865+
}
866+
867+
agents, err := q.getWorkspaceAgentsByResourceIDsNoLock(ctx, resourceIDs)
868+
if err != nil {
869+
continue
870+
}
871+
872+
for _, agent := range agents {
873+
app, err := q.getWorkspaceAppByAgentIDAndSlugNoLock(ctx, database.GetWorkspaceAppByAgentIDAndSlugParams{
874+
AgentID: agent.ID,
875+
Slug: slug,
876+
})
877+
if err != nil {
878+
continue
879+
}
880+
return app, nil
881+
}
882+
}
883+
884+
return database.WorkspaceApp{}, sql.ErrNoRows
885+
}
886+
802887
func (*FakeQuerier) AcquireLock(_ context.Context, _ int64) error {
803888
return xerrors.New("AcquireLock must only be called within a transaction")
804889
}
@@ -2896,119 +2981,214 @@ func (q *FakeQuerier) GetTemplateAppInsights(ctx context.Context, arg database.G
28962981
q.mutex.RLock()
28972982
defer q.mutex.RUnlock()
28982983

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

2933-
app, _ := q.getWorkspaceAppByAgentIDAndSlugNoLock(ctx, database.GetWorkspaceAppByAgentIDAndSlugParams{
2934-
AgentID: s.AgentID,
2935-
Slug: s.SlugOrPort,
2936-
})
3068+
// json_each
3069+
for appName, appUsage := range stat.AppUsageMins {
3070+
// LEFT JOIN LATERAL
3071+
app, _ := q.getLatestWorkspaceAppByTemplateIDUserIDSlugNoLock(ctx, stat.TemplateID, stat.UserID, appName)
29373072

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

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

2964-
appUsageTemplateIDs := make(map[appKey]map[uuid.UUID]struct{})
2965-
appUsageUserIDs := make(map[appKey]map[uuid.UUID]struct{})
2966-
appUsage := make(map[appKey]int64)
2967-
for uniqueKey, usage := range appUsageIntervalsByUserAgentApp {
2968-
for _, seconds := range usage {
2969-
if appUsageTemplateIDs[uniqueKey.AppKey] == nil {
2970-
appUsageTemplateIDs[uniqueKey.AppKey] = make(map[uuid.UUID]struct{})
2971-
}
2972-
appUsageTemplateIDs[uniqueKey.AppKey][uniqueKey.TemplateID] = struct{}{}
2973-
if appUsageUserIDs[uniqueKey.AppKey] == nil {
2974-
appUsageUserIDs[uniqueKey.AppKey] = make(map[uuid.UUID]struct{})
3120+
templateRows := make(map[appGroupBy]templateRow)
3121+
for _, row := range appInsightRows {
3122+
key := appGroupBy{
3123+
AppName: row.AppName,
3124+
DisplayName: row.DisplayName,
3125+
Icon: row.Icon,
3126+
IsApp: row.IsApp,
3127+
}
3128+
row, ok := templateRows[key]
3129+
if !ok {
3130+
row = templateRow{
3131+
appGroupBy: key,
29753132
}
2976-
appUsageUserIDs[uniqueKey.AppKey][uniqueKey.UserID] = struct{}{}
2977-
appUsage[uniqueKey.AppKey] += seconds
29783133
}
3134+
row.TemplateIDs = append(row.TemplateIDs, row.TemplateIDs...)
3135+
templateRows[key] = row
29793136
}
29803137

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

3179+
var rows []database.GetTemplateAppInsightsRow
3180+
for k, gr := range groupedRows {
29983181
rows = append(rows, database.GetTemplateAppInsightsRow{
2999-
TemplateIDs: templateIDs,
3000-
ActiveUserIDs: activeUserIDs,
3001-
AccessMethod: appKey.AccessMethod,
3002-
SlugOrPort: appKey.SlugOrPort,
3003-
DisplayName: appKey.DisplayName,
3004-
Icon: appKey.Icon,
3005-
IsApp: appKey.Slug != "",
3006-
UsageSeconds: usage,
3182+
TemplateIDs: uniqueSortedUUIDs(gr.TemplateIDs),
3183+
ActiveUsers: int64(len(uniqueSortedUUIDs(gr.ActiveUserIDs))),
3184+
DisplayName: k.DisplayName,
3185+
SlugOrPort: k.AppName,
3186+
UsageSeconds: gr.UsageSeconds,
30073187
})
30083188
}
30093189

30103190
// NOTE(mafredri): Add sorting if we decide on how to handle PostgreSQL collations.
3011-
// ORDER BY access_method, slug_or_port, display_name, icon, is_app
3191+
// ORDER BY slug_or_port, display_name, icon, is_app
30123192
return rows, nil
30133193
}
30143194

0 commit comments

Comments
 (0)