1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
|
import sqlite3
import threading
from contextlib import contextmanager
from logging import getLogger
from os import makedirs
from os.path import abspath, basename, dirname, expanduser, isabs, join
from pathlib import Path
from tempfile import gettempdir
from typing import Iterable, Tuple, Type, Union
from . import BaseCache, BaseStorage, get_valid_kwargs
logger = getLogger(__name__)
class DbCache(BaseCache):
"""SQLite cache backend.
Args:
db_path: Database file path (expands user paths and creates parent dirs)
use_temp: Store database in a temp directory (e.g., ``/tmp/http_cache.sqlite``).
Note: if ``db_path`` is an absolute path, this option will be ignored.
fast_save: Speedup cache saving up to 50 times but with possibility of data loss.
See :py:class:`.DbDict` for more info
kwargs: Additional keyword arguments for :py:func:`sqlite3.connect`
"""
def __init__(
self,
db_path: Union[Path, str] = 'http_cache',
use_temp: bool = False,
fast_save: bool = False,
**kwargs,
):
super().__init__(**kwargs)
self.responses = DbPickleDict(db_path, table_name='responses', use_temp=use_temp, fast_save=fast_save, **kwargs)
self.redirects = DbDict(db_path, table_name='redirects', use_temp=use_temp, **kwargs)
def delete_all(self, keys):
"""Remove multiple responses and their associated redirects from the cache, with additional cleanup"""
self.responses.delete_all(keys=keys)
self.responses.vacuum()
self.redirects.delete_all(keys=keys)
self.redirects.delete_all(values=keys)
self.redirects.vacuum()
class DbDict(BaseStorage):
"""A dictionary-like interface for SQLite.
It's possible to create multiply DbDict instances, which will be stored as separate
tables in one database::
d1 = DbDict('test', 'table1')
d2 = DbDict('test', 'table2')
d3 = DbDict('test', 'table3')
All data will be stored in separate tables in the file ``test.sqlite``.
Args:
db_path: Database file path
table_name: Table name
fast_save: Use `'PRAGMA synchronous = 0;' <http://www.sqlite.org/pragma.html#pragma_synchronous>`_
to speed up cache saving, but with the potential for data loss
timeout: Timeout for acquiring a database lock
"""
def __init__(self, db_path, table_name='http_cache', fast_save=False, use_temp: bool = False, **kwargs):
kwargs.setdefault('suppress_warnings', True)
super().__init__(**kwargs)
self.connection_kwargs = get_valid_kwargs(sqlite_template, kwargs)
self.db_path = _get_db_path(db_path, use_temp)
self.fast_save = fast_save
self.table_name = table_name
self._can_commit = True
self._local_context = threading.local()
with sqlite3.connect(self.db_path, **self.connection_kwargs) as con:
self._create_table(con)
# Initial CREATE TABLE must happen in shared connection; subsequent queries will use thread-local connections
def _create_table(self, connection):
connection.execute(f'CREATE TABLE IF NOT EXISTS {self.table_name} (key PRIMARY KEY, value)')
@contextmanager
def connection(self, commit=False) -> sqlite3.Connection:
"""Get a thread-local database connection"""
if not hasattr(self._local_context, 'con'):
logger.debug(f'Opening connection to {self.db_path}:{self.table_name}')
self._local_context.con = sqlite3.connect(self.db_path, **self.connection_kwargs)
if self.fast_save:
self._local_context.con.execute('PRAGMA synchronous = 0;')
yield self._local_context.con
if commit and self._can_commit:
self._local_context.con.commit()
def __del__(self):
if hasattr(self._local_context, 'con'):
self._local_context.con.close()
@contextmanager
def bulk_commit(self):
"""Context manager used to speed up insertion of a large number of records
Example:
>>> d1 = DbDict('test')
>>> with d1.bulk_commit():
... for i in range(1000):
... d1[i] = i * 2
"""
self._can_commit = False
try:
yield
if hasattr(self._local_context, 'con'):
self._local_context.con.commit()
finally:
self._can_commit = True
def __getitem__(self, key):
with self.connection() as con:
row = con.execute(f'SELECT value FROM {self.table_name} WHERE key=?', (key,)).fetchone()
# raise error after the with block, otherwise the connection will be locked
if not row:
raise KeyError
return row[0]
def __setitem__(self, key, item):
with self.connection(commit=True) as con:
con.execute(
f'INSERT OR REPLACE INTO {self.table_name} (key,value) VALUES (?,?)',
(key, item),
)
def __delitem__(self, key):
with self.connection(commit=True) as con:
cur = con.execute(f'DELETE FROM {self.table_name} WHERE key=?', (key,))
if not cur.rowcount:
raise KeyError
def __iter__(self):
with self.connection() as con:
for row in con.execute(f'SELECT key FROM {self.table_name}'):
yield row[0]
def __len__(self):
with self.connection() as con:
return con.execute(f'SELECT COUNT(key) FROM {self.table_name}').fetchone()[0]
def clear(self):
with self.connection(commit=True) as con:
con.execute(f'DROP TABLE IF EXISTS {self.table_name}')
self._create_table(con)
con.execute('VACUUM')
def delete_all(self, keys=None, values=None):
"""Delete multiple records, either by keys or values"""
if not keys and not values:
return
column = 'key' if keys else 'value'
marks, args = _format_sequence(keys or values)
statement = f'DELETE FROM {self.table_name} WHERE {column} IN ({marks})'
with self.connection(commit=True) as con:
con.execute(statement, args)
def vacuum(self):
with self.connection(commit=True) as con:
con.execute('VACUUM')
class DbPickleDict(DbDict):
"""Same as :class:`DbDict`, but pickles values before saving"""
def __setitem__(self, key, item):
super().__setitem__(key, sqlite3.Binary(self.serialize(item)))
def __getitem__(self, key):
return self.deserialize(super().__getitem__(key))
def _format_sequence(values) -> Tuple[str, Iterable]:
"""Get SQL parameter marks for a sequence-based query, and ensure value is a sequence"""
if not isinstance(values, Iterable):
values = [values]
return ','.join(['?'] * len(values)), values
def _get_db_path(db_path: Union[Path, str], use_temp: bool) -> str:
"""Get resolved path for database file"""
# Save to a temp directory, if specified
if use_temp and not isabs(db_path):
db_path = join(gettempdir(), db_path)
# Expand relative and user paths (~/*), and add file extension if not specified
db_path = abspath(expanduser(str(db_path)))
if '.' not in basename(db_path):
db_path += '.sqlite'
# Make sure parent dirs exist
makedirs(dirname(db_path), exist_ok=True)
return db_path
def sqlite_template(
timeout: float = 5.0,
detect_types: int = 0,
isolation_level: str = None,
check_same_thread: bool = True,
factory: Type = None,
cached_statements: int = 100,
uri: bool = False,
):
"""Template function to get an accurate signature for the builtin :py:func:`sqlite3.connect`"""
|
