Add functions to check for tablespaces, schemas and extensions in db

Author: joto

src/CMakeLists.txt

@@ -26,6 +26,7 @@ target_sources(osm2pgsql_lib PRIVATE
   output-pgsql.cpp
   output.cpp
   pgsql.cpp
+  pgsql-capabilities.cpp
   pgsql-helper.cpp
   progress-display.cpp
   reprojection.cpp

src/pgsql-capabilities.cpp

@@ -0,0 +1,66 @@
+/**
+ * SPDX-License-Identifier: GPL-2.0-or-later
+ *
+ * This file is part of osm2pgsql (https://osm2pgsql.org/).
+ *
+ * Copyright (C) 2006-2022 by the osm2pgsql developer community.
+ * For a full list of authors see the git log.
+ */
+
+#include "format.hpp"
+#include "pgsql.hpp"
+#include "pgsql-helper.hpp"
+
+#include <set>
+#include <string>
+
+static std::set<std::string> init_set_from_table(pg_conn_t const &db_connection,
+                                                 char const *table,
+                                                 char const *column,
+                                                 char const *condition)
+{
+    std::set<std::string> values;
+
+    auto const res = db_connection.query(
+        PGRES_TUPLES_OK,
+        "SELECT {} FROM {} WHERE {}"_format(column, table, condition));
+    for (int i = 0; i < res.num_tuples(); ++i) {
+        values.insert(res.get_value_as_string(i, 0));
+    }
+
+    return values;
+}
+
+bool has_extension(pg_conn_t const &db_connection, std::string const &value)
+{
+    static const std::set<std::string> values = init_set_from_table(
+        db_connection, "pg_catalog.pg_extension", "extname", "true");
+
+    return values.count(value);
+}
+
+bool has_schema(pg_conn_t const &db_connection, std::string const &value)
+{
+    static const std::set<std::string> values = init_set_from_table(
+        db_connection, "pg_catalog.pg_namespace", "nspname",
+        "nspname !~ '^pg_' AND nspname <> 'information_schema'");
+
+    if (value.empty()) {
+        return true;
+    }
+
+    return values.count(value);
+}
+
+bool has_tablespace(pg_conn_t const &db_connection, std::string const &value)
+{
+    static const std::set<std::string> values =
+        init_set_from_table(db_connection, "pg_catalog.pg_tablespace",
+                            "spcname", "spcname != 'pg_global'");
+
+    if (value.empty()) {
+        return true;
+    }
+
+    return values.count(value);
+}

src/pgsql-capabilities.hpp

@@ -0,0 +1,21 @@
+#ifndef OSM2PGSQL_PGSQL_CAPABILITIES_HPP
+#define OSM2PGSQL_PGSQL_CAPABILITIES_HPP
+
+/**
+ * SPDX-License-Identifier: GPL-2.0-or-later
+ *
+ * This file is part of osm2pgsql (https://osm2pgsql.org/).
+ *
+ * Copyright (C) 2006-2022 by the osm2pgsql developer community.
+ * For a full list of authors see the git log.
+ */
+
+#include "pgsql.hpp"
+
+#include <string>
+
+bool has_extension(pg_conn_t const &db_connection, std::string const &value);
+bool has_schema(pg_conn_t const &db_connection, std::string const &value);
+bool has_tablespace(pg_conn_t const &db_connection, std::string const &value);
+
+#endif // OSM2PGSQL_PGSQL_CAPABILITIES_HPP

tests/CMakeLists.txt

@@ -72,6 +72,7 @@ set_test(test-output-pgsql-z_order)
 set_test(test-parse-osmium LABELS NoDB)
 set_test(test-persistent-cache LABELS NoDB)
 set_test(test-pgsql)
+set_test(test-pgsql-capabilities)
 set_test(test-reprojection LABELS NoDB)
 set_test(test-taginfo LABELS NoDB)
 set_test(test-tile LABELS NoDB)

tests/test-pgsql-capabilities.cpp

@@ -0,0 +1,39 @@
+/**
+ * SPDX-License-Identifier: GPL-2.0-or-later
+ *
+ * This file is part of osm2pgsql (https://osm2pgsql.org/).
+ *
+ * Copyright (C) 2006-2022 by the osm2pgsql developer community.
+ * For a full list of authors see the git log.
+ */
+
+#include <catch.hpp>
+
+#include "common-import.hpp"
+
+#include "pgsql-capabilities.hpp"
+
+static testing::db::import_t const db;
+
+TEST_CASE("has_extension() should work")
+{
+    auto const conn = db.db().connect();
+    REQUIRE(has_extension(conn, "postgis"));
+    REQUIRE_FALSE(has_schema(conn, "xzxzxzxz"));
+}
+
+TEST_CASE("has_schema() should work")
+{
+    auto const conn = db.db().connect();
+    REQUIRE(has_schema(conn, "public"));
+    REQUIRE_FALSE(has_schema(conn, "xzxzxzxz"));
+    REQUIRE_FALSE(has_schema(conn, "pg_toast"));
+}
+
+TEST_CASE("has_tablespace() should work")
+{
+    auto const conn = db.db().connect();
+    REQUIRE(has_tablespace(conn, "pg_default"));
+    REQUIRE_FALSE(has_tablespace(conn, "xzxzxzxz"));
+    REQUIRE_FALSE(has_tablespace(conn, "pg_global"));
+}