fix: add flight info action

Author: rustyconover

src/query_farm_server_base/parameter_types.py

@@ -324,6 +324,28 @@ class TableFunctionFlightInfo(BaseModel):
     )
 
 
+class FlightInfo(BaseModel):
+    """
+    Parameters for a table function flight info request.
+    """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)  # for Pydantic v2
+
+    # The descriptor of the table function.
+    descriptor: flight.FlightDescriptor
+
+    at_unit: str | None = None
+    at_value: str | None = None
+
+    _validate_flight_descriptor = field_validator("descriptor", mode="before")(
+        deserialize_flight_descriptor
+    )
+
+
+def flight_info(action: flight.Action) -> FlightInfo:
+    return unpack_with_model(action, FlightInfo)
+
+
 def table_function_flight_info(action: flight.Action) -> TableFunctionFlightInfo:
     return unpack_with_model(action, TableFunctionFlightInfo)
 

src/query_farm_server_base/server.py

@@ -116,6 +116,7 @@ class ActionType(str, Enum):
     ENDPOINTS = "endpoints"
     LIST_SCHEMAS = "list_schemas"
     TABLE_FUNCTION_FLIGHT_INFO = "table_function_flight_info"
+    FLIGHT_INFO = "flight_info"
 
 
 @dataclass
@@ -263,6 +264,13 @@ def __init__(
                 False,
                 False,
             ),
+            ActionType.TABLE_FUNCTION_FLIGHT_INFO: ActionHandlerSpec(
+                self.action_flight_info,
+                parameter_types.flight_info,
+                lambda x: x.serialize(),
+                False,
+                False,
+            ),
             ActionType.LIST_SCHEMAS: ActionHandlerSpec(
                 self.action_list_schemas,
                 parameter_types.list_schemas,
@@ -536,6 +544,14 @@ def action_table_function_flight_info(
     ) -> flight.FlightInfo:
         self._unimplemented_action(ActionType.TABLE_FUNCTION_FLIGHT_INFO)
 
+    def action_flight_info(
+        self,
+        *,
+        context: CallContext[AccountType, TokenType],
+        parameters: parameter_types.TableFunctionFlightInfo,
+    ) -> flight.FlightInfo:
+        self._unimplemented_action(ActionType.FLIGHT_INFO)
+
     @abstractmethod
     def action_catalog_version(
         self,