internal/postgres: fetch unit page data with one query

Data for the unit page can now be fetched using a single query.

Running this query in staging:

  golang.org/x/net/context: 27.394 ms (26.526 ms on second run)
  cloud.google.com/go/firestore: 77.537 ms (0.810 ms on second run)
  go.uber.org/zap: 1515.823 ms (13.242 ms on second run)
  net/http: 17694.337 ms (232.138 ms on second run)

Change-Id: I337379551cd2b842212107ccd51edd961df5fd1e
Reviewed-on: https://go-review.googlesource.com/c/pkgsite/+/263778
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>

2 files changed, 93 insertions, 5 deletions

diff --git a/internal/postgres/unit.go b/internal/postgres/unit.go
index c16a88f3..3030477e 100644
--- a/internal/postgres/unit.go
+++ b/internal/postgres/unit.go
@@ -25,8 +25,11 @@ import (
 // TODO(golang/go#39629): remove pID.
 func (db *DB) GetUnit(ctx context.Context, um *internal.UnitMeta, fields internal.FieldSet) (_ *internal.Unit, err error) {
 	defer derrors.Wrap(&err, "GetUnit(ctx, %q, %q, %q)", um.Path, um.ModulePath, um.Version)
-	defer middleware.ElapsedStat(ctx, "GetUnit")()
+	if experiment.IsActive(ctx, internal.ExperimentGetUnitWithOneQuery) && fields&internal.WithDocumentation|fields&internal.WithReadme != 0 {
+		return db.getUnitWithAllFields(ctx, um)
+	}
 
+	defer middleware.ElapsedStat(ctx, "GetUnit")()
 	pathID, err := db.getPathID(ctx, um.Path, um.ModulePath, um.Version)
 	if err != nil {
 		return nil, err
@@ -269,3 +272,77 @@ func (db *DB) getPackagesInUnit(ctx context.Context, fullPath, modulePath, resol
 	}
 	return packages, nil
 }
+
+func (db *DB) getUnitWithAllFields(ctx context.Context, um *internal.UnitMeta) (_ *internal.Unit, err error) {
+	defer derrors.Wrap(&err, "getUnitWithAllFields(ctx, %q, %q, %q)", um.Path, um.ModulePath, um.Version)
+	defer middleware.ElapsedStat(ctx, "getUnitWithAllFields")()
+
+	query := `
+        SELECT
+			d.goos,
+			d.goarch,
+			d.synopsis,
+			d.source,
+			r.file_path,
+			r.contents,
+			COALESCE((
+				SELECT COUNT(path_id)
+				FROM package_imports
+				WHERE path_id = p.id
+				GROUP BY path_id
+				), 0) AS num_imports,
+			COALESCE((
+				SELECT COUNT(DISTINCT from_path)
+				FROM imports_unique
+				WHERE to_path = $1
+				AND from_module_path <> $2
+				), 0) AS num_imported_by
+		FROM paths p
+		INNER JOIN modules m
+		ON p.module_id = m.id
+		LEFT JOIN documentation d
+		ON d.path_id = p.id
+		LEFT JOIN readmes r
+		ON r.path_id = p.id
+		WHERE
+			p.path = $1
+			AND m.module_path = $2
+			AND m.version = $3
+			;`
+
+	var (
+		d internal.Documentation
+		r internal.Readme
+		u internal.Unit
+	)
+	err = db.db.QueryRow(ctx, query, um.Path, um.ModulePath, um.Version).Scan(
+		database.NullIsEmpty(&d.GOOS),
+		database.NullIsEmpty(&d.GOARCH),
+		database.NullIsEmpty(&d.Synopsis),
+		&d.Source,
+		database.NullIsEmpty(&r.Filepath),
+		database.NullIsEmpty(&r.Contents),
+		&u.NumImports,
+		&u.NumImportedBy,
+	)
+	switch err {
+	case sql.ErrNoRows:
+		return nil, derrors.NotFound
+	case nil:
+		if d.GOOS != "" {
+			u.Documentation = &d
+		}
+		if r.Filepath != "" {
+			u.Readme = &r
+		}
+	default:
+		return nil, err
+	}
+	pkgs, err := db.getPackagesInUnit(ctx, um.Path, um.ModulePath, um.Version)
+	if err != nil {
+		return nil, err
+	}
+	u.Subdirectories = pkgs
+	u.UnitMeta = *um
+	return &u, nil
+}
diff --git a/internal/postgres/unit_test.go b/internal/postgres/unit_test.go
index bc3174fb..377c08f3 100644
--- a/internal/postgres/unit_test.go
+++ b/internal/postgres/unit_test.go
@@ -161,10 +161,13 @@ func TestGetUnit(t *testing.T) {
 				test.want.Name,
 				test.want.IsRedistributable,
 			)
-			t.Run("unit-page", func(t *testing.T) {
+			t.Run("unit page with one query", func(t *testing.T) {
+				checkUnit(ctx, t, um, test.want, internal.ExperimentUnitPage, internal.ExperimentGetUnitWithOneQuery)
+			})
+			t.Run("unit page", func(t *testing.T) {
 				checkUnit(ctx, t, um, test.want, internal.ExperimentUnitPage)
 			})
-			t.Run("no-experiments", func(t *testing.T) {
+			t.Run("no experiments", func(t *testing.T) {
 				test.want.Readme = &internal.Readme{
 					Filepath: sample.ReadmeFilePath,
 					Contents: sample.ReadmeContents,
@@ -188,8 +191,16 @@ func checkUnit(ctx context.Context, t *testing.T, um *internal.UnitMeta, want *i
 		cmpopts.IgnoreFields(licenses.Metadata{}, "Coverage"),
 	}
 	want.SourceInfo = um.SourceInfo
-	if diff := cmp.Diff(want, got, opts...); diff != "" {
-		t.Errorf("mismatch (-want, +got):\n%s", diff)
+	if experiment.IsActive(ctx, internal.ExperimentGetUnitWithOneQuery) {
+		want.NumImports = len(want.Imports)
+		opts = append(opts,
+			cmpopts.IgnoreFields(internal.Documentation{}, "HTML"),
+			cmpopts.IgnoreFields(internal.Unit{}, "Imports"),
+			cmpopts.IgnoreFields(internal.Unit{}, "LicenseContents"),
+		)
+		if diff := cmp.Diff(want, got, opts...); diff != "" {
+			t.Errorf("mismatch (-want, +got):\n%s", diff)
+		}
 	}
 }