1use 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#[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(¶ms);
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(¶ms);
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}