pgLab/pglab/catalog/models/DatabasesTableModel.cpp

#include "DatabasesTableModel.h"
#include "CustomDataRole.h"
#include "OpenDatabase.h"
#include "ScopeGuard.h"
#include "catalog/PgDatabaseCatalog.h"
#include "catalog/PgDatabaseContainer.h"
#include "catalog/PgAuthIdContainer.h"
#include <QtConcurrent>
#include "Pgsql_Connection.h"

using namespace Pgsql;

DatabasesTableModel::DatabasesTableModel(std::shared_ptr<OpenDatabase> opendatabase, QObject *parent)
	: BaseTableModel(parent)
    , openDatabase(opendatabase)
{
}

void DatabasesTableModel::setDatabaseList(std::shared_ptr<const PgDatabaseCatalog> cat)
{
	beginResetModel();
    SCOPE_EXIT { endResetModel(); };

    databases.clear();

    std::map<Oid, int> oidIndex;
    m_catalog = cat;
    auto dats = cat->databases();
    databases.reserve(dats->count());
    for (const auto& d : *dats)
    {
        databases.emplace_back(d);
        oidIndex.emplace(d.oid(), static_cast<int>(databases.size() - 1));
    }
    StartLoadDatabaseSizes(std::move(oidIndex));
}

QVariant DatabasesTableModel::headerData(int section, Qt::Orientation orientation, int role) const
{
	QVariant v;
    if (orientation == Qt::Horizontal)
    {
        if (role == Qt::DisplayRole)
        {
			switch (section) {
			case NameCol:
				v = tr("Name");
				break;
			case DbaCol:
				v = tr("Owner");
				break;
			case EncodingCol:
				v = tr("Encoding");
				break;
			case CollateCol:
				v = tr("Collation");
				break;
			case CTypeCol:
				v = tr("CType");
				break;
			case IsTemplateCol:
				v = tr("Is template");
				break;
			case AllowConnCol:
				v = tr("Can connect");
				break;
			case ConnLimitCol:
				v = tr("Conn. limit");
				break;
			case TablespaceCol:
				v = tr("Tablespace");
				break;
            case CommentCol:
                v = tr("Comment");
                break;
            case SizeCol:
                v = tr("Size");
                break;
			case AclCol:
				v = tr("ACL");
				break;
			}
		}
	}

	return v;
}

int DatabasesTableModel::rowCount(const QModelIndex &) const
{
	int result = 0;
    result = static_cast<int>(databases.size());
	return result;
}

int DatabasesTableModel::columnCount(const QModelIndex &) const
{
	int result = COL_COUNT;

	return result;
}

Oid DatabasesTableModel::getType(int column) const
{
	Oid oid;
    switch (column)
    {
	case AllowConnCol:
	case IsTemplateCol:
        oid = Pgsql::bool_oid;
		break;
	case ConnLimitCol:
        oid = Pgsql::int4_oid;
		break;
    case SizeCol:
        oid = Pgsql::int8_oid;
        break;
    case AclCol:
	case CollateCol:
	case CTypeCol:
	case EncodingCol:
	case DbaCol:
	case NameCol:
	case TablespaceCol:
    case CommentCol:
        oid = Pgsql::varchar_oid;
		break;
	default:
        oid = InvalidOid;
	}
	return oid;
}

QVariant DatabasesTableModel::getData(const QModelIndex &index) const
{
	QVariant v;
    const Database &d = databases[index.row()];
    const PgDatabase &db = d.database;
    switch (index.column()) {
    case NameCol:
        v = db.objectName();
        break;
    case DbaCol:
        v = db.ownerName();
        break;
    case EncodingCol:
        v = db.encodingString;
        break;
    case CollateCol:
        v = db.collate;
        break;
    case CTypeCol:
        v = db.ctype;
        break;
    case IsTemplateCol:
        v = db.isTemplate;
        break;
    case AllowConnCol:
        v = db.allowConn;
        break;
    case ConnLimitCol:
        v = db.connLimit;
        break;
    case TablespaceCol:
        v = getTablespaceDisplayString(*m_catalog, db.tablespace);
        break;
    case CommentCol:
        v = db.description;
        break;
    case SizeCol:
        if (d.size.totalBytes >= 0)
            v = d.size.totalBytes;
        break;
    case AclCol:
        v = db.aclString();
        break;
    }
	return v;
}

QVariant DatabasesTableModel::getDataMeaning(const QModelIndex &index) const
{
    if (index.column() == SizeCol)
    {
        return static_cast<int>(DataMeaning::Bytes);
    }
    return BaseTableModel::getDataMeaning(index);
}

void DatabasesTableModel::StartLoadDatabaseSizes(std::map<Oid, int> oidIndex)
{
    QPointer p(this);
    QtConcurrent::run([this]
    {
        return QueryDatabaseSizes();
    })
    .then(qApp, [p, oi = std::move(oidIndex)] (DatabaseSizes sizes)
    {
        if (p)
        {
            p.data()->PopulateSizes(std::move(oi), std::move(sizes));
        }
    });
}

DatabasesTableModel::DatabaseSizes DatabasesTableModel::QueryDatabaseSizes()
{
    std::string q =
        "SELECT pg_database.oid, "
        "  case when has_database_privilege(current_role, pg_database.oid, \n"
        "    'connect') then pg_database_size(pg_database.oid) else -1 end \n"
        "FROM pg_database";

    Pgsql::Connection connection;
    connection.connect(openDatabase->config().connectionString());
    Pgsql::Result sizesResult = connection.query(q.c_str());

    DatabaseSizes result;
    result.reserve(sizesResult.rows());
    for (auto& row : sizesResult)
    {
        DatabaseSize ds;
        ds.oid = row.get(0);
        ds.totalBytes = row.get(1);
        result.push_back(ds);
    }
    return result;
}

void DatabasesTableModel::PopulateSizes(std::map<Oid, int> oidIndex, DatabaseSizes sizes)
{
    for (auto & s : sizes)
    {
        auto findIter = oidIndex.find(s.oid);
        if (findIter != oidIndex.end())
        {
            databases[findIter->second].size = s;
        }
    }
    emit dataChanged(
        createIndex(0, SizeCol),
        createIndex(static_cast<int>(databases.size()), SizeCol)
        );
}