pgLab/pglab/ConnectionListModel.cpp

#include "ConnectionListModel.h"
#include "ScopeGuard.h"

#include <botan/cryptobox.h>
#include <QDir>
#include <QException>
#include <QMimeData>
#include <QSettings>
#include <QSqlDatabase>
#include <QSqlError>
#include <QSqlQuery>
#include <QString>
#include <QStringBuilder>
#include <QStandardPaths>
#include <QStringBuilder>
#include <unordered_set>


namespace {

	const char * const q_create_table_conngroup =
R"__(
CREATE TABLE IF NOT EXISTS conngroup (
	conngroup_id  INTEGER PRIMARY KEY,
	gname         TEXT    NOT NULL UNIQUE
);)__";

	const char * const q_create_table_connection =
R"__(
CREATE TABLE IF NOT EXISTS connection (
	uuid          TEXT    PRIMARY KEY,
	cname         TEXT,
	conngroup_id  INTEGER NOT NULL,
	host          TEXT    ,
	hostaddr      TEXT    ,
	port          INTEGER NOT NULL,
	user          TEXT    ,
	dbname        TEXT    ,
	sslmode       INTEGER NOT NULL,
	sslcert       TEXT    ,
	sslkey        TEXT    ,
	sslrootcert   TEXT    ,
	sslcrl        TEXT    ,
	password      TEXT
);)__";

    const char * const q_create_table_migrations =
        R"__(
CREATE TABLE IF NOT EXISTS _migration (
    migration_id  TEXT PRIMARY KEY
);)__";

    const char * const q_load_migrations_present =
        R"__(
SELECT migration_id
FROM _migration;)__";


// Keeping migration function name and id DRY
#define APPLY_MIGRATION(id) ApplyMigration(#id, &MigrationDirector::id)

    class MigrationDirector {
    public:
        explicit MigrationDirector(SQLiteConnection &db)
            : db(db)
        {
        }

        void Execute()
        {
            InitConnectionTables();
            present = LoadMigrations();
            APPLY_MIGRATION(M20250215_0933_Parameters);
        }

    private:
        SQLiteConnection &db;
        std::unordered_set<QString> present;

        void M20250215_0933_Parameters()
        {
            db.Exec(R"__(
CREATE TABLE connection_parameter (
    connection_uuid TEXT,
    pname TEXT,
    pvalue TEXT NOT NULL,

    PRIMARY KEY(connection_uuid, pname)
);)__");

            db.Exec(R"__(
INSERT INTO connection_parameter (connection_uuid, pname, pvalue)
SELECT uuid, 'sslmode' AS pname, CASE
           WHEN sslmode = '0' THEN 'disable'
           WHEN sslmode = '1' THEN 'allow'
           WHEN sslmode = '2' THEN 'prefer'
           WHEN sslmode = '3' THEN 'require'
           WHEN sslmode = '4' THEN 'verify_ca'
           WHEN sslmode = '5' THEN 'verify_full'
       END AS pvalue
FROM connection
WHERE sslmode is not null and sslmode between 0 and 5)__");


            for (QString key : { "host", "hostaddr", "user", "dbname", "sslcert", "sslkey", "sslrootcert", "sslcrl" })
            {
                db.Exec(
                    "INSERT INTO connection_parameter (connection_uuid, pname, pvalue)"
                    "  SELECT uuid, '" % key % "', " % key % "\n"
                                          "   FROM connection\n"
                                          "   WHERE " % key % " IS NOT NULL and " % key % " <> '';");
            }
            db.Exec(R"__(
INSERT INTO connection_parameter (connection_uuid, pname, pvalue)
  SELECT uuid, 'port', port
  FROM connection
  WHERE port IS NOT NULL;
)__");

            for (QString column : { "host", "hostaddr", "user", "dbname", "sslmode", "sslcert", "sslkey", "sslrootcert", "sslcrl", "port" })
            {
                // sqlite does not seem to support dropping more then one column per alter table
                db.Exec("ALTER TABLE connection DROP COLUMN " % column % ";");
            }
        }

        void ApplyMigration(QString migration_id, void (MigrationDirector::*func)())
        {
            if (!present.contains(migration_id))
            {
                SQLiteTransaction tx(db);
                (this->*func)();
                RegisterMigration(migration_id);
                tx.Commit();
            }
        }

        std::unordered_set<QString> LoadMigrations()
        {
            std::unordered_set<QString> result;

            auto stmt = db.Prepare(q_load_migrations_present);
            while (stmt.Step())
            {
                result.insert(stmt.ColumnText(0));
            }
            return result;
        }

        void RegisterMigration(QString migrationId)
        {
            auto stmt = db.Prepare("INSERT INTO _migration VALUES (?1);");
            stmt.Bind(1, migrationId);
            stmt.Step();
        }

        void InitConnectionTables()
        {
            // Original schema
            db.Exec(q_create_table_conngroup);
            db.Exec(q_create_table_connection);
            // Start using migrations
            db.Exec(q_create_table_migrations);
        }
    };


    void SaveConnectionConfig(SQLiteConnection &db, const ConnectionConfig &cc, int conngroup_id)
	{
        const char * const q_insert_or_replace_into_connection =
            R"__(INSERT OR REPLACE INTO connection
    VALUES (?1, ?2, ?3, ?4);
)__"	;

        QByteArray b64; // needs to stay in scope until query is executed
        SQLiteTransaction tx(db);
        SQLitePreparedStatement stmt = db.Prepare(q_insert_or_replace_into_connection);
        stmt.Bind(1, cc.uuid().toString());
        stmt.Bind(2, cc.name());
        stmt.Bind(3, conngroup_id);
        auto& encodedPassword = cc.encodedPassword();
        if (!encodedPassword.isEmpty())
        {
            b64 = encodedPassword.toBase64(QByteArray::Base64Encoding);
            stmt.Bind(4, b64.data(), b64.length());
        }
        stmt.Step();
	}

} // end of unnamed namespace

ConnectionTreeModel::ConnectionTreeModel(QObject *parent, SQLiteConnection &db)
	: QAbstractItemModel(parent)
	, m_db(db)
{
}

void ConnectionTreeModel::load()
{
    //InitConnectionTables(m_db);
    MigrationDirector md(m_db);
    md.Execute();

    loadGroups();
    loadConnections();
}

void ConnectionTreeModel::loadGroups()
{
    auto stmt = m_db.Prepare("SELECT conngroup_id, gname FROM conngroup;");
    while (stmt.Step())
    {
        auto g = std::make_shared<ConnectionGroup>();
        g->conngroup_id = stmt.ColumnInteger(0);
        g->name = stmt.ColumnText(1);
        m_groups.push_back(g);
    }
}

void ConnectionTreeModel::loadConnections()
{
    auto stmt = m_db.Prepare(
        "SELECT uuid, cname, conngroup_id, password "
        "FROM connection ORDER BY conngroup_id, cname;");

    while (stmt.Step()) {
        auto cc = std::make_shared<ConnectionConfig>();
        cc->setUuid(QUuid::fromString(stmt.ColumnText(0)));
        cc->setName(stmt.ColumnText(1));
        cc->setEncodedPassword(stmt.ColumnCharPtr(3));
        loadConnectionParameters(*cc);

        int group_id = stmt.ColumnInteger(2);
        auto find_res = std::find_if(m_groups.begin(), m_groups.end(),
                                     [group_id] (auto item) { return item->conngroup_id == group_id; });
        if (find_res != m_groups.end()) {
            (*find_res)->add(cc);
        }
        else {
            throw std::runtime_error("conngroup missing");
        }
    }
}

void ConnectionTreeModel::loadConnectionParameters(ConnectionConfig &cc)
{
    auto stmt = m_db.Prepare(
        "SELECT pname, pvalue \n"
        "FROM connection_parameter\n"
        "WHERE connection_uuid=?1");
    stmt.Bind(1, cc.uuid().toString());
    while (stmt.Step())
    {
        cc.setParameter(
            stmt.ColumnText(0),
            stmt.ColumnText(1)
            );
    }
}

QVariant ConnectionTreeModel::data(const QModelIndex &index, int role) const
{
	// Code below assumes two level tree groups/connections
	// it will fail for nested groups
	QVariant v;

	auto privdata =  static_cast<ConnectionNode*>(index.internalPointer());
	if (auto group = dynamic_cast<ConnectionGroup*>(privdata); group != nullptr) {
		// This is a group
		if (role == Qt::DisplayRole) {
			if (index.column() == Name) {
				v = group->name;
			}
		}
	}
	else if (auto conn = dynamic_cast<ConnectionConfig*>(privdata); conn != nullptr) {
		// This is a connection
		if (role == Qt::DisplayRole) {
			switch (index.column()) {
			case Name: v = conn->name();	break;
			case Host: v = conn->host(); break;
			case Port: v = conn->port(); break;
			case User: v = conn->user(); break;
			case DbName: v= conn->dbname(); break;
			}
		}
	}

	return v;
}

QVariant ConnectionTreeModel::headerData(int section, Qt::Orientation orientation, int role) const
{
	QVariant v;
	if (orientation == Qt::Horizontal) {
		if (role == Qt::DisplayRole) {
			switch (section) {
			case Name: v = tr("Name");	break;
			case Host: v = tr("Host"); break;
			case Port: v = tr("Port"); break;
			case User: v = tr("User"); break;
			case DbName: v= tr("Database"); break;
			}
		}
	}
	return v;
}

QModelIndex ConnectionTreeModel::index(int row, int column, const QModelIndex &parent) const
{
	if (!hasIndex(row, column, parent))
		return {};

	const ConnectionNode *node = nullptr;
	if (parent.isValid()) {
		auto privdata = static_cast<ConnectionNode*>(parent.internalPointer());
		if (auto group = dynamic_cast<ConnectionGroup*>(privdata); group != nullptr) {
			node = group->connections().at(row).get();
		}
		else {
			throw std::logic_error("Should never ask for a child index of a connectionconfig");
		}
	}
	else {
		node = m_groups[row].get();
	}
	return createIndex(row, column, const_cast<ConnectionNode*>(node));
}

QModelIndex ConnectionTreeModel::parent(const QModelIndex &index) const
{
	if (!index.isValid())
		return {};

	auto privdata =  static_cast<ConnectionNode*>(index.internalPointer());
	if (auto group = dynamic_cast<ConnectionGroup*>(privdata); group != nullptr) {
		return {};
	}
	else if (auto config = dynamic_cast<ConnectionConfig*>(privdata); config != nullptr) {
		auto p = config->parent();
		auto find_res = std::find_if(m_groups.begin(), m_groups.end(), [p] (auto item) -> bool { return *p == *item; });
		if (find_res != m_groups.end()) {
			return createIndex(find_res - m_groups.begin(), 0,
							   const_cast<ConnectionGroup*>(config->parent()));
		}
	}
	throw std::logic_error("Should never get here");
}

int ConnectionTreeModel::rowCount(const QModelIndex &parent) const
{
	int result = 0;
	if (parent.isValid()) {
		auto privdata =  static_cast<ConnectionNode*>(parent.internalPointer());
		if (auto group = dynamic_cast<ConnectionGroup*>(privdata); group != nullptr) {
			result = group->connections().size();
		}
		else if (auto config = dynamic_cast<ConnectionConfig*>(privdata); config != nullptr) {
			result = 0;
		}
	}
	else {
		result = m_groups.size();
	}
	return result;
}

int ConnectionTreeModel::columnCount(const QModelIndex &) const
{
	return ColCount;
}

bool ConnectionTreeModel::removeRows(int row, int count, const QModelIndex &parent)
{
	if (parent.isValid() && count == 1) {
		// should be a group
		auto grp = m_groups[parent.row()];
		for (int i = 0; i < count; ++i) {
			QUuid uuid = grp->connections().at(row + i)->uuid();
            auto stmt = m_db.Prepare(
				"DELETE FROM connection "
                "  WHERE uuid=?0");
            stmt.Bind(0, uuid.toString());
            stmt.Step();
		}
		beginRemoveRows(parent, row, row + count - 1);
		SCOPE_EXIT { endRemoveRows(); };
		grp->erase(row, count);
		return true;
	}
	return false;
}

void ConnectionTreeModel::save(const QString &group_name, const ConnectionConfig &cc)
{
	auto [grp_idx, conn_idx] = findConfig(cc.uuid());
	if (grp_idx >= 0) {
		auto grp = m_groups[grp_idx];
		if (grp->name == group_name) {
			// update config
			grp->update(conn_idx, cc);
			// send change event
			auto node = grp->connections().at(conn_idx);
			dataChanged(
				createIndex(conn_idx, 0, node.get()),
				createIndex(conn_idx, ColCount-1, node.get()));
			saveToDb(*node);
			return;
		}
		else {
			auto parent = createIndex(grp_idx, 0, grp.get());
			beginRemoveRows(parent, conn_idx, conn_idx);
			SCOPE_EXIT { endRemoveRows(); };
			grp->erase(conn_idx);
		}
	}
	// Here we can assume we have to find the new group or create a new group
	// because if the connection was in the right group the function has already returned.
	// We assume the model is in sync with the DB as the DB should not be shared!
	int new_grp_idx = findGroup(group_name);
	if (new_grp_idx < 0) {
        // Group not found we are g
        new_grp_idx = addGroup(group_name);
	}
	auto new_grp = m_groups[new_grp_idx];

	auto parent = createIndex(new_grp_idx, 0, new_grp.get());
	auto idx = new_grp->connections().size();
	beginInsertRows(parent, idx, idx);
	SCOPE_EXIT { endInsertRows(); };
	auto node = std::make_shared<ConnectionConfig>(cc);
    new_grp->add(node);
    saveToDb(*node);
}

void ConnectionTreeModel::save(const ConnectionConfig &cc)
{
    saveToDb(cc);
}

void ConnectionTreeModel::clearAllPasswords()
{
    for (auto group : m_groups)
        for (auto cc : group->connections())
        {
            cc->setEncodedPassword({});
            saveToDb(*cc);
        }
}

std::tuple<int, int> ConnectionTreeModel::findConfig(const QUuid uuid) const
{
	int group_idx = -1, connection_idx = -1;

	for (int grp_idx = 0; grp_idx < m_groups.size(); ++grp_idx) {
		auto && grp = m_groups[grp_idx];
		auto && conns = grp->connections();
		auto find_res = std::find_if(conns.begin(), conns.end(),
			[&uuid] (auto item) -> bool { return item->uuid() == uuid; });
		if (find_res != conns.end()) {
			group_idx = grp_idx;
			connection_idx = find_res - conns.begin();
			break;
		}
	}

	return { group_idx, connection_idx };
}

int ConnectionTreeModel::findGroup(const QString &name) const
{
	for (int idx = 0; idx < m_groups.size(); ++idx) {
		if (m_groups[idx]->name == name) return idx;
	}
	return -1;
}

int ConnectionTreeModel::addGroup(const QString &group_name)
{
    auto stmt = m_db.Prepare("INSERT INTO conngroup (gname) VALUES (?1)");
    stmt.Bind(1, group_name);

    auto cg = std::make_shared<ConnectionGroup>();
    cg->conngroup_id = m_db.LastInsertRowId();
	cg->name = group_name;

	int row = m_groups.size();
	beginInsertRows({}, row, row);
	SCOPE_EXIT { endInsertRows(); };
	m_groups.push_back(cg);
	return row;
}

void ConnectionTreeModel::removeGroup(int row)
{
	beginRemoveRows({}, row, row);
	SCOPE_EXIT { endRemoveRows(); };
	auto id = m_groups[row]->conngroup_id;

    auto stmt = m_db.Prepare("DELETE FROM connection WHERE conngroup_id=?1");
    stmt.Bind(1, id);
    stmt.Step();

    stmt = m_db.Prepare("DELETE FROM conngroup WHERE conngroup_id=?1");
    stmt.Bind(1, id);
    stmt.Step();

	m_groups.remove(row);
}

int ConnectionTreeModel::findGroup(int conngroup_id) const
{
	auto find_res = std::find_if(m_groups.begin(), m_groups.end(),
		[conngroup_id] (auto item) { return item->conngroup_id == conngroup_id; });
	if (find_res == m_groups.end())
		return -1;
	return find_res - m_groups.begin();
}

ConnectionConfig *ConnectionTreeModel::getConfigFromModelIndex(QModelIndex index)
{
	if (!index.isValid())
		return nullptr;
	auto node = static_cast<ConnectionNode*>(index.internalPointer());
	return dynamic_cast<ConnectionConfig*>(node);
}

ConnectionGroup *ConnectionTreeModel::getGroupFromModelIndex(QModelIndex index)
{
	if (!index.isValid())
		return nullptr;
	auto node = static_cast<ConnectionNode*>(index.internalPointer());
	return dynamic_cast<ConnectionGroup*>(node);
}

void ConnectionTreeModel::saveToDb(const ConnectionConfig &cc)
{
    SaveConnectionConfig(m_db, cc, cc.parent()->conngroup_id);
}


Qt::DropActions ConnectionTreeModel::supportedDropActions() const
{
	return Qt::MoveAction;
}

Qt::DropActions ConnectionTreeModel::supportedDragActions() const
{
	return Qt::MoveAction;
}

Qt::ItemFlags ConnectionTreeModel::flags(const QModelIndex &index) const
{
	Qt::ItemFlags defaultFlags = QAbstractItemModel::flags(index);

	ConnectionConfig* cfg = getConfigFromModelIndex(index);

	if (cfg)
		return Qt::ItemIsDragEnabled | defaultFlags;
	else
		return Qt::ItemIsDropEnabled | defaultFlags;
}

//bool ConnectionTreeModel::insertRows(int row, int count, const QModelIndex &parent)
//{
//	return false;
//}

//bool ConnectionTreeModel::moveRows(const QModelIndex &sourceParent, int sourceRow, int count, const QModelIndex &destinationParent, int destinationChild)
//{
//	return false;
//}
namespace {
	const auto mimeType = "application/vnd.pgLab.connection";
}

QStringList ConnectionTreeModel::mimeTypes() const
{
	return { mimeType };
}

QMimeData *ConnectionTreeModel::mimeData(const QModelIndexList &indexes) const
{
	QMimeData *mimeData = new QMimeData;
	QByteArray encodedData;

	QDataStream stream(&encodedData, QIODevice::WriteOnly);

	for (const QModelIndex &index : indexes) {
		if (index.isValid()) {
			QString text = data(index, Qt::DisplayRole).toString();
			stream << text;
		}
	}

	mimeData->setData(mimeType, encodedData);
	return mimeData;
}

bool ConnectionTreeModel::dropMimeData(const QMimeData *data, Qt::DropAction action, int row, int column, const QModelIndex &parent)
{

	return false;
}