Add roles support

Author: Flgado

README.md

@@ -312,6 +312,11 @@ c := &Config{
 dsn, err := c.FormatDSN()
 ```
 
+**Query parameter example (overrides DSN roles):**
+```go
+rows, err := db.Query(query, sql.Named("X-Trino-Role", "catalog1:role1;catalog2:role2;catalog3:ALL"))
+```
+
 #### Examples
 
 ```

trino/trino.go

@@ -261,6 +261,14 @@ func ParseDSN(dsn string) (*Config, error) {
 		}
 	}
 
+	if roles := query.Get("roles"); roles != "" {
+		var err error
+		config.Roles, err = parseMapParameter(roles, "role", mapEntrySeparator, mapKeySeparator)
+		if err != nil {
+			return nil, err
+		}
+	}
+
 	if clientTags := query.Get("clientTags"); clientTags != "" {
 		config.ClientTags = strings.Split(clientTags, commaSeparator)
 	}
@@ -469,6 +477,39 @@ var (
 	_ driver.ConnPrepareContext = &Conn{}
 )
 
+// formatRolesFromMap formats roles from a map into the Trino header format
+func formatRolesFromMap(rolesMap map[string]string) string {
+	var formattedRoles []string
+	for catalog, role := range rolesMap {
+		formattedRoles = append(formattedRoles, formatRoleEntry(catalog, role))
+	}
+	sort.Strings(formattedRoles)
+	return strings.Join(formattedRoles, commaSeparator)
+}
+
+// parseAndFormatRoles parses roles from DSN format (catalog1:roleA;catalog2:roleB)
+// and formats them into the Trino header format
+func parseAndFormatRoles(rolesString string) (string, error) {
+	var formattedRoles []string
+	for _, entry := range strings.Split(rolesString, mapEntrySeparator) {
+		parts := strings.SplitN(entry, mapKeySeparator, 2)
+		if len(parts) != 2 {
+			return "", fmt.Errorf("invalid role entry: %q", entry)
+		}
+		formattedRoles = append(formattedRoles, formatRoleEntry(parts[0], parts[1]))
+	}
+	sort.Strings(formattedRoles)
+	return strings.Join(formattedRoles, commaSeparator), nil
+}
+
+// formatRoleEntry formats a single catalog role entry into Trino header format
+func formatRoleEntry(catalog, role string) string {
+	if role == "ALL" || role == "NONE" {
+		return fmt.Sprintf("%s=%s", catalog, role)
+	}
+	return fmt.Sprintf("%s=ROLE{%q}", catalog, role)
+}
+
 func newConn(dsn string) (*Conn, error) {
 	conf, err := ParseDSN(dsn)
 	if err != nil {
@@ -555,26 +596,11 @@ func newConn(dsn string) (*Conn, error) {
 		c.httpHeaders.Add(trinoTagsHeader, strings.Join(tags, commaSeparator))
 	}
 
-	var formatedRoles []string
-	if rolesStr := query.Get("roles"); rolesStr != "" {
-		splitRoles := strings.Split(rolesStr, mapEntrySeparator)
-		if len(splitRoles) == 0 {
-			return nil, fmt.Errorf("trino: roles cannot be empty")
+	if conf.Roles != nil {
+		rolesHeader := formatRolesFromMap(conf.Roles)
+		if rolesHeader != "" {
+			c.httpHeaders.Add(trinoRoleHeader, rolesHeader)
 		}
-		for _, role := range splitRoles {
-			splitRole := strings.Split(role, ":")
-			if len(splitRole) != 2 {
-				return nil, fmt.Errorf("trino: invalid role format: %q", role)
-			}
-			if splitRole[1] == "ALL" || splitRole[1] == "NONE" {
-				formatedRoles = append(formatedRoles, fmt.Sprintf("%s=%s", splitRole[0], splitRole[1]))
-				continue
-			}
-
-			formatedRoles = append(formatedRoles, fmt.Sprintf("%s=ROLE{%q}", splitRole[0], splitRole[1]))
-		}
-
-		c.httpHeaders.Add(trinoRoleHeader, strings.Join(formatedRoles, commaSeparator))
 	}
 
 	for k, v := range map[string]string{
@@ -1221,6 +1247,15 @@ func (st *driverStmt) exec(ctx context.Context, args []driver.NamedValue) (*stmt
 					st.user = headerValue
 				}
 
+				if arg.Name == trinoRoleHeader {
+					formattedRoles, err := parseAndFormatRoles(headerValue)
+					if err != nil {
+						return nil, err
+					}
+					st.conn.httpHeaders.Set(trinoRoleHeader, formattedRoles)
+					headerValue = formattedRoles
+				}
+
 				hs.Add(arg.Name, headerValue)
 			} else {
 				if st.conn.useExplicitPrepare && hs.Get(preparedStatementHeader) == "" {

trino/trino_test.go

@@ -83,6 +83,7 @@ func TestParseDSNToConfig(t *testing.T) {
 				DisableExplicitPrepare:     true,
 				ForwardAuthorizationHeader: true,
 				QueryTimeout:               &[]time.Duration{5 * time.Minute}[0],
+				Roles:                      map[string]string{"catalog1": "role1", "catalog2": "role2"},
 			},
 		},
 		{
@@ -158,7 +159,8 @@ func TestParseDSNToConfigAllFieldsHandled(t *testing.T) {
 		"accessToken=jwt-token-here&" +
 		"explicitPrepare=false&" +
 		"forwardAuthorizationHeader=true&" +
-		"query_timeout=5m30s"
+		"query_timeout=5m30s&" +
+		"roles=catalog1%3Arole1%3Bcatalog2%3Arole2"
 
 	config, err := ParseDSN(complexDSN)
 	require.NoError(t, err)
@@ -207,6 +209,7 @@ func TestParseDSNToConfigAllFieldsHandled(t *testing.T) {
 	assert.Equal(t, true, config.ForwardAuthorizationHeader)
 	assert.NotNil(t, config.QueryTimeout)
 	assert.Equal(t, 5*time.Minute+30*time.Second, *config.QueryTimeout)
+	assert.Equal(t, map[string]string{"catalog1": "role1", "catalog2": "role2"}, config.Roles)
 }
 
 func TestConfigFormatDSNTags(t *testing.T) {
@@ -1344,29 +1347,59 @@ func TestQueryCancellation(t *testing.T) {
 	assert.EqualError(t, err, ErrQueryCancelled.Error(), "unexpected error")
 }
 
-func TestTrinoRoleHeaderSent(t *testing.T) {
-	var receivedHeader string
+func TestTrinoRoleHeader(t *testing.T) {
+	tests := []struct {
+		name           string
+		roles          map[string]string
+		namedArg       string
+		expectedHeader string
+	}{
+		{
+			name:           "Roles from config",
+			roles:          map[string]string{"catalog1": "role1", "catalog2": "role2"},
+			namedArg:       "",
+			expectedHeader: `catalog1=ROLE{"role1"},catalog2=ROLE{"role2"}`,
+		},
+		{
+			name:           "Override roles with named argument",
+			roles:          map[string]string{"catalog1": "role1"},
+			namedArg:       `catalog3:role3;catalog4:role4;catalog5:ALL`,
+			expectedHeader: `catalog3=ROLE{"role3"},catalog4=ROLE{"role4"},catalog5=ALL`,
+		},
+	}
 
-	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		receivedHeader = r.Header.Get(trinoRoleHeader)
-	}))
-	t.Cleanup(ts.Close)
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			var receivedHeader string
+			var serverURL string
+			ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				receivedHeader = r.Header.Get(trinoRoleHeader)
+				w.Header().Set("Content-Type", "application/json")
+				w.WriteHeader(http.StatusOK)
+				_, _ = w.Write([]byte(`{"id":"1","nextUri":"` + serverURL + `/1"}`))
+			}))
+			serverURL = ts.URL
+			t.Cleanup(ts.Close)
 
-	c := &Config{
-		ServerURI:         ts.URL,
-		SessionProperties: map[string]string{"query_priority": "1"},
-		Roles:             map[string]string{"catalog1": "role1", "catalog2": "role2"},
-	}
+			c := &Config{
+				ServerURI: ts.URL,
+				Roles:     tt.roles,
+			}
 
-	dsn, err := c.FormatDSN()
-	require.NoError(t, err)
-	db, err := sql.Open("trino", dsn)
-	require.NoError(t, err)
+			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)
+			if tt.namedArg != "" {
+				_, _ = db.Query("SELECT 1", sql.Named("X-Trino-Role", tt.namedArg))
+			} else {
+				_, _ = db.Query("SELECT 1")
+			}
 
-	assert.Equal(t, `catalog1=ROLE{"role1"},catalog2=ROLE{"role2"}`, receivedHeader, "expected X-Trino-Role header to be set")
+			assert.Equal(t, tt.expectedHeader, receivedHeader, "expected X-Trino-Role header to match")
+		})
+	}
 }
 
 func TestQueryFailure(t *testing.T) {