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
googleapis · Aug 9, 2022 · 8d0d3ff · 8d0d3ff
1 parent d8b9331
commit 8d0d3ff
Show file tree

Hide file tree

Showing 5 changed files with 63 additions and 6 deletions.
diff --git a/bigquery/storage/apiv1/info.go 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
@@ -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
@@ -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
@@ -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
@@ -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,15 +127,15 @@ func TestWriterOptions(t *testing.T) {
 			options: []WriterOption{
 				WithType(PendingStream),
 				WithMaxInflightBytes(5),
-				WithTraceID("id"),
+				WithTraceID("traceid"),
 			},
 			want: func() *ManagedStream {
 				ms := &ManagedStream{
 					streamSettings: defaultStreamSettings(),
 				}
 				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
 			}(),
 		},