Skip to content

Commit a574cd3

Browse files
sfc-gh-ext-simba-lfsfc-gh-ext-simba-lf
committed
Merge branch 'master' of https://github.com/snowflakedb/snowflake-connector-net into NO-SNOW-Fix-Multistatement-Tests
2 parents b89c607 + 0df25bf commit a574cd3

7 files changed

Lines changed: 167 additions & 3 deletions

File tree

‎Snowflake.Data.Tests/Mock/MockRestSessionExpired.cs‎

Lines changed: 45 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,11 @@ class MockRestSessionExpired : IMockRestRequester
1212
{
1313
static private readonly String EXPIRED_SESSION_TOKEN = "session_expired_token";
1414

15+
static private readonly String NEW_SESSION_TOKEN = "new_session_token";
16+
1517
static private readonly String TOKEN_FMT = "Snowflake Token=\"{0}\"";
1618

17-
static private readonly int SESSION_EXPIRED_CODE = 390112;
19+
static internal readonly int SESSION_EXPIRED_CODE = 390112;
1820

1921
public string FirstTimeRequestID;
2022

@@ -54,7 +56,7 @@ public Task<T> PostAsync<T>(IRestRequest request, CancellationToken cancellation
5456
};
5557
return Task.FromResult<T>((T)(object)queryExecResponse);
5658
}
57-
else if (sfRequest.authorizationToken.Equals(String.Format(TOKEN_FMT, "new_session_token")))
59+
else if (sfRequest.authorizationToken.Equals(String.Format(TOKEN_FMT, NEW_SESSION_TOKEN)))
5860
{
5961
SecondTimeRequestID = ExtractRequestID(sfRequest.Url.Query);
6062
QueryExecResponse queryExecResponse = new QueryExecResponse
@@ -93,7 +95,7 @@ public Task<T> PostAsync<T>(IRestRequest request, CancellationToken cancellation
9395
success = true,
9496
data = new RenewSessionResponseData()
9597
{
96-
sessionToken = "new_session_token",
98+
sessionToken = NEW_SESSION_TOKEN,
9799
masterToken = "new_master_token"
98100
}
99101
});
@@ -116,6 +118,46 @@ public T Get<T>(IRestRequest request)
116118

117119
public Task<T> GetAsync<T>(IRestRequest request, CancellationToken cancellationToken)
118120
{
121+
SFRestRequest sfRequest = (SFRestRequest)request;
122+
if (sfRequest.Url.ToString().Contains("retryId"))
123+
{
124+
QueryExecResponse queryExecResponse = new QueryExecResponse
125+
{
126+
success = false,
127+
code = SESSION_EXPIRED_CODE
128+
};
129+
return Task.FromResult<T>((T)(object)queryExecResponse);
130+
}
131+
if (sfRequest.authorizationToken.Equals(String.Format(TOKEN_FMT, EXPIRED_SESSION_TOKEN)))
132+
{
133+
QueryExecResponse queryExecResponse = new QueryExecResponse
134+
{
135+
success = false,
136+
code = SESSION_EXPIRED_CODE
137+
};
138+
return Task.FromResult<T>((T)(object)queryExecResponse);
139+
}
140+
if (sfRequest.authorizationToken.Equals(String.Format(TOKEN_FMT, NEW_SESSION_TOKEN)))
141+
{
142+
QueryExecResponse queryExecResponse = new QueryExecResponse
143+
{
144+
success = true,
145+
data = new QueryExecResponseData
146+
{
147+
rowSet = new string[,] { { "abc" } },
148+
rowType = new List<ExecResponseRowType>()
149+
{
150+
new ExecResponseRowType
151+
{
152+
name = "colOne",
153+
type = SFDataType.TEXT.ToString()
154+
}
155+
},
156+
parameters = new List<NameValueParameter>()
157+
}
158+
};
159+
return Task.FromResult<T>((T)(object)queryExecResponse);
160+
}
119161
return Task.FromResult<T>((T)(object)null);
120162
}
121163

‎Snowflake.Data.Tests/UnitTests/SFStatementTest.cs‎

100755100644
Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
using Snowflake.Data.Core;
44
using NUnit.Framework;
55
using System;
6+
using System.Threading.Tasks;
67
using Snowflake.Data.Core.Session;
78

89
namespace Snowflake.Data.Tests.UnitTests
@@ -29,6 +30,56 @@ public void TestSessionRenew()
2930
Assert.AreEqual(restRequester.FirstTimeRequestID, restRequester.SecondTimeRequestID);
3031
}
3132

33+
[Test]
34+
public void TestSessionRenewGetResultWithId()
35+
{
36+
Mock.MockRestSessionExpired restRequester = new Mock.MockRestSessionExpired();
37+
SFSession sfSession = new SFSession("account=test;user=test;password=test", new SessionPropertiesContext(), restRequester);
38+
sfSession.Open();
39+
SFStatement statement = new SFStatement(sfSession);
40+
SFBaseResultSet resultSet = statement.GetResultWithId("mockId");
41+
Assert.AreEqual(true, resultSet.Next());
42+
Assert.AreEqual("abc", resultSet.GetString(0));
43+
Assert.AreEqual("new_session_token", sfSession.sessionToken);
44+
Assert.AreEqual("new_master_token", sfSession.masterToken);
45+
}
46+
47+
[Test]
48+
public void TestSessionRenewGetResultWithIdOnlyRetries3Times()
49+
{
50+
Mock.MockRestSessionExpired restRequester = new Mock.MockRestSessionExpired();
51+
SFSession sfSession = new SFSession("account=test;user=test;password=test", new SessionPropertiesContext(), restRequester);
52+
sfSession.Open();
53+
SFStatement statement = new SFStatement(sfSession);
54+
var thrown = Assert.Throws<SnowflakeDbException>(() => statement.GetResultWithId("retryId"));
55+
Assert.AreEqual(thrown.ErrorCode, Mock.MockRestSessionExpired.SESSION_EXPIRED_CODE);
56+
}
57+
58+
[Test]
59+
public async Task TestSessionRenewGetResultWithIdAsync()
60+
{
61+
Mock.MockRestSessionExpired restRequester = new Mock.MockRestSessionExpired();
62+
SFSession sfSession = new SFSession("account=test;user=test;password=test", new SessionPropertiesContext(), restRequester);
63+
await sfSession.OpenAsync(CancellationToken.None);
64+
SFStatement statement = new SFStatement(sfSession);
65+
SFBaseResultSet resultSet = await statement.GetResultWithIdAsync("mockId", CancellationToken.None);
66+
Assert.AreEqual(true, resultSet.Next());
67+
Assert.AreEqual("abc", resultSet.GetString(0));
68+
Assert.AreEqual("new_session_token", sfSession.sessionToken);
69+
Assert.AreEqual("new_master_token", sfSession.masterToken);
70+
}
71+
72+
[Test]
73+
public async Task TestSessionRenewGetResultWithIdOnlyRetries3TimesAsync()
74+
{
75+
Mock.MockRestSessionExpired restRequester = new Mock.MockRestSessionExpired();
76+
SFSession sfSession = new SFSession("account=test;user=test;password=test", new SessionPropertiesContext(), restRequester);
77+
await sfSession.OpenAsync(CancellationToken.None);
78+
SFStatement statement = new SFStatement(sfSession);
79+
var thrown = Assert.ThrowsAsync<SnowflakeDbException>(async () => await statement.GetResultWithIdAsync("retryId", CancellationToken.None));
80+
Assert.AreEqual(thrown.ErrorCode, Mock.MockRestSessionExpired.SESSION_EXPIRED_CODE);
81+
}
82+
3283
// Mock test for session renew during query execution
3384
[Test]
3485
public void TestSessionRenewDuringQueryExec()

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

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
using System;
2+
using System.Collections.Generic;
23
using System.Text.RegularExpressions;
34
using NUnit.Framework;
45
using Snowflake.Data.Client;
@@ -280,6 +281,45 @@ public void TestConnectionCachePoolingDisabledForNewAuthenticators(string connec
280281
Assert.IsFalse(isSessionReturnedToPool);
281282
}
282283

284+
[Test]
285+
public void TestShouldClearQueryContextCacheOnReturningToConnectionCache()
286+
{
287+
// arrange
288+
var session = CreateSessionWithCurrentStartTime("account=testAccount;user=testUser;password=testPassword");
289+
var pool = SessionPool.CreateSessionCache();
290+
var contextElement = new QueryContextElement(123, DateTimeOffset.UtcNow.ToUnixTimeMilliseconds(), 1, "context");
291+
var context = new ResponseQueryContext { Entries = new List<ResponseQueryContextElement> { new(contextElement) } };
292+
session.UpdateQueryContextCache(context);
293+
Assert.AreEqual(1, session.GetQueryContextRequest().Entries.Count);
294+
295+
// act
296+
var isSessionReturnedToPool = pool.AddSession(session, false);
297+
298+
// assert
299+
Assert.IsTrue(isSessionReturnedToPool);
300+
Assert.AreEqual(0, session.GetQueryContextRequest().Entries.Count);
301+
}
302+
303+
[Test]
304+
public void TestShouldClearQueryContextCacheOnReturningToConnectionPool()
305+
{
306+
// arrange
307+
var connectionString = "account=testAccount;user=testUser;password=testPassword";
308+
var session = CreateSessionWithCurrentStartTime(connectionString);
309+
var pool = SessionPool.CreateSessionPool(connectionString, null, null, null);
310+
var contextElement = new QueryContextElement(123, DateTimeOffset.UtcNow.ToUnixTimeMilliseconds(), 1, "context");
311+
var context = new ResponseQueryContext { Entries = new List<ResponseQueryContextElement> { new(contextElement) } };
312+
session.UpdateQueryContextCache(context);
313+
Assert.AreEqual(1, session.GetQueryContextRequest().Entries.Count);
314+
315+
// act
316+
var isSessionReturnedToPool = pool.AddSession(session, false);
317+
318+
// assert
319+
Assert.IsTrue(isSessionReturnedToPool);
320+
Assert.AreEqual(0, session.GetQueryContextRequest().Entries.Count);
321+
}
322+
283323
private SFSession CreateSessionWithCurrentStartTime(string connectionString)
284324
{
285325
var session = new SFSession(connectionString, new SessionPropertiesContext());

‎Snowflake.Data/Core/QueryContextCache.cs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -189,6 +189,14 @@ public void ClearCache()
189189
_logger.Debug($"clearCache() returns. Number of entries in cache now {_cacheSet.Count}");
190190
}
191191

192+
public void ClearCacheSafely()
193+
{
194+
lock (_qccLock)
195+
{
196+
ClearCache();
197+
}
198+
}
199+
192200
public void SetCapacity(int cap)
193201
{
194202
// check without locking first for performance reason

‎Snowflake.Data/Core/SFStatement.cs‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,8 @@ class SFStatement
113113

114114
private const int SF_QUERY_IN_PROGRESS_ASYNC = 333334;
115115

116+
private const int GetResultWithIdMaxRetriesCount = 3;
117+
116118
private string _requestId;
117119

118120
private readonly object _requestIdLock = new object();
@@ -573,6 +575,12 @@ internal async Task<SFBaseResultSet> GetResultWithIdAsync(string resultId, Cance
573575
var req = BuildResultRequestWithId(resultId);
574576
QueryExecResponse response = null;
575577
response = await _restRequester.GetAsync<QueryExecResponse>(req, cancellationToken).ConfigureAwait(false);
578+
for (var retryCount = 0; retryCount < GetResultWithIdMaxRetriesCount && SessionExpired(response); retryCount++)
579+
{
580+
await SfSession.renewSessionAsync(cancellationToken).ConfigureAwait(false);
581+
req.authorizationToken = string.Format(SF_AUTHORIZATION_SNOWFLAKE_FMT, SfSession.sessionToken);
582+
response = await _restRequester.GetAsync<QueryExecResponse>(req, cancellationToken).ConfigureAwait(false);
583+
}
576584
return BuildResultSet(response, cancellationToken);
577585
}
578586

@@ -581,6 +589,12 @@ internal SFBaseResultSet GetResultWithId(string resultId)
581589
var req = BuildResultRequestWithId(resultId);
582590
QueryExecResponse response = null;
583591
response = _restRequester.Get<QueryExecResponse>(req);
592+
for (var retryCount = 0; retryCount < GetResultWithIdMaxRetriesCount && SessionExpired(response); retryCount++)
593+
{
594+
SfSession.renewSession();
595+
req.authorizationToken = string.Format(SF_AUTHORIZATION_SNOWFLAKE_FMT, SfSession.sessionToken);
596+
response = _restRequester.Get<QueryExecResponse>(req);
597+
}
584598
return BuildResultSet(response, CancellationToken.None);
585599
}
586600

‎Snowflake.Data/Core/Session/SFSession.cs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -528,6 +528,14 @@ internal void UpdateSessionParameterMap(List<NameValueParameter> parameterList)
528528
}
529529
}
530530

531+
internal void ClearQueryContextCache()
532+
{
533+
if (!_disableQueryContextCache)
534+
{
535+
_queryContextCache.ClearCacheSafely();
536+
}
537+
}
538+
531539
internal void UpdateQueryContextCache(ResponseQueryContext queryContext)
532540
{
533541
if (!_disableQueryContextCache)

‎Snowflake.Data/Core/Session/SessionPool.cs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -527,6 +527,7 @@ internal bool AddSession(SFSession session, bool ensureMinPoolSize)
527527
return false;
528528
}
529529

530+
session.ClearQueryContextCache();
530531
var result = ReturnSessionToPool(session, ensureMinPoolSize);
531532
var wasSessionReturnedToPool = result.Item1;
532533
var sessionCreationTokens = result.Item2;

0 commit comments

Comments
 (0)