SNOW-726736 [HTAP] Add GetQueryId function in DbCommand (#760)

### Description sdk issue121 Add GetQueryId() function in SnowflakeDbCommand, as well as the existing one in SnowflakeDbDataReader.
snowflakedb · Sep 11, 2023 · a686265 · a686265
1 parent c363799
commit a686265
Show file tree

Hide file tree

Showing 3 changed files with 107 additions and 0 deletions.
diff --git a/Snowflake.Data.Tests/IntegrationTests/SFDbCommandIT.cs b/Snowflake.Data.Tests/IntegrationTests/SFDbCommandIT.cs
@@ -829,5 +829,91 @@ public void testExecuteLargeQueryWithGcsDownscopedToken()
                 conn.Close();
             }
         }
+
+        [Test]
+        public void TestGetQueryId()
+        {
+            using (SnowflakeDbConnection conn = new SnowflakeDbConnection())
+            {
+                conn.ConnectionString = ConnectionString;
+                conn.Open();
+
+                // query id is null when no query executed
+                SnowflakeDbCommand command = (SnowflakeDbCommand)conn.CreateCommand();
+                string queryId = command.GetQueryId();
+                Assert.IsNull(queryId);
+
+                // query id from ExecuteNonQuery
+                command.CommandText = "create or replace temporary table testgetqueryid(cola string)";
+                command.ExecuteNonQuery();
+                queryId = command.GetQueryId();
+                Assert.IsNotEmpty(queryId);
+
+                // query id from ExecuteReader
+                command.CommandText = "show tables like 'testgetqueryid'";
+                SnowflakeDbDataReader reader = (SnowflakeDbDataReader)command.ExecuteReader();
+                queryId = command.GetQueryId();
+                Assert.IsNotEmpty(queryId);
+                Assert.AreEqual(queryId, reader.GetQueryId());
+                Assert.IsTrue(reader.Read());
+
+                // query id from insert query
+                command.CommandText = "insert into testgetqueryid values('test')";
+                command.ExecuteNonQuery();
+                queryId = command.GetQueryId();
+                Assert.IsNotEmpty(queryId);
+
+                // query id from select query
+                command.CommandText = "select * from testgetqueryid";
+                reader = (SnowflakeDbDataReader)command.ExecuteReader();
+                queryId = command.GetQueryId();
+                Assert.IsNotEmpty(queryId);
+                Assert.AreEqual(queryId, reader.GetQueryId());
+                Assert.IsTrue(reader.Read());
+                Assert.AreEqual("test", reader.GetString(0));
+
+                // query id from different DbCommand instance
+                SnowflakeDbCommand command2 = (SnowflakeDbCommand)conn.CreateCommand();
+                string queryId2 = command2.GetQueryId();
+                Assert.IsNull(queryId2);
+                command2.CommandText = "select 'test2'";
+                SnowflakeDbDataReader reader2 = (SnowflakeDbDataReader)command2.ExecuteReader();
+                queryId2 = command2.GetQueryId();
+                Assert.IsNotEmpty(queryId2);
+                Assert.AreEqual(queryId2, reader2.GetQueryId());
+                // each DbCommand instance has it's own query Id.
+                Assert.AreNotEqual(queryId2, queryId);
+                Assert.IsTrue(reader2.Read());
+                Assert.AreEqual("test2", reader2.GetString(0));
+
+                // use query Id to get the result
+                command.CommandText = $"select * from table(result_scan('{queryId}'))";
+                reader = (SnowflakeDbDataReader)command.ExecuteReader();
+                Assert.IsTrue(reader.Read());
+                Assert.AreEqual("test", reader.GetString(0));
+
+                command2.CommandText = $"select * from table(result_scan('{queryId2}'))";
+                reader2 = (SnowflakeDbDataReader)command2.ExecuteReader();
+                Assert.IsTrue(reader2.Read());
+                Assert.AreEqual("test2", reader2.GetString(0));
+
+                // query id from failed query
+                command.CommandText = "select * from table_not_exists";
+                try
+                {
+                    reader = (SnowflakeDbDataReader)command.ExecuteReader();
+                    Assert.Fail();
+                }
+                catch (SnowflakeDbException e)
+                {
+                    Assert.AreEqual(2003, e.ErrorCode);
+                }
+
+                queryId = command.GetQueryId();
+                Assert.IsNotEmpty(queryId);
+
+                conn.Close();
+            }
+        }
     }
 }
diff --git a/Snowflake.Data/Client/SnowflakeDbCommand.cs b/Snowflake.Data/Client/SnowflakeDbCommand.cs
@@ -238,6 +238,15 @@ public override void Prepare()
             throw new NotImplementedException();
         }
 
+        public string GetQueryId()
+        {
+            if (sfStatement != null)
+            {
+                return sfStatement.GetQueryId();
+            }
+            return null;
+        }
+
         protected override DbParameter CreateDbParameter()
         {
             return new SnowflakeDbParameter();

diff --git a/Snowflake.Data/Core/SFStatement.cs b/Snowflake.Data/Core/SFStatement.cs
@@ -60,6 +60,9 @@ class SFStatement
 
         string _bindStage = null;
 
+        // the query id of the last query
+        string _lastQueryId = null;
+
         internal SFStatement(SFSession session)
         {
             SfSession = session;
@@ -180,6 +183,10 @@ private void CleanUpCancellationTokenSources()
 
         private SFBaseResultSet BuildResultSet(QueryExecResponse response, CancellationToken cancellationToken)
         {
+            if ((response.data != null) && (response.data.queryId != null))
+            {
+                _lastQueryId = response.data.queryId;
+            }
             if (response.success)
             {
                 if ((response.data.resultIds != null) && (response.data.resultIds.Length > 0))
@@ -786,5 +793,10 @@ await ExecuteAsyncHelper<PutGetExecResponse, PutGetResponseData>(
 
             return fileTransferAgent.result();
         }
+
+        internal string GetQueryId()
+        {
+            return _lastQueryId;
+        }
     }
 }