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 EVENT TRIGGER name ON event …`.
73    CreateEventTrigger(protobuf::CreateEventTrigStmt),
74    /// `ALTER EVENT TRIGGER name {ENABLE|DISABLE|ENABLE REPLICA|ENABLE ALWAYS}`.
75    AlterEventTrigger(protobuf::AlterEventTrigStmt),
76}
77
78impl Statement {
79    /// Classify a node into the supported whitelist, or return
80    /// [`ParseError::UnsupportedObjectKind`] for anything else.
81    #[allow(clippy::too_many_lines)] // one arm per supported `NodeEnum` variant; the dispatch table is the function.
82    pub fn classify(node: NodeEnum, location: SourceLocation) -> Result<Self, ParseError> {
83        match node {
84            NodeEnum::CreateSchemaStmt(s) => Ok(Self::CreateSchema(s)),
85            NodeEnum::CreateStmt(s) => Ok(Self::CreateTable(s)),
86            NodeEnum::CreateSeqStmt(s) => Ok(Self::CreateSequence(s)),
87            NodeEnum::IndexStmt(s) => Ok(Self::CreateIndex(*s)),
88            NodeEnum::AlterTableStmt(s) => {
89                if is_attach_partition_stmt(&s) {
90                    Ok(Self::AlterTableAttachPartition(s))
91                } else {
92                    Ok(Self::AlterTable(s))
93                }
94            }
95            NodeEnum::CommentStmt(s) => Ok(Self::Comment(*s)),
96            NodeEnum::ViewStmt(s) => Ok(Self::CreateView(*s)),
97            NodeEnum::CreateEnumStmt(s) => Ok(Self::CreateEnum(s)),
98            NodeEnum::CreateDomainStmt(s) => Ok(Self::CreateDomain(*s)),
99            NodeEnum::CompositeTypeStmt(s) => Ok(Self::CreateCompositeType(s)),
100            NodeEnum::CreateRangeStmt(s) => Ok(Self::CreateRange(s)),
101            NodeEnum::CreateFunctionStmt(s) => {
102                if s.is_procedure {
103                    Ok(Self::CreateProcedure(*s))
104                } else {
105                    Ok(Self::CreateFunction(*s))
106                }
107            }
108            NodeEnum::CreateExtensionStmt(s) => Ok(Self::CreateExtension(s)),
109            NodeEnum::CreateTrigStmt(s) => Ok(Self::CreateTrigger(*s)),
110            NodeEnum::GrantStmt(s) => Ok(Self::Grant(s)),
111            NodeEnum::AlterOwnerStmt(s) => Ok(Self::AlterOwner(*s)),
112            NodeEnum::AlterDefaultPrivilegesStmt(s) => Ok(Self::AlterDefaultPrivileges(s)),
113            NodeEnum::CreatePolicyStmt(s) => Ok(Self::CreatePolicy(*s)),
114            NodeEnum::CreatePublicationStmt(s) => Ok(Self::CreatePublication(s)),
115            NodeEnum::AlterPublicationStmt(s) => Ok(Self::AlterPublication(s)),
116            NodeEnum::CreateSubscriptionStmt(s) => Ok(Self::CreateSubscription(s)),
117            NodeEnum::AlterSubscriptionStmt(s) => Ok(Self::AlterSubscription(s)),
118            NodeEnum::CreateStatsStmt(s) => Ok(Self::CreateStatistics(s)),
119            NodeEnum::AlterStatsStmt(s) => Ok(Self::AlterStatistics(*s)),
120            NodeEnum::CreateEventTrigStmt(s) => Ok(Self::CreateEventTrigger(s)),
121            NodeEnum::AlterEventTrigStmt(s) => Ok(Self::AlterEventTrigger(s)),
122            NodeEnum::DefineStmt(s) => {
123                let kind = ObjectType::try_from(s.kind).unwrap_or(ObjectType::Undefined);
124                if matches!(kind, ObjectType::ObjectCollation) {
125                    Ok(Self::CreateCollation(s))
126                } else {
127                    Err(unsupported(location, "CREATE AGGREGATE/OPERATOR/etc."))
128                }
129            }
130            NodeEnum::RenameStmt(s) => {
131                use pg_query::protobuf::ObjectType;
132                let rename_type =
133                    ObjectType::try_from(s.rename_type).unwrap_or(ObjectType::Undefined);
134                if matches!(rename_type, ObjectType::ObjectPublication) {
135                    return Err(ParseError::Structural {
136                        location,
137                        message: "ALTER PUBLICATION … RENAME TO is not supported in pgevolve \
138                                  (pgevolve never models renames)"
139                            .into(),
140                    });
141                }
142                if matches!(rename_type, ObjectType::ObjectSubscription) {
143                    // Recover the subscription name for the error (subname is in
144                    // RenameStmt.relation for subscriptions).
145                    let sub_name = s
146                        .relation
147                        .as_ref()
148                        .map_or_else(|| String::from("<unknown>"), |r| r.relname.clone());
149                    let id = crate::identifier::Identifier::from_unquoted(&sub_name)
150                        .unwrap_or_else(|_| {
151                            crate::identifier::Identifier::from_unquoted("subscription")
152                                .expect("static identifier")
153                        });
154                    return Err(ParseError::SubscriptionRenameNotSupported(id, location));
155                }
156                if matches!(rename_type, ObjectType::ObjectEventTrigger) {
157                    // RenameStmt for an event trigger encodes the (old) name in
158                    // `s.object` as a bare String node.
159                    let name = s
160                        .object
161                        .as_ref()
162                        .and_then(|o| o.node.as_ref())
163                        .and_then(|n| match n {
164                            NodeEnum::String(st) => Some(st.sval.clone()),
165                            _ => None,
166                        })
167                        .unwrap_or_else(|| String::from("<unknown>"));
168                    let id =
169                        crate::identifier::Identifier::from_unquoted(&name).unwrap_or_else(|_| {
170                            crate::identifier::Identifier::from_unquoted("event_trigger")
171                                .expect("static identifier")
172                        });
173                    return Err(ParseError::EventTriggerRenameNotSupported(id, location));
174                }
175                if matches!(rename_type, ObjectType::ObjectStatisticExt) {
176                    // RenameStmt for statistics encodes the name in `s.subname` (the
177                    // new name) and `s.relation` (the old qualified name).
178                    let schema = s
179                        .relation
180                        .as_ref()
181                        .map_or_else(String::new, |r| r.schemaname.clone());
182                    let name = s
183                        .relation
184                        .as_ref()
185                        .map_or_else(|| String::from("<unknown>"), |r| r.relname.clone());
186                    let schema_id = crate::identifier::Identifier::from_unquoted(&schema)
187                        .unwrap_or_else(|_| {
188                            crate::identifier::Identifier::from_unquoted("unknown")
189                                .expect("static identifier")
190                        });
191                    let name_id = crate::identifier::Identifier::from_unquoted(&name)
192                        .unwrap_or_else(|_| {
193                            crate::identifier::Identifier::from_unquoted("statistic")
194                                .expect("static identifier")
195                        });
196                    let qname = crate::identifier::QualifiedName::new(schema_id, name_id);
197                    return Err(ParseError::StatisticRenameNotSupported(qname, location));
198                }
199                Err(unsupported(location, "RENAME"))
200            }
201            NodeEnum::AlterPolicyStmt(_) => Err(ParseError::Structural {
202                location,
203                message: "ALTER POLICY in source is not supported — policy modifications \
204                          happen via diff; use CREATE POLICY in source"
205                    .into(),
206            }),
207            NodeEnum::DropStmt(s) => {
208                let kind = ObjectType::try_from(s.remove_type).unwrap_or(ObjectType::Undefined);
209                if matches!(kind, ObjectType::ObjectPolicy) {
210                    return Err(ParseError::Structural {
211                        location,
212                        message: "DROP POLICY in source is not supported — drops happen \
213                                  via diff"
214                            .into(),
215                    });
216                }
217                Err(unsupported(location, "DROP"))
218            }
219            NodeEnum::AlterExtensionStmt(_) => Err(ParseError::Structural {
220                location,
221                message: "ALTER EXTENSION is not supported in source files — \
222                          declare the desired state via CREATE EXTENSION"
223                    .into(),
224            }),
225            NodeEnum::CreateTableAsStmt(s) => {
226                // `CREATE TABLE ... AS SELECT` and `CREATE MATERIALIZED VIEW`
227                // both arrive as CreateTableAsStmt. Route by the objtype field.
228                match ObjectType::try_from(s.objtype) {
229                    Ok(ObjectType::ObjectMatview) => Ok(Self::CreateMaterializedView(*s)),
230                    _ => Err(unsupported(location, "CREATE TABLE AS SELECT")),
231                }
232            }
233            other => Err(unsupported(location, friendly_kind(&other))),
234        }
235    }
236}
237
238/// Returns `true` when the `AlterTableStmt` is exactly one `ATTACH PARTITION`
239/// sub-command so the classifier can route it to the dedicated variant.
240fn is_attach_partition_stmt(stmt: &protobuf::AlterTableStmt) -> bool {
241    use pg_query::protobuf::AlterTableType;
242    if stmt.cmds.len() != 1 {
243        return false;
244    }
245    match &stmt.cmds[0].node {
246        Some(NodeEnum::AlterTableCmd(c)) => AlterTableType::try_from(c.subtype)
247            .ok()
248            .is_some_and(|t| matches!(t, AlterTableType::AtAttachPartition)),
249        _ => false,
250    }
251}
252
253const fn unsupported(location: SourceLocation, kind: &'static str) -> ParseError {
254    ParseError::UnsupportedObjectKind { location, kind }
255}
256
257/// Translate non-whitelisted node kinds to a human-readable label for diagnostics.
258#[allow(clippy::match_same_arms)]
259const fn friendly_kind(node: &NodeEnum) -> &'static str {
260    match node {
261        NodeEnum::ViewStmt(_) => "CREATE VIEW", // never reached — routed above
262        NodeEnum::CreateEnumStmt(_) => "CREATE TYPE ... AS ENUM", // never reached — routed above
263        NodeEnum::CreateFunctionStmt(_) => "CREATE FUNCTION/PROCEDURE", // never reached — routed above
264        NodeEnum::CreateTrigStmt(_) => "CREATE TRIGGER", // never reached — routed above
265        NodeEnum::CreateRangeStmt(_) => "CREATE TYPE ... AS RANGE", // never reached — routed above
266        NodeEnum::CreateExtensionStmt(_) => "CREATE EXTENSION", // never reached — routed above
267        NodeEnum::CreatePolicyStmt(_) => "CREATE POLICY", // never reached — routed above
268        NodeEnum::CreateForeignTableStmt(_) => "CREATE FOREIGN TABLE",
269        NodeEnum::CreateFdwStmt(_) => "CREATE FOREIGN DATA WRAPPER",
270        NodeEnum::CreateForeignServerStmt(_) => "CREATE FOREIGN SERVER",
271        NodeEnum::CreateRoleStmt(_) => "CREATE ROLE",
272        NodeEnum::GrantStmt(_) => "GRANT/REVOKE", // never reached — routed above
273        NodeEnum::GrantRoleStmt(_) => "GRANT/REVOKE ROLE",
274        NodeEnum::CreateTableAsStmt(_) => "CREATE TABLE AS",
275        NodeEnum::CreateAmStmt(_) => "CREATE ACCESS METHOD",
276        NodeEnum::CreateEventTrigStmt(_) => "CREATE EVENT TRIGGER",
277        NodeEnum::CreateOpClassStmt(_) => "CREATE OPERATOR CLASS",
278        NodeEnum::CreateOpFamilyStmt(_) => "CREATE OPERATOR FAMILY",
279        NodeEnum::CreatePlangStmt(_) => "CREATE LANGUAGE",
280        NodeEnum::CreateStatsStmt(_) => "CREATE STATISTICS",
281        NodeEnum::CreateUserMappingStmt(_) => "CREATE USER MAPPING",
282        NodeEnum::CreateTableSpaceStmt(_) => "CREATE TABLESPACE",
283        NodeEnum::CreatedbStmt(_) => "CREATE DATABASE",
284        NodeEnum::DefineStmt(_) => "CREATE AGGREGATE/OPERATOR/etc.",
285        NodeEnum::RuleStmt(_) => "CREATE RULE",
286        NodeEnum::DropStmt(_) => "DROP",
287        NodeEnum::TruncateStmt(_) => "TRUNCATE",
288        NodeEnum::SelectStmt(_) => "SELECT",
289        NodeEnum::InsertStmt(_) => "INSERT",
290        NodeEnum::UpdateStmt(_) => "UPDATE",
291        NodeEnum::DeleteStmt(_) => "DELETE",
292        NodeEnum::CopyStmt(_) => "COPY",
293        NodeEnum::TransactionStmt(_) => "BEGIN/COMMIT/ROLLBACK",
294        _ => "this statement kind",
295    }
296}
297
298#[cfg(test)]
299mod tests {
300    use super::*;
301    use std::path::PathBuf;
302
303    fn loc() -> SourceLocation {
304        SourceLocation::new(PathBuf::from("test.sql"), 1, 1)
305    }
306
307    fn first_node(sql: &str) -> NodeEnum {
308        let parsed = pg_query::parse(sql).expect("pg_query parses");
309        parsed
310            .protobuf
311            .stmts
312            .into_iter()
313            .next()
314            .and_then(|raw| raw.stmt)
315            .and_then(|n| n.node)
316            .expect("at least one statement")
317    }
318
319    #[test]
320    fn create_table_classifies() {
321        let node = first_node("CREATE TABLE app.users (id integer);");
322        let stmt = Statement::classify(node, loc()).unwrap();
323        assert!(matches!(stmt, Statement::CreateTable(_)));
324    }
325
326    #[test]
327    fn create_schema_classifies() {
328        let node = first_node("CREATE SCHEMA app;");
329        let stmt = Statement::classify(node, loc()).unwrap();
330        assert!(matches!(stmt, Statement::CreateSchema(_)));
331    }
332
333    #[test]
334    fn create_index_classifies() {
335        let node = first_node("CREATE INDEX users_email_idx ON app.users (email);");
336        let stmt = Statement::classify(node, loc()).unwrap();
337        assert!(matches!(stmt, Statement::CreateIndex(_)));
338    }
339
340    #[test]
341    fn create_sequence_classifies() {
342        let node = first_node("CREATE SEQUENCE app.seq1;");
343        let stmt = Statement::classify(node, loc()).unwrap();
344        assert!(matches!(stmt, Statement::CreateSequence(_)));
345    }
346
347    #[test]
348    fn alter_table_classifies() {
349        let node = first_node(
350            "ALTER TABLE app.invoices ADD CONSTRAINT fk1 FOREIGN KEY (cid) REFERENCES app.c(id);",
351        );
352        let stmt = Statement::classify(node, loc()).unwrap();
353        assert!(matches!(stmt, Statement::AlterTable(_)));
354    }
355
356    #[test]
357    fn comment_classifies() {
358        let node = first_node("COMMENT ON TABLE app.users IS 'people';");
359        let stmt = Statement::classify(node, loc()).unwrap();
360        assert!(matches!(stmt, Statement::Comment(_)));
361    }
362
363    #[test]
364    fn create_view_classifies() {
365        let node = first_node("CREATE VIEW app.v AS SELECT 1;");
366        let stmt = Statement::classify(node, loc()).unwrap();
367        assert!(matches!(stmt, Statement::CreateView(_)));
368    }
369
370    #[test]
371    fn create_materialized_view_classifies() {
372        let node = first_node(
373            "CREATE MATERIALIZED VIEW app.mv AS SELECT count(*) FROM app.t WITH NO DATA;",
374        );
375        let stmt = Statement::classify(node, loc()).unwrap();
376        assert!(matches!(stmt, Statement::CreateMaterializedView(_)));
377    }
378
379    #[test]
380    fn create_table_as_select_unsupported() {
381        let node = first_node("CREATE TABLE app.t2 AS SELECT 1;");
382        let err = Statement::classify(node, loc()).unwrap_err();
383        assert!(matches!(err, ParseError::UnsupportedObjectKind { .. }));
384    }
385
386    #[test]
387    fn create_function_classifies() {
388        let node =
389            first_node("CREATE FUNCTION app.f() RETURNS integer AS $$ SELECT 1; $$ LANGUAGE sql;");
390        let stmt = Statement::classify(node, loc()).unwrap();
391        assert!(matches!(stmt, Statement::CreateFunction(_)));
392    }
393
394    #[test]
395    fn create_procedure_classifies() {
396        let node =
397            first_node("CREATE PROCEDURE app.p() LANGUAGE plpgsql AS $$ BEGIN NULL; END $$;");
398        let stmt = Statement::classify(node, loc()).unwrap();
399        assert!(matches!(stmt, Statement::CreateProcedure(_)));
400    }
401
402    #[test]
403    fn create_extension_classifies() {
404        let node = first_node("CREATE EXTENSION pgcrypto;");
405        let stmt = Statement::classify(node, loc()).unwrap();
406        assert!(matches!(stmt, Statement::CreateExtension(_)));
407    }
408
409    #[test]
410    fn alter_extension_rejected() {
411        let node = first_node("ALTER EXTENSION pgcrypto UPDATE;");
412        let err = Statement::classify(node, loc()).unwrap_err();
413        assert!(matches!(err, ParseError::Structural { .. }));
414    }
415
416    #[test]
417    fn select_unsupported() {
418        let node = first_node("SELECT 1;");
419        let err = Statement::classify(node, loc()).unwrap_err();
420        assert!(matches!(err, ParseError::UnsupportedObjectKind { .. }));
421    }
422}
pgevolve_core/parse/statement.rs

pgevolve_core/parse/
statement.rs
