rana-cli/wakatime/offlinequeue.py
2018-09-20 10:14:25 -07:00

129 lines
3.2 KiB
Python

# -*- 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')