async def create_blob(
*,
db: Session = Depends(get_db),
ensemble_id: UUID,
name: str,
realization_index: Optional[int] = None,
) -> None:
"""
Create a record which points to a blob on Azure Blob Storage.
"""
ensemble = db.query(ds.Ensemble).filter_by(id=ensemble_id).one()
file_obj = ds.File(
filename="test",
mimetype="mime/type",
)
if HAS_AZURE_BLOB_STORAGE:
key = f"{name}@{realization_index}@{uuid4()}"
blob = azure_blob_container.get_blob_client(key)
file_obj.az_container = (azure_blob_container.container_name, )
file_obj.az_blob = (key, )
else:
pass
db.add(file_obj)
_create_record(
db,
ensemble,
name,
ds.RecordType.file,
realization_index=realization_index,
file=file_obj,
)
async def finalize_blob(
*,
db: Session = Depends(get_db),
ensemble_id: UUID,
name: str,
realization_index: Optional[int] = None,
) -> None:
"""
Commit all staged blocks to a blob record
"""
ensemble = db.query(ds.Ensemble).filter_by(id=ensemble_id).one()
record_obj = (db.query(
ds.Record).filter_by(realization_index=realization_index).join(
ds.RecordInfo).filter_by(ensemble_pk=ensemble.pk, name=name).one())
submitted_blocks = list(
db.query(ds.FileBlock).filter_by(
record_name=name,
ensemble_pk=ensemble.pk,
realization_index=realization_index,
).all())
if HAS_AZURE_BLOB_STORAGE:
key = record_obj.file.az_blob
blob = azure_blob_container.get_blob_client(key)
block_ids = [
block.block_id
for block in sorted(submitted_blocks, key=lambda x: x.block_index)
]
await blob.commit_block_list(block_ids)
else:
data = b"".join([block.content for block in submitted_blocks])
record_obj.file.content = data
async def finalize_blob(self, submitted_blocks: List[ds.FileBlock],
record: ds.Record) -> None:
blob = azure_blob_container.get_blob_client(record.file.az_blob)
block_ids = [
block.block_id
for block in sorted(submitted_blocks, key=lambda x: x.block_index)
]
await blob.commit_block_list(block_ids)
def create_blob(self) -> ds.File:
key = f"{self._name}@{self._realization_index}@{uuid4()}"
blob = azure_blob_container.get_blob_client(key)
return ds.File(
filename="test",
mimetype="mime/type",
az_container=azure_blob_container.container_name,
az_blob=key,
)
async def upload_file(
self,
file: UploadFile,
) -> ds.File:
key = f"{self._name}@{self._realization_index}@{uuid4()}"
blob = azure_blob_container.get_blob_client(key)
await blob.upload_blob(file.file)
return ds.File(
filename=file.filename,
mimetype=file.content_type,
az_container=azure_blob_container.container_name,
az_blob=key,
)
async def get_content(self, record: ds.Record) -> Response:
blob = azure_blob_container.get_blob_client(record.file.az_blob)
download = await blob.download_blob()
async def chunk_generator() -> AsyncGenerator[bytes, None]:
async for chunk in download.chunks():
yield chunk
return StreamingResponse(
chunk_generator(),
media_type=record.file.mimetype,
headers={
"Content-Disposition":
f'attachment; filename="{record.file.filename}"'
},
)
async def stage_blob(
self,
record: ds.Record,
request: Request,
block_index: int,
) -> ds.FileBlock:
block_id = str(uuid4())
blob = azure_blob_container.get_blob_client(record.file.az_blob)
await blob.stage_block(block_id, await request.body())
return ds.FileBlock(
ensemble_pk=record.ensemble_pk,
block_id=block_id,
block_index=block_index,
record_name=self._name,
realization_index=self._realization_index,
)
async def post_ensemble_record_file(
*,
db: Session = Depends(get_db),
ensemble_id: UUID,
name: str,
realization_index: Optional[int] = None,
file: UploadFile = File(...),
) -> None:
"""
Assign an arbitrary file to the given `name` record.
"""
ensemble = _get_and_assert_ensemble(db, ensemble_id, name,
realization_index)
file_obj = ds.File(
filename=file.filename,
mimetype=file.content_type,
)
if HAS_AZURE_BLOB_STORAGE:
key = f"{name}@{realization_index}@{uuid4()}"
blob = azure_blob_container.get_blob_client(key)
await blob.upload_blob(file.file)
file_obj.az_container = azure_blob_container.container_name
file_obj.az_blob = key
else:
file_obj.content = await file.read()
db.add(file_obj)
_create_record(
db,
ensemble,
name,
ds.RecordType.file,
realization_index=realization_index,
file=file_obj,
)
async def add_block(
*,
db: Session = Depends(get_db),
ensemble_id: UUID,
name: str,
realization_index: Optional[int] = None,
request: Request,
block_index: int,
) -> None:
"""
Stage blocks to an existing azure blob record.
"""
ensemble = db.query(ds.Ensemble).filter_by(id=ensemble_id).one()
block_id = str(uuid4())
file_block_obj = ds.FileBlock(
ensemble=ensemble,
block_id=block_id,
block_index=block_index,
record_name=name,
realization_index=realization_index,
)
record_obj = (db.query(
ds.Record).filter_by(realization_index=realization_index).join(
ds.RecordInfo).filter_by(ensemble_pk=ensemble.pk, name=name).one())
if HAS_AZURE_BLOB_STORAGE:
key = record_obj.file.az_blob
blob = azure_blob_container.get_blob_client(key)
await blob.stage_block(block_id, await request.body())
else:
file_block_obj.content = await request.body()
db.add(file_block_obj)
db.commit()
async def _get_record_data(
record: ds.Record,
accept: Optional[str],
realization_index: Optional[int] = None) -> Response:
type_ = record.record_info.record_type
if type_ == ds.RecordType.f64_matrix:
if realization_index is None:
content = record.f64_matrix.content
else:
content = record.f64_matrix.content[realization_index]
if accept == "application/x-numpy":
from numpy.lib.format import write_array
stream = io.BytesIO()
write_array(stream, np.array(content))
return Response(
content=stream.getvalue(),
media_type="application/x-numpy",
)
if accept == "text/csv":
data = pd.DataFrame(content)
labels = record.f64_matrix.labels
if labels is not None and realization_index is None:
data.columns = labels[0]
data.index = labels[1]
elif labels is not None and realization_index is not None:
# The output is such that rows are realizations. Because
# `content` is a 1d list in this case, it treats each element as
# its own row. We transpose the data so that all of the data
# falls on the same row.
data = data.T
data.columns = labels[0]
data.index = [realization_index]
return Response(
content=data.to_csv().encode(),
media_type="text/csv",
)
else:
return content
if type_ == ds.RecordType.file:
f = record.file
if f.content is not None:
return Response(
content=f.content,
media_type=f.mimetype,
headers={
"Content-Disposition":
f'attachment; filename="{f.filename}"'
},
)
elif f.az_container is not None and f.az_blob is not None:
blob = azure_blob_container.get_blob_client(f.az_blob)
download = await blob.download_blob()
async def chunk_generator() -> AsyncGenerator[bytes, None]:
async for chunk in download.chunks():
yield chunk
return StreamingResponse(
chunk_generator(),
media_type=f.mimetype,
headers={
"Content-Disposition":
f'attachment; filename="{f.filename}"'
},
)
raise NotImplementedError(
f"Getting record data for type {type_} and Accept header {accept} not implemented"
)
