Limit size of file to 128 MB

pkg/unstructured/source.go

@@ -35,6 +35,8 @@ import (
 	"github.com/redhat-data-and-ai/unstructured-data-controller/pkg/gdrive"
 )
 
+const maxFileSize int64 = 128 << 20 // 128 MB — Snowflake external stage limit
+
 type DataSource interface {
 	// SyncFilesToFilestore will store all files from the source to the filestore and return the list of file paths
 	SyncFilesToFilestore(ctx context.Context, fs *filestore.FileStore) ([]RawFileMetadata, error)
@@ -65,6 +67,11 @@ func (s *S3BucketSource) SyncFilesToFilestore(ctx context.Context, fs *filestore
 		if strings.HasSuffix(*object.Key, "/") {
 			continue
 		}
+		if object.Size != nil && *object.Size > maxFileSize {
+			logger.Info("WARNING: skipping file exceeding max file size limit",
+				"key", *object.Key, "sizeMB", *object.Size/(1<<20))
+			continue
+		}
 		file := RawFileMetadata{
 			FilePath: s.filestorePath(*object.Key),
 			UID:      *object.ETag,
@@ -287,6 +294,12 @@ func (g *GDriveSource) SyncFilesToFilestore(ctx context.Context, fs *filestore.F
 			if record.MimeType == "application/vnd.google-apps.folder" {
 				continue
 			}
+			if record.FileSize > 0 && record.FileSize > maxFileSize {
+				logger.Info("WARNING: skipping file exceeding max file size limit",
+					"fileID", record.FileID, "fileName", record.FileName,
+					"sizeMB", record.FileSize/(1<<20))
+				continue
+			}
 			if !seen[record.FileID] {
 				seen[record.FileID] = true
 				fileRecords = append(fileRecords, record)
@@ -433,11 +446,17 @@ func (g *GDriveSource) storeFile(
 		}
 		defer func() { _ = reader.Close() }()
 
-		data, err := io.ReadAll(reader)
+		data, err := io.ReadAll(io.LimitReader(reader, maxFileSize+1))
 		if err != nil {
 			return false, fmt.Errorf("failed to read file %s: %w", fileID, err)
 		}
 
+		if int64(len(data)) > maxFileSize {
+			logger.Info("WARNING: skipping file exceeding max file size limit",
+				"fileID", fileID, "sizeMB", len(data)/(1<<20))
+			return false, nil
+		}
+
 		if err := fs.Store(ctx, filePath, data); err != nil {
 			return false, fmt.Errorf("failed to store file %s: %w", fileID, err)
 		}