Commit 6ad22a4f authored by Price, Zach's avatar Price, Zach
Browse files

Include stage_arm_data cli

parent 5826bf03

app/ADL/client/cli.py

0 → 100644
+80 −0
Original line number Diff line number Diff line 
import argparse

from datetime import datetime

from pathlib import Path

import re



import typer

from typer import Argument, Option, BadParameter, echo, Typer



from ADL.client import Archive, DATASTREAM_REGEX



app = Typer()



def validate_datastream(datastream: str):

    if not re.match(DATASTREAM_REGEX, datastream):

        raise BadParameter(f"Datastream must match the regex {DATASTREAM_REGEX}")

    return datastream



valid_datetime_formats = [

    "%Y-%m-%d",

    "%Y%m%d",

    "%Y%m%d%H%M%S",

    "%Y-%m-%dT%H:%M:%S",

    "%Y-%m-%d %H:%M:%S",

    "%Y-%m-%d %H:%M:%S.%f",

]



datastream_arg = Option(

    default=...,

    callback=validate_datastream,

    help="ARM Datastream to download files from",

)



start_arg = Option(

    default="1970-01-01T00:00:00",

    help="Include files with data AFTER this date/time",

    formats=valid_datetime_formats,

)



end_arg = Option(

    default=str(datetime.now()),

    help="Include files with data BEFORE this date/time",

    formats=valid_datetime_formats,

)



dry_run_arg = Option(False, "--dry-run")



to_arg = Option(

    default=Path.cwd(),

    help="Directory to which data will be staged",

    resolve_path=True,

    file_okay=False,

    writable=True,

)



username_arg = Option(

    default="",

    envvar="USER",

)



password_arg = Option(

    default="",

    prompt_required=True,

)



@app.command()

def main(datastream: str = datastream_arg,

         start: datetime = start_arg,

         end: datetime = end_arg,

         dryrun: bool = dry_run_arg,

         to: Path = to_arg,

         username: str = username_arg,

         password: str = password_arg,):

    af = Archive(username=username, password=password)[datastream][start:end]

    echo(f'Downloading {af.file_count:,} files to {to}')

    if not dryrun:

        af.download(to)





if __name__ == "__main__":

    app()

app/ADL/client/config.py

+1 −1
Original line number Diff line number Diff line
@@ -5,7 +5,7 @@ from pydantic import AnyHttpUrl, BaseSettings 


class Settings(BaseSettings):

    debug: bool = True

    endpoint: AnyHttpUrl = 'http://localhost:8000/v1_0'

    endpoint: AnyHttpUrl = 'https://api.k8s.arm.gov/v1_0'

    download_chunk_size: int = 8192

    download_threads: int = None

app/setup.py

+4 −1
Original line number Diff line number Diff line
@@ -25,6 +25,9 @@ setup( 
            'hsi=ADL.api.fs:HPSSFileSystem',

            'guc=ADL.api.fs:GUCFileSystem',

            'fullpathcache=ADL.api.fs:FullPathSimpleCache',

        ]

        ],

        'console_scripts': [

            'stage_arm_data = ADL.client.cli:app',

        ],

    }

)