From 1fb15331c53d0f1fb747ffd3d24cada09987c77f Mon Sep 17 00:00:00 2001
From: Andrew Nester <andrew.nester@databricks.com>
Date: Thu, 25 Jan 2024 11:55:17 +0000
Subject: [PATCH] Bring back `--json` flag for workspace-conf set-status
 command (#1151)

## Changes
--json flag was removed from this command when MustUseJson / CanUseJson
generator functions were introduced which did not take requests types of
map.

This PR bring the flag back.

Relies on this Go SDK change:
https://github.com/databricks/databricks-sdk-go/pull/786
---
 .../workspace-conf/workspace-conf.go          | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/cmd/workspace/workspace-conf/workspace-conf.go b/cmd/workspace/workspace-conf/workspace-conf.go
index 687c31ad..99207ffa 100755
--- a/cmd/workspace/workspace-conf/workspace-conf.go
+++ b/cmd/workspace/workspace-conf/workspace-conf.go
@@ -3,8 +3,11 @@
 package workspace_conf
 
 import (
+	"fmt"
+
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdio"
+	"github.com/databricks/cli/libs/flags"
 	"github.com/databricks/databricks-sdk-go/service/settings"
 	"github.com/spf13/cobra"
 )
@@ -106,8 +109,10 @@ func newSetStatus() *cobra.Command {
 	cmd := &cobra.Command{}
 
 	var setStatusReq settings.WorkspaceConf
+	var setStatusJson flags.JsonFlag
 
 	// TODO: short flags
+	cmd.Flags().Var(&setStatusJson, "json", `either inline JSON string or @path/to/file.json with request body`)
 
 	cmd.Use = "set-status"
 	cmd.Short = `Enable/disable features.`
@@ -118,16 +123,20 @@ func newSetStatus() *cobra.Command {
 
 	cmd.Annotations = make(map[string]string)
 
-	cmd.Args = func(cmd *cobra.Command, args []string) error {
-		check := cobra.ExactArgs(0)
-		return check(cmd, args)
-	}
-
 	cmd.PreRunE = root.MustWorkspaceClient
 	cmd.RunE = func(cmd *cobra.Command, args []string) (err error) {
 		ctx := cmd.Context()
 		w := root.WorkspaceClient(ctx)
 
+		if cmd.Flags().Changed("json") {
+			err = setStatusJson.Unmarshal(&setStatusReq)
+			if err != nil {
+				return err
+			}
+		} else {
+			return fmt.Errorf("please provide command input in JSON format by specifying the --json flag")
+		}
+
 		err = w.WorkspaceConf.SetStatus(ctx, setStatusReq)
 		if err != nil {
 			return err