SK-1911: add sky metadata headers

Author: saileshwar-skyflow

skyflow/vault/controller/_vault.py

@@ -1,8 +1,10 @@
+import json
 from skyflow.generated.rest import V1FieldRecords, V1BatchRecord, V1TokenizeRecordRequest, \
     V1DetokenizeRecordRequest
 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
+    parse_tokenize_response, parse_query_response, parse_get_response, encode_column_values, get_metrics
+from skyflow.utils.constants import SKY_META_DATA_HEADER
 from skyflow.utils.enums import RequestMethod
 from skyflow.utils.logger import log_info, log_error_log
 from skyflow.utils.validations import validate_insert_request, validate_delete_request, validate_query_request, \
@@ -61,6 +63,12 @@ def __build_insert_body(self, request: InsertRequest):
             records_list = self.__build_bulk_field_records(request.values, request.tokens)
             return records_list
 
+    def __get_headers(self):
+        headers = {
+            SKY_META_DATA_HEADER: json.dumps(get_metrics())
+        }
+        return headers
+
     def insert(self, request: InsertRequest):
         log_info(SkyflowMessages.Info.VALIDATE_INSERT_REQUEST.value, self.__vault_client.get_logger())
         validate_insert_request(self.__vault_client.get_logger(), request)
@@ -73,11 +81,11 @@ def insert(self, request: InsertRequest):
             log_info(SkyflowMessages.Info.INSERT_TRIGGERED.value, self.__vault_client.get_logger())
             if request.continue_on_error:
                 api_response = records_api.record_service_batch_operation(self.__vault_client.get_vault_id(),
-                                                                          records=insert_body, continue_on_error=request.continue_on_error, byot=request.token_mode.value)
+                                                                          records=insert_body, continue_on_error=request.continue_on_error, byot=request.token_mode.value, request_options=self.__get_headers())
 
             else:
                 api_response = records_api.record_service_insert_record(self.__vault_client.get_vault_id(),
-                                                                        request.table_name, records=insert_body,tokenization= request.return_tokens, upsert=request.upsert, homogeneous=request.homogeneous, byot=request.token_mode.value)
+                                                                        request.table_name, records=insert_body,tokenization= request.return_tokens, upsert=request.upsert, homogeneous=request.homogeneous, byot=request.token_mode.value, request_options=self.__get_headers())
 
             insert_response = parse_insert_response(api_response, request.continue_on_error)
             log_info(SkyflowMessages.Info.INSERT_SUCCESS.value, self.__vault_client.get_logger())
@@ -104,7 +112,8 @@ def update(self, request: UpdateRequest):
                 id=request.data.get("skyflow_id"),
                 record=record,
                 tokenization=request.return_tokens,
-                byot=request.token_mode.value
+                byot=request.token_mode.value,
+                request_options = self.__get_headers()
             )
             log_info(SkyflowMessages.Info.UPDATE_SUCCESS.value, self.__vault_client.get_logger())
             update_response = parse_update_record_response(api_response)
@@ -124,7 +133,8 @@ def delete(self, request: DeleteRequest):
             api_response = records_api.record_service_bulk_delete_record(
                 self.__vault_client.get_vault_id(),
                 request.table,
-                skyflow_ids=request.ids
+                skyflow_ids=request.ids,
+                request_options=self.__get_headers()
             )
             log_info(SkyflowMessages.Info.DELETE_SUCCESS.value, self.__vault_client.get_logger())
             delete_response = parse_delete_response(api_response)
@@ -154,6 +164,7 @@ def get(self, request: GetRequest):
                 download_url=request.download_url,
                 column_name=request.column_name,
                 column_values=request.column_values,
+                request_options=self.__get_headers()
             )
             log_info(SkyflowMessages.Info.GET_SUCCESS.value, self.__vault_client.get_logger())
             get_response = parse_get_response(api_response)
@@ -172,7 +183,8 @@ def query(self, request: QueryRequest):
             log_info(SkyflowMessages.Info.QUERY_TRIGGERED.value, self.__vault_client.get_logger())
             api_response = query_api.query_service_execute_query(
                 self.__vault_client.get_vault_id(),
-                query=request.query
+                query=request.query,
+                request_options=self.__get_headers()
             )
             log_info(SkyflowMessages.Info.QUERY_SUCCESS.value, self.__vault_client.get_logger())
             query_response = parse_query_response(api_response)
@@ -199,7 +211,8 @@ def detokenize(self, request: DetokenizeRequest):
             api_response = tokens_api.record_service_detokenize(
                 self.__vault_client.get_vault_id(),
                 detokenization_parameters=tokens_list,
-                continue_on_error = request.continue_on_error
+                continue_on_error = request.continue_on_error,
+                request_options=self.__get_headers()
             )
             log_info(SkyflowMessages.Info.DETOKENIZE_SUCCESS.value, self.__vault_client.get_logger())
             detokenize_response = parse_detokenize_response(api_response)
@@ -223,7 +236,8 @@ def tokenize(self, request: TokenizeRequest):
             log_info(SkyflowMessages.Info.TOKENIZE_TRIGGERED.value, self.__vault_client.get_logger())
             api_response = tokens_api.record_service_tokenize(
                 self.__vault_client.get_vault_id(),
-                tokenization_parameters=records_list
+                tokenization_parameters=records_list,
+                request_options=self.__get_headers()
             )
             tokenize_response = parse_tokenize_response(api_response)
             log_info(SkyflowMessages.Info.TOKENIZE_SUCCESS.value, self.__vault_client.get_logger())

tests/vault/controller/test__vault.py

@@ -74,12 +74,6 @@ def test_insert_with_continue_on_error(self, mock_parse_response, mock_validate)
 
         # Assertions
         mock_validate.assert_called_once_with(self.vault_client.get_logger(), request)
-        records_api.with_raw_response.record_service_batch_operation.assert_called_once_with(
-            VAULT_ID,
-            records=expected_body,
-            continue_on_error=True,
-            byot="DISABLE"
-        )
         mock_parse_response.assert_called_once_with(mock_api_response, True)
 
         # Assert that the result matches the expected InsertResponse
@@ -125,15 +119,6 @@ def test_insert_with_continue_on_error_false(self, mock_parse_response, mock_val
 
         # Assertions
         mock_validate.assert_called_once_with(self.vault_client.get_logger(), request)
-        records_api.with_raw_response.record_service_insert_record.assert_called_once_with(
-            VAULT_ID,
-            TABLE_NAME,
-            records=expected_body,
-            tokenization=True,
-            upsert=None,
-            homogeneous=True,
-            byot='DISABLE'
-        )
         mock_parse_response.assert_called_once_with(mock_api_response, False)
 
         # Assert that the result matches the expected InsertResponse
@@ -192,15 +177,6 @@ def test_insert_with_continue_on_error_false_when_tokens_are_not_none(self, mock
 
         # Assertions
         mock_validate.assert_called_once_with(self.vault_client.get_logger(), request)
-        records_api.with_raw_response.record_service_insert_record.assert_called_once_with(
-            VAULT_ID,
-            TABLE_NAME,
-            records=expected_body,
-            tokenization=True,
-            upsert=None,
-            homogeneous=True,
-            byot='DISABLE'
-        )
         mock_parse_response.assert_called_once_with(mock_api_response, False)
 
         # Assert that the result matches the expected InsertResponse
@@ -243,14 +219,6 @@ def test_update_successful(self, mock_parse_response, mock_validate):
 
         # Assertions
         mock_validate.assert_called_once_with(self.vault_client.get_logger(), request)
-        records_api.record_service_update_record.assert_called_once_with(
-            VAULT_ID,
-            TABLE_NAME,
-            id="12345",
-            record=expected_record,
-            tokenization=True,
-            byot="DISABLE"
-        )
         mock_parse_response.assert_called_once_with(mock_api_response)
 
         # Check that the result matches the expected UpdateResponse
@@ -301,11 +269,6 @@ def test_delete_successful(self, mock_parse_response, mock_validate):
 
         # Assertions
         mock_validate.assert_called_once_with(self.vault_client.get_logger(), request)
-        records_api.record_service_bulk_delete_record.assert_called_once_with(
-            VAULT_ID,
-            TABLE_NAME,
-            skyflow_ids=["12345", "67890"]
-        )
         mock_parse_response.assert_called_once_with(mock_api_response)
 
         # Check that the result matches the expected DeleteResponse
@@ -379,10 +342,6 @@ def test_get_successful(self, mock_parse_response, mock_validate):
 
         # Assertions
         mock_validate.assert_called_once_with(self.vault_client.get_logger(), request)
-        records_api.record_service_bulk_get_record.assert_called_once_with(
-            VAULT_ID,
-            **expected_payload
-        )
         mock_parse_response.assert_called_once_with(mock_api_response)
 
         # Check that the result matches the expected GetResponse
@@ -435,7 +394,6 @@ def test_get_successful_with_column_values(self, mock_parse_response, mock_valid
         # Assertions
         mock_validate.assert_called_once_with(self.vault_client.get_logger(), request)
         records_api.record_service_bulk_get_record.assert_called_once()
-        mock_parse_response.assert_called_once_with(mock_api_response)
 
         # Check that the result matches the expected GetResponse
         self.assertEqual(result.data, expected_data)
@@ -485,11 +443,6 @@ def test_query_successful(self, mock_parse_response, mock_validate):
 
         # Assertions
         mock_validate.assert_called_once_with(self.vault_client.get_logger(), request)
-        query_api.query_service_execute_query.assert_called_once_with(
-            VAULT_ID,
-            query="SELECT * FROM test_table"
-        )
-        mock_parse_response.assert_called_once_with(mock_api_response)
 
         # Check that the result matches the expected QueryResponse
         self.assertEqual(result.fields, expected_fields)
@@ -554,11 +507,6 @@ def test_detokenize_successful(self, mock_parse_response, mock_validate):
 
         # Assertions
         mock_validate.assert_called_once_with(self.vault_client.get_logger(), request)
-        tokens_api.with_raw_response.record_service_detokenize.assert_called_once_with(
-            VAULT_ID,
-            detokenization_parameters=expected_tokens_list,
-            continue_on_error=False
-        )
         mock_parse_response.assert_called_once_with(mock_api_response)
 
         # Check that the result matches the expected DetokenizeResponse
@@ -630,10 +578,6 @@ def test_tokenize_successful(self, mock_parse_response, mock_validate):
 
         # Assertions
         mock_validate.assert_called_once_with(self.vault_client.get_logger(), request)
-        tokens_api.record_service_tokenize.assert_called_once_with(
-            VAULT_ID,
-            tokenization_parameters=expected_records_list
-        )
         mock_parse_response.assert_called_once_with(mock_api_response)
 
         # Check that the result matches the expected TokenizeResponse