Merge remote-tracking branch 'origin/main' into grafana-git-ui-sync

pkg/services/accesscontrol/database/policies.go

@@ -42,6 +42,9 @@ func GetAccessPolicies(ctx context.Context, orgID int64, sql *session.SessionDB,
 	if err != nil {
 		return nil, err
 	}
+	defer func() {
+		_ = rows.Close()
+	}()
 
 	created := time.Now()
 	updated := time.Now()

pkg/services/accesscontrol/dualwrite/collectors.go

@@ -115,7 +115,7 @@ func folderTreeCollector(store db.DB) legacyTupleCollector {
 func managedPermissionsCollector(store db.DB, kind string) legacyTupleCollector {
 	return func(ctx context.Context, orgID int64) (map[string]map[string]*openfgav1.TupleKey, error) {
 		query := `
-			SELECT u.uid as user_uid, t.uid as team_uid, p.action, p.kind, p.identifier, r.org_id, br.role as basic_role_name
+			SELECT u.uid as user_uid, u.is_service_account as is_service_account, t.uid as team_uid, p.action, p.kind, p.identifier, r.org_id, br.role as basic_role_name
 			FROM permission p
 			INNER JOIN role r ON p.role_id = r.id
 			LEFT JOIN user_role ur ON r.id = ur.role_id
@@ -128,12 +128,13 @@ func managedPermissionsCollector(store db.DB, kind string) legacyTupleCollector
 			AND p.kind = ?
 		`
 		type Permission struct {
-			Action        string `xorm:"action"`
-			Kind          string
-			Identifier    string
-			UserUID       string `xorm:"user_uid"`
-			TeamUID       string `xorm:"team_uid"`
-			BasicRoleName string `xorm:"basic_role_name"`
+			Action           string `xorm:"action"`
+			Kind             string
+			Identifier       string
+			UserUID          string `xorm:"user_uid"`
+			IsServiceAccount bool   `xorm:"is_service_account"`
+			TeamUID          string `xorm:"team_uid"`
+			BasicRoleName    string `xorm:"basic_role_name"`
 		}
 
 		var permissions []Permission
@@ -149,7 +150,9 @@ func managedPermissionsCollector(store db.DB, kind string) legacyTupleCollector
 
 		for _, p := range permissions {
 			var subject string
-			if len(p.UserUID) > 0 {
+			if len(p.UserUID) > 0 && p.IsServiceAccount {
+				subject = zanzana.NewTupleEntry(zanzana.TypeServiceAccount, p.UserUID, "")
+			} else if len(p.UserUID) > 0 {
 				subject = zanzana.NewTupleEntry(zanzana.TypeUser, p.UserUID, "")
 			} else if len(p.TeamUID) > 0 {
 				subject = zanzana.NewTupleEntry(zanzana.TypeTeam, p.TeamUID, zanzana.RelationTeamMember)
@@ -198,16 +201,19 @@ func tupleStringWithoutCondition(tuple *openfgav1.TupleKey) string {
 func basicRoleBindingsCollector(store db.DB) legacyTupleCollector {
 	return func(ctx context.Context, orgID int64) (map[string]map[string]*openfgav1.TupleKey, error) {
 		query := `
-			SELECT ou.org_id, u.uid as user_uid, ou.role as org_role
+			SELECT
+				ou.org_id, u.uid as user_uid,
+				u.is_service_account as is_service_account,
+				ou.role as org_role
 			FROM org_user ou
 			LEFT JOIN ` + store.GetDialect().Quote("user") + ` u ON u.id = ou.user_id
 			WHERE ou.org_id = ?
-			AND NOT u.is_service_account
 		`
 		// FIXME: handle service admin role
 		type Binding struct {
-			UserUID string `xorm:"user_uid"`
-			OrgRole string `xorm:"org_role"`
+			UserUID          string `xorm:"user_uid"`
+			IsServiceAccount bool   `xorm:"is_service_account"`
+			OrgRole          string `xorm:"org_role"`
 		}
 
 		var bindings []Binding
@@ -222,8 +228,13 @@ func basicRoleBindingsCollector(store db.DB) legacyTupleCollector {
 		tuples := make(map[string]map[string]*openfgav1.TupleKey)
 
 		for _, b := range bindings {
+			userType := zanzana.TypeUser
+			if b.IsServiceAccount {
+				userType = zanzana.TypeServiceAccount
+			}
+
 			tuple := &openfgav1.TupleKey{
-				User:     zanzana.NewTupleEntry(zanzana.TypeUser, b.UserUID, ""),
+				User:     zanzana.NewTupleEntry(userType, b.UserUID, ""),
 				Relation: zanzana.RelationAssignee,
 				Object:   zanzana.NewTupleEntry(zanzana.TypeRole, zanzana.TranslateBasicRole(b.OrgRole), ""),
 			}
@@ -286,7 +297,7 @@ func teamRoleBindingsCollector(store db.DB) legacyTupleCollector {
 func userRoleBindingsCollector(store db.DB) legacyTupleCollector {
 	return func(ctx context.Context, orgID int64) (map[string]map[string]*openfgav1.TupleKey, error) {
 		query := `
-			SELECT u.uid AS user_uid, r.uid AS role_uid
+			SELECT u.uid AS user_uid, u.is_service_account as is_service_account, r.uid AS role_uid
 			FROM user_role ur
 			INNER JOIN ` + store.GetDialect().Quote("user") + ` u ON ur.user_id = u.id
 			INNER JOIN role r ON ur.role_id = r.id
@@ -294,8 +305,9 @@ func userRoleBindingsCollector(store db.DB) legacyTupleCollector {
 			AND r.name NOT LIKE 'managed:%'
 		`
 		type Binding struct {
-			UserUID string `xorm:"user_uid"`
-			RoleUID string `xorm:"role_uid"`
+			UserUID          string `xorm:"user_uid"`
+			IsServiceAccount bool   `xorm:"is_service_account"`
+			RoleUID          string `xorm:"role_uid"`
 		}
 
 		var bindings []Binding
@@ -310,8 +322,13 @@ func userRoleBindingsCollector(store db.DB) legacyTupleCollector {
 		tuples := make(map[string]map[string]*openfgav1.TupleKey)
 
 		for _, b := range bindings {
+			userType := zanzana.TypeUser
+			if b.IsServiceAccount {
+				userType = zanzana.TypeServiceAccount
+			}
+
 			tuple := &openfgav1.TupleKey{
-				User:     zanzana.NewTupleEntry(zanzana.TypeUser, b.UserUID, ""),
+				User:     zanzana.NewTupleEntry(userType, b.UserUID, ""),
 				Relation: zanzana.RelationAssignee,
 				Object:   zanzana.NewTupleEntry(zanzana.TypeRole, b.RoleUID, ""),
 			}
@@ -397,7 +414,8 @@ func zanzanaCollector(relations []string) zanzanaTupleCollector {
 
 			for c != "" {
 				res, err := client.Read(ctx, &authzextv1.ReadRequest{
-					Namespace: namespace,
+					ContinuationToken: c,
+					Namespace:         namespace,
 					TupleKey: &authzextv1.ReadRequestTupleKey{
 						Object:   object,
 						Relation: relation,

pkg/services/auth/auth.go

@@ -74,12 +74,15 @@ type CreateTokenCommand struct {
 }
 
 // UserTokenService are used for generating and validating user tokens
+//
+//go:generate mockery --name UserTokenService --structname MockUserAuthTokenService --outpkg authtest --filename auth_token_service_mock.go --output ./authtest/
 type UserTokenService interface {
 	CreateToken(ctx context.Context, cmd *CreateTokenCommand) (*UserToken, error)
 	LookupToken(ctx context.Context, unhashedToken string) (*UserToken, error)
 	GetTokenByExternalSessionID(ctx context.Context, externalSessionID int64) (*UserToken, error)
-	GetExternalSession(ctx context.Context, extSessionID int64) (*ExternalSession, error)
+	GetExternalSession(ctx context.Context, externalSessionID int64) (*ExternalSession, error)
 	FindExternalSessions(ctx context.Context, query *ListExternalSessionQuery) ([]*ExternalSession, error)
+	UpdateExternalSession(ctx context.Context, externalSessionID int64, cmd *UpdateExternalSessionCommand) error
 
 	// RotateToken will always rotate a valid token
 	RotateToken(ctx context.Context, cmd RotateCommand) (*UserToken, error)

pkg/services/auth/authimpl/auth_token.go

@@ -248,14 +248,18 @@ func (s *UserAuthTokenService) GetTokenByExternalSessionID(ctx context.Context,
 	return &userToken, err
 }
 
-func (s *UserAuthTokenService) GetExternalSession(ctx context.Context, extSessionID int64) (*auth.ExternalSession, error) {
-	return s.externalSessionStore.Get(ctx, extSessionID)
+func (s *UserAuthTokenService) GetExternalSession(ctx context.Context, externalSessionID int64) (*auth.ExternalSession, error) {
+	return s.externalSessionStore.Get(ctx, externalSessionID)
 }
 
 func (s *UserAuthTokenService) FindExternalSessions(ctx context.Context, query *auth.ListExternalSessionQuery) ([]*auth.ExternalSession, error) {
 	return s.externalSessionStore.List(ctx, query)
 }
 
+func (s *UserAuthTokenService) UpdateExternalSession(ctx context.Context, externalSessionID int64, cmd *auth.UpdateExternalSessionCommand) error {
+	return s.externalSessionStore.Update(ctx, externalSessionID, cmd)
+}
+
 func (s *UserAuthTokenService) RotateToken(ctx context.Context, cmd auth.RotateCommand) (*auth.UserToken, error) {
 	if cmd.UnHashedToken == "" {
 		return nil, auth.ErrInvalidSessionToken

pkg/services/auth/authimpl/external_session_store.go

@@ -27,11 +27,11 @@ func provideExternalSessionStore(sqlStore db.DB, secretService secrets.Service,
 	}
 }
 
-func (s *store) Get(ctx context.Context, extSessionID int64) (*auth.ExternalSession, error) {
+func (s *store) Get(ctx context.Context, ID int64) (*auth.ExternalSession, error) {
 	ctx, span := s.tracer.Start(ctx, "externalsession.Get")
 	defer span.End()
 
-	externalSession := &auth.ExternalSession{ID: extSessionID}
+	externalSession := &auth.ExternalSession{ID: ID}
 
 	err := s.sqlStore.WithDbSession(ctx, func(sess *db.Session) error {
 		found, err := sess.Get(externalSession)
@@ -150,6 +150,45 @@ func (s *store) Create(ctx context.Context, extSession *auth.ExternalSession) er
 	return nil
 }
 
+func (s *store) Update(ctx context.Context, ID int64, cmd *auth.UpdateExternalSessionCommand) error {
+	ctx, span := s.tracer.Start(ctx, "externalsession.Update")
+	defer span.End()
+
+	var err error
+	externalSession := &auth.ExternalSession{}
+
+	externalSession.AccessToken, err = s.encryptAndEncode(cmd.Token.AccessToken)
+	if err != nil {
+		return err
+	}
+
+	externalSession.RefreshToken, err = s.encryptAndEncode(cmd.Token.RefreshToken)
+	if err != nil {
+		return err
+	}
+
+	var secretIdToken string
+	if idToken, ok := cmd.Token.Extra("id_token").(string); ok && idToken != "" {
+		secretIdToken, err = s.encryptAndEncode(idToken)
+		if err != nil {
+			return err
+		}
+		externalSession.IDToken = secretIdToken
+	}
+
+	externalSession.ExpiresAt = cmd.Token.Expiry
+
+	err = s.sqlStore.WithDbSession(ctx, func(sess *db.Session) error {
+		_, err := sess.ID(ID).Cols("access_token", "refresh_token", "id_token", "expires_at").Update(externalSession)
+		return err
+	})
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
 func (s *store) Delete(ctx context.Context, ID int64) error {
 	ctx, span := s.tracer.Start(ctx, "externalsession.Delete")
 	defer span.End()

pkg/services/auth/authtest/auth_token_service_mock.go

@@ -0,0 +1,384 @@
+// Code generated by mockery v2.42.1. DO NOT EDIT.
+
+package authtest
+
+import (
+	context "context"
+
+	auth "github.com/grafana/grafana/pkg/services/auth"
+
+	mock "github.com/stretchr/testify/mock"
+
+	usertoken "github.com/grafana/grafana/pkg/models/usertoken"
+)
+
+// MockUserAuthTokenService is an autogenerated mock type for the UserTokenService type
+type MockUserAuthTokenService struct {
+	mock.Mock
+}
+
+// ActiveTokenCount provides a mock function with given fields: ctx, userID
+func (_m *MockUserAuthTokenService) ActiveTokenCount(ctx context.Context, userID *int64) (int64, error) {
+	ret := _m.Called(ctx, userID)
+
+	if len(ret) == 0 {
+		panic("no return value specified for ActiveTokenCount")
+	}
+
+	var r0 int64
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, *int64) (int64, error)); ok {
+		return rf(ctx, userID)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, *int64) int64); ok {
+		r0 = rf(ctx, userID)
+	} else {
+		r0 = ret.Get(0).(int64)
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, *int64) error); ok {
+		r1 = rf(ctx, userID)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// CreateToken provides a mock function with given fields: ctx, cmd
+func (_m *MockUserAuthTokenService) CreateToken(ctx context.Context, cmd *auth.CreateTokenCommand) (*usertoken.UserToken, error) {
+	ret := _m.Called(ctx, cmd)
+
+	if len(ret) == 0 {
+		panic("no return value specified for CreateToken")
+	}
+
+	var r0 *usertoken.UserToken
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, *auth.CreateTokenCommand) (*usertoken.UserToken, error)); ok {
+		return rf(ctx, cmd)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, *auth.CreateTokenCommand) *usertoken.UserToken); ok {
+		r0 = rf(ctx, cmd)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*usertoken.UserToken)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, *auth.CreateTokenCommand) error); ok {
+		r1 = rf(ctx, cmd)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// FindExternalSessions provides a mock function with given fields: ctx, query
+func (_m *MockUserAuthTokenService) FindExternalSessions(ctx context.Context, query *auth.ListExternalSessionQuery) ([]*auth.ExternalSession, error) {
+	ret := _m.Called(ctx, query)
+
+	if len(ret) == 0 {
+		panic("no return value specified for FindExternalSessions")
+	}
+
+	var r0 []*auth.ExternalSession
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, *auth.ListExternalSessionQuery) ([]*auth.ExternalSession, error)); ok {
+		return rf(ctx, query)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, *auth.ListExternalSessionQuery) []*auth.ExternalSession); ok {
+		r0 = rf(ctx, query)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).([]*auth.ExternalSession)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, *auth.ListExternalSessionQuery) error); ok {
+		r1 = rf(ctx, query)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// GetExternalSession provides a mock function with given fields: ctx, externalSessionID
+func (_m *MockUserAuthTokenService) GetExternalSession(ctx context.Context, externalSessionID int64) (*auth.ExternalSession, error) {
+	ret := _m.Called(ctx, externalSessionID)
+
+	if len(ret) == 0 {
+		panic("no return value specified for GetExternalSession")
+	}
+
+	var r0 *auth.ExternalSession
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, int64) (*auth.ExternalSession, error)); ok {
+		return rf(ctx, externalSessionID)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, int64) *auth.ExternalSession); ok {
+		r0 = rf(ctx, externalSessionID)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*auth.ExternalSession)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, int64) error); ok {
+		r1 = rf(ctx, externalSessionID)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// GetTokenByExternalSessionID provides a mock function with given fields: ctx, externalSessionID
+func (_m *MockUserAuthTokenService) GetTokenByExternalSessionID(ctx context.Context, externalSessionID int64) (*usertoken.UserToken, error) {
+	ret := _m.Called(ctx, externalSessionID)
+
+	if len(ret) == 0 {
+		panic("no return value specified for GetTokenByExternalSessionID")
+	}
+
+	var r0 *usertoken.UserToken
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, int64) (*usertoken.UserToken, error)); ok {
+		return rf(ctx, externalSessionID)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, int64) *usertoken.UserToken); ok {
+		r0 = rf(ctx, externalSessionID)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*usertoken.UserToken)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, int64) error); ok {
+		r1 = rf(ctx, externalSessionID)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// GetUserRevokedTokens provides a mock function with given fields: ctx, userID
+func (_m *MockUserAuthTokenService) GetUserRevokedTokens(ctx context.Context, userID int64) ([]*usertoken.UserToken, error) {
+	ret := _m.Called(ctx, userID)
+
+	if len(ret) == 0 {
+		panic("no return value specified for GetUserRevokedTokens")
+	}
+
+	var r0 []*usertoken.UserToken
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, int64) ([]*usertoken.UserToken, error)); ok {
+		return rf(ctx, userID)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, int64) []*usertoken.UserToken); ok {
+		r0 = rf(ctx, userID)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).([]*usertoken.UserToken)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, int64) error); ok {
+		r1 = rf(ctx, userID)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// GetUserToken provides a mock function with given fields: ctx, userID, userTokenID
+func (_m *MockUserAuthTokenService) GetUserToken(ctx context.Context, userID int64, userTokenID int64) (*usertoken.UserToken, error) {
+	ret := _m.Called(ctx, userID, userTokenID)
+
+	if len(ret) == 0 {
+		panic("no return value specified for GetUserToken")
+	}
+
+	var r0 *usertoken.UserToken
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, int64, int64) (*usertoken.UserToken, error)); ok {
+		return rf(ctx, userID, userTokenID)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, int64, int64) *usertoken.UserToken); ok {
+		r0 = rf(ctx, userID, userTokenID)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*usertoken.UserToken)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, int64, int64) error); ok {
+		r1 = rf(ctx, userID, userTokenID)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// GetUserTokens provides a mock function with given fields: ctx, userID
+func (_m *MockUserAuthTokenService) GetUserTokens(ctx context.Context, userID int64) ([]*usertoken.UserToken, error) {
+	ret := _m.Called(ctx, userID)
+
+	if len(ret) == 0 {
+		panic("no return value specified for GetUserTokens")
+	}
+
+	var r0 []*usertoken.UserToken
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, int64) ([]*usertoken.UserToken, error)); ok {
+		return rf(ctx, userID)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, int64) []*usertoken.UserToken); ok {
+		r0 = rf(ctx, userID)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).([]*usertoken.UserToken)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, int64) error); ok {
+		r1 = rf(ctx, userID)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// LookupToken provides a mock function with given fields: ctx, unhashedToken
+func (_m *MockUserAuthTokenService) LookupToken(ctx context.Context, unhashedToken string) (*usertoken.UserToken, error) {
+	ret := _m.Called(ctx, unhashedToken)
+
+	if len(ret) == 0 {
+		panic("no return value specified for LookupToken")
+	}
+
+	var r0 *usertoken.UserToken
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, string) (*usertoken.UserToken, error)); ok {
+		return rf(ctx, unhashedToken)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, string) *usertoken.UserToken); ok {
+		r0 = rf(ctx, unhashedToken)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*usertoken.UserToken)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, string) error); ok {
+		r1 = rf(ctx, unhashedToken)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// RevokeAllUserTokens provides a mock function with given fields: ctx, userID
+func (_m *MockUserAuthTokenService) RevokeAllUserTokens(ctx context.Context, userID int64) error {
+	ret := _m.Called(ctx, userID)
+
+	if len(ret) == 0 {
+		panic("no return value specified for RevokeAllUserTokens")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context, int64) error); ok {
+		r0 = rf(ctx, userID)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// RevokeToken provides a mock function with given fields: ctx, token, soft
+func (_m *MockUserAuthTokenService) RevokeToken(ctx context.Context, token *usertoken.UserToken, soft bool) error {
+	ret := _m.Called(ctx, token, soft)
+
+	if len(ret) == 0 {
+		panic("no return value specified for RevokeToken")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context, *usertoken.UserToken, bool) error); ok {
+		r0 = rf(ctx, token, soft)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// RotateToken provides a mock function with given fields: ctx, cmd
+func (_m *MockUserAuthTokenService) RotateToken(ctx context.Context, cmd auth.RotateCommand) (*usertoken.UserToken, error) {
+	ret := _m.Called(ctx, cmd)
+
+	if len(ret) == 0 {
+		panic("no return value specified for RotateToken")
+	}
+
+	var r0 *usertoken.UserToken
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, auth.RotateCommand) (*usertoken.UserToken, error)); ok {
+		return rf(ctx, cmd)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, auth.RotateCommand) *usertoken.UserToken); ok {
+		r0 = rf(ctx, cmd)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*usertoken.UserToken)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, auth.RotateCommand) error); ok {
+		r1 = rf(ctx, cmd)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// UpdateExternalSession provides a mock function with given fields: ctx, externalSessionID, cmd
+func (_m *MockUserAuthTokenService) UpdateExternalSession(ctx context.Context, externalSessionID int64, cmd *auth.UpdateExternalSessionCommand) error {
+	ret := _m.Called(ctx, externalSessionID, cmd)
+
+	if len(ret) == 0 {
+		panic("no return value specified for UpdateExternalSession")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context, int64, *auth.UpdateExternalSessionCommand) error); ok {
+		r0 = rf(ctx, externalSessionID, cmd)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// NewMockUserAuthTokenService creates a new instance of MockUserAuthTokenService. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewMockUserAuthTokenService(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *MockUserAuthTokenService {
+	mock := &MockUserAuthTokenService{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}

pkg/services/auth/authtest/external_session_store_mock.go

@@ -147,6 +147,24 @@ func (_m *MockExternalSessionStore) List(ctx context.Context, query *auth.ListEx
 	return r0, r1
 }
 
+// Update provides a mock function with given fields: ctx, ID, cmd
+func (_m *MockExternalSessionStore) Update(ctx context.Context, ID int64, cmd *auth.UpdateExternalSessionCommand) error {
+	ret := _m.Called(ctx, ID, cmd)
+
+	if len(ret) == 0 {
+		panic("no return value specified for Update")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context, int64, *auth.UpdateExternalSessionCommand) error); ok {
+		r0 = rf(ctx, ID, cmd)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
 // NewMockExternalSessionStore creates a new instance of MockExternalSessionStore. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
 // The first argument is typically a *testing.T value.
 func NewMockExternalSessionStore(t interface {

pkg/services/auth/authtest/testing.go

@@ -14,12 +14,15 @@ import (
 	"github.com/grafana/grafana/pkg/services/login"
 )
 
+var _ auth.UserTokenService = (*FakeUserAuthTokenService)(nil)
+
 type FakeUserAuthTokenService struct {
 	CreateTokenProvider                 func(ctx context.Context, cmd *auth.CreateTokenCommand) (*auth.UserToken, error)
 	RotateTokenProvider                 func(ctx context.Context, cmd auth.RotateCommand) (*auth.UserToken, error)
 	GetTokenByExternalSessionIDProvider func(ctx context.Context, externalSessionID int64) (*auth.UserToken, error)
 	GetExternalSessionProvider          func(ctx context.Context, externalSessionID int64) (*auth.ExternalSession, error)
 	FindExternalSessionsProvider        func(ctx context.Context, query *auth.ListExternalSessionQuery) ([]*auth.ExternalSession, error)
+	UpdateExternalSessionProvider       func(ctx context.Context, externalSessionID int64, cmd *auth.UpdateExternalSessionCommand) error
 	TryRotateTokenProvider              func(ctx context.Context, token *auth.UserToken, clientIP net.IP, userAgent string) (bool, *auth.UserToken, error)
 	LookupTokenProvider                 func(ctx context.Context, unhashedToken string) (*auth.UserToken, error)
 	RevokeTokenProvider                 func(ctx context.Context, token *auth.UserToken, soft bool) error
@@ -98,6 +101,10 @@ func (s *FakeUserAuthTokenService) FindExternalSessions(ctx context.Context, que
 	return s.FindExternalSessionsProvider(context.Background(), query)
 }
 
+func (s *FakeUserAuthTokenService) UpdateExternalSession(ctx context.Context, externalSessionID int64, cmd *auth.UpdateExternalSessionCommand) error {
+	return s.UpdateExternalSessionProvider(context.Background(), externalSessionID, cmd)
+}
+
 func (s *FakeUserAuthTokenService) LookupToken(ctx context.Context, unhashedToken string) (*auth.UserToken, error) {
 	return s.LookupTokenProvider(context.Background(), unhashedToken)
 }
@@ -149,16 +156,6 @@ func (ts *FakeOAuthTokenService) IsOAuthPassThruEnabled(*datasources.DataSource)
 	return ts.passThruEnabled
 }
 
-func (ts *FakeOAuthTokenService) HasOAuthEntry(context.Context, identity.Requester) (*login.UserAuth, bool, error) {
-	if ts.ExpectedAuthUser != nil {
-		return ts.ExpectedAuthUser, true, nil
-	}
-	if error, ok := ts.ExpectedErrors["HasOAuthEntry"]; ok {
-		return nil, false, error
-	}
-	return nil, false, nil
-}
-
 func (ts *FakeOAuthTokenService) InvalidateOAuthTokens(ctx context.Context, usr *login.UserAuth) error {
 	ts.ExpectedAuthUser.OAuthAccessToken = ""
 	ts.ExpectedAuthUser.OAuthRefreshToken = ""

pkg/services/auth/external_session.go

@@ -3,6 +3,8 @@ package auth
 import (
 	"context"
 	"time"
+
+	"golang.org/x/oauth2"
 )
 
 type ExternalSession struct {
@@ -43,6 +45,10 @@ func (e *ExternalSession) Clone() *ExternalSession {
 	}
 }
 
+type UpdateExternalSessionCommand struct {
+	Token *oauth2.Token
+}
+
 type ListExternalSessionQuery struct {
 	ID        int64
 	NameID    string
@@ -57,6 +63,8 @@ type ExternalSessionStore interface {
 	List(ctx context.Context, query *ListExternalSessionQuery) ([]*ExternalSession, error)
 	// Create creates a new external session for a user
 	Create(ctx context.Context, extSesion *ExternalSession) error
+	// Update updates an external session
+	Update(ctx context.Context, ID int64, cmd *UpdateExternalSessionCommand) error
 	// Delete deletes an external session
 	Delete(ctx context.Context, ID int64) error
 	// DeleteExternalSessionsByUserID deletes an external session

pkg/services/authn/authn.go

@@ -178,7 +178,7 @@ type RedirectClient interface {
 // that should happen during logout and supports client specific redirect URL.
 type LogoutClient interface {
 	Client
-	Logout(ctx context.Context, user identity.Requester) (*Redirect, bool)
+	Logout(ctx context.Context, user identity.Requester, sessionToken *usertoken.UserToken) (*Redirect, bool)
 }
 
 type SSOSettingsAwareClient interface {

pkg/services/authn/authnimpl/registration.go

@@ -37,7 +37,7 @@ func ProvideRegistration(
 	jwtService auth.JWTVerifierService, userProtectionService login.UserProtectionService,
 	loginAttempts loginattempt.Service, quotaService quota.Service,
 	authInfoService login.AuthInfoService, renderService rendering.Service,
-	features *featuremgmt.FeatureManager, oauthTokenService oauthtoken.OAuthTokenService,
+	features featuremgmt.FeatureToggles, oauthTokenService oauthtoken.OAuthTokenService,
 	socialService social.Service, cache *remotecache.RemoteCache,
 	ldapService service.LDAP, settingsProviderService setting.Provider,
 	tracer tracing.Tracer, tempUserService tempuser.Service, notificationService notifications.Service,
@@ -108,13 +108,13 @@ func ProvideRegistration(
 	}
 
 	// FIXME (jguer): move to User package
-	userSync := sync.ProvideUserSync(userService, userProtectionService, authInfoService, quotaService, tracer)
+	userSync := sync.ProvideUserSync(userService, userProtectionService, authInfoService, quotaService, tracer, features)
 	orgSync := sync.ProvideOrgSync(userService, orgService, accessControlService, cfg, tracer)
 	authnSvc.RegisterPostAuthHook(userSync.SyncUserHook, 10)
 	authnSvc.RegisterPostAuthHook(userSync.EnableUserHook, 20)
 	authnSvc.RegisterPostAuthHook(orgSync.SyncOrgRolesHook, 30)
 	authnSvc.RegisterPostAuthHook(userSync.SyncLastSeenHook, 130)
-	authnSvc.RegisterPostAuthHook(sync.ProvideOAuthTokenSync(oauthTokenService, sessionService, socialService, tracer).SyncOauthTokenHook, 60)
+	authnSvc.RegisterPostAuthHook(sync.ProvideOAuthTokenSync(oauthTokenService, sessionService, socialService, tracer, features).SyncOauthTokenHook, 60)
 	authnSvc.RegisterPostAuthHook(userSync.FetchSyncedUserHook, 100)
 
 	rbacSync := sync.ProvideRBACSync(accessControlService, tracer, permRegistry)

pkg/services/authn/authnimpl/service.go

@@ -322,7 +322,7 @@ func (s *Service) Logout(ctx context.Context, user identity.Requester, sessionTo
 			goto Default
 		}
 
-		clientRedirect, ok := logoutClient.Logout(ctx, user)
+		clientRedirect, ok := logoutClient.Logout(ctx, user, sessionToken)
 		if !ok {
 			goto Default
 		}

pkg/services/authn/authnimpl/service_test.go

@@ -507,7 +507,7 @@ func TestService_Logout(t *testing.T) {
 			expectedRedirect: &authn.Redirect{URL: "http://idp.com/logout"},
 			client: &authntest.MockClient{
 				NameFunc: func() string { return "auth.client.azuread" },
-				LogoutFunc: func(ctx context.Context, _ identity.Requester) (*authn.Redirect, bool) {
+				LogoutFunc: func(ctx context.Context, _ identity.Requester, sessionToken *usertoken.UserToken) (*authn.Redirect, bool) {
 					return &authn.Redirect{URL: "http://idp.com/logout"}, true
 				},
 			},

pkg/services/authn/authnimpl/sync/oauth_token_sync.go

@@ -17,12 +17,15 @@ import (
 	"github.com/grafana/grafana/pkg/login/social"
 	"github.com/grafana/grafana/pkg/services/auth"
 	"github.com/grafana/grafana/pkg/services/authn"
+	"github.com/grafana/grafana/pkg/services/featuremgmt"
 	"github.com/grafana/grafana/pkg/services/oauthtoken"
 )
 
 const maxOAuthTokenCacheTTL = 5 * time.Minute
 
-func ProvideOAuthTokenSync(service oauthtoken.OAuthTokenService, sessionService auth.UserTokenService, socialService social.Service, tracer tracing.Tracer) *OAuthTokenSync {
+func ProvideOAuthTokenSync(service oauthtoken.OAuthTokenService, sessionService auth.UserTokenService, socialService social.Service, tracer tracing.Tracer,
+	features featuremgmt.FeatureToggles,
+) *OAuthTokenSync {
 	return &OAuthTokenSync{
 		log.New("oauth_token.sync"),
 		service,
@@ -31,6 +34,7 @@ func ProvideOAuthTokenSync(service oauthtoken.OAuthTokenService, sessionService
 		new(singleflight.Group),
 		tracer,
 		localcache.New(maxOAuthTokenCacheTTL, 15*time.Minute),
+		features,
 	}
 }
 
@@ -42,6 +46,7 @@ type OAuthTokenSync struct {
 	singleflightGroup *singleflight.Group
 	tracer            tracing.Tracer
 	cache             *localcache.CacheService
+	features          featuremgmt.FeatureToggles
 }
 
 func (s *OAuthTokenSync) SyncOauthTokenHook(ctx context.Context, id *authn.Identity, _ *authn.Request) error {
@@ -72,6 +77,10 @@ func (s *OAuthTokenSync) SyncOauthTokenHook(ctx context.Context, id *authn.Ident
 	ctxLogger := s.log.FromContext(ctx).New("userID", userID)
 
 	cacheKey := fmt.Sprintf("token-check-%s", id.GetID())
+	if s.features.IsEnabledGlobally(featuremgmt.FlagImprovedExternalSessionHandling) {
+		cacheKey = fmt.Sprintf("token-check-%s-%d", id.GetID(), id.SessionToken.Id)
+	}
+
 	if _, ok := s.cache.Get(cacheKey); ok {
 		ctxLogger.Debug("Expiration check has been cached, no need to refresh")
 		return nil
@@ -83,7 +92,7 @@ func (s *OAuthTokenSync) SyncOauthTokenHook(ctx context.Context, id *authn.Ident
 		updateCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 15*time.Second)
 		defer cancel()
 
-		token, refreshErr := s.service.TryTokenRefresh(updateCtx, id)
+		token, refreshErr := s.service.TryTokenRefresh(updateCtx, id, id.SessionToken)
 		if refreshErr != nil {
 			if errors.Is(refreshErr, context.Canceled) {
 				return nil, nil

pkg/services/authn/authnimpl/sync/oauth_token_sync_test.go

@@ -20,6 +20,9 @@ import (
 	"github.com/grafana/grafana/pkg/services/auth"
 	"github.com/grafana/grafana/pkg/services/auth/authtest"
 	"github.com/grafana/grafana/pkg/services/authn"
+	"github.com/grafana/grafana/pkg/services/contexthandler/ctxkey"
+	contextmodel "github.com/grafana/grafana/pkg/services/contexthandler/model"
+	"github.com/grafana/grafana/pkg/services/featuremgmt"
 	"github.com/grafana/grafana/pkg/services/login"
 	"github.com/grafana/grafana/pkg/services/oauthtoken/oauthtokentest"
 )
@@ -85,7 +88,7 @@ func TestOAuthTokenSync_SyncOAuthTokenHook(t *testing.T) {
 			)
 
 			service := &oauthtokentest.MockOauthTokenService{
-				TryTokenRefreshFunc: func(ctx context.Context, usr identity.Requester) (*oauth2.Token, error) {
+				TryTokenRefreshFunc: func(ctx context.Context, usr identity.Requester, _ *auth.UserToken) (*oauth2.Token, error) {
 					tryRefreshCalled = true
 					return nil, tt.expectedTryRefreshErr
 				},
@@ -116,9 +119,13 @@ func TestOAuthTokenSync_SyncOAuthTokenHook(t *testing.T) {
 				singleflightGroup: new(singleflight.Group),
 				tracer:            tracing.InitializeTracerForTest(),
 				cache:             localcache.New(maxOAuthTokenCacheTTL, 15*time.Minute),
+				features:          featuremgmt.WithFeatures(),
 			}
 
-			err := sync.SyncOauthTokenHook(context.Background(), tt.identity, nil)
+			ctx := context.Background()
+			reqCtx := context.WithValue(ctx, ctxkey.Key{}, &contextmodel.ReqContext{UserToken: nil})
+
+			err := sync.SyncOauthTokenHook(reqCtx, tt.identity, nil)
 			assert.ErrorIs(t, err, tt.expectedErr)
 			assert.Equal(t, tt.expectTryRefreshTokenCalled, tryRefreshCalled)
 			assert.Equal(t, tt.expectRevokeTokenCalled, revokeTokenCalled)

pkg/services/authn/authnimpl/sync/user_sync.go

@@ -12,6 +12,7 @@ import (
 	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/infra/tracing"
 	"github.com/grafana/grafana/pkg/services/authn"
+	"github.com/grafana/grafana/pkg/services/featuremgmt"
 	"github.com/grafana/grafana/pkg/services/login"
 	"github.com/grafana/grafana/pkg/services/org"
 	"github.com/grafana/grafana/pkg/services/quota"
@@ -51,7 +52,9 @@ var (
 	errSignupNotAllowed  = errors.New("system administrator has disabled signup")
 )
 
-func ProvideUserSync(userService user.Service, userProtectionService login.UserProtectionService, authInfoService login.AuthInfoService, quotaService quota.Service, tracer tracing.Tracer) *UserSync {
+func ProvideUserSync(userService user.Service, userProtectionService login.UserProtectionService, authInfoService login.AuthInfoService,
+	quotaService quota.Service, tracer tracing.Tracer, features featuremgmt.FeatureToggles,
+) *UserSync {
 	return &UserSync{
 		userService:           userService,
 		authInfoService:       authInfoService,
@@ -59,6 +62,7 @@ func ProvideUserSync(userService user.Service, userProtectionService login.UserP
 		quotaService:          quotaService,
 		log:                   log.New("user.sync"),
 		tracer:                tracer,
+		features:              features,
 	}
 }
 
@@ -69,6 +73,7 @@ type UserSync struct {
 	quotaService          quota.Service
 	log                   log.Logger
 	tracer                tracing.Tracer
+	features              featuremgmt.FeatureToggles
 }
 
 // SyncUserHook syncs a user with the database
@@ -223,21 +228,30 @@ func (s *UserSync) upsertAuthConnection(ctx context.Context, userID int64, ident
 	// This can happen when: using multiple auth client where the same user exists in several or
 	// changing to new auth client
 	if createConnection {
-		return s.authInfoService.SetAuthInfo(ctx, &login.SetAuthInfoCommand{
+		setAuthInfoCmd := &login.SetAuthInfoCommand{
 			UserId:     userID,
 			AuthModule: identity.AuthenticatedBy,
 			AuthId:     identity.AuthID,
-			OAuthToken: identity.OAuthToken,
-		})
+		}
+
+		if !s.features.IsEnabledGlobally(featuremgmt.FlagImprovedExternalSessionHandling) {
+			setAuthInfoCmd.OAuthToken = identity.OAuthToken
+		}
+		return s.authInfoService.SetAuthInfo(ctx, setAuthInfoCmd)
 	}
 
-	s.log.FromContext(ctx).Debug("Updating auth connection for user", "id", identity.ID)
-	return s.authInfoService.UpdateAuthInfo(ctx, &login.UpdateAuthInfoCommand{
+	updateAuthInfoCmd := &login.UpdateAuthInfoCommand{
 		UserId:     userID,
 		AuthId:     identity.AuthID,
 		AuthModule: identity.AuthenticatedBy,
-		OAuthToken: identity.OAuthToken,
-	})
+	}
+
+	if !s.features.IsEnabledGlobally(featuremgmt.FlagImprovedExternalSessionHandling) {
+		updateAuthInfoCmd.OAuthToken = identity.OAuthToken
+	}
+
+	s.log.FromContext(ctx).Debug("Updating auth connection for user", "id", identity.ID)
+	return s.authInfoService.UpdateAuthInfo(ctx, updateAuthInfoCmd)
 }
 
 func (s *UserSync) updateUserAttributes(ctx context.Context, usr *user.User, id *authn.Identity, userAuth *login.UserAuth) error {

pkg/services/authn/authnimpl/sync/user_sync_test.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/grafana/grafana/pkg/infra/tracing"
 	"github.com/grafana/grafana/pkg/services/authn"
+	"github.com/grafana/grafana/pkg/services/featuremgmt"
 	"github.com/grafana/grafana/pkg/services/login"
 	"github.com/grafana/grafana/pkg/services/login/authinfoimpl"
 	"github.com/grafana/grafana/pkg/services/login/authinfotest"
@@ -45,7 +46,9 @@ func TestUserSync_SyncUserHook(t *testing.T) {
 			AuthModule: "oauth",
 			AuthId:     "2032",
 			UserId:     1,
-			Id:         1}}
+			Id:         1,
+		},
+	}
 
 	userService := &usertest.FakeUserService{ExpectedUser: &user.User{
 		ID:    1,
@@ -434,7 +437,7 @@ func TestUserSync_SyncUserHook(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			s := ProvideUserSync(tt.fields.userService, userProtection, tt.fields.authInfoService, tt.fields.quotaService, tracing.InitializeTracerForTest())
+			s := ProvideUserSync(tt.fields.userService, userProtection, tt.fields.authInfoService, tt.fields.quotaService, tracing.InitializeTracerForTest(), featuremgmt.WithFeatures())
 			err := s.SyncUserHook(tt.args.ctx, tt.args.id, nil)
 			if tt.wantErr {
 				require.Error(t, err)

pkg/services/authn/authntest/mock.go

@@ -85,7 +85,7 @@ type MockClient struct {
 	TestFunc            func(ctx context.Context, r *authn.Request) bool
 	PriorityFunc        func() uint
 	HookFunc            func(ctx context.Context, identity *authn.Identity, r *authn.Request) error
-	LogoutFunc          func(ctx context.Context, user identity.Requester) (*authn.Redirect, bool)
+	LogoutFunc          func(ctx context.Context, user identity.Requester, sessionToken *usertoken.UserToken) (*authn.Redirect, bool)
 	IdentityTypeFunc    func() claims.IdentityType
 	ResolveIdentityFunc func(ctx context.Context, orgID int64, typ claims.IdentityType, id string) (*authn.Identity, error)
 }
@@ -133,9 +133,9 @@ func (m MockClient) Hook(ctx context.Context, identity *authn.Identity, r *authn
 	return nil
 }
 
-func (m *MockClient) Logout(ctx context.Context, user identity.Requester) (*authn.Redirect, bool) {
+func (m *MockClient) Logout(ctx context.Context, user identity.Requester, sessionToken *usertoken.UserToken) (*authn.Redirect, bool) {
 	if m.LogoutFunc != nil {
-		return m.LogoutFunc(ctx, user)
+		return m.LogoutFunc(ctx, user, sessionToken)
 	}
 	return nil, false
 }

pkg/services/authn/clients/oauth.go

@@ -18,6 +18,7 @@ import (
 	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/login/social"
 	"github.com/grafana/grafana/pkg/login/social/connectors"
+	"github.com/grafana/grafana/pkg/services/auth"
 	"github.com/grafana/grafana/pkg/services/authn"
 	"github.com/grafana/grafana/pkg/services/featuremgmt"
 	"github.com/grafana/grafana/pkg/services/login"
@@ -257,8 +258,8 @@ func (c *OAuth) RedirectURL(ctx context.Context, r *authn.Request) (*authn.Redir
 	}, nil
 }
 
-func (c *OAuth) Logout(ctx context.Context, user identity.Requester) (*authn.Redirect, bool) {
-	token := c.oauthService.GetCurrentOAuthToken(ctx, user)
+func (c *OAuth) Logout(ctx context.Context, user identity.Requester, sessionToken *auth.UserToken) (*authn.Redirect, bool) {
+	token := c.oauthService.GetCurrentOAuthToken(ctx, user, sessionToken)
 
 	userID, err := identity.UserIdentifier(user.GetID())
 	if err != nil {
@@ -268,7 +269,7 @@ func (c *OAuth) Logout(ctx context.Context, user identity.Requester) (*authn.Red
 
 	ctxLogger := c.log.FromContext(ctx).New("userID", userID)
 
-	if err := c.oauthService.InvalidateOAuthTokens(ctx, user); err != nil {
+	if err := c.oauthService.InvalidateOAuthTokens(ctx, user, sessionToken); err != nil {
 		ctxLogger.Error("Failed to invalidate tokens", "error", err)
 	}
 

pkg/services/authn/clients/oauth_test.go

@@ -18,6 +18,7 @@ import (
 	"github.com/grafana/grafana/pkg/apimachinery/identity"
 	"github.com/grafana/grafana/pkg/login/social"
 	"github.com/grafana/grafana/pkg/login/social/socialtest"
+	"github.com/grafana/grafana/pkg/services/auth"
 	"github.com/grafana/grafana/pkg/services/authn"
 	"github.com/grafana/grafana/pkg/services/featuremgmt"
 	"github.com/grafana/grafana/pkg/services/login"
@@ -469,7 +470,7 @@ func TestOAuth_Logout(t *testing.T) {
 			)
 
 			mockService := &oauthtokentest.MockOauthTokenService{
-				GetCurrentOauthTokenFunc: func(_ context.Context, _ identity.Requester) *oauth2.Token {
+				GetCurrentOauthTokenFunc: func(_ context.Context, _ identity.Requester, _ *auth.UserToken) *oauth2.Token {
 					getTokenCalled = true
 					token := &oauth2.Token{
 						AccessToken: "some.access.token",
@@ -479,7 +480,7 @@ func TestOAuth_Logout(t *testing.T) {
 						"id_token": "some.id.token",
 					})
 				},
-				InvalidateOAuthTokensFunc: func(_ context.Context, _ identity.Requester) error {
+				InvalidateOAuthTokensFunc: func(_ context.Context, _ identity.Requester, _ *auth.UserToken) error {
 					invalidateTokenCalled = true
 					return nil
 				},
@@ -490,7 +491,7 @@ func TestOAuth_Logout(t *testing.T) {
 			}
 			c := ProvideOAuth(authn.ClientWithPrefix("azuread"), tt.cfg, mockService, fakeSocialSvc, &setting.OSSImpl{Cfg: tt.cfg}, featuremgmt.WithFeatures())
 
-			redirect, ok := c.Logout(context.Background(), &authn.Identity{ID: "1", Type: claims.TypeUser})
+			redirect, ok := c.Logout(context.Background(), &authn.Identity{ID: "1", Type: claims.TypeUser}, nil)
 
 			assert.Equal(t, tt.expectedOK, ok)
 			if tt.expectedOK {

pkg/services/authn/grpcutils/config.go

@@ -3,6 +3,8 @@ package grpcutils
 import (
 	"fmt"
 
+	"github.com/spf13/pflag"
+
 	"github.com/grafana/grafana/pkg/setting"
 )
 
@@ -26,6 +28,11 @@ type GrpcServerConfig struct {
 	AllowedAudiences []string
 	Mode             Mode
 	LegacyFallback   bool
+	AllowInsecure    bool
+}
+
+func (c *GrpcServerConfig) AddFlags(fs *pflag.FlagSet) {
+	fs.StringVar(&c.SigningKeysURL, "grpc-server-authentication.signing-keys-url", "", "gRPC server authentication signing keys URL")
 }
 
 func ReadGrpcServerConfig(cfg *setting.Cfg) (*GrpcServerConfig, error) {
@@ -41,6 +48,7 @@ func ReadGrpcServerConfig(cfg *setting.Cfg) (*GrpcServerConfig, error) {
 		AllowedAudiences: section.Key("allowed_audiences").Strings(","),
 		Mode:             mode,
 		LegacyFallback:   section.Key("legacy_fallback").MustBool(true),
+		AllowInsecure:    cfg.Env == setting.Dev,
 	}, nil
 }
 

pkg/services/authn/grpcutils/grpc_authenticator.go

@@ -27,11 +27,7 @@ func NewInProcGrpcAuthenticator() *authnlib.GrpcAuthenticator {
 	)
 }
 
-func NewGrpcAuthenticator(cfg *setting.Cfg, tracer tracing.Tracer) (*authnlib.GrpcAuthenticator, error) {
-	authCfg, err := ReadGrpcServerConfig(cfg)
-	if err != nil {
-		return nil, err
-	}
+func NewGrpcAuthenticator(authCfg *GrpcServerConfig, tracer tracing.Tracer) (*authnlib.GrpcAuthenticator, error) {
 	grpcAuthCfg := authnlib.GrpcAuthenticatorConfig{
 		KeyRetrieverConfig: authnlib.KeyRetrieverConfig{
 			SigningKeysURL: authCfg.SigningKeysURL,
@@ -42,7 +38,7 @@ func NewGrpcAuthenticator(cfg *setting.Cfg, tracer tracing.Tracer) (*authnlib.Gr
 	}
 
 	client := http.DefaultClient
-	if cfg.Env == setting.Dev {
+	if authCfg.AllowInsecure {
 		// allow insecure connections in development mode to facilitate testing
 		client = &http.Client{Transport: &http.Transport{TLSClientConfig: &tls.Config{InsecureSkipVerify: true}}}
 	}
@@ -87,7 +83,7 @@ func NewGrpcAuthenticatorWithFallback(cfg *setting.Cfg, reg prometheus.Registere
 		return nil, err
 	}
 
-	authenticator, err := NewGrpcAuthenticator(cfg, tracer)
+	authenticator, err := NewGrpcAuthenticator(authCfg, tracer)
 	if err != nil {
 		return nil, err
 	}
@@ -111,42 +107,49 @@ func FallbackUsed(ctx context.Context) bool {
 func (f *AuthenticatorWithFallback) Authenticate(ctx context.Context) (context.Context, error) {
 	ctx, span := f.tracer.Start(ctx, "grpcutils.AuthenticatorWithFallback.Authenticate")
 	defer span.End()
-	span.SetAttributes(attribute.Bool("fallback_used", false))
+
 	// Try to authenticate with the new authenticator first
+	span.SetAttributes(attribute.Bool("fallback_used", false))
 	newCtx, err := f.authenticator.Authenticate(ctx)
-	if err != nil {
-		// In case of error, fallback to the legacy authenticator
-		newCtx, err = f.fallback.Authenticate(ctx)
-		f.metrics.fallbackCounter.WithLabelValues(fmt.Sprintf("%t", err == nil)).Inc()
-		span.SetAttributes(attribute.Bool("fallback_used", true))
+	if err == nil {
+		// fallback not used, authentication successful
+		f.metrics.requestsTotal.WithLabelValues("false", "true").Inc()
+		return newCtx, nil
+	}
+
+	// In case of error, fallback to the legacy authenticator
+	span.SetAttributes(attribute.Bool("fallback_used", true))
+	newCtx, err = f.fallback.Authenticate(ctx)
+	if newCtx != nil {
 		newCtx = context.WithValue(newCtx, contextFallbackKey{}, true)
 	}
+	f.metrics.requestsTotal.WithLabelValues("true", fmt.Sprintf("%t", err == nil)).Inc()
 	return newCtx, err
 }
 
 const (
 	metricsNamespace = "grafana"
-	metricsSubSystem = "grpc_authenticator"
+	metricsSubSystem = "grpc_authenticator_with_fallback"
 )
 
 type metrics struct {
-	fallbackCounter *prometheus.CounterVec
+	requestsTotal *prometheus.CounterVec
 }
 
 func newMetrics(reg prometheus.Registerer) *metrics {
 	m := &metrics{
-		fallbackCounter: prometheus.NewCounterVec(
+		requestsTotal: prometheus.NewCounterVec(
 			prometheus.CounterOpts{
 				Namespace: metricsNamespace,
 				Subsystem: metricsSubSystem,
-				Name:      "fallback_total",
-				Help:      "Number of times the fallback authenticator was used",
-			}, []string{"result"}),
+				Name:      "requests_total",
+				Help:      "Number requests using the authenticator with fallback",
+			}, []string{"fallback_used", "result"}),
 	}
 
 	if reg != nil {
 		once.Do(func() {
-			reg.MustRegister(m.fallbackCounter)
+			reg.MustRegister(m.requestsTotal)
 		})
 	}
 

pkg/services/authz/zanzana/common/info.go

@@ -7,14 +7,15 @@ import (
 )
 
 type TypeInfo struct {
-	Type string
+	Type      string
+	Relations []string
 }
 
 var typedResources = map[string]TypeInfo{
 	FormatGroupResource(
 		folderalpha1.FolderResourceInfo.GroupResource().Group,
 		folderalpha1.FolderResourceInfo.GroupResource().Resource,
-	): {Type: "folder"},
+	): {Type: "folder", Relations: append(ResourceRelations, RelationCreate)},
 }
 
 func GetTypeInfo(group, resource string) (TypeInfo, bool) {

pkg/services/authz/zanzana/common/tuple.go

@@ -10,12 +10,13 @@ import (
 )
 
 const (
-	TypeUser      string = "user"
-	TypeTeam      string = "team"
-	TypeRole      string = "role"
-	TypeFolder    string = "folder"
-	TypeResource  string = "resource"
-	TypeNamespace string = "namespace"
+	TypeUser           string = "user"
+	TypeServiceAccount string = "service-account"
+	TypeTeam           string = "team"
+	TypeRole           string = "role"
+	TypeFolder         string = "folder"
+	TypeResource       string = "resource"
+	TypeNamespace      string = "namespace"
 )
 
 const (
@@ -50,7 +51,6 @@ const (
 var ResourceRelations = []string{
 	RelationRead,
 	RelationWrite,
-	RelationCreate,
 	RelationDelete,
 	RelationPermissionsRead,
 	RelationPermissionsWrite,
@@ -58,6 +58,7 @@ var ResourceRelations = []string{
 
 var FolderRelations = append(
 	ResourceRelations,
+	RelationCreate,
 	RelationFolderResourceRead,
 	RelationFolderResourceWrite,
 	RelationFolderResourceCreate,

pkg/services/authz/zanzana/schema/schema_core.fga

@@ -2,28 +2,30 @@ module core
 
 type namespace
   relations
-    define view: [user, team#member, role#assignee] or edit
-    define edit: [user, team#member, role#assignee] or admin
-    define admin: [user, team#member, role#assignee]
+    define view: [user, service-account, team#member, role#assignee] or edit
+    define edit: [user, service-account, team#member, role#assignee] or admin
+    define admin: [user, service-account, team#member, role#assignee]
 
-    define read: [user, team#member, role#assignee] or view
-    define create: [user, team#member, role#assignee] or edit
-    define write: [user, team#member, role#assignee] or edit
-    define delete: [user, team#member, role#assignee] or edit
-    define permissions_read: [user, team#member, role#assignee] or admin
-    define permissions_write: [user, team#member, role#assignee] or admin
+    define read: [user, service-account, team#member, role#assignee] or view
+    define create: [user, service-account, team#member, role#assignee] or edit
+    define write: [user, service-account, team#member, role#assignee] or edit
+    define delete: [user, service-account, team#member, role#assignee] or edit
+    define permissions_read: [user, service-account, team#member, role#assignee] or admin
+    define permissions_write: [user, service-account, team#member, role#assignee] or admin
 
 type user
 
+type service-account
+
 type role
   relations
-    define assignee: [user, team#member, role#assignee]
+    define assignee: [user, service-account, team#member, role#assignee]
 
 type team
   relations
     # Action sets
-    define admin: [user]
-    define member: [user] or admin
+    define admin: [user, service-account]
+    define member: [user, service-account] or admin
 
     define read: [role#assignee] or member
     define write: [role#assignee] or admin

pkg/services/authz/zanzana/schema/schema_folder.fga

@@ -5,13 +5,13 @@ type folder
     define parent: [folder]
 
     # Action sets
-    define view: [user, team#member, role#assignee] or edit or view from parent
-    define edit: [user, team#member, role#assignee] or admin or edit from parent
-    define admin: [user, team#member, role#assignee] or admin from parent
+    define view: [user, service-account, team#member, role#assignee] or edit or view from parent
+    define edit: [user, service-account, team#member, role#assignee] or admin or edit from parent
+    define admin: [user, service-account, team#member, role#assignee] or admin from parent
 
-    define read: [user, team#member, role#assignee] or view or read from parent
-    define create: [user, team#member, role#assignee] or edit or create from parent
-    define write: [user, team#member, role#assignee] or edit or write from parent
-    define delete: [user, team#member, role#assignee] or edit or delete from parent
-    define permissions_read: [user, team#member, role#assignee] or admin or permissions_read from parent
-    define permissions_write: [user, team#member, role#assignee] or admin or permissions_write from parent
+    define read: [user, service-account, team#member, role#assignee] or view or read from parent
+    define create: [user, service-account, team#member, role#assignee] or edit or create from parent
+    define write: [user, service-account, team#member, role#assignee] or edit or write from parent
+    define delete: [user, service-account, team#member, role#assignee] or edit or delete from parent
+    define permissions_read: [user, service-account, team#member, role#assignee] or admin or permissions_read from parent
+    define permissions_write: [user, service-account, team#member, role#assignee] or admin or permissions_write from parent

pkg/services/authz/zanzana/schema/schema_resource.fga

@@ -2,29 +2,28 @@ module resource
 
 extend type folder
   relations
-    define resource_view: [user, team#member, role#assignee] or resource_edit or resource_view from parent
-    define resource_edit: [user, team#member, role#assignee] or resource_admin or resource_edit from parent
-    define resource_admin: [user, team#member, role#assignee] or resource_admin from parent
+    define resource_view: [user, service-account, team#member, role#assignee] or resource_edit or resource_view from parent
+    define resource_edit: [user, service-account, team#member, role#assignee] or resource_admin or resource_edit from parent
+    define resource_admin: [user, service-account, team#member, role#assignee] or resource_admin from parent
 
-    define resource_read: [user with folder_group_filter, team#member with folder_group_filter, role#assignee with folder_group_filter] or resource_view or resource_read from parent
-    define resource_create: [user with folder_group_filter, team#member with folder_group_filter, role#assignee with folder_group_filter] or resource_edit or resource_create from parent
-    define resource_write: [user with folder_group_filter, team#member with folder_group_filter, role#assignee with folder_group_filter] or resource_edit or resource_write from parent
-    define resource_delete: [user with folder_group_filter, team#member with folder_group_filter, role#assignee with folder_group_filter] or resource_edit or resource_delete from parent
-    define resource_permissions_read: [user with folder_group_filter, team#member with folder_group_filter, role#assignee with folder_group_filter] or resource_admin or resource_permissions_read from parent
-    define resource_permissions_write: [user with folder_group_filter, team#member with folder_group_filter, role#assignee with folder_group_filter] or resource_admin or resource_permissions_write from parent
+    define resource_read: [user with folder_group_filter, service-account with folder_group_filter, team#member with folder_group_filter, role#assignee with folder_group_filter] or resource_view or resource_read from parent
+    define resource_create: [user with folder_group_filter, service-account with folder_group_filter, team#member with folder_group_filter, role#assignee with folder_group_filter] or resource_edit or resource_create from parent
+    define resource_write: [user with folder_group_filter, service-account with folder_group_filter, team#member with folder_group_filter, role#assignee with folder_group_filter] or resource_edit or resource_write from parent
+    define resource_delete: [user with folder_group_filter, service-account with folder_group_filter, team#member with folder_group_filter, role#assignee with folder_group_filter] or resource_edit or resource_delete from parent
+    define resource_permissions_read: [user with folder_group_filter, service-account with folder_group_filter, team#member with folder_group_filter, role#assignee with folder_group_filter] or resource_admin or resource_permissions_read from parent
+    define resource_permissions_write: [user with folder_group_filter, service-account with folder_group_filter, team#member with folder_group_filter, role#assignee with folder_group_filter] or resource_admin or resource_permissions_write from parent
 
 type resource
   relations
-    define view: [user with group_filter, team#member with group_filter, role#assignee with group_filter] or edit
-    define edit: [user  with group_filter, team#member with group_filter, role#assignee with group_filter] or admin
-    define admin: [user with group_filter, team#member with group_filter, role#assignee with group_filter]
+    define view: [user with group_filter, service-account with group_filter, team#member with group_filter, role#assignee with group_filter] or edit
+    define edit: [user with group_filter, service-account with group_filter, team#member with group_filter, role#assignee with group_filter] or admin
+    define admin: [user with group_filter, service-account with group_filter, team#member with group_filter, role#assignee with group_filter]
 
-    define read: [user with group_filter, team#member with group_filter, role#assignee with group_filter] or view
-    define create: [user with group_filter, team#member with group_filter, role#assignee with group_filter] or edit
-    define write: [user with group_filter, team#member with group_filter, role#assignee with group_filter] or edit
-    define delete: [user with group_filter, team#member with group_filter, role#assignee with group_filter] or edit
-    define permissions_read: [user with group_filter, team#member with group_filter, role#assignee with group_filter] or admin
-    define permissions_write: [user with group_filter, team#member with group_filter, role#assignee with group_filter] or admin
+    define read: [user with group_filter, service-account with group_filter, team#member with group_filter, role#assignee with group_filter] or view
+    define write: [user with group_filter, service-account with group_filter, team#member with group_filter, role#assignee with group_filter] or edit
+    define delete: [user with group_filter, service-account with group_filter, team#member with group_filter, role#assignee with group_filter] or edit
+    define permissions_read: [user with group_filter, service-account with group_filter, team#member with group_filter, role#assignee with group_filter] or admin
+    define permissions_write: [user with group_filter, service-account with group_filter, team#member with group_filter, role#assignee with group_filter] or admin
 
 condition group_filter(requested_group: string, group_resource: string) {
   requested_group == group_resource

pkg/services/authz/zanzana/server/server.go

@@ -2,12 +2,15 @@ package server
 
 import (
 	"sync"
+	"time"
 
+	"github.com/fullstorydev/grpchan/inprocgrpc"
 	authzv1 "github.com/grafana/authlib/authz/proto/v1"
 	openfgav1 "github.com/openfga/api/proto/openfga/v1"
 	"github.com/openfga/language/pkg/go/transformer"
 	"go.opentelemetry.io/otel"
 
+	"github.com/grafana/grafana/pkg/infra/localcache"
 	"github.com/grafana/grafana/pkg/infra/log"
 	authzextv1 "github.com/grafana/grafana/pkg/services/authz/proto/v1"
 	"github.com/grafana/grafana/pkg/setting"
@@ -17,6 +20,8 @@ const (
 	resourceType     = "resource"
 	namespaceType    = "namespace"
 	folderTypePrefix = "folder:"
+
+	cacheCleanInterval = 2 * time.Minute
 )
 
 var _ authzv1.AuthzServiceServer = (*Server)(nil)
@@ -28,12 +33,15 @@ type Server struct {
 	authzv1.UnimplementedAuthzServiceServer
 	authzextv1.UnimplementedAuthzExtentionServiceServer
 
-	openfga openfgav1.OpenFGAServiceServer
+	openfga       openfgav1.OpenFGAServiceServer
+	openfgaClient openfgav1.OpenFGAServiceClient
 
+	cfg      setting.ZanzanaSettings
 	logger   log.Logger
 	modules  []transformer.ModuleFile
 	stores   map[string]storeInfo
 	storesMU *sync.Mutex
+	cache    *localcache.CacheService
 }
 
 type storeInfo struct {
@@ -56,14 +64,21 @@ func WithSchema(modules []transformer.ModuleFile) ServerOption {
 }
 
 func NewAuthzServer(cfg *setting.Cfg, openfga openfgav1.OpenFGAServiceServer) (*Server, error) {
-	return NewAuthz(openfga)
+	return NewAuthz(cfg, openfga)
 }
 
-func NewAuthz(openfga openfgav1.OpenFGAServiceServer, opts ...ServerOption) (*Server, error) {
+func NewAuthz(cfg *setting.Cfg, openfga openfgav1.OpenFGAServiceServer, opts ...ServerOption) (*Server, error) {
+	channel := &inprocgrpc.Channel{}
+	openfgav1.RegisterOpenFGAServiceServer(channel, openfga)
+	openFGAClient := openfgav1.NewOpenFGAServiceClient(channel)
+
 	s := &Server{
-		openfga:  openfga,
-		storesMU: &sync.Mutex{},
-		stores:   make(map[string]storeInfo),
+		openfga:       openfga,
+		openfgaClient: openFGAClient,
+		storesMU:      &sync.Mutex{},
+		stores:        make(map[string]storeInfo),
+		cfg:           cfg.Zanzana,
+		cache:         localcache.New(cfg.Zanzana.CheckQueryCacheTTL, cacheCleanInterval),
 	}
 
 	for _, o := range opts {

pkg/services/authz/zanzana/server/server_batch_check.go

@@ -60,6 +60,8 @@ func (s *Server) batchCheckItem(
 		if err != nil {
 			return nil, err
 		}
+
+		allowed = res.GetAllowed()
 		groupResourceAccess[groupResource] = res.GetAllowed()
 	}
 

pkg/services/authz/zanzana/server/server_batch_check_test.go

@@ -13,10 +13,10 @@ import (
 )
 
 func testBatchCheck(t *testing.T, server *Server) {
-	newReq := func(subject, group, resource string, items []*authzextv1.BatchCheckItem) *authzextv1.BatchCheckRequest {
+	newReq := func(subject, verb, group, resource string, items []*authzextv1.BatchCheckItem) *authzextv1.BatchCheckRequest {
 		for i, item := range items {
 			items[i] = &authzextv1.BatchCheckItem{
-				Verb:     utils.VerbGet,
+				Verb:     verb,
 				Group:    group,
 				Resource: resource,
 				Name:     item.GetName(),
@@ -32,89 +32,116 @@ func testBatchCheck(t *testing.T, server *Server) {
 	}
 
 	t.Run("user:1 should only be able to read resource:dashboard.grafana.app/dashboards/1", func(t *testing.T) {
-		groupPrefix := zanzana.FormatGroupResource(dashboardGroup, dashboardResource)
-		res, err := server.BatchCheck(context.Background(), newReq("user:1", dashboardGroup, dashboardResource, []*authzextv1.BatchCheckItem{
+		groupResource := zanzana.FormatGroupResource(dashboardGroup, dashboardResource)
+		res, err := server.BatchCheck(context.Background(), newReq("user:1", utils.VerbGet, dashboardGroup, dashboardResource, []*authzextv1.BatchCheckItem{
 			{Name: "1", Folder: "1"},
 			{Name: "2", Folder: "2"},
 		}))
 		require.NoError(t, err)
-		require.Len(t, res.Groups[groupPrefix].Items, 2)
+		require.Len(t, res.Groups[groupResource].Items, 2)
 
-		assert.True(t, res.Groups[groupPrefix].Items["1"])
-		assert.False(t, res.Groups[groupPrefix].Items["2"])
+		assert.True(t, res.Groups[groupResource].Items["1"])
+		assert.False(t, res.Groups[groupResource].Items["2"])
 	})
 
 	t.Run("user:2 should be able to read resource:dashboard.grafana.app/dashboards/{1,2} through namespace", func(t *testing.T) {
-		groupPrefix := zanzana.FormatGroupResource(dashboardGroup, dashboardResource)
-		res, err := server.BatchCheck(context.Background(), newReq("user:2", dashboardGroup, dashboardResource, []*authzextv1.BatchCheckItem{
+		groupResource := zanzana.FormatGroupResource(dashboardGroup, dashboardResource)
+		res, err := server.BatchCheck(context.Background(), newReq("user:2", utils.VerbGet, dashboardGroup, dashboardResource, []*authzextv1.BatchCheckItem{
 			{Name: "1", Folder: "1"},
 			{Name: "2", Folder: "2"},
 		}))
 		require.NoError(t, err)
-		assert.Len(t, res.Groups[groupPrefix].Items, 2)
+		assert.Len(t, res.Groups[groupResource].Items, 2)
 	})
 
 	t.Run("user:3 should be able to read resource:dashboard.grafana.app/dashboards/1 with set relation", func(t *testing.T) {
-		groupPrefix := zanzana.FormatGroupResource(dashboardGroup, dashboardResource)
-		res, err := server.BatchCheck(context.Background(), newReq("user:3", dashboardGroup, dashboardResource, []*authzextv1.BatchCheckItem{
+		groupResource := zanzana.FormatGroupResource(dashboardGroup, dashboardResource)
+		res, err := server.BatchCheck(context.Background(), newReq("user:3", utils.VerbGet, dashboardGroup, dashboardResource, []*authzextv1.BatchCheckItem{
 			{Name: "1", Folder: "1"},
 			{Name: "2", Folder: "2"},
 		}))
 		require.NoError(t, err)
-		require.Len(t, res.Groups[groupPrefix].Items, 2)
+		require.Len(t, res.Groups[groupResource].Items, 2)
 
-		assert.True(t, res.Groups[groupPrefix].Items["1"])
-		assert.False(t, res.Groups[groupPrefix].Items["2"])
+		assert.True(t, res.Groups[groupResource].Items["1"])
+		assert.False(t, res.Groups[groupResource].Items["2"])
 	})
 
 	t.Run("user:4 should be able to read all dashboard.grafana.app/dashboards in folder 1 and 3", func(t *testing.T) {
-		groupPrefix := zanzana.FormatGroupResource(dashboardGroup, dashboardResource)
-		res, err := server.BatchCheck(context.Background(), newReq("user:4", dashboardGroup, dashboardResource, []*authzextv1.BatchCheckItem{
+		groupResource := zanzana.FormatGroupResource(dashboardGroup, dashboardResource)
+		res, err := server.BatchCheck(context.Background(), newReq("user:4", utils.VerbGet, dashboardGroup, dashboardResource, []*authzextv1.BatchCheckItem{
 			{Name: "1", Folder: "1"},
 			{Name: "2", Folder: "3"},
 			{Name: "3", Folder: "2"},
 		}))
 		require.NoError(t, err)
-		require.Len(t, res.Groups[groupPrefix].Items, 3)
+		require.Len(t, res.Groups[groupResource].Items, 3)
 
-		assert.True(t, res.Groups[groupPrefix].Items["1"])
-		assert.True(t, res.Groups[groupPrefix].Items["2"])
-		assert.False(t, res.Groups[groupPrefix].Items["3"])
+		assert.True(t, res.Groups[groupResource].Items["1"])
+		assert.True(t, res.Groups[groupResource].Items["2"])
+		assert.False(t, res.Groups[groupResource].Items["3"])
 	})
 
 	t.Run("user:5 should be able to read resource:dashboard.grafana.app/dashboards/1 through folder with set relation", func(t *testing.T) {
-		groupPrefix := zanzana.FormatGroupResource(dashboardGroup, dashboardResource)
-		res, err := server.BatchCheck(context.Background(), newReq("user:5", dashboardGroup, dashboardResource, []*authzextv1.BatchCheckItem{
+		groupResource := zanzana.FormatGroupResource(dashboardGroup, dashboardResource)
+		res, err := server.BatchCheck(context.Background(), newReq("user:5", utils.VerbGet, dashboardGroup, dashboardResource, []*authzextv1.BatchCheckItem{
 			{Name: "1", Folder: "1"},
 			{Name: "2", Folder: "2"},
 		}))
 		require.NoError(t, err)
-		require.Len(t, res.Groups[groupPrefix].Items, 2)
+		require.Len(t, res.Groups[groupResource].Items, 2)
 
-		assert.True(t, res.Groups[groupPrefix].Items["1"])
-		assert.False(t, res.Groups[groupPrefix].Items["2"])
+		assert.True(t, res.Groups[groupResource].Items["1"])
+		assert.False(t, res.Groups[groupResource].Items["2"])
 	})
 
 	t.Run("user:6 should be able to read folder 1", func(t *testing.T) {
-		groupPrefix := zanzana.FormatGroupResource(folderGroup, folderResource)
-		res, err := server.BatchCheck(context.Background(), newReq("user:6", folderGroup, folderResource, []*authzextv1.BatchCheckItem{
+		groupResource := zanzana.FormatGroupResource(folderGroup, folderResource)
+		res, err := server.BatchCheck(context.Background(), newReq("user:6", utils.VerbGet, folderGroup, folderResource, []*authzextv1.BatchCheckItem{
 			{Name: "1"},
 			{Name: "2"},
 		}))
 		require.NoError(t, err)
-		require.Len(t, res.Groups[groupPrefix].Items, 2)
+		require.Len(t, res.Groups[groupResource].Items, 2)
 
-		assert.True(t, res.Groups[groupPrefix].Items["1"])
-		assert.False(t, res.Groups[groupPrefix].Items["2"])
+		assert.True(t, res.Groups[groupResource].Items["1"])
+		assert.False(t, res.Groups[groupResource].Items["2"])
 	})
 
 	t.Run("user:7 should be able to read folder {1,2} through namespace access", func(t *testing.T) {
-		groupPrefix := zanzana.FormatGroupResource(folderGroup, folderResource)
-		res, err := server.BatchCheck(context.Background(), newReq("user:7", folderGroup, folderResource, []*authzextv1.BatchCheckItem{
+		groupResource := zanzana.FormatGroupResource(folderGroup, folderResource)
+		res, err := server.BatchCheck(context.Background(), newReq("user:7", utils.VerbGet, folderGroup, folderResource, []*authzextv1.BatchCheckItem{
 			{Name: "1"},
 			{Name: "2"},
 		}))
 		require.NoError(t, err)
-		require.Len(t, res.Groups[groupPrefix].Items, 2)
+		require.Len(t, res.Groups[groupResource].Items, 2)
+		require.True(t, res.Groups[groupResource].Items["1"])
+		require.True(t, res.Groups[groupResource].Items["2"])
+	})
+
+	t.Run("user:8 should be able to read all resoruce:dashboard.grafana.app/dashboards in folder 6 through folder 5", func(t *testing.T) {
+		groupResource := zanzana.FormatGroupResource(dashboardGroup, dashboardResource)
+		res, err := server.BatchCheck(context.Background(), newReq("user:8", utils.VerbGet, dashboardGroup, dashboardResource, []*authzextv1.BatchCheckItem{
+			{Name: "10", Folder: "6"},
+			{Name: "20", Folder: "6"},
+		}))
+		require.NoError(t, err)
+		require.Len(t, res.Groups[groupResource].Items, 2)
+		require.True(t, res.Groups[groupResource].Items["10"])
+		require.True(t, res.Groups[groupResource].Items["20"])
+	})
+
+	t.Run("user:9 should be able to create dashboards in folder 6 through folder 5", func(t *testing.T) {
+		groupResource := zanzana.FormatGroupResource(dashboardGroup, dashboardResource)
+		res, err := server.BatchCheck(context.Background(), newReq("user:9", utils.VerbCreate, dashboardGroup, dashboardResource, []*authzextv1.BatchCheckItem{
+			{Name: "10", Folder: "6"},
+			{Name: "20", Folder: "6"},
+		}))
+		require.NoError(t, err)
+		t.Log(res.Groups)
+		require.Len(t, res.Groups[groupResource].Items, 2)
+		require.True(t, res.Groups[groupResource].Items["10"])
+		require.True(t, res.Groups[groupResource].Items["20"])
 	})
 }

pkg/services/authz/zanzana/server/server_capabilities.go

@@ -21,7 +21,7 @@ func (s *Server) Capabilities(ctx context.Context, r *authzextv1.CapabilitiesReq
 
 func (s *Server) capabilitiesTyped(ctx context.Context, r *authzextv1.CapabilitiesRequest, info common.TypeInfo, store *storeInfo) (*authzextv1.CapabilitiesResponse, error) {
 	out := make([]string, 0, len(common.ResourceRelations))
-	for _, relation := range common.ResourceRelations {
+	for _, relation := range info.Relations {
 		res, err := s.checkNamespace(ctx, r.GetSubject(), relation, r.GetGroup(), r.GetResource(), store)
 		if err != nil {
 			return nil, err

pkg/services/authz/zanzana/server/server_capabilities_test.go

@@ -50,7 +50,7 @@ func testCapabilities(t *testing.T, server *Server) {
 	t.Run("user:5 should be able to read, write, create and delete resource:dashboards.grafana.app/dashboards/1 through folder with set relation", func(t *testing.T) {
 		res, err := server.Capabilities(context.Background(), newReq("user:5", dashboardGroup, dashboardResource, "1", "1"))
 		require.NoError(t, err)
-		assert.Equal(t, []string{common.RelationRead, common.RelationWrite, common.RelationCreate, common.RelationDelete}, res.GetCapabilities())
+		assert.Equal(t, []string{common.RelationRead, common.RelationWrite, common.RelationDelete}, res.GetCapabilities())
 	})
 
 	t.Run("user:6 should be able to read folder 1 ", func(t *testing.T) {

pkg/services/authz/zanzana/server/server_check.go

@@ -83,28 +83,32 @@ func (s *Server) checkTyped(ctx context.Context, subject, relation, name string,
 func (s *Server) checkGeneric(ctx context.Context, subject, relation, group, resource, name, folder string, store *storeInfo) (*authzv1.CheckResponse, error) {
 	groupResource := structpb.NewStringValue(common.FormatGroupResource(group, resource))
 
-	// Check if subject has direct access to resource
-	res, err := s.openfga.Check(ctx, &openfgav1.CheckRequest{
-		StoreId:              store.ID,
-		AuthorizationModelId: store.ModelID,
-		TupleKey: &openfgav1.CheckRequestTupleKey{
-			User:     subject,
-			Relation: relation,
-			Object:   common.NewResourceIdent(group, resource, name),
-		},
-		Context: &structpb.Struct{
-			Fields: map[string]*structpb.Value{
-				"requested_group": groupResource,
+	// Create relation can only exist on namespace or folder level.
+	// So we skip direct resource access check.
+	if relation != common.RelationCreate {
+		// Check if subject has direct access to resource
+		res, err := s.openfga.Check(ctx, &openfgav1.CheckRequest{
+			StoreId:              store.ID,
+			AuthorizationModelId: store.ModelID,
+			TupleKey: &openfgav1.CheckRequestTupleKey{
+				User:     subject,
+				Relation: relation,
+				Object:   common.NewResourceIdent(group, resource, name),
 			},
-		},
-	})
+			Context: &structpb.Struct{
+				Fields: map[string]*structpb.Value{
+					"requested_group": groupResource,
+				},
+			},
+		})
 
-	if err != nil {
-		return nil, err
-	}
+		if err != nil {
+			return nil, err
+		}
 
-	if res.GetAllowed() {
-		return &authzv1.CheckResponse{Allowed: true}, nil
+		if res.GetAllowed() {
+			return &authzv1.CheckResponse{Allowed: true}, nil
+		}
 	}
 
 	if folder == "" {
@@ -112,7 +116,7 @@ func (s *Server) checkGeneric(ctx context.Context, subject, relation, group, res
 	}
 
 	// Check if subject has access as a sub resource for the folder
-	res, err = s.openfga.Check(ctx, &openfgav1.CheckRequest{
+	res, err := s.openfga.Check(ctx, &openfgav1.CheckRequest{
 		StoreId:              store.ID,
 		AuthorizationModelId: store.ModelID,
 		TupleKey: &openfgav1.CheckRequestTupleKey{

pkg/services/authz/zanzana/server/server_check_test.go

@@ -12,11 +12,11 @@ import (
 )
 
 func testCheck(t *testing.T, server *Server) {
-	newRead := func(subject, group, resource, folder, name string) *authzv1.CheckRequest {
+	newReq := func(subject, verb, group, resource, folder, name string) *authzv1.CheckRequest {
 		return &authzv1.CheckRequest{
 			Namespace: namespace,
 			Subject:   subject,
-			Verb:      utils.VerbGet,
+			Verb:      verb,
 			Group:     group,
 			Resource:  resource,
 			Name:      name,
@@ -25,85 +25,91 @@ func testCheck(t *testing.T, server *Server) {
 	}
 
 	t.Run("user:1 should only be able to read resource:dashboard.grafana.app/dashboards/1", func(t *testing.T) {
-		res, err := server.Check(context.Background(), newRead("user:1", dashboardGroup, dashboardResource, "1", "1"))
+		res, err := server.Check(context.Background(), newReq("user:1", utils.VerbGet, dashboardGroup, dashboardResource, "1", "1"))
 		require.NoError(t, err)
 		assert.True(t, res.GetAllowed())
 
 		// sanity check
-		res, err = server.Check(context.Background(), newRead("user:1", dashboardGroup, dashboardResource, "1", "2"))
+		res, err = server.Check(context.Background(), newReq("user:1", utils.VerbGet, dashboardGroup, dashboardResource, "1", "2"))
 		require.NoError(t, err)
 		assert.False(t, res.GetAllowed())
 	})
 
 	t.Run("user:2 should be able to read resource:dashboard.grafana.app/dashboards/1 through namespace", func(t *testing.T) {
-		res, err := server.Check(context.Background(), newRead("user:2", dashboardGroup, dashboardResource, "1", "1"))
+		res, err := server.Check(context.Background(), newReq("user:2", utils.VerbGet, dashboardGroup, dashboardResource, "1", "1"))
 		require.NoError(t, err)
 		assert.True(t, res.GetAllowed())
 	})
 
 	t.Run("user:3 should be able to read resource:dashboard.grafana.app/dashboards/1 with set relation", func(t *testing.T) {
-		res, err := server.Check(context.Background(), newRead("user:3", dashboardGroup, dashboardResource, "1", "1"))
+		res, err := server.Check(context.Background(), newReq("user:3", utils.VerbGet, dashboardGroup, dashboardResource, "1", "1"))
 		require.NoError(t, err)
 		assert.True(t, res.GetAllowed())
 
 		// sanity check
-		res, err = server.Check(context.Background(), newRead("user:3", dashboardGroup, dashboardResource, "1", "2"))
+		res, err = server.Check(context.Background(), newReq("user:3", utils.VerbGet, dashboardGroup, dashboardResource, "1", "2"))
 		require.NoError(t, err)
 		assert.False(t, res.GetAllowed())
 	})
 
 	t.Run("user:4 should be able to read all dashboard.grafana.app/dashboards in folder 1 and 3", func(t *testing.T) {
-		res, err := server.Check(context.Background(), newRead("user:4", dashboardGroup, dashboardResource, "1", "1"))
+		res, err := server.Check(context.Background(), newReq("user:4", utils.VerbGet, dashboardGroup, dashboardResource, "1", "1"))
 		require.NoError(t, err)
 		assert.True(t, res.GetAllowed())
 
-		res, err = server.Check(context.Background(), newRead("user:4", dashboardGroup, dashboardResource, "3", "2"))
+		res, err = server.Check(context.Background(), newReq("user:4", utils.VerbGet, dashboardGroup, dashboardResource, "3", "2"))
 		require.NoError(t, err)
 		assert.True(t, res.GetAllowed())
 
 		// sanity check
-		res, err = server.Check(context.Background(), newRead("user:4", dashboardGroup, dashboardResource, "1", "2"))
+		res, err = server.Check(context.Background(), newReq("user:4", utils.VerbGet, dashboardGroup, dashboardResource, "1", "2"))
 		require.NoError(t, err)
 		assert.True(t, res.GetAllowed())
 
-		res, err = server.Check(context.Background(), newRead("user:4", dashboardGroup, dashboardResource, "2", "2"))
+		res, err = server.Check(context.Background(), newReq("user:4", utils.VerbGet, dashboardGroup, dashboardResource, "2", "2"))
 		require.NoError(t, err)
 		assert.False(t, res.GetAllowed())
 	})
 
 	t.Run("user:5 should be able to read resource:dashboard.grafana.app/dashboards/1 through folder with set relation", func(t *testing.T) {
-		res, err := server.Check(context.Background(), newRead("user:5", dashboardGroup, dashboardResource, "1", "1"))
+		res, err := server.Check(context.Background(), newReq("user:5", utils.VerbGet, dashboardGroup, dashboardResource, "1", "1"))
 		require.NoError(t, err)
 		assert.True(t, res.GetAllowed())
 	})
 
 	t.Run("user:6 should be able to read folder 1 ", func(t *testing.T) {
-		res, err := server.Check(context.Background(), newRead("user:6", folderGroup, folderResource, "", "1"))
+		res, err := server.Check(context.Background(), newReq("user:6", utils.VerbGet, folderGroup, folderResource, "", "1"))
 		require.NoError(t, err)
 		assert.True(t, res.GetAllowed())
 	})
 
 	t.Run("user:7 should be able to read folder one through namespace access", func(t *testing.T) {
-		res, err := server.Check(context.Background(), newRead("user:7", folderGroup, folderResource, "", "1"))
+		res, err := server.Check(context.Background(), newReq("user:7", utils.VerbGet, folderGroup, folderResource, "", "1"))
 		require.NoError(t, err)
 		assert.True(t, res.GetAllowed())
 
-		res, err = server.Check(context.Background(), newRead("user:7", folderGroup, folderResource, "", "10"))
+		res, err = server.Check(context.Background(), newReq("user:7", utils.VerbGet, folderGroup, folderResource, "", "10"))
 		require.NoError(t, err)
 		assert.True(t, res.GetAllowed())
 	})
 
 	t.Run("user:8 should be able to read all resoruce:dashboard.grafana.app/dashboar in folder 6 through folder 5", func(t *testing.T) {
-		res, err := server.Check(context.Background(), newRead("user:8", dashboardGroup, dashboardResource, "6", "10"))
+		res, err := server.Check(context.Background(), newReq("user:8", utils.VerbGet, dashboardGroup, dashboardResource, "6", "10"))
 		require.NoError(t, err)
 		assert.True(t, res.GetAllowed())
 
-		res, err = server.Check(context.Background(), newRead("user:8", dashboardGroup, dashboardResource, "5", "11"))
+		res, err = server.Check(context.Background(), newReq("user:8", utils.VerbGet, dashboardGroup, dashboardResource, "5", "11"))
 		require.NoError(t, err)
 		assert.True(t, res.GetAllowed())
 
-		res, err = server.Check(context.Background(), newRead("user:8", folderGroup, folderResource, "4", "12"))
+		res, err = server.Check(context.Background(), newReq("user:8", utils.VerbGet, folderGroup, folderResource, "4", "12"))
 		require.NoError(t, err)
 		assert.False(t, res.GetAllowed())
 	})
+
+	t.Run("user:9 should be able to create dashboards in folder 5", func(t *testing.T) {
+		res, err := server.Check(context.Background(), newReq("user:9", utils.VerbCreate, dashboardGroup, dashboardResource, "5", ""))
+		require.NoError(t, err)
+		assert.True(t, res.GetAllowed())
+	})
 }

pkg/services/authz/zanzana/server/server_list.go

@@ -47,9 +47,16 @@ func (s *Server) List(ctx context.Context, r *authzextv1.ListRequest) (*authzext
 	return s.listGeneric(ctx, r.GetSubject(), relation, r.GetGroup(), r.GetResource(), store)
 }
 
+func (s *Server) listObjects(ctx context.Context, req *openfgav1.ListObjectsRequest) (*openfgav1.ListObjectsResponse, error) {
+	if s.cfg.UseStreamedListObjects {
+		return s.streamedListObjects(ctx, req)
+	}
+	return s.openfga.ListObjects(ctx, req)
+}
+
 func (s *Server) listTyped(ctx context.Context, subject, relation string, info common.TypeInfo, store *storeInfo) (*authzextv1.ListResponse, error) {
 	// List all resources user has access too
-	listRes, err := s.openfga.ListObjects(ctx, &openfgav1.ListObjectsRequest{
+	listRes, err := s.listObjects(ctx, &openfgav1.ListObjectsRequest{
 		StoreId:              store.ID,
 		AuthorizationModelId: store.ModelID,
 		Type:                 info.Type,
@@ -69,7 +76,7 @@ func (s *Server) listGeneric(ctx context.Context, subject, relation, group, reso
 	groupResource := structpb.NewStringValue(common.FormatGroupResource(group, resource))
 
 	// 1. List all folders subject has access to resource type in
-	folders, err := s.openfga.ListObjects(ctx, &openfgav1.ListObjectsRequest{
+	folders, err := s.listObjects(ctx, &openfgav1.ListObjectsRequest{
 		StoreId:              store.ID,
 		AuthorizationModelId: store.ModelID,
 		Type:                 common.TypeFolder,
@@ -86,7 +93,7 @@ func (s *Server) listGeneric(ctx context.Context, subject, relation, group, reso
 	}
 
 	// 2. List all resource directly assigned to subject
-	direct, err := s.openfga.ListObjects(ctx, &openfgav1.ListObjectsRequest{
+	direct, err := s.listObjects(ctx, &openfgav1.ListObjectsRequest{
 		StoreId:              store.ID,
 		AuthorizationModelId: store.ModelID,
 		Type:                 common.TypeResource,

pkg/services/authz/zanzana/server/server_list_streamed.go

@@ -0,0 +1,100 @@
+package server
+
+import (
+	"context"
+	"encoding/base64"
+	"errors"
+	"hash/fnv"
+	"io"
+
+	openfgav1 "github.com/openfga/api/proto/openfga/v1"
+)
+
+func (s *Server) streamedListObjects(ctx context.Context, req *openfgav1.ListObjectsRequest) (*openfgav1.ListObjectsResponse, error) {
+	if !s.cfg.CheckQueryCache {
+		return s.listObjectsWithStream(ctx, req)
+	}
+	return s.streamedListObjectsCached(ctx, req)
+}
+
+func (s *Server) streamedListObjectsCached(ctx context.Context, req *openfgav1.ListObjectsRequest) (*openfgav1.ListObjectsResponse, error) {
+	ctx, span := tracer.Start(ctx, "authzServer.streamedListObjectsCached")
+	defer span.End()
+
+	reqHash, err := getRequestHash(req)
+	if err != nil {
+		return nil, err
+	}
+
+	if res, ok := s.cache.Get(reqHash); ok {
+		return res.(*openfgav1.ListObjectsResponse), nil
+	}
+
+	res, err := s.listObjectsWithStream(ctx, req)
+	if err != nil {
+		return nil, err
+	}
+	s.cache.Set(reqHash, res, 0)
+	return res, nil
+}
+
+func (s *Server) listObjectsWithStream(ctx context.Context, req *openfgav1.ListObjectsRequest) (*openfgav1.ListObjectsResponse, error) {
+	ctx, span := tracer.Start(ctx, "authzServer.listObjectsWithStream")
+	defer span.End()
+
+	r := &openfgav1.StreamedListObjectsRequest{
+		StoreId:              req.GetStoreId(),
+		AuthorizationModelId: req.GetAuthorizationModelId(),
+		Type:                 req.GetType(),
+		Relation:             req.GetRelation(),
+		User:                 req.GetUser(),
+		Context:              req.GetContext(),
+	}
+
+	clientStream, err := s.openfgaClient.StreamedListObjects(ctx, r)
+	if err != nil {
+		return nil, err
+	}
+
+	done := make(chan struct{})
+	var streamedObjectIDs []string
+	var streamingErr error
+	var streamingResp *openfgav1.StreamedListObjectsResponse
+	go func() {
+		for {
+			streamingResp, streamingErr = clientStream.Recv()
+			if streamingErr == nil {
+				streamedObjectIDs = append(streamedObjectIDs, streamingResp.GetObject())
+			} else {
+				if errors.Is(streamingErr, io.EOF) {
+					streamingErr = nil
+				}
+				break
+			}
+		}
+		done <- struct{}{}
+	}()
+	<-done
+
+	if streamingErr != nil {
+		return nil, streamingErr
+	}
+
+	return &openfgav1.ListObjectsResponse{
+		Objects: streamedObjectIDs,
+	}, nil
+}
+
+func getRequestHash(req *openfgav1.ListObjectsRequest) (string, error) {
+	if req == nil {
+		return "", errors.New("request must not be empty")
+	}
+
+	hash := fnv.New64a()
+	_, err := hash.Write([]byte(req.String()))
+	if err != nil {
+		return "", err
+	}
+
+	return base64.StdEncoding.EncodeToString(hash.Sum(nil)), nil
+}

pkg/services/authz/zanzana/server/server_test.go

@@ -68,7 +68,7 @@ func setup(t *testing.T, testDB db.DB, cfg *setting.Cfg) *Server {
 	openfga, err := NewOpenFGA(&cfg.Zanzana, store, log.NewNopLogger())
 	require.NoError(t, err)
 
-	srv, err := NewAuthz(openfga)
+	srv, err := NewAuthz(cfg, openfga)
 	require.NoError(t, err)
 
 	storeInf, err := srv.getStoreInfo(context.Background(), namespace)
@@ -92,7 +92,8 @@ func setup(t *testing.T, testDB db.DB, cfg *setting.Cfg) *Server {
 				common.NewNamespaceResourceTuple("user:7", "read", folderGroup, folderResource),
 				common.NewFolderParentTuple("5", "4"),
 				common.NewFolderParentTuple("6", "5"),
-				common.NewFolderResourceTuple("user:8", "view", dashboardGroup, dashboardResource, "5"),
+				common.NewFolderResourceTuple("user:8", "edit", dashboardGroup, dashboardResource, "5"),
+				common.NewFolderResourceTuple("user:9", "create", dashboardGroup, dashboardResource, "5"),
 			},
 		},
 	})

pkg/services/authz/zanzana/zanzana.go

@@ -11,12 +11,13 @@ import (
 )
 
 const (
-	TypeUser      = common.TypeUser
-	TypeTeam      = common.TypeTeam
-	TypeRole      = common.TypeRole
-	TypeFolder    = common.TypeFolder
-	TypeResource  = common.TypeResource
-	TypeNamespace = common.TypeNamespace
+	TypeUser           = common.TypeUser
+	TypeServiceAccount = common.TypeServiceAccount
+	TypeTeam           = common.TypeTeam
+	TypeRole           = common.TypeRole
+	TypeFolder         = common.TypeFolder
+	TypeResource       = common.TypeResource
+	TypeNamespace      = common.TypeNamespace
 )
 
 const (

pkg/services/datasources/service/legacy.go

@@ -9,6 +9,7 @@ import (
 	data "github.com/grafana/grafana-plugin-sdk-go/experimental/apis/data/v0alpha1"
 
 	"github.com/grafana/grafana/pkg/apimachinery/identity"
+	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/services/datasources"
 )
 
@@ -38,6 +39,7 @@ type cachingLegacyDataSourceLookup struct {
 	retriever DataSourceRetriever
 	cache     map[string]cachedValue
 	cacheMu   sync.Mutex
+	log       log.Logger
 }
 
 type cachedValue struct {
@@ -49,15 +51,18 @@ func ProvideLegacyDataSourceLookup(p *Service) LegacyDataSourceLookup {
 	return &cachingLegacyDataSourceLookup{
 		retriever: p,
 		cache:     make(map[string]cachedValue),
+		log:       log.New("legacy-datasource-lookup"),
 	}
 }
 
 func (s *cachingLegacyDataSourceLookup) GetDataSourceFromDeprecatedFields(ctx context.Context, name string, id int64) (*data.DataSourceRef, error) {
 	if id == 0 && name == "" {
+		s.log.Error("missing id and name in GetDataSourceFromDeprecatedFields")
 		return nil, fmt.Errorf("either name or ID must be set")
 	}
 	user, err := identity.GetRequester(ctx)
 	if err != nil {
+		s.log.Error("failed to get user from context after getRequester", "error", err)
 		return nil, err
 	}
 	key := fmt.Sprintf("%d/%s/%d", user.GetOrgID(), name, id)
@@ -74,6 +79,9 @@ func (s *cachingLegacyDataSourceLookup) GetDataSourceFromDeprecatedFields(ctx co
 		Name:  name,
 		ID:    id,
 	})
+	if err != nil {
+		s.log.Error("failed to get datasource from retriever", "error", err)
+	}
 	if errors.Is(err, datasources.ErrDataSourceNotFound) && name != "" {
 		ds, err = s.retriever.GetDataSource(ctx, &datasources.GetDataSourceQuery{
 			OrgID: user.GetOrgID(),

pkg/services/featuremgmt/registry.go

@@ -1070,6 +1070,14 @@ var (
 			FrontendOnly: false,
 			Owner:        grafanaAlertingSquad,
 		},
+		{
+			Name:              "scopeApi",
+			Description:       "In-development feature flag for the scope api using the app platform.",
+			Stage:             FeatureStageExperimental,
+			Owner:             grafanaAppPlatformSquad,
+			HideFromAdminPage: true,
+			Expression:        "false",
+		},
 		{
 			Name:        "promQLScope",
 			Description: "In-development feature that will allow injection of labels into prometheus queries.",
@@ -1658,6 +1666,13 @@ var (
 			Owner:        grafanaAlertingSquad,
 			Expression:   "true", // enabled by default
 		},
+		{
+			Name:        "azureMonitorEnableUserAuth",
+			Description: "Enables user auth for Azure Monitor datasource only",
+			Stage:       FeatureStageGeneralAvailability,
+			Owner:       grafanaPartnerPluginsSquad,
+			Expression:  "true", // Enabled by default for now
+		},
 		{
 			Name:         "alertingNotificationsStepMode",
 			Description:  "Enables simplified step mode in the notifications section",

pkg/services/featuremgmt/toggles_gen.csv

@@ -141,6 +141,7 @@ onPremToCloudMigrations,preview,@grafana/grafana-operator-experience-squad,false
 onPremToCloudMigrationsAlerts,experimental,@grafana/grafana-operator-experience-squad,false,false,false
 onPremToCloudMigrationsAuthApiMig,experimental,@grafana/grafana-operator-experience-squad,false,false,false
 alertingSaveStatePeriodic,privatePreview,@grafana/alerting-squad,false,false,false
+scopeApi,experimental,@grafana/grafana-app-platform-squad,false,false,false
 promQLScope,GA,@grafana/observability-metrics,false,false,false
 logQLScope,privatePreview,@grafana/observability-logs,false,false,false
 sqlExpressions,experimental,@grafana/grafana-app-platform-squad,false,false,false
@@ -221,4 +222,5 @@ crashDetection,experimental,@grafana/observability-traces-and-profiling,false,fa
 jaegerBackendMigration,experimental,@grafana/oss-big-tent,false,false,false
 reportingUseRawTimeRange,preview,@grafana/sharing-squad,false,false,false
 alertingUIOptimizeReducer,GA,@grafana/alerting-squad,false,false,true
+azureMonitorEnableUserAuth,GA,@grafana/partner-datasources,false,false,false
 alertingNotificationsStepMode,experimental,@grafana/alerting-squad,false,false,true

pkg/services/featuremgmt/toggles_gen.go

@@ -575,6 +575,10 @@ const (
 	// Writes the state periodically to the database, asynchronous to rule evaluation
 	FlagAlertingSaveStatePeriodic = "alertingSaveStatePeriodic"
 
+	// FlagScopeApi
+	// In-development feature flag for the scope api using the app platform.
+	FlagScopeApi = "scopeApi"
+
 	// FlagPromQLScope
 	// In-development feature that will allow injection of labels into prometheus queries.
 	FlagPromQLScope = "promQLScope"
@@ -895,6 +899,10 @@ const (
 	// Enables removing the reducer from the alerting UI when creating a new alert rule and using instant query
 	FlagAlertingUIOptimizeReducer = "alertingUIOptimizeReducer"
 
+	// FlagAzureMonitorEnableUserAuth
+	// Enables user auth for Azure Monitor datasource only
+	FlagAzureMonitorEnableUserAuth = "azureMonitorEnableUserAuth"
+
 	// FlagAlertingNotificationsStepMode
 	// Enables simplified step mode in the notifications section
 	FlagAlertingNotificationsStepMode = "alertingNotificationsStepMode"

pkg/services/featuremgmt/toggles_gen.json

@@ -660,6 +660,22 @@
         "expression": "false"
       }
     },
+    {
+      "metadata": {
+        "name": "azureMonitorEnableUserAuth",
+        "resourceVersion": "1732189410576",
+        "creationTimestamp": "2024-11-21T11:42:29Z",
+        "annotations": {
+          "grafana.app/updatedTimestamp": "2024-11-21 11:43:30.576196 +0000 UTC"
+        }
+      },
+      "spec": {
+        "description": "Enables user auth for Azure Monitor datasource only",
+        "stage": "GA",
+        "codeowner": "@grafana/partner-datasources",
+        "expression": "true"
+      }
+    },
     {
       "metadata": {
         "name": "azureMonitorLogLimit",
@@ -3130,6 +3146,20 @@
         "frontend": true
       }
     },
+    {
+      "metadata": {
+        "name": "scopeApi",
+        "resourceVersion": "1732690644377",
+        "creationTimestamp": "2024-11-27T06:57:24Z"
+      },
+      "spec": {
+        "description": "In-development feature flag for the scope api using the app platform.",
+        "stage": "experimental",
+        "codeowner": "@grafana/grafana-app-platform-squad",
+        "hideFromAdminPage": true,
+        "expression": "false"
+      }
+    },
     {
       "metadata": {
         "name": "scopeFilters",

pkg/services/login/authinfoimpl/store.go

@@ -92,7 +92,6 @@ func (s *Store) GetUserLabels(ctx context.Context, query login.GetUserLabelsQuer
 	err := s.sqlStore.WithDbSession(ctx, func(sess *db.Session) error {
 		return sess.Table("user_auth").In("user_id", params).OrderBy("created").Find(&userAuths)
 	})
-
 	if err != nil {
 		return nil, err
 	}
@@ -187,7 +186,9 @@ func (s *Store) UpdateAuthInfo(ctx context.Context, cmd *login.UpdateAuthInfoCom
 	}
 
 	return s.sqlStore.WithTransactionalDbSession(ctx, func(sess *db.Session) error {
-		upd, err := sess.MustCols("o_auth_expiry").Where("user_id = ? AND auth_module = ?", cmd.UserId, cmd.AuthModule).Update(authUser)
+		upd, err := sess.MustCols("o_auth_expiry", "o_auth_access_token", "o_auth_refresh_token", "o_auth_id_token", "o_auth_token_type").
+			Where("user_id = ? AND auth_module = ?", cmd.UserId, cmd.AuthModule).
+			Update(authUser)
 
 		s.logger.Debug("Updated user_auth", "user_id", cmd.UserId, "auth_id", cmd.AuthId, "auth_module", cmd.AuthModule, "rows", upd)
 
@@ -198,7 +199,6 @@ func (s *Store) UpdateAuthInfo(ctx context.Context, cmd *login.UpdateAuthInfoCom
 				"SELECT id FROM user_auth WHERE user_id = ? AND auth_module = ? AND auth_id = ?",
 				cmd.UserId, cmd.AuthModule, cmd.AuthId,
 			).Get(&id)
-
 			if err != nil {
 				return err
 			}
@@ -220,7 +220,7 @@ func (s *Store) UpdateAuthInfo(ctx context.Context, cmd *login.UpdateAuthInfoCom
 
 func (s *Store) DeleteUserAuthInfo(ctx context.Context, userID int64) error {
 	return s.sqlStore.WithDbSession(ctx, func(sess *db.Session) error {
-		var rawSQL = "DELETE FROM user_auth WHERE user_id = ?"
+		rawSQL := "DELETE FROM user_auth WHERE user_id = ?"
 		_, err := sess.Exec(rawSQL, userID)
 		return err
 	})

pkg/services/loginattempt/loginattemptimpl/login_attempt.go

@@ -11,10 +11,7 @@ import (
 	"github.com/grafana/grafana/pkg/setting"
 )
 
-const (
-	maxInvalidLoginAttempts int64 = 5
-	loginAttemptsWindow           = time.Minute * 5
-)
+const loginAttemptsWindow = time.Minute * 5
 
 func ProvideService(db db.DB, cfg *setting.Cfg, lock *serverlock.ServerLockService) *Service {
 	return &Service{
@@ -80,7 +77,7 @@ func (s *Service) Validate(ctx context.Context, username string) (bool, error) {
 		return false, err
 	}
 
-	if count >= maxInvalidLoginAttempts {
+	if count >= s.cfg.BruteForceLoginProtectionMaxAttempts {
 		return false, nil
 	}
 

pkg/services/loginattempt/loginattemptimpl/login_attempt_test.go

@@ -12,6 +12,8 @@ import (
 )
 
 func TestService_Validate(t *testing.T) {
+	const maxInvalidLoginAttempts = 5
+
 	testCases := []struct {
 		name          string
 		loginAttempts int64
@@ -64,6 +66,7 @@ func TestService_Validate(t *testing.T) {
 	for _, tt := range testCases {
 		t.Run(tt.name, func(t *testing.T) {
 			cfg := setting.NewCfg()
+			cfg.BruteForceLoginProtectionMaxAttempts = maxInvalidLoginAttempts
 			cfg.DisableBruteForceLoginProtection = tt.disabled
 			service := &Service{
 				store: fakeStore{
@@ -84,6 +87,7 @@ func TestLoginAttempts(t *testing.T) {
 	ctx := context.Background()
 	cfg := setting.NewCfg()
 	cfg.DisableBruteForceLoginProtection = false
+	cfg.BruteForceLoginProtectionMaxAttempts = 5
 	db := db.InitTestDB(t)
 	service := ProvideService(db, cfg, nil)
 

pkg/services/oauthtoken/oauth_token.go

@@ -19,7 +19,9 @@ import (
 	"github.com/grafana/grafana/pkg/infra/serverlock"
 	"github.com/grafana/grafana/pkg/infra/tracing"
 	"github.com/grafana/grafana/pkg/login/social"
+	"github.com/grafana/grafana/pkg/services/auth"
 	"github.com/grafana/grafana/pkg/services/datasources"
+	"github.com/grafana/grafana/pkg/services/featuremgmt"
 	"github.com/grafana/grafana/pkg/services/login"
 	"github.com/grafana/grafana/pkg/services/user"
 	"github.com/grafana/grafana/pkg/setting"
@@ -39,28 +41,33 @@ type Service struct {
 	Cfg             *setting.Cfg
 	SocialService   social.Service
 	AuthInfoService login.AuthInfoService
+	sessionService  auth.UserTokenService
+	features        featuremgmt.FeatureToggles
 	serverLock      *serverlock.ServerLockService
 	tracer          tracing.Tracer
 
 	tokenRefreshDuration *prometheus.HistogramVec
 }
 
+var _ OAuthTokenService = (*Service)(nil)
+
 //go:generate mockery --name OAuthTokenService --structname MockService --outpkg oauthtokentest --filename service_mock.go --output ./oauthtokentest/
 type OAuthTokenService interface {
-	GetCurrentOAuthToken(context.Context, identity.Requester) *oauth2.Token
+	GetCurrentOAuthToken(context.Context, identity.Requester, *auth.UserToken) *oauth2.Token
 	IsOAuthPassThruEnabled(*datasources.DataSource) bool
-	HasOAuthEntry(context.Context, identity.Requester) (*login.UserAuth, bool, error)
-	TryTokenRefresh(context.Context, identity.Requester) (*oauth2.Token, error)
-	InvalidateOAuthTokens(context.Context, identity.Requester) error
+	TryTokenRefresh(context.Context, identity.Requester, *auth.UserToken) (*oauth2.Token, error)
+	InvalidateOAuthTokens(context.Context, identity.Requester, *auth.UserToken) error
 }
 
 func ProvideService(socialService social.Service, authInfoService login.AuthInfoService, cfg *setting.Cfg, registerer prometheus.Registerer,
-	serverLockService *serverlock.ServerLockService, tracer tracing.Tracer,
+	serverLockService *serverlock.ServerLockService, tracer tracing.Tracer, sessionService auth.UserTokenService, features featuremgmt.FeatureToggles,
 ) *Service {
 	return &Service{
 		AuthInfoService:      authInfoService,
+		sessionService:       sessionService,
 		Cfg:                  cfg,
 		SocialService:        socialService,
+		features:             features,
 		serverLock:           serverLockService,
 		tokenRefreshDuration: newTokenRefreshDurationMetric(registerer),
 		tracer:               tracer,
@@ -68,7 +75,7 @@ func ProvideService(socialService social.Service, authInfoService login.AuthInfo
 }
 
 // GetCurrentOAuthToken returns the OAuth token, if any, for the authenticated user. Will try to refresh the token if it has expired.
-func (o *Service) GetCurrentOAuthToken(ctx context.Context, usr identity.Requester) *oauth2.Token {
+func (o *Service) GetCurrentOAuthToken(ctx context.Context, usr identity.Requester, sessionToken *auth.UserToken) *oauth2.Token {
 	ctx, span := o.tracer.Start(ctx, "oauthtoken.GetCurrentOAuthToken")
 	defer span.End()
 
@@ -93,30 +100,54 @@ func (o *Service) GetCurrentOAuthToken(ctx context.Context, usr identity.Request
 
 	ctxLogger = ctxLogger.New("userID", userID)
 
-	authInfo, ok, _ := o.HasOAuthEntry(ctx, usr)
-	if !ok {
+	if !strings.HasPrefix(usr.GetAuthenticatedBy(), "oauth_") {
+		ctxLogger.Warn("The specified user's auth provider is not oauth",
+			"authmodule", usr.GetAuthenticatedBy())
 		return nil
 	}
 
-	if err := checkOAuthRefreshToken(authInfo); err != nil {
-		if errors.Is(err, ErrNoRefreshTokenFound) {
-			return buildOAuthTokenFromAuthInfo(authInfo)
+	var persistedToken *oauth2.Token
+	if o.features.IsEnabledGlobally(featuremgmt.FlagImprovedExternalSessionHandling) {
+		externalSession, err := o.sessionService.GetExternalSession(ctx, sessionToken.ExternalSessionId)
+		if err != nil {
+			if errors.Is(err, auth.ErrExternalSessionNotFound) {
+				return nil
+			}
+			ctxLogger.Error("Failed to fetch external session", "error", err)
+			return nil
 		}
 
-		return nil
-	}
+		persistedToken = buildOAuthTokenFromExternalSession(externalSession)
 
-	persistedToken := buildOAuthTokenFromAuthInfo(authInfo)
+		if persistedToken.RefreshToken == "" {
+			return persistedToken
+		}
+	} else {
+		authInfo, ok, _ := o.hasOAuthEntry(ctx, usr)
+		if !ok {
+			return nil
+		}
+
+		if err := checkOAuthRefreshToken(authInfo); err != nil {
+			if errors.Is(err, ErrNoRefreshTokenFound) {
+				return buildOAuthTokenFromAuthInfo(authInfo)
+			}
+
+			return nil
+		}
+
+		persistedToken = buildOAuthTokenFromAuthInfo(authInfo)
+	}
 
 	refreshNeeded := needTokenRefresh(ctx, persistedToken)
 	if !refreshNeeded {
 		return persistedToken
 	}
 
-	token, err := o.TryTokenRefresh(ctx, usr)
+	token, err := o.TryTokenRefresh(ctx, usr, sessionToken)
 	if err != nil {
 		if errors.Is(err, ErrNoRefreshTokenFound) {
-			return buildOAuthTokenFromAuthInfo(authInfo)
+			return persistedToken
 		}
 
 		return nil
@@ -130,9 +161,9 @@ func (o *Service) IsOAuthPassThruEnabled(ds *datasources.DataSource) bool {
 	return IsOAuthPassThruEnabled(ds)
 }
 
-// HasOAuthEntry returns true and the UserAuth object when OAuth info exists for the specified User
-func (o *Service) HasOAuthEntry(ctx context.Context, usr identity.Requester) (*login.UserAuth, bool, error) {
-	ctx, span := o.tracer.Start(ctx, "oauthtoken.HasOAuthEntry")
+// hasOAuthEntry returns true and the UserAuth object when OAuth info exists for the specified User
+func (o *Service) hasOAuthEntry(ctx context.Context, usr identity.Requester) (*login.UserAuth, bool, error) {
+	ctx, span := o.tracer.Start(ctx, "oauthtoken.hasOAuthEntry")
 	defer span.End()
 
 	if usr == nil || usr.IsNil() {
@@ -167,12 +198,19 @@ func (o *Service) HasOAuthEntry(ctx context.Context, usr identity.Requester) (*l
 	if !strings.Contains(authInfo.AuthModule, "oauth") {
 		return nil, false, nil
 	}
+
+	// An extra check to ensure that the user has an OAuth token
+	// It's required to handle the case when the `improvedExternalSessionHandling` feature flag gets disabled
+	if authInfo.OAuthAccessToken == "" {
+		ctxLogger.Debug("No access token found for user")
+		return nil, false, fmt.Errorf("no access token found for user %d", userID)
+	}
 	return authInfo, true, nil
 }
 
 // TryTokenRefresh returns an error in case the OAuth token refresh was unsuccessful
 // It uses a server lock to prevent getting the Refresh Token multiple times for a given User
-func (o *Service) TryTokenRefresh(ctx context.Context, usr identity.Requester) (*oauth2.Token, error) {
+func (o *Service) TryTokenRefresh(ctx context.Context, usr identity.Requester, sessionToken *auth.UserToken) (*oauth2.Token, error) {
 	ctx, span := o.tracer.Start(ctx, "oauthtoken.TryTokenRefresh")
 	defer span.End()
 
@@ -218,6 +256,9 @@ func (o *Service) TryTokenRefresh(ctx context.Context, usr identity.Requester) (
 	}
 
 	lockKey := fmt.Sprintf("oauth-refresh-token-%d", userID)
+	if o.features.IsEnabledGlobally(featuremgmt.FlagImprovedExternalSessionHandling) {
+		lockKey = fmt.Sprintf("oauth-refresh-token-%d-%d", userID, sessionToken.ExternalSessionId)
+	}
 
 	lockTimeConfig := serverlock.LockTimeConfig{
 		MaxInterval: 30 * time.Second,
@@ -242,15 +283,32 @@ func (o *Service) TryTokenRefresh(ctx context.Context, usr identity.Requester) (
 
 		ctxLogger.Debug("Serverlock request for getting a new access token", "key", lockKey)
 
-		authInfo, exists, err := o.HasOAuthEntry(ctx, usr)
-		if !exists {
+		var persistedToken *oauth2.Token
+		var externalSession *auth.ExternalSession
+		if o.features.IsEnabledGlobally(featuremgmt.FlagImprovedExternalSessionHandling) {
+			externalSession, err = o.sessionService.GetExternalSession(ctx, sessionToken.ExternalSessionId)
 			if err != nil {
-				ctxLogger.Debug("Failed to fetch oauth entry", "error", err)
+				if errors.Is(err, auth.ErrExternalSessionNotFound) {
+					ctxLogger.Error("External session was not found for user", "error", err)
+					return
+				}
+				ctxLogger.Error("Failed to fetch external session", "error", err)
+				return
 			}
-			return
-		}
 
-		persistedToken := buildOAuthTokenFromAuthInfo(authInfo)
+			persistedToken = buildOAuthTokenFromExternalSession(externalSession)
+		} else {
+			authInfo, exists, err := o.hasOAuthEntry(ctx, usr)
+			if !exists {
+				if err != nil {
+					ctxLogger.Debug("Failed to fetch oauth entry", "error", err)
+					cmdErr = err
+				}
+				return
+			}
+
+			persistedToken = buildOAuthTokenFromAuthInfo(authInfo)
+		}
 
 		needRefresh := needTokenRefresh(ctx, persistedToken)
 		if !needRefresh {
@@ -259,7 +317,7 @@ func (o *Service) TryTokenRefresh(ctx context.Context, usr identity.Requester) (
 			return
 		}
 
-		newToken, cmdErr = o.tryGetOrRefreshOAuthToken(ctx, persistedToken, usr)
+		newToken, cmdErr = o.tryGetOrRefreshOAuthToken(ctx, persistedToken, usr, sessionToken)
 	}, retryOpt)
 	if lockErr != nil {
 		ctxLogger.Error("Failed to obtain token refresh lock", "error", lockErr)
@@ -274,45 +332,27 @@ func (o *Service) TryTokenRefresh(ctx context.Context, usr identity.Requester) (
 	return newToken, cmdErr
 }
 
-func buildOAuthTokenFromAuthInfo(authInfo *login.UserAuth) *oauth2.Token {
-	token := &oauth2.Token{
-		AccessToken:  authInfo.OAuthAccessToken,
-		Expiry:       authInfo.OAuthExpiry,
-		RefreshToken: authInfo.OAuthRefreshToken,
-		TokenType:    authInfo.OAuthTokenType,
-	}
-
-	if authInfo.OAuthIdToken != "" {
-		token = token.WithExtra(map[string]any{"id_token": authInfo.OAuthIdToken})
-	}
-
-	return token
-}
-
-func checkOAuthRefreshToken(authInfo *login.UserAuth) error {
-	if !strings.Contains(authInfo.AuthModule, "oauth") {
-		logger.Warn("The specified user's auth provider is not oauth",
-			"authmodule", authInfo.AuthModule, "userid", authInfo.UserId)
-		return ErrNotAnOAuthProvider
-	}
-
-	if authInfo.OAuthRefreshToken == "" {
-		logger.Warn("No refresh token available",
-			"authmodule", authInfo.AuthModule, "userid", authInfo.UserId)
-		return ErrNoRefreshTokenFound
-	}
-
-	return nil
-}
-
 // InvalidateOAuthTokens invalidates the OAuth tokens (access_token, refresh_token) and sets the Expiry to default/zero
-func (o *Service) InvalidateOAuthTokens(ctx context.Context, usr identity.Requester) error {
+func (o *Service) InvalidateOAuthTokens(ctx context.Context, usr identity.Requester, sessionToken *auth.UserToken) error {
 	userID, err := usr.GetInternalID()
 	if err != nil {
 		logger.Error("Failed to convert user id to int", "id", usr.GetID(), "error", err)
 		return err
 	}
 
+	ctxLogger := logger.FromContext(ctx).New("userID", userID)
+
+	if o.features.IsEnabledGlobally(featuremgmt.FlagImprovedExternalSessionHandling) {
+		err := o.sessionService.UpdateExternalSession(ctx, sessionToken.ExternalSessionId, &auth.UpdateExternalSessionCommand{
+			Token: &oauth2.Token{},
+		})
+		if err != nil {
+			ctxLogger.Error("Failed to update external session", "error", err)
+			return err
+		}
+	}
+
+	// TODO: Should this run regardless of the feature flag?
 	return o.AuthInfoService.UpdateAuthInfo(ctx, &login.UpdateAuthInfoCommand{
 		UserId:     userID,
 		AuthModule: usr.GetAuthenticatedBy(),
@@ -325,7 +365,7 @@ func (o *Service) InvalidateOAuthTokens(ctx context.Context, usr identity.Reques
 	})
 }
 
-func (o *Service) tryGetOrRefreshOAuthToken(ctx context.Context, persistedToken *oauth2.Token, usr identity.Requester) (*oauth2.Token, error) {
+func (o *Service) tryGetOrRefreshOAuthToken(ctx context.Context, persistedToken *oauth2.Token, usr identity.Requester, sessionToken *auth.UserToken) (*oauth2.Token, error) {
 	ctx, span := o.tracer.Start(ctx, "oauthtoken.tryGetOrRefreshOAuthToken")
 	defer span.End()
 
@@ -374,7 +414,7 @@ func (o *Service) tryGetOrRefreshOAuthToken(ctx context.Context, persistedToken
 			"provider", usr.GetAuthenticatedBy(), "error", err)
 
 		// token refresh failed, invalidate the old token
-		if err := o.InvalidateOAuthTokens(ctx, usr); err != nil {
+		if err := o.InvalidateOAuthTokens(ctx, usr, sessionToken); err != nil {
 			ctxLogger.Warn("Failed to invalidate OAuth tokens", "authID", usr.GetAuthID(), "error", err)
 		}
 
@@ -399,10 +439,20 @@ func (o *Service) tryGetOrRefreshOAuthToken(ctx context.Context, persistedToken
 			)
 		}
 
-		if err := o.AuthInfoService.UpdateAuthInfo(ctx, updateAuthCommand); err != nil {
-			ctxLogger.Error("Failed to update auth info during token refresh", "authID", usr.GetAuthID(), "error", err)
-			return token, err
+		if o.features.IsEnabledGlobally(featuremgmt.FlagImprovedExternalSessionHandling) {
+			if err := o.sessionService.UpdateExternalSession(ctx, sessionToken.ExternalSessionId, &auth.UpdateExternalSessionCommand{
+				Token: token,
+			}); err != nil {
+				ctxLogger.Error("Failed to update external session during token refresh", "error", err)
+				return token, err
+			}
+		} else {
+			if err := o.AuthInfoService.UpdateAuthInfo(ctx, updateAuthCommand); err != nil {
+				ctxLogger.Error("Failed to update auth info during token refresh", "authID", usr.GetAuthID(), "error", err)
+				return token, err
+			}
 		}
+
 		ctxLogger.Debug("Updated oauth info for user")
 	}
 
@@ -467,6 +517,51 @@ func needTokenRefresh(ctx context.Context, persistedToken *oauth2.Token) bool {
 	return true
 }
 
+func buildOAuthTokenFromAuthInfo(authInfo *login.UserAuth) *oauth2.Token {
+	token := &oauth2.Token{
+		AccessToken:  authInfo.OAuthAccessToken,
+		Expiry:       authInfo.OAuthExpiry,
+		RefreshToken: authInfo.OAuthRefreshToken,
+		TokenType:    authInfo.OAuthTokenType,
+	}
+
+	if authInfo.OAuthIdToken != "" {
+		token = token.WithExtra(map[string]any{"id_token": authInfo.OAuthIdToken})
+	}
+
+	return token
+}
+
+func buildOAuthTokenFromExternalSession(externalSession *auth.ExternalSession) *oauth2.Token {
+	token := &oauth2.Token{
+		AccessToken:  externalSession.AccessToken,
+		Expiry:       externalSession.ExpiresAt,
+		RefreshToken: externalSession.RefreshToken,
+	}
+
+	if externalSession.IDToken != "" {
+		token = token.WithExtra(map[string]any{"id_token": externalSession.IDToken})
+	}
+
+	return token
+}
+
+func checkOAuthRefreshToken(authInfo *login.UserAuth) error {
+	if !strings.Contains(authInfo.AuthModule, "oauth") {
+		logger.Warn("The specified user's auth provider is not oauth",
+			"authmodule", authInfo.AuthModule, "userid", authInfo.UserId)
+		return ErrNotAnOAuthProvider
+	}
+
+	if authInfo.OAuthRefreshToken == "" {
+		logger.Warn("No refresh token available",
+			"authmodule", authInfo.AuthModule, "userid", authInfo.UserId)
+		return ErrNoRefreshTokenFound
+	}
+
+	return nil
+}
+
 // GetIDTokenExpiry extracts the expiry time from the ID token
 func GetIDTokenExpiry(token *oauth2.Token) (time.Time, error) {
 	idToken, ok := token.Extra("id_token").(string)

pkg/services/oauthtoken/oauth_token_test.go

@@ -3,32 +3,29 @@ package oauthtoken
 import (
 	"context"
 	"errors"
-	"reflect"
 	"testing"
 	"time"
 
 	"github.com/grafana/authlib/claims"
-	"github.com/prometheus/client_golang/prometheus"
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/mock"
-	"golang.org/x/oauth2"
-
 	"github.com/grafana/grafana/pkg/apimachinery/identity"
 	"github.com/grafana/grafana/pkg/infra/db"
-	"github.com/grafana/grafana/pkg/infra/remotecache"
 	"github.com/grafana/grafana/pkg/infra/serverlock"
 	"github.com/grafana/grafana/pkg/infra/tracing"
 	"github.com/grafana/grafana/pkg/login/social"
 	"github.com/grafana/grafana/pkg/login/social/socialtest"
+	"github.com/grafana/grafana/pkg/models/usertoken"
+	"github.com/grafana/grafana/pkg/services/auth"
+	"github.com/grafana/grafana/pkg/services/auth/authtest"
 	"github.com/grafana/grafana/pkg/services/authn"
+	"github.com/grafana/grafana/pkg/services/featuremgmt"
 	"github.com/grafana/grafana/pkg/services/login"
-	"github.com/grafana/grafana/pkg/services/login/authinfoimpl"
 	"github.com/grafana/grafana/pkg/services/login/authinfotest"
-	"github.com/grafana/grafana/pkg/services/secrets/fakes"
-	secretsManager "github.com/grafana/grafana/pkg/services/secrets/manager"
-	"github.com/grafana/grafana/pkg/services/user"
 	"github.com/grafana/grafana/pkg/setting"
 	"github.com/grafana/grafana/pkg/tests/testsuite"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"golang.org/x/oauth2"
 )
 
 const EXPIRED_ID_TOKEN = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJodHRwczovL2V4YW1wbGUuY29tIiwic3ViIjoiMTIzNDU2Nzg5MCIsImF1ZCI6InlvdXItY2xpZW50LWlkIiwiZXhwIjoxNjAwMDAwMDAwLCJpYXQiOjE2MDAwMDAwMDAsIm5hbWUiOiJKb2huIERvZSIsImVtYWlsIjoiam9obkBleGFtcGxlLmNvbSJ9.c2lnbmF0dXJl" // #nosec G101 not a hardcoded credential
@@ -39,104 +36,6 @@ func TestMain(m *testing.M) {
 	testsuite.Run(m)
 }
 
-func TestService_HasOAuthEntry(t *testing.T) {
-	testCases := []struct {
-		name            string
-		user            *user.SignedInUser
-		want            *login.UserAuth
-		wantExist       bool
-		wantErr         bool
-		err             error
-		getAuthInfoErr  error
-		getAuthInfoUser login.UserAuth
-	}{
-		{
-			name:      "returns false without an error in case user is nil",
-			user:      nil,
-			want:      nil,
-			wantExist: false,
-			wantErr:   false,
-		},
-		{
-			name:           "returns false and an error in case GetAuthInfo returns an error",
-			user:           &user.SignedInUser{UserID: 1},
-			want:           nil,
-			wantExist:      false,
-			wantErr:        true,
-			getAuthInfoErr: errors.New("error"),
-		},
-		{
-			name:           "returns false without an error in case auth entry is not found",
-			user:           &user.SignedInUser{UserID: 1},
-			want:           nil,
-			wantExist:      false,
-			wantErr:        false,
-			getAuthInfoErr: user.ErrUserNotFound,
-		},
-		{
-			name:            "returns false without an error in case the auth entry is not oauth",
-			user:            &user.SignedInUser{UserID: 1},
-			want:            nil,
-			wantExist:       false,
-			wantErr:         false,
-			getAuthInfoUser: login.UserAuth{AuthModule: "auth_saml"},
-		},
-		{
-			name:            "returns true when the auth entry is found",
-			user:            &user.SignedInUser{UserID: 1},
-			want:            &login.UserAuth{AuthModule: login.GenericOAuthModule},
-			wantExist:       true,
-			wantErr:         false,
-			getAuthInfoUser: login.UserAuth{AuthModule: login.GenericOAuthModule},
-		},
-	}
-	for _, tc := range testCases {
-		tc := tc
-		t.Run(tc.name, func(t *testing.T) {
-			srv, authInfoStore, _ := setupOAuthTokenService(t)
-			authInfoStore.ExpectedOAuth = &tc.getAuthInfoUser
-			authInfoStore.ExpectedError = tc.getAuthInfoErr
-
-			entry, exists, err := srv.HasOAuthEntry(context.Background(), tc.user)
-
-			if tc.wantErr {
-				assert.Error(t, err)
-			}
-
-			if tc.want != nil {
-				assert.True(t, reflect.DeepEqual(tc.want, entry))
-			}
-			assert.Equal(t, tc.wantExist, exists)
-		})
-	}
-}
-
-func setupOAuthTokenService(t *testing.T) (*Service, *FakeAuthInfoStore, *socialtest.MockSocialConnector) {
-	t.Helper()
-
-	socialConnector := &socialtest.MockSocialConnector{}
-	socialService := &socialtest.FakeSocialService{
-		ExpectedConnector: socialConnector,
-		ExpectedAuthInfoProvider: &social.OAuthInfo{
-			UseRefreshToken: true,
-		},
-	}
-
-	authInfoStore := &FakeAuthInfoStore{ExpectedOAuth: &login.UserAuth{}}
-	authInfoService := authinfoimpl.ProvideService(authInfoStore, remotecache.NewFakeCacheStorage(), secretsManager.SetupTestService(t, fakes.NewFakeSecretsStore()))
-
-	store := db.InitTestDB(t)
-
-	return &Service{
-		Cfg:                  setting.NewCfg(),
-		SocialService:        socialService,
-		AuthInfoService:      authInfoService,
-		serverLock:           serverlock.ProvideService(store, tracing.InitializeTracerForTest()),
-		tokenRefreshDuration: newTokenRefreshDurationMetric(prometheus.NewRegistry()),
-		tracer:               tracing.InitializeTracerForTest(),
-	}, authInfoStore, socialConnector
-}
-
 type FakeAuthInfoStore struct {
 	login.Store
 	ExpectedError error
@@ -379,10 +278,12 @@ func TestService_TryTokenRefresh(t *testing.T) {
 				prometheus.NewRegistry(),
 				env.serverLock,
 				tracing.InitializeTracerForTest(),
+				nil,
+				featuremgmt.WithFeatures(),
 			)
 
 			// token refresh
-			actualToken, err := env.service.TryTokenRefresh(context.Background(), tt.identity)
+			actualToken, err := env.service.TryTokenRefresh(context.Background(), tt.identity, nil)
 
 			if tt.expectedErr != nil {
 				assert.ErrorIs(t, err, tt.expectedErr)
@@ -407,6 +308,275 @@ func TestService_TryTokenRefresh(t *testing.T) {
 	}
 }
 
+func TestService_TryTokenRefresh_WithExternalSessions(t *testing.T) {
+	unexpiredToken := &oauth2.Token{
+		AccessToken:  "testaccess",
+		RefreshToken: "testrefresh",
+		Expiry:       time.Now().Add(time.Hour),
+		TokenType:    "Bearer",
+	}
+	unexpiredTokenWithIDToken := unexpiredToken.WithExtra(map[string]interface{}{
+		"id_token": UNEXPIRED_ID_TOKEN,
+	})
+
+	expiredToken := &oauth2.Token{
+		AccessToken:  "testaccess",
+		RefreshToken: "testrefresh",
+		Expiry:       time.Now().Add(-time.Hour),
+		TokenType:    "Bearer",
+	}
+
+	userIdentity := &authn.Identity{
+		AuthenticatedBy: login.GenericOAuthModule,
+		ID:              "1234",
+		Type:            claims.TypeUser,
+	}
+
+	type environment struct {
+		sessionService  *authtest.MockUserAuthTokenService
+		serverLock      *serverlock.ServerLockService
+		socialConnector *socialtest.MockSocialConnector
+		socialService   *socialtest.FakeSocialService
+
+		service *Service
+	}
+
+	type testCase struct {
+		desc          string
+		identity      identity.Requester
+		setup         func(env *environment)
+		expectedToken *oauth2.Token
+		expectedErr   error
+	}
+
+	tests := []testCase{
+		{
+			desc: "should skip sync when identity is nil",
+		},
+		{
+			desc:     "should skip sync when identity is not a user",
+			identity: &authn.Identity{ID: "1", Type: claims.TypeServiceAccount},
+		},
+		{
+			desc:     "should skip token refresh and return nil if namespace and id cannot be converted to user ID",
+			identity: &authn.Identity{ID: "invalid", Type: claims.TypeUser},
+		},
+		{
+			desc:     "should skip token refresh if there's an unexpected error while looking up the user oauth entry, additionally, no error should be returned",
+			identity: userIdentity,
+			setup: func(env *environment) {
+				env.sessionService.On("GetExternalSession", mock.Anything, int64(1)).Return(nil, assert.AnError).Once()
+
+				env.socialService.ExpectedAuthInfoProvider = &social.OAuthInfo{
+					UseRefreshToken: true,
+				}
+			},
+		},
+		// Kinda impossible to happen, can only happen after the feature is enabled and logged in users don't have their external sessions set
+		{
+			desc:     "should skip token refresh if the user doesn't have an external session",
+			identity: userIdentity,
+			setup: func(env *environment) {
+				env.sessionService.On("GetExternalSession", mock.Anything, int64(1)).Return(nil, auth.ErrExternalSessionNotFound).Once()
+
+				env.socialService.ExpectedAuthInfoProvider = &social.OAuthInfo{
+					UseRefreshToken: true,
+				}
+			},
+		},
+		{
+			desc:     "should skip token refresh when no oauth provider was found",
+			identity: userIdentity,
+			setup: func(env *environment) {
+				env.socialService.ExpectedAuthInfoProvider = nil
+			},
+		},
+		{
+			desc:     "should skip token refresh when oauth provider token handling is disabled (UseRefreshToken is false)",
+			identity: userIdentity,
+			setup: func(env *environment) {
+				env.socialService.ExpectedAuthInfoProvider = &social.OAuthInfo{
+					UseRefreshToken: false,
+				}
+			},
+		},
+		{
+			desc:     "should skip token refresh when the token is still valid and no id token is present",
+			identity: userIdentity,
+			setup: func(env *environment) {
+				env.sessionService.On("GetExternalSession", mock.Anything, int64(1)).Return(&auth.ExternalSession{
+					ID:           1,
+					UserID:       1,
+					AccessToken:  unexpiredTokenWithIDToken.AccessToken,
+					RefreshToken: unexpiredTokenWithIDToken.RefreshToken,
+					ExpiresAt:    unexpiredTokenWithIDToken.Expiry,
+				}, nil).Once()
+
+				env.socialService.ExpectedAuthInfoProvider = &social.OAuthInfo{
+					UseRefreshToken: true,
+				}
+			},
+			expectedToken: unexpiredToken,
+		},
+		{
+			desc:     "should not do token refresh if access token or id token have not expired yet",
+			identity: userIdentity,
+			setup: func(env *environment) {
+				env.sessionService.On("GetExternalSession", mock.Anything, int64(1)).Return(&auth.ExternalSession{
+					ID:           1,
+					UserID:       1,
+					AccessToken:  unexpiredTokenWithIDToken.AccessToken,
+					RefreshToken: unexpiredTokenWithIDToken.RefreshToken,
+					ExpiresAt:    unexpiredTokenWithIDToken.Expiry,
+					IDToken:      UNEXPIRED_ID_TOKEN,
+				}, nil).Once()
+
+				env.socialService.ExpectedAuthInfoProvider = &social.OAuthInfo{
+					UseRefreshToken: true,
+				}
+			},
+			expectedToken: unexpiredTokenWithIDToken,
+		},
+		{
+			desc:     "should skip token refresh when there is no refresh token",
+			identity: userIdentity,
+			setup: func(env *environment) {
+				env.sessionService.On("GetExternalSession", mock.Anything, int64(1)).Return(&auth.ExternalSession{
+					ID:           1,
+					UserID:       1,
+					AccessToken:  unexpiredTokenWithIDToken.AccessToken,
+					RefreshToken: "",
+					ExpiresAt:    unexpiredTokenWithIDToken.Expiry,
+				}, nil).Once()
+
+				env.socialService.ExpectedAuthInfoProvider = &social.OAuthInfo{
+					UseRefreshToken: true,
+				}
+			},
+			expectedToken: &oauth2.Token{
+				AccessToken:  unexpiredTokenWithIDToken.AccessToken,
+				RefreshToken: "",
+				Expiry:       unexpiredTokenWithIDToken.Expiry,
+			},
+		},
+		{
+			desc: "should refresh token when the access token is expired",
+			identity: &authn.Identity{
+				AuthenticatedBy: login.GenericOAuthModule,
+				ID:              "1",
+				Type:            claims.TypeUser,
+			},
+			setup: func(env *environment) {
+				env.sessionService.On("GetExternalSession", mock.Anything, int64(1)).Return(&auth.ExternalSession{
+					ID:           1,
+					UserID:       1,
+					AccessToken:  expiredToken.AccessToken,
+					IDToken:      UNEXPIRED_ID_TOKEN,
+					RefreshToken: expiredToken.RefreshToken,
+					ExpiresAt:    expiredToken.Expiry,
+				}, nil).Once()
+
+				env.sessionService.On("UpdateExternalSession", mock.Anything, int64(1), mock.MatchedBy(verifyUpdateExternalSessionCommand(unexpiredTokenWithIDToken))).Return(nil).Once()
+
+				env.socialConnector.On("TokenSource", mock.Anything, mock.Anything).Return(oauth2.StaticTokenSource(unexpiredTokenWithIDToken)).Once()
+
+				env.socialService.ExpectedAuthInfoProvider = &social.OAuthInfo{
+					UseRefreshToken: true,
+				}
+			},
+			expectedToken: unexpiredTokenWithIDToken,
+		},
+		{
+			desc:     "should refresh token when the id token is expired",
+			identity: userIdentity,
+			setup: func(env *environment) {
+				env.sessionService.On("GetExternalSession", mock.Anything, int64(1)).Return(&auth.ExternalSession{
+					ID:           1,
+					UserID:       1,
+					AccessToken:  unexpiredTokenWithIDToken.AccessToken,
+					RefreshToken: unexpiredTokenWithIDToken.RefreshToken,
+					ExpiresAt:    unexpiredTokenWithIDToken.Expiry,
+					IDToken:      EXPIRED_ID_TOKEN,
+				}, nil).Once()
+
+				env.socialService.ExpectedAuthInfoProvider = &social.OAuthInfo{
+					UseRefreshToken: true,
+				}
+
+				env.sessionService.On("UpdateExternalSession", mock.Anything, int64(1), mock.MatchedBy(verifyUpdateExternalSessionCommand(unexpiredTokenWithIDToken))).Return(nil).Once()
+
+				env.socialConnector.On("TokenSource", mock.Anything, mock.Anything).Return(oauth2.StaticTokenSource(unexpiredTokenWithIDToken)).Once()
+			},
+			expectedToken: unexpiredTokenWithIDToken,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.desc, func(t *testing.T) {
+			socialConnector := socialtest.NewMockSocialConnector(t)
+
+			store := db.InitTestDB(t)
+
+			env := environment{
+				sessionService:  authtest.NewMockUserAuthTokenService(t),
+				serverLock:      serverlock.ProvideService(store, tracing.InitializeTracerForTest()),
+				socialConnector: socialConnector,
+				socialService: &socialtest.FakeSocialService{
+					ExpectedConnector: socialConnector,
+				},
+			}
+
+			if tt.setup != nil {
+				tt.setup(&env)
+			}
+
+			env.service = ProvideService(
+				env.socialService,
+				nil,
+				setting.NewCfg(),
+				prometheus.NewRegistry(),
+				env.serverLock,
+				tracing.InitializeTracerForTest(),
+				env.sessionService,
+				featuremgmt.WithFeatures(featuremgmt.FlagImprovedExternalSessionHandling),
+			)
+
+			// token refresh
+			actualToken, err := env.service.TryTokenRefresh(context.Background(), tt.identity, &usertoken.UserToken{ExternalSessionId: 1})
+
+			if tt.expectedErr != nil {
+				assert.ErrorIs(t, err, tt.expectedErr)
+				return
+			}
+
+			assert.NoError(t, err)
+
+			if tt.expectedToken == nil {
+				assert.Nil(t, actualToken)
+				return
+			}
+
+			assert.Equal(t, tt.expectedToken.AccessToken, actualToken.AccessToken)
+			assert.Equal(t, tt.expectedToken.RefreshToken, actualToken.RefreshToken)
+			assert.Equal(t, tt.expectedToken.Expiry, actualToken.Expiry)
+			if tt.expectedToken.Extra("id_token") != nil {
+				assert.Equal(t, tt.expectedToken.Extra("id_token").(string), actualToken.Extra("id_token").(string))
+			} else {
+				assert.Nil(t, actualToken.Extra("id_token"))
+			}
+		})
+	}
+}
+
+func verifyUpdateExternalSessionCommand(token *oauth2.Token) func(*auth.UpdateExternalSessionCommand) bool {
+	return func(cmd *auth.UpdateExternalSessionCommand) bool {
+		idToken := cmd.Token.Extra("id_token")
+		return cmd.Token.AccessToken == token.AccessToken &&
+			cmd.Token.RefreshToken == token.RefreshToken &&
+			cmd.Token.Expiry == token.Expiry &&
+			idToken == token.Extra("id_token")
+	}
+}
+
 func TestOAuthTokenSync_needTokenRefresh(t *testing.T) {
 	tests := []struct {
 		name                     string

pkg/services/oauthtoken/oauthtokentest/mock.go

@@ -6,21 +6,20 @@ import (
 	"golang.org/x/oauth2"
 
 	"github.com/grafana/grafana/pkg/apimachinery/identity"
+	"github.com/grafana/grafana/pkg/services/auth"
 	"github.com/grafana/grafana/pkg/services/datasources"
-	"github.com/grafana/grafana/pkg/services/login"
 )
 
 type MockOauthTokenService struct {
-	GetCurrentOauthTokenFunc   func(ctx context.Context, usr identity.Requester) *oauth2.Token
+	GetCurrentOauthTokenFunc   func(ctx context.Context, usr identity.Requester, sessionToken *auth.UserToken) *oauth2.Token
 	IsOAuthPassThruEnabledFunc func(ds *datasources.DataSource) bool
-	HasOAuthEntryFunc          func(ctx context.Context, usr identity.Requester) (*login.UserAuth, bool, error)
-	InvalidateOAuthTokensFunc  func(ctx context.Context, usr identity.Requester) error
-	TryTokenRefreshFunc        func(ctx context.Context, usr identity.Requester) (*oauth2.Token, error)
+	InvalidateOAuthTokensFunc  func(ctx context.Context, usr identity.Requester, sessionToken *auth.UserToken) error
+	TryTokenRefreshFunc        func(ctx context.Context, usr identity.Requester, sessionToken *auth.UserToken) (*oauth2.Token, error)
 }
 
-func (m *MockOauthTokenService) GetCurrentOAuthToken(ctx context.Context, usr identity.Requester) *oauth2.Token {
+func (m *MockOauthTokenService) GetCurrentOAuthToken(ctx context.Context, usr identity.Requester, sessionToken *auth.UserToken) *oauth2.Token {
 	if m.GetCurrentOauthTokenFunc != nil {
-		return m.GetCurrentOauthTokenFunc(ctx, usr)
+		return m.GetCurrentOauthTokenFunc(ctx, usr, sessionToken)
 	}
 	return nil
 }
@@ -32,23 +31,16 @@ func (m *MockOauthTokenService) IsOAuthPassThruEnabled(ds *datasources.DataSourc
 	return false
 }
 
-func (m *MockOauthTokenService) HasOAuthEntry(ctx context.Context, usr identity.Requester) (*login.UserAuth, bool, error) {
-	if m.HasOAuthEntryFunc != nil {
-		return m.HasOAuthEntryFunc(ctx, usr)
-	}
-	return nil, false, nil
-}
-
-func (m *MockOauthTokenService) InvalidateOAuthTokens(ctx context.Context, usr identity.Requester) error {
+func (m *MockOauthTokenService) InvalidateOAuthTokens(ctx context.Context, usr identity.Requester, sessionToken *auth.UserToken) error {
 	if m.InvalidateOAuthTokensFunc != nil {
-		return m.InvalidateOAuthTokensFunc(ctx, usr)
+		return m.InvalidateOAuthTokensFunc(ctx, usr, sessionToken)
 	}
 	return nil
 }
 
-func (m *MockOauthTokenService) TryTokenRefresh(ctx context.Context, usr identity.Requester) (*oauth2.Token, error) {
+func (m *MockOauthTokenService) TryTokenRefresh(ctx context.Context, usr identity.Requester, sessionToken *auth.UserToken) (*oauth2.Token, error) {
 	if m.TryTokenRefreshFunc != nil {
-		return m.TryTokenRefreshFunc(ctx, usr)
+		return m.TryTokenRefreshFunc(ctx, usr, sessionToken)
 	}
 	return nil, nil
 }

pkg/services/oauthtoken/oauthtokentest/oauthtokentest.go

@@ -6,8 +6,8 @@ import (
 	"golang.org/x/oauth2"
 
 	"github.com/grafana/grafana/pkg/apimachinery/identity"
+	"github.com/grafana/grafana/pkg/services/auth"
 	"github.com/grafana/grafana/pkg/services/datasources"
-	"github.com/grafana/grafana/pkg/services/login"
 	"github.com/grafana/grafana/pkg/services/oauthtoken"
 )
 
@@ -21,7 +21,7 @@ func ProvideService() *Service {
 	return &Service{}
 }
 
-func (s *Service) GetCurrentOAuthToken(context.Context, identity.Requester) *oauth2.Token {
+func (s *Service) GetCurrentOAuthToken(context.Context, identity.Requester, *auth.UserToken) *oauth2.Token {
 	return s.Token
 }
 
@@ -29,14 +29,10 @@ func (s *Service) IsOAuthPassThruEnabled(ds *datasources.DataSource) bool {
 	return oauthtoken.IsOAuthPassThruEnabled(ds)
 }
 
-func (s *Service) HasOAuthEntry(context.Context, identity.Requester) (*login.UserAuth, bool, error) {
-	return nil, false, nil
-}
-
-func (s *Service) TryTokenRefresh(context.Context, identity.Requester) (*oauth2.Token, error) {
+func (s *Service) TryTokenRefresh(context.Context, identity.Requester, *auth.UserToken) (*oauth2.Token, error) {
 	return s.Token, nil
 }
 
-func (s *Service) InvalidateOAuthTokens(context.Context, identity.Requester) error {
+func (s *Service) InvalidateOAuthTokens(context.Context, identity.Requester, *auth.UserToken) error {
 	return nil
 }

pkg/services/oauthtoken/oauthtokentest/service_mock.go

@@ -1,4 +1,4 @@
-// Code generated by mockery v2.40.1. DO NOT EDIT.
+// Code generated by mockery v2.42.1. DO NOT EDIT.
 
 package oauthtokentest
 
@@ -8,11 +8,11 @@ import (
 	identity "github.com/grafana/grafana/pkg/apimachinery/identity"
 	datasources "github.com/grafana/grafana/pkg/services/datasources"
 
-	login "github.com/grafana/grafana/pkg/services/login"
-
 	mock "github.com/stretchr/testify/mock"
 
 	oauth2 "golang.org/x/oauth2"
+
+	usertoken "github.com/grafana/grafana/pkg/models/usertoken"
 )
 
 // MockService is an autogenerated mock type for the OAuthTokenService type
@@ -20,17 +20,17 @@ type MockService struct {
 	mock.Mock
 }
 
-// GetCurrentOAuthToken provides a mock function with given fields: _a0, _a1
-func (_m *MockService) GetCurrentOAuthToken(_a0 context.Context, _a1 identity.Requester) *oauth2.Token {
-	ret := _m.Called(_a0, _a1)
+// GetCurrentOAuthToken provides a mock function with given fields: _a0, _a1, _a2
+func (_m *MockService) GetCurrentOAuthToken(_a0 context.Context, _a1 identity.Requester, _a2 *usertoken.UserToken) *oauth2.Token {
+	ret := _m.Called(_a0, _a1, _a2)
 
 	if len(ret) == 0 {
 		panic("no return value specified for GetCurrentOAuthToken")
 	}
 
 	var r0 *oauth2.Token
-	if rf, ok := ret.Get(0).(func(context.Context, identity.Requester) *oauth2.Token); ok {
-		r0 = rf(_a0, _a1)
+	if rf, ok := ret.Get(0).(func(context.Context, identity.Requester, *usertoken.UserToken) *oauth2.Token); ok {
+		r0 = rf(_a0, _a1, _a2)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).(*oauth2.Token)
@@ -40,54 +40,17 @@ func (_m *MockService) GetCurrentOAuthToken(_a0 context.Context, _a1 identity.Re
 	return r0
 }
 
-// HasOAuthEntry provides a mock function with given fields: _a0, _a1
-func (_m *MockService) HasOAuthEntry(_a0 context.Context, _a1 identity.Requester) (*login.UserAuth, bool, error) {
-	ret := _m.Called(_a0, _a1)
-
-	if len(ret) == 0 {
-		panic("no return value specified for HasOAuthEntry")
-	}
-
-	var r0 *login.UserAuth
-	var r1 bool
-	var r2 error
-	if rf, ok := ret.Get(0).(func(context.Context, identity.Requester) (*login.UserAuth, bool, error)); ok {
-		return rf(_a0, _a1)
-	}
-	if rf, ok := ret.Get(0).(func(context.Context, identity.Requester) *login.UserAuth); ok {
-		r0 = rf(_a0, _a1)
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(*login.UserAuth)
-		}
-	}
-
-	if rf, ok := ret.Get(1).(func(context.Context, identity.Requester) bool); ok {
-		r1 = rf(_a0, _a1)
-	} else {
-		r1 = ret.Get(1).(bool)
-	}
-
-	if rf, ok := ret.Get(2).(func(context.Context, identity.Requester) error); ok {
-		r2 = rf(_a0, _a1)
-	} else {
-		r2 = ret.Error(2)
-	}
-
-	return r0, r1, r2
-}
-
-// InvalidateOAuthTokens provides a mock function with given fields: _a0, _a1
-func (_m *MockService) InvalidateOAuthTokens(_a0 context.Context, _a1 *login.UserAuth) error {
-	ret := _m.Called(_a0, _a1)
+// InvalidateOAuthTokens provides a mock function with given fields: _a0, _a1, _a2
+func (_m *MockService) InvalidateOAuthTokens(_a0 context.Context, _a1 identity.Requester, _a2 *usertoken.UserToken) error {
+	ret := _m.Called(_a0, _a1, _a2)
 
 	if len(ret) == 0 {
 		panic("no return value specified for InvalidateOAuthTokens")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(context.Context, *login.UserAuth) error); ok {
-		r0 = rf(_a0, _a1)
+	if rf, ok := ret.Get(0).(func(context.Context, identity.Requester, *usertoken.UserToken) error); ok {
+		r0 = rf(_a0, _a1, _a2)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -113,22 +76,34 @@ func (_m *MockService) IsOAuthPassThruEnabled(_a0 *datasources.DataSource) bool
 	return r0
 }
 
-// TryTokenRefresh provides a mock function with given fields: _a0, _a1
-func (_m *MockService) TryTokenRefresh(_a0 context.Context, _a1 identity.Requester) error {
-	ret := _m.Called(_a0, _a1)
+// TryTokenRefresh provides a mock function with given fields: _a0, _a1, _a2
+func (_m *MockService) TryTokenRefresh(_a0 context.Context, _a1 identity.Requester, _a2 *usertoken.UserToken) (*oauth2.Token, error) {
+	ret := _m.Called(_a0, _a1, _a2)
 
 	if len(ret) == 0 {
 		panic("no return value specified for TryTokenRefresh")
 	}
 
-	var r0 error
-	if rf, ok := ret.Get(0).(func(context.Context, identity.Requester) error); ok {
-		r0 = rf(_a0, _a1)
+	var r0 *oauth2.Token
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, identity.Requester, *usertoken.UserToken) (*oauth2.Token, error)); ok {
+		return rf(_a0, _a1, _a2)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, identity.Requester, *usertoken.UserToken) *oauth2.Token); ok {
+		r0 = rf(_a0, _a1, _a2)
 	} else {
-		r0 = ret.Error(0)
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*oauth2.Token)
+		}
 	}
 
-	return r0
+	if rf, ok := ret.Get(1).(func(context.Context, identity.Requester, *usertoken.UserToken) error); ok {
+		r1 = rf(_a0, _a1, _a2)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
 }
 
 // NewMockService creates a new instance of MockService. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.

pkg/services/playlist/playlistimpl/xorm_store.go

@@ -222,6 +222,11 @@ func (s *sqlStore) ListAll(ctx context.Context, orgId int64) ([]playlist.Playlis
 	if err != nil {
 		return nil, err
 	}
+
+	defer func() {
+		_ = rows.Close()
+	}()
+
 	for rows.Next() {
 		err = rows.Scan(&playlistId, &itemType, &itemValue)
 		if err != nil {

pkg/services/pluginsintegration/clientmiddleware/oauthtoken_middleware.go

@@ -50,7 +50,7 @@ func (m *OAuthTokenMiddleware) applyToken(ctx context.Context, pCtx backend.Plug
 	}
 
 	if m.oAuthTokenService.IsOAuthPassThruEnabled(ds) {
-		if token := m.oAuthTokenService.GetCurrentOAuthToken(ctx, reqCtx.SignedInUser); token != nil {
+		if token := m.oAuthTokenService.GetCurrentOAuthToken(ctx, reqCtx.SignedInUser, reqCtx.UserToken); token != nil {
 			authorizationHeader := fmt.Sprintf("%s %s", token.Type(), token.AccessToken)
 			idTokenHeader := ""
 

pkg/services/serviceaccounts/database/store.go

@@ -509,11 +509,30 @@ func (s *ServiceAccountsStoreImpl) CreateServiceAccountFromApikey(ctx context.Co
 		IsServiceAccount: true,
 	}
 
+	// maximum number of attempts for creating a service account
+	attempts := 10
+
 	return s.sqlStore.InTransaction(ctx, func(tctx context.Context) error {
 		newSA, errCreateSA := s.userService.CreateServiceAccount(tctx, &cmd)
+		if errCreateSA != nil {
+			if errors.Is(errCreateSA, serviceaccounts.ErrServiceAccountAlreadyExists) {
+				// The service account we tried to create already exists with that login name. We will attempt to create
+				// a unique service account by adding suffixes to the initial login name (e.g. -001, -002, ... , -010).
+				for i := 1; errCreateSA != nil && i <= attempts; i++ {
+					serviceAccountName := fmt.Sprintf("%s-%03d", key.Name, i)
+					cmd.Login = generateLogin(prefix, key.OrgID, serviceAccountName)
+					newSA, errCreateSA = s.userService.CreateServiceAccount(tctx, &cmd)
+					if errCreateSA != nil && !errors.Is(errCreateSA, serviceaccounts.ErrServiceAccountAlreadyExists) {
+						break
+					}
+				}
+			}
+		}
+
 		if errCreateSA != nil {
 			return fmt.Errorf("failed to create service account: %w", errCreateSA)
 		}
+
 		return s.assignApiKeyToServiceAccount(tctx, key.ID, newSA.ID)
 	})
 }

pkg/services/serviceaccounts/database/store_test.go

@@ -324,26 +324,85 @@ func TestIntegrationStore_MigrateApiKeys(t *testing.T) {
 		t.Skip("skipping test in short mode")
 	}
 	cases := []struct {
-		desc        string
-		key         tests.TestApiKey
-		expectedErr error
+		desc            string
+		serviceAccounts []user.CreateUserCommand
+		key             tests.TestApiKey
+		expectedLogin   string
+		expectedErr     error
 	}{
 		{
-			desc:        "api key should be migrated to service account token",
-			key:         tests.TestApiKey{Name: "Test1", Role: org.RoleEditor, OrgId: 1},
-			expectedErr: nil,
+			desc:            "api key should be migrated to service account token",
+			serviceAccounts: []user.CreateUserCommand{},
+			key:             tests.TestApiKey{Name: "test1", Role: org.RoleEditor, OrgId: 1},
+			expectedLogin:   "sa-autogen-1-test1",
+			expectedErr:     nil,
+		},
+		{
+			desc: "api key should be migrated to service account token on second attempt",
+			serviceAccounts: []user.CreateUserCommand{
+				{Login: "sa-autogen-1-test2"},
+			},
+			key:           tests.TestApiKey{Name: "test2", Role: org.RoleEditor, OrgId: 1},
+			expectedLogin: "sa-autogen-1-test2-001",
+			expectedErr:   nil,
+		},
+		{
+			desc: "api key should be migrated to service account token on last attempt (the 10th)",
+			serviceAccounts: []user.CreateUserCommand{
+				{Login: "sa-autogen-1-test3"},
+				{Login: "sa-autogen-1-test3-001"},
+				{Login: "sa-autogen-1-test3-002"},
+				{Login: "sa-autogen-1-test3-003"},
+				{Login: "sa-autogen-1-test3-004"},
+				{Login: "sa-autogen-1-test3-005"},
+				{Login: "sa-autogen-1-test3-006"},
+				{Login: "sa-autogen-1-test3-007"},
+				{Login: "sa-autogen-1-test3-008"},
+				{Login: "sa-autogen-1-test3-009"},
+			},
+			key:           tests.TestApiKey{Name: "test3", Role: org.RoleEditor, OrgId: 1},
+			expectedLogin: "sa-autogen-1-test3-010",
+			expectedErr:   nil,
+		},
+		{
+			desc: "api key should not be migrated to service account token because all attempts failed",
+			serviceAccounts: []user.CreateUserCommand{
+				{Login: "sa-autogen-1-test4"},
+				{Login: "sa-autogen-1-test4-001"},
+				{Login: "sa-autogen-1-test4-002"},
+				{Login: "sa-autogen-1-test4-003"},
+				{Login: "sa-autogen-1-test4-004"},
+				{Login: "sa-autogen-1-test4-005"},
+				{Login: "sa-autogen-1-test4-006"},
+				{Login: "sa-autogen-1-test4-007"},
+				{Login: "sa-autogen-1-test4-008"},
+				{Login: "sa-autogen-1-test4-009"},
+				{Login: "sa-autogen-1-test4-010"},
+			},
+			key:         tests.TestApiKey{Name: "test4", Role: org.RoleEditor, OrgId: 1},
+			expectedErr: serviceaccounts.ErrServiceAccountAlreadyExists,
 		},
 	}
 
 	for _, c := range cases {
 		t.Run(c.desc, func(t *testing.T) {
 			db, store := setupTestDatabase(t)
+
 			store.cfg.AutoAssignOrg = true
 			store.cfg.AutoAssignOrgId = 1
 			store.cfg.AutoAssignOrgRole = "Viewer"
 			_, err := store.orgService.CreateWithMember(context.Background(), &org.CreateOrgCommand{Name: "main"})
 			require.NoError(t, err)
+
 			key := tests.SetupApiKey(t, db, store.cfg, c.key)
+
+			for _, sa := range c.serviceAccounts {
+				sa.IsServiceAccount = true
+				sa.OrgID = key.OrgID
+				_, err := store.userService.CreateServiceAccount(context.Background(), &sa)
+				require.NoError(t, err)
+			}
+
 			err = store.MigrateApiKey(context.Background(), key.OrgID, key.ID)
 			if c.expectedErr != nil {
 				require.ErrorIs(t, err, c.expectedErr)
@@ -352,7 +411,7 @@ func TestIntegrationStore_MigrateApiKeys(t *testing.T) {
 
 				q := serviceaccounts.SearchOrgServiceAccountsQuery{
 					OrgID: key.OrgID,
-					Query: "",
+					Query: c.expectedLogin,
 					Page:  1,
 					Limit: 50,
 					SignedInUser: &user.SignedInUser{
@@ -370,6 +429,7 @@ func TestIntegrationStore_MigrateApiKeys(t *testing.T) {
 				require.Equal(t, int64(1), serviceAccounts.TotalCount)
 				saMigrated := serviceAccounts.ServiceAccounts[0]
 				require.Equal(t, string(key.Role), saMigrated.Role)
+				require.Equal(t, c.expectedLogin, saMigrated.Login)
 
 				tokens, err := store.ListTokens(context.Background(), &serviceaccounts.GetSATokensQuery{
 					OrgID:            &key.OrgID,