Skip to content

Commit a781cf9

Browse files
Gabriel39Gabriel39
committed
[fix](flight-sql) Check execution schemas and preserve prepared command behavior
1 parent 61a1c5d commit a781cf9

7 files changed

Lines changed: 314 additions & 9 deletions

File tree

‎fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/HelpCommand.java‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,29 @@ public ShowResultSet doRun(ConnectContext ctx, StmtExecutor executor) throws Exc
125125
return resultSet;
126126
}
127127

128+
/** Resolve the HELP result shape without materializing topic text or category rows. */
129+
public ShowResultSetMetaData getMetaData(ConnectContext ctx) throws AnalysisException {
130+
if (Strings.isNullOrEmpty(mark)) {
131+
throw new AnalysisException("Help empty info.");
132+
}
133+
HelpModule module = HelpModule.getInstance();
134+
// HELP has three result shapes; a topic match takes precedence over keywords and categories.
135+
HelpTopic topic = module.getTopic(mark);
136+
if (topic == null) {
137+
List<String> topics = module.listTopicByKeyword(mark);
138+
if (topics.size() > 1) {
139+
return KEYWORD_META_DATA;
140+
}
141+
if (topics.size() == 1) {
142+
topic = module.getTopic(topics.get(0));
143+
}
144+
}
145+
if (topic != null) {
146+
return TOPIC_META_DATA;
147+
}
148+
return module.listCategoryByName(mark).size() > 1 ? CATEGORY_META_DATA : KEYWORD_META_DATA;
149+
}
150+
128151
@Override
129152
public ShowResultSetMetaData getMetaData() {
130153
return ShowResultSetMetaData.builder().build();

‎fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ShowSnapshotCommand.java‎

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,16 @@ public ShowResultSetMetaData getMetaData() {
109109
return builder.build();
110110
}
111111

112+
/** Resolve the snapshot header without accessing the remote repository. */
113+
public ShowResultSetMetaData getMetaData(ConnectContext ctx) throws UserException {
114+
// Both filters must be analyzed before choosing the detail header, even during Prepare.
115+
if (!validate(ctx)) {
116+
throw new AnalysisException("Where clause should looks like: SNAPSHOT = 'your_snapshot_name'"
117+
+ " [AND TIMESTAMP = '2018-04-18-19-19-10'] [AND SNAPSHOTTYPE = 'remote' | 'local']");
118+
}
119+
return getMetaData();
120+
}
121+
112122
private boolean analyzeSubExpression(Expression expr) {
113123
if (expr instanceof ComparisonPredicate) {
114124
return analyzeComparisonPredicate((ComparisonPredicate) expr);
@@ -194,18 +204,14 @@ protected boolean validate(ConnectContext ctx) throws UserException {
194204
* handle show backup
195205
*/
196206
private ShowResultSet handleShowSnapshot(ConnectContext ctx, StmtExecutor executor) throws Exception {
197-
boolean valid = validate(ctx);
198-
if (!valid) {
199-
throw new AnalysisException("Where clause should looks like: SNAPSHOT = 'your_snapshot_name'"
200-
+ " [AND TIMESTAMP = '2018-04-18-19-19-10'] [AND SNAPSHOTTYPE = 'remote' | 'local']");
201-
}
207+
ShowResultSetMetaData metadata = getMetaData(ctx);
202208

203209
Repository repo = Env.getCurrentEnv().getBackupHandler().getRepoMgr().getRepo(repoName);
204210
if (repo == null) {
205211
throw new AnalysisException("Repository " + repoName + " does not exist");
206212
}
207213
List<List<String>> snapshotInfos = repo.getSnapshotInfos(snapshotName, timestamp);
208-
return new ShowResultSet(getMetaData(), snapshotInfos);
214+
return new ShowResultSet(metadata, snapshotInfos);
209215
}
210216

211217
@Override

‎fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/DorisFlightSqlProducer.java‎

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020

2121
package org.apache.doris.service.arrowflight;
2222

23+
import org.apache.doris.common.Status;
2324
import org.apache.doris.common.util.DebugUtil;
2425
import org.apache.doris.common.util.Util;
2526
import org.apache.doris.mysql.MysqlCommand;
@@ -329,8 +330,30 @@ public FlightInfo getFlightInfoPreparedStatement(final CommandPreparedStatementQ
329330
final CallContext context, final FlightDescriptor descriptor) {
330331
ConnectContext connection = flightSessionsManager.getConnectContext(context.peerIdentity());
331332
synchronized (connection) {
332-
return executeQueryStatement(context.peerIdentity(), connection,
333-
preparedQuery(connection, context, command).getLeft(), descriptor);
333+
Pair<String, Schema> prepared = preparedQuery(connection, context, command);
334+
FlightInfo info = executeQueryStatement(context.peerIdentity(), connection, prepared.getLeft(), descriptor);
335+
String id = command.getPreparedStatementHandle().toStringUtf8()
336+
.substring(context.peerIdentity().length() + 1);
337+
// Another session's DDL can change the result after reanalysis but before execution
338+
// acquires table locks. Reject the actual schema before publishing a DoGet ticket.
339+
if (!prepared.getRight().equals(info.getSchema())) {
340+
connection.removePreparedQuery(id);
341+
try {
342+
connection.cancelQuery(Status.CANCELLED);
343+
} catch (Exception e) {
344+
LOG.warn("Failed to cancel Flight query after prepared schema changed", e);
345+
} finally {
346+
connection.closeFlightSqlDeferredExecutors();
347+
connection.getFlightSqlChannel().reset();
348+
connection.clearFlightSqlEndpointsLocations();
349+
}
350+
throw CallStatus.NOT_FOUND.withDescription("Prepared statement schema changed; prepare again")
351+
.toRuntimeException();
352+
}
353+
// A successful USE/SWITCH in this statement may change its own namespace. Rebind
354+
// only this handle; unrelated handles must still reject an external namespace change.
355+
connection.addPreparedQuery(id, prepared.getLeft(), prepared.getRight());
356+
return info;
334357
}
335358
}
336359

‎fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/FlightSqlQuerySchema.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
import org.apache.doris.nereids.trees.plans.commands.DeleteFromCommand;
4747
import org.apache.doris.nereids.trees.plans.commands.DescribeCommand;
4848
import org.apache.doris.nereids.trees.plans.commands.ExplainCommand;
49+
import org.apache.doris.nereids.trees.plans.commands.HelpCommand;
4950
import org.apache.doris.nereids.trees.plans.commands.KillCommand;
5051
import org.apache.doris.nereids.trees.plans.commands.ReplayCommand;
5152
import org.apache.doris.nereids.trees.plans.commands.ShowCreateTableCommand;
@@ -54,6 +55,7 @@
5455
import org.apache.doris.nereids.trees.plans.commands.ShowProcCommand;
5556
import org.apache.doris.nereids.trees.plans.commands.ShowPythonPackagesCommand;
5657
import org.apache.doris.nereids.trees.plans.commands.ShowQueryStatsCommand;
58+
import org.apache.doris.nereids.trees.plans.commands.ShowSnapshotCommand;
5759
import org.apache.doris.nereids.trees.plans.commands.ShowTableCommand;
5860
import org.apache.doris.nereids.trees.plans.commands.TransactionCommand;
5961
import org.apache.doris.nereids.trees.plans.commands.UpdateCommand;
@@ -241,7 +243,11 @@ private static ResultSetMetaData commandMetadata(ConnectContext context, Command
241243
throw CallStatus.UNIMPLEMENTED.withDescription("Command schema requires execution-time metadata")
242244
.toRuntimeException();
243245
}
244-
if (command instanceof ShowTableCommand) {
246+
if (command instanceof HelpCommand) {
247+
return ((HelpCommand) command).getMetaData(context);
248+
} else if (command instanceof ShowSnapshotCommand) {
249+
return ((ShowSnapshotCommand) command).getMetaData(context);
250+
} else if (command instanceof ShowTableCommand) {
245251
((ShowTableCommand) command).validate(context);
246252
} else if (command instanceof ShowCreateTableCommand) {
247253
return ((ShowCreateTableCommand) command).getMetaData(context);

‎fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/SchemaOnlyCommandMetadataTest.java‎

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
package org.apache.doris.nereids.trees.plans.commands;
19+
20+
import org.apache.doris.catalog.Env;
21+
import org.apache.doris.common.AnalysisException;
22+
import org.apache.doris.mysql.privilege.PrivPredicate;
23+
import org.apache.doris.nereids.parser.NereidsParser;
24+
import org.apache.doris.qe.ConnectContext;
25+
import org.apache.doris.qe.ResultSetMetaData;
26+
import org.apache.doris.qe.help.HelpModule;
27+
import org.apache.doris.qe.help.HelpTopic;
28+
29+
import org.junit.jupiter.api.Assertions;
30+
import org.junit.jupiter.api.Test;
31+
import org.mockito.MockedStatic;
32+
import org.mockito.Mockito;
33+
34+
import java.util.Arrays;
35+
import java.util.Collections;
36+
import java.util.List;
37+
import java.util.stream.Collectors;
38+
39+
class SchemaOnlyCommandMetadataTest {
40+
private static List<String> names(ResultSetMetaData metadata) {
41+
return metadata.getColumns().stream().map(column -> column.getName()).collect(Collectors.toList());
42+
}
43+
44+
@Test
45+
void helpTopicDoesNotReadDescriptionOrExample() throws Exception {
46+
HelpModule module = Mockito.mock(HelpModule.class);
47+
HelpTopic topic = Mockito.mock(HelpTopic.class);
48+
Mockito.when(module.getTopic("topic")).thenReturn(topic);
49+
try (MockedStatic<HelpModule> ignored = Mockito.mockStatic(HelpModule.class)) {
50+
ignored.when(HelpModule::getInstance).thenReturn(module);
51+
HelpCommand command = new HelpCommand("topic");
52+
Assertions.assertEquals(Arrays.asList("name", "description", "example"), names(command.getMetaData(null)));
53+
Mockito.verifyNoInteractions(topic);
54+
Assertions.assertEquals(names(command.doRun(null, null).getMetaData()), names(command.getMetaData(null)));
55+
}
56+
}
57+
58+
@Test
59+
void helpKeywordWithOneTopicUsesTopicHeader() throws Exception {
60+
HelpModule module = Mockito.mock(HelpModule.class);
61+
Mockito.when(module.listTopicByKeyword("keyword")).thenReturn(Collections.singletonList("topic"));
62+
Mockito.when(module.getTopic("topic")).thenReturn(Mockito.mock(HelpTopic.class));
63+
try (MockedStatic<HelpModule> ignored = Mockito.mockStatic(HelpModule.class)) {
64+
ignored.when(HelpModule::getInstance).thenReturn(module);
65+
HelpCommand command = new HelpCommand("keyword");
66+
Assertions.assertEquals(Arrays.asList("name", "description", "example"), names(command.getMetaData(null)));
67+
Assertions.assertEquals(names(command.doRun(null, null).getMetaData()), names(command.getMetaData(null)));
68+
}
69+
}
70+
71+
@Test
72+
void helpKeywordAndCategoryHeadersMatchExecution() throws Exception {
73+
HelpModule module = Mockito.mock(HelpModule.class);
74+
Mockito.when(module.listTopicByKeyword("keyword")).thenReturn(Arrays.asList("first", "second"));
75+
Mockito.when(module.listCategoryByName("categories")).thenReturn(Arrays.asList("first", "second"));
76+
Mockito.when(module.listCategoryByName("category")).thenReturn(Collections.singletonList("first"));
77+
try (MockedStatic<HelpModule> ignored = Mockito.mockStatic(HelpModule.class)) {
78+
ignored.when(HelpModule::getInstance).thenReturn(module);
79+
for (String mark : Arrays.asList("keyword", "categories", "category", "missing")) {
80+
HelpCommand command = new HelpCommand(mark);
81+
List<String> expected = mark.equals("categories")
82+
? Arrays.asList("source_category_name", "name", "is_it_category")
83+
: Arrays.asList("name", "is_it_category");
84+
Assertions.assertEquals(expected, names(command.getMetaData(null)));
85+
Assertions.assertEquals(names(command.doRun(null, null).getMetaData()), names(command.getMetaData(null)));
86+
}
87+
}
88+
}
89+
90+
@Test
91+
void invalidHelpAndSnapshotFiltersAreRejected() throws Exception {
92+
Assertions.assertThrows(AnalysisException.class, () -> new HelpCommand("").getMetaData(null));
93+
Env env = Mockito.mock(Env.class, Mockito.RETURNS_DEEP_STUBS);
94+
Mockito.when(env.getAccessManager().checkGlobalPriv(Mockito.nullable(ConnectContext.class),
95+
Mockito.eq(PrivPredicate.ADMIN))).thenReturn(true);
96+
try (MockedStatic<Env> ignored = Mockito.mockStatic(Env.class)) {
97+
ignored.when(Env::getCurrentEnv).thenReturn(env);
98+
ShowSnapshotCommand command = (ShowSnapshotCommand) new NereidsParser()
99+
.parseSingle("SHOW SNAPSHOT ON example_repo WHERE TIMESTAMP = '2026-01-01-00-00-00'");
100+
Assertions.assertThrows(AnalysisException.class, () -> command.getMetaData(null));
101+
Mockito.verify(env, Mockito.never()).getBackupHandler();
102+
}
103+
}
104+
105+
@Test
106+
void filteredSnapshotsResolveWithoutRepositoryAccess() throws Exception {
107+
Env env = Mockito.mock(Env.class, Mockito.RETURNS_DEEP_STUBS);
108+
Mockito.when(env.getAccessManager().checkGlobalPriv(Mockito.nullable(ConnectContext.class),
109+
Mockito.eq(PrivPredicate.ADMIN))).thenReturn(true);
110+
try (MockedStatic<Env> ignored = Mockito.mockStatic(Env.class)) {
111+
ignored.when(Env::getCurrentEnv).thenReturn(env);
112+
for (String filter : Arrays.asList("", " WHERE SNAPSHOT = 'sample'",
113+
" WHERE SNAPSHOT = 'sample' AND TIMESTAMP = '2026-01-01-00-00-00'")) {
114+
ShowSnapshotCommand command = (ShowSnapshotCommand) new NereidsParser()
115+
.parseSingle("SHOW SNAPSHOT ON example_repo" + filter);
116+
Assertions.assertEquals(filter.contains("TIMESTAMP")
117+
? ShowSnapshotCommand.SNAPSHOT_DETAIL : ShowSnapshotCommand.SNAPSHOT_ALL,
118+
names(command.getMetaData(null)));
119+
}
120+
Mockito.verify(env, Mockito.never()).getBackupHandler();
121+
}
122+
}
123+
}

‎fe/fe-core/src/test/java/org/apache/doris/service/arrowflight/DorisFlightSqlSchemaTest.java‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import org.apache.doris.analysis.UserIdentity;
2121
import org.apache.doris.catalog.Column;
2222
import org.apache.doris.common.FeConstants;
23+
import org.apache.doris.common.Status;
2324
import org.apache.doris.common.proc.ProcNodeInterface;
2425
import org.apache.doris.common.proc.ProcService;
2526
import org.apache.doris.datasource.lance.LanceExternalCatalog;
@@ -605,4 +606,86 @@ void explainAndReplayExposeHeadersWithoutRunningPlanner() throws Exception {
605606
FlightRuntimeException.class, () -> schema("EXPLAIN PLAN PROCESS SELECT 1")).status().code());
606607
}
607608

609+
@Test
610+
void preparedNamespaceTransitionIsReusable() throws Exception {
611+
String database = connectContext.getDatabase();
612+
try {
613+
for (String query : Arrays.asList("USE information_schema",
614+
"USE information_schema; SELECT TABLE_NAME FROM tables WHERE 1=0")) {
615+
connectContext.setDatabase(database);
616+
Schema expected = schema(query);
617+
ActionCreatePreparedStatementResult result = prepare(query);
618+
CommandPreparedStatementQuery command = CommandPreparedStatementQuery.newBuilder()
619+
.setPreparedStatementHandle(result.getPreparedStatementHandle()).build();
620+
try (MockedConstruction<FlightSqlConnectProcessor> processors = Mockito.mockConstruction(
621+
FlightSqlConnectProcessor.class, (processor, context) -> {
622+
Mockito.doAnswer(invocation -> {
623+
connectContext.setDatabase("information_schema");
624+
connectContext.setReturnResultFromLocal(false);
625+
return null;
626+
}).when(processor).handleQuery(Mockito.anyString());
627+
Mockito.when(processor.getArrowSchema()).thenReturn(expected);
628+
})) {
629+
for (int i = 0; i < 2; ++i) {
630+
Assertions.assertEquals(expected, producer.getFlightInfoPreparedStatement(command,
631+
callContext, FlightDescriptor.command(new byte[0])).getSchema());
632+
Assertions.assertEquals(expected, producer.getSchemaPreparedStatement(command,
633+
callContext, FlightDescriptor.command(new byte[0])).getSchema());
634+
}
635+
} finally {
636+
String handle = result.getPreparedStatementHandle().toStringUtf8();
637+
connectContext.removePreparedQuery(handle.substring(handle.indexOf(':') + 1));
638+
}
639+
}
640+
} finally {
641+
connectContext.setDatabase(database);
642+
}
643+
}
644+
645+
@Test
646+
void executionSchemaDriftExpiresHandleAndCleansResults() throws Exception {
647+
Schema changed = new Schema(Arrays.asList(Field.nullable("id", new ArrowType.Int(64, true)),
648+
Field.nullable("new_column", new ArrowType.Utf8())));
649+
StmtExecutor previous = connectContext.getExecutor();
650+
try {
651+
for (boolean remote : Arrays.asList(true, false)) {
652+
ActionCreatePreparedStatementResult result = prepare("SELECT id FROM schema_input");
653+
CommandPreparedStatementQuery command = CommandPreparedStatementQuery.newBuilder()
654+
.setPreparedStatementHandle(result.getPreparedStatementHandle()).build();
655+
StmtExecutor executor = Mockito.mock(StmtExecutor.class);
656+
try (MockedConstruction<FlightSqlConnectProcessor> processors = Mockito.mockConstruction(
657+
FlightSqlConnectProcessor.class, (processor, context) -> {
658+
// Simulate replanning after another session changes the table, after schema reanalysis.
659+
Mockito.doAnswer(invocation -> {
660+
connectContext.setReturnResultFromLocal(!remote);
661+
connectContext.setExecutor(executor);
662+
connectContext.addFlightSqlDeferredExecutor(executor);
663+
return null;
664+
}).when(processor).handleQuery(Mockito.anyString());
665+
Mockito.when(processor.getArrowSchema()).thenReturn(changed);
666+
})) {
667+
FlightRuntimeException error = Assertions.assertThrows(FlightRuntimeException.class,
668+
() -> producer.getFlightInfoPreparedStatement(command, callContext,
669+
FlightDescriptor.command(new byte[0])));
670+
Assertions.assertEquals(FlightStatusCode.NOT_FOUND, error.status().code());
671+
Mockito.verify(executor).cancel(Mockito.any(Status.class));
672+
Mockito.verify(executor).finalizeArrowFlightQuery();
673+
Assertions.assertEquals(0, connectContext.getFlightSqlChannel().resultNum());
674+
Assertions.assertTrue(connectContext.getFlightSqlEndpointsLocations().isEmpty());
675+
String handle = result.getPreparedStatementHandle().toStringUtf8();
676+
Assertions.assertNull(connectContext.getPreparedQuery(handle.substring(handle.indexOf(':') + 1)));
677+
connectContext.closeFlightSqlDeferredExecutors();
678+
Mockito.verify(executor, Mockito.times(1)).finalizeArrowFlightQuery();
679+
} finally {
680+
connectContext.closeFlightSqlDeferredExecutors();
681+
connectContext.getFlightSqlChannel().reset();
682+
String handle = result.getPreparedStatementHandle().toStringUtf8();
683+
connectContext.removePreparedQuery(handle.substring(handle.indexOf(':') + 1));
684+
}
685+
}
686+
} finally {
687+
connectContext.setExecutor(previous);
688+
}
689+
}
690+
608691
}

0 commit comments

Comments
 (0)