#include "DbStructureModel.h"

#include "sqlitedb.h"

#include "sqlitetablemodel.h"

#include "Settings.h"

#include <QTreeWidgetItem>

#include <QMimeData>

#include <QMessageBox>

#include <QApplication>

DbStructureModel::DbStructureModel(DBBrowserDB& db, QObject* parent)

: QAbstractItemModel(parent),

m_db(db)

{

// Create root item and use its columns to store the header strings

QStringList header;

header << tr("Name") << tr("Object") << tr("Type") << tr("Schema") << tr("Database");

rootItem = new QTreeWidgetItem(header);

browsablesRootItem = nullptr;

}

DbStructureModel::~DbStructureModel()

{

delete rootItem;

}

int DbStructureModel::columnCount(const QModelIndex&) const

{

return rootItem->columnCount();

}

QVariant DbStructureModel::data(const QModelIndex& index, int role) const

{

if(!index.isValid())

return QVariant();

// Get the item the index points at

QTreeWidgetItem* item = static_cast<QTreeWidgetItem*>(index.internalPointer());

// Depending on the role either return the text or the icon

switch(role)

{

case Qt::DisplayRole:

// For the display role and the browsabled branch of the tree we want to show the column name including the schema name if necessary (i.e.

// for schemata != "main"). For the normal structure branch of the tree we don't want to add the schema name because it's already obvious from

// the position of the item in the tree.

if(index.column() == ColumnName && item->parent() == browsablesRootItem)

return sqlb::ObjectIdentifier(item->text(ColumnSchema), item->text(ColumnName)).toDisplayString();

else

return Settings::getValue("db", "hideschemalinebreaks").toBool() ? item->text(index.column()).replace("\n", " ").simplified() : item->text(index.column());

case Qt::EditRole:

case Qt::ToolTipRole: // Don't modify the text when it's supposed to be shown in a tooltip

return item->text(index.column());

case Qt::DecorationRole:

return item->icon(index.column());

default:

return QVariant();

}

}

Qt::ItemFlags DbStructureModel::flags(const QModelIndex &index) const

{

if(!index.isValid())

return Qt::ItemIsDropEnabled;

// All items are enabled and selectable

Qt::ItemFlags flags = Qt::ItemIsEnabled | Qt::ItemIsSelectable | Qt::ItemIsDropEnabled;

// Only enable dragging for entire table objects

QString type = data(index.sibling(index.row(), ColumnObjectType), Qt::DisplayRole).toString();

if(type == "table" || type == "view" || type == "index" || type == "trigger")

flags |= Qt::ItemIsDragEnabled;

return flags;

}

QVariant DbStructureModel::headerData(int section, Qt::Orientation orientation, int role) const

{

// Get the header string from the root item

if(orientation == Qt::Horizontal && role == Qt::DisplayRole)

return rootItem->data(section, role);

return QVariant();

}

QModelIndex DbStructureModel::index(int row, int column, const QModelIndex& parent) const

{

if(!hasIndex(row, column, parent))

return QModelIndex();

QTreeWidgetItem *parentItem;

if(!parent.isValid())

parentItem = rootItem;

else

parentItem = static_cast<QTreeWidgetItem*>(parent.internalPointer());

QTreeWidgetItem* childItem = parentItem->child(row);

if(childItem)

return createIndex(row, column, childItem);

else

return QModelIndex();

}

QModelIndex DbStructureModel::parent(const QModelIndex& index) const

{

if(!index.isValid())

return QModelIndex();

QTreeWidgetItem* childItem = static_cast<QTreeWidgetItem*>(index.internalPointer());

QTreeWidgetItem* parentItem = childItem->parent();

if(parentItem == rootItem)

return QModelIndex();

else

return createIndex(0, 0, parentItem);

}

int DbStructureModel::rowCount(const QModelIndex& parent) const

{

if(parent.column() > 0)

return 0;

if(!parent.isValid())

return rootItem->childCount();

else

return static_cast<QTreeWidgetItem*>(parent.internalPointer())->childCount();

}

void DbStructureModel::reloadData()

{

beginResetModel();

// Remove all data except for the root item

while(rootItem->childCount())

delete rootItem->child(0);

// Return here if no DB is opened

if(!m_db.isOpen())

{

endResetModel();

return;

}

// Create the nodes for browsables and for tables, indices, views and triggers. The idea here is to basically have two trees in one model:

// In the root node there are two nodes: 'browsables' and 'all'. The first node contains a list of a all browsable objects, i.e. views and tables.

// The seconds node contains four sub-nodes (tables, indices, views and triggers), each containing a list of objects of that type.

// This way we only have to have and only have to update one model and can use it in all sorts of places, just by setting a different root node.

browsablesRootItem = new QTreeWidgetItem(rootItem);

browsablesRootItem->setIcon(ColumnName, QIcon(QString(":/icons/view")));

browsablesRootItem->setText(ColumnName, tr("Browsables"));

// Make sure to always load the main schema first

QTreeWidgetItem* itemAll = new QTreeWidgetItem(rootItem);

itemAll->setIcon(ColumnName, QIcon(QString(":/icons/database")));

itemAll->setText(ColumnName, tr("All"));

buildTree(itemAll, "main");

// Add the temporary database as a node if it isn't empty. Make sure it's always second if it exists.

if(!m_db.schemata["temp"].isEmpty())

{

QTreeWidgetItem* itemTemp = new QTreeWidgetItem(itemAll);

itemTemp->setIcon(ColumnName, QIcon(QString(":/icons/database")));

itemTemp->setText(ColumnName, tr("Temporary"));

buildTree(itemTemp, "temp");

}

// Now load all the other schemata last

for(auto it=m_db.schemata.constBegin();it!=m_db.schemata.constEnd();++it)

{

// Don't load the main and temp schema again

if(it.key() != "main" && it.key() != "temp")

{

QTreeWidgetItem* itemSchema = new QTreeWidgetItem(itemAll);

itemSchema->setIcon(ColumnName, QIcon(QString(":/icons/database")));

itemSchema->setText(ColumnName, it.key());

buildTree(itemSchema, it.key());

}

}

// Refresh the view

endResetModel();

}

QStringList DbStructureModel::mimeTypes() const

{

QStringList types;

types << "text/plain";

return types;

}

QMimeData* DbStructureModel::mimeData(const QModelIndexList& indices) const

{

// Loop through selected indices

QByteArray d;

for(const QModelIndex& index : indices)

{

// Only export data for valid indices and only for the SQL column, i.e. only once per row

if(index.isValid() && index.column() == ColumnSQL)

{

// Add the SQL code used to create the object

d = d.append(data(index, Qt::DisplayRole).toString() + ";\n");

// If it is a table also add the content

if(data(index.sibling(index.row(), ColumnObjectType), Qt::DisplayRole).toString() == "table")

{

SqliteTableModel tableModel(m_db);

sqlb::ObjectIdentifier objid(data(index.sibling(index.row(), ColumnSchema), Qt::DisplayRole).toString(),

data(index.sibling(index.row(), ColumnName), Qt::DisplayRole).toString());

tableModel.setTable(objid);

tableModel.waitForFetchingFinished();

for(int i=0; i < tableModel.rowCount(); ++i)

{

QString insertStatement = "INSERT INTO " + objid.toString() + " VALUES(";

for(int j=1; j < tableModel.columnCount(); ++j)

insertStatement += QString("'%1',").arg(tableModel.data(tableModel.index(i, j)).toString());

insertStatement.chop(1);

insertStatement += ");\n";

d = d.append(insertStatement);

}

}

}

}

// Create the MIME data object

QMimeData* mime = new QMimeData();

mime->setProperty("db_file", m_db.currentFile()); // Also save the file name to avoid dropping an object on the same database as it comes from

mime->setData("text/plain", d);

return mime;

}

bool DbStructureModel::dropMimeData(const QMimeData* data, Qt::DropAction action, int, int, const QModelIndex&)

{

if(action == Qt::IgnoreAction)

return true;

if(!data->hasFormat("text/plain"))

return false;

if(data->property("db_file") == m_db.currentFile())

return false;

// Get data

QByteArray d = data->data("text/plain");

// Try to execute the SQL statement

if(m_db.executeMultiSQL(d, true, true))

{

m_db.updateSchema();

reloadData();

return true;

} else {

QMessageBox::warning(nullptr, QApplication::applicationName(), m_db.lastError());

return false;

}

}

void DbStructureModel::buildTree(QTreeWidgetItem* parent, const QString& schema)

{

// Build a map from object type to tree node to simplify finding the correct tree node later

QMap<QString, QTreeWidgetItem*> typeToParentItem;

// Get object map for the given schema

objectMap objmap = m_db.schemata[schema];

// Prepare tree

QTreeWidgetItem* itemTables = new QTreeWidgetItem(parent);

itemTables->setIcon(ColumnName, QIcon(QString(":/icons/table")));

itemTables->setText(ColumnName, tr("Tables (%1)").arg(objmap.values("table").count()));

typeToParentItem.insert("table", itemTables);

QTreeWidgetItem* itemIndices = new QTreeWidgetItem(parent);

itemIndices->setIcon(ColumnName, QIcon(QString(":/icons/index")));

itemIndices->setText(ColumnName, tr("Indices (%1)").arg(objmap.values("index").count()));

typeToParentItem.insert("index", itemIndices);

QTreeWidgetItem* itemViews = new QTreeWidgetItem(parent);

itemViews->setIcon(ColumnName, QIcon(QString(":/icons/view")));

itemViews->setText(ColumnName, tr("Views (%1)").arg(objmap.values("view").count()));

typeToParentItem.insert("view", itemViews);

QTreeWidgetItem* itemTriggers = new QTreeWidgetItem(parent);

itemTriggers->setIcon(ColumnName, QIcon(QString(":/icons/trigger")));

itemTriggers->setText(ColumnName, tr("Triggers (%1)").arg(objmap.values("trigger").count()));

typeToParentItem.insert("trigger", itemTriggers);

// Get all database objects and sort them by their name

QMultiMap<QString, sqlb::ObjectPtr> dbobjs;

for(auto it=objmap.constBegin(); it != objmap.constEnd(); ++it)

dbobjs.insert((*it)->name(), (*it));

// Add the database objects to the tree nodes

for(auto it=dbobjs.constBegin();it!=dbobjs.constEnd();++it)

{

// Object node

QTreeWidgetItem* item = addNode(typeToParentItem.value(sqlb::Object::typeToString((*it)->type())), *it, schema);

// If it is a table or view add the field nodes, add an extra node for the browsable section

if((*it)->type() == sqlb::Object::Types::Table || (*it)->type() == sqlb::Object::Types::View)

addNode(browsablesRootItem, *it, schema);

// Add field nodes if there are any

sqlb::FieldInfoList fieldList = (*it)->fieldInformation();

if(!fieldList.empty())

{

QStringList pk_columns;

if((*it)->type() == sqlb::Object::Types::Table)

{

sqlb::FieldVector pk = (*it).dynamicCast<sqlb::Table>()->primaryKey();

for(const sqlb::FieldPtr& pk_col : pk)

pk_columns.push_back(pk_col->name());

}

for(const sqlb::FieldInfo& field : fieldList)

{

QTreeWidgetItem *fldItem = new QTreeWidgetItem(item);

fldItem->setText(ColumnName, field.name);

fldItem->setText(ColumnObjectType, "field");

fldItem->setText(ColumnDataType, field.type);

fldItem->setText(ColumnSQL, field.sql);

fldItem->setText(ColumnSchema, schema);

if(pk_columns.contains(field.name))

fldItem->setIcon(ColumnName, QIcon(":/icons/field_key"));

else

fldItem->setIcon(ColumnName, QIcon(":/icons/field"));

}

}

}

}

QTreeWidgetItem* DbStructureModel::addNode(QTreeWidgetItem* parent, const sqlb::ObjectPtr& object, const QString& schema)

{

QString type = sqlb::Object::typeToString(object->type());

QTreeWidgetItem *item = new QTreeWidgetItem(parent);

item->setIcon(ColumnName, QIcon(QString(":/icons/%1").arg(type)));

item->setText(ColumnName, object->name());

item->setText(ColumnObjectType, type);

item->setText(ColumnSQL, object->originalSql());

item->setText(ColumnSchema, schema);

return item;

}