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};
#[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 {
pub async fn table_session(&self, default_ref: &str) -> Result<TableSession, Error> {
TableSession::open(self.clone(), default_ref.parse()?).await
}
}
impl Snapshot {
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)),
})
}
pub fn default_ref(&self) -> &ResourceRef {
&self.default_ref
}
pub fn snapshot(&self) -> &Snapshot {
&self.snapshot
}
pub fn head(&self) -> Option<BranchHead> {
BranchHead::of(&self.snapshot)
}
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)),
})
}
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())
}
}
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)
}
}