use chrono::{DateTime, Utc};
use crate::model::metric::{MetricKey, MetricPoint, Step};
use crate::model::run::{Run, RunId, RunStatus};
use crate::model::types::ProjectId;
use crate::storage::rows::{StoredRun, status_as_str};
use crate::storage::time::{current_timestamp, timestamp_as_rfc3339};
use crate::storage::{StorageError, percent_encode_metric_key};
pub struct NativeWriteStore<'connection> {
connection: &'connection duckdb::Connection,
}
impl<'connection> NativeWriteStore<'connection> {
pub const fn new(connection: &'connection duckdb::Connection) -> Self {
Self { connection }
}
pub fn create_run(
&self,
project_id: &ProjectId,
name: &str,
run_id: Option<RunId>,
) -> Result<Run, StorageError> {
let run_id = run_id.unwrap_or_else(|| RunId::from_string(uuid::Uuid::new_v4().to_string()));
if self.run_exists(&run_id)? {
return Err(StorageError::RunAlreadyExists {
run_id: run_id.as_str().to_owned(),
});
}
let now = current_timestamp("created_at")?;
self.connection.execute(
"INSERT INTO seex_runs
(run_id, project_id, name, status, created_at, started_at, finished_at)
VALUES (?, ?, ?, ?, ?, ?, NULL)",
(
run_id.as_str(),
project_id.as_str(),
name,
status_as_str(RunStatus::Running),
timestamp_as_rfc3339(now),
timestamp_as_rfc3339(now),
),
)?;
Ok(Run {
run_id,
project_id: project_id.clone(),
name: name.to_owned(),
status: RunStatus::Running,
created_at: now,
started_at: now,
finished_at: None,
})
}
pub fn resume_run(&self, run_id: &RunId) -> Result<Run, StorageError> {
let result = self.connection.query_row(
"SELECT run_id, project_id, name, status,
epoch_ms(created_at::TIMESTAMPTZ),
epoch_ms(started_at::TIMESTAMPTZ),
epoch_ms(finished_at::TIMESTAMPTZ)
FROM seex_runs
WHERE run_id = ?",
[run_id.as_str()],
|row| {
Ok(StoredRun {
run_id: row.get(0)?,
project_id: row.get(1)?,
name: row.get(2)?,
status: row.get(3)?,
created_at_millis: row.get(4)?,
started_at_millis: row.get(5)?,
finished_at_millis: row.get(6)?,
})
},
);
let stored = match result {
Ok(stored) => stored,
Err(duckdb::Error::QueryReturnedNoRows) => {
return Err(StorageError::RunNotFound {
run_id: run_id.as_str().to_owned(),
});
}
Err(source) => return Err(source.into()),
};
stored.into_run()
}
pub fn log_metric_at_step(
&self,
run_id: &RunId,
metric_key: &MetricKey,
step: Step,
value_f64: f64,
) -> Result<MetricPoint, StorageError> {
let timestamp = current_timestamp("timestamp")?;
let ingested_at = current_timestamp("ingested_at")?;
self.log_metric_at_step_with_timestamps(
run_id,
metric_key,
step,
value_f64,
timestamp,
ingested_at,
)
}
pub fn log_metric_at_step_with_timestamps(
&self,
run_id: &RunId,
metric_key: &MetricKey,
step: Step,
value_f64: f64,
timestamp: DateTime<Utc>,
ingested_at: DateTime<Utc>,
) -> Result<MetricPoint, StorageError> {
self.connection.execute_batch("BEGIN TRANSACTION")?;
let directory_result: Result<(), StorageError> = (|| {
self.connection.execute(
"INSERT INTO seex_metric_keys(metric_key)
SELECT ? WHERE NOT EXISTS (SELECT 1 FROM seex_metric_keys WHERE metric_key = ?)",
[metric_key.as_str(), metric_key.as_str()],
)?;
self.connection.execute(
"INSERT INTO seex_series_directory(run_id, metric_key, status)
SELECT ?, ?, 'pending' WHERE NOT EXISTS (
SELECT 1 FROM seex_series_directory WHERE run_id = ? AND metric_key = ?
)",
[
run_id.as_str(),
metric_key.as_str(),
run_id.as_str(),
metric_key.as_str(),
],
)?;
Ok(())
})();
if let Err(error) = directory_result {
let _ = self.connection.execute_batch("ROLLBACK");
return Err(error);
}
self.connection.execute_batch("COMMIT")?;
self.connection.execute(
"INSERT INTO dl.metric_points
(run_id, metric_key, metric_key_encoded, step, timestamp, value_f64, ingested_at)
VALUES (?, ?, ?, ?, ?, ?, ?)",
(
run_id.as_str(),
metric_key.as_str(),
percent_encode_metric_key(metric_key.as_str()),
step.value(),
timestamp_as_rfc3339(timestamp),
value_f64,
timestamp_as_rfc3339(ingested_at),
),
)?;
self.connection.execute(
"UPDATE seex_series_directory SET status = 'published'
WHERE run_id = ? AND metric_key = ? AND status = 'pending'",
[run_id.as_str(), metric_key.as_str()],
)?;
Ok(MetricPoint {
run_id: run_id.clone(),
metric_key: metric_key.clone(),
step,
timestamp,
value_f64,
ingested_at,
})
}
pub fn rebuild_metric_aggregates_for_run(&self, run_id: &RunId) -> Result<(), StorageError> {
self.connection.execute(
"DELETE FROM seex_metric_aggregates
WHERE run_id = ?",
[run_id.as_str()],
)?;
self.connection.execute(
"INSERT INTO seex_metric_aggregates
(run_id, metric_key, effective_count, last_step, last_value_f64,
min_value_f64, max_value_f64)
SELECT run_id, metric_key, count(*), max(step),
-- SQLite cannot persist IEEE non-finite values in scalar
-- columns; chart diagnostics still inspect the source rows.
CASE WHEN isfinite(arg_max(value_f64, step))
THEN arg_max(value_f64, step) ELSE 0.0 END,
CASE WHEN isfinite(min(value_f64))
THEN min(value_f64) ELSE 0.0 END,
CASE WHEN isfinite(max(value_f64))
THEN max(value_f64) ELSE 0.0 END
FROM dl.metric_points
WHERE run_id = ?
GROUP BY run_id, metric_key",
[run_id.as_str()],
)?;
Ok(())
}
fn run_exists(&self, run_id: &RunId) -> Result<bool, StorageError> {
let count: i64 = self.connection.query_row(
"SELECT count(*) FROM seex_runs WHERE run_id = ?",
[run_id.as_str()],
|row| row.get(0),
)?;
Ok(count > 0)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn summaries_count_all_observations_running_and_finalized()
-> Result<(), Box<dyn std::error::Error>> {
let connection = duckdb::Connection::open_in_memory()?;
connection.execute_batch(
"CREATE SCHEMA dl;
CREATE TABLE dl.metric_points (run_id VARCHAR, metric_key VARCHAR,
step BIGINT, value_f64 DOUBLE);
INSERT INTO dl.metric_points VALUES
('run-1', 'loss', 1, -10.0), ('run-1', 'loss', 1, 10.0),
('run-1', 'loss', 2, 3.0), ('run-1', 'loss', 2, 4.0);
CREATE TABLE seex_runs (run_id VARCHAR, status VARCHAR);
INSERT INTO seex_runs VALUES ('run-1', 'running');
CREATE TABLE seex_metric_aggregates (run_id VARCHAR, metric_key VARCHAR,
effective_count UBIGINT, last_step BIGINT, last_value_f64 DOUBLE,
min_value_f64 DOUBLE, max_value_f64 DOUBLE);",
)?;
let run_id = RunId::from_string("run-1");
let key = MetricKey::from_string("loss");
let reader = crate::storage::ProjectMetricReader::new(&connection);
let running = reader.list_metrics(&run_id, RunStatus::Running)?;
assert_eq!(
running,
reader.query_metric_summaries(std::slice::from_ref(&run_id), &key)?
);
NativeWriteStore::new(&connection).rebuild_metric_aggregates_for_run(&run_id)?;
connection.execute_batch("UPDATE seex_runs SET status = 'finished'")?;
let finalized = reader.list_metrics(&run_id, RunStatus::Finished)?;
let summaries = reader.query_metric_summaries(std::slice::from_ref(&run_id), &key)?;
assert_eq!(finalized, summaries);
assert_eq!(running.len(), 1);
for aggregate in running.iter().chain(&finalized) {
assert_eq!(aggregate.effective_count, 4);
assert_eq!(aggregate.last_step, Step::new(2));
assert!([3.0, 4.0].contains(&aggregate.last_value_f64));
assert_eq!(aggregate.min_value_f64, -10.0);
assert_eq!(aggregate.max_value_f64, 10.0);
}
Ok(())
}
#[test]
fn percent_encode_metric_key_preserves_unreserved_ascii() {
assert_eq!(percent_encode_metric_key("AZaz09._~-"), "AZaz09._~-");
}
#[test]
fn percent_encode_metric_key_encodes_path_and_utf8_bytes() {
assert_eq!(
percent_encode_metric_key("train/loss 零"),
"train%2Floss%20%E9%9B%B6"
);
}
}