Exemple #1
0
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,
    )
Exemple #2
0
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
Exemple #3
0
 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)
Exemple #4
0
    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,
        )
Exemple #5
0
    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,
        )
Exemple #6
0
    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}"'
            },
        )
Exemple #7
0
    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,
        )
Exemple #8
0
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,
    )
Exemple #9
0
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()
Exemple #10
0
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"
    )