Support command level URL query string parameters (#510)

DarkWanderer · Aug 25, 2024 · 4e199c6 · 4e199c6
1 parent a564737
commit 4e199c6
Show file tree

Hide file tree

Showing 6 changed files with 170 additions and 19 deletions.
diff --git a/ClickHouse.Client.Tests/Misc/UriBuilderTests.cs b/ClickHouse.Client.Tests/Misc/UriBuilderTests.cs
@@ -14,12 +14,16 @@ public void ShouldSetUriParametersCorrectly()
         var builder = new ClickHouseUriBuilder(new Uri("http://some.server:123"))
         {
             Database = "DATABASE",
-            CustomParameters = new Dictionary<string, object> { { "a", 1 }, { "b", "c" } },
+            ConnectionQueryStringParameters = new Dictionary<string, object> { { "a", 1 }, { "b", "c" } },
+            CommandQueryStringParameters = new Dictionary<string, object> { { "c", 1 }, { "d", "c" } },
             UseCompression = false,
             Sql = "SELECT 1",
-            SessionId = "SESSION"
+            SessionId = "SESSION",
+            QueryId = "QUERY",
         };
 
+        builder.AddSqlQueryParameter("sqlParameterName", "sqlParameterValue");
+
         var result = new Uri(builder.ToString());
         var @params = HttpUtility.ParseQueryString(result.Query);
 
@@ -29,8 +33,129 @@ public void ShouldSetUriParametersCorrectly()
         Assert.AreEqual("SELECT 1", @params.Get("query"));
         Assert.AreEqual("1", @params.Get("a"));
         Assert.AreEqual("c", @params.Get("b"));
+        Assert.AreEqual("1", @params.Get("c"));
+        Assert.AreEqual("c", @params.Get("d"));
         Assert.AreEqual("SESSION", @params.Get("session_id"));
         Assert.AreEqual("false", @params.Get("enable_http_compression"));
+        Assert.AreEqual("QUERY", @params.Get("query_id"));
+        Assert.AreEqual("sqlParameterValue", @params.Get("param_sqlParameterName"));
+    }
+
+    [Test]
+    public void CommandQueryStringParametersShouldOverrideConnectionParameters()
+    {
+        var builder = new ClickHouseUriBuilder(new Uri("http://some.server:123"))
+        {
+            ConnectionQueryStringParameters = new Dictionary<string, object> { { "a", 1 } },
+            CommandQueryStringParameters = new Dictionary<string, object> { { "a", 2 } },
+        };
+
+        var result = new Uri(builder.ToString());
+        var @params = HttpUtility.ParseQueryString(result.Query);
+
+        Assert.AreEqual("2", @params.Get("a"));
+    }
+
+    [Test]
+    public void ConnectionQueryStringParametersShouldOverrideCommonParameters()
+    {
+        var builder = new ClickHouseUriBuilder(new Uri("http://some.server:123"))
+        {
+            Database = "DATABASE",
+            UseCompression = false,
+            Sql = "SELECT 1",
+            SessionId = "SESSION",
+            QueryId = "QUERY",
+            ConnectionQueryStringParameters = new Dictionary<string, object>
+            {
+                { "database", "overrided" },
+                { "enable_http_compression", "overrided" },
+                { "query", "overrided" },
+                { "session_id", "overrided" },
+                { "query_id", "overrided" },
+            },
+        };
+
+        builder.AddSqlQueryParameter("sqlParameterName", "sqlParameterValue");
+
+        var result = new Uri(builder.ToString());
+        var @params = HttpUtility.ParseQueryString(result.Query);
+
+        Assert.AreEqual("overrided", @params.Get("database"));
+        Assert.AreEqual("overrided", @params.Get("enable_http_compression"));
+        Assert.AreEqual("overrided", @params.Get("query"));
+        Assert.AreEqual("overrided", @params.Get("session_id"));
+        Assert.AreEqual("overrided", @params.Get("query_id"));
+    }
+
+    [Test]
+    public void ConnectionQueryStringParametersShouldOverrideSqlQueryParameters()
+    {
+        var builder = new ClickHouseUriBuilder(new Uri("http://some.server:123"))
+        {
+            ConnectionQueryStringParameters = new Dictionary<string, object>
+            {
+                { "param_sqlParameterName", "overrided" },
+            },
+        };
+
+        builder.AddSqlQueryParameter("sqlParameterName", "sqlParameterValue");
+
+        var result = new Uri(builder.ToString());
+        var @params = HttpUtility.ParseQueryString(result.Query);
+
+        Assert.AreEqual("overrided", @params.Get("param_sqlParameterName"));
+    }
+
+    [Test]
+    public void CommandQueryStringParametersShouldOverrideCommonParameters()
+    {
+        var builder = new ClickHouseUriBuilder(new Uri("http://some.server:123"))
+        {
+            Database = "DATABASE",
+            UseCompression = false,
+            Sql = "SELECT 1",
+            SessionId = "SESSION",
+            QueryId = "QUERY",
+            CommandQueryStringParameters = new Dictionary<string, object>
+            {
+                { "database", "overrided" },
+                { "enable_http_compression", "overrided" },
+                { "query", "overrided" },
+                { "session_id", "overrided" },
+                { "query_id", "overrided" },
+            },
+        };
+
+        builder.AddSqlQueryParameter("sqlParameterName", "sqlParameterValue");
+
+        var result = new Uri(builder.ToString());
+        var @params = HttpUtility.ParseQueryString(result.Query);
+
+        Assert.AreEqual("overrided", @params.Get("database"));
+        Assert.AreEqual("overrided", @params.Get("enable_http_compression"));
+        Assert.AreEqual("overrided", @params.Get("query"));
+        Assert.AreEqual("overrided", @params.Get("session_id"));
+        Assert.AreEqual("overrided", @params.Get("query_id"));
+    }
+
+    [Test]
+    public void CommandQueryStringParametersShouldOverrideSqlQueryParameters()
+    {
+        var builder = new ClickHouseUriBuilder(new Uri("http://some.server:123"))
+        {
+            CommandQueryStringParameters = new Dictionary<string, object>
+            {
+                { "param_sqlParameterName", "overrided" },
+            },
+        };
+
+        builder.AddSqlQueryParameter("sqlParameterName", "sqlParameterValue");
+
+        var result = new Uri(builder.ToString());
+        var @params = HttpUtility.ParseQueryString(result.Query);
+
+        Assert.AreEqual("overrided", @params.Get("param_sqlParameterName"));
     }
 #endif
 }
diff --git a/ClickHouse.Client.Tests/SQL/SqlSimpleSelectTests.cs b/ClickHouse.Client.Tests/SQL/SqlSimpleSelectTests.cs
@@ -18,11 +18,9 @@ namespace ClickHouse.Client.Tests.SQL;
 public class SqlSimpleSelectTests : IDisposable
 {
     private readonly ClickHouseConnection connection;
-    private readonly bool useCompression;
 
     public SqlSimpleSelectTests(bool useCompression)
     {
-        this.useCompression = useCompression;
         connection = TestUtilities.GetTestClickHouseConnection(useCompression);
     }
 
@@ -217,9 +215,8 @@ public async Task ShouldGetElapsedQueryStats()
     [FromVersion(23, 7)]
     public async Task ShouldGetResultQueryStats()
     {
-        using var bufferingConnection = TestUtilities.GetTestClickHouseConnection(useCompression);
-        bufferingConnection.CustomSettings.Add("wait_end_of_query", 1);
-        var command = bufferingConnection.CreateCommand();
+        var command = connection.CreateCommand();
+        command.CustomSettings.Add("wait_end_of_query", 1);
         command.CommandText = "SELECT * FROM system.numbers LIMIT 100";
         using var reader = await command.ExecuteReaderAsync();
         var stats = command.QueryStats;

diff --git a/ClickHouse.Client/ADO/ClickHouseCommand.cs b/ClickHouse.Client/ADO/ClickHouseCommand.cs
@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Generic;
 using System.Data;
 using System.Data.Common;
 using System.Linq;
@@ -22,6 +23,7 @@ public class ClickHouseCommand : DbCommand, IClickHouseCommand, IDisposable
 {
     private readonly CancellationTokenSource cts = new CancellationTokenSource();
     private readonly ClickHouseParameterCollection commandParameters = new ClickHouseParameterCollection();
+    private Dictionary<string, object> customSettings;
     private ClickHouseConnection connection;
 
     public ClickHouseCommand()
@@ -52,6 +54,12 @@ public ClickHouseCommand(ClickHouseConnection connection)
 
     public QueryStats QueryStats { get; private set; }
 
+    /// <summary>
+    /// Gets collection of custom settings which will be passed as URL query string parameters.
+    /// </summary>
+    /// <remarks>Not thread-safe.</remarks>
+    public IDictionary<string, object> CustomSettings => customSettings ??= new Dictionary<string, object>();
+
     protected override DbConnection DbConnection
     {
         get => connection;
@@ -156,8 +164,8 @@ private async Task<HttpResponseMessage> PostSqlQueryAsync(string sqlQuery, Cance
         var uriBuilder = connection.CreateUriBuilder();
         await connection.EnsureOpenAsync().ConfigureAwait(false); // Preserve old behavior
 
-        if (!string.IsNullOrEmpty(QueryId))
-            uriBuilder.CustomParameters.Add("query_id", QueryId);
+        uriBuilder.QueryId = QueryId;
+        uriBuilder.CommandQueryStringParameters = customSettings;
 
         using var postMessage = connection.UseFormDataParameters
             ? BuildHttpRequestMessageWithFormData(
@@ -186,7 +194,9 @@ private HttpRequestMessage BuildHttpRequestMessageWithQueryParams(string sqlQuer
             sqlQuery = commandParameters.ReplacePlaceholders(sqlQuery);
             foreach (ClickHouseDbParameter parameter in commandParameters)
             {
-                uriBuilder.AddQueryParameter(parameter.ParameterName, HttpParameterFormatter.Format(parameter, connection.TypeSettings));
+                uriBuilder.AddSqlQueryParameter(
+                    parameter.ParameterName,
+                    HttpParameterFormatter.Format(parameter, connection.TypeSettings));
             }
         }
 

diff --git a/ClickHouse.Client/ADO/ClickHouseConnection.cs b/ClickHouse.Client/ADO/ClickHouseConnection.cs
@@ -355,7 +355,8 @@ internal static Version ParseVersion(string versionString)
         Database = database,
         SessionId = session,
         UseCompression = UseCompression,
-        CustomParameters = customSettings.ToDictionary(kvp => kvp.Key, kvp => kvp.Value),
+        ConnectionQueryStringParameters = customSettings
+            .ToDictionary(kvp => kvp.Key, kvp => kvp.Value),
         Sql = sql,
     };
 

diff --git a/ClickHouse.Client/ClickHouseUriBuilder.cs b/ClickHouse.Client/ClickHouseUriBuilder.cs
@@ -8,7 +8,7 @@ namespace ClickHouse.Client;
 
 internal class ClickHouseUriBuilder
 {
-    private readonly IDictionary<string, string> queryParameters = new Dictionary<string, string>();
+    private readonly IDictionary<string, string> sqlQueryParameters = new Dictionary<string, string>();
 
     public ClickHouseUriBuilder(Uri baseUri)
     {
@@ -25,29 +25,44 @@ public ClickHouseUriBuilder(Uri baseUri)
 
     public string SessionId { get; set; }
 
+    public string QueryId { get; set; }
+
     public static string DefaultFormat => "RowBinaryWithNamesAndTypes";
 
-    public IDictionary<string, object> CustomParameters { get; set; }
+    public IDictionary<string, object> ConnectionQueryStringParameters { get; set; }
+
+    public IDictionary<string, object> CommandQueryStringParameters { get; set; }
 
-    public bool AddQueryParameter(string name, string value) => DictionaryExtensions.TryAdd(queryParameters, name, value);
+    public bool AddSqlQueryParameter(string name, string value) =>
+        DictionaryExtensions.TryAdd(sqlQueryParameters, name, value);
 
     public override string ToString()
     {
         var parameters = HttpUtility.ParseQueryString(string.Empty); // NameValueCollection but a special one
-        parameters.Set("enable_http_compression", UseCompression.ToString(CultureInfo.InvariantCulture).ToLowerInvariant());
+        parameters.Set(
+            "enable_http_compression",
+            UseCompression.ToString(CultureInfo.InvariantCulture).ToLowerInvariant());
         parameters.Set("default_format", DefaultFormat);
         parameters.SetOrRemove("database", Database);
         parameters.SetOrRemove("session_id", SessionId);
         parameters.SetOrRemove("query", Sql);
+        parameters.SetOrRemove("query_id", QueryId);
 
-        foreach (var parameter in queryParameters)
+        foreach (var parameter in sqlQueryParameters)
             parameters.Set("param_" + parameter.Key, parameter.Value);
 
-        if (CustomParameters != null)
+        if (ConnectionQueryStringParameters != null)
         {
-            foreach (var parameter in CustomParameters)
+            foreach (var parameter in ConnectionQueryStringParameters)
                 parameters.Set(parameter.Key, Convert.ToString(parameter.Value, CultureInfo.InvariantCulture));
         }
+
+        if (CommandQueryStringParameters != null)
+        {
+            foreach (var parameter in CommandQueryStringParameters)
+                parameters.Set(parameter.Key, Convert.ToString(parameter.Value, CultureInfo.InvariantCulture));
+        }
+
         var uriBuilder = new UriBuilder(BaseUri) { Query = parameters.ToString() };
         return uriBuilder.ToString();
     }

diff --git a/ClickHouse.Client/IClickHouseCommand.cs b/ClickHouse.Client/IClickHouseCommand.cs
@@ -1,4 +1,5 @@
-using System.Data;
+using System.Collections.Generic;
+using System.Data;
 using System.Threading;
 using System.Threading.Tasks;
 using ClickHouse.Client.ADO;
@@ -11,4 +12,6 @@ public interface IClickHouseCommand : IDbCommand
     new ClickHouseDbParameter CreateParameter();
 
     Task<ClickHouseRawResult> ExecuteRawResultAsync(CancellationToken cancellationToken);
+
+    IDictionary<string, object> CustomSettings { get; }
 }