geoseeq 0.2.1__py3-none-any.whl → 0.2.3__py3-none-any.whl

geoseeq/cli/download.py

@@ -5,7 +5,7 @@ from os.path import dirname, join
 
 import click
 import pandas as pd
-
+from multiprocessing import Pool
 from .shared_params import (
     handle_project_id,
     project_id_arg,
@@ -14,13 +14,14 @@ from .shared_params import (
     use_common_state,
     flatten_list_of_els_and_files
 )
-from geoseeq.result.utils import _download_head
+from geoseeq.result.file_download import download_url
 from geoseeq.utils import download_ftp
 from geoseeq.blob_constructors import (
     sample_result_file_from_uuid,
     project_result_file_from_uuid,
 )
 from geoseeq.knex import GeoseeqNotFoundError
+from .progress_bar import PBarManager
 from .utils import convert_size
 
 logger = logging.getLogger('geoseeq_api')
@@ -83,8 +84,16 @@ def cli_download_metadata(state, sample_ids):
     click.echo("Metadata successfully downloaded for samples.", err=True)
 
 
+def _download_one_file(args):
+    url, file_path, pbar = args
+    return download_url(url, filename=file_path, progress_tracker=pbar)
+
+
+cores_option = click.option('--cores', default=1, help='Number of downloads to run in parallel')
+
 @cli_download.command("files")
 @use_common_state
+@cores_option
 @click.option("--target-dir", default=".")
 @click.option('--yes/--confirm', default=False, help='Skip confirmation prompts')
 @click.option("--download/--urls-only", default=True, help="Download files or just print urls")
@@ -98,6 +107,7 @@ def cli_download_metadata(state, sample_ids):
 @sample_ids_arg
 def cli_download_files(
     state,
+    cores,
     sample_name_includes,
     target_dir,
     yes,
@@ -186,23 +196,32 @@ def cli_download_files(
         if not yes:
             click.confirm('Do you want to download these files?', abort=True)
 
+        download_args = []
+        pbars = PBarManager()
         for fname, url in response["links"].items():
             click.echo(f"Downloading file {fname}")
             file_path = join(target_dir, fname)
             makedirs(dirname(file_path), exist_ok=True)
-            if url.startswith("ftp"):
-                download_ftp(url, file_path)
-            else:
-                _download_head(url, file_path)
+            pbar = pbars.get_new_bar(file_path)
+            download_args.append((url, file_path, pbar))
+            if cores == 1:
+                download_url(url, filename=file_path, progress_tracker=pbar)
+
+        if cores > 1:
+            with Pool(cores) as p:
+                for _ in p.imap_unordered(_download_one_file, download_args):
+                    pass
 
 
 @cli_download.command("ids")
 @use_common_state
+@cores_option
 @click.option("--target-dir", default=".")
 @click.option('--yes/--confirm', default=False, help='Skip confirmation prompts')
 @click.option("--download/--urls-only", default=True, help="Download files or just print urls")
+@click.option('--head', default=None, type=int, help='Download the first N bytes of each file')
 @click.argument("ids", nargs=-1)
-def cli_download_ids(state, target_dir, yes, download, ids):
+def cli_download_ids(state, cores, target_dir, yes, download, head, ids):
     """Download a files from GeoSeeq based on their UUID or GeoSeeq Resource Number (GRN).
 
     This command downloads files directly based on their ID. This is used for "manual"
@@ -228,6 +247,7 @@ def cli_download_ids(state, target_dir, yes, download, ids):
     ---
     """
     result_file_ids = flatten_list_of_els_and_files(ids)
+    cores = max(cores, len(result_file_ids))  # don't use more cores than files
     knex = state.get_knex()
     result_files = []
     for result_id in result_file_ids:
@@ -249,8 +269,18 @@ def cli_download_ids(state, target_dir, yes, download, ids):
     if not yes:
         click.confirm('Do you want to download these files?', abort=True)
 
+    download_args = []
+    pbars = PBarManager()
     for result_file in result_files:
         click.echo(f"Downloading file {result_file.get_referenced_filename()}")
         file_path = join(target_dir, result_file.get_referenced_filename())
         makedirs(dirname(file_path), exist_ok=True)
-        result_file.download(file_path)
+        pbar = pbars.get_new_bar(file_path)
+        download_args.append((result_file, file_path, pbar))
+        if cores == 1:
+            result_file.download(file_path, progress_tracker=pbar, head=head)
+
+    if cores > 1:
+        with Pool(cores) as p:
+            for _ in p.imap_unordered(_download_one_file, download_args):
+                pass

geoseeq/cli/main.py

@@ -31,7 +31,7 @@ main.add_command(cli_upload)
 @main.command()
 def version():
     """Print the version of the Geoseeq API being used."""
-    click.echo('0.2.1')  # remember to update setup
+    click.echo('0.2.3')  # remember to update setup
 
 
 @main.group('advanced')

geoseeq/cli/progress_bar.py

@@ -0,0 +1,28 @@
+from tqdm import tqdm
+from os.path import basename
+
+class TQBar:
+
+    def __init__(self, pos, desc) -> None:
+        self.n_bars = 0
+        self.pos = pos
+        self.desc = desc
+        self.bar = None
+
+    def set_num_chunks(self, n_chunks):
+        self.n_bars = n_chunks
+        self.bar = tqdm(total=n_chunks, position=self.pos, desc=self.desc, leave=False)
+
+    def update(self, chunk_num):
+        self.bar.update(chunk_num)
+
+
+class PBarManager:
+
+    def __init__(self):
+        self.n_bars = 0
+        self.pbars = []
+
+    def get_new_bar(self, filepath):
+        self.n_bars += 1
+        return TQBar(self.n_bars, basename(filepath))

geoseeq/cli/upload/upload_reads.py

@@ -1,5 +1,4 @@
 import logging
-from tqdm import tqdm
 import click
 import requests
 from os.path import basename
@@ -20,39 +19,11 @@ from geoseeq.cli.shared_params import (
 )
 
 from geoseeq.constants import FASTQ_MODULE_NAMES
-
-
+from geoseeq.cli.progress_bar import PBarManager
 
 logger = logging.getLogger('geoseeq_api')
 
 
-class TQBar:
-
-    def __init__(self, pos, desc) -> None:
-        self.n_bars = 0
-        self.pos = pos
-        self.desc = desc
-        self.bar = None
-
-    def set_num_chunks(self, n_chunks):
-        self.n_bars = n_chunks
-        self.bar = tqdm(total=n_chunks, position=self.pos, desc=self.desc, leave=False)
-
-    def update(self, chunk_num):
-        self.bar.update(chunk_num)
-
-
-class PBarManager:
-
-    def __init__(self):
-        self.n_bars = 0
-        self.pbars = []
-
-    def get_new_bar(self, filepath):
-        self.n_bars += 1
-        return TQBar(self.n_bars, basename(filepath))
-
-
 def _make_in_process_logger(log_level):
     logger = logging.getLogger('geoseeq_api')
     logger.setLevel(log_level)

geoseeq/result/file_download.py

@@ -1,12 +1,65 @@
 
 import urllib.request
+import logging
+import requests
 from os.path import basename, getsize, join
 from pathlib import Path
 from tempfile import NamedTemporaryFile
 
 from geoseeq.utils import download_ftp
+from geoseeq.constants import FIVE_MB
+
+logger = logging.getLogger("geoseeq_api")  # Same name as calling module
+
+
+def _download_head(url, filename, head=None, progress_tracker=None):
+    headers = None
+    if head and head > 0:
+        headers = {"Range": f"bytes=0-{head}"}
+    response = requests.get(url, stream=True, headers=headers)
+    total_size_in_bytes = int(response.headers.get('content-length', 0))
+    if progress_tracker: progress_tracker.set_num_chunks(total_size_in_bytes)
+    block_size = FIVE_MB
+    with open(filename, 'wb') as file:
+        for data in response.iter_content(block_size):
+            if progress_tracker: progress_tracker.update(len(data))
+            file.write(data)
+    return filename
+
+
+def _download_generic(url, filename, head=None):
+    urllib.request.urlretrieve(url, filename)
+    return filename
+
+
+def guess_download_kind(url):
+    if 'azure' in url:
+        return 'azure'
+    elif 's3' in url:
+        return 's3'
+    elif 'ftp' in url:
+        return 'ftp'
+    else:
+        return 'generic'
+
+
+def download_url(url, kind='guess', filename=None, head=None, progress_tracker=None):
+    """Return a local filepath to the downloaded file. Download the file."""
+    if kind == 'guess':
+        kind = guess_download_kind(url)
+        logger.info(f"Guessed download kind: {kind} for {url}")
+    logger.info(f"Downloading {kind} file to {filename}")
+    if kind == 'generic':
+        return _download_generic(url, filename, head=head)
+    elif kind == 's3':
+        return _download_head(url, filename, head=head, progress_tracker=progress_tracker)
+    elif kind == 'azure':
+        return _download_head(url, filename, head=head)
+    elif kind == 'ftp':
+        return download_ftp(url, filename, head=head)
+    else:
+        raise ValueError(f"Unknown download kind: {kind}")
 
-from .utils import *
 
 
 class ResultFileDownload:
@@ -15,21 +68,21 @@ class ResultFileDownload:
     def get_download_url(self):
         """Return a URL that can be used to download the file for this result."""
         blob_type = self.stored_data.get("__type__", "").lower()
-        if blob_type not in ["s3", "sra"]:
-            raise TypeError("Cannot fetch a file for a BLOB type result field.")
-        if blob_type == "s3":
+        if blob_type not in ["s3", "sra", "ftp", "azure"]:
+            raise ValueError(f'Unknown URL type: "{blob_type}"')
+        key = 'url' if 'url' in self.stored_data else 'uri'
+        if blob_type in ["s3", "azure"]:
             try:
                 url = self.stored_data["presigned_url"]
             except KeyError:
-                url = self.stored_data["uri"]
+                url = self.stored_data[key]
             if url.startswith("s3://"):
                 url = self.stored_data["endpoint_url"] + "/" + url[5:]
             return url
-        elif blob_type == "sra":
-            url = self.stored_data["url"]
-            return url
+        else:
+            return self.stored_data[key]
 
-    def download_file(self, filename=None, cache=True, head=None):
+    def download(self, filename=None, cache=True, head=None, progress_tracker=None):
         """Return a local filepath to the file this result points to."""
         if not filename:
             self._temp_filename = True
@@ -39,57 +92,11 @@ class ResultFileDownload:
         blob_type = self.stored_data.get("__type__", "").lower()
         if cache and self._cached_filename:
             return self._cached_filename
-        if blob_type == "s3":
-            return self._download_s3(filename, cache, head=head)
-        elif blob_type == "sra":
-            return self._download_sra(filename, cache)
-        elif blob_type == "ftp":
-            return self._download_ftp(filename, cache)
-        elif blob_type == "azure":
-            return self._download_azure(filename, cache, head=head)
-        else:
-            raise TypeError("Cannot fetch a file for a BLOB type result field.")
-
-    def _download_s3(self, filename, cache, head=None):
-        logger.info(f"Downloading S3 file to {filename}")
-        try:
-            url = self.stored_data["presigned_url"]
-        except KeyError:
-            key = 'uri' if 'uri' in self.stored_data else 'url'
-            url = self.stored_data[key]
-        if url.startswith("s3://"):
-            url = self.stored_data["endpoint_url"] + "/" + url[5:]
-        _download_head(url, filename, head=head) 
-        if cache:
-            self._cached_filename = filename
-        return filename
-
-    def _download_azure(self, filename, cache, head=None):
-        logger.info(f"Downloading Azure file to {filename}")
-        try:
-            url = self.stored_data["presigned_url"]
-        except KeyError:
-            key = 'uri' if 'uri' in self.stored_data else 'url'
-            url = self.stored_data[key]
-        _download_head(url, filename, head=head)
-        if cache:
-            self._cached_filename = filename
-        return filename
-
-    def _download_sra(self, filename, cache):
-        return self._download_generic_url(filename, cache)
-
-    def _download_ftp(self, filename, cache, head=None):
-        logger.info(f"Downloading FTP file to {filename}")
-        key = 'url' if 'url' in self.stored_data else 'uri'
-        download_ftp(self.stored_data[key], filename, head=head)
-        return filename
-
-    def _download_generic_url(self, filename, cache):
-        logger.info(f"Downloading generic URL file to {filename}")
-        key = 'url' if 'url' in self.stored_data else 'uri'
-        url = self.stored_data[key]
-        urllib.request.urlretrieve(url, filename)
+        url = self.get_download_url()
+        filepath = download_url(
+            url, blob_type, filename,
+            head=head, progress_tracker=progress_tracker
+        )
         if cache:
-            self._cached_filename = filename
-        return filename
+            self._cached_filename = filepath
+        return filepath

geoseeq/result/file_upload.py

@@ -108,7 +108,7 @@ class ResultFileUpload:
             for num, url in enumerate(list(urls.values())):
                 response_part = self._upload_one_part(file_chunker, url, num, max_retries, session)
                 complete_parts.append(response_part)
-                progress_tracker.update(file_chunker.get_chunk_size(num))
+                if progress_tracker: progress_tracker.update(file_chunker.get_chunk_size(num))
                 logger.info(f'Uploaded part {num + 1} of {len(urls)} for "{file_chunker.filepath}"')
             return complete_parts
         
@@ -123,7 +123,7 @@ class ResultFileUpload:
             for future in as_completed(futures):
                 response_part = future.result()
                 complete_parts.append(response_part)
-                progress_tracker.update(file_chunker.get_chunk_size(response_part["PartNumber"] - 1))
+                if progress_tracker: progress_tracker.update(file_chunker.get_chunk_size(response_part["PartNumber"] - 1))
                 logger.info(
                     f'Uploaded part {response_part["PartNumber"]} of {len(urls)} for "{file_chunker.filepath}"'
                 )
@@ -137,7 +137,7 @@ class ResultFileUpload:
         chunk_size=FIVE_MB,
         max_retries=3,
         session=None,
-        progress_tracker=lambda x: None,
+        progress_tracker=None,
         threads=1,
     ):
         """Upload a file to S3 using the multipart upload process."""
@@ -146,7 +146,7 @@ class ResultFileUpload:
         logger.info(f'Starting upload for "{filepath}"')
         complete_parts = []
         file_chunker = FileChunker(filepath, chunk_size).load_all_chunks()
-        progress_tracker.set_num_chunks(file_chunker.file_size)
+        if progress_tracker: progress_tracker.set_num_chunks(file_chunker.file_size)
         complete_parts = self._upload_parts(file_chunker, urls, max_retries, session, progress_tracker, threads)
         self._finish_multipart_upload(upload_id, complete_parts)
         logger.info(f'Finished Upload for "{filepath}"')

geoseeq/result/utils.py

@@ -15,22 +15,6 @@ from geoseeq.utils import download_ftp, md5_checksum
 
 logger = logging.getLogger("geoseeq_api")  # Same name as calling module
 logger.addHandler(logging.NullHandler())  # No output unless configured by calling program
-
-
-def _download_head(url, filename, head=None):
-    if head and head > 0:
-        opener = urllib.request.build_opener()
-        if head:
-            opener.addheaders = [('Range', f'bytes=0-{head}')]
-        urllib.request.install_opener(opener)
-    try:
-        urllib.request.urlretrieve(url, filename)  # can throw 416 error if head is too large
-    except urllib.error.HTTPError as e:
-        if e.code == 416:
-            logger.warning(f"HEAD request failed, trying again without HEAD.")
-            _download_head(url, filename, head=None)
-        else:
-            raise e
     
 
 def diff_dicts(blob1, blob2):

{geoseeq-0.2.1.dist-info → geoseeq-0.2.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: geoseeq
-Version: 0.2.1
+Version: 0.2.3
 Summary: # Geoseeq API Client
 Author: David C. Danko
 Author-email: dcdanko@biotia.io

{geoseeq-0.2.1.dist-info → geoseeq-0.2.3.dist-info}/RECORD

@@ -18,10 +18,11 @@ geoseeq/cli/constants.py,sha256=Do5AUf9lMO9_P8KpFJ3XwwFBAWsxSjZ6sx9_QEGyC_c,176
 geoseeq/cli/copy.py,sha256=ocJ-T6vZBTPELUSBA65KndzAugrBWMq-3ux0dbffacE,1997
 geoseeq/cli/create.py,sha256=nMixc5DgLdHS8JhJE8Jt1LAvDprQdwojaOQ2bvvI-QU,6686
 geoseeq/cli/delete.py,sha256=uvdmBqxBTquwLQAdP-fNpSfOd2hjImh8hMe-AnBd12Q,919
-geoseeq/cli/download.py,sha256=a__zCwCBCgBfdh604tpfeJqF8WGZid-vV7ZPKqXPTKY,9753
+geoseeq/cli/download.py,sha256=T4jB9lZCaksoZFGigk7sJX1_QyFR6Zt9MfNUWTrUjB8,10930
 geoseeq/cli/fastq_utils.py,sha256=-bmeQLaiMBm57zWOF0R5OlWTU0_3sh1JBC1RYw2BOFM,3083
 geoseeq/cli/list.py,sha256=BvqwUJeOGvEVy7O6H6xYR2raYoyvR0yjvY2iklOCWCk,1712
-geoseeq/cli/main.py,sha256=wclMAB3Z7a1-oBOqF-0qGy84kHiwDXZM_xJY8mb3a7M,1231
+geoseeq/cli/main.py,sha256=mgNakRaRrA_jTqU_OF7LkGR9vUKBI-Yom4pGK_IFK4g,1231
+geoseeq/cli/progress_bar.py,sha256=LrhMTbTXHaRSX99NJQzRncreDXl6edov3Zo-Y3-d5_c,644
 geoseeq/cli/user.py,sha256=fARv_5vx-QYT765Hsjgwr6J5ddA_viCPQdv9iUoVX6w,925
 geoseeq/cli/utils.py,sha256=f0wX5-19uDrYZ8d-aBLplqfbGOkr5BH8x3AMSAyPCes,2873
 geoseeq/cli/view.py,sha256=YmGcN4uGFuXWd_OA6Pi0KYvT-z4AD3nUwaOLewwJhBw,645
@@ -33,7 +34,7 @@ geoseeq/cli/shared_params/obj_getters.py,sha256=wmT37kwwcguy_JxqE7G9UH5mwcSZ0byw
 geoseeq/cli/shared_params/opts_and_args.py,sha256=kxqqsbXoGg58GE1VtIySb2TFw8KlRodsvh4xXWEVBRw,1352
 geoseeq/cli/upload/__init__.py,sha256=QqBh0WXD9BaEDcKsKpb-WgWwQ0WzZdQUgYXYwyowW_c,340
 geoseeq/cli/upload/upload.py,sha256=JC5AH9CNR7dKryZzkxSf7DX9ZqPer7-fpvahNMtlVQQ,6153
-geoseeq/cli/upload/upload_reads.py,sha256=xSoGPK_-9N3jTu-nNYgb0T_B27D-Dh8nbJ_V8HVTtls,7404
+geoseeq/cli/upload/upload_reads.py,sha256=2y1rN1zLDpEj0rwKIDLjVnNA6BeXxM6ZDYn6lNtIsbM,6834
 geoseeq/contrib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 geoseeq/contrib/ncbi/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 geoseeq/contrib/ncbi/api.py,sha256=WQeLoGA_-Zha-QeSO8_i7HpvXyD8UkV0qc5okm11KiA,1056
@@ -42,11 +43,11 @@ geoseeq/contrib/ncbi/cli.py,sha256=j9zEcaZPTryK3a4xluRxigcJKDhRpRxbp3KZSx-Bfhk,2
 geoseeq/contrib/ncbi/setup_logging.py,sha256=Tp1bY1U0f-o739aHpvVYriG2qdd1lFvCYBXZeXQgt-w,175
 geoseeq/result/__init__.py,sha256=IFHIyRV8ZzuKIfwfze1SXgcKwNMcSgMAknLHMkwjXIU,356
 geoseeq/result/bioinfo.py,sha256=QQtbyogrdro9avJSN0713sxLVnVeA24mFw3hWtKDKyw,1782
-geoseeq/result/file_download.py,sha256=S9QDSQ0rAJ2xLCw4Sup_PHsjN3GG8ljlq6zsxc4cMy0,3676
-geoseeq/result/file_upload.py,sha256=6ztnHubunURBR343ulLLNVqfJTK4jOQpHRzWZHU8uE0,6371
+geoseeq/result/file_download.py,sha256=XoWht-Al0HkvQxkDnU7IKqtisNBTrseDIV9IPwuDcuY,3617
+geoseeq/result/file_upload.py,sha256=6HhsxeTsMHMrXjRdhL5k8bAMIxx7fBYQ6oesFKOcMqs,6424
 geoseeq/result/result_file.py,sha256=YMCAivUXgHUNEhH6yTGCHJh9gXHd8OLjjbnccdvW7iw,7138
 geoseeq/result/result_folder.py,sha256=LhCD75837cip3Sjv4thl-Kqv3oyLJI-NVRPwqThktvM,7660
-geoseeq/result/utils.py,sha256=H7IBNv8WbHem_WnU2wFdJj-S840KxUMbzIk_gOhj71k,3354
+geoseeq/result/utils.py,sha256=C-CxGzB3WddlnRiqFSkrY78I_m0yFgNqsTBRzGU-y8Q,2772
 geoseeq/vc/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 geoseeq/vc/checksum.py,sha256=y8rh1asUZNbE_NLiFO0-9hImLNiTOc2YXQBRKORWK7k,710
 geoseeq/vc/cli.py,sha256=VlAvBoXckENKAvREZ5nKXl2F6Z6Glh_mKBDiSfaT_lQ,2853
@@ -59,9 +60,9 @@ geoseeq/vc/vc_stub.py,sha256=BNFr7CZtpvLPe0MOqgB0zYL037vRCSwd2A_Fcy9UTKA,3160
 tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/test_api_client.py,sha256=JafItPjtL8QnhZs7Vcl5bg4E6N5rLDYMqbrKaWwGjWs,14444
 tests/test_work_orders.py,sha256=CB6QbGGMNeEePZiOVioiSQP2dQsyMnsy4L31AQMRs3A,498
-geoseeq-0.2.1.dist-info/LICENSE,sha256=IuhIl1XCxXLPLJT_coN1CNqQU4Khlq7x4IdW7ioOJD8,1067
-geoseeq-0.2.1.dist-info/METADATA,sha256=8aC8pIOpKRUZ_oNZJrQdHTMnNW5TzUx_Jd8FZV7nhI0,454
-geoseeq-0.2.1.dist-info/WHEEL,sha256=5sUXSg9e4bi7lTLOHcm6QEYwO5TIF1TNbTSVFVjcJcc,92
-geoseeq-0.2.1.dist-info/entry_points.txt,sha256=yF-6KDM8zXib4Al0qn49TX-qM7PUkWUIcYtsgt36rjM,45
-geoseeq-0.2.1.dist-info/top_level.txt,sha256=zZk7mmeaqAYqFJG8nq2DTgSQPbflRjJwkDIhNURPDEU,14
-geoseeq-0.2.1.dist-info/RECORD,,
+geoseeq-0.2.3.dist-info/LICENSE,sha256=IuhIl1XCxXLPLJT_coN1CNqQU4Khlq7x4IdW7ioOJD8,1067
+geoseeq-0.2.3.dist-info/METADATA,sha256=f5sqGnpDcExJyNFt6DoEJ7BXjp32M125L91Y3kojrow,454
+geoseeq-0.2.3.dist-info/WHEEL,sha256=5sUXSg9e4bi7lTLOHcm6QEYwO5TIF1TNbTSVFVjcJcc,92
+geoseeq-0.2.3.dist-info/entry_points.txt,sha256=yF-6KDM8zXib4Al0qn49TX-qM7PUkWUIcYtsgt36rjM,45
+geoseeq-0.2.3.dist-info/top_level.txt,sha256=zZk7mmeaqAYqFJG8nq2DTgSQPbflRjJwkDIhNURPDEU,14
+geoseeq-0.2.3.dist-info/RECORD,,