Merge pull request #3798 from dolthub/bh/user-specific-sess-vars

Support per user session variables in the server config

go/cmd/dolt/commands/sql.go

@@ -362,7 +362,7 @@ func getMultiRepoEnv(ctx context.Context, apr *argparser.ArgParseResults, dEnv *
 	multiDir, multiDbMode := apr.GetValue(multiDBDirFlag)
 	if multiDbMode {
 		var err error
-		mrEnv, err = env.LoadMultiEnvFromDir(ctx, env.GetCurrentUserHomeDir, dEnv.Config.WriteableConfig(), dEnv.FS, multiDir, cmd.VersionStr)
+		mrEnv, err = env.LoadMultiEnvFromDir(ctx, env.GetCurrentUserHomeDir, dEnv.Config.WriteableConfig(), dEnv.FS, multiDir, cmd.VersionStr, dEnv.IgnoreLockFile)
 		if err != nil {
 			return nil, errhand.VerboseErrorFromError(err)
 		}

go/cmd/dolt/commands/sqlserver/server.go

@@ -95,7 +95,7 @@ func Serve(
 			}
 
 			// TODO: this should be the global config, probably?
-			mrEnv, err = env.MultiEnvForDirectory(ctx, dEnv.Config.WriteableConfig(), fs, dEnv.Version)
+			mrEnv, err = env.MultiEnvForDirectory(ctx, dEnv.Config.WriteableConfig(), fs, dEnv.Version, dEnv.IgnoreLockFile)
 			if err != nil {
 				return err, nil
 			}
@@ -117,7 +117,7 @@ func Serve(
 		}
 
 		// TODO: this should be the global config, probably?
-		mrEnv, err = env.LoadMultiEnv(ctx, env.GetCurrentUserHomeDir, dEnv.Config.WriteableConfig(), fs, version, dbNamesAndPaths...)
+		mrEnv, err = env.LoadMultiEnv(ctx, env.GetCurrentUserHomeDir, dEnv.Config.WriteableConfig(), fs, version, dEnv.IgnoreLockFile, dbNamesAndPaths...)
 
 		if err != nil {
 			return err, nil
@@ -158,7 +158,7 @@ func Serve(
 	mySQLServer, startError = server.NewServer(
 		serverConf,
 		sqlEngine.GetUnderlyingEngine(),
-		newSessionBuilder(sqlEngine),
+		newSessionBuilder(sqlEngine, serverConfig),
 		listener,
 	)
 
@@ -222,7 +222,13 @@ func portInUse(hostPort string) bool {
 	return false
 }
 
-func newSessionBuilder(se *engine.SqlEngine) server.SessionBuilder {
+func newSessionBuilder(se *engine.SqlEngine, config ServerConfig) server.SessionBuilder {
+	userToSessionVars := make(map[string]map[string]string)
+	userVars := config.UserVars()
+	for _, curr := range userVars {
+		userToSessionVars[curr.Name] = curr.Vars
+	}
+
 	return func(ctx context.Context, conn *mysql.Conn, host string) (sql.Session, error) {
 		mysqlSess, err := server.DefaultSessionBuilder(ctx, conn, host)
 		if err != nil {
@@ -233,7 +239,27 @@ func newSessionBuilder(se *engine.SqlEngine) server.SessionBuilder {
 			return nil, fmt.Errorf("unknown GMS base session type")
 		}
 
-		return se.NewDoltSession(ctx, mysqlBaseSess)
+		dsess, err := se.NewDoltSession(ctx, mysqlBaseSess)
+		if err != nil {
+			return nil, err
+		}
+
+		varsForUser := userToSessionVars[conn.User]
+		if len(varsForUser) > 0 {
+			sqlCtx, err := se.NewContext(ctx)
+			if err != nil {
+				return nil, err
+			}
+
+			for key, val := range varsForUser {
+				err = dsess.InitSessionVariable(sqlCtx, key, val)
+				if err != nil {
+					return nil, err
+				}
+			}
+		}
+
+		return dsess, nil
 	}
 }
 

go/cmd/dolt/commands/sqlserver/serverconfig.go

@@ -125,6 +125,8 @@ type ServerConfig interface {
 	// PrivilegeFilePath returns the path to the file which contains all needed privilege information in the form of a
 	// JSON string.
 	PrivilegeFilePath() string
+	// UserVars is an array containing user specific session variables
+	UserVars() []UserSessionVars
 }
 
 type commandLineServerConfig struct {
@@ -241,6 +243,10 @@ func (cfg *commandLineServerConfig) PrivilegeFilePath() string {
 	return cfg.privilegeFilePath
 }
 
+func (cfg *commandLineServerConfig) UserVars() []UserSessionVars {
+	return nil
+}
+
 // DatabaseNamesAndPaths returns an array of env.EnvNameAndPathObjects corresponding to the databases to be loaded in
 // a multiple db configuration. If nil is returned the server will look for a database in the current directory and
 // give it a name automatically.

go/cmd/dolt/commands/sqlserver/yaml_config.go

@@ -107,6 +107,11 @@ type MetricsYAMLConfig struct {
 	Port   *int              `yaml:"port"`
 }
 
+type UserSessionVars struct {
+	Name string            `yaml:"name"`
+	Vars map[string]string `yaml:"vars"`
+}
+
 // YAMLConfig is a ServerConfig implementation which is read from a yaml file
 type YAMLConfig struct {
 	LogLevelStr       *string               `yaml:"log_level"`
@@ -118,6 +123,7 @@ type YAMLConfig struct {
 	DataDirStr        *string               `yaml:"data_dir"`
 	MetricsConfig     MetricsYAMLConfig     `yaml:"metrics"`
 	PrivilegeFile     *string               `yaml:"privilege_file"`
+	Vars              []UserSessionVars     `yaml:"user_session_vars"`
 }
 
 var _ ServerConfig = YAMLConfig{}
@@ -324,6 +330,14 @@ func (cfg YAMLConfig) PrivilegeFilePath() string {
 	return ""
 }
 
+func (cfg YAMLConfig) UserVars() []UserSessionVars {
+	if cfg.Vars != nil {
+		return cfg.Vars
+	}
+
+	return nil
+}
+
 // QueryParallelism returns the parallelism that should be used by the go-mysql-server analyzer
 func (cfg YAMLConfig) QueryParallelism() int {
 	if cfg.PerformanceConfig.QueryParallelism == nil {

go/cmd/dolt/commands/sqlserver/yaml_config_test.go

@@ -59,6 +59,18 @@ metrics:
         label1: value1
         label2: 2
         label3: true
+
+user_session_vars:
+    - name: user0
+      vars:
+          var1: val0_1
+          var2: val0_2
+          var3: val0_3
+    - name: user1
+      vars:
+          var1: val1_1
+          var2: val1_2
+          var4: val1_4
 `
 
 	expected := serverConfigAsYAMLConfig(DefaultServerConfig())
@@ -82,6 +94,24 @@ metrics:
 		},
 	}
 	expected.DataDirStr = strPtr("some nonsense")
+	expected.Vars = []UserSessionVars{
+		{
+			Name: "user0",
+			Vars: map[string]string{
+				"var1": "val0_1",
+				"var2": "val0_2",
+				"var3": "val0_3",
+			},
+		},
+		{
+			Name: "user1",
+			Vars: map[string]string{
+				"var1": "val1_1",
+				"var2": "val1_2",
+				"var4": "val1_4",
+			},
+		},
+	}
 
 	config, err := NewYamlConfig([]byte(testStr))
 	require.NoError(t, err)

go/cmd/dolt/dolt.go

@@ -123,6 +123,8 @@ const stdInFlag = "--stdin"
 const stdOutFlag = "--stdout"
 const stdErrFlag = "--stderr"
 const stdOutAndErrFlag = "--out-and-err"
+const ignoreLocksFlag = "--ignore-lock-file"
+
 const cpuProf = "cpu"
 const memProf = "mem"
 const blockingProf = "blocking"
@@ -138,6 +140,7 @@ func runMain() int {
 	args := os.Args[1:]
 
 	csMetrics := false
+	ignoreLockFile := false
 	if len(args) > 0 {
 		var doneDebugFlags bool
 		for !doneDebugFlags && len(args) > 0 {
@@ -274,6 +277,10 @@ func runMain() int {
 				csMetrics = true
 				args = args[1:]
 
+			case ignoreLocksFlag:
+				ignoreLockFile = true
+				args = args[1:]
+
 			case featureVersionFlag:
 				if featureVersion, err := strconv.Atoi(args[1]); err == nil {
 					doltdb.DoltFeatureVersion = doltdb.FeatureVersion(featureVersion)
@@ -297,6 +304,7 @@ func runMain() int {
 
 	ctx := context.Background()
 	dEnv := env.Load(ctx, env.GetCurrentUserHomeDir, filesys.LocalFS, doltdb.LocalDirDoltDB, Version)
+	dEnv.IgnoreLockFile = ignoreLockFile
 
 	root, err := env.GetCurrentUserHomeDir()
 	if err != nil {

go/go.sum

@@ -238,6 +238,7 @@ github.com/go-logfmt/logfmt v0.5.0/go.mod h1:wCYkCAKZfumFQihp8CzCvQ3paCTfi41vtzG
 github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A=
 github.com/go-logr/logr v1.2.3 h1:2DntVwHkVopvECVRSlL5PSo9eG+cAkDCuckLubN+rq0=
 github.com/go-logr/logr v1.2.3/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A=
+github.com/go-logr/logr v1.2.3/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A=
 github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
 github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
 github.com/go-ole/go-ole v1.2.1/go.mod h1:7FAglXiTm7HKlQRDeOQ6ZNUHidzCWXuZWq/1dTyBNF8=

go/libraries/doltcore/env/environment.go

@@ -98,6 +98,8 @@ type DoltEnv struct {
 	FS     filesys.Filesys
 	urlStr string
 	hdp    HomeDirProvider
+
+	IgnoreLockFile bool
 }
 
 // Load loads the DoltEnv for the current directory of the cli
@@ -1288,12 +1290,20 @@ func (dEnv *DoltEnv) LockFile() string {
 
 // IsLocked returns true if this database's lockfile exists
 func (dEnv *DoltEnv) IsLocked() bool {
+	if dEnv.IgnoreLockFile {
+		return false
+	}
+
 	ok, _ := dEnv.FS.Exists(dEnv.LockFile())
 	return ok
 }
 
 // Lock writes this database's lockfile or errors if it already exists
 func (dEnv *DoltEnv) Lock() error {
+	if dEnv.IgnoreLockFile {
+		return nil
+	}
+
 	if dEnv.IsLocked() {
 		return ErrActiveServerLock.New(dEnv.LockFile())
 	}
@@ -1302,5 +1312,9 @@ func (dEnv *DoltEnv) Lock() error {
 
 // Unlock deletes this database's lockfile
 func (dEnv *DoltEnv) Unlock() error {
+	if dEnv.IgnoreLockFile {
+		return nil
+	}
+
 	return dEnv.FS.DeleteFile(dEnv.LockFile())
 }

go/libraries/doltcore/env/multi_repo_env.go

@@ -45,9 +45,10 @@ type EnvNameAndPath struct {
 
 // MultiRepoEnv is a type used to store multiple environments which can be retrieved by name
 type MultiRepoEnv struct {
-	envs []NamedEnv
-	fs   filesys.Filesys
-	cfg  config.ReadWriteConfig
+	envs           []NamedEnv
+	fs             filesys.Filesys
+	cfg            config.ReadWriteConfig
+	ignoreLockFile bool
 }
 
 type NamedEnv struct {
@@ -141,6 +142,10 @@ func (mrEnv *MultiRepoEnv) GetWorkingRoots(ctx context.Context) (map[string]*dol
 
 // IsLocked returns true if any env is locked
 func (mrEnv *MultiRepoEnv) IsLocked() (bool, string) {
+	if mrEnv.ignoreLockFile {
+		return false, ""
+	}
+
 	for _, e := range mrEnv.envs {
 		if e.env.IsLocked() {
 			return true, e.env.LockFile()
@@ -152,6 +157,10 @@ func (mrEnv *MultiRepoEnv) IsLocked() (bool, string) {
 // Lock locks all child envs. If an error is returned, all
 // child envs will be returned with their initial lock state.
 func (mrEnv *MultiRepoEnv) Lock() error {
+	if mrEnv.ignoreLockFile {
+		return nil
+	}
+
 	if ok, f := mrEnv.IsLocked(); ok {
 		return ErrActiveServerLock.New(f)
 	}
@@ -169,6 +178,10 @@ func (mrEnv *MultiRepoEnv) Lock() error {
 
 // Unlock unlocks all child envs.
 func (mrEnv *MultiRepoEnv) Unlock() error {
+	if mrEnv.ignoreLockFile {
+		return nil
+	}
+
 	var err, retErr error
 	for _, e := range mrEnv.envs {
 		err = e.env.Unlock()
@@ -221,7 +234,7 @@ func getRepoRootDir(path, pathSeparator string) string {
 func DoltEnvAsMultiEnv(ctx context.Context, dEnv *DoltEnv) (*MultiRepoEnv, error) {
 	if !dEnv.Valid() {
 		cfg := dEnv.Config.WriteableConfig()
-		return MultiEnvForDirectory(ctx, cfg, dEnv.FS, dEnv.Version)
+		return MultiEnvForDirectory(ctx, cfg, dEnv.FS, dEnv.Version, dEnv.IgnoreLockFile)
 	}
 
 	dbName := "dolt"
@@ -295,11 +308,13 @@ func MultiEnvForDirectory(
 	config config.ReadWriteConfig,
 	fs filesys.Filesys,
 	version string,
+	ignoreLockFile bool,
 ) (*MultiRepoEnv, error) {
 	mrEnv := &MultiRepoEnv{
-		envs: make([]NamedEnv, 0),
-		fs:   fs,
-		cfg:  config,
+		envs:           make([]NamedEnv, 0),
+		fs:             fs,
+		cfg:            config,
+		ignoreLockFile: ignoreLockFile,
 	}
 
 	// If there are other directories in the directory, try to load them as additional databases
@@ -334,6 +349,7 @@ func LoadMultiEnv(
 	cfg config.ReadWriteConfig,
 	fs filesys.Filesys,
 	version string,
+	ignoreLockFile bool,
 	envNamesAndPaths ...EnvNameAndPath,
 ) (*MultiRepoEnv, error) {
 	nameToPath := make(map[string]string)
@@ -352,9 +368,10 @@ func LoadMultiEnv(
 	}
 
 	mrEnv := &MultiRepoEnv{
-		envs: make([]NamedEnv, 0),
-		fs:   fs,
-		cfg:  cfg,
+		envs:           make([]NamedEnv, 0),
+		fs:             fs,
+		cfg:            cfg,
+		ignoreLockFile: ignoreLockFile,
 	}
 
 	for name, path := range nameToPath {
@@ -419,6 +436,7 @@ func LoadMultiEnvFromDir(
 	cfg config.ReadWriteConfig,
 	fs filesys.Filesys,
 	path, version string,
+	ignoreLockFile bool,
 ) (*MultiRepoEnv, error) {
 	envNamesAndPaths, err := DBNamesAndPathsFromDir(fs, path)
 
@@ -431,7 +449,7 @@ func LoadMultiEnvFromDir(
 		return nil, errhand.VerboseErrorFromError(err)
 	}
 
-	return LoadMultiEnv(ctx, hdp, cfg, multiDbDirFs, version, envNamesAndPaths...)
+	return LoadMultiEnv(ctx, hdp, cfg, multiDbDirFs, version, ignoreLockFile, envNamesAndPaths...)
 }
 
 func dirToDBName(dirName string) string {

go/libraries/doltcore/env/multi_repo_env_test.go

@@ -183,7 +183,7 @@ func TestLoadMultiEnv(t *testing.T) {
 		envNamesAndPaths[i] = EnvNameAndPath{name, filepath.Join(rootPath, name)}
 	}
 
-	mrEnv, err := LoadMultiEnv(context.Background(), hdp, config.NewEmptyMapConfig(), filesys.LocalFS, "test", envNamesAndPaths...)
+	mrEnv, err := LoadMultiEnv(context.Background(), hdp, config.NewEmptyMapConfig(), filesys.LocalFS, "test", false, envNamesAndPaths...)
 	require.NoError(t, err)
 
 	for _, name := range names {
@@ -205,7 +205,7 @@ func TestLoadMultiEnvFromDir(t *testing.T) {
 	}
 
 	rootPath, hdp, envs := initMultiEnv(t, "TestLoadMultiEnvFromDir", names)
-	mrEnv, err := LoadMultiEnvFromDir(context.Background(), hdp, config.NewEmptyMapConfig(), filesys.LocalFS, rootPath, "test")
+	mrEnv, err := LoadMultiEnvFromDir(context.Background(), hdp, config.NewEmptyMapConfig(), filesys.LocalFS, rootPath, "test", false)
 	require.NoError(t, err)
 
 	assert.Len(t, mrEnv.envs, len(names))

go/libraries/doltcore/sqle/dsess/session.go

@@ -58,6 +58,8 @@ const (
 	ReplicateHeadsKey        = "dolt_replicate_heads"
 	ReplicateAllHeadsKey     = "dolt_replicate_all_heads"
 	AsyncReplicationKey      = "dolt_async_replication"
+	AwsCredsFileKey          = "aws_credentials_file"
+	AwsCredsProfileKey       = "aws_credentials_profile"
 	// Transactions merges will stomp if either if the below keys are set
 	TransactionMergeStompKey    = "dolt_transaction_merge_stomp"
 	TransactionMergeStompEnvKey = "DOLT_TRANSACTION_MERGE_STOMP"

go/libraries/doltcore/sqle/dsess/variables.go

@@ -137,6 +137,22 @@ func defineSystemVariables(name string) {
 				Type:              sql.NewSystemStringType(DefaultBranchKey(name)),
 				Default:           "",
 			},
+			{
+				Name:              AwsCredsFileKey,
+				Scope:             sql.SystemVariableScope_Session,
+				Dynamic:           false,
+				SetVarHintApplies: false,
+				Type:              sql.NewSystemStringType(AwsCredsFileKey),
+				Default:           nil,
+			},
+			{
+				Name:              AwsCredsProfileKey,
+				Scope:             sql.SystemVariableScope_Session,
+				Dynamic:           false,
+				SetVarHintApplies: false,
+				Type:              sql.NewSystemStringType(AwsCredsProfileKey),
+				Default:           nil,
+			},
 		})
 	}
 }

integration-tests/bats/sql-server.bats

@@ -21,6 +21,55 @@ teardown() {
     teardown_common
 }
 
+
+
+@test "sql-server: user session variables from config" {
+  cd repo1
+  echo "
+privilege_file: privs.json
+user_session_vars:
+- name: user0
+  vars:
+    aws_credentials_file: /Users/user0/.aws/config
+    aws_credentials_profile: default
+- name: user1
+  vars:
+    aws_credentials_file: /Users/user1/.aws/config
+    aws_credentials_profile: lddev" > server.yaml
+
+    dolt sql --privilege-file=privs.json -q "CREATE USER dolt@'127.0.0.1'"
+    dolt sql --privilege-file=privs.json -q "CREATE USER user0@'127.0.0.1' IDENTIFIED BY 'pass0'"
+    dolt sql --privilege-file=privs.json -q "CREATE USER user1@'127.0.0.1' IDENTIFIED BY 'pass1'"
+    dolt sql --privilege-file=privs.json -q "CREATE USER user2@'127.0.0.1' IDENTIFIED BY 'pass2'"
+
+    start_sql_server_with_config "" server.yaml
+
+    run dolt sql-client --host=127.0.0.1 --port=$PORT --user=user0  --password=pass0<<SQL
+SELECT @@aws_credentials_file, @@aws_credentials_profile;
+SQL
+    echo $output
+    [[ "$output" =~ /Users/user0/.aws/config.*default ]] || false
+
+    run dolt sql-client --host=127.0.0.1 --port=$PORT --user=user1 --password=pass1<<SQL
+SELECT @@aws_credentials_file, @@aws_credentials_profile;
+SQL
+    echo $output
+    [[ "$output" =~ /Users/user1/.aws/config.*lddev ]] || false
+
+    run dolt sql-client --host=127.0.0.1 --port=$PORT --user=user2 --password=pass2<<SQL
+SELECT @@aws_credentials_file, @@aws_credentials_profile;
+SQL
+    echo $output
+    [[ "$output" =~ NULL.*NULL ]] || false
+
+    run dolt sql-client --host=127.0.0.1 --port=$PORT --user=user2 --password=pass2<<SQL
+SET @@aws_credentials_file="/Users/should_fail";
+SQL
+    echo $output
+    [[ "$output" =~ "Variable 'aws_credentials_file' is a read only variable" ]] || false
+}
+
+
 @test "sql-server: port in use" {
     cd repo1
 
@@ -1345,4 +1394,4 @@ databases:
     let PORT="$$ % (65536-1024) + 1024"
     run dolt sql-server -P $PORT
     [ "$status" -eq 1 ]
-}
+}