Implement CancelFlightInfo

kou · Jun 21, 2023 · 3c8f943 · 3c8f943
1 parent b440cab
commit 3c8f943
Showing 1 changed file with 9 additions and 7 deletions.
diff --git a/go/arrow/internal/flight_integration/scenario.go b/go/arrow/internal/flight_integration/scenario.go
@@ -2179,22 +2179,24 @@ func (m *flightSqlScenarioTester) BeginTransaction(context.Context, flightsql.Ac
 	return []byte(transactionID), nil
 }
 
-func (m *flightSqlScenarioTester) CancelQuery(_ context.Context, request flightsql.ActionCancelQueryRequest) (flightsql.CancelResult, error) {
-	if err := assertEq(1, len(request.GetInfo().Endpoint)); err != nil {
-		return flightsql.CancelResultUnspecified, fmt.Errorf("%w: expected 1 endpoint for CancelQuery", err)
+func (m *flightSqlScenarioTester) CancelFlightInfo(_ context.Context, info *flight.FlightInfo) (flight.CancelFlightInfoResult, error) {
+	result := flight.CancelFlightInfoResult{Status: flight.CancelStatusUnspecified}
+	if err := assertEq(1, len(info.Endpoint)); err != nil {
+		return result, fmt.Errorf("%w: expected 1 endpoint for CancelQuery", err)
 	}
 
-	endpoint := request.GetInfo().Endpoint[0]
+	endpoint := info.Endpoint[0]
 	tkt, err := flightsql.GetStatementQueryTicket(endpoint.Ticket)
 	if err != nil {
-		return flightsql.CancelResultUnspecified, err
+		return result, err
 	}
 
 	if err := assertEq([]byte("PLAN HANDLE"), tkt.GetStatementHandle()); err != nil {
-		return flightsql.CancelResultUnspecified, fmt.Errorf("%w: unexpected ticket in CancelQuery", err)
+		return result, fmt.Errorf("%w: unexpected ticket in CancelQuery", err)
 	}
 
-	return flightsql.CancelResultCancelled, nil
+	result.Status = flight.CancelStatusCancelled
+	return result, nil
 }
 
 func (m *flightSqlScenarioTester) EndSavepoint(_ context.Context, request flightsql.ActionEndSavepointRequest) error {