Added default no limit on storage

Signed-off-by: Yuvraj <[email protected]>
flyteorg · Jul 16, 2021 · 46dd797 · 46dd797
1 parent 92b25c6
commit 46dd797
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 2 deletions.
diff --git a/storage/config.go b/storage/config.go
@@ -33,7 +33,7 @@ var (
 	defaultConfig = &Config{
 		Type: TypeS3,
 		Limits: LimitsConfig{
-			GetLimitMegabytes: 2,
+			GetLimitMegabytes: 0,
 		},
 		Connection: ConnectionConfig{
 			Region:   "us-east-1",

diff --git a/storage/stow_store.go b/storage/stow_store.go
@@ -230,7 +230,7 @@ func (s *StowStore) ReadRaw(ctx context.Context, reference DataReference) (io.Re
 		return nil, err
 	}
 
-	if sizeMbs := sizeBytes / MiB; sizeMbs > GetConfig().Limits.GetLimitMegabytes {
+	if sizeMbs := sizeBytes / MiB; sizeMbs > GetConfig().Limits.GetLimitMegabytes && GetConfig().Limits.GetLimitMegabytes != 0 {
 		return nil, errors.Errorf(ErrExceedsLimit, "limit exceeded. %vmb > %vmb.", sizeMbs, GetConfig().Limits.GetLimitMegabytes)
 	}
 

diff --git a/storage/stow_store_test.go b/storage/stow_store_test.go
@@ -181,6 +181,8 @@ func TestStowStore_ReadRaw(t *testing.T) {
 	t.Run("Exceeds limit", func(t *testing.T) {
 		testScope := promutils.NewTestScope()
 		fn := fQNFn["s3"]
+		GetConfig().Limits.GetLimitMegabytes = 2
+
 		s, err := NewStowRawStore(fn(container), &mockStowLoc{
 			ContainerCb: func(id string) (stow.Container, error) {
 				if id == container {
@@ -207,6 +209,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"]