feat: Better handling of REMOVING status + balance update fixes + cost deletion on FORGOTTEN / REMOVED messages

Author: amalcaraz

src/aleph/config.py

@@ -41,7 +41,6 @@ def get_defaults():
                 "cron": {
                     # Interval between cron job trackers runs, expressed in hours.
                     "period": 0.5,  # 30 mins
-                    # "period": 0.002,  # 10 seconds
                 },
             },
             "cache": {

src/aleph/db/accessors/balances.py

@@ -8,6 +8,7 @@
 from sqlalchemy.sql import Select
 
 from aleph.db.models import AlephBalanceDb
+from aleph.toolkit.timestamp import utc_now
 from aleph.types.db_session import DbSession
 
 
@@ -141,6 +142,8 @@ def update_balances(
     table from the temporary one.
     """
 
+    last_update = utc_now()
+
     session.execute(
         "CREATE TEMPORARY TABLE temp_balances AS SELECT * FROM balances WITH NO DATA"  # type: ignore[arg-type]
     )
@@ -152,21 +155,21 @@ def update_balances(
     csv_balances = StringIO(
         "\n".join(
             [
-                f"{address};{chain.value};{dapp or ''};{balance};{eth_height}"
+                f"{address};{chain.value};{dapp or ''};{balance};{eth_height};{last_update}"
                 for address, balance in balances.items()
             ]
         )
     )
     cursor.copy_expert(
-        "COPY temp_balances(address, chain, dapp, balance, eth_height) FROM STDIN WITH CSV DELIMITER ';'",
+        "COPY temp_balances(address, chain, dapp, balance, eth_height, last_update) FROM STDIN WITH CSV DELIMITER ';'",
         csv_balances,
     )
     session.execute(
         """
-        INSERT INTO balances(address, chain, dapp, balance, eth_height)
-            (SELECT address, chain, dapp, balance, eth_height FROM temp_balances)
+        INSERT INTO balances(address, chain, dapp, balance, eth_height, last_update)
+            (SELECT address, chain, dapp, balance, eth_height, last_update FROM temp_balances)
             ON CONFLICT ON CONSTRAINT balances_address_chain_dapp_uindex DO UPDATE
-            SET balance = excluded.balance, eth_height = excluded.eth_height
+            SET balance = excluded.balance, eth_height = excluded.eth_height, last_update = (CASE WHEN excluded.balance <> balances.balance THEN excluded.last_update ELSE balances.last_update END)
             WHERE excluded.eth_height > balances.eth_height
         """  # type: ignore[arg-type]
     )
@@ -177,8 +180,10 @@ def update_balances(
     session.execute("DROP TABLE temp_balances")  # type: ignore[arg-type]
 
 
-def get_updated_balances(session: DbSession, last_update: dt.datetime):
-    select_stmt = select(AlephBalanceDb.address, AlephBalanceDb.balance).filter(
-        AlephBalanceDb.last_update >= last_update
+def get_updated_balance_accounts(session: DbSession, last_update: dt.datetime):
+    select_stmt = (
+        select(AlephBalanceDb.address)
+        .where(AlephBalanceDb.last_update >= last_update)
+        .distinct()
     )
-    return (session.execute(select_stmt)).all()
+    return (session.execute(select_stmt)).scalars().all()

src/aleph/db/accessors/cost.py

@@ -2,14 +2,16 @@
 from typing import Iterable, List, Optional
 
 from aleph_message.models import PaymentType
-from sqlalchemy import asc, func, select
+from sqlalchemy import asc, delete, func, select
 from sqlalchemy.dialects.postgresql import insert
 from sqlalchemy.sql import Insert
 
 from aleph.db.models import ChainTxDb, message_confirmations
 from aleph.db.models.account_costs import AccountCostsDb
+from aleph.db.models.messages import MessageStatusDb
 from aleph.toolkit.costs import format_cost
 from aleph.types.db_session import DbSession
+from aleph.types.message_status import MessageStatus
 
 
 def get_total_cost_for_address(
@@ -94,3 +96,21 @@ def make_costs_upsert_query(costs: List[AccountCostsDb]) -> Insert:
             "cost_stream": upsert_stmt.excluded.cost_stream,
         },
     )
+
+
+def delete_costs_for_message(session: DbSession, item_hash: str) -> None:
+    delete_stmt = delete(AccountCostsDb).where(AccountCostsDb.item_hash == item_hash)
+    session.execute(delete_stmt)
+
+
+def delete_costs_for_forgotten_and_deleted_messages(session: DbSession) -> None:
+    delete_stmt = (
+        delete(AccountCostsDb)
+        .where(AccountCostsDb.item_hash == MessageStatusDb.item_hash)
+        .where(
+            (MessageStatusDb.status == MessageStatus.FORGOTTEN)
+            | (MessageStatusDb.status == MessageStatus.REMOVED)
+        )
+        .execution_options(synchronize_session=False)
+    )
+    session.execute(delete_stmt)

src/aleph/db/accessors/files.py

@@ -8,6 +8,7 @@
 from aleph.types.db_session import DbSession
 from aleph.types.files import FileTag, FileType
 from aleph.types.sort_order import SortOrder
+from aleph.utils import make_file_tag
 
 from ..models.files import (
     ContentFilePinDb,
@@ -112,36 +113,90 @@ def insert_grace_period_file_pin(
     file_hash: str,
     created: dt.datetime,
     delete_by: dt.datetime,
+    item_hash: Optional[str] = None,
+    owner: Optional[str] = None,
+    ref: Optional[str] = None,
 ) -> None:
     insert_stmt = insert(GracePeriodFilePinDb).values(
+        item_hash=item_hash,
         file_hash=file_hash,
+        owner=owner,
+        ref=ref,
         created=created,
         type=FilePinType.GRACE_PERIOD,
         delete_by=delete_by,
     )
     session.execute(insert_stmt)
 
 
-def upsert_grace_period_file_pin(
+# TODO: Improve performance
+def update_file_pin_grace_period(
     session: DbSession,
-    file_hash: str,
-    created: dt.datetime,
+    item_hash: str,
     delete_by: Union[dt.datetime, None],
 ) -> None:
-    insert_stmt = insert(GracePeriodFilePinDb).values(
-        file_hash=file_hash,
-        created=created,
-        type=FilePinType.GRACE_PERIOD,
-        delete_by=delete_by,
-    )
-    upsert_stmt = insert_stmt.on_conflict_do_update(
-        constraint="file_pins_item_hash_type_key",
-        set_={
-            "created": created,
-            "delete_by": delete_by,
-        },
+    if delete_by is None:
+        delete_stmt = (
+            delete(GracePeriodFilePinDb)
+            .where(GracePeriodFilePinDb.item_hash == item_hash)
+            .returning(
+                GracePeriodFilePinDb.file_hash,
+                GracePeriodFilePinDb.owner,
+                GracePeriodFilePinDb.ref,
+                GracePeriodFilePinDb.created,
+            )
+        )
+
+        grace_period = session.execute(delete_stmt).first()
+        if grace_period is None:
+            return
+
+        file_hash, owner, ref, created = grace_period
+
+        insert_message_file_pin(
+            session=session,
+            item_hash=item_hash,
+            file_hash=file_hash,
+            owner=owner,
+            ref=ref,
+            created=created,
+        )
+    else:
+        delete_stmt = (
+            delete(MessageFilePinDb)
+            .where(MessageFilePinDb.item_hash == item_hash)
+            .returning(
+                MessageFilePinDb.file_hash,
+                MessageFilePinDb.owner,
+                MessageFilePinDb.ref,
+                MessageFilePinDb.created,
+            )
+        )
+
+        message_pin = session.execute(delete_stmt).first()
+        if message_pin is None:
+            return
+
+        file_hash, owner, ref, created = message_pin
+
+        insert_grace_period_file_pin(
+            session=session,
+            item_hash=item_hash,
+            file_hash=file_hash,
+            owner=owner,
+            ref=ref,
+            created=created,
+            delete_by=delete_by,
+        )
+
+    refresh_file_tag(
+        session=session,
+        tag=make_file_tag(
+            owner=owner,
+            ref=ref,
+            item_hash=item_hash,
+        ),
     )
-    session.execute(upsert_stmt)
 
 
 def delete_grace_period_file_pins(session: DbSession, datetime: dt.datetime) -> None:
@@ -235,6 +290,10 @@ def get_file_tag(session: DbSession, tag: FileTag) -> Optional[FileTagDb]:
     return session.execute(select_stmt).scalar()
 
 
+def file_pin_exists(session: DbSession, item_hash: str) -> bool:
+    return FilePinDb.exists(session=session, where=FilePinDb.item_hash == item_hash)
+
+
 def file_tag_exists(session: DbSession, tag: FileTag) -> bool:
     return FileTagDb.exists(session=session, where=FileTagDb.tag == tag)
 

src/aleph/db/accessors/messages.py

@@ -9,6 +9,7 @@
 from sqlalchemy.sql import Insert, Select
 from sqlalchemy.sql.elements import literal
 
+from aleph.db.accessors.cost import delete_costs_for_message
 from aleph.toolkit.timestamp import coerce_to_datetime, utc_now
 from aleph.types.channel import Channel
 from aleph.types.db_session import DbSession
@@ -49,6 +50,13 @@ def message_exists(session: DbSession, item_hash: str) -> bool:
     )
 
 
+def get_one_message_by_item_hash(
+    session: DbSession, item_hash: str
+) -> Optional[RejectedMessageDb]:
+    select_stmt = select(MessageDb).where(MessageDb.item_hash == item_hash)
+    return session.execute(select_stmt).scalar_one_or_none()
+
+
 def make_matching_messages_query(
     hashes: Optional[Sequence[ItemHash]] = None,
     addresses: Optional[Sequence[str]] = None,
@@ -413,6 +421,11 @@ def forget_message(
     )
     session.execute(delete(MessageDb).where(MessageDb.item_hash == item_hash))
 
+    delete_costs_for_message(
+        session=session,
+        item_hash=item_hash,
+    )
+
 
 def append_to_forgotten_by(
     session: DbSession, forgotten_message_hash: str, forget_message_hash: str

src/aleph/db/models/files.py

@@ -77,6 +77,9 @@ class FilePinDb(Base):
     owner = Column(String, nullable=True, index=True)
     item_hash = Column(String, nullable=True)
 
+    # Allow to recover MESSAGE pins refs marked for removing from grace period entries
+    ref = Column(String, nullable=True)
+
     file: StoredFileDb = relationship(StoredFileDb, back_populates="pins")
 
     __mapper_args__: Dict[str, Any] = {
@@ -94,7 +97,7 @@ class TxFilePinDb(FilePinDb):
 
 
 class MessageFilePinDb(FilePinDb):
-    ref = Column(String, nullable=True)
+    # ref = Column(String, nullable=True)
 
     __mapper_args__ = {
         "polymorphic_identity": FilePinType.MESSAGE.value,

src/aleph/handlers/content/forget.py

@@ -179,6 +179,8 @@ async def _forget_item_hash(
 
         if message_status.status == MessageStatus.REJECTED:
             logger.info("Message %s was rejected, nothing to do.", item_hash)
+        if message_status.status == MessageStatus.REMOVED:
+            logger.info("Message %s was removed, nothing to do.", item_hash)
         if message_status.status == MessageStatus.FORGOTTEN:
             logger.info("Message %s is already forgotten, nothing to do.", item_hash)
             append_to_forgotten_by(

src/aleph/handlers/content/store.py

@@ -8,7 +8,7 @@
 import asyncio
 import datetime as dt
 import logging
-from typing import List, Optional, Set
+from typing import List, Set
 
 import aioipfs
 from aleph_message.models import ItemHash, ItemType, StoreContent
@@ -37,7 +37,7 @@
 from aleph.toolkit.costs import are_store_and_program_free
 from aleph.toolkit.timestamp import timestamp_to_datetime, utc_now
 from aleph.types.db_session import DbSession
-from aleph.types.files import FileTag, FileType
+from aleph.types.files import FileType
 from aleph.types.message_status import (
     FileUnavailable,
     InsufficientBalanceException,
@@ -46,7 +46,7 @@
     StoreCannotUpdateStoreWithRef,
     StoreRefNotFound,
 )
-from aleph.utils import item_type_from_hash
+from aleph.utils import item_type_from_hash, make_file_tag
 
 LOGGER = logging.getLogger(__name__)
 
@@ -60,36 +60,6 @@ def _get_store_content(message: MessageDb) -> StoreContent:
     return content
 
 
-def make_file_tag(owner: str, ref: Optional[str], item_hash: str) -> FileTag:
-    """
-    Builds the file tag corresponding to a STORE message.
-
-    The file tag can be set to two different values:
-    * if the `ref` field is not set, the tag will be set to <item_hash>.
-    * if the `ref` field is set, two cases: if `ref` is an item hash, the tag is
-      the value of the ref field. If it is a user-defined value, the tag is
-      <owner>/<ref>.
-
-    :param owner: Owner of the file.
-    :param ref: Value of the `ref` field of the message content.
-    :param item_hash: Item hash of the message.
-    :return: The computed file tag.
-    """
-
-    # When the user does not specify a ref, we use the item hash.
-    if ref is None:
-        return FileTag(item_hash)
-
-    # If ref is an item hash, return it as is
-    try:
-        _item_hash = ItemHash(ref)
-        return FileTag(ref)
-    except ValueError:
-        pass
-
-    return FileTag(f"{owner}/{ref}")
-
-
 class StoreMessageHandler(ContentHandler):
     def __init__(self, storage_service: StorageService, grace_period: int):
         self.storage_service = storage_service