use std::collections::HashMap; use windmill_api_auth::ApiAuthed; use windmill_store::resources::try_get_resource_from_db_as; use windmill_trigger::trigger_helpers::TriggerJobArgs; use chrono::Utc; use itertools::Itertools; 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::{Deserialize, Deserializer, Serialize}; use serde_json::value::RawValue; use sqlx::FromRow; use windmill_common::{ db::UserDB, error::{to_anyhow, Error, Result}, triggers::TriggerKind, utils::empty_as_none, DB, }; mod bool; mod converter; pub mod handler; mod hex; pub mod listener; mod mapper; mod relation; mod replication_message; #[derive(Clone, Copy)] pub struct PostgresTrigger; impl TriggerJobArgs for PostgresTrigger { type Payload = HashMap>; const TRIGGER_KIND: TriggerKind = TriggerKind::Postgres; fn v1_payload_fn(payload: &HashMap>) -> HashMap> { payload.to_owned() } } #[derive(Debug, Clone, FromRow, Serialize, Deserialize)] pub struct PostgresConfig { pub postgres_resource_path: String, pub replication_slot_name: String, pub publication_name: String, #[serde(default, skip_serializing_if = "Option::is_none")] pub basic_mode: Option, } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct PostgresConfigRequest { postgres_resource_path: String, #[serde(default)] replication_slot_name: String, #[serde(default)] publication_name: String, publication: Option, } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct TestPostgresConfig { pub postgres_resource_path: String, } fn check_if_valid_relation<'de, D>( relations: D, ) -> std::result::Result>, D::Error> where D: Deserializer<'de>, { let relations: Option> = Option::deserialize(relations)?; let mut track_all_table_in_schema = false; let mut track_specific_columns_in_table = false; match relations { Some(relations) => { for relation in relations.iter() { if relation.schema_name.is_empty() { return Err(serde::de::Error::custom( "Schema Name must not be empty".to_string(), )); } if !track_all_table_in_schema && relation.table_to_track.is_empty() { track_all_table_in_schema = true; continue; } for table_to_track in relation.table_to_track.iter() { if table_to_track.table_name.trim().is_empty() { return Err(serde::de::Error::custom( "Table name must not be empty".to_string(), )); } if !track_specific_columns_in_table && table_to_track.columns_name.is_some() { track_specific_columns_in_table = true; } } if track_all_table_in_schema && track_specific_columns_in_table { return Err(serde::de::Error::custom("Incompatible tracking options. Schema-level tracking and specific table tracking with column selection cannot be used together. Refer to the documentation for valid configurations.")); } } if !relations .iter() .map(|relation| relation.schema_name.as_str()) .all_unique() { return Err(serde::de::Error::custom( "You cannot choose a schema more than one time".to_string(), )); } Ok(Some(relations)) } None => Ok(None), } } fn check_if_valid_transaction_type<'de, D>( transaction_type: D, ) -> std::result::Result, D::Error> where D: Deserializer<'de>, { let mut transaction_type: Vec = Vec::deserialize(transaction_type)?; if transaction_type.len() > 3 { return Err(serde::de::Error::custom( "More than 3 transaction type which is not authorized, you are only allowed to those 3 transaction types: Insert, Update and Delete" .to_string(), )); } transaction_type.sort_unstable(); transaction_type.dedup(); for transaction in transaction_type.iter() { match transaction.to_lowercase().as_ref() { "insert" => {}, "update" => {}, "delete" => {}, _ => { return Err(serde::de::Error::custom( "Only the following transaction types are allowed: Insert, Update and Delete (case insensitive)" .to_string(), )) } } } Ok(transaction_type) } #[derive(Debug, Clone, Deserialize, Serialize)] pub struct PublicationData { #[serde(default, deserialize_with = "check_if_valid_relation")] pub table_to_track: Option>, #[serde(deserialize_with = "check_if_valid_transaction_type")] pub transaction_to_track: Vec, } impl PublicationData { pub fn new( table_to_track: Option>, transaction_to_track: Vec, ) -> PublicationData { PublicationData { table_to_track, transaction_to_track } } } // Slot list struct #[derive(FromRow, Debug, Serialize)] pub struct SlotList { pub slot_name: Option, pub active: Option, } // Slot struct #[derive(Debug, Serialize, Deserialize)] pub struct Slot { pub name: String, } // Template script struct #[derive(Debug, Deserialize)] pub struct TemplateScript { pub postgres_resource_path: String, #[serde(deserialize_with = "check_if_valid_relation")] pub relations: Option>, pub language: Language, } // Language enum #[derive(Deserialize, Debug)] pub enum Language { #[serde(rename = "typescript", alias = "Typescript")] Typescript, } // Test postgres struct #[derive(Serialize, Deserialize)] pub struct TestPostgres { pub postgres_resource_path: String, } // PostgreSQL publication replication struct #[derive(Serialize, Deserialize)] pub struct PostgresPublicationReplication { pub publication_name: String, pub replication_slot_name: String, } impl PostgresPublicationReplication { pub fn new( publication_name: String, replication_slot_name: String, ) -> PostgresPublicationReplication { PostgresPublicationReplication { publication_name, replication_slot_name } } } pub 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"#; #[derive(FromRow, Serialize, Deserialize, Debug)] pub struct Postgres { pub user: String, pub password: String, pub host: String, pub port: Option, pub dbname: String, #[serde(default)] pub sslmode: String, #[serde(default, deserialize_with = "empty_as_none")] pub root_certificate_pem: Option, } #[derive(Debug, Clone, FromRow, Serialize, Deserialize)] pub struct TableToTrack { pub table_name: String, #[serde(default, deserialize_with = "empty_as_none")] pub where_clause: Option, #[serde(default, deserialize_with = "empty_as_none")] pub columns_name: Option>, } impl TableToTrack { pub fn new( table_name: String, where_clause: Option, columns_name: Option>, ) -> TableToTrack { TableToTrack { table_name, where_clause, columns_name } } } #[derive(Debug, Clone, FromRow, Serialize, Deserialize)] pub struct Relations { pub schema_name: String, pub table_to_track: Vec, } impl Relations { pub fn new(schema_name: String, table_to_track: Vec) -> Relations { Relations { schema_name, table_to_track } } pub fn add_new_table(&mut self, table_to_track: TableToTrack) { self.table_to_track.push(table_to_track); } } fn build_tls_connector( ssl_mode: SslMode, root_certificate_pem: Option<&String>, ) -> Result> { 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 { 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 { tracing::info!("Successfully connected to PostgreSQL database for trigger execution"); if let Err(e) = connection.await { tracing::debug!("Error during PostgreSQL trigger connection: {:#?}", e); }; tracing::info!("PostgreSQL trigger connection closed"); }); 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, db: &DB, postgres_resource_path: &str, w_id: &str, logical_mode: bool, ) -> Result { let database = try_get_resource_from_db_as::(&authed, user_db, db, postgres_resource_path, w_id) .await?; Ok(get_raw_postgres_connection(&database, logical_mode).await?) } pub async fn get_default_pg_connection( authed: ApiAuthed, user_db: Option, db: &DB, postgres_resource_path: &str, w_id: &str, ) -> Result { get_pg_connection(authed, user_db, db, postgres_resource_path, w_id, false).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(()) } pub async fn check_if_valid_publication_for_postgres_version( pg_connection: &Client, table_to_track: Option<&[Relations]>, ) -> Result { use crate::handler::get_postgres_version_internal; let postgres_version = get_postgres_version_internal(pg_connection).await?; let pg_14 = postgres_version.starts_with("14"); if pg_14 { let unsupported_publication = table_to_track .and_then(|relations| { relations.iter().find(|relation| { let invalid_relation = relation.table_to_track.iter().find(|table_to_track| { table_to_track.where_clause.is_some() || table_to_track.columns_name.is_some() }); relation.table_to_track.is_empty() || invalid_relation.is_some() }) }) .is_some(); if unsupported_publication { return Err(Error::BadRequest( "Your PostgreSQL database is running version 14, which does not support the following publication features: \ - WHERE clause filtering, \ - selective column tracking, and \ - tracking all tables within a schema.\n\ These features are only available in PostgreSQL 15 and above.".to_string(), )); } } Ok(pg_14) } pub async fn create_pg_publication( pg_connection: &Client, publication_name: &str, table_to_track: Option<&[Relations]>, transaction_to_track: &[String], ) -> Result<()> { let pg_14 = check_if_valid_publication_for_postgres_version(pg_connection, table_to_track).await?; let mut query = String::from("CREATE PUBLICATION "); query.push_str("e_identifier(publication_name)); match table_to_track { Some(database_component) if !database_component.is_empty() => { query.push_str(" FOR"); let mut first = true; for (i, schema) in database_component.iter().enumerate() { if schema.table_to_track.is_empty() { query.push_str(" TABLES IN SCHEMA "); query.push_str("e_identifier(&schema.schema_name)); } else { if pg_14 && first { query.push_str(" TABLE ONLY "); first = false } else if !pg_14 { query.push_str(" TABLE ONLY "); } for (j, table) in schema.table_to_track.iter().enumerate() { let table_name = quote_identifier(&table.table_name); let schema_name = quote_identifier(&schema.schema_name); let full_name = format!("{}.{}", &schema_name, &table_name); query.push_str(&full_name); if let Some(columns) = table.columns_name.as_ref() { query.push_str(" ("); let columns = columns .iter() .map(|column| quote_identifier(column)) .join(", "); query.push_str(&columns); query.push_str(")"); } if let Some(where_clause) = &table.where_clause { query.push_str(" WHERE ("); query.push_str(where_clause); query.push(')'); } if j + 1 != schema.table_to_track.len() { query.push_str(", "); } } } if i < database_component.len() - 1 { query.push_str(", "); } } } _ => { query.push_str(" FOR ALL TABLES "); } }; if !transaction_to_track.is_empty() { let transactions = || transaction_to_track.iter().join(", "); query.push_str(" WITH (publish = '"); query.push_str(&transactions()); query.push_str("');"); } pg_connection .execute(&query, &[]) .await .map_err(to_anyhow)?; Ok(()) } 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("ed_publication_name); pg_connection .execute(&query, &[]) .await .map_err(to_anyhow)?; Ok(()) } pub fn generate_random_string() -> String { let timestamp = Utc::now().timestamp_millis().to_string(); let mut rng = rand::rng(); let charset = "abcdefghijklmnopqrstuvwxyz0123456789"; let random_part = (0..10) .map(|_| { charset .chars() .nth(rng.random_range(0..charset.len())) .unwrap() }) .collect::(); format!("{}_{}", timestamp, random_part) } #[cfg(test)] mod tests { use super::*; #[test] fn test_valid_publication_data() { let json = r#"{ "table_to_track": [ { "schema_name": "public", "table_to_track": [ {"table_name": "users"} ] } ], "transaction_to_track": ["insert", "update"] }"#; let result: std::result::Result = serde_json::from_str(json); assert!(result.is_ok()); } #[test] fn test_publication_data_empty_schema_name() { let json = r#"{ "table_to_track": [ { "schema_name": "", "table_to_track": [] } ], "transaction_to_track": ["insert"] }"#; let result: std::result::Result = serde_json::from_str(json); assert!(result.is_err()); } #[test] fn test_publication_data_empty_table_name() { let json = r#"{ "table_to_track": [ { "schema_name": "public", "table_to_track": [ {"table_name": " "} ] } ], "transaction_to_track": ["insert"] }"#; let result: std::result::Result = serde_json::from_str(json); assert!(result.is_err()); } #[test] fn test_publication_data_invalid_transaction_type() { let json = r#"{ "transaction_to_track": ["insert", "truncate"] }"#; let result: std::result::Result = serde_json::from_str(json); assert!(result.is_err()); } #[test] fn test_publication_data_too_many_transaction_types() { let json = r#"{ "transaction_to_track": ["insert", "update", "delete", "insert"] }"#; let result: std::result::Result = serde_json::from_str(json); assert!(result.is_err()); } #[test] fn test_publication_data_duplicate_schema_names() { let json = r#"{ "table_to_track": [ {"schema_name": "public", "table_to_track": [{"table_name": "a"}]}, {"schema_name": "public", "table_to_track": [{"table_name": "b"}]} ], "transaction_to_track": ["insert"] }"#; let result: std::result::Result = serde_json::from_str(json); assert!(result.is_err()); } #[test] fn test_publication_data_all_tables_in_schema() { let json = r#"{ "table_to_track": [ {"schema_name": "public", "table_to_track": []} ], "transaction_to_track": ["insert"] }"#; let result: std::result::Result = serde_json::from_str(json); assert!(result.is_ok()); } #[test] fn test_publication_data_incompatible_tracking() { let json = r#"{ "table_to_track": [ {"schema_name": "schema1", "table_to_track": []}, {"schema_name": "schema2", "table_to_track": [{"table_name": "t1", "columns_name": ["col1"]}]} ], "transaction_to_track": ["insert"] }"#; let result: std::result::Result = serde_json::from_str(json); assert!(result.is_err()); } #[test] fn test_postgres_config_serialization() { let config = PostgresConfig { postgres_resource_path: "f/db/postgres".to_string(), replication_slot_name: "slot_1".to_string(), publication_name: "pub_1".to_string(), basic_mode: Some(false), }; let json = serde_json::to_value(&config).unwrap(); assert_eq!(json["postgres_resource_path"], "f/db/postgres"); assert_eq!(json["replication_slot_name"], "slot_1"); } #[test] fn test_generate_random_string_format() { let s = generate_random_string(); assert!(s.contains('_')); let parts: Vec<&str> = s.split('_').collect(); assert_eq!(parts.len(), 2); assert_eq!(parts[1].len(), 10); } #[test] fn test_relations_add_table() { let mut rel = Relations::new("public".to_string(), vec![]); rel.add_new_table(TableToTrack::new("users".to_string(), None, None)); assert_eq!(rel.table_to_track.len(), 1); assert_eq!(rel.table_to_track[0].table_name, "users"); } #[test] fn test_table_to_track_with_where_clause() { let tt = TableToTrack::new( "orders".to_string(), Some("status = 'active'".to_string()), None, ); assert_eq!(tt.where_clause, Some("status = 'active'".to_string())); } #[test] fn test_table_to_track_with_columns() { let tt = TableToTrack::new( "users".to_string(), None, Some(vec!["id".to_string(), "email".to_string()]), ); assert_eq!(tt.columns_name.as_ref().unwrap().len(), 2); } }