diff options
Diffstat (limited to 'lib')
| -rw-r--r-- | lib/rbot/journal/mongo.rb | 120 | ||||
| -rw-r--r-- | lib/rbot/journal/postgres.rb | 66 |
2 files changed, 176 insertions, 10 deletions
diff --git a/lib/rbot/journal/mongo.rb b/lib/rbot/journal/mongo.rb new file mode 100644 index 00000000..24e9cfcc --- /dev/null +++ b/lib/rbot/journal/mongo.rb @@ -0,0 +1,120 @@ +# encoding: UTF-8 +#-- vim:sw=2:et +#++ +# +# :title: journal backend for mongoDB + +require 'mongo' +require 'json' + +module Irc +class Bot +module Journal + + module Storage + + class MongoStorage < AbstractStorage + attr_reader :client + + def initialize(opts={}) + Mongo::Logger.logger.level = Logger::WARN + @uri = opts[:uri] || 'mongodb://127.0.0.1:27017/rbot' + @client = Mongo::Client.new(@uri) + @collection = @client['journal'] + log 'journal storage: mongodb connected to ' + @uri + + drop if opts[:drop] + @collection.indexes.create_one({topic: 1}) + end + + def ensure_index(key) + @collection.indexes.create_one({'payload.'+key => 1}) + end + + def insert(m) + @collection.insert_one({ + '_id' => m.id, + 'topic' => m.topic, + 'timestamp' => m.timestamp, + 'payload' => m.payload + }) + end + + def find(query=nil, limit=100, offset=0) + query_cursor(query).skip(offset).limit(limit).map do |document| + JournalMessage.new(id: document['_id'], timestamp: document['timestamp'].localtime, + topic: document['topic'], payload: document['payload'].to_h) + end + end + + # returns the number of messages that match the query + def count(query=nil) + query_cursor(query).count + end + + def remove(query=nil) + query_cursor(query).delete_many + end + + def drop + @collection.drop + end + + def query_cursor(query) + unless query + return @collection.find() + end + + query_and = [] + + # ID query OR condition + unless query.id.empty? + query_and << { + '$or' => query.id.map { |_id| + {'_id' => _id} + } + } + end + + unless query.topic.empty? + query_and << { + '$or' => query.topic.map { |topic| + if topic.include?('*') + pattern = topic.gsub('.', '\.').gsub('*', '.*') + {'topic' => {'$regex' => pattern}} + else + {'topic' => topic} + end + } + } + end + + if query.timestamp[:from] or query.timestamp[:to] + where = {} + if query.timestamp[:from] + where['$gte'] = query.timestamp[:from] + end + if query.timestamp[:to] + where['$lte'] = query.timestamp[:to] + end + query_and << {'timestamp' => where} + end + + unless query.payload.empty? + query_and << { + '$or' => query.payload.map { |key, value| + key = 'payload.' + key + {key => value} + } + } + end + + @collection.find({ + '$and' => query_and + }) + end + end + end +end # Journal +end # Bot +end # Irc diff --git a/lib/rbot/journal/postgres.rb b/lib/rbot/journal/postgres.rb index 52f5fb36..e63aefee 100644 --- a/lib/rbot/journal/postgres.rb +++ b/lib/rbot/journal/postgres.rb @@ -7,6 +7,28 @@ require 'pg' require 'json' +# as a replacement for CREATE INDEX IF NOT EXIST that is not in postgres. +# define function to be able to create an index in case it doesnt exist: +# source: http://stackoverflow.com/a/26012880 +CREATE_INDEX = <<-EOT +CREATE OR REPLACE FUNCTION create_index(table_name text, index_name text, column_name text) RETURNS void AS $$ +declare + l_count integer; +begin + select count(*) + into l_count + from pg_indexes + where schemaname = 'public' + and tablename = lower(table_name) + and indexname = lower(index_name); + + if l_count = 0 then + execute 'create index ' || index_name || ' on ' || table_name || '(' || column_name || ')'; + end if; +end; +$$ LANGUAGE plpgsql; +EOT + module Irc class Bot module Journal @@ -17,9 +39,10 @@ module Journal attr_reader :conn def initialize(opts={}) - @uri = opts[:uri] || 'postgresql://localhost/rbot_journal' + @uri = opts[:uri] || 'postgresql://localhost/rbot' @conn = PG.connect(@uri) @conn.exec('set client_min_messages = warning') + @conn.exec(CREATE_INDEX) @version = @conn.exec('SHOW server_version;')[0]['server_version'] @version.gsub!(/^(\d+\.\d+)$/, '\1.0') @@ -35,6 +58,7 @@ module Journal drop if opts[:drop] create_table + create_index('topic_index', 'topic') end def create_table @@ -46,6 +70,23 @@ module Journal payload %s NOT NULL)' % [@jsonb ? 'JSONB' : 'JSON']) end + def create_index(index_name, column_name) + debug 'journal postges backend: create index %s for %s' % [ + index_name, column_name] + @conn.exec_params('SELECT create_index($1, $2, $3)', [ + 'journal', index_name, column_name]) + end + + def create_payload_index(key) + index_name = 'idx_payload_' + key.gsub('.', '_') + column = sql_payload_selector(key) + create_index(index_name, column) + end + + def ensure_index(key) + create_payload_index(key) + end + def insert(m) @conn.exec_params('INSERT INTO journal VALUES ($1, $2, $3, $4);', [m.id, m.topic, m.timestamp, JSON.generate(m.payload)]) @@ -95,6 +136,19 @@ module Journal @conn.exec('DROP TABLE journal;') rescue nil end + def sql_payload_selector(key) + selector = 'payload' + k = key.to_s.split('.') + k.each_index { |i| + if i >= k.length-1 + selector += '->>\'%s\'' % [@conn.escape_string(k[i])] + else + selector += '->\'%s\'' % [@conn.escape_string(k[i])] + end + } + selector + end + def query_to_sql(query) params = [] placeholder = Proc.new do |value| @@ -142,15 +196,7 @@ module Journal unless query.payload.empty? list = [] query.payload.each_pair do |key, value| - selector = 'payload' - k = key.to_s.split('.') - k.each_index { |i| - if i >= k.length-1 - selector += '->>\'%s\'' % [@conn.escape_string(k[i])] - else - selector += '->\'%s\'' % [@conn.escape_string(k[i])] - end - } + selector = sql_payload_selector(key) list << selector + ' = ' + placeholder.call(value) end sql[:list] << { |