github.com/apache/beam/sdks/v2@v2.48.2/go/container/tools/provision.go (about)

     1  // Licensed to the Apache Software Foundation (ASF) under one or more
     2  // contributor license agreements.  See the NOTICE file distributed with
     3  // this work for additional information regarding copyright ownership.
     4  // The ASF licenses this file to You under the Apache License, Version 2.0
     5  // (the "License"); you may not use this file except in compliance with
     6  // the License.  You may obtain a copy of the License at
     7  //
     8  //    http://www.apache.org/licenses/LICENSE-2.0
     9  //
    10  // Unless required by applicable law or agreed to in writing, software
    11  // distributed under the License is distributed on an "AS IS" BASIS,
    12  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  // See the License for the specific language governing permissions and
    14  // limitations under the License.
    15  
    16  // Package tools contains utilities for Beam bootloader containers, such as
    17  // for obtaining runtime provision information -- such as pipeline options.
    18  // or for logging to the log service.
    19  //
    20  // For Beam Internal use.
    21  package tools
    22  
    23  import (
    24  	"context"
    25  	"encoding/json"
    26  	"errors"
    27  	"fmt"
    28  	"time"
    29  
    30  	fnpb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/fnexecution_v1"
    31  	"github.com/apache/beam/sdks/v2/go/pkg/beam/util/grpcx"
    32  	"github.com/golang/protobuf/jsonpb"
    33  	google_pb "github.com/golang/protobuf/ptypes/struct"
    34  )
    35  
    36  // ProvisionInfo returns the runtime provisioning info for the worker.
    37  func ProvisionInfo(ctx context.Context, endpoint string) (*fnpb.ProvisionInfo, error) {
    38  	cc, err := grpcx.Dial(ctx, endpoint, 2*time.Minute)
    39  	if err != nil {
    40  		return nil, err
    41  	}
    42  	defer cc.Close()
    43  
    44  	client := fnpb.NewProvisionServiceClient(cc)
    45  
    46  	resp, err := client.GetProvisionInfo(ctx, &fnpb.GetProvisionInfoRequest{})
    47  	if err != nil {
    48  		return nil, fmt.Errorf("failed to get manifest: %w", err)
    49  	}
    50  	if resp.GetInfo() == nil {
    51  		return nil, errors.New("empty manifest")
    52  	}
    53  	return resp.GetInfo(), nil
    54  }
    55  
    56  // OptionsToProto converts pipeline options to a proto struct via JSON.
    57  func OptionsToProto(v any) (*google_pb.Struct, error) {
    58  	data, err := json.Marshal(v)
    59  	if err != nil {
    60  		return nil, err
    61  	}
    62  	return JSONToProto(string(data))
    63  }
    64  
    65  // JSONToProto converts JSON-encoded pipeline options to a proto struct.
    66  func JSONToProto(data string) (*google_pb.Struct, error) {
    67  	var out google_pb.Struct
    68  	if err := jsonpb.UnmarshalString(string(data), &out); err != nil {
    69  		return nil, err
    70  	}
    71  	return &out, nil
    72  }
    73  
    74  // ProtoToOptions converts pipeline options from a proto struct via JSON.
    75  func ProtoToOptions(opt *google_pb.Struct, v any) error {
    76  	data, err := ProtoToJSON(opt)
    77  	if err != nil {
    78  		return err
    79  	}
    80  	return json.Unmarshal([]byte(data), v)
    81  }
    82  
    83  // ProtoToJSON converts pipeline options from a proto struct to JSON.
    84  func ProtoToJSON(opt *google_pb.Struct) (string, error) {
    85  	if opt == nil {
    86  		return "{}", nil
    87  	}
    88  	return (&jsonpb.Marshaler{}).MarshalToString(opt)
    89  }