feat: handle secondary roles in the session (#83)

* rename SessionConf to SessionState

* save session state

* restrict the secondary roles

* add more comments

Author: flaneur2020

query.go

@@ -27,14 +27,14 @@ type DataField struct {
 }
 
 type QueryResponse struct {
-	ID        string         `json:"id"`
-	SessionID string         `json:"session_id"`
-	Session   *SessionConfig `json:"session"`
-	Schema    []DataField    `json:"schema"`
-	Data      [][]string     `json:"data"`
-	State     string         `json:"state"`
-	Error     *QueryError    `json:"error"`
-	Stats     QueryStats     `json:"stats"`
+	ID        string        `json:"id"`
+	SessionID string        `json:"session_id"`
+	Session   *SessionState `json:"session"`
+	Schema    []DataField   `json:"schema"`
+	Data      [][]string    `json:"data"`
+	State     string        `json:"state"`
+	Error     *QueryError   `json:"error"`
+	Stats     QueryStats    `json:"stats"`
 	// TODO: Affect rows
 	StatsURI string `json:"stats_uri"`
 	FinalURI string `json:"final_uri"`
@@ -62,7 +62,7 @@ type QueryRequest struct {
 	// We use client session instead of server session with session_id
 	// SessionID  string            `json:"session_id,omitempty"`
 
-	Session    *SessionConfig    `json:"session,omitempty"`
+	Session    *SessionState     `json:"session,omitempty"`
 	SQL        string            `json:"sql"`
 	Pagination *PaginationConfig `json:"pagination,omitempty"`
 
@@ -80,9 +80,10 @@ type PaginationConfig struct {
 	MaxRowsPerPage  int64 `json:"max_rows_per_page,omitempty"`
 }
 
-type SessionConfig struct {
-	Database string `json:"database,omitempty"`
-	Role     string `json:"role,omitempty"`
+type SessionState struct {
+	Database       string    `json:"database,omitempty"`
+	Role           string    `json:"role,omitempty"`
+	SecondaryRoles *[]string `json:"secondary_roles,omitempty"`
 
 	// Since we use client session, this should not be used
 	// KeepServerSessionSecs uint64            `json:"keep_server_session_secs,omitempty"`

query_test.go

@@ -0,0 +1,36 @@
+package godatabend
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/test-go/testify/require"
+)
+
+func Test_SessionState(t *testing.T) {
+	ss := &SessionState{
+		Database:       "db1",
+		Role:           "",
+		SecondaryRoles: nil,
+		Settings:       map[string]string{},
+	}
+	buf, err := json.Marshal(ss)
+	require.NoError(t, err)
+	assert.Equal(t, `{"database":"db1"}`, string(buf))
+
+	buf = []byte(`{"database":"db1", "secondary_roles": []}`)
+	err = json.Unmarshal(buf, ss)
+	require.NoError(t, err)
+	assert.Equal(t, []string{}, *ss.SecondaryRoles)
+
+	buf = []byte(`{"database":"db1", "secondary_roles": null}`)
+	err = json.Unmarshal(buf, ss)
+	require.NoError(t, err)
+	assert.Nil(t, ss.SecondaryRoles)
+
+	buf = []byte(`{"database":"db1"}`)
+	err = json.Unmarshal(buf, ss)
+	require.NoError(t, err)
+	assert.Nil(t, ss.SecondaryRoles)
+}

restful.go

@@ -63,17 +63,19 @@ func NewDefaultCopyOptions() map[string]string {
 type APIClient struct {
 	cli *http.Client
 
-	apiEndpoint       string
-	host              string
-	tenant            string
-	warehouse         string
-	database          string
-	user              string
-	password          string
-	role              string
-	accessTokenLoader AccessTokenLoader
-	sessionSettings   map[string]string
+	apiEndpoint     string
+	host            string
+	tenant          string
+	warehouse       string
+	database        string
+	user            string
+	password        string
+	role            string
+	secondaryRoles  *[]string
+	sessionSettings map[string]string
+
 	statsTracker      QueryStatsTracker
+	accessTokenLoader AccessTokenLoader
 
 	WaitTimeSeconds      int64
 	MaxRowsInBuffer      int64
@@ -92,6 +94,19 @@ func NewAPIClientFromConfig(cfg *Config) *APIClient {
 	default:
 		apiScheme = "https"
 	}
+
+	// if role is set in config, we'd prefer to limit it as the only effective role,
+	// so you could limit the privileges by setting a role with limited privileges.
+	// however this can be overridden by executing `SET SECONDARY ROLES ALL` in the
+	// query.
+	// secondaryRoles now have two viable values:
+	// - nil: means enabling ALL the granted roles of the user
+	// - []string{}: means enabling NONE of the granted roles
+	var secondaryRoles *[]string
+	if len(cfg.Role) > 0 {
+		secondaryRoles = &[]string{}
+	}
+
 	return &APIClient{
 		cli: &http.Client{
 			Timeout: cfg.Timeout,
@@ -104,6 +119,7 @@ func NewAPIClientFromConfig(cfg *Config) *APIClient {
 		user:              cfg.User,
 		password:          cfg.Password,
 		role:              cfg.Role,
+		secondaryRoles:    secondaryRoles,
 		accessTokenLoader: initAccessTokenLoader(cfg),
 		sessionSettings:   cfg.Params,
 		statsTracker:      cfg.StatsTracker,
@@ -274,11 +290,12 @@ func (c *APIClient) getPagenationConfig() *PaginationConfig {
 	}
 }
 
-func (c *APIClient) getSessionConfig() *SessionConfig {
-	return &SessionConfig{
-		Database: c.database,
-		Role:     c.role,
-		Settings: c.sessionSettings,
+func (c *APIClient) getSessionState() *SessionState {
+	return &SessionState{
+		Database:       c.database,
+		Role:           c.role,
+		SecondaryRoles: c.secondaryRoles,
+		Settings:       c.sessionSettings,
 	}
 }
 
@@ -290,7 +307,7 @@ func (c *APIClient) DoQuery(ctx context.Context, query string, args []driver.Val
 	request := QueryRequest{
 		SQL:        q,
 		Pagination: c.getPagenationConfig(),
-		Session:    c.getSessionConfig(),
+		Session:    c.getSessionState(),
 	}
 
 	path := "/v1/query"
@@ -302,12 +319,12 @@ func (c *APIClient) DoQuery(ctx context.Context, query string, args []driver.Val
 	if result.Error != nil {
 		return nil, errors.Wrap(result.Error, "query error")
 	}
-	c.applySessionConfig(&result)
+	c.applySessionState(&result)
 	c.trackStats(&result)
 	return &result, nil
 }
 
-func (c *APIClient) applySessionConfig(response *QueryResponse) {
+func (c *APIClient) applySessionState(response *QueryResponse) {
 	if response.Session == nil {
 		return
 	}
@@ -317,6 +334,7 @@ func (c *APIClient) applySessionConfig(response *QueryResponse) {
 	if len(response.Session.Role) > 0 {
 		c.role = response.Session.Role
 	}
+	c.secondaryRoles = response.Session.SecondaryRoles
 	if response.Session.Settings != nil {
 		newSessionSettings := map[string]string{}
 		for k, v := range response.Session.Settings {
@@ -462,7 +480,7 @@ func (c *APIClient) InsertWithStage(ctx context.Context, sql string, stage *Stag
 	request := QueryRequest{
 		SQL:        sql,
 		Pagination: c.getPagenationConfig(),
-		Session:    c.getSessionConfig(),
+		Session:    c.getSessionState(),
 		StageAttachment: &StageAttachmentConfig{
 			Location:          stage.String(),
 			FileFormatOptions: fileFormatOptions,

restful_test.go

@@ -22,7 +22,7 @@ func TestMakeHeadersUserPassword(t *testing.T) {
 	assert.Nil(t, err)
 	assert.Equal(t, headers["Authorization"], []string{"Basic cm9vdDpyb290"})
 	assert.Equal(t, headers["X-Databend-Tenant"], []string{"default"})
-	session := c.getSessionConfig()
+	session := c.getSessionState()
 	assert.Equal(t, session.Role, "role1")
 }