plugins/mod_storage_sql_ejabberd.lua

1
local setmetatable = setmetatable;
2
local error = error;
3
local unpack = unpack;
4
local module = module;
5
local tostring = tostring;
6
local pairs, next = pairs, next;
7
local prosody = prosody;
8
local assert = assert;
9
local require = require;
10
local st = require "util.stanza";
11
local DBI = require "DBI";
12
 
13
-- connect to db
14
local params = module:get_option("sql_ejabberd") or error("No sql_ejabberd config option");
15
local database;
16
do
17
	module:log("debug", "Opening database: %s", "dbi:"..params.driver..":"..params.database);
18
	prosody.unlock_globals();
19
	local dbh, err = DBI.Connect(
20
		params.driver, params.database,
21
		params.username, params.password,
22
		params.host, params.port
23
	);
24
	prosody.lock_globals();
25
	assert(dbh, err);
26
	dbh:autocommit(true);
27
	database = dbh;
28
end
29
 
30
-- initialize db
31
local ejabberd_init = module:require("ejabberd_init");
32
ejabberd_init.init(database);
33
 
34
local sqlcache = {};
35
local function prepare(sql)
36
	module:log("debug", "query: %s", sql);
37
	local err;
38
	local r = sqlcache[sql];
39
	if not r then
40
		r, err = database:prepare(sql);
41
		if not r then error("Unable to prepare SQL statement: "..err); end
42
		sqlcache[sql] = r;
43
	end
44
	return r;
45
end
46
 
47
local _parse_xml = module:require("xmlparse");
48
local function parse_xml(str)
49
	local s = _parse_xml(str);
50
	if s and not s.gsub then
51
		return st.preserialize(s);
52
	end
53
end
54
local function unparse_xml(s)
55
	return tostring(st.deserialize(s));
56
end
57
 
58
 
59
local handlers = {};
60
 
61
handlers.accounts = {
62
	get = function(self, user)
63
		local select = self:query("select password from users where username=? and host=?", user, self.host);
64
		local row = select and select:fetch();
65
		if row then return { password = row[1] }; end
66
	end;
67
	set = function(self, user, data)
68
		if data and data.password then
69
			return self:modify("update users set password=? where username=? and host=?", data.password, user, self.host)
70
				or self:modify("insert into users (username, host, password) values (?, ?, ?)", user, self.host, data.password);
71
		else
72
			return self:modify("delete from users where username=? and host=?", user, self.host);
73
		end
74
	end;
75
};
76
handlers.vcard = {
77
	get = function(self, user)
78
		local select = self:query("select vcard from vcard where username=? and host=?", user, self.host);
79
		local row = select and select:fetch();
80
		if row then return parse_xml(row[1]); end
81
	end;
82
	set = function(self, user, data)
83
		if data then
84
			data = unparse_xml(data);
85
			return self:modify("update vcard set vcard=? where username=? and host=?", data, user, self.host)
86
				or self:modify("insert into vcard (username, host, vcard) values (?, ?, ?)", user, self.host, data);
87
		else
88
			return self:modify("delete from vcard where username=? and host=?", user, self.host);
89
		end
90
	end;
91
};
92
handlers.private = {
93
	get = function(self, user)
94
		local select = self:query("select namespace,data from private_storage where username=? and host=?", user, self.host);
95
		if select then
96
			local data = {};
97
			for row in select:rows() do
98
				data[row[1]] = parse_xml(row[2]);
99
			end
100
			return data;
101
		end
102
	end;
103
	set = function(self, user, data)
104
		if data then
105
			self:modify("delete from private_storage where username=? and host=?", user, self.host);
106
			for namespace,text in pairs(data) do
107
				self:modify("insert into private_storage (username, host, namespace, data) values (?, ?, ?, ?)", user, self.host, namespace, unparse_xml(text));
108
			end
109
			return true;
110
		else
111
			return self:modify("delete from private_storage where username=? and host=?", user, self.host);
112
		end
113
	end;
114
	-- TODO map_set, map_get
115
};
116
local subscription_map = { N = "none", B = "both", F = "from", T = "to" };
117
local subscription_map_reverse = { none = "N", both = "B", from = "F", to = "T" };
118
handlers.roster = {
119
	get = function(self, user)
120
		local select = self:query("select jid,nick,subscription,ask,server,subscribe,type from rosterusers where username=?", user);
121
		if select then
122
			local roster = { pending = {} };
123
			for row in select:rows() do
124
				local jid,nick,subscription,ask,server,subscribe,typ = unpack(row);
125
				local item = { groups = {} };
126
				if nick == "" then nick = nil; end
127
				item.nick = nick;
128
				item.subscription = subscription_map[subscription];
129
				if ask == "N" then ask = nil;
130
				elseif ask == "O" then ask = "subscribe"
131
				elseif ask == "I" then roster.pending[jid] = true; ask = nil;
132
				elseif ask == "B" then roster.pending[jid] = true; ask = "subscribe";
133
				else module:log("debug", "bad roster_item.ask: %s", ask); ask = nil; end
134
				item.ask = ask;
135
				roster[jid] = item;
136
			end
137
			
138
			select = self:query("select jid,grp from rostergroups where username=?", user);
139
			if select then
140
				for row in select:rows() do
141
					local jid,grp = unpack(row);
142
					if roster[jid] then roster[jid].groups[grp] = true; end
143
				end
144
			end
145
			select = self:query("select version from roster_version where username=?", user);
146
			local row = select and select:fetch();
147
			if row then
148
				roster[false] = { version = row[1]; };
149
			end
150
			return roster;
151
		end
152
	end;
153
	set = function(self, user, data)
154
		if data and next(data) ~= nil then
155
			self:modify("delete from rosterusers where username=?", user);
156
			self:modify("delete from rostergroups where username=?", user);
157
			self:modify("delete from roster_version where username=?", user);
158
			local done = {};
159
			local pending = data.pending or {};
160
			for jid,item in pairs(data) do
161
				if jid and jid ~= "pending" then
162
					local subscription = subscription_map_reverse[item.subscription];
163
					local ask;
164
					if pending[jid] then
165
						if item.ask then ask = "B"; else ask = "I"; end
166
					else
167
						if item.ask then ask = "O"; else ask = "N"; end
168
					end
169
					local r = self:modify("insert into rosterusers (username,jid,nick,subscription,ask,askmessage,server,subscribe) values (?, ?, ?, ?, ?, '', '', '')", user, jid, item.nick or "", subscription, ask);
170
					if not r then module:log("debug", "--- :( %s", tostring(r)); end
171
					done[jid] = true;
172
					for group in pairs(item.groups) do
173
						self:modify("insert into rostergroups (username,jid,grp) values (?, ?, ?)", user, jid, group);
174
					end
175
				end
176
			end
177
			for jid in pairs(pending) do
178
				if not done[jid] then
179
					self:modify("insert into rosterusers (username,jid,nick,subscription,ask,askmessage,server,subscribe) values (?, ?, ?, ?, ?. ''. ''. '')", user, jid, "", "N", "I");
180
				end
181
			end
182
			local version = data[false] and data[false].version;
183
			if version then
184
				self:modify("insert into roster_version (username,version) values (?, ?)", user, version);
185
			end
186
			return true;
187
		else
188
			self:modify("delete from rosterusers where username=?", user);
189
			self:modify("delete from rostergroups where username=?", user);
190
			self:modify("delete from roster_version where username=?", user);
191
		end
192
	end;
193
};
194
 
195
-----------------------------
196
local driver = {};
197
driver.__index = driver;
198
 
199
function driver:query(sql, ...)
200
	local stmt,err = prepare(sql);
201
	if not stmt then
202
		module:log("error", "Failed to prepare SQL [[%s]], error: %s", sql, err);
203
		return nil, err;
204
	end
205
	local ok, err = stmt:execute(...);
206
	if not ok then
207
		module:log("error", "Failed to execute SQL [[%s]], error: %s", sql, err);
208
		return nil, err;
209
	end
210
	return stmt;
211
end
212
function driver:modify(sql, ...)
213
	local stmt, err = self:query(sql, ...);
214
	if stmt and stmt:affected() > 0 then return stmt; end
215
	return nil, err;
216
end
217
 
218
function driver:open(datastore, typ)
219
	local instance = setmetatable({ host = module.host, datastore = datastore }, self);
220
	local handler = handlers[datastore];
221
	if not handler then return nil; end
222
	for key,val in pairs(handler) do
223
		instance[key] = val;
224
	end
225
	if instance.init then instance:init(); end
226
	return instance;
227
end
228
 
229
-----------------------------
230
 
231
module:add_item("data-driver", driver);