use crate::parser::{
analyzer::{AnalysisContext, AnalyzerError},
ast::{Column, ScalarExpr},
};
pub struct WildcardResolver;
impl WildcardResolver {
pub fn expand_wildcard(
expr: &ScalarExpr,
ctx: &AnalysisContext,
) -> Result<Vec<ScalarExpr>, AnalyzerError> {
match expr {
ScalarExpr::WildCard => {
let mut result = Vec::new();
for (visible_collection, backing) in &ctx.collections {
if let Some(schema) = ctx.schema_of_collection_ref(backing) {
for (col, _fi) in schema.fields {
result.push(ScalarExpr::Column(Column::WithCollection {
collection: visible_collection.clone(),
name: col,
}));
}
} else {
return Err(AnalyzerError::UnknownCollection(backing.clone()));
}
}
Ok(result)
}
ScalarExpr::WildCardWithCollection(collection) => {
let backing = ctx
.collections
.get(collection)
.ok_or_else(|| AnalyzerError::UnknownCollection(collection.clone()))?;
let schema = ctx
.schema_of_collection_ref(backing)
.ok_or_else(|| AnalyzerError::UnknownCollection(backing.clone()))?;
let mut result = Vec::new();
for (col, _fi) in schema.fields {
result.push(ScalarExpr::Column(Column::WithCollection {
collection: collection.clone(),
name: col,
}));
}
Ok(result)
}
_ => Ok(vec![expr.clone()]),
}
}
pub fn expand_projection(
exprs: &[ScalarExpr],
cx: &AnalysisContext,
) -> Result<Vec<ScalarExpr>, AnalyzerError> {
let mut result = Vec::new();
for e in exprs {
let expanded = Self::expand_wildcard(e, cx)?;
result.extend(expanded);
}
Ok(result)
}
}
#[cfg(test)]
mod tests {
use crate::{
JsonPrimitive, SchemaDict,
database::{FieldInfo, SchemaProvider},
parser::ast::Function,
};
use super::*;
use indexmap::IndexMap;
struct DummySchemas {
by_name: std::collections::HashMap<String, SchemaDict>,
}
impl DummySchemas {
fn new() -> Self {
Self {
by_name: std::collections::HashMap::new(),
}
}
fn with(mut self, name: &str, fields: Vec<(&str, JsonPrimitive, bool)>) -> Self {
let mut m = IndexMap::new();
for (k, ty, nullable) in fields {
m.insert(k.to_string(), FieldInfo { ty, nullable });
}
self.by_name
.insert(name.to_string(), SchemaDict { fields: m });
self
}
}
impl SchemaProvider for DummySchemas {
fn schema_of(&self, backing_collection: &str) -> Option<SchemaDict> {
self.by_name.get(backing_collection).cloned()
}
}
fn ctx_with<'a>(
sp: &'a DummySchemas,
pairs: &'a [(&'a str, Option<&'a str>)],
) -> AnalysisContext<'a> {
let mut ctx = AnalysisContext::new(sp);
for (backing, alias) in pairs {
ctx.add_collection(alias.unwrap_or(backing).to_string(), (*backing).to_string());
}
ctx
}
#[test]
fn star_expands_all_visible_collections_and_all_fields() {
let sp = DummySchemas::new()
.with(
"t1",
vec![
("id", JsonPrimitive::Int, false),
("name", JsonPrimitive::String, false),
],
)
.with("t2", vec![("x", JsonPrimitive::Int, false)]);
let ctx = ctx_with(&sp, &[("t1", None), ("t2", None)]);
let out = WildcardResolver::expand_wildcard(&ScalarExpr::WildCard, &ctx).expect("expand *");
let mut t1_cols = vec![];
let mut t2_cols = vec![];
for e in out {
match e {
ScalarExpr::Column(Column::WithCollection { collection, name })
if collection == "t1" =>
{
t1_cols.push(name)
}
ScalarExpr::Column(Column::WithCollection { collection, name })
if collection == "t2" =>
{
t2_cols.push(name)
}
other => panic!("unexpected expr from wildcard: {other:?}"),
}
}
assert_eq!(
t1_cols,
vec!["id".to_string(), "name".to_string()],
"field order for t1 preserved"
);
assert_eq!(
t2_cols,
vec!["x".to_string()],
"field order for t2 preserved"
);
}
#[test]
fn table_star_expands_only_that_collection() {
let sp = DummySchemas::new()
.with(
"users",
vec![
("id", JsonPrimitive::Int, false),
("name", JsonPrimitive::String, false),
],
)
.with("orders", vec![("total", JsonPrimitive::Float, false)]);
let ctx = ctx_with(&sp, &[("users", Some("u")), ("orders", Some("o"))]);
let out = WildcardResolver::expand_wildcard(
&ScalarExpr::WildCardWithCollection("u".into()),
&ctx,
)
.expect("expand u.*");
let cols: Vec<(String, String)> = out
.into_iter()
.map(|e| match e {
ScalarExpr::Column(Column::WithCollection { collection, name }) => {
(collection, name)
}
other => panic!("expected Column::WithCollection, got {other:?}"),
})
.collect();
assert_eq!(
cols,
vec![("u".into(), "id".into()), ("u".into(), "name".into()),]
);
}
#[test]
fn non_wildcard_expression_is_returned_as_singleton() {
let expr = ScalarExpr::Function(Function {
name: "count".into(),
args: vec![ScalarExpr::WildCard],
distinct: false,
});
let sp = DummySchemas::new();
let ctx = ctx_with(&sp, &[]);
let out = WildcardResolver::expand_wildcard(&expr, &ctx).expect("no-op");
assert_eq!(out.len(), 1);
assert_eq!(out[0], expr);
}
#[test]
fn table_star_unknown_visible_collection_errors() {
let sp = DummySchemas::new().with("t", vec![("a", JsonPrimitive::Int, false)]);
let ctx = ctx_with(&sp, &[("t", None)]);
let err = WildcardResolver::expand_wildcard(
&ScalarExpr::WildCardWithCollection("v".into()),
&ctx,
);
match err {
Err(AnalyzerError::UnknownCollection(c)) => assert_eq!(c, "v"),
other => panic!("expected UnknownCollection(\"v\"), got {other:?}"),
}
}
#[test]
fn star_errors_when_a_visible_collection_has_no_schema() {
let sp = DummySchemas::new(); let ctx = ctx_with(&sp, &[("t", Some("v"))]);
let err = WildcardResolver::expand_wildcard(&ScalarExpr::WildCard, &ctx);
match err {
Err(AnalyzerError::UnknownCollection(b)) => assert_eq!(b, "t"),
other => panic!("expected UnknownCollection(\"t\"), got {other:?}"),
}
}
#[test]
fn table_star_on_empty_schema_returns_empty_vec() {
let sp = {
let m = IndexMap::new();
let empty_schema = SchemaDict { fields: m };
let mut by = std::collections::HashMap::new();
by.insert("empty".to_string(), empty_schema);
DummySchemas { by_name: by }
};
let ctx = ctx_with(&sp, &[("empty", None)]);
let out = WildcardResolver::expand_wildcard(
&ScalarExpr::WildCardWithCollection("empty".into()),
&ctx,
)
.expect("expand empty.*");
assert!(
out.is_empty(),
"expansion over empty schema should be empty"
);
}
#[test]
fn expand_projection_expands_wildcards_and_preserves_non_wildcards() {
let sp = DummySchemas::new().with(
"people",
vec![
("id", JsonPrimitive::Int, false),
("name", JsonPrimitive::String, false),
],
);
let ctx = ctx_with(&sp, &[("people", Some("p"))]);
let exprs = vec![
ScalarExpr::WildCardWithCollection("p".into()),
ScalarExpr::Column(Column::Name {
name: "literal".into(),
}),
];
let expanded =
WildcardResolver::expand_projection(&exprs, &ctx).expect("projection should expand");
assert_eq!(expanded.len(), 3);
assert_eq!(
expanded[0],
ScalarExpr::Column(Column::WithCollection {
collection: "p".into(),
name: "id".into()
})
);
assert_eq!(
expanded[1],
ScalarExpr::Column(Column::WithCollection {
collection: "p".into(),
name: "name".into()
})
);
assert_eq!(expanded[2], exprs[1]);
}
#[test]
fn table_star_errors_when_visible_collection_schema_is_missing() {
let sp = DummySchemas::new();
let ctx = ctx_with(&sp, &[("people", Some("p"))]);
let err = WildcardResolver::expand_wildcard(
&ScalarExpr::WildCardWithCollection("p".into()),
&ctx,
)
.unwrap_err();
assert_eq!(err, AnalyzerError::UnknownCollection("people".into()));
}
}