Skip to content

Commit 8309f9e

Browse files
committed
Add CREATE/DROP RESOURCE and WORKLOAD statement support
Implement parsing and explain output for ClickHouse resource and workload management statements including CREATE RESOURCE, DROP RESOURCE, CREATE WORKLOAD, and DROP WORKLOAD.
1 parent 3f861b6 commit 8309f9e

5 files changed

Lines changed: 171 additions & 64 deletions

File tree

ast/ast.go

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -960,6 +960,45 @@ func (s *ShowCreateRoleQuery) Pos() token.Position { return s.Position }
960960
func (s *ShowCreateRoleQuery) End() token.Position { return s.Position }
961961
func (s *ShowCreateRoleQuery) statementNode() {}
962962

963+
// CreateResourceQuery represents a CREATE RESOURCE statement.
964+
type CreateResourceQuery struct {
965+
Position token.Position `json:"-"`
966+
Name string `json:"name"`
967+
}
968+
969+
func (c *CreateResourceQuery) Pos() token.Position { return c.Position }
970+
func (c *CreateResourceQuery) End() token.Position { return c.Position }
971+
func (c *CreateResourceQuery) statementNode() {}
972+
973+
// DropResourceQuery represents a DROP RESOURCE statement.
974+
type DropResourceQuery struct {
975+
Position token.Position `json:"-"`
976+
}
977+
978+
func (d *DropResourceQuery) Pos() token.Position { return d.Position }
979+
func (d *DropResourceQuery) End() token.Position { return d.Position }
980+
func (d *DropResourceQuery) statementNode() {}
981+
982+
// CreateWorkloadQuery represents a CREATE WORKLOAD statement.
983+
type CreateWorkloadQuery struct {
984+
Position token.Position `json:"-"`
985+
Name string `json:"name"`
986+
Parent string `json:"parent,omitempty"` // Parent workload name (after IN)
987+
}
988+
989+
func (c *CreateWorkloadQuery) Pos() token.Position { return c.Position }
990+
func (c *CreateWorkloadQuery) End() token.Position { return c.Position }
991+
func (c *CreateWorkloadQuery) statementNode() {}
992+
993+
// DropWorkloadQuery represents a DROP WORKLOAD statement.
994+
type DropWorkloadQuery struct {
995+
Position token.Position `json:"-"`
996+
}
997+
998+
func (d *DropWorkloadQuery) Pos() token.Position { return d.Position }
999+
func (d *DropWorkloadQuery) End() token.Position { return d.Position }
1000+
func (d *DropWorkloadQuery) statementNode() {}
1001+
9631002
// CreateIndexQuery represents a CREATE INDEX statement.
9641003
type CreateIndexQuery struct {
9651004
Position token.Position `json:"-"`

internal/explain/explain.go

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -160,6 +160,24 @@ func Node(sb *strings.Builder, node interface{}, depth int) {
160160
} else {
161161
fmt.Fprintf(sb, "%sSHOW CREATE ROLE query\n", indent)
162162
}
163+
case *ast.CreateResourceQuery:
164+
fmt.Fprintf(sb, "%sCreateResourceQuery %s (children 1)\n", indent, n.Name)
165+
childIndent := indent + " "
166+
explainIdentifier(sb, &ast.Identifier{Parts: []string{n.Name}}, childIndent)
167+
case *ast.DropResourceQuery:
168+
fmt.Fprintf(sb, "%sDropResourceQuery\n", indent)
169+
case *ast.CreateWorkloadQuery:
170+
childIndent := indent + " "
171+
if n.Parent != "" {
172+
fmt.Fprintf(sb, "%sCreateWorkloadQuery %s (children 2)\n", indent, n.Name)
173+
explainIdentifier(sb, &ast.Identifier{Parts: []string{n.Name}}, childIndent)
174+
explainIdentifier(sb, &ast.Identifier{Parts: []string{n.Parent}}, childIndent)
175+
} else {
176+
fmt.Fprintf(sb, "%sCreateWorkloadQuery %s (children 1)\n", indent, n.Name)
177+
explainIdentifier(sb, &ast.Identifier{Parts: []string{n.Name}}, childIndent)
178+
}
179+
case *ast.DropWorkloadQuery:
180+
fmt.Fprintf(sb, "%sDropWorkloadQuery\n", indent)
163181
case *ast.ShowGrantsQuery:
164182
fmt.Fprintf(sb, "%sShowGrantsQuery\n", indent)
165183
case *ast.GrantQuery:

parser/parser.go

Lines changed: 112 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -139,6 +139,14 @@ func (p *Parser) parseStatement() ast.Statement {
139139
if p.peek.Token == token.IDENT && strings.ToUpper(p.peek.Value) == "ROLE" {
140140
return p.parseDropRole()
141141
}
142+
// Check for DROP RESOURCE
143+
if p.peek.Token == token.IDENT && strings.ToUpper(p.peek.Value) == "RESOURCE" {
144+
return p.parseDropResource()
145+
}
146+
// Check for DROP WORKLOAD
147+
if p.peek.Token == token.IDENT && strings.ToUpper(p.peek.Value) == "WORKLOAD" {
148+
return p.parseDropWorkload()
149+
}
142150
return p.parseDrop()
143151
case token.ALTER:
144152
// Check for ALTER USER
@@ -1377,7 +1385,13 @@ func (p *Parser) parseCreate() ast.Statement {
13771385
case "ROLE":
13781386
// CREATE ROLE
13791387
return p.parseCreateRole(pos)
1380-
case "RESOURCE", "WORKLOAD", "QUOTA":
1388+
case "RESOURCE":
1389+
// CREATE RESOURCE
1390+
return p.parseCreateResource(pos)
1391+
case "WORKLOAD":
1392+
// CREATE WORKLOAD
1393+
return p.parseCreateWorkload(pos)
1394+
case "QUOTA":
13811395
// Skip these statements - just consume tokens until semicolon
13821396
p.parseCreateGeneric(create)
13831397
default:
@@ -2308,6 +2322,103 @@ func (p *Parser) parseShowCreateRole(pos token.Position) *ast.ShowCreateRoleQuer
23082322
return query
23092323
}
23102324

2325+
func (p *Parser) parseCreateResource(pos token.Position) *ast.CreateResourceQuery {
2326+
query := &ast.CreateResourceQuery{
2327+
Position: pos,
2328+
}
2329+
2330+
// Skip RESOURCE
2331+
if p.currentIs(token.IDENT) && strings.ToUpper(p.current.Value) == "RESOURCE" {
2332+
p.nextToken()
2333+
}
2334+
2335+
// Get resource name
2336+
if p.currentIs(token.IDENT) || p.current.Token.IsKeyword() {
2337+
query.Name = p.current.Value
2338+
p.nextToken()
2339+
}
2340+
2341+
// Skip the rest of the statement (resource definition, etc.)
2342+
for !p.currentIs(token.EOF) && !p.currentIs(token.SEMICOLON) {
2343+
p.nextToken()
2344+
}
2345+
2346+
return query
2347+
}
2348+
2349+
func (p *Parser) parseDropResource() *ast.DropResourceQuery {
2350+
query := &ast.DropResourceQuery{
2351+
Position: p.current.Pos,
2352+
}
2353+
2354+
p.nextToken() // skip DROP
2355+
2356+
// Skip RESOURCE
2357+
if p.currentIs(token.IDENT) && strings.ToUpper(p.current.Value) == "RESOURCE" {
2358+
p.nextToken()
2359+
}
2360+
2361+
// Skip the rest of the statement (IF EXISTS, name, etc.)
2362+
for !p.currentIs(token.EOF) && !p.currentIs(token.SEMICOLON) {
2363+
p.nextToken()
2364+
}
2365+
2366+
return query
2367+
}
2368+
2369+
func (p *Parser) parseCreateWorkload(pos token.Position) *ast.CreateWorkloadQuery {
2370+
query := &ast.CreateWorkloadQuery{
2371+
Position: pos,
2372+
}
2373+
2374+
// Skip WORKLOAD
2375+
if p.currentIs(token.IDENT) && strings.ToUpper(p.current.Value) == "WORKLOAD" {
2376+
p.nextToken()
2377+
}
2378+
2379+
// Get workload name
2380+
if p.currentIs(token.IDENT) || p.current.Token.IsKeyword() {
2381+
query.Name = p.current.Value
2382+
p.nextToken()
2383+
}
2384+
2385+
// Check for IN (parent workload)
2386+
if p.currentIs(token.IN) {
2387+
p.nextToken()
2388+
if p.currentIs(token.IDENT) || p.current.Token.IsKeyword() {
2389+
query.Parent = p.current.Value
2390+
p.nextToken()
2391+
}
2392+
}
2393+
2394+
// Skip the rest of the statement (SETTINGS, etc.)
2395+
for !p.currentIs(token.EOF) && !p.currentIs(token.SEMICOLON) {
2396+
p.nextToken()
2397+
}
2398+
2399+
return query
2400+
}
2401+
2402+
func (p *Parser) parseDropWorkload() *ast.DropWorkloadQuery {
2403+
query := &ast.DropWorkloadQuery{
2404+
Position: p.current.Pos,
2405+
}
2406+
2407+
p.nextToken() // skip DROP
2408+
2409+
// Skip WORKLOAD
2410+
if p.currentIs(token.IDENT) && strings.ToUpper(p.current.Value) == "WORKLOAD" {
2411+
p.nextToken()
2412+
}
2413+
2414+
// Skip the rest of the statement (IF EXISTS, name, etc.)
2415+
for !p.currentIs(token.EOF) && !p.currentIs(token.SEMICOLON) {
2416+
p.nextToken()
2417+
}
2418+
2419+
return query
2420+
}
2421+
23112422
func (p *Parser) parseCreateDictionary(create *ast.CreateQuery) {
23122423
// Handle IF NOT EXISTS
23132424
if p.currentIs(token.IF) {
Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,5 @@
11
{
22
"explain_todo": {
3-
"stmt1": true,
4-
"stmt3": true,
5-
"stmt4": true,
6-
"stmt6": true,
7-
"stmt7": true,
8-
"stmt9": true
3+
"stmt3": true
94
}
105
}
Lines changed: 1 addition & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -1,57 +1 @@
1-
{
2-
"explain_todo": {
3-
"stmt1": true,
4-
"stmt10": true,
5-
"stmt11": true,
6-
"stmt12": true,
7-
"stmt13": true,
8-
"stmt14": true,
9-
"stmt15": true,
10-
"stmt16": true,
11-
"stmt17": true,
12-
"stmt18": true,
13-
"stmt19": true,
14-
"stmt2": true,
15-
"stmt20": true,
16-
"stmt21": true,
17-
"stmt22": true,
18-
"stmt23": true,
19-
"stmt24": true,
20-
"stmt25": true,
21-
"stmt26": true,
22-
"stmt27": true,
23-
"stmt28": true,
24-
"stmt29": true,
25-
"stmt30": true,
26-
"stmt31": true,
27-
"stmt32": true,
28-
"stmt33": true,
29-
"stmt34": true,
30-
"stmt35": true,
31-
"stmt36": true,
32-
"stmt37": true,
33-
"stmt38": true,
34-
"stmt39": true,
35-
"stmt4": true,
36-
"stmt40": true,
37-
"stmt41": true,
38-
"stmt42": true,
39-
"stmt43": true,
40-
"stmt44": true,
41-
"stmt45": true,
42-
"stmt46": true,
43-
"stmt47": true,
44-
"stmt48": true,
45-
"stmt49": true,
46-
"stmt5": true,
47-
"stmt50": true,
48-
"stmt51": true,
49-
"stmt52": true,
50-
"stmt53": true,
51-
"stmt54": true,
52-
"stmt6": true,
53-
"stmt7": true,
54-
"stmt8": true,
55-
"stmt9": true
56-
}
57-
}
1+
{}

0 commit comments

Comments
 (0)