From 13f6b4686a7745f4b3f2cfe8c9767a76fbb85c40 Mon Sep 17 00:00:00 2001
From: Silvan <silvan.reusser@gmail.com>
Date: Tue, 25 Apr 2023 10:12:53 +0200
Subject: [PATCH] fix(setup): smaller transactions (#5742)

---
 cmd/setup/10.go      | 48 ++++++++++++++++++++++----------------------
 cmd/setup/10.sql     |  2 --
 cmd/setup/config.go  | 20 +++++++++---------
 cmd/setup/setup.go   |  4 ++--
 cmd/setup/steps.yaml |  2 ++
 5 files changed, 38 insertions(+), 38 deletions(-)

diff --git a/cmd/setup/10.go b/cmd/setup/10.go
index 77db7c6c5e..4efccc6537 100644
--- a/cmd/setup/10.go
+++ b/cmd/setup/10.go
@@ -2,8 +2,11 @@ package setup
 
 import (
 	"context"
+	"database/sql"
 	_ "embed"
+	"time"
 
+	"github.com/cockroachdb/cockroach-go/v2/crdb"
 	"github.com/zitadel/logging"
 
 	"github.com/zitadel/zitadel/internal/database"
@@ -15,35 +18,32 @@ var (
 )
 
 type CorrectCreationDate struct {
-	dbClient *database.DB
+	dbClient  *database.DB
+	FailAfter time.Duration
 }
 
 func (mig *CorrectCreationDate) Execute(ctx context.Context) (err error) {
-	tx, err := mig.dbClient.Begin()
-	if err != nil {
-		return err
-	}
-	if mig.dbClient.Type() == "cockroach" {
-		if _, err := tx.Exec("SET experimental_enable_temp_tables=on"); err != nil {
-			return err
-		}
-	}
-	defer func() {
-		if err != nil {
-			logging.OnError(tx.Rollback()).Debug("rollback failed")
-			return
-		}
-		err = tx.Commit()
-	}()
+	ctx, cancel := context.WithTimeout(ctx, mig.FailAfter)
+	defer cancel()
+
 	for {
-		res, err := tx.ExecContext(ctx, correctCreationDate10)
-		if err != nil {
-			return err
-		}
-		affected, _ := res.RowsAffected()
-		logging.WithFields("count", affected).Info("creation dates changed")
-		if affected == 0 {
+		var affected int64
+		err = crdb.ExecuteTx(ctx, mig.dbClient.DB, nil, func(tx *sql.Tx) error {
+			if mig.dbClient.Type() == "cockroach" {
+				if _, err := tx.Exec("SET experimental_enable_temp_tables=on"); err != nil {
+					return err
+				}
+			}
+			res, err := tx.ExecContext(ctx, correctCreationDate10)
+			if err != nil {
+				return err
+			}
+			affected, _ = res.RowsAffected()
+			logging.WithFields("count", affected).Info("creation dates changed")
 			return nil
+		})
+		if affected == 0 || err != nil {
+			return err
 		}
 	}
 }
diff --git a/cmd/setup/10.sql b/cmd/setup/10.sql
index 27853f55a7..591d0b3d88 100644
--- a/cmd/setup/10.sql
+++ b/cmd/setup/10.sql
@@ -5,8 +5,6 @@ CREATE temporary TABLE IF NOT EXISTS wrong_events (
     , next_cd TIMESTAMPTZ
 );
 
-TRUNCATE wrong_events;
-
 INSERT INTO wrong_events (
     SELECT * FROM (
         SELECT
diff --git a/cmd/setup/config.go b/cmd/setup/config.go
index eb55ba492a..dc75d54048 100644
--- a/cmd/setup/config.go
+++ b/cmd/setup/config.go
@@ -56,16 +56,16 @@ func MustNewConfig(v *viper.Viper) *Config {
 }
 
 type Steps struct {
-	s1ProjectionTable         *ProjectionTable
-	s2AssetsTable             *AssetTable
-	FirstInstance             *FirstInstance
-	s4EventstoreIndexes       *EventstoreIndexesNew
-	s5LastFailed              *LastFailed
-	s6OwnerRemoveColumns      *OwnerRemoveColumns
-	s7LogstoreTables          *LogstoreTables
-	s8AuthTokens              *AuthTokenIndexes
-	s9EventstoreIndexes2      *EventstoreIndexesNew
-	s10EventstoreCreationDate *CorrectCreationDate
+	s1ProjectionTable    *ProjectionTable
+	s2AssetsTable        *AssetTable
+	FirstInstance        *FirstInstance
+	s4EventstoreIndexes  *EventstoreIndexesNew
+	s5LastFailed         *LastFailed
+	s6OwnerRemoveColumns *OwnerRemoveColumns
+	s7LogstoreTables     *LogstoreTables
+	s8AuthTokens         *AuthTokenIndexes
+	s9EventstoreIndexes2 *EventstoreIndexesNew
+	CorrectCreationDate  *CorrectCreationDate
 }
 
 type encryptionKeyConfig struct {
diff --git a/cmd/setup/setup.go b/cmd/setup/setup.go
index a7deef0071..e90ad85f94 100644
--- a/cmd/setup/setup.go
+++ b/cmd/setup/setup.go
@@ -88,7 +88,7 @@ func Setup(config *Config, steps *Steps, masterKey string) {
 	steps.s7LogstoreTables = &LogstoreTables{dbClient: dbClient.DB, username: config.Database.Username(), dbType: config.Database.Type()}
 	steps.s8AuthTokens = &AuthTokenIndexes{dbClient: dbClient}
 	steps.s9EventstoreIndexes2 = New09(dbClient)
-	steps.s10EventstoreCreationDate = &CorrectCreationDate{dbClient: dbClient}
+	steps.CorrectCreationDate.dbClient = dbClient
 
 	err = projection.Create(ctx, dbClient, eventstoreClient, config.Projections, nil, nil)
 	logging.OnError(err).Fatal("unable to start projections")
@@ -124,7 +124,7 @@ func Setup(config *Config, steps *Steps, masterKey string) {
 	logging.OnError(err).Fatal("unable to migrate step 8")
 	err = migration.Migrate(ctx, eventstoreClient, steps.s9EventstoreIndexes2)
 	logging.OnError(err).Fatal("unable to migrate step 9")
-	err = migration.Migrate(ctx, eventstoreClient, steps.s10EventstoreCreationDate)
+	err = migration.Migrate(ctx, eventstoreClient, steps.CorrectCreationDate)
 	logging.OnError(err).Fatal("unable to migrate step 10")
 
 	for _, repeatableStep := range repeatableSteps {
diff --git a/cmd/setup/steps.yaml b/cmd/setup/steps.yaml
index c362f80e84..f5cb4b0c2a 100644
--- a/cmd/setup/steps.yaml
+++ b/cmd/setup/steps.yaml
@@ -30,3 +30,5 @@ FirstInstance:
       MachineKey:
         ExpirationDate:
         Type:
+CorrectCreationDate:
+  FailAfter: 5m