lix 0.18.0

#![allow(clippy::cloned_ref_to_slice_refs, clippy::match_same_arms)]

use std::collections::BTreeSet;
use std::sync::Arc;

use datafusion::prelude::SessionContext;

use crate::LixError;
use crate::branch::BranchRefReader;

mod branch;
mod change;
mod columns;
mod commit_ancestry;
pub(crate) use commit_ancestry::commit_ancestry_schema;
mod diff;
mod directory;
pub(crate) use diff::relation_diff_schema;
mod file;
mod mainline;
pub(crate) use mainline::relation_history_schema;
#[cfg(test)]
pub(crate) use mainline::take_mainline_work;
pub(crate) fn log_schema() -> datafusion::arrow::datatypes::SchemaRef {
    mainline::metadata_schema(false)
}
mod schema;
mod state_at;
pub(crate) use state_at::relation_state_schema;
#[cfg(test)]
pub(crate) use state_at::{arm_state_at_traversal_probe, take_state_at_traversal_probe};
mod spec;
pub(crate) use spec::{PhysicalScanKey, SpecScanExec, StatementScanKey};
mod upsert;
#[cfg(test)]
pub(crate) use upsert::take_upsert_source_batches;
mod values;

use crate::sql2::catalog::{
    PublicCatalog, PublicSurfaceClass, PublicSurfaceContract, PublicSurfaceKind,
};
use crate::sql2::session::SqlWriteSessionOptions;
use crate::sql2::{SqlExecutionContext, SqlWriteContext};

pub(crate) const READ_TABLE_FUNCTION_NAMES: &[&str] = &[
    "lix_diff",
    "lix_as_of",
    "lix_log",
    "lix_history",
    "lix_commit_ancestry",
];

pub(crate) use directory::execute_exact_lix_directory_root_listing;
pub(crate) use file::{
    ExactLixFileReadColumn, ExactLixFileReadSelector, FastLixFilePathWriteConflict,
    execute_exact_lix_file_batch_read, execute_exact_lix_file_id_manifest_batch_read,
    execute_exact_lix_file_read, execute_exact_lix_file_root_listing,
    execute_exact_lix_file_size_batch_read,
    execute_fast_lix_file_content_update_by_id,
    execute_fast_lix_file_content_update_by_id_with_metadata, execute_fast_lix_file_id_path_writes,
    execute_fast_lix_file_path_writes, execute_fast_lix_file_prepared_path_write,
};
pub(crate) use schema::{
    execute_exact_schema_batch_read, execute_exact_schema_point_read,
    revalidate_schema_amended_rows,
};
pub(crate) use spec::{DmlReturning, SpecWriteTarget, WriteTargetRegistry};
pub(crate) use upsert::{UpsertAction, excluded_field_name};

pub(crate) async fn register_read<C>(
    session: &SessionContext,
    ctx: &C,
    branch_ref: Arc<dyn BranchRefReader>,
    active_branch_commit_id: Option<String>,
    selection: &ProviderSelection,
) -> Result<(), LixError>
where
    C: SqlExecutionContext + ?Sized,
{
    let catalog = if selection.requires_visible_schemas() {
        ctx.public_catalog().await?
    } else {
        Arc::clone(PublicCatalog::fixed_system_shared())
    };
    crate::sql2::udfs::register_row_ref_function(session, Arc::clone(&catalog));
    register_read_table_functions(
        session,
        ctx,
        active_branch_commit_id.clone(),
        Arc::clone(&catalog),
        selection,
    )?;
    register_read_from_catalog(
        session,
        ctx,
        branch_ref,
        &catalog,
        ReadProviderScope::All,
        selection,
    )?;
    register_information_schema(session, selection, catalog)
}

/// Registers just the selected read-only relations needed while planning a
/// write. Writable surfaces are installed separately from the transaction's
/// write overlay.
pub(crate) fn register_write_read_relations<C>(
    session: &SessionContext,
    ctx: &C,
    catalog: Arc<PublicCatalog>,
    relation_names: BTreeSet<String>,
) -> Result<(), LixError>
where
    C: SqlExecutionContext + ?Sized,
{
    let selection = ProviderSelection::Only {
        names: relation_names,
        history_relations: BTreeSet::new(),
    };
    let branch_ref = Arc::new(super::branch_ref::CachingBranchRefReader::new(ctx.branch_ref()));
    register_read_from_catalog(
        session,
        ctx,
        branch_ref,
        &catalog,
        ReadProviderScope::ReadOnly,
        &selection,
    )
}

/// Narrows DataFusion's relation references to read-only catalog surfaces.
/// `None` means DataFusion could not resolve the reference list, so retain the
/// conservative provider-selection behavior.
pub(crate) fn write_read_relation_selection(
    catalog: &PublicCatalog,
    selection: &ProviderSelection,
    relation_names: Option<&BTreeSet<String>>,
) -> BTreeSet<String> {
    catalog
        .surfaces()
        .filter(|surface| {
            matches!(surface.class, PublicSurfaceClass::Relation(_))
                && !is_write_surface(surface)
                && selection.includes(surface)
                && relation_names.is_none_or(|names| names.contains(&surface.name))
        })
        .map(|surface| surface.name.clone())
        .collect()
}

/// Installs the `information_schema` views only for statements that can reach
/// them.
///
/// `read_provider_selection` widens to [`ProviderSelection::All`] for every
/// `information_schema`-qualified reference and every `SHOW` form, so a narrowed
/// selection provably never resolves an information-schema table. Registering
/// the schema anyway cost one catalog write lock and one `SchemaProvider`
/// allocation on every ordinary statement.
fn register_information_schema(
    session: &SessionContext,
    selection: &ProviderSelection,
    catalog: Arc<PublicCatalog>,
) -> Result<(), LixError> {
    if !matches!(
        selection,
        ProviderSelection::All | ProviderSelection::AllWithHistory(_)
    ) {
        return Ok(());
    }
    crate::sql2::information_schema::register(session, catalog)
}

/// Snapshot-local providers needed to plan already-bound SQL.
///
/// Ordinary reads use DataFusion's resolver, including its CTE scoping and
/// identifier normalization rules. A narrow AST walk additionally extracts the
/// plan-time relation literal from `lix_history(...)`; provider construction
/// cannot depend on a runtime parameter because each relation has a different
/// result schema. Bound target-only writes select their known target directly.
/// Providers and plans remain scoped to the current storage snapshot.
#[derive(Clone, Debug, PartialEq, Eq)]
pub(crate) enum ProviderSelection {
    /// Register every surface when catalog-wide visibility is part of the SQL
    /// semantics (notably `information_schema` and rewritten `SHOW` queries),
    /// or when reference resolution cannot prove a narrower set is sufficient.
    All,
    /// Register every ordinary surface while constructing history providers
    /// only for relation literals that occur in the statement.
    AllWithHistory(BTreeSet<String>),
    /// Register the union of concrete table names referenced by the statements.
    Only {
        names: BTreeSet<String>,
        history_relations: BTreeSet<String>,
    },
    /// Register only the concrete table names referenced by the statements,
    /// while loading the complete visible catalog for runtime metadata.
    OnlyWithVisibleSchemas {
        names: BTreeSet<String>,
        history_relations: BTreeSet<String>,
    },
}

impl ProviderSelection {
    fn includes(&self, surface: &PublicSurfaceContract) -> bool {
        match self {
            Self::All | Self::AllWithHistory(_) => true,
            Self::Only { names, .. } | Self::OnlyWithVisibleSchemas { names, .. } => {
                names.contains(&surface.name)
            }
        }
    }

    fn requested_history_relations(&self) -> Option<&BTreeSet<String>> {
        match self {
            Self::All => Some(empty_history_relations()),
            Self::AllWithHistory(history_relations) => Some(history_relations),
            Self::Only {
                history_relations, ..
            }
            | Self::OnlyWithVisibleSchemas {
                history_relations, ..
            } => Some(history_relations),
        }
    }

    /// Whether resolving this selection requires the storage-backed catalog.
    ///
    /// Table-free reads and references satisfied by the immutable system catalog
    /// can install providers without scanning `lix_registered_schema` rows.
    /// Runtime registration rejects schema keys whose generated table names
    /// would shadow these fixed providers.
    /// `All`, `OnlyWithVisibleSchemas`, and every unknown name load the full
    /// visible catalog so information-schema, custom rows, and normal
    /// unknown-table errors keep their current semantics.
    fn requires_visible_schemas(&self) -> bool {
        match self {
            Self::All | Self::AllWithHistory(_) | Self::OnlyWithVisibleSchemas { .. } => true,
            Self::Only {
                names,
                history_relations,
            } => {
                names
                    .iter()
                    .any(|name| PublicCatalog::fixed_system().surface(name).is_none())
                    || history_relations.iter().any(|name| {
                        PublicCatalog::fixed_system()
                            .history_relation(name)
                            .is_none()
                    })
            }
        }
    }
}

/// Install read-only table functions shared by ordinary reads, transaction
/// reads, and INSERT query sources. Relation providers are registered
/// separately so write sessions can keep their transaction-overlay sources.
pub(crate) fn register_read_table_functions<C>(
    session: &SessionContext,
    ctx: &C,
    active_branch_commit_id: Option<String>,
    catalog: Arc<PublicCatalog>,
    selection: &ProviderSelection,
) -> Result<(), LixError>
where
    C: SqlExecutionContext + ?Sized,
{
    if catalog
        .surface("lix_diff")
        .is_some_and(|surface| selection.includes(surface))
    {
        ctx.note_unvalidated_read("lix_diff");
        diff::register_diff_function(
            session,
            ctx.changelog_query_source(),
            Arc::clone(&catalog),
            ctx.read_interest_registry(),
            ctx.blob_reader(),
            ctx.filesystem_path_index().historical_cache(),
        );
    }
    if catalog
        .surface("lix_as_of")
        .is_some_and(|surface| selection.includes(surface))
    {
        ctx.note_unvalidated_read("lix_as_of");
        state_at::register_state_at_function(
            session,
            ctx.changelog_query_source(),
            Arc::clone(&catalog),
            ctx.active_branch_id().to_string(),
            ctx.blob_reader(),
        );
    }
    if ["lix_log", "lix_history"].iter().any(|name| {
        catalog
            .surface(name)
            .is_some_and(|surface| selection.includes(surface))
    }) {
        ctx.note_unvalidated_read("lix_history");
        mainline::register_functions(
            session,
            ctx.changelog_query_source(),
            Arc::clone(&catalog),
            ctx.blob_reader(),
        );
    }
    if let Some(surface) = catalog
        .surface("lix_commit_ancestry")
        .filter(|surface| selection.includes(surface))
    {
        let active_branch_commit_id = active_branch_commit_id.ok_or_else(|| {
            LixError::branch_not_found(
                ctx.active_branch_id(),
                "register lix_commit_ancestry",
                "active branch",
            )
        })?;
        ctx.note_unvalidated_read(&surface.name);
        commit_ancestry::register_commit_ancestry_function(
            session,
            &surface.name,
            active_branch_commit_id,
            ctx.commit_graph(),
        );
    }
    Ok(())
}

pub(crate) fn read_provider_selection(
    state: &datafusion::execution::session_state::SessionState,
    statements: &[datafusion::sql::parser::Statement],
) -> ProviderSelection {
    let mut names = BTreeSet::new();
    let mut history_relations = BTreeSet::new();
    let mut requires_all = false;
    let mut requires_visible_schemas = false;
    // Resolving references only reads the SQL parser configuration, so the
    // statement's pooled session state is used directly instead of cloning the
    // live one.
    for statement in statements {
        collect_history_relation_literals(statement, &mut history_relations);
        collect_dynamic_relation_literals(statement, &mut names, &mut requires_visible_schemas);
        if statement_requires_all_providers(statement) {
            requires_all = true;
            continue;
        }
        let Ok(references) = state.resolve_table_references(statement) else {
            requires_all = true;
            continue;
        };
        for reference in references {
            if reference.schema() == Some("information_schema") {
                requires_all = true;
            }
            names.insert(reference.table().to_string());
        }
    }
    if requires_all {
        return all_provider_selection(history_relations);
    }
    if requires_visible_schemas {
        return ProviderSelection::OnlyWithVisibleSchemas {
            names,
            history_relations,
        };
    }
    ProviderSelection::Only {
        names,
        history_relations,
    }
}

fn all_provider_selection(history_relations: BTreeSet<String>) -> ProviderSelection {
    if history_relations.is_empty() {
        ProviderSelection::All
    } else {
        ProviderSelection::AllWithHistory(history_relations)
    }
}

fn empty_history_relations() -> &'static BTreeSet<String> {
    static EMPTY: std::sync::OnceLock<BTreeSet<String>> = std::sync::OnceLock::new();
    EMPTY.get_or_init(BTreeSet::new)
}

fn collect_history_relation_literals(
    statement: &datafusion::sql::parser::Statement,
    relations: &mut BTreeSet<String>,
) {
    use std::ops::ControlFlow;

    use datafusion::sql::parser::Statement as DataFusionStatement;
    use datafusion::sql::sqlparser::ast::{
        Expr as SqlExpr, FunctionArg, FunctionArgExpr, TableFactor, Value as SqlValue, Visit,
        Visitor,
    };

    struct HistoryRelationVisitor<'a>(&'a mut BTreeSet<String>);

    impl Visitor for HistoryRelationVisitor<'_> {
        type Break = ();

        fn pre_visit_table_factor(
            &mut self,
            table_factor: &TableFactor,
        ) -> ControlFlow<Self::Break> {
            let TableFactor::Table {
                name,
                args: Some(arguments),
                ..
            } = table_factor
            else {
                return ControlFlow::Continue(());
            };
            if !crate::sql2::parse::object_name_is_public_function(name, "lix_history") {
                return ControlFlow::Continue(());
            }
            let Some(FunctionArg::Unnamed(FunctionArgExpr::Expr(SqlExpr::Value(value)))) =
                arguments.args.first()
            else {
                return ControlFlow::Continue(());
            };
            if let SqlValue::SingleQuotedString(relation_name) = &value.value {
                self.0.insert(relation_name.clone());
            }
            ControlFlow::Continue(())
        }
    }

    match statement {
        DataFusionStatement::Statement(statement) => {
            let _ = statement.visit(&mut HistoryRelationVisitor(relations));
        }
        DataFusionStatement::Explain(explain) => {
            collect_history_relation_literals(explain.statement.as_ref(), relations);
        }
        _ => {}
    }
}

/// Runtime relation arguments must participate in snapshot-local catalog
/// selection even though DataFusion only resolves the function names
/// themselves. A literal relation can select one surface; a dynamic
/// `lix_row_ref` relation requires the complete visible catalog.
fn collect_dynamic_relation_literals(
    statement: &datafusion::sql::parser::Statement,
    relations: &mut BTreeSet<String>,
    requires_visible_schemas: &mut bool,
) {
    use std::ops::ControlFlow;

    use datafusion::sql::parser::Statement as DataFusionStatement;
    use datafusion::sql::sqlparser::ast::{
        Expr as SqlExpr, FunctionArg, FunctionArgExpr, TableFactor, Value as SqlValue, Visit,
        Visitor,
    };

    struct DiffRelationVisitor<'a> {
        relations: &'a mut BTreeSet<String>,
        requires_visible_schemas: &'a mut bool,
    }

    impl Visitor for DiffRelationVisitor<'_> {
        type Break = ();

        fn pre_visit_expr(&mut self, expression: &SqlExpr) -> ControlFlow<Self::Break> {
            let SqlExpr::Function(function) = expression else {
                return ControlFlow::Continue(());
            };
            let datafusion::sql::sqlparser::ast::FunctionArguments::List(arguments) =
                &function.args
            else {
                return ControlFlow::Continue(());
            };
            let Some(first_argument) = arguments.args.first() else {
                return ControlFlow::Continue(());
            };
            if !crate::sql2::parse::object_name_is_public_function(&function.name, "lix_row_ref")
            {
                return ControlFlow::Continue(());
            }
            match first_argument {
                FunctionArg::Unnamed(FunctionArgExpr::Expr(SqlExpr::Value(value))) => {
                    if let SqlValue::SingleQuotedString(relation_name) = &value.value {
                        self.relations.insert(relation_name.clone());
                    } else {
                        // A parameter or another expression can name a
                        // runtime-registered relation. Load the complete
                        // visible catalog before DataFusion plans the UDF.
                        *self.requires_visible_schemas = true;
                    }
                }
                _ => *self.requires_visible_schemas = true,
            }
            ControlFlow::Continue(())
        }

        fn pre_visit_table_factor(
            &mut self,
            table_factor: &TableFactor,
        ) -> ControlFlow<Self::Break> {
            let TableFactor::Table {
                name,
                args: Some(arguments),
                ..
            } = table_factor
            else {
                return ControlFlow::Continue(());
            };
            let is_history =
                crate::sql2::parse::object_name_is_public_function(name, "lix_history");
            if !crate::sql2::parse::object_name_is_public_function(name, "lix_diff")
                && !crate::sql2::parse::object_name_is_public_function(name, "lix_as_of")
                && !is_history
            {
                return ControlFlow::Continue(());
            }
            let Some(FunctionArg::Unnamed(FunctionArgExpr::Expr(SqlExpr::Value(value)))) =
                arguments.args.first()
            else {
                *self.requires_visible_schemas = true;
                return ControlFlow::Continue(());
            };
            if let SqlValue::SingleQuotedString(relation_name) = &value.value {
                // History functions consume their relation name through the
                // separate history provider selection. Only runtime relation
                // schemas need to load the visible catalog here.
                if !is_history
                    || PublicCatalog::fixed_system()
                        .history_relation(relation_name)
                        .is_none()
                {
                    self.relations.insert(relation_name.clone());
                }
            } else {
                // The table function's schema is fixed during DataFusion
                // planning. Bound relation names are substituted before
                // planning, but provider selection runs before that bind, so
                // a non-literal relation must load the visible schema catalog.
                *self.requires_visible_schemas = true;
            }
            ControlFlow::Continue(())
        }
    }

    match statement {
        DataFusionStatement::Statement(statement) => {
            let _ = statement.visit(&mut DiffRelationVisitor {
                relations,
                requires_visible_schemas,
            });
        }
        DataFusionStatement::Explain(explain) => {
            collect_dynamic_relation_literals(
                explain.statement.as_ref(),
                relations,
                requires_visible_schemas,
            );
        }
        _ => {}
    }
}

fn statement_requires_all_providers(statement: &datafusion::sql::parser::Statement) -> bool {
    use datafusion::sql::parser::Statement as DataFusionStatement;
    use datafusion::sql::sqlparser::ast::Statement as SqlStatement;

    fn sql_statement_requires_all_providers(statement: &SqlStatement) -> bool {
        match statement {
            SqlStatement::ShowFunctions { .. }
            | SqlStatement::ShowVariable { .. }
            | SqlStatement::ShowStatus { .. }
            | SqlStatement::ShowVariables { .. }
            | SqlStatement::ShowCreate { .. }
            | SqlStatement::ShowColumns { .. }
            | SqlStatement::ShowTables { .. }
            | SqlStatement::ShowCollation { .. } => true,
            SqlStatement::Explain { statement, .. } => {
                sql_statement_requires_all_providers(statement)
            }
            _ => false,
        }
    }

    match statement {
        DataFusionStatement::Statement(statement) => {
            sql_statement_requires_all_providers(statement)
        }
        DataFusionStatement::Explain(explain) => {
            statement_requires_all_providers(explain.statement.as_ref())
        }
        _ => false,
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum ReadProviderScope {
    All,
    ReadOnly,
}

impl ReadProviderScope {
    fn includes(self, surface: &PublicSurfaceContract) -> bool {
        self == Self::All || !is_write_surface(surface)
    }
}

fn is_write_surface(surface: &PublicSurfaceContract) -> bool {
    surface.capabilities.insert || surface.capabilities.update || surface.capabilities.delete
}

fn register_read_from_catalog<C>(
    session: &SessionContext,
    ctx: &C,
    branch_ref: Arc<dyn BranchRefReader>,
    catalog: &Arc<PublicCatalog>,
    scope: ReadProviderScope,
    selection: &ProviderSelection,
) -> Result<(), LixError>
where
    C: SqlExecutionContext + ?Sized,
{
    if let Some(requested) = selection.requested_history_relations() {
        for relation_name in requested {
            if catalog.history_relation(relation_name).is_none() {
                return Err(LixError::new(
                    LixError::CODE_UNSUPPORTED_SQL,
                    format!("lix_history does not support relation '{relation_name}'"),
                ));
            }
        }
    }
    for surface in catalog.surfaces() {
        if !scope.includes(surface) || !selection.includes(surface) {
            continue;
        }
        match &surface.kind {
            PublicSurfaceKind::Branch => {
                ctx.note_unvalidated_read(&surface.name);
                branch::register_lix_branch_read_provider(
                    session,
                    &surface.name,
                    ctx.hot_state(),
                    Arc::clone(&branch_ref),
                )?;
            }
            PublicSurfaceKind::Change => {
                ctx.note_unvalidated_read(&surface.name);
                change::register_lix_change_read_provider(
                    session,
                    &surface.name,
                    ctx.changelog_query_source(),
                    ctx.hot_state().is_partial_replica(),
                )?;
            }
            PublicSurfaceKind::CommitAncestryFunction => {}
            PublicSurfaceKind::File => {
                file::register_lix_file_active_provider(
                    session,
                    &surface.name,
                    ctx.active_branch_id(),
                    ctx.hot_state(),
                    ctx.filesystem_path_index(),
                    Arc::clone(&branch_ref),
                    ctx.blob_reader(),
                    ctx.plugin_host(),
                    ctx.functions(),
                    ctx.session_file_views(),
                )?;
            }
            PublicSurfaceKind::Directory => {
                directory::register_lix_directory_active_provider(
                    session,
                    &surface.name,
                    ctx.active_branch_id(),
                    ctx.hot_state(),
                    ctx.filesystem_path_index(),
                    Arc::clone(&branch_ref),
                    ctx.functions(),
                )?;
            }
            PublicSurfaceKind::SchemaBase { .. }
            | PublicSurfaceKind::LogFunction
            | PublicSurfaceKind::HistoryFunction
            | PublicSurfaceKind::DiffFunction
            | PublicSurfaceKind::CheckpointFunction
            | PublicSurfaceKind::RecoveryFunction
            | PublicSurfaceKind::StateAtFunction => {}
        }
    }
    schema::register_row_providers(
        session,
        ctx.active_branch_id(),
        ctx.hot_state(),
        ctx.row_snapshot_reader(),
        Arc::clone(&branch_ref),
        catalog,
        scope == ReadProviderScope::All,
        selection,
    )?;

    Ok(())
}

pub(crate) async fn register_write(
    session: &SessionContext,
    write_ctx: SqlWriteContext,
    branch_ref: Arc<dyn BranchRefReader>,
    options: SqlWriteSessionOptions,
    catalog: Arc<PublicCatalog>,
    selection: &ProviderSelection,
) -> Result<(), LixError> {
    crate::sql2::udfs::register_row_ref_function(session, Arc::clone(&catalog));
    register_write_from_catalog(session, write_ctx, branch_ref, options, &catalog, selection)
        .await?;
    register_information_schema(session, selection, catalog)
}

pub(crate) async fn register_transaction<C>(
    session: &SessionContext,
    read_ctx: &C,
    read_branch_ref: Arc<dyn BranchRefReader>,
    active_branch_commit_id: Option<String>,
    write_ctx: SqlWriteContext,
    write_branch_ref: Arc<dyn BranchRefReader>,
    options: SqlWriteSessionOptions,
    selection: &ProviderSelection,
) -> Result<(), LixError>
where
    C: SqlExecutionContext + ?Sized,
{
    // Both capabilities project the same transaction-scoped schema snapshot.
    // Reuse that immutable metadata, then install read-only providers from the
    // committed read capability and writable providers from the overlay.
    let catalog = write_ctx.public_catalog()?;
    crate::sql2::udfs::register_row_ref_function(session, Arc::clone(&catalog));
    register_read_table_functions(
        session,
        read_ctx,
        active_branch_commit_id.clone(),
        Arc::clone(&catalog),
        selection,
    )?;
    register_read_from_catalog(
        session,
        read_ctx,
        read_branch_ref,
        &catalog,
        ReadProviderScope::ReadOnly,
        selection,
    )?;
    // Writable branch surfaces are served by the write providers below, but a
    // read of them still observes branch heads without row-level validation.
    for surface in catalog.surfaces() {
        if selection.includes(surface) && matches!(surface.kind, PublicSurfaceKind::Branch) {
            read_ctx.note_unvalidated_read(&surface.name);
        }
    }
    register_write_from_catalog(
        session,
        write_ctx,
        write_branch_ref,
        options,
        &catalog,
        selection,
    )
    .await?;
    register_information_schema(session, selection, catalog)
}

async fn register_write_from_catalog(
    session: &SessionContext,
    write_ctx: SqlWriteContext,
    branch_ref: Arc<dyn BranchRefReader>,
    options: SqlWriteSessionOptions,
    catalog: &PublicCatalog,
    selection: &ProviderSelection,
) -> Result<(), LixError> {
    for surface in catalog.surfaces() {
        if !selection.includes(surface) {
            continue;
        }
        match &surface.kind {
            PublicSurfaceKind::Branch => {
                branch::register_write_provider(
                    session,
                    &surface.name,
                    write_ctx.clone(),
                    Arc::clone(&branch_ref),
                )
                .await?;
            }
            PublicSurfaceKind::File => {
                file::register_active_write_provider(
                    session,
                    &surface.name,
                    write_ctx.clone(),
                    Arc::clone(&branch_ref),
                    options.clone(),
                )
                .await?;
            }
            PublicSurfaceKind::Directory => {
                directory::register_active_write_provider(
                    session,
                    &surface.name,
                    write_ctx.clone(),
                    Arc::clone(&branch_ref),
                )
                .await?;
            }
            PublicSurfaceKind::Change
            | PublicSurfaceKind::LogFunction
            | PublicSurfaceKind::HistoryFunction
            | PublicSurfaceKind::DiffFunction
            | PublicSurfaceKind::CheckpointFunction
            | PublicSurfaceKind::RecoveryFunction
            | PublicSurfaceKind::StateAtFunction
            | PublicSurfaceKind::CommitAncestryFunction => {}
            PublicSurfaceKind::SchemaBase { .. } => {}
        }
    }
    schema::register_row_write_providers(session, write_ctx, branch_ref, catalog, selection)
        .await?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeSet;
    use std::sync::Arc;

    use async_trait::async_trait;
    use serde_json::json;

    use datafusion::arrow::datatypes::{DataType, SchemaRef};
    use datafusion::prelude::SessionContext;

    use crate::LixError;
    use crate::branch::{BranchHead, BranchRefReader};
    use crate::changelog::CommitId;
    use crate::hot_state::{HotStateReader, HotStateScanRequest};
    use crate::sql2::catalog::{PublicCatalog, derive_schema_surface_spec_from_schema};

    use super::{
        ProviderSelection, ReadProviderScope, branch, change, directory, file, is_write_surface,
        read_provider_selection, relation_history_schema, schema,
    };

    fn selection_for_sql(sql: &[&str]) -> ProviderSelection {
        let statements = sql
            .iter()
            .map(|sql| crate::sql2::parse_statement(sql).expect("SQL should parse"))
            .collect::<Vec<_>>();
        read_provider_selection(&SessionContext::new().state(), &statements)
    }

    fn selected_names(names: &[&str]) -> ProviderSelection {
        ProviderSelection::Only {
            names: names.iter().map(|name| (*name).to_string()).collect(),
            history_relations: BTreeSet::new(),
        }
    }

    fn selected_names_with_visible_schemas(names: &[&str]) -> ProviderSelection {
        ProviderSelection::OnlyWithVisibleSchemas {
            names: names.iter().map(|name| (*name).to_string()).collect(),
            history_relations: BTreeSet::new(),
        }
    }

    #[test]
    fn referenced_provider_selection_uses_datafusion_cte_and_set_operation_resolution() {
        let selection = selection_for_sql(&["WITH shadowed AS (\
                 SELECT id FROM lix_key_value \
                 WHERE EXISTS (SELECT 1 FROM lix_file)\
             ) \
             SELECT left_side.id \
             FROM shadowed AS left_side \
             JOIN (\
                 SELECT row_pk FROM lix_change \
                 UNION ALL \
                 SELECT row_pk FROM lix_change\
               ) AS right_side \
               ON false \
             JOIN public.\"lix_directory\" AS directory_a ON true \
             JOIN public.\"lix_directory\" AS directory_b ON true"]);

        assert_eq!(
            selection,
            selected_names(&["lix_change", "lix_directory", "lix_file", "lix_key_value"])
        );
    }

    #[test]
    fn referenced_provider_selection_excludes_shadowed_and_recursive_cte_names() {
        assert_eq!(
            selection_for_sql(&["WITH lix_file AS (SELECT id FROM lix_key_value) \
                 SELECT * FROM lix_file",]),
            selected_names(&["lix_key_value"])
        );
        assert_eq!(
            selection_for_sql(&["WITH RECURSIVE walk(id) AS (\
                     SELECT id FROM lix_branch \
                     UNION ALL \
                     SELECT branch.id FROM lix_branch AS branch \
                     JOIN walk ON branch.id = walk.id\
                 ) \
                 SELECT * FROM walk",]),
            selected_names(&["lix_branch"])
        );
    }

    #[test]
    fn referenced_provider_selection_unions_batches_and_preserves_unknown_names() {
        assert_eq!(
            selection_for_sql(&[
                "SELECT * FROM lix_file",
                "SELECT * FROM public.lix_key_value JOIN \"UnknownTable\" ON true",
            ]),
            selected_names(&["UnknownTable", "lix_file", "lix_key_value"])
        );
    }

    #[test]
    fn referenced_provider_selection_registers_none_for_table_free_queries() {
        assert_eq!(
            selection_for_sql(&["SELECT 1, uuidv7()"]),
            ProviderSelection::Only {
                names: BTreeSet::new(),
                history_relations: BTreeSet::new(),
            }
        );
    }

    #[test]
    fn history_provider_selection_keeps_the_literal_relation_separate_from_the_function() {
        assert_eq!(
            selection_for_sql(&["SELECT id FROM lix_history('lix_file', $1)"]),
            ProviderSelection::Only {
                names: BTreeSet::from(["lix_history".to_string()]),
                history_relations: BTreeSet::from(["lix_file".to_string()]),
            }
        );
        assert!(
            selection_for_sql(&["SELECT * FROM lix_history('runtime_note')"])
                .requires_visible_schemas()
        );
        assert_eq!(
            selection_for_sql(&[
                "SELECT * FROM information_schema.tables",
                "SELECT * FROM lix_history('lix_file')",
            ]),
            ProviderSelection::AllWithHistory(BTreeSet::from(["lix_file".to_string()]))
        );

        for sql in [
            "SELECT * FROM LIX_HISTORY('lix_file')",
            "SELECT * FROM public.lix_history('lix_file')",
            "SELECT * FROM datafusion.public.lix_history('lix_file')",
            "SELECT * FROM \"lix_history\"('lix_file')",
        ] {
            let selection = selection_for_sql(&[sql]);
            assert!(
                selection
                    .requested_history_relations()
                    .is_some_and(|relations| relations.contains("lix_file")),
                "{sql} should select the literal history provider: {selection:?}",
            );
        }

        assert!(
            !selection_for_sql(&["SELECT * FROM \"LIX_HISTORY\"('lix_file')"])
                .requested_history_relations()
                .is_some_and(|relations| relations.contains("lix_file")),
            "quoted identifiers retain their case",
        );
        assert!(
            !selection_for_sql(&["SELECT * FROM \"PUBLIC\".lix_history('lix_file')"])
                .requested_history_relations()
                .is_some_and(|relations| relations.contains("lix_file")),
            "quoted schema identifiers retain their case",
        );
    }

    #[test]
    fn diff_provider_selection_loads_runtime_relation_schema_for_dynamic_side_columns() {
        assert_eq!(
            selection_for_sql(&["SELECT to_value FROM lix_diff('lix_key_value', $1, $2)"]),
            selected_names(&["lix_diff", "lix_key_value"]),
        );
        assert!(
            selection_for_sql(&["SELECT * FROM lix_diff('runtime_note', $1, $2)"])
                .requires_visible_schemas()
        );
        assert_eq!(
            selection_for_sql(&["SELECT * FROM lix_diff($1, $2, $3)"]),
            selected_names_with_visible_schemas(&["lix_diff"]),
            "a bound relation name must load the visible catalog before planning",
        );
        assert_eq!(
            selection_for_sql(&["SELECT * FROM lix_as_of($1, $2)"]),
            selected_names_with_visible_schemas(&["lix_as_of"]),
        );
        assert_eq!(
            selection_for_sql(&["SELECT * FROM lix_history($1)"]),
            selected_names_with_visible_schemas(&["lix_history"]),
        );
    }

    #[test]
    fn row_ref_provider_selection_loads_visible_catalog_for_dynamic_relation() {
        assert_eq!(
            selection_for_sql(&["SELECT lix_row_ref($1, NULL, $2)"]),
            selected_names_with_visible_schemas(&[]),
        );
        assert_eq!(
            selection_for_sql(&["SELECT lix_row_ref(CAST($1 AS TEXT), NULL, $2)"]),
            selected_names_with_visible_schemas(&[]),
        );
        assert_eq!(
            selection_for_sql(&["SELECT lix_row_ref('lix_file', NULL, $1)"]),
            selected_names(&["lix_file"]),
        );
        let selection = selection_for_sql(&["SELECT lix_row_ref($1, NULL, $2) FROM lix_file"]);
        assert_eq!(selection, selected_names_with_visible_schemas(&["lix_file"]));
        assert!(selection.requires_visible_schemas());
        assert!(selection.includes(PublicCatalog::fixed_system().surface("lix_file").unwrap()));
        assert!(!selection.includes(
            PublicCatalog::fixed_system()
                .surface("lix_directory")
                .unwrap()
        ));
    }

    #[test]
    fn referenced_provider_selection_keeps_catalog_wide_information_schema_semantics() {
        assert_eq!(
            selection_for_sql(&["SELECT * FROM information_schema.tables"]),
            ProviderSelection::All
        );
        assert_eq!(selection_for_sql(&["SHOW TABLES"]), ProviderSelection::All);
    }

    #[test]
    fn visible_schema_loading_boundary_is_conservative() {
        assert!(!selection_for_sql(&["SELECT 1"]).requires_visible_schemas());
        assert!(!selection_for_sql(&["SELECT * FROM lix_key_value"]).requires_visible_schemas());
        assert!(
            !selection_for_sql(&["SELECT * FROM lix_history('lix_key_value')"])
                .requires_visible_schemas()
        );
        assert!(
            !selection_for_sql(&["SELECT * FROM lix_key_value JOIN lix_file ON false"])
                .requires_visible_schemas()
        );
        assert!(selection_for_sql(&["SELECT * FROM custom_row"]).requires_visible_schemas());
        assert!(
            selection_for_sql(&["SELECT * FROM lix_key_value JOIN custom_row ON false",])
                .requires_visible_schemas()
        );
        assert!(
            selection_for_sql(&["SELECT * FROM information_schema.tables"])
                .requires_visible_schemas()
        );
    }

    #[test]
    fn referenced_provider_selection_filters_transaction_capabilities_symmetrically() {
        let catalog = PublicCatalog::from_visible_schemas(&[]).expect("catalog should build");
        let selection = selection_for_sql(&[
            "SELECT * FROM lix_file JOIN lix_history('lix_file') AS history ON false",
        ]);

        let committed_read_names = catalog
            .surfaces()
            .filter(|surface| {
                ReadProviderScope::ReadOnly.includes(surface) && selection.includes(surface)
            })
            .map(|surface| surface.name.as_str())
            .collect::<Vec<_>>();
        let overlay_write_names = catalog
            .surfaces()
            .filter(|surface| is_write_surface(surface) && selection.includes(surface))
            .map(|surface| surface.name.as_str())
            .collect::<Vec<_>>();

        assert_eq!(committed_read_names, vec!["lix_history"]);
        assert_eq!(overlay_write_names, vec!["lix_file"]);
    }

    #[test]
    fn transaction_registration_partitions_provider_construction_once() {
        let schema = json!({
            "$schema": "https://lix.dev/schema-v1.json",
            "key": "phase8_row",
            "columns": [
                { "name": "id", "type": "text", "nullable": false },
            ],
            "primary_key": ["id"],
        });
        let catalog = PublicCatalog::from_visible_schemas(&[schema]).expect("catalog should build");

        let read_only = catalog
            .surfaces()
            .filter(|surface| ReadProviderScope::ReadOnly.includes(surface))
            .map(|surface| surface.name.as_str())
            .collect::<Vec<_>>();
        let writable = catalog
            .surfaces()
            .filter(|surface| is_write_surface(surface))
            .map(|surface| surface.name.as_str())
            .collect::<Vec<_>>();
        let all_read = catalog
            .surfaces()
            .filter(|surface| ReadProviderScope::All.includes(surface))
            .count();

        assert_eq!(
            read_only,
            vec![
                "lix_apply",
                "lix_as_of",
                "lix_change",
                "lix_commit_ancestry",
                "lix_create_checkpoint",
                "lix_diff",
                "lix_history",
                "lix_log",
                "lix_redo",
                "lix_restore",
                "lix_revert",
                "lix_revert_range",
                "lix_undo",
            ]
        );
        assert_eq!(
            writable,
            vec![
                "lix_branch",
                "lix_directory",
                "lix_file",
                "phase8_row",
            ]
        );
        assert_eq!(read_only.len() + writable.len(), catalog.surfaces().count());
        assert_eq!(all_read + writable.len(), 21, "construction count");
        assert_eq!(read_only.len() + writable.len(), 17, "surface count");
    }

    #[test]
    fn target_write_selection_reduces_provider_construction_count_to_one() {
        let catalog = PublicCatalog::from_visible_schemas(&[]).expect("catalog should build");
        let all_writable = catalog
            .surfaces()
            .filter(|surface| is_write_surface(surface))
            .count();
        let selection = selected_names(&["lix_file"]);
        let selected_writable = catalog
            .surfaces()
            .filter(|surface| is_write_surface(surface) && selection.includes(surface))
            .map(|surface| surface.name.as_str())
            .collect::<Vec<_>>();

        assert_eq!(all_writable, 3, "standalone write count");
        assert_eq!(selected_writable, vec!["lix_file"]);
    }

    #[test]
    fn provider_history_schemas_match_catalog_contract_order() {
        let catalog = PublicCatalog::from_visible_schemas(&[]).expect("catalog should build");

        assert_surface_schema_matches_provider_schema(
            &catalog,
            "lix_file",
            file::lix_file_schema(),
        );
        assert_surface_schema_matches_provider_schema(
            &catalog,
            "lix_directory",
            directory::lix_directory_schema(),
        );
        assert_surface_schema_matches_provider_schema(
            &catalog,
            "lix_branch",
            branch::lix_branch_schema(),
        );
        assert_surface_schema_matches_provider_schema(
            &catalog,
            "lix_change",
            change::lix_change_schema(),
        );
        assert_history_schema_matches_provider_schema(
            &catalog,
            "lix_file",
            relation_history_schema(&catalog, "lix_file").expect("file history schema"),
        );
        assert_history_schema_matches_provider_schema(
            &catalog,
            "lix_directory",
            relation_history_schema(&catalog, "lix_directory").expect("directory history schema"),
        );
    }

    #[test]
    fn file_content_surfaces_use_large_binary() {
        let catalog = PublicCatalog::from_visible_schemas(&[]).expect("catalog should build");

        for (surface_name, column, schema) in [
            ("lix_file", "content", catalog.surface_schema("lix_file")),
            (
                "lix_history('lix_file')",
                "to_content",
                catalog.history_relation_schema("lix_file"),
            ),
        ] {
            let schema = schema.unwrap_or_else(|| panic!("{surface_name} should be in catalog"));
            let content_field = schema
                .field_with_name(column)
                .unwrap_or_else(|_| panic!("{surface_name}.content should exist"));

            assert_eq!(
                content_field.data_type(),
                &DataType::LargeBinary,
                "{surface_name}.content should avoid Arrow Binary's 32-bit offset limit",
            );
        }
    }

    #[tokio::test]
    async fn provider_row_schemas_match_catalog_contract_order() {
        let schema = json!({
            "$schema": "https://lix.dev/schema-v1.json",
            "key": "phase8_row",
            "columns": [
                { "name": "id", "type": "text", "nullable": false },
                { "name": "count", "type": "int8", "nullable": true },
                { "name": "body", "type": "jsonb", "nullable": true },
            ],
            "primary_key": ["id"],
        });
        let catalog =
            PublicCatalog::from_visible_schemas(&[schema.clone()]).expect("catalog should build");
        let _spec = derive_schema_surface_spec_from_schema(&schema).expect("schema should derive");
        let session = SessionContext::new();
        schema::register_row_providers(
            &session,
            "01920000-0000-7000-8000-0000000000a1",
            Arc::new(EmptyHotStateReader),
            None,
            Arc::new(EmptyBranchRefReader),
            &catalog,
            true,
            &ProviderSelection::All,
        )
        .expect("row providers should register");

        assert_registered_table_schema_matches_catalog(&session, &catalog, "phase8_row").await;
    }

    async fn assert_registered_table_schema_matches_catalog(
        session: &SessionContext,
        catalog: &PublicCatalog,
        surface_name: &str,
    ) {
        let provider = session
            .table_provider(surface_name)
            .await
            .unwrap_or_else(|error| panic!("{surface_name} provider should load: {error}"));
        assert_surface_schema_matches_provider_schema(catalog, surface_name, provider.schema());
    }

    fn assert_surface_schema_matches_provider_schema(
        catalog: &PublicCatalog,
        surface_name: &str,
        provider_schema: SchemaRef,
    ) {
        let surface = catalog
            .surface(surface_name)
            .unwrap_or_else(|| panic!("{surface_name} should be in catalog"));
        let catalog_column_names = surface
            .columns
            .iter()
            .map(|column| column.name.as_str())
            .collect::<Vec<_>>();
        let provider_field_names = provider_schema
            .fields()
            .iter()
            .map(|field| field.name().as_str())
            .collect::<Vec<_>>();
        assert_eq!(
            catalog_column_names, provider_field_names,
            "{surface_name} column order"
        );

        let catalog_schema = catalog
            .surface_schema(surface_name)
            .unwrap_or_else(|| panic!("{surface_name} should be in catalog"));
        assert_eq!(
            catalog_schema.fields(),
            provider_schema.fields(),
            "{surface_name}"
        );
    }

    fn assert_history_schema_matches_provider_schema(
        catalog: &PublicCatalog,
        relation_name: &str,
        provider_schema: SchemaRef,
    ) {
        let contract = catalog
            .history_relation(relation_name)
            .unwrap_or_else(|| panic!("{relation_name} history should be in catalog"));
        let catalog_columns = contract
            .columns
            .iter()
            .filter(|column| column.is_public())
            .map(|column| column.name.as_str())
            .collect::<Vec<_>>();
        let provider_columns = provider_schema
            .fields()
            .iter()
            .map(|field| field.name().as_str())
            .collect::<Vec<_>>();
        assert_eq!(
            catalog_columns, provider_columns,
            "{relation_name} history columns"
        );
    }

    struct EmptyHotStateReader;

    #[async_trait]
    impl HotStateReader for EmptyHotStateReader {
        async fn load_exact_batch(
            &self,
            request: &crate::hot_state::HotStateExactBatchRequest,
        ) -> Result<crate::hot_state::MaterializedHotStateExactBatch, LixError> {
            crate::hot_state::load_exact_batch_via_scan_for_test(self, request).await
        }

        async fn scan_batch(
            &self,
            _request: &HotStateScanRequest,
        ) -> Result<crate::hot_state::MaterializedHotStateBatch, LixError> {
            Ok(Vec::new().into())
        }
    }

    struct EmptyBranchRefReader;

    #[async_trait]
    impl BranchRefReader for EmptyBranchRefReader {
        async fn load_head(&self, branch_id: &str) -> Result<Option<BranchHead>, LixError> {
            Ok(Some(BranchHead {
                working_base_commit_id: None,
                branch_id: branch_id.to_string(),
                commit_id: CommitId::for_test_label(&format!("commit-{branch_id}")),
            }))
        }

        async fn scan_heads(&self) -> Result<Vec<BranchHead>, LixError> {
            Ok(Vec::new().into())
        }
    }
}

pub(crate) use diff::prepare_native_diff_interest;
pub(crate) use file::{
    prepare_native_file_content_inputs, prepare_native_file_content_interest,
    prepare_native_file_metadata_interest,
};