Commit 14aa185

use crate::expr::{Exists, InSubquery, SetComparison, Sort};

use datafusion_common::tree_node::{

Transformed, TreeNode, TreeNodeContainer, TreeNodeRecursion,

Constraints, DFSchema, DFSchemaRef, Result, SchemaReference, TableReference, plan_err,

impl CreateFunction {

/// Creates a function definition, validating placeholders in its body and

/// argument defaults, including expressions inside subqueries.

///

/// Placeholders must be positional (`$1`, `$2`, ...) and reference a

/// declared argument. SQL argument names must be resolved before calling

/// this constructor.

pub fn try_new(

or_replace: bool,

temporary: bool,

name: String,

args: Option<Vec<OperateFunctionArg>>,

return_type: Option<DataType>,

params: CreateFunctionBody,

) -> Result<Self> {

let arg_count = args.as_ref().map_or(0, |declared| declared.len());

for default_expr in args

.iter()

.flatten()

.filter_map(|arg| arg.default_expr.as_ref())

{

validate_function_body_placeholders(default_expr, arg_count)?;

}

if let Some(body) = &params.function_body {

validate_function_body_placeholders(body, arg_count)?;

}

Ok(Self {

or_replace,

temporary,

name,

args,

return_type,

params,

schema: DFSchemaRef::new(DFSchema::empty()),

})

}

fn validate_function_body_placeholders(expr: &Expr, arg_count: usize) -> Result<()> {

expr.apply(|expr| {

if let Expr::Placeholder(placeholder) = expr {

match placeholder

.id

.strip_prefix('$')

.and_then(|id| id.parse::<usize>().ok())

{

// In range, e.g. `$2` with two declared arguments

Some(idx) if (1..=arg_count).contains(&idx) => {}

// Out of range, e.g. `$3` with two declared arguments

Some(_) => {

return plan_err!(

"Invalid placeholder, out of range: {}",

placeholder.id

);

}

// Named or malformed placeholder

None => {

return plan_err!("Unknown placeholder: {}", placeholder.id);

}

}

}

// `Expr::apply` does not descend into subqueries, so walk their

// plans explicitly to validate placeholders inside them

if let Some(subquery) = match expr {

Expr::ScalarSubquery(subquery) => Some(&subquery.subquery),

Expr::Exists(Exists { subquery, .. }) => Some(&subquery.subquery),

Expr::InSubquery(InSubquery { subquery, .. }) => Some(&subquery.subquery),

Expr::SetComparison(SetComparison { subquery, .. }) => {

Some(&subquery.subquery)

}

_ => None,

} {

// Nested subqueries are handled by this expression visitor

subquery.apply(|plan| {

plan.apply_expressions(|e| {

validate_function_body_placeholders(e, arg_count)?;

Ok(TreeNodeRecursion::Continue)

})

})?;

}

Ok(TreeNodeRecursion::Continue)

})?;

Ok(())

use super::{CreateFunction, CreateFunctionBody, OperateFunctionArg};

use crate::expr_fn::placeholder;

CreateCatalog, CreateExternalCatalog, DdlStatement, DropCatalog, DropView, Expr,

LogicalPlanBuilder, Subquery,

use arrow::datatypes::DataType;

use std::{cmp::Ordering, collections::HashMap, sync::Arc};

#[test]

fn create_function_try_new_validates_placeholders() {

let create = |body, default_expr| {

CreateFunction::try_new(

false,

false,

"f".to_string(),

Some(vec![OperateFunctionArg {

name: None,

data_type: DataType::Int64,

default_expr,

}]),

Some(DataType::Int64),

CreateFunctionBody {

language: None,

behavior: None,

function_body: body,

},

)

};

assert!(create(Some(placeholder("$1")), Some(placeholder("$1"))).is_ok());

assert!(create(None, None).is_ok());

for (body, default_expr) in [

(Some(placeholder("$2")), None),

(None, Some(placeholder("$2"))),

] {

let err = create(body, default_expr).unwrap_err();

assert_eq!(

err.strip_backtrace(),

"Error during planning: Invalid placeholder, out of range: $2"

);

}

let mut nested_body = placeholder("$2");

for _ in 0..16 {

let plan = LogicalPlanBuilder::empty(true)

.project(vec![nested_body.alias("v")])

.unwrap()

.build()

.unwrap();

nested_body = Expr::ScalarSubquery(Subquery {

subquery: Arc::new(plan),

outer_ref_columns: vec![],

spans: Default::default(),

});

}

let err = create(Some(nested_body), None).unwrap_err();

assert_eq!(

err.strip_backtrace(),

"Error during planning: Invalid placeholder, out of range: $2"

);

}

// CreateFunction::try_new validates function placeholders; types stay

// permissive so PREPARE can defer inference to bind time.

let statement =

DdlStatement::CreateFunction(Box::new(CreateFunction::try_new(

or_replace,

temporary,

name,

args,

return_type.map(|f| f.data_type().clone()),

params,

)?));

CREATE FUNCTION foo (DOUBLE) RETURNS DOUBLE RETURN $1;

# Invalid placeholders are rejected during planning, before a function

# factory is consulted (see #25038)

statement error DataFusion error: Error during planning: Invalid placeholder, out of range: \$2

# a zero-argument function cannot reference arguments either

statement error DataFusion error: Error during planning: Invalid placeholder, out of range: \$1

CREATE FUNCTION foo_zero () RETURNS DOUBLE RETURN $1;

# placeholders inside subqueries are validated too

statement error DataFusion error: Error during planning: Invalid placeholder, out of range: \$9

CREATE FUNCTION foo_subquery (DOUBLE) RETURNS DOUBLE RETURN $1 + (SELECT $9);

# placeholders in argument default expressions are validated too

statement error DataFusion error: Error during planning: Invalid placeholder, out of range: \$9

CREATE FUNCTION foo_default (a DOUBLE, b DOUBLE DEFAULT $9) RETURNS DOUBLE RETURN $1 + $2;

# out of range placeholder in a subquery's LIMIT clause

statement error DataFusion error: Error during planning: Invalid placeholder, out of range: \$9

CREATE FUNCTION foo_limit (DOUBLE) RETURNS BOOLEAN RETURN $1 IN (SELECT v FROM (VALUES (1.0), (2.0)) AS t(v) ORDER BY v LIMIT $9);

# Named arguments are resolved before placeholder validation

statement error DataFusion error: Error during planning: Invalid placeholder, out of range: \$3

CREATE FUNCTION foo_named (a DOUBLE, b DOUBLE) RETURNS DOUBLE RETURN $a + $3;

statement error DataFusion error: Error during planning: Unknown placeholder: \$a

CREATE FUNCTION foo_unknown () RETURNS DOUBLE RETURN $a;

# Defaults use positional placeholders

statement error DataFusion error: Error during planning: Unknown placeholder: \$a

CREATE FUNCTION foo_named_default (a DOUBLE, b DOUBLE DEFAULT $a) RETURNS DOUBLE RETURN $1 + $2;

statement error DataFusion error: Error during planning: Invalid placeholder, out of range: \$9

CREATE FUNCTION foo_exists (DOUBLE) RETURNS BOOLEAN RETURN EXISTS (SELECT $9);

statement error DataFusion error: Error during planning: Invalid placeholder, out of range: \$9

CREATE FUNCTION foo_in (DOUBLE) RETURNS BOOLEAN RETURN $1 IN (SELECT $9);

statement error DataFusion error: Error during planning: Invalid placeholder, out of range: \$9

CREATE FUNCTION foo_any (DOUBLE) RETURNS BOOLEAN RETURN $1 = ANY (SELECT $9);

statement error DataFusion error: Error during planning: Invalid placeholder, out of range: \$9

CREATE FUNCTION foo_nested (DOUBLE) RETURNS BOOLEAN RETURN EXISTS (SELECT 1 WHERE (SELECT $9) > 0.0);

# Valid definitions reach the function factory

statement error DataFusion error: Invalid or Unsupported Configuration: Function factory has not been configured

CREATE FUNCTION foo_valid (DOUBLE, DOUBLE) RETURNS DOUBLE RETURN $1 + $2;

statement error DataFusion error: Invalid or Unsupported Configuration: Function factory has not been configured

CREATE FUNCTION foo_valid_default (a DOUBLE, b DOUBLE DEFAULT $1) RETURNS DOUBLE RETURN $1 + $2;

statement error DataFusion error: Invalid or Unsupported Configuration: Function factory has not been configured

CREATE FUNCTION foo_valid_subquery (DOUBLE) RETURNS BOOLEAN RETURN $1 IN (SELECT v FROM (VALUES (1.0)) AS t(v) WHERE v > $1);