4

xinyiZzz · Sep 13, 2023 · 809bf30 · 809bf30
1 parent 2dc22e3
commit 809bf30
Show file tree

Hide file tree

Showing 6 changed files with 93 additions and 100 deletions.
diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java b/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java
@@ -2674,7 +2674,7 @@ public void executeArrowFlightQuery(FlightStatementContext flightStatementContex
             } catch (Exception e) {
                 queryScheduleSpan.recordException(e);
                 LOG.warn("Failed to coord exec Arrow Flight SQL, because: {}", e.getMessage(), e);
-                throw new InternalQueryExecutionException(e.getMessage() + Util.getRootCauseMessage(e), e);
+                throw new RuntimeException(e.getMessage() + Util.getRootCauseMessage(e), e);
             } finally {
                 queryScheduleSpan.end();
             }

diff --git a/...re/src/main/java/org/apache/doris/service/arrowflight/FlightServerBasicAuthValidator.java b/...re/src/main/java/org/apache/doris/service/arrowflight/FlightServerBasicAuthValidator.java
@@ -14,6 +14,9 @@
 // KIND, either express or implied.  See the License for the
 // specific language governing permissions and limitations
 // under the License.
+// This file is copied from
+// https://github.com/dremio/dremio-oss/blob/master/services/arrow-flight/src/main/java/com/dremio/service/flight/ServerCookieMiddleware.java
+// and modified by Doris
 
 package org.apache.doris.service.arrowflight;
 
@@ -25,6 +28,7 @@
 /**
  * authentication specialized implementation of BasicAuthValidator. Authenticates with provided
  * credentials.
+ * TODO
  */
 public class FlightServerBasicAuthValidator implements BasicServerAuthHandler.BasicAuthValidator {
 

diff --git a/...e/arrowflight/ServerCookieMiddleware.java → ...wflight/FlightServerCookieMiddleware.java b/...e/arrowflight/ServerCookieMiddleware.java → ...wflight/FlightServerCookieMiddleware.java
@@ -14,6 +14,9 @@
 // KIND, either express or implied.  See the License for the
 // specific language governing permissions and limitations
 // under the License.
+// This file is copied from
+// https://github.com/dremio/dremio-oss/blob/master/services/arrow-flight/src/main/java/com/dremio/service/flight/ServerCookieMiddleware.java
+// and modified by Doris
 
 package org.apache.doris.service.arrowflight;
 
@@ -30,31 +33,30 @@
 
 /**
  * ServerCookieMiddleware allows a FlightServer to retrieve cookies from the request as well as set outgoing cookies
+ * TODO
  */
-public final class ServerCookieMiddleware implements FlightServerMiddleware {
+public final class FlightServerCookieMiddleware implements FlightServerMiddleware {
     private RequestContext requestContext;
     private Map<String, String> cookieValues;
     private final CallHeaders incomingHeaders;
 
-    /**
-    * Factory to construct @see com.dremio.service.flight.ServerCookieMiddlewares
-    */
-    public static class Factory implements FlightServerMiddleware.Factory<ServerCookieMiddleware> {
+    public static class Factory implements FlightServerMiddleware.Factory<FlightServerCookieMiddleware> {
         /**
         * Construct a factory for receiving call headers.
         */
         public Factory() {
         }
 
         @Override
-        public ServerCookieMiddleware onCallStarted(CallInfo callInfo, CallHeaders incomingHeaders,
+        public FlightServerCookieMiddleware onCallStarted(CallInfo callInfo, CallHeaders incomingHeaders,
                                                     RequestContext context) {
-            return new ServerCookieMiddleware(callInfo, incomingHeaders, context);
+            return new FlightServerCookieMiddleware(callInfo, incomingHeaders, context);
         }
     }
 
 
-    private ServerCookieMiddleware(CallInfo callInfo, CallHeaders incomingHeaders, RequestContext requestContext) {
+    private FlightServerCookieMiddleware(CallInfo callInfo, CallHeaders incomingHeaders,
+            RequestContext requestContext) {
         this.incomingHeaders = incomingHeaders;
         this.requestContext = requestContext;
         this.cookieValues = new HashMap<String, String>();

diff --git a/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/FlightSqlService.java b/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/FlightSqlService.java
@@ -36,27 +36,27 @@ public class FlightSqlService {
     private final FlightServer flightServer;
     private volatile boolean running;
     public static final String FLIGHT_CLIENT_PROPERTIES_MIDDLEWARE = "client-properties-middleware";
-    public static final FlightServerMiddleware.Key<ServerCookieMiddleware> FLIGHT_CLIENT_PROPERTIES_MIDDLEWARE_KEY
+    public static final FlightServerMiddleware.Key<FlightServerCookieMiddleware> FLIGHT_CLIENT_PROPERTIES_MIDDLEWARE_KEY
             = FlightServerMiddleware.Key.of(FLIGHT_CLIENT_PROPERTIES_MIDDLEWARE);
 
     public FlightSqlService(int port) {
         BufferAllocator allocator = new RootAllocator();
         Location location = Location.forGrpcInsecure("0.0.0.0", port);
         FlightSqlServiceImpl producer = new FlightSqlServiceImpl(location);
         flightServer = FlightServer.builder(allocator, location, producer)
-            .middleware(FLIGHT_CLIENT_PROPERTIES_MIDDLEWARE_KEY,
-                new ServerCookieMiddleware.Factory())
-            .authHandler(new BasicServerAuthHandler(new FlightServerBasicAuthValidator())).build();
+                .middleware(FLIGHT_CLIENT_PROPERTIES_MIDDLEWARE_KEY,
+                        new FlightServerCookieMiddleware.Factory())
+                .authHandler(new BasicServerAuthHandler(new FlightServerBasicAuthValidator())).build();
     }
 
-    // start Flightsql protocol service
-    // return true if success, otherwise false
+    // start Arrow Flight SQL service, return true if success, otherwise false
     public boolean start() {
         try {
             flightServer.start();
-            LOG.info("Flightsql network service is started.");
+            running = true;
+            LOG.info("Arrow Flight SQL service is started.");
         } catch (IOException e) {
-            LOG.warn("Open Flightsql network service failed.", e);
+            LOG.error("Start Arrow Flight SQL service failed.", e);
             return false;
         }
         return true;
@@ -65,11 +65,10 @@ public boolean start() {
     public void stop() {
         if (running) {
             running = false;
-            // close server channel, make accept throw exception
             try {
                 flightServer.close();
             } catch (InterruptedException e) {
-                LOG.warn("close server channel failed.", e);
+                LOG.warn("close Arrow Flight SQL server failed.", e);
             }
         }
     }