You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by cu...@apache.org on 2025/07/03 14:51:30 UTC
(arrow-adbc) branch main updated: feat(csharp/src/Drivers/Databricks): Add W3C trace context (#3082)
This is an automated email from the ASF dual-hosted git repository.
curth pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-adbc.git
The following commit(s) were added to refs/heads/main by this push:
new 1b04fca51 feat(csharp/src/Drivers/Databricks): Add W3C trace context (#3082)
1b04fca51 is described below
commit 1b04fca517b5ab2182419dfc05772407508dcc3b
Author: Jade Wang <11...@users.noreply.github.com>
AuthorDate: Thu Jul 3 07:51:25 2025 -0700
feat(csharp/src/Drivers/Databricks): Add W3C trace context (#3082)
## Summary
- Adds W3C trace context propagation to the Databricks C# driver via
HTTP headers
- Makes trace propagation configurable with options for enablement,
header names, and trace state
- Integrates with existing ADBC tracing infrastructure to support
distributed tracing
## Details
This PR implements W3C Trace Context (traceparent/tracestate)
propagation
for the Databricks C# driver. When distributed tracing is enabled, the
driver now propagates trace context to Databricks servers, enabling
end-to-end request tracing across the distributed system.
### Implementation
- Added `TracingDelegatingHandler` to inject trace headers into HTTP
requests
- Integrated handler into the existing HTTP pipeline in
`DatabricksConnection`
- Added configuration options:
- `adbc.databricks.trace_propagation.enabled` (default: true)
- `adbc.databricks.trace_propagation.header_name` (default:
"traceparent")
- `adbc.databricks.trace_propagation.state_enabled` (default: false)
### Behavior
1. If `Activity.Current` exists, uses its trace context
2. Falls back to trace parent from connection options
(`adbc.telemetry.trace_parent`)
3. Optionally includes trace state header when enabled
4. Supports custom header names for compatibility with different systems
## Test plan
- [x] Added comprehensive unit tests for `TracingDelegatingHandler`
- [x] Added tests for configuration parsing in
`DatabricksConnectionTest`
- [x] Verified trace header propagation with mock HTTP handlers
- [x] Tested with and without active activities
- [x] Tested custom header names and trace state propagation
- [ ] Manual testing with actual Databricks instance
---
.../src/Drivers/Databricks/DatabricksConnection.cs | 52 +++-
.../src/Drivers/Databricks/DatabricksParameters.cs | 21 ++
.../Drivers/Databricks/TracingDelegatingHandler.cs | 101 +++++++
.../Drivers/Databricks/DatabricksConnectionTest.cs | 35 +++
.../Databricks/DatabricksTestConfiguration.cs | 9 +
.../Databricks/DatabricksTestEnvironment.cs | 12 +
.../Databricks/TracingDelegatingHandlerTest.cs | 318 +++++++++++++++++++++
7 files changed, 547 insertions(+), 1 deletion(-)
diff --git a/csharp/src/Drivers/Databricks/DatabricksConnection.cs b/csharp/src/Drivers/Databricks/DatabricksConnection.cs
index 18beb5c2c..00e596eb8 100644
--- a/csharp/src/Drivers/Databricks/DatabricksConnection.cs
+++ b/csharp/src/Drivers/Databricks/DatabricksConnection.cs
@@ -29,6 +29,7 @@ using Apache.Arrow.Adbc.Drivers.Apache.Hive2.Client;
using Apache.Arrow.Adbc.Drivers.Apache.Spark;
using Apache.Arrow.Adbc.Drivers.Databricks.Auth;
using Apache.Arrow.Adbc.Drivers.Databricks.CloudFetch;
+using Apache.Arrow.Adbc.Tracing;
using Apache.Arrow.Ipc;
using Apache.Hive.Service.Rpc.Thrift;
using Thrift.Protocol;
@@ -58,6 +59,11 @@ namespace Apache.Arrow.Adbc.Drivers.Databricks
private const int DefaultTemporarilyUnavailableRetryTimeout = 900;
private bool _useDescTableExtended = true;
+ // Trace propagation configuration
+ private bool _tracePropagationEnabled = true;
+ private string _traceParentHeaderName = "traceparent";
+ private bool _traceStateEnabled = false;
+
// Default namespace
private TNamespace? _defaultNamespace;
@@ -199,6 +205,43 @@ namespace Apache.Arrow.Adbc.Drivers.Databricks
ns.SchemaName = defaultSchema;
_defaultNamespace = ns;
}
+
+ // Parse trace propagation options
+ if (Properties.TryGetValue(DatabricksParameters.TracePropagationEnabled, out string? tracePropagationEnabledStr))
+ {
+ if (bool.TryParse(tracePropagationEnabledStr, out bool tracePropagationEnabled))
+ {
+ _tracePropagationEnabled = tracePropagationEnabled;
+ }
+ else
+ {
+ throw new ArgumentException($"Parameter '{DatabricksParameters.TracePropagationEnabled}' value '{tracePropagationEnabledStr}' could not be parsed. Valid values are 'true' and 'false'.");
+ }
+ }
+
+ if (Properties.TryGetValue(DatabricksParameters.TraceParentHeaderName, out string? traceParentHeaderName))
+ {
+ if (!string.IsNullOrWhiteSpace(traceParentHeaderName))
+ {
+ _traceParentHeaderName = traceParentHeaderName;
+ }
+ else
+ {
+ throw new ArgumentException($"Parameter '{DatabricksParameters.TraceParentHeaderName}' cannot be empty.");
+ }
+ }
+
+ if (Properties.TryGetValue(DatabricksParameters.TraceStateEnabled, out string? traceStateEnabledStr))
+ {
+ if (bool.TryParse(traceStateEnabledStr, out bool traceStateEnabled))
+ {
+ _traceStateEnabled = traceStateEnabled;
+ }
+ else
+ {
+ throw new ArgumentException($"Parameter '{DatabricksParameters.TraceStateEnabled}' value '{traceStateEnabledStr}' could not be parsed. Valid values are 'true' and 'false'.");
+ }
+ }
}
/// <summary>
@@ -259,9 +302,16 @@ namespace Apache.Arrow.Adbc.Drivers.Databricks
protected override HttpMessageHandler CreateHttpHandler()
{
HttpMessageHandler baseHandler = base.CreateHttpHandler();
+
+ // Add tracing handler to propagate W3C trace context if enabled
+ if (_tracePropagationEnabled)
+ {
+ baseHandler = new TracingDelegatingHandler(baseHandler, this, _traceParentHeaderName, _traceStateEnabled);
+ }
+
if (TemporarilyUnavailableRetry)
{
- // Add OAuth handler if OAuth authentication is being used
+ // Add retry handler for 503 responses
baseHandler = new RetryHttpHandler(baseHandler, TemporarilyUnavailableRetryTimeout);
}
diff --git a/csharp/src/Drivers/Databricks/DatabricksParameters.cs b/csharp/src/Drivers/Databricks/DatabricksParameters.cs
index 26e35c9f3..861b7b324 100644
--- a/csharp/src/Drivers/Databricks/DatabricksParameters.cs
+++ b/csharp/src/Drivers/Databricks/DatabricksParameters.cs
@@ -179,6 +179,27 @@ namespace Apache.Arrow.Adbc.Drivers.Databricks
/// Default value is true if not specified.
/// </summary>
public const string UseDescTableExtended = "adbc.databricks.use_desc_table_extended";
+
+ /// <summary>
+ /// Whether to propagate trace parent headers in HTTP requests.
+ /// Default value is true if not specified.
+ /// When enabled, the driver will add W3C Trace Context headers to all HTTP requests.
+ /// </summary>
+ public const string TracePropagationEnabled = "adbc.databricks.trace_propagation.enabled";
+
+ /// <summary>
+ /// The name of the HTTP header to use for trace parent propagation.
+ /// Default value is "traceparent" (W3C standard) if not specified.
+ /// This allows customization for systems that use different header names.
+ /// </summary>
+ public const string TraceParentHeaderName = "adbc.databricks.trace_propagation.header_name";
+
+ /// <summary>
+ /// Whether to include trace state header in HTTP requests.
+ /// Default value is false if not specified.
+ /// When enabled, the driver will also propagate the tracestate header if available.
+ /// </summary>
+ public const string TraceStateEnabled = "adbc.databricks.trace_propagation.state_enabled";
}
/// <summary>
diff --git a/csharp/src/Drivers/Databricks/TracingDelegatingHandler.cs b/csharp/src/Drivers/Databricks/TracingDelegatingHandler.cs
new file mode 100644
index 000000000..f07c0dad8
--- /dev/null
+++ b/csharp/src/Drivers/Databricks/TracingDelegatingHandler.cs
@@ -0,0 +1,101 @@
+/*
+* Licensed to the Apache Software Foundation (ASF) under one or more
+* contributor license agreements. See the NOTICE file distributed with
+* this work for additional information regarding copyright ownership.
+* The ASF licenses this file to You under the Apache License, Version 2.0
+* (the "License"); you may not use this file except in compliance with
+* the License. You may obtain a copy of the License at
+*
+* http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+*/
+
+using System.Diagnostics;
+using System.Net.Http;
+using System.Threading;
+using System.Threading.Tasks;
+using Apache.Arrow.Adbc.Tracing;
+
+namespace Apache.Arrow.Adbc.Drivers.Databricks
+{
+ /// <summary>
+ /// A delegating handler that adds W3C Trace Context headers to outgoing HTTP requests.
+ /// </summary>
+ internal class TracingDelegatingHandler : DelegatingHandler
+ {
+ private readonly IActivityTracer _activityTracer;
+ private readonly string _traceParentHeaderName;
+ private readonly bool _includeTraceState;
+
+ /// <summary>
+ /// Initializes a new instance of the <see cref="TracingDelegatingHandler"/> class.
+ /// </summary>
+ /// <param name="innerHandler">The inner handler.</param>
+ /// <param name="activityTracer">The activity tracer that provides trace context.</param>
+ /// <param name="traceParentHeaderName">The name of the trace parent header. Defaults to "traceparent".</param>
+ /// <param name="includeTraceState">Whether to include trace state header. Defaults to false.</param>
+ public TracingDelegatingHandler(
+ HttpMessageHandler innerHandler,
+ IActivityTracer activityTracer,
+ string traceParentHeaderName = "traceparent",
+ bool includeTraceState = false)
+ : base(innerHandler)
+ {
+ _activityTracer = activityTracer;
+ _traceParentHeaderName = traceParentHeaderName;
+ _includeTraceState = includeTraceState;
+ }
+
+ /// <summary>
+ /// Sends an HTTP request with W3C trace context headers.
+ /// </summary>
+ /// <param name="request">The HTTP request message to send.</param>
+ /// <param name="cancellationToken">The cancellation token.</param>
+ /// <returns>The task object representing the asynchronous operation.</returns>
+ protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
+ {
+ // Get the current activity or use the trace parent from the connection
+ Activity? currentActivity = Activity.Current;
+ string? traceParentValue = null;
+ string? traceStateValue = null;
+
+ if (currentActivity != null)
+ {
+ // Use the current activity's W3C trace parent format
+ traceParentValue = currentActivity.Id;
+
+ // Get trace state if enabled
+ if (_includeTraceState)
+ {
+ traceStateValue = currentActivity.TraceStateString;
+ }
+ }
+ else if (!string.IsNullOrEmpty(_activityTracer.TraceParent))
+ {
+ // Fall back to the trace parent set on the connection
+ traceParentValue = _activityTracer.TraceParent;
+
+ // Note: We don't have trace state from the connection, only from Activity
+ }
+
+ // Add the trace parent header if we have a value
+ if (!string.IsNullOrEmpty(traceParentValue))
+ {
+ request.Headers.TryAddWithoutValidation(_traceParentHeaderName, traceParentValue);
+
+ // Add trace state header if we have a value and it's enabled
+ if (_includeTraceState && !string.IsNullOrEmpty(traceStateValue))
+ {
+ request.Headers.TryAddWithoutValidation("tracestate", traceStateValue);
+ }
+ }
+
+ return await base.SendAsync(request, cancellationToken);
+ }
+ }
+}
diff --git a/csharp/test/Drivers/Databricks/DatabricksConnectionTest.cs b/csharp/test/Drivers/Databricks/DatabricksConnectionTest.cs
index 9b30ab6a1..72171d781 100644
--- a/csharp/test/Drivers/Databricks/DatabricksConnectionTest.cs
+++ b/csharp/test/Drivers/Databricks/DatabricksConnectionTest.cs
@@ -320,6 +320,11 @@ namespace Apache.Arrow.Adbc.Tests.Drivers.Databricks
// Tests for EnableMultipleCatalogSupport parameter
Add(new(new() { [SparkParameters.Type] = SparkServerTypeConstants.Http, [SparkParameters.HostName] = "valid.server.com", [AdbcOptions.Username] = "user", [AdbcOptions.Password] = "myPassword", [DatabricksParameters.EnableMultipleCatalogSupport] = "notabool" }, typeof(ArgumentException)));
+
+ // Tests for trace propagation configuration parameters
+ Add(new(new() { [SparkParameters.Type] = SparkServerTypeConstants.Http, [SparkParameters.HostName] = "valid.server.com", [AdbcOptions.Username] = "user", [AdbcOptions.Password] = "myPassword", [DatabricksParameters.TracePropagationEnabled] = "notabool" }, typeof(ArgumentException)));
+ Add(new(new() { [SparkParameters.Type] = SparkServerTypeConstants.Http, [SparkParameters.HostName] = "valid.server.com", [AdbcOptions.Username] = "user", [AdbcOptions.Password] = "myPassword", [DatabricksParameters.TraceParentHeaderName] = "" }, typeof(ArgumentException)));
+ Add(new(new() { [SparkParameters.Type] = SparkServerTypeConstants.Http, [SparkParameters.HostName] = "valid.server.com", [AdbcOptions.Username] = "user", [AdbcOptions.Password] = "myPassword", [DatabricksParameters.TraceStateEnabled] = "notabool" }, typeof(ArgumentException)));
}
}
@@ -403,5 +408,35 @@ namespace Apache.Arrow.Adbc.Tests.Drivers.Databricks
$"Test passed for inputCatalog={inputCatalog}, inputSchema={inputSchema}, enableMultipleCatalogSupport={enableMultipleCatalogSupport}. " +
$"Runtime catalog={catalogFromRuntime}, schema={schemaFromRuntime}, Statement catalog={dbStatement.CatalogName}");
}
+
+ /// <summary>
+ /// Tests that trace propagation configuration is correctly parsed and stored.
+ /// </summary>
+ [SkippableTheory]
+ [InlineData("true", "traceparent", "false")]
+ [InlineData("false", "X-Trace-Id", "true")]
+ [InlineData("TRUE", "X-Custom-Trace", "FALSE")]
+ [InlineData("False", "custom-header", "True")]
+ public void TracePropagationConfigurationTest(string tracePropagationEnabled, string traceParentHeaderName, string traceStateEnabled)
+ {
+ // Arrange
+ var testConfig = (DatabricksTestConfiguration)TestConfiguration.Clone();
+ testConfig.TracePropagationEnabled = tracePropagationEnabled;
+ testConfig.TraceParentHeaderName = traceParentHeaderName;
+ testConfig.TraceStateEnabled = traceStateEnabled;
+
+ // Act
+ using var connection = NewConnection(testConfig);
+
+ // Assert
+ Assert.NotNull(connection);
+ Assert.IsType<DatabricksConnection>(connection);
+
+ // We can't directly test the private fields, but we can verify the connection was created successfully
+ // with the provided configuration. The actual functionality is tested in TracingDelegatingHandlerTest.
+ OutputHelper?.WriteLine(
+ $"Connection created successfully with tracePropagationEnabled={tracePropagationEnabled}, " +
+ $"traceParentHeaderName={traceParentHeaderName}, traceStateEnabled={traceStateEnabled}");
+ }
}
}
diff --git a/csharp/test/Drivers/Databricks/DatabricksTestConfiguration.cs b/csharp/test/Drivers/Databricks/DatabricksTestConfiguration.cs
index ae1a4838e..e41ffa46c 100644
--- a/csharp/test/Drivers/Databricks/DatabricksTestConfiguration.cs
+++ b/csharp/test/Drivers/Databricks/DatabricksTestConfiguration.cs
@@ -36,5 +36,14 @@ namespace Apache.Arrow.Adbc.Tests.Drivers.Databricks
[JsonPropertyName("enableMultipleCatalogSupport")]
public string EnableMultipleCatalogSupport { get; set; } = string.Empty;
+
+ [JsonPropertyName("tracePropagationEnabled"), JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingDefault)]
+ public string TracePropagationEnabled { get; set; } = string.Empty;
+
+ [JsonPropertyName("traceParentHeaderName"), JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingDefault)]
+ public string TraceParentHeaderName { get; set; } = string.Empty;
+
+ [JsonPropertyName("traceStateEnabled"), JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingDefault)]
+ public string TraceStateEnabled { get; set; } = string.Empty;
}
}
diff --git a/csharp/test/Drivers/Databricks/DatabricksTestEnvironment.cs b/csharp/test/Drivers/Databricks/DatabricksTestEnvironment.cs
index a2a9ebbf5..6f9c08473 100644
--- a/csharp/test/Drivers/Databricks/DatabricksTestEnvironment.cs
+++ b/csharp/test/Drivers/Databricks/DatabricksTestEnvironment.cs
@@ -145,6 +145,18 @@ namespace Apache.Arrow.Adbc.Tests.Drivers.Databricks
{
parameters.Add(DatabricksParameters.EnableMultipleCatalogSupport, testConfiguration.EnableMultipleCatalogSupport!);
}
+ if (!string.IsNullOrEmpty(testConfiguration.TracePropagationEnabled))
+ {
+ parameters.Add(DatabricksParameters.TracePropagationEnabled, testConfiguration.TracePropagationEnabled!);
+ }
+ if (!string.IsNullOrEmpty(testConfiguration.TraceParentHeaderName))
+ {
+ parameters.Add(DatabricksParameters.TraceParentHeaderName, testConfiguration.TraceParentHeaderName!);
+ }
+ if (!string.IsNullOrEmpty(testConfiguration.TraceStateEnabled))
+ {
+ parameters.Add(DatabricksParameters.TraceStateEnabled, testConfiguration.TraceStateEnabled!);
+ }
if (testConfiguration.HttpOptions != null)
{
if (testConfiguration.HttpOptions.Tls != null)
diff --git a/csharp/test/Drivers/Databricks/TracingDelegatingHandlerTest.cs b/csharp/test/Drivers/Databricks/TracingDelegatingHandlerTest.cs
new file mode 100644
index 000000000..794e5c99a
--- /dev/null
+++ b/csharp/test/Drivers/Databricks/TracingDelegatingHandlerTest.cs
@@ -0,0 +1,318 @@
+/*
+* Licensed to the Apache Software Foundation (ASF) under one or more
+* contributor license agreements. See the NOTICE file distributed with
+* this work for additional information regarding copyright ownership.
+* The ASF licenses this file to You under the Apache License, Version 2.0
+* (the "License"); you may not use this file except in compliance with
+* the License. You may obtain a copy of the License at
+*
+* http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+*/
+
+using System;
+using System.Diagnostics;
+using System.Linq;
+using System.Net;
+using System.Net.Http;
+using System.Threading;
+using System.Threading.Tasks;
+using Apache.Arrow.Adbc.Tracing;
+using Xunit;
+
+namespace Apache.Arrow.Adbc.Drivers.Databricks.Tests
+{
+ /// <summary>
+ /// Tests for the TracingDelegatingHandler class.
+ /// </summary>
+ public class TracingDelegatingHandlerTest
+ {
+ /// <summary>
+ /// Mock HTTP message handler for testing.
+ /// </summary>
+ private class MockHttpMessageHandler : HttpMessageHandler
+ {
+ public HttpRequestMessage? CapturedRequest { get; private set; }
+ public HttpResponseMessage ResponseToReturn { get; set; } = new HttpResponseMessage(HttpStatusCode.OK);
+
+ protected override Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
+ {
+ CapturedRequest = request;
+ return Task.FromResult(ResponseToReturn);
+ }
+ }
+
+ /// <summary>
+ /// Mock activity tracer for testing.
+ /// </summary>
+ private class MockActivityTracer : IActivityTracer
+ {
+ public string? TraceParent { get; set; }
+ public ActivityTrace Trace => new ActivityTrace("TestSource", "1.0.0", TraceParent);
+ public string AssemblyVersion => "1.0.0";
+ public string AssemblyName => "TestAssembly";
+ }
+
+ [Fact]
+ public async Task SendAsync_WithCurrentActivity_AddsTraceParentHeader()
+ {
+ // Arrange
+ var mockHandler = new MockHttpMessageHandler();
+ var mockTracer = new MockActivityTracer();
+ var tracingHandler = new TracingDelegatingHandler(mockHandler, mockTracer);
+ var httpClient = new HttpClient(tracingHandler);
+
+ // Create an activity with a listener to ensure it's created
+ Activity.DefaultIdFormat = ActivityIdFormat.W3C;
+ ActivitySource.AddActivityListener(new ActivityListener
+ {
+ ShouldListenTo = _ => true,
+ Sample = (ref ActivityCreationOptions<ActivityContext> _) => ActivitySamplingResult.AllData
+ });
+
+ using var activitySource = new ActivitySource("TestActivitySource");
+ using var activity = activitySource.StartActivity("TestActivity", ActivityKind.Client);
+
+ // Skip test if activity creation is not supported in this environment
+ if (activity == null)
+ {
+ // Test with trace parent from tracer instead
+ var fallbackTraceParent = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01";
+ mockTracer.TraceParent = fallbackTraceParent;
+
+ var request = new HttpRequestMessage(HttpMethod.Get, "https://test.databricks.com/api/test");
+ await httpClient.SendAsync(request);
+
+ Assert.NotNull(mockHandler.CapturedRequest);
+ Assert.True(mockHandler.CapturedRequest.Headers.Contains("traceparent"));
+ var traceParentValue = mockHandler.CapturedRequest.Headers.GetValues("traceparent").FirstOrDefault();
+ Assert.Equal(fallbackTraceParent, traceParentValue);
+ return;
+ }
+
+ var expectedTraceParent = activity.Id;
+
+ // Act
+ var request2 = new HttpRequestMessage(HttpMethod.Get, "https://test.databricks.com/api/test");
+ await httpClient.SendAsync(request2);
+
+ // Assert
+ Assert.NotNull(mockHandler.CapturedRequest);
+ Assert.True(mockHandler.CapturedRequest.Headers.Contains("traceparent"));
+ var actualTraceParent = mockHandler.CapturedRequest.Headers.GetValues("traceparent").FirstOrDefault();
+ Assert.Equal(expectedTraceParent, actualTraceParent);
+ }
+
+ [Fact]
+ public async Task SendAsync_WithTracerTraceParent_AddsTraceParentHeader()
+ {
+ // Arrange
+ var mockHandler = new MockHttpMessageHandler();
+ var expectedTraceParent = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01";
+ var mockTracer = new MockActivityTracer { TraceParent = expectedTraceParent };
+ var tracingHandler = new TracingDelegatingHandler(mockHandler, mockTracer);
+ var httpClient = new HttpClient(tracingHandler);
+
+ // Act
+ var request = new HttpRequestMessage(HttpMethod.Get, "https://test.databricks.com/api/test");
+ await httpClient.SendAsync(request);
+
+ // Assert
+ Assert.NotNull(mockHandler.CapturedRequest);
+ Assert.True(mockHandler.CapturedRequest.Headers.Contains("traceparent"));
+ var traceParentValue = mockHandler.CapturedRequest.Headers.GetValues("traceparent").FirstOrDefault();
+ Assert.Equal(expectedTraceParent, traceParentValue);
+ }
+
+ [Fact]
+ public async Task SendAsync_WithoutTraceContext_DoesNotAddHeader()
+ {
+ // Arrange
+ var mockHandler = new MockHttpMessageHandler();
+ var mockTracer = new MockActivityTracer { TraceParent = null };
+ var tracingHandler = new TracingDelegatingHandler(mockHandler, mockTracer);
+ var httpClient = new HttpClient(tracingHandler);
+
+ // Act
+ var request = new HttpRequestMessage(HttpMethod.Get, "https://test.databricks.com/api/test");
+ await httpClient.SendAsync(request);
+
+ // Assert
+ Assert.NotNull(mockHandler.CapturedRequest);
+ Assert.False(mockHandler.CapturedRequest.Headers.Contains("traceparent"));
+ }
+
+ [Fact]
+ public async Task SendAsync_CurrentActivityTakesPrecedenceOverTracer()
+ {
+ // Arrange
+ var mockHandler = new MockHttpMessageHandler();
+ var tracerTraceParent = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01";
+ var mockTracer = new MockActivityTracer { TraceParent = tracerTraceParent };
+ var tracingHandler = new TracingDelegatingHandler(mockHandler, mockTracer);
+ var httpClient = new HttpClient(tracingHandler);
+
+ // Create an activity with a listener to ensure it's created
+ Activity.DefaultIdFormat = ActivityIdFormat.W3C;
+ ActivitySource.AddActivityListener(new ActivityListener
+ {
+ ShouldListenTo = _ => true,
+ Sample = (ref ActivityCreationOptions<ActivityContext> _) => ActivitySamplingResult.AllData
+ });
+
+ using var activitySource = new ActivitySource("TestActivitySource");
+ using var activity = activitySource.StartActivity("TestActivity", ActivityKind.Client);
+
+ // Skip test if activity creation is not supported in this environment
+ if (activity == null)
+ {
+ // Can't test precedence without an activity - just verify tracer parent works
+ var request = new HttpRequestMessage(HttpMethod.Get, "https://test.databricks.com/api/test");
+ await httpClient.SendAsync(request);
+
+ Assert.NotNull(mockHandler.CapturedRequest);
+ Assert.True(mockHandler.CapturedRequest.Headers.Contains("traceparent"));
+ var traceParentValue = mockHandler.CapturedRequest.Headers.GetValues("traceparent").FirstOrDefault();
+ Assert.Equal(tracerTraceParent, traceParentValue);
+ return;
+ }
+
+ var expectedTraceParent = activity.Id;
+
+ // Act
+ var request2 = new HttpRequestMessage(HttpMethod.Get, "https://test.databricks.com/api/test");
+ await httpClient.SendAsync(request2);
+
+ // Assert
+ Assert.NotNull(mockHandler.CapturedRequest);
+ Assert.True(mockHandler.CapturedRequest.Headers.Contains("traceparent"));
+ var actualTraceParent = mockHandler.CapturedRequest.Headers.GetValues("traceparent").FirstOrDefault();
+ Assert.Equal(expectedTraceParent, actualTraceParent);
+ Assert.NotEqual(tracerTraceParent, actualTraceParent);
+ }
+
+ [Fact]
+ public async Task SendAsync_PreservesOtherHeaders()
+ {
+ // Arrange
+ var mockHandler = new MockHttpMessageHandler();
+ var mockTracer = new MockActivityTracer { TraceParent = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" };
+ var tracingHandler = new TracingDelegatingHandler(mockHandler, mockTracer);
+ var httpClient = new HttpClient(tracingHandler);
+
+ // Act
+ var request = new HttpRequestMessage(HttpMethod.Get, "https://test.databricks.com/api/test");
+ request.Headers.Add("Authorization", "Bearer test-token");
+ request.Headers.Add("X-Custom-Header", "custom-value");
+ await httpClient.SendAsync(request);
+
+ // Assert
+ Assert.NotNull(mockHandler.CapturedRequest);
+ Assert.True(mockHandler.CapturedRequest.Headers.Contains("traceparent"));
+ Assert.True(mockHandler.CapturedRequest.Headers.Contains("Authorization"));
+ Assert.True(mockHandler.CapturedRequest.Headers.Contains("X-Custom-Header"));
+ Assert.Equal("Bearer test-token", mockHandler.CapturedRequest.Headers.GetValues("Authorization").FirstOrDefault());
+ Assert.Equal("custom-value", mockHandler.CapturedRequest.Headers.GetValues("X-Custom-Header").FirstOrDefault());
+ }
+
+ [Fact]
+ public async Task SendAsync_WithCustomHeaderName_UsesCustomHeader()
+ {
+ // Arrange
+ var mockHandler = new MockHttpMessageHandler();
+ var traceParent = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01";
+ var mockTracer = new MockActivityTracer { TraceParent = traceParent };
+ var customHeaderName = "X-Trace-Id";
+ var tracingHandler = new TracingDelegatingHandler(mockHandler, mockTracer, customHeaderName);
+ var httpClient = new HttpClient(tracingHandler);
+
+ // Act
+ var request = new HttpRequestMessage(HttpMethod.Get, "https://test.databricks.com/api/test");
+ await httpClient.SendAsync(request);
+
+ // Assert
+ Assert.NotNull(mockHandler.CapturedRequest);
+ Assert.False(mockHandler.CapturedRequest.Headers.Contains("traceparent")); // Default header should not be present
+ Assert.True(mockHandler.CapturedRequest.Headers.Contains(customHeaderName));
+ var headerValue = mockHandler.CapturedRequest.Headers.GetValues(customHeaderName).FirstOrDefault();
+ Assert.Equal(traceParent, headerValue);
+ }
+
+ [Fact]
+ public async Task SendAsync_WithTraceStateEnabled_IncludesTraceState()
+ {
+ // Arrange
+ var mockHandler = new MockHttpMessageHandler();
+ var mockTracer = new MockActivityTracer();
+ var tracingHandler = new TracingDelegatingHandler(mockHandler, mockTracer, "traceparent", includeTraceState: true);
+ var httpClient = new HttpClient(tracingHandler);
+
+ // Create an activity with trace state
+ Activity.DefaultIdFormat = ActivityIdFormat.W3C;
+ ActivitySource.AddActivityListener(new ActivityListener
+ {
+ ShouldListenTo = _ => true,
+ Sample = (ref ActivityCreationOptions<ActivityContext> _) => ActivitySamplingResult.AllData
+ });
+
+ using var activitySource = new ActivitySource("TestActivitySource");
+ using var activity = activitySource.StartActivity("TestActivity", ActivityKind.Client);
+
+ if (activity != null)
+ {
+ activity.TraceStateString = "vendor1=value1,vendor2=value2";
+
+ // Act
+ var request = new HttpRequestMessage(HttpMethod.Get, "https://test.databricks.com/api/test");
+ await httpClient.SendAsync(request);
+
+ // Assert
+ Assert.NotNull(mockHandler.CapturedRequest);
+ Assert.True(mockHandler.CapturedRequest.Headers.Contains("traceparent"));
+ Assert.True(mockHandler.CapturedRequest.Headers.Contains("tracestate"));
+ var traceStateValue = mockHandler.CapturedRequest.Headers.GetValues("tracestate").FirstOrDefault();
+ Assert.Equal("vendor1=value1,vendor2=value2", traceStateValue);
+ }
+ }
+
+ [Fact]
+ public async Task SendAsync_WithTraceStateDisabled_DoesNotIncludeTraceState()
+ {
+ // Arrange
+ var mockHandler = new MockHttpMessageHandler();
+ var mockTracer = new MockActivityTracer();
+ var tracingHandler = new TracingDelegatingHandler(mockHandler, mockTracer, "traceparent", includeTraceState: false);
+ var httpClient = new HttpClient(tracingHandler);
+
+ // Create an activity with trace state
+ Activity.DefaultIdFormat = ActivityIdFormat.W3C;
+ ActivitySource.AddActivityListener(new ActivityListener
+ {
+ ShouldListenTo = _ => true,
+ Sample = (ref ActivityCreationOptions<ActivityContext> _) => ActivitySamplingResult.AllData
+ });
+
+ using var activitySource = new ActivitySource("TestActivitySource");
+ using var activity = activitySource.StartActivity("TestActivity", ActivityKind.Client);
+
+ if (activity != null)
+ {
+ activity.TraceStateString = "vendor1=value1,vendor2=value2";
+
+ // Act
+ var request = new HttpRequestMessage(HttpMethod.Get, "https://test.databricks.com/api/test");
+ await httpClient.SendAsync(request);
+
+ // Assert
+ Assert.NotNull(mockHandler.CapturedRequest);
+ Assert.True(mockHandler.CapturedRequest.Headers.Contains("traceparent"));
+ Assert.False(mockHandler.CapturedRequest.Headers.Contains("tracestate")); // Should not include trace state
+ }
+ }
+ }
+}