rana-cli/wakatime/offlinequeue.py

# -*- coding: utf-8 -*-
"""
    wakatime.offlinequeue
    ~~~~~~~~~~~~~~~~~~~~~

    Queue for saving heartbeats while offline.

    :copyright: (c) 2014 Alan Hamlett.
    :license: BSD, see LICENSE for more details.
"""


import logging
import os
from time import sleep

from .compat import json
from .heartbeat import Heartbeat


try:
    import sqlite3
    HAS_SQL = True
except ImportError:  # pragma: nocover
    HAS_SQL = False


log = logging.getLogger('WakaTime')


class Queue(object):
    db_file = '.wakatime.db'
    table_name = 'heartbeat_2'

    args = None
    configs = None

    def __init__(self, args, configs):
        self.args = args
        self.configs = configs

    def connect(self):
        conn = sqlite3.connect(self._get_db_file(), isolation_level=None)
        c = conn.cursor()
        c.execute('''CREATE TABLE IF NOT EXISTS {0} (
            id text,
            heartbeat text)
        '''.format(self.table_name))
        return (conn, c)

    def push(self, heartbeat):
        if not HAS_SQL:
            return
        try:
            conn, c = self.connect()
            data = {
                'id': heartbeat.get_id(),
                'heartbeat': heartbeat.json(),
            }
            c.execute('INSERT INTO {0} VALUES (:id,:heartbeat)'.format(self.table_name), data)
            conn.commit()
            conn.close()
        except sqlite3.Error:
            log.traceback()

    def pop(self):
        if not HAS_SQL:
            return None
        tries = 3
        wait = 0.1
        try:
            conn, c = self.connect()
        except sqlite3.Error:
            log.traceback(logging.DEBUG)
            return None

        heartbeat = None

        loop = True
        while loop and tries > -1:
            try:
                c.execute('BEGIN IMMEDIATE')
                c.execute('SELECT * FROM {0} LIMIT 1'.format(self.table_name))
                row = c.fetchone()
                if row is not None:
                    id = row[0]
                    heartbeat = Heartbeat(json.loads(row[1]), self.args, self.configs, _clone=True)
                    c.execute('DELETE FROM {0} WHERE id=?'.format(self.table_name), [id])
                conn.commit()
                loop = False
            except sqlite3.Error:
                log.traceback(logging.DEBUG)
                sleep(wait)
                tries -= 1
        try:
            conn.close()
        except sqlite3.Error:
            log.traceback(logging.DEBUG)
        return heartbeat

    def push_many(self, heartbeats):
        for heartbeat in heartbeats:
            self.push(heartbeat)

    def pop_many(self, limit=None):
        if limit is None:
            limit = 5

        heartbeats = []

        count = 0
        while count < limit:
            heartbeat = self.pop()
            if not heartbeat:
                break
            heartbeats.append(heartbeat)
            count += 1
            if count % 10 == 0:
                yield heartbeats
                heartbeats = []

        if heartbeats:
            yield heartbeats

    def _get_db_file(self):
        home = '~'
        if os.environ.get('WAKATIME_HOME'):
            home = os.environ.get('WAKATIME_HOME')
        return os.path.join(os.path.expanduser(home), '.wakatime.db')