Fix read raw limit (#4370)

Signed-off-by: Hongxin Liang <[email protected]>
flyteorg · Nov 6, 2023 · 939eb8b · 939eb8b
1 parent eccf993
commit 939eb8b
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/flytestdlib/storage/stow_store.go b/flytestdlib/storage/stow_store.go
@@ -265,8 +265,8 @@ func (s *StowStore) ReadRaw(ctx context.Context, reference DataReference) (io.Re
 	}
 
 	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)
+		if sizeBytes > GetConfig().Limits.GetLimitMegabytes*MiB {
+			return nil, errors.Errorf(ErrExceedsLimit, "limit exceeded. %.6fmb > %vmb.", float64(sizeBytes)/float64(MiB), GetConfig().Limits.GetLimitMegabytes)
 		}
 	}
 

diff --git a/flytestdlib/storage/stow_store_test.go b/flytestdlib/storage/stow_store_test.go
@@ -274,7 +274,7 @@ func TestStowStore_ReadRaw(t *testing.T) {
 			},
 		}, nil, false, metrics)
 		assert.NoError(t, err)
-		dataReference := writeTestFileWithSize(ctx, t, s, "s3://container/path", 3*MiB)
+		dataReference := writeTestFileWithSize(ctx, t, s, "s3://container/path", 2*MiB+1)
 		_, err = s.ReadRaw(ctx, dataReference)
 		assert.Error(t, err)
 		assert.True(t, IsExceedsLimit(err))