Class SqliteMetadataStore

mypy/metastore.py:179–278

class SqliteMetadataStore(MetadataStore):
    def __init__(
        self, cache_dir_prefix: str, set_journal_mode: bool = False, num_shards: int = 1
    ) -> None:
        # We check startswith instead of equality because the version
        # will have already been appended by the time the cache dir is
        # passed here.
        self.dbs: list[sqlite3.Connection] = []
        self.num_shards = num_shards
        self.dirty_shards: set[int] = set()
        if cache_dir_prefix.startswith(os.devnull):
            return

        os.makedirs(cache_dir_prefix, exist_ok=True)
        if num_shards <= 1:
            self.dbs.append(
                connect_db(os_path_join(cache_dir_prefix, "cache.db"), set_journal_mode)
            )
        else:
            for i in range(num_shards):
                self.dbs.append(
                    connect_db(os_path_join(cache_dir_prefix, f"cache.{i}.db"), set_journal_mode)
                )

    def _shard_index(self, name: str) -> int:
        if self.num_shards <= 1:
            return 0
        return hash_path_stem(name) % self.num_shards

    def _db_for(self, name: str) -> sqlite3.Connection:
        if not self.dbs:
            raise FileNotFoundError()
        return self.dbs[self._shard_index(name)]

    def _query(self, name: str, field: str) -> Any:
        # Raises FileNotFound for consistency with the file system version
        db = self._db_for(name)
        cur = db.execute(f"SELECT {field} FROM files2 WHERE path = ?", (name,))
        results = cur.fetchall()
        if not results:
            raise FileNotFoundError()
        assert len(results) == 1
        return results[0][0]

    def getmtime(self, name: str) -> float:
        mtime = self._query(name, "mtime")
        assert isinstance(mtime, float)
        return mtime

    def read(self, name: str) -> bytes:
        data = self._query(name, "data")
        assert isinstance(data, bytes)
        return data

    def write(self, name: str, data: bytes, mtime: float | None = None) -> bool:
        import sqlite3

        if not self.dbs: