risingwave_frontend/handler/
create_aggregate.rs

1// Copyright 2024 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_expr::sig::{CreateOptions, UdfKind};
19use risingwave_pb::catalog::Function;
20use risingwave_pb::catalog::function::{AggregateFunction, Kind};
21use risingwave_sqlparser::ast::DataType as AstDataType;
22
23use super::*;
24use crate::handler::create_function::reject_variant_in_udf_signature;
25use crate::{Binder, bind_data_type};
26
27pub async fn handle_create_aggregate(
28    handler_args: HandlerArgs,
29    or_replace: bool,
30    if_not_exists: bool,
31    name: ObjectName,
32    args: Vec<OperateFunctionArg>,
33    returns: AstDataType,
34    params: CreateFunctionBody,
35) -> Result<RwPgResponse> {
36    if or_replace {
37        bail_not_implemented!("CREATE OR REPLACE AGGREGATE");
38    }
39
40    let udf_config = handler_args.session.env().udf_config();
41
42    // e.g., `language [ python / java / ...etc]`
43    let language = match params.language {
44        Some(lang) => {
45            let lang = lang.real_value().to_lowercase();
46            match &*lang {
47                "python" if udf_config.enable_embedded_python_udf => lang,
48                "javascript" if udf_config.enable_embedded_javascript_udf => lang,
49                "python" | "javascript" => {
50                    return Err(ErrorCode::InvalidParameterValue(format!(
51                        "{} UDF is not enabled in configuration",
52                        lang
53                    ))
54                    .into());
55                }
56                _ => {
57                    return Err(ErrorCode::InvalidParameterValue(format!(
58                        "language {} is not supported",
59                        lang
60                    ))
61                    .into());
62                }
63            }
64        }
65        None => return Err(ErrorCode::InvalidParameterValue("no language".into()).into()),
66    };
67
68    let runtime = match params.runtime {
69        Some(_) => {
70            return Err(ErrorCode::InvalidParameterValue(
71                "runtime selection is currently not supported".to_owned(),
72            )
73            .into());
74        }
75        None => None,
76    };
77
78    let return_type = bind_data_type(&returns)?;
79
80    let mut arg_names = vec![];
81    let mut arg_types = vec![];
82    for arg in args {
83        arg_names.push(arg.name.map_or("".to_owned(), |n| n.real_value()));
84        arg_types.push(bind_data_type(&arg.data_type)?);
85    }
86
87    reject_variant_in_udf_signature(&return_type, &arg_types, "aggregate function")?;
88
89    // resolve database and schema id
90    let session = &handler_args.session;
91    let db_name = &session.database();
92    let (schema_name, function_name) = Binder::resolve_schema_qualified_name(db_name, &name)?;
93    let (database_id, schema_id) = session.get_database_and_schema_id_for_create(schema_name)?;
94
95    // check if the function exists in the catalog
96    if let Either::Right(resp) = session.check_function_name_duplicated(
97        StatementType::CREATE_FUNCTION,
98        name,
99        &arg_types,
100        if_not_exists,
101    )? {
102        return Ok(resp);
103    }
104
105    let link = match &params.using {
106        Some(CreateFunctionUsing::Link(l)) => Some(l.as_str()),
107        _ => None,
108    };
109    let base64_decoded = match &params.using {
110        Some(CreateFunctionUsing::Base64(encoded)) => {
111            use base64::prelude::{BASE64_STANDARD, Engine};
112            let bytes = BASE64_STANDARD
113                .decode(encoded)
114                .context("invalid base64 encoding")?;
115            Some(bytes)
116        }
117        _ => None,
118    };
119
120    let create_fn = risingwave_expr::sig::find_udf_impl(&language, None, link)?.create_fn;
121    let output = create_fn(CreateOptions {
122        kind: UdfKind::Aggregate,
123        name: &function_name,
124        arg_names: &arg_names,
125        arg_types: &arg_types,
126        return_type: &return_type,
127        as_: params.as_.as_ref().map(|s| s.as_str()),
128        using_link: link,
129        using_base64_decoded: base64_decoded.as_deref(),
130    })?;
131
132    let function = Function {
133        id: FunctionId::placeholder(),
134        schema_id,
135        database_id,
136        name: function_name,
137        kind: Some(Kind::Aggregate(AggregateFunction {})),
138        arg_names,
139        arg_types: arg_types.into_iter().map(|t| t.into()).collect(),
140        return_type: Some(return_type.into()),
141        language,
142        runtime,
143        name_in_runtime: Some(output.name_in_runtime),
144        link: link.map(|s| s.to_owned()),
145        body: output.body,
146        compressed_binary: output.compressed_binary,
147        owner: session.user_id(),
148        always_retry_on_network_error: false,
149        is_async: None,
150        is_batched: None,
151        created_at_epoch: None,
152        created_at_cluster_version: None,
153    };
154
155    let catalog_writer = session.catalog_writer()?;
156    catalog_writer.create_function(function).await?;
157
158    Ok(PgResponse::empty_result(StatementType::CREATE_AGGREGATE))
159}
risingwave_frontend/handler/create_aggregate.rs

risingwave_frontend/handler/
create_aggregate.rs
