risingwave_planner_test/
lib.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
15#![allow(clippy::derive_partial_eq_without_eq)]
16#![feature(used_with_arg)]
17
18//! Data-driven tests.
19
20risingwave_expr_impl::enable!();
21
22mod resolve_id;
23
24use std::collections::{BTreeMap, HashSet};
25use std::path::{Path, PathBuf};
26use std::sync::Arc;
27
28use anyhow::{Result, anyhow, bail};
29pub use resolve_id::*;
30use risingwave_expr::sig::{CreateFunctionOutput, UDF_IMPLS, UdfImplDescriptor};
31use risingwave_frontend::handler::util::SourceSchemaCompatExt;
32use risingwave_frontend::handler::{
33    HandlerArgs, create_function, create_index, create_mv, create_schema, create_source,
34    create_sql_function, create_table, create_view, drop_table, explain, variable,
35};
36use risingwave_frontend::optimizer::backfill_order_strategy::explain_backfill_order_in_dot_format;
37use risingwave_frontend::optimizer::plan_node::ConventionMarker;
38use risingwave_frontend::session::SessionImpl;
39use risingwave_frontend::test_utils::{LocalFrontend, create_proto_file, get_explain_output};
40use risingwave_frontend::{
41    Binder, Explain, FrontendOpts, OptimizerContext, OptimizerContextRef, PlanRef, Planner,
42    WithOptionsSecResolved, build_graph, explain_stream_graph,
43};
44use risingwave_license::{LicenseKey, LicenseManager};
45use risingwave_sqlparser::ast::{
46    AstOption, BackfillOrderStrategy, DropMode, EmitMode, ExplainOptions, ObjectName, Statement,
47};
48use risingwave_sqlparser::parser::Parser;
49use serde::{Deserialize, Serialize};
50use thiserror_ext::AsReport;
51
52/// Planner tests are built without the `udf` feature of `risingwave_expr_impl`, so
53/// `language javascript` UDFs are created with this implementation, which only records the
54/// function body. UDFs are never executed during planning.
55#[linkme::distributed_slice(UDF_IMPLS)]
56static PLANNER_TEST_JAVASCRIPT_UDF: UdfImplDescriptor = UdfImplDescriptor {
57    match_fn: |language, _runtime, _link| language == "javascript",
58    create_fn: |opts| {
59        Ok(CreateFunctionOutput {
60            name_in_runtime: opts.name.to_owned(),
61            body: opts.as_.map(str::to_owned),
62            compressed_binary: None,
63        })
64    },
65    build_fn: |_| bail!("UDFs are not executed in planner tests"),
66};
67
68#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Hash, Eq)]
69#[serde(deny_unknown_fields, rename_all = "snake_case")]
70pub enum TestType {
71    /// The result of an `EXPLAIN` statement.
72    ///
73    /// This field is used when `sql` is an `EXPLAIN` statement.
74    /// In this case, all other fields are invalid.
75    ExplainOutput,
76
77    /// The original logical plan
78    LogicalPlan,
79    /// Logical plan with optimization `.gen_optimized_logical_plan_for_batch()`
80    OptimizedLogicalPlanForBatch,
81    /// Logical plan with optimization `.gen_optimized_logical_plan_for_stream()`
82    OptimizedLogicalPlanForStream,
83
84    /// Distributed batch plan `.gen_batch_query_plan()`
85    BatchPlan,
86    /// Proto JSON of generated batch plan
87    BatchPlanProto,
88    /// Batch plan for local execution `.gen_batch_local_plan()`
89    BatchLocalPlan,
90    /// Batch plan for local execution `.gen_batch_distributed_plan()`
91    BatchDistributedPlan,
92
93    /// Create MV plan `.gen_create_mv_plan()`
94    StreamPlan,
95    /// Create MV fragments plan
96    StreamDistPlan,
97    /// Create MV plan with EOWC semantics `.gen_create_mv_plan(.., EmitMode::OnWindowClose)`
98    EowcStreamPlan,
99    /// Create MV fragments plan with EOWC semantics
100    EowcStreamDistPlan,
101    /// Create Backfill Order Plan
102    BackfillOrderPlan,
103
104    /// Create sink plan (assumes blackhole sink)
105    /// TODO: Other sinks
106    SinkPlan,
107
108    BinderError,
109    PlannerError,
110    OptimizerError,
111    BatchError,
112    BatchLocalError,
113    StreamError,
114    EowcStreamError,
115}
116
117pub fn check(actual: Vec<TestCaseResult>, expect: expect_test::ExpectFile) {
118    let actual = serde_yaml::to_string(&actual).unwrap();
119    expect.assert_eq(&format!("# This file is automatically generated. See `src/frontend/planner_test/README.md` for more information.\n{}",actual));
120}
121
122#[serde_with::skip_serializing_none]
123#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Default)]
124#[serde(deny_unknown_fields)]
125pub struct TestInput {
126    /// Id of the test case, used in before.
127    pub id: Option<String>,
128    /// A brief description of the test case.
129    pub name: Option<String>,
130    /// Before running the SQL statements, the test runner will execute the specified test cases
131    pub before: Option<Vec<String>>,
132    /// The resolved statements of the before ids
133    #[serde(skip_serializing)]
134    before_statements: Option<Vec<String>>,
135    /// The SQL statements
136    pub sql: String,
137}
138
139#[serde_with::skip_serializing_none]
140#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Default)]
141#[serde(deny_unknown_fields)]
142pub struct TestCase {
143    #[serde(flatten)]
144    pub input: TestInput,
145
146    // TODO: these should also be in TestInput, but it affects ordering. So next PR
147    /// Support using file content or file location to create source.
148    pub create_source: Option<CreateConnector>,
149    /// Support using file content or file location to create table with connector.
150    pub create_table_with_connector: Option<CreateConnector>,
151    /// Provide config map to frontend
152    pub with_config_map: Option<BTreeMap<String, String>>,
153
154    /// Specify what output fields to check
155    pub expected_outputs: HashSet<TestType>,
156}
157
158impl TestCase {
159    pub fn id(&self) -> &Option<String> {
160        &self.input.id
161    }
162
163    pub fn name(&self) -> &Option<String> {
164        &self.input.name
165    }
166
167    pub fn before(&self) -> &Option<Vec<String>> {
168        &self.input.before
169    }
170
171    pub fn before_statements(&self) -> &Option<Vec<String>> {
172        &self.input.before_statements
173    }
174
175    pub fn sql(&self) -> &String {
176        &self.input.sql
177    }
178
179    pub fn create_source(&self) -> &Option<CreateConnector> {
180        &self.create_source
181    }
182
183    pub fn create_table_with_connector(&self) -> &Option<CreateConnector> {
184        &self.create_table_with_connector
185    }
186
187    pub fn with_config_map(&self) -> &Option<BTreeMap<String, String>> {
188        &self.with_config_map
189    }
190}
191
192#[serde_with::skip_serializing_none]
193#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Default)]
194#[serde(deny_unknown_fields)]
195pub struct CreateConnector {
196    format: String,
197    encode: String,
198    name: String,
199    file: Option<String>,
200    is_table: Option<bool>,
201}
202
203#[serde_with::skip_serializing_none]
204#[derive(Debug, PartialEq, Serialize, Deserialize, Default)]
205#[serde(deny_unknown_fields)]
206pub struct TestCaseResult {
207    #[serde(flatten)]
208    pub input: TestInput,
209
210    /// The original logical plan
211    pub logical_plan: Option<String>,
212
213    /// Logical plan with optimization `.gen_optimized_logical_plan_for_batch()`
214    pub optimized_logical_plan_for_batch: Option<String>,
215
216    /// Logical plan with optimization `.gen_optimized_logical_plan_for_stream()`
217    pub optimized_logical_plan_for_stream: Option<String>,
218
219    /// Distributed batch plan `.gen_batch_query_plan()`
220    pub batch_plan: Option<String>,
221
222    /// Proto JSON of generated batch plan
223    pub batch_plan_proto: Option<String>,
224
225    /// Batch plan for local execution `.gen_batch_local_plan()`
226    pub batch_local_plan: Option<String>,
227
228    /// Batch plan for distributed execution `.gen_batch_distributed_plan()`
229    pub batch_distributed_plan: Option<String>,
230
231    /// Generate sink plan
232    pub sink_plan: Option<String>,
233
234    /// Create MV plan `.gen_create_mv_plan()`
235    pub stream_plan: Option<String>,
236
237    /// Create MV fragments plan
238    pub stream_dist_plan: Option<String>,
239
240    /// Create MV plan with EOWC semantics `.gen_create_mv_plan(.., EmitMode::OnWindowClose)`
241    pub eowc_stream_plan: Option<String>,
242
243    /// Create MV fragments plan with EOWC semantics
244    pub eowc_stream_dist_plan: Option<String>,
245
246    /// Create Backfill Order Plan
247    pub backfill_order_plan: Option<String>,
248
249    /// Error of binder
250    pub binder_error: Option<String>,
251
252    /// Error of planner
253    pub planner_error: Option<String>,
254
255    /// Error of optimizer
256    pub optimizer_error: Option<String>,
257
258    /// Error of `.gen_batch_query_plan()`
259    pub batch_error: Option<String>,
260
261    /// Error of `.gen_batch_local_plan()`
262    pub batch_local_error: Option<String>,
263
264    /// Error of `.gen_stream_plan()`
265    pub stream_error: Option<String>,
266
267    /// Error of `.gen_stream_plan()` with `emit_on_window_close = true`
268    pub eowc_stream_error: Option<String>,
269
270    /// Error of `.gen_sink_plan()`
271    pub sink_error: Option<String>,
272
273    /// The result of an `EXPLAIN` statement.
274    ///
275    /// This field is used when `sql` is an `EXPLAIN` statement.
276    /// In this case, all other fields are invalid.
277    pub explain_output: Option<String>,
278
279    // TODO: these should also be in TestInput, but it affects ordering. So next PR
280    /// Support using file content or file location to create source.
281    pub create_source: Option<CreateConnector>,
282    /// Support using file content or file location to create table with connector.
283    pub create_table_with_connector: Option<CreateConnector>,
284    /// Provide config map to frontend
285    pub with_config_map: Option<BTreeMap<String, String>>,
286}
287
288impl TestCase {
289    /// Run the test case, and return the expected output.
290    pub async fn run(&self, do_check_result: bool) -> Result<TestCaseResult> {
291        let session = {
292            let frontend = LocalFrontend::new(FrontendOpts::default()).await;
293            frontend.session_ref()
294        };
295
296        if let Some(config_map) = self.with_config_map() {
297            for (key, val) in config_map {
298                session.set_config(key, val.to_owned()).unwrap();
299            }
300        }
301
302        let placeholder_empty_vec = vec![];
303
304        // Since temp file will be deleted when it goes out of scope, so create source in advance.
305        Box::pin(self.do_create_source(session.clone())).await?;
306        Box::pin(self.do_create_table_with_connector(session.clone())).await?;
307
308        let mut result: Option<TestCaseResult> = None;
309        for sql in self
310            .before_statements()
311            .as_ref()
312            .unwrap_or(&placeholder_empty_vec)
313            .iter()
314            .chain(std::iter::once(self.sql()))
315        {
316            result = Box::pin(self.run_sql(
317                Arc::from(sql.to_owned()),
318                session.clone(),
319                do_check_result,
320                result,
321            ))
322            .await?;
323        }
324
325        let mut result = result.unwrap_or_default();
326        result.input = self.input.clone();
327        result.create_source.clone_from(self.create_source());
328        result
329            .create_table_with_connector
330            .clone_from(self.create_table_with_connector());
331        result.with_config_map.clone_from(self.with_config_map());
332
333        Ok(result)
334    }
335
336    #[inline(always)]
337    fn create_connector_sql(
338        is_table: bool,
339        connector_name: String,
340        connector_format: String,
341        connector_encode: String,
342    ) -> String {
343        let object_to_create = if is_table { "TABLE" } else { "SOURCE" };
344        format!(
345            r#"CREATE {} {}
346    WITH (connector = 'kafka', kafka.topic = 'abc', kafka.brokers = 'localhost:1001')
347    FORMAT {} ENCODE {} (message = '.test.TestRecord', schema.location = 'file://"#,
348            object_to_create, connector_name, connector_format, connector_encode
349        )
350    }
351
352    async fn do_create_table_with_connector(
353        &self,
354        session: Arc<SessionImpl>,
355    ) -> Result<Option<TestCaseResult>> {
356        match self.create_table_with_connector().clone() {
357            Some(connector) => {
358                if let Some(content) = connector.file {
359                    let sql = Self::create_connector_sql(
360                        true,
361                        connector.name,
362                        connector.format,
363                        connector.encode,
364                    );
365                    let temp_file = create_proto_file(content.as_str());
366                    Box::pin(self.run_sql(
367                        Arc::from(sql + temp_file.path().to_str().unwrap() + "')"),
368                        session.clone(),
369                        false,
370                        None,
371                    ))
372                    .await
373                } else {
374                    panic!(
375                        "{:?} create table with connector must include `file` for the file content",
376                        self.id()
377                    );
378                }
379            }
380            None => Ok(None),
381        }
382    }
383
384    // If testcase have create source info, run sql to create source.
385    // Support create source by file content or file location.
386    async fn do_create_source(&self, session: Arc<SessionImpl>) -> Result<Option<TestCaseResult>> {
387        match self.create_source().clone() {
388            Some(source) => {
389                if let Some(content) = source.file {
390                    let sql = Self::create_connector_sql(
391                        false,
392                        source.name,
393                        source.format,
394                        source.encode,
395                    );
396                    let temp_file = create_proto_file(content.as_str());
397                    Box::pin(self.run_sql(
398                        Arc::from(sql + temp_file.path().to_str().unwrap() + "')"),
399                        session.clone(),
400                        false,
401                        None,
402                    ))
403                    .await
404                } else {
405                    panic!(
406                        "{:?} create source must include `file` for the file content",
407                        self.id()
408                    );
409                }
410            }
411            None => Ok(None),
412        }
413    }
414
415    async fn run_sql(
416        &self,
417        sql: Arc<str>,
418        session: Arc<SessionImpl>,
419        do_check_result: bool,
420        mut result: Option<TestCaseResult>,
421    ) -> Result<Option<TestCaseResult>> {
422        let statements = Parser::parse_sql(&sql).unwrap();
423        for stmt in statements {
424            // TODO: `sql` may contain multiple statements here.
425            let handler_args = HandlerArgs::new(session.clone(), &stmt, sql.clone())?;
426            let _guard = session.txn_begin_implicit();
427            match stmt.clone() {
428                Statement::Query(_)
429                | Statement::Insert { .. }
430                | Statement::Delete { .. }
431                | Statement::Update { .. } => {
432                    if result.is_some() {
433                        panic!("two queries in one test case");
434                    }
435                    let explain_options = ExplainOptions {
436                        verbose: true,
437                        ..Default::default()
438                    };
439                    let context = OptimizerContext::new(
440                        HandlerArgs::new(session.clone(), &stmt, sql.clone())?,
441                        explain_options,
442                    );
443                    let ret = self.apply_query(&stmt, context.into())?;
444                    if do_check_result {
445                        check_result(self, &ret)?;
446                    }
447                    result = Some(ret);
448                }
449                Statement::CreateTable {
450                    name,
451                    columns,
452                    constraints,
453                    if_not_exists,
454                    format_encode,
455                    source_watermarks,
456                    append_only,
457                    on_conflict,
458                    with_version_columns,
459                    cdc_table_info,
460                    include_column_options,
461                    wildcard_idx,
462                    webhook_info,
463                    engine,
464                    ..
465                } => {
466                    let format_encode = format_encode.map(|schema| schema.into_v2_with_warning());
467
468                    Box::pin(create_table::handle_create_table(
469                        handler_args,
470                        name,
471                        columns,
472                        wildcard_idx,
473                        constraints,
474                        if_not_exists,
475                        format_encode,
476                        source_watermarks,
477                        append_only,
478                        on_conflict,
479                        with_version_columns
480                            .iter()
481                            .map(|x| x.real_value())
482                            .collect(),
483                        cdc_table_info,
484                        include_column_options,
485                        webhook_info,
486                        engine,
487                    ))
488                    .await?;
489                }
490                Statement::CreateSource { stmt } => {
491                    if let Err(error) =
492                        create_source::handle_create_source(handler_args, stmt).await
493                    {
494                        let actual_result = TestCaseResult {
495                            planner_error: Some(error.to_report_string()),
496                            ..Default::default()
497                        };
498
499                        check_result(self, &actual_result)?;
500                        result = Some(actual_result);
501                    }
502                }
503                Statement::CreateIndex {
504                    name,
505                    table_name,
506                    method,
507                    columns,
508                    include,
509                    distributed_by,
510                    if_not_exists,
511                    // TODO: support unique and if_not_exist in planner test
512                    ..
513                } => {
514                    create_index::handle_create_index(
515                        handler_args,
516                        if_not_exists,
517                        name,
518                        table_name,
519                        method,
520                        columns,
521                        include,
522                        distributed_by,
523                    )
524                    .await?;
525                }
526                Statement::CreateView {
527                    materialized: true,
528                    or_replace: false,
529                    if_not_exists,
530                    name,
531                    query,
532                    columns,
533                    emit_mode,
534                    ..
535                } => {
536                    create_mv::handle_create_mv(
537                        handler_args,
538                        if_not_exists,
539                        name,
540                        *query,
541                        columns,
542                        emit_mode,
543                    )
544                    .await?;
545                }
546                Statement::CreateView {
547                    materialized: false,
548                    or_replace: false,
549                    if_not_exists,
550                    name,
551                    query,
552                    columns,
553                    ..
554                } => {
555                    create_view::handle_create_view(
556                        handler_args,
557                        if_not_exists,
558                        name,
559                        columns,
560                        *query,
561                    )
562                    .await?;
563                }
564                Statement::Drop(drop_statement) => {
565                    drop_table::handle_drop_table(
566                        handler_args,
567                        drop_statement.object_name,
568                        drop_statement.if_exists,
569                        matches!(drop_statement.drop_mode, AstOption::Some(DropMode::Cascade)),
570                    )
571                    .await?;
572                }
573                Statement::SetVariable {
574                    local: _,
575                    variable,
576                    value,
577                } => {
578                    variable::handle_set(handler_args, variable, value).unwrap();
579                }
580                Statement::Explain {
581                    analyze,
582                    statement,
583                    options,
584                } => {
585                    if result.is_some() {
586                        panic!("two queries in one test case");
587                    }
588                    let ret = match Box::pin(explain::handle_explain(
589                        handler_args,
590                        *statement,
591                        options,
592                        analyze,
593                    ))
594                    .await
595                    {
596                        Ok(rsp) => TestCaseResult {
597                            explain_output: Some(get_explain_output(rsp).await),
598                            ..Default::default()
599                        },
600                        Err(error) => TestCaseResult {
601                            planner_error: Some(error.to_report_string()),
602                            ..Default::default()
603                        },
604                    };
605                    if do_check_result {
606                        check_result(self, &ret)?;
607                    }
608                    result = Some(ret);
609                }
610                Statement::CreateSchema {
611                    schema_name,
612                    if_not_exists,
613                    owner,
614                } => {
615                    create_schema::handle_create_schema(
616                        handler_args,
617                        schema_name,
618                        if_not_exists,
619                        owner,
620                    )
621                    .await?;
622                }
623                Statement::CreateFunction {
624                    or_replace,
625                    temporary,
626                    if_not_exists,
627                    name,
628                    args,
629                    returns,
630                    params,
631                    with_options,
632                } => {
633                    if params
634                        .language
635                        .as_ref()
636                        .is_some_and(|language| language.real_value().eq_ignore_ascii_case("sql"))
637                    {
638                        create_sql_function::handle_create_sql_function(
639                            handler_args,
640                            or_replace,
641                            temporary,
642                            if_not_exists,
643                            name,
644                            args,
645                            returns,
646                            params,
647                        )
648                        .await?;
649                    } else {
650                        create_function::handle_create_function(
651                            handler_args,
652                            or_replace,
653                            temporary,
654                            if_not_exists,
655                            name,
656                            args,
657                            returns,
658                            params,
659                            with_options,
660                        )
661                        .await?;
662                    }
663                }
664                _ => return Err(anyhow!("Unsupported statement type")),
665            }
666        }
667        Ok(result)
668    }
669
670    fn apply_query(
671        &self,
672        stmt: &Statement,
673        context: OptimizerContextRef,
674    ) -> Result<TestCaseResult> {
675        let session = context.session_ctx().clone();
676        let mut ret = TestCaseResult::default();
677
678        let bound = {
679            let mut binder = Binder::new_for_batch(&session);
680            match binder.bind(stmt.clone()) {
681                Ok(bound) => bound,
682                Err(err) => {
683                    ret.binder_error = Some(err.to_report_string_pretty());
684                    return Ok(ret);
685                }
686            }
687        };
688
689        let mut planner = Planner::new_for_stream(context.clone());
690
691        let plan_root = match planner.plan(bound) {
692            Ok(plan_root) => {
693                if self.expected_outputs.contains(&TestType::LogicalPlan) {
694                    ret.logical_plan =
695                        Some(explain_plan(&plan_root.clone().into_unordered_subplan()));
696                }
697                plan_root
698            }
699            Err(err) => {
700                ret.planner_error = Some(err.to_report_string_pretty());
701                return Ok(ret);
702            }
703        };
704
705        if self
706            .expected_outputs
707            .contains(&TestType::OptimizedLogicalPlanForBatch)
708            || self.expected_outputs.contains(&TestType::OptimizerError)
709        {
710            let plan_root = plan_root.clone();
711            let optimized_logical_plan_for_batch =
712                match plan_root.gen_optimized_logical_plan_for_batch() {
713                    Ok(optimized_logical_plan_for_batch) => optimized_logical_plan_for_batch,
714                    Err(err) => {
715                        ret.optimizer_error = Some(err.to_report_string_pretty());
716                        return Ok(ret);
717                    }
718                };
719
720            // Only generate optimized_logical_plan_for_batch if it is specified in test case
721            if self
722                .expected_outputs
723                .contains(&TestType::OptimizedLogicalPlanForBatch)
724            {
725                ret.optimized_logical_plan_for_batch =
726                    Some(explain_plan(&optimized_logical_plan_for_batch.plan));
727            }
728        }
729
730        if self
731            .expected_outputs
732            .contains(&TestType::OptimizedLogicalPlanForStream)
733            || self.expected_outputs.contains(&TestType::OptimizerError)
734        {
735            let plan_root = plan_root.clone();
736            let optimized_logical_plan_for_stream =
737                match plan_root.gen_optimized_logical_plan_for_stream() {
738                    Ok(optimized_logical_plan_for_stream) => optimized_logical_plan_for_stream,
739                    Err(err) => {
740                        ret.optimizer_error = Some(err.to_report_string_pretty());
741                        return Ok(ret);
742                    }
743                };
744
745            // Only generate optimized_logical_plan_for_stream if it is specified in test case
746            if self
747                .expected_outputs
748                .contains(&TestType::OptimizedLogicalPlanForStream)
749            {
750                ret.optimized_logical_plan_for_stream =
751                    Some(explain_plan(&optimized_logical_plan_for_stream.plan));
752            }
753        }
754
755        'batch: {
756            if self.expected_outputs.contains(&TestType::BatchPlan)
757                || self.expected_outputs.contains(&TestType::BatchPlanProto)
758                || self.expected_outputs.contains(&TestType::BatchError)
759            {
760                let plan_root = plan_root.clone();
761                let batch_plan = match plan_root.gen_batch_plan() {
762                    Ok(batch_plan) => match batch_plan.gen_batch_distributed_plan() {
763                        Ok(batch_plan) => batch_plan,
764                        Err(err) => {
765                            ret.batch_error = Some(err.to_report_string_pretty());
766                            break 'batch;
767                        }
768                    },
769                    Err(err) => {
770                        ret.batch_error = Some(err.to_report_string_pretty());
771                        break 'batch;
772                    }
773                };
774
775                // Only generate batch_plan if it is specified in test case
776                if self.expected_outputs.contains(&TestType::BatchPlan) {
777                    ret.batch_plan = Some(explain_plan(&batch_plan));
778                }
779
780                // Only generate batch_plan_proto if it is specified in test case
781                if self.expected_outputs.contains(&TestType::BatchPlanProto) {
782                    ret.batch_plan_proto = Some(serde_yaml::to_string(
783                        &batch_plan.to_batch_prost_identity(false)?,
784                    )?);
785                }
786            }
787        }
788
789        'local_batch: {
790            if self.expected_outputs.contains(&TestType::BatchLocalPlan)
791                || self.expected_outputs.contains(&TestType::BatchError)
792            {
793                let plan_root = plan_root.clone();
794                let batch_plan = match plan_root.gen_batch_plan() {
795                    Ok(batch_plan) => match batch_plan.gen_batch_local_plan() {
796                        Ok(batch_plan) => batch_plan,
797                        Err(err) => {
798                            ret.batch_error = Some(err.to_report_string_pretty());
799                            break 'local_batch;
800                        }
801                    },
802                    Err(err) => {
803                        ret.batch_error = Some(err.to_report_string_pretty());
804                        break 'local_batch;
805                    }
806                };
807
808                // Only generate batch_plan if it is specified in test case
809                if self.expected_outputs.contains(&TestType::BatchLocalPlan) {
810                    ret.batch_local_plan = Some(explain_plan(&batch_plan));
811                }
812            }
813        }
814
815        'distributed_batch: {
816            if self
817                .expected_outputs
818                .contains(&TestType::BatchDistributedPlan)
819                || self.expected_outputs.contains(&TestType::BatchError)
820            {
821                let plan_root = plan_root.clone();
822                let batch_plan = match plan_root.gen_batch_plan() {
823                    Ok(batch_plan) => match batch_plan.gen_batch_distributed_plan() {
824                        Ok(batch_plan) => batch_plan,
825                        Err(err) => {
826                            ret.batch_error = Some(err.to_report_string_pretty());
827                            break 'distributed_batch;
828                        }
829                    },
830                    Err(err) => {
831                        ret.batch_error = Some(err.to_report_string_pretty());
832                        break 'distributed_batch;
833                    }
834                };
835
836                // Only generate batch_plan if it is specified in test case
837                if self
838                    .expected_outputs
839                    .contains(&TestType::BatchDistributedPlan)
840                {
841                    ret.batch_distributed_plan = Some(explain_plan(&batch_plan));
842                }
843            }
844        }
845
846        {
847            // stream
848            for (
849                emit_mode,
850                plan,
851                ret_plan_str,
852                dist_plan,
853                ret_dist_plan_str,
854                error,
855                ret_error_str,
856            ) in [
857                (
858                    EmitMode::Immediately,
859                    self.expected_outputs.contains(&TestType::StreamPlan),
860                    &mut ret.stream_plan,
861                    self.expected_outputs.contains(&TestType::StreamDistPlan),
862                    &mut ret.stream_dist_plan,
863                    self.expected_outputs.contains(&TestType::StreamError),
864                    &mut ret.stream_error,
865                ),
866                (
867                    EmitMode::OnWindowClose,
868                    self.expected_outputs.contains(&TestType::EowcStreamPlan),
869                    &mut ret.eowc_stream_plan,
870                    self.expected_outputs
871                        .contains(&TestType::EowcStreamDistPlan),
872                    &mut ret.eowc_stream_dist_plan,
873                    self.expected_outputs.contains(&TestType::EowcStreamError),
874                    &mut ret.eowc_stream_error,
875                ),
876            ] {
877                if !plan && !dist_plan && !error {
878                    continue;
879                }
880
881                let q = if let Statement::Query(q) = stmt {
882                    q.as_ref().clone()
883                } else {
884                    return Err(anyhow!("expect a query"));
885                };
886
887                let (stream_plan, table) = match create_mv::explain_create_mv_plan(
888                    &session,
889                    context.clone(),
890                    q,
891                    ObjectName(vec!["test".into()]),
892                    vec![],
893                    Some(emit_mode),
894                ) {
895                    Ok(r) => r,
896                    Err(err) => {
897                        *ret_error_str = Some(err.to_report_string_pretty());
898                        continue;
899                    }
900                };
901
902                // Only generate stream_plan if it is specified in test case
903                if plan {
904                    *ret_plan_str = Some(explain_plan(&stream_plan));
905                }
906
907                // Only generate stream_dist_plan if it is specified in test case
908                if dist_plan {
909                    let graph = build_graph(stream_plan.clone(), None)?;
910                    *ret_dist_plan_str =
911                        Some(explain_stream_graph(&graph, Some(table.to_prost()), false));
912                }
913
914                if self.expected_outputs.contains(&TestType::BackfillOrderPlan) {
915                    match explain_backfill_order_in_dot_format(
916                        &session,
917                        BackfillOrderStrategy::Auto,
918                        stream_plan,
919                    ) {
920                        Ok(formatted_order_plan) => {
921                            ret.backfill_order_plan = Some(formatted_order_plan);
922                        }
923                        Err(err) => {
924                            *ret_error_str = Some(err.to_report_string_pretty());
925                        }
926                    }
927                }
928            }
929        }
930
931        'sink: {
932            if self.expected_outputs.contains(&TestType::SinkPlan) {
933                let plan_root = plan_root;
934                let sink_name = "sink_test";
935                let mut options = BTreeMap::new();
936                options.insert("connector".to_owned(), "blackhole".to_owned());
937                options.insert("type".to_owned(), "append-only".to_owned());
938                // let options = WithOptionsSecResolved::without_secrets(options);
939                let options = WithOptionsSecResolved::without_secrets(options);
940                let format_desc = (&options).try_into().unwrap();
941                match plan_root.gen_sink_plan(
942                    sink_name.to_owned(),
943                    format!("CREATE SINK {sink_name} AS {}", stmt),
944                    options,
945                    false,
946                    "test_db".into(),
947                    "test_table".into(),
948                    format_desc,
949                    false,
950                    false,
951                    false,
952                    None,
953                    None,
954                    false,
955                    None,
956                ) {
957                    Ok(sink_plan) => {
958                        ret.sink_plan = Some(explain_plan(&sink_plan.into()));
959                        break 'sink;
960                    }
961                    Err(err) => {
962                        ret.sink_error = Some(err.to_report_string_pretty());
963                        break 'sink;
964                    }
965                }
966            }
967        }
968
969        Ok(ret)
970    }
971}
972
973fn explain_plan(plan: &PlanRef<impl ConventionMarker>) -> String {
974    plan.explain_to_string()
975}
976
977/// Checks that the result matches `test_case.expected_outputs`.
978///
979/// We don't check the result matches here.
980fn check_result(test_case: &TestCase, actual: &TestCaseResult) -> Result<()> {
981    macro_rules! check {
982        ($field:ident) => {
983            paste::paste! {
984                let case_contains = test_case.expected_outputs.contains(&TestType:: [< $field:camel >]  );
985                let actual_contains = &actual.$field;
986                match (case_contains, actual_contains) {
987                    (false, None) | (true, Some(_)) => {},
988                    (false, Some(e)) => return Err(anyhow!("unexpected {}: {}", stringify!($field), e)),
989                    (true, None) => return Err(anyhow!(
990                        "expected {}, but there's no such result during execution",
991                        stringify!($field)
992                    )),
993                }
994            }
995        };
996    }
997
998    check!(binder_error);
999    check!(planner_error);
1000    check!(optimizer_error);
1001    check!(batch_error);
1002    check!(batch_local_error);
1003    check!(stream_error);
1004    check!(eowc_stream_error);
1005
1006    check!(logical_plan);
1007    check!(optimized_logical_plan_for_batch);
1008    check!(optimized_logical_plan_for_stream);
1009    check!(batch_plan);
1010    check!(batch_local_plan);
1011    check!(stream_plan);
1012    check!(stream_dist_plan);
1013    check!(eowc_stream_plan);
1014    check!(eowc_stream_dist_plan);
1015    check!(batch_plan_proto);
1016    check!(sink_plan);
1017
1018    check!(explain_output);
1019
1020    Ok(())
1021}
1022
1023/// `/tests/testdata` directory.
1024pub fn test_data_dir() -> PathBuf {
1025    std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
1026        .join("tests")
1027        .join("testdata")
1028}
1029
1030pub async fn run_test_file(file_path: &Path, file_content: &str) -> Result<()> {
1031    LicenseManager::get().refresh(LicenseKey::default().as_ref());
1032
1033    let file_name = file_path.file_name().unwrap().to_str().unwrap();
1034    println!("-- running {file_name} --");
1035
1036    let mut failed_num = 0;
1037    let cases: Vec<TestCase> = serde_yaml::from_str(file_content).map_err(|e| {
1038        let context = if let Some(loc) = e.location() {
1039            format!(
1040                "failed to parse yaml at {}:{}:{}",
1041                file_path.display(),
1042                loc.line(),
1043                loc.column()
1044            )
1045        } else {
1046            "failed to parse yaml".to_owned()
1047        };
1048        anyhow::anyhow!(e).context(context)
1049    })?;
1050    let cases = resolve_testcase_id(cases).expect("failed to resolve");
1051    let mut outputs = vec![];
1052
1053    for (i, c) in cases.into_iter().enumerate() {
1054        println!(
1055            "Running test #{i} (id: {}), SQL:\n{}",
1056            c.id().clone().unwrap_or_else(|| "<none>".to_owned()),
1057            c.sql()
1058        );
1059        match Box::pin(c.run(true)).await {
1060            Ok(case) => {
1061                outputs.push(case);
1062            }
1063            Err(e) => {
1064                eprintln!(
1065                    "Test #{i} (id: {}) failed, SQL:\n{}\nError: {}",
1066                    c.id().clone().unwrap_or_else(|| "<none>".to_owned()),
1067                    c.sql(),
1068                    e.as_report()
1069                );
1070                failed_num += 1;
1071            }
1072        }
1073    }
1074
1075    let output_path = test_data_dir().join("output").join(file_name);
1076    check(outputs, expect_test::expect_file![output_path]);
1077
1078    if failed_num > 0 {
1079        println!("\n");
1080        bail!(format!("{} test cases failed", failed_num));
1081    }
1082    Ok(())
1083}
risingwave_planner_test/lib.rs

risingwave_planner_test/
lib.rs
