spiral/

table_session.rs

//! Repeatable resource resolution in one project.

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

use parking_lot::Mutex;
use tokio::sync::OnceCell;

use crate::catalog_types::{ProjectBranchName, VersionId};
use crate::{BranchHead, Error, Project, ResourceAddress, ResourceRef, Snapshot, Table};

/// A default snapshot and branch heads pinned on first use.
///
/// Clones share pins. Refresh creates a separate session, leaving existing
/// tables, scans, and sessions unchanged.
#[derive(Clone)]
pub struct TableSession {
    project: Project,
    default_ref: ResourceRef,
    snapshot: Snapshot,
    branches: Arc<Mutex<HashMap<ProjectBranchName, Arc<OnceCell<Snapshot>>>>>,
}

impl Debug for TableSession {
    fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("TableSession")
            .field("project", self.project.id())
            .field("default_ref", &self.default_ref)
            .field("commit", &self.snapshot.commit_id())
            .finish()
    }
}

impl Project {
    /// Open a table session with a branch or `c`-prefixed commit as its default ref.
    pub async fn table_session(&self, default_ref: &str) -> Result<TableSession, Error> {
        TableSession::open(self.clone(), default_ref.parse()?).await
    }
}

impl Snapshot {
    /// Resolve qualified resource addresses with this exact commit as the default.
    pub fn table_session(&self) -> TableSession {
        TableSession {
            project: self.project().clone(),
            default_ref: ResourceRef::Commit(self.commit_id()),
            snapshot: self.clone(),
            branches: Arc::new(Mutex::new(HashMap::new())),
        }
    }
}

impl TableSession {
    async fn open(project: Project, default_ref: ResourceRef) -> Result<Self, Error> {
        let snapshot = match &default_ref {
            ResourceRef::Branch(branch) => {
                project.branch(branch.as_str()).await?.snapshot().clone()
            }
            ResourceRef::Commit(commit) => project.snapshot_at(*commit),
            ResourceRef::Version(_) => {
                return Err(Error::invalid(
                    "a session default must be a branch or project commit, not a table version",
                ));
            }
        };
        let mut branches = HashMap::new();
        if let ResourceRef::Branch(branch) = &default_ref {
            branches.insert(
                branch.clone(),
                Arc::new(OnceCell::new_with(Some(snapshot.clone()))),
            );
        }
        Ok(Self {
            project,
            default_ref,
            snapshot,
            branches: Arc::new(Mutex::new(branches)),
        })
    }

    /// Default ref selected when the session was opened.
    pub fn default_ref(&self) -> &ResourceRef {
        &self.default_ref
    }

    /// Immutable default snapshot.
    pub fn snapshot(&self) -> &Snapshot {
        &self.snapshot
    }

    /// The pinned head of the default branch, which commits publish onto; `None` when the
    /// default ref is an exact commit.
    pub fn head(&self) -> Option<BranchHead> {
        BranchHead::of(&self.snapshot)
    }

    /// Start a separate session at a newer head of the default branch, such as a published
    /// commit's.
    ///
    /// The default branch name is preserved and other branch pins are reset.
    /// Existing tables and sessions keep their original immutable bindings.
    pub fn with_head(&self, head: BranchHead) -> Result<Self, Error> {
        let ResourceRef::Branch(branch) = &self.default_ref else {
            return Err(Error::invalid("session must select a default branch"));
        };
        let snapshot = head.snapshot().clone();
        if !Arc::ptr_eq(
            &self.project.client.connection,
            &snapshot.project.client.connection,
        ) || self.project.id != snapshot.project.id
            || self.snapshot.branch_id != Some(head.branch_id())
        {
            return Err(Error::invalid("head must be of this session's branch"));
        }
        let branches = HashMap::from([(
            branch.clone(),
            Arc::new(OnceCell::new_with(Some(snapshot.clone()))),
        )]);
        Ok(Self {
            project: self.project.clone(),
            default_ref: self.default_ref.clone(),
            snapshot,
            branches: Arc::new(Mutex::new(branches)),
        })
    }

    /// Open a new session with fresh branch pins and the same default ref.
    pub async fn refresh(&self) -> Result<Self, Error> {
        if matches!(self.default_ref, ResourceRef::Branch(_)) {
            Self::open(self.project.clone(), self.default_ref.clone()).await
        } else {
            Ok(self.snapshot.table_session())
        }
    }

    /// Resolve `path`, `path@branch`, `path@c<commit>`, or `path@v<version>`.
    ///
    /// Exact versions identify their resource through the default snapshot's
    /// path. Names and refs never cross into scandal.
    pub async fn table(&self, address: &str) -> Result<Table, Error> {
        let address: ResourceAddress = address.parse()?;
        let snapshot = match address.reference() {
            None | Some(ResourceRef::Version(_)) => self.snapshot.clone(),
            Some(ResourceRef::Commit(commit)) => self.project.snapshot_at(*commit),
            Some(ResourceRef::Branch(branch)) => {
                let pin = {
                    let mut branches = self.branches.lock();
                    Arc::clone(
                        branches
                            .entry(branch.clone())
                            .or_insert_with(|| Arc::new(OnceCell::new())),
                    )
                };
                pin.get_or_try_init(|| async {
                    Ok::<_, Error>(
                        self.project
                            .branch(branch.as_str())
                            .await?
                            .snapshot()
                            .clone(),
                    )
                })
                .await?
                .clone()
            }
        };
        let mut table = snapshot.table_at_path(address.path()).await?;
        if let Some(ResourceRef::Version(version)) = address.reference() {
            table = self
                .project
                .client
                .table_by_id(table.identity().resource_id, version.as_uuid());
        }
        let version = VersionId::new(table.identity().version_id).map_err(Error::invalid)?;
        let address = address.at_version(version);
        if let Some(base) = &mut table.base {
            base.address = Some(address.clone());
        }
        table.address = Some(address);
        Ok(table)
    }
}