kalugin66/windmill
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: use rust-postgres client instead of sqlx for postgres trigger (#5853)

Browse Source 
* use rust-postgres client instead of sqlx

* fix

* Update backend/windmill-api/src/postgres_triggers/mod.rs

Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>

* fix import

---------

Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>
This commit is contained in:
dieriba
2025-06-04 18:53:30 +02:00
committed by GitHub
parent 10534657f4
commit 55eaf3f2bc
4 changed files with 512 additions and 520 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
backend/windmill-api/src/capture.rs
View File

@@ -55,12 +55,9 @@ use crate::mqtt_triggers::{MqttClientVersion, MqttV3Config, MqttV5Config, Subscr
use crate::nats_triggers_oss::NatsTriggerConfigConnection;
#[cfg(feature = "postgres_trigger")]
use {
crate::postgres_triggers::{
create_logical_replication_slot, create_pg_publication, generate_random_string,
get_pg_connection, PublicationData,
},
sqlx::Connection,
use crate::postgres_triggers::{
create_logical_replication_slot, create_pg_publication, generate_random_string,
get_default_pg_connection, PublicationData,
};
use crate::{
@@ -304,13 +301,15 @@ async fn set_postgres_trigger_config(
user_db: UserDB,
mut capture_config: NewCaptureConfig,
) -> Result<NewCaptureConfig> {
use windmill_common::error::to_anyhow;
let Some(TriggerConfig::Postgres(postgres_config)) = capture_config.trigger_config.as_mut()
else {
return Err(Error::BadRequest("Invalid postgres config".to_string()));
};
if postgres_config.basic_mode.unwrap_or(false) {
let mut pg_connection = get_pg_connection(
let mut pg_connection = get_default_pg_connection(
authed,
Some(user_db),
&db,
@@ -319,22 +318,26 @@ async fn set_postgres_trigger_config(
)
.await?;
let mut tx = pg_connection.begin().await?;
let tx = pg_connection.transaction().await.map_err(to_anyhow)?;
let publication_name = format!("windmill_capture_{}", generate_random_string());
let replication_slot_name = publication_name.clone();
create_logical_replication_slot(&mut tx, &replication_slot_name).await?;
create_logical_replication_slot(tx.client(), &replication_slot_name)
.await
.map_err(to_anyhow)?;
create_pg_publication(
&mut tx,
tx.client(),
&publication_name,
postgres_config.publication.table_to_track.as_deref(),
&postgres_config.publication.transaction_to_track,
)
.await?;
.await
.map_err(to_anyhow)?;
tx.commit().await.map_err(to_anyhow)?;
tx.commit().await?;
postgres_config.publication_name = Some(publication_name);
postgres_config.replication_slot_name = Some(replication_slot_name);
} else {

665
backend/windmill-api/src/postgres_triggers/handler.rs
View File

File diff suppressed because it is too large Load Diff

204
backend/windmill-api/src/postgres_triggers/mod.rs
View File

@@ -7,15 +7,13 @@ use crate::{
};
use chrono::Utc;
use itertools::Itertools;
use pg_escape::{quote_identifier, quote_literal};
use native_tls::{Certificate, TlsConnector};
use pg_escape::quote_identifier;
use rand::Rng;
use rust_postgres::{config::SslMode, Client, Config, NoTls};
use rust_postgres_native_tls::MakeTlsConnector;
use serde_json::value::RawValue;
use sqlx::{
postgres::{PgConnectOptions, PgSslMode},
Connection, PgConnection,
};
use std::collections::HashMap;
use std::str::FromStr;
use axum::{
routing::{delete, get, post},
@@ -33,7 +31,7 @@ use handler::{
};
use windmill_common::{
db::UserDB,
error::{Error, Result},
error::{to_anyhow, Error, Result},
utils::StripPath,
};
mod bool;
@@ -52,76 +50,144 @@ const ERROR_REPLICATION_SLOT_NOT_EXISTS: &str = r#"The replication slot associat
const ERROR_PUBLICATION_NAME_NOT_EXISTS: &str = r#"The publication associated with this trigger no longer exists. Recreate a new publication or select an existing one in the advanced tab, or delete and recreate a new trigger"#;
fn build_tls_connector(
ssl_mode: SslMode,
root_certificate_pem: Option<&String>,
) -> Result<Option<MakeTlsConnector>> {
let get_tls_builder_for_verify = |root_certificate: Option<&String>| {
let mut builder = TlsConnector::builder();
if let Some(root_certificate) = root_certificate {
let root_certificate_pem =
Certificate::from_pem(root_certificate.as_bytes()).map_err(to_anyhow)?;
builder.add_root_certificate(root_certificate_pem);
}
Ok::<_, Error>(builder)
};
let connector = match ssl_mode {
SslMode::Disable => return Ok(None),
SslMode::Require | SslMode::Prefer => {
let mut builder = TlsConnector::builder();
builder.danger_accept_invalid_certs(true);
builder.danger_accept_invalid_hostnames(true);
builder
}
SslMode::VerifyCa => {
let mut builder = get_tls_builder_for_verify(root_certificate_pem)?;
builder.danger_accept_invalid_hostnames(true);
builder
}
SslMode::VerifyFull => {
let builder = get_tls_builder_for_verify(root_certificate_pem)?;
builder
}
_ => unreachable!(),
};
Ok(Some(MakeTlsConnector::new(
connector.build().map_err(to_anyhow)?,
)))
}
pub async fn get_raw_postgres_connection(
database: &Postgres,
logical_mode: bool,
) -> Result<Client> {
let ssl_mode = match database.sslmode.as_ref() {
"disable" => SslMode::Disable,
"" | "prefer" | "allow" => SslMode::Prefer,
"require" => SslMode::Require,
"verify-ca" => SslMode::VerifyCa,
"verify-full" => SslMode::VerifyFull,
ssl_mode => {
return Err(Error::BadRequest(
format!("Invalid ssl mode for postgres: {}, please put a valid ssl_mode among the following available ssl mode: ['disable', 'allow', 'prefer', 'verify-ca', 'verify-full']", ssl_mode),
))
}
};
let mut config = Config::new();
config
.dbname(&database.dbname)
.host(&database.host)
.user(&database.user)
.ssl_mode(ssl_mode);
if logical_mode {
config.replication_mode(rust_postgres::config::ReplicationMode::Logical);
}
if let Some(port) = database.port {
config.port(port);
};
if !database.password.is_empty() {
config.password(&database.password);
}
let connector = build_tls_connector(ssl_mode, database.root_certificate_pem.as_ref())?;
let client = if let Some(connector) = connector {
let (client, connection) = config.connect(connector).await.map_err(to_anyhow)?;
tokio::spawn(async move {
if let Err(e) = connection.await {
tracing::debug!("{:#?}", e);
};
tracing::info!("Successfully Connected into database");
});
client
} else {
let (client, connection) = config.connect(NoTls).await.map_err(to_anyhow)?;
tokio::spawn(async move {
if let Err(e) = connection.await {
tracing::debug!("{:#?}", e);
};
tracing::info!("Successfully Connected into database");
});
client
};
Ok(client)
}
pub async fn get_pg_connection(
authed: ApiAuthed,
user_db: Option<UserDB>,
db: &DB,
postgres_resource_path: &str,
w_id: &str,
) -> Result<PgConnection> {
logical_mode: bool,
) -> Result<Client> {
let database =
try_get_resource_from_db_as::<Postgres>(authed, user_db, db, postgres_resource_path, w_id)
.await?;
Ok(get_raw_postgres_connection(&database).await?)
Ok(get_raw_postgres_connection(&database, logical_mode).await?)
}
pub async fn get_raw_postgres_connection(db: &Postgres) -> Result<PgConnection> {
let options = {
let sslmode = if !db.sslmode.is_empty() {
PgSslMode::from_str(&db.sslmode)?
} else {
PgSslMode::Prefer
};
let options = {
let inner_options = PgConnectOptions::new()
.host(&db.host)
.database(&db.dbname)
.ssl_mode(sslmode)
.username(&db.user);
if let Some(port) = db.port {
inner_options.port(port)
} else {
inner_options
}
};
let options = if let Some(root_certificate_pem) = &db.root_certificate_pem {
options.ssl_root_cert_from_pem(root_certificate_pem.as_bytes().to_vec())
} else {
options
};
if !db.password.is_empty() {
options.password(&db.password)
} else {
options
}
};
Ok(PgConnection::connect_with(&options).await?)
pub async fn get_default_pg_connection(
authed: ApiAuthed,
user_db: Option<UserDB>,
db: &DB,
postgres_resource_path: &str,
w_id: &str,
) -> Result<Client> {
get_pg_connection(authed, user_db, db, postgres_resource_path, w_id, false).await
}
pub async fn create_logical_replication_slot(
pg_connection: &mut PgConnection,
name: &str,
) -> Result<()> {
let query = format!(
r#"
SELECT
*
FROM
pg_create_logical_replication_slot({}, 'pgoutput');"#,
quote_literal(&name)
);
sqlx::query(&query).execute(pg_connection).await?;
pub async fn create_logical_replication_slot(tx: &Client, slot_name: &str) -> Result<()> {
tx.execute(
&format!("SELECT * FROM pg_create_logical_replication_slot($1, 'pgoutput')"),
&[&slot_name],
)
.await
.map_err(to_anyhow)?;
Ok(())
}
async fn check_if_valid_publication_for_postgres_version(
pg_connection: &mut PgConnection,
pg_connection: &Client,
table_to_track: Option<&[Relations]>,
) -> Result<bool> {
let postgres_version = get_postgres_version_internal(pg_connection).await?;
@@ -155,7 +221,7 @@ async fn check_if_valid_publication_for_postgres_version(
}
pub async fn create_pg_publication(
pg_connection: &mut PgConnection,
pg_connection: &Client,
publication_name: &str,
table_to_track: Option<&[Relations]>,
transaction_to_track: &[String],
@@ -177,7 +243,7 @@ pub async fn create_pg_publication(
} else {
if pg_14 && first {
query.push_str(" TABLE ONLY ");
first = false;
first = false
} else if !pg_14 {
query.push_str(" TABLE ONLY ");
}
@@ -224,20 +290,22 @@ pub async fn create_pg_publication(
query.push_str("');");
}
sqlx::query(&query).execute(pg_connection).await?;
pg_connection
.execute(&query, &[])
.await
.map_err(to_anyhow)?;
Ok(())
}
pub async fn drop_publication(
pg_connection: &mut PgConnection,
publication_name: &str,
) -> Result<()> {
pub async fn drop_publication(pg_connection: &Client, publication_name: &str) -> Result<()> {
let mut query = String::from("DROP PUBLICATION IF EXISTS ");
let quoted_publication_name = quote_identifier(publication_name);
query.push_str(&quoted_publication_name);
sqlx::query(&query).execute(pg_connection).await?;
pg_connection
.execute(&query, &[])
.await
.map_err(to_anyhow)?;
Ok(())
}

136
backend/windmill-api/src/postgres_triggers/trigger.rs
View File

@@ -19,25 +19,27 @@ use crate::{
use bytes::{BufMut, Bytes, BytesMut};
use chrono::TimeZone;
use futures::{pin_mut, SinkExt, StreamExt};
use native_tls::{Certificate, TlsConnector};
use pg_escape::{quote_identifier, quote_literal};
use rand::seq::SliceRandom;
use rust_postgres::{config::SslMode, Client, Config, CopyBothDuplex, NoTls, SimpleQueryMessage};
use rust_postgres_native_tls::MakeTlsConnector;
use rust_postgres::{Client, CopyBothDuplex, SimpleQueryMessage};
use serde::Deserialize;
use serde_json::value::RawValue;
use sqlx::types::Json as SqlxJson;
use windmill_common::{
db::UserDB, error, triggers::TriggerKind, utils::report_critical_error, worker::to_raw_value,
db::UserDB,
error::{self, to_anyhow},
triggers::TriggerKind,
utils::report_critical_error,
worker::to_raw_value,
INSTANCE_NAME,
};
use super::{
drop_publication, get_pg_connection,
drop_publication, get_default_pg_connection, get_raw_postgres_connection,
handler::{drop_logical_replication_slot, Postgres, PostgresTrigger},
replication_message::PrimaryKeepAliveBody,
ERROR_PUBLICATION_NAME_NOT_EXISTS, ERROR_REPLICATION_SLOT_NOT_EXISTS,
Error, ERROR_PUBLICATION_NAME_NOT_EXISTS, ERROR_REPLICATION_SLOT_NOT_EXISTS,
};
pub struct LogicalReplicationSettings {
@@ -69,109 +71,11 @@ impl RowExist for Vec<SimpleQueryMessage> {
}
}
#[derive(thiserror::Error, Debug)]
enum Error {
#[error("Error from database: {0}")]
Postgres(#[from] rust_postgres::Error),
#[error("Error : {0}")]
Common(#[from] windmill_common::error::Error),
#[error("Tls Error: {0}")]
Tls(#[from] native_tls::Error),
}
fn build_tls_connector(
ssl_mode: SslMode,
root_certificate_pem: Option<&String>,
) -> Result<Option<MakeTlsConnector>, Error> {
let get_tls_builder_for_verify = |root_certificate: Option<&String>| {
let mut builder = TlsConnector::builder();
if let Some(root_certificate) = root_certificate {
let root_certificate_pem = Certificate::from_pem(root_certificate.as_bytes()).map_err(|e| {
Error::Common(error::Error::BadConfig(format!("Invalid Certs: {e:#}")))
})?;
builder.add_root_certificate(root_certificate_pem);
}
Ok::<_, Error>(builder)
};
let connector = match ssl_mode {
SslMode::Disable => return Ok(None),
SslMode::Require | SslMode::Prefer => {
let mut builder = TlsConnector::builder();
builder.danger_accept_invalid_certs(true);
builder.danger_accept_invalid_hostnames(true);
builder
}
SslMode::VerifyCa => {
let mut builder = get_tls_builder_for_verify(root_certificate_pem)?;
builder.danger_accept_invalid_hostnames(true);
builder
}
SslMode::VerifyFull => {
let builder = get_tls_builder_for_verify(root_certificate_pem)?;
builder
}
_ => unreachable!(),
};
Ok(Some(MakeTlsConnector::new(connector.build()?)))
}
pub struct PostgresSimpleClient(Client);
impl PostgresSimpleClient {
async fn new(database: &Postgres) -> Result<Self, Error> {
let ssl_mode = match database.sslmode.as_ref() {
"disable" => SslMode::Disable,
"" | "prefer" | "allow" => SslMode::Prefer,
"require" => SslMode::Require,
"verify-ca" => SslMode::VerifyCa,
"verify-full" => SslMode::VerifyFull,
ssl_mode => {
return Err(Error::Common(windmill_common::error::Error::BadRequest(
format!("Invalid ssl mode for postgres: {}, please put a valid ssl_mode among the following avalible ssl mode: ['disable', 'allow', 'prefer', 'verify-ca', 'verify-full']", ssl_mode),
)))
}
};
let mut config = Config::new();
config
.dbname(&database.dbname)
.host(&database.host)
.user(&database.user)
.ssl_mode(ssl_mode)
.replication_mode(rust_postgres::config::ReplicationMode::Logical);
if let Some(port) = database.port {
config.port(port);
};
if !database.password.is_empty() {
config.password(&database.password);
}
let connector = build_tls_connector(ssl_mode, database.root_certificate_pem.as_ref())?;
let client = if let Some(connector) = connector {
let (client, connection) = config.connect(connector).await?;
tokio::spawn(async move {
if let Err(e) = connection.await {
tracing::debug!("{:#?}", e);
};
tracing::info!("Successfully Connected into database");
});
client
} else {
let (client, connection) = config.connect(NoTls).await?;
tokio::spawn(async move {
if let Err(e) = connection.await {
tracing::debug!("{:#?}", e);
};
tracing::info!("Successfully Connected into database");
});
client
};
let client = get_raw_postgres_connection(database, true).await?;
Ok(PostgresSimpleClient(client))
}
@@ -202,7 +106,8 @@ impl PostgresSimpleClient {
Ok((
self.0
.copy_both_simple::<bytes::Bytes>(query.as_str())
.await?,
.await
.map_err(to_anyhow)?,
LogicalReplicationSettings::new(false),
))
}
@@ -527,12 +432,13 @@ impl PostgresConfig {
"SELECT pubname FROM pg_publication WHERE pubname = {}",
quote_literal(&publication_name)
))
.await?;
.await
.map_err(to_anyhow)?;
if !publication.row_exist() {
return Err(Error::Common(error::Error::BadConfig(
return Err(Error::BadConfig(
ERROR_PUBLICATION_NAME_NOT_EXISTS.to_string(),
)));
));
}
let replication_slot = client
@@ -540,17 +446,19 @@ impl PostgresConfig {
"SELECT slot_name FROM pg_replication_slots WHERE slot_name = {}",
quote_literal(&replication_slot_name)
))
.await?;
.await
.map_err(to_anyhow)?;
if !replication_slot.row_exist() {
return Err(Error::Common(error::Error::BadConfig(
return Err(Error::BadConfig(
ERROR_REPLICATION_SLOT_NOT_EXISTS.to_string(),
)));
));
}
let (logical_replication_stream, logical_replication_settings) = client
.get_logical_replication_stream(&publication_name, &replication_slot_name)
.await?;
.await
.map_err(to_anyhow)?;
Ok((logical_replication_stream, logical_replication_settings))
}
@@ -585,7 +493,7 @@ impl PostgresConfig {
let user_db = UserDB::new(db.clone());
let mut pg_connection = get_pg_connection(
let mut pg_connection = get_default_pg_connection(
authed.clone(),
Some(user_db.clone()),
&db,