spiral-sdk 0.1.11

//! Immutable table versions.

use std::fmt::{self, Debug, Formatter};
use std::sync::Arc;

use uuid::Uuid;

use crate::api::tasks::{
    TableMaintenanceDefinition, TaskDefinition, TaskReadTarget, TaskReadTargetLatestOnBranch,
};
use crate::arrow_input::IntoArrowStream;
use crate::catalog_types::{ProjectBranchId, ProjectCommitId, ProjectId, ResourceId};
use crate::client::Connection;
use crate::commit::BindingMutation;
use crate::registration::register;
use crate::scandal::types::{PreparedWriteReceipt, TableWriteResult};
use crate::task_handle::submit;
use crate::{
    Error, Expression, Project, ResourceAddress, Run, Scan, ScanOptions, TaskHandle, WorkItem,
};

/// Stable identity of one immutable table version.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct TableIdentity {
    /// Arbor resource identifier.
    pub resource_id: Uuid,
    /// Exact immutable version identifier.
    pub version_id: Uuid,
}

/// An exact immutable table version.
///
/// Writing a table produces its next version without publishing it: nothing changes in the
/// Catalog until a [`Commit`](crate::Commit) stages the new version and publishes.
#[derive(Clone)]
pub struct Table {
    pub(crate) connection: Arc<Connection>,
    identity: TableIdentity,
    pub(crate) base: Option<CatalogBase>,
    pub(crate) address: Option<ResourceAddress>,
}

/// The version a Catalog commit binds for a table's resource. Versions written from that
/// table keep it, so a commit on that commit's branch can swap it for the newest one.
#[derive(Clone)]
pub(crate) struct CatalogBase {
    pub(crate) project_id: ProjectId,
    pub(crate) commit_id: ProjectCommitId,
    /// The branch that commit was pinned from, or `None` for an exact-commit snapshot.
    pub(crate) branch_id: Option<ProjectBranchId>,
    pub(crate) version_id: Uuid,
    /// The path that resolved the bound version, when it was opened by path.
    pub(crate) address: Option<ResourceAddress>,
}

/// Rows deleted from a table, and the version without them.
#[derive(Clone, Debug)]
pub struct DeletedRows {
    /// The version without the deleted rows; the original table when none matched.
    pub table: Table,
    /// How many visible rows were deleted.
    pub rows: u64,
}

impl Debug for Table {
    fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("Table")
            .field("identity", &self.identity)
            .finish()
    }
}

impl Project {
    /// Write the first version of a new table from `input`.
    ///
    /// The table has no name until [`Commit::create_table`](crate::Commit::create_table)
    /// binds it and the commit publishes.
    pub async fn write_table(&self, input: impl IntoArrowStream) -> Result<Table, Error> {
        let result = self
            .client
            .connection
            .tables()?
            .create(self.id.as_str(), input.into_write_input()?)
            .await?;
        Ok(self
            .client
            .table_by_id(result.resource_id, result.version_id))
    }
}

impl Table {
    pub(crate) fn new(connection: Arc<Connection>, identity: TableIdentity) -> Self {
        Self {
            connection,
            identity,
            base: None,
            address: None,
        }
    }

    /// Exact immutable identity selected by this handle.
    pub fn identity(&self) -> TableIdentity {
        self.identity
    }

    /// Canonical `path@v<version>` provenance, when opened through a resource path.
    /// Tables opened by IDs, and versions written from a table, have no resource path.
    pub fn resolved_address(&self) -> Option<&ResourceAddress> {
        self.address.as_ref()
    }

    /// Submit maintenance of this table's latest version on the branch it was opened from.
    ///
    /// Maintenance compacts and reclusters, publishing each step as a new version; its
    /// result lists them. Returns `None` when maintenance of this table and branch is
    /// already waiting to run.
    pub async fn maintain(&self, run: Run) -> Result<Option<TaskHandle>, Error> {
        let (project_id, branch_id) = self
            .base
            .as_ref()
            .and_then(|base| Some((&base.project_id, base.branch_id?)))
            .ok_or_else(|| Error::invalid("maintaining a table needs one opened from a branch"))?;
        let definition = TaskDefinition::TableMaintenance(TableMaintenanceDefinition {
            read_target: TaskReadTarget::LatestOnBranch(TaskReadTargetLatestOnBranch {
                branch_id: branch_id.as_uuid(),
            }),
            table_resource_id: self.identity.resource_id,
        });
        submit(
            &self.connection.tasks,
            project_id.as_str().into(),
            definition,
            None,
            run,
        )
        .await
    }

    /// Open a scan with fixed projection, predicate, schema, and execution binding.
    ///
    /// Metadata is resolved once. Batches are read when its streams are polled.
    pub async fn scan(&self, options: ScanOptions) -> Result<Scan, Error> {
        let mut scan = self
            .connection
            .tables()?
            .open_scan(self.identity, options)
            .await?;
        scan.address = self.address.clone();
        Ok(scan)
    }

    /// Write the next version of this table: its rows followed by `input`.
    pub async fn append(&self, input: impl IntoArrowStream) -> Result<Table, Error> {
        self.write(input, false).await
    }

    /// Write the next version of this table with only the rows of `input`.
    pub async fn overwrite(&self, input: impl IntoArrowStream) -> Result<Table, Error> {
        self.write(input, true).await
    }

    /// Write the next version of this table without the visible rows matching `filter`, or
    /// without any rows when there is no filter.
    ///
    /// Retained rows keep their stable row IDs. When no row matches, nothing is written and
    /// the result holds this table. Unsupported predicates fail before anything is written.
    pub async fn delete_rows(&self, filter: Option<Expression>) -> Result<DeletedRows, Error> {
        let scan = self
            .scan(ScanOptions {
                columns: Some(Vec::new()),
                filter: filter.clone(),
                ..Default::default()
            })
            .await?;
        let work = scan
            .splits()
            .iter()
            .enumerate()
            .map(|(split, item)| WorkItem {
                split,
                rows: item.rows.clone(),
            })
            .collect::<Vec<_>>();
        let rows = scan.count_rows(&work).await?.into_iter().sum();
        if rows == 0 {
            return Ok(DeletedRows {
                table: self.clone(),
                rows,
            });
        }
        let result = self
            .connection
            .tables()?
            .delete(self.identity, filter, rows)
            .await?;
        Ok(DeletedRows {
            table: self.successor(result)?,
            rows,
        })
    }

    /// Register a version of this table that a write prepared but whose registration outcome
    /// was unknown ([`Error::RegistrationUnknown`]), without uploading its rows again.
    pub async fn recover_write(&self, receipt: PreparedWriteReceipt) -> Result<Table, Error> {
        if receipt.resource_id != self.identity.resource_id
            || receipt.base_version != Some(self.identity.version_id)
        {
            return Err(Error::invalid("receipt was not written from this table"));
        }
        let result = register(&self.connection.arbor, receipt).await?;
        self.successor(result)
    }

    /// The binding that selected this table's Catalog base, to remove from the project root.
    pub(crate) fn unbinding(&self) -> Result<BindingMutation, Error> {
        let (parent, name) = self
            .base
            .as_ref()
            .and_then(|base| base.address.as_ref())
            .ok_or_else(|| Error::invalid("table was not opened by path"))?
            .root_binding()?;
        Ok(BindingMutation {
            parent,
            name,
            expected: Some(ResourceId::new(self.identity.resource_id).map_err(Error::invalid)?),
            replacement: None,
        })
    }

    async fn write(&self, input: impl IntoArrowStream, overwrite: bool) -> Result<Self, Error> {
        let result = self
            .connection
            .tables()?
            .write(self.identity, input.into_write_input()?, overwrite)
            .await?;
        self.successor(result)
    }

    /// The version `result` wrote from this table, keeping this table's Catalog base.
    pub(crate) fn successor(&self, result: TableWriteResult) -> Result<Self, Error> {
        if result.resource_id != self.identity.resource_id {
            return Err(Error::invalid("write changed resource identity"));
        }
        Ok(Self {
            connection: Arc::clone(&self.connection),
            identity: TableIdentity {
                resource_id: result.resource_id,
                version_id: result.version_id,
            },
            base: self.base.clone(),
            address: None,
        })
    }
}