From 8490fd77e31f093ddafc8bcc65327e32d31ab7b9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Piotr=20Jamr=C3=B3z?= <pm.jamroz@gmail.com>
Date: Wed, 6 Apr 2022 08:35:58 +0200
Subject: [PATCH] Users: clean up OrgId when all user orgs are removed (#46003)

* Clean up orgId when user organization is removed

* Add a test for removing user org

* Fix linting errors

* Update comment

* Fix linting errors

* Make removing user org more explicit
---
 pkg/models/user.go                      |  2 +-
 pkg/services/sqlstore/org_users.go      |  6 ++++
 pkg/services/sqlstore/org_users_test.go | 47 +++++++++++++++++++++++++
 pkg/services/sqlstore/user.go           | 10 ++++++
 4 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/pkg/models/user.go b/pkg/models/user.go
index 343ce25e85d..4d4b2bace22 100644
--- a/pkg/models/user.go
+++ b/pkg/models/user.go
@@ -233,7 +233,7 @@ type UserProfileDTO struct {
 	Name           string          `json:"name"`
 	Login          string          `json:"login"`
 	Theme          string          `json:"theme"`
-	OrgId          int64           `json:"orgId"`
+	OrgId          int64           `json:"orgId,omitempty"`
 	IsGrafanaAdmin bool            `json:"isGrafanaAdmin"`
 	IsDisabled     bool            `json:"isDisabled"`
 	IsExternal     bool            `json:"isExternal"`
diff --git a/pkg/services/sqlstore/org_users.go b/pkg/services/sqlstore/org_users.go
index 5d695bc4e1d..24f1a9e8119 100644
--- a/pkg/services/sqlstore/org_users.go
+++ b/pkg/services/sqlstore/org_users.go
@@ -307,6 +307,12 @@ func (ss *SQLStore) RemoveOrgUser(ctx context.Context, cmd *models.RemoveOrgUser
 			}
 
 			cmd.UserWasDeleted = true
+		} else {
+			// no orgs, but keep the user -> clean up orgId
+			err = removeUserOrg(sess, user.Id)
+			if err != nil {
+				return err
+			}
 		}
 
 		return nil
diff --git a/pkg/services/sqlstore/org_users_test.go b/pkg/services/sqlstore/org_users_test.go
index a018393fd47..b68eb0d17c3 100644
--- a/pkg/services/sqlstore/org_users_test.go
+++ b/pkg/services/sqlstore/org_users_test.go
@@ -144,6 +144,53 @@ func TestSQLStore_SearchOrgUsers(t *testing.T) {
 	}
 }
 
+func TestSQLStore_RemoveOrgUser(t *testing.T) {
+	store := InitTestDB(t)
+
+	// create org and admin
+	_, err := store.CreateUser(context.Background(), models.CreateUserCommand{
+		Login: "admin",
+		OrgId: 1,
+	})
+	require.NoError(t, err)
+
+	// create a user with no org
+	_, err = store.CreateUser(context.Background(), models.CreateUserCommand{
+		Login:        "user",
+		OrgId:        1,
+		SkipOrgSetup: true,
+	})
+	require.NoError(t, err)
+
+	// assign the user to the org
+	err = store.AddOrgUser(context.Background(), &models.AddOrgUserCommand{
+		Role:   "Viewer",
+		OrgId:  1,
+		UserId: 2,
+	})
+	require.NoError(t, err)
+
+	// assert the org has been assigned
+	user := &models.GetUserByIdQuery{Id: 2}
+	err = store.GetUserById(context.Background(), user)
+	require.NoError(t, err)
+	require.Equal(t, user.Result.OrgId, int64(1))
+
+	// remove the user org
+	err = store.RemoveOrgUser(context.Background(), &models.RemoveOrgUserCommand{
+		UserId:                   2,
+		OrgId:                    1,
+		ShouldDeleteOrphanedUser: false,
+	})
+	require.NoError(t, err)
+
+	// assert the org has been removed
+	user = &models.GetUserByIdQuery{Id: 2}
+	err = store.GetUserById(context.Background(), user)
+	require.NoError(t, err)
+	require.Equal(t, user.Result.OrgId, int64(0))
+}
+
 func seedOrgUsers(t *testing.T, store *SQLStore, numUsers int) {
 	t.Helper()
 	// Seed users
diff --git a/pkg/services/sqlstore/user.go b/pkg/services/sqlstore/user.go
index fa943100445..c9ce4d9b098 100644
--- a/pkg/services/sqlstore/user.go
+++ b/pkg/services/sqlstore/user.go
@@ -445,6 +445,16 @@ func setUsingOrgInTransaction(sess *DBSession, userID int64, orgID int64) error
 	return err
 }
 
+func removeUserOrg(sess *DBSession, userID int64) error {
+	user := models.User{
+		Id:    userID,
+		OrgId: 0,
+	}
+
+	_, err := sess.ID(userID).MustCols("org_id").Update(&user)
+	return err
+}
+
 func (ss *SQLStore) GetUserProfile(ctx context.Context, query *models.GetUserProfileQuery) error {
 	return ss.WithDbSession(ctx, func(sess *DBSession) error {
 		var user models.User