fix(bigquery): do not override ctx within inner span creation (#11119)

When creating child spans (mostly inside `runWithRetry` func), parent `context.Context` was being override by the child span `context.Context`, making the parent span to be missing on traces. This PR changes the `ctx` assignment inside those `runWithRetry` to always create a new `context.Context` just for that child span, similar to how it was done in other places like https://togithub.com/googleapis/google-cloud-go/blob/cc88fe1c8d834903e2b5f8d3a7cc9bbd65c70b29/bigquery/job.go#L248

Fixes #10961

Author: alvarowolfx

bigquery/inserter.go

@@ -181,9 +181,9 @@ func (u *Inserter) putMulti(ctx context.Context, src []ValueSaver) error {
 	setClientHeader(call.Header())
 	var res *bq.TableDataInsertAllResponse
 	err = runWithRetry(ctx, func() (err error) {
-		ctx = trace.StartSpan(ctx, "bigquery.tabledata.insertAll")
+		sCtx := trace.StartSpan(ctx, "bigquery.tabledata.insertAll")
 		res, err = call.Do()
-		trace.EndSpan(ctx, err)
+		trace.EndSpan(sCtx, err)
 		return err
 	})
 	if err != nil {

bigquery/model.go

@@ -86,9 +86,9 @@ func (m *Model) Metadata(ctx context.Context) (mm *ModelMetadata, err error) {
 	setClientHeader(req.Header())
 	var model *bq.Model
 	err = runWithRetry(ctx, func() (err error) {
-		ctx = trace.StartSpan(ctx, "bigquery.models.get")
+		sCtx := trace.StartSpan(ctx, "bigquery.models.get")
 		model, err = req.Do()
-		trace.EndSpan(ctx, err)
+		trace.EndSpan(sCtx, err)
 		return err
 	})
 	if err != nil {
@@ -113,9 +113,9 @@ func (m *Model) Update(ctx context.Context, mm ModelMetadataToUpdate, etag strin
 	}
 	var res *bq.Model
 	if err := runWithRetry(ctx, func() (err error) {
-		ctx = trace.StartSpan(ctx, "bigquery.models.patch")
+		sCtx := trace.StartSpan(ctx, "bigquery.models.patch")
 		res, err = call.Do()
-		trace.EndSpan(ctx, err)
+		trace.EndSpan(sCtx, err)
 		return err
 	}); err != nil {
 		return nil, err

bigquery/routine.go

@@ -97,9 +97,9 @@ func (r *Routine) Metadata(ctx context.Context) (rm *RoutineMetadata, err error)
 	setClientHeader(req.Header())
 	var routine *bq.Routine
 	err = runWithRetry(ctx, func() (err error) {
-		ctx = trace.StartSpan(ctx, "bigquery.routines.get")
+		sCtx := trace.StartSpan(ctx, "bigquery.routines.get")
 		routine, err = req.Do()
-		trace.EndSpan(ctx, err)
+		trace.EndSpan(sCtx, err)
 		return err
 	})
 	if err != nil {
@@ -131,9 +131,9 @@ func (r *Routine) Update(ctx context.Context, upd *RoutineMetadataToUpdate, etag
 	}
 	var res *bq.Routine
 	if err := runWithRetry(ctx, func() (err error) {
-		ctx = trace.StartSpan(ctx, "bigquery.routines.update")
+		sCtx := trace.StartSpan(ctx, "bigquery.routines.update")
 		res, err = call.Do()
-		trace.EndSpan(ctx, err)
+		trace.EndSpan(sCtx, err)
 		return err
 	}); err != nil {
 		return nil, err

bigquery/table.go

@@ -721,9 +721,9 @@ func (t *Table) Create(ctx context.Context, tm *TableMetadata) (err error) {
 	req := t.c.bqs.Tables.Insert(t.ProjectID, t.DatasetID, table).Context(ctx)
 	setClientHeader(req.Header())
 	return runWithRetry(ctx, func() (err error) {
-		ctx = trace.StartSpan(ctx, "bigquery.tables.insert")
+		sCtx := trace.StartSpan(ctx, "bigquery.tables.insert")
 		_, err = req.Do()
-		trace.EndSpan(ctx, err)
+		trace.EndSpan(sCtx, err)
 		return err
 	})
 }
@@ -960,9 +960,9 @@ func (t *Table) Delete(ctx context.Context) (err error) {
 	setClientHeader(call.Header())
 
 	return runWithRetry(ctx, func() (err error) {
-		ctx = trace.StartSpan(ctx, "bigquery.tables.delete")
+		sCtx := trace.StartSpan(ctx, "bigquery.tables.delete")
 		err = call.Do()
-		trace.EndSpan(ctx, err)
+		trace.EndSpan(sCtx, err)
 		return err
 	})
 }
@@ -1026,9 +1026,9 @@ func (t *Table) Update(ctx context.Context, tm TableMetadataToUpdate, etag strin
 	}
 	var res *bq.Table
 	if err := runWithRetry(ctx, func() (err error) {
-		ctx = trace.StartSpan(ctx, "bigquery.tables.patch")
+		sCtx := trace.StartSpan(ctx, "bigquery.tables.patch")
 		res, err = tpc.call.Do()
-		trace.EndSpan(ctx, err)
+		trace.EndSpan(sCtx, err)
 		return err
 	}); err != nil {
 		return nil, err

bigquery/trace_integration_test.go

@@ -94,6 +94,37 @@ func TestIntegration_OpenCensusTracing(t *testing.T) {
 			},
 			wantSpans: []string{"bigquery.tables.get", "cloud.google.com/go/bigquery.Table.Metadata"},
 		},
+		{
+			description: "dataset and table insert/update/delete",
+			callF: func(ctx context.Context) {
+				ds := client.Dataset(datasetIDs.New())
+				if err := ds.Create(ctx, nil); err != nil {
+					t.Fatalf("creating dataset %s: %v", ds.DatasetID, err)
+				}
+				defer ds.DeleteWithContents(ctx)
+				tbl := ds.Table(tableIDs.New())
+				tm := &TableMetadata{
+					Schema: schema,
+				}
+				if err := tbl.Create(ctx, tm); err != nil {
+					t.Fatalf("creating table %s: %v", tbl.TableID, err)
+				}
+				defer tbl.Delete(ctx)
+				md := TableMetadataToUpdate{
+					Description: "new description",
+				}
+				if _, err := tbl.Update(ctx, md, ""); err != nil {
+					t.Fatalf("updating table %s: %v", tbl.TableID, err)
+				}
+			},
+			wantSpans: []string{
+				"cloud.google.com/go/bigquery.Dataset.Create",
+				"bigquery.tables.insert", "cloud.google.com/go/bigquery.Table.Create",
+				"bigquery.tables.patch", "cloud.google.com/go/bigquery.Table.Update",
+				"bigquery.datasets.delete", "cloud.google.com/go/bigquery.Dataset.Delete",
+				"bigquery.tables.delete", "cloud.google.com/go/bigquery.Table.Delete",
+			},
+		},
 	} {
 		t.Run(tc.description, func(t *testing.T) {
 			exporter := &testExporter{}