[FSTORE-1661] Fix Arrow Flight certificate registration (#454)

Author: aversey

python/hsfs/core/arrow_flight_client.py

@@ -39,7 +39,6 @@
 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
 
 
@@ -67,7 +66,7 @@ def close() -> None:
 
 def _disable_feature_query_service_client():
     global _arrow_flight_instance
-    _logger.debug("Disabling Hopsworks Feature Query Service Client.")
+    _logger.debug("Disabling Hopsworks Query Service Client.")
     if _arrow_flight_instance is None:
         _arrow_flight_instance.ArrowFlightClient(disabled_for_session=True)
     else:
@@ -90,12 +89,10 @@ def _is_no_commits_found_error(exception):
     ) and "No commits found" in str(exception)
 
 
-def _should_retry_healthcheck_or_certificate_registration(exception):
+def _should_retry_healthcheck(exception):
     return (
         isinstance(exception, pyarrow._flight.FlightUnavailableError)
         or isinstance(exception, pyarrow._flight.FlightTimedOutError)
-        # not applicable for healthcheck, only certificate registration
-        or _is_feature_query_service_queue_full_error(exception)
     )
 
 
@@ -151,8 +148,8 @@ class ArrowFlightClient:
         StorageConnector.SNOWFLAKE,
         StorageConnector.BIGQUERY,
     ]
-    READ_ERROR = "Could not read data using Hopsworks Feature Query Service."
-    WRITE_ERROR = 'Could not write data using Hopsworks Feature Query Service. If the issue persists, use write_options={"use_spark": True} instead.'
+    READ_ERROR = "Could not read data using Hopsworks Query Service."
+    WRITE_ERROR = 'Could not write data using Hopsworks Query Service. If the issue persists, use write_options={"use_spark": True} instead.'
     DEFAULTING_TO_DIFFERENT_SERVICE_WARNING = (
         "Defaulting to Spark execution for this call."
     )
@@ -162,7 +159,7 @@ class ArrowFlightClient:
     DEFAULT_GRPC_MIN_RECONNECT_BACKOFF_MS = 2000
 
     def __init__(self, disabled_for_session: bool = False):
-        _logger.debug("Initializing Hopsworks Feature Query Service Client.")
+        _logger.debug("Initializing Hopsworks Query Service Client.")
         self._timeout: float = ArrowFlightClient.DEFAULT_TIMEOUT_SECONDS
         self._health_check_timeout: float = (
             ArrowFlightClient.DEFAULT_HEALTHCHECK_TIMEOUT_SECONDS
@@ -179,36 +176,34 @@ def __init__(self, disabled_for_session: bool = False):
 
         self._client = client.get_instance()
         self._variable_api: VariableApi = VariableApi()
+        self._certificates_json: Optional[str] = None
 
         try:
             self._check_cluster_service_enabled()
             self._host_url = self._retrieve_host_url()
 
             if self._enabled_on_cluster:
-                _logger.debug(
-                    "Hopsworks Feature Query Service is enabled on the cluster."
-                )
+                _logger.debug("Hopsworks Query Service is enabled on the cluster.")
                 self._initialize_flight_client()
             else:
                 _logger.debug(
-                    "Hopsworks Feature Query Service Client is not enabled on the cluster or a cluster variable is misconfigured."
+                    "Hopsworks Query Service Client is not enabled on the cluster or a cluster variable is misconfigured."
                 )
                 self._disable_for_session()
                 return
         except Exception as e:
-            _logger.debug("Failed to connect to Hopsworks Feature Query Service")
+            _logger.debug("Failed to connect to Hopsworks Query Service")
             _logger.exception(e)
             self._disable_for_session(str(e))
             return
 
         try:
             self._health_check()
-            self._register_certificates()
         except Exception as e:
-            _logger.debug("Failed to connect to Hopsworks Feature Query Service.")
+            _logger.debug("Failed to connect to Hopsworks Query Service.")
             _logger.exception(e)
             warnings.warn(
-                f"Failed to connect to Hopsworks Feature Query Service, got {str(e)}."
+                f"Failed to connect to Hopsworks Query Service, got {str(e)}."
                 + ArrowFlightClient.DEFAULTING_TO_DIFFERENT_SERVICE_WARNING
                 + ArrowFlightClient.CLIENT_WILL_STAY_ACTIVE_WARNING,
                 stacklevel=1,
@@ -218,13 +213,13 @@ def __init__(self, disabled_for_session: bool = False):
     def _check_cluster_service_enabled(self) -> None:
         try:
             _logger.debug(
-                "Connecting to Hopsworks Cluster to check if Feature Query Service is enabled."
+                "Connecting to Hopsworks Cluster to check if Hopsworks Query Service is enabled."
             )
             self._enabled_on_cluster = self._variable_api.get_flyingduck_enabled()
         except Exception as e:
             # if feature flag cannot be retrieved, assume it is disabled
             _logger.debug(
-                "Unable to fetch Hopsworks Feature Query Service (HQFS) flag, disabling HFQS client."
+                "Unable to fetch Hopsworks Query Service flag, disabling its client."
             )
             _logger.exception(e)
             self._enabled_on_cluster = False
@@ -240,13 +235,11 @@ def _retrieve_host_url(self) -> Optional[str]:
             service_discovery_domain = self._variable_api.get_service_discovery_domain()
             if service_discovery_domain == "":
                 raise FeatureStoreException(
-                    "Client could not get Feature Query Service hostname from service_discovery_domain. "
+                    "Client could not get Hopsworks Query Service hostname from service_discovery_domain. "
                     "The variable is either not set or empty in Hopsworks cluster configuration."
                 )
             host_url = f"grpc+tls://flyingduck.service.{service_discovery_domain}:5005"
-        _logger.debug(
-            f"Connecting to Hopsworks Feature Query Service on host {host_url}"
-        )
+        _logger.debug(f"Connecting to Hopsworks Query Service on host {host_url}")
         return host_url
 
     def _disable_for_session(
@@ -255,17 +248,17 @@ def _disable_for_session(
         self._disabled_for_session = True
         if on_purpose:
             warnings.warn(
-                "Hopsworks Feature Query Service will be disabled for this session.",
+                "Hopsworks Query Service will be disabled for this session.",
                 stacklevel=1,
             )
         if self._enabled_on_cluster:
             warnings.warn(
-                "Hospworks Feature Query Service is disabled on cluster. Contact your administrator to enable it.",
+                "Hospworks Query Service is disabled on cluster. Contact your administrator to enable it.",
                 stacklevel=1,
             )
         else:
             warnings.warn(
-                f"Client initialisation failed: {message}. Hopsworks Feature Query Service will be disabled for this session."
+                f"Client initialisation failed: {message}. Hopsworks Query Service will be disabled for this session."
                 "If you believe this is a transient error, you can call `(hopsworks.)hsfs.reset_offline_query_service_client()`"
                 " to re-enable it.",
                 stacklevel=1,
@@ -291,10 +284,10 @@ def _initialize_flight_client(self):
     @retry(
         wait_exponential_multiplier=1000,
         stop_max_attempt_number=5,
-        retry_on_exception=_should_retry_healthcheck_or_certificate_registration,
+        retry_on_exception=_should_retry_healthcheck,
     )
     def _health_check(self):
-        _logger.debug("Performing healthcheck of Hopsworks Feature Query Service.")
+        _logger.debug("Performing healthcheck of Hopsworks Query Service.")
         action = pyarrow.flight.Action("healthcheck", b"")
         options = pyarrow.flight.FlightCallOptions(timeout=self.health_check_timeout)
         list(self._connection.do_action(action, options=options))
@@ -303,17 +296,17 @@ def _health_check(self):
     def _should_be_used(self):
         if not self._enabled_on_cluster:
             _logger.debug(
-                "Hopsworks Feature Query Service not used as it is disabled on the cluster."
+                "Hopsworks Query Service not used as it is disabled on the cluster."
             )
             return False
 
         if self._disabled_for_session:
             _logger.debug(
-                "Hopsworks Feature Query Service client failed to initialise and is disabled for the session."
+                "Hopsworks Query Service client failed to initialise and is disabled for the session."
             )
             return False
 
-        _logger.debug("Using Hopsworks Feature Query Service.")
+        _logger.debug("Using Hopsworks Query Service.")
         return True
 
     def _extract_certs(self):
@@ -332,29 +325,16 @@ def _encode_certs(self, path):
             content = f.read()
             return base64.b64encode(content).decode("utf-8")
 
-    @retry(
-        wait_exponential_multiplier=1000,
-        stop_max_attempt_number=3,
-        retry_on_exception=_should_retry_healthcheck_or_certificate_registration,
-    )
-    def _register_certificates(self):
+    def _certificates(self):
         kstore = self._encode_certs(self._client._get_jks_key_store_path())
         tstore = self._encode_certs(self._client._get_jks_trust_store_path())
         cert_key = self._client._cert_key
-        certificates_json = json.dumps(
-            {"kstore": kstore, "tstore": tstore, "cert_key": cert_key}
-        ).encode("ascii")
-        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 Feature Query Service."
-        )
-        self._connection.do_action(action, options=options)
-        _logger.debug("Client certificates registered.")
+        return {"kstore": kstore, "tstore": tstore, "cert_key": cert_key}
+
+    def _certificates_header(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)
 
     def _handle_afs_exception(user_message="None"):
         def decorator(func):
@@ -366,15 +346,9 @@ 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 (
-                        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):
+                    if _is_feature_query_service_queue_full_error(e):
                         raise FeatureStoreException(
-                            "Hopsworks Feature Query Service is busy right now. Please try again later."
+                            "Hopsworks Query Service is busy right now. Please try again later."
                         ) from e
                     elif _is_no_commits_found_error(e):
                         raise FeatureStoreException(str(e).split("Details:")[0]) from e
@@ -409,7 +383,9 @@ def _get_dataset(self, descriptor, timeout=None, dataframe_type="pandas"):
             timeout = self.timeout
         info = self.get_flight_info(descriptor)
         _logger.debug("Retrieved flight info: %s. Fetching dataset.", str(info))
-        options = pyarrow.flight.FlightCallOptions(timeout=timeout)
+        options = pyarrow.flight.FlightCallOptions(
+            timeout=timeout, headers=[self._certificates_header()]
+        )
         reader = self._connection.do_get(info.endpoints[0].ticket, options)
         _logger.debug("Dataset fetched. Converting to dataframe %s.", dataframe_type)
         if dataframe_type.lower() == "polars":
@@ -479,7 +455,9 @@ def create_training_dataset(
                 if arrow_flight_config
                 else self.timeout
             )
-            options = pyarrow.flight.FlightCallOptions(timeout=timeout)
+            options = pyarrow.flight.FlightCallOptions(
+                timeout=timeout, headers=[self._certificates_header()]
+            )
             for result in self._connection.do_action(action, options):
                 return result.body.to_pybytes()
         except pyarrow.lib.ArrowIOError as e:
@@ -516,7 +494,7 @@ def is_enabled(self):
 
     @property
     def timeout(self) -> Union[int, float]:
-        """Timeout in seconds for Hopsworks Feature Query Service do_get or do_action operations, not including the healthcheck."""
+        """Timeout in seconds for Hopsworks Query Service do_get or do_action operations, not including the healthcheck."""
         return self._timeout
 
     @timeout.setter
@@ -534,7 +512,7 @@ def health_check_timeout(self, value: Union[int, float]) -> None:
 
     @property
     def host_url(self) -> Optional[str]:
-        """URL of Hopsworks Feature Query Service."""
+        """URL of Hopsworks Query Service."""
         return self._host_url
 
     @host_url.setter