Implement ResultSourceMapping

pglablib/CMakeLists.txt

@@ -32,6 +32,7 @@ add_library(pglablib STATIC
     SqlFormattingUtils.cpp
     catalog/PgKeywordList.cpp
     QueryGenerator.cpp
+    ResultSourceMapping.cpp
     catalog/PgObject.cpp
     catalog/PgTablespace.cpp
     catalog/PgTablespaceContainer.cpp
@@ -107,6 +108,7 @@ add_library(pglablib STATIC
     catalog/PgCatalogTypes.h
     catalog/PgKeywordList.h
     QueryGenerator.h
+    ResultSourceMapping.h
     catalog/PgObject.h
     catalog/PgTablespace.h
     catalog/PgTablespaceContainer.h

pglablib/ResultSourceMapping.cpp

@@ -0,0 +1,139 @@
+#include "ResultSourceMapping.h"
+
+#include "Pgsql_Result.h"
+#include "catalog/PgDatabaseCatalog.h"
+#include "catalog/PgConstraintContainer.h"
+
+#include <algorithm>
+#include <stdexcept>
+
+namespace {
+
+/// Whether a non-key column of the table is selected more than once.
+/// A table that appears twice in a join (through aliases) shows up like that, and libpq can not tell the instances apart.
+bool hasAmbiguousColumn(const std::vector<ResultColumnSource> &columns, Oid tableOid, const std::vector<int> &pkAttnums)
+{
+    for (const auto &col : columns) {
+        if (col.tableOid != tableOid || std::find(pkAttnums.begin(), pkAttnums.end(), col.attnum) != pkAttnums.end())
+            continue;
+        auto count = std::count_if(columns.begin(), columns.end(), [&](const ResultColumnSource &other) {
+            return other.tableOid == tableOid && other.attnum == col.attnum;
+        });
+        if (count > 1)
+            return true;
+    }
+    return false;
+}
+
+}
+
+ResultSourceMapping::ResultSourceMapping(std::vector<ResultColumnSource> columns, const std::vector<TablePrimaryKey> &primaryKeys)
+    : m_columns(std::move(columns))
+{
+    for (const auto &col : m_columns) {
+        if (col.tableOid == InvalidOid || findTable(col.tableOid))
+            continue;
+
+        TableInfo info;
+        info.tableOid = col.tableOid;
+
+        auto pk = std::find_if(primaryKeys.begin(), primaryKeys.end(), [&](const TablePrimaryKey &key) {
+            return key.tableOid == col.tableOid;
+        });
+        if (pk != primaryKeys.end() && !pk->attnums.empty()) {
+            bool complete = true;
+            for (int attnum : pk->attnums) {
+                auto found = std::find_if(m_columns.begin(), m_columns.end(), [&](const ResultColumnSource &source) {
+                    return source.tableOid == col.tableOid && source.attnum == attnum;
+                });
+                if (found == m_columns.end()) {
+                    complete = false;
+                    break;
+                }
+                info.pkColumns.push_back(static_cast<int>(std::distance(m_columns.begin(), found)));
+            }
+            info.editable = complete && !hasAmbiguousColumn(m_columns, col.tableOid, pk->attnums);
+        }
+        if (!info.editable)
+            info.pkColumns.clear();
+
+        m_tables.push_back(info);
+    }
+}
+
+int ResultSourceMapping::columnCount() const
+{
+    return static_cast<int>(m_columns.size());
+}
+
+Oid ResultSourceMapping::tableOid(int col) const
+{
+    if (col < 0 || col >= columnCount())
+        throw std::range_error("column index out of range");
+    return m_columns[col].tableOid;
+}
+
+int ResultSourceMapping::attnum(int col) const
+{
+    if (col < 0 || col >= columnCount())
+        throw std::range_error("column index out of range");
+    return m_columns[col].attnum;
+}
+
+bool ResultSourceMapping::isEditable(int col) const
+{
+    if (col < 0 || col >= columnCount())
+        throw std::range_error("column index out of range");
+
+    const auto &source = m_columns[col];
+    if (source.attnum <= 0)
+        return false;
+
+    const TableInfo *info = findTable(source.tableOid);
+    return info && info->editable;
+}
+
+std::vector<int> ResultSourceMapping::primaryKeyColumns(Oid tableOid) const
+{
+    const TableInfo *info = findTable(tableOid);
+    return info ? info->pkColumns : std::vector<int>{};
+}
+
+const ResultSourceMapping::TableInfo *ResultSourceMapping::findTable(Oid tableOid) const
+{
+    auto found = std::find_if(m_tables.begin(), m_tables.end(), [tableOid](const TableInfo &info) {
+        return info.tableOid == tableOid;
+    });
+    return found == m_tables.end() ? nullptr : &*found;
+}
+
+ResultSourceMapping createResultSourceMapping(const Pgsql::Result &result, const PgDatabaseCatalog &catalog)
+{
+    std::vector<ResultColumnSource> columns;
+    std::vector<TablePrimaryKey> primaryKeys;
+
+    for (int col = 0; col < result.cols(); ++col) {
+        ResultColumnSource source;
+        source.tableOid = result.ftable(col);
+        source.attnum = result.ftableCol(col);
+        columns.push_back(source);
+
+        if (source.tableOid == InvalidOid)
+            continue;
+
+        bool known = std::any_of(primaryKeys.begin(), primaryKeys.end(), [&](const TablePrimaryKey &key) {
+            return key.tableOid == source.tableOid;
+        });
+        if (known)
+            continue;
+
+        TablePrimaryKey pk;
+        pk.tableOid = source.tableOid;
+        auto constraint = catalog.constraints()->getPrimaryForRelation(source.tableOid);
+        if (constraint)
+            pk.attnums.assign(constraint->key.begin(), constraint->key.end());
+        primaryKeys.push_back(pk);
+    }
+
+    return ResultSourceMapping(std::move(columns), primaryKeys);
+}

pglablib/ResultSourceMapping.h

@@ -0,0 +1,61 @@
+#ifndef RESULTSOURCEMAPPING_H
+#define RESULTSOURCEMAPPING_H
+
+#include <libpq-fe.h>
+
+#include <vector>
+
+class PgDatabaseCatalog;
+
+namespace Pgsql {
+    class Result;
+}
+
+/// The table column a result column is read from, as reported by libpq (PQftable/PQftablecol).
+/// tableOid is InvalidOid and attnum is 0 for expression columns.
+struct ResultColumnSource {
+    Oid tableOid = InvalidOid;
+    int attnum = 0;
+};
+
+/// The primary key of a table, attnums in key order.
+struct TablePrimaryKey {
+    Oid tableOid = InvalidOid;
+    std::vector<int> attnums;
+};
+
+/// Maps the columns of a query result to the table columns they come from and decides which
+/// cells can be edited. A cell is editable when it maps to a user column of a table whose full
+/// primary key is present in the result and which is not ambiguous.
+class ResultSourceMapping {
+public:
+    ResultSourceMapping(std::vector<ResultColumnSource> columns, const std::vector<TablePrimaryKey> &primaryKeys);
+
+    int columnCount() const;
+    /// Table the column is read from, InvalidOid for expression columns.
+    Oid tableOid(int col) const;
+    /// attnum of the column in its table, 0 for expression columns.
+    int attnum(int col) const;
+    bool isEditable(int col) const;
+    /// Result column indices holding the primary key of the table, in key order.
+    /// Empty when the table is not editable.
+    std::vector<int> primaryKeyColumns(Oid tableOid) const;
+
+private:
+    struct TableInfo {
+        Oid tableOid = InvalidOid;
+        bool editable = false;
+        std::vector<int> pkColumns;
+    };
+
+    std::vector<ResultColumnSource> m_columns;
+    std::vector<TableInfo> m_tables;
+
+    const TableInfo *findTable(Oid tableOid) const;
+};
+
+/// Builds the mapping for a query result, looking up the primary keys of the source tables in the catalog.
+ResultSourceMapping createResultSourceMapping(const Pgsql::Result &result, const PgDatabaseCatalog &catalog);
+
+
+#endif // RESULTSOURCEMAPPING_H