use super::dispatch::compile_stmt;
use super::{
compile_expr, compile_qualified_name, extract_string, render_relation_component, Expr, Node,
NodeEnum, Result, SQLError, Statement,
};
struct CompiledFunctionTypeName {
name: String,
reference: Option<crate::ast::RoutineColumnTypeReference>,
}
fn compile_function_type_name(
t: &pg_query::protobuf::TypeName,
) -> Result<CompiledFunctionTypeName> {
let mut components = t
.names
.iter()
.map(extract_string)
.collect::<Result<Vec<_>>>()?;
if !t.pct_type
&& components
.first()
.is_some_and(|component| component.eq_ignore_ascii_case("pg_catalog"))
{
components.remove(0);
}
if components.is_empty() {
return Err(SQLError::Internal(
"function type has no name components".into(),
));
}
let reference = if t.pct_type {
let reference = match components.as_slice() {
[relation, column] => {
crate::ast::RoutineColumnTypeReference::new(None, relation.clone(), column.clone())
}
[schema, relation, column] => crate::ast::RoutineColumnTypeReference::new(
Some(schema.clone()),
relation.clone(),
column.clone(),
),
_ => {
return Err(SQLError::TypeMismatch(
"%TYPE requires a relation and column reference".into(),
))
}
};
Some(reference)
} else {
None
};
let mut name = components
.iter()
.map(|component| render_relation_component(component))
.collect::<Vec<_>>()
.join(".");
if !t.pct_type && t.array_bounds.is_empty() && components.len() == 1 {
if let Some(element) = crate::ast::builtin_array_element_name(&components[0]) {
name = format!("{element}[]");
}
}
if t.pct_type {
name.push_str("%type");
}
for _ in &t.array_bounds {
name.push_str("[]");
}
Ok(CompiledFunctionTypeName { name, reference })
}
fn def_elem_string(elem: &pg_query::protobuf::DefElem) -> Result<String> {
match elem.arg.as_ref().and_then(|a| a.node.as_ref()) {
Some(NodeEnum::String(s)) => Ok(s.sval.clone()),
other => Err(SQLError::TypeMismatch(format!(
"option `{}` expects a string, got {other:?}",
elem.defname
))),
}
}
pub(super) fn compile_create_function(
stmt: &pg_query::protobuf::CreateFunctionStmt,
) -> Result<crate::ast::CreateFunction> {
use crate::ast::{
CreateFunction, FunctionBody, FunctionParam, FunctionParamMode, FunctionReturns,
FunctionVolatility,
};
use pg_query::protobuf::FunctionParameterMode;
let keyword = if stmt.is_procedure {
"CREATE PROCEDURE"
} else {
"CREATE FUNCTION"
};
let name = compile_qualified_name(&stmt.funcname, keyword)?;
let mut params: Vec<FunctionParam> = Vec::with_capacity(stmt.parameters.len());
let mut has_table_param = false;
for p in &stmt.parameters {
let Some(NodeEnum::FunctionParameter(fp)) = p.node.as_ref() else {
return Err(SQLError::Internal(format!(
"{keyword}: malformed parameter"
)));
};
let mode = match fp.mode() {
FunctionParameterMode::FuncParamIn | FunctionParameterMode::FuncParamDefault => {
FunctionParamMode::In
}
FunctionParameterMode::FuncParamOut => FunctionParamMode::Out,
FunctionParameterMode::FuncParamInout => FunctionParamMode::InOut,
FunctionParameterMode::FuncParamTable => {
has_table_param = true;
FunctionParamMode::Table
}
FunctionParameterMode::FuncParamVariadic => {
return Err(SQLError::Unsupported(format!(
"{keyword}: VARIADIC parameters"
)));
}
FunctionParameterMode::Undefined => {
return Err(SQLError::Internal(format!(
"{keyword}: parameter mode missing"
)));
}
};
let compiled_type = fp
.arg_type
.as_ref()
.map(compile_function_type_name)
.transpose()?
.ok_or_else(|| SQLError::Internal(format!("{keyword}: parameter without type")))?;
let default = match fp.defexpr.as_ref() {
Some(node) => Some(compile_expr(node)?),
None => None,
};
params.push(FunctionParam {
name: fp.name.clone(),
type_name: compiled_type.name,
type_reference: compiled_type.reference,
mode,
default,
});
}
let mut saw_default = false;
for p in ¶ms {
if !matches!(p.mode, FunctionParamMode::In | FunctionParamMode::InOut) {
continue;
}
if p.default.is_some() {
saw_default = true;
} else if saw_default {
return Err(SQLError::Unsupported(
"input parameters after one with a default value must also have defaults".into(),
));
}
}
let (returns, return_type_reference) = if has_table_param {
(FunctionReturns::Table, None)
} else {
match stmt.return_type.as_ref() {
None => (FunctionReturns::None, None),
Some(t) => {
let compiled = compile_function_type_name(t)?;
let returns = if t.setof {
FunctionReturns::SetOf {
type_name: compiled.name,
}
} else {
FunctionReturns::Scalar {
type_name: compiled.name,
}
};
(returns, compiled.reference)
}
}
};
if params.iter().any(|param| param.type_name == "refcursor") {
return Err(SQLError::Unsupported(format!(
"{keyword}: refcursor parameters require session portal state"
)));
}
if matches!(
&returns,
FunctionReturns::Scalar { type_name } | FunctionReturns::SetOf { type_name }
if type_name == "refcursor"
) {
return Err(SQLError::Unsupported(format!(
"{keyword}: refcursor returns require session portal state"
)));
}
let mut language = String::new();
let mut volatility = FunctionVolatility::Volatile;
let mut strict = false;
let mut source: Option<String> = None;
for opt in &stmt.options {
let Some(NodeEnum::DefElem(elem)) = opt.node.as_ref() else {
return Err(SQLError::Internal(format!("{keyword}: malformed option")));
};
match elem.defname.to_ascii_lowercase().as_str() {
"language" => {
language = def_elem_string(elem)?.to_ascii_lowercase();
}
"volatility" => {
volatility = match def_elem_string(elem)?.as_str() {
"immutable" => FunctionVolatility::Immutable,
"stable" => FunctionVolatility::Stable,
"volatile" => FunctionVolatility::Volatile,
other => {
return Err(SQLError::TypeMismatch(format!(
"{keyword}: invalid volatility `{other}`"
)));
}
};
}
"strict" => {
strict = match elem.arg.as_ref().and_then(|a| a.node.as_ref()) {
Some(NodeEnum::Boolean(value)) => value.boolval,
other => {
return Err(SQLError::TypeMismatch(format!(
"{keyword}: STRICT expects a boolean, got {other:?}"
)));
}
};
}
"as" => {
let items: Vec<String> = match elem.arg.as_ref().and_then(|a| a.node.as_ref()) {
Some(NodeEnum::List(list)) => list
.items
.iter()
.map(extract_string)
.collect::<Result<Vec<_>>>()?,
Some(NodeEnum::String(s)) => vec![s.sval.clone()],
other => {
return Err(SQLError::TypeMismatch(format!(
"{keyword}: AS expects a string body, got {other:?}"
)));
}
};
match items.len() {
1 => source = items.into_iter().next(),
_ => {
return Err(SQLError::Unsupported(format!(
"{keyword}: AS 'obj_file', 'link_symbol' bodies"
)));
}
}
}
"window" => {
return Err(SQLError::Unsupported(format!(
"{keyword}: WINDOW functions"
)));
}
other => {
return Err(SQLError::Unsupported(format!(
"{keyword}: option `{other}` is not supported"
)));
}
}
}
let body = match (source, stmt.sql_body.as_deref()) {
(Some(src), None) => FunctionBody::Source(src),
(None, Some(node)) => FunctionBody::Statements(compile_sql_standard_body(node)?),
(Some(_), Some(_)) => {
return Err(SQLError::Unsupported(format!(
"{keyword}: both AS body and SQL-standard body"
)));
}
(None, None) => {
return Err(SQLError::Unsupported(format!(
"{keyword}: no function body"
)));
}
};
if language.is_empty() {
if matches!(body, FunctionBody::Statements(_)) {
language = "sql".into();
} else {
return Err(SQLError::Unsupported(format!(
"{keyword}: no language specified"
)));
}
}
Ok(CreateFunction {
name,
or_replace: stmt.replace,
is_procedure: stmt.is_procedure,
params,
returns,
return_type_reference,
language,
body,
volatility,
strict,
})
}
pub(super) fn compile_sql_standard_body(node: &Node) -> Result<Vec<Statement>> {
let Some(inner) = node.node.as_ref() else {
return Err(SQLError::Internal("empty SQL function body".into()));
};
match inner {
NodeEnum::ReturnStmt(ret) => {
let value = ret
.returnval
.as_deref()
.ok_or_else(|| SQLError::Internal("RETURN without a value".into()))?;
Ok(vec![select_of_expr(compile_expr(value)?)])
}
NodeEnum::List(list) => {
let mut out = Vec::with_capacity(list.items.len());
for item in &list.items {
let item_inner = item.node.as_ref().ok_or_else(|| {
SQLError::Internal("SQL function body contains an empty statement".into())
})?;
match item_inner {
NodeEnum::List(stmts) => {
for s in &stmts.items {
out.push(compile_stmt(s)?);
}
}
NodeEnum::ReturnStmt(ret) => {
let value = ret
.returnval
.as_deref()
.ok_or_else(|| SQLError::Internal("RETURN without a value".into()))?;
out.push(select_of_expr(compile_expr(value)?));
}
_ => out.push(compile_stmt(item)?),
}
}
Ok(out)
}
other => Err(SQLError::Unsupported(format!(
"SQL function body node {other:?}"
))),
}
}
fn select_of_expr(expr: Expr) -> Statement {
Statement::Select(Box::new(crate::ast::SelectStmt {
projections: vec![crate::ast::Projection { expr, alias: None }],
values: Vec::new(),
from: None,
r#where: None,
group_by: Vec::new(),
grouping_sets: Vec::new(),
having: None,
order_by: Vec::new(),
limit: None,
offset: None,
with: Vec::new(),
set_op: None,
distinct: false,
distinct_on: Vec::new(),
locking: Vec::new(),
}))
}
pub(super) fn compile_do(stmt: &pg_query::protobuf::DoStmt) -> Result<Statement> {
let mut language = "plpgsql".to_string();
let mut body: Option<String> = None;
for arg in &stmt.args {
let Some(NodeEnum::DefElem(elem)) = arg.node.as_ref() else {
return Err(SQLError::Internal("DO contains a malformed option".into()));
};
match elem.defname.to_ascii_lowercase().as_str() {
"as" => body = Some(def_elem_string(elem)?),
"language" => {
language = def_elem_string(elem)?.to_ascii_lowercase();
}
other => {
return Err(SQLError::Unsupported(format!(
"DO option `{other}` is not supported"
)));
}
}
}
let body = body.ok_or_else(|| SQLError::Internal("DO without a body".into()))?;
Ok(Statement::DoBlock { language, body })
}
pub(super) fn compile_call(stmt: &pg_query::protobuf::CallStmt) -> Result<Statement> {
let call = stmt
.funccall
.as_ref()
.ok_or_else(|| SQLError::Internal("CALL without a function".into()))?;
let name = compile_qualified_name(&call.funcname, "CALL")?;
let args = call
.args
.iter()
.map(compile_expr)
.collect::<Result<Vec<_>>>()?;
Ok(Statement::Call { name, args })
}
pub(super) fn compile_drop_function(
stmt: &pg_query::protobuf::DropStmt,
is_procedure: bool,
) -> Result<Statement> {
use crate::ast::{DropFunctionItem, DropFunctionStmt};
let mut items = Vec::new();
for object in &stmt.objects {
let Some(NodeEnum::ObjectWithArgs(owa)) = object.node.as_ref() else {
return Err(SQLError::Unsupported(
"DROP FUNCTION target is not a function signature".into(),
));
};
let name = compile_qualified_name(
&owa.objname,
if is_procedure {
"DROP PROCEDURE"
} else {
"DROP FUNCTION"
},
)?;
let arg_types = if owa.args_unspecified {
None
} else {
Some(
owa.objargs
.iter()
.map(|arg| match arg.node.as_ref() {
Some(NodeEnum::TypeName(t)) => {
compile_function_type_name(t).map(|compiled| compiled.name)
}
other => Err(SQLError::Unsupported(format!(
"DROP FUNCTION argument type node {other:?}"
))),
})
.collect::<Result<Vec<_>>>()?,
)
};
items.push(DropFunctionItem { name, arg_types });
}
if items.is_empty() {
return Err(SQLError::Internal("DROP FUNCTION without target".into()));
}
Ok(Statement::DropFunction(DropFunctionStmt {
is_procedure,
if_exists: stmt.missing_ok,
cascade: matches!(
stmt.behavior(),
pg_query::protobuf::DropBehavior::DropCascade
),
items,
}))
}