Merge pull request #2289 from joto/primary-key

Add option to create primary key on flex tables
osm2pgsql-dev · Jan 13, 2025 · 8063f60 · 8063f60
2 parents 86084e2 + 4d04f50
commit 8063f60
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 2 deletions.
diff --git a/src/flex-lua-table.cpp b/src/flex-lua-table.cpp
@@ -112,7 +112,10 @@ void parse_create_index(lua_State *lua_state, flex_table_t *table)
         table->set_always_build_id_index();
     } else if (create_index == "unique") {
         table->set_always_build_id_index();
-        table->set_build_unique_id_index();
+        table->set_build_unique_id_index(false);
+    } else if (create_index == "primary_key") {
+        table->set_always_build_id_index();
+        table->set_build_unique_id_index(true);
     } else if (create_index != "auto") {
         throw fmt_error("Unknown value '{}' for 'create_index' field of ids",
                         create_index);

diff --git a/src/flex-table.cpp b/src/flex-table.cpp
@@ -223,6 +223,15 @@ std::string flex_table_t::build_sql_column_list() const
 
 std::string flex_table_t::build_sql_create_id_index() const
 {
+    if (m_primary_key_index) {
+        auto ts = tablespace_clause(index_tablespace());
+        if (!ts.empty()) {
+            ts = " USING INDEX" + ts;
+        }
+        return fmt::format("ALTER TABLE {} ADD PRIMARY KEY ({}){}", full_name(),
+                           id_column_names(), ts);
+    }
+
     return fmt::format("CREATE {}INDEX ON {} USING BTREE ({}) {}",
                        m_build_unique_id_index ? "UNIQUE " : "", full_name(),
                        id_column_names(),

diff --git a/src/flex-table.hpp b/src/flex-table.hpp
@@ -194,9 +194,10 @@ class flex_table_t
         return m_always_build_id_index;
     }
 
-    void set_build_unique_id_index() noexcept
+    void set_build_unique_id_index(bool as_primary_key) noexcept
     {
         m_build_unique_id_index = true;
+        m_primary_key_index = as_primary_key;
     }
 
     bool build_unique_id_index() const noexcept
@@ -265,6 +266,9 @@ class flex_table_t
     /// Build the index as a unique index.
     bool m_build_unique_id_index = false;
 
+    /// Index should be a primary key.
+    bool m_primary_key_index = false;
+
 }; // class flex_table_t
 
 class table_connection_t

diff --git a/tests/bdd/flex/lua-table-definitions.feature b/tests/bdd/flex/lua-table-definitions.feature
@@ -99,6 +99,40 @@ Feature: Table definitions in Lua file
         When running osm2pgsql flex
         Then table foo has 1562 rows
 
+    Scenario: Unique index is okay
+        Given the input file 'liechtenstein-2013-08-03.osm.pbf'
+        And the lua style
+            """
+            local t = osm2pgsql.define_table({
+                name = 'foo',
+                ids = { type = 'node', id_column = 'node_id', index = 'unique' },
+                columns = {}
+            })
+
+            function osm2pgsql.process_node(object)
+                t:insert({})
+            end
+            """
+        When running osm2pgsql flex
+        Then table foo has 1562 rows
+
+    Scenario: Primary key is okay
+        Given the input file 'liechtenstein-2013-08-03.osm.pbf'
+        And the lua style
+            """
+            local t = osm2pgsql.define_table({
+                name = 'foo',
+                ids = { type = 'node', id_column = 'node_id', index = 'primary_key' },
+                columns = {}
+            })
+
+            function osm2pgsql.process_node(object)
+                t:insert({})
+            end
+            """
+        When running osm2pgsql flex
+        Then table foo has 1562 rows
+
     Scenario: Can not create two tables with the same name
         Given the input file 'liechtenstein-2013-08-03.osm.pbf'
         And the lua style