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

refactor: 移除 shadow results 相关模块和接口

Browse Source
This commit is contained in:
fawney19
2026-04-14 09:46:16 +08:00
parent 5fd3240fcf
commit 62e0a0338d
30 changed files with 115 additions and 1991 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
crates/aether-data/src/backends/mod.rs
View File

@@ -141,10 +141,8 @@ mod tests {
assert!(backends.read().provider_catalog().is_none());
assert!(backends.read().usage().is_none());
assert!(backends.read().video_tasks().is_none());
assert!(backends.read().shadow_results().is_none());
assert!(backends.transactions().postgres().is_none());
assert!(backends.workers().redis().is_none());
assert!(backends.write().shadow_results().is_none());
assert!(backends.write().settlement().is_none());
assert!(backends.write().usage().is_none());
}
@@ -184,9 +182,7 @@ mod tests {
assert!(backends.read().usage().is_some());
assert!(backends.read().video_tasks().is_some());
assert!(backends.read().wallets().is_some());
assert!(backends.read().shadow_results().is_some());
assert!(backends.transactions().postgres().is_some());
assert!(backends.write().shadow_results().is_some());
assert!(backends.write().auth_modules().is_some());
assert!(backends.write().gemini_file_mappings().is_some());
assert!(backends.write().management_tokens().is_some());
@@ -222,7 +218,6 @@ mod tests {
assert!(backends.read().global_models().is_none());
assert!(backends.read().oauth_providers().is_none());
assert!(backends.transactions().postgres().is_none());
assert!(backends.write().shadow_results().is_none());
assert!(backends.write().settlement().is_none());
assert!(backends.write().usage().is_none());
assert!(backends.config().redis.is_some());

13
crates/aether-data/src/backends/postgres.rs
View File

@@ -47,9 +47,6 @@ use crate::repository::quota::{
ProviderQuotaReadRepository, ProviderQuotaWriteRepository, SqlxProviderQuotaRepository,
};
use crate::repository::settlement::{SettlementWriteRepository, SqlxSettlementRepository};
use crate::repository::shadow_results::{
ShadowResultReadRepository, ShadowResultWriteRepository, SqlxShadowResultRepository,
};
use crate::repository::system::{AdminSystemStats, StoredSystemConfigEntry};
use crate::repository::usage::{
SqlxUsageReadRepository, UsageReadRepository, UsageWriteRepository,
@@ -287,18 +284,10 @@ impl PostgresBackend {
PostgresLeaseRunner::new(self.transaction_runner(), config)
}
pub fn shadow_result_write_repository(&self) -> Arc<dyn ShadowResultWriteRepository> {
Arc::new(SqlxShadowResultRepository::new(self.pool_clone()))
}
pub fn provider_quota_write_repository(&self) -> Arc<dyn ProviderQuotaWriteRepository> {
Arc::new(SqlxProviderQuotaRepository::new(self.pool_clone()))
}
pub fn shadow_result_read_repository(&self) -> Arc<dyn ShadowResultReadRepository> {
Arc::new(SqlxShadowResultRepository::new(self.pool_clone()))
}
pub async fn find_system_config_value(
&self,
key: &str,
@@ -465,8 +454,6 @@ mod tests {
let _lease_runner = backend
.lease_runner(PostgresLeaseRunnerConfig::default())
.expect("lease runner should build");
let _shadow_result_reader = backend.shadow_result_read_repository();
let _shadow_result_writer = backend.shadow_result_write_repository();
let _provider_quota_writer = backend.provider_quota_write_repository();
}
}

10
crates/aether-data/src/backends/read.rs
View File

@@ -15,7 +15,6 @@ use crate::repository::oauth_providers::OAuthProviderReadRepository;
use crate::repository::provider_catalog::ProviderCatalogReadRepository;
use crate::repository::proxy_nodes::ProxyNodeReadRepository;
use crate::repository::quota::ProviderQuotaReadRepository;
use crate::repository::shadow_results::ShadowResultReadRepository;
use crate::repository::usage::UsageReadRepository;
use crate::repository::users::UserReadRepository;
use crate::repository::video_tasks::VideoTaskReadRepository;
@@ -40,7 +39,6 @@ pub struct DataReadRepositories {
users: Option<Arc<dyn UserReadRepository>>,
video_tasks: Option<Arc<dyn VideoTaskReadRepository>>,
wallets: Option<Arc<dyn WalletReadRepository>>,
shadow_results: Option<Arc<dyn ShadowResultReadRepository>>,
}
impl fmt::Debug for DataReadRepositories {
@@ -69,7 +67,6 @@ impl fmt::Debug for DataReadRepositories {
.field("has_users", &self.users.is_some())
.field("has_video_tasks", &self.video_tasks.is_some())
.field("has_wallets", &self.wallets.is_some())
.field("has_shadow_results", &self.shadow_results.is_some())
.finish()
}
}
@@ -96,7 +93,6 @@ impl DataReadRepositories {
users: postgres.map(PostgresBackend::user_read_repository),
video_tasks: postgres.map(PostgresBackend::video_task_read_repository),
wallets: postgres.map(PostgresBackend::wallet_read_repository),
shadow_results: postgres.map(PostgresBackend::shadow_result_read_repository),
}
}
@@ -170,10 +166,6 @@ impl DataReadRepositories {
self.wallets.clone()
}
pub fn shadow_results(&self) -> Option<Arc<dyn ShadowResultReadRepository>> {
self.shadow_results.clone()
}
pub fn has_any(&self) -> bool {
self.auth_api_keys.is_some()
|| self.announcements.is_some()
@@ -192,7 +184,6 @@ impl DataReadRepositories {
|| self.users.is_some()
|| self.video_tasks.is_some()
|| self.wallets.is_some()
|| self.shadow_results.is_some()
}
}
@@ -235,6 +226,5 @@ mod tests {
assert!(read.usage().is_some());
assert!(read.video_tasks().is_some());
assert!(read.wallets().is_some());
assert!(read.shadow_results().is_some());
}
}

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

@@ -14,7 +14,6 @@ use crate::repository::provider_catalog::ProviderCatalogWriteRepository;
use crate::repository::proxy_nodes::ProxyNodeWriteRepository;
use crate::repository::quota::ProviderQuotaWriteRepository;
use crate::repository::settlement::SettlementWriteRepository;
use crate::repository::shadow_results::ShadowResultWriteRepository;
use crate::repository::usage::UsageWriteRepository;
use crate::repository::video_tasks::VideoTaskWriteRepository;
use crate::repository::wallet::WalletWriteRepository;
@@ -24,7 +23,6 @@ pub struct DataWriteRepositories {
announcements: Option<Arc<dyn AnnouncementWriteRepository>>,
auth_api_keys: Option<Arc<dyn AuthApiKeyWriteRepository>>,
auth_modules: Option<Arc<dyn AuthModuleWriteRepository>>,
shadow_results: Option<Arc<dyn ShadowResultWriteRepository>>,
request_candidates: Option<Arc<dyn RequestCandidateWriteRepository>>,
gemini_file_mappings: Option<Arc<dyn GeminiFileMappingWriteRepository>>,
global_models: Option<Arc<dyn GlobalModelWriteRepository>>,
@@ -45,7 +43,6 @@ impl fmt::Debug for DataWriteRepositories {
.field("has_announcements", &self.announcements.is_some())
.field("has_auth_api_keys", &self.auth_api_keys.is_some())
.field("has_auth_modules", &self.auth_modules.is_some())
.field("has_shadow_results", &self.shadow_results.is_some())
.field("has_request_candidates", &self.request_candidates.is_some())
.field(
"has_gemini_file_mappings",
@@ -71,7 +68,6 @@ impl DataWriteRepositories {
announcements: postgres.map(PostgresBackend::announcement_write_repository),
auth_api_keys: postgres.map(PostgresBackend::auth_api_key_write_repository),
auth_modules: postgres.map(PostgresBackend::auth_module_write_repository),
shadow_results: postgres.map(PostgresBackend::shadow_result_write_repository),
request_candidates: postgres.map(PostgresBackend::request_candidate_write_repository),
gemini_file_mappings: postgres
.map(PostgresBackend::gemini_file_mapping_write_repository),
@@ -88,10 +84,6 @@ impl DataWriteRepositories {
}
}
pub fn shadow_results(&self) -> Option<Arc<dyn ShadowResultWriteRepository>> {
self.shadow_results.clone()
}
pub fn announcements(&self) -> Option<Arc<dyn AnnouncementWriteRepository>> {
self.announcements.clone()
}
@@ -156,7 +148,6 @@ impl DataWriteRepositories {
self.announcements.is_some()
|| self.auth_api_keys.is_some()
|| self.auth_modules.is_some()
|| self.shadow_results.is_some()
|| self.request_candidates.is_some()
|| self.gemini_file_mappings.is_some()
|| self.global_models.is_some()
@@ -179,7 +170,7 @@ mod tests {
use crate::postgres::PostgresPoolConfig;
#[tokio::test]
async fn builds_shadow_result_writer_from_postgres_backend() {
async fn builds_write_repositories_from_postgres_backend() {
let backend = PostgresBackend::from_config(PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
@@ -198,7 +189,6 @@ mod tests {
assert!(write.announcements().is_some());
assert!(write.auth_api_keys().is_some());
assert!(write.auth_modules().is_some());
assert!(write.shadow_results().is_some());
assert!(write.request_candidates().is_some());
assert!(write.gemini_file_mappings().is_some());
assert!(write.global_models().is_some());

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

@@ -14,7 +14,6 @@ pub mod provider_oauth;
pub mod proxy_nodes;
pub mod quota;
pub mod settlement;
pub mod shadow_results;
pub mod system;
pub mod usage;
pub mod users;

170
crates/aether-data/src/repository/shadow_results/memory.rs
View File

@@ -1,170 +0,0 @@
use std::collections::BTreeMap;
use std::sync::RwLock;
use async_trait::async_trait;
use super::types::{
ShadowResultLookupKey, ShadowResultReadRepository, ShadowResultWriteRepository,
StoredShadowResult, UpsertShadowResult,
};
use crate::DataLayerError;
#[derive(Debug, Default)]
pub struct InMemoryShadowResultRepository {
results: RwLock<BTreeMap<(String, String), StoredShadowResult>>,
}
#[async_trait]
impl ShadowResultReadRepository for InMemoryShadowResultRepository {
async fn find(
&self,
key: ShadowResultLookupKey<'_>,
) -> Result<Option<StoredShadowResult>, DataLayerError> {
let results = self.results.read().expect("shadow result repository lock");
Ok(match key {
ShadowResultLookupKey::TraceFingerprint {
trace_id,
request_fingerprint,
} => results
.get(&(trace_id.to_string(), request_fingerprint.to_string()))
.cloned(),
})
}
async fn list_recent(&self, limit: usize) -> Result<Vec<StoredShadowResult>, DataLayerError> {
if limit == 0 {
return Ok(Vec::new());
}
let mut results = self
.results
.read()
.expect("shadow result repository lock")
.values()
.cloned()
.collect::<Vec<_>>();
results.sort_by(|left, right| right.updated_at_unix_secs.cmp(&left.updated_at_unix_secs));
results.truncate(limit);
Ok(results)
}
}
#[async_trait]
impl ShadowResultWriteRepository for InMemoryShadowResultRepository {
async fn upsert(
&self,
result: UpsertShadowResult,
) -> Result<StoredShadowResult, DataLayerError> {
let stored = result.into_stored();
let mut results = self.results.write().expect("shadow result repository lock");
results.insert(
(stored.trace_id.clone(), stored.request_fingerprint.clone()),
stored.clone(),
);
Ok(stored)
}
}
#[cfg(test)]
mod tests {
use super::InMemoryShadowResultRepository;
use crate::repository::shadow_results::{
ShadowResultLookupKey, ShadowResultMatchStatus, ShadowResultReadRepository,
ShadowResultWriteRepository, UpsertShadowResult,
};
fn sample_result(
trace_id: &str,
request_fingerprint: &str,
updated_at_unix_secs: u64,
) -> UpsertShadowResult {
UpsertShadowResult {
trace_id: trace_id.to_string(),
request_fingerprint: request_fingerprint.to_string(),
request_id: Some(format!("req-{trace_id}")),
route_family: Some("openai".to_string()),
route_kind: Some("chat".to_string()),
candidate_id: Some("cand-1".to_string()),
rust_result_digest: Some("rust-digest".to_string()),
python_result_digest: Some("python-digest".to_string()),
match_status: ShadowResultMatchStatus::Match,
status_code: Some(200),
error_message: None,
created_at_unix_ms: updated_at_unix_secs.saturating_sub(10),
updated_at_unix_secs,
}
}
#[tokio::test]
async fn reads_result_by_trace_and_fingerprint() {
let repo = InMemoryShadowResultRepository::default();
repo.upsert(sample_result("trace-1", "fp-1", 100))
.await
.expect("upsert should succeed");
assert!(repo
.find(ShadowResultLookupKey::TraceFingerprint {
trace_id: "trace-1",
request_fingerprint: "fp-1",
})
.await
.expect("find should succeed")
.is_some());
}
#[tokio::test]
async fn list_recent_returns_results_in_descending_update_order() {
let repo = InMemoryShadowResultRepository::default();
repo.upsert(sample_result("trace-1", "fp-1", 100))
.await
.expect("upsert should succeed");
repo.upsert(sample_result("trace-2", "fp-2", 200))
.await
.expect("upsert should succeed");
let recent = repo
.list_recent(10)
.await
.expect("list recent should succeed");
assert_eq!(recent.len(), 2);
assert_eq!(recent[0].trace_id, "trace-2");
assert_eq!(recent[1].trace_id, "trace-1");
}
#[tokio::test]
async fn upsert_replaces_existing_shadow_result() {
let repo = InMemoryShadowResultRepository::default();
repo.upsert(sample_result("trace-1", "fp-1", 100))
.await
.expect("upsert should succeed");
repo.upsert(UpsertShadowResult {
trace_id: "trace-1".to_string(),
request_fingerprint: "fp-1".to_string(),
request_id: Some("req-trace-1".to_string()),
route_family: Some("openai".to_string()),
route_kind: Some("chat".to_string()),
candidate_id: Some("cand-2".to_string()),
rust_result_digest: Some("rust-digest-2".to_string()),
python_result_digest: Some("python-digest-2".to_string()),
match_status: ShadowResultMatchStatus::Mismatch,
status_code: Some(502),
error_message: Some("mismatch".to_string()),
created_at_unix_ms: 100,
updated_at_unix_secs: 200,
})
.await
.expect("upsert should succeed");
let stored = repo
.find(ShadowResultLookupKey::TraceFingerprint {
trace_id: "trace-1",
request_fingerprint: "fp-1",
})
.await
.expect("find should succeed")
.expect("stored result should exist");
assert_eq!(stored.request_id.as_deref(), Some("req-trace-1"));
assert_eq!(stored.candidate_id.as_deref(), Some("cand-2"));
assert_eq!(stored.match_status, ShadowResultMatchStatus::Mismatch);
}
}

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

@@ -1,12 +0,0 @@
mod memory;
mod record;
mod sql;
mod types;
pub use memory::InMemoryShadowResultRepository;
pub use record::{merge_shadow_result_sample, RecordShadowResultSample, ShadowResultSampleOrigin};
pub use sql::SqlxShadowResultRepository;
pub use types::{
ShadowResultLookupKey, ShadowResultMatchStatus, ShadowResultReadRepository,
ShadowResultRepository, ShadowResultWriteRepository, StoredShadowResult, UpsertShadowResult,
};

188
crates/aether-data/src/repository/shadow_results/record.rs
View File

@@ -1,188 +0,0 @@
use super::types::{ShadowResultMatchStatus, StoredShadowResult, UpsertShadowResult};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ShadowResultSampleOrigin {
Rust,
Python,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RecordShadowResultSample {
pub trace_id: String,
pub request_fingerprint: String,
pub request_id: Option<String>,
pub route_family: Option<String>,
pub route_kind: Option<String>,
pub candidate_id: Option<String>,
pub origin: ShadowResultSampleOrigin,
pub result_digest: String,
pub status_code: Option<u16>,
pub error_message: Option<String>,
pub recorded_at_unix_secs: u64,
}
pub fn merge_shadow_result_sample(
existing: Option<&StoredShadowResult>,
sample: RecordShadowResultSample,
) -> UpsertShadowResult {
let RecordShadowResultSample {
trace_id,
request_fingerprint,
request_id,
route_family,
route_kind,
candidate_id,
origin,
result_digest,
status_code,
error_message,
recorded_at_unix_secs,
} = sample;
let (rust_result_digest, python_result_digest) = match origin {
ShadowResultSampleOrigin::Rust => (
Some(result_digest),
existing.and_then(|row| row.python_result_digest.clone()),
),
ShadowResultSampleOrigin::Python => (
existing.and_then(|row| row.rust_result_digest.clone()),
Some(result_digest),
),
};
let match_status = resolve_match_status(
rust_result_digest.as_deref(),
python_result_digest.as_deref(),
);
UpsertShadowResult {
trace_id,
request_fingerprint,
request_id: request_id.or_else(|| existing.and_then(|row| row.request_id.clone())),
route_family: route_family.or_else(|| existing.and_then(|row| row.route_family.clone())),
route_kind: route_kind.or_else(|| existing.and_then(|row| row.route_kind.clone())),
candidate_id: candidate_id.or_else(|| existing.and_then(|row| row.candidate_id.clone())),
rust_result_digest,
python_result_digest,
match_status,
status_code: status_code.or(existing.and_then(|row| row.status_code)),
error_message: resolve_error_message(existing, error_message, match_status),
created_at_unix_ms: existing
.map(|row| row.created_at_unix_ms)
.unwrap_or(recorded_at_unix_secs),
updated_at_unix_secs: recorded_at_unix_secs,
}
}
fn resolve_match_status(
rust_result_digest: Option<&str>,
python_result_digest: Option<&str>,
) -> ShadowResultMatchStatus {
match (rust_result_digest, python_result_digest) {
(Some(rust_digest), Some(python_digest)) if rust_digest == python_digest => {
ShadowResultMatchStatus::Match
}
(Some(_), Some(_)) => ShadowResultMatchStatus::Mismatch,
_ => ShadowResultMatchStatus::Pending,
}
}
fn resolve_error_message(
existing: Option<&StoredShadowResult>,
error_message: Option<String>,
match_status: ShadowResultMatchStatus,
) -> Option<String> {
if match_status == ShadowResultMatchStatus::Mismatch {
error_message
.or_else(|| existing.and_then(|row| row.error_message.clone()))
.or_else(|| Some("shadow result digest mismatch".to_string()))
} else {
error_message.or_else(|| existing.and_then(|row| row.error_message.clone()))
}
}
#[cfg(test)]
mod tests {
use super::{merge_shadow_result_sample, RecordShadowResultSample, ShadowResultSampleOrigin};
use crate::repository::shadow_results::{ShadowResultMatchStatus, UpsertShadowResult};
fn rust_sample(result_digest: &str, recorded_at_unix_secs: u64) -> RecordShadowResultSample {
RecordShadowResultSample {
trace_id: "trace-1".to_string(),
request_fingerprint: "fp-1".to_string(),
request_id: Some("req-1".to_string()),
route_family: Some("openai".to_string()),
route_kind: Some("chat".to_string()),
candidate_id: None,
origin: ShadowResultSampleOrigin::Rust,
result_digest: result_digest.to_string(),
status_code: Some(200),
error_message: None,
recorded_at_unix_secs,
}
}
fn python_sample(result_digest: &str, recorded_at_unix_secs: u64) -> RecordShadowResultSample {
RecordShadowResultSample {
trace_id: "trace-1".to_string(),
request_fingerprint: "fp-1".to_string(),
request_id: Some("req-1".to_string()),
route_family: Some("openai".to_string()),
route_kind: Some("chat".to_string()),
candidate_id: None,
origin: ShadowResultSampleOrigin::Python,
result_digest: result_digest.to_string(),
status_code: Some(200),
error_message: None,
recorded_at_unix_secs,
}
}
fn stored(upsert: UpsertShadowResult) -> crate::repository::shadow_results::StoredShadowResult {
upsert.into_stored()
}
#[test]
fn keeps_pending_until_both_samples_exist() {
let merged = merge_shadow_result_sample(None, rust_sample("digest-1", 100));
assert_eq!(merged.match_status, ShadowResultMatchStatus::Pending);
assert_eq!(merged.request_id.as_deref(), Some("req-1"));
assert_eq!(merged.rust_result_digest.as_deref(), Some("digest-1"));
assert!(merged.python_result_digest.is_none());
}
#[test]
fn marks_match_when_rust_and_python_digests_are_equal() {
let existing = stored(merge_shadow_result_sample(
None,
rust_sample("digest-1", 100),
));
let merged = merge_shadow_result_sample(Some(&existing), python_sample("digest-1", 200));
assert_eq!(merged.match_status, ShadowResultMatchStatus::Match);
assert_eq!(merged.created_at_unix_ms, 100);
assert_eq!(merged.updated_at_unix_secs, 200);
assert_eq!(merged.request_id.as_deref(), Some("req-1"));
assert_eq!(merged.rust_result_digest.as_deref(), Some("digest-1"));
assert_eq!(merged.python_result_digest.as_deref(), Some("digest-1"));
}
#[test]
fn marks_mismatch_when_rust_and_python_digests_differ() {
let existing = stored(merge_shadow_result_sample(
None,
rust_sample("digest-1", 100),
));
let merged = merge_shadow_result_sample(Some(&existing), python_sample("digest-2", 200));
assert_eq!(merged.match_status, ShadowResultMatchStatus::Mismatch);
assert_eq!(
merged.error_message.as_deref(),
Some("shadow result digest mismatch")
);
assert_eq!(merged.request_id.as_deref(), Some("req-1"));
assert_eq!(merged.rust_result_digest.as_deref(), Some("digest-1"));
assert_eq!(merged.python_result_digest.as_deref(), Some("digest-2"));
}
}

291
crates/aether-data/src/repository/shadow_results/sql.rs
View File

@@ -1,291 +0,0 @@
use async_trait::async_trait;
use futures_util::future::BoxFuture;
use futures_util::TryStreamExt;
use sqlx::{PgPool, Row};
use super::types::{
ShadowResultLookupKey, ShadowResultMatchStatus, ShadowResultReadRepository,
ShadowResultWriteRepository, StoredShadowResult, UpsertShadowResult,
};
use crate::postgres::PostgresTransactionRunner;
use crate::{error::SqlxResultExt, DataLayerError};
const FIND_BY_TRACE_FINGERPRINT_SQL: &str = r#"
SELECT
trace_id,
request_fingerprint,
NULL::TEXT AS request_id,
route_family,
route_kind,
candidate_id,
rust_result_digest,
python_result_digest,
match_status,
status_code,
error_message,
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_ms,
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
FROM gateway_shadow_results
WHERE trace_id = $1 AND request_fingerprint = $2
LIMIT 1
"#;
const LIST_RECENT_SQL: &str = r#"
SELECT
trace_id,
request_fingerprint,
NULL::TEXT AS request_id,
route_family,
route_kind,
candidate_id,
rust_result_digest,
python_result_digest,
match_status,
status_code,
error_message,
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_ms,
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
FROM gateway_shadow_results
ORDER BY updated_at DESC
LIMIT $1
"#;
const UPSERT_SQL: &str = r#"
INSERT INTO gateway_shadow_results (
trace_id,
request_fingerprint,
route_family,
route_kind,
candidate_id,
rust_result_digest,
python_result_digest,
match_status,
status_code,
error_message,
created_at,
updated_at
) VALUES (
$1,
$2,
$3,
$4,
$5,
$6,
$7,
$8,
$9,
$10,
TO_TIMESTAMP($11::double precision),
TO_TIMESTAMP($12::double precision)
)
ON CONFLICT (trace_id, request_fingerprint)
DO UPDATE SET
route_family = EXCLUDED.route_family,
route_kind = EXCLUDED.route_kind,
candidate_id = EXCLUDED.candidate_id,
rust_result_digest = EXCLUDED.rust_result_digest,
python_result_digest = EXCLUDED.python_result_digest,
match_status = EXCLUDED.match_status,
status_code = EXCLUDED.status_code,
error_message = EXCLUDED.error_message,
updated_at = EXCLUDED.updated_at
RETURNING
trace_id,
request_fingerprint,
NULL::TEXT AS request_id,
route_family,
route_kind,
candidate_id,
rust_result_digest,
python_result_digest,
match_status,
status_code,
error_message,
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_ms,
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
"#;
#[derive(Debug, Clone)]
pub struct SqlxShadowResultRepository {
pool: PgPool,
tx_runner: PostgresTransactionRunner,
}
impl SqlxShadowResultRepository {
pub fn new(pool: PgPool) -> Self {
let tx_runner = PostgresTransactionRunner::new(pool.clone());
Self { pool, tx_runner }
}
pub fn pool(&self) -> &PgPool {
&self.pool
}
pub fn transaction_runner(&self) -> &PostgresTransactionRunner {
&self.tx_runner
}
pub async fn find(
&self,
key: ShadowResultLookupKey<'_>,
) -> Result<Option<StoredShadowResult>, DataLayerError> {
match key {
ShadowResultLookupKey::TraceFingerprint {
trace_id,
request_fingerprint,
} => {
self.find_by_trace_fingerprint(trace_id, request_fingerprint)
.await
}
}
}
pub async fn find_by_trace_fingerprint(
&self,
trace_id: &str,
request_fingerprint: &str,
) -> Result<Option<StoredShadowResult>, DataLayerError> {
let row = sqlx::query(FIND_BY_TRACE_FINGERPRINT_SQL)
.bind(trace_id)
.bind(request_fingerprint)
.fetch_optional(&self.pool)
.await
.map_postgres_err()?;
row.as_ref().map(map_shadow_result_row).transpose()
}
pub async fn list_recent(
&self,
limit: usize,
) -> Result<Vec<StoredShadowResult>, DataLayerError> {
if limit == 0 {
return Ok(Vec::new());
}
let mut rows = sqlx::query(LIST_RECENT_SQL)
.bind(i64::try_from(limit).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"invalid recent shadow result limit: {limit}"
))
})?)
.fetch(&self.pool);
let mut items = Vec::new();
while let Some(row) = rows.try_next().await.map_postgres_err()? {
items.push(map_shadow_result_row(&row)?);
}
Ok(items)
}
pub async fn upsert(
&self,
result: UpsertShadowResult,
) -> Result<StoredShadowResult, DataLayerError> {
self.tx_runner
.run_read_write(|tx| {
Box::pin(async move {
let row = sqlx::query(UPSERT_SQL)
.bind(&result.trace_id)
.bind(&result.request_fingerprint)
.bind(&result.route_family)
.bind(&result.route_kind)
.bind(&result.candidate_id)
.bind(&result.rust_result_digest)
.bind(&result.python_result_digest)
.bind(match_status_to_database(result.match_status))
.bind(result.status_code.map(i32::from))
.bind(&result.error_message)
.bind(result.created_at_unix_ms as f64)
.bind(result.updated_at_unix_secs as f64)
.fetch_one(&mut **tx)
.await
.map_postgres_err()?;
map_shadow_result_row(&row)
}) as BoxFuture<'_, Result<StoredShadowResult, DataLayerError>>
})
.await
}
}
#[async_trait]
impl ShadowResultReadRepository for SqlxShadowResultRepository {
async fn find(
&self,
key: ShadowResultLookupKey<'_>,
) -> Result<Option<StoredShadowResult>, DataLayerError> {
Self::find(self, key).await
}
async fn list_recent(&self, limit: usize) -> Result<Vec<StoredShadowResult>, DataLayerError> {
Self::list_recent(self, limit).await
}
}
#[async_trait]
impl ShadowResultWriteRepository for SqlxShadowResultRepository {
async fn upsert(
&self,
result: UpsertShadowResult,
) -> Result<StoredShadowResult, DataLayerError> {
Self::upsert(self, result).await
}
}
fn match_status_to_database(status: ShadowResultMatchStatus) -> &'static str {
match status {
ShadowResultMatchStatus::Pending => "pending",
ShadowResultMatchStatus::Match => "match",
ShadowResultMatchStatus::Mismatch => "mismatch",
ShadowResultMatchStatus::Error => "error",
}
}
fn map_shadow_result_row(
row: &sqlx::postgres::PgRow,
) -> Result<StoredShadowResult, DataLayerError> {
let match_status = ShadowResultMatchStatus::from_database(
row.try_get::<String, _>("match_status")
.map_postgres_err()?
.as_str(),
)?;
StoredShadowResult::new(
row.try_get("trace_id").map_postgres_err()?,
row.try_get("request_fingerprint").map_postgres_err()?,
row.try_get("request_id").map_postgres_err()?,
row.try_get("route_family").map_postgres_err()?,
row.try_get("route_kind").map_postgres_err()?,
row.try_get("candidate_id").map_postgres_err()?,
row.try_get("rust_result_digest").map_postgres_err()?,
row.try_get("python_result_digest").map_postgres_err()?,
match_status,
row.try_get("status_code").map_postgres_err()?,
row.try_get("error_message").map_postgres_err()?,
row.try_get("created_at_unix_ms").map_postgres_err()?,
row.try_get("updated_at_unix_secs").map_postgres_err()?,
)
}
#[cfg(test)]
mod tests {
use super::SqlxShadowResultRepository;
use crate::postgres::{PostgresPoolConfig, PostgresPoolFactory};
#[tokio::test]
async fn repository_constructs_from_lazy_pool() {
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
})
.expect("factory should build");
let pool = factory.connect_lazy().expect("pool should build");
let repository = SqlxShadowResultRepository::new(pool);
let _ = repository.pool();
let _ = repository.transaction_runner();
}
}

227
crates/aether-data/src/repository/shadow_results/types.rs
View File

@@ -1,227 +0,0 @@
use async_trait::async_trait;
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum ShadowResultMatchStatus {
Pending,
Match,
Mismatch,
Error,
}
impl ShadowResultMatchStatus {
pub fn from_database(value: &str) -> Result<Self, crate::DataLayerError> {
match value.trim().to_ascii_lowercase().as_str() {
"pending" => Ok(Self::Pending),
"match" => Ok(Self::Match),
"mismatch" => Ok(Self::Mismatch),
"error" => Ok(Self::Error),
other => Err(crate::DataLayerError::UnexpectedValue(format!(
"unsupported gateway_shadow_results.match_status: {other}"
))),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct StoredShadowResult {
pub trace_id: String,
pub request_fingerprint: String,
pub request_id: Option<String>,
pub route_family: Option<String>,
pub route_kind: Option<String>,
pub candidate_id: Option<String>,
pub rust_result_digest: Option<String>,
pub python_result_digest: Option<String>,
pub match_status: ShadowResultMatchStatus,
pub status_code: Option<u16>,
pub error_message: Option<String>,
pub created_at_unix_ms: u64,
pub updated_at_unix_secs: u64,
}
impl StoredShadowResult {
#[allow(clippy::too_many_arguments)]
pub fn new(
trace_id: String,
request_fingerprint: String,
request_id: Option<String>,
route_family: Option<String>,
route_kind: Option<String>,
candidate_id: Option<String>,
rust_result_digest: Option<String>,
python_result_digest: Option<String>,
match_status: ShadowResultMatchStatus,
status_code: Option<i32>,
error_message: Option<String>,
created_at_unix_ms: i64,
updated_at_unix_secs: i64,
) -> Result<Self, crate::DataLayerError> {
let status_code = status_code
.map(|value| {
u16::try_from(value).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!("invalid status_code: {value}"))
})
})
.transpose()?;
let created_at_unix_ms = u64::try_from(created_at_unix_ms).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid created_at_unix_ms: {created_at_unix_ms}"
))
})?;
let updated_at_unix_secs = u64::try_from(updated_at_unix_secs).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid updated_at_unix_secs: {updated_at_unix_secs}"
))
})?;
Ok(Self {
trace_id,
request_fingerprint,
request_id,
route_family,
route_kind,
candidate_id,
rust_result_digest,
python_result_digest,
match_status,
status_code,
error_message,
created_at_unix_ms,
updated_at_unix_secs,
})
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UpsertShadowResult {
pub trace_id: String,
pub request_fingerprint: String,
pub request_id: Option<String>,
pub route_family: Option<String>,
pub route_kind: Option<String>,
pub candidate_id: Option<String>,
pub rust_result_digest: Option<String>,
pub python_result_digest: Option<String>,
pub match_status: ShadowResultMatchStatus,
pub status_code: Option<u16>,
pub error_message: Option<String>,
pub created_at_unix_ms: u64,
pub updated_at_unix_secs: u64,
}
impl UpsertShadowResult {
pub fn into_stored(self) -> StoredShadowResult {
StoredShadowResult {
trace_id: self.trace_id,
request_fingerprint: self.request_fingerprint,
request_id: self.request_id,
route_family: self.route_family,
route_kind: self.route_kind,
candidate_id: self.candidate_id,
rust_result_digest: self.rust_result_digest,
python_result_digest: self.python_result_digest,
match_status: self.match_status,
status_code: self.status_code,
error_message: self.error_message,
created_at_unix_ms: self.created_at_unix_ms,
updated_at_unix_secs: self.updated_at_unix_secs,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ShadowResultLookupKey<'a> {
TraceFingerprint {
trace_id: &'a str,
request_fingerprint: &'a str,
},
}
#[async_trait]
pub trait ShadowResultReadRepository: Send + Sync {
async fn find(
&self,
key: ShadowResultLookupKey<'_>,
) -> Result<Option<StoredShadowResult>, crate::DataLayerError>;
async fn list_recent(
&self,
limit: usize,
) -> Result<Vec<StoredShadowResult>, crate::DataLayerError>;
}
#[async_trait]
pub trait ShadowResultWriteRepository: Send + Sync {
async fn upsert(
&self,
result: UpsertShadowResult,
) -> Result<StoredShadowResult, crate::DataLayerError>;
}
pub trait ShadowResultRepository:
ShadowResultReadRepository + ShadowResultWriteRepository + Send + Sync
{
}
impl<T> ShadowResultRepository for T where
T: ShadowResultReadRepository + ShadowResultWriteRepository + Send + Sync
{
}
#[cfg(test)]
mod tests {
use super::{ShadowResultMatchStatus, StoredShadowResult};
#[test]
fn parses_match_status_from_database_text() {
assert_eq!(
ShadowResultMatchStatus::from_database("match").expect("status should parse"),
ShadowResultMatchStatus::Match
);
}
#[test]
fn rejects_invalid_database_status() {
assert!(ShadowResultMatchStatus::from_database("mystery").is_err());
}
#[test]
fn rejects_invalid_numeric_fields() {
assert!(StoredShadowResult::new(
"trace-1".to_string(),
"fp-1".to_string(),
None,
None,
None,
None,
None,
None,
ShadowResultMatchStatus::Pending,
Some(-1),
None,
1,
1,
)
.is_err());
}
#[test]
fn rejects_negative_updated_at_values() {
assert!(StoredShadowResult::new(
"trace-1".to_string(),
"fp-1".to_string(),
None,
None,
None,
None,
None,
None,
ShadowResultMatchStatus::Pending,
Some(200),
None,
1,
-1,
)
.is_err());
}
}