Store server version in OnceCell

Author: abhi-airspace-intelligence

etl/src/replication/client.rs

@@ -11,6 +11,7 @@ use std::collections::HashMap;
 use std::fmt;
 use std::io::BufReader;
 use std::sync::Arc;
+use tokio::sync::OnceCell;
 use tokio_postgres::error::SqlState;
 use tokio_postgres::tls::MakeTlsConnect;
 use tokio_postgres::{
@@ -162,7 +163,7 @@ impl PgReplicationSlotTransaction {
 /// and streaming changes from the database.
 #[derive(Debug, Clone)]
 pub struct PgReplicationClient {
-    client: Arc<Client>,
+    client: Arc<(Client, OnceCell<i32>)>,
 }
 
 impl PgReplicationClient {
@@ -177,6 +178,15 @@ impl PgReplicationClient {
         }
     }
 
+
+    // Convenience method to avoid having to access the client directly.
+    async fn simple_query(
+        &self,
+        query: &str,
+    ) -> Result<Vec<SimpleQueryMessage>, tokio_postgres::Error> {
+        self.client.0.simple_query(query).await
+    }
+
     /// Establishes a connection to Postgres without TLS encryption.
     ///
     /// The connection is configured for logical replication mode.
@@ -190,7 +200,7 @@ impl PgReplicationClient {
         info!("successfully connected to postgres without tls");
 
         Ok(PgReplicationClient {
-            client: Arc::new(client),
+            client: Arc::new((client, OnceCell::new())),
         })
     }
 
@@ -221,7 +231,7 @@ impl PgReplicationClient {
         info!("successfully connected to postgres with tls");
 
         Ok(PgReplicationClient {
-            client: Arc::new(client),
+            client: Arc::new((client, OnceCell::new())),
         })
     }
 
@@ -253,7 +263,7 @@ impl PgReplicationClient {
             quote_literal(slot_name)
         );
 
-        let results = self.client.simple_query(&query).await?;
+        let results = self.simple_query(&query).await?;
         for result in results {
             if let SimpleQueryMessage::Row(row) = result {
                 let confirmed_flush_lsn = Self::get_row_value::<PgLsn>(
@@ -315,7 +325,7 @@ impl PgReplicationClient {
             quote_identifier(slot_name)
         );
 
-        match self.client.simple_query(&query).await {
+        match self.simple_query(&query).await {
             Ok(_) => {
                 info!("successfully deleted replication slot '{}'", slot_name);
 
@@ -353,7 +363,7 @@ impl PgReplicationClient {
             "select 1 as exists from pg_publication where pubname = {};",
             quote_literal(publication)
         );
-        for msg in self.client.simple_query(&publication_exists_query).await? {
+        for msg in self.simple_query(&publication_exists_query).await? {
             if let SimpleQueryMessage::Row(_) = msg {
                 return Ok(true);
             }
@@ -373,7 +383,7 @@ impl PgReplicationClient {
         );
 
         let mut table_names = vec![];
-        for msg in self.client.simple_query(&publication_query).await? {
+        for msg in self.simple_query(&publication_query).await? {
             if let SimpleQueryMessage::Row(row) = msg {
                 let schema =
                     Self::get_row_value::<String>(&row, "schemaname", "pg_publication_tables")
@@ -403,7 +413,7 @@ impl PgReplicationClient {
         );
 
         let mut table_ids = vec![];
-        for msg in self.client.simple_query(&publication_query).await? {
+        for msg in self.simple_query(&publication_query).await? {
             if let SimpleQueryMessage::Row(row) = msg {
                 // For the sake of simplicity, we refer to the table oid as table id.
                 let table_id = Self::get_row_value::<TableId>(&row, "oid", "pg_class").await?;
@@ -441,7 +451,11 @@ impl PgReplicationClient {
             options
         );
 
-        let copy_stream = self.client.copy_both_simple::<bytes::Bytes>(&query).await?;
+        let copy_stream = self
+            .client
+            .0
+            .copy_both_simple::<bytes::Bytes>(&query)
+            .await?;
         let stream = LogicalReplicationStream::new(copy_stream);
 
         Ok(stream)
@@ -452,23 +466,22 @@ impl PgReplicationClient {
     /// The transaction doesn't make any assumptions about the snapshot in use, since this is a
     /// concern of the statements issued within the transaction.
     async fn begin_tx(&self) -> EtlResult<()> {
-        self.client
-            .simple_query("begin read only isolation level repeatable read;")
+        self.simple_query("begin read only isolation level repeatable read;")
             .await?;
 
         Ok(())
     }
 
     /// Commits the current transaction.
     async fn commit_tx(&self) -> EtlResult<()> {
-        self.client.simple_query("commit;").await?;
+        self.simple_query("commit;").await?;
 
         Ok(())
     }
 
     /// Rolls back the current transaction.
     async fn rollback_tx(&self) -> EtlResult<()> {
-        self.client.simple_query("rollback;").await?;
+        self.simple_query("rollback;").await?;
 
         Ok(())
     }
@@ -495,7 +508,7 @@ impl PgReplicationClient {
             quote_identifier(slot_name),
             snapshot_option
         );
-        match self.client.simple_query(&query).await {
+        match self.simple_query(&query).await {
             Ok(results) => {
                 for result in results {
                     if let SimpleQueryMessage::Row(row) = result {
@@ -595,7 +608,7 @@ impl PgReplicationClient {
             where c.oid = {table_id}",
         );
 
-        for message in self.client.simple_query(&table_info_query).await? {
+        for message in self.simple_query(&table_info_query).await? {
             if let SimpleQueryMessage::Row(row) = message {
                 let schema_name =
                     Self::get_row_value::<String>(&row, "schema_name", "pg_namespace").await?;
@@ -626,7 +639,7 @@ impl PgReplicationClient {
         publication: Option<&str>,
     ) -> EtlResult<Vec<ColumnSchema>> {
         let (pub_cte, pub_pred) = if let Some(publication) = publication {
-            let is_pg14_or_earlier = self.is_postgres_14_or_earlier().await?;
+            let is_pg14_or_earlier = self.get_server_version().await.unwrap_or(0) < 150000;
 
             if !is_pg14_or_earlier {
                 (
@@ -690,7 +703,7 @@ impl PgReplicationClient {
 
         let mut column_schemas = vec![];
 
-        for message in self.client.simple_query(&column_info_query).await? {
+        for message in self.simple_query(&column_info_query).await? {
             if let SimpleQueryMessage::Row(row) = message {
                 let name = Self::get_row_value::<String>(&row, "attname", "pg_attribute").await?;
                 let type_oid = Self::get_row_value::<u32>(&row, "atttypid", "pg_attribute").await?;
@@ -716,24 +729,36 @@ impl PgReplicationClient {
         Ok(column_schemas)
     }
 
-    async fn is_postgres_14_or_earlier(&self) -> EtlResult<bool> {
-        let version_query = "SHOW server_version_num";
-
-        for message in self.client.simple_query(version_query).await? {
-            if let SimpleQueryMessage::Row(row) = message {
-                let version_str =
-                    Self::get_row_value::<String>(&row, "server_version_num", "server_version_num")
+    /// Gets the PostgreSQL server version.
+    ///
+    /// Returns the version in the format: MAJOR * 10000 + MINOR * 100 + PATCH
+    /// For example: PostgreSQL 14.2 = 140200, PostgreSQL 15.1 = 150100
+    async fn get_server_version(&self) -> EtlResult<i32> {
+        let version = self
+            .client
+            .1
+            .get_or_try_init(|| async {
+                let version_query = "SHOW server_version_num";
+
+                for message in self.simple_query(version_query).await? {
+                    if let SimpleQueryMessage::Row(row) = message {
+                        let version_str = Self::get_row_value::<String>(
+                            &row,
+                            "server_version_num",
+                            "server_version_num",
+                        )
                         .await?;
-                let server_version: i32 = version_str.parse().unwrap_or(0);
+                        let version: i32 = version_str.parse().unwrap_or(0);
+                        return Ok::<_, EtlError>(version);
+                    }
+                }
 
-                // PostgreSQL version format is typically: MAJOR * 10000 + MINOR * 100 + PATCH
-                // For version 14.x.x, this would be 140000 + minor * 100 + patch
-                // For version 15.x.x, this would be 150000 + minor * 100 + patch
-                return Ok(server_version < 150000);
-            }
-        }
+                // If we can't determine, return 0 (which will be treated as very old version)
+                Ok(0)
+            })
+            .await?;
 
-        Ok(false)
+        Ok(*version)
     }
 
     /// Creates a COPY stream for reading data from a table using its OID.
@@ -759,7 +784,7 @@ impl PgReplicationClient {
             column_list
         );
 
-        let stream = self.client.copy_out_simple(&copy_query).await?;
+        let stream = self.client.0.copy_out_simple(&copy_query).await?;
 
         Ok(stream)
     }