def copy_from_dbfs(dbfs_path_src, dst, overwrite):
# Will truncate the local file if it exists.
try:
get_file(dbfs_path_src, dst, overwrite)
except LocalFileExistsException:
click.echo(
('The host destination {} already exists. You should provide the '
'--overwrite flag.').format(dst))
def test_get_file(tmpdir):
with mock.patch(
'databricks_cli.dbfs.api.get_dbfs_client') as get_dbfs_client:
api_mock = get_dbfs_client.return_value
api_mock.get_status.return_value = TEST_FILE_JSON
api_mock.read.return_value = {
'bytes_read': 1,
'data': b64encode('x'),
}
test_file_path = os.path.join(tmpdir.strpath, 'test')
api.get_file(TEST_DBFS_PATH, test_file_path, True)
with open(test_file_path, 'r') as f:
assert f.read() == 'x'
def copy_from_dbfs_recursive(dbfs_path_src, dst, overwrite):
if os.path.isfile(dst):
click.echo('{} exists as a file. Skipping this subtree {}'.format(
dst, repr(dbfs_path_src)))
return
elif not os.path.isdir(dst):
os.makedirs(dst)
for dbfs_src_file_info in list_files(dbfs_path_src):
cur_dbfs_src = dbfs_src_file_info.dbfs_path
cur_dst = os.path.join(dst, cur_dbfs_src.basename)
if dbfs_src_file_info.is_dir:
copy_from_dbfs_recursive(cur_dbfs_src, cur_dst, overwrite)
else:
try:
get_file(cur_dbfs_src, cur_dst, overwrite)
click.echo('{} -> {}'.format(cur_dbfs_src, cur_dst))
except LocalFileExistsException:
click.echo(
('{} already exists locally as {}. Skip. To overwrite, you'
+ 'should provide the --overwrite flag.').format(
cur_dbfs_src, cur_dst))
def test_get_file_check_overwrite(tmpdir):
test_file_path = os.path.join(tmpdir.strpath, 'test')
with open(test_file_path, 'w') as f:
f.write('test')
with pytest.raises(LocalFileExistsException):
api.get_file(TEST_DBFS_PATH, test_file_path, False)
def copy_from_dbfs_non_recursive(dbfs_path_src, dst, overwrite):
# Munge dst path in case dst is a dir
if os.path.isdir(dst):
dst = os.path.join(dst, dbfs_path_src.basename)
get_file(dbfs_path_src, dst, overwrite)