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))