Skip to content

Commit 717cc27

Browse files
gopalldbclaude
andauthored
Force runAsync=false for metadata SHOW commands in Thrift client (#1258)
## Summary - Metadata SHOW commands routed through `executeStatement()` via `UseQueryForMetadata` were incorrectly executing asynchronously because `getRequest()` always forces `runAsync=true` in production - Added `StatementType` parameter to `getRequest()` and skip the `runAsync=true` override when `statementType == METADATA`, aligning Thrift client behavior with the SEA client - Updated existing test expectations and added new tests verifying `runAsync` is `false` for METADATA and `true` for SQL statements ## Test plan - [x] `DatabricksThriftServiceClientTest` — 43 tests pass (includes 2 new tests) - [x] `DatabricksMetadataQueryClientTest` — 44 tests pass - [x] `DatabricksSessionTest` — 17 tests pass - [ ] Verify metadata SHOW commands execute synchronously in production with `UseQueryForMetadata` enabled 🤖 Generated with [Claude Code](https://claude.com/claude-code) NO_CHANGELOG=true --------- Signed-off-by: Gopal Lal <gopal.lal@databricks.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
1 parent 0531451 commit 717cc27

2 files changed

Lines changed: 70 additions & 5 deletions

File tree

src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftServiceClient.java

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -162,7 +162,8 @@ public DatabricksResultSet executeStatement(
162162

163163
DatabricksThreadContextHolder.setStatementType(statementType);
164164

165-
TExecuteStatementReq request = getRequest(sql, parameters, session, parentStatement, false);
165+
TExecuteStatementReq request =
166+
getRequest(sql, parameters, session, parentStatement, false, statementType);
166167

167168
return thriftAccessor.execute(request, parentStatement, session, statementType);
168169
}
@@ -181,7 +182,8 @@ public DatabricksResultSet executeStatementAsync(
181182
"public DatabricksResultSet executeStatementAsync(String sql = {%s}, Compute cluster = {%s}, Map<Integer, ImmutableSqlParameter> parameters = {%s})",
182183
sql, computeResource.toString(), parameters.toString()));
183184

184-
TExecuteStatementReq request = getRequest(sql, parameters, session, parentStatement, true);
185+
TExecuteStatementReq request =
186+
getRequest(sql, parameters, session, parentStatement, true, StatementType.SQL);
185187

186188
return thriftAccessor.executeAsync(request, parentStatement, session, StatementType.SQL);
187189
}
@@ -204,7 +206,8 @@ private TExecuteStatementReq getRequest(
204206
Map<Integer, ImmutableSqlParameter> parameters,
205207
IDatabricksSession session,
206208
IDatabricksStatementInternal parentStatement,
207-
boolean runAsync)
209+
boolean runAsync,
210+
StatementType statementType)
208211
throws SQLException {
209212
DatabricksThreadContextHolder.setSessionId(session.getSessionId());
210213
TSparkArrowTypes arrowNativeTypes = new TSparkArrowTypes().setTimestampAsArrow(true);
@@ -254,7 +257,9 @@ private TExecuteStatementReq getRequest(
254257
request.setResultRowLimit(maxRows);
255258
}
256259

257-
if (runAsync || !DriverUtil.isRunningAgainstFake()) {
260+
if (statementType == StatementType.METADATA) {
261+
request.setRunAsync(false);
262+
} else if (runAsync || !DriverUtil.isRunningAgainstFake()) {
258263
request.setRunAsync(true);
259264
}
260265

src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftServiceClientTest.java

Lines changed: 61 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -707,8 +707,8 @@ void testListFunctionsWithSQLEnabled() throws SQLException {
707707
.setCanDecompressLZ4Result(true)
708708
.setCanDownloadResult(true)
709709
.setQueryTimeout(0)
710+
.setRunAsync(false)
710711
.setParameters(Collections.emptyList())
711-
.setRunAsync(true)
712712
.setUseArrowNativeTypes(arrowNativeTypes);
713713
when(thriftAccessor.execute(executeStatementReq, null, session, StatementType.METADATA))
714714
.thenReturn(resultSet);
@@ -1055,6 +1055,66 @@ public void testNonDecimalType() {
10551055
assertEquals(2, result.getOrdinal());
10561056
}
10571057

1058+
@Test
1059+
void testExecuteStatementWithMetadataTypeDoesNotSetRunAsync() throws SQLException {
1060+
when(connectionContext.shouldEnableArrow()).thenReturn(true);
1061+
lenient().when(connectionContext.isCloudFetchEnabled()).thenReturn(true);
1062+
DatabricksThriftServiceClient client =
1063+
new DatabricksThriftServiceClient(thriftAccessor, connectionContext);
1064+
when(session.getSessionInfo()).thenReturn(SESSION_INFO);
1065+
1066+
when(thriftAccessor.execute(
1067+
any(TExecuteStatementReq.class), eq(null), eq(session), eq(StatementType.METADATA)))
1068+
.thenReturn(resultSet);
1069+
1070+
client.executeStatement(
1071+
"SHOW TABLES",
1072+
CLUSTER_COMPUTE,
1073+
Collections.emptyMap(),
1074+
StatementType.METADATA,
1075+
session,
1076+
null,
1077+
null);
1078+
1079+
ArgumentCaptor<TExecuteStatementReq> requestCaptor =
1080+
ArgumentCaptor.forClass(TExecuteStatementReq.class);
1081+
verify(thriftAccessor)
1082+
.execute(requestCaptor.capture(), eq(null), eq(session), eq(StatementType.METADATA));
1083+
TExecuteStatementReq request = requestCaptor.getValue();
1084+
1085+
assertFalse(request.isRunAsync(), "Expected runAsync to be false for METADATA statement type");
1086+
}
1087+
1088+
@Test
1089+
void testExecuteStatementWithSqlTypeStillSetsRunAsync() throws SQLException {
1090+
when(connectionContext.shouldEnableArrow()).thenReturn(true);
1091+
lenient().when(connectionContext.isCloudFetchEnabled()).thenReturn(true);
1092+
DatabricksThriftServiceClient client =
1093+
new DatabricksThriftServiceClient(thriftAccessor, connectionContext);
1094+
when(session.getSessionInfo()).thenReturn(SESSION_INFO);
1095+
1096+
when(thriftAccessor.execute(
1097+
any(TExecuteStatementReq.class), eq(null), eq(session), eq(StatementType.SQL)))
1098+
.thenReturn(resultSet);
1099+
1100+
client.executeStatement(
1101+
"SELECT 1",
1102+
CLUSTER_COMPUTE,
1103+
Collections.emptyMap(),
1104+
StatementType.SQL,
1105+
session,
1106+
null,
1107+
null);
1108+
1109+
ArgumentCaptor<TExecuteStatementReq> requestCaptor =
1110+
ArgumentCaptor.forClass(TExecuteStatementReq.class);
1111+
verify(thriftAccessor)
1112+
.execute(requestCaptor.capture(), eq(null), eq(session), eq(StatementType.SQL));
1113+
TExecuteStatementReq request = requestCaptor.getValue();
1114+
1115+
assertTrue(request.isRunAsync(), "Expected runAsync to be true for SQL statement type");
1116+
}
1117+
10581118
@Test
10591119
public void testNullValue() {
10601120
ImmutableSqlParameter parameter =

0 commit comments

Comments
 (0)