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

risingwave_planner_test/
lib.rs
