From 8d0d3ff31af3ed45725e07336fc4c045f2053234 Mon Sep 17 00:00:00 2001 From: shollyman Date: Tue, 9 Aug 2022 16:05:00 -0500 Subject: [PATCH] chore(bigquery/storage/managedwriter): add client instrumentation (#6480) * chore(bigquery/storage/managedwriter): add client instrumentation adds instrumentation for this client to aid in debugging scenarios with the write backend --- bigquery/storage/apiv1/info.go | 33 +++++++++++++++++++ bigquery/storage/managedwriter/client.go | 2 ++ .../storage/managedwriter/managed_stream.go | 11 ++++++- bigquery/storage/managedwriter/options.go | 2 +- .../storage/managedwriter/options_test.go | 21 +++++++++--- 5 files changed, 63 insertions(+), 6 deletions(-) create mode 100644 bigquery/storage/apiv1/info.go diff --git a/bigquery/storage/apiv1/info.go b/bigquery/storage/apiv1/info.go new file mode 100644 index 000000000000..7669185f650c --- /dev/null +++ b/bigquery/storage/apiv1/info.go @@ -0,0 +1,33 @@ +// Copyright 2022 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package storage + +// SetGoogleClientInfo sets the name and version of the application in +// the `x-goog-api-client` header passed on each request. Also passes any +// provided key-value pairs. Intended for use by Google-written clients. +// +// Internal use only. +func (c *BigQueryWriteClient) SetGoogleClientInfo(keyval ...string) { + c.setGoogleClientInfo(keyval...) +} + +// SetGoogleClientInfo sets the name and version of the application in +// the `x-goog-api-client` header passed on each request. Also passes any +// provided key-value pairs. Intended for use by Google-written clients. +// +// Internal use only. +func (c *BigQueryReadClient) SetGoogleClientInfo(keyval ...string) { + c.setGoogleClientInfo(keyval...) +} diff --git a/bigquery/storage/managedwriter/client.go b/bigquery/storage/managedwriter/client.go index f517a555ce42..035ffb7f7f57 100644 --- a/bigquery/storage/managedwriter/client.go +++ b/bigquery/storage/managedwriter/client.go @@ -20,6 +20,7 @@ import ( "runtime" "strings" + "cloud.google.com/go/bigquery/internal" storage "cloud.google.com/go/bigquery/storage/apiv1" "cloud.google.com/go/internal/detect" "github.com/googleapis/gax-go/v2" @@ -59,6 +60,7 @@ func NewClient(ctx context.Context, projectID string, opts ...option.ClientOptio if err != nil { return nil, err } + rawClient.SetGoogleClientInfo("gccl", internal.Version) // Handle project autodetection. projectID, err = detect.ProjectID(ctx, projectID, "", opts...) diff --git a/bigquery/storage/managedwriter/managed_stream.go b/bigquery/storage/managedwriter/managed_stream.go index 87b6d99aa105..2131d9337f15 100644 --- a/bigquery/storage/managedwriter/managed_stream.go +++ b/bigquery/storage/managedwriter/managed_stream.go @@ -21,6 +21,7 @@ import ( "io" "sync" + "cloud.google.com/go/bigquery/internal" "github.com/googleapis/gax-go/v2" "go.opencensus.io/tag" storagepb "google.golang.org/genproto/googleapis/cloud/bigquery/storage/v1" @@ -126,10 +127,18 @@ func defaultStreamSettings() *streamSettings { streamType: DefaultStream, MaxInflightRequests: 1000, MaxInflightBytes: 0, - TraceID: "", + TraceID: buildTraceID(""), } } +func buildTraceID(id string) string { + base := fmt.Sprintf("go-managedwriter:%s", internal.Version) + if id != "" { + return fmt.Sprintf("%s %s", base, id) + } + return base +} + // StreamName returns the corresponding write stream ID being managed by this writer. func (ms *ManagedStream) StreamName() string { return ms.streamSettings.streamID diff --git a/bigquery/storage/managedwriter/options.go b/bigquery/storage/managedwriter/options.go index 2e7168cc7c92..50164a846384 100644 --- a/bigquery/storage/managedwriter/options.go +++ b/bigquery/storage/managedwriter/options.go @@ -70,7 +70,7 @@ func WithMaxInflightBytes(n int) WriterOption { // This is generally for diagnostic purposes only. func WithTraceID(traceID string) WriterOption { return func(ms *ManagedStream) { - ms.streamSettings.TraceID = traceID + ms.streamSettings.TraceID = buildTraceID(traceID) } } diff --git a/bigquery/storage/managedwriter/options_test.go b/bigquery/storage/managedwriter/options_test.go index 5d81035ff29d..03551a9600d0 100644 --- a/bigquery/storage/managedwriter/options_test.go +++ b/bigquery/storage/managedwriter/options_test.go @@ -15,9 +15,11 @@ package managedwriter import ( + "fmt" "sync" "testing" + "cloud.google.com/go/bigquery/internal" "github.com/google/go-cmp/cmp" "github.com/googleapis/gax-go/v2" "google.golang.org/grpc" @@ -64,13 +66,24 @@ func TestWriterOptions(t *testing.T) { }(), }, { - desc: "WithTracePrefix", + desc: "WithTraceID", options: []WriterOption{WithTraceID("foo")}, want: func() *ManagedStream { ms := &ManagedStream{ streamSettings: defaultStreamSettings(), } - ms.streamSettings.TraceID = "foo" + ms.streamSettings.TraceID = fmt.Sprintf("go-managedwriter:%s foo", internal.Version) + return ms + }(), + }, + { + desc: "WithoutTraceID", + options: []WriterOption{}, + want: func() *ManagedStream { + ms := &ManagedStream{ + streamSettings: defaultStreamSettings(), + } + ms.streamSettings.TraceID = fmt.Sprintf("go-managedwriter:%s", internal.Version) return ms }(), }, @@ -114,7 +127,7 @@ func TestWriterOptions(t *testing.T) { options: []WriterOption{ WithType(PendingStream), WithMaxInflightBytes(5), - WithTraceID("id"), + WithTraceID("traceid"), }, want: func() *ManagedStream { ms := &ManagedStream{ @@ -122,7 +135,7 @@ func TestWriterOptions(t *testing.T) { } ms.streamSettings.MaxInflightBytes = 5 ms.streamSettings.streamType = PendingStream - ms.streamSettings.TraceID = "id" + ms.streamSettings.TraceID = fmt.Sprintf("go-managedwriter:%s traceid", internal.Version) return ms }(), },