Gabriel39 commented on code in PR #68631: URL: https://github.com/apache/doris/pull/68631#discussion_r4132316828
########## fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/FlightSqlQuerySchema.java: ########## @@ -0,0 +1,301 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package org.apache.doris.service.arrowflight; + +import org.apache.doris.analysis.StatementBase; +import org.apache.doris.catalog.AggStateType; +import org.apache.doris.catalog.ArrayType; +import org.apache.doris.catalog.Column; +import org.apache.doris.catalog.MapType; +import org.apache.doris.catalog.PrimitiveType; +import org.apache.doris.catalog.ScalarType; +import org.apache.doris.catalog.StructField; +import org.apache.doris.catalog.StructType; +import org.apache.doris.catalog.Type; +import org.apache.doris.datasource.CatalogIf; +import org.apache.doris.datasource.es.EsExternalCatalog; +import org.apache.doris.mysql.MysqlCommand; +import org.apache.doris.mysql.privilege.PrivPredicate; +import org.apache.doris.nereids.CascadesContext; +import org.apache.doris.nereids.StatementContext; +import org.apache.doris.nereids.glue.LogicalPlanAdapter; +import org.apache.doris.nereids.parser.NereidsParser; +import org.apache.doris.nereids.parser.SqlDialectHelper; +import org.apache.doris.nereids.rules.rewrite.CheckPrivileges; +import org.apache.doris.nereids.trees.expressions.Slot; +import org.apache.doris.nereids.trees.plans.Plan; +import org.apache.doris.nereids.trees.plans.PrepareCommandPlanner; +import org.apache.doris.nereids.trees.plans.commands.Command; +import org.apache.doris.nereids.trees.plans.commands.DescribeCommand; +import org.apache.doris.nereids.trees.plans.commands.ShowCreateTableCommand; +import org.apache.doris.nereids.trees.plans.commands.ShowDataCommand; +import org.apache.doris.nereids.trees.plans.commands.ShowPartitionsCommand; +import org.apache.doris.nereids.trees.plans.commands.ShowProcCommand; +import org.apache.doris.nereids.trees.plans.commands.ShowPythonPackagesCommand; +import org.apache.doris.nereids.trees.plans.commands.ShowQueryStatsCommand; +import org.apache.doris.nereids.trees.plans.commands.ShowTableCommand; +import org.apache.doris.nereids.trees.plans.commands.use.SwitchCommand; +import org.apache.doris.nereids.trees.plans.commands.use.UseCommand; +import org.apache.doris.qe.ConnectContext; +import org.apache.doris.qe.QueryState; +import org.apache.doris.qe.ResultSetMetaData; +import org.apache.doris.qe.SessionVariable; +import org.apache.doris.qe.StmtExecutor; +import org.apache.doris.qe.VariableMgr; + +import org.apache.arrow.flight.CallStatus; +import org.apache.arrow.util.AutoCloseables; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** Resolves result metadata without scheduling fragments or evaluating query expressions. */ +final class FlightSqlQuerySchema { + private FlightSqlQuerySchema() { + } + + static Schema analyze(ConnectContext context, String query) throws Exception { + synchronized (context) { + ConnectContext previousThreadContext = ConnectContext.get(); + StatementContext previousStatement = context.getStatementContext(); + SessionVariable previousSession = context.getSessionVariable(); + QueryState previousState = context.getState(); + StmtExecutor previousExecutor = context.getExecutor(); + String previousCatalog = context.getDefaultCatalog(); + String previousDatabase = context.getDatabase(); + List<StatementBase> statements = Collections.emptyList(); + try { + context.setThreadLocalInfo(); + context.setCommand(MysqlCommand.COM_QUERY); + // Parsing SET_VAR hints already mutates session variables. Isolate them even when parsing fails. + context.setSessionVariable(VariableMgr.cloneSessionVariable(previousSession)); + context.setState(new QueryState()); + context.setExecutor(null); + context.setStatementContext(null); + // Match execution's HTTP/plugin conversion before the dialect parser sees the SQL. + String converted = SqlDialectHelper.convertSqlByDialect(query, context.getSessionVariable()); + statements = new NereidsParser().parseSQL(converted, context.getSessionVariable()); + Map<String, String> scopedDatabases = new HashMap<>(); + if (statements.isEmpty()) { + throw CallStatus.UNIMPLEMENTED.withDescription( + "Schema discovery requires a statement").toRuntimeException(); + } + // JDBC clients commonly prefix their query with USE. Resolve that namespace only within this scope. + for (int i = 0; i < statements.size() - 1; ++i) { + Plan prefix = ((LogicalPlanAdapter) statements.get(i)).getLogicalPlan(); + if (!(prefix instanceof UseCommand) && !(prefix instanceof SwitchCommand)) { + throw CallStatus.UNIMPLEMENTED.withDescription( + "Schema discovery only supports USE or SWITCH before the result statement") + .toRuntimeException(); + } + resolveNamespace(context, prefix, scopedDatabases); + } + LogicalPlanAdapter statement = (LogicalPlanAdapter) statements.get(statements.size() - 1); + StatementContext statementContext = statement.getStatementContext(); + context.setStatementContext(statementContext); + statementContext.setParsedStatement(statement); + if (!statementContext.getPlaceholders().isEmpty()) { + throw CallStatus.UNIMPLEMENTED.withDescription( + "Flight SQL parameter binding is not supported").toRuntimeException(); + } + List<Field> fields = new ArrayList<>(); + Plan plan = statement.getLogicalPlan(); + if (plan instanceof Command) { + resolveNamespace(context, plan, scopedDatabases); + ResultSetMetaData metadata = commandMetadata(context, (Command) plan); + if (metadata == null) { + throw CallStatus.UNIMPLEMENTED.withDescription("Command result metadata is unavailable") + .toRuntimeException(); + } + // FE-local result sets are serialized as nullable strings by FlightSqlChannel. + for (Column column : metadata.getColumns()) { + fields.add(Field.nullable(column.getName(), new ArrowType.Utf8())); + } + if (fields.isEmpty()) { + switch (((Command) plan).stmtType()) { + case SET: + case USE: + case SWITCH: + case CREATE: + case ALTER: Review Comment: Confirmed and fixed in 37f71c9306c. Top-level CREATE/DROP INDEX targeting a Lance catalog now expose the JobId header by inspecting the parsed operations and catalog type, without table resolution, validation, admission, or ID allocation. Tests cover CREATE, CREATE IF NOT EXISTS, DROP IF EXISTS, and the ordinary internal-index StatusResult header. ########## fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/FlightSqlQuerySchema.java: ########## @@ -0,0 +1,301 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package org.apache.doris.service.arrowflight; + +import org.apache.doris.analysis.StatementBase; +import org.apache.doris.catalog.AggStateType; +import org.apache.doris.catalog.ArrayType; +import org.apache.doris.catalog.Column; +import org.apache.doris.catalog.MapType; +import org.apache.doris.catalog.PrimitiveType; +import org.apache.doris.catalog.ScalarType; +import org.apache.doris.catalog.StructField; +import org.apache.doris.catalog.StructType; +import org.apache.doris.catalog.Type; +import org.apache.doris.datasource.CatalogIf; +import org.apache.doris.datasource.es.EsExternalCatalog; +import org.apache.doris.mysql.MysqlCommand; +import org.apache.doris.mysql.privilege.PrivPredicate; +import org.apache.doris.nereids.CascadesContext; +import org.apache.doris.nereids.StatementContext; +import org.apache.doris.nereids.glue.LogicalPlanAdapter; +import org.apache.doris.nereids.parser.NereidsParser; +import org.apache.doris.nereids.parser.SqlDialectHelper; +import org.apache.doris.nereids.rules.rewrite.CheckPrivileges; +import org.apache.doris.nereids.trees.expressions.Slot; +import org.apache.doris.nereids.trees.plans.Plan; +import org.apache.doris.nereids.trees.plans.PrepareCommandPlanner; +import org.apache.doris.nereids.trees.plans.commands.Command; +import org.apache.doris.nereids.trees.plans.commands.DescribeCommand; +import org.apache.doris.nereids.trees.plans.commands.ShowCreateTableCommand; +import org.apache.doris.nereids.trees.plans.commands.ShowDataCommand; +import org.apache.doris.nereids.trees.plans.commands.ShowPartitionsCommand; +import org.apache.doris.nereids.trees.plans.commands.ShowProcCommand; +import org.apache.doris.nereids.trees.plans.commands.ShowPythonPackagesCommand; +import org.apache.doris.nereids.trees.plans.commands.ShowQueryStatsCommand; +import org.apache.doris.nereids.trees.plans.commands.ShowTableCommand; +import org.apache.doris.nereids.trees.plans.commands.use.SwitchCommand; +import org.apache.doris.nereids.trees.plans.commands.use.UseCommand; +import org.apache.doris.qe.ConnectContext; +import org.apache.doris.qe.QueryState; +import org.apache.doris.qe.ResultSetMetaData; +import org.apache.doris.qe.SessionVariable; +import org.apache.doris.qe.StmtExecutor; +import org.apache.doris.qe.VariableMgr; + +import org.apache.arrow.flight.CallStatus; +import org.apache.arrow.util.AutoCloseables; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** Resolves result metadata without scheduling fragments or evaluating query expressions. */ +final class FlightSqlQuerySchema { + private FlightSqlQuerySchema() { + } + + static Schema analyze(ConnectContext context, String query) throws Exception { + synchronized (context) { + ConnectContext previousThreadContext = ConnectContext.get(); + StatementContext previousStatement = context.getStatementContext(); + SessionVariable previousSession = context.getSessionVariable(); + QueryState previousState = context.getState(); + StmtExecutor previousExecutor = context.getExecutor(); + String previousCatalog = context.getDefaultCatalog(); + String previousDatabase = context.getDatabase(); + List<StatementBase> statements = Collections.emptyList(); + try { + context.setThreadLocalInfo(); + context.setCommand(MysqlCommand.COM_QUERY); + // Parsing SET_VAR hints already mutates session variables. Isolate them even when parsing fails. + context.setSessionVariable(VariableMgr.cloneSessionVariable(previousSession)); + context.setState(new QueryState()); + context.setExecutor(null); + context.setStatementContext(null); + // Match execution's HTTP/plugin conversion before the dialect parser sees the SQL. + String converted = SqlDialectHelper.convertSqlByDialect(query, context.getSessionVariable()); + statements = new NereidsParser().parseSQL(converted, context.getSessionVariable()); + Map<String, String> scopedDatabases = new HashMap<>(); + if (statements.isEmpty()) { + throw CallStatus.UNIMPLEMENTED.withDescription( + "Schema discovery requires a statement").toRuntimeException(); + } + // JDBC clients commonly prefix their query with USE. Resolve that namespace only within this scope. + for (int i = 0; i < statements.size() - 1; ++i) { + Plan prefix = ((LogicalPlanAdapter) statements.get(i)).getLogicalPlan(); + if (!(prefix instanceof UseCommand) && !(prefix instanceof SwitchCommand)) { + throw CallStatus.UNIMPLEMENTED.withDescription( + "Schema discovery only supports USE or SWITCH before the result statement") + .toRuntimeException(); + } + resolveNamespace(context, prefix, scopedDatabases); + } + LogicalPlanAdapter statement = (LogicalPlanAdapter) statements.get(statements.size() - 1); + StatementContext statementContext = statement.getStatementContext(); + context.setStatementContext(statementContext); + statementContext.setParsedStatement(statement); + if (!statementContext.getPlaceholders().isEmpty()) { + throw CallStatus.UNIMPLEMENTED.withDescription( + "Flight SQL parameter binding is not supported").toRuntimeException(); + } + List<Field> fields = new ArrayList<>(); + Plan plan = statement.getLogicalPlan(); + if (plan instanceof Command) { + resolveNamespace(context, plan, scopedDatabases); + ResultSetMetaData metadata = commandMetadata(context, (Command) plan); + if (metadata == null) { + throw CallStatus.UNIMPLEMENTED.withDescription("Command result metadata is unavailable") + .toRuntimeException(); + } + // FE-local result sets are serialized as nullable strings by FlightSqlChannel. + for (Column column : metadata.getColumns()) { + fields.add(Field.nullable(column.getName(), new ArrowType.Utf8())); + } + if (fields.isEmpty()) { + switch (((Command) plan).stmtType()) { + case SET: + case USE: + case SWITCH: + case CREATE: + case ALTER: + case DROP: + case TRUNCATE: + // Only known no-row command categories have the protocol OK schema. + fields.add(Field.nullable("StatusResult", new ArrowType.Utf8())); + break; + default: + throw CallStatus.UNIMPLEMENTED.withDescription( + "Result metadata is unavailable without executing this command") + .toRuntimeException(); + } + } + } else { + PrepareCommandPlanner planner = new PrepareCommandPlanner(statementContext); + planner.plan(statement, context.getSessionVariable().toThrift()); + CascadesContext cascades = planner.getCascadesContext(); + Plan analyzed = cascades.getRewritePlan(); + // PrepareCommandPlanner stops before the rewrite phase that normally checks privileges. + new CheckPrivileges().rewriteRoot(analyzed, cascades.getCurrentJobContext()); + for (Slot slot : analyzed.getOutput()) { + fields.add(field(slot.getName(), slot.getDataType().toCatalogDataType(), slot.nullable(), + true, context.getSessionVariable().getTimeZone())); + } + } + return new Schema(fields); + } finally { + try { + List<AutoCloseable> resources = new ArrayList<>(); + for (StatementBase statement : statements) { + if (statement instanceof LogicalPlanAdapter) { + resources.add(((LogicalPlanAdapter) statement).getStatementContext()); + } + } + // A parser failure can leave a context that was never added to the returned list. + StatementContext current = context.getStatementContext(); + if (current != null && current != previousStatement && !resources.contains(current)) { + resources.add(current); + } + AutoCloseables.close(resources); + } finally { + try { + context.setStatementContext(previousStatement); + context.setSessionVariable(previousSession); + context.setState(previousState); + context.setExecutor(previousExecutor); + if (!previousCatalog.equals(context.getDefaultCatalog()) + || !previousDatabase.equals(context.getDatabase())) { + context.changeDefaultCatalog(previousCatalog); + context.setDatabase(previousDatabase); + } + } finally { + context.setCommand(MysqlCommand.COM_SLEEP); + if (previousThreadContext == null) { + ConnectContext.remove(); + } else { + previousThreadContext.setThreadLocalInfo(); + } + } + } + } + } + } + + private static ResultSetMetaData commandMetadata(ConnectContext context, Command command) throws Exception { + // These getters depend on execution-time state or remote responses. Do not advertise a + // guessed schema, or run the command merely to discover it. + if (command instanceof ShowPythonPackagesCommand || command instanceof DescribeCommand + || command instanceof ShowDataCommand || command instanceof ShowPartitionsCommand + || command instanceof ShowQueryStatsCommand) { + throw CallStatus.UNIMPLEMENTED.withDescription("Command schema requires execution-time metadata") + .toRuntimeException(); + } + if (command instanceof ShowTableCommand) { + ((ShowTableCommand) command).validate(context); + } else if (command instanceof ShowCreateTableCommand) { + return ((ShowCreateTableCommand) command).getMetaData(context); + } else if (command instanceof ShowProcCommand) { + return ((ShowProcCommand) command).getMetaData(context); + } + return command.getResultSetMetaData(); Review Comment: Confirmed and fixed in 37f71c9306c. Both procedure SHOW commands now override getResultSetMetaData and reuse their fixed getMetaData headers. Added FE and Flight RPC coverage for SHOW PROCEDURE STATUS and SHOW CREATE PROCEDURE without running procedure lookup. ########## fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/FlightSqlQuerySchema.java: ########## @@ -0,0 +1,301 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package org.apache.doris.service.arrowflight; + +import org.apache.doris.analysis.StatementBase; +import org.apache.doris.catalog.AggStateType; +import org.apache.doris.catalog.ArrayType; +import org.apache.doris.catalog.Column; +import org.apache.doris.catalog.MapType; +import org.apache.doris.catalog.PrimitiveType; +import org.apache.doris.catalog.ScalarType; +import org.apache.doris.catalog.StructField; +import org.apache.doris.catalog.StructType; +import org.apache.doris.catalog.Type; +import org.apache.doris.datasource.CatalogIf; +import org.apache.doris.datasource.es.EsExternalCatalog; +import org.apache.doris.mysql.MysqlCommand; +import org.apache.doris.mysql.privilege.PrivPredicate; +import org.apache.doris.nereids.CascadesContext; +import org.apache.doris.nereids.StatementContext; +import org.apache.doris.nereids.glue.LogicalPlanAdapter; +import org.apache.doris.nereids.parser.NereidsParser; +import org.apache.doris.nereids.parser.SqlDialectHelper; +import org.apache.doris.nereids.rules.rewrite.CheckPrivileges; +import org.apache.doris.nereids.trees.expressions.Slot; +import org.apache.doris.nereids.trees.plans.Plan; +import org.apache.doris.nereids.trees.plans.PrepareCommandPlanner; +import org.apache.doris.nereids.trees.plans.commands.Command; +import org.apache.doris.nereids.trees.plans.commands.DescribeCommand; +import org.apache.doris.nereids.trees.plans.commands.ShowCreateTableCommand; +import org.apache.doris.nereids.trees.plans.commands.ShowDataCommand; +import org.apache.doris.nereids.trees.plans.commands.ShowPartitionsCommand; +import org.apache.doris.nereids.trees.plans.commands.ShowProcCommand; +import org.apache.doris.nereids.trees.plans.commands.ShowPythonPackagesCommand; +import org.apache.doris.nereids.trees.plans.commands.ShowQueryStatsCommand; +import org.apache.doris.nereids.trees.plans.commands.ShowTableCommand; +import org.apache.doris.nereids.trees.plans.commands.use.SwitchCommand; +import org.apache.doris.nereids.trees.plans.commands.use.UseCommand; +import org.apache.doris.qe.ConnectContext; +import org.apache.doris.qe.QueryState; +import org.apache.doris.qe.ResultSetMetaData; +import org.apache.doris.qe.SessionVariable; +import org.apache.doris.qe.StmtExecutor; +import org.apache.doris.qe.VariableMgr; + +import org.apache.arrow.flight.CallStatus; +import org.apache.arrow.util.AutoCloseables; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** Resolves result metadata without scheduling fragments or evaluating query expressions. */ +final class FlightSqlQuerySchema { + private FlightSqlQuerySchema() { + } + + static Schema analyze(ConnectContext context, String query) throws Exception { + synchronized (context) { + ConnectContext previousThreadContext = ConnectContext.get(); + StatementContext previousStatement = context.getStatementContext(); + SessionVariable previousSession = context.getSessionVariable(); + QueryState previousState = context.getState(); + StmtExecutor previousExecutor = context.getExecutor(); + String previousCatalog = context.getDefaultCatalog(); + String previousDatabase = context.getDatabase(); + List<StatementBase> statements = Collections.emptyList(); + try { + context.setThreadLocalInfo(); + context.setCommand(MysqlCommand.COM_QUERY); + // Parsing SET_VAR hints already mutates session variables. Isolate them even when parsing fails. + context.setSessionVariable(VariableMgr.cloneSessionVariable(previousSession)); + context.setState(new QueryState()); + context.setExecutor(null); + context.setStatementContext(null); + // Match execution's HTTP/plugin conversion before the dialect parser sees the SQL. + String converted = SqlDialectHelper.convertSqlByDialect(query, context.getSessionVariable()); + statements = new NereidsParser().parseSQL(converted, context.getSessionVariable()); + Map<String, String> scopedDatabases = new HashMap<>(); + if (statements.isEmpty()) { + throw CallStatus.UNIMPLEMENTED.withDescription( + "Schema discovery requires a statement").toRuntimeException(); + } + // JDBC clients commonly prefix their query with USE. Resolve that namespace only within this scope. + for (int i = 0; i < statements.size() - 1; ++i) { + Plan prefix = ((LogicalPlanAdapter) statements.get(i)).getLogicalPlan(); + if (!(prefix instanceof UseCommand) && !(prefix instanceof SwitchCommand)) { + throw CallStatus.UNIMPLEMENTED.withDescription( + "Schema discovery only supports USE or SWITCH before the result statement") + .toRuntimeException(); + } + resolveNamespace(context, prefix, scopedDatabases); + } + LogicalPlanAdapter statement = (LogicalPlanAdapter) statements.get(statements.size() - 1); + StatementContext statementContext = statement.getStatementContext(); + context.setStatementContext(statementContext); + statementContext.setParsedStatement(statement); + if (!statementContext.getPlaceholders().isEmpty()) { + throw CallStatus.UNIMPLEMENTED.withDescription( + "Flight SQL parameter binding is not supported").toRuntimeException(); + } + List<Field> fields = new ArrayList<>(); + Plan plan = statement.getLogicalPlan(); + if (plan instanceof Command) { + resolveNamespace(context, plan, scopedDatabases); + ResultSetMetaData metadata = commandMetadata(context, (Command) plan); + if (metadata == null) { + throw CallStatus.UNIMPLEMENTED.withDescription("Command result metadata is unavailable") + .toRuntimeException(); + } + // FE-local result sets are serialized as nullable strings by FlightSqlChannel. + for (Column column : metadata.getColumns()) { + fields.add(Field.nullable(column.getName(), new ArrowType.Utf8())); + } + if (fields.isEmpty()) { + switch (((Command) plan).stmtType()) { + case SET: + case USE: + case SWITCH: + case CREATE: + case ALTER: + case DROP: + case TRUNCATE: + // Only known no-row command categories have the protocol OK schema. + fields.add(Field.nullable("StatusResult", new ArrowType.Utf8())); + break; + default: + throw CallStatus.UNIMPLEMENTED.withDescription( Review Comment: Confirmed and fixed in 37f71c9306c. Concrete no-row DML, KILL, and transaction commands now expose StatusResult. WARM UP SELECT overrides getResultSetMetaData with its six statistics columns, which take precedence over the no-row fallback. Added FE and Flight RPC tests for INSERT, UPDATE, DELETE, MERGE, KILL, BEGIN/COMMIT/ROLLBACK and WARM UP SELECT; FE coverage also includes INSERT OVERWRITE. Schema discovery does not execute these commands. ########## fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/FlightSqlQuerySchema.java: ########## @@ -0,0 +1,301 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package org.apache.doris.service.arrowflight; + +import org.apache.doris.analysis.StatementBase; +import org.apache.doris.catalog.AggStateType; +import org.apache.doris.catalog.ArrayType; +import org.apache.doris.catalog.Column; +import org.apache.doris.catalog.MapType; +import org.apache.doris.catalog.PrimitiveType; +import org.apache.doris.catalog.ScalarType; +import org.apache.doris.catalog.StructField; +import org.apache.doris.catalog.StructType; +import org.apache.doris.catalog.Type; +import org.apache.doris.datasource.CatalogIf; +import org.apache.doris.datasource.es.EsExternalCatalog; +import org.apache.doris.mysql.MysqlCommand; +import org.apache.doris.mysql.privilege.PrivPredicate; +import org.apache.doris.nereids.CascadesContext; +import org.apache.doris.nereids.StatementContext; +import org.apache.doris.nereids.glue.LogicalPlanAdapter; +import org.apache.doris.nereids.parser.NereidsParser; +import org.apache.doris.nereids.parser.SqlDialectHelper; +import org.apache.doris.nereids.rules.rewrite.CheckPrivileges; +import org.apache.doris.nereids.trees.expressions.Slot; +import org.apache.doris.nereids.trees.plans.Plan; +import org.apache.doris.nereids.trees.plans.PrepareCommandPlanner; +import org.apache.doris.nereids.trees.plans.commands.Command; +import org.apache.doris.nereids.trees.plans.commands.DescribeCommand; +import org.apache.doris.nereids.trees.plans.commands.ShowCreateTableCommand; +import org.apache.doris.nereids.trees.plans.commands.ShowDataCommand; +import org.apache.doris.nereids.trees.plans.commands.ShowPartitionsCommand; +import org.apache.doris.nereids.trees.plans.commands.ShowProcCommand; +import org.apache.doris.nereids.trees.plans.commands.ShowPythonPackagesCommand; +import org.apache.doris.nereids.trees.plans.commands.ShowQueryStatsCommand; +import org.apache.doris.nereids.trees.plans.commands.ShowTableCommand; +import org.apache.doris.nereids.trees.plans.commands.use.SwitchCommand; +import org.apache.doris.nereids.trees.plans.commands.use.UseCommand; +import org.apache.doris.qe.ConnectContext; +import org.apache.doris.qe.QueryState; +import org.apache.doris.qe.ResultSetMetaData; +import org.apache.doris.qe.SessionVariable; +import org.apache.doris.qe.StmtExecutor; +import org.apache.doris.qe.VariableMgr; + +import org.apache.arrow.flight.CallStatus; +import org.apache.arrow.util.AutoCloseables; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** Resolves result metadata without scheduling fragments or evaluating query expressions. */ +final class FlightSqlQuerySchema { + private FlightSqlQuerySchema() { + } + + static Schema analyze(ConnectContext context, String query) throws Exception { + synchronized (context) { + ConnectContext previousThreadContext = ConnectContext.get(); + StatementContext previousStatement = context.getStatementContext(); + SessionVariable previousSession = context.getSessionVariable(); + QueryState previousState = context.getState(); + StmtExecutor previousExecutor = context.getExecutor(); + String previousCatalog = context.getDefaultCatalog(); + String previousDatabase = context.getDatabase(); + List<StatementBase> statements = Collections.emptyList(); + try { + context.setThreadLocalInfo(); + context.setCommand(MysqlCommand.COM_QUERY); + // Parsing SET_VAR hints already mutates session variables. Isolate them even when parsing fails. + context.setSessionVariable(VariableMgr.cloneSessionVariable(previousSession)); + context.setState(new QueryState()); + context.setExecutor(null); + context.setStatementContext(null); + // Match execution's HTTP/plugin conversion before the dialect parser sees the SQL. + String converted = SqlDialectHelper.convertSqlByDialect(query, context.getSessionVariable()); + statements = new NereidsParser().parseSQL(converted, context.getSessionVariable()); + Map<String, String> scopedDatabases = new HashMap<>(); + if (statements.isEmpty()) { + throw CallStatus.UNIMPLEMENTED.withDescription( + "Schema discovery requires a statement").toRuntimeException(); + } + // JDBC clients commonly prefix their query with USE. Resolve that namespace only within this scope. + for (int i = 0; i < statements.size() - 1; ++i) { + Plan prefix = ((LogicalPlanAdapter) statements.get(i)).getLogicalPlan(); + if (!(prefix instanceof UseCommand) && !(prefix instanceof SwitchCommand)) { + throw CallStatus.UNIMPLEMENTED.withDescription( + "Schema discovery only supports USE or SWITCH before the result statement") + .toRuntimeException(); + } + resolveNamespace(context, prefix, scopedDatabases); + } + LogicalPlanAdapter statement = (LogicalPlanAdapter) statements.get(statements.size() - 1); + StatementContext statementContext = statement.getStatementContext(); + context.setStatementContext(statementContext); + statementContext.setParsedStatement(statement); + if (!statementContext.getPlaceholders().isEmpty()) { + throw CallStatus.UNIMPLEMENTED.withDescription( + "Flight SQL parameter binding is not supported").toRuntimeException(); + } + List<Field> fields = new ArrayList<>(); + Plan plan = statement.getLogicalPlan(); + if (plan instanceof Command) { + resolveNamespace(context, plan, scopedDatabases); + ResultSetMetaData metadata = commandMetadata(context, (Command) plan); + if (metadata == null) { + throw CallStatus.UNIMPLEMENTED.withDescription("Command result metadata is unavailable") + .toRuntimeException(); + } + // FE-local result sets are serialized as nullable strings by FlightSqlChannel. + for (Column column : metadata.getColumns()) { + fields.add(Field.nullable(column.getName(), new ArrowType.Utf8())); + } + if (fields.isEmpty()) { + switch (((Command) plan).stmtType()) { + case SET: + case USE: + case SWITCH: + case CREATE: + case ALTER: + case DROP: + case TRUNCATE: + // Only known no-row command categories have the protocol OK schema. + fields.add(Field.nullable("StatusResult", new ArrowType.Utf8())); + break; + default: + throw CallStatus.UNIMPLEMENTED.withDescription( Review Comment: Confirmed and fixed in 37f71c9306c. Ordinary EXPLAIN and PLAN REPLAYER now expose their fixed Flight headers without running the planner or creating a dump. EXPLAIN PLAN PROCESS remains explicitly UNIMPLEMENTED because StmtExecutor currently serializes that result only for MySQL; advertising its three columns for Flight would be incorrect. Added FE and Flight RPC coverage for the supported headers and a test for the PLAN PROCESS boundary. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
