Added default no limit on storage (#93)

* Added default no limit on storage Signed-off-by: Yuvraj <[email protected]>
flyteorg · Jul 26, 2021 · 9f637fd · 9f637fd
1 parent 101aa2a
commit 9f637fd
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 4 deletions.
diff --git a/storage/stow_store.go b/storage/stow_store.go
@@ -7,6 +7,8 @@ import (
 	"sync"
 	"time"
 
+	"github.com/flyteorg/flytestdlib/errors"
+
 	"github.com/aws/aws-sdk-go/aws/awserr"
 	s32 "github.com/aws/aws-sdk-go/service/s3"
 	"github.com/graymeta/stow/azure"
@@ -20,8 +22,6 @@ import (
 	"github.com/flyteorg/flytestdlib/logger"
 	"github.com/flyteorg/flytestdlib/promutils/labeled"
 
-	"github.com/flyteorg/flytestdlib/errors"
-
 	"github.com/flyteorg/flytestdlib/promutils"
 
 	"github.com/graymeta/stow"
@@ -230,8 +230,10 @@ func (s *StowStore) ReadRaw(ctx context.Context, reference DataReference) (io.Re
 		return nil, err
 	}
 
-	if sizeMbs := sizeBytes / MiB; sizeMbs > GetConfig().Limits.GetLimitMegabytes {
-		return nil, errors.Errorf(ErrExceedsLimit, "limit exceeded. %vmb > %vmb.", sizeMbs, GetConfig().Limits.GetLimitMegabytes)
+	if GetConfig().Limits.GetLimitMegabytes != 0 {
+		if sizeMbs := sizeBytes / MiB; sizeMbs > GetConfig().Limits.GetLimitMegabytes {
+			return nil, errors.Errorf(ErrExceedsLimit, "limit exceeded. %vmb > %vmb.", sizeMbs, GetConfig().Limits.GetLimitMegabytes)
+		}
 	}
 
 	return item.Open()

diff --git a/storage/stow_store_test.go b/storage/stow_store_test.go
@@ -181,6 +181,7 @@ func TestStowStore_ReadRaw(t *testing.T) {
 	t.Run("Exceeds limit", func(t *testing.T) {
 		testScope := promutils.NewTestScope()
 		fn := fQNFn["s3"]
+
 		s, err := NewStowRawStore(fn(container), &mockStowLoc{
 			ContainerCb: func(id string) (stow.Container, error) {
 				if id == container {
@@ -207,6 +208,35 @@ func TestStowStore_ReadRaw(t *testing.T) {
 		assert.NotNil(t, errors.Cause(err))
 	})
 
+	t.Run("No Limit", func(t *testing.T) {
+		testScope := promutils.NewTestScope()
+		fn := fQNFn["s3"]
+		GetConfig().Limits.GetLimitMegabytes = 0
+
+		s, err := NewStowRawStore(fn(container), &mockStowLoc{
+			ContainerCb: func(id string) (stow.Container, error) {
+				if id == container {
+					return newMockStowContainer(container), nil
+				}
+				return nil, fmt.Errorf("container is not supported")
+			},
+			CreateContainerCb: func(name string) (stow.Container, error) {
+				if name == container {
+					return newMockStowContainer(container), nil
+				}
+				return nil, fmt.Errorf("container is not supported")
+			},
+		}, false, testScope)
+		assert.NoError(t, err)
+		err = s.WriteRaw(context.TODO(), DataReference("s3://container/path"), 3*MiB, Options{}, bytes.NewReader([]byte{}))
+		assert.NoError(t, err)
+		metadata, err := s.Head(context.TODO(), DataReference("s3://container/path"))
+		assert.NoError(t, err)
+		assert.True(t, metadata.Exists())
+		_, err = s.ReadRaw(context.TODO(), DataReference("s3://container/path"))
+		assert.Nil(t, err)
+	})
+
 	t.Run("Happy Path multi-container enabled", func(t *testing.T) {
 		testScope := promutils.NewTestScope()
 		fn := fQNFn["s3"]