update: pipeline, search, clip, embedding fixes
This commit is contained in:
1
src/core/identity/mod.rs
Normal file
1
src/core/identity/mod.rs
Normal file
@@ -0,0 +1 @@
|
||||
pub mod storage;
|
||||
513
src/core/identity/storage.rs
Normal file
513
src/core/identity/storage.rs
Normal file
@@ -0,0 +1,513 @@
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tracing::warn;
|
||||
|
||||
use crate::core::config::OUTPUT_DIR;
|
||||
use crate::core::db::PostgresDb;
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct IdentityFile {
|
||||
pub version: u32,
|
||||
pub identity_uuid: String,
|
||||
pub name: String,
|
||||
pub identity_type: Option<String>,
|
||||
pub source: Option<String>,
|
||||
pub status: Option<String>,
|
||||
pub tmdb_id: Option<i32>,
|
||||
pub tmdb_profile: Option<String>,
|
||||
pub metadata: serde_json::Value,
|
||||
pub file_bindings: Vec<FileBinding>,
|
||||
pub created_at: String,
|
||||
pub updated_at: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct FileBinding {
|
||||
pub file_uuid: String,
|
||||
pub trace_ids: Vec<i32>,
|
||||
pub face_count: i64,
|
||||
}
|
||||
|
||||
pub fn identities_root() -> PathBuf {
|
||||
PathBuf::from(&*OUTPUT_DIR).join("identities")
|
||||
}
|
||||
|
||||
pub fn identity_dir(uuid: &str) -> PathBuf {
|
||||
identities_root().join(uuid)
|
||||
}
|
||||
|
||||
pub fn identity_file_path(uuid: &str) -> PathBuf {
|
||||
identity_dir(uuid).join("identity.json")
|
||||
}
|
||||
|
||||
pub fn index_path() -> PathBuf {
|
||||
identities_root().join("_index.json")
|
||||
}
|
||||
|
||||
pub fn read_identity_file(uuid: &str) -> Result<IdentityFile> {
|
||||
let path = identity_file_path(uuid);
|
||||
let content = std::fs::read_to_string(&path)
|
||||
.with_context(|| format!("Identity file not found: {} ({})", uuid, path.display()))?;
|
||||
serde_json::from_str(&content)
|
||||
.with_context(|| format!("Invalid identity.json: {}", uuid))
|
||||
}
|
||||
|
||||
pub fn write_identity_file(file: &IdentityFile) -> Result<()> {
|
||||
let dir = identity_dir(&file.identity_uuid);
|
||||
std::fs::create_dir_all(&dir)
|
||||
.with_context(|| format!("Failed to create identity dir: {}", dir.display()))?;
|
||||
|
||||
let path = dir.join("identity.json");
|
||||
let json = serde_json::to_string_pretty(file)
|
||||
.with_context(|| format!("Failed to serialize identity: {}", file.identity_uuid))?;
|
||||
std::fs::write(&path, &json)
|
||||
.with_context(|| format!("Failed to write identity.json: {}", path.display()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn delete_identity_file(uuid: &str) -> Result<()> {
|
||||
let path = identity_file_path(uuid);
|
||||
if path.exists() {
|
||||
std::fs::remove_file(&path)
|
||||
.with_context(|| format!("Failed to delete identity.json: {}", path.display()))?;
|
||||
}
|
||||
let dir = identity_dir(uuid);
|
||||
if dir.exists() {
|
||||
std::fs::remove_dir(&dir).ok();
|
||||
}
|
||||
remove_from_index(uuid).ok();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn list_identity_uuids() -> Result<Vec<String>> {
|
||||
let root = identities_root();
|
||||
if !root.is_dir() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let mut uuids = Vec::new();
|
||||
for entry in std::fs::read_dir(&root)
|
||||
.with_context(|| format!("Failed to read identities dir: {}", root.display()))?
|
||||
{
|
||||
let entry = entry?;
|
||||
let name = entry.file_name().to_string_lossy().to_string();
|
||||
if entry.file_type().map(|t| t.is_dir()).unwrap_or(false)
|
||||
&& name.len() == 32
|
||||
&& name.chars().all(|c| c.is_ascii_hexdigit())
|
||||
{
|
||||
uuids.push(name);
|
||||
}
|
||||
}
|
||||
uuids.sort();
|
||||
Ok(uuids)
|
||||
}
|
||||
|
||||
pub fn count_identity_files() -> usize {
|
||||
list_identity_uuids().map(|v| v.len()).unwrap_or(0)
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
struct IndexFile {
|
||||
version: u32,
|
||||
updated_at: String,
|
||||
entries: HashMap<String, String>,
|
||||
}
|
||||
|
||||
fn read_index_inner() -> Result<IndexFile> {
|
||||
let path = index_path();
|
||||
if !path.exists() {
|
||||
return Ok(IndexFile {
|
||||
version: 1,
|
||||
updated_at: chrono::Utc::now().to_rfc3339(),
|
||||
entries: HashMap::new(),
|
||||
});
|
||||
}
|
||||
let content = std::fs::read_to_string(&path)
|
||||
.with_context(|| format!("Failed to read index: {}", path.display()))?;
|
||||
serde_json::from_str(&content)
|
||||
.with_context(|| format!("Invalid _index.json: {}", path.display()))
|
||||
}
|
||||
|
||||
pub fn read_index() -> Result<HashMap<String, String>> {
|
||||
read_index_inner().map(|idx| idx.entries)
|
||||
}
|
||||
|
||||
pub fn update_index(uuid: &str, name: &str) -> Result<()> {
|
||||
let mut idx = read_index_inner()?;
|
||||
idx.entries.insert(uuid.to_string(), name.to_string());
|
||||
idx.updated_at = chrono::Utc::now().to_rfc3339();
|
||||
let root = identities_root();
|
||||
std::fs::create_dir_all(&root)?;
|
||||
let json = serde_json::to_string_pretty(&idx)?;
|
||||
std::fs::write(index_path(), &json)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn remove_from_index(uuid: &str) -> Result<()> {
|
||||
let mut idx = read_index_inner()?;
|
||||
idx.entries.remove(uuid);
|
||||
idx.updated_at = chrono::Utc::now().to_rfc3339();
|
||||
let json = serde_json::to_string_pretty(&idx)?;
|
||||
std::fs::write(index_path(), &json)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn rebuild_index() -> Result<usize> {
|
||||
let uuids = list_identity_uuids()?;
|
||||
let mut entries = HashMap::new();
|
||||
for uuid in &uuids {
|
||||
match read_identity_file(uuid) {
|
||||
Ok(file) => {
|
||||
entries.insert(uuid.clone(), file.name);
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("[identity-storage] Skipping {} in index rebuild: {}", uuid, e);
|
||||
}
|
||||
}
|
||||
}
|
||||
let idx = IndexFile {
|
||||
version: 1,
|
||||
updated_at: chrono::Utc::now().to_rfc3339(),
|
||||
entries,
|
||||
};
|
||||
let root = identities_root();
|
||||
std::fs::create_dir_all(&root)?;
|
||||
let json = serde_json::to_string_pretty(&idx)?;
|
||||
std::fs::write(index_path(), &json)?;
|
||||
Ok(uuids.len())
|
||||
}
|
||||
|
||||
pub async fn save_identity_file_by_pool(pool: &sqlx::PgPool, uuid: &str) -> Result<()> {
|
||||
let identity_table = crate::core::db::schema::table_name("identities");
|
||||
let fd_table = crate::core::db::schema::table_name("face_detections");
|
||||
|
||||
let clean = uuid.replace('-', "");
|
||||
let record = sqlx::query_as::<_, crate::core::db::IdentityDetailRecord>(
|
||||
&format!(
|
||||
"SELECT id, uuid::text, name, identity_type, source, status, metadata, reference_data, \
|
||||
NULL::real[] as voice_embedding, NULL::real[] as identity_embedding, \
|
||||
face_embedding::real[] as face_embedding, \
|
||||
tmdb_id, tmdb_profile, created_at::timestamptz as created_at, NULL::timestamptz as updated_at \
|
||||
FROM {} WHERE REPLACE(uuid::text, '-', '') = $1",
|
||||
identity_table
|
||||
)
|
||||
)
|
||||
.bind(&clean)
|
||||
.fetch_optional(pool)
|
||||
.await?
|
||||
.with_context(|| format!("Identity not found in DB: {}", uuid))?;
|
||||
|
||||
let identity_uuid = record.uuid.clone();
|
||||
|
||||
let binding_rows = sqlx::query_as::<_, (String, Vec<i32>, i64)>(
|
||||
&format!(
|
||||
"SELECT fd.file_uuid, COALESCE(array_agg(DISTINCT fd.trace_id) FILTER (WHERE fd.trace_id IS NOT NULL), '{{}}'::int[]), COUNT(*)::bigint \
|
||||
FROM {} fd WHERE fd.identity_id = $1 GROUP BY fd.file_uuid ORDER BY fd.file_uuid",
|
||||
fd_table
|
||||
)
|
||||
)
|
||||
.bind(record.id)
|
||||
.fetch_all(pool)
|
||||
.await?;
|
||||
|
||||
let file_bindings: Vec<FileBinding> = binding_rows
|
||||
.into_iter()
|
||||
.map(|(fu, tids, cnt)| FileBinding {
|
||||
file_uuid: fu,
|
||||
trace_ids: tids,
|
||||
face_count: cnt,
|
||||
})
|
||||
.collect();
|
||||
|
||||
let fmt_time = |dt: Option<chrono::DateTime<chrono::Utc>>| -> String {
|
||||
dt.map(|d| d.to_rfc3339())
|
||||
.unwrap_or_else(|| chrono::Utc::now().to_rfc3339())
|
||||
};
|
||||
|
||||
let file = IdentityFile {
|
||||
version: 1,
|
||||
identity_uuid,
|
||||
name: record.name,
|
||||
identity_type: record.identity_type,
|
||||
source: record.source,
|
||||
status: record.status,
|
||||
tmdb_id: record.tmdb_id,
|
||||
tmdb_profile: record.tmdb_profile,
|
||||
metadata: record.metadata,
|
||||
file_bindings,
|
||||
created_at: fmt_time(record.created_at),
|
||||
updated_at: fmt_time(record.updated_at),
|
||||
};
|
||||
|
||||
write_identity_file(&file)?;
|
||||
update_index(&file.identity_uuid, &file.name)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub fn list_identity_uuids_at(base: &std::path::Path) -> Result<Vec<String>> {
|
||||
let root = base.join("identities");
|
||||
if !root.is_dir() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let mut uuids = Vec::new();
|
||||
for entry in std::fs::read_dir(&root)? {
|
||||
let entry = entry?;
|
||||
let name = entry.file_name().to_string_lossy().to_string();
|
||||
if entry.file_type().map(|t| t.is_dir()).unwrap_or(false)
|
||||
&& name.len() == 32
|
||||
&& name.chars().all(|c| c.is_ascii_hexdigit())
|
||||
{
|
||||
uuids.push(name);
|
||||
}
|
||||
}
|
||||
uuids.sort();
|
||||
Ok(uuids)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub fn identity_dir_at(base: &std::path::Path, uuid: &str) -> std::path::PathBuf {
|
||||
base.join("identities").join(uuid)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub fn identity_file_path_at(base: &std::path::Path, uuid: &str) -> std::path::PathBuf {
|
||||
identity_dir_at(base, uuid).join("identity.json")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub fn index_path_at(base: &std::path::Path) -> std::path::PathBuf {
|
||||
base.join("identities").join("_index.json")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub fn read_identity_file_at(base: &std::path::Path, uuid: &str) -> Result<IdentityFile> {
|
||||
let path = identity_file_path_at(base, uuid);
|
||||
let content = std::fs::read_to_string(&path)?;
|
||||
serde_json::from_str(&content).map_err(Into::into)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub fn write_identity_file_at(base: &std::path::Path, file: &IdentityFile) -> Result<()> {
|
||||
let dir = identity_dir_at(base, &file.identity_uuid);
|
||||
std::fs::create_dir_all(&dir)?;
|
||||
let json = serde_json::to_string_pretty(file)?;
|
||||
std::fs::write(dir.join("identity.json"), &json)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub fn update_index_at(base: &std::path::Path, uuid: &str, name: &str) -> Result<()> {
|
||||
use std::collections::HashMap;
|
||||
let index_path = index_path_at(base);
|
||||
let mut entries: HashMap<String, String> = if index_path.exists() {
|
||||
let content = std::fs::read_to_string(&index_path)?;
|
||||
let v: serde_json::Value = serde_json::from_str(&content).unwrap_or_default();
|
||||
v["entries"].as_object()
|
||||
.map(|obj| obj.iter().map(|(k, v)| (k.clone(), v.as_str().unwrap_or("").to_string())).collect())
|
||||
.unwrap_or_default()
|
||||
} else {
|
||||
HashMap::new()
|
||||
};
|
||||
entries.insert(uuid.to_string(), name.to_string());
|
||||
std::fs::create_dir_all(base.join("identities"))?;
|
||||
let json = serde_json::to_string_pretty(&serde_json::json!({
|
||||
"version": 1, "updated_at": chrono::Utc::now().to_rfc3339(), "entries": entries
|
||||
}))?;
|
||||
std::fs::write(&index_path, &json)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn save_identity_file(db: &PostgresDb, uuid: &str) -> Result<()> {
|
||||
let record = db.get_identity_by_uuid(uuid).await?
|
||||
.with_context(|| format!("Identity not found in DB: {}", uuid))?;
|
||||
|
||||
let identity_uuid = record.uuid.clone();
|
||||
|
||||
let binding_rows = sqlx::query_as::<_, (String, Vec<i32>, i64)>(
|
||||
"SELECT fd.file_uuid, COALESCE(array_agg(DISTINCT fd.trace_id) FILTER (WHERE fd.trace_id IS NOT NULL), '{}'::int[]), COUNT(*)::bigint \
|
||||
FROM face_detections fd \
|
||||
WHERE fd.identity_id = $1 \
|
||||
GROUP BY fd.file_uuid \
|
||||
ORDER BY fd.file_uuid"
|
||||
)
|
||||
.bind(record.id)
|
||||
.fetch_all(db.pool())
|
||||
.await
|
||||
.with_context(|| format!("Failed to query bindings for identity: {}", identity_uuid))?;
|
||||
|
||||
let file_bindings: Vec<FileBinding> = binding_rows
|
||||
.into_iter()
|
||||
.map(|(fu, tids, cnt)| FileBinding {
|
||||
file_uuid: fu,
|
||||
trace_ids: tids,
|
||||
face_count: cnt,
|
||||
})
|
||||
.collect();
|
||||
|
||||
let fmt_time = |dt: Option<chrono::DateTime<chrono::Utc>>| -> String {
|
||||
dt.map(|d| d.to_rfc3339())
|
||||
.unwrap_or_else(|| chrono::Utc::now().to_rfc3339())
|
||||
};
|
||||
|
||||
let file = IdentityFile {
|
||||
version: 1,
|
||||
identity_uuid,
|
||||
name: record.name,
|
||||
identity_type: record.identity_type,
|
||||
source: record.source,
|
||||
status: record.status,
|
||||
tmdb_id: record.tmdb_id,
|
||||
tmdb_profile: record.tmdb_profile,
|
||||
metadata: record.metadata,
|
||||
file_bindings,
|
||||
created_at: fmt_time(record.created_at),
|
||||
updated_at: fmt_time(record.updated_at),
|
||||
};
|
||||
|
||||
write_identity_file(&file)?;
|
||||
update_index(&file.identity_uuid, &file.name)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::path::Path;
|
||||
|
||||
fn sample_identity() -> IdentityFile {
|
||||
IdentityFile {
|
||||
version: 1,
|
||||
identity_uuid: "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa".to_string(),
|
||||
name: "Test Person".to_string(),
|
||||
identity_type: Some("people".to_string()),
|
||||
source: Some("tmdb".to_string()),
|
||||
status: Some("confirmed".to_string()),
|
||||
tmdb_id: Some(112),
|
||||
tmdb_profile: Some("https://image.tmdb.org/t/p/w185/test.jpg".to_string()),
|
||||
metadata: serde_json::json!({"tmdb_character": "Test Role"}),
|
||||
file_bindings: vec![FileBinding {
|
||||
file_uuid: "ffffffffffffffffffffffffffffffff".to_string(),
|
||||
trace_ids: vec![1, 2, 3],
|
||||
face_count: 5,
|
||||
}],
|
||||
created_at: "2026-05-16T00:00:00+00:00".to_string(),
|
||||
updated_at: "2026-05-16T01:00:00+00:00".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_serde_roundtrip() {
|
||||
let file = sample_identity();
|
||||
let json = serde_json::to_string_pretty(&file).unwrap();
|
||||
let parsed: IdentityFile = serde_json::from_str(&json).unwrap();
|
||||
assert_eq!(parsed.name, "Test Person");
|
||||
assert_eq!(parsed.identity_uuid, "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa");
|
||||
assert_eq!(parsed.tmdb_id, Some(112));
|
||||
assert_eq!(parsed.file_bindings.len(), 1);
|
||||
assert_eq!(parsed.file_bindings[0].face_count, 5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_identity_dir_path() {
|
||||
let uuid = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa";
|
||||
let p = identity_dir(uuid);
|
||||
assert!(p.to_string_lossy().ends_with(&format!("identities/{}", uuid)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_identity_file_path() {
|
||||
let uuid = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa";
|
||||
let p = identity_file_path(uuid);
|
||||
assert!(p.to_string_lossy().ends_with("identity.json"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_index_path() {
|
||||
let p = index_path();
|
||||
assert!(p.to_string_lossy().ends_with("_index.json"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_identity_dir_at() {
|
||||
let base = Path::new("/tmp/test_base");
|
||||
let uuid = "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb";
|
||||
let p = identity_dir_at(base, uuid);
|
||||
assert_eq!(p, Path::new("/tmp/test_base/identities/bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_identity_file_path_at() {
|
||||
let base = Path::new("/tmp/test_base");
|
||||
let uuid = "cccccccccccccccccccccccccccccccc";
|
||||
let p = identity_file_path_at(base, uuid);
|
||||
assert_eq!(
|
||||
p,
|
||||
Path::new("/tmp/test_base/identities/cccccccccccccccccccccccccccccccc/identity.json")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_write_then_read_identity_file_at() {
|
||||
let tmp = std::env::temp_dir().join("momentry_test_write_read");
|
||||
let _ = std::fs::remove_dir_all(&tmp);
|
||||
let base = &tmp;
|
||||
|
||||
let file = sample_identity();
|
||||
write_identity_file_at(base, &file).unwrap();
|
||||
|
||||
let read = read_identity_file_at(base, &file.identity_uuid).unwrap();
|
||||
assert_eq!(read.name, file.name);
|
||||
assert_eq!(read.source, file.source);
|
||||
assert_eq!(read.tmdb_id, file.tmdb_id);
|
||||
assert_eq!(read.file_bindings[0].face_count, file.file_bindings[0].face_count);
|
||||
|
||||
let _ = std::fs::remove_dir_all(&tmp);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_update_and_read_index_at() {
|
||||
let tmp = std::env::temp_dir().join("momentry_test_index");
|
||||
let _ = std::fs::remove_dir_all(&tmp);
|
||||
let base = &tmp;
|
||||
|
||||
update_index_at(base, "aaa", "Alice").unwrap();
|
||||
update_index_at(base, "bbb", "Bob").unwrap();
|
||||
|
||||
let idx_path = index_path_at(base);
|
||||
let content = std::fs::read_to_string(&idx_path).unwrap();
|
||||
let parsed: serde_json::Value = serde_json::from_str(&content).unwrap();
|
||||
let entries = parsed["entries"].as_object().unwrap();
|
||||
assert_eq!(entries.len(), 2);
|
||||
assert_eq!(entries["aaa"], "Alice");
|
||||
assert_eq!(entries["bbb"], "Bob");
|
||||
|
||||
let _ = std::fs::remove_dir_all(&tmp);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_list_identity_uuids_at() {
|
||||
let tmp = std::env::temp_dir().join("momentry_test_list");
|
||||
let _ = std::fs::remove_dir_all(&tmp);
|
||||
let base = &tmp;
|
||||
|
||||
std::fs::create_dir_all(base.join("identities").join("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa")).unwrap();
|
||||
std::fs::create_dir_all(base.join("identities").join("bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb")).unwrap();
|
||||
std::fs::create_dir_all(base.join("identities").join("cccccccccccccccccccccccccccccccc")).unwrap();
|
||||
std::fs::create_dir_all(base.join("identities").join("not_a_uuid")).unwrap();
|
||||
std::fs::create_dir_all(base.join("identities").join("short")).unwrap();
|
||||
|
||||
let uuids = list_identity_uuids_at(base).unwrap();
|
||||
assert_eq!(uuids.len(), 3);
|
||||
assert!(uuids.contains(&"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa".to_string()));
|
||||
assert!(uuids.contains(&"bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb".to_string()));
|
||||
assert!(uuids.contains(&"cccccccccccccccccccccccccccccccc".to_string()));
|
||||
|
||||
let _ = std::fs::remove_dir_all(&tmp);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user