set-auth (#475)

Summary:

- Simple `set` semantics for `auth` struct members, accessible through for example `set "$.auth.google.sub"='[email protected]'`
- Support for `sub` field in auth context, denoting subject.
- Motivating use case is email address based impersonation of google users by service account principals.
- Added robot test `Set Statement Update Auth Scenario Working`.
- Added robot test `Busted Auth Throws Error Then Set Statement Update Auth Scenario Working`.

Author: general-kroll-4-life

.vscode/launch.json

@@ -150,6 +150,8 @@
         "select listing.Identifier as key_policy_id, listing.region from aws.cloud_control.resources listing where listing.data__TypeName = 'AWS::KMS::Key' and listing.region in ('us-east-1', 'ap-southeast-1') order by key_policy_id;",
         "create or replace materialized view de_gen_01 as select json_extract_path_text(detail.Properties, 'KeyPolicy', 'Id') as key_policy_id, json_extract_path_text(detail.Properties, 'Tags') as key_tags, json_extract_path_text(detail.Properties, 'KeyUsage') as key_usage, json_extract_path_text(detail.Properties, 'Origin') as key_origin, case when json_extract_path_text(detail.Properties, 'MultiRegion') = 'true' then 1 else 0 end as key_is_multi_region, detail.region from aws.cloud_control.resources listing inner join aws.cloud_control.resource detail on detail.data__Identifier = listing.Identifier and detail.region = listing.region where listing.data__TypeName = 'AWS::KMS::Key' and listing.region IN ('us-east-1', 'ap-southeast-1') and detail.data__TypeName = 'AWS::KMS::Key' order by key_policy_id ASC ; select key_policy_id, key_tags, key_usage, key_origin, key_is_multi_region, region from de_gen_01 order by key_policy_id ASC ; drop materialized view if exists de_gen_01 ;",
         "show methods in azure.dev_center.customization_tasks;",
+        "set global \"a.b\"=1;",
+        "set global \"$.auth.google.sub\"='[email protected]';",
       ],
       "default": "show providers;"
     },

go.mod

@@ -21,7 +21,7 @@ require (
 	github.com/spf13/cobra v1.4.0
 	github.com/spf13/pflag v1.0.5
 	github.com/spf13/viper v1.10.1
-	github.com/stackql/any-sdk v0.0.3-beta09
+	github.com/stackql/any-sdk v0.0.3-beta17
 	github.com/stackql/go-suffix-map v0.0.1-alpha01
 	github.com/stackql/psql-wire v0.1.1-alpha07
 	github.com/stackql/stackql-parser v0.0.14-alpha04

go.sum

@@ -471,8 +471,8 @@ github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
 github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 github.com/spf13/viper v1.10.1 h1:nuJZuYpG7gTj/XqiUwg8bA0cp1+M2mC3J4g5luUYBKk=
 github.com/spf13/viper v1.10.1/go.mod h1:IGlFPqhNAPKRxohIzWpI5QEy4kuI7tcl5WvR+8qy1rU=
-github.com/stackql/any-sdk v0.0.3-beta09 h1:1Ddl3cpkaLC+3XXkEyAkKeunhMbYwRkWE4iPnTr+TfE=
-github.com/stackql/any-sdk v0.0.3-beta09/go.mod h1:CIMFo3fC2ScpqzkzeCkzUQQuzYA1VuqpG0p1EZXN+wY=
+github.com/stackql/any-sdk v0.0.3-beta17 h1:eajJfNOZBvWwaUD9WdsS81PAKHFxyHndmpdOo3P867A=
+github.com/stackql/any-sdk v0.0.3-beta17/go.mod h1:CIMFo3fC2ScpqzkzeCkzUQQuzYA1VuqpG0p1EZXN+wY=
 github.com/stackql/go-suffix-map v0.0.1-alpha01 h1:TDUDS8bySu41Oo9p0eniUeCm43mnRM6zFEd6j6VUaz8=
 github.com/stackql/go-suffix-map v0.0.1-alpha01/go.mod h1:QAi+SKukOyf4dBtWy8UMy+hsXXV+yyEE4vmBkji2V7g=
 github.com/stackql/psql-wire v0.1.1-alpha07 h1:LQWVUlx4Bougk6dztDNG5tmXxpIVeeTSsInTj801xCs=

internal/stackql/dbmsinternal/dbmsinternal.go

@@ -2,6 +2,7 @@ package dbmsinternal
 
 import (
 	"regexp"
+	"strings"
 
 	"github.com/stackql/stackql-parser/go/vt/sqlparser"
 	"github.com/stackql/stackql/internal/stackql/constants"
@@ -94,7 +95,7 @@ func (pgr *standardDBMSInternalRouter) CanRoute(node sqlparser.SQLNode) (constan
 		}
 		return pgr.analyzeSelect(node)
 	case *sqlparser.Set:
-		return pgr.affirmativeExec()
+		return pgr.analyzeSet(node)
 	case *sqlparser.Show:
 		return pgr.analyzeShow(node)
 	case *sqlparser.Begin, *sqlparser.Commit, *sqlparser.Rollback:
@@ -157,6 +158,15 @@ func (pgr *standardDBMSInternalRouter) analyzeShow(node *sqlparser.Show) (consta
 	return pgr.negative()
 }
 
+func (pgr *standardDBMSInternalRouter) analyzeSet(node *sqlparser.Set) (constants.BackendQueryType, bool) {
+	for _, n := range node.Exprs {
+		if strings.HasPrefix(n.Name.GetRawVal(), "$") {
+			return pgr.negative()
+		}
+	}
+	return pgr.affirmativeExec()
+}
+
 func (pgr *standardDBMSInternalRouter) ExprIsRoutable(node sqlparser.SQLNode) bool {
 	switch node := node.(type) {
 	case sqlparser.TableExpr:

internal/stackql/dto/auth_ctx.go

@@ -7,6 +7,16 @@ import (
 	"strings"
 )
 
+type AuthContexts map[string]*AuthCtx
+
+func (as AuthContexts) Clone() AuthContexts {
+	rv := make(AuthContexts)
+	for k, v := range as {
+		rv[k] = v.Clone()
+	}
+	return rv
+}
+
 type AuthCtx struct {
 	Scopes                  []string       `json:"scopes,omitempty" yaml:"scopes,omitempty"`
 	SQLCfg                  *SQLBackendCfg `json:"sqlDataSource" yaml:"sqlDataSource"`
@@ -28,6 +38,7 @@ type AuthCtx struct {
 	EnvVarPassword          string         `json:"password_var" yaml:"password_var"`
 	EncodedBasicCredentials string         `json:"-" yaml:"-"`
 	Successor               *AuthCtx       `json:"successor" yaml:"successor"`
+	Subject                 string         `json:"sub" yaml:"sub"`
 	Active                  bool           `json:"-" yaml:"-"`
 	Location                string         `json:"location" yaml:"location"`
 	Name                    string         `json:"name" yaml:"name"`

internal/stackql/handler/handler.go

@@ -9,6 +9,7 @@ import (
 	"sync"
 
 	"github.com/stackql/any-sdk/anysdk"
+	"github.com/stackql/any-sdk/pkg/jsonpath"
 	"github.com/stackql/any-sdk/pkg/nomenclature"
 	"github.com/stackql/stackql/internal/stackql/acid/tsm"
 	"github.com/stackql/stackql/internal/stackql/acid/txn_context"
@@ -54,7 +55,7 @@ type HandlerContext interface { //nolint:revive // don't mind stuttering this on
 	GetProviders() map[string]provider.IProvider
 	GetControlAttributes() sqlcontrol.ControlAttributes
 	GetCurrentProvider() string
-	GetAuthContexts() map[string]*dto.AuthCtx
+	GetAuthContexts() dto.AuthContexts
 	GetRegistry() anysdk.RegistryAPI
 	GetErrorPresentation() string
 	GetOutfile() io.Writer
@@ -94,6 +95,9 @@ type HandlerContext interface { //nolint:revive // don't mind stuttering this on
 	GetExportNamespace() string
 
 	GetDataFlowCfg() dto.DataFlowCfg
+
+	//
+	SetConfigAtPath(path string, rhs interface{}, scope string) error
 }
 
 type standardHandlerContext struct {
@@ -108,7 +112,7 @@ type standardHandlerContext struct {
 	providers           map[string]provider.IProvider
 	controlAttributes   sqlcontrol.ControlAttributes
 	currentProvider     string
-	authContexts        map[string]*dto.AuthCtx
+	authContexts        dto.AuthContexts
 	sqlDataSources      map[string]sql_datasource.SQLDataSource
 	registry            anysdk.RegistryAPI
 	errorPresentation   string
@@ -190,7 +194,7 @@ func (hc *standardHandlerContext) GetControlAttributes() sqlcontrol.ControlAttri
 }
 func (hc *standardHandlerContext) GetCurrentProvider() string { return hc.currentProvider }
 
-func (hc *standardHandlerContext) GetAuthContexts() map[string]*dto.AuthCtx {
+func (hc *standardHandlerContext) GetAuthContexts() dto.AuthContexts {
 	hc.authMapMutex.Lock()
 	defer hc.authMapMutex.Unlock()
 	return hc.authContexts
@@ -405,6 +409,41 @@ func (hc *standardHandlerContext) updateAuthContextIfNotExists(providerName stri
 	hc.authContexts[providerName] = authCtx
 }
 
+func (hc *standardHandlerContext) SetConfigAtPath(path string, rhs interface{}, scope string) error {
+	return hc.setConfigAtPath(path, rhs, scope)
+}
+
+func (hc *standardHandlerContext) setConfigAtPath(path string, rhs interface{}, scope string) error {
+	searchPath, searchPathErr := composeSystemSearchPath(path)
+	if searchPathErr != nil {
+		return searchPathErr
+	}
+	system := searchPath.GetSystem()
+	remainder := searchPath.GetRemainder()
+	switch system {
+	case dto.AuthCtxKey:
+		return hc.setAuthContextAtPath(remainder, rhs, scope)
+	default:
+		return fmt.Errorf("system '%s' not supported", system)
+	}
+}
+
+func (hc *standardHandlerContext) setAuthContextAtPath(path string, rhs interface{}, scope string) error {
+	hc.authMapMutex.Lock()
+	defer hc.authMapMutex.Unlock()
+	// TODO: review
+	if scope == "local" {
+		hc.authContexts = hc.authContexts.Clone()
+	}
+	searchPath, searchPathErr := composeSystemSearchPath(path)
+	if searchPathErr != nil {
+		return searchPathErr
+	}
+	authCtx := hc.authContexts[searchPath.GetSystem()]
+	rv := jsonpath.Set(authCtx, searchPath.GetRemainder(), rhs)
+	return rv
+}
+
 func (hc *standardHandlerContext) GetNamespaceCollection() tablenamespace.Collection {
 	return hc.namespaceCollection
 }
@@ -521,6 +560,7 @@ func GetHandlerCtx(
 func transformOpenapiStackqlAuthToLocal(authDTO anysdk.AuthDTO) *dto.AuthCtx {
 	rv := &dto.AuthCtx{
 		Scopes:                  authDTO.GetScopes(),
+		Subject:                 authDTO.GetSubject(),
 		Type:                    authDTO.GetType(),
 		ValuePrefix:             authDTO.GetValuePrefix(),
 		KeyID:                   authDTO.GetKeyID(),

internal/stackql/handler/path_searcher.go

@@ -0,0 +1,44 @@
+package handler
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/stackql/any-sdk/pkg/jsonpath"
+)
+
+type simpleSystemPathSearcher struct {
+	system    string
+	remainder string
+}
+
+type systemPathSearcher interface {
+	GetSystem() string
+	GetRemainder() string
+}
+
+func composeSystemSearchPath(path string) (systemPathSearcher, error) {
+	pSplit, splitErr := jsonpath.SplitSearchPath(path)
+	if splitErr != nil {
+		return nil, splitErr
+	}
+	if len(pSplit) < 1 {
+		return nil, fmt.Errorf("path '%s' is insufficient", path)
+	}
+	remainder := ""
+	if len(pSplit) > 1 {
+		remainder = strings.TrimPrefix(path, pSplit[0]+".")
+	}
+	return &simpleSystemPathSearcher{
+		system:    pSplit[0],
+		remainder: remainder,
+	}, nil
+}
+
+func (ss *simpleSystemPathSearcher) GetSystem() string {
+	return ss.system
+}
+
+func (ss *simpleSystemPathSearcher) GetRemainder() string {
+	return ss.remainder
+}

internal/stackql/planbuilder/plan_builder.go

@@ -1,12 +1,14 @@
 package planbuilder
 
 import (
+	"encoding/json"
 	"fmt"
 	"sort"
 	"strconv"
 	"strings"
 
 	"github.com/stackql/any-sdk/anysdk"
+
 	"github.com/stackql/stackql/internal/stackql/acid/txn_context"
 	"github.com/stackql/stackql/internal/stackql/astanalysis/routeanalysis"
 	"github.com/stackql/stackql/internal/stackql/handler"
@@ -123,7 +125,7 @@ func (pgb *standardPlanGraphBuilder) createInstructionFor(pbi planbuilderinput.P
 		_, _, err := pgb.handleSelect(pbi)
 		return err
 	case *sqlparser.Set:
-		return pgb.nop(pbi)
+		return pgb.handleSet(pbi)
 	case *sqlparser.SetTransaction:
 		return pgb.nop(pbi)
 	case *sqlparser.Show:
@@ -159,6 +161,46 @@ func (pgb *standardPlanGraphBuilder) nop(pbi planbuilderinput.PlanBuilderInput)
 	return err
 }
 
+func setLogic(pbi planbuilderinput.PlanBuilderInput, setExpr *sqlparser.SetExpr) error {
+	lhsRaw := setExpr.Name.GetRawVal()
+	lhsTrimmed := strings.TrimPrefix(lhsRaw, "$.")
+	if lhsTrimmed == lhsRaw {
+		return nil
+	}
+	exprStr := strings.Trim(sqlparser.String(setExpr.Expr), "'")
+	exprObj := map[string]interface{}{}
+	deserErr := json.Unmarshal([]byte(exprStr), &exprObj)
+	if deserErr != nil {
+		rawRv := pbi.GetHandlerCtx().SetConfigAtPath(lhsTrimmed, exprStr, setExpr.Scope)
+		return rawRv
+	}
+	rv := pbi.GetHandlerCtx().SetConfigAtPath(lhsTrimmed, exprObj, setExpr.Scope)
+	return rv
+}
+
+func (pgb *standardPlanGraphBuilder) handleSet(pbi planbuilderinput.PlanBuilderInput) error {
+	primitiveGenerator := pgb.rootPrimitiveGenerator
+	err := primitiveGenerator.AnalyzeStatement(pbi)
+	if err != nil {
+		return err
+	}
+	setStmt, ok := pbi.GetSet()
+	if !ok {
+		return fmt.Errorf("could not cast node of type '%T' to required Set", pbi.GetStatement())
+	}
+	if len(setStmt.Exprs) < 1 {
+		return fmt.Errorf("no set expressions found")
+	}
+	pr := primitive.NewLocalPrimitive(
+		//nolint:revive // acceptable for now
+		func(pc primitive.IPrimitiveCtx) internaldto.ExecutorOutput {
+			err = setLogic(pbi, setStmt.Exprs[0])
+			return internaldto.NewExecutorOutput(nil, nil, nil, nil, err)
+		})
+	pgb.planGraphHolder.GetPrimitiveGraph().CreatePrimitiveNode(pr)
+	return nil
+}
+
 func (pgb *standardPlanGraphBuilder) pgInternal(pbi planbuilderinput.PlanBuilderInput) error {
 	primitiveGenerator := pgb.rootPrimitiveGenerator
 	err := primitiveGenerator.AnalyzePGInternal(pbi)

internal/stackql/planbuilderinput/plan_builder_input.go

@@ -45,6 +45,7 @@ type PlanBuilderInput interface {
 	GetUnion() (*sqlparser.Union, bool)
 	GetUpdate() (*sqlparser.Update, bool)
 	GetUse() (*sqlparser.Use, bool)
+	GetSet() (*sqlparser.Set, bool)
 	IsTccSetAheadOfTime() bool
 	SetIsTccSetAheadOfTime(bool)
 	SetPrepStmtOffset(int)
@@ -384,6 +385,11 @@ func (pbi *StandardPlanBuilderInput) GetUse() (*sqlparser.Use, bool) {
 	return rv, ok
 }
 
+func (pbi *StandardPlanBuilderInput) GetSet() (*sqlparser.Set, bool) {
+	rv, ok := pbi.stmt.(*sqlparser.Set)
+	return rv, ok
+}
+
 func (pbi *StandardPlanBuilderInput) GetUpdate() (*sqlparser.Update, bool) {
 	rv, ok := pbi.stmt.(*sqlparser.Update)
 	return rv, ok

internal/stackql/primitivegenerator/statement_analyzer.go

@@ -82,7 +82,7 @@ func (pb *standardPrimitiveGenerator) AnalyzeStatement(
 	case *sqlparser.Select:
 		return pb.analyzeSelect(pbi)
 	case *sqlparser.Set:
-		return iqlerror.GetStatementNotSupportedError("SET")
+		return pb.analyzeSet(pbi)
 	case *sqlparser.SetTransaction:
 		return iqlerror.GetStatementNotSupportedError("SET TRANSACTION")
 	case *sqlparser.Show:
@@ -118,6 +118,15 @@ func (pb *standardPrimitiveGenerator) analyzeUse(
 	return nil
 }
 
+func (pb *standardPrimitiveGenerator) analyzeSet(
+	pbi planbuilderinput.PlanBuilderInput) error {
+	_, ok := pbi.GetSet()
+	if !ok {
+		return fmt.Errorf("could not cast node of type '%T' to required Set", pbi.GetStatement())
+	}
+	return nil
+}
+
 //nolint:govet,funlen // this is a beast
 func (pb *standardPrimitiveGenerator) analyzeUnion(
 	pbi planbuilderinput.PlanBuilderInput) error {