from __future__ import annotations import contextlib import copy import itertools import logging import shutil import sqlite3 from collections import defaultdict from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Iterator, Sequence from contextlib import asynccontextmanager, contextmanager from dataclasses import dataclass, field, replace from hashlib import sha256 from pathlib import Path from typing import Any, BinaryIO import aiosqlite import anyio.to_thread import chia_rs.datalayer import zstd from chia_rs.datalayer import ( DeltaFileCache, DeltaReader, KeyAlreadyPresentError, KeyId, MerkleBlob, ProofOfInclusion, TreeIndex, ValueId, ) from chia_rs.sized_bytes import bytes32 from chia_rs.sized_ints import int64 from chia.data_layer.data_layer_errors import KeyNotFoundError, MerkleBlobNotFoundError, TreeGenerationIncrementingError from chia.data_layer.data_layer_util import ( DiffData, InsertResult, InternalNode, KeysPaginationData, KeysValuesCompressed, KeysValuesPaginationData, KVDiffPaginationData, Node, NodeType, OperationType, Root, SerializedNode, ServerInfo, Side, Status, Subscription, TerminalNode, Unspecified, get_delta_filename_path, get_hashes_for_page, internal_hash, key_hash, leaf_hash, unspecified, ) from chia.util.batches import to_batches from chia.util.cpu import available_logical_cores from chia.util.db_wrapper import SQLITE_MAX_VARIABLE_NUMBER, DBWrapper2 from chia.util.log_exceptions import log_exceptions from chia.util.lru_cache import LRUCache log = logging.getLogger(__name__) # TODO: review exceptions for values that shouldn't be displayed # TODO: pick exception types other than Exception KeyOrValueId = int64 default_prefer_file_kv_blob_length: int = 4096 @dataclass class DataStore: """A key/value store with the pairs being terminal nodes in a CLVM object tree.""" db_wrapper: DBWrapper2 recent_merkle_blobs: LRUCache[bytes32, MerkleBlob] merkle_blobs_path: Path key_value_blobs_path: Path unconfirmed_keys_values: dict[bytes32, list[bytes32]] = field(default_factory=dict) prefer_db_kv_blob_length: int = default_prefer_file_kv_blob_length @classmethod @contextlib.asynccontextmanager async def managed( cls, database: str | Path, merkle_blobs_path: Path, key_value_blobs_path: Path, uri: bool = False, sql_log_path: Path | None = None, cache_capacity: int = 1, prefer_db_kv_blob_length: int = default_prefer_file_kv_blob_length, ) -> AsyncIterator[DataStore]: async with DBWrapper2.managed( database=database, uri=uri, journal_mode="WAL", # Setting to FULL despite other locations being configurable. If there are # performance issues we can consider other the implications of other options. synchronous="FULL", # If foreign key checking gets turned off, please add corresponding check # methods and enable foreign key checking in the tests. foreign_keys=True, row_factory=aiosqlite.Row, log_path=sql_log_path, ) as db_wrapper: recent_merkle_blobs: LRUCache[bytes32, MerkleBlob] = LRUCache(capacity=cache_capacity) self = cls( db_wrapper=db_wrapper, recent_merkle_blobs=recent_merkle_blobs, merkle_blobs_path=merkle_blobs_path, key_value_blobs_path=key_value_blobs_path, prefer_db_kv_blob_length=prefer_db_kv_blob_length, ) async with db_wrapper.writer() as writer: await writer.execute( f""" CREATE TABLE IF NOT EXISTS root( tree_id BLOB NOT NULL CHECK(length(tree_id) == 32), generation INTEGER NOT NULL CHECK(generation >= 0), node_hash BLOB, status INTEGER NOT NULL CHECK( {" OR ".join(f"status == {status}" for status in Status)} ), PRIMARY KEY(tree_id, generation) ) """ ) await writer.execute( """ CREATE TABLE IF NOT EXISTS subscriptions( tree_id BLOB NOT NULL CHECK(length(tree_id) == 32), url TEXT, ignore_till INTEGER, num_consecutive_failures INTEGER, from_wallet tinyint CHECK(from_wallet == 0 OR from_wallet == 1), PRIMARY KEY(tree_id, url) ) """ ) await writer.execute( """ CREATE TABLE IF NOT EXISTS schema( version_id TEXT PRIMARY KEY, applied_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ) """ ) await writer.execute( """ CREATE INDEX IF NOT EXISTS node_hash ON root(node_hash) """ ) await writer.execute( """ CREATE TABLE IF NOT EXISTS ids( kv_id INTEGER PRIMARY KEY, hash BLOB NOT NULL CHECK(length(store_id) == 32), blob BLOB, store_id BLOB NOT NULL CHECK(length(store_id) == 32) ) """ ) await writer.execute( """ CREATE TABLE IF NOT EXISTS nodes( store_id BLOB NOT NULL CHECK(length(store_id) == 32), hash BLOB NOT NULL, root_hash BLOB NOT NULL, generation INTEGER NOT NULL CHECK(generation >= 0), idx INTEGER NOT NULL, PRIMARY KEY(store_id, hash) ) """ ) await writer.execute( """ CREATE UNIQUE INDEX IF NOT EXISTS ids_hash_index ON ids(hash, store_id) """ ) await writer.execute( """ CREATE INDEX IF NOT EXISTS nodes_generation_index ON nodes(generation) """ ) yield self @asynccontextmanager async def transaction(self) -> AsyncIterator[None]: async with self.db_wrapper.writer(): yield async def insert_into_data_store_from_file( self, store_id: bytes32, root_hash: bytes32 | None, filename: Path, delta_reader: DeltaReader | None = None, ) -> DeltaReader | None: async with self.db_wrapper.writer(): with self.manage_kv_files(store_id): if root_hash is None: merkle_blob = MerkleBlob(b"") else: root = await self.get_tree_root(store_id=store_id) if delta_reader is None: delta_reader = DeltaReader(internal_nodes={}, leaf_nodes={}) if root.node_hash is not None: delta_reader.collect_from_merkle_blob( self.get_merkle_path(store_id=store_id, root_hash=root.node_hash), indexes=[TreeIndex(0)], ) internal_nodes, terminal_nodes = await self.read_from_file(filename, store_id) delta_reader.add_internal_nodes(internal_nodes) delta_reader.add_leaf_nodes(terminal_nodes) missing_hashes = await anyio.to_thread.run_sync(delta_reader.get_missing_hashes, root_hash) if len(missing_hashes) > 0: # TODO: consider adding transactions around this code merkle_blob_queries = await self.build_merkle_blob_queries_for_missing_hashes( missing_hashes, store_id ) if len(merkle_blob_queries) > 0: jobs = [ (self.get_merkle_path(store_id=store_id, root_hash=old_root_hash), indexes) for old_root_hash, indexes in merkle_blob_queries.items() ] await anyio.to_thread.run_sync(delta_reader.collect_from_merkle_blobs, jobs) await self.build_cache_and_collect_missing_hashes(root, root_hash, store_id, delta_reader) merkle_blob = delta_reader.create_merkle_blob_and_filter_unused_nodes(root_hash, set()) # Don't store these blob objects into cache, since their data structures are not calculated yet. await self.insert_root_from_merkle_blob(merkle_blob, store_id, Status.COMMITTED, update_cache=False) return delta_reader async def build_merkle_blob_queries_for_missing_hashes( self, missing_hashes: set[bytes32], store_id: bytes32, ) -> defaultdict[bytes32, list[TreeIndex]]: queries = defaultdict[bytes32, list[TreeIndex]](list) batch_size = min(500, SQLITE_MAX_VARIABLE_NUMBER - 10) async with self.db_wrapper.reader() as reader: for batch in to_batches(missing_hashes, batch_size): placeholders = ",".join(["?"] * len(batch.entries)) query = f""" SELECT hash, root_hash, idx FROM nodes WHERE store_id = ? AND hash IN ({placeholders}) LIMIT {len(batch.entries)} """ async with reader.execute(query, (store_id, *batch.entries)) as cursor: rows = await cursor.fetchall() for row in rows: root_hash_blob = bytes32(row["root_hash"]) index = TreeIndex(row["idx"]) queries[root_hash_blob].append(index) return queries async def build_cache_and_collect_missing_hashes( self, root: Root, root_hash: bytes32, store_id: bytes32, delta_reader: DeltaReader, ) -> None: missing_hashes = delta_reader.get_missing_hashes(root_hash) if len(missing_hashes) == 0: return async with self.db_wrapper.reader() as reader: cursor = await reader.execute( # TODO: the INDEXED BY seems like it shouldn't be needed, figure out why it is # https://sqlite.org/lang_indexedby.html: admonished to omit all use of INDEXED BY # https://sqlite.org/queryplanner-ng.html#howtofix "SELECT MAX(generation) FROM nodes INDEXED BY nodes_generation_index WHERE store_id = ?", (store_id,), ) generation_row = await cursor.fetchone() if generation_row is None or generation_row[0] is None: current_generation = 0 else: current_generation = generation_row[0] generations: Sequence[int] = [current_generation] while missing_hashes: if current_generation >= root.generation: raise Exception("Invalid delta file, cannot find all the required hashes") current_generation = generations[-1] + 1 batch_size = available_logical_cores() generations = range( current_generation, min(current_generation + batch_size, root.generation), ) jobs: list[tuple[bytes32, Path]] = [] generations_by_root_hash: dict[bytes32, int] = {} for generation in generations: generation_root = await self.get_tree_root(store_id=store_id, generation=generation) if generation_root.node_hash is None: # no need to process an empty generation continue path = self.get_merkle_path(store_id=store_id, root_hash=generation_root.node_hash) jobs.append((generation_root.node_hash, path)) generations_by_root_hash[generation_root.node_hash] = generation found = await anyio.to_thread.run_sync( delta_reader.collect_and_return_from_merkle_blobs, jobs, missing_hashes, ) async with self.db_wrapper.writer() as writer: await writer.executemany( """ INSERT OR IGNORE INTO nodes(store_id, hash, root_hash, generation, idx) VALUES (?, ?, ?, ?, ?) """, ( (store_id, hash, root_hash, generations_by_root_hash[root_hash], index.raw) for root_hash, map in found for hash, index in map.items() ), ) missing_hashes = delta_reader.get_missing_hashes(root_hash) log.info(f"Missing hashes: added old hashes from generation {current_generation}") async def read_from_file( self, filename: Path, store_id: bytes32 ) -> tuple[dict[bytes32, tuple[bytes32, bytes32]], dict[bytes32, tuple[KeyId, ValueId]]]: internal_nodes: dict[bytes32, tuple[bytes32, bytes32]] = {} terminal_nodes: dict[bytes32, tuple[KeyId, ValueId]] = {} with open(filename, "rb") as reader: async with self.db_wrapper.writer() as writer: while True: chunk = b"" while len(chunk) < 4: size_to_read = 4 - len(chunk) cur_chunk = reader.read(size_to_read) if cur_chunk is None or cur_chunk == b"": if size_to_read < 4: raise Exception("Incomplete read of length.") break chunk += cur_chunk if chunk == b"": break size = int.from_bytes(chunk, byteorder="big") serialize_nodes_bytes = b"" while len(serialize_nodes_bytes) < size: size_to_read = size - len(serialize_nodes_bytes) cur_chunk = reader.read(size_to_read) if cur_chunk is None or cur_chunk == b"": raise Exception("Incomplete read of blob.") serialize_nodes_bytes += cur_chunk serialized_node = SerializedNode.from_bytes(serialize_nodes_bytes) node_type = NodeType.TERMINAL if serialized_node.is_terminal else NodeType.INTERNAL if node_type == NodeType.INTERNAL: node_hash = internal_hash(bytes32(serialized_node.value1), bytes32(serialized_node.value2)) internal_nodes[node_hash] = (bytes32(serialized_node.value1), bytes32(serialized_node.value2)) else: kid, vid = await self.add_key_value( serialized_node.value1, serialized_node.value2, store_id, writer=writer, ) node_hash = leaf_hash(serialized_node.value1, serialized_node.value2) terminal_nodes[node_hash] = (kid, vid) return internal_nodes, terminal_nodes async def migrate_db(self, server_files_location: Path) -> None: async with self.db_wrapper.reader() as reader: cursor = await reader.execute("SELECT * FROM schema") rows = await cursor.fetchall() all_versions = {"v1.0", "v2.0"} for row in rows: version = row["version_id"] if version not in all_versions: raise Exception("Unknown version") if version == "v2.0": log.info(f"Found DB schema version {version}. No migration needed.") return version = "v2.0" old_tables = ["node", "root", "ancestors"] all_stores = await self.get_store_ids() all_roots: list[list[Root]] = [] for store_id in all_stores: try: root = await self.get_tree_root(store_id=store_id) roots = await self.get_roots_between(store_id, 1, root.generation) all_roots.append([*roots, root]) except Exception as e: if "unable to find root for id, generation" in str(e): log.error(f"Cannot find roots for {store_id}. Skipping it.") log.info(f"Initiating migration to version {version}. Found {len(all_roots)} stores to migrate") async with self.db_wrapper.writer(foreign_key_enforcement_enabled=False) as writer: await writer.execute( f""" CREATE TABLE IF NOT EXISTS new_root( tree_id BLOB NOT NULL CHECK(length(tree_id) == 32), generation INTEGER NOT NULL CHECK(generation >= 0), node_hash BLOB, status INTEGER NOT NULL CHECK( {" OR ".join(f"status == {status}" for status in Status)} ), PRIMARY KEY(tree_id, generation) ) """ ) for old_table in old_tables: await writer.execute(f"DROP TABLE IF EXISTS {old_table}") await writer.execute("ALTER TABLE new_root RENAME TO root") await writer.execute("INSERT INTO schema (version_id) VALUES (?)", (version,)) log.info(f"Initialized new DB schema {version}.") total_generations = 0 synced_generations = 0 for roots in all_roots: assert len(roots) > 0 total_generations += len(roots) for roots in all_roots: store_id = roots[0].store_id await self.create_tree(store_id=store_id, status=Status.COMMITTED) expected_synced_generations = synced_generations + len(roots) for root in roots: recovery_filename: Path | None = None for group_by_store in (True, False): filename = get_delta_filename_path( server_files_location, store_id, bytes32.zeros if root.node_hash is None else root.node_hash, root.generation, group_by_store, ) if filename.exists(): log.info(f"Found filename {filename}. Recovering data from it") recovery_filename = filename break if recovery_filename is None: log.error(f"Cannot find any recovery file for root {root}") break try: await self.insert_into_data_store_from_file(store_id, root.node_hash, recovery_filename) synced_generations += 1 log.info( f"Successfully recovered root from {filename}. " f"Total roots processed: {(synced_generations / total_generations * 100):.2f}%" ) except Exception as e: log.error(f"Cannot recover data from {filename}: {e}") break if synced_generations < expected_synced_generations: log.error( f"Could not recover {expected_synced_generations - synced_generations} generations. " f"Consider resyncing the store {store_id} once the migration is complete." ) # Reset the counter as if we synced correctly, so the percentages add to 100% at the end. synced_generations = expected_synced_generations async def get_merkle_blob( self, store_id: bytes32, root_hash: bytes32 | None, read_only: bool = False, update_cache: bool = True, ) -> MerkleBlob: if root_hash is None: return MerkleBlob(blob=b"") if self.recent_merkle_blobs.get_capacity() == 0: update_cache = False existing_blob = self.recent_merkle_blobs.get(root_hash) if existing_blob is not None: return existing_blob if read_only else copy.deepcopy(existing_blob) try: with log_exceptions(log=log, message="Error while getting merkle blob"): path = self.get_merkle_path(store_id=store_id, root_hash=root_hash) # TODO: consider file-system based locking of either the file or the store directory merkle_blob = MerkleBlob.from_path(path) except Exception as e: raise MerkleBlobNotFoundError(root_hash=root_hash) from e if update_cache: self.recent_merkle_blobs.put(root_hash, copy.deepcopy(merkle_blob)) return merkle_blob def get_bytes_path(self, bytes_: bytes) -> Path: raw = bytes_.hex() segment_sizes = [2, 2, 2] start = 0 segments = [] for size in segment_sizes: segments.append(raw[start : start + size]) start += size return Path(*segments, raw) def get_merkle_path(self, store_id: bytes32, root_hash: bytes32 | None) -> Path: store_root = self.merkle_blobs_path.joinpath(store_id.hex()) if root_hash is None: return store_root return store_root.joinpath(self.get_bytes_path(bytes_=root_hash)) def get_key_value_path(self, store_id: bytes32, blob_hash: bytes32 | None) -> Path: store_root = self.key_value_blobs_path.joinpath(store_id.hex()) if blob_hash is None: return store_root return store_root.joinpath(self.get_bytes_path(bytes_=blob_hash)) async def insert_root_from_merkle_blob( self, merkle_blob: MerkleBlob, store_id: bytes32, status: Status, old_root: Root | None = None, update_cache: bool = True, ) -> Root: if not merkle_blob.empty(): merkle_blob.calculate_lazy_hashes() if self.recent_merkle_blobs.get_capacity() == 0: update_cache = False root_hash = merkle_blob.get_root_hash() if old_root is not None and old_root.node_hash == root_hash: raise ValueError("Changelist resulted in no change to tree data") if root_hash is not None: log.info(f"inserting merkle blob: {len(merkle_blob)} bytes {root_hash.hex()}") blob_path = self.get_merkle_path(store_id=store_id, root_hash=merkle_blob.get_root_hash()) if not blob_path.exists(): blob_path.parent.mkdir(parents=True, exist_ok=True) # TODO: consider file-system based locking of either the file or the store directory merkle_blob.to_path(blob_path) if update_cache: self.recent_merkle_blobs.put(root_hash, copy.deepcopy(merkle_blob)) return await self._insert_root(store_id, root_hash, status) def _use_file_for_new_kv_blob(self, blob: bytes) -> bool: return len(blob) > self.prefer_db_kv_blob_length async def get_kvid(self, blob: bytes, store_id: bytes32) -> KeyOrValueId | None: blob_hash = bytes32(sha256(blob).digest()) async with self.db_wrapper.reader() as reader: cursor = await reader.execute( "SELECT kv_id FROM ids WHERE hash = ? AND store_id = ?", ( blob_hash, store_id, ), ) row = await cursor.fetchone() if row is None: return None return KeyOrValueId(row[0]) def get_blob_from_file(self, blob_hash: bytes32, store_id: bytes32) -> bytes: # TODO: seems that zstd needs hinting # TODO: consider file-system based locking of either the file or the store directory return zstd.decompress(self.get_key_value_path(store_id=store_id, blob_hash=blob_hash).read_bytes()) # type: ignore[no-any-return] async def get_blob_from_kvid(self, kv_id: KeyOrValueId, store_id: bytes32) -> bytes | None: async with self.db_wrapper.reader() as reader: cursor = await reader.execute( "SELECT hash, blob FROM ids WHERE kv_id = ? AND store_id = ?", ( kv_id, store_id, ), ) row = await cursor.fetchone() if row is None: return None blob: bytes = row["blob"] if blob is not None: return blob blob_hash = bytes32(row["hash"]) return self.get_blob_from_file(blob_hash, store_id) async def get_terminal_node(self, kid: KeyId, vid: ValueId, store_id: bytes32) -> TerminalNode: key = await self.get_blob_from_kvid(kid.raw, store_id) value = await self.get_blob_from_kvid(vid.raw, store_id) if key is None or value is None: raise Exception("Cannot find the key/value pair") return TerminalNode(hash=leaf_hash(key, value), key=key, value=value) async def add_kvid(self, blob: bytes, store_id: bytes32, writer: aiosqlite.Connection) -> KeyOrValueId: use_file = self._use_file_for_new_kv_blob(blob) blob_hash = bytes32(sha256(blob).digest()) if use_file: table_blob = None else: table_blob = blob try: row = await writer.execute_insert( "INSERT INTO ids (hash, blob, store_id) VALUES (?, ?, ?)", ( blob_hash, table_blob, store_id, ), ) except sqlite3.IntegrityError as e: if "UNIQUE constraint failed" in str(e): kv_id = await self.get_kvid(blob, store_id) if kv_id is None: raise Exception("Internal error") from e return kv_id raise if row is None: raise Exception("Internal error") if use_file: path = self.get_key_value_path(store_id=store_id, blob_hash=blob_hash) path.parent.mkdir(parents=True, exist_ok=True) self.unconfirmed_keys_values[store_id].append(blob_hash) # TODO: consider file-system based locking of either the file or the store directory path.write_bytes(zstd.compress(blob)) return KeyOrValueId(row[0]) def delete_unconfirmed_kvids(self, store_id: bytes32) -> None: for blob_hash in self.unconfirmed_keys_values[store_id]: with log_exceptions(log=log, consume=True): path = self.get_key_value_path(store_id=store_id, blob_hash=blob_hash) try: path.unlink() except FileNotFoundError: log.error(f"Cannot find key/value path {path} for hash {blob_hash}") del self.unconfirmed_keys_values[store_id] def confirm_all_kvids(self, store_id: bytes32) -> None: del self.unconfirmed_keys_values[store_id] @contextmanager def manage_kv_files(self, store_id: bytes32) -> Iterator[None]: if store_id not in self.unconfirmed_keys_values: self.unconfirmed_keys_values[store_id] = [] else: raise Exception("Internal error: unconfirmed keys values cache not cleaned") try: yield except: self.delete_unconfirmed_kvids(store_id) raise else: self.confirm_all_kvids(store_id) async def add_key_value( self, key: bytes, value: bytes, store_id: bytes32, writer: aiosqlite.Connection ) -> tuple[KeyId, ValueId]: kid = KeyId(await self.add_kvid(key, store_id, writer=writer)) vid = ValueId(await self.add_kvid(value, store_id, writer=writer)) return (kid, vid) async def get_terminal_node_by_hash( self, node_hash: bytes32, store_id: bytes32, root_hash: bytes32 | Unspecified = unspecified, ) -> TerminalNode: resolved_root_hash: bytes32 | None if root_hash is unspecified: root = await self.get_tree_root(store_id=store_id) resolved_root_hash = root.node_hash else: resolved_root_hash = root_hash merkle_blob = await self.get_merkle_blob(store_id=store_id, root_hash=resolved_root_hash) kid, vid = merkle_blob.get_node_by_hash(node_hash) return await self.get_terminal_node(kid, vid, store_id) async def get_terminal_nodes_by_hashes( self, node_hashes: list[bytes32], store_id: bytes32, root_hash: bytes32 | Unspecified = unspecified, ) -> list[TerminalNode]: resolved_root_hash: bytes32 | None if root_hash is unspecified: root = await self.get_tree_root(store_id=store_id) resolved_root_hash = root.node_hash else: resolved_root_hash = root_hash merkle_blob = await self.get_merkle_blob(store_id=store_id, root_hash=resolved_root_hash) kv_ids: list[tuple[KeyId, ValueId]] = [] for node_hash in node_hashes: kid, vid = merkle_blob.get_node_by_hash(node_hash) kv_ids.append((kid, vid)) kv_ids_unpacked = (KeyOrValueId(id.raw) for kv_id in kv_ids for id in kv_id) table_blobs = await self.get_table_blobs(kv_ids_unpacked, store_id) terminal_nodes: list[TerminalNode] = [] for kid, vid in kv_ids: terminal_nodes.append(self.get_terminal_node_from_table_blobs(kid, vid, table_blobs, store_id)) return terminal_nodes async def get_existing_hashes(self, node_hashes: list[bytes32], store_id: bytes32) -> set[bytes32]: result: set[bytes32] = set() batch_size = min(500, SQLITE_MAX_VARIABLE_NUMBER - 10) async with self.db_wrapper.reader() as reader: for i in range(0, len(node_hashes), batch_size): chunk = node_hashes[i : i + batch_size] placeholders = ",".join(["?"] * len(chunk)) query = f"SELECT hash FROM nodes WHERE store_id = ? AND hash IN ({placeholders}) LIMIT {len(chunk)}" async with reader.execute(query, (store_id, *chunk)) as cursor: rows = await cursor.fetchall() result.update(row["hash"] for row in rows) return result async def add_node_hashes(self, store_id: bytes32, generation: int | None = None) -> None: root = await self.get_tree_root(store_id=store_id, generation=generation) if root.node_hash is None: return merkle_blob = await self.get_merkle_blob( store_id=store_id, root_hash=root.node_hash, read_only=True, update_cache=False ) hash_to_index = merkle_blob.get_hashes_indexes() existing_hashes = await self.get_existing_hashes(list(hash_to_index.keys()), store_id) async with self.db_wrapper.writer() as writer: await writer.executemany( """ INSERT INTO nodes(store_id, hash, root_hash, generation, idx) VALUES (?, ?, ?, ?, ?) """, ( (store_id, hash, root.node_hash, root.generation, index.raw) for hash, index in hash_to_index.items() if hash not in existing_hashes ), ) async def _insert_root( self, store_id: bytes32, node_hash: bytes32 | None, status: Status, generation: int | None = None, ) -> Root: async with self.db_wrapper.writer_maybe_transaction() as writer: if generation is None: try: existing_generation = await self.get_tree_generation(store_id=store_id) except Exception as e: if not str(e).startswith("No generations found for store ID:"): raise generation = 0 else: generation = existing_generation + 1 new_root = Root( store_id=store_id, node_hash=node_hash, generation=generation, status=status, ) await writer.execute( """ INSERT INTO root(tree_id, generation, node_hash, status) VALUES(:tree_id, :generation, :node_hash, :status) """, new_root.to_row(), ) return new_root async def get_pending_root(self, store_id: bytes32) -> Root | None: async with self.db_wrapper.reader() as reader: cursor = await reader.execute( """ SELECT * FROM root WHERE tree_id == :tree_id AND status IN (:pending_status, :pending_batch_status) LIMIT 2 """, { "tree_id": store_id, "pending_status": Status.PENDING.value, "pending_batch_status": Status.PENDING_BATCH.value, }, ) row = await cursor.fetchone() if row is None: return None maybe_extra_result = await cursor.fetchone() if maybe_extra_result is not None: raise Exception(f"multiple pending roots found for id: {store_id.hex()}") return Root.from_row(row=row) async def clear_pending_roots(self, store_id: bytes32) -> Root | None: async with self.db_wrapper.writer() as writer: pending_root = await self.get_pending_root(store_id=store_id) if pending_root is not None: await writer.execute( "DELETE FROM root WHERE tree_id == :tree_id AND status IN (:pending_status, :pending_batch_status)", { "tree_id": store_id, "pending_status": Status.PENDING.value, "pending_batch_status": Status.PENDING_BATCH.value, }, ) return pending_root async def shift_root_generations(self, store_id: bytes32, shift_size: int) -> None: async with self.db_wrapper.writer(): root = await self.get_tree_root(store_id=store_id) for _ in range(shift_size): await self._insert_root(store_id=store_id, node_hash=root.node_hash, status=Status.COMMITTED) async def change_root_status(self, root: Root, status: Status = Status.PENDING) -> None: async with self.db_wrapper.writer() as writer: await writer.execute( "UPDATE root SET status = ? WHERE tree_id=? and generation = ?", ( status.value, root.store_id, root.generation, ), ) async def check(self) -> None: for check in self._checks: await check(self) async def _check_roots_are_incrementing(self) -> None: async with self.db_wrapper.reader() as reader: cursor = await reader.execute("SELECT * FROM root ORDER BY tree_id, generation") roots = [Root.from_row(row=row) async for row in cursor] roots_by_tree: dict[bytes32, list[Root]] = defaultdict(list) for root in roots: roots_by_tree[root.store_id].append(root) bad_trees = [] for store_id, roots in roots_by_tree.items(): current_generation = roots[-1].generation expected_generations = list(range(current_generation + 1)) actual_generations = [root.generation for root in roots] if actual_generations != expected_generations: bad_trees.append(store_id) if len(bad_trees) > 0: raise TreeGenerationIncrementingError(store_ids=bad_trees) _checks: tuple[Callable[[DataStore], Awaitable[None]], ...] = (_check_roots_are_incrementing,) async def create_tree(self, store_id: bytes32, status: Status = Status.PENDING) -> bool: await self._insert_root(store_id=store_id, node_hash=None, status=status) return True async def table_is_empty(self, store_id: bytes32) -> bool: tree_root = await self.get_tree_root(store_id=store_id) return tree_root.node_hash is None async def get_store_ids(self) -> set[bytes32]: async with self.db_wrapper.reader() as reader: cursor = await reader.execute("SELECT DISTINCT tree_id FROM root") store_ids = {bytes32(row["tree_id"]) async for row in cursor} return store_ids async def get_tree_generation(self, store_id: bytes32) -> int: async with self.db_wrapper.reader() as reader: cursor = await reader.execute( "SELECT MAX(generation) FROM root WHERE tree_id == :tree_id AND status == :status", {"tree_id": store_id, "status": Status.COMMITTED.value}, ) row = await cursor.fetchone() if row is not None: generation: int | None = row["MAX(generation)"] if generation is not None: return generation raise Exception(f"No generations found for store ID: {store_id.hex()}") async def get_tree_root(self, store_id: bytes32, generation: int | None = None) -> Root: async with self.db_wrapper.reader() as reader: if generation is None: generation = await self.get_tree_generation(store_id=store_id) cursor = await reader.execute( """ SELECT * FROM root WHERE tree_id == :tree_id AND generation == :generation AND status == :status LIMIT 1 """, {"tree_id": store_id, "generation": generation, "status": Status.COMMITTED.value}, ) row = await cursor.fetchone() if row is None: raise Exception(f"unable to find root for id, generation: {store_id.hex()}, {generation}") return Root.from_row(row=row) async def get_all_pending_batches_roots(self) -> list[Root]: async with self.db_wrapper.reader() as reader: cursor = await reader.execute( """ SELECT * FROM root WHERE status == :status """, {"status": Status.PENDING_BATCH.value}, ) roots = [Root.from_row(row=row) async for row in cursor] store_ids = [root.store_id for root in roots] if len(set(store_ids)) != len(store_ids): raise Exception("Internal error: multiple pending batches for a store") return roots async def store_id_exists(self, store_id: bytes32) -> bool: async with self.db_wrapper.reader() as reader: cursor = await reader.execute( "SELECT 1 FROM root WHERE tree_id == :tree_id AND status == :status LIMIT 1", {"tree_id": store_id, "status": Status.COMMITTED.value}, ) row = await cursor.fetchone() if row is None: return False return True async def get_roots_between(self, store_id: bytes32, generation_begin: int, generation_end: int) -> list[Root]: async with self.db_wrapper.reader() as reader: cursor = await reader.execute( "SELECT * FROM root WHERE tree_id == :tree_id " "AND generation >= :generation_begin AND generation < :generation_end ORDER BY generation ASC", {"tree_id": store_id, "generation_begin": generation_begin, "generation_end": generation_end}, ) roots = [Root.from_row(row=row) async for row in cursor] return roots async def get_ancestors( self, node_hash: bytes32, store_id: bytes32, root_hash: bytes32 | None = None, generation: int | None = None, ) -> list[InternalNode]: async with self.db_wrapper.reader(): if root_hash is None: root = await self.get_tree_root(store_id=store_id, generation=generation) root_hash = root.node_hash if root_hash is None: raise Exception(f"Root hash is unspecified for store ID: {store_id.hex()}") merkle_blob = await self.get_merkle_blob(store_id=store_id, root_hash=root_hash) reference_kid, _ = merkle_blob.get_node_by_hash(node_hash) reference_index = merkle_blob.get_key_index(reference_kid) lineage = merkle_blob.get_lineage_with_indexes(reference_index) result: list[InternalNode] = [] for index, node in itertools.islice(lineage, 1, None): assert isinstance(node, chia_rs.datalayer.InternalNode) result.append( InternalNode( hash=node.hash, left_hash=merkle_blob.get_hash_at_index(node.left), right_hash=merkle_blob.get_hash_at_index(node.right), ) ) return result def get_terminal_node_from_table_blobs( self, kid: KeyId, vid: ValueId, table_blobs: dict[KeyOrValueId, tuple[bytes32, bytes | None]], store_id: bytes32, ) -> TerminalNode: key = table_blobs[KeyOrValueId(kid.raw)][1] if key is None: key_hash = table_blobs[KeyOrValueId(kid.raw)][0] key = self.get_blob_from_file(key_hash, store_id) value = table_blobs[KeyOrValueId(vid.raw)][1] if value is None: value_hash = table_blobs[KeyOrValueId(vid.raw)][0] value = self.get_blob_from_file(value_hash, store_id) return TerminalNode(hash=leaf_hash(key, value), key=key, value=value) async def get_keys_values( self, store_id: bytes32, root_hash: bytes32 | Unspecified = unspecified, ) -> list[TerminalNode]: async with self.db_wrapper.reader(): resolved_root_hash: bytes32 | None if root_hash is unspecified: root = await self.get_tree_root(store_id=store_id) resolved_root_hash = root.node_hash else: resolved_root_hash = root_hash try: merkle_blob = await self.get_merkle_blob(store_id=store_id, root_hash=resolved_root_hash) except MerkleBlobNotFoundError: return [] kv_ids = merkle_blob.get_keys_values() kv_ids_unpacked = (KeyOrValueId(id.raw) for pair in kv_ids.items() for id in pair) table_blobs = await self.get_table_blobs(kv_ids_unpacked, store_id) terminal_nodes: list[TerminalNode] = [] for kid, vid in kv_ids.items(): terminal_nodes.append(self.get_terminal_node_from_table_blobs(kid, vid, table_blobs, store_id)) return terminal_nodes async def get_keys_values_compressed( self, store_id: bytes32, root_hash: bytes32 | Unspecified = unspecified, ) -> KeysValuesCompressed: async with self.db_wrapper.reader(): resolved_root_hash: bytes32 | None if root_hash is unspecified: root = await self.get_tree_root(store_id=store_id) resolved_root_hash = root.node_hash else: resolved_root_hash = root_hash keys_values_hashed: dict[bytes32, bytes32] = {} key_hash_to_length: dict[bytes32, int] = {} leaf_hash_to_length: dict[bytes32, int] = {} if resolved_root_hash is not None: try: merkle_blob = await self.get_merkle_blob(store_id=store_id, root_hash=resolved_root_hash) except MerkleBlobNotFoundError: return KeysValuesCompressed({}, {}, {}, resolved_root_hash) kv_ids = merkle_blob.get_keys_values() kv_ids_unpacked = (KeyOrValueId(id.raw) for pair in kv_ids.items() for id in pair) table_blobs = await self.get_table_blobs(kv_ids_unpacked, store_id) for kid, vid in kv_ids.items(): node = self.get_terminal_node_from_table_blobs(kid, vid, table_blobs, store_id) keys_values_hashed[key_hash(node.key)] = leaf_hash(node.key, node.value) key_hash_to_length[key_hash(node.key)] = len(node.key) leaf_hash_to_length[leaf_hash(node.key, node.value)] = len(node.key) + len(node.value) return KeysValuesCompressed(keys_values_hashed, key_hash_to_length, leaf_hash_to_length, resolved_root_hash) async def get_keys_paginated( self, store_id: bytes32, page: int, max_page_size: int, root_hash: bytes32 | Unspecified = unspecified, ) -> KeysPaginationData: keys_values_compressed = await self.get_keys_values_compressed(store_id, root_hash) pagination_data = get_hashes_for_page(page, keys_values_compressed.key_hash_to_length, max_page_size) keys: list[bytes] = [] leaf_hashes: list[bytes32] = [] for hash in pagination_data.hashes: leaf_hash = keys_values_compressed.keys_values_hashed[hash] leaf_hashes.append(leaf_hash) nodes = await self.get_terminal_nodes_by_hashes(leaf_hashes, store_id, root_hash) keys = [node.key for node in nodes] return KeysPaginationData( pagination_data.total_pages, pagination_data.total_bytes, keys, keys_values_compressed.root_hash, ) async def get_keys_values_paginated( self, store_id: bytes32, page: int, max_page_size: int, root_hash: bytes32 | Unspecified = unspecified, ) -> KeysValuesPaginationData: keys_values_compressed = await self.get_keys_values_compressed(store_id, root_hash) pagination_data = get_hashes_for_page(page, keys_values_compressed.leaf_hash_to_length, max_page_size) keys_values = await self.get_terminal_nodes_by_hashes(pagination_data.hashes, store_id, root_hash) return KeysValuesPaginationData( pagination_data.total_pages, pagination_data.total_bytes, keys_values, keys_values_compressed.root_hash, ) async def get_kv_diff_paginated( self, store_id: bytes32, page: int, max_page_size: int, # NOTE: empty is expressed as zeros hash1: bytes32, hash2: bytes32, ) -> KVDiffPaginationData: old_pairs = await self.get_keys_values_compressed(store_id, hash1) if len(old_pairs.keys_values_hashed) == 0 and hash1 != bytes32.zeros: raise Exception(f"Unable to diff: Can't find keys and values for {hash1}") new_pairs = await self.get_keys_values_compressed(store_id, hash2) if len(new_pairs.keys_values_hashed) == 0 and hash2 != bytes32.zeros: raise Exception(f"Unable to diff: Can't find keys and values for {hash2}") old_pairs_leaf_hashes = {v for v in old_pairs.keys_values_hashed.values()} new_pairs_leaf_hashes = {v for v in new_pairs.keys_values_hashed.values()} insertions = {k for k in new_pairs_leaf_hashes if k not in old_pairs_leaf_hashes} deletions = {k for k in old_pairs_leaf_hashes if k not in new_pairs_leaf_hashes} lengths = {} for hash in insertions: lengths[hash] = new_pairs.leaf_hash_to_length[hash] for hash in deletions: lengths[hash] = old_pairs.leaf_hash_to_length[hash] pagination_data = get_hashes_for_page(page, lengths, max_page_size) kv_diff: list[DiffData] = [] insertion_hashes: list[bytes32] = [] deletion_hashes: list[bytes32] = [] for hash in pagination_data.hashes: if hash in insertions: insertion_hashes.append(hash) else: deletion_hashes.append(hash) if hash2 != bytes32.zeros: insertion_nodes = await self.get_terminal_nodes_by_hashes(insertion_hashes, store_id, hash2) else: insertion_nodes = [] if hash1 != bytes32.zeros: deletion_nodes = await self.get_terminal_nodes_by_hashes(deletion_hashes, store_id, hash1) else: deletion_nodes = [] for node in insertion_nodes: kv_diff.append(DiffData(OperationType.INSERT, node.key, node.value)) for node in deletion_nodes: kv_diff.append(DiffData(OperationType.DELETE, node.key, node.value)) return KVDiffPaginationData( pagination_data.total_pages, pagination_data.total_bytes, kv_diff, ) async def autoinsert( self, key: bytes, value: bytes, store_id: bytes32, status: Status = Status.PENDING, root: Root | None = None, ) -> InsertResult: return await self.insert( key=key, value=value, store_id=store_id, reference_node_hash=None, side=None, status=status, root=root, ) async def get_keys( self, store_id: bytes32, root_hash: bytes32 | Unspecified = unspecified, ) -> list[bytes]: async with self.db_wrapper.reader(): if root_hash is unspecified: root = await self.get_tree_root(store_id=store_id) resolved_root_hash = root.node_hash else: resolved_root_hash = root_hash try: merkle_blob = await self.get_merkle_blob(store_id=store_id, root_hash=resolved_root_hash) except MerkleBlobNotFoundError: return [] kv_ids = merkle_blob.get_keys_values() raw_key_ids = (KeyOrValueId(id.raw) for id in kv_ids.keys()) table_blobs = await self.get_table_blobs(raw_key_ids, store_id) keys: list[bytes] = [] for kid in kv_ids.keys(): blob_hash, blob = table_blobs[KeyOrValueId(kid.raw)] if blob is None: blob = self.get_blob_from_file(blob_hash, store_id) keys.append(blob) return keys def get_reference_kid_side(self, merkle_blob: MerkleBlob, seed: bytes32) -> tuple[KeyId, Side]: side_seed = bytes(seed)[0] side = Side.LEFT if side_seed < 128 else Side.RIGHT reference_node = merkle_blob.get_random_leaf_node(seed) kid = reference_node.key return (kid, side) async def get_terminal_node_from_kid(self, merkle_blob: MerkleBlob, kid: KeyId, store_id: bytes32) -> TerminalNode: index = merkle_blob.get_key_index(kid) raw_node = merkle_blob.get_raw_node(index) assert isinstance(raw_node, chia_rs.datalayer.LeafNode) return await self.get_terminal_node(raw_node.key, raw_node.value, store_id) async def get_terminal_node_for_seed(self, seed: bytes32, store_id: bytes32) -> TerminalNode | None: root = await self.get_tree_root(store_id=store_id) if root is None or root.node_hash is None: return None merkle_blob = await self.get_merkle_blob(store_id=store_id, root_hash=root.node_hash) assert not merkle_blob.empty() kid, _ = self.get_reference_kid_side(merkle_blob, seed) return await self.get_terminal_node_from_kid(merkle_blob, kid, store_id) async def insert( self, key: bytes, value: bytes, store_id: bytes32, reference_node_hash: bytes32 | None = None, side: Side | None = None, status: Status = Status.PENDING, root: Root | None = None, ) -> InsertResult: async with self.db_wrapper.writer() as writer: with self.manage_kv_files(store_id): if root is None: root = await self.get_tree_root(store_id=store_id) merkle_blob = await self.get_merkle_blob(store_id=store_id, root_hash=root.node_hash) kid, vid = await self.add_key_value(key, value, store_id, writer=writer) hash = leaf_hash(key, value) reference_kid = None if reference_node_hash is not None: reference_kid, _ = merkle_blob.get_node_by_hash(reference_node_hash) was_empty = root.node_hash is None if not was_empty and reference_kid is None: if side is not None: raise Exception("Side specified without reference node hash") seed = leaf_hash(key=key, value=value) reference_kid, side = self.get_reference_kid_side(merkle_blob, seed) merkle_blob.insert(kid, vid, hash, reference_kid, side) new_root = await self.insert_root_from_merkle_blob(merkle_blob, store_id, status) return InsertResult(node_hash=hash, root=new_root) async def delete( self, key: bytes, store_id: bytes32, status: Status = Status.PENDING, root: Root | None = None, ) -> Root | None: async with self.db_wrapper.writer(): if root is None: root = await self.get_tree_root(store_id=store_id) merkle_blob = await self.get_merkle_blob(store_id=store_id, root_hash=root.node_hash) kid = await self.get_kvid(key, store_id) if kid is not None: merkle_blob.delete(KeyId(kid)) new_root = await self.insert_root_from_merkle_blob(merkle_blob, store_id, status) return new_root async def upsert( self, key: bytes, new_value: bytes, store_id: bytes32, status: Status = Status.PENDING, root: Root | None = None, ) -> InsertResult: async with self.db_wrapper.writer() as writer: with self.manage_kv_files(store_id): if root is None: root = await self.get_tree_root(store_id=store_id) merkle_blob = await self.get_merkle_blob(store_id=store_id, root_hash=root.node_hash) kid, vid = await self.add_key_value(key, new_value, store_id, writer=writer) hash = leaf_hash(key, new_value) merkle_blob.upsert(kid, vid, hash) new_root = await self.insert_root_from_merkle_blob(merkle_blob, store_id, status) return InsertResult(node_hash=hash, root=new_root) async def insert_batch( self, store_id: bytes32, changelist: list[dict[str, Any]], status: Status = Status.PENDING, enable_batch_autoinsert: bool = True, ) -> bytes32 | None: async with self.db_wrapper.writer() as writer: with self.manage_kv_files(store_id): old_root = await self.get_tree_root(store_id=store_id) pending_root = await self.get_pending_root(store_id=store_id) if pending_root is not None: if pending_root.status == Status.PENDING_BATCH: # We have an unfinished batch, continue the current batch on top of it. if pending_root.generation != old_root.generation + 1: raise Exception("Internal error") old_root = pending_root await self.clear_pending_roots(store_id) else: raise Exception("Internal error") merkle_blob = await self.get_merkle_blob(store_id=store_id, root_hash=old_root.node_hash) key_hash_frequency: dict[bytes32, int] = {} first_action: dict[bytes32, str] = {} last_action: dict[bytes32, str] = {} for change in changelist: key = change["key"] hash = key_hash(key) key_hash_frequency[hash] = key_hash_frequency.get(hash, 0) + 1 if hash not in first_action: first_action[hash] = change["action"] last_action[hash] = change["action"] batch_keys_values: list[tuple[KeyId, ValueId]] = [] batch_hashes: list[bytes32] = [] for change in changelist: if change["action"] == "insert": key = change["key"] value = change["value"] reference_node_hash = change.get("reference_node_hash", None) side = change.get("side", None) reference_kid: KeyId | None = None if reference_node_hash is not None: reference_kid, _ = merkle_blob.get_node_by_hash(reference_node_hash) key_hashed = key_hash(key) kid, vid = await self.add_key_value(key, value, store_id, writer=writer) try: merkle_blob.get_key_index(kid) except chia_rs.datalayer.UnknownKeyError: pass else: raise KeyAlreadyPresentError(kid) hash = leaf_hash(key, value) if reference_node_hash is None and side is None: if enable_batch_autoinsert and reference_kid is None: if key_hash_frequency[key_hashed] == 1 or ( key_hash_frequency[key_hashed] == 2 and first_action[key_hashed] == "delete" ): batch_keys_values.append((kid, vid)) batch_hashes.append(hash) continue if not merkle_blob.empty(): seed = leaf_hash(key=key, value=value) reference_kid, side = self.get_reference_kid_side(merkle_blob, seed) merkle_blob.insert(kid, vid, hash, reference_kid, side) elif change["action"] == "delete": key = change["key"] deletion_kid = await self.get_kvid(key, store_id) if deletion_kid is not None: merkle_blob.delete(KeyId(deletion_kid)) elif change["action"] == "upsert": key = change["key"] new_value = change["value"] kid, vid = await self.add_key_value(key, new_value, store_id, writer=writer) hash = leaf_hash(key, new_value) merkle_blob.upsert(kid, vid, hash) else: raise Exception(f"Operation in batch is not insert or delete: {change}") if len(batch_keys_values) > 0: merkle_blob.batch_insert(batch_keys_values, batch_hashes) new_root = await self.insert_root_from_merkle_blob(merkle_blob, store_id, status, old_root) return new_root.node_hash async def get_node_by_key( self, key: bytes, store_id: bytes32, root_hash: bytes32 | Unspecified = unspecified, ) -> TerminalNode: async with self.db_wrapper.reader(): resolved_root_hash: bytes32 | None if root_hash is unspecified: root = await self.get_tree_root(store_id=store_id) resolved_root_hash = root.node_hash else: resolved_root_hash = root_hash try: merkle_blob = await self.get_merkle_blob(store_id=store_id, root_hash=resolved_root_hash) except MerkleBlobNotFoundError: raise KeyNotFoundError(key=key) kvid = await self.get_kvid(key, store_id) if kvid is None: raise KeyNotFoundError(key=key) kid = KeyId(kvid) return await self.get_terminal_node_from_kid(merkle_blob, kid, store_id) async def get_tree_as_nodes(self, store_id: bytes32) -> Node: async with self.db_wrapper.reader(): root = await self.get_tree_root(store_id=store_id) # TODO: consider actual proper behavior assert root.node_hash is not None merkle_blob = await self.get_merkle_blob(store_id=store_id, root_hash=root.node_hash) nodes = merkle_blob.get_nodes_with_indexes() hash_to_node: dict[bytes32, Node] = {} tree_node: Node for _, node in reversed(nodes): if isinstance(node, chia_rs.datalayer.InternalNode): left_hash = merkle_blob.get_hash_at_index(node.left) right_hash = merkle_blob.get_hash_at_index(node.right) tree_node = InternalNode.from_child_nodes( left=hash_to_node[left_hash], right=hash_to_node[right_hash] ) else: assert isinstance(node, chia_rs.datalayer.LeafNode) tree_node = await self.get_terminal_node(node.key, node.value, store_id) hash_to_node[node.hash] = tree_node root_node = hash_to_node[root.node_hash] return root_node async def get_proof_of_inclusion_by_hash( self, node_hash: bytes32, store_id: bytes32, root_hash: bytes32 | None = None, ) -> ProofOfInclusion: if root_hash is None: root = await self.get_tree_root(store_id=store_id) root_hash = root.node_hash merkle_blob = await self.get_merkle_blob(store_id=store_id, root_hash=root_hash) kid, _ = merkle_blob.get_node_by_hash(node_hash) return merkle_blob.get_proof_of_inclusion(kid) async def get_proof_of_inclusion_by_key( self, key: bytes, store_id: bytes32, ) -> ProofOfInclusion: root = await self.get_tree_root(store_id=store_id) merkle_blob = await self.get_merkle_blob(store_id=store_id, root_hash=root.node_hash) kvid = await self.get_kvid(key, store_id) if kvid is None: raise Exception(f"Cannot find key: {key.hex()}") kid = KeyId(kvid) return merkle_blob.get_proof_of_inclusion(kid) async def get_nodes_for_file( self, root: Root, node_hash: bytes32, store_id: bytes32, deltas_only: bool, delta_file_cache: DeltaFileCache, tree_nodes: list[SerializedNode], ) -> None: if deltas_only: if delta_file_cache.seen_previous_hash(node_hash): return raw_index = delta_file_cache.get_index(node_hash) raw_node = delta_file_cache.get_raw_node(raw_index) if isinstance(raw_node, chia_rs.datalayer.InternalNode): left_hash = delta_file_cache.get_hash_at_index(raw_node.left) right_hash = delta_file_cache.get_hash_at_index(raw_node.right) await self.get_nodes_for_file(root, left_hash, store_id, deltas_only, delta_file_cache, tree_nodes) await self.get_nodes_for_file(root, right_hash, store_id, deltas_only, delta_file_cache, tree_nodes) tree_nodes.append(SerializedNode(False, bytes(left_hash), bytes(right_hash))) elif isinstance(raw_node, chia_rs.datalayer.LeafNode): tree_nodes.append( SerializedNode( True, raw_node.key.to_bytes(), raw_node.value.to_bytes(), ) ) else: raise Exception(f"Node is neither InternalNode nor TerminalNode: {raw_node}") async def get_table_blobs( self, kv_ids_iter: Iterable[KeyOrValueId], store_id: bytes32 ) -> dict[KeyOrValueId, tuple[bytes32, bytes | None]]: result: dict[KeyOrValueId, tuple[bytes32, bytes | None]] = {} batch_size = min(500, SQLITE_MAX_VARIABLE_NUMBER - 10) kv_ids = list(dict.fromkeys(kv_ids_iter)) async with self.db_wrapper.reader() as reader: for i in range(0, len(kv_ids), batch_size): chunk = kv_ids[i : i + batch_size] placeholders = ",".join(["?"] * len(chunk)) query = f""" SELECT hash, blob, kv_id FROM ids WHERE store_id = ? AND kv_id IN ({placeholders}) LIMIT {len(chunk)} """ async with reader.execute(query, (store_id, *chunk)) as cursor: rows = await cursor.fetchall() result.update({row["kv_id"]: (row["hash"], row["blob"]) for row in rows}) if len(result) != len(kv_ids): raise Exception("Cannot retrieve all the requested kv_ids") return result async def write_tree_to_file( self, root: Root, node_hash: bytes32, store_id: bytes32, deltas_only: bool, writer: BinaryIO, ) -> None: if node_hash == bytes32.zeros: return with log_exceptions(log=log, message="Error while getting merkle blob"): root_path = self.get_merkle_path(store_id=store_id, root_hash=root.node_hash) delta_file_cache = DeltaFileCache(root_path) if root.generation > 0: previous_root = await self.get_tree_root(store_id=store_id, generation=root.generation - 1) if previous_root.node_hash is not None: with log_exceptions(log=log, message="Error while getting previous merkle blob"): previous_root_path = self.get_merkle_path(store_id=store_id, root_hash=previous_root.node_hash) delta_file_cache.load_previous_hashes(previous_root_path) tree_nodes: list[SerializedNode] = [] await self.get_nodes_for_file(root, node_hash, store_id, deltas_only, delta_file_cache, tree_nodes) kv_ids = ( KeyOrValueId.from_bytes(raw_id) for node in tree_nodes if node.is_terminal for raw_id in (node.value1, node.value2) ) table_blobs = await self.get_table_blobs(kv_ids, store_id) for node in tree_nodes: if node.is_terminal: blobs = [] for raw_id in (node.value1, node.value2): id = KeyOrValueId.from_bytes(raw_id) blob_hash, blob = table_blobs[id] if blob is None: blob = self.get_blob_from_file(blob_hash, store_id) blobs.append(blob) to_write = bytes(SerializedNode(True, blobs[0], blobs[1])) else: to_write = bytes(node) writer.write(len(to_write).to_bytes(4, byteorder="big")) writer.write(to_write) async def update_subscriptions_from_wallet(self, store_id: bytes32, new_urls: list[str]) -> None: async with self.db_wrapper.writer() as writer: cursor = await writer.execute( "SELECT * FROM subscriptions WHERE from_wallet == 1 AND tree_id == :tree_id", { "tree_id": store_id, }, ) old_urls = [row["url"] async for row in cursor] cursor = await writer.execute( "SELECT * FROM subscriptions WHERE from_wallet == 0 AND tree_id == :tree_id", { "tree_id": store_id, }, ) from_subscriptions_urls = {row["url"] async for row in cursor} additions = {url for url in new_urls if url not in old_urls} removals = [url for url in old_urls if url not in new_urls] for url in removals: await writer.execute( "DELETE FROM subscriptions WHERE url == :url AND tree_id == :tree_id", { "url": url, "tree_id": store_id, }, ) for url in additions: if url not in from_subscriptions_urls: await writer.execute( "INSERT INTO subscriptions(tree_id, url, ignore_till, num_consecutive_failures, from_wallet) " "VALUES (:tree_id, :url, 0, 0, 1)", { "tree_id": store_id, "url": url, }, ) async def subscribe(self, subscription: Subscription) -> None: async with self.db_wrapper.writer() as writer: # Add a fake subscription, so we always have the store_id, even with no URLs. await writer.execute( "INSERT INTO subscriptions(tree_id, url, ignore_till, num_consecutive_failures, from_wallet) " "VALUES (:tree_id, NULL, NULL, NULL, 0)", { "tree_id": subscription.store_id, }, ) all_subscriptions = await self.get_subscriptions() old_subscription = next( ( old_subscription for old_subscription in all_subscriptions if old_subscription.store_id == subscription.store_id ), None, ) old_urls = set() if old_subscription is not None: old_urls = {server_info.url for server_info in old_subscription.servers_info} new_servers = [server_info for server_info in subscription.servers_info if server_info.url not in old_urls] for server_info in new_servers: await writer.execute( "INSERT INTO subscriptions(tree_id, url, ignore_till, num_consecutive_failures, from_wallet) " "VALUES (:tree_id, :url, :ignore_till, :num_consecutive_failures, 0)", { "tree_id": subscription.store_id, "url": server_info.url, "ignore_till": server_info.ignore_till, "num_consecutive_failures": server_info.num_consecutive_failures, }, ) async def remove_subscriptions(self, store_id: bytes32, urls: list[str]) -> None: async with self.db_wrapper.writer() as writer: for url in urls: await writer.execute( "DELETE FROM subscriptions WHERE tree_id == :tree_id AND url == :url", { "tree_id": store_id, "url": url, }, ) async def unsubscribe(self, store_id: bytes32) -> None: async with self.db_wrapper.writer() as writer: await writer.execute( "DELETE FROM subscriptions WHERE tree_id == :tree_id", {"tree_id": store_id}, ) await writer.execute( "DELETE FROM ids WHERE store_id == :store_id", {"store_id": store_id}, ) await writer.execute( "DELETE FROM nodes WHERE store_id == :store_id", {"store_id": store_id}, ) with contextlib.suppress(FileNotFoundError): shutil.rmtree(self.get_merkle_path(store_id=store_id, root_hash=None)) with contextlib.suppress(FileNotFoundError): shutil.rmtree(self.get_key_value_path(store_id=store_id, blob_hash=None)) async def rollback_to_generation(self, store_id: bytes32, target_generation: int) -> None: async with self.db_wrapper.writer() as writer: await writer.execute( "DELETE FROM root WHERE tree_id == :tree_id AND generation > :target_generation", {"tree_id": store_id, "target_generation": target_generation}, ) await writer.execute( "DELETE FROM nodes WHERE store_id == :store_id AND generation > :target_generation", {"store_id": store_id, "target_generation": target_generation}, ) async def update_server_info(self, store_id: bytes32, server_info: ServerInfo) -> None: async with self.db_wrapper.writer() as writer: await writer.execute( "UPDATE subscriptions SET ignore_till = :ignore_till, " "num_consecutive_failures = :num_consecutive_failures WHERE tree_id = :tree_id AND url = :url", { "ignore_till": server_info.ignore_till, "num_consecutive_failures": server_info.num_consecutive_failures, "tree_id": store_id, "url": server_info.url, }, ) async def received_incorrect_file(self, store_id: bytes32, server_info: ServerInfo, timestamp: int) -> None: SEVEN_DAYS_BAN = 7 * 24 * 60 * 60 new_server_info = replace( server_info, num_consecutive_failures=server_info.num_consecutive_failures + 1, ignore_till=max(server_info.ignore_till, timestamp + SEVEN_DAYS_BAN), ) await self.update_server_info(store_id, new_server_info) async def received_correct_file(self, store_id: bytes32, server_info: ServerInfo) -> None: new_server_info = replace( server_info, num_consecutive_failures=0, ) await self.update_server_info(store_id, new_server_info) async def server_misses_file(self, store_id: bytes32, server_info: ServerInfo, timestamp: int) -> ServerInfo: # Max banned time is 1 hour. BAN_TIME_BY_MISSING_COUNT = [5 * 60] * 3 + [15 * 60] * 3 + [30 * 60] * 2 + [60 * 60] index = min(server_info.num_consecutive_failures, len(BAN_TIME_BY_MISSING_COUNT) - 1) new_server_info = replace( server_info, num_consecutive_failures=server_info.num_consecutive_failures + 1, ignore_till=max(server_info.ignore_till, timestamp + BAN_TIME_BY_MISSING_COUNT[index]), ) await self.update_server_info(store_id, new_server_info) return new_server_info async def get_available_servers_for_store(self, store_id: bytes32, timestamp: int) -> list[ServerInfo]: subscriptions = await self.get_subscriptions() subscription = next((subscription for subscription in subscriptions if subscription.store_id == store_id), None) if subscription is None: return [] servers_info = [] for server_info in subscription.servers_info: if timestamp > server_info.ignore_till: servers_info.append(server_info) return servers_info async def get_subscriptions(self) -> list[Subscription]: subscriptions: list[Subscription] = [] async with self.db_wrapper.reader() as reader: cursor = await reader.execute( "SELECT * from subscriptions", ) async for row in cursor: store_id = bytes32(row["tree_id"]) url = row["url"] ignore_till = row["ignore_till"] num_consecutive_failures = row["num_consecutive_failures"] subscription = next( (subscription for subscription in subscriptions if subscription.store_id == store_id), None ) if subscription is None: if url is not None and num_consecutive_failures is not None and ignore_till is not None: subscriptions.append( Subscription(store_id, [ServerInfo(url, num_consecutive_failures, ignore_till)]) ) else: subscriptions.append(Subscription(store_id, [])) elif url is not None and num_consecutive_failures is not None and ignore_till is not None: new_servers_info = subscription.servers_info new_servers_info.append(ServerInfo(url, num_consecutive_failures, ignore_till)) new_subscription = replace(subscription, servers_info=new_servers_info) subscriptions.remove(subscription) subscriptions.append(new_subscription) return subscriptions async def get_kv_diff( self, store_id: bytes32, # NOTE: empty is expressed as zeros hash_1: bytes32, hash_2: bytes32, ) -> set[DiffData]: async with self.db_wrapper.reader(): old_pairs = set(await self.get_keys_values(store_id, hash_1)) if len(old_pairs) == 0 and hash_1 != bytes32.zeros: raise Exception(f"Unable to diff: Can't find keys and values for {hash_1}") new_pairs = set(await self.get_keys_values(store_id, hash_2)) if len(new_pairs) == 0 and hash_2 != bytes32.zeros: raise Exception(f"Unable to diff: Can't find keys and values for {hash_2}") insertions = { DiffData(type=OperationType.INSERT, key=node.key, value=node.value) for node in new_pairs if node not in old_pairs } deletions = { DiffData(type=OperationType.DELETE, key=node.key, value=node.value) for node in old_pairs if node not in new_pairs } return set.union(insertions, deletions)