internal: conditionally return documentation in GetUnit

- Large documentations blobs should avoid being returned unless
  explicitly requested. Add a WithDocsSource to enable retrieval of
  unit information without the large documentation source information.
- This does not functionally change existing GetUnit callsites.

Change-Id: Ie983b7eaeb61bb92f2716ae64a3d1e1ea738dbe2
Reviewed-on: https://go-review.googlesource.com/c/pkgsite/+/756480
Reviewed-by: Jonathan Amsterdam <jba@google.com>
kokoro-CI: kokoro <noreply+kokoro@google.com>
Auto-Submit: Ethan Lee <ethanalee@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>

10 files changed, 44 insertions, 17 deletions

diff --git a/internal/api/api.go b/internal/api/api.go
index 97c3b5a1..2b0f3f4f 100644
--- a/internal/api/api.go
+++ b/internal/api/api.go
@@ -74,6 +74,9 @@ func ServePackage(w http.ResponseWriter, r *http.Request, ds internal.DataSource
 	if params.Imports {
 		fs |= internal.WithImports
 	}
+	if params.Doc != "" {
+		fs |= internal.WithDocsSource
+	}
 
 	bc := internal.BuildContext{GOOS: params.GOOS, GOARCH: params.GOARCH}
 	var unit *internal.Unit
diff --git a/internal/fetchdatasource/fetchdatasource_test.go b/internal/fetchdatasource/fetchdatasource_test.go
index 05cb121a..223c38bd 100644
--- a/internal/fetchdatasource/fetchdatasource_test.go
+++ b/internal/fetchdatasource/fetchdatasource_test.go
@@ -227,7 +227,7 @@ func TestBypass(t *testing.T) {
 					if err != nil {
 						t.Fatal(err)
 					}
-					got, err := ds.GetUnit(ctx, um, 0, internal.BuildContext{})
+					got, err := ds.GetUnit(ctx, um, internal.WithDocsSource, internal.BuildContext{})
 					if err != nil {
 						t.Fatal(err)
 					}
@@ -495,7 +495,7 @@ func TestGetUnit(t *testing.T) {
 				Path:       test.path,
 				ModuleInfo: internal.ModuleInfo{ModulePath: test.modulePath},
 			}
-			got, err := ds.GetUnit(ctx, um, 0, internal.BuildContext{})
+			got, err := ds.GetUnit(ctx, um, internal.WithDocsSource, internal.BuildContext{})
 			if !test.wantLoaded {
 				if err == nil {
 					t.Fatal("returned not loaded module")
diff --git a/internal/frontend/main.go b/internal/frontend/main.go
index 55867fd3..a08ede8d 100644
--- a/internal/frontend/main.go
+++ b/internal/frontend/main.go
@@ -114,7 +114,7 @@ func fetchMainDetails(ctx context.Context, ds internal.DataSource, um *internal.
 	requestedVersion string, expandReadme bool, bc internal.BuildContext) (_ *MainDetails, err error) {
 	defer stats.Elapsed(ctx, "fetchMainDetails")()
 
-	unit, err := ds.GetUnit(ctx, um, internal.WithMain, bc)
+	unit, err := ds.GetUnit(ctx, um, internal.WithMain|internal.WithDocsSource, bc)
 	if err != nil {
 		return nil, err
 	}
diff --git a/internal/postgres/unit.go b/internal/postgres/unit.go
index 2ead3294..765d541e 100644
--- a/internal/postgres/unit.go
+++ b/internal/postgres/unit.go
@@ -235,7 +235,7 @@ func (db *DB) GetUnit(ctx context.Context, um *internal.UnitMeta, fields interna
 
 	u := &internal.Unit{UnitMeta: *um}
 	if fields&internal.WithMain != 0 {
-		u, err = db.getUnitWithAllFields(ctx, um, bc)
+		u, err = db.getUnitWithAllFields(ctx, um, fields, bc)
 		if err != nil {
 			return nil, err
 		}
@@ -413,7 +413,7 @@ func getPackagesInUnit(ctx context.Context, db *database.DB, fullPath, modulePat
 	return packages, nil
 }
 
-func (db *DB) getUnitWithAllFields(ctx context.Context, um *internal.UnitMeta, bc internal.BuildContext) (_ *internal.Unit, err error) {
+func (db *DB) getUnitWithAllFields(ctx context.Context, um *internal.UnitMeta, fields internal.FieldSet, bc internal.BuildContext) (_ *internal.Unit, err error) {
 	defer derrors.WrapStack(&err, "getUnitWithAllFields(ctx, %q, %q, %q)", um.Path, um.ModulePath, um.Version)
 	defer stats.Elapsed(ctx, "getUnitWithAllFields")()
 
@@ -470,13 +470,17 @@ func (db *DB) getUnitWithAllFields(ctx context.Context, um *internal.UnitMeta, b
 			break
 		}
 	}
-	// Get README, documentation and import counts.
-	query := `
+	docSelect := "CAST(NULL AS bytea),"
+	if fields&internal.WithDocsSource != 0 {
+		docSelect = "d.source,"
+	}
+
+	query := fmt.Sprintf(`
         SELECT
 			r.file_path,
 			r.contents,
 			d.synopsis,
-			d.source,
+			%s
 			COALESCE((
 				SELECT COUNT(unit_id)
 				FROM imports
@@ -501,7 +505,7 @@ func (db *DB) getUnitWithAllFields(ctx context.Context, um *internal.UnitMeta, b
         ) d
 		ON d.unit_id = u.id
 		WHERE u.id = $2
-	`
+	`, docSelect)
 	var (
 		r internal.Readme
 		u internal.Unit
@@ -546,7 +550,7 @@ func (db *DB) getUnitWithAllFields(ctx context.Context, um *internal.UnitMeta, b
 	u.Licenses = licenseMetas
 	u.IsRedistributable = isRedistributable
 
-	if um.IsPackage() && !um.IsCommand() && doc.Source != nil {
+	if um.IsPackage() && !um.IsCommand() && bcMatched.GOOS != "" {
 		u.SymbolHistory, err = GetSymbolHistoryForBuildContext(ctx, db.db, pathID, um.ModulePath, bcMatched)
 		if err != nil {
 			return nil, err
diff --git a/internal/postgres/unit_test.go b/internal/postgres/unit_test.go
index e8dc3046..82dbd9d9 100644
--- a/internal/postgres/unit_test.go
+++ b/internal/postgres/unit_test.go
@@ -672,7 +672,11 @@ func TestGetUnitFieldSet(t *testing.T) {
 	cleanFields := func(u *internal.Unit, fields internal.FieldSet) {
 		// Add/remove fields based on the FieldSet specified.
 		if fields&internal.WithMain != 0 {
-			u.Documentation = []*internal.Documentation{sample.Doc}
+			if fields&internal.WithDocsSource != 0 {
+				u.Documentation = []*internal.Documentation{sample.Doc}
+			} else {
+				u.Documentation = []*internal.Documentation{{GOOS: "all", GOARCH: "all", Synopsis: sample.Doc.Synopsis}}
+			}
 			u.BuildContexts = []internal.BuildContext{internal.BuildContextAll}
 			u.Readme = readme
 			u.NumImports = len(sample.Imports())
@@ -707,6 +711,11 @@ func TestGetUnitFieldSet(t *testing.T) {
 			want:   unit("a.com/m/dir/p", "a.com/m", "v1.2.3", "", readme, []string{}),
 		},
 		{
+			name:   "WithMain|WithDocsSource",
+			fields: internal.WithMain | internal.WithDocsSource,
+			want:   unit("a.com/m/dir/p", "a.com/m", "v1.2.3", "", readme, []string{}),
+		},
+		{
 			name:   "WithImports",
 			fields: internal.WithImports,
 			want:   unitNoLicenses("a.com/m/dir/p", "a.com/m", "v1.2.3", "", nil, []string{}),
@@ -777,7 +786,7 @@ func TestGetUnitBuildContext(t *testing.T) {
 	} {
 		t.Run(fmt.Sprintf("%s-%s", test.goos, test.goarch), func(t *testing.T) {
 			bc := internal.BuildContext{GOOS: test.goos, GOARCH: test.goarch}
-			u, err := testDB.GetUnit(ctx, um, internal.WithMain, bc)
+			u, err := testDB.GetUnit(ctx, um, internal.WithMain|internal.WithDocsSource, bc)
 			if err != nil {
 				t.Fatal(err)
 			}
diff --git a/internal/testing/fakedatasource/fakedatasource.go b/internal/testing/fakedatasource/fakedatasource.go
index 754a00cb..30c71a18 100644
--- a/internal/testing/fakedatasource/fakedatasource.go
+++ b/internal/testing/fakedatasource/fakedatasource.go
@@ -139,8 +139,18 @@ func (ds *FakeDataSource) GetUnit(ctx context.Context, um *internal.UnitMeta, fi
 	// Since we cache the module and its units, we have to copy this unit before we modify it.
 	// It can be a shallow copy, since we're only modifying the Unit.Documentation field.
 	u2 := *u
-	if d := matchingDoc(u.Documentation, bc); d != nil {
-		u2.Documentation = []*internal.Documentation{d}
+	if fields&internal.WithDocsSource != 0 {
+		if d := matchingDoc(u.Documentation, bc); d != nil {
+			u2.Documentation = []*internal.Documentation{d}
+		} else {
+			u2.Documentation = nil
+		}
+	} else if fields&internal.WithMain != 0 {
+		if d := matchingDoc(u.Documentation, bc); d != nil {
+			u2.Documentation = []*internal.Documentation{{GOOS: d.GOOS, GOARCH: d.GOARCH, Synopsis: d.Synopsis}}
+		} else {
+			u2.Documentation = nil
+		}
 	} else {
 		u2.Documentation = nil
 	}
diff --git a/internal/unit.go b/internal/unit.go
index 0a6b3ccc..97df181b 100644
--- a/internal/unit.go
+++ b/internal/unit.go
@@ -113,4 +113,5 @@ const (
 	WithMain FieldSet = 1 << iota
 	WithImports
 	WithLicenses
+	WithDocsSource
 )
diff --git a/internal/worker/fetch_test.go b/internal/worker/fetch_test.go
index f26c7acc..b99ed7ea 100644
--- a/internal/worker/fetch_test.go
+++ b/internal/worker/fetch_test.go
@@ -213,7 +213,7 @@ func TestFetchAndUpdateState(t *testing.T) {
 				t.Fatalf("testDB.GetUnitMeta(ctx, %q, %q) mismatch (-want +got):\n%s", test.modulePath, test.version, diff)
 			}
 
-			gotPkg, err := testDB.GetUnit(ctx, got, internal.WithMain, internal.BuildContext{})
+			gotPkg, err := testDB.GetUnit(ctx, got, internal.WithMain|internal.WithDocsSource, internal.BuildContext{})
 			if err != nil {
 				t.Fatal(err)
 			}
diff --git a/internal/worker/fetcherror_test.go b/internal/worker/fetcherror_test.go
index e33cef48..86a7cc40 100644
--- a/internal/worker/fetcherror_test.go
+++ b/internal/worker/fetcherror_test.go
@@ -405,7 +405,7 @@ func checkPackage(ctx context.Context, t *testing.T, pkgPath string) {
 		t.Fatalf("testDB.GetUnitMeta(%q, %q, %q): isPackage = false; want = true",
 			pkgPath, internal.UnknownModulePath, sample.VersionString)
 	}
-	dir, err := testDB.GetUnit(ctx, um, internal.WithMain, internal.BuildContext{})
+	dir, err := testDB.GetUnit(ctx, um, internal.WithMain|internal.WithDocsSource, internal.BuildContext{})
 	if err != nil {
 		t.Fatal(err)
 	}
diff --git a/internal/worker/refetch_test.go b/internal/worker/refetch_test.go
index 35f66a4e..ef4bdc11 100644
--- a/internal/worker/refetch_test.go
+++ b/internal/worker/refetch_test.go
@@ -131,7 +131,7 @@ func TestReFetch(t *testing.T) {
 		t.Fatalf("testDB.GetUnitMeta(ctx, %q, %q) mismatch (-want +got):\n%s", want.ModulePath, want.Version, diff)
 	}
 
-	gotPkg, err := testDB.GetUnit(ctx, got, internal.WithMain, internal.BuildContext{})
+	gotPkg, err := testDB.GetUnit(ctx, got, internal.WithMain|internal.WithDocsSource, internal.BuildContext{})
 	if err != nil {
 		t.Fatal(err)
 	}