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