Use sessions for downloaders, and allow to add auth headers

datahugger/__main__.py

@@ -95,6 +95,14 @@ def main():
         "May appear multiple times.",
     )
 
+    parser.add_argument(
+        "--auth-token",
+        help="Add an authentication token to the downloader. This will send an "
+             "Authorization header with the value of the token "
+             "(i.e. Authorization Bearer <token>).",
+        default=None,
+    )
+
     # version
     parser.add_argument(
         "-V",
@@ -119,6 +127,7 @@ def main():
             progress=args.progress,
             print_only=args.print_only,
             params=args.params,
+            token=args.auth_token,
         )
 
     except DOIError as doi_err:

datahugger/api.py

@@ -59,6 +59,7 @@ def info(
     progress=True,
     print_only=False,
     params=None,
+    token=None,
 ):
     """Get info on the content of the dataset.
 
@@ -83,6 +84,8 @@ def info(
         the actual files (Dry run). Default: False.
     params: dict
         Extra parameters for the request.
+    token: str
+        The authentication token for the service to be sent as an Authorization header.
 
     Returns
     -------
@@ -103,6 +106,7 @@ def info(
         progress=progress,
         print_only=print_only,
         params=params,
+        token=token,
     )
 
 
@@ -116,6 +120,7 @@ def get(
     progress=True,
     print_only=False,
     params=None,
+    token=None,
 ):
     """Get the content of repository.
 
@@ -145,6 +150,8 @@ def get(
         the actual files (Dry run). Default: False.
     params: dict
         Extra parameters for the request.
+    token: str
+        The authentication token for the service to be sent as an Authorization header.
 
     Returns
     -------
@@ -162,6 +169,7 @@ def get(
         progress=progress,
         print_only=print_only,
         params=params,
+        token=token,
     )
 
     return service.download(output_folder)

datahugger/base.py

@@ -57,6 +57,7 @@ def __init__(
         checksum=False,
         print_only=False,
         params=None,
+        token=None,
     ):
         super().__init__()
         self.resource = resource
@@ -67,6 +68,11 @@ def __init__(
         self.checksum = checksum
         self.print_only = print_only
         self.params = params
+        self.token = token
+
+        self.session = requests.Session()
+        if self.token:
+            self.session.headers["Authorization"] = f"Bearer {self.token}"
 
     def _get_attr_attr(self, record, jsonp):
         try:
@@ -160,7 +166,7 @@ def download_file(
 
         if not self.print_only:
             logging.info(f"Downloading file {file_link}")
-            res = requests.get(file_link, stream=True)
+            res = self.session.get(file_link, stream=True)
             res.raise_for_status()
 
             output_fp = Path(output_folder, file_name)
@@ -198,7 +204,7 @@ def _parse_url(self, url):
             raise ValueError(f"Failed to parse URL '{url}'") from err
 
     def _unpack_single_folder(self, zip_url, output_folder):
-        r = requests.get(zip_url)
+        r = self.session.get(zip_url)
         r.raise_for_status()
 
         z = zipfile.ZipFile(io.BytesIO(r.content))
@@ -286,7 +292,7 @@ def _get_files_recursive(self, url, folder_name=None, base_url=None):
         result = []
 
         # get the data from URL
-        res = requests.get(url)
+        res = self.session.get(url)
         res.raise_for_status()
         response = res.json()