update comment

http-api/routes/users/__init__.py

@@ -8,28 +8,33 @@ from aws_lambda_powertools.event_handler.api_gateway import (
     Router,
 )
 from elasticsearch import Elasticsearch
-from layercake.dynamodb import DynamoDBPersistenceLayer
+from layercake.dynamodb import (
+    ComposeKey,
+    DynamoDBCollection,
+    DynamoDBPersistenceLayer,
+    KeyPair,
+    PartitionKey,
+)
 from pydantic import UUID4, BaseModel, StringConstraints
 
 import elastic
-from dynamodb import KeyLoc, get_records
-from http_models import RecordResponse, SearchResponse
 from models import User
 from settings import ELASTIC_CONN, USER_TABLE
 
 router = Router()
 dynamodb_client = boto3.client('dynamodb')
 user_layer = DynamoDBPersistenceLayer(USER_TABLE, dynamodb_client)
+collect = DynamoDBCollection(user_layer)
 elastic_client = Elasticsearch(**ELASTIC_CONN)
 
 
 @router.get('/', compress=True, tags=['User'], summary='Get users')
-def get_users() -> SearchResponse:
+def get_users():
     event = router.current_event
     query = event.get_query_string_value('query', '{}')
     page_size = event.get_query_string_value('page_size', '25')
 
-    return elastic.search(  # type: ignore
+    return elastic.search(
         index=USER_TABLE,
         page_size=int(page_size),
         query=json.loads(query),
@@ -53,38 +58,30 @@ def patch_reset(id: str, payload: ResetPasswordPayload):
 
 
 @router.get('/<id>/emails', compress=True, tags=['User'], summary='Get user emails')
-def get_emails(id: str) -> RecordResponse:
+def get_emails(id: str):
     start_key = router.current_event.get_query_string_value('start_key', None)
 
-    return get_records(  # type: ignore
-        pk=id,
-        prefix='emails',
+    return collect.get_items(
+        key=KeyPair(id, 'emails'),
         start_key=start_key,
-        persistence_layer=user_layer,
     )
 
 
 @router.get('/<id>/logs', compress=True, tags=['User'], summary='Get user logs')
-def get_logs(id: str) -> RecordResponse:
+def get_logs(id: str):
     start_key = router.current_event.get_query_string_value('start_key', None)
 
-    return get_records(  # type: ignore
-        pk=id,
-        prefix='log',
-        delimiter=':',
-        keyloc_prefix=KeyLoc.PK,
+    return collect.get_items(
+        key=PartitionKey(ComposeKey(id, prefix='log', delimiter=':')),
         start_key=start_key,
-        persistence_layer=user_layer,
     )
 
 
 @router.get('/<id>/orgs', compress=True, tags=['User'], summary='Get user orgs')
-def get_orgs(id: str) -> RecordResponse:
+def get_orgs(id: str):
     start_key = router.current_event.get_query_string_value('start_key', None)
 
-    return get_records(  # type: ignore
-        pk=id,
-        prefix='orgs',
+    return collect.get_items(
+        key=KeyPair(id, 'orgs'),
         start_key=start_key,
-        persistence_layer=user_layer,
     )