pgLab/pglab/catalog/tables/TablesTableModel.cpp

#include "TablesTableModel.h"
#include "OpenDatabase.h"
#include "ScopeGuard.h"
#include "catalog/PgDatabaseCatalog.h"
#include "catalog/PgClass.h"
#include "catalog/PgClassContainer.h"
#include "catalog/PgNamespace.h"
#include "catalog/PgInheritsContainer.h"
//#include "Pgsql_declare.h"
#include "ui/catalog/tables/TableTreeBuilder.h"
#include "CustomDataRole.h"
#include <QBrush>
#include <QtConcurrent>
#include "Pgsql_Connection.h"

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

void TablesTableModel::setNamespaceFilter(NamespaceFilter nsf)
{
	m_namespaceFilter = nsf;
	refresh();
}

void TablesTableModel::setCatalog(std::shared_ptr<const PgDatabaseCatalog> cat)
{
	if (cat != m_catalog) {
		m_catalog = cat;
		refreshConnection = connect(m_catalog.get(), &PgDatabaseCatalog::refreshed,
									this, &TablesTableModel::refresh);
	}
	refresh();
}

bool TableLike(RelKind relkind)
{
    switch (relkind) {
    case RelKind::Table:
    case RelKind::View:
    case RelKind::MaterializedView:
    case RelKind::ForeignTable:
    case RelKind::PartitionedTable:
        return true;
    default:
        return false;
    }
}


void TablesTableModel::refresh()
{
	beginResetModel();
	SCOPE_EXIT { endResetModel(); };

	if (!m_catalog)
		return;

	auto classes = m_catalog->classes();
    rootNode.reset();

    auto nsfilter = GetNamespaceFilterLambda();
    std::map<Oid, PgClass> temp;
    for (const auto &e : *classes)
        if (TableLike(e.kind) && nsfilter(e))
            temp.emplace(e.oid(), e);

    TableTreeBuilder treeBuilder(temp, *m_catalog->inherits());
    auto [rn, oidIndex] = treeBuilder.Build();
    rootNode = rn;

    StartLoadTableSizes(std::move(oidIndex));
}

QVariant TablesTableModel::headerData(int section, Qt::Orientation orientation, int role) const
{
	if (orientation == Qt::Horizontal) {
		if (role == Qt::DisplayRole) {
			switch (section) {
			case NameCol: return tr("Name");
			case NamespaceCol: return tr("Schema");
			case KindCol: return tr("Kind");
			case OwnerCol: return tr("Owner");
			case TablespaceCol: return tr("Tablespace");
			case OptionsCol: return tr("Options");
			case AclCol: return tr("ACL");
            case CommentCol: return tr("Comment");
            case TotalSizeCol: return tr("Total size");
            case TableSizeCol: return tr("Table size");
            case IndexSizeCol: return tr("Index size");
            case ToastSizeCol: return tr("TOAST size");
            }
		}
	}
	return QVariant();
}

int TablesTableModel::rowCount(const QModelIndex &parent) const
{
    if (parent.isValid())
    {
        const TableNode *parentItem = static_cast<const TableNode*>(parent.internalPointer());
        return static_cast<int>(parentItem->children.size());
    }
    if (rootNode != nullptr)
        return static_cast<int>(rootNode->children.size());

    return 0;
}

int TablesTableModel::columnCount(const QModelIndex &) const
{
    return colCount;
}

QModelIndex TablesTableModel::index(int row, int column, const QModelIndex &parent) const
{
    if (parent.isValid())
    {
        const TableNode *parentItem = static_cast<const TableNode*>(parent.internalPointer());
        return createIndex(row, column, parentItem->getChildPtr(row));
    }
    return createIndex(row, column, rootNode->getChildPtr(row));
}

QModelIndex TablesTableModel::parent(const QModelIndex &index) const
{
    if (index.isValid())
    {
        const TableNode *item = static_cast<const TableNode*>(index.internalPointer());
        auto parent = item->parent.lock();
        if (parent && parent != rootNode)
        {
            return createIndex(parent->myIndex, 0, parent.get());
        }
    }
    return {};
}

Oid TablesTableModel::getType(int column) const
{
	Oid oid;
	switch (column) {
    case TotalSizeCol:
    case TableSizeCol:
    case IndexSizeCol:
    case ToastSizeCol:
        oid = Pgsql::int8_oid;
        break;
	case TablespaceCol:
	case OwnerCol:
	case NameCol:
	case NamespaceCol:
	case KindCol:
	case OptionsCol:
	case AclCol:
    case CommentCol:
	default:
		oid = Pgsql::varchar_oid;
	}
	return oid;
}

QVariant TablesTableModel::getData(const QModelIndex &index) const
{
//    const auto &table = rootNode->children[index.row()];
    const TableNode* table = nodeFromIndex(index);
    const auto &t = table->_class;
    const auto &s = table->sizes;

	switch (index.column()) {
        case NameCol: return t.objectName();
        case NamespaceCol: return t.nsName();
        case KindCol: return t.typeName();
        case OwnerCol: return t.ownerName();
        case TablespaceCol: return getTablespaceDisplayString(*m_catalog, t.tablespace);
        case OptionsCol: break;
        case AclCol: return t.aclString();
        case CommentCol: return t.description;
        case TotalSizeCol: return s.oid == 0 ? QVariant() : s.totalBytes;
        case TableSizeCol: return s.oid == 0 ? QVariant() : s.totalBytes - s.indexBytes - s.toastBytes;
        case IndexSizeCol: return s.oid == 0 ? QVariant() : s.indexBytes;
        case ToastSizeCol: return s.oid == 0 ? QVariant() : s.toastBytes;
    }

	return QVariant();
}

QVariant TablesTableModel::data(const QModelIndex &index, int role) const
{
	if (role == Qt::DisplayRole)
		return getData(index);
	else if (role == CustomDataTypeRole)
		return getType(index.column());
    else if (role == CustomDataMeaningRole)
    {
        switch (index.column()) {
        case TotalSizeCol:
        case TableSizeCol:
        case IndexSizeCol:
        case ToastSizeCol:
            return static_cast<int>(DataMeaningBytes);
        default:
            return static_cast<int>(DataMeaningNormal);
        }
    }

	return QVariant();
}

void TablesTableModel::StartLoadTableSizes(OidClassIndex oidIndex)
{
    QPointer p(this);
    QtConcurrent::run([this]
    {
        return QueryTableSizes();
    })
    .then(qApp, [p, oidIndex] (TableSizes sizes)
    {
        if (p)
            p.data()->PopulateSizes(oidIndex, sizes);
    });
}

TablesTableModel::TableSizes TablesTableModel::QueryTableSizes() const
{
    std::string nsfilter;
    switch (m_namespaceFilter)
    {
    case NamespaceFilter::User:
        nsfilter = "(NOT (nspname LIKE 'pg_%' OR nspname='information_schema'))";
        break;
    case NamespaceFilter::PgCatalog:
        nsfilter = "nspname = 'pg_catalog'";
        break;
    case NamespaceFilter::InformationSchema:
        nsfilter = "nspname = 'information_schema'";
        break;
    }

    std::string q =
        "SELECT pg_class.oid, "
        "  pg_total_relation_size(pg_class.oid) AS total_bytes, "
        "  CASE WHEN relkind='r' THEN pg_indexes_size(pg_class.oid) ELSE 0 END AS index_bytes, "
        "  CASE WHEN relkind='r' THEN pg_total_relation_size(reltoastrelid) ELSE 0 END AS toast_bytes "
        "\nFROM pg_catalog.pg_class \n"
        "    JOIN pg_namespace ON (relnamespace = pg_namespace.oid)"
        "    LEFT JOIN pg_catalog.pg_description AS d ON (objoid=pg_class.oid AND objsubid=0) \n"
        "WHERE relkind IN ('r', 'p', 'm') AND " + nsfilter;

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

    TablesTableModel::TableSizes sizes;
    sizes.reserve(sizesResult.rows());
    for (const Pgsql::Row& row : sizesResult)
    {
        TableSize size;
        size.oid = row.get(0);
        size.totalBytes = row.get(1);
        size.indexBytes = row.get(2);
        size.toastBytes = row.get(3);
        sizes.push_back(size);
    }
    return sizes;
}

void TablesTableModel::PopulateSizes(
    const OidClassIndex &oidIndex,
    const std::vector<TableSize> &sizes
)
{
    for (auto s : sizes)
    {
        auto findIter = oidIndex.find(s.oid);
        if (findIter != oidIndex.end())
        {
            findIter->second->sizes = s;
        }
    }
    emit dataChanged(
        index(0, TotalSizeCol),
        index(static_cast<int>(rootNode->children.size()), ToastSizeCol)
    );
}

const TableNode* TablesTableModel::nodeFromIndex(const QModelIndex &index)
{
    return static_cast<const TableNode*>(index.internalPointer());
}

std::function<bool (const PgClass &)> TablesTableModel::GetNamespaceFilterLambda()
{
    switch (m_namespaceFilter) {
    case NamespaceFilter::User:
        return [] (const PgClass &c) { return !c.ns().isSystemCatalog(); };

    case NamespaceFilter::PgCatalog:
        return [] (const PgClass &c) { return c.ns().objectName() == "pg_catalog"; };

    case NamespaceFilter::InformationSchema:
        return [] (const PgClass &c) { return c.ns().objectName() == "information_schema"; };
    }

}