remora/crates/remora-target/src/store.rs

use std::collections::HashMap;
use std::num::NonZeroUsize;
use std::time::Duration;

use futures_util::TryStreamExt;
use opendal::services::S3;
use opendal::{Metadata, Operator, Writer};
use remora_domain::OpaqueId;
use remora_domain::ports::StorageAdapter;

use crate::body::{OpenDalTargetBody, StoredWrite};
use crate::config::TargetStoreConfig;
use crate::digest::{ContentHasher, ContentSha256, check_expected, hash_bytes};
use crate::error::{TargetError, map_opendal, redact_sensitive};
use crate::identity::{
    CompletedPart, MultipartIdentity, ObjectIdentity, REMORA_SHA256_METADATA_KEY,
    attempt_manifest_key, attempt_object_key, validate_attempt_key,
};
use crate::multipart::MultipartClient;

/// Request for a create-only object write under a unique reserved key.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct CreateObjectRequest {
    key: String,
    expected_content_sha256: Option<ContentSha256>,
    chunk_size: Option<NonZeroUsize>,
}

impl CreateObjectRequest {
    pub fn parse(key: impl Into<String>) -> Result<Self, TargetError> {
        let key = key.into();
        validate_attempt_key(&key)?;
        Ok(Self {
            key,
            expected_content_sha256: None,
            chunk_size: None,
        })
    }

    #[must_use]
    pub fn for_attempt_object(
        attempt_id: &OpaqueId,
        object_id: &OpaqueId,
        relative_path: Option<&str>,
    ) -> Self {
        Self {
            key: attempt_object_key(attempt_id, object_id, relative_path),
            expected_content_sha256: None,
            chunk_size: None,
        }
    }

    #[must_use]
    pub fn for_attempt_manifest(attempt_id: &OpaqueId, manifest_id: &OpaqueId) -> Self {
        Self {
            key: attempt_manifest_key(attempt_id, manifest_id),
            expected_content_sha256: None,
            chunk_size: None,
        }
    }

    #[must_use]
    pub fn with_expected_content_sha256(mut self, expected: ContentSha256) -> Self {
        self.expected_content_sha256 = Some(expected);
        self
    }

    #[must_use]
    pub fn with_chunk_size(mut self, chunk_size: NonZeroUsize) -> Self {
        self.chunk_size = Some(chunk_size);
        self
    }

    #[must_use]
    pub fn key(&self) -> &str {
        &self.key
    }
}

/// Stored object returned after a create-only write completes.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct StoredObject {
    pub identity: ObjectIdentity,
    /// Independently computed SHA-256 of the written bytes.
    pub content_sha256: ContentSha256,
    pub size: u64,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ObjectMetadata {
    pub identity: ObjectIdentity,
    pub content_sha256: Option<ContentSha256>,
    pub size: u64,
}

/// Time-limited exact-object access URL.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct TimeLimitedAccess {
    pub method: String,
    pub url: String,
    pub expires_in: Duration,
}

/// Recoverable multipart upload bound to one exact target identity.
pub struct MultipartSession {
    identity: MultipartIdentity,
    expected_sha256: Option<ContentSha256>,
}

impl std::fmt::Debug for MultipartSession {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("MultipartSession")
            .field("identity", &self.identity)
            .field("expected_sha256", &self.expected_sha256)
            .finish_non_exhaustive()
    }
}

impl MultipartSession {
    #[must_use]
    pub fn identity(&self) -> &MultipartIdentity {
        &self.identity
    }
}

/// OpenDAL-backed target storage adapter for attempt-scoped object writes.
#[derive(Clone)]
pub struct TargetStore {
    config: TargetStoreConfig,
    operator: Operator,
    multipart: MultipartClient,
}

impl TargetStore {
    pub fn connect(config: TargetStoreConfig) -> Result<Self, TargetError> {
        opendal::install_default();
        let mut builder = S3::default()
            .bucket(&config.bucket)
            .endpoint(&config.endpoint)
            .region(&config.region)
            .access_key_id(&config.access_key_id)
            .secret_access_key(&config.secret_access_key)
            .disable_config_load()
            .disable_ec2_metadata();
        if !config.root.is_empty() {
            builder = builder.root(&config.root);
        }
        let operator = Operator::new(builder)
            .map_err(|error| TargetError::validation(redact_sensitive(error.to_string())))?;
        let multipart = MultipartClient::new(config.clone())?;
        Ok(Self {
            config,
            operator,
            multipart,
        })
    }

    pub fn connect_from_env() -> Result<Self, TargetError> {
        Self::connect(TargetStoreConfig::from_env()?)
    }

    #[must_use]
    pub fn config(&self) -> &TargetStoreConfig {
        &self.config
    }

    /// Open a create-only streaming writer for one attempt object.
    pub async fn open_create_only(
        &self,
        request: &CreateObjectRequest,
    ) -> Result<OpenDalTargetBody, TargetError> {
        validate_attempt_key(&request.key)?;
        self.require_absent(&request.key).await?;
        let writer = self.open_writer(request).await?;
        Ok(OpenDalTargetBody::new(
            writer,
            self.new_object_identity(&request.key, None),
            request.expected_content_sha256.clone(),
        ))
    }

    /// Create-only put of a complete buffer.
    ///
    /// Independently hashes `body`, persists that digest as user metadata, and
    /// retains the version ID from the write response when present.
    pub async fn put_create_only(
        &self,
        request: &CreateObjectRequest,
        body: impl Into<bytes::Bytes>,
    ) -> Result<StoredObject, TargetError> {
        validate_attempt_key(&request.key)?;
        self.require_absent(&request.key).await?;
        let body = body.into();
        let computed = hash_bytes(&body);
        check_expected(request.expected_content_sha256.as_ref(), &computed)?;
        let size = body.len() as u64;
        let write = self
            .operator
            .write_with(&request.key, body)
            .if_not_exists(true)
            .user_metadata(sha256_metadata(&computed));
        let metadata = write.await.map_err(map_opendal)?;
        Ok(self.stored_from_write(&request.key, &computed, &metadata, size))
    }

    /// Exact identity read. Version is required when present on the identity.
    pub async fn read_exact(&self, identity: &ObjectIdentity) -> Result<bytes::Bytes, TargetError> {
        self.ensure_object_identity(identity)?;
        let mut read = self.operator.read_with(&identity.key);
        if let Some(version) = identity.version_id.as_deref() {
            read = read.version(version);
        }
        let buffer = read.await.map_err(map_opendal)?;
        Ok(buffer.to_bytes())
    }

    /// Exact identity delete. Missing objects are success for cleanup.
    pub async fn delete_exact(&self, identity: &ObjectIdentity) -> Result<(), TargetError> {
        self.ensure_object_identity(identity)?;
        let mut delete = self.operator.delete_with(&identity.key);
        if let Some(version) = identity.version_id.as_deref() {
            delete = delete.version(version);
        }
        match delete.await {
            Ok(()) => Ok(()),
            Err(error) if matches!(error.kind(), opendal::ErrorKind::NotFound) => Ok(()),
            Err(error) => Err(map_opendal(error)),
        }
    }

    /// Head an exact identity, returning size and retained SHA-256 metadata.
    pub async fn head_exact(
        &self,
        identity: &ObjectIdentity,
    ) -> Result<ObjectMetadata, TargetError> {
        self.ensure_object_identity(identity)?;
        let mut stat = self.operator.stat_with(&identity.key);
        if let Some(version) = identity.version_id.as_deref() {
            stat = stat.version(version);
        }
        let metadata = stat.await.map_err(map_opendal)?;
        // Prefer exact version retained on the identity when the head omits it.
        let version_id = metadata
            .version()
            .map(str::to_owned)
            .or_else(|| identity.version_id.clone());
        let content_sha256 = observed_sha256(&metadata)?;
        Ok(ObjectMetadata {
            identity: ObjectIdentity {
                target_store_id: self.config.target_store_id.clone(),
                bucket: self.config.bucket.clone(),
                key: identity.key.clone(),
                version_id,
            },
            content_sha256,
            size: metadata.content_length(),
        })
    }

    /// Stream and independently hash an exact identity without buffering it.
    pub async fn verify_exact(
        &self,
        identity: &ObjectIdentity,
    ) -> Result<ObjectMetadata, TargetError> {
        let (content_sha256, size) = self.hash_exact(identity).await?;
        Ok(ObjectMetadata {
            identity: identity.clone(),
            content_sha256: Some(content_sha256),
            size,
        })
    }

    /// Issue a time-limited URL bound to one exact object identity.
    pub async fn presign_read(
        &self,
        identity: &ObjectIdentity,
        expires_in: Duration,
    ) -> Result<TimeLimitedAccess, TargetError> {
        self.ensure_object_identity(identity)?;
        if expires_in.is_zero() {
            return Err(TargetError::validation(
                "time-limited access expiry must be positive",
            ));
        }
        let mut read_options = opendal::options::ReadOptions::default();
        if let Some(version) = identity.version_id.as_deref() {
            read_options.version = Some(version.to_owned());
        }
        let signed = self
            .operator
            .presign_read_options(&identity.key, expires_in, read_options)
            .await
            .map_err(map_opendal)?;
        Ok(TimeLimitedAccess {
            method: signed.method().as_str().to_owned(),
            url: signed.uri().to_string(),
            expires_in,
        })
    }

    /// Begin a create-only multipart upload under a reserved key.
    ///
    /// The completed exact object is independently stream-hashed after S3
    /// accepts the freshly re-listed part set.
    pub async fn begin_multipart(
        &self,
        request: &CreateObjectRequest,
    ) -> Result<MultipartSession, TargetError> {
        self.require_absent(&request.key).await?;
        let identity = self.multipart.initiate(&request.key, None).await?;
        Ok(MultipartSession {
            identity,
            expected_sha256: request.expected_content_sha256.clone(),
        })
    }

    /// Reconstruct a multipart session from its durable identity after loss.
    pub async fn resume_multipart(
        &self,
        identity: MultipartIdentity,
        expected_content_sha256: Option<ContentSha256>,
    ) -> Result<(MultipartSession, Vec<CompletedPart>), TargetError> {
        self.ensure_multipart_identity(&identity)?;
        validate_attempt_key(&identity.key)?;
        let parts = self.multipart.list_parts(&identity).await?;
        Ok((
            MultipartSession {
                identity,
                expected_sha256: expected_content_sha256,
            },
            parts,
        ))
    }

    pub async fn upload_part(
        &self,
        session: &MultipartSession,
        part_number: u32,
        body: bytes::Bytes,
    ) -> Result<CompletedPart, TargetError> {
        self.ensure_multipart_identity(&session.identity)?;
        self.multipart
            .upload_part(&session.identity, part_number, body)
            .await
    }

    pub async fn list_parts(
        &self,
        session: &MultipartSession,
    ) -> Result<Vec<CompletedPart>, TargetError> {
        self.ensure_multipart_identity(&session.identity)?;
        self.multipart.list_parts(&session.identity).await
    }

    pub async fn complete_multipart(
        &self,
        session: MultipartSession,
        parts: &[CompletedPart],
    ) -> Result<StoredObject, TargetError> {
        self.ensure_multipart_identity(&session.identity)?;
        self.require_absent(&session.identity.key).await?;
        let relisted = self.multipart.list_parts(&session.identity).await?;
        if relisted != parts {
            return Err(TargetError::integrity(
                "multipart completion parts do not match exact target re-list",
            ));
        }
        let version_id = self
            .multipart
            .complete(&session.identity, parts, true)
            .await?;
        let object = self.new_object_identity(&session.identity.key, version_id);
        let (computed, size) = self.hash_exact(&object).await?;
        if let Err(error) = check_expected(session.expected_sha256.as_ref(), &computed) {
            let _ = self.delete_exact(&object).await;
            return Err(error);
        }
        Ok(StoredObject {
            identity: object,
            content_sha256: computed,
            size,
        })
    }

    pub async fn abort_multipart(&self, session: MultipartSession) -> Result<(), TargetError> {
        self.ensure_multipart_identity(&session.identity)?;
        self.multipart.abort(&session.identity).await
    }

    async fn open_writer(&self, request: &CreateObjectRequest) -> Result<Writer, TargetError> {
        let chunk = request
            .chunk_size
            .map(NonZeroUsize::get)
            .unwrap_or(5 * 1024 * 1024);
        self.operator
            .writer_with(&request.key)
            .if_not_exists(true)
            .chunk(chunk)
            .await
            .map_err(map_opendal)
    }

    async fn require_absent(&self, key: &str) -> Result<(), TargetError> {
        match self.operator.stat(key).await {
            Ok(_) => Err(TargetError::AlreadyExists),
            Err(error) if matches!(error.kind(), opendal::ErrorKind::NotFound) => Ok(()),
            Err(error) => Err(map_opendal(error)),
        }
    }

    async fn hash_exact(
        &self,
        identity: &ObjectIdentity,
    ) -> Result<(ContentSha256, u64), TargetError> {
        self.ensure_object_identity(identity)?;
        let mut reader = self.operator.reader_with(&identity.key).chunk(1024 * 1024);
        if let Some(version) = identity.version_id.as_deref() {
            reader = reader.version(version);
        }
        let reader = reader.await.map_err(map_opendal)?;
        let mut stream = reader.into_bytes_stream(..).await.map_err(map_opendal)?;
        let mut hasher = ContentHasher::new();
        let mut size = 0u64;
        while let Some(chunk) = stream
            .try_next()
            .await
            .map_err(|error| TargetError::transient(redact_sensitive(error.to_string())))?
        {
            hasher.update(&chunk);
            size = size.saturating_add(u64::try_from(chunk.len()).unwrap_or(u64::MAX));
        }
        Ok((hasher.finalize(), size))
    }

    fn stored_from_write(
        &self,
        key: &str,
        content_sha256: &ContentSha256,
        metadata: &Metadata,
        written_size: u64,
    ) -> StoredObject {
        StoredObject {
            identity: self.new_object_identity(key, metadata.version().map(str::to_owned)),
            content_sha256: content_sha256.clone(),
            size: metadata.content_length().max(written_size),
        }
    }

    fn new_object_identity(&self, key: &str, version_id: Option<String>) -> ObjectIdentity {
        ObjectIdentity {
            target_store_id: self.config.target_store_id.clone(),
            bucket: self.config.bucket.clone(),
            key: key.to_owned(),
            version_id: version_id.filter(|value| !value.is_empty()),
        }
    }

    fn ensure_object_identity(&self, identity: &ObjectIdentity) -> Result<(), TargetError> {
        if identity.target_store_id != self.config.target_store_id {
            return Err(TargetError::validation(
                "object identity target store does not match configured store",
            ));
        }
        if identity.bucket != self.config.bucket {
            return Err(TargetError::validation(
                "object identity bucket does not match configured target",
            ));
        }
        Ok(())
    }

    fn ensure_multipart_identity(&self, identity: &MultipartIdentity) -> Result<(), TargetError> {
        if identity.target_store_id != self.config.target_store_id {
            return Err(TargetError::validation(
                "multipart identity target store does not match configured store",
            ));
        }
        if identity.bucket != self.config.bucket {
            return Err(TargetError::validation(
                "multipart identity bucket does not match configured target",
            ));
        }
        Ok(())
    }
}

#[cfg(feature = "test-utils")]
impl TargetStore {
    /// Overwrite put without create-only guards (fixture version matrix only).
    pub async fn put_overwrite(
        &self,
        request: &CreateObjectRequest,
        body: impl Into<bytes::Bytes>,
    ) -> Result<StoredObject, TargetError> {
        validate_attempt_key(&request.key)?;
        let body = body.into();
        let computed = hash_bytes(&body);
        check_expected(request.expected_content_sha256.as_ref(), &computed)?;
        let size = body.len() as u64;
        let write = self
            .operator
            .write_with(&request.key, body)
            .user_metadata(sha256_metadata(&computed));
        let metadata = write.await.map_err(map_opendal)?;
        Ok(self.stored_from_write(&request.key, &computed, &metadata, size))
    }

    /// Enable bucket versioning when supported (fixture only).
    pub async fn enable_bucket_versioning(&self) -> Result<bool, TargetError> {
        self.multipart.enable_bucket_versioning().await
    }

    /// Delete every version under a prefix (fixture teardown only).
    pub async fn purge_versions_with_prefix(&self, prefix: &str) -> Result<usize, TargetError> {
        let versions = self.multipart.list_object_versions(prefix).await?;
        let count = versions.len();
        for version in versions {
            self.multipart
                .delete_object_version(&version.key, &version.version_id)
                .await?;
        }
        Ok(count)
    }
}

impl StorageAdapter for TargetStore {
    type Error = TargetError;
    type Request = CreateObjectRequest;
    type StoredObject = StoredObject;

    async fn create(&self, request: &Self::Request) -> Result<Self::StoredObject, Self::Error> {
        self.put_create_only(request, bytes::Bytes::new()).await
    }
}

fn sha256_metadata(digest: &ContentSha256) -> HashMap<String, String> {
    HashMap::from([(
        REMORA_SHA256_METADATA_KEY.to_owned(),
        digest.as_str().to_owned(),
    )])
}

fn observed_sha256(metadata: &Metadata) -> Result<Option<ContentSha256>, TargetError> {
    let Some(map) = metadata.user_metadata() else {
        return Ok(None);
    };
    map.get(REMORA_SHA256_METADATA_KEY)
        .or_else(|| {
            map.iter()
                .find(|(key, _)| key.eq_ignore_ascii_case(REMORA_SHA256_METADATA_KEY))
                .map(|(_, value)| value)
        })
        .map(ContentSha256::parse)
        .transpose()
}

/// Finish helper re-exported for callers that hold a streaming body.
pub async fn finish_create_only(body: OpenDalTargetBody) -> Result<StoredWrite, TargetError> {
    body.finish_object().await
}