SK-2777: update client re-init logic

Author: saileshwar-skyflow

skyflow/error/_skyflow_error.py

@@ -15,5 +15,4 @@ def __init__(self,
         self.http_status = http_status if http_status else SkyflowMessages.HttpStatus.BAD_REQUEST.value
         self.details = details
         self.request_id = request_id
-        log_error(message, http_code, request_id, grpc_code, http_status, details)
         super().__init__()

skyflow/utils/_utils.py

@@ -30,26 +30,18 @@
 invalid_input_error_code = SkyflowMessages.ErrorCodes.INVALID_INPUT.value
 
 def get_credentials(config_level_creds = None, common_skyflow_creds = None, logger = None):
-    dotenv.load_dotenv()
-    dotenv_path = dotenv.find_dotenv(usecwd=True)
-    if dotenv_path:
-        load_dotenv(dotenv_path)
-    env_skyflow_credentials = os.getenv("SKYFLOW_CREDENTIALS")
     if config_level_creds:
         return config_level_creds
     if common_skyflow_creds:
         return common_skyflow_creds
+    dotenv_path = dotenv.find_dotenv(usecwd=True)
+    if dotenv_path:
+        load_dotenv(dotenv_path)
+    env_skyflow_credentials = os.getenv("SKYFLOW_CREDENTIALS")
     if env_skyflow_credentials:
-        env_skyflow_credentials.strip()
-        try:
-            env_creds = env_skyflow_credentials.replace('\n', '\\n')
-            return {
-                'credentials_string': env_creds
-            }
-        except json.JSONDecodeError:
-            raise SkyflowError(SkyflowMessages.Error.INVALID_JSON_FORMAT_IN_CREDENTIALS_ENV.value, invalid_input_error_code)
-    else:
-        raise SkyflowError(SkyflowMessages.Error.INVALID_CREDENTIALS.value, invalid_input_error_code)
+        env_creds = env_skyflow_credentials.strip().replace('\n', '\\n')
+        return {'credentials_string': env_creds}
+    raise SkyflowError(SkyflowMessages.Error.INVALID_CREDENTIALS.value, invalid_input_error_code)
 
 def validate_api_key(api_key: str, logger = None) -> bool:
     if len(api_key) != 42:

skyflow/utils/validations/_validations.py

@@ -122,8 +122,8 @@ def validate_credentials(logger, credentials, config_id_type=None, config_id=Non
         )
         if is_expired(credentials.get("token"), logger):
             raise SkyflowError(
-                SkyflowMessages.Error.INVALID_CREDENTIALS_TOKEN.value.format(config_id_type, config_id)
-                if config_id_type and config_id else SkyflowMessages.Error.INVALID_CREDENTIALS_TOKEN.value,
+                SkyflowMessages.Error.EXPIRED_TOKEN.value
+                if config_id_type and config_id else SkyflowMessages.Error.EXPIRED_TOKEN.value,
                 invalid_input_error_code
             )
     elif "api_key" in credentials:
@@ -389,7 +389,7 @@ def validate_deidentify_file_request(logger, request: DeidentifyFileRequest):
     if hasattr(request, 'wait_time') and request.wait_time is not None:
         if not isinstance(request.wait_time, (int, float)):
             raise SkyflowError(SkyflowMessages.Error.INVALID_WAIT_TIME.value, invalid_input_error_code)
-        if request.wait_time < 0 and request.wait_time > 64:
+        if request.wait_time < 0 or request.wait_time > 64:
             raise SkyflowError(SkyflowMessages.Error.WAIT_TIME_GREATER_THEN_64.value, invalid_input_error_code)
 
 def validate_insert_request(logger, request):

skyflow/vault/client/client.py

@@ -14,6 +14,9 @@ def __init__(self, config):
         self.__logger = None
         self.__is_config_updated = False
         self.__bearer_token = None
+        self.__credentials = None
+        self.__vault_url = None
+        self.__is_static_token = None
 
     def set_common_skyflow_credentials(self, credentials):
         self.__common_skyflow_credentials = credentials
@@ -23,16 +26,29 @@ def set_logger(self, log_level, logger):
         self.__logger = logger
 
     def initialize_client_configuration(self):
-        credentials = get_credentials(self.__config.get("credentials"), self.__common_skyflow_credentials, logger = self.__logger)
-        token = self.get_bearer_token(credentials)
-        vault_url = get_vault_url(self.__config.get("cluster_id"),
-                                  self.__config.get("env"),
-                                  self.__config.get("vault_id"),
-                                  logger = self.__logger)
-        self.initialize_api_client(vault_url, token)
+        if self.__api_client is not None and not self.__is_config_updated:
+            if self.__is_static_token:
+                return
+            if self.__bearer_token is not None and not is_expired(self.__bearer_token):
+                return
+
+        needs_reinit = self.__api_client is None or self.__is_config_updated
+        if needs_reinit:
+            self.__credentials = get_credentials(self.__config.get("credentials"), self.__common_skyflow_credentials, logger=self.__logger)
+            self.__vault_url = get_vault_url(self.__config.get("cluster_id"),
+                                             self.__config.get("env"),
+                                             self.__config.get("vault_id"),
+                                             logger=self.__logger)
+            self.__is_static_token = 'token' in self.__credentials or 'api_key' in self.__credentials
+        token = self.get_bearer_token(self.__credentials)
+        if needs_reinit:
+            self.initialize_api_client(self.__vault_url, token)
 
     def initialize_api_client(self, vault_url, token):
-        self.__api_client =  Skyflow(base_url=vault_url, token=token)
+        self.__api_client = Skyflow(
+            base_url=vault_url,
+            token=lambda: self.__bearer_token if self.__bearer_token else token,
+        )
 
     def get_records_api(self):
         return self.__api_client.records
@@ -63,11 +79,10 @@ def get_bearer_token(self, credentials):
             "ctx": self.__config.get("ctx")
         }
 
-        if self.__bearer_token is None or self.__is_config_updated:
+        if self.__bearer_token is None or self.__is_config_updated or is_expired(self.__bearer_token):
             if 'path' in credentials:
-                path = credentials.get("path")
                 self.__bearer_token, _ = generate_bearer_token(
-                    path,
+                    credentials.get("path"),
                     options,
                     self.__logger
                 )
@@ -83,10 +98,6 @@ def get_bearer_token(self, credentials):
         else:
             log_info(SkyflowMessages.Info.REUSE_BEARER_TOKEN.value, self.__logger)
 
-        if is_expired(self.__bearer_token):
-            self.__is_config_updated = True
-            raise SyntaxError(SkyflowMessages.Error.EXPIRED_TOKEN.value, SkyflowMessages.ErrorCodes.INVALID_INPUT.value)
-
         return self.__bearer_token
 
     def update_config(self, config):