Refactor: use async sessions factory instead of sync

Author: 1yam

src/aleph/api_entrypoint.py

@@ -7,7 +7,7 @@
 
 import aleph.config
 from aleph.chains.signature_verifier import SignatureVerifier
-from aleph.db.connection import make_engine, make_session_factory
+from aleph.db.connection import make_async_engine, make_async_session_factory
 from aleph.services.cache.node_cache import NodeCache
 from aleph.services.ipfs import IpfsService
 from aleph.services.p2p import init_p2p_client
@@ -34,12 +34,12 @@ async def configure_aiohttp_app(
     with sentry_sdk.start_transaction(name="init-api-server"):
         p2p_client = await init_p2p_client(config, service_name="api-server-aiohttp")
 
-        engine = make_engine(
+        engine = make_async_engine(
             config,
             echo=config.logging.level.value == logging.DEBUG,
             application_name="aleph-api",
         )
-        session_factory = make_session_factory(engine)
+        session_factory = make_async_session_factory(engine)
 
         node_cache = NodeCache(
             redis_host=config.redis.host.value, redis_port=config.redis.port.value

src/aleph/chains/nuls2.py

@@ -29,7 +29,7 @@
 from aleph.schemas.chains.tx_context import TxContext
 from aleph.schemas.pending_messages import BasePendingMessage
 from aleph.toolkit.timestamp import utc_now
-from aleph.types.db_session import DbSessionFactory
+from aleph.types.db_session import AsyncDbSessionFactory
 from aleph.utils import run_in_executor
 
 from ..db.models import ChainTxDb
@@ -80,7 +80,7 @@ async def verify_signature(self, message: BasePendingMessage) -> bool:
 class Nuls2Connector(ChainWriter):
     def __init__(
         self,
-        session_factory: DbSessionFactory,
+        session_factory: AsyncDbSessionFactory,
         pending_tx_publisher: PendingTxPublisher,
         chain_data_service: ChainDataService,
     ):
@@ -90,8 +90,8 @@ def __init__(
 
     async def get_last_height(self, sync_type: ChainEventType) -> int:
         """Returns the last height for which we already have the nuls data."""
-        with self.session_factory() as session:
-            last_height = get_last_height(
+        async with self.session_factory() as session:
+            last_height = await get_last_height(
                 session=session, chain=Chain.NULS2, sync_type=sync_type
             )
 
@@ -133,15 +133,15 @@ async def _request_transactions(
                 LOGGER.info("Incoming logic data is not JSON, ignoring. %r" % ldata)
 
         if last_height:
-            with self.session_factory() as session:
-                upsert_chain_sync_status(
+            async with self.session_factory() as session:
+                await upsert_chain_sync_status(
                     session=session,
                     chain=Chain.NULS2,
                     sync_type=ChainEventType.SYNC,
                     height=last_height,
                     update_datetime=utc_now(),
                 )
-                session.commit()
+                await session.commit()
 
     async def fetcher(self, config: Config):
         last_stored_height = await self.get_last_height(sync_type=ChainEventType.SYNC)
@@ -158,11 +158,11 @@ async def fetcher(self, config: Config):
                     tx = ChainTxDb.from_sync_tx_context(
                         tx_context=context, tx_data=jdata
                     )
-                    with self.session_factory() as db_session:
+                    async with self.session_factory() as db_session:
                         await self.pending_tx_publisher.add_and_publish_pending_tx(
                             session=db_session, tx=tx
                         )
-                        db_session.commit()
+                        await db_session.commit()
 
                 await asyncio.sleep(10)
 
@@ -182,9 +182,9 @@ async def packer(self, config: Config):
         nonce = await get_nonce(server, address, chain_id)
 
         while True:
-            with self.session_factory() as session:
-                if (count_pending_txs(session=session, chain=Chain.NULS2)) or (
-                    count_pending_messages(session=session, chain=Chain.NULS2)
+            async with self.session_factory() as session:
+                if (await count_pending_txs(session=session, chain=Chain.NULS2)) or (
+                    await count_pending_messages(session=session, chain=Chain.NULS2)
                 ):
                     await asyncio.sleep(30)
                     continue
@@ -195,7 +195,7 @@ async def packer(self, config: Config):
                     i = 0
 
                 messages = list(
-                    get_unconfirmed_messages(
+                    await get_unconfirmed_messages(
                         session=session, limit=10000, chain=Chain.ETH
                     )
                 )
@@ -208,7 +208,7 @@ async def packer(self, config: Config):
                     )
                 )
                 # Required to apply update to the files table in get_chaindata
-                session.commit()
+                await session.commit()
 
                 content = sync_event_payload.json()
                 tx = await prepare_transfer_tx(

src/aleph/commands.py

@@ -26,7 +26,11 @@
 from aleph.chains.chain_data_service import ChainDataService, PendingTxPublisher
 from aleph.chains.connector import ChainConnector
 from aleph.cli.args import parse_args
-from aleph.db.connection import make_db_url, make_engine, make_session_factory
+from aleph.db.connection import (
+    make_async_engine,
+    make_async_session_factory,
+    make_db_url,
+)
 from aleph.exceptions import InvalidConfigException, KeyNotFoundException
 from aleph.jobs import start_jobs
 from aleph.jobs.cron.balance_job import BalanceCronJob
@@ -124,12 +128,12 @@ async def main(args: List[str]) -> None:
         run_db_migrations(config)
     LOGGER.info("Database initialized.")
 
-    engine = make_engine(
+    engine = make_async_engine(
         config,
         echo=args.loglevel == logging.DEBUG,
         application_name="aleph-conn-manager",
     )
-    session_factory = make_session_factory(engine)
+    session_factory = make_async_session_factory(engine)
 
     setup_logging(args.loglevel)
 

src/aleph/handlers/message_handler.py

@@ -8,6 +8,7 @@
 from aleph_message.models import ItemHash, ItemType, MessageType
 from configmanager import Config
 from pydantic import ValidationError
+from sqlalchemy import select
 from sqlalchemy.dialects.postgresql import insert
 
 from aleph.chains.signature_verifier import SignatureVerifier
@@ -39,7 +40,7 @@
 from aleph.schemas.pending_messages import parse_message
 from aleph.storage import StorageService
 from aleph.toolkit.timestamp import timestamp_to_datetime
-from aleph.types.db_session import DbSession, DbSessionFactory
+from aleph.types.db_session import AsyncDbSession, AsyncDbSessionFactory
 from aleph.types.files import FileType
 from aleph.types.message_processing_result import ProcessedMessage, RejectedMessage
 from aleph.types.message_status import (
@@ -122,7 +123,7 @@ async def fetch_pending_message(
 
         return validated_message
 
-    async def fetch_related_content(self, session: DbSession, message: MessageDb):
+    async def fetch_related_content(self, session: AsyncDbSession, message: MessageDb):
         content_handler = self.get_content_handler(message.type)
 
         try:
@@ -135,7 +136,7 @@ async def fetch_related_content(self, session: DbSession, message: MessageDb):
             ) from e
 
     async def load_fetched_content(
-        self, session: DbSession, pending_message: PendingMessageDb
+        self, session: AsyncDbSession, pending_message: PendingMessageDb
     ) -> PendingMessageDb:
         if pending_message.item_type != ItemType.inline:
             pending_message.fetched = False
@@ -161,7 +162,7 @@ class MessagePublisher(BaseMessageHandler):
 
     def __init__(
         self,
-        session_factory: DbSessionFactory,
+        session_factory: AsyncDbSessionFactory,
         storage_service: StorageService,
         config: Config,
         pending_message_exchange: aio_pika.abc.AbstractExchange,
@@ -191,21 +192,22 @@ async def add_pending_message(
         origin: Optional[MessageOrigin] = MessageOrigin.P2P,
     ) -> Optional[PendingMessageDb]:
         # TODO: this implementation is just messy, improve it.
-        with self.session_factory() as session:
+        async with self.session_factory() as session:
             try:
                 # we don't check signatures yet.
                 message = parse_message(message_dict)
             except InvalidMessageException as e:
                 LOGGER.warning(e)
-                reject_new_pending_message(
+                await reject_new_pending_message(
                     session=session,
                     pending_message=message_dict,
                     exception=e,
                     tx_hash=tx_hash,
                 )
-                session.commit()
+                await session.commit()
                 return None
 
+            # TODO: fix this pydanticV2 issue
             pending_message = PendingMessageDb.from_obj(
                 message,
                 reception_time=reception_time,
@@ -220,25 +222,24 @@ async def add_pending_message(
                 )
             except InvalidMessageException as e:
                 LOGGER.warning("Invalid message: %s - %s", message.item_hash, str(e))
-                reject_new_pending_message(
+                await reject_new_pending_message(
                     session=session,
                     pending_message=message_dict,
                     exception=e,
                     tx_hash=tx_hash,
                 )
-                session.commit()
+                await session.commit()
                 return None
 
             # Check if there are an already existing record
-            existing_message = (
-                session.query(PendingMessageDb)
-                .filter_by(
-                    sender=pending_message.sender,
-                    item_hash=pending_message.item_hash,
-                    signature=pending_message.signature,
-                )
-                .one_or_none()
+            stmt = select(PendingMessageDb).filter_by(
+                sender=pending_message.sender,
+                item_hash=pending_message.item_hash,
+                signature=pending_message.signature,
             )
+
+            existing_message = (await session.execute(stmt)).scalars().one_or_none()
+
             if existing_message:
                 return existing_message
 
@@ -255,9 +256,9 @@ async def add_pending_message(
             )
 
             try:
-                session.execute(upsert_message_status_stmt)
-                session.execute(insert_pending_message_stmt)
-                session.commit()
+                await session.execute(upsert_message_status_stmt)
+                await session.execute(insert_pending_message_stmt)
+                await session.commit()
             except sqlalchemy.exc.IntegrityError:
                 # Handle the unique constraint violation.
                 LOGGER.warning("Duplicate pending message detected trying to save it.")
@@ -269,14 +270,14 @@ async def add_pending_message(
                     pending_message.item_hash,
                     str(e),
                 )
-                session.rollback()
-                reject_new_pending_message(
+                await session.rollback()
+                await reject_new_pending_message(
                     session=session,
                     pending_message=message_dict,
                     exception=e,
                     tx_hash=tx_hash,
                 )
-                session.commit()
+                await session.commit()
                 return None
 
             await self._publish_pending_message(pending_message)
@@ -307,53 +308,56 @@ async def verify_signature(self, pending_message: PendingMessageDb):
 
     @staticmethod
     async def confirm_existing_message(
-        session: DbSession,
+        session: AsyncDbSession,
         existing_message: MessageDb,
         pending_message: PendingMessageDb,
     ):
         if pending_message.signature != existing_message.signature:
             raise InvalidSignature(f"Invalid signature for {pending_message.item_hash}")
 
-        delete_pending_message(session=session, pending_message=pending_message)
+        await delete_pending_message(session=session, pending_message=pending_message)
         if tx_hash := pending_message.tx_hash:
-            session.execute(
+            await session.execute(
                 make_confirmation_upsert_query(
                     item_hash=pending_message.item_hash, tx_hash=tx_hash
                 )
             )
 
     @staticmethod
     async def confirm_existing_forgotten_message(
-        session: DbSession,
+        session: AsyncDbSession,
         forgotten_message: ForgottenMessageDb,
         pending_message: PendingMessageDb,
     ):
         if pending_message.signature != forgotten_message.signature:
             raise InvalidSignature(f"Invalid signature for {pending_message.item_hash}")
 
-        delete_pending_message(session=session, pending_message=pending_message)
+        await delete_pending_message(session=session, pending_message=pending_message)
 
     async def insert_message(
-        self, session: DbSession, pending_message: PendingMessageDb, message: MessageDb
+        self,
+        session: AsyncDbSession,
+        pending_message: PendingMessageDb,
+        message: MessageDb,
     ):
-        session.execute(make_message_upsert_query(message))
+        await session.execute(make_message_upsert_query(message))
         if message.item_type != ItemType.inline:
-            upsert_file(
+            await upsert_file(
                 session=session,
                 file_hash=message.item_hash,
                 size=message.size,
                 file_type=FileType.FILE,
             )
-            insert_content_file_pin(
+            await insert_content_file_pin(
                 session=session,
                 file_hash=message.item_hash,
                 owner=message.sender,
                 item_hash=message.item_hash,
                 created=timestamp_to_datetime(message.content["time"]),
             )
 
-        delete_pending_message(session=session, pending_message=pending_message)
-        session.execute(
+        await delete_pending_message(session=session, pending_message=pending_message)
+        await session.execute(
             make_message_status_upsert_query(
                 item_hash=message.item_hash,
                 new_status=MessageStatus.PROCESSED,
@@ -363,18 +367,18 @@ async def insert_message(
         )
 
         if tx_hash := pending_message.tx_hash:
-            session.execute(
+            await session.execute(
                 make_confirmation_upsert_query(
                     item_hash=message.item_hash, tx_hash=tx_hash
                 )
             )
 
     async def insert_costs(
-        self, session: DbSession, costs: List[AccountCostsDb], message: MessageDb
+        self, session: AsyncDbSession, costs: List[AccountCostsDb], message: MessageDb
     ):
         if len(costs) > 0:
             insert_stmt = make_costs_upsert_query(costs)
-            session.execute(insert_stmt)
+            await session.execute(insert_stmt)
 
     async def verify_message(self, pending_message: PendingMessageDb) -> MessageDb:
         await self.verify_signature(pending_message=pending_message)
@@ -385,7 +389,7 @@ async def verify_message(self, pending_message: PendingMessageDb) -> MessageDb:
         return validated_message
 
     async def process(
-        self, session: DbSession, pending_message: PendingMessageDb
+        self, session: AsyncDbSession, pending_message: PendingMessageDb
     ) -> ProcessedMessage | RejectedMessage:
         """
         Process a pending message.
@@ -401,7 +405,7 @@ async def process(
         """
 
         # Note: Check if message already exists (and confirm it)
-        existing_message = get_message_by_item_hash(
+        existing_message = await get_message_by_item_hash(
             session=session, item_hash=ItemHash(pending_message.item_hash)
         )
         if existing_message:
@@ -414,7 +418,7 @@ async def process(
 
         # Note: Check if message is already forgotten (and confirm it)
         # this is to avoid race conditions when a confirmation arrives after the FORGET message has been preocessed
-        forgotten_message = get_forgotten_message(
+        forgotten_message = await get_forgotten_message(
             session=session, item_hash=ItemHash(pending_message.item_hash)
         )
         if forgotten_message: