use std::time::Duration;
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use scylla::frame::response::result::{CqlValue, Row};
use uuid::Uuid;
use super::cassandra::{CassandraCanonicalStore, now_unix_ms};
use super::system_store::{
DeadLetterGroup, PendingTaskMetric, ProjectionClaimFilter, ProjectionOperation,
ProjectionTaskInsert, ProjectionTaskRow, ProjectionTaskStatus, ProjectionTaskStore,
ProjectionTaskSummary, SystemStoreError, SystemStoreResult,
};
pub(super) fn col(row: &Row, idx: usize) -> Option<&CqlValue> {
row.columns.get(idx).and_then(|c| c.as_ref())
}
pub(super) fn get_text(row: &Row, idx: usize) -> String {
col(row, idx)
.and_then(|v| v.as_text())
.cloned()
.unwrap_or_default()
}
pub(super) fn get_i32(row: &Row, idx: usize) -> i32 {
col(row, idx)
.and_then(|v| v.as_int().or_else(|| v.as_bigint().map(|b| b as i32)))
.unwrap_or(0)
}
pub(super) fn get_i64(row: &Row, idx: usize) -> i64 {
col(row, idx)
.and_then(|v| v.as_bigint().or_else(|| v.as_int().map(i64::from)))
.unwrap_or(0)
}
pub(super) fn get_dt(row: &Row, idx: usize) -> DateTime<Utc> {
get_opt_dt(row, idx).unwrap_or_else(Utc::now)
}
pub(super) fn get_opt_dt(row: &Row, idx: usize) -> Option<DateTime<Utc>> {
match col(row, idx) {
Some(v) => {
let millis = v
.as_cql_timestamp()
.map(|t| t.0)
.or_else(|| v.as_bigint())?;
DateTime::<Utc>::from_timestamp_millis(millis)
}
None => None,
}
}
pub(super) fn get_uuid(row: &Row, idx: usize) -> SystemStoreResult<Uuid> {
let raw = col(row, idx).and_then(|v| v.as_text()).ok_or_else(|| {
SystemStoreError::InvalidInput(format!("cassandra row missing text uuid at column {idx}"))
})?;
Uuid::parse_str(raw).map_err(|e| {
SystemStoreError::InvalidInput(format!(
"cassandra row column {idx} is not a uuid '{raw}': {e}"
))
})
}
pub(super) fn get_json(row: &Row, idx: usize, default: serde_json::Value) -> serde_json::Value {
match col(row, idx).and_then(|v| v.as_text()) {
Some(s) => serde_json::from_str(s).unwrap_or(default),
None => default,
}
}
pub(super) fn cass_err(op: &str, err: impl std::fmt::Display) -> SystemStoreError {
SystemStoreError::Io {
backend: "cassandra",
source: format!("{op}: {err}"),
}
}
pub(super) fn cql_ts(millis: i64) -> scylla::frame::value::CqlTimestamp {
scylla::frame::value::CqlTimestamp(millis)
}
const TASK_COLS: &str = "task_id, idempotency_key, project_id, target_backend, target_instance, \
projection_kind, resource_name, operation, source_row_key, target_options, \
source_payload, source_checksum, status, retry_count, last_error, \
created_at, updated_at, next_retry_at, completed_at";
fn row_to_projection_task(row: &Row) -> SystemStoreResult<ProjectionTaskRow> {
let task_id = get_uuid(row, 0)?;
let operation_str = get_text(row, 7);
let operation = ProjectionOperation::parse(&operation_str).ok_or_else(|| {
SystemStoreError::InvalidInput(format!(
"unknown projection operation '{operation_str}' in cassandra row"
))
})?;
let status_str = get_text(row, 12);
let status = ProjectionTaskStatus::parse(&status_str).ok_or_else(|| {
SystemStoreError::InvalidInput(format!(
"unknown projection status '{status_str}' in cassandra row"
))
})?;
Ok(ProjectionTaskRow {
task_id,
idempotency_key: get_text(row, 1),
project_id: get_text(row, 2),
target_backend: get_text(row, 3),
target_instance: get_text(row, 4),
projection_kind: get_text(row, 5),
resource_name: get_text(row, 6),
operation,
source_row_key: get_json(row, 8, serde_json::Value::Null),
target_options: get_json(row, 9, serde_json::Value::Null),
source_payload: get_json(row, 10, serde_json::Value::Null),
source_checksum: get_text(row, 11),
status,
retry_count: get_i32(row, 13),
last_error: get_text(row, 14),
created_at: get_dt(row, 15),
updated_at: get_dt(row, 16),
next_retry_at: get_opt_dt(row, 17),
completed_at: get_opt_dt(row, 18),
})
}
impl CassandraCanonicalStore {
fn projection_table(&self) -> String {
self.qualified("udb_projection_tasks")
}
fn idem_table(&self) -> String {
self.qualified("udb_projection_idem")
}
}
#[async_trait]
impl ProjectionTaskStore for CassandraCanonicalStore {
fn backend_label(&self) -> &'static str {
"cassandra"
}
async fn ensure_projection_tables(&self) -> SystemStoreResult<()> {
self.ensure_keyspace()
.await
.map_err(|e| cass_err("ensure_projection_tables keyspace", e))?;
let tasks_ddl = format!(
"CREATE TABLE IF NOT EXISTS {tbl} ( \
task_id text PRIMARY KEY, \
idempotency_key text, \
project_id text, \
manifest_checksum text, \
message_type text, \
source_schema text, \
source_table text, \
source_row_key text, \
operation text, \
target_backend text, \
target_instance text, \
projection_kind text, \
resource_name text, \
target_options text, \
source_payload text, \
source_checksum text, \
status text, \
retry_count int, \
last_error text, \
created_at timestamp, \
updated_at timestamp, \
next_retry_at timestamp, \
completed_at timestamp \
)",
tbl = self.projection_table(),
);
self.client()
.cql_execute(&tasks_ddl, ())
.await
.map_err(|e| cass_err("ensure_projection_tables tasks", e))?;
let idem_ddl = format!(
"CREATE TABLE IF NOT EXISTS {tbl} ( idempotency_key text PRIMARY KEY, task_id text )",
tbl = self.idem_table(),
);
self.client()
.cql_execute(&idem_ddl, ())
.await
.map_err(|e| cass_err("ensure_projection_tables idem", e))?;
Ok(())
}
async fn enqueue_projection_task(
&self,
task: &ProjectionTaskInsert,
) -> SystemStoreResult<Uuid> {
use scylla::statement::SerialConsistency;
let new_id = Uuid::new_v4();
let idem_insert = format!(
"INSERT INTO {tbl} (idempotency_key, task_id) VALUES (?, ?) IF NOT EXISTS",
tbl = self.idem_table(),
);
let applied = self
.client()
.cql_lwt_applied(
&idem_insert,
(task.idempotency_key.as_str(), new_id.to_string()),
SerialConsistency::Serial,
)
.await
.map_err(|e| cass_err("enqueue_projection_task idem LWT", e))?;
if !applied {
let lookup = format!(
"SELECT task_id FROM {tbl} WHERE idempotency_key = ?",
tbl = self.idem_table(),
);
let rows = self
.client()
.cql_query_rows(&lookup, (task.idempotency_key.as_str(),))
.await
.map_err(|e| cass_err("enqueue_projection_task idem lookup", e))?;
let row = rows.first().ok_or_else(|| {
cass_err(
"enqueue_projection_task",
"idempotency row vanished after not-applied LWT",
)
})?;
return get_uuid(row, 0);
}
let now = now_unix_ms();
let insert = format!(
"INSERT INTO {tbl} ( \
task_id, idempotency_key, project_id, manifest_checksum, message_type, \
source_schema, source_table, source_row_key, operation, \
target_backend, target_instance, projection_kind, resource_name, \
target_options, source_payload, source_checksum, \
status, retry_count, last_error, created_at, updated_at \
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
tbl = self.projection_table(),
);
use scylla::frame::response::result::CqlValue;
let params: Vec<CqlValue> = vec![
CqlValue::Text(new_id.to_string()),
CqlValue::Text(task.idempotency_key.clone()),
CqlValue::Text(task.project_id.clone()),
CqlValue::Text(task.manifest_checksum.clone()),
CqlValue::Text(task.message_type.clone()),
CqlValue::Text(task.source_schema.clone()),
CqlValue::Text(task.source_table.clone()),
CqlValue::Text(task.source_row_key.to_string()),
CqlValue::Text(task.operation.as_str().to_string()),
CqlValue::Text(task.target_backend.clone()),
CqlValue::Text(task.target_instance.clone()),
CqlValue::Text(task.projection_kind.clone()),
CqlValue::Text(task.resource_name.clone()),
CqlValue::Text(task.target_options.to_string()),
CqlValue::Text(task.source_payload.to_string()),
CqlValue::Text(task.source_checksum.clone()),
CqlValue::Text(ProjectionTaskStatus::Pending.as_str().to_string()),
CqlValue::Int(0),
CqlValue::Text(String::new()),
CqlValue::Timestamp(cql_ts(now)),
CqlValue::Timestamp(cql_ts(now)),
];
self.client()
.cql_execute(&insert, params)
.await
.map_err(|e| cass_err("enqueue_projection_task insert", e))?;
Ok(new_id)
}
async fn claim_projection_tasks(
&self,
filter: &ProjectionClaimFilter,
) -> SystemStoreResult<Vec<ProjectionTaskRow>> {
use scylla::statement::SerialConsistency;
if filter.batch_size <= 0 {
return Ok(Vec::new());
}
let scan = format!(
"SELECT {TASK_COLS} FROM {tbl} ALLOW FILTERING",
tbl = self.projection_table(),
);
let rows = self
.client()
.cql_query_rows(&scan, ())
.await
.map_err(|e| cass_err("claim_projection_tasks scan", e))?;
let now = Utc::now();
let mut candidates: Vec<ProjectionTaskRow> = Vec::new();
for row in &rows {
let task = row_to_projection_task(row)?;
if !task.status.is_claimable() {
continue;
}
if task.retry_count >= filter.max_retries {
continue;
}
if let Some(p) = &filter.project_id {
if &task.project_id != p {
continue;
}
}
if let Some(b) = &filter.target_backend {
if &task.target_backend != b {
continue;
}
}
if let Some(i) = &filter.target_instance {
if &task.target_instance != i {
continue;
}
}
if task.status == ProjectionTaskStatus::Failed {
if let Some(nra) = task.next_retry_at {
if nra > now {
continue;
}
}
}
candidates.push(task);
}
candidates.sort_by_key(|t| t.created_at);
let claim_now = now_unix_ms();
let mut out = Vec::new();
for mut task in candidates {
if out.len() as i64 >= filter.batch_size {
break;
}
let cas = format!(
"UPDATE {tbl} SET status = ?, updated_at = ? WHERE task_id = ? \
IF status IN ('PENDING','FAILED')",
tbl = self.projection_table(),
);
let applied = self
.client()
.cql_lwt_applied(
&cas,
(
ProjectionTaskStatus::InProgress.as_str(),
cql_ts(claim_now),
task.task_id.to_string(),
),
SerialConsistency::Serial,
)
.await
.map_err(|e| cass_err("claim_projection_tasks LWT", e))?;
if applied {
task.status = ProjectionTaskStatus::InProgress;
task.updated_at =
DateTime::<Utc>::from_timestamp_millis(claim_now).unwrap_or_else(Utc::now);
out.push(task);
}
}
Ok(out)
}
async fn mark_projection_task_completed(&self, task_id: Uuid) -> SystemStoreResult<()> {
let now = now_unix_ms();
let sql = format!(
"UPDATE {tbl} SET status = ?, completed_at = ?, next_retry_at = ?, updated_at = ? \
WHERE task_id = ?",
tbl = self.projection_table(),
);
self.client()
.cql_execute(
&sql,
(
ProjectionTaskStatus::Completed.as_str(),
cql_ts(now),
None::<scylla::frame::value::CqlTimestamp>,
cql_ts(now),
task_id.to_string(),
),
)
.await
.map_err(|e| cass_err("mark_projection_task_completed", e))?;
Ok(())
}
async fn mark_projection_task_failed(
&self,
task_id: Uuid,
new_retry_count: i32,
new_status: ProjectionTaskStatus,
error: &str,
) -> SystemStoreResult<()> {
if !matches!(
new_status,
ProjectionTaskStatus::Failed | ProjectionTaskStatus::DeadLetter
) {
return Err(SystemStoreError::InvalidInput(format!(
"mark_projection_task_failed only accepts FAILED or DEAD_LETTER, got {}",
new_status.as_str()
)));
}
let now = now_unix_ms();
let sql = format!(
"UPDATE {tbl} SET status = ?, retry_count = ?, last_error = ?, \
next_retry_at = ?, updated_at = ? WHERE task_id = ?",
tbl = self.projection_table(),
);
self.client()
.cql_execute(
&sql,
(
new_status.as_str(),
new_retry_count,
error,
None::<scylla::frame::value::CqlTimestamp>,
cql_ts(now),
task_id.to_string(),
),
)
.await
.map_err(|e| cass_err("mark_projection_task_failed", e))?;
Ok(())
}
async fn requeue_dead_letter_tasks(
&self,
target_backend: Option<&str>,
) -> SystemStoreResult<i64> {
let scan = format!(
"SELECT {TASK_COLS} FROM {tbl} ALLOW FILTERING",
tbl = self.projection_table(),
);
let rows = self
.client()
.cql_query_rows(&scan, ())
.await
.map_err(|e| cass_err("requeue_dead_letter_tasks scan", e))?;
let now = now_unix_ms();
let mut n = 0i64;
for row in &rows {
let task = row_to_projection_task(row)?;
if task.status != ProjectionTaskStatus::DeadLetter {
continue;
}
if let Some(b) = target_backend {
if &task.target_backend != b {
continue;
}
}
let upd = format!(
"UPDATE {tbl} SET status = ?, retry_count = ?, last_error = ?, \
next_retry_at = ?, updated_at = ? WHERE task_id = ?",
tbl = self.projection_table(),
);
self.client()
.cql_execute(
&upd,
(
ProjectionTaskStatus::Pending.as_str(),
0_i32,
"",
None::<scylla::frame::value::CqlTimestamp>,
cql_ts(now),
task.task_id.to_string(),
),
)
.await
.map_err(|e| cass_err("requeue_dead_letter_tasks update", e))?;
n += 1;
}
Ok(n)
}
async fn reset_stale_in_progress_tasks(&self, stale_after: Duration) -> SystemStoreResult<i64> {
let scan = format!(
"SELECT {TASK_COLS} FROM {tbl} ALLOW FILTERING",
tbl = self.projection_table(),
);
let rows = self
.client()
.cql_query_rows(&scan, ())
.await
.map_err(|e| cass_err("reset_stale_in_progress_tasks scan", e))?;
let cutoff = Utc::now()
- chrono::Duration::from_std(stale_after).unwrap_or_else(|_| chrono::Duration::zero());
let now = now_unix_ms();
let mut n = 0i64;
for row in &rows {
let task = row_to_projection_task(row)?;
if task.status != ProjectionTaskStatus::InProgress {
continue;
}
if task.updated_at >= cutoff {
continue;
}
let upd = format!(
"UPDATE {tbl} SET status = ?, last_error = ?, updated_at = ? WHERE task_id = ?",
tbl = self.projection_table(),
);
self.client()
.cql_execute(
&upd,
(
ProjectionTaskStatus::Pending.as_str(),
"stale in-progress reconciliation",
cql_ts(now),
task.task_id.to_string(),
),
)
.await
.map_err(|e| cass_err("reset_stale_in_progress_tasks update", e))?;
n += 1;
}
Ok(n)
}
async fn pending_task_metrics(&self, limit: i64) -> SystemStoreResult<Vec<PendingTaskMetric>> {
let scan = format!(
"SELECT {TASK_COLS} FROM {tbl} ALLOW FILTERING",
tbl = self.projection_table(),
);
let rows = self
.client()
.cql_query_rows(&scan, ())
.await
.map_err(|e| cass_err("pending_task_metrics scan", e))?;
use std::collections::HashMap;
let mut groups: HashMap<(String, String, String, String), (i64, DateTime<Utc>)> =
HashMap::new();
for row in &rows {
let task = row_to_projection_task(row)?;
if !matches!(
task.status,
ProjectionTaskStatus::Pending | ProjectionTaskStatus::Failed
) {
continue;
}
let key = (
task.project_id.clone(),
task.target_backend.clone(),
task.target_instance.clone(),
task.projection_kind.clone(),
);
let entry = groups.entry(key).or_insert((0, task.created_at));
entry.0 += 1;
if task.created_at < entry.1 {
entry.1 = task.created_at;
}
}
let now = Utc::now();
let mut out: Vec<PendingTaskMetric> = groups
.into_iter()
.map(
|((project_id, target_backend, target_instance, projection_kind), (n, oldest))| {
let age = (now - oldest).num_milliseconds() as f64 / 1000.0;
PendingTaskMetric {
project_id,
target_backend,
target_instance,
projection_kind,
pending: n,
oldest_age_seconds: age.max(0.0),
}
},
)
.collect();
out.truncate(limit.max(1) as usize);
Ok(out)
}
async fn dead_letter_groups(&self, limit: i64) -> SystemStoreResult<Vec<DeadLetterGroup>> {
let scan = format!(
"SELECT source_table, target_backend, target_instance, status \
FROM {tbl} ALLOW FILTERING",
tbl = self.projection_table(),
);
let rows = self
.client()
.cql_query_rows(&scan, ())
.await
.map_err(|e| cass_err("dead_letter_groups scan", e))?;
use std::collections::HashMap;
let mut groups: HashMap<(String, String, String), i64> = HashMap::new();
for row in &rows {
if get_text(row, 3) != ProjectionTaskStatus::DeadLetter.as_str() {
continue;
}
*groups
.entry((get_text(row, 0), get_text(row, 1), get_text(row, 2)))
.or_insert(0) += 1;
}
let mut out: Vec<DeadLetterGroup> = groups
.into_iter()
.map(
|((source_table, target_backend, target_instance), dead_count)| DeadLetterGroup {
source_table,
target_backend,
target_instance,
dead_count,
},
)
.collect();
out.truncate(limit.max(1) as usize);
Ok(out)
}
async fn requeue_dead_letter_by_source(
&self,
source_table: &str,
target_backend: &str,
target_instance: &str,
) -> SystemStoreResult<i64> {
let scan = format!(
"SELECT task_id, source_table, target_backend, target_instance, status \
FROM {tbl} ALLOW FILTERING",
tbl = self.projection_table(),
);
let rows = self
.client()
.cql_query_rows(&scan, ())
.await
.map_err(|e| cass_err("requeue_dead_letter_by_source scan", e))?;
let now = now_unix_ms();
let mut n = 0i64;
for row in &rows {
let status = get_text(row, 4);
if status != ProjectionTaskStatus::DeadLetter.as_str() {
continue;
}
if get_text(row, 1) != source_table
|| get_text(row, 2) != target_backend
|| get_text(row, 3) != target_instance
{
continue;
}
let task_id = get_text(row, 0);
let upd = format!(
"UPDATE {tbl} SET status = ?, retry_count = ?, last_error = ?, \
next_retry_at = ?, updated_at = ? WHERE task_id = ?",
tbl = self.projection_table(),
);
self.client()
.cql_execute(
&upd,
(
ProjectionTaskStatus::Pending.as_str(),
0_i32,
"reconciliation repair",
None::<scylla::frame::value::CqlTimestamp>,
cql_ts(now),
task_id,
),
)
.await
.map_err(|e| cass_err("requeue_dead_letter_by_source update", e))?;
n += 1;
}
Ok(n)
}
async fn pending_projection_task_count(
&self,
idempotency_keys: &[String],
) -> SystemStoreResult<i64> {
if idempotency_keys.is_empty() {
return Ok(0);
}
let mut pending = 0i64;
for key in idempotency_keys {
let lookup = format!(
"SELECT task_id FROM {tbl} WHERE idempotency_key = ?",
tbl = self.idem_table(),
);
let idem_rows = self
.client()
.cql_query_rows(&lookup, (key.as_str(),))
.await
.map_err(|e| cass_err("pending_projection_task_count lookup", e))?;
let Some(idem_row) = idem_rows.first() else {
continue;
};
let task_id = get_text(idem_row, 0);
let status_sql = format!(
"SELECT status FROM {tbl} WHERE task_id = ?",
tbl = self.projection_table(),
);
let status_rows = self
.client()
.cql_query_rows(&status_sql, (task_id,))
.await
.map_err(|e| cass_err("pending_projection_task_count status", e))?;
let Some(status_row) = status_rows.first() else {
continue;
};
let status = get_text(status_row, 0);
let settled = matches!(
ProjectionTaskStatus::parse(&status),
Some(ProjectionTaskStatus::Completed)
| Some(ProjectionTaskStatus::DeadLetter)
| Some(ProjectionTaskStatus::Failed)
);
if !settled {
pending += 1;
}
}
Ok(pending)
}
async fn projection_task_summary(&self) -> SystemStoreResult<ProjectionTaskSummary> {
let scan = format!(
"SELECT status FROM {tbl} ALLOW FILTERING",
tbl = self.projection_table(),
);
let rows = self
.client()
.cql_query_rows(&scan, ())
.await
.map_err(|e| cass_err("projection_task_summary scan", e))?;
let mut s = ProjectionTaskSummary::default();
for row in &rows {
let status = get_text(row, 0);
match ProjectionTaskStatus::parse(&status) {
Some(ProjectionTaskStatus::Pending) => s.pending += 1,
Some(ProjectionTaskStatus::InProgress) => s.in_progress += 1,
Some(ProjectionTaskStatus::Completed) => s.completed += 1,
Some(ProjectionTaskStatus::Failed) => s.failed += 1,
Some(ProjectionTaskStatus::DeadLetter) => s.dead_letter += 1,
None => {
return Err(SystemStoreError::InvalidInput(format!(
"unknown projection status '{status}' in cassandra summary"
)));
}
}
}
Ok(s)
}
}