From bccc73f1dc29179bcb7db9c80d7dc54aaa7d97d9 Mon Sep 17 00:00:00 2001 From: Raphael Michel Date: Wed, 3 Jul 2019 13:35:26 +0200 Subject: [PATCH] Optimized command-line exports --- src/pretix/base/exporter.py | 87 ++++++++++++------- src/pretix/base/exporters/invoices.py | 11 ++- src/pretix/base/management/commands/export.py | 58 +++++++++++++ 3 files changed, 120 insertions(+), 36 deletions(-) create mode 100644 src/pretix/base/management/commands/export.py diff --git a/src/pretix/base/exporter.py b/src/pretix/base/exporter.py index 64ed12eb78..e8d725e19c 100644 --- a/src/pretix/base/exporter.py +++ b/src/pretix/base/exporter.py @@ -71,6 +71,8 @@ class BaseExporter: :type form_data: dict :param form_data: The form data of the export details form + :param output_file: You can optionally accept a parameter that will be given a file handle to write the + output to. In this case, you can return None instead of the file content. Note: If you use a ``ModelChoiceField`` (or a ``ModelMultipleChoiceField``), the ``form_data`` will not contain the model instance but only it's primary key (or @@ -111,14 +113,20 @@ class ListExporter(BaseExporter): def get_filename(self): return 'export.csv' - def _render_csv(self, form_data, **kwargs): - output = io.StringIO() - writer = csv.writer(output, **kwargs) - for line in self.iterate_list(form_data): - writer.writerow(line) - return self.get_filename() + '.csv', 'text/csv', output.getvalue().encode("utf-8") + def _render_csv(self, form_data, output_file=None, **kwargs): + if output_file: + writer = csv.writer(output_file, **kwargs) + for line in self.iterate_list(form_data): + writer.writerow(line) + return self.get_filename() + '.csv', 'text/csv', None + else: + output = io.StringIO() + writer = csv.writer(output, **kwargs) + for line in self.iterate_list(form_data): + writer.writerow(line) + return self.get_filename() + '.csv', 'text/csv', output.getvalue().encode("utf-8") - def _render_xlsx(self, form_data): + def _render_xlsx(self, form_data, output_file=None): wb = Workbook() ws = wb.get_active_sheet() try: @@ -129,20 +137,24 @@ class ListExporter(BaseExporter): for j, val in enumerate(line): ws.cell(row=i + 1, column=j + 1).value = str(val) if not isinstance(val, KNOWN_TYPES) else val - with tempfile.NamedTemporaryFile(suffix='.xlsx') as f: - wb.save(f.name) - f.seek(0) - return self.get_filename() + '.xlsx', 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet', f.read() + if output_file: + wb.save(output_file) + return self.get_filename() + '.xlsx', 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet', None + else: + with tempfile.NamedTemporaryFile(suffix='.xlsx') as f: + wb.save(f.name) + f.seek(0) + return self.get_filename() + '.xlsx', 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet', f.read() - def render(self, form_data: dict) -> Tuple[str, str, bytes]: + def render(self, form_data: dict, output_file=None) -> Tuple[str, str, bytes]: if form_data.get('_format') == 'xlsx': - return self._render_xlsx(form_data) + return self._render_xlsx(form_data, output_file=output_file) elif form_data.get('_format') == 'default': - return self._render_csv(form_data, quoting=csv.QUOTE_NONNUMERIC, delimiter=',') + return self._render_csv(form_data, quoting=csv.QUOTE_NONNUMERIC, delimiter=',', output_file=output_file) elif form_data.get('_format') == 'csv-excel': - return self._render_csv(form_data, dialect='excel') + return self._render_csv(form_data, dialect='excel', output_file=output_file) elif form_data.get('_format') == 'semicolon': - return self._render_csv(form_data, dialect='excel', delimiter=';') + return self._render_csv(form_data, dialect='excel', delimiter=';', output_file=output_file) class MultiSheetListExporter(ListExporter): @@ -180,14 +192,20 @@ class MultiSheetListExporter(ListExporter): def iterate_sheet(self, form_data, sheet): raise NotImplementedError() # noqa - def _render_sheet_csv(self, form_data, sheet, **kwargs): - output = io.StringIO() - writer = csv.writer(output, **kwargs) - for line in self.iterate_sheet(form_data, sheet): - writer.writerow(line) - return self.get_filename() + '.csv', 'text/csv', output.getvalue().encode("utf-8") + def _render_sheet_csv(self, form_data, sheet, output_file=None, **kwargs): + if output_file: + writer = csv.writer(output_file, **kwargs) + for line in self.iterate_sheet(form_data, sheet): + writer.writerow(line) + return self.get_filename() + '.csv', 'text/csv', None + else: + output = io.StringIO() + writer = csv.writer(output, **kwargs) + for line in self.iterate_sheet(form_data, sheet): + writer.writerow(line) + return self.get_filename() + '.csv', 'text/csv', output.getvalue().encode("utf-8") - def _render_xlsx(self, form_data): + def _render_xlsx(self, form_data, output_file=None): wb = Workbook() ws = wb.get_active_sheet() wb.remove(ws) @@ -197,19 +215,24 @@ class MultiSheetListExporter(ListExporter): for j, val in enumerate(line): ws.cell(row=i + 1, column=j + 1).value = str(val) if not isinstance(val, KNOWN_TYPES) else val - with tempfile.NamedTemporaryFile(suffix='.xlsx') as f: - wb.save(f.name) - f.seek(0) - return self.get_filename() + '.xlsx', 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet', f.read() + if output_file: + wb.save(output_file) + return self.get_filename() + '.xlsx', 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet', None + else: + with tempfile.NamedTemporaryFile(suffix='.xlsx') as f: + wb.save(f.name) + f.seek(0) + return self.get_filename() + '.xlsx', 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet', f.read() - def render(self, form_data: dict) -> Tuple[str, str, bytes]: + def render(self, form_data: dict, output_file=None) -> Tuple[str, str, bytes]: if form_data.get('_format') == 'xlsx': - return self._render_xlsx(form_data) + return self._render_xlsx(form_data, output_file=output_file) elif ':' in form_data.get('_format'): sheet, f = form_data.get('_format').split(':') if f == 'default': - return self._render_sheet_csv(form_data, sheet, quoting=csv.QUOTE_NONNUMERIC, delimiter=',') + return self._render_sheet_csv(form_data, sheet, quoting=csv.QUOTE_NONNUMERIC, delimiter=',', + output_file=output_file) elif f == 'excel': - return self._render_sheet_csv(form_data, sheet, dialect='excel') + return self._render_sheet_csv(form_data, sheet, dialect='excel', output_file=output_file) elif f == 'semicolon': - return self._render_sheet_csv(form_data, sheet, dialect='excel', delimiter=';') + return self._render_sheet_csv(form_data, sheet, dialect='excel', delimiter=';', output_file=output_file) diff --git a/src/pretix/base/exporters/invoices.py b/src/pretix/base/exporters/invoices.py index 08ee71e083..31b242bb95 100644 --- a/src/pretix/base/exporters/invoices.py +++ b/src/pretix/base/exporters/invoices.py @@ -20,7 +20,7 @@ class InvoiceExporter(BaseExporter): identifier = 'invoices' verbose_name = _('All invoices') - def render(self, form_data: dict): + def render(self, form_data: dict, output_file=None): qs = self.event.invoices.filter(shredded=False) if form_data.get('payment_provider'): @@ -47,7 +47,7 @@ class InvoiceExporter(BaseExporter): with tempfile.TemporaryDirectory() as d: any = False - with ZipFile(os.path.join(d, 'tmp.zip'), 'w') as zipf: + with ZipFile(output_file or os.path.join(d, 'tmp.zip'), 'w') as zipf: for i in qs: try: if not i.file: @@ -68,8 +68,11 @@ class InvoiceExporter(BaseExporter): if not any: return None - with open(os.path.join(d, 'tmp.zip'), 'rb') as zipf: - return '{}_invoices.zip'.format(self.event.slug), 'application/zip', zipf.read() + if output_file: + return '{}_invoices.zip'.format(self.event.slug), 'application/zip', None + else: + with open(os.path.join(d, 'tmp.zip'), 'rb') as zipf: + return '{}_invoices.zip'.format(self.event.slug), 'application/zip', zipf.read() @property def export_form_fields(self): diff --git a/src/pretix/base/management/commands/export.py b/src/pretix/base/management/commands/export.py new file mode 100644 index 0000000000..d1d481ad9d --- /dev/null +++ b/src/pretix/base/management/commands/export.py @@ -0,0 +1,58 @@ +import json +import sys + +from django.core.management.base import BaseCommand +from django.utils.timezone import override +from django_scopes import scope + +from pretix.base.i18n import language +from pretix.base.models import Event, Organizer +from pretix.base.signals import register_data_exporters + + +class Command(BaseCommand): + help = "Run an exporter to get data out of pretix" + + def add_arguments(self, parser): + parser.add_argument('organizer_slug', nargs=1, type=str) + parser.add_argument('event_slug', nargs=1, type=str) + parser.add_argument('export_provider', nargs=1, type=str) + parser.add_argument('output_file', nargs=1, type=str) + parser.add_argument('--parameters', action='store', type=str, help='JSON-formatted parameters') + + def handle(self, *args, **options): + try: + o = Organizer.objects.get(slug=options['organizer_slug'][0]) + except Organizer.DoesNotExist: + self.stderr.write(self.style.ERROR('Organizer not found.')) + sys.exit(1) + + with scope(organizer=o): + try: + e = o.events.get(slug=options['event_slug'][0]) + except Event.DoesNotExist: + self.stderr.write(self.style.ERROR('Event not found.')) + sys.exit(1) + + with language(e.settings.locale), override(e.settings.timezone): + responses = register_data_exporters.send(e) + for receiver, response in responses: + ex = response(e) + if ex.identifier == options['export_provider'][0]: + params = json.loads(options.get('parameters') or '{}') + with open(options['output_file'][0], 'wb') as f: + try: + ex.render(form_data=params, output_file=f) + except TypeError: + self.stderr.write(self.style.WARNING( + 'Provider does not support direct file writing, need to buffer export in memory.')) + d = ex.render(form_data=params) + if d is None: + self.stderr.write(self.style.ERROR('Empty export.')) + sys.exit(2) + f.write(d[2]) + + sys.exit(0) + + self.stderr.write(self.style.ERROR('Export provider not found.')) + sys.exit(1)