Skip to content

Commit 63317c6

Browse files
SK-1736: Fix inconsistencies
1 parent 8291443 commit 63317c6

File tree

2 files changed

+19
-5
lines changed

2 files changed

+19
-5
lines changed

skyflow/service_account/_utils.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
from skyflow.error import SkyflowError
66
from skyflow.generated.rest.models import V1GetAuthTokenRequest
77
from skyflow.service_account.client.auth_client import AuthClient
8-
from skyflow.utils.logger import log_error, log_info, log_error_log, Logger
8+
from skyflow.utils.logger import log_info, log_error_log
99
from skyflow.utils import get_base_url, format_scope, SkyflowMessages
1010

1111

@@ -115,18 +115,18 @@ def get_signed_jwt(options, client_id, key_id, token_uri, private_key, logger):
115115

116116
def get_signed_tokens(credentials_obj, options):
117117
try:
118-
expiry_time = time.time() + options.get("time_to_live", 60)
118+
expiry_time = int(time.time()) + options.get("time_to_live", 60)
119119
prefix = "signed_token_"
120120

121121
if options and options.get("data_tokens"):
122122
for token in options["data_tokens"]:
123123
claims = {
124124
"iss": "sdk",
125125
"key": credentials_obj.get("keyID"),
126-
"aud": credentials_obj.get("tokenURI"),
127126
"exp": expiry_time,
128127
"sub": credentials_obj.get("clientID"),
129-
"tok": token
128+
"tok": token,
129+
"iat": int(time.time()),
130130
}
131131

132132
if "ctx" in options:

skyflow/vault/controller/_vault.py

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
V1DetokenizePayload, V1TokenizeRecordRequest, V1TokenizePayload, QueryServiceExecuteQueryBody, \
33
RecordServiceBulkDeleteRecordBody, RecordServiceUpdateRecordBody, RecordServiceBatchOperationBody, V1BatchRecord, \
44
BatchRecordMethod
5-
from skyflow.generated.rest.exceptions import BadRequestException, UnauthorizedException
5+
from skyflow.generated.rest.exceptions import BadRequestException, UnauthorizedException, ForbiddenException
66
from skyflow.utils import SkyflowMessages, parse_insert_response, \
77
handle_exception, parse_update_record_response, parse_delete_response, parse_detokenize_response, \
88
parse_tokenize_response, parse_query_response, parse_get_response, encode_column_values
@@ -105,6 +105,8 @@ def insert(self, request: InsertRequest):
105105
handle_exception(e, self.__vault_client.get_logger())
106106
except UnauthorizedException as e:
107107
handle_exception(e, self.__vault_client.get_logger())
108+
except ForbiddenException as e:
109+
handle_exception(e, self.__vault_client.get_logger())
108110

109111
def update(self, request: UpdateRequest):
110112
log_info(SkyflowMessages.Info.VALIDATE_UPDATE_REQUEST.value, self.__vault_client.get_logger())
@@ -132,6 +134,8 @@ def update(self, request: UpdateRequest):
132134
handle_exception(e, self.__vault_client.get_logger())
133135
except UnauthorizedException as e:
134136
handle_exception(e, self.__vault_client.get_logger())
137+
except ForbiddenException as e:
138+
handle_exception(e, self.__vault_client.get_logger())
135139

136140
def delete(self, request: DeleteRequest):
137141
log_info(SkyflowMessages.Info.VALIDATING_DELETE_REQUEST.value, self.__vault_client.get_logger())
@@ -157,6 +161,8 @@ def delete(self, request: DeleteRequest):
157161
log_error_log(SkyflowMessages.ErrorLogs.DELETE_REQUEST_REJECTED.value,
158162
logger=self.__vault_client.get_logger())
159163
handle_exception(e, self.__vault_client.get_logger())
164+
except ForbiddenException as e:
165+
handle_exception(e, self.__vault_client.get_logger())
160166

161167
def get(self, request: GetRequest):
162168
log_info(SkyflowMessages.Info.VALIDATE_GET_REQUEST.value, self.__vault_client.get_logger())
@@ -190,6 +196,8 @@ def get(self, request: GetRequest):
190196
except UnauthorizedException as e:
191197
log_error_log(SkyflowMessages.ErrorLogs.GET_REQUEST_REJECTED.value, self.__vault_client.get_logger())
192198
handle_exception(e, self.__vault_client.get_logger())
199+
except ForbiddenException as e:
200+
handle_exception(e, self.__vault_client.get_logger())
193201

194202
def query(self, request: QueryRequest):
195203
log_info(SkyflowMessages.Info.VALIDATING_QUERY_REQUEST.value, self.__vault_client.get_logger())
@@ -213,6 +221,8 @@ def query(self, request: QueryRequest):
213221
except UnauthorizedException as e:
214222
log_error_log(SkyflowMessages.ErrorLogs.QUERY_REQUEST_REJECTED.value, self.__vault_client.get_logger())
215223
handle_exception(e, self.__vault_client.get_logger())
224+
except ForbiddenException as e:
225+
handle_exception(e, self.__vault_client.get_logger())
216226

217227
def detokenize(self, request: DetokenizeRequest):
218228
log_info(SkyflowMessages.Info.VALIDATE_DETOKENIZE_REQUEST.value, self.__vault_client.get_logger())
@@ -241,6 +251,8 @@ def detokenize(self, request: DetokenizeRequest):
241251
log_error_log(SkyflowMessages.ErrorLogs.DETOKENIZE_REQUEST_REJECTED.value,
242252
logger=self.__vault_client.get_logger())
243253
handle_exception(e, self.__vault_client.get_logger())
254+
except ForbiddenException as e:
255+
handle_exception(e, self.__vault_client.get_logger())
244256

245257
def tokenize(self, request: TokenizeRequest):
246258
log_info(SkyflowMessages.Info.VALIDATING_TOKENIZE_REQUEST.value, self.__vault_client.get_logger())
@@ -270,3 +282,5 @@ def tokenize(self, request: TokenizeRequest):
270282
log_error_log(SkyflowMessages.ErrorLogs.TOKENIZE_REQUEST_REJECTED.value,
271283
logger=self.__vault_client.get_logger())
272284
handle_exception(e, self.__vault_client.get_logger())
285+
except ForbiddenException as e:
286+
handle_exception(e, self.__vault_client.get_logger())

0 commit comments

Comments
 (0)