Ensure keys are qualified

Signed-off-by: Abhi Agarwal <[email protected]>

Author: abhi-airspace-intelligence

etl-destinations/src/deltalake/expr.rs

@@ -529,4 +529,71 @@ mod tests {
         let res = qualify_primary_keys(primary_keys, "source", "target");
         assert!(res.is_none());
     }
+
+    #[test]
+    fn test_qualify_primary_keys_case_sensitivity() {
+        let primary_keys = vec![
+            Expr::Column(Column::new_unqualified("id")),
+            Expr::Column(Column::new_unqualified("NAME")),
+        ];
+        let result = qualify_primary_keys(primary_keys, "source", "target").unwrap();
+
+        assert_debug_snapshot!(result, @r#"
+        BinaryExpr(
+            BinaryExpr {
+                left: BinaryExpr(
+                    BinaryExpr {
+                        left: Column(
+                            Column {
+                                relation: Some(
+                                    Bare {
+                                        table: "source",
+                                    },
+                                ),
+                                name: "id",
+                            },
+                        ),
+                        op: Eq,
+                        right: Column(
+                            Column {
+                                relation: Some(
+                                    Bare {
+                                        table: "target",
+                                    },
+                                ),
+                                name: "id",
+                            },
+                        ),
+                    },
+                ),
+                op: And,
+                right: BinaryExpr(
+                    BinaryExpr {
+                        left: Column(
+                            Column {
+                                relation: Some(
+                                    Bare {
+                                        table: "source",
+                                    },
+                                ),
+                                name: "NAME",
+                            },
+                        ),
+                        op: Eq,
+                        right: Column(
+                            Column {
+                                relation: Some(
+                                    Bare {
+                                        table: "target",
+                                    },
+                                ),
+                                name: "NAME",
+                            },
+                        ),
+                    },
+                ),
+            },
+        )
+        "#);
+    }
 }

etl-destinations/src/deltalake/operations/merge.rs

@@ -1,6 +1,6 @@
 use deltalake::DeltaTableError;
 use deltalake::datafusion::common::Column;
-use deltalake::datafusion::prelude::{SessionContext, col};
+use deltalake::datafusion::prelude::SessionContext;
 use deltalake::operations::merge::MergeBuilder;
 use deltalake::{DeltaResult, DeltaTable, datafusion::prelude::Expr};
 use etl::types::{TableRow as PgTableRow, TableSchema as PgTableSchema};
@@ -10,6 +10,10 @@ use crate::deltalake::config::DeltaTableConfig;
 use crate::deltalake::expr::qualify_primary_keys;
 use crate::deltalake::schema::postgres_to_arrow_schema;
 
+pub(crate) fn source_qualified_column_expr(column_name: &str, source_alias: &str) -> Expr {
+    Expr::Column(Column::new(Some(source_alias), column_name))
+}
+
 pub async fn merge_to_table(
     table: &mut DeltaTable,
     config: &DeltaTableConfig,
@@ -55,12 +59,18 @@ pub async fn merge_to_table(
         .with_target_alias("target")
         .when_not_matched_insert(|insert| {
             all_columns.iter().fold(insert, |insert, &column| {
-                insert.set(column.to_string(), col(format!("source.{column}")))
+                insert.set(
+                    column.to_string(),
+                    source_qualified_column_expr(column, "source"),
+                )
             })
         })?
         .when_matched_update(|update| {
             all_columns.iter().fold(update, |update, &column| {
-                update.update(column.to_string(), col(format!("source.{column}")))
+                update.update(
+                    column.to_string(),
+                    source_qualified_column_expr(column, "source"),
+                )
             })
         })?;
 
@@ -73,3 +83,45 @@ pub async fn merge_to_table(
     *table = merged_table;
     Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use insta::assert_debug_snapshot;
+
+    #[test]
+    fn source_qualified_column_expr_preserves_case_and_alias() {
+        let expr = source_qualified_column_expr("CASESensitivecolumn", "source");
+
+        assert_debug_snapshot!(expr, @r#"
+        Column(
+            Column {
+                relation: Some(
+                    Bare {
+                        table: "source",
+                    },
+                ),
+                name: "CASESensitivecolumn",
+            },
+        )
+        "#);
+    }
+
+    #[test]
+    fn source_qualified_column_expr_handles_lowercase() {
+        let expr = source_qualified_column_expr("lowercasecolumn", "source");
+
+        assert_debug_snapshot!(expr, @r#"
+        Column(
+            Column {
+                relation: Some(
+                    Bare {
+                        table: "source",
+                    },
+                ),
+                name: "lowercasecolumn",
+            },
+        )
+        "#);
+    }
+}