Function dump

beanhub_cli/connect/main.py:289–398

(
    env: Environment,
    repo: str | None,
    sync: bool,
    output_accounts: str | None,
    unsafe_tar_extract: bool,
)

@pass_env
@handle_api_exception(logger)
def dump(
    env: Environment,
    repo: str | None,
    sync: bool,
    output_accounts: str | None,
    unsafe_tar_extract: bool,
):
    if not hasattr(tarfile, "data_filter") and not unsafe_tar_extract:
        logger.error(
            "You need to use Python >= 3.11 in order to safely unpack the downloaded tar file, or you need to pass "
            "in --unsafe-tar-extract argument to allow unsafe tar file extracting"
        )
        sys.exit(-1)
    config = ensure_auth_config(api_base_url=env.api_base_url, repo=repo)
    if sync:
        run_sync(env, config)

    private_key = PrivateKey.generate()
    public_key = private_key.public_key.encode(URLSafeBase64Encoder).decode("ascii")

    with make_auth_client(base_url=env.api_base_url, token=config.token) as client:
        client.raise_on_unexpected_status = True
        resp: CreateDumpRequestResponse = create_dump_request.sync(
            body=CreateDumpRequestRequest(
                public_key=public_key, output_accounts=output_accounts is not None
            ),
            username=config.username,
            repo_name=config.repo,
            client=client,
        )
        dump_id = resp.id
        logger.info(
            "Created dump [green]%s[/] with public_key [green]%s[/], waiting for updates ...",
            dump_id,
            public_key,
            extra={"markup": True, "highlighter": None},
        )

        while True:
            time.sleep(5)
            resp: GetDumpRequestResponse = get_dump_request.sync(
                dump_request_id=dump_id,
                username=config.username,
                repo_name=config.repo,
                client=client,
            )
            if resp.state == DumpRequestState.FAILED:
                logger.error("Failed to dump with error: %s", resp.error_message)
                sys.exit(-1)
            elif resp.state == DumpRequestState.COMPLETE:
                break
            else:
                logger.debug("State is %s, keep polling...", resp.state)

    download_url = resp.download_url
    sealed_box = SealedBox(private_key)
    encryption_key = json.loads(
        sealed_box.decrypt(URLSafeBase64Encoder.decode(resp.encryption_key))