github.com/google/cloudprober@v0.11.3/contrib/gcp/cmd/bigquery_probe.go (about)

     1  // Copyright 2020 The Cloudprober 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  /*
    16  Bigquery_probe is an external Cloudprober probe suitable for blackbox-probing
    17  the GCP BigQuery API.
    18  
    19  This binary assumes that the environment variable
    20  $GOOGLE_APPLICATION_CREDENTIALS has been set following the instructions at
    21  https://cloud.google.com/docs/authentication/production
    22  */
    23  package main
    24  
    25  import (
    26  	"context"
    27  	"fmt"
    28  
    29  	"flag"
    30  	"github.com/golang/glog"
    31  	"github.com/golang/protobuf/proto"
    32  	"github.com/google/cloudprober/contrib/gcp/bigquery"
    33  	serverpb "github.com/google/cloudprober/probes/external/proto"
    34  	"github.com/google/cloudprober/probes/external/serverutils"
    35  )
    36  
    37  var (
    38  	projectID  = flag.String("project_id", "", "GCP project ID to connect to.")
    39  	serverMode = flag.Bool("server_mode", false, "Whether to run in server mode.")
    40  	table      = flag.String("table", "", "Table to probe (specified as \"dataset.table\"). "+
    41  		"If empty, the probe will simply try to connect to BQ and execute 'SELECT 1'. "+
    42  		"If running in server mode, the 'table' option from the configuration will override this flag.")
    43  )
    44  
    45  func parseProbeRequest(request *serverpb.ProbeRequest) map[string]string {
    46  	m := make(map[string]string)
    47  	for _, opt := range request.Options {
    48  		m[*opt.Name] = *opt.Value
    49  	}
    50  	return m
    51  }
    52  
    53  func main() {
    54  	flag.Parse()
    55  
    56  	if *projectID == "" {
    57  		glog.Exitf("--project_id must be specified")
    58  	}
    59  
    60  	dstTable := *table
    61  	ctx := context.Background()
    62  	runner, err := bigquery.NewRunner(ctx, *projectID)
    63  	if err != nil {
    64  		glog.Fatal(err)
    65  	}
    66  
    67  	if *serverMode {
    68  		serverutils.Serve(func(request *serverpb.ProbeRequest, reply *serverpb.ProbeReply) {
    69  
    70  			opts := parseProbeRequest(request)
    71  			if val, ok := opts["table"]; ok {
    72  				dstTable = val
    73  				glog.Infof("--table set to %q by ProbeRequest config", val)
    74  			}
    75  			payload, err := bigquery.Probe(ctx, runner, dstTable)
    76  			reply.Payload = proto.String(payload)
    77  			if err != nil {
    78  				reply.ErrorMessage = proto.String(err.Error())
    79  			}
    80  		})
    81  	}
    82  
    83  	payload, err := bigquery.Probe(ctx, runner, dstTable)
    84  	if err != nil {
    85  		glog.Fatal(err)
    86  	}
    87  	fmt.Println(payload)
    88  
    89  }