use std::collections::HashSet;
use std::sync::Arc;
use async_trait::async_trait;
use rusqlite::OptionalExtension;
use uuid::Uuid;
use khive_storage::attachment::{Attachment, AttachmentSubstrate};
use khive_storage::entity::{Entity, EntityFilter};
use khive_storage::error::{StorageError, WriterTaskRequestState};
use khive_storage::types::{
BatchWriteSummary, DeleteMode, Page, PageRequest, SeekCursor, SeekPage, SqlStatement, SqlValue,
};
use khive_storage::EntityStore;
use khive_storage::StorageCapability;
use crate::error::SqliteError;
use crate::pool::ConnectionPool;
use crate::sql_bridge::bind_params;
use crate::stores::attachment::{attachment_upsert_statement, delete_record_attachments_statement};
use crate::writer_task::{execute_wrapped_transaction, WriterTaskHandle};
fn map_err(e: rusqlite::Error, op: &'static str) -> StorageError {
StorageError::driver(StorageCapability::Entities, op, e)
}
fn map_sqlite_err(e: SqliteError, op: &'static str) -> StorageError {
StorageError::driver(StorageCapability::Entities, op, e)
}
const NAMESPACE_COUNT_CHUNK_SIZE: usize = 500;
const ENTITY_SELECT_COLUMNS: &str =
"entities.id, entities.namespace, entities.kind, entities.entity_type, entities.name, \
entities.description, entities.properties, entities.tags, entities.created_at, \
entities.updated_at, entities.deleted_at, entities.merged_into, entities.merge_event_id, \
(SELECT attachment.content_ref FROM attachments AS attachment \
WHERE attachment.record_uuid = entities.id \
AND attachment.substrate = 'entity' AND attachment.role = 'content') AS content_ref, entities.version";
pub fn entity_upsert_statement(entity: &Entity) -> SqlStatement {
let mut statement = entity_write_statement(entity, "INSERT", "entity-upsert");
statement.sql.push_str(
" ON CONFLICT(id) DO UPDATE SET namespace=excluded.namespace, kind=excluded.kind, \
entity_type=excluded.entity_type, name=excluded.name, description=excluded.description, \
properties=excluded.properties, tags=excluded.tags, created_at=excluded.created_at, \
updated_at=excluded.updated_at, deleted_at=excluded.deleted_at, \
merged_into=excluded.merged_into, merge_event_id=excluded.merge_event_id, \
version=entities.version+1",
);
statement
}
pub fn entity_insert_statement(entity: &Entity) -> SqlStatement {
entity_write_statement(entity, "INSERT", "entity-insert")
}
fn entity_write_statement(entity: &Entity, insert: &str, label: &str) -> SqlStatement {
let properties_str = entity
.properties
.as_ref()
.map(|v| serde_json::to_string(v).unwrap_or_default());
let tags_str = serde_json::to_string(&entity.tags).unwrap_or_else(|_| "[]".to_string());
SqlStatement {
sql: format!(
"{insert} INTO entities \
(id, namespace, kind, entity_type, name, description, properties, tags, \
created_at, updated_at, deleted_at, merged_into, merge_event_id) \
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13)"
),
params: vec![
SqlValue::Text(entity.id.to_string()),
SqlValue::Text(entity.namespace.clone()),
SqlValue::Text(entity.kind.clone()),
match &entity.entity_type {
Some(t) => SqlValue::Text(t.clone()),
None => SqlValue::Null,
},
SqlValue::Text(entity.name.clone()),
match &entity.description {
Some(d) => SqlValue::Text(d.clone()),
None => SqlValue::Null,
},
match properties_str {
Some(p) => SqlValue::Text(p),
None => SqlValue::Null,
},
SqlValue::Text(tags_str),
SqlValue::Integer(entity.created_at),
SqlValue::Integer(entity.updated_at),
match entity.deleted_at {
Some(d) => SqlValue::Integer(d),
None => SqlValue::Null,
},
match entity.merged_into {
Some(u) => SqlValue::Text(u.to_string()),
None => SqlValue::Null,
},
match entity.merge_event_id {
Some(u) => SqlValue::Text(u.to_string()),
None => SqlValue::Null,
},
],
label: Some(label.to_string()),
}
}
pub fn entity_insert_if_absent_statement(entity: &Entity) -> SqlStatement {
let mut statement = entity_upsert_statement(entity);
statement.sql = "INSERT INTO entities \
(id, namespace, kind, entity_type, name, description, properties, tags, \
created_at, updated_at, deleted_at, merged_into, merge_event_id) \
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13) \
ON CONFLICT DO NOTHING"
.to_string();
statement.label = Some("entity-insert-if-absent".to_string());
statement
}
pub fn entity_replace_if_unchanged_statement(
entity: &Entity,
expected_updated_at: i64,
expected_deleted_at: Option<i64>,
) -> SqlStatement {
let properties_str = entity
.properties
.as_ref()
.map(|v| serde_json::to_string(v).unwrap_or_default());
let tags_str = serde_json::to_string(&entity.tags).unwrap_or_else(|_| "[]".to_string());
SqlStatement {
sql: "UPDATE entities SET \
namespace = ?1, kind = ?2, entity_type = ?3, name = ?4, description = ?5, \
properties = ?6, tags = ?7, updated_at = ?8, deleted_at = ?9, \
merged_into = ?10, merge_event_id = ?11, version = version + 1 \
WHERE id = ?12 AND updated_at = ?13 AND deleted_at IS ?14 \
AND ?8 > updated_at AND version = ?15"
.to_string(),
params: vec![
SqlValue::Text(entity.namespace.clone()),
SqlValue::Text(entity.kind.clone()),
match &entity.entity_type {
Some(t) => SqlValue::Text(t.clone()),
None => SqlValue::Null,
},
SqlValue::Text(entity.name.clone()),
match &entity.description {
Some(d) => SqlValue::Text(d.clone()),
None => SqlValue::Null,
},
match properties_str {
Some(p) => SqlValue::Text(p),
None => SqlValue::Null,
},
SqlValue::Text(tags_str),
SqlValue::Integer(entity.updated_at),
match entity.deleted_at {
Some(d) => SqlValue::Integer(d),
None => SqlValue::Null,
},
match entity.merged_into {
Some(u) => SqlValue::Text(u.to_string()),
None => SqlValue::Null,
},
match entity.merge_event_id {
Some(u) => SqlValue::Text(u.to_string()),
None => SqlValue::Null,
},
SqlValue::Text(entity.id.to_string()),
SqlValue::Integer(expected_updated_at),
match expected_deleted_at {
Some(value) => SqlValue::Integer(value),
None => SqlValue::Null,
},
SqlValue::Integer(entity.version),
],
label: Some("entity-replace-if-unchanged".to_string()),
}
}
pub fn entity_soft_delete_statement(id: Uuid, deleted_at: i64) -> SqlStatement {
SqlStatement {
sql: "UPDATE entities SET deleted_at = ?1, version = version + 1 WHERE id = ?2 AND deleted_at IS NULL".to_string(),
params: vec![
SqlValue::Integer(deleted_at),
SqlValue::Text(id.to_string()),
],
label: Some("entity-delete-soft".to_string()),
}
}
pub fn entity_hard_delete_statement(id: Uuid) -> SqlStatement {
SqlStatement {
sql: "DELETE FROM entities WHERE id = ?1".to_string(),
params: vec![SqlValue::Text(id.to_string())],
label: Some("entity-delete-hard".to_string()),
}
}
pub struct SqlEntityStore {
pool: Arc<ConnectionPool>,
writer_task: Option<WriterTaskHandle>,
}
impl SqlEntityStore {
pub fn new(pool: Arc<ConnectionPool>, _is_file_backed: bool) -> Self {
let writer_task = pool.writer_task_handle().ok().flatten();
Self { pool, writer_task }
}
fn current_writer_task(
&self,
operation: &'static str,
) -> Result<Option<WriterTaskHandle>, StorageError> {
self.pool
.writer_task_for_write(self.writer_task.as_ref(), operation)
}
async fn with_writer<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
where
F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
R: Send + 'static,
{
if let Some(writer_task) = self.current_writer_task(op)? {
return writer_task
.send_bounded(move |conn| f(conn).map_err(|e| map_err(e, op)))
.await;
}
self.pool
.record_direct_route(crate::timeout_sink::Site::DirectRouteEntity);
let pool = Arc::clone(&self.pool);
tokio::task::spawn_blocking(move || {
let guard = pool.try_writer().map_err(|e| map_sqlite_err(e, op))?;
f(guard.conn()).map_err(|e| map_err(e, op))
})
.await
.map_err(|e| StorageError::driver(StorageCapability::Entities, op, e))?
}
async fn with_writer_tx<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
where
F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
R: Send + 'static,
{
if let Some(writer_task) = self.current_writer_task(op)? {
return writer_task
.send_bounded(move |conn| f(conn).map_err(|error| map_err(error, op)))
.await;
}
self.pool
.record_direct_route(crate::timeout_sink::Site::DirectRouteEntity);
let pool = Arc::clone(&self.pool);
tokio::task::spawn_blocking(move || {
let guard = pool
.try_writer()
.map_err(|error| map_sqlite_err(error, op))?;
let conn = guard.conn();
if !conn.is_autocommit() {
pool.retire_pooled_writer(conn);
return Err(StorageError::WriterTaskTerminated {
request_state: WriterTaskRequestState::SideEffectsUnknown,
});
}
if let Err(begin_error) = conn.execute_batch("BEGIN IMMEDIATE") {
if !conn.is_autocommit() {
pool.retire_pooled_writer(conn);
return Err(StorageError::WriterTaskTerminated {
request_state: WriterTaskRequestState::SideEffectsUnknown,
});
}
return Err(map_err(begin_error, op));
}
let (result, terminal_state) = execute_wrapped_transaction(conn, op, move |conn| {
f(conn).map_err(|error| map_err(error, op))
});
if terminal_state.is_some() {
pool.retire_pooled_writer(conn);
}
result
})
.await
.map_err(|error| StorageError::driver(StorageCapability::Entities, op, error))?
}
async fn with_reader<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
where
F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
R: Send + 'static,
{
super::run_pooled_store_read(
Arc::clone(&self.pool),
StorageCapability::Entities,
op,
move |conn| f(conn).map_err(|error| map_err(error, op)),
)
.await
}
}
fn read_entity(row: &rusqlite::Row<'_>) -> Result<Entity, rusqlite::Error> {
let id_str: String = row.get(0)?;
let namespace: String = row.get(1)?;
let kind: String = row.get(2)?;
let entity_type: Option<String> = row.get(3)?;
let name: String = row.get(4)?;
let description: Option<String> = row.get(5)?;
let properties_str: Option<String> = row.get(6)?;
let tags_str: String = row.get(7)?;
let created_at: i64 = row.get(8)?;
let updated_at: i64 = row.get(9)?;
let deleted_at: Option<i64> = row.get(10)?;
let merged_into_str: Option<String> = row.get(11)?;
let merge_event_id_str: Option<String> = row.get(12)?;
let content_ref: Option<String> = row.get(13)?;
let version: i64 = row.get(14)?;
let id = parse_uuid(&id_str)?;
let properties = properties_str
.map(|s| {
serde_json::from_str(&s).map_err(|e| {
rusqlite::Error::FromSqlConversionFailure(
6,
rusqlite::types::Type::Text,
Box::new(e),
)
})
})
.transpose()?;
let tags: Vec<String> = serde_json::from_str(&tags_str).map_err(|e| {
rusqlite::Error::FromSqlConversionFailure(7, rusqlite::types::Type::Text, Box::new(e))
})?;
let merged_into = merged_into_str
.as_deref()
.map(Uuid::parse_str)
.transpose()
.map_err(|e| {
rusqlite::Error::FromSqlConversionFailure(10, rusqlite::types::Type::Text, Box::new(e))
})?;
let merge_event_id = merge_event_id_str
.as_deref()
.map(Uuid::parse_str)
.transpose()
.map_err(|e| {
rusqlite::Error::FromSqlConversionFailure(11, rusqlite::types::Type::Text, Box::new(e))
})?;
Ok(Entity {
id,
namespace,
kind,
entity_type,
name,
description,
properties,
tags,
created_at,
updated_at,
version,
deleted_at,
merged_into,
merge_event_id,
content_ref,
})
}
fn batch_upsert_entities(
conn: &rusqlite::Connection,
entities: &[Entity],
attempted: u64,
) -> Result<BatchWriteSummary, rusqlite::Error> {
let mut summary = BatchWriteSummary {
attempted,
..BatchWriteSummary::default()
};
for (index, entity) in entities.iter().enumerate() {
let id_str = entity.id.to_string();
let statement = entity_upsert_statement(entity);
let result = (|| {
let mut prepared = conn.prepare(&statement.sql)?;
bind_params(&mut prepared, &statement.params)?;
prepared.raw_execute()
})();
match result {
Ok(_) => summary.affected = summary.affected.saturating_add(1),
Err(e) => {
let (class, retryability) = super::classify_batch_sqlite_error(&e);
summary.record_failure(index, Some(id_str), class, retryability, e.to_string());
}
}
}
Ok(summary)
}
fn parse_uuid(s: &str) -> Result<Uuid, rusqlite::Error> {
Uuid::parse_str(s).map_err(|e| {
rusqlite::Error::FromSqlConversionFailure(0, rusqlite::types::Type::Text, Box::new(e))
})
}
fn escape_like(input: &str) -> String {
let mut out = String::with_capacity(input.len());
for c in input.chars() {
if matches!(c, '\\' | '%' | '_') {
out.push('\\');
}
out.push(c);
}
out
}
fn build_entity_where(
namespace: &str,
filter: &EntityFilter,
) -> (String, Vec<Box<dyn rusqlite::types::ToSql>>) {
let (ns_condition, ns_params): (String, Vec<Box<dyn rusqlite::types::ToSql>>) =
if !filter.namespaces.is_empty() {
let placeholders: Vec<String> = (1..=filter.namespaces.len())
.map(|i| format!("?{i}"))
.collect();
let params: Vec<Box<dyn rusqlite::types::ToSql>> = filter
.namespaces
.iter()
.map(|ns| -> Box<dyn rusqlite::types::ToSql> { Box::new(ns.clone()) })
.collect();
(
format!("namespace IN ({})", placeholders.join(", ")),
params,
)
} else {
(
"namespace = ?1".to_string(),
vec![Box::new(namespace.to_string())],
)
};
let mut conditions: Vec<String> = vec![ns_condition, "deleted_at IS NULL".to_string()];
let mut params: Vec<Box<dyn rusqlite::types::ToSql>> = ns_params;
if !filter.ids.is_empty() {
let placeholders: Vec<String> = filter
.ids
.iter()
.map(|id| {
params.push(Box::new(id.to_string()));
format!("?{}", params.len())
})
.collect();
conditions.push(format!("id IN ({})", placeholders.join(", ")));
}
if !filter.kinds.is_empty() {
let placeholders: Vec<String> = filter
.kinds
.iter()
.map(|k| {
params.push(Box::new(k.clone()));
format!("?{}", params.len())
})
.collect();
conditions.push(format!("kind IN ({})", placeholders.join(", ")));
}
let type_scope = conditions.join(" AND ");
let type_predicate = |scope: &str, placeholders: &str| {
if filter.legacy_entity_type_fallback {
format!(
"id IN (SELECT id FROM entities WHERE {scope} \
AND entity_type IN ({placeholders}) \
UNION ALL SELECT id FROM entities WHERE {scope} \
AND entity_type IS NULL AND json_valid(properties) \
AND json_type(properties, '$.type') = 'text' \
AND json_extract(properties, '$.type') IN ({placeholders}))"
)
} else {
format!("entity_type IN ({placeholders})")
}
};
if !filter.entity_types.is_empty() {
let placeholders: Vec<String> = filter
.entity_types
.iter()
.map(|t| {
params.push(Box::new(t.clone()));
format!("?{}", params.len())
})
.collect();
conditions.push(type_predicate(&type_scope, &placeholders.join(", ")));
}
if !filter.entity_types_by_kind.is_empty() {
let mut groups = Vec::new();
for (kind, types) in &filter.entity_types_by_kind {
if types.is_empty() {
continue;
}
params.push(Box::new(kind.clone()));
let kind_param = params.len();
let placeholders = types
.iter()
.map(|value| {
params.push(Box::new(value.clone()));
format!("?{}", params.len())
})
.collect::<Vec<_>>()
.join(", ");
let scope = format!("{type_scope} AND kind = ?{kind_param}");
let predicate = type_predicate(&scope, &placeholders);
groups.push(format!("(kind = ?{kind_param} AND {predicate})"));
}
conditions.push(if groups.is_empty() {
"0".to_string()
} else {
format!("({})", groups.join(" OR "))
});
}
if let Some(ref prefix) = filter.name_prefix {
params.push(Box::new(format!("{}%", escape_like(prefix))));
conditions.push(format!("name LIKE ?{} ESCAPE '\\'", params.len()));
}
if let Some(ref exact) = filter.name_exact {
params.push(Box::new(exact.clone()));
conditions.push(format!("name = ?{} COLLATE BINARY", params.len()));
}
if !filter.names_ci.is_empty() {
let placeholders: Vec<String> = filter
.names_ci
.iter()
.map(|n| {
params.push(Box::new(n.to_ascii_lowercase()));
format!("?{}", params.len())
})
.collect();
conditions.push(format!("LOWER(name) IN ({})", placeholders.join(", ")));
}
if !filter.tags_any.is_empty() {
let placeholders: Vec<String> = filter
.tags_any
.iter()
.map(|t| {
params.push(Box::new(t.to_lowercase()));
format!("?{}", params.len())
})
.collect();
conditions.push(format!(
"EXISTS (SELECT 1 FROM json_each(tags) WHERE LOWER(json_each.value) IN ({}))",
placeholders.join(", ")
));
}
let clause = format!(" WHERE {}", conditions.join(" AND "));
(clause, params)
}
fn build_candidate_entity_query(
columns: &str,
where_sql: &str,
candidate_param_indices: &[usize],
order_by: &str,
limit_idx: usize,
offset_idx: usize,
) -> String {
let candidate_rows = candidate_param_indices
.iter()
.map(|idx| format!("(?{idx})"))
.collect::<Vec<_>>()
.join(", ");
format!(
"WITH candidates(folded_name) AS (VALUES {candidate_rows}), \
matched_entities(entity_id) AS (\
SELECT (\
SELECT id FROM entities{where_sql} \
AND LOWER(name) = candidates.folded_name LIMIT 1\
) FROM candidates\
) \
SELECT {columns} FROM entities \
JOIN matched_entities ON entities.id = matched_entities.entity_id \
ORDER BY {order_by} LIMIT ?{limit_idx} OFFSET ?{offset_idx}"
)
}
fn is_complete_id_lookup(filter: &EntityFilter, page: &PageRequest) -> bool {
!filter.ids.is_empty()
&& filter.kinds.is_empty()
&& filter.entity_types.is_empty()
&& filter.entity_types_by_kind.is_empty()
&& filter.name_prefix.is_none()
&& filter.name_exact.is_none()
&& filter.tags_any.is_empty()
&& filter.names_ci.is_empty()
&& page.offset == 0
&& usize::try_from(page.limit).ok() == Some(filter.ids.len())
}
#[async_trait]
impl EntityStore for SqlEntityStore {
async fn upsert_entity(&self, entity: Entity) -> Result<(), StorageError> {
let statement = entity_upsert_statement(&entity);
self.with_writer("upsert_entity", move |conn| {
let mut stmt = conn.prepare(&statement.sql)?;
bind_params(&mut stmt, &statement.params)?;
stmt.raw_execute()?;
Ok(())
})
.await
}
async fn insert_entity_if_absent(&self, entity: Entity) -> Result<bool, StorageError> {
let statement = entity_insert_if_absent_statement(&entity);
self.with_writer("insert_entity_if_absent", move |conn| {
let mut stmt = conn.prepare(&statement.sql)?;
bind_params(&mut stmt, &statement.params)?;
Ok(stmt.raw_execute()? > 0)
})
.await
}
async fn upsert_entity_with_attachments(
&self,
entity: Entity,
attachments: Vec<Attachment>,
) -> Result<(), StorageError> {
let entity_id = entity.id;
let entity_statement = entity_upsert_statement(&entity);
let mut attachment_statements = Vec::with_capacity(attachments.len());
for attachment in attachments {
attachment.validate()?;
if attachment.record_uuid != entity_id
|| attachment.substrate != AttachmentSubstrate::Entity
{
return Err(StorageError::InvalidInput {
capability: StorageCapability::Attachments,
operation: "upsert_entity_with_attachments".into(),
message: format!(
"attachment {} must target entity {entity_id}",
attachment.role
),
});
}
attachment_statements.push(attachment_upsert_statement(&attachment)?);
}
self.with_writer_tx("upsert_entity_with_attachments", move |conn| {
let mut entity_stmt = conn.prepare(&entity_statement.sql)?;
bind_params(&mut entity_stmt, &entity_statement.params)?;
entity_stmt.raw_execute()?;
drop(entity_stmt);
for statement in attachment_statements {
let mut stmt = conn.prepare(&statement.sql)?;
bind_params(&mut stmt, &statement.params)?;
stmt.raw_execute()?;
}
Ok(())
})
.await
}
async fn upsert_entities(
&self,
entities: Vec<Entity>,
) -> Result<BatchWriteSummary, StorageError> {
let attempted = entities.len() as u64;
if let Some(writer_task) = self.current_writer_task("upsert_entities")? {
return writer_task
.send_bounded(move |conn| {
batch_upsert_entities(conn, &entities, attempted)
.map_err(|e| map_err(e, "upsert_entities"))
})
.await;
}
let origin = self.pool.origin();
self.with_writer("upsert_entities", move |conn| {
conn.execute_batch("BEGIN IMMEDIATE")?;
let _tx_handle = khive_storage::tx_registry::register_scoped(
Some("entity_upsert_batch".to_string()),
origin,
);
let summary = batch_upsert_entities(conn, &entities, attempted)?;
if let Err(e) = conn.execute_batch("COMMIT") {
let _ = conn.execute_batch("ROLLBACK");
return Err(e);
}
Ok(summary)
})
.await
}
async fn replace_entity_if_unchanged(
&self,
entity: Entity,
expected_updated_at: i64,
expected_deleted_at: Option<i64>,
) -> Result<bool, StorageError> {
let statement = entity_replace_if_unchanged_statement(
&entity,
expected_updated_at,
expected_deleted_at,
);
self.with_writer("replace_entity_if_unchanged", move |conn| {
let mut stmt = conn.prepare(&statement.sql)?;
bind_params(&mut stmt, &statement.params)?;
Ok(stmt.raw_execute()? > 0)
})
.await
}
async fn get_entity(&self, id: Uuid) -> Result<Option<Entity>, StorageError> {
let id_str = id.to_string();
self.with_reader("get_entity", move |conn| {
let sql = format!(
"SELECT {ENTITY_SELECT_COLUMNS} FROM entities \
WHERE entities.id = ?1 AND entities.deleted_at IS NULL"
);
let mut stmt = conn.prepare(&sql)?;
let mut rows = stmt.query(rusqlite::params![id_str])?;
match rows.next()? {
Some(row) => Ok(Some(read_entity(row)?)),
None => Ok(None),
}
})
.await
}
async fn entity_sequence(&self, id: Uuid) -> Result<Option<i64>, StorageError> {
let id = id.to_string();
self.with_reader("entity_sequence", move |conn| {
conn.query_row(
"SELECT seq FROM entities_seq WHERE entity_id = ?1",
rusqlite::params![id],
|row| row.get(0),
)
.optional()
})
.await
}
async fn delete_entity(&self, id: Uuid, mode: DeleteMode) -> Result<bool, StorageError> {
match mode {
DeleteMode::Soft => {
let now = chrono::Utc::now().timestamp_micros();
let statement = entity_soft_delete_statement(id, now);
self.with_writer("delete_entity_soft", move |conn| {
let mut stmt = conn.prepare(&statement.sql)?;
bind_params(&mut stmt, &statement.params)?;
Ok(stmt.raw_execute()? > 0)
})
.await
}
DeleteMode::Hard => {
let entity_statement = entity_hard_delete_statement(id);
let attachment_statement =
delete_record_attachments_statement(id, AttachmentSubstrate::Entity);
self.with_writer_tx("delete_entity_hard", move |conn| {
let mut entity_stmt = conn.prepare(&entity_statement.sql)?;
bind_params(&mut entity_stmt, &entity_statement.params)?;
let deleted = entity_stmt.raw_execute()? > 0;
drop(entity_stmt);
if deleted {
let mut attachment_stmt = conn.prepare(&attachment_statement.sql)?;
bind_params(&mut attachment_stmt, &attachment_statement.params)?;
attachment_stmt.raw_execute()?;
}
Ok(deleted)
})
.await
}
}
}
async fn query_entities(
&self,
namespace: &str,
filter: EntityFilter,
page: PageRequest,
) -> Result<Page<Entity>, StorageError> {
let namespace = namespace.to_string();
let skip_total = is_complete_id_lookup(&filter, &page);
let limit_i64 = i64::from(page.limit);
let offset_i64 = i64::try_from(page.offset).map_err(|_| StorageError::InvalidInput {
capability: StorageCapability::Entities,
operation: "query_entities".into(),
message: format!(
"PageRequest: offset must be <= i64::MAX, got {}",
page.offset
),
})?;
self.with_reader("query_entities", move |conn| {
let total = if filter.names_ci.is_empty() && !skip_total {
let (count_sql, count_params) = build_entity_where(&namespace, &filter);
let sql = format!("SELECT COUNT(*) FROM entities{count_sql}");
let mut stmt = conn.prepare(&sql)?;
let param_refs: Vec<&dyn rusqlite::types::ToSql> =
count_params.iter().map(|p| p.as_ref()).collect();
Some(stmt.query_row(param_refs.as_slice(), |row| row.get::<_, i64>(0))? as u64)
} else {
None
};
let mut lookup_filter = filter.clone();
lookup_filter.names_ci.clear();
let effective_filter = if filter.names_ci.is_empty() {
&filter
} else {
&lookup_filter
};
let (where_sql, mut data_params) = build_entity_where(&namespace, effective_filter);
let candidate_param_indices = if filter.names_ci.is_empty() {
Vec::new()
} else {
let mut candidates: Vec<String> = filter
.names_ci
.iter()
.map(|name| name.to_ascii_lowercase())
.collect();
candidates.sort_unstable();
candidates.dedup();
candidates
.into_iter()
.map(|candidate| {
data_params.push(Box::new(candidate));
data_params.len()
})
.collect()
};
let order_by = if let Some(ref prefix) = filter.name_prefix {
data_params.push(Box::new(prefix.to_ascii_lowercase()));
format!(
"CASE WHEN LOWER(name) = ?{} THEN 0 ELSE 1 END, created_at DESC, id DESC",
data_params.len()
)
} else {
"created_at DESC, id DESC".to_string()
};
data_params.push(Box::new(limit_i64));
data_params.push(Box::new(offset_i64));
let limit_idx = data_params.len() - 1;
let offset_idx = data_params.len();
let columns = ENTITY_SELECT_COLUMNS;
let data_sql = if filter.names_ci.is_empty() {
format!(
"SELECT {columns} FROM entities{where_sql} \
ORDER BY {order_by} LIMIT ?{limit_idx} OFFSET ?{offset_idx}"
)
} else {
build_candidate_entity_query(
columns,
&where_sql,
&candidate_param_indices,
&order_by,
limit_idx,
offset_idx,
)
};
let mut stmt = conn.prepare(&data_sql)?;
let param_refs: Vec<&dyn rusqlite::types::ToSql> =
data_params.iter().map(|p| p.as_ref()).collect();
let rows = stmt.query_map(param_refs.as_slice(), read_entity)?;
let mut items = Vec::new();
for row in rows {
items.push(row?);
}
Ok(Page { items, total })
})
.await
}
async fn query_entities_after(
&self,
namespace: &str,
filter: EntityFilter,
after: Option<SeekCursor>,
limit: u32,
) -> Result<SeekPage<Entity>, StorageError> {
if limit == 0 {
return Ok(SeekPage::default());
}
if !filter.names_ci.is_empty() {
return Err(StorageError::InvalidInput {
capability: StorageCapability::Entities,
operation: "query_entities_after".into(),
message: "names_ci candidate folding is not compatible with seek pagination".into(),
});
}
let namespace = namespace.to_string();
let limit_usize = limit as usize;
let probe_limit_i64 = i64::from(limit) + 1;
self.with_reader("query_entities_after", move |conn| {
let (mut where_sql, mut params) = build_entity_where(&namespace, &filter);
if let Some(cursor) = after {
params.push(Box::new(cursor.sequence));
where_sql.push_str(&format!(" AND entities_seq.seq > ?{}", params.len()));
}
params.push(Box::new(probe_limit_i64));
let limit_idx = params.len();
let columns = ENTITY_SELECT_COLUMNS;
let join_kind = if filter.kinds.is_empty() {
"CROSS JOIN"
} else {
"JOIN"
};
let sql = format!(
"SELECT {columns}, entities_seq.seq FROM entities_seq \
{join_kind} entities ON entities.id = entities_seq.entity_id{where_sql} \
ORDER BY entities_seq.seq ASC LIMIT ?{limit_idx}"
);
let mut stmt = conn.prepare(&sql)?;
let param_refs: Vec<&dyn rusqlite::types::ToSql> =
params.iter().map(|param| param.as_ref()).collect();
let rows = stmt.query_map(param_refs.as_slice(), |row| {
Ok((read_entity(row)?, row.get::<_, i64>(15)?))
})?;
let mut entries = rows.collect::<Result<Vec<_>, _>>()?;
let has_more = entries.len() > limit_usize;
if has_more {
entries.truncate(limit_usize);
}
let next_after = if has_more {
entries.last().map(|(entity, sequence)| SeekCursor {
sequence: *sequence,
id: entity.id,
})
} else {
None
};
let items = entries.into_iter().map(|(entity, _)| entity).collect();
Ok(SeekPage { items, next_after })
})
.await
}
async fn get_entity_including_deleted(&self, id: Uuid) -> Result<Option<Entity>, StorageError> {
let id_str = id.to_string();
self.with_reader("get_entity_including_deleted", move |conn| {
let sql =
format!("SELECT {ENTITY_SELECT_COLUMNS} FROM entities WHERE entities.id = ?1");
let mut stmt = conn.prepare(&sql)?;
let mut rows = stmt.query(rusqlite::params![id_str])?;
match rows.next()? {
Some(row) => Ok(Some(read_entity(row)?)),
None => Ok(None),
}
})
.await
}
async fn count_entities(
&self,
namespace: &str,
filter: EntityFilter,
) -> Result<u64, StorageError> {
let namespace = namespace.to_string();
self.with_reader("count_entities", move |conn| {
if filter.namespaces.is_empty() {
let (where_sql, params) = build_entity_where(&namespace, &filter);
let sql = format!("SELECT COUNT(*) FROM entities{}", where_sql);
let mut stmt = conn.prepare(&sql)?;
let param_refs: Vec<&dyn rusqlite::types::ToSql> =
params.iter().map(|p| p.as_ref()).collect();
let count: i64 = stmt.query_row(param_refs.as_slice(), |row| row.get(0))?;
return Ok(count as u64);
}
let deduped_namespaces: Vec<String> = filter
.namespaces
.iter()
.cloned()
.collect::<HashSet<_>>()
.into_iter()
.collect();
let mut total = 0;
for chunk in deduped_namespaces.chunks(NAMESPACE_COUNT_CHUNK_SIZE) {
let chunk_filter = EntityFilter {
namespaces: chunk.to_vec(),
..filter.clone()
};
let (where_sql, params) = build_entity_where(&namespace, &chunk_filter);
let sql = format!("SELECT COUNT(*) FROM entities{}", where_sql);
let mut stmt = conn.prepare(&sql)?;
let param_refs: Vec<&dyn rusqlite::types::ToSql> =
params.iter().map(|p| p.as_ref()).collect();
let count: i64 = stmt.query_row(param_refs.as_slice(), |row| row.get(0))?;
total += count as u64;
}
Ok(total)
})
.await
}
}
const ENTITIES_DDL: &str = include_str!("../../sql/entities-ddl.sql");
pub(crate) fn ensure_entities_schema(conn: &rusqlite::Connection) -> Result<(), rusqlite::Error> {
conn.execute_batch(ENTITIES_DDL)
}
#[cfg(test)]
#[path = "entity_tests.rs"]
mod tests;