Add fallback from cancelQuery to cancelFlightInfo

kou · Jun 21, 2023 · b440cab · b440cab
1 parent f5cb9ab
commit b440cab
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 1 deletion.
diff --git a/java/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlProducer.java b/java/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlProducer.java
@@ -457,7 +457,38 @@ default void closeFlightInfo(FlightInfo info, CallContext context, StreamListene
    */
   @Deprecated
   default void cancelQuery(FlightInfo info, CallContext context, StreamListener<CancelResult> listener) {
-    listener.onError(CallStatus.UNIMPLEMENTED.toRuntimeException());
+    cancelFlightInfo(info, context, new StreamListener<CancelStatus>() {
+      @Override
+      public void onNext(CancelStatus val) {
+        switch (val) {
+          case UNSPECIFIED:
+            listener.onNext(CancelResult.UNSPECIFIED);
+            break;
+          case CANCELLED:
+            listener.onNext(CancelResult.CANCELLED);
+            break;
+          case CANCELLING:
+            listener.onNext(CancelResult.CANCELLING);
+            break;
+          case NOT_CANCELLABLE:
+            listener.onNext(CancelResult.NOT_CANCELLABLE);
+            break;
+          default:
+            // XXX: CheckStyle requires a default clause which arguably makes the code worse.
+            throw new AssertionError("Unknown enum variant " + val);
+        }
+      }
+
+      @Override
+      public void onError(Throwable t) {
+        listener.onError(t);
+      }
+
+      @Override
+      public void onCompleted() {
+        listener.onCompleted();
+      }
+    });
   }
 
   /**

diff --git a/java/flight/flight-sql/src/test/java/org/apache/arrow/flight/TestFlightSql.java b/java/flight/flight-sql/src/test/java/org/apache/arrow/flight/TestFlightSql.java
@@ -26,6 +26,8 @@
 import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.CoreMatchers.notNullValue;
 import static org.hamcrest.CoreMatchers.nullValue;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
@@ -991,4 +993,33 @@ public void testQueryWithNoResultsShouldNotHang() throws Exception {
       );
     }
   }
+
+  @Test
+  public void testCancelFlightInfo() {
+    FlightInfo info = sqlClient.getSqlInfo();
+    FlightRuntimeException fre = assertThrows(FlightRuntimeException.class, () -> sqlClient.cancelFlightInfo(info));
+    assertEquals(FlightStatusCode.UNIMPLEMENTED, fre.status().code());
+  }
+
+  @Test
+  public void testCancelQuery() {
+    FlightInfo info = sqlClient.getSqlInfo();
+    FlightRuntimeException fre = assertThrows(FlightRuntimeException.class, () -> sqlClient.cancelQuery(info));
+    assertEquals(FlightStatusCode.UNIMPLEMENTED, fre.status().code());
+  }
+
+  @Test
+  public void testCloseInfo() {
+    FlightInfo info = sqlClient.getSqlInfo();
+    FlightRuntimeException fre = assertThrows(FlightRuntimeException.class, () -> sqlClient.closeFlightInfo(info));
+    assertEquals(FlightStatusCode.UNIMPLEMENTED, fre.status().code());
+  }
+
+  @Test
+  public void testRefreshEndpoint() {
+    FlightInfo info = sqlClient.getSqlInfo();
+    FlightRuntimeException fre = assertThrows(FlightRuntimeException.class,
+        () -> sqlClient.refreshFlightEndpoint(info.getEndpoints().get(0)));
+    assertEquals(FlightStatusCode.UNIMPLEMENTED, fre.status().code());
+  }
 }