refactor(persistence/sql): move connection to context to enable transactions

cmd/serve.go

@@ -21,6 +21,7 @@ import (
 
 	"github.com/fsnotify/fsnotify"
 	"github.com/gobuffalo/packr/v2"
+
 	"github.com/ory/viper"
 	"github.com/ory/x/flagx"
 	"github.com/ory/x/viperx"

courier/courier_test.go

@@ -95,8 +95,8 @@ func TestSMTP(t *testing.T) {
 	}
 
 	smtp, api := runTestSMTP(t)
-	t.Logf("SMTP URL: %s",smtp)
-	t.Logf("API URL: %s",api)
+	t.Logf("SMTP URL: %s", smtp)
+	t.Logf("API URL: %s", api)
 
 	conf, reg := internal.NewRegistryDefault(t)
 	viper.Set(configuration.ViperKeyCourierSMTPURL, smtp)

persistence/reference.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"io"
 
+	"github.com/gobuffalo/pop/v5"
+
 	"github.com/ory/kratos/courier"
 	"github.com/ory/kratos/identity"
 	"github.com/ory/kratos/selfservice/errorx"
@@ -33,4 +35,6 @@ type Persister interface {
 	MigrationStatus(c context.Context, b io.Writer) error
 	MigrateDown(c context.Context, steps int) error
 	MigrateUp(c context.Context) error
+	GetConnection(ctx context.Context) *pop.Connection
+	Transaction(ctx context.Context, callback func(connection *pop.Connection) error) error
 }

persistence/sql/persister.go

@@ -41,26 +41,26 @@ func NewPersister(r persisterDependencies, conf configuration.Provider, c *pop.C
 	return &Persister{c: c, mb: m, cf: conf, r: r}, nil
 }
 
-func (p *Persister) MigrationStatus(c context.Context, w io.Writer) error {
+func (p *Persister) MigrationStatus(ctx context.Context, w io.Writer) error {
 	return errors.WithStack(p.mb.Status(w))
 }
 
-func (p *Persister) MigrateDown(c context.Context, steps int) error {
+func (p *Persister) MigrateDown(ctx context.Context, steps int) error {
 	return errors.WithStack(p.mb.Down(steps))
 }
 
-func (p *Persister) MigrateUp(c context.Context) error {
+func (p *Persister) MigrateUp(ctx context.Context) error {
 	return errors.WithStack(p.mb.Up())
 }
 
-func (p *Persister) Close(c context.Context) error {
-	return errors.WithStack(p.c.Close())
+func (p *Persister) Close(ctx context.Context) error {
+	return errors.WithStack(p.GetConnection(ctx).Close())
 }
 
-func (p *Persister) Ping(c context.Context) error {
+func (p *Persister) Ping(ctx context.Context) error {
 	type pinger interface {
 		Ping() error
 	}
 
-	return errors.WithStack(p.c.Store.(pinger).Ping())
+	return errors.WithStack(p.GetConnection(ctx).Store.(pinger).Ping())
 }

persistence/sql/persister_courier.go

@@ -16,12 +16,12 @@ var _ courier.Persister = new(Persister)
 
 func (p *Persister) AddMessage(ctx context.Context, m *courier.Message) error {
 	m.Status = courier.MessageStatusQueued
-	return sqlcon.HandleError(p.c.Create(m)) // do not create eager to avoid identity injection.
+	return sqlcon.HandleError(p.GetConnection(ctx).Create(m)) // do not create eager to avoid identity injection.
 }
 
 func (p *Persister) NextMessages(ctx context.Context, limit uint8) ([]courier.Message, error) {
 	var m []courier.Message
-	if err := p.c.
+	if err := p.GetConnection(ctx).
 		Eager().
 		Where("status != ?", courier.MessageStatusSent).
 		Order("created_at ASC").Limit(int(limit)).All(&m); err != nil {
@@ -40,7 +40,7 @@ func (p *Persister) NextMessages(ctx context.Context, limit uint8) ([]courier.Me
 
 func (p *Persister) LatestQueuedMessage(ctx context.Context) (*courier.Message, error) {
 	var m courier.Message
-	if err := p.c.
+	if err := p.GetConnection(ctx).
 		Eager().
 		Where("status != ?", courier.MessageStatusSent).
 		Order("created_at DESC").First(&m); err != nil {
@@ -54,7 +54,7 @@ func (p *Persister) LatestQueuedMessage(ctx context.Context) (*courier.Message,
 }
 
 func (p *Persister) SetMessageStatus(ctx context.Context, id uuid.UUID, ms courier.MessageStatus) error {
-	count, err := p.c.RawQuery("UPDATE courier_messages SET status = ? WHERE id = ?", ms, id).ExecWithCount()
+	count, err := p.GetConnection(ctx).RawQuery("UPDATE courier_messages SET status = ? WHERE id = ?", ms, id).ExecWithCount()
 	if err != nil {
 		return sqlcon.HandleError(err)
 	}

persistence/sql/persister_errorx.go

@@ -33,7 +33,7 @@ func (p *Persister) Add(ctx context.Context, csrfToken string, errs ...error) (u
 		WasSeen:   false,
 	}
 
-	if err := p.c.Create(c); err != nil {
+	if err := p.GetConnection(ctx).Create(c); err != nil {
 		return uuid.Nil, sqlcon.HandleError(err)
 	}
 
@@ -42,11 +42,11 @@ func (p *Persister) Add(ctx context.Context, csrfToken string, errs ...error) (u
 
 func (p *Persister) Read(ctx context.Context, id uuid.UUID) (*errorx.ErrorContainer, error) {
 	var ec errorx.ErrorContainer
-	if err := p.c.Find(&ec, id); err != nil {
+	if err := p.GetConnection(ctx).Find(&ec, id); err != nil {
 		return nil, sqlcon.HandleError(err)
 	}
 
-	if err := p.c.RawQuery("UPDATE selfservice_errors SET was_seen = true, seen_at = ? WHERE id = ?", time.Now().UTC(), id).Exec(); err != nil {
+	if err := p.GetConnection(ctx).RawQuery("UPDATE selfservice_errors SET was_seen = true, seen_at = ? WHERE id = ?", time.Now().UTC(), id).Exec(); err != nil {
 		return nil, sqlcon.HandleError(err)
 	}
 
@@ -55,9 +55,9 @@ func (p *Persister) Read(ctx context.Context, id uuid.UUID) (*errorx.ErrorContai
 
 func (p *Persister) Clear(ctx context.Context, olderThan time.Duration, force bool) (err error) {
 	if force {
-		err = p.c.RawQuery("DELETE FROM selfservice_errors WHERE seen_at < ?", olderThan).Exec()
+		err = p.GetConnection(ctx).RawQuery("DELETE FROM selfservice_errors WHERE seen_at < ?", olderThan).Exec()
 	} else {
-		err = p.c.RawQuery("DELETE FROM selfservice_errors WHERE was_seen=true AND seen_at < ?", time.Now().UTC().Add(-olderThan)).Exec()
+		err = p.GetConnection(ctx).RawQuery("DELETE FROM selfservice_errors WHERE was_seen=true AND seen_at < ?", time.Now().UTC().Add(-olderThan)).Exec()
 	}
 
 	return sqlcon.HandleError(err)

persistence/sql/persister_identity.go

@@ -28,15 +28,15 @@ var _ identity.PrivilegedPool = new(Persister)
 
 func (p *Persister) FindByCredentialsIdentifier(ctx context.Context, ct identity.CredentialsType, match string) (*identity.Identity, *identity.Credentials, error) {
 	var cts []identity.CredentialsTypeTable
-	if err := p.c.All(&cts); err != nil {
+	if err := p.GetConnection(ctx).All(&cts); err != nil {
 		return nil, nil, sqlcon.HandleError(err)
 	}
 
 	var find struct {
 		IdentityID uuid.UUID `db:"identity_id"`
 	}
 
-	if err := p.c.RawQuery(`SELECT
+	if err := p.GetConnection(ctx).RawQuery(`SELECT
     ic.identity_id
 FROM identity_credentials ic
          INNER JOIN identity_credential_types ict on ic.identity_credential_type_id = ict.id
@@ -148,7 +148,7 @@ func (p *Persister) CreateIdentity(ctx context.Context, i *identity.Identity) er
 		return err
 	}
 
-	return sqlcon.HandleError(p.c.Transaction(func(tx *pop.Connection) error {
+	return sqlcon.HandleError(p.GetConnection(ctx).Transaction(func(tx *pop.Connection) error {
 		if err := tx.Create(i); err != nil {
 			return err
 		}
@@ -165,7 +165,7 @@ func (p *Persister) ListIdentities(ctx context.Context, limit, offset int) ([]id
 	is := make([]identity.Identity, 0)
 
 	/* #nosec G201 TableName is static */
-	if err := sqlcon.HandleError(p.c.RawQuery(fmt.Sprintf("SELECT * FROM %s LIMIT ? OFFSET ?", new(identity.Identity).TableName()), limit, offset).All(&is)); err != nil {
+	if err := sqlcon.HandleError(p.GetConnection(ctx).RawQuery(fmt.Sprintf("SELECT * FROM %s LIMIT ? OFFSET ?", new(identity.Identity).TableName()), limit, offset).All(&is)); err != nil {
 		return nil, err
 	}
 
@@ -183,7 +183,7 @@ func (p *Persister) UpdateIdentity(ctx context.Context, i *identity.Identity) er
 		return err
 	}
 
-	return sqlcon.HandleError(p.c.Transaction(func(tx *pop.Connection) error {
+	return sqlcon.HandleError(p.GetConnection(ctx).Transaction(func(tx *pop.Connection) error {
 		if count, err := tx.Where("id = ?", i.ID).Count(i); err != nil {
 			return err
 		} else if count == 0 {
@@ -214,7 +214,7 @@ func (p *Persister) UpdateIdentity(ctx context.Context, i *identity.Identity) er
 
 func (p *Persister) DeleteIdentity(ctx context.Context, id uuid.UUID) error {
 	/* #nosec G201 TableName is static */
-	count, err := p.c.RawQuery(fmt.Sprintf("DELETE FROM %s WHERE id = ?", new(identity.Identity).TableName()), id).ExecWithCount()
+	count, err := p.GetConnection(ctx).RawQuery(fmt.Sprintf("DELETE FROM %s WHERE id = ?", new(identity.Identity).TableName()), id).ExecWithCount()
 	if err != nil {
 		return sqlcon.HandleError(err)
 	}
@@ -226,7 +226,7 @@ func (p *Persister) DeleteIdentity(ctx context.Context, id uuid.UUID) error {
 
 func (p *Persister) GetIdentity(ctx context.Context, id uuid.UUID) (*identity.Identity, error) {
 	var i identity.Identity
-	if err := p.c.Find(&i, id); err != nil {
+	if err := p.GetConnection(ctx).Find(&i, id); err != nil {
 		return nil, sqlcon.HandleError(err)
 	}
 	i.Credentials = nil
@@ -239,19 +239,19 @@ func (p *Persister) GetIdentity(ctx context.Context, id uuid.UUID) (*identity.Id
 
 func (p *Persister) GetIdentityConfidential(ctx context.Context, id uuid.UUID) (*identity.Identity, error) {
 	var i identity.Identity
-	if err := p.c.Eager().Find(&i, id); err != nil {
+	if err := p.GetConnection(ctx).Eager().Find(&i, id); err != nil {
 		return nil, sqlcon.HandleError(err)
 	}
 
 	var cts []identity.CredentialsTypeTable
-	if err := p.c.All(&cts); err != nil {
+	if err := p.GetConnection(ctx).All(&cts); err != nil {
 		return nil, sqlcon.HandleError(err)
 	}
 
 	i.Credentials = map[identity.CredentialsType]identity.Credentials{}
 	for _, creds := range i.CredentialsCollection {
 		var cs identity.CredentialIdentifierCollection
-		if err := p.c.Where("identity_credential_id = ?", creds.ID).All(&cs); err != nil {
+		if err := p.GetConnection(ctx).Where("identity_credential_id = ?", creds.ID).All(&cs); err != nil {
 			return nil, sqlcon.HandleError(err)
 		}
 
@@ -277,7 +277,7 @@ func (p *Persister) GetIdentityConfidential(ctx context.Context, id uuid.UUID) (
 
 func (p *Persister) FindAddressByCode(ctx context.Context, code string) (*identity.VerifiableAddress, error) {
 	var address identity.VerifiableAddress
-	if err := p.c.Where("code = ?", code).First(&address); err != nil {
+	if err := p.GetConnection(ctx).Where("code = ?", code).First(&address); err != nil {
 		return nil, sqlcon.HandleError(err)
 	}
 
@@ -286,7 +286,7 @@ func (p *Persister) FindAddressByCode(ctx context.Context, code string) (*identi
 
 func (p *Persister) FindAddressByValue(ctx context.Context, via identity.VerifiableAddressType, value string) (*identity.VerifiableAddress, error) {
 	var address identity.VerifiableAddress
-	if err := p.c.Where("via = ? AND value = ?", via, value).First(&address); err != nil {
+	if err := p.GetConnection(ctx).Where("via = ? AND value = ?", via, value).First(&address); err != nil {
 		return nil, sqlcon.HandleError(err)
 	}
 
@@ -299,7 +299,7 @@ func (p *Persister) VerifyAddress(ctx context.Context, code string) error {
 		return err
 	}
 
-	return sqlcon.HandleError(p.c.RawQuery(
+	return sqlcon.HandleError(p.GetConnection(ctx).RawQuery(
 		/* #nosec G201 TableName is static */
 		fmt.Sprintf(
 			"UPDATE %s SET status = ?, verified = true, verified_at = ?, code = ? WHERE code = ?",
@@ -313,7 +313,7 @@ func (p *Persister) VerifyAddress(ctx context.Context, code string) error {
 }
 
 func (p *Persister) UpdateVerifiableAddress(ctx context.Context, address *identity.VerifiableAddress) error {
-	return sqlcon.HandleError(p.c.Update(address))
+	return sqlcon.HandleError(p.GetConnection(ctx).Update(address))
 }
 
 func (p *Persister) validateIdentity(i *identity.Identity) error {

persistence/sql/persister_login.go

@@ -3,6 +3,8 @@ package sql
 import (
 	"context"
 
+	"github.com/gobuffalo/pop/v5"
+
 	"github.com/gofrs/uuid"
 
 	"github.com/ory/x/sqlcon"
@@ -14,35 +16,39 @@ import (
 var _ login.RequestPersister = new(Persister)
 
 func (p *Persister) CreateLoginRequest(ctx context.Context, r *login.Request) error {
-	return p.c.Eager().Create(r)
+	return p.GetConnection(ctx).Eager().Create(r)
 }
 
 func (p *Persister) GetLoginRequest(ctx context.Context, id uuid.UUID) (*login.Request, error) {
+	conn := p.GetConnection(ctx)
 	var r login.Request
-	if err := p.c.Eager().Find(&r, id); err != nil {
+	if err := conn.Eager().Find(&r, id); err != nil {
 		return nil, sqlcon.HandleError(err)
 	}
 
-	if err := (&r).AfterFind(p.c); err != nil {
+	if err := (&r).AfterFind(conn); err != nil {
 		return nil, err
 	}
 
 	return &r, nil
 }
 
 func (p *Persister) UpdateLoginRequest(ctx context.Context, id uuid.UUID, ct identity.CredentialsType, rm *login.RequestMethod) error {
-	rr, err := p.GetLoginRequest(ctx, id)
-	if err != nil {
-		return err
-	}
-
-	method, ok := rr.Methods[ct]
-	if !ok {
-		rm.RequestID = rr.ID
-		rm.Method = ct
-		return p.c.Save(rm)
-	}
-
-	method.Config = rm.Config
-	return p.c.Save(method)
+	return p.Transaction(ctx, func(tx *pop.Connection) error {
+		ctx := WithTransaction(ctx, tx)
+		rr, err := p.GetLoginRequest(ctx, id)
+		if err != nil {
+			return err
+		}
+
+		method, ok := rr.Methods[ct]
+		if !ok {
+			rm.RequestID = rr.ID
+			rm.Method = ct
+			return tx.Save(rm)
+		}
+
+		method.Config = rm.Config
+		return tx.Save(method)
+	})
 }

persistence/sql/persister_profile.go

@@ -14,17 +14,17 @@ var _ profile.RequestPersister = new(Persister)
 
 func (p *Persister) CreateProfileRequest(ctx context.Context, r *profile.Request) error {
 	r.IdentityID = r.Identity.ID
-	return sqlcon.HandleError(p.c.Create(r)) // This must not be eager or identities will be created / updated
+	return sqlcon.HandleError(p.GetConnection(ctx).Create(r)) // This must not be eager or identities will be created / updated
 }
 
 func (p *Persister) GetProfileRequest(ctx context.Context, id uuid.UUID) (*profile.Request, error) {
 	var r profile.Request
-	if err := p.c.Eager().Find(&r, id); err != nil {
+	if err := p.GetConnection(ctx).Eager().Find(&r, id); err != nil {
 		return nil, sqlcon.HandleError(err)
 	}
 	return &r, nil
 }
 
 func (p *Persister) UpdateProfileRequest(ctx context.Context, r *profile.Request) error {
-	return sqlcon.HandleError(p.c.Update(r)) // This must not be eager or identities will be created / updated
+	return sqlcon.HandleError(p.GetConnection(ctx).Update(r)) // This must not be eager or identities will be created / updated
 }

persistence/sql/persister_registration.go

@@ -12,16 +12,16 @@ import (
 )
 
 func (p *Persister) CreateRegistrationRequest(ctx context.Context, r *registration.Request) error {
-	return p.c.Eager().Create(r)
+	return p.GetConnection(ctx).Eager().Create(r)
 }
 
 func (p *Persister) GetRegistrationRequest(ctx context.Context, id uuid.UUID) (*registration.Request, error) {
 	var r registration.Request
-	if err := p.c.Eager().Find(&r, id); err != nil {
+	if err := p.GetConnection(ctx).Eager().Find(&r, id); err != nil {
 		return nil, sqlcon.HandleError(err)
 	}
 
-	if err := (&r).AfterFind(p.c); err != nil {
+	if err := (&r).AfterFind(p.GetConnection(ctx)); err != nil {
 		return nil, err
 	}
 
@@ -38,9 +38,9 @@ func (p *Persister) UpdateRegistrationRequest(ctx context.Context, id uuid.UUID,
 	if !ok {
 		rm.RequestID = rr.ID
 		rm.Method = ct
-		return p.c.Save(rm)
+		return p.GetConnection(ctx).Save(rm)
 	}
 
 	method.Config = rm.Config
-	return p.c.Save(method)
+	return p.GetConnection(ctx).Save(method)
 }