diff --git a/wren-core/core/src/logical_plan/analyze/model_generation.rs b/wren-core/core/src/logical_plan/analyze/model_generation.rs index 972eb9bab..efbe7e9da 100644 --- a/wren-core/core/src/logical_plan/analyze/model_generation.rs +++ b/wren-core/core/src/logical_plan/analyze/model_generation.rs @@ -117,8 +117,9 @@ impl ModelGenerationRule { // The filter should be on on the top of the model plan // and the model plan should be another subquery alias if let Some(filter) = rls_filter { - builder = - builder.alias(model_plan.plan_name())?.filter(filter)?; + builder = builder + .alias(quoted(model_plan.plan_name()))? + .filter(filter)?; // Following the DataFusion planning behavior, we need to // add a projection behind the filter to ensure the unparsing is correct. let indices = 0..builder.schema().fields().len(); diff --git a/wren-core/core/src/mdl/function.rs b/wren-core/core/src/mdl/function.rs index 37d1457ea..a8ed27283 100644 --- a/wren-core/core/src/mdl/function.rs +++ b/wren-core/core/src/mdl/function.rs @@ -366,6 +366,7 @@ impl WindowUDFImpl for ByPassWindowFunction { #[cfg(test)] mod test { + use std::slice::from_ref; use std::sync::Arc; use crate::mdl::function::{ @@ -598,9 +599,7 @@ mod test { DataType::List(Arc::new(Field::new("element", DataType::Int32, false))); assert_eq!(udf.name, "test"); assert_eq!( - udf.return_type - .to_data_type(std::slice::from_ref(&list_type)) - .unwrap(), + udf.return_type.to_data_type(from_ref(&list_type)).unwrap(), DataType::Int32 ); assert_eq!( diff --git a/wren-core/core/src/mdl/mod.rs b/wren-core/core/src/mdl/mod.rs index 4640a0cbf..f4585d443 100644 --- a/wren-core/core/src/mdl/mod.rs +++ b/wren-core/core/src/mdl/mod.rs @@ -2523,6 +2523,46 @@ mod test { Ok(()) } + #[tokio::test] + async fn test_rlac_unicode_model_column_name() -> Result<()> { + let ctx = SessionContext::new(); + let manifest = ManifestBuilder::new() + .catalog("wren") + .schema("test") + .model( + ModelBuilder::new("VTU藝人") + .table_reference("artist") + .column(ColumnBuilder::new("名字", "string").build()) + .column(ColumnBuilder::new("組別", "string").build()) + .column(ColumnBuilder::new("訂閱數", "int").build()) + .add_row_level_access_control( + "rule", + vec![SessionProperty::new_required("預定組別A")], + "\"組別\" = @預定組別A", + ) + .build(), + ) + .build(); + let headers = Arc::new(build_headers(&[( + "預定組別A".to_string(), + Some("'JP'".to_string()), + )])); + + let analyzed_mdl = Arc::new(AnalyzedWrenMDL::analyze( + manifest, + Arc::clone(&headers), + Mode::Unparse, + )?); + + let sql = r#"SELECT "名字", "組別", "訂閱數" FROM "VTU藝人""#; + assert_snapshot!( + transform_sql_with_ctx(&ctx, Arc::clone(&analyzed_mdl), &[], headers, sql) + .await?, + @r#"SELECT "VTU藝人"."名字", "VTU藝人"."組別", "VTU藝人"."訂閱數" FROM (SELECT "VTU藝人"."名字", "VTU藝人"."組別", "VTU藝人"."訂閱數" FROM (SELECT "VTU藝人"."名字", "VTU藝人"."組別", "VTU藝人"."訂閱數" FROM (SELECT __source."名字" AS "名字", __source."組別" AS "組別", __source."訂閱數" AS "訂閱數" FROM artist AS __source) AS "VTU藝人") AS "VTU藝人" WHERE "VTU藝人"."組別" = 'JP') AS "VTU藝人""# + ); + Ok(()) + } + #[tokio::test] async fn test_clac_with_required_properties() -> Result<()> { let ctx = SessionContext::new();