[fix](flight-sql) Return analyzed result schemas from Prepare and GetSchema (#69867)

### What problem does this PR solve?

Issue Number: None

Related PR: #68631, #68768

Problem Summary:

**In short.** This ports branch-4.1 #68631 (also on branch-4.2) to
master. A Flight SQL `CreatePreparedStatement` now returns the analyzed
result schema of the statement, `GetSchema` is implemented for both
statements and prepared statements, and a prepared execution is checked
against the schema that Prepare advertised. Before this PR, Prepare
always returned a fixed one-column placeholder, and `GetSchema` answered
UNIMPLEMENTED. The analysis runs in the FE only and does not plan or
execute the query.

Original author: @Gabriel39

**Background.** Flight SQL clients learn the shape of a result before
they fetch it. The Arrow Flight SQL JDBC driver prepares every statement
and builds `ResultSetMetaData` from the dataset schema that
`CreatePreparedStatement` returns. ADBC exposes the same schema as
`execute_schema`, and `GetSchema` is the RPC that asks for it without
running anything.

**The problem, and what it cost.** On master:

| | master today | with this PR |
|---|---|---|
| `CreatePreparedStatement` dataset schema | always `ResultMeta: Utf8`,
whatever the statement | the statement's columns, names and Arrow types
as BE / FE will send them; statements whose schema cannot be told before
they run, and results with an untyped NULL column, keep the placeholder
(see below) |
| `GetSchema` for a statement or a prepared statement | UNIMPLEMENTED |
the analyzed schema, or UNIMPLEMENTED for a statement that keeps the
placeholder |
| Invalid SQL | fails only at execution | fails at Prepare with
INVALID_ARGUMENT |
| A query whose privilege check fails | fails only at execution | still
fails only at execution: Prepare keeps the placeholder, because only the
check on the executed plan decides (see below), and GetSchema answers
UNIMPLEMENTED without a schema |
| Prepared handle used after the session's catalog / database changed,
or after the result schema changed (e.g. another session's DDL) | runs
the unqualified SQL in whatever namespace the session has now, with
whatever columns come back | refused with NOT_FOUND "prepare again", and
the result left by that execution is dropped |
| Prepared handle presented by another session | resolved by the
identity written in the handle | must start with the caller's own
identity (from #68768) |
| TIMEV2 / VARBINARY columns in GetTables and Prepare schemas | `Null`,
while BE sends `float64` / `binary` | `Float64` / `Binary`, as BE sends
them |

A client that trusts the advertised schema saw one string column for
every query, and could not ask for a schema at all without executing.

**How this PR fixes it.**

- `arrowflight/FlightSqlQuerySchema` analyzes the statement on a clone
of the session variables: a USE / SWITCH prefix is resolved in a
temporary namespace, the analysis `StatementContext` is closed on the
same thread, and the connector statement scope it may have opened is
closed explicitly. Each field gets the same `doris_type` markers and
native Variant layout the execution side uses
(`FlightSqlSchemaHelper.withDorisTypeMetadata` / `nativeVariantField`),
and TIMESTAMPTZ carries the canonical session time zone that BE stamps
on its batches.
- Privileges are checked on the analyzed plan after the CTEs that
execution inlines have been inlined (`PullUpCteAnchor`, `CTEInline`, as
the rewriter does before its own check). A CTE used once is therefore
checked only on the columns the query reads, and a CTE used more than
once on every column it selects, which is also how execution checks
them. The rewriter normalizes the plan further before it checks
privileges. For example, it drops an ORDER BY under a derived table, so
the analyzed plan can read a column that the executed plan does not.
When the analyzed plan is denied, the statement keeps the placeholder
and execution decides. No schema is disclosed in either case.
- FE-side commands whose result header does not depend on running them
now report that header without running: HELP, SHOW SNAPSHOT, SHOW CREATE
TABLE (one header selector shared with `doRun`, including master's
stream, plugin-view and system-table cases), SHOW PROC `/` and
`/current_queries`, SHOW CONFIG, and EXPLAIN PLAN PROCESS (Rule / Before
/ After).
- `ShowStoragePolicyCommand.getMetaData()` returned the four-column
USING header for every form, while plain SHOW STORAGE POLICY answers the
seven columns of `StoragePolicy.STORAGE_META_DATA`. It now returns the
header of the form that runs. This also corrects the column definitions
that MySQL `COM_STMT_PREPARE` reports for the statement. No other SHOW
command declares a header that differs from the one it returns.
- WARM UP SELECT keeps the placeholder. It answers OK instead of its
statistics whenever planning proves its source empty: a predicate that
folds to false, no partition left after pruning, or partitions without
data.
- A result with an untyped NULL column keeps the placeholder too, at any
depth (`SELECT NULL`, `SELECT id, NULL FROM t`, `ARRAY(NULL)`).
`DorisArrowTypeMapping` gives NULL_TYPE Arrow `Null`, but BE builds such
a column as a `UInt8` flagged as a NULL literal and sends it as nullable
`Bool`, so an advertised `Null` would expire the handle at every
execution. A typed NULL (`CAST(NULL AS VARCHAR(20))`) is described as
its type.
- The prepared query kept in `FlightProtocolAdapter` records the SQL,
the namespace and the advertised schema. A prepared execution
re-analyzes the statement and compares the result with the advertised
schema before it runs, and compares the schema it actually produced
before it publishes a DoGet ticket. It is the session's next request
whether or not these checks pass: before them, it drops what the
previous request left, as a statement's execution does, while Prepare
and GetSchema are metadata requests and leave it. On a mismatch, the
handle is removed, the execution is cancelled, and the deferred
coordinator, the channel result and the endpoints are dropped. A handle
whose statement can no longer be described when it is re-analyzed, for
example after a column grant was revoked, expires the same way, and the
statement prepared again keeps the placeholder. A USE / SWITCH inside
the statement rebinds only that handle.
- `DorisArrowTypeMapping` maps TIMEV2 to `Float64` and VARBINARY to
`Binary`, the types BE sends. TIMESTAMPTZ keeps the fixed `"UTC"` in
GetTables, and AGG_STATE keeps `Null` there.
- `FeServiceClient` rejects a remote Doris metadata timeout whose
milliseconds do not fit in an int, instead of letting it wrap around.

Where master differs from branch-4.x:

- branch-4.x answers UNIMPLEMENTED at Prepare for statements whose
schema it cannot tell before they run: DESC, SHOW DATA / PARTITIONS /
QUERY STATS / PYTHON PACKAGES, other PROC nodes, COPY INTO, WARM UP
CLUSTER, multi-statement requests that do not start with USE / SWITCH,
`?` placeholders, and result types the FE cannot describe. Python ADBC
falls back to executing without a prepared statement, but the Arrow
Flight SQL JDBC driver prepares every statement and has no fallback, so
on master these statements would fail over JDBC.
`test_arrow_flight_session_lifecycle` runs DESC and `SET @multi = 2;
SELECT ...` through JDBC. Here, these statements keep today's
placeholder schema, and their handles are not checked at execution.
- branch-4.x checks privileges on the analyzed plan as it is, so a CTE
used once is checked on every column it selects, and any denial fails
Prepare. Here, CTEs are inlined first, and a remaining denial keeps the
placeholder.
- branch-4.x advertises the statistics header for WARM UP SELECT and
then rejects, as a schema change, an execution whose source turned out
empty. Here, WARM UP SELECT keeps the placeholder.
- branch-4.x advertises Arrow `Null` for an untyped NULL column, which
BE sends as `Bool`, so every prepared execution of such a query is
rejected as a schema change. Here, such a result keeps the placeholder.
- The analysis and the checks run under the session's command lock
(`FlightProtocolAdapter.runCommand` / `callCommand`) instead of
`synchronized (connectContext)`, and the namespace is saved and restored
with `getNameSpaceContext()` / `restoreNameSpaceContext()`.
- EXPLAIN PLAN PROCESS has a result over Flight on master, so it gets a
static header instead of UNIMPLEMENTED. PLAN REPLAYER and SHOW CREATE
PROCEDURE do not exist on master and are not carried over.
- SWITCH resolves the default database through the connector provider
(`Env.defaultDatabaseOnUse`, now shared with `Env.changeCatalog`)
instead of `EsExternalCatalog.DEFAULT_DB`.
- TIMESTAMPTZ fields carry the canonical session time zone that BE
stamps, where branch-4.x uses the raw session value.
- Only the two handle checks from #68768 that do not depend on parameter
binding are taken: the caller-prefix check and rebinding only the handle
that ran the USE / SWITCH. Parameter binding itself is not part of this
PR.
- `GetTableTypes`, the other half of #68631, is implemented by #67899.

Cost: every Prepare, every prepared execution and every `GetSchema` runs
one more FE analysis of the statement.

**Results.** On a local cluster built from this branch, every suite
under `arrow_flight_sql_p0` passes: the new `test_schema_only` and
`test_get_tables_schema`, plus the seven suites that run through Flight
JDBC, `test_arrow_flight_session_lifecycle` among them.
`datatype_p0/uuid/test_uuid_arrow_flight` passes too. `test_schema_only`
checks the advertised schema of each kind of statement against what its
execution actually returns. That covers queries, the side-effect-free
frontend-side commands executed as prepared statements (SHOW STORAGE
POLICY among them), a WARM UP SELECT on an empty table (executed in
cloud mode, the only mode that warms up internal tables), and a
`group_concat_state` AGG_STATE, which BE sends as Utf8.

Co-authored-by: Gabriel <liwenqiang@selectdb.com>
This commit is contained in:
Mingyu Chen (Rayner)andGabriel authored and GitHub committed 2026-11-04 10:22:17 +08:00
1 parent 30ce25dc1f
commit 1adeb3a680
21 files changed
+2725 -100

No files matched your search

@@ -47,10 +47,12 @@ import java.util.Map;
* columns from one of them reads the batches BE emits as that type. The values mirror
* {@code convert_to_arrow_type} in {@code be/src/format/arrow/arrow_row_batch.cpp}, which alone
* decides the shape of the data on the wire; the cells known to disagree with it (a literal
* {@code "UTC"} zone on TIMESTAMPTZ, {@code Null} for TIMEV2 / VARBINARY / AGG_STATE) are kept
* deliberately until the BE Arrow type layer is reworked, and are then corrected here, once, against
* a golden shared with BE (#67577). {@code DorisArrowTypeMappingTest} records every cell as it
* stands, so a change to any of them is a change to that test as well.
* {@code "UTC"} zone on TIMESTAMPTZ, {@code Null} for AGG_STATE) are kept deliberately until the BE
* Arrow type layer is reworked, and are then corrected here, once, against a golden shared with BE
* (#67577). TIMEV2 and VARBINARY already follow BE (float64 and binary): a prepared statement
* advertises its result schema from here and is checked against the schema of the batches BE emits.
* {@code DorisArrowTypeMappingTest} records every cell as it stands, so a change to any of them is a
* change to that test as well.
*/
public final class DorisArrowTypeMapping {
@@ -77,6 +79,7 @@ public final class DorisArrowTypeMapping {
case FLOAT:
return new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE);
case DOUBLE:
case TIMEV2:
return new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE);
case LARGEINT:
case VARCHAR:
@@ -125,6 +128,7 @@ public final class DorisArrowTypeMapping {
case HLL:
case BITMAP:
case QUANTILE_STATE:
case VARBINARY:
return new ArrowType.Binary();
case MAP:
return new ArrowType.Map(false);
@@ -20,7 +20,9 @@
package org.apache.doris.arrowflight;
import org.apache.doris.arrowflight.FlightSqlQuerySchema.UndiscoverableSchemaException;
import org.apache.doris.arrowflight.protocol.FlightProtocolAdapter;
import org.apache.doris.arrowflight.protocol.FlightProtocolAdapter.PreparedQuery;
import org.apache.doris.arrowflight.results.FlightSqlEndpointsLocation;
import org.apache.doris.arrowflight.results.FlightSqlResultCacheEntry;
import org.apache.doris.arrowflight.sessions.FlightSessionsManager;
@@ -29,9 +31,11 @@ import org.apache.doris.common.Status;
import org.apache.doris.common.util.DebugUtil;
import org.apache.doris.common.util.Util;
import org.apache.doris.mysql.MysqlCommand;
import org.apache.doris.nereids.glue.LogicalPlanAdapter;
import org.apache.doris.qe.ConnectContext;
import org.apache.doris.qe.QueryState;
import org.apache.doris.qe.QueryState.MysqlStateType;
import org.apache.doris.qe.StmtExecutor;
import org.apache.doris.thrift.TStatusCode;
import org.apache.doris.thrift.TUniqueId;
@@ -93,6 +97,8 @@ import org.apache.arrow.util.AutoCloseables;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.ipc.WriteChannel;
import org.apache.arrow.vector.ipc.message.MessageSerializer;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.Schema;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
@@ -102,7 +108,6 @@ import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.Channels;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Base64;
import java.util.Collections;
import java.util.List;
@@ -132,6 +137,10 @@ public class DorisFlightSqlProducer implements FlightSqlProducer, AutoCloseable
private final SqlInfoBuilder sqlInfoBuilder;
private final FlightSessionsManager flightSessionsManager;
private final ExecutorService executorService = Executors.newFixedThreadPool(100);
// What Prepare advertises for a statement whose result schema it cannot tell: a client takes a
// prepared statement with an empty dataset schema for an update (arrow#38911).
private static final Schema PLACEHOLDER_DATASET_SCHEMA = new Schema(
Collections.singletonList(Field.nullable("ResultMeta", new ArrowType.Utf8())));
public DorisFlightSqlProducer(final Location location, FlightSessionsManager flightSessionsManager) {
this.location = location;
@@ -202,10 +211,8 @@ public class DorisFlightSqlProducer implements FlightSqlProducer, AutoCloseable
final StreamListener<Result> listener) {
executorService.submit(() -> {
try {
String[] handleParts = request.getPreparedStatementHandle().toStringUtf8().split(":");
String executedPeerIdentity = handleParts[0];
String preparedStatementId = handleParts[1];
ConnectContext connectContext = flightSessionsManager.getConnectContext(executedPeerIdentity);
String preparedStatementId = preparedStatementId(context, request.getPreparedStatementHandle());
ConnectContext connectContext = flightSessionsManager.getConnectContext(context.peerIdentity());
FlightProtocolAdapter.of(connectContext).runCommand(connectContext,
() -> connectContext.removePreparedQuery(preparedStatementId));
} catch (final Throwable e) {
@@ -216,7 +223,21 @@ public class DorisFlightSqlProducer implements FlightSqlProducer, AutoCloseable
});
}
private FlightInfo executeQueryStatement(String peerIdentity, ConnectContext connectContext, String query,
// What a request answers GetFlightInfo with, and the executor of its last statement and where that
// statement left its result -- taken before the processor's close clears them from the session.
private static final class ExecutedQuery {
private final FlightInfo info;
private final StmtExecutor executor;
private final boolean resultFromLocal;
private ExecutedQuery(FlightInfo info, StmtExecutor executor, boolean resultFromLocal) {
this.info = info;
this.executor = executor;
this.resultFromLocal = resultFromLocal;
}
}
private ExecutedQuery executeQueryStatement(String peerIdentity, ConnectContext connectContext, String query,
final FlightDescriptor descriptor) {
try {
Preconditions.checkState(null != connectContext);
@@ -232,6 +253,7 @@ public class DorisFlightSqlProducer implements FlightSqlProducer, AutoCloseable
throw new RuntimeException("after executeQueryStatement handleQuery");
}
StmtExecutor executor = connectContext.getExecutor();
if (connectContext.isReturnResultFromLocal()) {
// set/use etc. stmt returns an OK result by default.
if (connectContext.getFlightSqlChannel().resultNum() == 0) {
@@ -242,9 +264,9 @@ public class DorisFlightSqlProducer implements FlightSqlProducer, AutoCloseable
final ByteString handle = ByteString.copyFromUtf8(peerIdentity + ":" + queryId);
TicketStatementQuery ticketStatement = TicketStatementQuery.newBuilder()
.setStatementHandle(handle).build();
return getFlightInfoForSchema(ticketStatement, descriptor,
return new ExecutedQuery(getFlightInfoForSchema(ticketStatement, descriptor,
connectContext.getFlightSqlChannel().getResult(queryId).getVectorSchemaRoot()
.getSchema());
.getSchema()), executor, true);
} else {
// A Flight Sql request can only contain one statement that returns result,
// otherwise expected thrown exception during execution.
@@ -257,9 +279,10 @@ public class DorisFlightSqlProducer implements FlightSqlProducer, AutoCloseable
peerIdentity + ":" + DebugUtil.printId(connectContext.queryId()));
TicketStatementQuery ticketStatement = TicketStatementQuery.newBuilder()
.setStatementHandle(handle).build();
return getFlightInfoForSchema(ticketStatement, descriptor, connectContext.getFlightSqlChannel()
.getResult(DebugUtil.printId(connectContext.queryId())).getVectorSchemaRoot()
.getSchema());
return new ExecutedQuery(getFlightInfoForSchema(ticketStatement, descriptor,
connectContext.getFlightSqlChannel()
.getResult(DebugUtil.printId(connectContext.queryId()))
.getVectorSchemaRoot().getSchema()), executor, true);
}
} else {
// Now only query stmt will pull results from BE.
@@ -309,7 +332,8 @@ public class DorisFlightSqlProducer implements FlightSqlProducer, AutoCloseable
endpoints.add(new FlightEndpoint(ticket, location));
}
// TODO Set in BE callback after query end, Client will not callback.
return new FlightInfo(flightSQLConnectProcessor.getArrowSchema(), descriptor, endpoints, -1, -1);
return new ExecutedQuery(new FlightInfo(flightSQLConnectProcessor.getArrowSchema(), descriptor,
endpoints, -1, -1), executor, false);
}
}
} catch (Throwable e) {
@@ -350,7 +374,7 @@ public class DorisFlightSqlProducer implements FlightSqlProducer, AutoCloseable
ConnectContext connectContext = flightSessionsManager.getConnectContext(context.peerIdentity());
return FlightProtocolAdapter.of(connectContext).callCommand(connectContext,
() -> executeQueryStatement(context.peerIdentity(), connectContext, request.getQuery(),
descriptor));
descriptor).info);
} catch (FlightRuntimeException e) {
// Already carries the status meant for the client - UNAVAILABLE from the session's
// command lock or from an incremental window that is not ready (queryFailure, with its
@@ -368,19 +392,151 @@ public class DorisFlightSqlProducer implements FlightSqlProducer, AutoCloseable
@Override
public FlightInfo getFlightInfoPreparedStatement(final CommandPreparedStatementQuery command,
final CallContext context, final FlightDescriptor descriptor) {
String[] handleParts = command.getPreparedStatementHandle().toStringUtf8().split(":");
String executedPeerIdentity = handleParts[0];
String preparedStatementId = handleParts[1];
ConnectContext connectContext = flightSessionsManager.getConnectContext(executedPeerIdentity);
String preparedStatementId = preparedStatementId(context, command.getPreparedStatementHandle());
ConnectContext connectContext = flightSessionsManager.getConnectContext(context.peerIdentity());
return FlightProtocolAdapter.of(connectContext).callCommand(connectContext,
() -> executeQueryStatement(executedPeerIdentity, connectContext,
connectContext.getPreparedQuery(preparedStatementId), descriptor));
() -> executePreparedStatement(context.peerIdentity(), connectContext, preparedStatementId,
descriptor));
}
private FlightInfo executePreparedStatement(String peerIdentity, ConnectContext connectContext,
String preparedStatementId, FlightDescriptor descriptor) {
// A prepared execution is the session's next request even when its handle fails the checks below, so
// what the previous request left is dropped before them, as executeQueryStatement does for a statement
// (which then finds nothing left). Prepare and GetSchema are metadata requests and leave it.
FlightProtocolAdapter.of(connectContext).beginRequest();
PreparedQuery prepared = preparedQuery(connectContext, preparedStatementId);
if (prepared.getSchema() == null) {
// Prepare could not tell this statement's result schema and advertised the placeholder, so
// its executions have no schema to be checked against.
return executeQueryStatement(peerIdentity, connectContext, prepared.getSql(), descriptor).info;
}
Schema schema = checkPreparedSchema(connectContext, preparedStatementId, prepared);
ExecutedQuery executed = executeQueryStatement(peerIdentity, connectContext, prepared.getSql(), descriptor);
// Another session's DDL can change the result after reanalysis but before execution
// acquires table locks. Reject the actual schema before publishing a DoGet ticket.
boolean matches = executed.resultFromLocal
// A frontend-side result is a header of nullable strings, exactly as Prepare advertised it.
? schema.equals(executed.info.getSchema())
// BE labels may carry the SQL text of an expression or a TVF prefix, and constant folding
// can narrow nullability: compare the final planner's labels and the positional types.
: FlightSqlQuerySchema.matchesExecutionSchema(schema, executed.info.getSchema(),
((LogicalPlanAdapter) executed.executor.getParsedStmt()).getColLabels());
if (!matches) {
connectContext.removePreparedQuery(preparedStatementId);
try {
executed.executor.cancel(new Status(TStatusCode.CANCELLED, "prepared statement schema changed"));
} catch (Exception e) {
LOG.warn("Failed to cancel Flight query after prepared schema changed", e);
} finally {
// No DoGet will pull this result: drop what the request left behind as the next request
// would -- its deferred coordinator, the result cached on the channel and its endpoints.
FlightProtocolAdapter.of(connectContext).beginRequest();
}
throw CallStatus.NOT_FOUND.withDescription("Prepared statement schema changed; prepare again")
.toRuntimeException();
}
// A successful USE or SWITCH in this statement moves its namespace. Rebind only this handle;
// the other handles still reject a namespace change made outside them.
prepared.rebind(connectContext.getDefaultCatalog(), connectContext.getDatabase());
return executed.info;
}
@Override
public SchemaResult getSchemaStatement(final CommandStatementQuery command, final CallContext context,
final FlightDescriptor descriptor) {
throw CallStatus.UNIMPLEMENTED.withDescription("getSchemaStatement unimplemented").toRuntimeException();
ConnectContext connectContext = flightSessionsManager.getConnectContext(context.peerIdentity());
return FlightProtocolAdapter.of(connectContext).callCommand(connectContext,
() -> new SchemaResult(discoverSchema(connectContext, command.getQuery())));
}
@Override
public SchemaResult getSchemaPreparedStatement(final CommandPreparedStatementQuery command,
final CallContext context, final FlightDescriptor descriptor) {
String preparedStatementId = preparedStatementId(context, command.getPreparedStatementHandle());
ConnectContext connectContext = flightSessionsManager.getConnectContext(context.peerIdentity());
return FlightProtocolAdapter.of(connectContext).callCommand(connectContext, () -> {
PreparedQuery prepared = preparedQuery(connectContext, preparedStatementId);
if (prepared.getSchema() == null) {
throw CallStatus.UNIMPLEMENTED.withDescription(
"The result schema of this statement is only known once it runs").toRuntimeException();
}
return new SchemaResult(checkPreparedSchema(connectContext, preparedStatementId, prepared));
});
}
// The id of a prepared statement this session created: a handle is "<peer identity>:<id>".
private static String preparedStatementId(CallContext context, ByteString handleBytes) {
String prefix = context.peerIdentity() + ":";
String handle = handleBytes.toStringUtf8();
if (!handle.startsWith(prefix) || handle.length() == prefix.length()) {
throw CallStatus.INVALID_ARGUMENT.withDescription("Invalid prepared statement handle").toRuntimeException();
}
return handle.substring(prefix.length());
}
private static PreparedQuery preparedQuery(ConnectContext connectContext, String preparedStatementId) {
PreparedQuery prepared = connectContext.getPreparedQuery(preparedStatementId);
if (prepared == null) {
throw CallStatus.NOT_FOUND.withDescription("Prepared statement expired; prepare again")
.toRuntimeException();
}
// A checked handle must not run its unqualified SQL in another namespace than its schema was told in.
if (prepared.getSchema() != null
&& !prepared.isPreparedIn(connectContext.getDefaultCatalog(), connectContext.getDatabase())) {
connectContext.removePreparedQuery(preparedStatementId);
throw CallStatus.NOT_FOUND
.withDescription("Prepared statement expired; prepare again in the current namespace")
.toRuntimeException();
}
return prepared;
}
// Execution parses the SQL again in the current session. Never silently replace the schema Prepare
// advertised when a setting such as sql_mode or time_zone changes the result. The check compares what
// Prepare would advertise now, so a handle also expires when only running its statement could tell the
// schema any more (a column grant revoked since Prepare, a DDL that turned a result column into an untyped
// NULL); prepared again, the statement gets the placeholder and runs unchecked.
private static Schema checkPreparedSchema(ConnectContext connectContext, String preparedStatementId,
PreparedQuery prepared) {
Schema schema = advertisedSchema(connectContext, prepared.getSql());
if (!prepared.getSchema().equals(schema)) {
connectContext.removePreparedQuery(preparedStatementId);
throw CallStatus.NOT_FOUND.withDescription("Prepared statement schema changed; prepare again")
.toRuntimeException();
}
return schema;
}
// The schema Prepare advertises and records for query: null stands for the placeholder of a statement
// whose result schema only running it can tell.
private static Schema advertisedSchema(ConnectContext connectContext, String query) {
try {
return analyzeSchema(connectContext, query);
} catch (UndiscoverableSchemaException e) {
return null;
}
}
// The schema GetSchema answers for query.
private static Schema discoverSchema(ConnectContext connectContext, String query) {
try {
return analyzeSchema(connectContext, query);
} catch (UndiscoverableSchemaException e) {
throw CallStatus.UNIMPLEMENTED.withDescription(e.getMessage()).withCause(e).toRuntimeException();
}
}
private static Schema analyzeSchema(ConnectContext connectContext, String query)
throws UndiscoverableSchemaException {
try {
return FlightSqlQuerySchema.analyze(connectContext, query);
} catch (UndiscoverableSchemaException | FlightRuntimeException e) {
throw e;
} catch (Exception e) {
throw CallStatus.INVALID_ARGUMENT.withDescription("Cannot determine query schema: " + e.getMessage())
.withCause(e).toRuntimeException();
}
}
@Override
@@ -426,43 +582,28 @@ public class DorisFlightSqlProducer implements FlightSqlProducer, AutoCloseable
private void createPreparedStatement(final ActionCreatePreparedStatementRequest request, String peerIdentity,
ConnectContext connectContext, final StreamListener<Result> listener) {
String preparedStatementId = UUID.randomUUID().toString();
try {
connectContext.setCommand(MysqlCommand.COM_QUERY);
final String query = request.getQuery();
String preparedStatementId = UUID.randomUUID().toString();
// ADBC ExecuteSchema reads this dataset schema directly without calling GetSchema. Analyze
// before registering a handle so that a statement that cannot run leaves none behind. A
// statement whose schema only running it can tell keeps the placeholder: the Arrow Flight
// SQL JDBC driver prepares every statement and does not fall back when Prepare refuses one.
Schema schema = advertisedSchema(connectContext, query);
final ByteString handle = ByteString.copyFromUtf8(peerIdentity + ":" + preparedStatementId);
connectContext.addPreparedQuery(preparedStatementId, query);
// Close the temporary VectorSchemaRoot after extracting its Schema, otherwise the
// off-heap buffers backing its vectors are leaked on every prepare (FE direct memory leak).
final Schema parameterSchema;
try (VectorSchemaRoot emptyVectorSchemaRoot =
new VectorSchemaRoot(new ArrayList<>(), new ArrayList<>())) {
parameterSchema = emptyVectorSchemaRoot.getSchema();
}
// TODO FE does not have the ability to convert root fragment output expr into arrow schema.
// However, the metaData schema returned by createPreparedStatement is usually not used by the client,
// but it cannot be empty, otherwise it will be mistaken by the client as an updata statement.
// see: https://github.com/apache/arrow/issues/38911
final Schema metaData;
try (VectorSchemaRoot metaSchemaRoot = connectContext.getFlightSqlChannel()
.createOneOneSchemaRoot("ResultMeta", "UNIMPLEMENTED")) {
metaData = metaSchemaRoot.getSchema();
}
listener.onNext(new Result(
Any.pack(buildCreatePreparedStatementResult(handle, parameterSchema, metaData)).toByteArray()));
} catch (Exception e) {
String errMsg = "create prepared statement failed, " + e.getMessage() + ", " + Util.getRootCauseMessage(
e) + ", error code: " + connectContext.getState().getErrorCode() + ", error msg: "
+ connectContext.getState().getErrorMessage();
LOG.error(errMsg, e);
listener.onError(CallStatus.INTERNAL.withDescription(errMsg).withCause(e).toRuntimeException());
Result result = new Result(Any.pack(buildCreatePreparedStatementResult(handle,
new Schema(Collections.emptyList()), schema == null ? PLACEHOLDER_DATASET_SCHEMA : schema))
.toByteArray());
connectContext.addPreparedQuery(preparedStatementId, new PreparedQuery(query, schema,
connectContext.getDefaultCatalog(), connectContext.getDatabase()));
listener.onNext(result);
} catch (final Throwable e) {
connectContext.removePreparedQuery(preparedStatementId);
LOG.warn("create prepared statement failed", e);
listener.onError(e instanceof FlightRuntimeException ? e
: CallStatus.INTERNAL.withDescription("Create prepared statement failed: " + e.getMessage())
.withCause(e).toRuntimeException());
return;
} catch (final Throwable t) {
listener.onError(CallStatus.INTERNAL.withDescription("Unknown error: " + t).toRuntimeException());
return;
} finally {
connectContext.setCommand(MysqlCommand.COM_SLEEP);
}
listener.onCompleted();
}
@@ -0,0 +1,447 @@
// 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.StatementBase;
import org.apache.doris.arrow.DorisArrowTypeMapping;
import org.apache.doris.catalog.AggStateType;
import org.apache.doris.catalog.ArrayType;
import org.apache.doris.catalog.Column;
import org.apache.doris.catalog.Env;
import org.apache.doris.catalog.MapType;
import org.apache.doris.catalog.NameSpaceContext;
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.common.util.TimeUtils;
import org.apache.doris.datasource.CatalogIf;
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.exceptions.AnalysisException;
import org.apache.doris.nereids.glue.LogicalPlanAdapter;
import org.apache.doris.nereids.jobs.JobContext;
import org.apache.doris.nereids.parser.NereidsParser;
import org.apache.doris.nereids.parser.SqlDialectHelper;
import org.apache.doris.nereids.rules.rewrite.CTEInline;
import org.apache.doris.nereids.rules.rewrite.CheckPrivileges;
import org.apache.doris.nereids.rules.rewrite.PullUpCteAnchor;
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.DeleteFromCommand;
import org.apache.doris.nereids.trees.plans.commands.DescribeCommand;
import org.apache.doris.nereids.trees.plans.commands.ExplainCommand;
import org.apache.doris.nereids.trees.plans.commands.HelpCommand;
import org.apache.doris.nereids.trees.plans.commands.KillCommand;
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.ShowSnapshotCommand;
import org.apache.doris.nereids.trees.plans.commands.ShowTableCommand;
import org.apache.doris.nereids.trees.plans.commands.TransactionCommand;
import org.apache.doris.nereids.trees.plans.commands.UpdateCommand;
import org.apache.doris.nereids.trees.plans.commands.insert.BatchInsertIntoTableCommand;
import org.apache.doris.nereids.trees.plans.commands.insert.InsertIntoTVFCommand;
import org.apache.doris.nereids.trees.plans.commands.insert.InsertIntoTableCommand;
import org.apache.doris.nereids.trees.plans.commands.insert.InsertOverwriteTableCommand;
import org.apache.doris.nereids.trees.plans.commands.insert.WarmupSelectCommand;
import org.apache.doris.nereids.trees.plans.commands.merge.MergeIntoCommand;
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.ShowResultSetMetaData;
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;
import java.util.Objects;
/**
* Resolves the result schema of a Flight SQL statement without scheduling fragments or evaluating
* query expressions: what Prepare advertises and GetSchema answers, and what a prepared execution
* is checked against.
*/
final class FlightSqlQuerySchema {
private FlightSqlQuerySchema() {
}
/**
* The statement's result schema can only be told by running it: DESCRIBE and the other commands
* whose result columns depend on what they find (WARM UP SELECT answers OK instead of its
* statistics when its source turns out empty), a request whose statements before the last are
* not USE or SWITCH, a query with parameter placeholders, a result type this frontend cannot
* describe as the backend sends it (an untyped NULL column among them, at any depth), a query
* that the privilege check of its analyzed plan denies (only the check of the executed plan
* decides, see checkPrivileges). Prepare advertises the placeholder schema for such a statement
* and does not check its executions; GetSchema answers UNIMPLEMENTED.
*/
static final class UndiscoverableSchemaException extends Exception {
UndiscoverableSchemaException(String message) {
super(message);
}
}
/**
* The result schema of {@code query} in the session's current namespace and settings. It runs as
* a command of the session ({@link org.apache.doris.arrowflight.protocol.FlightProtocolAdapter#callCommand},
* which makes the context the thread's own), and puts back everything of the session it borrows.
*/
static Schema analyze(ConnectContext context, String query) throws Exception {
StatementContext previousStatement = context.getStatementContext();
SessionVariable previousSession = context.getSessionVariable();
QueryState previousState = context.getState();
StmtExecutor previousExecutor = context.getExecutor();
NameSpaceContext previousNamespace = context.getNameSpaceContext();
List<StatementBase> statements = Collections.emptyList();
try {
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());
try {
statements = new NereidsParser().parseSQL(converted, context.getSessionVariable());
} catch (Exception convertedError) {
if (!context.getSessionVariable().isRetryOriginSqlOnConvertFail() || converted.equals(query)) {
throw convertedError;
}
// Match execution's parse fallback while discarding any failed parser context.
StatementContext failed = context.getStatementContext();
if (failed != null) {
failed.close();
context.setStatementContext(null);
}
statements = new NereidsParser().parseSQL(query, context.getSessionVariable());
}
if (statements.isEmpty()) {
throw new UndiscoverableSchemaException("Schema discovery requires a statement");
}
Map<String, String> scopedDatabases = new HashMap<>();
// 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 new UndiscoverableSchemaException(
"Schema discovery only supports USE or SWITCH before the result statement");
}
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 new UndiscoverableSchemaException("Flight SQL parameter binding is not supported");
}
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 new UndiscoverableSchemaException("Command result metadata is unavailable");
}
// 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 new UndiscoverableSchemaException(
"Result metadata is unavailable without executing this command");
}
}
} else {
// BindResultSink preserves SQL column labels only in query mode, just as
// StmtExecutor does; command mode would infer synthetic aliases instead.
context.getState().setIsQuery(true);
PrepareCommandPlanner planner = new PrepareCommandPlanner(statementContext);
planner.plan(statement, context.getSessionVariable().toThrift());
CascadesContext cascades = planner.getCascadesContext();
Plan analyzed = cascades.getRewritePlan();
checkPrivileges(analyzed, cascades.getCurrentJobContext());
// The zone BE stamps on TIMESTAMPTZ: the one it is sent, see CoordinatorContext.
String timezone = TimeUtils.getCanonicalTimeZone(context.getSessionVariable().getTimeZone());
for (Slot slot : analyzed.getOutput()) {
fields.add(field(slot.getName(), slot.getDataType().toCatalogDataType(), slot.nullable(),
timezone));
}
}
return new Schema(fields);
} finally {
try {
List<AutoCloseable> resources = new ArrayList<>();
for (StatementBase statement : statements) {
addAnalysisResources(resources, ((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)) {
addAnalysisResources(resources, current);
}
AutoCloseables.close(resources);
} finally {
context.setStatementContext(previousStatement);
context.setSessionVariable(previousSession);
context.setState(previousState);
context.setExecutor(previousExecutor);
// Restoring the captured namespace does not contact a catalog.
context.restoreNameSpaceContext(previousNamespace);
context.setCommand(MysqlCommand.COM_SLEEP);
}
}
}
/**
* Checks privileges on {@code analyzed}: PrepareCommandPlanner stops before the rewrite phase that
* checks them for execution. As that phase does before its check, the CTEs execution inlines are
* inlined first, so a CTE used once is checked on the columns the query reads instead of on every
* column it selects. The rewriter normalizes the plan further before its check (it drops an ORDER
* BY under a derived table, for one), so the analyzed plan can still read a column the executed
* plan does not: a denial here does not show that execution is denied, and only running the
* statement can tell. Such a statement is undiscoverable rather than refused, and execution
* decides as it did before Prepare analyzed anything.
*/
private static void checkPrivileges(Plan analyzed, JobContext jobContext) throws UndiscoverableSchemaException {
Plan inlined = new CTEInline().rewriteRoot(new PullUpCteAnchor().rewriteRoot(analyzed, jobContext),
jobContext);
try {
new CheckPrivileges().rewriteRoot(inlined, jobContext);
} catch (AnalysisException e) {
throw new UndiscoverableSchemaException("Only execution can decide the privileges of this query: "
+ e.getMessage());
}
}
private static void addAnalysisResources(List<AutoCloseable> resources, StatementContext statementContext) {
resources.add(statementContext);
// close() leaves the connector statement scope of a Flight statement to the query-finish
// cleanup of the query that runs it. Nothing runs an analysis, so its scope is closed here.
resources.add(statementContext::closeConnectorStatementScope);
}
static boolean matchesExecutionSchema(Schema prepared, Schema actual, List<String> columnLabels) {
List<Field> expectedFields = prepared.getFields();
List<Field> actualFields = actual.getFields();
if (columnLabels == null || expectedFields.size() != actualFields.size()
|| expectedFields.size() != columnLabels.size()
|| !prepared.getCustomMetadata().equals(actual.getCustomMetadata())) {
return false;
}
for (int i = 0; i < expectedFields.size(); ++i) {
Field expected = expectedFields.get(i);
// Compare the final planner's semantic labels: BE labels may instead contain SQL
// text or type_name_index. Ignoring all names would hide a concurrent column rename.
if (!expected.getName().equals(columnLabels.get(i))
|| !matchesExecutionField(expected, actualFields.get(i))) {
return false;
}
}
return true;
}
private static boolean matchesExecutionField(Field expected, Field actual) {
// Rewrites can prove an expression non-null (e.g. a folded CAST). Such narrowing is
// compatible with Prepare's nullable field; the reverse violates its advertised contract.
if ((!expected.isNullable() && actual.isNullable()) || !expected.getType().equals(actual.getType())
|| !expected.getMetadata().equals(actual.getMetadata())
|| !Objects.equals(expected.getFieldType().getDictionary(), actual.getFieldType().getDictionary())
|| expected.getChildren().size() != actual.getChildren().size()) {
return false;
}
for (int i = 0; i < expected.getChildren().size(); ++i) {
Field child = expected.getChildren().get(i);
Field actualChild = actual.getChildren().get(i);
if (!child.getName().equals(actualChild.getName()) || !matchesExecutionField(child, actualChild)) {
return false;
}
}
return true;
}
private static ResultSetMetaData commandMetadata(ConnectContext context, Command command) throws Exception {
// These getters depend on execution-time state or remote responses, and WARM UP SELECT answers
// OK instead of its statistics when planning proves its source empty (a predicate folded to false,
// no partition left, or none with data). 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 || command instanceof WarmupSelectCommand) {
throw new UndiscoverableSchemaException("Command schema requires execution-time metadata");
}
if (command instanceof HelpCommand) {
return ((HelpCommand) command).getMetaData(context);
} else if (command instanceof ShowSnapshotCommand) {
return ((ShowSnapshotCommand) command).getMetaData(context);
} else 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);
}
if (command instanceof ExplainCommand) {
// The headers StmtExecutor sends for the plan process and for the plan text.
if (((ExplainCommand) command).showPlanProcess()) {
return ShowResultSetMetaData.builder()
.addColumn(new Column("Rule", Type.STRING))
.addColumn(new Column("Before", Type.STRING))
.addColumn(new Column("After", Type.STRING))
.build();
}
return stringMetadata("Explain String(Nereids Planner)");
}
ResultSetMetaData metadata = command.getResultSetMetaData();
// Concrete DML commands return OK; a subclass that declares its own header keeps it.
if (metadata != null && metadata.getColumnCount() == 0 && (command instanceof InsertIntoTableCommand
|| command instanceof InsertOverwriteTableCommand || command instanceof BatchInsertIntoTableCommand
|| command instanceof InsertIntoTVFCommand || command instanceof UpdateCommand
|| command instanceof DeleteFromCommand || command instanceof MergeIntoCommand
|| command instanceof KillCommand || command instanceof TransactionCommand)) {
return stringMetadata("StatusResult");
}
return metadata;
}
private static ResultSetMetaData stringMetadata(String name) {
return ShowResultSetMetaData.builder().addColumn(new Column(name, Type.STRING)).build();
}
private static void resolveNamespace(ConnectContext context, Plan plan, Map<String, String> scopedDatabases)
throws Exception {
if (plan instanceof UseCommand) {
UseCommand use = (UseCommand) plan;
String catalog = use.getCatalogName() == null ? context.getDefaultCatalog() : use.getCatalogName();
CatalogIf catalogObject = context.getCatalog(catalog);
if (catalogObject == null || !context.getEnv().getAccessManager()
.checkDbPriv(context, catalog, use.getDatabaseName(), PrivPredicate.SHOW)) {
throw CallStatus.UNAUTHORIZED.withDescription("Database access denied").toRuntimeException();
}
catalogObject.getDbOrAnalysisException(use.getDatabaseName());
if (use.getCatalogName() != null && !context.getDatabase().isEmpty()) {
scopedDatabases.put(context.getDefaultCatalog(), context.getDatabase());
}
context.changeDefaultCatalog(catalog);
context.setDatabase(use.getDatabaseName());
} else if (plan instanceof SwitchCommand) {
String catalog = ((SwitchCommand) plan).getCatalogName();
CatalogIf catalogObject = context.getCatalog(catalog);
if (catalogObject == null || !context.getEnv().getAccessManager()
.checkCtlPriv(context, catalog, PrivPredicate.SHOW)) {
throw CallStatus.UNAUTHORIZED.withDescription("Catalog access denied").toRuntimeException();
}
// Mirror Env.changeCatalog, keeping remembered databases local to this analysis.
if (!context.getDatabase().isEmpty()) {
scopedDatabases.put(context.getDefaultCatalog(), context.getDatabase());
}
String database = scopedDatabases.getOrDefault(catalog, context.getLastDBOfCatalog(catalog));
context.changeDefaultCatalog(catalog);
if (database != null && !database.isEmpty()) {
context.setDatabase(database);
}
Env.defaultDatabaseOnUse(catalogObject).ifPresent(context::setDatabase);
}
}
/**
* The Arrow field BE sends for a result column of {@code type} ({@code convert_to_arrow_type} in
* be/src/format/arrow/arrow_row_batch.cpp), with the {@code doris_type} markers execution adds to
* it, at every depth ({@link FlightSqlSchemaHelper#withDorisTypeMetadata}).
*/
private static Field field(String name, Type type, boolean nullable, String timezone)
throws UndiscoverableSchemaException {
return FlightSqlSchemaHelper.withDorisTypeMetadata(arrowField(name, type, nullable, timezone), type);
}
private static Field arrowField(String name, Type type, boolean nullable, String timezone)
throws UndiscoverableSchemaException {
// group_concat uses IAggregateFunction's string serialization, unlike fixed-size states
// such as sum/count. Match that BE wire type instead of treating every AGG_STATE as Null.
if (type instanceof AggStateType && "group_concat".equals(((AggStateType) type).getFunctionName())) {
type = Type.STRING;
}
PrimitiveType primitive = type.getPrimitiveType();
if (primitive == PrimitiveType.VARIANT) {
return FlightSqlSchemaHelper.nativeVariantField(name, nullable, Collections.emptyMap());
}
int precision = type instanceof ScalarType ? ((ScalarType) type).getScalarPrecision() : 0;
int scale = type instanceof ScalarType ? ((ScalarType) type).getScalarScale() : 0;
ArrowType arrowType = DorisArrowTypeMapping.toArrowType(primitive, precision, scale);
if (primitive == PrimitiveType.TIMESTAMPTZ) {
// BE stamps the session time zone, where the shared mapping still says "UTC" (#67577).
arrowType = new ArrowType.Timestamp(((ArrowType.Timestamp) arrowType).getUnit(), timezone);
}
// NULL_TYPE too: BE builds an untyped NULL column as a UInt8 data type flagged as a NULL literal
// (data_type_factory.cpp), so convert_to_arrow_type sends it as Bool, never as Null.
if (arrowType instanceof ArrowType.Null) {
throw new UndiscoverableSchemaException("Unsupported Arrow result type: " + type);
}
List<Field> children = new ArrayList<>();
if (type instanceof ArrayType) {
// BE constructs ListType and MapType from data types, so item/value fields are nullable.
children.add(arrowField("item", ((ArrayType) type).getItemType(), true, timezone));
} else if (type instanceof MapType) {
MapType map = (MapType) type;
children.add(new Field("entries", FieldType.notNullable(new ArrowType.Struct()), Arrays.asList(
arrowField("key", map.getKeyType(), false, timezone),
arrowField("value", map.getValueType(), true, timezone))));
} else if (type instanceof StructType) {
for (StructField child : ((StructType) type).getFields()) {
children.add(arrowField(child.getName(), child.getType(), child.getContainsNull(), timezone));
}
}
return new Field(name, new FieldType(nullable, arrowType, null), children);
}
}
@@ -41,6 +41,7 @@ import com.google.common.base.Strings;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Lists;
import org.apache.arrow.flight.CallStatus;
import org.apache.arrow.vector.types.pojo.Schema;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
@@ -51,6 +52,7 @@ import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.ReentrantLock;
@@ -91,7 +93,8 @@ public class FlightProtocolAdapter implements ProtocolAdapter {
private final String peerIdentity;
private final FlightSqlChannel channel = new FlightSqlChannel();
private final Map<String, String> preparedQuerys = new HashMap<>();
// Only used by the commands of the session, under the command lock; teardown does not touch it.
private final Map<String, PreparedQuery> preparedQuerys = new HashMap<>();
private String runningQuery;
private final List<FlightSqlEndpointsLocation> endpointsLocations = Lists.newArrayList();
// How many of endpointsLocations were registered before the statement being executed
@@ -372,11 +375,11 @@ public class FlightProtocolAdapter implements ProtocolAdapter {
return channel;
}
public void addPreparedQuery(String preparedStatementId, String preparedQuery) {
public void addPreparedQuery(String preparedStatementId, PreparedQuery preparedQuery) {
preparedQuerys.put(preparedStatementId, preparedQuery);
}
public String getPreparedQuery(String preparedStatementId) {
public PreparedQuery getPreparedQuery(String preparedStatementId) {
return preparedQuerys.get(preparedStatementId);
}
@@ -557,6 +560,45 @@ public class FlightProtocolAdapter implements ProtocolAdapter {
}
/** The body of a command run by {@link #runCommand}. */
/**
* A prepared statement of the session: its SQL, the namespace it was prepared in, and the result
* schema Prepare advertised for it -- null when that schema could not be told without running the
* statement, whose executions are then not checked against one.
*/
public static final class PreparedQuery {
private final String sql;
private final Schema schema;
private String catalog;
private String database;
public PreparedQuery(String sql, Schema schema, String catalog, String database) {
this.sql = sql;
this.schema = schema;
this.catalog = catalog;
this.database = database;
}
public String getSql() {
return sql;
}
public Schema getSchema() {
return schema;
}
// Unqualified names in the SQL resolve against the session's namespace, so the schema holds
// only where it was told.
public boolean isPreparedIn(String catalog, String database) {
return Objects.equals(this.catalog, catalog) && Objects.equals(this.database, database);
}
// A USE or SWITCH of the statement itself moves its namespace with it.
public void rebind(String catalog, String database) {
this.catalog = catalog;
this.database = database;
}
}
@FunctionalInterface
public interface SessionAction<E extends Exception> {
void run() throws E;
@@ -6692,6 +6692,22 @@ public class Env {
this.analysisManager.createAnalyze(command, isProxy);
}
/**
* The database a session lands in when it switches to {@code catalog}, if the catalog names one: a data
* source whose metadata model has no database layer can name the single database Doris presents for it;
* the connector owns that name. Asked of the provider, not the connector, because switching to a catalog
* must not force it to initialize.
*/
public static Optional<String> defaultDatabaseOnUse(CatalogIf catalog) {
Map<String, String> catalogProps = catalog.getProperties();
String catalogType = catalogProps.get(CatalogMgr.CATALOG_TYPE_PROP);
if (catalogType == null) {
return Optional.empty();
}
return ConnectorFactory.findProvider(catalogType, catalogProps)
.flatMap(ConnectorProvider::defaultDatabaseOnUse);
}
// Switch catalog of this session
public void changeCatalog(ConnectContext ctx, String catalogName) throws DdlException {
CatalogIf catalogIf = catalogMgr.getCatalogOrDdlException(catalogName);
@@ -6709,17 +6725,9 @@ public class Env {
if (StringUtils.isNotEmpty(lastDb)) {
ctx.setDatabase(lastDb);
}
// A data source whose metadata model has no database layer can name the single database Doris
// presents for it; the connector owns that name. Asked of the provider, not the connector, because
// switching to a catalog must not force it to initialize. Applied after the remembered database on
// purpose: such a catalog has exactly one database, so there is nothing else to return to.
Map<String, String> catalogProps = catalogIf.getProperties();
String catalogType = catalogProps.get(CatalogMgr.CATALOG_TYPE_PROP);
if (catalogType != null) {
ConnectorFactory.findProvider(catalogType, catalogProps)
.flatMap(ConnectorProvider::defaultDatabaseOnUse)
.ifPresent(ctx::setDatabase);
}
// Applied after the remembered database on purpose: such a catalog has exactly one database, so
// there is nothing else to return to.
defaultDatabaseOnUse(catalogIf).ifPresent(ctx::setDatabase);
} catch (RuntimeException | Error e) {
// A failed remembered/default database lookup must not leave a partially switched session.
ctx.restoreNameSpaceContext(previousNamespace);
@@ -93,7 +93,9 @@ public class FeServiceClient {
this.user = user;
this.password = password;
this.retryCount = retryCount;
this.timeoutMs = (int) TimeUnit.SECONDS.toMillis(timeoutS);
// Catalog metadata timeouts are in seconds, the pooled Thrift sockets take int milliseconds: reject a
// timeout that does not fit instead of letting it wrap around.
this.timeoutMs = Math.toIntExact(TimeUnit.SECONDS.toMillis(timeoutS));
}
private List<TNetworkAddress> getAddresses() {
@@ -125,6 +125,29 @@ public class HelpCommand extends ShowCommand {
return resultSet;
}
/** Resolve the HELP result shape without materializing topic text or category rows. */
public ShowResultSetMetaData getMetaData(ConnectContext ctx) throws AnalysisException {
if (Strings.isNullOrEmpty(mark)) {
throw new AnalysisException("Help empty info.");
}
HelpModule module = HelpModule.getInstance();
// HELP has three result shapes; a topic match takes precedence over keywords and categories.
HelpTopic topic = module.getTopic(mark);
if (topic == null) {
List<String> topics = module.listTopicByKeyword(mark);
if (topics.size() > 1) {
return KEYWORD_META_DATA;
}
if (topics.size() == 1) {
topic = module.getTopic(topics.get(0));
}
}
if (topic != null) {
return TOPIC_META_DATA;
}
return module.listCategoryByName(mark).size() > 1 ? CATEGORY_META_DATA : KEYWORD_META_DATA;
}
@Override
public ShowResultSetMetaData getMetaData() {
return ShowResultSetMetaData.builder().build();
@@ -81,6 +81,12 @@ public class ShowConfigCommand extends Command implements NoForward {
this.isShowSingleBackend = true;
}
@Override
public ShowResultSetMetaData getResultSetMetaData() {
// Config commands produce rows even though their statement category is OTHER.
return getMetaData(nodeType == NodeType.FRONTEND ? FE_TITLE_NAMES : BE_TITLE_NAMES);
}
private ShowResultSetMetaData getMetaData(ImmutableList<String> metaNames) {
ShowResultSetMetaData.Builder builder = ShowResultSetMetaData.builder();
for (String title : metaNames) {
@@ -138,6 +138,26 @@ public class ShowCreateTableCommand extends ShowCommand {
return ShowResultSetMetaData.builder().build();
}
/** Resolve the result header without rendering the table DDL. */
public ShowResultSetMetaData getMetaData(ConnectContext ctx) throws AnalysisException {
// The table kind determines the header; rendering DDL is unnecessary for schema discovery.
validate(ctx);
DatabaseIf db = ctx.getEnv().getCatalogMgr().getCatalogOrAnalysisException(tblNameInfo.getCtl())
.getDbOrAnalysisException(tblNameInfo.getDb());
return metaDataOf(redirectSysTableToSource(resolveShowCreateTarget(db)));
}
// The header doRun returns for a table whose DDL it renders. A plugin-driven table, view or not,
// gets META_DATA here as from its own branches in doRun.
private static ShowResultSetMetaData metaDataOf(TableIf table) {
if (table instanceof View) {
return VIEW_META_DATA;
} else if (table instanceof BaseTableStream) {
return STREAM_META_DATA;
}
return table.getType() != Table.TableType.MATERIALIZED_VIEW ? META_DATA : MATERIALIZED_VIEW_META_DATA;
}
@Override
public ShowResultSet doRun(ConnectContext ctx, StmtExecutor executor) throws Exception {
validate(ctx);
@@ -182,16 +202,10 @@ public class ShowCreateTableCommand extends ShowCommand {
if (table instanceof View) {
rows.add(Lists.newArrayList(table.getName(), createTableStmt.get(0), "utf8mb4", "utf8mb4_0900_bin"));
return new ShowResultSet(VIEW_META_DATA, rows);
} else if (table instanceof BaseTableStream) {
rows.add(Lists.newArrayList(Util.getTempTableDisplayName(table.getName()), createTableStmt.get(0)));
return new ShowResultSet(STREAM_META_DATA, rows);
} else {
rows.add(Lists.newArrayList(Util.getTempTableDisplayName(table.getName()), createTableStmt.get(0)));
return (table.getType() != Table.TableType.MATERIALIZED_VIEW
? new ShowResultSet(META_DATA, rows)
: new ShowResultSet(MATERIALIZED_VIEW_META_DATA, rows));
}
return new ShowResultSet(metaDataOf(table), rows);
} finally {
table.readUnlock();
}
@@ -24,6 +24,7 @@ import org.apache.doris.catalog.ScalarType;
import org.apache.doris.common.AnalysisException;
import org.apache.doris.common.ErrorCode;
import org.apache.doris.common.ErrorReport;
import org.apache.doris.common.proc.CurrentQueryStatisticsProcDir;
import org.apache.doris.common.proc.ProcNodeInterface;
import org.apache.doris.common.proc.ProcResult;
import org.apache.doris.common.proc.ProcService;
@@ -35,6 +36,7 @@ import org.apache.doris.qe.ShowResultSet;
import org.apache.doris.qe.ShowResultSetMetaData;
import org.apache.doris.qe.StmtExecutor;
import java.util.Collections;
import java.util.List;
/**
@@ -74,6 +76,28 @@ public class ShowProcCommand extends ShowCommand {
return ShowResultSetMetaData.builder().build();
}
/** Resolve a static PROC header, or return null when discovery would require fetching rows. */
public ShowResultSetMetaData getMetaData(ConnectContext ctx) throws AnalysisException {
// PROC headers belong to the resolved node and require the same privilege as reading it.
if (!Env.getCurrentEnv().getAccessManager().checkGlobalPriv(ctx, PrivPredicate.ADMIN_OR_NODE)) {
ErrorReport.reportAnalysisException(ErrorCode.ERR_SPECIFIC_ACCESS_DENIED_ERROR, "ADMIN");
}
// Even opening nested nodes can perform remote lookups. Only use known static headers.
List<String> names;
if ("/".equals(path)) {
names = Collections.singletonList("name");
} else if ("/current_queries".equals(path)) {
names = CurrentQueryStatisticsProcDir.TITLE_NAMES;
} else {
return null;
}
ShowResultSetMetaData.Builder builder = ShowResultSetMetaData.builder();
for (String name : names) {
builder.addColumn(new Column(name, ScalarType.createVarchar(30)));
}
return builder.build();
}
private ShowResultSet handleShowProc(ConnectContext ctx, StmtExecutor executor) throws Exception {
if (!Env.getCurrentEnv().getAccessManager()
.checkGlobalPriv(ctx, PrivPredicate.ADMIN_OR_NODE)) {
@@ -109,6 +109,16 @@ public class ShowSnapshotCommand extends ShowCommand {
return builder.build();
}
/** Resolve the snapshot header without accessing the remote repository. */
public ShowResultSetMetaData getMetaData(ConnectContext ctx) throws UserException {
// Both filters must be analyzed before choosing the detail header, even during Prepare.
if (!validate(ctx)) {
throw new AnalysisException("Where clause should looks like: SNAPSHOT = 'your_snapshot_name'"
+ " [AND TIMESTAMP = '2018-04-18-19-19-10'] [AND SNAPSHOTTYPE = 'remote' | 'local']");
}
return getMetaData();
}
private boolean analyzeSubExpression(Expression expr) {
if (expr instanceof ComparisonPredicate) {
return analyzeComparisonPredicate((ComparisonPredicate) expr);
@@ -194,18 +204,14 @@ public class ShowSnapshotCommand extends ShowCommand {
* handle show backup
*/
private ShowResultSet handleShowSnapshot(ConnectContext ctx, StmtExecutor executor) throws Exception {
boolean valid = validate(ctx);
if (!valid) {
throw new AnalysisException("Where clause should looks like: SNAPSHOT = 'your_snapshot_name'"
+ " [AND TIMESTAMP = '2018-04-18-19-19-10'] [AND SNAPSHOTTYPE = 'remote' | 'local']");
}
ShowResultSetMetaData metadata = getMetaData(ctx);
Repository repo = Env.getCurrentEnv().getBackupHandler().getRepoMgr().getRepo(repoName);
if (repo == null) {
throw new AnalysisException("Repository " + repoName + " does not exist");
}
List<List<String>> snapshotInfos = repo.getSnapshotInfos(snapshotName, timestamp);
return new ShowResultSet(getMetaData(), snapshotInfos);
return new ShowResultSet(metadata, snapshotInfos);
}
@Override
@@ -25,6 +25,7 @@ import org.apache.doris.common.ErrorReport;
import org.apache.doris.mysql.privilege.PrivPredicate;
import org.apache.doris.nereids.trees.plans.PlanType;
import org.apache.doris.nereids.trees.plans.visitor.PlanVisitor;
import org.apache.doris.policy.StoragePolicy;
import org.apache.doris.qe.ConnectContext;
import org.apache.doris.qe.ShowResultSet;
import org.apache.doris.qe.ShowResultSetMetaData;
@@ -55,7 +56,8 @@ public class ShowStoragePolicyCommand extends ShowCommand {
@Override
public ShowResultSetMetaData getMetaData() {
return USING_META_DATA;
// The header doRun answers with: the objects using the policies, or the policies themselves.
return isUsing ? USING_META_DATA : StoragePolicy.STORAGE_META_DATA;
}
@Override
@@ -54,6 +54,14 @@ public class UseCommand extends Command implements NoForward {
this.databaseName = databaseName;
}
public String getCatalogName() {
return catalogName;
}
public String getDatabaseName() {
return databaseName;
}
@Override
public void run(ConnectContext ctx, StmtExecutor executor) throws Exception {
validate(ctx);
@@ -31,6 +31,7 @@ import org.apache.doris.analysis.StringLiteral;
import org.apache.doris.analysis.TimeStampNsLiteral;
import org.apache.doris.analysis.UserIdentity;
import org.apache.doris.arrowflight.protocol.FlightProtocolAdapter;
import org.apache.doris.arrowflight.protocol.FlightProtocolAdapter.PreparedQuery;
import org.apache.doris.arrowflight.results.FlightSqlChannel;
import org.apache.doris.arrowflight.results.FlightSqlEndpointsLocation;
import org.apache.doris.authentication.Principal;
@@ -874,11 +875,11 @@ public class ConnectContext {
}
// Arrow Flight SQL only.
public void addPreparedQuery(String preparedStatementId, String preparedQuery) {
public void addPreparedQuery(String preparedStatementId, PreparedQuery preparedQuery) {
FlightProtocolAdapter.of(this).addPreparedQuery(preparedStatementId, preparedQuery);
}
public String getPreparedQuery(String preparedStatementId) {
public PreparedQuery getPreparedQuery(String preparedStatementId) {
return FlightProtocolAdapter.of(this).getPreparedQuery(preparedStatementId);
}
@@ -151,10 +151,10 @@ public class DorisArrowTypeMappingTest {
row(PrimitiveType.HLL, new ArrowType.Binary()),
row(PrimitiveType.BITMAP, new ArrowType.Binary()),
row(PrimitiveType.QUANTILE_STATE, new ArrowType.Binary()),
// BE emits float64 for TIMEV2 and binary for VARBINARY and AGG_STATE; the schema says
// Null for all three. Kept as is (#67577).
row(PrimitiveType.TIMEV2, 18, 0, new ArrowType.Null()),
row(PrimitiveType.VARBINARY, new ArrowType.Null()),
// As BE emits them: float64 for TIMEV2 and binary for VARBINARY.
row(PrimitiveType.TIMEV2, 18, 0, new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)),
row(PrimitiveType.VARBINARY, new ArrowType.Binary()),
// BE emits binary for AGG_STATE; the schema says Null. Kept as is (#67577).
row(PrimitiveType.AGG_STATE, new ArrowType.Null()),
// The element types of a complex column live in the field's children, not in its type.
row(PrimitiveType.ARRAY, new ArrowType.List()),
@@ -34,6 +34,8 @@ import org.apache.doris.qe.StmtExecutor;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.protobuf.Any;
import com.google.protobuf.ByteString;
import org.apache.arrow.flight.ActionType;
import org.apache.arrow.flight.CallStatus;
import org.apache.arrow.flight.CloseSessionRequest;
@@ -55,7 +57,11 @@ import org.apache.arrow.flight.SetSessionOptionsResult;
import org.apache.arrow.flight.SetSessionOptionsResult.ErrorValue;
import org.apache.arrow.flight.sql.FlightSqlUtils;
import org.apache.arrow.flight.sql.impl.FlightSql.ActionCreatePreparedStatementRequest;
import org.apache.arrow.flight.sql.impl.FlightSql.ActionCreatePreparedStatementResult;
import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementQuery;
import org.apache.arrow.vector.ipc.ReadChannel;
import org.apache.arrow.vector.ipc.message.MessageSerializer;
import org.apache.arrow.vector.types.pojo.Schema;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.BeforeEach;
@@ -63,6 +69,8 @@ import org.junit.jupiter.api.Test;
import org.mockito.MockedConstruction;
import org.mockito.Mockito;
import java.io.ByteArrayInputStream;
import java.nio.channels.Channels;
import java.util.ArrayList;
import java.util.List;
import java.util.Set;
@@ -173,15 +181,17 @@ public class DorisFlightSqlProducerTest {
* Regression test for the FE direct-memory leak in {@code createPreparedStatement}
* (issue apache/doris#65305, fixed by PR #65311).
*
* <p>Before the fix, each prepare allocated a {@link org.apache.arrow.vector.VectorSchemaRoot}
* <p>Before that fix, each prepare allocated a {@link org.apache.arrow.vector.VectorSchemaRoot}
* via {@code FlightSqlChannel.createOneOneSchemaRoot("ResultMeta", ...)} and read only its
* {@code Schema}, never closing the root, leaking one off-heap {@code VarCharVector} buffer per
* prepare. Because Arrow Flight has no parameter binding, every client query triggers a fresh
* prepare, so the leak grows monotonically until {@code MaxDirectMemorySize} is exhausted.
*
* <p>This test drives {@code createPreparedStatement} against a real {@link FlightSqlChannel} many
* times and asserts the channel's Arrow allocator holds zero bytes afterwards. It goes red if the
* try-with-resources that closes the temporary roots is removed.
* <p>Prepare now advertises the analyzed result schema; that {@code ResultMeta} placeholder is left
* for a statement whose schema only running it can tell, such as a request that sets a variable
* before its query. This test drives {@code createPreparedStatement} for such statements against a
* real {@link FlightSqlChannel} many times and asserts that each one gets the placeholder and the
* channel's Arrow allocator holds zero bytes afterwards.
*/
@Test
public void createPreparedStatementDoesNotLeakChannelAllocator() throws Exception {
@@ -224,7 +234,17 @@ public class DorisFlightSqlProducerTest {
StreamListener<Result> listener = new StreamListener<Result>() {
@Override
public void onNext(Result val) {
// discard the placeholder prepared-statement result
try {
ByteString dataset = Any.parseFrom(val.getBody())
.unpack(ActionCreatePreparedStatementResult.class).getDatasetSchema();
Schema schema = MessageSerializer.deserializeSchema(new ReadChannel(
Channels.newChannel(new ByteArrayInputStream(dataset.toByteArray()))));
if (!"ResultMeta".equals(schema.getFields().get(0).getName())) {
errors.incrementAndGet();
}
} catch (Exception e) {
errors.incrementAndGet();
}
}
@Override
@@ -239,14 +259,15 @@ public class DorisFlightSqlProducerTest {
}
};
ActionCreatePreparedStatementRequest request = ActionCreatePreparedStatementRequest.newBuilder()
.setQuery("select * from t where id = " + i).build();
.setQuery("set @id = " + i + "; select * from t where id = @id").build();
producer.createPreparedStatement(request, callContext, listener);
Assertions.assertTrue(finished.await(30, TimeUnit.SECONDS), "createPreparedStatement #" + i + " did not finish in time");
}
// Guard against a false pass: if a prepare failed before reaching the allocation, no buffer
// Guard against a false pass: if a prepare failed before reaching the placeholder, no buffer
// would be leaked and the memory assertion below could not detect a regression.
Assertions.assertEquals(0, errors.get(), "no createPreparedStatement call should fail");
Assertions.assertEquals(0, errors.get(),
"every createPreparedStatement call should answer with the placeholder schema");
// Every temporary VectorSchemaRoot must have been closed, so the channel's Arrow allocator
// is back to zero. Reverting the fix leaves `rounds` ResultMeta buffers allocated here.
Assertions.assertEquals(0L, channel.getAllocatedMemory(), "createPreparedStatement leaked off-heap memory in the channel allocator");
@@ -0,0 +1,183 @@
// 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.Expr;
import org.apache.doris.analysis.ExprToThriftVisitor;
import org.apache.doris.arrow.DorisArrowTypeMapping;
import org.apache.doris.arrowflight.protocol.FlightProtocolAdapter;
import org.apache.doris.catalog.Env;
import org.apache.doris.catalog.ScalarType;
import org.apache.doris.common.FeConstants;
import org.apache.doris.nereids.glue.LogicalPlanAdapter;
import org.apache.doris.qe.ConnectContext;
import org.apache.doris.qe.QueryState.MysqlStateType;
import org.apache.doris.qe.VariableMgr;
import org.apache.doris.thrift.TExprNode;
import org.apache.doris.thrift.TExprNodeType;
import org.apache.doris.utframe.TestWithFeService;
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 org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
class FlightSqlExecutionSchemaTest extends TestWithFeService {
private ConnectContext flightContext;
private boolean previousUnitTest;
@Override
protected void runBeforeAll() throws Exception {
previousUnitTest = FeConstants.runningUnitTest;
FeConstants.runningUnitTest = true;
createDatabase("execution_schema_test");
connectContext.setDatabase("execution_schema_test");
createTable("CREATE TABLE schema_input (id INT NOT NULL, encoded STRING NULL) "
+ "DISTRIBUTED BY HASH(id) BUCKETS 1 PROPERTIES ('replication_num' = '1')");
flightContext = ConnectContext.forFlight("schema-peer");
flightContext.setCurrentUserIdentity(connectContext.getCurrentUserIdentity());
flightContext.setRemoteIP(connectContext.getRemoteIP());
flightContext.setEnv(Env.getCurrentEnv());
flightContext.setSessionVariable(VariableMgr.cloneSessionVariable(connectContext.getSessionVariable()));
flightContext.setDatabase("execution_schema_test");
}
@Override
protected void runAfterAll() throws Exception {
FeConstants.runningUnitTest = previousUnitTest;
flightContext.getFlightSqlChannel().close();
}
@Test
void unaliasedLiteralRetainsExecutionLabel() throws Exception {
assertExecutionSchema("SELECT 1", false, false);
assertExecutionSchema("SELECT 1 AS id", false, false);
}
@Test
void unaliasedFunctionRetainsExecutionLabel() throws Exception {
assertExecutionSchema("SELECT unhex(encoded) FROM schema_input WHERE id <= 4 ORDER BY id", false, false);
}
@Test
void constantFoldingNarrowsNullableCast() throws Exception {
assertExecutionSchema("SELECT CAST(CAST('2026-01-01' AS DATETIME(6)) AS STRING)", false, true);
}
@Test
void nullableColumnRemainsCompatibleWithFilter() throws Exception {
assertExecutionSchema("SELECT encoded FROM schema_input WHERE encoded IS NOT NULL", false, false);
}
@Test
void tableValuedFunctionUsesQualifiedWireLabel() throws Exception {
Schema actual = assertExecutionSchema("SELECT * FROM numbers('number' = '1')", true, false);
Assertions.assertTrue(actual.getFields().get(0).getName().startsWith("_tvf_numbers."));
}
private Schema assertExecutionSchema(String sql, boolean differentLabel, boolean narrowedNullability)
throws Exception {
// What a command of the session does first: the session's context is the thread's.
flightContext.setThreadLocalInfo();
try (FlightSqlConnectProcessor processor = new FlightSqlConnectProcessor(flightContext)) {
Schema prepared = FlightSqlQuerySchema.analyze(flightContext, sql);
// Execution marks the query state before binding output labels. A bare planner call
// misses that transition and can reproduce Prepare's incorrect command-style aliases.
processor.handleQuery(sql);
Assertions.assertNotEquals(MysqlStateType.ERR, flightContext.getState().getStateType(),
flightContext.getState().getErrorMessage());
Assertions.assertFalse(flightContext.isReturnResultFromLocal());
Assertions.assertNotNull(flightContext.getExecutor());
LogicalPlanAdapter statement = (LogicalPlanAdapter) flightContext.getExecutor().getParsedStmt();
List<Field> wireFields = new ArrayList<>();
for (Expr expression : statement.getResultExprs()) {
TExprNode node = ExprToThriftVisitor.treeToThrift(expression).getNodes().get(0);
Assertions.assertEquals(TExprNodeType.SLOT_REF, node.getNodeType());
ScalarType type = (ScalarType) expression.getType();
ArrowType arrowType = DorisArrowTypeMapping.toArrowType(type.getPrimitiveType(),
type.getScalarPrecision(), type.getScalarScale());
// BE derives these fields from the serialized output expressions, not analyzed Slot names.
wireFields.add(new Field(node.getLabel(),
new FieldType(node.isIsNullable(), arrowType, null), Collections.emptyList()));
}
Schema actual = new Schema(wireFields);
Assertions.assertEquals(differentLabel,
!prepared.getFields().get(0).getName().equals(actual.getFields().get(0).getName()));
Assertions.assertEquals(narrowedNullability,
prepared.getFields().get(0).isNullable() && !actual.getFields().get(0).isNullable());
Assertions.assertTrue(FlightSqlQuerySchema.matchesExecutionSchema(
prepared, actual, statement.getColLabels()));
return actual;
} finally {
FlightProtocolAdapter.of(flightContext).beginRequest();
connectContext.setThreadLocalInfo();
}
}
@Test
void sameArityRenameIsRejectedDespiteIdenticalWireTypes() {
Schema schema = schema(Field.nullable("old_name", new ArrowType.Utf8()));
Assertions.assertFalse(FlightSqlQuerySchema.matchesExecutionSchema(schema, schema,
Collections.singletonList("new_name")));
}
@Test
void differentTypesAreRejected() {
Schema prepared = schema(Field.nullable("value", new ArrowType.Int(32, true)));
Schema actual = schema(Field.nullable("value", new ArrowType.Int(64, true)));
Assertions.assertFalse(FlightSqlQuerySchema.matchesExecutionSchema(prepared, actual,
Collections.singletonList("value")));
}
@Test
void nestedFieldNamesAreNotTransportLabels() {
Schema prepared = schema(new Field("value", FieldType.nullable(new ArrowType.Struct()),
Collections.singletonList(Field.nullable("old_name", new ArrowType.Utf8()))));
Schema actual = schema(new Field("value", FieldType.nullable(new ArrowType.Struct()),
Collections.singletonList(Field.nullable("new_name", new ArrowType.Utf8()))));
Assertions.assertFalse(FlightSqlQuerySchema.matchesExecutionSchema(prepared, actual,
Collections.singletonList("value")));
}
@Test
void typeMetadataChangesAreRejected() {
Schema prepared = schema(new Field("value", new FieldType(true, new ArrowType.Utf8(), null,
Collections.singletonMap("doris_type", "LARGEINT")), Collections.emptyList()));
Schema actual = schema(Field.nullable("value", new ArrowType.Utf8()));
Assertions.assertFalse(FlightSqlQuerySchema.matchesExecutionSchema(prepared, actual,
Collections.singletonList("value")));
}
@Test
void nullableWideningIsRejected() {
Schema prepared = schema(Field.notNullable("value", new ArrowType.Utf8()));
Schema actual = schema(Field.nullable("value", new ArrowType.Utf8()));
Assertions.assertFalse(FlightSqlQuerySchema.matchesExecutionSchema(prepared, actual,
Collections.singletonList("value")));
}
private static Schema schema(Field field) {
return new Schema(Collections.singletonList(field));
}
}
@@ -0,0 +1,83 @@
// 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.datasource.doris;
import org.apache.doris.common.ClientPool;
import org.apache.doris.thrift.FrontendService;
import org.apache.doris.thrift.TGetBackendMetaResult;
import org.apache.doris.thrift.TNetworkAddress;
import org.apache.doris.thrift.TStatus;
import org.apache.doris.thrift.TStatusCode;
import org.apache.thrift.protocol.TBinaryProtocol;
import org.apache.thrift.transport.TSocket;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;
import org.mockito.Mockito;
import java.net.InetAddress;
import java.net.ServerSocket;
import java.util.Collections;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
class FeServiceClientTest {
@Test
void rejectsTimeoutThatCannotFitInSocketMilliseconds() {
Assertions.assertThrows(ArithmeticException.class, () -> new FeServiceClient("remote_catalog",
Collections.emptyList(), "test_user", "", 1, Integer.MAX_VALUE));
}
@Test
void metadataReadTimeoutUsesSeconds() throws Exception {
FrontendService.Iface handler = Mockito.mock(FrontendService.Iface.class);
Mockito.when(handler.getBackendMeta(Mockito.any())).thenAnswer(invocation -> {
// Ordinary metadata latency must remain within the configured seconds-based timeout.
Thread.sleep(100);
return new TGetBackendMetaResult().setStatus(new TStatus(TStatusCode.OK))
.setBackends(Collections.emptyList());
});
ExecutorService executor = Executors.newSingleThreadExecutor();
try (ServerSocket server = new ServerSocket(0, 1, InetAddress.getLoopbackAddress())) {
server.setSoTimeout(5000);
Future<?> response = executor.submit(() -> {
try (TSocket socket = new TSocket(server.accept())) {
TBinaryProtocol protocol = new TBinaryProtocol(socket);
new FrontendService.Processor<>(handler).process(protocol, protocol);
} catch (Exception e) {
throw new RuntimeException(e);
}
});
TNetworkAddress address = new TNetworkAddress(server.getInetAddress().getHostAddress(),
server.getLocalPort());
try {
FeServiceClient client = new FeServiceClient("remote_catalog", Collections.singletonList(address),
"test_user", "", 1, 10);
Assertions.assertTrue(client.listBackends().isEmpty());
response.get(5, TimeUnit.SECONDS);
} finally {
ClientPool.frontendPool.clearPool(address);
}
} finally {
executor.shutdownNow();
Assertions.assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS));
}
}
}
@@ -0,0 +1,147 @@
// 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.nereids.trees.plans.commands;
import org.apache.doris.catalog.Env;
import org.apache.doris.common.AnalysisException;
import org.apache.doris.mysql.privilege.PrivPredicate;
import org.apache.doris.nereids.parser.NereidsParser;
import org.apache.doris.policy.PolicyMgr;
import org.apache.doris.policy.StoragePolicy;
import org.apache.doris.qe.ConnectContext;
import org.apache.doris.qe.ResultSetMetaData;
import org.apache.doris.qe.help.HelpModule;
import org.apache.doris.qe.help.HelpTopic;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;
import org.mockito.MockedStatic;
import org.mockito.Mockito;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;
class SchemaOnlyCommandMetadataTest {
private static List<String> names(ResultSetMetaData metadata) {
return metadata.getColumns().stream().map(column -> column.getName()).collect(Collectors.toList());
}
@Test
void helpTopicDoesNotReadDescriptionOrExample() throws Exception {
HelpModule module = Mockito.mock(HelpModule.class);
HelpTopic topic = Mockito.mock(HelpTopic.class);
Mockito.when(module.getTopic("topic")).thenReturn(topic);
try (MockedStatic<HelpModule> ignored = Mockito.mockStatic(HelpModule.class)) {
ignored.when(HelpModule::getInstance).thenReturn(module);
HelpCommand command = new HelpCommand("topic");
Assertions.assertEquals(Arrays.asList("name", "description", "example"), names(command.getMetaData(null)));
Mockito.verifyNoInteractions(topic);
Assertions.assertEquals(names(command.doRun(null, null).getMetaData()), names(command.getMetaData(null)));
}
}
@Test
void helpKeywordWithOneTopicUsesTopicHeader() throws Exception {
HelpModule module = Mockito.mock(HelpModule.class);
Mockito.when(module.listTopicByKeyword("keyword")).thenReturn(Collections.singletonList("topic"));
Mockito.when(module.getTopic("topic")).thenReturn(Mockito.mock(HelpTopic.class));
try (MockedStatic<HelpModule> ignored = Mockito.mockStatic(HelpModule.class)) {
ignored.when(HelpModule::getInstance).thenReturn(module);
HelpCommand command = new HelpCommand("keyword");
Assertions.assertEquals(Arrays.asList("name", "description", "example"), names(command.getMetaData(null)));
Assertions.assertEquals(names(command.doRun(null, null).getMetaData()), names(command.getMetaData(null)));
}
}
@Test
void helpKeywordAndCategoryHeadersMatchExecution() throws Exception {
HelpModule module = Mockito.mock(HelpModule.class);
Mockito.when(module.listTopicByKeyword("keyword")).thenReturn(Arrays.asList("first", "second"));
Mockito.when(module.listCategoryByName("categories")).thenReturn(Arrays.asList("first", "second"));
Mockito.when(module.listCategoryByName("category")).thenReturn(Collections.singletonList("first"));
try (MockedStatic<HelpModule> ignored = Mockito.mockStatic(HelpModule.class)) {
ignored.when(HelpModule::getInstance).thenReturn(module);
for (String mark : Arrays.asList("keyword", "categories", "category", "missing")) {
HelpCommand command = new HelpCommand(mark);
List<String> expected = mark.equals("categories")
? Arrays.asList("source_category_name", "name", "is_it_category")
: Arrays.asList("name", "is_it_category");
Assertions.assertEquals(expected, names(command.getMetaData(null)));
Assertions.assertEquals(names(command.doRun(null, null).getMetaData()),
names(command.getMetaData(null)));
}
}
}
@Test
void invalidHelpAndSnapshotFiltersAreRejected() throws Exception {
Assertions.assertThrows(AnalysisException.class, () -> new HelpCommand("").getMetaData(null));
Env env = Mockito.mock(Env.class, Mockito.RETURNS_DEEP_STUBS);
Mockito.when(env.getAccessManager().checkGlobalPriv(Mockito.nullable(ConnectContext.class),
Mockito.eq(PrivPredicate.ADMIN))).thenReturn(true);
try (MockedStatic<Env> ignored = Mockito.mockStatic(Env.class)) {
ignored.when(Env::getCurrentEnv).thenReturn(env);
ShowSnapshotCommand command = (ShowSnapshotCommand) new NereidsParser()
.parseSingle("SHOW SNAPSHOT ON example_repo WHERE TIMESTAMP = '2026-01-01-00-00-00'");
Assertions.assertThrows(AnalysisException.class, () -> command.getMetaData(null));
Mockito.verify(env, Mockito.never()).getBackupHandler();
}
}
@Test
void filteredSnapshotsResolveWithoutRepositoryAccess() throws Exception {
Env env = Mockito.mock(Env.class, Mockito.RETURNS_DEEP_STUBS);
Mockito.when(env.getAccessManager().checkGlobalPriv(Mockito.nullable(ConnectContext.class),
Mockito.eq(PrivPredicate.ADMIN))).thenReturn(true);
try (MockedStatic<Env> ignored = Mockito.mockStatic(Env.class)) {
ignored.when(Env::getCurrentEnv).thenReturn(env);
for (String filter : Arrays.asList("", " WHERE SNAPSHOT = 'sample'",
" WHERE SNAPSHOT = 'sample' AND TIMESTAMP = '2026-01-01-00-00-00'")) {
ShowSnapshotCommand command = (ShowSnapshotCommand) new NereidsParser()
.parseSingle("SHOW SNAPSHOT ON example_repo" + filter);
Assertions.assertEquals(filter.contains("TIMESTAMP")
? ShowSnapshotCommand.SNAPSHOT_DETAIL : ShowSnapshotCommand.SNAPSHOT_ALL,
names(command.getMetaData(null)));
}
Mockito.verify(env, Mockito.never()).getBackupHandler();
}
}
@Test
void storagePolicyHeadersMatchExecution() throws Exception {
Env env = Mockito.mock(Env.class, Mockito.RETURNS_DEEP_STUBS);
Mockito.when(env.getAccessManager().checkGlobalPriv(Mockito.nullable(ConnectContext.class),
Mockito.eq(PrivPredicate.ADMIN))).thenReturn(true);
Mockito.when(env.getPolicyMgr()).thenReturn(new PolicyMgr());
Mockito.when(env.getInternalCatalog().getDbs()).thenReturn(Collections.emptyList());
try (MockedStatic<Env> ignored = Mockito.mockStatic(Env.class)) {
ignored.when(Env::getCurrentEnv).thenReturn(env);
for (String sql : Arrays.asList("SHOW STORAGE POLICY", "SHOW STORAGE POLICY USING",
"SHOW STORAGE POLICY USING FOR example_policy")) {
ShowStoragePolicyCommand command = (ShowStoragePolicyCommand) new NereidsParser().parseSingle(sql);
Assertions.assertEquals(names(command.doRun(null, null).getMetaData()), names(command.getMetaData()),
sql);
}
// Without USING it lists the policies themselves, not the objects that use them.
Assertions.assertEquals(names(StoragePolicy.STORAGE_META_DATA), names(((ShowStoragePolicyCommand)
new NereidsParser().parseSingle("SHOW STORAGE POLICY")).getMetaData()));
}
}
}
@@ -0,0 +1,405 @@
// 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.
import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.CallOption
import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.CallOptions
import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.CloseSessionRequest
import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.FlightClient
import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.FlightRuntimeException
import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.FlightStatusCode
import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.Location
import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.sql.FlightSqlClient
import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.memory.RootAllocator
import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.vector.types.TimeUnit
import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.vector.types.pojo.ArrowType
suite("test_schema_only", "arrow_flight_sql") {
def config = context.config.otherConfigs
def location = Location.forGrpcInsecure(config.get("extArrowFlightSqlHost"),
Integer.parseInt(config.get("extArrowFlightSqlPort")))
new RootAllocator().withCloseable { allocator ->
FlightClient.builder(allocator, location).build().withCloseable { flight ->
def token = flight.authenticateBasicToken(config.get("extArrowFlightSqlUser"),
config.get("extArrowFlightSqlPassword")).get()
CallOption[] options = [token, CallOptions.timeout(30, java.util.concurrent.TimeUnit.SECONDS)]
def client = new FlightSqlClient(flight)
String showTimeout = "SHOW VARIABLES LIKE 'query_timeout'"
def readVariable = { String name ->
def info = client.execute("SHOW VARIABLES LIKE '${name}'", options)
assertEquals(1, info.getEndpoints().size())
def values = []
flight.getStream(info.getEndpoints().get(0).getTicket(), options).withCloseable { stream ->
while (stream.next()) {
def root = stream.getRoot()
for (int i = 0; i < root.getRowCount(); ++i) {
values.add(root.getVector("Value").getObject(i).toString())
}
}
}
assertEquals(1, values.size())
return values[0]
}
def readTimeout = { readVariable("query_timeout") }
def prepareSchema = { String query ->
def prepared = client.prepare(query, options)
try {
// ADBC ExecuteSchema consumes this schema before executing or fetching FlightInfo.
def schema = prepared.getResultSetSchema()
assertTrue(schema.getFields()*.getName() != ["ResultMeta"])
assertEquals(schema, prepared.fetchSchema(options).getSchema())
return schema
} finally {
prepared.close(options)
}
}
try {
client.execute("USE information_schema", options)
// Prepare must bind unaliased expressions in query mode, as execution does.
// Comparing metadata-only calls misses synthetic aliases that execution rejects.
["1", "unhex('616263')", "CAST(CAST('2026-01-01' AS DATETIME(6)) AS STRING)"].each { expression ->
def prepared = client.prepare("SELECT ${expression}", options)
try {
def advertised = prepared.getResultSetSchema()
assertEquals([expression], advertised.getFields()*.getName())
def executed = prepared.execute(options).getSchema()
assertEquals(advertised.getFields()*.getName(), executed.getFields()*.getName())
assertEquals(advertised.getFields()*.getType(), executed.getFields()*.getType())
assertEquals(advertised, prepared.fetchSchema(options).getSchema())
} finally {
prepared.close(options)
}
}
def cases = [
["SELECT CAST(1 AS BIGINT) AS id, CAST(NULL AS VARCHAR(20)) AS name",
["id", "name"], [new ArrowType.Int(64, true), new ArrowType.Utf8()]],
["""SELECT CAST('2025-01-02 03:04:05.123456' AS DATETIME(6)) AS event_time,
ARRAY('first', 'second') AS tags,
MAP('key', CAST(1 AS BIGINT)) AS properties,
NAMED_STRUCT('amount', CAST(1 AS DECIMAL(18,4))) AS record""",
["event_time", "tags", "properties", "record"],
[new ArrowType.Timestamp(TimeUnit.MICROSECOND, null), new ArrowType.List(),
new ArrowType.Map(false), new ArrowType.Struct()]],
["SELECT TABLE_CATALOG, TABLE_SCHEMA, TABLE_NAME FROM information_schema.tables WHERE 1=0",
["TABLE_CATALOG", "TABLE_SCHEMA", "TABLE_NAME"],
[new ArrowType.Utf8(), new ArrowType.Utf8(), new ArrowType.Utf8()]],
["SELECT group_concat_state('x') AS state", ["state"], [new ArrowType.Utf8()]],
[showTimeout, ["Variable_name", "Value", "Default_Value", "Changed"],
[new ArrowType.Utf8(), new ArrowType.Utf8(), new ArrowType.Utf8(), new ArrowType.Utf8()]]
]
[
["SHOW PROCEDURE STATUS", ["ProcedureName", "CatalogId", "DbId", "DbName", "PackageName",
"OwnerName", "CreateTime", "ModifyTime"]],
["EXPLAIN SELECT 1", ["Explain String(Nereids Planner)"]],
["EXPLAIN PLAN PROCESS SELECT 1", ["Rule", "Before", "After"]],
["SHOW STORAGE POLICY", ["PolicyName", "Id", "Version", "Type", "StorageResource",
"CooldownDatetime", "CooldownTtl"]],
["SHOW STORAGE POLICY USING", ["PolicyName", "Database", "Table", "Partitions"]],
["INSERT INTO schema_source VALUES (1)", ["StatusResult"]],
["UPDATE schema_source SET id=1", ["StatusResult"]],
["DELETE FROM schema_source WHERE id=1", ["StatusResult"]],
["MERGE INTO schema_target t USING schema_source s ON t.id=s.id WHEN MATCHED THEN DELETE",
["StatusResult"]],
["KILL QUERY 123", ["StatusResult"]],
["BEGIN", ["StatusResult"]], ["COMMIT", ["StatusResult"]], ["ROLLBACK", ["StatusResult"]]
].each { entry ->
def schema = prepareSchema(entry[0])
assertEquals(entry[1], schema.getFields()*.getName())
assertEquals(schema, client.getExecuteSchema(entry[0], options).getSchema())
assertTrue(schema.getFields().every { it.getType() == new ArrowType.Utf8() })
}
// A frontend-side result is checked against its prepared schema in full, so the header Prepare
// tells without running a command must be the one running it sends: run the commands above that
// have no side effects, and the other frontend-side ones this suite prepares.
def executeAsPrepared = { String query ->
def prepared = client.prepare(query, options)
try {
def advertised = prepared.getResultSetSchema()
assertTrue(advertised.getFields()*.getName() != ["ResultMeta"], query)
def info = prepared.execute(options)
assertEquals(advertised, info.getSchema(), query)
flight.getStream(info.getEndpoints()[0].getTicket(), options).withCloseable { stream ->
while (stream.next()) {
assertEquals(advertised, stream.getRoot().getSchema(), query)
}
}
} finally {
prepared.close(options)
}
}
["SHOW PROCEDURE STATUS", "EXPLAIN SELECT 1", "EXPLAIN PLAN PROCESS SELECT 1", "SHOW STORAGE POLICY",
"SHOW STORAGE POLICY USING", "SHOW STORAGE POLICY USING FOR schema_only_missing_policy",
showTimeout, "SHOW FRONTEND CONFIG", "SHOW PROC '/'"].each { query ->
executeAsPrepared(query)
}
// Re-parsing a prepared statement must not silently change its advertised result type.
String previousMode = readVariable("sql_mode")
try {
[true, false].each { fetchSchema ->
client.execute("SET sql_mode=''", options)
def prepared = client.prepare("SELECT 1 || 2 AS x", options)
try {
assertEquals(new ArrowType.Bool(), prepared.getResultSetSchema().getFields()[0].getType())
client.execute("SET sql_mode='PIPES_AS_CONCAT'", options)
try {
if (fetchSchema) {
prepared.fetchSchema(options)
} else {
prepared.execute(options)
}
assertTrue(false, "A changed result schema must expire the prepared handle")
} catch (FlightRuntimeException error) {
assertEquals(FlightStatusCode.NOT_FOUND, error.status().code())
}
} finally {
prepared.close(options)
}
}
} finally {
client.execute("SET sql_mode='${previousMode}'", options)
}
cases.eachWithIndex { entry, index ->
def schema = prepareSchema(entry[0])
assertEquals(entry[1], schema.getFields()*.getName())
assertEquals(entry[2], schema.getFields()*.getType())
assertEquals(schema, client.getExecuteSchema(entry[0], options).getSchema())
if (index == 0) {
assertTrue(!schema.getFields()[0].isNullable())
assertTrue(schema.getFields()[1].isNullable())
} else if (index == 1) {
assertEquals(new ArrowType.Utf8(), schema.getFields()[1].getChildren()[0].getType())
def entries = schema.getFields()[2].getChildren()[0]
assertTrue(!entries.isNullable())
assertTrue(!entries.getChildren()[0].isNullable())
assertEquals(new ArrowType.Int(64, true), entries.getChildren()[1].getType())
assertEquals(new ArrowType.Decimal(18, 4, 128),
schema.getFields()[3].getChildren()[0].getType())
}
}
// Resolving command metadata must not execute SET or leak query-scoped SET_VAR hints.
def originalTimeout = readTimeout()
def changedTimeout = originalTimeout == "17" ? "19" : "17"
def setQuery = "SET query_timeout=${changedTimeout}".toString()
assertEquals(["StatusResult"], prepareSchema(setQuery).getFields()*.getName())
client.getExecuteSchema(setQuery, options)
assertEquals(originalTimeout, readTimeout())
def hintedQuery = "SELECT /*+ SET_VAR(query_timeout=${changedTimeout}) */ 1 AS id".toString()
prepareSchema(hintedQuery)
client.getExecuteSchema(hintedQuery, options)
assertEquals(originalTimeout, readTimeout())
["SELEC 1", "SELECT definitely_missing_schema_column"].each { query ->
[true, false].each { preparedRoute ->
FlightRuntimeException failure = null
try {
if (preparedRoute) {
prepareSchema(query)
} else {
client.getExecuteSchema(query, options)
}
} catch (FlightRuntimeException e) {
failure = e
}
assertTrue(failure != null, "Schema discovery accepted invalid SQL: ${query}")
assertEquals(FlightStatusCode.INVALID_ARGUMENT, failure.status().code())
assertEquals(["id", "name"], prepareSchema(cases[0][0]).getFields()*.getName())
assertEquals(originalTimeout, readTimeout())
}
}
assertEquals(6, prepareSchema("SHOW FRONTEND CONFIG").getFields().size())
assertEquals(prepareSchema("SHOW PROC '/'"), client.getExecuteSchema("SHOW PROC '/'", options).getSchema())
// A statement whose result schema only running it can tell keeps the placeholder on Prepare,
// since the Arrow Flight SQL JDBC driver prepares every statement and does not fall back when
// Prepare refuses one; GetSchema, asked for that schema itself, answers UNIMPLEMENTED. So does a
// result with an untyped NULL column, at any depth: BE sends it as nullable Bool, not as Null.
["SHOW PYTHON PACKAGES IN '3.11'", "SHOW QUERY STATS", "DESC information_schema.tables",
"WARM UP SELECT * FROM schema_source", "SET @schema_only = 1; SELECT @schema_only AS v",
"SELECT NULL AS n", "SELECT ARRAY(NULL) AS a"].each { query ->
def placeholder = client.prepare(query, options)
try {
assertEquals(["ResultMeta"], placeholder.getResultSetSchema().getFields()*.getName())
FlightRuntimeException failure = null
try {
placeholder.fetchSchema(options)
} catch (FlightRuntimeException e) {
failure = e
}
assertTrue(failure != null, "GetSchema answered for a placeholder handle: ${query}")
assertEquals(FlightStatusCode.UNIMPLEMENTED, failure.status().code())
} finally {
placeholder.close(options)
}
FlightRuntimeException failure = null
try {
client.getExecuteSchema(query, options)
} catch (FlightRuntimeException e) {
failure = e
}
assertTrue(failure != null, "GetSchema answered without running: ${query}")
assertEquals(FlightStatusCode.UNIMPLEMENTED, failure.status().code())
}
// Such a handle still runs, unchecked, as it did before Prepare analyzed anything.
def unchecked = client.prepare("SET @schema_only = 7; SELECT @schema_only AS v", options)
try {
def info = unchecked.execute(options)
def values = []
// The query's result is on the backends: read each endpoint from where it says it is.
info.getEndpoints().each { endpoint ->
FlightClient.builder(allocator, endpoint.getLocations()[0]).build().withCloseable { be ->
be.getStream(endpoint.getTicket(), options).withCloseable { stream ->
while (stream.next()) {
def root = stream.getRoot()
for (int i = 0; i < root.getRowCount(); ++i) {
values.add(root.getVector("v").getObject(i).toString())
}
}
}
}
}
assertEquals(["7"], values)
} finally {
unchecked.close(options)
}
// A query such as `SELECT org_id, NULL FROM t` runs unchecked the same way, every time, instead of
// having its handle expired by the Bool that BE sends for the NULL column.
def nulls = client.prepare(
"SELECT k, NULL AS n FROM (SELECT 1 AS k UNION ALL SELECT 2 AS k) t ORDER BY k", options)
try {
assertEquals(["ResultMeta"], nulls.getResultSetSchema().getFields()*.getName())
2.times {
def info = nulls.execute(options)
def keys = []
def nullCount = 0
info.getEndpoints().each { endpoint ->
FlightClient.builder(allocator, endpoint.getLocations()[0]).build().withCloseable { be ->
be.getStream(endpoint.getTicket(), options).withCloseable { stream ->
while (stream.next()) {
def root = stream.getRoot()
for (int i = 0; i < root.getRowCount(); ++i) {
keys.add(root.getVector("k").getObject(i).toString())
nullCount += root.getVector("n").isNull(i) ? 1 : 0
}
}
}
}
}
assertEquals(["1", "2"], keys.sort())
assertEquals(2, nullCount)
}
} finally {
nulls.close(options)
}
// WARM UP SELECT answers OK instead of its statistics when planning proves its source empty, as for a
// table without data: its placeholder handle runs that unchecked rather than refusing it.
client.execute("CREATE DATABASE IF NOT EXISTS ${context.dbName}", options)
client.execute("DROP TABLE IF EXISTS ${context.dbName}.schema_only_warmup_source", options)
client.execute("""CREATE TABLE ${context.dbName}.schema_only_warmup_source (id INT)
DUPLICATE KEY(id) DISTRIBUTED BY HASH(id) BUCKETS 1
PROPERTIES ("replication_num" = "1")""", options)
def warmup = client.prepare("WARM UP SELECT * FROM ${context.dbName}.schema_only_warmup_source", options)
try {
assertEquals(["ResultMeta"], warmup.getResultSetSchema().getFields()*.getName())
// Only cloud mode warms up internal tables; elsewhere WARM UP SELECT reads external tables only.
if (isCloudMode()) {
assertEquals(["StatusResult"], warmup.execute(options).getSchema().getFields()*.getName())
}
} finally {
warmup.close(options)
}
// The serialized group_concat state is advertised as UTF8; running it checks that against the
// type BE sends for it, as every prepared execution does.
def state = client.prepare("SELECT group_concat_state('x') AS state", options)
try {
def advertised = state.getResultSetSchema()
assertEquals([new ArrowType.Utf8()], advertised.getFields()*.getType())
assertEquals(advertised.getFields()*.getType(), state.execute(options).getSchema().getFields()*.getType())
} finally {
state.close(options)
}
// Executing an explicit namespace change must not invalidate that same handle.
["USE mysql", "USE mysql; SHOW VARIABLES LIKE 'query_timeout'"].each { query ->
client.execute("USE information_schema", options)
def changing = client.prepare(query, options)
try {
def expected = changing.getResultSetSchema()
2.times {
def info = changing.execute(options)
assertEquals(expected, info.getSchema())
assertEquals(expected, changing.fetchSchema(options).getSchema())
flight.getStream(info.getEndpoints()[0].getTicket(), options).withCloseable { stream ->
while (stream.next()) {
assertEquals(expected, stream.getRoot().getSchema())
}
}
}
} finally {
changing.close(options)
}
}
def help = client.prepare("HELP 'schema_only_unknown_help_topic'", options)
try {
def expected = help.getResultSetSchema()
assertEquals(["name", "is_it_category"], expected.getFields()*.getName())
assertEquals(expected, help.fetchSchema(options).getSchema())
def info = help.execute(options)
assertEquals(expected, info.getSchema())
flight.getStream(info.getEndpoints()[0].getTicket(), options).withCloseable { stream ->
while (stream.next()) {
assertEquals(expected, stream.getRoot().getSchema())
}
}
} finally {
help.close(options)
}
assertEquals(["Snapshot", "Timestamp", "Status"],
prepareSchema("SHOW SNAPSHOT ON schema_repository").getFields()*.getName())
assertEquals(["Snapshot", "Timestamp", "Database", "Details", "Status"],
prepareSchema("SHOW SNAPSHOT ON schema_repository WHERE SNAPSHOT='s' AND TIMESTAMP='t'")
.getFields()*.getName())
client.execute("USE information_schema", options)
def prepared = client.prepare("SELECT TABLE_NAME FROM tables WHERE 1=0", options)
try {
client.execute("USE mysql", options)
// A prepared header must not silently describe a different database's table.
[true, false].each { schemaRoute ->
FlightRuntimeException failure = null
try {
if (schemaRoute) {
prepared.fetchSchema(options)
} else {
prepared.execute(options)
}
} catch (FlightRuntimeException e) {
failure = e
}
assertTrue(failure != null)
assertEquals(FlightStatusCode.NOT_FOUND, failure.status().code())
}
} finally {
prepared.close(options)
}
} finally {
client.closeSession(new CloseSessionRequest(), options)
}
}
}
}