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
49fn lsn_to_position_bytes(lsn: u64) -> Bytes {
50 Bytes::from(lsn.to_be_bytes().to_vec())
51}
52
53pub struct PostgresBootstrapProvider {
58 config: PostgresConfig,
59}
60
61impl PostgresBootstrapProvider {
62 pub fn new(postgres_config: PostgresBootstrapConfig) -> Self {
64 Self {
65 config: PostgresConfig::from_bootstrap_config(postgres_config),
66 }
67 }
68
69 pub fn builder() -> PostgresBootstrapProviderBuilder {
71 PostgresBootstrapProviderBuilder::new()
72 }
73}
74
75pub struct PostgresBootstrapProviderBuilder {
92 host: String,
93 port: u16,
94 database: String,
95 user: String,
96 password: String,
97 tables: Vec<String>,
98 slot_name: String,
99 publication_name: String,
100 ssl_mode: SslMode,
101 table_keys: Vec<TableKeyConfig>,
102}
103
104impl PostgresBootstrapProviderBuilder {
105 pub fn new() -> Self {
107 Self {
108 host: "localhost".to_string(), port: 5432,
110 database: String::new(),
111 user: String::new(),
112 password: String::new(),
113 tables: Vec::new(),
114 slot_name: "drasi_slot".to_string(),
115 publication_name: "drasi_pub".to_string(),
116 ssl_mode: SslMode::Disable,
117 table_keys: Vec::new(),
118 }
119 }
120
121 pub fn with_host(mut self, host: impl Into<String>) -> Self {
123 self.host = host.into();
124 self
125 }
126
127 pub fn with_port(mut self, port: u16) -> Self {
129 self.port = port;
130 self
131 }
132
133 pub fn with_database(mut self, database: impl Into<String>) -> Self {
135 self.database = database.into();
136 self
137 }
138
139 pub fn with_user(mut self, user: impl Into<String>) -> Self {
141 self.user = user.into();
142 self
143 }
144
145 pub fn with_password(mut self, password: impl Into<String>) -> Self {
147 self.password = password.into();
148 self
149 }
150
151 pub fn with_tables(mut self, tables: Vec<String>) -> Self {
153 self.tables = tables;
154 self
155 }
156
157 pub fn with_table(mut self, table: impl Into<String>) -> Self {
159 self.tables.push(table.into());
160 self
161 }
162
163 pub fn with_slot_name(mut self, slot_name: impl Into<String>) -> Self {
165 self.slot_name = slot_name.into();
166 self
167 }
168
169 pub fn with_publication_name(mut self, publication_name: impl Into<String>) -> Self {
171 self.publication_name = publication_name.into();
172 self
173 }
174
175 pub fn with_ssl_mode(mut self, ssl_mode: SslMode) -> Self {
177 self.ssl_mode = ssl_mode;
178 self
179 }
180
181 pub fn with_table_keys(mut self, table_keys: Vec<TableKeyConfig>) -> Self {
183 self.table_keys = table_keys;
184 self
185 }
186
187 pub fn with_table_key(mut self, table: impl Into<String>, key_columns: Vec<String>) -> Self {
189 self.table_keys.push(TableKeyConfig {
190 table: table.into(),
191 key_columns,
192 });
193 self
194 }
195
196 pub fn build(self) -> PostgresBootstrapProvider {
198 let config = PostgresBootstrapConfig {
199 host: self.host,
200 port: self.port,
201 database: self.database,
202 user: self.user,
203 password: self.password,
204 tables: self.tables,
205 slot_name: self.slot_name,
206 publication_name: self.publication_name,
207 ssl_mode: self.ssl_mode,
208 table_keys: self.table_keys,
209 };
210 PostgresBootstrapProvider::new(config)
211 }
212}
213
214impl Default for PostgresBootstrapProviderBuilder {
215 fn default() -> Self {
216 Self::new()
217 }
218}
219
220#[async_trait]
221impl BootstrapProvider for PostgresBootstrapProvider {
222 async fn bootstrap(
223 &self,
224 request: BootstrapRequest,
225 context: &BootstrapContext,
226 event_tx: drasi_lib::channels::BootstrapEventSender,
227 _settings: Option<&drasi_lib::config::SourceSubscriptionSettings>,
228 ) -> Result<BootstrapResult> {
229 info!(
230 "Starting PostgreSQL bootstrap for query '{}' with {} node labels and {} relation labels",
231 request.query_id,
232 request.node_labels.len(),
233 request.relation_labels.len()
234 );
235
236 let mut handler =
238 PostgresBootstrapHandler::new(self.config.clone(), context.source_id.clone());
239
240 let query_id = request.query_id.clone();
242
243 let (count, source_position) = handler.execute(request, context, event_tx).await?;
245
246 info!("Completed PostgreSQL bootstrap for query {query_id}: sent {count} records");
247
248 Ok(BootstrapResult {
249 event_count: count,
250 source_position: Some(source_position),
251 })
252 }
253}
254
255#[derive(Debug, Clone)]
257struct PostgresConfig {
258 pub host: String,
259 pub port: u16,
260 pub database: String,
261 pub user: String,
262 pub password: String,
263 #[allow(dead_code)]
264 pub tables: Vec<String>,
265 #[allow(dead_code)]
266 pub slot_name: String,
267 #[allow(dead_code)]
268 pub publication_name: String,
269 #[allow(dead_code)]
270 pub ssl_mode: SslMode,
271 pub table_keys: Vec<TableKeyConfig>,
272}
273
274impl PostgresConfig {
275 fn from_bootstrap_config(postgres_config: PostgresBootstrapConfig) -> Self {
276 PostgresConfig {
277 host: postgres_config.host.clone(),
278 port: postgres_config.port,
279 database: postgres_config.database.clone(),
280 user: postgres_config.user.clone(),
281 password: postgres_config.password.clone(),
282 tables: postgres_config.tables.clone(),
283 slot_name: postgres_config.slot_name.clone(),
284 publication_name: postgres_config.publication_name.clone(),
285 ssl_mode: postgres_config.ssl_mode,
286 table_keys: postgres_config.table_keys.clone(),
287 }
288 }
289}
290
291struct PostgresBootstrapHandler {
293 config: PostgresConfig,
294 source_id: String,
295 table_primary_keys: HashMap<String, Vec<String>>,
297}
298
299impl PostgresBootstrapHandler {
300 fn new(config: PostgresConfig, source_id: String) -> Self {
301 Self {
302 config,
303 source_id,
304 table_primary_keys: HashMap::new(),
305 }
306 }
307
308 async fn execute(
310 &mut self,
311 request: BootstrapRequest,
312 context: &BootstrapContext,
313 event_tx: drasi_lib::channels::BootstrapEventSender,
314 ) -> Result<(usize, Bytes)> {
315 info!(
316 "Bootstrap: Connecting to PostgreSQL at {}:{}",
317 self.config.host, self.config.port
318 );
319
320 let mut client = self.connect().await?;
322
323 self.query_primary_keys(&client).await?;
325
326 info!("Bootstrap: Connected, creating snapshot transaction...");
327 let (transaction, snapshot_lsn) = self.create_snapshot(&mut client).await?;
329 let source_position = lsn_to_position_bytes(snapshot_lsn);
330
331 info!("Bootstrap snapshot created at LSN: {snapshot_lsn:x}");
332
333 let tables = self.resolve_tables(&request, &transaction).await?;
335 info!(
336 "Resolved {} labels to {} tables",
337 request.node_labels.len() + request.relation_labels.len(),
338 tables.len()
339 );
340
341 let mut total_count = 0;
343 for table in &tables {
344 let count = self
345 .bootstrap_table(&transaction, table, context, &event_tx)
346 .await?;
347 info!("Bootstrapped {count} rows from table '{table}'");
348 total_count += count;
349 }
350
351 transaction.commit().await?;
353
354 info!("Bootstrap completed: {total_count} total elements sent");
355 Ok((total_count, source_position))
356 }
357
358 async fn connect(&self) -> Result<Client> {
360 let connection_string = format!(
361 "host={} port={} user={} password={} dbname={}",
362 self.config.host,
363 self.config.port,
364 self.config.user,
365 self.config.password,
366 self.config.database
367 );
368
369 let (client, connection) = tokio_postgres::connect(&connection_string, NoTls).await?;
370
371 tokio::spawn(async move {
373 if let Err(e) = connection.await {
374 error!("PostgreSQL connection error: {e}");
375 }
376 });
377
378 Ok(client)
379 }
380
381 async fn create_snapshot<'a>(&self, client: &'a mut Client) -> Result<(Transaction<'a>, u64)> {
383 let transaction = client
385 .build_transaction()
386 .isolation_level(tokio_postgres::IsolationLevel::RepeatableRead)
387 .start()
388 .await?;
389
390 let row = transaction
392 .query_one("SELECT pg_current_wal_lsn()::text", &[])
393 .await?;
394 let lsn: String = row.get(0);
395 let lsn = parse_lsn(&lsn).context("Failed to parse PostgreSQL snapshot LSN")?;
396
397 Ok((transaction, lsn))
398 }
399
400 async fn resolve_tables(
404 &self,
405 request: &BootstrapRequest,
406 transaction: &Transaction<'_>,
407 ) -> Result<Vec<String>> {
408 let mut tables = Vec::new();
409
410 let all_labels: Vec<String> = request
412 .node_labels
413 .iter()
414 .chain(request.relation_labels.iter())
415 .cloned()
416 .collect();
417
418 for label in all_labels {
419 if self.table_exists(transaction, &label).await? {
420 tables.push(label);
421 } else {
422 warn!("Table '{label}' does not exist, skipping");
423 }
424 }
425
426 Ok(tables)
427 }
428
429 async fn table_exists(&self, transaction: &Transaction<'_>, table_name: &str) -> Result<bool> {
431 let row = transaction
432 .query_one(
433 "SELECT EXISTS (
434 SELECT 1 FROM information_schema.tables
435 WHERE table_schema = 'public'
436 AND table_name = $1
437 )",
438 &[&table_name],
439 )
440 .await?;
441
442 Ok(row.get(0))
443 }
444
445 async fn bootstrap_table(
447 &self,
448 transaction: &Transaction<'_>,
449 table: &str,
450 context: &BootstrapContext,
451 event_tx: &drasi_lib::channels::BootstrapEventSender,
452 ) -> Result<usize> {
453 debug!("Starting bootstrap of table '{table}'");
454
455 let columns = self.get_table_columns(transaction, table).await?;
457
458 let query = format!("SELECT * FROM \"{}\"", table.replace('"', "\"\""));
460 let rows = transaction.query(&query, &[]).await?;
461
462 let mut count = 0;
463 let mut batch = Vec::new();
464 let batch_size = 1000;
465
466 for row in rows {
467 let source_change = self.row_to_source_change(&row, table, &columns).await?;
468
469 batch.push(SourceChangeEvent {
470 source_id: self.source_id.clone(),
471 change: source_change,
472 timestamp: chrono::Utc::now(),
473 sequence: None,
474 });
475
476 if batch.len() >= batch_size {
477 self.send_batch(&mut batch, context, event_tx).await?;
478 count += batch_size;
479 }
480 }
481
482 if !batch.is_empty() {
484 count += batch.len();
485 self.send_batch(&mut batch, context, event_tx).await?;
486 }
487
488 Ok(count)
489 }
490
491 async fn get_table_columns(
493 &self,
494 transaction: &Transaction<'_>,
495 table_name: &str,
496 ) -> Result<Vec<ColumnInfo>> {
497 let rows = transaction
498 .query(
499 "SELECT column_name,
500 CASE
501 WHEN data_type = 'character varying' THEN 1043
502 WHEN data_type = 'integer' THEN 23
503 WHEN data_type = 'bigint' THEN 20
504 WHEN data_type = 'smallint' THEN 21
505 WHEN data_type = 'text' THEN 25
506 WHEN data_type = 'boolean' THEN 16
507 WHEN data_type = 'numeric' THEN 1700
508 WHEN data_type = 'real' THEN 700
509 WHEN data_type = 'double precision' THEN 701
510 WHEN data_type = 'timestamp without time zone' THEN 1114
511 WHEN data_type = 'timestamp with time zone' THEN 1184
512 WHEN data_type = 'date' THEN 1082
513 WHEN data_type = 'uuid' THEN 2950
514 WHEN data_type = 'json' THEN 114
515 WHEN data_type = 'jsonb' THEN 3802
516 ELSE 25 -- Default to text
517 END as type_oid
518 FROM information_schema.columns
519 WHERE table_schema = 'public' AND table_name = $1
520 ORDER BY ordinal_position",
521 &[&table_name],
522 )
523 .await?;
524
525 let mut columns = Vec::new();
526 for row in rows {
527 columns.push(ColumnInfo {
528 name: row.get(0),
529 type_oid: row.get::<_, i32>(1),
530 });
531 }
532
533 Ok(columns)
534 }
535
536 async fn query_primary_keys(&mut self, client: &Client) -> Result<()> {
538 info!("Querying primary key information from PostgreSQL system catalogs");
539
540 let query = r#"
541 SELECT
542 n.nspname as schema_name,
543 c.relname as table_name,
544 a.attname as column_name
545 FROM pg_constraint con
546 JOIN pg_class c ON con.conrelid = c.oid
547 JOIN pg_namespace n ON c.relnamespace = n.oid
548 JOIN pg_attribute a ON a.attrelid = c.oid
549 WHERE con.contype = 'p' -- Primary key constraint
550 AND a.attnum = ANY(con.conkey)
551 AND n.nspname NOT IN ('pg_catalog', 'information_schema')
552 ORDER BY n.nspname, c.relname, array_position(con.conkey, a.attnum)
553 "#;
554
555 let rows = client.query(query, &[]).await?;
556
557 let mut primary_keys: HashMap<String, Vec<String>> = HashMap::new();
558
559 for row in rows {
560 let schema: &str = row.get(0);
561 let table: &str = row.get(1);
562 let column: &str = row.get(2);
563
564 let table_key = if schema == "public" {
566 table.to_string()
567 } else {
568 format!("{schema}.{table}")
569 };
570
571 primary_keys
572 .entry(table_key.clone())
573 .or_default()
574 .push(column.to_string());
575
576 debug!("Found primary key column '{column}' for table '{table_key}'");
577 }
578
579 for table_key_config in &self.config.table_keys {
581 let table_name = &table_key_config.table;
582 let key_columns = &table_key_config.key_columns;
583
584 if !key_columns.is_empty() {
585 info!(
586 "Using user-configured key columns for table '{table_name}': {key_columns:?}"
587 );
588 primary_keys.insert(table_name.clone(), key_columns.clone());
589 }
590 }
591
592 self.table_primary_keys = primary_keys.clone();
594
595 info!("Found primary keys for {} tables", primary_keys.len());
596 for (table, keys) in &primary_keys {
597 info!("Table '{table}' primary key columns: {keys:?}");
598 }
599
600 Ok(())
601 }
602
603 async fn row_to_source_change(
605 &self,
606 row: &Row,
607 table: &str,
608 columns: &[ColumnInfo],
609 ) -> Result<SourceChange> {
610 let mut properties = ElementPropertyMap::new();
611
612 let pk_columns = self.table_primary_keys.get(table);
614
615 let mut pk_values = Vec::new();
617
618 for (idx, column) in columns.iter().enumerate() {
619 let is_pk = pk_columns
621 .map(|pks| pks.contains(&column.name))
622 .unwrap_or(false);
623
624 let element_value = match column.type_oid {
626 16 => {
627 if let Ok(Some(val)) = row.try_get::<_, Option<bool>>(idx) {
629 drasi_core::models::ElementValue::Bool(val)
630 } else {
631 drasi_core::models::ElementValue::Null
632 }
633 }
634 21 | 23 | 20 => {
635 if let Ok(Some(val)) = row.try_get::<_, Option<i64>>(idx) {
637 drasi_core::models::ElementValue::Integer(val)
638 } else if let Ok(Some(val)) = row.try_get::<_, Option<i32>>(idx) {
639 drasi_core::models::ElementValue::Integer(val as i64)
640 } else if let Ok(Some(val)) = row.try_get::<_, Option<i16>>(idx) {
641 drasi_core::models::ElementValue::Integer(val as i64)
642 } else {
643 drasi_core::models::ElementValue::Null
644 }
645 }
646 700 | 701 => {
647 if let Ok(Some(val)) = row.try_get::<_, Option<f64>>(idx) {
649 drasi_core::models::ElementValue::Float(ordered_float::OrderedFloat(val))
650 } else if let Ok(Some(val)) = row.try_get::<_, Option<f32>>(idx) {
651 drasi_core::models::ElementValue::Float(ordered_float::OrderedFloat(
652 val as f64,
653 ))
654 } else {
655 drasi_core::models::ElementValue::Null
656 }
657 }
658 1700 => {
659 if let Ok(Some(val)) = row.try_get::<_, Option<rust_decimal::Decimal>>(idx) {
661 drasi_core::models::ElementValue::Float(ordered_float::OrderedFloat(
662 val.to_string().parse::<f64>().unwrap_or(0.0),
663 ))
664 } else {
665 drasi_core::models::ElementValue::Null
666 }
667 }
668 25 | 1043 | 19 => {
669 if let Ok(Some(val)) = row.try_get::<_, Option<String>>(idx) {
671 drasi_core::models::ElementValue::String(std::sync::Arc::from(val))
672 } else {
673 drasi_core::models::ElementValue::Null
674 }
675 }
676 1114 | 1184 => {
677 if let Ok(Some(val)) = row.try_get::<_, Option<chrono::NaiveDateTime>>(idx) {
679 drasi_core::models::ElementValue::String(std::sync::Arc::from(
680 val.to_string(),
681 ))
682 } else if let Ok(Some(val)) =
683 row.try_get::<_, Option<chrono::DateTime<chrono::Utc>>>(idx)
684 {
685 drasi_core::models::ElementValue::String(std::sync::Arc::from(
686 val.to_string(),
687 ))
688 } else {
689 drasi_core::models::ElementValue::Null
690 }
691 }
692 _ => {
693 if let Ok(Some(val)) = row.try_get::<_, Option<String>>(idx) {
695 drasi_core::models::ElementValue::String(std::sync::Arc::from(val))
696 } else {
697 drasi_core::models::ElementValue::Null
698 }
699 }
700 };
701
702 if is_pk && !matches!(element_value, drasi_core::models::ElementValue::Null) {
704 let value_str = match &element_value {
705 drasi_core::models::ElementValue::Integer(i) => i.to_string(),
706 drasi_core::models::ElementValue::Float(f) => f.to_string(),
707 drasi_core::models::ElementValue::String(s) => s.to_string(),
708 drasi_core::models::ElementValue::Bool(b) => b.to_string(),
709 _ => format!("{element_value:?}"),
710 };
711 pk_values.push(value_str);
712 }
713
714 properties.insert(&column.name, element_value);
715 }
716
717 let elem_id = if !pk_values.is_empty() {
720 format!("{}:{}", table, pk_values.join("_"))
722 } else if pk_columns.is_none() || pk_columns.map(|pks| pks.is_empty()).unwrap_or(true) {
723 warn!(
725 "No primary key found for table '{table}'. Consider adding 'table_keys' configuration."
726 );
727 format!("{}:{}", table, uuid::Uuid::new_v4())
729 } else {
730 format!("{}:{}", table, uuid::Uuid::new_v4())
732 };
733
734 let metadata = ElementMetadata {
735 reference: ElementReference::new(&self.source_id, &elem_id),
736 labels: Arc::from(vec![Arc::from(table)]),
737 effective_from: chrono::Utc::now().timestamp_millis() as u64,
738 };
739
740 let element = Element::Node {
741 metadata,
742 properties,
743 };
744
745 Ok(SourceChange::Insert { element })
746 }
747
748 async fn send_batch(
750 &self,
751 batch: &mut Vec<SourceChangeEvent>,
752 context: &BootstrapContext,
753 event_tx: &drasi_lib::channels::BootstrapEventSender,
754 ) -> Result<()> {
755 for event in batch.drain(..) {
756 let sequence = context.next_sequence();
758
759 let bootstrap_event = drasi_lib::channels::BootstrapEvent {
760 source_id: event.source_id,
761 change: event.change,
762 timestamp: event.timestamp,
763 sequence,
764 };
765 event_tx.send(bootstrap_event).await.map_err(|e| {
766 anyhow!("Failed to send bootstrap event to channel (channel may be closed): {e}")
767 })?;
768 }
769 Ok(())
770 }
771}
772
773#[derive(Debug)]
774struct ColumnInfo {
775 name: String,
776 type_oid: i32,
777}
778
779#[cfg(test)]
780mod tests {
781 use drasi_core::models::validate_effective_from;
782
783 #[test]
789 fn effective_from_uses_milliseconds() {
790 let effective_from = chrono::Utc::now().timestamp_millis() as u64;
791 assert!(
792 validate_effective_from(effective_from).is_ok(),
793 "Postgres bootstrapper effective_from ({effective_from}) should be in millisecond range"
794 );
795 }
796
797 #[test]
799 fn effective_from_rejects_nanoseconds_pattern() {
800 let bad_effective_from = chrono::Utc::now().timestamp_nanos_opt().unwrap() as u64;
802 assert!(
803 validate_effective_from(bad_effective_from).is_err(),
804 "Nanosecond timestamp ({bad_effective_from}) should be rejected"
805 );
806 }
807}