libsql/bottomless-cli/src/replicator_extras.rs

use anyhow::Result;
use aws_sdk_s3::error::SdkError;
use aws_sdk_s3::types::ObjectAttributes;
use aws_sdk_s3::Client;
use aws_smithy_types::date_time::Format;
use chrono::{NaiveDate, NaiveDateTime, NaiveTime};
use tokio::io::BufReader;

pub(crate) struct Replicator {
    inner: bottomless::replicator::Replicator,
}

impl std::ops::Deref for Replicator {
    type Target = bottomless::replicator::Replicator;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl std::ops::DerefMut for Replicator {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner
    }
}

fn uuid_to_datetime(uuid: &uuid::Uuid) -> chrono::NaiveDateTime {
    let timestamp = bottomless::replicator::Replicator::generation_to_timestamp(uuid);
    let (seconds, _) = timestamp
        .as_ref()
        .map(uuid::Timestamp::to_unix)
        .unwrap_or_default();
    chrono::DateTime::from_timestamp_millis((seconds * 1000) as i64)
        .unwrap()
        .naive_utc()
}

pub(crate) async fn detect_db(client: &Client, bucket: &str, namespace: &str) -> Option<String> {
    let namespace = namespace.to_owned() + ":";
    let response = client
        .list_objects()
        .bucket(bucket)
        .set_delimiter(Some("/".to_string()))
        .prefix(namespace.clone())
        .send()
        .await
        .ok()?;

    let prefix = response.common_prefixes().first()?.prefix()?;
    // 38 is the length of the uuid part
    if let Some('-') = prefix.chars().nth(prefix.len().saturating_sub(38)) {
        let ns_db = &prefix[..prefix.len().saturating_sub(38)];
        Some(ns_db.strip_prefix(&namespace).unwrap_or(ns_db).to_owned())
    } else {
        None
    }
}

impl Replicator {
    pub async fn new(db: String) -> Result<Self> {
        let inner = bottomless::replicator::Replicator::new(db).await?;
        Ok(Replicator { inner })
    }

    pub(crate) async fn print_snapshot_summary(&self, generation: &uuid::Uuid) -> Result<()> {
        match self
            .client
            .get_object_attributes()
            .bucket(&self.bucket)
            .key(format!("{}-{}/db.gz", self.db_name, generation))
            .object_attributes(ObjectAttributes::ObjectSize)
            .send()
            .await
        {
            Ok(attrs) => {
                println!("\tmain database snapshot:");
                println!("\t\tobject size:   {}", attrs.object_size().unwrap());
                println!(
                    "\t\tlast modified: {}",
                    attrs
                        .last_modified()
                        .map(|s| s.fmt(Format::DateTime).unwrap_or_else(|e| e.to_string()))
                        .as_deref()
                        .unwrap_or("never")
                );
            }
            Err(SdkError::ServiceError(err)) if err.err().is_no_such_key() => {
                println!("\tno main database snapshot file found")
            }
            Err(e) => println!("\tfailed to fetch main database snapshot info: {e}"),
        };
        Ok(())
    }

    pub(crate) async fn list_generations(
        &self,
        limit: Option<u64>,
        older_than: Option<chrono::NaiveDate>,
        newer_than: Option<chrono::NaiveDate>,
        verbose: bool,
    ) -> Result<()> {
        let mut next_marker = None;
        let mut limit = limit.unwrap_or(u64::MAX);
        loop {
            let mut list_request = self
                .client
                .list_objects()
                .bucket(&self.bucket)
                .set_delimiter(Some("/".to_string()))
                .prefix(&self.db_name);

            if let Some(marker) = next_marker {
                list_request = list_request.marker(marker)
            }

            if verbose {
                println!("Database {}:", self.db_name);
            }

            let response = list_request.send().await?;
            let prefixes = response.common_prefixes();

            if prefixes.is_empty() {
                println!("No generations found");
                return Ok(());
            }

            for prefix in prefixes {
                if let Some(prefix) = &prefix.prefix {
                    let prefix = &prefix[self.db_name.len() + 1..prefix.len() - 1];
                    let uuid = uuid::Uuid::try_parse(prefix)?;
                    let datetime = uuid_to_datetime(&uuid);
                    if datetime.date() < newer_than.unwrap_or(chrono::NaiveDate::MIN) {
                        continue;
                    }
                    if datetime.date() > older_than.unwrap_or(chrono::NaiveDate::MAX) {
                        continue;
                    }
                    println!("{} (created: {})", uuid, datetime.and_utc().to_rfc3339());
                    if verbose {
                        let counter = self.get_remote_change_counter(&uuid).await?;
                        let consistent_frame = self.get_last_consistent_frame(&uuid).await?;
                        let m = self.get_metadata(&uuid).await?;
                        let parent = self.get_dependency(&uuid).await?;
                        println!("\tcreated at (UTC):     {datetime}");
                        println!("\tchange counter:       {counter:?}");
                        println!("\tconsistent WAL frame: {consistent_frame}");
                        if let Some((page_size, crc)) = m {
                            println!("\tpage size:            {}", page_size);
                            println!("\tWAL frame checksum:   {:X}-{:X}", crc.0, crc.1);
                        }
                        if let Some(prev_gen) = parent {
                            println!("\tprevious generation:  {}", prev_gen);
                        }
                        self.print_snapshot_summary(&uuid).await?;
                        println!()
                    }
                }
                limit -= 1;
                if limit == 0 {
                    return Ok(());
                }
            }

            next_marker = response.next_marker().map(|s| s.to_owned());
            if next_marker.is_none() {
                return Ok(());
            }
        }
    }

    pub(crate) async fn remove_many(&self, older_than: NaiveDate, verbose: bool) -> Result<()> {
        let older_than = NaiveDateTime::new(older_than, NaiveTime::MIN);
        let delete_all = self.inner.delete_all(Some(older_than)).await?;
        if verbose {
            println!("Tombstoned {} at {}", self.inner.db_path, older_than);
        }
        let removed_generations = delete_all.commit().await?;
        if verbose {
            println!(
                "Removed {} generations of {} up to {}",
                removed_generations, self.inner.db_path, older_than
            );
        }
        Ok(())
    }

    pub(crate) async fn remove(&self, generation: uuid::Uuid, verbose: bool) -> Result<()> {
        let mut removed = 0;
        let mut next_marker = None;
        loop {
            let mut list_request = self
                .client
                .list_objects()
                .bucket(&self.bucket)
                .prefix(format!("{}-{}/", &self.db_name, generation));

            if let Some(marker) = next_marker {
                list_request = list_request.marker(marker)
            }

            let response = list_request.send().await?;
            let objs = response.contents();

            if objs.is_empty() {
                if verbose {
                    println!("No objects found")
                }
                return Ok(());
            }

            for obj in objs {
                if let Some(key) = obj.key() {
                    if verbose {
                        println!("Removing {key}")
                    }
                    self.client
                        .delete_object()
                        .bucket(&self.bucket)
                        .key(key)
                        .send()
                        .await?;
                    removed += 1;
                }
            }

            next_marker = response.next_marker().map(|s| s.to_owned());
            if next_marker.is_none() {
                if verbose {
                    println!("Removed {removed} snapshot generations");
                }
                return Ok(());
            }
        }
    }

    pub(crate) async fn list_generation(&self, generation: uuid::Uuid) -> Result<()> {
        let res = self
            .client
            .list_objects()
            .bucket(&self.bucket)
            .prefix(format!("{}-{}/", &self.db_name, generation))
            .max_keys(1)
            .send()
            .await?;

        if res.contents().is_empty() {
            anyhow::bail!("Generation {} not found for {}", generation, &self.db_name)
        }

        let counter = self.get_remote_change_counter(&generation).await?;
        let consistent_frame = self.get_last_consistent_frame(&generation).await?;
        let meta = self.get_metadata(&generation).await?;
        let dep = self.get_dependency(&generation).await?;
        println!("Generation {} for {}", generation, self.db_name);
        println!("\tcreated at:           {}", uuid_to_datetime(&generation));
        println!("\tchange counter:       {counter:?}");
        println!("\tconsistent WAL frame: {consistent_frame}");
        if let Some((page_size, crc)) = meta {
            println!("\tpage size:            {}", page_size);
            println!("\tWAL frame checksum:   {:X}-{:X}", crc.0, crc.1);
        }
        if let Some(prev_gen) = dep {
            println!("\tprevious generation:  {}", prev_gen);
        }
        self.print_snapshot_summary(&generation).await?;
        Ok(())
    }

    pub async fn restore_from_local_snapshot(
        from_dir: impl AsRef<std::path::Path>,
        db: &mut tokio::fs::File,
    ) -> Result<bool> {
        let from_dir = from_dir.as_ref();
        use bottomless::replicator::CompressionKind;
        use tokio::io::AsyncWriteExt;

        let algos_to_try = &[
            CompressionKind::Gzip,
            CompressionKind::Zstd,
            CompressionKind::None,
        ];

        for algo in algos_to_try {
            let main_db_path = match algo {
                CompressionKind::None => from_dir.join("db.db"),
                CompressionKind::Gzip => from_dir.join("db.gz"),
                CompressionKind::Zstd => from_dir.join("db.zstd"),
            };
            if let Ok(mut db_file) = tokio::fs::File::open(&main_db_path).await {
                let db_size = match algo {
                    CompressionKind::None => tokio::io::copy(&mut db_file, db).await?,
                    CompressionKind::Gzip => {
                        let mut decompress_reader =
                            async_compression::tokio::bufread::GzipDecoder::new(
                                tokio::io::BufReader::new(db_file),
                            );
                        tokio::io::copy(&mut decompress_reader, db).await?
                    }
                    CompressionKind::Zstd => {
                        let mut decompress_reader =
                            async_compression::tokio::bufread::ZstdDecoder::new(
                                tokio::io::BufReader::new(db_file),
                            );
                        tokio::io::copy(&mut decompress_reader, db).await?
                    }
                };
                db.flush().await?;

                tracing::info!("Restored the main database file ({} bytes)", db_size);
                return Ok(true);
            }
        }
        Ok(false)
    }

    pub async fn apply_wal_from_local_generation(
        from_dir: impl AsRef<std::path::Path>,
        db: &mut tokio::fs::File,
        page_size: u32,
        checksum: (u32, u32),
    ) -> Result<u32> {
        use bottomless::transaction_cache::TransactionPageCache;
        use tokio::io::AsyncWriteExt;

        const SWAP_AFTER: u32 = 65536;
        const TMP_RESTORE_DIR: &str = ".bottomless.restore.tmp";

        let from_dir = from_dir.as_ref();
        let mut page_buf = {
            let mut v = Vec::with_capacity(page_size as usize);
            v.spare_capacity_mut();
            unsafe { v.set_len(page_size as usize) };
            v
        };

        let objs = {
            let mut objs = Vec::new();
            let mut dir = tokio::fs::read_dir(from_dir).await.unwrap();
            while let Some(entry) = dir.next_entry().await.unwrap() {
                let path = entry.path();
                if let Some(file_name) = path.file_name() {
                    if let Some(file_name) = file_name.to_str() {
                        if file_name.ends_with(".gz")
                            || file_name.ends_with(".zstd")
                            || file_name.ends_with(".raw")
                        {
                            objs.push(path);
                        }
                    }
                }
            }
            objs.sort();
            objs.into_iter()
        };

        let mut last_received_frame_no = 0;
        let mut pending_pages =
            TransactionPageCache::new(SWAP_AFTER, page_size, TMP_RESTORE_DIR.into());

        let mut checksum: Option<(u32, u32)> = Some(checksum);
        for obj in objs {
            let key = obj.file_name().unwrap().to_str().unwrap();
            tracing::debug!("Loading {}", key);

            let (first_frame_no, _last_frame_no, _timestamp, compression_kind) =
                match bottomless::replicator::Replicator::parse_frame_range(&format!("/{key}")) {
                    Some(result) => result,
                    None => {
                        if key != "db.gz" && key != "db.zstd" && key != "db.db" {
                            tracing::warn!("Failed to parse frame/page from key {}", key);
                        }
                        continue;
                    }
                };
            if first_frame_no != last_received_frame_no + 1 {
                tracing::warn!("Missing series of consecutive frames. Last applied frame: {}, next found: {}. Stopping the restoration process",
                        last_received_frame_no, first_frame_no);
                break;
            }
            // read frame from the file - from_dir and `obj` dir entry compose the path to it
            let frame = tokio::fs::File::open(&obj).await?;
            let frame_buf_reader = BufReader::new(frame);

            let mut frameno = first_frame_no;
            let mut reader = bottomless::read::BatchReader::new(
                frameno,
                frame_buf_reader,
                page_size as usize,
                compression_kind,
            );

            while let Some(frame) = reader.next_frame_header().await? {
                let pgno = frame.pgno();
                reader.next_page(&mut page_buf).await?;
                if let Some(ck) = checksum {
                    checksum = match frame.verify(ck, &page_buf) {
                        Ok(checksum) => Some(checksum),
                        Err(e) => {
                            println!("ERROR: failed to verify checksum of page {pgno}: {e}, continuing anyway. Checksum will no longer be validated");
                            tracing::error!("Failed to verify checksum of page {pgno}: {e}, continuing anyway. Checksum will no longer be validated");
                            None
                        }
                    };
                }
                pending_pages.insert(pgno, &page_buf).await?;
                if frame.is_committed() {
                    let pending_pages = std::mem::replace(
                        &mut pending_pages,
                        TransactionPageCache::new(SWAP_AFTER, page_size, TMP_RESTORE_DIR.into()),
                    );
                    pending_pages.flush(db).await?;
                }
                frameno += 1;
                last_received_frame_no += 1;
            }
            db.flush().await?;
        }
        Ok(last_received_frame_no)
    }

    pub async fn get_local_metadata(
        from_dir: impl AsRef<std::path::Path>,
    ) -> Result<Option<(u32, (u32, u32))>> {
        use bytes::Buf;

        if let Ok(data) = tokio::fs::read(from_dir.as_ref().join(".meta")).await {
            let mut data = bytes::Bytes::from(data);
            let page_size = data.get_u32();
            let crc = (data.get_u32(), data.get_u32());
            Ok(Some((page_size, crc)))
        } else {
            Ok(None)
        }
    }
}