Skip to content

Commit 3790802

Browse files
westarlewestarle
committed
feat(bigquery): add T4 tracing telemetry to Dataset and Model operations
1 parent 14b9656 commit 3790802

File tree

7 files changed

+440
-59
lines changed

7 files changed

+440
-59
lines changed

‎bigquery/bigquery.go‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ import (
3333
bq "google.golang.org/api/bigquery/v2"
3434
"google.golang.org/api/googleapi"
3535
"google.golang.org/api/option"
36+
"google.golang.org/api/option/internaloption"
3637
)
3738

3839
const (
@@ -83,6 +84,18 @@ func NewClient(ctx context.Context, projectID string, opts ...option.ClientOptio
8384
option.WithScopes(Scope),
8485
option.WithUserAgent(fmt.Sprintf("%s/%s", userAgentPrefix, internal.Version)),
8586
}
87+
88+
if gax.IsFeatureEnabled("TRACING") {
89+
o = append(o, internaloption.WithTelemetryAttributes(map[string]string{
90+
"gcp.client.repo": "googleapis/google-cloud-go",
91+
"gcp.client.version": internal.Version,
92+
"gcp.client.service": "bigquery.googleapis.com",
93+
"gcp.client.artifact": "cloud.google.com/go/bigquery",
94+
"gcp.client.language": "go",
95+
"url.domain": "bigquery.googleapis.com",
96+
}))
97+
}
98+
8699
o = append(o, opts...)
87100
bqs, err := bq.NewService(ctx, o...)
88101
if err != nil {

‎bigquery/dataset.go‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -231,6 +231,7 @@ func (d *Dataset) CreateWithOptions(ctx context.Context, md *DatasetMetadata, op
231231
if ds.Location == "" {
232232
ds.Location = d.c.Location
233233
}
234+
ctx = setProjectItemTraceMetadata(ctx, d.ProjectID, "datasets")
234235
call := d.c.bqs.Datasets.Insert(d.ProjectID, ds).Context(ctx)
235236
setClientHeader(call.Header())
236237
if cOpt.accessPolicyVersion != nil {
@@ -307,6 +308,7 @@ func (d *Dataset) deleteInternal(ctx context.Context, deleteContents bool) (err
307308
ctx = trace.StartSpan(ctx, "cloud.google.com/go/bigquery.Dataset.Delete")
308309
defer func() { trace.EndSpan(ctx, err) }()
309310

311+
ctx = setDatasetTraceMetadata(ctx, d.ProjectID, d.DatasetID)
310312
call := d.c.bqs.Datasets.Delete(d.ProjectID, d.DatasetID).Context(ctx).DeleteContents(deleteContents)
311313
setClientHeader(call.Header())
312314
return runWithRetry(ctx, func() (err error) {
@@ -333,6 +335,7 @@ func (d *Dataset) MetadataWithOptions(ctx context.Context, opts ...DatasetOption
333335
o(cOpt)
334336
}
335337

338+
ctx = setDatasetTraceMetadata(ctx, d.ProjectID, d.DatasetID)
336339
call := d.c.bqs.Datasets.Get(d.ProjectID, d.DatasetID).Context(ctx)
337340
setClientHeader(call.Header())
338341
if cOpt.accessPolicyVersion != nil {
@@ -504,6 +507,7 @@ func (d *Dataset) UpdateWithOptions(ctx context.Context, dm DatasetMetadataToUpd
504507
return nil, err
505508
}
506509

510+
ctx = setDatasetTraceMetadata(ctx, d.ProjectID, d.DatasetID)
507511
call := d.c.bqs.Datasets.Patch(d.ProjectID, d.DatasetID, ds).Context(ctx)
508512
setClientHeader(call.Header())
509513
if etag != "" {
@@ -652,6 +656,7 @@ func (it *TableIterator) PageInfo() *iterator.PageInfo { return it.pageInfo }
652656

653657
// listTables exists to aid testing.
654658
var listTables = func(it *TableIterator, pageSize int, pageToken string) (*bq.TableList, error) {
659+
it.ctx = setDatasetItemTraceMetadata(it.ctx, it.dataset.ProjectID, it.dataset.DatasetID, "tables")
655660
call := it.dataset.c.bqs.Tables.List(it.dataset.ProjectID, it.dataset.DatasetID).
656661
PageToken(pageToken).
657662
Context(it.ctx)
@@ -739,6 +744,7 @@ func (it *ModelIterator) PageInfo() *iterator.PageInfo { return it.pageInfo }
739744

740745
// listTables exists to aid testing.
741746
var listModels = func(it *ModelIterator, pageSize int, pageToken string) (*bq.ListModelsResponse, error) {
747+
it.ctx = setDatasetItemTraceMetadata(it.ctx, it.dataset.ProjectID, it.dataset.DatasetID, "models")
742748
call := it.dataset.c.bqs.Models.List(it.dataset.ProjectID, it.dataset.DatasetID).
743749
PageToken(pageToken).
744750
Context(it.ctx)
@@ -828,6 +834,7 @@ func (it *RoutineIterator) PageInfo() *iterator.PageInfo { return it.pageInfo }
828834

829835
// listRoutines exists to aid testing.
830836
var listRoutines = func(it *RoutineIterator, pageSize int, pageToken string) (*bq.ListRoutinesResponse, error) {
837+
it.ctx = setDatasetItemTraceMetadata(it.ctx, it.dataset.ProjectID, it.dataset.DatasetID, "routines")
831838
call := it.dataset.c.bqs.Routines.List(it.dataset.ProjectID, it.dataset.DatasetID).
832839
PageToken(pageToken).
833840
Context(it.ctx)
@@ -930,6 +937,7 @@ func (it *DatasetIterator) Next() (*Dataset, error) {
930937

931938
// for testing
932939
var listDatasets = func(it *DatasetIterator, pageSize int, pageToken string) (*bq.DatasetList, error) {
940+
it.ctx = setProjectItemTraceMetadata(it.ctx, it.ProjectID, "datasets")
933941
call := it.c.bqs.Datasets.List(it.ProjectID).
934942
Context(it.ctx).
935943
PageToken(pageToken).

‎bigquery/model.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,7 @@ func (m *Model) Metadata(ctx context.Context) (mm *ModelMetadata, err error) {
8282
ctx = trace.StartSpan(ctx, "cloud.google.com/go/bigquery.Model.Metadata")
8383
defer func() { trace.EndSpan(ctx, err) }()
8484

85+
ctx = setModelTraceMetadata(ctx, m.ProjectID, m.DatasetID, m.ModelID)
8586
req := m.c.bqs.Models.Get(m.ProjectID, m.DatasetID, m.ModelID).Context(ctx)
8687
setClientHeader(req.Header())
8788
var model *bq.Model
@@ -106,6 +107,7 @@ func (m *Model) Update(ctx context.Context, mm ModelMetadataToUpdate, etag strin
106107
if err != nil {
107108
return nil, err
108109
}
110+
ctx = setModelTraceMetadata(ctx, m.ProjectID, m.DatasetID, m.ModelID)
109111
call := m.c.bqs.Models.Patch(m.ProjectID, m.DatasetID, m.ModelID, bqm).Context(ctx)
110112
setClientHeader(call.Header())
111113
if etag != "" {
@@ -128,6 +130,7 @@ func (m *Model) Delete(ctx context.Context) (err error) {
128130
ctx = trace.StartSpan(ctx, "cloud.google.com/go/bigquery.Model.Delete")
129131
defer func() { trace.EndSpan(ctx, err) }()
130132

133+
ctx = setModelTraceMetadata(ctx, m.ProjectID, m.DatasetID, m.ModelID)
131134
req := m.c.bqs.Models.Delete(m.ProjectID, m.DatasetID, m.ModelID).Context(ctx)
132135
setClientHeader(req.Header())
133136
return req.Do()

‎bigquery/trace.go‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
// Copyright 2026 Google LLC
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 bigquery
16+
17+
import (
18+
"context"
19+
"fmt"
20+
"strings"
21+
22+
"github.com/googleapis/gax-go/v2"
23+
"github.com/googleapis/gax-go/v2/callctx"
24+
)
25+
26+
// setTraceMetadata sets the trace metadata into the context if tracing is enabled.
27+
func setTraceMetadata(ctx context.Context, resourceName, urlTemplate string) context.Context {
28+
if gax.IsFeatureEnabled("TRACING") {
29+
return callctx.WithTelemetryContext(ctx, "resource_name", resourceName, "url_template", urlTemplate)
30+
}
31+
return ctx
32+
}
33+
34+
// datasetResourceName constructs the standard resource name for a dataset.
35+
// E.g., "//bigquery.googleapis.com/projects/{project}/datasets/{dataset}"
36+
func datasetResourceName(projectID, datasetID string) string {
37+
return fmt.Sprintf("//bigquery.googleapis.com/projects/%s/datasets/%s", projectID, datasetID)
38+
}
39+
40+
// modelResourceName constructs the standard resource name for a model.
41+
// E.g., "//bigquery.googleapis.com/projects/{project}/datasets/{dataset}/models/{model}"
42+
func modelResourceName(projectID, datasetID, modelID string) string {
43+
return fmt.Sprintf("//bigquery.googleapis.com/projects/%s/datasets/%s/models/%s", projectID, datasetID, modelID)
44+
}
45+
46+
func fullyQualifiedDatasetResourceName(projectID, datasetID string) string {
47+
if strings.HasPrefix(datasetID, "projects/") {
48+
// Handle fully qualified names
49+
return fmt.Sprintf("//bigquery.googleapis.com/%s", datasetID)
50+
}
51+
return datasetResourceName(projectID, datasetID)
52+
}
53+
54+
func setProjectItemTraceMetadata(ctx context.Context, projectID, childType string) context.Context {
55+
return setTraceMetadata(ctx,
56+
fmt.Sprintf("//bigquery.googleapis.com/projects/%s", projectID),
57+
fmt.Sprintf("/bigquery/v2/projects/{projectId}/%s", childType))
58+
}
59+
60+
func setDatasetTraceMetadata(ctx context.Context, projectID, datasetID string) context.Context {
61+
return setTraceMetadata(ctx,
62+
fullyQualifiedDatasetResourceName(projectID, datasetID),
63+
"/bigquery/v2/projects/{projectId}/datasets/{datasetId}")
64+
}
65+
66+
func setDatasetItemTraceMetadata(ctx context.Context, projectID, datasetID, childType string) context.Context {
67+
return setTraceMetadata(ctx,
68+
fullyQualifiedDatasetResourceName(projectID, datasetID),
69+
fmt.Sprintf("/bigquery/v2/projects/{projectId}/datasets/{datasetId}/%s", childType))
70+
}
71+
72+
func setModelTraceMetadata(ctx context.Context, projectID, datasetID, modelID string) context.Context {
73+
return setTraceMetadata(ctx,
74+
modelResourceName(projectID, datasetID, modelID),
75+
"/bigquery/v2/projects/{projectId}/datasets/{datasetId}/models/{modelId}")
76+
}

0 commit comments

Comments
 (0)