logadmin: use generated iterators

Change the logadmin iterators to use the iterators
from the underlying generated client.

Change-Id: Ie4be1d645ae0058bea88a1bc438e1bfa82945349
Reviewed-on: https://code-review.googlesource.com/8997
Reviewed-by: Jun Mukai <mukai@google.com>
Reviewed-by: Michael Darakananda <pongad@google.com>
diff --git a/logging/logadmin/logadmin.go b/logging/logadmin/logadmin.go
index 1d5c792..6ccd0e0 100644
--- a/logging/logadmin/logadmin.go
+++ b/logging/logadmin/logadmin.go
@@ -40,7 +40,6 @@
 	vkit "cloud.google.com/go/logging/apiv2"
 	"cloud.google.com/go/logging/internal"
 	"github.com/golang/protobuf/ptypes"
-	gax "github.com/googleapis/gax-go"
 	"golang.org/x/net/context"
 	"google.golang.org/api/iterator"
 	"google.golang.org/api/option"
@@ -207,9 +206,7 @@
 // NewClient. This may be overridden by passing a ProjectIDs option. Requires ReadScope or AdminScope.
 func (c *Client) Entries(ctx context.Context, opts ...EntriesOption) *EntryIterator {
 	it := &EntryIterator{
-		ctx:    ctx,
-		client: c.lClient,
-		req:    listLogEntriesRequest(c.projectID, opts),
+		it: c.lClient.ListLogEntries(ctx, listLogEntriesRequest(c.projectID, opts)),
 	}
 	it.pageInfo, it.nextFunc = iterator.NewPageInfo(
 		it.fetch,
@@ -230,11 +227,9 @@
 
 // An EntryIterator iterates over log entries.
 type EntryIterator struct {
-	ctx      context.Context
-	client   *vkit.Client
+	it       *vkit.LogEntryIterator
 	pageInfo *iterator.PageInfo
 	nextFunc func() error
-	req      *logpb.ListLogEntriesRequest
 	items    []*logging.Entry
 }
 
@@ -254,28 +249,18 @@
 }
 
 func (it *EntryIterator) fetch(pageSize int, pageToken string) (string, error) {
-	// TODO(jba): Do this a nicer way if the generated code supports one.
-	// TODO(jba): If the above TODO can't be done, find a way to pass metadata in the call.
-	client := logpb.NewLoggingServiceV2Client(it.client.Connection())
-	var res *logpb.ListLogEntriesResponse
-	err := gax.Invoke(it.ctx, func(ctx context.Context) error {
-		it.req.PageSize = trunc32(pageSize)
-		it.req.PageToken = pageToken
-		var err error
-		res, err = client.ListLogEntries(ctx, it.req)
-		return err
-	}, it.client.CallOptions.ListLogEntries...)
-	if err != nil {
-		return "", err
-	}
-	for _, ep := range res.Entries {
-		e, err := fromLogEntry(ep)
+	return iterFetch(pageSize, pageToken, it.it.PageInfo(), func() error {
+		item, err := it.it.Next()
 		if err != nil {
-			return "", err
+			return err
+		}
+		e, err := fromLogEntry(item)
+		if err != nil {
+			return err
 		}
 		it.items = append(it.items, e)
-	}
-	return res.NextPageToken, nil
+		return nil
+	})
 }
 
 func trunc32(i int) int32 {
@@ -328,3 +313,20 @@
 		Resource:    le.Resource,
 	}, nil
 }
+
+// Common fetch code for iterators that are backed by vkit iterators.
+func iterFetch(pageSize int, pageToken string, pi *iterator.PageInfo, next func() error) (string, error) {
+	pi.MaxSize = pageSize
+	pi.Token = pageToken
+	// Get one item, which will fill the buffer.
+	if err := next(); err != nil {
+		return "", err
+	}
+	// Collect the rest of the buffer.
+	for pi.Remaining() > 0 {
+		if err := next(); err != nil {
+			return "", err
+		}
+	}
+	return pi.Token, nil
+}
diff --git a/logging/logadmin/metrics.go b/logging/logadmin/metrics.go
index 51c8a89..9374ac4 100644
--- a/logging/logadmin/metrics.go
+++ b/logging/logadmin/metrics.go
@@ -18,7 +18,6 @@
 	"fmt"
 
 	vkit "cloud.google.com/go/logging/apiv2"
-	gax "github.com/googleapis/gax-go"
 	"golang.org/x/net/context"
 	"google.golang.org/api/iterator"
 	logpb "google.golang.org/genproto/googleapis/logging/v2"
@@ -95,9 +94,7 @@
 // Requires ReadScope or AdminScope.
 func (c *Client) Metrics(ctx context.Context) *MetricIterator {
 	it := &MetricIterator{
-		ctx:    ctx,
-		client: c.mClient,
-		req:    &logpb.ListLogMetricsRequest{Parent: c.parent()},
+		it: c.mClient.ListLogMetrics(ctx, &logpb.ListLogMetricsRequest{Parent: c.parent()}),
 	}
 	it.pageInfo, it.nextFunc = iterator.NewPageInfo(
 		it.fetch,
@@ -108,11 +105,9 @@
 
 // A MetricIterator iterates over Metrics.
 type MetricIterator struct {
-	ctx      context.Context
-	client   *vkit.MetricsClient
+	it       *vkit.LogMetricIterator
 	pageInfo *iterator.PageInfo
 	nextFunc func() error
-	req      *logpb.ListLogMetricsRequest
 	items    []*Metric
 }
 
@@ -132,24 +127,14 @@
 }
 
 func (it *MetricIterator) fetch(pageSize int, pageToken string) (string, error) {
-	// TODO(jba): Do this a nicer way if the generated code supports one.
-	// TODO(jba): If the above TODO can't be done, find a way to pass metadata in the call.
-	client := logpb.NewMetricsServiceV2Client(it.client.Connection())
-	var res *logpb.ListLogMetricsResponse
-	err := gax.Invoke(it.ctx, func(ctx context.Context) error {
-		it.req.PageSize = trunc32(pageSize)
-		it.req.PageToken = pageToken
-		var err error
-		res, err = client.ListLogMetrics(ctx, it.req)
-		return err
-	}, it.client.CallOptions.ListLogMetrics...)
-	if err != nil {
-		return "", err
-	}
-	for _, sp := range res.Metrics {
-		it.items = append(it.items, fromLogMetric(sp))
-	}
-	return res.NextPageToken, nil
+	return iterFetch(pageSize, pageToken, it.it.PageInfo(), func() error {
+		item, err := it.it.Next()
+		if err != nil {
+			return err
+		}
+		it.items = append(it.items, fromLogMetric(item))
+		return nil
+	})
 }
 
 func toLogMetric(m *Metric) *logpb.LogMetric {
diff --git a/logging/logadmin/resources.go b/logging/logadmin/resources.go
index 7345fd6..79e8fdb 100644
--- a/logging/logadmin/resources.go
+++ b/logging/logadmin/resources.go
@@ -16,7 +16,6 @@
 
 import (
 	vkit "cloud.google.com/go/logging/apiv2"
-	gax "github.com/googleapis/gax-go"
 	"golang.org/x/net/context"
 	"google.golang.org/api/iterator"
 	mrpb "google.golang.org/genproto/googleapis/api/monitoredres"
@@ -30,9 +29,8 @@
 // See https://cloud.google.com/logging/docs/api/v2/resource-list for a list of monitored resources.
 func (c *Client) ResourceDescriptors(ctx context.Context) *ResourceDescriptorIterator {
 	it := &ResourceDescriptorIterator{
-		ctx:    ctx,
-		client: c.lClient,
-		req:    &logpb.ListMonitoredResourceDescriptorsRequest{},
+		it: c.lClient.ListMonitoredResourceDescriptors(ctx,
+			&logpb.ListMonitoredResourceDescriptorsRequest{}),
 	}
 	it.pageInfo, it.nextFunc = iterator.NewPageInfo(
 		it.fetch,
@@ -43,11 +41,9 @@
 
 // ResourceDescriptorIterator is an iterator over MonitoredResourceDescriptors.
 type ResourceDescriptorIterator struct {
-	ctx      context.Context
-	client   *vkit.Client
+	it       *vkit.MonitoredResourceDescriptorIterator
 	pageInfo *iterator.PageInfo
 	nextFunc func() error
-	req      *logpb.ListMonitoredResourceDescriptorsRequest
 	items    []*mrpb.MonitoredResourceDescriptor
 }
 
@@ -67,20 +63,12 @@
 }
 
 func (it *ResourceDescriptorIterator) fetch(pageSize int, pageToken string) (string, error) {
-	// TODO(jba): Do this a nicer way if the generated code supports one.
-	// TODO(jba): If the above TODO can't be done, find a way to pass metadata in the call.
-	client := logpb.NewLoggingServiceV2Client(it.client.Connection())
-	var res *logpb.ListMonitoredResourceDescriptorsResponse
-	err := gax.Invoke(it.ctx, func(ctx context.Context) error {
-		it.req.PageSize = trunc32(pageSize)
-		it.req.PageToken = pageToken
-		var err error
-		res, err = client.ListMonitoredResourceDescriptors(ctx, it.req)
-		return err
-	}, it.client.CallOptions.ListMonitoredResourceDescriptors...)
-	if err != nil {
-		return "", err
-	}
-	it.items = append(it.items, res.ResourceDescriptors...)
-	return res.NextPageToken, nil
+	return iterFetch(pageSize, pageToken, it.it.PageInfo(), func() error {
+		item, err := it.it.Next()
+		if err != nil {
+			return err
+		}
+		it.items = append(it.items, item)
+		return nil
+	})
 }
diff --git a/logging/logadmin/sinks.go b/logging/logadmin/sinks.go
index 52353c5..588c7af 100644
--- a/logging/logadmin/sinks.go
+++ b/logging/logadmin/sinks.go
@@ -18,7 +18,6 @@
 	"fmt"
 
 	vkit "cloud.google.com/go/logging/apiv2"
-	gax "github.com/googleapis/gax-go"
 	"golang.org/x/net/context"
 	"google.golang.org/api/iterator"
 	logpb "google.golang.org/genproto/googleapis/logging/v2"
@@ -109,9 +108,7 @@
 // Requires ReadScope or AdminScope.
 func (c *Client) Sinks(ctx context.Context) *SinkIterator {
 	it := &SinkIterator{
-		ctx:    ctx,
-		client: c.sClient,
-		req:    &logpb.ListSinksRequest{Parent: c.parent()},
+		it: c.sClient.ListSinks(ctx, &logpb.ListSinksRequest{Parent: c.parent()}),
 	}
 	it.pageInfo, it.nextFunc = iterator.NewPageInfo(
 		it.fetch,
@@ -122,11 +119,9 @@
 
 // A SinkIterator iterates over Sinks.
 type SinkIterator struct {
-	ctx      context.Context
-	client   *vkit.ConfigClient
+	it       *vkit.LogSinkIterator
 	pageInfo *iterator.PageInfo
 	nextFunc func() error
-	req      *logpb.ListSinksRequest
 	items    []*Sink
 }
 
@@ -146,24 +141,14 @@
 }
 
 func (it *SinkIterator) fetch(pageSize int, pageToken string) (string, error) {
-	// TODO(jba): Do this a nicer way if the generated code supports one.
-	// TODO(jba): If the above TODO can't be done, find a way to pass metadata in the call.
-	client := logpb.NewConfigServiceV2Client(it.client.Connection())
-	var res *logpb.ListSinksResponse
-	err := gax.Invoke(it.ctx, func(ctx context.Context) error {
-		it.req.PageSize = trunc32(pageSize)
-		it.req.PageToken = pageToken
-		var err error
-		res, err = client.ListSinks(ctx, it.req)
-		return err
-	}, it.client.CallOptions.ListSinks...)
-	if err != nil {
-		return "", err
-	}
-	for _, sp := range res.Sinks {
-		it.items = append(it.items, fromLogSink(sp))
-	}
-	return res.NextPageToken, nil
+	return iterFetch(pageSize, pageToken, it.it.PageInfo(), func() error {
+		item, err := it.it.Next()
+		if err != nil {
+			return err
+		}
+		it.items = append(it.items, fromLogSink(item))
+		return nil
+	})
 }
 
 func toLogSink(s *Sink) *logpb.LogSink {