Move timeout

Kaggle · Jul 23, 2020 · 2259a07 · 2259a07
1 parent 4e07647
commit 2259a07
Show file tree

Hide file tree

Showing 5 changed files with 11 additions and 11 deletions.
diff --git a/patches/kaggle_datasets.py b/patches/kaggle_datasets.py
@@ -5,6 +5,7 @@
 
 class KaggleDatasets:
     GET_GCS_PATH_ENDPOINT = '/requests/CopyDatasetVersionToKnownGcsBucketRequest'
+    TIMEOUT_SECS = 600
 
     # Integration types for GCS
     AUTO_ML = 1
@@ -20,5 +21,5 @@ def get_gcs_path(self, dataset_dir: str = None) -> str:
             'MountSlug': dataset_dir,
             'IntegrationType': integration_type,
         }
-        result = self.web_client.make_post_request(data, self.GET_GCS_PATH_ENDPOINT)
+        result = self.web_client.make_post_request(data, self.GET_GCS_PATH_ENDPOINT, self.TIMEOUT_SECS)
         return result['destinationBucket']
diff --git a/patches/kaggle_secrets.py b/patches/kaggle_secrets.py
@@ -9,7 +9,10 @@
 from enum import Enum, unique
 from typing import Optional, Tuple
 from kaggle_web_client import KaggleWebClient
-from kaggle_web_client import (CredentialError, BackendError, ValidationError)
+from kaggle_web_client import (CredentialError, BackendError)
+
+class ValidationError(Exception):
+    pass
 
 class NotFoundError(Exception):
     pass

diff --git a/patches/kaggle_web_client.py b/patches/kaggle_web_client.py
@@ -17,11 +17,7 @@ class BackendError(Exception):
     pass
 
 
-class ValidationError(Exception):
-    pass
-
 class KaggleWebClient:
-    TIMEOUT_SECS = 600
 
     def __init__(self):
         url_base_override = os.getenv(_KAGGLE_URL_BASE_ENV_VAR_NAME)
@@ -38,14 +34,14 @@ def __init__(self):
             'X-Kaggle-Authorization': f'Bearer {self.jwt_token}',
         }
 
-    def make_post_request(self, data: dict, endpoint: str) -> dict:
+    def make_post_request(self, data: dict, endpoint: str, timeout: int = TIMEOUT_SECS) -> dict:
         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=self.TIMEOUT_SECS) as response:
+            with urllib.request.urlopen(req, timeout=timeout) as response:
                 response_json = json.loads(response.read())
                 if not response_json.get('wasSuccessful') or 'result' not in response_json:
                     raise BackendError(

diff --git a/tests/test_datasets.py b/tests/test_datasets.py
@@ -9,7 +9,7 @@
 from kaggle_web_client import (KaggleWebClient,
                             _KAGGLE_URL_BASE_ENV_VAR_NAME,
                             _KAGGLE_USER_SECRETS_TOKEN_ENV_VAR_NAME,
-                            CredentialError, BackendError, ValidationError)
+                            CredentialError, BackendError)
 from kaggle_datasets import KaggleDatasets, _KAGGLE_TPU_NAME_ENV_VAR_NAME
 
 _TEST_JWT = 'test-secrets-key'

diff --git a/tests/test_user_secrets.py b/tests/test_user_secrets.py
@@ -11,10 +11,10 @@
 from google.auth.exceptions import DefaultCredentialsError
 from google.cloud import bigquery
 from kaggle_secrets import (GcpTarget, UserSecretsClient,
-                            NotFoundError)
+                            NotFoundError, ValidationError)
 from kaggle_web_client import (_KAGGLE_URL_BASE_ENV_VAR_NAME,
                             _KAGGLE_USER_SECRETS_TOKEN_ENV_VAR_NAME,
-                            CredentialError, BackendError, ValidationError)
+                            CredentialError, BackendError)
 
 _TEST_JWT = 'test-secrets-key'