fix: harden remote filters and async retries

christian-smith · christian-smith · commit 7ba7c7cefd89 · 2025-11-26T06:52:07.000Z
diff --git a/src/main.rs b/src/main.rs
@@ -9,6 +9,9 @@ use database_replicator::commands;
 #[command(about = "Universal database-to-PostgreSQL replication CLI", long_about = None)]
 #[command(version)]
 struct Cli {
+    /// Allow self-signed TLS certificates. Also honors SEREN_ALLOW_SELF_SIGNED_CERTS=1
+    #[arg(long = "allow-self-signed-certs", global = true, default_value_t = false)]
+    allow_self_signed_certs: bool,
     #[command(subcommand)]
     command: Commands,
 }
@@ -181,6 +184,24 @@ async fn main() -> anyhow::Result<()> {
 
     let cli = Cli::parse();
 
+    // Honor allow-self-signed flag or env var by setting env consumed in postgres::connection
+    let allow_self_signed_env = std::env::var("SEREN_ALLOW_SELF_SIGNED_CERTS")
+        .ok()
+        .map(|v| matches!(v.to_lowercase().as_str(), "1" | "true" | "yes"))
+        .unwrap_or(false)
+        // Backward compatibility with older env name
+        || std::env::var("SEREN_ALLOW_INVALID_CERTS")
+            .ok()
+            .map(|v| matches!(v.to_lowercase().as_str(), "1" | "true" | "yes"))
+            .unwrap_or(false);
+
+    if cli.allow_self_signed_certs || allow_self_signed_env {
+        // Set both names for compatibility
+        std::env::set_var("SEREN_ALLOW_SELF_SIGNED_CERTS", "1");
+        std::env::set_var("SEREN_ALLOW_INVALID_CERTS", "1");
+        tracing::warn!("Allowing self-signed/invalid TLS certificates (insecure)");
+    }
+
     match cli.command {
         Commands::Validate {
             source,
@@ -402,7 +423,7 @@ async fn init_remote(
     drop_existing: bool,
     no_sync: bool,
     seren_api: String,
-    _job_timeout: u64,
+    job_timeout: u64,
 ) -> anyhow::Result<()> {
     use database_replicator::migration;
     use database_replicator::postgres;
@@ -466,6 +487,8 @@ async fn init_remote(
     } else {
         Some(FilterSpec {
             include_databases,
+            exclude_databases,
+            include_tables,
             exclude_tables,
         })
     };
@@ -481,7 +504,12 @@ async fn init_remote(
         "estimated_size_bytes".to_string(),
         serde_json::Value::Number(serde_json::Number::from(estimated_size_bytes)),
     );
-    // Note: "yes" and "job_timeout" are client-side only options, not sent to server
+    // Optional timeout hint for remote orchestrator
+    options.insert(
+        "job_timeout_seconds".to_string(),
+        serde_json::Value::Number(serde_json::Number::from(job_timeout as i64)),
+    );
+    // Note: "yes" is client-side only, not sent to server
 
     let job_spec = JobSpec {
         version: "1.0".to_string(),
diff --git a/src/migration/dump.rs b/src/migration/dump.rs
@@ -65,6 +65,7 @@ pub async fn dump_globals(source_url: &str, output_path: &str) -> Result<()> {
         Duration::from_secs(1), // Start with 1 second delay
         "pg_dumpall (dump globals)",
     )
+    .await
     .context(
         "pg_dumpall failed to dump global objects.\n\
          \n\
@@ -172,6 +173,7 @@ pub async fn dump_schema(
         Duration::from_secs(1), // Start with 1 second delay
         "pg_dump (dump schema)",
     )
+    .await
     .with_context(|| {
         format!(
             "pg_dump failed to dump schema for database '{}'.\n\
@@ -299,6 +301,7 @@ pub async fn dump_data(
         Duration::from_secs(1), // Start with 1 second delay
         "pg_dump (dump data)",
     )
+    .await
     .with_context(|| {
         format!(
             "pg_dump failed to dump data for database '{}'.\n\
diff --git a/src/migration/restore.rs b/src/migration/restore.rs
@@ -62,7 +62,8 @@ pub async fn restore_globals(target_url: &str, input_path: &str) -> Result<()> {
         3,                      // Max 3 retries
         Duration::from_secs(1), // Start with 1 second delay
         "psql (restore globals)",
-    );
+    )
+    .await;
 
     // Handle result - don't fail on warnings for global objects
     match result {
@@ -136,6 +137,7 @@ pub async fn restore_schema(target_url: &str, input_path: &str) -> Result<()> {
         Duration::from_secs(1), // Start with 1 second delay
         "psql (restore schema)",
     )
+    .await
     .context(
         "Schema restoration failed.\n\
          \n\
@@ -228,6 +230,7 @@ pub async fn restore_data(target_url: &str, input_path: &str) -> Result<()> {
         Duration::from_secs(1), // Start with 1 second delay
         "pg_restore (restore data)",
     )
+    .await
     .context(
         "Data restoration failed.\n\
          \n\
diff --git a/src/mysql/mod.rs b/src/mysql/mod.rs
@@ -44,11 +44,7 @@ pub fn validate_mysql_url(connection_string: &str) -> Result<String> {
     }
 
     if !connection_string.starts_with("mysql://") {
-        bail!(
-            "Invalid MySQL connection string '{}'. \
-             Must start with 'mysql://'",
-            connection_string
-        );
+        bail!("Invalid MySQL connection string. Must start with 'mysql://'");
     }
 
     tracing::debug!("Validated MySQL connection string");
diff --git a/src/postgres/connection.rs b/src/postgres/connection.rs
@@ -130,10 +130,24 @@ pub async fn connect(connection_string: &str) -> Result<Client> {
     )?;
 
     // Set up TLS connector for cloud connections
-    // TEMPORARY: Accept invalid certs to debug TLS issues
-    // TODO: Remove this once we identify the certificate validation issue
-    let tls_connector = TlsConnector::builder()
-        .danger_accept_invalid_certs(true)
+    // By default, require valid certificates. Allow opt-in for self-signed/invalid certs via env.
+    let allow_self_signed = std::env::var("SEREN_ALLOW_SELF_SIGNED_CERTS")
+        .ok()
+        .map(|v| matches!(v.to_lowercase().as_str(), "1" | "true" | "yes"))
+        .unwrap_or(false)
+        // Backward compatibility with older env name
+        || std::env::var("SEREN_ALLOW_INVALID_CERTS")
+            .ok()
+            .map(|v| matches!(v.to_lowercase().as_str(), "1" | "true" | "yes"))
+            .unwrap_or(false);
+
+    let mut tls_builder = TlsConnector::builder();
+    if allow_self_signed {
+        tracing::warn!("Accepting self-signed/invalid TLS certificates");
+        tls_builder.danger_accept_invalid_certs(true);
+    }
+
+    let tls_connector = tls_builder
         .build()
         .context("Failed to build TLS connector")?;
     let tls = MakeTlsConnector::new(tls_connector);
diff --git a/src/remote/models.rs b/src/remote/models.rs
@@ -19,6 +19,8 @@ pub struct JobSpec {
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct FilterSpec {
     pub include_databases: Option<Vec<String>>,
+    pub exclude_databases: Option<Vec<String>>,
+    pub include_tables: Option<Vec<String>>,
     pub exclude_tables: Option<Vec<String>>,
 }
 
diff --git a/src/replication/monitor.rs b/src/replication/monitor.rs
@@ -34,41 +34,47 @@ pub async fn get_replication_lag(
     client: &Client,
     subscription_name: Option<&str>,
 ) -> Result<Vec<SourceReplicationStats>> {
-    let query = if let Some(sub_name) = subscription_name {
-        format!(
-            "SELECT
-                application_name,
-                state,
-                sent_lsn::text,
-                write_lsn::text,
-                flush_lsn::text,
-                replay_lsn::text,
-                EXTRACT(EPOCH FROM write_lag) * 1000 as write_lag_ms,
-                EXTRACT(EPOCH FROM flush_lag) * 1000 as flush_lag_ms,
-                EXTRACT(EPOCH FROM replay_lag) * 1000 as replay_lag_ms
-            FROM pg_stat_replication
-            WHERE application_name = '{}'",
-            sub_name
-        )
-    } else {
-        "SELECT
-            application_name,
-            state,
-            sent_lsn::text,
-            write_lsn::text,
-            flush_lsn::text,
-            replay_lsn::text,
-            EXTRACT(EPOCH FROM write_lag) * 1000 as write_lag_ms,
-            EXTRACT(EPOCH FROM flush_lag) * 1000 as flush_lag_ms,
-            EXTRACT(EPOCH FROM replay_lag) * 1000 as replay_lag_ms
-        FROM pg_stat_replication"
-            .to_string()
-    };
+    if let Some(name) = subscription_name {
+        crate::utils::validate_postgres_identifier(name).context("Invalid subscription name")?;
+    }
 
-    let rows = client
-        .query(&query, &[])
-        .await
-        .context("Failed to query replication statistics")?;
+    let rows = if let Some(sub_name) = subscription_name {
+        client
+            .query(
+                "SELECT
+                    application_name,
+                    state,
+                    sent_lsn::text,
+                    write_lsn::text,
+                    flush_lsn::text,
+                    replay_lsn::text,
+                    EXTRACT(EPOCH FROM write_lag) * 1000 as write_lag_ms,
+                    EXTRACT(EPOCH FROM flush_lag) * 1000 as flush_lag_ms,
+                    EXTRACT(EPOCH FROM replay_lag) * 1000 as replay_lag_ms
+                FROM pg_stat_replication
+                WHERE application_name = $1",
+                &[&sub_name],
+            )
+            .await
+    } else {
+        client
+            .query(
+                "SELECT
+                    application_name,
+                    state,
+                    sent_lsn::text,
+                    write_lsn::text,
+                    flush_lsn::text,
+                    replay_lsn::text,
+                    EXTRACT(EPOCH FROM write_lag) * 1000 as write_lag_ms,
+                    EXTRACT(EPOCH FROM flush_lag) * 1000 as flush_lag_ms,
+                    EXTRACT(EPOCH FROM replay_lag) * 1000 as replay_lag_ms
+                FROM pg_stat_replication",
+                &[],
+            )
+            .await
+    }
+    .context("Failed to query replication statistics")?;
 
     let mut stats = Vec::new();
     for row in rows {
@@ -94,33 +100,39 @@ pub async fn get_subscription_status(
     client: &Client,
     subscription_name: Option<&str>,
 ) -> Result<Vec<SubscriptionStats>> {
-    let query = if let Some(sub_name) = subscription_name {
-        format!(
-            "SELECT
-                subname,
-                pid,
-                received_lsn::text,
-                latest_end_lsn::text,
-                srsubstate
-            FROM pg_stat_subscription
-            WHERE subname = '{}'",
-            sub_name
-        )
-    } else {
-        "SELECT
-            subname,
-            pid,
-            received_lsn::text,
-            latest_end_lsn::text,
-            srsubstate
-        FROM pg_stat_subscription"
-            .to_string()
-    };
+    if let Some(name) = subscription_name {
+        crate::utils::validate_postgres_identifier(name).context("Invalid subscription name")?;
+    }
 
-    let rows = client
-        .query(&query, &[])
-        .await
-        .context("Failed to query subscription statistics")?;
+    let rows = if let Some(sub_name) = subscription_name {
+        client
+            .query(
+                "SELECT
+                    subname,
+                    pid,
+                    received_lsn::text,
+                    latest_end_lsn::text,
+                    srsubstate
+                FROM pg_stat_subscription
+                WHERE subname = $1",
+                &[&sub_name],
+            )
+            .await
+    } else {
+        client
+            .query(
+                "SELECT
+                    subname,
+                    pid,
+                    received_lsn::text,
+                    latest_end_lsn::text,
+                    srsubstate
+                FROM pg_stat_subscription",
+                &[],
+            )
+            .await
+    }
+    .context("Failed to query subscription statistics")?;
 
     let mut stats = Vec::new();
     for row in rows {
diff --git a/src/replication/subscription.rs b/src/replication/subscription.rs
@@ -50,9 +50,12 @@ pub async fn create_subscription(
         "  To avoid storing passwords, configure .pgpass on the target PostgreSQL server"
     );
 
+    // Escape single quotes in connection string to avoid breaking SQL literal
+    let escaped_conn = source_connection_string.replace('\'', "''");
+
     let query = format!(
         "CREATE SUBSCRIPTION \"{}\" CONNECTION '{}' PUBLICATION \"{}\"",
-        subscription_name, source_connection_string, publication_name
+        subscription_name, escaped_conn, publication_name
     );
 
     match client.execute(&query, &[]).await {
diff --git a/src/utils.rs b/src/utils.rs
@@ -271,11 +271,11 @@ where
 ///     3,  // Try up to 3 times
 ///     Duration::from_secs(1),  // Start with 1s delay
 ///     "psql"
-/// )?;
+/// ).await?;
 /// # Ok(())
 /// # }
 /// ```
-pub fn retry_subprocess_with_backoff<F>(
+pub async fn retry_subprocess_with_backoff<F>(
     mut operation: F,
     max_retries: u32,
     initial_delay: Duration,
@@ -328,7 +328,7 @@ where
                         last_error.as_ref().unwrap(),
                         delay
                     );
-                    std::thread::sleep(delay);
+                    tokio::time::sleep(delay).await;
                     delay *= 2; // Exponential backoff
                 }
             }
@@ -561,7 +561,18 @@ pub fn validate_source_target_different(source_url: &str, target_url: &str) -> R
         && source_parts.user == target_parts.user
     {
         bail!(
-            "Source and target URLs point to the same database!\\n\\\n             \\n\\\n             This would cause DATA LOSS - the target would overwrite the source.\\n\\\n             \\n\\\n             Source: {}@{}:{}/{}\\n\\\n             Target: {}@{}:{}/{}\\n\\\n             \\n\\\n             Please ensure source and target are different databases.\\n\\\n             Common causes:\\n\\\n             - Copy-paste error in connection strings\\n\\\n             - Wrong environment variables (e.g., SOURCE_URL == TARGET_URL)\\n\\\n             - Typo in database name or host",
+            "Source and target URLs point to the same database!\n\
+             \n\
+             This would cause DATA LOSS - the target would overwrite the source.\n\
+             \n\
+             Source: {}@{}:{}/{}\n\
+             Target: {}@{}:{}/{}\n\
+             \n\
+             Please ensure source and target are different databases.\n\
+             Common causes:\n\
+             - Copy-paste error in connection strings\n\
+             - Wrong environment variables (e.g., SOURCE_URL == TARGET_URL)\n\
+             - Typo in database name or host",
             source_parts.user.as_deref().unwrap_or("(no user)"),
             source_parts.host,
             source_parts.port,

Original file line number	Diff line number	Diff line change
`@@ -44,11 +44,7 @@ pub fn validate_mysql_url(connection_string: &str) -> Result<String> {`
`44`	`44`	`}`
`45`	`45`
`46`	`46`	`if !connection_string.starts_with("mysql://") {`
`47`		`- bail!(`
`48`		`- "Invalid MySQL connection string '{}'. \`
`49`		`- Must start with 'mysql://'",`
`50`		`- connection_string`
`51`		`- );`
	`47`	`+ bail!("Invalid MySQL connection string. Must start with 'mysql://'");`
`52`	`48`	`}`
`53`	`49`
`54`	`50`	`tracing::debug!("Validated MySQL connection string");`
Original file line number	Diff line number	Diff line change
`@@ -19,6 +19,8 @@ pub struct JobSpec {`
`19`	`19`	`#[derive(Debug, Clone, Serialize, Deserialize)]`
`20`	`20`	`pub struct FilterSpec {`
`21`	`21`	`pub include_databases: Option<Vec<String>>,`
	`22`	`+ pub exclude_databases: Option<Vec<String>>,`
	`23`	`+ pub include_tables: Option<Vec<String>>,`
`22`	`24`	`pub exclude_tables: Option<Vec<String>>,`
`23`	`25`	`}`
`24`	`26`