diff --git a/pkg/cmd/smith/main.go b/pkg/cmd/smith/main.go index adb95a360959..a663f7f85153 100644 --- a/pkg/cmd/smith/main.go +++ b/pkg/cmd/smith/main.go @@ -66,6 +66,7 @@ var ( "DisableCrossJoins": sqlsmith.DisableCrossJoins(), "DisableDDLs": sqlsmith.DisableDDLs(), "DisableDecimals": sqlsmith.DisableDecimals(), + "DisableDoBlocks": sqlsmith.DisableDoBlocks(), "DisableEverything": sqlsmith.DisableEverything(), "DisableIndexHints": sqlsmith.DisableIndexHints(), "DisableInsertSelect": sqlsmith.DisableInsertSelect(), diff --git a/pkg/internal/sqlsmith/plpgsql.go b/pkg/internal/sqlsmith/plpgsql.go index e92f3c46fea4..c644d1a83ec8 100644 --- a/pkg/internal/sqlsmith/plpgsql.go +++ b/pkg/internal/sqlsmith/plpgsql.go @@ -141,6 +141,7 @@ var ( {2, makePLpgSQLReturn}, {2, makePLpgSQLIf}, {2, makePLpgSQLWhile}, + {2, makeDoBlockWithinScope}, {2, makePLpgSQLForLoop}, {5, makePLpgSQLNull}, {10, makePLpgSQLAssign}, @@ -395,3 +396,20 @@ func (s *plpgsqlBlockScope) addVariable(name string, typ *types.T, constant bool s.constants[name] = struct{}{} } } + +func (s *Smither) makeDoBlockTreeStmt() (*tree.DoBlock, bool) { + scope := makeBlockScope(0, types.Unknown, tree.RoutineVolatile) + doBlock, ok := makeDoBlockWithinScope(s, scope) + if !ok { + return nil, false + } + return &tree.DoBlock{Code: doBlock.(*ast.DoBlock)}, true +} + +func makeDoBlockWithinScope(s *Smither, scope plpgsqlBlockScope) (ast.Statement, bool) { + if s.disableDoBlock { + return nil, false + } + block := s.makePLpgSQLBlock(scope) + return &ast.DoBlock{Block: block}, true +} diff --git a/pkg/internal/sqlsmith/relational.go b/pkg/internal/sqlsmith/relational.go index 1218dbc4c927..9c9d67c0830c 100644 --- a/pkg/internal/sqlsmith/relational.go +++ b/pkg/internal/sqlsmith/relational.go @@ -110,9 +110,11 @@ var ( {1, makeImport}, {1, makeCreateStats}, {1, makeSetSessionCharacteristics}, + {1, makeDoBlock}, } nonMutatingStatements = []statementWeight{ {10, makeSelect}, + {1, makeDoBlock}, } allStatements = append(mutatingStatements, nonMutatingStatements...) @@ -917,6 +919,10 @@ func makeCreateFunc(s *Smither) (tree.Statement, bool) { return s.makeCreateFunc() } +func makeDoBlock(s *Smither) (tree.Statement, bool) { + return s.makeDoBlockTreeStmt() +} + func makeDropFunc(s *Smither) (tree.Statement, bool) { if s.disableUDFCreation { return nil, false diff --git a/pkg/internal/sqlsmith/sqlsmith.go b/pkg/internal/sqlsmith/sqlsmith.go index 06e69ceef437..19395e8775a0 100644 --- a/pkg/internal/sqlsmith/sqlsmith.go +++ b/pkg/internal/sqlsmith/sqlsmith.go @@ -118,7 +118,11 @@ type Smither struct { // disableUDFCreation indicates whether we're not allowed to create UDFs. // It follows that if we haven't created any UDFs, we have no UDFs to invoke // too. - disableUDFCreation bool + disableUDFCreation bool + // disableDoBlock indicates whether we're not allowed to create DO blocks, + // both as top level statements and inside plpgsql function body + // definition. + disableDoBlock bool disableIsolationChange bool bulkSrv *httptest.Server @@ -612,6 +616,11 @@ var DisableUDFs = simpleOption("disable udfs", func(s *Smither) { s.disableUDFCreation = true }) +// DisableDoBlocks causes the Smither to disable DO blocks. +var DisableDoBlocks = simpleOption("disable do block", func(s *Smither) { + s.disableDoBlock = true +}) + // DisableIsolationChange causes the Smither to disable stmts that modify the // txn isolation level. var DisableIsolationChange = simpleOption("disable isolation change", func(s *Smither) {