cli/data/

snapshot_storage.rs

// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Storage abstraction for Export/Import V2.
//!
//! This module provides a unified interface for reading and writing snapshot data
//! to various storage backends (S3, OSS, GCS, Azure Blob, local filesystem).

use std::collections::BTreeSet;
use std::path::Component;

use async_trait::async_trait;
use common_datasource::packed_writer::WRITE_BYTES;
use futures::TryStreamExt;
use object_store::services::{Azblob, Fs, Gcs, Oss, S3};
use object_store::util::{with_instrument_layers, with_retry_layers};
use object_store::{
    AzblobConnection, ErrorKind, GcsConnection, ObjectStore, OssConnection, S3Connection,
};
use snafu::ResultExt;
use url::Url;

use crate::common::ObjectStoreConfig;
use crate::data::export_v2::error::{
    BuildObjectStoreSnafu, InvalidUriSnafu, ManifestParseSnafu, ManifestSerializeSnafu, Result,
    SnapshotNotFoundSnafu, StorageOperationSnafu, TextDecodeSnafu, UnsupportedSchemeSnafu,
    UrlParseSnafu,
};
use crate::data::export_v2::manifest::{MANIFEST_FILE, Manifest};
#[cfg(test)]
use crate::data::export_v2::schema::SchemaDefinition;
use crate::data::export_v2::schema::{SCHEMA_DIR, SCHEMAS_FILE, SchemaSnapshot};
use crate::data::path::data_dir_for_schema_chunk;

struct RemoteLocation {
    bucket_or_container: String,
    root: String,
}

/// URI schemes supported for snapshot storage.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StorageScheme {
    /// Amazon S3.
    S3,
    /// Alibaba Cloud OSS.
    Oss,
    /// Google Cloud Storage.
    Gcs,
    /// Azure Blob Storage.
    Azblob,
    /// Local filesystem (file://).
    File,
}

impl StorageScheme {
    /// Parses storage scheme from URI.
    pub fn from_uri(uri: &str) -> Result<Self> {
        let url = Url::parse(uri).context(UrlParseSnafu)?;

        match url.scheme() {
            "s3" => Ok(Self::S3),
            "oss" => Ok(Self::Oss),
            "gs" | "gcs" => Ok(Self::Gcs),
            "azblob" => Ok(Self::Azblob),
            "file" => Ok(Self::File),
            scheme => UnsupportedSchemeSnafu { scheme }.fail(),
        }
    }
}

/// Extracts bucket/container and root path from a URI.
fn extract_remote_location_with_root_policy(
    uri: &str,
    allow_empty_root: bool,
) -> Result<RemoteLocation> {
    let url = Url::parse(uri).context(UrlParseSnafu)?;
    let bucket_or_container = url.host_str().unwrap_or("").to_string();
    if bucket_or_container.is_empty() {
        return InvalidUriSnafu {
            uri,
            reason: "URI must include bucket/container in host",
        }
        .fail();
    }

    let root = url.path().trim_start_matches('/').to_string();
    if root.is_empty() && !allow_empty_root {
        return InvalidUriSnafu {
            uri,
            reason: "snapshot URI must include a non-empty path after the bucket/container",
        }
        .fail();
    }

    Ok(RemoteLocation {
        bucket_or_container,
        root,
    })
}

/// Validates that a URI has a proper scheme.
///
/// Rejects bare paths (e.g., `/tmp/backup`, `./backup`) because:
/// - Schema export (CLI) and data export (server) run in different processes
/// - Using bare paths would split the snapshot across machines
///
/// Supported URI schemes:
/// - `s3://bucket/path` - Amazon S3
/// - `oss://bucket/path` - Alibaba Cloud OSS
/// - `gs://bucket/path` - Google Cloud Storage
/// - `azblob://container/path` - Azure Blob Storage
/// - `file:///absolute/path` - Local filesystem
pub fn validate_uri(uri: &str) -> Result<StorageScheme> {
    // Must have a scheme
    if !uri.contains("://") {
        return InvalidUriSnafu {
            uri,
            reason: "URI must have a scheme (e.g., s3://, file://). Bare paths are not supported.",
        }
        .fail();
    }

    StorageScheme::from_uri(uri)
}

/// Validates a URI for snapshot-scoped destructive operations.
///
/// Unlike read-only parent scans, destructive commands must target a concrete
/// snapshot directory instead of a bucket/container root or filesystem root.
/// Remote storage buckets/containers already provide namespace isolation, so a
/// non-empty object prefix is enough; local filesystem paths require at least
/// two non-root path segments to avoid deleting broad system directories.
pub fn validate_snapshot_uri(uri: &str) -> Result<StorageScheme> {
    let scheme = validate_uri(uri)?;
    reject_query_or_fragment(uri)?;
    match scheme {
        StorageScheme::File => validate_file_snapshot_uri(uri)?,
        StorageScheme::S3 | StorageScheme::Oss | StorageScheme::Gcs | StorageScheme::Azblob => {
            extract_remote_location_with_root_policy(uri, false)?;
        }
    }
    Ok(scheme)
}

fn reject_query_or_fragment(uri: &str) -> Result<()> {
    let url = Url::parse(uri).context(UrlParseSnafu)?;
    if url.query().is_some() || url.fragment().is_some() {
        return InvalidUriSnafu {
            uri,
            reason: "snapshot URI must not include query or fragment",
        }
        .fail();
    }

    Ok(())
}

fn validate_file_snapshot_uri(uri: &str) -> Result<()> {
    if has_explicit_dot_segment(uri) {
        return InvalidUriSnafu {
            uri,
            reason: "file snapshot URI must not contain '.' or '..' path segments",
        }
        .fail();
    }

    let path = extract_file_path_from_uri(uri)?;
    let mut normal_component_count = 0;

    // This is only a path-shape guard for destructive operations. It does not
    // resolve symlinks. Drive prefixes and root separators also do not count
    // toward depth; delete still relies on the manifest check and explicit
    // confirmation before removing the rooted storage prefix.
    for component in std::path::Path::new(&path).components() {
        match component {
            Component::Normal(_) => normal_component_count += 1,
            Component::CurDir | Component::ParentDir => {
                return InvalidUriSnafu {
                    uri,
                    reason: "file snapshot URI must not contain '.' or '..' path segments",
                }
                .fail();
            }
            Component::Prefix(_) | Component::RootDir => {}
        }
    }

    if normal_component_count < 2 {
        return InvalidUriSnafu {
            uri,
            reason: "file snapshot URI must point to a directory at least two levels deep",
        }
        .fail();
    }

    Ok(())
}

fn has_explicit_dot_segment(uri: &str) -> bool {
    // Defense in depth: catch dot segments at the raw URI level before
    // `Url::to_file_path()` can normalize them away. The `Path::components()`
    // check below still runs because URL decoding can reintroduce them.
    let without_fragment = uri.split_once('#').map_or(uri, |(path, _)| path);
    let path = without_fragment
        .split_once('?')
        .map_or(without_fragment, |(path, _)| path);

    path.split('/')
        .any(|segment| segment == "." || segment == "..")
}

fn schema_index_path() -> String {
    format!("{}/{}", SCHEMA_DIR, SCHEMAS_FILE)
}

/// Extracts the absolute filesystem path from a file:// URI.
fn extract_file_path_from_uri(uri: &str) -> Result<String> {
    let url = Url::parse(uri).context(UrlParseSnafu)?;

    match url.host_str() {
        Some(host) if !host.is_empty() && host != "localhost" => InvalidUriSnafu {
            uri,
            reason: "file:// URI must use an absolute path like file:///tmp/backup",
        }
        .fail(),
        _ => url
            .to_file_path()
            .map_err(|_| {
                InvalidUriSnafu {
                    uri,
                    reason: "file:// URI must use an absolute path like file:///tmp/backup",
                }
                .build()
            })
            .map(|path| path.to_string_lossy().into_owned()),
    }
}

async fn ensure_snapshot_exists(storage: &OpenDalStorage) -> Result<()> {
    if storage.exists().await? {
        Ok(())
    } else {
        SnapshotNotFoundSnafu {
            uri: storage.target_uri.as_str(),
        }
        .fail()
    }
}

/// Snapshot storage abstraction.
///
/// Provides operations for reading and writing snapshot data to various storage backends.
#[async_trait]
pub trait SnapshotStorage: Send + Sync {
    /// Checks if a snapshot exists at this location (manifest.json exists).
    async fn exists(&self) -> Result<bool>;

    /// Reads the manifest file.
    async fn read_manifest(&self) -> Result<Manifest>;

    /// Returns file length, or None for a missing path or directory.
    async fn file_size(&self, path: &str) -> Result<Option<u64>>;

    /// Writes the manifest file.
    async fn write_manifest(&self, manifest: &Manifest) -> Result<()>;

    /// Writes the schema index to schema/schemas.json.
    async fn write_schema(&self, schema: &SchemaSnapshot) -> Result<()>;

    /// Writes a text file to a relative path under the snapshot root.
    async fn write_text(&self, path: &str, content: &str) -> Result<()>;

    /// Reads a text file from a relative path under the snapshot root.
    async fn read_text(&self, path: &str) -> Result<String>;

    /// Creates a directory-like prefix under the snapshot root when needed by the backend.
    async fn create_dir_all(&self, path: &str) -> Result<()>;

    /// Lists files recursively under a relative prefix.
    async fn list_files_recursive(&self, prefix: &str) -> Result<Vec<String>>;

    /// Checks fresh output directories or removes a terminated, unfinished chunk's files.
    async fn prepare_export_chunk(
        &self,
        schemas: &[String],
        chunk_id: u32,
        resume: bool,
        packed: bool,
    ) -> Result<()> {
        let _ = (schemas, chunk_id, resume, packed);
        InvalidUriSnafu {
            uri: "snapshot",
            reason: "storage does not support preparing export chunks",
        }
        .fail()
    }

    /// Deletes the entire snapshot (for --force).
    async fn delete_snapshot(&self) -> Result<()>;
}

/// OpenDAL-based implementation of SnapshotStorage.
pub struct OpenDalStorage {
    object_store: ObjectStore,
    target_uri: String,
}

impl OpenDalStorage {
    fn new_operator_rooted(object_store: ObjectStore, target_uri: &str) -> Self {
        Self {
            object_store,
            target_uri: target_uri.to_string(),
        }
    }

    fn finish_local_store(object_store: ObjectStore) -> ObjectStore {
        with_instrument_layers(object_store, false)
    }

    fn finish_remote_store(object_store: ObjectStore) -> ObjectStore {
        with_instrument_layers(with_retry_layers(object_store), false)
    }

    fn ensure_backend_enabled(uri: &str, enabled: bool, reason: &'static str) -> Result<()> {
        if enabled {
            Ok(())
        } else {
            InvalidUriSnafu { uri, reason }.fail()
        }
    }

    fn validate_remote_config<E: std::fmt::Display>(
        uri: &str,
        backend: &str,
        result: std::result::Result<(), E>,
    ) -> Result<()> {
        result.map_err(|error| {
            InvalidUriSnafu {
                uri,
                reason: format!("invalid {} config: {}", backend, error),
            }
            .build()
        })
    }

    /// Creates a new storage from a file:// URI.
    pub fn from_file_uri(uri: &str) -> Result<Self> {
        let path = extract_file_path_from_uri(uri)?;

        let builder = Fs::default().root(&path);
        let object_store = ObjectStore::new(builder).context(BuildObjectStoreSnafu)?;
        Ok(Self::new_operator_rooted(
            Self::finish_local_store(object_store),
            uri,
        ))
    }

    fn from_file_uri_with_config(uri: &str, storage: &ObjectStoreConfig) -> Result<Self> {
        if storage.enable_s3 || storage.enable_oss || storage.enable_gcs || storage.enable_azblob {
            return InvalidUriSnafu {
                uri,
                reason: "file:// cannot be used with remote storage flags",
            }
            .fail();
        }

        Self::from_file_uri(uri)
    }

    fn from_s3_uri(uri: &str, storage: &ObjectStoreConfig) -> Result<Self> {
        Self::from_s3_uri_with_root_policy(uri, storage, false)
    }

    fn from_s3_uri_with_root_policy(
        uri: &str,
        storage: &ObjectStoreConfig,
        allow_empty_root: bool,
    ) -> Result<Self> {
        Self::ensure_backend_enabled(
            uri,
            storage.enable_s3,
            "s3:// requires --s3 and related options",
        )?;

        let location = extract_remote_location_with_root_policy(uri, allow_empty_root)?;
        let mut config = storage.s3.clone();
        config.s3_bucket = location.bucket_or_container;
        config.s3_root = location.root;
        Self::validate_remote_config(uri, "s3", config.validate())?;

        let conn: S3Connection = config.into();
        let object_store = ObjectStore::new(S3::from(&conn)).context(BuildObjectStoreSnafu)?;
        Ok(Self::new_operator_rooted(
            Self::finish_remote_store(object_store),
            uri,
        ))
    }

    fn from_oss_uri(uri: &str, storage: &ObjectStoreConfig) -> Result<Self> {
        Self::from_oss_uri_with_root_policy(uri, storage, false)
    }

    fn from_oss_uri_with_root_policy(
        uri: &str,
        storage: &ObjectStoreConfig,
        allow_empty_root: bool,
    ) -> Result<Self> {
        Self::ensure_backend_enabled(
            uri,
            storage.enable_oss,
            "oss:// requires --oss and related options",
        )?;

        let location = extract_remote_location_with_root_policy(uri, allow_empty_root)?;
        let mut config = storage.oss.clone();
        config.oss_bucket = location.bucket_or_container;
        config.oss_root = location.root;
        Self::validate_remote_config(uri, "oss", config.validate())?;

        let conn: OssConnection = config.into();
        let object_store = ObjectStore::new(Oss::from(&conn)).context(BuildObjectStoreSnafu)?;
        Ok(Self::new_operator_rooted(
            Self::finish_remote_store(object_store),
            uri,
        ))
    }

    fn from_gcs_uri(uri: &str, storage: &ObjectStoreConfig) -> Result<Self> {
        Self::from_gcs_uri_with_root_policy(uri, storage, false)
    }

    fn from_gcs_uri_with_root_policy(
        uri: &str,
        storage: &ObjectStoreConfig,
        allow_empty_root: bool,
    ) -> Result<Self> {
        Self::ensure_backend_enabled(
            uri,
            storage.enable_gcs,
            "gs:// or gcs:// requires --gcs and related options",
        )?;

        let location = extract_remote_location_with_root_policy(uri, allow_empty_root)?;
        let mut config = storage.gcs.clone();
        config.gcs_bucket = location.bucket_or_container;
        config.gcs_root = location.root;
        // GCS validate() rejects empty root, unlike S3/OSS/Azblob.
        if allow_empty_root && config.gcs_root.is_empty() {
            Self::validate_gcs_parent_config(uri, &config)?;
        } else {
            Self::validate_remote_config(uri, "gcs", config.validate())?;
        }

        let conn: GcsConnection = config.into();
        let object_store = ObjectStore::new(Gcs::from(&conn)).context(BuildObjectStoreSnafu)?;
        Ok(Self::new_operator_rooted(
            Self::finish_remote_store(object_store),
            uri,
        ))
    }

    fn validate_gcs_parent_config(
        uri: &str,
        config: &crate::common::PrefixedGcsConnection,
    ) -> Result<()> {
        if config.gcs_bucket.is_empty() {
            return InvalidUriSnafu {
                uri,
                reason: "invalid gcs config: GCS bucket must be set when --gcs is enabled.",
            }
            .fail();
        }
        if config.gcs_scope.is_empty() {
            return InvalidUriSnafu {
                uri,
                reason: "invalid gcs config: GCS scope must be set when --gcs is enabled.",
            }
            .fail();
        }
        Ok(())
    }

    fn from_azblob_uri(uri: &str, storage: &ObjectStoreConfig) -> Result<Self> {
        Self::from_azblob_uri_with_root_policy(uri, storage, false)
    }

    fn from_azblob_uri_with_root_policy(
        uri: &str,
        storage: &ObjectStoreConfig,
        allow_empty_root: bool,
    ) -> Result<Self> {
        Self::ensure_backend_enabled(
            uri,
            storage.enable_azblob,
            "azblob:// requires --azblob and related options",
        )?;

        let location = extract_remote_location_with_root_policy(uri, allow_empty_root)?;
        let mut config = storage.azblob.clone();
        config.azblob_container = location.bucket_or_container;
        config.azblob_root = location.root;
        Self::validate_remote_config(uri, "azblob", config.validate())?;

        let conn: AzblobConnection = config.into();
        let object_store = ObjectStore::new(Azblob::from(&conn)).context(BuildObjectStoreSnafu)?;
        Ok(Self::new_operator_rooted(
            Self::finish_remote_store(object_store),
            uri,
        ))
    }

    /// Creates a new storage from a URI and object store config.
    pub fn from_uri(uri: &str, storage: &ObjectStoreConfig) -> Result<Self> {
        match StorageScheme::from_uri(uri)? {
            StorageScheme::File => Self::from_file_uri_with_config(uri, storage),
            StorageScheme::S3 => Self::from_s3_uri(uri, storage),
            StorageScheme::Oss => Self::from_oss_uri(uri, storage),
            StorageScheme::Gcs => Self::from_gcs_uri(uri, storage),
            StorageScheme::Azblob => Self::from_azblob_uri(uri, storage),
        }
    }

    /// Creates storage rooted at a snapshot parent URI.
    ///
    /// Parent-oriented commands such as `export-v2 list` may scan bucket/container
    /// roots. Snapshot-oriented commands must keep using `from_uri`, which rejects
    /// empty remote roots to avoid unsafe snapshot operations at bucket scope.
    pub fn from_parent_uri(uri: &str, storage: &ObjectStoreConfig) -> Result<Self> {
        match StorageScheme::from_uri(uri)? {
            StorageScheme::File => Self::from_file_uri_with_config(uri, storage),
            StorageScheme::S3 => Self::from_s3_uri_with_root_policy(uri, storage, true),
            StorageScheme::Oss => Self::from_oss_uri_with_root_policy(uri, storage, true),
            StorageScheme::Gcs => Self::from_gcs_uri_with_root_policy(uri, storage, true),
            StorageScheme::Azblob => Self::from_azblob_uri_with_root_policy(uri, storage, true),
        }
    }

    /// Reads a file as bytes.
    async fn read_file(&self, path: &str) -> Result<Vec<u8>> {
        let data = self
            .object_store
            .read(path)
            .await
            .context(StorageOperationSnafu {
                operation: format!("read {}", path),
            })?;
        Ok(data.to_vec())
    }

    /// Reads a file as bytes if it exists.
    pub(crate) async fn read_file_if_exists(&self, path: &str) -> Result<Option<Vec<u8>>> {
        match self.object_store.read(path).await {
            Ok(data) => Ok(Some(data.to_vec())),
            Err(error) if error.kind() == ErrorKind::NotFound => Ok(None),
            Err(error) => Err(error).context(StorageOperationSnafu {
                operation: format!("read {}", path),
            }),
        }
    }

    /// Writes bytes to a file.
    async fn write_file(&self, path: &str, data: Vec<u8>) -> Result<()> {
        let mut writer = self
            .object_store
            .writer_with(path)
            .chunk(WRITE_BYTES)
            .concurrent(1)
            .await
            .with_context(|_| StorageOperationSnafu {
                operation: format!("open writer {}", path),
            })?;
        let data = object_store::Buffer::from(data);
        let result = async {
            for chunk in data.chunks(WRITE_BYTES) {
                writer.write(chunk).await?;
            }
            writer.close().await.map(|_| ())
        }
        .await;
        if result.is_err()
            && let Err(error) = writer.abort().await
        {
            common_telemetry::warn!(error; "Failed to abort snapshot metadata write");
        }
        result.with_context(|_| StorageOperationSnafu {
            operation: format!("write {}", path),
        })
    }

    /// Checks if a file exists using stat.
    pub(crate) async fn file_exists(&self, path: &str) -> Result<bool> {
        match self.object_store.stat(path).await {
            Ok(metadata) => Ok(!metadata.is_dir()),
            Err(e) if e.kind() == object_store::ErrorKind::NotFound => Ok(false),
            Err(e) => Err(e).context(StorageOperationSnafu {
                operation: format!("check exists {}", path),
            }),
        }
    }

    /// Iterates files recursively under a relative prefix without materializing
    /// the full listing first.
    pub(crate) async fn for_each_file_recursive<F>(&self, prefix: &str, mut f: F) -> Result<()>
    where
        F: FnMut(String) -> Result<()>,
    {
        let mut lister = match self.object_store.lister_with(prefix).recursive(true).await {
            Ok(lister) => lister,
            Err(error) if error.kind() == ErrorKind::NotFound => return Ok(()),
            Err(error) => {
                return Err(error).context(StorageOperationSnafu {
                    operation: format!("list {}", prefix),
                });
            }
        };

        while let Some(entry) = lister.try_next().await.context(StorageOperationSnafu {
            operation: format!("list {}", prefix),
        })? {
            if entry.metadata().is_dir() {
                continue;
            }
            f(entry.path().to_string())?;
        }

        Ok(())
    }

    /// Lists direct child directory names under the storage root.
    pub(crate) async fn list_direct_child_dirs(&self) -> Result<Vec<String>> {
        let mut lister = match self.object_store.lister_with("/").recursive(false).await {
            Ok(lister) => lister,
            Err(error) if error.kind() == ErrorKind::NotFound => return Ok(Vec::new()),
            Err(error) => {
                return Err(error).context(StorageOperationSnafu {
                    operation: "list /",
                });
            }
        };

        let mut dirs = BTreeSet::new();
        while let Some(entry) = lister.try_next().await.context(StorageOperationSnafu {
            operation: "list /",
        })? {
            let path = entry.path().trim_matches('/');
            if path.is_empty() {
                continue;
            }

            if entry.metadata().is_dir()
                && let Some(name) = path.split('/').next()
            {
                dirs.insert(name.to_string());
            }
        }

        Ok(dirs.into_iter().collect())
    }

    #[cfg(test)]
    pub async fn read_schema(&self) -> Result<SchemaSnapshot> {
        let schemas_path = schema_index_path();
        let schemas: Vec<SchemaDefinition> = if self.file_exists(&schemas_path).await? {
            let data = self.read_file(&schemas_path).await?;
            serde_json::from_slice(&data).context(ManifestParseSnafu)?
        } else {
            vec![]
        };

        Ok(SchemaSnapshot { schemas })
    }
}

#[async_trait]
impl SnapshotStorage for OpenDalStorage {
    async fn exists(&self) -> Result<bool> {
        self.file_exists(MANIFEST_FILE).await
    }

    async fn read_manifest(&self) -> Result<Manifest> {
        ensure_snapshot_exists(self).await?;

        let data = self.read_file(MANIFEST_FILE).await?;
        Manifest::from_json(&data).context(ManifestParseSnafu)
    }

    async fn file_size(&self, path: &str) -> Result<Option<u64>> {
        match self.object_store.stat(path).await {
            Ok(metadata) => Ok(metadata.is_file().then(|| metadata.content_length())),
            Err(error) if error.kind() == ErrorKind::NotFound => Ok(None),
            Err(error) => Err(error).with_context(|_| StorageOperationSnafu {
                operation: format!("stat {path}"),
            }),
        }
    }

    async fn write_manifest(&self, manifest: &Manifest) -> Result<()> {
        let data = serde_json::to_vec_pretty(manifest).context(ManifestSerializeSnafu)?;
        self.write_file(MANIFEST_FILE, data).await
    }

    async fn write_schema(&self, schema: &SchemaSnapshot) -> Result<()> {
        let schemas_path = schema_index_path();
        let schemas_data =
            serde_json::to_vec_pretty(&schema.schemas).context(ManifestSerializeSnafu)?;
        self.write_file(&schemas_path, schemas_data).await
    }

    async fn write_text(&self, path: &str, content: &str) -> Result<()> {
        self.write_file(path, content.as_bytes().to_vec()).await
    }

    async fn read_text(&self, path: &str) -> Result<String> {
        let data = self.read_file(path).await?;
        String::from_utf8(data).context(TextDecodeSnafu)
    }

    async fn create_dir_all(&self, path: &str) -> Result<()> {
        self.object_store
            .create_dir(path)
            .await
            .context(StorageOperationSnafu {
                operation: format!("create dir {}", path),
            })
    }

    async fn list_files_recursive(&self, prefix: &str) -> Result<Vec<String>> {
        let mut files = Vec::new();
        self.for_each_file_recursive(prefix, |path| {
            files.push(path);
            Ok(())
        })
        .await?;
        Ok(files)
    }

    async fn prepare_export_chunk(
        &self,
        schemas: &[String],
        chunk_id: u32,
        resume: bool,
        packed: bool,
    ) -> Result<()> {
        let mut files = Vec::new();
        for schema in schemas {
            let prefix = data_dir_for_schema_chunk(schema, chunk_id);
            let mut entries = match self.object_store.lister_with(&prefix).recursive(true).await {
                Ok(entries) => entries,
                Err(error) if error.kind() == ErrorKind::NotFound => continue,
                Err(error) => {
                    return Err(error).context(StorageOperationSnafu {
                        operation: format!("list {prefix}"),
                    });
                }
            };
            while let Some(entry) = entries.try_next().await.context(StorageOperationSnafu {
                operation: format!("list {prefix}"),
            })? {
                let path = entry.path();
                if path == prefix && entry.metadata().is_dir() {
                    continue;
                }
                let name = path.strip_prefix(&prefix).unwrap_or("");
                if !resume
                    || entry.metadata().is_dir()
                    || !(if packed {
                        valid_packed_chunk_filename(name)
                    } else {
                        valid_chunk_filename(name)
                    })
                {
                    return InvalidUriSnafu {
                        uri: path,
                        reason: "expected an empty new chunk or recognized files in an owned unfinished chunk",
                    }.fail();
                }
                files.push(path.to_string());
            }
        }
        // Validate every schema before deleting any files; COPY starts only after this returns.
        for path in files {
            match self.object_store.delete(&path).await {
                Ok(()) => {}
                Err(error) if error.kind() == ErrorKind::NotFound => {}
                Err(error) => {
                    return Err(error).context(StorageOperationSnafu {
                        operation: format!("delete unfinished chunk file {path}"),
                    });
                }
            }
        }
        Ok(())
    }

    async fn delete_snapshot(&self) -> Result<()> {
        self.object_store
            .delete_with("/")
            .recursive(true)
            .await
            .context(StorageOperationSnafu {
                operation: "delete snapshot",
            })
    }
}

fn valid_packed_chunk_filename(name: &str) -> bool {
    name == common_datasource::packed_snapshot::PACK_INDEX_FILE
        || [("pack-", ".bin"), ("table-", ".parquet")]
            .iter()
            .any(|(prefix, suffix)| {
                name.strip_prefix(prefix)
                    .and_then(|n| n.strip_suffix(suffix))
                    .is_some_and(|id| !id.is_empty() && id.bytes().all(|b| b.is_ascii_digit()))
            })
}

fn valid_chunk_filename(name: &str) -> bool {
    name.strip_suffix(".parquet")
        .is_some_and(|stem| !stem.is_empty())
        && !name.contains(['/', '\\', '\0'])
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;
    use std::path::Path;
    use std::sync::Arc;

    use object_store::ObjectStore;
    use object_store::layers::mock::MockLayerBuilder;
    use object_store::services::Fs;
    use tempfile::tempdir;
    use url::Url;

    use super::*;
    use crate::data::export_v2::error::Error;
    use crate::data::export_v2::manifest::{DataFormat, TimeRange};
    use crate::data::export_v2::schema::SchemaDefinition;

    fn make_storage_with_rooted_fs(dir: &std::path::Path) -> OpenDalStorage {
        let object_store = ObjectStore::new(Fs::default().root(dir.to_str().unwrap())).unwrap();
        OpenDalStorage::new_operator_rooted(
            OpenDalStorage::finish_local_store(object_store),
            Url::from_directory_path(dir).unwrap().as_ref(),
        )
    }

    #[tokio::test]
    async fn snapshot_metadata_submissions_are_bounded_and_invoke_abort_on_failure() {
        use std::sync::Mutex;
        use std::sync::atomic::{AtomicUsize, Ordering};

        use object_store::layers::mock::{Metadata, ObservedWriter, WriterObservations, oio};

        struct FailWriter {
            inner: oio::Writer,
            writes: usize,
            fail: Arc<AtomicUsize>,
        }
        impl oio::Write for FailWriter {
            async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> {
                self.writes += 1;
                if self.fail.load(Ordering::SeqCst) == 1 && self.writes == 2 {
                    return Err(object_store::Error::new(
                        ErrorKind::Unexpected,
                        "injected metadata write failure",
                    ));
                }
                self.inner.write(bytes).await
            }
            async fn close(&mut self) -> object_store::Result<Metadata> {
                if self.fail.load(Ordering::SeqCst) == 2 {
                    return Err(object_store::Error::new(
                        ErrorKind::Unexpected,
                        "injected metadata close failure",
                    ));
                }
                self.inner.close().await
            }
            async fn abort(&mut self) -> object_store::Result<()> {
                self.inner.abort().await
            }
        }

        let dir = tempdir().unwrap();
        let mut storage = make_storage_with_rooted_fs(dir.path());
        let observations = Arc::new(Mutex::new(WriterObservations::default()));
        let fail = Arc::new(AtomicUsize::new(0));
        storage.object_store = storage.object_store.layer(
            MockLayerBuilder::default()
                .writer_factory(Arc::new({
                    let observations = observations.clone();
                    let fail = fail.clone();
                    move |_, _, inner| {
                        Box::new(ObservedWriter::new(
                            Box::new(FailWriter {
                                inner,
                                writes: 0,
                                fail: fail.clone(),
                            }),
                            observations.clone(),
                        ))
                    }
                }))
                .build()
                .unwrap(),
        );
        for length in [2 * WRITE_BYTES + 17, 1, 0] {
            *observations.lock().unwrap() = WriterObservations::default();
            let data = vec![42; length];
            storage.write_file("metadata", data.clone()).await.unwrap();
            assert_eq!(storage.read_file("metadata").await.unwrap(), data);
            let observed = observations.lock().unwrap();
            assert_eq!(observed.write_sizes.iter().sum::<usize>(), length);
            assert!(observed.write_sizes.iter().all(|size| *size <= WRITE_BYTES));
            assert_eq!(observed.close_calls, 1);
            assert_eq!(observed.abort_calls, 0);
        }
        for failure in [1, 2] {
            *observations.lock().unwrap() = WriterObservations::default();
            fail.store(failure, Ordering::SeqCst);
            assert!(
                storage
                    .write_file("failed", vec![42; 2 * WRITE_BYTES + 17])
                    .await
                    .is_err()
            );
            assert_eq!(
                observations.lock().unwrap().write_sizes.len(),
                if failure == 1 { 2 } else { 3 }
            );
            let observed = observations.lock().unwrap();
            assert_eq!(observed.close_calls, usize::from(failure == 2));
            assert_eq!(observed.abort_calls, 1);
        }
    }

    #[test]
    fn test_validate_uri_valid() {
        assert_eq!(validate_uri("s3://bucket/path").unwrap(), StorageScheme::S3);
        assert_eq!(
            validate_uri("oss://bucket/path").unwrap(),
            StorageScheme::Oss
        );
        assert_eq!(
            validate_uri("gs://bucket/path").unwrap(),
            StorageScheme::Gcs
        );
        assert_eq!(
            validate_uri("gcs://bucket/path").unwrap(),
            StorageScheme::Gcs
        );
        assert_eq!(
            validate_uri("azblob://container/path").unwrap(),
            StorageScheme::Azblob
        );
        assert_eq!(
            validate_uri("file:///tmp/backup").unwrap(),
            StorageScheme::File
        );
    }

    #[test]
    fn test_validate_uri_invalid() {
        // Bare paths should be rejected
        assert!(validate_uri("/tmp/backup").is_err());
        assert!(validate_uri("./backup").is_err());
        assert!(validate_uri("backup").is_err());

        // Unknown schemes
        assert!(validate_uri("ftp://server/path").is_err());
    }

    #[test]
    fn test_extract_remote_location_requires_non_empty_root() {
        assert!(extract_remote_location_with_root_policy("s3://bucket", false).is_err());
        assert!(extract_remote_location_with_root_policy("s3://bucket/", false).is_err());
        assert!(extract_remote_location_with_root_policy("oss://bucket", false).is_err());
        assert!(extract_remote_location_with_root_policy("gs://bucket", false).is_err());
        assert!(extract_remote_location_with_root_policy("azblob://container", false).is_err());
    }

    #[test]
    fn test_extract_remote_location_allows_empty_root_when_permitted() {
        let location = extract_remote_location_with_root_policy("s3://bucket", true).unwrap();
        assert_eq!(location.bucket_or_container, "bucket");
        assert_eq!(location.root, "");

        let location =
            extract_remote_location_with_root_policy("azblob://container/", true).unwrap();
        assert_eq!(location.bucket_or_container, "container");
        assert_eq!(location.root, "");
    }

    #[test]
    fn test_parent_storage_allows_s3_bucket_root() {
        let mut storage = ObjectStoreConfig {
            enable_s3: true,
            ..Default::default()
        };
        storage.s3.s3_region = Some("us-east-1".to_string());

        assert!(OpenDalStorage::from_uri("s3://bucket", &storage).is_err());
        assert!(OpenDalStorage::from_parent_uri("s3://bucket", &storage).is_ok());
    }

    #[test]
    fn test_validate_snapshot_uri_rejects_dangerous_roots() {
        assert!(validate_snapshot_uri("s3://bucket").is_err());
        assert!(validate_snapshot_uri("s3://bucket/").is_err());
        assert!(validate_snapshot_uri("oss://bucket").is_err());
        assert!(validate_snapshot_uri("gs://bucket").is_err());
        assert!(validate_snapshot_uri("azblob://container").is_err());
        assert!(validate_snapshot_uri("s3://bucket/snapshot?version=1").is_err());
        assert!(validate_snapshot_uri("file:///tmp/backup#fragment").is_err());
        assert!(validate_snapshot_uri("file:///").is_err());
        assert!(validate_snapshot_uri("file:///tmp").is_err());
        assert!(validate_snapshot_uri("file:///tmp/backup/.").is_err());
        assert!(validate_snapshot_uri("file:///tmp/backup/..").is_err());
    }

    #[test]
    fn test_validate_snapshot_uri_accepts_snapshot_paths() {
        assert_eq!(
            validate_snapshot_uri("s3://bucket/snapshots/prod").unwrap(),
            StorageScheme::S3
        );

        let dir = tempdir().unwrap();
        let snapshot = dir.path().join("snapshot");
        std::fs::create_dir_all(&snapshot).unwrap();
        let uri = Url::from_directory_path(snapshot).unwrap().to_string();
        assert_eq!(validate_snapshot_uri(&uri).unwrap(), StorageScheme::File);
    }

    #[cfg(windows)]
    #[test]
    fn test_validate_snapshot_uri_windows_drive_prefix_depth() {
        assert!(validate_snapshot_uri("file:///C:/").is_err());
        assert!(validate_snapshot_uri("file:///C:/Users").is_err());
        assert!(validate_snapshot_uri("file:///C:/Users/snapshot").is_ok());
    }

    #[cfg(not(windows))]
    #[test]
    fn test_extract_path_from_uri_unix_examples() {
        assert_eq!(
            extract_file_path_from_uri("file:///tmp/backup").unwrap(),
            "/tmp/backup"
        );
        assert_eq!(
            extract_file_path_from_uri("file://localhost/tmp/backup").unwrap(),
            "/tmp/backup"
        );
        assert_eq!(
            extract_file_path_from_uri("file:///tmp/my%20backup").unwrap(),
            "/tmp/my backup"
        );
        assert_eq!(
            extract_file_path_from_uri("file://localhost/tmp/my%20backup").unwrap(),
            "/tmp/my backup"
        );
    }

    #[test]
    fn test_extract_file_path_from_uri_rejects_file_host() {
        assert!(extract_file_path_from_uri("file://tmp/backup").is_err());
    }

    #[test]
    fn test_extract_file_path_from_uri_round_trips_directory_url() {
        let dir = tempdir().unwrap();
        let uri = Url::from_directory_path(dir.path()).unwrap().to_string();
        let path = extract_file_path_from_uri(&uri).unwrap();

        assert_eq!(Path::new(&path), dir.path());
    }

    #[tokio::test]
    async fn test_read_manifest_reports_requested_uri() {
        let dir = tempdir().unwrap();
        let uri = Url::from_directory_path(dir.path()).unwrap().to_string();
        let storage = OpenDalStorage::from_file_uri(&uri).unwrap();

        let error = storage.read_manifest().await.unwrap_err().to_string();

        assert!(error.contains(uri.as_str()));
    }

    #[tokio::test]
    async fn test_manifest_round_trip() {
        let dir = tempdir().unwrap();
        let storage = make_storage_with_rooted_fs(dir.path());

        let manifest = Manifest::new_full(
            "greptime".to_string(),
            vec!["public".to_string()],
            TimeRange::unbounded(),
            DataFormat::Parquet,
        );

        storage.write_manifest(&manifest).await.unwrap();
        let loaded = storage.read_manifest().await.unwrap();

        assert_eq!(loaded.catalog, manifest.catalog);
        assert_eq!(loaded.schemas, manifest.schemas);
        assert_eq!(loaded.schema_only, manifest.schema_only);
        assert_eq!(loaded.format, manifest.format);
        assert_eq!(loaded.snapshot_id, manifest.snapshot_id);
    }

    #[tokio::test]
    async fn test_schema_round_trip() {
        let dir = tempdir().unwrap();
        let storage = make_storage_with_rooted_fs(dir.path());

        let mut snapshot = SchemaSnapshot::new();
        snapshot.add_schema(SchemaDefinition {
            catalog: "greptime".to_string(),
            name: "test_db".to_string(),
            options: HashMap::from([("ttl".to_string(), "7d".to_string())]),
        });

        storage.write_schema(&snapshot).await.unwrap();
        let loaded = storage.read_schema().await.unwrap();

        assert_eq!(loaded, snapshot);
    }

    #[tokio::test]
    async fn test_text_round_trip() {
        let dir = tempdir().unwrap();
        let storage = make_storage_with_rooted_fs(dir.path());
        let content = "CREATE TABLE metrics (ts TIMESTAMP TIME INDEX);";

        storage
            .write_text("schema/ddl/public.sql", content)
            .await
            .unwrap();
        let loaded = storage.read_text("schema/ddl/public.sql").await.unwrap();

        assert_eq!(loaded, content);
    }

    #[tokio::test]
    async fn test_read_text_rejects_invalid_utf8() {
        let dir = tempdir().unwrap();
        let storage = make_storage_with_rooted_fs(dir.path());

        storage
            .write_file("schema/ddl/public.sql", vec![0xff, 0xfe, 0xfd])
            .await
            .unwrap();

        let error = storage
            .read_text("schema/ddl/public.sql")
            .await
            .unwrap_err();
        assert!(error.to_string().contains("UTF-8"));
    }

    #[tokio::test]
    async fn test_exists_follows_manifest_presence() {
        let dir = tempdir().unwrap();
        let storage = make_storage_with_rooted_fs(dir.path());

        assert!(!storage.exists().await.unwrap());

        storage
            .write_manifest(&Manifest::new_schema_only(
                "greptime".to_string(),
                vec!["public".to_string()],
            ))
            .await
            .unwrap();

        assert!(storage.exists().await.unwrap());
    }

    #[tokio::test]
    async fn test_prepare_export_chunk_preserves_other_chunks_and_schema_artifacts() {
        let dir = tempdir().unwrap();
        let storage = make_storage_with_rooted_fs(dir.path());
        for path in [
            "data/public/1/keep.parquet",
            "data/public/2/cpu.v1.parquet",
            "data/other/2/a.parquet",
            "schema/ddl/public.sql",
        ] {
            storage.write_text(path, "original").await.unwrap();
        }
        let schemas = vec!["public".to_string(), "other".to_string()];
        assert!(
            storage
                .prepare_export_chunk(&schemas, 2, false, false)
                .await
                .is_err()
        );
        assert!(
            storage
                .file_exists("data/public/2/cpu.v1.parquet")
                .await
                .unwrap()
        );
        storage
            .prepare_export_chunk(&schemas, 2, true, false)
            .await
            .unwrap();
        storage
            .prepare_export_chunk(&schemas, 2, true, false)
            .await
            .unwrap();
        assert!(
            !storage
                .file_exists("data/public/2/cpu.v1.parquet")
                .await
                .unwrap()
        );
        assert!(!storage.file_exists("data/other/2/a.parquet").await.unwrap());
        assert_eq!(
            storage
                .read_text("data/public/1/keep.parquet")
                .await
                .unwrap(),
            "original"
        );
        assert_eq!(
            storage.read_text("schema/ddl/public.sql").await.unwrap(),
            "original"
        );
    }

    #[tokio::test]
    async fn test_prepare_export_chunk_rejects_foreign_or_nested_entries_before_deletion() {
        for foreign in ["notes.txt", "nested/a.parquet", "empty/"] {
            let dir = tempdir().unwrap();
            let storage = make_storage_with_rooted_fs(dir.path());
            storage
                .write_text("data/public/2/cpu.parquet", "keep")
                .await
                .unwrap();
            let path = format!("data/other/2/{foreign}");
            if foreign.ends_with('/') {
                storage.create_dir_all(&path).await.unwrap();
            } else {
                storage.write_text(&path, "foreign").await.unwrap();
            }
            assert!(
                storage
                    .prepare_export_chunk(&["public".into(), "other".into()], 2, true, false)
                    .await
                    .is_err()
            );
            assert_eq!(
                storage
                    .read_text("data/public/2/cpu.parquet")
                    .await
                    .unwrap(),
                "keep"
            );
        }
    }

    #[tokio::test]
    async fn packed_retry_preserves_foreign_files_and_completed_chunks() {
        let dir = tempdir().unwrap();
        let storage = make_storage_with_rooted_fs(dir.path());
        let owned = ["pack-000000.bin", "table-42.parquet", "pack-index.json"];
        for name in owned {
            storage
                .write_text(&format!("data/public/2/{name}"), "partial")
                .await
                .unwrap();
        }
        for path in [
            "data/public/1/pack-000000.bin",
            "data/public/2/foreign.parquet",
        ] {
            storage.write_text(path, "keep").await.unwrap();
        }
        let schemas = ["public".into()];
        assert!(
            storage
                .prepare_export_chunk(&schemas, 2, true, true)
                .await
                .is_err()
        );
        for name in owned {
            assert!(
                storage
                    .file_exists(&format!("data/public/2/{name}"))
                    .await
                    .unwrap()
            );
        }
        assert_eq!(
            storage
                .read_text("data/public/2/foreign.parquet")
                .await
                .unwrap(),
            "keep"
        );
        storage
            .object_store
            .delete("data/public/2/foreign.parquet")
            .await
            .unwrap();
        storage
            .prepare_export_chunk(&schemas, 2, true, true)
            .await
            .unwrap();
        assert!(
            storage
                .list_files_recursive("data/public/2/")
                .await
                .unwrap()
                .is_empty()
        );
        assert_eq!(
            storage
                .read_text("data/public/1/pack-000000.bin")
                .await
                .unwrap(),
            "keep"
        );
    }

    #[tokio::test]
    async fn test_prepare_export_chunk_reports_delete_failure() {
        let dir = tempdir().unwrap();
        let mut storage = make_storage_with_rooted_fs(dir.path());
        storage.object_store = storage.object_store.layer(
            MockLayerBuilder::default()
                // OpenDAL's unit deleter returns Unsupported for every deletion.
                .deleter_factory(Arc::new(|_| Box::new(())))
                .build()
                .unwrap(),
        );
        storage
            .write_text("data/public/2/a.parquet", "keep")
            .await
            .unwrap();
        let error = storage
            .prepare_export_chunk(&["public".into()], 2, true, false)
            .await
            .unwrap_err();
        let Error::StorageOperation {
            operation, error, ..
        } = error
        else {
            panic!("expected storage operation error, got {error:?}");
        };
        assert_eq!(
            operation,
            "delete unfinished chunk file data/public/2/a.parquet"
        );
        assert_eq!(error.kind(), ErrorKind::Unsupported);
        assert_eq!(
            storage.read_text("data/public/2/a.parquet").await.unwrap(),
            "keep"
        );
    }

    #[tokio::test]
    async fn test_prepare_export_chunk_encodes_schema_boundary() {
        let dir = tempdir().unwrap();
        let storage = make_storage_with_rooted_fs(dir.path());
        let schema = "../other".to_string();
        let owned = format!("{}a.parquet", data_dir_for_schema_chunk(&schema, 2));
        storage.write_text(&owned, "remove").await.unwrap();
        storage
            .write_text("data/other/2/a.parquet", "keep")
            .await
            .unwrap();
        storage
            .prepare_export_chunk(&[schema], 2, true, false)
            .await
            .unwrap();
        assert!(!storage.file_exists(&owned).await.unwrap());
        assert_eq!(
            storage.read_text("data/other/2/a.parquet").await.unwrap(),
            "keep"
        );
    }

    #[tokio::test]
    async fn test_delete_snapshot_only_removes_rooted_contents() {
        let parent = tempdir().unwrap();
        let snapshot_root = parent.path().join("snapshot");
        let sibling = parent.path().join("sibling");
        std::fs::create_dir_all(&snapshot_root).unwrap();
        std::fs::create_dir_all(&sibling).unwrap();
        std::fs::write(snapshot_root.join("manifest.json"), b"{}").unwrap();
        std::fs::write(sibling.join("keep.txt"), b"keep").unwrap();

        let storage = make_storage_with_rooted_fs(&snapshot_root);
        storage.delete_snapshot().await.unwrap();

        assert!(!snapshot_root.join("manifest.json").exists());
        assert!(sibling.join("keep.txt").exists());
    }
}