i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-01 17:00:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: add routing profile scheduling policies

Browse Source
This commit is contained in:
fawney19
2026-05-18 11:03:49 +08:00
parent a2f91b4108
commit 92813e6122
124 changed files with 11681 additions and 578 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
crates/aether-data-contracts/Cargo.toml
View File

@@ -8,6 +8,7 @@ description = "Shared data contracts and repository traits for Aether Rust servi
[dependencies]
aether-ai-formats.workspace = true
aether-routing-core.workspace = true
async-trait.workspace = true
chrono.workspace = true
serde.workspace = true

1
crates/aether-data-contracts/src/repository/mod.rs
View File

@@ -6,6 +6,7 @@ pub mod global_models;
pub mod pool_scores;
pub mod provider_catalog;
pub mod quota;
pub mod routing_profiles;
pub mod settlement;
pub mod usage;
pub mod video_tasks;

9
crates/aether-data-contracts/src/repository/routing_profiles/mod.rs Normal file
View File

@@ -0,0 +1,9 @@
mod types;
pub use aether_routing_core::{RoutingGroupBindingSubject, RoutingGroupConfig, RoutingGroupRecord};
pub use types::{
CreateRoutingGroupBindingRecord, CreateRoutingGroupRecord, CreateRoutingGroupVersionRecord,
RoutingGroupBindingQuery, RoutingGroupLookupKey, RoutingGroupReadRepository,
RoutingGroupWriteRepository, StoredRoutingGroup, StoredRoutingGroupBinding,
StoredRoutingGroupVersion, UpdateRoutingGroupBindingRecord, UpdateRoutingGroupRecord,
};

238
crates/aether-data-contracts/src/repository/routing_profiles/types.rs Normal file
View File

@@ -0,0 +1,238 @@
use async_trait::async_trait;
use serde_json::Value;
use aether_routing_core::RoutingGroupBindingSubject;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StoredRoutingGroup {
pub id: String,
pub name: String,
pub description: Option<String>,
pub enabled: bool,
pub is_system_default: bool,
pub config_json: Value,
pub version: i64,
pub created_at: i64,
pub updated_at: i64,
pub published_at: Option<i64>,
}
impl StoredRoutingGroup {
pub fn new(record: CreateRoutingGroupRecord) -> Result<Self, crate::DataLayerError> {
validate_non_empty(&record.id, "routing_groups.id")?;
validate_non_empty(&record.name, "routing_groups.name")?;
validate_config_object(&record.config_json, "routing_groups.config_json")?;
Ok(Self {
id: record.id,
name: record.name,
description: record.description,
enabled: record.enabled,
is_system_default: record.is_system_default,
config_json: record.config_json,
version: record.version.max(1),
created_at: record.created_at,
updated_at: record.updated_at,
published_at: record.published_at,
})
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CreateRoutingGroupRecord {
pub id: String,
pub name: String,
pub description: Option<String>,
pub enabled: bool,
pub is_system_default: bool,
pub config_json: Value,
pub version: i64,
pub created_at: i64,
pub updated_at: i64,
pub published_at: Option<i64>,
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct UpdateRoutingGroupRecord {
pub name: Option<String>,
pub description: Option<Option<String>>,
pub enabled: Option<bool>,
pub is_system_default: Option<bool>,
pub config_json: Option<Value>,
pub version: Option<i64>,
pub updated_at: i64,
pub published_at: Option<Option<i64>>,
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct UpdateRoutingGroupBindingRecord {
pub group_id: Option<String>,
pub subject_type: Option<RoutingGroupBindingSubject>,
pub subject_id: Option<String>,
pub is_default: Option<bool>,
pub allow_explicit_select: Option<bool>,
pub updated_at: i64,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StoredRoutingGroupBinding {
pub id: String,
pub group_id: String,
pub subject_type: RoutingGroupBindingSubject,
pub subject_id: String,
pub is_default: bool,
pub allow_explicit_select: bool,
pub created_at: i64,
pub updated_at: i64,
}
impl StoredRoutingGroupBinding {
pub fn new(record: CreateRoutingGroupBindingRecord) -> Result<Self, crate::DataLayerError> {
validate_non_empty(&record.id, "routing_group_bindings.id")?;
validate_non_empty(&record.group_id, "routing_group_bindings.group_id")?;
validate_non_empty(&record.subject_id, "routing_group_bindings.subject_id")?;
Ok(Self {
id: record.id,
group_id: record.group_id,
subject_type: record.subject_type,
subject_id: record.subject_id,
is_default: record.is_default,
allow_explicit_select: record.allow_explicit_select,
created_at: record.created_at,
updated_at: record.updated_at,
})
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CreateRoutingGroupBindingRecord {
pub id: String,
pub group_id: String,
pub subject_type: RoutingGroupBindingSubject,
pub subject_id: String,
pub is_default: bool,
pub allow_explicit_select: bool,
pub created_at: i64,
pub updated_at: i64,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StoredRoutingGroupVersion {
pub id: String,
pub group_id: String,
pub version: i64,
pub config_json: Value,
pub created_at: i64,
pub created_by: Option<String>,
}
impl StoredRoutingGroupVersion {
pub fn new(record: CreateRoutingGroupVersionRecord) -> Result<Self, crate::DataLayerError> {
validate_non_empty(&record.id, "routing_group_versions.id")?;
validate_non_empty(&record.group_id, "routing_group_versions.group_id")?;
validate_config_object(&record.config_json, "routing_group_versions.config_json")?;
Ok(Self {
id: record.id,
group_id: record.group_id,
version: record.version.max(1),
config_json: record.config_json,
created_at: record.created_at,
created_by: record.created_by,
})
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CreateRoutingGroupVersionRecord {
pub id: String,
pub group_id: String,
pub version: i64,
pub config_json: Value,
pub created_at: i64,
pub created_by: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum RoutingGroupLookupKey<'a> {
Id(&'a str),
Name(&'a str),
SystemDefault,
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct RoutingGroupBindingQuery {
pub group_id: Option<String>,
pub subject_type: Option<RoutingGroupBindingSubject>,
pub subject_id: Option<String>,
}
#[async_trait]
pub trait RoutingGroupReadRepository: Send + Sync {
async fn list_routing_groups(&self) -> Result<Vec<StoredRoutingGroup>, crate::DataLayerError>;
async fn find_routing_group(
&self,
lookup: RoutingGroupLookupKey<'_>,
) -> Result<Option<StoredRoutingGroup>, crate::DataLayerError>;
async fn list_routing_group_bindings(
&self,
query: &RoutingGroupBindingQuery,
) -> Result<Vec<StoredRoutingGroupBinding>, crate::DataLayerError>;
async fn list_routing_group_versions(
&self,
group_id: &str,
) -> Result<Vec<StoredRoutingGroupVersion>, crate::DataLayerError>;
}
#[async_trait]
pub trait RoutingGroupWriteRepository: Send + Sync {
async fn create_routing_group(
&self,
record: CreateRoutingGroupRecord,
) -> Result<StoredRoutingGroup, crate::DataLayerError>;
async fn update_routing_group(
&self,
id: &str,
patch: UpdateRoutingGroupRecord,
) -> Result<Option<StoredRoutingGroup>, crate::DataLayerError>;
async fn delete_routing_group(&self, id: &str) -> Result<bool, crate::DataLayerError>;
async fn create_routing_group_binding(
&self,
record: CreateRoutingGroupBindingRecord,
) -> Result<StoredRoutingGroupBinding, crate::DataLayerError>;
async fn delete_routing_group_binding(&self, id: &str) -> Result<bool, crate::DataLayerError>;
async fn update_routing_group_binding(
&self,
id: &str,
patch: UpdateRoutingGroupBindingRecord,
) -> Result<Option<StoredRoutingGroupBinding>, crate::DataLayerError>;
async fn create_routing_group_version(
&self,
record: CreateRoutingGroupVersionRecord,
) -> Result<StoredRoutingGroupVersion, crate::DataLayerError>;
}
fn validate_non_empty(value: &str, field: &str) -> Result<(), crate::DataLayerError> {
if value.trim().is_empty() {
return Err(crate::DataLayerError::InvalidInput(format!(
"{field} is empty"
)));
}
Ok(())
}
fn validate_config_object(value: &Value, field: &str) -> Result<(), crate::DataLayerError> {
if !value.is_object() {
return Err(crate::DataLayerError::InvalidInput(format!(
"{field} must be a JSON object"
)));
}
Ok(())
}

41
crates/aether-data/migrations/mysql/20260516000000_add_routing_profiles.sql Normal file
View File

@@ -0,0 +1,41 @@
CREATE TABLE IF NOT EXISTS routing_groups (
`id` VARCHAR(64) NOT NULL,
`name` VARCHAR(255) NOT NULL,
`description` LONGTEXT,
`enabled` TINYINT(1) NOT NULL DEFAULT 1,
`is_system_default` TINYINT(1) NOT NULL DEFAULT 0,
`config_json` JSON NOT NULL,
`version` BIGINT NOT NULL DEFAULT 1,
`created_at` BIGINT NOT NULL,
`updated_at` BIGINT NOT NULL,
`published_at` BIGINT,
PRIMARY KEY (`id`),
UNIQUE KEY routing_groups_name_key (`name`),
KEY routing_groups_system_default_idx (`is_system_default`, `enabled`)
);
CREATE TABLE IF NOT EXISTS routing_group_bindings (
`id` VARCHAR(64) NOT NULL,
`group_id` VARCHAR(64) NOT NULL,
`subject_type` VARCHAR(32) NOT NULL,
`subject_id` VARCHAR(64) NOT NULL,
`is_default` TINYINT(1) NOT NULL DEFAULT 0,
`allow_explicit_select` TINYINT(1) NOT NULL DEFAULT 1,
`created_at` BIGINT NOT NULL,
`updated_at` BIGINT NOT NULL,
PRIMARY KEY (`id`),
KEY routing_group_bindings_group_id_idx (`group_id`),
KEY routing_group_bindings_subject_idx (`subject_type`, `subject_id`)
);
CREATE TABLE IF NOT EXISTS routing_group_versions (
`id` VARCHAR(64) NOT NULL,
`group_id` VARCHAR(64) NOT NULL,
`version` BIGINT NOT NULL,
`config_json` JSON NOT NULL,
`created_at` BIGINT NOT NULL,
`created_by` VARCHAR(64),
PRIMARY KEY (`id`),
UNIQUE KEY routing_group_versions_group_version_key (`group_id`, `version`),
KEY routing_group_versions_group_id_idx (`group_id`)
);

82
crates/aether-data/migrations/postgres/20260516000000_add_routing_profiles.sql Normal file
View File

@@ -0,0 +1,82 @@
CREATE TABLE IF NOT EXISTS public.routing_groups (
id character varying(64) NOT NULL,
name character varying(255) NOT NULL,
description text,
enabled boolean DEFAULT true NOT NULL,
is_system_default boolean DEFAULT false NOT NULL,
config_json jsonb NOT NULL,
version bigint DEFAULT 1 NOT NULL,
created_at bigint NOT NULL,
updated_at bigint NOT NULL,
published_at bigint
);
DO $$
BEGIN
IF NOT EXISTS (
SELECT 1 FROM pg_constraint WHERE conname = 'routing_groups_pkey'
) THEN
ALTER TABLE ONLY public.routing_groups
ADD CONSTRAINT routing_groups_pkey PRIMARY KEY (id);
END IF;
IF NOT EXISTS (
SELECT 1 FROM pg_constraint WHERE conname = 'routing_groups_name_key'
) THEN
ALTER TABLE ONLY public.routing_groups
ADD CONSTRAINT routing_groups_name_key UNIQUE (name);
END IF;
END $$;
CREATE INDEX IF NOT EXISTS routing_groups_system_default_idx
ON public.routing_groups USING btree (is_system_default, enabled);
CREATE TABLE IF NOT EXISTS public.routing_group_bindings (
id character varying(64) NOT NULL,
group_id character varying(64) NOT NULL,
subject_type character varying(32) NOT NULL,
subject_id character varying(64) NOT NULL,
is_default boolean DEFAULT false NOT NULL,
allow_explicit_select boolean DEFAULT true NOT NULL,
created_at bigint NOT NULL,
updated_at bigint NOT NULL
);
DO $$
BEGIN
IF NOT EXISTS (
SELECT 1 FROM pg_constraint WHERE conname = 'routing_group_bindings_pkey'
) THEN
ALTER TABLE ONLY public.routing_group_bindings
ADD CONSTRAINT routing_group_bindings_pkey PRIMARY KEY (id);
END IF;
END $$;
CREATE INDEX IF NOT EXISTS routing_group_bindings_group_id_idx
ON public.routing_group_bindings USING btree (group_id);
CREATE INDEX IF NOT EXISTS routing_group_bindings_subject_idx
ON public.routing_group_bindings USING btree (subject_type, subject_id);
CREATE TABLE IF NOT EXISTS public.routing_group_versions (
id character varying(64) NOT NULL,
group_id character varying(64) NOT NULL,
version bigint NOT NULL,
config_json jsonb NOT NULL,
created_at bigint NOT NULL,
created_by character varying(64)
);
DO $$
BEGIN
IF NOT EXISTS (
SELECT 1 FROM pg_constraint WHERE conname = 'routing_group_versions_pkey'
) THEN
ALTER TABLE ONLY public.routing_group_versions
ADD CONSTRAINT routing_group_versions_pkey PRIMARY KEY (id);
END IF;
IF NOT EXISTS (
SELECT 1 FROM pg_constraint WHERE conname = 'routing_group_versions_group_version_key'
) THEN
ALTER TABLE ONLY public.routing_group_versions
ADD CONSTRAINT routing_group_versions_group_version_key UNIQUE (group_id, version);
END IF;
END $$;
CREATE INDEX IF NOT EXISTS routing_group_versions_group_id_idx
ON public.routing_group_versions USING btree (group_id);

45
crates/aether-data/migrations/sqlite/20260516000000_add_routing_profiles.sql Normal file
View File

@@ -0,0 +1,45 @@
CREATE TABLE IF NOT EXISTS routing_groups (
id TEXT PRIMARY KEY NOT NULL,
name TEXT NOT NULL,
description TEXT,
enabled INTEGER NOT NULL DEFAULT 1,
is_system_default INTEGER NOT NULL DEFAULT 0,
config_json TEXT NOT NULL,
version INTEGER NOT NULL DEFAULT 1,
created_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL,
published_at INTEGER,
UNIQUE (name)
);
CREATE INDEX IF NOT EXISTS routing_groups_system_default_idx
ON routing_groups (is_system_default, enabled);
CREATE TABLE IF NOT EXISTS routing_group_bindings (
id TEXT PRIMARY KEY NOT NULL,
group_id TEXT NOT NULL,
subject_type TEXT NOT NULL,
subject_id TEXT NOT NULL,
is_default INTEGER NOT NULL DEFAULT 0,
allow_explicit_select INTEGER NOT NULL DEFAULT 1,
created_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS routing_group_bindings_group_id_idx
ON routing_group_bindings (group_id);
CREATE INDEX IF NOT EXISTS routing_group_bindings_subject_idx
ON routing_group_bindings (subject_type, subject_id);
CREATE TABLE IF NOT EXISTS routing_group_versions (
id TEXT PRIMARY KEY NOT NULL,
group_id TEXT NOT NULL,
version INTEGER NOT NULL,
config_json TEXT NOT NULL,
created_at INTEGER NOT NULL,
created_by TEXT,
UNIQUE (group_id, version)
);
CREATE INDEX IF NOT EXISTS routing_group_versions_group_id_idx
ON routing_group_versions (group_id);

42
crates/aether-data/schema/generated/mysql/baseline/002_provider_catalog.sql
View File

@@ -383,3 +383,45 @@ CREATE TABLE IF NOT EXISTS global_models (
UNIQUE KEY global_models_name_key (`name`)
);
CREATE TABLE IF NOT EXISTS routing_groups (
`id` VARCHAR(64) NOT NULL,
`name` VARCHAR(255) NOT NULL,
`description` LONGTEXT,
`enabled` TINYINT(1) NOT NULL DEFAULT 1,
`is_system_default` TINYINT(1) NOT NULL DEFAULT 0,
`config_json` JSON NOT NULL,
`version` BIGINT NOT NULL DEFAULT 1,
`created_at` BIGINT NOT NULL,
`updated_at` BIGINT NOT NULL,
`published_at` BIGINT,
PRIMARY KEY (`id`),
UNIQUE KEY routing_groups_name_key (`name`),
KEY routing_groups_system_default_idx (`is_system_default`, `enabled`)
);
CREATE TABLE IF NOT EXISTS routing_group_bindings (
`id` VARCHAR(64) NOT NULL,
`group_id` VARCHAR(64) NOT NULL,
`subject_type` VARCHAR(32) NOT NULL,
`subject_id` VARCHAR(64) NOT NULL,
`is_default` TINYINT(1) NOT NULL DEFAULT 0,
`allow_explicit_select` TINYINT(1) NOT NULL DEFAULT 1,
`created_at` BIGINT NOT NULL,
`updated_at` BIGINT NOT NULL,
PRIMARY KEY (`id`),
KEY routing_group_bindings_group_id_idx (`group_id`),
KEY routing_group_bindings_subject_idx (`subject_type`, `subject_id`)
);
CREATE TABLE IF NOT EXISTS routing_group_versions (
`id` VARCHAR(64) NOT NULL,
`group_id` VARCHAR(64) NOT NULL,
`version` BIGINT NOT NULL,
`config_json` JSON NOT NULL,
`created_at` BIGINT NOT NULL,
`created_by` VARCHAR(64),
PRIMARY KEY (`id`),
UNIQUE KEY routing_group_versions_group_version_key (`group_id`, `version`),
KEY routing_group_versions_group_id_idx (`group_id`)
);

45
crates/aether-data/schema/generated/postgres/baseline/002_provider_catalog.sql
View File

@@ -396,3 +396,48 @@ CREATE TABLE IF NOT EXISTS public.global_models (
ALTER TABLE ONLY public.global_models ADD CONSTRAINT global_models_pkey PRIMARY KEY (id);
ALTER TABLE ONLY public.global_models ADD CONSTRAINT global_models_name_key UNIQUE (name);
CREATE TABLE IF NOT EXISTS public.routing_groups (
id character varying(64) NOT NULL,
name character varying(255) NOT NULL,
description text,
enabled boolean DEFAULT true NOT NULL,
is_system_default boolean DEFAULT false NOT NULL,
config_json jsonb NOT NULL,
version bigint DEFAULT 1 NOT NULL,
created_at bigint NOT NULL,
updated_at bigint NOT NULL,
published_at bigint
);
ALTER TABLE ONLY public.routing_groups ADD CONSTRAINT routing_groups_pkey PRIMARY KEY (id);
ALTER TABLE ONLY public.routing_groups ADD CONSTRAINT routing_groups_name_key UNIQUE (name);
CREATE INDEX IF NOT EXISTS routing_groups_system_default_idx ON public.routing_groups USING btree (is_system_default, enabled);
CREATE TABLE IF NOT EXISTS public.routing_group_bindings (
id character varying(64) NOT NULL,
group_id character varying(64) NOT NULL,
subject_type character varying(32) NOT NULL,
subject_id character varying(64) NOT NULL,
is_default boolean DEFAULT false NOT NULL,
allow_explicit_select boolean DEFAULT true NOT NULL,
created_at bigint NOT NULL,
updated_at bigint NOT NULL
);
ALTER TABLE ONLY public.routing_group_bindings ADD CONSTRAINT routing_group_bindings_pkey PRIMARY KEY (id);
CREATE INDEX IF NOT EXISTS routing_group_bindings_group_id_idx ON public.routing_group_bindings USING btree (group_id);
CREATE INDEX IF NOT EXISTS routing_group_bindings_subject_idx ON public.routing_group_bindings USING btree (subject_type, subject_id);
CREATE TABLE IF NOT EXISTS public.routing_group_versions (
id character varying(64) NOT NULL,
group_id character varying(64) NOT NULL,
version bigint NOT NULL,
config_json jsonb NOT NULL,
created_at bigint NOT NULL,
created_by character varying(64)
);
ALTER TABLE ONLY public.routing_group_versions ADD CONSTRAINT routing_group_versions_pkey PRIMARY KEY (id);
ALTER TABLE ONLY public.routing_group_versions ADD CONSTRAINT routing_group_versions_group_version_key UNIQUE (group_id, version);
CREATE INDEX IF NOT EXISTS routing_group_versions_group_id_idx ON public.routing_group_versions USING btree (group_id);

39
crates/aether-data/schema/generated/sqlite/baseline/002_provider_catalog.sql
View File

@@ -370,3 +370,42 @@ CREATE TABLE IF NOT EXISTS global_models (
UNIQUE (name)
);
CREATE TABLE IF NOT EXISTS routing_groups (
id TEXT PRIMARY KEY NOT NULL,
name TEXT NOT NULL,
description TEXT,
enabled INTEGER NOT NULL DEFAULT 1,
is_system_default INTEGER NOT NULL DEFAULT 0,
config_json TEXT NOT NULL,
version INTEGER NOT NULL DEFAULT 1,
created_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL,
published_at INTEGER,
UNIQUE (name)
);
CREATE INDEX IF NOT EXISTS routing_groups_system_default_idx ON routing_groups (is_system_default, enabled);
CREATE TABLE IF NOT EXISTS routing_group_bindings (
id TEXT PRIMARY KEY NOT NULL,
group_id TEXT NOT NULL,
subject_type TEXT NOT NULL,
subject_id TEXT NOT NULL,
is_default INTEGER NOT NULL DEFAULT 0,
allow_explicit_select INTEGER NOT NULL DEFAULT 1,
created_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS routing_group_bindings_group_id_idx ON routing_group_bindings (group_id);
CREATE INDEX IF NOT EXISTS routing_group_bindings_subject_idx ON routing_group_bindings (subject_type, subject_id);
CREATE TABLE IF NOT EXISTS routing_group_versions (
id TEXT PRIMARY KEY NOT NULL,
group_id TEXT NOT NULL,
version INTEGER NOT NULL,
config_json TEXT NOT NULL,
created_at INTEGER NOT NULL,
created_by TEXT,
UNIQUE (group_id, version)
);
CREATE INDEX IF NOT EXISTS routing_group_versions_group_id_idx ON routing_group_versions (group_id);

152
crates/aether-data/schema/logical/002_provider_catalog.toml
View File

@@ -1694,3 +1694,155 @@ type = "unix_seconds"
[[table.global_models.uniques]]
name = "global_models_name_key"
columns = ["name"]
[table.routing_groups]
domain = "provider_catalog"
order = 110
primary_key = ["id"]
[[table.routing_groups.columns]]
name = "id"
type = "text_id"
length = 64
[[table.routing_groups.columns]]
name = "name"
type = "text"
length = 255
[[table.routing_groups.columns]]
name = "description"
type = "long_text"
nullable = true
[[table.routing_groups.columns]]
name = "enabled"
type = "bool"
default = true
[[table.routing_groups.columns]]
name = "is_system_default"
type = "bool"
default = false
[[table.routing_groups.columns]]
name = "config_json"
type = "json"
[[table.routing_groups.columns]]
name = "version"
type = "int64"
default = 1
[[table.routing_groups.columns]]
name = "created_at"
type = "unix_seconds"
[[table.routing_groups.columns]]
name = "updated_at"
type = "unix_seconds"
[[table.routing_groups.columns]]
name = "published_at"
type = "unix_seconds"
nullable = true
[[table.routing_groups.uniques]]
name = "routing_groups_name_key"
columns = ["name"]
[[table.routing_groups.indexes]]
name = "routing_groups_system_default_idx"
columns = ["is_system_default", "enabled"]
[table.routing_group_bindings]
domain = "provider_catalog"
order = 111
primary_key = ["id"]
[[table.routing_group_bindings.columns]]
name = "id"
type = "text_id"
length = 64
[[table.routing_group_bindings.columns]]
name = "group_id"
type = "text_id"
length = 64
[[table.routing_group_bindings.columns]]
name = "subject_type"
type = "text"
length = 32
[[table.routing_group_bindings.columns]]
name = "subject_id"
type = "text_id"
length = 64
[[table.routing_group_bindings.columns]]
name = "is_default"
type = "bool"
default = false
[[table.routing_group_bindings.columns]]
name = "allow_explicit_select"
type = "bool"
default = true
[[table.routing_group_bindings.columns]]
name = "created_at"
type = "unix_seconds"
[[table.routing_group_bindings.columns]]
name = "updated_at"
type = "unix_seconds"
[[table.routing_group_bindings.indexes]]
name = "routing_group_bindings_group_id_idx"
columns = ["group_id"]
[[table.routing_group_bindings.indexes]]
name = "routing_group_bindings_subject_idx"
columns = ["subject_type", "subject_id"]
[table.routing_group_versions]
domain = "provider_catalog"
order = 112
primary_key = ["id"]
[[table.routing_group_versions.columns]]
name = "id"
type = "text_id"
length = 64
[[table.routing_group_versions.columns]]
name = "group_id"
type = "text_id"
length = 64
[[table.routing_group_versions.columns]]
name = "version"
type = "int64"
[[table.routing_group_versions.columns]]
name = "config_json"
type = "json"
[[table.routing_group_versions.columns]]
name = "created_at"
type = "unix_seconds"
[[table.routing_group_versions.columns]]
name = "created_by"
type = "text_id"
length = 64
nullable = true
[[table.routing_group_versions.uniques]]
name = "routing_group_versions_group_version_key"
columns = ["group_id", "version"]
[[table.routing_group_versions.indexes]]
name = "routing_group_versions_group_id_idx"
columns = ["group_id"]

11
crates/aether-data/src/backend/mysql.rs
View File

@@ -50,6 +50,9 @@ use crate::repository::proxy_nodes::{
use crate::repository::quota::{
MysqlProviderQuotaRepository, ProviderQuotaReadRepository, ProviderQuotaWriteRepository,
};
use crate::repository::routing_profiles::{
MysqlRoutingGroupRepository, RoutingGroupReadRepository, RoutingGroupWriteRepository,
};
use crate::repository::settlement::{MysqlSettlementRepository, SettlementWriteRepository};
use crate::repository::usage::{
MysqlUsageReadRepository, MysqlUsageWriteRepository, UsageReadRepository, UsageWriteRepository,
@@ -195,6 +198,14 @@ impl MysqlBackend {
Arc::new(MysqlPoolMemberScoreRepository::new(self.pool_clone()))
}
pub fn routing_group_read_repository(&self) -> Arc<dyn RoutingGroupReadRepository> {
Arc::new(MysqlRoutingGroupRepository::new(self.pool_clone()))
}
pub fn routing_group_write_repository(&self) -> Arc<dyn RoutingGroupWriteRepository> {
Arc::new(MysqlRoutingGroupRepository::new(self.pool_clone()))
}
pub fn proxy_node_read_repository(&self) -> Arc<dyn ProxyNodeReadRepository> {
Arc::new(MysqlProxyNodeReadRepository::new(self.pool_clone()))
}

11
crates/aether-data/src/backend/postgres.rs
View File

@@ -52,6 +52,9 @@ use crate::repository::proxy_nodes::{
use crate::repository::quota::{
ProviderQuotaReadRepository, ProviderQuotaWriteRepository, SqlxProviderQuotaRepository,
};
use crate::repository::routing_profiles::{
PostgresRoutingGroupRepository, RoutingGroupReadRepository, RoutingGroupWriteRepository,
};
use crate::repository::settlement::{SettlementWriteRepository, SqlxSettlementRepository};
use crate::repository::usage::{
SqlxUsageReadRepository, UsageReadRepository, UsageWriteRepository,
@@ -206,6 +209,14 @@ impl PostgresBackend {
Arc::new(PostgresPoolMemberScoreRepository::new(self.pool_clone()))
}
pub fn routing_group_read_repository(&self) -> Arc<dyn RoutingGroupReadRepository> {
Arc::new(PostgresRoutingGroupRepository::new(self.pool_clone()))
}
pub fn routing_group_write_repository(&self) -> Arc<dyn RoutingGroupWriteRepository> {
Arc::new(PostgresRoutingGroupRepository::new(self.pool_clone()))
}
pub fn provider_quota_read_repository(&self) -> Arc<dyn ProviderQuotaReadRepository> {
Arc::new(SqlxProviderQuotaRepository::new(self.pool_clone()))
}

12
crates/aether-data/src/backend/read.rs
View File

@@ -18,6 +18,7 @@ use crate::repository::pool_scores::PoolScoreReadRepository;
use crate::repository::provider_catalog::ProviderCatalogReadRepository;
use crate::repository::proxy_nodes::ProxyNodeReadRepository;
use crate::repository::quota::ProviderQuotaReadRepository;
use crate::repository::routing_profiles::RoutingGroupReadRepository;
use crate::repository::usage::UsageReadRepository;
use crate::repository::users::UserReadRepository;
use crate::repository::video_tasks::VideoTaskReadRepository;
@@ -41,6 +42,7 @@ pub struct DataReadRepositories {
request_candidates: Option<Arc<dyn RequestCandidateReadRepository>>,
provider_catalog: Option<Arc<dyn ProviderCatalogReadRepository>>,
provider_quotas: Option<Arc<dyn ProviderQuotaReadRepository>>,
routing_groups: Option<Arc<dyn RoutingGroupReadRepository>>,
usage: Option<Arc<dyn UsageReadRepository>>,
users: Option<Arc<dyn UserReadRepository>>,
video_tasks: Option<Arc<dyn VideoTaskReadRepository>>,
@@ -72,6 +74,7 @@ impl fmt::Debug for DataReadRepositories {
.field("has_request_candidates", &self.request_candidates.is_some())
.field("has_provider_catalog", &self.provider_catalog.is_some())
.field("has_provider_quotas", &self.provider_quotas.is_some())
.field("has_routing_groups", &self.routing_groups.is_some())
.field("has_usage", &self.usage.is_some())
.field("has_users", &self.users.is_some())
.field("has_video_tasks", &self.video_tasks.is_some())
@@ -151,6 +154,10 @@ impl DataReadRepositories {
.map(PostgresBackend::provider_quota_read_repository)
.or_else(|| mysql.map(MysqlBackend::provider_quota_read_repository))
.or_else(|| sqlite.map(SqliteBackend::provider_quota_read_repository)),
routing_groups: postgres
.map(PostgresBackend::routing_group_read_repository)
.or_else(|| mysql.map(MysqlBackend::routing_group_read_repository))
.or_else(|| sqlite.map(SqliteBackend::routing_group_read_repository)),
usage: postgres
.map(PostgresBackend::usage_read_repository)
.or_else(|| mysql.map(MysqlBackend::usage_read_repository))
@@ -241,6 +248,10 @@ impl DataReadRepositories {
self.provider_quotas.clone()
}
pub fn routing_groups(&self) -> Option<Arc<dyn RoutingGroupReadRepository>> {
self.routing_groups.clone()
}
pub fn usage(&self) -> Option<Arc<dyn UsageReadRepository>> {
self.usage.clone()
}
@@ -274,6 +285,7 @@ impl DataReadRepositories {
|| self.request_candidates.is_some()
|| self.provider_catalog.is_some()
|| self.provider_quotas.is_some()
|| self.routing_groups.is_some()
|| self.usage.is_some()
|| self.users.is_some()
|| self.video_tasks.is_some()

11
crates/aether-data/src/backend/sqlite.rs
View File

@@ -50,6 +50,9 @@ use crate::repository::proxy_nodes::{
use crate::repository::quota::{
ProviderQuotaReadRepository, ProviderQuotaWriteRepository, SqliteProviderQuotaRepository,
};
use crate::repository::routing_profiles::{
RoutingGroupReadRepository, RoutingGroupWriteRepository, SqliteRoutingGroupRepository,
};
use crate::repository::settlement::{SettlementWriteRepository, SqliteSettlementRepository};
use crate::repository::usage::{
SqliteUsageReadRepository, SqliteUsageWriteRepository, UsageReadRepository,
@@ -208,6 +211,14 @@ impl SqliteBackend {
Arc::new(SqlitePoolMemberScoreRepository::new(self.pool_clone()))
}
pub fn routing_group_read_repository(&self) -> Arc<dyn RoutingGroupReadRepository> {
Arc::new(SqliteRoutingGroupRepository::new(self.pool_clone()))
}
pub fn routing_group_write_repository(&self) -> Arc<dyn RoutingGroupWriteRepository> {
Arc::new(SqliteRoutingGroupRepository::new(self.pool_clone()))
}
pub fn proxy_node_read_repository(&self) -> Arc<dyn ProxyNodeReadRepository> {
Arc::new(SqliteProxyNodeReadRepository::new(self.pool_clone()))
}

12
crates/aether-data/src/backend/write.rs
View File

@@ -15,6 +15,7 @@ use crate::repository::pool_scores::PoolMemberScoreWriteRepository;
use crate::repository::provider_catalog::ProviderCatalogWriteRepository;
use crate::repository::proxy_nodes::ProxyNodeWriteRepository;
use crate::repository::quota::ProviderQuotaWriteRepository;
use crate::repository::routing_profiles::RoutingGroupWriteRepository;
use crate::repository::settlement::SettlementWriteRepository;
use crate::repository::usage::UsageWriteRepository;
use crate::repository::video_tasks::VideoTaskWriteRepository;
@@ -35,6 +36,7 @@ pub struct DataWriteRepositories {
proxy_nodes: Option<Arc<dyn ProxyNodeWriteRepository>>,
provider_catalog: Option<Arc<dyn ProviderCatalogWriteRepository>>,
provider_quotas: Option<Arc<dyn ProviderQuotaWriteRepository>>,
routing_groups: Option<Arc<dyn RoutingGroupWriteRepository>>,
settlement: Option<Arc<dyn SettlementWriteRepository>>,
usage: Option<Arc<dyn UsageWriteRepository>>,
video_tasks: Option<Arc<dyn VideoTaskWriteRepository>>,
@@ -60,6 +62,7 @@ impl fmt::Debug for DataWriteRepositories {
.field("has_proxy_nodes", &self.proxy_nodes.is_some())
.field("has_provider_catalog", &self.provider_catalog.is_some())
.field("has_provider_quotas", &self.provider_quotas.is_some())
.field("has_routing_groups", &self.routing_groups.is_some())
.field("has_settlement", &self.settlement.is_some())
.field("has_usage", &self.usage.is_some())
.field("has_video_tasks", &self.video_tasks.is_some())
@@ -127,6 +130,10 @@ impl DataWriteRepositories {
.map(PostgresBackend::provider_quota_write_repository)
.or_else(|| mysql.map(MysqlBackend::provider_quota_write_repository))
.or_else(|| sqlite.map(SqliteBackend::provider_quota_write_repository)),
routing_groups: postgres
.map(PostgresBackend::routing_group_write_repository)
.or_else(|| mysql.map(MysqlBackend::routing_group_write_repository))
.or_else(|| sqlite.map(SqliteBackend::routing_group_write_repository)),
settlement: postgres
.map(PostgresBackend::settlement_write_repository)
.or_else(|| mysql.map(MysqlBackend::settlement_write_repository))
@@ -203,6 +210,10 @@ impl DataWriteRepositories {
self.provider_quotas.clone()
}
pub fn routing_groups(&self) -> Option<Arc<dyn RoutingGroupWriteRepository>> {
self.routing_groups.clone()
}
pub fn provider_catalog(&self) -> Option<Arc<dyn ProviderCatalogWriteRepository>> {
self.provider_catalog.clone()
}
@@ -233,6 +244,7 @@ impl DataWriteRepositories {
|| self.proxy_nodes.is_some()
|| self.provider_catalog.is_some()
|| self.provider_quotas.is_some()
|| self.routing_groups.is_some()
|| self.settlement.is_some()
|| self.usage.is_some()
|| self.video_tasks.is_some()

3
crates/aether-data/src/lifecycle/bootstrap/postgres.rs
View File

@@ -7,7 +7,7 @@ use tracing::info;
// Generated by build.rs from schema/bootstrap/postgres.
pub(crate) static EMPTY_DATABASE_SNAPSHOT_SQL: &str =
include_str!(concat!(env!("OUT_DIR"), "/empty_database_snapshot.sql"));
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260515000000;
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260516000000;
const PUBLIC_BASE_TABLE_COUNT_SQL: &str = r#"
SELECT COUNT(*)::BIGINT
@@ -29,6 +29,7 @@ WHERE table_schema = 'public'
'oauth_providers',
'provider_api_keys',
'proxy_nodes',
'routing_groups',
'user_groups',
'usage_routing_snapshots',
'usage_settlement_snapshots'

4
crates/aether-data/src/lifecycle/migrate/tests.rs
View File

@@ -305,6 +305,7 @@ fn empty_database_snapshot_covers_current_cutoff_versions() {
20260512090000,
20260512110000,
20260515000000,
20260516000000,
]
);
}
@@ -587,6 +588,7 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
20260512000000,
20260512090000,
20260512110000,
20260516000000,
]
);
assert_eq!(
@@ -603,6 +605,7 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
20260512000000,
20260512090000,
20260512110000,
20260516000000,
]
);
}
@@ -1119,6 +1122,7 @@ fn pending_migrations_from_applied_skips_versions_already_applied() {
20260512090000,
20260512110000,
20260515000000,
20260516000000,
]
);
}

1
crates/aether-data/src/repository/mod.rs
View File

@@ -21,6 +21,7 @@ pub mod provider_catalog;
pub mod provider_oauth;
pub mod proxy_nodes;
pub mod quota;
pub mod routing_profiles;
pub mod settlement;
pub mod system;
pub mod usage;

342
crates/aether-data/src/repository/routing_profiles/memory.rs Normal file
View File

@@ -0,0 +1,342 @@
use std::collections::BTreeMap;
use std::sync::RwLock;
use async_trait::async_trait;
use super::{
CreateRoutingGroupBindingRecord, CreateRoutingGroupRecord, CreateRoutingGroupVersionRecord,
RoutingGroupBindingQuery, RoutingGroupLookupKey, RoutingGroupReadRepository,
RoutingGroupWriteRepository, StoredRoutingGroup, StoredRoutingGroupBinding,
StoredRoutingGroupVersion, UpdateRoutingGroupBindingRecord, UpdateRoutingGroupRecord,
};
use crate::DataLayerError;
#[derive(Debug, Default)]
pub struct InMemoryRoutingGroupRepository {
groups: RwLock<BTreeMap<String, StoredRoutingGroup>>,
bindings: RwLock<BTreeMap<String, StoredRoutingGroupBinding>>,
versions: RwLock<BTreeMap<String, StoredRoutingGroupVersion>>,
}
impl InMemoryRoutingGroupRepository {
pub fn seed<I, B, V>(groups: I, bindings: B, versions: V) -> Self
where
I: IntoIterator<Item = StoredRoutingGroup>,
B: IntoIterator<Item = StoredRoutingGroupBinding>,
V: IntoIterator<Item = StoredRoutingGroupVersion>,
{
Self {
groups: RwLock::new(
groups
.into_iter()
.map(|item| (item.id.clone(), item))
.collect(),
),
bindings: RwLock::new(
bindings
.into_iter()
.map(|item| (item.id.clone(), item))
.collect(),
),
versions: RwLock::new(
versions
.into_iter()
.map(|item| (item.id.clone(), item))
.collect(),
),
}
}
}
#[async_trait]
impl RoutingGroupReadRepository for InMemoryRoutingGroupRepository {
async fn list_routing_groups(&self) -> Result<Vec<StoredRoutingGroup>, DataLayerError> {
let mut groups = self
.groups
.read()
.expect("routing group repository lock")
.values()
.cloned()
.collect::<Vec<_>>();
groups.sort_by(|left, right| left.name.cmp(&right.name).then(left.id.cmp(&right.id)));
Ok(groups)
}
async fn find_routing_group(
&self,
lookup: RoutingGroupLookupKey<'_>,
) -> Result<Option<StoredRoutingGroup>, DataLayerError> {
let groups = self.groups.read().expect("routing group repository lock");
Ok(match lookup {
RoutingGroupLookupKey::Id(id) => groups.get(id).cloned(),
RoutingGroupLookupKey::Name(name) => {
groups.values().find(|group| group.name == name).cloned()
}
RoutingGroupLookupKey::SystemDefault => groups
.values()
.find(|group| group.is_system_default && group.enabled)
.cloned(),
})
}
async fn list_routing_group_bindings(
&self,
query: &RoutingGroupBindingQuery,
) -> Result<Vec<StoredRoutingGroupBinding>, DataLayerError> {
let mut rows = self
.bindings
.read()
.expect("routing group binding repository lock")
.values()
.filter(|row| {
query
.group_id
.as_ref()
.is_none_or(|group_id| &row.group_id == group_id)
&& query
.subject_type
.as_ref()
.is_none_or(|subject_type| &row.subject_type == subject_type)
&& query
.subject_id
.as_ref()
.is_none_or(|subject_id| &row.subject_id == subject_id)
})
.cloned()
.collect::<Vec<_>>();
rows.sort_by(|left, right| {
left.created_at
.cmp(&right.created_at)
.then(left.id.cmp(&right.id))
});
Ok(rows)
}
async fn list_routing_group_versions(
&self,
group_id: &str,
) -> Result<Vec<StoredRoutingGroupVersion>, DataLayerError> {
let mut rows = self
.versions
.read()
.expect("routing group version repository lock")
.values()
.filter(|row| row.group_id == group_id)
.cloned()
.collect::<Vec<_>>();
rows.sort_by(|left, right| {
right
.version
.cmp(&left.version)
.then(right.created_at.cmp(&left.created_at))
});
Ok(rows)
}
}
#[async_trait]
impl RoutingGroupWriteRepository for InMemoryRoutingGroupRepository {
async fn create_routing_group(
&self,
record: CreateRoutingGroupRecord,
) -> Result<StoredRoutingGroup, DataLayerError> {
let group = StoredRoutingGroup::new(record)?;
self.groups
.write()
.expect("routing group repository lock")
.insert(group.id.clone(), group.clone());
Ok(group)
}
async fn update_routing_group(
&self,
id: &str,
patch: UpdateRoutingGroupRecord,
) -> Result<Option<StoredRoutingGroup>, DataLayerError> {
let mut groups = self.groups.write().expect("routing group repository lock");
let Some(group) = groups.get_mut(id) else {
return Ok(None);
};
if let Some(name) = patch.name {
if name.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"routing_groups.name is empty".to_string(),
));
}
group.name = name;
}
if let Some(description) = patch.description {
group.description = description;
}
if let Some(enabled) = patch.enabled {
group.enabled = enabled;
}
if let Some(is_system_default) = patch.is_system_default {
group.is_system_default = is_system_default;
}
if let Some(config_json) = patch.config_json {
if !config_json.is_object() {
return Err(DataLayerError::InvalidInput(
"routing_groups.config_json must be a JSON object".to_string(),
));
}
group.config_json = config_json;
}
if let Some(version) = patch.version {
group.version = version.max(1);
}
if let Some(published_at) = patch.published_at {
group.published_at = published_at;
}
group.updated_at = patch.updated_at;
Ok(Some(group.clone()))
}
async fn delete_routing_group(&self, id: &str) -> Result<bool, DataLayerError> {
Ok(self
.groups
.write()
.expect("routing group repository lock")
.remove(id)
.is_some())
}
async fn create_routing_group_binding(
&self,
record: CreateRoutingGroupBindingRecord,
) -> Result<StoredRoutingGroupBinding, DataLayerError> {
let binding = StoredRoutingGroupBinding::new(record)?;
self.bindings
.write()
.expect("routing group binding repository lock")
.insert(binding.id.clone(), binding.clone());
Ok(binding)
}
async fn delete_routing_group_binding(&self, id: &str) -> Result<bool, DataLayerError> {
Ok(self
.bindings
.write()
.expect("routing group binding repository lock")
.remove(id)
.is_some())
}
async fn update_routing_group_binding(
&self,
id: &str,
patch: UpdateRoutingGroupBindingRecord,
) -> Result<Option<StoredRoutingGroupBinding>, DataLayerError> {
let mut bindings = self
.bindings
.write()
.expect("routing group binding repository lock");
let Some(binding) = bindings.get_mut(id) else {
return Ok(None);
};
if let Some(group_id) = patch.group_id {
if group_id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"routing_group_bindings.group_id is empty".to_string(),
));
}
binding.group_id = group_id;
}
if let Some(subject_type) = patch.subject_type {
binding.subject_type = subject_type;
}
if let Some(subject_id) = patch.subject_id {
if subject_id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"routing_group_bindings.subject_id is empty".to_string(),
));
}
binding.subject_id = subject_id;
}
if let Some(is_default) = patch.is_default {
binding.is_default = is_default;
}
if let Some(allow_explicit_select) = patch.allow_explicit_select {
binding.allow_explicit_select = allow_explicit_select;
}
binding.updated_at = patch.updated_at;
Ok(Some(binding.clone()))
}
async fn create_routing_group_version(
&self,
record: CreateRoutingGroupVersionRecord,
) -> Result<StoredRoutingGroupVersion, DataLayerError> {
let version = StoredRoutingGroupVersion::new(record)?;
self.versions
.write()
.expect("routing group version repository lock")
.insert(version.id.clone(), version.clone());
Ok(version)
}
}
#[cfg(test)]
mod tests {
use aether_data_contracts::repository::routing_profiles::RoutingGroupBindingSubject;
use serde_json::json;
use super::*;
#[tokio::test]
async fn stores_groups_bindings_and_versions() {
let repository = InMemoryRoutingGroupRepository::default();
let group = repository
.create_routing_group(CreateRoutingGroupRecord {
id: "group-1".to_string(),
name: "default".to_string(),
description: None,
enabled: true,
is_system_default: true,
config_json: json!({}),
version: 1,
created_at: 1,
updated_at: 1,
published_at: None,
})
.await
.expect("group should store");
assert_eq!(
repository
.find_routing_group(RoutingGroupLookupKey::SystemDefault)
.await
.unwrap()
.as_ref()
.map(|group| group.id.as_str()),
Some(group.id.as_str())
);
repository
.create_routing_group_binding(CreateRoutingGroupBindingRecord {
id: "binding-1".to_string(),
group_id: "group-1".to_string(),
subject_type: RoutingGroupBindingSubject::ApiKey,
subject_id: "api-key-1".to_string(),
is_default: true,
allow_explicit_select: true,
created_at: 1,
updated_at: 1,
})
.await
.unwrap();
assert_eq!(
repository
.list_routing_group_bindings(&RoutingGroupBindingQuery {
subject_type: Some(RoutingGroupBindingSubject::ApiKey),
subject_id: Some("api-key-1".to_string()),
group_id: None,
})
.await
.unwrap()
.len(),
1
);
}
}

16
crates/aether-data/src/repository/routing_profiles/mod.rs Normal file
View File

@@ -0,0 +1,16 @@
mod memory;
mod mysql;
mod postgres;
mod sqlite;
pub(crate) use aether_data_contracts::repository::routing_profiles::{
CreateRoutingGroupBindingRecord, CreateRoutingGroupRecord, CreateRoutingGroupVersionRecord,
RoutingGroupBindingQuery, RoutingGroupBindingSubject, RoutingGroupLookupKey,
RoutingGroupReadRepository, RoutingGroupWriteRepository, StoredRoutingGroup,
StoredRoutingGroupBinding, StoredRoutingGroupVersion, UpdateRoutingGroupBindingRecord,
UpdateRoutingGroupRecord,
};
pub use memory::InMemoryRoutingGroupRepository;
pub use mysql::MysqlRoutingGroupRepository;
pub use postgres::PostgresRoutingGroupRepository;
pub use sqlite::SqliteRoutingGroupRepository;

417
crates/aether-data/src/repository/routing_profiles/mysql.rs Normal file
View File

@@ -0,0 +1,417 @@
use async_trait::async_trait;
use serde_json::Value;
use sqlx::{mysql::MySqlRow, Row};
use super::postgres::{
apply_binding_patch, apply_group_patch, binding_subject_from_database,
binding_subject_to_database,
};
use super::*;
use crate::driver::mysql::MysqlPool;
use crate::error::SqlResultExt;
use crate::DataLayerError;
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 MysqlRoutingGroupRepository {
pool: MysqlPool,
}
impl MysqlRoutingGroupRepository {
pub fn new(pool: MysqlPool) -> Self {
Self { pool }
}
async fn reload_group(&self, id: &str) -> Result<Option<StoredRoutingGroup>, DataLayerError> {
self.find_routing_group(RoutingGroupLookupKey::Id(id)).await
}
async fn find_binding_by_id(
&self,
id: &str,
) -> Result<Option<StoredRoutingGroupBinding>, 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 MysqlRoutingGroupRepository {
async fn list_routing_groups(&self) -> Result<Vec<StoredRoutingGroup>, 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<Option<StoredRoutingGroup>, 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<Vec<StoredRoutingGroupBinding>, 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 list_routing_group_versions(
&self,
group_id: &str,
) -> Result<Vec<StoredRoutingGroupVersion>, 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 MysqlRoutingGroupRepository {
async fn create_routing_group(
&self,
record: CreateRoutingGroupRecord,
) -> Result<StoredRoutingGroup, DataLayerError> {
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<Option<StoredRoutingGroup>, 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<bool, DataLayerError> {
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<StoredRoutingGroupBinding, DataLayerError> {
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<bool, DataLayerError> {
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<Option<StoredRoutingGroupBinding>, 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<StoredRoutingGroupVersion, DataLayerError> {
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: &MySqlRow) -> Result<StoredRoutingGroup, DataLayerError> {
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: &MySqlRow) -> Result<StoredRoutingGroupBinding, DataLayerError> {
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: &MySqlRow) -> Result<StoredRoutingGroupVersion, DataLayerError> {
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<String, DataLayerError> {
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<Value, DataLayerError> {
serde_json::from_str(&value).map_err(|err| {
DataLayerError::UnexpectedValue(format!("{field_name} contains invalid JSON: {err}"))
})
}

483
crates/aether-data/src/repository/routing_profiles/postgres.rs Normal file
View File

@@ -0,0 +1,483 @@
use async_trait::async_trait;
use futures_util::TryStreamExt;
use sqlx::{postgres::PgRow, PgPool, Row};
use super::*;
use crate::error::SqlxResultExt;
use crate::DataLayerError;
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 PostgresRoutingGroupRepository {
pool: PgPool,
}
impl PostgresRoutingGroupRepository {
pub fn new(pool: PgPool) -> Self {
Self { pool }
}
async fn reload_group(&self, id: &str) -> Result<Option<StoredRoutingGroup>, DataLayerError> {
self.find_routing_group(RoutingGroupLookupKey::Id(id)).await
}
async fn find_binding_by_id(
&self,
id: &str,
) -> Result<Option<StoredRoutingGroupBinding>, DataLayerError> {
let row = sqlx::query(&format!(
"{ROUTING_GROUP_BINDING_SELECT} WHERE id = $1 LIMIT 1"
))
.bind(id)
.fetch_optional(&self.pool)
.await
.map_postgres_err()?;
row.as_ref().map(map_binding_row).transpose()
}
}
#[async_trait]
impl RoutingGroupReadRepository for PostgresRoutingGroupRepository {
async fn list_routing_groups(&self) -> Result<Vec<StoredRoutingGroup>, DataLayerError> {
let sql = format!("{ROUTING_GROUP_SELECT} ORDER BY name ASC, id ASC");
let mut rows = sqlx::query(&sql).fetch(&self.pool);
let mut groups = Vec::new();
while let Some(row) = rows.try_next().await.map_postgres_err()? {
groups.push(map_group_row(&row)?);
}
Ok(groups)
}
async fn find_routing_group(
&self,
lookup: RoutingGroupLookupKey<'_>,
) -> Result<Option<StoredRoutingGroup>, DataLayerError> {
let row = match lookup {
RoutingGroupLookupKey::Id(id) => sqlx::query(&format!(
"{ROUTING_GROUP_SELECT} WHERE id = $1 LIMIT 1"
))
.bind(id)
.fetch_optional(&self.pool)
.await
.map_postgres_err()?,
RoutingGroupLookupKey::Name(name) => sqlx::query(&format!(
"{ROUTING_GROUP_SELECT} WHERE name = $1 LIMIT 1"
))
.bind(name)
.fetch_optional(&self.pool)
.await
.map_postgres_err()?,
RoutingGroupLookupKey::SystemDefault => sqlx::query(&format!(
"{ROUTING_GROUP_SELECT} WHERE is_system_default = TRUE AND enabled = TRUE ORDER BY updated_at DESC, id ASC LIMIT 1"
))
.fetch_optional(&self.pool)
.await
.map_postgres_err()?,
};
row.as_ref().map(map_group_row).transpose()
}
async fn list_routing_group_bindings(
&self,
query: &RoutingGroupBindingQuery,
) -> Result<Vec<StoredRoutingGroupBinding>, DataLayerError> {
let sql = format!(
r#"
{ROUTING_GROUP_BINDING_SELECT}
WHERE ($1::text IS NULL OR group_id = $1)
AND ($2::text IS NULL OR subject_type = $2)
AND ($3::text IS NULL OR subject_id = $3)
ORDER BY created_at ASC, id ASC
"#
);
let mut rows = sqlx::query(&sql)
.bind(query.group_id.as_deref())
.bind(query.subject_type.map(binding_subject_to_database))
.bind(query.subject_id.as_deref())
.fetch(&self.pool);
let mut bindings = Vec::new();
while let Some(row) = rows.try_next().await.map_postgres_err()? {
bindings.push(map_binding_row(&row)?);
}
Ok(bindings)
}
async fn list_routing_group_versions(
&self,
group_id: &str,
) -> Result<Vec<StoredRoutingGroupVersion>, DataLayerError> {
let sql = format!(
"{ROUTING_GROUP_VERSION_SELECT} WHERE group_id = $1 ORDER BY version DESC, created_at DESC, id ASC"
);
let mut rows = sqlx::query(&sql).bind(group_id).fetch(&self.pool);
let mut versions = Vec::new();
while let Some(row) = rows.try_next().await.map_postgres_err()? {
versions.push(map_version_row(&row)?);
}
Ok(versions)
}
}
#[async_trait]
impl RoutingGroupWriteRepository for PostgresRoutingGroupRepository {
async fn create_routing_group(
&self,
record: CreateRoutingGroupRecord,
) -> Result<StoredRoutingGroup, DataLayerError> {
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 ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
"#,
)
.bind(&group.id)
.bind(&group.name)
.bind(&group.description)
.bind(group.enabled)
.bind(group.is_system_default)
.bind(&group.config_json)
.bind(group.version)
.bind(group.created_at)
.bind(group.updated_at)
.bind(group.published_at)
.execute(&self.pool)
.await
.map_postgres_err()?;
Ok(group)
}
async fn update_routing_group(
&self,
id: &str,
patch: UpdateRoutingGroupRecord,
) -> Result<Option<StoredRoutingGroup>, 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 = $2,
description = $3,
enabled = $4,
is_system_default = $5,
config_json = $6,
version = $7,
updated_at = $8,
published_at = $9
WHERE id = $1
"#,
)
.bind(id)
.bind(&group.name)
.bind(&group.description)
.bind(group.enabled)
.bind(group.is_system_default)
.bind(&group.config_json)
.bind(group.version)
.bind(group.updated_at)
.bind(group.published_at)
.execute(&self.pool)
.await
.map_postgres_err()?;
Ok(Some(group))
}
async fn delete_routing_group(&self, id: &str) -> Result<bool, DataLayerError> {
let mut tx = self.pool.begin().await.map_postgres_err()?;
sqlx::query("DELETE FROM routing_group_bindings WHERE group_id = $1")
.bind(id)
.execute(&mut *tx)
.await
.map_postgres_err()?;
sqlx::query("DELETE FROM routing_group_versions WHERE group_id = $1")
.bind(id)
.execute(&mut *tx)
.await
.map_postgres_err()?;
let rows_affected = sqlx::query("DELETE FROM routing_groups WHERE id = $1")
.bind(id)
.execute(&mut *tx)
.await
.map_postgres_err()?
.rows_affected();
tx.commit().await.map_postgres_err()?;
Ok(rows_affected > 0)
}
async fn create_routing_group_binding(
&self,
record: CreateRoutingGroupBindingRecord,
) -> Result<StoredRoutingGroupBinding, DataLayerError> {
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 ($1, $2, $3, $4, $5, $6, $7, $8)
"#,
)
.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_postgres_err()?;
Ok(binding)
}
async fn delete_routing_group_binding(&self, id: &str) -> Result<bool, DataLayerError> {
Ok(
sqlx::query("DELETE FROM routing_group_bindings WHERE id = $1")
.bind(id)
.execute(&self.pool)
.await
.map_postgres_err()?
.rows_affected()
> 0,
)
}
async fn update_routing_group_binding(
&self,
id: &str,
patch: UpdateRoutingGroupBindingRecord,
) -> Result<Option<StoredRoutingGroupBinding>, 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 = $2,
subject_type = $3,
subject_id = $4,
is_default = $5,
allow_explicit_select = $6,
updated_at = $7
WHERE id = $1
"#,
)
.bind(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)
.execute(&self.pool)
.await
.map_postgres_err()?;
Ok(Some(binding))
}
async fn create_routing_group_version(
&self,
record: CreateRoutingGroupVersionRecord,
) -> Result<StoredRoutingGroupVersion, DataLayerError> {
let version = StoredRoutingGroupVersion::new(record)?;
sqlx::query(
r#"
INSERT INTO routing_group_versions (
id, group_id, version, config_json, created_at, created_by
)
VALUES ($1, $2, $3, $4, $5, $6)
"#,
)
.bind(&version.id)
.bind(&version.group_id)
.bind(version.version)
.bind(&version.config_json)
.bind(version.created_at)
.bind(&version.created_by)
.execute(&self.pool)
.await
.map_postgres_err()?;
Ok(version)
}
}
pub(super) fn apply_group_patch(
group: &mut StoredRoutingGroup,
patch: UpdateRoutingGroupRecord,
) -> Result<(), DataLayerError> {
if let Some(name) = patch.name {
if name.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"routing_groups.name is empty".to_string(),
));
}
group.name = name;
}
if let Some(description) = patch.description {
group.description = description;
}
if let Some(enabled) = patch.enabled {
group.enabled = enabled;
}
if let Some(is_system_default) = patch.is_system_default {
group.is_system_default = is_system_default;
}
if let Some(config_json) = patch.config_json {
if !config_json.is_object() {
return Err(DataLayerError::InvalidInput(
"routing_groups.config_json must be a JSON object".to_string(),
));
}
group.config_json = config_json;
}
if let Some(version) = patch.version {
group.version = version.max(1);
}
if let Some(published_at) = patch.published_at {
group.published_at = published_at;
}
group.updated_at = patch.updated_at;
Ok(())
}
pub(super) fn apply_binding_patch(
binding: &mut StoredRoutingGroupBinding,
patch: UpdateRoutingGroupBindingRecord,
) -> Result<(), DataLayerError> {
if let Some(group_id) = patch.group_id {
if group_id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"routing_group_bindings.group_id is empty".to_string(),
));
}
binding.group_id = group_id;
}
if let Some(subject_type) = patch.subject_type {
binding.subject_type = subject_type;
}
if let Some(subject_id) = patch.subject_id {
if subject_id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"routing_group_bindings.subject_id is empty".to_string(),
));
}
binding.subject_id = subject_id;
}
if let Some(is_default) = patch.is_default {
binding.is_default = is_default;
}
if let Some(allow_explicit_select) = patch.allow_explicit_select {
binding.allow_explicit_select = allow_explicit_select;
}
binding.updated_at = patch.updated_at;
Ok(())
}
pub(super) fn binding_subject_to_database(subject: RoutingGroupBindingSubject) -> &'static str {
match subject {
RoutingGroupBindingSubject::User => "user",
RoutingGroupBindingSubject::ApiKey => "api_key",
RoutingGroupBindingSubject::UserGroup => "user_group",
}
}
pub(super) fn binding_subject_from_database(
value: String,
) -> Result<RoutingGroupBindingSubject, DataLayerError> {
match value.as_str() {
"user" => Ok(RoutingGroupBindingSubject::User),
"api_key" => Ok(RoutingGroupBindingSubject::ApiKey),
"user_group" => Ok(RoutingGroupBindingSubject::UserGroup),
_ => Err(DataLayerError::UnexpectedValue(format!(
"invalid routing_group_bindings.subject_type: {value}"
))),
}
}
fn map_group_row(row: &PgRow) -> Result<StoredRoutingGroup, DataLayerError> {
Ok(StoredRoutingGroup {
id: row.try_get("id").map_postgres_err()?,
name: row.try_get("name").map_postgres_err()?,
description: row.try_get("description").map_postgres_err()?,
enabled: row.try_get("enabled").map_postgres_err()?,
is_system_default: row.try_get("is_system_default").map_postgres_err()?,
config_json: row.try_get("config_json").map_postgres_err()?,
version: row.try_get("version").map_postgres_err()?,
created_at: row.try_get("created_at").map_postgres_err()?,
updated_at: row.try_get("updated_at").map_postgres_err()?,
published_at: row.try_get("published_at").map_postgres_err()?,
})
}
fn map_binding_row(row: &PgRow) -> Result<StoredRoutingGroupBinding, DataLayerError> {
Ok(StoredRoutingGroupBinding {
id: row.try_get("id").map_postgres_err()?,
group_id: row.try_get("group_id").map_postgres_err()?,
subject_type: binding_subject_from_database(
row.try_get("subject_type").map_postgres_err()?,
)?,
subject_id: row.try_get("subject_id").map_postgres_err()?,
is_default: row.try_get("is_default").map_postgres_err()?,
allow_explicit_select: row.try_get("allow_explicit_select").map_postgres_err()?,
created_at: row.try_get("created_at").map_postgres_err()?,
updated_at: row.try_get("updated_at").map_postgres_err()?,
})
}
fn map_version_row(row: &PgRow) -> Result<StoredRoutingGroupVersion, DataLayerError> {
Ok(StoredRoutingGroupVersion {
id: row.try_get("id").map_postgres_err()?,
group_id: row.try_get("group_id").map_postgres_err()?,
version: row.try_get("version").map_postgres_err()?,
config_json: row.try_get("config_json").map_postgres_err()?,
created_at: row.try_get("created_at").map_postgres_err()?,
created_by: row.try_get("created_by").map_postgres_err()?,
})
}

524
crates/aether-data/src/repository/routing_profiles/sqlite.rs Normal file
View File

@@ -0,0 +1,524 @@
use async_trait::async_trait;
use serde_json::Value;
use sqlx::{sqlite::SqliteRow, Row};
use super::postgres::{
apply_binding_patch, apply_group_patch, binding_subject_from_database,
binding_subject_to_database,
};
use super::*;
use crate::driver::sqlite::SqlitePool;
use crate::error::SqlResultExt;
use crate::DataLayerError;
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<Option<StoredRoutingGroup>, DataLayerError> {
self.find_routing_group(RoutingGroupLookupKey::Id(id)).await
}
async fn find_binding_by_id(
&self,
id: &str,
) -> Result<Option<StoredRoutingGroupBinding>, 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<Vec<StoredRoutingGroup>, 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<Option<StoredRoutingGroup>, 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<Vec<StoredRoutingGroupBinding>, 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 list_routing_group_versions(
&self,
group_id: &str,
) -> Result<Vec<StoredRoutingGroupVersion>, 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<StoredRoutingGroup, DataLayerError> {
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<Option<StoredRoutingGroup>, 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<bool, DataLayerError> {
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<StoredRoutingGroupBinding, DataLayerError> {
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<bool, DataLayerError> {
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<Option<StoredRoutingGroupBinding>, 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<StoredRoutingGroupVersion, DataLayerError> {
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<StoredRoutingGroup, DataLayerError> {
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<StoredRoutingGroupBinding, DataLayerError> {
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<StoredRoutingGroupVersion, DataLayerError> {
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<String, DataLayerError> {
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<Value, DataLayerError> {
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::lifecycle::migrate::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
);
}
}

13
crates/aether-routing-core/Cargo.toml Normal file
View File

@@ -0,0 +1,13 @@
[package]
name = "aether-routing-core"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
description = "Pure routing profile policy, mutation, ranking overlay, and trace primitives for Aether"
[dependencies]
regex.workspace = true
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true

91
crates/aether-routing-core/src/actions.rs Normal file
View File

@@ -0,0 +1,91 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RoutingRulePhase {
#[default]
ClientRequest,
ProviderRequest,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RoutingSetPriorityMode {
#[default]
Provider,
GlobalKey,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RoutingSchedulingMode {
#[default]
CacheAffinity,
LoadBalance,
FixedOrder,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case", tag = "op")]
pub enum RoutingJsonPatchOperation {
Add { path: String, value: Value },
Replace { path: String, value: Value },
Remove { path: String },
}
impl RoutingJsonPatchOperation {
pub fn path(&self) -> &str {
match self {
Self::Add { path, .. } | Self::Replace { path, .. } | Self::Remove { path } => path,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case", tag = "op")]
pub enum RoutingHeaderPatch {
Set { name: String, value: String },
Remove { name: String },
}
impl RoutingHeaderPatch {
pub fn name(&self) -> &str {
match self {
Self::Set { name, .. } | Self::Remove { name } => name,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case", tag = "type")]
pub enum RoutingAction {
RestrictModels {
models: Vec<String>,
},
RestrictProviders {
provider_ids: Vec<String>,
},
RestrictKeys {
key_ids: Vec<String>,
},
SetScheduling {
priority_mode: Option<RoutingSetPriorityMode>,
scheduling_mode: Option<RoutingSchedulingMode>,
keep_priority_on_conversion: Option<bool>,
},
SetProviderPriority {
provider_id: String,
priority: i32,
},
SetKeyPriority {
key_id: String,
priority: i32,
},
JsonPatchBody {
patch: Vec<RoutingJsonPatchOperation>,
},
PatchHeaders {
patch: Vec<RoutingHeaderPatch>,
},
}

232
crates/aether-routing-core/src/conditions.rs Normal file
View File

@@ -0,0 +1,232 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RoutingConditionOp {
Eq,
Ne,
In,
Contains,
Exists,
Prefix,
Suffix,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum RoutingCondition {
All {
all: Vec<RoutingCondition>,
},
Any {
any: Vec<RoutingCondition>,
},
Not {
not: Box<RoutingCondition>,
},
Predicate {
field: String,
op: RoutingConditionOp,
#[serde(default, skip_serializing_if = "Option::is_none")]
value: Option<Value>,
},
Empty {},
}
impl Default for RoutingCondition {
fn default() -> Self {
Self::Empty {}
}
}
#[derive(Debug, Clone)]
pub struct RoutingConditionContext<'a> {
pub model: &'a str,
pub api_format: &'a str,
pub user_id: Option<&'a str>,
pub api_key_id: Option<&'a str>,
pub headers: &'a Value,
pub body: &'a Value,
}
impl RoutingCondition {
pub fn matches(&self, context: &RoutingConditionContext<'_>) -> bool {
match self {
Self::All { all } => all.iter().all(|condition| condition.matches(context)),
Self::Any { any } => any.iter().any(|condition| condition.matches(context)),
Self::Not { not } => !not.matches(context),
Self::Predicate { field, op, value } => {
let actual = resolve_field(context, field);
compare_condition(actual, *op, value.as_ref())
}
Self::Empty {} => true,
}
}
}
fn resolve_field(context: &RoutingConditionContext<'_>, field: &str) -> Option<Value> {
let normalized = field.trim();
match normalized {
"model" => return Some(Value::String(context.model.to_string())),
"api_format" | "client_api_format" => {
return Some(Value::String(context.api_format.to_string()))
}
"user_id" => {
return context
.user_id
.map(|value| Value::String(value.to_string()))
}
"api_key_id" => {
return context
.api_key_id
.map(|value| Value::String(value.to_string()))
}
_ => {}
}
if let Some(path) = normalized.strip_prefix("headers.") {
return lookup_dotted_path(context.headers, path).cloned();
}
if let Some(path) = normalized.strip_prefix("body.") {
return lookup_dotted_path(context.body, path).cloned();
}
None
}
fn lookup_dotted_path<'a>(root: &'a Value, path: &str) -> Option<&'a Value> {
let mut current = root;
for part in path.split('.').filter(|part| !part.is_empty()) {
match current {
Value::Object(map) => current = map.get(part)?,
Value::Array(items) => {
let index = part.parse::<usize>().ok()?;
current = items.get(index)?;
}
_ => return None,
}
}
Some(current)
}
fn compare_condition(
actual: Option<Value>,
op: RoutingConditionOp,
expected: Option<&Value>,
) -> bool {
match op {
RoutingConditionOp::Exists => actual.is_some(),
RoutingConditionOp::Eq => actual
.as_ref()
.zip(expected)
.is_some_and(|(actual, expected)| values_equal(actual, expected)),
RoutingConditionOp::Ne => actual
.as_ref()
.zip(expected)
.is_none_or(|(actual, expected)| !values_equal(actual, expected)),
RoutingConditionOp::In => {
actual
.as_ref()
.zip(expected)
.is_some_and(|(actual, expected)| {
expected
.as_array()
.is_some_and(|items| items.iter().any(|item| values_equal(actual, item)))
})
}
RoutingConditionOp::Contains => {
actual
.as_ref()
.zip(expected)
.is_some_and(|(actual, expected)| {
let Some(expected) = expected.as_str() else {
return false;
};
value_as_string(actual).is_some_and(|actual| actual.contains(expected))
})
}
RoutingConditionOp::Prefix => {
actual
.as_ref()
.zip(expected)
.is_some_and(|(actual, expected)| {
let Some(expected) = expected.as_str() else {
return false;
};
value_as_string(actual).is_some_and(|actual| actual.starts_with(expected))
})
}
RoutingConditionOp::Suffix => {
actual
.as_ref()
.zip(expected)
.is_some_and(|(actual, expected)| {
let Some(expected) = expected.as_str() else {
return false;
};
value_as_string(actual).is_some_and(|actual| actual.ends_with(expected))
})
}
}
}
fn values_equal(left: &Value, right: &Value) -> bool {
match (value_as_string(left), value_as_string(right)) {
(Some(left), Some(right)) => left == right,
_ => left == right,
}
}
fn value_as_string(value: &Value) -> Option<String> {
match value {
Value::String(value) => Some(value.clone()),
Value::Number(value) => Some(value.to_string()),
Value::Bool(value) => Some(value.to_string()),
_ => None,
}
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::*;
fn context<'a>(headers: &'a Value, body: &'a Value) -> RoutingConditionContext<'a> {
RoutingConditionContext {
model: "gpt-5",
api_format: "openai:chat",
user_id: Some("user-1"),
api_key_id: Some("key-1"),
headers,
body,
}
}
#[test]
fn matches_all_body_and_header_predicates() {
let condition = RoutingCondition::All {
all: vec![
RoutingCondition::Predicate {
field: "model".to_string(),
op: RoutingConditionOp::Eq,
value: Some(json!("gpt-5")),
},
RoutingCondition::Predicate {
field: "headers.x-app".to_string(),
op: RoutingConditionOp::Eq,
value: Some(json!("coding")),
},
RoutingCondition::Predicate {
field: "body.reasoning_effort".to_string(),
op: RoutingConditionOp::In,
value: Some(json!(["high", "xhigh"])),
},
],
};
let headers = json!({"x-app":"coding"});
assert!(condition.matches(&context(&headers, &json!({"reasoning_effort":"high"}))));
assert!(!condition.matches(&context(&headers, &json!({"reasoning_effort":"low"}))));
}
}

36
crates/aether-routing-core/src/lib.rs Normal file
View File

@@ -0,0 +1,36 @@
mod actions;
mod conditions;
mod model;
mod mutations;
mod policy;
mod ranking;
mod trace;
mod validation;
pub use actions::{
RoutingAction, RoutingHeaderPatch, RoutingJsonPatchOperation, RoutingRulePhase,
RoutingSchedulingMode, RoutingSetPriorityMode,
};
pub use conditions::{RoutingCondition, RoutingConditionContext, RoutingConditionOp};
pub use model::{
RoutingGroupBinding, RoutingGroupBindingSubject, RoutingGroupConfig, RoutingGroupRecord,
RoutingGroupVersionRecord, RoutingModelPolicy, RoutingPoolPolicyOverride, RoutingRule,
RoutingSchedulingPreset,
};
pub use mutations::{
apply_json_patch_operations, validate_header_patch, validate_json_patch_operations,
HeaderMutation, MutationError, MutationPlan,
};
pub use policy::{
resolve_routing_policy, MatchedRoutingRule, ResolvedRoutingPolicy, RoutingPolicyError,
RoutingPolicyInput,
};
pub use ranking::{
rank_vector_for_candidate, CandidateKind, RankingOverlay, RoutingCandidateFacts,
RoutingCandidateRankVector, ROUTING_PRIORITY_UNSPECIFIED,
};
pub use trace::{
RoutingCandidateTrace, RoutingDecisionTrace, RoutingPatchSummary, RoutingPoolExpansionTrace,
RoutingRuntimeFacts,
};
pub use validation::{validate_routing_group_config, RoutingValidationError};

131
crates/aether-routing-core/src/model.rs Normal file
View File

@@ -0,0 +1,131 @@
use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::actions::{
RoutingAction, RoutingRulePhase, RoutingSchedulingMode, RoutingSetPriorityMode,
};
use crate::conditions::RoutingCondition;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RoutingSchedulingPreset {
pub preset: String,
#[serde(default = "default_true")]
pub enabled: bool,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub mode: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
pub struct RoutingPoolPolicyOverride {
#[serde(default)]
pub scheduling_presets: Vec<RoutingSchedulingPreset>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
pub struct RoutingDefaultPolicy {
#[serde(default)]
pub priority_mode: RoutingSetPriorityMode,
#[serde(default)]
pub scheduling_mode: RoutingSchedulingMode,
#[serde(default)]
pub keep_priority_on_conversion: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
pub struct RoutingModelPolicy {
pub model: String,
#[serde(default)]
pub allowed_providers: Vec<String>,
#[serde(default)]
pub allowed_keys: Vec<String>,
#[serde(default)]
pub provider_priority_overrides: BTreeMap<String, i32>,
#[serde(default)]
pub key_priority_overrides: BTreeMap<String, i32>,
#[serde(default)]
pub pool_priority_overrides: BTreeMap<String, i32>,
#[serde(default)]
pub pool_policy_overrides: BTreeMap<String, RoutingPoolPolicyOverride>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RoutingRule {
pub id: String,
#[serde(default)]
pub priority: i32,
#[serde(default = "default_true")]
pub enabled: bool,
#[serde(default)]
pub phase: RoutingRulePhase,
#[serde(default)]
pub conditions: RoutingCondition,
#[serde(default)]
pub actions: Vec<RoutingAction>,
#[serde(default)]
pub stop_processing: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
pub struct RoutingGroupConfig {
#[serde(default)]
pub allowed_models: Vec<String>,
#[serde(default)]
pub default_policy: RoutingDefaultPolicy,
#[serde(default)]
pub model_policies: Vec<RoutingModelPolicy>,
#[serde(default)]
pub rules: Vec<RoutingRule>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RoutingGroupRecord {
pub id: String,
pub name: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub description: Option<String>,
pub enabled: bool,
pub is_system_default: bool,
pub config_json: Value,
pub version: i64,
pub created_at: i64,
pub updated_at: i64,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub published_at: Option<i64>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RoutingGroupBindingSubject {
User,
ApiKey,
UserGroup,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RoutingGroupBinding {
pub id: String,
pub group_id: String,
pub subject_type: RoutingGroupBindingSubject,
pub subject_id: String,
pub is_default: bool,
pub allow_explicit_select: bool,
pub created_at: i64,
pub updated_at: i64,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RoutingGroupVersionRecord {
pub id: String,
pub group_id: String,
pub version: i64,
pub config_json: Value,
pub created_at: i64,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub created_by: Option<String>,
}
fn default_true() -> bool {
true
}

236
crates/aether-routing-core/src/mutations.rs Normal file
View File

@@ -0,0 +1,236 @@
use std::collections::BTreeSet;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use thiserror::Error;
use crate::actions::{RoutingHeaderPatch, RoutingJsonPatchOperation};
const RESERVED_HEADERS: &[&str] = &[
"authorization",
"x-api-key",
"api-key",
"cookie",
"set-cookie",
"x-aether-trace-id",
"x-aether-internal",
"x-aether-scheduler-group",
];
#[derive(Debug, Error, Clone, PartialEq, Eq)]
pub enum MutationError {
#[error("json patch path must be an absolute JSON pointer: {0}")]
InvalidJsonPointer(String),
#[error("json patch cannot target reserved path: {0}")]
ReservedJsonPath(String),
#[error("json patch target does not exist: {0}")]
MissingTarget(String),
#[error("json patch parent is not an object: {0}")]
InvalidParent(String),
#[error("header patch targets reserved header: {0}")]
ReservedHeader(String),
#[error("header patch has invalid header name: {0}")]
InvalidHeaderName(String),
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct HeaderMutation {
pub set: Vec<(String, String)>,
pub remove: Vec<String>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct MutationPlan {
pub body_patch: Vec<RoutingJsonPatchOperation>,
pub header_patch: Vec<RoutingHeaderPatch>,
}
impl MutationPlan {
pub fn is_empty(&self) -> bool {
self.body_patch.is_empty() && self.header_patch.is_empty()
}
}
pub fn validate_json_patch_operations(
operations: &[RoutingJsonPatchOperation],
) -> Result<(), MutationError> {
for operation in operations {
let path = operation.path();
validate_json_pointer(path)?;
if is_reserved_json_path(path) {
return Err(MutationError::ReservedJsonPath(path.to_string()));
}
}
Ok(())
}
pub fn apply_json_patch_operations(
body: &mut Value,
operations: &[RoutingJsonPatchOperation],
) -> Result<(), MutationError> {
validate_json_patch_operations(operations)?;
for operation in operations {
match operation {
RoutingJsonPatchOperation::Add { path, value } => {
set_json_pointer(body, path, value.clone(), true)?;
}
RoutingJsonPatchOperation::Replace { path, value } => {
set_json_pointer(body, path, value.clone(), false)?;
}
RoutingJsonPatchOperation::Remove { path } => {
remove_json_pointer(body, path)?;
}
}
}
Ok(())
}
pub fn validate_header_patch(patch: &[RoutingHeaderPatch]) -> Result<(), MutationError> {
let reserved = RESERVED_HEADERS.iter().copied().collect::<BTreeSet<_>>();
for item in patch {
let name = item.name().trim().to_ascii_lowercase();
if name.is_empty()
|| name
.chars()
.any(|ch| !(ch.is_ascii_alphanumeric() || ch == '-'))
{
return Err(MutationError::InvalidHeaderName(item.name().to_string()));
}
if reserved.contains(name.as_str()) {
return Err(MutationError::ReservedHeader(item.name().to_string()));
}
}
Ok(())
}
fn validate_json_pointer(path: &str) -> Result<(), MutationError> {
if !path.starts_with('/') {
return Err(MutationError::InvalidJsonPointer(path.to_string()));
}
Ok(())
}
fn is_reserved_json_path(path: &str) -> bool {
matches!(
path,
"/authorization"
| "/api_key"
| "/provider_secret"
| "/upstream_url"
| "/upstream_base_url"
| "/auth"
)
}
fn set_json_pointer(
root: &mut Value,
pointer: &str,
value: Value,
allow_create: bool,
) -> Result<(), MutationError> {
let tokens = pointer_tokens(pointer);
if tokens.is_empty() {
*root = value;
return Ok(());
}
let (parents, leaf) = tokens.split_at(tokens.len() - 1);
let parent = descend_mut(root, parents, pointer)?;
match parent {
Value::Object(map) => {
if !allow_create && !map.contains_key(&leaf[0]) {
return Err(MutationError::MissingTarget(pointer.to_string()));
}
map.insert(leaf[0].clone(), value);
Ok(())
}
_ => Err(MutationError::InvalidParent(pointer.to_string())),
}
}
fn remove_json_pointer(root: &mut Value, pointer: &str) -> Result<(), MutationError> {
let tokens = pointer_tokens(pointer);
if tokens.is_empty() {
*root = Value::Null;
return Ok(());
}
let (parents, leaf) = tokens.split_at(tokens.len() - 1);
let parent = descend_mut(root, parents, pointer)?;
match parent {
Value::Object(map) => map
.remove(&leaf[0])
.map(|_| ())
.ok_or_else(|| MutationError::MissingTarget(pointer.to_string())),
_ => Err(MutationError::InvalidParent(pointer.to_string())),
}
}
fn descend_mut<'a>(
root: &'a mut Value,
tokens: &[String],
pointer: &str,
) -> Result<&'a mut Value, MutationError> {
let mut current = root;
for token in tokens {
match current {
Value::Object(map) => {
current = map
.get_mut(token)
.ok_or_else(|| MutationError::MissingTarget(pointer.to_string()))?;
}
Value::Null => {
*current = Value::Object(Map::new());
if let Value::Object(map) = current {
current = map
.entry(token.clone())
.or_insert_with(|| Value::Object(Map::new()));
}
}
_ => return Err(MutationError::InvalidParent(pointer.to_string())),
}
}
Ok(current)
}
fn pointer_tokens(pointer: &str) -> Vec<String> {
pointer
.trim_start_matches('/')
.split('/')
.filter(|part| !part.is_empty())
.map(|part| part.replace("~1", "/").replace("~0", "~"))
.collect()
}
#[cfg(test)]
mod tests {
use serde_json::json;
use crate::actions::RoutingJsonPatchOperation;
use super::*;
#[test]
fn applies_body_patch() {
let mut body = json!({"metadata":{}});
apply_json_patch_operations(
&mut body,
&[RoutingJsonPatchOperation::Add {
path: "/metadata/routing".to_string(),
value: json!("high"),
}],
)
.expect("patch should apply");
assert_eq!(body["metadata"]["routing"], json!("high"));
}
#[test]
fn rejects_reserved_headers() {
assert_eq!(
validate_header_patch(&[RoutingHeaderPatch::Set {
name: "authorization".to_string(),
value: "secret".to_string()
}]),
Err(MutationError::ReservedHeader("authorization".to_string()))
);
}
}

423
crates/aether-routing-core/src/policy.rs Normal file
View File

@@ -0,0 +1,423 @@
use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use thiserror::Error;
use crate::actions::{
RoutingAction, RoutingRulePhase, RoutingSchedulingMode, RoutingSetPriorityMode,
};
use crate::conditions::RoutingConditionContext;
use crate::model::{RoutingGroupConfig, RoutingModelPolicy, RoutingPoolPolicyOverride};
use crate::mutations::{validate_header_patch, validate_json_patch_operations, MutationPlan};
use crate::ranking::RankingOverlay;
#[derive(Debug, Error, Clone, PartialEq, Eq)]
pub enum RoutingPolicyError {
#[error("routing group config is invalid: {0}")]
InvalidConfig(String),
#[error("model is not allowed by routing group: {0}")]
ModelNotAllowed(String),
#[error("mutation action is invalid: {0}")]
InvalidMutation(String),
}
#[derive(Debug, Clone)]
pub struct RoutingPolicyInput<'a> {
pub group_id: Option<&'a str>,
pub group_version: Option<i64>,
pub selection_source: &'a str,
pub requested_model: &'a str,
pub resolved_model: &'a str,
pub api_format: &'a str,
pub user_id: Option<&'a str>,
pub api_key_id: Option<&'a str>,
pub headers: &'a Value,
pub body: &'a Value,
pub phase: RoutingRulePhase,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct MatchedRoutingRule {
pub id: String,
pub priority: i32,
pub phase: RoutingRulePhase,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ResolvedRoutingPolicy {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub group_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub group_version: Option<i64>,
pub selection_source: String,
pub requested_model: String,
pub resolved_model: String,
pub priority_mode: RoutingSetPriorityMode,
pub scheduling_mode: RoutingSchedulingMode,
pub keep_priority_on_conversion: bool,
pub ranking_overlay: RankingOverlay,
pub mutation_plan: MutationPlan,
#[serde(default)]
pub pool_policy_overrides: BTreeMap<String, RoutingPoolPolicyOverride>,
#[serde(default)]
pub matched_rules: Vec<MatchedRoutingRule>,
}
pub fn resolve_routing_policy(
config: &RoutingGroupConfig,
input: RoutingPolicyInput<'_>,
) -> Result<ResolvedRoutingPolicy, RoutingPolicyError> {
if !model_allowed(&config.allowed_models, input.requested_model)
&& !model_allowed(&config.allowed_models, input.resolved_model)
{
return Err(RoutingPolicyError::ModelNotAllowed(
input.requested_model.to_string(),
));
}
let mut policy = ResolvedRoutingPolicy {
group_id: input.group_id.map(str::to_string),
group_version: input.group_version,
selection_source: input.selection_source.to_string(),
requested_model: input.requested_model.to_string(),
resolved_model: input.resolved_model.to_string(),
priority_mode: config.default_policy.priority_mode,
scheduling_mode: config.default_policy.scheduling_mode,
keep_priority_on_conversion: config.default_policy.keep_priority_on_conversion,
ranking_overlay: RankingOverlay::default(),
mutation_plan: MutationPlan::default(),
pool_policy_overrides: BTreeMap::new(),
matched_rules: Vec::new(),
};
for model_policy in matching_model_policies(config, input.requested_model, input.resolved_model)
{
apply_model_policy(&mut policy, model_policy);
}
let condition_context = RoutingConditionContext {
model: input.requested_model,
api_format: input.api_format,
user_id: input.user_id,
api_key_id: input.api_key_id,
headers: input.headers,
body: input.body,
};
let mut rules = config
.rules
.iter()
.filter(|rule| rule.enabled && rule.phase == input.phase)
.collect::<Vec<_>>();
rules.sort_by(|left, right| {
left.priority
.cmp(&right.priority)
.then(left.id.cmp(&right.id))
});
for rule in rules {
if !rule.conditions.matches(&condition_context) {
continue;
}
for action in &rule.actions {
apply_action(
&mut policy,
action,
input.requested_model,
input.resolved_model,
)?;
}
policy.matched_rules.push(MatchedRoutingRule {
id: rule.id.clone(),
priority: rule.priority,
phase: rule.phase,
});
if rule.stop_processing {
break;
}
}
Ok(policy)
}
fn apply_model_policy(policy: &mut ResolvedRoutingPolicy, model_policy: &RoutingModelPolicy) {
if !model_policy.allowed_providers.is_empty() {
policy.ranking_overlay.allowed_providers = model_policy.allowed_providers.clone();
}
if !model_policy.allowed_keys.is_empty() {
policy.ranking_overlay.allowed_keys = model_policy.allowed_keys.clone();
}
policy.ranking_overlay.provider_priority_overrides.extend(
model_policy
.provider_priority_overrides
.iter()
.map(|(key, value)| (key.clone(), *value)),
);
policy.ranking_overlay.key_priority_overrides.extend(
model_policy
.key_priority_overrides
.iter()
.map(|(key, value)| (key.clone(), *value)),
);
policy.ranking_overlay.pool_priority_overrides.extend(
model_policy
.pool_priority_overrides
.iter()
.map(|(key, value)| (key.clone(), *value)),
);
policy
.pool_policy_overrides
.extend(model_policy.pool_policy_overrides.clone());
}
fn apply_action(
policy: &mut ResolvedRoutingPolicy,
action: &RoutingAction,
requested_model: &str,
resolved_model: &str,
) -> Result<(), RoutingPolicyError> {
match action {
RoutingAction::RestrictModels { models } => {
if !model_allowed(models, requested_model) && !model_allowed(models, resolved_model) {
return Err(RoutingPolicyError::ModelNotAllowed(
requested_model.to_string(),
));
}
}
RoutingAction::RestrictProviders { provider_ids } => {
policy.ranking_overlay.allowed_providers = provider_ids.clone();
}
RoutingAction::RestrictKeys { key_ids } => {
policy.ranking_overlay.allowed_keys = key_ids.clone();
}
RoutingAction::SetScheduling {
priority_mode,
scheduling_mode,
keep_priority_on_conversion,
} => {
if let Some(priority_mode) = priority_mode {
policy.priority_mode = *priority_mode;
}
if let Some(scheduling_mode) = scheduling_mode {
policy.scheduling_mode = *scheduling_mode;
}
if let Some(keep_priority_on_conversion) = keep_priority_on_conversion {
policy.keep_priority_on_conversion = *keep_priority_on_conversion;
}
}
RoutingAction::SetProviderPriority {
provider_id,
priority,
} => {
policy
.ranking_overlay
.provider_priority_overrides
.insert(provider_id.clone(), *priority);
}
RoutingAction::SetKeyPriority { key_id, priority } => {
policy
.ranking_overlay
.key_priority_overrides
.insert(key_id.clone(), *priority);
}
RoutingAction::JsonPatchBody { patch } => {
validate_json_patch_operations(patch)
.map_err(|error| RoutingPolicyError::InvalidMutation(error.to_string()))?;
policy.mutation_plan.body_patch.extend(patch.clone());
}
RoutingAction::PatchHeaders { patch } => {
validate_header_patch(patch)
.map_err(|error| RoutingPolicyError::InvalidMutation(error.to_string()))?;
policy.mutation_plan.header_patch.extend(patch.clone());
}
}
Ok(())
}
fn matching_model_policies<'a>(
config: &'a RoutingGroupConfig,
requested_model: &str,
resolved_model: &str,
) -> Vec<&'a RoutingModelPolicy> {
config
.model_policies
.iter()
.filter(|policy| {
model_pattern_matches(&policy.model, requested_model)
|| model_pattern_matches(&policy.model, resolved_model)
})
.collect()
}
fn model_allowed(patterns: &[String], requested_model: &str) -> bool {
patterns.is_empty()
|| patterns
.iter()
.any(|pattern| model_pattern_matches(pattern, requested_model))
}
fn model_pattern_matches(pattern: &str, value: &str) -> bool {
let pattern = pattern.trim();
if pattern == "*" {
return true;
}
if let Some(prefix) = pattern.strip_suffix('*') {
return value.starts_with(prefix);
}
pattern == value
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use serde_json::json;
use crate::actions::{RoutingJsonPatchOperation, RoutingRulePhase};
use crate::conditions::{RoutingCondition, RoutingConditionOp};
use crate::model::{RoutingDefaultPolicy, RoutingRule};
use super::*;
#[test]
fn resolves_model_policy_and_matching_rule() {
let config = RoutingGroupConfig {
allowed_models: vec!["gpt-*".to_string()],
default_policy: RoutingDefaultPolicy::default(),
model_policies: vec![RoutingModelPolicy {
model: "gpt-5".to_string(),
allowed_providers: vec!["provider-a".to_string()],
provider_priority_overrides: BTreeMap::from([("provider-a".to_string(), 0)]),
pool_priority_overrides: BTreeMap::from([("provider-a".to_string(), 3)]),
..RoutingModelPolicy::default()
}],
rules: vec![RoutingRule {
id: "high".to_string(),
priority: 10,
enabled: true,
phase: RoutingRulePhase::ClientRequest,
conditions: RoutingCondition::Predicate {
field: "body.reasoning_effort".to_string(),
op: RoutingConditionOp::Eq,
value: Some(json!("high")),
},
actions: vec![RoutingAction::JsonPatchBody {
patch: vec![RoutingJsonPatchOperation::Add {
path: "/metadata/routing".to_string(),
value: json!("high"),
}],
}],
stop_processing: false,
}],
};
let policy = resolve_routing_policy(
&config,
RoutingPolicyInput {
group_id: Some("group-1"),
group_version: Some(1),
selection_source: "explicit",
requested_model: "gpt-5",
resolved_model: "gpt-5",
api_format: "openai:chat",
user_id: Some("user-1"),
api_key_id: Some("api-key-1"),
headers: &json!({}),
body: &json!({"reasoning_effort":"high"}),
phase: RoutingRulePhase::ClientRequest,
},
)
.expect("policy should resolve");
assert_eq!(policy.ranking_overlay.allowed_providers, vec!["provider-a"]);
assert_eq!(
policy
.ranking_overlay
.provider_priority_overrides
.get("provider-a"),
Some(&0)
);
assert_eq!(
policy
.ranking_overlay
.pool_priority_overrides
.get("provider-a"),
Some(&3)
);
assert_eq!(policy.matched_rules.len(), 1);
assert_eq!(policy.mutation_plan.body_patch.len(), 1);
}
#[test]
fn rejects_disallowed_model() {
let config = RoutingGroupConfig {
allowed_models: vec!["gpt-5".to_string()],
..RoutingGroupConfig::default()
};
let err = resolve_routing_policy(
&config,
RoutingPolicyInput {
group_id: None,
group_version: None,
selection_source: "test",
requested_model: "claude",
resolved_model: "claude",
api_format: "openai:chat",
user_id: None,
api_key_id: None,
headers: &json!({}),
body: &json!({}),
phase: RoutingRulePhase::ClientRequest,
},
)
.unwrap_err();
assert_eq!(
err,
RoutingPolicyError::ModelNotAllowed("claude".to_string())
);
}
#[test]
fn restrict_model_action_rejects_matching_request() {
let config = RoutingGroupConfig {
allowed_models: vec!["*".to_string()],
rules: vec![RoutingRule {
id: "restrict".to_string(),
priority: 1,
enabled: true,
phase: RoutingRulePhase::ClientRequest,
conditions: RoutingCondition::default(),
actions: vec![RoutingAction::RestrictModels {
models: vec!["gpt-5".to_string()],
}],
stop_processing: false,
}],
..RoutingGroupConfig::default()
};
let err = resolve_routing_policy(
&config,
RoutingPolicyInput {
group_id: None,
group_version: None,
selection_source: "test",
requested_model: "claude",
resolved_model: "claude",
api_format: "openai:chat",
user_id: None,
api_key_id: None,
headers: &json!({}),
body: &json!({}),
phase: RoutingRulePhase::ClientRequest,
},
)
.unwrap_err();
assert_eq!(
err,
RoutingPolicyError::ModelNotAllowed("claude".to_string())
);
}
}

181
crates/aether-routing-core/src/ranking.rs Normal file
View File

@@ -0,0 +1,181 @@
use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
pub const ROUTING_PRIORITY_UNSPECIFIED: i32 = i32::MAX;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum CandidateKind {
Provider,
PoolGroup,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct RankingOverlay {
#[serde(default)]
pub allowed_providers: Vec<String>,
#[serde(default)]
pub allowed_keys: Vec<String>,
#[serde(default)]
pub provider_priority_overrides: BTreeMap<String, i32>,
#[serde(default)]
pub key_priority_overrides: BTreeMap<String, i32>,
#[serde(default)]
pub pool_priority_overrides: BTreeMap<String, i32>,
}
impl RankingOverlay {
pub fn provider_priority(&self, provider_id: &str, fallback: i32) -> i32 {
self.provider_priority_overrides
.get(provider_id)
.copied()
.unwrap_or(fallback)
}
pub fn key_priority(&self, key_id: &str, fallback: i32) -> i32 {
self.key_priority_overrides
.get(key_id)
.copied()
.unwrap_or(fallback)
}
pub fn provider_priority_or_unspecified(&self, provider_id: &str) -> i32 {
self.provider_priority_overrides
.get(provider_id)
.copied()
.unwrap_or(ROUTING_PRIORITY_UNSPECIFIED)
}
pub fn key_priority_or_unspecified(&self, key_id: &str) -> i32 {
self.key_priority_overrides
.get(key_id)
.copied()
.unwrap_or(ROUTING_PRIORITY_UNSPECIFIED)
}
pub fn pool_priority_or_unspecified(&self, provider_id: &str) -> i32 {
self.pool_priority_overrides
.get(provider_id)
.copied()
.unwrap_or(ROUTING_PRIORITY_UNSPECIFIED)
}
pub fn provider_allowed(&self, provider_id: &str) -> bool {
self.allowed_providers.is_empty()
|| self
.allowed_providers
.iter()
.any(|item| item == provider_id)
}
pub fn key_allowed(&self, key_id: &str) -> bool {
self.allowed_keys.is_empty() || self.allowed_keys.iter().any(|item| item == key_id)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RoutingCandidateFacts {
pub candidate_kind: CandidateKind,
pub provider_id: String,
pub endpoint_id: String,
pub model_id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub key_id: Option<String>,
pub provider_priority: i32,
pub key_priority: i32,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RoutingCandidateRankVector {
pub provider_priority_before: i32,
pub provider_priority_after: i32,
pub key_priority_before: i32,
pub key_priority_after: i32,
}
pub fn rank_vector_for_candidate(
overlay: &RankingOverlay,
facts: &RoutingCandidateFacts,
) -> RoutingCandidateRankVector {
RoutingCandidateRankVector {
provider_priority_before: facts.provider_priority,
provider_priority_after: overlay.provider_priority_or_unspecified(&facts.provider_id),
key_priority_before: facts.key_priority,
key_priority_after: match facts.candidate_kind {
CandidateKind::Provider => facts
.key_id
.as_deref()
.map(|key_id| overlay.key_priority_or_unspecified(key_id))
.unwrap_or(ROUTING_PRIORITY_UNSPECIFIED),
CandidateKind::PoolGroup => overlay.pool_priority_or_unspecified(&facts.provider_id),
},
}
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use super::*;
#[test]
fn overlay_applies_provider_and_key_priority() {
let overlay = RankingOverlay {
provider_priority_overrides: BTreeMap::from([("provider-a".to_string(), 2)]),
key_priority_overrides: BTreeMap::from([("key-a".to_string(), 5)]),
..RankingOverlay::default()
};
let facts = RoutingCandidateFacts {
candidate_kind: CandidateKind::Provider,
provider_id: "provider-a".to_string(),
endpoint_id: "endpoint-a".to_string(),
model_id: "model-a".to_string(),
key_id: Some("key-a".to_string()),
provider_priority: 10,
key_priority: 20,
};
let vector = rank_vector_for_candidate(&overlay, &facts);
assert_eq!(vector.provider_priority_after, 2);
assert_eq!(vector.key_priority_after, 5);
}
#[test]
fn rank_vector_marks_missing_routing_priorities_unspecified() {
let facts = RoutingCandidateFacts {
candidate_kind: CandidateKind::Provider,
provider_id: "provider-a".to_string(),
endpoint_id: "endpoint-a".to_string(),
model_id: "model-a".to_string(),
key_id: Some("key-a".to_string()),
provider_priority: 10,
key_priority: 20,
};
let vector = rank_vector_for_candidate(&RankingOverlay::default(), &facts);
assert_eq!(vector.provider_priority_after, ROUTING_PRIORITY_UNSPECIFIED);
assert_eq!(vector.key_priority_after, ROUTING_PRIORITY_UNSPECIFIED);
}
#[test]
fn rank_vector_uses_pool_priority_for_pool_groups() {
let overlay = RankingOverlay {
pool_priority_overrides: BTreeMap::from([("provider-a".to_string(), 4)]),
..RankingOverlay::default()
};
let facts = RoutingCandidateFacts {
candidate_kind: CandidateKind::PoolGroup,
provider_id: "provider-a".to_string(),
endpoint_id: "endpoint-a".to_string(),
model_id: "model-a".to_string(),
key_id: None,
provider_priority: 10,
key_priority: 20,
};
let vector = rank_vector_for_candidate(&overlay, &facts);
assert_eq!(vector.key_priority_after, 4);
}
}

86
crates/aether-routing-core/src/trace.rs Normal file
View File

@@ -0,0 +1,86 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::actions::{RoutingSchedulingMode, RoutingSetPriorityMode};
use crate::ranking::{CandidateKind, RoutingCandidateRankVector};
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct RoutingPatchSummary {
#[serde(default)]
pub body_paths: Vec<String>,
#[serde(default)]
pub header_names: Vec<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub failed_action: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RoutingCandidateTrace {
pub candidate_kind: CandidateKind,
pub provider_id: String,
pub endpoint_id: String,
pub model_id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub key_id: Option<String>,
pub ranking_vector: RoutingCandidateRankVector,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub skip_reason: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub selected_order: Option<u32>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RoutingPoolExpansionTrace {
pub pool_group_id: String,
pub key_id: String,
#[serde(default)]
pub pool_ranking_vector: Vec<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub pool_skip_reason: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub selected_order: Option<u32>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct RoutingRuntimeFacts {
#[serde(default)]
pub cache_affinity_hit: bool,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub sticky_key_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub load_balance_seed: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub scheduler_mode: Option<RoutingSchedulingMode>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub priority_mode: Option<RoutingSetPriorityMode>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct RoutingDecisionTrace {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub group_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub group_version: Option<i64>,
pub selection_source: String,
#[serde(default)]
pub selected_rules: Vec<String>,
pub original_model: String,
pub resolved_model: String,
pub client_api_format: String,
#[serde(default)]
pub client_request_patch_summary: RoutingPatchSummary,
#[serde(default)]
pub provider_request_patch_summary: RoutingPatchSummary,
#[serde(default)]
pub global_candidates: Vec<RoutingCandidateTrace>,
#[serde(default)]
pub pool_expansion: Vec<RoutingPoolExpansionTrace>,
#[serde(default)]
pub runtime_facts: RoutingRuntimeFacts,
}
impl RoutingDecisionTrace {
pub fn to_extra_data_value(&self) -> Value {
serde_json::json!({ "routing_trace": self })
}
}

54
crates/aether-routing-core/src/validation.rs Normal file
View File

@@ -0,0 +1,54 @@
use std::collections::BTreeSet;
use thiserror::Error;
use crate::model::RoutingGroupConfig;
use crate::mutations::{validate_header_patch, validate_json_patch_operations};
use crate::RoutingAction;
#[derive(Debug, Error, Clone, PartialEq, Eq)]
pub enum RoutingValidationError {
#[error("routing rule id is empty")]
EmptyRuleId,
#[error("duplicate routing rule id: {0}")]
DuplicateRuleId(String),
#[error("routing model policy selector is empty")]
EmptyModelSelector,
#[error("invalid mutation action: {0}")]
InvalidMutation(String),
}
pub fn validate_routing_group_config(
config: &RoutingGroupConfig,
) -> Result<(), RoutingValidationError> {
let mut rule_ids = BTreeSet::new();
for model_policy in &config.model_policies {
if model_policy.model.trim().is_empty() {
return Err(RoutingValidationError::EmptyModelSelector);
}
}
for rule in &config.rules {
if rule.id.trim().is_empty() {
return Err(RoutingValidationError::EmptyRuleId);
}
if !rule_ids.insert(rule.id.clone()) {
return Err(RoutingValidationError::DuplicateRuleId(rule.id.clone()));
}
for action in &rule.actions {
match action {
RoutingAction::JsonPatchBody { patch } => {
validate_json_patch_operations(patch).map_err(|error| {
RoutingValidationError::InvalidMutation(error.to_string())
})?;
}
RoutingAction::PatchHeaders { patch } => {
validate_header_patch(patch).map_err(|error| {
RoutingValidationError::InvalidMutation(error.to_string())
})?;
}
_ => {}
}
}
}
Ok(())
}

32
crates/aether-scheduler-core/src/request_candidate.rs
View File

@@ -36,6 +36,7 @@ pub struct SchedulerRequestCandidateReportContext {
pub priority_slot: Option<i32>,
pub promoted_by: Option<String>,
pub demoted_by: Option<String>,
pub routing_trace: Option<Value>,
}
#[derive(Debug, Clone, PartialEq)]
@@ -83,6 +84,7 @@ struct ReportCandidateExtraDataInput {
priority_slot: Option<i32>,
promoted_by: Option<String>,
demoted_by: Option<String>,
routing_trace: Option<Value>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -180,6 +182,10 @@ pub fn parse_request_candidate_report_context(
priority_slot: i32_field(report_context, "priority_slot"),
promoted_by: string_field(report_context, "promoted_by"),
demoted_by: string_field(report_context, "demoted_by"),
routing_trace: report_context
.get("routing_trace")
.cloned()
.filter(|value| !value.is_null()),
})
}
@@ -221,6 +227,7 @@ pub fn resolve_report_request_candidate_slot(
priority_slot,
promoted_by,
demoted_by,
routing_trace,
} = metadata;
let request_id = request_id?;
let synthesized_extra_data = build_report_candidate_extra_data(ReportCandidateExtraDataInput {
@@ -245,6 +252,7 @@ pub fn resolve_report_request_candidate_slot(
priority_slot,
promoted_by,
demoted_by,
routing_trace,
});
let created_at_unix_ms = matched_candidate
.as_ref()
@@ -368,6 +376,7 @@ pub fn build_execution_request_candidate_seed(
priority_slot: metadata.priority_slot,
promoted_by: metadata.promoted_by,
demoted_by: metadata.demoted_by,
routing_trace: metadata.routing_trace,
})
});
append_seed_extra_data_from_report_context(&mut extra_data, &context);
@@ -477,6 +486,7 @@ pub fn build_local_request_candidate_status_record(
priority_slot: metadata.priority_slot,
promoted_by: metadata.promoted_by.clone(),
demoted_by: metadata.demoted_by.clone(),
routing_trace: metadata.routing_trace.clone(),
});
let created_at_unix_ms = started_at_unix_ms.or(finished_at_unix_ms);
@@ -717,6 +727,7 @@ fn build_report_candidate_extra_data(input: ReportCandidateExtraDataInput) -> Op
priority_slot,
promoted_by,
demoted_by,
routing_trace,
} = input;
let mut extra_data = Map::with_capacity(8);
extra_data.insert("gateway_execution_runtime".to_string(), Value::Bool(true));
@@ -812,6 +823,9 @@ fn build_report_candidate_extra_data(input: ReportCandidateExtraDataInput) -> Op
if let Some(demoted_by) = demoted_by {
extra_data.insert("demoted_by".to_string(), Value::String(demoted_by));
}
if let Some(routing_trace) = routing_trace {
extra_data.insert("routing_trace".to_string(), routing_trace);
}
(!extra_data.is_empty()).then_some(Value::Object(extra_data))
}
@@ -1138,7 +1152,15 @@ mod tests {
"ranking_index": 2,
"priority_slot": 7,
"promoted_by": "cached_affinity",
"demoted_by": "cross_format"
"demoted_by": "cross_format",
"routing_trace": {
"group_id": "routing-group-1",
"pool_expansion": [{
"pool_group_id": "pool-1",
"key_id": "key-1",
"selected_order": 0
}]
}
})),
status_update: SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
@@ -1228,6 +1250,14 @@ mod tests {
.and_then(|value| value.get("demoted_by")),
Some(&json!("cross_format"))
);
assert_eq!(
record
.extra_data
.as_ref()
.and_then(|value| value.get("routing_trace"))
.and_then(|value| value.get("group_id")),
Some(&json!("routing-group-1"))
);
}
#[test]