internal/postgres: add WithLicenses to GetUnit

GetUnit now supports the WithLicenses fieldset, which fetches license
contents for the unit.

For golang/go#39629

Change-Id: Ic51baf1d36e75c7ad05c59d232f400a756fb7a94
Reviewed-on: https://go-review.googlesource.com/c/pkgsite/+/252318
Run-TryBot: Julie Qiu <julie@golang.org>
TryBot-Result: kokoro <noreply+kokoro@google.com>
Reviewed-by: Jonathan Amsterdam <jba@google.com>

6 files changed, 48 insertions, 25 deletions

diff --git a/internal/discovery.go b/internal/discovery.go
index 92cd52f0..822c4eb4 100644
--- a/internal/discovery.go
+++ b/internal/discovery.go
@@ -268,6 +268,7 @@ const (
 	WithReadme FieldSet = 1 << iota
 	WithDocumentation
 	WithImports
+	WithLicenses
 )
 
 // LegacyDirectory represents a directory in a module version, and all of the
diff --git a/internal/postgres/insert_module_test.go b/internal/postgres/insert_module_test.go
index acaeaee4..fd74d67e 100644
--- a/internal/postgres/insert_module_test.go
+++ b/internal/postgres/insert_module_test.go
@@ -109,28 +109,28 @@ func checkModule(ctx context.Context, t *testing.T, want *internal.Module) {
 		}
 	}
 
-	for _, dir := range want.Units {
-		got, err := testDB.GetUnit(ctx, &dir.UnitMeta, internal.AllFields)
+	for _, wantu := range want.Units {
+		got, err := testDB.GetUnit(ctx, &wantu.UnitMeta, internal.AllFields)
 		if err != nil {
 			t.Fatal(err)
 		}
 		// TODO(golang/go#38513): remove once we start displaying
 		// READMEs for directories instead of the top-level module.
-		dir.Readme = &internal.Readme{
+		wantu.Readme = &internal.Readme{
 			Filepath: sample.ReadmeFilePath,
 			Contents: sample.ReadmeContents,
 		}
-		if dir.Package != nil {
-			dir.Name = dir.Package.Name
+		wantu.LicenseContents = sample.Licenses
+		if wantu.Package != nil {
+			wantu.Name = wantu.Package.Name
 		}
-		wantd := dir
 		opts := cmp.Options{
 			cmpopts.IgnoreFields(internal.LegacyModuleInfo{}, "LegacyReadmeFilePath"),
 			cmpopts.IgnoreFields(internal.LegacyModuleInfo{}, "LegacyReadmeContents"),
 			cmpopts.IgnoreFields(licenses.Metadata{}, "Coverage"),
 			cmp.AllowUnexported(source.Info{}, safehtml.HTML{}),
 		}
-		if diff := cmp.Diff(wantd, got, opts); diff != "" {
+		if diff := cmp.Diff(wantu, got, opts); diff != "" {
 			t.Errorf("mismatch (-want +got):\n%s", diff)
 		}
 	}
diff --git a/internal/postgres/unit.go b/internal/postgres/unit.go
index a8070541..59871efd 100644
--- a/internal/postgres/unit.go
+++ b/internal/postgres/unit.go
@@ -121,6 +121,13 @@ func (db *DB) GetUnit(ctx context.Context, um *internal.UnitMeta, fields interna
 			u.Imports = imports
 		}
 	}
+	if fields&internal.WithLicenses != 0 {
+		lics, err := db.getLicenses(ctx, u.Path, u.ModulePath, pathID)
+		if err != nil {
+			return nil, err
+		}
+		u.LicenseContents = lics
+	}
 	if fields == internal.AllFields {
 		if u.Name != "" {
 			if u.Package == nil {
diff --git a/internal/postgres/unit_test.go b/internal/postgres/unit_test.go
index a5154c8e..d2ffe817 100644
--- a/internal/postgres/unit_test.go
+++ b/internal/postgres/unit_test.go
@@ -363,6 +363,9 @@ func TestGetUnitFieldSet(t *testing.T) {
 		if fields&internal.WithImports != 0 {
 			u.Imports = sample.Imports
 		}
+		if fields&internal.WithLicenses == 0 {
+			u.LicenseContents = nil
+		}
 		if u.Package != nil {
 			u.Package.Name = u.Name
 		}
@@ -386,6 +389,12 @@ func TestGetUnitFieldSet(t *testing.T) {
 				nil, nil),
 		},
 		{
+			name:   "WithLicenses",
+			fields: internal.WithLicenses,
+			want: unit("a.com/m/dir/p", "a.com/m", "v1.2.3",
+				nil, nil),
+		},
+		{
 			name:   "WithReadme",
 			fields: internal.WithReadme,
 			want: unit("a.com/m/dir/p", "a.com/m", "v1.2.3",
@@ -431,8 +440,9 @@ func unit(path, modulePath, version string, readme *internal.Readme, pkg *intern
 			IsRedistributable: true,
 			Licenses:          sample.LicenseMetadata,
 		},
-		Readme:  readme,
-		Package: pkg,
+		LicenseContents: sample.Licenses,
+		Readme:          readme,
+		Package:         pkg,
 	}
 	if pkg != nil {
 		u.Name = pkg.Name
diff --git a/internal/testing/sample/sample.go b/internal/testing/sample/sample.go
index bbe70832..1096093d 100644
--- a/internal/testing/sample/sample.go
+++ b/internal/testing/sample/sample.go
@@ -217,8 +217,8 @@ func AddPackage(m *internal.Module, p *internal.LegacyPackage) *internal.Module
 	}
 	for pth := p.Path; len(pth) > minLen; pth = path.Dir(pth) {
 		found := false
-		for _, d := range m.Units {
-			if d.Path == pth {
+		for _, u := range m.Units {
+			if u.Path == pth {
 				found = true
 				break
 			}
@@ -230,13 +230,13 @@ func AddPackage(m *internal.Module, p *internal.LegacyPackage) *internal.Module
 	return m
 }
 
-func AddUnit(m *internal.Module, d *internal.Unit) {
+func AddUnit(m *internal.Module, u *internal.Unit) {
 	for _, e := range m.Units {
-		if e.Path == d.Path {
+		if e.Path == u.Path {
 			panic(fmt.Sprintf("module already has path %q", e.Path))
 		}
 	}
-	m.Units = append(m.Units, d)
+	m.Units = append(m.Units, u)
 }
 
 func AddLicense(m *internal.Module, lic *licenses.License) {
@@ -245,9 +245,9 @@ func AddLicense(m *internal.Module, lic *licenses.License) {
 	if dir == "." {
 		dir = ""
 	}
-	for _, d := range m.Units {
-		if strings.TrimPrefix(d.Path, m.ModulePath+"/") == dir {
-			d.Licenses = append(d.Licenses, lic.Metadata)
+	for _, u := range m.Units {
+		if strings.TrimPrefix(u.Path, m.ModulePath+"/") == dir {
+			u.Licenses = append(u.Licenses, lic.Metadata)
 		}
 	}
 }
@@ -259,20 +259,24 @@ func UnitEmpty(path, modulePath, version string) *internal.Unit {
 }
 
 func UnitForModuleRoot(m *internal.LegacyModuleInfo, licenses []*licenses.Metadata) *internal.Unit {
-	d := &internal.Unit{UnitMeta: *UnitMeta(m.ModulePath, m.ModulePath, m.Version, "", m.IsRedistributable)}
+	u := &internal.Unit{
+		UnitMeta:        *UnitMeta(m.ModulePath, m.ModulePath, m.Version, "", m.IsRedistributable),
+		LicenseContents: Licenses,
+	}
 	if m.LegacyReadmeFilePath != "" {
-		d.Readme = &internal.Readme{
+		u.Readme = &internal.Readme{
 			Filepath: m.LegacyReadmeFilePath,
 			Contents: m.LegacyReadmeContents,
 		}
 	}
-	return d
+	return u
 }
 
 func UnitForPackage(pkg *internal.LegacyPackage, modulePath, version string) *internal.Unit {
 	return &internal.Unit{
-		UnitMeta: *UnitMeta(pkg.Path, modulePath, version, pkg.Name, pkg.IsRedistributable),
-		Imports:  pkg.Imports,
+		UnitMeta:        *UnitMeta(pkg.Path, modulePath, version, pkg.Name, pkg.IsRedistributable),
+		Imports:         pkg.Imports,
+		LicenseContents: Licenses,
 		Package: &internal.Package{
 			Name: pkg.Name,
 			Path: pkg.Path,
diff --git a/internal/unit.go b/internal/unit.go
index b28bdeb0..8a37cc3c 100644
--- a/internal/unit.go
+++ b/internal/unit.go
@@ -45,9 +45,10 @@ func (um *UnitMeta) IsModule() bool {
 // contains other units, licenses and/or READMEs."
 type Unit struct {
 	UnitMeta
-	Readme  *Readme
-	Package *Package
-	Imports []string
+	Readme          *Readme
+	Package         *Package
+	Imports         []string
+	LicenseContents []*licenses.License
 }
 
 // Documentation is the rendered documentation for a given package