internal/postgres: rewrite GetUnitMeta query

The query for GetUnitMeta is refactored behind an experiment flag, in
order to fix QueryRowContext timeout issues from the existing query.

Initial experiments on dev show the following improvements for the
first invocation:

cloud.google.com/go: 881ms -> 22ms
golang.org/x/tools: 3423ms -> 160ms

The second call took <40ms on both queries.

Change-Id: I7f356e29ee3354545a866c75e304b0d69cd43364
Reviewed-on: https://go-review.googlesource.com/c/pkgsite/+/276297
Trust: Julie Qiu <julie@golang.org>
Run-TryBot: Julie Qiu <julie@golang.org>
Reviewed-by: Jonathan Amsterdam <jba@google.com>
diff --git a/internal/experiment.go b/internal/experiment.go
index 74cf726..ee2d62d 100644
--- a/internal/experiment.go
+++ b/internal/experiment.go
@@ -7,6 +7,7 @@
 
 const (
 	ExperimentAutocomplete       = "autocomplete"
+	ExperimentGetUnitMetaQuery   = "get-unit-meta-query"
 	ExperimentGoldmark           = "goldmark"
 	ExperimentReadmeOutline      = "readme-outline"
 	ExperimentUnitSidebarDetails = "unit-sidebar-details"
@@ -16,6 +17,7 @@
 // a description of each experiment.
 var Experiments = map[string]string{
 	ExperimentAutocomplete:       "Enable autocomplete with search.",
+	ExperimentGetUnitMetaQuery:   "Enable the new get unit meta query, which reads from the paths table.",
 	ExperimentGoldmark:           "Enable the usage of rendering markdown using goldmark instead of blackfriday.",
 	ExperimentReadmeOutline:      "Enable the readme outline in the side nav.",
 	ExperimentUnitSidebarDetails: "Enable the details section in the right sidebar.",
diff --git a/internal/postgres/unit.go b/internal/postgres/unit.go
index e5278a5..4c56e79 100644
--- a/internal/postgres/unit.go
+++ b/internal/postgres/unit.go
@@ -15,6 +15,7 @@
 	"golang.org/x/pkgsite/internal"
 	"golang.org/x/pkgsite/internal/database"
 	"golang.org/x/pkgsite/internal/derrors"
+	"golang.org/x/pkgsite/internal/experiment"
 	"golang.org/x/pkgsite/internal/middleware"
 	"golang.org/x/pkgsite/internal/stdlib"
 )
@@ -32,7 +33,15 @@
 	defer derrors.Wrap(&err, "DB.GetUnitMeta(ctx, %q, %q, %q)", fullPath, requestedModulePath, requestedVersion)
 	defer middleware.ElapsedStat(ctx, "GetUnitMeta")()
 
-	q, args, err := legacyGetUnitMetaQuery(fullPath, requestedModulePath, requestedVersion).ToSql()
+	var (
+		q    string
+		args []interface{}
+	)
+	if experiment.IsActive(ctx, internal.ExperimentGetUnitMetaQuery) {
+		q, args, err = getUnitMetaQuery(fullPath, requestedModulePath, requestedVersion).PlaceholderFormat(squirrel.Dollar).ToSql()
+	} else {
+		q, args, err = legacyGetUnitMetaQuery(fullPath, requestedModulePath, requestedVersion).PlaceholderFormat(squirrel.Dollar).ToSql()
+	}
 	if err != nil {
 		return nil, fmt.Errorf("squirrel.ToSql: %v", err)
 	}
@@ -71,6 +80,52 @@
 	}
 }
 
+func getUnitMetaQuery(fullPath, requestedModulePath, requestedVersion string) squirrel.SelectBuilder {
+	query := squirrel.Select(
+		"m.module_path",
+		"m.version",
+		"m.commit_time",
+		"m.source_info",
+		"u.name",
+		"u.redistributable",
+		"u.license_types",
+		"u.license_paths",
+	)
+	if requestedVersion != internal.LatestVersion {
+		query = query.From("modules m").Join("units u on u.module_id = m.id").Where(squirrel.Eq{"u.path": fullPath})
+		if requestedModulePath != internal.UnknownModulePath {
+			query = query.Where(squirrel.Eq{"m.module_path": requestedModulePath})
+		}
+		if internal.DefaultBranches[requestedVersion] {
+			query = query.Join("version_map vm ON m.id = vm.module_id").Where("vm.requested_version = ? ", requestedVersion)
+		} else if requestedVersion != internal.LatestVersion {
+			query = query.Where(squirrel.Eq{"version": requestedVersion})
+		}
+		return orderByLatest(query).Limit(1)
+	}
+
+	// Use a nested select to fetch the latest version of the unit, then JOIN
+	// on units to fetch other relevant information. This allows us to use the
+	// index on units.id and paths.path to get the latest path. We can then
+	// look up only the relevant information from the units table.
+	nestedSelect := orderByLatest(squirrel.Select(
+		"m.id",
+		"m.module_path",
+		"m.version",
+		"m.commit_time",
+		"m.source_info",
+		"u.id AS unit_id",
+	).From("modules m").
+		Join("units u ON u.module_id = m.id").
+		Join("paths p ON p.id = u.path_id").
+		Where(squirrel.Eq{"p.path": fullPath}))
+	if requestedModulePath != internal.UnknownModulePath {
+		nestedSelect = nestedSelect.Where(squirrel.Eq{"m.module_path": requestedModulePath})
+	}
+	nestedSelect = nestedSelect.Limit(1)
+	return query.From("units u").JoinClause(nestedSelect.Prefix("JOIN (").Suffix(") m ON u.id = m.unit_id"))
+}
+
 func legacyGetUnitMetaQuery(fullPath, requestedModulePath, requestedVersion string) squirrel.SelectBuilder {
 	query := squirrel.Select(
 		"m.module_path",