[FSTORE-1661] Fix backwards compatibility (#469)

Author: aversey

python/hsfs/core/arrow_flight_client.py

@@ -39,6 +39,7 @@
 from hsfs.constructor import query
 from hsfs.core.variable_api import VariableApi
 from hsfs.storage_connector import StorageConnector
+from pyarrow.flight import FlightServerError
 from retrying import retry
 
 
@@ -96,6 +97,13 @@ def _should_retry_healthcheck(exception):
     )
 
 
+def _should_retry_certificate_registration(exception):
+    return (
+        _should_retry_healthcheck(exception)
+        or _is_feature_query_service_queue_full_error(exception)
+    )
+
+
 # Avoid unnecessary client init
 def is_data_format_supported(data_format: str, read_options: Optional[Dict[str, Any]]):
     if data_format not in ArrowFlightClient.SUPPORTED_FORMATS:
@@ -199,6 +207,12 @@ def __init__(self, disabled_for_session: bool = False):
 
         try:
             self._health_check()
+            if "get-version" in [action.type for action in self._connection.list_actions()]:
+                self._server_version = self._get_server_version()
+            else:
+                self._server_version = None
+            if self._server_version is None:
+                self._register_certificates()
         except Exception as e:
             _logger.debug("Failed to connect to Hopsworks Query Service.")
             _logger.exception(e)
@@ -293,6 +307,20 @@ def _health_check(self):
         list(self._connection.do_action(action, options=options))
         _logger.debug("Healthcheck succeeded.")
 
+    @retry(
+        wait_exponential_multiplier=1000,
+        stop_max_attempt_number=3,
+        retry_on_exception=_should_retry,
+    )
+    def _get_server_version(self):
+        _logger.debug("Acquiring the server version of Hopsworks Query Service.")
+        action = pyarrow.flight.Action("get-version", b"")
+        options = pyarrow.flight.FlightCallOptions(timeout=self.health_check_timeout)
+        for res in self._connection.do_action(action, options=options):
+            version = res.body.to_pybytes()
+            _logger.debug(f"The HQS server is of version {version}.")
+            return version
+
     def _should_be_used(self):
         if not self._enabled_on_cluster:
             _logger.debug(
@@ -331,10 +359,32 @@ def _certificates(self):
         cert_key = self._client._cert_key
         return {"kstore": kstore, "tstore": tstore, "cert_key": cert_key}
 
-    def _certificates_header(self):
+    def _make_certificates_json(self):
         if self._certificates_json is None:
             self._certificates_json = json.dumps(self._certificates()).encode("utf-8")
-        return (b"x-certificates-json", self._certificates_json)
+        return self._certificates_json
+
+    def _certificates_headers(self):
+        if self._server_version is None:
+            return []
+        return [(b"x-certificates-json", self._make_certificates_json())]
+
+    @retry(
+        wait_exponential_multiplier=1000,
+        stop_max_attempt_number=3,
+        retry_on_exception=_should_retry_certificate_registration,
+    )
+    def _register_certificates(self):
+        certificates_json = self._make_certificates_json()
+        certificates_json_buf = pyarrow.py_buffer(certificates_json)
+        action = pyarrow.flight.Action(
+            "register-client-certificates", certificates_json_buf
+        )
+        # Registering certificates queue time occasionally spike.
+        options = pyarrow.flight.FlightCallOptions(timeout=self.health_check_timeout)
+        _logger.debug("Registering client certificates with Hopsworks Query Service.")
+        self._connection.do_action(action, options=options)
+        _logger.debug("Client certificates registered.")
 
     def _handle_afs_exception(user_message="None"):
         def decorator(func):
@@ -346,7 +396,13 @@ def afs_error_handler_wrapper(instance, *args, **kw):
                     message = str(e)
                     _logger.debug("Caught exception in %s: %s", func.__name__, message)
                     _logger.exception(e)
-                    if _is_feature_query_service_queue_full_error(e):
+                    if instance._server_version is None and (
+                        isinstance(e, FlightServerError)
+                        and "Please register client certificates first." in message
+                    ):
+                        instance._register_certificates()
+                        return func(instance, *args, **kw)
+                    elif _is_feature_query_service_queue_full_error(e):
                         raise FeatureStoreException(
                             "Hopsworks Query Service is busy right now. Please try again later."
                         ) from e
@@ -384,7 +440,7 @@ def _get_dataset(self, descriptor, timeout=None, dataframe_type="pandas"):
         info = self.get_flight_info(descriptor)
         _logger.debug("Retrieved flight info: %s. Fetching dataset.", str(info))
         options = pyarrow.flight.FlightCallOptions(
-            timeout=timeout, headers=[self._certificates_header()]
+            timeout=timeout, headers=self._certificates_headers()
         )
         reader = self._connection.do_get(info.endpoints[0].ticket, options)
         _logger.debug("Dataset fetched. Converting to dataframe %s.", dataframe_type)
@@ -456,7 +512,7 @@ def create_training_dataset(
                 else self.timeout
             )
             options = pyarrow.flight.FlightCallOptions(
-                timeout=timeout, headers=[self._certificates_header()]
+                timeout=timeout, headers=self._certificates_headers()
             )
             for result in self._connection.do_action(action, options):
                 return result.body.to_pybytes()