make pf data source map mutable

pkg/pf/internal/schemashim/datasource_map.go

@@ -17,39 +17,16 @@ package schemashim
 import (
 	"github.com/pulumi/pulumi-terraform-bridge/v3/pkg/pf/internal/runtypes"
 	shim "github.com/pulumi/pulumi-terraform-bridge/v3/pkg/tfshim"
+	"github.com/pulumi/pulumi-terraform-bridge/v3/pkg/tfshim/schema"
 )
 
-type schemaOnlyDataSourceMap struct {
-	dataSources runtypes.DataSources
-}
-
-var _ shim.ResourceMap = (*schemaOnlyDataSourceMap)(nil)
-
-func (m *schemaOnlyDataSourceMap) Len() int {
-	return len(m.dataSources.All())
-}
-
-func (m *schemaOnlyDataSourceMap) Get(key string) shim.Resource {
-	s := m.dataSources.Schema(runtypes.TypeName(key))
-	return &schemaOnlyDataSource{s}
-}
-
-func (m *schemaOnlyDataSourceMap) GetOk(key string) (shim.Resource, bool) {
-	if !m.dataSources.Has(runtypes.TypeName(key)) {
-		return nil, false
-	}
-	return m.Get(key), true
-}
-
-func (m *schemaOnlyDataSourceMap) Range(each func(key string, value shim.Resource) bool) {
-	for _, typeName := range m.dataSources.All() {
-		key := string(typeName)
-		if !each(key, m.Get(key)) {
-			return
-		}
+// Data Source map needs to support Set (mutability) for RenameDataSource.
+func newSchemaOnlyDataSourceMap(dataSources runtypes.DataSources) shim.ResourceMap {
+	m := schema.ResourceMap{}
+	for _, name := range dataSources.All() {
+		key := string(name)
+		v := dataSources.Schema(name)
+		m[key] = &schemaOnlyDataSource{v}
 	}
-}
-
-func (*schemaOnlyDataSourceMap) Set(key string, value shim.Resource) {
-	panic("Set not supported - is it possible to treat this as immutable?")
+	return m
 }

pkg/pf/internal/schemashim/provider.go

@@ -35,6 +35,7 @@ type SchemaOnlyProvider struct {
 	ctx         context.Context
 	tf          pfprovider.Provider
 	resourceMap shim.ResourceMap
+	dataSourceMap shim.ResourceMap
 }
 
 func (p *SchemaOnlyProvider) Server(ctx context.Context) (tfprotov6.ProviderServer, error) {
@@ -87,11 +88,7 @@ func (p *SchemaOnlyProvider) ResourcesMap() shim.ResourceMap {
 }
 
 func (p *SchemaOnlyProvider) DataSourcesMap() shim.ResourceMap {
-	dataSources, err := pfutils.GatherDatasources(context.TODO(), p.tf, NewSchemaMap)
-	if err != nil {
-		panic(err)
-	}
-	return &schemaOnlyDataSourceMap{dataSources}
+	return p.dataSourceMap
 }
 
 func (p *SchemaOnlyProvider) InternalValidate() error {

pkg/pf/internal/schemashim/schemashim.go

@@ -28,10 +28,16 @@ func ShimSchemaOnlyProvider(ctx context.Context, provider pfprovider.Provider) s
 	if err != nil {
 		panic(err)
 	}
+	dataSources, err := pfutils.GatherDatasources(ctx, provider, NewSchemaMap)
+	if err != nil {
+		panic(err)
+	}
 	resourceMap := newSchemaOnlyResourceMap(resources)
+	dataSourceMap := newSchemaOnlyDataSourceMap(dataSources)
 	return &SchemaOnlyProvider{
-		ctx:         ctx,
-		tf:          provider,
-		resourceMap: resourceMap,
+		ctx:           ctx,
+		tf:            provider,
+		resourceMap:   resourceMap,
+		dataSourceMap: dataSourceMap,
 	}
 }