feat(datastore): Adding reserve IDs support (#9027)

datastore/client.go

@@ -137,6 +137,17 @@ func (dc *datastoreClient) AllocateIds(ctx context.Context, in *pb.AllocateIdsRe
 	return res, err
 }
 
+func (dc *datastoreClient) ReserveIds(ctx context.Context, in *pb.ReserveIdsRequest, opts ...grpc.CallOption) (res *pb.ReserveIdsResponse, err error) {
+	ctx = trace.StartSpan(ctx, "cloud.google.com/go/datastore.datastoreClient.ReserveIds")
+	defer func() { trace.EndSpan(ctx, err) }()
+
+	err = dc.invoke(ctx, func(ctx context.Context) error {
+		res, err = dc.c.ReserveIds(ctx, in, opts...)
+		return err
+	})
+	return res, err
+}
+
 func (dc *datastoreClient) invoke(ctx context.Context, f func(ctx context.Context) error) error {
 	ctx = metadata.NewOutgoingContext(ctx, dc.md)
 	return cloudinternal.Retry(ctx, gax.Backoff{Initial: 100 * time.Millisecond}, func() (stop bool, err error) {

datastore/integration_test.go

@@ -1224,6 +1224,21 @@ func TestIntegration_Projection(t *testing.T) {
 	})
 }
 
+func TestIntegration_ReserveIDs(t *testing.T) {
+	ctx := context.Background()
+	client := newTestClient(ctx, t)
+	defer client.Close()
+
+	keys := make([]*Key, 3)
+	for i := range keys {
+		keys[i] = NameKey("ReserveIDs", "id-"+fmt.Sprint(i), nil)
+	}
+	err := client.ReserveIDs(ctx, keys)
+	if err != nil {
+		t.Fatalf("ReserveIDs failed: %v", err)
+	}
+}
+
 func TestIntegration_AllocateIDs(t *testing.T) {
 	ctx := context.Background()
 	client := newTestClient(ctx, t)

datastore/key.go

@@ -263,6 +263,22 @@ func (c *Client) AllocateIDs(ctx context.Context, keys []*Key) ([]*Key, error) {
 	return multiProtoToKey(resp.Keys)
 }
 
+// ReserveIDs accepts a slice of keys and prevents them from being auto-allocated by Datastore
+func (c *Client) ReserveIDs(ctx context.Context, keys []*Key) error {
+	if keys == nil {
+		return nil
+	}
+
+	req := &pb.ReserveIdsRequest{
+		ProjectId:  c.dataset,
+		DatabaseId: c.databaseID,
+		Keys:       multiKeyToProto(keys),
+	}
+
+	_, err := c.client.ReserveIds(ctx, req)
+	return err
+}
+
 // IncompleteKey creates a new incomplete key.
 // The supplied kind cannot be empty.
 // The namespace of the new key is empty.