risingwave_frontend/expr/
table_function.rs

1// Copyright 2022 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 std::sync::Arc;
16
17use anyhow::Context;
18use itertools::Itertools;
19use mysql_async::consts::ColumnType as MySqlColumnType;
20use mysql_async::prelude::*;
21use risingwave_common::array::arrow::IcebergArrowConvert;
22use risingwave_common::types::{DataType, ScalarImpl, StructType};
23use risingwave_connector::connector_common::sql_server::{
24    MssqlConnectionConfig, describe_mssql_query,
25};
26use risingwave_connector::connector_common::{PgConnectionConfig, create_pg_client};
27use risingwave_connector::source::iceberg::{
28    FileScanBackend, extract_bucket_and_file_name, get_parquet_fields, list_data_directory,
29    new_azblob_operator, new_gcs_operator, new_s3_operator,
30};
31use risingwave_pb::expr::PbTableFunction;
32pub use risingwave_pb::expr::table_function::PbType as TableFunctionType;
33use tokio_postgres::types::Type as TokioPgType;
34
35use super::{Expr, ExprImpl, ExprRewriter, Literal, RwResult, infer_type};
36use crate::catalog::function_catalog::{FunctionCatalog, FunctionKind};
37use crate::error::ErrorCode::BindError;
38use crate::expr::reject_impure;
39use crate::utils::FRONTEND_RUNTIME;
40
41pub(crate) const MSSQL_QUERY_ARG_COUNT: usize = 8;
42const MSSQL_QUERY_ENCRYPT_ARG_INDEX: usize = 6;
43const MSSQL_QUERY_TRUST_CERT_ARG_INDEX: usize = 7;
44
45/// A table function takes a row as input and returns a table. It is also known as Set-Returning
46/// Function.
47///
48/// See also [`TableFunction`](risingwave_expr::table_function::TableFunction) trait in expr crate
49/// and [`ProjectSetSelectItem`](risingwave_pb::expr::ProjectSetSelectItem).
50#[derive(Clone, Eq, PartialEq, Hash)]
51pub struct TableFunction {
52    pub args: Vec<ExprImpl>,
53    pub return_type: DataType,
54    pub function_type: TableFunctionType,
55    /// Catalog of user defined table function.
56    pub user_defined: Option<Arc<FunctionCatalog>>,
57}
58
59impl TableFunction {
60    /// Create a `TableFunction` expr with the return type inferred from `func_type` and types of
61    /// `inputs`.
62    pub fn new(func_type: TableFunctionType, mut args: Vec<ExprImpl>) -> RwResult<Self> {
63        let return_type = infer_type(func_type.into(), &mut args)?;
64        Ok(TableFunction {
65            args,
66            return_type,
67            function_type: func_type,
68            user_defined: None,
69        })
70    }
71
72    /// Create a user-defined `TableFunction`.
73    pub fn new_user_defined(catalog: Arc<FunctionCatalog>, args: Vec<ExprImpl>) -> Self {
74        let FunctionKind::Table = &catalog.kind else {
75            panic!("not a table function");
76        };
77        TableFunction {
78            args,
79            return_type: catalog.return_type.clone(),
80            function_type: TableFunctionType::UserDefined,
81            user_defined: Some(catalog),
82        }
83    }
84
85    /// A special table function which would be transformed into `LogicalFileScan` by `TableFunctionToFileScanRule` in the optimizer.
86    /// select * from `file_scan`('parquet', 's3', region, ak, sk, location)
87    pub fn new_file_scan(mut args: Vec<ExprImpl>) -> RwResult<Self> {
88        let return_type = {
89            // arguments:
90            // file format e.g. parquet
91            // storage type e.g. s3, gcs, azblob
92            // For s3: file_scan('parquet', 's3', s3_region, s3_access_key, s3_secret_key, file_location_or_directory)
93            // For gcs: file_scan('parquet', 'gcs', credential, file_location_or_directory)
94            // For azblob: file_scan('parquet', 'azblob', endpoint, account_name, account_key, file_location)
95            let mut eval_args: Vec<String> = vec![];
96            for arg in &args {
97                if arg.return_type() != DataType::Varchar {
98                    return Err(BindError(
99                        "file_scan function only accepts string arguments".to_owned(),
100                    )
101                    .into());
102                }
103                match arg.try_fold_const() {
104                    Some(Ok(value)) => {
105                        if value.is_none() {
106                            return Err(BindError(
107                                "file_scan function does not accept null arguments".to_owned(),
108                            )
109                            .into());
110                        }
111                        match value {
112                            Some(ScalarImpl::Utf8(s)) => {
113                                eval_args.push(s.to_string());
114                            }
115                            _ => {
116                                return Err(BindError(
117                                    "file_scan function only accepts string arguments".to_owned(),
118                                )
119                                .into());
120                            }
121                        }
122                    }
123                    Some(Err(err)) => {
124                        return Err(err);
125                    }
126                    None => {
127                        return Err(BindError(
128                            "file_scan function only accepts constant arguments".to_owned(),
129                        )
130                        .into());
131                    }
132                }
133            }
134
135            if (eval_args.len() != 4 && eval_args.len() != 6)
136                || (eval_args.len() == 4 && !"gcs".eq_ignore_ascii_case(&eval_args[1]))
137                || (eval_args.len() == 6
138                    && !"s3".eq_ignore_ascii_case(&eval_args[1])
139                    && !"azblob".eq_ignore_ascii_case(&eval_args[1]))
140            {
141                return Err(BindError(
142                "file_scan function supports three backends: s3, gcs, and azblob. Their formats are as follows: \n
143                    file_scan('parquet', 's3', s3_region, s3_access_key, s3_secret_key, file_location) \n
144                    file_scan('parquet', 'gcs', credential, service_account, file_location) \n
145                    file_scan('parquet', 'azblob', endpoint, account_name, account_key, file_location)"
146                        .to_owned(),
147                )
148                .into());
149            }
150            if !"parquet".eq_ignore_ascii_case(&eval_args[0]) {
151                return Err(BindError(
152                    "file_scan function only accepts 'parquet' as file format".to_owned(),
153                )
154                .into());
155            }
156
157            if !"s3".eq_ignore_ascii_case(&eval_args[1])
158                && !"gcs".eq_ignore_ascii_case(&eval_args[1])
159                && !"azblob".eq_ignore_ascii_case(&eval_args[1])
160            {
161                return Err(BindError(
162                    "file_scan function only accepts 's3', 'gcs' or 'azblob' as storage type"
163                        .to_owned(),
164                )
165                .into());
166            }
167
168            #[cfg(madsim)]
169            return Err(crate::error::ErrorCode::BindError(
170                "file_scan can't be used in the madsim mode".to_string(),
171            )
172            .into());
173
174            #[cfg(not(madsim))]
175            {
176                let (file_scan_backend, input_file_location) =
177                    if "s3".eq_ignore_ascii_case(&eval_args[1]) {
178                        (FileScanBackend::S3, eval_args[5].clone())
179                    } else if "gcs".eq_ignore_ascii_case(&eval_args[1]) {
180                        (FileScanBackend::Gcs, eval_args[3].clone())
181                    } else if "azblob".eq_ignore_ascii_case(&eval_args[1]) {
182                        (FileScanBackend::Azblob, eval_args[5].clone())
183                    } else {
184                        unreachable!();
185                    };
186                let op = match file_scan_backend {
187                    FileScanBackend::S3 => {
188                        let (bucket, _) = extract_bucket_and_file_name(
189                            &eval_args[5].clone(),
190                            &file_scan_backend,
191                        )?;
192
193                        let (s3_region, s3_endpoint) = match eval_args[2].starts_with("http") {
194                            true => ("us-east-1".to_owned(), eval_args[2].clone()), /* for minio, hard code region as not used but needed. */
195                            false => (
196                                eval_args[2].clone(),
197                                format!("https://{}.s3.{}.amazonaws.com", bucket, eval_args[2],),
198                            ),
199                        };
200                        new_s3_operator(
201                            s3_region,
202                            eval_args[3].clone(),
203                            eval_args[4].clone(),
204                            bucket,
205                            s3_endpoint,
206                        )?
207                    }
208                    FileScanBackend::Gcs => {
209                        let (bucket, _) =
210                            extract_bucket_and_file_name(&input_file_location, &file_scan_backend)?;
211
212                        new_gcs_operator(eval_args[2].clone(), bucket)?
213                    }
214                    FileScanBackend::Azblob => {
215                        let (bucket, _) =
216                            extract_bucket_and_file_name(&input_file_location, &file_scan_backend)?;
217
218                        new_azblob_operator(
219                            eval_args[2].clone(),
220                            eval_args[3].clone(),
221                            eval_args[4].clone(),
222                            bucket,
223                        )?
224                    }
225                };
226                let files = if input_file_location.ends_with('/') {
227                    let files = tokio::task::block_in_place(|| {
228                        FRONTEND_RUNTIME.block_on(async {
229                            let files = list_data_directory(
230                                op.clone(),
231                                input_file_location.clone(),
232                                &file_scan_backend,
233                            )
234                            .await?;
235
236                            Ok::<Vec<String>, anyhow::Error>(files)
237                        })
238                    })?;
239                    if files.is_empty() {
240                        return Err(BindError(
241                            "file_scan function only accepts non-empty directory".to_owned(),
242                        )
243                        .into());
244                    }
245
246                    Some(files)
247                } else {
248                    None
249                };
250                let schema = tokio::task::block_in_place(|| {
251                    FRONTEND_RUNTIME.block_on(async {
252                        let location = match files.as_ref() {
253                            Some(files) => files[0].clone(),
254                            None => input_file_location.clone(),
255                        };
256                        let (_, file_name) =
257                            extract_bucket_and_file_name(&location, &file_scan_backend)?;
258
259                        let fields = get_parquet_fields(op, file_name).await?;
260
261                        let mut rw_types = vec![];
262                        for field in &fields {
263                            rw_types.push((
264                                field.name().clone(),
265                                IcebergArrowConvert.type_from_field(field)?,
266                            ));
267                        }
268
269                        Ok::<risingwave_common::types::DataType, anyhow::Error>(DataType::Struct(
270                            StructType::new(rw_types),
271                        ))
272                    })
273                })?;
274
275                if let Some(files) = files {
276                    // if the file location is a directory, we need to remove the last argument and add all files in the directory as arguments
277                    match file_scan_backend {
278                        FileScanBackend::S3 => args.remove(5),
279                        FileScanBackend::Gcs => args.remove(3),
280                        FileScanBackend::Azblob => args.remove(5),
281                    };
282                    for file in files {
283                        args.push(ExprImpl::Literal(Box::new(Literal::new(
284                            Some(ScalarImpl::Utf8(file.into())),
285                            DataType::Varchar,
286                        ))));
287                    }
288                }
289
290                schema
291            }
292        };
293
294        Ok(TableFunction {
295            args,
296            return_type,
297            function_type: TableFunctionType::FileScan,
298            user_defined: None,
299        })
300    }
301
302    pub fn new_postgres_query(args: Vec<ExprImpl>) -> RwResult<Self> {
303        let evaled_args = args
304            .iter()
305            .map(expr_impl_to_string_fn)
306            .collect::<RwResult<Vec<_>>>()?;
307
308        #[cfg(madsim)]
309        {
310            return Err(crate::error::ErrorCode::BindError(
311                "postgres_query can't be used in the madsim mode".to_string(),
312            )
313            .into());
314        }
315
316        #[cfg(not(madsim))]
317        {
318            let schema = tokio::task::block_in_place(|| {
319                FRONTEND_RUNTIME.block_on(async {
320                    let ssl_mode = evaled_args
321                        .get(6)
322                        .and_then(|s| s.parse().ok())
323                        .unwrap_or_default();
324
325                    let ssl_root_cert = evaled_args
326                        .get(7)
327                        .and_then(|s| if s.is_empty() { None } else { Some(s.clone()) });
328
329                    let port = evaled_args[1]
330                        .parse::<u16>()
331                        .with_context(|| format!("invalid postgres port `{}`", evaled_args[1]))?;
332                    let pg_conn = PgConnectionConfig {
333                        host: evaled_args[0].clone(),
334                        port,
335                        user: evaled_args[2].clone(),
336                        password: evaled_args[3].clone(),
337                        database: evaled_args[4].clone(),
338                        ssl_mode,
339                        ssl_root_cert,
340                    };
341                    let client = create_pg_client(&pg_conn, None, None).await?;
342
343                    let statement = client.prepare(evaled_args[5].as_str()).await?;
344
345                    let mut rw_types = vec![];
346                    for column in statement.columns() {
347                        let name = column.name().to_owned();
348                        let data_type = match *column.type_() {
349                            TokioPgType::BOOL => DataType::Boolean,
350                            TokioPgType::INT2 => DataType::Int16,
351                            TokioPgType::INT4 => DataType::Int32,
352                            TokioPgType::INT8 => DataType::Int64,
353                            TokioPgType::FLOAT4 => DataType::Float32,
354                            TokioPgType::FLOAT8 => DataType::Float64,
355                            TokioPgType::NUMERIC => DataType::Decimal,
356                            TokioPgType::DATE => DataType::Date,
357                            TokioPgType::TIME => DataType::Time,
358                            TokioPgType::TIMESTAMP => DataType::Timestamp,
359                            TokioPgType::TIMESTAMPTZ => DataType::Timestamptz,
360                            TokioPgType::TEXT | TokioPgType::VARCHAR => DataType::Varchar,
361                            TokioPgType::INTERVAL => DataType::Interval,
362                            TokioPgType::JSONB => DataType::Jsonb,
363                            TokioPgType::BYTEA => DataType::Bytea,
364                            _ => {
365                                return Err(crate::error::ErrorCode::BindError(format!(
366                                    "unsupported column type: {}",
367                                    column.type_()
368                                ))
369                                .into());
370                            }
371                        };
372                        rw_types.push((name, data_type));
373                    }
374                    Ok::<risingwave_common::types::DataType, anyhow::Error>(DataType::Struct(
375                        StructType::new(rw_types),
376                    ))
377                })
378            })?;
379
380            Ok(TableFunction {
381                args,
382                return_type: schema,
383                function_type: TableFunctionType::PostgresQuery,
384                user_defined: None,
385            })
386        }
387    }
388
389    pub fn new_mysql_query(args: Vec<ExprImpl>) -> RwResult<Self> {
390        let evaled_args = args
391            .iter()
392            .map(expr_impl_to_string_fn)
393            .collect::<RwResult<Vec<_>>>()?;
394
395        #[cfg(madsim)]
396        {
397            return Err(crate::error::ErrorCode::BindError(
398                "postgres_query can't be used in the madsim mode".to_string(),
399            )
400            .into());
401        }
402
403        #[cfg(not(madsim))]
404        {
405            let schema = tokio::task::block_in_place(|| {
406                FRONTEND_RUNTIME.block_on(async {
407                    let database_opts: mysql_async::Opts = {
408                        let port = evaled_args[1]
409                            .parse::<u16>()
410                            .context("failed to parse port")?;
411                        mysql_async::OptsBuilder::default()
412                            .ip_or_hostname(evaled_args[0].clone())
413                            .tcp_port(port)
414                            .user(Some(evaled_args[2].clone()))
415                            .pass(Some(evaled_args[3].clone()))
416                            .db_name(Some(evaled_args[4].clone()))
417                            .into()
418                    };
419
420                    let pool = mysql_async::Pool::new(database_opts);
421                    let mut conn = pool
422                        .get_conn()
423                        .await
424                        .context("failed to connect to mysql in binder")?;
425
426                    let query = evaled_args[5].clone();
427                    let statement = conn
428                        .prep(query)
429                        .await
430                        .context("failed to prepare mysql_query in binder")?;
431
432                    let mut rw_types = vec![];
433
434                    for column in statement.columns().iter() {
435                        let name = column.name_str().to_string();
436                        let data_type = match column.column_type() {
437                            // Boolean types
438                            MySqlColumnType::MYSQL_TYPE_BIT if column.column_length() == 1 => {
439                                DataType::Boolean
440                            }
441
442                            // Numeric types
443                            // NOTE(kwannoel): Although `bool/boolean` is a synonym of TINY(1) in MySQL,
444                            // we treat it as Int16 here. It is better to be straightforward in our conversion.
445                            MySqlColumnType::MYSQL_TYPE_TINY => DataType::Int16,
446                            MySqlColumnType::MYSQL_TYPE_SHORT => DataType::Int16,
447                            MySqlColumnType::MYSQL_TYPE_INT24 => DataType::Int32,
448                            MySqlColumnType::MYSQL_TYPE_LONG => DataType::Int32,
449                            MySqlColumnType::MYSQL_TYPE_LONGLONG => DataType::Int64,
450                            MySqlColumnType::MYSQL_TYPE_FLOAT => DataType::Float32,
451                            MySqlColumnType::MYSQL_TYPE_DOUBLE => DataType::Float64,
452                            MySqlColumnType::MYSQL_TYPE_NEWDECIMAL => DataType::Decimal,
453                            MySqlColumnType::MYSQL_TYPE_DECIMAL => DataType::Decimal,
454
455                            // Date time types
456                            MySqlColumnType::MYSQL_TYPE_YEAR => DataType::Int32,
457                            MySqlColumnType::MYSQL_TYPE_DATE => DataType::Date,
458                            MySqlColumnType::MYSQL_TYPE_NEWDATE => DataType::Date,
459                            MySqlColumnType::MYSQL_TYPE_TIME => DataType::Time,
460                            MySqlColumnType::MYSQL_TYPE_TIME2 => DataType::Time,
461                            MySqlColumnType::MYSQL_TYPE_DATETIME => DataType::Timestamp,
462                            MySqlColumnType::MYSQL_TYPE_DATETIME2 => DataType::Timestamp,
463                            MySqlColumnType::MYSQL_TYPE_TIMESTAMP => DataType::Timestamptz,
464                            MySqlColumnType::MYSQL_TYPE_TIMESTAMP2 => DataType::Timestamptz,
465
466                            // String types
467                            MySqlColumnType::MYSQL_TYPE_VARCHAR => DataType::Varchar,
468                            // mysql_async does not have explicit `varbinary` and `binary` types,
469                            // we need to check the `ColumnFlags` to distinguish them.
470                            MySqlColumnType::MYSQL_TYPE_STRING
471                            | MySqlColumnType::MYSQL_TYPE_VAR_STRING => {
472                                if column
473                                    .flags()
474                                    .contains(mysql_async::consts::ColumnFlags::BINARY_FLAG)
475                                {
476                                    DataType::Bytea
477                                } else {
478                                    DataType::Varchar
479                                }
480                            }
481
482                            // JSON types
483                            MySqlColumnType::MYSQL_TYPE_JSON => DataType::Jsonb,
484
485                            // Binary types
486                            MySqlColumnType::MYSQL_TYPE_BIT
487                            | MySqlColumnType::MYSQL_TYPE_BLOB
488                            | MySqlColumnType::MYSQL_TYPE_TINY_BLOB
489                            | MySqlColumnType::MYSQL_TYPE_MEDIUM_BLOB
490                            | MySqlColumnType::MYSQL_TYPE_LONG_BLOB => DataType::Bytea,
491
492                            MySqlColumnType::MYSQL_TYPE_UNKNOWN
493                            | MySqlColumnType::MYSQL_TYPE_TYPED_ARRAY
494                            | MySqlColumnType::MYSQL_TYPE_ENUM
495                            | MySqlColumnType::MYSQL_TYPE_SET
496                            | MySqlColumnType::MYSQL_TYPE_GEOMETRY
497                            | MySqlColumnType::MYSQL_TYPE_VECTOR
498                            | MySqlColumnType::MYSQL_TYPE_NULL => {
499                                return Err(crate::error::ErrorCode::BindError(format!(
500                                    "unsupported column type: {:?}",
501                                    column.column_type()
502                                ))
503                                .into());
504                            }
505                        };
506                        rw_types.push((name, data_type));
507                    }
508                    Ok::<risingwave_common::types::DataType, anyhow::Error>(DataType::Struct(
509                        StructType::new(rw_types),
510                    ))
511                })
512            })?;
513
514            Ok(TableFunction {
515                args,
516                return_type: schema,
517                function_type: TableFunctionType::MysqlQuery,
518                user_defined: None,
519            })
520        }
521    }
522
523    /// Bind a `mssql_query(...)` table function call.
524    ///
525    /// `args` is the **pre-resolved** argument vector produced by the
526    /// binder (`Binder::bind_external_db_query_args`), so this function
527    /// only does the bind-time schema discovery — it does not look up the
528    /// source catalog or check privileges. See `bind_external_db_query_args`
529    /// for the inline / source-reference calling forms and the
530    /// privilege-check ordering.
531    ///
532    /// The returned `TableFunction` has a `return_type` that is a
533    /// `DataType::Struct` discovered at bind time by calling
534    /// `describe_mssql_query` against the user-provided query. The
535    /// `encrypt` and `trust_cert` flags are parsed strictly (any value
536    /// other than case-insensitive `true` / `false` is a bind error).
537    pub fn new_mssql_query(mut args: Vec<ExprImpl>) -> RwResult<Self> {
538        if args.len() != MSSQL_QUERY_ARG_COUNT {
539            return Err(BindError(format!(
540                "mssql_query expects {MSSQL_QUERY_ARG_COUNT} resolved arguments, got {}",
541                args.len()
542            ))
543            .into());
544        }
545        let evaled_args = args
546            .iter()
547            .map(expr_impl_to_string_fn)
548            .collect::<RwResult<Vec<_>>>()?;
549
550        let encrypt = evaled_args[MSSQL_QUERY_ENCRYPT_ARG_INDEX]
551            .trim()
552            .to_ascii_lowercase()
553            .parse::<bool>()
554            .context("invalid `encrypt` value for mssql_query: expected 'true' or 'false'")?;
555        let trust_cert = evaled_args[MSSQL_QUERY_TRUST_CERT_ARG_INDEX]
556            .trim()
557            .to_ascii_lowercase()
558            .parse::<bool>()
559            .context("invalid `trust_cert` value for mssql_query: expected 'true' or 'false'")?;
560        // These flags are strings at the SQL boundary to preserve the public
561        // function signature, but booleans everywhere in the internal plan.
562        // This also ensures they are parsed exactly once.
563        args[MSSQL_QUERY_ENCRYPT_ARG_INDEX] = ExprImpl::literal_bool(encrypt);
564        args[MSSQL_QUERY_TRUST_CERT_ARG_INDEX] = ExprImpl::literal_bool(trust_cert);
565
566        #[cfg(madsim)]
567        {
568            return Err(crate::error::ErrorCode::BindError(
569                "mssql_query can't be used in the madsim mode".to_string(),
570            )
571            .into());
572        }
573
574        #[cfg(not(madsim))]
575        {
576            // Schema discovery is async and yields on `FRONTEND_RUNTIME` while
577            // waiting for network I/O. The binder API itself is synchronous,
578            // so `block_in_place` tells the caller's Tokio runtime to move its
579            // other tasks off this worker while it waits for the async result.
580            let schema = tokio::task::block_in_place(|| {
581                FRONTEND_RUNTIME.block_on(async {
582                    let conn_config = MssqlConnectionConfig {
583                        host: evaled_args[0].clone(),
584                        port: evaled_args[1].parse().with_context(|| {
585                            format!("invalid sql server port `{}`", evaled_args[1])
586                        })?,
587                        user: evaled_args[2].clone(),
588                        password: evaled_args[3].clone(),
589                        database: evaled_args[4].clone(),
590                        encrypt,
591                        trust_cert,
592                    };
593
594                    let rw_types = describe_mssql_query(&conn_config, &evaled_args[5]).await?;
595
596                    Ok::<risingwave_common::types::DataType, anyhow::Error>(DataType::Struct(
597                        StructType::new(rw_types),
598                    ))
599                })
600            })?;
601
602            Ok(TableFunction {
603                args,
604                return_type: schema,
605                function_type: TableFunctionType::MssqlQuery,
606                user_defined: None,
607            })
608        }
609    }
610
611    /// This is a highly specific _internal_ table function meant to scan and aggregate
612    /// `backfill_table_id`, `row_count` for all MVs which are still being created.
613    pub fn new_internal_backfill_progress() -> Self {
614        TableFunction {
615            args: vec![],
616            return_type: DataType::Struct(StructType::new(vec![
617                ("job_id".to_owned(), DataType::Int32),
618                ("fragment_id".to_owned(), DataType::Int32),
619                ("backfill_state_table_id".to_owned(), DataType::Int32),
620                ("current_row_count".to_owned(), DataType::Int64),
621                ("min_epoch".to_owned(), DataType::Int64),
622            ])),
623            function_type: TableFunctionType::InternalBackfillProgress,
624            user_defined: None,
625        }
626    }
627
628    pub fn new_internal_source_backfill_progress() -> Self {
629        TableFunction {
630            args: vec![],
631            return_type: DataType::Struct(StructType::new(vec![
632                ("job_id".to_owned(), DataType::Int32),
633                ("fragment_id".to_owned(), DataType::Int32),
634                ("backfill_state_table_id".to_owned(), DataType::Int32),
635                ("partition_id".to_owned(), DataType::Varchar),
636                ("backfill_progress".to_owned(), DataType::Jsonb),
637            ])),
638            function_type: TableFunctionType::InternalSourceBackfillProgress,
639            user_defined: None,
640        }
641    }
642
643    pub fn new_internal_get_channel_delta_stats(args: Vec<ExprImpl>) -> Self {
644        Self {
645            args,
646            return_type: DataType::Struct(StructType::new(vec![
647                ("upstream_fragment_id".to_owned(), DataType::Int32),
648                ("downstream_fragment_id".to_owned(), DataType::Int32),
649                ("backpressure_rate".to_owned(), DataType::Float64),
650                ("recv_throughput".to_owned(), DataType::Float64),
651                ("send_throughput".to_owned(), DataType::Float64),
652            ])),
653            function_type: TableFunctionType::InternalGetChannelDeltaStats,
654            user_defined: None,
655        }
656    }
657
658    pub fn to_protobuf(&self) -> PbTableFunction {
659        PbTableFunction {
660            function_type: self.function_type as i32,
661            args: self.args.iter().map(|c| c.to_expr_proto()).collect_vec(),
662            return_type: Some(self.return_type.to_protobuf()),
663            udf: self.user_defined.as_ref().map(|c| c.as_ref().into()),
664        }
665    }
666
667    /// Serialize the table function. Returns an error if this will result in an impure table
668    /// function on a retract stream, which may lead to inconsistent results.
669    pub fn to_protobuf_checked_pure(&self, retract: bool) -> crate::error::Result<PbTableFunction> {
670        if retract {
671            reject_impure(self.clone(), "table function")?;
672        }
673
674        let args = self
675            .args
676            .iter()
677            .map(|arg| arg.to_expr_proto_checked_pure(retract, "table function argument"))
678            .collect::<crate::error::Result<Vec<_>>>()?;
679
680        Ok(PbTableFunction {
681            function_type: self.function_type as i32,
682            args,
683            return_type: Some(self.return_type.to_protobuf()),
684            udf: self.user_defined.as_ref().map(|c| c.as_ref().into()),
685        })
686    }
687
688    /// Get the name of the table function.
689    pub fn name(&self) -> String {
690        match self.function_type {
691            TableFunctionType::UserDefined => self.user_defined.as_ref().unwrap().name.clone(),
692            t => t.as_str_name().to_lowercase(),
693        }
694    }
695
696    pub fn rewrite(self, rewriter: &mut impl ExprRewriter) -> Self {
697        Self {
698            args: self
699                .args
700                .into_iter()
701                .map(|e| rewriter.rewrite_expr(e))
702                .collect(),
703            ..self
704        }
705    }
706}
707
708impl std::fmt::Debug for TableFunction {
709    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
710        if f.alternate() {
711            f.debug_struct("FunctionCall")
712                .field("function_type", &self.function_type)
713                .field("return_type", &self.return_type)
714                .field("args", &self.args)
715                .finish()
716        } else {
717            let func_name = format!("{:?}", self.function_type);
718            let mut builder = f.debug_tuple(&func_name);
719            self.args.iter().for_each(|child| {
720                builder.field(child);
721            });
722            builder.finish()
723        }
724    }
725}
726
727impl Expr for TableFunction {
728    fn return_type(&self) -> DataType {
729        self.return_type.clone()
730    }
731
732    fn try_to_expr_proto(&self) -> Result<risingwave_pb::expr::ExprNode, String> {
733        Err("Table function should not be converted to ExprNode".to_owned())
734    }
735}
736
737pub(crate) fn expr_impl_to_string_fn(arg: &ExprImpl) -> RwResult<String> {
738    match arg.try_fold_const() {
739        Some(Ok(value)) => {
740            let Some(scalar) = value else {
741                return Err(BindError(
742                    "postgres_query / mysql_query / mssql_query do not accept null arguments"
743                        .to_owned(),
744                )
745                .into());
746            };
747            Ok(scalar.into_utf8().to_string())
748        }
749        Some(Err(err)) => Err(err),
750        None => Err(BindError(
751            "postgres_query / mysql_query / mssql_query only accept constant arguments".to_owned(),
752        )
753        .into()),
754    }
755}
risingwave_frontend/expr/table_function.rs

risingwave_frontend/expr/
table_function.rs
