Write unit tests, disable some tests if PG<=15,

abhi-airspace-intelligence · abhi-airspace-intelligence · commit 49b53b676476 · 2025-08-29T13:10:32.000-04:00
diff --git a/etl-postgres/src/replication/db.rs b/etl-postgres/src/replication/db.rs
@@ -1,3 +1,5 @@
+use std::num::NonZeroI32;
+
 use etl_config::shared::{IntoConnectOptions, PgConnectionConfig};
 use sqlx::{PgPool, Row, postgres::PgPoolOptions};
 use thiserror::Error;
@@ -68,3 +70,107 @@ pub async fn get_table_name_from_oid(
         None => Err(TableLookupError::TableNotFound(table_id)),
     }
 }
+
+/// Extracts the PostgreSQL server version from a version string.
+///
+/// This function parses version strings like "15.5 (Homebrew)" or "14.2"
+/// and converts them to the numeric format used by PostgreSQL.
+///
+/// Returns the version in the format: MAJOR * 10000 + MINOR * 100 + PATCH
+/// For example: PostgreSQL 14.2 = 140200, PostgreSQL 15.1 = 150100
+///
+/// Returns `None` if the version string cannot be parsed or results in zero.
+pub fn extract_server_version(server_version_str: impl AsRef<str>) -> Option<NonZeroI32> {
+    // Parse version string like "15.5 (Homebrew)" or "14.2"
+    let version_part = server_version_str
+        .as_ref()
+        .split_whitespace()
+        .next()
+        .unwrap_or("0.0");
+
+    let version_components: Vec<&str> = version_part.split('.').collect();
+
+    let major = version_components
+        .first()
+        .and_then(|v| v.parse::<i32>().ok())
+        .unwrap_or(0);
+    let minor = version_components
+        .get(1)
+        .and_then(|v| v.parse::<i32>().ok())
+        .unwrap_or(0);
+    let patch = version_components
+        .get(2)
+        .and_then(|v| v.parse::<i32>().ok())
+        .unwrap_or(0);
+
+    let version = major * 10000 + minor * 100 + patch;
+
+    NonZeroI32::new(version)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_extract_server_version_basic_versions() {
+        assert_eq!(extract_server_version("15.5"), NonZeroI32::new(150500));
+        assert_eq!(extract_server_version("14.2"), NonZeroI32::new(140200));
+        assert_eq!(extract_server_version("13.0"), NonZeroI32::new(130000));
+        assert_eq!(extract_server_version("16.1"), NonZeroI32::new(160100));
+    }
+
+    #[test]
+    fn test_extract_server_version_with_suffixes() {
+        assert_eq!(
+            extract_server_version("15.5 (Homebrew)"),
+            NonZeroI32::new(150500)
+        );
+        assert_eq!(
+            extract_server_version("14.2 on x86_64-pc-linux-gnu"),
+            NonZeroI32::new(140200)
+        );
+        assert_eq!(
+            extract_server_version("13.7 (Ubuntu 13.7-1.pgdg20.04+1)"),
+            NonZeroI32::new(130700)
+        );
+        assert_eq!(
+            extract_server_version("16.0 devel"),
+            NonZeroI32::new(160000)
+        );
+    }
+
+    #[test]
+    fn test_extract_server_version_patch_versions() {
+        // Test versions with patch numbers
+        assert_eq!(extract_server_version("15.5.1"), NonZeroI32::new(150501));
+        assert_eq!(extract_server_version("14.10.3"), NonZeroI32::new(141003));
+        assert_eq!(extract_server_version("13.12.25"), NonZeroI32::new(131225));
+    }
+
+    #[test]
+    fn test_extract_server_version_invalid_inputs() {
+        // Test invalid inputs that should return None
+        assert_eq!(extract_server_version(""), None);
+        assert_eq!(extract_server_version("invalid"), None);
+        assert_eq!(extract_server_version("not.a.version"), None);
+        assert_eq!(extract_server_version("PostgreSQL"), None);
+        assert_eq!(extract_server_version("   "), None);
+    }
+
+    #[test]
+    fn test_extract_server_version_zero_versions() {
+        assert_eq!(extract_server_version("0.0.0"), None);
+        assert_eq!(extract_server_version("0.0"), None);
+    }
+
+    #[test]
+    fn test_extract_server_version_whitespace_handling() {
+        assert_eq!(extract_server_version("  15.5  "), NonZeroI32::new(150500));
+        assert_eq!(
+            extract_server_version("15.5\t(Homebrew)"),
+            NonZeroI32::new(150500)
+        );
+        assert_eq!(extract_server_version("15.5\n"), NonZeroI32::new(150500));
+    }
+}
diff --git a/etl-postgres/src/tokio/test_utils.rs b/etl-postgres/src/tokio/test_utils.rs
@@ -1,9 +1,12 @@
+use std::num::NonZeroI32;
+
 use etl_config::shared::{IntoConnectOptions, PgConnectionConfig};
 use tokio::runtime::Handle;
 use tokio_postgres::types::{ToSql, Type};
 use tokio_postgres::{Client, GenericClient, NoTls, Transaction};
 use tracing::info;
 
+use crate::replication::extract_server_version;
 use crate::types::{ColumnSchema, TableId, TableName};
 
 /// Table modification operations for ALTER TABLE statements.
@@ -34,10 +37,15 @@ pub enum TableModification<'a> {
 pub struct PgDatabase<G> {
     pub config: PgConnectionConfig,
     pub client: Option<G>,
+    server_version: Option<NonZeroI32>,
     destroy_on_drop: bool,
 }
 
 impl<G: GenericClient> PgDatabase<G> {
+    pub fn server_version(&self) -> Option<NonZeroI32> {
+        self.server_version
+    }
+
     /// Creates a Postgres publication for the specified tables.
     ///
     /// Sets up logical replication by creating a publication that includes
@@ -71,19 +79,51 @@ impl<G: GenericClient> PgDatabase<G> {
         publication_name: &str,
         schema: Option<&str>,
     ) -> Result<(), tokio_postgres::Error> {
-        let create_publication_query = match schema {
-            Some(schema_name) => format!(
-                "create publication {} for tables in schema {}",
-                publication_name, schema_name
-            ),
-            None => format!("create publication {} for all tables", publication_name),
-        };
-
-        self.client
-            .as_ref()
-            .unwrap()
-            .execute(&create_publication_query, &[])
-            .await?;
+        let client = self.client.as_ref().unwrap();
+
+        if let Some(server_version) = self.server_version
+            && server_version.get() >= 150000
+        {
+            // PostgreSQL 15+ supports FOR ALL TABLES IN SCHEMA syntax
+            let create_publication_query = match schema {
+                Some(schema_name) => format!(
+                    "create publication {} for tables in schema {}",
+                    publication_name, schema_name
+                ),
+                None => format!("create publication {} for all tables", publication_name),
+            };
+
+            client.execute(&create_publication_query, &[]).await?;
+        } else {
+            // PostgreSQL 14 and earlier: create publication and add tables individually
+            match schema {
+                Some(schema_name) => {
+                    let create_pub_query = format!("create publication {}", publication_name);
+                    client.execute(&create_pub_query, &[]).await?;
+
+                    let tables_query = format!(
+                        "select schemaname, tablename from pg_tables where schemaname = '{}'",
+                        schema_name
+                    );
+                    let rows = client.query(&tables_query, &[]).await?;
+
+                    for row in rows {
+                        let schema: String = row.get(0);
+                        let table: String = row.get(1);
+                        let add_table_query = format!(
+                            "alter publication {} add table {}.{}",
+                            publication_name, schema, table
+                        );
+                        client.execute(&add_table_query, &[]).await?;
+                    }
+                }
+                None => {
+                    let create_publication_query =
+                        format!("create publication {} for all tables", publication_name);
+                    client.execute(&create_publication_query, &[]).await?;
+                }
+            }
+        }
 
         Ok(())
     }
@@ -369,7 +409,8 @@ impl PgDatabase<Client> {
 
         Self {
             config,
-            client: Some(client),
+            client: Some(client.0),
+            server_version: client.1,
             destroy_on_drop: true,
         }
     }
@@ -386,7 +427,8 @@ impl PgDatabase<Client> {
 
         Self {
             config,
-            client: Some(client),
+            client: Some(client.0),
+            server_version: client.1,
             destroy_on_drop: true,
         }
     }
@@ -401,6 +443,7 @@ impl PgDatabase<Client> {
         PgDatabase {
             config: self.config.clone(),
             client: Some(transaction),
+            server_version: self.server_version,
             destroy_on_drop: false,
         }
     }
@@ -450,7 +493,7 @@ pub fn id_column_schema() -> ColumnSchema {
 ///
 /// # Panics
 /// Panics if connection or database creation fails.
-pub async fn create_pg_database(config: &PgConnectionConfig) -> Client {
+pub async fn create_pg_database(config: &PgConnectionConfig) -> (Client, Option<NonZeroI32>) {
     // Create the database via a single connection
     let (client, connection) = {
         let config: tokio_postgres::Config = config.without_db();
@@ -474,14 +517,16 @@ pub async fn create_pg_database(config: &PgConnectionConfig) -> Client {
         .expect("Failed to create database");
 
     // Connects to the actual Postgres database
-    connect_to_pg_database(config).await
+    let (client, server_version) = connect_to_pg_database(config).await;
+
+    (client, server_version)
 }
 
 /// Connects to an existing Postgres database.
 ///
 /// Establishes a client connection to the database specified in the configuration.
 /// Assumes the database already exists.
-pub async fn connect_to_pg_database(config: &PgConnectionConfig) -> Client {
+pub async fn connect_to_pg_database(config: &PgConnectionConfig) -> (Client, Option<NonZeroI32>) {
     // Create a new client connected to the created database
     let (client, connection) = {
         let config: tokio_postgres::Config = config.with_db();
@@ -490,6 +535,9 @@ pub async fn connect_to_pg_database(config: &PgConnectionConfig) -> Client {
             .await
             .expect("Failed to connect to Postgres")
     };
+    let server_version = connection
+        .parameter("server_version")
+        .and_then(extract_server_version);
 
     // Spawn the connection on a new task
     tokio::spawn(async move {
@@ -498,7 +546,7 @@ pub async fn connect_to_pg_database(config: &PgConnectionConfig) -> Client {
         }
     });
 
-    client
+    (client, server_version)
 }
 
 /// Drops a Postgres database and cleans up all resources.
diff --git a/etl/src/replication/client.rs b/etl/src/replication/client.rs
@@ -2,6 +2,7 @@ use crate::error::{ErrorKind, EtlError, EtlResult};
 use crate::utils::tokio::MakeRustlsConnect;
 use crate::{bail, etl_error};
 use etl_config::shared::{IntoConnectOptions, PgConnectionConfig};
+use etl_postgres::replication::extract_server_version;
 use etl_postgres::types::convert_type_oid_to_type;
 use etl_postgres::types::{ColumnSchema, TableId, TableName, TableSchema};
 use pg_escape::{quote_identifier, quote_literal};
@@ -13,7 +14,6 @@ use std::io::BufReader;
 use std::num::NonZeroI32;
 use std::sync::Arc;
 
-use tokio::io::{AsyncRead, AsyncWrite};
 use tokio_postgres::error::SqlState;
 use tokio_postgres::tls::MakeTlsConnect;
 use tokio_postgres::{
@@ -190,7 +190,9 @@ impl PgReplicationClient {
 
         let (client, connection) = config.connect(NoTls).await?;
 
-        let server_version = Self::extract_server_version(&connection);
+        let server_version = connection
+            .parameter("server_version")
+            .and_then(extract_server_version);
 
         spawn_postgres_connection::<NoTls>(connection);
 
@@ -225,7 +227,9 @@ impl PgReplicationClient {
 
         let (client, connection) = config.connect(MakeRustlsConnect::new(tls_config)).await?;
 
-        let server_version = Self::extract_server_version(&connection);
+        let server_version = connection
+            .parameter("server_version")
+            .and_then(extract_server_version);
 
         spawn_postgres_connection::<MakeRustlsConnect>(connection);
 
@@ -732,49 +736,6 @@ impl PgReplicationClient {
 
         Ok(column_schemas)
     }
-
-    /// Extracts the PostgreSQL server version from connection parameters.
-    ///
-    /// This method should be called during connection establishment to extract
-    /// the server version from the parameter status messages sent by the server.
-    ///
-    /// Returns the version in the format: MAJOR * 10000 + MINOR * 100 + PATCH
-    /// This matches the format used by `SELECT version()`.
-    /// For example: PostgreSQL 14.2 = 140200, PostgreSQL 15.1 = 150100
-    fn extract_server_version<S>(connection: &Connection<Socket, S>) -> Option<NonZeroI32>
-    where
-        S: AsyncRead + AsyncWrite + Unpin + Send,
-    {
-        if let Some(server_version_str) = connection.parameter("server_version") {
-            // Parse version string like "15.5 (Homebrew)" or "14.2"
-            let version_part = server_version_str
-                .split_whitespace()
-                .next()
-                .unwrap_or("0.0");
-
-            let version_components: Vec<&str> = version_part.split('.').collect();
-
-            let major = version_components
-                .first()
-                .and_then(|v| v.parse::<i32>().ok())
-                .unwrap_or(0);
-            let minor = version_components
-                .get(1)
-                .and_then(|v| v.parse::<i32>().ok())
-                .unwrap_or(0);
-            let patch = version_components
-                .get(2)
-                .and_then(|v| v.parse::<i32>().ok())
-                .unwrap_or(0);
-
-            let version = major * 10000 + minor * 100 + patch;
-
-            NonZeroI32::new(version)
-        } else {
-            None
-        }
-    }
-
     /// Creates a COPY stream for reading data from a table using its OID.
     ///
     /// The stream will include only the specified columns and use text format.
diff --git a/etl/tests/integration/pipeline_test.rs b/etl/tests/integration/pipeline_test.rs
@@ -190,6 +190,13 @@ async fn publication_changes_are_correctly_handled() {
     table_1_done.notified().await;
     table_2_done.notified().await;
 
+    if let Some(server_version) = database.server_version()
+        && server_version.get() <= 150000
+    {
+        println!("Skipping test for PostgreSQL version <= 15");
+        return;
+    }
+
     // Insert one row in each table and wait for two insert events.
     let inserts_notify = destination
         .wait_for_events_count(vec![(EventType::Insert, 2)])
@@ -353,6 +360,13 @@ async fn publication_for_all_tables_in_schema_ignores_new_tables_until_restart()
     assert!(table_schemas.contains_key(&table_1_id));
     assert!(!table_schemas.contains_key(&table_2_id));
 
+    if let Some(server_version) = database.server_version()
+        && server_version.get() <= 150000
+    {
+        println!("Skipping test for PostgreSQL version <= 15");
+        return;
+    }
+
     // We restart the pipeline and verify that the new table is now processed.
     let mut pipeline = create_pipeline(
         &database.config,
