BUF_TESTING_PUBLIC_REGISTRY

private/buf/bufcli/cache.go

@@ -132,6 +132,7 @@ func newModuleDataProvider(
 	delegateModuleDataProvider := bufmoduleapi.NewModuleDataProvider(
 		container.Logger(),
 		clientProvider,
+		newGraphProvider(container, clientProvider),
 	)
 	// No symlinks.
 	storageosProvider := storageos.NewProvider()

private/buf/bufcli/controller.go

@@ -50,7 +50,7 @@ func NewController(
 		container.Logger(),
 		tracing.NewTracer(container.Tracer()),
 		container,
-		bufmoduleapi.NewGraphProvider(container.Logger(), clientProvider),
+		newGraphProvider(container, clientProvider),
 		bufmoduleapi.NewModuleKeyProvider(container.Logger(), clientProvider),
 		moduleDataProvider,
 		commitProvider,

private/buf/bufcli/env.go

@@ -37,6 +37,10 @@ const (
 	// implemented now, all 1000 .proto file are copied. You could get smarter with caching
 	// at a per-file level.
 	copyToInMemoryEnvKey = "BUF_BETA_COPY_FILES_TO_MEMORY"
+
+	// This should only be used for testing. This is not part of Buf's API, and should
+	// never be documented or part of Buf's contract.
+	publicRegistryEnvKey = "BUF_TESTING_PUBLIC_REGISTRY"
 )
 
 var (

private/buf/bufcli/graph_provider.go

@@ -27,10 +27,17 @@ func NewGraphProvider(container appext.Container) (bufmodule.GraphProvider, erro
 	if err != nil {
 		return nil, err
 	}
+	return newGraphProvider(container, bufapi.NewClientProvider(clientConfig)), nil
+}
+
+func newGraphProvider(
+	container appext.Container,
+	clientProvider bufapi.ClientProvider,
+) bufmodule.GraphProvider {
 	return bufmoduleapi.NewGraphProvider(
 		container.Logger(),
-		bufapi.NewClientProvider(
-			clientConfig,
-		),
-	), nil
+		clientProvider,
+		// OK if empty
+		bufmoduleapi.GraphProviderWithPublicRegistry(container.Env(publicRegistryEnvKey)),
+	)
 }

private/buf/bufcli/uploader.go

@@ -27,10 +27,17 @@ func NewUploader(container appext.Container) (bufmodule.Uploader, error) {
 	if err != nil {
 		return nil, err
 	}
+	return newUploader(container, bufapi.NewClientProvider(clientConfig)), nil
+}
+
+func newUploader(
+	container appext.Container,
+	clientProvider bufapi.ClientProvider,
+) bufmodule.Uploader {
 	return bufmoduleapi.NewUploader(
 		container.Logger(),
-		bufapi.NewClientProvider(
-			clientConfig,
-		),
-	), nil
+		clientProvider,
+		// OK if empty
+		bufmoduleapi.UploaderWithPublicRegistry(container.Env(publicRegistryEnvKey)),
+	)
 }

private/bufpkg/bufmodule/bufmoduleapi/graph_provider.go

@@ -39,8 +39,25 @@ func NewGraphProvider(
 		bufapi.ModuleServiceClientProvider
 		bufapi.OwnerServiceClientProvider
 	},
+	options ...GraphProviderOption,
 ) bufmodule.GraphProvider {
-	return newGraphProvider(logger, clientProvider)
+	return newGraphProvider(logger, clientProvider, options...)
+}
+
+// GraphProviderOption is an option for a new GraphProvider.
+type GraphProviderOption func(*graphProvider)
+
+// GraphProviderWithPublicRegistry returns a new GraphProviderOption that specifies
+// the hostname of the public registry. By default this is "buf.build", however in testing,
+// this may be something else. This is needed to discern which which registry to make calls
+// against in the case where there is >1 registries represented in the ModuleKeys - we always
+// want to call the non-public registry.
+func GraphProviderWithPublicRegistry(publicRegistry string) GraphProviderOption {
+	return func(graphProvider *graphProvider) {
+		if publicRegistry != "" {
+			graphProvider.publicRegistry = publicRegistry
+		}
+	}
 }
 
 // *** PRIVATE ***
@@ -53,6 +70,7 @@ type graphProvider struct {
 		bufapi.ModuleServiceClientProvider
 		bufapi.OwnerServiceClientProvider
 	}
+	publicRegistry string
 }
 
 func newGraphProvider(
@@ -63,11 +81,17 @@ func newGraphProvider(
 		bufapi.ModuleServiceClientProvider
 		bufapi.OwnerServiceClientProvider
 	},
+	options ...GraphProviderOption,
 ) *graphProvider {
-	return &graphProvider{
+	graphProvider := &graphProvider{
 		logger:         logger,
 		clientProvider: clientProvider,
+		publicRegistry: defaultPublicRegistry,
+	}
+	for _, option := range options {
+		option(graphProvider)
 	}
+	return graphProvider
 }
 
 func (a *graphProvider) GetGraphForModuleKeys(
@@ -160,7 +184,7 @@ func (a *graphProvider) getProtoLegacyFederationGraphForModuleKeys(
 	moduleKeys []bufmodule.ModuleKey,
 	digestType bufmodule.DigestType,
 ) (*federationv1beta1.Graph, error) {
-	primaryRegistry, secondaryRegistry, err := getPrimarySecondaryRegistry(moduleKeys)
+	primaryRegistry, secondaryRegistry, err := getPrimarySecondaryRegistry(moduleKeys, a.publicRegistry)
 	if err != nil {
 		return nil, err
 	}

private/bufpkg/bufmodule/bufmoduleapi/module_data_provider.go

@@ -43,8 +43,9 @@ func NewModuleDataProvider(
 		bufapi.ModuleServiceClientProvider
 		bufapi.OwnerServiceClientProvider
 	},
+	graphProvider bufmodule.GraphProvider,
 ) bufmodule.ModuleDataProvider {
-	return newModuleDataProvider(logger, clientProvider)
+	return newModuleDataProvider(logger, clientProvider, graphProvider)
 }
 
 // *** PRIVATE ***
@@ -62,19 +63,15 @@ func newModuleDataProvider(
 	logger *zap.Logger,
 	clientProvider interface {
 		bufapi.DownloadServiceClientProvider
-		bufapi.GraphServiceClientProvider
-		bufapi.LegacyFederationGraphServiceClientProvider
 		bufapi.ModuleServiceClientProvider
 		bufapi.OwnerServiceClientProvider
 	},
+	graphProvider bufmodule.GraphProvider,
 ) *moduleDataProvider {
 	return &moduleDataProvider{
 		logger:         logger,
 		clientProvider: clientProvider,
-		graphProvider: NewGraphProvider(
-			logger,
-			clientProvider,
-		),
+		graphProvider:  graphProvider,
 	}
 }
 

private/bufpkg/bufmodule/bufmoduleapi/registry.go

@@ -25,7 +25,7 @@ import (
 
 // If we ever get to a case where we're supporting legacy federation, and we're moving buf.build,
 // we have way bigger problems than this hardcoded variable.
-const publicRegistry = "buf.build"
+const defaultPublicRegistry = "buf.build"
 
 type hasModuleFullName interface {
 	ModuleFullName() bufmodule.ModuleFullName
@@ -44,7 +44,7 @@ type hasModuleFullName interface {
 // beyond a non-public registry depending on buf.build.
 //
 // This is used to support legacy federation.
-func getPrimarySecondaryRegistry[T hasModuleFullName](s []T) (string, string, error) {
+func getPrimarySecondaryRegistry[T hasModuleFullName](s []T, publicRegistry string) (string, string, error) {
 	if len(s) == 0 {
 		return "", "", syserror.New("must have at least one value in getPrimarySecondaryRegistry")
 	}
@@ -80,7 +80,7 @@ func getPrimarySecondaryRegistry[T hasModuleFullName](s []T) (string, string, er
 	}
 }
 
-func validateDepRegistries(primaryRegistry string, depRegistries []string) error {
+func validateDepRegistries(primaryRegistry string, depRegistries []string, publicRegistry string) error {
 	switch len(depRegistries) {
 	case 0:
 		return nil

private/bufpkg/bufmodule/bufmoduleapi/uploader.go

@@ -39,8 +39,25 @@ func NewUploader(
 		bufapi.ModuleServiceClientProvider
 		bufapi.UploadServiceClientProvider
 	},
+	options ...UploaderOption,
 ) bufmodule.Uploader {
-	return newUploader(logger, clientProvider)
+	return newUploader(logger, clientProvider, options...)
+}
+
+// UploaderOption is an option for a new Uploader.
+type UploaderOption func(*uploader)
+
+// UploaderWithPublicRegistry returns a new UploaderOption that specifies
+// the hostname of the public registry. By default this is "buf.build", however in testing,
+// this may be something else. This is needed to discern which which registry to make calls
+// against in the case where there is >1 registries represented in the ModuleKeys - we always
+// want to call the non-public registry.
+func UploaderWithPublicRegistry(publicRegistry string) UploaderOption {
+	return func(uploader *uploader) {
+		if publicRegistry != "" {
+			uploader.publicRegistry = publicRegistry
+		}
+	}
 }
 
 // *** PRIVATE ***
@@ -52,6 +69,7 @@ type uploader struct {
 		bufapi.ModuleServiceClientProvider
 		bufapi.UploadServiceClientProvider
 	}
+	publicRegistry string
 }
 
 func newUploader(
@@ -61,11 +79,17 @@ func newUploader(
 		bufapi.ModuleServiceClientProvider
 		bufapi.UploadServiceClientProvider
 	},
+	options ...UploaderOption,
 ) *uploader {
-	return &uploader{
+	uploader := &uploader{
 		logger:         logger,
 		clientProvider: clientProvider,
+		publicRegistry: defaultPublicRegistry,
+	}
+	for _, option := range options {
+		option(uploader)
 	}
+	return uploader
 }
 
 func (a *uploader) Upload(
@@ -151,7 +175,7 @@ func (a *uploader) Upload(
 			},
 		),
 	)
-	if err := validateDepRegistries(primaryRegistry, remoteDepRegistries); err != nil {
+	if err := validateDepRegistries(primaryRegistry, remoteDepRegistries, a.publicRegistry); err != nil {
 		return nil, err
 	}