use std::{
fmt,
fs::File,
io::{self, Write, stdout},
path::PathBuf,
};
use anyhow::{Result, bail};
use clap::{ArgGroup, Args, Subcommand};
use io_pimdir::{PimdirItem, PimdirRetainedItem, PimdirStore, codec::PimdirAction};
use io_replica::{
collection::ReplicaCollectionId,
object::ReplicaHash,
placement::{ReplicaFlags, ReplicaLevel, ReplicaLinkId, ReplicaMeta},
};
use log::warn;
use pimalaya_cli::{
printer::Printer,
table::{Cell, ContentArrangement, Table, presets::UTF8_FULL},
};
use serde::Serialize;
use crate::cli::{StoreFlags, bytes, confirm, now, or_dash, report};
const DEFAULT_LIMIT: usize = 50;
#[derive(Debug, Subcommand)]
pub enum ItemCommand {
List(ItemListCommand),
Show(ItemShowCommand),
Export(ItemExportCommand),
Restore(ItemRestoreCommand),
Purge(ItemPurgeCommand),
}
impl ItemCommand {
pub fn execute(self, printer: &mut impl Printer, store: &StoreFlags) -> Result<()> {
match self {
Self::List(cmd) => cmd.execute(printer, store),
Self::Show(cmd) => cmd.execute(printer, store),
Self::Export(cmd) => cmd.execute(printer, store),
Self::Restore(cmd) => cmd.execute(printer, store),
Self::Purge(cmd) => cmd.execute(printer, store),
}
}
}
#[derive(Debug, Args)]
pub struct ItemListCommand {
#[arg(value_name = "COLLECTION")]
pub collection: String,
#[arg(long)]
pub retained: bool,
#[arg(long, value_name = "CURSOR")]
pub after: Option<String>,
#[arg(long, short = 'n', value_name = "COUNT", default_value_t = DEFAULT_LIMIT)]
pub limit: usize,
}
impl ItemListCommand {
pub fn execute(self, printer: &mut impl Printer, store: &StoreFlags) -> Result<()> {
let store = store.read()?;
let limit = self.limit.max(1);
let probe = limit.saturating_add(1);
let mut rows: Vec<ItemRow> = if self.retained {
let after = match &self.after {
None => None,
Some(cursor) => Some(cursor.parse::<i64>().map_err(|_| {
anyhow::anyhow!("--after takes a seq with --retained, got {cursor:?}")
})?),
};
let collection = ReplicaCollectionId(self.collection.clone());
store
.list_retained(&collection, after, probe)
.map_err(report)?
.into_iter()
.map(|item| ItemRow::retained(&self.collection, &item))
.collect()
} else {
store
.list_items(&self.collection, self.after.as_deref(), probe)
.map_err(report)?
.into_iter()
.map(|item| ItemRow::live(&self.collection, &item))
.collect()
};
let truncated = rows.len() > limit;
rows.truncate(limit);
let next = truncated
.then(|| rows.last().map(|row| row.cursor(self.retained)))
.flatten();
printer.out(ItemsOutput {
collection: self.collection,
retained: self.retained,
items: rows,
next,
})
}
}
#[derive(Debug, Args)]
pub struct ItemShowCommand {
#[arg(value_name = "SEQ")]
pub seq: i64,
#[arg(long, short = 'c', value_name = "COLLECTION")]
pub collection: Option<String>,
}
impl ItemShowCommand {
pub fn execute(self, printer: &mut impl Printer, store: &StoreFlags) -> Result<()> {
let db = store.db().ok();
let read = store.read()?;
let found = locate(&read, self.seq, self.collection.as_deref())?;
if found.is_empty() {
bail!("no item with seq {} in this store", self.seq);
}
let placements = found
.into_iter()
.map(|found| {
let mut row = found.row();
row.size = row
.object
.as_deref()
.and_then(|hash| db.as_ref().and_then(|db| db.object_size(hash).ok()))
.flatten();
row
})
.collect();
printer.out(ItemShowOutput {
seq: self.seq,
placements,
})
}
}
#[derive(Debug, Args)]
pub struct ItemExportCommand {
#[arg(value_name = "SEQ")]
pub seq: i64,
#[arg(long, short = 'c', value_name = "COLLECTION")]
pub collection: Option<String>,
#[arg(long, short = 'o', value_name = "PATH")]
pub output: Option<PathBuf>,
}
impl ItemExportCommand {
pub fn execute(self, printer: &mut impl Printer, store: &StoreFlags) -> Result<()> {
if printer.is_json() && self.output.is_none() {
bail!("a body is raw bytes, not JSON: write it to a file with --output");
}
let read = store.read()?;
let found = one(
locate(&read, self.seq, self.collection.as_deref())?,
self.seq,
)?;
let Some(hash) = found.object() else {
bail!(
"seq {} holds no body in {} (its detail level is {})",
self.seq,
found.collection,
level_name(found.level())
);
};
let blobs = store.blobs()?;
let Some(mut reader) = blobs.reader(&ReplicaHash(hash.clone()))? else {
bail!(
"the body of seq {} is missing from the blob store (hash {hash}); run `pimdir check`",
self.seq
);
};
let hash = hash.clone();
let collection = found.collection.clone();
match self.output {
None => {
let mut out = stdout().lock();
io::copy(&mut reader, &mut out)?;
out.flush()?;
Ok(())
}
Some(path) => {
let mut file = File::create(&path)?;
let written = io::copy(&mut reader, &mut file)?;
file.flush()?;
printer.out(ItemExportOutput {
seq: self.seq,
collection,
hash,
bytes: written,
path,
})
}
}
}
}
#[derive(Debug, Args)]
pub struct ItemRestoreCommand {
#[arg(value_name = "SEQ")]
pub seq: i64,
#[arg(long, short = 'c', value_name = "COLLECTION")]
pub collection: Option<String>,
}
impl ItemRestoreCommand {
pub fn execute(self, printer: &mut impl Printer, store: &StoreFlags) -> Result<()> {
let read = store.read()?;
let found = one(
locate(&read, self.seq, self.collection.as_deref())?,
self.seq,
)?;
let FoundItem::Retained(item) = &found.item else {
bail!(
"seq {} is already live in {}: only a retained item can be restored",
self.seq,
found.collection
);
};
drop(read);
let mut owner = store.owner()?;
let action = PimdirAction::Add {
link_id: Some(ReplicaLinkId(item.link_id.clone())),
flags: item.flags.clone(),
object: item.object_hash.clone().map(ReplicaHash),
meta: item.meta.clone().map(ReplicaMeta),
handle: None,
};
let id = store
.producer()?
.enqueue(&found.collection, &action, None, &now())
.map_err(report)?;
let status = match owner.drain_collection(&found.collection) {
Err(io_pimdir::PimdirError::Busy) => RestoreStatus::Queued,
Err(err) => return Err(report(err)),
Ok(_) => match owner
.get_item(&found.collection, self.seq)
.map_err(report)?
{
Some(_) => RestoreStatus::Applied,
None => RestoreStatus::Refused,
},
};
printer.out(ItemRestoreOutput {
seq: self.seq,
collection: found.collection,
link_id: item.link_id.clone(),
action: id,
status,
})
}
}
#[derive(Debug, Args)]
#[command(group(ArgGroup::new("target").required(true).args(["seq", "older_than", "all"])))]
pub struct ItemPurgeCommand {
#[arg(value_name = "SEQ")]
pub seq: Option<i64>,
#[arg(long, value_name = "DURATION")]
pub older_than: Option<humantime::Duration>,
#[arg(long)]
pub all: bool,
#[arg(long, short = 'c', value_name = "COLLECTION")]
pub collection: Option<String>,
#[arg(long, short = 'y')]
pub yes: bool,
}
impl ItemPurgeCommand {
pub fn execute(self, printer: &mut impl Printer, store: &StoreFlags) -> Result<()> {
if let Some(seq) = self.seq {
return self.purge_one(printer, store, seq);
}
let cutoff = match self.older_than {
Some(age) => crate::cli::cutoff(*age),
None => now(),
};
let preview = store
.db()
.and_then(|db| db.retained_before(&cutoff))
.map_err(|err| warn!("cannot preview what this purge would destroy: {err}"))
.ok();
let question = match preview {
Some((items, size)) => format!(
"Destroy {items} retained item(s) and up to {} for good?",
bytes(size)
),
None => format!("Destroy every item retained before {cutoff} for good?"),
};
confirm(printer, self.yes, &question)?;
let purged = store
.owner_any_source()?
.purge_retained_before(&cutoff)
.map_err(report)?;
printer.out(ItemPurgeOutput {
cutoff: Some(cutoff),
items: purged.items,
bytes: purged.bytes,
})
}
fn purge_one(&self, printer: &mut impl Printer, store: &StoreFlags, seq: i64) -> Result<()> {
let read = store.read()?;
let found = one(locate(&read, seq, self.collection.as_deref())?, seq)?;
let FoundItem::Retained(item) = &found.item else {
bail!(
"seq {seq} is live in {}: purge only destroys retained items, remove it first",
found.collection
);
};
let size = item.size.unwrap_or(0);
drop(read);
confirm(
printer,
self.yes,
&format!(
"Destroy retained item {seq} ({}) in {} for good?",
bytes(size),
found.collection
),
)?;
let collection = ReplicaCollectionId(found.collection.clone());
let purged = store
.owner_any_source()?
.purge(&collection, seq)
.map_err(report)?;
if !purged {
let collection = &found.collection;
bail!("seq {seq} was not purged: it is no longer retained in {collection}");
}
printer.out(ItemPurgeOutput {
cutoff: None,
items: 1,
bytes: size,
})
}
}
enum FoundItem {
Live(PimdirItem),
Retained(PimdirRetainedItem),
}
struct Found {
collection: String,
item: FoundItem,
}
impl Found {
fn object(&self) -> Option<&String> {
match &self.item {
FoundItem::Live(item) => item.object.as_ref().map(|hash| &hash.0),
FoundItem::Retained(item) => item.object_hash.as_ref(),
}
}
fn level(&self) -> ReplicaLevel {
match &self.item {
FoundItem::Live(item) => item.level,
FoundItem::Retained(item) => item.level,
}
}
fn row(&self) -> ItemRow {
match &self.item {
FoundItem::Live(item) => ItemRow::live(&self.collection, item),
FoundItem::Retained(item) => ItemRow::retained(&self.collection, item),
}
}
}
fn locate(store: &PimdirStore, seq: i64, collection: Option<&str>) -> Result<Vec<Found>> {
let collections: Vec<String> = match collection {
Some(collection) => vec![collection.to_string()],
None => store
.list_collections()
.map_err(report)?
.into_iter()
.map(|collection| collection.id)
.collect(),
};
let mut found = Vec::new();
for collection in collections {
if let Some(item) = store.get_item(&collection, seq).map_err(report)? {
found.push(Found {
collection,
item: FoundItem::Live(item),
});
continue;
}
if let Some(item) = retained(store, &collection, seq)? {
found.push(Found {
collection,
item: FoundItem::Retained(item),
});
}
}
Ok(found)
}
fn retained(store: &PimdirStore, collection: &str, seq: i64) -> Result<Option<PimdirRetainedItem>> {
let collection = ReplicaCollectionId(collection.to_string());
let page = store
.list_retained(&collection, Some(seq - 1), 1)
.map_err(report)?;
Ok(page.into_iter().find(|item| item.seq == seq))
}
fn one(found: Vec<Found>, seq: i64) -> Result<Found> {
match found.len() {
0 => bail!("no item with seq {seq} in this store"),
1 => Ok(found.into_iter().next().unwrap()),
_ => {
let collections: Vec<&str> = found.iter().map(|f| f.collection.as_str()).collect();
bail!(
"seq {seq} is placed in several collections ({}): pick one with --collection",
collections.join(", ")
)
}
}
}
fn level_name(level: ReplicaLevel) -> &'static str {
match level {
ReplicaLevel::Probed => "probed",
ReplicaLevel::Meta => "meta",
ReplicaLevel::Full => "full",
}
}
fn flag_list(flags: &ReplicaFlags) -> Option<Vec<String>> {
Some(flags.known()?.iter().cloned().collect())
}
#[derive(Debug, Serialize)]
pub struct ItemRow {
pub collection: String,
pub seq: i64,
pub link_id: String,
pub flags: Option<Vec<String>>,
pub level: &'static str,
pub object: Option<String>,
pub size: Option<u64>,
pub meta: Option<String>,
pub retained_at: Option<String>,
pub retained_by: Option<String>,
}
impl ItemRow {
fn live(collection: &str, item: &PimdirItem) -> Self {
Self {
collection: collection.to_string(),
seq: item.seq,
link_id: item.link_id.0.clone(),
flags: flag_list(&item.flags),
level: level_name(item.level),
object: item.object.as_ref().map(|hash| hash.0.clone()),
size: None,
meta: item.meta.as_ref().map(|meta| meta.0.clone()),
retained_at: None,
retained_by: None,
}
}
fn retained(collection: &str, item: &PimdirRetainedItem) -> Self {
Self {
collection: collection.to_string(),
seq: item.seq,
link_id: item.link_id.clone(),
flags: flag_list(&item.flags),
level: level_name(item.level),
object: item.object_hash.clone(),
size: item.size,
meta: item.meta.clone(),
retained_at: Some(item.retained_at.clone()),
retained_by: item.retained_by.clone(),
}
}
fn cursor(&self, retained: bool) -> String {
if retained {
self.seq.to_string()
} else {
self.link_id.clone()
}
}
}
#[derive(Debug, Serialize)]
pub struct ItemsOutput {
pub collection: String,
pub retained: bool,
pub items: Vec<ItemRow>,
pub next: Option<String>,
}
impl fmt::Display for ItemsOutput {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
if self.items.is_empty() {
let what = if self.retained { "retained" } else { "live" };
return writeln!(f, "No {what} item in {}", self.collection);
}
let mut table = Table::new();
let mut header = vec![
Cell::new("SEQ"),
Cell::new("LINK ID"),
Cell::new("FLAGS"),
Cell::new("LEVEL"),
Cell::new("OBJECT"),
];
if self.retained {
header.push(Cell::new("RETAINED AT"));
header.push(Cell::new("BY"));
}
table
.load_style(UTF8_FULL)
.set_content_arrangement(ContentArrangement::Dynamic)
.set_header(header);
for item in &self.items {
let mut row = vec![
Cell::new(item.seq),
Cell::new(&item.link_id),
Cell::new(or_dash(
item.flags
.as_ref()
.map(|flags| flags.join(" "))
.filter(|f| !f.is_empty())
.as_deref(),
)),
Cell::new(item.level),
Cell::new(or_dash(item.object.as_deref())),
];
if self.retained {
row.push(Cell::new(or_dash(item.retained_at.as_deref())));
row.push(Cell::new(or_dash(item.retained_by.as_deref())));
}
table.add_row(row);
}
writeln!(f, "{table}")?;
if let Some(next) = &self.next {
writeln!(f, "More items follow: continue with --after {next}")?;
}
Ok(())
}
}
#[derive(Debug, Serialize)]
pub struct ItemShowOutput {
pub seq: i64,
pub placements: Vec<ItemRow>,
}
impl fmt::Display for ItemShowOutput {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
for (index, item) in self.placements.iter().enumerate() {
if index > 0 {
writeln!(f)?;
}
writeln!(f, "Item {} in {}", item.seq, item.collection)?;
writeln!(f, " - link id: {}", item.link_id)?;
writeln!(
f,
" - flags: {}",
or_dash(
item.flags
.as_ref()
.map(|flags| flags.join(" "))
.filter(|flags| !flags.is_empty())
.as_deref()
)
)?;
writeln!(f, " - level: {}", item.level)?;
writeln!(f, " - object: {}", or_dash(item.object.as_deref()))?;
if let Some(size) = item.size {
writeln!(f, " - size: {}", bytes(size))?;
}
if let Some(at) = &item.retained_at {
writeln!(f, " - retained at: {at}")?;
writeln!(
f,
" - retained by: {}",
or_dash(item.retained_by.as_deref())
)?;
}
writeln!(f, " - meta: {}", or_dash(item.meta.as_deref()))?;
}
Ok(())
}
}
#[derive(Debug, Serialize)]
pub struct ItemExportOutput {
pub seq: i64,
pub collection: String,
pub hash: String,
pub bytes: u64,
pub path: PathBuf,
}
impl fmt::Display for ItemExportOutput {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
writeln!(
f,
"Wrote {} of item {} to {}",
bytes(self.bytes),
self.seq,
self.path.display()
)
}
}
#[derive(Clone, Copy, Debug, Serialize)]
#[serde(rename_all = "kebab-case")]
pub enum RestoreStatus {
Applied,
Queued,
Refused,
}
#[derive(Debug, Serialize)]
pub struct ItemRestoreOutput {
pub seq: i64,
pub collection: String,
pub link_id: String,
pub action: i64,
pub status: RestoreStatus,
}
impl fmt::Display for ItemRestoreOutput {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self.status {
RestoreStatus::Applied => writeln!(
f,
"Restored item {} into {} (applied)",
self.seq, self.collection
),
RestoreStatus::Queued => writeln!(
f,
"Restore of item {} into {} queued as action {}: another writer holds the store lock, so it applies at the next sync",
self.seq, self.collection, self.action
),
RestoreStatus::Refused => writeln!(
f,
"Restore of item {} into {} was not applied: see `pimdir queue list --parked`",
self.seq, self.collection
),
}
}
}
#[derive(Debug, Serialize)]
pub struct ItemPurgeOutput {
pub cutoff: Option<String>,
pub items: usize,
pub bytes: u64,
}
impl fmt::Display for ItemPurgeOutput {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
writeln!(
f,
"Purged {} item(s), reclaiming {}",
self.items,
bytes(self.bytes)
)
}
}