eth-portfolio 0.5.7__cp312-cp312-win32.whl

eth_portfolio/_db/utils.py

@@ -0,0 +1,616 @@
+import threading
+from asyncio import create_task, gather, get_event_loop, sleep
+from contextlib import suppress
+from decimal import getcontext
+from functools import lru_cache
+from logging import getLogger
+from typing import Any
+
+import evmspec
+import y._db.common
+import y._db.config as config
+from a_sync import PruningThreadPoolExecutor, a_sync
+from eth_typing import ChecksumAddress, HexAddress
+from evmspec.data import _decode_hook
+from msgspec import ValidationError, json
+from multicall.utils import get_event_loop
+from pony.orm import (
+    BindingError,
+    OperationalError,
+    TransactionIntegrityError,
+    commit,
+    db_session,
+    flush,
+    select,
+)
+from y import ENVIRONMENT_VARIABLES as ENVS
+from y._db.entities import db
+from y.constants import CHAINID
+from y.exceptions import reraise_excs_with_extra_context
+
+from eth_portfolio._db import entities
+from eth_portfolio._db.decorators import break_locks, requery_objs_on_diff_tx_err
+from eth_portfolio._db.entities import AddressExtended, BlockExtended, TokenExtended
+from eth_portfolio._decimal import Decimal
+from eth_portfolio.structs import InternalTransfer, TokenTransfer, Transaction, TransactionRLP
+from eth_portfolio.typing import _P, _T, Fn
+
+logger = getLogger(__name__)
+
+
+def __bind() -> None:
+    try:
+        db.bind(**config.connection_settings)
+    except BindingError as e:
+        if not str(e).startswith("Database object was already bound to"):
+            raise e
+
+
+__bind()
+
+try:
+    db.generate_mapping(create_tables=True)
+except OperationalError as e:
+    if not str(e).startswith("no such column:"):
+        raise
+    raise OperationalError(
+        "Since eth-portfolio extends the ypricemagic database with additional column definitions, you will need to delete your ypricemagic database at ~/.ypricemagic and rerun this script"
+    ) from e
+
+from y import ERC20
+from y._db.decorators import retry_locked
+from y._db.entities import Address, Block, Chain, insert
+
+# The db must be bound before we do this since we're adding some new columns to the tables defined in ypricemagic
+from y._db.utils import ensure_chain, get_chain
+from y._db.utils.price import _set_price
+from y._db.utils.traces import insert_trace
+from y.constants import EEE_ADDRESS
+from y.exceptions import NonStandardERC20
+
+_big_pool_size = 4 if ENVS.DB_PROVIDER == "sqlite" else 10
+_small_pool_size = 2 if ENVS.DB_PROVIDER == "sqlite" else 4
+
+_block_executor = PruningThreadPoolExecutor(_small_pool_size, "eth-portfolio block")
+_token_executor = PruningThreadPoolExecutor(_small_pool_size, "eth-portfolio token")
+_address_executor = PruningThreadPoolExecutor(_small_pool_size, "eth-portfolio address")
+_transaction_read_executor = PruningThreadPoolExecutor(
+    _big_pool_size, "eth-portfolio-transaction-read"
+)
+_transaction_write_executor = PruningThreadPoolExecutor(
+    _small_pool_size, "eth-portfolio-transaction-write"
+)
+_token_transfer_read_executor = PruningThreadPoolExecutor(
+    _big_pool_size, "eth-portfolio-token-transfer-read"
+)
+_token_transfer_write_executor = PruningThreadPoolExecutor(
+    _small_pool_size, "eth-portfolio-token-transfer-write"
+)
+_internal_transfer_read_executor = PruningThreadPoolExecutor(
+    _big_pool_size, "eth-portfolio-internal-transfer read"
+)
+_internal_transfer_write_executor = PruningThreadPoolExecutor(
+    _small_pool_size, "eth-portfolio-internal-transfer write"
+)
+
+
+def robust_db_session(fn: Fn[_P, _T]) -> Fn[_P, _T]:
+    return retry_locked(break_locks(db_session(fn)))
+
+
+db_session_cached = lambda func: retry_locked(
+    lru_cache(maxsize=None)(db_session(retry_locked(func)))
+)
+
+
+@a_sync(default="async", executor=_block_executor)
+@robust_db_session
+def get_block(block: int) -> BlockExtended:
+    if b := BlockExtended.get(chain=CHAINID, number=block):
+        return b
+    elif b := Block.get(chain=CHAINID, number=block):
+        if isinstance(b, BlockExtended):
+            # in case of race cndtn
+            return b
+        raise ValueError(b, b.number, b.CHAINID)
+        hash = b.hash
+        ts = b.timestamp
+        prices = [(price.token.address, price.price) for price in b.prices]
+        logs = [json.decode(log.raw) for log in b.logs]
+        traces = [json.decode(trace.raw) for trace in b.traces]
+        for p in b.prices:
+            p.delete()
+        for l in b.logs:
+            l.delete()
+        for t in b.traces:
+            t.delete()
+        flush()
+        b.delete()
+        commit()
+        b = insert(
+            type=BlockExtended,
+            chain=get_chain(sync=True),
+            number=block,
+            hash=hash,
+            timestamp=ts,
+        )
+        try:
+            # for log in logs:
+            #    insert_log(log)
+            for trace in traces:
+                insert_trace(trace)
+        except Exception as e:
+            e.args = (
+                "This is really bad. Might need to nuke your db if you value your logs/traces",
+                *e.args,
+            )
+            raise
+        for token, price in prices:
+            _set_price(token, price, sync=True)
+    asdasd = get_chain(sync=True)
+    if not isinstance(asdasd, Chain):
+        raise TypeError(asdasd)
+    commit()
+    if b := insert(type=BlockExtended, chain=asdasd, number=block):
+        return b
+    return BlockExtended.get(chain=CHAINID, number=block)
+
+
+@a_sync(default="async", executor=_block_executor)
+@db_session_cached
+def ensure_block(block: int) -> None:
+    get_block(block, sync=True)
+
+
+# TODO refactor this out, async is annoying sometimes
+# process = ProcessPoolExecutor(
+#  1,
+#  # NOTE: come on apple, what are you dooooin?
+#  mp_context=get_context('fork'),
+# )
+
+
+def is_token(address: ChecksumAddress) -> bool:
+    if address == EEE_ADDRESS:
+        return False
+    # with suppress(NonStandardERC20):
+    #    erc = ERC20(address)
+    #    if all(erc.symbol, erc.name, erc.total_supply(), erc.scale):
+    #    #if all(erc._symbol(), erc._name(), erc.total_supply(), erc._scale()):
+    #        return True
+    # return False
+    return get_event_loop().run_until_complete(_is_token(address))
+
+
+async def _is_token(address: HexAddress) -> bool:
+    # just breaking a weird lock, dont mind me
+    if retval := await get_event_loop().run_in_executor(process, __is_token, address):  # type: ignore [name-defined]
+        logger.debug("%s is token")
+    else:
+        logger.debug("%s is not token")
+    return retval
+
+
+def __is_token(address: HexAddress) -> bool:
+    with suppress(NonStandardERC20):
+        erc = ERC20(address, asynchronous=True)
+        if all(
+            get_event_loop().run_until_complete(
+                gather(erc._symbol(), erc._name(), erc.total_supply_readable())
+            )
+        ):
+            return True
+    return False
+
+
+@a_sync(default="async", executor=_address_executor)
+@robust_db_session
+def get_address(address: ChecksumAddress) -> AddressExtended:
+    entity_type = TokenExtended
+    entity = entities.Address.get(chain=CHAINID, address=address)
+    """ TODO: fix this later
+    entity = entities.Address.get(chain=chain, address=address)
+    if isinstance(entity, (Token, TokenExtended)):
+        entity_type = TokenExtended
+    elif isinstance(entity, (Contract, ContractExtended)):
+        entity_type = ContractExtended
+    elif isinstance(entity, (Address, AddressExtended)):
+        entity_type = AddressExtended
+    elif entity is None:
+        # TODO: this logic should live in ypm, prob
+        entity_type = AddressExtended if not is_contract(address) else TokenExtended if is_token(address) else ContractExtended
+    else:
+        raise NotImplementedError(entity, entity_type)
+        
+    if isinstance(entity, entity_type):
+        return entity
+    
+    elif entity:
+        logger.debug("deleting %s", entity)
+        entity.delete()
+        commit()
+    """
+    if entity := entities.Address.get(chain=CHAINID, address=address):
+        return entity
+
+    ensure_chain()
+    return insert(type=entity_type, chain=CHAINID, address=address) or entity_type.get(
+        chain=CHAINID, address=address
+    )
+
+
+@a_sync(default="async", executor=_address_executor)
+@db_session_cached
+def ensure_address(address: ChecksumAddress) -> None:
+    get_address(address, sync=True)
+
+
+@a_sync(default="async", executor=_address_executor)
+@db_session_cached
+def ensure_addresses(*addresses: ChecksumAddress) -> None:
+    for address in addresses:
+        ensure_address(address, sync=True)
+
+
+@a_sync(default="async", executor=_token_executor)
+@robust_db_session
+def get_token(address: ChecksumAddress) -> TokenExtended:
+    if t := TokenExtended.get(chain=CHAINID, address=address):
+        return t
+    kwargs = {}
+    if t := Address.get(chain=CHAINID, address=address):
+        if isinstance(t, TokenExtended):
+            # double check due to possible race cntdn
+            return t
+        """
+        with suppress(TypeError):
+            if t.notes:
+                kwargs['notes'] = t.notes
+        if isinstance(t, Contract):
+            with suppress(TypeError):
+                if t.deployer:
+                    kwargs['deployer'] = t.deployer
+            with suppress(TypeError):
+                if t.deploy_block:
+                    kwargs['deploy_block'] = t.deploy_block
+        if isinstance(t, Token):
+            with suppress(TypeError):
+                if t.symbol:
+                    kwargs['symbol'] = t.symbol
+            with suppress(TypeError):
+                if t.name:
+                    kwargs['name'] = t.name
+            with suppress(TypeError):
+                if t.bucket:
+                    kwargs['bucket'] = t.bucket
+        """
+
+        try:
+            flush()
+            t.delete()
+            commit()
+        except KeyError as e:
+            raise KeyError(f"cant delete {t}") from e
+
+    ensure_chain()
+    commit()
+    return insert(
+        type=TokenExtended, chain=CHAINID, address=address, **kwargs
+    ) or TokenExtended.get(chain=CHAINID, address=address)
+
+
+@a_sync(default="async", executor=_token_executor)
+@db_session_cached
+def ensure_token(token_address: ChecksumAddress) -> None:
+    get_token(token_address, sync=True)
+
+
+async def get_transaction(sender: ChecksumAddress, nonce: int) -> Transaction | None:
+    startup_txs = await transactions_known_at_startup(CHAINID, sender)
+    data = startup_txs.pop(nonce, None) or await __get_transaction_bytes_from_db(sender, nonce)
+    if data:
+        await _yield_to_loop()
+        return decode_transaction(data)
+
+
+_decoded = 0
+
+
+async def _yield_to_loop() -> None:
+    """dont let the event loop get congested, let your rpc begin work asap"""
+    global _decoded
+    _decoded += 1
+    if _decoded % 1000 == 0:
+        await sleep(0)
+
+
+@a_sync(default="async", executor=_transaction_read_executor)
+@robust_db_session
+def __get_transaction_bytes_from_db(sender: ChecksumAddress, nonce: int) -> bytes | None:
+    entity: entities.Transaction
+    if entity := entities.Transaction.get(from_address=(CHAINID, sender), nonce=nonce):
+        return entity.raw
+
+
+def decode_transaction(data: bytes) -> Transaction | TransactionRLP:
+    try:
+        try:
+            return json.decode(data, type=Transaction, dec_hook=_decode_hook)
+        except ValidationError as e:
+            if str(e) == "Object missing required field `type` - at `$[2]`":
+                return json.decode(data, type=TransactionRLP, dec_hook=_decode_hook)
+            raise
+    except Exception as e:
+        e.args = *e.args, json.decode(data)
+        raise
+
+
+@a_sync(default="async", executor=_transaction_write_executor)
+@robust_db_session
+def delete_transaction(transaction: Transaction) -> None:
+    if entity := entities.Transaction.get(**transaction.__db_primary_key__):
+        entity.delete()
+
+
+async def insert_transaction(transaction: Transaction) -> None:
+    # Make sure these are in the db so below we can call them and use the results all in one transaction
+    # NOTE: this create task -> await coro -> await task pattern is faster than a 2-task gather
+    block_task = create_task(ensure_block(transaction.block_number))
+    if to_address := transaction.to_address:
+        await ensure_addresses(to_address, transaction.from_address)
+    else:
+        await ensure_address(transaction.from_address)
+    await block_task
+    await _insert_transaction(transaction)
+
+
+@a_sync(default="async", executor=_transaction_write_executor)
+@requery_objs_on_diff_tx_err
+@robust_db_session
+def _insert_transaction(transaction: Transaction) -> None:
+    with reraise_excs_with_extra_context(transaction):
+        entities.Transaction(
+            **transaction.__db_primary_key__,
+            block=(CHAINID, transaction.block_number),
+            transaction_index=transaction.transaction_index,
+            hash=transaction.hash.hex(),
+            to_address=(CHAINID, transaction.to_address) if transaction.to_address else None,
+            value=transaction.value,
+            price=transaction.price,
+            value_usd=transaction.value_usd,
+            type=getattr(transaction, "type", None),
+            gas=transaction.gas,
+            gas_price=transaction.gas_price,
+            max_fee_per_gas=getattr(transaction, "max_fee_per_gas", None),
+            max_priority_fee_per_gas=getattr(transaction, "max_priority_fee_per_gas", None),
+            raw=json.encode(transaction, enc_hook=enc_hook),
+        )
+
+
+@a_sync(default="async", executor=_internal_transfer_read_executor)
+@robust_db_session
+def get_internal_transfer(trace: evmspec.FilterTrace) -> InternalTransfer | None:
+    block = trace.blockNumber
+    entity: entities.InternalTransfer
+    if entity := entities.InternalTransfer.get(
+        block=(CHAINID, block),
+        transaction_index=trace.transactionPosition,
+        hash=trace.transactionHash,
+        type=trace.type.name,
+        call_type=trace.callType,
+        from_address=(CHAINID, trace.sender),
+        to_address=(CHAINID, trace.to),
+        value=trace.value.scaled,
+        trace_address=(CHAINID, trace.traceAddress),
+        gas=trace.gas,
+        gas_used=trace.gasUsed if "gasUsed" in trace else None,
+        input=trace.input,
+        output=trace.output,
+        subtraces=trace.subtraces,
+        address=(CHAINID, trace.address),
+    ):
+        return json.decode(entity.raw, type=InternalTransfer, dec_hook=_decode_hook)
+
+
+@a_sync(default="async", executor=_internal_transfer_write_executor)
+@robust_db_session
+def delete_internal_transfer(transfer: InternalTransfer) -> None:
+    if entity := entities.InternalTransfer.get(
+        block=(CHAINID, transfer.block_number),
+        transaction_index=transfer.transaction_index,
+        hash=transfer.hash,
+        type=transfer.type,
+        call_type=transfer.call_type,
+        from_address=(CHAINID, transfer.from_address),
+        to_address=(CHAINID, transfer.to_address),
+        value=transfer.value,
+        trace_address=(CHAINID, transfer.trace_address),
+        gas=transfer.gas,
+        gas_used=transfer.gas_used,
+        input=transfer.input,
+        output=transfer.output,
+        subtraces=transfer.subtraces,
+        address=(CHAINID, transfer.address),
+    ):
+        entity.delete()
+
+
+async def insert_internal_transfer(transfer: InternalTransfer) -> None:
+    # NOTE: this create task -> await coro -> await task pattern is faster than a 2-task gather
+    block_task = create_task(ensure_block(transfer.block_number))
+    if to_address := getattr(transfer, "to_address", None):
+        await ensure_addresses(to_address, transfer.from_address)
+    else:
+        await ensure_address(transfer.from_address)
+    await block_task
+    await _insert_internal_transfer(transfer)
+
+
+@a_sync(default="async", executor=_internal_transfer_write_executor)
+@robust_db_session
+def _insert_internal_transfer(transfer: InternalTransfer) -> None:
+    entities.InternalTransfer(
+        block=(CHAINID, transfer.block_number),
+        transaction_index=transfer.transaction_index,
+        hash=transfer.hash,
+        type=transfer.type,
+        call_type=transfer.call_type,
+        from_address=(CHAINID, transfer.from_address),
+        to_address=(CHAINID, transfer.to_address),
+        value=transfer.value,
+        price=transfer.price,
+        value_usd=transfer.value_usd,
+        trace_address=str(transfer.trace_address),
+        gas=transfer.gas,
+        gas_used=transfer.gas_used,
+        raw=json.encode(transfer, enc_hook=enc_hook),
+    )
+
+
+async def get_token_transfer(transfer: evmspec.Log) -> TokenTransfer | None:
+    pk = {
+        "block": (CHAINID, transfer.blockNumber),
+        "transaction_index": transfer.transactionIndex,
+        "log_index": transfer.logIndex,
+    }
+    startup_xfers = await token_transfers_known_at_startup(CHAINID)
+    data_at_startup = startup_xfers.pop(tuple(pk.values()), None)
+    try:
+        data = data_at_startup or await __get_token_transfer_bytes_from_db(pk)
+    except TypeError:
+        # NoneType object is not subscriptable
+        # I'm not sure why this happens on occasion, I think it can occur if the script is interrupted mid-insert
+        # Doesn't really matter why, we can just fetch from the chain and insert the correct data now.
+        return None
+
+    if data:
+        await _yield_to_loop()
+        with reraise_excs_with_extra_context(data):
+            return json.decode(data, type=TokenTransfer, dec_hook=_decode_hook)
+
+
+@a_sync(default="async", executor=_token_transfer_read_executor)
+@robust_db_session
+def __get_token_transfer_bytes_from_db(pk: dict) -> bytes | None:
+    entity: entities.TokenTransfer
+    if entity := entities.TokenTransfer.get(**pk):
+        return entity.raw
+
+
+_TPK = tuple[tuple[int, ChecksumAddress], int]
+
+
+@a_sync(default="async", executor=_transaction_read_executor, ram_cache_maxsize=None)
+@robust_db_session
+def transactions_known_at_startup(chainid: int, from_address: ChecksumAddress) -> dict[_TPK, bytes]:
+    return dict(
+        select(
+            (t.nonce, t.raw)
+            for t in entities.Transaction  # type: ignore [attr-defined]
+            if t.from_address.chain.id == chainid and t.from_address.address == from_address
+        )
+    )
+
+
+_TokenTransferPK = tuple[tuple[int, int], int, int]
+
+
+@a_sync(default="async", executor=_transaction_read_executor, ram_cache_maxsize=None)
+@robust_db_session
+def token_transfers_known_at_startup(chainid: int) -> dict[_TokenTransferPK, bytes]:
+    block: int
+    tx_index: int
+    log_index: int
+    raw: bytes
+
+    transfers = {}
+    for block, tx_index, log_index, raw in select(
+        (t.block.number, t.transaction_index, t.log_index, t.raw)
+        for t in entities.TokenTransfer  # type: ignore [attr-defined]
+        if t.block.chain.id == chainid
+    ):
+        pk = ((chainid, block), tx_index, log_index)
+        transfers[pk] = raw
+    return transfers
+
+
+@a_sync(default="async", executor=_token_transfer_write_executor)
+@robust_db_session
+def delete_token_transfer(token_transfer: TokenTransfer) -> None:
+    if entity := entities.TokenTransfer.get(
+        block=(CHAINID, token_transfer.block_number),
+        transaction_index=token_transfer.transaction_index,
+        log_index=token_transfer.log_index,
+    ):
+        entity.delete()
+
+
+async def insert_token_transfer(token_transfer: TokenTransfer) -> None:
+    # two tasks and a coroutine like this should be faster than gather
+    block_task = create_task(ensure_block(token_transfer.block_number))
+    while True:
+        try:
+            token_task = create_task(ensure_token(token_transfer.token_address))
+        except KeyError:
+            # This KeyError comes from a bug in cachetools.ttl_cache
+            # TODO: move this handler into evmspec
+            pass
+        else:
+            break
+
+    while True:
+        try:
+            address_coro = ensure_addresses(token_transfer.to_address, token_transfer.from_address)
+        except KeyError:
+            # This KeyError comes from a bug in cachetools.ttl_cache
+            # TODO: move this handler into evmspec
+            pass
+        else:
+            break
+
+    await address_coro
+    await block_task
+    await token_task
+    await _insert_token_transfer(token_transfer)
+
+
+__tt_lock = threading.Lock()
+
+
+@a_sync(default="async", executor=_token_transfer_write_executor)
+@requery_objs_on_diff_tx_err
+@robust_db_session
+def _insert_token_transfer(token_transfer: TokenTransfer) -> None:
+    with __tt_lock:
+        # we need to temporarily set the Context for Decimal objects to ensure we have enough precision for our postgres column
+        decimal_context = getcontext()
+        decimal_precision = decimal_context.prec
+        decimal_context.prec = 38
+        try:
+            entities.TokenTransfer(
+                block=(CHAINID, token_transfer.block_number),
+                transaction_index=token_transfer.transaction_index,
+                log_index=token_transfer.log_index,
+                hash=token_transfer.hash.hex(),
+                token=(CHAINID, token_transfer.token_address),
+                from_address=(CHAINID, token_transfer.from_address),
+                to_address=(CHAINID, token_transfer.to_address),
+                value=token_transfer.value,
+                price=token_transfer.price,
+                value_usd=token_transfer.value_usd,
+                raw=json.encode(token_transfer, enc_hook=enc_hook),
+            )
+            commit()
+        except TransactionIntegrityError:
+            pass  # most likely non-issue, debug later if needed
+        finally:
+            decimal_context.prec = decimal_precision
+
+
+def enc_hook(obj: Any) -> Any:
+    try:
+        return y._db.common.enc_hook(obj)
+    except TypeError:
+        if type(obj) is Decimal:
+            return obj.jsonify()
+        raise TypeError(type(obj), obj) from None