Merge pull request #2102 from dolthub/zachmu/default-db

Set the default DB from the `DOLTGRES_DB` user var

Author: zachmu

docker-entrypoint.sh

@@ -272,19 +272,6 @@ docker_process_init_files() {
   done
 }
 
-# create_database_from_env creates a database if the DATABASE environment variable is set.
-# It retrieves the database name from environment the env var DATABASE
-# and attempts to create the database using exec_sql.
-create_database_from_env() {
-  local database
-  database=$(get_env_var "DB")
-
-  if [ -n "$database" ]; then
-    note "Creating database '${database}'"
-    exec_sql "Failed to create database '$database': " "CREATE DATABASE IF NOT EXISTS \"$database\";"
-  fi
-}
-
 # is_port_open checks if a TCP port is open on a given host.
 # Arguments:
 #   $1 - Host (IP or hostname)
@@ -310,12 +297,14 @@ start_server() {
   local start_time
   start_time=$(date +%s)
 
-  # If either of the doltgres or postgres specific user and password env vars are set, export them
+  # If any of the doltgres or postgres specific user, password, and DB env vars are set, export them
   # to the ones that doltgres understands so they can be used during initialization.
   user=$(get_env_var "USER")
   password=$(get_env_var "PASSWORD")
+  db=$(get_env_var "DB")
   export DOLTGRES_USER=$user
   export DOLTGRES_PASSWORD=$password
+  export DOLTGRES_DB=$db
 
   SERVER_PID=-1
 
@@ -419,8 +408,6 @@ _main() {
 
   start_server "$@"
 
-  create_database_from_env
-
   if [[ ! -f $INIT_COMPLETED ]]; then
     if ls /docker-entrypoint-initdb.d/* >/dev/null 2>&1; then
       docker_process_init_files /docker-entrypoint-initdb.d/*

server/auth/init.go

@@ -37,15 +37,15 @@ func Init(dEnv *env.DoltEnv) {
 
 // GetSuperUserAndPassword returns the superuser and password for the server to use, as defined in the environment
 func GetSuperUserAndPassword() (string, string) {
-	password := "password"
-	if envPassword := os.Getenv(doltgresPasswordEnvVar); envPassword != "" {
-		password = envPassword
-	}
-
 	user := "postgres"
 	if envUser := os.Getenv(doltgresUserEnvVar); envUser != "" {
 		user = envUser
 	}
 
+	password := "password"
+	if envPassword := os.Getenv(doltgresPasswordEnvVar); envPassword != "" {
+		password = envPassword
+	}
+
 	return user, password
 }

server/server.go

@@ -18,6 +18,7 @@ import (
 	"context"
 	"fmt"
 	_ "net/http/pprof"
+	"os"
 	"path/filepath"
 
 	"github.com/cockroachdb/errors"
@@ -49,9 +50,10 @@ import (
 const (
 	Version = "0.54.4"
 
-	DefUserName  = "postres"
-	DefUserEmail = "postgres@somewhere.com"
-	DoltgresDir  = "postgres"
+	DefUserName         = "postres"
+	DefUserEmail        = "postgres@somewhere.com"
+	DefUserEmailFmt     = "%s@somewhere.com"
+	DefaultDbNameEnvVar = "DOLTGRES_DB"
 )
 
 func init() {
@@ -121,9 +123,10 @@ func runServer(ctx context.Context, cfg *servercfg.DoltgresConfig, dEnv *env.Dol
 	}
 
 	// We need a username and password for many SQL commands, so set defaults if they don't exist
+	user, _ := auth.GetSuperUserAndPassword()
 	dEnv.Config.SetFailsafes(map[string]string{
-		config.UserNameKey:  DefUserName,
-		config.UserEmailKey: DefUserEmail,
+		config.UserNameKey:  user,
+		config.UserEmailKey: fmt.Sprintf(DefUserEmailFmt, user),
 	})
 
 	// Reload the dolt environment with the correct data dir that was specified in the configuration.
@@ -134,19 +137,16 @@ func runServer(ctx context.Context, cfg *servercfg.DoltgresConfig, dEnv *env.Dol
 	}
 	dEnv = env.Load(ctx, dEnv.GetUserHomeDir, dataDirFs, doltdb.LocalDirDoltDB, dEnv.Version)
 
-	// Automatically initialize a doltgres database if necessary
-	// TODO: probably should only do this if there are no databases in the data dir already
-	createDoltgresDatabase := false
-	if exists, isDirectory := dataDirFs.Exists(DoltgresDir); !exists {
-		createDoltgresDatabase = true
-	} else if !isDirectory {
-		workingDir, _ := dataDirFs.Abs(".")
-		// The else branch means that there's a Doltgres item, so we need to error if it's a file since we
-		// enforce the creation of a Doltgres database/directory, which would create a name conflict with the file
-		return nil, errors.Errorf("Attempted to create the default `postgres` database at `%s`, but a file with "+
-			"the same name was found. Either remove the file, change the directory using the `--data-dir` argument, "+
-			"or change the environment variable `%s` so that it points to a different directory.", workingDir, servercfg.DOLTGRES_DATA_DIR)
+	// Determine whether we need to initialize the default database
+	initializeDefaultDatabase := true
+	mrEnv, err := env.MultiEnvForDirectory(ctx, dEnv.Config.WriteableConfig(), dataDirFs, Version, dEnv)
+	if err != nil {
+		return nil, err
 	}
+	mrEnv.Iter(func(_ string, _ *env.DoltEnv) (stop bool, err error) {
+		initializeDefaultDatabase = false
+		return true, nil
+	})
 
 	controller := svcs.NewController()
 	newCtx, cancelF := context.WithCancel(ctx)
@@ -178,8 +178,8 @@ func runServer(ctx context.Context, cfg *servercfg.DoltgresConfig, dEnv *env.Dol
 		return nil, err
 	}
 
-	if createDoltgresDatabase {
-		err = createDatabase(ssCfg, "postgres")
+	if initializeDefaultDatabase {
+		err = createDefaultDatabase(ssCfg)
 		if err != nil {
 			return nil, err
 		}
@@ -194,10 +194,11 @@ func runServer(ctx context.Context, cfg *servercfg.DoltgresConfig, dEnv *env.Dol
 	return controller, nil
 }
 
-// createDatabase creates the database named on the local server using the configuration values to connect, returning
+// createDefaultDatabase creates the database named on the local server using the configuration values to connect, returning
 // any error
-func createDatabase(cfg doltservercfg.ServerConfig, dbName string) error {
+func createDefaultDatabase(cfg doltservercfg.ServerConfig) error {
 	user, password := auth.GetSuperUserAndPassword()
+	dbName := getDefaultDatabaseName(user)
 
 	dsn := fmt.Sprintf("postgres://%s:%s@localhost:%d", user, password, cfg.Port())
 
@@ -213,6 +214,17 @@ func createDatabase(cfg doltservercfg.ServerConfig, dbName string) error {
 	return err
 }
 
+// getDefaultDatabaseName returns the name of the default database to create on first server start.
+// If the environment variable DOLTGRES_DB is set, that value is used. Otherwise, the username is used.
+// The username is in turn configured with the environment variable DOLTGRES_USER, defaulting to "postgres".
+func getDefaultDatabaseName(userName string) string {
+	defaultDbName := os.Getenv(DefaultDbNameEnvVar)
+	if defaultDbName != "" {
+		return defaultDbName
+	}
+	return userName
+}
+
 // startReplication begins the background thread that replicates from Postgres, if one is configured.
 func startReplication(cfg *servercfg.DoltgresConfig, ssCfg doltservercfg.ServerConfig) (*logrepl.LogicalReplicator, error) {
 	if cfg.PostgresReplicationConfig == nil {

testing/bats/doltgres.bats

@@ -354,25 +354,44 @@ EOF
 
     [ ! -d "auth.db" ]
 
-    start_sql_server postgres log.txt myuser mypass
+    start_sql_server "" log.txt myuser mypass
     cat log.txt
 
-    query_server_for_user_and_pass myuser mypass -c "create table myTable (a int);"
+    # db matches user name since DOLTGRES_DB was not set
+    query_server_for_user_and_pass myuser mypass myuser -c "create table myTable (a int);"
 
-    run query_server_for_user_and_pass myuser mypass -c "insert into mytable values (1), (2)"
+    run query_server_for_user_and_pass myuser mypass myuser -c "insert into mytable values (1), (2)"
     [ "$status" -eq 0 ]
     [[ "$output" =~ "INSERT" ]] || false
 
-    run query_server_for_user_and_pass postgres password -c "insert into mytable values (1), (2)"
+    run query_server_for_user_and_pass postgres password myuser -c "insert into mytable values (1), (2)"
+    [ "$status" -ne 0 ]
+}
+
+@test 'doltgres: default db via env' {
+    [ ! -d "auth.db" ]
+    
+    start_sql_server mydb log.txt myuser
+    cat log.txt
+
+    query_server_for_user_and_pass myuser password mydb -c "create table myTable (a int);"
+
+    run query_server_for_user_and_pass myuser password mydb -c "insert into mytable values (1), (2)"
+    [ "$status" -eq 0 ]
+    [[ "$output" =~ "INSERT" ]] || false
+
+    run query_server_for_user_and_pass postgres password mydb -c "insert into mytable values (1), (2)"
     [ "$status" -ne 0 ]
 }
 
 query_server_for_user_and_pass() {
     user=$1
     pass=$2
+    db=$3
+    shift
     shift
     shift
 
     nativevar PGPASSWORD "$pass" /w
-    psql -U "$user" -h localhost -p $PORT "$@" postgres
+    psql -U "$user" -h localhost -p $PORT "$@" $db
 }

testing/bats/setup/query-server-common.bash

@@ -19,10 +19,12 @@ wait_for_connection() {
   nativevar PGPASSWORD "$PASSWORD" /w
 
   while [ $SECONDS -lt $end_time ]; do
-    run psql -U $USERNAME -h localhost -p $port -c "SELECT 1;" postgres
+    run psql -U $USERNAME -h localhost -p $port -c "SELECT 1;" $DEFAULT_DB
     if [ $status -eq 0 ]; then
       echo "Connected successfully!"
       return 0
+    else
+      echo "$output"
     fi
     sleep 1
   done
@@ -33,17 +35,32 @@ wait_for_connection() {
 
 start_sql_server() {
     DEFAULT_DB="$1"
-    DEFAULT_DB="${DEFAULT_DB:=postgres}"
     logFile=$2
     USERNAME=$3
-    USERNAME="${USERNAME:=postgres}"
     PASSWORD=$4
-    PASSWORD="${PASSWORD:=password}"
 
-    nativevar DEFAULT_DB "$DEFAULT_DB" /w
-    nativevar PGPASSWORD "$PASSWORD" /w
-    nativevar DOLTGRES_PASSWORD "$PASSWORD" /w
-    nativevar DOLTGRES_USER "$USERNAME" /w
+    if [ -n "$DEFAULT_DB" ]; then
+        nativevar DOLTGRES_DB "$DEFAULT_DB" /w
+    fi
+
+    if [ -n "$PASSWORD" ]; then
+        nativevar PGPASSWORD "password" /w
+        nativevar DOLTGRES_PASSWORD "$PASSWORD" /w
+    else
+        nativevar PGPASSWORD "$PASSWORD" /w
+        PASSWORD="password"
+    fi
+
+    if [ -n "$USERNAME" ]; then
+        nativevar DOLTGRES_USER "$USERNAME" /w
+        if [ -z "$DEFAULT_DB" ]; then
+            DEFAULT_DB="$USERNAME"
+        fi
+    else
+        USERNAME="postgres"
+    fi
+
+    DEFAULT_DB="${DEFAULT_DB:=postgres}"
 
     PORT=$( definePORT )
     CONFIG=$( defineCONFIG $PORT )