1use std::collections::HashSet;
4use std::sync::Arc;
5
6use async_trait::async_trait;
7use rusqlite::OptionalExtension;
8use uuid::Uuid;
9
10use khive_storage::attachment::{Attachment, AttachmentSubstrate};
11use khive_storage::entity::{Entity, EntityFilter, EntityTombstones, EntityTypeCounts};
12use khive_storage::error::StorageError;
13use khive_storage::types::{
14 BatchWriteSummary, DeleteMode, Page, PageRequest, SeekCursor, SeekPage, SqlStatement, SqlValue,
15};
16use khive_storage::EntityStore;
17use khive_storage::StorageCapability;
18
19use crate::error::SqliteError;
20use crate::pool::ConnectionPool;
21use crate::sql_bridge::bind_params;
22use crate::stores::attachment::{attachment_upsert_statement, delete_record_attachments_statement};
23use crate::writer_task::{execute_wrapped_transaction, WriterTaskHandle};
24
25fn map_err(e: rusqlite::Error, op: &'static str) -> StorageError {
26 StorageError::driver(StorageCapability::Entities, op, e)
27}
28
29fn map_sqlite_err(e: SqliteError, op: &'static str) -> StorageError {
30 e.into_storage_error(StorageCapability::Entities, op)
31}
32
33const NAMESPACE_COUNT_CHUNK_SIZE: usize = 500;
34
35const ENTITIES_COUNT_BY_TYPE_SQL: &str = include_str!("../../sql/entities-count-by-type.sql");
36
37const ENTITY_SELECT_COLUMNS: &str =
38 "entities.id, entities.namespace, entities.kind, entities.entity_type, entities.name, \
39 entities.description, entities.properties, entities.tags, entities.created_at, \
40 entities.updated_at, entities.deleted_at, entities.merged_into, entities.merge_event_id, \
41 (SELECT attachment.content_ref FROM attachments AS attachment \
42 WHERE attachment.record_uuid = entities.id \
43 AND attachment.substrate = 'entity' AND attachment.role = 'content') AS content_ref, entities.version";
44
45pub fn entity_upsert_statement(entity: &Entity) -> SqlStatement {
61 let mut statement = entity_write_statement(entity, "INSERT", "entity-upsert");
62 statement.sql.push_str(
63 " ON CONFLICT(id) DO UPDATE SET namespace=excluded.namespace, kind=excluded.kind, \
64 entity_type=excluded.entity_type, name=excluded.name, description=excluded.description, \
65 properties=excluded.properties, tags=excluded.tags, created_at=excluded.created_at, \
66 updated_at=excluded.updated_at, deleted_at=excluded.deleted_at, \
67 merged_into=excluded.merged_into, merge_event_id=excluded.merge_event_id, \
68 version=entities.version+1",
69 );
70 statement
71}
72
73fn entity_write_statement(entity: &Entity, insert: &str, label: &str) -> SqlStatement {
74 let properties_str = entity
75 .properties
76 .as_ref()
77 .map(|v| serde_json::to_string(v).unwrap_or_default());
78 let tags_str = serde_json::to_string(&entity.tags).unwrap_or_else(|_| "[]".to_string());
79 SqlStatement {
80 sql: format!(
81 "{insert} INTO entities \
82 (id, namespace, kind, entity_type, name, description, properties, tags, \
83 created_at, updated_at, deleted_at, merged_into, merge_event_id) \
84 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13)"
85 ),
86 params: vec![
87 SqlValue::Text(entity.id.to_string()),
88 SqlValue::Text(entity.namespace.clone()),
89 SqlValue::Text(entity.kind.clone()),
90 match &entity.entity_type {
91 Some(t) => SqlValue::Text(t.clone()),
92 None => SqlValue::Null,
93 },
94 SqlValue::Text(entity.name.clone()),
95 match &entity.description {
96 Some(d) => SqlValue::Text(d.clone()),
97 None => SqlValue::Null,
98 },
99 match properties_str {
100 Some(p) => SqlValue::Text(p),
101 None => SqlValue::Null,
102 },
103 SqlValue::Text(tags_str),
104 SqlValue::Integer(entity.created_at),
105 SqlValue::Integer(entity.updated_at),
106 match entity.deleted_at {
107 Some(d) => SqlValue::Integer(d),
108 None => SqlValue::Null,
109 },
110 match entity.merged_into {
111 Some(u) => SqlValue::Text(u.to_string()),
112 None => SqlValue::Null,
113 },
114 match entity.merge_event_id {
115 Some(u) => SqlValue::Text(u.to_string()),
116 None => SqlValue::Null,
117 },
118 ],
119 label: Some(label.to_string()),
120 }
121}
122
123pub fn entity_insert_if_absent_statement(entity: &Entity) -> SqlStatement {
127 let mut statement = entity_upsert_statement(entity);
128 statement.sql = "INSERT INTO entities \
129 (id, namespace, kind, entity_type, name, description, properties, tags, \
130 created_at, updated_at, deleted_at, merged_into, merge_event_id) \
131 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13) \
132 ON CONFLICT DO NOTHING"
133 .to_string();
134 statement.label = Some("entity-insert-if-absent".to_string());
135 statement
136}
137
138pub fn entity_replace_if_unchanged_statement(
149 entity: &Entity,
150 expected_updated_at: i64,
151 expected_deleted_at: Option<i64>,
152) -> SqlStatement {
153 let properties_str = entity
154 .properties
155 .as_ref()
156 .map(|v| serde_json::to_string(v).unwrap_or_default());
157 let tags_str = serde_json::to_string(&entity.tags).unwrap_or_else(|_| "[]".to_string());
158 SqlStatement {
159 sql: "UPDATE entities SET \
160 namespace = ?1, kind = ?2, entity_type = ?3, name = ?4, description = ?5, \
161 properties = ?6, tags = ?7, updated_at = ?8, deleted_at = ?9, \
162 merged_into = ?10, merge_event_id = ?11, version = version + 1 \
163 WHERE id = ?12 AND updated_at = ?13 AND deleted_at IS ?14 \
164 AND ?8 > updated_at AND version = ?15"
165 .to_string(),
166 params: vec![
167 SqlValue::Text(entity.namespace.clone()),
168 SqlValue::Text(entity.kind.clone()),
169 match &entity.entity_type {
170 Some(t) => SqlValue::Text(t.clone()),
171 None => SqlValue::Null,
172 },
173 SqlValue::Text(entity.name.clone()),
174 match &entity.description {
175 Some(d) => SqlValue::Text(d.clone()),
176 None => SqlValue::Null,
177 },
178 match properties_str {
179 Some(p) => SqlValue::Text(p),
180 None => SqlValue::Null,
181 },
182 SqlValue::Text(tags_str),
183 SqlValue::Integer(entity.updated_at),
184 match entity.deleted_at {
185 Some(d) => SqlValue::Integer(d),
186 None => SqlValue::Null,
187 },
188 match entity.merged_into {
189 Some(u) => SqlValue::Text(u.to_string()),
190 None => SqlValue::Null,
191 },
192 match entity.merge_event_id {
193 Some(u) => SqlValue::Text(u.to_string()),
194 None => SqlValue::Null,
195 },
196 SqlValue::Text(entity.id.to_string()),
197 SqlValue::Integer(expected_updated_at),
198 match expected_deleted_at {
199 Some(value) => SqlValue::Integer(value),
200 None => SqlValue::Null,
201 },
202 SqlValue::Integer(entity.version),
203 ],
204 label: Some("entity-replace-if-unchanged".to_string()),
205 }
206}
207
208pub fn entity_soft_delete_statement(id: Uuid, deleted_at: i64) -> SqlStatement {
210 SqlStatement {
211 sql: "UPDATE entities SET deleted_at = ?1, version = version + 1 WHERE id = ?2 AND deleted_at IS NULL".to_string(),
212 params: vec![
213 SqlValue::Integer(deleted_at),
214 SqlValue::Text(id.to_string()),
215 ],
216 label: Some("entity-delete-soft".to_string()),
217 }
218}
219
220pub fn entity_hard_delete_statement(id: Uuid) -> SqlStatement {
223 SqlStatement {
224 sql: "DELETE FROM entities WHERE id = ?1".to_string(),
225 params: vec![SqlValue::Text(id.to_string())],
226 label: Some("entity-delete-hard".to_string()),
227 }
228}
229
230pub struct SqlEntityStore {
236 pool: Arc<ConnectionPool>,
237 writer_task: Option<WriterTaskHandle>,
238}
239
240#[derive(Clone, Copy, PartialEq, Eq)]
241enum EntityPageMode {
242 ExactTotal,
243 CountFree,
244}
245
246impl EntityPageMode {
247 fn operation(self) -> &'static str {
248 match self {
249 Self::ExactTotal => "query_entities",
250 Self::CountFree => "query_entities_count_free",
251 }
252 }
253}
254
255impl SqlEntityStore {
256 pub fn new(pool: Arc<ConnectionPool>, _is_file_backed: bool) -> Self {
274 let writer_task = pool.writer_task_handle().ok().flatten();
282
283 Self { pool, writer_task }
284 }
285
286 fn current_writer_task(
287 &self,
288 operation: &'static str,
289 ) -> Result<Option<WriterTaskHandle>, StorageError> {
290 self.pool
291 .writer_task_for_write(self.writer_task.as_ref(), operation)
292 }
293
294 async fn with_writer<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
311 where
312 F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
313 R: Send + 'static,
314 {
315 if let Some(writer_task) = self.current_writer_task(op)? {
316 return writer_task
317 .send_bounded(move |conn| f(conn).map_err(|e| map_err(e, op)))
318 .await;
319 }
320
321 self.pool
322 .record_direct_route(crate::timeout_sink::Site::DirectRouteEntity);
323 let pool = Arc::clone(&self.pool);
324 tokio::task::spawn_blocking(move || {
325 let guard = pool
326 .autocommit_write_unit()
327 .map_err(|e| map_sqlite_err(e, op))?;
328 f(guard.conn())
329 .map_err(|e| map_err(e, op))
330 .inspect_err(|error| pool.record_direct_writer_error(error))
331 })
332 .await
333 .map_err(|e| StorageError::driver(StorageCapability::Entities, op, e))?
334 }
335
336 async fn with_writer_tx<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
340 where
341 F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
342 R: Send + 'static,
343 {
344 if let Some(writer_task) = self.current_writer_task(op)? {
345 return writer_task
346 .send_bounded(move |conn| f(conn).map_err(|error| map_err(error, op)))
347 .await;
348 }
349
350 self.pool
351 .record_direct_route(crate::timeout_sink::Site::DirectRouteEntity);
352 let pool = Arc::clone(&self.pool);
353 tokio::task::spawn_blocking(move || {
354 let guard = pool
355 .transaction_write_unit()
356 .map_err(|error| map_sqlite_err(error, op))
357 .inspect_err(|error| pool.record_direct_writer_error(error))?;
358 let conn = guard.conn();
359 let (result, terminal_state) = execute_wrapped_transaction(conn, op, move |conn| {
360 f(conn).map_err(|error| map_err(error, op))
361 });
362 if terminal_state.is_some() {
363 pool.retire_pooled_writer(conn);
364 }
365 result.inspect_err(|error| pool.record_direct_writer_error(error))
366 })
367 .await
368 .map_err(|error| StorageError::driver(StorageCapability::Entities, op, error))?
369 }
370
371 async fn with_reader<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
372 where
373 F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
374 R: Send + 'static,
375 {
376 super::run_pooled_store_read(
377 Arc::clone(&self.pool),
378 StorageCapability::Entities,
379 op,
380 move |conn| f(conn).map_err(|error| map_err(error, op)),
381 )
382 .await
383 }
384
385 async fn with_list_reader<F, R>(
386 &self,
387 op: &'static str,
388 index: Option<&'static str>,
389 mut read: F,
390 ) -> Result<R, StorageError>
391 where
392 F: FnMut(&rusqlite::Connection, bool) -> Result<R, rusqlite::Error> + Send + 'static,
393 R: Send + 'static,
394 {
395 let dispatch = tracing::dispatcher::get_default(Clone::clone);
396 self.with_reader(op, move |conn| match read(conn, true) {
397 Ok(value) => Ok(value),
398 Err(error) => {
399 let Some(index) = index.filter(|index| missing_list_index(&error, index)) else {
400 return Err(error);
401 };
402 tracing::dispatcher::with_default(&dispatch, || {
403 tracing::warn!(
404 index,
405 operation = op,
406 "entity list index missing; retrying without forced index"
407 );
408 });
409 read(conn, false)
410 }
411 })
412 .await
413 }
414
415 async fn query_entities_page(
416 &self,
417 namespace: &str,
418 filter: EntityFilter,
419 page: PageRequest,
420 mode: EntityPageMode,
421 ) -> Result<Page<Entity>, StorageError> {
422 let operation = mode.operation();
423 super::validate_json_equality_paths(
424 &filter.property_equalities,
425 StorageCapability::Entities,
426 operation,
427 )?;
428 let namespace = namespace.to_string();
429 let skip_total = mode == EntityPageMode::CountFree || is_complete_id_lookup(&filter, &page);
430 let limit_i64 = i64::from(page.limit);
431 let offset_i64 = i64::try_from(page.offset).map_err(|_| StorageError::InvalidInput {
432 capability: StorageCapability::Entities,
433 operation: operation.into(),
434 message: format!(
435 "PageRequest: offset must be <= i64::MAX, got {}",
436 page.offset
437 ),
438 })?;
439
440 let streaming = mode == EntityPageMode::CountFree && is_entity_streaming_page(&filter);
441 let index = if streaming {
442 entity_count_free_index(&filter)
443 } else {
444 entity_list_index(&filter)
445 };
446 let read = move |conn: &rusqlite::Connection, force_index| {
447 let total = if filter.names_ci.is_empty() && !skip_total {
448 let (count_sql, count_params) = build_entity_where(&namespace, &filter);
449 let sql = entity_list_query(
450 &filter,
451 build_entity_count_query(&filter, &count_sql),
452 force_index,
453 );
454 let mut stmt = conn.prepare(&sql)?;
455 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
456 count_params.iter().map(|p| p.as_ref()).collect();
457 Some(stmt.query_row(param_refs.as_slice(), |row| row.get::<_, i64>(0))? as u64)
458 } else {
459 None
460 };
461
462 let mut lookup_filter = filter.clone();
463 lookup_filter.names_ci.clear();
464 let effective_filter = if filter.names_ci.is_empty() {
465 &filter
466 } else {
467 &lookup_filter
468 };
469 let (where_sql, mut data_params) = if streaming {
470 build_entity_streaming_where(&namespace, effective_filter)
471 } else {
472 build_entity_where(&namespace, effective_filter)
473 };
474
475 let candidate_param_indices = if filter.names_ci.is_empty() {
476 Vec::new()
477 } else {
478 let mut candidates: Vec<String> = filter
479 .names_ci
480 .iter()
481 .map(|name| name.to_ascii_lowercase())
482 .collect();
483 candidates.sort_unstable();
484 candidates.dedup();
485 candidates
486 .into_iter()
487 .map(|candidate| {
488 data_params.push(Box::new(candidate));
489 data_params.len()
490 })
491 .collect()
492 };
493
494 let order_by = if let Some(ref prefix) = filter.name_prefix {
500 data_params.push(Box::new(prefix.to_ascii_lowercase()));
501 format!(
502 "CASE WHEN LOWER(name) = ?{} THEN 0 ELSE 1 END, created_at DESC, id DESC",
503 data_params.len()
504 )
505 } else {
506 "created_at DESC, id DESC".to_string()
514 };
515
516 data_params.push(Box::new(limit_i64));
517 data_params.push(Box::new(offset_i64));
518
519 let limit_idx = data_params.len() - 1;
520 let offset_idx = data_params.len();
521
522 let columns = ENTITY_SELECT_COLUMNS;
523 let data_sql = if streaming {
524 build_entity_count_free_page_query(
525 columns,
526 effective_filter,
527 &where_sql,
528 &order_by,
529 limit_idx,
530 offset_idx,
531 )
532 } else if filter.names_ci.is_empty() {
533 build_entity_page_query(
534 columns,
535 effective_filter,
536 &where_sql,
537 &order_by,
538 limit_idx,
539 offset_idx,
540 )
541 } else {
542 build_candidate_entity_query(
543 columns,
544 effective_filter,
545 &where_sql,
546 &candidate_param_indices,
547 &order_by,
548 limit_idx,
549 offset_idx,
550 )
551 };
552
553 let data_sql = if streaming {
554 entity_count_free_query(effective_filter, data_sql, force_index)
555 } else {
556 entity_list_query(effective_filter, data_sql, force_index)
557 };
558 let mut stmt = conn.prepare(&data_sql)?;
559 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
560 data_params.iter().map(|p| p.as_ref()).collect();
561 let rows = stmt.query_map(param_refs.as_slice(), read_entity)?;
562
563 let mut items = Vec::new();
564 for row in rows {
565 items.push(row?);
566 }
567
568 Ok(Page { items, total })
569 };
570 self.with_list_reader(operation, index, read).await
571 }
572}
573
574fn read_entity(row: &rusqlite::Row<'_>) -> Result<Entity, rusqlite::Error> {
579 let id_str: String = row.get(0)?;
580 let namespace: String = row.get(1)?;
581 let kind: String = row.get(2)?;
582 let entity_type: Option<String> = row.get(3)?;
583 let name: String = row.get(4)?;
584 let description: Option<String> = row.get(5)?;
585 let properties_str: Option<String> = row.get(6)?;
586 let tags_str: String = row.get(7)?;
587 let created_at: i64 = row.get(8)?;
588 let updated_at: i64 = row.get(9)?;
589 let deleted_at: Option<i64> = row.get(10)?;
590 let merged_into_str: Option<String> = row.get(11)?;
591 let merge_event_id_str: Option<String> = row.get(12)?;
592 let content_ref: Option<String> = row.get(13)?;
593 let version: i64 = row.get(14)?;
594
595 let id = parse_uuid(&id_str)?;
596
597 let properties = properties_str
598 .map(|s| {
599 serde_json::from_str(&s).map_err(|e| {
600 rusqlite::Error::FromSqlConversionFailure(
601 6,
602 rusqlite::types::Type::Text,
603 Box::new(e),
604 )
605 })
606 })
607 .transpose()?;
608
609 let tags: Vec<String> = serde_json::from_str(&tags_str).map_err(|e| {
610 rusqlite::Error::FromSqlConversionFailure(7, rusqlite::types::Type::Text, Box::new(e))
611 })?;
612
613 let merged_into = merged_into_str
614 .as_deref()
615 .map(Uuid::parse_str)
616 .transpose()
617 .map_err(|e| {
618 rusqlite::Error::FromSqlConversionFailure(10, rusqlite::types::Type::Text, Box::new(e))
619 })?;
620
621 let merge_event_id = merge_event_id_str
622 .as_deref()
623 .map(Uuid::parse_str)
624 .transpose()
625 .map_err(|e| {
626 rusqlite::Error::FromSqlConversionFailure(11, rusqlite::types::Type::Text, Box::new(e))
627 })?;
628
629 Ok(Entity {
630 id,
631 namespace,
632 kind,
633 entity_type,
634 name,
635 description,
636 properties,
637 tags,
638 created_at,
639 updated_at,
640 version,
641 deleted_at,
642 merged_into,
643 merge_event_id,
644 content_ref,
645 })
646}
647
648fn batch_upsert_entities(
658 conn: &rusqlite::Connection,
659 entities: &[Entity],
660 attempted: u64,
661) -> Result<BatchWriteSummary, rusqlite::Error> {
662 let mut summary = BatchWriteSummary {
663 attempted,
664 ..BatchWriteSummary::default()
665 };
666
667 for (index, entity) in entities.iter().enumerate() {
668 let id_str = entity.id.to_string();
669 let statement = entity_upsert_statement(entity);
670 let result = (|| {
671 let mut prepared = conn.prepare(&statement.sql)?;
672 bind_params(&mut prepared, &statement.params)?;
673 prepared.raw_execute()
674 })();
675 match result {
676 Ok(_) => summary.affected = summary.affected.saturating_add(1),
677 Err(e) => {
678 let (class, retryability) = super::classify_batch_sqlite_error(&e);
679 summary.record_failure(index, Some(id_str), class, retryability, e.to_string());
680 }
681 }
682 }
683
684 Ok(summary)
685}
686
687fn parse_uuid(s: &str) -> Result<Uuid, rusqlite::Error> {
688 Uuid::parse_str(s).map_err(|e| {
689 rusqlite::Error::FromSqlConversionFailure(0, rusqlite::types::Type::Text, Box::new(e))
690 })
691}
692
693fn build_entity_where(
694 namespace: &str,
695 filter: &EntityFilter,
696) -> (String, Vec<Box<dyn rusqlite::types::ToSql>>) {
697 build_entity_where_with_mode(namespace, filter, false)
698}
699
700fn build_entity_streaming_where(
704 namespace: &str,
705 filter: &EntityFilter,
706) -> (String, Vec<Box<dyn rusqlite::types::ToSql>>) {
707 let mut filter = filter.clone();
708 for values in [
709 &mut filter.namespaces,
710 &mut filter.kinds,
711 &mut filter.entity_types,
712 ] {
713 values.sort_unstable();
714 values.dedup();
715 }
716 for values in filter.entity_types_by_kind.values_mut() {
717 values.sort_unstable();
718 values.dedup();
719 }
720 build_entity_where_with_mode(namespace, &filter, true)
721}
722
723fn build_entity_where_with_mode(
724 namespace: &str,
725 filter: &EntityFilter,
726 row_local: bool,
727) -> (String, Vec<Box<dyn rusqlite::types::ToSql>>) {
728 let (ns_condition, ns_params): (String, Vec<Box<dyn rusqlite::types::ToSql>>) =
732 if !filter.namespaces.is_empty() {
733 let placeholders: Vec<String> = (1..=filter.namespaces.len())
734 .map(|i| format!("?{i}"))
735 .collect();
736 let params: Vec<Box<dyn rusqlite::types::ToSql>> = filter
737 .namespaces
738 .iter()
739 .map(|ns| -> Box<dyn rusqlite::types::ToSql> { Box::new(ns.clone()) })
740 .collect();
741 (
742 format!("namespace IN ({})", placeholders.join(", ")),
743 params,
744 )
745 } else {
746 (
747 "namespace = ?1".to_string(),
748 vec![Box::new(namespace.to_string())],
749 )
750 };
751
752 let mut conditions: Vec<String> = vec![ns_condition];
753 match filter.tombstones {
754 EntityTombstones::Live => conditions.push("deleted_at IS NULL".into()),
755 EntityTombstones::All => {}
756 EntityTombstones::Only => conditions.push("deleted_at IS NOT NULL".into()),
757 }
758 let mut params: Vec<Box<dyn rusqlite::types::ToSql>> = ns_params;
759
760 if !filter.ids.is_empty() {
761 let placeholders: Vec<String> = filter
762 .ids
763 .iter()
764 .map(|id| {
765 params.push(Box::new(id.to_string()));
766 format!("?{}", params.len())
767 })
768 .collect();
769 conditions.push(format!("id IN ({})", placeholders.join(", ")));
770 }
771
772 if !filter.kinds.is_empty() {
773 let placeholders: Vec<String> = filter
774 .kinds
775 .iter()
776 .map(|k| {
777 params.push(Box::new(k.clone()));
778 format!("?{}", params.len())
779 })
780 .collect();
781 conditions.push(format!("kind IN ({})", placeholders.join(", ")));
782 }
783
784 let type_scope = conditions.join(" AND ");
785 let type_predicate = |scope: &str, placeholders: &str| {
786 if filter.legacy_entity_type_fallback && row_local {
787 format!(
790 "CASE WHEN entity_type IS NOT NULL THEN entity_type IN ({placeholders}) \
791 WHEN json_valid(properties) THEN \
792 CASE WHEN json_type(properties, '$.type') = 'text' \
793 THEN json_extract(properties, '$.type') IN ({placeholders}) \
794 ELSE 0 END \
795 ELSE 0 END"
796 )
797 } else if filter.legacy_entity_type_fallback {
798 format!(
803 "id IN (SELECT id FROM entities WHERE {scope} \
804 AND entity_type IN ({placeholders}) \
805 UNION ALL SELECT id FROM entities WHERE {scope} \
806 AND entity_type IS NULL AND json_valid(properties) \
807 AND json_type(properties, '$.type') = 'text' \
808 AND json_extract(properties, '$.type') IN ({placeholders}))"
809 )
810 } else {
811 format!("entity_type IN ({placeholders})")
812 }
813 };
814 if !filter.entity_types.is_empty() {
815 let placeholders: Vec<String> = filter
816 .entity_types
817 .iter()
818 .map(|t| {
819 params.push(Box::new(t.clone()));
820 format!("?{}", params.len())
821 })
822 .collect();
823 conditions.push(type_predicate(&type_scope, &placeholders.join(", ")));
824 }
825
826 if !filter.entity_types_by_kind.is_empty() {
827 let mut groups = Vec::new();
828 for (kind, types) in &filter.entity_types_by_kind {
829 if types.is_empty() {
830 continue;
831 }
832 params.push(Box::new(kind.clone()));
833 let kind_param = params.len();
834 let placeholders = types
835 .iter()
836 .map(|value| {
837 params.push(Box::new(value.clone()));
838 format!("?{}", params.len())
839 })
840 .collect::<Vec<_>>()
841 .join(", ");
842 let scope = format!("{type_scope} AND kind = ?{kind_param}");
843 let predicate = type_predicate(&scope, &placeholders);
844 groups.push(format!("(kind = ?{kind_param} AND {predicate})"));
845 }
846 conditions.push(if groups.is_empty() {
847 "0".to_string()
848 } else {
849 format!("({})", groups.join(" OR "))
850 });
851 }
852
853 if let Some(ref prefix) = filter.name_prefix {
854 params.push(Box::new(format!(
855 "{}%",
856 khive_types::escape_like_literal(prefix)
857 )));
858 conditions.push(format!("name LIKE ?{} ESCAPE '\\'", params.len()));
859 }
860
861 if let Some(ref exact) = filter.name_exact {
862 params.push(Box::new(exact.clone()));
863 conditions.push(format!("name = ?{} COLLATE BINARY", params.len()));
868 }
869
870 if !filter.names_ci.is_empty() {
871 let placeholders: Vec<String> = filter
874 .names_ci
875 .iter()
876 .map(|n| {
877 params.push(Box::new(n.to_ascii_lowercase()));
878 format!("?{}", params.len())
879 })
880 .collect();
881 conditions.push(format!("LOWER(name) IN ({})", placeholders.join(", ")));
882 }
883
884 if !filter.tags_any.is_empty() {
885 let placeholders: Vec<String> = filter
886 .tags_any
887 .iter()
888 .map(|t| {
889 params.push(Box::new(t.to_lowercase()));
892 format!("?{}", params.len())
893 })
894 .collect();
895 conditions.push(format!(
896 "EXISTS (SELECT 1 FROM json_each(tags) WHERE LOWER(json_each.value) IN ({}))",
897 placeholders.join(", ")
898 ));
899 }
900
901 super::append_json_equalities(
902 &mut conditions,
903 &mut params,
904 "properties",
905 &filter.property_equalities,
906 );
907 let clause = format!(" WHERE {}", conditions.join(" AND "));
908 (clause, params)
909}
910
911fn entity_read_source(filter: &EntityFilter) -> &'static str {
916 if filter.ids.is_empty() {
917 "entities"
918 } else {
919 "entities INDEXED BY sqlite_autoindex_entities_1"
920 }
921}
922
923fn build_entity_count_query(filter: &EntityFilter, where_sql: &str) -> String {
924 let source = entity_list_source(filter);
925 format!("SELECT COUNT(*) FROM {source}{where_sql}")
926}
927
928fn build_entity_page_query(
929 columns: &str,
930 filter: &EntityFilter,
931 where_sql: &str,
932 order_by: &str,
933 limit_idx: usize,
934 offset_idx: usize,
935) -> String {
936 let source = entity_list_source(filter);
937 build_entity_page_query_from_source(columns, source, where_sql, order_by, limit_idx, offset_idx)
938}
939
940fn build_entity_count_free_page_query(
941 columns: &str,
942 filter: &EntityFilter,
943 where_sql: &str,
944 order_by: &str,
945 limit_idx: usize,
946 offset_idx: usize,
947) -> String {
948 let source = entity_count_free_source(filter);
949 build_entity_page_query_from_source(columns, source, where_sql, order_by, limit_idx, offset_idx)
950}
951
952fn build_entity_page_query_from_source(
953 columns: &str,
954 source: &str,
955 where_sql: &str,
956 order_by: &str,
957 limit_idx: usize,
958 offset_idx: usize,
959) -> String {
960 format!(
961 "SELECT {columns} FROM {source}{where_sql} \
962 ORDER BY {order_by} LIMIT ?{limit_idx} OFFSET ?{offset_idx}"
963 )
964}
965
966fn entity_list_source(filter: &EntityFilter) -> &'static str {
970 if filter.tombstones != EntityTombstones::Live
971 || !filter.property_equalities.is_empty()
972 || !filter.ids.is_empty()
973 || !filter.kinds.is_empty()
974 || !filter.entity_types_by_kind.is_empty()
975 || filter.legacy_entity_type_fallback
976 || filter.name_prefix.is_some()
977 || filter.name_exact.is_some()
978 || !filter.names_ci.is_empty()
979 || !filter.tags_any.is_empty()
980 {
981 return entity_read_source(filter);
982 }
983 if !filter.entity_types.is_empty() {
984 "entities INDEXED BY idx_entities_live_namespace_type_order"
985 } else if filter.namespaces.len() <= 1 {
986 "entities INDEXED BY idx_entities_live_namespace_order"
987 } else {
988 "entities"
989 }
990}
991
992fn entity_list_index(filter: &EntityFilter) -> Option<&'static str> {
993 if !filter.ids.is_empty() {
994 return None;
995 }
996 entity_list_source(filter).strip_prefix("entities INDEXED BY ")
997}
998
999fn is_entity_streaming_page(filter: &EntityFilter) -> bool {
1000 filter.ids.is_empty()
1001 && filter.name_prefix.is_none()
1002 && filter.name_exact.is_none()
1003 && filter.names_ci.is_empty()
1004}
1005
1006fn has_one_distinct_value(values: &[String]) -> bool {
1007 values
1008 .first()
1009 .is_some_and(|first| values.iter().all(|value| value == first))
1010}
1011
1012fn entity_count_free_source(filter: &EntityFilter) -> &'static str {
1016 if filter.tombstones != EntityTombstones::Live {
1017 return entity_read_source(filter);
1018 }
1019 if !is_entity_streaming_page(filter) {
1020 return entity_list_source(filter);
1021 }
1022 if !filter.namespaces.is_empty() && !has_one_distinct_value(&filter.namespaces) {
1023 return "entities";
1024 }
1025 if !filter.entity_types_by_kind.is_empty()
1030 && filter.entity_types_by_kind.values().all(Vec::is_empty)
1031 {
1032 return "entities";
1033 }
1034
1035 let mut kind_groups = filter
1036 .entity_types_by_kind
1037 .iter()
1038 .filter(|(_, types)| !types.is_empty());
1039 let one_kind_group = kind_groups.next().is_some() && kind_groups.next().is_none();
1040 if has_one_distinct_value(&filter.kinds) || one_kind_group {
1041 "entities INDEXED BY idx_entities_live_namespace_kind_order"
1042 } else if !filter.legacy_entity_type_fallback && has_one_distinct_value(&filter.entity_types) {
1043 "entities INDEXED BY idx_entities_live_namespace_type_order"
1046 } else {
1047 "entities INDEXED BY idx_entities_live_namespace_order"
1048 }
1049}
1050
1051fn entity_count_free_index(filter: &EntityFilter) -> Option<&'static str> {
1052 if !filter.ids.is_empty() {
1053 return None;
1054 }
1055 entity_count_free_source(filter).strip_prefix("entities INDEXED BY ")
1056}
1057
1058fn missing_list_index(error: &rusqlite::Error, index: &str) -> bool {
1059 if error
1060 .sqlite_error()
1061 .is_none_or(|detail| detail.extended_code != rusqlite::ffi::SQLITE_ERROR)
1062 {
1063 return false;
1064 }
1065 let message = match error {
1066 rusqlite::Error::SqliteFailure(_, Some(message)) => message.as_str(),
1067 rusqlite::Error::SqlInputError { msg, .. } => msg.as_str(),
1068 _ => return false,
1069 };
1070 message.strip_prefix("no such index: ") == Some(index)
1071}
1072
1073fn entity_list_query(filter: &EntityFilter, sql: String, force_index: bool) -> String {
1074 entity_query_with_index(
1075 sql,
1076 entity_list_source(filter),
1077 entity_list_index(filter),
1078 force_index,
1079 )
1080}
1081
1082fn entity_count_free_query(filter: &EntityFilter, sql: String, force_index: bool) -> String {
1083 entity_query_with_index(
1084 sql,
1085 entity_count_free_source(filter),
1086 entity_count_free_index(filter),
1087 force_index,
1088 )
1089}
1090
1091fn entity_query_with_index(
1092 sql: String,
1093 source: &str,
1094 index: Option<&str>,
1095 force_index: bool,
1096) -> String {
1097 if force_index || index.is_none() {
1098 sql
1099 } else {
1100 sql.replacen(source, "entities", 1)
1101 }
1102}
1103
1104fn build_candidate_entity_query(
1105 columns: &str,
1106 filter: &EntityFilter,
1107 where_sql: &str,
1108 candidate_param_indices: &[usize],
1109 order_by: &str,
1110 limit_idx: usize,
1111 offset_idx: usize,
1112) -> String {
1113 let source = entity_read_source(filter);
1114 let candidate_rows = candidate_param_indices
1115 .iter()
1116 .map(|idx| format!("(?{idx})"))
1117 .collect::<Vec<_>>()
1118 .join(", ");
1119
1120 format!(
1121 "WITH candidates(folded_name) AS (VALUES {candidate_rows}), \
1122 matched_entities(entity_id) AS (\
1123 SELECT (\
1124 SELECT id FROM {source}{where_sql} \
1125 AND LOWER(name) = candidates.folded_name LIMIT 1\
1126 ) FROM candidates\
1127 ) \
1128 SELECT {columns} FROM entities \
1129 JOIN matched_entities ON entities.id = matched_entities.entity_id \
1130 ORDER BY {order_by} LIMIT ?{limit_idx} OFFSET ?{offset_idx}"
1131 )
1132}
1133
1134fn build_entity_cursor_query(
1135 columns: &str,
1136 filter: &EntityFilter,
1137 where_sql: &str,
1138 limit_idx: usize,
1139) -> String {
1140 let source = entity_read_source(filter);
1145 let from_clause = if !filter.ids.is_empty() {
1146 format!("{source} CROSS JOIN entities_seq ON entities.id = entities_seq.entity_id")
1147 } else {
1148 "entities_seq CROSS JOIN entities INDEXED BY sqlite_autoindex_entities_1 \
1149 ON entities.id = entities_seq.entity_id"
1150 .to_string()
1151 };
1152 format!(
1153 "SELECT {columns}, entities_seq.seq FROM {from_clause}{where_sql} \
1154 ORDER BY entities_seq.seq ASC LIMIT ?{limit_idx}"
1155 )
1156}
1157
1158fn is_complete_id_lookup(filter: &EntityFilter, page: &PageRequest) -> bool {
1159 !filter.ids.is_empty()
1160 && filter.property_equalities.is_empty()
1161 && filter.kinds.is_empty()
1162 && filter.entity_types.is_empty()
1163 && filter.entity_types_by_kind.is_empty()
1164 && filter.name_prefix.is_none()
1165 && filter.name_exact.is_none()
1166 && filter.tags_any.is_empty()
1167 && filter.names_ci.is_empty()
1168 && page.offset == 0
1169 && usize::try_from(page.limit).ok() == Some(filter.ids.len())
1170}
1171
1172#[async_trait]
1177impl EntityStore for SqlEntityStore {
1178 async fn upsert_entity(&self, entity: Entity) -> Result<(), StorageError> {
1179 let statement = entity_upsert_statement(&entity);
1180 self.with_writer("upsert_entity", move |conn| {
1181 let mut stmt = conn.prepare(&statement.sql)?;
1182 bind_params(&mut stmt, &statement.params)?;
1183 stmt.raw_execute()?;
1184 Ok(())
1185 })
1186 .await
1187 }
1188
1189 async fn insert_entity_if_absent(&self, entity: Entity) -> Result<bool, StorageError> {
1190 let statement = entity_insert_if_absent_statement(&entity);
1191 self.with_writer("insert_entity_if_absent", move |conn| {
1192 let mut stmt = conn.prepare(&statement.sql)?;
1193 bind_params(&mut stmt, &statement.params)?;
1194 Ok(stmt.raw_execute()? > 0)
1195 })
1196 .await
1197 }
1198
1199 async fn upsert_entity_with_attachments(
1200 &self,
1201 entity: Entity,
1202 attachments: Vec<Attachment>,
1203 ) -> Result<(), StorageError> {
1204 let entity_id = entity.id;
1205 let entity_statement = entity_upsert_statement(&entity);
1206 let mut attachment_statements = Vec::with_capacity(attachments.len());
1207 for attachment in attachments {
1208 attachment.validate()?;
1209 if attachment.record_uuid != entity_id
1210 || attachment.substrate != AttachmentSubstrate::Entity
1211 {
1212 return Err(StorageError::InvalidInput {
1213 capability: StorageCapability::Attachments,
1214 operation: "upsert_entity_with_attachments".into(),
1215 message: format!(
1216 "attachment {} must target entity {entity_id}",
1217 attachment.role
1218 ),
1219 });
1220 }
1221 attachment_statements.push(attachment_upsert_statement(&attachment)?);
1222 }
1223
1224 self.with_writer_tx("upsert_entity_with_attachments", move |conn| {
1225 let mut entity_stmt = conn.prepare(&entity_statement.sql)?;
1226 bind_params(&mut entity_stmt, &entity_statement.params)?;
1227 entity_stmt.raw_execute()?;
1228 drop(entity_stmt);
1229
1230 for statement in attachment_statements {
1231 let mut stmt = conn.prepare(&statement.sql)?;
1232 bind_params(&mut stmt, &statement.params)?;
1233 stmt.raw_execute()?;
1234 }
1235 Ok(())
1236 })
1237 .await
1238 }
1239
1240 async fn upsert_entities(
1241 &self,
1242 entities: Vec<Entity>,
1243 ) -> Result<BatchWriteSummary, StorageError> {
1244 let attempted = entities.len() as u64;
1245
1246 if let Some(writer_task) = self.current_writer_task("upsert_entities")? {
1253 return writer_task
1254 .send_bounded(move |conn| {
1255 batch_upsert_entities(conn, &entities, attempted)
1256 .map_err(|e| map_err(e, "upsert_entities"))
1257 })
1258 .await;
1259 }
1260
1261 let origin = self.pool.origin();
1265 self.with_writer("upsert_entities", move |conn| {
1266 conn.execute_batch("BEGIN IMMEDIATE")?;
1267 let _tx_handle = khive_storage::tx_registry::register_scoped(
1268 Some("entity_upsert_batch".to_string()),
1269 origin,
1270 );
1271
1272 let summary = batch_upsert_entities(conn, &entities, attempted)?;
1273
1274 if let Err(e) = conn.execute_batch("COMMIT") {
1275 let _ = conn.execute_batch("ROLLBACK");
1276 return Err(e);
1277 }
1278 Ok(summary)
1279 })
1280 .await
1281 }
1282
1283 async fn replace_entity_if_unchanged(
1284 &self,
1285 entity: Entity,
1286 expected_updated_at: i64,
1287 expected_deleted_at: Option<i64>,
1288 ) -> Result<bool, StorageError> {
1289 let statement = entity_replace_if_unchanged_statement(
1290 &entity,
1291 expected_updated_at,
1292 expected_deleted_at,
1293 );
1294 self.with_writer("replace_entity_if_unchanged", move |conn| {
1295 let mut stmt = conn.prepare(&statement.sql)?;
1296 bind_params(&mut stmt, &statement.params)?;
1297 Ok(stmt.raw_execute()? > 0)
1298 })
1299 .await
1300 }
1301
1302 async fn get_entity(&self, id: Uuid) -> Result<Option<Entity>, StorageError> {
1303 let id_str = id.to_string();
1304
1305 self.with_reader("get_entity", move |conn| {
1306 let sql = format!(
1307 "SELECT {ENTITY_SELECT_COLUMNS} FROM entities \
1308 WHERE entities.id = ?1 AND entities.deleted_at IS NULL"
1309 );
1310 let mut stmt = conn.prepare(&sql)?;
1311 let mut rows = stmt.query(rusqlite::params![id_str])?;
1312 match rows.next()? {
1313 Some(row) => Ok(Some(read_entity(row)?)),
1314 None => Ok(None),
1315 }
1316 })
1317 .await
1318 }
1319
1320 async fn entity_sequence(&self, id: Uuid) -> Result<Option<i64>, StorageError> {
1321 let id = id.to_string();
1322 self.with_reader("entity_sequence", move |conn| {
1323 conn.query_row(
1324 "SELECT seq FROM entities_seq WHERE entity_id = ?1",
1325 rusqlite::params![id],
1326 |row| row.get(0),
1327 )
1328 .optional()
1329 })
1330 .await
1331 }
1332
1333 async fn delete_entity(&self, id: Uuid, mode: DeleteMode) -> Result<bool, StorageError> {
1334 match mode {
1335 DeleteMode::Soft => {
1336 let now = chrono::Utc::now().timestamp_micros();
1337 let statement = entity_soft_delete_statement(id, now);
1338 self.with_writer("delete_entity_soft", move |conn| {
1339 let mut stmt = conn.prepare(&statement.sql)?;
1340 bind_params(&mut stmt, &statement.params)?;
1341 Ok(stmt.raw_execute()? > 0)
1342 })
1343 .await
1344 }
1345 DeleteMode::Hard => {
1346 let entity_statement = entity_hard_delete_statement(id);
1347 let attachment_statement =
1348 delete_record_attachments_statement(id, AttachmentSubstrate::Entity);
1349 self.with_writer_tx("delete_entity_hard", move |conn| {
1350 let mut entity_stmt = conn.prepare(&entity_statement.sql)?;
1351 bind_params(&mut entity_stmt, &entity_statement.params)?;
1352 let deleted = entity_stmt.raw_execute()? > 0;
1353 drop(entity_stmt);
1354 if deleted {
1355 let mut attachment_stmt = conn.prepare(&attachment_statement.sql)?;
1356 bind_params(&mut attachment_stmt, &attachment_statement.params)?;
1357 attachment_stmt.raw_execute()?;
1358 }
1359 Ok(deleted)
1360 })
1361 .await
1362 }
1363 }
1364 }
1365
1366 async fn query_entities(
1367 &self,
1368 namespace: &str,
1369 filter: EntityFilter,
1370 page: PageRequest,
1371 ) -> Result<Page<Entity>, StorageError> {
1372 self.query_entities_page(namespace, filter, page, EntityPageMode::ExactTotal)
1373 .await
1374 }
1375
1376 async fn query_entities_count_free(
1377 &self,
1378 namespace: &str,
1379 filter: EntityFilter,
1380 page: PageRequest,
1381 ) -> Result<Page<Entity>, StorageError> {
1382 self.query_entities_page(namespace, filter, page, EntityPageMode::CountFree)
1383 .await
1384 }
1385
1386 async fn query_entities_after(
1387 &self,
1388 namespace: &str,
1389 filter: EntityFilter,
1390 after: Option<SeekCursor>,
1391 limit: u32,
1392 ) -> Result<SeekPage<Entity>, StorageError> {
1393 super::validate_json_equality_paths(
1394 &filter.property_equalities,
1395 StorageCapability::Entities,
1396 "query_entities_after",
1397 )?;
1398 if limit == 0 {
1399 return Ok(SeekPage::default());
1400 }
1401 if !filter.names_ci.is_empty() {
1402 return Err(StorageError::InvalidInput {
1403 capability: StorageCapability::Entities,
1404 operation: "query_entities_after".into(),
1405 message: "names_ci candidate folding is not compatible with seek pagination".into(),
1406 });
1407 }
1408
1409 let namespace = namespace.to_string();
1410 let limit_usize = limit as usize;
1411 let probe_limit_i64 = i64::from(limit) + 1;
1412 self.with_reader("query_entities_after", move |conn| {
1413 let (mut where_sql, mut params) = if filter.ids.is_empty() {
1414 build_entity_streaming_where(&namespace, &filter)
1415 } else {
1416 build_entity_where(&namespace, &filter)
1417 };
1418 if let Some(cursor) = after {
1419 params.push(Box::new(cursor.sequence));
1420 where_sql.push_str(&format!(" AND entities_seq.seq > ?{}", params.len()));
1421 }
1422 params.push(Box::new(probe_limit_i64));
1423 let limit_idx = params.len();
1424
1425 let columns = ENTITY_SELECT_COLUMNS;
1426 let sql = build_entity_cursor_query(columns, &filter, &where_sql, limit_idx);
1427 let mut stmt = conn.prepare(&sql)?;
1428 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
1429 params.iter().map(|param| param.as_ref()).collect();
1430 let rows = stmt.query_map(param_refs.as_slice(), |row| {
1431 Ok((read_entity(row)?, row.get::<_, i64>(15)?))
1432 })?;
1433 let mut entries = rows.collect::<Result<Vec<_>, _>>()?;
1434 let has_more = entries.len() > limit_usize;
1435 if has_more {
1436 entries.truncate(limit_usize);
1437 }
1438 let next_after = if has_more {
1439 entries.last().map(|(entity, sequence)| SeekCursor {
1440 sequence: *sequence,
1441 id: entity.id,
1442 })
1443 } else {
1444 None
1445 };
1446 let items = entries.into_iter().map(|(entity, _)| entity).collect();
1447 Ok(SeekPage { items, next_after })
1448 })
1449 .await
1450 }
1451
1452 async fn get_entity_including_deleted(&self, id: Uuid) -> Result<Option<Entity>, StorageError> {
1453 let id_str = id.to_string();
1454
1455 self.with_reader("get_entity_including_deleted", move |conn| {
1456 let sql =
1457 format!("SELECT {ENTITY_SELECT_COLUMNS} FROM entities WHERE entities.id = ?1");
1458 let mut stmt = conn.prepare(&sql)?;
1459 let mut rows = stmt.query(rusqlite::params![id_str])?;
1460 match rows.next()? {
1461 Some(row) => Ok(Some(read_entity(row)?)),
1462 None => Ok(None),
1463 }
1464 })
1465 .await
1466 }
1467
1468 async fn count_entities_by_type(
1469 &self,
1470 namespaces: &[String],
1471 ) -> Result<Option<EntityTypeCounts>, StorageError> {
1472 let namespaces =
1473 serde_json::to_string(namespaces).map_err(|error| StorageError::Serialization {
1474 capability: StorageCapability::Entities,
1475 message: error.to_string(),
1476 })?;
1477 self.with_reader("count_entities_by_type", move |conn| {
1478 let mut statement = conn.prepare(ENTITIES_COUNT_BY_TYPE_SQL)?;
1479 let rows = statement.query_map(rusqlite::params![namespaces], |row| {
1480 let count: i64 = row.get(1)?;
1481 let count = u64::try_from(count).map_err(|error| {
1482 rusqlite::Error::FromSqlConversionFailure(
1483 1,
1484 rusqlite::types::Type::Integer,
1485 Box::new(error),
1486 )
1487 })?;
1488 Ok((row.get::<_, Option<String>>(0)?, count))
1489 })?;
1490 rows.collect::<Result<Vec<_>, _>>().map(Some)
1491 })
1492 .await
1493 }
1494
1495 async fn count_entities(
1496 &self,
1497 namespace: &str,
1498 filter: EntityFilter,
1499 ) -> Result<u64, StorageError> {
1500 super::validate_json_equality_paths(
1501 &filter.property_equalities,
1502 StorageCapability::Entities,
1503 "count_entities",
1504 )?;
1505 let namespace = namespace.to_string();
1506
1507 let index = entity_list_index(&filter);
1508 self.with_list_reader("count_entities", index, move |conn, force_index| {
1509 if filter.namespaces.is_empty() {
1510 let (where_sql, params) = build_entity_where(&namespace, &filter);
1511 let sql = entity_list_query(
1512 &filter,
1513 build_entity_count_query(&filter, &where_sql),
1514 force_index,
1515 );
1516 let mut stmt = conn.prepare(&sql)?;
1517 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
1518 params.iter().map(|p| p.as_ref()).collect();
1519 let count: i64 = stmt.query_row(param_refs.as_slice(), |row| row.get(0))?;
1520 return Ok(count as u64);
1521 }
1522
1523 let deduped_namespaces: Vec<String> = filter
1524 .namespaces
1525 .iter()
1526 .cloned()
1527 .collect::<HashSet<_>>()
1528 .into_iter()
1529 .collect();
1530
1531 let mut total = 0;
1532 for chunk in deduped_namespaces.chunks(NAMESPACE_COUNT_CHUNK_SIZE) {
1533 let chunk_filter = EntityFilter {
1534 namespaces: chunk.to_vec(),
1535 ..filter.clone()
1536 };
1537 let (where_sql, params) = build_entity_where(&namespace, &chunk_filter);
1538 let sql = entity_list_query(
1539 &chunk_filter,
1540 build_entity_count_query(&chunk_filter, &where_sql),
1541 force_index && index.is_some(),
1542 );
1543 let mut stmt = conn.prepare(&sql)?;
1544 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
1545 params.iter().map(|p| p.as_ref()).collect();
1546 let count: i64 = stmt.query_row(param_refs.as_slice(), |row| row.get(0))?;
1547 total += count as u64;
1548 }
1549 Ok(total)
1550 })
1551 .await
1552 }
1553}
1554
1555const ENTITIES_DDL: &str = include_str!("../../sql/entities-ddl.sql");
1560
1561pub(crate) fn ensure_entities_schema(conn: &rusqlite::Connection) -> Result<(), rusqlite::Error> {
1562 conn.execute_batch(ENTITIES_DDL)
1563}
1564
1565#[cfg(test)]
1566#[path = "entity_tests.rs"]
1567mod tests;
1568
1569#[cfg(test)]
1570#[path = "entity_type_counts_tests.rs"]
1571mod entity_type_counts_tests;
1572
1573#[cfg(test)]
1574#[path = "entity_busy_tests.rs"]
1575mod direct_busy_tests;
1576
1577#[cfg(test)]
1578#[path = "entity_list_plan_tests.rs"]
1579mod entity_list_plan_tests;