use std::collections::HashMap;
use std::fmt;
use std::num::NonZeroU32;
use std::sync::Arc;
use chrono::{DateTime, Utc};
use rusqlite::types::Value;
use eventuary_core::partition::{PartitionHasher, PartitionKeyResolver, PartitionStrategy};
use eventuary_core::{Error, Result, SerializedEvent, SerializedPayload};
use crate::database::SqliteConn;
use crate::event_log::{SqliteEventLogSchema, SqliteEventLogSchemaConfig};
use crate::relation::SqliteRelationName;
pub struct SqlitePartitionBackfillConfig {
pub events_relation: SqliteRelationName,
pub partition_count: NonZeroU32,
pub key_resolver: Arc<dyn PartitionKeyResolver>,
pub hasher: Arc<dyn PartitionHasher>,
pub batch_size: usize,
}
impl SqlitePartitionBackfillConfig {
pub fn new(
events_relation: SqliteRelationName,
partition_count: NonZeroU32,
key_resolver: impl PartitionKeyResolver + 'static,
hasher: impl PartitionHasher + 'static,
batch_size: usize,
) -> Result<Self> {
if batch_size == 0 {
return Err(Error::Config(
"partition backfill batch size must be greater than zero".to_owned(),
));
}
Ok(Self {
events_relation,
partition_count,
key_resolver: Arc::new(key_resolver),
hasher: Arc::new(hasher),
batch_size,
})
}
}
impl fmt::Debug for SqlitePartitionBackfillConfig {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("SqlitePartitionBackfillConfig")
.field("events_relation", &self.events_relation)
.field("partition_count", &self.partition_count)
.field("batch_size", &self.batch_size)
.finish()
}
}
#[derive(Debug, Clone, Default)]
pub struct BackfillReport {
pub rows_updated: u64,
pub batches: u32,
}
pub struct SqlitePartitionBackfill {
conn: SqliteConn,
config: SqlitePartitionBackfillConfig,
}
impl SqlitePartitionBackfill {
pub fn new(conn: SqliteConn, config: SqlitePartitionBackfillConfig) -> Self {
Self { conn, config }
}
pub fn connect(conn: SqliteConn, config: SqlitePartitionBackfillConfig) -> Result<Self> {
Self::prepare_schema(&conn, &config)?;
Ok(Self::new(conn, config))
}
pub fn prepare_schema(conn: &SqliteConn, config: &SqlitePartitionBackfillConfig) -> Result<()> {
let guard = conn.lock().map_err(|e| Error::Store(e.to_string()))?;
SqliteEventLogSchema::prepare(
&guard,
&SqliteEventLogSchemaConfig {
events_relation: config.events_relation.clone(),
},
)
}
pub fn schema_sql(config: &SqlitePartitionBackfillConfig) -> String {
SqliteEventLogSchema::schema_sql(&SqliteEventLogSchemaConfig {
events_relation: config.events_relation.clone(),
})
}
pub async fn run(&self) -> Result<BackfillReport> {
let events = self.config.events_relation.render();
let batch_size = self.config.batch_size;
let partition_count = self.config.partition_count;
let key_resolver = Arc::clone(&self.config.key_resolver);
let hasher = Arc::clone(&self.config.hasher);
let conn = Arc::clone(&self.conn);
let fetch_sql = format!(
"SELECT sequence, id, organization, namespace, topic, event_key, payload, \
content_type, metadata, timestamp, version, parent_id, correlation_id, causation_id \
FROM {events} \
WHERE partition_id IS NULL \
ORDER BY sequence \
LIMIT ?1"
);
let update_sql = format!(
"UPDATE {events} \
SET partition_key = ?1, \
partition_hash = ?2, \
partition_id = ?3, \
partition_count = ?4, \
partition_strategy = ?5 \
WHERE sequence = ?6 AND partition_id IS NULL"
);
tokio::task::spawn_blocking(move || {
let mut report = BackfillReport::default();
let mut guard = conn.lock().map_err(|e| Error::Store(e.to_string()))?;
loop {
let rows = {
let mut stmt = guard
.prepare(&fetch_sql)
.map_err(|e| Error::Store(e.to_string()))?;
let mapped = stmt
.query_map([Value::Integer(batch_size as i64)], decode_row)
.map_err(|e| Error::Store(e.to_string()))?;
let mut out: Vec<(i64, SerializedEvent)> = Vec::new();
for row in mapped {
out.push(row.map_err(|e| Error::Store(e.to_string()))?);
}
out
};
if rows.is_empty() {
break;
}
let tx = guard
.transaction()
.map_err(|e| Error::Store(e.to_string()))?;
for (sequence, serialized) in &rows {
let event = serialized.to_event()?;
let partition_key = key_resolver.partition_key(&event)?;
let partition_hash = hasher.hash(&partition_key);
let partition = hasher.partition_for(&partition_key, partition_count);
let partition_strategy = PartitionStrategy::new(hasher.strategy())?;
let updated = tx
.execute(
&update_sql,
rusqlite::params![
partition_key.as_str(),
partition_hash.to_sql_i64(),
partition.id() as i64,
partition.count() as i64,
partition_strategy.as_str(),
*sequence,
],
)
.map_err(|e| Error::Store(e.to_string()))?;
report.rows_updated += updated as u64;
}
tx.commit().map_err(|e| Error::Store(e.to_string()))?;
report.batches += 1;
}
Ok(report)
})
.await
.map_err(|e| Error::Store(format!("blocking task panicked: {e}")))?
}
}
fn decode_row(row: &rusqlite::Row<'_>) -> rusqlite::Result<(i64, SerializedEvent)> {
let sequence: i64 = row.get(0)?;
let id: String = row.get(1)?;
let organization: String = row.get(2)?;
let namespace: String = row.get(3)?;
let topic: String = row.get(4)?;
let key: String = row.get(5)?;
let payload_str: String = row.get(6)?;
let _content_type: String = row.get(7)?;
let metadata_str: String = row.get(8)?;
let timestamp_str: String = row.get(9)?;
let version: i64 = row.get(10)?;
let parent_id: Option<String> = row.get(11)?;
let correlation_id: Option<String> = row.get(12)?;
let causation_id: Option<String> = row.get(13)?;
let id = uuid::Uuid::parse_str(&id).map_err(|e| {
rusqlite::Error::FromSqlConversionFailure(
1,
rusqlite::types::Type::Text,
Box::new(std::io::Error::other(format!("decode id: {e}"))),
)
})?;
let parent_id = parent_id
.as_deref()
.map(uuid::Uuid::parse_str)
.transpose()
.map_err(|e| {
rusqlite::Error::FromSqlConversionFailure(
11,
rusqlite::types::Type::Text,
Box::new(std::io::Error::other(format!("decode parent_id: {e}"))),
)
})?;
let payload: SerializedPayload = serde_json::from_str(&payload_str).map_err(|e| {
rusqlite::Error::FromSqlConversionFailure(
6,
rusqlite::types::Type::Text,
Box::new(std::io::Error::other(format!("decode payload: {e}"))),
)
})?;
let metadata: HashMap<String, String> = serde_json::from_str(&metadata_str).map_err(|e| {
rusqlite::Error::FromSqlConversionFailure(
8,
rusqlite::types::Type::Text,
Box::new(std::io::Error::other(format!("decode metadata: {e}"))),
)
})?;
let timestamp = DateTime::parse_from_rfc3339(×tamp_str)
.map(|d| d.with_timezone(&Utc))
.map_err(|e| {
rusqlite::Error::FromSqlConversionFailure(
9,
rusqlite::types::Type::Text,
Box::new(std::io::Error::other(format!("decode timestamp: {e}"))),
)
})?;
Ok((
sequence,
SerializedEvent {
id,
organization,
namespace,
topic,
payload,
metadata,
timestamp,
version: version as u64,
key,
parent_id,
correlation_id,
causation_id,
},
))
}