SK-1736: Fix inconsistencies

skyflowapi · Dec 12, 2024 · 63317c6 · 63317c6
1 parent 8291443
commit 63317c6
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 5 deletions.
diff --git a/skyflow/service_account/_utils.py b/skyflow/service_account/_utils.py
@@ -5,7 +5,7 @@
 from skyflow.error import SkyflowError
 from skyflow.generated.rest.models import V1GetAuthTokenRequest
 from skyflow.service_account.client.auth_client import AuthClient
-from skyflow.utils.logger import log_error, log_info, log_error_log, Logger
+from skyflow.utils.logger import log_info, log_error_log
 from skyflow.utils import get_base_url, format_scope, SkyflowMessages
 
 
@@ -115,18 +115,18 @@ def get_signed_jwt(options, client_id, key_id, token_uri, private_key, logger):
 
 def get_signed_tokens(credentials_obj, options):
     try:
-        expiry_time = time.time() + options.get("time_to_live", 60)
+        expiry_time = int(time.time()) + options.get("time_to_live", 60)
         prefix = "signed_token_"
 
         if options and options.get("data_tokens"):
             for token in options["data_tokens"]:
                 claims = {
                     "iss": "sdk",
                     "key": credentials_obj.get("keyID"),
-                    "aud": credentials_obj.get("tokenURI"),
                     "exp": expiry_time,
                     "sub": credentials_obj.get("clientID"),
-                    "tok": token
+                    "tok": token,
+                    "iat": int(time.time()),
                 }
 
                 if "ctx" in options:

diff --git a/skyflow/vault/controller/_vault.py b/skyflow/vault/controller/_vault.py
@@ -2,7 +2,7 @@
     V1DetokenizePayload, V1TokenizeRecordRequest, V1TokenizePayload, QueryServiceExecuteQueryBody, \
     RecordServiceBulkDeleteRecordBody, RecordServiceUpdateRecordBody, RecordServiceBatchOperationBody, V1BatchRecord, \
     BatchRecordMethod
-from skyflow.generated.rest.exceptions import BadRequestException, UnauthorizedException
+from skyflow.generated.rest.exceptions import BadRequestException, UnauthorizedException, ForbiddenException
 from skyflow.utils import SkyflowMessages, parse_insert_response, \
     handle_exception, parse_update_record_response, parse_delete_response, parse_detokenize_response, \
     parse_tokenize_response, parse_query_response, parse_get_response, encode_column_values
@@ -105,6 +105,8 @@ def insert(self, request: InsertRequest):
             handle_exception(e, self.__vault_client.get_logger())
         except UnauthorizedException as e:
             handle_exception(e, self.__vault_client.get_logger())
+        except ForbiddenException as e:
+            handle_exception(e, self.__vault_client.get_logger())
 
     def update(self, request: UpdateRequest):
         log_info(SkyflowMessages.Info.VALIDATE_UPDATE_REQUEST.value, self.__vault_client.get_logger())
@@ -132,6 +134,8 @@ def update(self, request: UpdateRequest):
             handle_exception(e, self.__vault_client.get_logger())
         except UnauthorizedException as e:
             handle_exception(e, self.__vault_client.get_logger())
+        except ForbiddenException as e:
+            handle_exception(e, self.__vault_client.get_logger())
 
     def delete(self, request: DeleteRequest):
         log_info(SkyflowMessages.Info.VALIDATING_DELETE_REQUEST.value, self.__vault_client.get_logger())
@@ -157,6 +161,8 @@ def delete(self, request: DeleteRequest):
             log_error_log(SkyflowMessages.ErrorLogs.DELETE_REQUEST_REJECTED.value,
                           logger=self.__vault_client.get_logger())
             handle_exception(e, self.__vault_client.get_logger())
+        except ForbiddenException as e:
+            handle_exception(e, self.__vault_client.get_logger())
 
     def get(self, request: GetRequest):
         log_info(SkyflowMessages.Info.VALIDATE_GET_REQUEST.value, self.__vault_client.get_logger())
@@ -190,6 +196,8 @@ def get(self, request: GetRequest):
         except UnauthorizedException as e:
             log_error_log(SkyflowMessages.ErrorLogs.GET_REQUEST_REJECTED.value, self.__vault_client.get_logger())
             handle_exception(e, self.__vault_client.get_logger())
+        except ForbiddenException as e:
+            handle_exception(e, self.__vault_client.get_logger())
 
     def query(self, request: QueryRequest):
         log_info(SkyflowMessages.Info.VALIDATING_QUERY_REQUEST.value, self.__vault_client.get_logger())
@@ -213,6 +221,8 @@ def query(self, request: QueryRequest):
         except UnauthorizedException as e:
             log_error_log(SkyflowMessages.ErrorLogs.QUERY_REQUEST_REJECTED.value, self.__vault_client.get_logger())
             handle_exception(e, self.__vault_client.get_logger())
+        except ForbiddenException as e:
+            handle_exception(e, self.__vault_client.get_logger())
 
     def detokenize(self, request: DetokenizeRequest):
         log_info(SkyflowMessages.Info.VALIDATE_DETOKENIZE_REQUEST.value, self.__vault_client.get_logger())
@@ -241,6 +251,8 @@ def detokenize(self, request: DetokenizeRequest):
             log_error_log(SkyflowMessages.ErrorLogs.DETOKENIZE_REQUEST_REJECTED.value,
                           logger=self.__vault_client.get_logger())
             handle_exception(e, self.__vault_client.get_logger())
+        except ForbiddenException as e:
+            handle_exception(e, self.__vault_client.get_logger())
 
     def tokenize(self, request: TokenizeRequest):
         log_info(SkyflowMessages.Info.VALIDATING_TOKENIZE_REQUEST.value, self.__vault_client.get_logger())
@@ -270,3 +282,5 @@ def tokenize(self, request: TokenizeRequest):
             log_error_log(SkyflowMessages.ErrorLogs.TOKENIZE_REQUEST_REJECTED.value,
                           logger=self.__vault_client.get_logger())
             handle_exception(e, self.__vault_client.get_logger())
+        except ForbiddenException as e:
+            handle_exception(e, self.__vault_client.get_logger())