fix: check health status in status command

ory · Mar 9, 2021 · 21c64d4 · 21c64d4
1 parent ca5ccb9
commit 21c64d4
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 11 deletions.
diff --git a/cmd/status/root.go b/cmd/status/root.go
@@ -63,13 +63,45 @@ func newStatusCmd() *cobra.Command {
 
 			c := grpcHealthV1.NewHealthClient(conn)
 
-			resp, err := c.Check(cmd.Context(), &grpcHealthV1.HealthCheckRequest{})
-			if err != nil {
-				_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "Unable to get a check response: %+v\n", err)
-				return cmdx.FailSilently(cmd)
+			var status interface {
+				GetStatus() grpcHealthV1.HealthCheckResponse_ServingStatus
+			}
+			if block {
+				wc, err := c.Watch(cmd.Context(), &grpcHealthV1.HealthCheckRequest{})
+				if err != nil {
+					_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "Could not start watching the status: %+v\n", err)
+					return cmdx.FailSilently(cmd)
+				}
+				//_ = wc.CloseSend()
+
+				for {
+					select {
+					case <-cmd.Context().Done():
+						return nil
+					default:
+					}
+
+					status, err = wc.Recv()
+					if err != nil {
+						_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "Could not watch the status: %+v\n", err)
+						return cmdx.FailSilently(cmd)
+					}
+
+					if status.GetStatus() == grpcHealthV1.HealthCheckResponse_SERVING {
+						break
+					}
+
+					_, _ = loudPrinter.Printf("Got the status %s, waiting until %s.\n", status.GetStatus(), grpcHealthV1.HealthCheckResponse_SERVING)
+				}
+			} else {
+				status, err = c.Check(cmd.Context(), &grpcHealthV1.HealthCheckRequest{})
+				if err != nil {
+					_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "Unable to get a check response: %+v\n", err)
+					return cmdx.FailSilently(cmd)
+				}
 			}
 
-			_, _ = fmt.Fprintln(cmd.OutOrStdout(), resp.Status.String())
+			_, _ = fmt.Fprintln(cmd.OutOrStdout(), status.GetStatus().String())
 			return nil
 		},
 	}

diff --git a/internal/driver/registry_default.go b/internal/driver/registry_default.go
@@ -7,6 +7,7 @@ import (
 	grpc_logrus "github.com/grpc-ecosystem/go-grpc-middleware/logging/logrus"
 	"github.com/ory/x/reqlog"
 	"github.com/urfave/negroni"
+	"google.golang.org/grpc/reflection"
 
 	grpcMiddleware "github.com/grpc-ecosystem/go-grpc-middleware"
 	"github.com/julienschmidt/httprouter"
@@ -239,6 +240,7 @@ func (r *RegistryDefault) ReadGRPCServer() *grpc.Server {
 	)
 
 	grpcHealthV1.RegisterHealthServer(s, r.HealthServer())
+	reflection.Register(s)
 
 	for _, h := range r.allHandlers() {
 		h.RegisterReadGRPC(s)
@@ -264,6 +266,7 @@ func (r *RegistryDefault) WriteGRPCServer() *grpc.Server {
 	)
 
 	grpcHealthV1.RegisterHealthServer(s, r.HealthServer())
+	reflection.Register(s)
 
 	for _, h := range r.allHandlers() {
 		h.RegisterWriteGRPC(s)

diff --git a/internal/x/registry.go b/internal/x/registry.go
@@ -1,8 +1,6 @@
 package x
 
 import (
-	"github.com/gorilla/sessions"
-
 	"github.com/ory/herodot"
 	"github.com/ory/x/logrusx"
 )
@@ -14,7 +12,3 @@ type LoggerProvider interface {
 type WriterProvider interface {
 	Writer() herodot.Writer
 }
-
-type RegistryCookieStore interface {
-	CookieStore() sessions.Store
-}