uqa_sql/compiler/

types.rs

//
// Unified Query Algebra
//
// Copyright (c) 2023-2026 Cognica, Inc.
//

//! Column type and foreign-key constraint lowering helpers.

use pg_query::protobuf::Node;
use pg_query::NodeEnum;

use crate::ast::{ColumnType, IntervalFields, RangeSubtype};
use crate::error::{Result, SQLError};

use super::tree::extract_string;

/// Parser-normalized type identity used by `PostgreSQL`'s `regtype` and `regprocedure` input functions. Components retain the parser's distinction between aliases such as unquoted `integer` (normalized to `pg_catalog.int4`) and a quoted type named `"integer"`. Modifier values do not participate in OID lookup; their presence is retained for declaration validation.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ParsedRegtypeName {
    pub names: Vec<String>,
    pub array_dimensions: usize,
    pub has_type_modifiers: bool,
}

/// Parser-normalized routine name and optional exact input-type signature used by `PostgreSQL`'s `regproc` and `regprocedure` input functions.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ParsedRegprocedureName {
    pub names: Vec<String>,
    pub argument_types: Option<Vec<ParsedRegtypeName>>,
}

const POSTGRES_IDENTIFIER_MAX_BYTES: usize = 63;
const POSTGRES_FUNCTION_MAX_ARGUMENTS: usize = 100;

fn scanner_isspace(byte: u8) -> bool {
    matches!(byte, b' ' | b'\t' | b'\n' | b'\r' | 0x0b | 0x0c)
}

fn truncate_postgres_identifier(mut identifier: String) -> String {
    if identifier.len() <= POSTGRES_IDENTIFIER_MAX_BYTES {
        return identifier;
    }
    let mut end = POSTGRES_IDENTIFIER_MAX_BYTES;
    while !identifier.is_char_boundary(end) {
        end -= 1;
    }
    identifier.truncate(end);
    identifier
}

/// Parse the dotted identifier strings consumed by `PostgreSQL`'s `reg*` input functions. This follows `SplitIdentifierString`: surrounding component whitespace is ignored, quoted components collapse doubled quotes without case folding, unquoted components extend to a dot or whitespace and use ASCII case folding, and every component is clipped to `PostgreSQL`'s 63-byte identifier limit.
#[must_use]
pub fn parse_regobject_name(input: &str) -> Option<Vec<String>> {
    let bytes = input.as_bytes();
    let mut offset = 0usize;
    while bytes.get(offset).is_some_and(|byte| scanner_isspace(*byte)) {
        offset += 1;
    }
    if offset == bytes.len() {
        return None;
    }

    let mut names = Vec::new();
    loop {
        let component = if bytes[offset] == b'"' {
            offset += 1;
            let mut quoted = String::new();
            loop {
                let relative = bytes[offset..].iter().position(|byte| *byte == b'"')?;
                let quote = offset + relative;
                quoted.push_str(&input[offset..quote]);
                offset = quote + 1;
                if bytes.get(offset) == Some(&b'"') {
                    quoted.push('"');
                    offset += 1;
                    continue;
                }
                break;
            }
            quoted
        } else {
            let start = offset;
            while bytes
                .get(offset)
                .is_some_and(|byte| *byte != b'.' && !scanner_isspace(*byte))
            {
                offset += 1;
            }
            if offset == start {
                return None;
            }
            input[start..offset].to_ascii_lowercase()
        };
        names.push(truncate_postgres_identifier(component));

        while bytes.get(offset).is_some_and(|byte| scanner_isspace(*byte)) {
            offset += 1;
        }
        match bytes.get(offset) {
            None => return Some(names),
            Some(b'.') => {
                offset += 1;
                while bytes.get(offset).is_some_and(|byte| scanner_isspace(*byte)) {
                    offset += 1;
                }
                if offset == bytes.len() {
                    return None;
                }
            }
            Some(_) => return None,
        }
    }
}

/// Parse exactly one `PostgreSQL` type-name string without accepting adjacent SQL expressions. `None` is the soft-failure shape used for inputs such as an empty string or `SETOF integer`; lexical and type-name syntax errors are retained for the caller.
pub fn parse_regtype_name(input: &str) -> Result<Option<ParsedRegtypeName>> {
    if input.bytes().all(scanner_isspace) {
        return Ok(None);
    }
    let parsed = pg_query::parse_with_mode(input, pg_query::ParseMode::TypeName)?;
    let [raw] = parsed.protobuf.stmts.as_slice() else {
        return Ok(None);
    };
    let Some(NodeEnum::List(names)) = raw.stmt.as_ref().and_then(|node| node.node.as_ref()) else {
        return Ok(None);
    };
    let names = names
        .items
        .iter()
        .map(extract_string)
        .collect::<Result<Vec<_>>>()?;
    if names.is_empty() {
        return Ok(None);
    }
    let scanned = pg_query::scan(input)?;
    let tokens = scanned
        .tokens
        .iter()
        .filter_map(|token| pg_query::protobuf::Token::try_from(token.token).ok())
        .collect::<Vec<_>>();
    if tokens.contains(&pg_query::protobuf::Token::Setof) {
        return Ok(None);
    }
    let bracket_dimensions = tokens
        .iter()
        .filter(|token| **token == pg_query::protobuf::Token::Ascii91)
        .count();
    let array_dimensions = bracket_dimensions.max(usize::from(
        tokens.contains(&pg_query::protobuf::Token::Array),
    ));
    Ok(Some(ParsedRegtypeName {
        names,
        array_dimensions,
        has_type_modifiers: tokens.contains(&pg_query::protobuf::Token::Ascii40),
    }))
}

/// Parse one routine-name string with either an omitted signature (`regproc`) or an exact signature (`regprocedure`). The caller owns the soft-error and cross-database policy because the two SQL input functions intentionally differ from ordinary DDL.
pub fn parse_regprocedure_name(input: &str) -> Result<Option<ParsedRegprocedureName>> {
    let mut in_quote = false;
    let left_parenthesis = input.bytes().enumerate().find_map(|(offset, byte)| {
        if byte == b'"' {
            in_quote = !in_quote;
            None
        } else if byte == b'(' && !in_quote {
            Some(offset)
        } else {
            None
        }
    });
    let Some(left_parenthesis) = left_parenthesis else {
        return parse_regobject_name(input)
            .map(|names| {
                Some(ParsedRegprocedureName {
                    names,
                    argument_types: None,
                })
            })
            .ok_or_else(|| regprocedure_syntax_error("expected a left parenthesis".into()));
    };
    let Some(names) = parse_regobject_name(&input[..left_parenthesis]) else {
        return Ok(None);
    };

    let bytes = input.as_bytes();
    let mut end = bytes.len();
    while end > left_parenthesis + 1 && scanner_isspace(bytes[end - 1]) {
        end -= 1;
    }
    if end <= left_parenthesis + 1 || bytes[end - 1] != b')' {
        return Err(regprocedure_syntax_error(format!(
            "expected a right parenthesis in routine identity \"{input}\""
        )));
    }
    let arguments = &input[left_parenthesis + 1..end - 1];
    let argument_types = parse_regprocedure_argument_types(input, arguments)?;

    Ok(Some(ParsedRegprocedureName {
        names,
        argument_types: Some(argument_types),
    }))
}

fn parse_regprocedure_argument_types(
    input: &str,
    arguments: &str,
) -> Result<Vec<ParsedRegtypeName>> {
    let argument_bytes = arguments.as_bytes();
    let mut argument_types = Vec::new();
    let mut offset = 0usize;
    let mut had_comma = false;
    loop {
        while argument_bytes
            .get(offset)
            .is_some_and(|byte| scanner_isspace(*byte))
        {
            offset += 1;
        }
        if offset == argument_bytes.len() {
            if had_comma {
                return Err(regprocedure_syntax_error(format!(
                    "expected a type name in routine identity \"{input}\""
                )));
            }
            break;
        }

        let start = offset;
        let mut quoted = false;
        let mut nesting = 0i32;
        while let Some(byte) = argument_bytes.get(offset).copied() {
            if byte == b'"' {
                quoted = !quoted;
            } else if byte == b',' && !quoted && nesting == 0 {
                break;
            } else if !quoted {
                match byte {
                    b'(' | b'[' => nesting += 1,
                    b')' | b']' => nesting -= 1,
                    _ => {}
                }
            }
            offset += 1;
        }
        if quoted || nesting != 0 {
            return Err(regprocedure_syntax_error(format!(
                "improper type name in routine identity \"{input}\""
            )));
        }
        let mut type_end = offset;
        while type_end > start && scanner_isspace(argument_bytes[type_end - 1]) {
            type_end -= 1;
        }
        let Some(type_name) = parse_regtype_name(&arguments[start..type_end])? else {
            return Err(SQLError::Parse(format!(
                "invalid type name in routine identity \"{input}\""
            )));
        };
        if argument_types.len() == POSTGRES_FUNCTION_MAX_ARGUMENTS {
            return Err(SQLError::Routine {
                sqlstate: "54023".into(),
                message: format!("too many arguments in routine identity \"{input}\""),
            });
        }
        argument_types.push(type_name);
        had_comma = argument_bytes.get(offset) == Some(&b',');
        if had_comma {
            offset += 1;
        }
    }

    Ok(argument_types)
}

fn regprocedure_syntax_error(message: String) -> SQLError {
    SQLError::Routine {
        sqlstate: "22P02".into(),
        message,
    }
}

pub(super) fn compile_foreign_key_action(raw: &str) -> Result<crate::ast::ForeignKeyAction> {
    use crate::ast::ForeignKeyAction;
    match raw.as_bytes().first().copied() {
        None | Some(0) | Some(b'a') => Ok(ForeignKeyAction::NoAction),
        Some(b'r') => Ok(ForeignKeyAction::Restrict),
        Some(b'c') => Ok(ForeignKeyAction::Cascade),
        Some(b'n') => Ok(ForeignKeyAction::SetNull),
        Some(b'd') => Ok(ForeignKeyAction::SetDefault),
        Some(other) => Err(SQLError::Unsupported(format!(
            "unsupported FOREIGN KEY action byte {other:?}"
        ))),
    }
}

pub(super) fn compile_foreign_key_match(raw: &str) -> Result<crate::ast::ForeignKeyMatch> {
    use crate::ast::ForeignKeyMatch;
    match raw.as_bytes().first().copied() {
        None | Some(0) | Some(b's') => Ok(ForeignKeyMatch::Simple),
        Some(b'f') => Ok(ForeignKeyMatch::Full),
        Some(b'p') => Err(SQLError::Unsupported(
            "FOREIGN KEY MATCH PARTIAL is not implemented by PostgreSQL".into(),
        )),
        Some(other) => Err(SQLError::Unsupported(format!(
            "unsupported FOREIGN KEY match byte {other:?}"
        ))),
    }
}

pub(super) fn validate_foreign_key_set_columns(
    local_columns: &[String],
    set_columns: &[String],
    raw_delete_action: &str,
) -> Result<()> {
    if set_columns.is_empty() {
        return Ok(());
    }
    let action = compile_foreign_key_action(raw_delete_action)?;
    if !matches!(
        action,
        crate::ast::ForeignKeyAction::SetNull | crate::ast::ForeignKeyAction::SetDefault
    ) {
        return Err(SQLError::Unsupported(
            "FOREIGN KEY column lists are only valid for ON DELETE SET NULL/DEFAULT".into(),
        ));
    }
    for col in set_columns {
        if !local_columns.iter().any(|local| local == col) {
            return Err(SQLError::Unsupported(format!(
                "FOREIGN KEY SET column `{col}` is not part of the local key"
            )));
        }
    }
    Ok(())
}

pub(super) fn raw_type_name(col: &pg_query::protobuf::ColumnDef) -> Result<Option<String>> {
    let Some(type_name) = col.type_name.as_ref() else {
        return Ok(None);
    };
    let names = type_name
        .names
        .iter()
        .map(extract_string)
        .collect::<Result<Vec<_>>>()?;
    Ok(names.last().map(|name| name.to_lowercase()))
}

pub(super) fn compile_type_name(col: &pg_query::protobuf::ColumnDef) -> Result<ColumnType> {
    let Some(type_name) = col.type_name.as_ref() else {
        return Err(SQLError::Internal(format!(
            "column `{}` has no type",
            col.colname
        )));
    };
    compile_pg_type_name(type_name, &col.colname)
}

#[expect(
    clippy::too_many_lines,
    reason = "ordered PostgreSQL lowering preserves syntax and error precedence"
)]
pub(super) fn compile_pg_type_name(
    type_name: &pg_query::protobuf::TypeName,
    column_name: &str,
) -> Result<ColumnType> {
    let names = type_name
        .names
        .iter()
        .map(extract_string)
        .collect::<Result<Vec<_>>>()?;
    let raw = names
        .last()
        .ok_or_else(|| {
            SQLError::Internal(format!(
                "type name for `{column_name}` has no name components"
            ))
        })?
        .to_lowercase();
    let bind_named = names
        .first()
        .is_some_and(|schema| names.len() > 1 && schema != "pg_catalog")
        || names.last().is_some_and(|name| *name != raw)
        || matches!(raw.as_str(), "integer" | "smallint" | "bigint" | "boolean");
    let base = if bind_named {
        compile_named_type(&names, type_name)
    } else {
        match raw.as_str() {
            "smallint" | "int2" | "smallserial" | "serial2" => Ok(ColumnType::SmallInteger),
            "int" | "int4" | "integer" | "serial" | "serial4" => Ok(ColumnType::Integer),
            "bigint" | "int8" | "bigserial" | "serial8" => Ok(ColumnType::BigInteger),
            "oid" => Ok(ColumnType::Oid),
            "xid" => Ok(ColumnType::Xid),
            "void" => Ok(ColumnType::Void),
            "text" => Ok(ColumnType::Text),
            "name" => Ok(ColumnType::Name),
            "uuid" => Ok(ColumnType::Uuid),
            "varchar" | "character varying" => {
                if type_name.typmods.len() > 1 {
                    return Err(SQLError::TypeMismatch(format!(
                        "CHARACTER VARYING accepts at most one length modifier, got {}",
                        type_name.typmods.len()
                    )));
                }
                let length = type_name
                    .typmods
                    .first()
                    .map(|node| expect_positive_character_length(node, "varchar"))
                    .transpose()?;
                Ok(ColumnType::Varchar(length))
            }
            "char" => Ok(ColumnType::InternalChar),
            "character" | "bpchar" => {
                if type_name.typmods.len() > 1 {
                    return Err(SQLError::TypeMismatch(format!(
                        "CHARACTER accepts at most one length modifier, got {}",
                        type_name.typmods.len()
                    )));
                }
                let length = type_name
                    .typmods
                    .first()
                    .map(|node| expect_positive_character_length(node, "bpchar"))
                    .transpose()?
                    .unwrap_or(1);
                Ok(ColumnType::Character(length))
            }
            "bool" | "boolean" => Ok(ColumnType::Boolean),
            "real" | "float4" => Ok(ColumnType::Real),
            "float8" | "double" | "double precision" => Ok(ColumnType::DoublePrecision),
            "numeric" | "decimal" => {
                if type_name.typmods.len() > 2 {
                    return Err(SQLError::TypeMismatch(format!(
                        "NUMERIC accepts at most precision and scale, got {} modifiers",
                        type_name.typmods.len()
                    )));
                }
                let mut typmods_iter = type_name.typmods.iter();
                let precision = typmods_iter
                    .next()
                    .map(|n| {
                        let value = expect_integer_const(n)?;
                        if !(1..=1000).contains(&value) {
                            return Err(SQLError::TypeMismatch(format!(
                                "NUMERIC precision must be between 1 and 1000, got {value}"
                            )));
                        }
                        Ok(value as u32)
                    })
                    .transpose()?;
                let scale = typmods_iter
                    .next()
                    .map(|n| {
                        let value = expect_integer_const(n)?;
                        if !(-1000..=1000).contains(&value) {
                            return Err(SQLError::TypeMismatch(format!(
                                "NUMERIC scale must be between -1000 and 1000, got {value}"
                            )));
                        }
                        Ok(value as i32)
                    })
                    .transpose()?;
                // PostgreSQL semantics: NUMERIC(precision) without an
                // explicit scale defaults to scale=0, rounding to integers.
                let scale = scale.or(precision.map(|_| 0));
                Ok(ColumnType::Numeric { precision, scale })
            }
            "date" => Ok(ColumnType::Date),
            "time" | "time without time zone" => Ok(ColumnType::Time),
            "timetz" | "time with time zone" => Ok(ColumnType::TimeTz),
            "timestamp" | "datetime" | "timestamp without time zone" => Ok(ColumnType::Timestamp),
            "timestamptz" | "timestamp with time zone" => Ok(ColumnType::TimestampTz),
            "interval" => {
                let fields = type_name
                    .typmods
                    .first()
                    .map(expect_integer_const)
                    .transpose()?
                    .unwrap_or(32767);
                let fields = IntervalFields::from_modifier_mask(fields)
                    .ok_or_else(|| SQLError::TypeMismatch("invalid interval fields".into()))?;
                let precision = type_name
                    .typmods
                    .get(1)
                    .map(expect_integer_const)
                    .transpose()?;
                ColumnType::with_interval_modifiers(fields, precision)
            }
            "int4range" => Ok(ColumnType::Range(RangeSubtype::Integer)),
            "int8range" => Ok(ColumnType::Range(RangeSubtype::BigInteger)),
            "numrange" => Ok(ColumnType::Range(RangeSubtype::Numeric)),
            "daterange" => Ok(ColumnType::Range(RangeSubtype::Date)),
            "tsrange" => Ok(ColumnType::Range(RangeSubtype::Timestamp)),
            "tstzrange" => Ok(ColumnType::Range(RangeSubtype::TimestampTz)),
            "int4multirange" => Ok(ColumnType::Multirange(RangeSubtype::Integer)),
            "int8multirange" => Ok(ColumnType::Multirange(RangeSubtype::BigInteger)),
            "nummultirange" => Ok(ColumnType::Multirange(RangeSubtype::Numeric)),
            "datemultirange" => Ok(ColumnType::Multirange(RangeSubtype::Date)),
            "tsmultirange" => Ok(ColumnType::Multirange(RangeSubtype::Timestamp)),
            "tstzmultirange" => Ok(ColumnType::Multirange(RangeSubtype::TimestampTz)),
            "json" => Ok(ColumnType::Json),
            "jsonb" => Ok(ColumnType::JsonB),
            "bytea" => Ok(ColumnType::Bytea),
            "regproc" => Ok(ColumnType::Regproc),
            "regprocedure" => Ok(ColumnType::Regprocedure),
            "regclass" => Ok(ColumnType::Regclass),
            "regnamespace" => Ok(ColumnType::Regnamespace),
            "regrole" => Ok(ColumnType::Regrole),
            "regtype" => Ok(ColumnType::Regtype),
            "pg_node_tree" => Ok(ColumnType::PgNodeTree),
            "aclitem" => Ok(ColumnType::AclItem),
            "int2vector" => Ok(ColumnType::Int2Vector),
            "oidvector" => Ok(ColumnType::OidVector),
            "anyarray" => Ok(ColumnType::AnyArray),
            "record" => Ok(ColumnType::Record),
            "vector" => {
                // VECTOR(N): the dimension is the only typmod argument.
                let [arg] = type_name.typmods.as_slice() else {
                    return Err(SQLError::Unsupported(
                        "VECTOR requires exactly one dimension".into(),
                    ));
                };
                let raw_dim = expect_integer_const(arg)?;
                let dim = u32::try_from(raw_dim).map_err(|_| {
                    SQLError::TypeMismatch(format!(
                        "VECTOR dimension must be between 1 and {}, got {raw_dim}",
                        u32::MAX
                    ))
                })?;
                if dim == 0 {
                    return Err(SQLError::TypeMismatch(
                        "VECTOR dimension must be greater than zero".into(),
                    ));
                }
                Ok(ColumnType::Vector(dim))
            }
            "tensor" => {
                // TENSOR(N): an array of N-dimensional vectors.
                let [arg] = type_name.typmods.as_slice() else {
                    return Err(SQLError::Unsupported(
                        "TENSOR requires exactly one dimension".into(),
                    ));
                };
                let raw_dim = expect_integer_const(arg)?;
                let dim = u32::try_from(raw_dim).map_err(|_| {
                    SQLError::TypeMismatch(format!(
                        "TENSOR dimension must be between 1 and {}, got {raw_dim}",
                        u32::MAX
                    ))
                })?;
                if dim == 0 {
                    return Err(SQLError::TypeMismatch(
                        "TENSOR dimension must be greater than zero".into(),
                    ));
                }
                Ok(ColumnType::Tensor(dim))
            }
            _ => compile_named_type(&names, type_name),
        }
    }?;
    let base = if matches!(
        base,
        ColumnType::Time | ColumnType::TimeTz | ColumnType::Timestamp | ColumnType::TimestampTz
    ) {
        if type_name.typmods.len() > 1 {
            return Err(SQLError::TypeMismatch(
                "invalid temporal type modifier".into(),
            ));
        }
        base.with_temporal_precision(
            type_name
                .typmods
                .first()
                .map(expect_integer_const)
                .transpose()?,
        )?
    } else {
        base
    };
    if matches!(base, ColumnType::Void) && !type_name.array_bounds.is_empty() {
        return Err(SQLError::Routine {
            sqlstate: "42704".into(),
            message: "type \"void[]\" does not exist".into(),
        });
    }
    Ok(type_name
        .array_bounds
        .iter()
        .fold(base, |element, _| ColumnType::Array(Box::new(element))))
}

/// Resolve a type reference without interpreting column-only serial pseudo-types.
pub(super) fn compile_pg_type_reference(
    type_name: &pg_query::protobuf::TypeName,
    context: &str,
) -> Result<ColumnType> {
    let names = type_name
        .names
        .iter()
        .map(extract_string)
        .collect::<Result<Vec<_>>>()?;
    if names.last().is_some_and(|name| {
        matches!(
            name.as_str(),
            "serial" | "serial2" | "serial4" | "serial8" | "smallserial" | "bigserial"
        )
    }) {
        let base = compile_named_type(&names, type_name)?;
        return Ok(type_name
            .array_bounds
            .iter()
            .fold(base, |element, _| ColumnType::Array(Box::new(element))));
    }
    compile_pg_type_name(type_name, context)
}

fn compile_named_type(
    names: &[String],
    type_name: &pg_query::protobuf::TypeName,
) -> Result<ColumnType> {
    let mut name = names
        .iter()
        .map(|name| format!("\"{}\"", name.replace('"', "\"\"")))
        .collect::<Vec<_>>()
        .join(".");
    if !type_name.typmods.is_empty() {
        let modifiers = type_name
            .typmods
            .iter()
            .map(expect_integer_const)
            .collect::<Result<Vec<_>>>()?;
        name.push('(');
        name.push_str(
            &modifiers
                .iter()
                .map(ToString::to_string)
                .collect::<Vec<_>>()
                .join(","),
        );
        name.push(')');
    }
    Ok(ColumnType::Named(name))
}

fn expect_positive_character_length(node: &Node, type_name: &str) -> Result<u32> {
    let length = expect_integer_const(node)?;
    u32::try_from(length)
        .ok()
        .filter(|length| *length > 0)
        .ok_or_else(|| SQLError::Routine {
            sqlstate: "22023".into(),
            message: format!("length for type {type_name} must be at least 1"),
        })
}

fn expect_integer_const(node: &Node) -> Result<i64> {
    let Some(inner) = node.node.as_ref() else {
        return Err(SQLError::Internal("missing const node".into()));
    };
    match inner {
        NodeEnum::AConst(c) => match &c.val {
            Some(pg_query::protobuf::a_const::Val::Ival(i)) => Ok(i64::from(i.ival)),
            Some(pg_query::protobuf::a_const::Val::Fval(f)) => {
                f.fval.parse::<i64>().map_err(|_| {
                    SQLError::TypeMismatch(format!(
                        "type modifier must be an integer, got `{}`",
                        f.fval
                    ))
                })
            }
            other => Err(SQLError::Internal(format!(
                "expected integer constant, got {other:?}"
            ))),
        },
        _ => Err(SQLError::Internal(format!(
            "expected A_Const, got {inner:?}"
        ))),
    }
}