Skip to main content

pgevolve_core/parse/
statement.rs

1//! Classify a top-level `pg_query` statement node into the v0.1 whitelist.
2//!
3//! Anything outside the whitelist is rejected with [`ParseError::UnsupportedObjectKind`]
4//! so that source-loading fails loudly instead of silently dropping unsupported DDL.
5
6use pg_query::NodeEnum;
7use pg_query::protobuf;
8use pg_query::protobuf::ObjectType;
9
10use crate::parse::error::{ParseError, SourceLocation};
11
12/// One supported top-level statement kind. Each variant carries the typed protobuf
13/// message that the corresponding builder will consume.
14#[derive(Debug, Clone)]
15pub enum Statement {
16    /// `CREATE SCHEMA ...`.
17    CreateSchema(protobuf::CreateSchemaStmt),
18    /// `CREATE TABLE ...`.
19    CreateTable(protobuf::CreateStmt),
20    /// `CREATE SEQUENCE ...`.
21    CreateSequence(protobuf::CreateSeqStmt),
22    /// `CREATE [UNIQUE] INDEX ...`.
23    CreateIndex(protobuf::IndexStmt),
24    /// `ALTER TABLE ...` — only `ADD CONSTRAINT FOREIGN KEY` is allowed in source DDL.
25    AlterTable(protobuf::AlterTableStmt),
26    /// `COMMENT ON ...`.
27    Comment(protobuf::CommentStmt),
28    /// `CREATE [OR REPLACE] VIEW ...`.
29    CreateView(protobuf::ViewStmt),
30    /// `CREATE MATERIALIZED VIEW ...`.
31    CreateMaterializedView(protobuf::CreateTableAsStmt),
32    /// `CREATE TYPE ... AS ENUM (...)`.
33    CreateEnum(protobuf::CreateEnumStmt),
34    /// `CREATE DOMAIN ... AS ...`.
35    CreateDomain(protobuf::CreateDomainStmt),
36    /// `CREATE TYPE ... AS (...)` (composite).
37    CreateCompositeType(protobuf::CompositeTypeStmt),
38    /// `CREATE TYPE ... AS RANGE (...)`.
39    CreateRange(protobuf::CreateRangeStmt),
40    /// `CREATE [OR REPLACE] FUNCTION ...`.
41    CreateFunction(protobuf::CreateFunctionStmt),
42    /// `CREATE [OR REPLACE] PROCEDURE ...`.
43    CreateProcedure(protobuf::CreateFunctionStmt),
44    /// `CREATE EXTENSION ...`.
45    CreateExtension(protobuf::CreateExtensionStmt),
46    /// `CREATE [CONSTRAINT] TRIGGER ...`.
47    CreateTrigger(protobuf::CreateTrigStmt),
48    /// `ALTER TABLE parent ATTACH PARTITION child FOR VALUES ...`.
49    AlterTableAttachPartition(protobuf::AlterTableStmt),
50    /// `GRANT priv ON obj TO grantee` — object-level grant (not REVOKE, not role grant).
51    Grant(protobuf::GrantStmt),
52    /// `ALTER <kind> name OWNER TO role`.
53    AlterOwner(protobuf::AlterOwnerStmt),
54    /// `ALTER DEFAULT PRIVILEGES ... GRANT ...`.
55    AlterDefaultPrivileges(protobuf::AlterDefaultPrivilegesStmt),
56    /// `CREATE POLICY name ON table ...`.
57    CreatePolicy(protobuf::CreatePolicyStmt),
58    /// `CREATE PUBLICATION name ...`.
59    CreatePublication(protobuf::CreatePublicationStmt),
60    /// `ALTER PUBLICATION name ...`.
61    AlterPublication(protobuf::AlterPublicationStmt),
62    /// `CREATE SUBSCRIPTION name CONNECTION ... PUBLICATION ...`.
63    CreateSubscription(protobuf::CreateSubscriptionStmt),
64    /// `ALTER SUBSCRIPTION name ...`.
65    AlterSubscription(protobuf::AlterSubscriptionStmt),
66    /// `CREATE STATISTICS name … ON … FROM …`.
67    CreateStatistics(protobuf::CreateStatsStmt),
68    /// `ALTER STATISTICS name SET STATISTICS n`.
69    AlterStatistics(protobuf::AlterStatsStmt),
70    /// `CREATE COLLATION qname (provider = …, locale = …, …)`.
71    CreateCollation(protobuf::DefineStmt),
72    /// `CREATE AGGREGATE qname(args) (SFUNC = …, STYPE = …, …)`.
73    CreateAggregate(protobuf::DefineStmt),
74    /// `CREATE EVENT TRIGGER name ON event …`.
75    CreateEventTrigger(protobuf::CreateEventTrigStmt),
76    /// `ALTER EVENT TRIGGER name {ENABLE|DISABLE|ENABLE REPLICA|ENABLE ALWAYS}`.
77    AlterEventTrigger(protobuf::AlterEventTrigStmt),
78    /// `CREATE CAST (src AS tgt) WITH FUNCTION …` / `WITH INOUT` / `WITHOUT FUNCTION`.
79    CreateCast(protobuf::CreateCastStmt),
80    /// `CREATE TEXT SEARCH DICTIONARY qname (TEMPLATE = …, …)`.
81    CreateTsDictionary(protobuf::DefineStmt),
82    /// `CREATE TEXT SEARCH CONFIGURATION qname (PARSER = …)`.
83    CreateTsConfiguration(protobuf::DefineStmt),
84    /// `ALTER TEXT SEARCH DICTIONARY name (opts)`.
85    AlterTsDictionary(protobuf::AlterTsDictionaryStmt),
86    /// `ALTER TEXT SEARCH CONFIGURATION name ADD/ALTER/DROP/REPLACE MAPPING …`.
87    AlterTsConfiguration(protobuf::AlterTsConfigurationStmt),
88}
89
90impl Statement {
91    /// Classify a node into the supported whitelist, or return
92    /// [`ParseError::UnsupportedObjectKind`] for anything else.
93    #[allow(clippy::too_many_lines)] // one arm per supported `NodeEnum` variant; the dispatch table is the function.
94    pub fn classify(node: NodeEnum, location: SourceLocation) -> Result<Self, ParseError> {
95        match node {
96            NodeEnum::CreateSchemaStmt(s) => Ok(Self::CreateSchema(s)),
97            NodeEnum::CreateStmt(s) => Ok(Self::CreateTable(s)),
98            NodeEnum::CreateSeqStmt(s) => Ok(Self::CreateSequence(s)),
99            NodeEnum::IndexStmt(s) => Ok(Self::CreateIndex(*s)),
100            NodeEnum::AlterTableStmt(s) => {
101                if is_attach_partition_stmt(&s) {
102                    Ok(Self::AlterTableAttachPartition(s))
103                } else {
104                    Ok(Self::AlterTable(s))
105                }
106            }
107            NodeEnum::CommentStmt(s) => Ok(Self::Comment(*s)),
108            NodeEnum::ViewStmt(s) => Ok(Self::CreateView(*s)),
109            NodeEnum::CreateEnumStmt(s) => Ok(Self::CreateEnum(s)),
110            NodeEnum::CreateDomainStmt(s) => Ok(Self::CreateDomain(*s)),
111            NodeEnum::CompositeTypeStmt(s) => Ok(Self::CreateCompositeType(s)),
112            NodeEnum::CreateRangeStmt(s) => Ok(Self::CreateRange(s)),
113            NodeEnum::CreateFunctionStmt(s) => {
114                if s.is_procedure {
115                    Ok(Self::CreateProcedure(*s))
116                } else {
117                    Ok(Self::CreateFunction(*s))
118                }
119            }
120            NodeEnum::CreateExtensionStmt(s) => Ok(Self::CreateExtension(s)),
121            NodeEnum::CreateTrigStmt(s) => Ok(Self::CreateTrigger(*s)),
122            NodeEnum::GrantStmt(s) => Ok(Self::Grant(s)),
123            NodeEnum::AlterOwnerStmt(s) => Ok(Self::AlterOwner(*s)),
124            NodeEnum::AlterDefaultPrivilegesStmt(s) => Ok(Self::AlterDefaultPrivileges(s)),
125            NodeEnum::CreatePolicyStmt(s) => Ok(Self::CreatePolicy(*s)),
126            NodeEnum::CreatePublicationStmt(s) => Ok(Self::CreatePublication(s)),
127            NodeEnum::AlterPublicationStmt(s) => Ok(Self::AlterPublication(s)),
128            NodeEnum::CreateSubscriptionStmt(s) => Ok(Self::CreateSubscription(s)),
129            NodeEnum::AlterSubscriptionStmt(s) => Ok(Self::AlterSubscription(s)),
130            NodeEnum::CreateStatsStmt(s) => Ok(Self::CreateStatistics(s)),
131            NodeEnum::AlterStatsStmt(s) => Ok(Self::AlterStatistics(*s)),
132            NodeEnum::CreateEventTrigStmt(s) => Ok(Self::CreateEventTrigger(s)),
133            NodeEnum::AlterEventTrigStmt(s) => Ok(Self::AlterEventTrigger(s)),
134            NodeEnum::CreateCastStmt(s) => Ok(Self::CreateCast(s)),
135            NodeEnum::DefineStmt(s) => {
136                let kind = ObjectType::try_from(s.kind).unwrap_or(ObjectType::Undefined);
137                match kind {
138                    ObjectType::ObjectCollation => Ok(Self::CreateCollation(s)),
139                    ObjectType::ObjectAggregate => Ok(Self::CreateAggregate(s)),
140                    ObjectType::ObjectTsdictionary => Ok(Self::CreateTsDictionary(s)),
141                    ObjectType::ObjectTsconfiguration => Ok(Self::CreateTsConfiguration(s)),
142                    _ => Err(unsupported(location, "CREATE OPERATOR/etc.")),
143                }
144            }
145            NodeEnum::RenameStmt(s) => {
146                use pg_query::protobuf::ObjectType;
147                let rename_type =
148                    ObjectType::try_from(s.rename_type).unwrap_or(ObjectType::Undefined);
149                if matches!(rename_type, ObjectType::ObjectPublication) {
150                    return Err(ParseError::Structural {
151                        location,
152                        message: "ALTER PUBLICATION … RENAME TO is not supported in pgevolve \
153                                  (pgevolve never models renames)"
154                            .into(),
155                    });
156                }
157                if matches!(rename_type, ObjectType::ObjectSubscription) {
158                    // Recover the subscription name for the error (subname is in
159                    // RenameStmt.relation for subscriptions).
160                    let sub_name = s
161                        .relation
162                        .as_ref()
163                        .map_or_else(|| String::from("<unknown>"), |r| r.relname.clone());
164                    let id = crate::identifier::Identifier::from_unquoted(&sub_name)
165                        .unwrap_or_else(|_| {
166                            crate::identifier::Identifier::from_unquoted("subscription")
167                                .expect("static identifier")
168                        });
169                    return Err(ParseError::SubscriptionRenameNotSupported(id, location));
170                }
171                if matches!(rename_type, ObjectType::ObjectEventTrigger) {
172                    // RenameStmt for an event trigger encodes the (old) name in
173                    // `s.object` as a bare String node.
174                    let name = s
175                        .object
176                        .as_ref()
177                        .and_then(|o| o.node.as_ref())
178                        .and_then(|n| match n {
179                            NodeEnum::String(st) => Some(st.sval.clone()),
180                            _ => None,
181                        })
182                        .unwrap_or_else(|| String::from("<unknown>"));
183                    let id =
184                        crate::identifier::Identifier::from_unquoted(&name).unwrap_or_else(|_| {
185                            crate::identifier::Identifier::from_unquoted("event_trigger")
186                                .expect("static identifier")
187                        });
188                    return Err(ParseError::EventTriggerRenameNotSupported(id, location));
189                }
190                if matches!(rename_type, ObjectType::ObjectAggregate) {
191                    return Err(ParseError::Structural {
192                        location,
193                        message: "ALTER AGGREGATE … RENAME TO is not supported in pgevolve \
194                                  (pgevolve never models renames)"
195                            .into(),
196                    });
197                }
198                if matches!(rename_type, ObjectType::ObjectStatisticExt) {
199                    // RenameStmt for statistics encodes the name in `s.subname` (the
200                    // new name) and `s.relation` (the old qualified name).
201                    let schema = s
202                        .relation
203                        .as_ref()
204                        .map_or_else(String::new, |r| r.schemaname.clone());
205                    let name = s
206                        .relation
207                        .as_ref()
208                        .map_or_else(|| String::from("<unknown>"), |r| r.relname.clone());
209                    let schema_id = crate::identifier::Identifier::from_unquoted(&schema)
210                        .unwrap_or_else(|_| {
211                            crate::identifier::Identifier::from_unquoted("unknown")
212                                .expect("static identifier")
213                        });
214                    let name_id = crate::identifier::Identifier::from_unquoted(&name)
215                        .unwrap_or_else(|_| {
216                            crate::identifier::Identifier::from_unquoted("statistic")
217                                .expect("static identifier")
218                        });
219                    let qname = crate::identifier::QualifiedName::new(schema_id, name_id);
220                    return Err(ParseError::StatisticRenameNotSupported(qname, location));
221                }
222                Err(unsupported(location, "RENAME"))
223            }
224            NodeEnum::AlterPolicyStmt(_) => Err(ParseError::Structural {
225                location,
226                message: "ALTER POLICY in source is not supported — policy modifications \
227                          happen via diff; use CREATE POLICY in source"
228                    .into(),
229            }),
230            NodeEnum::AlterTsdictionaryStmt(s) => Ok(Self::AlterTsDictionary(s)),
231            NodeEnum::AlterTsconfigurationStmt(s) => Ok(Self::AlterTsConfiguration(s)),
232            NodeEnum::DropStmt(s) => {
233                let kind = ObjectType::try_from(s.remove_type).unwrap_or(ObjectType::Undefined);
234                if matches!(kind, ObjectType::ObjectPolicy) {
235                    return Err(ParseError::Structural {
236                        location,
237                        message: "DROP POLICY in source is not supported — drops happen \
238                                  via diff"
239                            .into(),
240                    });
241                }
242                if matches!(kind, ObjectType::ObjectAggregate) {
243                    return Err(ParseError::Structural {
244                        location,
245                        message: "DROP AGGREGATE in source is not supported — drops happen \
246                                  via diff"
247                            .into(),
248                    });
249                }
250                if matches!(kind, ObjectType::ObjectCast) {
251                    return Err(ParseError::Structural {
252                        location,
253                        message: "DROP CAST in source is not supported — drops happen via diff"
254                            .into(),
255                    });
256                }
257                if matches!(kind, ObjectType::ObjectTsdictionary) {
258                    return Err(ParseError::Structural {
259                        location,
260                        message: "DROP TEXT SEARCH DICTIONARY in source is not supported — \
261                                  drops happen via diff"
262                            .into(),
263                    });
264                }
265                if matches!(kind, ObjectType::ObjectTsconfiguration) {
266                    return Err(ParseError::Structural {
267                        location,
268                        message: "DROP TEXT SEARCH CONFIGURATION in source is not supported — \
269                                  drops happen via diff"
270                            .into(),
271                    });
272                }
273                Err(unsupported(location, "DROP"))
274            }
275            NodeEnum::AlterExtensionStmt(_) => Err(ParseError::Structural {
276                location,
277                message: "ALTER EXTENSION is not supported in source files — \
278                          declare the desired state via CREATE EXTENSION"
279                    .into(),
280            }),
281            NodeEnum::CreateTableAsStmt(s) => {
282                // `CREATE TABLE ... AS SELECT` and `CREATE MATERIALIZED VIEW`
283                // both arrive as CreateTableAsStmt. Route by the objtype field.
284                match ObjectType::try_from(s.objtype) {
285                    Ok(ObjectType::ObjectMatview) => Ok(Self::CreateMaterializedView(*s)),
286                    _ => Err(unsupported(location, "CREATE TABLE AS SELECT")),
287                }
288            }
289            other => Err(unsupported(location, friendly_kind(&other))),
290        }
291    }
292}
293
294/// Returns `true` when the `AlterTableStmt` is exactly one `ATTACH PARTITION`
295/// sub-command so the classifier can route it to the dedicated variant.
296fn is_attach_partition_stmt(stmt: &protobuf::AlterTableStmt) -> bool {
297    use pg_query::protobuf::AlterTableType;
298    if stmt.cmds.len() != 1 {
299        return false;
300    }
301    match &stmt.cmds[0].node {
302        Some(NodeEnum::AlterTableCmd(c)) => AlterTableType::try_from(c.subtype)
303            .ok()
304            .is_some_and(|t| matches!(t, AlterTableType::AtAttachPartition)),
305        _ => false,
306    }
307}
308
309const fn unsupported(location: SourceLocation, kind: &'static str) -> ParseError {
310    ParseError::UnsupportedObjectKind { location, kind }
311}
312
313/// Translate non-whitelisted node kinds to a human-readable label for diagnostics.
314#[allow(clippy::match_same_arms)]
315const fn friendly_kind(node: &NodeEnum) -> &'static str {
316    match node {
317        NodeEnum::ViewStmt(_) => "CREATE VIEW", // never reached — routed above
318        NodeEnum::CreateEnumStmt(_) => "CREATE TYPE ... AS ENUM", // never reached — routed above
319        NodeEnum::CreateFunctionStmt(_) => "CREATE FUNCTION/PROCEDURE", // never reached — routed above
320        NodeEnum::CreateTrigStmt(_) => "CREATE TRIGGER", // never reached — routed above
321        NodeEnum::CreateRangeStmt(_) => "CREATE TYPE ... AS RANGE", // never reached — routed above
322        NodeEnum::CreateExtensionStmt(_) => "CREATE EXTENSION", // never reached — routed above
323        NodeEnum::CreatePolicyStmt(_) => "CREATE POLICY", // never reached — routed above
324        NodeEnum::CreateForeignTableStmt(_) => "CREATE FOREIGN TABLE",
325        NodeEnum::CreateFdwStmt(_) => "CREATE FOREIGN DATA WRAPPER",
326        NodeEnum::CreateForeignServerStmt(_) => "CREATE FOREIGN SERVER",
327        NodeEnum::CreateRoleStmt(_) => "CREATE ROLE",
328        NodeEnum::GrantStmt(_) => "GRANT/REVOKE", // never reached — routed above
329        NodeEnum::GrantRoleStmt(_) => "GRANT/REVOKE ROLE",
330        NodeEnum::CreateTableAsStmt(_) => "CREATE TABLE AS",
331        NodeEnum::CreateAmStmt(_) => "CREATE ACCESS METHOD",
332        NodeEnum::CreateEventTrigStmt(_) => "CREATE EVENT TRIGGER",
333        NodeEnum::CreateOpClassStmt(_) => "CREATE OPERATOR CLASS",
334        NodeEnum::CreateOpFamilyStmt(_) => "CREATE OPERATOR FAMILY",
335        NodeEnum::CreatePlangStmt(_) => "CREATE LANGUAGE",
336        NodeEnum::CreateStatsStmt(_) => "CREATE STATISTICS",
337        NodeEnum::CreateUserMappingStmt(_) => "CREATE USER MAPPING",
338        NodeEnum::CreateTableSpaceStmt(_) => "CREATE TABLESPACE",
339        NodeEnum::CreatedbStmt(_) => "CREATE DATABASE",
340        NodeEnum::DefineStmt(_) => "CREATE AGGREGATE/OPERATOR/etc.",
341        NodeEnum::RuleStmt(_) => "CREATE RULE",
342        NodeEnum::DropStmt(_) => "DROP",
343        NodeEnum::TruncateStmt(_) => "TRUNCATE",
344        NodeEnum::SelectStmt(_) => "SELECT",
345        NodeEnum::InsertStmt(_) => "INSERT",
346        NodeEnum::UpdateStmt(_) => "UPDATE",
347        NodeEnum::DeleteStmt(_) => "DELETE",
348        NodeEnum::CopyStmt(_) => "COPY",
349        NodeEnum::TransactionStmt(_) => "BEGIN/COMMIT/ROLLBACK",
350        _ => "this statement kind",
351    }
352}
353
354#[cfg(test)]
355mod tests {
356    use super::*;
357    use std::path::PathBuf;
358
359    fn loc() -> SourceLocation {
360        SourceLocation::new(PathBuf::from("test.sql"), 1, 1)
361    }
362
363    fn first_node(sql: &str) -> NodeEnum {
364        let parsed = pg_query::parse(sql).expect("pg_query parses");
365        parsed
366            .protobuf
367            .stmts
368            .into_iter()
369            .next()
370            .and_then(|raw| raw.stmt)
371            .and_then(|n| n.node)
372            .expect("at least one statement")
373    }
374
375    #[test]
376    fn create_table_classifies() {
377        let node = first_node("CREATE TABLE app.users (id integer);");
378        let stmt = Statement::classify(node, loc()).unwrap();
379        assert!(matches!(stmt, Statement::CreateTable(_)));
380    }
381
382    #[test]
383    fn create_schema_classifies() {
384        let node = first_node("CREATE SCHEMA app;");
385        let stmt = Statement::classify(node, loc()).unwrap();
386        assert!(matches!(stmt, Statement::CreateSchema(_)));
387    }
388
389    #[test]
390    fn create_index_classifies() {
391        let node = first_node("CREATE INDEX users_email_idx ON app.users (email);");
392        let stmt = Statement::classify(node, loc()).unwrap();
393        assert!(matches!(stmt, Statement::CreateIndex(_)));
394    }
395
396    #[test]
397    fn create_sequence_classifies() {
398        let node = first_node("CREATE SEQUENCE app.seq1;");
399        let stmt = Statement::classify(node, loc()).unwrap();
400        assert!(matches!(stmt, Statement::CreateSequence(_)));
401    }
402
403    #[test]
404    fn alter_table_classifies() {
405        let node = first_node(
406            "ALTER TABLE app.invoices ADD CONSTRAINT fk1 FOREIGN KEY (cid) REFERENCES app.c(id);",
407        );
408        let stmt = Statement::classify(node, loc()).unwrap();
409        assert!(matches!(stmt, Statement::AlterTable(_)));
410    }
411
412    #[test]
413    fn comment_classifies() {
414        let node = first_node("COMMENT ON TABLE app.users IS 'people';");
415        let stmt = Statement::classify(node, loc()).unwrap();
416        assert!(matches!(stmt, Statement::Comment(_)));
417    }
418
419    #[test]
420    fn create_view_classifies() {
421        let node = first_node("CREATE VIEW app.v AS SELECT 1;");
422        let stmt = Statement::classify(node, loc()).unwrap();
423        assert!(matches!(stmt, Statement::CreateView(_)));
424    }
425
426    #[test]
427    fn create_materialized_view_classifies() {
428        let node = first_node(
429            "CREATE MATERIALIZED VIEW app.mv AS SELECT count(*) FROM app.t WITH NO DATA;",
430        );
431        let stmt = Statement::classify(node, loc()).unwrap();
432        assert!(matches!(stmt, Statement::CreateMaterializedView(_)));
433    }
434
435    #[test]
436    fn create_table_as_select_unsupported() {
437        let node = first_node("CREATE TABLE app.t2 AS SELECT 1;");
438        let err = Statement::classify(node, loc()).unwrap_err();
439        assert!(matches!(err, ParseError::UnsupportedObjectKind { .. }));
440    }
441
442    #[test]
443    fn create_function_classifies() {
444        let node =
445            first_node("CREATE FUNCTION app.f() RETURNS integer AS $$ SELECT 1; $$ LANGUAGE sql;");
446        let stmt = Statement::classify(node, loc()).unwrap();
447        assert!(matches!(stmt, Statement::CreateFunction(_)));
448    }
449
450    #[test]
451    fn create_procedure_classifies() {
452        let node =
453            first_node("CREATE PROCEDURE app.p() LANGUAGE plpgsql AS $$ BEGIN NULL; END $$;");
454        let stmt = Statement::classify(node, loc()).unwrap();
455        assert!(matches!(stmt, Statement::CreateProcedure(_)));
456    }
457
458    #[test]
459    fn create_extension_classifies() {
460        let node = first_node("CREATE EXTENSION pgcrypto;");
461        let stmt = Statement::classify(node, loc()).unwrap();
462        assert!(matches!(stmt, Statement::CreateExtension(_)));
463    }
464
465    #[test]
466    fn alter_extension_rejected() {
467        let node = first_node("ALTER EXTENSION pgcrypto UPDATE;");
468        let err = Statement::classify(node, loc()).unwrap_err();
469        assert!(matches!(err, ParseError::Structural { .. }));
470    }
471
472    #[test]
473    fn select_unsupported() {
474        let node = first_node("SELECT 1;");
475        let err = Statement::classify(node, loc()).unwrap_err();
476        assert!(matches!(err, ParseError::UnsupportedObjectKind { .. }));
477    }
478}