Merge pull request #227 from lorenzorubi-db/master

chores: added timeout, verbose, repos parent directory

Author: gregwood-db

dbclient/WorkspaceClient.py

@@ -922,7 +922,7 @@ def create_repo(self, repo_str, error_logger, checkpoint_repo_set):
                 parent_directory = re.sub(r"^RESOURCE_DOES_NOT_EXIST: Parent directory ", '', resp.get('message'))
                 parent_directory = re.sub(r" does not exist.$", '', parent_directory)
                 if re.fullmatch(
-                    r'/Repos/.+\@.+', parent_directory
+                    r'/Repos/.+[^/]', parent_directory
                 ):
                     logging.info(f"Creating parent directory {parent_directory}")
                     resp2 = self.post('/workspace/mkdirs', {"path": parent_directory})

dbclient/dbclient.py

@@ -190,17 +190,21 @@ def req_session(self):
             self._local.session = session
         return self._local.session
 
-    def get(self, endpoint, json_params=None, version='2.0', print_json=False, do_not_throw=False):
+    def get(self, endpoint, json_params=None, version='2.0', print_json=False, do_not_throw=False, timeout=10):
         if version:
             ver = version
         while True:
             full_endpoint = self._url + '/api/{0}'.format(ver) + endpoint
             if self.is_verbose():
                 print("Get: {0}".format(full_endpoint))
             if json_params:
-                raw_results = self.req_session().get(full_endpoint, headers=self._token, params=json_params, verify=self._verify_ssl)
+                raw_results = self.req_session().get(
+                    full_endpoint, headers=self._token, params=json_params, verify=self._verify_ssl, timeout=timeout
+                )
             else:
-                raw_results = self.req_session().get(full_endpoint, headers=self._token, verify=self._verify_ssl)
+                raw_results = self.req_session().get(
+                    full_endpoint, headers=self._token, verify=self._verify_ssl, timeout=timeout
+                )
 
             if self._should_retry_with_new_token(raw_results):
                 continue
@@ -218,7 +222,7 @@ def get(self, endpoint, json_params=None, version='2.0', print_json=False, do_no
             results['http_status_code'] = http_status_code
             return results
 
-    def http_req(self, http_type, endpoint, json_params, version='2.0', print_json=False, files_json=None):
+    def http_req(self, http_type, endpoint, json_params, version='2.0', print_json=False, files_json=None, timeout=10):
         if version:
             ver = version
         while True:
@@ -228,17 +232,25 @@ def http_req(self, http_type, endpoint, json_params, version='2.0', print_json=F
             if json_params:
                 if http_type == 'post':
                     if files_json:
-                        raw_results = self.req_session().post(full_endpoint, headers=self._token,
-                                                    data=json_params, files=files_json, verify=self._verify_ssl)
+                        raw_results = self.req_session().post(
+                            full_endpoint, headers=self._token, data=json_params, files=files_json,
+                            verify=self._verify_ssl, timeout=timeout
+                        )
                     else:
-                        raw_results = self.req_session().post(full_endpoint, headers=self._token,
-                                                    json=json_params, verify=self._verify_ssl)
+                        raw_results = self.req_session().post(
+                            full_endpoint, headers=self._token, json=json_params, verify=self._verify_ssl,
+                            timeout=timeout
+                        )
                 if http_type == 'put':
-                    raw_results = self.req_session().put(full_endpoint, headers=self._token,
-                                               json=json_params, verify=self._verify_ssl)
+                    raw_results = self.req_session().put(
+                        full_endpoint, headers=self._token, json=json_params, verify=self._verify_ssl,
+                        timeout=timeout
+                    )
                 if http_type == 'patch':
-                    raw_results = self.req_session().patch(full_endpoint, headers=self._token,
-                                                 json=json_params, verify=self._verify_ssl)
+                    raw_results = self.req_session().patch(
+                        full_endpoint, headers=self._token, json=json_params, verify=self._verify_ssl,
+                        timeout=timeout
+                    )
             else:
                 print("Must have a payload in json_args param.")
                 return {}

dbclient/parser.py

@@ -461,6 +461,9 @@ def get_pipeline_parser() -> argparse.ArgumentParser:
     parser.add_argument('--silent', action='store_true', default=False,
                         help='Silent all logging of export operations.')
 
+    parser.add_argument('--verbose', action='store_true', default=False,
+                        help='Verbose logging')
+
     parser.add_argument('--no-ssl-verification', action='store_true',
                         help='Set Verify=False when making http requests.')
 
@@ -550,4 +553,5 @@ def get_pipeline_parser() -> argparse.ArgumentParser:
 
     parser.add_argument('--groups-to-keep', nargs='+', type=str, default=[],
                         help='List of groups (and therefore users/notebooks) to keep if specified')
+
     return parser

migration_pipeline.py

@@ -55,6 +55,8 @@ def build_pipeline(args) -> Pipeline:
     client_config['base_dir'] = client_config['export_dir']
     client_config['export_dir'] = os.path.join(client_config['base_dir'], session) + '/'
 
+    client_config['verbose'] = args.verbose
+
     if not args.dry_run:
         os.makedirs(client_config['export_dir'], exist_ok=True)