mirror of
https://github.com/coder/coder.git
synced 2026-06-03 04:58:23 +00:00
84de391f26
AI seat tracking inserted as heartbeat into usage table.
32 lines
812 B
Go
32 lines
812 B
Go
package usage
|
|
|
|
import (
|
|
"context"
|
|
"time"
|
|
|
|
"golang.org/x/xerrors"
|
|
|
|
"github.com/coder/coder/v2/coderd/database"
|
|
"github.com/coder/coder/v2/coderd/database/dbauthz"
|
|
"github.com/coder/coder/v2/coderd/usage/usagetypes"
|
|
)
|
|
|
|
const (
|
|
AISeatsInterval = 4 * time.Hour
|
|
)
|
|
|
|
// AISeatsHeartbeat returns a HeartbeatFunc that queries the active
|
|
// AI seat count and emits it as an HBAISeats heartbeat event.
|
|
func AISeatsHeartbeat(db database.Store) HeartbeatFunc {
|
|
return func(ctx context.Context) (usagetypes.HeartbeatEvent, error) {
|
|
//nolint:gocritic // We are a publisher in this function
|
|
ctx = dbauthz.AsUsagePublisher(ctx)
|
|
count, err := db.GetActiveAISeatCount(ctx)
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("get active AI seat count: %w", err)
|
|
}
|
|
|
|
return usagetypes.HBAISeats{Count: count}, nil
|
|
}
|
|
}
|