path: root/requests_cache/backends/sqlite.py

"""SQLite cache backend. For usage details, see :ref:`Backends: SQLite <sqlite>`.

.. automodsumm:: requests_cache.backends.sqlite
   :classes-only:
   :nosignatures:
"""
import sqlite3
import threading
from contextlib import contextmanager
from logging import getLogger
from os import unlink
from os.path import getsize, isfile
from pathlib import Path
from tempfile import gettempdir
from time import time
from typing import Collection, Iterator, List, Optional, Tuple, Type, Union

from platformdirs import user_cache_dir

from .._utils import chunkify, get_valid_kwargs
from ..models.response import CachedResponse
from ..policy import ExpirationTime
from ..serializers import SerializerType, pickle_serializer
from . import BaseCache, BaseStorage
from .base import VT

MEMORY_URI = 'file::memory:?cache=shared'
SQLITE_MAX_VARIABLE_NUMBER = 999
AnyPath = Union[Path, str]
logger = getLogger(__name__)


class SQLiteCache(BaseCache):
    """SQLite cache backend.

    Args:
        db_path: Database file path
        use_cache_dir: Store datebase in a user cache directory (e.g., `~/.cache/http_cache.sqlite`)
        use_temp: Store database in a temp directory (e.g., ``/tmp/http_cache.sqlite``)
        use_memory: Store database in memory instead of in a file
        fast_save: Significantly increases cache write performance, but with the possibility of data
            loss. See `pragma: synchronous <https://www.sqlite.org/pragma.html#pragma_synchronous>`_
            for details.
        wal: Use `Write Ahead Logging <https://sqlite.org/wal.html>`_, so readers do not block writers.
        kwargs: Additional keyword arguments for :py:func:`sqlite3.connect`
    """

    def __init__(
        self,
        db_path: AnyPath = 'http_cache',
        serializer: Optional[SerializerType] = None,
        **kwargs,
    ):
        super().__init__(cache_name=str(db_path), **kwargs)
        # Only override serializer if a non-None value is specified
        skwargs = {'serializer': serializer, **kwargs} if serializer else kwargs
        self.responses: SQLiteDict = SQLiteDict(db_path, table_name='responses', **skwargs)
        self.redirects: SQLiteDict = SQLiteDict(
            db_path, table_name='redirects', serializer=None, **kwargs
        )

    @property
    def db_path(self) -> AnyPath:
        return self.responses.db_path

    def clear(self):
        """Delete all items from the cache. If this fails due to a corrupted cache or other I/O
        error, this will  attempt to delete the cache file and re-initialize.
        """
        try:
            super().clear()
        except Exception:
            logger.exception('Failed to clear cache')
            if isfile(self.responses.db_path):
                unlink(self.responses.db_path)
            self.responses.init_db()
            self.redirects.init_db()

    # A more efficient SQLite implementation of :py:meth:`BaseCache.delete`
    def delete(
        self,
        *keys: str,
        expired: bool = False,
        **kwargs,
    ):
        if keys:
            self.responses.bulk_delete(keys)
        if expired:
            self._delete_expired()

        # For any remaining conditions, use base implementation
        if kwargs:
            with self.responses._lock, self.redirects._lock:
                return super().delete(**kwargs)
        else:
            self._prune_redirects()

        self.responses.vacuum()
        self.redirects.vacuum()

    def _delete_expired(self):
        """A more efficient implementation of deleting expired responses in SQL"""
        with self.responses.connection(commit=True) as con:
            con.execute(
                f'DELETE FROM {self.responses.table_name} WHERE expires <= ?', (round(time()),)
            )

    def _prune_redirects(self):
        """A more efficient implementation of removing invalid redirects in SQL"""
        with self.redirects.connection(commit=True) as conn:
            t1 = self.redirects.table_name
            t2 = self.responses.table_name
            conn.execute(
                f'DELETE FROM {t1} WHERE key IN ('
                f'    SELECT {t1}.key FROM {t1}'
                f'    LEFT JOIN {t2} ON {t2}.key = {t1}.value'
                f'    WHERE {t2}.key IS NULL'
                ')'
            )

    def count(self, expired: bool = True) -> int:
        """Count number of responses, optionally excluding expired

        Args:
            expired: Set to ``False`` to count only unexpired responses
        """
        return self.responses.count(expired=expired)

    # A more efficient implementation of :py:meth:`BaseCache.filter` to make use of indexes
    def filter(
        self,
        valid: bool = True,
        expired: bool = True,
        invalid: bool = False,
        older_than: ExpirationTime = None,
    ) -> Iterator[CachedResponse]:
        if valid and not invalid:
            return self.responses.sorted(expired=expired)
        else:
            return super().filter(
                valid=valid, expired=expired, invalid=invalid, older_than=older_than
            )

    # A more efficient implementation of :py:meth:`BaseCache.recreate_keys
    def recreate_keys(self):
        with self.responses.bulk_commit():
            super().recreate_keys()

    def sorted(
        self,
        key: str = 'expires',
        reversed: bool = False,
        limit: Optional[int] = None,
        expired: bool = True,
    ) -> Iterator[CachedResponse]:
        """Get cached responses, with sorting and other query options.

        Args:
            key: Key to sort by; either 'expires', 'size', or 'key'
            reversed: Sort in descending order
            limit: Maximum number of responses to return
            expired: Set to ``False`` to exclude expired responses
        """
        return self.responses.sorted(key, reversed, limit, expired)


class SQLiteDict(BaseStorage):
    """A dictionary-like interface for SQLite"""

    def __init__(
        self,
        db_path,
        table_name='http_cache',
        fast_save=False,
        serializer: Optional[SerializerType] = pickle_serializer,
        use_cache_dir: bool = False,
        use_memory: bool = False,
        use_temp: bool = False,
        wal: bool = False,
        **kwargs,
    ):
        super().__init__(serializer=serializer, **kwargs)
        self._can_commit = True
        self._connection: Optional[sqlite3.Connection] = None
        self._lock = threading.RLock()
        self.connection_kwargs = get_valid_kwargs(sqlite_template, kwargs)
        self.connection_kwargs.setdefault('check_same_thread', False)
        if use_memory:
            self.connection_kwargs['uri'] = True
        self.db_path = _get_sqlite_cache_path(db_path, use_cache_dir, use_temp, use_memory)
        self.fast_save = fast_save
        self.table_name = table_name
        self.wal = wal
        self.init_db()

    def init_db(self):
        """Initialize the database, if it hasn't already been"""
        self.close()
        with self.connection(commit=True) as con:
            # Add new column to tables created before 1.0
            try:
                con.execute(f'ALTER TABLE {self.table_name} ADD COLUMN expires INTEGER')
            except sqlite3.OperationalError:
                pass

            con.execute(
                f'CREATE TABLE IF NOT EXISTS {self.table_name} ('
                '    key TEXT PRIMARY KEY,'
                '    value BLOB, '
                '    expires INTEGER'
                ')'
            )
            con.execute(f'CREATE INDEX IF NOT EXISTS expires_idx ON {self.table_name}(expires)')

    @contextmanager
    def connection(self, commit=False) -> Iterator[sqlite3.Connection]:
        """Get a thread-local database connection"""
        if not self._connection:
            logger.debug(f'Opening connection to {self.db_path}:{self.table_name}')
            self._connection = sqlite3.connect(self.db_path, **self.connection_kwargs)
            if self.fast_save:
                self._connection.execute('PRAGMA synchronous = 0;')
            if self.wal:
                self._connection.execute('PRAGMA journal_mode = wal')

        # Any write operations need to be run in serial
        if commit and self._can_commit:
            self._lock.acquire()
        yield self._connection
        if commit and self._can_commit:
            self._connection.commit()
            self._lock.release()

    def close(self):
        """Close any active connections"""
        if self._connection:
            self._connection.close()
            self._connection = None

    @contextmanager
    def bulk_commit(self):
        """Insert a large number of records within a single transaction

        Example:

            >>> d1 = SQLiteDict('test')
            >>> with d1.bulk_commit():
            ...     for i in range(1000):
            ...         d1[i] = i * 2

        """
        with self._lock:
            self._can_commit = False
            try:
                yield
                self._connection.commit()
            finally:
                self._can_commit = True

    def __del__(self):
        self.close()

    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 __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 self.deserialize(key, row[0])

    def __setitem__(self, key, value):
        # If available, set expiration as a timestamp in unix format
        expires = getattr(value, 'expires_unix', None)
        value = self.serialize(value)
        with self.connection(commit=True) as con:
            con.execute(
                f'INSERT OR REPLACE INTO {self.table_name} (key,value,expires) VALUES (?,?,?)',
                (key, value, expires),
            )

    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):
        return self.count()

    def bulk_delete(self, keys=None, values=None):
        """Delete multiple items from the cache, without raising errors for any missing items.
        Supports deleting by either key or by value.
        """
        if not keys and not values:
            return

        column = 'key' if keys else 'value'
        with self.connection(commit=True) as con:
            # Split into small enough chunks for SQLite to handle
            for chunk in chunkify(keys or values, max_size=SQLITE_MAX_VARIABLE_NUMBER):
                marks, args = _format_sequence(chunk)
                statement = f'DELETE FROM {self.table_name} WHERE {column} IN ({marks})'
                con.execute(statement, args)

    def clear(self):
        with self._lock:
            with self.connection(commit=True) as con:
                con.execute(f'DROP TABLE IF EXISTS {self.table_name}')
            self.init_db()
            self.vacuum()

    def count(self, expired: bool = True) -> int:
        """Count number of responses, optionally excluding expired"""
        filter_expr = ''
        params: Tuple = ()
        if not expired:
            filter_expr = 'WHERE expires is null or expires > ?'
            params = (time(),)
        query = f'SELECT COUNT(key) FROM {self.table_name} {filter_expr}'

        with self.connection() as con:
            return con.execute(query, params).fetchone()[0]

    def serialize(self, value: VT):
        value = super().serialize(value)
        return sqlite3.Binary(value) if isinstance(value, bytes) else value

    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: Optional[int] = None,
        expired: bool = True,
    ) -> Iterator[CachedResponse]:
        """Get cache values in sorted order; see :py:meth:`.SQLiteCache.sorted` for usage details"""
        # Get sort key, direction, and limit
        if key not in ['expires', 'size', 'key']:
            raise ValueError(f'Invalid sort key: {key}')
        if key == 'size':
            key = 'LENGTH(value)'
        direction = 'DESC' if reversed else 'ASC'
        limit_expr = f'LIMIT {limit}' if limit else ''

        # Filter out expired items, if specified
        filter_expr = ''
        params: Tuple = ()
        if not expired:
            filter_expr = 'WHERE expires is null or expires > ?'
            params = (time(),)

        with self.connection(commit=True) as con:
            for row in con.execute(
                f'SELECT key, value FROM {self.table_name} {filter_expr}'
                f'  ORDER BY {key} {direction} {limit_expr}',
                params,
            ):
                result = self.deserialize(row[0], row[1])
                # Omit any results that can't be deserialized
                if result:
                    yield result

    def vacuum(self):
        with self.connection(commit=True) as con:
            con.execute('VACUUM')


class SQLiteLRUDict(SQLiteDict):
    """SQLite backend that stays within a certain size limit"""

    def __init__(self, *args, max_size: int = True, **kwargs):
        super().__init__(*args, **kwargs)
        self.max_size = max_size

        # Maximim number of rows to remove in a single prune call
        self.max_remove_per_prune = SQLITE_MAX_VARIABLE_NUMBER

        # Minimum interval for pruning cache
        self.prune_interval = 10
        self.last_prune = time()

        # % of max size to target with each prune call (higher % results in more frequent pruning)
        self.prune_target_size_percent = 0.9

    def init_db(self):
        """Add last_used column and index, if it doesn't already exist"""
        super().init_db()
        with self._lock, self.connection() as con:
            try:
                con.execute(f'ALTER TABLE {self.table_name} ADD COLUMN last_used INTEGER')
            except sqlite3.OperationalError:
                pass
            con.execute(f'CREATE INDEX IF NOT EXISTS last_used_idx ON {self.table_name}(last_used)')

    def __getitem__(self, key):
        """Set last_used when retrieving an item"""
        value = super().__getitem__(key)
        with self.connection(commit=True) as con:
            con.execute(
                f'UPDATE {self.table_name} SET last_used=? WHERE key=?',
                (round(time()), key),
            )
        return value

    def __setitem__(self, key, value):
        """Set an item and remove any expired responses"""
        # If available, set expiration as a timestamp in unix format
        expires = getattr(value, 'expires_unix', None)
        now = round(time())
        value = self.serialize(value)
        with self.connection(commit=True) as con:
            con.execute(
                f'INSERT OR REPLACE INTO {self.table_name} '
                '(key,value,expires,last_used) VALUES (?,?,?,?)',
                (key, value, expires, now),
            )
            con.execute(f'DELETE FROM {self.table_name} WHERE expires <= ?', (now,))

        self._prune()

    def _check_prune(self) -> bool:
        return time() - self.last_prune > self.prune_interval and self.size() > self.max_size

    # WIP; just one of multiple approaches
    def _prune(self):
        if not self._check_prune():
            return

        with self.connection(commit=True) as con:
            # Get a batch of least recently used rows (up to a max query size)
            rows = con.execute(
                f'SELECT key, LENGTH(value) from {self.table_name} '
                'ORDER BY last_used ASC LIMIT ?',
                (self.max_remove_per_prune,),
            )

            # Figure out how many rows we need to delete to reach a target size of 90% of capacity
            delete_keys = []
            delete_size = self.size() - int(self.max_size * self.prune_target_size_percent)
            pending_delete_size = 0
            for key, row_size in rows:
                delete_keys.append(key)
                pending_delete_size += row_size
                if pending_delete_size >= delete_size:
                    break

            # Delete selected rows
            marks, args = _format_sequence(delete_keys)
            con.execute(f'DELETE FROM {self.table_name} WHERE key IN ({marks})', args)

        self.last_prune = time()


def _format_sequence(values: Collection) -> Tuple[str, List]:
    """Get SQL parameter marks for a sequence-based query"""
    return ','.join(['?'] * len(values)), list(values)


def _get_sqlite_cache_path(
    db_path: AnyPath, use_cache_dir: bool, use_temp: bool, use_memory: bool = False
) -> AnyPath:
    """Get a resolved path for a SQLite database file (or memory URI)"""
    # Use an in-memory database, if specified
    db_path = str(db_path)
    if use_memory:
        return MEMORY_URI
    elif ':memory:' in db_path or 'mode=memory' in db_path:
        return db_path

    # Add file extension if not specified
    if not Path(db_path).suffix:
        db_path += '.sqlite'
    return get_cache_path(db_path, use_cache_dir, use_temp)


def get_cache_path(db_path: AnyPath, use_cache_dir: bool = False, use_temp: bool = False) -> Path:
    """Get a resolved cache path"""
    db_path = Path(db_path)

    # Save to platform-specific temp or user cache directory, if specified
    if use_cache_dir and not db_path.is_absolute():
        db_path = Path(user_cache_dir()) / db_path
    elif use_temp and not db_path.is_absolute():
        db_path = Path(gettempdir()) / db_path

    # Expand relative and user paths (~), make parent dir(s), and better error if parent is a file
    db_path = db_path.expanduser().absolute()
    try:
        db_path.parent.mkdir(parents=True, exist_ok=True)
    except FileExistsError:
        raise FileExistsError(
            f'Parent path exists and is not a directory: {db_path.parent}.'
            'Please either delete the file or choose a different path.'
        )
    return db_path


def sqlite_template(
    timeout: float = 5.0,
    detect_types: int = 0,
    isolation_level: Optional[str] = None,
    check_same_thread: bool = True,
    factory: Optional[Type] = None,
    cached_statements: int = 100,
    uri: bool = False,
):
    """Template function to get an accurate signature for the builtin :py:func:`sqlite3.connect`"""