perf(enterprise): remove expensive GetWorkspaces query from entitlements (#19747)

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

This PR addresses the significant database load issue where the
`GetWorkspaces` query was causing performance problems in the license
entitlements code.
This commit is contained in:
Kacper Sawicki
2025-09-09 15:46:11 +02:00
committed by GitHub
parent d527f91f47
commit 3074547322
2 changed files with 8 additions and 49 deletions
+8 -33
View File
@@ -96,17 +96,6 @@ func Entitlements(
return codersdk.Entitlements{}, xerrors.Errorf("query active user count: %w", err)
}
// nolint:gocritic // Getting external workspaces is a system function.
externalWorkspaces, err := db.GetWorkspaces(dbauthz.AsSystemRestricted(ctx), database.GetWorkspacesParams{
HasExternalAgent: sql.NullBool{
Bool: true,
Valid: true,
},
})
if err != nil {
return codersdk.Entitlements{}, xerrors.Errorf("query external workspaces: %w", err)
}
// nolint:gocritic // Getting external templates is a system function.
externalTemplates, err := db.GetTemplatesWithFilter(dbauthz.AsSystemRestricted(ctx), database.GetTemplatesWithFilterParams{
HasExternalAgent: sql.NullBool{
@@ -119,11 +108,10 @@ func Entitlements(
}
entitlements, err := LicensesEntitlements(ctx, now, licenses, enablements, keys, FeatureArguments{
ActiveUserCount: activeUserCount,
ReplicaCount: replicaCount,
ExternalAuthCount: externalAuthCount,
ExternalWorkspaceCount: int64(len(externalWorkspaces)),
ExternalTemplateCount: int64(len(externalTemplates)),
ActiveUserCount: activeUserCount,
ReplicaCount: replicaCount,
ExternalAuthCount: externalAuthCount,
ExternalTemplateCount: int64(len(externalTemplates)),
ManagedAgentCountFn: func(ctx context.Context, startTime time.Time, endTime time.Time) (int64, error) {
// This is not super accurate, as the start and end times will be
// truncated to the date in UTC timezone. This is an optimization
@@ -149,11 +137,10 @@ func Entitlements(
}
type FeatureArguments struct {
ActiveUserCount int64
ReplicaCount int
ExternalAuthCount int
ExternalWorkspaceCount int64
ExternalTemplateCount int64
ActiveUserCount int64
ReplicaCount int
ExternalAuthCount int
ExternalTemplateCount int64
// Unfortunately, managed agent count is not a simple count of the current
// state of the world, but a count between two points in time determined by
// the licenses.
@@ -477,18 +464,6 @@ func LicensesEntitlements(
}
}
if featureArguments.ExternalWorkspaceCount > 0 {
feature := entitlements.Features[codersdk.FeatureWorkspaceExternalAgent]
switch feature.Entitlement {
case codersdk.EntitlementNotEntitled:
entitlements.Errors = append(entitlements.Errors,
"You have external workspaces but your license is not entitled to this feature.")
case codersdk.EntitlementGracePeriod:
entitlements.Warnings = append(entitlements.Warnings,
"You have external workspaces but your license is expired.")
}
}
if featureArguments.ExternalTemplateCount > 0 {
feature := entitlements.Features[codersdk.FeatureWorkspaceExternalAgent]
switch feature.Entitlement {
-16
View File
@@ -843,9 +843,6 @@ func TestEntitlements(t *testing.T) {
return true
})).
Return(int64(175), nil)
mDB.EXPECT().
GetWorkspaces(gomock.Any(), gomock.Any()).
Return([]database.GetWorkspacesRow{}, nil)
mDB.EXPECT().
GetTemplatesWithFilter(gomock.Any(), gomock.Any()).
Return([]database.Template{}, nil)
@@ -1236,19 +1233,6 @@ func TestLicenseEntitlements(t *testing.T) {
assert.Equal(t, int64(200), *feature.Actual)
},
},
{
Name: "ExternalWorkspace",
Licenses: []*coderdenttest.LicenseOptions{
enterpriseLicense().UserLimit(100),
},
Arguments: license.FeatureArguments{
ExternalWorkspaceCount: 1,
},
AssertEntitlements: func(t *testing.T, entitlements codersdk.Entitlements) {
assert.Equal(t, codersdk.EntitlementEntitled, entitlements.Features[codersdk.FeatureWorkspaceExternalAgent].Entitlement)
assert.True(t, entitlements.Features[codersdk.FeatureWorkspaceExternalAgent].Enabled)
},
},
{
Name: "ExternalTemplate",
Licenses: []*coderdenttest.LicenseOptions{