mz_sql/ast/
transform.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Provides a publicly available interface to transform our SQL ASTs.
11
12use std::collections::{BTreeMap, BTreeSet};
13
14use mz_ore::str::StrExt;
15use mz_repr::CatalogItemId;
16use mz_sql_parser::ast::CreateTableFromSourceStatement;
17
18use crate::ast::visit::{self, Visit};
19use crate::ast::visit_mut::{self, VisitMut};
20use crate::ast::{
21    AstInfo, CreateConnectionStatement, CreateIndexStatement, CreateMaterializedViewStatement,
22    CreateSecretStatement, CreateSinkStatement, CreateSourceStatement, CreateSubsourceStatement,
23    CreateTableStatement, CreateViewStatement, CreateWebhookSourceStatement, Expr, Ident, Query,
24    Raw, RawDataType, RawItemName, Statement, UnresolvedItemName, ViewDefinition,
25};
26use crate::names::FullItemName;
27
28/// Given a [`Statement`] rewrites all references of the schema name `cur_schema_name` to
29/// `new_schema_name`.
30pub fn create_stmt_rename_schema_refs(
31    create_stmt: &mut Statement<Raw>,
32    database: &str,
33    cur_schema: &str,
34    new_schema: &str,
35) -> Result<(), (String, String)> {
36    match create_stmt {
37        stmt @ Statement::CreateConnection(_)
38        | stmt @ Statement::CreateDatabase(_)
39        | stmt @ Statement::CreateSchema(_)
40        | stmt @ Statement::CreateWebhookSource(_)
41        | stmt @ Statement::CreateSource(_)
42        | stmt @ Statement::CreateSubsource(_)
43        | stmt @ Statement::CreateSink(_)
44        | stmt @ Statement::CreateView(_)
45        | stmt @ Statement::CreateMaterializedView(_)
46        | stmt @ Statement::CreateTable(_)
47        | stmt @ Statement::CreateTableFromSource(_)
48        | stmt @ Statement::CreateIndex(_)
49        | stmt @ Statement::CreateType(_)
50        | stmt @ Statement::CreateSecret(_) => {
51            let mut visitor = CreateSqlRewriteSchema {
52                database,
53                cur_schema,
54                new_schema,
55                error: None,
56            };
57            visitor.visit_statement_mut(stmt);
58
59            if let Some(e) = visitor.error.take() {
60                Err(e)
61            } else {
62                Ok(())
63            }
64        }
65        stmt => {
66            unreachable!("Internal error: only catalog items need to update item refs. {stmt:?}")
67        }
68    }
69}
70
71struct CreateSqlRewriteSchema<'a> {
72    database: &'a str,
73    cur_schema: &'a str,
74    new_schema: &'a str,
75    error: Option<(String, String)>,
76}
77
78impl<'a> CreateSqlRewriteSchema<'a> {
79    fn maybe_rewrite_idents(&mut self, name: &mut [Ident]) {
80        match name {
81            [schema, item] if schema.as_str() == self.cur_schema => {
82                // TODO(parkmycar): I _think_ when the database component is not specified we can
83                // always infer we're using the current database. But I'm not positive, so for now
84                // we'll bail in this case.
85                if self.error.is_none() {
86                    self.error = Some((schema.to_string(), item.to_string()));
87                }
88            }
89            [database, schema, _item] => {
90                if database.as_str() == self.database && schema.as_str() == self.cur_schema {
91                    *schema = Ident::new_unchecked(self.new_schema);
92                }
93            }
94            _ => (),
95        }
96    }
97}
98
99impl<'a, 'ast> VisitMut<'ast, Raw> for CreateSqlRewriteSchema<'a> {
100    fn visit_expr_mut(&mut self, e: &'ast mut Expr<Raw>) {
101        match e {
102            Expr::Identifier(id) => {
103                // The last ID component is a column name that should not be
104                // considered in the rewrite.
105                let i = id.len() - 1;
106                self.maybe_rewrite_idents(&mut id[..i]);
107            }
108            Expr::QualifiedWildcard(id) => {
109                self.maybe_rewrite_idents(id);
110            }
111            _ => visit_mut::visit_expr_mut(self, e),
112        }
113    }
114
115    fn visit_unresolved_item_name_mut(
116        &mut self,
117        unresolved_item_name: &'ast mut UnresolvedItemName,
118    ) {
119        self.maybe_rewrite_idents(&mut unresolved_item_name.0);
120    }
121
122    fn visit_data_type_mut(&mut self, data_type: &'ast mut RawDataType) {
123        // The generated `visit_data_type_mut` is a no-op because `DataType` is an
124        // associated type the generic visitor treats opaquely, so we must descend
125        // by hand. A type is referenced by a schema-qualified name precisely in
126        // these data type positions (a cast, a column type, or a nested element
127        // type), so without this a schema rename leaves stale references in any
128        // dependent that uses one of the schema's types.
129        match data_type {
130            RawDataType::Array(element_type) | RawDataType::List(element_type) => {
131                self.visit_data_type_mut(element_type)
132            }
133            RawDataType::Map {
134                key_type,
135                value_type,
136            } => {
137                self.visit_data_type_mut(key_type);
138                self.visit_data_type_mut(value_type);
139            }
140            RawDataType::Other { name, .. } => self.visit_item_name_mut(name),
141        }
142    }
143
144    fn visit_item_name_mut(
145        &mut self,
146        item_name: &'ast mut <mz_sql_parser::ast::Raw as AstInfo>::ItemName,
147    ) {
148        match item_name {
149            RawItemName::Name(n) | RawItemName::Id(_, n, _) => self.maybe_rewrite_idents(&mut n.0),
150        }
151    }
152}
153
154/// Changes the `name` used in an item's `CREATE` statement. To complete a
155/// rename operation, you must also call `create_stmt_rename_refs` on all dependent
156/// items.
157pub fn create_stmt_rename(create_stmt: &mut Statement<Raw>, to_item_name: String) {
158    // TODO(sploiselle): Support renaming schemas and databases.
159    match create_stmt {
160        Statement::CreateIndex(CreateIndexStatement { name, .. }) => {
161            *name = Some(Ident::new_unchecked(to_item_name));
162        }
163        Statement::CreateSink(CreateSinkStatement {
164            name: Some(name), ..
165        })
166        | Statement::CreateSource(CreateSourceStatement { name, .. })
167        | Statement::CreateSubsource(CreateSubsourceStatement { name, .. })
168        | Statement::CreateView(CreateViewStatement {
169            definition: ViewDefinition { name, .. },
170            ..
171        })
172        | Statement::CreateMaterializedView(CreateMaterializedViewStatement { name, .. })
173        | Statement::CreateTable(CreateTableStatement { name, .. })
174        | Statement::CreateTableFromSource(CreateTableFromSourceStatement { name, .. })
175        | Statement::CreateSecret(CreateSecretStatement { name, .. })
176        | Statement::CreateConnection(CreateConnectionStatement { name, .. })
177        | Statement::CreateWebhookSource(CreateWebhookSourceStatement { name, .. }) => {
178            // The last name in an ItemName is the item name. The item name
179            // does not have a fixed index.
180            // TODO: https://github.com/MaterializeInc/database-issues/issues/1721
181            let item_name_len = name.0.len() - 1;
182            name.0[item_name_len] = Ident::new_unchecked(to_item_name);
183        }
184        item => unreachable!("Internal error: only catalog items can be renamed {item:?}"),
185    }
186}
187
188/// Updates all references of `from_name` in `create_stmt` to `to_name` or
189/// errors if request is ambiguous.
190///
191/// Requests are considered ambiguous if `create_stmt` is a
192/// `Statement::CreateView`, and any of the following apply to its `query`:
193/// - `to_name.item` is used as an [`Ident`] in `query`.
194/// - `from_name.item` does not unambiguously refer to an item in the query,
195///   e.g. it is also used as a schema, or not all references to the item are
196///   sufficiently qualified.
197/// - `to_name.item` does not unambiguously refer to an item in the query after
198///   the rename. Right now, given the first condition, this is just a coherence
199///   check, but will be more meaningful once the first restriction is lifted.
200pub fn create_stmt_rename_refs(
201    create_stmt: &mut Statement<Raw>,
202    from_name: FullItemName,
203    to_item_name: String,
204) -> Result<(), String> {
205    let from_item = UnresolvedItemName::from(from_name.clone());
206    let maybe_update_item_name = |item_name: &mut UnresolvedItemName| {
207        if item_name.0 == from_item.0 {
208            // The last name in an ItemName is the item name. The item name
209            // does not have a fixed index.
210            // TODO: https://github.com/MaterializeInc/database-issues/issues/1721
211            let item_name_len = item_name.0.len() - 1;
212            item_name.0[item_name_len] = Ident::new_unchecked(to_item_name.clone());
213        }
214    };
215
216    // TODO(sploiselle): Support renaming schemas and databases.
217    match create_stmt {
218        Statement::CreateIndex(CreateIndexStatement { on_name, .. }) => {
219            maybe_update_item_name(on_name.name_mut());
220        }
221        Statement::CreateSink(CreateSinkStatement { from, .. }) => {
222            maybe_update_item_name(from.name_mut());
223        }
224        Statement::CreateTableFromSource(CreateTableFromSourceStatement { source, .. }) => {
225            maybe_update_item_name(source.name_mut());
226        }
227        Statement::CreateView(CreateViewStatement {
228            definition: ViewDefinition { query, .. },
229            ..
230        }) => {
231            rewrite_query(from_name, to_item_name, query)?;
232        }
233        Statement::CreateMaterializedView(CreateMaterializedViewStatement {
234            replacement_for,
235            query,
236            ..
237        }) => {
238            if let Some(target) = replacement_for {
239                maybe_update_item_name(target.name_mut());
240            }
241            rewrite_query(from_name, to_item_name, query)?;
242        }
243        Statement::CreateSource(_)
244        | Statement::CreateSubsource(_)
245        | Statement::CreateTable(_)
246        | Statement::CreateSecret(_)
247        | Statement::CreateConnection(_)
248        | Statement::CreateWebhookSource(_) => {}
249        item => {
250            unreachable!("Internal error: only catalog items need to update item refs {item:?}")
251        }
252    }
253
254    Ok(())
255}
256
257/// Rewrites `query`'s references of `from` to `to` or errors if too ambiguous.
258fn rewrite_query(from: FullItemName, to: String, query: &mut Query<Raw>) -> Result<(), String> {
259    let from_ident = Ident::new_unchecked(from.item.clone());
260    let to_ident = Ident::new_unchecked(to);
261    let qual_depth =
262        QueryIdentAgg::determine_qual_depth(&from_ident, Some(to_ident.clone()), query)?;
263    CreateSqlRewriter::rewrite_query_with_qual_depth(from, to_ident.clone(), qual_depth, query);
264    // Ensure that our rewrite didn't didn't introduce ambiguous
265    // references to `to_name`.
266    match QueryIdentAgg::determine_qual_depth(&to_ident, None, query) {
267        Ok(_) => Ok(()),
268        Err(e) => Err(e),
269    }
270}
271
272fn ambiguous_err(n: &Ident, t: &str) -> String {
273    format!(
274        "{} potentially used ambiguously as item and {}",
275        n.as_str().quoted(),
276        t
277    )
278}
279
280/// Visits a [`Query`], assessing catalog item [`Ident`]s' use of a specified `Ident`.
281struct QueryIdentAgg<'a> {
282    /// The name whose usage you want to assess.
283    name: &'a Ident,
284    /// Tracks all second-level qualifiers used on `name` in a `BTreeMap`, as
285    /// well as any third-level qualifiers used on those second-level qualifiers
286    /// in a `BTreeSet`.
287    qualifiers: BTreeMap<Ident, BTreeSet<Ident>>,
288    /// Tracks the least qualified instance of `name` seen.
289    min_qual_depth: usize,
290    /// Provides an option to fail the visit if encounters a specified `Ident`.
291    fail_on: Option<Ident>,
292    err: Option<String>,
293}
294
295impl<'a> QueryIdentAgg<'a> {
296    /// Determines the depth of qualification needed to unambiguously reference
297    /// catalog items in a [`Query`].
298    ///
299    /// Includes an option to fail if a given `Ident` is encountered.
300    ///
301    /// `Result`s of `Ok(usize)` indicate that `name` can be unambiguously
302    /// referred to with `usize` parts, e.g. 2 requires schema and item name
303    /// qualification.
304    ///
305    /// `Result`s of `Err` indicate that we cannot unambiguously reference
306    /// `name` or encountered `fail_on`, if it's provided.
307    fn determine_qual_depth(
308        name: &Ident,
309        fail_on: Option<Ident>,
310        query: &Query<Raw>,
311    ) -> Result<usize, String> {
312        let mut v = QueryIdentAgg {
313            qualifiers: BTreeMap::new(),
314            min_qual_depth: usize::MAX,
315            err: None,
316            name,
317            fail_on,
318        };
319
320        // Aggregate identities in `v`.
321        v.visit_query(query);
322        // Not possible to have a qualification depth of 0;
323        assert!(v.min_qual_depth > 0);
324
325        if let Some(e) = v.err {
326            return Err(e);
327        }
328
329        // Check if there was more than one 3rd-level (e.g.
330        // database) qualification used for any reference to `name`.
331        let req_depth = if v.qualifiers.values().any(|v| v.len() > 1) {
332            3
333        // Check if there was more than one 2nd-level (e.g. schema)
334        // qualification used for any reference to `name`.
335        } else if v.qualifiers.len() > 1 {
336            2
337        } else {
338            1
339        };
340
341        if v.min_qual_depth < req_depth {
342            Err(format!(
343                "{} is not sufficiently qualified to support renaming",
344                name.as_str().quoted()
345            ))
346        } else {
347            Ok(req_depth)
348        }
349    }
350
351    // Assesses `v` for uses of `self.name` and `self.fail_on`.
352    fn check_failure(&mut self, v: &[Ident]) {
353        // Fail if we encounter `self.fail_on`.
354        if let Some(f) = &self.fail_on {
355            if v.iter().any(|i| i == f) {
356                self.err = Some(format!(
357                    "found reference to {}; cannot rename {} to any identity \
358                    used in any existing view definitions",
359                    f.as_str().quoted(),
360                    self.name.as_str().quoted()
361                ));
362            }
363        }
364    }
365}
366
367impl<'a, 'ast> Visit<'ast, Raw> for QueryIdentAgg<'a> {
368    fn visit_expr(&mut self, e: &'ast Expr<Raw>) {
369        match e {
370            Expr::Identifier(i) => {
371                self.check_failure(i);
372                if let Some(p) = i.iter().rposition(|e| e == self.name) {
373                    if p == i.len() - 1 {
374                        // `self.name` used as a column if it's in the final
375                        // position here, e.g. `SELECT view.col FROM ...`
376                        self.err = Some(ambiguous_err(self.name, "column"));
377                        return;
378                    }
379                    self.min_qual_depth = std::cmp::min(p + 1, self.min_qual_depth);
380                }
381            }
382            Expr::QualifiedWildcard(i) => {
383                self.check_failure(i);
384                if let Some(p) = i.iter().rposition(|e| e == self.name) {
385                    self.min_qual_depth = std::cmp::min(p + 1, self.min_qual_depth);
386                }
387            }
388            _ => visit::visit_expr(self, e),
389        }
390    }
391
392    fn visit_ident(&mut self, ident: &'ast Ident) {
393        self.check_failure(std::slice::from_ref(ident));
394        // This is an unqualified item using `self.name`, e.g. an alias, which
395        // we cannot unambiguously resolve.
396        if ident == self.name {
397            self.err = Some(ambiguous_err(self.name, "alias or column"));
398        }
399    }
400
401    fn visit_unresolved_item_name(&mut self, unresolved_item_name: &'ast UnresolvedItemName) {
402        let names = &unresolved_item_name.0;
403        self.check_failure(names);
404        // Every item is used as an `ItemName` at least once, which
405        // lets use track all items named `self.name`.
406        if let Some(p) = names.iter().rposition(|e| e == self.name) {
407            // Name used as last element of `<db>.<schema>.<item>`
408            if p == names.len() - 1 && names.len() == 3 {
409                self.qualifiers
410                    .entry(names[1].clone())
411                    .or_default()
412                    .insert(names[0].clone());
413                self.min_qual_depth = std::cmp::min(3, self.min_qual_depth);
414            } else {
415                // Any other use is a database or schema
416                self.err = Some(ambiguous_err(self.name, "database, schema, or function"))
417            }
418        }
419    }
420
421    fn visit_item_name(&mut self, item_name: &'ast <Raw as AstInfo>::ItemName) {
422        match item_name {
423            RawItemName::Name(n) | RawItemName::Id(_, n, _) => self.visit_unresolved_item_name(n),
424        }
425    }
426}
427
428struct CreateSqlRewriter {
429    from: Vec<Ident>,
430    to: Ident,
431}
432
433impl CreateSqlRewriter {
434    fn rewrite_query_with_qual_depth(
435        from_name: FullItemName,
436        to_name: Ident,
437        qual_depth: usize,
438        query: &mut Query<Raw>,
439    ) {
440        let from = match qual_depth {
441            1 => vec![Ident::new_unchecked(from_name.item)],
442            2 => vec![
443                Ident::new_unchecked(from_name.schema),
444                Ident::new_unchecked(from_name.item),
445            ],
446            3 => vec![
447                Ident::new_unchecked(from_name.database.to_string()),
448                Ident::new_unchecked(from_name.schema),
449                Ident::new_unchecked(from_name.item),
450            ],
451            _ => unreachable!(),
452        };
453        let mut v = CreateSqlRewriter { from, to: to_name };
454        v.visit_query_mut(query);
455    }
456
457    fn maybe_rewrite_idents(&self, name: &mut [Ident]) {
458        if name.len() > 0 && name.ends_with(&self.from) {
459            name[name.len() - 1] = self.to.clone();
460        }
461    }
462}
463
464impl<'ast> VisitMut<'ast, Raw> for CreateSqlRewriter {
465    fn visit_expr_mut(&mut self, e: &'ast mut Expr<Raw>) {
466        match e {
467            Expr::Identifier(id) => {
468                // The last ID component is a column name that should not be
469                // considered in the rewrite.
470                let i = id.len() - 1;
471                self.maybe_rewrite_idents(&mut id[..i]);
472            }
473            Expr::QualifiedWildcard(id) => {
474                self.maybe_rewrite_idents(id);
475            }
476            _ => visit_mut::visit_expr_mut(self, e),
477        }
478    }
479    fn visit_unresolved_item_name_mut(
480        &mut self,
481        unresolved_item_name: &'ast mut UnresolvedItemName,
482    ) {
483        self.maybe_rewrite_idents(&mut unresolved_item_name.0);
484    }
485    fn visit_item_name_mut(
486        &mut self,
487        item_name: &'ast mut <mz_sql_parser::ast::Raw as AstInfo>::ItemName,
488    ) {
489        match item_name {
490            RawItemName::Name(n) | RawItemName::Id(_, n, _) => self.maybe_rewrite_idents(&mut n.0),
491        }
492    }
493}
494
495/// Updates all `CatalogItemId`s from the keys of `ids` to the values of `ids` within `create_stmt`.
496pub fn create_stmt_replace_ids(
497    create_stmt: &mut Statement<Raw>,
498    ids: &BTreeMap<CatalogItemId, CatalogItemId>,
499) {
500    let mut id_replacer = CreateSqlIdReplacer { ids };
501    id_replacer.visit_statement_mut(create_stmt);
502}
503
504struct CreateSqlIdReplacer<'a> {
505    ids: &'a BTreeMap<CatalogItemId, CatalogItemId>,
506}
507
508impl<'ast> VisitMut<'ast, Raw> for CreateSqlIdReplacer<'_> {
509    fn visit_item_name_mut(
510        &mut self,
511        item_name: &'ast mut <mz_sql_parser::ast::Raw as AstInfo>::ItemName,
512    ) {
513        match item_name {
514            RawItemName::Id(id, _, _) => {
515                let old_id = match id.parse() {
516                    Ok(old_id) => old_id,
517                    Err(e) => panic!("invalid persisted global id {id}: {e}"),
518                };
519                if let Some(new_id) = self.ids.get(&old_id) {
520                    *id = new_id.to_string();
521                }
522            }
523            RawItemName::Name(_) => {}
524        }
525    }
526}
mz_sql/ast/transform.rs

mz_sql/ast/
transform.rs
