SNOW-2139007 Clear query context cache before returning session to pool (#1186)

Author: sfc-gh-knozderko

Snowflake.Data.Tests/UnitTests/Session/SessionPoolTest.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Generic;
 using System.Text.RegularExpressions;
 using NUnit.Framework;
 using Snowflake.Data.Client;
@@ -280,6 +281,45 @@ public void TestConnectionCachePoolingDisabledForNewAuthenticators(string connec
             Assert.IsFalse(isSessionReturnedToPool);
         }
 
+        [Test]
+        public void TestShouldClearQueryContextCacheOnReturningToConnectionCache()
+        {
+            // arrange
+            var session = CreateSessionWithCurrentStartTime("account=testAccount;user=testUser;password=testPassword");
+            var pool = SessionPool.CreateSessionCache();
+            var contextElement = new QueryContextElement(123, DateTimeOffset.UtcNow.ToUnixTimeMilliseconds(), 1, "context");
+            var context = new ResponseQueryContext { Entries = new List<ResponseQueryContextElement> { new(contextElement) } };
+            session.UpdateQueryContextCache(context);
+            Assert.AreEqual(1, session.GetQueryContextRequest().Entries.Count);
+
+            // act
+            var isSessionReturnedToPool = pool.AddSession(session, false);
+
+            // assert
+            Assert.IsTrue(isSessionReturnedToPool);
+            Assert.AreEqual(0, session.GetQueryContextRequest().Entries.Count);
+        }
+
+        [Test]
+        public void TestShouldClearQueryContextCacheOnReturningToConnectionPool()
+        {
+            // arrange
+            var connectionString = "account=testAccount;user=testUser;password=testPassword";
+            var session = CreateSessionWithCurrentStartTime(connectionString);
+            var pool = SessionPool.CreateSessionPool(connectionString, null, null, null);
+            var contextElement = new QueryContextElement(123, DateTimeOffset.UtcNow.ToUnixTimeMilliseconds(), 1, "context");
+            var context = new ResponseQueryContext { Entries = new List<ResponseQueryContextElement> { new(contextElement) } };
+            session.UpdateQueryContextCache(context);
+            Assert.AreEqual(1, session.GetQueryContextRequest().Entries.Count);
+
+            // act
+            var isSessionReturnedToPool = pool.AddSession(session, false);
+
+            // assert
+            Assert.IsTrue(isSessionReturnedToPool);
+            Assert.AreEqual(0, session.GetQueryContextRequest().Entries.Count);
+        }
+
         private SFSession CreateSessionWithCurrentStartTime(string connectionString)
         {
             var session = new SFSession(connectionString, new SessionPropertiesContext());

Snowflake.Data/Core/QueryContextCache.cs

@@ -189,6 +189,14 @@ public void ClearCache()
             _logger.Debug($"clearCache() returns. Number of entries in cache now {_cacheSet.Count}");
         }
 
+        public void ClearCacheSafely()
+        {
+            lock (_qccLock)
+            {
+                ClearCache();
+            }
+        }
+
         public void SetCapacity(int cap)
         {
             // check without locking first for performance reason

Snowflake.Data/Core/Session/SFSession.cs

@@ -528,6 +528,14 @@ internal void UpdateSessionParameterMap(List<NameValueParameter> parameterList)
             }
         }
 
+        internal void ClearQueryContextCache()
+        {
+            if (!_disableQueryContextCache)
+            {
+                _queryContextCache.ClearCacheSafely();
+            }
+        }
+
         internal void UpdateQueryContextCache(ResponseQueryContext queryContext)
         {
             if (!_disableQueryContextCache)

Snowflake.Data/Core/Session/SessionPool.cs

@@ -527,6 +527,7 @@ internal bool AddSession(SFSession session, bool ensureMinPoolSize)
                 return false;
             }
 
+            session.ClearQueryContextCache();
             var result = ReturnSessionToPool(session, ensureMinPoolSize);
             var wasSessionReturnedToPool = result.Item1;
             var sessionCreationTokens = result.Item2;