From 1622a57f41dfb95f6b14da0823743f45ecccf876 Mon Sep 17 00:00:00 2001
From: Michael Mandrus <michael.mandrus@grafana.com>
Date: Wed, 8 Jun 2022 13:03:44 -0400
Subject: [PATCH] ensure atomicity by adding secret management as a step to sql
 operations and rolling back if necessary

---
 pkg/models/datasource.go                      |  7 ++++
 .../datasources/service/datasource_service.go | 34 +++++++++++--------
 pkg/services/sqlstore/datasource.go           | 24 +++++++++++++
 3 files changed, 51 insertions(+), 14 deletions(-)

diff --git a/pkg/models/datasource.go b/pkg/models/datasource.go
index dc2b3e2614a..52298d505b9 100644
--- a/pkg/models/datasource.go
+++ b/pkg/models/datasource.go
@@ -104,6 +104,7 @@ type AddDataSourceCommand struct {
 	UserId                  int64             `json:"-"`
 	ReadOnly                bool              `json:"-"`
 	EncryptedSecureJsonData map[string][]byte `json:"-"`
+	UpdateSecretFn          UpdateSecretFn    `json:"-"`
 
 	Result *DataSource `json:"-"`
 }
@@ -129,6 +130,7 @@ type UpdateDataSourceCommand struct {
 	Id                      int64             `json:"-"`
 	ReadOnly                bool              `json:"-"`
 	EncryptedSecureJsonData map[string][]byte `json:"-"`
+	UpdateSecretFn          UpdateSecretFn    `json:"-"`
 
 	Result *DataSource `json:"-"`
 }
@@ -143,8 +145,13 @@ type DeleteDataSourceCommand struct {
 	OrgID int64
 
 	DeletedDatasourcesCount int64
+
+	UpdateSecretFn UpdateSecretFn
 }
 
+// Function for updating secrets along with datasources, to ensure atomicity
+type UpdateSecretFn func() error
+
 // ---------------------
 // QUERIES
 
diff --git a/pkg/services/datasources/service/datasource_service.go b/pkg/services/datasources/service/datasource_service.go
index 1a12960103d..e7c7869d49c 100644
--- a/pkg/services/datasources/service/datasource_service.go
+++ b/pkg/services/datasources/service/datasource_service.go
@@ -152,17 +152,16 @@ func (s *Service) AddDataSource(ctx context.Context, cmd *models.AddDataSourceCo
 		return err
 	}
 
-	if err := s.SQLStore.AddDataSource(ctx, cmd); err != nil {
-		return err
-	}
-
 	secret, err := json.Marshal(cmd.SecureJsonData)
 	if err != nil {
 		return err
 	}
 
-	err = s.SecretsStore.Set(ctx, cmd.OrgId, cmd.Name, secretType, string(secret))
-	if err != nil {
+	cmd.UpdateSecretFn = func() error {
+		return s.SecretsStore.Set(ctx, cmd.OrgId, cmd.Name, secretType, string(secret))
+	}
+
+	if err := s.SQLStore.AddDataSource(ctx, cmd); err != nil {
 		return err
 	}
 
@@ -188,11 +187,15 @@ func (s *Service) AddDataSource(ctx context.Context, cmd *models.AddDataSourceCo
 }
 
 func (s *Service) DeleteDataSource(ctx context.Context, cmd *models.DeleteDataSourceCommand) error {
+	cmd.UpdateSecretFn = func() error {
+		return s.SecretsStore.Del(ctx, cmd.OrgID, cmd.Name, secretType)
+	}
+
 	err := s.SQLStore.DeleteDataSource(ctx, cmd)
 	if err != nil {
 		return err
 	}
-	return s.SecretsStore.Del(ctx, cmd.OrgID, cmd.Name, secretType)
+	return nil
 }
 
 func (s *Service) UpdateDataSource(ctx context.Context, cmd *models.UpdateDataSourceCommand) error {
@@ -212,11 +215,19 @@ func (s *Service) UpdateDataSource(ctx context.Context, cmd *models.UpdateDataSo
 		return err
 	}
 
-	err = s.SQLStore.UpdateDataSource(ctx, cmd)
+	secret, err := json.Marshal(cmd.SecureJsonData)
 	if err != nil {
 		return err
 	}
 
+	cmd.UpdateSecretFn = func() error {
+		return s.SecretsStore.Set(ctx, cmd.OrgId, cmd.Name, secretType, string(secret))
+	}
+
+	if err = s.SQLStore.UpdateDataSource(ctx, cmd); err != nil {
+		return err
+	}
+
 	if query.Result.Name != cmd.Name {
 		err = s.SecretsStore.Rename(ctx, cmd.OrgId, query.Result.Name, secretType, cmd.Name)
 		if err != nil {
@@ -224,12 +235,7 @@ func (s *Service) UpdateDataSource(ctx context.Context, cmd *models.UpdateDataSo
 		}
 	}
 
-	secret, err := json.Marshal(cmd.SecureJsonData)
-	if err != nil {
-		return err
-	}
-
-	return s.SecretsStore.Set(ctx, cmd.OrgId, cmd.Name, secretType, string(secret))
+	return nil
 }
 
 func (s *Service) GetDefaultDataSource(ctx context.Context, query *models.GetDefaultDataSourceQuery) error {
diff --git a/pkg/services/sqlstore/datasource.go b/pkg/services/sqlstore/datasource.go
index 1c831bec3ab..9d47ad292d9 100644
--- a/pkg/services/sqlstore/datasource.go
+++ b/pkg/services/sqlstore/datasource.go
@@ -116,6 +116,14 @@ func (ss *SQLStore) DeleteDataSource(ctx context.Context, cmd *models.DeleteData
 			}
 		}
 
+		if cmd.UpdateSecretFn != nil {
+			if err := cmd.UpdateSecretFn(); err != nil {
+				sess.Rollback()
+				sqlog.Error("Failed to update datasource secrets -- rolling back update", "UID", cmd.UID, "name", cmd.Name, "orgId", cmd.OrgID)
+				return err
+			}
+		}
+
 		// Publish data source deletion event
 		sess.publishAfterCommit(&events.DataSourceDeleted{
 			Timestamp: time.Now(),
@@ -181,6 +189,14 @@ func (ss *SQLStore) AddDataSource(ctx context.Context, cmd *models.AddDataSource
 			return err
 		}
 
+		if cmd.UpdateSecretFn != nil {
+			if err := cmd.UpdateSecretFn(); err != nil {
+				sess.Rollback()
+				sqlog.Error("Failed to update datasource secrets -- rolling back update", "name", cmd.Name, "type", cmd.Type, "orgId", cmd.OrgId)
+				return err
+			}
+		}
+
 		cmd.Result = ds
 
 		sess.publishAfterCommit(&events.DataSourceCreated{
@@ -263,6 +279,14 @@ func (ss *SQLStore) UpdateDataSource(ctx context.Context, cmd *models.UpdateData
 
 		err = updateIsDefaultFlag(ds, sess)
 
+		if cmd.UpdateSecretFn != nil {
+			if err := cmd.UpdateSecretFn(); err != nil {
+				sess.Rollback()
+				sqlog.Error("Failed to update datasource secrets -- rolling back update", "UID", cmd.Uid, "name", cmd.Name, "type", cmd.Type, "orgId", cmd.OrgId)
+				return err
+			}
+		}
+
 		cmd.Result = ds
 		return err
 	})