coder · Emyrk · Apr 9, 2024 · Apr 4, 2024 · Apr 4, 2024 · Apr 4, 2024
diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go
@@ -2527,20 +2527,6 @@ func (q *querier) InsertWorkspaceAgentScripts(ctx context.Context, arg database.
 	return q.db.InsertWorkspaceAgentScripts(ctx, arg)
 }
 
-func (q *querier) InsertWorkspaceAgentStat(ctx context.Context, arg database.InsertWorkspaceAgentStatParams) (database.WorkspaceAgentStat, error) {
-	// TODO: This is a workspace agent operation. Should users be able to query this?
-	// Not really sure what this is for.
-	workspace, err := q.db.GetWorkspaceByID(ctx, arg.WorkspaceID)
-	if err != nil {
-		return database.WorkspaceAgentStat{}, err
-	}
-	err = q.authorizeContext(ctx, rbac.ActionUpdate, workspace)
-	if err != nil {
-		return database.WorkspaceAgentStat{}, err
-	}
-	return q.db.InsertWorkspaceAgentStat(ctx, arg)
-}
-
 func (q *querier) InsertWorkspaceAgentStats(ctx context.Context, arg database.InsertWorkspaceAgentStatsParams) error {
 	if err := q.authorizeContext(ctx, rbac.ActionCreate, rbac.ResourceSystem); err != nil {
 		return err

diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go
@@ -1512,12 +1512,6 @@ func (s *MethodTestSuite) TestWorkspace() {
 			AutomaticUpdates: database.AutomaticUpdatesAlways,
 		}).Asserts(w, rbac.ActionUpdate)
 	}))
-	s.Run("InsertWorkspaceAgentStat", s.Subtest(func(db database.Store, check *expects) {
-		ws := dbgen.Workspace(s.T(), db, database.Workspace{})
-		check.Args(database.InsertWorkspaceAgentStatParams{
-			WorkspaceID: ws.ID,
-		}).Asserts(ws, rbac.ActionUpdate)
-	}))
 	s.Run("UpdateWorkspaceAppHealthByID", s.Subtest(func(db database.Store, check *expects) {
 		ws := dbgen.Workspace(s.T(), db, database.Workspace{})
 		build := dbgen.WorkspaceBuild(s.T(), db, database.WorkspaceBuild{WorkspaceID: ws.ID, JobID: uuid.New()})

diff --git a/coderd/database/dbgen/dbgen.go b/coderd/database/dbgen/dbgen.go
@@ -707,27 +707,49 @@ func WorkspaceAgentStat(t testing.TB, db database.Store, orig database.Workspace
 	if orig.ConnectionsByProto == nil {
 		orig.ConnectionsByProto = json.RawMessage([]byte("{}"))
 	}
-	scheme, err := db.InsertWorkspaceAgentStat(genCtx, database.InsertWorkspaceAgentStatParams{
-		ID:                          takeFirst(orig.ID, uuid.New()),
-		CreatedAt:                   takeFirst(orig.CreatedAt, dbtime.Now()),
-		UserID:                      takeFirst(orig.UserID, uuid.New()),
-		TemplateID:                  takeFirst(orig.TemplateID, uuid.New()),
-		WorkspaceID:                 takeFirst(orig.WorkspaceID, uuid.New()),
-		AgentID:                     takeFirst(orig.AgentID, uuid.New()),
-		ConnectionsByProto:          orig.ConnectionsByProto,
-		ConnectionCount:             takeFirst(orig.ConnectionCount, 0),
-		RxPackets:                   takeFirst(orig.RxPackets, 0),
-		RxBytes:                     takeFirst(orig.RxBytes, 0),
-		TxPackets:                   takeFirst(orig.TxPackets, 0),
-		TxBytes:                     takeFirst(orig.TxBytes, 0),
-		SessionCountVSCode:          takeFirst(orig.SessionCountVSCode, 0),
-		SessionCountJetBrains:       takeFirst(orig.SessionCountJetBrains, 0),
-		SessionCountReconnectingPTY: takeFirst(orig.SessionCountReconnectingPTY, 0),
-		SessionCountSSH:             takeFirst(orig.SessionCountSSH, 0),
-		ConnectionMedianLatencyMS:   takeFirst(orig.ConnectionMedianLatencyMS, 0),
-	})
+	jsonProto := []byte(fmt.Sprintf("[%s]", orig.ConnectionsByProto))
+
+	params := database.InsertWorkspaceAgentStatsParams{
+		ID:                          []uuid.UUID{takeFirst(orig.ID, uuid.New())},
+		CreatedAt:                   []time.Time{takeFirst(orig.CreatedAt, dbtime.Now())},
+		UserID:                      []uuid.UUID{takeFirst(orig.UserID, uuid.New())},
+		TemplateID:                  []uuid.UUID{takeFirst(orig.TemplateID, uuid.New())},
+		WorkspaceID:                 []uuid.UUID{takeFirst(orig.WorkspaceID, uuid.New())},
+		AgentID:                     []uuid.UUID{takeFirst(orig.AgentID, uuid.New())},
+		ConnectionsByProto:          jsonProto,
+		ConnectionCount:             []int64{takeFirst(orig.ConnectionCount, 0)},
+		RxPackets:                   []int64{takeFirst(orig.RxPackets, 0)},
+		RxBytes:                     []int64{takeFirst(orig.RxBytes, 0)},
+		TxPackets:                   []int64{takeFirst(orig.TxPackets, 0)},
+		TxBytes:                     []int64{takeFirst(orig.TxBytes, 0)},
+		SessionCountVSCode:          []int64{takeFirst(orig.SessionCountVSCode, 0)},
+		SessionCountJetBrains:       []int64{takeFirst(orig.SessionCountJetBrains, 0)},
+		SessionCountReconnectingPTY: []int64{takeFirst(orig.SessionCountReconnectingPTY, 0)},
+		SessionCountSSH:             []int64{takeFirst(orig.SessionCountSSH, 0)},
+		ConnectionMedianLatencyMS:   []float64{takeFirst(orig.ConnectionMedianLatencyMS, 0)},
+	}
+	err := db.InsertWorkspaceAgentStats(genCtx, params)
 	require.NoError(t, err, "insert workspace agent stat")
-	return scheme
+
+	return database.WorkspaceAgentStat{
+		ID:                          params.ID[0],
+		CreatedAt:                   params.CreatedAt[0],
+		UserID:                      params.UserID[0],
+		AgentID:                     params.AgentID[0],
+		WorkspaceID:                 params.WorkspaceID[0],
+		TemplateID:                  params.TemplateID[0],
+		ConnectionsByProto:          orig.ConnectionsByProto,
+		ConnectionCount:             params.ConnectionCount[0],
+		RxPackets:                   params.RxPackets[0],
+		RxBytes:                     params.RxBytes[0],
+		TxPackets:                   params.TxPackets[0],
+		TxBytes:                     params.TxBytes[0],
+		ConnectionMedianLatencyMS:   params.ConnectionMedianLatencyMS[0],
+		SessionCountVSCode:          params.SessionCountVSCode[0],
+		SessionCountJetBrains:       params.SessionCountJetBrains[0],
+		SessionCountReconnectingPTY: params.SessionCountReconnectingPTY[0],
+		SessionCountSSH:             params.SessionCountSSH[0],
+	}
 }
 
 func OAuth2ProviderApp(t testing.TB, db database.Store, seed database.OAuth2ProviderApp) database.OAuth2ProviderApp {

diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go
@@ -6457,37 +6457,6 @@ func (q *FakeQuerier) InsertWorkspaceAgentScripts(_ context.Context, arg databas
 	return scripts, nil
 }
 
-func (q *FakeQuerier) InsertWorkspaceAgentStat(_ context.Context, p database.InsertWorkspaceAgentStatParams) (database.WorkspaceAgentStat, error) {
-	if err := validateDatabaseType(p); err != nil {
-		return database.WorkspaceAgentStat{}, err
-	}
-
-	q.mutex.Lock()
-	defer q.mutex.Unlock()
-
-	stat := database.WorkspaceAgentStat{
-		ID:                          p.ID,
-		CreatedAt:                   p.CreatedAt,
-		WorkspaceID:                 p.WorkspaceID,
-		AgentID:                     p.AgentID,
-		UserID:                      p.UserID,
-		ConnectionsByProto:          p.ConnectionsByProto,
-		ConnectionCount:             p.ConnectionCount,
-		RxPackets:                   p.RxPackets,
-		RxBytes:                     p.RxBytes,
-		TxPackets:                   p.TxPackets,
-		TxBytes:                     p.TxBytes,
-		TemplateID:                  p.TemplateID,
-		SessionCountVSCode:          p.SessionCountVSCode,
-		SessionCountJetBrains:       p.SessionCountJetBrains,
-		SessionCountReconnectingPTY: p.SessionCountReconnectingPTY,
-		SessionCountSSH:             p.SessionCountSSH,
-		ConnectionMedianLatencyMS:   p.ConnectionMedianLatencyMS,
-	}
-	q.workspaceAgentStats = append(q.workspaceAgentStats, stat)
-	return stat, nil
-}
-
 func (q *FakeQuerier) InsertWorkspaceAgentStats(_ context.Context, arg database.InsertWorkspaceAgentStatsParams) error {
 	err := validateDatabaseType(arg)
 	if err != nil {

diff --git a/coderd/database/dbmetrics/dbmetrics.go b/coderd/database/dbmetrics/dbmetrics.go
diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go
diff --git a/coderd/database/dbrollup/dbrollup_test.go b/coderd/database/dbrollup/dbrollup_test.go
@@ -143,8 +143,8 @@ func TestRollupTemplateUsageStats(t *testing.T) {
 	db, ps := dbtestutil.NewDB(t, dbtestutil.WithDumpOnFailure())
 	logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
 
-	anHourAgo := dbtime.Now().Add(-time.Hour).Truncate(time.Hour)
-	anHourAndSixMonthsAgo := anHourAgo.AddDate(0, -6, 0)
+	anHourAgo := dbtime.Now().Add(-time.Hour).Truncate(time.Hour).UTC()
+	anHourAndSixMonthsAgo := anHourAgo.AddDate(0, -6, 0).UTC()
 
 	var (
 		org   = dbgen.Organization(t, db, database.Organization{})
@@ -242,6 +242,12 @@ func TestRollupTemplateUsageStats(t *testing.T) {
 	require.NoError(t, err)
 	require.Len(t, stats, 1)
 
+	// I do not know a better way to do this. Our database runs in a *random*
+	// timezone. So the returned time is in a random timezone and fails on the
+	// equal even though they are the same time if converted back to the same timezone.
+	stats[0].EndTime = stats[0].EndTime.UTC()
+	stats[0].StartTime = stats[0].StartTime.UTC()
+
 	require.Equal(t, database.TemplateUsageStat{
 		TemplateID:          tpl.ID,
 		UserID:              user.ID,

diff --git a/coderd/database/querier.go b/coderd/database/querier.go
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/workspaceagentstats.sql b/coderd/database/queries/workspaceagentstats.sql
@@ -1,27 +1,3 @@
--- name: InsertWorkspaceAgentStat :one
-INSERT INTO
-	workspace_agent_stats (
-		id,
-		created_at,
-		user_id,
-		workspace_id,
-		template_id,
-		agent_id,
-		connections_by_proto,
-		connection_count,
-		rx_packets,
-		rx_bytes,
-		tx_packets,
-		tx_bytes,
-		session_count_vscode,
-		session_count_jetbrains,
-		session_count_reconnecting_pty,
-		session_count_ssh,
-		connection_median_latency_ms
-	)
-VALUES
-	($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17) RETURNING *;
-
 -- name: InsertWorkspaceAgentStats :exec
 INSERT INTO
 	workspace_agent_stats (

diff --git a/coderd/metricscache/metricscache_test.go b/coderd/metricscache/metricscache_test.go
@@ -3,6 +3,7 @@ package metricscache_test
 import (
 	"context"
 	"database/sql"
+	"encoding/json"
 	"testing"
 	"time"
 
@@ -280,14 +281,25 @@ func TestCache_DeploymentStats(t *testing.T) {
 	})
 	defer cache.Close()
 
-	_, err := db.InsertWorkspaceAgentStat(context.Background(), database.InsertWorkspaceAgentStatParams{
-		ID:                 uuid.New(),
-		AgentID:            uuid.New(),
-		CreatedAt:          dbtime.Now(),
-		ConnectionCount:    1,
-		RxBytes:            1,
-		TxBytes:            1,
-		SessionCountVSCode: 1,
+	err := db.InsertWorkspaceAgentStats(context.Background(), database.InsertWorkspaceAgentStatsParams{
+		ID:                 []uuid.UUID{uuid.New()},
+		CreatedAt:          []time.Time{dbtime.Now()},
+		WorkspaceID:        []uuid.UUID{uuid.New()},
+		UserID:             []uuid.UUID{uuid.New()},
+		TemplateID:         []uuid.UUID{uuid.New()},
+		AgentID:            []uuid.UUID{uuid.New()},
+		ConnectionsByProto: json.RawMessage(`[{}]`),
+
+		RxPackets:                   []int64{0},
+		RxBytes:                     []int64{1},
+		TxPackets:                   []int64{0},
+		TxBytes:                     []int64{1},
+		ConnectionCount:             []int64{1},
+		SessionCountVSCode:          []int64{1},
+		SessionCountJetBrains:       []int64{0},
+		SessionCountReconnectingPTY: []int64{0},
+		SessionCountSSH:             []int64{0},
+		ConnectionMedianLatencyMS:   []float64{10},
 	})
 	require.NoError(t, err)
 

diff --git a/coderd/prometheusmetrics/prometheusmetrics_test.go b/coderd/prometheusmetrics/prometheusmetrics_test.go
@@ -11,7 +11,6 @@ import (
 	"testing"
 	"time"
 
-	"github.com/coder/coder/v2/cryptorand"
 	"github.com/google/uuid"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/stretchr/testify/assert"
@@ -32,6 +31,7 @@ import (
 	"github.com/coder/coder/v2/coderd/prometheusmetrics"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/agentsdk"
+	"github.com/coder/coder/v2/cryptorand"
 	"github.com/coder/coder/v2/provisioner/echo"
 	"github.com/coder/coder/v2/provisionersdk/proto"
 	"github.com/coder/coder/v2/tailnet"