Merge pull request 'refactor(custody): replace the store abstraction with direct database functions' (#109) from cleverwild/rrvrwqun into feat-shamir-custody

Reviewed-on: #109

server/crates/arbiter-server/migrations/2026-02-14-171124-0000_init/up.sql

@@ -53,7 +53,7 @@ create table if not exists operator_identity (
 create unique index if not exists uniq_operator_identity_public_key on operator_identity (public_key);
 
 create table if not exists operator (
-    id integer primary key references operator_identity(id) on delete restrict, -- same id as operator_identity
+    id integer not null primary key references operator_identity(id) on delete restrict, -- same id as operator_identity
 
     share blob not null,
     share_nonce blob not null,

server/crates/arbiter-server/src/db/custody.rs

@@ -1,18 +1,13 @@
 //! Storage for Shamir custody material: the reconstruction threshold and the
 //! per-operator encrypted shares of the vault seal key.
-//!
-//! Every query lives here so that the actors above hold no Diesel code of their
-//! own. The functions borrow the caller's connection instead of taking one from
-//! the pool, which lets the vault write custody material inside the same
-//! transaction that stores the root key.
 
 use std::collections::HashMap;
 
-use diesel::{ExpressionMethods as _, QueryDsl, sqlite::Sqlite};
+use diesel::{ExpressionMethods as _, QueryDsl, SelectableHelper as _, sqlite::Sqlite};
 use diesel_async::{AsyncConnection, RunQueryDsl};
 
 use crate::db::{
-    models::{OperatorId, SqliteTimestamp},
+    models::{Operator, OperatorId, SqliteTimestamp},
     schema,
 };
 
@@ -55,7 +50,7 @@ pub async fn write_record(
     for (operator_id, share) in &record.shares {
         diesel::replace_into(schema::operator::table)
             .values((
-                schema::operator::id.eq(Some(*operator_id)),
+                schema::operator::id.eq(*operator_id),
                 schema::operator::share.eq(&share.ciphertext),
                 schema::operator::share_nonce.eq(&share.nonce),
                 schema::operator::share_salt.eq(&share.salt),
@@ -90,41 +85,29 @@ pub async fn threshold(conn: &mut impl AsyncConnection<Backend = Sqlite>) -> Res
         .ok_or(Error::BrokenThreshold)
 }
 
-/// One row of the share query: operator id, ciphertext, nonce, salt.
-type ShareRow = (Option<OperatorId>, Vec<u8>, Vec<u8>, Vec<u8>);
-
 /// Load the shares of `operators` in one query, in the order requested.
 pub async fn shares(
     conn: &mut impl AsyncConnection<Backend = Sqlite>,
     operators: &[OperatorId],
 ) -> Result<Vec<EncryptedShare>, Error> {
-    let wanted: Vec<Option<OperatorId>> = operators.iter().copied().map(Some).collect();
+    let rows: Vec<Operator> = schema::operator::table
-
+        .filter(schema::operator::id.eq_any(operators))
-    let rows: Vec<ShareRow> = schema::operator::table
+        .select(Operator::as_select())
-        .filter(schema::operator::id.eq_any(wanted))
-        .select((
-            schema::operator::id,
-            schema::operator::share,
-            schema::operator::share_nonce,
-            schema::operator::share_salt,
-        ))
         .load(conn)
         .await?;
 
     let mut found: HashMap<OperatorId, EncryptedShare> = rows
         .into_iter()
-        .filter_map(|(id, ciphertext, nonce, salt)| {
+        .map(|row| {
-            id.map(|id| {
             (
-                    id,
+                row.id,
                 EncryptedShare {
-                        ciphertext,
+                    ciphertext: row.share,
-                        nonce,
+                    nonce: row.share_nonce,
-                        salt,
+                    salt: row.share_salt,
                 },
             )
         })
-        })
         .collect();
 
     operators

server/crates/arbiter-server/src/db/models.rs

@@ -292,7 +292,7 @@ pub struct OperatorClient {
     pub updated_at: SqliteTimestamp,
 }
 
-#[derive(Queryable, Debug)]
+#[derive(Queryable, Debug, Selectable)]
 #[diesel(table_name = schema::operator, check_for_backend(Sqlite))]
 pub struct Operator {
     pub id: OperatorId,

server/crates/arbiter-server/src/db/schema.rs

@@ -155,7 +155,7 @@ diesel::table! {
 
 diesel::table! {
     operator (id) {
-        id -> Nullable<Integer>,
+        id -> Integer,
         share -> Binary,
         share_nonce -> Binary,
         share_salt -> Binary,