use warble::{Additivity, ContextLoader, DimensionInfo, LineageGraph, MetricInfo, ModelInfo};
use wren_core_base::mdl::manifest::Manifest;
use crate::lineage;
use crate::project::{assemble, LoadError, ProjectSources};
pub struct MdlContext {
parseable: bool,
parse_error: Option<String>,
metrics: Vec<MetricInfo>,
dimensions: Vec<DimensionInfo>,
time_dimensions: Vec<DimensionInfo>,
models: Vec<ModelInfo>,
lineage: LineageGraph,
lineage_diagnostics: Vec<String>,
}
impl MdlContext {
pub fn from_sources(sources: &ProjectSources) -> Self {
match assemble(sources) {
Ok(loaded) => Self::from_manifest_and_consumers(&loaded.manifest, sources),
Err(_) => Self::unparseable(),
}
}
pub fn try_from_sources(sources: &ProjectSources) -> Result<Self, LoadError> {
assemble(sources).map(|loaded| Self::from_manifest_and_consumers(&loaded.manifest, sources))
}
fn from_manifest_and_consumers(manifest: &Manifest, sources: &ProjectSources) -> Self {
let mut ctx = Self::from_manifest(manifest);
lineage::extend_with_consumers(
&mut ctx.lineage,
manifest,
sources,
&mut ctx.lineage_diagnostics,
);
ctx
}
pub fn unparseable() -> Self {
Self::unparseable_with_error(None)
}
pub fn unparseable_with_error(parse_error: Option<String>) -> Self {
MdlContext {
parseable: false,
parse_error,
metrics: Vec::new(),
dimensions: Vec::new(),
time_dimensions: Vec::new(),
models: Vec::new(),
lineage: LineageGraph::default(),
lineage_diagnostics: Vec::new(),
}
}
pub fn from_manifest(manifest: &Manifest) -> Self {
let mut metrics = Vec::new();
let mut dimensions = Vec::new();
let mut time_dimensions = Vec::new();
let mut models = Vec::new();
for cube in &manifest.cubes {
for measure in &cube.measures {
metrics.push(MetricInfo {
name: measure.name.clone(),
owner: cube.name.clone(),
declared: true,
additivity: Some(infer_additivity(&measure.expression)),
});
}
for dim in &cube.dimensions {
dimensions.push(DimensionInfo {
name: dim.name.clone(),
owner: cube.name.clone(),
is_temporal: false,
});
}
for tdim in &cube.time_dimensions {
let d = DimensionInfo {
name: tdim.name.clone(),
owner: cube.name.clone(),
is_temporal: true,
};
time_dimensions.push(d.clone());
dimensions.push(d);
}
}
for model in &manifest.models {
let mut has_timestamp = false;
let mut column_names = Vec::new();
for col in model.columns.iter().filter(|c| !c.is_hidden) {
column_names.push(col.name.clone());
if col.relationship.is_some() {
continue;
}
if is_temporal_type(&col.r#type) {
has_timestamp = true;
let d = DimensionInfo {
name: col.name.clone(),
owner: model.name.clone(),
is_temporal: true,
};
time_dimensions.push(d.clone());
dimensions.push(d);
} else if is_numeric_type(&col.r#type) {
metrics.push(MetricInfo {
name: col.name.clone(),
owner: model.name.clone(),
declared: false,
additivity: None,
});
} else {
dimensions.push(DimensionInfo {
name: col.name.clone(),
owner: model.name.clone(),
is_temporal: false,
});
}
}
models.push(ModelInfo {
name: model.name.clone(),
has_timestamp,
columns: column_names,
});
}
let (lineage, lineage_diagnostics) = lineage::build(manifest);
MdlContext {
parseable: true,
parse_error: None,
metrics,
dimensions,
time_dimensions,
models,
lineage,
lineage_diagnostics,
}
}
}
impl ContextLoader for MdlContext {
fn is_parseable(&self) -> bool {
self.parseable
}
fn parse_error(&self) -> Option<&str> {
self.parse_error.as_deref()
}
fn metrics(&self) -> &[MetricInfo] {
&self.metrics
}
fn dimensions(&self) -> &[DimensionInfo] {
&self.dimensions
}
fn time_dimensions(&self) -> &[DimensionInfo] {
&self.time_dimensions
}
fn models(&self) -> &[ModelInfo] {
&self.models
}
fn lineage(&self) -> &LineageGraph {
&self.lineage
}
fn lineage_diagnostics(&self) -> &[String] {
&self.lineage_diagnostics
}
}
pub fn infer_additivity(expression: &str) -> Additivity {
let upper = expression.to_uppercase();
if upper.contains("DISTINCT") || upper.contains('/') {
return Additivity::NonAdditive;
}
match leading_function(&upper).as_deref() {
Some("SUM") | Some("COUNT") => Additivity::Additive,
_ => Additivity::NonAdditive,
}
}
fn leading_function(upper_expr: &str) -> Option<String> {
let open = upper_expr.find('(')?;
let name = upper_expr[..open].trim();
if name.is_empty() || !name.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') {
return None;
}
Some(name.to_string())
}
fn normalize_type(t: &str) -> String {
let base = t.split('(').next().unwrap_or(t).trim();
base.to_uppercase()
}
fn is_temporal_type(t: &str) -> bool {
let t = normalize_type(t);
t == "DATE" || t == "DATETIME" || t.starts_with("TIMESTAMP")
}
fn is_numeric_type(t: &str) -> bool {
matches!(
normalize_type(t).as_str(),
"INT"
| "INTEGER"
| "BIGINT"
| "SMALLINT"
| "TINYINT"
| "DECIMAL"
| "NUMERIC"
| "DOUBLE"
| "FLOAT"
| "REAL"
)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn additivity_heuristic() {
assert_eq!(infer_additivity("SUM(amount)"), Additivity::Additive);
assert_eq!(infer_additivity("sum(amount)"), Additivity::Additive);
assert_eq!(infer_additivity("COUNT(*)"), Additivity::Additive);
assert_eq!(
infer_additivity("COUNT(DISTINCT customer_id)"),
Additivity::NonAdditive
);
assert_eq!(infer_additivity("AVG(amount)"), Additivity::NonAdditive);
assert_eq!(infer_additivity("MIN(amount)"), Additivity::NonAdditive);
assert_eq!(infer_additivity("MAX(amount)"), Additivity::NonAdditive);
assert_eq!(infer_additivity("SUM(a) / SUM(b)"), Additivity::NonAdditive);
assert_eq!(infer_additivity("weird(x)"), Additivity::NonAdditive);
assert_eq!(infer_additivity("bare_column"), Additivity::NonAdditive);
}
#[test]
fn cubeless_manifest_cannot_answer_metric_additive() {
use warble::ContextLoader;
use wren_core_base::mdl::manifest::Manifest;
let json = r#"{
"catalog":"wren","schema":"public",
"models":[{"name":"orders","tableReference":{"schema":"main","table":"orders"},
"columns":[{"name":"amount","type":"DOUBLE"},{"name":"status","type":"TEXT"}]}],
"relationships":[],"cubes":[],"views":[]
}"#;
let manifest: Manifest = serde_json::from_str(json).unwrap();
let ctx = MdlContext::from_manifest(&manifest);
assert!(!ctx.metrics().is_empty(), "amount is an implicit metric");
assert!(
!ctx.can_answer("metric_additive"),
"no declared measure ⇒ additivity unanswerable"
);
assert!(ctx
.metrics()
.iter()
.all(|m| !m.declared && m.additivity.is_none()));
}
#[test]
fn type_classification() {
assert!(is_temporal_type("DATE"));
assert!(is_temporal_type("timestamp"));
assert!(is_temporal_type("TIMESTAMP WITH TIME ZONE"));
assert!(!is_temporal_type("TEXT"));
assert!(is_numeric_type("INT"));
assert!(is_numeric_type("BIGINT"));
assert!(is_numeric_type("DECIMAL(10,2)"));
assert!(!is_numeric_type("TEXT"));
assert!(!is_numeric_type("DATE"));
}
#[test]
fn mdl_context_cannot_answer_raw_shape_predicates() {
use crate::MdlContext;
use wren_core_base::mdl::manifest::Manifest;
let json = r#"{
"catalog":"wren","schema":"public",
"models":[],"relationships":[],"cubes":[],"views":[]
}"#;
let manifest: Manifest = serde_json::from_str(json).unwrap();
let ctx = MdlContext::from_manifest(&manifest);
assert_eq!(ctx.source_introspectable(), None);
assert_eq!(ctx.raw_docs_readable(), None);
assert!(!ctx.can_answer("source_introspectable"));
assert!(!ctx.can_answer("raw_docs_readable"));
}
}