From d138abb6278ebb232e120bee0fb956a0f2816b8d Mon Sep 17 00:00:00 2001
From: hackerman <3372410+aeneasr@users.noreply.github.com>
Date: Fri, 6 Oct 2023 12:13:44 +0200
Subject: [PATCH] fix: remove slow queries from update identities (#3553)

---
 go.mod                                        |  5 ++--
 go.sum                                        |  8 +++----
 .../sql/identity/persister_identity.go        | 11 +++------
 persistence/sql/update/update.go              | 24 ++++++-------------
 script/testenv.sh                             |  2 +-
 5 files changed, 17 insertions(+), 33 deletions(-)

diff --git a/go.mod b/go.mod
index dd4618fbb2b1..3b98b3d283ac 100644
--- a/go.mod
+++ b/go.mod
@@ -5,8 +5,7 @@ go 1.19
 replace (
 	github.com/bradleyjkemp/cupaloy/v2 => github.com/aeneasr/cupaloy/v2 v2.6.1-0.20210924214125-3dfdd01210a3
 
-	github.com/go-sql-driver/mysql => github.com/go-sql-driver/mysql v1.7.0
-
+	github.com/go-sql-driver/mysql => github.com/go-sql-driver/mysql v1.7.2-0.20231005084435-37980127edfb
 	github.com/gorilla/sessions => github.com/ory/sessions v1.2.2-0.20220110165800-b09c17334dc2
 
 	github.com/mattn/go-sqlite3 => github.com/mattn/go-sqlite3 v1.14.7-0.20210414154423-1157a4212dcb
@@ -78,7 +77,7 @@ require (
 	github.com/ory/jsonschema/v3 v3.0.8
 	github.com/ory/mail/v3 v3.0.0
 	github.com/ory/nosurf v1.2.7
-	github.com/ory/x v0.0.589
+	github.com/ory/x v0.0.590
 	github.com/phayes/freeport v0.0.0-20180830031419-95f893ade6f2
 	github.com/pkg/errors v0.9.1
 	github.com/pquerna/otp v1.4.0
diff --git a/go.sum b/go.sum
index fb96c9cdac9f..1493216aec55 100644
--- a/go.sum
+++ b/go.sum
@@ -282,8 +282,8 @@ github.com/go-playground/universal-translator v0.17.0 h1:icxd5fm+REJzpZx7ZfpaD87
 github.com/go-playground/universal-translator v0.17.0/go.mod h1:UkSxE5sNxxRwHyU+Scu5vgOQjsIJAF8j9muTVoKLVtA=
 github.com/go-playground/validator/v10 v10.4.1 h1:pH2c5ADXtd66mxoE0Zm9SUhxE20r7aM3F26W0hOn+GE=
 github.com/go-playground/validator/v10 v10.4.1/go.mod h1:nlOn6nFhuKACm19sB/8EGNn9GlaMV7XkbRSipzJ0Ii4=
-github.com/go-sql-driver/mysql v1.7.0 h1:ueSltNNllEqE3qcWBTD0iQd3IpL/6U+mJxLkazJ7YPc=
-github.com/go-sql-driver/mysql v1.7.0/go.mod h1:OXbVy3sEdcQ2Doequ6Z5BW6fXNQTmx+9S1MCJN5yJMI=
+github.com/go-sql-driver/mysql v1.7.2-0.20231005084435-37980127edfb h1:R5sscofA9Cyd0h2DNMbR8BHYVKj1ZTOgUah1PoU4OVU=
+github.com/go-sql-driver/mysql v1.7.2-0.20231005084435-37980127edfb/go.mod h1:6gYm/zDt3ahdnMVTPeT/LfoBFsws1qZm5yI6FmVjB14=
 github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY=
 github.com/go-stack/stack v1.8.1/go.mod h1:dcoOX6HbPZSZptuspn9bctJ+N/CnF5gGygcUP3XYfe4=
 github.com/go-swagger/go-swagger v0.30.3 h1:HuzvdMRed/9Q8vmzVcfNBQByZVtT79DNZxZ18OprdoI=
@@ -847,8 +847,8 @@ github.com/ory/nosurf v1.2.7 h1:YrHrbSensQyU6r6HT/V5+HPdVEgrOTMJiLoJABSBOp4=
 github.com/ory/nosurf v1.2.7/go.mod h1:d4L3ZBa7Amv55bqxCBtCs63wSlyaiCkWVl4vKf3OUxA=
 github.com/ory/sessions v1.2.2-0.20220110165800-b09c17334dc2 h1:zm6sDvHy/U9XrGpixwHiuAwpp0Ock6khSVHkrv6lQQU=
 github.com/ory/sessions v1.2.2-0.20220110165800-b09c17334dc2/go.mod h1:dk2InVEVJ0sfLlnXv9EAgkf6ecYs/i80K/zI+bUmuGM=
-github.com/ory/x v0.0.589 h1:ZNQ+nBzTCm3jI2ZZY/1kGWSE4jEtyvDYWu0ScfLgzac=
-github.com/ory/x v0.0.589/go.mod h1:ksLBEd6iW6czGpE6eNA0gCIxO1FFeqIxCZgsgwNrzMM=
+github.com/ory/x v0.0.590 h1:t0+XlSlDw5pzZhdAxOB8uFp1Dp+MStPRTG8Nn/fm1PE=
+github.com/ory/x v0.0.590/go.mod h1:ksLBEd6iW6czGpE6eNA0gCIxO1FFeqIxCZgsgwNrzMM=
 github.com/pascaldekloe/goe v0.0.0-20180627143212-57f6aae5913c/go.mod h1:lzWF7FIEvWOWxwDKqyGYQf6ZUaNfKdP144TG7ZOy1lc=
 github.com/pascaldekloe/goe v0.1.0 h1:cBOtyMzM9HTpWjXfbbunk26uA6nG3a8n06Wieeh0MwY=
 github.com/pascaldekloe/goe v0.1.0/go.mod h1:lzWF7FIEvWOWxwDKqyGYQf6ZUaNfKdP144TG7ZOy1lc=
diff --git a/persistence/sql/identity/persister_identity.go b/persistence/sql/identity/persister_identity.go
index 99469f365f1b..3e8335c8027a 100644
--- a/persistence/sql/identity/persister_identity.go
+++ b/persistence/sql/identity/persister_identity.go
@@ -311,7 +311,6 @@ func updateAssociation[T interface {
 	var inDB []T
 	if err := p.GetConnection(ctx).
 		Where("identity_id = ? AND nid = ?", i.ID, p.NetworkID(ctx)).
-		Order("id ASC").
 		All(&inDB); err != nil {
 		return sqlcon.HandleError(err)
 	}
@@ -791,16 +790,16 @@ func (p *IdentityPersister) UpdateIdentity(ctx context.Context, i *identity.Iden
 
 	i.NID = p.NetworkID(ctx)
 	return sqlcon.HandleError(p.Transaction(ctx, func(ctx context.Context, tx *pop.Connection) error {
-		if count, err := tx.Where("id = ? AND nid = ?", i.ID, p.NetworkID(ctx)).Count(i); err != nil {
+		// This returns "ErrNoRows" if the identity does not exist
+		if err := update.Generic(WithTransaction(ctx, tx), tx, p.r.Tracer(ctx).Tracer(), i); err != nil {
 			return err
-		} else if count == 0 {
-			return sql.ErrNoRows
 		}
 
 		p.normalizeAllAddressess(ctx, i)
 		if err := updateAssociation(ctx, p, i, i.RecoveryAddresses); err != nil {
 			return err
 		}
+
 		if err := updateAssociation(ctx, p, i, i.VerifiableAddresses); err != nil {
 			return err
 		}
@@ -814,10 +813,6 @@ func (p *IdentityPersister) UpdateIdentity(ctx context.Context, i *identity.Iden
 			return sqlcon.HandleError(err)
 		}
 
-		if err := update.Generic(WithTransaction(ctx, tx), tx, p.r.Tracer(ctx).Tracer(), i); err != nil {
-			return err
-		}
-
 		return sqlcon.HandleError(p.createIdentityCredentials(ctx, tx, i))
 	}))
 }
diff --git a/persistence/sql/update/update.go b/persistence/sql/update/update.go
index 2d7582060f8e..27bdf2c9cc93 100644
--- a/persistence/sql/update/update.go
+++ b/persistence/sql/update/update.go
@@ -42,22 +42,7 @@ func Generic(ctx context.Context, c *pop.Connection, tracer trace.Tracer, v Mode
 	}
 
 	//#nosec G201 -- TableName is static
-	stmt := fmt.Sprintf("SELECT COUNT(id) FROM %s AS %s WHERE %s.id = ? AND %s.nid = ?",
-		quoter.Quote(model.TableName()),
-		model.Alias(),
-		model.Alias(),
-		model.Alias(),
-	)
-
-	var count int
-	if err := c.Store.GetContext(ctx, &count, c.Dialect.TranslateSQL(stmt), v.GetID(), v.GetNID()); err != nil {
-		return sqlcon.HandleError(err)
-	} else if count == 0 {
-		return errors.WithStack(sqlcon.ErrNoRows)
-	}
-
-	//#nosec G201 -- TableName is static
-	stmt = fmt.Sprintf("UPDATE %s AS %s SET %s WHERE %s AND %s.nid = :nid",
+	stmt := fmt.Sprintf("UPDATE %s AS %s SET %s WHERE %s AND %s.nid = :nid",
 		quoter.Quote(model.TableName()),
 		model.Alias(),
 		cols.Writeable().QuotedUpdateString(quoter),
@@ -65,8 +50,13 @@ func Generic(ctx context.Context, c *pop.Connection, tracer trace.Tracer, v Mode
 		model.Alias(),
 	)
 
-	if _, err := c.Store.NamedExecContext(ctx, stmt, v); err != nil {
+	if result, err := c.Store.NamedExecContext(ctx, stmt, v); err != nil {
 		return sqlcon.HandleError(err)
+	} else if affected, err := result.RowsAffected(); err != nil {
+		return sqlcon.HandleError(err)
+	} else if affected == 0 {
+		return errors.WithStack(sqlcon.ErrNoRows)
 	}
+
 	return nil
 }
diff --git a/script/testenv.sh b/script/testenv.sh
index f7a5467a0f7f..671c5abc1b43 100755
--- a/script/testenv.sh
+++ b/script/testenv.sh
@@ -1,7 +1,7 @@
 #!/bin/bash
 
 docker rm -f kratos_test_database_mysql kratos_test_database_postgres kratos_test_database_cockroach kratos_test_hydra || true
-docker run --platform linux/amd64 --name kratos_test_database_mysql -p 3444:3306 -e MYSQL_ROOT_PASSWORD=secret -d mysql:8.0.26
+docker run --platform linux/amd64 --name kratos_test_database_mysql -p 3444:3306 -e MYSQL_ROOT_PASSWORD=secret -d mysql:8.0.34
 docker run --platform linux/amd64 --name kratos_test_database_postgres -p 3445:5432 -e POSTGRES_PASSWORD=secret -e POSTGRES_DB=postgres -d postgres:11.8 postgres -c log_statement=all
 docker run --platform linux/amd64 --name kratos_test_database_cockroach -p 3446:26257 -p 3447:8080 -d cockroachdb/cockroach:v22.2.6 start-single-node --insecure
 docker run --platform linux/amd64 --name kratos_test_hydra -p 4444:4444 -p 4445:4445 -d -e DSN=memory -e URLS_SELF_ISSUER=http://localhost:4444/ -e URLS_LOGIN=http://localhost:4446/login -e URLS_CONSENT=http://localhost:4446/consent oryd/hydra:v2.0.2 serve all --dev