Skip to content

Commit 2cd823e

Browse files
birschick-bqbirschick-bq
authored
chore(csharp): increase tracing granularity (#150)
## What's Changed - increases tracing granularity around most important steps in query execution.
1 parent d35ac0c commit 2cd823e

File tree

2 files changed

+33
-23
lines changed

2 files changed

+33
-23
lines changed

‎csharp/src/BigQueryConnection.cs‎

Lines changed: 15 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ public class BigQueryConnection : TracingConnection, ITokenProtectedResource
5353
{
5454
readonly Dictionary<string, string> properties;
5555
readonly HttpClient httpClient;
56+
const string ClassName = nameof(BigQueryConnection);
5657
const string infoDriverName = "ADBC BigQuery Driver";
5758
const string infoVendorName = "BigQuery";
5859
// Note: this needs to be set before the constructor runs
@@ -268,7 +269,7 @@ internal BigQueryClient Open(string? projectId = null)
268269

269270
Client = client;
270271
return client;
271-
});
272+
}, ClassName + "." + nameof(Open));
272273
}
273274

274275
internal void SetCredential()
@@ -342,7 +343,7 @@ internal void SetCredential()
342343
{
343344
throw new ArgumentException($"{authenticationType} is not a valid authenticationType");
344345
}
345-
});
346+
}, ClassName + "." + nameof(SetCredential));
346347
}
347348

348349
public override void SetOption(string key, string value)
@@ -357,7 +358,7 @@ public override void SetOption(string key, string value)
357358
{
358359
UpdateClientToken();
359360
}
360-
});
361+
}, ClassName + "." + nameof(SetOption));
361362
}
362363

363364
/// <summary>
@@ -505,7 +506,7 @@ public override IArrowArrayStream GetInfo(IReadOnlyList<AdbcInfoCode> codes)
505506
StandardSchemas.GetInfoSchema.Validate(dataArrays);
506507

507508
return new BigQueryInfoArrowStream(StandardSchemas.GetInfoSchema, dataArrays);
508-
});
509+
}, ClassName + "." + nameof(GetInfo));
509510
}
510511

511512
public override IArrowArrayStream GetObjects(
@@ -529,7 +530,7 @@ public override IArrowArrayStream GetObjects(
529530
{
530531
throw new AdbcException(googleEx!.Message, AdbcStatusCode.Unauthorized, ex);
531532
}
532-
});
533+
}, ClassName + "." + nameof(GetObjects));
533534
}
534535

535536
/// <summary>
@@ -578,7 +579,7 @@ internal void UpdateClientToken()
578579
{
579580
throw new AdbcException(googleEx!.Message, AdbcStatusCode.Unauthorized, ex);
580581
}
581-
});
582+
}, ClassName + "." + nameof(ExecuteQuery));
582583
}
583584

584585
internal static bool IsUnauthorizedException(Exception ex, out GoogleApiException? googleEx)
@@ -648,7 +649,7 @@ private IArrowArray[] GetCatalogs(
648649
StandardSchemas.GetObjectsSchema.Validate(dataArrays);
649650

650651
return dataArrays;
651-
});
652+
}, ClassName + "." + nameof(GetCatalogs));
652653
}
653654

654655
private StructArray GetDbSchemas(
@@ -712,7 +713,7 @@ private StructArray GetDbSchemas(
712713
length,
713714
dataArrays,
714715
nullBitmapBuffer.Build());
715-
});
716+
}, ClassName + "." + nameof(GetDbSchemas));
716717
}
717718

718719
private StructArray GetTableSchemas(
@@ -806,7 +807,7 @@ private StructArray GetTableSchemas(
806807
length,
807808
dataArrays,
808809
nullBitmapBuffer.Build());
809-
});
810+
}, ClassName + "." + nameof(GetTableSchemas));
810811
}
811812

812813
private StructArray GetColumnSchema(
@@ -922,7 +923,7 @@ private StructArray GetColumnSchema(
922923
length,
923924
dataArrays,
924925
nullBitmapBuffer.Build());
925-
});
926+
}, ClassName + "." + nameof(GetColumnSchema));
926927
}
927928

928929
private StructArray GetConstraintSchema(
@@ -994,7 +995,7 @@ private StructArray GetConstraintSchema(
994995
length,
995996
dataArrays,
996997
nullBitmapBuffer.Build());
997-
});
998+
}, ClassName + "." + nameof(GetConstraintSchema));
998999
}
9991000

10001001
private StringArray GetConstraintColumnNames(
@@ -1022,7 +1023,7 @@ private StringArray GetConstraintColumnNames(
10221023
}
10231024

10241025
return constraintColumnNamesBuilder.Build();
1025-
});
1026+
}, ClassName + "." + nameof(GetConstraintColumnNames));
10261027
}
10271028

10281029
private StructArray GetConstraintsUsage(
@@ -1078,7 +1079,7 @@ private StructArray GetConstraintsUsage(
10781079
length,
10791080
dataArrays,
10801081
nullBitmapBuffer.Build());
1081-
});
1082+
}, ClassName + "." + nameof(GetConstraintsUsage));
10821083
}
10831084

10841085
private string PatternToRegEx(string? pattern)
@@ -1196,7 +1197,7 @@ public override Schema GetTableSchema(string? catalog, string? dbSchema, string
11961197
}
11971198

11981199
return new Schema(fields, null);
1199-
});
1200+
}, ClassName + "." + nameof(GetTableSchema));
12001201
}
12011202

12021203
private Field DescToField(BigQueryRow row)

‎csharp/src/BigQueryStatement.cs‎

Lines changed: 18 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ namespace AdbcDrivers.BigQuery
5050
/// </summary>
5151
class BigQueryStatement : TracingStatement, ITokenProtectedResource, IDisposable
5252
{
53+
private const string ClassName = nameof(BigQueryStatement);
5354
readonly BigQueryConnection bigQueryConnection;
5455
readonly CancellationRegistry cancellationRegistry;
5556

@@ -131,6 +132,7 @@ private async Task<QueryResult> ExecuteQueryInternalAsync()
131132
getQueryResultsOptions.Timeout = TimeSpan.FromSeconds(seconds);
132133
activity?.AddBigQueryParameterTag(BigQueryParameters.GetQueryResultsOptionsTimeout, seconds);
133134
}
135+
activity?.AddBigQueryParameterTag(BigQueryParameters.ClientTimeout, Client.Service.HttpClient.Timeout.Seconds);
134136

135137
using JobCancellationContext cancellationContext = new JobCancellationContext(cancellationRegistry, job);
136138

@@ -140,9 +142,12 @@ private async Task<QueryResult> ExecuteQueryInternalAsync()
140142
{
141143
return await ExecuteCancellableJobAsync(cancellationContext, activity, async (context) =>
142144
{
143-
// if the authentication token was reset, then we need a new job with the latest token
144-
context.Job = await Client.GetJobAsync(jobReference, cancellationToken: context.CancellationToken).ConfigureAwait(false);
145-
return await context.Job.GetQueryResultsAsync(getQueryResultsOptions, cancellationToken: context.CancellationToken).ConfigureAwait(false);
145+
return await this.TraceActivityAsync(async activity =>
146+
{
147+
// if the authentication token was reset, then we need a new job with the latest token
148+
context.Job = await Client.GetJobAsync(jobReference, cancellationToken: context.CancellationToken).ConfigureAwait(false);
149+
return await context.Job.GetQueryResultsAsync(getQueryResultsOptions, cancellationToken: context.CancellationToken).ConfigureAwait(false);
150+
}, ClassName + "." + nameof(ExecuteQueryInternalAsync) + "." + nameof(BigQueryJob.GetQueryResultsAsync));
146151
}).ConfigureAwait(false);
147152
};
148153

@@ -235,8 +240,11 @@ private async Task<QueryResult> ExecuteQueryInternalAsync()
235240
{
236241
return await ExecuteCancellableJobAsync(cancellationContext, activity, async (context) =>
237242
{
238-
// Cancelling this step may leave the server with unread streams.
239-
return await GetArrowReaders(clientMgr, table, results.TableReference.ProjectId, maxStreamCount, activity, context.CancellationToken).ConfigureAwait(false);
243+
return await this.TraceActivityAsync(async activity =>
244+
{
245+
// Cancelling this step may leave the server with unread streams.
246+
return await GetArrowReaders(clientMgr, table, results.TableReference.ProjectId, maxStreamCount, activity, context.CancellationToken).ConfigureAwait(false);
247+
}, ClassName + "." + nameof(ExecuteQueryInternalAsync) + "." + nameof(GetArrowReaders));
240248
}).ConfigureAwait(false);
241249
};
242250
IEnumerable<IArrowReader> readers = await ExecuteWithRetriesAsync(getArrowReadersFunc, activity, cancellationContext.CancellationToken).ConfigureAwait(false);
@@ -245,7 +253,7 @@ private async Task<QueryResult> ExecuteQueryInternalAsync()
245253
IArrowArrayStream stream = new MultiArrowReader(this, TranslateSchema(results.Schema), readers, new CancellationContext(cancellationRegistry));
246254
activity?.AddTag(SemanticConventions.Db.Response.ReturnedRows, totalRows);
247255
return new QueryResult(totalRows, stream);
248-
});
256+
}, ClassName + "." + nameof(ExecuteQueryInternalAsync));
249257
}
250258

251259
private Task<QueryResult> ExecuteMetadataCommandQuery(Activity? activity)
@@ -706,7 +714,7 @@ public override void Cancel()
706714
this.TraceActivity(_ =>
707715
{
708716
this.cancellationRegistry.CancelAll();
709-
});
717+
}, ClassName + "." + nameof(Cancel));
710718
}
711719

712720
public override void Dispose()
@@ -756,7 +764,7 @@ private async Task<UpdateResult> ExecuteUpdateInternalAsync()
756764

757765
activity?.AddTag(SemanticConventions.Db.Response.ReturnedRows, updatedRows);
758766
return new UpdateResult(updatedRows);
759-
});
767+
}, ClassName + "." + nameof(ExecuteUpdateInternalAsync));
760768
}
761769

762770
private Schema TranslateSchema(TableSchema schema)
@@ -1166,6 +1174,7 @@ public void Dispose()
11661174

11671175
private class MultiArrowReader : TracingReader
11681176
{
1177+
private const string ClassName = BigQueryStatement.ClassName + ".MultiArrowReader";
11691178
private static readonly string s_assemblyName = BigQueryUtils.GetAssemblyName(typeof(BigQueryStatement));
11701179
private static readonly string s_assemblyVersion = BigQueryUtils.GetAssemblyVersion(typeof(BigQueryStatement));
11711180

@@ -1222,7 +1231,7 @@ public MultiArrowReader(BigQueryStatement statement, Schema schema, IEnumerable<
12221231

12231232
this.reader = null;
12241233
}
1225-
});
1234+
}, ClassName + "." + nameof(ReadNextRecordBatchAsync));
12261235
}
12271236

12281237
protected override void Dispose(bool disposing)

0 commit comments

Comments
 (0)