feat(native): run pivot proofs, opt-out checks and semantic completion through the native session

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

@@ -1,8 +1,11 @@
 //! Python's `analyze_dynamic_column_contract` over the crate's parse of the model SQL.
 
 use std::collections::{HashMap, HashSet};
+use std::panic::{AssertUnwindSafe, catch_unwind};
 
 use polyglot_sql::{ComplexityGuardOptions, Dialect, Expression, ParseOptions};
+use rayon::ThreadPool;
+use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
 use serde_json::{Map, Value, json};
 use sqlbuild_core::panics::main::catch_compiler_panic::catch_compiler_panic;
 
@@ -17,7 +20,7 @@ use crate::assembly::analysis_session::constants::{
     UNKNOWN_NULLABILITY, UNKNOWN_TYPE,
 };
 use crate::assembly::analysis_session::models::{
-    ColumnFact, ContractProof, DynamicFamily, PivotOutcome,
+    ColumnFact, ContractProof, DynamicFamily, PivotModel, PivotOutcome, PivotTables,
 };
 use crate::assembly::analysis_session::types::{Pairs, Shapes};
 
@@ -33,6 +36,47 @@ pub(crate) struct PivotFacts<'a> {
     pub(crate) families_by_table: &'a [(String, Vec<DynamicFamily>)],
 }
 
+/// Each model's proof in order; a model whose proof panics is deferred to Python.
+pub(crate) fn pivot_outcomes(models: &[PivotModel], tables: &PivotTables) -> Vec<PivotOutcome> {
+    let facts = pivot_facts(tables);
+    models
+        .iter()
+        .map(|model| guarded_outcome(model, &facts))
+        .collect()
+}
+
+/// Each model's proof in order on `pool`; a model whose proof panics is deferred to Python.
+pub(crate) fn pooled_pivot_outcomes(
+    pool: &ThreadPool,
+    models: &[PivotModel],
+    tables: &PivotTables,
+) -> Vec<PivotOutcome> {
+    let facts = pivot_facts(tables);
+    pool.install(|| {
+        models
+            .par_iter()
+            .map(|model| guarded_outcome(model, &facts))
+            .collect()
+    })
+}
+
+fn pivot_facts(tables: &PivotTables) -> PivotFacts<'_> {
+    PivotFacts {
+        dialect: &tables.dialect,
+        column_types: &tables.column_types,
+        authoritative_types: &tables.authoritative_types,
+        column_nullability: &tables.column_nullability,
+        families_by_table: &tables.families_by_table,
+    }
+}
+
+fn guarded_outcome(model: &PivotModel, facts: &PivotFacts<'_>) -> PivotOutcome {
+    catch_unwind(AssertUnwindSafe(|| {
+        pivot_outcome(&model.sql, &model.families, facts)
+    }))
+    .unwrap_or(PivotOutcome::Deferred)
+}
+
 /// One CTE body by folded name; None for a CTE with column aliases.
 type CteMap<'a> = HashMap<String, Option<&'a Value>>;
 

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

@@ -1,7 +1,9 @@
 //! The session's phases: Python's uncached model analysis, wave by wave.
 
 use std::collections::HashMap;
+use std::sync::Arc;
 
+use rayon::ThreadPool;
 use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
 
 use crate::assembly::analysis_session::_helpers::catalog_state::SessionCatalog;
@@ -21,8 +23,9 @@ use crate::assembly::analysis_session::constants::{
     DEFERRAL_ANALYSIS, DEFERRAL_ENRICHMENT, UNKNOWN_NULLABILITY, UNKNOWN_TYPE,
 };
 use crate::assembly::analysis_session::models::{
-    AnalysisSession, Awaiting, ColumnFact, Deferral, DeferredAnalysis, LineageFacts, ModelAnalysis,
-    ModelOutcome, ModelRequest, Phase, PivotOutcome, SessionOutcome, SessionRequest, SessionStep,
+    AnalysisSession, Awaiting, ColumnFact, Deferral, DeferredAnalysis, FinishedSession,
+    LineageFacts, LineageRow, ModelAnalysis, ModelOutcome, ModelRequest, Phase, PivotOutcome,
+    PivotTables, SessionModelFacts, SessionOutcome, SessionRequest, SessionStep,
 };
 use crate::assembly::analysis_session::types::{Pairs, Shapes};
 use crate::semantic_validation::models::ProjectCatalog;
@@ -148,23 +151,43 @@ impl AnalysisSession {
     }
 
     /// Every model's outcome and the binding catalog changes Python records.
-    pub(crate) fn finish(self) -> Result<SessionOutcome, String> {
+    pub(crate) fn finish(self) -> Result<(SessionOutcome, FinishedSession), String> {
         if !matches!(self.phase, Phase::Done) {
             return Err("the session has not finished".to_owned());
         }
         let dynamic_contracts: Vec<PivotOutcome> = self.dynamic_contracts()?;
+        let pool: Result<Arc<ThreadPool>, String> = self.catalog.native.analysis_pool();
         let models: Vec<ModelOutcome> = self
             .outcomes
             .into_iter()
             .collect::<Option<_>>()
             .ok_or("the session left a model unanalysed")?;
         let (schema_additions, analysis_names) = self.catalog.into_changes();
-        Ok(SessionOutcome {
+        let request: SessionRequest = self.request;
+        let mut model_facts: HashMap<String, SessionModelFacts> = HashMap::new();
+        for (model, outcome) in request.models.iter().zip(&models) {
+            if let Some(facts) = session_model_facts(&outcome.analysis) {
+                model_facts.insert(model.name.clone(), facts);
+            }
+        }
+        let finished = FinishedSession {
+            tables: PivotTables {
+                dialect: request.dialect,
+                column_types: request.column_types,
+                authoritative_types: request.complete_schemas,
+                column_nullability: request.column_nullability,
+                families_by_table: request.dynamic_families_by_table,
+            },
+            pool,
+            models: model_facts,
+        };
+        let outcome = SessionOutcome {
             models,
             schema_additions,
             analysis_names,
             dynamic_contracts,
-        })
+        };
+        Ok((outcome, finished))
     }
 
     /// Python's dynamic pivot proof of each model over the relation facts before analysis.
@@ -595,3 +618,28 @@ fn native_answer(model: usize, columns: Vec<ColumnFact>, has_star: bool) -> Defe
         binding_validated: false,
     }
 }
+
+/// A compiled model's output names and lineage from a successful analysis with native lineage.
+fn session_model_facts(analysis: &ModelAnalysis) -> Option<SessionModelFacts> {
+    let rows: &Vec<LineageRow> = match &analysis.lineage {
+        LineageFacts::Native(rows) | LineageFacts::NativeEnrichment(rows) => rows,
+        LineageFacts::PythonAnalysis | LineageFacts::PythonEnrichment => return None,
+    };
+    analysis.analysis_succeeded.then(|| SessionModelFacts {
+        columns: analysis.columns.as_deref().map(column_names),
+        lineage: rows.iter().map(output_sources).collect(),
+    })
+}
+
+fn column_names(columns: &[ColumnFact]) -> Vec<String> {
+    columns.iter().map(|column| column.name.clone()).collect()
+}
+
+/// One lineage row as `(output column, [(resource name, column name)])`.
+fn output_sources(row: &LineageRow) -> (String, Vec<(String, String)>) {
+    let mut sources: Vec<(String, String)> = Vec::with_capacity(row.sources.len());
+    for (_, name, column) in &row.sources {
+        sources.push((name.clone(), column.clone()));
+    }
+    (row.output_column.clone(), sources)
+}

crates/sqlbuild-analysis/src/assembly/analysis_session/main/finish_analysis_session.rs

@@ -1,8 +1,10 @@
 //! Collect a finished native model analysis session's outcomes.
 
-use crate::assembly::analysis_session::models::{AnalysisSession, SessionOutcome};
+use crate::assembly::analysis_session::models::{AnalysisSession, FinishedSession, SessionOutcome};
 
-/// Every model's outcome and the binding catalog changes Python records.
-pub fn finish_analysis_session(session: AnalysisSession) -> Result<SessionOutcome, String> {
+/// Every model's outcome and catalog changes, and the session kept for later pivot proofs.
+pub fn finish_analysis_session(
+    session: AnalysisSession,
+) -> Result<(SessionOutcome, FinishedSession), String> {
     session.finish()
 }

crates/sqlbuild-analysis/src/assembly/analysis_session/main/finished_fact_models.rs

@@ -0,0 +1,8 @@
+//! The models whose output names and lineage a finished analysis session kept.
+
+use crate::assembly::analysis_session::models::FinishedSession;
+
+/// Names of the models whose facts the session kept, in no particular order.
+pub fn finished_fact_models(session: &FinishedSession) -> Vec<String> {
+    session.models.keys().cloned().collect()
+}

crates/sqlbuild-analysis/src/assembly/analysis_session/main/finished_model_facts.rs

@@ -0,0 +1,11 @@
+//! The output names and lineage a finished analysis session kept for compiled models.
+
+use crate::assembly::analysis_session::models::{FinishedSession, SessionModelFacts};
+
+/// The named model's facts, or None where the session kept none for it.
+pub fn finished_model_facts<'a>(
+    session: &'a FinishedSession,
+    name: &str,
+) -> Option<&'a SessionModelFacts> {
+    session.models.get(name)
+}

crates/sqlbuild-analysis/src/assembly/analysis_session/main/mod.rs

@@ -1,6 +1,9 @@
 pub mod expression_shapes;
 pub mod finish_analysis_session;
-pub mod prove_dynamic_contract;
+pub mod finished_fact_models;
+pub mod finished_model_facts;
+pub mod prove_dynamic_contracts;
+pub mod prove_finished_dynamic_contracts;
 pub mod provide_deferred_analyses;
 pub mod run_analysis_session;
 pub mod start_analysis_session;

crates/sqlbuild-analysis/src/assembly/analysis_session/main/prove_dynamic_contracts.rs

@@ -0,0 +1,17 @@
+//! Python's `analyze_dynamic_column_contract` for models no analysis session analysed.
+
+use crate::assembly::analysis_session::_helpers::dynamic_pivot::pivot_outcomes;
+use crate::assembly::analysis_session::constants::PIVOT_WORKER_STACK_BYTES;
+use crate::assembly::analysis_session::models::{PivotBatchRequest, PivotOutcome};
+
+/// Each model's proof in order; a model whose proof panics is deferred to Python.
+pub fn prove_dynamic_contracts(request: &PivotBatchRequest) -> Result<Vec<PivotOutcome>, String> {
+    std::thread::scope(|scope| {
+        std::thread::Builder::new()
+            .stack_size(PIVOT_WORKER_STACK_BYTES)
+            .spawn_scoped(scope, || pivot_outcomes(&request.models, &request.tables))
+            .map_err(|error| error.to_string())?
+            .join()
+            .map_err(|_| "the native dynamic pivot proofs panicked".to_owned())
+    })
+}

crates/sqlbuild-analysis/src/assembly/analysis_session/main/prove_finished_dynamic_contracts.rs

@@ -0,0 +1,13 @@
+//! Python's `analyze_dynamic_column_contract` for models outside a finished analysis session.
+
+use crate::assembly::analysis_session::_helpers::dynamic_pivot::pooled_pivot_outcomes;
+use crate::assembly::analysis_session::models::{FinishedSession, PivotModel, PivotOutcome};
+
+/// Each model's proof in order from the session's own tables; a panicking proof is deferred.
+pub fn prove_finished_dynamic_contracts(
+    session: &FinishedSession,
+    models: &[PivotModel],
+) -> Result<Vec<PivotOutcome>, String> {
+    let pool = session.pool.as_ref().map_err(Clone::clone)?;
+    Ok(pooled_pivot_outcomes(pool, models, &session.tables))
+}

crates/sqlbuild-analysis/src/assembly/analysis_session/models.rs

@@ -1,10 +1,13 @@
 //! Plain-data request, deferrals and outcomes of the native model analysis session.
 
 use std::collections::HashMap;
+use std::sync::Arc;
+
+use rayon::ThreadPool;
 
 use crate::assembly::analysis_session::_helpers::catalog_state::SessionCatalog;
 use crate::assembly::analysis_session::_helpers::mappings::ShapeTable;
-use crate::assembly::analysis_session::types::{Pairs, Shapes};
+use crate::assembly::analysis_session::types::{OutputSources, Pairs, Shapes};
 use crate::semantic_validation::types::DiagnosticRow;
 
 /// One `ref`, `source`, `seed`, `table_fn` or `udf` call a model's SQL makes.
@@ -61,18 +64,45 @@ pub struct ContractProof {
     pub bare_dynamic_pivot: bool,
 }
 
-/// One model's dynamic pivot proof inputs, for a model the session does not analyse.
+/// The relation facts every dynamic pivot proof reads: Python's tables before analysis.
 #[derive(Debug, Clone, Default)]
-pub struct PivotRequest {
+pub struct PivotTables {
     pub dialect: String,
     pub column_types: Shapes,
     pub authoritative_types: Shapes,
     pub column_nullability: Shapes,
     pub families_by_table: Vec<(String, Vec<DynamicFamily>)>,
+}
+
+/// One model's pivot SQL and declared families.
+#[derive(Debug, Clone, Default)]
+pub struct PivotModel {
     pub sql: String,
     pub families: Vec<DynamicFamily>,
 }
 
+/// Dynamic pivot proofs for models no session analysed.
+#[derive(Debug, Clone, Default)]
+pub struct PivotBatchRequest {
+    pub tables: PivotTables,
+    pub models: Vec<PivotModel>,
+}
+
+/// A finished session's pivot tables, analysis pool and model facts, kept for later stages.
+#[derive(Debug)]
+pub struct FinishedSession {
+    pub(crate) tables: PivotTables,
+    pub(crate) pool: Result<Arc<ThreadPool>, String>,
+    pub(crate) models: HashMap<String, SessionModelFacts>,
+}
+
+/// A compiled model's output names and lineage from an analysis the session completed natively.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct SessionModelFacts {
+    pub columns: Option<Vec<String>>,
+    pub lineage: OutputSources,
+}
+
 /// A model's dynamic pivot proof: none declared, left to Python, or proven natively.
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum PivotOutcome {