use async_trait::async_trait; use serde_json::Value; use sqlx::{sqlite::SqliteRow, Row}; use aether_data_contracts::repository::routing_profiles::*; use aether_data_contracts::DataLayerError; use crate::error::SqlResultExt; use crate::pool::SqlitePool; const ROUTING_GROUP_SELECT: &str = r#" SELECT id, name, description, enabled, is_system_default, config_json, version, created_at, updated_at, published_at FROM routing_groups "#; const ROUTING_GROUP_BINDING_SELECT: &str = r#" SELECT id, group_id, subject_type, subject_id, is_default, allow_explicit_select, created_at, updated_at FROM routing_group_bindings "#; const ROUTING_GROUP_VERSION_SELECT: &str = r#" SELECT id, group_id, version, config_json, created_at, created_by FROM routing_group_versions "#; #[derive(Debug, Clone)] pub struct SqliteRoutingGroupRepository { pool: SqlitePool, } impl SqliteRoutingGroupRepository { pub fn new(pool: SqlitePool) -> Self { Self { pool } } async fn reload_group(&self, id: &str) -> Result, DataLayerError> { self.find_routing_group(RoutingGroupLookupKey::Id(id)).await } async fn find_binding_by_id( &self, id: &str, ) -> Result, DataLayerError> { let row = sqlx::query(&format!( "{ROUTING_GROUP_BINDING_SELECT} WHERE id = ? LIMIT 1" )) .bind(id) .fetch_optional(&self.pool) .await .map_sql_err()?; row.as_ref().map(map_binding_row).transpose() } } #[async_trait] impl RoutingGroupReadRepository for SqliteRoutingGroupRepository { async fn list_routing_groups(&self) -> Result, DataLayerError> { let rows = sqlx::query(&format!("{ROUTING_GROUP_SELECT} ORDER BY name ASC, id ASC")) .fetch_all(&self.pool) .await .map_sql_err()?; rows.iter().map(map_group_row).collect() } async fn find_routing_group( &self, lookup: RoutingGroupLookupKey<'_>, ) -> Result, DataLayerError> { let row = match lookup { RoutingGroupLookupKey::Id(id) => sqlx::query(&format!( "{ROUTING_GROUP_SELECT} WHERE id = ? LIMIT 1" )) .bind(id) .fetch_optional(&self.pool) .await .map_sql_err()?, RoutingGroupLookupKey::Name(name) => sqlx::query(&format!( "{ROUTING_GROUP_SELECT} WHERE name = ? LIMIT 1" )) .bind(name) .fetch_optional(&self.pool) .await .map_sql_err()?, RoutingGroupLookupKey::SystemDefault => sqlx::query(&format!( "{ROUTING_GROUP_SELECT} WHERE is_system_default = 1 AND enabled = 1 ORDER BY updated_at DESC, id ASC LIMIT 1" )) .fetch_optional(&self.pool) .await .map_sql_err()?, }; row.as_ref().map(map_group_row).transpose() } async fn list_routing_group_bindings( &self, query: &RoutingGroupBindingQuery, ) -> Result, DataLayerError> { let rows = sqlx::query(&format!( r#" {ROUTING_GROUP_BINDING_SELECT} WHERE (? IS NULL OR group_id = ?) AND (? IS NULL OR subject_type = ?) AND (? IS NULL OR subject_id = ?) ORDER BY created_at ASC, id ASC "# )) .bind(query.group_id.as_deref()) .bind(query.group_id.as_deref()) .bind(query.subject_type.map(binding_subject_to_database)) .bind(query.subject_type.map(binding_subject_to_database)) .bind(query.subject_id.as_deref()) .bind(query.subject_id.as_deref()) .fetch_all(&self.pool) .await .map_sql_err()?; rows.iter().map(map_binding_row).collect() } async fn has_any_routing_group_binding(&self) -> Result { let row = sqlx::query("SELECT 1 FROM routing_group_bindings LIMIT 1") .fetch_optional(&self.pool) .await .map_sql_err()?; Ok(row.is_some()) } async fn list_routing_group_versions( &self, group_id: &str, ) -> Result, DataLayerError> { let rows = sqlx::query(&format!( "{ROUTING_GROUP_VERSION_SELECT} WHERE group_id = ? ORDER BY version DESC, created_at DESC, id ASC" )) .bind(group_id) .fetch_all(&self.pool) .await .map_sql_err()?; rows.iter().map(map_version_row).collect() } } #[async_trait] impl RoutingGroupWriteRepository for SqliteRoutingGroupRepository { async fn create_routing_group( &self, record: CreateRoutingGroupRecord, ) -> Result { let group = StoredRoutingGroup::new(record)?; sqlx::query( r#" INSERT INTO routing_groups ( id, name, description, enabled, is_system_default, config_json, version, created_at, updated_at, published_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) "#, ) .bind(&group.id) .bind(&group.name) .bind(&group.description) .bind(group.enabled) .bind(group.is_system_default) .bind(json_to_string( &group.config_json, "routing_groups.config_json", )?) .bind(group.version) .bind(group.created_at) .bind(group.updated_at) .bind(group.published_at) .execute(&self.pool) .await .map_sql_err()?; Ok(group) } async fn update_routing_group( &self, id: &str, patch: UpdateRoutingGroupRecord, ) -> Result, DataLayerError> { let Some(mut group) = self.reload_group(id).await? else { return Ok(None); }; apply_group_patch(&mut group, patch)?; sqlx::query( r#" UPDATE routing_groups SET name = ?, description = ?, enabled = ?, is_system_default = ?, config_json = ?, version = ?, updated_at = ?, published_at = ? WHERE id = ? "#, ) .bind(&group.name) .bind(&group.description) .bind(group.enabled) .bind(group.is_system_default) .bind(json_to_string( &group.config_json, "routing_groups.config_json", )?) .bind(group.version) .bind(group.updated_at) .bind(group.published_at) .bind(id) .execute(&self.pool) .await .map_sql_err()?; Ok(Some(group)) } async fn delete_routing_group(&self, id: &str) -> Result { let mut tx = self.pool.begin().await.map_sql_err()?; sqlx::query("DELETE FROM routing_group_bindings WHERE group_id = ?") .bind(id) .execute(&mut *tx) .await .map_sql_err()?; sqlx::query("DELETE FROM routing_group_versions WHERE group_id = ?") .bind(id) .execute(&mut *tx) .await .map_sql_err()?; let rows_affected = sqlx::query("DELETE FROM routing_groups WHERE id = ?") .bind(id) .execute(&mut *tx) .await .map_sql_err()? .rows_affected(); tx.commit().await.map_sql_err()?; Ok(rows_affected > 0) } async fn create_routing_group_binding( &self, record: CreateRoutingGroupBindingRecord, ) -> Result { let binding = StoredRoutingGroupBinding::new(record)?; sqlx::query( r#" INSERT INTO routing_group_bindings ( id, group_id, subject_type, subject_id, is_default, allow_explicit_select, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?) "#, ) .bind(&binding.id) .bind(&binding.group_id) .bind(binding_subject_to_database(binding.subject_type)) .bind(&binding.subject_id) .bind(binding.is_default) .bind(binding.allow_explicit_select) .bind(binding.created_at) .bind(binding.updated_at) .execute(&self.pool) .await .map_sql_err()?; Ok(binding) } async fn delete_routing_group_binding(&self, id: &str) -> Result { Ok( sqlx::query("DELETE FROM routing_group_bindings WHERE id = ?") .bind(id) .execute(&self.pool) .await .map_sql_err()? .rows_affected() > 0, ) } async fn update_routing_group_binding( &self, id: &str, patch: UpdateRoutingGroupBindingRecord, ) -> Result, DataLayerError> { let Some(mut binding) = self.find_binding_by_id(id).await? else { return Ok(None); }; apply_binding_patch(&mut binding, patch)?; sqlx::query( r#" UPDATE routing_group_bindings SET group_id = ?, subject_type = ?, subject_id = ?, is_default = ?, allow_explicit_select = ?, updated_at = ? WHERE id = ? "#, ) .bind(&binding.group_id) .bind(binding_subject_to_database(binding.subject_type)) .bind(&binding.subject_id) .bind(binding.is_default) .bind(binding.allow_explicit_select) .bind(binding.updated_at) .bind(id) .execute(&self.pool) .await .map_sql_err()?; Ok(Some(binding)) } async fn create_routing_group_version( &self, record: CreateRoutingGroupVersionRecord, ) -> Result { let version = StoredRoutingGroupVersion::new(record)?; sqlx::query( r#" INSERT INTO routing_group_versions ( id, group_id, version, config_json, created_at, created_by ) VALUES (?, ?, ?, ?, ?, ?) "#, ) .bind(&version.id) .bind(&version.group_id) .bind(version.version) .bind(json_to_string( &version.config_json, "routing_group_versions.config_json", )?) .bind(version.created_at) .bind(&version.created_by) .execute(&self.pool) .await .map_sql_err()?; Ok(version) } } fn map_group_row(row: &SqliteRow) -> Result { Ok(StoredRoutingGroup { id: row.try_get("id").map_sql_err()?, name: row.try_get("name").map_sql_err()?, description: row.try_get("description").map_sql_err()?, enabled: row.try_get("enabled").map_sql_err()?, is_system_default: row.try_get("is_system_default").map_sql_err()?, config_json: json_from_string( row.try_get("config_json").map_sql_err()?, "routing_groups.config_json", )?, version: row.try_get("version").map_sql_err()?, created_at: row.try_get("created_at").map_sql_err()?, updated_at: row.try_get("updated_at").map_sql_err()?, published_at: row.try_get("published_at").map_sql_err()?, }) } fn map_binding_row(row: &SqliteRow) -> Result { Ok(StoredRoutingGroupBinding { id: row.try_get("id").map_sql_err()?, group_id: row.try_get("group_id").map_sql_err()?, subject_type: binding_subject_from_database(row.try_get("subject_type").map_sql_err()?)?, subject_id: row.try_get("subject_id").map_sql_err()?, is_default: row.try_get("is_default").map_sql_err()?, allow_explicit_select: row.try_get("allow_explicit_select").map_sql_err()?, created_at: row.try_get("created_at").map_sql_err()?, updated_at: row.try_get("updated_at").map_sql_err()?, }) } fn map_version_row(row: &SqliteRow) -> Result { Ok(StoredRoutingGroupVersion { id: row.try_get("id").map_sql_err()?, group_id: row.try_get("group_id").map_sql_err()?, version: row.try_get("version").map_sql_err()?, config_json: json_from_string( row.try_get("config_json").map_sql_err()?, "routing_group_versions.config_json", )?, created_at: row.try_get("created_at").map_sql_err()?, created_by: row.try_get("created_by").map_sql_err()?, }) } fn json_to_string(value: &Value, field_name: &str) -> Result { serde_json::to_string(value).map_err(|err| { DataLayerError::UnexpectedValue(format!("{field_name} contains unserializable JSON: {err}")) }) } fn json_from_string(value: String, field_name: &str) -> Result { serde_json::from_str(&value).map_err(|err| { DataLayerError::UnexpectedValue(format!("{field_name} contains invalid JSON: {err}")) }) } #[cfg(test)] mod tests { use serde_json::json; use super::*; use crate::run_migrations as run_sqlite_migrations; #[tokio::test] async fn sqlite_routing_group_repository_round_trips() { let pool = sqlx::sqlite::SqlitePoolOptions::new() .max_connections(1) .connect("sqlite::memory:") .await .expect("sqlite pool should connect"); run_sqlite_migrations(&pool) .await .expect("sqlite migrations should run"); let repository = SqliteRoutingGroupRepository::new(pool); repository .create_routing_group(CreateRoutingGroupRecord { id: "routing-group-1".to_string(), name: "default".to_string(), description: Some("initial".to_string()), enabled: true, is_system_default: true, config_json: json!({"allowed_models": ["gpt-*"]}), version: 1, created_at: 10, updated_at: 10, published_at: None, }) .await .expect("group should create"); let system_default = repository .find_routing_group(RoutingGroupLookupKey::SystemDefault) .await .expect("group lookup should succeed") .expect("system default should exist"); assert_eq!(system_default.id, "routing-group-1"); repository .update_routing_group( "routing-group-1", UpdateRoutingGroupRecord { description: Some(None), version: Some(2), updated_at: 20, published_at: Some(Some(20)), ..UpdateRoutingGroupRecord::default() }, ) .await .expect("group should update"); let binding = repository .create_routing_group_binding(CreateRoutingGroupBindingRecord { id: "binding-1".to_string(), group_id: "routing-group-1".to_string(), subject_type: RoutingGroupBindingSubject::ApiKey, subject_id: "api-key-1".to_string(), is_default: true, allow_explicit_select: true, created_at: 10, updated_at: 10, }) .await .expect("binding should create"); assert_eq!(binding.subject_type, RoutingGroupBindingSubject::ApiKey); assert_eq!( repository .list_routing_group_bindings(&RoutingGroupBindingQuery { group_id: Some("routing-group-1".to_string()), subject_type: Some(RoutingGroupBindingSubject::ApiKey), subject_id: Some("api-key-1".to_string()), }) .await .expect("bindings should list") .len(), 1 ); repository .create_routing_group_version(CreateRoutingGroupVersionRecord { id: "version-1".to_string(), group_id: "routing-group-1".to_string(), version: 2, config_json: json!({"allowed_models": ["gpt-*"]}), created_at: 20, created_by: Some("admin".to_string()), }) .await .expect("version should create"); assert_eq!( repository .list_routing_group_versions("routing-group-1") .await .expect("versions should list") .len(), 1 ); } }