diff --git a/cmd/health.go b/cmd/health.go index a9405c1..db7b83d 100644 --- a/cmd/health.go +++ b/cmd/health.go @@ -14,14 +14,20 @@ 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() { @@ -29,7 +35,7 @@ func init() { } // 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,19 +45,19 @@ 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(".") } @@ -59,8 +65,8 @@ func checkHealth(requiredStorageNodes int) error { } } -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 {