cleanup

Author: pcrespov

packages/postgres-database/src/simcore_postgres_database/utils_products.py

@@ -1,20 +1,19 @@
-""" Common functions to access products table
-
-"""
+"""Common functions to access products table"""
 
 import warnings
 
 import sqlalchemy as sa
+from aiopg.sa.connection import SAConnection
+from sqlalchemy.ext.asyncio import AsyncConnection
 
-from ._protocols import AiopgConnection, DBConnection
 from .models.groups import GroupType, groups
 from .models.products import products
 
 # NOTE: outside this module, use instead packages/models-library/src/models_library/users.py
 _GroupID = int
 
 
-async def get_default_product_name(conn: DBConnection) -> str:
+async def get_default_product_name(conn: AsyncConnection | SAConnection) -> str:
     """The first row in the table is considered as the default product
 
     :: raises ValueError if undefined
@@ -31,15 +30,17 @@ async def get_default_product_name(conn: DBConnection) -> str:
 
 
 async def get_product_group_id(
-    connection: DBConnection, product_name: str
+    connection: SAConnection, product_name: str
 ) -> _GroupID | None:
     group_id = await connection.scalar(
         sa.select(products.c.group_id).where(products.c.name == product_name)
     )
     return None if group_id is None else _GroupID(group_id)
 
 
-async def execute_get_or_create_product_group(conn, product_name: str) -> int:
+async def execute_get_or_create_product_group(
+    conn: AsyncConnection, product_name: str
+) -> int:
     #
     # NOTE: Separated so it can be used in asyncpg and aiopg environs while both
     #       coexist
@@ -73,7 +74,7 @@ async def execute_get_or_create_product_group(conn, product_name: str) -> int:
 
 
 async def get_or_create_product_group(
-    connection: AiopgConnection, product_name: str
+    connection: SAConnection, product_name: str
 ) -> _GroupID:
     """
     Returns group_id of a product. Creates it if undefined