cmd/health: update health command to support more DBs

Health command was created specifically for the nodes table. This update allows it to be used for other tables as well. Use the flags to pass in the table name and expected amount of records in the DB.

Change-Id: If215d580cde74e04e2c2554bf8422208543b8bff

cmd/health.go

@@ -14,22 +14,28 @@ import (
 	"github.com/zeebo/errs/v2"
 )
 
+var table string
+var number int
+
 func healthCmd() *cobra.Command {
-	return &cobra.Command{
+	cmd := &cobra.Command{
 		Use:   "health",
 		Short: "wait until cluster is healthy (10 storagenodes are registered in the db)",
 		RunE: func(cmd *cobra.Command, args []string) error {
-			return checkHealth(10)
+			return checkHealth(table, number)
 		},
 	}
+	cmd.PersistentFlags().StringVarP(&table, "table", "t", "nodes", "table to use for health check")
+	cmd.PersistentFlags().IntVarP(&number, "number", "n", 10, "number of entries to expect in the table")
+	return cmd
 }
 
 func init() {
 	rootCmd.AddCommand(healthCmd())
 }
 
 // checkHealth polls the database until all storagenodes are checked in.
-func checkHealth(requiredStorageNodes int) error {
+func checkHealth(table string, records int) error {
 	prevCount := -1
 	for {
 		time.Sleep(1 * time.Second)
@@ -39,28 +45,28 @@ func checkHealth(requiredStorageNodes int) error {
 			continue
 		}
 
-		count, err := registeredNodeCount(db)
+		count, err := dbRecordCount(db, table)
 		_ = db.Close()
 		if err != nil {
-			fmt.Printf("Couldn't query database for nodes: %s\n", err.Error())
+			fmt.Printf("Couldn't query database for records: %s\n", err.Error())
 			continue
 		}
-		if count == requiredStorageNodes {
+		if count == records {
 			fmt.Println()
-			fmt.Println("Storj cluster is healthy")
+			fmt.Println(table, "has", records, "records")
 			return nil
 		}
 		if count != prevCount {
-			fmt.Printf("Found only %d storagenodes in the database ", count)
+			fmt.Printf("Found only %d records in the database ", count)
 		} else {
 			fmt.Print(".")
 		}
 		prevCount = count
 	}
 }
 
-func registeredNodeCount(db *sql.DB) (int, error) {
-	row := db.QueryRow("select count(*) from nodes")
+func dbRecordCount(db *sql.DB, table string) (int, error) {
+	row := db.QueryRow("select count(*) from " + table)
 	var count int
 	err := row.Scan(&count)
 	if err != nil {