go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/buildbucket/cli/prop.go (about)

     1  // Copyright 2019 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package cli
    16  
    17  import (
    18  	"flag"
    19  	"fmt"
    20  	"os"
    21  	"strings"
    22  
    23  	"github.com/golang/protobuf/jsonpb"
    24  
    25  	"go.chromium.org/luci/common/data/stringset"
    26  
    27  	"google.golang.org/protobuf/encoding/protojson"
    28  	"google.golang.org/protobuf/types/known/structpb"
    29  )
    30  
    31  type propertiesFlag struct {
    32  	props         *structpb.Struct
    33  	first         bool
    34  	explicitlySet stringset.Set
    35  }
    36  
    37  // PropertiesFlag returns a flag.Getter that can read property values into props.
    38  //
    39  // If a flag value starts with @, properties are read from the JSON file at the
    40  // path that follows @. Example:
    41  //
    42  //	-p @my_properties.json
    43  //
    44  // This form can be used only in the first flag value.
    45  //
    46  // Otherwise, a flag value must have name=value form.
    47  // If the property value is valid JSON, then it is parsed as JSON;
    48  // otherwise treated as a string. Example:
    49  //
    50  //	-p foo=1 -p 'bar={"a": 2}'
    51  //
    52  // Different property names can be specified multiple times.
    53  //
    54  // Panics if props is nil.
    55  // String() of the return value panics if marshaling of props to JSON fails.
    56  func PropertiesFlag(props *structpb.Struct) flag.Getter {
    57  	if props == nil {
    58  		panic("props is nil")
    59  	}
    60  	return &propertiesFlag{
    61  		props: props,
    62  		first: true,
    63  		// We don't expect a lot of properties.
    64  		// After a certain number, it is easier to pass properties via a file.
    65  		// Choose an arbitrary number.
    66  		explicitlySet: stringset.New(4),
    67  	}
    68  }
    69  
    70  func (f *propertiesFlag) String() string {
    71  	// https://godoc.org/flag#Value says that String() may be called with a
    72  	// zero-valued receiver.
    73  	if f == nil || f.props == nil {
    74  		return ""
    75  	}
    76  
    77  	ret, err := (&jsonpb.Marshaler{}).MarshalToString(f.props)
    78  	if err != nil {
    79  		panic(err)
    80  	}
    81  	return ret
    82  }
    83  
    84  func (f *propertiesFlag) Get() any {
    85  	return f.props
    86  }
    87  
    88  func (f *propertiesFlag) Set(s string) error {
    89  	first := f.first
    90  	f.first = false
    91  
    92  	if strings.HasPrefix(s, "@") {
    93  		if !first {
    94  			return fmt.Errorf("value that starts with @ must be the first value for the flag")
    95  		}
    96  		fileName := s[1:]
    97  
    98  		file, err := os.Open(fileName)
    99  		if err != nil {
   100  			return err
   101  		}
   102  		defer file.Close()
   103  		return jsonpb.Unmarshal(file, f.props)
   104  	}
   105  
   106  	parts := strings.SplitN(s, "=", 2)
   107  	if len(parts) == 1 {
   108  		return fmt.Errorf("invalid property %q: no equals sign", s)
   109  	}
   110  	name := strings.TrimSpace(parts[0])
   111  	value := strings.TrimSpace(parts[1])
   112  
   113  	if f.explicitlySet.Has(name) {
   114  		return fmt.Errorf("duplicate property %q", name)
   115  	}
   116  	f.explicitlySet.Add(name)
   117  
   118  	// Try parsing as JSON.
   119  	// Note: jsonpb cannot unmarshal structpb.Value from JSON natively,
   120  	// so we have to wrap JSON value in an object.
   121  	wrappedJSON := fmt.Sprintf(`{"a": %s}`, value)
   122  	buf := &structpb.Struct{}
   123  	if f.props.Fields == nil {
   124  		f.props.Fields = map[string]*structpb.Value{}
   125  	}
   126  	if err := protojson.Unmarshal([]byte(wrappedJSON), buf); err == nil {
   127  		f.props.Fields[name] = buf.Fields["a"]
   128  	} else {
   129  		// Treat as string.
   130  		f.props.Fields[name] = &structpb.Value{
   131  			Kind: &structpb.Value_StringValue{StringValue: value},
   132  		}
   133  	}
   134  	return nil
   135  }