Add not_null and create_only options to table columns definitions

joto · joto · commit 264c95883cf0 · 2020-02-18T16:54:22.000+01:00
diff --git a/docs/flex.md b/docs/flex.md
@@ -115,6 +115,12 @@ Instead of the above types you can use any SQL type you want. If you do that
 you have to supply the PostgreSQL string representation for that type when
 adding data to such columns (or Lua nil to set the column to `NULL`).
 
+When defining a column you can add the following options:
+* `not_null = true`: Make this a `NOT NULL` column.
+* `create_only = true`: Add the column to the `CREATE TABLE` command, but
+  do not try to fill this column when adding data. This can be useful for
+  `SERIAL` columns or when you want to fill in the column later yourself.
+
 ### Processing callbacks
 
 You are expected to define one or more of the following functions:
diff --git a/flex-config/data-types.lua b/flex-config/data-types.lua
@@ -5,7 +5,12 @@
 
 local highways = osm2pgsql.define_way_table('highways', {
     { column = 'name',     type = 'text' },
-    { column = 'type',     type = 'text' },
+    -- We always need a highway type, so we can declare the column as NOT NULL
+    { column = 'type',     type = 'text', not_null = true },
+
+    -- Add a SERIAL column and tell osm2pgsql not to fill it (PostgreSQL will
+    -- do that for us)
+    { column = 'id',       type = 'serial', create_only = true },
 
     -- type "direction" is special, see below
     { column = 'oneway',   type = 'direction' },
diff --git a/src/flex-table-column.cpp b/src/flex-table-column.cpp
@@ -143,7 +143,7 @@ std::string flex_table_column_t::sql_modifiers() const
 {
     std::string modifiers;
 
-    if ((m_flags & table_column_flags::not_null) != 0U) {
+    if (m_not_null) {
         modifiers += "NOT NULL ";
     }
 
diff --git a/src/flex-table-column.hpp b/src/flex-table-column.hpp
@@ -38,12 +38,6 @@ enum class table_column_type : uint8_t
     id_num
 };
 
-enum table_column_flags : uint8_t
-{
-    none = 0,
-    not_null = 1
-};
-
 /**
  * A column in a flex_table_t.
  */
@@ -56,8 +50,6 @@ class flex_table_column_t
 
     table_column_type type() const noexcept { return m_type; }
 
-    table_column_flags flags() const noexcept { return m_flags; }
-
     bool is_point_column() const noexcept
     {
         return (m_type == table_column_type::point) ||
@@ -85,11 +77,13 @@ class flex_table_column_t
 
     std::string const &type_name() const noexcept { return m_type_name; }
 
-    void set_not_null_constraint() noexcept
-    {
-        m_flags = static_cast<table_column_flags>(m_flags |
-                                                  table_column_flags::not_null);
-    }
+    bool not_null() const noexcept { return m_not_null; }
+
+    bool create_only() const noexcept { return m_create_only; }
+
+    void set_not_null(bool value = true) noexcept { m_not_null = value; }
+
+    void set_create_only(bool value = true) noexcept { m_create_only = value; }
 
     std::string sql_type_name(int srid) const;
     std::string sql_modifiers() const;
@@ -111,8 +105,11 @@ class flex_table_column_t
      */
     table_column_type m_type;
 
-    /// Flags like NOT NULL.
-    table_column_flags m_flags = table_column_flags::none;
+    /// NOT NULL constraint
+    bool m_not_null = false;
+
+    /// Column will be created but not filled by osm2pgsql.
+    bool m_create_only = false;
 };
 
 #endif // OSM2PGSQL_FLEX_TABLE_COLUMN_HPP
diff --git a/src/flex-table.cpp b/src/flex-table.cpp
@@ -75,7 +75,7 @@ flex_table_column_t &flex_table_t::add_column(std::string const &name,
 
     if (column.is_geometry_column()) {
         m_geom_column = m_columns.size() - 1;
-        column.set_not_null_constraint();
+        column.set_not_null();
     }
 
     return column;
@@ -127,10 +127,12 @@ std::string flex_table_t::build_sql_column_list() const
 
     std::string result;
     for (auto const &column : m_columns) {
-        result += '"';
-        result += column.name();
-        result += '"';
-        result += ',';
+        if (!column.create_only()) {
+            result += '"';
+            result += column.name();
+            result += '"';
+            result += ',';
+        }
     }
     result.resize(result.size() - 1);
 
diff --git a/src/lua-utils.cpp b/src/lua-utils.cpp
@@ -114,3 +114,24 @@ char const *luaX_get_table_string(lua_State *lua_state, char const *key,
     }
     return lua_tostring(lua_state, -1);
 }
+
+bool luaX_get_table_bool(lua_State *lua_state, char const *key, int table_index,
+                         char const *error_msg, bool default_value)
+{
+    assert(lua_state);
+    assert(key);
+    assert(error_msg);
+    lua_getfield(lua_state, table_index, key);
+    auto const type = lua_type(lua_state, -1);
+
+    if (type == LUA_TNIL) {
+        return default_value;
+    }
+
+    if (type == LUA_TBOOLEAN) {
+        return lua_toboolean(lua_state, -1);
+    }
+
+    throw std::runtime_error{
+        "{} must contain a '{}' boolean field"_format(error_msg, key)};
+}
diff --git a/src/lua-utils.hpp b/src/lua-utils.hpp
@@ -46,4 +46,7 @@ void luaX_add_table_array(lua_State *lua_state, char const *key,
 char const *luaX_get_table_string(lua_State *lua_state, char const *key,
                                   int table_index, char const *error_msg);
 
+bool luaX_get_table_bool(lua_State *lua_state, char const *key, int table_index,
+                         char const *error_msg, bool default_value);
+
 #endif // OSM2PGSQL_FLEX_LUA_HPP
diff --git a/src/output-flex.cpp b/src/output-flex.cpp
@@ -326,6 +326,9 @@ void output_flex_t::write_row(flex_table_t *table, osmium::item_type id_type,
     auto *copy_mgr = table->copy_mgr();
 
     for (auto const &column : *table) {
+        if (column.create_only()) {
+            continue;
+        }
         if (column.type() == table_column_type::id_type) {
             copy_mgr->add_column(type_to_char(id_type));
         } else if (column.type() == table_column_type::id_num) {
@@ -509,7 +512,7 @@ void output_flex_t::setup_id_columns(flex_table_t *table)
         }
         lua_pop(lua_state(), 1); // type_column
         auto &column = table->add_column(type_column_name, "id_type");
-        column.set_not_null_constraint();
+        column.set_not_null();
         table->set_id_type(osmium::item_type::undefined);
     } else {
         throw std::runtime_error{"Unknown ids type: " + type};
@@ -520,7 +523,7 @@ void output_flex_t::setup_id_columns(flex_table_t *table)
     check_name(name, "column");
 
     auto &column = table->add_column(name, "id_num");
-    column.set_not_null_constraint();
+    column.set_not_null();
     lua_pop(lua_state(), 3); // id_column, type, ids
 }
 
@@ -551,9 +554,15 @@ void output_flex_t::setup_flex_table_columns(flex_table_t *table)
             luaX_get_table_string(lua_state(), "column", -2, "Column entry");
         check_name(name, "column");
 
-        table->add_column(name, type);
+        auto &column = table->add_column(name, type);
+
+        column.set_not_null(luaX_get_table_bool(lua_state(), "not_null", -3,
+                                                "Entry 'not_null'", false));
+
+        column.set_create_only(luaX_get_table_bool(
+            lua_state(), "create_only", -4, "Entry 'create_only'", false));
 
-        lua_pop(lua_state(), 3); // column, type, table
+        lua_pop(lua_state(), 5); // create_only, not_null, column, type, table
         ++num_columns;
     }
 
@@ -670,6 +679,8 @@ int output_flex_t::table_columns()
                            column.sql_type_name(table.srid()).c_str());
         luaX_add_table_str(lua_state(), "sql_modifiers",
                            column.sql_modifiers().c_str());
+        luaX_add_table_bool(lua_state(), "not_null", column.not_null());
+        luaX_add_table_bool(lua_state(), "create_only", column.create_only());
 
         lua_rawset(lua_state(), -3);
     }

Original file line number	Diff line number	Diff line change
`@@ -143,7 +143,7 @@ std::string flex_table_column_t::sql_modifiers() const`
`143`	`143`	`{`
`144`	`144`	`std::string modifiers;`
`145`	`145`
`146`		`- if ((m_flags & table_column_flags::not_null) != 0U) {`
	`146`	`+ if (m_not_null) {`
`147`	`147`	`modifiers += "NOT NULL ";`
`148`	`148`	`}`
`149`	`149`