Skip to content

Commit 196c188

Browse files
deansheatherdeansheather
committed
refactor(usage): move license selection logic into tallymansdk.New
- Move getBestLicenseJWT logic from publisher.go to tallymansdk.New() - Add NewOptions struct for tallymansdk client creation - Export ErrNoLicenseSupportsPublishing for reuse in other contexts - Simplify publisher.publishOnce by delegating license selection to SDK
1 parent 7ee1daa commit 196c188

File tree

2 files changed

+106
-75
lines changed

2 files changed

+106
-75
lines changed

‎enterprise/coderd/usage/publisher.go‎

Lines changed: 11 additions & 73 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@ import (
1919
"github.com/coder/coder/v2/coderd/pproflabel"
2020
"github.com/coder/coder/v2/coderd/usage/usagetypes"
2121
"github.com/coder/coder/v2/cryptorand"
22-
"github.com/coder/coder/v2/enterprise/coderd/license"
2322
"github.com/coder/coder/v2/enterprise/coderd/usage/tallymansdk"
2423
"github.com/coder/quartz"
2524
)
@@ -33,8 +32,6 @@ const (
3332
tallymanPublishBatchSize = 100
3433
)
3534

36-
var errUsagePublishingDisabled = xerrors.New("usage publishing is not enabled by any license")
37-
3835
// Publisher publishes usage events ***somewhere***.
3936
type Publisher interface {
4037
// Close closes the publisher and waits for it to finish.
@@ -203,11 +200,17 @@ func (p *tallymanPublisher) publish(ctx context.Context, deploymentID uuid.UUID)
203200
// publishOnce publishes up to tallymanPublishBatchSize usage events to
204201
// tallyman. It returns the number of successfully published events.
205202
func (p *tallymanPublisher) publishOnce(ctx context.Context, deploymentID uuid.UUID) (int, error) {
206-
licenseJwt, err := p.getBestLicenseJWT(ctx)
207-
if xerrors.Is(err, errUsagePublishingDisabled) {
203+
sdkClient, err := tallymansdk.New(ctx, tallymansdk.NewOptions{
204+
DB: p.db,
205+
DeploymentID: deploymentID,
206+
LicenseKeys: p.licenseKeys,
207+
BaseURL: p.baseURL,
208+
HTTPClient: p.httpClient,
209+
})
210+
if xerrors.Is(err, tallymansdk.ErrNoLicenseSupportsPublishing) {
208211
return 0, nil
209212
} else if err != nil {
210-
return 0, xerrors.Errorf("find usage publishing license: %w", err)
213+
return 0, xerrors.Errorf("create tallyman client: %w", err)
211214
}
212215

213216
events, err := p.db.SelectUsageEventsForPublishing(ctx, dbtime.Time(p.clock.Now()))
@@ -246,7 +249,7 @@ func (p *tallymanPublisher) publishOnce(ctx context.Context, deploymentID uuid.U
246249
return 0, xerrors.Errorf("duplicate event IDs found in events for publishing")
247250
}
248251

249-
resp, err := p.sendPublishRequest(ctx, deploymentID, licenseJwt, tallymanReq)
252+
resp, err := p.sendPublishRequest(ctx, sdkClient, tallymanReq)
250253
allFailed := err != nil
251254
if err != nil {
252255
p.log.Warn(ctx, "failed to send publish request to tallyman", slog.F("count", len(events)), slog.Error(err))
@@ -336,72 +339,7 @@ func (p *tallymanPublisher) publishOnce(ctx context.Context, deploymentID uuid.U
336339
return len(resp.AcceptedEvents), returnErr
337340
}
338341

339-
// getBestLicenseJWT returns the best license JWT to use for the request. The
340-
// criteria is as follows:
341-
// - The license must be valid and active (after nbf, before exp)
342-
// - The license must have usage publishing enabled
343-
// The most recently issued (iat) license is chosen.
344-
//
345-
// If no licenses are found or none have usage publishing enabled,
346-
// errUsagePublishingDisabled is returned.
347-
func (p *tallymanPublisher) getBestLicenseJWT(ctx context.Context) (string, error) {
348-
licenses, err := p.db.GetUnexpiredLicenses(ctx)
349-
if err != nil {
350-
return "", xerrors.Errorf("get unexpired licenses: %w", err)
351-
}
352-
if len(licenses) == 0 {
353-
return "", errUsagePublishingDisabled
354-
}
355-
356-
type licenseJWTWithClaims struct {
357-
Claims *license.Claims
358-
Raw string
359-
}
360-
361-
var bestLicense licenseJWTWithClaims
362-
for _, dbLicense := range licenses {
363-
claims, err := license.ParseClaims(dbLicense.JWT, p.licenseKeys)
364-
if err != nil {
365-
p.log.Warn(ctx, "failed to parse license claims", slog.F("license_id", dbLicense.ID), slog.Error(err))
366-
continue
367-
}
368-
if claims.AccountType != license.AccountTypeSalesforce {
369-
// Non-Salesforce accounts cannot be tracked as they do not have a
370-
// trusted Salesforce opportunity ID encoded in the license.
371-
continue
372-
}
373-
if !claims.PublishUsageData {
374-
// Publishing is disabled.
375-
continue
376-
}
377-
378-
// Otherwise, if it's issued more recently, it's the best license.
379-
// IssuedAt is verified to be non-nil in license.ParseClaims.
380-
if bestLicense.Claims == nil || claims.IssuedAt.Time.After(bestLicense.Claims.IssuedAt.Time) {
381-
bestLicense = licenseJWTWithClaims{
382-
Claims: claims,
383-
Raw: dbLicense.JWT,
384-
}
385-
}
386-
}
387-
388-
if bestLicense.Raw == "" {
389-
return "", errUsagePublishingDisabled
390-
}
391-
392-
return bestLicense.Raw, nil
393-
}
394-
395-
func (p *tallymanPublisher) sendPublishRequest(ctx context.Context, deploymentID uuid.UUID, licenseJwt string, req usagetypes.TallymanV1IngestRequest) (usagetypes.TallymanV1IngestResponse, error) {
396-
// Create a new SDK client for this request.
397-
// We create it per-request since the license key may change.
398-
sdkClient := tallymansdk.New(
399-
p.baseURL,
400-
licenseJwt,
401-
deploymentID,
402-
tallymansdk.WithHTTPClient(p.httpClient),
403-
)
404-
342+
func (*tallymanPublisher) sendPublishRequest(ctx context.Context, sdkClient *tallymansdk.Client, req usagetypes.TallymanV1IngestRequest) (usagetypes.TallymanV1IngestResponse, error) {
405343
return sdkClient.PublishUsageEvents(ctx, req)
406344
}
407345

‎enterprise/coderd/usage/tallymansdk/client.go‎

Lines changed: 95 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package tallymansdk
33
import (
44
"bytes"
55
"context"
6+
"crypto/ed25519"
67
"encoding/json"
78
"io"
89
"net/http"
@@ -12,14 +13,32 @@ import (
1213
"golang.org/x/xerrors"
1314

1415
"github.com/coder/coder/v2/buildinfo"
16+
"github.com/coder/coder/v2/coderd/database"
1517
"github.com/coder/coder/v2/coderd/usage/usagetypes"
18+
"github.com/coder/coder/v2/enterprise/coderd/license"
1619
)
1720

1821
const (
1922
// DefaultURL is the default URL for the Tallyman API.
2023
DefaultURL = "https://tallyman-prod.coder.com"
2124
)
2225

26+
var ErrNoLicenseSupportsPublishing = xerrors.New("usage publishing is not enabled by any license")
27+
28+
// NewOptions contains options for creating a new Tallyman client.
29+
type NewOptions struct {
30+
// DB is the database store for querying licenses and deployment ID.
31+
DB database.Store
32+
// DeploymentID is the deployment ID. If uuid.Nil, it will be fetched from the database.
33+
DeploymentID uuid.UUID
34+
// LicenseKeys is a map of license keys for verifying license JWTs.
35+
LicenseKeys map[string]ed25519.PublicKey
36+
// BaseURL is the base URL for the Tallyman API. If nil, DefaultURL is used.
37+
BaseURL *url.URL
38+
// HTTPClient is the HTTP client to use for requests. If nil, http.DefaultClient is used.
39+
HTTPClient *http.Client
40+
}
41+
2342
// Client is a client for the Tallyman API.
2443
type Client struct {
2544
// URL is the base URL for the Tallyman API.
@@ -35,8 +54,8 @@ type Client struct {
3554
// ClientOption is a functional option for configuring the Client.
3655
type ClientOption func(*Client)
3756

38-
// New creates a new Tallyman API client.
39-
func New(baseURL *url.URL, licenseKey string, deploymentID uuid.UUID, opts ...ClientOption) *Client {
57+
// NewWithAuth creates a new Tallyman API client with explicit authentication.
58+
func NewWithAuth(baseURL *url.URL, licenseKey string, deploymentID uuid.UUID, opts ...ClientOption) *Client {
4059
if baseURL == nil {
4160
baseURL, _ = url.Parse(DefaultURL)
4261
}
@@ -64,6 +83,80 @@ func WithHTTPClient(httpClient *http.Client) ClientOption {
6483
}
6584
}
6685

86+
// New creates a new Tallyman API client by looking up the best license from the database.
87+
// It selects the most recently issued license that:
88+
// - Is unexpired
89+
// - Has AccountType == AccountTypeSalesforce
90+
// - Has PublishUsageData enabled
91+
//
92+
// If opts.DeploymentID is uuid.Nil, it will be fetched from the database.
93+
// If no suitable license is found, it returns ErrNoLicenseSupportsPublishing.
94+
func New(ctx context.Context, opts NewOptions) (*Client, error) {
95+
// Fetch deployment ID if not provided
96+
deploymentID := opts.DeploymentID
97+
if deploymentID == uuid.Nil {
98+
deploymentIDStr, err := opts.DB.GetDeploymentID(ctx)
99+
if err != nil {
100+
return nil, xerrors.Errorf("get deployment ID: %w", err)
101+
}
102+
deploymentID, err = uuid.Parse(deploymentIDStr)
103+
if err != nil {
104+
return nil, xerrors.Errorf("parse deployment ID %q: %w", deploymentIDStr, err)
105+
}
106+
}
107+
108+
licenses, err := opts.DB.GetUnexpiredLicenses(ctx)
109+
if err != nil {
110+
return nil, xerrors.Errorf("get unexpired licenses: %w", err)
111+
}
112+
if len(licenses) == 0 {
113+
return nil, ErrNoLicenseSupportsPublishing
114+
}
115+
116+
type licenseJWTWithClaims struct {
117+
Claims *license.Claims
118+
Raw string
119+
}
120+
121+
var bestLicense licenseJWTWithClaims
122+
for _, dbLicense := range licenses {
123+
claims, err := license.ParseClaims(dbLicense.JWT, opts.LicenseKeys)
124+
if err != nil {
125+
// Skip licenses that can't be parsed
126+
continue
127+
}
128+
if claims.AccountType != license.AccountTypeSalesforce {
129+
// Non-Salesforce accounts cannot be tracked
130+
continue
131+
}
132+
if !claims.PublishUsageData {
133+
// Publishing is disabled
134+
continue
135+
}
136+
137+
// Select the most recently issued license
138+
// IssuedAt is verified to be non-nil in license.ParseClaims
139+
if bestLicense.Claims == nil || claims.IssuedAt.Time.After(bestLicense.Claims.IssuedAt.Time) {
140+
bestLicense = licenseJWTWithClaims{
141+
Claims: claims,
142+
Raw: dbLicense.JWT,
143+
}
144+
}
145+
}
146+
147+
if bestLicense.Raw == "" {
148+
return nil, ErrNoLicenseSupportsPublishing
149+
}
150+
151+
// Set default HTTP client if not provided
152+
httpClient := opts.HTTPClient
153+
if httpClient == nil {
154+
httpClient = http.DefaultClient
155+
}
156+
157+
return NewWithAuth(opts.BaseURL, bestLicense.Raw, deploymentID, WithHTTPClient(httpClient)), nil
158+
}
159+
67160
// Request makes an HTTP request to the Tallyman API.
68161
// It sets the authentication headers and User-Agent automatically.
69162
func (c *Client) Request(ctx context.Context, method, path string, body interface{}) (*http.Response, error) {

0 commit comments

Comments
 (0)