Added support for multiple schedules (#765)

cuvva · Jul 10, 2023 · f38556b · f38556b
1 parent 827a80d
commit f38556b
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 7 deletions.
diff --git a/tools/oncallcalc/app/generate_report.go b/tools/oncallcalc/app/generate_report.go
@@ -8,7 +8,7 @@ import (
 )
 
 // GenerateRota generates a rota based on a schedule and month
-func (a *App) GenerateRota(scheduleID string, year int, month time.Month) (oncallcalc.Rota, interface{}, error) {
+func (a *App) GenerateRota(scheduleIDs []string, year int, month time.Month) (oncallcalc.Rota, interface{}, error) {
 	loc, err := time.LoadLocation("Europe/London")
 
 	if err != nil {
@@ -18,7 +18,7 @@ func (a *App) GenerateRota(scheduleID string, year int, month time.Month) (oncal
 	monthStart := time.Date(year, month, 1, 0, 0, 0, 0, loc)
 	monthEnd := monthStart.AddDate(0, 1, 0)
 
-	pdshifts, badShift, err := a.ListPagerDutyShifts(scheduleID, monthStart, monthEnd)
+	pdshifts, badShift, err := a.ListPagerDutyShifts(scheduleIDs, monthStart, monthEnd)
 	if err != nil {
 		return nil, badShift, err
 	}

diff --git a/tools/oncallcalc/app/list_pd_shifts.go b/tools/oncallcalc/app/list_pd_shifts.go
@@ -17,15 +17,15 @@ type PDShift struct {
 }
 
 // ListPagerDutyShifts returns a slice of shifts pagerduty is tracking for a schedule
-func (a *App) ListPagerDutyShifts(scheduleID string, start, end time.Time) ([]*PDShift, interface{}, error) {
+func (a *App) ListPagerDutyShifts(scheduleIDs []string, start, end time.Time) ([]*PDShift, interface{}, error) {
 	var shifts []*PDShift
 
 	fmt.Println("Requesting on call dates for ", start.Format(time.RFC3339), " until ", end.Format(time.RFC3339))
 
 	// @TODO context should be injected from the cobra entry point
 	res, err := a.pagerduty.ListOnCallsWithContext(context.Background(), pagerduty.ListOnCallOptions{
 		Includes:    []string{"users"},
-		ScheduleIDs: []string{scheduleID},
+		ScheduleIDs: scheduleIDs,
 		Since:       start.Format(time.RFC3339),
 		Until:       end.Format(time.RFC3339),
 		Limit:       100, // TODO implement pagination if we require more values

diff --git a/tools/oncallcalc/commands/generate_report.go b/tools/oncallcalc/commands/generate_report.go
@@ -11,12 +11,12 @@ import (
 	"github.com/spf13/cobra"
 )
 
-var ScheduleID string
+var ScheduleIDs []string
 var TimeIn string
 var Verbose bool
 
 func init() {
-	GenerateReportCmd.Flags().StringVarP(&ScheduleID, "schedule_id", "s", "PKICNIO", "Schedule ID from PagerDuty")
+	GenerateReportCmd.Flags().StringArrayVarP(&ScheduleIDs, "schedule_id", "s", []string{"PKICNIO", "PKW2AKK"}, "Schedule IDs from PagerDuty")
 	GenerateReportCmd.Flags().StringVarP(&TimeIn, "time", "t", "Jan 2020", "Which month and year should we look at?")
 	GenerateReportCmd.Flags().BoolVarP(&Verbose, "verbose", "v", false, "Verbose mode?")
 }
@@ -49,7 +49,7 @@ var GenerateReportCmd = &cobra.Command{
 			return err
 		}
 
-		rota, debug, err := app.GenerateRota(ScheduleID, timeIn.Year(), timeIn.Month())
+		rota, debug, err := app.GenerateRota(ScheduleIDs, timeIn.Year(), timeIn.Month())
 		if err != nil {
 			if Verbose {
 				cmd.Printf("%+v\n", debug)