risingwave_sqlsmith/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#![feature(box_patterns)]
16#![feature(register_tool)]
17#![register_tool(rw)]
18#![allow(rw::format_error)] // test code
19
20risingwave_expr_impl::enable!();
21
22use std::collections::{HashMap, HashSet};
23
24use anyhow::{Result, bail};
25use config::Configuration;
26use itertools::Itertools;
27use rand::Rng;
28use rand::prelude::IndexedRandom;
29use risingwave_sqlparser::ast::{
30 BinaryOperator, ColumnOption, Expr, Join, JoinConstraint, JoinOperator, Statement,
31 TableConstraint,
32};
33use risingwave_sqlparser::parser::Parser;
34
35use crate::sql_gen::SqlGenerator;
36
37pub mod config;
38pub mod reducer;
39mod sql_gen;
40pub mod sqlreduce;
41pub mod test_runners;
42mod utils;
43pub mod validation;
44pub use validation::is_permissible_error;
45
46pub use crate::sql_gen::{Table, print_function_table};
47
48/// Generate a random SQL string.
49pub fn sql_gen(rng: &mut impl Rng, tables: Vec<Table>, config: &Configuration) -> String {
50 let mut r#gen = SqlGenerator::new(rng, tables, config.clone());
51 format!("{}", r#gen.gen_batch_query_stmt())
52}
53
54/// Generate `INSERT`
55pub fn insert_sql_gen(
56 rng: &mut impl Rng,
57 tables: Vec<Table>,
58 count: usize,
59 config: &Configuration,
60) -> Vec<String> {
61 let mut r#gen = SqlGenerator::new(rng, vec![], config.clone());
62 tables
63 .into_iter()
64 .map(|table| format!("{}", r#gen.generate_insert_statement(&table, count)))
65 .collect()
66}
67
68/// Generate a random CREATE MATERIALIZED VIEW sql string.
69/// These are derived from `tables`.
70pub fn mview_sql_gen<R: Rng>(
71 rng: &mut R,
72 tables: Vec<Table>,
73 name: &str,
74 config: &Configuration,
75) -> (String, Table) {
76 let mut r#gen = SqlGenerator::new_for_mview(rng, tables, config.clone());
77 let (mview, table) = r#gen.gen_mview_stmt(name);
78 (mview.to_string(), table)
79}
80
81pub fn differential_sql_gen<R: Rng>(
82 rng: &mut R,
83 tables: Vec<Table>,
84 name: &str,
85 config: &Configuration,
86) -> Result<(String, String, Table)> {
87 let mut r#gen = SqlGenerator::new_for_mview(rng, tables, config.clone());
88 let (stream, table) = r#gen.gen_mview_stmt(name);
89 let batch = match stream {
90 Statement::CreateView { ref query, .. } => query.to_string(),
91 _ => bail!("Differential pair should be mview statement!"),
92 };
93 Ok((batch, stream.to_string(), table))
94}
95
96/// TODO(noel): Eventually all session variables should be fuzzed.
97/// For now we start of with a few hardcoded configs.
98/// Some config need workarounds, for instance `QUERY_MODE`,
99/// which can lead to stack overflow
100/// (a simple workaround is limit length of
101/// generated query when `QUERY_MODE=local`.
102pub fn session_sql_gen<R: Rng>(rng: &mut R) -> String {
103 [
104 "SET ENABLE_TWO_PHASE_AGG TO TRUE",
105 "SET ENABLE_TWO_PHASE_AGG TO FALSE",
106 "SET RW_FORCE_TWO_PHASE_AGG TO TRUE",
107 "SET RW_FORCE_TWO_PHASE_AGG TO FALSE",
108 ]
109 .choose(rng)
110 .unwrap()
111 .to_string()
112}
113
114pub fn generate_update_statements<R: Rng>(
115 rng: &mut R,
116 tables: &[Table],
117 inserts: &[Statement],
118 config: &Configuration,
119) -> Result<Vec<Statement>> {
120 let mut r#gen = SqlGenerator::new(rng, vec![], config.clone());
121 r#gen.generate_update_statements(tables, inserts)
122}
123
124/// Parse SQL
125/// FIXME(Noel): Introduce error type for sqlsmith for this.
126pub fn parse_sql<S: AsRef<str>>(sql: S) -> Vec<Statement> {
127 let sql = sql.as_ref();
128 Parser::parse_sql(sql).unwrap_or_else(|_| panic!("Failed to parse SQL: {}", sql))
129}
130
131/// Extract relevant info from CREATE TABLE statement, to construct a Table
132pub fn create_table_statement_to_table(statement: &Statement) -> Table {
133 match statement {
134 Statement::CreateTable {
135 name,
136 columns,
137 constraints,
138 append_only,
139 source_watermarks,
140 ..
141 } => {
142 let column_name_to_index_mapping: HashMap<_, _> = columns
143 .iter()
144 .enumerate()
145 .map(|(i, c)| (&c.name, i))
146 .collect();
147 let mut pk_indices = HashSet::new();
148 for (i, column) in columns.iter().enumerate() {
149 let is_primary_key = column
150 .options
151 .iter()
152 .any(|option| option.option == ColumnOption::Unique { is_primary: true });
153 if is_primary_key {
154 pk_indices.insert(i);
155 }
156 }
157 for constraint in constraints {
158 if let TableConstraint::Unique {
159 columns,
160 is_primary: true,
161 ..
162 } = constraint
163 {
164 for column in columns {
165 let pk_index = column_name_to_index_mapping.get(column).unwrap();
166 pk_indices.insert(*pk_index);
167 }
168 }
169 }
170 let mut pk_indices = pk_indices.into_iter().collect_vec();
171 pk_indices.sort_unstable();
172 Table::new_for_base_table(
173 name.0[0].real_value(),
174 columns.iter().map(|c| c.clone().into()).collect(),
175 pk_indices,
176 *append_only,
177 source_watermarks.clone(),
178 )
179 }
180 _ => panic!(
181 "Only CREATE TABLE statements permitted, received: {}",
182 statement
183 ),
184 }
185}
186
187pub fn parse_create_table_statements(sql: impl AsRef<str>) -> (Vec<Table>, Vec<Statement>) {
188 let statements = parse_sql(&sql);
189 let tables = statements
190 .iter()
191 .map(create_table_statement_to_table)
192 .collect();
193 (tables, statements)
194}
195
196#[cfg(test)]
197mod tests {
198 use std::fmt::Debug;
199
200 use expect_test::{Expect, expect};
201
202 use super::*;
203
204 fn check(actual: impl Debug, expect: Expect) {
205 let actual = format!("{:#?}", actual);
206 expect.assert_eq(&actual);
207 }
208
209 #[test]
210 fn test_parse_create_table_statements_no_pk() {
211 let test_string = "
212CREATE TABLE t(v1 int);
213CREATE TABLE t2(v1 int, v2 bool);
214CREATE TABLE t3(v1 int, v2 bool, v3 smallint);
215 ";
216 check(
217 parse_create_table_statements(test_string),
218 expect![[r#"
219 (
220 [
221 Table {
222 name: "t",
223 columns: [
224 Column {
225 name: ObjectName(
226 [
227 Ident {
228 value: "v1",
229 quote_style: None,
230 },
231 ],
232 ),
233 data_type: Int32,
234 },
235 ],
236 pk_indices: [],
237 is_base_table: true,
238 is_append_only: false,
239 source_watermarks: [],
240 },
241 Table {
242 name: "t2",
243 columns: [
244 Column {
245 name: ObjectName(
246 [
247 Ident {
248 value: "v1",
249 quote_style: None,
250 },
251 ],
252 ),
253 data_type: Int32,
254 },
255 Column {
256 name: ObjectName(
257 [
258 Ident {
259 value: "v2",
260 quote_style: None,
261 },
262 ],
263 ),
264 data_type: Boolean,
265 },
266 ],
267 pk_indices: [],
268 is_base_table: true,
269 is_append_only: false,
270 source_watermarks: [],
271 },
272 Table {
273 name: "t3",
274 columns: [
275 Column {
276 name: ObjectName(
277 [
278 Ident {
279 value: "v1",
280 quote_style: None,
281 },
282 ],
283 ),
284 data_type: Int32,
285 },
286 Column {
287 name: ObjectName(
288 [
289 Ident {
290 value: "v2",
291 quote_style: None,
292 },
293 ],
294 ),
295 data_type: Boolean,
296 },
297 Column {
298 name: ObjectName(
299 [
300 Ident {
301 value: "v3",
302 quote_style: None,
303 },
304 ],
305 ),
306 data_type: Int16,
307 },
308 ],
309 pk_indices: [],
310 is_base_table: true,
311 is_append_only: false,
312 source_watermarks: [],
313 },
314 ],
315 [
316 CreateTable {
317 or_replace: false,
318 temporary: false,
319 if_not_exists: false,
320 name: ObjectName(
321 [
322 Ident {
323 value: "t",
324 quote_style: None,
325 },
326 ],
327 ),
328 columns: [
329 ColumnDef {
330 name: Ident {
331 value: "v1",
332 quote_style: None,
333 },
334 data_type: Some(
335 Int,
336 ),
337 collation: None,
338 options: [],
339 },
340 ],
341 wildcard_idx: None,
342 constraints: [],
343 with_options: [],
344 format_encode: None,
345 source_watermarks: [],
346 append_only: false,
347 on_conflict: None,
348 with_version_columns: [],
349 query: None,
350 cdc_table_info: None,
351 include_column_options: [],
352 webhook_info: None,
353 engine: Hummock,
354 },
355 CreateTable {
356 or_replace: false,
357 temporary: false,
358 if_not_exists: false,
359 name: ObjectName(
360 [
361 Ident {
362 value: "t2",
363 quote_style: None,
364 },
365 ],
366 ),
367 columns: [
368 ColumnDef {
369 name: Ident {
370 value: "v1",
371 quote_style: None,
372 },
373 data_type: Some(
374 Int,
375 ),
376 collation: None,
377 options: [],
378 },
379 ColumnDef {
380 name: Ident {
381 value: "v2",
382 quote_style: None,
383 },
384 data_type: Some(
385 Boolean,
386 ),
387 collation: None,
388 options: [],
389 },
390 ],
391 wildcard_idx: None,
392 constraints: [],
393 with_options: [],
394 format_encode: None,
395 source_watermarks: [],
396 append_only: false,
397 on_conflict: None,
398 with_version_columns: [],
399 query: None,
400 cdc_table_info: None,
401 include_column_options: [],
402 webhook_info: None,
403 engine: Hummock,
404 },
405 CreateTable {
406 or_replace: false,
407 temporary: false,
408 if_not_exists: false,
409 name: ObjectName(
410 [
411 Ident {
412 value: "t3",
413 quote_style: None,
414 },
415 ],
416 ),
417 columns: [
418 ColumnDef {
419 name: Ident {
420 value: "v1",
421 quote_style: None,
422 },
423 data_type: Some(
424 Int,
425 ),
426 collation: None,
427 options: [],
428 },
429 ColumnDef {
430 name: Ident {
431 value: "v2",
432 quote_style: None,
433 },
434 data_type: Some(
435 Boolean,
436 ),
437 collation: None,
438 options: [],
439 },
440 ColumnDef {
441 name: Ident {
442 value: "v3",
443 quote_style: None,
444 },
445 data_type: Some(
446 SmallInt,
447 ),
448 collation: None,
449 options: [],
450 },
451 ],
452 wildcard_idx: None,
453 constraints: [],
454 with_options: [],
455 format_encode: None,
456 source_watermarks: [],
457 append_only: false,
458 on_conflict: None,
459 with_version_columns: [],
460 query: None,
461 cdc_table_info: None,
462 include_column_options: [],
463 webhook_info: None,
464 engine: Hummock,
465 },
466 ],
467 )"#]],
468 );
469 }
470
471 #[test]
472 fn test_parse_create_table_statements_with_pk() {
473 let test_string = "
474CREATE TABLE t(v1 int PRIMARY KEY);
475CREATE TABLE t2(v1 int, v2 smallint PRIMARY KEY);
476CREATE TABLE t3(v1 int PRIMARY KEY, v2 smallint PRIMARY KEY);
477CREATE TABLE t4(v1 int PRIMARY KEY, v2 smallint PRIMARY KEY, v3 bool PRIMARY KEY);
478";
479 check(
480 parse_create_table_statements(test_string),
481 expect![[r#"
482 (
483 [
484 Table {
485 name: "t",
486 columns: [
487 Column {
488 name: ObjectName(
489 [
490 Ident {
491 value: "v1",
492 quote_style: None,
493 },
494 ],
495 ),
496 data_type: Int32,
497 },
498 ],
499 pk_indices: [
500 0,
501 ],
502 is_base_table: true,
503 is_append_only: false,
504 source_watermarks: [],
505 },
506 Table {
507 name: "t2",
508 columns: [
509 Column {
510 name: ObjectName(
511 [
512 Ident {
513 value: "v1",
514 quote_style: None,
515 },
516 ],
517 ),
518 data_type: Int32,
519 },
520 Column {
521 name: ObjectName(
522 [
523 Ident {
524 value: "v2",
525 quote_style: None,
526 },
527 ],
528 ),
529 data_type: Int16,
530 },
531 ],
532 pk_indices: [
533 1,
534 ],
535 is_base_table: true,
536 is_append_only: false,
537 source_watermarks: [],
538 },
539 Table {
540 name: "t3",
541 columns: [
542 Column {
543 name: ObjectName(
544 [
545 Ident {
546 value: "v1",
547 quote_style: None,
548 },
549 ],
550 ),
551 data_type: Int32,
552 },
553 Column {
554 name: ObjectName(
555 [
556 Ident {
557 value: "v2",
558 quote_style: None,
559 },
560 ],
561 ),
562 data_type: Int16,
563 },
564 ],
565 pk_indices: [
566 0,
567 1,
568 ],
569 is_base_table: true,
570 is_append_only: false,
571 source_watermarks: [],
572 },
573 Table {
574 name: "t4",
575 columns: [
576 Column {
577 name: ObjectName(
578 [
579 Ident {
580 value: "v1",
581 quote_style: None,
582 },
583 ],
584 ),
585 data_type: Int32,
586 },
587 Column {
588 name: ObjectName(
589 [
590 Ident {
591 value: "v2",
592 quote_style: None,
593 },
594 ],
595 ),
596 data_type: Int16,
597 },
598 Column {
599 name: ObjectName(
600 [
601 Ident {
602 value: "v3",
603 quote_style: None,
604 },
605 ],
606 ),
607 data_type: Boolean,
608 },
609 ],
610 pk_indices: [
611 0,
612 1,
613 2,
614 ],
615 is_base_table: true,
616 is_append_only: false,
617 source_watermarks: [],
618 },
619 ],
620 [
621 CreateTable {
622 or_replace: false,
623 temporary: false,
624 if_not_exists: false,
625 name: ObjectName(
626 [
627 Ident {
628 value: "t",
629 quote_style: None,
630 },
631 ],
632 ),
633 columns: [
634 ColumnDef {
635 name: Ident {
636 value: "v1",
637 quote_style: None,
638 },
639 data_type: Some(
640 Int,
641 ),
642 collation: None,
643 options: [
644 ColumnOptionDef {
645 name: None,
646 option: Unique {
647 is_primary: true,
648 },
649 },
650 ],
651 },
652 ],
653 wildcard_idx: None,
654 constraints: [],
655 with_options: [],
656 format_encode: None,
657 source_watermarks: [],
658 append_only: false,
659 on_conflict: None,
660 with_version_columns: [],
661 query: None,
662 cdc_table_info: None,
663 include_column_options: [],
664 webhook_info: None,
665 engine: Hummock,
666 },
667 CreateTable {
668 or_replace: false,
669 temporary: false,
670 if_not_exists: false,
671 name: ObjectName(
672 [
673 Ident {
674 value: "t2",
675 quote_style: None,
676 },
677 ],
678 ),
679 columns: [
680 ColumnDef {
681 name: Ident {
682 value: "v1",
683 quote_style: None,
684 },
685 data_type: Some(
686 Int,
687 ),
688 collation: None,
689 options: [],
690 },
691 ColumnDef {
692 name: Ident {
693 value: "v2",
694 quote_style: None,
695 },
696 data_type: Some(
697 SmallInt,
698 ),
699 collation: None,
700 options: [
701 ColumnOptionDef {
702 name: None,
703 option: Unique {
704 is_primary: true,
705 },
706 },
707 ],
708 },
709 ],
710 wildcard_idx: None,
711 constraints: [],
712 with_options: [],
713 format_encode: None,
714 source_watermarks: [],
715 append_only: false,
716 on_conflict: None,
717 with_version_columns: [],
718 query: None,
719 cdc_table_info: None,
720 include_column_options: [],
721 webhook_info: None,
722 engine: Hummock,
723 },
724 CreateTable {
725 or_replace: false,
726 temporary: false,
727 if_not_exists: false,
728 name: ObjectName(
729 [
730 Ident {
731 value: "t3",
732 quote_style: None,
733 },
734 ],
735 ),
736 columns: [
737 ColumnDef {
738 name: Ident {
739 value: "v1",
740 quote_style: None,
741 },
742 data_type: Some(
743 Int,
744 ),
745 collation: None,
746 options: [
747 ColumnOptionDef {
748 name: None,
749 option: Unique {
750 is_primary: true,
751 },
752 },
753 ],
754 },
755 ColumnDef {
756 name: Ident {
757 value: "v2",
758 quote_style: None,
759 },
760 data_type: Some(
761 SmallInt,
762 ),
763 collation: None,
764 options: [
765 ColumnOptionDef {
766 name: None,
767 option: Unique {
768 is_primary: true,
769 },
770 },
771 ],
772 },
773 ],
774 wildcard_idx: None,
775 constraints: [],
776 with_options: [],
777 format_encode: None,
778 source_watermarks: [],
779 append_only: false,
780 on_conflict: None,
781 with_version_columns: [],
782 query: None,
783 cdc_table_info: None,
784 include_column_options: [],
785 webhook_info: None,
786 engine: Hummock,
787 },
788 CreateTable {
789 or_replace: false,
790 temporary: false,
791 if_not_exists: false,
792 name: ObjectName(
793 [
794 Ident {
795 value: "t4",
796 quote_style: None,
797 },
798 ],
799 ),
800 columns: [
801 ColumnDef {
802 name: Ident {
803 value: "v1",
804 quote_style: None,
805 },
806 data_type: Some(
807 Int,
808 ),
809 collation: None,
810 options: [
811 ColumnOptionDef {
812 name: None,
813 option: Unique {
814 is_primary: true,
815 },
816 },
817 ],
818 },
819 ColumnDef {
820 name: Ident {
821 value: "v2",
822 quote_style: None,
823 },
824 data_type: Some(
825 SmallInt,
826 ),
827 collation: None,
828 options: [
829 ColumnOptionDef {
830 name: None,
831 option: Unique {
832 is_primary: true,
833 },
834 },
835 ],
836 },
837 ColumnDef {
838 name: Ident {
839 value: "v3",
840 quote_style: None,
841 },
842 data_type: Some(
843 Boolean,
844 ),
845 collation: None,
846 options: [
847 ColumnOptionDef {
848 name: None,
849 option: Unique {
850 is_primary: true,
851 },
852 },
853 ],
854 },
855 ],
856 wildcard_idx: None,
857 constraints: [],
858 with_options: [],
859 format_encode: None,
860 source_watermarks: [],
861 append_only: false,
862 on_conflict: None,
863 with_version_columns: [],
864 query: None,
865 cdc_table_info: None,
866 include_column_options: [],
867 webhook_info: None,
868 engine: Hummock,
869 },
870 ],
871 )"#]],
872 );
873 }
874
875 #[test]
876 fn test_parse_create_table_statements_with_append_only_and_watermark() {
877 let test_string = r#"
878 CREATE TABLE t1 (
879 v1 INT,
880 ts TIMESTAMP,
881 WATERMARK FOR ts AS ts - INTERVAL '5' SECOND
882 ) APPEND ONLY;
883
884 CREATE TABLE t2 (
885 v1 INT,
886 ts TIMESTAMP
887 ) APPEND ONLY;
888
889 CREATE TABLE t3 (
890 v1 INT,
891 ts TIMESTAMP
892 );
893 "#;
894
895 let (tables, _) = parse_create_table_statements(test_string);
896
897 // Check t1
898 let t1 = &tables[0];
899 assert!(t1.is_append_only);
900 assert_eq!(t1.source_watermarks.len(), 1);
901 assert_eq!(t1.source_watermarks[0].column.real_value(), "ts");
902
903 // Check t2
904 let t2 = &tables[1];
905 assert!(t2.is_append_only);
906 assert_eq!(t2.source_watermarks.len(), 0);
907
908 // Check t3
909 let t3 = &tables[2];
910 assert!(!t3.is_append_only);
911 assert_eq!(t3.source_watermarks.len(), 0);
912 }
913}