#include "ExportDataDialog.h"

#include "ui_ExportDataDialog.h"

#include "sqlitedb.h"

#include "Settings.h"

#include "sqlite.h"

#include "FileDialog.h"

#include <QFile>

#include <QTextStream>

#include <QMessageBox>

#include <QTextCodec>

#include <json.hpp>

using json = nlohmann::json;

ExportDataDialog::ExportDataDialog(DBBrowserDB& db, ExportFormats format, QWidget* parent, const QString& query, const sqlb::ObjectIdentifier& selection)

: QDialog(parent),

ui(new Ui::ExportDataDialog),

pdb(db),

m_format(format),

m_sQuery(query)

{

// Create UI

ui->setupUi(this);

// Show different option widgets depending on the export format

ui->stackFormat->setCurrentIndex(format);

if(format == ExportFormatJson) {

setWindowTitle(tr("Export data as JSON"));

}

// Retrieve the saved dialog preferences

ui->checkHeader->setChecked(Settings::getValue("exportcsv", "firstrowheader").toBool());

setSeparatorChar(Settings::getValue("exportcsv", "separator").toInt());

setQuoteChar(Settings::getValue("exportcsv", "quotecharacter").toInt());

setNewLineString(Settings::getValue("exportcsv", "newlinecharacters").toString());

ui->checkPrettyPrint->setChecked(Settings::getValue("exportjson", "prettyprint").toBool());

// Update the visible/hidden status of the "Other" line edit fields

showCustomCharEdits();

// If a SQL query was specified hide the table combo box. If not fill it with tables to export

if(query.isEmpty())

{

// Get list of tables to export

for(const auto& it : pdb.schemata)

{

const auto tables = it.second.equal_range("table");

const auto views = it.second.equal_range("view");

std::map<std::string, sqlb::ObjectPtr> objects;

for(auto jt=tables.first;jt!=tables.second;++jt)

objects.insert({jt->second->name(), jt->second});

for(auto jt=views.first;jt!=views.second;++jt)

objects.insert({jt->second->name(), jt->second});

for(const auto& jt : objects)

{

sqlb::ObjectIdentifier obj(it.first, jt.second->name());

QListWidgetItem* item = new QListWidgetItem(QIcon(QString(":icons/%1").arg(QString::fromStdString(sqlb::Object::typeToString(jt.second->type())))), QString::fromStdString(obj.toDisplayString()));

item->setData(Qt::UserRole, QString::fromStdString(obj.toSerialised()));

ui->listTables->addItem(item);

}

}

// Sort list of tables and select the table specified in the selection parameter or alternatively the first one

ui->listTables->model()->sort(0);

if(selection.isEmpty())

{

ui->listTables->setCurrentItem(ui->listTables->item(0));

} else {

for(int i=0;i<ui->listTables->count();i++)

{

if(sqlb::ObjectIdentifier(ui->listTables->item(i)->data(Qt::UserRole).toString().toStdString()) == selection)

{

ui->listTables->setCurrentRow(i);

break;

}

}

}

} else {

// Hide table combo box

ui->labelTable->setVisible(false);

ui->listTables->setVisible(false);

resize(minimumSize());

}

}

ExportDataDialog::~ExportDataDialog()

{

delete ui;

}

bool ExportDataDialog::exportQuery(const QString& sQuery, const QString& sFilename)

{

switch(m_format)

{

case ExportFormatCsv:

return exportQueryCsv(sQuery, sFilename);

case ExportFormatJson:

return exportQueryJson(sQuery, sFilename);

}

return false;

}

bool ExportDataDialog::exportQueryCsv(const QString& sQuery, const QString& sFilename)

{

// Prepare the quote and separating characters

QChar quoteChar = currentQuoteChar();

QString quotequoteChar = QString(quoteChar) + quoteChar;

QChar sepChar = currentSeparatorChar();

QString newlineStr = currentNewLineString();

// Chars that require escaping

std::string special_chars = newlineStr.toStdString() + sepChar.toLatin1() + quoteChar.toLatin1();

// Open file

QFile file(sFilename);

if(file.open(QIODevice::WriteOnly))

{

// Open text stream to the file

QTextStream stream(&file);

QByteArray utf8Query = sQuery.toUtf8();

sqlite3_stmt *stmt;

auto pDb = pdb.get(tr("exporting CSV"));

int status = sqlite3_prepare_v2(pDb.get(), utf8Query.data(), utf8Query.size(), &stmt, nullptr);

if(SQLITE_OK == status)

{

if(ui->checkHeader->isChecked())

{

int columns = sqlite3_column_count(stmt);

for (int i = 0; i < columns; ++i)

{

QString content = QString::fromUtf8(sqlite3_column_name(stmt, i));

if(content.toStdString().find_first_of(special_chars) != std::string::npos)

stream << quoteChar << content.replace(quoteChar, quotequoteChar) << quoteChar;

else

stream << content;

if(i != columns - 1)

// Only output the separator value if sepChar isn't 0,

// as that's used to indicate no separator character

// should be used

if(!sepChar.isNull())

stream << sepChar;

}

stream << newlineStr;

}

QApplication::setOverrideCursor(Qt::WaitCursor);

int columns = sqlite3_column_count(stmt);

size_t counter = 0;

while(sqlite3_step(stmt) == SQLITE_ROW)

{

for (int i = 0; i < columns; ++i)

{

QString content = QString::fromUtf8(

reinterpret_cast<const char*>(sqlite3_column_blob(stmt, i)),

sqlite3_column_bytes(stmt, i));

// If no quote char is set but the content contains a line break, we enforce some quote characters. This probably isn't entirely correct

// but still better than having the line breaks unquoted and effectively outputting a garbage file.

if(quoteChar.isNull() && content.contains(newlineStr))

stream << '"' << content.replace('"', "\"\"") << '"';

// If the content needs to be quoted, quote it. But only if a quote char has been specified

else if(!quoteChar.isNull() && content.toStdString().find_first_of(special_chars) != std::string::npos)

stream << quoteChar << content.replace(quoteChar, quotequoteChar) << quoteChar;

// If it doesn't need to be quoted, don't quote it

else

stream << content;

if(i != columns - 1)

// Only output the separator value if sepChar isn't 0,

// as that's used to indicate no separator character

// should be used

if(!sepChar.isNull())

stream << sepChar;

}

stream << newlineStr;

if(counter % 1000 == 0)

qApp->processEvents();

counter++;

}

}

sqlite3_finalize(stmt);

QApplication::restoreOverrideCursor();

qApp->processEvents();

// Done writing the file

file.close();

} else {

QMessageBox::warning(this, QApplication::applicationName(),

tr("Could not open output file: %1").arg(sFilename));

return false;

}

return true;

}

bool ExportDataDialog::exportQueryJson(const QString& sQuery, const QString& sFilename)

{

// Open file

QFile file(sFilename);

if(file.open(QIODevice::WriteOnly))

{

QByteArray utf8Query = sQuery.toUtf8();

sqlite3_stmt *stmt;

auto pDb = pdb.get(tr("exporting JSON"));

int status = sqlite3_prepare_v2(pDb.get(), utf8Query.data(), utf8Query.size(), &stmt, nullptr);

json json_table;

if(SQLITE_OK == status)

{

QApplication::setOverrideCursor(Qt::WaitCursor);

int columns = sqlite3_column_count(stmt);

size_t counter = 0;

std::vector<std::string> column_names;

while(sqlite3_step(stmt) == SQLITE_ROW)

{

// Get column names if we didn't do so before

if(!column_names.size())

{

for(int i=0;i<columns;++i)

column_names.push_back(sqlite3_column_name(stmt, i));

}

json json_row;

for(int i=0;i<columns;++i)

{

int type = sqlite3_column_type(stmt, i);

std::string column_name = column_names[static_cast<size_t>(i)];

switch (type) {

case SQLITE_INTEGER: {

qint64 content = sqlite3_column_int64(stmt, i);

json_row[column_name] = content;

break;

}

case SQLITE_FLOAT: {

double content = sqlite3_column_double(stmt, i);

json_row[column_name] = content;

break;

}

case SQLITE_NULL: {

json_row[column_name] = nullptr;

break;

}

case SQLITE_TEXT: {

QString content = QString::fromUtf8(

reinterpret_cast<const char*>(sqlite3_column_text(stmt, i)),

sqlite3_column_bytes(stmt, i));

json_row[column_name] = content.toStdString();

break;

}

case SQLITE_BLOB: {

QByteArray content(reinterpret_cast<const char*>(sqlite3_column_blob(stmt, i)),

sqlite3_column_bytes(stmt, i));

QTextCodec *codec = QTextCodec::codecForName("UTF-8");

QString string = codec->toUnicode(content.toBase64(QByteArray::Base64Encoding));

json_row[column_name] = string.toStdString();

break;

}

}

}

json_table.push_back(json_row);

if(counter % 1000 == 0)

qApp->processEvents();

counter++;

}

}

sqlite3_finalize(stmt);

// Create JSON document

file.write(json_table.dump(ui->checkPrettyPrint->isChecked() ? 4 : -1).c_str());

QApplication::restoreOverrideCursor();

qApp->processEvents();

// Done writing the file

file.close();

} else {

QMessageBox::warning(this, QApplication::applicationName(),

tr("Could not open output file: %1").arg(sFilename));

return false;

}

return true;

}

void ExportDataDialog::accept()

{

QStringList file_dialog_filter;

QString default_file_extension;

switch(m_format)

{

case ExportFormatCsv:

file_dialog_filter << FILE_FILTER_CSV

<< FILE_FILTER_TSV

<< FILE_FILTER_DSV

<< FILE_FILTER_TXT

<< FILE_FILTER_ALL;

default_file_extension = FILE_EXT_CSV_DEFAULT;

break;

case ExportFormatJson:

file_dialog_filter << FILE_FILTER_JSON

<< FILE_FILTER_TXT

<< FILE_FILTER_ALL;

default_file_extension = FILE_EXT_JSON_DEFAULT;

break;

}

if(!m_sQuery.isEmpty())

{

// called from sqlexecute query tab

QString sFilename = FileDialog::getSaveFileName(

CreateDataFile,

this,

tr("Choose a filename to export data"),

file_dialog_filter.join(";;"));

if(sFilename.isEmpty())

{

close();

return;

}

exportQuery(m_sQuery, sFilename);

} else {

// called from the File export menu

QList<QListWidgetItem*> selectedItems = ui->listTables->selectedItems();

if(selectedItems.isEmpty())

{

QMessageBox::warning(this, QApplication::applicationName(),

tr("Please select at least 1 table."));

return;

}

// Get filename

QStringList filenames;

if(selectedItems.size() == 1)

{

QString fileName = FileDialog::getSaveFileName(

CreateDataFile,

this,

tr("Choose a filename to export data"),

file_dialog_filter.join(";;"),

selectedItems.at(0)->text() + default_file_extension);

if(fileName.isEmpty())

{

close();

return;

}

filenames << fileName;

} else {

// ask for folder

QString exportfolder = FileDialog::getExistingDirectory(

CreateDataFile,

this,

tr("Choose a directory"),

QFileDialog::ShowDirsOnly | QFileDialog::DontResolveSymlinks);

if(exportfolder.isEmpty())

{

close();

return;

}

for(const QListWidgetItem* item : selectedItems)

filenames << QDir(exportfolder).filePath(item->text() + default_file_extension);

}

// Only if the user hasn't clicked the cancel button

for(int i = 0; i < selectedItems.size(); ++i)

{

// if we are called from execute sql tab, query is already set

// and we only export 1 select

QString sQuery = QString("SELECT * FROM %1;").arg(QString::fromStdString(sqlb::ObjectIdentifier(selectedItems.at(i)->data(Qt::UserRole).toString().toStdString()).toString()));

exportQuery(sQuery, filenames.at(i));

}

}

// Save the dialog preferences for future use

Settings::setValue("exportcsv", "firstrowheader", ui->checkHeader->isChecked());

Settings::setValue("exportjson", "prettyprint", ui->checkPrettyPrint->isChecked());

Settings::setValue("exportcsv", "separator", currentSeparatorChar());

Settings::setValue("exportcsv", "quotecharacter", currentQuoteChar());

Settings::setValue("exportcsv", "newlinecharacters", currentNewLineString());

// Notify the user the export has completed

QMessageBox::information(this, QApplication::applicationName(), tr("Export completed."));

QDialog::accept();

}

void ExportDataDialog::showCustomCharEdits()

{

// Retrieve selection info for the quote, separator, and newline widgets

int quoteIndex = ui->comboQuoteCharacter->currentIndex();

int quoteCount = ui->comboQuoteCharacter->count();

int sepIndex = ui->comboFieldSeparator->currentIndex();

int sepCount = ui->comboFieldSeparator->count();

int newLineIndex = ui->comboNewLineString->currentIndex();

int newLineCount = ui->comboNewLineString->count();

// Determine which will have their 'Other' line edit widget visible

bool quoteVisible = quoteIndex == (quoteCount - 1);

bool sepVisible = sepIndex == (sepCount - 1);

bool newLineVisible = newLineIndex == (newLineCount - 1);

// Update the visibility of the 'Other' line edit widgets

ui->editCustomQuote->setVisible(quoteVisible);

ui->editCustomSeparator->setVisible(sepVisible);

ui->editCustomNewLine->setVisible(newLineVisible);

}

void ExportDataDialog::setQuoteChar(const QChar& c)

{

QComboBox* combo = ui->comboQuoteCharacter;

// Set the combo and/or Other box to the correct selection

switch (c.toLatin1()) {

case '"':

combo->setCurrentIndex(0); // First option is a quote character

break;

case '\'':

combo->setCurrentIndex(1); // Second option is a single quote character

break;

case 0:

combo->setCurrentIndex(2); // Third option is blank (no character)

break;

default:

// For everything else, set the combo box to option 3 ('Other') and

// place the desired string into the matching edit line box

combo->setCurrentIndex(3);

if(!c.isNull())

{

// Don't set it if/when it's the 0 flag value

ui->editCustomQuote->setText(c);

}

break;

}

}

char ExportDataDialog::currentQuoteChar() const

{

QComboBox* combo = ui->comboQuoteCharacter;

switch (combo->currentIndex()) {

case 0:

return '"'; // First option is a quote character

case 1:

return '\''; // Second option is a single quote character

case 2:

return 0; // Third option is a blank (no character)

default:

// The 'Other' option was selected, so check if the matching edit

// line widget contains something

int customQuoteLength = ui->editCustomQuote->text().length();

if (customQuoteLength > 0) {

// Yes it does. Return its first character

char customQuoteChar = ui->editCustomQuote->text().at(0).toLatin1();

return customQuoteChar;

} else {

// No it doesn't, so return 0 to indicate it was empty

return 0;

}

}

}

void ExportDataDialog::setSeparatorChar(const QChar& c)

{

QComboBox* combo = ui->comboFieldSeparator;

// Set the combo and/or Other box to the correct selection

switch (c.toLatin1()) {

case ',':

combo->setCurrentIndex(0); // First option is a comma character

break;

case ';':

combo->setCurrentIndex(1); // Second option is a semi-colon character

break;

case '\t':

combo->setCurrentIndex(2); // Third option is a tab character

break;

case '|':

combo->setCurrentIndex(3); // Fourth option is a pipe symbol

break;

default:

// For everything else, set the combo box to option 3 ('Other') and

// place the desired string into the matching edit line box

combo->setCurrentIndex(4);

// Only put the separator character in the matching line edit box if

// it's not the flag value of 0, which is for indicating its empty

if(!c.isNull())

ui->editCustomSeparator->setText(c);

break;

}

}

char ExportDataDialog::currentSeparatorChar() const

{

QComboBox* combo = ui->comboFieldSeparator;

switch (combo->currentIndex()) {

case 0:

return ','; // First option is a comma character

case 1:

return ';'; // Second option is a semi-colon character

case 2:

return '\t'; // Third option is a tab character

case 3:

return '|'; // Fourth option is a pipe character

default:

// The 'Other' option was selected, so check if the matching edit

// line widget contains something

int customSeparatorLength = ui->editCustomSeparator->text().length();

if (customSeparatorLength > 0) {

// Yes it does. Return its first character

char customSeparatorChar = ui->editCustomSeparator->text().at(0).toLatin1();

return customSeparatorChar;

} else {

// No it doesn't, so return 0 to indicate it was empty

return 0;

}

}

}

void ExportDataDialog::setNewLineString(const QString& s)

{

QComboBox* combo = ui->comboNewLineString;

// Set the combo and/or Other box to the correct selection

if (s == "\r\n") {

// For Windows style newlines, set the combo box to option 0

combo->setCurrentIndex(0);

} else if (s == "\n") {

// For Unix style newlines, set the combo box to option 1

combo->setCurrentIndex(1);

} else {

// For everything else, set the combo box to option 2 ('Other') and

// place the desired string into the matching edit line box

combo->setCurrentIndex(2);

ui->editCustomNewLine->setText(s);

}

}

QString ExportDataDialog::currentNewLineString() const

{

QComboBox* combo = ui->comboNewLineString;

switch (combo->currentIndex()) {

case 0:

// Windows style newlines

return QString("\r\n");

case 1:

// Unix style newlines

return QString("\n");

default:

// Return the text from the 'Other' box

return QString(ui->editCustomNewLine->text().toLatin1());

}

}