Source

rethinklog / rethinklog.py

Full commit
"""
RethinkLog - Log messages to a rethinkdb instance.

Copyright (C) 2013 Josh VanderLinden <arch@cloudlery.com>

CHANGELOG {{{

* 0.0.3 - 2013-02-07

    - general code reorganization
    - configuration changes withour requiring a reload

* 0.0.2 - 2013-02-07

    - bulk inserts
    - disconnect from rethinkdb on unload

* 0.0.1 - 2013-02-07

    - initial plugin

}}}

TODOS {{{

* handle error cases when RethinkDB is inaccessible
* handle errors inserting records (retries?)
* show most recent logs for a given buffer when opened

}}}

"""

from collections import deque
from os.path import basename, splitext

from rethinkdb import r

from weechat import buffer_get_string as get, config_get_plugin as cfg
import weechat

__name__ = 'RethinkLog'
__author__ = 'codekoala'
__version__ = '0.0.3'
__license__ = 'GPL3'
__description__ = 'Log messages to a RethinkDB cluster'
__modname__ = basename(splitext(__file__)[0])

OK = weechat.WEECHAT_RC_OK
ERR = weechat.WEECHAT_RC_ERROR


class RethinkLogger(object):
    """Log messages to a RethinkDB cluster"""

    DEFAULTS = dict(
        host='localhost',
        port='28015',
        db='weechat',
        table='messages',
        sync_freq='10000',
        create_db='on',
        create_table='on',
    )

    def __init__(self):
        self.sync_hook = None
        self.cnx = None
        self.queue = deque()

        # initialize config
        for o, v in RethinkLogger.DEFAULTS.items():
            if not cfg(o):
                weechat.config_set_plugin(o, v)

        weechat.hook_print('', '', '', 1, 'log', '')
        weechat.hook_config('plugins.var.python.%s.*' % __modname__,
                            'reconfigure', '')

        self.start_timer()

    @property
    def r(self):
        """
        Return a cursor to the database. Connect to RethinkDB if necessary.
        """

        if self.cnx is None:
            host = cfg('host')
            port = int(cfg('port'))
            self.p('Connecting to RethinkDB: %s:%s' % (host, port))
            self.cnx = r.connect(host, port)

        # create the database?
        db_name = cfg('db')
        if db_name not in r.db_list().run() and cfg('create_db') == 'on':
            self.p('Creating RethinkLog database: %s' % (db_name,))
            r.db_create(db_name).run()

        db = r.db(db_name)

        # create the table?
        table_name = cfg('table')
        if table_name not in db.table_list().run() and cfg('create_table') == 'on':
            self.p('Creating RethinkLog table: %s' % (table_name,))
            db.table_create(table_name).run()

        return db.table(table_name)

    def p(self, msg, buf=''):
        """Print a message in weechat"""

        weechat.prnt(buf, msg)

    def log(self, data, buf, time, tags, display, hilight, prefix, msg):
        """Queue a message to be written to the database."""

        return_code = OK

        try:
            buffer_name = get(buf, 'short_name') or get(buf, 'name')

            record = dict(
                data=data.strip(),
                buffer=buffer_name.strip(),
                timestamp=int(time),
                tags=[t.strip() for t in tags.split(',') if t],
                display=bool(int(display)),
                hilight=bool(int(hilight)),
                prefix=prefix.strip(),
                message=msg.strip().decode('utf-8')
            )

            self.queue.append(record)
        except Exception as e:
            self.p('Failed to queue message for insertion: %s' % str(e))
            return_code = ERR

        return return_code

    def flush(self, data, remaining_calls):
        """Send multiple records to RethinkDB at one time."""

        return_code = OK

        try:
            self.r.insert(self.queue).run()
        except Exception as e:
            self.p('Failed to log messages to rethinkdb: %s' % str(e))
            return_code = ERR
        else:
            self.queue.clear()

        return return_code

    def reconfigure(self, data, option, value):
        """Handle a configuration change without reloading"""

        # disconnect from the database when the host or port change
        if option in ('host', 'port'):
            self.p('RethinkLog database information changed')
            self.close()

        if option == 'sync_freq':
            self.p('RethinkLog sync interval updated')
            weechat.unhook(self.sync_hook)
            self.start_timer()

        return OK

    def start_timer(self):
        """
        Start a timer to automatically synchronize messages with the
        database.
        """

        interval = int(cfg('sync_freq'))
        self.p('Starting RethinkLog timer to sync every %sms' % (interval,))
        self.sync_hook = weechat.hook_timer(interval, 0, 0, 'flush', '')

    def close(self):
        """Disconnect from the RethinkDB"""

        if self.cnx:
            self.cnx.close()

        return OK


weechat.register(__name__, __author__, __version__, __license__,
                 __description__, 'close', '')


# instantiate and expose useful methods
plugin = RethinkLogger()

log = plugin.log
flush = plugin.flush
reconfigure = plugin.reconfigure
close = plugin.close

# vim:et ts=4 sw=4 fdm=marker: