sublime-rana/packages/wakatime/session_cache.py

113 lines
2.9 KiB
Python
Raw Normal View History

2015-05-12 22:03:23 +00:00
# -*- 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
2015-09-29 10:11:25 +00:00
except ImportError: # pragma: nocover
2015-05-12 22:03:23 +00:00
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):
2017-10-29 18:32:03 +00:00
db_file = '.wakatime.db'
table_name = 'session'
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')
2015-05-12 22:03:23 +00:00
def connect(self):
2017-10-29 18:32:03 +00:00
conn = sqlite3.connect(self.get_db_file(), isolation_level=None)
2015-05-12 22:03:23 +00:00
c = conn.cursor()
2017-10-29 18:32:03 +00:00
c.execute('''CREATE TABLE IF NOT EXISTS {0} (
2015-05-12 22:03:23 +00:00
value BLOB)
2017-10-29 18:32:03 +00:00
'''.format(self.table_name))
2015-05-12 22:03:23 +00:00
return (conn, c)
def save(self, session):
"""Saves a requests.Session object for the next heartbeat process.
"""
2015-09-29 10:11:25 +00:00
if not HAS_SQL: # pragma: nocover
2015-05-12 22:03:23 +00:00
return
try:
conn, c = self.connect()
2017-10-29 18:32:03 +00:00
c.execute('DELETE FROM {0}'.format(self.table_name))
2015-05-12 22:03:23 +00:00
values = {
2015-08-28 22:05:55 +00:00
'value': sqlite3.Binary(pickle.dumps(session, protocol=2)),
2015-05-12 22:03:23 +00:00
}
2017-10-29 18:32:03 +00:00
c.execute('INSERT INTO {0} VALUES (:value)'.format(self.table_name), values)
2015-05-12 22:03:23 +00:00
conn.commit()
conn.close()
2015-09-29 10:11:25 +00:00
except: # pragma: nocover
2016-09-02 08:50:54 +00:00
log.traceback(logging.DEBUG)
2015-05-12 22:03:23 +00:00
def get(self):
"""Returns a requests.Session object.
Gets Session from sqlite3 cache or creates a new Session.
"""
2015-09-29 10:11:25 +00:00
if not HAS_SQL: # pragma: nocover
2015-05-12 22:03:23 +00:00
return requests.session()
try:
conn, c = self.connect()
except:
2016-09-02 08:50:54 +00:00
log.traceback(logging.DEBUG)
2015-05-12 22:03:23 +00:00
return requests.session()
session = None
try:
c.execute('BEGIN IMMEDIATE')
2017-10-29 18:32:03 +00:00
c.execute('SELECT value FROM {0} LIMIT 1'.format(self.table_name))
2015-05-12 22:03:23 +00:00
row = c.fetchone()
if row is not None:
session = pickle.loads(row[0])
2015-09-29 10:11:25 +00:00
except: # pragma: nocover
2016-09-02 08:50:54 +00:00
log.traceback(logging.DEBUG)
2015-05-12 22:03:23 +00:00
try:
conn.close()
2015-09-29 10:11:25 +00:00
except: # pragma: nocover
2016-09-02 08:50:54 +00:00
log.traceback(logging.DEBUG)
2015-05-12 22:03:23 +00:00
return session if session is not None else requests.session()
def delete(self):
"""Clears all cached Session objects.
"""
2015-09-29 10:11:25 +00:00
if not HAS_SQL: # pragma: nocover
2015-05-12 22:03:23 +00:00
return
try:
conn, c = self.connect()
2017-10-29 18:32:03 +00:00
c.execute('DELETE FROM {0}'.format(self.table_name))
2015-05-12 22:03:23 +00:00
conn.commit()
conn.close()
except:
2016-09-02 08:50:54 +00:00
log.traceback(logging.DEBUG)