its asyncing time

Author: juliansalvador727

gs/backend/config/config.py

@@ -46,4 +46,4 @@ def getenv(config: str) -> str:
 
 DATABASE_CONNECTION_STRING: Final[
     str
-] = f"postgresql+psycopg2://{GS_DATABASE_USER}:{GS_DATABASE_PASSWORD}@{GS_DATABASE_LOCATION}:{GS_DATABASE_PORT}/{GS_DATABASE_NAME}"
+] = f"postgresql+asyncpg://{GS_DATABASE_USER}:{GS_DATABASE_PASSWORD}@{GS_DATABASE_LOCATION}:{GS_DATABASE_PORT}/{GS_DATABASE_NAME}"

gs/backend/data/data_wrappers/abstract_wrapper.py

@@ -18,53 +18,54 @@ class AbstractWrapper(ABC, Generic[T, PK]):
 
     model: type[T]
 
-    def get_all(self) -> list[T]:
+    async def get_all(self) -> list[T]:
         """
         Get all data wrapper for the unspecified model
 
         :return: a list of all model instances
         """
-        with get_db_session() as session:
-            return list(session.exec(select(self.model)).all())
+        async with get_db_session() as session:
+            result = await session.exec(select(self.model))
+            return list(result.all())
 
-    def get_by_id(self, obj_id: PK) -> T:
+    async def get_by_id(self, obj_id: PK) -> T:
         """
         Retrieve data wrapper for the unspecified model
 
         :param obj_id: PK of the model instance to be retrieved
         :return: the retrieved instance
         """
-        with get_db_session() as session:
-            obj = session.get(self.model, obj_id)
+        async with get_db_session() as session:
+            obj = await session.get(self.model, obj_id)
             if not obj:
                 raise ValueError(f"{self.model.__name__} with ID {obj_id} not found.")
             return obj
 
-    def create(self, data: dict[str, Any]) -> T:
+    async def create(self, data: dict[str, Any]) -> T:
         """
         Post data wrapper for the unspecified model
 
         :param data: the JSON object of the model instance to be created
         :return: the newly created instance
         """
-        with get_db_session() as session:
+        async with get_db_session() as session:
             obj = self.model(**data)
             session.add(obj)
-            session.commit()
-            session.refresh(obj)
+            await session.commit()
+            await session.refresh(obj)
             return obj
 
-    def delete_by_id(self, obj_id: PK) -> T:
+    async def delete_by_id(self, obj_id: PK) -> T:
         """
         Delete data wrapper for the unspecified model
 
         :param obj_id: PK of the model instance to be deleted
         :return: the deleted instance
         """
-        with get_db_session() as session:
-            obj = session.get(self.model, obj_id)
+        async with get_db_session() as session:
+            obj = await session.get(self.model, obj_id)
             if not obj:
                 raise ValueError(f"{self.model.__name__} with ID {obj_id} not found.")
             session.delete(obj)
-            session.commit()
+            await session.commit()
             return obj

gs/backend/data/data_wrappers/wrappers.py

@@ -101,17 +101,17 @@ class CommandsWrapper(AbstractWrapper[Commands, UUID]):
 
     model = Commands
 
-    def retrieve_floating_commands(self) -> list[Commands]:
+    async def retrieve_floating_commands(self) -> list[Commands]:
         """
         Retrieves all commands which do not have a valid entry in
         the packet_commands table.
         A command which is not valid is considered as any command whose ID
         does not match with any command_id in the packet_commands table
         """
-        packet_commands = PacketCommandsWrapper().get_all()
+        packet_commands = await PacketCommandsWrapper().get_all()
         packet_ids = {packet_command.command_id for packet_command in packet_commands}
 
-        commands = self.get_all()
+        commands = await self.get_all()
         floating_commands = [fc for fc in commands if fc.id not in packet_ids]
 
         return floating_commands

gs/backend/data/database/engine.py

@@ -1,22 +1,22 @@
-from sqlalchemy import Engine
-from sqlmodel import Session, create_engine, text
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine
 
 from gs.backend.config.config import DATABASE_CONNECTION_STRING
 from gs.backend.data.tables.aro_user_tables import ARO_USER_SCHEMA_NAME
 from gs.backend.data.tables.main_tables import MAIN_SCHEMA_NAME
 from gs.backend.data.tables.transactional_tables import TRANSACTIONAL_SCHEMA_NAME
 
 
-def get_db_engine() -> Engine:
+def get_db_engine() -> AsyncEngine:
     """
     Creates the database engine
 
     :return: engine
     """
-    return create_engine(DATABASE_CONNECTION_STRING)
+    return create_async_engine(DATABASE_CONNECTION_STRING)
 
 
-def get_db_session() -> Session:
+async def get_db_session() -> AsyncSession:
     """
     Creates the database session.
 
@@ -25,22 +25,22 @@ def get_db_session() -> Session:
     :return: session
     """
     engine = get_db_engine()
-    with Session(engine) as session:
-        return session
+    async with AsyncSession(engine) as session:
+        yield session
 
 
-def _create_schemas(session: Session) -> None:
+async def _create_schemas(session: AsyncSession) -> None:
     """
     Creates the schemas in the database.
 
     :param session: The session for which to create the schemas
     """
-    connection = session.connection()
+    connection = await session.connection()
     schemas = [MAIN_SCHEMA_NAME, TRANSACTIONAL_SCHEMA_NAME, ARO_USER_SCHEMA_NAME]
     for schema in schemas:
         # sqlalchemy doesn't check if the schema exists before attempting to create one
-        connection.execute(text(f"CREATE SCHEMA IF NOT EXISTS {schema}"))
-    connection.commit()
+        await connection.execute(text(f"CREATE SCHEMA IF NOT EXISTS {schema}"))
+    await connection.commit()
 
 
 '''Deprecated method to create tables, now handled by Alembic migrations
@@ -58,11 +58,11 @@ def _create_tables(session: Session) -> None:
 '''
 
 
-def setup_database(session: Session) -> None:
+async def setup_database(session: AsyncSession) -> None:
     """
     Creates the schemas for the session.
     Table creation is now handled by Alembic migrations
 
     :param session: The session for which to create the schemas
     """
-    _create_schemas(session)
+    await _create_schemas(session)

gs/backend/data/resources/utils.py

@@ -1,4 +1,5 @@
-from sqlmodel import Session, select
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlmodel import select
 
 from gs.backend.data.resources.callsigns import callsigns
 from gs.backend.data.resources.main_commands import main_commands
@@ -7,34 +8,34 @@
 from gs.backend.data.tables.main_tables import MainCommand, MainTelemetry
 
 
-def add_main_commands(session: Session) -> None:
+async def add_main_commands(session: AsyncSession) -> None:
     """
     Setup the main commands to the database
     """
     query = select(MainCommand).limit(1)  # Check if the db is empty
-    result = session.exec(query).first()
-    if not result:
+    result = await session.exec(query)
+    if not result.first():
         session.add_all(main_commands())
-        session.commit()
+        await session.commit()
 
 
-def add_callsigns(session: Session) -> None:
+async def add_callsigns(session: AsyncSession) -> None:
     """
     Setup the valid callsigns to the database
     """
     query = select(AROUserCallsigns).limit(1)
-    result = session.exec(query).first()
-    if not result:
+    result = await session.exec(query)
+    if not result.first():
         session.add_all(callsigns())
-        session.commit()
+        await session.commit()
 
 
-def add_telemetry(session: Session) -> None:
+async def add_telemetry(session: AsyncSession) -> None:
     """
     Setup the main telemetry to the database
     """
     query = select(MainTelemetry).limit(1)  # Check if the db is empty
-    result = session.exec(query).first()
-    if not result:
+    result = await session.exec(query)
+    if not result.first():
         session.add_all(main_telemetry())
-        session.commit()
+        await session.commit()

gs/backend/migrate.py

@@ -1,3 +1,4 @@
+import asyncio
 import sys
 
 from gs.backend.data.database.engine import get_db_session
@@ -13,26 +14,36 @@
 individually.
 """
 
-if __name__ == "__main__":
+
+async def main() -> None:
+    """Main async function to run migrations"""
     if len(sys.argv) > 2:
         raise ValueError(f"Invalid input. Expected at most 1 argument, received {len(sys.argv)}")
     elif len(sys.argv[1:]) == 0:
-        print("Migrating callsign data...")
-        add_callsigns(get_db_session())
-        print("Migrating main command data...")
-        add_main_commands(get_db_session())
-        print("Migrating telemetry data...")
-        add_telemetry(get_db_session())
+        async with get_db_session() as session:
+            print("Migrating callsign data...")
+            await add_callsigns(session)
+            print("Migrating main command data...")
+            await add_main_commands(session)
+            print("Migrating telemetry data...")
+            await add_telemetry(session)
     else:
         match sys.argv[1]:
             case "callsigns":
-                print("Migrating callsign data...")
-                add_callsigns(get_db_session())
+                async with get_db_session() as session:
+                    print("Migrating callsign data...")
+                    await add_callsigns(session)
             case "commands":
-                print("Migrating main command data...")
-                add_main_commands(get_db_session())
+                async with get_db_session() as session:
+                    print("Migrating main command data...")
+                    await add_main_commands(session)
             case "telemetries":
-                print("Migrating telemetry data...")
-                add_telemetry(get_db_session())
+                async with get_db_session() as session:
+                    print("Migrating telemetry data...")
+                    await add_telemetry(session)
             case _:
                 raise ValueError("Invalid input. Optional arguments include 'callsigns', 'commands', or 'telemetries'.")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())