Add support for catalog roles

README.md

@@ -249,6 +249,30 @@ Default: nil
 
 The `queryTimeout` parameter sets a timeout for the query. If the query takes longer than the timeout, it will be cancelled. If it is not set the default context timeout will be used.
 
+
+#### `roles`
+
+```
+Type: string 
+Format: roles=catalog1:role1;catalog2=role2 
+Valid values: A semicolon-separated list of catalog-to-role assignments, where each assignment maps a catalog to a role. 
+Default: empty
+```
+The roles parameter defines authorization roles to assume for one or more catalogs during the Trino session.
+
+You can assign roles either as a map of catalog-to-role pairs or a string direcly in the dns connection.
+
+##### Example
+``` go
+c := &Config{
+ServerURI: "https://foobar@localhost:8090",
+SessionProperties: map[string]string{"query_priority": "1"},
+Roles: map[string]string{"catalog1": "role1", "catalog2": "role2"},
+}
+
+dsn, err := c.FormatDSN()
+```
+
 #### Examples
 
 ```
@@ -259,6 +283,11 @@ http://user@localhost:8080?source=hello&catalog=default&schema=foobar
 https://user@localhost:8443?session_properties=query_max_run_time=10m,query_priority=2
 ```
 
+
+```
+http://user@localhost:8080?source=hello&catalog=default&schema=foobar&roles=catalog1:role1;catalog2:role2
+```
+
 ## Data types
 
 ### Query arguments

trino/integration_test.go

@@ -49,6 +49,7 @@ import (
 "github.com/golang-jwt/jwt/v5"
 dt "github.com/ory/dockertest/v3"
 docker "github.com/ory/dockertest/v3/docker"
+"github.com/stretchr/testify/require"
 )
 
 const (
@@ -1024,6 +1025,77 @@ func TestIntegrationNoResults(t *testing.T) {
 t.Fatal(err)
 }
 }
+func TestRoleHeaderSupport(t *testing.T) {
+tests := []struct {
+name string
+config Config
+rawDSN string
+expectError bool
+errorSubstr string
+}{
+{
+name: "Valid roles via Config",
+config: Config{
+ServerURI: *integrationServerFlag,
+Roles: map[string]string{"tpch": "role1", "memory": "role2"},
+},
+expectError: false,
+},
+{
+name: "Valid roles via DSN, not encoded url",
+rawDSN: *integrationServerFlag + "?roles=tpch:role1;memory:role2",
+expectError: false,
+},
+{
+name: "Valid roles via DSN, url encoded",
+rawDSN: *integrationServerFlag + "?roles%3Dtpch%3Arole1%3Bmemory%3Arole2",
+expectError: false,
+},
+{
+name: "Non-existent catalog role",
+config: Config{
+ServerURI: *integrationServerFlag,
+Roles: map[string]string{"not-exist-catalog": "role1"},
+},
+expectError: true,
+errorSubstr: "USER_ERROR: Catalog",
+},
+{
+name: "Invalid role format missing ROLE{}",
+rawDSN: *integrationServerFlag + "?roles=catolog%3Drole1",
+expectError: true,
+errorSubstr: "Invalid role format: catolog=role1",
+},
+}
+
+for _, tt := range tests {
+t.Run(tt.name, func(t *testing.T) {
+var dns string
+var err error
+
+if tt.rawDSN != "" {
+dns = tt.rawDSN
+} else {
+dns, err = tt.config.FormatDSN()
+if err != nil {
+t.Fatal(err)
+}
+}
+
+db := integrationOpen(t, dns)
+_, err = db.Query("SELECT 1")
+
+if tt.expectError {
+require.Error(t, err)
+if tt.errorSubstr != "" {
+require.Contains(t, err.Error(), tt.errorSubstr)
+}
+} else {
+require.NoError(t, err)
+}
+})
+}
+}
 
 func TestIntegrationQueryParametersSelect(t *testing.T) {
 scenarios := []struct {

trino/trino.go

@@ -126,6 +126,7 @@ const (
 trinoSetSessionHeader = trinoHeaderPrefix + `Set-Session`
 trinoClearSessionHeader = trinoHeaderPrefix + `Clear-Session`
 trinoSetRoleHeader = trinoHeaderPrefix + `Set-Role`
+trinoRoleHeader = trinoHeaderPrefix + `Role`
 trinoExtraCredentialHeader = trinoHeaderPrefix + `Extra-Credential`
 
 trinoProgressCallbackParam = trinoHeaderPrefix + `Progress-Callback`
@@ -153,6 +154,7 @@ const (
 
 mapKeySeparator = ":"
 mapEntrySeparator = ";"
+mapCommaSeparator = ","
 )
 
 var (
@@ -194,6 +196,7 @@ type Config struct {
 AccessToken string // An access token (JWT) for authentication (optional)
 ForwardAuthorizationHeader bool // Allow forwarding the `accessToken` named query parameter in the authorization header, overwriting the `AccessToken` option, if set (optional)
 QueryTimeout *time.Duration // Configurable timeout for query (optional)
+Roles map[string]string // Roles (optional)
 }
 
 // FormatDSN returns a DSN string from the configuration.
@@ -214,6 +217,14 @@ func (c *Config) FormatDSN() (string, error) {
 credkv = append(credkv, k+mapKeySeparator+v)
 }
 }
+
+var roles []string
+if c.Roles != nil {
+for k, v := range c.Roles {
+roles = append(roles, fmt.Sprintf("%s=ROLE{%q}", k, v))
+}
+}
+
 source := c.Source
 if source == "" {
 source = "trino-go-client"
@@ -284,6 +295,7 @@ func (c *Config) FormatDSN() (string, error) {
 "extra_credentials": strings.Join(credkv, mapEntrySeparator),
 "custom_client": c.CustomClientName,
 accessTokenConfig: c.AccessToken,
+"roles": strings.Join(roles, mapCommaSeparator),
 } {
 if v != "" {
 query[k] = []string{v}
@@ -307,6 +319,7 @@ type Conn struct {
 useExplicitPrepare bool
 forwardAuthorizationHeader bool
 queryTimeout *time.Duration
+Roles string
 }
 
 var (
@@ -390,6 +403,26 @@ func newConn(dsn string) (*Conn, error) {
 queryTimeout = &d
 }
 
+var formatedRoles string
+if rolesStr := query.Get("roles"); rolesStr != "" {
+if !strings.Contains(rolesStr, "=ROLE{") {
+roles := []string{}
+rolesToFormat := strings.Split(rolesStr, ";")
+
+for _, role := range rolesToFormat {
+parts := strings.Split(role, ":")
+if len(parts) != 2 {
+return nil, fmt.Errorf("Invalid role format: %s", role)
+}
+roles = append(roles, fmt.Sprintf("%s=ROLE{%q}", parts[0], parts[1]))
+}
+
+formatedRoles = strings.Join(roles, mapCommaSeparator)
+} else {
+formatedRoles = rolesStr
+}
+}
+
 c := &Conn{
 baseURL: serverURL.Scheme + "://" + serverURL.Host,
 httpClient: *httpClient,
@@ -400,6 +433,7 @@ func newConn(dsn string) (*Conn, error) {
 useExplicitPrepare: useExplicitPrepare,
 forwardAuthorizationHeader: forwardAuthorizationHeader,
 queryTimeout: queryTimeout,
+Roles: formatedRoles,
 }
 
 var user string
@@ -931,6 +965,10 @@ func (st *driverStmt) exec(ctx context.Context, args []driver.NamedValue) (*stmt
 // Ensure the server returns timestamps preserving their precision, without truncating them to timestamp(3).
 hs.Add("X-Trino-Client-Capabilities", "PARAMETRIC_DATETIME")
 
+if st.conn.Roles != "" {
+hs.Add(trinoRoleHeader, st.conn.Roles)
+}
+
 if len(args) > 0 {
 var ss []string
 for _, arg := range args {

trino/trino_test.go

@@ -203,6 +203,46 @@ func TestKerberosConfig(t *testing.T) {
 assert.Equal(t, want, dsn)
 }
 
+func TestFormatDSNWithRoles(t *testing.T) {
+tests := []struct {
+name string
+config *Config
+wantDSN string
+expectError bool
+}{
+{
+name: "Multiple catalog roles",
+config: &Config{
+ServerURI: "https://foobar@localhost:8090",
+SessionProperties: map[string]string{"query_priority": "1"},
+Roles: map[string]string{"catalog1": "role1", "catalog2": "role2"},
+},
+wantDSN: "https://foobar@localhost:8090?roles=catalog1%3DROLE%7B%22role1%22%7D%2Ccatalog2%3DROLE%7B%22role2%22%7D&session_properties=query_priority%3A1&source=trino-go-client",
+},
+{
+name: "Single catalog role",
+config: &Config{
+ServerURI: "https://foobar@localhost:8090",
+SessionProperties: map[string]string{"query_priority": "1"},
+Roles: map[string]string{"catalog1": "role1"},
+},
+wantDSN: "https://foobar@localhost:8090?roles=catalog1%3DROLE%7B%22role1%22%7D&session_properties=query_priority%3A1&source=trino-go-client",
+},
+}
+
+for _, tt := range tests {
+t.Run(tt.name, func(t *testing.T) {
+dsn, err := tt.config.FormatDSN()
+if tt.expectError {
+require.Error(t, err)
+} else {
+require.NoError(t, err)
+assert.Equal(t, tt.wantDSN, dsn)
+}
+})
+}
+}
+
 func TestInvalidKerberosConfig(t *testing.T) {
 c := &Config{
 ServerURI: "http://foobar@localhost:8090",
@@ -1098,6 +1138,31 @@ func TestQueryCancellation(t *testing.T) {
 assert.EqualError(t, err, ErrQueryCancelled.Error(), "unexpected error")
 }
 
+func TestTrinoRoleHeaderSent(t *testing.T) {
+var receivedHeader string
+
+ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+receivedHeader = r.Header.Get(trinoRoleHeader)
+}))
+t.Cleanup(ts.Close)
+
+c := &Config{
+ServerURI: ts.URL,
+SessionProperties: map[string]string{"query_priority": "1"},
+Roles: map[string]string{"catalog1": "role1", "catalog2": "role2"},
+}
+
+dsn, err := c.FormatDSN()
+require.NoError(t, err)
+db, err := sql.Open("trino", dsn)
+require.NoError(t, err)
+
+_, _ = db.Query("SHOW TABLES")
+require.NoError(t, err)
+
+assert.Equal(t, `catalog1=ROLE{"role1"},catalog2=ROLE{"role2"}`, receivedHeader, "expected X-Trino-Role header to be set")
+}
+
 func TestQueryFailure(t *testing.T) {
 ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 w.WriteHeader(http.StatusInternalServerError)