SK-1731: Added service account samples and vault api samples

skyflowapi · Nov 12, 2024 · 1e40d2f · 1e40d2f
1 parent c73aa65
commit 1e40d2f
Show file tree

Hide file tree

Showing 19 changed files with 107 additions and 5 deletions.
diff --git a/samples/service_account/scoped_token_generation_example.py b/samples/service_account/scoped_token_generation_example.py
@@ -0,0 +1,27 @@
+from skyflow.service_account import generate_bearer_token, generate_bearer_token_from_creds, is_expired
+
+file_path = 'CREDENTIALS_FILE_PATH'
+bearer_token = ''
+
+skyflow_credentials_string = '{"clientID":"<YOUR_CLIENT_ID>","clientName":"<YOUR_CLIENT_NAME>","tokenURI":"<YOUR_TOKEN_URI>","keyID":"<YOUR_KEY_ID>","privateKey":"<YOUR_PRIVATE_KEY>"}'
+
+
+# Generate bearer token from credentials file path
+options = {
+    "ctx": "<CONTEXT_ID>"
+}
+
+if is_expired(bearer_token):
+    bearer_token, token_type = generate_bearer_token(
+        "<YOUR_CREDENTIALS_FILE_PATH>", options
+    )
+
+    print(bearer_token, token_type)
+
+
+# Generate bearer token from credentials string
+if is_expired(bearer_token):
+    bearer_token, token_type = generate_bearer_token_from_creds(skyflow_credentials_string, options)
+
+    print(bearer_token, token_type)
+
diff --git a/samples/service_account/signed_token_generation_example.py b/samples/service_account/signed_token_generation_example.py
@@ -0,0 +1,27 @@
+from skyflow.service_account import is_expired, generate_signed_data_tokens, generate_signed_data_tokens_from_creds
+
+file_path = 'CREDENTIALS_FILE_PATH'
+bearer_token = ''
+
+skyflow_credentials_string = '{"clientID":"<YOUR_CLIENT_ID>","clientName":"<YOUR_CLIENT_NAME>","tokenURI":"<YOUR_TOKEN_URI>","keyID":"<YOUR_KEY_ID>","privateKey":"<YOUR_PRIVATE_KEY>"}'
+
+
+
+options = {
+    "ctx": "CONTEX_ID",
+    "data_tokens": ["DATA_TOKEN1", "DATA_TOKEN2"],
+    "time_to_live": 90 # in seconds
+}
+
+# Generate bearer token from credentials file path
+if is_expired(bearer_token):
+    actual_token, signed_token = generate_signed_data_tokens(
+        "<YOUR_CREDENTIALS_FILE_PATH>", options
+    )
+
+
+# Generate bearer token from credentials string
+if is_expired(bearer_token):
+    actual_token, signed_token = generate_signed_data_tokens_from_creds(
+        "<CREDENTIALS_STRING>", options
+    )
diff --git a/samples/service_account/token_generation_example.py b/samples/service_account/token_generation_example.py
@@ -0,0 +1,23 @@
+from skyflow.service_account import generate_bearer_token, generate_bearer_token_from_creds, is_expired
+
+file_path = 'CREDENTIALS_FILE_PATH'
+bearer_token = ''
+
+skyflow_credentials_string = '{"clientID":"<YOUR_CLIENT_ID>","clientName":"<YOUR_CLIENT_NAME>","tokenURI":"<YOUR_TOKEN_URI>","keyID":"<YOUR_KEY_ID>","privateKey":"<YOUR_PRIVATE_KEY>"}'
+
+
+# Generate bearer token from credentials file path
+if is_expired(bearer_token):
+    bearer_token, token_type = generate_bearer_token(
+        "<YOUR_CREDENTIALS_FILE_PATH>"
+    )
+
+    print(bearer_token, token_type)
+
+
+# Generate bearer token from credentials string
+if is_expired(bearer_token):
+    bearer_token, token_type = generate_bearer_token_from_creds(skyflow_credentials_string)
+
+    print(bearer_token, token_type)
+
diff --git a/samples/service_account/token_generation_with_context_example.py b/samples/service_account/token_generation_with_context_example.py
@@ -0,0 +1,26 @@
+from skyflow.service_account import generate_bearer_token, generate_bearer_token_from_creds, is_expired
+
+file_path = 'CREDENTIALS_FILE_PATH'
+bearer_token = ''
+
+skyflow_credentials_string = '{"clientID":"<YOUR_CLIENT_ID>","clientName":"<YOUR_CLIENT_NAME>","tokenURI":"<YOUR_TOKEN_URI>","keyID":"<YOUR_KEY_ID>","privateKey":"<YOUR_PRIVATE_KEY>"}'
+
+
+# Generate bearer token from credentials file path
+options = {
+    "role_ids": ["ROLE_ID1", "ROLE_ID2"]
+}
+if is_expired(bearer_token):
+    bearer_token, token_type = generate_bearer_token(
+        "<YOUR_CREDENTIALS_FILE_PATH>", options
+    )
+
+    print(bearer_token, token_type)
+
+
+# Generate bearer token from credentials string
+if is_expired(bearer_token):
+    bearer_token, token_type = generate_bearer_token_from_creds(skyflow_credentials_string, options)
+
+    print(bearer_token, token_type)
+
diff --git a/samples/client_operations.py → samples/vault_api/client_operations.py b/samples/client_operations.py → samples/vault_api/client_operations.py
diff --git a/samples/credentials_options.py → samples/vault_api/credentials_options.py b/samples/credentials_options.py → samples/vault_api/credentials_options.py
diff --git a/samples/delete_records.py → samples/vault_api/delete_records.py b/samples/delete_records.py → samples/vault_api/delete_records.py
diff --git a/samples/detokenize_records.py → samples/vault_api/detokenize_records.py b/samples/detokenize_records.py → samples/vault_api/detokenize_records.py
diff --git a/samples/get_column_values.py → samples/vault_api/get_column_values.py b/samples/get_column_values.py → samples/vault_api/get_column_values.py
diff --git a/samples/get_records.py → samples/vault_api/get_records.py b/samples/get_records.py → samples/vault_api/get_records.py
diff --git a/samples/insert_byot.py → samples/vault_api/insert_byot.py b/samples/insert_byot.py → samples/vault_api/insert_byot.py
diff --git a/samples/insert_records.py → samples/vault_api/insert_records.py b/samples/insert_records.py → samples/vault_api/insert_records.py
diff --git a/samples/invoke_connection.py → samples/vault_api/invoke_connection.py b/samples/invoke_connection.py → samples/vault_api/invoke_connection.py
diff --git a/samples/query_records.py → samples/vault_api/query_records.py b/samples/query_records.py → samples/vault_api/query_records.py
diff --git a/samples/tokenize_records.py → samples/vault_api/tokenize_records.py b/samples/tokenize_records.py → samples/vault_api/tokenize_records.py
diff --git a/samples/update_record.py → samples/vault_api/update_record.py b/samples/update_record.py → samples/vault_api/update_record.py
diff --git a/skyflow/service_account/__init__.py b/skyflow/service_account/__init__.py
@@ -1 +1 @@
-from ._utils import generate_bearer_token, generate_bearer_token_from_creds, is_expired
+from ._utils import generate_bearer_token, generate_bearer_token_from_creds, is_expired, generate_signed_data_tokens, generate_signed_data_tokens_from_creds
diff --git a/skyflow/service_account/_utils.py b/skyflow/service_account/_utils.py
@@ -177,4 +177,4 @@ def get_signed_data_token_response_object(signed_token, actual_token):
         "token": actual_token,
         "signed_token": signed_token
     }
-    return response_object
+    return response_object.get("token"), response_object.get("signed_token")
diff --git a/tests/service_account/test__utils.py b/tests/service_account/test__utils.py
@@ -97,13 +97,12 @@ def test_get_signed_jwt_invalid_format(self, mock_jwt_encode):
             get_signed_jwt({}, "client_id", "key_id", "token_uri", "private_key", None)
         self.assertEqual(context.exception.message, SkyflowMessages.Error.JWT_INVALID_FORMAT.value)
 
-
     def test_get_signed_data_token_response_object(self):
         token = "sample_token"
         signed_token = "signed_sample_token"
         response = get_signed_data_token_response_object(signed_token, token)
-        self.assertEqual(response["token"], token)
-        self.assertEqual(response["signed_token"], signed_token)
+        self.assertEqual(response[0], token)
+        self.assertEqual(response[1], signed_token)
 
     def test_generate_signed_data_tokens_from_file_path(self):
         creds_path = os.path.join(os.path.dirname(__file__), "valid_credentials.json")
Original file line number	Diff line number	Diff line change
		@@ -1 +1 @@
		from ._utils import generate_bearer_token, generate_bearer_token_from_creds, is_expired
		from ._utils import generate_bearer_token, generate_bearer_token_from_creds, is_expired, generate_signed_data_tokens, generate_signed_data_tokens_from_creds