morningman commented on code in PR #67966: URL: https://github.com/apache/doris/pull/67966#discussion_r4013706275
########## fe/fe-core/src/main/java/org/apache/doris/arrowflight/FlightSessionOptions.java: ########## @@ -0,0 +1,312 @@ +// 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.arrowflight; + +import org.apache.doris.analysis.SetType; +import org.apache.doris.common.AnalysisException; +import org.apache.doris.common.ErrorCode; +import org.apache.doris.common.util.Util; +import org.apache.doris.nereids.util.SqlLiteralUtils; +import org.apache.doris.qe.ConnectContext; +import org.apache.doris.qe.QueryState.MysqlStateType; +import org.apache.doris.qe.SqlModeHelper; +import org.apache.doris.qe.VarAttrDef; +import org.apache.doris.qe.VariableMgr; +import org.apache.doris.qe.VariableMgr.VarContext; + +import com.google.common.annotations.VisibleForTesting; +import org.apache.arrow.flight.SessionOptionValue; +import org.apache.arrow.flight.SessionOptionValueFactory; +import org.apache.arrow.flight.SessionOptionValueVisitor; +import org.apache.arrow.flight.SetSessionOptionsResult; +import org.apache.arrow.flight.SetSessionOptionsResult.ErrorValue; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +/** + * The session options of an Arrow Flight SQL session, as the SetSessionOptions and GetSessionOptions + * actions see them. + * + * <p>An option stands for the statement of the session that sets it: {@code catalog} is the current + * catalog ({@code SWITCH}), {@code schema} the current database ({@code USE}), and any other name is + * the session variable {@code SHOW VARIABLES} lists under that name ({@code SET SESSION}). Those two + * names are what the ADBC Flight SQL driver puts on the wire for {@code adbc.connection.catalog} and + * {@code adbc.connection.db_schema}, and what the Flight SQL JDBC driver sends for its + * {@code catalog} property; every other option name the drivers pass through as given. Setting an + * option runs its statement as a command of the session, so it is checked, audited and takes effect + * exactly as if the client had sent the statement. Reading the options back gives what + * {@code SHOW VARIABLES} shows, every value as a string: the one representation {@code SET} accepts + * back, whatever the Java type of the variable behind it. The names an option goes by are exactly + * the names that are read back, spelled the same: a client looks an option up in the result of + * GetSessionOptions by the name it set it under. And what is read back can be set back: the empty + * value, Flight's way of unsetting an option, puts a variable back to its default and the session + * back into no database, as it started -- the state GetSessionOptions reports as an empty + * {@code schema}, which sets it too. A session is always in some catalog, so {@code catalog} has no + * empty value. + * + * <p>The result of setting an option is one of the three {@link ErrorValue}s per name and nothing + * else, so the reason a value was refused only reaches the frontend log. + */ +public final class FlightSessionOptions { + private static final Logger LOG = LogManager.getLogger(FlightSessionOptions.class); + + /** The current catalog; set with {@code SWITCH}. */ + public static final String CATALOG = "catalog"; + /** The current database; set with {@code USE}. */ + public static final String SCHEMA = "schema"; + + private FlightSessionOptions() { + } + + /** + * Sets the options of one request, each on its own, and returns the error of every option that + * could not be set; an option absent from the result was set. {@code catalog} goes first and + * {@code schema} second, since {@code USE} names a database of the current catalog; the others + * follow in name order, so that a request applies the same way every time. Runs as a command of + * the session, under its command lock. + */ + public static Map<String, SetSessionOptionsResult.Error> set(ConnectContext ctx, + Map<String, SessionOptionValue> options) { + List<String> names = new ArrayList<>(options.keySet()); + Collections.sort(names); + names.remove(SCHEMA); + names.remove(CATALOG); + if (options.containsKey(SCHEMA)) { + names.add(0, SCHEMA); + } + if (options.containsKey(CATALOG)) { + names.add(0, CATALOG); + } + Map<String, SetSessionOptionsResult.Error> errors = new LinkedHashMap<>(); + for (String name : names) { + ErrorValue error = setOne(ctx, name, options.get(name)); + if (error != null) { + errors.put(name, new SetSessionOptionsResult.Error(error)); + } + } + return errors; + } + + /** Sets one option and returns why it could not be, or null when it was set. */ + @VisibleForTesting + static ErrorValue setOne(ConnectContext ctx, String name, SessionOptionValue value) { + if (CATALOG.equals(name)) { + return switchCatalog(ctx, value); + } + if (SCHEMA.equals(name)) { + return useDatabase(ctx, value); + } + return setVariable(ctx, name, value); + } + + /** + * The options of the session: the current catalog, the current database (an empty string when + * the session is in none, the value that puts it back there) and every session variable + * {@code SHOW VARIABLES} would list, in its text. + */ + public static Map<String, SessionOptionValue> get(ConnectContext ctx) { + Map<String, SessionOptionValue> options = new LinkedHashMap<>(); + options.put(CATALOG, SessionOptionValueFactory.makeSessionOptionValue(ctx.getDefaultCatalog())); + String database = ctx.getDatabase(); + options.put(SCHEMA, SessionOptionValueFactory.makeSessionOptionValue(database == null ? "" : database)); + for (List<String> row : VariableMgr.dump(SetType.SESSION, ctx.getSessionVariable(), null)) { + // A row is name, value, default value, changed. + options.put(row.get(0), SessionOptionValueFactory.makeSessionOptionValue(row.get(1))); + } + return options; + } + + private static ErrorValue switchCatalog(ConnectContext ctx, SessionOptionValue value) { + String catalog = value.acceptVisitor(STRING_VALUE); + // A session is always in some catalog, so there is no catalog the empty value could stand for. + if (catalog == null || catalog.isEmpty()) { + return ErrorValue.INVALID_VALUE; + } + // SWITCH checks the name's format before anything else and its error code would not reach + // here (see runStatement); the check is one of the value alone, so it is made here first. + try { + Util.checkCatalogAllRules(catalog); + } catch (AnalysisException e) { + LOG.warn("session option {} of Arrow Flight SQL connection {} could not be set, catalog name {} " + + "is malformed: {}", CATALOG, ctx.getConnectionId(), catalog, e.getMessage()); + return ErrorValue.INVALID_VALUE; + } + return runStatement(ctx, CATALOG, "SWITCH " + quoteIdentifier(catalog), ErrorCode.ERR_UNKNOWN_CATALOG); + } + + private static ErrorValue useDatabase(ConnectContext ctx, SessionOptionValue value) { + String database = value.acceptVisitor(STRING_VALUE); + if (database == null) { + return ErrorValue.INVALID_VALUE; + } + // No database: the state the session started in, which GetSessionOptions reports as the + // empty string and the empty value asks for back (Flight's way of unsetting an option; what + // the ADBC driver sends to erase one). No statement leads there -- there is no USE of + // nothing -- so it is not one the session runs, and there is nothing in it to check. + if (database.isEmpty()) { + ctx.clearDatabase(); Review Comment: Fixed in 674e491f9a9: the empty value and the empty string go through the gate every statement passes first -- inside a transaction they are refused with `ERROR`, the code `USE` and `SWITCH` produce there, and nothing is mutated (`FlightSessionOptions.useDatabase`). Test: `FlightSessionOptionsTest.testNoOptionChangesTheSessionInsideATransaction`, a real `BEGIN` / `ROLLBACK` on the session, both values, next to the refused `USE`, `SWITCH` and `SET` of the same transaction. For the record, the transaction itself was never at risk: it is fenced by database, and an unqualified DML after a clear fails with "No database selected" rather than writing elsewhere. The gap was one of consistency between the option that runs a statement and the one that does not. ########## fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java: ########## @@ -1076,7 +1121,13 @@ public void finalizeQuery() { // received after unregisterQuery(), causing the instance profile to be lost, so we should wait // for the profile before unregisterQuery(). updateProfile(true); - QeProcessorImpl.INSTANCE.unregisterQuery(context.queryId()); + QeProcessorImpl.INSTANCE.unregisterQuery(queryId()); Review Comment: Pre-existing, and untouched by this PR: the retry-with-replan path is not part of it. `handleQueryWithRetry`'s finally finalizes only when the result is on the frontend -- `if (context.isReturnResultFromLocal()) finalizeQuery();` -- and that line is the same on apache/master (`StmtExecutor.java:1207` there). On a Flight session the result has been on the backends before `registerQuery` on master as well: before #67900 `context.setReturnResultFromLocal(false)` preceded `handleQueryStmt()`, now `beforeQuery` precedes `registerQuery`. So a replan-retried attempt (cloud only, `NEED_REPLAN_ERRORS`) leaves its registration behind on master exactly as it does here, and `queryId()` returns `context.queryId()` for every non-deferred executor, as `finalizeQuery` always did. This PR's captured id covers the one case it introduced a gap for: a query finalized after the session ran another statement. Unregistering each failed attempt under its own id is a fix to the retry path -- finalize the attempt when `executeAndSendResult` fails with the result on the backends, since nothing will ever pull it -- and will be a separate PR. ########## fe/fe-core/src/main/java/org/apache/doris/arrowflight/protocol/FlightProtocolAdapter.java: ########## @@ -356,49 +380,88 @@ public void addDeferredExecutor(StmtExecutor executor) { } } + /** + * Takes every deferred executor out of the list and finalizes it: the next request of the + * session ({@link #beginRequest}), which stands in for the end of the previous query's DoGet, and + * the session's teardown. + */ public void closeDeferredExecutors() { - List<StmtExecutor> toClose; + List<StmtExecutor> taken; synchronized (deferredExecutors) { if (deferredExecutors.isEmpty()) { return; } - toClose = new ArrayList<>(deferredExecutors); + taken = new ArrayList<>(deferredExecutors); deferredExecutors.clear(); } - for (StmtExecutor deferredExecutor : toClose) { - try { - deferredExecutor.finalizeArrowFlightQuery(); - } catch (Throwable t) { - LOG.warn("failed to finalize deferred arrow flight executor", t); - } - } + finalizeDeferredExecutors(taken); } /** - * How long, in seconds, a sleeping connection may keep its deferred executors before the - * timeout checker finalizes them without killing the connection - * (Config.arrow_flight_deferred_query_idle_timeout_second). A Flight client that opens a - * session per query and never closes it would otherwise pin each deferred query's query queue - * slot and query registration until wait_timeout (8h by default). The bound is never shorter - * than the execution timeout the deferred query was run with: the client may still be pulling - * that query's results from the BE, which still needs the batch split source the coordinator - * holds. Returns -1 when the bound is disabled or nothing is deferred. + * Takes out of the list the deferred executors whose own bound has passed by {@code now} and + * returns them for the caller to finalize: the timeout checker's half of the exactly-once rule + * in the class comment. Each executor is judged and removed in the same critical section, by + * its own deadline -- when its query started plus {@link #deferredBoundMs} -- so that of the + * executors of one multi-statement request, each with a start and an execution timeout of its Review Comment: Fixed in 674e491f9a9: `finishStatement` counts a result left on the backends (`beforeQuery`, i.e. `returnResultFromLocal == false`) the same as one cached on the channel, so a query that is not the last statement of its request is refused with `ERR_ARROW_FLIGHT_SQL_MUST_ONLY_RESULT_STMT` before the next statement runs; a statement that failed produced no result and keeps its own error. `SELECT; SELECT` (equal schemas or not -- the schema comparison is never reached) and `SELECT; SET` are refused, and the SET does not run. Unit test in `FlightProtocolAdapterTest.testOnlyTheLastStatementOfARequestMayReturnAResult`, regression in `test_arrow_flight_session_lifecycle` ยง6 (both pairs, and `@multi` left untouched by the refused SET). Note that the check has only ever counted frontend-side results, since before #67900; this PR had not touched it. -- 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]
