devtools/cmd/db: setup db scripts

We currently have two different setups for our databases with local
development, one using the CLI tools and one with Go scripts. Instead, a
script is now added at devtools/cmd/db to consolidate places where logic
is implemented.

Existing scripts are still maintained for ergonomics, but they just
invoke the relevant devtools/cmd/db command.

The devtools/cmd/db script will also be used in a later CL to simplify
the docker setup. We no longer need the seeddb and migrate containers,
since they can be reduced to a Go binary.

Change-Id: I6d9fcdd452ab024bcffd91bf3f6f15598f4e502b
Reviewed-on: https://go-review.googlesource.com/c/pkgsite/+/333937
Trust: Julie Qiu <julie@golang.org>
Run-TryBot: Julie Qiu <julie@golang.org>
TryBot-Result: kokoro <noreply+kokoro@google.com>
Reviewed-by: Jonathan Amsterdam <jba@google.com>

4 files changed, 107 insertions, 41 deletions

diff --git a/devtools/cmd/db/main.go b/devtools/cmd/db/main.go
new file mode 100644
index 00000000..6e6757c5
--- /dev/null
+++ b/devtools/cmd/db/main.go
@@ -0,0 +1,91 @@
+// Copyright 2021 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// The seeddb command is used to populates a database with an initial set of
+// modules.
+package main
+
+import (
+	"context"
+	"flag"
+	"fmt"
+	"os"
+	"strings"
+
+	_ "github.com/jackc/pgx/v4/stdlib" // for pgx driver
+	"golang.org/x/pkgsite/internal/config"
+	"golang.org/x/pkgsite/internal/database"
+	"golang.org/x/pkgsite/internal/log"
+)
+
+func main() {
+	flag.Usage = func() {
+		fmt.Fprintf(flag.CommandLine.Output(), "usage: db [cmd] [dbname]\n")
+		fmt.Fprintf(flag.CommandLine.Output(), "  create [dbname]: creates a new database and run migrations\n")
+		fmt.Fprintf(flag.CommandLine.Output(), "  drop [dbname]: drops database\n")
+		fmt.Fprintf(flag.CommandLine.Output(), "  truncate [dbname]: truncates all tables in database\n")
+		fmt.Fprintf(flag.CommandLine.Output(), "  recreate [dbname]: drop, create and run migrations\n")
+		fmt.Fprintf(flag.CommandLine.Output(), "dbname is set using $GO_DISCOVERY_DATABASE_NAME. ")
+		fmt.Fprintf(flag.CommandLine.Output(), "See doc/postgres.md for details.\n")
+		flag.PrintDefaults()
+	}
+
+	flag.Parse()
+	if flag.NArg() != 2 {
+		flag.Usage()
+		os.Exit(1)
+	}
+
+	ctx := context.Background()
+	cfg, err := config.Init(ctx)
+	if err != nil {
+		log.Fatal(ctx, err)
+	}
+	log.SetLevel(cfg.LogLevel)
+
+	// Wrap the postgres driver with our own wrapper, which adds OpenCensus instrumentation.
+	ddb, err := database.Open("pgx", cfg.DBConnInfo(), "dbadmin")
+	if err != nil {
+		log.Fatalf(ctx, "database.Open for host %s failed with %v", cfg.DBHost, err)
+	}
+	defer ddb.Close()
+
+	if err := run(ctx, ddb, flag.Args()[0], flag.Args()[1]); err != nil {
+		log.Fatal(ctx, err)
+	}
+}
+
+func run(ctx context.Context, db *database.DB, cmd, dbName string) error {
+	switch cmd {
+	case "create":
+		return createDB(dbName)
+	case "migrate":
+		_, err := database.TryToMigrate(dbName)
+		return err
+	case "drop":
+		err := database.DropDB(dbName)
+		if err != nil && strings.HasSuffix(err.Error(), "does not exist") {
+			fmt.Printf("%q does not exist\n", dbName)
+			return nil
+		}
+		return nil
+	case "recreate":
+		if err := database.DropDB(dbName); err != nil {
+			return err
+		}
+		return createDB(dbName)
+	case "truncate":
+		return database.ResetDB(ctx, db)
+	default:
+		return fmt.Errorf("unsupported arg: %q", cmd)
+	}
+}
+
+func createDB(dbName string) error {
+	if err := database.CreateDB(dbName); err != nil {
+		return err
+	}
+	_, err := database.TryToMigrate(dbName)
+	return err
+}
diff --git a/devtools/create_local_db.sh b/devtools/create_local_db.sh
index d2aa653f..c354424f 100755
--- a/devtools/create_local_db.sh
+++ b/devtools/create_local_db.sh
@@ -8,29 +8,5 @@
 
 source devtools/lib.sh || { echo "Are you at repo root?"; exit 1; }
 
-database_user="postgres"
-if [[ $GO_DISCOVERY_DATABASE_USER != "" ]]; then
-  database_user=$GO_DISCOVERY_DATABASE_USER
-fi
-
-database_password=""
-if [[ $GO_DISCOVERY_DATABASE_PASSWORD != "" ]]; then
-  database_password=$GO_DISCOVERY_DATABASE_PASSWORD
-fi
-
-database_host="localhost"
-if [[ $GO_DISCOVERY_DATABASE_HOST != "" ]]; then
-  database_host=$GO_DISCOVERY_DATABASE_HOST
-fi
-
-database_name='discovery-db'
-if [[ $GO_DISCOVERY_DATABASE_NAME != "" ]]; then
-  database_name=$GO_DISCOVERY_DATABASE_NAME
-fi
-
-echo "CREATE DATABASE \"$database_name\" \
-        OWNER = $database_user \
-        TEMPLATE=template0 \
-        LC_COLLATE = 'C' \
-        LC_CTYPE = 'C';" | psql postgresql://$database_user:$database_password@$database_host:5432?sslmode=disable
-
+go run devtools/cmd/db/main.go create "discovery-db"
+./devtools/migrate_db.sh up
diff --git a/devtools/drop_test_dbs.sh b/devtools/drop_test_dbs.sh
index 24447846..6553a1a6 100755
--- a/devtools/drop_test_dbs.sh
+++ b/devtools/drop_test_dbs.sh
@@ -5,17 +5,16 @@
 # license that can be found in the LICENSE file.
 
 # Drop all test databases, when migrations are beyond repair.
-
-run_pg() {
-  PGPASSWORD="${GO_DISCOVERY_DATABASE_TEST_PASSWORD}" \
-    psql -U postgres -h localhost -c "$1"
-}
-
-run_pg "DROP DATABASE IF EXISTS discovery_frontend_test;"
-run_pg "DROP DATABASE IF EXISTS discovery_integration_test;"
-run_pg "DROP DATABASE IF EXISTS discovery_postgres_test;"
-run_pg "DROP DATABASE IF EXISTS discovery_worker_test;"
-run_pg 'DROP DATABASE IF EXISTS "discovery_postgres_test-0";'
-run_pg 'DROP DATABASE IF EXISTS "discovery_postgres_test-1";'
-run_pg 'DROP DATABASE IF EXISTS "discovery_postgres_test-2";'
-run_pg 'DROP DATABASE IF EXISTS "discovery_postgres_test-3";'
+for dbname in \
+    discovery_frontend_test \
+    discovery_frontend_test \
+    discovery_integration_test \
+    discovery_postgres_test \
+    discovery_worker_test \
+    "discovery_postgres_test-0" \
+    "discovery_postgres_test-1" \
+    "discovery_postgres_test-2" \
+    "discovery_postgres_test-3"
+do
+    GO_DISCOVERY_LOG_LEVEL=info go run devtools/cmd/db/main.go drop $dbname
+done
diff --git a/devtools/recreate_db.sh b/devtools/recreate_db.sh
index 25e9dfd3..efcc7d41 100755
--- a/devtools/recreate_db.sh
+++ b/devtools/recreate_db.sh
@@ -11,4 +11,4 @@ source devtools/lib.sh || { echo "Are you at repo root?"; exit 1; }
 docker-compose -f devtools/docker/compose.yaml stop
 docker-compose -f devtools/docker/compose.yaml down --remove-orphans
 docker-compose -f devtools/docker/compose.yaml up -d db
-docker-compose -f devtools/docker/compose.yaml run migrate
+./devtools/create_local_db.sh