import re
from os.path import normpath, join, isabs
import shutil
from glob import glob
from fnmatch import fnmatch
import logging
import boto3
from swiftclient.multithreading import OutputManager
from swiftclient.service import SwiftError, SwiftService
logger = logging.getLogger(__name__)
class RegistrationError(Exception):
pass
class Source:
def __init__(self, name: str=None):
self.name = name
def list_files(self, path, glob_pattern=None):
raise NotImplementedError
def get_file(self, path, target_path):
raise NotImplementedError
def get_vsi_env_and_path(self, path):
raise NotImplementedError
class SwiftSource(Source):
def __init__(self, name=None, username=None, password=None, tenant_name=None,
tenant_id=None, region_name=None, user_domain_id=None,
user_domain_name=None, auth_url=None, auth_version=None,
container=None):
super().__init__(name)
self.username = username
self.password = password
self.tenant_name = tenant_name
self.tenant_id = tenant_id
self.region_name = region_name
self.user_domain_id = user_domain_id
self.user_domain_name = user_domain_name
self.auth_url = auth_url
self.auth_version = auth_version # TODO: assume 3
self.container = container
def get_service(self):
return SwiftService(options={
"os_username": self.username,
"os_password": self.password,
"os_tenant_name": self.tenant_name,
"os_tenant_id": self.tenant_id,
"os_region_name": self.region_name,
"os_auth_url": self.auth_url,
"auth_version": self.auth_version,
"os_user_domain_id": self.user_domain_id,
"os_user_domain_name": self.user_domain_name,
})
def get_container_and_path(self, path: str):
container = self.container
if container is None:
parts = (path[1:] if path.startswith('/') else path).split('/')
container, path = parts[0], '/'.join(parts[1:])
return container, path
def list_files(self, path, glob_pattern=None):
container, path = self.get_container_and_path(path)
with self.get_service() as swift:
pages = swift.list(
container=container,
options={"prefix": path},
)
filenames = []
for page in pages:
if page["success"]:
# at least two files present -> pass validation
for item in page["listing"]:
if glob_pattern is None or fnmatch(item['name'], glob_pattern):
filenames.append(item['name'])
else:
raise page['error']
return filenames
def get_file(self, path, target_path):
container, path = self.get_container_and_path(path)
with self.get_service() as swift:
results = swift.download(
container,
[path],
options={
'out_file': target_path
}
)
for result in results:
if not result["success"]:
raise Exception('Failed to download %s' % path)
def get_vsi_env_and_path(self, path):
container, path = self.get_container_and_path(path)
return {
'OS_IDENTITY_API_VERSION': self.auth_version,
'OS_AUTH_URL': self.auth_url,
'OS_USERNAME': self.username,
'OS_PASSWORD': self.password,
'OS_USER_DOMAIN_NAME': self.user_domain_name,
# 'OS_PROJECT_NAME': self.tena,
# 'OS_PROJECT_DOMAIN_NAME': ,
'OS_REGION_NAME': self.region_name,
}, f'/vsiswift/{container}/{path}'
class S3Source(Source):
def __init__(self, name=None, bucket_name=None, secret_access_key=None, access_key_id=None, endpoint_url=None, strip_bucket=True, **client_kwargs):
super().__init__(name)
# see https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html#boto3.session.Session.client
# for client_kwargs
self.bucket_name = bucket_name
self.secret_access_key=secret_access_key
self.access_key_id=access_key_id
self.endpoint_url = endpoint_url
self.strip_bucket = strip_bucket
self.client = boto3.client(
's3',
aws_secret_access_key=secret_access_key,
aws_access_key_id=access_key_id,
endpoint_url=endpoint_url,
**client_kwargs,
)
def get_bucket_and_key(self, path: str):
bucket = self.bucket_name
if bucket is None:
parts = (path[1:] if path.startswith('/') else path).split('/')
bucket, path = parts[0], '/'.join(parts[1:])
elif self.strip_bucket:
parts = (path[1:] if path.startswith('/') else path).split('/')
if parts[0] == bucket:
parts.pop(0)
path = '/'.join(parts)
return bucket, path
def list_files(self, path, glob_pattern=None):
bucket, key = self.get_bucket_and_key(path)
logger.info(f'Listing S3 files for bucket {bucket} and prefix {key}')
response = self.client.list_objects_v2(
Bucket=bucket,
Prefix=key,
)
return [
f"{bucket}/{item['Key']}"
for item in response['Contents']
if glob_pattern is None or fnmatch(item['Key'], glob_pattern)
]
def get_file(self, path, target_path):
bucket, key = self.get_bucket_and_key(path)
logger.info(f'Retrieving file from S3 {bucket}/{key} to be stored at {target_path}')
self.client.download_file(bucket, key, target_path)
def get_vsi_env_and_path(self, path: str, streaming: bool=False):
bucket, key = self.get_bucket_and_key(path)
return {
'AWS_SECRET_ACCESS_KEY': self.secret_access_key,
'AWS_ACCESS_KEY_ID': self.access_key_id,
'AWS_S3_ENDPOINT': self.endpoint_url,
}, f'/{"vsis3" if not streaming else "vsis3_streaming"}/{bucket}/{key}'
class LocalSource(Source):
def __init__(self, name, root_directory):
super().__init__(name)
self.root_directory = root_directory
def _join_path(self, path):
path = normpath(path)
if isabs(path):
path = path[1:]
return join(self.root_directory, path)
def list_files(self, path, glob_pattern=None):
if glob_pattern is not None:
return glob(join(self._join_path(path), glob_pattern))
else:
return glob(join(self._join_path(path), '*'))
def get_file(self, path, target_path):
shutil.copy(self._join_path(path), target_path)
def get_vsi_env_and_path(self, path):
return {}, self._join_path(path)
SOURCE_TYPES = {
'swift': SwiftSource,
's3': S3Source,
'local': LocalSource,
}
def get_source(config: dict, path: str) -> Source:
cfg_sources = config['sources']
for cfg_source in cfg_sources:
if cfg_source['filter']:
if re.match(cfg_source['filter'], path):
break
else:
break
else:
# no source found
raise RegistrationError(f'Could not find a suitable source for the path {path}')
return SOURCE_TYPES[cfg_source['type']](
cfg_source['name'],
*cfg_source.get('args', []),
**cfg_source.get('kwargs', {})
)