Files
keep/src/meta_plugin/digest.rs
Andrew Phillips a0fcb86fce refactor: update response types to MetaPluginResponse
Co-authored-by: aider (openai/andrew/openrouter/mistralai/mistral-medium-3.1) <aider@aider.chat>
2025-08-26 16:17:34 -03:00

341 lines
10 KiB
Rust

use anyhow::Result;
use sha2::{Digest, Sha256};
use std::time::Instant;
use rusqlite::Connection;
use crate::meta_plugin::MetaPlugin;
#[derive(Debug, Clone, Default)]
pub struct DigestSha256MetaPlugin {
hasher: Sha256,
meta_name: String,
item_id: Option<i64>,
outputs: std::collections::HashMap<String, serde_yaml::Value>,
options: std::collections::HashMap<String, serde_yaml::Value>,
}
impl DigestSha256MetaPlugin {
pub fn new(
_options: Option<std::collections::HashMap<String, serde_yaml::Value>>,
outputs: Option<std::collections::HashMap<String, serde_yaml::Value>>,
) -> DigestSha256MetaPlugin {
// Start with default options
let mut final_options = std::collections::HashMap::new();
if let Some(opts) = _options {
for (key, value) in opts {
final_options.insert(key, value);
}
}
// Start with default outputs
let mut final_outputs = std::collections::HashMap::new();
let default_outputs = vec!["digest_sha256".to_string()];
for output_name in default_outputs {
final_outputs.insert(output_name.clone(), serde_yaml::Value::String(output_name));
}
if let Some(outs) = outputs {
for (key, value) in outs {
final_outputs.insert(key, value);
}
}
DigestSha256MetaPlugin {
hasher: Sha256::new(),
meta_name: "digest_sha256".to_string(),
item_id: None,
outputs: final_outputs,
options: final_options,
}
}
pub fn new_simple() -> DigestSha256MetaPlugin {
Self::new(None, None)
}
}
impl MetaPlugin for DigestSha256MetaPlugin {
fn is_internal(&self) -> bool {
true
}
fn initialize(&mut self, item_id: i64) -> Result<MetaPluginResponse> {
self.item_id = Some(item_id);
Ok(MetaPluginResponse::default())
}
fn finalize(&mut self) -> Result<MetaPluginResponse> {
let mut metadata = Vec::new();
if let Some(item_id) = self.item_id {
// Finalize the hash
let hash_result = self.hasher.finalize_reset();
let hex_string = format!("{:x}", hash_result);
// Create metadata to be stored
if let Some(meta) = self.create_meta(item_id, "digest_sha256", hex_string) {
metadata.push(meta);
}
}
Ok(MetaPluginResponse {
metadata: Some(metadata),
is_finalized: true,
})
}
fn update(&mut self, data: &[u8]) -> Result<MetaPluginResponse> {
self.hasher.update(data);
Ok(MetaPluginResponse::default())
}
fn meta_name(&self) -> String {
self.meta_name.clone()
}
fn outputs(&self) -> &std::collections::HashMap<String, serde_yaml::Value> {
&self.outputs
}
fn outputs_mut(&mut self) -> &mut std::collections::HashMap<String, serde_yaml::Value> {
&mut self.outputs
}
fn default_outputs(&self) -> Vec<String> {
vec!["digest_sha256".to_string()]
}
fn default_options(&self) -> std::collections::HashMap<String, serde_yaml::Value> {
std::collections::HashMap::new()
}
fn options(&self) -> &std::collections::HashMap<String, serde_yaml::Value> {
&self.options
}
fn options_mut(&mut self) -> &mut std::collections::HashMap<String, serde_yaml::Value> {
&mut self.options
}
}
#[derive(Debug, Clone, Default)]
pub struct ReadTimeMetaPlugin {
start_time: Option<Instant>,
meta_name: String,
outputs: std::collections::HashMap<String, serde_yaml::Value>,
options: std::collections::HashMap<String, serde_yaml::Value>,
}
impl ReadTimeMetaPlugin {
pub fn new(
_options: Option<std::collections::HashMap<String, serde_yaml::Value>>,
outputs: Option<std::collections::HashMap<String, serde_yaml::Value>>,
) -> ReadTimeMetaPlugin {
// Start with default options
let mut final_options = std::collections::HashMap::new();
if let Some(opts) = _options {
for (key, value) in opts {
final_options.insert(key, value);
}
}
// Start with default outputs
let mut final_outputs = std::collections::HashMap::new();
let default_outputs = Self::default().default_outputs();
for output_name in default_outputs {
final_outputs.insert(output_name.clone(), serde_yaml::Value::String(output_name));
}
if let Some(outs) = outputs {
for (key, value) in outs {
final_outputs.insert(key, value);
}
}
ReadTimeMetaPlugin {
start_time: None,
meta_name: "read_time".to_string(),
outputs: final_outputs,
options: final_options,
}
}
pub fn new_simple() -> ReadTimeMetaPlugin {
Self::new(None, None)
}
}
impl MetaPlugin for ReadTimeMetaPlugin {
fn is_internal(&self) -> bool {
true
}
fn finalize(&mut self) -> Result<MetaPluginResponse> {
let mut metadata = Vec::new();
if let Some(start_time) = self.start_time {
if let Some(item_id) = self.item_id {
let duration = start_time.elapsed();
let duration_str = format!("{:.3} seconds", duration.as_secs_f64());
if let Some(meta) = self.create_meta(item_id, "read_time", duration_str) {
metadata.push(meta);
}
}
}
Ok(MetaPluginResponse {
metadata: Some(metadata),
is_finalized: true,
})
}
fn update(&mut self, _data: &[u8]) -> Result<MetaPluginResponse> {
if self.start_time.is_none() {
self.start_time = Some(Instant::now());
}
Ok(MetaPluginResponse::default())
}
fn meta_name(&self) -> String {
self.meta_name.clone()
}
fn outputs(&self) -> &std::collections::HashMap<String, serde_yaml::Value> {
&self.outputs
}
fn outputs_mut(&mut self) -> &mut std::collections::HashMap<String, serde_yaml::Value> {
&mut self.outputs
}
fn default_outputs(&self) -> Vec<String> {
vec!["read_time".to_string()]
}
fn default_options(&self) -> std::collections::HashMap<String, serde_yaml::Value> {
std::collections::HashMap::new()
}
fn options(&self) -> &std::collections::HashMap<String, serde_yaml::Value> {
&self.options
}
fn options_mut(&mut self) -> &mut std::collections::HashMap<String, serde_yaml::Value> {
&mut self.options
}
}
#[derive(Debug, Clone, Default)]
pub struct ReadRateMetaPlugin {
start_time: Option<Instant>,
bytes_read: u64,
meta_name: String,
outputs: std::collections::HashMap<String, serde_yaml::Value>,
options: std::collections::HashMap<String, serde_yaml::Value>,
}
impl ReadRateMetaPlugin {
pub fn new(
_options: Option<std::collections::HashMap<String, serde_yaml::Value>>,
outputs: Option<std::collections::HashMap<String, serde_yaml::Value>>,
) -> ReadRateMetaPlugin {
// Start with default options
let mut final_options = std::collections::HashMap::new();
if let Some(opts) = _options {
for (key, value) in opts {
final_options.insert(key, value);
}
}
// Start with default outputs
let mut final_outputs = std::collections::HashMap::new();
let default_outputs = Self::default().default_outputs();
for output_name in default_outputs {
final_outputs.insert(output_name.clone(), serde_yaml::Value::String(output_name));
}
if let Some(outs) = outputs {
for (key, value) in outs {
final_outputs.insert(key, value);
}
}
ReadRateMetaPlugin {
start_time: None,
bytes_read: 0,
meta_name: "read_rate".to_string(),
outputs: final_outputs,
options: final_options,
}
}
pub fn new_simple() -> ReadRateMetaPlugin {
Self::new(None, None)
}
}
impl MetaPlugin for ReadRateMetaPlugin {
fn is_internal(&self) -> bool {
true
}
fn finalize(&mut self) -> Result<MetaPluginResponse> {
let mut metadata = Vec::new();
if let Some(start_time) = self.start_time {
if let Some(item_id) = self.item_id {
let duration = start_time.elapsed();
let rate = if duration.as_secs_f64() > 0.0 {
format!("{:.2} KB/s", (self.bytes_read as f64 / 1024.0) / duration.as_secs_f64())
} else {
"N/A".to_string()
};
if let Some(meta) = self.create_meta(item_id, "read_rate", rate) {
metadata.push(meta);
}
}
}
Ok(MetaPluginResponse {
metadata: Some(metadata),
is_finalized: true,
})
}
fn update(&mut self, data: &[u8]) -> Result<MetaPluginResponse> {
if self.start_time.is_none() {
self.start_time = Some(Instant::now());
}
self.bytes_read += data.len() as u64;
Ok(MetaPluginResponse::default())
}
fn meta_name(&self) -> String {
self.meta_name.clone()
}
fn outputs(&self) -> &std::collections::HashMap<String, serde_yaml::Value> {
&self.outputs
}
fn outputs_mut(&mut self) -> &mut std::collections::HashMap<String, serde_yaml::Value> {
&mut self.outputs
}
fn default_outputs(&self) -> Vec<String> {
vec!["read_rate".to_string()]
}
fn default_options(&self) -> std::collections::HashMap<String, serde_yaml::Value> {
std::collections::HashMap::new()
}
fn options(&self) -> &std::collections::HashMap<String, serde_yaml::Value> {
&self.options
}
fn options_mut(&mut self) -> &mut std::collections::HashMap<String, serde_yaml::Value> {
&mut self.options
}
}