Skip to content

Commit 6defca9

Browse files
westarlewestarle
committed
feat(bigquery): add T4 tracing telemetry to Table and Routine operations
1 parent 48c70e2 commit 6defca9

4 files changed

Lines changed: 165 additions & 34 deletions

File tree

‎bigquery/routine.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,7 @@ func (r *Routine) Create(ctx context.Context, rm *RoutineMetadata) (err error) {
8181
DatasetId: r.DatasetID,
8282
RoutineId: r.RoutineID,
8383
}
84+
ctx = setDatasetItemTraceMetadata(ctx, r.ProjectID, r.DatasetID, "routines")
8485
req := r.c.bqs.Routines.Insert(r.ProjectID, r.DatasetID, routine).Context(ctx)
8586
setClientHeader(req.Header())
8687
_, err = req.Do()
@@ -92,6 +93,7 @@ func (r *Routine) Metadata(ctx context.Context) (rm *RoutineMetadata, err error)
9293
ctx = trace.StartSpan(ctx, "cloud.google.com/go/bigquery.Routine.Metadata")
9394
defer func() { trace.EndSpan(ctx, err) }()
9495

96+
ctx = setRoutineTraceMetadata(ctx, r.ProjectID, r.DatasetID, r.RoutineID)
9597
req := r.c.bqs.Routines.Get(r.ProjectID, r.DatasetID, r.RoutineID).Context(ctx)
9698
setClientHeader(req.Header())
9799
var routine *bq.Routine
@@ -123,6 +125,7 @@ func (r *Routine) Update(ctx context.Context, upd *RoutineMetadataToUpdate, etag
123125
RoutineId: r.RoutineID,
124126
}
125127

128+
ctx = setRoutineTraceMetadata(ctx, r.ProjectID, r.DatasetID, r.RoutineID)
126129
call := r.c.bqs.Routines.Update(r.ProjectID, r.DatasetID, r.RoutineID, bqr).Context(ctx)
127130
setClientHeader(call.Header())
128131
if etag != "" {
@@ -145,6 +148,7 @@ func (r *Routine) Delete(ctx context.Context) (err error) {
145148
ctx = trace.StartSpan(ctx, "cloud.google.com/go/bigquery.Model.Delete")
146149
defer func() { trace.EndSpan(ctx, err) }()
147150

151+
ctx = setRoutineTraceMetadata(ctx, r.ProjectID, r.DatasetID, r.RoutineID)
148152
req := r.c.bqs.Routines.Delete(r.ProjectID, r.DatasetID, r.RoutineID).Context(ctx)
149153
setClientHeader(req.Header())
150154
return req.Do()

‎bigquery/table.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -796,6 +796,7 @@ func (t *Table) Create(ctx context.Context, tm *TableMetadata) (err error) {
796796
TableId: t.TableID,
797797
}
798798

799+
ctx = setDatasetItemTraceMetadata(ctx, t.ProjectID, t.DatasetID, "tables")
799800
req := t.c.bqs.Tables.Insert(t.ProjectID, t.DatasetID, table).Context(ctx)
800801
setClientHeader(req.Header())
801802
return runWithRetry(ctx, func() (err error) {
@@ -935,6 +936,7 @@ func (t *Table) Metadata(ctx context.Context, opts ...TableMetadataOption) (md *
935936
ctx = trace.StartSpan(ctx, "cloud.google.com/go/bigquery.Table.Metadata")
936937
defer func() { trace.EndSpan(ctx, err) }()
937938

939+
ctx = setTableTraceMetadata(ctx, t.ProjectID, t.DatasetID, t.TableID)
938940
tgc := &tableGetCall{
939941
call: t.c.bqs.Tables.Get(t.ProjectID, t.DatasetID, t.TableID).Context(ctx),
940942
}
@@ -1028,6 +1030,7 @@ func (t *Table) Delete(ctx context.Context) (err error) {
10281030
ctx = trace.StartSpan(ctx, "cloud.google.com/go/bigquery.Table.Delete")
10291031
defer func() { trace.EndSpan(ctx, err) }()
10301032

1033+
ctx = setTableTraceMetadata(ctx, t.ProjectID, t.DatasetID, t.TableID)
10311034
call := t.c.bqs.Tables.Delete(t.ProjectID, t.DatasetID, t.TableID).Context(ctx)
10321035
setClientHeader(call.Header())
10331036

@@ -1084,6 +1087,7 @@ func (t *Table) Update(ctx context.Context, tm TableMetadataToUpdate, etag strin
10841087
return nil, err
10851088
}
10861089

1090+
ctx = setTableTraceMetadata(ctx, t.ProjectID, t.DatasetID, t.TableID)
10871091
tpc := &tablePatchCall{
10881092
call: t.c.bqs.Tables.Patch(t.ProjectID, t.DatasetID, t.TableID, bqt).Context(ctx),
10891093
}

‎bigquery/trace.go‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,3 +83,33 @@ func setModelTraceMetadata(ctx context.Context, projectID, datasetID, modelID st
8383
modelResourceName(projectID, datasetID, modelID),
8484
"/bigquery/v2/projects/{projectId}/datasets/{datasetId}/models/{modelId}")
8585
}
86+
87+
// tableResourceName constructs the standard resource name for a table.
88+
// E.g., "//bigquery.googleapis.com/projects/{project}/datasets/{dataset}/tables/{table}"
89+
func tableResourceName(projectID, datasetID, tableID string) string {
90+
return fmt.Sprintf("//bigquery.googleapis.com/projects/%s/datasets/%s/tables/%s", projectID, datasetID, tableID)
91+
}
92+
93+
// routineResourceName constructs the standard resource name for a routine.
94+
// E.g., "//bigquery.googleapis.com/projects/{project}/datasets/{dataset}/routines/{routine}"
95+
func routineResourceName(projectID, datasetID, routineID string) string {
96+
return fmt.Sprintf("//bigquery.googleapis.com/projects/%s/datasets/%s/routines/%s", projectID, datasetID, routineID)
97+
}
98+
99+
func setTableTraceMetadata(ctx context.Context, projectID, datasetID, tableID string) context.Context {
100+
if !gax.IsFeatureEnabled("TRACING") {
101+
return ctx
102+
}
103+
return setTraceMetadata(ctx,
104+
tableResourceName(projectID, datasetID, tableID),
105+
"/bigquery/v2/projects/{projectId}/datasets/{datasetId}/tables/{tableId}")
106+
}
107+
108+
func setRoutineTraceMetadata(ctx context.Context, projectID, datasetID, routineID string) context.Context {
109+
if !gax.IsFeatureEnabled("TRACING") {
110+
return ctx
111+
}
112+
return setTraceMetadata(ctx,
113+
routineResourceName(projectID, datasetID, routineID),
114+
"/bigquery/v2/projects/{projectId}/datasets/{datasetId}/routines/{routineId}")
115+
}

‎bigquery/trace_test.go‎

Lines changed: 127 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -173,6 +173,118 @@ func TestTracingTelemetryAttributes(t *testing.T) {
173173
wantURLTemplate: "/bigquery/v2/projects/{projectId}/datasets/{datasetId}/models/{modelId}",
174174
wantAttempts: 1,
175175
},
176+
{
177+
name: "Table_Create",
178+
callFunc: func(ctx context.Context, client *Client) {
179+
_ = client.Dataset("test-dataset").Table("test-table").Create(ctx, &TableMetadata{})
180+
},
181+
mockResponse: `{"tableReference": {"projectId": "test-project", "datasetId": "test-dataset", "tableId": "test-table"}}`,
182+
mockStatusCodes: []int{http.StatusOK},
183+
wantResourceName: "//bigquery.googleapis.com/projects/test-project/datasets/test-dataset",
184+
wantURLTemplate: "/bigquery/v2/projects/{projectId}/datasets/{datasetId}/tables",
185+
wantAttempts: 1,
186+
},
187+
{
188+
name: "Table_Metadata",
189+
callFunc: func(ctx context.Context, client *Client) {
190+
_, _ = client.Dataset("test-dataset").Table("test-table").Metadata(ctx)
191+
},
192+
mockResponse: `{"tableReference": {"projectId": "test-project", "datasetId": "test-dataset", "tableId": "test-table"}}`,
193+
mockStatusCodes: []int{http.StatusOK},
194+
wantResourceName: "//bigquery.googleapis.com/projects/test-project/datasets/test-dataset/tables/test-table",
195+
wantURLTemplate: "/bigquery/v2/projects/{projectId}/datasets/{datasetId}/tables/{tableId}",
196+
wantAttempts: 1,
197+
},
198+
{
199+
name: "Table_Update",
200+
callFunc: func(ctx context.Context, client *Client) {
201+
_, _ = client.Dataset("test-dataset").Table("test-table").Update(ctx, TableMetadataToUpdate{}, "")
202+
},
203+
mockResponse: `{"tableReference": {"projectId": "test-project", "datasetId": "test-dataset", "tableId": "test-table"}}`,
204+
mockStatusCodes: []int{http.StatusOK},
205+
wantResourceName: "//bigquery.googleapis.com/projects/test-project/datasets/test-dataset/tables/test-table",
206+
wantURLTemplate: "/bigquery/v2/projects/{projectId}/datasets/{datasetId}/tables/{tableId}",
207+
wantAttempts: 1,
208+
},
209+
{
210+
name: "Table_Delete",
211+
callFunc: func(ctx context.Context, client *Client) {
212+
_ = client.Dataset("test-dataset").Table("test-table").Delete(ctx)
213+
},
214+
mockResponse: `{}`,
215+
mockStatusCodes: []int{http.StatusOK},
216+
wantResourceName: "//bigquery.googleapis.com/projects/test-project/datasets/test-dataset/tables/test-table",
217+
wantURLTemplate: "/bigquery/v2/projects/{projectId}/datasets/{datasetId}/tables/{tableId}",
218+
wantAttempts: 1,
219+
},
220+
{
221+
name: "Dataset_Tables",
222+
callFunc: func(ctx context.Context, client *Client) {
223+
it := client.Dataset("test-dataset").Tables(ctx)
224+
_, _ = it.Next()
225+
},
226+
mockResponse: `{"tables": [{"tableReference": {"projectId": "test-project", "datasetId": "test-dataset", "tableId": "test-table"}}]}`,
227+
mockStatusCodes: []int{http.StatusOK},
228+
wantResourceName: "//bigquery.googleapis.com/projects/test-project/datasets/test-dataset",
229+
wantURLTemplate: "/bigquery/v2/projects/{projectId}/datasets/{datasetId}/tables",
230+
wantAttempts: 1,
231+
},
232+
{
233+
name: "Routine_Create",
234+
callFunc: func(ctx context.Context, client *Client) {
235+
_ = client.Dataset("test-dataset").Routine("test-routine").Create(ctx, &RoutineMetadata{})
236+
},
237+
mockResponse: `{"routineReference": {"projectId": "test-project", "datasetId": "test-dataset", "routineId": "test-routine"}}`,
238+
mockStatusCodes: []int{http.StatusOK},
239+
wantResourceName: "//bigquery.googleapis.com/projects/test-project/datasets/test-dataset",
240+
wantURLTemplate: "/bigquery/v2/projects/{projectId}/datasets/{datasetId}/routines",
241+
wantAttempts: 1,
242+
},
243+
{
244+
name: "Routine_Metadata",
245+
callFunc: func(ctx context.Context, client *Client) {
246+
_, _ = client.Dataset("test-dataset").Routine("test-routine").Metadata(ctx)
247+
},
248+
mockResponse: `{"routineReference": {"projectId": "test-project", "datasetId": "test-dataset", "routineId": "test-routine"}}`,
249+
mockStatusCodes: []int{http.StatusOK},
250+
wantResourceName: "//bigquery.googleapis.com/projects/test-project/datasets/test-dataset/routines/test-routine",
251+
wantURLTemplate: "/bigquery/v2/projects/{projectId}/datasets/{datasetId}/routines/{routineId}",
252+
wantAttempts: 1,
253+
},
254+
{
255+
name: "Routine_Update",
256+
callFunc: func(ctx context.Context, client *Client) {
257+
_, _ = client.Dataset("test-dataset").Routine("test-routine").Update(ctx, &RoutineMetadataToUpdate{}, "")
258+
},
259+
mockResponse: `{"routineReference": {"projectId": "test-project", "datasetId": "test-dataset", "routineId": "test-routine"}}`,
260+
mockStatusCodes: []int{http.StatusOK},
261+
wantResourceName: "//bigquery.googleapis.com/projects/test-project/datasets/test-dataset/routines/test-routine",
262+
wantURLTemplate: "/bigquery/v2/projects/{projectId}/datasets/{datasetId}/routines/{routineId}",
263+
wantAttempts: 1,
264+
},
265+
{
266+
name: "Routine_Delete",
267+
callFunc: func(ctx context.Context, client *Client) {
268+
_ = client.Dataset("test-dataset").Routine("test-routine").Delete(ctx)
269+
},
270+
mockResponse: `{}`,
271+
mockStatusCodes: []int{http.StatusOK},
272+
wantResourceName: "//bigquery.googleapis.com/projects/test-project/datasets/test-dataset/routines/test-routine",
273+
wantURLTemplate: "/bigquery/v2/projects/{projectId}/datasets/{datasetId}/routines/{routineId}",
274+
wantAttempts: 1,
275+
},
276+
{
277+
name: "Dataset_Routines",
278+
callFunc: func(ctx context.Context, client *Client) {
279+
it := client.Dataset("test-dataset").Routines(ctx)
280+
_, _ = it.Next()
281+
},
282+
mockResponse: `{"routines": [{"routineReference": {"projectId": "test-project", "datasetId": "test-dataset", "routineId": "test-routine"}}]}`,
283+
mockStatusCodes: []int{http.StatusOK},
284+
wantResourceName: "//bigquery.googleapis.com/projects/test-project/datasets/test-dataset",
285+
wantURLTemplate: "/bigquery/v2/projects/{projectId}/datasets/{datasetId}/routines",
286+
wantAttempts: 1,
287+
},
176288
{
177289
name: "Retry_Dataset_Metadata",
178290
callFunc: func(ctx context.Context, client *Client) {
@@ -232,44 +344,25 @@ func TestTracingTelemetryAttributes(t *testing.T) {
232344
if strings.Contains(span.Name, tt.wantURLTemplate) {
233345
networkSpans++
234346

235-
foundRes := false
236-
foundURL := false
237-
foundArtifact := false
238-
foundLanguage := false
239-
foundDomain := false
347+
expectedAttributes := map[attribute.Key]string{
348+
"gcp.resource.destination.id": tt.wantResourceName,
349+
"url.template": tt.wantURLTemplate,
350+
"gcp.client.artifact": "cloud.google.com/go/bigquery",
351+
"gcp.client.language": "go",
352+
"url.domain": "bigquery.googleapis.com",
353+
}
240354

355+
actualAttributes := make(map[attribute.Key]string, len(span.Attributes))
241356
for _, attr := range span.Attributes {
242-
if attr.Key == attribute.Key("gcp.resource.destination.id") && attr.Value.AsString() == tt.wantResourceName {
243-
foundRes = true
244-
}
245-
if attr.Key == attribute.Key("url.template") && attr.Value.AsString() == tt.wantURLTemplate {
246-
foundURL = true
247-
}
248-
if attr.Key == attribute.Key("gcp.client.artifact") && attr.Value.AsString() == "cloud.google.com/go/bigquery" {
249-
foundArtifact = true
250-
}
251-
if attr.Key == attribute.Key("gcp.client.language") && attr.Value.AsString() == "go" {
252-
foundLanguage = true
253-
}
254-
if attr.Key == attribute.Key("url.domain") && attr.Value.AsString() == "bigquery.googleapis.com" {
255-
foundDomain = true
256-
}
357+
actualAttributes[attr.Key] = attr.Value.AsString()
257358
}
258359

259-
if !foundRes {
260-
t.Errorf("missing gcp.resource.destination.id attribute on network span attempt")
261-
}
262-
if !foundURL {
263-
t.Errorf("missing url.template attribute on network span attempt")
264-
}
265-
if !foundArtifact {
266-
t.Errorf("missing gcp.client.artifact attribute on network span attempt")
267-
}
268-
if !foundLanguage {
269-
t.Errorf("missing gcp.client.language attribute on network span attempt")
270-
}
271-
if !foundDomain {
272-
t.Errorf("missing url.domain attribute on network span attempt")
360+
for key, wantValue := range expectedAttributes {
361+
if gotValue, ok := actualAttributes[key]; !ok {
362+
t.Errorf("attribute %q not found on span", key)
363+
} else if gotValue != wantValue {
364+
t.Errorf("attribute %q: got %q, want %q", key, gotValue, wantValue)
365+
}
273366
}
274367
}
275368
}

0 commit comments

Comments
 (0)
⚡