rana-cli/wakatime/offlinequeue.py

131 lines
3.3 KiB
Python
Raw Normal View History

2014-05-26 00:14:38 +00:00
# -*- coding: utf-8 -*-
"""
wakatime.offlinequeue
~~~~~~~~~~~~~~~~~~~~~
2014-05-26 00:14:38 +00:00
Queue for saving heartbeats while offline.
2014-05-26 00:14:38 +00:00
:copyright: (c) 2014 Alan Hamlett.
:license: BSD, see LICENSE for more details.
"""
import logging
import os
from time import sleep
2017-11-09 06:54:33 +00:00
from .compat import json
from .constants import DEFAULT_SYNC_OFFLINE_ACTIVITY, HEARTBEATS_PER_REQUEST
2017-11-09 06:54:33 +00:00
from .heartbeat import Heartbeat
try:
import sqlite3
HAS_SQL = True
2015-09-07 01:22:05 +00:00
except ImportError: # pragma: nocover
HAS_SQL = False
2014-05-26 00:14:38 +00:00
log = logging.getLogger('WakaTime')
2014-05-26 00:14:38 +00:00
class Queue(object):
db_file = '.wakatime.db'
2017-11-09 06:54:33 +00:00
table_name = 'heartbeat_2'
2014-05-26 00:14:38 +00:00
2017-11-09 06:54:33 +00:00
args = None
configs = None
def __init__(self, args, configs):
self.args = args
self.configs = configs
2014-05-26 00:14:38 +00:00
def connect(self):
2017-11-09 06:54:33 +00:00
conn = sqlite3.connect(self._get_db_file(), isolation_level=None)
2014-05-26 00:14:38 +00:00
c = conn.cursor()
c.execute('''CREATE TABLE IF NOT EXISTS {0} (
2017-11-09 06:54:33 +00:00
id text,
heartbeat text)
'''.format(self.table_name))
2014-05-26 00:14:38 +00:00
return (conn, c)
2017-11-09 06:54:33 +00:00
def push(self, heartbeat):
if not HAS_SQL:
return
2014-05-26 21:05:54 +00:00
try:
conn, c = self.connect()
2017-11-09 06:54:33 +00:00
data = {
'id': heartbeat.get_id(),
'heartbeat': heartbeat.json(),
2014-05-26 21:05:54 +00:00
}
2017-11-09 06:54:33 +00:00
c.execute('INSERT INTO {0} VALUES (:id,:heartbeat)'.format(self.table_name), data)
2014-05-26 21:05:54 +00:00
conn.commit()
conn.close()
2014-05-27 03:02:30 +00:00
except sqlite3.Error:
2016-06-16 06:02:56 +00:00
log.traceback()
2014-05-26 00:14:38 +00:00
def pop(self):
2017-11-09 06:54:33 +00:00
if not HAS_SQL:
return None
2014-05-26 00:14:38 +00:00
tries = 3
wait = 0.1
2014-05-26 21:05:54 +00:00
try:
conn, c = self.connect()
2014-05-27 03:02:30 +00:00
except sqlite3.Error:
log.traceback(logging.DEBUG)
2014-05-26 21:05:54 +00:00
return None
2017-11-09 06:54:33 +00:00
heartbeat = None
2014-05-26 00:14:38 +00:00
loop = True
while loop and tries > -1:
try:
c.execute('BEGIN IMMEDIATE')
c.execute('SELECT * FROM {0} LIMIT 1'.format(self.table_name))
2014-05-26 00:14:38 +00:00
row = c.fetchone()
if row is not None:
2017-11-09 06:54:33 +00:00
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])
2014-05-26 00:14:38 +00:00
conn.commit()
loop = False
2017-11-09 06:54:33 +00:00
except sqlite3.Error:
log.traceback(logging.DEBUG)
2014-05-26 00:14:38 +00:00
sleep(wait)
tries -= 1
2014-05-26 21:05:54 +00:00
try:
conn.close()
2017-11-09 06:54:33 +00:00
except sqlite3.Error:
log.traceback(logging.DEBUG)
return heartbeat
2017-11-09 06:54:33 +00:00
def push_many(self, heartbeats):
for heartbeat in heartbeats:
self.push(heartbeat)
def pop_many(self, limit=None):
if limit is None:
limit = DEFAULT_SYNC_OFFLINE_ACTIVITY
2017-11-09 06:54:33 +00:00
heartbeats = []
count = 0
2018-09-20 16:08:09 +00:00
while count < limit:
2017-11-09 06:54:33 +00:00
heartbeat = self.pop()
if not heartbeat:
break
heartbeats.append(heartbeat)
count += 1
if count % HEARTBEATS_PER_REQUEST == 0:
yield heartbeats
heartbeats = []
2017-11-09 06:54:33 +00:00
if heartbeats:
yield heartbeats
2017-11-09 06:54:33 +00:00
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')