From 97995db05784d05b71284afe72b6fdeffa6826e2 Mon Sep 17 00:00:00 2001 From: Todd Meng Date: Wed, 16 Apr 2025 12:01:19 -0700 Subject: [PATCH 1/5] Add support for client user-agent entry in Spark driver --- .../Apache/Spark/SparkHttpConnection.cs | 15 +- .../Drivers/Apache/Spark/SparkParameters.cs | 1 + .../Spark/SparkHttpConnectionUserAgentTest.cs | 129 ++++++++++++++++++ 3 files changed, 143 insertions(+), 2 deletions(-) create mode 100644 csharp/test/Drivers/Apache/Spark/SparkHttpConnectionUserAgentTest.cs diff --git a/csharp/src/Drivers/Apache/Spark/SparkHttpConnection.cs b/csharp/src/Drivers/Apache/Spark/SparkHttpConnection.cs index 9cabd1ac41..730a9d78c4 100644 --- a/csharp/src/Drivers/Apache/Spark/SparkHttpConnection.cs +++ b/csharp/src/Drivers/Apache/Spark/SparkHttpConnection.cs @@ -35,7 +35,7 @@ namespace Apache.Arrow.Adbc.Drivers.Apache.Spark { internal class SparkHttpConnection : SparkConnection { - private static readonly string s_userAgent = $"{DriverName.Replace(" ", "")}/{ProductVersionDefault}"; + private static readonly string s_baseUserAgent = $"{DriverName.Replace(" ", "")}/{ProductVersionDefault}"; private const string BasicAuthenticationScheme = "Basic"; private const string BearerAuthenticationScheme = "Bearer"; @@ -164,7 +164,7 @@ protected override TTransport CreateTransport() HttpClient httpClient = new(httpClientHandler); httpClient.BaseAddress = baseAddress; httpClient.DefaultRequestHeaders.Authorization = authenticationHeaderValue; - httpClient.DefaultRequestHeaders.UserAgent.ParseAdd(s_userAgent); + httpClient.DefaultRequestHeaders.UserAgent.ParseAdd(GetUserAgent()); httpClient.DefaultRequestHeaders.AcceptEncoding.Clear(); httpClient.DefaultRequestHeaders.AcceptEncoding.Add(new StringWithQualityHeaderValue("identity")); httpClient.DefaultRequestHeaders.ExpectContinue = false; @@ -252,5 +252,16 @@ protected internal override Task GetRowSetAsync(TGetPrimaryKeysResp res internal override SparkServerType ServerType => SparkServerType.Http; protected override int ColumnMapIndexOffset => 1; + + private string GetUserAgent() + { + // Check if a client has provided a user-agent entry + if (Properties.TryGetValue(SparkParameters.UserAgentEntry, out string userAgentEntry) && !string.IsNullOrWhiteSpace(userAgentEntry)) + { + return $"{s_baseUserAgent} {userAgentEntry}"; + } + + return s_baseUserAgent; + } } } diff --git a/csharp/src/Drivers/Apache/Spark/SparkParameters.cs b/csharp/src/Drivers/Apache/Spark/SparkParameters.cs index 66d329814f..99275d7d0e 100644 --- a/csharp/src/Drivers/Apache/Spark/SparkParameters.cs +++ b/csharp/src/Drivers/Apache/Spark/SparkParameters.cs @@ -33,6 +33,7 @@ public class SparkParameters public const string Type = "adbc.spark.type"; public const string DataTypeConv = "adbc.spark.data_type_conv"; public const string ConnectTimeoutMilliseconds = "adbc.spark.connect_timeout_ms"; + public const string UserAgentEntry = "adbc.spark.user_agent_entry"; } public static class SparkAuthTypeConstants diff --git a/csharp/test/Drivers/Apache/Spark/SparkHttpConnectionUserAgentTest.cs b/csharp/test/Drivers/Apache/Spark/SparkHttpConnectionUserAgentTest.cs new file mode 100644 index 0000000000..3c15622ec9 --- /dev/null +++ b/csharp/test/Drivers/Apache/Spark/SparkHttpConnectionUserAgentTest.cs @@ -0,0 +1,129 @@ +/* +* 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.Collections.Generic; +using System.Net.Http; +using System.Reflection; +using Apache.Arrow.Adbc.Drivers.Apache.Spark; +using Xunit; + +namespace Apache.Arrow.Adbc.Tests.Drivers.Apache.Spark +{ + /// + /// Tests for the SparkHttpConnection user agent functionality. + /// + public class SparkHttpConnectionUserAgentTest + { + [Fact] + public void UserAgentEntry_WhenNotProvided_UsesBaseUserAgent() + { + // Arrange + var properties = new Dictionary + { + [SparkParameters.Type] = SparkServerTypeConstants.Http, + [SparkParameters.HostName] = "valid.server.com", + [SparkParameters.Path] = "/path", + [SparkParameters.AuthType] = SparkAuthTypeConstants.None + }; + + // Act + string userAgent = GetUserAgentFromConnection(properties); + + // Assert + Assert.Equal("ADBCSparkDriver/1.0.0", userAgent); + } + + [Fact] + public void UserAgentEntry_WhenProvided_AppendsToBaseUserAgent() + { + // Arrange + var properties = new Dictionary + { + [SparkParameters.Type] = SparkServerTypeConstants.Http, + [SparkParameters.HostName] = "valid.server.com", + [SparkParameters.Path] = "/path", + [SparkParameters.AuthType] = SparkAuthTypeConstants.None, + [SparkParameters.UserAgentEntry] = "PowerBI" + }; + + // Act + string userAgent = GetUserAgentFromConnection(properties); + + // Assert + Assert.Equal("ADBCSparkDriver/1.0.0 PowerBI", userAgent); + } + + [Fact] + public void UserAgentEntry_WhenEmpty_UsesBaseUserAgent() + { + // Arrange + var properties = new Dictionary + { + [SparkParameters.Type] = SparkServerTypeConstants.Http, + [SparkParameters.HostName] = "valid.server.com", + [SparkParameters.Path] = "/path", + [SparkParameters.AuthType] = SparkAuthTypeConstants.None, + [SparkParameters.UserAgentEntry] = "" + }; + + // Act + string userAgent = GetUserAgentFromConnection(properties); + + // Assert + Assert.Equal("ADBCSparkDriver/1.0.0", userAgent); + } + + [Fact] + public void UserAgentEntry_WhenWhitespace_UsesBaseUserAgent() + { + // Arrange + var properties = new Dictionary + { + [SparkParameters.Type] = SparkServerTypeConstants.Http, + [SparkParameters.HostName] = "valid.server.com", + [SparkParameters.Path] = "/path", + [SparkParameters.AuthType] = SparkAuthTypeConstants.None, + [SparkParameters.UserAgentEntry] = " " + }; + + // Act + string userAgent = GetUserAgentFromConnection(properties); + + // Assert + Assert.Equal("ADBCSparkDriver/1.0.0", userAgent); + } + + private string GetUserAgentFromConnection(Dictionary properties) + { + // Create the connection + var connection = new SparkHttpConnection(properties); + + // Use reflection to access the private GetUserAgent method + var getUserAgentMethod = typeof(SparkHttpConnection).GetMethod("GetUserAgent", + BindingFlags.NonPublic | BindingFlags.Instance); + + if (getUserAgentMethod == null) + { + throw new InvalidOperationException("GetUserAgent method not found"); + } + + // Invoke the method + return (string)getUserAgentMethod.Invoke(connection, null); + } + } +} \ No newline at end of file From 8057762928dafb16a13f99b041b9da532115ccf5 Mon Sep 17 00:00:00 2001 From: Todd Meng Date: Wed, 16 Apr 2025 12:13:04 -0700 Subject: [PATCH 2/5] Fix nullable reference type warnings --- csharp/src/Drivers/Apache/Spark/SparkHttpConnection.cs | 2 +- .../Drivers/Apache/Spark/SparkHttpConnectionUserAgentTest.cs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/csharp/src/Drivers/Apache/Spark/SparkHttpConnection.cs b/csharp/src/Drivers/Apache/Spark/SparkHttpConnection.cs index 730a9d78c4..0c1356d839 100644 --- a/csharp/src/Drivers/Apache/Spark/SparkHttpConnection.cs +++ b/csharp/src/Drivers/Apache/Spark/SparkHttpConnection.cs @@ -256,7 +256,7 @@ protected internal override Task GetRowSetAsync(TGetPrimaryKeysResp res private string GetUserAgent() { // Check if a client has provided a user-agent entry - if (Properties.TryGetValue(SparkParameters.UserAgentEntry, out string userAgentEntry) && !string.IsNullOrWhiteSpace(userAgentEntry)) + if (Properties.TryGetValue(SparkParameters.UserAgentEntry, out string? userAgentEntry) && !string.IsNullOrWhiteSpace(userAgentEntry)) { return $"{s_baseUserAgent} {userAgentEntry}"; } diff --git a/csharp/test/Drivers/Apache/Spark/SparkHttpConnectionUserAgentTest.cs b/csharp/test/Drivers/Apache/Spark/SparkHttpConnectionUserAgentTest.cs index 3c15622ec9..27c2cc5871 100644 --- a/csharp/test/Drivers/Apache/Spark/SparkHttpConnectionUserAgentTest.cs +++ b/csharp/test/Drivers/Apache/Spark/SparkHttpConnectionUserAgentTest.cs @@ -123,7 +123,7 @@ private string GetUserAgentFromConnection(Dictionary properties) } // Invoke the method - return (string)getUserAgentMethod.Invoke(connection, null); + return (string)getUserAgentMethod.Invoke(connection, null)!; } } } \ No newline at end of file From 0f70d050b6bf75995a68ccd96cd7c5efdd9425cf Mon Sep 17 00:00:00 2001 From: toddmeng-db Date: Mon, 21 Apr 2025 23:48:39 -0700 Subject: [PATCH 3/5] add thrift version --- .../Apache/Spark/SparkHttpConnection.cs | 24 +++++++++++++-- .../Spark/SparkHttpConnectionUserAgentTest.cs | 29 +++++++++++++++---- 2 files changed, 46 insertions(+), 7 deletions(-) diff --git a/csharp/src/Drivers/Apache/Spark/SparkHttpConnection.cs b/csharp/src/Drivers/Apache/Spark/SparkHttpConnection.cs index 0c1356d839..a5ecffaf27 100644 --- a/csharp/src/Drivers/Apache/Spark/SparkHttpConnection.cs +++ b/csharp/src/Drivers/Apache/Spark/SparkHttpConnection.cs @@ -255,13 +255,33 @@ protected internal override Task GetRowSetAsync(TGetPrimaryKeysResp res private string GetUserAgent() { + // Build the base user agent string with Thrift version + string thriftVersion = GetThriftVersion(); + string thriftComponent = string.IsNullOrEmpty(thriftVersion) ? "Thrift" : $"Thrift/{thriftVersion}"; + string baseUserAgent = $"{DriverName.Replace(" ", "")}/{ProductVersionDefault} {thriftComponent}"; // consider using ProductVersion instead of ProductVersionDefault + // Check if a client has provided a user-agent entry if (Properties.TryGetValue(SparkParameters.UserAgentEntry, out string? userAgentEntry) && !string.IsNullOrWhiteSpace(userAgentEntry)) { - return $"{s_baseUserAgent} {userAgentEntry}"; + return $"{baseUserAgent} {userAgentEntry}"; } - return s_baseUserAgent; + return baseUserAgent; + } + + private string GetThriftVersion() + { + try + { + var thriftAssembly = typeof(TProtocol).Assembly; + var version = thriftAssembly.GetName().Version; + return version != null ? $"{version.Major}.{version.Minor}.{version.Build}" : ""; + } + catch + { + // Return empty string if there's any issue retrieving the assembly version + return ""; + } } } } diff --git a/csharp/test/Drivers/Apache/Spark/SparkHttpConnectionUserAgentTest.cs b/csharp/test/Drivers/Apache/Spark/SparkHttpConnectionUserAgentTest.cs index 27c2cc5871..3b42b27165 100644 --- a/csharp/test/Drivers/Apache/Spark/SparkHttpConnectionUserAgentTest.cs +++ b/csharp/test/Drivers/Apache/Spark/SparkHttpConnectionUserAgentTest.cs @@ -30,7 +30,7 @@ namespace Apache.Arrow.Adbc.Tests.Drivers.Apache.Spark public class SparkHttpConnectionUserAgentTest { [Fact] - public void UserAgentEntry_WhenNotProvided_UsesBaseUserAgent() + public void UserAgentEntry_WhenNotProvided_UsesBaseUserAgentWithThrift() { // Arrange var properties = new Dictionary @@ -45,7 +45,7 @@ public void UserAgentEntry_WhenNotProvided_UsesBaseUserAgent() string userAgent = GetUserAgentFromConnection(properties); // Assert - Assert.Equal("ADBCSparkDriver/1.0.0", userAgent); + Assert.Matches(@"ADBCSparkDriver/[\d\.]+ Thrift(/[\d\.]+)?", userAgent); } [Fact] @@ -65,7 +65,7 @@ public void UserAgentEntry_WhenProvided_AppendsToBaseUserAgent() string userAgent = GetUserAgentFromConnection(properties); // Assert - Assert.Equal("ADBCSparkDriver/1.0.0 PowerBI", userAgent); + Assert.Matches(@"ADBCSparkDriver/[\d\.]+ Thrift(/[\d\.]+)? PowerBI", userAgent); } [Fact] @@ -85,7 +85,7 @@ public void UserAgentEntry_WhenEmpty_UsesBaseUserAgent() string userAgent = GetUserAgentFromConnection(properties); // Assert - Assert.Equal("ADBCSparkDriver/1.0.0", userAgent); + Assert.Matches(@"ADBCSparkDriver/[\d\.]+ Thrift(/[\d\.]+)?", userAgent); } [Fact] @@ -105,7 +105,26 @@ public void UserAgentEntry_WhenWhitespace_UsesBaseUserAgent() string userAgent = GetUserAgentFromConnection(properties); // Assert - Assert.Equal("ADBCSparkDriver/1.0.0", userAgent); + Assert.Matches(@"ADBCSparkDriver/[\d\.]+ Thrift(/[\d\.]+)?", userAgent); + } + + [Fact] + public void UserAgent_IncludesThriftComponent() + { + // Arrange + var properties = new Dictionary + { + [SparkParameters.Type] = SparkServerTypeConstants.Http, + [SparkParameters.HostName] = "valid.server.com", + [SparkParameters.Path] = "/path", + [SparkParameters.AuthType] = SparkAuthTypeConstants.None + }; + + // Act + string userAgent = GetUserAgentFromConnection(properties); + + // Assert + Assert.Contains("Thrift", userAgent); } private string GetUserAgentFromConnection(Dictionary properties) From b1bb50d70807d1a0589be37447498ac7c14c8df7 Mon Sep 17 00:00:00 2001 From: toddmeng-db Date: Mon, 21 Apr 2025 23:55:36 -0700 Subject: [PATCH 4/5] remove redudant readonly user agent --- csharp/src/Drivers/Apache/Spark/SparkHttpConnection.cs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/csharp/src/Drivers/Apache/Spark/SparkHttpConnection.cs b/csharp/src/Drivers/Apache/Spark/SparkHttpConnection.cs index a5ecffaf27..1f4748e135 100644 --- a/csharp/src/Drivers/Apache/Spark/SparkHttpConnection.cs +++ b/csharp/src/Drivers/Apache/Spark/SparkHttpConnection.cs @@ -35,7 +35,6 @@ namespace Apache.Arrow.Adbc.Drivers.Apache.Spark { internal class SparkHttpConnection : SparkConnection { - private static readonly string s_baseUserAgent = $"{DriverName.Replace(" ", "")}/{ProductVersionDefault}"; private const string BasicAuthenticationScheme = "Basic"; private const string BearerAuthenticationScheme = "Bearer"; @@ -258,7 +257,7 @@ private string GetUserAgent() // Build the base user agent string with Thrift version string thriftVersion = GetThriftVersion(); string thriftComponent = string.IsNullOrEmpty(thriftVersion) ? "Thrift" : $"Thrift/{thriftVersion}"; - string baseUserAgent = $"{DriverName.Replace(" ", "")}/{ProductVersionDefault} {thriftComponent}"; // consider using ProductVersion instead of ProductVersionDefault + string baseUserAgent = $"{DriverName.Replace(" ", "")}/{ProductVersionDefault} {thriftComponent}"; // Check if a client has provided a user-agent entry if (Properties.TryGetValue(SparkParameters.UserAgentEntry, out string? userAgentEntry) && !string.IsNullOrWhiteSpace(userAgentEntry)) From e846be3caf17ec3045f3cf98b081f1a75d5cbc1e Mon Sep 17 00:00:00 2001 From: toddmeng-db Date: Tue, 22 Apr 2025 11:28:16 -0700 Subject: [PATCH 5/5] lint fixes --- .../src/Drivers/Apache/Spark/SparkHttpConnection.cs | 8 ++++---- .../Apache/Spark/SparkHttpConnectionUserAgentTest.cs | 11 +++++------ 2 files changed, 9 insertions(+), 10 deletions(-) diff --git a/csharp/src/Drivers/Apache/Spark/SparkHttpConnection.cs b/csharp/src/Drivers/Apache/Spark/SparkHttpConnection.cs index 1f4748e135..75a81228b8 100644 --- a/csharp/src/Drivers/Apache/Spark/SparkHttpConnection.cs +++ b/csharp/src/Drivers/Apache/Spark/SparkHttpConnection.cs @@ -251,23 +251,23 @@ protected internal override Task GetRowSetAsync(TGetPrimaryKeysResp res internal override SparkServerType ServerType => SparkServerType.Http; protected override int ColumnMapIndexOffset => 1; - + private string GetUserAgent() { // Build the base user agent string with Thrift version string thriftVersion = GetThriftVersion(); string thriftComponent = string.IsNullOrEmpty(thriftVersion) ? "Thrift" : $"Thrift/{thriftVersion}"; string baseUserAgent = $"{DriverName.Replace(" ", "")}/{ProductVersionDefault} {thriftComponent}"; - + // Check if a client has provided a user-agent entry if (Properties.TryGetValue(SparkParameters.UserAgentEntry, out string? userAgentEntry) && !string.IsNullOrWhiteSpace(userAgentEntry)) { return $"{baseUserAgent} {userAgentEntry}"; } - + return baseUserAgent; } - + private string GetThriftVersion() { try diff --git a/csharp/test/Drivers/Apache/Spark/SparkHttpConnectionUserAgentTest.cs b/csharp/test/Drivers/Apache/Spark/SparkHttpConnectionUserAgentTest.cs index 3b42b27165..84871ad78f 100644 --- a/csharp/test/Drivers/Apache/Spark/SparkHttpConnectionUserAgentTest.cs +++ b/csharp/test/Drivers/Apache/Spark/SparkHttpConnectionUserAgentTest.cs @@ -131,18 +131,17 @@ private string GetUserAgentFromConnection(Dictionary properties) { // Create the connection var connection = new SparkHttpConnection(properties); - + // Use reflection to access the private GetUserAgent method - var getUserAgentMethod = typeof(SparkHttpConnection).GetMethod("GetUserAgent", - BindingFlags.NonPublic | BindingFlags.Instance); - + var getUserAgentMethod = typeof(SparkHttpConnection).GetMethod("GetUserAgent", BindingFlags.NonPublic | BindingFlags.Instance); + if (getUserAgentMethod == null) { throw new InvalidOperationException("GetUserAgent method not found"); } - + // Invoke the method return (string)getUserAgentMethod.Invoke(connection, null)!; } } -} \ No newline at end of file +}