diff --git a/src/include/statement_generator.hpp b/src/include/statement_generator.hpp index 57e8446..2d8f9e4 100644 --- a/src/include/statement_generator.hpp +++ b/src/include/statement_generator.hpp @@ -98,6 +98,9 @@ class StatementGenerator { unique_ptr GenerateComparison(); unique_ptr GeneratePositionalReference(); unique_ptr GenerateCase(); + unique_ptr GeneratePragma(); + unique_ptr GenerateCopyDatabase(); + unique_ptr GenerateExplain(); unique_ptr GenerateOrderBy(); unique_ptr GenerateOrderByAll(); diff --git a/src/statement_generator.cpp b/src/statement_generator.cpp index 455833b..1e967e2 100644 --- a/src/statement_generator.cpp +++ b/src/statement_generator.cpp @@ -13,14 +13,18 @@ #include "duckdb/parser/parsed_data/create_type_info.hpp" #include "duckdb/parser/parsed_data/create_view_info.hpp" #include "duckdb/parser/parsed_expression_iterator.hpp" +#include "duckdb/parser/statement/explain_statement.hpp" #include "duckdb/parser/query_node/select_node.hpp" #include "duckdb/parser/query_node/set_operation_node.hpp" #include "duckdb/parser/statement/attach_statement.hpp" #include "duckdb/parser/statement/create_statement.hpp" +#include "duckdb/parser/statement/copy_statement.hpp" +#include "duckdb/parser/statement/copy_database_statement.hpp" #include "duckdb/parser/statement/delete_statement.hpp" #include "duckdb/parser/statement/detach_statement.hpp" #include "duckdb/parser/statement/insert_statement.hpp" #include "duckdb/parser/statement/multi_statement.hpp" +#include "duckdb/parser/statement/pragma_statement.hpp" #include "duckdb/parser/statement/select_statement.hpp" #include "duckdb/parser/statement/set_statement.hpp" #include "duckdb/parser/statement/update_statement.hpp" @@ -106,12 +110,21 @@ unique_ptr StatementGenerator::GenerateStatement() { if (RandomPercentage(60)) { return GenerateStatement(StatementType::DETACH_STATEMENT); } + if (RandomPercentage(50)) { + return GenerateStatement(StatementType::PRAGMA_STATEMENT); + } if (RandomPercentage(30)) { return GenerateStatement(StatementType::SET_STATEMENT); } if (RandomPercentage(40)) { // 20 return GenerateStatement(StatementType::DELETE_STATEMENT); } + if (RandomPercentage(10)) { + return GenerateCopyDatabase(); + } + if (RandomPercentage(20)) { + return GenerateStatement(StatementType::EXPLAIN_STATEMENT); + } return GenerateStatement(StatementType::CREATE_STATEMENT); } @@ -130,6 +143,12 @@ unique_ptr StatementGenerator::GenerateStatement(StatementType typ return GenerateSet(); case StatementType::DELETE_STATEMENT: return GenerateDelete(); + case StatementType::PRAGMA_STATEMENT: + return GeneratePragma(); + case StatementType::COPY_DATABASE_STATEMENT: + return GenerateCopyDatabase(); + case StatementType::EXPLAIN_STATEMENT: + return GenerateExplain(); default: throw InternalException("Unsupported type"); } @@ -200,6 +219,60 @@ unique_ptr StatementGenerator::GenerateDelete() { return delete_statement; } +//===--------------------------------------------------------------------===// +// Generate Pragma +//===--------------------------------------------------------------------===// + +unique_ptr StatementGenerator::GeneratePragma() { + auto pragma_stmt = make_uniq(); + pragma_stmt->info = make_uniq(); + // getting a random pragma_function + if (!generator_context->pragma_functions.empty()) { + auto &entry = Choose(generator_context->pragma_functions).get(); + pragma_stmt->info->name = entry.name; + } else { + pragma_stmt->info->name = "enable_progress_bar"; + } + if (RandomPercentage(50)) { + pragma_stmt->info->parameters.push_back(GenerateConstant()); + } + return pragma_stmt; +} + +//===--------------------------------------------------------------------===// +// Copy Database Statement +//===--------------------------------------------------------------------===// + +unique_ptr StatementGenerator::GenerateCopyDatabase() { + auto from_db = GetRandomAttachedDataBase(); + auto to_db = string("db_") + RandomString(6); + auto mode = RandomPercentage(50) ? CopyDatabaseType::COPY_SCHEMA : CopyDatabaseType::COPY_DATA; + return make_uniq(std::move(from_db), std::move(to_db), mode); +} + +//===--------------------------------------------------------------------===// +// Explain Statement +//===--------------------------------------------------------------------===// + +unique_ptr StatementGenerator::GenerateExplain() { + unique_ptr payload; + if (RandomPercentage(70)) { + payload = GenerateStatement(StatementType::SELECT_STATEMENT); + } else { + payload = GenerateStatement(Choose({ + StatementType::DELETE_STATEMENT, + StatementType::UPDATE_STATEMENT, + StatementType::CREATE_STATEMENT + })); + } + auto stmt = make_uniq( + std::move(payload), + ExplainType::EXPLAIN_STANDARD, + ExplainFormat::DEFAULT + ); + return stmt; +} + //===--------------------------------------------------------------------===// // Generate Detach Info //===--------------------------------------------------------------------===//