DynamoDB

Lucas Belo · Lucas Belo · commit 096e99930a80 · 2024-10-07T10:35:54.000-03:00
diff --git a/services/tasks_api/models.py b/services/tasks_api/models.py
@@ -17,4 +17,4 @@ class Task:
 
     @classmethod
     def create(cls, id_, title, owner):
-        return cls(id_, title, TaskStatus.OPEN, owner)
+        return cls(id_, title, TaskStatus.OPEN, owner)
diff --git a/services/tasks_api/poetry.lock b/services/tasks_api/poetry.lock
diff --git a/services/tasks_api/pyproject.toml b/services/tasks_api/pyproject.toml
@@ -11,6 +11,7 @@ fastapi = "^0.115.0"
 uvicorn = "^0.31.0"
 httpx = "^0.27.2"
 mangum = "^0.19.0"
+boto3 = "1.21.45"
 
 
 [tool.poetry.group.dev.dependencies]
@@ -20,6 +21,7 @@ black = "^24.8.0"
 isort = "^5.13.2"
 flake8 = "^7.1.1"
 bandit = "^1.7.10"
+moto = "3.1.5"
 
 [build-system]
 requires = ["poetry-core"]
diff --git a/services/tasks_api/resources/dynamodb.yml b/services/tasks_api/resources/dynamodb.yml
@@ -0,0 +1,29 @@
+Resources:
+  TasksAPITable:
+    Type: AWS::DynamoDB::Table
+    Properties:
+      TableName: ${self:custom.tableName}
+      BillingMode: PAY_PER_REQUEST
+      AttributeDefinitions:
+        - AttributeName: PK
+          AttributeType: S
+        - AttributeName: SK
+          AttributeType: S
+        - AttributeName: GS1PK
+          AttributeType: S
+        - AttributeName: GS1SK
+          AttributeType: S
+      KeySchema:
+        - AttributeName: PK
+          KeyType: HASH
+        - AttributeName: SK
+          KeyType: RANGE
+      GlobalSecondaryIndexes:
+        - IndexName: GS1
+          KeySchema:
+            - AttributeName: GS1PK
+              KeyType: HASH
+            - AttributeName: GS1SK
+              KeyType: RANGE
+          Projection:
+            ProjectionType: ALL
diff --git a/services/tasks_api/serverless.yml b/services/tasks_api/serverless.yml
@@ -12,6 +12,23 @@ provider:
   logRetentionInDays: 30
   environment:
     APP_ENVIRONMENT: ${self:provider.stage}
+  iam:
+    role:
+      statements:
+        - Effect: Allow
+          Action:
+            - dynamodb:DescribeTable
+            - dynamodb:Query
+            - dynamodb:Scan
+            - dynamodb:GetItem
+            - dynamodb:PutItem
+            - dynamodb:UpdateItem
+            - dynamodb:DeleteItem
+          # Allow only access to the API's table and its indexes
+          Resource:
+            - "Fn::GetAtt": [ TasksAPITable, Arn ]
+            - "Fn::Join": ['/', ["Fn::GetAtt": [ TasksAPITable, Arn ], 'index', '*']]
+
 
 functions:
   API:
@@ -33,6 +50,11 @@ custom:
     noDeploy:
       - boto3  # already on Lambda
       - botocore  # already on Lambda
+  stage: ${opt:stage, self:provider.stage}
+  tableName: ${self:custom.stage}-tasks-api
 
 plugins:
-  - serverless-python-requirements
+  - serverless-python-requirements
+
+resources:
+  - ${file(resources/dynamodb.yml)}
diff --git a/services/tasks_api/store.py b/services/tasks_api/store.py
@@ -0,0 +1,79 @@
+import datetime
+from uuid import UUID
+
+import boto3
+from boto3.dynamodb.conditions import Key
+
+from models import Task, TaskStatus
+
+
+class TaskStore:
+    def __init__(self, table_name):
+        self.table_name = table_name
+
+    def add(self, task):
+        dynamodb = boto3.resource("dynamodb")
+        table = dynamodb.Table(self.table_name)
+        table.put_item(
+            Item={
+                "PK": f"#{task.owner}",
+                "SK": f"#{task.id}",
+                "GS1PK": f"#{task.owner}#{task.status.value}",
+                "GS1SK": f"#{datetime.datetime.utcnow().isoformat()}",
+                "id": str(task.id),
+                "title": task.title,
+                "status": task.status.value,
+                "owner": task.owner,
+            }
+        )
+
+    def get_by_id(self, task_id, owner):
+        dynamodb = boto3.resource("dynamodb")
+        table = dynamodb.Table(self.table_name)
+        record = table.get_item(
+            Key={
+                "PK": f"#{owner}",
+                "SK": f"#{task_id}",
+            },
+        )
+        return Task(
+            id=UUID(record["Item"]["id"]),
+            title=record["Item"]["title"],
+            owner=record["Item"]["owner"],
+            status=TaskStatus[record["Item"]["status"]],
+        )
+
+    def list_open(self, owner):
+        return self._list_by_status(owner, TaskStatus.OPEN)
+
+    def list_closed(self, owner):
+        return self._list_by_status(owner, TaskStatus.CLOSED)
+
+    def _list_by_status(self, owner, status):
+        dynamodb = boto3.resource("dynamodb")
+        table = dynamodb.Table(self.table_name)
+        last_key = None
+        query_kwargs = {
+            "IndexName": "GS1",
+            "KeyConditionExpression": Key("GS1PK").eq(f"#{owner}#{status.value}"),
+        }
+        tasks = []
+        while True:
+            if last_key is not None:
+                query_kwargs["ExclusiveStartKey"] = last_key
+            response = table.query(**query_kwargs)
+            tasks.extend(
+                [
+                    Task(
+                        id=UUID(record["id"]),
+                        title=record["title"],
+                        owner=record["owner"],
+                        status=TaskStatus[record["status"]],
+                    )
+                    for record in response["Items"]
+                ]
+            )
+            last_key = response.get("LastEvaluatedKey")
+            if last_key is None:
+                break
+        return tasks
diff --git a/services/tasks_api/tests.py b/services/tasks_api/tests.py
@@ -1,8 +1,14 @@
+import uuid
+
+import boto3
 import pytest
 from fastapi import status
+from moto import mock_dynamodb
 from starlette.testclient import TestClient
 
 from main import app
+from models import Task, TaskStatus
+from store import TaskStore
 
 
 @pytest.fixture
@@ -19,3 +25,72 @@ def test_health_check(client):
     response = client.get("/api/health-check/")
     assert response.status_code == status.HTTP_200_OK
     assert response.json() == {"message": "OK"}
+
+
+@pytest.fixture
+def dynamodb_table():
+    with mock_dynamodb():
+        client = boto3.client("dynamodb")
+        table_name = "test-table"
+        client.create_table(
+            AttributeDefinitions=[
+                {"AttributeName": "PK", "AttributeType": "S"},
+                {"AttributeName": "SK", "AttributeType": "S"},
+                {"AttributeName": "GS1PK", "AttributeType": "S"},
+                {"AttributeName": "GS1SK", "AttributeType": "S"},
+            ],
+            TableName=table_name,
+            KeySchema=[
+                {"AttributeName": "PK", "KeyType": "HASH"},
+                {"AttributeName": "SK", "KeyType": "RANGE"},
+            ],
+            BillingMode="PAY_PER_REQUEST",
+            GlobalSecondaryIndexes=[
+                {
+                    "IndexName": "GS1",
+                    "KeySchema": [
+                        {"AttributeName": "GS1PK", "KeyType": "HASH"},
+                        {"AttributeName": "GS1SK", "KeyType": "RANGE"},
+                    ],
+                    "Projection": {
+                        "ProjectionType": "ALL",
+                    },
+                },
+            ],
+        )
+        yield table_name
+
+
+def test_added_task_retrieved_by_id(dynamodb_table):
+    repository = TaskStore(table_name=dynamodb_table)
+    task = Task.create(uuid.uuid4(), "Clean your office", "john@doe.com")
+
+    repository.add(task)
+
+    assert repository.get_by_id(task_id=task.id, owner=task.owner) == task
+
+
+def test_open_tasks_listed(dynamodb_table):
+    repository = TaskStore(table_name=dynamodb_table)
+    open_task = Task.create(uuid.uuid4(), "Clean your office", "john@doe.com")
+    closed_task = Task(
+        uuid.uuid4(), "Clean your office", TaskStatus.CLOSED, "john@doe.com"
+    )
+
+    repository.add(open_task)
+    repository.add(closed_task)
+
+    assert repository.list_open(owner=open_task.owner) == [open_task]
+
+
+def test_closed_tasks_listed(dynamodb_table):
+    repository = TaskStore(table_name=dynamodb_table)
+    open_task = Task.create(uuid.uuid4(), "Clean your office", "john@doe.com")
+    closed_task = Task(
+        uuid.uuid4(), "Clean your office", TaskStatus.CLOSED, "john@doe.com"
+    )
+
+    repository.add(open_task)
+    repository.add(closed_task)
+
+    assert repository.list_closed(owner=open_task.owner) == [closed_task]
