def build_final_export(self, export_results):
base_name = safe_filename(self.export_instance.name or 'Export')
final_zip = self._get_zipfile_for_final_archive()
with final_zip:
pages = len(export_results)
for result in export_results:
if not result.success:
logger.error(
' Error in page %s so not added to final output',
result.page)
if os.path.exists(result.path):
raw_dump_path = result.path
logger.info(
' Adding raw dump of page %s to final output',
result.page)
destination = '{}/page_{}.json.gz'.format(
UNPROCESSED_PAGES_DIR, result.page)
final_zip.write(raw_dump_path, destination,
zipfile.ZIP_STORED)
os.remove(raw_dump_path)
continue
logger.info(' Adding page {} of {} to final file'.format(
result.page, pages))
if self.is_zip:
_add_compressed_page_to_zip(final_zip, result.page,
result.path)
else:
final_zip.write(result.path,
'{}_{}'.format(base_name, result.page))
return final_zip.filename
def compile_final_zip(self, error_pages, export_archive_path,
export_instance, successful_pages):
final_dir, orig_name = os.path.split(export_archive_path)
if not error_pages:
fd, final_path = tempfile.mkstemp()
else:
final_name = 'INCOMPLETE_{}_{}.zip'.format(
orig_name,
datetime.utcnow().isoformat())
final_path = os.path.join(final_dir, final_name)
print('Recompiling export')
export_name = safe_filename(export_instance.name or 'Export')
with zipfile.ZipFile(final_path,
mode='w',
compression=zipfile.ZIP_DEFLATED,
allowZip64=True) as final_zip:
for result in successful_pages:
print(' Adding page {} to final file'.format(result.page))
_add_compressed_page_to_zip(final_zip, result.page,
result.path)
print(
' Adding original export pages and unprocessed pages final file'
)
def _include_member(member):
# add original export pages and any raw data that we weren't able to process
add = member.startswith(export_name) or member in error_pages
if add:
print(' {}'.format(member))
return add
_copy_files_from_zip_to_zip(final_zip, export_archive_path,
_include_member)
return final_path
def open(self, export_instances):
"""
Open the _Writer for writing. This must be called before using _Writer.write().
Note that this function returns a context manager!
A _Writer can only be opened once.
"""
# Create and open a temp file
assert self._path is None
if len(export_instances) == 1:
name = export_instances[0].name or ''
else:
name = ''
name = safe_filename(name)
fd, self._path = tempfile.mkstemp()
with os.fdopen(fd, 'wb') as file:
# open the ExportWriter
headers = []
table_titles = {}
for instance in export_instances:
headers += [(t, (t.get_headers(
split_columns=instance.split_multiselects), ))
for t in instance.selected_tables]
table_titles.update(
{t: t.label
for t in instance.selected_tables})
self.writer.open(headers,
file,
table_titles=table_titles,
archive_basepath=name)
yield
self.writer.close()
def compile_final_zip(self, error_pages, export_archive_path, export_instance, successful_pages):
final_dir, orig_name = os.path.split(export_archive_path)
if not error_pages:
fd, final_path = tempfile.mkstemp()
else:
final_name = 'INCOMPLETE_{}_{}.zip'.format(orig_name, datetime.utcnow().isoformat())
final_path = os.path.join(final_dir, final_name)
print('Recompiling export')
export_name = safe_filename(export_instance.name or 'Export')
with zipfile.ZipFile(final_path, mode='w', compression=zipfile.ZIP_DEFLATED, allowZip64=True) as final_zip:
for result in successful_pages:
print(' Adding page {} to final file'.format(result.page))
_add_compressed_page_to_zip(final_zip, result.page, result.path)
print(' Adding original export pages and unprocessed pages final file')
def _include_member(member):
# add original export pages and any raw data that we weren't able to process
add = member.startswith(export_name) or member in error_pages
if add:
print(' {}'.format(member))
return add
_copy_files_from_zip_to_zip(final_zip, export_archive_path, _include_member)
return final_path
def _download_export(self, export_instance):
export_archive_path = '{}_{}.zip'.format(
safe_filename(export_instance.name or 'Export'),
datetime.utcnow().isoformat())
payload = export_instance.get_payload(stream=True)
with open(export_archive_path, 'w') as download:
shutil.copyfileobj(payload, download)
return export_archive_path
def _download_export(self, export_instance):
export_archive_path = '{}_{}.zip'.format(
safe_filename(export_instance.name or 'Export'),
datetime.utcnow().isoformat()
)
payload = export_instance.get_payload(stream=True)
with open(export_archive_path, 'wb') as download:
shutil.copyfileobj(payload, download)
return export_archive_path
def download_export(export_instance, download_path=None):
if not download_path:
export_archive_path = '{}_{}.zip'.format(
safe_filename(export_instance.name.encode('ascii', 'replace') or 'Export'),
datetime.utcnow().isoformat()
)
download_path = os.path.join(settings.SHARED_DRIVE_ROOT, export_archive_path)
if not os.path.exists(download_path):
payload = export_instance.get_payload(stream=True)
with open(download_path, 'wb') as download:
shutil.copyfileobj(payload, download)
return download_path
def download_saved_export(export_id, dest_dir=None):
# Downloads the latest saved export to shared-directory
dest_dir = (dest_dir or settings.SHARED_DRIVE_ROOT).rstrip()
export_instance = get_properly_wrapped_export_instance(export_id)
export_archive_path = '{}/{}_{}.zip'.format(
dest_dir,
safe_filename(
export_instance.name.encode('ascii', 'replace') or 'Export'),
datetime.utcnow().isoformat())
payload = export_instance.get_payload(stream=True)
print("Downloading Export to {}".format(export_archive_path))
with open(export_archive_path, 'w') as download:
shutil.copyfileobj(payload, download)
print("Download Finished!")
def download_saved_export(export_id, dest_dir=None):
# Downloads the latest saved export to shared-directory
dest_dir = (dest_dir or settings.SHARED_DRIVE_ROOT).rstrip()
export_instance = get_properly_wrapped_export_instance(export_id)
export_archive_path = '{}/{}_{}.zip'.format(
dest_dir,
safe_filename(export_instance.name.encode('ascii', 'replace') or 'Export'),
datetime.utcnow().isoformat()
)
payload = export_instance.get_payload(stream=True)
print("Downloading Export to {}".format(export_archive_path))
with open(export_archive_path, 'wb') as download:
shutil.copyfileobj(payload, download)
print("Download Finished!")
def open(self, export_instances):
"""
Open the _Writer for writing. This must be called before using _Writer.write().
Note that this function returns a context manager!
A _Writer can only be opened once.
"""
# Create and open a temp file
assert self._path is None
if len(export_instances) == 1:
name = export_instances[0].name or ''
else:
name = ''
name = safe_filename(name)
all_sheet_names = Counter(table.label for instance in export_instances
for table in instance.selected_tables)
fd, self._path = tempfile.mkstemp()
with os.fdopen(fd, 'wb') as file:
# open the ExportWriter
headers = []
table_titles = {}
for instance_index, instance in enumerate(export_instances):
headers += [(t, (t.get_headers(
split_columns=instance.split_multiselects), ))
for t in instance.selected_tables]
for table_index, table in enumerate(instance.selected_tables):
sheet_name = table.label or "Sheet{}".format(table_index +
1)
# If it's a bulk export and the sheet has the same name as another sheet,
# Prefix the sheet name with the export name
if len(export_instances
) > 1 and all_sheet_names[sheet_name] > 1:
sheet_name = "{}-{}".format(
instance.name
or "Export{}".format(instance_index + 1),
sheet_name)
table_titles[table] = sheet_name
self.writer.open(headers,
file,
table_titles=table_titles,
archive_basepath=name)
try:
yield
finally:
self.writer.close()
def open(self, export_instances):
"""
Open the _Writer for writing. This must be called before using _Writer.write().
Note that this function returns a context manager!
A _Writer can only be opened once.
"""
if len(export_instances) == 1:
name = export_instances[0].name or ''
else:
name = ''
name = safe_filename(name)
all_sheet_names = Counter(
table.label for instance in export_instances for table in instance.selected_tables
)
with open(self.path, 'wb') as file:
# open the ExportWriter
headers = []
table_titles = {}
for instance_index, instance in enumerate(export_instances):
headers += [
(t, (t.get_headers(split_columns=instance.split_multiselects),))
for t in instance.selected_tables
]
for table_index, table in enumerate(instance.selected_tables):
sheet_name = table.label or "Sheet{}".format(table_index + 1)
# If it's a bulk export and the sheet has the same name as another sheet,
# Prefix the sheet name with the export name
if len(export_instances) > 1 and all_sheet_names[sheet_name] > 1:
sheet_name = "{}-{}".format(
instance.name or "Export{}".format(instance_index + 1),
sheet_name
)
table_titles[table] = sheet_name
self.writer.open(headers, file, table_titles=table_titles, archive_basepath=name)
try:
yield
finally:
self.writer.close()
def build_final_export(self, export_results):
prefix = '{}{}_final_'.format(TEMP_FILE_PREFIX,
self.export_instance.get_id)
final_file_obj = tempfile.NamedTemporaryFile(prefix=prefix,
mode='wb',
delete=False)
base_name = safe_filename(self.export_instance.name or 'Export')
with zipfile.ZipFile(final_file_obj,
mode='w',
compression=zipfile.ZIP_DEFLATED,
allowZip64=True) as final_zip:
pages = len(export_results)
for result in export_results:
if not result.success:
logger.error(
' Error in page %s so not added to final output',
result.page)
if os.path.exists(result.path):
raw_dump_path = result.path
logger.info(
' Adding raw dump of page %s to final output',
result.page)
destination = '{}/page_{}.json.gz'.format(
UNPROCESSED_PAGES_DIR, result.page)
final_zip.write(raw_dump_path, destination,
zipfile.ZIP_STORED)
os.remove(raw_dump_path)
continue
logger.info(' Adding page {} of {} to final file'.format(
result.page, pages))
if self.is_zip:
_add_compressed_page_to_zip(final_zip, result.page,
result.path)
else:
final_zip.write(result.path,
'{}_{}'.format(base_name, result.page))
return final_file_obj.name
def _get_name(self, export_instances):
if len(export_instances) == 1:
name = export_instances[0].name or 'Export'
else:
name = 'BulkExport'
return safe_filename(name)
def test_format_and_set_as_header(self, filename, expected_filename):
self.assertEqual(safe_filename(filename), expected_filename)
self.assertWorksAsHeader(safe_filename_header(filename))
def test_add_extension(self):
filename = safe_filename('test', 'zip')
self.assertEqual(filename, 'test.zip')
def _get_name(self, export_instances):
if len(export_instances) == 1:
name = export_instances[0].name or 'Export'
else:
name = 'BulkExport'
return safe_filename(name)
