diff options
| -rw-r--r-- | requests_cache/backends/sqlite.py | 18 | ||||
| -rw-r--r-- | tests/integration/test_sqlite.py | 15 |
2 files changed, 28 insertions, 5 deletions
diff --git a/requests_cache/backends/sqlite.py b/requests_cache/backends/sqlite.py index 74e0d71..3652481 100644 --- a/requests_cache/backends/sqlite.py +++ b/requests_cache/backends/sqlite.py @@ -9,7 +9,7 @@ import threading from contextlib import contextmanager from logging import getLogger from os import unlink -from os.path import isfile +from os.path import getsize, isfile from pathlib import Path from tempfile import gettempdir from time import time @@ -266,6 +266,22 @@ class SQLiteDict(BaseStorage): con.execute(f"DELETE FROM {self.table_name} WHERE expires <= ?", (round(time()),)) self.vacuum() + def size(self) -> int: + """Return the size of the database, in bytes. For an in-memory database, this will be an + estimate based on page size. + """ + try: + return getsize(self.db_path) + except IOError: + return self._estimate_size() + + def _estimate_size(self) -> int: + """Estimate the current size of the database based on page count * size""" + with self.connection() as conn: + page_count = conn.execute('PRAGMA page_count').fetchone()[0] + page_size = conn.execute('PRAGMA page_size').fetchone()[0] + return page_count * page_size + def sorted( self, key: str = 'expires', reversed: bool = False, limit: int = None, exclude_expired=False ): diff --git a/tests/integration/test_sqlite.py b/tests/integration/test_sqlite.py index b696236..54b0ceb 100644 --- a/tests/integration/test_sqlite.py +++ b/tests/integration/test_sqlite.py @@ -219,12 +219,19 @@ class TestSQLiteDict(BaseStorageTest): assert prev_item is None or prev_item.expires < item.expires assert item.status_code % 2 == 0 - def test_filesize(self): - """Test approximate expected size of database file, in bytes""" - cache = self.init_cache() + @pytest.mark.parametrize( + 'db_path, use_temp', + [ + ('filesize_test', True), + (':memory:', False), + ], + ) + def test_size(self, db_path, use_temp): + """Test approximate expected size of a database, for both file-based and in-memory databases""" + cache = self.init_cache(db_path, use_temp=use_temp) for i in range(100): cache[f'key_{i}'] = f'value_{i}' - assert 50000 < cache.filesize() < 200000 + assert 10000 < cache.size() < 200000 class TestSQLiteCache(BaseCacheTest): |