feat: add Inspect option to driver

ory · Feb 7, 2023 · 8aa75e9 · 8aa75e9
1 parent e17f307
commit 8aa75e9
Show file tree

Hide file tree

Showing 3 changed files with 73 additions and 56 deletions.
diff --git a/driver/registry.go b/driver/registry.go
@@ -104,6 +104,7 @@ type Registry interface {
 	courier.PersistenceProvider
 
 	schema.HandlerProvider
+	schema.IdentityTraitsProvider
 
 	password2.ValidationProvider
 
@@ -177,6 +178,7 @@ type options struct {
 	skipNetworkInit bool
 	config          *config.Config
 	replaceTracer   func(*otelx.Tracer) *otelx.Tracer
+	inspect         func(Registry) error
 }
 
 type RegistryOption func(*options)
@@ -197,6 +199,12 @@ func ReplaceTracer(f func(*otelx.Tracer) *otelx.Tracer) func(o *options) {
 	}
 }
 
+func Inspect(f func(reg Registry) error) func(o *options) {
+	return func(o *options) {
+		o.inspect = f
+	}
+}
+
 func newOptions(os []RegistryOption) *options {
 	o := new(options)
 	for _, f := range os {

diff --git a/driver/registry_default.go b/driver/registry_default.go
@@ -586,68 +586,77 @@ func (m *RegistryDefault) Init(ctx context.Context, ctxer contextx.Contextualize
 	bc := backoff.NewExponentialBackOff()
 	bc.MaxElapsedTime = time.Minute * 5
 	bc.Reset()
-	return errors.WithStack(
-		backoff.Retry(func() error {
-			m.WithContextualizer(ctxer)
-
-			pool, idlePool, connMaxLifetime, connMaxIdleTime, cleanedDSN := sqlcon.ParseConnectionOptions(m.l, m.Config().DSN(ctx))
-			m.Logger().
-				WithField("pool", pool).
-				WithField("idlePool", idlePool).
-				WithField("connMaxLifetime", connMaxLifetime).
-				Debug("Connecting to SQL Database")
-			c, err := pop.NewConnection(&pop.ConnectionDetails{
-				URL:                       sqlcon.FinalizeDSN(m.l, cleanedDSN),
-				IdlePool:                  idlePool,
-				ConnMaxLifetime:           connMaxLifetime,
-				ConnMaxIdleTime:           connMaxIdleTime,
-				Pool:                      pool,
-				UseInstrumentedDriver:     m.Tracer(ctx).IsLoaded(),
-				InstrumentedDriverOptions: instrumentedDriverOpts,
-			})
-			if err != nil {
-				m.Logger().WithError(err).Warnf("Unable to connect to database, retrying.")
-				return errors.WithStack(err)
-			}
-			if err := c.Open(); err != nil {
-				m.Logger().WithError(err).Warnf("Unable to open database, retrying.")
-				return errors.WithStack(err)
-			}
-			p, err := sql.NewPersister(ctx, m, c)
-			if err != nil {
-				m.Logger().WithError(err).Warnf("Unable to initialize persister, retrying.")
-				return err
-			}
+	err := backoff.Retry(func() error {
+		m.WithContextualizer(ctxer)
+
+		pool, idlePool, connMaxLifetime, connMaxIdleTime, cleanedDSN := sqlcon.ParseConnectionOptions(m.l, m.Config().DSN(ctx))
+		m.Logger().
+			WithField("pool", pool).
+			WithField("idlePool", idlePool).
+			WithField("connMaxLifetime", connMaxLifetime).
+			Debug("Connecting to SQL Database")
+		c, err := pop.NewConnection(&pop.ConnectionDetails{
+			URL:                       sqlcon.FinalizeDSN(m.l, cleanedDSN),
+			IdlePool:                  idlePool,
+			ConnMaxLifetime:           connMaxLifetime,
+			ConnMaxIdleTime:           connMaxIdleTime,
+			Pool:                      pool,
+			UseInstrumentedDriver:     m.Tracer(ctx).IsLoaded(),
+			InstrumentedDriverOptions: instrumentedDriverOpts,
+		})
+		if err != nil {
+			m.Logger().WithError(err).Warnf("Unable to connect to database, retrying.")
+			return errors.WithStack(err)
+		}
+		if err := c.Open(); err != nil {
+			m.Logger().WithError(err).Warnf("Unable to open database, retrying.")
+			return errors.WithStack(err)
+		}
+		p, err := sql.NewPersister(ctx, m, c)
+		if err != nil {
+			m.Logger().WithError(err).Warnf("Unable to initialize persister, retrying.")
+			return err
+		}
 
-			if err := p.Ping(); err != nil {
-				m.Logger().WithError(err).Warnf("Unable to ping database, retrying.")
+		if err := p.Ping(); err != nil {
+			m.Logger().WithError(err).Warnf("Unable to ping database, retrying.")
+			return err
+		}
+
+		// if dsn is memory we have to run the migrations on every start
+		if dbal.IsMemorySQLite(m.Config().DSN(ctx)) || m.Config().DSN(ctx) == "memory" {
+			m.Logger().Infoln("Ory Kratos is running migrations on every startup as DSN is memory. This means your data is lost when Kratos terminates.")
+			if err := p.MigrateUp(ctx); err != nil {
+				m.Logger().WithError(err).Warnf("Unable to run migrations, retrying.")
 				return err
 			}
+		}
 
-			// if dsn is memory we have to run the migrations on every start
-			if dbal.IsMemorySQLite(m.Config().DSN(ctx)) || m.Config().DSN(ctx) == "memory" {
-				m.Logger().Infoln("Ory Kratos is running migrations on every startup as DSN is memory. This means your data is lost when Kratos terminates.")
-				if err := p.MigrateUp(ctx); err != nil {
-					m.Logger().WithError(err).Warnf("Unable to run migrations, retrying.")
-					return err
-				}
-			}
+		if o.skipNetworkInit {
+			m.persister = p
+			return nil
+		}
 
-			if o.skipNetworkInit {
-				m.persister = p
-				return nil
-			}
+		net, err := p.DetermineNetwork(ctx)
+		if err != nil {
+			m.Logger().WithError(err).Warnf("Unable to determine network, retrying.")
+			return err
+		}
 
-			net, err := p.DetermineNetwork(ctx)
-			if err != nil {
-				m.Logger().WithError(err).Warnf("Unable to determine network, retrying.")
-				return err
-			}
+		m.persister = p.WithNetworkID(net.ID)
+		return nil
+	}, bc)
 
-			m.persister = p.WithNetworkID(net.ID)
-			return nil
-		}, bc),
-	)
+	if err != nil {
+		return err
+	}
+
+	if o.inspect != nil {
+		if err := o.inspect(m); err != nil {
+			return errors.WithStack(err)
+		}
+	}
+	return nil
 }
 
 func (m *RegistryDefault) SetPersister(p persistence.Persister) {

diff --git a/persistence/sql/persister.go b/persistence/sql/persister.go
@@ -37,7 +37,7 @@ var migrations embed.FS
 
 type (
 	persisterDependencies interface {
-		IdentityTraitsSchemas(ctx context.Context) (schema.Schemas, error)
+		schema.IdentityTraitsProvider
 		identity.ValidationProvider
 		x.LoggingProvider
 		config.Provider