eventuary_postgres/
partition_backfill.rs

1use std::collections::HashMap;
2use std::fmt;
3use std::num::NonZeroU16;
4use std::sync::Arc;
5
6use chrono::{DateTime, Utc};
7use sqlx::{PgPool, Row};
8
9use eventuary_core::partition::{PartitionHasher, PartitionKeyResolver, PartitionStrategy};
10use eventuary_core::{Error, Result, SerializedEvent, SerializedPayload};
11
12use crate::event_log::{PgEventLogSchema, PgEventLogSchemaConfig};
13use crate::relation::PgRelationName;
14
15pub struct PgPartitionBackfillConfig {
16    pub events_relation: PgRelationName,
17    pub partition_count: NonZeroU16,
18    pub key_resolver: Arc<dyn PartitionKeyResolver>,
19    pub hasher: Arc<dyn PartitionHasher>,
20    pub batch_size: usize,
21}
22
23impl PgPartitionBackfillConfig {
24    pub fn new(
25        events_relation: PgRelationName,
26        partition_count: NonZeroU16,
27        key_resolver: impl PartitionKeyResolver + 'static,
28        hasher: impl PartitionHasher + 'static,
29        batch_size: usize,
30    ) -> Result<Self> {
31        if batch_size == 0 {
32            return Err(Error::Config(
33                "partition backfill batch size must be greater than zero".to_owned(),
34            ));
35        }
36
37        Ok(Self {
38            events_relation,
39            partition_count,
40            key_resolver: Arc::new(key_resolver),
41            hasher: Arc::new(hasher),
42            batch_size,
43        })
44    }
45}
46
47impl fmt::Debug for PgPartitionBackfillConfig {
48    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
49        f.debug_struct("PgPartitionBackfillConfig")
50            .field("events_relation", &self.events_relation)
51            .field("partition_count", &self.partition_count)
52            .field("batch_size", &self.batch_size)
53            .finish()
54    }
55}
56
57#[derive(Debug, Clone, Default)]
58pub struct BackfillReport {
59    pub rows_updated: u64,
60    pub batches: u32,
61}
62
63pub struct PgPartitionBackfill {
64    pool: PgPool,
65    config: PgPartitionBackfillConfig,
66}
67
68impl PgPartitionBackfill {
69    pub fn new(pool: PgPool, config: PgPartitionBackfillConfig) -> Self {
70        Self { pool, config }
71    }
72
73    pub async fn connect(pool: PgPool, config: PgPartitionBackfillConfig) -> Result<Self> {
74        Self::prepare_schema(&pool, &config).await?;
75        Ok(Self::new(pool, config))
76    }
77
78    pub async fn prepare_schema(pool: &PgPool, config: &PgPartitionBackfillConfig) -> Result<()> {
79        PgEventLogSchema::prepare(
80            pool,
81            &PgEventLogSchemaConfig {
82                events_relation: config.events_relation.clone(),
83            },
84        )
85        .await
86    }
87
88    pub fn schema_sql(config: &PgPartitionBackfillConfig) -> String {
89        PgEventLogSchema::schema_sql(&PgEventLogSchemaConfig {
90            events_relation: config.events_relation.clone(),
91        })
92    }
93
94    pub async fn run(&self) -> Result<BackfillReport> {
95        let events = self.config.events_relation.render();
96        let batch_size = self.config.batch_size;
97        let mut report = BackfillReport::default();
98
99        let fetch_sql = format!(
100            "SELECT sequence, id::text AS id_text, organization, namespace, topic, event_key, \
101             payload::text AS payload_text, content_type, metadata::text AS metadata_text, \
102             timestamp::text AS timestamp_text, version, parent_id::text AS parent_id_text, \
103             correlation_id, causation_id \
104             FROM {events} \
105             WHERE partition_id IS NULL \
106             ORDER BY sequence \
107             LIMIT $1",
108        );
109
110        let update_sql = format!(
111            "UPDATE {events} \
112             SET partition_key = $1, \
113                 partition_hash = $2, \
114                 partition_id = $3, \
115                 partition_count = $4, \
116                 partition_strategy = $5 \
117             WHERE sequence = $6 AND partition_id IS NULL",
118        );
119
120        loop {
121            let rows = sqlx::query(&fetch_sql)
122                .bind(batch_size as i64)
123                .fetch_all(&self.pool)
124                .await
125                .map_err(|e| Error::Store(e.to_string()))?;
126
127            if rows.is_empty() {
128                break;
129            }
130
131            let mut tx = self
132                .pool
133                .begin()
134                .await
135                .map_err(|e| Error::Store(e.to_string()))?;
136
137            for row in &rows {
138                let sequence: i64 = row.get("sequence");
139                let serialized = deserialize_row(row, sequence)?;
140                let event = serialized.to_event()?;
141
142                let partition_key = self.config.key_resolver.partition_key(&event)?;
143                let partition_hash = self.config.hasher.hash(&partition_key);
144                let partition = self
145                    .config
146                    .hasher
147                    .partition_for(&partition_key, self.config.partition_count);
148                let partition_strategy = PartitionStrategy::new(self.config.hasher.strategy())?;
149
150                let result = sqlx::query(&update_sql)
151                    .bind(partition_key.as_str())
152                    .bind(partition_hash.to_sql_i64())
153                    .bind(partition.id() as i32)
154                    .bind(partition.count() as i32)
155                    .bind(partition_strategy.as_str())
156                    .bind(sequence)
157                    .execute(&mut *tx)
158                    .await
159                    .map_err(|e| Error::Store(e.to_string()))?;
160
161                report.rows_updated += result.rows_affected();
162            }
163
164            tx.commit().await.map_err(|e| Error::Store(e.to_string()))?;
165            report.batches += 1;
166        }
167
168        Ok(report)
169    }
170}
171
172fn deserialize_row(row: &sqlx::postgres::PgRow, sequence: i64) -> Result<SerializedEvent> {
173    let id_text: String = row.get("id_text");
174    let id = uuid::Uuid::parse_str(&id_text)
175        .map_err(|e| Error::Serialization(format!("decode id: {e}")))?;
176    let parent_id = row
177        .get::<Option<String>, _>("parent_id_text")
178        .as_deref()
179        .map(uuid::Uuid::parse_str)
180        .transpose()
181        .map_err(|e| Error::Serialization(format!("decode parent_id: {e}")))?;
182    let payload_str: String = row.get("payload_text");
183    let payload: SerializedPayload = serde_json::from_str(&payload_str)
184        .map_err(|e| Error::Serialization(format!("decode payload: {e}")))?;
185    let metadata_str: String = row.get("metadata_text");
186    let metadata: HashMap<String, String> = serde_json::from_str(&metadata_str)
187        .map_err(|e| Error::Serialization(format!("decode metadata: {e}")))?;
188    let timestamp_str: String = row.get("timestamp_text");
189    let timestamp = parse_pg_timestamp(&timestamp_str).map_err(|e| {
190        Error::Serialization(format!("decode timestamp at sequence {sequence}: {e}"))
191    })?;
192    Ok(SerializedEvent {
193        id,
194        organization: row.get("organization"),
195        namespace: row.get("namespace"),
196        topic: row.get("topic"),
197        payload,
198        metadata,
199        timestamp,
200        version: row.get::<i64, _>("version") as u64,
201        key: row.get("event_key"),
202        parent_id,
203        correlation_id: row.get("correlation_id"),
204        causation_id: row.get("causation_id"),
205    })
206}
207
208fn parse_pg_timestamp(s: &str) -> std::result::Result<DateTime<Utc>, chrono::ParseError> {
209    if let Ok(dt) = DateTime::parse_from_rfc3339(s) {
210        return Ok(dt.with_timezone(&Utc));
211    }
212    DateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S%.f%#z").map(|dt| dt.with_timezone(&Utc))
213}
eventuary_postgres/partition_backfill.rs

eventuary_postgres/
partition_backfill.rs
