briefcase_core/storage/

lakefs.rs

use super::{FlushResult, SnapshotQuery, StorageBackend, StorageError};
use crate::models::{DecisionSnapshot, Snapshot};
#[cfg(feature = "networking")]
use base64::{engine::general_purpose, Engine as _};
#[cfg(feature = "networking")]
use reqwest::{
    header::{HeaderMap, HeaderValue, AUTHORIZATION, CONTENT_TYPE},
    Client,
};
use serde::{Deserialize, Serialize};
use std::sync::{Arc, Mutex};

#[cfg(all(feature = "async", feature = "networking"))]
#[derive(Clone)]
pub struct LakeFSBackend {
    client: Client,
    endpoint: String,
    repository: String, // "engagement" in Briefcase terminology
    branch: String,     // "workstream" in Briefcase terminology
    #[allow(dead_code)]
    access_key: String,
    #[allow(dead_code)]
    secret_key: String,
    pending_writes: Arc<Mutex<Vec<Snapshot>>>,
}

#[cfg(all(feature = "async", feature = "networking"))]
impl LakeFSBackend {
    pub fn new(config: LakeFSConfig) -> Result<Self, StorageError> {
        let mut headers = HeaderMap::new();

        // Create basic auth header
        let credentials = format!("{}:{}", config.access_key, config.secret_key);
        let encoded = general_purpose::STANDARD.encode(credentials.as_bytes());
        let auth_header = format!("Basic {}", encoded);

        headers.insert(
            AUTHORIZATION,
            HeaderValue::from_str(&auth_header).map_err(|e| {
                StorageError::ConnectionError(format!("Invalid auth header: {}", e))
            })?,
        );

        headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));

        let client = Client::builder()
            .default_headers(headers)
            .timeout(std::time::Duration::from_secs(30))
            .build()
            .map_err(|e| {
                StorageError::ConnectionError(format!("Failed to create HTTP client: {}", e))
            })?;

        Ok(Self {
            client,
            endpoint: config.endpoint.trim_end_matches('/').to_string(),
            repository: config.repository,
            branch: config.branch,
            access_key: config.access_key,
            secret_key: config.secret_key,
            pending_writes: Arc::new(Mutex::new(Vec::new())),
        })
    }

    /// Create a commit (checkpoint) with pending writes
    async fn create_commit(&self, message: &str) -> Result<String, StorageError> {
        let url = format!(
            "{}/repositories/{}/branches/{}/commits",
            self.endpoint, self.repository, self.branch
        );

        #[derive(Serialize)]
        struct CommitRequest {
            message: String,
            metadata: std::collections::HashMap<String, String>,
        }

        let mut metadata = std::collections::HashMap::new();
        metadata.insert("source".to_string(), "briefcase-ai".to_string());

        let request = CommitRequest {
            message: message.to_string(),
            metadata,
        };

        let response = self
            .client
            .post(&url)
            .json(&request)
            .send()
            .await
            .map_err(|e| {
                StorageError::ConnectionError(format!("Failed to create commit: {}", e))
            })?;

        let status = response.status();
        if !status.is_success() {
            let error_text = response.text().await.unwrap_or_default();
            return Err(StorageError::ConnectionError(format!(
                "Commit failed with status {}: {}",
                status, error_text
            )));
        }

        #[derive(Deserialize)]
        struct CommitResponse {
            id: String,
        }

        let commit_response: CommitResponse = response.json().await.map_err(|e| {
            StorageError::SerializationError(format!("Failed to parse commit response: {}", e))
        })?;

        Ok(commit_response.id)
    }

    /// Upload object to LakeFS
    async fn upload_object(&self, path: &str, data: &[u8]) -> Result<(), StorageError> {
        let url = format!(
            "{}/repositories/{}/branches/{}/objects",
            self.endpoint, self.repository, self.branch
        );

        let response = self
            .client
            .put(&url)
            .query(&[("path", path)])
            .header("Content-Type", "application/octet-stream")
            .body(data.to_vec())
            .send()
            .await
            .map_err(|e| {
                StorageError::ConnectionError(format!("Failed to upload object: {}", e))
            })?;

        let status = response.status();
        if !status.is_success() {
            let error_text = response.text().await.unwrap_or_default();
            return Err(StorageError::ConnectionError(format!(
                "Upload failed with status {}: {}",
                status, error_text
            )));
        }

        Ok(())
    }

    /// Download object from LakeFS
    async fn download_object(&self, path: &str) -> Result<Vec<u8>, StorageError> {
        let url = format!(
            "{}/repositories/{}/refs/{}/objects",
            self.endpoint, self.repository, self.branch
        );

        let response = self
            .client
            .get(&url)
            .query(&[("path", path)])
            .send()
            .await
            .map_err(|e| {
                StorageError::ConnectionError(format!("Failed to download object: {}", e))
            })?;

        if response.status() == reqwest::StatusCode::NOT_FOUND {
            return Err(StorageError::NotFound(format!(
                "Object not found: {}",
                path
            )));
        }

        let status = response.status();
        if !status.is_success() {
            let error_text = response.text().await.unwrap_or_default();
            return Err(StorageError::ConnectionError(format!(
                "Download failed with status {}: {}",
                status, error_text
            )));
        }

        let data = response.bytes().await.map_err(|e| {
            StorageError::ConnectionError(format!("Failed to read response: {}", e))
        })?;

        Ok(data.to_vec())
    }

    /// List objects with prefix
    async fn list_objects(&self, prefix: &str) -> Result<Vec<String>, StorageError> {
        let url = format!(
            "{}/repositories/{}/refs/{}/objects/ls",
            self.endpoint, self.repository, self.branch
        );

        let response = self
            .client
            .get(&url)
            .query(&[("prefix", prefix)])
            .send()
            .await
            .map_err(|e| StorageError::ConnectionError(format!("Failed to list objects: {}", e)))?;

        let status = response.status();
        if !status.is_success() {
            let error_text = response.text().await.unwrap_or_default();
            return Err(StorageError::ConnectionError(format!(
                "List failed with status {}: {}",
                status, error_text
            )));
        }

        #[derive(Deserialize)]
        struct ListResponse {
            results: Vec<ObjectInfo>,
        }

        #[derive(Deserialize)]
        struct ObjectInfo {
            path: String,
            #[serde(rename = "type")]
            object_type: String,
        }

        let list_response: ListResponse = response.json().await.map_err(|e| {
            StorageError::SerializationError(format!("Failed to parse list response: {}", e))
        })?;

        let paths = list_response
            .results
            .into_iter()
            .filter(|obj| obj.object_type == "object")
            .map(|obj| obj.path)
            .collect();

        Ok(paths)
    }

    /// Generate object path for snapshot
    fn snapshot_path(&self, snapshot_id: &str) -> String {
        format!("snapshots/{}.json", snapshot_id)
    }

    /// Generate object path for decision
    fn decision_path(&self, decision_id: &str) -> String {
        format!("decisions/{}.json", decision_id)
    }
}

#[cfg(all(feature = "async", feature = "networking"))]
#[async_trait::async_trait]
impl StorageBackend for LakeFSBackend {
    async fn save(&self, snapshot: &Snapshot) -> Result<String, StorageError> {
        let snapshot_id = snapshot.metadata.snapshot_id.to_string();

        // Add to pending writes instead of immediate upload
        {
            let mut pending = self.pending_writes.lock().unwrap();
            pending.push(snapshot.clone());
        }

        Ok(snapshot_id)
    }

    async fn save_decision(&self, decision: &DecisionSnapshot) -> Result<String, StorageError> {
        let decision_id = decision.metadata.snapshot_id.to_string();
        let path = self.decision_path(&decision_id);

        let json_data = serde_json::to_vec(decision)
            .map_err(|e| StorageError::SerializationError(e.to_string()))?;

        self.upload_object(&path, &json_data).await?;

        Ok(decision_id)
    }

    async fn load(&self, snapshot_id: &str) -> Result<Snapshot, StorageError> {
        let path = self.snapshot_path(snapshot_id);
        let data = self.download_object(&path).await?;

        let snapshot: Snapshot = serde_json::from_slice(&data)
            .map_err(|e| StorageError::SerializationError(e.to_string()))?;

        Ok(snapshot)
    }

    async fn load_decision(&self, decision_id: &str) -> Result<DecisionSnapshot, StorageError> {
        let path = self.decision_path(decision_id);
        let data = self.download_object(&path).await?;

        let decision: DecisionSnapshot = serde_json::from_slice(&data)
            .map_err(|e| StorageError::SerializationError(e.to_string()))?;

        Ok(decision)
    }

    async fn query(&self, query: SnapshotQuery) -> Result<Vec<Snapshot>, StorageError> {
        // List all snapshots
        let paths = self.list_objects("snapshots/").await?;

        let mut snapshots = Vec::new();
        let mut count = 0;
        let offset = query.offset.unwrap_or(0);
        let limit = query.limit.unwrap_or(usize::MAX);

        for path in paths {
            if let Some(filename) = path.split('/').next_back() {
                if let Some(snapshot_id) = filename.strip_suffix(".json") {
                    // Load snapshot to check filters
                    match self.load(snapshot_id).await {
                        Ok(snapshot) => {
                            // Apply filters
                            if self.matches_query(&snapshot, &query) {
                                if count >= offset {
                                    snapshots.push(snapshot);
                                    if snapshots.len() >= limit {
                                        break;
                                    }
                                }
                                count += 1;
                            }
                        }
                        Err(_) => continue, // Skip invalid snapshots
                    }
                }
            }
        }

        // Sort by timestamp (newest first)
        snapshots.sort_by(|a, b| b.metadata.timestamp.cmp(&a.metadata.timestamp));

        Ok(snapshots)
    }

    async fn delete(&self, _snapshot_id: &str) -> Result<bool, StorageError> {
        // LakeFS doesn't have a direct delete API for objects in branches
        // This would typically involve creating a new commit without the object
        // For simplicity, we'll return an error indicating this operation isn't supported
        Err(StorageError::PermissionDenied(
            "LakeFS doesn't support direct object deletion. Use branch operations instead."
                .to_string(),
        ))
    }

    async fn flush(&self) -> Result<FlushResult, StorageError> {
        let pending_snapshots = {
            let mut pending = self.pending_writes.lock().unwrap();
            let snapshots = pending.clone();
            pending.clear();
            snapshots
        };

        let mut bytes_written = 0;

        // Upload all pending snapshots
        for snapshot in &pending_snapshots {
            let snapshot_id = snapshot.metadata.snapshot_id.to_string();
            let path = self.snapshot_path(&snapshot_id);

            let json_data = serde_json::to_vec(snapshot)
                .map_err(|e| StorageError::SerializationError(e.to_string()))?;

            bytes_written += json_data.len();

            self.upload_object(&path, &json_data).await?;

            // Also upload individual decisions
            for decision in &snapshot.decisions {
                let decision_id = decision.metadata.snapshot_id.to_string();
                let decision_path = self.decision_path(&decision_id);

                let decision_data = serde_json::to_vec(decision)
                    .map_err(|e| StorageError::SerializationError(e.to_string()))?;

                bytes_written += decision_data.len();
                self.upload_object(&decision_path, &decision_data).await?;
            }
        }

        // Create commit
        let commit_message = format!("Briefcase AI flush: {} snapshots", pending_snapshots.len());
        let commit_id = self.create_commit(&commit_message).await?;

        Ok(FlushResult {
            snapshots_written: pending_snapshots.len(),
            bytes_written,
            checkpoint_id: Some(commit_id),
        })
    }

    async fn health_check(&self) -> Result<bool, StorageError> {
        let url = format!("{}/repositories/{}", self.endpoint, self.repository);

        let response =
            self.client.get(&url).send().await.map_err(|e| {
                StorageError::ConnectionError(format!("Health check failed: {}", e))
            })?;

        Ok(response.status().is_success())
    }
}

#[cfg(all(feature = "async", feature = "networking"))]
impl LakeFSBackend {
    /// Check if snapshot matches query filters
    fn matches_query(&self, snapshot: &Snapshot, query: &SnapshotQuery) -> bool {
        // Check time range
        if let Some(start_time) = query.start_time {
            if snapshot.metadata.timestamp < start_time {
                return false;
            }
        }

        if let Some(end_time) = query.end_time {
            if snapshot.metadata.timestamp > end_time {
                return false;
            }
        }

        // Check function name, module name, model name, tags in decisions
        if query.function_name.is_some()
            || query.module_name.is_some()
            || query.model_name.is_some()
            || query.tags.is_some()
        {
            let mut found_match = false;

            for decision in &snapshot.decisions {
                let mut decision_matches = true;

                if let Some(function_name) = &query.function_name {
                    if decision.function_name != *function_name {
                        decision_matches = false;
                    }
                }

                if let Some(module_name) = &query.module_name {
                    if decision.module_name.as_ref() != Some(module_name) {
                        decision_matches = false;
                    }
                }

                if let Some(model_name) = &query.model_name {
                    if let Some(model_params) = &decision.model_parameters {
                        if model_params.model_name != *model_name {
                            decision_matches = false;
                        }
                    } else {
                        decision_matches = false;
                    }
                }

                if let Some(query_tags) = &query.tags {
                    for (key, value) in query_tags {
                        if decision.tags.get(key) != Some(value) {
                            decision_matches = false;
                            break;
                        }
                    }
                }

                if decision_matches {
                    found_match = true;
                    break;
                }
            }

            if !found_match {
                return false;
            }
        }

        true
    }
}

#[derive(Debug, Clone)]
pub struct LakeFSConfig {
    pub endpoint: String,
    pub repository: String,
    pub branch: String,
    pub access_key: String,
    pub secret_key: String,
}

impl LakeFSConfig {
    pub fn new(
        endpoint: impl Into<String>,
        repository: impl Into<String>,
        branch: impl Into<String>,
        access_key: impl Into<String>,
        secret_key: impl Into<String>,
    ) -> Self {
        Self {
            endpoint: endpoint.into(),
            repository: repository.into(),
            branch: branch.into(),
            access_key: access_key.into(),
            secret_key: secret_key.into(),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::models::*;
    use serde_json::json;

    fn create_test_config() -> LakeFSConfig {
        LakeFSConfig::new(
            "http://localhost:8000",
            "briefcase-test",
            "main",
            "test_access_key",
            "test_secret_key",
        )
    }

    async fn create_test_snapshot() -> Snapshot {
        let input = Input::new("test_input", json!("value"), "string");
        let output = Output::new("test_output", json!("result"), "string");
        let model_params = ModelParameters::new("gpt-4");

        let decision = DecisionSnapshot::new("test_function")
            .with_module("test_module")
            .add_input(input)
            .add_output(output)
            .with_model_parameters(model_params)
            .add_tag("env", "test");

        let mut snapshot = Snapshot::new(SnapshotType::Session);
        snapshot.add_decision(decision);
        snapshot
    }

    #[tokio::test]
    async fn test_lakefs_config_creation() {
        let config = create_test_config();
        assert_eq!(config.endpoint, "http://localhost:8000");
        assert_eq!(config.repository, "briefcase-test");
        assert_eq!(config.branch, "main");
    }

    #[tokio::test]
    async fn test_object_paths() {
        let config = create_test_config();
        let backend = LakeFSBackend::new(config).unwrap();

        let snapshot_id = "test-snapshot-123";
        let decision_id = "test-decision-456";

        assert_eq!(
            backend.snapshot_path(snapshot_id),
            "snapshots/test-snapshot-123.json"
        );
        assert_eq!(
            backend.decision_path(decision_id),
            "decisions/test-decision-456.json"
        );
    }

    #[tokio::test]
    async fn test_query_matching() {
        let config = create_test_config();
        let backend = LakeFSBackend::new(config).unwrap();
        let snapshot = create_test_snapshot().await;

        // Test function name matching
        let query = SnapshotQuery::new().with_function_name("test_function");
        assert!(backend.matches_query(&snapshot, &query));

        let query = SnapshotQuery::new().with_function_name("other_function");
        assert!(!backend.matches_query(&snapshot, &query));

        // Test tag matching
        let query = SnapshotQuery::new().with_tag("env", "test");
        assert!(backend.matches_query(&snapshot, &query));

        let query = SnapshotQuery::new().with_tag("env", "prod");
        assert!(!backend.matches_query(&snapshot, &query));

        // Test model name matching
        let query = SnapshotQuery::new().with_model_name("gpt-4");
        assert!(backend.matches_query(&snapshot, &query));

        let query = SnapshotQuery::new().with_model_name("claude-3");
        assert!(!backend.matches_query(&snapshot, &query));
    }

    #[tokio::test]
    async fn test_pending_writes() {
        let config = create_test_config();
        let backend = LakeFSBackend::new(config).unwrap();
        let snapshot = create_test_snapshot().await;

        // Save should add to pending writes
        let snapshot_id = backend.save(&snapshot).await.unwrap();
        assert_eq!(snapshot_id, snapshot.metadata.snapshot_id.to_string());

        // Check pending writes
        {
            let pending = backend.pending_writes.lock().unwrap();
            assert_eq!(pending.len(), 1);
        }
    }

    // Note: Integration tests against real LakeFS would require a running instance
    // These would be better suited for a separate integration test suite
}