Skip to content

Commit 2815eeb

Browse files
committed
Add a UserSecrets.get_gcloud_credentials() method.
https://b.corp.google.com/issues/158133824
1 parent 4b8bc76 commit 2815eeb

File tree

2 files changed

+44
-4
lines changed

2 files changed

+44
-4
lines changed

patches/kaggle_secrets.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,9 @@ class BackendError(Exception):
3030
class ValidationError(Exception):
3131
pass
3232

33+
class NotFoundError(Exception):
34+
pass
35+
3336
@unique
3437
class GcpTarget(Enum):
3538
"""Enum class to store GCP targets."""
@@ -116,6 +119,22 @@ def get_secret(self, label) -> str:
116119
f'Unexpected response from the service. Response: {response_json}')
117120
return response_json['secret']
118121

122+
def get_gcloud_credential(self) -> str:
123+
"""Retrieves the Google Cloud SDK credential attached to the current
124+
kernel.
125+
Example usage:
126+
client = UserSecretsClient()
127+
credential_json = client.get_gcloud_credential()
128+
"""
129+
try:
130+
return self.get_secret("__gcloud_sdk_auth__")
131+
except BackendError as backend_error:
132+
message = str(backend_error.args)
133+
if message.find('No user secrets exist') != -1:
134+
raise NotFoundError('Google Cloud SDK credential not found.')
135+
else:
136+
raise
137+
119138
def get_bigquery_access_token(self) -> Tuple[str, Optional[datetime]]:
120139
"""Retrieves BigQuery access token information from the UserSecrets service.
121140

tests/test_user_secrets.py

Lines changed: 25 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
from kaggle_secrets import (_KAGGLE_URL_BASE_ENV_VAR_NAME,
1414
_KAGGLE_USER_SECRETS_TOKEN_ENV_VAR_NAME,
1515
CredentialError, GcpTarget, UserSecretsClient,
16-
BackendError, ValidationError)
16+
BackendError, NotFoundError, ValidationError)
1717

1818
_TEST_JWT = 'test-secrets-key'
1919

@@ -55,7 +55,7 @@ def get_response(self):
5555
if success:
5656
return {'result': {'secret': secret, 'secretType': 'refreshToken', 'secretProvider': 'google', 'expiresInSeconds': 3600}, 'wasSuccessful': "true"}
5757
else:
58-
return {'wasSuccessful': "false"}
58+
return {'wasSuccessful': "false", 'errors': ['No user secrets exist for kernel']}
5959

6060
env = EnvironmentVarGuard()
6161
env.set(_KAGGLE_USER_SECRETS_TOKEN_ENV_VAR_NAME, _TEST_JWT)
@@ -95,7 +95,7 @@ def call_get_secret():
9595
self._test_client(call_get_secret,
9696
'/requests/GetUserSecretByLabelRequest', {'Label': "secret_label", 'JWE': _TEST_JWT},
9797
secret=secret)
98-
98+
9999
def test_get_secret_handles_unsuccessful(self):
100100
def call_get_secret():
101101
client = UserSecretsClient()
@@ -112,7 +112,28 @@ def test_get_secret_validates_label(self):
112112
client = UserSecretsClient()
113113
with self.assertRaises(ValidationError):
114114
secret_response = client.get_secret("")
115-
115+
116+
def test_get_gcloud_secret_succeeds(self):
117+
secret = '{"client_id":"gcloud","type":"authorized_user"}'
118+
119+
def call_get_secret():
120+
client = UserSecretsClient()
121+
secret_response = client.get_gcloud_credential()
122+
self.assertEqual(secret_response, secret)
123+
self._test_client(call_get_secret,
124+
'/requests/GetUserSecretByLabelRequest', {'Label': "__gcloud_sdk_auth__", 'JWE': _TEST_JWT},
125+
secret=secret)
126+
127+
def test_get_gcloud_secret_handles_unsuccessful(self):
128+
def call_get_secret():
129+
client = UserSecretsClient()
130+
with self.assertRaises(NotFoundError):
131+
secret_response = client.get_gcloud_credential()
132+
self._test_client(call_get_secret,
133+
'/requests/GetUserSecretByLabelRequest', {'Label': "__gcloud_sdk_auth__", 'JWE': _TEST_JWT},
134+
success=False)
135+
136+
116137
@mock.patch('kaggle_secrets.datetime')
117138
def test_get_access_token_succeeds(self, mock_dt):
118139
secret = '12345'

0 commit comments

Comments
 (0)