Skip to content

Commit ebb8f6c

Browse files
author
Jade Wang
committed
fix(csharp): address PR review feedback - test isolation, assertions, resource cleanup
- Restore AsyncLocal for ExporterOverride to prevent parallel test interference - Add missing IsInternalCall assertion in InternalCallTests - Replace silent-pass with Skip.If when no telemetry captured in baseline tests - Use await instead of .Result for async calls in MetadataOperationTests - Add TimestampMillis to metadata operation telemetry Context - Cache Process.GetCurrentProcess() call in BuildSystemConfiguration - Move reader disposal to finally blocks in ChunkMetricsReaderTests Co-authored-by: Isaac
1 parent 662b363 commit ebb8f6c

File tree

6 files changed

+79
-62
lines changed

6 files changed

+79
-62
lines changed

‎csharp/src/DatabricksConnection.cs‎

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -527,6 +527,10 @@ public override IArrowArrayStream GetObjects(
527527
{
528528
WorkspaceId = telemetryContext.WorkspaceId,
529529
FrontendLogEventId = Guid.NewGuid().ToString(),
530+
Context = new Telemetry.Models.FrontendLogContext
531+
{
532+
TimestampMillis = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds(),
533+
},
530534
Entry = new Telemetry.Models.FrontendLogEntry
531535
{
532536
SqlDriverLog = telemetryLog
@@ -639,6 +643,10 @@ public override IArrowArrayStream GetTableTypes()
639643
{
640644
WorkspaceId = telemetryContext.WorkspaceId,
641645
FrontendLogEventId = Guid.NewGuid().ToString(),
646+
Context = new Telemetry.Models.FrontendLogContext
647+
{
648+
TimestampMillis = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds(),
649+
},
642650
Entry = new Telemetry.Models.FrontendLogEntry
643651
{
644652
SqlDriverLog = telemetryLog
@@ -971,6 +979,7 @@ private void InitializeTelemetry(Activity? activity = null)
971979
private Telemetry.Proto.DriverSystemConfiguration BuildSystemConfiguration()
972980
{
973981
var osVersion = System.Environment.OSVersion;
982+
var processName = System.Diagnostics.Process.GetCurrentProcess().ProcessName;
974983
return new Telemetry.Proto.DriverSystemConfiguration
975984
{
976985
DriverVersion = s_assemblyVersion,
@@ -983,16 +992,16 @@ private Telemetry.Proto.DriverSystemConfiguration BuildSystemConfiguration()
983992
RuntimeVendor = "Microsoft",
984993
LocaleName = System.Globalization.CultureInfo.CurrentCulture.Name,
985994
CharSetEncoding = System.Text.Encoding.Default.WebName,
986-
ProcessName = System.Diagnostics.Process.GetCurrentProcess().ProcessName,
987-
ClientAppName = GetClientAppName()
995+
ProcessName = processName,
996+
ClientAppName = GetClientAppName(processName)
988997
};
989998
}
990999

991-
private string GetClientAppName()
1000+
private string GetClientAppName(string processName)
9921001
{
9931002
// Check connection property first, fall back to process name
9941003
Properties.TryGetValue("adbc.databricks.client_app_name", out string? appName);
995-
return appName ?? Process.GetCurrentProcess().ProcessName;
1004+
return appName ?? processName;
9961005
}
9971006

9981007
private Telemetry.Proto.DriverConnectionParameters BuildDriverConnectionParams(bool isAuthenticated)

‎csharp/src/Telemetry/TelemetryClientManager.cs‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
*/
1616

1717
using System.Collections.Generic;
18+
using System.Threading;
1819
using System.Threading.Tasks;
1920

2021
namespace AdbcDrivers.Databricks.Telemetry
@@ -46,12 +47,19 @@ internal sealed class TelemetryClientManager
4647
private readonly Dictionary<string, TelemetryClientHolder> _clients = new Dictionary<string, TelemetryClientHolder>();
4748
private readonly object _lock = new object();
4849

50+
private static readonly AsyncLocal<ITelemetryExporter?> s_exporterOverride = new AsyncLocal<ITelemetryExporter?>();
51+
4952
/// <summary>
5053
/// Optional exporter override for testing. When set, newly created TelemetryClients
5154
/// use this exporter instead of the default DatabricksTelemetryExporter pipeline.
55+
/// Uses AsyncLocal to prevent interference across parallel xUnit test classes.
5256
/// Must be set before connections are opened and cleared after tests complete.
5357
/// </summary>
54-
internal static ITelemetryExporter? ExporterOverride { get; set; }
58+
internal static ITelemetryExporter? ExporterOverride
59+
{
60+
get => s_exporterOverride.Value;
61+
set => s_exporterOverride.Value = value;
62+
}
5563

5664
/// <summary>
5765
/// Internal constructor. Public API uses GetInstance() for the singleton.

‎csharp/test/E2E/Telemetry/ChunkMetricsReaderTests.cs‎

Lines changed: 16 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
using AdbcDrivers.Databricks.Reader.CloudFetch;
2222
using Apache.Arrow.Adbc;
2323
using Apache.Arrow.Adbc.Tests;
24+
using Apache.Arrow.Ipc;
2425
using Xunit;
2526
using Xunit.Abstractions;
2627

@@ -47,6 +48,7 @@ public ChunkMetricsReaderTests(ITestOutputHelper? outputHelper)
4748
public async Task Reader_GetChunkMetrics_ReturnsNonNull()
4849
{
4950
AdbcConnection? connection = null;
51+
Apache.Arrow.Ipc.IArrowArrayStream? reader = null;
5052

5153
try
5254
{
@@ -67,7 +69,7 @@ public async Task Reader_GetChunkMetrics_ReturnsNonNull()
6769
statement.SqlQuery = "SELECT * FROM range(1000000)";
6870

6971
var result = statement.ExecuteQuery();
70-
var reader = result.Stream;
72+
reader = result.Stream;
7173

7274
// Consume at least one batch to ensure chunks are downloaded
7375
var batch = await reader.ReadNextRecordBatchAsync();
@@ -87,11 +89,10 @@ public async Task Reader_GetChunkMetrics_ReturnsNonNull()
8789

8890
Assert.NotNull(chunkMetrics);
8991
OutputHelper?.WriteLine($"ChunkMetrics retrieved successfully from reader");
90-
91-
reader?.Dispose();
9292
}
9393
finally
9494
{
95+
reader?.Dispose();
9596
connection?.Dispose();
9697
}
9798
}
@@ -104,6 +105,7 @@ public async Task Reader_GetChunkMetrics_ReturnsNonNull()
104105
public async Task Reader_GetChunkMetrics_MatchesDownloaderValues()
105106
{
106107
AdbcConnection? connection = null;
108+
Apache.Arrow.Ipc.IArrowArrayStream? reader = null;
107109

108110
try
109111
{
@@ -121,7 +123,7 @@ public async Task Reader_GetChunkMetrics_MatchesDownloaderValues()
121123
statement.SqlQuery = "SELECT * FROM range(1000000)";
122124

123125
var result = statement.ExecuteQuery();
124-
var reader = result.Stream;
126+
reader = result.Stream;
125127

126128
// Consume several batches to ensure multiple chunks are processed
127129
int batchCount = 0;
@@ -166,11 +168,10 @@ public async Task Reader_GetChunkMetrics_MatchesDownloaderValues()
166168
OutputHelper?.WriteLine($" InitialChunkLatencyMs: {initialChunkLatencyMs}");
167169
OutputHelper?.WriteLine($" SlowestChunkLatencyMs: {slowestChunkLatencyMs}");
168170
OutputHelper?.WriteLine($" SumChunksDownloadTimeMs: {sumChunksDownloadTimeMs}");
169-
170-
reader?.Dispose();
171171
}
172172
finally
173173
{
174+
reader?.Dispose();
174175
connection?.Dispose();
175176
}
176177
}
@@ -183,6 +184,7 @@ public async Task Reader_GetChunkMetrics_MatchesDownloaderValues()
183184
public async Task Reader_GetChunkMetrics_AvailableAfterBatchConsumption()
184185
{
185186
AdbcConnection? connection = null;
187+
Apache.Arrow.Ipc.IArrowArrayStream? reader = null;
186188

187189
try
188190
{
@@ -200,7 +202,7 @@ public async Task Reader_GetChunkMetrics_AvailableAfterBatchConsumption()
200202
statement.SqlQuery = "SELECT * FROM range(1000000)";
201203

202204
var result = statement.ExecuteQuery();
203-
var reader = result.Stream;
205+
reader = result.Stream;
204206

205207
// Act - Consume all batches
206208
int totalBatches = 0;
@@ -235,11 +237,10 @@ public async Task Reader_GetChunkMetrics_AvailableAfterBatchConsumption()
235237
OutputHelper?.WriteLine($"Metrics available after full consumption:");
236238
OutputHelper?.WriteLine($" TotalChunksPresent: {totalChunksPresent}");
237239
OutputHelper?.WriteLine($" TotalChunksIterated: {totalChunksIterated}");
238-
239-
reader?.Dispose();
240240
}
241241
finally
242242
{
243+
reader?.Dispose();
243244
connection?.Dispose();
244245
}
245246
}
@@ -253,6 +254,7 @@ public async Task Reader_GetChunkMetrics_AvailableAfterBatchConsumption()
253254
public async Task Reader_GetChunkMetrics_ReflectsPartialConsumption()
254255
{
255256
AdbcConnection? connection = null;
257+
Apache.Arrow.Ipc.IArrowArrayStream? reader = null;
256258

257259
try
258260
{
@@ -270,7 +272,7 @@ public async Task Reader_GetChunkMetrics_ReflectsPartialConsumption()
270272
statement.SqlQuery = "SELECT * FROM range(2000000)"; // Large enough to ensure multiple chunks
271273

272274
var result = statement.ExecuteQuery();
273-
var reader = result.Stream;
275+
reader = result.Stream;
274276

275277
// Act - Consume only a few batches, not all
276278
int batchesToConsume = 3;
@@ -306,11 +308,10 @@ public async Task Reader_GetChunkMetrics_ReflectsPartialConsumption()
306308
OutputHelper?.WriteLine($" Batches consumed: {batchCount}");
307309
OutputHelper?.WriteLine($" TotalChunksPresent: {totalChunksPresent}");
308310
OutputHelper?.WriteLine($" TotalChunksIterated: {totalChunksIterated}");
309-
310-
reader?.Dispose();
311311
}
312312
finally
313313
{
314+
reader?.Dispose();
314315
connection?.Dispose();
315316
}
316317
}
@@ -323,6 +324,7 @@ public async Task Reader_GetChunkMetrics_ReflectsPartialConsumption()
323324
public async Task Reader_GetChunkMetrics_ConsistentAcrossMultipleCalls()
324325
{
325326
AdbcConnection? connection = null;
327+
Apache.Arrow.Ipc.IArrowArrayStream? reader = null;
326328

327329
try
328330
{
@@ -338,7 +340,7 @@ public async Task Reader_GetChunkMetrics_ConsistentAcrossMultipleCalls()
338340
statement.SqlQuery = "SELECT * FROM range(1000000)";
339341

340342
var result = statement.ExecuteQuery();
341-
var reader = result.Stream;
343+
reader = result.Stream;
342344

343345
// Consume some batches
344346
var batch = await reader.ReadNextRecordBatchAsync();
@@ -367,11 +369,10 @@ public async Task Reader_GetChunkMetrics_ConsistentAcrossMultipleCalls()
367369
Assert.Equal(iterated1, iterated2);
368370

369371
OutputHelper?.WriteLine("Metrics are consistent across multiple calls");
370-
371-
reader?.Dispose();
372372
}
373373
finally
374374
{
375+
reader?.Dispose();
375376
connection?.Dispose();
376377
}
377378
}

‎csharp/test/E2E/Telemetry/InternalCallTests.cs‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -86,8 +86,9 @@ public async Task InternalCall_UseSchema_IsMarkedAsInternal()
8686
return protoLog.SqlOperation?.OperationDetail != null;
8787
}).ToList();
8888

89-
// If there are multiple operations, check if any are internal
89+
// Check if any operations are marked as internal
9090
// Internal operations would have been from SetSchema()
91+
bool foundInternalCall = false;
9192
foreach (var log in useSchemaLogs)
9293
{
9394
var protoLog = TelemetryTestHelpers.GetProtoLog(log);
@@ -97,10 +98,18 @@ public async Task InternalCall_UseSchema_IsMarkedAsInternal()
9798
{
9899
OutputHelper?.WriteLine($"Found operation: StatementType={protoLog.SqlOperation.StatementType}, " +
99100
$"IsInternalCall={opDetail.IsInternalCall}");
101+
if (opDetail.IsInternalCall)
102+
{
103+
foundInternalCall = true;
104+
}
100105
}
101106
}
102107

103-
OutputHelper?.WriteLine($"✓ Captured {logs.Count} telemetry event(s)");
108+
// Assert that at least one log entry has IsInternalCall set to true
109+
Assert.True(foundInternalCall,
110+
"Expected at least one telemetry log entry with IsInternalCall == true from the internal USE SCHEMA operation");
111+
112+
OutputHelper?.WriteLine($"✓ Captured {logs.Count} telemetry event(s), found internal call: {foundInternalCall}");
104113
}
105114
finally
106115
{

‎csharp/test/E2E/Telemetry/MetadataOperationTests.cs‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ public async Task Telemetry_GetObjects_Catalogs_EmitsListCatalogs()
6060
columnNamePattern: null);
6161

6262
// Consume the stream
63-
while (stream.ReadNextRecordBatchAsync().Result != null) { }
63+
while (await stream.ReadNextRecordBatchAsync() != null) { }
6464

6565
// Wait for telemetry events
6666
var logs = await TelemetryTestHelpers.WaitForTelemetryEvents(exporter, expectedCount: 1, timeoutMs: 5000);
@@ -113,7 +113,7 @@ public async Task Telemetry_GetObjects_Schemas_EmitsListSchemas()
113113
columnNamePattern: null);
114114

115115
// Consume the stream
116-
while (stream.ReadNextRecordBatchAsync().Result != null) { }
116+
while (await stream.ReadNextRecordBatchAsync() != null) { }
117117

118118
// Wait for telemetry events
119119
var logs = await TelemetryTestHelpers.WaitForTelemetryEvents(exporter, expectedCount: 1, timeoutMs: 5000);
@@ -166,7 +166,7 @@ public async Task Telemetry_GetObjects_Tables_EmitsListTables()
166166
columnNamePattern: null);
167167

168168
// Consume the stream
169-
while (stream.ReadNextRecordBatchAsync().Result != null) { }
169+
while (await stream.ReadNextRecordBatchAsync() != null) { }
170170

171171
// Wait for telemetry events
172172
var logs = await TelemetryTestHelpers.WaitForTelemetryEvents(exporter, expectedCount: 1, timeoutMs: 5000);
@@ -219,7 +219,7 @@ public async Task Telemetry_GetObjects_Columns_EmitsListColumns()
219219
columnNamePattern: null);
220220

221221
// Consume the stream
222-
while (stream.ReadNextRecordBatchAsync().Result != null) { }
222+
while (await stream.ReadNextRecordBatchAsync() != null) { }
223223

224224
// Wait for telemetry events
225225
var logs = await TelemetryTestHelpers.WaitForTelemetryEvents(exporter, expectedCount: 1, timeoutMs: 5000);
@@ -266,7 +266,7 @@ public async Task Telemetry_GetTableTypes_EmitsListTableTypes()
266266
using var stream = connection.GetTableTypes();
267267

268268
// Consume the stream
269-
while (stream.ReadNextRecordBatchAsync().Result != null) { }
269+
while (await stream.ReadNextRecordBatchAsync() != null) { }
270270

271271
// Wait for telemetry events
272272
var logs = await TelemetryTestHelpers.WaitForTelemetryEvents(exporter, expectedCount: 1, timeoutMs: 5000);
@@ -331,7 +331,7 @@ public async Task Telemetry_GetObjects_AllDepths_EmitCorrectOperationType()
331331
columnNamePattern: null);
332332

333333
// Consume the stream
334-
while (stream.ReadNextRecordBatchAsync().Result != null) { }
334+
while (await stream.ReadNextRecordBatchAsync() != null) { }
335335

336336
// Flush telemetry
337337
if (connection is DatabricksConnection dbConn && dbConn.TelemetrySession?.TelemetryClient != null)

0 commit comments

Comments
 (0)