feat: make the polling extension opt-in

By default, the FlexConnect will conform to the Arrow Flight RPC spec.
However, if an opt-in header is present, it will use the polling
extension used by GoodData.
This allows for things like query cancellation.
Ideally, we would use the PollFlightInfo from the Arrow Flight RPC
but unfortunately it is not yet available in PyArrow.

JIRA: CQ-1124
risk: low

Author: no23reason

gooddata-flexconnect/gooddata_flexconnect/function/flight_methods.py

@@ -13,6 +13,7 @@
     FlightServerMethods,
     ServerContext,
     TaskExecutionResult,
+    TaskWaitTimeoutError,
     flight_server_methods,
 )
 
@@ -21,13 +22,20 @@
     CancelInvocation,
     RetryInvocation,
     SubmitInvocation,
-    extract_invocation_from_descriptor,
+    extract_pollable_invocation_from_descriptor,
+    extract_submit_invocation_from_descriptor,
 )
 from gooddata_flexconnect.function.function_registry import FlexConnectFunctionRegistry
 from gooddata_flexconnect.function.function_task import FlexConnectFunctionTask
 
 _LOGGER = structlog.get_logger("gooddata_flexconnect.rpc")
 
+POLLING_HEADER_NAME = "x-quiver-pollable"
+"""
+If this header is present on the get flight info call, the polling extension will be used.
+Otherwise the basic do get will be used.
+"""
+
 
 def _prepare_poll_error(task_id: str) -> pyarrow.flight.FlightError:
     return ErrorInfo.poll(
@@ -122,28 +130,69 @@ def _prepare_flight_info(
             total_bytes=-1,
         )
 
-    ###################################################################
-    # Implementation of Flight RPC methods
-    ###################################################################
-
-    def list_flights(
-        self, context: pyarrow.flight.ServerCallContext, criteria: bytes
-    ) -> Generator[pyarrow.flight.FlightInfo, None, None]:
+    def _get_flight_info_no_polling(
+        self,
+        context: pyarrow.flight.ServerCallContext,
+        descriptor: pyarrow.flight.FlightDescriptor,
+    ) -> pyarrow.flight.FlightInfo:
+        """
+        Basic DoGetInfo flow with no polling extension.
+        This conforms to the mainline Arrow Flight RPC specification.
+        """
         structlog.contextvars.bind_contextvars(peer=context.peer())
-        _LOGGER.info("list_flights", available_funs=self._registry.function_names)
+        invocation = extract_submit_invocation_from_descriptor(descriptor)
 
-        return (self._create_fun_info(fun) for fun in self._registry.functions.values())
+        task: Optional[FlexConnectFunctionTask] = None
 
-    def get_flight_info(
+        try:
+            task = self._prepare_task(context, invocation)
+            self._ctx.task_executor.submit(task)
+
+            try:
+                task_result = self._ctx.task_executor.wait_for_result(task.task_id, self._call_deadline)
+            except TaskWaitTimeoutError:
+                cancelled = self._ctx.task_executor.cancel(task.task_id)
+                _LOGGER.warning(
+                    "flexconnect_fun_call_timeout", task_id=task.task_id, fun=task.fun_name, cancelled=cancelled
+                )
+
+                raise ErrorInfo.for_reason(
+                    ErrorCode.TIMEOUT, f"GetFlightInfo timed out while waiting for task {task.task_id}."
+                ).to_timeout_error()
+
+            # if this bombs then there must be something really wrong because the task
+            # was clearly submitted and code was waiting for its completion. this invariant
+            # should not happen in this particular code path. The None return value may
+            # be applicable one day when polling is in use and a request comes to check whether
+            # particular task id finished
+            assert task_result is not None
+
+            return self._prepare_flight_info(task_id=task.task_id, task_result=task_result)
+        except Exception:
+            if task is not None:
+                _LOGGER.error(
+                    "get_flight_info_failed", task_id=task.task_id, fun=task.fun_name, exc_info=True, polling=False
+                )
+            else:
+                _LOGGER.error("flexconnect_fun_submit_failed", exc_info=True, polling=False)
+            raise
+
+    def _get_flight_info_polling(
         self,
         context: pyarrow.flight.ServerCallContext,
         descriptor: pyarrow.flight.FlightDescriptor,
     ) -> pyarrow.flight.FlightInfo:
+        """
+        DoGetInfo flow with polling extension.
+        This extends the mainline Arrow Flight RPC specification with polling capabilities using the RetryInfo
+        encoded into the FlightTimedOutError.extra_info.
+        Ideally, we would use the mainline PollFlightInfo, but that has yet to be implemented in the PyArrow library.
+        """
         structlog.contextvars.bind_contextvars(peer=context.peer())
+        invocation = extract_pollable_invocation_from_descriptor(descriptor)
 
         task_id: str
         fun_name: Optional[str] = None
-        invocation = extract_invocation_from_descriptor(descriptor)
 
         if isinstance(invocation, CancelInvocation):
             # cancel the given task and raise cancellation exception
@@ -166,7 +215,7 @@ def get_flight_info(
                 task_id = task.task_id
                 fun_name = task.fun_name
             except Exception:
-                _LOGGER.error("flexconnect_fun_submit_failed", exc_info=True)
+                _LOGGER.error("flexconnect_fun_submit_failed", exc_info=True, polling=True)
                 raise
         else:
             # can be replaced by assert_never when we are on 3.11
@@ -188,9 +237,36 @@ def get_flight_info(
             # how to poll for the results
             raise _prepare_poll_error(task_id)
         except Exception:
-            _LOGGER.error("get_flight_info_failed", task_id=task_id, fun=fun_name, exc_info=True)
+            _LOGGER.error("get_flight_info_failed", task_id=task_id, fun=fun_name, exc_info=True, polling=True)
             raise
 
+    ###################################################################
+    # Implementation of Flight RPC methods
+    ###################################################################
+
+    def list_flights(
+        self, context: pyarrow.flight.ServerCallContext, criteria: bytes
+    ) -> Generator[pyarrow.flight.FlightInfo, None, None]:
+        structlog.contextvars.bind_contextvars(peer=context.peer())
+        _LOGGER.info("list_flights", available_funs=self._registry.function_names)
+
+        return (self._create_fun_info(fun) for fun in self._registry.functions.values())
+
+    def get_flight_info(
+        self,
+        context: pyarrow.flight.ServerCallContext,
+        descriptor: pyarrow.flight.FlightDescriptor,
+    ) -> pyarrow.flight.FlightInfo:
+        structlog.contextvars.bind_contextvars(peer=context.peer())
+
+        headers = self.call_info_middleware(context).headers
+        allow_polling = headers.get(POLLING_HEADER_NAME) is not None
+
+        if allow_polling:
+            return self._get_flight_info_polling(context, descriptor)
+        else:
+            return self._get_flight_info_no_polling(context, descriptor)
+
     def do_get(
         self,
         context: pyarrow.flight.ServerCallContext,

gooddata-flexconnect/gooddata_flexconnect/function/function_invocation.py

@@ -53,42 +53,51 @@ class SubmitInvocation:
     """
 
 
-def extract_invocation_from_descriptor(
+def extract_submit_invocation_from_descriptor(descriptor: pyarrow.flight.FlightDescriptor) -> SubmitInvocation:
+    """
+    Given a flight descriptor, extract the invocation information from it.
+    Do not allow the polling-related variants.
+    """
+    try:
+        payload = orjson.loads(descriptor.command)
+    except Exception:
+        raise ErrorInfo.bad_argument(
+            "Incorrect FlexConnect function invocation. The invocation payload is not a valid JSON."
+        )
+
+    function_name = payload.get("functionName")
+    if function_name is None or not len(function_name):
+        raise ErrorInfo.bad_argument(
+            "Incorrect FlexConnect function invocation. The invocation payload does not specify 'functionName'."
+        )
+
+    parameters = payload.get("parameters") or {}
+    columns = parameters.get("columns")
+
+    return SubmitInvocation(
+        function_name=function_name, parameters=parameters, columns=columns, command=descriptor.command
+    )
+
+
+def extract_pollable_invocation_from_descriptor(
     descriptor: pyarrow.flight.FlightDescriptor,
 ) -> Union[RetryInvocation, CancelInvocation, SubmitInvocation]:
     """
     Given a flight descriptor, extract the invocation information from it.
+    Allow also the polling-related variants.
     """
-
     if descriptor.command is None or not len(descriptor.command):
         raise ErrorInfo.bad_argument(
             "Incorrect FlexConnect function invocation. Flight descriptor must contain command "
             "with the invocation payload."
         )
 
+    # we are in the polling-enabled realm: try parsing the retry and cancel descriptors first
     if descriptor.command.startswith(b"c:"):
         task_id = descriptor.command[2:].decode()
         return CancelInvocation(task_id)
     elif descriptor.command.startswith(b"r:"):
         task_id = descriptor.command[2:].decode()
         return RetryInvocation(task_id)
-    else:
-        try:
-            payload = orjson.loads(descriptor.command)
-        except Exception:
-            raise ErrorInfo.bad_argument(
-                "Incorrect FlexConnect function invocation. The invocation payload is not a valid JSON."
-            )
-
-        function_name = payload.get("functionName")
-        if function_name is None or not len(function_name):
-            raise ErrorInfo.bad_argument(
-                "Incorrect FlexConnect function invocation. The invocation payload does not specify 'functionName'."
-            )
-
-        parameters = payload.get("parameters") or {}
-        columns = parameters.get("columns")
-
-        return SubmitInvocation(
-            function_name=function_name, parameters=parameters, columns=columns, command=descriptor.command
-        )
+
+    return extract_submit_invocation_from_descriptor(descriptor)

gooddata-flexconnect/tests/server/test_flexconnect_server.py

@@ -3,12 +3,18 @@
 import orjson
 import pyarrow.flight
 import pytest
+from gooddata_flexconnect.function.flight_methods import POLLING_HEADER_NAME
 from gooddata_flight_server import ErrorCode, ErrorInfo, RetryInfo
 
 from tests.assert_error_info import assert_error_code
 from tests.server.conftest import flexconnect_server
 
 
+@pytest.fixture
+def call_options_with_polling():
+    return pyarrow.flight.FlightCallOptions(headers=[(POLLING_HEADER_NAME.encode(), b"true")])
+
+
 def test_basic_function():
     """
     This function should return immediately when called, no polling necessary.
@@ -94,7 +100,7 @@ def test_basic_function_tls(tls_ca_cert):
         assert data.column_names == ["col1", "col2", "col3"]
 
 
-def test_function_with_polling():
+def test_function_with_polling(call_options_with_polling):
     """
     Flight RPC implementation that invokes FlexConnect can return a polling info.
 
@@ -114,7 +120,7 @@ def test_function_with_polling():
         # the function is set to sleep a bit longer than the polling interval,
         # so the first iteration returns retry info in the exception
         with pytest.raises(pyarrow.flight.FlightTimedOutError) as e:
-            c.get_flight_info(descriptor)
+            c.get_flight_info(descriptor, call_options_with_polling)
 
         assert e.value is not None
         assert_error_code(ErrorCode.POLL, e.value)
@@ -124,21 +130,21 @@ def test_function_with_polling():
 
         # use the retry info to poll again for the result,
         # now it should be ready and returned normally
-        info = c.get_flight_info(retry_info.retry_descriptor)
+        info = c.get_flight_info(retry_info.retry_descriptor, call_options_with_polling)
         data: pyarrow.Table = c.do_get(info.endpoints[0].ticket).read_all()
 
         assert len(data) == 3
         assert data.column_names == ["col1", "col2", "col3"]
 
         # also check that trying to cancel already completed task results in cancelled with correct code
         with pytest.raises(pyarrow.flight.FlightCancelledError) as e:
-            c.get_flight_info(retry_info.cancel_descriptor)
+            c.get_flight_info(retry_info.cancel_descriptor, call_options_with_polling)
 
         assert e.value is not None
         assert_error_code(ErrorCode.COMMAND_CANCELLED, e.value)
 
 
-def test_function_with_call_deadline():
+def test_function_with_call_deadline(call_options_with_polling):
     """
     Flight RPC implementation that invokes FlexConnect can be setup with
     deadline for the invocation duration (done by GetFlightInfo).
@@ -163,7 +169,7 @@ def test_function_with_call_deadline():
 
         # the initial submit returns polling info
         with pytest.raises(pyarrow.flight.FlightTimedOutError) as e:
-            c.get_flight_info(descriptor)
+            c.get_flight_info(descriptor, call_options_with_polling)
 
         assert e.value is not None
         assert_error_code(ErrorCode.POLL, e.value)
@@ -173,19 +179,19 @@ def test_function_with_call_deadline():
 
         # the next poll still returns polling info
         with pytest.raises(pyarrow.flight.FlightTimedOutError) as e:
-            c.get_flight_info(retry_info.retry_descriptor)
+            c.get_flight_info(retry_info.retry_descriptor, call_options_with_polling)
 
         assert e.value is not None
         assert_error_code(ErrorCode.POLL, e.value)
 
         # the third one reaches the deadline so the Timeout code is returned instead
         with pytest.raises(pyarrow.flight.FlightTimedOutError) as e:
-            c.get_flight_info(retry_info.retry_descriptor)
+            c.get_flight_info(retry_info.retry_descriptor, call_options_with_polling)
 
         assert_error_code(ErrorCode.TIMEOUT, e.value)
 
 
-def test_function_with_cancellation():
+def test_function_with_cancellation(call_options_with_polling):
     """
     Run a long-running function and cancel it after one poll iteration.
     """
@@ -202,7 +208,7 @@ def test_function_with_cancellation():
 
         # the initial submit returns polling info
         with pytest.raises(pyarrow.flight.FlightTimedOutError) as e:
-            c.get_flight_info(descriptor)
+            c.get_flight_info(descriptor, call_options_with_polling)
 
         assert e.value is not None
         assert_error_code(ErrorCode.POLL, e.value)
@@ -212,14 +218,14 @@ def test_function_with_cancellation():
 
         # use the poll info to cancel the task
         with pytest.raises(pyarrow.flight.FlightCancelledError) as e:
-            c.get_flight_info(retry_info.cancel_descriptor)
+            c.get_flight_info(retry_info.cancel_descriptor, call_options_with_polling)
 
         assert e.value is not None
         assert_error_code(ErrorCode.COMMAND_CANCELLED, e.value)
 
         # even multiple cancellations return the same error
         with pytest.raises(pyarrow.flight.FlightCancelledError) as e:
-            c.get_flight_info(retry_info.cancel_descriptor)
+            c.get_flight_info(retry_info.cancel_descriptor, call_options_with_polling)
 
         assert e.value is not None
         assert_error_code(ErrorCode.COMMAND_CANCELLED, e.value)