125 lines
3.1 KiB
Python
125 lines
3.1 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 limit == 0 or count < limit:
|
|
heartbeat = self.pop()
|
|
if not heartbeat:
|
|
break
|
|
heartbeats.append(heartbeat)
|
|
count += 1
|
|
|
|
return 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')
|