Add after_nodes/ways/relations() callback functions in Lua

Author: joto

src/output-flex.cpp

@@ -804,6 +804,15 @@ int output_flex_t::expire_output_table()
     return 1;
 }
 
+void output_flex_t::call_lua_function(prepared_lua_function_t func)
+{
+    lua_pushvalue(lua_state(), func.index());
+    if (luaX_pcall(lua_state(), 0, func.nresults())) {
+        throw fmt_error("Failed to execute Lua function 'osm2pgsql.{}': {}.",
+                        func.name(), lua_tostring(lua_state(), -1));
+    }
+}
+
 void output_flex_t::call_lua_function(prepared_lua_function_t func,
                                       osmium::OSMObject const &object)
 {
@@ -823,6 +832,13 @@ void output_flex_t::call_lua_function(prepared_lua_function_t func,
     m_calling_context = calling_context::main;
 }
 
+void output_flex_t::get_mutex_and_call_lua_function(
+    prepared_lua_function_t func)
+{
+    std::lock_guard<std::mutex> const guard{lua_mutex};
+    call_lua_function(func);
+}
+
 void output_flex_t::get_mutex_and_call_lua_function(
     prepared_lua_function_t func, osmium::OSMObject const &object)
 {
@@ -960,11 +976,32 @@ void output_flex_t::sync()
     }
 }
 
-void output_flex_t::after_nodes() { flush_tables(m_table_connections); }
+void output_flex_t::after_nodes()
+{
+    if (m_after_nodes) {
+        get_mutex_and_call_lua_function(m_after_nodes);
+    }
+
+    flush_tables(m_table_connections);
+}
 
-void output_flex_t::after_ways() { flush_tables(m_table_connections); }
+void output_flex_t::after_ways()
+{
+    if (m_after_ways) {
+        get_mutex_and_call_lua_function(m_after_ways);
+    }
 
-void output_flex_t::after_relations() { flush_tables(m_table_connections); }
+    flush_tables(m_table_connections);
+}
+
+void output_flex_t::after_relations()
+{
+    if (m_after_relations) {
+        get_mutex_and_call_lua_function(m_after_relations);
+    }
+
+    flush_tables(m_table_connections);
+}
 
 void output_flex_t::stop()
 {
@@ -1138,7 +1175,9 @@ output_flex_t::output_flex_t(output_flex_t const *other,
   m_copy_thread(std::move(copy_thread)), m_lua_state(other->m_lua_state),
   m_process_node(other->m_process_node), m_process_way(other->m_process_way),
   m_process_relation(other->m_process_relation),
-  m_select_relation_members(other->m_select_relation_members)
+  m_select_relation_members(other->m_select_relation_members),
+  m_after_nodes(other->m_after_nodes), m_after_ways(other->m_after_ways),
+  m_after_relations(other->m_after_relations)
 {
     for (auto &table : *m_tables) {
         table.prepare(m_db_connection);
@@ -1369,6 +1408,12 @@ void output_flex_t::init_lua(std::string const &filename,
     m_select_relation_members = prepared_lua_function_t{
         lua_state(), calling_context::select_relation_members,
         "select_relation_members", 1};
+    m_after_nodes = prepared_lua_function_t{lua_state(), calling_context::main,
+                                            "after_nodes"};
+    m_after_ways = prepared_lua_function_t{lua_state(), calling_context::main,
+                                           "after_ways"};
+    m_after_relations = prepared_lua_function_t{
+        lua_state(), calling_context::main, "after_relations"};
 
     lua_remove(lua_state(), 1); // global "osm2pgsql"
 }

src/output-flex.hpp

@@ -181,14 +181,19 @@ class output_flex_t : public output_t
 private:
     void select_relation_members();
 
+    /// Call a Lua function that was "prepared" earlier.
+    void call_lua_function(prepared_lua_function_t func);
+
     /**
      * Call a Lua function that was "prepared" earlier with the OSMObject
      * as its only parameter.
      */
     void call_lua_function(prepared_lua_function_t func,
                            osmium::OSMObject const &object);
 
-    /// Aquire the lua_mutex and the call `call_lua_function()`.
+    /// Aquire the lua_mutex and then call `call_lua_function()`.
+    void get_mutex_and_call_lua_function(prepared_lua_function_t func);
+
     void get_mutex_and_call_lua_function(prepared_lua_function_t func,
                                          osmium::OSMObject const &object);
 
@@ -296,6 +301,9 @@ class output_flex_t : public output_t
     prepared_lua_function_t m_process_way{};
     prepared_lua_function_t m_process_relation{};
     prepared_lua_function_t m_select_relation_members{};
+    prepared_lua_function_t m_after_nodes{};
+    prepared_lua_function_t m_after_ways{};
+    prepared_lua_function_t m_after_relations{};
 
     calling_context m_calling_context = calling_context::main;
 

tests/bdd/flex/lua-callbacks.feature

@@ -0,0 +1,49 @@
+Feature: Lua callbacks are called
+
+    Scenario: Check access to osm2pgsql object from Lua
+        Given the input file 'liechtenstein-2013-08-03.osm.pbf'
+        And the lua style
+            """
+            local nodes = 0
+            local ways = 0
+            local relations = 0
+            local n = 0
+
+            osm2pgsql.define_node_table('dummy', {})
+
+            function osm2pgsql.process_node(object)
+                nodes = nodes + 1
+            end
+
+            function osm2pgsql.process_way(object)
+                ways = ways + 1
+            end
+
+            function osm2pgsql.process_relation(object)
+                relations = relations + 1
+            end
+
+            local function out()
+                print(n .. 'n=' .. nodes .. '@')
+                print(n .. 'w=' .. ways .. '@')
+                print(n .. 'r=' .. relations .. '@')
+                n = n + 1
+            end
+
+            osm2pgsql.after_nodes = out
+            osm2pgsql.after_ways = out
+            osm2pgsql.after_relations = out
+            """
+        When running osm2pgsql flex
+        Then the standard output contains
+            """
+            0n=1562@
+            0w=0@
+            0r=0@
+            1n=1562@
+            1w=7105@
+            1r=0@
+            2n=1562@
+            2w=7105@
+            2r=113@
+            """