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

feat: revamp analytics dashboards and harden database migrations

Browse Source 
Add dashboard and overview analytics, health monitoring, provider expense tracking, and announcement updates across the gateway and frontend.

Keep schema migrations free of historical backfills while preserving automatic backfill execution. Bound migration deadlines, run schema preparation before Compose replacement, and anonymize deleted dashboard users.

Include the current documentation cleanup and regression coverage.
This commit is contained in:
elky
2026-10-01 11:48:17 +08:00
parent 60b89cc840
commit 066ea87d72
327 changed files with 31728 additions and 20645 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-data/runtime/schema/README.md
+69
View File
@@ -109,3 +109,72 @@ remains byte-for-byte stable when composed:
The Rust migration tests compose these manifests too, so fragment drift is
caught during `cargo test -p aether-data split_baseline_sources_match_executable_migrations`.
## Statistics Migrations
The statistics release retains its applied migration history and includes
incremental upgrades for databases that ran the earlier overview and dashboard
definitions. Concurrent index operations remain separate because PostgreSQL
cannot run them inside a transaction.
| Version | Change |
|---|---|
| `20260911000000` | Overview facts, attribution, aggregate tables, and transaction-owned dirty-event queue. Attribution indexes are created while the new table is empty. |
| `20260917000000` | Original account-attribution migration, retained byte-for-byte for databases that already applied it. |
| `20260917000100` | Upgrade the original attribution trigger to the dirty-event queue before later concurrent index builds. |
| `20260918000000` | Create the replacement settlement covering index concurrently. |
| `20260918000100` | Drop the previous settlement covering index concurrently, after its replacement succeeds. |
| `20260919000000` | Dashboard aggregates, activation boundary, and retention support. |
| `20260920000000` | Create the credited-payment lookup index concurrently. |
| `20260920120000` | Provider expense records. |
| `20260921010000` | Add retention support to existing dashboard schemas; safe when the initial dashboard migration already includes it. |
| `20260921020000` | Add the attribution-owner lookup index concurrently on existing databases. |
| `20260921020100` | Create the usage metadata actor index concurrently. |
| `20261001000000` | Remove deleted-user attribution from dashboard activity on future user deletion; schema-only upgrade without rewriting historical rows. |
Do not remove an applied migration after folding its changes into an earlier
schema definition. Existing databases retain its version in `_sqlx_migrations`
and do not rerun earlier versions when their SQL changes. Preserve that history
and provide incremental migrations for any remaining schema differences.
These migrations do not backfill historical requests. Dashboard totals start at
the stored activation boundary. Background maintenance compacts dashboard minute
details older than 35 days in bounded batches, preserving cumulative totals and
the narrow activity counts; it does not delete source usage. JSONL backups include
the dashboard snapshot and its integrity manifest so retained totals can survive
restoration after source usage has expired.
Schema migrations and historical backfills remain separate phases. Normal `auto`
startup and `db prepare` still apply pending scripts from `backfills/postgres`
after schema migration; `verify-only` still requires both phases to be current.
The statistics schema migrations above do not embed a historical data rebuild.
The new dashboard's activation boundary is not moved by legacy backfills, so
they do not restore pre-activation dashboard totals.
Deleted users are excluded from dashboard active-user reads even when an older
version left orphan activity rows. The anonymization upgrade installs rules for
future deletions without cleaning old rows during migration; those old activity
rows age out through the existing 35-day retention task.
The overview worker can still rebuild a historical hour/day when normal writes
change facts in that bucket. That work runs after startup with bounded batches
and query deadlines; it is not a full historical rebuild during migration.
The migration runner defaults to a 1-second lock wait, a 10-second deadline per
transactional migration, and a 15-minute deadline per concurrent index migration.
Timeouts are configurable through `AETHER_POSTGRES_MIGRATION_LOCK_TIMEOUT_MS`,
`AETHER_POSTGRES_MIGRATION_TIMEOUT_MS`, and
`AETHER_POSTGRES_MIGRATION_CONCURRENT_TIMEOUT_MS`; none accepts zero. An independent
control connection attempts to terminate the migration session on failure or
cancellation. An interrupted concurrent index build can leave an invalid index;
the runner removes that index before retrying its migration.
For Compose deployments, `update.sh` applies schema migrations with the new image
before replacing the running app. A migration failure stops the update; already
committed migrations remain applied. Its `local-build` mode delegates to
`deploy.sh` and does not use this separate migration step. Allow for brief table
locks and I/O pressure from concurrent index scans during the upgrade. Keeping
historical backfills out of schema migrations does not make index creation
constant-time: concurrent indexes still scan existing rows and can take minutes
on a large database. The existing app stays running during the Compose migration
preflight.
crates/aether-data/runtime/schema/bootstrap/postgres/001_types_and_tables.sql
+1 -4
View File
@@ -18,10 +18,7 @@
-- Runs before 20260403000000_baseline.sql so that fresh databases have a
-- complete schema by the time baseline (a no-op handoff point) and all
-- later ADD COLUMN IF NOT EXISTS migrations execute.
SET statement_timeout = 0;
SET lock_timeout = 0;
-- Keep the migration runner's statement and lock deadlines in effect.
SET idle_in_transaction_session_timeout = 0;
SET client_encoding = 'UTF8';
crates/aether-data/runtime/schema/bootstrap/postgres/190_overview_analytics.sql
+276
View File
@@ -0,0 +1,276 @@
ALTER TABLE public.usage ADD COLUMN IF NOT EXISTS failure_origin text;
ALTER TABLE public.usage ADD COLUMN IF NOT EXISTS failure_stage text;
ALTER TABLE public.usage ADD COLUMN IF NOT EXISTS failure_reason text;
ALTER TABLE public.usage ADD COLUMN IF NOT EXISTS failure_schema_version integer;
ALTER TABLE public.usage_settlement_snapshots
ADD COLUMN IF NOT EXISTS quota_covered_amount_usd numeric(20,8),
ADD COLUMN IF NOT EXISTS wallet_consumed_amount_usd numeric(20,8),
ADD COLUMN IF NOT EXISTS wallet_debit_amount_usd numeric(20,8),
ADD COLUMN IF NOT EXISTS wallet_recharge_debit_usd numeric(20,8),
ADD COLUMN IF NOT EXISTS wallet_gift_debit_usd numeric(20,8),
ADD COLUMN IF NOT EXISTS wallet_overdraft_usd numeric(20,8),
ADD COLUMN IF NOT EXISTS allocation_schema_version integer,
ADD COLUMN IF NOT EXISTS allocation_status text;
CREATE TABLE IF NOT EXISTS public.usage_attribution_snapshots (
request_id text PRIMARY KEY,
actor_user_id text,
credential_owner_id text,
attribution_kind text NOT NULL DEFAULT 'unknown',
attribution_source text NOT NULL DEFAULT 'unknown',
record_kind text NOT NULL DEFAULT 'request',
parent_request_id text,
schema_version integer NOT NULL DEFAULT 1,
attribution_revision bigint NOT NULL DEFAULT 1,
recorded_at timestamptz NOT NULL DEFAULT NOW()
);
CREATE INDEX IF NOT EXISTS ix_usage_attribution_actor_request
ON public.usage_attribution_snapshots(actor_user_id, request_id);
-- This attribution table is new and empty; build both lookup indexes here.
CREATE INDEX IF NOT EXISTS ix_usage_attribution_owner_request
ON public.usage_attribution_snapshots(credential_owner_id, request_id);
CREATE TABLE IF NOT EXISTS public.stats_bucket_state (
projection_version text NOT NULL,
granularity text NOT NULL,
bucket_start timestamptz NOT NULL,
source_revision bigint NOT NULL DEFAULT 0,
built_revision bigint NOT NULL DEFAULT -1,
coverage_status text NOT NULL DEFAULT 'unbuilt',
built_at timestamptz,
last_error text,
last_failed_at timestamptz,
PRIMARY KEY (projection_version, granularity, bucket_start)
);
CREATE INDEX IF NOT EXISTS ix_stats_bucket_state_dirty
ON public.stats_bucket_state(bucket_start)
WHERE source_revision > built_revision;
CREATE TABLE IF NOT EXISTS public.stats_overview_hourly (
projection_version text NOT NULL,
bucket_start timestamptz NOT NULL,
dimensions jsonb NOT NULL,
metrics jsonb NOT NULL,
PRIMARY KEY(projection_version, bucket_start, dimensions)
);
CREATE TABLE IF NOT EXISTS public.stats_overview_daily (
projection_version text NOT NULL,
bucket_start timestamptz NOT NULL,
dimensions jsonb NOT NULL,
metrics jsonb NOT NULL,
PRIMARY KEY(projection_version, bucket_start, dimensions)
);
-- These triggers cover old writers, delayed settlement, and maintenance in the fact transaction.
-- Only future fact mutations enqueue work; no historical rows are backfilled.
-- Each writer owns its transaction's keys, so unrelated requests never contend
-- on the current hour/day's stats_bucket_state row.
CREATE TABLE IF NOT EXISTS public.stats_overview_dirty_events (
transaction_id bigint NOT NULL,
projection_version text NOT NULL,
granularity text NOT NULL,
bucket_start timestamptz NOT NULL,
unrecoverable boolean NOT NULL DEFAULT false,
PRIMARY KEY (transaction_id, projection_version, granularity, bucket_start)
);
CREATE INDEX IF NOT EXISTS ix_stats_overview_dirty_events_bucket
ON public.stats_overview_dirty_events(projection_version, granularity, bucket_start);
CREATE OR REPLACE FUNCTION public.overview_mark_usage_bucket() RETURNS trigger
LANGUAGE plpgsql AS $$
DECLARE old_time timestamptz; new_time timestamptz; bucket record;
BEGIN
IF TG_TABLE_NAME = 'usage' THEN
IF TG_OP <> 'INSERT' THEN old_time := OLD.created_at; END IF;
IF TG_OP <> 'DELETE' THEN new_time := NEW.created_at; END IF;
ELSE
IF TG_OP <> 'INSERT' THEN
SELECT created_at INTO old_time FROM public.usage WHERE request_id = OLD.request_id;
END IF;
IF TG_OP <> 'DELETE' THEN
SELECT created_at INTO new_time FROM public.usage WHERE request_id = NEW.request_id;
END IF;
END IF;
FOR bucket IN
SELECT DISTINCT g, date_trunc(g, t AT TIME ZONE 'UTC') AT TIME ZONE 'UTC' AS starts
FROM unnest(ARRAY[old_time, new_time]) t CROSS JOIN unnest(ARRAY['day','hour']) g
WHERE t IS NOT NULL ORDER BY g, starts
LOOP
INSERT INTO public.stats_overview_dirty_events
(transaction_id, projection_version, granularity, bucket_start, unrecoverable)
VALUES (txid_current(), 'overview-v2', bucket.g, bucket.starts,
TG_TABLE_NAME = 'usage' AND TG_OP = 'DELETE')
ON CONFLICT (transaction_id, projection_version, granularity, bucket_start)
DO UPDATE SET unrecoverable = stats_overview_dirty_events.unrecoverable OR EXCLUDED.unrecoverable;
END LOOP;
RETURN NULL;
END $$;
DROP TRIGGER IF EXISTS overview_usage_dirty ON public.usage;
CREATE TRIGGER overview_usage_dirty AFTER INSERT OR UPDATE OR DELETE ON public.usage
FOR EACH ROW EXECUTE FUNCTION public.overview_mark_usage_bucket();
DROP TRIGGER IF EXISTS overview_settlement_dirty ON public.usage_settlement_snapshots;
CREATE TRIGGER overview_settlement_dirty AFTER INSERT OR UPDATE OR DELETE ON public.usage_settlement_snapshots
FOR EACH ROW EXECUTE FUNCTION public.overview_mark_usage_bucket();
DROP TRIGGER IF EXISTS overview_attribution_dirty ON public.usage_attribution_snapshots;
CREATE TRIGGER overview_attribution_dirty AFTER INSERT OR UPDATE OR DELETE ON public.usage_attribution_snapshots
FOR EACH ROW EXECUTE FUNCTION public.overview_mark_usage_bucket();
CREATE OR REPLACE FUNCTION public.overview_capture_usage_identity() RETURNS trigger
LANGUAGE plpgsql AS $$
DECLARE actor text; owner text; kind text; source text; standalone boolean;
BEGIN
SELECT id INTO owner FROM public.users WHERE id=NEW.user_id AND NOT is_deleted;
SELECT COALESCE(k.is_standalone,
CASE WHEN jsonb_typeof(NEW.request_metadata::jsonb #> '{analytics_attribution,is_standalone}')='boolean'
THEN (NEW.request_metadata #>> '{analytics_attribution,is_standalone}')::boolean END,
CASE WHEN jsonb_typeof(NEW.request_metadata::jsonb->'api_key_is_standalone')='boolean'
THEN (NEW.request_metadata->>'api_key_is_standalone')::boolean END,
CASE WHEN a.attribution_source='user_account' THEN false
WHEN a.attribution_source='standalone_key' THEN true END,
CASE WHEN NEW.api_key_id IS NULL THEN false END)
INTO standalone FROM (SELECT 1) seed
LEFT JOIN public.api_keys k ON k.id=NEW.api_key_id
LEFT JOIN public.usage_attribution_snapshots a ON a.request_id=NEW.request_id;
kind := CASE WHEN owner IS NULL THEN 'unknown' WHEN standalone THEN 'standalone'
WHEN NOT standalone THEN 'employee' ELSE 'unknown' END;
actor := CASE WHEN kind='employee' THEN owner END;
source := CASE kind WHEN 'employee' THEN 'user_account' WHEN 'standalone' THEN 'standalone_key' ELSE 'unknown' END;
INSERT INTO public.usage_attribution_snapshots(request_id, actor_user_id, credential_owner_id,
attribution_kind, attribution_source, record_kind, parent_request_id, attribution_revision)
VALUES (NEW.request_id, actor, owner, kind, source,
COALESCE(NEW.request_metadata #>> '{analytics_attribution,record_kind}', 'request'),
NEW.request_metadata #>> '{analytics_attribution,parent_request_id}', 2)
ON CONFLICT (request_id) DO UPDATE SET actor_user_id = EXCLUDED.actor_user_id,
credential_owner_id = EXCLUDED.credential_owner_id, attribution_kind = EXCLUDED.attribution_kind,
attribution_source = EXCLUDED.attribution_source,
record_kind = COALESCE(NEW.request_metadata #>> '{analytics_attribution,record_kind}', usage_attribution_snapshots.record_kind),
parent_request_id = COALESCE(EXCLUDED.parent_request_id, usage_attribution_snapshots.parent_request_id),
attribution_revision = EXCLUDED.attribution_revision, recorded_at = NOW()
WHERE usage_attribution_snapshots.attribution_revision <= EXCLUDED.attribution_revision;
RETURN NULL;
END $$;
DROP TRIGGER IF EXISTS overview_usage_identity ON public.usage;
CREATE TRIGGER overview_usage_identity AFTER INSERT OR UPDATE OF user_id, api_key_id, request_metadata ON public.usage
FOR EACH ROW EXECUTE FUNCTION public.overview_capture_usage_identity();
CREATE OR REPLACE FUNCTION public.overview_capture_failure() RETURNS trigger
LANGUAGE plpgsql AS $$
BEGIN
IF NEW.status = 'completed' THEN
NEW.failure_origin := NULL; NEW.failure_stage := NULL; NEW.failure_reason := NULL;
NEW.failure_schema_version := NULL; RETURN NEW;
END IF;
IF NEW.request_metadata #>> '{analytics_failure,origin}' IN ('client','gateway','upstream','transport','unknown') THEN
NEW.failure_origin := NEW.request_metadata #>> '{analytics_failure,origin}';
NEW.failure_stage := NEW.request_metadata #>> '{analytics_failure,stage}';
NEW.failure_reason := NEW.request_metadata #>> '{analytics_failure,reason}';
NEW.failure_schema_version := 1;
END IF;
RETURN NEW;
END $$;
DROP TRIGGER IF EXISTS overview_usage_failure ON public.usage;
CREATE TRIGGER overview_usage_failure BEFORE INSERT OR UPDATE OF request_metadata, status ON public.usage
FOR EACH ROW EXECUTE FUNCTION public.overview_capture_failure();
CREATE OR REPLACE FUNCTION public.overview_anonymize_user() RETURNS trigger
LANGUAGE plpgsql AS $$
BEGIN
IF TG_OP = 'DELETE' OR NEW.is_deleted THEN
UPDATE public.usage_attribution_snapshots SET actor_user_id = NULL, credential_owner_id = NULL,
attribution_kind = 'unknown', attribution_source = 'unknown', attribution_revision = attribution_revision + 10
WHERE actor_user_id = OLD.id OR credential_owner_id = OLD.id;
UPDATE public.usage SET request_metadata = (request_metadata::jsonb #- '{analytics_attribution,actor_user_id}')::json
WHERE request_metadata #>> '{analytics_attribution,actor_user_id}' = OLD.id;
DELETE FROM public.stats_overview_hourly WHERE dimensions->>'actor_user_id'=OLD.id OR dimensions->>'credential_owner_id'=OLD.id;
DELETE FROM public.stats_overview_daily WHERE dimensions->>'actor_user_id'=OLD.id OR dimensions->>'credential_owner_id'=OLD.id;
END IF;
RETURN NULL;
END $$;
DROP TRIGGER IF EXISTS overview_user_anonymize ON public.users;
CREATE TRIGGER overview_user_anonymize AFTER DELETE OR UPDATE OF is_deleted ON public.users
FOR EACH ROW EXECUTE FUNCTION public.overview_anonymize_user();
CREATE OR REPLACE FUNCTION public.overview_delete_attribution() RETURNS trigger
LANGUAGE plpgsql AS $$
BEGIN
DELETE FROM public.usage_attribution_snapshots WHERE request_id = OLD.request_id;
RETURN OLD;
END $$;
DROP TRIGGER IF EXISTS overview_usage_delete_attribution ON public.usage;
CREATE TRIGGER overview_usage_delete_attribution BEFORE DELETE ON public.usage
FOR EACH ROW EXECUTE FUNCTION public.overview_delete_attribution();
CREATE OR REPLACE VIEW public.usage_analytics_facts_v1 AS
SELECT u.request_id, COALESCE(u.id, u.request_id) AS id, u.created_at,
CASE WHEN identity.owner_id IS NOT NULL AND identity.is_standalone=false THEN identity.owner_id END AS actor_user_id,
identity.owner_id AS credential_owner_id,
CASE WHEN identity.owner_id IS NULL THEN 'unknown' WHEN identity.is_standalone THEN 'standalone'
WHEN NOT identity.is_standalone THEN 'employee' ELSE 'unknown' END AS attribution_kind,
CASE WHEN identity.owner_id IS NULL THEN 'unknown' WHEN identity.is_standalone THEN 'standalone_key'
WHEN NOT identity.is_standalone THEN 'user_account' ELSE 'unknown' END AS attribution_source,
COALESCE(a.record_kind, 'request') AS record_kind, a.parent_request_id,
u.api_key_id, u.model, u.target_model, u.provider_id, u.provider_name,
u.api_format, u.endpoint_kind, u.request_type, u.is_stream, u.has_format_conversion,
u.status, u.status_code, u.error_category, u.failure_origin, u.failure_stage, u.failure_reason,
u.failure_schema_version, u.response_time_ms, u.first_byte_time_ms,
COALESCE(s.billing_status, u.billing_status) AS settlement_status,
COALESCE(u.request_metadata::jsonb->'usage_available', 'true'::jsonb) <> 'false'::jsonb AS usage_available,
COALESCE(u.request_metadata::jsonb->'usage_pricing_available', 'true'::jsonb) <> 'false'::jsonb
AND (s.billing_total_cost_usd IS NOT NULL OR COALESCE(s.billing_status, u.billing_status) = 'settled') AS pricing_available,
CASE WHEN COALESCE(u.request_metadata::jsonb->'usage_available', 'true'::jsonb) <> 'false'::jsonb
THEN b.input_tokens END AS input_tokens,
CASE WHEN COALESCE(u.request_metadata::jsonb->'usage_available', 'true'::jsonb) <> 'false'::jsonb
THEN b.output_tokens END AS output_tokens,
CASE WHEN COALESCE(u.request_metadata::jsonb->'usage_available', 'true'::jsonb) <> 'false'::jsonb
THEN b.total_tokens END AS total_tokens,
CASE WHEN COALESCE(u.request_metadata::jsonb->'usage_available', 'true'::jsonb) <> 'false'::jsonb
THEN b.cache_read_input_tokens END AS cache_read_input_tokens,
CASE WHEN COALESCE(u.request_metadata::jsonb->'usage_available', 'true'::jsonb) <> 'false'::jsonb
THEN b.cache_creation_input_tokens END AS cache_creation_input_tokens,
CASE WHEN COALESCE(u.request_metadata::jsonb->'usage_pricing_available', 'true'::jsonb) <> 'false'::jsonb
AND (s.billing_total_cost_usd IS NOT NULL OR COALESCE(s.billing_status, u.billing_status) = 'settled')
THEN round(COALESCE(s.billing_total_cost_usd::numeric, u.total_cost_usd::numeric), 8) END AS rated_amount,
CASE WHEN COALESCE(u.request_metadata::jsonb->'usage_pricing_available', 'true'::jsonb) <> 'false'::jsonb
AND (s.billing_actual_total_cost_usd IS NOT NULL OR COALESCE(s.billing_status, u.billing_status) = 'settled')
THEN round(COALESCE(s.billing_actual_total_cost_usd::numeric, u.actual_total_cost_usd::numeric), 8) END AS billable_amount,
s.quota_covered_amount_usd AS quota_covered_amount,
s.wallet_consumed_amount_usd AS wallet_consumed_amount,
s.wallet_debit_amount_usd AS wallet_debit_amount,
s.wallet_recharge_debit_usd AS wallet_recharge_debit_amount,
s.wallet_gift_debit_usd AS wallet_gift_debit_amount,
s.wallet_overdraft_usd AS wallet_overdraft_amount,
s.allocation_status, s.finalized_at AS settled_at,
CASE WHEN s.billing_total_cost_usd IS NOT NULL THEN 'settlement_snapshot' ELSE 'legacy_float' END AS amount_source,
b.upstream_is_stream,
CASE WHEN u.request_metadata #>> '{analytics_measurement,source}' IN ('reported','estimated','mixed')
THEN u.request_metadata #>> '{analytics_measurement,source}' ELSE 'unknown' END AS token_source,
CASE WHEN COALESCE(u.request_metadata::jsonb->'usage_available','true'::jsonb) <> 'false'::jsonb
AND COALESCE(u.request_metadata::jsonb->'usage_pricing_available','true'::jsonb) <> 'false'::jsonb
AND s.input_price_per_1m IS NOT NULL AND s.billing_cache_read_cost_usd IS NOT NULL
THEN round(s.input_price_per_1m::numeric * b.cache_read_input_tokens::numeric / 1000000,8) END AS cache_estimated_full_cost_amount,
CASE WHEN COALESCE(u.request_metadata::jsonb->'usage_available','true'::jsonb) <> 'false'::jsonb
AND COALESCE(u.request_metadata::jsonb->'usage_pricing_available','true'::jsonb) <> 'false'::jsonb
AND s.input_price_per_1m IS NOT NULL AND s.billing_cache_read_cost_usd IS NOT NULL
THEN round(s.billing_cache_read_cost_usd::numeric,8) END AS cache_read_cost_amount,
CASE WHEN COALESCE(u.request_metadata::jsonb->'usage_available','true'::jsonb) <> 'false'::jsonb
AND COALESCE(u.request_metadata::jsonb->'usage_pricing_available','true'::jsonb) <> 'false'::jsonb
AND s.input_price_per_1m IS NOT NULL AND s.billing_cache_creation_cost_usd IS NOT NULL
THEN round(s.billing_cache_creation_cost_usd::numeric,8) END AS cache_creation_cost_amount
FROM public.usage u
LEFT JOIN public.usage_settlement_snapshots s USING (request_id)
LEFT JOIN public.usage_attribution_snapshots a USING (request_id)
JOIN public.usage_billing_facts b USING (request_id)
LEFT JOIN public.api_keys k ON k.id=u.api_key_id
CROSS JOIN LATERAL (
SELECT CASE WHEN a.request_id IS NOT NULL THEN a.credential_owner_id
WHEN EXISTS (SELECT 1 FROM public.users WHERE id=u.user_id AND NOT is_deleted) THEN u.user_id END AS owner_id,
COALESCE(k.is_standalone,
CASE WHEN jsonb_typeof(u.request_metadata::jsonb #> '{analytics_attribution,is_standalone}')='boolean'
THEN (u.request_metadata #>> '{analytics_attribution,is_standalone}')::boolean END,
CASE WHEN jsonb_typeof(u.request_metadata::jsonb->'api_key_is_standalone')='boolean'
THEN (u.request_metadata->>'api_key_is_standalone')::boolean END,
CASE WHEN a.attribution_source='user_account' THEN false
WHEN a.attribution_source='standalone_key' THEN true END,
CASE WHEN u.api_key_id IS NULL THEN false END) AS is_standalone
) identity;
crates/aether-data/runtime/schema/bootstrap/postgres/210_provider_expenses.sql
+25
View File
@@ -0,0 +1,25 @@
-- Manual purchasing ledger; no foreign-key cascade may erase historical expenditures.
CREATE TABLE IF NOT EXISTS public.provider_expenses (
id text PRIMARY KEY,
client_request_id text NOT NULL UNIQUE,
provider_id text NOT NULL,
provider_name text NOT NULL,
kind text NOT NULL CHECK (kind IN ('recharge', 'subscription', 'other')),
amount numeric(20,8) NOT NULL CHECK (amount > 0),
currency text NOT NULL CHECK (currency ~ '^[A-Z]{3}$'),
paid_at timestamptz NOT NULL,
period_start timestamptz,
period_end timestamptz,
note text,
external_reference text,
created_by text,
created_at timestamptz NOT NULL DEFAULT NOW(),
voided_at timestamptz,
voided_by text,
CONSTRAINT provider_expenses_period_check CHECK (
(period_start IS NULL AND period_end IS NULL) OR
(period_start IS NOT NULL AND period_end IS NOT NULL AND period_start < period_end)
)
);
CREATE INDEX IF NOT EXISTS ix_provider_expenses_paid_at ON public.provider_expenses(paid_at, id);
CREATE INDEX IF NOT EXISTS ix_provider_expenses_provider_paid_at ON public.provider_expenses(provider_id, paid_at);
crates/aether-data/runtime/schema/bootstrap/postgres/manifest.txt
+1
View File
@@ -13,3 +13,4 @@
160_routing_profiles.sql
170_usage_legacy_body_ref_cleanup_index.sql
180_usage_stale_pending_cleanup_index.sql
210_provider_expenses.sql
crates/aether-data/runtime/schema/drivers/postgres/baseline/001_types_and_tables.sql
+1 -4
View File
@@ -18,10 +18,7 @@
-- Runs before 20260403000000_baseline.sql so that fresh databases have a
-- complete schema by the time baseline (a no-op handoff point) and all
-- later ADD COLUMN IF NOT EXISTS migrations execute.
SET statement_timeout = 0;
SET lock_timeout = 0;
-- Keep the migration runner's statement and lock deadlines in effect.
SET idle_in_transaction_session_timeout = 0;
SET client_encoding = 'UTF8';
crates/aether-data/runtime/schema/generated/postgres/baseline/005_wallet_billing.sql
+1
View File
@@ -104,6 +104,7 @@ ALTER TABLE ONLY public.payment_orders ADD CONSTRAINT uq_payment_orders_order_no
CREATE INDEX IF NOT EXISTS idx_payment_orders_wallet_created ON public.payment_orders USING btree (wallet_id, created_at);
CREATE INDEX IF NOT EXISTS idx_payment_orders_user_created ON public.payment_orders USING btree (user_id, created_at);
CREATE INDEX IF NOT EXISTS idx_payment_orders_status ON public.payment_orders USING btree (status);
CREATE INDEX IF NOT EXISTS idx_payment_orders_status_credited_user ON public.payment_orders USING btree (status, credited_at, user_id);
CREATE INDEX IF NOT EXISTS idx_payment_orders_gateway_order_id ON public.payment_orders USING btree (gateway_order_id);
CREATE UNIQUE INDEX IF NOT EXISTS uq_payment_orders_payment_method_gateway_order_id ON public.payment_orders USING btree (payment_method, gateway_order_id);
CREATE INDEX IF NOT EXISTS idx_payment_orders_kind_status ON public.payment_orders USING btree (order_kind, status);
crates/aether-data/runtime/schema/generated/postgres/baseline/006_usage.sql
+13 -1
View File
@@ -99,7 +99,11 @@ CREATE TABLE IF NOT EXISTS public.usage (
wallet_gift_balance_after double precision,
finalized_at bigint,
created_at_unix_ms bigint DEFAULT 0 NOT NULL,
updated_at_unix_secs bigint DEFAULT 0 NOT NULL
updated_at_unix_secs bigint DEFAULT 0 NOT NULL,
failure_origin text,
failure_stage text,
failure_reason text,
failure_schema_version integer
);
ALTER TABLE ONLY public.usage ADD CONSTRAINT usage_pkey PRIMARY KEY (request_id);
@@ -204,6 +208,14 @@ CREATE INDEX IF NOT EXISTS ix_usage_counter_deltas_request_kind ON public.usage_
CREATE TABLE IF NOT EXISTS public.usage_settlement_snapshots (
request_id character varying(128) NOT NULL,
billing_status character varying(64) NOT NULL,
quota_covered_amount_usd numeric(20,8),
wallet_consumed_amount_usd numeric(20,8),
wallet_debit_amount_usd numeric(20,8),
wallet_recharge_debit_usd numeric(20,8),
wallet_gift_debit_usd numeric(20,8),
wallet_overdraft_usd numeric(20,8),
allocation_schema_version integer,
allocation_status text,
wallet_id character varying(64),
wallet_balance_before double precision,
wallet_balance_after double precision,
crates/aether-data/runtime/schema/generated/postgres/baseline/009_overview.sql
+63
View File
@@ -0,0 +1,63 @@
-- Generated by aether-data-schema from crates/aether-data/runtime/schema/logical/*.toml.
-- Do not edit generated files directly; edit logical schema or explicit overrides instead.
CREATE TABLE IF NOT EXISTS public.stats_bucket_state (
projection_version text NOT NULL,
granularity text NOT NULL,
bucket_start timestamp with time zone NOT NULL,
source_revision bigint DEFAULT 0 NOT NULL,
built_revision bigint DEFAULT -1 NOT NULL,
coverage_status text DEFAULT 'unbuilt' NOT NULL,
built_at timestamp with time zone,
last_error text,
last_failed_at timestamp with time zone
);
ALTER TABLE ONLY public.stats_bucket_state ADD CONSTRAINT stats_bucket_state_pkey PRIMARY KEY (projection_version, granularity, bucket_start);
CREATE TABLE IF NOT EXISTS public.stats_overview_daily (
projection_version text NOT NULL,
bucket_start timestamp with time zone NOT NULL,
dimensions jsonb NOT NULL,
metrics jsonb NOT NULL
);
ALTER TABLE ONLY public.stats_overview_daily ADD CONSTRAINT stats_overview_daily_pkey PRIMARY KEY (projection_version, bucket_start, dimensions);
CREATE TABLE IF NOT EXISTS public.stats_overview_dirty_events (
transaction_id bigint NOT NULL,
projection_version text NOT NULL,
granularity text NOT NULL,
bucket_start timestamp with time zone NOT NULL,
unrecoverable boolean DEFAULT false NOT NULL
);
ALTER TABLE ONLY public.stats_overview_dirty_events ADD CONSTRAINT stats_overview_dirty_events_pkey PRIMARY KEY (transaction_id, projection_version, granularity, bucket_start);
CREATE INDEX IF NOT EXISTS ix_stats_overview_dirty_events_bucket ON public.stats_overview_dirty_events USING btree (projection_version, granularity, bucket_start);
CREATE TABLE IF NOT EXISTS public.stats_overview_hourly (
projection_version text NOT NULL,
bucket_start timestamp with time zone NOT NULL,
dimensions jsonb NOT NULL,
metrics jsonb NOT NULL
);
ALTER TABLE ONLY public.stats_overview_hourly ADD CONSTRAINT stats_overview_hourly_pkey PRIMARY KEY (projection_version, bucket_start, dimensions);
CREATE TABLE IF NOT EXISTS public.usage_attribution_snapshots (
request_id text NOT NULL,
actor_user_id text,
credential_owner_id text,
attribution_kind text DEFAULT 'unknown' NOT NULL,
attribution_source text DEFAULT 'unknown' NOT NULL,
record_kind text DEFAULT 'request' NOT NULL,
parent_request_id text,
schema_version integer DEFAULT 1 NOT NULL,
attribution_revision bigint DEFAULT 1 NOT NULL,
recorded_at timestamp with time zone DEFAULT NOW() NOT NULL
);
ALTER TABLE ONLY public.usage_attribution_snapshots ADD CONSTRAINT usage_attribution_snapshots_pkey PRIMARY KEY (request_id);
CREATE INDEX IF NOT EXISTS ix_usage_attribution_actor_request ON public.usage_attribution_snapshots USING btree (actor_user_id, request_id);
CREATE INDEX IF NOT EXISTS ix_usage_attribution_owner_request ON public.usage_attribution_snapshots USING btree (credential_owner_id, request_id);
crates/aether-data/runtime/schema/generated/postgres/baseline/010_dashboard.sql
+84
View File
@@ -0,0 +1,84 @@
-- Generated by aether-data-schema from crates/aether-data/runtime/schema/logical/*.toml.
-- Do not edit generated files directly; edit logical schema or explicit overrides instead.
CREATE TABLE IF NOT EXISTS public.dashboard_activity_hour (
bucket_start timestamp with time zone NOT NULL,
shard smallint NOT NULL,
request_count bigint NOT NULL
);
ALTER TABLE ONLY public.dashboard_activity_hour ADD CONSTRAINT dashboard_activity_hour_pkey PRIMARY KEY (bucket_start, shard);
CREATE TABLE IF NOT EXISTS public.dashboard_activity_minute (
bucket_start timestamp with time zone NOT NULL,
shard smallint NOT NULL,
request_count bigint NOT NULL
);
ALTER TABLE ONLY public.dashboard_activity_minute ADD CONSTRAINT dashboard_activity_minute_pkey PRIMARY KEY (bucket_start, shard);
CREATE TABLE IF NOT EXISTS public.dashboard_actor_minute (
bucket_start timestamp with time zone NOT NULL,
shard smallint NOT NULL,
actor_user_id character varying(255) NOT NULL,
request_count bigint NOT NULL
);
ALTER TABLE ONLY public.dashboard_actor_minute ADD CONSTRAINT dashboard_actor_minute_pkey PRIMARY KEY (bucket_start, shard, actor_user_id);
CREATE TABLE IF NOT EXISTS public.dashboard_request_contributions (
request_id character varying(100) NOT NULL,
created_at timestamp with time zone NOT NULL,
actor_user_id character varying(255),
metrics jsonb NOT NULL
);
ALTER TABLE ONLY public.dashboard_request_contributions ADD CONSTRAINT dashboard_request_contributions_pkey PRIMARY KEY (request_id);
CREATE TABLE IF NOT EXISTS public.dashboard_stats_minute (
bucket_start timestamp with time zone NOT NULL,
shard smallint NOT NULL,
metrics jsonb DEFAULT '{}'::jsonb NOT NULL
);
ALTER TABLE ONLY public.dashboard_stats_minute ADD CONSTRAINT dashboard_stats_minute_pkey PRIMARY KEY (bucket_start, shard);
CREATE TABLE IF NOT EXISTS public.dashboard_stats_pending (
transaction_id bigint NOT NULL,
request_id character varying(100) NOT NULL,
deleted_fact jsonb
);
ALTER TABLE ONLY public.dashboard_stats_pending ADD CONSTRAINT dashboard_stats_pending_pkey PRIMARY KEY (transaction_id, request_id);
CREATE TABLE IF NOT EXISTS public.dashboard_stats_state (
singleton boolean DEFAULT true NOT NULL,
stats_since timestamp with time zone NOT NULL,
contributions_cleanup_cursor character varying(100)
);
ALTER TABLE ONLY public.dashboard_stats_state ADD CONSTRAINT dashboard_stats_state_pkey PRIMARY KEY (singleton);
CREATE TABLE IF NOT EXISTS public.dashboard_stats_total (
shard smallint NOT NULL,
metrics jsonb DEFAULT '{}'::jsonb NOT NULL
);
ALTER TABLE ONLY public.dashboard_stats_total ADD CONSTRAINT dashboard_stats_total_pkey PRIMARY KEY (shard);
CREATE TABLE IF NOT EXISTS public.dashboard_user_anonymization_pending (
transaction_id bigint NOT NULL,
user_id character varying(255) NOT NULL
);
ALTER TABLE ONLY public.dashboard_user_anonymization_pending ADD CONSTRAINT dashboard_user_anonymization_pending_pkey PRIMARY KEY (transaction_id, user_id);
CREATE TABLE IF NOT EXISTS public.dashboard_user_events_minute (
bucket_start timestamp with time zone NOT NULL,
shard smallint NOT NULL,
created_count bigint DEFAULT 0 NOT NULL,
deleted_count bigint DEFAULT 0 NOT NULL
);
ALTER TABLE ONLY public.dashboard_user_events_minute ADD CONSTRAINT dashboard_user_events_minute_pkey PRIMARY KEY (bucket_start, shard);
crates/aether-data/runtime/schema/generated/postgres/baseline/010_provider_expenses.sql
+27
View File
@@ -0,0 +1,27 @@
-- Generated by aether-data-schema from crates/aether-data/runtime/schema/logical/*.toml.
-- Do not edit generated files directly; edit logical schema or explicit overrides instead.
CREATE TABLE IF NOT EXISTS public.provider_expenses (
id text NOT NULL,
client_request_id text NOT NULL,
provider_id text NOT NULL,
provider_name text NOT NULL,
kind text NOT NULL,
amount numeric(20,8) NOT NULL,
currency text NOT NULL,
paid_at timestamp with time zone NOT NULL,
period_start timestamp with time zone,
period_end timestamp with time zone,
note text,
external_reference text,
created_by text,
created_at timestamp with time zone DEFAULT NOW() NOT NULL,
voided_at timestamp with time zone,
voided_by text
);
ALTER TABLE ONLY public.provider_expenses ADD CONSTRAINT provider_expenses_pkey PRIMARY KEY (id);
ALTER TABLE ONLY public.provider_expenses ADD CONSTRAINT provider_expenses_client_request_id_key UNIQUE (client_request_id);
CREATE INDEX IF NOT EXISTS ix_provider_expenses_paid_at ON public.provider_expenses USING btree (paid_at, id);
CREATE INDEX IF NOT EXISTS ix_provider_expenses_provider_paid_at ON public.provider_expenses USING btree (provider_id, paid_at);
crates/aether-data/runtime/schema/generated/postgres/baseline/manifest.txt
+3
View File
@@ -9,3 +9,6 @@
006_usage.sql
007_stats.sql
008_background_tasks.sql
009_overview.sql
010_dashboard.sql
010_provider_expenses.sql
crates/aether-data/runtime/schema/logical/005_wallet_billing.toml
+4
View File
@@ -426,6 +426,10 @@ columns = ["user_id", "created_at"]
name = "idx_payment_orders_status"
columns = ["status"]
[[table.payment_orders.indexes]]
name = "idx_payment_orders_status_credited_user"
columns = ["status", "credited_at", "user_id"]
[[table.payment_orders.indexes]]
name = "idx_payment_orders_gateway_order_id"
columns = ["gateway_order_id"]
crates/aether-data/runtime/schema/logical/006_usage.toml
+66
View File
@@ -928,6 +928,26 @@ columns = ["processed_at", "created_at", "id"]
name = "ix_usage_counter_deltas_request_kind"
columns = ["request_id", "kind", "target_id"]
[[table.usage.columns]]
name = "failure_origin"
type = "text"
nullable = true
[[table.usage.columns]]
name = "failure_stage"
type = "text"
nullable = true
[[table.usage.columns]]
name = "failure_reason"
type = "text"
nullable = true
[[table.usage.columns]]
name = "failure_schema_version"
type = "int32"
nullable = true
[table.usage_settlement_snapshots]
domain = "usage"
order = 20
@@ -943,6 +963,52 @@ name = "billing_status"
type = "text"
length = 64
[[table.usage_settlement_snapshots.columns]]
name = "quota_covered_amount_usd"
type = "decimal_money"
nullable = true
driver.postgres.type = "numeric(20,8)"
[[table.usage_settlement_snapshots.columns]]
name = "wallet_consumed_amount_usd"
type = "decimal_money"
nullable = true
driver.postgres.type = "numeric(20,8)"
[[table.usage_settlement_snapshots.columns]]
name = "wallet_debit_amount_usd"
type = "decimal_money"
nullable = true
driver.postgres.type = "numeric(20,8)"
[[table.usage_settlement_snapshots.columns]]
name = "wallet_recharge_debit_usd"
type = "decimal_money"
nullable = true
driver.postgres.type = "numeric(20,8)"
[[table.usage_settlement_snapshots.columns]]
name = "wallet_gift_debit_usd"
type = "decimal_money"
nullable = true
driver.postgres.type = "numeric(20,8)"
[[table.usage_settlement_snapshots.columns]]
name = "wallet_overdraft_usd"
type = "decimal_money"
nullable = true
driver.postgres.type = "numeric(20,8)"
[[table.usage_settlement_snapshots.columns]]
name = "allocation_schema_version"
type = "int32"
nullable = true
[[table.usage_settlement_snapshots.columns]]
name = "allocation_status"
type = "text"
nullable = true
[[table.usage_settlement_snapshots.columns]]
name = "wallet_id"
type = "text_id"
crates/aether-data/runtime/schema/logical/009_overview.toml
+175
View File
@@ -0,0 +1,175 @@
[table.usage_attribution_snapshots]
domain = "usage"
primary_key = ["request_id"]
[[table.usage_attribution_snapshots.columns]]
name = "request_id"
type = "text"
[[table.usage_attribution_snapshots.columns]]
name = "actor_user_id"
type = "text"
nullable = true
[[table.usage_attribution_snapshots.columns]]
name = "credential_owner_id"
type = "text"
nullable = true
[[table.usage_attribution_snapshots.columns]]
name = "attribution_kind"
type = "text"
default = "unknown"
[[table.usage_attribution_snapshots.columns]]
name = "attribution_source"
type = "text"
default = "unknown"
[[table.usage_attribution_snapshots.columns]]
name = "record_kind"
type = "text"
default = "request"
[[table.usage_attribution_snapshots.columns]]
name = "parent_request_id"
type = "text"
nullable = true
[[table.usage_attribution_snapshots.columns]]
name = "schema_version"
type = "int32"
default = 1
[[table.usage_attribution_snapshots.columns]]
name = "attribution_revision"
type = "int64"
default = 1
[[table.usage_attribution_snapshots.columns]]
name = "recorded_at"
type = "timestamp"
default = { raw = "NOW()" }
[[table.usage_attribution_snapshots.indexes]]
name = "ix_usage_attribution_actor_request"
columns = ["actor_user_id", "request_id"]
[[table.usage_attribution_snapshots.indexes]]
name = "ix_usage_attribution_owner_request"
columns = ["credential_owner_id", "request_id"]
[table.stats_bucket_state]
domain = "stats"
primary_key = ["projection_version","granularity","bucket_start"]
[[table.stats_bucket_state.columns]]
name = "projection_version"
type = "text"
[[table.stats_bucket_state.columns]]
name = "granularity"
type = "text"
[[table.stats_bucket_state.columns]]
name = "bucket_start"
type = "timestamp"
[[table.stats_bucket_state.columns]]
name = "source_revision"
type = "int64"
default = 0
[[table.stats_bucket_state.columns]]
name = "built_revision"
type = "int64"
default = -1
[[table.stats_bucket_state.columns]]
name = "coverage_status"
type = "text"
default = "unbuilt"
[[table.stats_bucket_state.columns]]
name = "built_at"
type = "timestamp"
nullable = true
[[table.stats_bucket_state.columns]]
name = "last_error"
type = "text"
nullable = true
[[table.stats_bucket_state.columns]]
name = "last_failed_at"
type = "timestamp"
nullable = true
[table.stats_overview_hourly]
domain = "stats"
primary_key = ["projection_version","bucket_start","dimensions"]
[[table.stats_overview_hourly.columns]]
name = "projection_version"
type = "text"
[[table.stats_overview_hourly.columns]]
name = "bucket_start"
type = "timestamp"
[[table.stats_overview_hourly.columns]]
name = "dimensions"
type = "json"
[[table.stats_overview_hourly.columns]]
name = "metrics"
type = "json"
[table.stats_overview_daily]
domain = "stats"
primary_key = ["projection_version","bucket_start","dimensions"]
[[table.stats_overview_daily.columns]]
name = "projection_version"
type = "text"
[[table.stats_overview_daily.columns]]
name = "bucket_start"
type = "timestamp"
[[table.stats_overview_daily.columns]]
name = "dimensions"
type = "json"
[[table.stats_overview_daily.columns]]
name = "metrics"
type = "json"
[table.stats_overview_dirty_events]
domain = "stats"
primary_key = ["transaction_id", "projection_version", "granularity", "bucket_start"]
[[table.stats_overview_dirty_events.columns]]
name = "transaction_id"
type = "int64"
[[table.stats_overview_dirty_events.columns]]
name = "projection_version"
type = "text"
[[table.stats_overview_dirty_events.columns]]
name = "granularity"
type = "text"
[[table.stats_overview_dirty_events.columns]]
name = "bucket_start"
type = "timestamp"
[[table.stats_overview_dirty_events.columns]]
name = "unrecoverable"
type = "bool"
default = false
[[table.stats_overview_dirty_events.indexes]]
name = "ix_stats_overview_dirty_events_bucket"
columns = ["projection_version", "granularity", "bucket_start"]
crates/aether-data/runtime/schema/logical/010_dashboard.toml
+187
View File
@@ -0,0 +1,187 @@
# Future-only dashboard aggregates. Trigger functions and activation data live in
# migration 20260919000000_add_future_dashboard_summary.sql; no history is backfilled.
[table.dashboard_activity_hour]
domain = "stats"
primary_key = ["bucket_start", "shard"]
[[table.dashboard_activity_hour.columns]]
name = "bucket_start"
type = "timestamp"
[[table.dashboard_activity_hour.columns]]
name = "shard"
type = "int32"
driver.postgres.type = "smallint"
[[table.dashboard_activity_hour.columns]]
name = "request_count"
type = "int64"
[table.dashboard_stats_state]
domain = "stats"
primary_key = ["singleton"]
[[table.dashboard_stats_state.columns]]
name = "singleton"
type = "bool"
default = true
[[table.dashboard_stats_state.columns]]
name = "stats_since"
type = "timestamp"
[[table.dashboard_stats_state.columns]]
name = "contributions_cleanup_cursor"
type = "text"
length = 100
nullable = true
[table.dashboard_activity_minute]
domain = "stats"
primary_key = ["bucket_start", "shard"]
[[table.dashboard_activity_minute.columns]]
name = "bucket_start"
type = "timestamp"
[[table.dashboard_activity_minute.columns]]
name = "shard"
type = "int32"
driver.postgres.type = "smallint"
[[table.dashboard_activity_minute.columns]]
name = "request_count"
type = "int64"
[table.dashboard_request_contributions]
domain = "stats"
primary_key = ["request_id"]
[[table.dashboard_request_contributions.columns]]
name = "request_id"
type = "text"
length = 100
[[table.dashboard_request_contributions.columns]]
name = "created_at"
type = "timestamp"
[[table.dashboard_request_contributions.columns]]
name = "actor_user_id"
type = "text"
nullable = true
length = 255
[[table.dashboard_request_contributions.columns]]
name = "metrics"
type = "json"
[table.dashboard_stats_total]
domain = "stats"
primary_key = ["shard"]
[[table.dashboard_stats_total.columns]]
name = "shard"
type = "int32"
driver.postgres.type = "smallint"
[[table.dashboard_stats_total.columns]]
name = "metrics"
type = "json"
default = { raw = "'{}'::jsonb" }
[table.dashboard_stats_minute]
domain = "stats"
primary_key = ["bucket_start","shard"]
[[table.dashboard_stats_minute.columns]]
name = "bucket_start"
type = "timestamp"
[[table.dashboard_stats_minute.columns]]
name = "shard"
type = "int32"
driver.postgres.type = "smallint"
[[table.dashboard_stats_minute.columns]]
name = "metrics"
type = "json"
default = { raw = "'{}'::jsonb" }
[table.dashboard_actor_minute]
domain = "stats"
primary_key = ["bucket_start","shard","actor_user_id"]
[[table.dashboard_actor_minute.columns]]
name = "bucket_start"
type = "timestamp"
[[table.dashboard_actor_minute.columns]]
name = "shard"
type = "int32"
driver.postgres.type = "smallint"
[[table.dashboard_actor_minute.columns]]
name = "actor_user_id"
type = "text"
length = 255
[[table.dashboard_actor_minute.columns]]
name = "request_count"
type = "int64"
[table.dashboard_stats_pending]
domain = "stats"
primary_key = ["transaction_id","request_id"]
[[table.dashboard_stats_pending.columns]]
name = "transaction_id"
type = "int64"
[[table.dashboard_stats_pending.columns]]
name = "request_id"
type = "text"
length = 100
[[table.dashboard_stats_pending.columns]]
name = "deleted_fact"
type = "json"
nullable = true
[table.dashboard_user_events_minute]
domain = "stats"
primary_key = ["bucket_start","shard"]
[[table.dashboard_user_events_minute.columns]]
name = "bucket_start"
type = "timestamp"
[[table.dashboard_user_events_minute.columns]]
name = "shard"
type = "int32"
driver.postgres.type = "smallint"
[[table.dashboard_user_events_minute.columns]]
name = "created_count"
type = "int64"
default = 0
[[table.dashboard_user_events_minute.columns]]
name = "deleted_count"
type = "int64"
default = 0
# Transaction-local events only; no committed rows are exported or restored.
[table.dashboard_user_anonymization_pending]
domain = "stats"
primary_key = ["transaction_id", "user_id"]
[[table.dashboard_user_anonymization_pending.columns]]
name = "transaction_id"
type = "int64"
[[table.dashboard_user_anonymization_pending.columns]]
name = "user_id"
type = "text"
length = 255
crates/aether-data/runtime/schema/logical/010_provider_expenses.toml
+88
View File
@@ -0,0 +1,88 @@
[table.provider_expenses]
domain = "wallet_billing"
primary_key = ["id"]
[[table.provider_expenses.columns]]
name = "id"
type = "text"
[[table.provider_expenses.columns]]
name = "client_request_id"
type = "text"
[[table.provider_expenses.columns]]
name = "provider_id"
type = "text"
[[table.provider_expenses.columns]]
name = "provider_name"
type = "text"
[[table.provider_expenses.columns]]
name = "kind"
type = "text"
[[table.provider_expenses.columns]]
name = "amount"
type = "decimal_money"
driver.postgres.type = "numeric(20,8)"
[[table.provider_expenses.columns]]
name = "currency"
type = "text"
[[table.provider_expenses.columns]]
name = "paid_at"
type = "timestamp"
[[table.provider_expenses.columns]]
name = "period_start"
type = "timestamp"
nullable = true
[[table.provider_expenses.columns]]
name = "period_end"
type = "timestamp"
nullable = true
[[table.provider_expenses.columns]]
name = "note"
type = "text"
nullable = true
[[table.provider_expenses.columns]]
name = "external_reference"
type = "text"
nullable = true
[[table.provider_expenses.columns]]
name = "created_by"
type = "text"
nullable = true
[[table.provider_expenses.columns]]
name = "created_at"
type = "timestamp"
default = { raw = "NOW()" }
[[table.provider_expenses.columns]]
name = "voided_at"
type = "timestamp"
nullable = true
[[table.provider_expenses.columns]]
name = "voided_by"
type = "text"
nullable = true
[[table.provider_expenses.uniques]]
name = "provider_expenses_client_request_id_key"
columns = ["client_request_id"]
[[table.provider_expenses.indexes]]
name = "ix_provider_expenses_paid_at"
columns = ["paid_at", "id"]
[[table.provider_expenses.indexes]]
name = "ix_provider_expenses_provider_paid_at"
columns = ["provider_id", "paid_at"]
crates/aether-data/runtime/src/backend/maintenance.rs
+44
View File
@@ -168,6 +168,26 @@ impl DataBackends {
}
}
pub async fn rebuild_overview_buckets(
&self,
input: &StatsHourlyAggregationInput,
) -> Result<usize, DataLayerError> {
match self.sql_backend() {
Some(backend) => backend.rebuild_overview_buckets(input).await,
None => Ok(0),
}
}
pub async fn drain_overview_dirty_events(
&self,
now: chrono::DateTime<chrono::Utc>,
) -> Result<u64, DataLayerError> {
match self.sql_backend() {
Some(backend) => backend.drain_overview_dirty_events(now).await,
None => Ok(0),
}
}
pub async fn aggregate_stats_hourly(
&self,
input: &StatsHourlyAggregationInput,
@@ -417,6 +437,30 @@ impl<'a> SqlBackendRef<'a> {
}
}
async fn rebuild_overview_buckets(
self,
input: &StatsHourlyAggregationInput,
) -> Result<usize, DataLayerError> {
match self {
#[cfg(feature = "postgres")]
Self::Postgres(postgres) => postgres.rebuild_overview_buckets(input).await,
#[cfg(not(feature = "postgres"))]
Self::Disabled(_) => unreachable!("a SQL backend cannot exist without a driver"),
}
}
async fn drain_overview_dirty_events(
self,
now: chrono::DateTime<chrono::Utc>,
) -> Result<u64, DataLayerError> {
match self {
#[cfg(feature = "postgres")]
Self::Postgres(postgres) => postgres.drain_overview_dirty_events(now).await,
#[cfg(not(feature = "postgres"))]
Self::Disabled(_) => unreachable!("a SQL backend cannot exist without a driver"),
}
}
async fn aggregate_stats_hourly(
self,
input: &StatsHourlyAggregationInput,
crates/aether-data/runtime/src/backend/stats/postgres_hourly/mod.rs
+26
View File
@@ -13,6 +13,32 @@ mod sql;
use self::sql::*;
impl PostgresBackend {
pub async fn drain_overview_dirty_events(
&self,
now: DateTime<Utc>,
) -> Result<u64, DataLayerError> {
let repository = aether_data_postgres::SqlxUsageReadRepository::new(self.pool().clone());
let merged = repository.merge_overview_dirty_events().await?;
let retained = repository.maintain_dashboard_projection(now, 1_000).await?;
Ok(merged + u64::from(retained))
}
pub async fn rebuild_overview_buckets(
&self,
input: &StatsHourlyAggregationInput,
) -> Result<usize, DataLayerError> {
let repository = aether_data_postgres::SqlxUsageReadRepository::new(self.pool().clone());
let cleaned = repository
.maintain_dashboard_projection(input.aggregated_at, 1_000)
.await?;
let rebuilt = repository
.rebuild_overview_buckets(input.target_hour_utc + chrono::Duration::hours(1), 8)
.await?;
// Retention work uses the worker's existing bounded catch-up loop too,
// so a busy installation can retire more than one batch per hour.
Ok(rebuilt + usize::from(cleaned))
}
pub async fn aggregate_stats_hourly(
&self,
input: &StatsHourlyAggregationInput,
crates/aether-data/runtime/src/lifecycle/export.rs
+59
View File
@@ -15,6 +15,8 @@ use crate::{DataLayerError, DatabaseDriver, SqlDatabaseConfig};
#[cfg(feature = "postgres")]
mod postgres;
#[cfg(all(test, feature = "postgres"))]
mod dashboard_snapshot_tests;
#[cfg(all(test, feature = "postgres"))]
mod tests;
@@ -97,6 +99,38 @@ struct AuxiliaryTable {
}
const AUXILIARY_TABLES: &[AuxiliaryTable] = &[
AuxiliaryTable {
name: "dashboard_stats_state",
primary_key: &["singleton"],
},
AuxiliaryTable {
name: "dashboard_stats_total",
primary_key: &["shard"],
},
AuxiliaryTable {
name: "dashboard_stats_minute",
primary_key: &["bucket_start", "shard"],
},
AuxiliaryTable {
name: "dashboard_activity_hour",
primary_key: &["bucket_start", "shard"],
},
AuxiliaryTable {
name: "dashboard_activity_minute",
primary_key: &["bucket_start", "shard"],
},
AuxiliaryTable {
name: "dashboard_actor_minute",
primary_key: &["bucket_start", "shard", "actor_user_id"],
},
AuxiliaryTable {
name: "dashboard_user_events_minute",
primary_key: &["bucket_start", "shard"],
},
AuxiliaryTable {
name: "dashboard_request_contributions",
primary_key: &["request_id"],
},
AuxiliaryTable {
name: "audit_logs",
primary_key: &["id"],
@@ -181,6 +215,10 @@ const AUXILIARY_TABLES: &[AuxiliaryTable] = &[
name: "payment_gateway_configs",
primary_key: &["provider"],
},
AuxiliaryTable {
name: "provider_expenses",
primary_key: &["id"],
},
AuxiliaryTable {
name: "billing_plans",
primary_key: &["id"],
@@ -213,6 +251,10 @@ const AUXILIARY_TABLES: &[AuxiliaryTable] = &[
name: "usage_routing_snapshots",
primary_key: &["request_id"],
},
AuxiliaryTable {
name: "usage_attribution_snapshots",
primary_key: &["request_id"],
},
AuxiliaryTable {
name: "usage_counter_deltas",
primary_key: &["id"],
@@ -389,6 +431,22 @@ pub struct DataExportManifest {
pub created_at_unix_secs: u64,
pub source_driver: Option<DatabaseDriver>,
pub domains: Vec<ExportDomain>,
/// Complete dashboard projection, restored atomically rather than merged by row.
/// Older exports omit this field and retain their ordinary import behavior.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub dashboard_snapshot: Option<DashboardSnapshotManifest>,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct DashboardSnapshotManifest {
pub version: u32,
pub tables: BTreeMap<String, DashboardSnapshotTable>,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct DashboardSnapshotTable {
pub rows: usize,
pub sha256: String,
}
impl DataExportManifest {
@@ -405,6 +463,7 @@ impl DataExportManifest {
created_at_unix_secs,
source_driver,
domains,
dashboard_snapshot: None,
}
}
}
crates/aether-data/runtime/src/lifecycle/export/dashboard_snapshot_tests.rs
+241
View File
@@ -0,0 +1,241 @@
use super::*;
use crate::lifecycle::postgres_test_support::ManagedPostgresServer;
use sqlx::PgPool;
async fn migrate(pool: &PgPool) {
crate::lifecycle::migrate::prepare_database_for_startup(pool)
.await
.unwrap();
crate::lifecycle::migrate::run_migrations(pool)
.await
.unwrap();
}
async fn request(pool: &PgPool, id: &str, input: i32) {
sqlx::query("INSERT INTO usage(id,request_id,user_id,api_key_id,provider_name,model,status,billing_status,created_at,input_tokens,output_tokens,total_tokens) VALUES ($1,$1,'backup-user','backup-key','test','test','completed','settled',clock_timestamp(),$2,3,$2+3)")
.bind(id).bind(input).execute(pool).await.unwrap();
}
async fn dashboard_rows(pool: &PgPool) -> BTreeMap<String, Vec<Value>> {
let mut result = BTreeMap::new();
for table in AUXILIARY_TABLES
.iter()
.filter(|table| table.name.starts_with("dashboard_"))
{
let mut rows =
sqlx::query_scalar::<_, Value>(&format!("SELECT to_jsonb(t) FROM {} t", table.name))
.fetch_all(pool)
.await
.unwrap();
rows.sort_by_key(Value::to_string);
result.insert(table.name.to_owned(), rows);
}
result
}
async fn total(pool: &PgPool) -> i64 {
sqlx::query_scalar("SELECT COALESCE(sum((metrics->>'request_count')::bigint),0)::bigint FROM dashboard_stats_total").fetch_one(pool).await.unwrap()
}
async fn tokens(pool: &PgPool) -> i64 {
sqlx::query_scalar("SELECT COALESCE(sum((metrics->>'total_tokens')::bigint),0)::bigint FROM dashboard_stats_total")
.fetch_one(pool).await.unwrap()
}
#[tokio::test]
async fn postgres_dashboard_snapshot_roundtrip_preserves_purged_totals_and_future_updates() {
let Some(server) = ManagedPostgresServer::try_start().await.unwrap() else {
return;
};
let source = sqlx::postgres::PgPoolOptions::new()
.max_connections(2)
.connect(server.database_url())
.await
.unwrap();
migrate(&source).await;
sqlx::raw_sql("INSERT INTO users(id,username,email_verified) VALUES('backup-user','backup-user',false); INSERT INTO api_keys(id,user_id,key_hash) VALUES('backup-key','backup-user',repeat('e',64));").execute(&source).await.unwrap();
request(&source, "backup-kept", 11).await;
request(&source, "backup-purged", 23).await;
sqlx::query("DELETE FROM usage WHERE request_id='backup-purged'")
.execute(&source)
.await
.unwrap();
// Simulate the retention worker dropping a contribution whose raw request is gone.
sqlx::query("DELETE FROM dashboard_request_contributions WHERE request_id='backup-purged'")
.execute(&source)
.await
.unwrap();
assert_eq!(total(&source).await, 2);
assert_eq!(tokens(&source).await, 40);
let expected = dashboard_rows(&source).await;
let export = export_postgres_core_jsonl(&source, 1_800_000_000)
.await
.unwrap();
let plan = build_import_plan(&export).unwrap();
assert!(plan.manifest.dashboard_snapshot.is_some());
assert!(!plan
.rows(ExportDomain::Auxiliary)
.iter()
.any(|row| row.payload["__table"] == "dashboard_stats_pending"));
sqlx::query("CREATE DATABASE dashboard_restore_test")
.execute(&source)
.await
.unwrap();
let target_url = server
.database_url()
.strip_suffix("/postgres")
.unwrap()
.to_owned()
+ "/dashboard_restore_test";
let target = sqlx::postgres::PgPoolOptions::new()
.max_connections(1)
.connect(&target_url)
.await
.unwrap();
migrate(&target).await;
// The gateway initializes one verified local admin and a zero-value unlimited
// wallet before the operator can run restore. Preserve that account while
// replacing its one initialization event with the source statistics.
sqlx::raw_sql("INSERT INTO users(id,username,email_verified,role,auth_source,is_active) VALUES('bootstrap-admin','custom-root-name',true,'admin','local',true); INSERT INTO wallets(id,user_id,limit_mode,currency,status,created_at,updated_at) VALUES('bootstrap-wallet','bootstrap-admin','unlimited','USD','active',clock_timestamp(),clock_timestamp());")
.execute(&target).await.unwrap();
sqlx::query("UPDATE wallets SET balance=1,total_recharged=1 WHERE id='bootstrap-wallet'")
.execute(&target)
.await
.unwrap();
let error = import_postgres_jsonl(&target, &export).await.unwrap_err();
assert!(
error
.to_string()
.contains("conflicts with existing statistics"),
"{error}"
);
sqlx::query("UPDATE wallets SET balance=0,total_recharged=0 WHERE id='bootstrap-wallet'")
.execute(&target)
.await
.unwrap();
import_postgres_jsonl(&target, &export).await.unwrap();
assert_eq!(
sqlx::query_scalar::<_, i64>("SELECT count(*) FROM users WHERE id='bootstrap-admin'")
.fetch_one(&target)
.await
.unwrap(),
1,
"the initialized admin account is not removed"
);
assert_eq!(dashboard_rows(&target).await,expected,"restore includes exact activation timestamp, narrow history, user events and purged request counts");
assert_eq!(
sqlx::query_scalar::<_, i64>("SELECT count(*) FROM usage")
.fetch_one(&target)
.await
.unwrap(),
1
);
import_postgres_jsonl(&target, &export).await.unwrap();
assert_eq!(
dashboard_rows(&target).await,
expected,
"repeated import must not count users or requests twice"
);
let mut invalid_source = decode_jsonl(&export).unwrap();
for record in &mut invalid_source {
if let DataExportRecord::Row {
domain: ExportDomain::Users,
payload,
..
} = record
{
payload["auth_source"] = Value::String("invalid-auth-source".into());
}
}
assert!(
import_postgres_jsonl(&target, &encode_jsonl(&invalid_source).unwrap())
.await
.is_err()
);
assert_eq!(
dashboard_rows(&target).await,
expected,
"failure after the restore guard and deletes rolls the entire transaction back"
);
let guard: Option<String> =
sqlx::query_scalar("SELECT current_setting('aether.dashboard_restore',true)")
.fetch_one(&target)
.await
.unwrap();
assert_ne!(guard.as_deref(), Some("on"));
sqlx::query("UPDATE usage SET input_tokens=17,total_tokens=20 WHERE request_id='backup-kept'")
.execute(&target)
.await
.unwrap();
assert_eq!(
total(&target).await,
2,
"existing restored contribution updates by delta"
);
assert_eq!(tokens(&target).await, 46);
request(&target, "backup-next", 31).await;
assert_eq!(
total(&target).await,
3,
"future writes resume ordinary aggregation"
);
assert_eq!(tokens(&target).await, 80);
let before_conflict = dashboard_rows(&target).await;
let error = import_postgres_jsonl(&target, &export).await.unwrap_err();
assert!(
error
.to_string()
.contains("conflicts with existing statistics"),
"{error}"
);
assert_eq!(
dashboard_rows(&target).await,
before_conflict,
"conflicting snapshots must roll back without overwriting accumulated data"
);
let mut truncated = decode_jsonl(&export).unwrap();
truncated.retain(|row| !matches!(row,DataExportRecord::Row { payload,.. } if payload["__table"]=="dashboard_stats_total" && payload["shard"]==serde_json::json!(1)));
let error = import_postgres_jsonl(&target, &encode_jsonl(&truncated).unwrap())
.await
.unwrap_err();
assert!(error.to_string().contains("missing, truncated"), "{error}");
// A legacy file has neither aggregate rows nor the optional snapshot manifest.
// Its ordinary source import must still trigger aggregation.
let mut legacy = decode_jsonl(&export).unwrap();
if let DataExportRecord::Manifest { manifest } = &mut legacy[0] {
manifest.dashboard_snapshot = None;
}
legacy.retain(|row| !matches!(row,DataExportRecord::Row { payload,.. } if payload["__table"].as_str().is_some_and(|name| name.starts_with("dashboard_"))));
sqlx::query("CREATE DATABASE dashboard_legacy_restore_test")
.execute(&source)
.await
.unwrap();
let legacy_url = server
.database_url()
.strip_suffix("/postgres")
.unwrap()
.to_owned()
+ "/dashboard_legacy_restore_test";
let legacy_target = sqlx::postgres::PgPoolOptions::new()
.max_connections(2)
.connect(&legacy_url)
.await
.unwrap();
migrate(&legacy_target).await;
sqlx::query("UPDATE dashboard_stats_state SET stats_since='2000-01-01'")
.execute(&legacy_target)
.await
.unwrap();
let legacy_export = encode_jsonl(&legacy).unwrap();
import_postgres_jsonl(&legacy_target, &legacy_export)
.await
.unwrap();
assert_eq!(total(&legacy_target).await, 1);
import_postgres_jsonl(&legacy_target, &legacy_export)
.await
.unwrap();
assert_eq!(total(&legacy_target).await, 1);
legacy_target.close().await;
target.close().await;
source.close().await;
}
crates/aether-data/runtime/src/lifecycle/export/postgres.rs
+12
View File
@@ -1,5 +1,7 @@
use super::*;
mod dashboard_snapshot;
pub async fn export_postgres_core_jsonl(
pool: &crate::driver::postgres::PostgresPool,
created_at_unix_secs: u64,
@@ -17,6 +19,10 @@ pub async fn export_postgres_jsonl(
.execute(&mut *tx)
.await
.map_sql_err()?;
sqlx::query("SET LOCAL TIME ZONE 'UTC'")
.execute(&mut *tx)
.await
.map_sql_err()?;
let manifest = DataExportManifest::new(
created_at_unix_secs,
Some(DatabaseDriver::Postgres),
@@ -51,6 +57,7 @@ pub async fn export_postgres_jsonl(
}
}
dashboard_snapshot::attach_manifest(&mut records)?;
tx.commit().await.map_sql_err()?;
encode_jsonl(&records)
}
@@ -85,6 +92,11 @@ async fn import_postgres_plan_with_options(
) -> Result<usize, DataLayerError> {
let identity_scope = IdentityImportScope::from_plan(plan)?;
let mut tx = pool.begin().await.map_sql_err()?;
sqlx::query("SET LOCAL TIME ZONE 'UTC'")
.execute(&mut *tx)
.await
.map_sql_err()?;
dashboard_snapshot::prepare_restore(&mut tx, plan).await?;
let identity_state = capture_postgres_identity_import_state(&mut tx, &identity_scope).await?;
let mut imported = 0usize;
let mut column_cache = BTreeMap::<String, PostgresImportColumns>::new();
crates/aether-data/runtime/src/lifecycle/export/postgres/dashboard_snapshot.rs
+289
View File
@@ -0,0 +1,289 @@
use super::*;
fn tables() -> impl Iterator<Item = &'static AuxiliaryTable> {
AUXILIARY_TABLES
.iter()
.filter(|table| table.name.starts_with("dashboard_"))
}
fn fingerprint(rows: &[Value]) -> DashboardSnapshotTable {
// Sorting canonical JSON makes fingerprints independent of file row order.
let mut encoded = rows
.iter()
.map(|row| {
let mut canonical = row.clone();
canonical.sort_all_objects();
canonical.to_string()
})
.collect::<Vec<_>>();
encoded.sort();
let mut digest = Sha256::new();
for row in encoded {
digest.update(row.as_bytes());
digest.update(b"\n");
}
DashboardSnapshotTable {
rows: rows.len(),
sha256: format!("{:x}", digest.finalize()),
}
}
fn snapshot_rows<'a>(
rows: impl Iterator<Item = &'a ExportRow>,
) -> Result<BTreeMap<String, Vec<Value>>, DataLayerError> {
let mut result = tables()
.map(|table| (table.name.to_owned(), Vec::new()))
.collect::<BTreeMap<_, _>>();
for row in rows {
let (name, payload) = domain_payload_table(row, "auxiliary", None)?;
if let Some(entries) = result.get_mut(&name) {
entries.push(payload);
}
}
Ok(result)
}
pub(super) fn attach_manifest(records: &mut [DataExportRecord]) -> Result<(), DataLayerError> {
let Some(DataExportRecord::Manifest { manifest }) = records.first() else {
return Ok(());
};
if !manifest.domains.contains(&ExportDomain::Auxiliary) {
return Ok(());
}
let rows = records
.iter()
.filter_map(|record| match record {
DataExportRecord::Row {
domain: ExportDomain::Auxiliary,
id,
payload,
} => Some(ExportRow {
id: id.clone(),
payload: payload.clone(),
}),
_ => None,
})
.collect::<Vec<_>>();
let values = snapshot_rows(rows.iter())?;
let tables = values
.into_iter()
.map(|(name, values)| (name, fingerprint(&values)))
.collect();
if let Some(DataExportRecord::Manifest { manifest }) = records.first_mut() {
manifest.dashboard_snapshot = Some(DashboardSnapshotManifest { version: 1, tables });
}
Ok(())
}
fn invalid(detail: &str) -> DataLayerError {
DataLayerError::InvalidInput(format!("dashboard snapshot {detail}"))
}
pub(super) async fn prepare_restore(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
plan: &DataImportPlan,
) -> Result<(), DataLayerError> {
let values = snapshot_rows(plan.rows(ExportDomain::Auxiliary).iter())?;
let Some(manifest) = &plan.manifest.dashboard_snapshot else {
if values.values().any(|rows| !rows.is_empty()) {
return Err(invalid(
"requires its complete manifest; partial aggregate imports cannot be merged",
));
}
return Ok(()); // Legacy backups deliberately keep normal trigger behavior.
};
if manifest.version != 1
|| !plan.imports_domain(ExportDomain::Auxiliary)
|| manifest.tables.len() != values.len()
{
return Err(invalid(
"has an unsupported version or incomplete table inventory",
));
}
for (name, rows) in &values {
if manifest.tables.get(name) != Some(&fingerprint(rows)) {
return Err(invalid(&format!(
"table '{name}' is missing, truncated, or has changed"
)));
}
}
let state = &values["dashboard_stats_state"];
let totals = &values["dashboard_stats_total"];
let shards = totals
.iter()
.filter_map(|row| row.get("shard").and_then(Value::as_u64))
.collect::<BTreeSet<_>>();
if state.len() != 1
|| state[0].get("singleton") != Some(&Value::Bool(true))
|| state[0]
.get("stats_since")
.and_then(Value::as_str)
.and_then(parse_imported_datetime)
.is_none()
|| totals.len() != 16
|| shards != (0..16).collect()
{
return Err(invalid(
"must include one activation state and all 16 total shards",
));
}
validate_request_counts(&values)?;
// Exclude concurrent source writes before examining or replacing projections.
// Their triggers acquire projection locks in this same source-first order.
sqlx::query("LOCK TABLE public.users, public.usage, public.usage_settlement_snapshots, public.usage_attribution_snapshots IN SHARE ROW EXCLUSIVE MODE")
.execute(&mut **tx).await.map_sql_err()?;
// Match retention's shard -> minute -> activity -> actor -> event -> state -> ledger
// order, otherwise a maintenance pass could deadlock against the restore.
for table in [
"dashboard_stats_total",
"dashboard_stats_minute",
"dashboard_activity_minute",
"dashboard_actor_minute",
"dashboard_user_events_minute",
"dashboard_stats_state",
"dashboard_request_contributions",
"dashboard_activity_hour",
] {
sqlx::query(&format!(
"LOCK TABLE public.{table} IN SHARE ROW EXCLUSIVE MODE"
))
.execute(&mut **tx)
.await
.map_sql_err()?;
}
let bootstrap_only = bootstrap_admin_only(tx).await?;
let mut identical = true;
let mut empty = true;
for table in tables() {
let mut current = sqlx::query_scalar::<_, Value>(&format!(
"SELECT to_jsonb(t) FROM public.{} t",
table.name
))
.fetch_all(&mut **tx)
.await
.map_sql_err()?;
// The bounded cleanup cursor is operational progress, not a change to
// statistics; moving it alone must not make a repeated restore conflict.
if table.name == "dashboard_stats_state" {
for row in &mut current {
if let Some(object) = row.as_object_mut() {
if let Some(cursor) = state[0].get("contributions_cleanup_cursor") {
object.insert("contributions_cleanup_cursor".into(), cursor.clone());
}
}
}
}
identical &= manifest.tables.get(table.name) == Some(&fingerprint(&current));
empty &= match table.name {
"dashboard_stats_state" => true, // A freshly migrated database already has an activation timestamp.
"dashboard_stats_total" => current.iter().all(|row| {
row.get("metrics")
.and_then(Value::as_object)
.is_some_and(|metrics| metrics.values().all(|n| n.as_f64() == Some(0.0)))
}),
"dashboard_user_events_minute" => current.is_empty() || bootstrap_only,
_ => current.is_empty(),
};
}
if !empty && !identical {
return Err(invalid("conflicts with existing statistics; restore into an empty database. Complete statistics cannot be incrementally merged"));
}
if empty && !identical {
let since =
parse_imported_datetime(state[0]["stats_since"].as_str().expect("validated state"))
.expect("validated timestamp");
if sqlx::query_scalar::<_, bool>(
"SELECT EXISTS(SELECT 1 FROM public.usage WHERE created_at >= $1)",
)
.bind(since)
.fetch_one(&mut **tx)
.await
.map_sql_err()?
{
return Err(invalid("conflicts with existing requests within the restored activation period; restore into an empty database"));
}
}
// A local custom setting affects only dashboard triggers, not integrity or
// billing triggers, and is automatically reverted on both commit and rollback.
sqlx::query("SET LOCAL aether.dashboard_restore = 'on'")
.execute(&mut **tx)
.await
.map_sql_err()?;
for table in tables() {
sqlx::query(&format!("DELETE FROM public.{}", table.name))
.execute(&mut **tx)
.await
.map_sql_err()?;
}
Ok(())
}
async fn bootstrap_admin_only(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
) -> Result<bool, DataLayerError> {
// bootstrap_admin_from_env has no persistent marker and permits a configured
// username. Recognize only its otherwise untouched single-admin/zero-wallet
// state, never an installation with financial or request history.
sqlx::query("LOCK TABLE public.api_keys, public.wallets, public.payment_orders, public.wallet_transactions, public.provider_expenses IN SHARE ROW EXCLUSIVE MODE")
.execute(&mut **tx).await.map_sql_err()?;
sqlx::query_scalar(r#"
SELECT (SELECT count(*) FROM users)=1
AND (SELECT count(*) FROM wallets)=1
AND (SELECT count(*) FROM dashboard_user_events_minute)=1
AND NOT EXISTS(SELECT 1 FROM usage)
AND NOT EXISTS(SELECT 1 FROM api_keys)
AND NOT EXISTS(SELECT 1 FROM payment_orders)
AND NOT EXISTS(SELECT 1 FROM wallet_transactions)
AND NOT EXISTS(SELECT 1 FROM provider_expenses)
AND EXISTS(
SELECT 1 FROM users u JOIN wallets w ON w.user_id=u.id
JOIN dashboard_user_events_minute e
ON e.bucket_start=date_trunc('minute',u.created_at AT TIME ZONE 'UTC') AT TIME ZONE 'UTC'
AND e.shard=(hashtextextended(u.id,0)&15)::smallint
JOIN dashboard_stats_state s ON s.singleton
WHERE u.role='admin' AND u.auth_source='local' AND u.is_active
AND NOT u.is_deleted AND u.email_verified AND u.created_at>=s.stats_since
AND w.api_key_id IS NULL AND w.limit_mode='unlimited'
AND w.currency='USD' AND w.status='active'
AND w.balance=0 AND w.gift_balance=0 AND w.total_recharged=0
AND w.total_consumed=0 AND w.total_refunded=0 AND w.total_adjusted=0
AND e.created_count=1 AND e.deleted_count=0
)
"#).fetch_one(&mut **tx).await.map_sql_err()
}
fn validate_request_counts(values: &BTreeMap<String, Vec<Value>>) -> Result<(), DataLayerError> {
let mut totals = [0u64; 16];
for row in &values["dashboard_stats_total"] {
let shard = row["shard"]
.as_u64()
.ok_or_else(|| invalid("has an invalid shard"))? as usize;
totals[shard] = request_count(&row["metrics"]["request_count"])?;
}
let mut hours = [0u64; 16];
for row in &values["dashboard_activity_hour"] {
let shard = row["shard"]
.as_u64()
.filter(|n| *n < 16)
.ok_or_else(|| invalid("has an invalid activity shard"))? as usize;
hours[shard] = hours[shard]
.checked_add(request_count(&row["request_count"])?)
.ok_or_else(|| invalid("activity counts overflow"))?;
}
if hours != totals {
return Err(invalid(
"hourly activity and cumulative request counts disagree",
));
}
Ok(())
}
fn request_count(value: &Value) -> Result<u64, DataLayerError> {
if value.is_null() {
return Ok(0);
}
value
.as_u64()
.ok_or_else(|| invalid("contains an invalid request count"))
}
crates/aether-data/runtime/src/lifecycle/migrate/tests.rs
+258 -57
View File
@@ -28,6 +28,12 @@ use crate::lifecycle::bootstrap::postgres::{
};
mod policy_nulls;
mod overview_dirty_events;
mod provider_expenses;
mod migration_deadlines;
mod overview_migration_safety;
mod legacy_overview_upgrade;
mod dashboard_user_anonymization;
/// A clean PostgreSQL database is bootstrapped from the schema snapshot first;
/// migrations after the privacy/security frontier are intentionally left
@@ -1575,6 +1581,18 @@ fn pending_migrations_from_applied_skips_versions_already_applied() {
20260901000000,
20260903000000,
20260908000000,
20260911000000,
20260917000000,
20260917000100,
20260918000000,
20260918000100,
20260919000000,
20260920000000,
20260920120000,
20260921010000,
20260921020000,
20260921020100,
20261001000000,
]
);
}
@@ -1875,6 +1893,169 @@ WHERE id = 'metadata-migration-key'
.expect("provider migration fixture should clean up");
}
#[tokio::test]
async fn overview_migrations_support_fresh_and_legacy_install() {
for legacy in [false, true] {
let Some(server) = ManagedPostgresServer::try_start()
.await
.expect("overview PostgreSQL should start or skip")
else {
return;
};
if legacy {
let mut connection = PgConnection::connect(server.database_url()).await.unwrap();
connection.ensure_migrations_table().await.unwrap();
for migration in POSTGRES_MIGRATOR
.iter()
.filter(|migration| migration.version < 20260911000000)
{
connection.apply(migration).await.unwrap();
}
sqlx::raw_sql(
r#"
INSERT INTO users(id, username, email_verified)
VALUES ('overview-legacy-owner', 'overview-legacy-owner', false);
INSERT INTO api_keys(id, user_id, key_hash)
VALUES ('overview-legacy-key', 'overview-legacy-owner', repeat('e', 64));
INSERT INTO usage(id, request_id, user_id, api_key_id, provider_name, model,
status, billing_status, created_at)
VALUES ('overview-legacy-request', 'overview-legacy-request',
'overview-legacy-owner', 'overview-legacy-key', 'test', 'test',
'completed', 'settled', '2020-01-01 12:34:00+00');
"#,
)
.execute(&mut connection)
.await
.unwrap();
}
let pool = PgPool::connect(server.database_url()).await.unwrap();
prepare_and_apply_clean_postgres_database(&pool).await;
for table in [
"usage_attribution_snapshots",
"stats_bucket_state",
"stats_overview_hourly",
"stats_overview_daily",
] {
assert!(table_exists(&pool, table).await.unwrap());
}
assert!(!column_exists(&pool, "api_keys", "credential_kind")
.await
.unwrap());
assert!(
!column_exists(&pool, "usage_attribution_snapshots", "credential_kind")
.await
.unwrap()
);
assert!(column_exists(&pool, "stats_bucket_state", "last_failed_at")
.await
.unwrap());
let precision:(i32,i32)=sqlx::query_as("SELECT numeric_precision::integer,numeric_scale::integer FROM information_schema.columns WHERE table_schema='public' AND table_name='usage_settlement_snapshots' AND column_name='wallet_debit_amount_usd'").fetch_one(&pool).await.unwrap();
assert_eq!(precision, (20, 8));
if legacy {
let facts:(Option<String>,Option<String>,String,Option<String>)=sqlx::query_as("SELECT actor_user_id,credential_owner_id,attribution_source,wallet_debit_amount::text FROM usage_analytics_facts_v1 WHERE request_id='overview-legacy-request'").fetch_one(&pool).await.unwrap();
assert_eq!(
facts,
(
Some("overview-legacy-owner".into()),
Some("overview-legacy-owner".into()),
"user_account".into(),
None
)
);
}
let repo = aether_data_postgres::SqlxUsageReadRepository::new(pool.clone());
let target = chrono::DateTime::parse_from_rfc3339("2026-09-17T00:00:00Z")
.unwrap()
.with_timezone(&chrono::Utc);
for _ in 0..2 {
let counts: (i64, i64, i64, i64) = sqlx::query_as(
r#"
SELECT (SELECT COUNT(*) FROM usage_attribution_snapshots),
(SELECT COUNT(*) FROM stats_bucket_state),
(SELECT COUNT(*) FROM stats_overview_hourly),
(SELECT COUNT(*) FROM stats_overview_daily)
"#,
)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(counts, (0, 0, 0, 0));
assert_eq!(repo.rebuild_overview_buckets(target, 8).await.unwrap(), 0);
super::run_migrations(&pool).await.unwrap();
}
sqlx::raw_sql(
r#"
INSERT INTO users(id, username, email_verified)
VALUES ('overview-new-owner', 'overview-new-owner', false);
INSERT INTO api_keys(id, user_id, key_hash)
VALUES ('overview-new-key', 'overview-new-owner', repeat('f', 64));
INSERT INTO usage(id, request_id, user_id, api_key_id, provider_name, model,
status, billing_status, created_at, request_metadata)
VALUES ('overview-new-request', 'overview-new-request',
'overview-new-owner', 'overview-new-key', 'test', 'test',
'completed', 'settled', '2026-09-15 12:34:00+00',
'{"analytics_attribution":{"is_standalone":false}}'::json);
"#,
)
.execute(&pool)
.await
.unwrap();
let identity: (String, String, String, String) = sqlx::query_as(
r#"
SELECT actor_user_id, credential_owner_id, attribution_kind,
attribution_source
FROM usage_attribution_snapshots WHERE request_id = 'overview-new-request'
"#,
)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(
identity,
(
"overview-new-owner".into(),
"overview-new-owner".into(),
"employee".into(),
"user_account".into(),
)
);
assert_eq!(repo.rebuild_overview_buckets(target, 8).await.unwrap(), 2);
assert_eq!(repo.rebuild_overview_buckets(target, 8).await.unwrap(), 0);
let counts: (i64, i64, i64, i64, i64) = sqlx::query_as(
r#"
SELECT (SELECT COUNT(*) FROM usage_attribution_snapshots),
(SELECT COUNT(*) FROM usage_attribution_snapshots
WHERE request_id = 'overview-legacy-request'),
(SELECT COUNT(*) FROM stats_bucket_state),
(SELECT COUNT(*) FROM stats_overview_hourly),
(SELECT COUNT(*) FROM stats_overview_daily)
"#,
)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(counts, (1, 0, 2, 1, 1));
let unexpected_bucket_count: i64 = query_scalar(
r#"
SELECT COUNT(*) FROM stats_bucket_state
WHERE projection_version <> 'overview-v2'
OR (granularity, bucket_start) NOT IN (
('day', '2026-09-15 00:00:00+00'::timestamptz),
('hour', '2026-09-15 12:00:00+00'::timestamptz)
)
OR coverage_status <> 'complete'
OR source_revision <> built_revision
"#,
)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(unexpected_bucket_count, 0);
pool.close().await;
}
}
#[tokio::test]
async fn prepare_database_for_startup_bootstraps_clean_database() {
let Some(server) = ManagedPostgresServer::try_start()
@@ -2810,9 +2991,7 @@ WHERE request_id = 'billing-facts-cache-create'
}
#[tokio::test]
async fn postgres_migrations_repair_invalid_concurrent_cleanup_index() {
const MIGRATION_VERSION: i64 = 20260715000000;
async fn postgres_migrations_repair_invalid_concurrent_indexes() {
let Some(server) = ManagedPostgresServer::try_start()
.await
.expect("postgres migration retry test should start or skip")
@@ -2824,11 +3003,6 @@ async fn postgres_migrations_repair_invalid_concurrent_cleanup_index() {
.await
.expect("pool should connect");
prepare_and_apply_clean_postgres_database(&pool).await;
query("DROP INDEX CONCURRENTLY public.idx_usage_legacy_body_ref_cleanup_created_at")
.execute(&pool)
.await
.expect("snapshot cleanup index should exist");
query("CREATE TABLE public.concurrent_index_failure_fixture (value integer NOT NULL)")
.execute(&pool)
.await
@@ -2838,62 +3012,89 @@ async fn postgres_migrations_repair_invalid_concurrent_cleanup_index() {
.await
.expect("duplicate failure fixtures should be inserted");
query(
"CREATE UNIQUE INDEX CONCURRENTLY idx_usage_legacy_body_ref_cleanup_created_at ON public.concurrent_index_failure_fixture (value)",
)
.execute(&pool)
.await
.expect_err("duplicate values should leave a failed concurrent index build");
let invalid_index_exists: bool = query_scalar(
r#"
for (migration_version, index_name, table_name) in [
(
20260715000000_i64,
"idx_usage_legacy_body_ref_cleanup_created_at",
"public.usage",
),
(
20260918000000,
"idx_usage_settlement_dashboard_cover_v2",
"public.usage_settlement_snapshots",
),
(
20260920000000,
"idx_payment_orders_status_credited_user",
"public.payment_orders",
),
(
20260921020100,
"ix_usage_analytics_actor_metadata",
"public.usage",
),
] {
query(&format!("DROP INDEX CONCURRENTLY public.{index_name}"))
.execute(&pool)
.await
.expect("migrated index should exist");
query(&format!("CREATE UNIQUE INDEX CONCURRENTLY {index_name} ON public.concurrent_index_failure_fixture (value)"))
.execute(&pool)
.await
.expect_err("duplicate values should leave a failed concurrent index build");
let invalid_index_exists: bool = query_scalar(
r#"
SELECT EXISTS (
SELECT 1
FROM pg_catalog.pg_class AS index_relation
JOIN pg_catalog.pg_namespace AS index_namespace
ON index_namespace.oid = index_relation.relnamespace
JOIN pg_catalog.pg_index AS index_state
ON index_state.indexrelid = index_relation.oid
WHERE index_namespace.nspname = 'public'
AND index_relation.relname = 'idx_usage_legacy_body_ref_cleanup_created_at'
AND NOT index_state.indisvalid
)
"#,
SELECT 1 FROM pg_catalog.pg_index
WHERE indexrelid = to_regclass($1) AND NOT indisvalid
)"#,
)
.bind(format!("public.{index_name}"))
.fetch_one(&pool)
.await
.expect("failed concurrent index state should be readable");
assert!(invalid_index_exists);
query("DELETE FROM public._sqlx_migrations WHERE version = $1")
.bind(migration_version)
.execute(&pool)
.await
.expect("index migration stamp should be reset");
super::run_migrations(&pool)
.await
.expect("migration retry should replace the invalid index");
let valid_index_exists: bool = query_scalar(
r#"
SELECT EXISTS (
SELECT 1 FROM pg_catalog.pg_index
WHERE indexrelid = to_regclass($1) AND indrelid = $2::regclass AND indisvalid
)"#,
)
.bind(format!("public.{index_name}"))
.bind(table_name)
.fetch_one(&pool)
.await
.expect("rebuilt index state should be readable");
assert!(valid_index_exists);
}
let index_definition: String = query_scalar(
"SELECT pg_get_indexdef('public.idx_usage_settlement_dashboard_cover_v2'::regclass)",
)
.fetch_one(&pool)
.await
.expect("failed concurrent index state should be readable");
assert!(invalid_index_exists);
query("DELETE FROM public._sqlx_migrations WHERE version = $1")
.bind(MIGRATION_VERSION)
.execute(&pool)
.await
.expect("cleanup index migration stamp should be reset");
super::run_migrations(&pool)
.await
.expect("migration retry should replace the invalid index");
let valid_usage_index_exists: bool = query_scalar(
r#"
SELECT EXISTS (
SELECT 1
FROM pg_catalog.pg_class AS index_relation
JOIN pg_catalog.pg_namespace AS index_namespace
ON index_namespace.oid = index_relation.relnamespace
JOIN pg_catalog.pg_index AS index_state
ON index_state.indexrelid = index_relation.oid
WHERE index_namespace.nspname = 'public'
AND index_relation.relname = 'idx_usage_legacy_body_ref_cleanup_created_at'
AND index_state.indrelid = 'public.usage'::regclass
AND index_state.indisvalid
)
"#,
.expect("replacement index should be installed");
assert!(index_definition.contains("billing_status"));
assert!(index_definition.contains("allocation_status"));
let old_cover_exists: bool = query_scalar(
"SELECT to_regclass('public.idx_usage_settlement_dashboard_cover') IS NOT NULL",
)
.fetch_one(&pool)
.await
.expect("rebuilt cleanup index state should be readable");
assert!(valid_usage_index_exists);
.expect("old covering index state should be readable");
assert!(
!old_cover_exists,
"successful migration must retire the redundant old cover"
);
}
#[tokio::test]
crates/aether-data/runtime/src/lifecycle/migrate/tests/dashboard_user_anonymization.rs
+231
View File
@@ -0,0 +1,231 @@
use super::*;
use aether_data_contracts::repository::usage::UsageDashboardAnalyticsQuery;
const ANONYMIZATION_VERSION: i64 = 20261001000000;
const INSERT_USER: &str = "INSERT INTO users(id,username,email_verified) VALUES($1,$1,false)";
const INSERT_USAGE: &str = "INSERT INTO usage(id,request_id,user_id,model,provider_name,status,billing_status,created_at) VALUES($1,$1,$2,'anonymization','test','completed','settled',clock_timestamp())";
async fn assert_anonymous(pool: &PgPool, user: &str) {
let counts: (i64, i64) = sqlx::query_as(
"SELECT (SELECT count(*) FROM dashboard_actor_minute WHERE actor_user_id=$1), (SELECT count(*) FROM dashboard_request_contributions WHERE actor_user_id=$1)",
)
.bind(user)
.fetch_one(pool)
.await
.unwrap();
assert_eq!(counts, (0, 0), "retained identity for {user}");
}
#[tokio::test]
async fn dashboard_user_anonymization_preserves_totals_without_migration_backfill() {
let Some(server) = ManagedPostgresServer::try_start().await.unwrap() else {
return;
};
let mut connection = PgConnection::connect(server.database_url()).await.unwrap();
connection.ensure_migrations_table().await.unwrap();
for migration in POSTGRES_MIGRATOR
.iter()
.filter(|migration| migration.version < ANONYMIZATION_VERSION)
{
connection.apply(migration).await.unwrap();
}
let pool = PgPool::connect(server.database_url()).await.unwrap();
query(INSERT_USER)
.bind("legacy-deleted")
.execute(&pool)
.await
.unwrap();
query(INSERT_USAGE)
.bind("legacy-request")
.bind("legacy-deleted")
.execute(&pool)
.await
.unwrap();
query("DELETE FROM usage WHERE request_id='legacy-request'")
.execute(&pool)
.await
.unwrap();
query("DELETE FROM users WHERE id='legacy-deleted'")
.execute(&pool)
.await
.unwrap();
let actor_before: String =
query_scalar("SELECT jsonb_agg(to_jsonb(a))::text FROM dashboard_actor_minute a")
.fetch_one(&pool)
.await
.unwrap();
// An upgrade must succeed even while historical projection tables are
// inaccessible: install definitions without reading or rewriting their rows.
let mut blocked_history = pool.begin().await.unwrap();
query("LOCK TABLE dashboard_actor_minute, dashboard_request_contributions IN ACCESS EXCLUSIVE MODE")
.execute(&mut *blocked_history).await.unwrap();
query("SET lock_timeout='500ms'")
.execute(&mut connection)
.await
.unwrap();
let migration = POSTGRES_MIGRATOR
.iter()
.find(|migration| migration.version == ANONYMIZATION_VERSION)
.unwrap();
connection.apply(migration).await.unwrap();
blocked_history.rollback().await.unwrap();
let actor_after: String =
query_scalar("SELECT jsonb_agg(to_jsonb(a))::text FROM dashboard_actor_minute a")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(
actor_before, actor_after,
"migration must not rewrite old actors"
);
let repo = aether_data_postgres::SqlxUsageReadRepository::new(pool.clone());
let summary_query = UsageDashboardAnalyticsQuery {
timezone: "UTC".into(),
};
let legacy = repo.query_dashboard_summary(&summary_query).await.unwrap();
assert_eq!(
legacy.today.active_users, 0,
"old orphan actors must be excluded without backfill"
);
assert_eq!(legacy.total.request_count, 1);
for (user, purge, soft) in [
("hard-live", false, false),
("hard-purged", true, false),
("soft-live", false, true),
("soft-purged", true, true),
] {
query(INSERT_USER).bind(user).execute(&pool).await.unwrap();
query(INSERT_USAGE)
.bind(user)
.bind(user)
.execute(&pool)
.await
.unwrap();
if purge {
query("DELETE FROM usage WHERE request_id=$1")
.bind(user)
.execute(&pool)
.await
.unwrap();
}
let before = repo.query_dashboard_summary(&summary_query).await.unwrap();
let sql = if soft {
"UPDATE users SET is_deleted=true WHERE id=$1"
} else {
"DELETE FROM users WHERE id=$1"
};
query(sql).bind(user).execute(&pool).await.unwrap();
assert_anonymous(&pool, user).await;
let after = repo.query_dashboard_summary(&summary_query).await.unwrap();
assert_eq!(
after.total, before.total,
"deletion must preserve request totals"
);
assert_eq!(after.today.active_users, 0);
if !purge {
query("UPDATE usage SET response_time_ms=200 WHERE request_id=$1")
.bind(user)
.execute(&pool)
.await
.unwrap();
assert_anonymous(&pool, user).await;
}
if soft {
query("DELETE FROM users WHERE id=$1")
.bind(user)
.execute(&pool)
.await
.unwrap();
assert_anonymous(&pool, user).await;
}
}
// Exercise both deferred-event orders, plus explicitly immediate constraint
// triggers. A captured deleted_fact must never recreate a deleted actor.
for (user, user_first, immediate) in [
("deferred-request-first", false, false),
("deferred-user-first", true, false),
("immediate-user", true, true),
] {
query(INSERT_USER).bind(user).execute(&pool).await.unwrap();
let mut tx = pool.begin().await.unwrap();
if immediate {
query("SET CONSTRAINTS ALL IMMEDIATE")
.execute(&mut *tx)
.await
.unwrap();
}
query(INSERT_USAGE)
.bind(user)
.bind(user)
.execute(&mut *tx)
.await
.unwrap();
if user_first {
query("DELETE FROM users WHERE id=$1")
.bind(user)
.execute(&mut *tx)
.await
.unwrap();
}
query("DELETE FROM usage WHERE request_id=$1")
.bind(user)
.execute(&mut *tx)
.await
.unwrap();
if !user_first {
query("DELETE FROM users WHERE id=$1")
.bind(user)
.execute(&mut *tx)
.await
.unwrap();
}
tx.commit().await.unwrap();
assert_anonymous(&pool, user).await;
}
// A concurrent writer can still see a user after its deletion statement
// but before that transaction commits. Check both commit orders: the new
// actor must be rejected after the shard lock, or removed by the deleter.
for (user, writer_first) in [
("concurrent-delete-first", false),
("concurrent-writer-first", true),
] {
query(INSERT_USER).bind(user).execute(&pool).await.unwrap();
let mut deletion = pool.begin().await.unwrap();
query("DELETE FROM users WHERE id=$1")
.bind(user)
.execute(&mut *deletion)
.await
.unwrap();
let mut writer = pool.begin().await.unwrap();
query(INSERT_USAGE)
.bind(user)
.bind(user)
.execute(&mut *writer)
.await
.unwrap();
tokio::time::timeout(std::time::Duration::from_secs(5), async {
if writer_first {
writer.commit().await.unwrap();
deletion.commit().await.unwrap();
} else {
deletion.commit().await.unwrap();
writer.commit().await.unwrap();
}
})
.await
.expect("concurrent deletion and usage must not deadlock");
assert_anonymous(&pool, user).await;
}
let final_summary = repo.query_dashboard_summary(&summary_query).await.unwrap();
assert_eq!(final_summary.total.request_count, 10);
assert_eq!(final_summary.today.active_users, 0);
let invalid: (i64, i64, i64) = sqlx::query_as("SELECT (SELECT count(*) FROM dashboard_actor_minute WHERE request_count < 0), (SELECT count(*) FROM dashboard_stats_pending), (SELECT count(*) FROM dashboard_user_anonymization_pending)")
.fetch_one(&pool).await.unwrap();
assert_eq!(invalid, (0, 0, 0));
pool.close().await;
}
crates/aether-data/runtime/src/lifecycle/migrate/tests/legacy_overview_upgrade.rs
+211
View File
@@ -0,0 +1,211 @@
use super::*;
const ACCOUNT_ATTRIBUTION: i64 = 20260917000000;
const DIRTY_EVENTS: i64 = 20260917000100;
async fn rows_snapshot(pool: &PgPool, table: &str) -> String {
query_scalar(&format!(
"SELECT COALESCE(jsonb_agg(to_jsonb(t) ORDER BY to_jsonb(t)::text), '[]')::text FROM {table} t"
))
.fetch_one(pool)
.await
.unwrap()
}
#[tokio::test]
async fn legacy_overview_upgrade_preserves_applied_history_and_existing_statistics() {
let Some(server) = ManagedPostgresServer::try_start().await.unwrap() else {
return;
};
let mut connection = PgConnection::connect(server.database_url()).await.unwrap();
connection.ensure_migrations_table().await.unwrap();
for migration in POSTGRES_MIGRATOR.iter().filter(|migration| {
migration.version <= 20260920120000 && migration.version != DIRTY_EVENTS
}) {
connection.apply(migration).await.unwrap();
}
connection.close().await.unwrap();
let pool = PgPool::connect(server.database_url()).await.unwrap();
// The restored September 17 migration installs the historical direct-write
// trigger. Remove schema additions folded into the September 11 baseline
// so this exercises an already-running database that never received them.
sqlx::raw_sql(
r#"
DROP TABLE public.stats_overview_dirty_events;
DROP INDEX public.ix_usage_attribution_owner_request;
UPDATE public.dashboard_stats_state SET stats_since = clock_timestamp() - INTERVAL '1 day';
INSERT INTO users(id, username, email_verified)
VALUES ('legacy-upgrade-owner', 'legacy-upgrade-owner', false);
INSERT INTO api_keys(id, user_id, key_hash)
VALUES ('legacy-upgrade-key', 'legacy-upgrade-owner', repeat('a', 64));
INSERT INTO usage(id, request_id, user_id, api_key_id, provider_name, model,
status, billing_status, created_at, response_time_ms)
VALUES ('legacy-upgrade-request', 'legacy-upgrade-request',
'legacy-upgrade-owner', 'legacy-upgrade-key', 'test', 'test',
'completed', 'settled', clock_timestamp() - INTERVAL '1 minute', 100);
INSERT INTO stats_overview_hourly(projection_version, bucket_start, dimensions, metrics)
SELECT projection_version, bucket_start, '{}'::jsonb, '{"request_count":1}'::jsonb
FROM stats_bucket_state WHERE granularity = 'hour';
INSERT INTO stats_overview_daily(projection_version, bucket_start, dimensions, metrics)
SELECT projection_version, bucket_start, '{}'::jsonb, '{"request_count":1}'::jsonb
FROM stats_bucket_state WHERE granularity = 'day';
UPDATE stats_bucket_state SET built_revision=source_revision, coverage_status='complete';
"#,
)
.execute(&pool)
.await
.unwrap();
// Model the old September 19 schema after creating real dashboard totals.
// No fact is changed while its later retention helpers are absent.
sqlx::raw_sql(
r#"
DROP FUNCTION public.dashboard_ensure_activity_minute(timestamptz, smallint);
DROP TABLE public.dashboard_activity_minute;
ALTER TABLE public.dashboard_stats_state DROP COLUMN contributions_cleanup_cursor;
UPDATE _sqlx_migrations SET checksum=decode(
'1dd622827b22ae0540f5e43232e7419727aa02d0742649af7e045185e9a13f68655ef7b31007bd0bf6e9e63177022815', 'hex')
WHERE version=20260911000000;
UPDATE _sqlx_migrations SET checksum=decode(
'c64293c5d95fba6c3c43fff764a89da0225b386cfbef14743ab585e1db012fea35fe2cb2eee132e0fb5dac834c0410f4', 'hex')
WHERE version=20260919000000;
"#,
)
.execute(&pool)
.await
.unwrap();
let historical_records = rows_snapshot(&pool, "_sqlx_migrations").await;
let preserved_tables = [
"users",
"api_keys",
"usage",
"usage_attribution_snapshots",
"stats_bucket_state",
"stats_overview_hourly",
"stats_overview_daily",
"dashboard_request_contributions",
"dashboard_stats_total",
"dashboard_stats_minute",
"dashboard_activity_hour",
"dashboard_actor_minute",
];
let mut original_rows = Vec::new();
for table in preserved_tables {
original_rows.push((table, rows_snapshot(&pool, table).await));
}
let account_record: String =
query_scalar("SELECT to_jsonb(m)::text FROM _sqlx_migrations m WHERE version=$1")
.bind(ACCOUNT_ATTRIBUTION)
.fetch_one(&pool)
.await
.unwrap();
let pending = prepare_database_for_startup(&pool).await.unwrap();
assert_eq!(
pending
.iter()
.map(|migration| migration.version)
.collect::<Vec<_>>(),
vec![DIRTY_EVENTS, 20260921010000, 20260921020000, 20260921020100, 20261001000000]
);
assert_eq!(
rows_snapshot(&pool, "_sqlx_migrations").await,
historical_records
);
for _ in 0..2 {
super::super::run_migrations(&pool).await.unwrap();
assert!(super::super::pending_migrations(&pool)
.await
.unwrap()
.is_empty());
for (table, expected) in &original_rows {
assert_eq!(&rows_snapshot(&pool, table).await, expected, "{table}");
}
let after: String =
query_scalar("SELECT to_jsonb(m)::text FROM _sqlx_migrations m WHERE version=$1")
.bind(ACCOUNT_ATTRIBUTION)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(
after, account_record,
"the historical migration must not be restamped"
);
}
assert!(table_exists(&pool, "stats_overview_dirty_events")
.await
.unwrap());
assert!(table_exists(&pool, "dashboard_activity_minute")
.await
.unwrap());
assert!(column_exists(
&pool,
"dashboard_stats_state",
"contributions_cleanup_cursor"
)
.await
.unwrap());
for index in [
"ix_usage_attribution_owner_request",
"ix_usage_analytics_actor_metadata",
] {
let valid: bool =
query_scalar("SELECT indisvalid FROM pg_index WHERE indexrelid=to_regclass($1)")
.bind(index)
.fetch_one(&pool)
.await
.unwrap();
assert!(valid, "{index}");
}
let empty_projections: (i64, i64) = sqlx::query_as(
"SELECT (SELECT count(*) FROM stats_overview_dirty_events), (SELECT count(*) FROM dashboard_activity_minute)",
)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(empty_projections, (0, 0), "upgrade must not backfill usage");
// Correcting an existing request must seed the retained activity counter
// from its old detailed minute, without counting that request twice.
query("UPDATE usage SET response_time_ms=200 WHERE request_id='legacy-upgrade-request'")
.execute(&pool)
.await
.unwrap();
let corrected: (i64, i64, i64) = sqlx::query_as(
"SELECT (SELECT sum((metrics->>'request_count')::bigint)::bigint FROM dashboard_stats_total), (SELECT sum(request_count)::bigint FROM dashboard_activity_minute), (SELECT sum((metrics->>'response_sum_ms')::bigint)::bigint FROM dashboard_stats_total)",
)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(corrected, (1, 1, 200));
query(
"INSERT INTO usage(id,request_id,user_id,api_key_id,provider_name,model,status,billing_status,created_at,response_time_ms) SELECT 'after-upgrade','after-upgrade',user_id,api_key_id,provider_name,model,status,billing_status,created_at,300 FROM usage WHERE request_id='legacy-upgrade-request'",
)
.execute(&pool)
.await
.unwrap();
let written: (i64, i64, i64) = sqlx::query_as(
"SELECT (SELECT sum((metrics->>'request_count')::bigint)::bigint FROM dashboard_stats_total), (SELECT sum(request_count)::bigint FROM dashboard_activity_minute), (SELECT count(*) FROM stats_overview_dirty_events)",
)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(written, (2, 2, 4));
let repo = aether_data_postgres::SqlxUsageReadRepository::new(pool.clone());
assert_eq!(
repo.rebuild_overview_buckets(chrono::Utc::now() + chrono::Duration::days(1), 8)
.await
.unwrap(),
2
);
let rebuilt: (i64, i64) = sqlx::query_as(
"SELECT (SELECT sum((metrics->>'request_count')::bigint)::bigint FROM stats_overview_hourly), (SELECT count(*) FROM stats_overview_dirty_events)",
)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(rebuilt, (2, 0));
pool.close().await;
}
crates/aether-data/runtime/src/lifecycle/migrate/tests/migration_deadlines.rs
+318
View File
@@ -0,0 +1,318 @@
use super::*;
use sqlx::postgres::PgPoolOptions;
use std::time::{Duration, Instant};
const OVERVIEW_MIGRATION: i64 = 20260911000000;
async fn legacy_connection(server: &ManagedPostgresServer) -> PgConnection {
let mut connection = PgConnection::connect(server.database_url()).await.unwrap();
connection.ensure_migrations_table().await.unwrap();
for migration in POSTGRES_MIGRATOR
.iter()
.filter(|migration| migration.version < OVERVIEW_MIGRATION)
{
connection.apply(migration).await.unwrap();
}
connection
}
async fn wait_for_settlement_ddl(connection: &mut PgConnection) {
tokio::time::timeout(Duration::from_secs(5), async {
loop {
let waiting: bool = query_scalar(
"SELECT EXISTS (SELECT 1 FROM pg_locks WHERE relation='public.usage_settlement_snapshots'::regclass AND mode='AccessExclusiveLock' AND NOT granted)",
)
.fetch_one(&mut *connection)
.await
.unwrap();
if waiting {
return;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("migration should reach the blocked settlement ALTER TABLE");
}
async fn assert_overview_migration_rolled_back(pool: &PgPool) {
assert!(
!column_exists(pool, "usage", "failure_origin")
.await
.unwrap(),
"the earlier ALTER TABLE must roll back with the blocked statement"
);
let stamped: bool =
query_scalar("SELECT EXISTS (SELECT 1 FROM public._sqlx_migrations WHERE version=$1)")
.bind(OVERVIEW_MIGRATION)
.fetch_one(pool)
.await
.unwrap();
assert!(
!stamped,
"a failed migration must not receive a success stamp"
);
}
#[tokio::test]
async fn migration_deadlines_release_queued_usage_work_and_roll_back_before_retry() {
let Some(server) = ManagedPostgresServer::try_start().await.unwrap() else {
return;
};
let mut observer = legacy_connection(&server).await;
let pool = PgPoolOptions::new()
.max_connections(1)
.after_connect(|connection, _| {
Box::pin(async move {
query("SET statement_timeout='30s'")
.execute(&mut *connection)
.await?;
query("SET lock_timeout='3s'")
.execute(&mut *connection)
.await?;
Ok(())
})
})
.connect(server.database_url())
.await
.unwrap();
let original_pid: i32 = query_scalar("SELECT pg_backend_pid()")
.fetch_one(&pool)
.await
.unwrap();
let mut business = PgConnection::connect(server.database_url()).await.unwrap();
let mut business_transaction = business.begin().await.unwrap();
query("LOCK TABLE public.usage_settlement_snapshots IN ACCESS SHARE MODE")
.execute(&mut *business_transaction)
.await
.unwrap();
let mut queued_business = PgConnection::connect(server.database_url()).await.unwrap();
// The usage ALTERs take their table exclusively, then settlement ALTER
// waits behind an existing reader. Business writes wait on the usage lock
// already held by this transaction and must resume when it rolls back.
let started = Instant::now();
let (migration_result, ()) = tokio::join!(super::super::run_migrations(&pool), async {
wait_for_settlement_ddl(&mut observer).await;
let acquired_first_lock: bool = query_scalar(
"SELECT EXISTS (SELECT 1 FROM pg_locks WHERE relation='public.usage'::regclass AND mode='AccessExclusiveLock' AND granted)",
)
.fetch_one(&mut observer)
.await
.unwrap();
assert!(
acquired_first_lock,
"the migration must have already changed the first table"
);
let queued_write = query("INSERT INTO public.usage(id,request_id,model,provider_name,status,billing_status,created_at) VALUES ('migration-live-request','migration-live-request','test','test','completed','settled',NOW())")
.execute(&mut queued_business);
tokio::pin!(queued_write);
tokio::select! {
result = &mut queued_write => panic!("the business write should initially queue behind the DDL: {result:?}"),
() = tokio::time::sleep(Duration::from_millis(100)) => {}
}
let written = tokio::time::timeout(Duration::from_secs(4), queued_write)
.await
.expect("queued business work must resume after the migration's lock timeout")
.unwrap();
assert_eq!(written.rows_affected(), 1);
});
let error =
migration_result.expect_err("busy settlement table must defer this upgrade attempt");
assert!(
error.to_string().contains("lock timeout"),
"the failure should identify the bounded lock wait: {error}"
);
assert!(started.elapsed() < Duration::from_secs(5));
assert_overview_migration_rolled_back(&pool).await;
let (new_pid, statement_timeout, lock_timeout): (i32, String, String) = sqlx::query_as(
"SELECT pg_backend_pid(), current_setting('statement_timeout'), current_setting('lock_timeout')",
)
.fetch_one(&pool)
.await
.unwrap();
assert_ne!(
new_pid, original_pid,
"failed migration connection must be discarded"
);
assert_eq!(statement_timeout, "30s");
assert_eq!(lock_timeout, "3s");
business_transaction.rollback().await.unwrap();
super::super::run_migrations(&pool).await.unwrap();
assert!(super::super::pending_migrations(&pool)
.await
.unwrap()
.is_empty());
let success: bool =
query_scalar("SELECT success FROM public._sqlx_migrations WHERE version=$1")
.bind(OVERVIEW_MIGRATION)
.fetch_one(&pool)
.await
.unwrap();
assert!(
success,
"a later quiet retry must succeed without manual stamp repair"
);
assert_eq!(
query_scalar::<_, i64>(
"SELECT count(*) FROM public.usage WHERE request_id='migration-live-request'"
)
.fetch_one(&pool)
.await
.unwrap(),
1,
"the resumed business write must survive the upgrade retry"
);
let settings: (String, String) = sqlx::query_as(
"SELECT current_setting('statement_timeout'), current_setting('lock_timeout')",
)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(settings, ("30s".into(), "3s".into()));
// Startup preparation also takes SQLx's advisory migration lock. Another
// upgrade process must not leave startup waiting indefinitely for it.
observer.lock().await.unwrap();
let preparation =
tokio::time::timeout(Duration::from_secs(4), prepare_database_for_startup(&pool))
.await
.expect("startup preparation must bound its advisory-lock wait");
let preparation_error = preparation.expect_err("another migration owns the advisory lock");
assert!(preparation_error.to_string().contains("lock timeout"));
observer.unlock().await.unwrap();
assert!(prepare_database_for_startup(&pool)
.await
.unwrap()
.is_empty());
pool.close().await;
}
#[tokio::test]
async fn migration_deadlines_caller_cancellation_releases_ddl_locks_and_rolls_back() {
let Some(server) = ManagedPostgresServer::try_start().await.unwrap() else {
return;
};
let mut connection = legacy_connection(&server).await;
sqlx::raw_sql(
r#"
CREATE FUNCTION public.test_cancelled_migration_ddl() RETURNS event_trigger LANGUAGE plpgsql AS $$
BEGIN PERFORM pg_sleep(8); END $$;
CREATE EVENT TRIGGER test_cancelled_migration_ddl ON ddl_command_end
WHEN TAG IN ('ALTER TABLE') EXECUTE FUNCTION public.test_cancelled_migration_ddl();
"#,
)
.execute(&mut connection)
.await
.unwrap();
let pool = PgPoolOptions::new()
.max_connections(1)
.connect(server.database_url())
.await
.unwrap();
// Keep this future in an inner scope: leaving it actually drops the entire
// migration operation, rather than merely dropping a pinned reference.
{
let migrate = super::super::run_migrations(&pool);
tokio::pin!(migrate);
tokio::select! {
result = &mut migrate => panic!("the caller must cancel before the slow migration finishes: {result:?}"),
() = async {
tokio::time::timeout(Duration::from_secs(4), async {
loop {
let running: bool = query_scalar(
"SELECT EXISTS (SELECT 1 FROM pg_locks l JOIN pg_stat_activity a USING(pid) WHERE l.relation='public.usage'::regclass AND l.mode='AccessExclusiveLock' AND l.granted AND a.wait_event='PgSleep')",
)
.fetch_one(&mut connection)
.await
.unwrap();
if running {
return;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
}).await.expect("migration must acquire its first DDL lock before caller cancellation");
} => {}
}
}
tokio::time::timeout(
Duration::from_secs(4),
query_scalar::<_, i64>("SELECT count(*) FROM public.usage").fetch_one(&pool),
)
.await
.expect("dropping the caller future must stop the server-side statement and release DDL locks")
.unwrap();
assert_overview_migration_rolled_back(&pool).await;
sqlx::raw_sql(
"DROP EVENT TRIGGER test_cancelled_migration_ddl; DROP FUNCTION public.test_cancelled_migration_ddl()",
)
.execute(&mut connection)
.await
.unwrap();
super::super::run_migrations(&pool).await.unwrap();
assert!(super::super::pending_migrations(&pool)
.await
.unwrap()
.is_empty());
pool.close().await;
}
#[tokio::test]
async fn migration_deadlines_bound_the_whole_transaction_not_only_each_statement() {
let Some(server) = ManagedPostgresServer::try_start().await.unwrap() else {
return;
};
let mut connection = legacy_connection(&server).await;
// Each ALTER finishes within the ten-second statement limit, but the
// entire migration exceeds it. A per-statement timeout alone is insufficient.
sqlx::raw_sql(
r#"
CREATE FUNCTION public.test_slow_migration_ddl() RETURNS event_trigger LANGUAGE plpgsql AS $$
BEGIN PERFORM pg_sleep(3); END $$;
CREATE EVENT TRIGGER test_slow_migration_ddl ON ddl_command_end
WHEN TAG IN ('ALTER TABLE') EXECUTE FUNCTION public.test_slow_migration_ddl();
"#,
)
.execute(&mut connection)
.await
.unwrap();
let pool = PgPoolOptions::new()
.max_connections(1)
.connect(server.database_url())
.await
.unwrap();
let started = Instant::now();
let result = tokio::time::timeout(Duration::from_secs(15), super::super::run_migrations(&pool))
.await
.expect("the whole migration deadline must fire before all slow statements finish");
assert!(result.is_err(), "the over-budget migration must fail");
assert!(
started.elapsed() >= Duration::from_secs(9),
"upgrade failed before its deadline: {result:?}"
);
assert!(started.elapsed() < Duration::from_secs(15));
// The backend can still be unwinding the statement when its socket closes.
// Reading the first locked table proves cancellation released the DDL lock.
tokio::time::timeout(
Duration::from_secs(4),
query_scalar::<_, i64>("SELECT count(*) FROM public.usage").fetch_one(&pool),
)
.await
.expect("DDL locks must be released when the migration connection closes")
.unwrap();
assert_overview_migration_rolled_back(&pool).await;
sqlx::raw_sql(
"DROP EVENT TRIGGER test_slow_migration_ddl; DROP FUNCTION public.test_slow_migration_ddl()",
)
.execute(&mut connection)
.await
.unwrap();
super::super::run_migrations(&pool).await.unwrap();
assert!(super::super::pending_migrations(&pool)
.await
.unwrap()
.is_empty());
pool.close().await;
}
crates/aether-data/runtime/src/lifecycle/migrate/tests/overview_dirty_events.rs
+154
View File
@@ -0,0 +1,154 @@
use super::*;
use aether_data_contracts::repository::usage::{UsageAnalyticsQuery, UsageAnalyticsView};
use chrono::{TimeZone, Utc};
#[tokio::test]
async fn overview_dirty_events_keep_independent_usage_writes_concurrent_and_reads_fresh() {
let Some(server) = ManagedPostgresServer::try_start()
.await
.expect("overview PostgreSQL should start or skip")
else {
return;
};
let pool = PgPool::connect(server.database_url()).await.unwrap();
prepare_and_apply_clean_postgres_database(&pool).await;
let at = Utc.with_ymd_and_hms(2026, 1, 2, 3, 0, 0).unwrap();
let repo = aether_data_postgres::SqlxUsageReadRepository::new(pool.clone());
let insert = "INSERT INTO usage(id,request_id,model,provider_name,status,billing_status,created_at,response_time_ms) VALUES($1,$1,'dirty-event-test','test','completed','settled',$2,100)";
let second_id: String = query_scalar(
"SELECT 'second-' || n FROM generate_series(1,100) n WHERE (hashtextextended('second-' || n,0) & 15) <> (hashtextextended('first',0) & 15) LIMIT 1",
)
.fetch_one(&pool)
.await
.unwrap();
// A stays open while B writes the same hour/day. Before the migration,
// B times out on the shared stats_bucket_state day row despite a distinct request.
let mut first = pool.begin().await.unwrap();
sqlx::query(insert)
.bind("first")
.bind(at)
.execute(&mut *first)
.await
.unwrap();
let mut second = pool.begin().await.unwrap();
sqlx::query("SET LOCAL lock_timeout='500ms'")
.execute(&mut *second)
.await
.unwrap();
sqlx::query(insert)
.bind(&second_id)
.bind(at)
.execute(&mut *second)
.await
.unwrap();
second.commit().await.unwrap();
assert_eq!(
repo.rebuild_overview_buckets(at + chrono::Duration::days(1), 8)
.await
.unwrap(),
2
);
let query = UsageAnalyticsQuery {
from_unix_ms: at.timestamp_millis() as u64,
to_unix_ms: (at + chrono::Duration::hours(1)).timestamp_millis() as u64,
timezone: "UTC".into(),
view: UsageAnalyticsView::Summary,
limit: 100,
..Default::default()
};
assert_eq!(
repo.query_usage_analytics(&query)
.await
.unwrap()
.summary
.request_count,
1
);
// Commit order differs from transaction-ID order. Consuming B must never
// acknowledge A, and A must invalidate the already-published clean projection.
first.commit().await.unwrap();
let pending = repo.query_usage_analytics(&query).await.unwrap();
assert_eq!(pending.summary.request_count, 2);
assert_eq!(pending.coverage.dirty_bucket_count, 1);
let queued: i64 = query_scalar("SELECT count(*) FROM stats_overview_dirty_events")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(queued, 2);
assert_eq!(
repo.rebuild_overview_buckets(at + chrono::Duration::days(1), 8)
.await
.unwrap(),
2
);
let rebuilt = repo.query_usage_analytics(&query).await.unwrap();
assert_eq!(rebuilt.summary, pending.summary);
assert_eq!(rebuilt.coverage.dirty_bucket_count, 0);
assert_eq!(
query_scalar::<_, i64>("SELECT count(*) FROM stats_overview_dirty_events")
.fetch_one(&pool)
.await
.unwrap(),
0
);
let mut rolled_back = pool.begin().await.unwrap();
sqlx::query(insert)
.bind("rollback")
.bind(at)
.execute(&mut *rolled_back)
.await
.unwrap();
rolled_back.rollback().await.unwrap();
assert_eq!(
query_scalar::<_, i64>("SELECT count(*) FROM stats_overview_dirty_events")
.fetch_one(&pool)
.await
.unwrap(),
0
);
// Repeated mutations in one transaction deduplicate, but deleting facts
// retains the irreversible-loss marker until it is merged into state.
let mut deleted = pool.begin().await.unwrap();
sqlx::query("UPDATE usage SET response_time_ms=200 WHERE request_id='first'")
.execute(&mut *deleted)
.await
.unwrap();
sqlx::query("DELETE FROM usage WHERE request_id='first'")
.execute(&mut *deleted)
.await
.unwrap();
deleted.commit().await.unwrap();
assert_eq!(
query_scalar::<_, i64>(
"SELECT count(*) FROM stats_overview_dirty_events WHERE unrecoverable"
)
.fetch_one(&pool)
.await
.unwrap(),
2
);
let lost = repo.query_usage_analytics(&query).await.unwrap();
assert_eq!(lost.summary.request_count, 1);
assert_eq!(lost.unrecoverable_bucket_count, 1);
assert_eq!(
repo.rebuild_overview_buckets(at + chrono::Duration::days(1), 8)
.await
.unwrap(),
1
);
let merged = repo.query_usage_analytics(&query).await.unwrap();
assert_eq!(merged.unrecoverable_bucket_count, 1);
assert_eq!(merged.summary, lost.summary);
assert_eq!(
query_scalar::<_, i64>("SELECT count(*) FROM stats_overview_dirty_events")
.fetch_one(&pool)
.await
.unwrap(),
0
);
pool.close().await;
}
crates/aether-data/runtime/src/lifecycle/migrate/tests/overview_migration_safety.rs
+163
View File
@@ -0,0 +1,163 @@
use super::*;
use aether_data_contracts::repository::usage::{UsageAnalyticsQuery, UsageAnalyticsView};
use chrono::{TimeZone, Utc};
const OVERVIEW_START: i64 = 20260911000000;
const BILLING_INDEX: i64 = 20260918000000;
const INSERT_USAGE: &str = "INSERT INTO usage(id,request_id,model,provider_name,status,billing_status,created_at) VALUES($1,$1,'overview-migration-test','test','completed','settled',$2)";
async fn apply_through_overview(connection: &mut PgConnection) {
connection.ensure_migrations_table().await.unwrap();
for migration in POSTGRES_MIGRATOR
.iter()
.filter(|migration| migration.version <= OVERVIEW_START)
{
connection.apply(migration).await.unwrap();
}
}
async fn assert_independent_same_bucket_writes(pool: &PgPool, prefix: &str) {
let at = Utc.with_ymd_and_hms(2026, 1, 2, 3, 0, 0).unwrap();
let mut first = pool.begin().await.unwrap();
query(INSERT_USAGE)
.bind(format!("{prefix}-first"))
.bind(at)
.execute(&mut *first)
.await
.unwrap();
let mut second = pool.begin().await.unwrap();
query("SET LOCAL lock_timeout='500ms'")
.execute(&mut *second)
.await
.unwrap();
query(INSERT_USAGE)
.bind(format!("{prefix}-second"))
.bind(at)
.execute(&mut *second)
.await
.expect("another request in the same hour/day must not wait for the first transaction");
second.commit().await.unwrap();
first.commit().await.unwrap();
}
async fn is_stamped(pool: &PgPool, version: i64) -> bool {
query_scalar("SELECT EXISTS(SELECT 1 FROM _sqlx_migrations WHERE version=$1 AND success)")
.bind(version)
.fetch_one(pool)
.await
.unwrap()
}
async fn queue_and_state(pool: &PgPool) -> (String, String) {
sqlx::query_as(
r#"
SELECT
(SELECT COALESCE(jsonb_agg(to_jsonb(e) ORDER BY transaction_id,projection_version,granularity,bucket_start),'[]')::text FROM stats_overview_dirty_events e),
(SELECT COALESCE(jsonb_agg(to_jsonb(s) ORDER BY projection_version,granularity,bucket_start),'[]')::text FROM stats_bucket_state s)
"#,
)
.fetch_one(pool)
.await
.unwrap()
}
#[tokio::test]
async fn overview_queue_is_safe_from_its_first_installation() {
let Some(server) = ManagedPostgresServer::try_start().await.unwrap() else {
return;
};
let mut connection = PgConnection::connect(server.database_url()).await.unwrap();
apply_through_overview(&mut connection).await;
let pool = PgPool::connect(server.database_url()).await.unwrap();
// The initial installation must be safe before any subsequent migration,
// including the potentially long concurrent index build, has completed.
assert_independent_same_bucket_writes(&pool, "first-install").await;
let (queued, obsolete, direct): (i64, i64, i64) = sqlx::query_as(
"SELECT (SELECT count(*) FROM stats_overview_dirty_events), (SELECT count(*) FROM stats_overview_dirty_events WHERE projection_version <> 'overview-v2'), (SELECT count(*) FROM stats_bucket_state)",
)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!((queued, obsolete, direct), (4, 0, 0));
assert!(!is_stamped(&pool, BILLING_INDEX).await);
pool.close().await;
}
#[tokio::test]
async fn overview_queue_survives_later_index_failure_and_retry_without_losing_data() {
let Some(server) = ManagedPostgresServer::try_start().await.unwrap() else {
return;
};
let mut connection = PgConnection::connect(server.database_url()).await.unwrap();
apply_through_overview(&mut connection).await;
let pool = PgPool::connect(server.database_url()).await.unwrap();
let at = Utc.with_ymd_and_hms(2026, 1, 2, 3, 0, 0).unwrap();
query(INSERT_USAGE)
.bind("before-index-failure")
.bind(at)
.execute(&pool)
.await
.unwrap();
let repo = aether_data_postgres::SqlxUsageReadRepository::new(pool.clone());
assert_eq!(
repo.rebuild_overview_buckets(at + chrono::Duration::days(1), 8)
.await
.unwrap(),
2
);
let original_state = queue_and_state(&pool).await.1;
let mut blocker = pool.begin().await.unwrap();
query("LOCK TABLE usage_settlement_snapshots IN SHARE MODE")
.execute(&mut *blocker)
.await
.unwrap();
let error = super::super::run_migrations(&pool)
.await
.expect_err("the concurrent index must encounter the held relation lock");
assert!(error.to_string().contains("lock timeout"), "{error}");
assert!(is_stamped(&pool, OVERVIEW_START).await);
assert!(!is_stamped(&pool, BILLING_INDEX).await);
// A failed later migration must leave the safe trigger committed and able
// to accept independent business writes until the upgrade can be retried.
assert_independent_same_bucket_writes(&pool, "failed-upgrade").await;
let after_failure = queue_and_state(&pool).await;
assert_eq!(after_failure.1, original_state);
assert_eq!(
query_scalar::<_, i64>("SELECT count(*) FROM stats_overview_dirty_events")
.fetch_one(&pool)
.await
.unwrap(),
4
);
blocker.rollback().await.unwrap();
super::super::run_migrations(&pool).await.unwrap();
assert_eq!(queue_and_state(&pool).await, after_failure);
assert!(super::super::pending_migrations(&pool)
.await
.unwrap()
.is_empty());
assert_eq!(
repo.rebuild_overview_buckets(at + chrono::Duration::days(1), 8)
.await
.unwrap(),
2
);
let result = repo
.query_usage_analytics(&UsageAnalyticsQuery {
from_unix_ms: at.timestamp_millis() as u64,
to_unix_ms: (at + chrono::Duration::hours(1)).timestamp_millis() as u64,
timezone: "UTC".into(),
view: UsageAnalyticsView::Summary,
limit: 100,
..Default::default()
})
.await
.unwrap();
assert_eq!(result.summary.request_count, 3);
assert_eq!(result.coverage.dirty_bucket_count, 0);
pool.close().await;
}
crates/aether-data/runtime/src/lifecycle/migrate/tests/provider_expenses.rs
+115
View File
@@ -0,0 +1,115 @@
use super::*;
use aether_data_contracts::repository::billing::{
AdminBillingMutationOutcome, BillingReadRepository, ProviderExpenseInput, ProviderExpenseQuery,
};
#[tokio::test]
async fn postgres_provider_expense_ledger_migration_preserves_exact_sums_idempotency_and_voids() {
let Some(server) = ManagedPostgresServer::try_start()
.await
.expect("local postgres should start or skip")
else {
return;
};
let pool = PgPool::connect(server.database_url()).await.unwrap();
// Exercise only this self-contained migration against an isolated database.
sqlx::raw_sql(include_str!(
"../../../../schema/bootstrap/postgres/210_provider_expenses.sql"
))
.execute(&pool)
.await
.unwrap();
let repository = crate::repository::billing::SqlxBillingReadRepository::new(pool.clone());
let input = ProviderExpenseInput {
client_request_id: uuid::Uuid::new_v4().to_string(),
provider_id: "deleted-later".into(),
provider_name: "Supplier".into(),
kind: "recharge".into(),
amount: "0.10000000".into(),
currency: "USD".into(),
paid_at_unix_ms: 1000,
period_start_unix_ms: None,
period_end_unix_ms: None,
note: None,
external_reference: None,
created_by: Some("admin".into()),
};
let (a, b) = tokio::join!(
repository.create_provider_expense(&input),
repository.create_provider_expense(&input)
);
let (AdminBillingMutationOutcome::Applied(a), AdminBillingMutationOutcome::Applied(b)) =
(a.unwrap(), b.unwrap())
else {
panic!("expected applied")
};
assert_eq!(a.id, b.id);
let mut subscription = input.clone();
subscription.client_request_id = uuid::Uuid::new_v4().to_string();
subscription.amount = "0.20000000".into();
subscription.kind = "subscription".into();
let AdminBillingMutationOutcome::Applied(second) = repository
.create_provider_expense(&subscription)
.await
.unwrap()
else {
panic!("expected applied")
};
let mut cny = input.clone();
cny.client_request_id = uuid::Uuid::new_v4().to_string();
cny.currency = "CNY".into();
cny.amount = "7.00000000".into();
repository.create_provider_expense(&cny).await.unwrap();
let query = ProviderExpenseQuery {
from_unix_ms: 0,
to_unix_ms: 2000,
limit: 1,
offset: 1,
};
let page = repository
.list_provider_expenses(&query)
.await
.unwrap()
.unwrap();
assert_eq!(page.total, 3);
assert_eq!(page.items.len(), 1);
assert_eq!(page.totals[0].currency, "CNY");
assert_eq!(page.totals[1].amount, "0.30000000");
assert_eq!(page.totals[1].subscription_amount, "0.20000000");
let mut conflict = input.clone();
conflict.amount = "5.00000000".into();
assert!(matches!(
repository.create_provider_expense(&conflict).await.unwrap(),
AdminBillingMutationOutcome::Invalid(_)
));
let AdminBillingMutationOutcome::Applied(voided) = repository
.void_provider_expense(&second.id, Some("operator-1"))
.await
.unwrap()
else {
panic!("expected void")
};
let AdminBillingMutationOutcome::Applied(again) = repository
.void_provider_expense(&second.id, Some("operator-2"))
.await
.unwrap()
else {
panic!("expected repeat void")
};
assert_eq!(voided, again);
let page = repository
.list_provider_expenses(&query)
.await
.unwrap()
.unwrap();
assert_eq!(page.total, 2);
assert_eq!(page.totals[1].amount, "0.10000000");
let raw_count: i64 = sqlx::query_scalar("SELECT count(*) FROM provider_expenses")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(raw_count, 3);
let invalid=sqlx::query("INSERT INTO provider_expenses(id,client_request_id,provider_id,provider_name,kind,amount,currency,paid_at) VALUES('bad','bad','p','P','recharge',-1,'USD',NOW())").execute(&pool).await;
assert!(invalid.is_err());
pool.close().await;
}
crates/aether-data/runtime/src/repository/announcements/memory.rs
+148 -2
View File
@@ -8,7 +8,8 @@ use uuid::Uuid;
use crate::DataLayerError;
use aether_data_contracts::repository::announcements::{
AnnouncementListQuery, AnnouncementReadRepository, AnnouncementWriteRepository,
CreateAnnouncementRecord, StoredAnnouncement, StoredAnnouncementPage, UpdateAnnouncementRecord,
CreateAnnouncementRecord, StoredAnnouncement, StoredAnnouncementPage, StoredUserAnnouncement,
StoredUserAnnouncementPage, UpdateAnnouncementRecord, UserAnnouncementListQuery,
};
#[derive(Debug, Default)]
@@ -105,6 +106,65 @@ impl AnnouncementReadRepository for InMemoryAnnouncementReadRepository {
Ok(StoredAnnouncementPage { items, total })
}
async fn list_user_announcements(
&self,
user_id: &str,
query: &UserAnnouncementListQuery,
) -> Result<StoredUserAnnouncementPage, DataLayerError> {
query.validate()?;
let announcements = self
.announcements
.read()
.expect("announcement repository lock");
let reads = self
.announcement_reads
.read()
.expect("announcement reads repository lock");
let mut unread_count = 0;
let mut items = announcements
.iter()
.filter(|announcement| {
announcement.is_active
&& announcement
.start_time_unix_secs
.is_none_or(|value| value <= query.now_unix_secs)
&& announcement
.end_time_unix_secs
.is_none_or(|value| value >= query.now_unix_secs)
})
.filter_map(|announcement| {
let is_read = reads.contains(&(user_id.to_string(), announcement.id.clone()));
if !is_read {
unread_count += 1;
}
(!query.unread_only || !is_read).then_some((announcement, is_read))
})
.collect::<Vec<_>>();
items.sort_by(|(left, _), (right, _)| {
right
.is_pinned
.cmp(&left.is_pinned)
.then_with(|| right.priority.cmp(&left.priority))
.then_with(|| right.created_at_unix_ms.cmp(&left.created_at_unix_ms))
.then_with(|| left.id.cmp(&right.id))
});
let total = items.len() as u64;
let items = items
.into_iter()
.skip(query.offset)
.take(query.limit)
.map(|(announcement, is_read)| StoredUserAnnouncement {
announcement: announcement.clone(),
is_read,
})
.collect();
Ok(StoredUserAnnouncementPage {
items,
total,
unread_count,
})
}
async fn count_unread_active_announcements(
&self,
user_id: &str,
@@ -294,9 +354,95 @@ mod tests {
use super::InMemoryAnnouncementReadRepository;
use crate::repository::announcements::{
AnnouncementReadRepository, AnnouncementWriteRepository, CreateAnnouncementRecord,
StoredAnnouncement, UpdateAnnouncementRecord,
StoredAnnouncement, UpdateAnnouncementRecord, UserAnnouncementListQuery,
};
#[tokio::test]
async fn personal_announcement_page_preserves_global_unread_and_visibility() {
let now = 1_800_000_000;
let announcements = [
("pinned", true, true, 0, None, None),
("normal-a", true, false, 20, Some(now), Some(now)),
("normal-b", true, false, 20, None, None),
("draft", false, false, 99, None, None),
("future", true, false, 99, Some(now + 1), None),
("expired", true, false, 99, None, Some(now - 1)),
]
.into_iter()
.map(|(id, active, pinned, priority, start, end)| {
StoredAnnouncement::new(
id.into(),
id.into(),
"content".into(),
"info".into(),
priority,
active,
pinned,
false,
None,
None,
start,
end,
now,
now,
)
.unwrap()
});
let repository = InMemoryAnnouncementReadRepository::seed_with_reads(
announcements,
[("reader".into(), "pinned".into())],
);
let mut query = UserAnnouncementListQuery {
unread_only: false,
offset: 0,
limit: 1,
now_unix_secs: now as u64,
};
let first = repository
.list_user_announcements("reader", &query)
.await
.unwrap();
assert_eq!((first.total, first.unread_count), (3, 2));
assert_eq!(first.items[0].announcement.id, "pinned");
assert!(first.items[0].is_read);
query.unread_only = true;
query.offset = 1;
let second = repository
.list_user_announcements("reader", &query)
.await
.unwrap();
assert_eq!((second.total, second.unread_count), (2, 2));
assert_eq!(second.items[0].announcement.id, "normal-b");
assert!(!second.items[0].is_read);
query.offset = i64::MAX as usize;
let empty = repository
.list_user_announcements("reader", &query)
.await
.unwrap();
assert!(empty.items.is_empty());
assert_eq!((empty.total, empty.unread_count), (2, 2));
let other = repository
.list_user_announcements("other", &query)
.await
.unwrap();
assert_eq!((other.total, other.unread_count), (3, 3));
repository
.mark_announcement_as_read("reader", "normal-a", now as u64)
.await
.unwrap();
query.offset = 0;
let after_read = repository
.list_user_announcements("reader", &query)
.await
.unwrap();
assert_eq!((after_read.total, after_read.unread_count), (1, 1));
query.limit = 101;
assert!(repository
.list_user_announcements("reader", &query)
.await
.is_err());
}
#[tokio::test]
async fn reads_seeded_announcements() {
let repository = InMemoryAnnouncementReadRepository::seed(vec![StoredAnnouncement::new(
crates/aether-data/runtime/src/repository/announcements/mod.rs
+2 -1
View File
@@ -2,7 +2,8 @@ mod memory;
pub use aether_data_contracts::repository::announcements::{
AnnouncementListQuery, AnnouncementReadRepository, AnnouncementWriteRepository,
CreateAnnouncementRecord, StoredAnnouncement, StoredAnnouncementPage, UpdateAnnouncementRecord,
CreateAnnouncementRecord, StoredAnnouncement, StoredAnnouncementPage, StoredUserAnnouncement,
StoredUserAnnouncementPage, UpdateAnnouncementRecord, UserAnnouncementListQuery,
};
#[cfg(feature = "postgres")]
pub use aether_data_postgres::SqlxAnnouncementReadRepository;
crates/aether-data/runtime/src/repository/auth/memory.rs
+55
View File
@@ -229,6 +229,24 @@ impl InMemoryAuthApiKeySnapshotRepository {
.unwrap_or(0)
}
pub fn standalone_flags(&self) -> BTreeMap<String, bool> {
let index = self
.index
.read()
.expect("auth api key snapshot repository lock");
index
.export_by_api_key_id
.iter()
.map(|(id, record)| (id.clone(), record.is_standalone))
.chain(
index
.by_api_key_id
.iter()
.map(|(id, snapshot)| (id.clone(), snapshot.api_key_is_standalone)),
)
.collect()
}
pub fn snapshot_lookup_count(&self, api_key_id: &str) -> usize {
self.index
.read()
@@ -2101,6 +2119,43 @@ mod tests {
.is_some());
}
#[tokio::test]
async fn standalone_flags_cover_export_only_keys_and_remove_deleted_keys() {
let mut standalone = sample_snapshot("standalone-key", "user-1");
standalone.api_key_is_standalone = true;
let repository = InMemoryAuthApiKeySnapshotRepository::seed([
(None, sample_snapshot("member-key", "user-1")),
(None, standalone),
]);
let mut export = repository
.list_export_api_keys_by_ids(&["standalone-key".into()])
.await
.unwrap()
.remove(0);
export.api_key_id = "export-only-key".into();
let repository = repository.with_export_records([export]);
assert_eq!(
repository.standalone_flags(),
std::collections::BTreeMap::from([
("member-key".into(), false),
("standalone-key".into(), true),
("export-only-key".into(), true),
])
);
assert!(repository
.delete_standalone_api_key("standalone-key")
.await
.unwrap());
assert_eq!(
repository.standalone_flags(),
std::collections::BTreeMap::from([
("member-key".into(), false),
("export-only-key".into(), true),
])
);
}
#[tokio::test]
async fn reads_auth_snapshot_by_all_supported_keys() {
let repository = InMemoryAuthApiKeySnapshotRepository::seed(vec![(
crates/aether-data/runtime/src/repository/billing/memory.rs
+248 -3
View File
@@ -16,6 +16,7 @@ type BillingContextMap = BTreeMap<BillingContextKey, StoredBillingModelContext>;
#[derive(Debug, Default)]
pub struct InMemoryBillingReadRepository {
provider_expenses: RwLock<BTreeMap<String, super::ProviderExpenseRecord>>,
by_key: RwLock<BillingContextMap>,
gateway_configs_by_provider: RwLock<BTreeMap<String, PaymentGatewayConfigRecord>>,
billing_plans_by_id: RwLock<BTreeMap<String, BillingPlanRecord>>,
@@ -23,6 +24,20 @@ pub struct InMemoryBillingReadRepository {
}
impl InMemoryBillingReadRepository {
pub fn seed_user_plan_entitlements(
items: impl IntoIterator<Item = UserPlanEntitlementRecord>,
) -> Self {
Self {
entitlements_by_id: RwLock::new(
items
.into_iter()
.map(|item| (item.id.clone(), item))
.collect(),
),
..Self::default()
}
}
pub fn seed<I>(items: I) -> Self
where
I: IntoIterator<Item = StoredBillingModelContext>,
@@ -39,6 +54,7 @@ impl InMemoryBillingReadRepository {
);
}
Self {
provider_expenses: RwLock::default(),
by_key: RwLock::new(by_key),
gateway_configs_by_provider: RwLock::new(BTreeMap::new()),
billing_plans_by_id: RwLock::new(BTreeMap::new()),
@@ -325,6 +341,68 @@ impl BillingReadRepository for InMemoryBillingReadRepository {
Ok(AdminBillingMutationOutcome::Applied(record))
}
async fn list_provider_expenses(
&self,
query: &super::ProviderExpenseQuery,
) -> Result<Option<super::ProviderExpensePage>, DataLayerError> {
let guard = self
.provider_expenses
.read()
.expect("provider expense store should lock");
super::provider_expense_memory_page(guard.values().cloned(), query).map(Some)
}
async fn create_provider_expense(
&self,
input: &super::ProviderExpenseInput,
) -> Result<AdminBillingMutationOutcome<super::ProviderExpenseRecord>, DataLayerError> {
if let Err(detail) = input.validate() {
return Ok(AdminBillingMutationOutcome::Invalid(detail));
}
let mut guard = self
.provider_expenses
.write()
.expect("provider expense store should lock");
if let Some(existing) = guard
.values()
.find(|r| r.entry.client_request_id == input.client_request_id)
{
return Ok(if existing.entry.same_request_as(input) {
AdminBillingMutationOutcome::Applied(existing.clone())
} else {
AdminBillingMutationOutcome::Invalid(
"client_request_id was already used for another expense".into(),
)
});
}
let record = super::ProviderExpenseRecord {
id: uuid::Uuid::new_v4().to_string(),
entry: input.clone(),
created_at_unix_ms: chrono::Utc::now().timestamp_millis().max(0) as u64,
voided_at_unix_ms: None,
voided_by: None,
};
guard.insert(record.id.clone(), record.clone());
Ok(AdminBillingMutationOutcome::Applied(record))
}
async fn void_provider_expense(
&self,
id: &str,
operator: Option<&str>,
) -> Result<AdminBillingMutationOutcome<super::ProviderExpenseRecord>, DataLayerError> {
let mut guard = self
.provider_expenses
.write()
.expect("provider expense store should lock");
let Some(record) = guard.get_mut(id) else {
return Ok(AdminBillingMutationOutcome::NotFound);
};
if record.voided_at_unix_ms.is_none() {
record.voided_at_unix_ms = Some(chrono::Utc::now().timestamp_millis().max(0) as u64);
record.voided_by = operator.map(str::to_owned);
}
Ok(AdminBillingMutationOutcome::Applied(record.clone()))
}
async fn list_billing_plans(
&self,
include_disabled: bool,
@@ -435,6 +513,15 @@ impl BillingReadRepository for InMemoryBillingReadRepository {
async fn list_user_plan_entitlements(
&self,
user_id: &str,
) -> Result<Option<Vec<UserPlanEntitlementRecord>>, DataLayerError> {
self.list_user_plan_entitlements_with_history(user_id, false)
.await
}
async fn list_user_plan_entitlements_with_history(
&self,
user_id: &str,
include_inactive: bool,
) -> Result<Option<Vec<UserPlanEntitlementRecord>>, DataLayerError> {
let now = current_unix_secs();
let mut items = self
@@ -444,12 +531,12 @@ impl BillingReadRepository for InMemoryBillingReadRepository {
.values()
.filter(|item| {
item.user_id == user_id
&& item.status == "active"
&& item.expires_at_unix_secs > now
&& (include_inactive
|| (item.status == "active" && item.expires_at_unix_secs > now))
})
.cloned()
.collect::<Vec<_>>();
items.sort_by_key(|item| item.expires_at_unix_secs);
items.sort_by_key(|item| (item.expires_at_unix_secs, item.created_at_unix_secs));
Ok(Some(items))
}
@@ -592,6 +679,59 @@ mod tests {
.expect("billing context should build")
}
#[tokio::test]
async fn list_user_plan_entitlements_history_includes_inactive_only_for_selected_user() {
let repository = InMemoryBillingReadRepository::default();
let now = super::current_unix_secs();
for (id, user_id, status, expires_at) in [
("active", "user-1", "active", now + 3600),
("expired", "user-1", "active", now - 60),
("revoked", "user-1", "revoked", now + 3600),
("replaced", "user-1", "replaced", now + 3600),
("another-user", "user-2", "revoked", now + 3600),
] {
repository.entitlements_by_id.write().unwrap().insert(
id.to_string(),
super::UserPlanEntitlementRecord {
id: id.to_string(),
user_id: user_id.to_string(),
plan_id: "plan-1".to_string(),
payment_order_id: format!("order-{id}"),
status: status.to_string(),
starts_at_unix_secs: now - 120,
expires_at_unix_secs: expires_at,
entitlements_snapshot: json!([]),
created_at_unix_secs: now - 120,
updated_at_unix_secs: now - 60,
},
);
}
let active = repository
.list_user_plan_entitlements("user-1")
.await
.unwrap()
.unwrap();
assert_eq!(active.len(), 1);
assert_eq!(active[0].id, "active");
let explicit_active = repository
.list_user_plan_entitlements_with_history("user-1", false)
.await
.unwrap()
.unwrap();
assert_eq!(active, explicit_active);
let history = repository
.list_user_plan_entitlements_with_history("user-1", true)
.await
.unwrap()
.unwrap();
assert_eq!(history.len(), 4);
assert!(history.iter().all(|item| item.user_id == "user-1"));
for expected in ["active", "expired", "revoked", "replaced"] {
assert!(history.iter().any(|item| item.id == expected));
}
}
#[tokio::test]
async fn falls_back_to_provider_without_key_scope() {
let repository = InMemoryBillingReadRepository::seed(vec![sample_context()]);
@@ -779,3 +919,108 @@ mod tests {
assert_eq!(after, expected);
}
}
#[cfg(test)]
mod provider_expense_tests {
use super::*;
use crate::repository::billing::{ProviderExpenseInput, ProviderExpenseQuery};
fn input() -> ProviderExpenseInput {
ProviderExpenseInput {
client_request_id: uuid::Uuid::new_v4().to_string(),
provider_id: "provider-1".into(),
provider_name: "Supplier".into(),
kind: "recharge".into(),
amount: "12.34000000".into(),
currency: "USD".into(),
paid_at_unix_ms: 1000,
period_start_unix_ms: None,
period_end_unix_ms: None,
note: Some("test".into()),
external_reference: None,
created_by: Some("admin-1".into()),
}
}
#[tokio::test]
async fn provider_expense_retries_and_void_are_idempotent_without_erasing_audit() {
let repo = InMemoryBillingReadRepository::default();
let input = input();
let AdminBillingMutationOutcome::Applied(first) =
repo.create_provider_expense(&input).await.unwrap()
else {
panic!("expected record")
};
let mut retried = input.clone();
retried.provider_name = "Renamed supplier".into();
let AdminBillingMutationOutcome::Applied(second) =
repo.create_provider_expense(&retried).await.unwrap()
else {
panic!("expected retry")
};
assert_eq!(first, second);
retried.amount = "99".into();
assert!(matches!(
repo.create_provider_expense(&retried).await.unwrap(),
AdminBillingMutationOutcome::Invalid(_)
));
let query = ProviderExpenseQuery {
from_unix_ms: 0,
to_unix_ms: 2000,
limit: 20,
offset: 0,
};
assert_eq!(
repo.list_provider_expenses(&query)
.await
.unwrap()
.unwrap()
.total,
1
);
let AdminBillingMutationOutcome::Applied(voided) = repo
.void_provider_expense(&first.id, Some("admin-2"))
.await
.unwrap()
else {
panic!("expected void")
};
let AdminBillingMutationOutcome::Applied(again) = repo
.void_provider_expense(&first.id, Some("admin-3"))
.await
.unwrap()
else {
panic!("expected retry void")
};
assert_eq!(voided, again);
assert_eq!(again.voided_by.as_deref(), Some("admin-2"));
let page = repo.list_provider_expenses(&query).await.unwrap().unwrap();
assert_eq!(page.total, 0);
assert!(page.totals.is_empty());
let AdminBillingMutationOutcome::Applied(after_void) =
repo.create_provider_expense(&input).await.unwrap()
else {
panic!("expected original tombstone")
};
assert_eq!(after_void, again);
}
#[tokio::test]
async fn provider_expense_duplicate_submissions_record_once() {
let repo = InMemoryBillingReadRepository::default();
let input = input();
let (a, b) = tokio::join!(
repo.create_provider_expense(&input),
repo.create_provider_expense(&input)
);
let (AdminBillingMutationOutcome::Applied(a), AdminBillingMutationOutcome::Applied(b)) =
(a.unwrap(), b.unwrap())
else {
panic!("expected records")
};
assert_eq!(a.id, b.id);
let mut invalid = input.clone();
invalid.period_start_unix_ms = Some(100);
assert!(matches!(
repo.create_provider_expense(&invalid).await.unwrap(),
AdminBillingMutationOutcome::Invalid(_)
));
}
}
crates/aether-data/runtime/src/repository/settlement/memory.rs
+2 -4
View File
@@ -10,7 +10,7 @@ use super::{
ReserveUsagePolicyCostInput, ReserveUsagePolicyCostOutcome, ReserveUsagePolicyRequestInput,
ReserveUsagePolicyRequestOutcome, SettlementWriteRepository, StoredUsagePolicyCostReservation,
StoredUsagePolicyRequestAdmission, StoredUsageSettlement, UsagePolicyCostReservationState,
UsagePolicyRequestAdmissionState, UsageSettlementInput, SETTLEMENT_EPSILON_USD,
UsagePolicyRequestAdmissionState, UsageSettlementInput,
};
use crate::repository::wallet::{InMemoryWalletRepository, StoredWalletSnapshot};
use crate::DataLayerError;
@@ -511,9 +511,7 @@ impl SettlementWriteRepository for InMemorySettlementRepository {
settlement.wallet_recharge_balance_after = Some(wallet.balance);
settlement.wallet_gift_balance_after = Some(wallet.gift_balance);
settlement.wallet_balance_after = Some(wallet.balance + wallet.gift_balance);
} else if final_billing_status == "settled"
&& billable_cost_usd > SETTLEMENT_EPSILON_USD
{
} else if final_billing_status == "settled" && billable_cost_usd > 0.0 {
final_billing_status = "insufficient_quota".to_string();
settlement.billing_status = final_billing_status.clone();
}
crates/aether-data/runtime/src/repository/usage/memory.rs
+150 -5
View File
@@ -44,17 +44,33 @@ use super::{
use crate::repository::auth::InMemoryAuthApiKeySnapshotRepository;
use crate::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use crate::DataLayerError;
mod analytics;
mod dashboard_summary;
#[derive(Debug, Default)]
pub struct InMemoryUsageReadRepository {
dashboard_projection: RwLock<dashboard_summary::DashboardProjection>,
by_request_id: RwLock<BTreeMap<String, StoredRequestUsageAudit>>,
detached_bodies: RwLock<BTreeMap<String, Value>>,
provider_usage_windows: RwLock<Vec<StoredProviderUsageWindow>>,
auth_api_keys: Option<Arc<InMemoryAuthApiKeySnapshotRepository>>,
provider_catalog: Option<Arc<InMemoryProviderCatalogReadRepository>>,
analytics_users: RwLock<Vec<aether_data_contracts::repository::users::StoredUserSummary>>,
analytics_candidates:
RwLock<Vec<aether_data_contracts::repository::candidates::StoredRequestCandidate>>,
analytics_allocations: RwLock<
BTreeMap<String, aether_data_contracts::repository::usage::UsageAnalyticsAllocation>,
>,
}
impl InMemoryUsageReadRepository {
fn analytics_key_flags(&self) -> BTreeMap<String, bool> {
self.auth_api_keys
.as_ref()
.map(|repository| repository.standalone_flags())
.unwrap_or_default()
}
pub fn seed<I>(items: I) -> Self
where
I: IntoIterator<Item = StoredRequestUsageAudit>,
@@ -66,11 +82,15 @@ impl InMemoryUsageReadRepository {
by_request_id.insert(item.request_id.clone(), item);
}
Self {
dashboard_projection: Default::default(),
by_request_id: RwLock::new(by_request_id),
detached_bodies: RwLock::new(BTreeMap::new()),
provider_usage_windows: RwLock::new(Vec::new()),
auth_api_keys: None,
provider_catalog: None,
analytics_users: Default::default(),
analytics_candidates: Default::default(),
analytics_allocations: Default::default(),
}
}
@@ -119,11 +139,15 @@ impl InMemoryUsageReadRepository {
by_request_id.insert(request_id, item);
}
Self {
dashboard_projection: Default::default(),
by_request_id: RwLock::new(by_request_id),
detached_bodies: RwLock::new(detached_bodies),
provider_usage_windows: RwLock::new(Vec::new()),
auth_api_keys: None,
provider_catalog: None,
analytics_users: Default::default(),
analytics_candidates: Default::default(),
analytics_allocations: Default::default(),
}
}
@@ -132,11 +156,15 @@ impl InMemoryUsageReadRepository {
I: IntoIterator<Item = StoredProviderUsageWindow>,
{
Self {
dashboard_projection: self.dashboard_projection,
by_request_id: self.by_request_id,
detached_bodies: self.detached_bodies,
provider_usage_windows: RwLock::new(items.into_iter().collect()),
auth_api_keys: self.auth_api_keys,
provider_catalog: self.provider_catalog,
analytics_users: self.analytics_users,
analytics_candidates: self.analytics_candidates,
analytics_allocations: self.analytics_allocations,
}
}
@@ -244,7 +272,48 @@ fn usage_has_admin_unknown_model_or_provider(item: &StoredRequestUsageAudit) ->
usage_admin_unknown_label(&item.model) || usage_admin_unknown_label(&item.provider_name)
}
fn usage_matches_list_query(item: &StoredRequestUsageAudit, query: &UsageAuditListQuery) -> bool {
fn usage_matches_list_query(
item: &StoredRequestUsageAudit,
query: &UsageAuditListQuery,
keys: &BTreeMap<String, bool>,
) -> bool {
if query
.provider_id
.as_ref()
.is_some_and(|id| item.provider_id.as_ref() != Some(id))
|| query
.endpoint_kind
.as_ref()
.is_some_and(|value| item.endpoint_kind.as_ref() != Some(value))
|| query
.request_type
.as_ref()
.is_some_and(|value| item.request_type.as_ref() != Some(value))
|| query
.slow_threshold_ms
.is_some_and(|value| item.response_time_ms.is_none_or(|latency| latency < value))
|| query
.has_format_conversion
.is_some_and(|value| item.has_format_conversion != value)
|| query
.api_key_id
.as_ref()
.is_some_and(|id| item.api_key_id.as_ref() != Some(id))
|| query
.request_id
.as_ref()
.is_some_and(|id| item.request_id != *id)
|| query
.attribution_kind
.as_deref()
.is_some_and(|kind| analytics::attribution(item, keys) != kind)
|| query
.actor_user_id
.as_deref()
.is_some_and(|id| analytics::actor(item, keys) != Some(id))
{
return false;
}
// The field is historically named `created_at_unix_ms`, but usage audit rows
// across gateway handlers, SQL repositories and tests are stored as epoch seconds.
if let Some(created_from_unix_secs) = query.created_from_unix_secs {
@@ -332,7 +401,45 @@ fn usage_matches_list_query(item: &StoredRequestUsageAudit, query: &UsageAuditLi
fn usage_matches_keyword_search_query(
item: &StoredRequestUsageAudit,
query: &UsageAuditKeywordSearchQuery,
keys: &BTreeMap<String, bool>,
) -> bool {
if query
.provider_id
.as_ref()
.is_some_and(|id| item.provider_id.as_ref() != Some(id))
|| query
.endpoint_kind
.as_ref()
.is_some_and(|value| item.endpoint_kind.as_ref() != Some(value))
|| query
.request_type
.as_ref()
.is_some_and(|value| item.request_type.as_ref() != Some(value))
|| query
.slow_threshold_ms
.is_some_and(|value| item.response_time_ms.is_none_or(|latency| latency < value))
|| query
.has_format_conversion
.is_some_and(|value| item.has_format_conversion != value)
|| query
.api_key_id
.as_ref()
.is_some_and(|id| item.api_key_id.as_ref() != Some(id))
|| query
.request_id
.as_ref()
.is_some_and(|id| item.request_id != *id)
|| query
.attribution_kind
.as_deref()
.is_some_and(|kind| analytics::attribution(item, keys) != kind)
|| query
.actor_user_id
.as_deref()
.is_some_and(|id| analytics::actor(item, keys) != Some(id))
{
return false;
}
if let Some(created_from_unix_secs) = query.created_from_unix_secs {
if item.created_at_unix_ms < created_from_unix_secs {
return false;
@@ -1127,6 +1234,38 @@ fn usage_provider_aggregation_identity(
#[async_trait]
impl UsageReadRepository for InMemoryUsageReadRepository {
async fn query_dashboard_summary(
&self,
query: &aether_data_contracts::repository::usage::UsageDashboardAnalyticsQuery,
) -> Result<aether_data_contracts::repository::usage::StoredDashboardSummary, DataLayerError> {
self.dashboard_summary_query(query)
}
async fn query_dashboard_analytics(
&self,
query: &aether_data_contracts::repository::usage::UsageDashboardAnalyticsQuery,
) -> Result<
aether_data_contracts::repository::usage::StoredUsageDashboardAnalytics,
DataLayerError,
> {
self.dashboard_analytics_query(query)
}
async fn query_usage_analytics(
&self,
query: &aether_data_contracts::repository::usage::UsageAnalyticsQuery,
) -> Result<aether_data_contracts::repository::usage::StoredUsageAnalytics, DataLayerError>
{
self.analytics_query(query)
}
async fn summarize_health_observations(
&self,
query: &aether_data_contracts::repository::usage::HealthObservationQuery,
) -> Result<aether_data_contracts::repository::usage::HealthObservationSummary, DataLayerError>
{
self.health_observations(query)
}
async fn find_by_id(
&self,
id: &str,
@@ -1205,12 +1344,13 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
&self,
query: &UsageAuditListQuery,
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
let keys = self.analytics_key_flags();
let mut items: Vec<_> = self
.by_request_id
.read()
.expect("usage repository lock")
.values()
.filter(|item| usage_matches_list_query(item, query))
.filter(|item| usage_matches_list_query(item, query, &keys))
.cloned()
.collect();
sort_usage_items(&mut items, query.newest_first);
@@ -1231,12 +1371,13 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
&self,
query: &UsageAuditKeywordSearchQuery,
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
let keys = self.analytics_key_flags();
let mut items: Vec<_> = self
.by_request_id
.read()
.expect("usage repository lock")
.values()
.filter(|item| usage_matches_keyword_search_query(item, query))
.filter(|item| usage_matches_keyword_search_query(item, query, &keys))
.cloned()
.collect();
sort_usage_items(&mut items, query.newest_first);
@@ -1254,12 +1395,13 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
}
async fn count_usage_audits(&self, query: &UsageAuditListQuery) -> Result<u64, DataLayerError> {
let keys = self.analytics_key_flags();
Ok(self
.by_request_id
.read()
.expect("usage repository lock")
.values()
.filter(|item| usage_matches_list_query(item, query))
.filter(|item| usage_matches_list_query(item, query, &keys))
.count() as u64)
}
@@ -1267,12 +1409,13 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
&self,
query: &UsageAuditKeywordSearchQuery,
) -> Result<u64, DataLayerError> {
let keys = self.analytics_key_flags();
Ok(self
.by_request_id
.read()
.expect("usage repository lock")
.values()
.filter(|item| usage_matches_keyword_search_query(item, query))
.filter(|item| usage_matches_keyword_search_query(item, query, &keys))
.count() as u64)
}
@@ -3281,6 +3424,8 @@ impl UsageWriteRepository for InMemoryUsageReadRepository {
finalized_at_unix_secs: usage.finalized_at_unix_secs,
};
self.dashboard_projection.write().expect("dashboard projection lock")
.record(&stored, &self.analytics_key_flags());
by_request_id.insert(stored.request_id.clone(), stored.clone());
if let Some(auth_api_keys) = self.auth_api_keys.as_ref() {
let before_contribution = existing.as_ref().and_then(api_key_usage_contribution);
crates/aether-data/runtime/src/repository/usage/memory/analytics.rs
+1080
View File
File diff suppressed because it is too large Load Diff
crates/aether-data/runtime/src/repository/usage/memory/dashboard_summary.rs
+228
View File
@@ -0,0 +1,228 @@
use super::{
analytics, usage_cache_creation_tokens, usage_total_input_context, usage_total_tokens,
InMemoryUsageReadRepository, StoredRequestUsageAudit,
};
use aether_data_contracts::{repository::usage::*, DataLayerError};
use chrono::{DateTime, NaiveDate, Utc};
use std::collections::{BTreeMap, BTreeSet};
#[derive(Debug)]
pub(super) struct DashboardProjection {
pub since: DateTime<Utc>,
entries: BTreeMap<String, Contribution>,
}
impl Default for DashboardProjection {
fn default() -> Self {
Self {
since: Utc::now(),
entries: BTreeMap::new(),
}
}
}
#[derive(Debug)]
struct Contribution {
at: DateTime<Utc>,
actor: Option<String>,
metrics: DashboardSummaryMetrics,
billable_units: Option<i128>,
}
impl DashboardProjection {
pub fn record(&mut self, row: &StoredRequestUsageAudit, keys: &BTreeMap<String, bool>) {
let Some(at) = DateTime::from_timestamp(row.created_at_unix_ms as i64, 0) else {
return;
};
if at < self.since {
return;
}
if row
.request_metadata
.as_ref()
.and_then(|m| m.pointer("/analytics_attribution/record_kind"))
.and_then(|v| v.as_str())
== Some("session")
{
self.entries.remove(&row.request_id);
return;
}
let available = |key| {
row.request_metadata
.as_ref()
.and_then(|m| m.get(key))
.and_then(|v| v.as_bool())
!= Some(false)
};
let usage = available(USAGE_AVAILABLE_METADATA_KEY);
let priced =
available(USAGE_PRICING_AVAILABLE_METADATA_KEY) && row.billing_status == "settled";
let stream = row
.request_metadata
.as_ref()
.and_then(|m| m.get("upstream_is_stream"))
.and_then(|v| v.as_bool())
.unwrap_or(row.is_stream);
let metrics = DashboardSummaryMetrics {
request_count: 1,
input_tokens: if usage { row.input_tokens } else { 0 },
output_tokens: if usage { row.output_tokens } else { 0 },
total_tokens: if usage {
if row.total_tokens > 0 {
row.total_tokens
} else {
usage_total_tokens(row)
}
} else {
0
},
usage_available_count: u64::from(usage),
pricing_available_count: u64::from(priced),
cache_read_tokens: if usage {
row.cache_read_input_tokens
} else {
0
},
cache_creation_tokens: if usage {
usage_cache_creation_tokens(row)
} else {
0
},
cache_input_tokens: if usage {
usage_total_input_context(row)
} else {
0
},
first_byte_sum_ms: row.first_byte_time_ms.unwrap_or(0) as f64,
first_byte_sample_count: u64::from(row.first_byte_time_ms.is_some()),
response_sum_ms: row.response_time_ms.unwrap_or(0) as f64,
response_sample_count: u64::from(row.response_time_ms.is_some()),
stream_requests: u64::from(stream),
standard_requests: u64::from(!stream),
..Default::default()
};
self.entries.insert(
row.request_id.clone(),
Contribution {
at,
actor: analytics::actor(row, keys).map(str::to_owned),
metrics,
billable_units: priced
.then(|| (row.actual_total_cost_usd * 100_000_000.0).round() as i128),
},
);
}
}
fn sum_metrics<'a>(rows: impl Iterator<Item = &'a Contribution>) -> DashboardSummaryMetrics {
let mut sum = DashboardSummaryMetrics::default();
let mut users = BTreeSet::new();
let mut units = 0_i128;
for row in rows {
let m = &row.metrics;
sum.request_count += m.request_count;
sum.input_tokens += m.input_tokens;
sum.output_tokens += m.output_tokens;
sum.total_tokens += m.total_tokens;
sum.usage_available_count += m.usage_available_count;
sum.pricing_available_count += m.pricing_available_count;
sum.cache_read_tokens += m.cache_read_tokens;
sum.cache_creation_tokens += m.cache_creation_tokens;
sum.cache_input_tokens += m.cache_input_tokens;
sum.first_byte_sum_ms += m.first_byte_sum_ms;
sum.first_byte_sample_count += m.first_byte_sample_count;
sum.response_sum_ms += m.response_sum_ms;
sum.response_sample_count += m.response_sample_count;
sum.stream_requests += m.stream_requests;
sum.standard_requests += m.standard_requests;
units += row.billable_units.unwrap_or(0);
if let Some(actor) = row.actor.as_deref() {
users.insert(actor);
}
}
sum.active_users = users.len() as u64;
if sum.request_count == 0 || sum.pricing_available_count > 0 {
sum.billable_amount = Some(format!(
"{}{}.{:08}",
if units < 0 { "-" } else { "" },
units.abs() / 100_000_000,
units.abs() % 100_000_000
));
}
sum
}
impl InMemoryUsageReadRepository {
/// Test/embedded initialization boundary; seeded older audit rows stay excluded.
pub fn with_dashboard_stats_since(self, since: DateTime<Utc>) -> Self {
let keys = self.analytics_key_flags();
let mut projection = self
.dashboard_projection
.write()
.expect("dashboard projection lock");
projection.since = since;
projection.entries.clear();
for row in self
.by_request_id
.read()
.expect("usage repository lock")
.values()
{
projection.record(row, &keys);
}
drop(projection);
self
}
pub(super) fn dashboard_summary_query(
&self,
query: &UsageDashboardAnalyticsQuery,
) -> Result<StoredDashboardSummary, DataLayerError> {
query.validate()?;
let projection = self.dashboard_projection.read().map_err(|_| {
DataLayerError::UnexpectedValue("dashboard projection lock poisoned".into())
})?;
let now = Utc::now();
let today_from = query.today_start(now)?.max(projection.since);
let tz = query
.timezone
.parse::<chrono_tz::Tz>()
.map_err(|_| DataLayerError::InvalidInput("invalid timezone".into()))?;
let rows = || projection.entries.values().filter(|row| row.at < now);
let today = sum_metrics(rows().filter(|row| row.at >= today_from));
let total = sum_metrics(rows());
let mut days = BTreeMap::<NaiveDate, u64>::new();
for row in rows() {
*days
.entry(row.at.with_timezone(&tz).date_naive())
.or_default() += 1;
}
let active_days = days.len() as u64;
let local_today = now.with_timezone(&tz).date_naive();
let consecutive_active_days =
dashboard_consecutive_active_days(days.keys().copied(), local_today);
let earliest_day = local_today - chrono::Duration::days(364);
let activity_days = days
.into_iter()
.filter(|(day, _)| day >= &earliest_day)
.map(|(date, requests)| DashboardActivityDay {
date: date.to_string(),
requests,
})
.collect();
let users = self
.analytics_users
.read()
.map_err(|_| DataLayerError::UnexpectedValue("users lock poisoned".into()))?;
Ok(StoredDashboardSummary {
stats_since: projection.since.to_rfc3339(),
generated_at: now.to_rfc3339(),
timezone: query.timezone.clone(),
today_from: today_from.to_rfc3339(),
window_seconds: (now - today_from).num_milliseconds().max(0) as f64 / 1000.0,
today,
total,
users: DashboardUserCounts {
total: users.iter().filter(|user| !user.is_deleted).count() as u64,
..Default::default()
},
active_days,
consecutive_active_days,
activity_days,
})
}
}
crates/aether-data/runtime/src/repository/usage/memory/tests.rs
+576
View File
@@ -22,6 +22,486 @@ use aether_data_contracts::repository::usage::{
};
use serde_json::json;
#[tokio::test]
async fn overview_model_performance_merges_provider_samples_without_pagination() {
use aether_data_contracts::repository::usage::*;
let at = chrono::DateTime::parse_from_rfc3339("2026-09-12T10:05:00Z").unwrap();
let mut records = Vec::new();
for (index, provider, first_byte, response_time, output_tokens) in [
(0, "provider-a", 100, 1100, 100),
(1, "provider-a", 300, 1300, 200),
(2, "provider-b", 500, 2500, 400),
] {
let mut row = sample_usage(&format!("model-sample-{index}"), at.timestamp());
row.provider_id = Some(provider.into());
row.model = "shared-model".into();
row.target_model = Some(format!("{provider}-deployment"));
row.first_byte_time_ms = Some(first_byte);
row.response_time_ms = Some(response_time);
row.output_tokens = output_tokens;
row.is_stream = true;
records.push(row);
}
let mut failed = sample_usage("model-failed", at.timestamp());
failed.model = "shared-model".into();
failed.status = "failed".into();
failed.first_byte_time_ms = None;
failed.response_time_ms = None;
records.push(failed);
let mut pending = sample_usage("model-pending", at.timestamp());
pending.model = "pending-model".into();
pending.status = "pending".into();
pending.first_byte_time_ms = None;
pending.response_time_ms = None;
records.push(pending);
let repo = InMemoryUsageReadRepository::seed(records);
let query = UsageAnalyticsQuery {
from_unix_ms: (at - chrono::Duration::minutes(5)).timestamp_millis() as u64,
to_unix_ms: (at + chrono::Duration::minutes(55)).timestamp_millis() as u64,
timezone: "UTC".into(),
view: UsageAnalyticsView::Performance,
limit: 1,
offset: 1,
..Default::default()
};
let result = repo.query_usage_analytics(&query).await.unwrap();
assert_eq!(result.model_rows.len(), 2);
assert_eq!(result.model_rows[0].id.as_deref(), Some("shared-model"));
assert!(result
.model_rows
.iter()
.all(|row| row.bucket_start.is_none()));
let metrics = &result.model_rows[0].metrics;
assert_eq!(metrics.request_count, 4);
assert_eq!(metrics.successful_request_count, 3);
assert_eq!(metrics.failed_request_count, 1);
assert_eq!(metrics.first_byte_sample_count, 3);
assert_eq!(metrics.first_byte_sum_ms, 900.0);
assert_eq!(metrics.latency_sample_count, 3);
assert_eq!(metrics.latency_sum_ms, 4900.0);
assert_eq!(metrics.output_tps_sample_count, 3);
assert_eq!(metrics.output_tps_sum, 500.0);
assert_eq!(result.model_rows[1].metrics.first_byte_sample_count, 0);
assert_eq!(result.model_rows[1].metrics.in_flight_request_count, 1);
assert_eq!(
result
.model_rows
.iter()
.map(|row| row.metrics.request_count)
.sum::<u64>(),
result.summary.request_count
);
let filtered = repo
.query_usage_analytics(&UsageAnalyticsQuery {
model: Some("shared-model".into()),
..query
})
.await
.unwrap();
assert_eq!(filtered.model_rows, vec![result.model_rows[0].clone()]);
}
#[tokio::test]
async fn overview_memory_chart_hour_buckets_are_utc_in_half_hour_zones() {
use aether_data_contracts::repository::usage::*;
let at = chrono::DateTime::parse_from_rfc3339("2026-09-12T10:05:00Z").unwrap();
let repo = InMemoryUsageReadRepository::seed([sample_usage("hour-zone", at.timestamp())]);
let query = UsageAnalyticsQuery {
from_unix_ms: (at - chrono::Duration::minutes(5)).timestamp_millis() as u64,
to_unix_ms: (at + chrono::Duration::minutes(55)).timestamp_millis() as u64,
timezone: "Asia/Kolkata".into(),
view: UsageAnalyticsView::DashboardCharts,
granularity: UsageAnalyticsGranularity::Hour,
limit: 1,
..Default::default()
};
let charts = repo.query_usage_analytics(&query).await.unwrap();
assert_eq!(charts.summary.request_count, 1);
assert_eq!(charts.rows.len(), 1);
assert_eq!(charts.rows[0].metrics.request_count, 1);
assert_eq!(
charts.rows[0].bucket_start,
charts.model_rows[0].bucket_start
);
}
#[tokio::test]
async fn overview_memory_chart_days_survive_skipped_midnight() {
use aether_data_contracts::repository::usage::*;
let moments = [
"2026-09-05T12:00:00Z",
"2026-09-06T12:00:00Z",
"2026-09-07T12:00:00Z",
];
let repo = InMemoryUsageReadRepository::seed(moments.map(|moment| {
sample_usage(
moment,
chrono::DateTime::parse_from_rfc3339(moment)
.unwrap()
.timestamp(),
)
}));
let query = UsageAnalyticsQuery {
from_unix_ms: chrono::DateTime::parse_from_rfc3339("2026-09-05T04:00:00Z")
.unwrap()
.timestamp_millis() as u64,
to_unix_ms: chrono::DateTime::parse_from_rfc3339("2026-09-08T03:00:00Z")
.unwrap()
.timestamp_millis() as u64,
timezone: "America/Santiago".into(),
view: UsageAnalyticsView::DashboardCharts,
limit: 1,
..Default::default()
};
let charts = repo.query_usage_analytics(&query).await.unwrap();
assert_eq!(charts.summary.request_count, 3);
assert_eq!(charts.rows.len(), 3);
assert_eq!(charts.model_rows.len(), 3);
for (series, model) in charts.rows.iter().zip(&charts.model_rows) {
assert_eq!(series.metrics.request_count, 1);
assert_eq!(series.bucket_start, model.bucket_start);
assert_eq!(
series.metrics.billable_amount,
model.metrics.billable_amount
);
}
}
#[tokio::test]
async fn overview_memory_dashboard_keeps_all_history_and_chart_dimensions() {
use aether_data_contracts::repository::usage::*;
let now = chrono::Utc::now();
let mut old = sample_usage(
"old-dashboard",
(now - chrono::Duration::days(800)).timestamp(),
);
old.actual_total_cost_usd = 2.0;
old.model = "old-model".into();
let mut current = sample_usage("current-dashboard", now.timestamp());
current.actual_total_cost_usd = 0.5;
current.model = "new-model".into();
let repo = InMemoryUsageReadRepository::seed([old, current]);
let dashboard = repo
.query_dashboard_analytics(&UsageDashboardAnalyticsQuery {
timezone: "Asia/Shanghai".into(),
})
.await
.unwrap();
assert_eq!(dashboard.today.summary.request_count, 1);
assert_eq!(dashboard.total.summary.request_count, 2);
assert_eq!(
dashboard.today.summary.billable_amount.as_deref(),
Some("0.50000000")
);
assert_eq!(
dashboard.total.summary.billable_amount.as_deref(),
Some("2.50000000")
);
assert_eq!(dashboard.today.read_revision, dashboard.total.read_revision);
assert_eq!(dashboard.history_complete, None);
let charts = repo
.query_usage_analytics(&UsageAnalyticsQuery {
from_unix_ms: (now - chrono::Duration::days(1)).timestamp_millis() as u64,
to_unix_ms: (now + chrono::Duration::seconds(1)).timestamp_millis() as u64,
timezone: "Asia/Shanghai".into(),
view: UsageAnalyticsView::DashboardCharts,
limit: 1,
..Default::default()
})
.await
.unwrap();
assert_eq!(charts.model_rows.len(), 1);
assert_eq!(charts.model_rows[0].id.as_deref(), Some("new-model"));
assert_eq!(charts.provider_rows.len(), 1);
assert_eq!(
charts.model_rows[0].metrics.billable_amount,
charts.summary.billable_amount
);
}
#[tokio::test]
async fn overview_memory_dashboard_total_keeps_card_coverage_without_historical_diagnostics() {
use aether_data_contracts::repository::usage::*;
let now = chrono::Utc::now();
let mut known = sample_usage("dashboard-total-known", now.timestamp());
known.request_metadata = Some(json!({"analytics_attribution":{"is_standalone":false}}));
let mut missing = sample_usage(
"dashboard-total-missing",
(now - chrono::Duration::days(800)).timestamp(),
);
missing.billing_status = "pending".into();
missing.request_metadata =
Some(json!({"usage_available":false,"usage_pricing_available":false}));
let mut session = sample_usage("dashboard-total-session", now.timestamp());
session.request_metadata = Some(json!({"analytics_attribution":{"record_kind":"session"}}));
let future = sample_usage(
"dashboard-total-future",
(now + chrono::Duration::days(1)).timestamp(),
);
let repo = InMemoryUsageReadRepository::seed([known, missing, session, future])
.with_analytics_allocations([UsageAnalyticsAllocation {
request_id: "dashboard-total-known".into(),
complete: true,
wallet_debit_amount: Some("0.18000000".into()),
..Default::default()
}]);
let dashboard = repo
.query_dashboard_analytics(&UsageDashboardAnalyticsQuery {
timezone: "UTC".into(),
})
.await
.unwrap();
let total = dashboard.total.summary;
assert_eq!(total.request_count, 2);
assert_eq!(total.total_tokens, 150);
assert_eq!(total.billable_amount.as_deref(), Some("0.18000000"));
assert_eq!(total.usage_available_count, 1);
assert_eq!(total.pricing_available_count, 1);
assert_eq!(total.settled_count, 1);
assert_eq!(total.allocation_available_count, 1);
assert!(total.latency_p95_ms.is_none());
assert!(total.wallet_debit_amount.is_none());
assert_eq!(dashboard.today.summary.latency_p95_ms, Some(420.0));
assert_eq!(dashboard.today.summary.successful_request_count, 1);
assert_eq!(
dashboard.today.summary.wallet_debit_amount.as_deref(),
Some("0.18000000")
);
}
#[tokio::test]
async fn overview_memory_employee_roster_and_allocations_are_global() {
use aether_data_contracts::repository::usage::*;
use aether_data_contracts::repository::users::StoredUserSummary;
let mut usage = sample_usage("overview-request", 1_700_000_000);
usage.user_id = Some("alice".into());
usage.request_metadata = Some(json!({"analytics_attribution":{"is_standalone":false}}));
usage.billing_status = "settled".into();
usage.actual_total_cost_usd = 0.00000001;
let repo = InMemoryUsageReadRepository::seed([usage])
.with_analytics_users([
StoredUserSummary::new(
"alice".into(),
"Alice".into(),
None,
"user".into(),
true,
false,
)
.unwrap(),
StoredUserSummary::new(
"zero".into(),
"Zero".into(),
None,
"user".into(),
true,
false,
)
.unwrap(),
])
.with_analytics_allocations([UsageAnalyticsAllocation {
request_id: "overview-request".into(),
quota_covered_amount: Some("0.00000000".into()),
wallet_consumed_amount: Some("0.00000001".into()),
wallet_debit_amount: Some("0.00000000".into()),
complete: true,
..Default::default()
}]);
let mut query = UsageAnalyticsQuery {
from_unix_ms: 1_700_000_000_000,
to_unix_ms: 1_700_000_060_000,
timezone: "UTC".into(),
view: UsageAnalyticsView::Users,
limit: 1,
descending: true,
..Default::default()
};
let first = repo.query_usage_analytics(&query).await.unwrap();
assert_eq!(first.total, 2);
assert_eq!(first.users[0].user_id, "alice");
assert_eq!(first.user_summary.as_ref().unwrap().user_count, 2);
assert_eq!(first.user_summary.as_ref().unwrap().active_user_count, 1);
assert!(first.user_finance_summary.is_none());
assert!(first.user_payments.is_none());
assert!(first.users[0].finance.is_none());
assert_eq!(
first.summary.wallet_consumed_amount.as_deref(),
Some("0.00000001")
);
assert_eq!(
first.summary.wallet_debit_amount.as_deref(),
Some("0.00000000")
);
query.offset = 1;
let second = repo.query_usage_analytics(&query).await.unwrap();
assert_eq!(second.users[0].user_id, "zero");
assert_eq!(second.users[0].metrics.request_count, 0);
assert_eq!(second.user_summary, first.user_summary);
let searched = repo
.query_usage_analytics(&UsageAnalyticsQuery {
search: Some("Zero".into()),
offset: 0,
..query.clone()
})
.await
.unwrap();
assert_eq!(searched.user_summary.as_ref().unwrap().user_count, 1);
assert_eq!(searched.user_summary.as_ref().unwrap().active_user_count, 0);
assert_eq!(searched.summary.request_count, 0);
query.from_unix_ms = query.to_unix_ms;
query.to_unix_ms += 60_000;
let outside = repo.query_usage_analytics(&query).await.unwrap();
assert_eq!(outside.summary.request_count, 0);
}
#[tokio::test]
async fn overview_memory_employee_is_grouped_by_account_owner() {
use aether_data_contracts::repository::{usage::*, users::StoredUserSummary};
let mut usage = sample_usage("trusted-request", 1_700_000_000);
usage.user_id = Some("owner".into());
usage.request_metadata =
Some(json!({"analytics_attribution":{"is_standalone":false,"actor_user_id":"actor"}}));
let repo = InMemoryUsageReadRepository::seed([usage]).with_analytics_users(
["owner", "actor"].map(|id| {
StoredUserSummary::new(id.into(), id.into(), None, "user".into(), true, false).unwrap()
}),
);
let query = UsageAnalyticsQuery {
from_unix_ms: 1_700_000_000_000,
to_unix_ms: 1_700_000_060_000,
timezone: "UTC".into(),
view: UsageAnalyticsView::Users,
attribution_kind: Some("employee".into()),
has_usage: Some(true),
limit: 100,
..Default::default()
};
let result = repo.query_usage_analytics(&query).await.unwrap();
assert_eq!(result.total, 1);
assert_eq!(result.users[0].user_id, "owner");
let detail = repo
.query_usage_analytics(&UsageAnalyticsQuery {
actor_user_id: Some("owner".into()),
..query
})
.await
.unwrap();
assert_eq!(detail.users[0], result.users[0]);
}
#[tokio::test]
async fn overview_memory_legacy_requests_use_existing_key_account_flags() {
use aether_data_contracts::repository::{usage::*, users::StoredUserSummary};
let snapshots = [("member-key", false), ("standalone-key", true)].map(|(id, standalone)| {
(
None,
StoredAuthApiKeySnapshot::new(
"user-1".into(),
"alice".into(),
None,
"user".into(),
"local".into(),
true,
false,
None,
None,
None,
id.into(),
None,
true,
false,
standalone,
None,
None,
None,
None,
None,
None,
)
.unwrap(),
)
});
let keys = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(snapshots));
let rows = ["member-key", "standalone-key", "deleted-key"].map(|id| {
let mut usage = sample_usage(id, 1_700_000_000);
usage.api_key_id = Some(id.into());
usage.request_metadata = None;
usage
});
let repo = InMemoryUsageReadRepository::seed(rows)
.with_auth_api_key_repository(keys)
.with_analytics_users([StoredUserSummary::new(
"user-1".into(),
"alice".into(),
None,
"user".into(),
true,
false,
)
.unwrap()]);
let mut query = UsageAnalyticsQuery {
from_unix_ms: 1_700_000_000_000,
to_unix_ms: 1_700_000_060_000,
timezone: "UTC".into(),
view: UsageAnalyticsView::Consumption,
limit: 100,
..Default::default()
};
let result = repo.query_usage_analytics(&query).await.unwrap();
assert_eq!(result.summary.request_count, 3);
assert_eq!(result.summary.usage_active_users, 1);
assert_eq!(result.summary.trusted_attribution_count, 1);
for (key, kind, source, user) in [
("member-key", "employee", "user_account", Some("user-1")),
("standalone-key", "standalone", "standalone_key", None),
("deleted-key", "unknown", "unknown", None),
] {
let row = result
.consumption
.iter()
.find(|row| row.request_id == key)
.unwrap();
assert_eq!(row.attribution_kind, kind);
assert_eq!(row.attribution_source, source);
assert_eq!(row.user_id.as_deref(), user);
assert_eq!(row.credential_owner_id.as_deref(), Some("user-1"));
}
query.view = UsageAnalyticsView::Users;
query.attribution_kind = Some("employee".into());
let employees = repo.query_usage_analytics(&query).await.unwrap();
assert_eq!(employees.users[0].metrics.request_count, 1);
assert_eq!(employees.summary.request_count, 1);
}
#[tokio::test]
async fn overview_memory_health_does_not_treat_upstream_cancellation_as_client() {
use aether_data_contracts::repository::usage::*;
let mut upstream = sample_usage("upstream-cancel", 1_700_000_000);
upstream.status = "cancelled".into();
upstream.request_metadata =
Some(json!({"analytics_failure":{"origin":"upstream","reason":"provider_cancelled"}}));
let mut client = upstream.clone();
client.request_id = "client-cancel".into();
client.request_metadata =
Some(json!({"analytics_failure":{"origin":"client","reason":"downstream_disconnect"}}));
let repo = InMemoryUsageReadRepository::seed([upstream, client]);
let summary = repo
.summarize_health_observations(&HealthObservationQuery {
from_unix_ms: 1_700_000_000_000,
to_unix_ms: 1_700_000_060_000,
object_kind: HealthObservationObjectKind::Model,
object_values: None,
segments: 4,
})
.await
.unwrap();
assert_eq!(summary.overall.request_count, 2);
assert_eq!(summary.overall.service_failed_count, 1);
assert_eq!(summary.overall.excluded_count, 1);
}
fn sample_usage(request_id: &str, created_at_unix_ms: i64) -> StoredRequestUsageAudit {
StoredRequestUsageAudit::new(
"usage-1".to_string(),
@@ -2763,3 +3243,99 @@ async fn summarize_usage_provider_performance_computes_tps_and_top_provider_time
assert_eq!(without_timeline.providers, summary.providers);
assert!(without_timeline.timeline.is_empty());
}
#[tokio::test]
async fn future_dashboard_summary_excludes_old_rows_and_preserves_canonical_cache_samples() {
use aether_data_contracts::repository::usage::UsageDashboardAnalyticsQuery;
let now = chrono::Utc::now();
let since = now - chrono::Duration::minutes(5);
let mut included = sample_usage(
"future-summary",
(now - chrono::Duration::seconds(1)).timestamp(),
);
included.api_key_id = None;
included.api_format = Some("claude:messages".into());
included.endpoint_api_format = Some("claude:messages".into());
included.input_tokens = 80;
included.output_tokens = 20;
included.total_tokens = 150;
included.cache_read_input_tokens = 40;
included.cache_creation_input_tokens = 10;
included.response_time_ms = Some(800);
included.first_byte_time_ms = Some(100);
included.is_stream = false;
included.request_metadata = Some(json!({"upstream_is_stream": true}));
included.actual_total_cost_usd = 0.12345678;
included.billing_status = "settled".into();
let mut excluded = included.clone();
excluded.request_id = "before-activation".into();
excluded.created_at_unix_ms = (since - chrono::Duration::days(500)).timestamp() as u64;
let repo =
InMemoryUsageReadRepository::seed([included, excluded]).with_dashboard_stats_since(since);
let first = repo
.query_dashboard_summary(&UsageDashboardAnalyticsQuery {
timezone: "Asia/Kathmandu".into(),
})
.await
.unwrap();
assert_eq!(first.total.request_count, 1);
assert_eq!(first.total.total_tokens, 150);
assert_eq!(first.total.billable_amount.as_deref(), Some("0.12345678"));
assert_eq!(first.today.cache_input_tokens, 130);
assert_eq!(first.today.cache_read_tokens, 40);
assert_eq!(first.today.first_byte_sample_count, 1);
assert_eq!(first.today.response_sample_count, 1);
assert_eq!(first.today.stream_requests, 1);
assert_eq!(first.today.standard_requests, 0);
assert_eq!(first.active_days, 1);
assert_eq!(first.consecutive_active_days, 1);
assert_eq!(
first.activity_days.iter().map(|d| d.requests).sum::<u64>(),
1
);
// Audit retention does not own the additive projection.
repo.by_request_id.write().unwrap().clear();
let retained = repo
.query_dashboard_summary(&UsageDashboardAnalyticsQuery {
timezone: "Asia/Kathmandu".into(),
})
.await
.unwrap();
assert_eq!(retained.total, first.total);
assert_eq!(retained.stats_since, since.to_rfc3339());
}
#[tokio::test]
async fn future_dashboard_summary_streak_uses_local_days_before_heatmap_truncation() {
use aether_data_contracts::repository::usage::UsageDashboardAnalyticsQuery;
use chrono::TimeZone;
let now = chrono::Utc::now();
let tz = chrono_tz::Asia::Kathmandu;
let today = now.with_timezone(&tz).date_naive();
let rows = (1..=400).map(|offset| {
let day = today - chrono::Duration::days(offset);
let at = tz
.from_local_datetime(&day.and_hms_opt(12, 0, 0).unwrap())
.single()
.unwrap();
sample_usage(&format!("streak-{offset}"), at.timestamp())
});
let repo = InMemoryUsageReadRepository::seed(rows)
.with_dashboard_stats_since(now - chrono::Duration::days(401));
let summary = repo
.query_dashboard_summary(&UsageDashboardAnalyticsQuery {
timezone: tz.to_string(),
})
.await
.unwrap();
assert_eq!(summary.today.request_count, 0);
assert_eq!(summary.active_days, 400);
assert_eq!(summary.consecutive_active_days, 400);
assert_eq!(summary.activity_days.len(), 364);
assert_eq!(
summary.activity_days.last().unwrap().date,
today.pred_opt().unwrap().to_string()
);
}
crates/aether-data/runtime/src/repository/wallet/memory.rs
+32
View File
@@ -831,6 +831,12 @@ impl WalletReadRepository for InMemoryWalletRepository {
.as_deref()
.is_none_or(|expected| wallet.status == expected)
})
.filter(|wallet| {
query
.user_id
.as_deref()
.is_none_or(|expected| wallet.user_id.as_deref() == Some(expected))
})
.filter(|wallet| match query.owner_type.as_deref() {
Some("user") => wallet.user_id.is_some(),
Some("api_key") => wallet.api_key_id.is_some(),
@@ -3418,6 +3424,7 @@ mod tests {
let page = repository
.list_admin_wallets(&AdminWalletListQuery {
user_id: None,
status: Some("active".to_string()),
owner_type: Some("api_key".to_string()),
limit: 1,
@@ -3430,6 +3437,31 @@ mod tests {
assert_eq!(page.items.len(), 1);
assert_eq!(page.items[0].id, "wallet-3");
assert_eq!(page.items[0].updated_at_unix_secs, Some(110));
let query = AdminWalletListQuery {
user_id: Some("user-2".to_string()),
owner_type: Some("user".to_string()),
limit: 1,
..Default::default()
};
let selected = repository.list_admin_wallets(&query).await.unwrap();
assert_eq!(selected.total, 1);
assert_eq!(selected.items[0].id, "wallet-2");
let missing = repository
.list_admin_wallets(&AdminWalletListQuery {
user_id: Some("missing-user".to_string()),
..query.clone()
})
.await
.unwrap();
assert_eq!(missing.total, 0);
assert!(missing.items.is_empty());
let beyond_page = repository
.list_admin_wallets(&AdminWalletListQuery { offset: 1, ..query })
.await
.unwrap();
assert_eq!(beyond_page.total, 1);
assert!(beyond_page.items.is_empty());
}
#[tokio::test]