bzl: wrap migrations schemas/describe scripts (#56389)

Co-authored-by: William Bezuidenhout <william.bezuidenhout@sourcegraph.com>

.aspect/bazelrc/ci.sourcegraph.bazelrc

@@ -42,3 +42,6 @@ build --test_env=INCLUDE_ADMIN_ONBOARDING=false
 
 # Used for container_structure_tests
 build --test_env=DOCKER_HOST
+
+# Used by migration rules
+build --action_env=PGUSER=postgres

.github/workflows/sg-setup.yml

@@ -24,7 +24,7 @@ jobs:
       - name: Install Go
         uses: actions/setup-go@v2
         with:
-          go-version: 1.19.8
+          go-version: 1.20.5
 
       - name: Install asdf plugins
         uses: asdf-vm/actions/install@v1

dev/migrations.bzl

@@ -0,0 +1,90 @@
+CMD_PREAMBLE = """set -e
+
+export HOME=$(pwd)
+export SG_FORCE_REPO_ROOT=$(pwd)
+
+if [ -n "$PG_UTILS_PATH" ]; then
+    PATH="$PG_UTILS_PATH:$PATH"
+fi
+
+if [ -z "$PGUSER" ]; then
+    export PGUSER="sourcegraph"
+fi
+
+if [ -z "$CODEINTEL_PGUSER" ]; then
+    export CODEINTEL_PGUSER="$PGUSER"
+fi
+
+if [ -z "$CODEINSIGHTS_PGUSER" ]; then
+    export CODEINSIGHTS_PGUSER="$PGUSER"
+fi
+"""
+
+def _migration_impl(ctx):
+    ctx.actions.run_shell(
+        inputs = ctx.files.srcs,
+        outputs = [ctx.outputs.out],
+        progress_message = "Running squash migration for %s" % ctx.attr.db,
+        use_default_shell_env = True,
+        execution_requirements = {"requires-network": "1"},
+        command = """{cmd_preamble}
+
+        trap "dropdb --if-exists sg-squasher-{db} && echo 'temp db sg-squasher-{db} dropped'" EXIT
+
+        {sg} migration squash-all -skip-teardown -db {db} -f {output_file}
+        """.format(
+            cmd_preamble = CMD_PREAMBLE,
+            sg = ctx.executable._sg.path,
+            db = ctx.attr.db,
+            output_file = ctx.outputs.out.path,
+        ),
+        tools = ctx.attr._sg[DefaultInfo].default_runfiles.files
+    )
+
+migration = rule(
+    implementation = _migration_impl,
+    attrs = {
+        "srcs": attr.label_list(allow_files= True, mandatory= True),
+        "db": attr.string(mandatory = True),
+        "out": attr.output(mandatory = True),
+        "_sg": attr.label(executable = True, default = "//dev/sg:sg", cfg = "exec"),
+    },
+)
+
+def _describe_impl(ctx):
+    ctx.actions.run_shell(
+        inputs = ctx.files.srcs,
+        outputs = [ctx.outputs.out],
+        progress_message = "Running describe migration for %s" % ctx.attr.db,
+        use_default_shell_env = True,
+        execution_requirements = {"requires-network": "1"},
+        command = """{cmd_preamble}
+
+        export PGDATABASE="_describe_{name}"
+        dropdb --if-exists $PGDATABASE
+        createdb "$PGDATABASE"
+        trap "dropdb --if-exists $PGDATABASE" exit
+
+        {sg} migration describe -db {db} --format={format} -force -out {output_file}
+        """.format(
+            cmd_preamble = CMD_PREAMBLE,
+            sg = ctx.executable._sg.path,
+            db = ctx.attr.db,
+            format = ctx.attr.format,
+            output_file = ctx.outputs.out.path,
+            name = ctx.attr.name,
+        ),
+        tools = ctx.attr._sg[DefaultInfo].default_runfiles.files,
+    )
+
+describe = rule(
+    implementation = _describe_impl,
+    attrs = {
+        "srcs": attr.label_list(allow_files= True, mandatory= True),
+        "db": attr.string(mandatory = True),
+        "format": attr.string(mandatory = True),
+        "out": attr.output(mandatory = True),
+        "_sg": attr.label(executable = True, default = "//dev/sg:sg", cfg = "exec"),
+    },
+)
+

dev/sg/dependencies/BUILD.bazel

@@ -21,6 +21,7 @@ go_library(
         "//dev/sg/root",
         "//internal/database/postgresdsn",
         "//lib/errors",
+        "//lib/output",
         "@com_github_gomodule_redigo//redis",
         "@com_github_grafana_regexp//:regexp",
         "@com_github_jackc_pgx_v4//:pgx",

dev/sg/dependencies/helpers.go

@@ -1,6 +1,7 @@
 package dependencies
 
 import (
+	"bufio"
 	"context"
 	"encoding/json"
 	"fmt"
@@ -8,11 +9,13 @@ import (
 	"net/url"
 	"os"
 	"os/user"
+	"path"
 	"path/filepath"
 	"strings"
 	"time"
 
 	"github.com/gomodule/redigo/redis"
+	"github.com/grafana/regexp"
 	"github.com/jackc/pgx/v4"
 
 	"github.com/sourcegraph/run"
@@ -426,3 +429,98 @@ func forceASDFPluginAdd(ctx context.Context, plugin string, source string) error
 	}
 	return errors.Wrap(err, "asdf plugin-add")
 }
+
+// pgUtilsPathRe is the regexp used to check what value user.bazelrc defines for
+// the PG_UTILS_PATH env var.
+var pgUtilsPathRe = regexp.MustCompile(`build --action_env=PG_UTILS_PATH=(.*)$`)
+
+// userBazelRcPath is the path to a git ignored file that contains Bazel flags
+// specific to the current machine that are required in certain cases.
+var userBazelRcPath = ".aspect/bazelrc/user.bazelrc"
+
+// checkPGUtilsPath ensures that a PG_UTILS_PATH is being defined in .aspect/bazelrc/user.bazelrc
+// if it's needed. For example, on Linux hosts, it's usually located in /usr/bin, which is
+// perfectly fine. But on Mac machines, it's either in the homebrew PATH or on a different
+// location if the user installed Posgres through the Postgresql desktop app.
+func checkPGUtilsPath(ctx context.Context, out *std.Output, args CheckArgs) error {
+	// Check for standard PATH location, that is available inside Bazel when
+	// inheriting the shell environment. That is just /usr/bin, not /usr/local/bin.
+	_, err := os.Stat("/usr/bin/createdb")
+	if err == nil {
+		// If we have createdb in /usr/bin/, nothing to do, it will work outside the box.
+		return nil
+	}
+
+	// Check for the presence of git ignored user.bazelrc, that is specific to local
+	// environment. Because createdb is not under /usr/bin, we have to create that file
+	// and define the PG_UTILS_PATH for migration rules.
+	_, err = os.Stat(userBazelRcPath)
+	if err != nil {
+		if os.IsNotExist(err) {
+			return errors.Wrapf(err, "%s doesn't exist", userBazelRcPath)
+		}
+		return errors.Wrapf(err, "unexpected error with %s", userBazelRcPath)
+	}
+
+	// If it exists, we check if the injected PATH actually contains createdb as intended.
+	// If not, we'll raise an error for sg setup to correct.
+	f, err := os.Open(userBazelRcPath)
+	if err != nil {
+		return errors.Wrapf(err, "can't open %s", userBazelRcPath)
+	}
+	defer f.Close()
+
+	err, pgUtilsPath := parsePgUtilsPathInUserBazelrc(f)
+	if err != nil {
+		return errors.Wrapf(err, "can't parse %s", userBazelRcPath)
+	}
+
+	// If the file exists, but doesn't reference PG_UTILS_PATH, that's an error as well.
+	if pgUtilsPath == "" {
+		return errors.Newf("%s doesn't define PG_UTILS_PATH", userBazelRcPath)
+	}
+
+	// Check that this path contains createdb as expected.
+	if err := checkPgUtilsPathIncludesBinaries(pgUtilsPath); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// parsePgUtilsPathInUserBazelrc extracts the defined path to the createdb postgresql
+// utilities that are used in a the Bazel migration rules.
+func parsePgUtilsPathInUserBazelrc(r io.Reader) (error, string) {
+	scanner := bufio.NewScanner(r)
+	for scanner.Scan() {
+		line := scanner.Text()
+		matches := pgUtilsPathRe.FindStringSubmatch(line)
+		if len(matches) > 1 {
+			return nil, matches[1]
+		}
+	}
+	return scanner.Err(), ""
+}
+
+// checkPgUtilsPathIncludesBinaries ensures that the given path contains createdb as expected.
+func checkPgUtilsPathIncludesBinaries(pgUtilsPath string) error {
+	_, err := os.Stat(path.Join(pgUtilsPath, "createdb"))
+	if err != nil {
+		if os.IsNotExist(err) {
+			return errors.Wrap(err, "currently defined PG_UTILS_PATH doesn't include createdb")
+		}
+		return errors.Wrap(err, "currently defined PG_UTILS_PATH is incorrect")
+	}
+	return nil
+}
+
+// guessPgUtilsPath infers from the environment where the createdb binary
+// is located and returns its parent folder, so it can be used to extend
+// PATH for the migrations Bazel rules.
+func guessPgUtilsPath(ctx context.Context) (error, string) {
+	str, err := usershell.Run(ctx, "which", "createdb").String()
+	if err != nil {
+		return err, ""
+	}
+	return nil, filepath.Dir(str)
+}

dev/sg/dependencies/mac.go

@@ -2,12 +2,16 @@ package dependencies
 
 import (
 	"context"
+	"fmt"
+	"os"
 	"time"
 
 	"github.com/sourcegraph/sourcegraph/dev/sg/internal/check"
 	"github.com/sourcegraph/sourcegraph/dev/sg/internal/std"
 	"github.com/sourcegraph/sourcegraph/dev/sg/internal/usershell"
 	"github.com/sourcegraph/sourcegraph/dev/sg/root"
+	"github.com/sourcegraph/sourcegraph/lib/errors"
+	"github.com/sourcegraph/sourcegraph/lib/output"
 )
 
 const (
@@ -208,6 +212,70 @@ If you're not sure: use the recommended commands to install PostgreSQL.`,
 					`createdb --owner=sourcegraph --encoding=UTF8 --template=template0 sourcegraph`,
 				),
 			},
+			{
+				Name:        "Path to pg utilities (createdb, etc ...)",
+				Enabled:     disableInCI(), // will never pass in CI.
+				Check:       checkPGUtilsPath,
+				Description: `Bazel need to know where the createdb, pg_dump binaries are located, we need to ensure they are accessible\nand possibly indicate where they are located if non default.`,
+				Fix: func(ctx context.Context, cio check.IO, args CheckArgs) error {
+					_, err := root.RepositoryRoot()
+					if err != nil {
+						return errors.Wrap(err, "This check requires sg setup to be run inside sourcegraph/sourcegraph the repository.")
+					}
+
+					// Check if we need to create a user.bazelrc or not
+					_, err = os.Stat(userBazelRcPath)
+					if err != nil {
+						if os.IsNotExist(err) {
+							// It doesn't exist, so we create a new one.
+							f, err := os.Create(".aspect/bazelrc/user.bazelrc")
+							if err != nil {
+								return errors.Wrap(err, "cannot create user.bazelrc to inject PG_UTILS_PATH")
+							}
+							defer f.Close()
+
+							// Try guessing the path to the createdb postgres utilities.
+							err, pgUtilsPath := guessPgUtilsPath(ctx)
+							if err != nil {
+								return err
+							}
+							_, err = fmt.Fprintf(f, "build --action_env=PG_UTILS_PATH=%s\n", pgUtilsPath)
+
+							// Inform the user of what happened, so it's not dark magic.
+							cio.Write(fmt.Sprintf("Guessed PATH for pg utils (createdb,...) to be %q\nCreated %s.", pgUtilsPath, userBazelRcPath))
+							return err
+						}
+
+						// File exists, but we got a different error. Can't continue, bubble up the error.
+						return errors.Wrapf(err, "unexpected error with %s", userBazelRcPath)
+					}
+
+					// If we didn't create it, open the existing one.
+					f, err := os.Open(userBazelRcPath)
+					if err != nil {
+						return errors.Wrapf(err, "cannot open existing %s", userBazelRcPath)
+					}
+					defer f.Close()
+
+					// Parse the path it contains.
+					err, pgUtilsPath := parsePgUtilsPathInUserBazelrc(f)
+					if err != nil {
+						return err
+					}
+
+					// Ensure that path is correct, if not tell the user about it.
+					err = checkPgUtilsPathIncludesBinaries(pgUtilsPath)
+					if err != nil {
+						cio.WriteLine(output.Styled(output.StyleWarning, "--- Manual action needed ---"))
+						cio.WriteLine(output.Styled(output.StyleYellow, fmt.Sprintf("➡️  PG_UTILS_PATH=%q defined in %s doesn't include createdb. Please correct the file manually.", pgUtilsPath, userBazelRcPath)))
+						cio.WriteLine(output.Styled(output.StyleWarning, "Please make sure that this file contains:"))
+						cio.WriteLine(output.Styled(output.StyleWarning, "`build --action_env=PG_UTILS_PATH=[PATH TO PARENT FOLDER OF WHERE createdb IS LOCATED`"))
+						cio.WriteLine(output.Styled(output.StyleWarning, "--- Manual action needed ---"))
+						return err
+					}
+					return nil
+				},
+			},
 		},
 	},
 	{

docker-images/postgres-12-alpine/BUILD.bazel

@@ -5,9 +5,7 @@ load("//dev:oci_defs.bzl", "image_repository")
 
 filegroup(
     name = "config",
-    srcs = glob(
-        ["rootfs/*"],
-    ) + ["config/postgresql.conf.sample"],
+    srcs = ["//docker-images/postgres-12-alpine/rootfs:files"] + glob(["config/postgresql.conf.sample"]),
 )
 
 pkg_tar(
@@ -35,15 +33,14 @@ oci_image(
     },
     tars = [":config_tar"],
     user = "postgres",
-    visibility = [
-        "//docker-images/codeintel-db:__pkg__",
-    ],
+    visibility = ["//visibility:public"],
 )
 
 oci_tarball(
     name = "image_tarball",
     image = ":image",
     repo_tags = ["postgres-12:candidate"],
+    visibility = ["//visibility:public"],
 )
 
 container_structure_test(

docker-images/postgres-12-alpine/rootfs/BUILD.bazel

@@ -3,3 +3,11 @@ package(default_visibility = ["//visibility:public"])
 exports_files([
     "reindex.sh",
 ])
+
+filegroup(
+    name = "files",
+    srcs = glob(
+        ["**"],
+        ["*.bazel"],
+    ),
+)

internal/database/BUILD.bazel

@@ -1,5 +1,8 @@
 load("//dev:go_defs.bzl", "go_test")
 load("@io_bazel_rules_go//go:def.bzl", "go_library")
+load("//dev:write_generated_to_source_files.bzl", "write_generated_to_source_files")
+load("//dev:migrations.bzl", "describe", "migration")
+load("@bazel_skylib//rules:build_test.bzl", "build_test")
 
 go_library(
     name = "database",
@@ -324,3 +327,112 @@ go_test(
         "@tools_gotest//assert",
     ],
 )
+
+migration(
+    name = "frontend-migration",
+    srcs = ["//migrations/frontend:sql"],
+    out = "migrations/frontend/squashed.sql",
+    db = "frontend",
+)
+
+migration(
+    name = "codeinsights-migration",
+    srcs = ["//migrations/codeinsights:sql"],
+    out = "migrations/codeinsights/squashed.sql",
+    db = "codeinsights",
+)
+
+migration(
+    name = "codeintel-migration",
+    srcs = ["//migrations/codeintel:sql"],
+    out = "migrations/codeintel/squashed.sql",
+    db = "codeintel",
+)
+
+describe(
+    name = "frontend-describe-md",
+    srcs = [
+        "//migrations/codeinsights:sql",
+        "//migrations/codeintel:sql",
+        "//migrations/frontend:sql",
+    ],
+    out = "schema.md",
+    db = "frontend",
+    format = "psql",
+)
+
+describe(
+    name = "codeinsights-describe-md",
+    srcs = [
+        "//migrations/codeinsights:sql",
+        "//migrations/codeintel:sql",
+        "//migrations/frontend:sql",
+    ],
+    out = "schema.codeinsights.md",
+    db = "codeinsights",
+    format = "psql",
+)
+
+describe(
+    name = "codeintel-describe-md",
+    srcs = [
+        "//migrations/codeinsights:sql",
+        "//migrations/codeintel:sql",
+        "//migrations/frontend:sql",
+    ],
+    out = "schema.codeintel.md",
+    db = "codeintel",
+    format = "psql",
+)
+
+describe(
+    name = "frontend-describe-json",
+    srcs = [
+        "//migrations/codeinsights:sql",
+        "//migrations/codeintel:sql",
+        "//migrations/frontend:sql",
+    ],
+    out = "schema.json",
+    db = "frontend",
+    format = "json",
+)
+
+describe(
+    name = "codeinsights-describe-json",
+    srcs = [
+        "//migrations/codeinsights:sql",
+        "//migrations/codeintel:sql",
+        "//migrations/frontend:sql",
+    ],
+    out = "schema.codeinsights.json",
+    db = "codeinsights",
+    format = "json",
+)
+
+describe(
+    name = "codeintel-describe-json",
+    srcs = [
+        "//migrations/codeinsights:sql",
+        "//migrations/codeintel:sql",
+        "//migrations/frontend:sql",
+    ],
+    out = "schema.codeintel.json",
+    db = "codeintel",
+    format = "json",
+)
+
+build_test(
+    name = "schema_generation_tests",
+    tags = ["requires-network"],
+    targets = [
+        "frontend-migration",
+        "codeintel-migration",
+        "codeinsights-migration",
+        "frontend-describe-md",
+        "codeintel-describe-md",
+        "codeinsights-describe-md",
+        "frontend-describe-json",
+        "codeintel-describe-json",
+        "codeinsights-describe-json",
+    ],
+)

migrations/codeinsights/BUILD.bazel

@@ -0,0 +1,5 @@
+filegroup(
+    name = "sql",
+    srcs = glob(["**/*"]),
+    visibility = ["//visibility:public"],
+)

migrations/codeintel/BUILD.bazel

@@ -0,0 +1,5 @@
+filegroup(
+    name = "sql",
+    srcs = glob(["**/*"]),
+    visibility = ["//visibility:public"],
+)

migrations/frontend/BUILD.bazel

@@ -0,0 +1,5 @@
+filegroup(
+    name = "sql",
+    srcs = glob(["**/*"]),
+    visibility = ["//visibility:public"],
+)