feat: add option to add custom health checks (#1225)

internal/driver/registry_default.go

@@ -76,6 +76,7 @@ type (
 		defaultHttpMiddlewares    []func(rw http.ResponseWriter, r *http.Request, next http.HandlerFunc)
 		grpcTransportCredentials  credentials.TransportCredentials
 		defaultMigrationOptions   []popx.MigrationBoxOption
+		healthReadyCheckers       healthx.ReadyCheckers
 	}
 	ReadHandler interface {
 		RegisterReadRoutes(r *x.ReadRouter)
@@ -119,7 +120,10 @@ func (r *RegistryDefault) Config(ctx context.Context) *config.Config {
 
 func (r *RegistryDefault) HealthHandler() *healthx.Handler {
 	if r.healthH == nil {
-		r.healthH = healthx.NewHandler(r.Writer(), config.Version, healthx.ReadyCheckers{})
+		if r.healthReadyCheckers == nil {
+			r.healthReadyCheckers = healthx.ReadyCheckers{}
+		}
+		r.healthH = healthx.NewHandler(r.Writer(), config.Version, r.healthReadyCheckers)
 	}
 
 	return r.healthH

internal/driver/registry_factory.go

@@ -15,26 +15,21 @@ import (
 	"testing"
 
 	"github.com/ory/x/configx"
+	"github.com/ory/x/logrusx"
 	"github.com/ory/x/otelx"
 	"github.com/ory/x/tlsx"
+	"github.com/pkg/errors"
+	"github.com/sirupsen/logrus"
+	"github.com/spf13/pflag"
+	"github.com/stretchr/testify/require"
 	"go.opentelemetry.io/otel/trace"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
 
-	"github.com/ory/keto/ketoctx"
-
+	"github.com/ory/keto/internal/driver/config"
 	"github.com/ory/keto/internal/namespace"
 	"github.com/ory/keto/internal/x/dbx"
-
-	"github.com/sirupsen/logrus"
-
-	"github.com/pkg/errors"
-
-	"github.com/ory/x/logrusx"
-	"github.com/spf13/pflag"
-	"github.com/stretchr/testify/require"
-
-	"github.com/ory/keto/internal/driver/config"
+	"github.com/ory/keto/ketoctx"
 )
 
 // createFile writes the content to a temporary file, returning the path.
@@ -88,6 +83,7 @@ func NewDefaultRegistry(ctx context.Context, flags *pflag.FlagSet, withoutNetwor
 		defaultStreamInterceptors: options.GRPCStreamInterceptors(),
 		defaultHttpMiddlewares:    options.HTTPMiddlewares(),
 		defaultMigrationOptions:   options.MigrationOptions(),
+		healthReadyCheckers:       options.ReadyCheckers(),
 	}
 
 	init := r.Init
@@ -150,7 +146,7 @@ func WithGRPCStreamInterceptors(i ...grpc.StreamServerInterceptor) TestRegistryO
 	}
 }
 func WithTracer(tracer trace.Tracer) TestRegistryOption {
-	return func(t testing.TB, r *RegistryDefault) {
+	return func(_ testing.TB, r *RegistryDefault) {
 		r.tracer = new(otelx.Tracer).WithOTLP(tracer)
 	}
 }

ketoctx/options.go

@@ -6,6 +6,7 @@ package ketoctx
 import (
 	"net/http"
 
+	"github.com/ory/x/healthx"
 	"github.com/ory/x/logrusx"
 	"github.com/ory/x/popx"
 	"google.golang.org/grpc"
@@ -19,46 +20,67 @@ type (
 		grpcUnaryInterceptors  []grpc.UnaryServerInterceptor
 		grpcStreamInterceptors []grpc.StreamServerInterceptor
 		migrationOpts          []popx.MigrationBoxOption
+		readyCheckers          healthx.ReadyCheckers
 	}
 	Option func(o *opts)
 )
 
+// WithLogger sets the logger.
 func WithLogger(l *logrusx.Logger) Option {
 	return func(o *opts) {
 		o.logger = l
 	}
 }
 
+// WithContextualizer sets the contextualizer.
 func WithContextualizer(ctxer Contextualizer) Option {
 	return func(o *opts) {
 		o.contextualizer = ctxer
 	}
 }
 
+// WithHTTPMiddlewares adds HTTP middlewares to the list of HTTP middlewares.
 func WithHTTPMiddlewares(m ...func(rw http.ResponseWriter, r *http.Request, next http.HandlerFunc)) Option {
 	return func(o *opts) {
 		o.httpMiddlewares = m
 	}
 }
 
+// WithGRPCUnaryInterceptors adds gRPC unary interceptors to the list of gRPC
+// interceptors.
 func WithGRPCUnaryInterceptors(i ...grpc.UnaryServerInterceptor) Option {
 	return func(o *opts) {
 		o.grpcUnaryInterceptors = i
 	}
 }
 
+// WithGRPCStreamInterceptors adds gRPC stream interceptors to the list of gRPC
+// stream interceptors.
 func WithGRPCStreamInterceptors(i ...grpc.StreamServerInterceptor) Option {
 	return func(o *opts) {
 		o.grpcStreamInterceptors = i
 	}
 }
 
+// WithMigrationOptions adds migration options to the list of migration options.
 func WithMigrationOptions(o ...popx.MigrationBoxOption) Option {
 	return func(opts *opts) {
 		opts.migrationOpts = o
 	}
 }
 
+// WithReadinessCheck adds a new readness health checker to the list of
+// checkers. Can be called multiple times. If the name is already taken, the
+// checker will be overwritten.
+func WithReadinessCheck(name string, rc healthx.ReadyChecker) Option {
+	return func(o *opts) {
+		if o.readyCheckers == nil {
+			o.readyCheckers = make(healthx.ReadyCheckers)
+		}
+		o.readyCheckers[name] = rc
+	}
+}
+
 func (o *opts) Logger() *logrusx.Logger {
 	return o.logger
 }
@@ -83,6 +105,10 @@ func (o *opts) MigrationOptions() []popx.MigrationBoxOption {
 	return o.migrationOpts
 }
 
+func (o *opts) ReadyCheckers() healthx.ReadyCheckers {
+	return o.readyCheckers
+}
+
 func Options(options ...Option) *opts {
 	o := &opts{
 		contextualizer: &DefaultContextualizer{},