Sync: Move ActiveTokenCount to a new service (#52991)

* Move ActiveTokenCount to a new service

* Fixing tests

* fix `RootSystem`

Co-authored-by: Artur Wierzbicki <artur.wierzbicki@grafana.com>

pkg/models/user_token.go

@@ -71,12 +71,15 @@ type UserTokenService interface {
 	TryRotateToken(ctx context.Context, token *UserToken, clientIP net.IP, userAgent string) (bool, error)
 	RevokeToken(ctx context.Context, token *UserToken, soft bool) error
 	RevokeAllUserTokens(ctx context.Context, userId int64) error
-	ActiveTokenCount(ctx context.Context) (int64, error)
 	GetUserToken(ctx context.Context, userId, userTokenId int64) (*UserToken, error)
 	GetUserTokens(ctx context.Context, userId int64) ([]*UserToken, error)
 	GetUserRevokedTokens(ctx context.Context, userId int64) ([]*UserToken, error)
 }
 
+type ActiveTokenService interface {
+	ActiveTokenCount(ctx context.Context) (int64, error)
+}
+
 type UserTokenBackgroundService interface {
 	registry.BackgroundService
 }

pkg/server/wire.go

@@ -6,6 +6,7 @@ package server
 import (
 	"github.com/google/wire"
 	sdkhttpclient "github.com/grafana/grafana-plugin-sdk-go/backend/httpclient"
+	"github.com/grafana/grafana/pkg/services/auth"
 	"github.com/grafana/grafana/pkg/services/playlist/playlistimpl"
 	"github.com/grafana/grafana/pkg/services/store/sanitizer"
 
@@ -222,6 +223,8 @@ var wireBasicSet = wire.NewSet(
 	influxdb.ProvideService,
 	wire.Bind(new(social.Service), new(*social.SocialService)),
 	oauthtoken.ProvideService,
+	auth.ProvideActiveAuthTokenService,
+	wire.Bind(new(models.ActiveTokenService), new(*auth.ActiveAuthTokenService)),
 	wire.Bind(new(oauthtoken.OAuthTokenService), new(*oauthtoken.Service)),
 	tempo.ProvideService,
 	loki.ProvideService,

pkg/services/auth/auth_token.go

@@ -42,14 +42,26 @@ type UserAuthTokenService struct {
 	log               log.Logger
 }
 
-func (s *UserAuthTokenService) ActiveTokenCount(ctx context.Context) (int64, error) {
+type ActiveAuthTokenService struct {
+	cfg      *setting.Cfg
+	sqlStore sqlstore.Store
+}
+
+func ProvideActiveAuthTokenService(cfg *setting.Cfg, sqlStore sqlstore.Store) *ActiveAuthTokenService {
+	return &ActiveAuthTokenService{
+		cfg:      cfg,
+		sqlStore: sqlStore,
+	}
+}
+
+func (a *ActiveAuthTokenService) ActiveTokenCount(ctx context.Context) (int64, error) {
 	var count int64
 	var err error
-	err = s.SQLStore.WithDbSession(ctx, func(dbSession *sqlstore.DBSession) error {
+	err = a.sqlStore.WithDbSession(ctx, func(dbSession *sqlstore.DBSession) error {
 		var model userAuthToken
 		count, err = dbSession.Where(`created_at > ? AND rotated_at > ? AND revoked_at = 0`,
-			s.createdAfterParam(),
+			getTime().Add(-a.cfg.LoginMaxLifetime).Unix(),
-			s.rotatedAfterParam()).
+			getTime().Add(-a.cfg.LoginMaxInactiveLifetime).Unix()).
 			Count(&model)
 
 		return err

pkg/services/auth/auth_token_test.go

@@ -41,7 +41,7 @@ func TestUserAuthToken(t *testing.T) {
 		userToken := createToken()
 
 		t.Run("Can count active tokens", func(t *testing.T) {
-			count, err := ctx.tokenService.ActiveTokenCount(context.Background())
+			count, err := ctx.activeTokenService.ActiveTokenCount(context.Background())
 			require.Nil(t, err)
 			require.Equal(t, int64(1), count)
 		})
@@ -209,7 +209,7 @@ func TestUserAuthToken(t *testing.T) {
 			require.Nil(t, notGood)
 
 			t.Run("should not find active token when expired", func(t *testing.T) {
-				count, err := ctx.tokenService.ActiveTokenCount(context.Background())
+				count, err := ctx.activeTokenService.ActiveTokenCount(context.Background())
 				require.Nil(t, err)
 				require.Equal(t, int64(0), count)
 			})
@@ -534,25 +534,35 @@ func createTestContext(t *testing.T) *testContext {
 	maxInactiveDurationVal, _ := time.ParseDuration("168h")
 	maxLifetimeDurationVal, _ := time.ParseDuration("720h")
 	sqlstore := sqlstore.InitTestDB(t)
+
+	cfg := &setting.Cfg{
+		LoginMaxInactiveLifetime:     maxInactiveDurationVal,
+		LoginMaxLifetime:             maxLifetimeDurationVal,
+		TokenRotationIntervalMinutes: 10,
+	}
+
 	tokenService := &UserAuthTokenService{
 		SQLStore: sqlstore,
-		Cfg: &setting.Cfg{
+		Cfg:      cfg,
-			LoginMaxInactiveLifetime:     maxInactiveDurationVal,
+		log:      log.New("test-logger"),
-			LoginMaxLifetime:             maxLifetimeDurationVal,
+	}
-			TokenRotationIntervalMinutes: 10,
+
-		},
+	activeTokenService := &ActiveAuthTokenService{
-		log: log.New("test-logger"),
+		cfg:      cfg,
+		sqlStore: sqlstore,
 	}
 
 	return &testContext{
-		sqlstore:     sqlstore,
+		sqlstore:           sqlstore,
-		tokenService: tokenService,
+		tokenService:       tokenService,
+		activeTokenService: activeTokenService,
 	}
 }
 
 type testContext struct {
-	sqlstore     *sqlstore.SQLStore
+	sqlstore           *sqlstore.SQLStore
-	tokenService *UserAuthTokenService
+	tokenService       *UserAuthTokenService
+	activeTokenService *ActiveAuthTokenService
 }
 
 func (c *testContext) getAuthTokenByID(id int64) (*userAuthToken, error) {

pkg/services/quota/quotaimpl/quota.go

@@ -13,17 +13,17 @@ import (
 
 type Service struct {
 	store            store
-	AuthTokenService models.UserTokenService
+	authTokenService models.ActiveTokenService
 	Cfg              *setting.Cfg
 	SQLStore         sqlstore.Store
 	Logger           log.Logger
 }
 
-func ProvideService(db db.DB, cfg *setting.Cfg, tokenService models.UserTokenService, ss *sqlstore.SQLStore) quota.Service {
+func ProvideService(db db.DB, cfg *setting.Cfg, tokenService models.ActiveTokenService, ss *sqlstore.SQLStore) quota.Service {
 	return &Service{
 		store:            &sqlStore{db: db},
 		Cfg:              cfg,
-		AuthTokenService: tokenService,
+		authTokenService: tokenService,
 		SQLStore:         ss,
 		Logger:           log.New("quota_service"),
 	}
@@ -71,7 +71,7 @@ func (s *Service) CheckQuotaReached(ctx context.Context, target string, scopePar
 				return true, nil
 			}
 			if target == "session" {
-				usedSessions, err := s.AuthTokenService.ActiveTokenCount(ctx)
+				usedSessions, err := s.authTokenService.ActiveTokenCount(ctx)
 				if err != nil {
 					return false, err
 				}

pkg/services/store/service.go

@@ -171,7 +171,7 @@ func ProvideService(
 		storages = append(storages,
 			newSQLStorage(RootStorageMeta{
 				Builtin: true,
-			}, RootResources,
+			}, RootSystem,
 				"System",
 				"Grafana system storage",
 				&StorageSQLConfig{}, sql, orgId))