Skip to content

fix(wren-core): resolve relationship handles by the model they point at #2661

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wilyan09007 wants to merge 2 commits into
base: main
Choose a base branch
from
+298 −5
Open

fix(wren-core): resolve relationship handles by the model they point at #2661

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6295a5a
fix(wren-core): resolve relationship handles by the model they point at
wilyan09007 Aug 11, 2026
728eee9
fix(wren-core): resolve a path component that shares its model's name
wilyan09007 Aug 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
60 changes: 60 additions & 0 deletions core/wren-core/core/src/mdl/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2122,6 +2122,66 @@ mod test {
Ok(())
}

/// A relationship handle names the join locally, so it may differ from the model it
/// points at. Here the handle `customer` on `orders` resolves to the model `customers`.
#[tokio::test]
async fn test_calculated_column_with_aliased_relationship_handle() -> Result<()> {
let ctx = create_wren_ctx(None, None);
let manifest = ManifestBuilder::new()
.catalog("wren")
.schema("test")
.model(
ModelBuilder::new("orders")
.table_reference("orders")
.column(ColumnBuilder::new("order_id", "int").build())
.column(ColumnBuilder::new("customer_id", "int").build())
.column(
ColumnBuilder::new_relationship(
"customer",
"customers",
"orders_customers",
)
.build(),
)
.column(
ColumnBuilder::new_calculated("customer_name", "string")
.expression("customer.name")
.build(),
)
.primary_key("order_id")
.build(),
)
.model(
ModelBuilder::new("customers")
.table_reference("customers")
.column(ColumnBuilder::new("customer_id", "int").build())
.column(ColumnBuilder::new("name", "string").build())
.primary_key("customer_id")
.build(),
)
.relationship(
RelationshipBuilder::new("orders_customers")
.model("orders")
.model("customers")
.join_type(JoinType::ManyToOne)
.condition("orders.customer_id = customers.customer_id")
.build(),
)
.build();
let analyzed_mdl = Arc::new(AnalyzedWrenMDL::analyze(
manifest,
Arc::new(HashMap::default()),
Mode::Unparse,
)?);

let sql = "SELECT order_id, customer_name FROM orders";
assert_snapshot!(
transform_sql_with_ctx(&ctx, Arc::clone(&analyzed_mdl), &[], Arc::new(HashMap::new()), sql).await?,
@r#"SELECT orders.order_id, orders.customer_name FROM (SELECT __relation__1."name" AS customer_name, __relation__1.order_id FROM (SELECT orders.customer_id, customers."name", orders.order_id FROM (SELECT customers.customer_id, customers."name" FROM (SELECT customers.customer_id, customers."name" FROM (SELECT __source.customer_id AS customer_id, __source."name" AS "name" FROM customers AS __source) AS customers) AS customers) AS customers RIGHT OUTER JOIN (SELECT __source.customer_id AS customer_id, __source.order_id AS order_id FROM orders AS __source) AS orders ON customers.customer_id = orders.customer_id) AS __relation__1) AS orders"#
);
Ok(())
}

#[tokio::test]
async fn test_rlac_with_requried_properties() -> Result<()> {
let ctx = create_wren_ctx(None, None);
Expand Down
243 changes: 238 additions & 5 deletions core/wren-core/core/src/mdl/utils.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ use std::sync::Arc;

use crate::logical_plan::utils::{from_qualified_name, try_map_data_type};
use crate::mdl::manifest::Model;
use crate::mdl::{AnalyzedWrenMDL, ColumnReference, Dataset, SessionStateRef};
use crate::mdl::{AnalyzedWrenMDL, ColumnReference, Dataset, SessionStateRef, WrenMDL};

pub fn to_expr_queue(column: Column) -> VecDeque<String> {
column.name.split('.').map(String::from).collect()
Expand Down Expand Up @@ -188,15 +188,22 @@ pub fn create_wren_calculated_field_expr(
.collect::<Vec<String>>();
// Remove all relationship fields from the expression. Only keep the target expression and its source table.
let expr = column_rf.column.expression.clone().unwrap();
let base_model = column_rf.dataset.name();
let session_state = session_state.read();
let mut expr = session_state
.sql_to_expr(&expr, &session_state.config_options().sql_parser.dialect)?;
let _ = visit_expressions_mut(&mut expr, |e| {
if let CompoundIdentifier(ids) = e {
let name_size = ids.len();
if name_size > 2 {
let slice = &ids[name_size - 2..name_size];
*e = CompoundIdentifier(slice.to_vec());
if let Some(resolved) =
resolve_relationship_path(&analyzed_wren_mdl.wren_mdl, base_model, ids)
{
*e = CompoundIdentifier(resolved);
} else {
let name_size = ids.len();
if name_size > 2 {
let slice = &ids[name_size - 2..name_size];
*e = CompoundIdentifier(slice.to_vec());
}
}
}
ControlFlow::<()>::Continue(())
Expand All @@ -216,6 +223,50 @@ pub fn create_wren_calculated_field_expr(
session_state.create_logical_expr(&expr.to_string(), &schema)
}

/// Rewrite a relationship path so that it is qualified by the model the path lands on.
///
/// A relationship column is a handle whose name is local to the model declaring it, so a
/// handle `customer` on `orders` may point at the model `customers`. The schema the
/// expression is planned against is built from the lineage's required fields and is
/// therefore qualified by model name, so each handle has to be translated to the model it
/// resolves to. On `orders`, `customer.name` becomes `customers.name`.
///
/// Every component before the column is a handle on the model reached so far, the first
/// one included: [`crate::mdl::lineage`] walks the path the same way and rejects a leading
/// component that is not a column of the base model, so a path cannot arrive here carrying
/// its own model as a qualifier. A handle that shares a name with the model declaring it
/// is therefore still a handle.
///
/// Returns `None` when the path traverses no relationship, leaving the identifier alone.
fn resolve_relationship_path(
wren_mdl: &WrenMDL,
base_model: &str,
ids: &[Ident],
) -> Option<Vec<Ident>> {
let (column, path) = ids.split_last()?;
if path.is_empty() {
return None;
}
let mut relation = base_model.to_string();
for ident in path {
let column_ref = wren_mdl.get_column_reference(&from_qualified_name(
wren_mdl,
&relation,
&ident.value,
))?;
let relationship =
wren_mdl.get_relationship(column_ref.column.relationship.as_ref()?)?;
// Resolve the related model the way [`crate::mdl::lineage`] does, so the qualifier
// matches the schema built from the required fields it collected.
relation = relationship
.models
.iter()
.find(|model| *model != &relation)?
.clone();
}
Some(vec![quoted_ident(&relation), column.clone()])
}

/// Create the Logical Expr for the remote column.
pub(crate) fn create_remote_expr_for_model(
expr: &str,
Expand Down Expand Up @@ -350,8 +401,12 @@ mod tests {

use datafusion::error::Result;
use datafusion::prelude::SessionContext;
use wren_core_base::mdl::JoinType;

use crate::logical_plan::utils::from_qualified_name;
use crate::mdl::builder::{
ColumnBuilder, ManifestBuilder, ModelBuilder, RelationshipBuilder,
};
use crate::mdl::context::Mode;
use crate::mdl::manifest::Manifest;
use crate::mdl::AnalyzedWrenMDL;
Expand Down Expand Up @@ -388,6 +443,184 @@ mod tests {
Ok(())
}

/// A relationship handle is a local alias, so its name need not match the model it
/// points at. Here the handle `customer` on `orders` resolves to the model `customers`.
#[test]
fn test_create_wren_expr_handle_name_differs_from_model() -> Result<()> {
let manifest = ManifestBuilder::new()
.catalog("wren")
.schema("test")
.model(
ModelBuilder::new("orders")
.table_reference("orders")
.column(ColumnBuilder::new("order_id", "integer").build())
.column(ColumnBuilder::new("customer_id", "integer").build())
.column(
ColumnBuilder::new_relationship(
"customer",
"customers",
"orders_customers",
)
.build(),
)
.column(
ColumnBuilder::new_calculated("customer_name", "varchar")
.expression("customer.name")
.build(),
)
.primary_key("order_id")
.build(),
)
.model(
ModelBuilder::new("customers")
.table_reference("customers")
.column(ColumnBuilder::new("customer_id", "integer").build())
.column(ColumnBuilder::new("name", "varchar").build())
.primary_key("customer_id")
.build(),
)
.relationship(
RelationshipBuilder::new("orders_customers")
.model("orders")
.model("customers")
.join_type(JoinType::ManyToOne)
.condition("orders.customer_id = customers.customer_id")
.build(),
)
.build();
let analyzed_mdl = Arc::new(AnalyzedWrenMDL::analyze(
manifest,
Arc::new(HashMap::default()),
Mode::Unparse,
)?);
let ctx = SessionContext::new();
let column_rf = analyzed_mdl
.wren_mdl
.qualified_references
.get(&from_qualified_name(
&analyzed_mdl.wren_mdl,
"orders",
"customer_name",
))
.unwrap();
let expr = super::create_wren_calculated_field_expr(
column_rf.clone(),
Arc::clone(&analyzed_mdl),
ctx.state_ref(),
)?;
// qualified by the related model, not by the handle that reached it
assert_eq!(expr.to_string(), "customers.name");
Ok(())
}

/// A handle may share a name with the model that declares it. It is still a handle:
/// every component before the column names a relationship on the model reached so
/// far, the first one included.
#[test]
fn test_create_wren_expr_handle_named_after_its_own_model() -> Result<()> {
let manifest = ManifestBuilder::new()
.catalog("wren")
.schema("test")
.model(
ModelBuilder::new("orders")
.table_reference("orders")
.column(ColumnBuilder::new("order_id", "integer").build())
.column(ColumnBuilder::new("customer_id", "integer").build())
.column(
ColumnBuilder::new_relationship(
"orders",
"customers",
"orders_customers",
)
.build(),
)
.column(
ColumnBuilder::new_calculated("customer_name", "varchar")
.expression("orders.name")
.build(),
)
.column(
ColumnBuilder::new_calculated("customer_city", "varchar")
.expression("orders.customers.name")
.build(),
)
.primary_key("order_id")
.build(),
)
.model(
ModelBuilder::new("customers")
.table_reference("customers")
.column(ColumnBuilder::new("customer_id", "integer").build())
.column(ColumnBuilder::new("city_id", "integer").build())
.column(ColumnBuilder::new("name", "varchar").build())
.column(
ColumnBuilder::new_relationship(
"customers",
"cities",
"customers_cities",
)
.build(),
)
.primary_key("customer_id")
.build(),
)
.model(
ModelBuilder::new("cities")
.table_reference("cities")
.column(ColumnBuilder::new("city_id", "integer").build())
.column(ColumnBuilder::new("name", "varchar").build())
.primary_key("city_id")
.build(),
)
.relationship(
RelationshipBuilder::new("orders_customers")
.model("orders")
.model("customers")
.join_type(JoinType::ManyToOne)
.condition("orders.customer_id = customers.customer_id")
.build(),
)
.relationship(
RelationshipBuilder::new("customers_cities")
.model("customers")
.model("cities")
.join_type(JoinType::ManyToOne)
.condition("customers.city_id = cities.city_id")
.build(),
)
.build();
let analyzed_mdl = Arc::new(AnalyzedWrenMDL::analyze(
manifest,
Arc::new(HashMap::default()),
Mode::Unparse,
)?);
let ctx = SessionContext::new();
let ctx_state = ctx.state_ref();
let resolve = |calculated: &str| -> Result<String> {
let column_rf = analyzed_mdl
.wren_mdl
.qualified_references
.get(&from_qualified_name(
&analyzed_mdl.wren_mdl,
"orders",
calculated,
))
.unwrap();
Ok(super::create_wren_calculated_field_expr(
column_rf.clone(),
Arc::clone(&analyzed_mdl),
Arc::clone(&ctx_state),
)?
.to_string())
};
// the leading `orders` is the handle declared on `orders`, not the model itself
assert_eq!(resolve("customer_name")?, "customers.name");
// `customers` has a `name` of its own, so stopping a hop short would plan against
// a real but wrong column rather than fail
assert_eq!(resolve("customer_city")?, "cities.name");
Ok(())
}

#[test]
fn test_create_wren_expr_non_relationship() -> Result<()> {
let test_data: PathBuf =
Expand Down
Loading