Add type hints for squeak entry db queries (#679)

squeaknode/admin/util.py

@@ -2,12 +2,13 @@ import logging
 
 from proto import squeak_admin_pb2
 from squeaknode.core.received_offer_with_peer import ReceivedOfferWithPeer
+from squeaknode.core.squeak_entry_with_profile import SqueakEntryWithProfile
 from squeaknode.core.util import get_hash
 
 logger = logging.getLogger(__name__)
 
 
-def squeak_entry_to_message(squeak_entry_with_profile):
+def squeak_entry_to_message(squeak_entry_with_profile: SqueakEntryWithProfile):
     if squeak_entry_with_profile is None:
         return None
     squeak_entry = squeak_entry_with_profile.squeak_entry
@@ -110,7 +111,7 @@ def sync_result_to_message(sync_result):
     )
 
 
-def squeak_entry_to_detail_message(squeak_entry_with_profile):
+def squeak_entry_to_detail_message(squeak_entry_with_profile: SqueakEntryWithProfile):
     if squeak_entry_with_profile is None:
         return None
     squeak_entry = squeak_entry_with_profile.squeak_entry

squeaknode/bitcoin/util.py

@@ -1,5 +1,5 @@
 from bitcoin.core import CBlockHeader
 
 
-def parse_block_header(header_bytes: bytes):
+def parse_block_header(header_bytes: bytes) -> CBlockHeader:
     return CBlockHeader.deserialize(header_bytes)

squeaknode/core/squeak_core.py

@@ -9,6 +9,7 @@ from squeak.core.elliptic import payment_point_bytes_from_scalar_bytes
 from squeak.core.signing import CSigningKey
 
 from squeaknode.bitcoin.blockchain_client import BlockchainClient
+from squeaknode.bitcoin.util import parse_block_header
 from squeaknode.core.offer import Offer
 from squeaknode.core.received_offer import ReceivedOffer
 from squeaknode.core.received_payment import ReceivedPayment
@@ -69,9 +70,10 @@ class SqueakCore:
             timestamp,
             replyto_hash,
         )
+        block_header = parse_block_header(block_info.block_header)
         return SqueakEntry(
             squeak=squeak,
-            block_header=block_info.block_header,
+            block_header=block_header,
         )
 
     def validate_squeak(self, squeak: CSqueak) -> SqueakEntry:
@@ -91,9 +93,10 @@ class SqueakCore:
             squeak.nBlockHeight)
         if squeak.hashBlock != block_info.block_hash:
             raise Exception("Block hash incorrect.")
+        block_header = parse_block_header(block_info.block_header)
         return SqueakEntry(
             squeak=squeak,
-            block_header=block_info.block_header,
+            block_header=block_header,
         )
 
     def get_best_block_height(self) -> int:

squeaknode/core/squeak_entry.py

@@ -1,8 +1,9 @@
 from typing import NamedTuple
 
+from bitcoin.core import CBlockHeader
 from squeak.core import CSqueak
 
 
 class SqueakEntry(NamedTuple):
     squeak: CSqueak
-    block_header: bytes
+    block_header: CBlockHeader

squeaknode/core/squeak_entry_with_profile.py

@@ -1,4 +1,5 @@
 from typing import NamedTuple
+from typing import Optional
 
 from squeaknode.core.squeak_entry import SqueakEntry
 from squeaknode.core.squeak_profile import SqueakProfile
@@ -6,4 +7,4 @@ from squeaknode.core.squeak_profile import SqueakProfile
 
 class SqueakEntryWithProfile(NamedTuple):
     squeak_entry: SqueakEntry
-    squeak_profile: SqueakProfile
+    squeak_profile: Optional[SqueakProfile]

squeaknode/db/squeak_db.py

@@ -3,8 +3,11 @@ from contextlib import contextmanager
 from datetime import datetime
 from datetime import timedelta
 from datetime import timezone
+from typing import List
+from typing import Optional
 
 import sqlalchemy
+from bitcoin.core import CBlockHeader
 from sqlalchemy import func
 from sqlalchemy import literal
 from sqlalchemy.sql import and_
@@ -74,7 +77,7 @@ class SqueakDb:
     def sent_offers(self):
         return self.models.sent_offers
 
-    def insert_squeak(self, squeak, block_header_bytes):
+    def insert_squeak(self, squeak: CSqueak, block_header: CBlockHeader):
         """ Insert a new squeak.
 
         Return the hash (bytes) of the inserted squeak.
@@ -92,7 +95,7 @@ class SqueakDb:
             n_time=squeak.nTime,
             author_address=str(squeak.GetAddress()),
             secret_key=secret_key_hex,
-            block_header=block_header_bytes,
+            block_header=block_header.serialize(),
         )
         with self.get_connection() as connection:
             try:
@@ -111,7 +114,7 @@ class SqueakDb:
             row = result.fetchone()
             return self._parse_squeak_entry(row)
 
-    def get_squeak_entry_with_profile(self, squeak_hash: bytes):
+    def get_squeak_entry_with_profile(self, squeak_hash: bytes) -> Optional[SqueakEntryWithProfile]:
         """ Get a squeak with the author profile. """
         s = (
             select([self.squeaks, self.profiles])
@@ -126,9 +129,11 @@ class SqueakDb:
         with self.get_connection() as connection:
             result = connection.execute(s)
             row = result.fetchone()
+            if row is None:
+                return None
             return self._parse_squeak_entry_with_profile(row)
 
-    def get_timeline_squeak_entries_with_profile(self):
+    def get_timeline_squeak_entries_with_profile(self) -> List[SqueakEntryWithProfile]:
         """ Get all followed squeaks. """
         s = (
             select([self.squeaks, self.profiles])
@@ -151,7 +156,7 @@ class SqueakDb:
 
     def get_squeak_entries_with_profile_for_address(
         self, address, min_block, max_block
-    ):
+    ) -> List[SqueakEntryWithProfile]:
         """ Get a squeak. """
         s = (
             select([self.squeaks, self.profiles])
@@ -175,7 +180,7 @@ class SqueakDb:
             rows = result.fetchall()
             return [self._parse_squeak_entry_with_profile(row) for row in rows]
 
-    def get_thread_ancestor_squeak_entries_with_profile(self, squeak_hash: bytes):
+    def get_thread_ancestor_squeak_entries_with_profile(self, squeak_hash: bytes) -> List[SqueakEntryWithProfile]:
         """ Get all reply ancestors of squeak hash. """
         ancestors = (
             select(
@@ -242,7 +247,7 @@ class SqueakDb:
         #     rows = curs.fetchall()
         #     return [self._parse_squeak_entry_with_profile(row) for row in rows]
 
-    def get_thread_reply_squeak_entries_with_profile(self, squeak_hash: bytes):
+    def get_thread_reply_squeak_entries_with_profile(self, squeak_hash: bytes) -> List[SqueakEntryWithProfile]:
         """ Get all replies for a squeak hash. """
         s = (
             select([self.squeaks, self.profiles])
@@ -1137,11 +1142,12 @@ class SqueakDb:
             following=row["following"],
         )
 
-    def _parse_squeak_entry_with_profile(self, row):
-        if row is None:
-            return None
+    def _parse_squeak_entry_with_profile(self, row) -> SqueakEntryWithProfile:
         squeak_entry = self._parse_squeak_entry(row)
-        squeak_profile = self._parse_squeak_profile(row)
+        if row["profile_id"] is None:
+            squeak_profile = None
+        else:
+            squeak_profile = self._parse_squeak_profile(row)
         return SqueakEntryWithProfile(
             squeak_entry=squeak_entry,
             squeak_profile=squeak_profile,

tests/core/test_squeak_core.py

@@ -35,7 +35,7 @@ class MockBitcoinClient(BlockchainClient):
     genesis_block_info = BlockInfo(
         block_height=0,
         block_hash=CoreMainParams.GENESIS_BLOCK.GetHash(),
-        block_header=CoreMainParams.GENESIS_BLOCK.serialize(),
+        block_header=CoreMainParams.GENESIS_BLOCK.get_header().serialize(),
     )
 
     def get_best_block_info(self) -> BlockInfo: