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

fix: More dangling references (#54)

* fix: More dangling references

* test: Add tests for remove_dangling_identifiers

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

@@ -46,6 +46,9 @@ impl QueryBuilder {
     pub fn take_body(&mut self) -> Option<Box<ast::SetExpr>> {
         self.body.take()
     }
+    pub fn get_order_by(&self) -> Vec<ast::OrderByExpr> {
+        self.order_by.clone()
+    }
     pub fn order_by(&mut self, value: Vec<ast::OrderByExpr>) -> &mut Self {
         self.order_by = value;
         self
@@ -150,6 +153,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
@@ -217,6 +223,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
@@ -304,7 +313,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
@@ -358,6 +369,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,69 @@ 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() {
+            let 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(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.
 ///
@@ -417,3 +417,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);
+        }
+    }
+}

datafusion/sql/tests/cases/plan_to_sql.rs

@@ -366,6 +366,80 @@ fn roundtrip_statement_with_dialect() -> Result<()> {
             parser_dialect: Box::new(GenericDialect {}),
             unparser_dialect: Box::new(UnparserDefaultDialect {}),
         },
+        TestStatementWithDialect {
+            sql: "select j1_id from (select ta.j1_id from j1 ta)",
+            expected:
+                // This seems like desirable behavior, but is actually hiding an underlying issue
+                // The re-written identifier is `ta`.`j1_id`, because `reconstuct_select_statement` runs before the derived projection
+                // and for some reason, the derived table alias is pre-set to `ta` for the top-level projection
+                "SELECT `j1_id` FROM (SELECT `ta`.`j1_id` FROM `j1` AS `ta`) AS `derived_projection`",
+            parser_dialect: Box::new(MySqlDialect {}),
+            unparser_dialect: Box::new(UnparserMySqlDialect {}),
+        },
+        TestStatementWithDialect {
+            sql: "select j1_id from (select ta.j1_id from j1 ta)",
+            expected:
+                "SELECT j1_id FROM (SELECT ta.j1_id FROM j1 AS ta)",
+            parser_dialect: Box::new(GenericDialect {}),
+            unparser_dialect: Box::new(UnparserDefaultDialect {}),
+        },
+        TestStatementWithDialect {
+            sql: "select j1_id from (select ta.j1_id from j1 ta) order by j1_id",
+            expected:
+                "SELECT j1_id FROM (SELECT ta.j1_id FROM j1 AS ta) ORDER BY j1_id ASC NULLS LAST",
+            parser_dialect: Box::new(GenericDialect {}),
+            unparser_dialect: Box::new(UnparserDefaultDialect {}),
+        },
+        // TODO: remove dangling identifiers from group by, filter, etc
+        // TestStatementWithDialect {
+        //     sql: "select j1_id from (select ta.j1_id from j1 ta) where j1_id = 1",
+        //     expected:
+        //         "SELECT j1_id FROM (SELECT ta.j1_id FROM j1 AS ta) WHERE (ta.j1_id = 1)",
+        //     parser_dialect: Box::new(GenericDialect {}),
+        //     unparser_dialect: Box::new(UnparserDefaultDialect {}),
+        // },
+        TestStatementWithDialect {
+            sql: "select j1_id from (select ta.j1_id from j1 ta) order by j1_id",
+            expected:
+                "SELECT `j1_id` FROM (SELECT `ta`.`j1_id` FROM `j1` AS `ta`) AS `derived_projection` ORDER BY `j1_id` ASC",
+            parser_dialect: Box::new(MySqlDialect {}),
+            unparser_dialect: Box::new(UnparserMySqlDialect {}),
+        },
+        TestStatementWithDialect {
+            sql: "select j1_id from (select ta.j1_id from j1 ta) AS tbl1",
+            expected:
+                "SELECT tbl1.j1_id FROM (SELECT ta.j1_id FROM j1 AS ta) AS tbl1",
+            parser_dialect: Box::new(GenericDialect {}),
+            unparser_dialect: Box::new(UnparserDefaultDialect {}),
+        },
+        TestStatementWithDialect {
+            sql: "select j1_id, j2_id from (select ta.j1_id from j1 ta) AS tbl1, (select ta.j1_id as j2_id from j1 ta) as tbl2",
+            expected:
+                "SELECT tbl1.j1_id, tbl2.j2_id FROM (SELECT ta.j1_id FROM j1 AS ta) AS tbl1 JOIN (SELECT ta.j1_id AS j2_id FROM j1 AS ta) AS tbl2 ON true",
+            parser_dialect: Box::new(GenericDialect {}),
+            unparser_dialect: Box::new(UnparserDefaultDialect {}),
+        },
+        TestStatementWithDialect {
+            sql: "select j1_id, j2_id from (select ta.j1_id from j1 ta) AS tbl1, (select ta.j1_id as j2_id from j1 ta) as tbl2",
+            expected:
+                "SELECT `tbl1`.`j1_id`, `tbl2`.`j2_id` FROM (SELECT `ta`.`j1_id` FROM `j1` AS `ta`) AS `tbl1` JOIN (SELECT `ta`.`j1_id` AS `j2_id` FROM `j1` AS `ta`) AS `tbl2` ON true",
+            parser_dialect: Box::new(MySqlDialect {}),
+            unparser_dialect: Box::new(UnparserMySqlDialect {}),
+        },
+        TestStatementWithDialect {
+            sql: "select j1_id, j2_id from (select ta.j1_id from j1 ta), (select ta.j1_id as j2_id from j1 ta)",
+            expected:
+                "SELECT `j1_id`, `j2_id` FROM (SELECT `ta`.`j1_id` FROM `j1` AS `ta`) AS `derived_projection` JOIN (SELECT `ta`.`j1_id` AS `j2_id` FROM `j1` AS `ta`) AS `derived_projection` ON true",
+            parser_dialect: Box::new(MySqlDialect {}),
+            unparser_dialect: Box::new(UnparserMySqlDialect {}),
+        },
+        TestStatementWithDialect {
+            sql: "select j1_id, j2_id from (select ta.j1_id from j1 ta), (select ta.j1_id AS j2_id from j1 ta)",
+            expected:
+                "SELECT j1_id, j2_id FROM (SELECT ta.j1_id FROM j1 AS ta) JOIN (SELECT ta.j1_id AS j2_id FROM j1 AS ta) ON true",
+            parser_dialect: Box::new(GenericDialect {}),
+            unparser_dialect: Box::new(UnparserDefaultDialect {}),
+        },
         TestStatementWithDialect {
             sql: "SELECT j1_string from j1 join j2 on j1.j1_id = j2.j2_id order by j1_id",
             expected: r#"SELECT j1.j1_string FROM j1 JOIN j2 ON (j1.j1_id = j2.j2_id) ORDER BY j1.j1_id ASC NULLS LAST"#,