velocia 0.3.1

//! Generic DynamoDB repository.
//!
//! Mirrors Python's `DynamoDBRepository` – a reusable abstraction over
//! common DynamoDB CRUD operations.  Enabled via the `dynamodb` feature.

#[cfg(feature = "dynamodb")]
use std::collections::HashMap;
#[cfg(feature = "dynamodb")]
use aws_sdk_dynamodb::{types::AttributeValue, Client as DynamoDbClient};
#[cfg(feature = "dynamodb")]
use tracing::debug;
#[cfg(feature = "dynamodb")]
use crate::error::{AgentKitError, Result};


// ── Repository ────────────────────────────────────────────────────────────────

#[cfg(feature = "dynamodb")]
pub struct DynamoDbRepository {
    client: DynamoDbClient,
    table_name: String,
}

#[cfg(feature = "dynamodb")]
impl DynamoDbRepository {
    /// Create a new repository.  AWS config is loaded from the environment.
    pub async fn new(
        table_name: impl Into<String>,
        endpoint_url: Option<String>,
        region: Option<String>,
        assumed_role: Option<String>,
    ) -> Result<Self> {
        let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest());

        if let Some(r) = region {
            loader = loader.region(aws_config::Region::new(r));
        }

        let sdk_config = loader.load().await;
        let mut builder = aws_sdk_dynamodb::config::Builder::from(&sdk_config);

        if let Some(ep) = endpoint_url {
            builder = builder.endpoint_url(ep);
        }

        // Optional role assumption for cross-account access.
        // TODO: Implement STS assume-role when `assumed_role` is set.

        Ok(Self {
            client: DynamoDbClient::from_conf(builder.build()),
            table_name: table_name.into(),
        })
    }

    // ── Write ─────────────────────────────────────────────────────────────────

    pub async fn put_item(&self, item: HashMap<String, serde_json::Value>) -> Result<()> {
        let dynamo_item: HashMap<String, AttributeValue> =
            item.into_iter().map(|(k, v)| (k, json_to_av(&v))).collect();

        self.client
            .put_item()
            .table_name(&self.table_name)
            .set_item(Some(dynamo_item))
            .send()
            .await
            .map_err(|e| AgentKitError::DynamoDB(e.to_string()))?;

        debug!("Item stored in '{}'", self.table_name);
        Ok(())
    }

    // ── Read ──────────────────────────────────────────────────────────────────

    pub async fn get_item(
        &self,
        key: HashMap<String, serde_json::Value>,
    ) -> Result<Option<HashMap<String, serde_json::Value>>> {
        let dynamo_key: HashMap<String, AttributeValue> =
            key.into_iter().map(|(k, v)| (k, json_to_av(&v))).collect();

        let resp = self
            .client
            .get_item()
            .table_name(&self.table_name)
            .set_key(Some(dynamo_key))
            .send()
            .await
            .map_err(|e| AgentKitError::DynamoDB(e.to_string()))?;

        Ok(resp.item.map(|item| item.into_iter().map(|(k, v)| (k, av_to_json(&v))).collect()))
    }

    pub async fn query(
        &self,
        key_condition: &str,
        expression_values: HashMap<String, serde_json::Value>,
    ) -> Result<Vec<HashMap<String, serde_json::Value>>> {
        let av_map: HashMap<String, AttributeValue> = expression_values
            .into_iter()
            .map(|(k, v)| (k, json_to_av(&v)))
            .collect();

        let resp = self
            .client
            .query()
            .table_name(&self.table_name)
            .key_condition_expression(key_condition)
            .set_expression_attribute_values(Some(av_map))
            .send()
            .await
            .map_err(|e| AgentKitError::DynamoDB(e.to_string()))?;

        Ok(resp
            .items
            .unwrap_or_default()
            .into_iter()
            .map(|row| row.into_iter().map(|(k, v)| (k, av_to_json(&v))).collect())
            .collect())
    }

    // ── Delete ────────────────────────────────────────────────────────────────

    pub async fn delete_item(&self, key: HashMap<String, serde_json::Value>) -> Result<()> {
        let dynamo_key: HashMap<String, AttributeValue> =
            key.into_iter().map(|(k, v)| (k, json_to_av(&v))).collect();

        self.client
            .delete_item()
            .table_name(&self.table_name)
            .set_key(Some(dynamo_key))
            .send()
            .await
            .map_err(|e| AgentKitError::DynamoDB(e.to_string()))?;
        Ok(())
    }
}

// ── Attribute value ↔ JSON conversion ─────────────────────────────────────────

#[cfg(feature = "dynamodb")]
pub fn json_to_av(v: &serde_json::Value) -> AttributeValue {
    match v {
        serde_json::Value::String(s) => AttributeValue::S(s.clone()),
        serde_json::Value::Number(n) => AttributeValue::N(n.to_string()),
        serde_json::Value::Bool(b) => AttributeValue::Bool(*b),
        serde_json::Value::Null => AttributeValue::Null(true),
        serde_json::Value::Array(arr) => {
            AttributeValue::L(arr.iter().map(json_to_av).collect())
        }
        serde_json::Value::Object(obj) => AttributeValue::M(
            obj.iter().map(|(k, v)| (k.clone(), json_to_av(v))).collect(),
        ),
    }
}

#[cfg(feature = "dynamodb")]
pub fn av_to_json(av: &AttributeValue) -> serde_json::Value {
    match av {
        AttributeValue::S(s) => serde_json::Value::String(s.clone()),
        AttributeValue::N(n) => {
            if let Ok(i) = n.parse::<i64>() {
                serde_json::json!(i)
            } else if let Ok(f) = n.parse::<f64>() {
                serde_json::json!(f)
            } else {
                serde_json::Value::String(n.clone())
            }
        }
        AttributeValue::Bool(b) => serde_json::Value::Bool(*b),
        AttributeValue::Null(_) => serde_json::Value::Null,
        AttributeValue::L(list) => {
            serde_json::Value::Array(list.iter().map(av_to_json).collect())
        }
        AttributeValue::M(map) => serde_json::Value::Object(
            map.iter().map(|(k, v)| (k.clone(), av_to_json(v))).collect(),
        ),
        _ => serde_json::Value::Null,
    }
}