From abe6a29a0e82d2739ecd295f387aff27e9eee7f7 Mon Sep 17 00:00:00 2001 From: Eugene Date: Wed, 13 Nov 2024 15:43:24 +0200 Subject: [PATCH] implements JSON data export support (#552) refs: #546 --- pgmanage/app/include/Spartacus/Utils.py | 122 +++++++++++------- .../src/components/QueryTab.vue | 1 + pgmanage/app/views/polling.py | 12 +- 3 files changed, 82 insertions(+), 53 deletions(-) diff --git a/pgmanage/app/include/Spartacus/Utils.py b/pgmanage/app/include/Spartacus/Utils.py index 83e08d753..685ab45e2 100644 --- a/pgmanage/app/include/Spartacus/Utils.py +++ b/pgmanage/app/include/Spartacus/Utils.py @@ -27,7 +27,8 @@ import openpyxl from collections import OrderedDict import tempfile - +import json +from django.core.serializers.json import DjangoJSONEncoder import app.include.Spartacus as Spartacus class Exception(Exception): @@ -35,76 +36,103 @@ class Exception(Exception): class DataFileWriter(object): - def __init__(self, p_filename, p_fieldnames=None, p_encoding='utf-8', p_delimiter=';', p_lineterminator='\n', skip_headers=False): - v_tmp = p_filename.split('.') - if len(v_tmp) > 1: - self.v_extension = v_tmp[-1].lower() + def __init__(self, filename, fieldnames=None, encoding='utf-8', delimiter=';', lineterminator='\n', skip_headers=False): + tmp = filename.split('.') + if len(tmp) > 1: + self.extension = tmp[-1].lower() else: - self.v_extension = 'csv' - if self.v_extension == 'txt' or self.v_extension == 'out': - self.v_extension = 'csv' - self.v_filename = p_filename - self.v_file = None - self.v_header = p_fieldnames # Can't be empty for CSV - self.v_encoding = p_encoding - self.v_delimiter = p_delimiter - self.v_lineterminator = p_lineterminator - self.v_currentrow = 1 - self.v_open = False + self.extension = 'csv' + if self.extension == 'txt' or self.extension == 'out': + self.extension = 'csv' + + self.filename = filename + self.file_handle = None + self.header = [] + self.currentrow = 1 + self.opened = False + self.writer = None + self.encoding = encoding + self.delimiter = delimiter + self.lineterminator = lineterminator self.skip_headers = skip_headers + for idx, field in enumerate(fieldnames): + if field == '?column?': + self.header.append(f'?column-{idx}') + else: + self.header.append(field) + def Open(self): try: - if self.v_extension == 'csv': - self.v_file = open(self.v_filename, 'w', encoding=self.v_encoding) - self.v_object = csv.writer(self.v_file, delimiter=self.v_delimiter, lineterminator=self.v_lineterminator) + if self.extension == 'csv': + self.file_handle = open(self.filename, 'w', encoding=self.encoding) + self.writer = csv.writer(self.file_handle, delimiter=self.delimiter, lineterminator=self.lineterminator) if not self.skip_headers: - self.v_object.writerow(self.v_header) - self.v_open = True - elif self.v_extension == 'xlsx': - self.v_object = openpyxl.Workbook(write_only=True) - self.v_open = True + self.writer.writerow(self.header) + self.opened = True + elif self.extension == 'xlsx': + self.writer = openpyxl.Workbook(write_only=True) + self.opened = True + elif self.extension == 'json': + self.file_handle = open(self.filename, 'w+', encoding=self.encoding) + print("[", file=self.file_handle, end='') + self.opened = True else: - raise Spartacus.Utils.Exception('File extension "{0}" not supported.'.format(self.v_extension)) + raise Spartacus.Utils.Exception('File extension "{0}" not supported.'.format(self.extension)) except Spartacus.Utils.Exception as exc: raise exc except Exception as exc: raise Spartacus.Utils.Exception(str(exc)) - def Write(self, p_datatable, p_sheetname=None): + def Write(self, p_datatable, p_hasmore=False): try: - if not self.v_open: + if not self.opened: raise Spartacus.Utils.Exception('You need to call Open() first.') - if self.v_extension == 'csv': - for v_row in p_datatable.Rows: - self.v_object.writerow(v_row) - else: - if self.v_currentrow == 1: - if p_sheetname: - v_worksheet = self.v_object.create_sheet(p_sheetname) - else: - v_worksheet = self.v_object.create_sheet() + if self.extension == 'csv': + for row in p_datatable.Rows: + self.writer.writerow(row) + elif self.extension == 'xlsx': + if self.currentrow == 1: + worksheet = self.writer.create_sheet() if not self.skip_headers: - v_worksheet.append(p_datatable.Columns) - self.v_currentrow = self.v_currentrow + 1 + worksheet.append(p_datatable.Columns) + self.currentrow = self.currentrow + 1 else: - v_worksheet = self.v_object.active + worksheet = self.writer.active for r in range(0, len(p_datatable.Rows)): - v_row = [] + row = [] for c in range(0, len(p_datatable.Columns)): - v_row.append(p_datatable.Rows[r][c]) - v_worksheet.append(v_row) - self.v_currentrow = self.v_currentrow + len(p_datatable.Rows) + row.append(p_datatable.Rows[r][c]) + worksheet.append(row) + self.currentrow = self.currentrow + len(p_datatable.Rows) + else: + for row in p_datatable.Rows: + print( + json.dumps( + dict(zip (self.header, row)), cls=DjangoJSONEncoder), + end=',\n', + file=self.file_handle + ) + except Spartacus.Utils.Exception as exc: raise exc except Exception as exc: raise Spartacus.Utils.Exception(str(exc)) def Flush(self): try: - if not self.v_open: + if not self.opened: raise Spartacus.Utils.Exception('You need to call Open() first.') - if self.v_extension == 'csv': - self.v_file.close() + if self.extension == 'csv': + self.file_handle.close() + elif self.extension == 'xlsx': + self.writer.save(self.filename) else: - self.v_object.save(self.v_filename) + pos = self.file_handle.tell() + # non-empty json file + if pos > 2: + # rewind 2 bytes back so the last ',\n' gets truncated + self.file_handle.seek(pos-2) + print("]", file=self.file_handle) + self.file_handle.close() + except Spartacus.Utils.Exception as exc: raise exc except Exception as exc: diff --git a/pgmanage/app/static/pgmanage_frontend/src/components/QueryTab.vue b/pgmanage/app/static/pgmanage_frontend/src/components/QueryTab.vue index c3584b6b6..ed3cf1927 100644 --- a/pgmanage/app/static/pgmanage_frontend/src/components/QueryTab.vue +++ b/pgmanage/app/static/pgmanage_frontend/src/components/QueryTab.vue @@ -185,6 +185,7 @@ export default { "csv-no_headers": "CSV(no headers)", xlsx: "XLSX", "xlsx-no_headers": "XLSX(no headers)", + json: "JSON" }, exportType: "csv", showFetchButtons: false, diff --git a/pgmanage/app/views/polling.py b/pgmanage/app/views/polling.py index 91aa0845e..a0ec5c0d9 100644 --- a/pgmanage/app/views/polling.py +++ b/pgmanage/app/views/polling.py @@ -153,15 +153,13 @@ def export_data( # cleaning temp folder clean_temp_folder() - if len(cmd_type.split("-")) == 2: - cmd_type = cmd_type.split("-")[0] - skip_headers = True - extension: str - if cmd_type == "export_csv": + if 'csv' in cmd_type: extension = "csv" - else: + elif 'xlsx' in cmd_type: extension = "xlsx" + else: + extension = "json" export_dir: str = settings.TEMP_DIR @@ -175,6 +173,7 @@ def export_data( data = database.v_connection.QueryBlock(sql_cmd, 1000, False, True) file_path: str = os.path.join(export_dir, file_name) + skip_headers = cmd_type in ['export_xlsx-no_headers', 'export_csv-no_headers'] file = Utils.DataFileWriter( file_path, data.Columns, encoding, delimiter, skip_headers=skip_headers ) @@ -1013,6 +1012,7 @@ def thread_query(self, args) -> None: "export_xlsx", "export_csv-no_headers", "export_xlsx-no_headers", + "export_json" ]: file_name, extension = export_data( sql_cmd=sql_cmd,