rana-cli/wakatime/session_cache.py
2017-11-08 22:54:33 -08:00

112 lines
2.9 KiB
Python

# -*- coding: utf-8 -*-
"""
wakatime.session_cache
~~~~~~~~~~~~~~~~~~~~~~
Persist requests.Session for multiprocess SSL handshake pooling.
:copyright: (c) 2015 Alan Hamlett.
:license: BSD, see LICENSE for more details.
"""
import logging
import os
import pickle
import sys
try:
import sqlite3
HAS_SQL = True
except ImportError: # pragma: nocover
HAS_SQL = False
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), 'packages'))
from .packages import requests
log = logging.getLogger('WakaTime')
class SessionCache(object):
db_file = '.wakatime.db'
table_name = 'session'
def connect(self):
conn = sqlite3.connect(self._get_db_file(), isolation_level=None)
c = conn.cursor()
c.execute('''CREATE TABLE IF NOT EXISTS {0} (
value BLOB)
'''.format(self.table_name))
return (conn, c)
def save(self, session):
"""Saves a requests.Session object for the next heartbeat process.
"""
if not HAS_SQL: # pragma: nocover
return
try:
conn, c = self.connect()
c.execute('DELETE FROM {0}'.format(self.table_name))
values = {
'value': sqlite3.Binary(pickle.dumps(session, protocol=2)),
}
c.execute('INSERT INTO {0} VALUES (:value)'.format(self.table_name), values)
conn.commit()
conn.close()
except: # pragma: nocover
log.traceback(logging.DEBUG)
def get(self):
"""Returns a requests.Session object.
Gets Session from sqlite3 cache or creates a new Session.
"""
if not HAS_SQL: # pragma: nocover
return requests.session()
try:
conn, c = self.connect()
except:
log.traceback(logging.DEBUG)
return requests.session()
session = None
try:
c.execute('BEGIN IMMEDIATE')
c.execute('SELECT value FROM {0} LIMIT 1'.format(self.table_name))
row = c.fetchone()
if row is not None:
session = pickle.loads(row[0])
except: # pragma: nocover
log.traceback(logging.DEBUG)
try:
conn.close()
except: # pragma: nocover
log.traceback(logging.DEBUG)
return session if session is not None else requests.session()
def delete(self):
"""Clears all cached Session objects.
"""
if not HAS_SQL: # pragma: nocover
return
try:
conn, c = self.connect()
c.execute('DELETE FROM {0}'.format(self.table_name))
conn.commit()
conn.close()
except:
log.traceback(logging.DEBUG)
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')