Skip to main content

pgsync_postgres/
schema.rs

1//! Schema bootstrap via `pg_dump --schema-only`, and target column verification.
2
3use crate::client::PgClient;
4use crate::errors::from_pg;
5use crate::url::{ConnectionInfo, TlsPolicy};
6use pgsync_core::config::{SafetyConfig, TableName, quote_ident};
7use pgsync_core::{Error, ErrorKind, Result};
8use std::collections::BTreeSet;
9use tracing::info;
10
11/// Run `pg_dump --schema-only` for `tables` and return SQL safe for `batch_execute`.
12pub async fn dump_schema(
13    pg_dump: Option<&str>,
14    info: &ConnectionInfo,
15    safety: &SafetyConfig,
16    tables: &[TableName],
17) -> Result<String> {
18    let bin = pg_dump.unwrap_or("pg_dump");
19    let mut cmd = tokio::process::Command::new(bin);
20    cmd.args([
21        "--schema-only",
22        "--no-owner",
23        "--no-privileges",
24        "--no-comments",
25        "--no-publications",
26        "--no-subscriptions",
27        "--no-security-labels",
28        "--no-tablespaces",
29        "--format=plain",
30    ]);
31    for t in tables {
32        cmd.arg("--table").arg(format!(
33            "{}.{}",
34            quote_ident(&t.schema),
35            quote_ident(&t.name)
36        ));
37    }
38    // Credentials go through the environment, never the argument list.
39    let url = format!(
40        "postgresql://{}@{}:{}/{}",
41        info.user, info.host, info.port, info.dbname
42    );
43    cmd.arg("--dbname").arg(&url);
44    cmd.env_remove("PGPASSWORD");
45    if let Some(pw) = info.password() {
46        cmd.env("PGPASSWORD", String::from_utf8_lossy(pw).into_owned());
47    }
48    let sslmode = match info.tls {
49        TlsPolicy::Disabled => "disable",
50        TlsPolicy::Preferred { .. } => "prefer",
51        TlsPolicy::Required { verify: true } => "verify-full",
52        TlsPolicy::Required { verify: false } => "require",
53    };
54    cmd.env("PGSSLMODE", sslmode);
55    // libpq defaults to ~/.postgresql/root.crt, not the system store; `system` (libpq 16+)
56    // matches what pgsync itself trusts when no CA file is configured.
57    let root = safety
58        .tls_ca_file
59        .as_deref()
60        .or(info.tls_ca_file.as_deref())
61        .unwrap_or("system");
62    // libpq refuses sslrootcert=system with sslmode=prefer, so only set it when verifying.
63    if matches!(info.tls, TlsPolicy::Required { verify: true }) {
64        cmd.env("PGSSLROOTCERT", root);
65    }
66    cmd.stdin(std::process::Stdio::null());
67    let out = cmd.output().await.map_err(|e| {
68        Error::with_source(
69            ErrorKind::Config,
70            format!("cannot run {bin} (set schema.pg_dump_path or schema.mode: existing)"),
71            e,
72        )
73    })?;
74    if !out.status.success() {
75        let stderr = String::from_utf8_lossy(&out.stderr);
76        return Err(Error::preflight(format!(
77            "{bin} failed ({}): {}",
78            out.status,
79            stderr.trim()
80        )));
81    }
82    let sql = String::from_utf8(out.stdout)
83        .map_err(|_| Error::internal("pg_dump output is not UTF-8"))?;
84    info!(tables = tables.len(), bytes = sql.len(), "schema dumped");
85    Ok(filter_dump(&sql))
86}
87
88/// Remove psql meta-commands and session settings that are not portable across server
89/// versions. Statements that define objects are kept verbatim.
90pub fn filter_dump(sql: &str) -> String {
91    let drop_prefixes = [
92        "SET statement_timeout",
93        "SET lock_timeout",
94        "SET idle_in_transaction_session_timeout",
95        "SET transaction_timeout",
96        "SET default_tablespace",
97        "SET default_table_access_method",
98        "SET row_security",
99        "SET check_function_bodies",
100        "SET xmloption",
101        "SET client_min_messages",
102    ];
103    sql.lines()
104        .filter(|l| !l.starts_with('\\'))
105        .filter(|l| !drop_prefixes.iter().any(|p| l.starts_with(p)))
106        .map(|l| format!("{l}\n"))
107        .collect()
108}
109
110/// Create the schemas and apply the dump in one transaction.
111pub async fn apply_schema(target: &PgClient, tables: &[TableName], sql: &str) -> Result<()> {
112    let schemas: BTreeSet<&str> = tables.iter().map(|t| t.schema.as_str()).collect();
113    let mut script = String::from("BEGIN;\n");
114    for s in schemas {
115        script.push_str(&format!(
116            "CREATE SCHEMA IF NOT EXISTS {};\n",
117            quote_ident(s)
118        ));
119    }
120    script.push_str(sql);
121    script.push_str("\nCOMMIT;\n");
122    target
123        .batch_execute(&script)
124        .await
125        .map_err(|e| from_pg("apply schema to target", e))?;
126    Ok(())
127}
128
129/// Target column metadata used by the sink to cast values.
130#[derive(Debug, Clone)]
131pub struct TargetColumn {
132    pub name: String,
133    /// Declared type with modifiers, e.g. `character varying(80)`; for comparison with the
134    /// source and for messages.
135    pub type_name: String,
136    /// Type without modifiers, e.g. `character varying`, used in `$n::type` casts. An
137    /// explicit cast to a modified type (`::varchar(80)`, `::numeric(12,2)`) silently
138    /// truncates or rounds; casting to the base type lets the column's own assignment
139    /// rules apply, which raise an error for values that do not fit, as COPY does.
140    pub cast_type: String,
141    pub generated: bool,
142    /// `GENERATED ALWAYS AS IDENTITY`: cannot be assigned by UPDATE, never changes on the source.
143    pub identity_always: bool,
144}
145
146pub async fn target_columns(
147    target: &PgClient,
148    schema: &str,
149    name: &str,
150) -> Result<Option<Vec<TargetColumn>>> {
151    let rows = target
152        .query(
153            "SELECT a.attname::text, format_type(a.atttypid, a.atttypmod), a.attgenerated <> '', a.attidentity = 'a', \
154             format_type(a.atttypid, NULL) \
155             FROM pg_attribute a JOIN pg_class c ON c.oid = a.attrelid JOIN pg_namespace n ON n.oid = c.relnamespace \
156             WHERE n.nspname = $1 AND c.relname = $2 AND c.relkind = 'r' AND a.attnum > 0 AND NOT a.attisdropped \
157             ORDER BY a.attnum",
158            &[&schema, &name],
159        )
160        .await
161        .map_err(|e| from_pg(&format!("target columns of {schema}.{name}"), e))?;
162    if rows.is_empty() {
163        return Ok(None);
164    }
165    Ok(Some(
166        rows.iter()
167            .map(|r| TargetColumn {
168                name: r.get(0),
169                type_name: r.get(1),
170                generated: r.get(2),
171                identity_always: r.get(3),
172                cast_type: r.get(4),
173            })
174            .collect(),
175    ))
176}
177
178/// Verify that every source column exists on the target with a compatible declaration.
179pub fn verify_columns(
180    table: &TableName,
181    source: &[crate::source::ColumnInfo],
182    target: &[TargetColumn],
183) -> Result<Vec<String>> {
184    let mut warnings = Vec::new();
185    for s in source.iter().filter(|c| !c.generated) {
186        match target.iter().find(|t| t.name == s.name) {
187            None => {
188                return Err(Error::preflight(format!(
189                    "{table}: column {:?} is missing on the target",
190                    s.name
191                )));
192            }
193            Some(t) if t.generated => {
194                return Err(Error::preflight(format!(
195                    "{table}: column {:?} is generated on the target but not on the source",
196                    s.name
197                )));
198            }
199            Some(t) if t.type_name != s.type_name => {
200                warnings.push(format!(
201                    "{table}.{}: source type {} vs target type {}",
202                    s.name, s.type_name, t.type_name
203                ));
204            }
205            Some(_) => {}
206        }
207    }
208    Ok(warnings)
209}
210
211#[cfg(test)]
212mod tests {
213    use super::*;
214
215    #[test]
216    fn filter_strips_meta_commands_and_timeouts() {
217        let dump = "\\restrict abc\nSET statement_timeout = 0;\nSET transaction_timeout = 0;\nSET client_encoding = 'UTF8';\nSELECT pg_catalog.set_config('search_path', '', false);\nCREATE TABLE public.t (\n    id bigint NOT NULL\n);\n\\unrestrict abc\n";
218        let out = filter_dump(dump);
219        assert_eq!(
220            out,
221            "SET client_encoding = 'UTF8';\nSELECT pg_catalog.set_config('search_path', '', false);\nCREATE TABLE public.t (\n    id bigint NOT NULL\n);\n"
222        );
223    }
224
225    #[test]
226    fn verify_columns_detects_missing_and_type_drift() {
227        let t = TableName::new("public", "t");
228        let src = vec![
229            crate::source::ColumnInfo {
230                attnum: 1,
231                name: "id".into(),
232                type_name: "bigint".into(),
233                not_null: true,
234                generated: false,
235                identity: "a".into(),
236            },
237            crate::source::ColumnInfo {
238                attnum: 2,
239                name: "v".into(),
240                type_name: "text".into(),
241                not_null: false,
242                generated: false,
243                identity: "".into(),
244            },
245            crate::source::ColumnInfo {
246                attnum: 3,
247                name: "g".into(),
248                type_name: "text".into(),
249                not_null: false,
250                generated: true,
251                identity: "".into(),
252            },
253        ];
254        let tgt = vec![
255            TargetColumn {
256                name: "id".into(),
257                type_name: "bigint".into(),
258                cast_type: "bigint".into(),
259                generated: false,
260                identity_always: true,
261            },
262            TargetColumn {
263                name: "v".into(),
264                type_name: "character varying(10)".into(),
265                cast_type: "character varying".into(),
266                generated: false,
267                identity_always: false,
268            },
269        ];
270        let w = verify_columns(&t, &src, &tgt).unwrap();
271        assert_eq!(w.len(), 1);
272        assert!(w[0].contains("text vs target type character varying"));
273        let err = verify_columns(&t, &src, &tgt[..1]).unwrap_err();
274        assert!(err.message().contains("missing on the target"));
275    }
276}