feat: cache prices, LogFilter and TraceFilter classes

requirements.txt

@@ -4,7 +4,7 @@ checksum_dict>=1.1.1
 dank_mids>=4.20.50
 eth-brownie>=1.18.1,<1.20
 eth_retry>=0.1.17,<0.2
-ez-a-sync>=0.7.4
+ez-a-sync>=0.9.0
 joblib>=1.0.1
 multicall>=0.8.1
 pony

tests/prices/test_chainlink.py

@@ -124,7 +124,7 @@ def test_chainlink_get_feed(token):
 async def test_chainlink_latest(token):
     if not await chainlink.get_price(token):
         feed = await chainlink.get_feed(token)
-        assert await feed.aggregator.coroutine() == ZERO_ADDRESS, 'no current price available'
+        assert await feed.contract.aggregator.coroutine() == ZERO_ADDRESS, 'no current price available'
 
 
 @mainnet_only
@@ -136,10 +136,12 @@ async def test_chainlink_before_registry(token):
     feed = await chainlink.get_feed(token, sync=False)
     if await contract_creation_block_async(feed.address) > test_block:
         pytest.skip('not applicable to feeds deployed after test block')
-    price = await chainlink.get_price(token, block=test_block)
+    print(type(chainlink.get_price))
+    price = chainlink.get_price(token, block=test_block)
+    price = await price
     if not price:
         feed = await chainlink.get_feed(token)
-        assert await feed.aggregator.coroutine() == ZERO_ADDRESS, 'no price available before registry'
+        assert await feed.contract.aggregator.coroutine() == ZERO_ADDRESS, 'no price available before registry'
 
 
 def test_chainlink_nonexistent():

y/_db/common.py

@@ -0,0 +1,235 @@
+
+import abc
+import asyncio
+import logging
+from concurrent.futures import ThreadPoolExecutor
+from typing import (Any, AsyncIterator, Callable, Generic, List, NoReturn,
+                    Optional, Type, TypeVar)
+
+from a_sync.iter import ASyncIterable
+from a_sync.primitives.executor import _AsyncExecutorMixin
+from a_sync.primitives.locks.counter import CounterLock
+from dank_mids.semaphores import BlockSemaphore
+from hexbytes import HexBytes
+from tqdm.asyncio import tqdm_asyncio
+from web3.datastructures import AttributeDict
+from web3.middleware.filter import block_ranges
+
+from y._db.exceptions import CacheNotPopulatedError
+from y.constants import BIG_VALUE, thread_pool_executor
+from y.utils.dank_mids import dank_w3
+
+T = TypeVar('T')
+S = TypeVar('S')
+M = TypeVar('M')
+
+logger = logging.getLogger(__name__)
+executor = ThreadPoolExecutor(16)
+
+def enc_hook(obj: Any) -> bytes:
+    if isinstance(obj, AttributeDict):
+        return dict(obj)
+    elif isinstance(obj, HexBytes):
+        return obj.hex()
+    raise NotImplementedError(obj)
+
+def dec_hook(typ: Type[T], obj: bytes) -> T:
+    if typ == HexBytes:
+        return HexBytes(obj)
+    raise ValueError(f"{typ} is not a valid type for decoding")
+
+class DiskCache(Generic[S, M], metaclass=abc.ABCMeta):
+    __slots__ = []
+    @abc.abstractmethod
+    def set_metadata(self, from_block: int, done_thru: int) -> None:
+        """Updates the cache metadata to indicate the cache is populated from block `from_block` to block `to_block`."""
+    @abc.abstractmethod
+    def is_cached_thru(self, from_block: int) -> int:
+        """Returns max cached block for this cache or 0 if not cached."""
+    @abc.abstractmethod
+    def select(self, from_block: int, to_block: int) -> List[S]:
+        """Selects all cached objects from block `from_block` to block `to_block`"""
+    def check_and_select(self, from_block: int, to_block: int) -> List[S]:
+        """
+        Selects all cached objects from block `from_block` to block `to_block` if the cache is fully populated.
+        Raises `CacheNotPopulatedError` if it is not.
+        """
+        if self.is_cached_thru(from_block) >= to_block:
+            return self.select(from_block, to_block)
+        else:
+            raise CacheNotPopulatedError(self, from_block, to_block)
+
+C = TypeVar('C', bound=DiskCache)
+
+class _DiskCachedMixin(Generic[T, C], metaclass=abc.ABCMeta):
+    __slots__ = 'is_reusable', '_cache', '_executor', '_objects', '_pruned'
+    def __init__(
+        self, 
+        executor: _AsyncExecutorMixin = thread_pool_executor,
+        is_reusable: bool = True,
+    ):
+        self.is_reusable = is_reusable
+        self._cache = None
+        self._executor = executor
+        self._objects: List[T] = []
+        self._pruned = 0
+    @abc.abstractproperty
+    def cache(self) -> C:
+        ...
+    @abc.abstractproperty
+    def insert_to_db(self) -> Callable[[T], None]:
+        ...
+    def _extend(self, objs) -> None:
+        """Override this to pre-process objects before storing."""
+        return self._objects.extend(objs)
+    def _remove(self, obj: T) -> None:
+        self._objects.remove(obj)
+        self._pruned += 1
+    async def _load_cache(self, from_block: int) -> int:
+        """
+        Loads cached logs from disk.
+        Returns max block of logs loaded from cache.
+        """
+        logger.debug('checking to see if %s is cached in local db', self)
+        if cached_thru := await self._executor.run(self.cache.is_cached_thru, from_block):
+            logger.debug('%s is cached thru block %s, loading from db', self, cached_thru)
+            self._extend(await self._executor.run(self.cache.select, from_block, cached_thru))
+            logger.info('%s loaded %s objects thru block %s from disk', self, len(self._objects), cached_thru)
+            return cached_thru
+        return from_block - 1
+
+class Filter(ASyncIterable[T], _DiskCachedMixin[T, C]):
+    __slots__ = 'from_block', 'to_block', '_chunk_size', '_chunks_per_batch', '_exc', '_interval', '_lock', '_semaphore', '_task', '_verbose'
+    def __init__(
+        self, 
+        from_block: int,
+        *, 
+        chunk_size: int = 10_000, 
+        chunks_per_batch: int = 20,
+        interval: int = 300, 
+        semaphore: Optional[BlockSemaphore] = None,
+        executor: _AsyncExecutorMixin = thread_pool_executor,
+        is_reusable: bool = True,
+        verbose: bool = False,
+    ):
+        self.from_block = from_block
+        self._chunk_size = chunk_size
+        self._chunks_per_batch = chunks_per_batch
+        self._exc = None
+        self._interval = interval
+        self._lock = CounterLock()
+        self._semaphore = semaphore
+        self._task = None
+        self._verbose = verbose
+        super().__init__(executor=executor, is_reusable=is_reusable)
+
+    def __aiter__(self) -> AsyncIterator[T]:
+        return self._objects_thru(block=None).__aiter__()
+
+    @abc.abstractmethod
+    async def _fetch_range(self, from_block: int, to_block: int) -> List[T]:
+        ...
+
+    @property
+    def semaphore(self) -> BlockSemaphore:
+        if self._semaphore is None:
+            self._semaphore = BlockSemaphore(self._chunks_per_batch)
+        return self._semaphore
+
+    def _get_block_for_obj(self, obj: T) -> int:
+        """Override this as needed for different object types"""
+        return obj['blockNumber']
+
+    async def _objects_thru(self, block: Optional[int]) -> AsyncIterator[T]:
+        self._ensure_task()
+        yielded = 0
+        done_thru = 0
+        while True:
+            if block is None or done_thru < block:
+                await self._lock.wait_for(done_thru + 1)
+            if self._exc:
+                raise self._exc
+            if to_yield := self._objects[yielded-self._pruned:]:
+                for obj in to_yield:
+                    if block and self._get_block_for_obj(obj) > block:
+                        return
+                    if not self.is_reusable:
+                        self._remove(obj)
+                    yield obj
+                    yielded += 1
+            elif block and done_thru > block:
+                return
+            done_thru = self._lock.value
+
+    async def __fetch(self) -> NoReturn:
+        try:
+            await self._fetch()
+        except Exception as e:
+            logger.exception(e)
+            self._exc = e
+            self._lock.set(BIG_VALUE)
+            raise e
+
+    async def _fetch(self) -> NoReturn:
+        """Override this if you want"""
+        await self._loop(self.from_block)
+
+    async def _fetch_range_wrapped(self, i: int, range_start: int, range_end: int) -> List[T]:
+        async with self.semaphore[range_end]:
+            logger.debug("fetching %s block %s to %s", self, range_start, range_end)
+            return i, range_end, await self._fetch_range(range_start, range_end)
+
+    async def _loop(self, from_block: int) -> NoReturn:
+        logger.debug('starting work loop for %s', self)
+        self._lock.set(await self._load_cache(from_block))
+        while True:
+            await self._load_new_objects(start_from_block=from_block)
+            await asyncio.sleep(self._interval)
+
+    async def _load_new_objects(self, to_block: Optional[int] = None, start_from_block: Optional[int] = None) -> None:
+        logger.debug('loading new objects for %s', self)
+        start = v + 1 if (v := self._lock.value) else start_from_block or self.from_block
+        end = to_block or await dank_w3.eth.block_number
+        await self._load_range(start, end)
+
+    async def _load_range(self, from_block: int, to_block: int) -> None:
+        logger.debug('loading block range %s to %s', from_block, to_block)
+        db_insert_tasks = []
+        cache_info_tasks = []
+        batches_yielded = 0
+        done = {}
+        as_completed = tqdm_asyncio.as_completed if self._verbose else asyncio.as_completed
+        coros = [
+            self._fetch_range_wrapped(i, start, end) 
+            for i, (start, end) 
+            in enumerate(block_ranges(from_block, to_block, self._chunk_size))
+            if i <= self._chunks_per_batch
+        ]
+        for objs in as_completed(coros, timeout=None):
+            i, end, objs = await objs
+            done[i] = end, objs
+            for i in range(len(coros)):
+                if batches_yielded > i:
+                    continue
+                if i not in done:
+                    if db_insert_tasks:
+                        await asyncio.gather(*db_insert_tasks)
+                        db_insert_tasks.clear()
+                    if cache_info_tasks:
+                        await cache_info_tasks[-1]
+                        cache_info_tasks.clear()
+                    break
+                end, objs = done.pop(i)
+                self._extend(objs)
+                db_insert_tasks.extend(self._executor.run(self.insert_to_db, obj) for obj in objs)
+                cache_info_tasks.append(self._executor.run(self.cache.set_metadata, from_block, end))
+                batches_yielded += 1
+                self._lock.set(end)
+
+    def _ensure_task(self) -> None:
+        if self._task is None:
+            logger.debug('creating task for %s', self)
+            self._task = asyncio.create_task(self.__fetch())
+        if self._task.done() and (e := self._task.exception()):
+            raise e
+

y/_db/entities.py

@@ -1,8 +1,12 @@
 
+from contextlib import suppress
 from datetime import datetime
+from decimal import Decimal
+from typing import Any
 
 from pony.orm import (Database, Optional, PrimaryKey, Required, Set,
-                      composite_key)
+                      TransactionIntegrityError, commit, composite_key,
+                      db_session)
 
 db = Database()
 
@@ -27,6 +31,8 @@ class Chain(db.Entity, _AsyncEntityMixin):
 
     blocks = Set("Block")
     addresses = Set("Address")
+    log_cached = Set("LogCacheInfo")
+    trace_caches = Set("TraceCacheInfo")
 
 class Block(db.Entity, _AsyncEntityMixin):
     _pk = PrimaryKey(int, auto=True)
@@ -38,7 +44,10 @@ class Block(db.Entity, _AsyncEntityMixin):
     composite_key(chain, number)
     composite_key(chain, hash)
 
+    prices = Set("Price", reverse="block", cascade_delete=False)
     contracts_deployed = Set("Contract", reverse="deploy_block")
+    logs = Set("Log", reverse="block", cascade_delete=False)
+    traces = Set("Trace", reverse="block", cascade_delete=False)
 
 class Address(db.Entity, _AsyncEntityMixin):
     _pk = PrimaryKey(int, auto=True)
@@ -51,11 +60,62 @@ class Address(db.Entity, _AsyncEntityMixin):
     contracts_deployed = Set("Contract", reverse="deployer")
 
 class Contract(Address):
-    deployer = Optional(Address, reverse='contracts_deployed', lazy=True)
-    deploy_block = Optional(Block, reverse='contracts_deployed', lazy=True)
+    deployer = Optional(Address, reverse='contracts_deployed', lazy=True, cascade_delete=False)
+    deploy_block = Optional(Block, reverse='contracts_deployed', lazy=True, cascade_delete=False)
 
 class Token(Contract):
     symbol = Optional(str, lazy=True)
     name = Optional(str, lazy=True)
     decimals = Optional(int, lazy=True)
-    bucket = Optional(str, lazy=True)
+    bucket = Optional(str, index=True, lazy=True)
+
+    prices = Set("Price", reverse="token")
+
+class Price(db.Entity):
+    dbid = PrimaryKey(int, auto=True)
+    block = Required(Block, index=True, lazy=True)
+    token = Required(Token, index=True, lazy=True)
+    composite_key(block, token)
+    price = Required(Decimal)
+
+class TraceCacheInfo(db.Entity):
+    chain = Required(Chain, index=True)
+    to_addresses = Required(bytes, index=True)
+    from_addresses = Required(bytes, index=True)
+    composite_key(chain, to_addresses, from_addresses)
+    cached_from = Required(int)
+    cached_thru = Required(int)
+
+class LogCacheInfo(db.Entity):
+    chain = Required(Chain, index=True)
+    address = Required(str, index=True)
+    topics = Required(bytes)
+    composite_key(chain, address, topics)
+    cached_from = Required(int)
+    cached_thru = Required(int)
+
+class Log(db.Entity):
+    block = Required(Block, index=True, lazy=True)
+    transaction_hash = Required(str, lazy=True)
+    log_index = Required(int, lazy=True)
+    composite_key(block, transaction_hash, log_index)
+
+    address = Required(str, index=True, lazy=True)
+    topic0 = Required(str, index=True, lazy=True)
+    topic1 = Optional(str, index=True, lazy=True)
+    topic2 = Optional(str, index=True, lazy=True)
+    topic3 = Optional(str, index=True, lazy=True)
+    raw = Required(bytes, lazy=True)
+
+class Trace(db.Entity):
+    block = Required(Block, index=True, lazy=True)
+    hash = Required(str, index=True, lazy=True)
+    from_address = Required(str, index=True, lazy=True)
+    to_address = Required(str, index=True, lazy=True)
+    raw = Required(bytes)
+
+@db_session
+def insert(type: db.Entity, **kwargs: Any) -> None:
+    with suppress(TransactionIntegrityError):
+        type(**kwargs)
+        commit()

y/_db/exceptions.py

@@ -0,0 +1,3 @@
+
+class CacheNotPopulatedError(Exception):
+    pass

y/_db/utils/__init__.py

@@ -0,0 +1,18 @@
+
+from pony.orm import BindingError, TransactionError
+
+from y._db.config import connection_settings
+from y._db.entities import db
+from y._db.utils.utils import get_chain, get_block
+
+try:
+    db.bind(**connection_settings, create_db=True)
+    db.generate_mapping(create_tables=True)
+except TransactionError as e:
+    if str(e) != "@db_session-decorated create_tables() function with `ddl` option cannot be called inside of another db_session":
+        raise e
+except BindingError as e:
+    if not str(e).startswith('Database object was already bound to'):
+        raise e
+
+__all__ = 'get_chain', 'get_block'