Let users customize the path to the scouting database

I want to write a script that runs the scouting webserver. For that to
work elegantly, I need to be able to customize the database path.
Otherwise it gets created in the runfiles tree.

Signed-off-by: Philipp Schrader <philipp.schrader@gmail.com>
Change-Id: I6f06f9b82ca2d3476c76d795ad57b06f7b360ba8
diff --git a/scouting/webserver/main.go b/scouting/webserver/main.go
index 668104b..3c699fd 100644
--- a/scouting/webserver/main.go
+++ b/scouting/webserver/main.go
@@ -3,9 +3,11 @@
 import (
 	"flag"
 	"fmt"
+	"io/ioutil"
 	"log"
 	"os"
 	"os/signal"
+	"path/filepath"
 	"syscall"
 
 	"github.com/frc971/971-Robot-Code/scouting/db"
@@ -14,12 +16,33 @@
 	"github.com/frc971/971-Robot-Code/scouting/webserver/static"
 )
 
+func getDefaultDatabasePath() string {
+	// If using `bazel run`, let's create the database in the root of the
+	// workspace.
+	workspaceDir := os.Getenv("BUILD_WORKSPACE_DIRECTORY")
+	if workspaceDir != "" {
+		return filepath.Join(workspaceDir, "scouting.db")
+	}
+	// If we're inside a `bazel test`, then we will create the database in
+	// a temporary directory.
+	testTempDir := os.Getenv("TEST_TMPDIR")
+	if testTempDir != "" {
+		tempDir, err := ioutil.TempDir(testTempDir, "db")
+		if err != nil {
+			log.Fatal("Failed to create temporary directory in TEST_TMPDIR: ", err)
+		}
+		return filepath.Join(tempDir, "scouting.db")
+	}
+	return filepath.Join(".", "scouting.db")
+}
+
 func main() {
 	portPtr := flag.Int("port", 8080, "The port number to bind to.")
 	dirPtr := flag.String("directory", ".", "The directory to serve at /.")
+	dbPathPtr := flag.String("database", getDefaultDatabasePath(), "The path to the database.")
 	flag.Parse()
 
-	database, err := db.NewDatabase()
+	database, err := db.NewDatabase(*dbPathPtr)
 	if err != nil {
 		log.Fatal("Failed to connect to database: ", err)
 	}