From af8df64997166ff3defce3c44060af1f8bdb2569 Mon Sep 17 00:00:00 2001
From: Pavel Kravtsov <kravtsovpo@gmail.com>
Date: Mon, 3 Jun 2024 00:28:27 +0700
Subject: [PATCH] Add session parameters

---
 .../ClickHouseConnectionBrokerStub.cs            |  5 +++++
 .../Context/ClickHouseConnectionBroker.cs        | 16 ++++++++++++----
 .../Context/ClickHouseContext.cs                 |  5 +++++
 .../Context/IClickHouseConnectionBroker.cs       |  2 ++
 4 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/src/ClickHouse.Facades.Testing/ClickHouseConnectionBrokerStub.cs b/src/ClickHouse.Facades.Testing/ClickHouseConnectionBrokerStub.cs
index 17c7610..1cf48b2 100644
--- a/src/ClickHouse.Facades.Testing/ClickHouseConnectionBrokerStub.cs
+++ b/src/ClickHouse.Facades.Testing/ClickHouseConnectionBrokerStub.cs
@@ -110,4 +110,9 @@ public Task<long> BulkInsertAsync(
 	{
 		throw new NotImplementedException();
 	}
+
+	public Task SetSessionParameter(string parameterName, object value)
+	{
+		return Task.CompletedTask;
+	}
 }
diff --git a/src/ClickHouse.Facades/Context/ClickHouseConnectionBroker.cs b/src/ClickHouse.Facades/Context/ClickHouseConnectionBroker.cs
index d1e7f59..7e479fd 100644
--- a/src/ClickHouse.Facades/Context/ClickHouseConnectionBroker.cs
+++ b/src/ClickHouse.Facades/Context/ClickHouseConnectionBroker.cs
@@ -52,7 +52,7 @@ public async Task<object> ExecuteScalarAsync(
 		Dictionary<string, object>? parameters,
 		CancellationToken cancellationToken)
 	{
-		ThrowIfNotConnected();
+		cancellationToken.ThrowIfCancellationRequested();
 
 		await using var command = CreateCommand();
 		command.CommandText = query;
@@ -66,7 +66,6 @@ public async Task<int> ExecuteNonQueryAsync(
 		Dictionary<string, object>? parameters,
 		CancellationToken cancellationToken)
 	{
-		ThrowIfNotConnected();
 		cancellationToken.ThrowIfCancellationRequested();
 
 		await using var command = CreateCommand();
@@ -81,7 +80,6 @@ public async Task<DbDataReader> ExecuteReaderAsync(
 		Dictionary<string, object>? parameters,
 		CancellationToken cancellationToken)
 	{
-		ThrowIfNotConnected();
 		cancellationToken.ThrowIfCancellationRequested();
 
 		await using var command = CreateCommand();
@@ -96,7 +94,6 @@ public DataTable ExecuteDataTable(
 		Dictionary<string, object>? parameters,
 		CancellationToken cancellationToken)
 	{
-		ThrowIfNotConnected();
 		cancellationToken.ThrowIfCancellationRequested();
 
 		using var command = CreateCommand();
@@ -154,6 +151,17 @@ public async Task<long> BulkInsertAsync(
 		return bulkCopyInterface.RowsWritten;
 	}
 
+	public Task SetSessionParameter(string parameterName, object value)
+	{
+		if (!_sessionEnabled)
+		{
+			throw new InvalidOperationException(
+				"Unable to set session parameter while sessions are not enabled in the current context.");
+		}
+
+		return _connection.ExecuteStatementAsync($"set {parameterName} = '{value}'");
+	}
+
 	private void SetParameters(ClickHouseCommand command, Dictionary<string, object>? parameters)
 	{
 		if (parameters == null)
diff --git a/src/ClickHouse.Facades/Context/ClickHouseContext.cs b/src/ClickHouse.Facades/Context/ClickHouseContext.cs
index 7bcfe20..d7b11f4 100644
--- a/src/ClickHouse.Facades/Context/ClickHouseContext.cs
+++ b/src/ClickHouse.Facades/Context/ClickHouseContext.cs
@@ -87,6 +87,11 @@ public void ChangeDatabase(string databaseName)
 		_connection!.ChangeDatabase(databaseName);
 	}
 
+	public Task SetSessionParameter(string parameterName, object value)
+	{
+		return _connectionBroker.SetSessionParameter(parameterName, value);
+	}
+
 	internal void Initialize(ClickHouseContextOptions<TContext> options)
 	{
 		ThrowIfInitialized();
diff --git a/src/ClickHouse.Facades/Context/IClickHouseConnectionBroker.cs b/src/ClickHouse.Facades/Context/IClickHouseConnectionBroker.cs
index e76a97f..d3da274 100644
--- a/src/ClickHouse.Facades/Context/IClickHouseConnectionBroker.cs
+++ b/src/ClickHouse.Facades/Context/IClickHouseConnectionBroker.cs
@@ -39,4 +39,6 @@ Task<long> BulkInsertAsync(
 		int batchSize,
 		int maxDegreeOfParallelism,
 		IReadOnlyCollection<string>? columnNames = null);
+
+	Task SetSessionParameter(string parameterName, object value);
 }