database/doris
2
0
Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

[pick](branch-2.1) pick #43281 (#44020)

Browse Source 
pick #43281
This commit is contained in:
Xinyi Zou
2024-11-16 21:53:21 +08:00
committed by GitHub
parent 44ed5d6120
commit ea61206233
26 changed files with 884 additions and 253 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/DorisFlightSqlProducer.java
View File

@ -27,6 +27,7 @@ import org.apache.doris.qe.ConnectContext;
import org.apache.doris.qe.QueryState.MysqlStateType;
import org.apache.doris.service.arrowflight.results.FlightSqlResultCacheEntry;
import org.apache.doris.service.arrowflight.sessions.FlightSessionsManager;
import org.apache.doris.thrift.TUniqueId;
import com.google.common.base.Preconditions;
import com.google.protobuf.Any;
@ -224,19 +225,36 @@ public class DorisFlightSqlProducer implements FlightSqlProducer, AutoCloseable
}
} else {
// Now only query stmt will pull results from BE.
final ByteString handle = ByteString.copyFromUtf8(
DebugUtil.printId(connectContext.getFinstId()) + ":" + query);
Schema schema = flightSQLConnectProcessor.fetchArrowFlightSchema(5000);
if (schema == null) {
throw CallStatus.INTERNAL.withDescription("fetch arrow flight schema is null")
.toRuntimeException();
}
TUniqueId queryId = connectContext.getFinstId();
// Ticket contains the IP and Brpc Port of the Doris BE node where the query result is located.
final ByteString handle = ByteString.copyFromUtf8(
DebugUtil.printId(queryId) + "&" + connectContext.getResultInternalServiceAddr().hostname
+ "&" + connectContext.getResultInternalServiceAddr().port + "&" + query);
TicketStatementQuery ticketStatement = TicketStatementQuery.newBuilder().setStatementHandle(handle)
.build();
Ticket ticket = new Ticket(Any.pack(ticketStatement).toByteArray());
// TODO Support multiple endpoints.
Location location = Location.forGrpcInsecure(connectContext.getResultFlightServerAddr().hostname,
connectContext.getResultFlightServerAddr().port);
Location location;
if (flightSQLConnectProcessor.getPublicAccessAddr().isSetHostname()) {
// In a production environment, it is often inconvenient to expose Doris BE nodes
// to the external network.
// However, a reverse proxy (such as nginx) can be added to all Doris BE nodes,
// and the external client will be randomly routed to a Doris BE node when connecting to nginx.
// The query results of Arrow Flight SQL will be randomly saved on a Doris BE node.
// If it is different from the Doris BE node randomly routed by nginx,
// data forwarding needs to be done inside the Doris BE node.
location = Location.forGrpcInsecure(flightSQLConnectProcessor.getPublicAccessAddr().hostname,
flightSQLConnectProcessor.getPublicAccessAddr().port);
} else {
location = Location.forGrpcInsecure(connectContext.getResultFlightServerAddr().hostname,
connectContext.getResultFlightServerAddr().port);
}
List<FlightEndpoint> endpoints = Collections.singletonList(new FlightEndpoint(ticket, location));
// TODO Set in BE callback after query end, Client will not callback.
return new FlightInfo(schema, descriptor, endpoints, -1, -1);

16
fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/FlightSqlConnectProcessor.java
View File

@ -53,11 +53,12 @@ import java.util.concurrent.TimeoutException;
/**
* Process one flgiht sql connection.
*
* <p>
* Must use try-with-resources.
*/
public class FlightSqlConnectProcessor extends ConnectProcessor implements AutoCloseable {
private static final Logger LOG = LogManager.getLogger(FlightSqlConnectProcessor.class);
private TNetworkAddress publicAccessAddr = new TNetworkAddress();
public FlightSqlConnectProcessor(ConnectContext context) {
super(context);
@ -66,6 +67,10 @@ public class FlightSqlConnectProcessor extends ConnectProcessor implements AutoC
context.setReturnResultFromLocal(true);
}
public TNetworkAddress getPublicAccessAddr() {
return publicAccessAddr;
}
public void prepare(MysqlCommand command) {
// set status of query to OK.
ctx.getState().reset();
@ -123,11 +128,12 @@ public class FlightSqlConnectProcessor extends ConnectProcessor implements AutoC
}
Status resultStatus = new Status(pResult.getStatus());
if (resultStatus.getErrorCode() != TStatusCode.OK) {
throw new RuntimeException(String.format("fetch arrow flight schema failed, finstId: %s, errmsg: %s",
DebugUtil.printId(tid), resultStatus.toString()));
throw new RuntimeException(String.format("fetch arrow flight schema failed, queryId: %s, errmsg: %s",
DebugUtil.printId(tid), resultStatus));
}
if (pResult.hasBeArrowFlightIp()) {
ctx.getResultFlightServerAddr().hostname = pResult.getBeArrowFlightIp().toStringUtf8();
if (pResult.hasBeArrowFlightIp() && pResult.hasBeArrowFlightPort()) {
publicAccessAddr.hostname = pResult.getBeArrowFlightIp().toStringUtf8();
publicAccessAddr.port = pResult.getBeArrowFlightPort();
}
if (pResult.hasSchema() && pResult.getSchema().size() > 0) {
RootAllocator rootAllocator = new RootAllocator(Integer.MAX_VALUE);