cmd/atlas: extendctx can return an error (#2918)

cmd/atlas/internal/cloudapi/client.go

@@ -246,6 +246,9 @@ func (c *Client) ReportMigration(ctx context.Context, input ReportMigrationInput
 	return payload.ReportMigration.URL, nil
 }
 
+// ErrUnauthorized is returned when the server returns a 401 status code.
+var ErrUnauthorized = errors.New(http.StatusText(http.StatusUnauthorized))
+
 func (c *Client) post(ctx context.Context, query string, vars, data any) error {
 	body, err := json.Marshal(struct {
 		Query     string `json:"query"`
@@ -267,7 +270,10 @@ func (c *Client) post(ctx context.Context, query string, vars, data any) error {
 		return err
 	}
 	defer req.Body.Close()
-	if res.StatusCode != http.StatusOK {
+	switch {
+	case res.StatusCode == http.StatusUnauthorized:
+		return ErrUnauthorized
+	case res.StatusCode != http.StatusOK:
 		var v struct {
 			Errors errlist `json:"errors,omitempty"`
 		}

cmd/atlas/main.go

@@ -49,9 +49,14 @@ func main() {
 		<-stop // will not block if no signal received due to main routine exiting
 		os.Exit(1)
 	}()
-	ctx, done := initialize(extendContext(ctx))
+	ctx, err := extendContext(ctx)
+	if err != nil {
+		fmt.Fprintln(os.Stderr, err)
+		os.Exit(1)
+	}
+	ctx, done := initialize(ctx)
 	update := checkForUpdate(ctx)
-	err := cmdapi.Root.ExecuteContext(ctx)
+	err = cmdapi.Root.ExecuteContext(ctx)
 	if u := update(); u != "" {
 		fmt.Fprintln(os.Stderr, u)
 	}

cmd/atlas/main_oss.go

@@ -10,7 +10,9 @@ import (
 	"context"
 )
 
-func extendContext(ctx context.Context) context.Context { return ctx }
+func extendContext(ctx context.Context) (context.Context, error) {
+	return ctx, nil
+}
 
 func vercheckEndpoint(context.Context) string {
 	return vercheckURL