sqdb-build: write out category section

store the names of categories to database

2 files changed, 51 insertions, 4 deletions

diff --git a/lua-squarkdb.c b/lua-squarkdb.c
index 80a0d32..dbac6d0 100644
--- a/lua-squarkdb.c
+++ b/lua-squarkdb.c
@@ -266,6 +266,40 @@ static int Lsqdb_map_strings(lua_State *L)
 	return 0;
 }
 
+static int Lsqdb_write_section(lua_State *L)
+{
+	struct sqdb *db;
+	uint32_t *ptr;
+	const char *section;
+	int i, tbllen, si = -1;
+
+	db = Lsqdb_checkarg(L, 1);
+	section = luaL_checkstring(L, 2);
+	luaL_checktype(L, 3, LUA_TTABLE);
+	tbllen = lua_objlen(L, 3);
+
+	for (i = 0; sqdb_section_names[i] && i < SQDB_SECTION_MAX; i++) {
+		if (strcmp(sqdb_section_names[i], section) == 0) {
+			si = 0;
+			break;
+		}
+	}
+	if (si < 0)
+		luaL_error(L, "Section name '%s' is invalid", section);
+
+	ptr = sqdb_section_create(db, i, sizeof(uint32_t) * tbllen);
+	if (ptr == NULL)
+		luaL_error(L, "Failed to create section '%s'", section);
+
+	for (i = 0; i < tbllen; i++) {
+		lua_rawgeti(L, 3, i + 1);
+		ptr[i] = lua_tointeger(L, -1);
+		lua_pop(L, 1);
+	}
+
+	return 0;
+}
+
 static const luaL_reg sqdb_meta_methods[] = {
 	{ "__gc",		Lsqdb_destroy },
 	{ NULL,			NULL }
@@ -278,6 +312,7 @@ static const luaL_reg squarkdb_methods[] = {
 	{ "create_index",	Lsqdb_create_index },
 	{ "assign_index",	Lsqdb_assign_index },
 	{ "map_strings",	Lsqdb_map_strings },
+	{ "write_section",	Lsqdb_write_section },
 	{ NULL,			NULL }
 };
 
diff --git a/sqdb-build.lua b/sqdb-build.lua
index 6cde349..fce1e7b 100755
--- a/sqdb-build.lua
+++ b/sqdb-build.lua
@@ -7,6 +7,7 @@ local all_domains = {}
 local all_ips = {}
 
 local all_categories = {}
+local all_categories_by_id = {}
 local num_categories = 0
 
 local strfind = string.find
@@ -42,9 +43,13 @@ local function get_category(category_text)
 	cat = all_categories[category_text]
 	if cat ~= nil then return cat end
 
-	num_categories = num_categories + 1
+	-- start category ID's from zero
 	cat = { desc=category_text, id=num_categories }
 	all_categories[category_text] = cat
+	num_categories = num_categories + 1
+
+	-- but index them from one
+	all_categories_by_id[num_categories] = cat
 
 	account_string(category_text)
 
@@ -169,11 +174,10 @@ local function enum_paths(cb, category, path, data)
 end
 
 local function enum_tree(cb, category, dns, data)
-	local cdns, cdata, fdns
-	local cat = data.category or category
+	local cdns, cdata, fdns, cat
 
 	if data.paths ~= nil then
-		enum_paths(cb, cat, dns, data.paths)
+		enum_paths(cb, category, dns, data.paths)
 	end
 	if data.children ~= nil then
 		for cdns, cdata in pairs(data.children) do
@@ -182,6 +186,7 @@ local function enum_tree(cb, category, dns, data)
 			else
 				fdns = cdns
 			end
+			cat = cdata.category or category
 			cb(fdns, dns, cdns, cat, data.children, data.paths)
 			enum_tree(cb, cat, fdns, cdata)
 		end
@@ -276,6 +281,7 @@ local function load_lists(conffile, part)
 end
 
 -- start by reading in all classification data
+get_category("unknown")
 load_lists("lists.conf", "domains")
 prune_tree(all_domains, nil)
 load_lists("lists.conf", "urls")
@@ -288,6 +294,12 @@ num_entries = db:generate_hash(function() enum_all(coroutine.yield) end)
 -- write string literals
 db:map_strings(all_strings)
 
+-- map category names and write the category section out
+for id, cdata in ipairs(all_categories_by_id) do
+	all_categories_by_id[id] = all_strings[cdata.desc]
+end
+db:write_section("categories", all_categories_by_id)
+
 -- create master index
 db:create_index(num_entries)
 enum_all(