Skip to content

Commit c0216bb

Browse files
shivanshuraj1333shivanshuraj1333
committed
feat: add Session api for stateful connection and set role support
1 parent 757e102 commit c0216bb

File tree

6 files changed

+943
-0
lines changed

6 files changed

+943
-0
lines changed

‎SET_ROLE_SUPPORT.md‎

Lines changed: 304 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,304 @@
1+
# SET ROLE Support for ClickHouse Go Driver
2+
3+
This document describes the implementation of SET ROLE support for the ClickHouse Go driver, addressing the feature request in [#1391](https://github.com/ClickHouse/clickhouse-go/discussions/1391) and [#1443](https://github.com/ClickHouse/clickhouse-go/issues/1443).
4+
5+
## Problem Statement
6+
7+
The current clickhouse-go driver uses connection pooling where each operation acquires a connection from the pool, executes the query, and releases it back to the pool. This design makes it impossible to maintain connection state across multiple operations, which is required for features like `SET ROLE`.
8+
9+
## Solution: Session Management
10+
11+
We've implemented a **Session Management** feature that allows users to acquire and hold a connection for multiple operations while maintaining connection state.
12+
13+
### Key Features
14+
15+
1. **Stateful Connections**: Sessions maintain connection state across multiple operations
16+
2. **Resource Management**: Proper connection pool integration with automatic cleanup
17+
3. **Error Handling**: Comprehensive error handling with specific error types
18+
4. **Debug Logging**: Full debug logging support for troubleshooting
19+
5. **Backward Compatibility**: Additive changes that don't break existing code
20+
21+
## API Design
22+
23+
### New Interface: Session
24+
25+
```go
26+
type Session interface {
27+
// Exec executes a query without returning results
28+
Exec(ctx context.Context, query string, args ...any) error
29+
// Query executes a query and returns rows
30+
Query(ctx context.Context, query string, args ...any) (Rows, error)
31+
// QueryRow executes a query and returns a single row
32+
QueryRow(ctx context.Context, query string, args ...any) Row
33+
// PrepareBatch prepares a batch for insertion
34+
PrepareBatch(ctx context.Context, query string, opts ...PrepareBatchOption) (Batch, error)
35+
// Ping checks if the connection is still alive
36+
Ping(ctx context.Context) error
37+
// Close releases the session back to the connection pool
38+
Close() error
39+
}
40+
```
41+
42+
### New Method: AcquireSession
43+
44+
```go
45+
// AcquireSession acquires a connection from the pool and returns a Session
46+
// that maintains connection state for multiple operations
47+
AcquireSession(ctx context.Context) (Session, error)
48+
```
49+
50+
## Usage Examples
51+
52+
### Basic SET ROLE Usage
53+
54+
```go
55+
package main
56+
57+
import (
58+
"context"
59+
"fmt"
60+
"log"
61+
"time"
62+
63+
"github.com/ClickHouse/clickhouse-go/v2"
64+
)
65+
66+
func main() {
67+
// Open connection
68+
conn, err := clickhouse.Open(&clickhouse.Options{
69+
Addr: []string{"localhost:9000"},
70+
Auth: clickhouse.Auth{
71+
Database: "default",
72+
Username: "default",
73+
Password: "",
74+
},
75+
Settings: clickhouse.Settings{
76+
"max_execution_time": 60,
77+
},
78+
DialTimeout: time.Second * 30,
79+
MaxOpenConns: 5,
80+
MaxIdleConns: 5,
81+
ConnMaxLifetime: time.Hour,
82+
ConnOpenStrategy: clickhouse.ConnOpenInOrder,
83+
Debug: false,
84+
})
85+
if err != nil {
86+
log.Fatal(err)
87+
}
88+
defer conn.Close()
89+
90+
// Acquire a session for stateful operations
91+
session, err := conn.AcquireSession(context.Background())
92+
if err != nil {
93+
log.Fatal(err)
94+
}
95+
defer session.Close()
96+
97+
// Set role for this session
98+
err = session.Exec(context.Background(), "SET ROLE admin")
99+
if err != nil {
100+
log.Fatal(err)
101+
}
102+
103+
// Execute queries with the role applied
104+
rows, err := session.Query(context.Background(), "SELECT currentUser(), currentRole()")
105+
if err != nil {
106+
log.Fatal(err)
107+
}
108+
defer rows.Close()
109+
110+
for rows.Next() {
111+
var user, role string
112+
err := rows.Scan(&user, &role)
113+
if err != nil {
114+
log.Fatal(err)
115+
}
116+
fmt.Printf("User: %s, Role: %s\n", user, role)
117+
}
118+
}
119+
```
120+
121+
### Session State Persistence
122+
123+
```go
124+
// Set session variables that persist across operations
125+
err = session.Exec(context.Background(), "SET max_memory_usage = 1000000")
126+
if err != nil {
127+
log.Fatal(err)
128+
}
129+
130+
// Verify the setting is applied
131+
rows, err := session.Query(context.Background(),
132+
"SELECT value FROM system.settings WHERE name = 'max_memory_usage'")
133+
if err != nil {
134+
log.Fatal(err)
135+
}
136+
defer rows.Close()
137+
138+
if rows.Next() {
139+
var value string
140+
err := rows.Scan(&value)
141+
if err != nil {
142+
log.Fatal(err)
143+
}
144+
fmt.Printf("Max memory usage: %s\n", value)
145+
}
146+
```
147+
148+
### Multiple Sessions Isolation
149+
150+
```go
151+
// Create multiple sessions - each maintains its own state
152+
session1, err := conn.AcquireSession(context.Background())
153+
if err != nil {
154+
log.Fatal(err)
155+
}
156+
defer session1.Close()
157+
158+
session2, err := conn.AcquireSession(context.Background())
159+
if err != nil {
160+
log.Fatal(err)
161+
}
162+
defer session2.Close()
163+
164+
// Set different roles in each session
165+
err = session1.Exec(context.Background(), "SET ROLE admin")
166+
if err != nil {
167+
log.Fatal(err)
168+
}
169+
170+
err = session2.Exec(context.Background(), "SET ROLE readonly")
171+
if err != nil {
172+
log.Fatal(err)
173+
}
174+
175+
// Each session maintains its own state
176+
rows1, err := session1.Query(context.Background(), "SELECT currentRole()")
177+
if err != nil {
178+
log.Fatal(err)
179+
}
180+
defer rows1.Close()
181+
182+
rows2, err := session2.Query(context.Background(), "SELECT currentRole()")
183+
if err != nil {
184+
log.Fatal(err)
185+
}
186+
defer rows2.Close()
187+
188+
// Verify different roles
189+
if rows1.Next() {
190+
var role1 string
191+
rows1.Scan(&role1)
192+
fmt.Printf("Session 1 role: %s\n", role1)
193+
}
194+
195+
if rows2.Next() {
196+
var role2 string
197+
rows2.Scan(&role2)
198+
fmt.Printf("Session 2 role: %s\n", role2)
199+
}
200+
```
201+
202+
### Error Handling
203+
204+
```go
205+
session, err := conn.AcquireSession(context.Background())
206+
if err != nil {
207+
log.Fatal(err)
208+
}
209+
defer session.Close()
210+
211+
// Close the session
212+
session.Close()
213+
214+
// These operations will return ErrSessionClosed
215+
err = session.Exec(context.Background(), "SELECT 1")
216+
if err != nil {
217+
fmt.Printf("Expected error: %v\n", err)
218+
}
219+
220+
_, err = session.Query(context.Background(), "SELECT 1")
221+
if err != nil {
222+
fmt.Printf("Expected error: %v\n", err)
223+
}
224+
```
225+
226+
## Error Types
227+
228+
The implementation introduces specific error types for better error handling:
229+
230+
```go
231+
var (
232+
ErrSessionClosed = errors.New("clickhouse: session is closed")
233+
ErrSessionNotSupported = errors.New("clickhouse: session operations not supported in this context")
234+
)
235+
```
236+
237+
## Resource Management
238+
239+
Sessions properly integrate with the connection pool:
240+
241+
1. **Acquisition**: Sessions acquire connections from the pool
242+
2. **State Maintenance**: Connections maintain state across operations
243+
3. **Release**: Sessions release connections back to the pool when closed
244+
4. **Cleanup**: Automatic cleanup on session close or error
245+
246+
## Debug Logging
247+
248+
Sessions support comprehensive debug logging:
249+
250+
```go
251+
conn, err := clickhouse.Open(&clickhouse.Options{
252+
// ... other options ...
253+
Debug: true,
254+
Debugf: func(format string, v ...any) {
255+
log.Printf("[SESSION] "+format, v...)
256+
},
257+
})
258+
```
259+
260+
Debug output includes:
261+
- Session acquisition and release
262+
- Query execution with SQL
263+
- Error conditions
264+
- Connection state changes
265+
266+
## Testing
267+
268+
Comprehensive tests are provided in `tests/set_role_test.go`:
269+
270+
- Basic session functionality
271+
- SET ROLE operations
272+
- Session state persistence
273+
- Error handling
274+
- Resource management
275+
- Connection pool integration
276+
277+
## Backward Compatibility
278+
279+
This implementation is fully backward compatible:
280+
281+
- No breaking changes to existing APIs
282+
- Sessions are additive functionality
283+
- Existing code continues to work unchanged
284+
- Connection pooling behavior unchanged for non-session operations
285+
286+
## Performance Considerations
287+
288+
- Sessions hold connections longer than regular operations
289+
- Use sessions only when stateful operations are required
290+
- Close sessions promptly to return connections to the pool
291+
- Consider connection pool size when using multiple sessions
292+
293+
## Future Enhancements
294+
295+
Potential future improvements:
296+
297+
1. **Batch Support**: Full batch operation support in sessions
298+
2. **Transaction Integration**: Better integration with database/sql transactions
299+
3. **Session Pooling**: Dedicated session pools for high-throughput scenarios
300+
4. **Configuration Options**: Session-specific configuration options
301+
302+
## Conclusion
303+
304+
This implementation provides a robust, well-tested solution for SET ROLE functionality while maintaining the high standards of the clickhouse-go driver. The design follows established patterns in the codebase and provides a clean, intuitive API for users.

0 commit comments

Comments
 (0)