def parse_wasbs_url(wasbs_url):
"""
Parses and validates a wasbs url.
Returns:
tuple(container, storage_account, path).
"""
try:
spec = parser.parse_wasbs_path(wasbs_url)
return spec.container, spec.storage_account, spec.path
except PolyaxonSchemaError as e:
raise PolyaxonStoresException("Connection error: %s" % e) from e
def parse_gcs_url(gcs_url):
"""
Parses and validates a google cloud storage url.
Returns:
tuple(bucket_name, blob).
"""
try:
spec = parser.parse_gcs_path(gcs_url)
return spec.bucket, spec.blob
except PolyaxonSchemaError as e:
raise PolyaxonStoresException(e)
def parse_s3_url(s3_url):
"""
Parses and validates an S3 url.
Returns:
tuple(bucket_name, key).
"""
try:
spec = parser.parse_s3_path(s3_url)
return spec.bucket, spec.key
except PolyaxonSchemaError as e:
raise PolyaxonStoresException("Connection error: %s" % e) from e
def get_key(self, key, bucket_name=None):
"""
Returns a boto3.s3.Object
Args:
key: `str`. the path to the key.
bucket_name: `str`. the name of the bucket.
"""
if not bucket_name:
(bucket_name, key) = self.parse_s3_url(key)
try:
obj = self.resource.Object(bucket_name, key)
obj.load()
return obj
except Exception as e:
raise PolyaxonStoresException("Connection error: %s" % e) from e
def upload_file(
self,
filename,
key,
bucket_name=None,
overwrite=True,
encrypt=False,
acl=None,
use_basename=True,
):
"""
Uploads a local file to S3.
Args:
filename: `str`. name of the file to upload.
key: `str`. S3 key that will point to the file.
bucket_name: `str`. Name of the bucket in which to store the file.
overwrite: `bool`. A flag to decide whether or not to overwrite the key
if it already exists. If replace is False and the key exists, an
error will be raised.
encrypt: `bool`. If True, the file will be encrypted on the server-side
by S3 and will be stored in an encrypted form while at rest in S3.
acl: `str`. ACL to use for uploading, e.g. "public-read".
use_basename: `bool`. whether or not to use the basename of the filename.
"""
if not bucket_name:
bucket_name, key = self.parse_s3_url(key)
if use_basename:
key = append_basename(key, filename)
if not overwrite and self.check_key(key, bucket_name):
raise PolyaxonStoresException(
"The key {} already exists.".format(key))
extra_args = {}
if encrypt:
extra_args["ServerSideEncryption"] = self.ENCRYPTION
if acl:
extra_args["ACL"] = acl
self.connection.upload_file(filename,
bucket_name,
key,
ExtraArgs=extra_args)
def list_files(subpath: str,
filepath: str = None,
connection_type: V1ConnectionType = None):
connection_type = connection_type or get_artifacts_connection()
validate_store(connection_type)
store_path = get_path(connection_type.store_path, subpath)
if filepath:
store_path = get_path(store_path, filepath)
store_manager = get_connection_from_type(connection_type=connection_type)
try:
results = store_manager.ls(store_path)
results["files"] = {f[0]: f[1] for f in results["files"]}
return results
except Exception:
raise PolyaxonStoresException(
"Run store path does not exists or bad configuration.")
def get_blob(self, blob, bucket_name=None):
"""
Get a file in Google Cloud Storage.
Args:
blob: `str`. the path to the object to check in the Google cloud storage bucket.
bucket_name: `str`. Name of the bucket in which the file is stored
"""
if not bucket_name:
bucket_name, blob = self.parse_gcs_url(blob)
bucket = self.get_bucket(bucket_name)
# Wrap google.cloud.storage's blob to raise if the file doesn't exist
obj = bucket.get_blob(blob)
if obj is None:
raise PolyaxonStoresException(
"File does not exist: {}".format(blob))
return obj
def validate_store(connection_type: V1ConnectionType):
if not connection_type or not connection_type.is_artifact:
raise PolyaxonStoresException(
"An artifact store type was not provided.")
def _download_blob(blob, local_path):
try:
blob.download_to_filename(local_path)
except (NotFound, GoogleAPIError) as e:
raise PolyaxonStoresException("Connection error: %s" % e) from e
def _delete_blob(blob):
try:
return blob.delete()
except (NotFound, GoogleAPIError) as e:
raise PolyaxonStoresException("Connection error: %s" % e) from e
def check_dirname_exists(path, is_dir=False):
if not is_dir:
path = os.path.dirname(os.path.abspath(path))
if not os.path.isdir(path):
raise PolyaxonStoresException(
"The parent path is not a directory {}".format(path))
