1use std::fmt;
2use std::num::NonZeroU32;
3use std::sync::Arc;
4
5use sqlx::PgPool;
6
7use eventuary_core::io::Writer;
8use eventuary_core::partition::{
9 PartitionHash, PartitionHasher, PartitionKey, PartitionKeyResolver, PartitionStrategy,
10};
11use eventuary_core::{Error, Event, Result, SerializedEvent};
12
13use crate::event_log::{PgEventLogSchema, PgEventLogSchemaConfig};
14use crate::relation::PgRelationName;
15
16#[derive(Clone, Default)]
17pub enum PgPartitioningConfig {
18 #[default]
19 Off,
20 Inline {
21 partition_count: NonZeroU32,
22 key_resolver: Arc<dyn PartitionKeyResolver>,
23 hasher: Arc<dyn PartitionHasher>,
24 },
25}
26
27impl PgPartitioningConfig {
28 pub fn inline(
29 count: NonZeroU32,
30 resolver: impl PartitionKeyResolver + 'static,
31 hasher: impl PartitionHasher + 'static,
32 ) -> Self {
33 Self::Inline {
34 partition_count: count,
35 key_resolver: Arc::new(resolver),
36 hasher: Arc::new(hasher),
37 }
38 }
39}
40
41impl fmt::Debug for PgPartitioningConfig {
42 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
43 match self {
44 Self::Off => write!(f, "PgPartitioningConfig::Off"),
45 Self::Inline {
46 partition_count, ..
47 } => f
48 .debug_struct("PgPartitioningConfig::Inline")
49 .field("partition_count", partition_count)
50 .finish(),
51 }
52 }
53}
54
55#[derive(Debug, Clone)]
56pub struct PgWriterConfig {
57 pub events_relation: PgRelationName,
58 pub partitioning: PgPartitioningConfig,
59}
60
61impl Default for PgWriterConfig {
62 fn default() -> Self {
63 Self {
64 events_relation: PgRelationName::new("events").expect("default events relation"),
65 partitioning: PgPartitioningConfig::Off,
66 }
67 }
68}
69
70pub struct PgWriter {
71 pool: PgPool,
72 insert_sql: String,
73 partitioning: PgPartitioningConfig,
74}
75
76impl PgWriter {
77 pub fn new(pool: PgPool) -> Self {
78 Self::new_with_config(pool, PgWriterConfig::default())
79 }
80
81 pub async fn connect(pool: PgPool, config: PgWriterConfig) -> Result<Self> {
82 Self::prepare_schema(&pool, &config).await?;
83 Ok(Self::new_with_config(pool, config))
84 }
85
86 pub async fn prepare_schema(pool: &PgPool, config: &PgWriterConfig) -> Result<()> {
87 PgEventLogSchema::prepare(
88 pool,
89 &PgEventLogSchemaConfig {
90 events_relation: config.events_relation.clone(),
91 },
92 )
93 .await
94 }
95
96 pub fn schema_sql(config: &PgWriterConfig) -> String {
97 PgEventLogSchema::schema_sql(&PgEventLogSchemaConfig {
98 events_relation: config.events_relation.clone(),
99 })
100 }
101
102 pub fn new_with_config(pool: PgPool, config: PgWriterConfig) -> Self {
103 let insert_sql = format!(
104 "INSERT INTO {events} \
105 (id, organization, namespace, topic, event_key, payload, content_type, metadata, \
106 timestamp, version, parent_id, correlation_id, causation_id, \
107 partition_key, partition_hash, partition_id, partition_count, partition_strategy) \
108 VALUES \
109 ($1::uuid, $2, $3, $4, $5, $6::jsonb, $7, $8::jsonb, $9::timestamptz, $10, \
110 $11::uuid, $12, $13, $14, $15, $16, $17, $18)",
111 events = config.events_relation.render(),
112 );
113 Self {
114 pool,
115 insert_sql,
116 partitioning: config.partitioning,
117 }
118 }
119
120 fn partition_data(&self, event: &Event) -> Result<PartitionData> {
121 match &self.partitioning {
122 PgPartitioningConfig::Off => Ok(PartitionData::default()),
123 PgPartitioningConfig::Inline {
124 partition_count,
125 key_resolver,
126 hasher,
127 } => {
128 let partition_key = key_resolver.partition_key(event)?;
129 let partition_hash = hasher.hash(&partition_key);
130 let partition = hasher.partition_for(&partition_key, *partition_count);
131 let partition_strategy = PartitionStrategy::new(hasher.strategy())?;
132 Ok(PartitionData {
133 partition_key: Some(partition_key),
134 partition_hash: Some(partition_hash),
135 partition_id: Some(partition.id() as i64),
136 partition_count: Some(partition.count() as i64),
137 partition_strategy: Some(partition_strategy),
138 })
139 }
140 }
141 }
142}
143
144impl Writer for PgWriter {
145 async fn write(&self, event: &Event) -> Result<()> {
146 let row = EventRow::from_event(event)?;
147 let pd = self.partition_data(event)?;
148
149 sqlx::query(&self.insert_sql)
150 .bind(&row.id)
151 .bind(&row.organization)
152 .bind(&row.namespace)
153 .bind(&row.topic)
154 .bind(&row.key)
155 .bind(&row.payload)
156 .bind(&row.content_type)
157 .bind(&row.metadata)
158 .bind(&row.timestamp)
159 .bind(row.version)
160 .bind(&row.parent_id)
161 .bind(&row.correlation_id)
162 .bind(&row.causation_id)
163 .bind(pd.partition_key.as_ref().map(|k| k.as_str()))
164 .bind(pd.partition_hash.map(|h| h.to_sql_i64()))
165 .bind(pd.partition_id)
166 .bind(pd.partition_count)
167 .bind(pd.partition_strategy.as_ref().map(|s| s.as_str()))
168 .execute(&self.pool)
169 .await
170 .map_err(|e| Error::Store(e.to_string()))?;
171
172 Ok(())
173 }
174
175 async fn write_all(&self, events: &[Event]) -> Result<()> {
176 if events.is_empty() {
177 return Ok(());
178 }
179 let mut tx = self
180 .pool
181 .begin()
182 .await
183 .map_err(|e| Error::Store(e.to_string()))?;
184 for event in events {
185 let row = EventRow::from_event(event)?;
186 let pd = self.partition_data(event)?;
187 sqlx::query(&self.insert_sql)
188 .bind(&row.id)
189 .bind(&row.organization)
190 .bind(&row.namespace)
191 .bind(&row.topic)
192 .bind(&row.key)
193 .bind(&row.payload)
194 .bind(&row.content_type)
195 .bind(&row.metadata)
196 .bind(&row.timestamp)
197 .bind(row.version)
198 .bind(&row.parent_id)
199 .bind(&row.correlation_id)
200 .bind(&row.causation_id)
201 .bind(pd.partition_key.as_ref().map(|k| k.as_str()))
202 .bind(pd.partition_hash.map(|h| h.to_sql_i64()))
203 .bind(pd.partition_id)
204 .bind(pd.partition_count)
205 .bind(pd.partition_strategy.as_ref().map(|s| s.as_str()))
206 .execute(&mut *tx)
207 .await
208 .map_err(|e| Error::Store(e.to_string()))?;
209 }
210 tx.commit().await.map_err(|e| Error::Store(e.to_string()))?;
211 Ok(())
212 }
213}
214
215#[derive(Default)]
216struct PartitionData {
217 partition_key: Option<PartitionKey>,
218 partition_hash: Option<PartitionHash>,
219 partition_id: Option<i64>,
220 partition_count: Option<i64>,
221 partition_strategy: Option<PartitionStrategy>,
222}
223
224struct EventRow {
225 id: String,
226 organization: String,
227 namespace: String,
228 topic: String,
229 key: String,
230 payload: String,
231 content_type: String,
232 metadata: String,
233 timestamp: String,
234 version: i64,
235 parent_id: Option<String>,
236 correlation_id: Option<String>,
237 causation_id: Option<String>,
238}
239
240impl EventRow {
241 fn from_event(event: &Event) -> Result<Self> {
242 let serialized = SerializedEvent::from_event(event)?;
243 let content_type = serialized.payload.content_type().to_string();
244 let payload = serde_json::to_string(&serialized.payload)
245 .map_err(|e| Error::Store(format!("encode payload: {e}")))?;
246 let metadata = serde_json::to_string(&serialized.metadata)
247 .map_err(|e| Error::Store(format!("encode metadata: {e}")))?;
248 Ok(Self {
249 id: serialized.id.to_string(),
250 organization: serialized.organization,
251 namespace: serialized.namespace,
252 topic: serialized.topic,
253 key: serialized.key,
254 payload,
255 content_type,
256 metadata,
257 timestamp: serialized.timestamp.to_rfc3339(),
258 version: serialized.version as i64,
259 parent_id: serialized.parent_id.map(|id| id.to_string()),
260 correlation_id: serialized.correlation_id,
261 causation_id: serialized.causation_id,
262 })
263 }
264}