chore: replace GetManagedAgentCount query with aggregate table (#19636)

- Removes GetManagedAgentCount query
- Adds new table `usage_events_daily` which stores aggregated usage
events by the type and UTC day
- Adds trigger to update the values in this table when a new row is
inserted into `usage_events`
- Adds a migration that adds `usage_events_daily` rows for existing data
in `usage_events`
- Adds tests for the trigger
- Adds tests for the backfill query in the migration

Since the `usage_events` table is unreleased currently, this migration
will do nothing on real deployments and will only affect preview
deployments such as dogfood.

Closes https://github.com/coder/internal/issues/943

coderd/database/queries.sql.go

@@ -4334,44 +4334,6 @@ func (q *sqlQuerier) GetLicenses(ctx context.Context) ([]License, error) {
 	return items, nil
 }
 
-const getManagedAgentCount = `-- name: GetManagedAgentCount :one
-SELECT
-	COUNT(DISTINCT wb.id) AS count
-FROM
-	workspace_builds AS wb
-JOIN
-	provisioner_jobs AS pj
-ON
-	wb.job_id = pj.id
-WHERE
-	wb.transition = 'start'::workspace_transition
-	AND wb.has_ai_task = true
-	-- Only count jobs that are pending, running or succeeded. Other statuses
-	-- like cancel(ed|ing), failed or unknown are not considered as managed
-	-- agent usage. These workspace builds are typically unusable anyway.
-	AND pj.job_status IN (
-		'pending'::provisioner_job_status,
-		'running'::provisioner_job_status,
-		'succeeded'::provisioner_job_status
-	)
-	-- Jobs are counted at the time they are created, not when they are
-	-- completed, as pending jobs haven't completed yet.
-	AND wb.created_at BETWEEN $1::timestamptz AND $2::timestamptz
-`
-
-type GetManagedAgentCountParams struct {
-	StartTime time.Time `db:"start_time" json:"start_time"`
-	EndTime   time.Time `db:"end_time" json:"end_time"`
-}
-
-// This isn't strictly a license query, but it's related to license enforcement.
-func (q *sqlQuerier) GetManagedAgentCount(ctx context.Context, arg GetManagedAgentCountParams) (int64, error) {
-	row := q.db.QueryRowContext(ctx, getManagedAgentCount, arg.StartTime, arg.EndTime)
-	var count int64
-	err := row.Scan(&count)
-	return count, err
-}
-
 const getUnexpiredLicenses = `-- name: GetUnexpiredLicenses :many
 SELECT id, uploaded_at, jwt, exp, uuid
 FROM licenses
@@ -13634,6 +13596,40 @@ func (q *sqlQuerier) DisableForeignKeysAndTriggers(ctx context.Context) error {
 	return err
 }
 
+const getTotalUsageDCManagedAgentsV1 = `-- name: GetTotalUsageDCManagedAgentsV1 :one
+SELECT
+    -- The first cast is necessary since you can't sum strings, and the second
+    -- cast is necessary to make sqlc happy.
+    COALESCE(SUM((usage_data->>'count')::bigint), 0)::bigint AS total_count
+FROM
+    usage_events_daily
+WHERE
+    event_type = 'dc_managed_agents_v1'
+    -- Parentheses are necessary to avoid sqlc from generating an extra
+    -- argument.
+    AND day BETWEEN date_trunc('day', ($1::timestamptz) AT TIME ZONE 'UTC')::date AND date_trunc('day', ($2::timestamptz) AT TIME ZONE 'UTC')::date
+`
+
+type GetTotalUsageDCManagedAgentsV1Params struct {
+	StartDate time.Time `db:"start_date" json:"start_date"`
+	EndDate   time.Time `db:"end_date" json:"end_date"`
+}
+
+// Gets the total number of managed agents created between two dates. Uses the
+// aggregate table to avoid large scans or a complex index on the usage_events
+// table.
+//
+// This has the trade off that we can't count accurately between two exact
+// timestamps. The provided timestamps will be converted to UTC and truncated to
+// the events that happened on and between the two dates. Both dates are
+// inclusive.
+func (q *sqlQuerier) GetTotalUsageDCManagedAgentsV1(ctx context.Context, arg GetTotalUsageDCManagedAgentsV1Params) (int64, error) {
+	row := q.db.QueryRowContext(ctx, getTotalUsageDCManagedAgentsV1, arg.StartDate, arg.EndDate)
+	var total_count int64
+	err := row.Scan(&total_count)
+	return total_count, err
+}
+
 const insertUsageEvent = `-- name: InsertUsageEvent :exec
 INSERT INTO
     usage_events (
@@ -13693,7 +13689,7 @@ WITH usage_events AS (
                     -- than an hour ago. This is so we can retry publishing
                     -- events where the replica exited or couldn't update the
                     -- row.
-                    -- The parenthesis around @now::timestamptz are necessary to
+                    -- The parentheses around @now::timestamptz are necessary to
                     -- avoid sqlc from generating an extra argument.
                     OR potential_event.publish_started_at < ($1::timestamptz) - INTERVAL '1 hour'
                 )
@@ -13701,7 +13697,7 @@ WITH usage_events AS (
                 -- always permanently reject these events anyways. This is to
                 -- avoid duplicate events being billed to customers, as
                 -- Metronome will only deduplicate events within 34 days.
-                -- Also, the same parenthesis thing here as above.
+                -- Also, the same parentheses thing here as above.
                 AND potential_event.created_at > ($1::timestamptz) - INTERVAL '30 days'
             ORDER BY potential_event.created_at ASC
             FOR UPDATE SKIP LOCKED