Skip to content

Commit 1a5256f

Browse files
sfc-gh-dszmolkasfc-gh-dszmolka
committed
Merge remote-tracking branch 'origin/SNOW-3344317-check-endpointOverride-scheme' into SNOW-3344317-check-endpointOverride-scheme
2 parents 895fed6 + 2bb3edf commit 1a5256f

File tree

7 files changed

+225
-5
lines changed

7 files changed

+225
-5
lines changed

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
- Removed the io.netty.tryReflectionSetAccessible system property setting as it's no longer needed with modern Arrow/Netty versions (snowflakedb/snowflake-jdbc#2563)
1212
- Fixed crash in getColumns operation when table contained unrecognised column type (snowflakedb/snowflake-jdbc#2568).
1313
- Fixed session expiration when multiple sessions have different heartbeat intervals (snowflakedb/snowflake-jdbc#2566).
14+
- Merge QueryContext from failed query responses
1415

1516
- v4.0.2
1617
- Fix expired session token renewal when polling results (snowflakedb/snowflake-jdbc#2489)

‎src/main/java/net/snowflake/client/internal/core/SFSession.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -129,7 +129,7 @@ public class SFSession extends SFBaseSession {
129129
private Telemetry telemetryClient;
130130
private SnowflakeConnectString sfConnStr;
131131
// The cache of query context sent from Cloud Service.
132-
private QueryContextCache qcc;
132+
QueryContextCache qcc;
133133

134134
// Max retries for outgoing http requests.
135135
private int maxHttpRetries = 7;

‎src/main/java/net/snowflake/client/internal/core/StmtUtil.java‎

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -443,6 +443,16 @@ public static StmtOutput execute(
443443
}
444444
}
445445

446+
static void updateQueryContextFromResponse(JsonNode responseJson, SFBaseSession session) {
447+
if (session == null) {
448+
return;
449+
}
450+
JsonNode queryContextNode = responseJson.path("data").path("queryContext");
451+
if (SnowflakeUtil.isJsonNodePresent(queryContextNode)) {
452+
session.setQueryContext(queryContextNode.toString());
453+
}
454+
}
455+
446456
private static void setServiceNameHeader(StmtInput stmtInput, HttpRequestBase httpRequest) {
447457
if (!isNullOrEmpty(stmtInput.serviceName)) {
448458
httpRequest.setHeader(SessionUtil.SF_HEADER_SERVICE_NAME, stmtInput.serviceName);
@@ -505,11 +515,12 @@ private static StmtOutput pollForOutput(
505515
}
506516
} else {
507517
retries = 0; // reset retry counter after a successful response
508-
}
509518

510-
if (pingPongResponseJson != null)
511-
// raise server side error as an exception if any
512-
{
519+
// Merge QueryContext before error checking so that QCC is updated
520+
// even for failed queries (e.g. DPO changes from aborted transactions).
521+
updateQueryContextFromResponse(pingPongResponseJson, session);
522+
523+
// raise server side error as an exception if any
513524
SnowflakeUtil.checkErrorAndThrowException(pingPongResponseJson);
514525
}
515526

‎src/main/java/net/snowflake/client/internal/jdbc/SnowflakeUtil.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1024,6 +1024,11 @@ public static boolean isNullOrEmpty(String str) {
10241024
return str == null || str.isEmpty();
10251025
}
10261026

1027+
/** Returns {@code true} when the node exists and carries a non-null value. */
1028+
public static boolean isJsonNodePresent(JsonNode node) {
1029+
return !node.isMissingNode() && !node.isNull();
1030+
}
1031+
10271032
/**
10281033
* Converts Byte array to hex string
10291034
*

‎src/test/java/net/snowflake/client/internal/core/StmtUtilTest.java‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,13 @@
11
package net.snowflake.client.internal.core;
22

3+
import static org.junit.jupiter.api.Assertions.assertEquals;
4+
import static org.junit.jupiter.api.Assertions.assertNotNull;
35
import static org.mockito.Mockito.any;
46
import static org.mockito.Mockito.mockStatic;
57
import static org.mockito.Mockito.times;
68

9+
import com.fasterxml.jackson.databind.JsonNode;
10+
import com.fasterxml.jackson.databind.ObjectMapper;
711
import java.util.HashMap;
812
import java.util.Map;
913
import java.util.Map.Entry;
@@ -92,6 +96,37 @@ public void testForwardedHeaders() throws Throwable {
9296
}
9397
}
9498

99+
private static final ObjectMapper MAPPER = new ObjectMapper();
100+
101+
/** SNOW-3063492 Verify QCC is merged from a failed query response */
102+
@Test
103+
public void testUpdateQueryContextFromFailedResponse() throws Exception {
104+
String failedResponse =
105+
"{\"data\":{\"errorCode\":\"200001\",\"sqlState\":\"22000\","
106+
+ "\"queryId\":\"test-query-id\","
107+
+ "\"queryContext\":{\"entries\":[{\"id\":0,\"timestamp\":123456789,\"priority\":0,\"context\":\"opaque\"}]}"
108+
+ "},\"code\":\"200001\","
109+
+ "\"message\":\"A primary key already exists.\","
110+
+ "\"success\":false}";
111+
112+
JsonNode responseJson = MAPPER.readTree(failedResponse);
113+
SFSession session = createSessionWithQCC();
114+
115+
StmtUtil.updateQueryContextFromResponse(responseJson, session);
116+
117+
QueryContextDTO qcc = session.getQueryContextDTO();
118+
assertNotNull(qcc, "QCC should be populated from failed response");
119+
assertEquals(1, qcc.getEntries().size());
120+
assertEquals(0, qcc.getEntries().get(0).getId());
121+
assertEquals(123456789L, qcc.getEntries().get(0).getTimestamp());
122+
}
123+
124+
private static SFSession createSessionWithQCC() {
125+
SFSession session = new SFSession();
126+
session.qcc = new QueryContextCache(session.getQueryContextCacheSize());
127+
return session;
128+
}
129+
95130
private SFLoginInput createLoginInput() {
96131
SFLoginInput input = new SFLoginInput();
97132
input.setServerUrl("MOCK_TEST_HOST");

‎src/test/java/net/snowflake/client/internal/jdbc/QueryContextWiremockIT.java‎

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
package net.snowflake.client.internal.jdbc;
2+
3+
import static org.junit.jupiter.api.Assertions.assertEquals;
4+
import static org.junit.jupiter.api.Assertions.assertNotNull;
5+
import static org.junit.jupiter.api.Assertions.assertThrows;
6+
7+
import java.sql.Connection;
8+
import java.sql.DriverManager;
9+
import java.sql.SQLException;
10+
import java.sql.Statement;
11+
import java.util.List;
12+
import java.util.Properties;
13+
import net.snowflake.client.category.TestTags;
14+
import net.snowflake.client.internal.api.implementation.connection.SnowflakeConnectionImpl;
15+
import net.snowflake.client.internal.core.QueryContextDTO;
16+
import net.snowflake.client.internal.core.QueryContextEntryDTO;
17+
import net.snowflake.client.internal.core.SFSession;
18+
import org.junit.jupiter.api.Tag;
19+
import org.junit.jupiter.api.Test;
20+
21+
@Tag(TestTags.OTHERS)
22+
public class QueryContextWiremockIT extends BaseWiremockTest {
23+
24+
private static final String QCC_FAILED_QUERY_MAPPING =
25+
"/wiremock/mappings/querycontext/qcc-merge-on-failed-query.json";
26+
27+
/**
28+
* SNOW-3063492: Verify that QueryContext is merged from the server response even when the query
29+
* fails. The WireMock stub returns a duplicate primary key error with queryContext entries.
30+
* Before the fix, the driver would throw the exception without merging QCC.
31+
*/
32+
@Test
33+
public void testQueryContextMergedOnFailedQuery() throws Exception {
34+
importMappingFromResources(QCC_FAILED_QUERY_MAPPING);
35+
36+
Properties props = getWiremockProps();
37+
String connectStr = String.format("jdbc:snowflake://%s:%s", WIREMOCK_HOST, wiremockHttpPort);
38+
39+
Connection conn = DriverManager.getConnection(connectStr, props);
40+
SFSession session = conn.unwrap(SnowflakeConnectionImpl.class).getSfSession();
41+
Statement stmt = conn.createStatement();
42+
43+
assertThrows(
44+
SQLException.class,
45+
() -> stmt.executeQuery("INSERT INTO hybrid_test VALUES ('key1', 'duplicate')"));
46+
47+
QueryContextDTO qcc = session.getQueryContextDTO();
48+
assertNotNull(qcc, "QueryContext should have been merged from the failed response");
49+
50+
List<QueryContextEntryDTO> entries = qcc.getEntries();
51+
assertNotNull(entries);
52+
assertEquals(2, entries.size());
53+
54+
assertEquals(0, entries.get(0).getId());
55+
assertEquals(1775206502642407L, entries.get(0).getTimestamp());
56+
assertEquals(0, entries.get(0).getPriority());
57+
58+
assertEquals(69924, entries.get(1).getId());
59+
assertEquals(1775206502558790L, entries.get(1).getTimestamp());
60+
assertEquals(1, entries.get(1).getPriority());
61+
}
62+
63+
private static Properties getWiremockProps() {
64+
Properties props = new Properties();
65+
props.put("account", "testaccount");
66+
props.put("user", "testuser");
67+
props.put("password", "testpassword");
68+
props.put("warehouse", "testwh");
69+
props.put("database", "testdb");
70+
props.put("schema", "testschema");
71+
props.put("ssl", "off");
72+
props.put("insecureMode", "true");
73+
return props;
74+
}
75+
}

‎src/test/resources/wiremock/mappings/querycontext/qcc-merge-on-failed-query.json‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
{
2+
"mappings": [
3+
{
4+
"scenarioName": "QCC merge on failed query",
5+
"requiredScenarioState": "Started",
6+
"newScenarioState": "Logged in",
7+
"request": {
8+
"urlPathPattern": "/session/v1/login-request.*",
9+
"method": "POST"
10+
},
11+
"response": {
12+
"status": 200,
13+
"jsonBody": {
14+
"data": {
15+
"masterToken": "master token",
16+
"token": "session token",
17+
"validityInSeconds": 3600,
18+
"masterValidityInSeconds": 14400,
19+
"displayUserName": "TEST_USER",
20+
"serverVersion": "8.48.0",
21+
"firstLogin": false,
22+
"remMeToken": null,
23+
"remMeValidityInSeconds": 0,
24+
"healthCheckInterval": 45,
25+
"newClientForUpgrade": "3.12.3",
26+
"sessionId": 1172562260498,
27+
"parameters": [
28+
{
29+
"name": "AUTOCOMMIT",
30+
"value": true
31+
}
32+
],
33+
"sessionInfo": {
34+
"databaseName": "TEST_DB",
35+
"schemaName": "TEST_SCHEMA",
36+
"warehouseName": "TEST_WH",
37+
"roleName": "ANALYST"
38+
},
39+
"idToken": null,
40+
"idTokenValidityInSeconds": 0,
41+
"responseData": null,
42+
"mfaToken": null,
43+
"mfaTokenValidityInSeconds": 0
44+
},
45+
"code": null,
46+
"message": null,
47+
"success": true
48+
}
49+
}
50+
},
51+
{
52+
"scenarioName": "QCC merge on failed query",
53+
"requiredScenarioState": "Logged in",
54+
"request": {
55+
"urlPathPattern": "/queries/v1/query-request.*",
56+
"method": "POST"
57+
},
58+
"response": {
59+
"status": 200,
60+
"headers": {
61+
"Content-Type": "application/json"
62+
},
63+
"jsonBody": {
64+
"data": {
65+
"errorCode": "200001",
66+
"age": 0,
67+
"sqlState": "22000",
68+
"queryId": "01c37557-c810-44d1-0000-011109f5513e",
69+
"queryContext": {
70+
"entries": [
71+
{
72+
"id": 0,
73+
"timestamp": 1775206502642407,
74+
"priority": 0,
75+
"context": "CMamhAI="
76+
},
77+
{
78+
"id": 69924,
79+
"timestamp": 1775206502558790,
80+
"priority": 1
81+
}
82+
]
83+
}
84+
},
85+
"code": "200001",
86+
"message": "A primary key already exists.",
87+
"success": false,
88+
"headers": null
89+
}
90+
}
91+
}
92+
]
93+
}

0 commit comments

Comments
 (0)