Skip to content

feat: add Session api for stateful connection and set role support #1597

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 4 commits into
base: main
Choose a base branch
from

Conversation

shivanshuraj1333
Copy link
Contributor

@shivanshuraj1333 shivanshuraj1333 commented Jul 13, 2025

fixes: #1443

please check the end to end implementation with this toy program https://github.com/shivanshuraj1333/clickhuose-test

make sure to update the go.mod to point to the changes present in this PR

Summary:

  • Problem: ClickHouse Go driver lacks support for SET ROLE functionality due to connection pooling, which doesn't maintain connection state across multiple statements.

  • Solution: Added session management API that provides stateful connections while maintaining backward compatibility.

Key Changes

  • New Session Interface: Added Session interface with SetRole(), UnsetRole(), and Close() methods
  • Connection Enhancement: Added AcquireSession() method to main connection interface

Implementation:

  • session struct for native driver
  • stdSession struct for standard SQL driver
  • Proper resource management and error handling

example usage:

// Acquire a stateful session
session, err := conn.AcquireSession()
if err != nil {
    return err
}
defer session.Close()

// Set role for this session
if err := session.SetRole("admin"); err != nil {
    return err
}

// Execute queries - role persists across operations
rows, err := session.Query("SELECT currentUser()")
// ... more operations

// Unset role when done
session.UnsetRole()

@shivanshuraj1333
Copy link
Contributor Author

@SpencerTorres need your help with reviews on this one.
Thanks!

@SpencerTorres
Copy link
Member

Nice! This is pretty much how I'd do it.

A few thoughts:

  • How does this behave when the protocol is HTTP? HTTP has a session_id parameter, so users on HTTP might be mislead to think that their session is unique, but each request is separate in this case. It might be good to override the HTTP session ID in the driver's context.Context wrapper.
  • The added documentation file is good, we should add it to our https://github.com/ClickHouse/clickhouse-docs repo so it goes on the docs website.
  • In relation to usage, I would definitely be very clear with the semantics of holding a session. The session MUST call .Close() or else the connections will leak. Also what happens when an error occurs within a session, such as the connection closing? Handling these cases should be clearly documented for users.
  • It's good to have this feature, but we don't want to lead users into writing dangerous code for things like setting roles. Perhaps the connection should be completely closed before returning it to the pool. This way there's no permission issues in cases of multi-tenant code, or at least we should document it clearly that they should open a whole new pool to get a truly clear session.

Let me know if you have any thoughts. I appreciate you submitting this and other PRs, these are useful additions!

@shivanshuraj1333
Copy link
Contributor Author

These are some great feedbacks, IMO if we want to role this out, we can do this Iteratively like below.

  1. => AFAIU
    Now, when using HTTP, driver does not automatically manage session_id for sessions, so session state only persists if the user manually sets and reuses the same session_id for all requests, else each HTTP request is stateless even when using the session API. Maybe we can document that and as an enhancement add a followup which updates the driver so that it will automatically generate and inject a unique session_id for each HTTP session. Does that sound right approach to you?

  2. => Sure

  3. => Shall we add these in docs?
    Always call .Close() on your session when done, to avoid leaking connections.
    If a session encounters a fatal error (e.g., connection closed), it becomes unusable and should be closed. Do not attempt to reuse a session after a fatal error.

  4. => Have made a slight change which closes the connection instead of returning it to the pool

lmk your thoughts, and thanks for the feedbacks!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Support SET ROLE for clickhouse.Conn
2 participants