Commit abd338a0 authored by Price, Zach's avatar Price, Zach
Browse files

Fix download auth and expand port range

parent 1607e56c
Loading
Loading
Loading
Loading
Loading
+13 −6
Original line number Diff line number Diff line
@@ -21,12 +21,19 @@ SITE_REGEX = r'^D?[a-z]{3}$'
DATASTREAM_REGEX = r'^D?([a-z]{3})\w*([A-Z][0-9]*)\.([a-z\d]\d)$'


def download_file(file, dest):
class BearerAuth(requests.auth.AuthBase):
    def __init__(self, token):
        self.token = token
    def __call__(self, r):
        r.headers["authorization"] = "Bearer " + self.token
        return r


def download_file(file, dest, token):
    destination = Path(dest).expanduser().resolve()
    if not destination.is_dir():
        raise ValueError(f'Download destination ({destination}) MUST be a directory. Files will be saved to destination/site/datastream/filename.cdf')

    print(f'attempting to save {file["url"]} to {destination}')
    final_path = destination.joinpath(file['name'])
    final_path.parent.mkdir(parents=True, exist_ok=True)

@@ -34,10 +41,10 @@ def download_file(file, dest):
        log.debug(f'Skipping {final_path} because it already exists and matches the expected size.')
        return

    with requests.get(file['url'], stream=True) as r:
    with requests.get(file['url'], stream=True, auth=BearerAuth(token)) as r:
        r.raise_for_status()
        with open(dest, mode='wb') as f:
            with tqdm(desc=dest, total=file['size'], unit='iB', unit_scale=True, unit_divisor=1024) as bar:
        with open(final_path, mode='wb') as f:
            with tqdm(desc=final_path.as_posix(), total=file['size'], unit='iB', unit_scale=True, unit_divisor=1024) as bar:
                for chunk in r.iter_content(chunk_size=config.download_chunk_size):
                    bar.update(f.write(chunk))

@@ -101,7 +108,7 @@ class Archive(object):

    def download(self, destination):
        with ThreadPoolExecutor(config.download_threads) as pool:
            futures = (pool.submit(download_file, file, destination) for file in self.selected_files)
            futures = (pool.submit(download_file, file, destination, self.auth_data['id_token']) for file in self.selected_files)
            for future in tqdm(as_completed(futures), desc=str(self)):
                future.result()

+1 −1
Original line number Diff line number Diff line
@@ -73,7 +73,7 @@ dataService:
    loadBalancerIPs:
    - 128.219.248.172
  guc:
    portRange: 10100-10110
    portRange: 10100-10200
    loadBalancerIP: 128.219.248.175
    cron:
      enabled: true