fix: FIR-29665: validate set statement (#67)

Co-authored-by: Stepan Burlakov <stepansergeevitch@gmail.com>
firebolt-db · Jan 31, 2024 · ccab497 · ccab497
1 parent b4a9de1
commit ccab497
Show file tree

Hide file tree

Showing 5 changed files with 75 additions and 7 deletions.
diff --git a/FireboltDotNetSdk.Tests/Integration/ExecuteTest.cs b/FireboltDotNetSdk.Tests/Integration/ExecuteTest.cs
@@ -270,6 +270,50 @@ private void ExecuteDateTimeTest(string connString, string commandText, string[]
             conn.Close();
         }
 
+        [TestCase("SET time_zone=Asia/Calcutta")]
+        [TestCase("set time_zone=Asia/Calcutta")]
+        public void SetGoodParameterTest(string setCommand)
+        {
+            using var conn = new FireboltConnection(USER_CONNECTION_STRING);
+            conn.Open();
+            HashSet<string> expectedParamerters = new HashSet<string>() { "time_zone=Asia/Calcutta" };
+            DbCommand command = conn.CreateCommand();
+            command.CommandText = setCommand;
+            command.ExecuteNonQuery();
+            Assert.That(conn.SetParamList, Is.EqualTo(expectedParamerters));
+            conn.Close();
+        }
+
+        [Test]
+        public void SetWrongParameterTest()
+        {
+            using var conn = new FireboltConnection(USER_CONNECTION_STRING);
+            conn.Open();
+            DbCommand command = conn.CreateCommand();
+            command.CommandText = "SET foo=bar";
+            string errorMessage = Assert.Throws<FireboltException>(() => command.ExecuteNonQuery()).Message;
+            Assert.That(errorMessage, Does.Contain("parameter foo is not allowed"));
+            Assert.That(conn.SetParamList, Is.EqualTo(new HashSet<string>()));
+            conn.Close();
+        }
+
+        [Test]
+        public void SetGoodThenWrongParameterTest()
+        {
+            using var conn = new FireboltConnection(USER_CONNECTION_STRING);
+            conn.Open();
+            HashSet<string> expectedParamerters = new HashSet<string>() { "time_zone=Asia/Calcutta" };
+            DbCommand command = conn.CreateCommand();
+            command.CommandText = "SET time_zone=Asia/Calcutta";
+            command.ExecuteNonQuery();
+            Assert.That(conn.SetParamList, Is.EqualTo(expectedParamerters));
+
+            command.CommandText = "SET foo=bar";
+            string errorMessage = Assert.Throws<FireboltException>(() => command.ExecuteNonQuery()).Message;
+            Assert.That(errorMessage, Does.Contain("parameter foo is not allowed"));
+            Assert.That(conn.SetParamList, Is.EqualTo(expectedParamerters));
+            conn.Close();
+        }
 
         [Test]
         public void ExecuteSelectBoolean()

diff --git a/FireboltDotNetSdk.Tests/Unit/FireboltCommandTest.cs b/FireboltDotNetSdk.Tests/Unit/FireboltCommandTest.cs
@@ -65,7 +65,7 @@ public class FireboltCommandTest
         [TestCase("SET param=1,param=2")]
         public void SetTest(string commandText)
         {
-            var connection = new FireboltConnection(mockConnectionString) { Client = new MockClient(""), EngineUrl = "engine" };
+            var connection = new FireboltConnection(mockConnectionString) { Client = new MockClient("{}"), EngineUrl = "engine" };
             var cs = new FireboltCommand(connection, commandText, new FireboltParameterCollection());
             Assert.IsEmpty(cs.SetParamList);
             cs.ExecuteNonQuery();
@@ -76,7 +76,7 @@ public void SetTest(string commandText)
         [TestCase("SET param=1,param=2")]
         public async Task SetTestAsync(string commandText)
         {
-            var connection = new FireboltConnection(mockConnectionString) { Client = new MockClient(""), EngineUrl = "engine" };
+            var connection = new FireboltConnection(mockConnectionString) { Client = new MockClient("{}"), EngineUrl = "engine" };
             var cs = new FireboltCommand(connection, commandText, new FireboltParameterCollection());
             Assert.IsEmpty(cs.SetParamList);
             await cs.ExecuteNonQueryAsync();
@@ -129,7 +129,7 @@ public void GetOriginalJsonDataTest()
         [TestCase("SET param=1,param=2")]
         public void ClearSetListTest(string commandText)
         {
-            var connection = new FireboltConnection(mockConnectionString) { Client = new MockClient(""), EngineUrl = "engine" };
+            var connection = new FireboltConnection(mockConnectionString) { Client = new MockClient("{}"), EngineUrl = "engine" };
             var cs = new FireboltCommand(connection, commandText, new FireboltParameterCollection());
 
             cs.ExecuteNonQuery();

diff --git a/FireboltDotNetSdk.Tests/Unit/FireboltConnectionTest.cs b/FireboltDotNetSdk.Tests/Unit/FireboltConnectionTest.cs
@@ -390,15 +390,18 @@ public void SuccessfulLogin()
             cs.Client = client;
             FireResponse.LoginResponse loginResponse = new FireResponse.LoginResponse("access_token", "3600", "Bearer");
             FireResponse.GetSystemEngineUrlResponse systemEngineResponse = new FireResponse.GetSystemEngineUrlResponse() { engineUrl = "api.test.firebolt.io" };
+            FireResponse.GetAccountIdByNameResponse accountIdResponse = new FireResponse.GetAccountIdByNameResponse() { id = "account_id" };
             httpClientMock.SetupSequence(p => p.SendAsync(It.IsAny<HttpRequestMessage>(), It.IsAny<CancellationToken>()))
             .ReturnsAsync(FireboltClientTest.GetResponseMessage(loginResponse, HttpStatusCode.OK)) // retrieve access token
             .ReturnsAsync(FireboltClientTest.GetResponseMessage(systemEngineResponse, HttpStatusCode.OK)) // get system engine URL
+            .ReturnsAsync(FireboltClientTest.GetResponseMessage(accountIdResponse, HttpStatusCode.OK)) // get account ID
+            .ReturnsAsync(FireboltClientTest.GetResponseMessage("{\"query\":{\"query_id\": \"1\"},\"meta\":[{\"type\": \"int\", \"name\": \"1\"}],\"data\":[[\"1\"]]}", HttpStatusCode.OK)) // select 1
             ;
             cs.Open(); // should succeed
             // Due to Open does not return value the only way to validate that everything passed well is to validate that SendAsync was called exactly twice:
             // 1. to retrive token
             // 2. to retrieve system engine URL
-            httpClientMock.Verify(m => m.SendAsync(It.IsAny<HttpRequestMessage>(), It.IsAny<CancellationToken>()), Times.Exactly(2));
+            httpClientMock.Verify(m => m.SendAsync(It.IsAny<HttpRequestMessage>(), It.IsAny<CancellationToken>()), Times.Exactly(4));
         }
 
         [Test]
@@ -420,12 +423,13 @@ public void SuccessfulLoginWithEngineName()
             .ReturnsAsync(FireboltClientTest.GetResponseMessage("{\"query\":{\"query_id\": \"1\"},\"meta\":[{\"type\": \"text\", \"name\": \"attached_to\"}],\"data\":[[\"db\"]]}", HttpStatusCode.OK)) // get Engine DB
             .ReturnsAsync(FireboltClientTest.GetResponseMessage("{\"query\":{\"query_id\": \"2\"},\"meta\":[{\"type\": \"t\", \"name\": \"database_name\"}],\"data\":[[\"db\"]]}", HttpStatusCode.OK)) // check whether the DB is accessible
             .ReturnsAsync(FireboltClientTest.GetResponseMessage("{\"query\":{\"query_id\": \"3\"}," + engineUrlMeta + ", " + engineUrlData + "}", HttpStatusCode.OK)) // get engine URL
+            .ReturnsAsync(FireboltClientTest.GetResponseMessage("{\"query\":{\"query_id\": \"1\"},\"meta\":[{\"type\": \"int\", \"name\": \"1\"}],\"data\":[[\"1\"]]}", HttpStatusCode.OK)) // select 1
             .ReturnsAsync(FireboltClientTest.GetResponseMessage("{\"query\":{\"query_id\": \"1\"},\"meta\":[{\"type\": \"string\"}, {\"name\": \"version()\"}],\"data\":[[\"1.2.3\"]]}", HttpStatusCode.OK)) // get version
             ;
             cs.Client = client;
             cs.Open(); // should succeed
             That(cs.ServerVersion, Is.EqualTo("1.2.3"));
-            httpClientMock.Verify(m => m.SendAsync(It.IsAny<HttpRequestMessage>(), It.IsAny<CancellationToken>()), Times.Exactly(7));
+            httpClientMock.Verify(m => m.SendAsync(It.IsAny<HttpRequestMessage>(), It.IsAny<CancellationToken>()), Times.Exactly(8));
         }
 
         [Test]

diff --git a/FireboltNETSDK/Client/FireboltCommand.cs b/FireboltNETSDK/Client/FireboltCommand.cs
@@ -164,10 +164,19 @@ private DbDataReader CreateDbDataReader(QueryResult? queryResult)
                 throw new FireboltException("Client is undefined. Initialize connection properly");
             }
             var engineUrl = Connection?.EngineUrl;
-            if (commandText.Trim().StartsWith("SET"))
+            if (commandText.Trim().ToUpper().StartsWith("SET"))
             {
                 commandText = commandText.Remove(0, 4).Trim();
                 SetParamList.Add(commandText);
+                try
+                {
+                    Connection?.ValidateConnection();
+                }
+                catch (FireboltException e)
+                {
+                    SetParamList.Remove(commandText);
+                    throw e;
+                }
                 return await Task.FromResult<string?>(null);
             }
             string newCommandText = commandText;

diff --git a/FireboltNETSDK/Client/FireboltConnection.cs b/FireboltNETSDK/Client/FireboltConnection.cs
@@ -269,6 +269,7 @@ public override async Task<bool> OpenAsync(CancellationToken cancellationToken)
                 EngineUrl = response.EngineUrl;
                 _isSystem = response.IsSystem;
                 _database = response.Database;
+                ValidateConnection();
                 OnSessionEstablished();
                 return EngineUrl != null;
             }
@@ -279,6 +280,11 @@ public override async Task<bool> OpenAsync(CancellationToken cancellationToken)
             }
         }
 
+        internal void ValidateConnection()
+        {
+            CreateDbCommand("SELECT 1").ExecuteScalar();
+        }
+
         private FireboltClient CreateClient()
         {
             var builder = new FireboltConnectionStringBuilder(_connectionString);
@@ -311,7 +317,12 @@ private FireboltClient CreateClient()
         /// <inheritdoc cref="CreateDbCommand()"/>
         protected override DbCommand CreateDbCommand()
         {
-            return new FireboltCommand(this, null, new FireboltParameterCollection());
+            return CreateDbCommand(null);
+        }
+
+        private DbCommand CreateDbCommand(string? command)
+        {
+            return new FireboltCommand(this, command, new FireboltParameterCollection());
         }
 
         internal void OnSessionEstablished()