Primary key check (#573)

Author: levkk

pgdog/src/backend/replication/logical/error.rs

@@ -2,7 +2,7 @@ use std::num::ParseIntError;
 
 use thiserror::Error;
 
-use crate::net::ErrorResponse;
+use crate::{backend::replication::publisher::PublicationTable, net::ErrorResponse};
 
 #[derive(Debug, Error)]
 pub enum Error {
@@ -65,6 +65,9 @@ pub enum Error {
 
     #[error("no replicas available for table sync")]
     NoReplicas,
+
+    #[error("table {0} doesn't have a primary key")]
+    NoPrimaryKey(PublicationTable),
 }
 
 impl From<ErrorResponse> for Error {

pgdog/src/backend/replication/logical/publisher/queries.rs

@@ -3,6 +3,8 @@
 //! TODO: I think these are Postgres-version specific, so we need to handle that
 //! later. These were fetched from CREATE SUBSCRIPTION ran on Postgres 17.
 //!
+use std::fmt::Display;
+
 use crate::{
     backend::Server,
     net::{DataRow, Format},
@@ -28,6 +30,12 @@ pub struct PublicationTable {
     pub attributes: String,
 }
 
+impl Display for PublicationTable {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "\"{}\".\"{}\"", self.schema, self.name)
+    }
+}
+
 impl PublicationTable {
     pub async fn load(
         publication: &str,

pgdog/src/backend/replication/logical/publisher/table.rs

@@ -49,6 +49,15 @@ impl Table {
         Ok(results)
     }
 
+    /// Check that the table supports replication.
+    pub fn valid(&self) -> Result<(), Error> {
+        if !self.columns.iter().any(|c| c.identity) {
+            return Err(Error::NoPrimaryKey(self.table.clone()));
+        }
+
+        Ok(())
+    }
+
     /// Upsert record into table.
     pub fn insert(&self, upsert: bool) -> String {
         let names = format!(

pgdog/src/backend/replication/logical/subscriber/stream.rs

@@ -13,7 +13,7 @@ use pg_query::{
     protobuf::{InsertStmt, ParseResult},
     NodeEnum,
 };
-use tracing::trace;
+use tracing::{debug, trace};
 
 use super::super::{publisher::Table, Error};
 use crate::{
@@ -95,8 +95,11 @@ impl Statement {
             .flatten()
             .flatten()
     }
-}
 
+    fn query(&self) -> &str {
+        &self.query
+    }
+}
 #[derive(Debug)]
 pub struct StreamSubscriber {
     /// Destination cluster.
@@ -303,10 +306,17 @@ impl StreamSubscriber {
         if let Some(table) = table {
             // Prepare queries for this table. Prepared statements
             // are much faster.
+
+            table.valid()?;
+
             let insert = Statement::new(&table.insert(false))?;
             let upsert = Statement::new(&table.insert(true))?;
 
             for server in &mut self.connections {
+                for stmt in &[&insert, &upsert] {
+                    debug!("preparing \"{}\" [{}]", stmt.query(), server.addr());
+                }
+
                 server
                     .send(&vec![insert.parse().into(), upsert.parse().into(), Sync.into()].into())
                     .await?;