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

fix: harden frontdoor and usage ingestion

Browse Source
This commit is contained in:
fawney19
2026-05-22 23:57:38 +08:00
parent 6ef6cbade2
commit d18b13a91a
25 changed files with 1363 additions and 252 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
crates/aether-data/migrations/postgres/20260522000000_decouple_usage_identity_foreign_keys.sql Normal file
View File

@@ -0,0 +1,21 @@
-- Usage is a historical fact table. Terminal usage events can arrive after
-- mutable catalog/auth rows have been disabled or deleted, so these snapshot
-- identity columns must not make ingestion depend on current dimension rows.
ALTER TABLE ONLY public.usage
DROP CONSTRAINT IF EXISTS usage_provider_id_fkey;
ALTER TABLE ONLY public.usage
DROP CONSTRAINT IF EXISTS usage_provider_endpoint_id_fkey;
ALTER TABLE ONLY public.usage
DROP CONSTRAINT IF EXISTS usage_provider_api_key_id_fkey;
ALTER TABLE ONLY public.usage
DROP CONSTRAINT IF EXISTS usage_api_key_id_fkey;
ALTER TABLE ONLY public.usage
DROP CONSTRAINT IF EXISTS usage_user_id_fkey;
ALTER TABLE ONLY public.usage
DROP CONSTRAINT IF EXISTS usage_wallet_id_fkey;

76
crates/aether-data/schema/bootstrap/postgres/005_foreign_keys.sql
View File

@@ -583,80 +583,12 @@ END $mig$;
--
-- Name: usage usage_api_key_id_fkey; Type: FK CONSTRAINT; Schema: public; Owner: -
-- Usage is a historical fact table. Snapshot identity columns such as
-- user_id/api_key_id/provider_endpoint_id/provider_api_key_id/wallet_id
-- intentionally do not carry foreign keys because terminal usage events may
-- arrive after those mutable dimension rows have been disabled or deleted.
--
DO $mig$ BEGIN
ALTER TABLE ONLY public.usage
ADD CONSTRAINT usage_api_key_id_fkey FOREIGN KEY (api_key_id) REFERENCES public.api_keys(id) ON DELETE SET NULL;
EXCEPTION
WHEN duplicate_object THEN NULL;
WHEN duplicate_table THEN NULL;
WHEN invalid_table_definition THEN NULL;
END $mig$;
--
-- Name: usage usage_provider_api_key_id_fkey; Type: FK CONSTRAINT; Schema: public; Owner: -
--
DO $mig$ BEGIN
ALTER TABLE ONLY public.usage
ADD CONSTRAINT usage_provider_api_key_id_fkey FOREIGN KEY (provider_api_key_id) REFERENCES public.provider_api_keys(id) ON DELETE SET NULL;
EXCEPTION
WHEN duplicate_object THEN NULL;
WHEN duplicate_table THEN NULL;
WHEN invalid_table_definition THEN NULL;
END $mig$;
--
-- Name: usage usage_provider_endpoint_id_fkey; Type: FK CONSTRAINT; Schema: public; Owner: -
--
DO $mig$ BEGIN
ALTER TABLE ONLY public.usage
ADD CONSTRAINT usage_provider_endpoint_id_fkey FOREIGN KEY (provider_endpoint_id) REFERENCES public.provider_endpoints(id) ON DELETE SET NULL;
EXCEPTION
WHEN duplicate_object THEN NULL;
WHEN duplicate_table THEN NULL;
WHEN invalid_table_definition THEN NULL;
END $mig$;
--
-- Name: usage usage_user_id_fkey; Type: FK CONSTRAINT; Schema: public; Owner: -
--
DO $mig$ BEGIN
ALTER TABLE ONLY public.usage
ADD CONSTRAINT usage_user_id_fkey FOREIGN KEY (user_id) REFERENCES public.users(id) ON DELETE SET NULL;
EXCEPTION
WHEN duplicate_object THEN NULL;
WHEN duplicate_table THEN NULL;
WHEN invalid_table_definition THEN NULL;
END $mig$;
--
-- Name: usage usage_wallet_id_fkey; Type: FK CONSTRAINT; Schema: public; Owner: -
--
DO $mig$ BEGIN
ALTER TABLE ONLY public.usage
ADD CONSTRAINT usage_wallet_id_fkey FOREIGN KEY (wallet_id) REFERENCES public.wallets(id) ON DELETE SET NULL;
EXCEPTION
WHEN duplicate_object THEN NULL;
WHEN duplicate_table THEN NULL;
WHEN invalid_table_definition THEN NULL;
END $mig$;
--
-- Name: user_model_usage_counts user_model_usage_counts_user_id_fkey; Type: FK CONSTRAINT; Schema: public; Owner: -
--

2
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 = 20260520010000;
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260522000000;
const PUBLIC_BASE_TABLE_COUNT_SQL: &str = r#"
SELECT COUNT(*)::BIGINT

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

@@ -312,6 +312,7 @@ fn empty_database_snapshot_covers_current_cutoff_versions() {
20260519130000,
20260520000000,
20260520010000,
20260522000000,
]
);
}
@@ -393,6 +394,34 @@ fn empty_database_snapshot_sql_includes_usage_body_blobs_and_audit_admin_role()
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("usage_count bigint DEFAULT 0 NOT NULL"));
}
#[test]
fn usage_identity_foreign_keys_are_decoupled_for_historical_ingestion() {
let migration = POSTGRES_MIGRATOR
.iter()
.find(|migration| migration.version == 20260522000000)
.expect("usage identity foreign key decoupling migration should be embedded");
for constraint in [
"usage_provider_id_fkey",
"usage_provider_endpoint_id_fkey",
"usage_provider_api_key_id_fkey",
"usage_api_key_id_fkey",
"usage_user_id_fkey",
"usage_wallet_id_fkey",
] {
assert!(
migration
.sql
.contains(format!("DROP CONSTRAINT IF EXISTS {constraint}").as_str()),
"migration should drop {constraint}"
);
assert!(
!EMPTY_DATABASE_SNAPSHOT_SQL.contains(format!("ADD CONSTRAINT {constraint}").as_str()),
"fresh bootstrap snapshot should not recreate {constraint}"
);
}
}
#[test]
fn empty_database_snapshot_sql_includes_payment_gateway_and_plans() {
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("payment_provider character varying(64)"));
@@ -1179,6 +1208,7 @@ fn pending_migrations_from_applied_skips_versions_already_applied() {
20260519130000,
20260520000000,
20260520010000,
20260522000000,
]
);
}

205
crates/aether-usage-runtime/src/worker.rs
View File

@@ -193,16 +193,83 @@ impl UsageQueueWorker {
let event = match UsageEvent::from_stream_fields(&entry.fields) {
Ok(event) => event,
Err(err) => {
warn!(
event_name = "usage_worker_entry_decode_dead_lettered",
log_type = "ops",
worker_consumer = %self.consumer,
worker_group = %self.config.consumer_group,
entry_id = %entry.id,
error = %err,
"usage worker moved malformed queue entry to dead letter"
);
self.queue.push_dead_letter(entry, &err.to_string()).await?;
return Ok(true);
}
};
self.recorder.record_usage_event(&event).await?;
Ok(true)
match self.recorder.record_usage_event(&event).await {
Ok(()) => Ok(true),
Err(err) if usage_event_record_error_is_permanent(&err) => {
warn!(
event_name = "usage_worker_entry_record_dead_lettered",
log_type = "ops",
worker_consumer = %self.consumer,
worker_group = %self.config.consumer_group,
entry_id = %entry.id,
request_id = %event.request_id,
event_type = ?event.event_type,
provider_name = %event.data.provider_name,
model = %event.data.model,
api_format = event.data.api_format.as_deref().unwrap_or(""),
provider_id = event.data.provider_id.as_deref().unwrap_or(""),
provider_endpoint_id = event.data.provider_endpoint_id.as_deref().unwrap_or(""),
provider_api_key_id = event.data.provider_api_key_id.as_deref().unwrap_or(""),
error = %err,
"usage worker moved non-retryable usage event to dead letter"
);
self.queue.push_dead_letter(entry, &err.to_string()).await?;
Ok(true)
}
Err(err) => {
warn!(
event_name = "usage_worker_entry_record_retryable_failed",
log_type = "ops",
worker_consumer = %self.consumer,
worker_group = %self.config.consumer_group,
entry_id = %entry.id,
request_id = %event.request_id,
event_type = ?event.event_type,
provider_name = %event.data.provider_name,
model = %event.data.model,
api_format = event.data.api_format.as_deref().unwrap_or(""),
provider_id = event.data.provider_id.as_deref().unwrap_or(""),
provider_endpoint_id = event.data.provider_endpoint_id.as_deref().unwrap_or(""),
provider_api_key_id = event.data.provider_api_key_id.as_deref().unwrap_or(""),
error = %err,
"usage worker will retry usage event after record failure"
);
Err(err)
}
}
}
}
fn usage_event_record_error_is_permanent(err: &DataLayerError) -> bool {
match err {
DataLayerError::InvalidConfiguration(_)
| DataLayerError::InvalidInput(_)
| DataLayerError::UnexpectedValue(_) => true,
DataLayerError::Postgres(message) | DataLayerError::Sql(message) => {
database_error_is_known_permanent(message)
}
DataLayerError::Redis(_) | DataLayerError::TimedOut(_) => false,
}
}
fn database_error_is_known_permanent(message: &str) -> bool {
message.contains("SQLSTATE 23503") || message.contains("violates foreign key constraint")
}
pub fn build_usage_queue_worker<T>(
runner: Arc<dyn RuntimeQueueStore>,
data: Arc<T>,
@@ -287,16 +354,23 @@ fn consumer_name() -> String {
#[cfg(test)]
mod tests {
use std::sync::Mutex;
use std::sync::{Arc, Mutex};
use aether_data_contracts::repository::settlement::{
StoredUsageSettlement, UsageSettlementInput,
};
use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UpsertUsageRecord};
use aether_data_contracts::DataLayerError;
use aether_runtime_state::{MemoryRuntimeStateConfig, RuntimeQueueStore, RuntimeState};
use async_trait::async_trait;
use super::{write_event_record, ManualProxyNodeCounter, UsageRecordWriter};
use crate::{UsageEvent, UsageEventData, UsageEventType, UsageSettlementWriter};
use super::{
usage_event_record_error_is_permanent, write_event_record, ManualProxyNodeCounter,
UsageEventRecorder, UsageQueueWorker, UsageRecordWriter,
};
use crate::{
UsageEvent, UsageEventData, UsageEventType, UsageRuntimeConfig, UsageSettlementWriter,
};
#[derive(Default)]
struct TestUsageStore {
@@ -304,6 +378,11 @@ mod tests {
settlements: Mutex<Vec<UsageSettlementInput>>,
}
#[derive(Default)]
struct SelectiveFailingRecorder {
calls: Mutex<Vec<String>>,
}
#[async_trait]
impl UsageRecordWriter for TestUsageStore {
async fn upsert_usage_record(
@@ -392,6 +471,22 @@ mod tests {
}
}
#[async_trait]
impl UsageEventRecorder for SelectiveFailingRecorder {
async fn record_usage_event(&self, event: &UsageEvent) -> Result<(), DataLayerError> {
self.calls
.lock()
.expect("calls lock")
.push(event.request_id.clone());
if event.request_id == "req-worker-poison" {
return Err(DataLayerError::UnexpectedValue(
"permanent test error".to_string(),
));
}
Ok(())
}
}
fn sample_event() -> UsageEvent {
UsageEvent::new(
UsageEventType::Completed,
@@ -436,4 +531,104 @@ mod tests {
assert_eq!(settlements.len(), 1);
assert_eq!(settlements[0].request_id, "req-worker-123");
}
#[test]
fn usage_event_record_error_classifies_permanent_failures() {
assert!(usage_event_record_error_is_permanent(
&DataLayerError::UnexpectedValue("bad payload".to_string())
));
assert!(usage_event_record_error_is_permanent(
&DataLayerError::Postgres(
"error returned from database: violates foreign key constraint (SQLSTATE 23503)"
.to_string()
)
));
assert!(!usage_event_record_error_is_permanent(
&DataLayerError::Redis("connection refused".to_string())
));
assert!(!usage_event_record_error_is_permanent(
&DataLayerError::TimedOut("postgres acquire".to_string())
));
}
#[tokio::test]
async fn process_entries_dead_letters_permanent_record_error_and_continues() {
let runner = Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default()));
let queue_runner: Arc<dyn RuntimeQueueStore> = runner.clone();
let recorder = Arc::new(SelectiveFailingRecorder::default());
let config = UsageRuntimeConfig {
enabled: true,
stream_key: "usage:test:worker:events".to_string(),
consumer_group: "usage:test:worker:group".to_string(),
dlq_stream_key: "usage:test:worker:dlq".to_string(),
consumer_batch_size: 10,
consumer_block_ms: 1,
..UsageRuntimeConfig::default()
};
let worker = UsageQueueWorker::new(queue_runner, recorder.clone(), config)
.expect("worker should build");
worker
.queue
.ensure_consumer_group()
.await
.expect("group should initialize");
let mut poison = sample_event();
poison.request_id = "req-worker-poison".to_string();
let mut ok = sample_event();
ok.request_id = "req-worker-ok".to_string();
worker
.queue
.enqueue(&poison)
.await
.expect("poison event should enqueue");
worker
.queue
.enqueue(&ok)
.await
.expect("ok event should enqueue");
let entries = worker
.queue
.read_group(&worker.consumer)
.await
.expect("events should read");
assert_eq!(entries.len(), 2);
worker
.process_entries(entries)
.await
.expect("permanent failure should not block batch");
assert_eq!(
recorder.calls.lock().expect("calls lock").as_slice(),
["req-worker-poison", "req-worker-ok"]
);
runner
.ensure_consumer_group(
"usage:test:worker:dlq",
"usage:test:worker:dlq-group",
"0-0",
)
.await
.expect("dlq group should initialize");
let dlq_entries = runner
.read_group(
"usage:test:worker:dlq",
"usage:test:worker:dlq-group",
"usage-test-dlq-consumer",
10,
Some(1),
)
.await
.expect("dlq should read");
assert_eq!(dlq_entries.len(), 1);
let payload = dlq_entries[0]
.fields
.get("payload")
.expect("dlq payload should exist");
assert!(payload.contains("req-worker-poison"));
assert!(payload.contains("permanent test error"));
}
}