Use kaggle-web-client for kaggle_secrets

Author: ifigotin

patches/kaggle_secrets.py

@@ -4,31 +4,12 @@
 (ie. BigQuery).
 """
 
-import json
 import os
-import socket
-import urllib.request
 from datetime import datetime, timedelta
 from enum import Enum, unique
 from typing import Optional, Tuple
-from urllib.error import HTTPError, URLError
-
-_KAGGLE_DEFAULT_URL_BASE = "https://www.kaggle.com"
-_KAGGLE_URL_BASE_ENV_VAR_NAME = "KAGGLE_URL_BASE"
-_KAGGLE_USER_SECRETS_TOKEN_ENV_VAR_NAME = "KAGGLE_USER_SECRETS_TOKEN"
-TIMEOUT_SECS = 40
-
-
-class CredentialError(Exception):
-    pass
-
-
-class BackendError(Exception):
-    pass
-
-
-class ValidationError(Exception):
-    pass
+from kaggle_web_client import KaggleWebClient
+from kaggle_web_client import (CredentialError, BackendError, ValidationError)
 
 class NotFoundError(Exception):
     pass
@@ -56,48 +37,9 @@ def service(self):
 class UserSecretsClient():
     GET_USER_SECRET_ENDPOINT = '/requests/GetUserSecretRequest'
     GET_USER_SECRET_BY_LABEL_ENDPOINT = '/requests/GetUserSecretByLabelRequest'
-    BIGQUERY_TARGET_VALUE = 1
 
     def __init__(self):
-        url_base_override = os.getenv(_KAGGLE_URL_BASE_ENV_VAR_NAME)
-        self.url_base = url_base_override or _KAGGLE_DEFAULT_URL_BASE
-        # Follow the OAuth 2.0 Authorization standard (https://tools.ietf.org/html/rfc6750)
-        self.jwt_token = os.getenv(_KAGGLE_USER_SECRETS_TOKEN_ENV_VAR_NAME)
-        if self.jwt_token is None:
-            raise CredentialError(
-                'A JWT Token is required to use the UserSecretsClient, '
-                f'but none found in environment variable {_KAGGLE_USER_SECRETS_TOKEN_ENV_VAR_NAME}')
-        self.headers = {'Content-type': 'application/json'}
-
-    def _make_post_request(self, data: dict, endpoint: str = GET_USER_SECRET_ENDPOINT) -> dict:
-        # TODO(b/148309982) This code and the code in the constructor should be
-        # removed and this class should use the new KaggleWebClient class instead.
-        url = f'{self.url_base}{endpoint}'
-        request_body = dict(data)
-        request_body['JWE'] = self.jwt_token
-        req = urllib.request.Request(url, headers=self.headers, data=bytes(
-            json.dumps(request_body), encoding="utf-8"))
-        try:
-            with urllib.request.urlopen(req, timeout=TIMEOUT_SECS) as response:
-                response_json = json.loads(response.read())
-                if not response_json.get('wasSuccessful') or 'result' not in response_json:
-                    raise BackendError(
-                        f'Unexpected response from the service. Response: {response_json}.')
-                return response_json['result']
-        except (URLError, socket.timeout) as e:
-            if isinstance(
-                    e, socket.timeout) or isinstance(
-                    e.reason, socket.timeout):
-                raise ConnectionError(
-                    'Timeout error trying to communicate with service. Please ensure internet is on.') from e
-            raise ConnectionError(
-                'Connection error trying to communicate with service.') from e
-        except HTTPError as e:
-            if e.code == 401 or e.code == 403:
-                raise CredentialError(
-                    f'Service responded with error code {e.code}.'
-                    ' Please ensure you have access to the resource.') from e
-            raise BackendError('Unexpected response from the service.') from e
+        self.web_client = KaggleWebClient()
 
     def get_secret(self, label) -> str:
         """Retrieves a user secret value by its label.
@@ -113,7 +55,7 @@ def get_secret(self, label) -> str:
         request_body = {
             'Label': label,
         }
-        response_json = self._make_post_request(request_body, self.GET_USER_SECRET_BY_LABEL_ENDPOINT)
+        response_json = self.web_client.make_post_request(request_body, self.GET_USER_SECRET_BY_LABEL_ENDPOINT)
         if 'secret' not in response_json:
             raise BackendError(
                 f'Unexpected response from the service. Response: {response_json}')
@@ -174,7 +116,7 @@ def _get_access_token(self, target: GcpTarget) -> Tuple[str, Optional[datetime]]
         request_body = {
             'Target': target.target
         }
-        response_json = self._make_post_request(request_body)
+        response_json = self.web_client.make_post_request(request_body, self.GET_USER_SECRET_ENDPOINT)
         if 'secret' not in response_json:
             raise BackendError(
                 f'Unexpected response from the service. Response: {response_json}')

patches/kaggle_web_client.py

@@ -2,14 +2,23 @@
 import os
 import socket
 import urllib.request
-from datetime import datetime, timedelta
-from enum import Enum, unique
-from typing import Optional, Tuple
 from urllib.error import HTTPError, URLError
-from kaggle_secrets import (_KAGGLE_DEFAULT_URL_BASE,
-                            _KAGGLE_URL_BASE_ENV_VAR_NAME,
-                            _KAGGLE_USER_SECRETS_TOKEN_ENV_VAR_NAME,
-                            CredentialError, BackendError, ValidationError)
+
+_KAGGLE_DEFAULT_URL_BASE = "https://www.kaggle.com"
+_KAGGLE_URL_BASE_ENV_VAR_NAME = "KAGGLE_URL_BASE"
+_KAGGLE_USER_SECRETS_TOKEN_ENV_VAR_NAME = "KAGGLE_USER_SECRETS_TOKEN"
+TIMEOUT_SECS = 40
+
+class CredentialError(Exception):
+    pass
+
+
+class BackendError(Exception):
+    pass
+
+
+class ValidationError(Exception):
+    pass
 
 class KaggleWebClient:
     TIMEOUT_SECS = 600