pgLab/pglab/crud/CrudModel.cpp

#include "CrudModel.h"
#include "OpenDatabase.h"
#include "catalog/PgDatabaseCatalog.h"
#include "catalog/PgAttribute.h"
#include "catalog/PgAttributeContainer.h"
#include "catalog/PgConstraintContainer.h"
#include "catalog/PgClassContainer.h"
#include "SqlFormattingUtils.h"
#include "Pgsql_oids.h"
#include "CustomDataRole.h"
#include <QtConcurrent>
#include <QFuture>
#include <QMessageBox>
#include "Pgsql_oids.h"
#include "Pgsql_PgException.h"
#include "Pgsql_Params.h"
#include "Pgsql_Transaction.h"
#include <string>
#include "ScopeGuard.h"

CrudModel::CrudModel(QObject *parent)
	: QAbstractTableModel(parent)
    , m_dbConn()
{
	qDebug("CrudModel created");
	connect(&m_dbConn, &ASyncDBConnection::onStateChanged, this, &CrudModel::connectionStateChanged);
}

CrudModel::~CrudModel()
{
	m_dbConn.closeConnection();
}

/*
 * Strategy
 * when ordered by primary key, offset and limit work very quickly so we can get away with not loading
 * everything.
 */
void CrudModel::setConfig(std::shared_ptr<OpenDatabase> db, const PgClass &table)
{
	m_database = db;
	m_table = table;
	m_wholeTable = true;
	callLoadData = true;
	auto dbconfig = m_database->config();
	m_dbConn.setupConnection(dbconfig);
}

void CrudModel::setResult(std::shared_ptr<OpenDatabase> db, std::shared_ptr<Pgsql::Result> result)
{
	m_database = db;
	m_table.reset();
	m_wholeTable = false;
	callLoadData = false;
	auto dbconfig = m_database->config();
	m_dbConn.setupConnection(dbconfig);
	loadIntoModel(result);
}

QVariant CrudModel::headerData(int section, Qt::Orientation orientation, int role) const
{
    if (section == 0)
        return tr("#");

    int column = section - PreColumnCount;
	QVariant r;
    if (role == Qt::DisplayRole)
    {
        if (orientation == Qt::Horizontal)
        {
            QString s(m_roData->getColName(column));
			s += "\n";
            s += getTypeDisplayString(*m_database->catalog(), getType(column));
			r = s;
		}
        else
            r = QString::number(column + 1);
	}
	return r;
}

int CrudModel::rowCount(const QModelIndex &/*parent*/) const
{
    return (int)m_rowMapping.size();
}

int CrudModel::columnCount(const QModelIndex &/*parent*/) const
{
	int col_count = m_roData ? m_roData->cols() : 0;
    return PreColumnCount + col_count;
}

Oid CrudModel::getType(int column) const
{
	return m_roData ? m_roData->type(column) : InvalidOid;
}

CrudModel::Value CrudModel::getLatestData(int columnIndex, int rowIndex) const
{
    if (m_roData)
    {
        auto row_mapping = m_rowMapping[rowIndex];

        std::optional<Value> val;
        if (row_mapping.pending)
            val = m_pendingRowList.getValue(columnIndex, row_mapping.rowKey);

        if (!val.has_value())
            val = getSavedData(row_mapping, columnIndex);

        return val.value_or(Value());
    }
    return {};
}

CrudModel::Value CrudModel::getSavedData(int columnIndex, int rowIndex) const
{
    if (m_roData)
    {
        auto row_mapping = m_rowMapping[rowIndex];
        return getSavedData(row_mapping, columnIndex);
    }
    return {};
}

CrudModel::Value CrudModel::getSavedData(const RowMapping &row_mapping, int columnIndex) const
{
    std::optional<Value> val;
    if (!val.has_value() && row_mapping.isModified())
        val = row_mapping.modifiedValue(columnIndex);

    if (!val.has_value() && row_mapping.rowKey < m_roData->rows())
        if (!m_roData->null(columnIndex, row_mapping.rowKey))
            val = std::string(m_roData->val(columnIndex, row_mapping.rowKey));

    return val.value_or(Value());
}

QVariant StringToVariant(const std::string &str, Oid oid)
{
    Pgsql::Value v(str.c_str(), oid);
    switch (oid) {
        default:
            return QString::fromStdString(str);

        case Pgsql::timestamp_oid:
        case Pgsql::timestamptz_oid:
            return v.operator QDateTime();

        case Pgsql::date_oid:
            return v.operator QDate();

        case Pgsql::time_oid:
        case Pgsql::timetz_oid:
            return v.operator QTime();

        case Pgsql::int2_oid:
            return v.operator int16_t();

        case Pgsql::int4_oid:
            return v.operator int32_t();

        case Pgsql::oid_oid:
        case Pgsql::int8_oid:
            return QVariant::fromValue(static_cast<int64_t>(v));

        case Pgsql::bool_oid:
            return v.operator bool();

        case Pgsql::float4_oid:
            return v.operator float();

        case Pgsql::float8_oid:
            return v.operator double();
    }

}

QVariant CrudModel::data(const QModelIndex &index, int role) const
{
    if (index.column() < PreColumnCount)
    {
        if (role == Qt::DisplayRole || role == CustomSortRole)
            return index.row() + 1;
        else if (role == CustomDataTypeRole)
            return Pgsql::int4_oid;

        return {};
    }

    Oid typ = getType(index.column() - PreColumnCount);
    if (role == Qt::EditRole || role == Qt::DisplayRole)
    {
        auto value = getLatestData(index);
        if (value)
        {
            if (role == Qt::EditRole)
                return QString::fromUtf8(value->c_str());
            else
            {
                if (typ == Pgsql::bool_oid)
                    return *value == "t";
                else
                {
                    QString s = QString::fromUtf8(value->c_str());
                    s.truncate(256);
                    return s;
                }
            }
        }
    }
    else if (role == CustomDataTypeRole)
        return typ;
    else if (role == CustomSortRole)
    {
        auto value = getLatestData(index);
        if (value)
            return StringToVariant(value->c_str(), typ);
    }

    return {};
}

void CrudModel::loadData()
{
    QString table_name = m_table->fullyQualifiedQuotedObjectName();
	std::string q = "SELECT * FROM ";
    q += table_name.toStdString();
    m_dbConn.send(q, [this] (Expected<std::shared_ptr<Pgsql::Result>> res, qint64)
    {
        if (res.valid())
        {
			auto dbres = res.get();
            if (dbres && *dbres)
				QMetaObject::invokeMethod(this, "loadIntoModel", Qt::QueuedConnection,
										  Q_ARG(std::shared_ptr<Pgsql::Result>, dbres));
		}
	});
}

void CrudModel::loadIntoModel(std::shared_ptr<Pgsql::Result> data)
{
	beginResetModel();
	m_pendingRowList.clear();
	m_roData = data;
    initializeColumnList();
    m_mapping.emplace(createResultSourceMapping(*m_roData, *m_database->catalog()));
    lastRowKey = data->rows() - 1;
	initRowMapping();
    if (m_wholeTable)
        appendNewRowInternal();
	endResetModel();
}

void CrudModel::initRowMapping()
{
    size_t cnt = m_roData->rows();
    m_rowMapping.clear();
    m_rowMapping.reserve(cnt + 1);
    for (int i = 0; i < cnt; ++i)
        m_rowMapping.emplace_back(i);
}

void CrudModel::connectionStateChanged()
{
    switch (m_dbConn.state()) {
	case ASyncDBConnection::State::NotConnected:
		break;
	case ASyncDBConnection::State::Connecting:
		break;
	case ASyncDBConnection::State::Connected:
        if (callLoadData)
        {
			callLoadData = false;
			loadData();
		}
		break;
	case ASyncDBConnection::State::QuerySend:

		break;
	case ASyncDBConnection::State::CancelSend:
		break;
	case ASyncDBConnection::State::Terminating:
		break;
	}
}

Qt::ItemFlags CrudModel::flags(const QModelIndex &index) const
{
    Qt::ItemFlags flags = Qt::ItemIsSelectable | Qt::ItemIsEnabled;
    if (index.column() < PreColumnCount)
        return flags;

    if (!columnIsReadOnly(index.column()))
		flags |= Qt::ItemIsEditable;

	return flags;
}

bool CrudModel::setData(const QModelIndex &index, const QVariant &value, int role)
{
    if (index.column() < PreColumnCount)
        return false;

    if (role == Qt::EditRole)
    {
		int grid_row = index.row();
        int col = index.column() - PreColumnCount;

		auto& row_mapping = m_rowMapping[grid_row];
		row_mapping.pending = true;

		Value val;
		std::string s = value.toString().toUtf8().data();
        if (!s.empty())
        {
			if (s == "''")
				s.clear();
			val = s;
		}
		m_pendingRowList.setValue(col, row_mapping.rowKey, val);

		emit dataChanged(index, index, QVector<int>() << role);
		return true;
	}
	return false;
}

std::tuple<bool, std::vector<CrudModel::Value>> CrudModel::saveRow(const PendingRow &pending_row, const RowMapping &row_mapping)
{
    auto data = pending_row.data();
    if (data.empty())
        return { true, {} };

    if (row_mapping.isNew())
    {
        QString buffer;
        Pgsql::Params params;
        std::tie(buffer, params) = createInsertQuery(pending_row);

        Pgsql::Connection db_update_conn;
        auto dbconfig = m_database->config();
        db_update_conn.connect(dbconfig.connectionString());
        try
        {
            auto result = db_update_conn.queryParam(buffer, params);
            if (result && result.rows() == 1)
            {
                std::vector<Value> values;
                auto row = *result.begin();
                for (auto v : row)
                {
                    if (v.null())
                        values.push_back(Value());
                    else
                        values.push_back(std::string(v.c_str()));
                }

                return { true, values };
            }
        }
        catch (const Pgsql::PgResultError &ex)
        {
            QMessageBox msgBox;
            msgBox.setText(ex.what());
            msgBox.exec();
        }
        return { false, {} };
    }

    return saveUpdate(pending_row, row_mapping);
}

std::tuple<bool, std::vector<CrudModel::Value>> CrudModel::saveUpdate(const PendingRow &pending_row, const RowMapping &row_mapping)
{
    std::set<Oid> tables;
    for (const auto &e : pending_row.data())
        tables.insert(m_mapping->tableOid(e.first));

    try
    {
        Pgsql::Connection db_update_conn;
        auto dbconfig = m_database->config();
        db_update_conn.connect(dbconfig.connectionString());
        auto tx = Pgsql::Transaction::startTransaction(db_update_conn);
        for (Oid table_oid : tables)
        {
            QString buffer;
            Pgsql::Params params;
            std::tie(buffer, params) = createUpdateQuery(table_oid, pending_row, row_mapping);
            auto result = tx.queryParam(buffer, params);
            if (!result || result.tuplesAffected() != 1)
            {
                QMessageBox::warning(nullptr, tr("Update failed"), tr("The change did not update exactly one row in %1, the row may have changed or is no longer unique.").arg(tableName(table_oid)));
                return { false, {} };
            }
        }
        tx.commit();

        std::vector<Value> values;
        values.reserve(m_roData->cols());
        for (int col = 0; col < m_roData->cols(); ++col)
        {
            auto &&pending = pending_row.data();
            auto found = pending.find(col);
            values.push_back(found != pending.end() ? found->second : getSavedData(row_mapping, col));
        }
        return { true, values };
    }
    catch (const Pgsql::PgResultError &ex)
    {
        QMessageBox msgBox;
        msgBox.setText(ex.what());
        msgBox.exec();
    }
    return { false, {} };
}

QString CrudModel::tableName(Oid table_oid) const
{
    auto table = m_database->catalog()->classes()->getByKey(table_oid);
    return table ? table->fullyQualifiedQuotedObjectName() : QString();
}

std::vector<int> CrudModel::resultColumnsOfTable(Oid table_oid) const
{
    std::vector<int> cols;
    for (int col = 0; col < m_roData->cols(); ++col)
        if (m_roData->ftable(col) == table_oid && m_roData->ftableCol(col) > 0)
            cols.push_back(col);
    return cols;
}

QString CrudModel::attributeName(int col) const
{
    if (col < static_cast<int>(columnList.size()) && columnList[col])
        return columnList[col]->name;
    return QString::fromUtf8(m_roData->getColName(col));
}

Pgsql::Params CrudModel::getPKeyParamsForRow(const RowMapping &row_mapping, Oid table_oid) const
{
    Pgsql::Params params;
    for (int col : m_mapping->primaryKeyColumns(table_oid))
    {
        auto s = getSavedData(row_mapping, col);
        params.add(s, getType(col));
    }
    return params;
}

std::tuple<QString, Pgsql::Params> CrudModel::createUpdateQuery(Oid table_oid, const PendingRow &pending_row, const RowMapping &row_mapping)
{
    Pgsql::Params params;
    QString table_name = tableName(table_oid);
    QString buffer;
    QTextStream q(&buffer);
    q << "UPDATE " << table_name << " AS d\n  SET ";
    int param = 0;
    for (auto& e : pending_row.data())
    {
        if (m_mapping->tableOid(e.first) != table_oid)
            continue;
        if (param > 0)
            q << ",";
        q << quoteIdent(attributeName(e.first)) << "=$" << ++param;
        params.add(e.second, getType(e.first));
    }

    q << "\nWHERE ";
    int i = 0;
    for (int col : m_mapping->primaryKeyColumns(table_oid))
    {
        if (i > 0)
            q << " AND ";
        q << quoteIdent(attributeName(col)) << "=$" << ++param;
        ++i;
    }
    params.addParams(getPKeyParamsForRow(row_mapping, table_oid));
    q.flush();
    return { buffer, params };
}

std::tuple<QString, Pgsql::Params> CrudModel::createInsertQuery(const PendingRow &pending_row)
{
	Pgsql::Params params;
	auto data = pending_row.data();
    QString table_name = m_table->fullyQualifiedQuotedObjectName();
	QString buffer;
	QTextStream q(&buffer);
	q << "INSERT INTO " << table_name << "(";

	bool first = true;
    for (const auto& e : data)
    {
        int num = e.first;
        if (first)
            first = false;
        else
            q << ",";
        q << quoteIdent(attributeName(num));
	}
	q << ") VALUES ($1";
	for (size_t p = 2; p <= data.size(); ++p)
		q << ",$" << p;
	q << ") RETURNING *";
    for (auto& e : data)
		params.add(e.second, getType(e.first));

	q.flush();
	return { buffer, params };
}


QString CrudModel::createDeleteStatement() const
{
	Pgsql::Params params;
	QString table_name = m_table->fullyQualifiedQuotedObjectName();
	QString buffer;
	QTextStream q(&buffer);
	q << "DELETE FROM " << table_name;
	q << "\nWHERE ";
	int i = 0;
    for (int col : m_mapping->primaryKeyColumns(m_table->oid()))
    {
		if (i > 0)
			q << " AND ";
		q << quoteIdent(attributeName(col)) << "=$" << ++i;
	}

	q.flush();
	return buffer;
}


bool CrudModel::savePendingChanges()
{
    while (!m_pendingRowList.m_rows.empty())
    {
		auto iter = m_pendingRowList.m_rows.begin();
        int pendingKey = iter->first;
        auto mapping_iter = std::find_if(m_rowMapping.begin(), m_rowMapping.end(),
                                         [pendingKey](const RowMapping &rhs) -> bool { return rhs.rowKey == pendingKey; });
        if (mapping_iter == m_rowMapping.end())
            return false;

        auto [ok, modified_row] = saveRow(iter->second, *mapping_iter);
        if (ok)
        {
			int rowKey = iter->first;
			m_pendingRowList.m_rows.erase(iter);

			auto iter = std::find_if(m_rowMapping.begin(), m_rowMapping.end(),
                                  [rowKey](const RowMapping &rhs) -> bool { return rhs.rowKey == rowKey; });
            if (iter != m_rowMapping.end())
            {
                iter->setModifiedRowData(modified_row);
				iter->pending = false;
                if (IsLastRow(iter))
					appendNewRow();
			}
		}
        else
			return false;
	}

	return true;
}

bool CrudModel::IsLastRow(RowMappingVector::iterator mapping_iter) const
{
    return mapping_iter == --m_rowMapping.end();
}

bool CrudModel::columnIsReadOnly(int column_index) const
{
    if (m_roData == nullptr)
        return true;

    int col = column_index - PreColumnCount;
    if (!m_mapping || !m_mapping->isEditable(col))
        return true;

    const auto &column = columnList[col];
    return !column || column->getIdentity() == PgAttribute::Identity::Always
            || column->getGenerated() != PgAttribute::Generated::None;
}

void CrudModel::initializeColumnList()
{
    columnList.clear();
    columnList.reserve(m_roData->cols());
    auto atts = m_database->catalog()->attributes();
    for (int col = 0; col < m_roData->cols(); ++col)
    {
        int attnum = m_roData->ftableCol(col);
        if (attnum <= 0)
        {
            columnList.emplace_back();
            continue;
        }
        auto columns = atts->getColumnsForRelation(m_roData->ftable(col));
        auto find_result = std::find_if(columns.begin(), columns.end(), [attnum](const PgAttribute &att) { return att.num == attnum; });
        if (find_result != columns.end())
            columnList.emplace_back(*find_result);
        else
            columnList.emplace_back();
    }
}

bool CrudModel::submit()
{
	return savePendingChanges();
}

void CrudModel::revert()
{

}

void CrudModel::appendNewRow()
{
    if (!m_wholeTable)
        return;

    int row = static_cast<int>(m_rowMapping.size());
	beginInsertRows(QModelIndex(), row, row);
    appendNewRowInternal();
    endInsertRows();
}

void CrudModel::appendNewRowInternal()
{
    m_rowMapping.emplace_back(allocNewRowKey(), std::vector<Value>(m_roData->cols()));
}

std::tuple<bool, QString> CrudModel::removeRows(const std::set<IntegerRange<int>> &row_ranges)
{
    if (row_ranges.empty())
        return { true, "" };
    if (row_ranges.rbegin()->end() > static_cast<int>(m_rowMapping.size()))
        return { false, "Range error" };
    if (!m_wholeTable || !m_table || m_mapping->primaryKeyColumns(m_table->oid()).empty())
        return { false, "Rows can only be removed from a table with a primary key" };
	// When removing rows there is no direct mapping anymore between the rows in the grid
	// and the rows in m_roData

	// Therefor we need an indirection to keep track of which rows are visible and which
	// grid row maps to what data row. Maybe we can also keep track where the current data
	// of the row is located original data, pending data or modified data

	// 1. Get PKEY and remove that row from table

    try
    {
		Pgsql::Connection db_update_conn;
		auto dbconfig = m_database->config();
		db_update_conn.connect(dbconfig.connectionString());

		// First delete rows in table
		QString delete_statement = createDeleteStatement();
		{
			auto tx = Pgsql::Transaction::startTransaction(db_update_conn);
            for (auto range : row_ranges)
            {
                for (int current_row = range.start(); current_row < range.end(); ++current_row)
                {
					auto&& mapping = m_rowMapping[static_cast<size_t>(current_row)];
					auto params = getPKeyParamsForRow(mapping, m_table->oid());
                    if (!params.empty())
						db_update_conn.queryParam(delete_statement, params);
				}
			}
			tx.commit();
			// If something goes wrong after this commit we should reload contents of model
		}
		// Then from model
        RemoveRangesOfRowsFromModel(row_ranges);
		return  { true, "" };
    }
    catch (const Pgsql::PgResultError &error)
    {
		return { false, QString::fromUtf8(error.details().messageDetail.c_str()) };
	}
}

bool CrudModel::removeRows(int row, int count, const QModelIndex &)
{
    if (m_rowMapping.empty())
        return false;

    IntegerRange<int> range(row, count);
	auto [res, message] = removeRows({ range });
	return res;
}

void CrudModel::RemoveRangesOfRowsFromModel(const std::set<IntegerRange<int>> &row_ranges)
{
    int rows_deleted = 0;
    for (auto range : row_ranges)
    {
        range.setStart(range.start() - rows_deleted);
        RemoveRangeOfRowsFromModel(range);
        rows_deleted += range.length();
    }
}

void CrudModel::RemoveRangeOfRowsFromModel(IntegerRange<int> range)
{
    beginRemoveRows(QModelIndex(), range.start(), range.end() - 1);
    SCOPE_EXIT { endRemoveRows(); };
    auto first = m_rowMapping.begin() + range.start();
    m_rowMapping.erase(first, first + range.length());
}