環境変数設定時にdeadlockするのを修正

traPtitech · Apr 2, 2024 · 07499d1 · 07499d1
1 parent 62dbba3
commit 07499d1
Showing 1 changed file with 18 additions and 6 deletions.
diff --git a/pkg/infrastructure/repository/environment.go b/pkg/infrastructure/repository/environment.go
@@ -44,24 +44,36 @@ func (r *environmentRepository) GetEnv(ctx context.Context, cond domain.GetEnvCo
 	return ds.Map(environments, repoconvert.ToDomainEnvironment), nil
 }
 
-func (r *environmentRepository) SetEnv(ctx context.Context, env *domain.Environment) error {
+func (r *environmentRepository) SetEnv(ctx context.Context, env *domain.Environment) (err error) {
 	// NOTE: sqlboiler does not recognize multiple column unique keys: https://github.com/volatiletech/sqlboiler/issues/328
 	tx, err := r.db.BeginTx(ctx, nil)
 	if err != nil {
 		return errors.Wrap(err, "failed to start transaction")
 	}
+	defer func() {
+		if err != nil {
+			tx.Rollback()
+		}
+	}()
 
-	me, err := models.Environments(
-		models.EnvironmentWhere.ApplicationID.EQ(env.ApplicationID),
-		models.EnvironmentWhere.Key.EQ(env.Key),
+	_, err = models.Applications(
+		qm.Select(models.ApplicationColumns.ID),
+		models.ApplicationWhere.ID.EQ(env.ApplicationID),
 		qm.For("UPDATE"),
 	).One(ctx, tx)
+	if err != nil {
+		return errors.Wrap(err, "failed to get application")
+	}
+
+	exists, err := models.Environments(
+		models.EnvironmentWhere.ApplicationID.EQ(env.ApplicationID),
+		models.EnvironmentWhere.Key.EQ(env.Key),
+	).Exists(ctx, tx)
 	if err != nil && !isNoRowsErr(err) {
 		return errors.Wrap(err, "failed to get environment")
 	}
-	exists := !isNoRowsErr(err)
 
-	me = repoconvert.FromDomainEnvironment(env)
+	me := repoconvert.FromDomainEnvironment(env)
 
 	if exists {
 		_, err = me.Update(ctx, tx, boil.Blacklist())