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

pgevolve_core/parse/
statement.rs
