fix: handle new agent stat format correctly (#14576)

---------

Co-authored-by: Ethan Dickson <ethan@coder.com>
This commit is contained in:
Garrett Delfosse
2024-09-20 01:52:14 +10:00
committed by GitHub
co-authored by Ethan Dickson
parent 37885e2e82
commit 922f4c545f
28 changed files with 1495 additions and 70 deletions
+12
View File
@@ -1447,6 +1447,10 @@ func (q *querier) GetDeploymentWorkspaceAgentStats(ctx context.Context, createdA
return q.db.GetDeploymentWorkspaceAgentStats(ctx, createdAfter)
}
func (q *querier) GetDeploymentWorkspaceAgentUsageStats(ctx context.Context, createdAt time.Time) (database.GetDeploymentWorkspaceAgentUsageStatsRow, error) {
return q.db.GetDeploymentWorkspaceAgentUsageStats(ctx, createdAt)
}
func (q *querier) GetDeploymentWorkspaceStats(ctx context.Context) (database.GetDeploymentWorkspaceStatsRow, error) {
return q.db.GetDeploymentWorkspaceStats(ctx)
}
@@ -2425,6 +2429,14 @@ func (q *querier) GetWorkspaceAgentStatsAndLabels(ctx context.Context, createdAf
return q.db.GetWorkspaceAgentStatsAndLabels(ctx, createdAfter)
}
func (q *querier) GetWorkspaceAgentUsageStats(ctx context.Context, createdAt time.Time) ([]database.GetWorkspaceAgentUsageStatsRow, error) {
return q.db.GetWorkspaceAgentUsageStats(ctx, createdAt)
}
func (q *querier) GetWorkspaceAgentUsageStatsAndLabels(ctx context.Context, createdAt time.Time) ([]database.GetWorkspaceAgentUsageStatsAndLabelsRow, error) {
return q.db.GetWorkspaceAgentUsageStatsAndLabels(ctx, createdAt)
}
// GetWorkspaceAgentsByResourceIDs
// The workspace/job is already fetched.
func (q *querier) GetWorkspaceAgentsByResourceIDs(ctx context.Context, ids []uuid.UUID) ([]database.WorkspaceAgent, error) {
+9
View File
@@ -2681,6 +2681,9 @@ func (s *MethodTestSuite) TestSystemFunctions() {
s.Run("GetDeploymentWorkspaceAgentStats", s.Subtest(func(db database.Store, check *expects) {
check.Args(time.Time{}).Asserts()
}))
s.Run("GetDeploymentWorkspaceAgentUsageStats", s.Subtest(func(db database.Store, check *expects) {
check.Args(time.Time{}).Asserts()
}))
s.Run("GetDeploymentWorkspaceStats", s.Subtest(func(db database.Store, check *expects) {
check.Args().Asserts()
}))
@@ -2717,9 +2720,15 @@ func (s *MethodTestSuite) TestSystemFunctions() {
s.Run("GetWorkspaceAgentStatsAndLabels", s.Subtest(func(db database.Store, check *expects) {
check.Args(time.Time{}).Asserts()
}))
s.Run("GetWorkspaceAgentUsageStatsAndLabels", s.Subtest(func(db database.Store, check *expects) {
check.Args(time.Time{}).Asserts()
}))
s.Run("GetWorkspaceAgentStats", s.Subtest(func(db database.Store, check *expects) {
check.Args(time.Time{}).Asserts()
}))
s.Run("GetWorkspaceAgentUsageStats", s.Subtest(func(db database.Store, check *expects) {
check.Args(time.Time{}).Asserts()
}))
s.Run("GetWorkspaceProxyByHostname", s.Subtest(func(db database.Store, check *expects) {
p, _ := dbgen.WorkspaceProxy(s.T(), db, database.WorkspaceProxy{
WildcardHostname: "*.example.com",
+2
View File
@@ -803,6 +803,7 @@ func WorkspaceAgentStat(t testing.TB, db database.Store, orig database.Workspace
SessionCountReconnectingPTY: []int64{takeFirst(orig.SessionCountReconnectingPTY, 0)},
SessionCountSSH: []int64{takeFirst(orig.SessionCountSSH, 0)},
ConnectionMedianLatencyMS: []float64{takeFirst(orig.ConnectionMedianLatencyMS, 0)},
Usage: []bool{takeFirst(orig.Usage, false)},
}
err := db.InsertWorkspaceAgentStats(genCtx, params)
require.NoError(t, err, "insert workspace agent stat")
@@ -825,6 +826,7 @@ func WorkspaceAgentStat(t testing.TB, db database.Store, orig database.Workspace
SessionCountJetBrains: params.SessionCountJetBrains[0],
SessionCountReconnectingPTY: params.SessionCountReconnectingPTY[0],
SessionCountSSH: params.SessionCountSSH[0],
Usage: params.Usage[0],
}
}
+283 -32
View File
@@ -7,6 +7,7 @@ import (
"encoding/json"
"errors"
"fmt"
"math"
"reflect"
"regexp"
"sort"
@@ -254,6 +255,19 @@ type data struct {
defaultProxyIconURL string
}
func tryPercentile(fs []float64, p float64) float64 {
if len(fs) == 0 {
return -1
}
sort.Float64s(fs)
pos := p * (float64(len(fs)) - 1) / 100
lower, upper := int(pos), int(math.Ceil(pos))
if lower == upper {
return fs[lower]
}
return fs[lower] + (fs[upper]-fs[lower])*(pos-float64(lower))
}
func validateDatabaseTypeWithValid(v reflect.Value) (handled bool, err error) {
if v.Kind() == reflect.Struct {
return false, nil
@@ -2533,20 +2547,68 @@ func (q *FakeQuerier) GetDeploymentWorkspaceAgentStats(_ context.Context, create
latencies = append(latencies, agentStat.ConnectionMedianLatencyMS)
}
tryPercentile := func(fs []float64, p float64) float64 {
if len(fs) == 0 {
return -1
}
sort.Float64s(fs)
return fs[int(float64(len(fs))*p/100)]
}
stat.WorkspaceConnectionLatency50 = tryPercentile(latencies, 50)
stat.WorkspaceConnectionLatency95 = tryPercentile(latencies, 95)
return stat, nil
}
func (q *FakeQuerier) GetDeploymentWorkspaceAgentUsageStats(_ context.Context, createdAt time.Time) (database.GetDeploymentWorkspaceAgentUsageStatsRow, error) {
q.mutex.RLock()
defer q.mutex.RUnlock()
stat := database.GetDeploymentWorkspaceAgentUsageStatsRow{}
sessions := make(map[uuid.UUID]database.WorkspaceAgentStat)
agentStatsCreatedAfter := make([]database.WorkspaceAgentStat, 0)
for _, agentStat := range q.workspaceAgentStats {
// WHERE workspace_agent_stats.created_at > $1
if agentStat.CreatedAt.After(createdAt) {
agentStatsCreatedAfter = append(agentStatsCreatedAfter, agentStat)
}
// WHERE
// created_at > $1
// AND created_at < date_trunc('minute', now()) -- Exclude current partial minute
// AND usage = true
if agentStat.Usage &&
(agentStat.CreatedAt.After(createdAt) || agentStat.CreatedAt.Equal(createdAt)) &&
agentStat.CreatedAt.Before(time.Now().Truncate(time.Minute)) {
val, ok := sessions[agentStat.AgentID]
if !ok {
sessions[agentStat.AgentID] = agentStat
} else if agentStat.CreatedAt.After(val.CreatedAt) {
sessions[agentStat.AgentID] = agentStat
} else if agentStat.CreatedAt.Truncate(time.Minute).Equal(val.CreatedAt.Truncate(time.Minute)) {
val.SessionCountVSCode += agentStat.SessionCountVSCode
val.SessionCountJetBrains += agentStat.SessionCountJetBrains
val.SessionCountReconnectingPTY += agentStat.SessionCountReconnectingPTY
val.SessionCountSSH += agentStat.SessionCountSSH
sessions[agentStat.AgentID] = val
}
}
}
latencies := make([]float64, 0)
for _, agentStat := range agentStatsCreatedAfter {
if agentStat.ConnectionMedianLatencyMS <= 0 {
continue
}
stat.WorkspaceRxBytes += agentStat.RxBytes
stat.WorkspaceTxBytes += agentStat.TxBytes
latencies = append(latencies, agentStat.ConnectionMedianLatencyMS)
}
stat.WorkspaceConnectionLatency50 = tryPercentile(latencies, 50)
stat.WorkspaceConnectionLatency95 = tryPercentile(latencies, 95)
for _, agentStat := range sessions {
stat.SessionCountVSCode += agentStat.SessionCountVSCode
stat.SessionCountJetBrains += agentStat.SessionCountJetBrains
stat.SessionCountReconnectingPTY += agentStat.SessionCountReconnectingPTY
stat.SessionCountSSH += agentStat.SessionCountSSH
}
return stat, nil
}
func (q *FakeQuerier) GetDeploymentWorkspaceStats(ctx context.Context) (database.GetDeploymentWorkspaceStatsRow, error) {
q.mutex.RLock()
defer q.mutex.RUnlock()
@@ -4238,14 +4300,6 @@ func (q *FakeQuerier) GetTemplateAverageBuildTime(ctx context.Context, arg datab
}
}
tryPercentile := func(fs []float64, p float64) float64 {
if len(fs) == 0 {
return -1
}
sort.Float64s(fs)
return fs[int(float64(len(fs))*p/100)]
}
var row database.GetTemplateAverageBuildTimeRow
row.Delete50, row.Delete95 = tryPercentile(deleteTimes, 50), tryPercentile(deleteTimes, 95)
row.Stop50, row.Stop95 = tryPercentile(stopTimes, 50), tryPercentile(stopTimes, 95)
@@ -5273,14 +5327,6 @@ func (q *FakeQuerier) GetUserLatencyInsights(_ context.Context, arg database.Get
seenTemplatesByUserID[stat.UserID] = uniqueSortedUUIDs(append(seenTemplatesByUserID[stat.UserID], stat.TemplateID))
}
tryPercentile := func(fs []float64, p float64) float64 {
if len(fs) == 0 {
return -1
}
sort.Float64s(fs)
return fs[int(float64(len(fs))*p/100)]
}
var rows []database.GetUserLatencyInsightsRow
for userID, latencies := range latenciesByUserID {
user, err := q.getUserByIDNoLock(userID)
@@ -5794,14 +5840,6 @@ func (q *FakeQuerier) GetWorkspaceAgentStats(_ context.Context, createdAfter tim
latenciesByAgent[agentStat.AgentID] = append(latenciesByAgent[agentStat.AgentID], agentStat.ConnectionMedianLatencyMS)
}
tryPercentile := func(fs []float64, p float64) float64 {
if len(fs) == 0 {
return -1
}
sort.Float64s(fs)
return fs[int(float64(len(fs))*p/100)]
}
for _, stat := range statByAgent {
stat.AggregatedFrom = minimumDateByAgent[stat.AgentID]
statByAgent[stat.AgentID] = stat
@@ -5893,6 +5931,218 @@ func (q *FakeQuerier) GetWorkspaceAgentStatsAndLabels(ctx context.Context, creat
return stats, nil
}
func (q *FakeQuerier) GetWorkspaceAgentUsageStats(_ context.Context, createdAt time.Time) ([]database.GetWorkspaceAgentUsageStatsRow, error) {
q.mutex.RLock()
defer q.mutex.RUnlock()
type agentStatsKey struct {
UserID uuid.UUID
AgentID uuid.UUID
WorkspaceID uuid.UUID
TemplateID uuid.UUID
}
type minuteStatsKey struct {
agentStatsKey
MinuteBucket time.Time
}
latestAgentStats := map[agentStatsKey]database.GetWorkspaceAgentUsageStatsRow{}
latestAgentLatencies := map[agentStatsKey][]float64{}
for _, agentStat := range q.workspaceAgentStats {
key := agentStatsKey{
UserID: agentStat.UserID,
AgentID: agentStat.AgentID,
WorkspaceID: agentStat.WorkspaceID,
TemplateID: agentStat.TemplateID,
}
if agentStat.CreatedAt.After(createdAt) {
val, ok := latestAgentStats[key]
if ok {
val.WorkspaceRxBytes += agentStat.RxBytes
val.WorkspaceTxBytes += agentStat.TxBytes
latestAgentStats[key] = val
} else {
latestAgentStats[key] = database.GetWorkspaceAgentUsageStatsRow{
UserID: agentStat.UserID,
AgentID: agentStat.AgentID,
WorkspaceID: agentStat.WorkspaceID,
TemplateID: agentStat.TemplateID,
AggregatedFrom: createdAt,
WorkspaceRxBytes: agentStat.RxBytes,
WorkspaceTxBytes: agentStat.TxBytes,
}
}
latencies, ok := latestAgentLatencies[key]
if !ok {
latestAgentLatencies[key] = []float64{agentStat.ConnectionMedianLatencyMS}
} else {
latestAgentLatencies[key] = append(latencies, agentStat.ConnectionMedianLatencyMS)
}
}
}
for key, latencies := range latestAgentLatencies {
val, ok := latestAgentStats[key]
if ok {
val.WorkspaceConnectionLatency50 = tryPercentile(latencies, 50)
val.WorkspaceConnectionLatency95 = tryPercentile(latencies, 95)
}
latestAgentStats[key] = val
}
type bucketRow struct {
database.GetWorkspaceAgentUsageStatsRow
MinuteBucket time.Time
}
minuteBuckets := make(map[minuteStatsKey]bucketRow)
for _, agentStat := range q.workspaceAgentStats {
if agentStat.Usage &&
(agentStat.CreatedAt.After(createdAt) || agentStat.CreatedAt.Equal(createdAt)) &&
agentStat.CreatedAt.Before(time.Now().Truncate(time.Minute)) {
key := minuteStatsKey{
agentStatsKey: agentStatsKey{
UserID: agentStat.UserID,
AgentID: agentStat.AgentID,
WorkspaceID: agentStat.WorkspaceID,
TemplateID: agentStat.TemplateID,
},
MinuteBucket: agentStat.CreatedAt.Truncate(time.Minute),
}
val, ok := minuteBuckets[key]
if ok {
val.SessionCountVSCode += agentStat.SessionCountVSCode
val.SessionCountJetBrains += agentStat.SessionCountJetBrains
val.SessionCountReconnectingPTY += agentStat.SessionCountReconnectingPTY
val.SessionCountSSH += agentStat.SessionCountSSH
minuteBuckets[key] = val
} else {
minuteBuckets[key] = bucketRow{
GetWorkspaceAgentUsageStatsRow: database.GetWorkspaceAgentUsageStatsRow{
UserID: agentStat.UserID,
AgentID: agentStat.AgentID,
WorkspaceID: agentStat.WorkspaceID,
TemplateID: agentStat.TemplateID,
SessionCountVSCode: agentStat.SessionCountVSCode,
SessionCountSSH: agentStat.SessionCountSSH,
SessionCountJetBrains: agentStat.SessionCountJetBrains,
SessionCountReconnectingPTY: agentStat.SessionCountReconnectingPTY,
},
MinuteBucket: agentStat.CreatedAt.Truncate(time.Minute),
}
}
}
}
// Get the latest minute bucket for each agent.
latestBuckets := make(map[uuid.UUID]bucketRow)
for key, bucket := range minuteBuckets {
latest, ok := latestBuckets[key.AgentID]
if !ok || key.MinuteBucket.After(latest.MinuteBucket) {
latestBuckets[key.AgentID] = bucket
}
}
for key, stat := range latestAgentStats {
bucket, ok := latestBuckets[stat.AgentID]
if ok {
stat.SessionCountVSCode = bucket.SessionCountVSCode
stat.SessionCountJetBrains = bucket.SessionCountJetBrains
stat.SessionCountReconnectingPTY = bucket.SessionCountReconnectingPTY
stat.SessionCountSSH = bucket.SessionCountSSH
}
latestAgentStats[key] = stat
}
return maps.Values(latestAgentStats), nil
}
func (q *FakeQuerier) GetWorkspaceAgentUsageStatsAndLabels(_ context.Context, createdAt time.Time) ([]database.GetWorkspaceAgentUsageStatsAndLabelsRow, error) {
q.mutex.RLock()
defer q.mutex.RUnlock()
type statsKey struct {
AgentID uuid.UUID
UserID uuid.UUID
WorkspaceID uuid.UUID
}
latestAgentStats := map[statsKey]database.WorkspaceAgentStat{}
maxConnMedianLatency := 0.0
for _, agentStat := range q.workspaceAgentStats {
key := statsKey{
AgentID: agentStat.AgentID,
UserID: agentStat.UserID,
WorkspaceID: agentStat.WorkspaceID,
}
// WHERE workspace_agent_stats.created_at > $1
// GROUP BY user_id, agent_id, workspace_id
if agentStat.CreatedAt.After(createdAt) {
val, ok := latestAgentStats[key]
if !ok {
val = agentStat
val.SessionCountJetBrains = 0
val.SessionCountReconnectingPTY = 0
val.SessionCountSSH = 0
val.SessionCountVSCode = 0
} else {
val.RxBytes += agentStat.RxBytes
val.TxBytes += agentStat.TxBytes
}
if agentStat.ConnectionMedianLatencyMS > maxConnMedianLatency {
val.ConnectionMedianLatencyMS = agentStat.ConnectionMedianLatencyMS
}
latestAgentStats[key] = val
}
// WHERE usage = true AND created_at > now() - '1 minute'::interval
// GROUP BY user_id, agent_id, workspace_id
if agentStat.Usage && agentStat.CreatedAt.After(time.Now().Add(-time.Minute)) {
val, ok := latestAgentStats[key]
if !ok {
latestAgentStats[key] = agentStat
} else {
val.SessionCountVSCode += agentStat.SessionCountVSCode
val.SessionCountJetBrains += agentStat.SessionCountJetBrains
val.SessionCountReconnectingPTY += agentStat.SessionCountReconnectingPTY
val.SessionCountSSH += agentStat.SessionCountSSH
val.ConnectionCount += agentStat.ConnectionCount
latestAgentStats[key] = val
}
}
}
stats := make([]database.GetWorkspaceAgentUsageStatsAndLabelsRow, 0, len(latestAgentStats))
for key, agentStat := range latestAgentStats {
user, err := q.getUserByIDNoLock(key.UserID)
if err != nil {
return nil, err
}
workspace, err := q.getWorkspaceByIDNoLock(context.Background(), key.WorkspaceID)
if err != nil {
return nil, err
}
agent, err := q.getWorkspaceAgentByIDNoLock(context.Background(), key.AgentID)
if err != nil {
return nil, err
}
stats = append(stats, database.GetWorkspaceAgentUsageStatsAndLabelsRow{
Username: user.Username,
AgentName: agent.Name,
WorkspaceName: workspace.Name,
RxBytes: agentStat.RxBytes,
TxBytes: agentStat.TxBytes,
SessionCountVSCode: agentStat.SessionCountVSCode,
SessionCountSSH: agentStat.SessionCountSSH,
SessionCountJetBrains: agentStat.SessionCountJetBrains,
SessionCountReconnectingPTY: agentStat.SessionCountReconnectingPTY,
ConnectionCount: agentStat.ConnectionCount,
ConnectionMedianLatencyMS: agentStat.ConnectionMedianLatencyMS,
})
}
return stats, nil
}
func (q *FakeQuerier) GetWorkspaceAgentsByResourceIDs(ctx context.Context, resourceIDs []uuid.UUID) ([]database.WorkspaceAgent, error) {
q.mutex.RLock()
defer q.mutex.RUnlock()
@@ -7641,6 +7891,7 @@ func (q *FakeQuerier) InsertWorkspaceAgentStats(_ context.Context, arg database.
SessionCountReconnectingPTY: arg.SessionCountReconnectingPTY[i],
SessionCountSSH: arg.SessionCountSSH[i],
ConnectionMedianLatencyMS: arg.ConnectionMedianLatencyMS[i],
Usage: arg.Usage[i],
}
q.workspaceAgentStats = append(q.workspaceAgentStats, stat)
}
+21
View File
@@ -613,6 +613,13 @@ func (m metricsStore) GetDeploymentWorkspaceAgentStats(ctx context.Context, crea
return row, err
}
func (m metricsStore) GetDeploymentWorkspaceAgentUsageStats(ctx context.Context, createdAt time.Time) (database.GetDeploymentWorkspaceAgentUsageStatsRow, error) {
start := time.Now()
r0, r1 := m.s.GetDeploymentWorkspaceAgentUsageStats(ctx, createdAt)
m.queryLatencies.WithLabelValues("GetDeploymentWorkspaceAgentUsageStats").Observe(time.Since(start).Seconds())
return r0, r1
}
func (m metricsStore) GetDeploymentWorkspaceStats(ctx context.Context) (database.GetDeploymentWorkspaceStatsRow, error) {
start := time.Now()
row, err := m.s.GetDeploymentWorkspaceStats(ctx)
@@ -1411,6 +1418,20 @@ func (m metricsStore) GetWorkspaceAgentStatsAndLabels(ctx context.Context, creat
return stats, err
}
func (m metricsStore) GetWorkspaceAgentUsageStats(ctx context.Context, createdAt time.Time) ([]database.GetWorkspaceAgentUsageStatsRow, error) {
start := time.Now()
r0, r1 := m.s.GetWorkspaceAgentUsageStats(ctx, createdAt)
m.queryLatencies.WithLabelValues("GetWorkspaceAgentUsageStats").Observe(time.Since(start).Seconds())
return r0, r1
}
func (m metricsStore) GetWorkspaceAgentUsageStatsAndLabels(ctx context.Context, createdAt time.Time) ([]database.GetWorkspaceAgentUsageStatsAndLabelsRow, error) {
start := time.Now()
r0, r1 := m.s.GetWorkspaceAgentUsageStatsAndLabels(ctx, createdAt)
m.queryLatencies.WithLabelValues("GetWorkspaceAgentUsageStatsAndLabels").Observe(time.Since(start).Seconds())
return r0, r1
}
func (m metricsStore) GetWorkspaceAgentsByResourceIDs(ctx context.Context, ids []uuid.UUID) ([]database.WorkspaceAgent, error) {
start := time.Now()
agents, err := m.s.GetWorkspaceAgentsByResourceIDs(ctx, ids)
+45
View File
@@ -1208,6 +1208,21 @@ func (mr *MockStoreMockRecorder) GetDeploymentWorkspaceAgentStats(arg0, arg1 any
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetDeploymentWorkspaceAgentStats", reflect.TypeOf((*MockStore)(nil).GetDeploymentWorkspaceAgentStats), arg0, arg1)
}
// GetDeploymentWorkspaceAgentUsageStats mocks base method.
func (m *MockStore) GetDeploymentWorkspaceAgentUsageStats(arg0 context.Context, arg1 time.Time) (database.GetDeploymentWorkspaceAgentUsageStatsRow, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetDeploymentWorkspaceAgentUsageStats", arg0, arg1)
ret0, _ := ret[0].(database.GetDeploymentWorkspaceAgentUsageStatsRow)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetDeploymentWorkspaceAgentUsageStats indicates an expected call of GetDeploymentWorkspaceAgentUsageStats.
func (mr *MockStoreMockRecorder) GetDeploymentWorkspaceAgentUsageStats(arg0, arg1 any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetDeploymentWorkspaceAgentUsageStats", reflect.TypeOf((*MockStore)(nil).GetDeploymentWorkspaceAgentUsageStats), arg0, arg1)
}
// GetDeploymentWorkspaceStats mocks base method.
func (m *MockStore) GetDeploymentWorkspaceStats(arg0 context.Context) (database.GetDeploymentWorkspaceStatsRow, error) {
m.ctrl.T.Helper()
@@ -2948,6 +2963,36 @@ func (mr *MockStoreMockRecorder) GetWorkspaceAgentStatsAndLabels(arg0, arg1 any)
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceAgentStatsAndLabels", reflect.TypeOf((*MockStore)(nil).GetWorkspaceAgentStatsAndLabels), arg0, arg1)
}
// GetWorkspaceAgentUsageStats mocks base method.
func (m *MockStore) GetWorkspaceAgentUsageStats(arg0 context.Context, arg1 time.Time) ([]database.GetWorkspaceAgentUsageStatsRow, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetWorkspaceAgentUsageStats", arg0, arg1)
ret0, _ := ret[0].([]database.GetWorkspaceAgentUsageStatsRow)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetWorkspaceAgentUsageStats indicates an expected call of GetWorkspaceAgentUsageStats.
func (mr *MockStoreMockRecorder) GetWorkspaceAgentUsageStats(arg0, arg1 any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceAgentUsageStats", reflect.TypeOf((*MockStore)(nil).GetWorkspaceAgentUsageStats), arg0, arg1)
}
// GetWorkspaceAgentUsageStatsAndLabels mocks base method.
func (m *MockStore) GetWorkspaceAgentUsageStatsAndLabels(arg0 context.Context, arg1 time.Time) ([]database.GetWorkspaceAgentUsageStatsAndLabelsRow, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetWorkspaceAgentUsageStatsAndLabels", arg0, arg1)
ret0, _ := ret[0].([]database.GetWorkspaceAgentUsageStatsAndLabelsRow)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetWorkspaceAgentUsageStatsAndLabels indicates an expected call of GetWorkspaceAgentUsageStatsAndLabels.
func (mr *MockStoreMockRecorder) GetWorkspaceAgentUsageStatsAndLabels(arg0, arg1 any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceAgentUsageStatsAndLabels", reflect.TypeOf((*MockStore)(nil).GetWorkspaceAgentUsageStatsAndLabels), arg0, arg1)
}
// GetWorkspaceAgentsByResourceIDs mocks base method.
func (m *MockStore) GetWorkspaceAgentsByResourceIDs(arg0 context.Context, arg1 []uuid.UUID) ([]database.WorkspaceAgent, error) {
m.ctrl.T.Helper()
+2 -1
View File
@@ -1394,7 +1394,8 @@ CREATE TABLE workspace_agent_stats (
session_count_vscode bigint DEFAULT 0 NOT NULL,
session_count_jetbrains bigint DEFAULT 0 NOT NULL,
session_count_reconnecting_pty bigint DEFAULT 0 NOT NULL,
session_count_ssh bigint DEFAULT 0 NOT NULL
session_count_ssh bigint DEFAULT 0 NOT NULL,
usage boolean DEFAULT false NOT NULL
);
CREATE TABLE workspace_agents (
@@ -0,0 +1 @@
ALTER TABLE workspace_agent_stats DROP COLUMN usage;
@@ -0,0 +1 @@
ALTER TABLE workspace_agent_stats ADD COLUMN usage boolean NOT NULL DEFAULT false;
+1
View File
@@ -2900,6 +2900,7 @@ type WorkspaceAgentStat struct {
SessionCountJetBrains int64 `db:"session_count_jetbrains" json:"session_count_jetbrains"`
SessionCountReconnectingPTY int64 `db:"session_count_reconnecting_pty" json:"session_count_reconnecting_pty"`
SessionCountSSH int64 `db:"session_count_ssh" json:"session_count_ssh"`
Usage bool `db:"usage" json:"usage"`
}
type WorkspaceApp struct {
+4
View File
@@ -141,6 +141,7 @@ type sqlcQuerier interface {
GetDeploymentDAUs(ctx context.Context, tzOffset int32) ([]GetDeploymentDAUsRow, error)
GetDeploymentID(ctx context.Context) (string, error)
GetDeploymentWorkspaceAgentStats(ctx context.Context, createdAt time.Time) (GetDeploymentWorkspaceAgentStatsRow, error)
GetDeploymentWorkspaceAgentUsageStats(ctx context.Context, createdAt time.Time) (GetDeploymentWorkspaceAgentUsageStatsRow, error)
GetDeploymentWorkspaceStats(ctx context.Context) (GetDeploymentWorkspaceStatsRow, error)
GetExternalAuthLink(ctx context.Context, arg GetExternalAuthLinkParams) (ExternalAuthLink, error)
GetExternalAuthLinksByUserID(ctx context.Context, userID uuid.UUID) ([]ExternalAuthLink, error)
@@ -299,6 +300,9 @@ type sqlcQuerier interface {
GetWorkspaceAgentScriptsByAgentIDs(ctx context.Context, ids []uuid.UUID) ([]WorkspaceAgentScript, error)
GetWorkspaceAgentStats(ctx context.Context, createdAt time.Time) ([]GetWorkspaceAgentStatsRow, error)
GetWorkspaceAgentStatsAndLabels(ctx context.Context, createdAt time.Time) ([]GetWorkspaceAgentStatsAndLabelsRow, error)
// `minute_buckets` could return 0 rows if there are no usage stats since `created_at`.
GetWorkspaceAgentUsageStats(ctx context.Context, createdAt time.Time) ([]GetWorkspaceAgentUsageStatsRow, error)
GetWorkspaceAgentUsageStatsAndLabels(ctx context.Context, createdAt time.Time) ([]GetWorkspaceAgentUsageStatsAndLabelsRow, error)
GetWorkspaceAgentsByResourceIDs(ctx context.Context, ids []uuid.UUID) ([]WorkspaceAgent, error)
GetWorkspaceAgentsCreatedAfter(ctx context.Context, createdAt time.Time) ([]WorkspaceAgent, error)
GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx context.Context, workspaceID uuid.UUID) ([]WorkspaceAgent, error)
+512
View File
@@ -100,6 +100,518 @@ func TestGetDeploymentWorkspaceAgentStats(t *testing.T) {
})
}
func TestGetDeploymentWorkspaceAgentUsageStats(t *testing.T) {
t.Parallel()
t.Run("Aggregates", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
authz := rbac.NewAuthorizer(prometheus.NewRegistry())
db = dbauthz.New(db, authz, slogtest.Make(t, &slogtest.Options{}), coderdtest.AccessControlStorePointer())
ctx := context.Background()
agentID := uuid.New()
// Since the queries exclude the current minute
insertTime := dbtime.Now().Add(-time.Minute)
// Old stats
dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
CreatedAt: insertTime.Add(-time.Minute),
AgentID: agentID,
TxBytes: 1,
RxBytes: 1,
ConnectionMedianLatencyMS: 1,
// Should be ignored
SessionCountSSH: 4,
SessionCountVSCode: 3,
})
dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
CreatedAt: insertTime.Add(-time.Minute),
AgentID: agentID,
SessionCountVSCode: 1,
Usage: true,
})
dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
CreatedAt: insertTime.Add(-time.Minute),
AgentID: agentID,
SessionCountReconnectingPTY: 1,
Usage: true,
})
// Latest stats
dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
CreatedAt: insertTime,
AgentID: agentID,
TxBytes: 1,
RxBytes: 1,
ConnectionMedianLatencyMS: 2,
// Should be ignored
SessionCountSSH: 3,
SessionCountVSCode: 1,
})
dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
CreatedAt: insertTime,
AgentID: agentID,
SessionCountVSCode: 1,
Usage: true,
})
dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
CreatedAt: insertTime,
AgentID: agentID,
SessionCountSSH: 1,
Usage: true,
})
stats, err := db.GetDeploymentWorkspaceAgentUsageStats(ctx, dbtime.Now().Add(-time.Hour))
require.NoError(t, err)
require.Equal(t, int64(2), stats.WorkspaceTxBytes)
require.Equal(t, int64(2), stats.WorkspaceRxBytes)
require.Equal(t, 1.5, stats.WorkspaceConnectionLatency50)
require.Equal(t, 1.95, stats.WorkspaceConnectionLatency95)
require.Equal(t, int64(1), stats.SessionCountVSCode)
require.Equal(t, int64(1), stats.SessionCountSSH)
require.Equal(t, int64(0), stats.SessionCountReconnectingPTY)
require.Equal(t, int64(0), stats.SessionCountJetBrains)
})
t.Run("NoUsage", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
authz := rbac.NewAuthorizer(prometheus.NewRegistry())
db = dbauthz.New(db, authz, slogtest.Make(t, &slogtest.Options{}), coderdtest.AccessControlStorePointer())
ctx := context.Background()
agentID := uuid.New()
// Since the queries exclude the current minute
insertTime := dbtime.Now().Add(-time.Minute)
dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
CreatedAt: insertTime,
AgentID: agentID,
TxBytes: 3,
RxBytes: 4,
ConnectionMedianLatencyMS: 2,
// Should be ignored
SessionCountSSH: 3,
SessionCountVSCode: 1,
})
stats, err := db.GetDeploymentWorkspaceAgentUsageStats(ctx, dbtime.Now().Add(-time.Hour))
require.NoError(t, err)
require.Equal(t, int64(3), stats.WorkspaceTxBytes)
require.Equal(t, int64(4), stats.WorkspaceRxBytes)
require.Equal(t, int64(0), stats.SessionCountVSCode)
require.Equal(t, int64(0), stats.SessionCountSSH)
require.Equal(t, int64(0), stats.SessionCountReconnectingPTY)
require.Equal(t, int64(0), stats.SessionCountJetBrains)
})
}
func TestGetWorkspaceAgentUsageStats(t *testing.T) {
t.Parallel()
t.Run("Aggregates", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
authz := rbac.NewAuthorizer(prometheus.NewRegistry())
db = dbauthz.New(db, authz, slogtest.Make(t, &slogtest.Options{}), coderdtest.AccessControlStorePointer())
ctx := context.Background()
// Since the queries exclude the current minute
insertTime := dbtime.Now().Add(-time.Minute)
agentID1 := uuid.New()
agentID2 := uuid.New()
workspaceID1 := uuid.New()
workspaceID2 := uuid.New()
templateID1 := uuid.New()
templateID2 := uuid.New()
userID1 := uuid.New()
userID2 := uuid.New()
// Old workspace 1 stats
dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
CreatedAt: insertTime.Add(-time.Minute),
AgentID: agentID1,
WorkspaceID: workspaceID1,
TemplateID: templateID1,
UserID: userID1,
TxBytes: 1,
RxBytes: 1,
ConnectionMedianLatencyMS: 1,
// Should be ignored
SessionCountVSCode: 3,
SessionCountSSH: 1,
})
dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
CreatedAt: insertTime.Add(-time.Minute),
AgentID: agentID1,
WorkspaceID: workspaceID1,
TemplateID: templateID1,
UserID: userID1,
SessionCountVSCode: 1,
Usage: true,
})
// Latest workspace 1 stats
dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
CreatedAt: insertTime,
AgentID: agentID1,
WorkspaceID: workspaceID1,
TemplateID: templateID1,
UserID: userID1,
TxBytes: 2,
RxBytes: 2,
ConnectionMedianLatencyMS: 1,
// Should be ignored
SessionCountVSCode: 3,
SessionCountSSH: 4,
})
dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
CreatedAt: insertTime,
AgentID: agentID1,
WorkspaceID: workspaceID1,
TemplateID: templateID1,
UserID: userID1,
SessionCountVSCode: 1,
Usage: true,
})
dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
CreatedAt: insertTime,
AgentID: agentID1,
WorkspaceID: workspaceID1,
TemplateID: templateID1,
UserID: userID1,
SessionCountJetBrains: 1,
Usage: true,
})
// Latest workspace 2 stats
dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
CreatedAt: insertTime,
AgentID: agentID2,
WorkspaceID: workspaceID2,
TemplateID: templateID2,
UserID: userID2,
TxBytes: 4,
RxBytes: 8,
ConnectionMedianLatencyMS: 1,
})
dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
CreatedAt: insertTime,
AgentID: agentID2,
WorkspaceID: workspaceID2,
TemplateID: templateID2,
UserID: userID2,
TxBytes: 2,
RxBytes: 3,
ConnectionMedianLatencyMS: 1,
// Should be ignored
SessionCountVSCode: 3,
SessionCountSSH: 4,
})
dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
CreatedAt: insertTime,
AgentID: agentID2,
WorkspaceID: workspaceID2,
TemplateID: templateID2,
UserID: userID2,
SessionCountSSH: 1,
Usage: true,
})
dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
CreatedAt: insertTime,
AgentID: agentID2,
WorkspaceID: workspaceID2,
TemplateID: templateID2,
UserID: userID2,
SessionCountJetBrains: 1,
Usage: true,
})
reqTime := dbtime.Now().Add(-time.Hour)
stats, err := db.GetWorkspaceAgentUsageStats(ctx, reqTime)
require.NoError(t, err)
ws1Stats, ws2Stats := stats[0], stats[1]
if ws1Stats.WorkspaceID != workspaceID1 {
ws1Stats, ws2Stats = ws2Stats, ws1Stats
}
require.Equal(t, int64(3), ws1Stats.WorkspaceTxBytes)
require.Equal(t, int64(3), ws1Stats.WorkspaceRxBytes)
require.Equal(t, int64(1), ws1Stats.SessionCountVSCode)
require.Equal(t, int64(1), ws1Stats.SessionCountJetBrains)
require.Equal(t, int64(0), ws1Stats.SessionCountSSH)
require.Equal(t, int64(0), ws1Stats.SessionCountReconnectingPTY)
require.Equal(t, int64(6), ws2Stats.WorkspaceTxBytes)
require.Equal(t, int64(11), ws2Stats.WorkspaceRxBytes)
require.Equal(t, int64(1), ws2Stats.SessionCountSSH)
require.Equal(t, int64(1), ws2Stats.SessionCountJetBrains)
require.Equal(t, int64(0), ws2Stats.SessionCountVSCode)
require.Equal(t, int64(0), ws2Stats.SessionCountReconnectingPTY)
})
t.Run("NoUsage", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
authz := rbac.NewAuthorizer(prometheus.NewRegistry())
db = dbauthz.New(db, authz, slogtest.Make(t, &slogtest.Options{}), coderdtest.AccessControlStorePointer())
ctx := context.Background()
// Since the queries exclude the current minute
insertTime := dbtime.Now().Add(-time.Minute)
agentID := uuid.New()
dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
CreatedAt: insertTime,
AgentID: agentID,
TxBytes: 3,
RxBytes: 4,
ConnectionMedianLatencyMS: 2,
// Should be ignored
SessionCountSSH: 3,
SessionCountVSCode: 1,
})
stats, err := db.GetWorkspaceAgentUsageStats(ctx, dbtime.Now().Add(-time.Hour))
require.NoError(t, err)
require.Len(t, stats, 1)
require.Equal(t, int64(3), stats[0].WorkspaceTxBytes)
require.Equal(t, int64(4), stats[0].WorkspaceRxBytes)
require.Equal(t, int64(0), stats[0].SessionCountVSCode)
require.Equal(t, int64(0), stats[0].SessionCountSSH)
require.Equal(t, int64(0), stats[0].SessionCountReconnectingPTY)
require.Equal(t, int64(0), stats[0].SessionCountJetBrains)
})
}
func TestGetWorkspaceAgentUsageStatsAndLabels(t *testing.T) {
t.Parallel()
t.Run("Aggregates", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
ctx := context.Background()
insertTime := dbtime.Now()
// Insert user, agent, template, workspace
user1 := dbgen.User(t, db, database.User{})
org := dbgen.Organization(t, db, database.Organization{})
job1 := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
OrganizationID: org.ID,
})
resource1 := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
JobID: job1.ID,
})
agent1 := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
ResourceID: resource1.ID,
})
template1 := dbgen.Template(t, db, database.Template{
OrganizationID: org.ID,
CreatedBy: user1.ID,
})
workspace1 := dbgen.Workspace(t, db, database.Workspace{
OwnerID: user1.ID,
OrganizationID: org.ID,
TemplateID: template1.ID,
})
user2 := dbgen.User(t, db, database.User{})
job2 := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
OrganizationID: org.ID,
})
resource2 := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
JobID: job2.ID,
})
agent2 := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
ResourceID: resource2.ID,
})
template2 := dbgen.Template(t, db, database.Template{
CreatedBy: user1.ID,
OrganizationID: org.ID,
})
workspace2 := dbgen.Workspace(t, db, database.Workspace{
OwnerID: user2.ID,
OrganizationID: org.ID,
TemplateID: template2.ID,
})
// Old workspace 1 stats
dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
CreatedAt: insertTime.Add(-time.Minute),
AgentID: agent1.ID,
WorkspaceID: workspace1.ID,
TemplateID: template1.ID,
UserID: user1.ID,
TxBytes: 1,
RxBytes: 1,
ConnectionMedianLatencyMS: 1,
// Should be ignored
SessionCountVSCode: 3,
SessionCountSSH: 1,
})
dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
CreatedAt: insertTime.Add(-time.Minute),
AgentID: agent1.ID,
WorkspaceID: workspace1.ID,
TemplateID: template1.ID,
UserID: user1.ID,
SessionCountVSCode: 1,
Usage: true,
})
// Latest workspace 1 stats
dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
CreatedAt: insertTime,
AgentID: agent1.ID,
WorkspaceID: workspace1.ID,
TemplateID: template1.ID,
UserID: user1.ID,
TxBytes: 2,
RxBytes: 2,
ConnectionMedianLatencyMS: 1,
// Should be ignored
SessionCountVSCode: 4,
SessionCountSSH: 3,
})
dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
CreatedAt: insertTime,
AgentID: agent1.ID,
WorkspaceID: workspace1.ID,
TemplateID: template1.ID,
UserID: user1.ID,
SessionCountJetBrains: 1,
Usage: true,
})
dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
CreatedAt: insertTime,
AgentID: agent1.ID,
WorkspaceID: workspace1.ID,
TemplateID: template1.ID,
UserID: user1.ID,
SessionCountReconnectingPTY: 1,
Usage: true,
})
// Latest workspace 2 stats
dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
CreatedAt: insertTime,
AgentID: agent2.ID,
WorkspaceID: workspace2.ID,
TemplateID: template2.ID,
UserID: user2.ID,
TxBytes: 4,
RxBytes: 8,
ConnectionMedianLatencyMS: 1,
})
dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
CreatedAt: insertTime,
AgentID: agent2.ID,
WorkspaceID: workspace2.ID,
TemplateID: template2.ID,
UserID: user2.ID,
SessionCountVSCode: 1,
Usage: true,
})
dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
CreatedAt: insertTime,
AgentID: agent2.ID,
WorkspaceID: workspace2.ID,
TemplateID: template2.ID,
UserID: user2.ID,
SessionCountSSH: 1,
Usage: true,
})
stats, err := db.GetWorkspaceAgentUsageStatsAndLabels(ctx, insertTime.Add(-time.Hour))
require.NoError(t, err)
require.Len(t, stats, 2)
require.Contains(t, stats, database.GetWorkspaceAgentUsageStatsAndLabelsRow{
Username: user1.Username,
AgentName: agent1.Name,
WorkspaceName: workspace1.Name,
TxBytes: 3,
RxBytes: 3,
SessionCountJetBrains: 1,
SessionCountReconnectingPTY: 1,
ConnectionMedianLatencyMS: 1,
})
require.Contains(t, stats, database.GetWorkspaceAgentUsageStatsAndLabelsRow{
Username: user2.Username,
AgentName: agent2.Name,
WorkspaceName: workspace2.Name,
RxBytes: 8,
TxBytes: 4,
SessionCountVSCode: 1,
SessionCountSSH: 1,
ConnectionMedianLatencyMS: 1,
})
})
t.Run("NoUsage", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
ctx := context.Background()
insertTime := dbtime.Now()
// Insert user, agent, template, workspace
user := dbgen.User(t, db, database.User{})
org := dbgen.Organization(t, db, database.Organization{})
job := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
OrganizationID: org.ID,
})
resource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
JobID: job.ID,
})
agent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
ResourceID: resource.ID,
})
template := dbgen.Template(t, db, database.Template{
OrganizationID: org.ID,
CreatedBy: user.ID,
})
workspace := dbgen.Workspace(t, db, database.Workspace{
OwnerID: user.ID,
OrganizationID: org.ID,
TemplateID: template.ID,
})
dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
CreatedAt: insertTime.Add(-time.Minute),
AgentID: agent.ID,
WorkspaceID: workspace.ID,
TemplateID: template.ID,
UserID: user.ID,
RxBytes: 4,
TxBytes: 5,
ConnectionMedianLatencyMS: 1,
// Should be ignored
SessionCountVSCode: 3,
SessionCountSSH: 1,
})
stats, err := db.GetWorkspaceAgentUsageStatsAndLabels(ctx, insertTime.Add(-time.Hour))
require.NoError(t, err)
require.Len(t, stats, 1)
require.Contains(t, stats, database.GetWorkspaceAgentUsageStatsAndLabelsRow{
Username: user.Username,
AgentName: agent.Name,
WorkspaceName: workspace.Name,
RxBytes: 4,
TxBytes: 5,
ConnectionMedianLatencyMS: 1,
})
})
}
func TestInsertWorkspaceAgentLogs(t *testing.T) {
t.Parallel()
if testing.Short() {
+328 -5
View File
@@ -12050,7 +12050,7 @@ WITH agent_stats AS (
coalesce(SUM(session_count_jetbrains), 0)::bigint AS session_count_jetbrains,
coalesce(SUM(session_count_reconnecting_pty), 0)::bigint AS session_count_reconnecting_pty
FROM (
SELECT id, created_at, user_id, agent_id, workspace_id, template_id, connections_by_proto, connection_count, rx_packets, rx_bytes, tx_packets, tx_bytes, connection_median_latency_ms, session_count_vscode, session_count_jetbrains, session_count_reconnecting_pty, session_count_ssh, ROW_NUMBER() OVER(PARTITION BY agent_id ORDER BY created_at DESC) AS rn
SELECT id, created_at, user_id, agent_id, workspace_id, template_id, connections_by_proto, connection_count, rx_packets, rx_bytes, tx_packets, tx_bytes, connection_median_latency_ms, session_count_vscode, session_count_jetbrains, session_count_reconnecting_pty, session_count_ssh, usage, ROW_NUMBER() OVER(PARTITION BY agent_id ORDER BY created_at DESC) AS rn
FROM workspace_agent_stats WHERE created_at > $1
) AS a WHERE a.rn = 1
)
@@ -12084,6 +12084,88 @@ func (q *sqlQuerier) GetDeploymentWorkspaceAgentStats(ctx context.Context, creat
return i, err
}
const getDeploymentWorkspaceAgentUsageStats = `-- name: GetDeploymentWorkspaceAgentUsageStats :one
WITH agent_stats AS (
SELECT
coalesce(SUM(rx_bytes), 0)::bigint AS workspace_rx_bytes,
coalesce(SUM(tx_bytes), 0)::bigint AS workspace_tx_bytes,
coalesce((PERCENTILE_CONT(0.5) WITHIN GROUP (ORDER BY connection_median_latency_ms)), -1)::FLOAT AS workspace_connection_latency_50,
coalesce((PERCENTILE_CONT(0.95) WITHIN GROUP (ORDER BY connection_median_latency_ms)), -1)::FLOAT AS workspace_connection_latency_95
FROM workspace_agent_stats
-- The greater than 0 is to support legacy agents that don't report connection_median_latency_ms.
WHERE workspace_agent_stats.created_at > $1 AND connection_median_latency_ms > 0
),
minute_buckets AS (
SELECT
agent_id,
date_trunc('minute', created_at) AS minute_bucket,
coalesce(SUM(session_count_vscode), 0)::bigint AS session_count_vscode,
coalesce(SUM(session_count_ssh), 0)::bigint AS session_count_ssh,
coalesce(SUM(session_count_jetbrains), 0)::bigint AS session_count_jetbrains,
coalesce(SUM(session_count_reconnecting_pty), 0)::bigint AS session_count_reconnecting_pty
FROM
workspace_agent_stats
WHERE
created_at >= $1
AND created_at < date_trunc('minute', now()) -- Exclude current partial minute
AND usage = true
GROUP BY
agent_id,
minute_bucket
),
latest_buckets AS (
SELECT DISTINCT ON (agent_id)
agent_id,
minute_bucket,
session_count_vscode,
session_count_jetbrains,
session_count_reconnecting_pty,
session_count_ssh
FROM
minute_buckets
ORDER BY
agent_id,
minute_bucket DESC
),
latest_agent_stats AS (
SELECT
coalesce(SUM(session_count_vscode), 0)::bigint AS session_count_vscode,
coalesce(SUM(session_count_ssh), 0)::bigint AS session_count_ssh,
coalesce(SUM(session_count_jetbrains), 0)::bigint AS session_count_jetbrains,
coalesce(SUM(session_count_reconnecting_pty), 0)::bigint AS session_count_reconnecting_pty
FROM
latest_buckets
)
SELECT workspace_rx_bytes, workspace_tx_bytes, workspace_connection_latency_50, workspace_connection_latency_95, session_count_vscode, session_count_ssh, session_count_jetbrains, session_count_reconnecting_pty FROM agent_stats, latest_agent_stats
`
type GetDeploymentWorkspaceAgentUsageStatsRow struct {
WorkspaceRxBytes int64 `db:"workspace_rx_bytes" json:"workspace_rx_bytes"`
WorkspaceTxBytes int64 `db:"workspace_tx_bytes" json:"workspace_tx_bytes"`
WorkspaceConnectionLatency50 float64 `db:"workspace_connection_latency_50" json:"workspace_connection_latency_50"`
WorkspaceConnectionLatency95 float64 `db:"workspace_connection_latency_95" json:"workspace_connection_latency_95"`
SessionCountVSCode int64 `db:"session_count_vscode" json:"session_count_vscode"`
SessionCountSSH int64 `db:"session_count_ssh" json:"session_count_ssh"`
SessionCountJetBrains int64 `db:"session_count_jetbrains" json:"session_count_jetbrains"`
SessionCountReconnectingPTY int64 `db:"session_count_reconnecting_pty" json:"session_count_reconnecting_pty"`
}
func (q *sqlQuerier) GetDeploymentWorkspaceAgentUsageStats(ctx context.Context, createdAt time.Time) (GetDeploymentWorkspaceAgentUsageStatsRow, error) {
row := q.db.QueryRowContext(ctx, getDeploymentWorkspaceAgentUsageStats, createdAt)
var i GetDeploymentWorkspaceAgentUsageStatsRow
err := row.Scan(
&i.WorkspaceRxBytes,
&i.WorkspaceTxBytes,
&i.WorkspaceConnectionLatency50,
&i.WorkspaceConnectionLatency95,
&i.SessionCountVSCode,
&i.SessionCountSSH,
&i.SessionCountJetBrains,
&i.SessionCountReconnectingPTY,
)
return i, err
}
const getTemplateDAUs = `-- name: GetTemplateDAUs :many
SELECT
(created_at at TIME ZONE cast($2::integer as text))::date as date,
@@ -12155,7 +12237,7 @@ WITH agent_stats AS (
coalesce(SUM(session_count_jetbrains), 0)::bigint AS session_count_jetbrains,
coalesce(SUM(session_count_reconnecting_pty), 0)::bigint AS session_count_reconnecting_pty
FROM (
SELECT id, created_at, user_id, agent_id, workspace_id, template_id, connections_by_proto, connection_count, rx_packets, rx_bytes, tx_packets, tx_bytes, connection_median_latency_ms, session_count_vscode, session_count_jetbrains, session_count_reconnecting_pty, session_count_ssh, ROW_NUMBER() OVER(PARTITION BY agent_id ORDER BY created_at DESC) AS rn
SELECT id, created_at, user_id, agent_id, workspace_id, template_id, connections_by_proto, connection_count, rx_packets, rx_bytes, tx_packets, tx_bytes, connection_median_latency_ms, session_count_vscode, session_count_jetbrains, session_count_reconnecting_pty, session_count_ssh, usage, ROW_NUMBER() OVER(PARTITION BY agent_id ORDER BY created_at DESC) AS rn
FROM workspace_agent_stats WHERE created_at > $1
) AS a WHERE a.rn = 1 GROUP BY a.user_id, a.agent_id, a.workspace_id, a.template_id
)
@@ -12238,7 +12320,7 @@ WITH agent_stats AS (
coalesce(SUM(connection_count), 0)::bigint AS connection_count,
coalesce(MAX(connection_median_latency_ms), 0)::float AS connection_median_latency_ms
FROM (
SELECT id, created_at, user_id, agent_id, workspace_id, template_id, connections_by_proto, connection_count, rx_packets, rx_bytes, tx_packets, tx_bytes, connection_median_latency_ms, session_count_vscode, session_count_jetbrains, session_count_reconnecting_pty, session_count_ssh, ROW_NUMBER() OVER(PARTITION BY agent_id ORDER BY created_at DESC) AS rn
SELECT id, created_at, user_id, agent_id, workspace_id, template_id, connections_by_proto, connection_count, rx_packets, rx_bytes, tx_packets, tx_bytes, connection_median_latency_ms, session_count_vscode, session_count_jetbrains, session_count_reconnecting_pty, session_count_ssh, usage, ROW_NUMBER() OVER(PARTITION BY agent_id ORDER BY created_at DESC) AS rn
FROM workspace_agent_stats
-- The greater than 0 is to support legacy agents that don't report connection_median_latency_ms.
WHERE created_at > $1 AND connection_median_latency_ms > 0
@@ -12319,6 +12401,243 @@ func (q *sqlQuerier) GetWorkspaceAgentStatsAndLabels(ctx context.Context, create
return items, nil
}
const getWorkspaceAgentUsageStats = `-- name: GetWorkspaceAgentUsageStats :many
WITH agent_stats AS (
SELECT
user_id,
agent_id,
workspace_id,
template_id,
MIN(created_at)::timestamptz AS aggregated_from,
coalesce(SUM(rx_bytes), 0)::bigint AS workspace_rx_bytes,
coalesce(SUM(tx_bytes), 0)::bigint AS workspace_tx_bytes,
coalesce((PERCENTILE_CONT(0.5) WITHIN GROUP (ORDER BY connection_median_latency_ms)), -1)::FLOAT AS workspace_connection_latency_50,
coalesce((PERCENTILE_CONT(0.95) WITHIN GROUP (ORDER BY connection_median_latency_ms)), -1)::FLOAT AS workspace_connection_latency_95
FROM workspace_agent_stats
-- The greater than 0 is to support legacy agents that don't report connection_median_latency_ms.
WHERE workspace_agent_stats.created_at > $1 AND connection_median_latency_ms > 0
GROUP BY user_id, agent_id, workspace_id, template_id
),
minute_buckets AS (
SELECT
agent_id,
date_trunc('minute', created_at) AS minute_bucket,
coalesce(SUM(session_count_vscode), 0)::bigint AS session_count_vscode,
coalesce(SUM(session_count_ssh), 0)::bigint AS session_count_ssh,
coalesce(SUM(session_count_jetbrains), 0)::bigint AS session_count_jetbrains,
coalesce(SUM(session_count_reconnecting_pty), 0)::bigint AS session_count_reconnecting_pty
FROM
workspace_agent_stats
WHERE
created_at >= $1
AND created_at < date_trunc('minute', now()) -- Exclude current partial minute
AND usage = true
GROUP BY
agent_id,
minute_bucket,
user_id,
agent_id,
workspace_id,
template_id
),
latest_buckets AS (
SELECT DISTINCT ON (agent_id)
agent_id,
session_count_vscode,
session_count_ssh,
session_count_jetbrains,
session_count_reconnecting_pty
FROM
minute_buckets
ORDER BY
agent_id,
minute_bucket DESC
)
SELECT user_id,
agent_stats.agent_id,
workspace_id,
template_id,
aggregated_from,
workspace_rx_bytes,
workspace_tx_bytes,
workspace_connection_latency_50,
workspace_connection_latency_95,
coalesce(latest_buckets.agent_id,agent_stats.agent_id) AS agent_id,
coalesce(session_count_vscode, 0)::bigint AS session_count_vscode,
coalesce(session_count_ssh, 0)::bigint AS session_count_ssh,
coalesce(session_count_jetbrains, 0)::bigint AS session_count_jetbrains,
coalesce(session_count_reconnecting_pty, 0)::bigint AS session_count_reconnecting_pty
FROM agent_stats LEFT JOIN latest_buckets ON agent_stats.agent_id = latest_buckets.agent_id
`
type GetWorkspaceAgentUsageStatsRow struct {
UserID uuid.UUID `db:"user_id" json:"user_id"`
AgentID uuid.UUID `db:"agent_id" json:"agent_id"`
WorkspaceID uuid.UUID `db:"workspace_id" json:"workspace_id"`
TemplateID uuid.UUID `db:"template_id" json:"template_id"`
AggregatedFrom time.Time `db:"aggregated_from" json:"aggregated_from"`
WorkspaceRxBytes int64 `db:"workspace_rx_bytes" json:"workspace_rx_bytes"`
WorkspaceTxBytes int64 `db:"workspace_tx_bytes" json:"workspace_tx_bytes"`
WorkspaceConnectionLatency50 float64 `db:"workspace_connection_latency_50" json:"workspace_connection_latency_50"`
WorkspaceConnectionLatency95 float64 `db:"workspace_connection_latency_95" json:"workspace_connection_latency_95"`
AgentID_2 uuid.UUID `db:"agent_id_2" json:"agent_id_2"`
SessionCountVSCode int64 `db:"session_count_vscode" json:"session_count_vscode"`
SessionCountSSH int64 `db:"session_count_ssh" json:"session_count_ssh"`
SessionCountJetBrains int64 `db:"session_count_jetbrains" json:"session_count_jetbrains"`
SessionCountReconnectingPTY int64 `db:"session_count_reconnecting_pty" json:"session_count_reconnecting_pty"`
}
// `minute_buckets` could return 0 rows if there are no usage stats since `created_at`.
func (q *sqlQuerier) GetWorkspaceAgentUsageStats(ctx context.Context, createdAt time.Time) ([]GetWorkspaceAgentUsageStatsRow, error) {
rows, err := q.db.QueryContext(ctx, getWorkspaceAgentUsageStats, createdAt)
if err != nil {
return nil, err
}
defer rows.Close()
var items []GetWorkspaceAgentUsageStatsRow
for rows.Next() {
var i GetWorkspaceAgentUsageStatsRow
if err := rows.Scan(
&i.UserID,
&i.AgentID,
&i.WorkspaceID,
&i.TemplateID,
&i.AggregatedFrom,
&i.WorkspaceRxBytes,
&i.WorkspaceTxBytes,
&i.WorkspaceConnectionLatency50,
&i.WorkspaceConnectionLatency95,
&i.AgentID_2,
&i.SessionCountVSCode,
&i.SessionCountSSH,
&i.SessionCountJetBrains,
&i.SessionCountReconnectingPTY,
); err != nil {
return nil, err
}
items = append(items, i)
}
if err := rows.Close(); err != nil {
return nil, err
}
if err := rows.Err(); err != nil {
return nil, err
}
return items, nil
}
const getWorkspaceAgentUsageStatsAndLabels = `-- name: GetWorkspaceAgentUsageStatsAndLabels :many
WITH agent_stats AS (
SELECT
user_id,
agent_id,
workspace_id,
coalesce(SUM(rx_bytes), 0)::bigint AS rx_bytes,
coalesce(SUM(tx_bytes), 0)::bigint AS tx_bytes
FROM workspace_agent_stats
WHERE workspace_agent_stats.created_at > $1
GROUP BY user_id, agent_id, workspace_id
), latest_agent_stats AS (
SELECT
agent_id,
coalesce(SUM(session_count_vscode), 0)::bigint AS session_count_vscode,
coalesce(SUM(session_count_ssh), 0)::bigint AS session_count_ssh,
coalesce(SUM(session_count_jetbrains), 0)::bigint AS session_count_jetbrains,
coalesce(SUM(session_count_reconnecting_pty), 0)::bigint AS session_count_reconnecting_pty,
coalesce(SUM(connection_count), 0)::bigint AS connection_count
FROM workspace_agent_stats
-- We only want the latest stats, but those stats might be
-- spread across multiple rows.
WHERE usage = true AND created_at > now() - '1 minute'::interval
GROUP BY user_id, agent_id, workspace_id
), latest_agent_latencies AS (
SELECT
agent_id,
coalesce(MAX(connection_median_latency_ms), 0)::float AS connection_median_latency_ms
FROM workspace_agent_stats
GROUP BY user_id, agent_id, workspace_id
)
SELECT
users.username, workspace_agents.name AS agent_name, workspaces.name AS workspace_name, rx_bytes, tx_bytes,
coalesce(session_count_vscode, 0)::bigint AS session_count_vscode,
coalesce(session_count_ssh, 0)::bigint AS session_count_ssh,
coalesce(session_count_jetbrains, 0)::bigint AS session_count_jetbrains,
coalesce(session_count_reconnecting_pty, 0)::bigint AS session_count_reconnecting_pty,
coalesce(connection_count, 0)::bigint AS connection_count,
connection_median_latency_ms
FROM
agent_stats
LEFT JOIN
latest_agent_stats
ON
agent_stats.agent_id = latest_agent_stats.agent_id
JOIN
latest_agent_latencies
ON
agent_stats.agent_id = latest_agent_latencies.agent_id
JOIN
users
ON
users.id = agent_stats.user_id
JOIN
workspace_agents
ON
workspace_agents.id = agent_stats.agent_id
JOIN
workspaces
ON
workspaces.id = agent_stats.workspace_id
`
type GetWorkspaceAgentUsageStatsAndLabelsRow struct {
Username string `db:"username" json:"username"`
AgentName string `db:"agent_name" json:"agent_name"`
WorkspaceName string `db:"workspace_name" json:"workspace_name"`
RxBytes int64 `db:"rx_bytes" json:"rx_bytes"`
TxBytes int64 `db:"tx_bytes" json:"tx_bytes"`
SessionCountVSCode int64 `db:"session_count_vscode" json:"session_count_vscode"`
SessionCountSSH int64 `db:"session_count_ssh" json:"session_count_ssh"`
SessionCountJetBrains int64 `db:"session_count_jetbrains" json:"session_count_jetbrains"`
SessionCountReconnectingPTY int64 `db:"session_count_reconnecting_pty" json:"session_count_reconnecting_pty"`
ConnectionCount int64 `db:"connection_count" json:"connection_count"`
ConnectionMedianLatencyMS float64 `db:"connection_median_latency_ms" json:"connection_median_latency_ms"`
}
func (q *sqlQuerier) GetWorkspaceAgentUsageStatsAndLabels(ctx context.Context, createdAt time.Time) ([]GetWorkspaceAgentUsageStatsAndLabelsRow, error) {
rows, err := q.db.QueryContext(ctx, getWorkspaceAgentUsageStatsAndLabels, createdAt)
if err != nil {
return nil, err
}
defer rows.Close()
var items []GetWorkspaceAgentUsageStatsAndLabelsRow
for rows.Next() {
var i GetWorkspaceAgentUsageStatsAndLabelsRow
if err := rows.Scan(
&i.Username,
&i.AgentName,
&i.WorkspaceName,
&i.RxBytes,
&i.TxBytes,
&i.SessionCountVSCode,
&i.SessionCountSSH,
&i.SessionCountJetBrains,
&i.SessionCountReconnectingPTY,
&i.ConnectionCount,
&i.ConnectionMedianLatencyMS,
); err != nil {
return nil, err
}
items = append(items, i)
}
if err := rows.Close(); err != nil {
return nil, err
}
if err := rows.Err(); err != nil {
return nil, err
}
return items, nil
}
const insertWorkspaceAgentStats = `-- name: InsertWorkspaceAgentStats :exec
INSERT INTO
workspace_agent_stats (
@@ -12338,7 +12657,8 @@ INSERT INTO
session_count_jetbrains,
session_count_reconnecting_pty,
session_count_ssh,
connection_median_latency_ms
connection_median_latency_ms,
usage
)
SELECT
unnest($1 :: uuid[]) AS id,
@@ -12357,7 +12677,8 @@ SELECT
unnest($14 :: bigint[]) AS session_count_jetbrains,
unnest($15 :: bigint[]) AS session_count_reconnecting_pty,
unnest($16 :: bigint[]) AS session_count_ssh,
unnest($17 :: double precision[]) AS connection_median_latency_ms
unnest($17 :: double precision[]) AS connection_median_latency_ms,
unnest($18 :: boolean[]) AS usage
`
type InsertWorkspaceAgentStatsParams struct {
@@ -12378,6 +12699,7 @@ type InsertWorkspaceAgentStatsParams struct {
SessionCountReconnectingPTY []int64 `db:"session_count_reconnecting_pty" json:"session_count_reconnecting_pty"`
SessionCountSSH []int64 `db:"session_count_ssh" json:"session_count_ssh"`
ConnectionMedianLatencyMS []float64 `db:"connection_median_latency_ms" json:"connection_median_latency_ms"`
Usage []bool `db:"usage" json:"usage"`
}
func (q *sqlQuerier) InsertWorkspaceAgentStats(ctx context.Context, arg InsertWorkspaceAgentStatsParams) error {
@@ -12399,6 +12721,7 @@ func (q *sqlQuerier) InsertWorkspaceAgentStats(ctx context.Context, arg InsertWo
pq.Array(arg.SessionCountReconnectingPTY),
pq.Array(arg.SessionCountSSH),
pq.Array(arg.ConnectionMedianLatencyMS),
pq.Array(arg.Usage),
)
return err
}
+189 -2
View File
@@ -17,7 +17,8 @@ INSERT INTO
session_count_jetbrains,
session_count_reconnecting_pty,
session_count_ssh,
connection_median_latency_ms
connection_median_latency_ms,
usage
)
SELECT
unnest(@id :: uuid[]) AS id,
@@ -36,7 +37,8 @@ SELECT
unnest(@session_count_jetbrains :: bigint[]) AS session_count_jetbrains,
unnest(@session_count_reconnecting_pty :: bigint[]) AS session_count_reconnecting_pty,
unnest(@session_count_ssh :: bigint[]) AS session_count_ssh,
unnest(@connection_median_latency_ms :: double precision[]) AS connection_median_latency_ms;
unnest(@connection_median_latency_ms :: double precision[]) AS connection_median_latency_ms,
unnest(@usage :: boolean[]) AS usage;
-- name: GetTemplateDAUs :many
SELECT
@@ -119,6 +121,60 @@ WITH agent_stats AS (
)
SELECT * FROM agent_stats, latest_agent_stats;
-- name: GetDeploymentWorkspaceAgentUsageStats :one
WITH agent_stats AS (
SELECT
coalesce(SUM(rx_bytes), 0)::bigint AS workspace_rx_bytes,
coalesce(SUM(tx_bytes), 0)::bigint AS workspace_tx_bytes,
coalesce((PERCENTILE_CONT(0.5) WITHIN GROUP (ORDER BY connection_median_latency_ms)), -1)::FLOAT AS workspace_connection_latency_50,
coalesce((PERCENTILE_CONT(0.95) WITHIN GROUP (ORDER BY connection_median_latency_ms)), -1)::FLOAT AS workspace_connection_latency_95
FROM workspace_agent_stats
-- The greater than 0 is to support legacy agents that don't report connection_median_latency_ms.
WHERE workspace_agent_stats.created_at > $1 AND connection_median_latency_ms > 0
),
minute_buckets AS (
SELECT
agent_id,
date_trunc('minute', created_at) AS minute_bucket,
coalesce(SUM(session_count_vscode), 0)::bigint AS session_count_vscode,
coalesce(SUM(session_count_ssh), 0)::bigint AS session_count_ssh,
coalesce(SUM(session_count_jetbrains), 0)::bigint AS session_count_jetbrains,
coalesce(SUM(session_count_reconnecting_pty), 0)::bigint AS session_count_reconnecting_pty
FROM
workspace_agent_stats
WHERE
created_at >= $1
AND created_at < date_trunc('minute', now()) -- Exclude current partial minute
AND usage = true
GROUP BY
agent_id,
minute_bucket
),
latest_buckets AS (
SELECT DISTINCT ON (agent_id)
agent_id,
minute_bucket,
session_count_vscode,
session_count_jetbrains,
session_count_reconnecting_pty,
session_count_ssh
FROM
minute_buckets
ORDER BY
agent_id,
minute_bucket DESC
),
latest_agent_stats AS (
SELECT
coalesce(SUM(session_count_vscode), 0)::bigint AS session_count_vscode,
coalesce(SUM(session_count_ssh), 0)::bigint AS session_count_ssh,
coalesce(SUM(session_count_jetbrains), 0)::bigint AS session_count_jetbrains,
coalesce(SUM(session_count_reconnecting_pty), 0)::bigint AS session_count_reconnecting_pty
FROM
latest_buckets
)
SELECT * FROM agent_stats, latest_agent_stats;
-- name: GetWorkspaceAgentStats :many
WITH agent_stats AS (
SELECT
@@ -148,6 +204,75 @@ WITH agent_stats AS (
)
SELECT * FROM agent_stats JOIN latest_agent_stats ON agent_stats.agent_id = latest_agent_stats.agent_id;
-- name: GetWorkspaceAgentUsageStats :many
WITH agent_stats AS (
SELECT
user_id,
agent_id,
workspace_id,
template_id,
MIN(created_at)::timestamptz AS aggregated_from,
coalesce(SUM(rx_bytes), 0)::bigint AS workspace_rx_bytes,
coalesce(SUM(tx_bytes), 0)::bigint AS workspace_tx_bytes,
coalesce((PERCENTILE_CONT(0.5) WITHIN GROUP (ORDER BY connection_median_latency_ms)), -1)::FLOAT AS workspace_connection_latency_50,
coalesce((PERCENTILE_CONT(0.95) WITHIN GROUP (ORDER BY connection_median_latency_ms)), -1)::FLOAT AS workspace_connection_latency_95
FROM workspace_agent_stats
-- The greater than 0 is to support legacy agents that don't report connection_median_latency_ms.
WHERE workspace_agent_stats.created_at > $1 AND connection_median_latency_ms > 0
GROUP BY user_id, agent_id, workspace_id, template_id
),
minute_buckets AS (
SELECT
agent_id,
date_trunc('minute', created_at) AS minute_bucket,
coalesce(SUM(session_count_vscode), 0)::bigint AS session_count_vscode,
coalesce(SUM(session_count_ssh), 0)::bigint AS session_count_ssh,
coalesce(SUM(session_count_jetbrains), 0)::bigint AS session_count_jetbrains,
coalesce(SUM(session_count_reconnecting_pty), 0)::bigint AS session_count_reconnecting_pty
FROM
workspace_agent_stats
WHERE
created_at >= $1
AND created_at < date_trunc('minute', now()) -- Exclude current partial minute
AND usage = true
GROUP BY
agent_id,
minute_bucket,
user_id,
agent_id,
workspace_id,
template_id
),
latest_buckets AS (
SELECT DISTINCT ON (agent_id)
agent_id,
session_count_vscode,
session_count_ssh,
session_count_jetbrains,
session_count_reconnecting_pty
FROM
minute_buckets
ORDER BY
agent_id,
minute_bucket DESC
)
SELECT user_id,
agent_stats.agent_id,
workspace_id,
template_id,
aggregated_from,
workspace_rx_bytes,
workspace_tx_bytes,
workspace_connection_latency_50,
workspace_connection_latency_95,
-- `minute_buckets` could return 0 rows if there are no usage stats since `created_at`.
coalesce(latest_buckets.agent_id,agent_stats.agent_id) AS agent_id,
coalesce(session_count_vscode, 0)::bigint AS session_count_vscode,
coalesce(session_count_ssh, 0)::bigint AS session_count_ssh,
coalesce(session_count_jetbrains, 0)::bigint AS session_count_jetbrains,
coalesce(session_count_reconnecting_pty, 0)::bigint AS session_count_reconnecting_pty
FROM agent_stats LEFT JOIN latest_buckets ON agent_stats.agent_id = latest_buckets.agent_id;
-- name: GetWorkspaceAgentStatsAndLabels :many
WITH agent_stats AS (
SELECT
@@ -199,3 +324,65 @@ JOIN
workspaces
ON
workspaces.id = agent_stats.workspace_id;
-- name: GetWorkspaceAgentUsageStatsAndLabels :many
WITH agent_stats AS (
SELECT
user_id,
agent_id,
workspace_id,
coalesce(SUM(rx_bytes), 0)::bigint AS rx_bytes,
coalesce(SUM(tx_bytes), 0)::bigint AS tx_bytes
FROM workspace_agent_stats
WHERE workspace_agent_stats.created_at > $1
GROUP BY user_id, agent_id, workspace_id
), latest_agent_stats AS (
SELECT
agent_id,
coalesce(SUM(session_count_vscode), 0)::bigint AS session_count_vscode,
coalesce(SUM(session_count_ssh), 0)::bigint AS session_count_ssh,
coalesce(SUM(session_count_jetbrains), 0)::bigint AS session_count_jetbrains,
coalesce(SUM(session_count_reconnecting_pty), 0)::bigint AS session_count_reconnecting_pty,
coalesce(SUM(connection_count), 0)::bigint AS connection_count
FROM workspace_agent_stats
-- We only want the latest stats, but those stats might be
-- spread across multiple rows.
WHERE usage = true AND created_at > now() - '1 minute'::interval
GROUP BY user_id, agent_id, workspace_id
), latest_agent_latencies AS (
SELECT
agent_id,
coalesce(MAX(connection_median_latency_ms), 0)::float AS connection_median_latency_ms
FROM workspace_agent_stats
GROUP BY user_id, agent_id, workspace_id
)
SELECT
users.username, workspace_agents.name AS agent_name, workspaces.name AS workspace_name, rx_bytes, tx_bytes,
coalesce(session_count_vscode, 0)::bigint AS session_count_vscode,
coalesce(session_count_ssh, 0)::bigint AS session_count_ssh,
coalesce(session_count_jetbrains, 0)::bigint AS session_count_jetbrains,
coalesce(session_count_reconnecting_pty, 0)::bigint AS session_count_reconnecting_pty,
coalesce(connection_count, 0)::bigint AS connection_count,
connection_median_latency_ms
FROM
agent_stats
LEFT JOIN
latest_agent_stats
ON
agent_stats.agent_id = latest_agent_stats.agent_id
JOIN
latest_agent_latencies
ON
agent_stats.agent_id = latest_agent_latencies.agent_id
JOIN
users
ON
users.id = agent_stats.user_id
JOIN
workspace_agents
ON
workspace_agents.id = agent_stats.agent_id
JOIN
workspaces
ON
workspaces.id = agent_stats.workspace_id;