internal/postgres: support README in subdirectories

When the unit-page experiment is on, GetUnit will return the README for
that path, instead of the one that the module root.

For golang/go#38513
For golang/go#39629

Change-Id: I7a7bcb5762c614e5abb828b3d530207d24562339
Reviewed-on: https://go-review.googlesource.com/c/pkgsite/+/258314
Run-TryBot: Julie Qiu <julie@golang.org>
Reviewed-by: Jonathan Amsterdam <jba@google.com>
Trust: Julie Qiu <julie@golang.org>
diff --git a/internal/postgres/unit.go b/internal/postgres/unit.go
index 84a588c..2e95e85 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/stdlib"
 )
 
@@ -30,8 +31,13 @@
 
 	u := &internal.Unit{UnitMeta: *um}
 	if fields&internal.WithReadme != 0 {
-		readme, err := db.getReadme(ctx, u.ModulePath, u.Version)
-		if err != nil && !errors.Is(err, derrors.NotFound) {
+		var readme *internal.Readme
+		if experiment.IsActive(ctx, internal.ExperimentUnitPage) {
+			readme, err = db.getReadme(ctx, pathID)
+		} else {
+			readme, err = db.getModuleReadme(ctx, u.ModulePath, u.Version)
+		}
+          	if err != nil && !errors.Is(err, derrors.NotFound) {
 			return nil, err
 		}
 		u.Readme = readme
@@ -131,11 +137,26 @@
 }
 
 // getReadme returns the README corresponding to the modulePath and version.
-func (db *DB) getReadme(ctx context.Context, modulePath, resolvedVersion string) (_ *internal.Readme, err error) {
-	defer derrors.Wrap(&err, "getReadme(ctx, %q, %q)", modulePath, resolvedVersion)
-	// TODO(golang/go#38513): update to query on PathID and query the readmes
-	// table directly once we start displaying READMEs for directories instead
-	// of the top-level module.
+func (db *DB) getReadme(ctx context.Context, pathID int) (_ *internal.Readme, err error) {
+	defer derrors.Wrap(&err, "getReadme(ctx, %d)", pathID)
+	var readme internal.Readme
+	err = db.db.QueryRow(ctx, `
+		SELECT file_path, contents
+		FROM readmes
+		WHERE path_id=$1;`, pathID).Scan(&readme.Filepath, &readme.Contents)
+	switch err {
+	case sql.ErrNoRows:
+		return nil, derrors.NotFound
+	case nil:
+		return &readme, nil
+	default:
+		return nil, err
+	}
+}
+
+// getModuleReadme returns the README corresponding to the modulePath and version.
+func (db *DB) getModuleReadme(ctx context.Context, modulePath, resolvedVersion string) (_ *internal.Readme, err error) {
+	defer derrors.Wrap(&err, "getModuleReadme(ctx, %q, %q)", modulePath, resolvedVersion)
 	var readme internal.Readme
 	err = db.db.QueryRow(ctx, `
 		SELECT file_path, contents
diff --git a/internal/postgres/unit_test.go b/internal/postgres/unit_test.go
index d0ee8e6..cdc83e5 100644
--- a/internal/postgres/unit_test.go
+++ b/internal/postgres/unit_test.go
@@ -13,6 +13,7 @@
 	"github.com/google/go-cmp/cmp/cmpopts"
 	"github.com/google/safehtml"
 	"golang.org/x/pkgsite/internal"
+	"golang.org/x/pkgsite/internal/experiment"
 	"golang.org/x/pkgsite/internal/licenses"
 	"golang.org/x/pkgsite/internal/source"
 	"golang.org/x/pkgsite/internal/stdlib"
@@ -160,29 +161,38 @@
 				test.want.Name,
 				test.want.IsRedistributable,
 			)
-			got, err := testDB.GetUnit(ctx, um, internal.AllFields)
-			if err != nil {
-				t.Fatal(err)
-			}
-			opts := []cmp.Option{
-				cmp.AllowUnexported(source.Info{}, safehtml.HTML{}),
-				// The packages table only includes partial license information; it omits the Coverage field.
-				cmpopts.IgnoreFields(licenses.Metadata{}, "Coverage"),
-			}
-			// TODO(golang/go#38513): remove once we start displaying
-			// READMEs for directories instead of the top-level module.
-			test.want.Readme = &internal.Readme{
-				Filepath: sample.ReadmeFilePath,
-				Contents: sample.ReadmeContents,
-			}
-			test.want.SourceInfo = um.SourceInfo
-			if diff := cmp.Diff(test.want, got, opts...); diff != "" {
-				t.Errorf("mismatch (-want, +got):\n%s", diff)
-			}
+			t.Run("unit-page", func(t *testing.T) {
+				checkUnit(ctx, t, um, test.want, internal.ExperimentUnitPage)
+			})
+			t.Run("no-experiments", func(t *testing.T) {
+				test.want.Readme = &internal.Readme{
+					Filepath: sample.ReadmeFilePath,
+					Contents: sample.ReadmeContents,
+				}
+				checkUnit(ctx, t, um, test.want)
+			})
 		})
 	}
 }
 
+func checkUnit(ctx context.Context, t *testing.T, um *internal.UnitMeta, want *internal.Unit, experiments ...string) {
+	t.Helper()
+	ctx = experiment.NewContext(ctx, experiments...)
+	got, err := testDB.GetUnit(ctx, um, internal.AllFields)
+	if err != nil {
+		t.Fatal(err)
+	}
+	opts := []cmp.Option{
+		cmp.AllowUnexported(source.Info{}, safehtml.HTML{}),
+		// The packages table only includes partial license information; it omits the Coverage field.
+		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)
+	}
+}
+
 func TestGetUnitFieldSet(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), testTimeout)
 	defer cancel()