Skip to main content

risingwave_frontend/optimizer/rule/
table_function_to_mssql_query_rule.rs

1// Copyright 2026 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 itertools::Itertools;
16use risingwave_common::catalog::{Field, Schema};
17use risingwave_common::types::{DataType, ScalarImpl};
18
19use super::prelude::{PlanRef, *};
20use crate::expr::{Expr, ExprImpl, MSSQL_QUERY_ARG_COUNT, TableFunctionType};
21use crate::optimizer::plan_node::generic::GenericPlanRef;
22use crate::optimizer::plan_node::{LogicalMssqlQuery, LogicalTableFunction, MssqlConnection};
23
24/// Optimizer rule that rewrites a `LogicalTableFunction` of type
25/// `MSSQL_QUERY` into a `LogicalMssqlQuery` plan node. Registered in the
26/// `TABLE_FUNCTION_TO_MSSQL_QUERY` stage of the logical optimization
27/// pipeline (after the equivalent `postgres_query` and `mysql_query`
28/// rules). Validates that exactly 8 arguments were supplied; the 2-arg
29/// source-reference form is handled earlier by the binder, not here.
30pub struct TableFunctionToMssqlQueryRule {}
31
32fn fold_const_varchar(expr: &ExprImpl) -> Option<String> {
33    let value = expr.try_fold_const()?.ok()??;
34    match value {
35        ScalarImpl::Utf8(value) => Some(value.to_string()),
36        _ => None,
37    }
38}
39
40fn fold_const_bool(expr: &ExprImpl) -> Option<bool> {
41    let value = expr.try_fold_const()?.ok()??;
42    match value {
43        ScalarImpl::Bool(value) => Some(value),
44        _ => None,
45    }
46}
47
48impl Rule<Logical> for TableFunctionToMssqlQueryRule {
49    /// Apply the rewrite: if the plan is a `LogicalTableFunction` of
50    /// type `MSSQL_QUERY`, build a `LogicalMssqlQuery` from its 8 inline
51    /// arguments. Returns `None` for any other plan.
52    fn apply(&self, plan: PlanRef) -> Option<PlanRef> {
53        let logical_table_function: &LogicalTableFunction = plan.as_logical_table_function()?;
54        if logical_table_function.table_function().function_type != TableFunctionType::MssqlQuery {
55            return None;
56        }
57        // The binder rejects this form with a user-facing error. Keep this
58        // check fallible as defense against malformed internally-built plans.
59        if logical_table_function.with_ordinality() {
60            return None;
61        }
62        let table_function_return_type = logical_table_function.table_function().return_type();
63
64        if let DataType::Struct(st) = table_function_return_type {
65            let fields = st
66                .iter()
67                .map(|(name, data_type)| Field::with_name(data_type.clone(), name.to_owned()))
68                .collect_vec();
69
70            let schema = Schema::new(fields);
71
72            let args = &logical_table_function.table_function().args;
73            if args.len() != MSSQL_QUERY_ARG_COUNT {
74                return None;
75            }
76
77            // `TableFunction::new_mssql_query` validates all arguments as
78            // constants and normalizes the two TLS flags to boolean literals.
79            // Re-check the invariant here instead of panicking if an internal
80            // caller constructs a malformed table function.
81            let hostname = fold_const_varchar(&args[0])?;
82            let port = fold_const_varchar(&args[1])?;
83            let username = fold_const_varchar(&args[2])?;
84            let password = fold_const_varchar(&args[3])?;
85            let database = fold_const_varchar(&args[4])?;
86            let query = fold_const_varchar(&args[5])?;
87            let encrypt = fold_const_bool(&args[6])?;
88            let trust_cert = fold_const_bool(&args[7])?;
89
90            Some(
91                LogicalMssqlQuery::new(
92                    logical_table_function.ctx(),
93                    schema,
94                    MssqlConnection {
95                        hostname,
96                        port,
97                        username,
98                        password,
99                        database,
100                        encrypt,
101                        trust_cert,
102                    },
103                    query,
104                )
105                .into(),
106            )
107        } else {
108            unreachable!("TableFunction return type should be struct")
109        }
110    }
111}
112
113impl TableFunctionToMssqlQueryRule {
114    /// Construct a boxed instance for registration in the
115    /// `TABLE_FUNCTION_TO_MSSQL_QUERY` stage.
116    pub fn create() -> BoxedRule {
117        Box::new(TableFunctionToMssqlQueryRule {})
118    }
119}