chore(bigquery/storage/managedwriter): internal refactor (flow controller, ids)

bigquery/storage/managedwriter/client.go

@@ -23,6 +23,7 @@ import (
 	"cloud.google.com/go/bigquery/internal"
 	storage "cloud.google.com/go/bigquery/storage/apiv1"
 	"cloud.google.com/go/internal/detect"
+	"github.com/google/uuid"
 	"github.com/googleapis/gax-go/v2"
 	"google.golang.org/api/option"
 	storagepb "google.golang.org/genproto/googleapis/cloud/bigquery/storage/v1"
@@ -38,6 +39,8 @@ import (
 // does not have the project ID encoded.
 const DetectProjectID = "*detect-project-id*"
 
+const managedstreamIDPrefix = "managedstream"
+
 // Client is a managed BigQuery Storage write client scoped to a single project.
 type Client struct {
 	rawClient *storage.BigQueryWriteClient
@@ -106,6 +109,7 @@ func (c *Client) buildManagedStream(ctx context.Context, streamFunc streamClient
 	ctx, cancel := context.WithCancel(ctx)
 
 	ms := &ManagedStream{
+		id:             newUUID(managedstreamIDPrefix),
 		streamSettings: defaultStreamSettings(),
 		c:              c,
 		ctx:            ctx,
@@ -232,3 +236,9 @@ func TableParentFromStreamName(streamName string) string {
 func TableParentFromParts(projectID, datasetID, tableID string) string {
 	return fmt.Sprintf("projects/%s/datasets/%s/tables/%s", projectID, datasetID, tableID)
 }
+
+// newUUID simplifies generating UUIDs for internal resources.
+func newUUID(prefix string) string {
+	id := uuid.New()
+	return fmt.Sprintf("%s_%s", prefix, id.String())
+}

bigquery/storage/managedwriter/flow_controller.go

@@ -50,6 +50,17 @@ func newFlowController(maxInserts, maxInsertBytes int) *flowController {
 	return fc
 }
 
+// copyFlowController is for creating a new flow controller based on
+// settings from another.  It does not copy flow state.
+func copyFlowController(in *flowController) *flowController {
+	var maxInserts, maxBytes int
+	if in != nil {
+		maxInserts = in.maxInsertCount
+		maxBytes = in.maxInsertBytes
+	}
+	return newFlowController(maxInserts, maxBytes)
+}
+
 // acquire blocks until one insert of size bytes can proceed or ctx is done.
 // It returns nil in the first case, or ctx.Err() in the second.
 //

bigquery/storage/managedwriter/flow_controller_test.go

@@ -257,3 +257,40 @@ func TestFlowControllerUnboundedBytes(t *testing.T) {
 		t.Error("got true, wanted false")
 	}
 }
+
+func TestCopyFlowController(t *testing.T) {
+	testcases := []struct {
+		description     string
+		in              *flowController
+		wantMaxRequests int
+		wantMaxBytes    int
+	}{
+		{
+			description:     "nil source",
+			wantMaxRequests: 0,
+			wantMaxBytes:    0,
+		},
+		{
+			description:     "no limit",
+			in:              newFlowController(0, 0),
+			wantMaxRequests: 0,
+			wantMaxBytes:    0,
+		},
+		{
+			description:     "bounded",
+			in:              newFlowController(10, 1024),
+			wantMaxRequests: 10,
+			wantMaxBytes:    1024,
+		},
+	}
+
+	for _, tc := range testcases {
+		fc := copyFlowController(tc.in)
+		if fc.maxInsertBytes != tc.wantMaxBytes {
+			t.Errorf("%s: max bytes mismatch, got %d want %d ", tc.description, fc.maxInsertBytes, tc.wantMaxBytes)
+		}
+		if fc.maxInsertCount != tc.wantMaxRequests {
+			t.Errorf("%s: max requests mismatch, got %d want %d ", tc.description, fc.maxInsertBytes, tc.wantMaxBytes)
+		}
+	}
+}

bigquery/storage/managedwriter/integration_test.go

@@ -187,6 +187,9 @@ func testDefaultStream(ctx context.Context, t *testing.T, mwClient *Client, bqCl
 	if err != nil {
 		t.Fatalf("NewManagedStream: %v", err)
 	}
+	if ms.id == "" {
+		t.Errorf("managed stream is missing ID")
+	}
 	validateTableConstraints(ctx, t, bqClient, testTable, "before send",
 		withExactRowCount(0))
 

bigquery/storage/managedwriter/managed_stream.go

@@ -72,6 +72,9 @@ func streamTypeToEnum(t StreamType) storagepb.WriteStream_Type {
 
 // ManagedStream is the abstraction over a single write stream.
 type ManagedStream struct {
+	// Unique id for the managedstream instance.
+	id string
+
 	streamSettings   *streamSettings
 	schemaDescriptor *descriptorpb.DescriptorProto
 	destinationTable string