Commit b33cf9c2 authored by Yakubov, Sergey's avatar Yakubov, Sergey
Browse files

add filesize to headers, set chunk szie as parameter

parent 8adba02b

remote_data_broker/brokers/filesys/broker.py

+17 −7
Original line number Diff line number Diff line
@@ -66,16 +66,29 @@ class FilesysBroker(RemoteDataBroker): 
        copy_complete(user_info, data.input_path, settings.rdb_storage_path + "/" + uid)

        return UploadResponse(uid=uid)



    def __get_download_cmd(self, data: DownloadData) -> List[str]:

    def __get_command(self, cmd: str, data: DownloadData) -> List[str]:

        token = data.token

        cmd = settings.rdb_cat_cmd

        cmd = cmd.replace("$token", token)

        if cmd:

            return cmd.split(" ") + [data.filename]

        return ["cat", data.filename]



    def __get_file_size(self, data: DownloadData) -> str:

        command = self.__get_command(settings.rdb_fsize_cmd, data)

        try:

            result = subprocess.run(command, check=True, stdout=subprocess.PIPE)

            return result.stdout.decode("utf-8")

        except Exception:

            raise HTTPException(status_code=500, detail="download failed")



    def download(self, data: DownloadData) -> StreamingResponse:

        command = self.__get_download_cmd(data)

        headers = {

            "Content-Disposition": f'attachment; filename="{os.path.basename(data.filename)}"',

        }

        if settings.rdb_send_filesize:

            fsize = self.__get_file_size(data)

            headers["File-Size"] = fsize.strip()

        command = self.__get_command(settings.rdb_cat_cmd, data)

        try:

            process = subprocess.Popen(

                command, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE
@@ -84,7 +97,7 @@ class FilesysBroker(RemoteDataBroker): 
            def generate() -> Generator:

                while True:

                    assert process.stdout is not None

                    chunk = process.stdout.read(4096)

                    chunk = process.stdout.read(settings.rdb_send_chunk_size)

                    if not chunk:

                        process.wait()

                        if process.returncode != 0:
@@ -95,9 +108,6 @@ class FilesysBroker(RemoteDataBroker): 
                        break

                    yield chunk



            headers = {

                "Content-Disposition": f'attachment; filename="{os.path.basename(data.filename)}"'

            }

            return StreamingResponse(

                headers=headers,

                content=generate(),

remote_data_broker/settings.py

+3 −0
Original line number Diff line number Diff line
@@ -6,6 +6,9 @@ class Settings(BaseSettings): 


    rdb_storage_path: str = "/datastore"

    rdb_cat_cmd: str = "cat"

    rdb_fsize_cmd: str = "stat --format=%s"

    rdb_send_filesize: bool = False

    rdb_send_chunk_size: int = 1024 * 1024 * 1024

    rdb_listen_port: int = 8000

    rdb_authorize: bool = False

    keycloak_url: str = ""