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

tests/pglabtests/CMakeLists.txt

@@ -17,6 +17,7 @@ add_executable(pglabtests
     tst_NameManglingRules.cpp
     tst_TypeMappings.cpp
     tst_Range.cpp
+    ResultSourceMappingTest.cpp
     )
 
 target_link_libraries(pglabtests PRIVATE

tests/pglabtests/ResultSourceMappingTest.cpp

@@ -0,0 +1,168 @@
+#include <gtest/gtest.h>
+#include <gmock/gmock-matchers.h>
+
+#include "ResultSourceMapping.h"
+
+
+namespace {
+
+constexpr Oid tableA = 1001;
+constexpr Oid tableB = 1002;
+
+ResultColumnSource src(Oid table, int attnum)
+{
+	ResultColumnSource s;
+	s.tableOid = table;
+	s.attnum = attnum;
+	return s;
+}
+
+ResultColumnSource expression()
+{
+	return ResultColumnSource{};
+}
+
+TablePrimaryKey pk(Oid table, std::vector<int> attnums)
+{
+	TablePrimaryKey key;
+	key.tableOid = table;
+	key.attnums = std::move(attnums);
+	return key;
+}
+
+}
+
+
+TEST(ResultSourceMappingTest, single_table_with_pk_all_columns_editable)
+{
+	ResultSourceMapping m({src(tableA, 1), src(tableA, 2), src(tableA, 3)}, {pk(tableA, {1})});
+
+	EXPECT_EQ(m.columnCount(), 3);
+	EXPECT_TRUE(m.isEditable(0));
+	EXPECT_TRUE(m.isEditable(1));
+	EXPECT_TRUE(m.isEditable(2));
+	EXPECT_EQ(m.tableOid(1), tableA);
+	EXPECT_EQ(m.attnum(2), 3);
+	EXPECT_THAT(m.primaryKeyColumns(tableA), ::testing::ElementsAre(0));
+}
+
+TEST(ResultSourceMappingTest, composite_pk_maps_key_columns_in_key_order)
+{
+	ResultSourceMapping m({src(tableA, 3), src(tableA, 1), src(tableA, 2)}, {pk(tableA, {1, 2})});
+
+	EXPECT_TRUE(m.isEditable(0));
+	EXPECT_THAT(m.primaryKeyColumns(tableA), ::testing::ElementsAre(1, 2));
+}
+
+TEST(ResultSourceMappingTest, join_with_both_pks_both_tables_editable)
+{
+	ResultSourceMapping m({src(tableA, 1), src(tableA, 2), src(tableB, 1), src(tableB, 2)},
+		{pk(tableA, {1}), pk(tableB, {1})});
+
+	for (int col = 0; col < m.columnCount(); ++col) {
+		EXPECT_TRUE(m.isEditable(col)) << "column " << col;
+	}
+	EXPECT_THAT(m.primaryKeyColumns(tableA), ::testing::ElementsAre(0));
+	EXPECT_THAT(m.primaryKeyColumns(tableB), ::testing::ElementsAre(2));
+}
+
+TEST(ResultSourceMappingTest, join_where_only_one_table_has_pk_in_result)
+{
+	ResultSourceMapping m({src(tableA, 1), src(tableA, 2), src(tableB, 2)},
+		{pk(tableA, {1}), pk(tableB, {1})});
+
+	EXPECT_TRUE(m.isEditable(0));
+	EXPECT_TRUE(m.isEditable(1));
+	EXPECT_FALSE(m.isEditable(2));
+	EXPECT_TRUE(m.primaryKeyColumns(tableB).empty());
+}
+
+TEST(ResultSourceMappingTest, pk_missing_from_result_is_read_only)
+{
+	ResultSourceMapping m({src(tableA, 2), src(tableA, 3)}, {pk(tableA, {1})});
+
+	EXPECT_FALSE(m.isEditable(0));
+	EXPECT_FALSE(m.isEditable(1));
+	EXPECT_TRUE(m.primaryKeyColumns(tableA).empty());
+}
+
+TEST(ResultSourceMappingTest, composite_pk_partially_present_is_read_only)
+{
+	ResultSourceMapping m({src(tableA, 1), src(tableA, 3)}, {pk(tableA, {1, 2})});
+
+	EXPECT_FALSE(m.isEditable(0));
+	EXPECT_FALSE(m.isEditable(1));
+}
+
+TEST(ResultSourceMappingTest, table_without_pk_is_read_only)
+{
+	ResultSourceMapping m({src(tableA, 1), src(tableA, 2)}, {});
+
+	EXPECT_FALSE(m.isEditable(0));
+	EXPECT_FALSE(m.isEditable(1));
+}
+
+TEST(ResultSourceMappingTest, expression_column_is_read_only)
+{
+	ResultSourceMapping m({src(tableA, 1), expression(), src(tableA, 2)}, {pk(tableA, {1})});
+
+	EXPECT_TRUE(m.isEditable(0));
+	EXPECT_FALSE(m.isEditable(1));
+	EXPECT_EQ(m.tableOid(1), InvalidOid);
+	EXPECT_EQ(m.attnum(1), 0);
+	EXPECT_TRUE(m.isEditable(2));
+}
+
+TEST(ResultSourceMappingTest, system_column_is_read_only)
+{
+	ResultSourceMapping m({src(tableA, 1), src(tableA, -2)}, {pk(tableA, {1})});
+
+	EXPECT_TRUE(m.isEditable(0));
+	EXPECT_FALSE(m.isEditable(1));
+}
+
+TEST(ResultSourceMappingTest, duplicate_pk_column_still_editable)
+{
+	ResultSourceMapping m({src(tableA, 1), src(tableA, 1), src(tableA, 2)}, {pk(tableA, {1})});
+
+	EXPECT_TRUE(m.isEditable(0));
+	EXPECT_TRUE(m.isEditable(1));
+	EXPECT_TRUE(m.isEditable(2));
+	EXPECT_THAT(m.primaryKeyColumns(tableA), ::testing::ElementsAre(0));
+}
+
+TEST(ResultSourceMappingTest, same_non_pk_column_twice_makes_table_read_only)
+{
+	ResultSourceMapping m({src(tableA, 1), src(tableA, 2), src(tableA, 2)}, {pk(tableA, {1})});
+
+	EXPECT_FALSE(m.isEditable(0));
+	EXPECT_FALSE(m.isEditable(1));
+	EXPECT_FALSE(m.isEditable(2));
+	EXPECT_TRUE(m.primaryKeyColumns(tableA).empty());
+}
+
+TEST(ResultSourceMappingTest, self_join_with_aliased_non_pk_columns_is_read_only)
+{
+	ResultSourceMapping m({src(tableA, 1), src(tableA, 1), src(tableA, 2), src(tableA, 2)}, {pk(tableA, {1})});
+
+	EXPECT_FALSE(m.isEditable(0));
+	EXPECT_FALSE(m.isEditable(2));
+}
+
+TEST(ResultSourceMappingTest, out_of_range_column_throws)
+{
+	ResultSourceMapping m({src(tableA, 1)}, {pk(tableA, {1})});
+
+	EXPECT_THROW(m.isEditable(1), std::range_error);
+	EXPECT_THROW(m.isEditable(-1), std::range_error);
+	EXPECT_THROW(m.tableOid(1), std::range_error);
+	EXPECT_THROW(m.attnum(-1), std::range_error);
+}
+
+TEST(ResultSourceMappingTest, empty_result_has_no_columns)
+{
+	ResultSourceMapping m({}, {});
+
+	EXPECT_EQ(m.columnCount(), 0);
+	EXPECT_TRUE(m.primaryKeyColumns(tableA).empty());
+}