use jdbc sintax on roles query parameter

Author: Flgado

README.md

@@ -253,17 +253,16 @@ The `queryTimeout` parameter sets a timeout for the query. If the query takes lo
 #### `roles`
 
 ```
-Type:           string
-Format:         roles=catalog1=ROLE{role1},catalog2=ROLE{role2}
-Valid values:   A comma-separated list of catalog-to-role assignments, where each assignment maps a catalog to a role.
+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. When a string is used, it applies the role to the `system` catalog by default.
+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",
@@ -272,23 +271,8 @@ c := &Config{
 }
 
 dsn, err := c.FormatDSN()
-// Result: 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
 ```
 
-**Example using a string (applies to system catalog)**
-
-``` go
-c := &Config{
-	ServerURI:         "https://foobar@localhost:8090",
-	SessionProperties: map[string]string{"query_priority": "1"},
-	Roles:             "admin", // equivalent to map[string]string{"system": "admin"}
-}
-
-dsn, err := c.FormatDSN()
-// Result: https://foobar@localhost:8090?roles=system%3DROLE%7B%22admin%22%7D&session_properties=query_priority%3A1&source=trino-go-client
-```
-
-
 #### Examples
 
 ```
@@ -299,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

@@ -1042,8 +1042,13 @@ func TestRoleHeaderSupport(t *testing.T) {
 			expectError: false,
 		},
 		{
-			name:        "Valid single role via DSN",
-			rawDSN:      *integrationServerFlag + "?roles=tpch%3DROLE%7Brole1%7D",
+			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,
 		},
 		{
@@ -1053,25 +1058,13 @@ func TestRoleHeaderSupport(t *testing.T) {
 				Roles:     map[string]string{"not-exist-catalog": "role1"},
 			},
 			expectError: true,
-			errorSubstr: "USER_ERROR: Catalog 'not-exist-catalog' not found",
-		},
-		{
-			name:        "Invalid role format with colon",
-			rawDSN:      *integrationServerFlag + "?roles=not-exist-catalog%3Arole1",
-			expectError: true,
-			errorSubstr: "Invalid X-Trino-Role header",
-		},
-		{
-			name:        "Invalid role format missing ROLE{}",
-			rawDSN:      *integrationServerFlag + "?roles=catolog%3Drole1",
-			expectError: true,
-			errorSubstr: "Invalid X-Trino-Role header",
+			errorSubstr: "USER_ERROR: Catalog",
 		},
 		{
 			name:        "Invalid role format missing ROLE{}",
 			rawDSN:      *integrationServerFlag + "?roles=catolog%3Drole1",
 			expectError: true,
-			errorSubstr: "Invalid X-Trino-Role header",
+			errorSubstr: "Invalid role format: catolog=role1",
 		},
 	}
 

trino/trino.go

@@ -155,8 +155,6 @@ const (
 	mapKeySeparator   = ":"
 	mapEntrySeparator = ";"
 	mapCommaSeparator = ","
-	mapRolesSeparator = "="
-	sistemRole        = "system"
 )
 
 var (
@@ -198,7 +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                      interface{}       // Roles (optional)
+	Roles                      map[string]string // Roles (optional)
 }
 
 // FormatDSN returns a DSN string from the configuration.
@@ -222,14 +220,8 @@ func (c *Config) FormatDSN() (string, error) {
 
 	var roles []string
 	if c.Roles != nil {
-		if v, ok := c.Roles.(string); ok {
-			roles = append(roles, sistemRole+mapRolesSeparator+fmt.Sprintf("ROLE{%q}", v))
-		} else if v, ok := c.Roles.(map[string]string); ok {
-			for k, v := range v {
-				roles = append(roles, k+mapRolesSeparator+fmt.Sprintf("ROLE{%q}", v))
-			}
-		} else {
-			return "", fmt.Errorf("Invalid roles type %T", c.Roles)
+		for k, v := range c.Roles {
+			roles = append(roles, fmt.Sprintf("%s=ROLE{%q}", k, v))
 		}
 	}
 
@@ -411,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,
@@ -421,7 +433,7 @@ func newConn(dsn string) (*Conn, error) {
 		useExplicitPrepare:         useExplicitPrepare,
 		forwardAuthorizationHeader: forwardAuthorizationHeader,
 		queryTimeout:               queryTimeout,
-		Roles:                      query.Get("roles"),
+		Roles:                      formatedRoles,
 	}
 
 	var user string

trino/trino_test.go

@@ -220,13 +220,13 @@ func TestFormatDSNWithRoles(t *testing.T) {
 			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: "Default system role as string",
+			name: "Single catalog role",
 			config: &Config{
 				ServerURI:         "https://foobar@localhost:8090",
 				SessionProperties: map[string]string{"query_priority": "1"},
-				Roles:             "role1",
+				Roles:             map[string]string{"catalog1": "role1"},
 			},
-			wantDSN: "https://foobar@localhost:8090?roles=system%3DROLE%7B%22role1%22%7D&session_properties=query_priority%3A1&source=trino-go-client",
+			wantDSN: "https://foobar@localhost:8090?roles=catalog1%3DROLE%7B%22role1%22%7D&session_properties=query_priority%3A1&source=trino-go-client",
 		},
 	}