14703640163

Committed 28 Apr 2025 08:29AM UTC coverage: 36.672%. Remained the same

Build # 14703640163

Build Type

Pull #106

github

Committed by

web-flow

Commit Message

Merge 011234590 into e9f9c3612

Pull Request Pull Request #106: fix incorrect syntax in start_replication command

Coverage Stats

0 of 2 new or added lines in 1 file covered. (0.0%)

2109 of 5751 relevant lines covered (36.67%)

15.2 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

0.0

/pg_replicate/src/clients/postgres.rs

use std::collections::HashMap;

use pg_escape::{quote_identifier, quote_literal};
use postgres_replication::LogicalReplicationStream;
use rustls::{pki_types::CertificateDer, ClientConfig};
use thiserror::Error;
use tokio_postgres::{
    config::{ReplicationMode, SslMode},
    types::{Kind, PgLsn, Type},
    Client as PostgresClient, Config, CopyOutStream, NoTls, SimpleQueryMessage,
};
use tokio_postgres_rustls::MakeRustlsConnect;
use tracing::{info, warn};

use crate::table::{ColumnSchema, TableId, TableName, TableSchema};

pub struct SlotInfo {
    pub confirmed_flush_lsn: PgLsn,
}

/// A client for Postgres logical replication
pub struct ReplicationClient {
    postgres_client: PostgresClient,
    in_txn: bool,
}

#[derive(Debug, Error)]
pub enum ReplicationClientError {
    #[error("tokio_postgres error: {0}")]
    TokioPostgresError(#[from] tokio_postgres::Error),

    #[error("column {0} is missing from table {1}")]
    MissingColumn(String, String),

    #[error("publication {0} doesn't exist")]
    MissingPublication(String),

    #[error("oid column is not a valid u32")]
    OidColumnNotU32,

    #[error("replica identity '{0}' not supported")]
    ReplicaIdentityNotSupported(String),

    #[error("type modifier column is not a valid u32")]
    TypeModifierColumnNotI32,

    #[error("column {0}'s type with oid {1} in relation {2} is not supported")]
    UnsupportedType(String, u32, String),

    #[error("table {0} doesn't exist")]
    MissingTable(TableName),

    #[error("not a valid PgLsn")]
    InvalidPgLsn,

    #[error("failed to create slot")]
    FailedToCreateSlot,

    #[error("rustls error: {0}")]
    RustlsError(#[from] rustls::Error),
}

impl ReplicationClient {
    /// Connect to a postgres database in logical replication mode without TLS
    pub async fn connect_no_tls(
        host: &str,
        port: u16,
        database: &str,
        username: &str,
        password: Option<String>,
    ) -> Result<ReplicationClient, ReplicationClientError> {
        info!("connecting to postgres without TLS");

        let mut config = Config::new();
        config
            .host(host)
            .port(port)
            .dbname(database)
            .user(username)
            .replication_mode(ReplicationMode::Logical);

        if let Some(password) = password {
            config.password(password);
        }

        let (postgres_client, connection) = config.connect(NoTls).await?;

        tokio::spawn(async move {
            info!("waiting for connection to terminate");
            if let Err(e) = connection.await {
                warn!("connection error: {}", e);
            }
        });

        info!("successfully connected to postgres");

        Ok(ReplicationClient {
            postgres_client,
            in_txn: false,
        })
    }

    /// Connect to a postgres database in logical replication mode with TLS
    pub async fn connect_tls(
        host: &str,
        port: u16,
        database: &str,
        username: &str,
        password: Option<String>,
        ssl_mode: SslMode,
        trusted_root_certs: Vec<CertificateDer<'static>>,
    ) -> Result<ReplicationClient, ReplicationClientError> {
        info!("connecting to postgres with TLS");

        let mut config = Config::new();
        config
            .host(host)
            .port(port)
            .dbname(database)
            .user(username)
            .ssl_mode(ssl_mode)
            .replication_mode(ReplicationMode::Logical);

        if let Some(password) = password {
            config.password(password);
        }

        let mut root_store = rustls::RootCertStore::empty();
        for trusted_root_cert in trusted_root_certs {
            root_store.add(trusted_root_cert)?;
        }
        let tls_config = ClientConfig::builder()
            .with_root_certificates(root_store)
            .with_no_client_auth();

        let tls = MakeRustlsConnect::new(tls_config);

        let (postgres_client, connection) = config.connect(tls).await?;

        tokio::spawn(async move {
            info!("waiting for connection to terminate");
            if let Err(e) = connection.await {
                warn!("connection error: {}", e);
            }
        });

        info!("successfully connected to postgres");

        Ok(ReplicationClient {
            postgres_client,
            in_txn: false,
        })
    }

    /// Starts a read-only trasaction with repeatable read isolation level
    pub async fn begin_readonly_transaction(&mut self) -> Result<(), ReplicationClientError> {
        self.postgres_client
            .simple_query("begin read only isolation level repeatable read;")
            .await?;
        self.in_txn = true;
        Ok(())
    }

    /// Commits a transaction
    pub async fn commit_txn(&mut self) -> Result<(), ReplicationClientError> {
        if self.in_txn {
            self.postgres_client.simple_query("commit;").await?;
            self.in_txn = false;
        }
        Ok(())
    }

    async fn rollback_txn(&mut self) -> Result<(), ReplicationClientError> {
        if self.in_txn {
            self.postgres_client.simple_query("rollback;").await?;
            self.in_txn = false;
        }
        Ok(())
    }

    /// Returns a [CopyOutStream] for a table
    pub async fn get_table_copy_stream(
        &self,
        table_name: &TableName,
        column_schemas: &[ColumnSchema],
    ) -> Result<CopyOutStream, ReplicationClientError> {
        let column_list = column_schemas
            .iter()
            .map(|col| quote_identifier(&col.name))
            .collect::<Vec<_>>()
            .join(", ");

        let copy_query = format!(
            r#"COPY {} ({column_list}) TO STDOUT WITH (FORMAT text);"#,
            table_name.as_quoted_identifier(),
        );

        let stream = self.postgres_client.copy_out_simple(&copy_query).await?;

        Ok(stream)
    }

    /// Returns a vector of columns of a table, optionally filtered by a publication's column list
    pub async fn get_column_schemas(
        &self,
        table_id: TableId,
        publication: Option<&str>,
    ) -> Result<Vec<ColumnSchema>, ReplicationClientError> {
        let (pub_cte, pub_pred) = if let Some(publication) = publication {
            (
                format!(
                    "with pub_attrs as (
                        select unnest(r.prattrs)
                        from pg_publication_rel r
                        left join pg_publication p on r.prpubid = p.oid
                        where p.pubname = {publication}
                        and r.prrelid = {table_id}
                    )",
                    publication = quote_literal(publication),
                ),
                "and (
                    case (select count(*) from pub_attrs)
                    when 0 then true
                    else (a.attnum in (select * from pub_attrs))
                    end
                )",
            )
        } else {
            ("".into(), "")
        };

        let column_info_query = format!(
            "{pub_cte}
            select a.attname,
                a.atttypid,
                a.atttypmod,
                a.attnotnull,
                coalesce(i.indisprimary, false) as primary
            from pg_attribute a
            left join pg_index i
                on a.attrelid = i.indrelid
                and a.attnum = any(i.indkey)
                and i.indisprimary = true
            where a.attnum > 0::int2
            and not a.attisdropped
            and a.attgenerated = ''
            and a.attrelid = {table_id}
            {pub_pred}
            order by a.attnum
            ",
        );

        let mut column_schemas = vec![];

        for message in self
            .postgres_client
            .simple_query(&column_info_query)
            .await?
        {
            if let SimpleQueryMessage::Row(row) = message {
                let name = row
                    .try_get("attname")?
                    .ok_or(ReplicationClientError::MissingColumn(
                        "attname".to_string(),
                        "pg_attribute".to_string(),
                    ))?
                    .to_string();

                let type_oid = row
                    .try_get("atttypid")?
                    .ok_or(ReplicationClientError::MissingColumn(
                        "atttypid".to_string(),
                        "pg_attribute".to_string(),
                    ))?
                    .parse()
                    .map_err(|_| ReplicationClientError::OidColumnNotU32)?;

                //TODO: For now we assume all types are simple, fix it later
                let typ = Type::from_oid(type_oid).unwrap_or(Type::new(
                    format!("unnamed(oid: {type_oid})"),
                    type_oid,
                    Kind::Simple,
                    "pg_catalog".to_string(),
                ));

                let modifier = row
                    .try_get("atttypmod")?
                    .ok_or(ReplicationClientError::MissingColumn(
                        "atttypmod".to_string(),
                        "pg_attribute".to_string(),
                    ))?
                    .parse()
                    .map_err(|_| ReplicationClientError::TypeModifierColumnNotI32)?;

                let nullable =
                    row.try_get("attnotnull")?
                        .ok_or(ReplicationClientError::MissingColumn(
                            "attnotnull".to_string(),
                            "pg_attribute".to_string(),
                        ))?
                        == "f";

                let primary =
                    row.try_get("primary")?
                        .ok_or(ReplicationClientError::MissingColumn(
                            "indisprimary".to_string(),
                            "pg_index".to_string(),
                        ))?
                        == "t";

                column_schemas.push(ColumnSchema {
                    name,
                    typ,
                    modifier,
                    nullable,
                    primary,
                })
            }
        }

        Ok(column_schemas)
    }

    pub async fn get_table_schemas(
        &self,
        table_names: &[TableName],
        publication: Option<&str>,
    ) -> Result<HashMap<TableId, TableSchema>, ReplicationClientError> {
        let mut table_schemas = HashMap::new();

        for table_name in table_names {
            let table_schema = self
                .get_table_schema(table_name.clone(), publication)
                .await?;
            if !table_schema.has_primary_keys() {
                warn!(
                    "table {} with id {} will not be copied because it has no primary key",
                    table_schema.table_name, table_schema.table_id
                );
                continue;
            }
            table_schemas.insert(table_schema.table_id, table_schema);
        }

        Ok(table_schemas)
    }

    async fn get_table_schema(
        &self,
        table_name: TableName,
        publication: Option<&str>,
    ) -> Result<TableSchema, ReplicationClientError> {
        let table_id = self
            .get_table_id(&table_name)
            .await?
            .ok_or(ReplicationClientError::MissingTable(table_name.clone()))?;
        let column_schemas = self.get_column_schemas(table_id, publication).await?;
        Ok(TableSchema {
            table_name,
            table_id,
            column_schemas,
        })
    }

    /// Returns the table id (called relation id in Postgres) of a table
    /// Also checks whether the replica identity is default or full and
    /// returns an error if not.
    pub async fn get_table_id(
        &self,
        table: &TableName,
    ) -> Result<Option<TableId>, ReplicationClientError> {
        let quoted_schema = quote_literal(&table.schema);
        let quoted_name = quote_literal(&table.name);

        let table_info_query = format!(
            "select c.oid,
                c.relreplident
            from pg_class c
            join pg_namespace n
                on (c.relnamespace = n.oid)
            where n.nspname = {}
                and c.relname = {}
            ",
            quoted_schema, quoted_name
        );

        for message in self.postgres_client.simple_query(&table_info_query).await? {
            if let SimpleQueryMessage::Row(row) = message {
                let replica_identity =
                    row.try_get("relreplident")?
                        .ok_or(ReplicationClientError::MissingColumn(
                            "relreplident".to_string(),
                            "pg_class".to_string(),
                        ))?;

                if !(replica_identity == "d" || replica_identity == "f") {
                    return Err(ReplicationClientError::ReplicaIdentityNotSupported(
                        replica_identity.to_string(),
                    ));
                }

                let oid: u32 = row
                    .try_get("oid")?
                    .ok_or(ReplicationClientError::MissingColumn(
                        "oid".to_string(),
                        "pg_class".to_string(),
                    ))?
                    .parse()
                    .map_err(|_| ReplicationClientError::OidColumnNotU32)?;
                return Ok(Some(oid));
            }
        }

        Ok(None)
    }

    /// Returns the slot info of an existing slot. The slot info currently only has the
    /// confirmed_flush_lsn column of the pg_replication_slots table.
    async fn get_slot(&self, slot_name: &str) -> Result<Option<SlotInfo>, ReplicationClientError> {
        let query = format!(
            r#"select confirmed_flush_lsn from pg_replication_slots where slot_name = {};"#,
            quote_literal(slot_name)
        );

        let query_result = self.postgres_client.simple_query(&query).await?;

        for res in &query_result {
            if let SimpleQueryMessage::Row(row) = res {
                let confirmed_flush_lsn = row
                    .get("confirmed_flush_lsn")
                    .ok_or(ReplicationClientError::MissingColumn(
                        "confirmed_flush_lsn".to_string(),
                        "pg_replication_slots".to_string(),
                    ))?
                    .parse()
                    .map_err(|_| ReplicationClientError::InvalidPgLsn)?;

                return Ok(Some(SlotInfo {
                    confirmed_flush_lsn,
                }));
            }
        }

        Ok(None)
    }

    /// Creates a logical replication slot. This will only succeed if the postgres connection
    /// is in logical replication mode. Otherwise it will fail with the following error:
    /// `syntax error at or near "CREATE_REPLICATION_SLOT"``
    ///
    /// Returns the consistent_point column as slot info.
    async fn create_slot(&self, slot_name: &str) -> Result<SlotInfo, ReplicationClientError> {
        let query = format!(
            r#"CREATE_REPLICATION_SLOT {} LOGICAL pgoutput USE_SNAPSHOT"#,
            quote_identifier(slot_name)
        );
        let results = self.postgres_client.simple_query(&query).await?;

        for result in results {
            if let SimpleQueryMessage::Row(row) = result {
                let consistent_point: PgLsn = row
                    .get("consistent_point")
                    .ok_or(ReplicationClientError::MissingColumn(
                        "consistent_point".to_string(),
                        "create_replication_slot".to_string(),
                    ))?
                    .parse()
                    .map_err(|_| ReplicationClientError::InvalidPgLsn)?;
                return Ok(SlotInfo {
                    confirmed_flush_lsn: consistent_point,
                });
            }
        }
        Err(ReplicationClientError::FailedToCreateSlot)
    }

    /// Either return the slot info of an existing slot or creates a new
    /// slot and returns its slot info.
    pub async fn get_or_create_slot(
        &mut self,
        slot_name: &str,
    ) -> Result<SlotInfo, ReplicationClientError> {
        if let Some(slot_info) = self.get_slot(slot_name).await? {
            Ok(slot_info)
        } else {
            self.rollback_txn().await?;
            self.begin_readonly_transaction().await?;
            Ok(self.create_slot(slot_name).await?)
        }
    }

    /// Returns all table names in a publication
    pub async fn get_publication_table_names(
        &self,
        publication: &str,
    ) -> Result<Vec<TableName>, ReplicationClientError> {
        let publication_query = format!(
            "select schemaname, tablename from pg_publication_tables where pubname = {};",
            quote_literal(publication)
        );

        let mut table_names = vec![];
        for msg in self
            .postgres_client
            .simple_query(&publication_query)
            .await?
        {
            if let SimpleQueryMessage::Row(row) = msg {
                let schema = row
                    .get(0)
                    .ok_or(ReplicationClientError::MissingColumn(
                        "schemaname".to_string(),
                        "pg_publication_tables".to_string(),
                    ))?
                    .to_string();

                let name = row
                    .get(1)
                    .ok_or(ReplicationClientError::MissingColumn(
                        "tablename".to_string(),
                        "pg_publication_tables".to_string(),
                    ))?
                    .to_string();

                table_names.push(TableName { schema, name })
            }
        }

        Ok(table_names)
    }

    pub async fn publication_exists(
        &self,
        publication: &str,
    ) -> Result<bool, ReplicationClientError> {
        let publication_exists_query = format!(
            "select 1 as exists from pg_publication where pubname = {};",
            quote_literal(publication)
        );
        for msg in self
            .postgres_client
            .simple_query(&publication_exists_query)
            .await?
        {
            if let SimpleQueryMessage::Row(_) = msg {
                return Ok(true);
            }
        }
        Ok(false)
    }

    pub async fn get_logical_replication_stream(
        &self,
        publication: &str,
        slot_name: &str,
        start_lsn: PgLsn,
    ) -> Result<LogicalReplicationStream, ReplicationClientError> {
        let options = format!(
            r#"("proto_version" '1', "publication_names" '{}')"#,
            quote_identifier(publication),
        );

        let query = format!(
            r#"START_REPLICATION SLOT {} LOGICAL {} {}"#,
            quote_identifier(slot_name),
            start_lsn,
            options
        );

        let copy_stream = self
            .postgres_client
            .copy_both_simple::<bytes::Bytes>(&query)
            .await?;

        let stream = LogicalReplicationStream::new(copy_stream);

        Ok(stream)
    }
}

1	use std::collections::HashMap;
2
3	use pg_escape::{quote_identifier, quote_literal};
4	use postgres_replication::LogicalReplicationStream;
5	use rustls::{pki_types::CertificateDer, ClientConfig};
6	use thiserror::Error;
7	use tokio_postgres::{
8	config::{ReplicationMode, SslMode},
9	types::{Kind, PgLsn, Type},
10	Client as PostgresClient, Config, CopyOutStream, NoTls, SimpleQueryMessage,
11	};
12	use tokio_postgres_rustls::MakeRustlsConnect;
13	use tracing::{info, warn};
14
15	use crate::table::{ColumnSchema, TableId, TableName, TableSchema};
16
17	pub struct SlotInfo {
18	pub confirmed_flush_lsn: PgLsn,
19	}
20
21	/// A client for Postgres logical replication
22	pub struct ReplicationClient {
23	postgres_client: PostgresClient,
24	in_txn: bool,
25	}
26
27	#[derive(Debug, Error)]
28	pub enum ReplicationClientError {
29	#[error("tokio_postgres error: {0}")]
30	TokioPostgresError(#[from] tokio_postgres::Error),
31
32	#[error("column {0} is missing from table {1}")]
33	MissingColumn(String, String),
34
35	#[error("publication {0} doesn't exist")]
36	MissingPublication(String),
37
38	#[error("oid column is not a valid u32")]
39	OidColumnNotU32,
40
41	#[error("replica identity '{0}' not supported")]
42	ReplicaIdentityNotSupported(String),
43
44	#[error("type modifier column is not a valid u32")]
45	TypeModifierColumnNotI32,
46
47	#[error("column {0}'s type with oid {1} in relation {2} is not supported")]
48	UnsupportedType(String, u32, String),
49
50	#[error("table {0} doesn't exist")]
51	MissingTable(TableName),
52
53	#[error("not a valid PgLsn")]
54	InvalidPgLsn,
55
56	#[error("failed to create slot")]
57	FailedToCreateSlot,
58
59	#[error("rustls error: {0}")]
60	RustlsError(#[from] rustls::Error),
61	}
62
63	impl ReplicationClient {
64	/// Connect to a postgres database in logical replication mode without TLS
65	pub async fn connect_no_tls(	×
66	host: &str,	×
67	port: u16,	×
68	database: &str,	×
69	username: &str,	×
70	password: Option<String>,	×
71	) -> Result<ReplicationClient, ReplicationClientError> {	×
72	info!("connecting to postgres without TLS");	×
73
74	let mut config = Config::new();	×
75	config	×
76	.host(host)	×
77	.port(port)	×
78	.dbname(database)	×
79	.user(username)	×
80	.replication_mode(ReplicationMode::Logical);	×
81
82	if let Some(password) = password {	×
83	config.password(password);	×
84	}	×
85
86	let (postgres_client, connection) = config.connect(NoTls).await?;	×
87
88	tokio::spawn(async move {	×
89	info!("waiting for connection to terminate");	×
90	if let Err(e) = connection.await {	×
91	warn!("connection error: {}", e);	×
92	}	×
93	});	×
94		×
95	info!("successfully connected to postgres");	×
96
97	Ok(ReplicationClient {	×
98	postgres_client,	×
99	in_txn: false,	×
100	})	×
101	}	×
102
103	/// Connect to a postgres database in logical replication mode with TLS
104	pub async fn connect_tls(	×
105	host: &str,	×
106	port: u16,	×
107	database: &str,	×
108	username: &str,	×
109	password: Option<String>,	×
110	ssl_mode: SslMode,	×
111	trusted_root_certs: Vec<CertificateDer<'static>>,	×
112	) -> Result<ReplicationClient, ReplicationClientError> {	×
113	info!("connecting to postgres with TLS");	×
114
115	let mut config = Config::new();	×
116	config	×
117	.host(host)	×
118	.port(port)	×
119	.dbname(database)	×
120	.user(username)	×
121	.ssl_mode(ssl_mode)	×
122	.replication_mode(ReplicationMode::Logical);	×
123
124	if let Some(password) = password {	×
125	config.password(password);	×
126	}	×
127
128	let mut root_store = rustls::RootCertStore::empty();	×
129	for trusted_root_cert in trusted_root_certs {	×
130	root_store.add(trusted_root_cert)?;	×
131	}
132	let tls_config = ClientConfig::builder()	×
133	.with_root_certificates(root_store)	×
134	.with_no_client_auth();	×
135		×
136	let tls = MakeRustlsConnect::new(tls_config);	×
137
138	let (postgres_client, connection) = config.connect(tls).await?;	×
139
140	tokio::spawn(async move {	×
141	info!("waiting for connection to terminate");	×
142	if let Err(e) = connection.await {	×
143	warn!("connection error: {}", e);	×
144	}	×
145	});	×
146		×
147	info!("successfully connected to postgres");	×
148
149	Ok(ReplicationClient {	×
150	postgres_client,	×
151	in_txn: false,	×
152	})	×
153	}	×
154
155	/// Starts a read-only trasaction with repeatable read isolation level
156	pub async fn begin_readonly_transaction(&mut self) -> Result<(), ReplicationClientError> {	×
157	self.postgres_client	×
158	.simple_query("begin read only isolation level repeatable read;")	×
159	.await?;	×
160	self.in_txn = true;	×
161	Ok(())	×
162	}	×
163
164	/// Commits a transaction
165	pub async fn commit_txn(&mut self) -> Result<(), ReplicationClientError> {	×
166	if self.in_txn {	×
167	self.postgres_client.simple_query("commit;").await?;	×
168	self.in_txn = false;	×
169	}	×
170	Ok(())	×
171	}	×
172
173	async fn rollback_txn(&mut self) -> Result<(), ReplicationClientError> {	×
174	if self.in_txn {	×
175	self.postgres_client.simple_query("rollback;").await?;	×
176	self.in_txn = false;	×
177	}	×
178	Ok(())	×
179	}	×
180
181	/// Returns a [CopyOutStream] for a table
182	pub async fn get_table_copy_stream(	×
183	&self,	×
184	table_name: &TableName,	×
185	column_schemas: &[ColumnSchema],	×
186	) -> Result<CopyOutStream, ReplicationClientError> {	×
187	let column_list = column_schemas	×
188	.iter()	×
189	.map(\|col\| quote_identifier(&col.name))	×
190	.collect::<Vec<_>>()	×
191	.join(", ");	×
192		×
193	let copy_query = format!(	×
194	r#"COPY {} ({column_list}) TO STDOUT WITH (FORMAT text);"#,	×
195	table_name.as_quoted_identifier(),	×
196	);	×
197
198	let stream = self.postgres_client.copy_out_simple(&copy_query).await?;	×
199
200	Ok(stream)	×
201	}	×
202
203	/// Returns a vector of columns of a table, optionally filtered by a publication's column list
204	pub async fn get_column_schemas(	×
205	&self,	×
206	table_id: TableId,	×
207	publication: Option<&str>,	×
208	) -> Result<Vec<ColumnSchema>, ReplicationClientError> {	×
209	let (pub_cte, pub_pred) = if let Some(publication) = publication {	×
210	(	×
211	format!(	×
212	"with pub_attrs as (	×
213	select unnest(r.prattrs)	×
214	from pg_publication_rel r	×
215	left join pg_publication p on r.prpubid = p.oid	×
216	where p.pubname = {publication}	×
217	and r.prrelid = {table_id}	×
218	)",	×
219	publication = quote_literal(publication),	×
220	),	×
221	"and (	×
222	case (select count(*) from pub_attrs)	×
223	when 0 then true	×
224	else (a.attnum in (select * from pub_attrs))	×
225	end	×
226	)",	×
227	)	×
228	} else {
229	("".into(), "")	×
230	};
231
232	let column_info_query = format!(	×
233	"{pub_cte}	×
234	select a.attname,	×
235	a.atttypid,	×
236	a.atttypmod,	×
237	a.attnotnull,	×
238	coalesce(i.indisprimary, false) as primary	×
239	from pg_attribute a	×
240	left join pg_index i	×
241	on a.attrelid = i.indrelid	×
242	and a.attnum = any(i.indkey)	×
243	and i.indisprimary = true	×
244	where a.attnum > 0::int2	×
245	and not a.attisdropped	×
246	and a.attgenerated = ''	×
247	and a.attrelid = {table_id}	×
248	{pub_pred}	×
249	order by a.attnum	×
250	",	×
251	);	×
252		×
253	let mut column_schemas = vec![];	×
254
255	for message in self	×
256	.postgres_client	×
257	.simple_query(&column_info_query)	×
258	.await?	×
259	{
260	if let SimpleQueryMessage::Row(row) = message {	×
261	let name = row	×
262	.try_get("attname")?	×
263	.ok_or(ReplicationClientError::MissingColumn(	×
264	"attname".to_string(),	×
265	"pg_attribute".to_string(),	×
266	))?	×
267	.to_string();	×
268
269	let type_oid = row	×
270	.try_get("atttypid")?	×
271	.ok_or(ReplicationClientError::MissingColumn(	×
272	"atttypid".to_string(),	×
273	"pg_attribute".to_string(),	×
274	))?	×
275	.parse()	×
276	.map_err(\|_\| ReplicationClientError::OidColumnNotU32)?;	×
277
278	//TODO: For now we assume all types are simple, fix it later
279	let typ = Type::from_oid(type_oid).unwrap_or(Type::new(	×
280	format!("unnamed(oid: {type_oid})"),	×
281	type_oid,	×
282	Kind::Simple,	×
283	"pg_catalog".to_string(),	×
284	));	×
285
286	let modifier = row	×
287	.try_get("atttypmod")?	×
288	.ok_or(ReplicationClientError::MissingColumn(	×
289	"atttypmod".to_string(),	×
290	"pg_attribute".to_string(),	×
291	))?	×
292	.parse()	×
293	.map_err(\|_\| ReplicationClientError::TypeModifierColumnNotI32)?;	×
294
295	let nullable =	×
296	row.try_get("attnotnull")?	×
297	.ok_or(ReplicationClientError::MissingColumn(	×
298	"attnotnull".to_string(),	×
299	"pg_attribute".to_string(),	×
300	))?	×
301	== "f";	×
302
303	let primary =	×
304	row.try_get("primary")?	×
305	.ok_or(ReplicationClientError::MissingColumn(	×
306	"indisprimary".to_string(),	×
307	"pg_index".to_string(),	×
308	))?	×
309	== "t";	×
310		×
311	column_schemas.push(ColumnSchema {	×
312	name,	×
313	typ,	×
314	modifier,	×
315	nullable,	×
316	primary,	×
317	})	×
318	}	×
319	}
320
321	Ok(column_schemas)	×
322	}	×
323
324	pub async fn get_table_schemas(	×
325	&self,	×
326	table_names: &[TableName],	×
327	publication: Option<&str>,	×
328	) -> Result<HashMap<TableId, TableSchema>, ReplicationClientError> {	×
329	let mut table_schemas = HashMap::new();	×
330
331	for table_name in table_names {	×
332	let table_schema = self	×
333	.get_table_schema(table_name.clone(), publication)	×
334	.await?;	×
335	if !table_schema.has_primary_keys() {	×
336	warn!(	×
337	"table {} with id {} will not be copied because it has no primary key",	×
338	table_schema.table_name, table_schema.table_id
339	);
340	continue;	×
341	}	×
342	table_schemas.insert(table_schema.table_id, table_schema);	×
343	}
344
345	Ok(table_schemas)	×
346	}	×
347
348	async fn get_table_schema(	×
349	&self,	×
350	table_name: TableName,	×
351	publication: Option<&str>,	×
352	) -> Result<TableSchema, ReplicationClientError> {	×
353	let table_id = self	×
354	.get_table_id(&table_name)	×
355	.await?	×
356	.ok_or(ReplicationClientError::MissingTable(table_name.clone()))?;	×
357	let column_schemas = self.get_column_schemas(table_id, publication).await?;	×
358	Ok(TableSchema {	×
359	table_name,	×
360	table_id,	×
361	column_schemas,	×
362	})	×
363	}	×
364
365	/// Returns the table id (called relation id in Postgres) of a table
366	/// Also checks whether the replica identity is default or full and
367	/// returns an error if not.
368	pub async fn get_table_id(	×
369	&self,	×
370	table: &TableName,	×
371	) -> Result<Option<TableId>, ReplicationClientError> {	×
372	let quoted_schema = quote_literal(&table.schema);	×
373	let quoted_name = quote_literal(&table.name);	×
374		×
375	let table_info_query = format!(	×
376	"select c.oid,	×
377	c.relreplident	×
378	from pg_class c	×
379	join pg_namespace n	×
380	on (c.relnamespace = n.oid)	×
381	where n.nspname = {}	×
382	and c.relname = {}	×
383	",	×
384	quoted_schema, quoted_name	×
385	);	×
386
387	for message in self.postgres_client.simple_query(&table_info_query).await? {	×
388	if let SimpleQueryMessage::Row(row) = message {	×
389	let replica_identity =	×
390	row.try_get("relreplident")?	×
391	.ok_or(ReplicationClientError::MissingColumn(	×
392	"relreplident".to_string(),	×
393	"pg_class".to_string(),	×
394	))?;	×
395
396	if !(replica_identity == "d" \|\| replica_identity == "f") {	×
397	return Err(ReplicationClientError::ReplicaIdentityNotSupported(	×
398	replica_identity.to_string(),	×
399	));	×
400	}	×
401
402	let oid: u32 = row	×
403	.try_get("oid")?	×
404	.ok_or(ReplicationClientError::MissingColumn(	×
405	"oid".to_string(),	×
406	"pg_class".to_string(),	×
407	))?	×
408	.parse()	×
409	.map_err(\|_\| ReplicationClientError::OidColumnNotU32)?;	×
410	return Ok(Some(oid));	×
411	}	×
412	}
413
414	Ok(None)	×
415	}	×
416
417	/// Returns the slot info of an existing slot. The slot info currently only has the
418	/// confirmed_flush_lsn column of the pg_replication_slots table.
419	async fn get_slot(&self, slot_name: &str) -> Result<Option<SlotInfo>, ReplicationClientError> {	×
420	let query = format!(	×
421	r#"select confirmed_flush_lsn from pg_replication_slots where slot_name = {};"#,	×
422	quote_literal(slot_name)	×
423	);	×
424
425	let query_result = self.postgres_client.simple_query(&query).await?;	×
426
427	for res in &query_result {	×
428	if let SimpleQueryMessage::Row(row) = res {	×
429	let confirmed_flush_lsn = row	×
430	.get("confirmed_flush_lsn")	×
431	.ok_or(ReplicationClientError::MissingColumn(	×
432	"confirmed_flush_lsn".to_string(),	×
433	"pg_replication_slots".to_string(),	×
434	))?	×
435	.parse()	×
436	.map_err(\|_\| ReplicationClientError::InvalidPgLsn)?;	×
437
438	return Ok(Some(SlotInfo {	×
439	confirmed_flush_lsn,	×
440	}));	×
441	}	×
442	}
443
444	Ok(None)	×
445	}	×
446
447	/// Creates a logical replication slot. This will only succeed if the postgres connection
448	/// is in logical replication mode. Otherwise it will fail with the following error:
449	/// `syntax error at or near "CREATE_REPLICATION_SLOT"``
450	///
451	/// Returns the consistent_point column as slot info.
452	async fn create_slot(&self, slot_name: &str) -> Result<SlotInfo, ReplicationClientError> {	×
453	let query = format!(	×
454	r#"CREATE_REPLICATION_SLOT {} LOGICAL pgoutput USE_SNAPSHOT"#,	×
455	quote_identifier(slot_name)	×
456	);	×
457	let results = self.postgres_client.simple_query(&query).await?;	×
458
459	for result in results {	×
460	if let SimpleQueryMessage::Row(row) = result {	×
461	let consistent_point: PgLsn = row	×
462	.get("consistent_point")	×
463	.ok_or(ReplicationClientError::MissingColumn(	×
464	"consistent_point".to_string(),	×
465	"create_replication_slot".to_string(),	×
466	))?	×
467	.parse()	×
468	.map_err(\|_\| ReplicationClientError::InvalidPgLsn)?;	×
469	return Ok(SlotInfo {	×
470	confirmed_flush_lsn: consistent_point,	×
471	});	×
472	}	×
473	}
474	Err(ReplicationClientError::FailedToCreateSlot)	×
475	}	×
476
477	/// Either return the slot info of an existing slot or creates a new
478	/// slot and returns its slot info.
479	pub async fn get_or_create_slot(	×
480	&mut self,	×
481	slot_name: &str,	×
482	) -> Result<SlotInfo, ReplicationClientError> {	×
483	if let Some(slot_info) = self.get_slot(slot_name).await? {	×
484	Ok(slot_info)	×
485	} else {
486	self.rollback_txn().await?;	×
487	self.begin_readonly_transaction().await?;	×
488	Ok(self.create_slot(slot_name).await?)	×
489	}
490	}	×
491
492	/// Returns all table names in a publication
493	pub async fn get_publication_table_names(	×
494	&self,	×
495	publication: &str,	×
496	) -> Result<Vec<TableName>, ReplicationClientError> {	×
497	let publication_query = format!(	×
498	"select schemaname, tablename from pg_publication_tables where pubname = {};",	×
499	quote_literal(publication)	×
500	);	×
501		×
502	let mut table_names = vec![];	×
503	for msg in self	×
504	.postgres_client	×
505	.simple_query(&publication_query)	×
506	.await?	×
507	{
508	if let SimpleQueryMessage::Row(row) = msg {	×
509	let schema = row	×
510	.get(0)	×
511	.ok_or(ReplicationClientError::MissingColumn(	×
512	"schemaname".to_string(),	×
513	"pg_publication_tables".to_string(),	×
514	))?	×
515	.to_string();	×
516
517	let name = row	×
518	.get(1)	×
519	.ok_or(ReplicationClientError::MissingColumn(	×
520	"tablename".to_string(),	×
521	"pg_publication_tables".to_string(),	×
522	))?	×
523	.to_string();	×
524		×
525	table_names.push(TableName { schema, name })	×
526	}	×
527	}
528
529	Ok(table_names)	×
530	}	×
531
532	pub async fn publication_exists(	×
533	&self,	×
534	publication: &str,	×
535	) -> Result<bool, ReplicationClientError> {	×
536	let publication_exists_query = format!(	×
537	"select 1 as exists from pg_publication where pubname = {};",	×
538	quote_literal(publication)	×
539	);	×
540	for msg in self	×
541	.postgres_client	×
542	.simple_query(&publication_exists_query)	×
543	.await?	×
544	{
545	if let SimpleQueryMessage::Row(_) = msg {	×
546	return Ok(true);	×
547	}	×
548	}
549	Ok(false)	×
550	}	×
551
552	pub async fn get_logical_replication_stream(	×
553	&self,	×
554	publication: &str,	×
555	slot_name: &str,	×
556	start_lsn: PgLsn,	×
557	) -> Result<LogicalReplicationStream, ReplicationClientError> {	×
558	let options = format!(	×
NEW 559	r#"("proto_version" '1', "publication_names" '{}')"#,	×
NEW 560	quote_identifier(publication),	×
561	);	×
562		×
563	let query = format!(	×
564	r#"START_REPLICATION SLOT {} LOGICAL {} {}"#,	×
565	quote_identifier(slot_name),	×
566	start_lsn,	×
567	options	×
568	);	×
569
570	let copy_stream = self	×
571	.postgres_client	×
572	.copy_both_simple::<bytes::Bytes>(&query)	×
573	.await?;	×
574
575	let stream = LogicalReplicationStream::new(copy_stream);	×
576		×
577	Ok(stream)	×
578	}	×
579	}

supabase / pg_replicate / 14703640163

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous