From ac638d3162376f52e4e98d7172655846e9655add Mon Sep 17 00:00:00 2001
From: Vaibhav Kamra <vkamra@alcion.ai>
Date: Fri, 16 Sep 2022 11:15:43 -0700
Subject: [PATCH] Provide interface to get stream size (#876)

## Description

Allows kopia data collection consumers to query stream size.

This is required for OneDrive that needs to set the file size prior to uploading file data.

## Type of change

<!--- Please check the type of change your PR introduces: --->
- [ ] :sunflower: Feature
- [ ] :bug: Bugfix
- [ ] :world_map: Documentation
- [ ] :robot: Test
- [ ] :computer: CI/Deployment
- [x] :hamster: Trivial/Minor

## Issue(s)

* #668

## Test Plan

<!-- How will this be tested prior to merging.-->
- [ ] :muscle: Manual
- [x] :zap: Unit test
- [ ] :green_heart: E2E
---
 src/internal/data/data_collection.go       | 12 +++++++++---
 src/internal/kopia/data_collection.go      |  5 +++++
 src/internal/kopia/data_collection_test.go |  6 ++++++
 src/internal/kopia/wrapper.go              |  1 +
 4 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/src/internal/data/data_collection.go b/src/internal/data/data_collection.go
index efd85e8a0..8e027663d 100644
--- a/src/internal/data/data_collection.go
+++ b/src/internal/data/data_collection.go
@@ -22,7 +22,7 @@ type Collection interface {
 	FullPath() path.Path
 }
 
-// DataStream represents a single item within a DataCollection
+// Stream represents a single item within a Collection
 // that can be consumed as a stream (it embeds io.Reader)
 type Stream interface {
 	// ToReader returns an io.Reader for the DataStream
@@ -31,8 +31,14 @@ type Stream interface {
 	UUID() string
 }
 
-// DataStreamInfo is used to provide service specific
-// information about the DataStream
+// StreamInfo is used to provide service specific
+// information about the Stream
 type StreamInfo interface {
 	Info() details.ItemInfo
 }
+
+// StreamSize is used to provide size
+// information about the Stream
+type StreamSize interface {
+	Size() int64
+}
diff --git a/src/internal/kopia/data_collection.go b/src/internal/kopia/data_collection.go
index 6c5b1696c..fed1bc002 100644
--- a/src/internal/kopia/data_collection.go
+++ b/src/internal/kopia/data_collection.go
@@ -38,6 +38,7 @@ func (kdc kopiaDataCollection) FullPath() path.Path {
 type kopiaDataStream struct {
 	reader io.ReadCloser
 	uuid   string
+	size   int64
 }
 
 func (kds kopiaDataStream) ToReader() io.ReadCloser {
@@ -47,3 +48,7 @@ func (kds kopiaDataStream) ToReader() io.ReadCloser {
 func (kds kopiaDataStream) UUID() string {
 	return kds.uuid
 }
+
+func (kds kopiaDataStream) Size() int64 {
+	return kds.size
+}
diff --git a/src/internal/kopia/data_collection_test.go b/src/internal/kopia/data_collection_test.go
index 97293d2a4..82788ee3f 100644
--- a/src/internal/kopia/data_collection_test.go
+++ b/src/internal/kopia/data_collection_test.go
@@ -66,6 +66,7 @@ func (suite *KopiaDataCollectionUnitSuite) TestReturnsStreams() {
 				&kopiaDataStream{
 					reader: io.NopCloser(bytes.NewReader(testData[0])),
 					uuid:   uuids[0],
+					size:   int64(len(testData[0])),
 				},
 			},
 		},
@@ -75,10 +76,12 @@ func (suite *KopiaDataCollectionUnitSuite) TestReturnsStreams() {
 				&kopiaDataStream{
 					reader: io.NopCloser(bytes.NewReader(testData[0])),
 					uuid:   uuids[0],
+					size:   int64(len(testData[0])),
 				},
 				&kopiaDataStream{
 					reader: io.NopCloser(bytes.NewReader(testData[1])),
 					uuid:   uuids[1],
+					size:   int64(len(testData[1])),
 				},
 			},
 		},
@@ -100,6 +103,9 @@ func (suite *KopiaDataCollectionUnitSuite) TestReturnsStreams() {
 				buf, err := ioutil.ReadAll(returnedStream.ToReader())
 				require.NoError(t, err)
 				assert.Equal(t, buf, testData[count])
+				require.Implements(t, (*data.StreamSize)(nil), returnedStream)
+				ss := returnedStream.(data.StreamSize)
+				assert.Equal(t, len(buf), int(ss.Size()))
 
 				count++
 			}
diff --git a/src/internal/kopia/wrapper.go b/src/internal/kopia/wrapper.go
index 0b39f1960..894e07809 100644
--- a/src/internal/kopia/wrapper.go
+++ b/src/internal/kopia/wrapper.go
@@ -564,6 +564,7 @@ func restoreSingleItem(
 			&kopiaDataStream{
 				uuid:   f.Name(),
 				reader: r,
+				size:   f.Size(),
 			},
 		},
 		path: itemDir,