File

-- Based on mod_mam_sql
-- Copyright (C) 2011-2012 Kim Alvefur
--
-- This file is MIT/X11 licensed.

local st = require "util.stanza";
local jid_bare = require "util.jid".bare;
local jid_split = require "util.jid".split;

local serialize = require"util.json".encode, require"util.json".decode;
local tostring = tostring;
local time_now = os.time;


local table_name = module:get_option("message_log_sql_table", pcall(require, "util.cache") and "messages" or "prosodyarchive");

local sql, setsql, getsql = {};
do -- SQL stuff
local connection;
local resolve_relative_path = require "core.configmanager".resolve_relative_path;
local params = module:get_option("message_log_sql", module:get_option("sql"));


local function test_connection()
	if not connection then return nil; end
	if connection:ping() then
		return true;
	else
		module:log("debug", "Database connection closed");
		connection = nil;
	end
end
local function connect()
	if not test_connection() then
		prosody.unlock_globals();
		local dbh, err = DBI.Connect(
			params.driver, params.database,
			params.username, params.password,
			params.host, params.port
		);
		prosody.lock_globals();
		if not dbh then
			module:log("debug", "Database connection failed: %s", tostring(err));
			return nil, err;
		end
		module:log("debug", "Successfully connected to database");
		dbh:autocommit(false); -- don't commit automatically
		connection = dbh;

	end
	return connection;
end


do -- process options to get a db connection
	local ok;
	prosody.unlock_globals();
	ok, DBI = pcall(require, "DBI");
	if not ok then
		package.loaded["DBI"] = {};
		module:log("error", "Failed to load the LuaDBI library for accessing SQL databases: %s", DBI);
		module:log("error", "More information on installing LuaDBI can be found at http://prosody.im/doc/depends#luadbi");
	end
	prosody.lock_globals();
	if not ok or not DBI.Connect then
		return; -- Halt loading of this module
	end

	params = params or { driver = "SQLite3" };

	if params.driver == "SQLite3" then
		params.database = resolve_relative_path(prosody.paths.data or ".", params.database or "prosody.sqlite");
	end

	assert(params.driver and params.database, "Both the SQL driver and the database need to be specified");

	assert(connect());

end

function getsql(sql, ...)
	if params.driver == "PostgreSQL" then
		sql = sql:gsub("`", "\"");
	end
	-- do prepared statement stuff
	local stmt, err = connection:prepare(sql);
	if not stmt and not test_connection() then error("connection failed"); end
	if not stmt then module:log("error", "QUERY FAILED: %s %s", err, debug.traceback()); return nil, err; end
	-- run query
	local ok, err = stmt:execute(...);
	if not ok and not test_connection() then error("connection failed"); end
	if not ok then return nil, err; end

	return stmt;
end
function setsql(sql, ...)
	local stmt, err = getsql(sql, ...);
	if not stmt then return stmt, err; end
	return stmt:affected();
end
function sql.rollback(...)
	if connection then connection:rollback(); end -- FIXME check for rollback error?
	return ...;
end
function sql.commit(...)
	local ok, err = connection:commit();
	if not ok then
		module:log("error", "SQL commit failed: %s", tostring(err));
		return nil, "SQL commit failed: "..tostring(err);
	end
	return ...;
end

end

-- Handle messages
local function message_handler(event, c2s)
	local origin, stanza = event.origin, event.stanza;
	local orig_type = stanza.attr.type or "normal";
	local orig_to = stanza.attr.to;
	local orig_from = stanza.attr.from;

	if not orig_from and c2s then
		orig_from = origin.full_jid;
	end
	orig_to = orig_to or orig_from; -- Weird corner cases

	-- Don't store messages of these types
	if orig_type == "error"
	or orig_type == "headline"
	or orig_type == "groupchat"
	or not stanza:get_child("body") then
		return;
		-- TODO Maybe headlines should be configurable?
	end

	local store_user, store_host = jid_split(c2s and orig_from or orig_to);
	local target_jid = c2s and orig_to or orig_from;
	local target_bare = jid_bare(target_jid);
	local _, _, target_resource = jid_split(target_jid);

	--local id = uuid();
	local when = time_now();
	-- And stash it
	local ok, err = setsql([[
	INSERT INTO `]]..table_name..[[`
	(`host`, `user`, `store`, `when`, `with`, `resource`, `stanza`)
	VALUES (?, ?, ?, ?, ?, ?, ?);
	]], store_host, store_user, "message_log", when, target_bare, target_resource, serialize(st.preserialize(stanza)))
	if ok then
		sql.commit();
	else
		module:log("error", "SQL error: %s", err);
		sql.rollback();
	end
end

local function c2s_message_handler(event)
	return message_handler(event, true);
end

-- Stanzas sent by local clients
module:hook("pre-message/bare", c2s_message_handler, 2);
module:hook("pre-message/full", c2s_message_handler, 2);
-- Stanszas to local clients
module:hook("message/bare", message_handler, 2);
module:hook("message/full", message_handler, 2);

-- In the telnet console, run:
-- >hosts["this host"].modules.mam_sql.environment.create_sql()
function create_sql()
	local stm = assert(getsql([[
	CREATE TABLE `]]..table_name..[[` (
		`host` TEXT,
		`user` TEXT,
		`store` TEXT,
		`id` INTEGER PRIMARY KEY AUTOINCREMENT,
		`when` INTEGER,
		`with` TEXT,
		`resource` TEXT,
		`stanza` TEXT
	);
	CREATE INDEX `hus` ON `]]..table_name..[[` (`host`, `user`, `store`);
	CREATE INDEX `with` ON `]]..table_name..[[` (`with`);
	CREATE INDEX `thetime` ON `]]..table_name..[[` (`when`);
	]]));
	stm:execute();
	sql.commit();
end