Adding custom User-Agent to .NET client (#11)

* Adding custom User-Agent to .NET client * Adding a missing parameter * Renaming UserAgent -> SpiceUserAgent
spiceai · Nov 25, 2024 · c244a9b · c244a9b
1 parent ae9a6ea
commit c244a9b
Show file tree

Hide file tree

Showing 5 changed files with 31 additions and 10 deletions.
diff --git a/Spice/src/Config/UserAgent.cs → Spice/src/Config/SpiceUserAgent.cs b/Spice/src/Config/UserAgent.cs → Spice/src/Config/SpiceUserAgent.cs
@@ -22,9 +22,9 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 
 namespace Spice.Config;
 
-internal static class UserAgent
+public static class SpiceUserAgent
 {
-    public static string agent()
+    public static string agent(string? client = null, string? clientVersion = null, string? clientSystem = null, string? clientExtension = null)
     {
         // get OS type, release and machine type (x86, x64, arm, etc)
         var os = Environment.OSVersion;
@@ -60,7 +60,11 @@ public static string agent()
         // get the runtime version
         var appVersion = typeof(SpiceClient).Assembly.GetName().Version;
 
+        var clientName = client ?? "spice-dotnet";
+        var clientVer = clientVersion ?? appVersion?.ToString();
+        var clientSys = clientSystem ?? $"{osTypeStr}/{osVersion} {osArchStr}";
+        var clientExt = $" {clientExtension}" ?? "";
         // return the user agent string
-        return $"spice-dotnet {appVersion} ({osTypeStr}/{osVersion} {osArchStr})";
+        return $"spice-dotnet/{clientVer} ({clientSys}){clientExt}";
     }
 }
diff --git a/Spice/src/Flight/SpiceFlightClient.cs b/Spice/src/Flight/SpiceFlightClient.cs
@@ -38,7 +38,7 @@ internal class SpiceFlightClient
     private readonly FlightClient _flightClient;
     private readonly AsyncRetryPolicy _retryPolicy;
 
-    private static GrpcChannelOptions GetGrpcChannelOptions(string? appId, string? apiKey)
+    private static GrpcChannelOptions GetGrpcChannelOptions(string? appId, string? apiKey, string? userAgent)
     {
         var options = new GrpcChannelOptions();
 
@@ -54,7 +54,8 @@ private static GrpcChannelOptions GetGrpcChannelOptions(string? appId, string? a
             }
         };
 
-        options.HttpClient.DefaultRequestHeaders.Add("X-Spice-User-Agent", UserAgent.agent());
+        var uaString = userAgent ?? SpiceUserAgent.agent();
+        options.HttpClient.DefaultRequestHeaders.Add("User-Agent", uaString);
 
         return options;
     }
@@ -64,7 +65,7 @@ private static GrpcChannelOptions GetGrpcChannelOptions(string? appId, string? a
         return responseHeaders.Get("authorization") ?? trailers.Get("authorization");
     }
 
-    internal SpiceFlightClient(string address, int maxRetries, string? appId, string? apiKey)
+    internal SpiceFlightClient(string address, int maxRetries, string? appId, string? apiKey, string? userAgent)
     {
         _retryPolicy = Policy.Handle<RpcException>(ex =>
                 ex.Status.StatusCode is StatusCode.Unavailable or StatusCode.DeadlineExceeded or StatusCode.Aborted
@@ -77,7 +78,7 @@ ex.Status.StatusCode is StatusCode.Unavailable or StatusCode.DeadlineExceeded or
                         $"Request failed. Waiting {timespan} before next retry. Retry attempt {retryAttempt}");
                 });
 
-        var options = GetGrpcChannelOptions(appId, apiKey);
+        var options = GetGrpcChannelOptions(appId, apiKey, userAgent);
 
         _flightClient = new FlightClient(GrpcChannel.ForAddress(address, options));
 

diff --git a/Spice/src/SpiceClient.cs b/Spice/src/SpiceClient.cs
@@ -38,6 +38,11 @@ public class SpiceClient
     /// </summary>
     public string? ApiKey { get; internal set; }
 
+    /// <summary>
+    /// Gets or sets the User-Agent string. This property is internal set and can be null.
+    /// </summary>
+    public string? UserAgent { get; internal set; }
+
     /// <summary>
     /// Gets or sets the flight address. This property is internal set and defaults to local flight endpoint.
     /// </summary>
@@ -53,7 +58,7 @@ public class SpiceClient
 
     internal void Init()
     {
-        FlightClient = new SpiceFlightClient(FlightAddress, MaxRetries, AppId, ApiKey);
+        FlightClient = new SpiceFlightClient(FlightAddress, MaxRetries, AppId, ApiKey, UserAgent);
     }
 
     /// <summary>

diff --git a/Spice/src/SpiceClientBuilder.cs b/Spice/src/SpiceClientBuilder.cs
@@ -80,6 +80,17 @@ public SpiceClientBuilder WithMaxRetries(int maxRetries)
         return this;
     }
 
+    /// <summary>
+    /// Sets the client's user agent.
+    /// </summary>
+    /// <param name="userAgent">User agent string</param>
+    /// <returns>The current instance of <see cref="SpiceClientBuilder"/> for method chaining.</returns>
+    public SpiceClientBuilder WithUserAgent(string userAgent)
+    {
+        _spiceClient.UserAgent = userAgent;
+        return this;
+    }
+
     /// <summary>
     /// Initiates <see cref="SpiceClient" /> with provided parameters.
     /// </summary>

diff --git a/SpiceTest/UserAgentTest.cs b/SpiceTest/UserAgentTest.cs
@@ -32,10 +32,10 @@ public async Task TestUserAgent()
     {
         await Task.Run(() =>
         {
-            var agent = UserAgent.agent();
+            var agent = SpiceUserAgent.agent();
 
             // test with a regex
-            Assert.IsTrue(Regex.IsMatch(agent, @"spice-dotnet \d+\.\d+\.\d+\.\d+ \((Unix|Windows|Darwin)/[\d\w\.\-_]+ (x86_64|aarch64|i386)\)"), agent);
+            Assert.IsTrue(Regex.IsMatch(agent, @"spice-dotnet/\d+\.\d+\.\d+\.\d+ \((Unix|Windows|Darwin)/[\d\w\.\-_]+ (x86_64|aarch64|i386)\)"), agent);
         });
     }
 }