somatize-core 0.5.1

//! Data Store: abstraction for moving data between workers.
//!
//! Separates WHERE data lives from HOW it's processed.
//! Workers use DataRef to reference data without materializing it.

/// Maximum payload size for inline WebSocket transport.
/// Payloads above this threshold are uploaded via HTTP bulk or DataStore.
pub const INLINE_THRESHOLD_BYTES: usize = 10 * 1024 * 1024; // 10 MB

// The S3 and Zarr backends live in `somatize-store`. They each own a
// tokio runtime, and a contract crate must not hand one to everything
// that depends on it.

use crate::cache::CacheKey;
use crate::error::{Result, SomaError};
use crate::value::Value;
use serde::{Deserialize, Serialize};

/// Metadata about a stored value, queryable without loading data.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StoreMeta {
    /// Total number of rows (`shape[0]` for tensors, 1 for scalar types).
    pub total_rows: usize,
    /// Remaining shape dimensions after the row axis (shape[1..] for tensors).
    pub shape_tail: Vec<usize>,
    /// Type tag: "tensor", "text", "json", "bytes", or "empty".
    pub dtype: String,
}

impl StoreMeta {
    /// Build metadata from an in-memory Value.
    pub fn from_value(value: &Value) -> Self {
        match value {
            Value::Tensor { shape, .. } => Self {
                total_rows: shape.first().copied().unwrap_or(0),
                shape_tail: shape.get(1..).unwrap_or_default().to_vec(),
                dtype: "tensor".into(),
            },
            Value::Text(_) => Self {
                total_rows: 1,
                shape_tail: vec![],
                dtype: "text".into(),
            },
            Value::Json(_) => Self {
                total_rows: 1,
                shape_tail: vec![],
                dtype: "json".into(),
            },
            Value::Bytes(b) | Value::Object(b) => Self {
                total_rows: b.len(),
                shape_tail: vec![],
                dtype: "bytes".into(),
            },
            Value::Empty => Self {
                total_rows: 0,
                shape_tail: vec![],
                dtype: "empty".into(),
            },
        }
    }
}

/// Slice rows `[start..start+len)` from a tensor value.
pub fn slice_tensor_rows(value: &Value, start: usize, len: usize) -> Result<Value> {
    match value {
        Value::Tensor { values, shape } => {
            if shape.is_empty() {
                return Err(SomaError::DataStore("cannot slice scalar tensor".into()));
            }
            let cols: usize = shape[1..].iter().product::<usize>().max(1);
            let row_start = start * cols;
            let row_end = (start + len) * cols;
            if row_end > values.len() {
                return Err(SomaError::DataStore(format!(
                    "row range {start}..{} out of bounds (total rows: {})",
                    start + len,
                    shape[0]
                )));
            }
            let mut new_shape = shape.clone();
            new_shape[0] = len;
            Ok(Value::tensor(
                values[row_start..row_end].to_vec(),
                new_shape,
            ))
        }
        _ => Err(SomaError::DataStore(
            "get_rows only works on Tensor values".into(),
        )),
    }
}

/// A reference to data that may live in different places.
/// Workers exchange DataRefs instead of raw data.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type")]
#[non_exhaustive]
pub enum DataRef {
    /// Data in local filesystem
    Local {
        /// Absolute path to the file holding the serialized value.
        path: String,
    },
    /// Data in S3-compatible object storage
    S3 {
        /// Bucket the object lives in.
        bucket: String,
        /// Object key within the bucket.
        key: String,
        /// AWS region, `None` for endpoints that don't need one.
        region: Option<String>,
    },
    /// Data in Soma cache (content-addressable)
    Cached {
        /// Key the value is cached under.
        cache_key: CacheKey,
    },
    /// Data available as a stream endpoint
    Stream {
        /// URL the stream can be read from.
        endpoint: String,
        /// Wire format of the streamed records.
        format: StreamFormat,
    },
    /// Data materialized inline (small values only)
    Inline {
        /// The value itself, carried in the reference.
        value: Value,
    },
    /// Data stored as a Zarr v3 array in object storage (chunked tensors).
    Zarr {
        /// Bucket the array lives in.
        bucket: String,
        /// Root path of the Zarr array (contains zarr.json + chunk objects).
        array_path: String,
        /// AWS region, `None` for endpoints that don't need one.
        region: Option<String>,
    },
}

/// Stream data format.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
#[serde(rename_all = "snake_case")]
#[non_exhaustive]
pub enum StreamFormat {
    /// Newline-delimited JSON, one record per line (the default).
    #[default]
    JsonLines,
    /// Comma-separated values.
    Csv,
    /// Apache Arrow IPC stream.
    Arrow,
    /// Length-prefixed protobuf messages.
    Protobuf,
}

/// Storage configuration for an investigation/pipeline.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type")]
#[non_exhaustive]
pub enum StorageConfig {
    /// Local filesystem (NFS, mounted volume)
    #[serde(rename = "local")]
    Local {
        /// Directory values are written under.
        base_path: String,
    },
    /// S3-compatible object storage
    #[serde(rename = "s3")]
    S3 {
        /// Bucket to store objects in.
        bucket: String,
        /// Key prefix all objects are written under.
        prefix: String,
        /// AWS region, `None` for endpoints that don't need one.
        region: Option<String>,
        /// Custom endpoint URL for non-AWS backends (MinIO, Ceph).
        endpoint: Option<String>,
    },
    /// Zarr v3 chunked storage on S3-compatible backend.
    #[serde(rename = "zarr")]
    Zarr {
        /// Bucket to store arrays in.
        bucket: String,
        /// Key prefix all arrays are written under.
        prefix: String,
        /// AWS region, `None` for endpoints that don't need one.
        region: Option<String>,
        /// Custom endpoint URL for non-AWS backends (MinIO, Ceph).
        endpoint: Option<String>,
        /// Rows per chunk (first dimension).
        chunk_rows: usize,
    },
}

impl Default for StorageConfig {
    fn default() -> Self {
        Self::Local {
            base_path: "/tmp/soma-data".to_string(),
        }
    }
}

/// The DataStore trait: put/get/stream data across workers.
///
/// Unlike CacheStore (which stores Values by CacheKey),
/// DataStore moves data between locations and supports streaming.
pub trait DataStore: Send + Sync {
    /// Store data and return a reference to it.
    fn put(&self, key: &CacheKey, data: &Value) -> Result<DataRef>;

    /// Retrieve data from a reference.
    fn get(&self, data_ref: &DataRef) -> Result<Value>;

    /// Check if data exists at a reference.
    fn exists(&self, data_ref: &DataRef) -> Result<bool>;

    /// Delete data at a reference.
    fn remove(&self, data_ref: &DataRef) -> Result<()>;

    /// Get the storage config.
    fn config(&self) -> &StorageConfig;

    /// Read a range of rows `[start..start+len)` from a tensor.
    /// Returns a `Value::Tensor` with `shape[0] == len`.
    /// Default impl downloads the full value and slices in memory.
    fn get_rows(&self, data_ref: &DataRef, start: usize, len: usize) -> Result<Value> {
        let value = self.get(data_ref)?;
        slice_tensor_rows(&value, start, len)
    }

    /// Get metadata about a stored value without reading the data.
    /// Default impl downloads the full value to extract metadata.
    fn meta(&self, data_ref: &DataRef) -> Result<StoreMeta> {
        let value = self.get(data_ref)?;
        Ok(StoreMeta::from_value(&value))
    }
}

/// Local filesystem data store.
pub struct LocalDataStore {
    config: StorageConfig,
    base_path: std::path::PathBuf,
}

impl LocalDataStore {
    /// Create a store rooted at `base_path`, creating the directory if
    /// needed. Creation failure is deliberately ignored here — the
    /// first `put` will surface it as a [`SomaError::DataStore`].
    pub fn new(base_path: impl Into<std::path::PathBuf>) -> Self {
        let base = base_path.into();
        std::fs::create_dir_all(&base).ok();
        Self {
            config: StorageConfig::Local {
                base_path: base.to_string_lossy().to_string(),
            },
            base_path: base,
        }
    }
}

impl DataStore for LocalDataStore {
    fn put(&self, key: &CacheKey, data: &Value) -> Result<DataRef> {
        let path = self.base_path.join(key.to_hex());
        let bytes = serde_json::to_vec(data)
            .map_err(|e| crate::error::SomaError::DataStore(e.to_string()))?;
        std::fs::write(&path, &bytes)
            .map_err(|e| crate::error::SomaError::DataStore(e.to_string()))?;
        Ok(DataRef::Local {
            path: path.to_string_lossy().to_string(),
        })
    }

    fn get(&self, data_ref: &DataRef) -> Result<Value> {
        match data_ref {
            DataRef::Local { path } => {
                let bytes = std::fs::read(path)
                    .map_err(|e| crate::error::SomaError::DataStore(e.to_string()))?;
                serde_json::from_slice(&bytes)
                    .map_err(|e| crate::error::SomaError::DataStore(e.to_string()))
            }
            DataRef::Cached { cache_key } => {
                let path = self.base_path.join(cache_key.to_hex());
                let bytes = std::fs::read(&path)
                    .map_err(|e| crate::error::SomaError::DataStore(e.to_string()))?;
                serde_json::from_slice(&bytes)
                    .map_err(|e| crate::error::SomaError::DataStore(e.to_string()))
            }
            DataRef::Inline { value } => Ok(value.clone()),
            _ => Err(crate::error::SomaError::DataStore(
                "Cannot get non-local DataRef from LocalDataStore".into(),
            )),
        }
    }

    fn exists(&self, data_ref: &DataRef) -> Result<bool> {
        match data_ref {
            DataRef::Local { path } => Ok(std::path::Path::new(path).exists()),
            DataRef::Cached { cache_key } => Ok(self.base_path.join(cache_key.to_hex()).exists()),
            DataRef::Inline { .. } => Ok(true),
            _ => Ok(false),
        }
    }

    fn remove(&self, data_ref: &DataRef) -> Result<()> {
        if let DataRef::Local { path } = data_ref {
            std::fs::remove_file(path).ok();
        }
        Ok(())
    }

    fn config(&self) -> &StorageConfig {
        &self.config
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn local_data_store_roundtrip() {
        let dir = std::env::temp_dir().join("soma-ds-test");
        let store = LocalDataStore::new(&dir);

        let key = CacheKey::hash_data(b"test_data");
        let value = Value::tensor(vec![1.0, 2.0, 3.0], vec![3]);

        let data_ref = store.put(&key, &value).unwrap();
        assert!(store.exists(&data_ref).unwrap());

        let retrieved = store.get(&data_ref).unwrap();
        let (data, _) = retrieved.as_tensor().unwrap();
        assert_eq!(data, &[1.0, 2.0, 3.0]);

        store.remove(&data_ref).unwrap();
        assert!(!store.exists(&data_ref).unwrap());

        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn inline_data_ref() {
        let dir = std::env::temp_dir().join("soma-ds-test-inline");
        let store = LocalDataStore::new(&dir);

        let data_ref = DataRef::Inline {
            value: Value::tensor(vec![42.0], vec![1]),
        };

        assert!(store.exists(&data_ref).unwrap());
        let v = store.get(&data_ref).unwrap();
        let (data, _) = v.as_tensor().unwrap();
        assert_eq!(data, &[42.0]);

        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn storage_config_serde() {
        let s3 = StorageConfig::S3 {
            bucket: "my-lab".into(),
            prefix: "experiments/".into(),
            region: Some("eu-west-1".into()),
            endpoint: None,
        };
        let json = serde_json::to_string(&s3).unwrap();
        assert!(json.contains("my-lab"));

        let local = StorageConfig::Local {
            base_path: "/data".into(),
        };
        let json = serde_json::to_string(&local).unwrap();
        assert!(json.contains("/data"));
    }

    #[test]
    fn data_ref_serde() {
        let refs = vec![
            DataRef::Local {
                path: "/tmp/x".into(),
            },
            DataRef::S3 {
                bucket: "b".into(),
                key: "k".into(),
                region: None,
            },
            DataRef::Cached {
                cache_key: CacheKey::hash_data(b"x"),
            },
            DataRef::Inline {
                value: Value::Empty,
            },
            DataRef::Zarr {
                bucket: "b".into(),
                array_path: "data/abc".into(),
                region: None,
            },
        ];
        for r in &refs {
            let json = serde_json::to_string(r).unwrap();
            let _: DataRef = serde_json::from_str(&json).unwrap();
        }
    }

    #[test]
    fn slice_tensor_rows_basic() {
        // 4 rows × 3 cols
        let v = Value::tensor(
            vec![
                1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0,
            ],
            vec![4, 3],
        );
        // Rows 1..3 → [[4,5,6], [7,8,9]]
        let sliced = slice_tensor_rows(&v, 1, 2).unwrap();
        let (data, shape) = sliced.as_tensor().unwrap();
        assert_eq!(shape, &[2, 3]);
        assert_eq!(data, &[4.0, 5.0, 6.0, 7.0, 8.0, 9.0]);
    }

    #[test]
    fn slice_tensor_rows_single() {
        let v = Value::tensor(vec![10.0, 20.0, 30.0], vec![3]);
        let sliced = slice_tensor_rows(&v, 1, 1).unwrap();
        let (data, shape) = sliced.as_tensor().unwrap();
        assert_eq!(shape, &[1]);
        assert_eq!(data, &[20.0]);
    }

    #[test]
    fn slice_tensor_rows_out_of_bounds() {
        let v = Value::tensor(vec![1.0, 2.0, 3.0], vec![3]);
        assert!(slice_tensor_rows(&v, 2, 5).is_err());
    }

    #[test]
    fn store_meta_from_tensor() {
        let v = Value::tensor(vec![0.0; 12], vec![4, 3]);
        let meta = StoreMeta::from_value(&v);
        assert_eq!(meta.total_rows, 4);
        assert_eq!(meta.shape_tail, vec![3]);
        assert_eq!(meta.dtype, "tensor");
    }

    #[test]
    fn store_meta_from_json() {
        let v = Value::json(serde_json::json!({"a": 1}));
        let meta = StoreMeta::from_value(&v);
        assert_eq!(meta.dtype, "json");
        assert_eq!(meta.total_rows, 1);
    }

    #[test]
    fn default_get_rows_on_local_store() {
        let dir = std::env::temp_dir().join("soma-ds-test-getrows");
        let store = LocalDataStore::new(&dir);

        let key = CacheKey::hash_data(b"rows_test");
        let value = Value::tensor(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![3, 2]);
        let data_ref = store.put(&key, &value).unwrap();

        // Read rows 1..2 via default impl (full get + slice)
        let sliced = store.get_rows(&data_ref, 1, 2).unwrap();
        let (data, shape) = sliced.as_tensor().unwrap();
        assert_eq!(shape, &[2, 2]);
        assert_eq!(data, &[3.0, 4.0, 5.0, 6.0]);

        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn default_meta_on_local_store() {
        let dir = std::env::temp_dir().join("soma-ds-test-meta");
        let store = LocalDataStore::new(&dir);

        let key = CacheKey::hash_data(b"meta_test");
        let value = Value::tensor(vec![0.0; 20], vec![5, 4]);
        let data_ref = store.put(&key, &value).unwrap();

        let meta = store.meta(&data_ref).unwrap();
        assert_eq!(meta.total_rows, 5);
        assert_eq!(meta.shape_tail, vec![4]);
        assert_eq!(meta.dtype, "tensor");

        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn zarr_storage_config_serde() {
        let zarr = StorageConfig::Zarr {
            bucket: "soma-research".into(),
            prefix: "data/".into(),
            region: None,
            endpoint: Some("s3.eu-central-003.backblazeb2.com".into()),
            chunk_rows: 1024,
        };
        let json = serde_json::to_string(&zarr).unwrap();
        assert!(json.contains("soma-research"));
        assert!(json.contains("1024"));
        let _: StorageConfig = serde_json::from_str(&json).unwrap();
    }
}