def test_key_and_connection(self):
from azure.storage.file import FileService
hook = AzureFileShareHook(wasb_conn_id='wasb_test_key')
self.assertEqual(hook.conn_id, 'wasb_test_key')
self.assertIsNone(hook._conn)
self.assertIsInstance(hook.get_conn(), FileService)
def test_get_file_to_stream(self, mock_service):
mock_instance = mock_service.return_value
hook = AzureFileShareHook(wasb_conn_id='wasb_test_sas_token')
hook.get_file_to_stream('stream', 'share', 'directory', 'file', max_connections=1)
mock_instance.get_file_to_stream.assert_called_once_with(
'share', 'directory', 'file', 'stream', max_connections=1
)
def test_load_string(self, mock_service):
mock_instance = mock_service.return_value
hook = AzureFileShareHook(wasb_conn_id='wasb_test_sas_token')
hook.load_string('big string', 'share', 'directory', 'file', timeout=1)
mock_instance.create_file_from_text.assert_called_once_with(
'share', 'directory', 'file', 'big string', timeout=1
)
def test_load_file(self, mock_service):
mock_instance = mock_service.return_value
hook = AzureFileShareHook(wasb_conn_id='wasb_test_sas_token')
hook.load_file('path', 'share', 'directory', 'file', max_connections=1)
mock_instance.create_file_from_path.assert_called_once_with(
'share', 'directory', 'file', 'path', max_connections=1
)
def test_create_directory(self, mock_service):
mock_instance = mock_service.return_value
hook = AzureFileShareHook(wasb_conn_id='wasb_test_sas_token')
hook.create_directory('share', 'directory', timeout=1)
mock_instance.create_directory.assert_called_once_with('share',
'directory',
timeout=1)
def test_key_and_connection(self):
from azure.storage.file import FileService
hook = AzureFileShareHook(wasb_conn_id='wasb_test_key')
assert hook.conn_id == 'wasb_test_key'
assert hook._conn is None
assert isinstance(hook.get_conn(), FileService)
def test_list_directories_and_files(self, mock_service):
mock_instance = mock_service.return_value
hook = AzureFileShareHook(wasb_conn_id='wasb_test_sas_token')
hook.list_directories_and_files('share', 'directory', timeout=1)
mock_instance.list_directories_and_files.assert_called_once_with(
'share', 'directory', timeout=1
)
def test_check_for_directory(self, mock_service):
mock_instance = mock_service.return_value
mock_instance.exists.return_value = True
hook = AzureFileShareHook(wasb_conn_id='wasb_test_sas_token')
assert hook.check_for_directory('share', 'directory', timeout=3)
mock_instance.exists.assert_called_once_with('share',
'directory',
timeout=3)
def test_list_files(self, mock_service):
mock_instance = mock_service.return_value
mock_instance.list_directories_and_files.return_value = [
File("file1"),
File("file2"),
Directory("dir1"),
Directory("dir2"),
]
hook = AzureFileShareHook(wasb_conn_id='wasb_test_sas_token')
files = hook.list_files('share', 'directory', timeout=1)
self.assertEqual(files, ["file1", 'file2'])
mock_instance.list_directories_and_files.assert_called_once_with('share', 'directory', timeout=1)
def upload_file_from_string(cls,
string_data: str,
share_name: str,
wasb_conn_id: str,
file_name: str,
directory: Optional[str] = None):
hook = AzureFileShareHook(wasb_conn_id=wasb_conn_id)
hook.load_string(
string_data=string_data,
share_name=share_name,
directory_name=directory,
file_name=file_name,
)
def create_fileshare():
"""Create a fileshare with directory"""
hook = AzureFileShareHook()
hook.create_share(NAME)
hook.create_directory(share_name=NAME, directory_name=DIRECTORY)
exists = hook.check_for_directory(share_name=NAME,
directory_name=DIRECTORY)
if not exists:
raise Exception
def test_sas_token(self):
from azure.storage.file import FileService
hook = AzureFileShareHook(wasb_conn_id='wasb_test_sas_token')
self.assertEqual(hook.conn_id, 'wasb_test_sas_token')
self.assertIsInstance(hook.get_conn(), FileService)
def test_sas_token(self):
from azure.storage.file import FileService
hook = AzureFileShareHook(wasb_conn_id='wasb_test_sas_token')
assert hook.conn_id == 'wasb_test_sas_token'
assert isinstance(hook.get_conn(), FileService)
def execute(self, context):
azure_fileshare_hook = AzureFileShareHook(self.wasb_conn_id)
files = azure_fileshare_hook.list_files(
share_name=self.share_name, directory_name=self.directory_name)
gcs_hook = GCSHook(
gcp_conn_id=self.gcp_conn_id,
delegate_to=self.delegate_to,
impersonation_chain=self.google_impersonation_chain,
)
dest_gcs_bucket, dest_gcs_object_prefix = _parse_gcs_url(self.dest_gcs)
# pylint: disable=too-many-nested-blocks
if not self.replace:
# if we are not replacing -> list all files in the GCS bucket
# and only keep those files which are present in
# S3 and not in Google Cloud Storage
existing_files_prefixed = gcs_hook.list(
dest_gcs_bucket, prefix=dest_gcs_object_prefix)
existing_files = []
# Remove the object prefix itself, an empty directory was found
if dest_gcs_object_prefix in existing_files_prefixed:
existing_files_prefixed.remove(dest_gcs_object_prefix)
# Remove the object prefix from all object string paths
for file in existing_files_prefixed:
if file.startswith(dest_gcs_object_prefix):
existing_files.append(file[len(dest_gcs_object_prefix):])
else:
existing_files.append(file)
files = list(set(files) - set(existing_files))
if files:
self.log.info('%s files are going to be synced.', len(files))
else:
self.log.info('There are no new files to sync. Have a nice day!')
for file in files:
with NamedTemporaryFile() as temp_file:
azure_fileshare_hook.get_file_to_stream(
stream=temp_file,
share_name=self.share_name,
directory_name=self.directory_name,
file_name=file,
)
temp_file.flush()
# There will always be a '/' before file because it is
# enforced at instantiation time
dest_gcs_object = dest_gcs_object_prefix + file
gcs_hook.upload(dest_gcs_bucket,
dest_gcs_object,
temp_file.name,
gzip=self.gzip)
if files:
self.log.info(
"All done, uploaded %d files to Google Cloud Storage.",
len(files))
else:
self.log.info(
'In sync, no files needed to be uploaded to Google Cloud Storage'
)
return files
def delete_share(cls, share_name: str, wasb_conn_id: str):
hook = AzureFileShareHook(wasb_conn_id=wasb_conn_id)
hook.delete_share(share_name)
def create_directory(cls, share_name: str, wasb_conn_id: str, directory: str):
hook = AzureFileShareHook(wasb_conn_id=wasb_conn_id)
hook.create_directory(share_name=share_name, directory_name=directory)
def test_delete_share(self, mock_service):
mock_instance = mock_service.return_value
hook = AzureFileShareHook(wasb_conn_id='wasb_test_sas_token')
hook.delete_share('my_share')
mock_instance.delete_share.assert_called_once_with('my_share')
def delete_fileshare():
"""Delete a fileshare"""
hook = AzureFileShareHook()
hook.delete_share(NAME)