Skip to main content

stateset_db/sqlite/
agent_identities.rs

1//! SQLite agent identity registry repository implementation
2
3use super::{
4    map_db_error, params_refs, parse_datetime_opt_row, parse_datetime_row, parse_enum_row,
5    parse_uuid_opt_row, parse_uuid_row,
6};
7use chrono::Utc;
8use r2d2::Pool;
9use r2d2_sqlite::SqliteConnectionManager;
10use rusqlite::OptionalExtension;
11use stateset_core::{
12    AgentIdentity, AgentIdentityFilter, AgentIdentityRepository, AgentMetadataEntry,
13    AgentWalletProofType, CommerceError, CreateAgentIdentity, Result, UpdateAgentIdentity,
14};
15use uuid::Uuid;
16
17/// SQLite implementation of `AgentIdentityRepository`
18#[derive(Debug)]
19pub struct SqliteAgentIdentityRepository {
20    pool: Pool<SqliteConnectionManager>,
21}
22
23impl SqliteAgentIdentityRepository {
24    #[must_use]
25    pub const fn new(pool: Pool<SqliteConnectionManager>) -> Self {
26        Self { pool }
27    }
28
29    fn conn(&self) -> Result<r2d2::PooledConnection<SqliteConnectionManager>> {
30        self.pool.get().map_err(|e| CommerceError::DatabaseError(e.to_string()))
31    }
32
33    fn row_to_agent_identity(row: &rusqlite::Row<'_>) -> rusqlite::Result<AgentIdentity> {
34        let wallet_proof_type = match row.get::<_, Option<String>>("wallet_proof_type")? {
35            Some(val) => Some(parse_enum_row(&val, "agent_identity", "wallet_proof_type")?),
36            None => None,
37        };
38
39        Ok(AgentIdentity {
40            id: parse_uuid_row(&row.get::<_, String>("id")?, "agent_identity", "id")?,
41            agent_registry: row.get("agent_registry")?,
42            agent_id: row.get("agent_id")?,
43            agent_uri: row.get("agent_uri")?,
44            agent_wallet: row.get("agent_wallet")?,
45            owner_address: row.get("owner_address")?,
46            agent_card_id: parse_uuid_opt_row(
47                row.get::<_, Option<String>>("agent_card_id")?,
48                "agent_identity",
49                "agent_card_id",
50            )?,
51            registration: row.get("registration")?,
52            registration_hash: row.get("registration_hash")?,
53            wallet_proof_type,
54            wallet_proof: row.get("wallet_proof")?,
55            wallet_proof_chain_id: row
56                .get::<_, Option<i64>>("wallet_proof_chain_id")?
57                .map(|n| n as u64),
58            wallet_proof_deadline: parse_datetime_opt_row(
59                row.get::<_, Option<String>>("wallet_proof_deadline")?,
60                "agent_identity",
61                "wallet_proof_deadline",
62            )?,
63            active: row.get::<_, i32>("active")? == 1,
64            created_at: parse_datetime_row(
65                &row.get::<_, String>("created_at")?,
66                "agent_identity",
67                "created_at",
68            )?,
69            updated_at: parse_datetime_row(
70                &row.get::<_, String>("updated_at")?,
71                "agent_identity",
72                "updated_at",
73            )?,
74        })
75    }
76}
77
78impl AgentIdentityRepository for SqliteAgentIdentityRepository {
79    fn register(&self, input: CreateAgentIdentity) -> Result<AgentIdentity> {
80        let conn = self.conn()?;
81        let now = Utc::now();
82        let id = Uuid::new_v4();
83        let active = input.active.unwrap_or(true);
84
85        conn.execute(
86            "INSERT INTO agent_identities (
87                id, agent_registry, agent_id, agent_uri, agent_wallet, owner_address,
88                agent_card_id, registration, registration_hash,
89                wallet_proof_type, wallet_proof, wallet_proof_chain_id, wallet_proof_deadline,
90                active, created_at, updated_at
91            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
92            rusqlite::params![
93                id.to_string(),
94                input.agent_registry,
95                input.agent_id,
96                input.agent_uri,
97                input.agent_wallet,
98                input.owner_address,
99                input.agent_card_id.map(|id| id.to_string()),
100                input.registration,
101                input.registration_hash,
102                input.wallet_proof_type.map(|t| t.to_string()),
103                input.wallet_proof,
104                input.wallet_proof_chain_id.map(|n| n as i64),
105                input.wallet_proof_deadline.map(|d| d.to_rfc3339()),
106                i32::from(active),
107                now.to_rfc3339(),
108                now.to_rfc3339(),
109            ],
110        )
111        .map_err(map_db_error)?;
112
113        self.get(&input.agent_registry, &input.agent_id)?.ok_or(CommerceError::NotFound)
114    }
115
116    fn get(&self, agent_registry: &str, agent_id: &str) -> Result<Option<AgentIdentity>> {
117        let conn = self.conn()?;
118        let mut stmt = conn
119            .prepare("SELECT * FROM agent_identities WHERE agent_registry = ? AND agent_id = ?")
120            .map_err(map_db_error)?;
121
122        stmt.query_row([agent_registry, agent_id], Self::row_to_agent_identity)
123            .optional()
124            .map_err(map_db_error)
125    }
126
127    fn get_by_wallet(&self, agent_wallet: &str) -> Result<Option<AgentIdentity>> {
128        let conn = self.conn()?;
129        let mut stmt = conn
130            .prepare("SELECT * FROM agent_identities WHERE agent_wallet = ?")
131            .map_err(map_db_error)?;
132
133        stmt.query_row([agent_wallet], Self::row_to_agent_identity).optional().map_err(map_db_error)
134    }
135
136    fn update(
137        &self,
138        agent_registry: &str,
139        agent_id: &str,
140        input: UpdateAgentIdentity,
141    ) -> Result<AgentIdentity> {
142        let conn = self.conn()?;
143        let existing = self.get(agent_registry, agent_id)?.ok_or(CommerceError::NotFound)?;
144
145        let agent_uri = input.agent_uri.unwrap_or(existing.agent_uri);
146        let agent_wallet = input.agent_wallet.or(existing.agent_wallet);
147        let owner_address = input.owner_address.or(existing.owner_address);
148        let agent_card_id = input.agent_card_id.or(existing.agent_card_id);
149        let registration = input.registration.or(existing.registration);
150        let registration_hash = input.registration_hash.or(existing.registration_hash);
151        let wallet_proof_type = input.wallet_proof_type.or(existing.wallet_proof_type);
152        let wallet_proof = input.wallet_proof.or(existing.wallet_proof);
153        let wallet_proof_chain_id = input.wallet_proof_chain_id.or(existing.wallet_proof_chain_id);
154        let wallet_proof_deadline = input.wallet_proof_deadline.or(existing.wallet_proof_deadline);
155        let active = input.active.unwrap_or(existing.active);
156
157        conn.execute(
158            "UPDATE agent_identities SET
159                agent_uri = ?, agent_wallet = ?, owner_address = ?, agent_card_id = ?,
160                registration = ?, registration_hash = ?,
161                wallet_proof_type = ?, wallet_proof = ?, wallet_proof_chain_id = ?, wallet_proof_deadline = ?,
162                active = ?, updated_at = ?
163             WHERE agent_registry = ? AND agent_id = ?",
164            rusqlite::params![
165                agent_uri,
166                agent_wallet,
167                owner_address,
168                agent_card_id.map(|id| id.to_string()),
169                registration,
170                registration_hash,
171                wallet_proof_type.map(|t| t.to_string()),
172                wallet_proof,
173                wallet_proof_chain_id.map(|n| n as i64),
174                wallet_proof_deadline.map(|d| d.to_rfc3339()),
175                i32::from(active),
176                Utc::now().to_rfc3339(),
177                agent_registry,
178                agent_id,
179            ],
180        )
181        .map_err(map_db_error)?;
182
183        self.get(agent_registry, agent_id)?.ok_or(CommerceError::NotFound)
184    }
185
186    fn set_agent_wallet(
187        &self,
188        agent_registry: &str,
189        agent_id: &str,
190        agent_wallet: &str,
191        proof_type: Option<AgentWalletProofType>,
192        proof: Option<&str>,
193        proof_chain_id: Option<u64>,
194        proof_deadline: Option<chrono::DateTime<Utc>>,
195    ) -> Result<AgentIdentity> {
196        let conn = self.conn()?;
197
198        conn.execute(
199            "UPDATE agent_identities SET
200                agent_wallet = ?, wallet_proof_type = ?, wallet_proof = ?,
201                wallet_proof_chain_id = ?, wallet_proof_deadline = ?, updated_at = ?
202             WHERE agent_registry = ? AND agent_id = ?",
203            rusqlite::params![
204                agent_wallet,
205                proof_type.map(|t| t.to_string()),
206                proof.map(std::string::ToString::to_string),
207                proof_chain_id.map(|n| n as i64),
208                proof_deadline.map(|d| d.to_rfc3339()),
209                Utc::now().to_rfc3339(),
210                agent_registry,
211                agent_id,
212            ],
213        )
214        .map_err(map_db_error)?;
215
216        self.get(agent_registry, agent_id)?.ok_or(CommerceError::NotFound)
217    }
218
219    fn clear_agent_wallet(&self, agent_registry: &str, agent_id: &str) -> Result<AgentIdentity> {
220        let conn = self.conn()?;
221
222        conn.execute(
223            "UPDATE agent_identities SET
224                agent_wallet = NULL, wallet_proof_type = NULL, wallet_proof = NULL,
225                wallet_proof_chain_id = NULL, wallet_proof_deadline = NULL, updated_at = ?
226             WHERE agent_registry = ? AND agent_id = ?",
227            rusqlite::params![Utc::now().to_rfc3339(), agent_registry, agent_id,],
228        )
229        .map_err(map_db_error)?;
230
231        self.get(agent_registry, agent_id)?.ok_or(CommerceError::NotFound)
232    }
233
234    fn list(&self, filter: AgentIdentityFilter) -> Result<Vec<AgentIdentity>> {
235        let conn = self.conn()?;
236        let mut conditions = vec!["1=1".to_string()];
237        let mut params: Vec<Box<dyn rusqlite::ToSql>> = vec![];
238
239        if let Some(ref registry) = filter.agent_registry {
240            conditions.push("agent_registry = ?".to_string());
241            params.push(Box::new(registry.clone()));
242        }
243        if let Some(ref agent_id) = filter.agent_id {
244            conditions.push("agent_id = ?".to_string());
245            params.push(Box::new(agent_id.clone()));
246        }
247        if let Some(ref wallet) = filter.agent_wallet {
248            conditions.push("agent_wallet = ?".to_string());
249            params.push(Box::new(wallet.clone()));
250        }
251        if let Some(ref owner) = filter.owner_address {
252            conditions.push("owner_address = ?".to_string());
253            params.push(Box::new(owner.clone()));
254        }
255        if let Some(ref card_id) = filter.agent_card_id {
256            conditions.push("agent_card_id = ?".to_string());
257            params.push(Box::new(card_id.to_string()));
258        }
259        if let Some(active) = filter.active {
260            conditions.push("active = ?".to_string());
261            params.push(Box::new(i32::from(active)));
262        }
263
264        let limit = filter.limit.unwrap_or(100).min(1000);
265        let offset = filter.offset.unwrap_or(0);
266
267        let sql = format!(
268            "SELECT * FROM agent_identities WHERE {} ORDER BY created_at DESC LIMIT ? OFFSET ?",
269            conditions.join(" AND ")
270        );
271        params.push(Box::new(i64::from(limit)));
272        params.push(Box::new(i64::from(offset)));
273
274        let param_refs = params_refs(&params);
275        let mut stmt = conn.prepare(&sql).map_err(map_db_error)?;
276        let rows = stmt
277            .query_map(rusqlite::params_from_iter(param_refs), Self::row_to_agent_identity)
278            .map_err(map_db_error)?;
279
280        let mut results = Vec::new();
281        for row in rows {
282            results.push(row.map_err(map_db_error)?);
283        }
284        Ok(results)
285    }
286
287    fn count(&self, filter: AgentIdentityFilter) -> Result<u64> {
288        let conn = self.conn()?;
289        let mut conditions = vec!["1=1".to_string()];
290        let mut params: Vec<Box<dyn rusqlite::ToSql>> = vec![];
291
292        if let Some(ref registry) = filter.agent_registry {
293            conditions.push("agent_registry = ?".to_string());
294            params.push(Box::new(registry.clone()));
295        }
296        if let Some(ref agent_id) = filter.agent_id {
297            conditions.push("agent_id = ?".to_string());
298            params.push(Box::new(agent_id.clone()));
299        }
300        if let Some(ref wallet) = filter.agent_wallet {
301            conditions.push("agent_wallet = ?".to_string());
302            params.push(Box::new(wallet.clone()));
303        }
304        if let Some(ref owner) = filter.owner_address {
305            conditions.push("owner_address = ?".to_string());
306            params.push(Box::new(owner.clone()));
307        }
308        if let Some(ref card_id) = filter.agent_card_id {
309            conditions.push("agent_card_id = ?".to_string());
310            params.push(Box::new(card_id.to_string()));
311        }
312        if let Some(active) = filter.active {
313            conditions.push("active = ?".to_string());
314            params.push(Box::new(i32::from(active)));
315        }
316
317        let sql =
318            format!("SELECT COUNT(*) FROM agent_identities WHERE {}", conditions.join(" AND "));
319
320        let param_refs = params_refs(&params);
321        let count: i64 = conn
322            .query_row(&sql, rusqlite::params_from_iter(param_refs), |row| row.get(0))
323            .map_err(map_db_error)?;
324
325        Ok(count as u64)
326    }
327
328    fn set_metadata(
329        &self,
330        agent_registry: &str,
331        agent_id: &str,
332        entry: AgentMetadataEntry,
333    ) -> Result<()> {
334        let conn = self.conn()?;
335        let now = Utc::now();
336
337        conn.execute(
338            "INSERT INTO agent_identity_metadata (
339                agent_registry, agent_id, metadata_key, metadata_value, created_at, updated_at
340            ) VALUES (?, ?, ?, ?, ?, ?)
341            ON CONFLICT(agent_registry, agent_id, metadata_key)
342            DO UPDATE SET metadata_value = excluded.metadata_value, updated_at = excluded.updated_at",
343            rusqlite::params![
344                agent_registry,
345                agent_id,
346                entry.metadata_key,
347                entry.metadata_value,
348                now.to_rfc3339(),
349                now.to_rfc3339(),
350            ],
351        )
352        .map_err(map_db_error)?;
353
354        Ok(())
355    }
356
357    fn get_metadata(
358        &self,
359        agent_registry: &str,
360        agent_id: &str,
361        metadata_key: &str,
362    ) -> Result<Option<Vec<u8>>> {
363        let conn = self.conn()?;
364        let mut stmt = conn
365            .prepare(
366                "SELECT metadata_value FROM agent_identity_metadata
367                 WHERE agent_registry = ? AND agent_id = ? AND metadata_key = ?",
368            )
369            .map_err(map_db_error)?;
370
371        stmt.query_row([agent_registry, agent_id, metadata_key], |row| row.get(0))
372            .optional()
373            .map_err(map_db_error)
374    }
375
376    fn delete_metadata(
377        &self,
378        agent_registry: &str,
379        agent_id: &str,
380        metadata_key: &str,
381    ) -> Result<()> {
382        let conn = self.conn()?;
383        conn.execute(
384            "DELETE FROM agent_identity_metadata WHERE agent_registry = ? AND agent_id = ? AND metadata_key = ?",
385            rusqlite::params![agent_registry, agent_id, metadata_key],
386        )
387        .map_err(map_db_error)?;
388        Ok(())
389    }
390}