snapshot of https://github.com/ory/kratos.git (master, 2026-09-15; blobs > 131072B omitted)

persistence/sql/persister_test.go

@@ -0,0 +1,513 @@
+// Copyright © 2023 Ory Corp
+// SPDX-License-Identifier: Apache-2.0
+
+package sql_test
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"regexp"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/cockroachdb/cockroach-go/v2/testserver"
+	"github.com/pkg/errors"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/sync/errgroup"
+
+	continuity "github.com/ory/kratos/continuity/test"
+	"github.com/ory/kratos/corpx"
+	courier "github.com/ory/kratos/courier/test"
+	"github.com/ory/kratos/driver/config"
+	ri "github.com/ory/kratos/identity"
+	identity "github.com/ory/kratos/identity/test"
+	"github.com/ory/kratos/persistence/sql"
+	"github.com/ory/kratos/persistence/sql/batch"
+	sqltesthelpers "github.com/ory/kratos/persistence/sql/testhelpers"
+	"github.com/ory/kratos/pkg"
+	"github.com/ory/kratos/pkg/testhelpers"
+	"github.com/ory/kratos/schema"
+	errorx "github.com/ory/kratos/selfservice/errorx/test"
+	lf "github.com/ory/kratos/selfservice/flow/login"
+	login "github.com/ory/kratos/selfservice/flow/login/test"
+	recovery "github.com/ory/kratos/selfservice/flow/recovery/test"
+	registration "github.com/ory/kratos/selfservice/flow/registration/test"
+	settings "github.com/ory/kratos/selfservice/flow/settings/test"
+	verification "github.com/ory/kratos/selfservice/flow/verification/test"
+	sessiontokenexchange "github.com/ory/kratos/selfservice/sessiontokenexchange/test"
+	code "github.com/ory/kratos/selfservice/strategy/code/test"
+	link "github.com/ory/kratos/selfservice/strategy/link/test"
+	session "github.com/ory/kratos/session/test"
+	"github.com/ory/kratos/x"
+	"github.com/ory/pop/v6"
+	"github.com/ory/pop/v6/logging"
+	"github.com/ory/x/configx"
+	"github.com/ory/x/dbal"
+	"github.com/ory/x/popx"
+	"github.com/ory/x/sqlcon"
+	"github.com/ory/x/sqlcon/dockertest"
+	"github.com/ory/x/sqlxx"
+	"github.com/ory/x/urlx"
+)
+
+func init() {
+	corpx.RegisterFakes()
+	pop.SetNowFunc(func() time.Time {
+		return time.Now().UTC().Round(time.Second)
+	})
+}
+
+func TestMain(m *testing.M) {
+	os.Exit(m.Run())
+}
+
+func pl(t testing.TB) func(lvl logging.Level, s string, args ...interface{}) {
+	return func(lvl logging.Level, s string, args ...interface{}) {
+		if pop.Debug == false {
+			return
+		}
+
+		if lvl == logging.SQL {
+			if len(args) > 0 {
+				xargs := make([]string, len(args))
+				for i, a := range args {
+					switch a.(type) {
+					case string:
+						xargs[i] = fmt.Sprintf("%q", a)
+					default:
+						xargs[i] = fmt.Sprintf("%v", a)
+					}
+				}
+				s = fmt.Sprintf("%s - %s | %s", lvl, s, xargs)
+			} else {
+				s = fmt.Sprintf("%s - %s", lvl, s)
+			}
+		} else {
+			s = fmt.Sprintf(s, args...)
+			s = fmt.Sprintf("%s - %s", lvl, s)
+		}
+		t.Log(s)
+	}
+}
+
+var dbNames = []string{"sqlite", "postgres", "mysql", "cockroach"}
+
+func setupDatabase(t testing.TB, name string) string {
+	var dsn string
+	switch name {
+	case "sqlite":
+		dsn = dbal.NewSQLiteTestDatabase(t)
+	case "postgres":
+		dsn = dockertest.RunTestPostgreSQLWithVersion(t, "16")
+	case "mysql":
+		dsn = dockertest.RunTestMySQLWithVersion(t, "8.4")
+	case "cockroach":
+		dsn = newLocalTestCRDBServer(t)
+	default:
+		t.Fatalf("unknown database: %s", name)
+	}
+
+	if name != "sqlite" {
+		require.EventuallyWithT(t, func(t *assert.CollectT) {
+			c, err := pop.NewConnection(&pop.ConnectionDetails{URL: dsn})
+			require.NoError(t, err)
+			require.NoError(t, c.Open())
+			dbName := "testdb" + strings.ReplaceAll(x.NewUUID().String(), "-", "")
+			require.NoError(t, c.RawQuery("CREATE DATABASE "+dbName).Exec())
+			dsn = regexp.MustCompile(`/[a-z0-9]+\?`).ReplaceAllString(dsn, "/"+dbName+"?")
+		}, 60*time.Second, 200*time.Millisecond)
+	}
+
+	t.Logf("Connecting to %s: %s", name, dsn)
+
+	_, reg := pkg.NewRegistryDefaultWithDSN(t, dsn)
+	p := reg.Persister().(*sql.Persister)
+
+	t.Logf("Applying %s migrations", name)
+	pop.SetLogger(pl(t))
+	require.NoError(t, p.MigrateUp(t.Context()))
+	t.Logf("%s migrations applied", name)
+	status, err := p.MigrationStatus(t.Context())
+	require.NoError(t, err)
+	require.False(t, status.HasPending())
+
+	t.Logf("Database %s initialized successfully", name)
+	return dsn
+}
+
+func TestPersister(t *testing.T) {
+	t.Parallel()
+
+	ctx := t.Context()
+	// Bake the default identity schema into each registry's base config below
+	// instead of layering it onto ctx. A context overlay makes
+	// TestConfigProvider.Config rebuild the whole koanf tree on every config read
+	// (the same reason identity.TestPool bakes PoolConfigValues), which the
+	// contracts here would otherwise pay on every read.
+	defaultSchema := configx.WithValues(testhelpers.DefaultIdentitySchemaConfig("file://./stub/identity.schema.json"))
+
+	for _, name := range dbNames {
+		t.Run(fmt.Sprintf("database=%s", name), func(t *testing.T) {
+			t.Parallel()
+
+			if name != "sqlite" && testing.Short() {
+				t.Skip("skipping non-sqlite under -short")
+			}
+
+			dsn := setupDatabase(t, name)
+			t.Logf("DSN: %s", dsn)
+
+			t.Run("racy identity creation", func(t *testing.T) {
+				t.Parallel()
+				dsn := dsn
+				// Just have a separate DB for sqlite to speed it up.
+				if name == "sqlite" {
+					dsn = dbal.NewSQLiteTestDatabase(t)
+					// Pre-run migrations so concurrent goroutine startups below
+					// don't race to apply the same migrations on a fresh DB.
+					pkg.NewRegistryDefaultWithDSN(t, dsn)
+				}
+
+				var wg sync.WaitGroup
+
+				for i := range 10 {
+					wg.Add(1)
+					go func() {
+						defer wg.Done()
+
+						_, reg := pkg.NewRegistryDefaultWithDSN(t, dsn, defaultSchema)
+						_, ps := testhelpers.NewNetwork(t, ctx, reg.Persister())
+						id := ri.NewIdentity("")
+						id.SetCredentials(ri.CredentialsTypePassword, ri.Credentials{
+							Type:        ri.CredentialsTypePassword,
+							Identifiers: []string{fmt.Sprintf("racy identity %d", i)},
+							Config:      sqlxx.JSONRawMessage(`{"foo":"bar"}`),
+						})
+						id.Traits = ri.Traits("{}")
+
+						require.NoError(t, ps.CreateIdentity(ctx, id))
+					}()
+				}
+
+				wg.Wait()
+			})
+
+			t.Run("case=credential types exist", func(t *testing.T) {
+				t.Parallel()
+				dsn := dsn
+				// Just have a separate DB for sqlite to speed it up.
+				if name == "sqlite" {
+					dsn = dbal.NewSQLiteTestDatabase(t)
+				}
+
+				_, reg := pkg.NewRegistryDefaultWithDSN(t, dsn, defaultSchema)
+				_, p := testhelpers.NewNetwork(t, ctx, reg.Persister())
+				for _, ct := range []ri.CredentialsType{ri.CredentialsTypeOIDC, ri.CredentialsTypePassword} {
+					require.NoError(t, p.(*sql.Persister).Connection(t.Context()).Where("name = ?", ct).First(&ri.CredentialsTypeTable{}))
+				}
+			})
+
+			t.Run("contract=identity.TestPool", func(t *testing.T) {
+				t.Parallel()
+				dsn := dsn
+				// Just have a separate DB for sqlite to speed it up.
+				if name == "sqlite" {
+					dsn = dbal.NewSQLiteTestDatabase(t)
+				}
+
+				// Unlike the sibling contracts (which bake only the single default
+				// identity schema via defaultSchema), this contract needs the full
+				// PoolConfigValues schema set baked into the registry's base config.
+				ctx := t.Context()
+				_, reg := pkg.NewRegistryDefaultWithDSN(t, dsn, configx.WithValues(identity.PoolConfigValues()))
+				_, p := testhelpers.NewNetwork(t, ctx, reg.Persister())
+				identity.TestPool(ctx, p, reg.IdentityManager(), name)(t)
+			})
+			t.Run("contract=registration.TestFlowPersister", func(t *testing.T) {
+				t.Parallel()
+				dsn := dsn
+				// Just have a separate DB for sqlite to speed it up.
+				if name == "sqlite" {
+					dsn = dbal.NewSQLiteTestDatabase(t)
+				}
+
+				_, reg := pkg.NewRegistryDefaultWithDSN(t, dsn, defaultSchema)
+				_, p := testhelpers.NewNetwork(t, ctx, reg.Persister())
+				registration.TestFlowPersister(ctx, p)(t)
+			})
+			t.Run("contract=errorx.TestPersister", func(t *testing.T) {
+				t.Parallel()
+				dsn := dsn
+				// Just have a separate DB for sqlite to speed it up.
+				if name == "sqlite" {
+					dsn = dbal.NewSQLiteTestDatabase(t)
+				}
+
+				_, reg := pkg.NewRegistryDefaultWithDSN(t, dsn, defaultSchema)
+				_, p := testhelpers.NewNetwork(t, ctx, reg.Persister())
+				errorx.TestPersister(ctx, p)(t)
+			})
+			t.Run("contract=login.TestFlowPersister", func(t *testing.T) {
+				t.Parallel()
+				dsn := dsn
+				// Just have a separate DB for sqlite to speed it up.
+				if name == "sqlite" {
+					dsn = dbal.NewSQLiteTestDatabase(t)
+				}
+
+				_, reg := pkg.NewRegistryDefaultWithDSN(t, dsn, defaultSchema)
+				_, p := testhelpers.NewNetwork(t, ctx, reg.Persister())
+				login.TestFlowPersister(ctx, p)(t)
+			})
+			t.Run("contract=settings.TestFlowPersister", func(t *testing.T) {
+				t.Parallel()
+				dsn := dsn
+				// Just have a separate DB for sqlite to speed it up.
+				if name == "sqlite" {
+					dsn = dbal.NewSQLiteTestDatabase(t)
+				}
+
+				_, reg := pkg.NewRegistryDefaultWithDSN(t, dsn, defaultSchema)
+				_, p := testhelpers.NewNetwork(t, ctx, reg.Persister())
+				settings.TestFlowPersister(ctx, p)(t)
+			})
+			t.Run("contract=session.TestPersister", func(t *testing.T) {
+				t.Parallel()
+				dsn := dsn
+				// Just have a separate DB for sqlite to speed it up.
+				if name == "sqlite" {
+					dsn = dbal.NewSQLiteTestDatabase(t)
+				}
+
+				_, reg := pkg.NewRegistryDefaultWithDSN(t, dsn, defaultSchema)
+				_, p := testhelpers.NewNetwork(t, ctx, reg.Persister())
+				session.TestPersister(ctx, reg.Config(), p)(t)
+			})
+			t.Run("contract=sessiontokenexchange.TestPersister", func(t *testing.T) {
+				t.Parallel()
+				dsn := dsn
+				// Just have a separate DB for sqlite to speed it up.
+				if name == "sqlite" {
+					dsn = dbal.NewSQLiteTestDatabase(t)
+				}
+
+				_, reg := pkg.NewRegistryDefaultWithDSN(t, dsn, defaultSchema)
+				_, p := testhelpers.NewNetwork(t, ctx, reg.Persister())
+				sessiontokenexchange.TestPersister(ctx, p)(t)
+			})
+			t.Run("contract=courier.TestPersister", func(t *testing.T) {
+				t.Parallel()
+				dsn := dsn
+				// Just have a separate DB for sqlite to speed it up.
+				if name == "sqlite" {
+					dsn = dbal.NewSQLiteTestDatabase(t)
+				}
+
+				_, reg := pkg.NewRegistryDefaultWithDSN(t, dsn, defaultSchema)
+				_, p := testhelpers.NewNetwork(t, ctx, reg.Persister())
+				upsert, insert := sqltesthelpers.DefaultNetworkWrapper(p)
+				courier.TestPersister(ctx, upsert, insert)(t)
+			})
+			t.Run("contract=verification.TestFlowPersister", func(t *testing.T) {
+				t.Parallel()
+				dsn := dsn
+				// Just have a separate DB for sqlite to speed it up.
+				if name == "sqlite" {
+					dsn = dbal.NewSQLiteTestDatabase(t)
+				}
+
+				_, reg := pkg.NewRegistryDefaultWithDSN(t, dsn, defaultSchema)
+				_, p := testhelpers.NewNetwork(t, ctx, reg.Persister())
+				verification.TestFlowPersister(ctx, p)(t)
+			})
+			t.Run("contract=recovery.TestFlowPersister", func(t *testing.T) {
+				t.Parallel()
+				dsn := dsn
+				// Just have a separate DB for sqlite to speed it up.
+				if name == "sqlite" {
+					dsn = dbal.NewSQLiteTestDatabase(t)
+				}
+
+				_, reg := pkg.NewRegistryDefaultWithDSN(t, dsn, defaultSchema)
+				_, p := testhelpers.NewNetwork(t, ctx, reg.Persister())
+				recovery.TestFlowPersister(ctx, p)(t)
+			})
+			t.Run("contract=link.TestPersister", func(t *testing.T) {
+				t.Parallel()
+				dsn := dsn
+				// Just have a separate DB for sqlite to speed it up.
+				if name == "sqlite" {
+					dsn = dbal.NewSQLiteTestDatabase(t)
+				}
+
+				_, reg := pkg.NewRegistryDefaultWithDSN(t, dsn, defaultSchema)
+				_, p := testhelpers.NewNetwork(t, ctx, reg.Persister())
+				link.TestPersister(ctx, p)(t)
+			})
+			t.Run("contract=code.TestPersister", func(t *testing.T) {
+				t.Parallel()
+				dsn := dsn
+				// Just have a separate DB for sqlite to speed it up.
+				if name == "sqlite" {
+					dsn = dbal.NewSQLiteTestDatabase(t)
+				}
+
+				_, reg := pkg.NewRegistryDefaultWithDSN(t, dsn, defaultSchema)
+				_, p := testhelpers.NewNetwork(t, ctx, reg.Persister())
+				code.TestPersister(ctx, p)(t)
+			})
+			t.Run("contract=continuity.TestPersister", func(t *testing.T) {
+				t.Parallel()
+				dsn := dsn
+				// Just have a separate DB for sqlite to speed it up.
+				if name == "sqlite" {
+					dsn = dbal.NewSQLiteTestDatabase(t)
+				}
+
+				_, reg := pkg.NewRegistryDefaultWithDSN(t, dsn, defaultSchema)
+				_, p := testhelpers.NewNetwork(t, ctx, reg.Persister())
+				continuity.TestPersister(ctx, p)(t)
+			})
+			t.Run("contract=batch.TestPersister", func(t *testing.T) {
+				t.Parallel()
+				dsn := dsn
+				// Just have a separate DB for sqlite to speed it up.
+				if name == "sqlite" {
+					dsn = dbal.NewSQLiteTestDatabase(t)
+				}
+
+				_, reg := pkg.NewRegistryDefaultWithDSN(t, dsn, defaultSchema)
+				_, p := testhelpers.NewNetwork(t, ctx, reg.Persister())
+				batch.TestPersister(ctx, reg.Tracer(ctx), p)(t)
+			})
+		})
+	}
+}
+
+func getErr(args ...interface{}) error {
+	if len(args) == 0 {
+		return nil
+	}
+	lastArg := args[len(args)-1]
+	if e, ok := lastArg.(error); ok {
+		return e
+	}
+	return nil
+}
+
+func TestPersister_Transaction(t *testing.T) {
+	t.Parallel()
+
+	_, reg := pkg.NewFastRegistryWithMocks(t)
+	p := reg.Persister()
+
+	t.Run("case=should not create identity because callback returned error", func(t *testing.T) {
+		i := &ri.Identity{
+			ID:     x.NewUUID(),
+			State:  ri.StateActive,
+			Traits: ri.Traits(`{}`),
+		}
+		errMessage := "failing because why not"
+		err := p.Transaction(t.Context(), func(_ context.Context, connection *pop.Connection) error {
+			require.NoError(t, connection.Create(i))
+			return errors.New(errMessage)
+		})
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), errMessage)
+		_, err = p.GetIdentity(t.Context(), i.ID, ri.ExpandNothing)
+		require.Error(t, err)
+		assert.Equal(t, sqlcon.ErrNoRows().Error(), err.Error())
+	})
+
+	t.Run("case=functions should use the context connection", func(t *testing.T) {
+		c := p.GetConnection(t.Context())
+		errMessage := "some stupid error you can't debug"
+		lr := &lf.Flow{
+			ID: x.NewUUID(),
+		}
+		err := c.Transaction(func(tx *pop.Connection) error {
+			ctx := popx.WithTransaction(t.Context(), tx)
+			require.NoError(t, p.CreateLoginFlow(ctx, lr), "%+v", lr)
+			require.NoError(t, getErr(p.GetLoginFlow(ctx, lr.ID)), "%+v", lr)
+			return errors.New(errMessage)
+		})
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), errMessage)
+		_, err = p.GetLoginFlow(t.Context(), lr.ID)
+		require.Error(t, err)
+		assert.Equal(t, sqlcon.ErrNoRows().Error(), err.Error())
+	})
+}
+
+func Benchmark_BatchCreateIdentities(b *testing.B) {
+	ctx := b.Context()
+	batchSizes := []int{1, 10, 100, 500, 800, 900, 1000, 2000, 3000}
+	parallelRequests := []int{1, 4, 8, 16}
+
+	for _, name := range dbNames {
+		b.Run(fmt.Sprintf("database=%s", name), func(b *testing.B) {
+			dsn := setupDatabase(b, name)
+			_, reg := pkg.NewRegistryDefaultWithDSN(b, dsn)
+			conf := reg.Config()
+			_, p := testhelpers.NewNetwork(b, ctx, reg.Persister())
+			multipleEmailsSchema := schema.Schema{
+				ID:     "multiple_emails",
+				URL:    urlx.ParseOrPanic("file://./stub/handler/multiple_emails.schema.json"),
+				RawURL: "file://./stub/identity-2.schema.json",
+			}
+			conf.MustSet(ctx, config.ViperKeyIdentitySchemas, []config.Schema{
+				{
+					ID:  multipleEmailsSchema.ID,
+					URL: multipleEmailsSchema.RawURL,
+				},
+			})
+			conf.MustSet(ctx, config.ViperKeyPublicBaseURL, "http://localhost/")
+
+			run := 0
+			for _, batchSize := range batchSizes {
+				b.Run(fmt.Sprintf("batch-size=%d", batchSize), func(b *testing.B) {
+					for _, paralellism := range parallelRequests {
+						b.Run(fmt.Sprintf("parallelism=%d", paralellism), func(b *testing.B) {
+							start := time.Now()
+							for i := 0; i < b.N; i++ {
+								wg := new(errgroup.Group)
+								for paralell := 0; paralell < paralellism; paralell++ {
+									paralell := paralell
+									wg.Go(func() error {
+										identities := make([]*ri.Identity, batchSize)
+										prefix := fmt.Sprintf("bench-insert-run-%d", run+paralell)
+										for j := range identities {
+											identities[j] = identity.NewTestIdentity(1, prefix, j)
+										}
+
+										return p.CreateIdentities(ctx, identities)
+									})
+								}
+								assert.NoError(b, wg.Wait())
+								run += paralellism
+							}
+							end := time.Now()
+							b.ReportMetric(float64(paralellism*batchSize*b.N), "identites_created")
+							b.ReportMetric(float64(paralellism*batchSize*b.N)/end.Sub(start).Seconds(), "identities/s")
+						})
+					}
+				})
+			}
+		})
+	}
+}
+
+func newLocalTestCRDBServer(t testing.TB) string {
+	ts, err := testserver.NewTestServer(testserver.CustomVersionOpt("v25.4.1"))
+	require.NoError(t, err)
+	t.Cleanup(ts.Stop)
+
+	require.NoError(t, ts.WaitForInit())
+
+	ts.PGURL().Scheme = "cockroach"
+	return ts.PGURL().String()
+}