Skip to content

Commit 0f70d05

Browse files
committed
add thrift version
1 parent 8057762 commit 0f70d05

2 files changed

Lines changed: 46 additions & 7 deletions

File tree

csharp/src/Drivers/Apache/Spark/SparkHttpConnection.cs

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -255,13 +255,33 @@ protected internal override Task<TRowSet> GetRowSetAsync(TGetPrimaryKeysResp res
255255

256256
private string GetUserAgent()
257257
{
258+
// Build the base user agent string with Thrift version
259+
string thriftVersion = GetThriftVersion();
260+
string thriftComponent = string.IsNullOrEmpty(thriftVersion) ? "Thrift" : $"Thrift/{thriftVersion}";
261+
string baseUserAgent = $"{DriverName.Replace(" ", "")}/{ProductVersionDefault} {thriftComponent}"; // consider using ProductVersion instead of ProductVersionDefault
262+
258263
// Check if a client has provided a user-agent entry
259264
if (Properties.TryGetValue(SparkParameters.UserAgentEntry, out string? userAgentEntry) && !string.IsNullOrWhiteSpace(userAgentEntry))
260265
{
261-
return $"{s_baseUserAgent} {userAgentEntry}";
266+
return $"{baseUserAgent} {userAgentEntry}";
262267
}
263268

264-
return s_baseUserAgent;
269+
return baseUserAgent;
270+
}
271+
272+
private string GetThriftVersion()
273+
{
274+
try
275+
{
276+
var thriftAssembly = typeof(TProtocol).Assembly;
277+
var version = thriftAssembly.GetName().Version;
278+
return version != null ? $"{version.Major}.{version.Minor}.{version.Build}" : "";
279+
}
280+
catch
281+
{
282+
// Return empty string if there's any issue retrieving the assembly version
283+
return "";
284+
}
265285
}
266286
}
267287
}

csharp/test/Drivers/Apache/Spark/SparkHttpConnectionUserAgentTest.cs

Lines changed: 24 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ namespace Apache.Arrow.Adbc.Tests.Drivers.Apache.Spark
3030
public class SparkHttpConnectionUserAgentTest
3131
{
3232
[Fact]
33-
public void UserAgentEntry_WhenNotProvided_UsesBaseUserAgent()
33+
public void UserAgentEntry_WhenNotProvided_UsesBaseUserAgentWithThrift()
3434
{
3535
// Arrange
3636
var properties = new Dictionary<string, string>
@@ -45,7 +45,7 @@ public void UserAgentEntry_WhenNotProvided_UsesBaseUserAgent()
4545
string userAgent = GetUserAgentFromConnection(properties);
4646

4747
// Assert
48-
Assert.Equal("ADBCSparkDriver/1.0.0", userAgent);
48+
Assert.Matches(@"ADBCSparkDriver/[\d\.]+ Thrift(/[\d\.]+)?", userAgent);
4949
}
5050

5151
[Fact]
@@ -65,7 +65,7 @@ public void UserAgentEntry_WhenProvided_AppendsToBaseUserAgent()
6565
string userAgent = GetUserAgentFromConnection(properties);
6666

6767
// Assert
68-
Assert.Equal("ADBCSparkDriver/1.0.0 PowerBI", userAgent);
68+
Assert.Matches(@"ADBCSparkDriver/[\d\.]+ Thrift(/[\d\.]+)? PowerBI", userAgent);
6969
}
7070

7171
[Fact]
@@ -85,7 +85,7 @@ public void UserAgentEntry_WhenEmpty_UsesBaseUserAgent()
8585
string userAgent = GetUserAgentFromConnection(properties);
8686

8787
// Assert
88-
Assert.Equal("ADBCSparkDriver/1.0.0", userAgent);
88+
Assert.Matches(@"ADBCSparkDriver/[\d\.]+ Thrift(/[\d\.]+)?", userAgent);
8989
}
9090

9191
[Fact]
@@ -105,7 +105,26 @@ public void UserAgentEntry_WhenWhitespace_UsesBaseUserAgent()
105105
string userAgent = GetUserAgentFromConnection(properties);
106106

107107
// Assert
108-
Assert.Equal("ADBCSparkDriver/1.0.0", userAgent);
108+
Assert.Matches(@"ADBCSparkDriver/[\d\.]+ Thrift(/[\d\.]+)?", userAgent);
109+
}
110+
111+
[Fact]
112+
public void UserAgent_IncludesThriftComponent()
113+
{
114+
// Arrange
115+
var properties = new Dictionary<string, string>
116+
{
117+
[SparkParameters.Type] = SparkServerTypeConstants.Http,
118+
[SparkParameters.HostName] = "valid.server.com",
119+
[SparkParameters.Path] = "/path",
120+
[SparkParameters.AuthType] = SparkAuthTypeConstants.None
121+
};
122+
123+
// Act
124+
string userAgent = GetUserAgentFromConnection(properties);
125+
126+
// Assert
127+
Assert.Contains("Thrift", userAgent);
109128
}
110129

111130
private string GetUserAgentFromConnection(Dictionary<string, string> properties)

0 commit comments

Comments
 (0)