Enable fetching package metadata from pulumi registry api

Author: sean1588

tools/resourcedocsgen/cmd/docs/registry.go

@@ -15,12 +15,12 @@
 package docs
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
 	"net/http"
 	"net/url"
-	"os"
 	"path/filepath"
 	"runtime"
 	"strings"
@@ -36,6 +36,7 @@ import (
 	pschema "github.com/pulumi/pulumi/pkg/v3/codegen/schema"
 	"github.com/pulumi/pulumi/sdk/v3/go/common/util/contract"
 	"github.com/pulumi/registry/tools/resourcedocsgen/pkg"
+	"github.com/pulumi/registry/tools/resourcedocsgen/pkg/registry/svc"
 	concpool "github.com/sourcegraph/conc/pool"
 )
 
@@ -132,89 +133,69 @@ func getSchemaFileURL(metadata pkg.PackageMeta) (string, error) {
 	return fmt.Sprintf("https://raw.githubusercontent.com/%s/%s/%s", repoSlug, metadata.Version, schemaFilePath), nil
 }
 
-func getRegistryPackagesPath(repoPath string) string {
-	return filepath.Join(repoPath, "themes", "default", "data", "registry", "packages")
-}
-
-func genResourceDocsForAllRegistryPackages(registryRepoPath, baseDocsOutDir, basePackageTreeJSONOutDir string) error {
-	registryPackagesPath := getRegistryPackagesPath(registryRepoPath)
-	metadataFiles, err := os.ReadDir(registryPackagesPath)
+func genResourceDocsForAllRegistryPackages(
+	ctx context.Context,
+	provider svc.PackageMetadataProvider,
+	baseDocsOutDir, basePackageTreeJSONOutDir string,
+) error {
+	metadataList, err := provider.ListPackageMetadata(ctx)
 	if err != nil {
-		return errors.Wrap(err, "reading the registry packages dir")
+		return errors.Wrap(err, "listing package metadata")
 	}
 
 	pool := concpool.New().WithErrors().WithMaxGoroutines(runtime.NumCPU())
-	for _, f := range metadataFiles {
-		f := f
+	for _, metadata := range metadataList {
 		pool.Go(func() error {
-			glog.Infof("=== starting %s ===\n", f.Name())
-			glog.Infoln("Processing metadata file")
-			metadataFilePath := filepath.Join(registryPackagesPath, f.Name())
-
-			b, err := os.ReadFile(metadataFilePath)
-			if err != nil {
-				return errors.Wrapf(err, "reading the metadata file %s", metadataFilePath)
-			}
-
-			var metadata pkg.PackageMeta
-			if err := yaml.Unmarshal(b, &metadata); err != nil {
-				return errors.Wrapf(err, "unmarshalling the metadata file %s", metadataFilePath)
-			}
-
+			glog.Infof("=== starting %s ===\n", metadata.Name)
 			docsOutDir := filepath.Join(baseDocsOutDir, metadata.Name, "api-docs")
 			err = genResourceDocsForPackageFromRegistryMetadata(metadata, docsOutDir, basePackageTreeJSONOutDir)
 			if err != nil {
-				return errors.Wrapf(err, "generating resource docs using metadata file info %s", f.Name())
+				return errors.Wrapf(err, "generating resource docs using metadata file info %s", metadata.Name)
 			}
 
-			glog.Infof("=== completed %s ===", f.Name())
+			glog.Infof("=== completed %s ===", metadata.Name)
 			return nil
 		})
 	}
-
 	return pool.Wait()
 }
 
 func resourceDocsFromRegistryCmd() *cobra.Command {
 	var baseDocsOutDir string
 	var basePackageTreeJSONOutDir string
 	var registryDir string
+	var useAPI bool
+	var apiURL string
 
 	cmd := &cobra.Command{
 		Use:   "registry [pkgName]",
 		Short: "Generate resource docs for a package from the registry",
 		Long: "Generate resource docs for all packages in the registry or specific packages. " +
 			"Pass a package name in the registry as an optional arg to generate docs only for that package.",
 		RunE: func(cmd *cobra.Command, args []string) error {
-			registryDir, err := filepath.Abs(registryDir)
-			if err != nil {
-				return errors.Wrap(err, "finding the cwd")
+			ctx := cmd.Context()
+			var provider svc.PackageMetadataProvider
+			if useAPI {
+				provider = svc.NewAPIProvider(apiURL)
+			} else {
+				provider = svc.NewFileSystemProvider(registryDir)
 			}
 
 			if len(args) > 0 {
 				glog.Infoln("Generating docs for a single package:", args[0])
-				registryPackagesPath := getRegistryPackagesPath(registryDir)
-				pkgName := args[0]
-				metadataFilePath := filepath.Join(registryPackagesPath, pkgName+".yaml")
-				b, err := os.ReadFile(metadataFilePath)
+				metadata, err := provider.GetPackageMetadata(ctx, args[0])
 				if err != nil {
-					return errors.Wrapf(err, "reading the metadata file %s", metadataFilePath)
-				}
-
-				var metadata pkg.PackageMeta
-				if err := yaml.Unmarshal(b, &metadata); err != nil {
-					return errors.Wrapf(err, "unmarshalling the metadata file %s", metadataFilePath)
+					return errors.Wrapf(err, "getting metadata for package %q", args[0])
 				}
 
 				docsOutDir := filepath.Join(baseDocsOutDir, metadata.Name, "api-docs")
-
 				err = genResourceDocsForPackageFromRegistryMetadata(metadata, docsOutDir, basePackageTreeJSONOutDir)
 				if err != nil {
-					return errors.Wrapf(err, "generating docs for package %q from registry metadata", pkgName)
+					return errors.Wrapf(err, "generating docs for package %q from registry metadata", args[0])
 				}
 			} else {
 				glog.Infoln("Generating docs for all packages in the registry...")
-				err := genResourceDocsForAllRegistryPackages(registryDir, baseDocsOutDir, basePackageTreeJSONOutDir)
+				err := genResourceDocsForAllRegistryPackages(ctx, provider, baseDocsOutDir, basePackageTreeJSONOutDir)
 				if err != nil {
 					return errors.Wrap(err, "generating docs for all packages from registry metadata")
 				}
@@ -234,6 +215,10 @@ func resourceDocsFromRegistryCmd() *cobra.Command {
 	cmd.Flags().StringVar(&registryDir, "registryDir",
 		".",
 		"The root of the pulumi/registry directory")
+	cmd.Flags().BoolVar(&useAPI, "use-api", false, "Use the Pulumi Registry API instead of local files")
+	cmd.Flags().StringVar(&apiURL, "api-url",
+		"https://api.pulumi.com/api/preview/registry",
+		"URL of the Pulumi Registry API")
 
 	return cmd
 }

tools/resourcedocsgen/pkg/registry/svc/metadata.go

@@ -0,0 +1,252 @@
+// Copyright 2024, Pulumi Corporation.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package svc
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"os"
+	"path/filepath"
+	"slices"
+	"strings"
+	"time"
+
+	"github.com/ghodss/yaml"
+	"github.com/pkg/errors"
+
+	"github.com/pulumi/registry/tools/resourcedocsgen/pkg"
+)
+
+// PackageMetadataProvider is an interface for providers that retrieve package metadata
+// from either the filesystem directory or the Pulumi Registry API.
+type PackageMetadataProvider interface {
+	// GetPackageMetadata returns metadata for a specific package
+	GetPackageMetadata(ctx context.Context, pkgName string) (pkg.PackageMeta, error)
+	// ListPackageMetadata returns metadata for all packages
+	ListPackageMetadata(ctx context.Context) ([]pkg.PackageMeta, error)
+}
+
+// fileSystemProvider implements PackageMetadataProvider using the local yaml data files
+// in the pulumi/registry repository.
+type fileSystemProvider struct {
+	registryDir string
+}
+
+// registryAPIProvider implements PackageMetadataProvider using the Pulumi API
+// to retrieve package metadata.
+type registryAPIProvider struct {
+	apiURL string
+	client *http.Client
+}
+
+// PackageMetadata represents the API response structure for package metadata
+// from the Pulumi Registry API.
+// TODO: will eventually be available in the "github.com/pulumi/pulumi/sdk/v3/go/common/apitype" package.
+type PackageMetadata struct {
+	Name          string    `json:"name"`
+	Publisher     string    `json:"publisher"`
+	Source        string    `json:"source"`
+	Version       string    `json:"version"`
+	Title         string    `json:"title,omitempty"`
+	Description   string    `json:"description,omitempty"`
+	LogoURL       string    `json:"logoUrl,omitempty"`
+	RepoURL       string    `json:"repoUrl,omitempty"`
+	Category      string    `json:"category,omitempty"`
+	IsFeatured    bool      `json:"isFeatured"`
+	PackageTypes  []string  `json:"packageTypes,omitempty"`
+	PackageStatus string    `json:"packageStatus"`
+	SchemaURL     string    `json:"schemaURL"`
+	CreatedAt     time.Time `json:"createdAt"`
+}
+
+// PackageListResponse represents the API response structure for package lists
+type PackageListResponse struct {
+	Packages []PackageMetadata `json:"packages"`
+}
+
+// NewFileSystemProvider creates a new PackageMetadataProvider that reads from the filesystem
+func NewFileSystemProvider(registryDir string) PackageMetadataProvider {
+	return &fileSystemProvider{
+		registryDir: registryDir,
+	}
+}
+
+// NewAPIProvider creates a new PackageMetadataProvider that reads from the Pulumi API
+func NewAPIProvider(apiURL string) PackageMetadataProvider {
+	return &registryAPIProvider{
+		apiURL: apiURL,
+		client: http.DefaultClient,
+	}
+}
+
+func getRegistryPackagesPath(repoPath string) string {
+	return filepath.Join(repoPath, "themes", "default", "data", "registry", "packages")
+}
+
+func convertAPIPackageToPackageMeta(apiPkg PackageMetadata) pkg.PackageMeta {
+	return pkg.PackageMeta{
+		Name:          apiPkg.Name,
+		Publisher:     apiPkg.Publisher,
+		Description:   apiPkg.Description,
+		LogoURL:       apiPkg.LogoURL,
+		RepoURL:       apiPkg.RepoURL,
+		Category:      pkg.PackageCategory(apiPkg.Category),
+		Featured:      apiPkg.IsFeatured,
+		Native:        slices.Contains(apiPkg.PackageTypes, "native"),
+		Component:     slices.Contains(apiPkg.PackageTypes, "component"),
+		PackageStatus: pkg.PackageStatus(apiPkg.PackageStatus),
+		SchemaFileURL: apiPkg.SchemaURL,
+		Version:       apiPkg.Version,
+		Title:         apiPkg.Title,
+		UpdatedOn:     apiPkg.CreatedAt.Unix(),
+	}
+}
+
+// GetPackageMetadata implements PackageMetadataProvider for fileSystemProvider
+func (p *fileSystemProvider) GetPackageMetadata(ctx context.Context, pkgName string) (pkg.PackageMeta, error) {
+	metadataFilePath := filepath.Join(getRegistryPackagesPath(p.registryDir), pkgName+".yaml")
+	b, err := os.ReadFile(metadataFilePath)
+	if err != nil {
+		return pkg.PackageMeta{}, errors.Wrapf(err, "reading the metadata file %s", metadataFilePath)
+	}
+
+	var metadata pkg.PackageMeta
+	if err := yaml.Unmarshal(b, &metadata); err != nil {
+		return pkg.PackageMeta{}, errors.Wrapf(err, "unmarshalling the metadata file %s", metadataFilePath)
+	}
+
+	return metadata, nil
+}
+
+// ListPackageMetadata implements PackageMetadataProvider for fileSystemProvider
+func (p *fileSystemProvider) ListPackageMetadata(ctx context.Context) ([]pkg.PackageMeta, error) {
+	registryPackagesPath := getRegistryPackagesPath(p.registryDir)
+	files, err := os.ReadDir(registryPackagesPath)
+	if err != nil {
+		return nil, errors.Wrapf(err, "reading directory %s", registryPackagesPath)
+	}
+
+	// Count YAML files to pre-allocate the slice mostly to appease the linter.
+	var metadataCount int
+	for _, file := range files {
+		if strings.HasSuffix(file.Name(), ".yaml") {
+			metadataCount++
+		}
+	}
+
+	metadataList := make([]pkg.PackageMeta, 0, metadataCount)
+	for _, file := range files {
+		if !strings.HasSuffix(file.Name(), ".yaml") {
+			continue
+		}
+
+		metadata, err := p.GetPackageMetadata(ctx, strings.TrimSuffix(file.Name(), ".yaml"))
+		if err != nil {
+			return nil, err
+		}
+		metadataList = append(metadataList, metadata)
+	}
+
+	return metadataList, nil
+}
+
+// GetPackageMetadata implements PackageMetadataProvider for registryAPIProvider
+func (p *registryAPIProvider) GetPackageMetadata(ctx context.Context, pkgName string) (pkg.PackageMeta, error) {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet,
+		fmt.Sprintf("%s/packages?name=%s", p.apiURL, pkgName), nil)
+	if err != nil {
+		return pkg.PackageMeta{}, errors.Wrapf(err, "creating request for package %s", pkgName)
+	}
+
+	resp, err := p.client.Do(req)
+	if err != nil {
+		return pkg.PackageMeta{}, errors.Wrapf(err, "fetching package metadata from API for %s", pkgName)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return pkg.PackageMeta{}, errors.Errorf("unexpected status code %d when fetching package metadata", resp.StatusCode)
+	}
+
+	var response PackageListResponse
+	if err := json.NewDecoder(resp.Body).Decode(&response); err != nil {
+		return pkg.PackageMeta{}, errors.Wrap(err, "decoding API response")
+	}
+
+	switch len(response.Packages) {
+	case 0:
+		return pkg.PackageMeta{}, errors.Errorf("no package found with name %s", pkgName)
+	case 1:
+		metadata := convertAPIPackageToPackageMeta(response.Packages[0])
+		return metadata, nil
+	default:
+		return pkg.PackageMeta{}, errors.Errorf("multiple packages found with name %s", pkgName)
+	}
+}
+
+// ListPackageMetadata implements PackageMetadataProvider for registryAPIProvider
+func (p *registryAPIProvider) ListPackageMetadata(ctx context.Context) ([]pkg.PackageMeta, error) {
+	var allPackages []pkg.PackageMeta
+	// Maximum allowed by the API (must be less than 500). Request up to 499 to
+	// account for pagination with minimum number of requests.
+	const limit = 499
+	continuationToken := ""
+
+	for {
+		url := fmt.Sprintf("%s/packages?limit=%d", p.apiURL, limit)
+		if continuationToken != "" {
+			url = fmt.Sprintf("%s&continuationToken=%s", url, continuationToken)
+		}
+
+		req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
+		if err != nil {
+			return nil, errors.Wrap(err, "creating request for package list")
+		}
+
+		resp, err := p.client.Do(req)
+		if err != nil {
+			return nil, errors.Wrap(err, "fetching package list from API")
+		}
+		defer resp.Body.Close()
+
+		if resp.StatusCode != http.StatusOK {
+			return nil, errors.Errorf("unexpected status code %d when fetching package list", resp.StatusCode)
+		}
+
+		var response struct {
+			Packages          []PackageMetadata `json:"packages"`
+			ContinuationToken *string           `json:"continuationToken"`
+		}
+		if err := json.NewDecoder(resp.Body).Decode(&response); err != nil {
+			return nil, errors.Wrap(err, "decoding API response")
+		}
+
+		for _, apiPkg := range response.Packages {
+			metadata := convertAPIPackageToPackageMeta(apiPkg)
+			allPackages = append(allPackages, metadata)
+		}
+
+		// If there's no continuation token, we've reached the end
+		if response.ContinuationToken == nil {
+			break
+		}
+
+		continuationToken = *response.ContinuationToken
+	}
+
+	return allPackages, nil
+}