mz_sql/pure/

postgres.rs

// Copyright Materialize, Inc. and contributors. All rights reserved.
//
// Use of this software is governed by the Business Source License
// included in the LICENSE file.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0.

//! Postgres utilities for SQL purification.

use std::collections::{BTreeMap, BTreeSet};

use mz_postgres_util::desc::PostgresTableDesc;
use mz_proto::RustType;
use mz_repr::{Datum, ReprColumnType, ReprScalarType, Row, SqlScalarType};
use mz_sql_parser::ast::display::AstDisplay;
use mz_sql_parser::ast::{
    ColumnDef, CreateSubsourceOption, CreateSubsourceOptionName, CreateSubsourceStatement,
    ExternalReferences, Ident, PgConfigOptionName, TableConstraint, UnresolvedItemName, Value,
    WithOptionValue,
};
use mz_storage_types::sources::SourceExportStatementDetails;
use mz_storage_types::sources::casts::{CastFunc, StorageScalarExpr};
use mz_storage_types::sources::postgres::CastType;
use prost::Message;
use tokio_postgres::Client;
use tokio_postgres::types::Oid;

use crate::names::{Aug, ResolvedItemName};
use crate::normalize;
use crate::plan::{PlanError, StatementContext};

use super::error::PgSourcePurificationError;
use super::references::RetrievedSourceReferences;
use super::{PartialItemName, PurifiedExportDetails, PurifiedSourceExport, SourceReferencePolicy};

/// Ensure that we have select permissions on all tables; we have to do this before we
/// start snapshotting because if we discover we cannot `COPY` from a table while
/// snapshotting, we break the entire source.
pub(super) async fn validate_requested_references_privileges(
    client: &Client,
    table_oids: &[Oid],
) -> Result<(), PlanError> {
    privileges::check_table_privileges(client, table_oids).await?;
    privileges::check_rls_privileges(client, table_oids).await?;
    replica_identity::check_replica_identity_full(client, table_oids).await?;

    Ok(())
}

/// Map a list of column references to a map of table oids to column names.
///
/// Additionally, modify `columns` so that they contain database-qualified
/// references to the columns.
pub(super) fn map_column_refs(
    retrieved_references: &RetrievedSourceReferences,
    columns: &mut [UnresolvedItemName],
    option_type: PgConfigOptionName,
) -> Result<BTreeMap<u32, BTreeSet<String>>, PlanError> {
    let mut cols_map: BTreeMap<u32, BTreeSet<String>> = BTreeMap::new();

    for name in columns {
        let (qual, col) = match name.0.split_last().expect("must have at least one element") {
            (col, []) => {
                return Err(PlanError::InvalidOptionValue {
                    option_name: option_type.to_ast_string_simple(),
                    err: Box::new(PlanError::UnderqualifiedColumnName(
                        col.as_str().to_string(),
                    )),
                });
            }
            (col, qual) => (qual.to_vec(), col.as_str().to_string()),
        };

        let resolved_reference = retrieved_references.resolve_name(&qual)?;
        let mut fully_qualified_name =
            resolved_reference
                .external_reference()
                .map_err(|e| PlanError::InvalidOptionValue {
                    option_name: option_type.to_ast_string_simple(),
                    err: Box::new(e.into()),
                })?;

        let desc = resolved_reference
            .postgres_desc()
            .expect("known to be postgres");

        if !desc.columns.iter().any(|column| column.name == col) {
            let column = mz_repr::ColumnName::from(col);
            let similar = desc
                .columns
                .iter()
                .filter_map(|c| {
                    let c_name = mz_repr::ColumnName::from(c.name.clone());
                    c_name.is_similar(&column).then_some(c_name)
                })
                .collect();
            return Err(PlanError::InvalidOptionValue {
                option_name: option_type.to_ast_string_simple(),
                err: Box::new(PlanError::UnknownColumn {
                    table: Some(
                        normalize::unresolved_item_name(fully_qualified_name)
                            .expect("known to be of valid len"),
                    ),
                    column,
                    similar,
                }),
            });
        }

        // Rewrite fully qualified name.
        let col_ident = Ident::new(col.as_str().to_string())?;
        fully_qualified_name.0.push(col_ident);
        *name = fully_qualified_name;

        let new = cols_map
            .entry(desc.oid)
            .or_default()
            .insert(col.as_str().to_string());

        if !new {
            return Err(PlanError::InvalidOptionValue {
                option_name: option_type.to_ast_string_simple(),
                err: Box::new(PlanError::UnexpectedDuplicateReference { name: name.clone() }),
            });
        }
    }

    Ok(cols_map)
}

pub fn generate_create_subsource_statements(
    scx: &StatementContext,
    source_name: ResolvedItemName,
    requested_subsources: BTreeMap<UnresolvedItemName, PurifiedSourceExport>,
) -> Result<Vec<CreateSubsourceStatement<Aug>>, PlanError> {
    // Aggregate all unrecognized types.
    let mut unsupported_cols = vec![];

    // Now that we have an explicit list of validated requested subsources we can create them
    let mut subsources = Vec::with_capacity(requested_subsources.len());

    for (subsource_name, purified_export) in requested_subsources {
        let PostgresExportStatementValues {
            columns,
            constraints,
            text_columns,
            exclude_columns,
            details,
            external_reference,
        } = generate_source_export_statement_values(scx, purified_export, &mut unsupported_cols)?;

        let mut with_options = vec![
            CreateSubsourceOption {
                name: CreateSubsourceOptionName::ExternalReference,
                value: Some(WithOptionValue::UnresolvedItemName(external_reference)),
            },
            CreateSubsourceOption {
                name: CreateSubsourceOptionName::Details,
                value: Some(WithOptionValue::Value(Value::String(hex::encode(
                    details.into_proto().encode_to_vec(),
                )))),
            },
        ];

        if let Some(text_columns) = text_columns {
            with_options.push(CreateSubsourceOption {
                name: CreateSubsourceOptionName::TextColumns,
                value: Some(WithOptionValue::Sequence(text_columns)),
            });
        }

        if let Some(exclude_columns) = exclude_columns {
            with_options.push(CreateSubsourceOption {
                name: CreateSubsourceOptionName::ExcludeColumns,
                value: Some(WithOptionValue::Sequence(exclude_columns)),
            });
        }

        // Create the subsource statement
        let subsource = CreateSubsourceStatement {
            name: subsource_name,
            columns,
            // We might not know the primary source's `GlobalId` yet; if not,
            // we'll fill it in once we generate it.
            of_source: Some(source_name.clone()),
            // TODO(petrosagg): nothing stops us from getting the constraints of the
            // upstream tables and mirroring them here which will lead to more optimization
            // opportunities if for example there is a primary key or an index.
            //
            // If we ever do that we must triple check that we will get notified *in the
            // replication stream*, if our assumptions change. Failure to do that could
            // mean that an upstream table that started with an index was then altered to
            // one without and now we're producing garbage data.
            constraints,
            if_not_exists: false,
            with_options,
        };
        subsources.push(subsource);
    }

    if !unsupported_cols.is_empty() {
        unsupported_cols.sort();
        Err(PgSourcePurificationError::UnrecognizedTypes {
            cols: unsupported_cols,
        })?;
    }

    Ok(subsources)
}

pub(super) struct PostgresExportStatementValues {
    pub(super) columns: Vec<ColumnDef<Aug>>,
    pub(super) constraints: Vec<TableConstraint<Aug>>,
    pub(super) text_columns: Option<Vec<WithOptionValue<Aug>>>,
    pub(super) exclude_columns: Option<Vec<WithOptionValue<Aug>>>,
    pub(super) details: SourceExportStatementDetails,
    pub(super) external_reference: UnresolvedItemName,
}

pub(super) fn generate_source_export_statement_values(
    scx: &StatementContext,
    purified_export: PurifiedSourceExport,
    unsupported_cols: &mut Vec<(String, mz_repr::adt::system::Oid)>,
) -> Result<PostgresExportStatementValues, PlanError> {
    let PurifiedExportDetails::Postgres {
        table,
        text_columns,
        exclude_columns,
    } = purified_export.details
    else {
        bail_internal!("purified export details must be postgres");
    };

    let text_column_set = BTreeSet::from_iter(text_columns.iter().flatten().map(Ident::as_str));
    let exclude_column_set =
        BTreeSet::from_iter(exclude_columns.iter().flatten().map(Ident::as_str));

    // Figure out the schema of the subsource
    let mut columns = vec![];
    for c in table.columns.iter() {
        let name = Ident::new(c.name.clone())?;

        if exclude_column_set.contains(c.name.as_str()) {
            continue;
        }

        let ty = if text_column_set.contains(c.name.as_str()) {
            mz_pgrepr::Type::Text
        } else {
            match mz_pgrepr::Type::from_oid_and_typmod(c.type_oid, c.type_mod) {
                Ok(t) => t,
                Err(_) => {
                    let mut full_name = purified_export.external_reference.0.clone();
                    full_name.push(name);
                    unsupported_cols.push((
                        UnresolvedItemName(full_name).to_ast_string_simple(),
                        mz_repr::adt::system::Oid(c.type_oid),
                    ));
                    continue;
                }
            }
        };

        let data_type = scx.resolve_type(ty)?;
        let mut options = vec![];

        if !c.nullable {
            options.push(mz_sql_parser::ast::ColumnOptionDef {
                name: None,
                option: mz_sql_parser::ast::ColumnOption::NotNull,
            });
        }

        columns.push(ColumnDef {
            name,
            data_type,
            collation: None,
            options,
        });
    }

    let mut constraints = vec![];
    for key in table.keys.clone() {
        let mut key_columns = vec![];
        let mut all_key_cols_included = true;

        for col_num in key.cols {
            match table.columns.iter().find(|col| col.col_num == col_num) {
                Some(col) => {
                    let ident = Ident::new(col.name.clone())?;
                    key_columns.push(ident);
                }
                None => {
                    all_key_cols_included = false;
                    break;
                }
            }
        }
        if !all_key_cols_included {
            continue;
        }

        let constraint = mz_sql_parser::ast::TableConstraint::Unique {
            name: Some(Ident::new(key.name)?),
            columns: key_columns,
            is_primary: key.is_primary,
            nulls_not_distinct: key.nulls_not_distinct,
        };

        // We take the first constraint available to be the primary key.
        if key.is_primary {
            constraints.insert(0, constraint);
        } else {
            constraints.push(constraint);
        }
    }
    let details = SourceExportStatementDetails::Postgres { table };

    let text_columns = text_columns.map(|mut columns| {
        columns.sort();
        columns
            .into_iter()
            .map(WithOptionValue::Ident::<Aug>)
            .collect()
    });

    let exclude_columns = exclude_columns.map(|mut columns| {
        columns.sort();
        columns
            .into_iter()
            .map(WithOptionValue::Ident::<Aug>)
            .collect()
    });

    Ok(PostgresExportStatementValues {
        columns,
        constraints,
        text_columns,
        exclude_columns,
        details,
        external_reference: purified_export.external_reference,
    })
}

pub(super) struct PurifiedSourceExports {
    pub(super) source_exports: BTreeMap<UnresolvedItemName, PurifiedSourceExport>,
    // NOTE(roshan): The text columns are already part of their
    // appropriate `source_exports` above, but these are returned to allow
    // round-tripping a `CREATE SOURCE` statement while we still allow creating
    // implicit subsources from `CREATE SOURCE`. Remove once
    // fully deprecating that feature and forcing users to use explicit
    // `CREATE TABLE .. FROM SOURCE` statements.
    pub(super) normalized_text_columns: Vec<WithOptionValue<Aug>>,
}

// Purify the requested external references, returning a set of purified
// source exports corresponding to external tables, and and additional
// fields necessary to generate relevant statements and update statement options
pub(super) async fn purify_source_exports(
    client: &Client,
    retrieved_references: &RetrievedSourceReferences,
    requested_references: &Option<ExternalReferences>,
    mut text_columns: Vec<UnresolvedItemName>,
    mut exclude_columns: Vec<UnresolvedItemName>,
    unresolved_source_name: &UnresolvedItemName,
    reference_policy: &SourceReferencePolicy,
) -> Result<PurifiedSourceExports, PlanError> {
    let requested_exports = match requested_references.as_ref() {
        Some(requested) if matches!(reference_policy, SourceReferencePolicy::NotAllowed) => {
            Err(PlanError::UseTablesForSources(requested.to_string()))?
        }
        Some(requested) => retrieved_references
            .requested_source_exports(Some(requested), unresolved_source_name)?,
        None => {
            if matches!(reference_policy, SourceReferencePolicy::Required) {
                Err(PgSourcePurificationError::RequiresExternalReferences)?
            }

            // If no external reference is specified, it does not make sense to include
            // text columns.
            if !text_columns.is_empty() {
                Err(
                    PgSourcePurificationError::UnnecessaryOptionsWithoutReferences(
                        "TEXT COLUMNS".to_string(),
                    ),
                )?
            }

            // If no external reference is specified, it does not make sense to include
            // exclude columns.
            if !exclude_columns.is_empty() {
                Err(
                    PgSourcePurificationError::UnnecessaryOptionsWithoutReferences(
                        "EXCLUDE COLUMNS".to_string(),
                    ),
                )?
            }

            return Ok(PurifiedSourceExports {
                source_exports: BTreeMap::new(),
                normalized_text_columns: vec![],
            });
        }
    };

    if requested_exports.is_empty() {
        sql_bail!(
            "[internal error]: Postgres reference {} did not match any tables",
            requested_references
                .as_ref()
                .unwrap()
                .to_ast_string_simple()
        );
    }

    super::validate_source_export_names(&requested_exports)?;

    let table_oids: Vec<_> = requested_exports
        .iter()
        .map(|r| r.meta.postgres_desc().expect("is postgres").oid)
        .collect();

    validate_requested_references_privileges(client, &table_oids).await?;

    let mut text_column_map = map_column_refs(
        retrieved_references,
        &mut text_columns,
        PgConfigOptionName::TextColumns,
    )?;
    let mut exclude_column_map = map_column_refs(
        retrieved_references,
        &mut exclude_columns,
        PgConfigOptionName::ExcludeColumns,
    )?;

    // Normalize options to contain full qualified values.
    text_columns.sort();
    text_columns.dedup();
    let normalized_text_columns: Vec<_> = text_columns
        .into_iter()
        .map(WithOptionValue::UnresolvedItemName)
        .collect();

    let source_exports = requested_exports
        .into_iter()
        .map(|r| {
            let mut desc = r.meta.postgres_desc().expect("known postgres").clone();
            let text_columns = text_column_map.remove(&desc.oid);
            let exclude_columns = exclude_column_map.remove(&desc.oid);

            if let Some(exclude_cols) = &exclude_columns {
                desc.columns.retain(|c| !exclude_cols.contains(&c.name));
            }

            if let (Some(text_cols), Some(exclude_cols)) = (&text_columns, &exclude_columns) {
                let intersection: Vec<_> = text_cols.intersection(exclude_cols).collect();
                if !intersection.is_empty() {
                    return Err(PgSourcePurificationError::DuplicatedColumnNames(
                        intersection.iter().map(|s| (*s).to_string()).collect(),
                    ));
                }
            }
            Ok((
                r.name,
                PurifiedSourceExport {
                    external_reference: r.external_reference,
                    details: PurifiedExportDetails::Postgres {
                        text_columns: text_columns.map(|v| {
                            v.into_iter()
                                .map(|s| Ident::new(s).expect("validated above"))
                                .collect()
                        }),
                        exclude_columns: exclude_columns.map(|v| {
                            v.into_iter()
                                .map(|s| Ident::new(s).expect("validated above"))
                                .collect()
                        }),
                        table: desc,
                    },
                },
            ))
        })
        .collect::<Result<BTreeMap<_, _>, _>>()?;

    if !text_column_map.is_empty() {
        // If any any item was not removed from the text_column_map, it wasn't being
        // added.
        let mut dangling_text_column_refs = vec![];
        let all_references = retrieved_references.all_references();

        for id in text_column_map.keys() {
            let desc = all_references
                .iter()
                .find_map(|reference| {
                    let desc = reference.postgres_desc().expect("is postgres");
                    if desc.oid == *id { Some(desc) } else { None }
                })
                .expect("validated when generating text columns");

            dangling_text_column_refs.push(PartialItemName {
                database: None,
                schema: Some(desc.namespace.clone()),
                item: desc.name.clone(),
            });
        }

        dangling_text_column_refs.sort();
        return Err(PlanError::from(
            PgSourcePurificationError::DanglingTextColumns {
                items: dangling_text_column_refs,
            },
        ));
    }

    if !exclude_column_map.is_empty() {
        // If any any item was not removed from the exclude_column_map, it wasn't being
        // added.
        let mut dangling_exclude_column_refs = vec![];
        let all_references = retrieved_references.all_references();

        for id in exclude_column_map.keys() {
            let desc = all_references
                .iter()
                .find_map(|reference| {
                    let desc = reference.postgres_desc().expect("is postgres");
                    if desc.oid == *id { Some(desc) } else { None }
                })
                .expect("validated when generating exclude columns");

            dangling_exclude_column_refs.push(PartialItemName {
                database: None,
                schema: Some(desc.namespace.clone()),
                item: desc.name.clone(),
            });
        }

        dangling_exclude_column_refs.sort();
        return Err(PlanError::from(
            PgSourcePurificationError::DanglingExcludeColumns {
                items: dangling_exclude_column_refs,
            },
        ));
    }

    Ok(PurifiedSourceExports {
        source_exports,
        normalized_text_columns,
    })
}

pub(crate) fn generate_column_casts(
    scx: &StatementContext,
    table: &PostgresTableDesc,
    text_columns: &Vec<Ident>,
) -> Result<Vec<(CastType, StorageScalarExpr)>, PlanError> {
    // Generate the cast expressions required to convert the text encoded columns into
    // the appropriate target types, creating a Vec<StorageScalarExpr>.
    // The postgres source reader will then eval each of those on the incoming rows.

    let text_columns = BTreeSet::from_iter(text_columns.iter().map(Ident::as_str));

    let mut table_cast = vec![];
    for (i, column) in table.columns.iter().enumerate() {
        let (cast_type, ty) = if text_columns.contains(column.name.as_str()) {
            // Treat the column as text if it was referenced in
            // `TEXT COLUMNS`. This is the only place we need to
            // perform this logic; even if the type is unsupported,
            // we'll be able to ingest its values as text in
            // storage.
            (CastType::Text, mz_pgrepr::Type::Text)
        } else {
            match mz_pgrepr::Type::from_oid_and_typmod(column.type_oid, column.type_mod) {
                Ok(t) => (CastType::Natural, t),
                // If this reference survived purification, we
                // do not expect it to be from a table that the
                // user will consume., i.e. expect this table to
                // be filtered out of table casts.
                Err(_) => {
                    table_cast.push((
                        CastType::Natural,
                        StorageScalarExpr::ErrorIfNull(
                            Box::new(StorageScalarExpr::Literal(
                                Row::pack_slice(&[Datum::Null]),
                                ReprColumnType {
                                    nullable: true,
                                    scalar_type: ReprScalarType::String,
                                },
                            )),
                            format!("Unsupported type with OID {}", column.type_oid),
                        ),
                    ));
                    continue;
                }
            }
        };

        let cast_expr = match pg_type_to_cast_func(scx, &ty) {
            Ok(None) => {
                // No cast needed (e.g. Text → String identity).
                StorageScalarExpr::Column(i)
            }
            Ok(Some(cast_func)) => {
                StorageScalarExpr::CallUnary(cast_func, Box::new(StorageScalarExpr::Column(i)))
            }
            Err(PlanError::TableContainsUningestableTypes { type_, .. }) => {
                // We expect only reg* types and similar to encounter
                // this. Users can ingest the data as text if they need
                // to. This is acceptable because we don't expect the
                // OIDs from an external PG source to be unilaterally
                // usable in resolving item names in MZ.
                return Err(PlanError::TableContainsUningestableTypes {
                    name: table.name.to_string(),
                    type_,
                    column: column.name.to_string(),
                });
            }
            Err(e) => return Err(e),
        };

        let cast = if column.nullable {
            cast_expr
        } else {
            // We must enforce nullability constraint on cast
            // because PG replication stream does not propagate
            // constraint changes and we want to error subsource if
            // e.g. the constraint is dropped and we don't notice
            // it.
            let message = format!(
                "PG column {}.{}.{} contained NULL data, despite having NOT NULL constraint",
                table.namespace, table.name, column.name
            );
            StorageScalarExpr::ErrorIfNull(Box::new(cast_expr), message)
        };

        table_cast.push((cast_type, cast));
    }
    Ok(table_cast)
}

/// Resolve a PG type to its corresponding `SqlScalarType` via the catalog.
fn resolve_pg_type_to_scalar_type(
    scx: &StatementContext,
    ty: &mz_pgrepr::Type,
) -> Result<SqlScalarType, PlanError> {
    let data_type = scx.resolve_type(ty.clone())?;
    crate::plan::query::scalar_type_from_sql(scx, &data_type)
}

/// Map a PG type to the corresponding `CastFunc` variant. Returns:
/// - `Ok(Some(func))` for types that need a cast
/// - `Ok(None)` for types that need no cast (Text → String identity)
/// - `Err(PlanError::TableContainsUningestableTypes { .. })` for types
///   that cannot be ingested. The error uses placeholder strings for
///   table/column name; callers with context should use
///   `pg_type_to_cast_func_or_uningestable` instead.
fn pg_type_to_cast_func(
    scx: &StatementContext,
    ty: &mz_pgrepr::Type,
) -> Result<Option<CastFunc>, PlanError> {
    use mz_pgrepr::Type;

    let cast_func = match ty {
        Type::Bool => CastFunc::CastStringToBool,
        Type::Bytea => CastFunc::CastStringToBytes,
        Type::Char => CastFunc::CastStringToPgLegacyChar,
        Type::Date => CastFunc::CastStringToDate,
        Type::Float4 => CastFunc::CastStringToFloat32,
        Type::Float8 => CastFunc::CastStringToFloat64,
        Type::Int2 => CastFunc::CastStringToInt16,
        Type::Int4 => CastFunc::CastStringToInt32,
        Type::Int8 => CastFunc::CastStringToInt64,
        Type::UInt2 => CastFunc::CastStringToUint16,
        Type::UInt4 => CastFunc::CastStringToUint32,
        Type::UInt8 => CastFunc::CastStringToUint64,
        Type::Interval { .. } => CastFunc::CastStringToInterval,
        Type::Jsonb => CastFunc::CastStringToJsonb,
        Type::Name => CastFunc::CastStringToPgLegacyName,
        Type::Numeric { .. } => {
            // Resolve through the catalog to get the repr NumericMaxScale type.
            let scalar_type = resolve_pg_type_to_scalar_type(scx, ty)?;
            match scalar_type {
                SqlScalarType::Numeric { max_scale } => CastFunc::CastStringToNumeric(max_scale),
                _ => unreachable!("Numeric must resolve to Numeric"),
            }
        }
        Type::Oid => CastFunc::CastStringToOid,
        Type::Text => return Ok(None),
        Type::BpChar { .. } => {
            // Resolve through the catalog to get the repr CharLength type.
            let scalar_type = resolve_pg_type_to_scalar_type(scx, ty)?;
            match scalar_type {
                SqlScalarType::Char { length } => CastFunc::CastStringToChar {
                    length,
                    fail_on_len: true,
                },
                _ => unreachable!("BpChar must resolve to Char"),
            }
        }
        Type::VarChar { .. } => {
            // Resolve through the catalog to get the repr VarCharMaxLength type.
            let scalar_type = resolve_pg_type_to_scalar_type(scx, ty)?;
            match scalar_type {
                SqlScalarType::VarChar { max_length } => CastFunc::CastStringToVarChar {
                    length: max_length,
                    fail_on_len: true,
                },
                _ => unreachable!("VarChar must resolve to VarChar"),
            }
        }
        Type::Time { .. } => {
            // Time precision is not yet fully supported; resolve_type strips precision.
            CastFunc::CastStringToTime
        }
        Type::Timestamp { .. } => {
            // Resolve through the catalog to get the repr TimestampPrecision type.
            let scalar_type = resolve_pg_type_to_scalar_type(scx, ty)?;
            match scalar_type {
                SqlScalarType::Timestamp { precision } => {
                    CastFunc::CastStringToTimestamp(precision)
                }
                _ => unreachable!("Timestamp must resolve to Timestamp"),
            }
        }
        Type::TimestampTz { .. } => {
            // Resolve through the catalog to get the repr TimestampPrecision type.
            let scalar_type = resolve_pg_type_to_scalar_type(scx, ty)?;
            match scalar_type {
                SqlScalarType::TimestampTz { precision } => {
                    CastFunc::CastStringToTimestampTz(precision)
                }
                _ => unreachable!("TimestampTz must resolve to TimestampTz"),
            }
        }
        Type::Uuid => CastFunc::CastStringToUuid,
        Type::Int2Vector => CastFunc::CastStringToInt2Vector,
        Type::MzTimestamp => CastFunc::CastStringToMzTimestamp,
        // JSON is ingested as JSONB (same as the old plan_cast path).
        Type::Json => CastFunc::CastStringToJsonb,
        Type::Array(elem) => {
            let return_ty = resolve_pg_type_to_scalar_type(scx, ty)?;
            let elem_cast = build_element_cast_expr(scx, elem)?;
            CastFunc::CastStringToArray {
                return_ty,
                cast_expr: Box::new(elem_cast),
            }
        }
        Type::List(elem) => {
            let return_ty = resolve_pg_type_to_scalar_type(scx, ty)?;
            let elem_cast = build_element_cast_expr(scx, elem)?;
            CastFunc::CastStringToList {
                return_ty,
                cast_expr: Box::new(elem_cast),
            }
        }
        Type::Map { value_type } => {
            let return_ty = resolve_pg_type_to_scalar_type(scx, ty)?;
            let value_cast = build_element_cast_expr(scx, value_type)?;
            CastFunc::CastStringToMap {
                return_ty,
                cast_expr: Box::new(value_cast),
            }
        }
        Type::Range { element_type } => {
            let return_ty = resolve_pg_type_to_scalar_type(scx, ty)?;
            let elem_cast = build_element_cast_expr(scx, element_type)?;
            CastFunc::CastStringToRange {
                return_ty,
                cast_expr: Box::new(elem_cast),
            }
        }
        // reg* types require subquery-based casts that storage cannot
        // evaluate. Users can ingest them as text via TEXT COLUMNS.
        Type::RegType | Type::RegClass | Type::RegProc => {
            return Err(PlanError::TableContainsUningestableTypes {
                name: String::new(),
                type_: ty.name().to_string(),
                column: String::new(),
            });
        }
        other => {
            return Err(PlanError::TableContainsUningestableTypes {
                name: String::new(),
                type_: other.name().to_string(),
                column: String::new(),
            });
        }
    };
    Ok(Some(cast_func))
}

/// Build the element cast expression for container types (Array, List, Map,
/// Range). The element expression operates on a single-column input row
/// containing the text-encoded element at column 0.
fn build_element_cast_expr(
    scx: &StatementContext,
    elem_ty: &mz_pgrepr::Type,
) -> Result<StorageScalarExpr, PlanError> {
    match pg_type_to_cast_func(scx, elem_ty)? {
        None => Ok(StorageScalarExpr::Column(0)),
        Some(cast_func) => Ok(StorageScalarExpr::CallUnary(
            cast_func,
            Box::new(StorageScalarExpr::Column(0)),
        )),
    }
}

mod privileges {
    use mz_postgres_util::{PostgresError, query, sql};

    use super::*;
    use crate::plan::PlanError;
    use crate::pure::PgSourcePurificationError;

    async fn check_schema_privileges(client: &Client, table_oids: &[Oid]) -> Result<(), PlanError> {
        let invalid_schema_privileges_rows = query(
            client,
            sql!(
                "
                WITH distinct_namespace AS (
                    SELECT
                        DISTINCT n.oid, n.nspname AS schema_name
                    FROM unnest($1::OID[]) AS oids (oid)
                    JOIN pg_class AS c ON c.oid = oids.oid
                    JOIN pg_namespace AS n ON c.relnamespace = n.oid
                )
                SELECT d.schema_name
                FROM distinct_namespace AS d
                WHERE
                    NOT has_schema_privilege(CURRENT_USER::TEXT, d.oid, 'usage')"
            ),
            &[&table_oids],
        )
        .await?;

        let mut invalid_schema_privileges = invalid_schema_privileges_rows
            .into_iter()
            .map(|row| row.get("schema_name"))
            .collect::<Vec<String>>();

        if invalid_schema_privileges.is_empty() {
            Ok(())
        } else {
            invalid_schema_privileges.sort();
            Err(PgSourcePurificationError::UserLacksUsageOnSchemas {
                schemas: invalid_schema_privileges,
            })?
        }
    }

    /// Ensure that the user specified in `config` has:
    ///
    /// -`SELECT` privileges for the identified `tables`.
    ///
    ///  `tables`'s elements should be of the structure `[<schema name>, <table name>]`.
    ///
    /// - `USAGE` privileges on the schemas references in `tables`.
    ///
    /// # Panics
    /// If `config` does not specify a user.
    pub async fn check_table_privileges(
        client: &Client,
        table_oids: &[Oid],
    ) -> Result<(), PlanError> {
        check_schema_privileges(client, table_oids).await?;

        let invalid_table_privileges_rows = query(
            client,
            sql!(
                "
            SELECT
                format('%I.%I', n.nspname, c.relname) AS schema_qualified_table_name
             FROM unnest($1::oid[]) AS oids (oid)
             JOIN
                 pg_class c ON c.oid = oids.oid
             JOIN
                 pg_namespace n ON c.relnamespace = n.oid
             WHERE NOT has_table_privilege(CURRENT_USER::text, c.oid, 'select')"
            ),
            &[&table_oids],
        )
        .await?;

        let mut invalid_table_privileges = invalid_table_privileges_rows
            .into_iter()
            .map(|row| row.get("schema_qualified_table_name"))
            .collect::<Vec<String>>();

        if invalid_table_privileges.is_empty() {
            Ok(())
        } else {
            invalid_table_privileges.sort();
            Err(PgSourcePurificationError::UserLacksSelectOnTables {
                tables: invalid_table_privileges,
            })?
        }
    }

    /// Ensure that the user specified in `config` can read data from tables if row level security
    /// (RLS) is enabled. If the user/role does not have the BYPASSRLS attribute set, there is
    /// the possibility that MZ may not be able to read all data during the snapshot, which would
    /// result in missing data.
    pub async fn check_rls_privileges(
        client: &Client,
        table_oids: &[Oid],
    ) -> Result<(), PlanError> {
        match mz_postgres_util::validate_no_rls_policies(client, table_oids).await {
            Ok(_) => Ok(()),
            Err(err) => match err {
                // This is a little gross to do, but PlanError::PostgresConnectionErr implements
                // From<PostgresError>, and the error in that case would be
                // "failed to connect to PostgreSQL database", which doesn't make any sense.
                PostgresError::BypassRLSRequired(tables) => {
                    Err(PgSourcePurificationError::BypassRLSRequired { tables })?
                }
                _ => Err(err)?,
            },
        }
    }
}

mod replica_identity {
    use mz_postgres_util::{query, sql};

    use super::*;
    use crate::plan::PlanError;
    use crate::pure::PgSourcePurificationError;

    /// Ensures that all provided OIDs are tables with `REPLICA IDENTITY FULL`.
    pub async fn check_replica_identity_full(
        client: &Client,
        table_oids: &[Oid],
    ) -> Result<(), PlanError> {
        let invalid_replica_identity_rows = query(
            client,
            sql!(
                "
            SELECT
                format('%I.%I', n.nspname, c.relname) AS schema_qualified_table_name
             FROM unnest($1::oid[]) AS oids (oid)
             JOIN
                 pg_class c ON c.oid = oids.oid
             JOIN
                 pg_namespace n ON c.relnamespace = n.oid
             WHERE relreplident != 'f' OR relreplident IS NULL;"
            ),
            &[&table_oids],
        )
        .await?;

        let mut invalid_replica_identity = invalid_replica_identity_rows
            .into_iter()
            .map(|row| row.get("schema_qualified_table_name"))
            .collect::<Vec<String>>();

        if invalid_replica_identity.is_empty() {
            Ok(())
        } else {
            invalid_replica_identity.sort();
            Err(PgSourcePurificationError::NotTablesWReplicaIdentityFull {
                items: invalid_replica_identity,
            })?
        }
    }
}