Skip to content

Commit 9fef4b2

Browse files
craig[bot]Vidit Bhatyuzefovich
committed
144042: drt: add generic framework for SQL object creation and cleanup r=nameisbhaskar a=vidit-bhat This patch introduces a unified framework that allows declarative registration of SQL object operations. Each operation specifies its object type, naming pattern, and creation statement, while cleanup is handled generically with a simple DROP statement. This framework supports `create-role`, `create-user`, `create-type`, and `create-table` operations out of the box. Epic: none Fixes: #138412, #138413, #138416, #138417, #138418, #138419 Release note: None 144058: colbuilder: don't use direct scans when buffered writes are enabled r=yuzefovich a=yuzefovich `txnWriteBuffer` currently doesn't support COL_BATCH_RESPONSE scan format, so if buffered writes are enabled, we won't use the direct scans. We could've relaxed this condition further to not use direct scans if at least some writes are actually buffered, but given this feature is in experimental state, we don't bother doing so for now. Epic: None Release note: None Co-authored-by: Vidit Bhat <[email protected]> Co-authored-by: Yahor Yuzefovich <[email protected]>
3 parents 5581496 + 3dd501e + bfc4772 commit 9fef4b2

File tree

4 files changed

+139
-0
lines changed

4 files changed

+139
-0
lines changed

‎pkg/cmd/roachtest/operations/BUILD.bazel‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ go_library(
2020
"register.go",
2121
"resize.go",
2222
"session_vars.go",
23+
"sql_objects.go",
2324
],
2425
importpath = "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/operations",
2526
visibility = ["//visibility:public"],

‎pkg/cmd/roachtest/operations/register.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ func RegisterOperations(r registry.Registry) {
2020
registerDiskStall(r)
2121
registerNodeKill(r)
2222
registerClusterSettings(r)
23+
registerCreateSQLOperations(r)
2324
registerBackupRestore(r)
2425
registerManualCompaction(r)
2526
registerResize(r)

‎pkg/cmd/roachtest/operations/sql_objects.go‎

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
// Copyright 2024 The Cockroach Authors.
2+
//
3+
// Use of this software is governed by the CockroachDB Software License
4+
// included in the /LICENSE file.
5+
6+
package operations
7+
8+
import (
9+
"context"
10+
"fmt"
11+
"strings"
12+
"time"
13+
14+
"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/cluster"
15+
"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/operation"
16+
"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/option"
17+
"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/registry"
18+
"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/roachtestflags"
19+
"github.com/cockroachdb/cockroach/pkg/util/randutil"
20+
)
21+
22+
// genericSQLCleanup is a generic cleanup struct that drops a SQL object
23+
// (e.g., ROLE, USER, TYPE, TABLE) using a simple DROP statement.
24+
// It assumes the object name was generated and is safe to drop unconditionally.
25+
type genericSQLCleanup struct {
26+
objectType string
27+
objectName string
28+
}
29+
30+
func (cl *genericSQLCleanup) Cleanup(
31+
ctx context.Context, o operation.Operation, c cluster.Cluster,
32+
) {
33+
conn := c.Conn(ctx, o.L(), 1, option.VirtualClusterName(roachtestflags.VirtualCluster))
34+
defer conn.Close()
35+
36+
o.Status(fmt.Sprintf("dropping %s %s", cl.objectType, cl.objectName))
37+
_, err := conn.ExecContext(ctx,
38+
fmt.Sprintf("DROP %s %s", cl.objectType, cl.objectName))
39+
if err != nil {
40+
o.Fatal(err)
41+
}
42+
}
43+
44+
func runSQLOperation(
45+
objectType string, namePrefix string, createSQL func(name string) string,
46+
) func(ctx context.Context, o operation.Operation, c cluster.Cluster) registry.OperationCleanup {
47+
return func(ctx context.Context, o operation.Operation, c cluster.Cluster) registry.OperationCleanup {
48+
conn := c.Conn(ctx, o.L(), 1, option.VirtualClusterName(roachtestflags.VirtualCluster))
49+
defer conn.Close()
50+
51+
rng, _ := randutil.NewPseudoRand()
52+
name := fmt.Sprintf("%s_%d", namePrefix, rng.Uint32())
53+
54+
createStmt := createSQL(name)
55+
56+
o.Status(fmt.Sprintf("creating %s %s", objectType, name))
57+
if _, err := conn.ExecContext(ctx, createStmt); err != nil {
58+
o.Fatal(err)
59+
}
60+
61+
o.Status(fmt.Sprintf("%s %s created", objectType, name))
62+
63+
return &genericSQLCleanup{
64+
objectType: objectType,
65+
objectName: name,
66+
}
67+
}
68+
}
69+
70+
// registerCreateSQLOperations registers a group of SQL object creation operations
71+
// (e.g., create-role, create-user, create-type, create-table) under a shared
72+
// framework that handles execution and cleanup consistently.
73+
func registerCreateSQLOperations(r registry.Registry) {
74+
objs := []struct {
75+
name string
76+
objectType string
77+
prefix string
78+
createSQL func(name string) string
79+
}{
80+
{
81+
name: "create-role",
82+
objectType: "ROLE",
83+
prefix: "roachtest_role",
84+
createSQL: func(name string) string {
85+
return fmt.Sprintf("CREATE ROLE %s", name)
86+
},
87+
},
88+
{
89+
name: "create-user",
90+
objectType: "USER",
91+
prefix: "roachtest_user",
92+
createSQL: func(name string) string {
93+
return fmt.Sprintf("CREATE USER %s", name)
94+
},
95+
},
96+
{
97+
name: "create-type",
98+
objectType: "TYPE",
99+
prefix: "roachtest_enum",
100+
createSQL: func(name string) string {
101+
enumVals := []string{"'one'", "'two'", "'three'"}
102+
return fmt.Sprintf("CREATE TYPE %s AS ENUM (%s)", name, strings.Join(enumVals, ", "))
103+
},
104+
},
105+
{
106+
name: "create-table",
107+
objectType: "TABLE",
108+
prefix: "roachtest_table",
109+
createSQL: func(name string) string {
110+
return fmt.Sprintf("CREATE TABLE %s (id INT PRIMARY KEY, val STRING)", name)
111+
},
112+
},
113+
}
114+
115+
for _, obj := range objs {
116+
r.AddOperation(registry.OperationSpec{
117+
Name: obj.name,
118+
Owner: registry.OwnerSQLFoundations,
119+
Timeout: time.Hour,
120+
CompatibleClouds: registry.AllClouds,
121+
CanRunConcurrently: registry.OperationCanRunConcurrently,
122+
Dependencies: []registry.OperationDependency{registry.OperationRequiresPopulatedDatabase},
123+
Run: runSQLOperation(obj.objectType, obj.prefix, obj.createSQL),
124+
})
125+
}
126+
}

‎pkg/sql/colexec/colbuilder/execplan.go‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -842,6 +842,17 @@ func NewColOperator(
842842
var resultTypes []*types.T
843843
if flowCtx.EvalCtx.SessionData().DirectColumnarScansEnabled {
844844
canUseDirectScan := func() bool {
845+
// txnWriteBuffer currently doesn't support
846+
// COL_BATCH_RESPONSE scan format, so if buffered writes are
847+
// enabled, we won't use the direct scans.
848+
//
849+
// We could've relaxed this condition further to not use
850+
// direct scans if at least some writes are actually
851+
// buffered, but given this feature is in experimental
852+
// state, we don't bother doing so for now.
853+
if flowCtx.Txn.BufferedWritesEnabled() {
854+
return false
855+
}
845856
// We currently don't use the direct scans if TraceKV is
846857
// enabled (due to not being able to tell the KV server
847858
// about it). One idea would be to include this boolean into

0 commit comments

Comments
 (0)