Swapping ProjectId with ProjectNubmer for DataStream APIs (#733)

Author: VardhanThigle

go.mod

@@ -8,6 +8,7 @@ require (
 	cloud.google.com/go/datastream v1.10.3
 	cloud.google.com/go/monitoring v1.16.3
 	cloud.google.com/go/pubsub v1.33.0
+	cloud.google.com/go/resourcemanager v1.9.4
 	cloud.google.com/go/spanner v1.53.0
 	cloud.google.com/go/storage v1.30.1
 	github.com/DATA-DOG/go-sqlmock v1.5.1

go.sum

@@ -46,6 +46,8 @@ cloud.google.com/go/pubsub v1.2.0/go.mod h1:jhfEVHT8odbXTkndysNHCcx0awwzvfOlguIA
 cloud.google.com/go/pubsub v1.3.1/go.mod h1:i+ucay31+CNRpDW4Lu78I4xXG+O1r/MAHgjpRVR+TSU=
 cloud.google.com/go/pubsub v1.33.0 h1:6SPCPvWav64tj0sVX/+npCBKhUi/UjJehy9op/V3p2g=
 cloud.google.com/go/pubsub v1.33.0/go.mod h1:f+w71I33OMyxf9VpMVcZbnG5KSUkCOUHYpFd5U1GdRc=
+cloud.google.com/go/resourcemanager v1.9.4 h1:JwZ7Ggle54XQ/FVYSBrMLOQIKoIT/uer8mmNvNLK51k=
+cloud.google.com/go/resourcemanager v1.9.4/go.mod h1:N1dhP9RFvo3lUfwtfLWVxfUWq8+KUQ+XLlHLH3BoFJ0=
 cloud.google.com/go/spanner v1.53.0 h1:/NzWQJ1MEhdRcffiutRKbW/AIGVKhcTeivWTDjEyCCo=
 cloud.google.com/go/spanner v1.53.0/go.mod h1:liG4iCeLqm5L3fFLU5whFITqP0e0orsAW1uUSrd4rws=
 cloud.google.com/go/storage v1.0.0/go.mod h1:IhtSnM/ZTZV8YYJWCY8RULGVqBDmpoyjwiyrjsg+URw=

streaming/streaming.go

@@ -31,6 +31,8 @@ import (
 	dataflowpb "google.golang.org/genproto/googleapis/dataflow/v1beta3"
 	"google.golang.org/protobuf/types/known/fieldmaskpb"
 
+	resourcemanager "cloud.google.com/go/resourcemanager/apiv3"
+	resourcemanagerpb "cloud.google.com/go/resourcemanager/apiv3/resourcemanagerpb"
 	"github.com/GoogleCloudPlatform/spanner-migration-tool/common/constants"
 	"github.com/GoogleCloudPlatform/spanner-migration-tool/common/utils"
 	"github.com/GoogleCloudPlatform/spanner-migration-tool/internal"
@@ -106,7 +108,7 @@ type DataflowCfg struct {
 
 type StreamingCfg struct {
 	DatastreamCfg DatastreamCfg            `json:"datastreamCfg"`
-	GcsCfg        GcsCfg             `json:"gcsCfg"`
+	GcsCfg        GcsCfg                   `json:"gcsCfg"`
 	DataflowCfg   DataflowCfg              `json:"dataflowCfg"`
 	TmpDir        string                   `json:"tmpDir"`
 	PubsubCfg     internal.PubsubResources `json:"pubsubCfg"`
@@ -463,7 +465,8 @@ func createNotificationOnBucket(ctx context.Context, storageClient *storage.Clie
 
 // LaunchStream populates the parameters from the streaming config and triggers a stream on Cloud Datastream.
 func LaunchStream(ctx context.Context, sourceProfile profiles.SourceProfile, dbList []profiles.LogicalShard, projectID string, datastreamCfg DatastreamCfg) error {
-	fmt.Println("Launching stream ", fmt.Sprintf("projects/%s/locations/%s", projectID, datastreamCfg.StreamLocation))
+	projectNumberResource := GetProjectNumberResource(ctx, fmt.Sprintf("projects/%s", projectID))
+	fmt.Println("Launching stream ", fmt.Sprintf("%s/locations/%s", projectNumberResource, datastreamCfg.StreamLocation))
 	dsClient, err := datastream.NewClient(ctx)
 	if err != nil {
 		return fmt.Errorf("datastream client can not be created: %v", err)
@@ -478,15 +481,15 @@ func LaunchStream(ctx context.Context, sourceProfile profiles.SourceProfile, dbL
 		FileFormat: &datastreampb.GcsDestinationConfig_AvroFileFormat{},
 	}
 	srcCfg := &datastreampb.SourceConfig{
-		SourceConnectionProfile: fmt.Sprintf("projects/%s/locations/%s/connectionProfiles/%s", projectID, datastreamCfg.SourceConnectionConfig.Location, datastreamCfg.SourceConnectionConfig.Name),
+		SourceConnectionProfile: fmt.Sprintf("%s/locations/%s/connectionProfiles/%s", projectNumberResource, datastreamCfg.SourceConnectionConfig.Location, datastreamCfg.SourceConnectionConfig.Name),
 	}
 	err = getSourceStreamConfig(srcCfg, sourceProfile, dbList, datastreamCfg)
 	if err != nil {
 		return fmt.Errorf("could not get source stream config: %v", err)
 	}
 
 	dstCfg := &datastreampb.DestinationConfig{
-		DestinationConnectionProfile: fmt.Sprintf("projects/%s/locations/%s/connectionProfiles/%s", projectID, datastreamCfg.DestinationConnectionConfig.Location, datastreamCfg.DestinationConnectionConfig.Name),
+		DestinationConnectionProfile: fmt.Sprintf("%s/locations/%s/connectionProfiles/%s", projectNumberResource, datastreamCfg.DestinationConnectionConfig.Location, datastreamCfg.DestinationConnectionConfig.Name),
 		DestinationStreamConfig:      &datastreampb.DestinationConfig_GcsDestinationConfig{GcsDestinationConfig: gcsDstCfg},
 	}
 	streamInfo := &datastreampb.Stream{
@@ -497,7 +500,7 @@ func LaunchStream(ctx context.Context, sourceProfile profiles.SourceProfile, dbL
 		BackfillStrategy:  &datastreampb.Stream_BackfillAll{BackfillAll: &datastreampb.Stream_BackfillAllStrategy{}},
 	}
 	createStreamRequest := &datastreampb.CreateStreamRequest{
-		Parent:   fmt.Sprintf("projects/%s/locations/%s", projectID, datastreamCfg.StreamLocation),
+		Parent:   fmt.Sprintf("%s/locations/%s", projectNumberResource, datastreamCfg.StreamLocation),
 		StreamId: datastreamCfg.StreamId,
 		Stream:   streamInfo,
 	}
@@ -518,7 +521,7 @@ func LaunchStream(ctx context.Context, sourceProfile profiles.SourceProfile, dbL
 	fmt.Println("Successfully created stream ", datastreamCfg.StreamId)
 
 	fmt.Print("Setting stream state to RUNNING...")
-	streamInfo.Name = fmt.Sprintf("projects/%s/locations/%s/streams/%s", projectID, datastreamCfg.StreamLocation, datastreamCfg.StreamId)
+	streamInfo.Name = fmt.Sprintf("%s/locations/%s/streams/%s", projectNumberResource, datastreamCfg.StreamLocation, datastreamCfg.StreamId)
 	updateStreamRequest := &datastreampb.UpdateStreamRequest{
 		UpdateMask: &fieldmaskpb.FieldMask{Paths: []string{"state"}},
 		Stream:     streamInfo,
@@ -756,6 +759,36 @@ func CreateStreamingConfig(pl profiles.DataShard) StreamingCfg {
 	return streamingCfg
 }
 
+// Maps Project-Id to ProjectNumber.
+var ProjectNumberResourceCache sync.Map
+
+// Returns a string that encodes the project number like `projects/12345`
+func GetProjectNumberResource(ctx context.Context, projectID string) string {
+	projectNumberResource, found := ProjectNumberResourceCache.Load(projectID)
+	if found {
+		return projectNumberResource.(string)
+	}
+
+	rmClient, err := resourcemanager.NewProjectsClient(ctx)
+	if err != nil {
+		logger.Log.Warn(fmt.Sprintf("Could not create resourcemanager client to query project number. Defaulting to ProjectId=%s. error=%v",
+			projectID, err))
+		return projectID
+	}
+	defer rmClient.Close()
+	req := resourcemanagerpb.GetProjectRequest{Name: projectID}
+	project, err := rmClient.GetProject(ctx, &req)
+	if err != nil {
+		logger.Log.Warn(fmt.Sprintf("Could not query resourcemanager to get project number. Defaulting to ProjectId=%s. error=%v",
+			projectID, err))
+		return projectID
+	}
+	projectNumberResource = project.GetName()
+	ProjectNumberResourceCache.Store(projectID, projectNumberResource)
+	return projectNumberResource.(string)
+
+}
+
 func StartDatastream(ctx context.Context, streamingCfg StreamingCfg, sourceProfile profiles.SourceProfile, targetProfile profiles.TargetProfile, schemaDetails map[string]internal.SchemaDetails) (StreamingCfg, error) {
 	driver := sourceProfile.Driver
 	var dbList []profiles.LogicalShard