public/grafana
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Optimize tokens count from service accounts query (#96663)

Browse Source 
* optimize tokens count from service accounts query

* add unit tests for tokens count

* skip broken test

* fix lint error

* rename Tokens to TokenCount
This commit is contained in:
Mihai Doarna
2024-11-20 14:52:23 +02:00
committed by GitHub
parent f1601b1c0f
commit cb7dd25f8a
6 changed files with 146 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/services/serviceaccounts/database/store.go
+37
View File
@@ -402,6 +402,43 @@ func (s *ServiceAccountsStoreImpl) SearchOrgServiceAccounts(ctx context.Context,
if err := sess.Find(&searchResult.ServiceAccounts); err != nil {
return err
}
// stop here if we don't want to count the number of tokens per service account
if !query.CountTokens {
return nil
}
// Fetch tokens count for each service account
accountIDs := make([]int64, len(searchResult.ServiceAccounts))
for i, serviceAccount := range searchResult.ServiceAccounts {
accountIDs[i] = serviceAccount.Id
}
tokensSess := dbSession.Table("api_key").
Select("service_account_id, COUNT(id) AS token_count").
Where("org_id = ?", query.OrgID).
In("service_account_id", accountIDs).
GroupBy("service_account_id")
type tokenCount struct {
AccountId int64 `xorm:"service_account_id"`
TokenCount int64 `xorm:"token_count"`
}
var tokens []tokenCount
if err = tokensSess.Find(&tokens); err != nil {
return err
}
tokenMap := make(map[int64]int64)
for _, token := range tokens {
tokenMap[token.AccountId] = token.TokenCount
}
for i, account := range searchResult.ServiceAccounts {
searchResult.ServiceAccounts[i].Tokens = tokenMap[account.Id]
}
return nil
})
if err != nil {
pkg/services/serviceaccounts/database/store_test.go
+89 -25
View File
@@ -488,41 +488,86 @@ func TestIntegrationServiceAccountsStoreImpl_SearchOrgServiceAccounts(t *testing
t.Skip("skipping test in short mode")
}
db, store := setupTestDatabase(t)
initUsers := []tests.TestUser{
{Name: "satest-1", Role: string(org.RoleViewer), Login: "sa-1-satest-1", IsServiceAccount: true},
{Name: "extsvc-test-1", Role: string(org.RoleNone), Login: "sa-1-extsvc-test-1", IsServiceAccount: true},
{Name: "usertest-2", Role: string(org.RoleEditor), Login: "usertest-2", IsServiceAccount: false},
{Name: "satest-3", Role: string(org.RoleEditor), Login: "sa-1-satest-3", IsServiceAccount: true},
{Name: "satest-4", Role: string(org.RoleAdmin), Login: "sa-1-satest-4", IsServiceAccount: true},
{Name: "extsvc-test-3", Role: string(org.RoleNone), Login: "sa-1-extsvc-test-3", IsServiceAccount: true},
{Name: "extsvc-test-4", Role: string(org.RoleNone), Login: "sa-1-extsvc-test-4", IsServiceAccount: true},
{Name: "extsvc-test-5", Role: string(org.RoleNone), Login: "sa-1-extsvc-test-5", IsServiceAccount: true},
{Name: "extsvc-test-6", Role: string(org.RoleNone), Login: "sa-1-extsvc-test-6", IsServiceAccount: true},
{Name: "extsvc-test-7", Role: string(org.RoleNone), Login: "sa-1-extsvc-test-7", IsServiceAccount: true},
{Name: "extsvc-test-8", Role: string(org.RoleNone), Login: "sa-1-extsvc-test-8", IsServiceAccount: true},
{Name: "satest-6", Role: string(org.RoleViewer), Login: "sa-1-satest-6", IsServiceAccount: true},
{Name: "satest-7", Role: string(org.RoleEditor), Login: "sa-1-satest-7", IsServiceAccount: true},
{Name: "satest-8", Role: string(org.RoleAdmin), Login: "sa-1-satest-8", IsServiceAccount: true},
}
db, store := setupTestDatabase(t)
orgID := tests.SetupUsersServiceAccounts(t, db, store.cfg, initUsers)
users, orgID := tests.SetupUsersServiceAccounts(t, db, store.cfg, initUsers)
apiKeys := []tests.TestApiKey{
{Name: "sa-01-apikey-01", OrgId: orgID, Key: "key01", IsExpired: false, ServiceAccountID: &users[0].ID},
{Name: "sa-01-apikey-02", OrgId: orgID, Key: "key02", IsExpired: false, ServiceAccountID: &users[0].ID},
{Name: "sa-01-apikey-03", OrgId: orgID, Key: "key03", IsExpired: false, ServiceAccountID: &users[0].ID},
{Name: "sa-02-apikey-01", OrgId: orgID, Key: "key04", IsExpired: false, ServiceAccountID: &users[2].ID},
{Name: "sa-02-apikey-02", OrgId: orgID, Key: "key05", IsExpired: false, ServiceAccountID: &users[2].ID},
{Name: "sa-03-apikey-01", OrgId: orgID, Key: "key06", IsExpired: false, ServiceAccountID: &users[3].ID},
}
tests.SetupApiKeys(t, db, store.cfg, apiKeys)
userWithPerm := &user.SignedInUser{
OrgID: orgID,
Permissions: map[int64]map[string][]string{orgID: {serviceaccounts.ActionRead: {serviceaccounts.ScopeAll}}},
}
expectedServiceAccount := func(i int, tokens int64) *serviceaccounts.ServiceAccountDTO {
return &serviceaccounts.ServiceAccountDTO{
Id: users[i].ID, UID: users[i].UID, Name: users[i].Name, Login: users[i].Login, OrgId: orgID, Role: "None", Tokens: tokens,
}
}
tt := []struct {
desc string
query *serviceaccounts.SearchOrgServiceAccountsQuery
expectedTotal int64 // Value of the result.TotalCount
expectedCount int // Length of the result.ServiceAccounts slice
expectedErr error
desc string
query *serviceaccounts.SearchOrgServiceAccountsQuery
expectedTotal int64 // Value of the result.TotalCount
expectedServiceAccounts []*serviceaccounts.ServiceAccountDTO
expectedErr error
}{
{
desc: "should list all service accounts",
desc: "should list all service accounts with tokens count",
query: &serviceaccounts.SearchOrgServiceAccountsQuery{
OrgID: orgID,
SignedInUser: userWithPerm,
Filter: serviceaccounts.FilterIncludeAll,
CountTokens: true,
},
expectedTotal: 7,
expectedServiceAccounts: []*serviceaccounts.ServiceAccountDTO{
expectedServiceAccount(0, 3),
expectedServiceAccount(2, 2),
expectedServiceAccount(3, 1),
expectedServiceAccount(4, 0),
expectedServiceAccount(5, 0),
expectedServiceAccount(6, 0),
expectedServiceAccount(7, 0),
},
},
{
desc: "should list all service accounts with no tokens count",
query: &serviceaccounts.SearchOrgServiceAccountsQuery{
OrgID: orgID,
SignedInUser: userWithPerm,
Filter: serviceaccounts.FilterIncludeAll,
},
expectedTotal: 7,
expectedCount: 7,
expectedServiceAccounts: []*serviceaccounts.ServiceAccountDTO{
expectedServiceAccount(0, 0),
expectedServiceAccount(2, 0),
expectedServiceAccount(3, 0),
expectedServiceAccount(4, 0),
expectedServiceAccount(5, 0),
expectedServiceAccount(6, 0),
expectedServiceAccount(7, 0),
},
},
{
desc: "should list no service accounts without permissions",
@@ -534,8 +579,8 @@ func TestIntegrationServiceAccountsStoreImpl_SearchOrgServiceAccounts(t *testing
},
Filter: serviceaccounts.FilterIncludeAll,
},
expectedTotal: 0,
expectedCount: 0,
expectedTotal: 0,
expectedServiceAccounts: []*serviceaccounts.ServiceAccountDTO{},
},
{
desc: "should list one service accounts with restricted permissions",
@@ -551,7 +596,10 @@ func TestIntegrationServiceAccountsStoreImpl_SearchOrgServiceAccounts(t *testing
Filter: serviceaccounts.FilterIncludeAll,
},
expectedTotal: 2,
expectedCount: 2,
expectedServiceAccounts: []*serviceaccounts.ServiceAccountDTO{
expectedServiceAccount(0, 0),
expectedServiceAccount(6, 0),
},
},
{
desc: "should list only external service accounts",
@@ -559,9 +607,15 @@ func TestIntegrationServiceAccountsStoreImpl_SearchOrgServiceAccounts(t *testing
OrgID: orgID,
SignedInUser: userWithPerm,
Filter: serviceaccounts.FilterOnlyExternal,
CountTokens: true,
},
expectedTotal: 4,
expectedCount: 4,
expectedServiceAccounts: []*serviceaccounts.ServiceAccountDTO{
expectedServiceAccount(0, 3),
expectedServiceAccount(2, 2),
expectedServiceAccount(3, 1),
expectedServiceAccount(4, 0),
},
},
{
desc: "should return service accounts with sa-1-satest login",
@@ -570,9 +624,14 @@ func TestIntegrationServiceAccountsStoreImpl_SearchOrgServiceAccounts(t *testing
Query: "sa-1-satest",
SignedInUser: userWithPerm,
Filter: serviceaccounts.FilterIncludeAll,
CountTokens: true,
},
expectedTotal: 3,
expectedCount: 3,
expectedServiceAccounts: []*serviceaccounts.ServiceAccountDTO{
expectedServiceAccount(5, 0),
expectedServiceAccount(6, 0),
expectedServiceAccount(7, 0),
},
},
{
desc: "should only count service accounts",
@@ -582,8 +641,8 @@ func TestIntegrationServiceAccountsStoreImpl_SearchOrgServiceAccounts(t *testing
Filter: serviceaccounts.FilterIncludeAll,
CountOnly: true,
},
expectedTotal: 7,
expectedCount: 0,
expectedTotal: 7,
expectedServiceAccounts: []*serviceaccounts.ServiceAccountDTO{},
},
{
desc: "should paginate result",
@@ -595,7 +654,9 @@ func TestIntegrationServiceAccountsStoreImpl_SearchOrgServiceAccounts(t *testing
Filter: serviceaccounts.FilterIncludeAll,
},
expectedTotal: 7,
expectedCount: 1,
expectedServiceAccounts: []*serviceaccounts.ServiceAccountDTO{
expectedServiceAccount(7, 0),
},
},
}
for _, tc := range tt {
@@ -609,7 +670,10 @@ func TestIntegrationServiceAccountsStoreImpl_SearchOrgServiceAccounts(t *testing
}
require.Equal(t, tc.expectedTotal, got.TotalCount)
require.Len(t, got.ServiceAccounts, tc.expectedCount)
require.Len(t, got.ServiceAccounts, len(tc.expectedServiceAccounts))
for i, sa := range got.ServiceAccounts {
require.EqualValues(t, tc.expectedServiceAccounts[i], sa)
}
})
}
}
@@ -628,7 +692,7 @@ func TestIntegrationServiceAccountsStoreImpl_EnableServiceAccounts(t *testing.T)
}
db, store := setupTestDatabase(t)
orgID := tests.SetupUsersServiceAccounts(t, db, store.cfg, initUsers)
_, orgID := tests.SetupUsersServiceAccounts(t, db, store.cfg, initUsers)
fetchStates := func() map[int64]bool {
sa1, err := store.RetrieveServiceAccount(ctx, &serviceaccounts.GetServiceAccountQuery{OrgID: orgID, ID: 1})