feat: type SQL macro calls everywhere, with optional returns, TYPE OF, ANY families and __cte

Cargo.toml

@@ -72,14 +72,14 @@ junction = { version = "2.1.0", default-features = false }
 mimalloc = { version = "0.1", default-features = false, features = ["local_dynamic_tls"] }
 proc-macro2 = { version = "1", features = ["span-locations"] }
 pyo3 = { version = "0.26", features = ["abi3-py312"] }
-polyglot-sql = { git = "https://github.com/chio-labs/polyglot.git", rev = "b18dc647e797e480d2faadf3af407077f8e063df", default-features = false, features = [
+polyglot-sql = { git = "https://github.com/chio-labs/polyglot.git", rev = "1ef5056c2b315e69bd13023bd90ffe60d94880a3", default-features = false, features = [
     "ast-tools",
     "builder",
     "all-dialects",
     "semantic",
     "stacker",
 ] }
-polyglot-sql-function-catalogs = { git = "https://github.com/chio-labs/polyglot.git", rev = "b18dc647e797e480d2faadf3af407077f8e063df", default-features = false, features = [
+polyglot-sql-function-catalogs = { git = "https://github.com/chio-labs/polyglot.git", rev = "1ef5056c2b315e69bd13023bd90ffe60d94880a3", default-features = false, features = [
     "dialect-duckdb",
 ] }
 blake3 = { version = "1", features = ["rayon"] }

crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/analysis_cache.rs

@@ -1,7 +1,5 @@
 //! Per-model analysis cache keys over everything an analysis reads, and stored outcomes.
 
-use crate::sql_macros::main::sql_macro_declared_outputs::sql_macro_declared_outputs;
-use crate::sql_macros::models::DeclaredOutput;
 use std::collections::{BTreeSet, HashMap};
 
 use sha2::{Digest, Sha256};
@@ -171,20 +169,6 @@ impl KeyHasher {
         }
         self.text(&model.pivot_sql);
         self.families(&model.dynamic_families);
-        let declared: Vec<DeclaredOutput> = sql_macro_declared_outputs(&model.sql_macro_checks);
-        self.count(declared.len());
-        for output in &declared {
-            match &output.cte {
-                Some(cte) => {
-                    self.flag(true);
-                    self.text(cte);
-                }
-                None => self.flag(false),
-            }
-            self.text(&output.name);
-            self.text(&output.returns);
-            self.count(output.branch);
-        }
     }
 
     fn families(&mut self, families: &[DynamicFamily]) {

crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/cte_facts.rs

@@ -8,10 +8,10 @@ use polyglot_sql::traversal::ExpressionWalk;
 use serde_json::{Map, Value};
 
 use crate::assembly::analysis_session::_helpers::dict_walk::{
-    PyValue, alias_or_name, arg, children, column_table_name, direct_tables, expression_args,
-    find_all, is_star, kind, name, name_payload, output_name, parse_one, projected_expression,
-    py_truthy, python_int, raw_name_payload, set_operation, to_dict, top_level_ctes,
-    unwrap_annotations, where_clause,
+    Assertion, PyValue, alias_or_name, arg, children, column_table_name, direct_tables,
+    expression_args, find_all, is_star, kind, name, name_payload, output_name, parse_one,
+    projected_expression, py_truthy, python_int, raw_name_payload, set_operation, to_dict,
+    top_level_ctes, typed_annotation, unwrap_annotations, where_clause,
 };
 use crate::assembly::analysis_session::_helpers::mappings::ShapeTable;
 use crate::assembly::analysis_session::constants::{
@@ -23,7 +23,7 @@ use crate::assembly::analysis_session::constants::{
     NULLABLE_NULLABILITY, NULLIF_FUNCTION_NAME, POLYGLOT_TYPE_NAMES, PYTHON_RULE, SAFE_CAST_RULE,
     SELECT_AST_KIND, SET_OPERATION_AST_KINDS, STRING_LITERAL_TYPE, SUBSTRING_AST_KIND,
     TABLE_AST_KIND, TEXT_TYPE, TIMESTAMP_TYPE_NAME, TIMESTAMP_TZ_TYPE, TRY_CAST_AST_KIND,
-    TYPE_PASSTHROUGH_AST_KINDS, UNKNOWN_NULLABILITY, VARCHAR_DATA_TYPES, WILDCARD,
+    TYPE_PASSTHROUGH_AST_KINDS, UNKNOWN_NULLABILITY, UNKNOWN_TYPE, VARCHAR_DATA_TYPES, WILDCARD,
 };
 use crate::assembly::analysis_session::constants::{
     CONFIDENCE_HIGH, CONFIDENCE_MEDIUM, CONFIDENCE_UNKNOWN, TRANSFORM_AGGREGATION, TRANSFORM_CAST,
@@ -35,8 +35,6 @@ use crate::assembly::analysis_session::types::{Pairs, Shapes};
 use crate::nullability_rules::main::conditional_result_nullability::conditional_result_nullability;
 use crate::nullability_rules::main::first_arg_nullability::first_arg_nullability;
 use crate::nullability_rules::main::safe_cast_nullability::safe_cast_nullability;
-use crate::sql_macros::main::sql_macro_declared_type::sql_macro_declared_type;
-use crate::sql_macros::models::DeclaredOutput;
 use crate::type_system::main::normalize_type::normalize_type;
 use crate::type_system::models::{NormalizedType, TypeFamily};
 
@@ -60,8 +58,6 @@ pub(crate) struct RecoveryProfile<'a> {
     pub(crate) rules: Option<&'a Pairs>,
     /// Runs the adapter's own rules, those with the `python` rule id.
     pub(crate) callback: Option<&'a NullabilityCallback>,
-    /// The model's select items that are one SQL macro call each, typed by its `returns`.
-    pub(crate) declared: &'a [DeclaredOutput],
 }
 
 /// One enrichment's recovery request.
@@ -133,108 +129,6 @@ pub(crate) struct LegacyAnalysis {
     pub(crate) lineage: Vec<LineageRow>,
 }
 
-/// The model query's output types its SQL macro call items change, given the relation types.
-pub(crate) fn declared_output_types(
-    cleaned_sql: &str,
-    types: &ShapeTable,
-    profile: &RecoveryProfile<'_>,
-) -> Fact<Vec<(String, String)>> {
-    let Some(parsed) = parse_one(cleaned_sql, profile.dialect)? else {
-        return Ok(Vec::new());
-    };
-    let facts: Relations = referenced_types(&parsed, types);
-    let declared: Vec<(String, String)> = Context {
-        profile,
-        branches: declared_branches(&parsed),
-    }
-    .query_types(&parsed, &facts)?;
-    let plain: Vec<(String, String)> = Context {
-        profile,
-        branches: Vec::new(),
-    }
-    .query_types(&parsed, &facts)?;
-    let main: HashSet<String> = main_output_columns(&parsed, profile.declared);
-    Ok(declared
-        .into_iter()
-        .filter(|(name, data_type)| {
-            dict_get(&plain, name) != Some(data_type) || main.contains(&casefold(name))
-        })
-        .collect())
-}
-
-/// Each select of the main query and top-level CTEs by address, with its scope and branch.
-fn declared_branches(root: &Expression) -> Vec<Branch> {
-    let mut branches: Vec<Branch> = Vec::new();
-    for (cte_name, has_column_aliases, body) in top_level_ctes(root) {
-        if !has_column_aliases
-            && unwrap_annotations(Some(body)).is_some_and(|body| top_level_ctes(body).is_empty())
-        {
-            for (branch, select) in set_branches(Some(body)).into_iter().enumerate() {
-                branches.push((
-                    std::ptr::from_ref(select),
-                    Some(cte_name.to_owned()),
-                    branch,
-                ));
-            }
-        }
-    }
-    for (branch, select) in set_branches(Some(root)).into_iter().enumerate() {
-        branches.push((std::ptr::from_ref(select), None, branch));
-    }
-    branches
-}
-
-/// The folded main-query output columns its SQL macro call items stand in, in any branch.
-fn main_output_columns(root: &Expression, declared: &[DeclaredOutput]) -> HashSet<String> {
-    let selects: Vec<&Expression> = set_branches(Some(root));
-    let by_name: bool = set_operation(root).is_some_and(|(_, _, by_name)| by_name);
-    let Some(first) = selects.first() else {
-        return HashSet::new();
-    };
-    let first_names: Vec<&str> = projection_names(first);
-    let mut columns: HashSet<String> = HashSet::new();
-    for output in declared.iter().filter(|output| output.cte.is_none()) {
-        if by_name {
-            columns.insert(casefold(&output.name));
-            continue;
-        }
-        let Some(select) = selects.get(output.branch) else {
-            continue;
-        };
-        let position: Option<usize> = projection_names(select)
-            .iter()
-            .position(|name| name.eq_ignore_ascii_case(&output.name));
-        if let Some(column) = position.and_then(|position| first_names.get(position)) {
-            columns.insert(casefold(column));
-        }
-    }
-    columns
-}
-
-/// The output name of each of `select`'s projections, in order.
-fn projection_names(select: &Expression) -> Vec<&str> {
-    select
-        .get_expressions()
-        .iter()
-        .map(|projection| output_name(unwrap_annotations(Some(projection))))
-        .collect()
-}
-
-/// The selects of `node`'s set operation, left to right, or `node` itself.
-fn set_branches(node: Option<&Expression>) -> Vec<&Expression> {
-    let Some(node) = unwrap_annotations(node) else {
-        return Vec::new();
-    };
-    match set_operation(node) {
-        Some((left, right, _)) => {
-            let mut branches: Vec<&Expression> = set_branches(Some(left));
-            branches.extend(set_branches(Some(right)));
-            branches
-        }
-        None => vec![node],
-    }
-}
-
 /// Python's legacy analysis of a model the native engine handed back, or why it defers.
 pub(crate) fn legacy_analysis(input: &LegacyInput<'_>) -> Fact<LegacyAnalysis> {
     let Some(parsed) = parse_one(input.cleaned_sql, input.profile.dialect)? else {
@@ -272,7 +166,6 @@ pub(crate) fn legacy_analysis(input: &LegacyInput<'_>) -> Fact<LegacyAnalysis> {
     let unqualified: Option<&(String, String)> = single_resource(&resources);
     let context = Context {
         profile: &input.profile,
-        branches: Vec::new(),
     };
     let (cte_types, cte_nullability) = if input.recover {
         (
@@ -543,7 +436,6 @@ fn passthrough_facts(
     }
     let context = Context {
         profile: &input.profile,
-        branches: Vec::new(),
     };
     recovery.types = context.passthrough_types(root, &input_types(input))?;
     recovery.nullability = context.passthrough_nullability(root, &input_nullability(input))?;
@@ -574,21 +466,24 @@ fn direct_cte_outputs(root: &Expression, cte_names: &HashSet<String>) -> HashSet
     outputs
 }
 
-/// A select's address, its top-level CTE or `None` for the main query, and its branch.
-type Branch = (*const Expression, Option<String>, usize);
-
 /// The profile every resolver reads.
 struct Context<'a, 'b> {
     profile: &'a RecoveryProfile<'b>,
-    /// The selects whose SQL macro call items take their declared types.
-    branches: Vec<Branch>,
 }
 
 impl Context<'_, '_> {
-    /// The model query's output types, through its top-level CTEs.
-    fn query_types(&self, root: &Expression, facts: &Relations) -> Fact<Vec<(String, String)>> {
+    /// Python's `_polyglot_cte_passthrough_types_from_parsed`.
+    fn passthrough_types(
+        &self,
+        root: &Expression,
+        facts: &Relations,
+    ) -> Fact<Vec<(String, String)>> {
+        let ctes: Vec<(&str, bool, &Expression)> = top_level_ctes(root);
+        if ctes.is_empty() {
+            return Ok(Vec::new());
+        }
         let mut relations: Relations = referenced_facts(root, facts);
-        for (cte_name, has_column_aliases, body) in top_level_ctes(root) {
+        for (cte_name, has_column_aliases, body) in ctes {
             if has_column_aliases {
                 continue;
             }
@@ -604,52 +499,6 @@ impl Context<'_, '_> {
         self.direct_select_types(Some(root), &relations)
     }
 
-    /// Python's `_polyglot_cte_passthrough_types_from_parsed`.
-    fn passthrough_types(
-        &self,
-        root: &Expression,
-        facts: &Relations,
-    ) -> Fact<Vec<(String, String)>> {
-        if top_level_ctes(root).is_empty() {
-            return Ok(Vec::new());
-        }
-        self.query_types(root, facts)
-    }
-
-    /// `types` of `select` with each of its SQL macro call items typed by its `returns`.
-    fn declared_types(
-        &self,
-        select: &Expression,
-        mut types: Vec<(String, String)>,
-    ) -> Vec<(String, String)> {
-        let address: *const Expression = std::ptr::from_ref(select);
-        let Some((_, cte, branch)) = self
-            .branches
-            .iter()
-            .find(|(branch_select, _, _)| *branch_select == address)
-        else {
-            return types;
-        };
-        for output in self.profile.declared {
-            let key: String = types
-                .iter()
-                .map(|(name, _)| name)
-                .find(|name| name.eq_ignore_ascii_case(&output.name))
-                .cloned()
-                .unwrap_or_else(|| output.name.clone());
-            let inferred: Option<&str> = dict_get(&types, &key).map(String::as_str);
-            if let Some(declared) = sql_macro_declared_type(
-                self.profile.declared,
-                (cte.as_deref(), *branch, &output.name),
-                inferred,
-                self.profile.dialect,
-            ) {
-                types = dict_set(types, &key, declared);
-            }
-        }
-        types
-    }
-
     /// Python's `_polyglot_select_output_types`.
     fn select_output_types(
         &self,
@@ -840,7 +689,7 @@ impl Context<'_, '_> {
                 inferred = dict_set(inferred, output, inferred_type);
             }
         }
-        Ok(self.declared_types(select, inferred))
+        Ok(inferred)
     }
 
     /// Python's `_polyglot_direct_expression_type`.
@@ -849,6 +698,24 @@ impl Context<'_, '_> {
         expression: Option<&Expression>,
         by_alias: &Relations,
     ) -> Fact<Option<String>> {
+        if let Some((asserted, inner)) = typed_annotation(expression) {
+            if let Some(inferred) = self
+                .expression_type_in(Some(inner), by_alias)?
+                .filter(|inferred| !upper(inferred).contains(UNKNOWN_TYPE))
+            {
+                return Ok(Some(inferred));
+            }
+            return match asserted {
+                Assertion::Fixed(asserted) => {
+                    let target: Value =
+                        serde_json::to_value(asserted).map_err(|error| error.to_string())?;
+                    Ok(target.as_object().and_then(cast_target_type))
+                }
+                Assertion::Like(like) => Ok(self
+                    .expression_type_in(Some(like), by_alias)?
+                    .filter(|like| !upper(like).contains(UNKNOWN_TYPE))),
+            };
+        }
         let expression: Option<&Expression> = unwrap_annotations(expression);
         let node_kind: &str = kind(expression);
         if let Some(inferred) = self.expression_type(expression)? {
@@ -1070,6 +937,7 @@ impl Context<'_, '_> {
         aliases: &[(String, &'static str)],
         scoped: &NullRelations,
     ) -> Fact<&'static str> {
+        let expression: Option<&Expression> = unwrap_annotations(expression);
         let node_kind: &str = kind(expression);
         match node_kind {
             NULL_AST_KIND => return Ok(NULLABLE_NULLABILITY),
@@ -1101,6 +969,7 @@ impl Context<'_, '_> {
 
     /// Python's `_infer_polyglot_shallow_nullability`.
     fn shallow_nullability(&self, expression: Option<&Expression>) -> Fact<&'static str> {
+        let expression: Option<&Expression> = unwrap_annotations(expression);
         let node_kind: &str = kind(expression);
         match node_kind {
             NULL_AST_KIND => return Ok(NULLABLE_NULLABILITY),