more

Author: sanderegg

services/web/server/src/simcore_service_webserver/login/_sql.py

@@ -1,9 +1,21 @@
-def find_one(conn, table, filter_, fields=None):
-    sql, values = find_one_sql(table, filter_, fields)
-    return conn.fetchrow(sql, *values)
+from typing import Any, Awaitable, Iterable
 
+import asyncpg
 
-def find_one_sql(table, filter_, fields=None):
+
+def find_one(
+    conn: asyncpg.pool.PoolConnectionProxy,
+    table: str,
+    filter_: dict[str, Any],
+    fields=None,
+) -> Awaitable[asyncpg.Record | None]:
+    query, values = find_one_sql(table, filter_, fields)
+    return conn.fetchrow(query, *values)
+
+
+def find_one_sql(
+    table: str, filter_: dict[str, Any], fields: Iterable[str] | None = None
+) -> tuple[str, list[Any]]:
     """
     >>> find_one_sql('tbl', {'foo': 10, 'bar': 'baz'})
     ('SELECT * FROM tbl WHERE bar=$1 AND foo=$2', ['baz', 10])
@@ -17,12 +29,19 @@ def find_one_sql(table, filter_, fields=None):
     return sql, values
 
 
-def insert(conn, table, data, returning="id"):
-    sql, values = insert_sql(table, data, returning)
-    return conn.fetchval(sql, *values)
+def insert(
+    conn: asyncpg.pool.PoolConnectionProxy,
+    table: str,
+    data: dict[str, Any],
+    returning: str = "id",
+) -> Awaitable[Any | None]:
+    query, values = insert_sql(table, data, returning)
+    return conn.fetchval(query, *values)
 
 
-def insert_sql(table, data, returning="id"):
+def insert_sql(
+    table: str, data: dict[str, Any], returning: str = "id"
+) -> tuple[str, list[Any]]:
     """
     >>> insert_sql('tbl', {'foo': 'bar', 'id': 1})
     ('INSERT INTO tbl (foo, id) VALUES ($1, $2) RETURNING id', ['bar', 1])
@@ -43,12 +62,16 @@ def insert_sql(table, data, returning="id"):
     return sql, values
 
 
-def update(conn, table, filter_, updates):
-    sql, values = update_sql(table, filter_, updates)
-    return conn.execute(sql, *values)
+def update(
+    conn: asyncpg.pool.PoolConnectionProxy, table: str, filter_, updates
+) -> Awaitable[str]:
+    query, values = update_sql(table, filter_, updates)
+    return conn.execute(query, *values)
 
 
-def update_sql(table, filter_, updates):
+def update_sql(
+    table: str, filter_: dict[str, Any], updates: dict[str, Any]
+) -> tuple[str, list[Any]]:
     """
     >>> update_sql('tbl', {'foo': 'a', 'bar': 1}, {'bar': 2, 'baz': 'b'})
     ('UPDATE tbl SET bar=$1, baz=$2 WHERE bar=$3 AND foo=$4', [2, 'b', 1, 'a'])
@@ -61,12 +84,14 @@ def update_sql(table, filter_, updates):
     return sql, up_vals + where_vals
 
 
-def delete(conn, table, filter_):
+def delete(
+    conn: asyncpg.pool.PoolConnectionProxy, table: str, filter_: dict[str, Any]
+) -> Awaitable[str]:
     sql, values = delete_sql(table, filter_)
     return conn.execute(sql, *values)
 
 
-def delete_sql(table, filter_):
+def delete_sql(table: str, filter_: dict[str, Any]) -> tuple[str, list[Any]]:
     """
     >>> delete_sql('tbl', {'foo': 10, 'bar': 'baz'})
     ('DELETE FROM tbl WHERE bar=$1 AND foo=$2', ['baz', 10])
@@ -77,7 +102,7 @@ def delete_sql(table, filter_):
     return sql, values
 
 
-def _pairs(keys, *, start=1, sep=" AND "):
+def _pairs(keys: Iterable[str], *, start: int = 1, sep: str = " AND ") -> str:
     """
     >>> _pairs(['foo', 'bar', 'baz'], sep=', ')
     'foo=$1, bar=$2, baz=$3'
@@ -87,7 +112,7 @@ def _pairs(keys, *, start=1, sep=" AND "):
     return sep.join(f"{k}=${i}" for i, k in enumerate(keys, start))
 
 
-def _placeholders(variables):
+def _placeholders(variables: Iterable[Any]) -> list[Any]:
     """Returns placeholders by number of variables
 
     >>> _placeholders(['foo', 'bar', 1])
@@ -96,7 +121,7 @@ def _placeholders(variables):
     return [f"${i}" for i, _ in enumerate(variables, 1)]
 
 
-def _split_dict(dic):
+def _split_dict(dic: dict[str, Any]) -> tuple[list[str], list[Any]]:
     """Split dict into sorted keys and values
 
     >>> _split_dict({'b': 2, 'a': 1})

services/web/server/src/simcore_service_webserver/login/storage.py

@@ -38,7 +38,11 @@ class ConfirmationTokenDict(BaseConfirmationTokenDict):
 
 class AsyncpgStorage:
     def __init__(
-        self, pool, *, user_table_name="users", confirmation_table_name="confirmations"
+        self,
+        pool: asyncpg.Pool,
+        *,
+        user_table_name: str = "users",
+        confirmation_table_name: str = "confirmations",
     ):
         self.pool = pool
         self.user_tbl = user_table_name
@@ -48,26 +52,29 @@ def __init__(
     # CRUD user
     #
 
-    async def get_user(self, with_data) -> asyncpg.Record:
+    async def get_user(self, with_data: dict[str, Any]) -> asyncpg.Record:
         async with self.pool.acquire() as conn:
-            return await _sql.find_one(conn, self.user_tbl, with_data)
+            user = await _sql.find_one(conn, self.user_tbl, with_data)
+            assert user  # nosec
+            return user
 
-    async def create_user(self, data: dict[str, Any]) -> asyncpg.Record:
+    async def create_user(self, data: dict[str, Any]) -> dict[str, Any]:
         async with self.pool.acquire() as conn:
             user_id = await _sql.insert(conn, self.user_tbl, data)
             new_user = await _sql.find_one(conn, self.user_tbl, {"id": user_id})
+            assert new_user  # nosec
             data.update(
                 id=new_user["id"],
                 created_at=new_user["created_at"],
                 primary_gid=new_user["primary_gid"],
             )
         return data
 
-    async def update_user(self, user, updates) -> asyncpg.Record:
+    async def update_user(self, user: dict[str, Any], updates: dict[str, Any]) -> None:
         async with self.pool.acquire() as conn:
             await _sql.update(conn, self.user_tbl, {"id": user["id"]}, updates)
 
-    async def delete_user(self, user):
+    async def delete_user(self, user: dict[str, Any]) -> None:
         async with self.pool.acquire() as conn:
             await _sql.delete(conn, self.user_tbl, {"id": user["id"]})
 
@@ -96,12 +103,14 @@ async def create_confirmation(
                 created_at=datetime.utcnow(),
             )
             c = await _sql.insert(
-                conn, self.confirm_tbl, confirmation, returning="code"
+                conn, self.confirm_tbl, dict(confirmation), returning="code"
             )
             assert numeric_code == c  # nosec
             return confirmation
 
-    async def get_confirmation(self, filter_dict) -> ConfirmationTokenDict | None:
+    async def get_confirmation(
+        self, filter_dict: dict[str, Any]
+    ) -> ConfirmationTokenDict | None:
         if "user" in filter_dict:
             filter_dict["user_id"] = filter_dict.pop("user")["id"]
         async with self.pool.acquire() as conn:
@@ -121,14 +130,14 @@ async def delete_confirmation(self, confirmation: ConfirmationTokenDict):
     #
 
     async def delete_confirmation_and_user(
-        self, user, confirmation: ConfirmationTokenDict
+        self, user: dict[str, Any], confirmation: ConfirmationTokenDict
     ):
         async with self.pool.acquire() as conn, conn.transaction():
             await _sql.delete(conn, self.confirm_tbl, {"code": confirmation["code"]})
             await _sql.delete(conn, self.user_tbl, {"id": user["id"]})
 
     async def delete_confirmation_and_update_user(
-        self, user_id, updates, confirmation: ConfirmationTokenDict
+        self, user_id: int, updates: dict[str, Any], confirmation: ConfirmationTokenDict
     ):
         async with self.pool.acquire() as conn, conn.transaction():
             await _sql.delete(conn, self.confirm_tbl, {"code": confirmation["code"]})