risingwave_frontend/handler/
create_function.rs

1// Copyright 2023 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use anyhow::Context;
16use either::Either;
17use risingwave_common::catalog::FunctionId;
18use risingwave_common::types::StructType;
19use risingwave_expr::sig::{CreateOptions, UdfKind};
20use risingwave_pb::catalog::PbFunction;
21use risingwave_pb::catalog::function::{Kind, ScalarFunction, TableFunction};
22
23use super::*;
24use crate::{Binder, bind_data_type};
25
26/// Non-SQL UDFs exchange data via Arrow, which does not support VARIANT yet.
27pub(crate) fn reject_variant_in_udf_signature(
28    return_type: &risingwave_common::types::DataType,
29    arg_types: &[risingwave_common::types::DataType],
30    kind: &str,
31) -> Result<()> {
32    if return_type.contains_variant() || arg_types.iter().any(|t| t.contains_variant()) {
33        return Err(ErrorCode::NotSupported(
34            format!("VARIANT type in {kind} signature"),
35            "VARIANT is not supported in UDFs yet".to_owned(),
36        )
37        .into());
38    }
39    Ok(())
40}
41
42pub async fn handle_create_function(
43    handler_args: HandlerArgs,
44    or_replace: bool,
45    temporary: bool,
46    if_not_exists: bool,
47    name: ObjectName,
48    args: Option<Vec<OperateFunctionArg>>,
49    returns: Option<CreateFunctionReturns>,
50    params: CreateFunctionBody,
51    with_options: CreateFunctionWithOptions,
52) -> Result<RwPgResponse> {
53    if or_replace {
54        bail_not_implemented!("CREATE OR REPLACE FUNCTION");
55    }
56    if temporary {
57        bail_not_implemented!("CREATE TEMPORARY FUNCTION");
58    }
59
60    let udf_config = handler_args.session.env().udf_config();
61
62    // e.g., `language [ python / javascript / ...etc]`
63    let language = match params.language {
64        Some(lang) => {
65            let lang = lang.real_value().to_lowercase();
66            match &*lang {
67                "java" => lang, // only support external UDF for Java
68                "python" if udf_config.enable_embedded_python_udf => lang,
69                "javascript" if udf_config.enable_embedded_javascript_udf => lang,
70                "rust" | "wasm" if udf_config.enable_embedded_wasm_udf => lang,
71                "python" | "javascript" | "rust" | "wasm" => {
72                    return Err(ErrorCode::InvalidParameterValue(format!(
73                        "{} UDF is not enabled in configuration",
74                        lang
75                    ))
76                    .into());
77                }
78                _ => {
79                    return Err(ErrorCode::InvalidParameterValue(format!(
80                        "language {} is not supported",
81                        lang
82                    ))
83                    .into());
84                }
85            }
86        }
87        // Empty language is acceptable since we only require the external server implements the
88        // correct protocol.
89        None => "".to_owned(),
90    };
91
92    let runtime = match params.runtime {
93        Some(_) => {
94            return Err(ErrorCode::InvalidParameterValue(
95                "runtime selection is currently not supported".to_owned(),
96            )
97            .into());
98        }
99        None => None,
100    };
101
102    let return_type;
103    let kind = match returns {
104        Some(CreateFunctionReturns::Value(data_type)) => {
105            return_type = bind_data_type(&data_type)?;
106            Kind::Scalar(ScalarFunction {})
107        }
108        Some(CreateFunctionReturns::Table(columns)) => {
109            if columns.len() == 1 {
110                // return type is the original type for single column
111                return_type = bind_data_type(&columns[0].data_type)?;
112            } else {
113                // return type is a struct for multiple columns
114                let it = columns
115                    .into_iter()
116                    .map(|c| bind_data_type(&c.data_type).map(|ty| (c.name.real_value(), ty)));
117                let fields = it.try_collect::<_, Vec<_>, _>()?;
118                return_type = StructType::new(fields).into();
119            }
120            Kind::Table(TableFunction {})
121        }
122        None => {
123            return Err(ErrorCode::InvalidParameterValue(
124                "return type must be specified".to_owned(),
125            )
126            .into());
127        }
128    };
129
130    let mut arg_names = vec![];
131    let mut arg_types = vec![];
132    for arg in args.unwrap_or_default() {
133        arg_names.push(arg.name.map_or("".to_owned(), |n| n.real_value()));
134        arg_types.push(bind_data_type(&arg.data_type)?);
135    }
136
137    reject_variant_in_udf_signature(&return_type, &arg_types, "function")?;
138
139    // resolve database and schema id
140    let session = &handler_args.session;
141    let db_name = &session.database();
142    let (schema_name, function_name) = Binder::resolve_schema_qualified_name(db_name, &name)?;
143    let (database_id, schema_id) = session.get_database_and_schema_id_for_create(schema_name)?;
144
145    // check if the function exists in the catalog
146    if let Either::Right(resp) = session.check_function_name_duplicated(
147        StatementType::CREATE_FUNCTION,
148        name,
149        &arg_types,
150        if_not_exists,
151    )? {
152        return Ok(resp);
153    }
154
155    let link = match &params.using {
156        Some(CreateFunctionUsing::Link(l)) => Some(l.as_str()),
157        _ => None,
158    };
159    let base64_decoded = match &params.using {
160        Some(CreateFunctionUsing::Base64(encoded)) => {
161            use base64::prelude::{BASE64_STANDARD, Engine};
162            let bytes = BASE64_STANDARD
163                .decode(encoded)
164                .context("invalid base64 encoding")?;
165            Some(bytes)
166        }
167        _ => None,
168    };
169
170    let create_fn =
171        risingwave_expr::sig::find_udf_impl(&language, runtime.as_deref(), link)?.create_fn;
172    let output = create_fn(CreateOptions {
173        kind: match kind {
174            Kind::Scalar(_) => UdfKind::Scalar,
175            Kind::Table(_) => UdfKind::Table,
176            Kind::Aggregate(_) => unreachable!(),
177        },
178        name: &function_name,
179        arg_names: &arg_names,
180        arg_types: &arg_types,
181        return_type: &return_type,
182        as_: params.as_.as_ref().map(|s| s.as_str()),
183        using_link: link,
184        using_base64_decoded: base64_decoded.as_deref(),
185    })?;
186
187    let function = PbFunction {
188        id: FunctionId::placeholder(),
189        schema_id,
190        database_id,
191        name: function_name,
192        kind: Some(kind),
193        arg_names,
194        arg_types: arg_types.into_iter().map(|t| t.into()).collect(),
195        return_type: Some(return_type.into()),
196        language,
197        runtime,
198        name_in_runtime: Some(output.name_in_runtime),
199        link: link.map(|s| s.to_owned()),
200        body: output.body,
201        compressed_binary: output.compressed_binary,
202        owner: session.user_id(),
203        always_retry_on_network_error: with_options
204            .always_retry_on_network_error
205            .unwrap_or_default(),
206        is_async: with_options.r#async,
207        is_batched: with_options.batch,
208        created_at_epoch: None,
209        created_at_cluster_version: None,
210    };
211
212    let catalog_writer = session.catalog_writer()?;
213    catalog_writer.create_function(function).await?;
214
215    Ok(PgResponse::empty_result(StatementType::CREATE_FUNCTION))
216}
risingwave_frontend/handler/create_function.rs

risingwave_frontend/handler/
create_function.rs
