2323logger = logging .getLogger (__name__ )
2424
2525
26- def hash_file (infile , hasher = None , blocksize = 65536 ):
26+ def hash_file (infile , hasher = None , blocksize = 8388608 ): # 8MB buffer (8 * 1024 * 1024)
2727 """
28- Incrementally generate file hashes (suitable for large files).
28+ Incrementally generate file hashes (suitable for very large files).
29+ Uses a large buffer size (8MB) for better performance with multi-GB files.
2930
3031 @param infile a python file-like object
3132 @param hasher a hasher function from hashlib library (e.g. md5, sha256, ...)
32- @param blocksize an integer of bytes to read into the hash at a time
33+ @param blocksize an integer of bytes to read into the hash at a time (default: 8MB)
3334 @return a hexdigest of the hash based on the hasher
3435 """
3536 if hasher is None :
3637 hasher = md5 ()
3738
39+ # Save current file position
40+ current_pos = infile .tell ()
41+
42+ # Go to beginning of file
43+ infile .seek (0 )
44+
45+ # Read and hash in chunks
3846 buf = infile .read (blocksize )
39-
4047 while len (buf ) > 0 :
4148 hasher .update (buf )
4249 buf = infile .read (blocksize )
4350
44- # Reset file buffer back to starting position
45- infile .seek (0 )
51+ # Reset file buffer back to original position
52+ infile .seek (current_pos )
4653
4754 return hasher .hexdigest ()
4855
@@ -67,31 +74,19 @@ def format_byte_size(n_bytes):
6774def decompress_file (infile , upload_filename ):
6875 """
6976 Opens gzip and bz2 files and returns the uncompressed data.
70- Uncompresses in pieces andand cuts off when limit is reached in order to
71- avoid 'zip bombs'.
72-
73- Solution suggested by Mark Adler:
74- http://stackoverflow.com/questions/13622706/how-to-protect-myself-from-a-gzip-or-bzip2-bomb
75-
76- The gist is that we decompress the file in chunks, counting the uncompressed
77- size of each and adding it to the total. If the filesize is OK, then
78- decompress the file and return data. This seems a little redundant, but
79- uncompressing the actual file in chunks and storing the data creates
80- crashing issues in the Docker container.
81-
82- TODO: Bz2 doesn't have the same ability as zlib to read in chunks, so a bomb is
83- just going to cause a memory error. So, given the extra time it takes to
84- decompress a bz2, I'm just assembling the chunks on the fly instead of
85- unzipping a second time. Investigate a workaround for this at some point.
77+ Optimized for handling large files (4GB+) efficiently.
78+ Uncompresses in pieces and cuts off when limit is reached to avoid 'zip bombs'.
8679
8780 @param infile file object containing the data to decompress
8881 @param upload_filename name of the file to decompress
8982 @return a temporary file object containing the decompressed data
9083 """
9184 app .logger .info ("Decompressing file %s" , upload_filename )
9285
93- decompressed_file = tempfile .NamedTemporaryFile ()
94- extension = os .path .splitext (upload_filename )[- 1 ]
86+ # Use the configured temporary directory for large file processing
87+ temp_dir = app .config .get ('TEMP_FILE_DIR' , tempfile .gettempdir ())
88+ decompressed_file = tempfile .NamedTemporaryFile (dir = temp_dir )
89+ extension = os .path .splitext (upload_filename )[- 1 ].lower ()
9590
9691 # Determine the correct open function based on the type of decompression
9792 if extension == ".gz" :
@@ -101,67 +96,107 @@ def decompress_file(infile, upload_filename):
10196 else :
10297 raise ValueError (f"Unknown file extension ({ extension } )
10398
99+ # Use a larger buffer size for better performance with large files
100+ buffer_size = 8 * 1024 * 1024 # 8MB buffer
101+
104102 try :
105103 data_length = 0
106104 reader = open_fn (infile )
107105
106+ # Stream decompression with a larger buffer
108107 while data_length < app .config ['MAX_CONTENT_LENGTH' ]:
109- buf = reader .read (1024 )
108+ buf = reader .read (buffer_size )
110109
111110 if len (buf ) == 0 :
112111 break
113112
114113 data_length += len (buf )
115114 decompressed_file .write (buf )
116115 else :
116+ # If we exit the loop without breaking, the file is too large
117117 decompressed_file .close ()
118118 return abort (
119119 400 , f"The decompressed file size is too large. "
120120 f"Max decompressed file size is: { format_byte_size (app .config ['MAX_CONTENT_LENGTH' ])}
121121 f"Filename: { upload_filename }
122122 )
123+
124+ # Close the reader to free resources
125+ reader .close ()
126+
123127 except (BadGzipFile , OSError , ValueError , TypeError , IOError , EOFError ) as err :
124128 decompressed_file .close ()
125129 raise ValueError (
126130 f'Failed to decompress { extension } { upload_filename } { str (err )}
127131 )
128132 except MemoryError :
129- # Noticed that some bzip bombs cause memory errors. There doesn't
130- # seem to be a ton of great ways around this.
133+ # Handle memory errors from bzip bombs
131134 decompressed_file .close ()
132135 return abort (
133- 400 , f"The decompressed file size is too large. "
136+ 400 , f"The decompressed file size is too large or caused a memory error . "
134137 f"Max decompressed file size is: { format_byte_size (app .config ['MAX_CONTENT_LENGTH' ])}
135138 f"Filename: { upload_filename }
136139 )
137140
138141 # Roll file pointer back to beginning of buffer now that decompression is complete
139142 decompressed_file .seek (0 )
143+ app .logger .info (f"Successfully decompressed { upload_filename } { format_byte_size (data_length )} )
140144
141145 return decompressed_file
142146
143147
144148def get_dataset_from_file (uploaded_file ):
145149 """
146150 Derives a netcdf4.Dataset object from the provided file upload dictionary.
151+ Optimized for handling large files (4GB+) efficiently.
147152
148- @param uploaded_file open file handle to the data to convert.
153+ @param uploaded_file open file handle or path to the data to convert.
149154 """
150155 app .logger .info ("Attempting to get dataset from uploaded file %s" , uploaded_file )
151-
152- datafile_name = uploaded_file .filename
156+
157+ # Handle both file paths and file objects
158+ if isinstance (uploaded_file , str ):
159+ # It's a file path
160+ datafile_name = os .path .basename (uploaded_file )
161+ is_path = True
162+ else :
163+ # It's a file object
164+ datafile_name = uploaded_file .filename
165+ is_path = False
166+
153167 check_valid_filename (datafile_name )
154-
155- # uploaded_file is an instance of werkzeug.FileStorage, which only allows us
156- # to read the file contents but not reference a location on disk.
157- # Copy the uploaded file to a named temporary file, so we can provide a disk
158- # location when initializing the netCDF.Dataset object below.
159-
160- # Since the CF suite checks the Dataset's filename for compliance,
161- # we need to make sure the original is incorporated into the name of the
162- # temporary file used to initialize the Dataset object below.
163- datafile = tempfile .NamedTemporaryFile (suffix = f"_{ datafile_name } )
164- datafile .write (uploaded_file .read ())
168+
169+ # For large files, we want to avoid loading the entire file into memory
170+ # Instead, we'll use a temporary file and stream the data in chunks
171+
172+ # Use the configured temporary directory for large file processing
173+ temp_dir = app .config .get ('TEMP_FILE_DIR' , tempfile .gettempdir ())
174+
175+ # Create a named temporary file in the specified directory
176+ datafile = tempfile .NamedTemporaryFile (dir = temp_dir , suffix = f"_{ datafile_name } )
177+
178+ # Stream the file in chunks to avoid memory issues
179+ if is_path :
180+ # If it's a path, open the file and copy it
181+ with open (uploaded_file , 'rb' ) as src_file :
182+ # Use a larger buffer size (8MB) for faster copying of large files
183+ buffer_size = 8 * 1024 * 1024 # 8MB buffer
184+ while True :
185+ buffer = src_file .read (buffer_size )
186+ if not buffer :
187+ break
188+ datafile .write (buffer )
189+ else :
190+ # If it's a file object, stream from it
191+ # Use a larger buffer size (8MB) for faster copying of large files
192+ buffer_size = 8 * 1024 * 1024 # 8MB buffer
193+ while True :
194+ buffer = uploaded_file .read (buffer_size )
195+ if not buffer :
196+ break
197+ datafile .write (buffer )
198+
199+ # Reset file pointer to beginning
165200 datafile .seek (0 )
166201
167202 # Calculate hash and file size now, prior to any potential file decompression
0 commit comments