Optimize collectibles DB query

- Use a proper query instead of returning all elements and do the processing on Python
- Related to safe-global/safe-pm#89

safe_transaction_service/history/models.py

@@ -23,7 +23,7 @@
 from django.contrib.postgres.fields import ArrayField
 from django.contrib.postgres.indexes import GinIndex
 from django.core.exceptions import ValidationError
-from django.db import IntegrityError, models, transaction
+from django.db import IntegrityError, connection, models, transaction
 from django.db.models import Case, Count, Index, JSONField, Max, Q, QuerySet
 from django.db.models.expressions import F, OuterRef, RawSQL, Subquery, Value, When
 from django.db.models.functions import Coalesce
@@ -44,6 +44,7 @@
     Keccak256Field,
     Uint256Field,
 )
+from gnosis.eth.utils import fast_to_checksum_address
 from gnosis.safe import SafeOperation
 from gnosis.safe.safe import SafeInfo
 from gnosis.safe.safe_signature import SafeSignature, SafeSignatureType
@@ -499,41 +500,45 @@ def to_erc721_transfer(self):
 
 
 class ERC721TransferManager(TokenTransferManager):
-    # TODO Optimize this
-    def erc721_owned_by(self, address: str) -> List[Tuple[str, int]]:
+    def erc721_owned_by(
+        self, address: ChecksumAddress
+    ) -> List[Tuple[ChecksumAddress, int]]:
         """
         Returns erc721 owned by address, removing the ones sent
 
         :return: List of tuples(token_address: str, token_id: int)
         """
-        # Get all the token history
-        erc721_events = self.to_or_from(address)
-        # Get tokens received and remove tokens transferred
-        tokens_in: Tuple[str, int] = []
-        tokens_out: Tuple[str, int] = []
-        for erc721_event in erc721_events:
-            token_address = erc721_event.address
-            token_id = erc721_event.token_id
-            if token_id is None:
-                logger.error(
-                    "TokenId for ERC721 info token=%s with owner=%s can never be None",
-                    token_address,
-                    address,
-                )
-                continue
-            if erc721_event.to == erc721_event._from:
-                continue  # Nice try ¯\_(ツ)_/¯
-
-            if erc721_event.to == address:
-                list_to_append = tokens_in
-            else:
-                list_to_append = tokens_out
-            list_to_append.append((token_address, token_id))
 
-        for token_out in tokens_out:  # Remove tokens sent from list
-            if token_out in tokens_in:
-                tokens_in.remove(token_out)
-        return tokens_in
+        with connection.cursor() as cursor:
+            hex_address = HexBytes(address)
+            # Queries all the ERC721 IN and all OUT and only returns the ones currently owned
+            cursor.execute(
+                """
+                SELECT Q1.address, Q1.token_id
+                FROM (SELECT address,
+                             token_id,
+                             Count(*) AS count
+                      FROM   history_erc721transfer
+                      WHERE  "to" = %s AND "to" != "_from"
+                      GROUP  BY address,
+                                token_id) Q1
+                     LEFT JOIN (SELECT address,
+                                       token_id,
+                                       Count(*) AS count
+                                FROM   history_erc721transfer
+                                WHERE  "_from" = %s AND "to" != "_from"
+                                GROUP  BY address,
+                                          token_id) Q2
+                            ON Q1.address = Q2.address
+                               AND Q1.token_id = Q2.token_id
+                WHERE Q1.count > COALESCE(Q2.count, 0)
+                """,
+                [hex_address, hex_address],
+            )
+            return [
+                (fast_to_checksum_address(bytes(address)), int(token_id))
+                for address, token_id in cursor.fetchall()
+            ]
 
 
 class ERC721TransferQuerySet(TokenTransferQuerySet):