refactor: Support context-based SQL transactions (#254)

Co-authored-by: hackerman <[email protected]>

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)
 }