From d77ea36e4668762a70faa858c4e45328218bc8b8 Mon Sep 17 00:00:00 2001
From: Arne Luenser <arne.luenser@ory.sh>
Date: Mon, 24 Jul 2023 15:28:15 +0200
Subject: [PATCH] feat: add `hydra migrate status` subcommand, with optional
 `--block` flag to wait for migrations to be finished by another instance

---
 cmd/cli/handler_migrate.go | 65 ++++++++++++++++++++++++++++++++++----
 cmd/migrate_status.go      | 27 ++++++++++++++++
 cmd/root.go                |  1 +
 3 files changed, 87 insertions(+), 6 deletions(-)
 create mode 100644 cmd/migrate_status.go

diff --git a/cmd/cli/handler_migrate.go b/cmd/cli/handler_migrate.go
index 6bb293d127f..0172584028f 100644
--- a/cmd/cli/handler_migrate.go
+++ b/cmd/cli/handler_migrate.go
@@ -13,7 +13,9 @@ import (
 	"path/filepath"
 	"regexp"
 	"strings"
+	"time"
 
+	"github.com/ory/x/popx"
 	"github.com/ory/x/servicelocatorx"
 
 	"github.com/pkg/errors"
@@ -28,6 +30,7 @@ import (
 
 	"github.com/ory/hydra/v2/driver"
 	"github.com/ory/hydra/v2/driver/config"
+	"github.com/ory/hydra/v2/persistence"
 	"github.com/ory/x/flagx"
 )
 
@@ -259,7 +262,7 @@ func (h *MigrateHandler) MigrateGen(cmd *cobra.Command, args []string) {
 	os.Exit(0)
 }
 
-func (h *MigrateHandler) MigrateSQL(cmd *cobra.Command, args []string) (err error) {
+func makePersister(cmd *cobra.Command, args []string) (p persistence.Persister, err error) {
 	var d driver.Registry
 
 	if flagx.MustGetBool(cmd, "read-from-env") {
@@ -275,16 +278,16 @@ func (h *MigrateHandler) MigrateSQL(cmd *cobra.Command, args []string) (err erro
 				driver.SkipNetworkInit(),
 			})
 		if err != nil {
-			return err
+			return nil, err
 		}
 		if len(d.Config().DSN()) == 0 {
 			_, _ = fmt.Fprintln(cmd.ErrOrStderr(), "When using flag -e, environment variable DSN must be set.")
-			return cmdx.FailSilently(cmd)
+			return nil, cmdx.FailSilently(cmd)
 		}
 	} else {
 		if len(args) != 1 {
 			_, _ = fmt.Fprintln(cmd.ErrOrStderr(), "Please provide the database URL.")
-			return cmdx.FailSilently(cmd)
+			return nil, cmdx.FailSilently(cmd)
 		}
 		d, err = driver.New(
 			cmd.Context(),
@@ -300,11 +303,17 @@ func (h *MigrateHandler) MigrateSQL(cmd *cobra.Command, args []string) (err erro
 				driver.SkipNetworkInit(),
 			})
 		if err != nil {
-			return err
+			return nil, err
 		}
 	}
+	return d.Persister(), nil
+}
 
-	p := d.Persister()
+func (h *MigrateHandler) MigrateSQL(cmd *cobra.Command, args []string) (err error) {
+	p, err := makePersister(cmd, args)
+	if err != nil {
+		return err
+	}
 	conn := p.Connection(context.Background())
 	if conn == nil {
 		_, _ = fmt.Fprintln(cmd.ErrOrStderr(), "Migrations can only be executed against a SQL-compatible driver but DSN is not a SQL source.")
@@ -349,3 +358,47 @@ func (h *MigrateHandler) MigrateSQL(cmd *cobra.Command, args []string) (err erro
 	_, _ = fmt.Fprintln(cmd.OutOrStdout(), "Successfully applied migrations!")
 	return nil
 }
+
+func (h *MigrateHandler) MigrateStatus(cmd *cobra.Command, args []string) error {
+	p, err := makePersister(cmd, args)
+	if err != nil {
+		return err
+	}
+	conn := p.Connection(context.Background())
+	if conn == nil {
+		_, _ = fmt.Fprintln(cmd.ErrOrStderr(), "Migrations can only be checked against a SQL-compatible driver but DSN is not a SQL source.")
+		return cmdx.FailSilently(cmd)
+	}
+
+	if err := conn.Open(); err != nil {
+		_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "Could not open the database connection:\n%+v\n", err)
+		return cmdx.FailSilently(cmd)
+	}
+
+	block := flagx.MustGetBool(cmd, "block")
+	ctx := cmd.Context()
+	s, err := p.MigrationStatus(ctx)
+	if err != nil {
+		_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "Could not get migration status: %+v\n", err)
+		return cmdx.FailSilently(cmd)
+	}
+
+	for block && s.HasPending() {
+		_, _ = fmt.Fprintf(cmd.OutOrStdout(), "Waiting for migrations to finish...\n")
+		for _, m := range s {
+			if m.State == popx.Pending {
+				_, _ = fmt.Fprintf(cmd.OutOrStdout(), " - %s\n", m.Name)
+			}
+		}
+		time.Sleep(time.Second)
+		s, err = p.MigrationStatus(ctx)
+		if err != nil {
+			_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "Could not get migration status: %+v\n", err)
+			return cmdx.FailSilently(cmd)
+		}
+	}
+
+	cmdx.PrintTable(cmd, s)
+	return nil
+
+}
diff --git a/cmd/migrate_status.go b/cmd/migrate_status.go
new file mode 100644
index 00000000000..397f5e86e48
--- /dev/null
+++ b/cmd/migrate_status.go
@@ -0,0 +1,27 @@
+// Copyright © 2023 Ory Corp
+// SPDX-License-Identifier: Apache-2.0
+
+package cmd
+
+import (
+	"github.com/ory/x/configx"
+	"github.com/ory/x/servicelocatorx"
+
+	"github.com/spf13/cobra"
+
+	"github.com/ory/hydra/v2/cmd/cli"
+	"github.com/ory/hydra/v2/driver"
+)
+
+func NewMigrateStatusCmd(slOpts []servicelocatorx.Option, dOpts []driver.OptionsModifier, cOpts []configx.OptionModifier) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "status",
+		Short: "Get the current migration status",
+		RunE:  cli.NewHandler(slOpts, dOpts, cOpts).Migration.MigrateStatus,
+	}
+
+	cmd.Flags().BoolP("read-from-env", "e", false, "If set, reads the database connection string from the environment variable DSN or config file key dsn.")
+	cmd.Flags().Bool("block", false, "Block until all migrations have been applied")
+
+	return cmd
+}
diff --git a/cmd/root.go b/cmd/root.go
index 24f6e8a68aa..6feabdb8103 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -73,6 +73,7 @@ func RegisterCommandRecursive(parent *cobra.Command, slOpts []servicelocatorx.Op
 	migrateCmd := NewMigrateCmd()
 	migrateCmd.AddCommand(NewMigrateGenCmd())
 	migrateCmd.AddCommand(NewMigrateSqlCmd(slOpts, dOpts, cOpts))
+	migrateCmd.AddCommand(NewMigrateStatusCmd(slOpts, dOpts, cOpts))
 
 	serveCmd := NewServeCmd()
 	serveCmd.AddCommand(NewServeAdminCmd(slOpts, dOpts, cOpts))