sqlsmith: support DO block for plpgsql

A `DO` block can be either a top-level sql statement or be part of the
plpgsql function body. In this commit we introduce the DO block for
both of these 2 usages. We also introduced `DisableDoBlocks` which
disable both usages of a DO block.

Release note: None

Author: ZhouXing19

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(),

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, _ := makePLpgSQLBlock(s, scope)
+	return &ast.DoBlock{Block: block.(*ast.Block)}, true
+}

pkg/internal/sqlsmith/relational.go

@@ -110,9 +110,11 @@ var (
 		{1, makeImport},
 		{1, makeCreateStats},
 		{1, makeSetSessionCharacteristics},
+		{5, makeDoBlock},
 	}
 	nonMutatingStatements = []statementWeight{
 		{10, makeSelect},
+		{2, 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

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) {