Add internal json/jsonb support to flex output

joto · joto · commit 94dae34b7aa1 · 2021-05-04T09:07:22.000+02:00
Table columns with type 'json' or 'jsonb' are now explicitly supported.
Lua booleans, strings, numbers, and tables can be written to those
columns and are automatically encoded as JSON on the fly.

This is a breaking change to the behaviour we had before where you had
to encode the JSON yourself in Lua before giving it to osm2pgsql. We
can still do this change as the flex backend is still marked
experimental and the behaviour is much cleaner this way then if we keep
backwards compatibility.
diff --git a/flex-config/places.lua b/flex-config/places.lua
@@ -2,16 +2,7 @@
 
 -- This example shows how you can use JSON and JSONB columns in your database.
 
--- You need the 'dkjson' JSON library for Lua. Install 'dkjson' with LuaRocks
--- or install from http://dkolf.de/src/dkjson-lua.fsl/home . Debian/Ubuntu
--- users can install the 'lua-dkjson' package.
-
--- Use JSON encoder
-local json = require('dkjson')
-
 local places = osm2pgsql.define_node_table('places', {
-    -- The jsonb column is handled by osm2pgsql just like a normal text
-    -- column. It is the job of the Lua script to put correct json there.
     { column = 'tags', type = 'jsonb' },
     { column = 'geom', type = 'point' },
 })
@@ -46,8 +37,13 @@ function osm2pgsql.process_node(object)
         object.tags.names = names
     end
 
+    -- The population should be stored as number, not as a string
+    if object.tags.population then
+        object.tags.population = tonumber(object.tags.population)
+    end
+
     places:add_row({
-        tags = json.encode(object.tags)
+        tags = object.tags
     })
 end
 
diff --git a/src/flex-table-column.cpp b/src/flex-table-column.cpp
@@ -22,7 +22,7 @@ struct column_type_lookup
     table_column_type type;
 };
 
-static std::array<column_type_lookup, 24> const column_types = {
+static std::array<column_type_lookup, 26> const column_types = {
     {{"sql", table_column_type::sql},
      {"text", table_column_type::text},
      {"boolean", table_column_type::boolean},
@@ -36,6 +36,8 @@ static std::array<column_type_lookup, 24> const column_types = {
      {"bigint", table_column_type::int8},
      {"real", table_column_type::real},
      {"hstore", table_column_type::hstore},
+     {"json", table_column_type::json},
+     {"jsonb", table_column_type::jsonb},
      {"direction", table_column_type::direction},
      {"geometry", table_column_type::geometry},
      {"point", table_column_type::point},
@@ -130,6 +132,10 @@ std::string flex_table_column_t::sql_type_name() const
         return "real";
     case table_column_type::hstore:
         return "hstore";
+    case table_column_type::json:
+        return "json";
+    case table_column_type::jsonb:
+        return "jsonb";
     case table_column_type::direction:
         return "int2";
     case table_column_type::geometry:
diff --git a/src/flex-table-column.hpp b/src/flex-table-column.hpp
@@ -31,6 +31,8 @@ enum class table_column_type : uint8_t
     real,
 
     hstore,
+    json,
+    jsonb,
 
     direction,
 
diff --git a/src/output-flex.cpp b/src/output-flex.cpp
@@ -30,6 +30,9 @@ extern "C"
 #include <lualib.h>
 }
 
+#include <rapidjson/stringbuffer.h>
+#include <rapidjson/writer.h>
+
 #include <cassert>
 #include <cstdio>
 #include <cstdlib>
@@ -318,6 +321,122 @@ static void write_double(db_copy_mgr_t<db_deleter_by_type_and_id_t> *copy_mgr,
     copy_mgr->add_column(value);
 }
 
+using json_writer_type = rapidjson::Writer<rapidjson::StringBuffer>;
+using table_register_type = std::vector<void const *>;
+
+/**
+ * Check that the value on the top of the Lua stack is a simple array.
+ * This means that all keys must be consecutive integers starting from 1.
+ */
+static bool is_lua_array(lua_State *lua_state)
+{
+    uint32_t n = 1;
+    lua_pushnil(lua_state);
+    while (lua_next(lua_state, -2) != 0) {
+        lua_pop(lua_state, 1); // remove value from stack
+        if (!lua_isinteger(lua_state, -1)) {
+            lua_pop(lua_state, 1);
+            return false;
+        }
+        int okay = 0;
+        auto const num = lua_tointegerx(lua_state, -1, &okay);
+        if (!okay || num != n++) {
+            lua_pop(lua_state, 1);
+            return false;
+        }
+    }
+
+    // An empty lua table could be both, we decide here that it is not stored
+    // as a JSON array but as a JSON object.
+    if (n == 1) {
+        return false;
+    }
+
+    return true;
+}
+
+static void write_json(json_writer_type *writer, lua_State *lua_state,
+                       table_register_type *tables);
+
+static void write_json_table(json_writer_type *writer, lua_State *lua_state,
+                             table_register_type *tables)
+{
+    void const *table_ptr = lua_topointer(lua_state, -1);
+    assert(table_ptr);
+    auto const it = std::find(tables->cbegin(), tables->cend(), table_ptr);
+    if (it != tables->cend()) {
+        throw std::runtime_error{"Loop detected in table"};
+    }
+    tables->push_back(table_ptr);
+
+    if (is_lua_array(lua_state)) {
+        writer->StartArray();
+        lua_pushnil(lua_state);
+        while (lua_next(lua_state, -2) != 0) {
+            write_json(writer, lua_state, tables);
+            lua_pop(lua_state, 1);
+        }
+        writer->EndArray();
+    } else {
+        writer->StartObject();
+        lua_pushnil(lua_state);
+        while (lua_next(lua_state, -2) != 0) {
+            int const ltype_key = lua_type(lua_state, -2);
+            if (ltype_key != LUA_TSTRING) {
+                throw std::runtime_error{
+                    "Incorrect data type '{}' as key."_format(
+                        lua_typename(lua_state, ltype_key))};
+            }
+            char const *const key = lua_tostring(lua_state, -2);
+            writer->Key(key);
+            write_json(writer, lua_state, tables);
+            lua_pop(lua_state, 1);
+        }
+        writer->EndObject();
+    }
+}
+
+static void write_json_number(json_writer_type *writer, lua_State *lua_state)
+{
+    int okay = 0;
+    auto const num = lua_tointegerx(lua_state, -1, &okay);
+    if (okay) {
+        writer->Int64(num);
+    } else {
+        writer->Double(lua_tonumber(lua_state, -1));
+    }
+}
+
+static void write_json(json_writer_type *writer, lua_State *lua_state,
+                       table_register_type *tables)
+{
+    assert(writer);
+    assert(lua_state);
+
+    int const ltype = lua_type(lua_state, -1);
+    switch (ltype) {
+    case LUA_TNIL:
+        writer->Null();
+        break;
+    case LUA_TBOOLEAN:
+        writer->Bool(lua_toboolean(lua_state, -1) != 0);
+        break;
+    case LUA_TNUMBER:
+        write_json_number(writer, lua_state);
+        break;
+    case LUA_TSTRING:
+        writer->String(lua_tostring(lua_state, -1));
+        break;
+    case LUA_TTABLE:
+        write_json_table(writer, lua_state, tables);
+        break;
+    default:
+        throw std::runtime_error{
+            "Invalid type '{}' for json/jsonb column."_format(
+                lua_typename(lua_state, ltype))};
+    }
+}
+
 void output_flex_t::write_column(
     db_copy_mgr_t<db_deleter_by_type_and_id_t> *copy_mgr,
     flex_table_column_t const &column)
@@ -467,6 +586,13 @@ void output_flex_t::write_column(
                 "Invalid type '{}' for hstore column."_format(
                     lua_typename(lua_state(), ltype))};
         }
+    } else if ((column.type() == table_column_type::json) ||
+               (column.type() == table_column_type::jsonb)) {
+        rapidjson::StringBuffer stream;
+        json_writer_type writer{stream};
+        table_register_type tables;
+        write_json(&writer, lua_state(), &tables);
+        copy_mgr->add_column(stream.GetString());
     } else if (column.type() == table_column_type::direction) {
         switch (ltype) {
         case LUA_TBOOLEAN:
diff --git a/tests/data/test_output_flex_types.lua b/tests/data/test_output_flex_types.lua
@@ -7,6 +7,7 @@ local test_table = osm2pgsql.define_node_table("nodes", {
     { column = 'tint8', type = 'int8' },
     { column = 'treal', type = 'real' },
     { column = 'thstr', type = 'hstore' },
+    { column = 'tjson', type = 'jsonb' },
     { column = 'tdirn', type = 'direction' },
     { column = 'tsqlt', type = 'varchar' },
 })
@@ -20,11 +21,11 @@ function osm2pgsql.process_node(object)
     end
     if test_type == 'boolean' then
         local row = { tbool = true, tint2 = true, tint4 = true,
-                      tint8 = true, tdirn = true }
+                      tint8 = true, tjson = true, tdirn = true }
         test_table:add_row(row)
 
         row = { tbool = false, tint2 = false, tint4 = false,
-                tint8 = false, tdirn = false }
+                tint8 = false, tjson = false, tdirn = false }
         test_table:add_row(row)
         return
     end
@@ -46,6 +47,7 @@ function osm2pgsql.process_node(object)
                 tint4 = n,
                 tint8 = n,
                 treal = n,
+                tjson = n,
                 tdirn = n,
                 tsqlt = n,
             }
@@ -129,8 +131,9 @@ function osm2pgsql.process_node(object)
         return
     end
     if test_type == 'table' then
-        test_table:add_row{ thstr = {} }
-        test_table:add_row{ thstr = { a = 'b', c = 'd' } }
+        local t = { a = 'b', c = 'd' }
+        test_table:add_row{ thstr = {}, tjson = {} }
+        test_table:add_row{ thstr = t, tjson = t }
         return
     end
     if test_type == 'table-hstore-fail' then
@@ -141,5 +144,24 @@ function osm2pgsql.process_node(object)
         test_table:add_row{ [object.tags.column] = {} }
         return
     end
+    if test_type == 'json' then
+        test_table:add_row{ tjson = {
+            astring = '123',
+            aninteger = 124,
+            anumber = 12.5,
+            atrue = true,
+            afalse = false,
+            anull = nil,
+            atable = { a = 'nested', tab = 'le' },
+            anarray = { 4, 3, 7 }
+        }}
+        return
+    end
+    if test_type == 'json-loop' then
+        local atable = { a = 'b' }
+        atable.c = atable
+        test_table:add_row{ tjson = atable }
+        return
+    end
 end
 
diff --git a/tests/test-output-flex-types.cpp b/tests/test-output-flex-types.cpp

-Original file line number
+Diff line change
     real,
     hstore,
 +    json,
 +    jsonb,
     direction,