1use anyhow::{anyhow, Context, Result};
18use async_trait::async_trait;
19use bytes::Bytes;
20use drasi_core::models::{
21 Element, ElementMetadata, ElementPropertyMap, ElementReference, SourceChange,
22};
23use log::{debug, error, info, warn};
24use std::collections::HashMap;
25use std::sync::Arc;
26use tokio_postgres::{Client, NoTls, Row, Transaction};
27
28use drasi_lib::bootstrap::{
29 BootstrapContext, BootstrapProvider, BootstrapRequest, BootstrapResult,
30};
31use drasi_lib::channels::SourceChangeEvent;
32
33pub use crate::config::{PostgresBootstrapConfig, SslMode, TableKeyConfig};
34
35fn parse_lsn(lsn_str: &str) -> Result<u64> {
36 let parts: Vec<&str> = lsn_str.split('/').collect();
37 if parts.len() != 2 {
38 return Err(anyhow!("Invalid LSN format: {lsn_str}"));
39 }
40
41 let high = u64::from_str_radix(parts[0], 16)
42 .with_context(|| format!("Invalid high bits in LSN: {lsn_str}"))?;
43 let low = u64::from_str_radix(parts[1], 16)
44 .with_context(|| format!("Invalid low bits in LSN: {lsn_str}"))?;
45
46 Ok((high << 32) | low)
47}
48
49pub fn snapshot_position_bytes(snapshot_lsn: u64) -> Bytes {
58 let mut buf = Vec::with_capacity(16);
59 buf.extend_from_slice(&snapshot_lsn.to_be_bytes());
60 buf.extend_from_slice(&u64::MAX.to_be_bytes());
61 Bytes::from(buf)
62}
63
64pub struct PostgresBootstrapProvider {
69 config: PostgresConfig,
70}
71
72impl PostgresBootstrapProvider {
73 pub fn new(postgres_config: PostgresBootstrapConfig) -> Self {
75 Self {
76 config: PostgresConfig::from_bootstrap_config(postgres_config),
77 }
78 }
79
80 pub fn builder() -> PostgresBootstrapProviderBuilder {
82 PostgresBootstrapProviderBuilder::new()
83 }
84}
85
86pub struct PostgresBootstrapProviderBuilder {
103 host: String,
104 port: u16,
105 database: String,
106 user: String,
107 password: String,
108 tables: Vec<String>,
109 slot_name: String,
110 publication_name: String,
111 ssl_mode: SslMode,
112 table_keys: Vec<TableKeyConfig>,
113}
114
115impl PostgresBootstrapProviderBuilder {
116 pub fn new() -> Self {
118 Self {
119 host: "localhost".to_string(), port: 5432,
121 database: String::new(),
122 user: String::new(),
123 password: String::new(),
124 tables: Vec::new(),
125 slot_name: "drasi_slot".to_string(),
126 publication_name: "drasi_pub".to_string(),
127 ssl_mode: SslMode::Disable,
128 table_keys: Vec::new(),
129 }
130 }
131
132 pub fn with_host(mut self, host: impl Into<String>) -> Self {
134 self.host = host.into();
135 self
136 }
137
138 pub fn with_port(mut self, port: u16) -> Self {
140 self.port = port;
141 self
142 }
143
144 pub fn with_database(mut self, database: impl Into<String>) -> Self {
146 self.database = database.into();
147 self
148 }
149
150 pub fn with_user(mut self, user: impl Into<String>) -> Self {
152 self.user = user.into();
153 self
154 }
155
156 pub fn with_password(mut self, password: impl Into<String>) -> Self {
158 self.password = password.into();
159 self
160 }
161
162 pub fn with_tables(mut self, tables: Vec<String>) -> Self {
164 self.tables = tables;
165 self
166 }
167
168 pub fn with_table(mut self, table: impl Into<String>) -> Self {
170 self.tables.push(table.into());
171 self
172 }
173
174 pub fn with_slot_name(mut self, slot_name: impl Into<String>) -> Self {
176 self.slot_name = slot_name.into();
177 self
178 }
179
180 pub fn with_publication_name(mut self, publication_name: impl Into<String>) -> Self {
182 self.publication_name = publication_name.into();
183 self
184 }
185
186 pub fn with_ssl_mode(mut self, ssl_mode: SslMode) -> Self {
188 self.ssl_mode = ssl_mode;
189 self
190 }
191
192 pub fn with_table_keys(mut self, table_keys: Vec<TableKeyConfig>) -> Self {
194 self.table_keys = table_keys;
195 self
196 }
197
198 pub fn with_table_key(mut self, table: impl Into<String>, key_columns: Vec<String>) -> Self {
200 self.table_keys.push(TableKeyConfig {
201 table: table.into(),
202 key_columns,
203 });
204 self
205 }
206
207 pub fn build(self) -> PostgresBootstrapProvider {
209 let config = PostgresBootstrapConfig {
210 host: self.host,
211 port: self.port,
212 database: self.database,
213 user: self.user,
214 password: self.password,
215 tables: self.tables,
216 slot_name: self.slot_name,
217 publication_name: self.publication_name,
218 ssl_mode: self.ssl_mode,
219 table_keys: self.table_keys,
220 };
221 PostgresBootstrapProvider::new(config)
222 }
223}
224
225impl Default for PostgresBootstrapProviderBuilder {
226 fn default() -> Self {
227 Self::new()
228 }
229}
230
231#[async_trait]
232impl BootstrapProvider for PostgresBootstrapProvider {
233 async fn bootstrap(
234 &self,
235 request: BootstrapRequest,
236 context: &BootstrapContext,
237 event_tx: drasi_lib::channels::BootstrapEventSender,
238 _settings: Option<&drasi_lib::config::SourceSubscriptionSettings>,
239 ) -> Result<BootstrapResult> {
240 info!(
241 "Starting PostgreSQL bootstrap for query '{}' with {} node labels and {} relation labels",
242 request.query_id,
243 request.node_labels.len(),
244 request.relation_labels.len()
245 );
246
247 let mut handler =
249 PostgresBootstrapHandler::new(self.config.clone(), context.source_id.clone());
250
251 let query_id = request.query_id.clone();
253
254 let (count, source_position) = handler.execute(request, context, event_tx).await?;
256
257 info!("Completed PostgreSQL bootstrap for query {query_id}: sent {count} records");
258
259 Ok(BootstrapResult {
260 event_count: count,
261 source_position: Some(source_position),
262 })
263 }
264}
265
266#[derive(Debug, Clone)]
268struct PostgresConfig {
269 pub host: String,
270 pub port: u16,
271 pub database: String,
272 pub user: String,
273 pub password: String,
274 #[allow(dead_code)]
275 pub tables: Vec<String>,
276 #[allow(dead_code)]
277 pub slot_name: String,
278 #[allow(dead_code)]
279 pub publication_name: String,
280 #[allow(dead_code)]
281 pub ssl_mode: SslMode,
282 pub table_keys: Vec<TableKeyConfig>,
283}
284
285impl PostgresConfig {
286 fn from_bootstrap_config(postgres_config: PostgresBootstrapConfig) -> Self {
287 PostgresConfig {
288 host: postgres_config.host.clone(),
289 port: postgres_config.port,
290 database: postgres_config.database.clone(),
291 user: postgres_config.user.clone(),
292 password: postgres_config.password.clone(),
293 tables: postgres_config.tables.clone(),
294 slot_name: postgres_config.slot_name.clone(),
295 publication_name: postgres_config.publication_name.clone(),
296 ssl_mode: postgres_config.ssl_mode,
297 table_keys: postgres_config.table_keys.clone(),
298 }
299 }
300}
301
302struct PostgresBootstrapHandler {
304 config: PostgresConfig,
305 source_id: String,
306 table_primary_keys: HashMap<String, Vec<String>>,
308}
309
310impl PostgresBootstrapHandler {
311 fn new(config: PostgresConfig, source_id: String) -> Self {
312 Self {
313 config,
314 source_id,
315 table_primary_keys: HashMap::new(),
316 }
317 }
318
319 async fn execute(
321 &mut self,
322 request: BootstrapRequest,
323 context: &BootstrapContext,
324 event_tx: drasi_lib::channels::BootstrapEventSender,
325 ) -> Result<(usize, Bytes)> {
326 info!(
327 "Bootstrap: Connecting to PostgreSQL at {}:{}",
328 self.config.host, self.config.port
329 );
330
331 let mut client = self.connect().await?;
333
334 self.query_primary_keys(&client).await?;
336
337 info!("Bootstrap: Connected, creating snapshot transaction...");
338 let (transaction, snapshot_lsn) = self.create_snapshot(&mut client).await?;
340 let source_position = snapshot_position_bytes(snapshot_lsn);
341
342 info!("Bootstrap snapshot created at LSN: {snapshot_lsn:x}");
343
344 let tables = self.resolve_tables(&request, &transaction).await?;
346 info!(
347 "Resolved {} labels to {} tables",
348 request.node_labels.len() + request.relation_labels.len(),
349 tables.len()
350 );
351
352 let mut total_count = 0;
354 for table in &tables {
355 let count = self
356 .bootstrap_table(&transaction, table, context, &event_tx)
357 .await?;
358 info!("Bootstrapped {count} rows from table '{table}'");
359 total_count += count;
360 }
361
362 transaction.commit().await?;
364
365 info!("Bootstrap completed: {total_count} total elements sent");
366 Ok((total_count, source_position))
367 }
368
369 async fn connect(&self) -> Result<Client> {
371 let connection_string = format!(
372 "host={} port={} user={} password={} dbname={}",
373 self.config.host,
374 self.config.port,
375 self.config.user,
376 self.config.password,
377 self.config.database
378 );
379
380 let (client, connection) = tokio_postgres::connect(&connection_string, NoTls).await?;
381
382 tokio::spawn(async move {
384 if let Err(e) = connection.await {
385 error!("PostgreSQL connection error: {e}");
386 }
387 });
388
389 Ok(client)
390 }
391
392 async fn create_snapshot<'a>(&self, client: &'a mut Client) -> Result<(Transaction<'a>, u64)> {
394 let transaction = client
396 .build_transaction()
397 .isolation_level(tokio_postgres::IsolationLevel::RepeatableRead)
398 .start()
399 .await?;
400
401 let row = transaction
403 .query_one("SELECT pg_current_wal_lsn()::text", &[])
404 .await?;
405 let lsn: String = row.get(0);
406 let lsn = parse_lsn(&lsn).context("Failed to parse PostgreSQL snapshot LSN")?;
407
408 Ok((transaction, lsn))
409 }
410
411 async fn resolve_tables(
415 &self,
416 request: &BootstrapRequest,
417 transaction: &Transaction<'_>,
418 ) -> Result<Vec<String>> {
419 let mut tables = Vec::new();
420
421 let all_labels: Vec<String> = request
423 .node_labels
424 .iter()
425 .chain(request.relation_labels.iter())
426 .cloned()
427 .collect();
428
429 for label in all_labels {
430 if self.table_exists(transaction, &label).await? {
431 tables.push(label);
432 } else {
433 warn!("Table '{label}' does not exist, skipping");
434 }
435 }
436
437 Ok(tables)
438 }
439
440 async fn table_exists(&self, transaction: &Transaction<'_>, table_name: &str) -> Result<bool> {
442 let row = transaction
443 .query_one(
444 "SELECT EXISTS (
445 SELECT 1 FROM information_schema.tables
446 WHERE table_schema = 'public'
447 AND table_name = $1
448 )",
449 &[&table_name],
450 )
451 .await?;
452
453 Ok(row.get(0))
454 }
455
456 async fn bootstrap_table(
458 &self,
459 transaction: &Transaction<'_>,
460 table: &str,
461 context: &BootstrapContext,
462 event_tx: &drasi_lib::channels::BootstrapEventSender,
463 ) -> Result<usize> {
464 debug!("Starting bootstrap of table '{table}'");
465
466 let columns = self.get_table_columns(transaction, table).await?;
468
469 let query = format!("SELECT * FROM \"{}\"", table.replace('"', "\"\""));
471 let rows = transaction.query(&query, &[]).await?;
472
473 let mut count = 0;
474 let mut batch = Vec::new();
475 let batch_size = 1000;
476
477 for row in rows {
478 let source_change = self.row_to_source_change(&row, table, &columns).await?;
479
480 batch.push(SourceChangeEvent {
481 source_id: self.source_id.clone(),
482 change: source_change,
483 timestamp: chrono::Utc::now(),
484 sequence: None,
485 });
486
487 if batch.len() >= batch_size {
488 self.send_batch(&mut batch, context, event_tx).await?;
489 count += batch_size;
490 }
491 }
492
493 if !batch.is_empty() {
495 count += batch.len();
496 self.send_batch(&mut batch, context, event_tx).await?;
497 }
498
499 Ok(count)
500 }
501
502 async fn get_table_columns(
504 &self,
505 transaction: &Transaction<'_>,
506 table_name: &str,
507 ) -> Result<Vec<ColumnInfo>> {
508 let rows = transaction
509 .query(
510 "SELECT column_name,
511 CASE
512 WHEN data_type = 'character varying' THEN 1043
513 WHEN data_type = 'integer' THEN 23
514 WHEN data_type = 'bigint' THEN 20
515 WHEN data_type = 'smallint' THEN 21
516 WHEN data_type = 'text' THEN 25
517 WHEN data_type = 'boolean' THEN 16
518 WHEN data_type = 'numeric' THEN 1700
519 WHEN data_type = 'real' THEN 700
520 WHEN data_type = 'double precision' THEN 701
521 WHEN data_type = 'timestamp without time zone' THEN 1114
522 WHEN data_type = 'timestamp with time zone' THEN 1184
523 WHEN data_type = 'date' THEN 1082
524 WHEN data_type = 'uuid' THEN 2950
525 WHEN data_type = 'json' THEN 114
526 WHEN data_type = 'jsonb' THEN 3802
527 ELSE 25 -- Default to text
528 END as type_oid
529 FROM information_schema.columns
530 WHERE table_schema = 'public' AND table_name = $1
531 ORDER BY ordinal_position",
532 &[&table_name],
533 )
534 .await?;
535
536 let mut columns = Vec::new();
537 for row in rows {
538 columns.push(ColumnInfo {
539 name: row.get(0),
540 type_oid: row.get::<_, i32>(1),
541 });
542 }
543
544 Ok(columns)
545 }
546
547 async fn query_primary_keys(&mut self, client: &Client) -> Result<()> {
549 info!("Querying primary key information from PostgreSQL system catalogs");
550
551 let query = r#"
552 SELECT
553 n.nspname as schema_name,
554 c.relname as table_name,
555 a.attname as column_name
556 FROM pg_constraint con
557 JOIN pg_class c ON con.conrelid = c.oid
558 JOIN pg_namespace n ON c.relnamespace = n.oid
559 JOIN pg_attribute a ON a.attrelid = c.oid
560 WHERE con.contype = 'p' -- Primary key constraint
561 AND a.attnum = ANY(con.conkey)
562 AND n.nspname NOT IN ('pg_catalog', 'information_schema')
563 ORDER BY n.nspname, c.relname, array_position(con.conkey, a.attnum)
564 "#;
565
566 let rows = client.query(query, &[]).await?;
567
568 let mut primary_keys: HashMap<String, Vec<String>> = HashMap::new();
569
570 for row in rows {
571 let schema: &str = row.get(0);
572 let table: &str = row.get(1);
573 let column: &str = row.get(2);
574
575 let table_key = if schema == "public" {
577 table.to_string()
578 } else {
579 format!("{schema}.{table}")
580 };
581
582 primary_keys
583 .entry(table_key.clone())
584 .or_default()
585 .push(column.to_string());
586
587 debug!("Found primary key column '{column}' for table '{table_key}'");
588 }
589
590 for table_key_config in &self.config.table_keys {
592 let table_name = &table_key_config.table;
593 let key_columns = &table_key_config.key_columns;
594
595 if !key_columns.is_empty() {
596 info!(
597 "Using user-configured key columns for table '{table_name}': {key_columns:?}"
598 );
599 primary_keys.insert(table_name.clone(), key_columns.clone());
600 }
601 }
602
603 self.table_primary_keys = primary_keys.clone();
605
606 info!("Found primary keys for {} tables", primary_keys.len());
607 for (table, keys) in &primary_keys {
608 info!("Table '{table}' primary key columns: {keys:?}");
609 }
610
611 Ok(())
612 }
613
614 async fn row_to_source_change(
616 &self,
617 row: &Row,
618 table: &str,
619 columns: &[ColumnInfo],
620 ) -> Result<SourceChange> {
621 let mut properties = ElementPropertyMap::new();
622
623 let pk_columns = self.table_primary_keys.get(table);
625
626 let mut pk_values = Vec::new();
628
629 for (idx, column) in columns.iter().enumerate() {
630 let is_pk = pk_columns
632 .map(|pks| pks.contains(&column.name))
633 .unwrap_or(false);
634
635 let element_value = match column.type_oid {
637 16 => {
638 if let Ok(Some(val)) = row.try_get::<_, Option<bool>>(idx) {
640 drasi_core::models::ElementValue::Bool(val)
641 } else {
642 drasi_core::models::ElementValue::Null
643 }
644 }
645 21 | 23 | 20 => {
646 if let Ok(Some(val)) = row.try_get::<_, Option<i64>>(idx) {
648 drasi_core::models::ElementValue::Integer(val)
649 } else if let Ok(Some(val)) = row.try_get::<_, Option<i32>>(idx) {
650 drasi_core::models::ElementValue::Integer(val as i64)
651 } else if let Ok(Some(val)) = row.try_get::<_, Option<i16>>(idx) {
652 drasi_core::models::ElementValue::Integer(val as i64)
653 } else {
654 drasi_core::models::ElementValue::Null
655 }
656 }
657 700 | 701 => {
658 if let Ok(Some(val)) = row.try_get::<_, Option<f64>>(idx) {
660 drasi_core::models::ElementValue::Float(ordered_float::OrderedFloat(val))
661 } else if let Ok(Some(val)) = row.try_get::<_, Option<f32>>(idx) {
662 drasi_core::models::ElementValue::Float(ordered_float::OrderedFloat(
663 val as f64,
664 ))
665 } else {
666 drasi_core::models::ElementValue::Null
667 }
668 }
669 1700 => {
670 if let Ok(Some(val)) = row.try_get::<_, Option<rust_decimal::Decimal>>(idx) {
672 drasi_core::models::ElementValue::Float(ordered_float::OrderedFloat(
673 val.to_string().parse::<f64>().unwrap_or(0.0),
674 ))
675 } else {
676 drasi_core::models::ElementValue::Null
677 }
678 }
679 25 | 1043 | 19 => {
680 if let Ok(Some(val)) = row.try_get::<_, Option<String>>(idx) {
682 drasi_core::models::ElementValue::String(std::sync::Arc::from(val))
683 } else {
684 drasi_core::models::ElementValue::Null
685 }
686 }
687 1114 | 1184 => {
688 if let Ok(Some(val)) = row.try_get::<_, Option<chrono::NaiveDateTime>>(idx) {
690 drasi_core::models::ElementValue::LocalDateTime(val)
691 } else if let Ok(Some(val)) =
692 row.try_get::<_, Option<chrono::DateTime<chrono::Utc>>>(idx)
693 {
694 drasi_core::models::ElementValue::ZonedDateTime(val.fixed_offset())
695 } else {
696 drasi_core::models::ElementValue::Null
697 }
698 }
699 _ => {
700 if let Ok(Some(val)) = row.try_get::<_, Option<String>>(idx) {
702 drasi_core::models::ElementValue::String(std::sync::Arc::from(val))
703 } else {
704 drasi_core::models::ElementValue::Null
705 }
706 }
707 };
708
709 if is_pk && !matches!(element_value, drasi_core::models::ElementValue::Null) {
711 let value_str = match &element_value {
712 drasi_core::models::ElementValue::Integer(i) => i.to_string(),
713 drasi_core::models::ElementValue::Float(f) => f.to_string(),
714 drasi_core::models::ElementValue::String(s) => s.to_string(),
715 drasi_core::models::ElementValue::Bool(b) => b.to_string(),
716 drasi_core::models::ElementValue::LocalDateTime(dt) => dt.to_string(),
717 drasi_core::models::ElementValue::ZonedDateTime(dt) => dt.to_rfc3339(),
718 _ => format!("{element_value:?}"),
719 };
720 pk_values.push(value_str);
721 }
722
723 properties.insert(&column.name, element_value);
724 }
725
726 let elem_id = if !pk_values.is_empty() {
729 format!("{}:{}", table, pk_values.join("_"))
731 } else if pk_columns.is_none() || pk_columns.map(|pks| pks.is_empty()).unwrap_or(true) {
732 warn!(
734 "No primary key found for table '{table}'. Consider adding 'table_keys' configuration."
735 );
736 format!("{}:{}", table, uuid::Uuid::new_v4())
738 } else {
739 format!("{}:{}", table, uuid::Uuid::new_v4())
741 };
742
743 let metadata = ElementMetadata {
744 reference: ElementReference::new(&self.source_id, &elem_id),
745 labels: Arc::from(vec![Arc::from(table)]),
746 effective_from: chrono::Utc::now().timestamp_millis() as u64,
747 };
748
749 let element = Element::Node {
750 metadata,
751 properties,
752 };
753
754 Ok(SourceChange::Insert { element })
755 }
756
757 async fn send_batch(
759 &self,
760 batch: &mut Vec<SourceChangeEvent>,
761 context: &BootstrapContext,
762 event_tx: &drasi_lib::channels::BootstrapEventSender,
763 ) -> Result<()> {
764 for event in batch.drain(..) {
765 let sequence = context.next_sequence();
767
768 let bootstrap_event = drasi_lib::channels::BootstrapEvent {
769 source_id: event.source_id,
770 change: event.change,
771 timestamp: event.timestamp,
772 sequence,
773 };
774 event_tx.send(bootstrap_event).await.map_err(|e| {
775 anyhow!("Failed to send bootstrap event to channel (channel may be closed): {e}")
776 })?;
777 }
778 Ok(())
779 }
780}
781
782#[derive(Debug)]
783struct ColumnInfo {
784 name: String,
785 type_oid: i32,
786}
787
788#[cfg(test)]
789mod tests {
790 use drasi_core::models::validate_effective_from;
791
792 #[test]
798 fn effective_from_uses_milliseconds() {
799 let effective_from = chrono::Utc::now().timestamp_millis() as u64;
800 assert!(
801 validate_effective_from(effective_from).is_ok(),
802 "Postgres bootstrapper effective_from ({effective_from}) should be in millisecond range"
803 );
804 }
805
806 #[test]
808 fn effective_from_rejects_nanoseconds_pattern() {
809 let bad_effective_from = chrono::Utc::now().timestamp_nanos_opt().unwrap() as u64;
811 assert!(
812 validate_effective_from(bad_effective_from).is_err(),
813 "Nanosecond timestamp ({bad_effective_from}) should be rejected"
814 );
815 }
816
817 #[test]
822 fn snapshot_position_bytes_is_16_byte_max_padded() {
823 let lsn = 0x0000_0000_0152_00b0u64;
824 let pos = super::snapshot_position_bytes(lsn);
825 assert_eq!(pos.len(), 16, "snapshot boundary must be 16 bytes");
826
827 let mut expected = Vec::with_capacity(16);
828 expected.extend_from_slice(&lsn.to_be_bytes());
829 expected.extend_from_slice(&u64::MAX.to_be_bytes());
830 assert_eq!(pos.as_ref(), expected.as_slice());
831 }
832}