spiral-sdk 0.1.7

//! A table write that uploads large blobs before the rows that reference them, and edits its
//! base version's rows by row ID.

use std::sync::Arc;

use arrow_schema::{DataType, Schema, SchemaRef};
use futures::future::BoxFuture;
use uuid::Uuid;

use crate::client::Connection;
use crate::registration::register;
use crate::scandal::RecordBatchStream;
use crate::scandal::types::FragmentReceipt;
use crate::scandal_backend::{AttemptWrite, RowEdit, reader_stream};
use crate::{Blob, BlobPack, BlobStream, Error, IntoArrowStream, Project, Table, TableIdentity};

/// The column that names rows by row ID: in the rows that [`TableWrite::patch_rows`] and
/// [`TableWrite::delete_rows`] take, and in the scans that the SDK's own edits and change reads
/// open. It is a non-null `uint64` column.
pub const ROW_ID_COLUMN: &str = "s:row_id";

/// One write of a new table's first version or of a table's next version, under its own write
/// attempt. Large payloads upload into the attempt first, as blobs; rows appended afterwards
/// carry them, encoded with [`Blob::array`]. A write of a table's next version can also patch
/// and delete the base version's rows by row ID. [`TableWrite::finish`] returns the
/// written version, one version for every part, which is published only once a
/// [`Commit`](crate::Commit) stages it and publishes.
pub struct TableWrite {
    /// The table this write continues, or `None` for a new table.
    base: Option<Table>,
    connection: Arc<Connection>,
    write: AttemptWrite,
}

impl Project {
    /// Begin the first version of a new table with rows of `schema`.
    pub async fn table_writer(&self, schema: SchemaRef) -> Result<TableWrite, Error> {
        let write = self
            .client
            .connection
            .tables()?
            .begin_create(self.id.as_str(), &schema)
            .await?;
        Ok(TableWrite {
            base: None,
            connection: Arc::clone(&self.client.connection),
            write,
        })
    }
}

impl Table {
    /// Begin the next version of this table: its rows, edited by the patches and deletions the
    /// write takes, followed by every appended row.
    pub async fn writer(&self) -> Result<TableWrite, Error> {
        let write = self
            .connection
            .tables()?
            .begin_write(self.identity(), false)
            .await?;
        Ok(TableWrite {
            base: Some(self.clone()),
            connection: Arc::clone(&self.connection),
            write,
        })
    }
}

impl TableWrite {
    /// The write attempt whose directory holds this write's files and uploads.
    pub fn attempt(&self) -> Uuid {
        self.write.attempt()
    }

    /// Upload `input` as one blob of this write, however large, in parts of 64 MiB sent several
    /// at once. Each part is retried on its own, and a failed upload is aborted. The blob's
    /// digest is the BLAKE3 digest of the whole payload.
    pub async fn upload_blob(&self, media_type: &str, input: BlobStream) -> Result<Blob, Error> {
        self.write.upload_blob(media_type, input).await
    }

    /// Begin a pack of payloads to store as blobs of this write, packed together as the inline
    /// payloads of appended rows are.
    pub fn blob_pack(&self) -> BlobPack<'_> {
        BlobPack::new(&self.write)
    }

    /// Append rows, which may carry blobs this write uploaded.
    pub async fn append(&mut self, input: impl IntoArrowStream) -> Result<(), Error> {
        self.write.append(input.into_write_input()?).await
    }

    /// Record new values for existing rows beside the data: each row names its target in a
    /// [`ROW_ID_COLUMN`] column, and its other columns are columns of the table, with their
    /// stored types, whose values it replaces. A row ID appears at most once per write.
    pub async fn patch_rows(&mut self, rows: impl IntoArrowStream) -> Result<(), Error> {
        self.edit_input(RowEdit::Patch, rows).await
    }

    /// Delete existing rows, named by their row IDs in a [`ROW_ID_COLUMN`] column, which is the
    /// input's only column. Each row must be visible in the base version.
    pub async fn delete_rows(&mut self, ids: impl IntoArrowStream) -> Result<(), Error> {
        self.edit_input(RowEdit::Delete, ids).await
    }

    /// Prepare the version from every part this write took and register it.
    pub async fn finish(self) -> Result<Table, Error> {
        let result = register(&self.connection.arbor, self.write.prepare().await?).await?;
        match self.base {
            Some(base) => base.successor(result),
            None => Ok(Table::new(
                self.connection,
                TableIdentity {
                    resource_id: result.resource_id,
                    version_id: result.version_id,
                },
            )),
        }
    }

    async fn edit_input(
        &mut self,
        edit: RowEdit,
        input: impl IntoArrowStream,
    ) -> Result<(), Error> {
        let reader = input.into_arrow_stream()?;
        self.edit_rows(edit, reader.schema(), reader_stream(reader))
            .await
    }

    /// Write `input` as one part of `edit` and keep its receipts.
    pub(crate) async fn edit_rows(
        &mut self,
        edit: RowEdit,
        schema: SchemaRef,
        input: RecordBatchStream,
    ) -> Result<(), Error> {
        let receipts = self.start_edit(edit, schema, input)?.await?;
        self.write.add_receipts(receipts);
        Ok(())
    }

    /// Start writing `input` as one part of `edit`, beside other parts. Keep the receipts the
    /// request returns with [`Self::add_receipts`].
    pub(crate) fn start_edit(
        &mut self,
        edit: RowEdit,
        schema: SchemaRef,
        input: RecordBatchStream,
    ) -> Result<BoxFuture<'static, Result<Vec<FragmentReceipt>, Error>>, Error> {
        if self.base.is_none() {
            return Err(Error::invalid("a new table has no rows to patch or delete"));
        }
        check_edit_schema(edit, &schema)?;
        self.write.edit_base()?;
        Ok(self.write.edit_rows(edit, ROW_ID_COLUMN, schema, input))
    }

    /// Keep the receipts of a part started with [`Self::start_edit`].
    pub(crate) fn add_receipts(&mut self, receipts: Vec<FragmentReceipt>) {
        self.write.add_receipts(receipts);
    }
}

/// Check that rows of `edit` name their targets in a `uint64` [`ROW_ID_COLUMN`] and carry the
/// columns the edit needs: values for patches, nothing else for deletions.
pub(crate) fn check_edit_schema(edit: RowEdit, schema: &Schema) -> Result<(), Error> {
    let row_ids = schema.field_with_name(ROW_ID_COLUMN).map_err(|_| {
        Error::invalid(format!(
            "rows that address existing rows need a `{ROW_ID_COLUMN}` column"
        ))
    })?;
    if row_ids.data_type() != &DataType::UInt64 {
        return Err(Error::invalid(format!(
            "`{ROW_ID_COLUMN}` holds uint64 row IDs, not {}",
            row_ids.data_type()
        )));
    }
    let values = schema.fields().len() - 1;
    match edit {
        RowEdit::Delete if values > 0 => Err(Error::invalid(format!(
            "deletions hold only the `{ROW_ID_COLUMN}` column"
        ))),
        RowEdit::Patch if values == 0 => Err(Error::invalid(
            "patched rows hold at least one column besides their row IDs",
        )),
        _ => Ok(()),
    }
}

#[cfg(test)]
mod tests {
    use arrow_schema::Field;

    use super::*;

    fn schema(fields: &[(&str, DataType)]) -> Schema {
        Schema::new(
            fields
                .iter()
                .map(|(name, data_type)| Field::new(*name, data_type.clone(), false))
                .collect::<Vec<_>>(),
        )
    }

    #[test]
    fn edit_rows_name_their_targets_and_carry_the_columns_their_edit_needs() {
        let ids = schema(&[(ROW_ID_COLUMN, DataType::UInt64)]);
        let values = schema(&[(ROW_ID_COLUMN, DataType::UInt64), ("label", DataType::Utf8)]);
        assert!(check_edit_schema(RowEdit::Delete, &ids).is_ok());
        assert!(check_edit_schema(RowEdit::Delete, &values).is_err());
        assert!(check_edit_schema(RowEdit::Patch, &values).is_ok());
        assert!(check_edit_schema(RowEdit::Patch, &ids).is_err());
        let signed = schema(&[(ROW_ID_COLUMN, DataType::Int64)]);
        assert!(check_edit_schema(RowEdit::Delete, &signed).is_err());
        assert!(check_edit_schema(RowEdit::Patch, &schema(&[("label", DataType::Utf8)])).is_err());
    }
}