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

Merge remote-tracking branch 'origin/aether-rust-pioneer' into payment-billing-plans

Browse Source 
# Conflicts:
#	crates/aether-data/src/lifecycle/bootstrap/postgres.rs
#	crates/aether-data/src/lifecycle/migrate/tests.rs
This commit is contained in:
Entropy.Xu
2026-05-13 01:29:18 +08:00
parent 10285c5eb9 4387a9cdd5
commit 3fb007a56d
185 changed files with 14980 additions and 3638 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
crates/aether-data/migrations/mysql/20260403000000_baseline.sql
View File

@@ -294,6 +294,36 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
KEY provider_api_keys_provider_id_idx (provider_id)
);
CREATE TABLE IF NOT EXISTS pool_member_scores (
id VARCHAR(192) PRIMARY KEY,
pool_kind VARCHAR(64) NOT NULL,
pool_id VARCHAR(64) NOT NULL,
member_kind VARCHAR(64) NOT NULL,
member_id VARCHAR(64) NOT NULL,
capability VARCHAR(64) NOT NULL,
scope_kind VARCHAR(64) NOT NULL,
scope_id VARCHAR(128),
score DOUBLE NOT NULL DEFAULT 0,
hard_state VARCHAR(64) NOT NULL DEFAULT 'unknown',
score_version BIGINT NOT NULL DEFAULT 1,
score_reason TEXT NOT NULL,
last_ranked_at BIGINT,
last_scheduled_at BIGINT,
last_success_at BIGINT,
last_failure_at BIGINT,
failure_count BIGINT NOT NULL DEFAULT 0,
last_probe_attempt_at BIGINT,
last_probe_success_at BIGINT,
last_probe_failure_at BIGINT,
probe_failure_count BIGINT NOT NULL DEFAULT 0,
probe_status VARCHAR(64) NOT NULL DEFAULT 'never',
updated_at BIGINT NOT NULL,
KEY pool_member_scores_rank_idx (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC),
KEY pool_member_scores_member_idx (pool_kind, pool_id, member_kind, member_id),
KEY pool_member_scores_probe_idx (pool_kind, pool_id, probe_status, last_probe_success_at),
KEY pool_member_scores_updated_at_idx (updated_at)
);
CREATE TABLE IF NOT EXISTS gemini_file_mappings (
id VARCHAR(64) PRIMARY KEY,
file_name VARCHAR(512) NOT NULL,

29
crates/aether-data/migrations/mysql/20260511130000_add_pool_member_scores.sql Normal file
View File

@@ -0,0 +1,29 @@
CREATE TABLE IF NOT EXISTS pool_member_scores (
id VARCHAR(192) PRIMARY KEY,
pool_kind VARCHAR(64) NOT NULL,
pool_id VARCHAR(64) NOT NULL,
member_kind VARCHAR(64) NOT NULL,
member_id VARCHAR(64) NOT NULL,
capability VARCHAR(64) NOT NULL,
scope_kind VARCHAR(64) NOT NULL,
scope_id VARCHAR(128),
score DOUBLE NOT NULL DEFAULT 0,
hard_state VARCHAR(64) NOT NULL DEFAULT 'unknown',
score_version BIGINT NOT NULL DEFAULT 1,
score_reason TEXT NOT NULL,
last_ranked_at BIGINT,
last_scheduled_at BIGINT,
last_success_at BIGINT,
last_failure_at BIGINT,
failure_count BIGINT NOT NULL DEFAULT 0,
last_probe_attempt_at BIGINT,
last_probe_success_at BIGINT,
last_probe_failure_at BIGINT,
probe_failure_count BIGINT NOT NULL DEFAULT 0,
probe_status VARCHAR(64) NOT NULL DEFAULT 'never',
updated_at BIGINT NOT NULL,
KEY pool_member_scores_rank_idx (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC),
KEY pool_member_scores_member_idx (pool_kind, pool_id, member_kind, member_id),
KEY pool_member_scores_probe_idx (pool_kind, pool_id, probe_status, last_probe_success_at),
KEY pool_member_scores_updated_at_idx (updated_at)
);

79
crates/aether-data/migrations/postgres/20260403000000_baseline.sql
View File

@@ -534,6 +534,38 @@ CREATE TABLE IF NOT EXISTS public.provider_api_keys (
--
-- Name: pool_member_scores; Type: TABLE; Schema: public; Owner: -
--
CREATE TABLE IF NOT EXISTS public.pool_member_scores (
id character varying(192) NOT NULL,
pool_kind character varying(64) NOT NULL,
pool_id character varying(64) NOT NULL,
member_kind character varying(64) NOT NULL,
member_id character varying(64) NOT NULL,
capability character varying(64) NOT NULL,
scope_kind character varying(64) NOT NULL,
scope_id character varying(128),
score double precision DEFAULT 0 NOT NULL,
hard_state character varying(64) DEFAULT 'unknown'::character varying NOT NULL,
score_version bigint DEFAULT 1 NOT NULL,
score_reason jsonb NOT NULL,
last_ranked_at bigint,
last_scheduled_at bigint,
last_success_at bigint,
last_failure_at bigint,
failure_count bigint DEFAULT 0 NOT NULL,
last_probe_attempt_at bigint,
last_probe_success_at bigint,
last_probe_failure_at bigint,
probe_failure_count bigint DEFAULT 0 NOT NULL,
probe_status character varying(64) DEFAULT 'never'::character varying NOT NULL,
updated_at bigint NOT NULL
);
--
-- Name: provider_endpoints; Type: TABLE; Schema: public; Owner: -
--
@@ -1659,6 +1691,21 @@ END $mig$;
--
-- Name: pool_member_scores pool_member_scores_pkey; Type: CONSTRAINT; Schema: public; Owner: -
--
DO $mig$ BEGIN
ALTER TABLE ONLY public.pool_member_scores
ADD CONSTRAINT pool_member_scores_pkey PRIMARY KEY (id);
EXCEPTION
WHEN duplicate_object THEN NULL;
WHEN duplicate_table THEN NULL;
WHEN invalid_table_definition THEN NULL;
END $mig$;
--
-- Name: provider_endpoints provider_endpoints_pkey; Type: CONSTRAINT; Schema: public; Owner: -
--
@@ -3401,6 +3448,38 @@ CREATE INDEX IF NOT EXISTS ix_provider_api_keys_id ON public.provider_api_keys U
--
-- Name: pool_member_scores_rank_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC);
--
-- Name: pool_member_scores_member_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, member_kind, member_id);
--
-- Name: pool_member_scores_probe_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, probe_status, last_probe_success_at);
--
-- Name: pool_member_scores_updated_at_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx ON public.pool_member_scores USING btree (updated_at);
--
-- Name: ix_provider_endpoints_id; Type: INDEX; Schema: public; Owner: -
--

40
crates/aether-data/migrations/postgres/20260511130000_add_pool_member_scores.sql Normal file
View File

@@ -0,0 +1,40 @@
CREATE TABLE IF NOT EXISTS public.pool_member_scores (
id character varying(192) PRIMARY KEY,
pool_kind character varying(64) NOT NULL,
pool_id character varying(64) NOT NULL,
member_kind character varying(64) NOT NULL,
member_id character varying(64) NOT NULL,
capability character varying(64) NOT NULL,
scope_kind character varying(64) NOT NULL,
scope_id character varying(128),
score double precision NOT NULL DEFAULT 0,
hard_state character varying(64) NOT NULL DEFAULT 'unknown',
score_version bigint NOT NULL DEFAULT 1,
score_reason jsonb NOT NULL,
last_ranked_at bigint,
last_scheduled_at bigint,
last_success_at bigint,
last_failure_at bigint,
failure_count bigint NOT NULL DEFAULT 0,
last_probe_attempt_at bigint,
last_probe_success_at bigint,
last_probe_failure_at bigint,
probe_failure_count bigint NOT NULL DEFAULT 0,
probe_status character varying(64) NOT NULL DEFAULT 'never',
updated_at bigint NOT NULL
);
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx
ON public.pool_member_scores USING btree
(pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC);
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx
ON public.pool_member_scores USING btree
(pool_kind, pool_id, member_kind, member_id);
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx
ON public.pool_member_scores USING btree
(pool_kind, pool_id, probe_status, last_probe_success_at);
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx
ON public.pool_member_scores USING btree (updated_at);

30
crates/aether-data/migrations/sqlite/20260403000000_baseline.sql
View File

@@ -301,6 +301,36 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
);
CREATE INDEX IF NOT EXISTS provider_api_keys_provider_id_idx ON provider_api_keys (provider_id);
CREATE TABLE IF NOT EXISTS pool_member_scores (
id TEXT PRIMARY KEY,
pool_kind TEXT NOT NULL,
pool_id TEXT NOT NULL,
member_kind TEXT NOT NULL,
member_id TEXT NOT NULL,
capability TEXT NOT NULL,
scope_kind TEXT NOT NULL,
scope_id TEXT,
score REAL NOT NULL DEFAULT 0,
hard_state TEXT NOT NULL DEFAULT 'unknown',
score_version INTEGER NOT NULL DEFAULT 1,
score_reason TEXT NOT NULL,
last_ranked_at INTEGER,
last_scheduled_at INTEGER,
last_success_at INTEGER,
last_failure_at INTEGER,
failure_count INTEGER NOT NULL DEFAULT 0,
last_probe_attempt_at INTEGER,
last_probe_success_at INTEGER,
last_probe_failure_at INTEGER,
probe_failure_count INTEGER NOT NULL DEFAULT 0,
probe_status TEXT NOT NULL DEFAULT 'never',
updated_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx ON pool_member_scores (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC);
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx ON pool_member_scores (pool_kind, pool_id, member_kind, member_id);
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx ON pool_member_scores (pool_kind, pool_id, probe_status, last_probe_success_at);
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx ON pool_member_scores (updated_at);
CREATE TABLE IF NOT EXISTS gemini_file_mappings (
id TEXT PRIMARY KEY,
file_name TEXT NOT NULL UNIQUE,

37
crates/aether-data/migrations/sqlite/20260511130000_add_pool_member_scores.sql Normal file
View File

@@ -0,0 +1,37 @@
CREATE TABLE IF NOT EXISTS pool_member_scores (
id TEXT PRIMARY KEY,
pool_kind TEXT NOT NULL,
pool_id TEXT NOT NULL,
member_kind TEXT NOT NULL,
member_id TEXT NOT NULL,
capability TEXT NOT NULL,
scope_kind TEXT NOT NULL,
scope_id TEXT,
score REAL NOT NULL DEFAULT 0,
hard_state TEXT NOT NULL DEFAULT 'unknown',
score_version INTEGER NOT NULL DEFAULT 1,
score_reason TEXT NOT NULL,
last_ranked_at INTEGER,
last_scheduled_at INTEGER,
last_success_at INTEGER,
last_failure_at INTEGER,
failure_count INTEGER NOT NULL DEFAULT 0,
last_probe_attempt_at INTEGER,
last_probe_success_at INTEGER,
last_probe_failure_at INTEGER,
probe_failure_count INTEGER NOT NULL DEFAULT 0,
probe_status TEXT NOT NULL DEFAULT 'never',
updated_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx
ON pool_member_scores (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC);
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx
ON pool_member_scores (pool_kind, pool_id, member_kind, member_id);
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx
ON pool_member_scores (pool_kind, pool_id, probe_status, last_probe_success_at);
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx
ON pool_member_scores (updated_at);

32
crates/aether-data/schema/bootstrap/postgres/001_types_and_tables.sql
View File

@@ -625,6 +625,38 @@ CREATE TABLE IF NOT EXISTS public.provider_api_keys (
--
-- Name: pool_member_scores; Type: TABLE; Schema: public; Owner: -
--
CREATE TABLE IF NOT EXISTS public.pool_member_scores (
id character varying(192) NOT NULL,
pool_kind character varying(64) NOT NULL,
pool_id character varying(64) NOT NULL,
member_kind character varying(64) NOT NULL,
member_id character varying(64) NOT NULL,
capability character varying(64) NOT NULL,
scope_kind character varying(64) NOT NULL,
scope_id character varying(128),
score double precision DEFAULT 0 NOT NULL,
hard_state character varying(64) DEFAULT 'unknown'::character varying NOT NULL,
score_version bigint DEFAULT 1 NOT NULL,
score_reason jsonb NOT NULL,
last_ranked_at bigint,
last_scheduled_at bigint,
last_success_at bigint,
last_failure_at bigint,
failure_count bigint DEFAULT 0 NOT NULL,
last_probe_attempt_at bigint,
last_probe_success_at bigint,
last_probe_failure_at bigint,
probe_failure_count bigint DEFAULT 0 NOT NULL,
probe_status character varying(64) DEFAULT 'never'::character varying NOT NULL,
updated_at bigint NOT NULL
);
--
-- Name: provider_endpoints; Type: TABLE; Schema: public; Owner: -
--

15
crates/aether-data/schema/bootstrap/postgres/003_constraints.sql
View File

@@ -297,6 +297,21 @@ END $mig$;
--
-- Name: pool_member_scores pool_member_scores_pkey; Type: CONSTRAINT; Schema: public; Owner: -
--
DO $mig$ BEGIN
ALTER TABLE ONLY public.pool_member_scores
ADD CONSTRAINT pool_member_scores_pkey PRIMARY KEY (id);
EXCEPTION
WHEN duplicate_object THEN NULL;
WHEN duplicate_table THEN NULL;
WHEN invalid_table_definition THEN NULL;
END $mig$;
--
-- Name: provider_endpoints provider_endpoints_pkey; Type: CONSTRAINT; Schema: public; Owner: -
--

32
crates/aether-data/schema/bootstrap/postgres/004_indexes.sql
View File

@@ -925,6 +925,38 @@ CREATE INDEX IF NOT EXISTS ix_provider_api_keys_id ON public.provider_api_keys U
--
-- Name: pool_member_scores_rank_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC);
--
-- Name: pool_member_scores_member_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, member_kind, member_id);
--
-- Name: pool_member_scores_probe_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, probe_status, last_probe_success_at);
--
-- Name: pool_member_scores_updated_at_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx ON public.pool_member_scores USING btree (updated_at);
--
-- Name: ix_provider_endpoints_id; Type: INDEX; Schema: public; Owner: -
--

30
crates/aether-data/schema/drivers/mysql/baseline/002_provider_catalog.sql
View File

@@ -127,6 +127,36 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
KEY provider_api_keys_provider_id_idx (provider_id)
);
CREATE TABLE IF NOT EXISTS pool_member_scores (
id VARCHAR(192) PRIMARY KEY,
pool_kind VARCHAR(64) NOT NULL,
pool_id VARCHAR(64) NOT NULL,
member_kind VARCHAR(64) NOT NULL,
member_id VARCHAR(64) NOT NULL,
capability VARCHAR(64) NOT NULL,
scope_kind VARCHAR(64) NOT NULL,
scope_id VARCHAR(128),
score DOUBLE NOT NULL DEFAULT 0,
hard_state VARCHAR(64) NOT NULL DEFAULT 'unknown',
score_version BIGINT NOT NULL DEFAULT 1,
score_reason TEXT NOT NULL,
last_ranked_at BIGINT,
last_scheduled_at BIGINT,
last_success_at BIGINT,
last_failure_at BIGINT,
failure_count BIGINT NOT NULL DEFAULT 0,
last_probe_attempt_at BIGINT,
last_probe_success_at BIGINT,
last_probe_failure_at BIGINT,
probe_failure_count BIGINT NOT NULL DEFAULT 0,
probe_status VARCHAR(64) NOT NULL DEFAULT 'never',
updated_at BIGINT NOT NULL,
KEY pool_member_scores_rank_idx (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC),
KEY pool_member_scores_member_idx (pool_kind, pool_id, member_kind, member_id),
KEY pool_member_scores_probe_idx (pool_kind, pool_id, probe_status, last_probe_success_at),
KEY pool_member_scores_updated_at_idx (updated_at)
);
CREATE TABLE IF NOT EXISTS gemini_file_mappings (
id VARCHAR(64) PRIMARY KEY,
file_name VARCHAR(512) NOT NULL,

32
crates/aether-data/schema/drivers/postgres/baseline/001_types_and_tables.sql
View File

@@ -534,6 +534,38 @@ CREATE TABLE IF NOT EXISTS public.provider_api_keys (
--
-- Name: pool_member_scores; Type: TABLE; Schema: public; Owner: -
--
CREATE TABLE IF NOT EXISTS public.pool_member_scores (
id character varying(192) NOT NULL,
pool_kind character varying(64) NOT NULL,
pool_id character varying(64) NOT NULL,
member_kind character varying(64) NOT NULL,
member_id character varying(64) NOT NULL,
capability character varying(64) NOT NULL,
scope_kind character varying(64) NOT NULL,
scope_id character varying(128),
score double precision DEFAULT 0 NOT NULL,
hard_state character varying(64) DEFAULT 'unknown'::character varying NOT NULL,
score_version bigint DEFAULT 1 NOT NULL,
score_reason jsonb NOT NULL,
last_ranked_at bigint,
last_scheduled_at bigint,
last_success_at bigint,
last_failure_at bigint,
failure_count bigint DEFAULT 0 NOT NULL,
last_probe_attempt_at bigint,
last_probe_success_at bigint,
last_probe_failure_at bigint,
probe_failure_count bigint DEFAULT 0 NOT NULL,
probe_status character varying(64) DEFAULT 'never'::character varying NOT NULL,
updated_at bigint NOT NULL
);
--
-- Name: provider_endpoints; Type: TABLE; Schema: public; Owner: -
--

15
crates/aether-data/schema/drivers/postgres/baseline/003_constraints.sql
View File

@@ -252,6 +252,21 @@ END $mig$;
--
-- Name: pool_member_scores pool_member_scores_pkey; Type: CONSTRAINT; Schema: public; Owner: -
--
DO $mig$ BEGIN
ALTER TABLE ONLY public.pool_member_scores
ADD CONSTRAINT pool_member_scores_pkey PRIMARY KEY (id);
EXCEPTION
WHEN duplicate_object THEN NULL;
WHEN duplicate_table THEN NULL;
WHEN invalid_table_definition THEN NULL;
END $mig$;
--
-- Name: provider_endpoints provider_endpoints_pkey; Type: CONSTRAINT; Schema: public; Owner: -
--

32
crates/aether-data/schema/drivers/postgres/baseline/004_indexes.sql
View File

@@ -869,6 +869,38 @@ CREATE INDEX IF NOT EXISTS ix_provider_api_keys_id ON public.provider_api_keys U
--
-- Name: pool_member_scores_rank_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC);
--
-- Name: pool_member_scores_member_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, member_kind, member_id);
--
-- Name: pool_member_scores_probe_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, probe_status, last_probe_success_at);
--
-- Name: pool_member_scores_updated_at_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx ON public.pool_member_scores USING btree (updated_at);
--
-- Name: ix_provider_endpoints_id; Type: INDEX; Schema: public; Owner: -
--

30
crates/aether-data/schema/drivers/sqlite/baseline/002_provider_catalog.sql
View File

@@ -134,6 +134,36 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
);
CREATE INDEX IF NOT EXISTS provider_api_keys_provider_id_idx ON provider_api_keys (provider_id);
CREATE TABLE IF NOT EXISTS pool_member_scores (
id TEXT PRIMARY KEY,
pool_kind TEXT NOT NULL,
pool_id TEXT NOT NULL,
member_kind TEXT NOT NULL,
member_id TEXT NOT NULL,
capability TEXT NOT NULL,
scope_kind TEXT NOT NULL,
scope_id TEXT,
score REAL NOT NULL DEFAULT 0,
hard_state TEXT NOT NULL DEFAULT 'unknown',
score_version INTEGER NOT NULL DEFAULT 1,
score_reason TEXT NOT NULL,
last_ranked_at INTEGER,
last_scheduled_at INTEGER,
last_success_at INTEGER,
last_failure_at INTEGER,
failure_count INTEGER NOT NULL DEFAULT 0,
last_probe_attempt_at INTEGER,
last_probe_success_at INTEGER,
last_probe_failure_at INTEGER,
probe_failure_count INTEGER NOT NULL DEFAULT 0,
probe_status TEXT NOT NULL DEFAULT 'never',
updated_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx ON pool_member_scores (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC);
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx ON pool_member_scores (pool_kind, pool_id, member_kind, member_id);
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx ON pool_member_scores (pool_kind, pool_id, probe_status, last_probe_success_at);
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx ON pool_member_scores (updated_at);
CREATE TABLE IF NOT EXISTS gemini_file_mappings (
id TEXT PRIMARY KEY,
file_name TEXT NOT NULL UNIQUE,

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

@@ -130,6 +130,37 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
KEY provider_api_keys_provider_id_idx (`provider_id`)
);
CREATE TABLE IF NOT EXISTS pool_member_scores (
`id` VARCHAR(192) NOT NULL,
`pool_kind` VARCHAR(64) NOT NULL,
`pool_id` VARCHAR(64) NOT NULL,
`member_kind` VARCHAR(64) NOT NULL,
`member_id` VARCHAR(64) NOT NULL,
`capability` VARCHAR(64) NOT NULL,
`scope_kind` VARCHAR(64) NOT NULL,
`scope_id` VARCHAR(128),
`score` DOUBLE NOT NULL DEFAULT 0,
`hard_state` VARCHAR(64) NOT NULL DEFAULT 'unknown',
`score_version` BIGINT NOT NULL DEFAULT 1,
`score_reason` JSON NOT NULL,
`last_ranked_at` BIGINT,
`last_scheduled_at` BIGINT,
`last_success_at` BIGINT,
`last_failure_at` BIGINT,
`failure_count` BIGINT NOT NULL DEFAULT 0,
`last_probe_attempt_at` BIGINT,
`last_probe_success_at` BIGINT,
`last_probe_failure_at` BIGINT,
`probe_failure_count` BIGINT NOT NULL DEFAULT 0,
`probe_status` VARCHAR(64) NOT NULL DEFAULT 'never',
`updated_at` BIGINT NOT NULL,
PRIMARY KEY (`id`),
KEY pool_member_scores_rank_idx (`pool_kind`, `pool_id`, `capability`, `scope_kind`, `scope_id`, `hard_state`, `score`),
KEY pool_member_scores_member_idx (`pool_kind`, `pool_id`, `member_kind`, `member_id`),
KEY pool_member_scores_probe_idx (`pool_kind`, `pool_id`, `probe_status`, `last_probe_success_at`),
KEY pool_member_scores_updated_at_idx (`updated_at`)
);
CREATE TABLE IF NOT EXISTS api_key_provider_mappings (
`id` VARCHAR(64) NOT NULL,
`api_key_id` VARCHAR(64) NOT NULL,

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

@@ -134,6 +134,38 @@ CREATE TABLE IF NOT EXISTS public.provider_api_keys (
ALTER TABLE ONLY public.provider_api_keys ADD CONSTRAINT provider_api_keys_pkey PRIMARY KEY (id);
CREATE INDEX IF NOT EXISTS provider_api_keys_provider_id_idx ON public.provider_api_keys USING btree (provider_id);
CREATE TABLE IF NOT EXISTS public.pool_member_scores (
id character varying(192) NOT NULL,
pool_kind character varying(64) NOT NULL,
pool_id character varying(64) NOT NULL,
member_kind character varying(64) NOT NULL,
member_id character varying(64) NOT NULL,
capability character varying(64) NOT NULL,
scope_kind character varying(64) NOT NULL,
scope_id character varying(128),
score double precision DEFAULT 0 NOT NULL,
hard_state character varying(64) DEFAULT 'unknown' NOT NULL,
score_version bigint DEFAULT 1 NOT NULL,
score_reason jsonb NOT NULL,
last_ranked_at bigint,
last_scheduled_at bigint,
last_success_at bigint,
last_failure_at bigint,
failure_count bigint DEFAULT 0 NOT NULL,
last_probe_attempt_at bigint,
last_probe_success_at bigint,
last_probe_failure_at bigint,
probe_failure_count bigint DEFAULT 0 NOT NULL,
probe_status character varying(64) DEFAULT 'never' NOT NULL,
updated_at bigint NOT NULL
);
ALTER TABLE ONLY public.pool_member_scores ADD CONSTRAINT pool_member_scores_pkey PRIMARY KEY (id);
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score);
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, member_kind, member_id);
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, probe_status, last_probe_success_at);
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx ON public.pool_member_scores USING btree (updated_at);
CREATE TABLE IF NOT EXISTS public.api_key_provider_mappings (
id character varying(64) NOT NULL,
api_key_id character varying(64) NOT NULL,

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

@@ -126,6 +126,36 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
);
CREATE INDEX IF NOT EXISTS provider_api_keys_provider_id_idx ON provider_api_keys (provider_id);
CREATE TABLE IF NOT EXISTS pool_member_scores (
id TEXT PRIMARY KEY NOT NULL,
pool_kind TEXT NOT NULL,
pool_id TEXT NOT NULL,
member_kind TEXT NOT NULL,
member_id TEXT NOT NULL,
capability TEXT NOT NULL,
scope_kind TEXT NOT NULL,
scope_id TEXT,
score REAL NOT NULL DEFAULT 0,
hard_state TEXT NOT NULL DEFAULT 'unknown',
score_version INTEGER NOT NULL DEFAULT 1,
score_reason TEXT NOT NULL,
last_ranked_at INTEGER,
last_scheduled_at INTEGER,
last_success_at INTEGER,
last_failure_at INTEGER,
failure_count INTEGER NOT NULL DEFAULT 0,
last_probe_attempt_at INTEGER,
last_probe_success_at INTEGER,
last_probe_failure_at INTEGER,
probe_failure_count INTEGER NOT NULL DEFAULT 0,
probe_status TEXT NOT NULL DEFAULT 'never',
updated_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx ON pool_member_scores (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score);
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx ON pool_member_scores (pool_kind, pool_id, member_kind, member_id);
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx ON pool_member_scores (pool_kind, pool_id, probe_status, last_probe_success_at);
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx ON pool_member_scores (updated_at);
CREATE TABLE IF NOT EXISTS api_key_provider_mappings (
id TEXT PRIMARY KEY NOT NULL,
api_key_id TEXT NOT NULL,

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

@@ -579,6 +579,143 @@ type = "unix_seconds"
name = "provider_api_keys_provider_id_idx"
columns = ["provider_id"]
[table.pool_member_scores]
domain = "provider_catalog"
order = 44
primary_key = ["id"]
[[table.pool_member_scores.columns]]
name = "id"
type = "text_id"
length = 192
[[table.pool_member_scores.columns]]
name = "pool_kind"
type = "text"
length = 64
[[table.pool_member_scores.columns]]
name = "pool_id"
type = "text_id"
length = 64
[[table.pool_member_scores.columns]]
name = "member_kind"
type = "text"
length = 64
[[table.pool_member_scores.columns]]
name = "member_id"
type = "text_id"
length = 64
[[table.pool_member_scores.columns]]
name = "capability"
type = "text"
length = 64
[[table.pool_member_scores.columns]]
name = "scope_kind"
type = "text"
length = 64
[[table.pool_member_scores.columns]]
name = "scope_id"
type = "text_id"
length = 128
nullable = true
[[table.pool_member_scores.columns]]
name = "score"
type = "float64"
default = 0
[[table.pool_member_scores.columns]]
name = "hard_state"
type = "text"
length = 64
default = "unknown"
[[table.pool_member_scores.columns]]
name = "score_version"
type = "int64"
default = 1
[[table.pool_member_scores.columns]]
name = "score_reason"
type = "json"
[[table.pool_member_scores.columns]]
name = "last_ranked_at"
type = "unix_seconds"
nullable = true
[[table.pool_member_scores.columns]]
name = "last_scheduled_at"
type = "unix_seconds"
nullable = true
[[table.pool_member_scores.columns]]
name = "last_success_at"
type = "unix_seconds"
nullable = true
[[table.pool_member_scores.columns]]
name = "last_failure_at"
type = "unix_seconds"
nullable = true
[[table.pool_member_scores.columns]]
name = "failure_count"
type = "int64"
default = 0
[[table.pool_member_scores.columns]]
name = "last_probe_attempt_at"
type = "unix_seconds"
nullable = true
[[table.pool_member_scores.columns]]
name = "last_probe_success_at"
type = "unix_seconds"
nullable = true
[[table.pool_member_scores.columns]]
name = "last_probe_failure_at"
type = "unix_seconds"
nullable = true
[[table.pool_member_scores.columns]]
name = "probe_failure_count"
type = "int64"
default = 0
[[table.pool_member_scores.columns]]
name = "probe_status"
type = "text"
length = 64
default = "never"
[[table.pool_member_scores.columns]]
name = "updated_at"
type = "unix_seconds"
[[table.pool_member_scores.indexes]]
name = "pool_member_scores_rank_idx"
columns = ["pool_kind", "pool_id", "capability", "scope_kind", "scope_id", "hard_state", "score"]
[[table.pool_member_scores.indexes]]
name = "pool_member_scores_member_idx"
columns = ["pool_kind", "pool_id", "member_kind", "member_id"]
[[table.pool_member_scores.indexes]]
name = "pool_member_scores_probe_idx"
columns = ["pool_kind", "pool_id", "probe_status", "last_probe_success_at"]
[[table.pool_member_scores.indexes]]
name = "pool_member_scores_updated_at_idx"
columns = ["updated_at"]
[table.api_key_provider_mappings]
domain = "provider_catalog"
order = 45

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

@@ -37,6 +37,9 @@ use crate::repository::management_tokens::{
use crate::repository::oauth_providers::{
MysqlOAuthProviderRepository, OAuthProviderReadRepository, OAuthProviderWriteRepository,
};
use crate::repository::pool_scores::{
MysqlPoolMemberScoreRepository, PoolMemberScoreWriteRepository, PoolScoreReadRepository,
};
use crate::repository::provider_catalog::{
MysqlProviderCatalogReadRepository, ProviderCatalogReadRepository,
ProviderCatalogWriteRepository,
@@ -184,6 +187,14 @@ impl MysqlBackend {
Arc::new(MysqlProviderCatalogReadRepository::new(self.pool_clone()))
}
pub fn pool_score_read_repository(&self) -> Arc<dyn PoolScoreReadRepository> {
Arc::new(MysqlPoolMemberScoreRepository::new(self.pool_clone()))
}
pub fn pool_score_write_repository(&self) -> Arc<dyn PoolMemberScoreWriteRepository> {
Arc::new(MysqlPoolMemberScoreRepository::new(self.pool_clone()))
}
pub fn proxy_node_read_repository(&self) -> Arc<dyn ProxyNodeReadRepository> {
Arc::new(MysqlProxyNodeReadRepository::new(self.pool_clone()))
}

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

@@ -39,6 +39,9 @@ use crate::repository::management_tokens::{
use crate::repository::oauth_providers::{
OAuthProviderReadRepository, OAuthProviderWriteRepository, SqlxOAuthProviderRepository,
};
use crate::repository::pool_scores::{
PoolMemberScoreWriteRepository, PoolScoreReadRepository, PostgresPoolMemberScoreRepository,
};
use crate::repository::provider_catalog::{
ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
SqlxProviderCatalogReadRepository,
@@ -195,6 +198,14 @@ impl PostgresBackend {
Arc::new(SqlxProviderCatalogReadRepository::new(self.pool_clone()))
}
pub fn pool_score_read_repository(&self) -> Arc<dyn PoolScoreReadRepository> {
Arc::new(PostgresPoolMemberScoreRepository::new(self.pool_clone()))
}
pub fn pool_score_write_repository(&self) -> Arc<dyn PoolMemberScoreWriteRepository> {
Arc::new(PostgresPoolMemberScoreRepository::new(self.pool_clone()))
}
pub fn provider_quota_read_repository(&self) -> Arc<dyn ProviderQuotaReadRepository> {
Arc::new(SqlxProviderQuotaRepository::new(self.pool_clone()))
}

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

@@ -14,6 +14,7 @@ use crate::repository::gemini_file_mappings::GeminiFileMappingReadRepository;
use crate::repository::global_models::GlobalModelReadRepository;
use crate::repository::management_tokens::ManagementTokenReadRepository;
use crate::repository::oauth_providers::OAuthProviderReadRepository;
use crate::repository::pool_scores::PoolScoreReadRepository;
use crate::repository::provider_catalog::ProviderCatalogReadRepository;
use crate::repository::proxy_nodes::ProxyNodeReadRepository;
use crate::repository::quota::ProviderQuotaReadRepository;
@@ -34,6 +35,7 @@ pub struct DataReadRepositories {
global_models: Option<Arc<dyn GlobalModelReadRepository>>,
management_tokens: Option<Arc<dyn ManagementTokenReadRepository>>,
oauth_providers: Option<Arc<dyn OAuthProviderReadRepository>>,
pool_scores: Option<Arc<dyn PoolScoreReadRepository>>,
proxy_nodes: Option<Arc<dyn ProxyNodeReadRepository>>,
minimal_candidate_selection: Option<Arc<dyn MinimalCandidateSelectionReadRepository>>,
request_candidates: Option<Arc<dyn RequestCandidateReadRepository>>,
@@ -61,6 +63,7 @@ impl fmt::Debug for DataReadRepositories {
.field("has_global_models", &self.global_models.is_some())
.field("has_management_tokens", &self.management_tokens.is_some())
.field("has_oauth_providers", &self.oauth_providers.is_some())
.field("has_pool_scores", &self.pool_scores.is_some())
.field("has_proxy_nodes", &self.proxy_nodes.is_some())
.field(
"has_minimal_candidate_selection",
@@ -124,6 +127,10 @@ impl DataReadRepositories {
.map(PostgresBackend::oauth_provider_read_repository)
.or_else(|| mysql.map(MysqlBackend::oauth_provider_read_repository))
.or_else(|| sqlite.map(SqliteBackend::oauth_provider_read_repository)),
pool_scores: postgres
.map(PostgresBackend::pool_score_read_repository)
.or_else(|| mysql.map(MysqlBackend::pool_score_read_repository))
.or_else(|| sqlite.map(SqliteBackend::pool_score_read_repository)),
proxy_nodes: postgres
.map(PostgresBackend::proxy_node_read_repository)
.or_else(|| mysql.map(MysqlBackend::proxy_node_read_repository))
@@ -208,6 +215,10 @@ impl DataReadRepositories {
self.oauth_providers.clone()
}
pub fn pool_scores(&self) -> Option<Arc<dyn PoolScoreReadRepository>> {
self.pool_scores.clone()
}
pub fn proxy_nodes(&self) -> Option<Arc<dyn ProxyNodeReadRepository>> {
self.proxy_nodes.clone()
}
@@ -257,6 +268,7 @@ impl DataReadRepositories {
|| self.global_models.is_some()
|| self.management_tokens.is_some()
|| self.oauth_providers.is_some()
|| self.pool_scores.is_some()
|| self.proxy_nodes.is_some()
|| self.minimal_candidate_selection.is_some()
|| self.request_candidates.is_some()

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

@@ -37,6 +37,9 @@ use crate::repository::management_tokens::{
use crate::repository::oauth_providers::{
OAuthProviderReadRepository, OAuthProviderWriteRepository, SqliteOAuthProviderRepository,
};
use crate::repository::pool_scores::{
PoolMemberScoreWriteRepository, PoolScoreReadRepository, SqlitePoolMemberScoreRepository,
};
use crate::repository::provider_catalog::{
ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
SqliteProviderCatalogReadRepository,
@@ -197,6 +200,14 @@ impl SqliteBackend {
Arc::new(SqliteProviderCatalogReadRepository::new(self.pool_clone()))
}
pub fn pool_score_read_repository(&self) -> Arc<dyn PoolScoreReadRepository> {
Arc::new(SqlitePoolMemberScoreRepository::new(self.pool_clone()))
}
pub fn pool_score_write_repository(&self) -> Arc<dyn PoolMemberScoreWriteRepository> {
Arc::new(SqlitePoolMemberScoreRepository::new(self.pool_clone()))
}
pub fn proxy_node_read_repository(&self) -> Arc<dyn ProxyNodeReadRepository> {
Arc::new(SqliteProxyNodeReadRepository::new(self.pool_clone()))
}

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

@@ -11,6 +11,7 @@ use crate::repository::gemini_file_mappings::GeminiFileMappingWriteRepository;
use crate::repository::global_models::GlobalModelWriteRepository;
use crate::repository::management_tokens::ManagementTokenWriteRepository;
use crate::repository::oauth_providers::OAuthProviderWriteRepository;
use crate::repository::pool_scores::PoolMemberScoreWriteRepository;
use crate::repository::provider_catalog::ProviderCatalogWriteRepository;
use crate::repository::proxy_nodes::ProxyNodeWriteRepository;
use crate::repository::quota::ProviderQuotaWriteRepository;
@@ -30,6 +31,7 @@ pub struct DataWriteRepositories {
global_models: Option<Arc<dyn GlobalModelWriteRepository>>,
management_tokens: Option<Arc<dyn ManagementTokenWriteRepository>>,
oauth_providers: Option<Arc<dyn OAuthProviderWriteRepository>>,
pool_scores: Option<Arc<dyn PoolMemberScoreWriteRepository>>,
proxy_nodes: Option<Arc<dyn ProxyNodeWriteRepository>>,
provider_catalog: Option<Arc<dyn ProviderCatalogWriteRepository>>,
provider_quotas: Option<Arc<dyn ProviderQuotaWriteRepository>>,
@@ -54,6 +56,7 @@ impl fmt::Debug for DataWriteRepositories {
.field("has_global_models", &self.global_models.is_some())
.field("has_management_tokens", &self.management_tokens.is_some())
.field("has_oauth_providers", &self.oauth_providers.is_some())
.field("has_pool_scores", &self.pool_scores.is_some())
.field("has_proxy_nodes", &self.proxy_nodes.is_some())
.field("has_provider_catalog", &self.provider_catalog.is_some())
.field("has_provider_quotas", &self.provider_quotas.is_some())
@@ -108,6 +111,10 @@ impl DataWriteRepositories {
.map(PostgresBackend::oauth_provider_write_repository)
.or_else(|| mysql.map(MysqlBackend::oauth_provider_write_repository))
.or_else(|| sqlite.map(SqliteBackend::oauth_provider_write_repository)),
pool_scores: postgres
.map(PostgresBackend::pool_score_write_repository)
.or_else(|| mysql.map(MysqlBackend::pool_score_write_repository))
.or_else(|| sqlite.map(SqliteBackend::pool_score_write_repository)),
proxy_nodes: postgres
.map(PostgresBackend::proxy_node_write_repository)
.or_else(|| mysql.map(MysqlBackend::proxy_node_write_repository))
@@ -184,6 +191,10 @@ impl DataWriteRepositories {
self.oauth_providers.clone()
}
pub fn pool_scores(&self) -> Option<Arc<dyn PoolMemberScoreWriteRepository>> {
self.pool_scores.clone()
}
pub fn proxy_nodes(&self) -> Option<Arc<dyn ProxyNodeWriteRepository>> {
self.proxy_nodes.clone()
}
@@ -218,6 +229,7 @@ impl DataWriteRepositories {
|| self.global_models.is_some()
|| self.management_tokens.is_some()
|| self.oauth_providers.is_some()
|| self.pool_scores.is_some()
|| self.proxy_nodes.is_some()
|| self.provider_catalog.is_some()
|| self.provider_quotas.is_some()

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

@@ -300,6 +300,7 @@ fn empty_database_snapshot_covers_current_cutoff_versions() {
20260510120000,
20260511000000,
20260511120000,
20260511130000,
20260512090000,
20260512110000,
]
@@ -579,8 +580,9 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
20260509120000,
20260510120000,
20260511120000,
20260511130000,
20260512090000,
20260512110000
20260512110000,
]
);
assert_eq!(
@@ -593,8 +595,9 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
20260509120000,
20260510120000,
20260511120000,
20260511130000,
20260512090000,
20260512110000
20260512110000,
]
);
}
@@ -1106,6 +1109,7 @@ fn pending_migrations_from_applied_skips_versions_already_applied() {
20260510120000,
20260511000000,
20260511120000,
20260511130000,
20260512090000,
20260512110000,
]

149
crates/aether-data/src/repository/candidate_selection/memory.rs
View File

@@ -4,8 +4,8 @@ use async_trait::async_trait;
use super::{
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsQuery,
StoredRequestedModelCandidateRowsQuery,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsByKeyIdsQuery,
StoredPoolKeyCandidateRowsQuery, StoredRequestedModelCandidateRowsQuery,
};
use crate::DataLayerError;
@@ -138,6 +138,39 @@ impl MinimalCandidateSelectionReadRepository for InMemoryMinimalCandidateSelecti
.take(query.limit as usize)
.collect())
}
async fn list_pool_key_rows_for_group_key_ids(
&self,
query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
if query.key_ids.is_empty() {
return Ok(Vec::new());
}
let key_order = query
.key_ids
.iter()
.enumerate()
.map(|(index, key_id)| (key_id.as_str(), index))
.collect::<std::collections::BTreeMap<_, _>>();
let mut rows = self
.list_for_exact_api_format(&query.api_format)
.await?
.into_iter()
.filter(|row| {
row.provider_id == query.provider_id
&& row.endpoint_id == query.endpoint_id
&& row.model_id == query.model_id
&& key_order.contains_key(row.key_id.as_str())
})
.collect::<Vec<_>>();
rows.sort_by(|left, right| {
key_order
.get(left.key_id.as_str())
.cmp(&key_order.get(right.key_id.as_str()))
.then(left.key_id.cmp(&right.key_id))
});
Ok(rows)
}
}
fn sort_pool_key_rows(
@@ -185,8 +218,10 @@ fn row_matches_requested_model(
requested_model_name: &str,
api_format: &str,
) -> bool {
row.global_model_name == requested_model_name
|| row.model_provider_model_name == requested_model_name
(row_has_available_provider_model(row, api_format)
&& row.global_model_name == requested_model_name)
|| (row_default_provider_model_name_available(row, api_format)
&& row.model_provider_model_name == requested_model_name)
|| row
.model_provider_model_mappings
.as_ref()
@@ -205,6 +240,60 @@ fn row_matches_requested_model(
})
}
fn row_has_available_provider_model(
row: &StoredMinimalCandidateSelectionRow,
api_format: &str,
) -> bool {
row_mapping_matches_scope(row, api_format)
|| row_default_provider_model_name_available(row, api_format)
}
fn row_default_provider_model_name_available(
row: &StoredMinimalCandidateSelectionRow,
api_format: &str,
) -> bool {
let Some(mappings) = row.model_provider_model_mappings.as_ref() else {
return true;
};
let mut has_explicit_default_mapping = false;
for mapping in mappings {
if mapping.name != row.model_provider_model_name {
continue;
}
has_explicit_default_mapping = true;
if mapping_scope_matches(mapping, row, api_format) {
return true;
}
}
!has_explicit_default_mapping
}
fn row_mapping_matches_scope(row: &StoredMinimalCandidateSelectionRow, api_format: &str) -> bool {
row.model_provider_model_mappings
.as_ref()
.is_some_and(|mappings| {
mappings
.iter()
.any(|mapping| mapping_scope_matches(mapping, row, api_format))
})
}
fn mapping_scope_matches(
mapping: &super::StoredProviderModelMapping,
row: &StoredMinimalCandidateSelectionRow,
api_format: &str,
) -> bool {
mapping.api_formats.as_ref().is_none_or(|formats| {
formats
.iter()
.any(|value| api_format_matches(value, api_format))
}) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| {
endpoint_ids
.iter()
.any(|endpoint_id| endpoint_id == &row.endpoint_id)
})
}
fn key_auth_channel_matches(row: &StoredMinimalCandidateSelectionRow, api_format: &str) -> bool {
let provider_type = row.provider_type.trim().to_ascii_lowercase();
let auth_type = row.key_auth_type.trim().to_ascii_lowercase();
@@ -244,7 +333,7 @@ mod tests {
use super::InMemoryMinimalCandidateSelectionReadRepository;
use crate::repository::candidate_selection::{
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsQuery,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredRequestedModelCandidateRowsQuery,
};
@@ -310,14 +399,12 @@ mod tests {
async fn filters_by_exact_api_format_and_requested_model_aliases() {
let mut mapped = sample_row("provider-1", "openai:chat", "gpt-4.1", 10);
mapped.model_provider_model_name = "provider-gpt-4.1".to_string();
mapped.model_provider_model_mappings = Some(vec![
crate::repository::candidate_selection::StoredProviderModelMapping {
name: "alias-gpt-4.1".to_string(),
priority: 0,
api_formats: Some(vec!["openai:chat".to_string()]),
endpoint_ids: None,
},
]);
mapped.model_provider_model_mappings = Some(vec![StoredProviderModelMapping {
name: "alias-gpt-4.1".to_string(),
priority: 0,
api_formats: Some(vec!["openai:chat".to_string()]),
endpoint_ids: None,
}]);
let repository = InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
mapped,
sample_row("provider-2", "openai:chat", "gpt-4.1-mini", 20),
@@ -332,6 +419,42 @@ mod tests {
assert_eq!(rows[0].provider_id, "provider-1");
}
#[tokio::test]
async fn requested_model_filter_respects_endpoint_scoped_default_mapping() {
let mut selected = sample_row("provider-1", "openai:chat", "deepseek-v4-pro", 10);
selected.endpoint_id = "endpoint-openai".to_string();
selected.model_provider_model_name = "deepseek-v4-pro".to_string();
selected.model_provider_model_mappings = Some(vec![StoredProviderModelMapping {
name: "deepseek-v4-pro".to_string(),
priority: 1,
api_formats: None,
endpoint_ids: Some(vec!["endpoint-openai".to_string()]),
}]);
let mut scoped_out = selected.clone();
scoped_out.provider_id = "provider-2".to_string();
scoped_out.endpoint_id = "endpoint-claude".to_string();
scoped_out.endpoint_api_format = "claude:messages".to_string();
scoped_out.key_id = "key-provider-2".to_string();
scoped_out.key_api_formats = Some(vec!["claude:messages".to_string()]);
let repository =
InMemoryMinimalCandidateSelectionReadRepository::seed(vec![scoped_out, selected]);
let rows = repository
.list_for_exact_api_format_and_requested_model("claude:messages", "deepseek-v4-pro")
.await
.expect("list should succeed");
assert!(rows.is_empty());
let rows = repository
.list_for_exact_api_format_and_requested_model("openai:chat", "deepseek-v4-pro")
.await
.expect("list should succeed");
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].endpoint_id, "endpoint-openai");
}
#[tokio::test]
async fn requested_model_page_returns_requested_slice_only() {
let mut rows = Vec::new();

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

@@ -7,8 +7,8 @@ mod sqlite;
pub(crate) use aether_data_contracts::repository::candidate_selection::{
MinimalCandidateSelectionReadRepository, MinimalCandidateSelectionRepository,
StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateOrder,
StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredRequestedModelCandidateRowsQuery,
StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery,
StoredProviderModelMapping, StoredRequestedModelCandidateRowsQuery,
};
pub use memory::InMemoryMinimalCandidateSelectionReadRepository;
pub use mysql::MysqlMinimalCandidateSelectionReadRepository;

108
crates/aether-data/src/repository/candidate_selection/mysql.rs
View File

@@ -5,7 +5,8 @@ use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row};
use super::{
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsByKeyIdsQuery,
StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredRequestedModelCandidateRowsQuery,
};
use crate::driver::mysql::MysqlPool;
@@ -196,6 +197,40 @@ impl MinimalCandidateSelectionReadRepository for MysqlMinimalCandidateSelectionR
.map(|item| item.row)
.collect())
}
async fn list_pool_key_rows_for_group_key_ids(
&self,
query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
if query.key_ids.is_empty() {
return Ok(Vec::new());
}
let key_order = query
.key_ids
.iter()
.enumerate()
.map(|(index, key_id)| (key_id.as_str(), index))
.collect::<BTreeMap<_, _>>();
let mut rows = self
.load_rows_for_api_format(&query.api_format)
.await?
.into_iter()
.filter(|row| {
row.row.provider_id == query.provider_id
&& row.row.endpoint_id == query.endpoint_id
&& row.row.model_id == query.model_id
&& key_order.contains_key(row.row.key_id.as_str())
})
.map(|item| item.row)
.collect::<Vec<_>>();
rows.sort_by(|left, right| {
key_order
.get(left.key_id.as_str())
.cmp(&key_order.get(right.key_id.as_str()))
.then(left.key_id.cmp(&right.key_id))
});
Ok(dedupe_candidate_selection_rows(rows))
}
}
fn select_pool_rows(rows: Vec<CandidateSelectionRow>) -> Vec<StoredMinimalCandidateSelectionRow> {
@@ -313,26 +348,75 @@ fn row_matches_requested_model(
requested_model_name: &str,
api_format: &str,
) -> bool {
row.global_model_name == requested_model_name
|| row.model_provider_model_name == requested_model_name
(row_has_available_provider_model(row, api_format)
&& row.global_model_name == requested_model_name)
|| (row_default_provider_model_name_available(row, api_format)
&& row.model_provider_model_name == requested_model_name)
|| row
.model_provider_model_mappings
.as_ref()
.is_some_and(|mappings| {
mappings.iter().any(|mapping| {
mapping.api_formats.as_ref().is_none_or(|formats| {
formats
.iter()
.any(|value| api_format_matches(value, api_format))
}) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| {
endpoint_ids
.iter()
.any(|endpoint_id| endpoint_id == &row.endpoint_id)
}) && mapping.name == requested_model_name
mapping_scope_matches(mapping, row, api_format)
&& mapping.name == requested_model_name
})
})
}
fn row_has_available_provider_model(
row: &StoredMinimalCandidateSelectionRow,
api_format: &str,
) -> bool {
row_mapping_matches_scope(row, api_format)
|| row_default_provider_model_name_available(row, api_format)
}
fn row_default_provider_model_name_available(
row: &StoredMinimalCandidateSelectionRow,
api_format: &str,
) -> bool {
let Some(mappings) = row.model_provider_model_mappings.as_ref() else {
return true;
};
let mut has_explicit_default_mapping = false;
for mapping in mappings {
if mapping.name != row.model_provider_model_name {
continue;
}
has_explicit_default_mapping = true;
if mapping_scope_matches(mapping, row, api_format) {
return true;
}
}
!has_explicit_default_mapping
}
fn row_mapping_matches_scope(row: &StoredMinimalCandidateSelectionRow, api_format: &str) -> bool {
row.model_provider_model_mappings
.as_ref()
.is_some_and(|mappings| {
mappings
.iter()
.any(|mapping| mapping_scope_matches(mapping, row, api_format))
})
}
fn mapping_scope_matches(
mapping: &super::StoredProviderModelMapping,
row: &StoredMinimalCandidateSelectionRow,
api_format: &str,
) -> bool {
mapping.api_formats.as_ref().is_none_or(|formats| {
formats
.iter()
.any(|value| api_format_matches(value, api_format))
}) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| {
endpoint_ids
.iter()
.any(|endpoint_id| endpoint_id == &row.endpoint_id)
})
}
fn key_auth_channel_matches(row: &CandidateSelectionRow, api_format: &str) -> bool {
let provider_type = row.row.provider_type.trim().to_ascii_lowercase();
let auth_type = row.row.key_auth_type.trim().to_ascii_lowercase();

167
crates/aether-data/src/repository/candidate_selection/postgres.rs
View File

@@ -5,7 +5,8 @@ use std::collections::BTreeSet;
use super::{
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsByKeyIdsQuery,
StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredRequestedModelCandidateRowsQuery,
};
use crate::{error::SqlxResultExt, DataLayerError};
@@ -535,6 +536,13 @@ fn pool_key_candidate_selection_sql(order: &StoredPoolKeyCandidateOrder) -> Stri
LIST_POOL_KEYS_FOR_GROUP_SQL.replace(default_order, &replacement)
}
fn pool_key_candidate_selection_by_key_ids_sql() -> String {
let default_order =
"ORDER BY\n pak.internal_priority ASC,\n pak.id ASC\nLIMIT $7\nOFFSET $8\n";
let replacement = "AND pak.id = ANY($7::text[])\nORDER BY\n array_position($7::text[], pak.id) ASC,\n pak.id ASC\n";
LIST_POOL_KEYS_FOR_GROUP_SQL.replace(default_order, replacement)
}
#[derive(Debug, Clone)]
pub struct SqlxMinimalCandidateSelectionReadRepository {
pool: PgPool,
@@ -704,6 +712,51 @@ impl SqlxMinimalCandidateSelectionReadRepository {
}
Ok(dedupe_candidate_selection_rows(rows))
}
pub async fn list_pool_key_rows_for_group_key_ids(
&self,
query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
if query.key_ids.is_empty() {
return Ok(Vec::new());
}
let mut rows = Vec::new();
let canonical_api_format = normalize_api_format(&query.api_format);
let storage_aliases = api_format_aliases(&canonical_api_format);
let sql_match_aliases = sql_match_aliases(&storage_aliases);
let sql = pool_key_candidate_selection_by_key_ids_sql();
for api_format in storage_aliases {
rows.extend(
Self::collect_query_rows(
sqlx::query(sql.as_str())
.bind(api_format)
.bind(query.provider_id.as_str())
.bind(query.endpoint_id.as_str())
.bind(query.model_id.as_str())
.bind(sql_match_aliases.clone())
.bind(canonical_api_format.clone())
.bind(query.key_ids.clone())
.fetch(&self.pool),
map_candidate_selection_row,
)
.await?,
);
}
let key_order = query
.key_ids
.iter()
.enumerate()
.map(|(index, key_id)| (key_id.as_str(), index))
.collect::<std::collections::BTreeMap<_, _>>();
let mut rows = dedupe_candidate_selection_rows(rows);
rows.sort_by(|left, right| {
key_order
.get(left.key_id.as_str())
.cmp(&key_order.get(right.key_id.as_str()))
.then(left.key_id.cmp(&right.key_id))
});
Ok(rows)
}
}
fn requested_model_selection_sql() -> String {
@@ -711,13 +764,110 @@ fn requested_model_selection_sql() -> String {
.replace(
"AND gm.name = $2",
r#"AND (
gm.name = $2
OR m.provider_model_name = $2
(
gm.name = $2
AND (
m.provider_model_mappings IS NULL
OR jsonb_typeof(m.provider_model_mappings) <> 'array'
OR EXISTS (
SELECT 1
FROM jsonb_array_elements(
CASE
WHEN jsonb_typeof(m.provider_model_mappings) = 'array'
THEN m.provider_model_mappings
ELSE '[]'::jsonb
END
) AS mapping(value)
WHERE (
mapping.value -> 'api_formats' IS NULL
OR jsonb_typeof(mapping.value -> 'api_formats') <> 'array'
OR EXISTS (
SELECT 1
FROM jsonb_array_elements_text(mapping.value -> 'api_formats') AS fmt(value)
WHERE LOWER(BTRIM(fmt.value)) = ANY($3::text[])
)
)
AND (
mapping.value -> 'endpoint_ids' IS NULL
OR jsonb_typeof(mapping.value -> 'endpoint_ids') <> 'array'
OR EXISTS (
SELECT 1
FROM jsonb_array_elements_text(mapping.value -> 'endpoint_ids') AS endpoint(value)
WHERE endpoint.value = pe.id
)
)
)
OR NOT EXISTS (
SELECT 1
FROM jsonb_array_elements(
CASE
WHEN jsonb_typeof(m.provider_model_mappings) = 'array'
THEN m.provider_model_mappings
ELSE '[]'::jsonb
END
) AS mapping(value)
WHERE mapping.value ->> 'name' = m.provider_model_name
)
)
)
OR (
m.provider_model_name = $2
AND (
m.provider_model_mappings IS NULL
OR jsonb_typeof(m.provider_model_mappings) <> 'array'
OR NOT EXISTS (
SELECT 1
FROM jsonb_array_elements(
CASE
WHEN jsonb_typeof(m.provider_model_mappings) = 'array'
THEN m.provider_model_mappings
ELSE '[]'::jsonb
END
) AS mapping(value)
WHERE mapping.value ->> 'name' = m.provider_model_name
)
OR EXISTS (
SELECT 1
FROM jsonb_array_elements(
CASE
WHEN jsonb_typeof(m.provider_model_mappings) = 'array'
THEN m.provider_model_mappings
ELSE '[]'::jsonb
END
) AS mapping(value)
WHERE mapping.value ->> 'name' = m.provider_model_name
AND (
mapping.value -> 'api_formats' IS NULL
OR jsonb_typeof(mapping.value -> 'api_formats') <> 'array'
OR EXISTS (
SELECT 1
FROM jsonb_array_elements_text(mapping.value -> 'api_formats') AS fmt(value)
WHERE LOWER(BTRIM(fmt.value)) = ANY($3::text[])
)
)
AND (
mapping.value -> 'endpoint_ids' IS NULL
OR jsonb_typeof(mapping.value -> 'endpoint_ids') <> 'array'
OR EXISTS (
SELECT 1
FROM jsonb_array_elements_text(mapping.value -> 'endpoint_ids') AS endpoint(value)
WHERE endpoint.value = pe.id
)
)
)
)
)
OR (
jsonb_typeof(m.provider_model_mappings) = 'array'
AND EXISTS (
SELECT 1
FROM jsonb_array_elements(m.provider_model_mappings) AS mapping(value)
FROM jsonb_array_elements(
CASE
WHEN jsonb_typeof(m.provider_model_mappings) = 'array'
THEN m.provider_model_mappings
ELSE '[]'::jsonb
END
) AS mapping(value)
WHERE mapping.value ->> 'name' = $2
AND (
mapping.value -> 'api_formats' IS NULL
@@ -820,6 +970,13 @@ impl MinimalCandidateSelectionReadRepository for SqlxMinimalCandidateSelectionRe
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
Self::list_pool_key_rows_for_group(self, query).await
}
async fn list_pool_key_rows_for_group_key_ids(
&self,
query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
Self::list_pool_key_rows_for_group_key_ids(self, query).await
}
}
fn map_candidate_selection_row(
@@ -1108,7 +1265,7 @@ mod tests {
let sql = requested_model_selection_sql();
assert!(sql.contains("m.provider_model_name = $2"));
assert!(sql.contains("jsonb_array_elements(m.provider_model_mappings)"));
assert!(sql.contains("jsonb_array_elements("));
assert!(!sql.contains("json_typeof(m.provider_model_mappings)"));
assert!(!sql.contains("json_array_elements_text(gm.config -> 'model_mappings')"));
assert!(sql.contains("ORDER BY\n global_model_name ASC,"));

108
crates/aether-data/src/repository/candidate_selection/sqlite.rs
View File

@@ -5,7 +5,8 @@ use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite};
use super::{
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsByKeyIdsQuery,
StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredRequestedModelCandidateRowsQuery,
};
use crate::driver::sqlite::SqlitePool;
@@ -196,6 +197,40 @@ impl MinimalCandidateSelectionReadRepository for SqliteMinimalCandidateSelection
.map(|item| item.row)
.collect())
}
async fn list_pool_key_rows_for_group_key_ids(
&self,
query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
if query.key_ids.is_empty() {
return Ok(Vec::new());
}
let key_order = query
.key_ids
.iter()
.enumerate()
.map(|(index, key_id)| (key_id.as_str(), index))
.collect::<BTreeMap<_, _>>();
let mut rows = self
.load_rows_for_api_format(&query.api_format)
.await?
.into_iter()
.filter(|row| {
row.row.provider_id == query.provider_id
&& row.row.endpoint_id == query.endpoint_id
&& row.row.model_id == query.model_id
&& key_order.contains_key(row.row.key_id.as_str())
})
.map(|item| item.row)
.collect::<Vec<_>>();
rows.sort_by(|left, right| {
key_order
.get(left.key_id.as_str())
.cmp(&key_order.get(right.key_id.as_str()))
.then(left.key_id.cmp(&right.key_id))
});
Ok(dedupe_candidate_selection_rows(rows))
}
}
fn select_pool_rows(rows: Vec<CandidateSelectionRow>) -> Vec<StoredMinimalCandidateSelectionRow> {
@@ -313,26 +348,75 @@ fn row_matches_requested_model(
requested_model_name: &str,
api_format: &str,
) -> bool {
row.global_model_name == requested_model_name
|| row.model_provider_model_name == requested_model_name
(row_has_available_provider_model(row, api_format)
&& row.global_model_name == requested_model_name)
|| (row_default_provider_model_name_available(row, api_format)
&& row.model_provider_model_name == requested_model_name)
|| row
.model_provider_model_mappings
.as_ref()
.is_some_and(|mappings| {
mappings.iter().any(|mapping| {
mapping.api_formats.as_ref().is_none_or(|formats| {
formats
.iter()
.any(|value| api_format_matches(value, api_format))
}) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| {
endpoint_ids
.iter()
.any(|endpoint_id| endpoint_id == &row.endpoint_id)
}) && mapping.name == requested_model_name
mapping_scope_matches(mapping, row, api_format)
&& mapping.name == requested_model_name
})
})
}
fn row_has_available_provider_model(
row: &StoredMinimalCandidateSelectionRow,
api_format: &str,
) -> bool {
row_mapping_matches_scope(row, api_format)
|| row_default_provider_model_name_available(row, api_format)
}
fn row_default_provider_model_name_available(
row: &StoredMinimalCandidateSelectionRow,
api_format: &str,
) -> bool {
let Some(mappings) = row.model_provider_model_mappings.as_ref() else {
return true;
};
let mut has_explicit_default_mapping = false;
for mapping in mappings {
if mapping.name != row.model_provider_model_name {
continue;
}
has_explicit_default_mapping = true;
if mapping_scope_matches(mapping, row, api_format) {
return true;
}
}
!has_explicit_default_mapping
}
fn row_mapping_matches_scope(row: &StoredMinimalCandidateSelectionRow, api_format: &str) -> bool {
row.model_provider_model_mappings
.as_ref()
.is_some_and(|mappings| {
mappings
.iter()
.any(|mapping| mapping_scope_matches(mapping, row, api_format))
})
}
fn mapping_scope_matches(
mapping: &super::StoredProviderModelMapping,
row: &StoredMinimalCandidateSelectionRow,
api_format: &str,
) -> bool {
mapping.api_formats.as_ref().is_none_or(|formats| {
formats
.iter()
.any(|value| api_format_matches(value, api_format))
}) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| {
endpoint_ids
.iter()
.any(|endpoint_id| endpoint_id == &row.endpoint_id)
})
}
fn key_auth_channel_matches(row: &CandidateSelectionRow, api_format: &str) -> bool {
let provider_type = row.row.provider_type.trim().to_ascii_lowercase();
let auth_type = row.row.key_auth_type.trim().to_ascii_lowercase();

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

@@ -16,6 +16,7 @@ pub mod gemini_file_mappings;
pub mod global_models;
pub mod management_tokens;
pub mod oauth_providers;
pub mod pool_scores;
pub mod provider_catalog;
pub mod provider_oauth;
pub mod proxy_nodes;

513
crates/aether-data/src/repository/pool_scores/memory.rs Normal file
View File

@@ -0,0 +1,513 @@
use std::collections::BTreeMap;
use std::sync::RwLock;
use async_trait::async_trait;
use super::{
score_with_delta, GetPoolMemberScoresByIdsQuery, ListPoolMemberProbeCandidatesQuery,
ListPoolMemberScoresQuery, ListRankedPoolMembersQuery, PoolMemberHardState, PoolMemberIdentity,
PoolMemberProbeAttempt, PoolMemberProbeResult, PoolMemberProbeStatus,
PoolMemberScheduleFeedback, PoolMemberScoreWriteRepository, PoolScoreReadRepository,
PoolScoreScope, StoredPoolMemberScore, UpsertPoolMemberScore,
};
use crate::repository::pool_scores::merge_score_reason_patch;
use crate::DataLayerError;
#[derive(Debug, Default)]
pub struct InMemoryPoolMemberScoreRepository {
scores: RwLock<BTreeMap<String, StoredPoolMemberScore>>,
}
impl InMemoryPoolMemberScoreRepository {
pub fn seed<I>(scores: I) -> Self
where
I: IntoIterator<Item = StoredPoolMemberScore>,
{
Self {
scores: RwLock::new(
scores
.into_iter()
.map(|score| (score.id.clone(), score))
.collect(),
),
}
}
fn matches_identity(score: &StoredPoolMemberScore, identity: &PoolMemberIdentity) -> bool {
score.pool_kind == identity.pool_kind
&& score.pool_id == identity.pool_id
&& score.member_kind == identity.member_kind
&& score.member_id == identity.member_id
}
fn matches_scope(score: &StoredPoolMemberScore, scope: Option<&PoolScoreScope>) -> bool {
let Some(scope) = scope else {
return true;
};
score.capability == scope.capability
&& score.scope_kind == scope.scope_kind
&& score.scope_id == scope.scope_id
}
fn sort_ranked(scores: &mut [StoredPoolMemberScore]) {
scores.sort_by(|left, right| {
right
.score
.partial_cmp(&left.score)
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| {
right
.last_ranked_at
.unwrap_or(0)
.cmp(&left.last_ranked_at.unwrap_or(0))
})
.then_with(|| left.member_id.cmp(&right.member_id))
.then_with(|| left.id.cmp(&right.id))
});
}
fn sort_probe(scores: &mut [StoredPoolMemberScore]) {
scores.sort_by(|left, right| {
probe_priority(left)
.cmp(&probe_priority(right))
.then_with(|| right.probe_failure_count.cmp(&left.probe_failure_count))
.then_with(|| {
left.last_probe_success_at
.unwrap_or(0)
.cmp(&right.last_probe_success_at.unwrap_or(0))
})
.then_with(|| {
left.last_scheduled_at
.unwrap_or(0)
.cmp(&right.last_scheduled_at.unwrap_or(0))
.reverse()
})
.then_with(|| left.member_id.cmp(&right.member_id))
});
}
}
#[async_trait]
impl PoolScoreReadRepository for InMemoryPoolMemberScoreRepository {
async fn list_ranked_pool_members(
&self,
query: &ListRankedPoolMembersQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let hard_states = query
.hard_states
.iter()
.copied()
.collect::<std::collections::BTreeSet<_>>();
let probe_statuses = query.probe_statuses.as_ref().map(|items| {
items
.iter()
.copied()
.collect::<std::collections::BTreeSet<_>>()
});
let mut scores = self
.scores
.read()
.expect("pool member score repository lock")
.values()
.filter(|score| {
score.pool_kind == query.pool_kind
&& score.pool_id == query.pool_id
&& score.capability == query.capability
&& score.scope_kind == query.scope_kind
&& score.scope_id == query.scope_id
&& (hard_states.is_empty() || hard_states.contains(&score.hard_state))
&& probe_statuses
.as_ref()
.is_none_or(|statuses| statuses.contains(&score.probe_status))
})
.cloned()
.collect::<Vec<_>>();
Self::sort_ranked(&mut scores);
Ok(scores
.into_iter()
.skip(query.offset)
.take(query.limit.max(1))
.collect())
}
async fn list_pool_member_scores(
&self,
query: &ListPoolMemberScoresQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let hard_states = query
.hard_states
.iter()
.copied()
.collect::<std::collections::BTreeSet<_>>();
let probe_statuses = query.probe_statuses.as_ref().map(|items| {
items
.iter()
.copied()
.collect::<std::collections::BTreeSet<_>>()
});
let mut scores = self
.scores
.read()
.expect("pool member score repository lock")
.values()
.filter(|score| {
score.pool_kind == query.pool_kind
&& score.pool_id == query.pool_id
&& query
.capability
.as_ref()
.is_none_or(|capability| score.capability == *capability)
&& query
.scope_kind
.as_ref()
.is_none_or(|scope_kind| score.scope_kind == *scope_kind)
&& query
.scope_id
.as_ref()
.is_none_or(|scope_id| score.scope_id.as_ref() == Some(scope_id))
&& (hard_states.is_empty() || hard_states.contains(&score.hard_state))
&& probe_statuses
.as_ref()
.is_none_or(|statuses| statuses.contains(&score.probe_status))
})
.cloned()
.collect::<Vec<_>>();
Self::sort_ranked(&mut scores);
Ok(scores
.into_iter()
.skip(query.offset)
.take(query.limit.max(1))
.collect())
}
async fn list_pool_member_probe_candidates(
&self,
query: &ListPoolMemberProbeCandidatesQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut scores = self
.scores
.read()
.expect("pool member score repository lock")
.values()
.filter(|score| {
score.pool_kind == query.pool_kind
&& score.pool_id == query.pool_id
&& query
.capability
.as_ref()
.is_none_or(|capability| score.capability == *capability)
&& matches!(
score.hard_state,
PoolMemberHardState::Available
| PoolMemberHardState::Unknown
| PoolMemberHardState::Cooldown
| PoolMemberHardState::QuotaExhausted
)
&& match score.probe_status {
PoolMemberProbeStatus::Never
| PoolMemberProbeStatus::Failed
| PoolMemberProbeStatus::Stale => true,
PoolMemberProbeStatus::Ok => score
.last_probe_success_at
.is_none_or(|ts| ts <= query.stale_before_unix_secs),
PoolMemberProbeStatus::InProgress => score
.last_probe_attempt_at
.is_none_or(|ts| ts <= query.stale_before_unix_secs),
}
})
.cloned()
.collect::<Vec<_>>();
Self::sort_probe(&mut scores);
Ok(scores.into_iter().take(query.limit.max(1)).collect())
}
async fn get_pool_member_scores_by_ids(
&self,
query: &GetPoolMemberScoresByIdsQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let ids = query
.ids
.iter()
.cloned()
.collect::<std::collections::BTreeSet<_>>();
let scores = self
.scores
.read()
.expect("pool member score repository lock")
.values()
.filter(|score| ids.contains(&score.id))
.cloned()
.collect::<Vec<_>>();
Ok(scores)
}
}
#[async_trait]
impl PoolMemberScoreWriteRepository for InMemoryPoolMemberScoreRepository {
async fn upsert_pool_member_score(
&self,
score: UpsertPoolMemberScore,
) -> Result<StoredPoolMemberScore, DataLayerError> {
score.validate()?;
let stored = score.into_stored();
self.scores
.write()
.expect("pool member score repository lock")
.insert(stored.id.clone(), stored.clone());
Ok(stored)
}
async fn mark_pool_member_probe_in_progress(
&self,
attempt: PoolMemberProbeAttempt,
) -> Result<usize, DataLayerError> {
let mut updated = 0;
let mut guard = self
.scores
.write()
.expect("pool member score repository lock");
for score in guard.values_mut() {
if !Self::matches_identity(score, &attempt.identity)
|| !Self::matches_scope(score, attempt.scope.as_ref())
{
continue;
}
score.last_probe_attempt_at = Some(attempt.attempted_at);
score.probe_status = PoolMemberProbeStatus::InProgress;
score.score_reason = merge_score_reason_patch(
score.score_reason.clone(),
attempt.score_reason_patch.clone(),
);
score.updated_at = attempt.attempted_at;
updated += 1;
}
Ok(updated)
}
async fn record_pool_member_probe_result(
&self,
result: PoolMemberProbeResult,
) -> Result<usize, DataLayerError> {
let mut updated = 0;
let mut guard = self
.scores
.write()
.expect("pool member score repository lock");
for score in guard.values_mut() {
if !Self::matches_identity(score, &result.identity)
|| !Self::matches_scope(score, result.scope.as_ref())
{
continue;
}
score.last_probe_attempt_at = Some(result.attempted_at);
score.probe_status = result.probe_status;
if result.succeeded {
score.last_probe_success_at = Some(result.attempted_at);
score.probe_failure_count = 0;
} else {
score.last_probe_failure_at = Some(result.attempted_at);
score.probe_failure_count = score.probe_failure_count.saturating_add(1);
}
if let Some(hard_state) = result.hard_state {
score.hard_state = hard_state;
}
score.score_reason = merge_score_reason_patch(
score.score_reason.clone(),
result.score_reason_patch.clone(),
);
score.updated_at = result.attempted_at;
updated += 1;
}
Ok(updated)
}
async fn record_pool_member_schedule_feedback(
&self,
feedback: PoolMemberScheduleFeedback,
) -> Result<usize, DataLayerError> {
let mut updated = 0;
let mut guard = self
.scores
.write()
.expect("pool member score repository lock");
for score in guard.values_mut() {
if !Self::matches_identity(score, &feedback.identity)
|| !Self::matches_scope(score, feedback.scope.as_ref())
{
continue;
}
score.last_scheduled_at = Some(feedback.scheduled_at);
match feedback.succeeded {
Some(true) => {
score.last_success_at = Some(feedback.scheduled_at);
}
Some(false) => {
score.last_failure_at = Some(feedback.scheduled_at);
score.failure_count = score.failure_count.saturating_add(1);
}
None => {}
}
if let Some(hard_state) = feedback.hard_state {
score.hard_state = hard_state;
}
score.score = score_with_delta(score.score, feedback.score_delta);
score.score_reason = merge_score_reason_patch(
score.score_reason.clone(),
feedback.score_reason_patch.clone(),
);
score.updated_at = feedback.scheduled_at;
updated += 1;
}
Ok(updated)
}
async fn mark_pool_member_hard_state(
&self,
identity: &PoolMemberIdentity,
scope: Option<&PoolScoreScope>,
hard_state: PoolMemberHardState,
updated_at: u64,
) -> Result<usize, DataLayerError> {
let mut updated = 0;
let mut guard = self
.scores
.write()
.expect("pool member score repository lock");
for score in guard.values_mut() {
if Self::matches_identity(score, identity) && Self::matches_scope(score, scope) {
score.hard_state = hard_state;
score.updated_at = updated_at;
updated += 1;
}
}
Ok(updated)
}
async fn delete_pool_member_scores_for_member(
&self,
identity: &PoolMemberIdentity,
) -> Result<usize, DataLayerError> {
let mut guard = self
.scores
.write()
.expect("pool member score repository lock");
let before = guard.len();
guard.retain(|_, score| !Self::matches_identity(score, identity));
Ok(before.saturating_sub(guard.len()))
}
}
fn probe_priority(score: &StoredPoolMemberScore) -> u8 {
if score.last_scheduled_at.is_some() && score.probe_status != PoolMemberProbeStatus::Ok {
return 0;
}
match score.hard_state {
PoolMemberHardState::QuotaExhausted => 1,
PoolMemberHardState::Unknown => 2,
_ if score.probe_status == PoolMemberProbeStatus::Stale => 3,
_ => 4,
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::repository::pool_scores::{
POOL_KIND_PROVIDER_KEY_POOL, POOL_MEMBER_KIND_PROVIDER_API_KEY,
POOL_SCORE_CAPABILITY_API_FORMAT, POOL_SCORE_SCOPE_KIND_MODEL,
};
fn score(id: &str, member_id: &str, value: f64) -> StoredPoolMemberScore {
StoredPoolMemberScore {
id: id.to_string(),
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
pool_id: "provider-1".to_string(),
member_kind: POOL_MEMBER_KIND_PROVIDER_API_KEY.to_string(),
member_id: member_id.to_string(),
capability: POOL_SCORE_CAPABILITY_API_FORMAT.to_string(),
scope_kind: POOL_SCORE_SCOPE_KIND_MODEL.to_string(),
scope_id: Some("model-1".to_string()),
score: value,
hard_state: PoolMemberHardState::Available,
score_version: 1,
score_reason: serde_json::json!({}),
last_ranked_at: Some(1),
last_scheduled_at: None,
last_success_at: None,
last_failure_at: None,
failure_count: 0,
last_probe_attempt_at: None,
last_probe_success_at: None,
last_probe_failure_at: None,
probe_failure_count: 0,
probe_status: PoolMemberProbeStatus::Never,
updated_at: 1,
}
}
#[tokio::test]
async fn lists_ranked_members_by_score() {
let repository = InMemoryPoolMemberScoreRepository::seed(vec![
score("score-1", "key-1", 0.2),
score("score-2", "key-2", 0.9),
]);
let rows = repository
.list_ranked_pool_members(&ListRankedPoolMembersQuery {
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
pool_id: "provider-1".to_string(),
capability: POOL_SCORE_CAPABILITY_API_FORMAT.to_string(),
scope_kind: POOL_SCORE_SCOPE_KIND_MODEL.to_string(),
scope_id: Some("model-1".to_string()),
hard_states: vec![PoolMemberHardState::Available],
probe_statuses: None,
offset: 0,
limit: 10,
})
.await
.expect("list should succeed");
assert_eq!(
rows.into_iter()
.map(|row| row.member_id)
.collect::<Vec<_>>(),
vec!["key-2".to_string(), "key-1".to_string()]
);
}
#[tokio::test]
async fn marks_probe_in_progress_without_incrementing_failure_count() {
let repository =
InMemoryPoolMemberScoreRepository::seed(vec![score("score-1", "key-1", 0.2)]);
let updated = repository
.mark_pool_member_probe_in_progress(PoolMemberProbeAttempt {
identity: PoolMemberIdentity::provider_api_key("provider-1", "key-1"),
scope: None,
attempted_at: 100,
score_reason_patch: Some(serde_json::json!({ "last_probe": "in_progress" })),
})
.await
.expect("mark should succeed");
assert_eq!(updated, 1);
let rows = repository
.list_pool_member_scores(&ListPoolMemberScoresQuery {
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
pool_id: "provider-1".to_string(),
capability: None,
scope_kind: None,
scope_id: None,
hard_states: Vec::new(),
probe_statuses: Some(vec![PoolMemberProbeStatus::InProgress]),
offset: 0,
limit: 10,
})
.await
.expect("list should succeed");
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].last_probe_attempt_at, Some(100));
assert_eq!(rows[0].probe_failure_count, 0);
}
}

54
crates/aether-data/src/repository/pool_scores/mod.rs Normal file
View File

@@ -0,0 +1,54 @@
pub use aether_data_contracts::repository::pool_scores::*;
mod memory;
mod mysql;
mod postgres;
mod sqlite;
pub use memory::InMemoryPoolMemberScoreRepository;
pub use mysql::MysqlPoolMemberScoreRepository;
pub use postgres::PostgresPoolMemberScoreRepository;
pub use sqlite::SqlitePoolMemberScoreRepository;
fn merge_score_reason_patch(
mut current: serde_json::Value,
patch: Option<serde_json::Value>,
) -> serde_json::Value {
let Some(patch) = patch else {
return current;
};
match (current.as_object_mut(), patch) {
(Some(current), serde_json::Value::Object(patch)) => {
for (key, value) in patch {
current.insert(key, value);
}
serde_json::Value::Object(current.clone())
}
(_, patch) => patch,
}
}
fn score_with_delta(score: f64, delta_basis_points: Option<i32>) -> f64 {
let delta = delta_basis_points.unwrap_or_default() as f64 / 10_000.0;
(score + delta).clamp(0.0, 1.0)
}
fn i64_from_u64(value: u64, field: &str) -> Result<i64, crate::DataLayerError> {
i64::try_from(value).map_err(|_| {
crate::DataLayerError::InvalidInput(format!("{field} exceeds signed 64-bit range"))
})
}
fn i64_opt_from_u64(value: Option<u64>, field: &str) -> Result<Option<i64>, crate::DataLayerError> {
value.map(|value| i64_from_u64(value, field)).transpose()
}
fn u64_from_i64(value: i64, field: &str) -> Result<u64, crate::DataLayerError> {
u64::try_from(value).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!("{field} is negative: {value}"))
})
}
fn u64_opt_from_i64(value: Option<i64>, field: &str) -> Result<Option<u64>, crate::DataLayerError> {
value.map(|value| u64_from_i64(value, field)).transpose()
}

586
crates/aether-data/src/repository/pool_scores/mysql.rs Normal file
View File

@@ -0,0 +1,586 @@
use async_trait::async_trait;
use sqlx::{mysql::MySqlRow, MySql, MySqlPool, QueryBuilder, Row};
use super::{
i64_from_u64, i64_opt_from_u64, score_with_delta, u64_from_i64, u64_opt_from_i64,
GetPoolMemberScoresByIdsQuery, ListPoolMemberProbeCandidatesQuery, ListPoolMemberScoresQuery,
ListRankedPoolMembersQuery, PoolMemberHardState, PoolMemberIdentity, PoolMemberProbeAttempt,
PoolMemberProbeResult, PoolMemberProbeStatus, PoolMemberScheduleFeedback,
PoolMemberScoreWriteRepository, PoolScoreReadRepository, PoolScoreScope, StoredPoolMemberScore,
UpsertPoolMemberScore,
};
use crate::error::SqlResultExt;
use crate::repository::pool_scores::merge_score_reason_patch;
use crate::DataLayerError;
const SCORE_COLUMNS: &str = r#"
SELECT
id,
pool_kind,
pool_id,
member_kind,
member_id,
capability,
scope_kind,
scope_id,
score,
hard_state,
score_version,
score_reason,
last_ranked_at,
last_scheduled_at,
last_success_at,
last_failure_at,
failure_count,
last_probe_attempt_at,
last_probe_success_at,
last_probe_failure_at,
probe_failure_count,
probe_status,
updated_at
FROM pool_member_scores
"#;
#[derive(Debug, Clone)]
pub struct MysqlPoolMemberScoreRepository {
pool: MySqlPool,
}
impl MysqlPoolMemberScoreRepository {
pub fn new(pool: MySqlPool) -> Self {
Self { pool }
}
async fn find_scores_by_identity(
&self,
identity: &PoolMemberIdentity,
scope: Option<&PoolScoreScope>,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<MySql>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(identity.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(identity.pool_id.clone())
.push(" AND member_kind = ")
.push_bind(identity.member_kind.clone())
.push(" AND member_id = ")
.push_bind(identity.member_id.clone());
if let Some(scope) = scope {
builder
.push(" AND capability = ")
.push_bind(scope.capability.clone())
.push(" AND scope_kind = ")
.push_bind(scope.scope_kind.clone());
if let Some(scope_id) = &scope.scope_id {
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
} else {
builder.push(" AND scope_id IS NULL");
}
}
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
rows.iter().map(map_score_row).collect()
}
}
#[async_trait]
impl PoolScoreReadRepository for MysqlPoolMemberScoreRepository {
async fn list_ranked_pool_members(
&self,
query: &ListRankedPoolMembersQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<MySql>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(query.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(query.pool_id.clone())
.push(" AND capability = ")
.push_bind(query.capability.clone())
.push(" AND scope_kind = ")
.push_bind(query.scope_kind.clone());
if let Some(scope_id) = &query.scope_id {
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
} else {
builder.push(" AND scope_id IS NULL");
}
if !query.hard_states.is_empty() {
builder.push(" AND hard_state IN (");
let mut separated = builder.separated(", ");
for state in &query.hard_states {
separated.push_bind(state.as_database());
}
separated.push_unseparated(")");
}
if let Some(statuses) = &query.probe_statuses {
if !statuses.is_empty() {
builder.push(" AND probe_status IN (");
let mut separated = builder.separated(", ");
for status in statuses {
separated.push_bind(status.as_database());
}
separated.push_unseparated(")");
}
}
builder
.push(" ORDER BY score DESC, last_ranked_at DESC, member_id ASC, id ASC")
.push(" LIMIT ")
.push_bind(i64_from_usize(query.limit.max(1), "pool score limit")?)
.push(" OFFSET ")
.push_bind(i64_from_usize(query.offset, "pool score offset")?);
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
rows.iter().map(map_score_row).collect()
}
async fn list_pool_member_scores(
&self,
query: &ListPoolMemberScoresQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<MySql>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(query.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(query.pool_id.clone());
if let Some(capability) = &query.capability {
builder
.push(" AND capability = ")
.push_bind(capability.clone());
}
if let Some(scope_kind) = &query.scope_kind {
builder
.push(" AND scope_kind = ")
.push_bind(scope_kind.clone());
}
if let Some(scope_id) = &query.scope_id {
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
}
if !query.hard_states.is_empty() {
builder.push(" AND hard_state IN (");
let mut separated = builder.separated(", ");
for state in &query.hard_states {
separated.push_bind(state.as_database());
}
separated.push_unseparated(")");
}
if let Some(statuses) = &query.probe_statuses {
if !statuses.is_empty() {
builder.push(" AND probe_status IN (");
let mut separated = builder.separated(", ");
for status in statuses {
separated.push_bind(status.as_database());
}
separated.push_unseparated(")");
}
}
builder
.push(" ORDER BY score DESC, last_ranked_at DESC, member_id ASC, id ASC")
.push(" LIMIT ")
.push_bind(i64_from_usize(query.limit.max(1), "pool score limit")?)
.push(" OFFSET ")
.push_bind(i64_from_usize(query.offset, "pool score offset")?);
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
rows.iter().map(map_score_row).collect()
}
async fn list_pool_member_probe_candidates(
&self,
query: &ListPoolMemberProbeCandidatesQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<MySql>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(query.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(query.pool_id.clone());
if let Some(capability) = &query.capability {
builder
.push(" AND capability = ")
.push_bind(capability.clone());
}
builder
.push(" AND hard_state IN ('available','unknown','cooldown','quota_exhausted')")
.push(" AND (probe_status IN ('never','failed','stale')")
.push(" OR (probe_status = 'ok' AND (last_probe_success_at IS NULL OR last_probe_success_at <= ")
.push_bind(i64_from_u64(
query.stale_before_unix_secs,
"pool probe stale_before_unix_secs",
)?)
.push("))")
.push(" OR (probe_status = 'in_progress' AND (last_probe_attempt_at IS NULL OR last_probe_attempt_at <= ")
.push_bind(i64_from_u64(
query.stale_before_unix_secs,
"pool probe stale_before_unix_secs",
)?)
.push(")))")
.push(
r#"
ORDER BY
CASE
WHEN last_scheduled_at IS NOT NULL AND probe_status <> 'ok' THEN 0
WHEN hard_state = 'quota_exhausted' THEN 1
WHEN hard_state = 'unknown' THEN 2
WHEN probe_status = 'stale' THEN 3
ELSE 4
END ASC,
probe_failure_count DESC,
COALESCE(last_probe_success_at, 0) ASC,
COALESCE(last_scheduled_at, 0) DESC,
member_id ASC
"#,
)
.push(" LIMIT ")
.push_bind(i64_from_usize(
query.limit.max(1),
"pool probe candidate limit",
)?);
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
rows.iter().map(map_score_row).collect()
}
async fn get_pool_member_scores_by_ids(
&self,
query: &GetPoolMemberScoresByIdsQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
if query.ids.is_empty() {
return Ok(Vec::new());
}
let mut builder = QueryBuilder::<MySql>::new(SCORE_COLUMNS);
builder.push(" WHERE id IN (");
let mut separated = builder.separated(", ");
for id in &query.ids {
separated.push_bind(id.clone());
}
separated.push_unseparated(")");
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
rows.iter().map(map_score_row).collect()
}
}
#[async_trait]
impl PoolMemberScoreWriteRepository for MysqlPoolMemberScoreRepository {
async fn upsert_pool_member_score(
&self,
score: UpsertPoolMemberScore,
) -> Result<StoredPoolMemberScore, DataLayerError> {
score.validate()?;
let stored = score.into_stored();
let score_reason = serde_json::to_string(&stored.score_reason)
.map_err(|err| DataLayerError::InvalidInput(err.to_string()))?;
sqlx::query(
r#"
INSERT INTO pool_member_scores (
id, pool_kind, pool_id, member_kind, member_id, capability, scope_kind, scope_id,
score, hard_state, score_version, score_reason, last_ranked_at, last_scheduled_at,
last_success_at, last_failure_at, failure_count, last_probe_attempt_at,
last_probe_success_at, last_probe_failure_at, probe_failure_count, probe_status, updated_at
) VALUES (
?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?
)
ON DUPLICATE KEY UPDATE
pool_kind = VALUES(pool_kind),
pool_id = VALUES(pool_id),
member_kind = VALUES(member_kind),
member_id = VALUES(member_id),
capability = VALUES(capability),
scope_kind = VALUES(scope_kind),
scope_id = VALUES(scope_id),
score = VALUES(score),
hard_state = VALUES(hard_state),
score_version = VALUES(score_version),
score_reason = VALUES(score_reason),
last_ranked_at = VALUES(last_ranked_at),
last_scheduled_at = COALESCE(VALUES(last_scheduled_at), last_scheduled_at),
last_success_at = COALESCE(VALUES(last_success_at), last_success_at),
last_failure_at = COALESCE(VALUES(last_failure_at), last_failure_at),
failure_count = VALUES(failure_count),
last_probe_attempt_at = COALESCE(VALUES(last_probe_attempt_at), last_probe_attempt_at),
last_probe_success_at = COALESCE(VALUES(last_probe_success_at), last_probe_success_at),
last_probe_failure_at = COALESCE(VALUES(last_probe_failure_at), last_probe_failure_at),
probe_failure_count = VALUES(probe_failure_count),
probe_status = VALUES(probe_status),
updated_at = VALUES(updated_at)
"#,
)
.bind(stored.id.as_str())
.bind(stored.pool_kind.as_str())
.bind(stored.pool_id.as_str())
.bind(stored.member_kind.as_str())
.bind(stored.member_id.as_str())
.bind(stored.capability.as_str())
.bind(stored.scope_kind.as_str())
.bind(stored.scope_id.as_deref())
.bind(stored.score)
.bind(stored.hard_state.as_database())
.bind(i64_from_u64(stored.score_version, "pool score version")?)
.bind(score_reason)
.bind(i64_opt_from_u64(
stored.last_ranked_at,
"pool score last_ranked_at",
)?)
.bind(i64_opt_from_u64(
stored.last_scheduled_at,
"pool score last_scheduled_at",
)?)
.bind(i64_opt_from_u64(
stored.last_success_at,
"pool score last_success_at",
)?)
.bind(i64_opt_from_u64(
stored.last_failure_at,
"pool score last_failure_at",
)?)
.bind(i64_from_u64(
stored.failure_count,
"pool score failure_count",
)?)
.bind(i64_opt_from_u64(
stored.last_probe_attempt_at,
"pool score last_probe_attempt_at",
)?)
.bind(i64_opt_from_u64(
stored.last_probe_success_at,
"pool score last_probe_success_at",
)?)
.bind(i64_opt_from_u64(
stored.last_probe_failure_at,
"pool score last_probe_failure_at",
)?)
.bind(i64_from_u64(
stored.probe_failure_count,
"pool score probe_failure_count",
)?)
.bind(stored.probe_status.as_database())
.bind(i64_from_u64(stored.updated_at, "pool score updated_at")?)
.execute(&self.pool)
.await
.map_sql_err()?;
Ok(stored)
}
async fn mark_pool_member_probe_in_progress(
&self,
attempt: PoolMemberProbeAttempt,
) -> Result<usize, DataLayerError> {
let rows = self
.find_scores_by_identity(&attempt.identity, attempt.scope.as_ref())
.await?;
let count = rows.len();
for mut row in rows {
row.last_probe_attempt_at = Some(attempt.attempted_at);
row.probe_status = PoolMemberProbeStatus::InProgress;
row.score_reason =
merge_score_reason_patch(row.score_reason, attempt.score_reason_patch.clone());
row.updated_at = attempt.attempted_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn record_pool_member_probe_result(
&self,
result: PoolMemberProbeResult,
) -> Result<usize, DataLayerError> {
let rows = self
.find_scores_by_identity(&result.identity, result.scope.as_ref())
.await?;
let count = rows.len();
for mut row in rows {
row.last_probe_attempt_at = Some(result.attempted_at);
row.probe_status = result.probe_status;
if result.succeeded {
row.last_probe_success_at = Some(result.attempted_at);
row.probe_failure_count = 0;
} else {
row.last_probe_failure_at = Some(result.attempted_at);
row.probe_failure_count = row.probe_failure_count.saturating_add(1);
}
if let Some(hard_state) = result.hard_state {
row.hard_state = hard_state;
}
row.score_reason =
merge_score_reason_patch(row.score_reason, result.score_reason_patch.clone());
row.updated_at = result.attempted_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn record_pool_member_schedule_feedback(
&self,
feedback: PoolMemberScheduleFeedback,
) -> Result<usize, DataLayerError> {
let rows = self
.find_scores_by_identity(&feedback.identity, feedback.scope.as_ref())
.await?;
let count = rows.len();
for mut row in rows {
row.last_scheduled_at = Some(feedback.scheduled_at);
match feedback.succeeded {
Some(true) => row.last_success_at = Some(feedback.scheduled_at),
Some(false) => {
row.last_failure_at = Some(feedback.scheduled_at);
row.failure_count = row.failure_count.saturating_add(1);
}
None => {}
}
if let Some(hard_state) = feedback.hard_state {
row.hard_state = hard_state;
}
row.score = score_with_delta(row.score, feedback.score_delta);
row.score_reason =
merge_score_reason_patch(row.score_reason, feedback.score_reason_patch.clone());
row.updated_at = feedback.scheduled_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn mark_pool_member_hard_state(
&self,
identity: &PoolMemberIdentity,
scope: Option<&PoolScoreScope>,
hard_state: PoolMemberHardState,
updated_at: u64,
) -> Result<usize, DataLayerError> {
let rows = self.find_scores_by_identity(identity, scope).await?;
let count = rows.len();
for mut row in rows {
row.hard_state = hard_state;
row.updated_at = updated_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn delete_pool_member_scores_for_member(
&self,
identity: &PoolMemberIdentity,
) -> Result<usize, DataLayerError> {
let result = sqlx::query(
r#"
DELETE FROM pool_member_scores
WHERE pool_kind = ? AND pool_id = ? AND member_kind = ? AND member_id = ?
"#,
)
.bind(identity.pool_kind.as_str())
.bind(identity.pool_id.as_str())
.bind(identity.member_kind.as_str())
.bind(identity.member_id.as_str())
.execute(&self.pool)
.await
.map_sql_err()?;
Ok(result.rows_affected() as usize)
}
}
fn map_score_row(row: &MySqlRow) -> Result<StoredPoolMemberScore, DataLayerError> {
let score_reason_raw: String = row.try_get("score_reason").map_sql_err()?;
Ok(StoredPoolMemberScore {
id: row.try_get("id").map_sql_err()?,
pool_kind: row.try_get("pool_kind").map_sql_err()?,
pool_id: row.try_get("pool_id").map_sql_err()?,
member_kind: row.try_get("member_kind").map_sql_err()?,
member_id: row.try_get("member_id").map_sql_err()?,
capability: row.try_get("capability").map_sql_err()?,
scope_kind: row.try_get("scope_kind").map_sql_err()?,
scope_id: row.try_get("scope_id").map_sql_err()?,
score: row.try_get("score").map_sql_err()?,
hard_state: PoolMemberHardState::from_database(
row.try_get::<String, _>("hard_state")
.map_sql_err()?
.as_str(),
)?,
score_version: u64_from_i64(
row.try_get("score_version").map_sql_err()?,
"pool_member_scores.score_version",
)?,
score_reason: serde_json::from_str(&score_reason_raw).unwrap_or(serde_json::Value::Null),
last_ranked_at: u64_opt_from_i64(
row.try_get("last_ranked_at").map_sql_err()?,
"pool_member_scores.last_ranked_at",
)?,
last_scheduled_at: u64_opt_from_i64(
row.try_get("last_scheduled_at").map_sql_err()?,
"pool_member_scores.last_scheduled_at",
)?,
last_success_at: u64_opt_from_i64(
row.try_get("last_success_at").map_sql_err()?,
"pool_member_scores.last_success_at",
)?,
last_failure_at: u64_opt_from_i64(
row.try_get("last_failure_at").map_sql_err()?,
"pool_member_scores.last_failure_at",
)?,
failure_count: u64_from_i64(
row.try_get("failure_count").map_sql_err()?,
"pool_member_scores.failure_count",
)?,
last_probe_attempt_at: u64_opt_from_i64(
row.try_get("last_probe_attempt_at").map_sql_err()?,
"pool_member_scores.last_probe_attempt_at",
)?,
last_probe_success_at: u64_opt_from_i64(
row.try_get("last_probe_success_at").map_sql_err()?,
"pool_member_scores.last_probe_success_at",
)?,
last_probe_failure_at: u64_opt_from_i64(
row.try_get("last_probe_failure_at").map_sql_err()?,
"pool_member_scores.last_probe_failure_at",
)?,
probe_failure_count: u64_from_i64(
row.try_get("probe_failure_count").map_sql_err()?,
"pool_member_scores.probe_failure_count",
)?,
probe_status: PoolMemberProbeStatus::from_database(
row.try_get::<String, _>("probe_status")
.map_sql_err()?
.as_str(),
)?,
updated_at: u64_from_i64(
row.try_get("updated_at").map_sql_err()?,
"pool_member_scores.updated_at",
)?,
})
}
fn upsert_from_stored(score: StoredPoolMemberScore) -> UpsertPoolMemberScore {
UpsertPoolMemberScore {
id: score.id,
identity: PoolMemberIdentity {
pool_kind: score.pool_kind,
pool_id: score.pool_id,
member_kind: score.member_kind,
member_id: score.member_id,
},
scope: PoolScoreScope {
capability: score.capability,
scope_kind: score.scope_kind,
scope_id: score.scope_id,
},
score: score.score,
hard_state: score.hard_state,
score_version: score.score_version,
score_reason: score.score_reason,
last_ranked_at: score.last_ranked_at,
last_scheduled_at: score.last_scheduled_at,
last_success_at: score.last_success_at,
last_failure_at: score.last_failure_at,
failure_count: score.failure_count,
last_probe_attempt_at: score.last_probe_attempt_at,
last_probe_success_at: score.last_probe_success_at,
last_probe_failure_at: score.last_probe_failure_at,
probe_failure_count: score.probe_failure_count,
probe_status: score.probe_status,
updated_at: score.updated_at,
}
}
fn i64_from_usize(value: usize, field: &str) -> Result<i64, DataLayerError> {
i64::try_from(value)
.map_err(|_| DataLayerError::InvalidInput(format!("{field} exceeds signed 64-bit range")))
}

588
crates/aether-data/src/repository/pool_scores/postgres.rs Normal file
View File

@@ -0,0 +1,588 @@
use async_trait::async_trait;
use sqlx::{postgres::PgRow, PgPool, Postgres, QueryBuilder, Row};
use super::{
i64_from_u64, i64_opt_from_u64, score_with_delta, u64_from_i64, u64_opt_from_i64,
GetPoolMemberScoresByIdsQuery, ListPoolMemberProbeCandidatesQuery, ListPoolMemberScoresQuery,
ListRankedPoolMembersQuery, PoolMemberHardState, PoolMemberIdentity, PoolMemberProbeAttempt,
PoolMemberProbeResult, PoolMemberProbeStatus, PoolMemberScheduleFeedback,
PoolMemberScoreWriteRepository, PoolScoreReadRepository, PoolScoreScope, StoredPoolMemberScore,
UpsertPoolMemberScore,
};
use crate::error::SqlxResultExt;
use crate::repository::pool_scores::merge_score_reason_patch;
use crate::DataLayerError;
const SCORE_COLUMNS: &str = r#"
SELECT
id,
pool_kind,
pool_id,
member_kind,
member_id,
capability,
scope_kind,
scope_id,
score,
hard_state,
score_version,
score_reason,
last_ranked_at,
last_scheduled_at,
last_success_at,
last_failure_at,
failure_count,
last_probe_attempt_at,
last_probe_success_at,
last_probe_failure_at,
probe_failure_count,
probe_status,
updated_at
FROM pool_member_scores
"#;
#[derive(Debug, Clone)]
pub struct PostgresPoolMemberScoreRepository {
pool: PgPool,
}
impl PostgresPoolMemberScoreRepository {
pub fn new(pool: PgPool) -> Self {
Self { pool }
}
async fn find_scores_by_identity(
&self,
identity: &PoolMemberIdentity,
scope: Option<&PoolScoreScope>,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<Postgres>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(identity.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(identity.pool_id.clone())
.push(" AND member_kind = ")
.push_bind(identity.member_kind.clone())
.push(" AND member_id = ")
.push_bind(identity.member_id.clone());
if let Some(scope) = scope {
builder
.push(" AND capability = ")
.push_bind(scope.capability.clone())
.push(" AND scope_kind = ")
.push_bind(scope.scope_kind.clone());
if let Some(scope_id) = &scope.scope_id {
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
} else {
builder.push(" AND scope_id IS NULL");
}
}
let rows = builder
.build()
.fetch_all(&self.pool)
.await
.map_postgres_err()?;
rows.iter().map(map_score_row).collect()
}
}
#[async_trait]
impl PoolScoreReadRepository for PostgresPoolMemberScoreRepository {
async fn list_ranked_pool_members(
&self,
query: &ListRankedPoolMembersQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<Postgres>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(query.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(query.pool_id.clone())
.push(" AND capability = ")
.push_bind(query.capability.clone())
.push(" AND scope_kind = ")
.push_bind(query.scope_kind.clone());
if let Some(scope_id) = &query.scope_id {
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
} else {
builder.push(" AND scope_id IS NULL");
}
if !query.hard_states.is_empty() {
builder.push(" AND hard_state IN (");
let mut separated = builder.separated(", ");
for state in &query.hard_states {
separated.push_bind(state.as_database());
}
separated.push_unseparated(")");
}
if let Some(statuses) = &query.probe_statuses {
if !statuses.is_empty() {
builder.push(" AND probe_status IN (");
let mut separated = builder.separated(", ");
for status in statuses {
separated.push_bind(status.as_database());
}
separated.push_unseparated(")");
}
}
builder
.push(" ORDER BY score DESC, last_ranked_at DESC NULLS LAST, member_id ASC, id ASC")
.push(" LIMIT ")
.push_bind(i64_from_usize(query.limit.max(1), "pool score limit")?)
.push(" OFFSET ")
.push_bind(i64_from_usize(query.offset, "pool score offset")?);
let rows = builder
.build()
.fetch_all(&self.pool)
.await
.map_postgres_err()?;
rows.iter().map(map_score_row).collect()
}
async fn list_pool_member_scores(
&self,
query: &ListPoolMemberScoresQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<Postgres>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(query.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(query.pool_id.clone());
if let Some(capability) = &query.capability {
builder
.push(" AND capability = ")
.push_bind(capability.clone());
}
if let Some(scope_kind) = &query.scope_kind {
builder
.push(" AND scope_kind = ")
.push_bind(scope_kind.clone());
}
if let Some(scope_id) = &query.scope_id {
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
}
if !query.hard_states.is_empty() {
builder.push(" AND hard_state IN (");
let mut separated = builder.separated(", ");
for state in &query.hard_states {
separated.push_bind(state.as_database());
}
separated.push_unseparated(")");
}
if let Some(statuses) = &query.probe_statuses {
if !statuses.is_empty() {
builder.push(" AND probe_status IN (");
let mut separated = builder.separated(", ");
for status in statuses {
separated.push_bind(status.as_database());
}
separated.push_unseparated(")");
}
}
builder
.push(" ORDER BY score DESC, last_ranked_at DESC NULLS LAST, member_id ASC, id ASC")
.push(" LIMIT ")
.push_bind(i64_from_usize(query.limit.max(1), "pool score limit")?)
.push(" OFFSET ")
.push_bind(i64_from_usize(query.offset, "pool score offset")?);
let rows = builder
.build()
.fetch_all(&self.pool)
.await
.map_postgres_err()?;
rows.iter().map(map_score_row).collect()
}
async fn list_pool_member_probe_candidates(
&self,
query: &ListPoolMemberProbeCandidatesQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<Postgres>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(query.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(query.pool_id.clone());
if let Some(capability) = &query.capability {
builder
.push(" AND capability = ")
.push_bind(capability.clone());
}
builder
.push(" AND hard_state IN ('available','unknown','cooldown','quota_exhausted')")
.push(" AND (probe_status IN ('never','failed','stale')")
.push(" OR (probe_status = 'ok' AND (last_probe_success_at IS NULL OR last_probe_success_at <= ")
.push_bind(i64_from_u64(
query.stale_before_unix_secs,
"pool probe stale_before_unix_secs",
)?)
.push("))")
.push(" OR (probe_status = 'in_progress' AND (last_probe_attempt_at IS NULL OR last_probe_attempt_at <= ")
.push_bind(i64_from_u64(
query.stale_before_unix_secs,
"pool probe stale_before_unix_secs",
)?)
.push(")))")
.push(
r#"
ORDER BY
CASE
WHEN last_scheduled_at IS NOT NULL AND probe_status <> 'ok' THEN 0
WHEN hard_state = 'quota_exhausted' THEN 1
WHEN hard_state = 'unknown' THEN 2
WHEN probe_status = 'stale' THEN 3
ELSE 4
END ASC,
probe_failure_count DESC,
COALESCE(last_probe_success_at, 0) ASC,
COALESCE(last_scheduled_at, 0) DESC,
member_id ASC
"#,
)
.push(" LIMIT ")
.push_bind(i64_from_usize(
query.limit.max(1),
"pool probe candidate limit",
)?);
let rows = builder
.build()
.fetch_all(&self.pool)
.await
.map_postgres_err()?;
rows.iter().map(map_score_row).collect()
}
async fn get_pool_member_scores_by_ids(
&self,
query: &GetPoolMemberScoresByIdsQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
if query.ids.is_empty() {
return Ok(Vec::new());
}
let mut builder = QueryBuilder::<Postgres>::new(SCORE_COLUMNS);
builder.push(" WHERE id IN (");
let mut separated = builder.separated(", ");
for id in &query.ids {
separated.push_bind(id.clone());
}
separated.push_unseparated(")");
let rows = builder
.build()
.fetch_all(&self.pool)
.await
.map_postgres_err()?;
rows.iter().map(map_score_row).collect()
}
}
#[async_trait]
impl PoolMemberScoreWriteRepository for PostgresPoolMemberScoreRepository {
async fn upsert_pool_member_score(
&self,
score: UpsertPoolMemberScore,
) -> Result<StoredPoolMemberScore, DataLayerError> {
score.validate()?;
let stored = score.clone().into_stored();
sqlx::query(
r#"
INSERT INTO pool_member_scores (
id, pool_kind, pool_id, member_kind, member_id, capability, scope_kind, scope_id,
score, hard_state, score_version, score_reason, last_ranked_at, last_scheduled_at,
last_success_at, last_failure_at, failure_count, last_probe_attempt_at,
last_probe_success_at, last_probe_failure_at, probe_failure_count, probe_status, updated_at
) VALUES (
$1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19,$20,$21,$22,$23
)
ON CONFLICT(id) DO UPDATE SET
pool_kind = EXCLUDED.pool_kind,
pool_id = EXCLUDED.pool_id,
member_kind = EXCLUDED.member_kind,
member_id = EXCLUDED.member_id,
capability = EXCLUDED.capability,
scope_kind = EXCLUDED.scope_kind,
scope_id = EXCLUDED.scope_id,
score = EXCLUDED.score,
hard_state = EXCLUDED.hard_state,
score_version = EXCLUDED.score_version,
score_reason = EXCLUDED.score_reason,
last_ranked_at = EXCLUDED.last_ranked_at,
last_scheduled_at = COALESCE(EXCLUDED.last_scheduled_at, pool_member_scores.last_scheduled_at),
last_success_at = COALESCE(EXCLUDED.last_success_at, pool_member_scores.last_success_at),
last_failure_at = COALESCE(EXCLUDED.last_failure_at, pool_member_scores.last_failure_at),
failure_count = EXCLUDED.failure_count,
last_probe_attempt_at = COALESCE(EXCLUDED.last_probe_attempt_at, pool_member_scores.last_probe_attempt_at),
last_probe_success_at = COALESCE(EXCLUDED.last_probe_success_at, pool_member_scores.last_probe_success_at),
last_probe_failure_at = COALESCE(EXCLUDED.last_probe_failure_at, pool_member_scores.last_probe_failure_at),
probe_failure_count = EXCLUDED.probe_failure_count,
probe_status = EXCLUDED.probe_status,
updated_at = EXCLUDED.updated_at
"#,
)
.bind(stored.id.as_str())
.bind(stored.pool_kind.as_str())
.bind(stored.pool_id.as_str())
.bind(stored.member_kind.as_str())
.bind(stored.member_id.as_str())
.bind(stored.capability.as_str())
.bind(stored.scope_kind.as_str())
.bind(stored.scope_id.as_deref())
.bind(stored.score)
.bind(stored.hard_state.as_database())
.bind(i64_from_u64(stored.score_version, "pool score version")?)
.bind(&stored.score_reason)
.bind(i64_opt_from_u64(stored.last_ranked_at, "pool score last_ranked_at")?)
.bind(i64_opt_from_u64(stored.last_scheduled_at, "pool score last_scheduled_at")?)
.bind(i64_opt_from_u64(stored.last_success_at, "pool score last_success_at")?)
.bind(i64_opt_from_u64(stored.last_failure_at, "pool score last_failure_at")?)
.bind(i64_from_u64(stored.failure_count, "pool score failure_count")?)
.bind(i64_opt_from_u64(
stored.last_probe_attempt_at,
"pool score last_probe_attempt_at",
)?)
.bind(i64_opt_from_u64(
stored.last_probe_success_at,
"pool score last_probe_success_at",
)?)
.bind(i64_opt_from_u64(
stored.last_probe_failure_at,
"pool score last_probe_failure_at",
)?)
.bind(i64_from_u64(
stored.probe_failure_count,
"pool score probe_failure_count",
)?)
.bind(stored.probe_status.as_database())
.bind(i64_from_u64(stored.updated_at, "pool score updated_at")?)
.execute(&self.pool)
.await
.map_postgres_err()?;
Ok(stored)
}
async fn mark_pool_member_probe_in_progress(
&self,
attempt: PoolMemberProbeAttempt,
) -> Result<usize, DataLayerError> {
let rows = self
.find_scores_by_identity(&attempt.identity, attempt.scope.as_ref())
.await?;
let count = rows.len();
for mut row in rows {
row.last_probe_attempt_at = Some(attempt.attempted_at);
row.probe_status = PoolMemberProbeStatus::InProgress;
row.score_reason =
merge_score_reason_patch(row.score_reason, attempt.score_reason_patch.clone());
row.updated_at = attempt.attempted_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn record_pool_member_probe_result(
&self,
result: PoolMemberProbeResult,
) -> Result<usize, DataLayerError> {
let rows = self
.find_scores_by_identity(&result.identity, result.scope.as_ref())
.await?;
let count = rows.len();
for mut row in rows {
row.last_probe_attempt_at = Some(result.attempted_at);
row.probe_status = result.probe_status;
if result.succeeded {
row.last_probe_success_at = Some(result.attempted_at);
row.probe_failure_count = 0;
} else {
row.last_probe_failure_at = Some(result.attempted_at);
row.probe_failure_count = row.probe_failure_count.saturating_add(1);
}
if let Some(hard_state) = result.hard_state {
row.hard_state = hard_state;
}
row.score_reason =
merge_score_reason_patch(row.score_reason, result.score_reason_patch.clone());
row.updated_at = result.attempted_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn record_pool_member_schedule_feedback(
&self,
feedback: PoolMemberScheduleFeedback,
) -> Result<usize, DataLayerError> {
let rows = self
.find_scores_by_identity(&feedback.identity, feedback.scope.as_ref())
.await?;
let count = rows.len();
for mut row in rows {
row.last_scheduled_at = Some(feedback.scheduled_at);
match feedback.succeeded {
Some(true) => row.last_success_at = Some(feedback.scheduled_at),
Some(false) => {
row.last_failure_at = Some(feedback.scheduled_at);
row.failure_count = row.failure_count.saturating_add(1);
}
None => {}
}
if let Some(hard_state) = feedback.hard_state {
row.hard_state = hard_state;
}
row.score = score_with_delta(row.score, feedback.score_delta);
row.score_reason =
merge_score_reason_patch(row.score_reason, feedback.score_reason_patch.clone());
row.updated_at = feedback.scheduled_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn mark_pool_member_hard_state(
&self,
identity: &PoolMemberIdentity,
scope: Option<&PoolScoreScope>,
hard_state: PoolMemberHardState,
updated_at: u64,
) -> Result<usize, DataLayerError> {
let rows = self.find_scores_by_identity(identity, scope).await?;
let count = rows.len();
for mut row in rows {
row.hard_state = hard_state;
row.updated_at = updated_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn delete_pool_member_scores_for_member(
&self,
identity: &PoolMemberIdentity,
) -> Result<usize, DataLayerError> {
let result = sqlx::query(
r#"
DELETE FROM pool_member_scores
WHERE pool_kind = $1 AND pool_id = $2 AND member_kind = $3 AND member_id = $4
"#,
)
.bind(identity.pool_kind.as_str())
.bind(identity.pool_id.as_str())
.bind(identity.member_kind.as_str())
.bind(identity.member_id.as_str())
.execute(&self.pool)
.await
.map_postgres_err()?;
Ok(result.rows_affected() as usize)
}
}
fn map_score_row(row: &PgRow) -> Result<StoredPoolMemberScore, DataLayerError> {
Ok(StoredPoolMemberScore {
id: row.try_get("id").map_postgres_err()?,
pool_kind: row.try_get("pool_kind").map_postgres_err()?,
pool_id: row.try_get("pool_id").map_postgres_err()?,
member_kind: row.try_get("member_kind").map_postgres_err()?,
member_id: row.try_get("member_id").map_postgres_err()?,
capability: row.try_get("capability").map_postgres_err()?,
scope_kind: row.try_get("scope_kind").map_postgres_err()?,
scope_id: row.try_get("scope_id").map_postgres_err()?,
score: row.try_get("score").map_postgres_err()?,
hard_state: PoolMemberHardState::from_database(
row.try_get::<String, _>("hard_state")
.map_postgres_err()?
.as_str(),
)?,
score_version: u64_from_i64(
row.try_get("score_version").map_postgres_err()?,
"pool_member_scores.score_version",
)?,
score_reason: row.try_get("score_reason").map_postgres_err()?,
last_ranked_at: u64_opt_from_i64(
row.try_get("last_ranked_at").map_postgres_err()?,
"pool_member_scores.last_ranked_at",
)?,
last_scheduled_at: u64_opt_from_i64(
row.try_get("last_scheduled_at").map_postgres_err()?,
"pool_member_scores.last_scheduled_at",
)?,
last_success_at: u64_opt_from_i64(
row.try_get("last_success_at").map_postgres_err()?,
"pool_member_scores.last_success_at",
)?,
last_failure_at: u64_opt_from_i64(
row.try_get("last_failure_at").map_postgres_err()?,
"pool_member_scores.last_failure_at",
)?,
failure_count: u64_from_i64(
row.try_get("failure_count").map_postgres_err()?,
"pool_member_scores.failure_count",
)?,
last_probe_attempt_at: u64_opt_from_i64(
row.try_get("last_probe_attempt_at").map_postgres_err()?,
"pool_member_scores.last_probe_attempt_at",
)?,
last_probe_success_at: u64_opt_from_i64(
row.try_get("last_probe_success_at").map_postgres_err()?,
"pool_member_scores.last_probe_success_at",
)?,
last_probe_failure_at: u64_opt_from_i64(
row.try_get("last_probe_failure_at").map_postgres_err()?,
"pool_member_scores.last_probe_failure_at",
)?,
probe_failure_count: u64_from_i64(
row.try_get("probe_failure_count").map_postgres_err()?,
"pool_member_scores.probe_failure_count",
)?,
probe_status: PoolMemberProbeStatus::from_database(
row.try_get::<String, _>("probe_status")
.map_postgres_err()?
.as_str(),
)?,
updated_at: u64_from_i64(
row.try_get("updated_at").map_postgres_err()?,
"pool_member_scores.updated_at",
)?,
})
}
fn upsert_from_stored(score: StoredPoolMemberScore) -> UpsertPoolMemberScore {
UpsertPoolMemberScore {
id: score.id,
identity: PoolMemberIdentity {
pool_kind: score.pool_kind,
pool_id: score.pool_id,
member_kind: score.member_kind,
member_id: score.member_id,
},
scope: PoolScoreScope {
capability: score.capability,
scope_kind: score.scope_kind,
scope_id: score.scope_id,
},
score: score.score,
hard_state: score.hard_state,
score_version: score.score_version,
score_reason: score.score_reason,
last_ranked_at: score.last_ranked_at,
last_scheduled_at: score.last_scheduled_at,
last_success_at: score.last_success_at,
last_failure_at: score.last_failure_at,
failure_count: score.failure_count,
last_probe_attempt_at: score.last_probe_attempt_at,
last_probe_success_at: score.last_probe_success_at,
last_probe_failure_at: score.last_probe_failure_at,
probe_failure_count: score.probe_failure_count,
probe_status: score.probe_status,
updated_at: score.updated_at,
}
}
fn i64_from_usize(value: usize, field: &str) -> Result<i64, DataLayerError> {
i64::try_from(value)
.map_err(|_| DataLayerError::InvalidInput(format!("{field} exceeds signed 64-bit range")))
}

571
crates/aether-data/src/repository/pool_scores/sqlite.rs Normal file
View File

@@ -0,0 +1,571 @@
use async_trait::async_trait;
use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite, SqlitePool};
use super::{
i64_from_u64, i64_opt_from_u64, score_with_delta, u64_from_i64, u64_opt_from_i64,
GetPoolMemberScoresByIdsQuery, ListPoolMemberProbeCandidatesQuery, ListPoolMemberScoresQuery,
ListRankedPoolMembersQuery, PoolMemberHardState, PoolMemberIdentity, PoolMemberProbeAttempt,
PoolMemberProbeResult, PoolMemberProbeStatus, PoolMemberScheduleFeedback,
PoolMemberScoreWriteRepository, PoolScoreReadRepository, PoolScoreScope, StoredPoolMemberScore,
UpsertPoolMemberScore,
};
use crate::error::SqlResultExt;
use crate::repository::pool_scores::merge_score_reason_patch;
use crate::DataLayerError;
const SCORE_COLUMNS: &str = r#"
SELECT
id,
pool_kind,
pool_id,
member_kind,
member_id,
capability,
scope_kind,
scope_id,
score,
hard_state,
score_version,
score_reason,
last_ranked_at,
last_scheduled_at,
last_success_at,
last_failure_at,
failure_count,
last_probe_attempt_at,
last_probe_success_at,
last_probe_failure_at,
probe_failure_count,
probe_status,
updated_at
FROM pool_member_scores
"#;
#[derive(Debug, Clone)]
pub struct SqlitePoolMemberScoreRepository {
pool: SqlitePool,
}
impl SqlitePoolMemberScoreRepository {
pub fn new(pool: SqlitePool) -> Self {
Self { pool }
}
async fn find_scores_by_identity(
&self,
identity: &PoolMemberIdentity,
scope: Option<&PoolScoreScope>,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<Sqlite>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(identity.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(identity.pool_id.clone())
.push(" AND member_kind = ")
.push_bind(identity.member_kind.clone())
.push(" AND member_id = ")
.push_bind(identity.member_id.clone());
if let Some(scope) = scope {
builder
.push(" AND capability = ")
.push_bind(scope.capability.clone())
.push(" AND scope_kind = ")
.push_bind(scope.scope_kind.clone());
if let Some(scope_id) = &scope.scope_id {
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
} else {
builder.push(" AND scope_id IS NULL");
}
}
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
rows.iter().map(map_score_row).collect()
}
}
#[async_trait]
impl PoolScoreReadRepository for SqlitePoolMemberScoreRepository {
async fn list_ranked_pool_members(
&self,
query: &ListRankedPoolMembersQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<Sqlite>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(query.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(query.pool_id.clone())
.push(" AND capability = ")
.push_bind(query.capability.clone())
.push(" AND scope_kind = ")
.push_bind(query.scope_kind.clone());
if let Some(scope_id) = &query.scope_id {
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
} else {
builder.push(" AND scope_id IS NULL");
}
if !query.hard_states.is_empty() {
builder.push(" AND hard_state IN (");
let mut separated = builder.separated(", ");
for state in &query.hard_states {
separated.push_bind(state.as_database());
}
separated.push_unseparated(")");
}
if let Some(statuses) = &query.probe_statuses {
if !statuses.is_empty() {
builder.push(" AND probe_status IN (");
let mut separated = builder.separated(", ");
for status in statuses {
separated.push_bind(status.as_database());
}
separated.push_unseparated(")");
}
}
builder
.push(" ORDER BY score DESC, last_ranked_at DESC, member_id ASC, id ASC")
.push(" LIMIT ")
.push_bind(i64_from_usize(query.limit.max(1), "pool score limit")?)
.push(" OFFSET ")
.push_bind(i64_from_usize(query.offset, "pool score offset")?);
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
rows.iter().map(map_score_row).collect()
}
async fn list_pool_member_scores(
&self,
query: &ListPoolMemberScoresQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<Sqlite>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(query.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(query.pool_id.clone());
if let Some(capability) = &query.capability {
builder
.push(" AND capability = ")
.push_bind(capability.clone());
}
if let Some(scope_kind) = &query.scope_kind {
builder
.push(" AND scope_kind = ")
.push_bind(scope_kind.clone());
}
if let Some(scope_id) = &query.scope_id {
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
}
if !query.hard_states.is_empty() {
builder.push(" AND hard_state IN (");
let mut separated = builder.separated(", ");
for state in &query.hard_states {
separated.push_bind(state.as_database());
}
separated.push_unseparated(")");
}
if let Some(statuses) = &query.probe_statuses {
if !statuses.is_empty() {
builder.push(" AND probe_status IN (");
let mut separated = builder.separated(", ");
for status in statuses {
separated.push_bind(status.as_database());
}
separated.push_unseparated(")");
}
}
builder
.push(" ORDER BY score DESC, last_ranked_at DESC, member_id ASC, id ASC")
.push(" LIMIT ")
.push_bind(i64_from_usize(query.limit.max(1), "pool score limit")?)
.push(" OFFSET ")
.push_bind(i64_from_usize(query.offset, "pool score offset")?);
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
rows.iter().map(map_score_row).collect()
}
async fn list_pool_member_probe_candidates(
&self,
query: &ListPoolMemberProbeCandidatesQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<Sqlite>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(query.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(query.pool_id.clone());
if let Some(capability) = &query.capability {
builder
.push(" AND capability = ")
.push_bind(capability.clone());
}
builder
.push(" AND hard_state IN ('available','unknown','cooldown','quota_exhausted')")
.push(" AND (probe_status IN ('never','failed','stale')")
.push(" OR (probe_status = 'ok' AND (last_probe_success_at IS NULL OR last_probe_success_at <= ")
.push_bind(i64_from_u64(
query.stale_before_unix_secs,
"pool probe stale_before_unix_secs",
)?)
.push("))")
.push(" OR (probe_status = 'in_progress' AND (last_probe_attempt_at IS NULL OR last_probe_attempt_at <= ")
.push_bind(i64_from_u64(
query.stale_before_unix_secs,
"pool probe stale_before_unix_secs",
)?)
.push(")))")
.push(
r#"
ORDER BY
CASE
WHEN last_scheduled_at IS NOT NULL AND probe_status <> 'ok' THEN 0
WHEN hard_state = 'quota_exhausted' THEN 1
WHEN hard_state = 'unknown' THEN 2
WHEN probe_status = 'stale' THEN 3
ELSE 4
END ASC,
probe_failure_count DESC,
COALESCE(last_probe_success_at, 0) ASC,
COALESCE(last_scheduled_at, 0) DESC,
member_id ASC
"#,
)
.push(" LIMIT ")
.push_bind(i64_from_usize(
query.limit.max(1),
"pool probe candidate limit",
)?);
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
rows.iter().map(map_score_row).collect()
}
async fn get_pool_member_scores_by_ids(
&self,
query: &GetPoolMemberScoresByIdsQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
if query.ids.is_empty() {
return Ok(Vec::new());
}
let mut builder = QueryBuilder::<Sqlite>::new(SCORE_COLUMNS);
builder.push(" WHERE id IN (");
let mut separated = builder.separated(", ");
for id in &query.ids {
separated.push_bind(id.clone());
}
separated.push_unseparated(")");
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
rows.iter().map(map_score_row).collect()
}
}
#[async_trait]
impl PoolMemberScoreWriteRepository for SqlitePoolMemberScoreRepository {
async fn upsert_pool_member_score(
&self,
score: UpsertPoolMemberScore,
) -> Result<StoredPoolMemberScore, DataLayerError> {
score.validate()?;
let stored = score.into_stored();
let score_reason = serde_json::to_string(&stored.score_reason)
.map_err(|err| DataLayerError::InvalidInput(err.to_string()))?;
sqlx::query(
r#"
INSERT INTO pool_member_scores (
id, pool_kind, pool_id, member_kind, member_id, capability, scope_kind, scope_id,
score, hard_state, score_version, score_reason, last_ranked_at, last_scheduled_at,
last_success_at, last_failure_at, failure_count, last_probe_attempt_at,
last_probe_success_at, last_probe_failure_at, probe_failure_count, probe_status, updated_at
) VALUES (
?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?
)
ON CONFLICT(id) DO UPDATE SET
pool_kind = excluded.pool_kind,
pool_id = excluded.pool_id,
member_kind = excluded.member_kind,
member_id = excluded.member_id,
capability = excluded.capability,
scope_kind = excluded.scope_kind,
scope_id = excluded.scope_id,
score = excluded.score,
hard_state = excluded.hard_state,
score_version = excluded.score_version,
score_reason = excluded.score_reason,
last_ranked_at = excluded.last_ranked_at,
last_scheduled_at = COALESCE(excluded.last_scheduled_at, pool_member_scores.last_scheduled_at),
last_success_at = COALESCE(excluded.last_success_at, pool_member_scores.last_success_at),
last_failure_at = COALESCE(excluded.last_failure_at, pool_member_scores.last_failure_at),
failure_count = excluded.failure_count,
last_probe_attempt_at = COALESCE(excluded.last_probe_attempt_at, pool_member_scores.last_probe_attempt_at),
last_probe_success_at = COALESCE(excluded.last_probe_success_at, pool_member_scores.last_probe_success_at),
last_probe_failure_at = COALESCE(excluded.last_probe_failure_at, pool_member_scores.last_probe_failure_at),
probe_failure_count = excluded.probe_failure_count,
probe_status = excluded.probe_status,
updated_at = excluded.updated_at
"#,
)
.bind(stored.id.as_str())
.bind(stored.pool_kind.as_str())
.bind(stored.pool_id.as_str())
.bind(stored.member_kind.as_str())
.bind(stored.member_id.as_str())
.bind(stored.capability.as_str())
.bind(stored.scope_kind.as_str())
.bind(stored.scope_id.as_deref())
.bind(stored.score)
.bind(stored.hard_state.as_database())
.bind(i64_from_u64(stored.score_version, "pool score version")?)
.bind(score_reason)
.bind(i64_opt_from_u64(stored.last_ranked_at, "pool score last_ranked_at")?)
.bind(i64_opt_from_u64(stored.last_scheduled_at, "pool score last_scheduled_at")?)
.bind(i64_opt_from_u64(stored.last_success_at, "pool score last_success_at")?)
.bind(i64_opt_from_u64(stored.last_failure_at, "pool score last_failure_at")?)
.bind(i64_from_u64(stored.failure_count, "pool score failure_count")?)
.bind(i64_opt_from_u64(
stored.last_probe_attempt_at,
"pool score last_probe_attempt_at",
)?)
.bind(i64_opt_from_u64(
stored.last_probe_success_at,
"pool score last_probe_success_at",
)?)
.bind(i64_opt_from_u64(
stored.last_probe_failure_at,
"pool score last_probe_failure_at",
)?)
.bind(i64_from_u64(
stored.probe_failure_count,
"pool score probe_failure_count",
)?)
.bind(stored.probe_status.as_database())
.bind(i64_from_u64(stored.updated_at, "pool score updated_at")?)
.execute(&self.pool)
.await
.map_sql_err()?;
Ok(stored)
}
async fn mark_pool_member_probe_in_progress(
&self,
attempt: PoolMemberProbeAttempt,
) -> Result<usize, DataLayerError> {
let rows = self
.find_scores_by_identity(&attempt.identity, attempt.scope.as_ref())
.await?;
let count = rows.len();
for mut row in rows {
row.last_probe_attempt_at = Some(attempt.attempted_at);
row.probe_status = PoolMemberProbeStatus::InProgress;
row.score_reason =
merge_score_reason_patch(row.score_reason, attempt.score_reason_patch.clone());
row.updated_at = attempt.attempted_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn record_pool_member_probe_result(
&self,
result: PoolMemberProbeResult,
) -> Result<usize, DataLayerError> {
let rows = self
.find_scores_by_identity(&result.identity, result.scope.as_ref())
.await?;
let count = rows.len();
for mut row in rows {
row.last_probe_attempt_at = Some(result.attempted_at);
row.probe_status = result.probe_status;
if result.succeeded {
row.last_probe_success_at = Some(result.attempted_at);
row.probe_failure_count = 0;
} else {
row.last_probe_failure_at = Some(result.attempted_at);
row.probe_failure_count = row.probe_failure_count.saturating_add(1);
}
if let Some(hard_state) = result.hard_state {
row.hard_state = hard_state;
}
row.score_reason =
merge_score_reason_patch(row.score_reason, result.score_reason_patch.clone());
row.updated_at = result.attempted_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn record_pool_member_schedule_feedback(
&self,
feedback: PoolMemberScheduleFeedback,
) -> Result<usize, DataLayerError> {
let rows = self
.find_scores_by_identity(&feedback.identity, feedback.scope.as_ref())
.await?;
let count = rows.len();
for mut row in rows {
row.last_scheduled_at = Some(feedback.scheduled_at);
match feedback.succeeded {
Some(true) => row.last_success_at = Some(feedback.scheduled_at),
Some(false) => {
row.last_failure_at = Some(feedback.scheduled_at);
row.failure_count = row.failure_count.saturating_add(1);
}
None => {}
}
if let Some(hard_state) = feedback.hard_state {
row.hard_state = hard_state;
}
row.score = score_with_delta(row.score, feedback.score_delta);
row.score_reason =
merge_score_reason_patch(row.score_reason, feedback.score_reason_patch.clone());
row.updated_at = feedback.scheduled_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn mark_pool_member_hard_state(
&self,
identity: &PoolMemberIdentity,
scope: Option<&PoolScoreScope>,
hard_state: PoolMemberHardState,
updated_at: u64,
) -> Result<usize, DataLayerError> {
let rows = self.find_scores_by_identity(identity, scope).await?;
let count = rows.len();
for mut row in rows {
row.hard_state = hard_state;
row.updated_at = updated_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn delete_pool_member_scores_for_member(
&self,
identity: &PoolMemberIdentity,
) -> Result<usize, DataLayerError> {
let result = sqlx::query(
r#"
DELETE FROM pool_member_scores
WHERE pool_kind = ? AND pool_id = ? AND member_kind = ? AND member_id = ?
"#,
)
.bind(identity.pool_kind.as_str())
.bind(identity.pool_id.as_str())
.bind(identity.member_kind.as_str())
.bind(identity.member_id.as_str())
.execute(&self.pool)
.await
.map_sql_err()?;
Ok(result.rows_affected() as usize)
}
}
fn map_score_row(row: &SqliteRow) -> Result<StoredPoolMemberScore, DataLayerError> {
let score_reason_raw: String = row.try_get("score_reason").map_sql_err()?;
Ok(StoredPoolMemberScore {
id: row.try_get("id").map_sql_err()?,
pool_kind: row.try_get("pool_kind").map_sql_err()?,
pool_id: row.try_get("pool_id").map_sql_err()?,
member_kind: row.try_get("member_kind").map_sql_err()?,
member_id: row.try_get("member_id").map_sql_err()?,
capability: row.try_get("capability").map_sql_err()?,
scope_kind: row.try_get("scope_kind").map_sql_err()?,
scope_id: row.try_get("scope_id").map_sql_err()?,
score: row.try_get("score").map_sql_err()?,
hard_state: PoolMemberHardState::from_database(
row.try_get::<String, _>("hard_state")
.map_sql_err()?
.as_str(),
)?,
score_version: u64_from_i64(
row.try_get("score_version").map_sql_err()?,
"pool_member_scores.score_version",
)?,
score_reason: serde_json::from_str(&score_reason_raw).unwrap_or(serde_json::Value::Null),
last_ranked_at: u64_opt_from_i64(
row.try_get("last_ranked_at").map_sql_err()?,
"pool_member_scores.last_ranked_at",
)?,
last_scheduled_at: u64_opt_from_i64(
row.try_get("last_scheduled_at").map_sql_err()?,
"pool_member_scores.last_scheduled_at",
)?,
last_success_at: u64_opt_from_i64(
row.try_get("last_success_at").map_sql_err()?,
"pool_member_scores.last_success_at",
)?,
last_failure_at: u64_opt_from_i64(
row.try_get("last_failure_at").map_sql_err()?,
"pool_member_scores.last_failure_at",
)?,
failure_count: u64_from_i64(
row.try_get("failure_count").map_sql_err()?,
"pool_member_scores.failure_count",
)?,
last_probe_attempt_at: u64_opt_from_i64(
row.try_get("last_probe_attempt_at").map_sql_err()?,
"pool_member_scores.last_probe_attempt_at",
)?,
last_probe_success_at: u64_opt_from_i64(
row.try_get("last_probe_success_at").map_sql_err()?,
"pool_member_scores.last_probe_success_at",
)?,
last_probe_failure_at: u64_opt_from_i64(
row.try_get("last_probe_failure_at").map_sql_err()?,
"pool_member_scores.last_probe_failure_at",
)?,
probe_failure_count: u64_from_i64(
row.try_get("probe_failure_count").map_sql_err()?,
"pool_member_scores.probe_failure_count",
)?,
probe_status: PoolMemberProbeStatus::from_database(
row.try_get::<String, _>("probe_status")
.map_sql_err()?
.as_str(),
)?,
updated_at: u64_from_i64(
row.try_get("updated_at").map_sql_err()?,
"pool_member_scores.updated_at",
)?,
})
}
fn upsert_from_stored(score: StoredPoolMemberScore) -> UpsertPoolMemberScore {
UpsertPoolMemberScore {
id: score.id,
identity: PoolMemberIdentity {
pool_kind: score.pool_kind,
pool_id: score.pool_id,
member_kind: score.member_kind,
member_id: score.member_id,
},
scope: PoolScoreScope {
capability: score.capability,
scope_kind: score.scope_kind,
scope_id: score.scope_id,
},
score: score.score,
hard_state: score.hard_state,
score_version: score.score_version,
score_reason: score.score_reason,
last_ranked_at: score.last_ranked_at,
last_scheduled_at: score.last_scheduled_at,
last_success_at: score.last_success_at,
last_failure_at: score.last_failure_at,
failure_count: score.failure_count,
last_probe_attempt_at: score.last_probe_attempt_at,
last_probe_success_at: score.last_probe_success_at,
last_probe_failure_at: score.last_probe_failure_at,
probe_failure_count: score.probe_failure_count,
probe_status: score.probe_status,
updated_at: score.updated_at,
}
}
fn i64_from_usize(value: usize, field: &str) -> Result<i64, DataLayerError> {
i64::try_from(value)
.map_err(|_| DataLayerError::InvalidInput(format!("{field} exceeds signed 64-bit range")))
}