From 6b3feb30119a282bdd2a3ee5097b57b8b78f7b0b Mon Sep 17 00:00:00 2001
From: Shikhar Gupta <shikgupt>
Date: Tue, 21 Jul 2026 15:39:01 +0530
Subject: [PATCH 1/2] Limit size of file to 128 MB

- Snowflake does not support files greater than 128 MB in stage
---
 pkg/unstructured/source.go | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/pkg/unstructured/source.go b/pkg/unstructured/source.go
index 67f86d7d..6c8e450a 100644
--- a/pkg/unstructured/source.go
+++ b/pkg/unstructured/source.go
@@ -65,6 +65,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("skipping file exceeding 128 MB size limit",
+				"key", *object.Key, "sizeMB", *object.Size/(1<<20))
+			continue
+		}
 		file := RawFileMetadata{
 			FilePath: s.filestorePath(*object.Key),
 			UID:      *object.ETag,
@@ -229,6 +234,8 @@ func (s *S3BucketSource) s3Key(filestorePath string) string {
 	return path.Join(s.Prefix, baseName)
 }
 
+const maxFileSize int64 = 128 << 20 // 128 MB — Snowflake external stage limit
+
 // GDriveSource implements DataSource for Google Drive folders.
 type GDriveSource struct {
 	GDriveClient        *gdrive.Client
@@ -438,6 +445,12 @@ func (g *GDriveSource) storeFile(
 			return false, fmt.Errorf("failed to read file %s: %w", fileID, err)
 		}
 
+		if int64(len(data)) > maxFileSize {
+			logger.Info("skipping file exceeding 128 MB after export",
+				"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)
 		}

From eb02601af5686d5b7370712a5e4be318971124bb Mon Sep 17 00:00:00 2001
From: Shikhar Gupta <shikgupt>
Date: Tue, 21 Jul 2026 17:44:41 +0530
Subject: [PATCH 2/2] Address comments

---
 pkg/unstructured/source.go | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/pkg/unstructured/source.go b/pkg/unstructured/source.go
index 6c8e450a..2d31391e 100644
--- a/pkg/unstructured/source.go
+++ b/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)
@@ -66,7 +68,7 @@ func (s *S3BucketSource) SyncFilesToFilestore(ctx context.Context, fs *filestore
 			continue
 		}
 		if object.Size != nil && *object.Size > maxFileSize {
-			logger.Info("skipping file exceeding 128 MB size limit",
+			logger.Info("WARNING: skipping file exceeding max file size limit",
 				"key", *object.Key, "sizeMB", *object.Size/(1<<20))
 			continue
 		}
@@ -234,8 +236,6 @@ func (s *S3BucketSource) s3Key(filestorePath string) string {
 	return path.Join(s.Prefix, baseName)
 }
 
-const maxFileSize int64 = 128 << 20 // 128 MB — Snowflake external stage limit
-
 // GDriveSource implements DataSource for Google Drive folders.
 type GDriveSource struct {
 	GDriveClient        *gdrive.Client
@@ -294,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)
@@ -440,13 +446,13 @@ 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("skipping file exceeding 128 MB after export",
+			logger.Info("WARNING: skipping file exceeding max file size limit",
 				"fileID", fileID, "sizeMB", len(data)/(1<<20))
 			return false, nil
 		}