fix: Ensure only tables or aliases that exist are projected (#52)

fix: More dangling references (#54)

UPSTREAM NOTE: This PR was attempted to be upstreamed in apache#13405 - but it was not accepted due to the complexity it brought. Phillip needs to figure out what a good solution that solves our problem and can be upstreamed is.

Author: peasee

datafusion/sql/src/unparser/ast.rs

@@ -53,6 +53,9 @@ impl QueryBuilder {
         self.order_by_kind = Some(value);
         self
     }
+    pub fn get_order_by(&self) -> Option<OrderByKind> {
+        self.order_by_kind.clone()
+    }
     pub fn limit(&mut self, value: Option<ast::Expr>) -> &mut Self {
         self.limit = value;
         self
@@ -161,6 +164,9 @@ impl SelectBuilder {
         self.top = value;
         self
     }
+    pub fn get_projection(&self) -> Vec<ast::SelectItem> {
+        self.projection.clone()
+    }
     pub fn projection(&mut self, value: Vec<ast::SelectItem>) -> &mut Self {
         self.projection = value;
         self
@@ -264,6 +270,9 @@ impl SelectBuilder {
         self.sort_by = value;
         self
     }
+    pub fn get_sort_by(&self) -> Vec<ast::Expr> {
+        self.sort_by.clone()
+    }
     pub fn having(&mut self, value: Option<ast::Expr>) -> &mut Self {
         self.having = value;
         self
@@ -356,7 +365,9 @@ impl TableWithJoinsBuilder {
         self.relation = Some(value);
         self
     }
-
+    pub fn get_joins(&self) -> Vec<ast::Join> {
+        self.joins.clone()
+    }
     pub fn joins(&mut self, value: Vec<ast::Join>) -> &mut Self {
         self.joins = value;
         self
@@ -411,6 +422,25 @@ impl RelationBuilder {
     pub fn has_relation(&self) -> bool {
         self.relation.is_some()
     }
+    pub fn get_name(&self) -> Option<String> {
+        match self.relation {
+            Some(TableFactorBuilder::Table(ref value)) => {
+                value.name.as_ref().map(|a| a.to_string())
+            }
+            _ => None,
+        }
+    }
+    pub fn get_alias(&self) -> Option<String> {
+        match self.relation {
+            Some(TableFactorBuilder::Table(ref value)) => {
+                value.alias.as_ref().map(|a| a.name.to_string())
+            }
+            Some(TableFactorBuilder::Derived(ref value)) => {
+                value.alias.as_ref().map(|a| a.name.to_string())
+            }
+            _ => None,
+        }
+    }
     pub fn table(&mut self, value: TableRelationBuilder) -> &mut Self {
         self.relation = Some(TableFactorBuilder::Table(value));
         self

datafusion/sql/src/unparser/plan.rs

@@ -22,7 +22,7 @@ use super::{
     },
     rewrite::{
         inject_column_aliases_into_subquery, normalize_union_schema,
-        rewrite_plan_for_sort_on_non_projected_fields,
+        remove_dangling_identifiers, rewrite_plan_for_sort_on_non_projected_fields,
         subquery_alias_inner_query_and_columns, TableAliasRewriter,
     },
     utils::{
@@ -209,10 +209,70 @@ impl Unparser<'_> {
             )]);
         }
 
+        // Construct a list of all the identifiers present in query sources
+        let mut all_idents = Vec::new();
+        if let Some(source_alias) = relation_builder.get_alias() {
+            all_idents.push(source_alias);
+        } else if let Some(source_name) = relation_builder.get_name() {
+            all_idents.push(source_name);
+        }
+
         let mut twj = select_builder.pop_from().unwrap();
+        twj.get_joins()
+            .iter()
+            .for_each(|join| match &join.relation {
+                ast::TableFactor::Table { alias, name, .. } => {
+                    if let Some(alias) = alias {
+                        all_idents.push(alias.name.to_string());
+                    } else {
+                        all_idents.push(name.to_string());
+                    }
+                }
+                ast::TableFactor::Derived { alias, .. } => {
+                    if let Some(alias) = alias {
+                        all_idents.push(alias.name.to_string());
+                    }
+                }
+                _ => {}
+            });
+
         twj.relation(relation_builder);
         select_builder.push_from(twj);
 
+        // Ensure that the projection contains references to sources that actually exist
+        let mut projection = select_builder.get_projection();
+        projection
+            .iter_mut()
+            .for_each(|select_item| match select_item {
+                ast::SelectItem::UnnamedExpr(ast::Expr::CompoundIdentifier(idents)) => {
+                    remove_dangling_identifiers(idents, &all_idents);
+                }
+                _ => {}
+            });
+
+        // Check the order by as well
+        if let Some(query) = query.as_mut() {
+            if let Some(OrderByKind::Expressions(mut order_by)) = query.get_order_by() {
+                order_by.iter_mut().for_each(|sort_item| {
+                    if let ast::Expr::CompoundIdentifier(idents) = &mut sort_item.expr {
+                        remove_dangling_identifiers(idents, &all_idents);
+                    }
+                });
+
+                query.order_by(OrderByKind::Expressions(order_by));
+            }
+        }
+
+        // Order by could be a sort in the select builder
+        let mut sort = select_builder.get_sort_by();
+        sort.iter_mut().for_each(|sort_item| {
+            if let ast::Expr::CompoundIdentifier(idents) = sort_item {
+                remove_dangling_identifiers(idents, &all_idents);
+            }
+        });
+
+        select_builder.projection(projection);
+
         Ok(SetExpr::Select(Box::new(select_builder.build()?)))
     }
 

datafusion/sql/src/unparser/rewrite.rs

@@ -25,7 +25,7 @@ use datafusion_common::{
 };
 use datafusion_expr::expr::{Alias, UNNEST_COLUMN_PREFIX};
 use datafusion_expr::{Expr, LogicalPlan, Projection, Sort, SortExpr};
-use sqlparser::ast::Ident;
+use sqlparser::ast::{display_separated, Ident};
 
 /// Normalize the schema of a union plan to remove qualifiers from the schema fields and sort expressions.
 ///
@@ -418,3 +418,68 @@ impl TreeNodeRewriter for TableAliasRewriter<'_> {
         }
     }
 }
+
+/// Takes an input list of identifiers and a list of identifiers that are available from relations or joins.
+/// Removes any table identifiers that are not present in the list of available identifiers, retains original column names.
+pub fn remove_dangling_identifiers(
+    idents: &mut Vec<Ident>,
+    available_idents: &Vec<String>,
+) -> () {
+    if idents.len() > 1 {
+        let ident_source = display_separated(
+            &idents
+                .clone()
+                .into_iter()
+                .take(idents.len() - 1)
+                .collect::<Vec<Ident>>(),
+            ".",
+        )
+        .to_string();
+        // If the identifier is not present in the list of all identifiers, it refers to a table that does not exist
+        if !available_idents.contains(&ident_source) {
+            let Some(last) = idents.last() else {
+                unreachable!("CompoundIdentifier must have a last element");
+            };
+            // Reset the identifiers to only the last element, which is the column name
+            *idents = vec![last.clone()];
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_remove_dangling_identifiers() {
+        let tests = vec![
+            (vec![], vec![Ident::new("column1".to_string())]),
+            (
+                vec!["table1.table2".to_string()],
+                vec![
+                    Ident::new("table1".to_string()),
+                    Ident::new("table2".to_string()),
+                    Ident::new("column1".to_string()),
+                ],
+            ),
+            (
+                vec!["table1".to_string()],
+                vec![Ident::new("column1".to_string())],
+            ),
+        ];
+
+        for test in tests {
+            let test_in = test.0;
+            let test_out = test.1;
+
+            let mut idents = vec![
+                Ident::new("table1".to_string()),
+                Ident::new("table2".to_string()),
+                Ident::new("column1".to_string()),
+            ];
+
+            remove_dangling_identifiers(&mut idents, &test_in);
+            assert_eq!(idents, test_out);
+        }
+    }
+}