Skip to content

Commit

Permalink
Downloading file to target directory
Browse files Browse the repository at this point in the history
  • Loading branch information
kkalinowski-reef committed Dec 6, 2023
1 parent ea82734 commit c777913
Show file tree
Hide file tree
Showing 4 changed files with 144 additions and 4 deletions.
50 changes: 46 additions & 4 deletions b2/console_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@
######################################################################
from __future__ import annotations

import tempfile

from b2._cli.autocomplete_cache import AUTOCOMPLETE # noqa

AUTOCOMPLETE.autocomplete_from_cache()
Expand Down Expand Up @@ -1606,10 +1608,50 @@ def _represent_legal_hold(cls, legal_hold: LegalHold):
def _print_file_attribute(self, label, value):
self._print((label + ':').ljust(20) + ' ' + value)

def get_local_output_filepath(self, filename: str) -> pathlib.Path:
def get_local_output_filepath(
self, filename: str, file_request: DownloadedFile
) -> pathlib.Path:
if filename == '-':
return STDOUT_FILEPATH
return pathlib.Path(filename)

output_filepath = pathlib.Path(filename)

# As longs as it's not a directory, we're overwriting everything.
if not output_filepath.is_dir():
return output_filepath

# If the output is directory, we're expected to download the file right there.
# Normally, we overwrite the target without asking any questions, but in this case
# user might be oblivious of the actual mistake he's about to commit.
# If he, e.g.: downloads file by ID, he might not know the name of the file
# and actually overwrite something unintended.
output_directory = output_filepath
output_filepath = output_directory / file_request.download_version.file_name
# If it doesn't exist, we stop worrying.
if not output_filepath.exists():
return output_filepath

# If it does exist, we make a unique file prefixed with the actual file name.
file_name_as_path = pathlib.Path(file_request.download_version.file_name)
file_name = file_name_as_path.stem
file_extension = file_name_as_path.suffix

# Default permissions are: readable and writable by this user only, executable by noone.
# This "temporary" file is not automatically removed, but still created in the safest way possible.
fd_handle, output_filepath_str = tempfile.mkstemp(
prefix=file_name,
suffix=file_extension,
dir=output_directory,
)
# Close the handle, so the file is not locked.
# This file is no longer 100% "safe", but that's acceptable.
os.close(fd_handle)

# "Normal" file created by Python has readable for everyone, writable for user only.
# We change the permissions, to match the default ones.
os.chmod(output_filepath_str, 0o644)

return pathlib.Path(output_filepath_str)


class DownloadFileBase(
Expand Down Expand Up @@ -1645,7 +1687,7 @@ def _run(self, args):
)

self._print_download_info(downloaded_file)
output_filepath = self.get_local_output_filepath(args.localFileName)
output_filepath = self.get_local_output_filepath(args.localFileName, downloaded_file)
downloaded_file.save_to(output_filepath)
self._print('Download finished')

Expand Down Expand Up @@ -1711,7 +1753,7 @@ def _run(self, args):
file_request = self.api.download_file_by_uri(
args.B2_URI, progress_listener=progress_listener, encryption=encryption_setting
)
output_filepath = self.get_local_output_filepath(target_filename)
output_filepath = self.get_local_output_filepath(target_filename, file_request)
file_request.save_to(output_filepath)
return 0

Expand Down
1 change: 1 addition & 0 deletions changelog.d/+downloading_to_directory.added.md
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Whenever target filename is a directory, file is downloaded into that directory.
41 changes: 41 additions & 0 deletions test/integration/test_b2_command_line.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
import json
import os
import os.path
import pathlib
import re
import sys
import time
Expand Down Expand Up @@ -2682,6 +2683,46 @@ def test_download_file_stdout(
).replace("\r", "") == sample_filepath.read_text()


def test_download_file_to_directory(
b2_tool, bucket_name, sample_filepath, tmp_path, uploaded_sample_file
):
downloads_directory = 'downloads'
target_directory = tmp_path / downloads_directory
target_directory.mkdir()
filename_as_path = pathlib.Path(uploaded_sample_file['fileName'])

sample_file_content = sample_filepath.read_text()
b2_tool.should_succeed(
[
'download-file',
'--quiet',
f"b2://{bucket_name}/{uploaded_sample_file['fileName']}",
str(target_directory),
],
)
downloaded_file = target_directory / filename_as_path
assert downloaded_file.read_text() == sample_file_content, \
f'{downloaded_file}, {downloaded_file.read_text()}, {sample_file_content}'

b2_tool.should_succeed(
[
'download-file',
'--quiet',
f"b2id://{uploaded_sample_file['fileId']}",
str(target_directory),
],
)
# A second file should be created.
new_files = [
filepath
for filepath in target_directory.glob(f'{filename_as_path.stem}*{filename_as_path.suffix}')
if filepath.name != filename_as_path.name
]
assert len(new_files) == 1, f'{new_files}'
assert new_files[0].read_text() == sample_file_content, \
f'{new_files}, {new_files[0].read_text()}, {sample_file_content}'


def test_cat(b2_tool, bucket_name, sample_filepath, tmp_path, uploaded_sample_file):
assert b2_tool.should_succeed(
['cat', f"b2://{bucket_name}/{uploaded_sample_file['fileName']}"],
Expand Down
56 changes: 56 additions & 0 deletions test/unit/test_console_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -1117,6 +1117,62 @@ def test_download_by_name_1_thread(self):
def test_download_by_name_10_threads(self):
self._test_download_threads(download_by='name', num_threads=10)

def _test_download_to_directory(self, download_by):
self._authorize_account()
self._create_my_bucket()

base_filename = 'file'
extension = '.txt'
source_filename = f'{base_filename}{extension}'

with TempDir() as temp_dir:
local_file = self._make_local_file(temp_dir, source_filename)
local_file_content = self._read_file(local_file)

self._run_command(
['upload-file', '--noProgress', 'my-bucket', local_file, source_filename],
remove_version=True,
)

command = [
'download-file-by-%s' % download_by,
'--noProgress',
]
command += ['9999'] if download_by == 'id' else ['my-bucket', source_filename]
target_directory = os.path.join(temp_dir, 'target')
os.mkdir(target_directory)
command += [target_directory]
self._run_command(
command,
expected_stderr=
f'WARNING: download-file-by-{download_by} command is deprecated. Use download-file instead.\n'
)
self.assertEqual(
local_file_content,
self._read_file(os.path.join(target_directory, source_filename))
)

# Download the file second time, to check the override behavior.
self._run_command(
command,
expected_stderr=
f'WARNING: download-file-by-{download_by} command is deprecated. Use download-file instead.\n'
)
# We should get another file.
target_directory_files = [
elem
for elem in pathlib.Path(target_directory).glob(f'{base_filename}*{extension}')
if elem.name != source_filename
]
assert len(target_directory_files) == 1, f'{target_directory_files}'
self.assertEqual(local_file_content, self._read_file(target_directory_files[0]))

def test_download_by_id_to_directory(self):
self._test_download_to_directory(download_by='id')

def test_download_by_name_to_directory(self):
self._test_download_to_directory(download_by='name')

def test_copy_file_by_id(self):
self._authorize_account()
self._create_my_bucket()
Expand Down

0 comments on commit c777913

Please sign in to comment.