1use std::collections::HashSet;
4use std::sync::Arc;
5
6use async_trait::async_trait;
7use rusqlite::OptionalExtension;
8use uuid::Uuid;
9
10use khive_storage::entity::{Entity, EntityFilter};
11use khive_storage::error::StorageError;
12use khive_storage::types::{
13 BatchWriteSummary, DeleteMode, Page, PageRequest, SeekCursor, SeekPage, SqlStatement, SqlValue,
14};
15use khive_storage::EntityStore;
16use khive_storage::StorageCapability;
17
18use crate::error::SqliteError;
19use crate::pool::ConnectionPool;
20use crate::sql_bridge::bind_params;
21use crate::writer_task::WriterTaskHandle;
22
23fn map_err(e: rusqlite::Error, op: &'static str) -> StorageError {
24 StorageError::driver(StorageCapability::Entities, op, e)
25}
26
27fn map_sqlite_err(e: SqliteError, op: &'static str) -> StorageError {
28 StorageError::driver(StorageCapability::Entities, op, e)
29}
30
31const NAMESPACE_COUNT_CHUNK_SIZE: usize = 500;
32
33pub fn entity_upsert_statement(entity: &Entity) -> SqlStatement {
49 let properties_str = entity
50 .properties
51 .as_ref()
52 .map(|v| serde_json::to_string(v).unwrap_or_default());
53 let tags_str = serde_json::to_string(&entity.tags).unwrap_or_else(|_| "[]".to_string());
54 SqlStatement {
55 sql: "INSERT OR REPLACE INTO entities \
56 (id, namespace, kind, entity_type, name, description, properties, tags, \
57 created_at, updated_at, deleted_at, merged_into, merge_event_id, content_ref) \
58 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14)"
59 .to_string(),
60 params: vec![
61 SqlValue::Text(entity.id.to_string()),
62 SqlValue::Text(entity.namespace.clone()),
63 SqlValue::Text(entity.kind.clone()),
64 match &entity.entity_type {
65 Some(t) => SqlValue::Text(t.clone()),
66 None => SqlValue::Null,
67 },
68 SqlValue::Text(entity.name.clone()),
69 match &entity.description {
70 Some(d) => SqlValue::Text(d.clone()),
71 None => SqlValue::Null,
72 },
73 match properties_str {
74 Some(p) => SqlValue::Text(p),
75 None => SqlValue::Null,
76 },
77 SqlValue::Text(tags_str),
78 SqlValue::Integer(entity.created_at),
79 SqlValue::Integer(entity.updated_at),
80 match entity.deleted_at {
81 Some(d) => SqlValue::Integer(d),
82 None => SqlValue::Null,
83 },
84 match entity.merged_into {
85 Some(u) => SqlValue::Text(u.to_string()),
86 None => SqlValue::Null,
87 },
88 match entity.merge_event_id {
89 Some(u) => SqlValue::Text(u.to_string()),
90 None => SqlValue::Null,
91 },
92 match &entity.content_ref {
93 Some(c) => SqlValue::Text(c.clone()),
94 None => SqlValue::Null,
95 },
96 ],
97 label: Some("entity-upsert".to_string()),
98 }
99}
100
101pub fn entity_soft_delete_statement(id: Uuid, deleted_at: i64) -> SqlStatement {
103 SqlStatement {
104 sql: "UPDATE entities SET deleted_at = ?1 WHERE id = ?2 AND deleted_at IS NULL".to_string(),
105 params: vec![
106 SqlValue::Integer(deleted_at),
107 SqlValue::Text(id.to_string()),
108 ],
109 label: Some("entity-delete-soft".to_string()),
110 }
111}
112
113pub fn entity_hard_delete_statement(id: Uuid) -> SqlStatement {
116 SqlStatement {
117 sql: "DELETE FROM entities WHERE id = ?1".to_string(),
118 params: vec![SqlValue::Text(id.to_string())],
119 label: Some("entity-delete-hard".to_string()),
120 }
121}
122
123pub struct SqlEntityStore {
128 pool: Arc<ConnectionPool>,
129 is_file_backed: bool,
130 writer_task: Option<WriterTaskHandle>,
131}
132
133impl SqlEntityStore {
134 pub fn new(pool: Arc<ConnectionPool>, is_file_backed: bool) -> Self {
151 let writer_task = pool.writer_task_handle().ok().flatten();
157
158 Self {
159 pool,
160 is_file_backed,
161 writer_task,
162 }
163 }
164
165 fn open_standalone_reader(&self) -> Result<rusqlite::Connection, StorageError> {
166 let config = self.pool.config();
167 let path = config.path.as_ref().ok_or_else(|| StorageError::Pool {
168 operation: "entity_reader".into(),
169 message: "in-memory databases do not support standalone connections".into(),
170 })?;
171
172 let conn = rusqlite::Connection::open_with_flags(
173 path,
174 rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY
175 | rusqlite::OpenFlags::SQLITE_OPEN_NO_MUTEX
176 | rusqlite::OpenFlags::SQLITE_OPEN_URI,
177 )
178 .map_err(|e| map_err(e, "open_entity_reader"))?;
179
180 conn.busy_timeout(config.busy_timeout)
181 .map_err(|e| map_err(e, "open_entity_reader"))?;
182 conn.pragma_update(None, "foreign_keys", "ON")
183 .map_err(|e| map_err(e, "open_entity_reader"))?;
184 conn.pragma_update(None, "synchronous", "NORMAL")
185 .map_err(|e| map_err(e, "open_entity_reader"))?;
186
187 Ok(conn)
188 }
189
190 async fn with_writer<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
207 where
208 F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
209 R: Send + 'static,
210 {
211 if let Some(writer_task) = &self.writer_task {
212 return writer_task
213 .send(move |conn| f(conn).map_err(|e| map_err(e, op)))
214 .await;
215 }
216
217 let pool = Arc::clone(&self.pool);
218 tokio::task::spawn_blocking(move || {
219 let guard = pool.try_writer().map_err(|e| map_sqlite_err(e, op))?;
220 f(guard.conn()).map_err(|e| map_err(e, op))
221 })
222 .await
223 .map_err(|e| StorageError::driver(StorageCapability::Entities, op, e))?
224 }
225
226 async fn with_reader<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
227 where
228 F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
229 R: Send + 'static,
230 {
231 if self.is_file_backed {
232 let conn = self.open_standalone_reader()?;
233 tokio::task::spawn_blocking(move || f(&conn).map_err(|e| map_err(e, op)))
234 .await
235 .map_err(|e| StorageError::driver(StorageCapability::Entities, op, e))?
236 } else {
237 let pool = Arc::clone(&self.pool);
238 tokio::task::spawn_blocking(move || {
239 let guard = pool.reader().map_err(|e| map_sqlite_err(e, op))?;
240 f(guard.conn()).map_err(|e| map_err(e, op))
241 })
242 .await
243 .map_err(|e| StorageError::driver(StorageCapability::Entities, op, e))?
244 }
245 }
246}
247
248fn read_entity(row: &rusqlite::Row<'_>) -> Result<Entity, rusqlite::Error> {
253 let id_str: String = row.get(0)?;
254 let namespace: String = row.get(1)?;
255 let kind: String = row.get(2)?;
256 let entity_type: Option<String> = row.get(3)?;
257 let name: String = row.get(4)?;
258 let description: Option<String> = row.get(5)?;
259 let properties_str: Option<String> = row.get(6)?;
260 let tags_str: String = row.get(7)?;
261 let created_at: i64 = row.get(8)?;
262 let updated_at: i64 = row.get(9)?;
263 let deleted_at: Option<i64> = row.get(10)?;
264 let merged_into_str: Option<String> = row.get(11)?;
265 let merge_event_id_str: Option<String> = row.get(12)?;
266 let content_ref: Option<String> = row.get(13)?;
267
268 let id = parse_uuid(&id_str)?;
269
270 let properties = properties_str
271 .map(|s| {
272 serde_json::from_str(&s).map_err(|e| {
273 rusqlite::Error::FromSqlConversionFailure(
274 6,
275 rusqlite::types::Type::Text,
276 Box::new(e),
277 )
278 })
279 })
280 .transpose()?;
281
282 let tags: Vec<String> = serde_json::from_str(&tags_str).map_err(|e| {
283 rusqlite::Error::FromSqlConversionFailure(7, rusqlite::types::Type::Text, Box::new(e))
284 })?;
285
286 let merged_into = merged_into_str
287 .as_deref()
288 .map(Uuid::parse_str)
289 .transpose()
290 .map_err(|e| {
291 rusqlite::Error::FromSqlConversionFailure(10, rusqlite::types::Type::Text, Box::new(e))
292 })?;
293
294 let merge_event_id = merge_event_id_str
295 .as_deref()
296 .map(Uuid::parse_str)
297 .transpose()
298 .map_err(|e| {
299 rusqlite::Error::FromSqlConversionFailure(11, rusqlite::types::Type::Text, Box::new(e))
300 })?;
301
302 Ok(Entity {
303 id,
304 namespace,
305 kind,
306 entity_type,
307 name,
308 description,
309 properties,
310 tags,
311 created_at,
312 updated_at,
313 deleted_at,
314 merged_into,
315 merge_event_id,
316 content_ref,
317 })
318}
319
320fn batch_upsert_entities(
330 conn: &rusqlite::Connection,
331 entities: &[Entity],
332 attempted: u64,
333) -> Result<BatchWriteSummary, rusqlite::Error> {
334 let mut affected = 0u64;
335 let mut failed = 0u64;
336 let mut first_error = String::new();
337
338 for entity in entities {
339 let id_str = entity.id.to_string();
340 let properties_str = entity
341 .properties
342 .as_ref()
343 .map(|v| serde_json::to_string(v).unwrap_or_default());
344 let tags_str = serde_json::to_string(&entity.tags).unwrap_or_else(|_| "[]".to_string());
345
346 let merged_into_str = entity.merged_into.map(|u| u.to_string());
347 let merge_event_id_str = entity.merge_event_id.map(|u| u.to_string());
348 match conn.execute(
349 "INSERT OR REPLACE INTO entities \
350 (id, namespace, kind, entity_type, name, description, properties, tags, \
351 created_at, updated_at, deleted_at, merged_into, merge_event_id, content_ref) \
352 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14)",
353 rusqlite::params![
354 id_str,
355 &entity.namespace,
356 entity.kind,
357 entity.entity_type,
358 entity.name,
359 entity.description,
360 properties_str,
361 tags_str,
362 entity.created_at,
363 entity.updated_at,
364 entity.deleted_at,
365 merged_into_str,
366 merge_event_id_str,
367 entity.content_ref,
368 ],
369 ) {
370 Ok(_) => affected += 1,
371 Err(e) => {
372 if first_error.is_empty() {
373 first_error = e.to_string();
374 }
375 failed += 1;
376 }
377 }
378 }
379
380 Ok(BatchWriteSummary {
381 attempted,
382 affected,
383 failed,
384 first_error,
385 })
386}
387
388fn parse_uuid(s: &str) -> Result<Uuid, rusqlite::Error> {
389 Uuid::parse_str(s).map_err(|e| {
390 rusqlite::Error::FromSqlConversionFailure(0, rusqlite::types::Type::Text, Box::new(e))
391 })
392}
393
394fn escape_like(input: &str) -> String {
400 let mut out = String::with_capacity(input.len());
401 for c in input.chars() {
402 if matches!(c, '\\' | '%' | '_') {
403 out.push('\\');
404 }
405 out.push(c);
406 }
407 out
408}
409
410fn build_entity_where(
411 namespace: &str,
412 filter: &EntityFilter,
413) -> (String, Vec<Box<dyn rusqlite::types::ToSql>>) {
414 let (ns_condition, ns_params): (String, Vec<Box<dyn rusqlite::types::ToSql>>) =
418 if !filter.namespaces.is_empty() {
419 let placeholders: Vec<String> = (1..=filter.namespaces.len())
420 .map(|i| format!("?{i}"))
421 .collect();
422 let params: Vec<Box<dyn rusqlite::types::ToSql>> = filter
423 .namespaces
424 .iter()
425 .map(|ns| -> Box<dyn rusqlite::types::ToSql> { Box::new(ns.clone()) })
426 .collect();
427 (
428 format!("namespace IN ({})", placeholders.join(", ")),
429 params,
430 )
431 } else {
432 (
433 "namespace = ?1".to_string(),
434 vec![Box::new(namespace.to_string())],
435 )
436 };
437
438 let mut conditions: Vec<String> = vec![ns_condition, "deleted_at IS NULL".to_string()];
439 let mut params: Vec<Box<dyn rusqlite::types::ToSql>> = ns_params;
440
441 if !filter.ids.is_empty() {
442 let placeholders: Vec<String> = filter
443 .ids
444 .iter()
445 .map(|id| {
446 params.push(Box::new(id.to_string()));
447 format!("?{}", params.len())
448 })
449 .collect();
450 conditions.push(format!("id IN ({})", placeholders.join(", ")));
451 }
452
453 if !filter.kinds.is_empty() {
454 let placeholders: Vec<String> = filter
455 .kinds
456 .iter()
457 .map(|k| {
458 params.push(Box::new(k.clone()));
459 format!("?{}", params.len())
460 })
461 .collect();
462 conditions.push(format!("kind IN ({})", placeholders.join(", ")));
463 }
464
465 if !filter.entity_types.is_empty() {
466 let placeholders: Vec<String> = filter
467 .entity_types
468 .iter()
469 .map(|t| {
470 params.push(Box::new(t.clone()));
471 format!("?{}", params.len())
472 })
473 .collect();
474 conditions.push(format!("entity_type IN ({})", placeholders.join(", ")));
475 }
476
477 if let Some(ref prefix) = filter.name_prefix {
478 params.push(Box::new(format!("{}%", escape_like(prefix))));
479 conditions.push(format!("name LIKE ?{} ESCAPE '\\'", params.len()));
480 }
481
482 if let Some(ref exact) = filter.name_exact {
483 params.push(Box::new(exact.clone()));
484 conditions.push(format!("name = ?{} COLLATE BINARY", params.len()));
489 }
490
491 if !filter.names_ci.is_empty() {
492 let placeholders: Vec<String> = filter
495 .names_ci
496 .iter()
497 .map(|n| {
498 params.push(Box::new(n.to_ascii_lowercase()));
499 format!("?{}", params.len())
500 })
501 .collect();
502 conditions.push(format!("LOWER(name) IN ({})", placeholders.join(", ")));
503 }
504
505 if !filter.tags_any.is_empty() {
506 let placeholders: Vec<String> = filter
507 .tags_any
508 .iter()
509 .map(|t| {
510 params.push(Box::new(t.to_lowercase()));
513 format!("?{}", params.len())
514 })
515 .collect();
516 conditions.push(format!(
517 "EXISTS (SELECT 1 FROM json_each(tags) WHERE LOWER(json_each.value) IN ({}))",
518 placeholders.join(", ")
519 ));
520 }
521
522 let clause = format!(" WHERE {}", conditions.join(" AND "));
523 (clause, params)
524}
525
526fn build_candidate_entity_query(
527 columns: &str,
528 where_sql: &str,
529 candidate_param_indices: &[usize],
530 order_by: &str,
531 limit_idx: usize,
532 offset_idx: usize,
533) -> String {
534 let candidate_rows = candidate_param_indices
535 .iter()
536 .map(|idx| format!("(?{idx})"))
537 .collect::<Vec<_>>()
538 .join(", ");
539
540 format!(
541 "WITH candidates(folded_name) AS (VALUES {candidate_rows}), \
542 matched_entities(entity_id) AS (\
543 SELECT (\
544 SELECT id FROM entities{where_sql} \
545 AND LOWER(name) = candidates.folded_name LIMIT 1\
546 ) FROM candidates\
547 ) \
548 SELECT {columns} FROM entities \
549 JOIN matched_entities ON entities.id = matched_entities.entity_id \
550 ORDER BY {order_by} LIMIT ?{limit_idx} OFFSET ?{offset_idx}"
551 )
552}
553
554fn is_complete_id_lookup(filter: &EntityFilter, page: &PageRequest) -> bool {
555 !filter.ids.is_empty()
556 && filter.kinds.is_empty()
557 && filter.entity_types.is_empty()
558 && filter.name_prefix.is_none()
559 && filter.name_exact.is_none()
560 && filter.tags_any.is_empty()
561 && filter.names_ci.is_empty()
562 && page.offset == 0
563 && usize::try_from(page.limit).ok() == Some(filter.ids.len())
564}
565
566#[async_trait]
571impl EntityStore for SqlEntityStore {
572 async fn upsert_entity(&self, entity: Entity) -> Result<(), StorageError> {
573 let statement = entity_upsert_statement(&entity);
574 self.with_writer("upsert_entity", move |conn| {
575 let mut stmt = conn.prepare(&statement.sql)?;
576 bind_params(&mut stmt, &statement.params)?;
577 stmt.raw_execute()?;
578 Ok(())
579 })
580 .await
581 }
582
583 async fn upsert_entities(
584 &self,
585 entities: Vec<Entity>,
586 ) -> Result<BatchWriteSummary, StorageError> {
587 let attempted = entities.len() as u64;
588
589 if let Some(writer_task) = &self.writer_task {
596 return writer_task
597 .send(move |conn| {
598 batch_upsert_entities(conn, &entities, attempted)
599 .map_err(|e| map_err(e, "upsert_entities"))
600 })
601 .await;
602 }
603
604 let origin = self.pool.origin();
608 self.with_writer("upsert_entities", move |conn| {
609 conn.execute_batch("BEGIN IMMEDIATE")?;
610 let _tx_handle = khive_storage::tx_registry::register_scoped(
611 Some("entity_upsert_batch".to_string()),
612 origin,
613 );
614
615 let summary = batch_upsert_entities(conn, &entities, attempted)?;
616
617 if let Err(e) = conn.execute_batch("COMMIT") {
618 let _ = conn.execute_batch("ROLLBACK");
619 return Err(e);
620 }
621 Ok(summary)
622 })
623 .await
624 }
625
626 async fn get_entity(&self, id: Uuid) -> Result<Option<Entity>, StorageError> {
627 let id_str = id.to_string();
628
629 self.with_reader("get_entity", move |conn| {
630 let mut stmt = conn.prepare(
631 "SELECT id, namespace, kind, entity_type, name, description, properties, tags, \
632 created_at, updated_at, deleted_at, merged_into, merge_event_id, content_ref \
633 FROM entities WHERE id = ?1 AND deleted_at IS NULL",
634 )?;
635 let mut rows = stmt.query(rusqlite::params![id_str])?;
636 match rows.next()? {
637 Some(row) => Ok(Some(read_entity(row)?)),
638 None => Ok(None),
639 }
640 })
641 .await
642 }
643
644 async fn entity_sequence(&self, id: Uuid) -> Result<Option<i64>, StorageError> {
645 let id = id.to_string();
646 self.with_reader("entity_sequence", move |conn| {
647 conn.query_row(
648 "SELECT seq FROM entities_seq WHERE entity_id = ?1",
649 rusqlite::params![id],
650 |row| row.get(0),
651 )
652 .optional()
653 })
654 .await
655 }
656
657 async fn delete_entity(&self, id: Uuid, mode: DeleteMode) -> Result<bool, StorageError> {
658 match mode {
659 DeleteMode::Soft => {
660 let now = chrono::Utc::now().timestamp_micros();
661 let statement = entity_soft_delete_statement(id, now);
662 self.with_writer("delete_entity_soft", move |conn| {
663 let mut stmt = conn.prepare(&statement.sql)?;
664 bind_params(&mut stmt, &statement.params)?;
665 Ok(stmt.raw_execute()? > 0)
666 })
667 .await
668 }
669 DeleteMode::Hard => {
670 let statement = entity_hard_delete_statement(id);
671 self.with_writer("delete_entity_hard", move |conn| {
672 let mut stmt = conn.prepare(&statement.sql)?;
673 bind_params(&mut stmt, &statement.params)?;
674 Ok(stmt.raw_execute()? > 0)
675 })
676 .await
677 }
678 }
679 }
680
681 async fn query_entities(
682 &self,
683 namespace: &str,
684 filter: EntityFilter,
685 page: PageRequest,
686 ) -> Result<Page<Entity>, StorageError> {
687 let namespace = namespace.to_string();
688 let skip_total = is_complete_id_lookup(&filter, &page);
689 let limit_i64 = i64::from(page.limit);
690 let offset_i64 = i64::try_from(page.offset).map_err(|_| StorageError::InvalidInput {
691 capability: StorageCapability::Entities,
692 operation: "query_entities".into(),
693 message: format!(
694 "PageRequest: offset must be <= i64::MAX, got {}",
695 page.offset
696 ),
697 })?;
698
699 self.with_reader("query_entities", move |conn| {
700 let total = if filter.names_ci.is_empty() && !skip_total {
701 let (count_sql, count_params) = build_entity_where(&namespace, &filter);
702 let sql = format!("SELECT COUNT(*) FROM entities{count_sql}");
703 let mut stmt = conn.prepare(&sql)?;
704 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
705 count_params.iter().map(|p| p.as_ref()).collect();
706 Some(stmt.query_row(param_refs.as_slice(), |row| row.get::<_, i64>(0))? as u64)
707 } else {
708 None
709 };
710
711 let mut lookup_filter = filter.clone();
712 lookup_filter.names_ci.clear();
713 let effective_filter = if filter.names_ci.is_empty() {
714 &filter
715 } else {
716 &lookup_filter
717 };
718 let (where_sql, mut data_params) = build_entity_where(&namespace, effective_filter);
719
720 let candidate_param_indices = if filter.names_ci.is_empty() {
721 Vec::new()
722 } else {
723 let mut candidates: Vec<String> = filter
724 .names_ci
725 .iter()
726 .map(|name| name.to_ascii_lowercase())
727 .collect();
728 candidates.sort_unstable();
729 candidates.dedup();
730 candidates
731 .into_iter()
732 .map(|candidate| {
733 data_params.push(Box::new(candidate));
734 data_params.len()
735 })
736 .collect()
737 };
738
739 let order_by = if let Some(ref prefix) = filter.name_prefix {
745 data_params.push(Box::new(prefix.to_ascii_lowercase()));
746 format!(
747 "CASE WHEN LOWER(name) = ?{} THEN 0 ELSE 1 END, created_at DESC",
748 data_params.len()
749 )
750 } else {
751 "created_at DESC".to_string()
752 };
753
754 data_params.push(Box::new(limit_i64));
755 data_params.push(Box::new(offset_i64));
756
757 let limit_idx = data_params.len() - 1;
758 let offset_idx = data_params.len();
759
760 let columns = "id, namespace, kind, entity_type, name, description, properties, tags, \
761 created_at, updated_at, deleted_at, merged_into, merge_event_id, content_ref";
762 let data_sql = if filter.names_ci.is_empty() {
763 format!(
764 "SELECT {columns} FROM entities{where_sql} \
765 ORDER BY {order_by} LIMIT ?{limit_idx} OFFSET ?{offset_idx}"
766 )
767 } else {
768 build_candidate_entity_query(
769 columns,
770 &where_sql,
771 &candidate_param_indices,
772 &order_by,
773 limit_idx,
774 offset_idx,
775 )
776 };
777
778 let mut stmt = conn.prepare(&data_sql)?;
779 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
780 data_params.iter().map(|p| p.as_ref()).collect();
781 let rows = stmt.query_map(param_refs.as_slice(), read_entity)?;
782
783 let mut items = Vec::new();
784 for row in rows {
785 items.push(row?);
786 }
787
788 Ok(Page { items, total })
789 })
790 .await
791 }
792
793 async fn query_entities_after(
794 &self,
795 namespace: &str,
796 filter: EntityFilter,
797 after: Option<SeekCursor>,
798 limit: u32,
799 ) -> Result<SeekPage<Entity>, StorageError> {
800 if limit == 0 {
801 return Ok(SeekPage::default());
802 }
803 if !filter.names_ci.is_empty() {
804 return Err(StorageError::InvalidInput {
805 capability: StorageCapability::Entities,
806 operation: "query_entities_after".into(),
807 message: "names_ci candidate folding is not compatible with seek pagination".into(),
808 });
809 }
810
811 let namespace = namespace.to_string();
812 let limit_usize = limit as usize;
813 let probe_limit_i64 = i64::from(limit) + 1;
814 self.with_reader("query_entities_after", move |conn| {
815 let (mut where_sql, mut params) = build_entity_where(&namespace, &filter);
816 if let Some(cursor) = after {
817 params.push(Box::new(cursor.sequence));
818 where_sql.push_str(&format!(" AND entities_seq.seq > ?{}", params.len()));
819 }
820 params.push(Box::new(probe_limit_i64));
821 let limit_idx = params.len();
822
823 let columns = "id, namespace, kind, entity_type, name, description, properties, tags, \
824 created_at, updated_at, deleted_at, merged_into, merge_event_id, content_ref";
825 let sql = format!(
829 "SELECT {columns}, entities_seq.seq FROM entities_seq \
830 CROSS JOIN entities ON entities.id = entities_seq.entity_id{where_sql} \
831 ORDER BY entities_seq.seq ASC LIMIT ?{limit_idx}"
832 );
833 let mut stmt = conn.prepare(&sql)?;
834 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
835 params.iter().map(|param| param.as_ref()).collect();
836 let rows = stmt.query_map(param_refs.as_slice(), |row| {
837 Ok((read_entity(row)?, row.get::<_, i64>(14)?))
838 })?;
839 let mut entries = rows.collect::<Result<Vec<_>, _>>()?;
840 let has_more = entries.len() > limit_usize;
841 if has_more {
842 entries.truncate(limit_usize);
843 }
844 let next_after = if has_more {
845 entries.last().map(|(entity, sequence)| SeekCursor {
846 sequence: *sequence,
847 id: entity.id,
848 })
849 } else {
850 None
851 };
852 let items = entries.into_iter().map(|(entity, _)| entity).collect();
853 Ok(SeekPage { items, next_after })
854 })
855 .await
856 }
857
858 async fn get_entity_including_deleted(&self, id: Uuid) -> Result<Option<Entity>, StorageError> {
859 let id_str = id.to_string();
860
861 self.with_reader("get_entity_including_deleted", move |conn| {
862 let mut stmt = conn.prepare(
863 "SELECT id, namespace, kind, entity_type, name, description, properties, tags, \
864 created_at, updated_at, deleted_at, merged_into, merge_event_id, content_ref \
865 FROM entities WHERE id = ?1",
866 )?;
867 let mut rows = stmt.query(rusqlite::params![id_str])?;
868 match rows.next()? {
869 Some(row) => Ok(Some(read_entity(row)?)),
870 None => Ok(None),
871 }
872 })
873 .await
874 }
875
876 async fn count_entities(
877 &self,
878 namespace: &str,
879 filter: EntityFilter,
880 ) -> Result<u64, StorageError> {
881 let namespace = namespace.to_string();
882
883 self.with_reader("count_entities", move |conn| {
884 if filter.namespaces.is_empty() {
885 let (where_sql, params) = build_entity_where(&namespace, &filter);
886 let sql = format!("SELECT COUNT(*) FROM entities{}", where_sql);
887 let mut stmt = conn.prepare(&sql)?;
888 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
889 params.iter().map(|p| p.as_ref()).collect();
890 let count: i64 = stmt.query_row(param_refs.as_slice(), |row| row.get(0))?;
891 return Ok(count as u64);
892 }
893
894 let deduped_namespaces: Vec<String> = filter
895 .namespaces
896 .iter()
897 .cloned()
898 .collect::<HashSet<_>>()
899 .into_iter()
900 .collect();
901
902 let mut total = 0;
903 for chunk in deduped_namespaces.chunks(NAMESPACE_COUNT_CHUNK_SIZE) {
904 let chunk_filter = EntityFilter {
905 namespaces: chunk.to_vec(),
906 ..filter.clone()
907 };
908 let (where_sql, params) = build_entity_where(&namespace, &chunk_filter);
909 let sql = format!("SELECT COUNT(*) FROM entities{}", where_sql);
910 let mut stmt = conn.prepare(&sql)?;
911 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
912 params.iter().map(|p| p.as_ref()).collect();
913 let count: i64 = stmt.query_row(param_refs.as_slice(), |row| row.get(0))?;
914 total += count as u64;
915 }
916 Ok(total)
917 })
918 .await
919 }
920}
921
922const ENTITIES_DDL: &str = include_str!("../../sql/entities-ddl.sql");
927
928pub(crate) fn ensure_entities_schema(conn: &rusqlite::Connection) -> Result<(), rusqlite::Error> {
929 conn.execute_batch(ENTITIES_DDL)
930}
931
932#[cfg(test)]
933#[path = "entity_tests.rs"]
934mod tests;