pgLab/pglab/catalog/models/ColumnTableModel.cpp

#include "catalog/models/ColumnTableModel.h"
#include "catalog/PgDatabaseCatalog.h"
#include "catalog/PgAttribute.h"
#include "catalog/PgAttributeContainer.h"
#include "catalog/PgClassContainer.h"
#include "catalog/PgConstraintContainer.h"
#include "catalog/PgCollation.h"
#include "catalog/PgCollationContainer.h"
#include "catalog/PgType.h"
#include "catalog/PgTypeContainer.h"
#include "catalog/PgIndexContainer.h"
#include "ScopeGuard.h"
#include "SqlFormattingUtils.h"
#include <QBrush>
#include <QStringBuilder>

namespace {

	// Filter dropped and system columns but keep the oid column
	bool ColumnFilterWithOidsPred(PgAttribute &e)
	{
		return e.isdropped || (e.num <= 0 && e.name != "oid");
	}

	// Filter dropped and system columns
	bool ColumnFilterWithoutOidsPred(PgAttribute &e)
	{
		return e.isdropped || e.num <= 0;
	}

}

void ColumnTableModel::setData(std::shared_ptr<const PgDatabaseCatalog> cat, const std::optional<PgClass> &table)
{
	if (cat != m_catalog) {
		m_catalog = cat;
		refreshConnection = connect(m_catalog.get(), &PgDatabaseCatalog::refreshed, this, &ColumnTableModel::refresh);
	}
	m_table = table;
	refresh();
}

QVariant ColumnTableModel::headerData(int section, Qt::Orientation orientation, int role) const
{
	QVariant v;
	if (orientation == Qt::Horizontal) {
		if (role == Qt::DisplayRole) {
			QString c;
			if (section >= colCount) {
				const auto &tbl_idx = m_indexes[section - colCount];
				if (tbl_idx.isprimary)
					c = "PK";
				else if (tbl_idx.isunique)
					c = "UIdx";
				else
					c = "Idx";
			}
			else {
				switch (section) {
				case AttnumCol:
					c = tr("#");
					break;
				case NameCol:
					c = tr("Name");
					break;
				case TypeCol:
					c = tr("Type");
					break;
				case NullCol:
					c = tr("N");
					break;
				case DefaultCol:
					c = tr("Default");
					break;
				case ForeignKeyCol:
					c = tr("Refs");
					break;
				case CollationCol:
					c = tr("Collation");
					break;
                case CommentCol:
                    c = tr("Comment");
                    break;
				}
			}
			v = c;
		}
		else if (role == Qt::ToolTipRole) {
			if (section >= colCount) {
				const auto &tbl_idx = m_indexes[section - colCount];
				//auto idx_cls = m_catalog->classes()->getByKey(tbl_idx.indexrelid);
				auto idx_class_name = tbl_idx.objectName(); // getClassDisplayString(*m_catalog, tbl_idx.indexrelid);
				QString s;
				if (tbl_idx.isprimary)
					s = tr("Primary key");
				else if (tbl_idx.isunique)
					s = tr("Unique index");
				else
					s = tr("Index");
				s +=  "\n" + idx_class_name;
				v = s;
			}
			else {
				switch (section) {
//				case NameCol:
//					c = tr("Name");
//					break;
//				case TypeCol:
//					c = tr("Type");
//					break;
				case NullCol:
					v = tr("N = Column is nullable");
					break;
				case DefaultCol:
					v = tr("Default value for the columns");
					break;
				case ForeignKeyCol:
					v = tr("Foreign key constraint for this column");
					break;
//				case CollationCol:
//					c = tr("Collation");
//					break;
				}
			}


		}
	}

	return v;
}

int ColumnTableModel::rowCount(const QModelIndex &/*parent*/) const
{
    return static_cast<int>(m_columns.size());
}

int ColumnTableModel::columnCount(const QModelIndex &/*parent*/) const
{
	return colCount + m_indexes.size();
}

Oid ColumnTableModel::getType(int /*column*/) const
{
	Oid oid = Pgsql::varchar_oid;
	return oid;
}

QVariant ColumnTableModel::getData(const QModelIndex &index) const
{
	QVariant v;
	const int row = index.row();
	const auto &t = m_columns[row];

	const int col = index.column();
	if (col >= colCount) {
		const auto &tbl_idx = m_indexes[col - colCount];
		int i = 1;
		for (auto attr : tbl_idx.key) {
			if (attr == t.num) {
				v = i;
				break;
			}
			++i;
		}
	}
	else {
		switch (col) {
		case AttnumCol:
			//s = QString::asprintf("%d", (int) t.num);
			//break;
			return static_cast<int>(t.num);
		case NameCol:
			v = t.name;
			break;
		case TypeCol:
			v = getTypeDisplayString(*m_catalog, t.typid, t.typmod);
			break;
		case NullCol:
			v = QString::fromStdU16String(t.notnull ? u"" : u"N");
			break;
		case DefaultCol:
            v = getDefaultString(t);
			break;
		case ForeignKeyCol:
			v = getFKey(t);
			break;
		case CollationCol:
			if (t.collation != InvalidOid) {
				auto&& col = m_catalog->collations()->getByKey(t.collation);
				if (col)
					v = col->objectName();
			}
			break;
        case CommentCol:
            v = t.description;
            break;
		}
	}

	return v;
}

QString ColumnTableModel::getFKey(const PgAttribute &column) const
{
	QString result;
	auto&& list = m_catalog->constraints()->getFKeyForTableColumn(column.relid, column.num);
	for (auto&& elem : list) {
		if (elem.key[0] == column.num) {
			//result = elem.name;
			result = getForeignKeyConstraintReferencesShort(*m_catalog, elem);
		}
	}
    return result;
}

QString ColumnTableModel::getDefaultString(const PgAttribute &column) const
{
    if (column.identity != '\0') {
        QString sql = "GENERATED ";
        switch (column.identity) {
            case 'a':
                sql += "ALWAYS";
                break;
            case 'd':
                sql += "BY DEFAULT";
                break;
        }
        sql += " AS IDENTITY";
        // TODO sequence options might be missing
        return sql;
    }

    if (column.generated != '\0')
    {
        return "GENERATED ALWAYS AS " % column.defaultValue % " STORED";
    }


    if (column.hasdef)
        return column.defaultValue;

    return {};
}

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

	if (m_table) {
		m_columns = m_catalog->attributes()->getColumnsForRelation(m_table->oid());
		// hide system and dropped columns
		auto column_filter_pred = m_table->hasoids ? ColumnFilterWithOidsPred : ColumnFilterWithoutOidsPred;
		auto si = std::remove_if(m_columns.begin(), m_columns.end(), column_filter_pred);
		// move columns to end and remove them
		m_columns.erase(si, m_columns.end());

		// sort remaining columns by order in table
		std::sort(m_columns.begin(), m_columns.end(),
				  [] (auto &l, auto &r) -> bool { return l.num < r.num; });
	}
	else
		m_columns.clear();


	if (m_table) {
		m_indexes = m_catalog->indexes()->getIndexesForTable(m_table->oid());
		std::sort(m_indexes.begin(), m_indexes.end(),
			[] (const auto &l, const auto &r) -> bool
			{
				return l.isprimary > r.isprimary
					|| (l.isprimary == r.isprimary && l.oid() < r.oid());
			});
	}
	else
		m_indexes.clear();

	emit
}

QVariant ColumnTableModel::data(const QModelIndex &index, int role) const
{
	if (role == Qt::ForegroundRole && index.column() == TypeCol) {
		QVariant v;
		const auto &t = m_columns[index.row()];
		if (t.typid == InvalidOid)
			v = QBrush(Qt::black);
		else {
			auto c = m_catalog->types()->getByKey(t.typid);
			switch (c->category) {
			case TypCategory::Boolean:
				v = QBrush(Qt::darkGreen);
				break;
			case TypCategory::Numeric:
				v = QBrush(Qt::darkBlue);
				break;
			case TypCategory::DateTime:
			case TypCategory::Timespan:
				v = QBrush(Qt::darkMagenta);
				break;
			case TypCategory::String:
				v = QBrush(Qt::darkYellow);
				break;
			case TypCategory::Array:
			case TypCategory::Composite:
			case TypCategory::Enum:
			case TypCategory::Geometric:
			case TypCategory::NetworkAddress:
			case TypCategory::Pseudo:
			case TypCategory::Range:
			case TypCategory::UserDefined:
			case TypCategory::BitString:
			case TypCategory::Unknown:
			default:
				v = QBrush(Qt::black);
			}
		}
		return v;
	}
	if (role == Qt::TextAlignmentRole) {
		QVariant v;
		int col = index.column();
		if (col == NullCol || col >= colCount)
            v = int(Qt::AlignCenter | Qt::AlignVCenter);
		else
            v = int(Qt::AlignLeft | Qt::AlignVCenter);
		return v;
	}
	return BaseTableModel::data(index, role);
}

const PgAttribute& ColumnTableModel::column(int row) const
{
	return m_columns.at(static_cast<size_t>(row));
}