add partition key to env var

This commit is contained in:
2025-03-21 12:14:00 -03:00
parent 1e874bf106
commit b22644a165
7 changed files with 170 additions and 27 deletions

View File

@@ -1,2 +1 @@
def hello() -> str:
return "Hello from layercake!"

View File

@@ -1,7 +1,8 @@
import os
from abc import ABC, abstractmethod
from datetime import datetime
from ipaddress import IPv4Address
from typing import Any, Type
from typing import Any, Type, TypedDict
from aws_lambda_powertools import Logger
from boto3.dynamodb.types import TypeDeserializer, TypeSerializer
@@ -10,6 +11,11 @@ from botocore.exceptions import ClientError
from .dateutils import now, timestamp
TZ = os.getenv('TZ', 'UTC')
DELIMITER = os.getenv('DELIMITER', '#')
LIMIT = int(os.getenv('LIMIT', 25))
PARTITION_KEY = os.getenv('PARTITION_KEY', 'pk')
SORT_KEY = os.getenv('SORT_KEY', 'sk')
logger = Logger(__name__)
@@ -40,11 +46,11 @@ def deserialize(value: dict) -> dict:
return {k: deserializer.deserialize(v) for k, v in value.items()}
def Key(
def ComposeKey(
keyparts: str | tuple[str, ...],
*,
prefix: str | None = None,
delimiter: str = '#',
delimiter: str = DELIMITER,
) -> str:
"""Creates a composite key by joining string parts with a specified delimiter.
If a prefix is provided, it is added at the beginning of the key parts.
@@ -67,14 +73,49 @@ def Key(
return delimiter.join(keyparts)
class KeyPair(dict):
class PrimaryKey(ABC, dict):
@abstractmethod
def expr_attr_name(self) -> dict: ...
@abstractmethod
def expr_attr_values(self) -> dict: ...
class PartitionKey(PrimaryKey):
"""Represents a partition key for DynamoDB queries"""
def __init__(self, pk: str) -> None:
super().__init__(**{PARTITION_KEY: pk})
def expr_attr_name(self) -> dict:
return {'#pk': PARTITION_KEY}
def expr_attr_values(self) -> dict:
return {':pk': self[PARTITION_KEY]}
class KeyPair(PrimaryKey):
"""Represents a composite key (partition key and sort key) for DynamoDB queries"""
def __init__(self, pk: str, sk: str) -> None:
super().__init__(id=pk, sk=sk)
super().__init__(**{PARTITION_KEY: pk, SORT_KEY: sk})
def __repr__(self) -> str:
pk, sk = self.values()
return f'KeyPair({pk!r}, {sk!r})'
def expr_attr_name(self) -> dict:
return {
'#pk': PARTITION_KEY,
'#sk': SORT_KEY,
}
def expr_attr_values(self) -> dict:
return {
':pk': self[PARTITION_KEY],
':sk': self[SORT_KEY],
}
class TransactItems:
"""
@@ -440,20 +481,40 @@ class DynamoDBCollection:
"""
Example
-------
Get an item using a composed sort key
**Get an item using a composed sort key**
collect = DynamoDBCollection(...)
collect.get_item(
key=KeyPair(
pk='5OxmMjL-ujoR5IMGegQz',
sk=Key('sergio@somosbeta.com.br', prefix='emails'),
KeyPair(
'b3511b5a-cb32-4833-a373-f8223f2088d4',
ComposeKey('sergio@somosbeta.com.br', prefix='emails'),
),
)
**Get items using a composed partition key**
collect = DynamoDBCollection(...)
collect.get_items(
PartitionKey(
ComposeKey('b3511b5a-cb32-4833-a373-f8223f2088d4', prefix='logs')
),
)
**Get items using a key pair**
collect = DynamoDBCollection(...)
collect.get_items(
KeyPair('b3511b5a-cb32-4833-a373-f8223f2088d4', 'emails),
)
"""
class MissingError(ValueError):
pass
class PaginatedResult(TypedDict):
items: list[dict]
last_key: str | None
def __init__(
self,
persistence_layer: DynamoDBPersistenceLayer,
@@ -533,4 +594,37 @@ class DynamoDBCollection:
expr_attr_values=expr_attr_values,
)
def get_items(self): ...
def get_items(
self,
key: PartitionKey | KeyPair,
*,
expr_attr_name: dict = {},
expr_attr_values: dict = {},
start_key: str | None = None,
filter_expr: str | None = None,
index_forward: bool = False,
limit: int = LIMIT,
) -> PaginatedResult:
key_cond_expr = (
'#pk = :pk AND begins_with(#sk, :sk)'
if isinstance(key, KeyPair)
else '#pk = :pk'
)
expr_attr_name.update(key.expr_attr_name())
expr_attr_values.update(key.expr_attr_values())
response = self.persistence_layer.query(
key_cond_expr=key_cond_expr,
expr_attr_name=expr_attr_name,
expr_attr_values=expr_attr_values,
filter_expr=filter_expr,
index_forward=index_forward,
limit=limit,
# start_key=start_key if start_key else {},
)
return {
'items': response['items'],
'last_key': response['last_key'] if 'last_key' in response else None,
}

View File

@@ -23,7 +23,12 @@ dependencies = [
]
[dependency-groups]
dev = ["pytest>=8.3.5", "pytest-cov>=6.0.0", "ruff>=0.11.1"]
dev = [
"pytest>=8.3.5",
"pytest-cov>=6.0.0",
"pytest-env>=1.1.5",
"ruff>=0.11.1",
]
[tool.pytest.ini_options]
addopts = "--cov --cov-report html -v"

3
layercake/pytest.ini Normal file
View File

@@ -0,0 +1,3 @@
[pytest]
env =
PARTITION_KEY=id

View File

@@ -1,2 +1,4 @@
{"update_date": {"S": "2024-02-08T16:42:33.776409-03:00"}, "create_date": {"S": "2019-03-25T00:00:00-03:00"}, "email_verified": {"BOOL": true}, "cognito:sub": {"S": "58efed8d-d276-41a8-8502-4ab8b5a6415e"}, "cpf": {"S": "07879819908"}, "sk": {"S": "0"}, "email": {"S": "sergio@somosbeta.com.br"}, "id": {"S": "5OxmMjL-ujoR5IMGegQz"}, "name": {"S": "S\u00e9rgio Rafael de Siqueira"}, "last_login": {"S": "2024-02-08T20:53:45.818126-03:00"}, "tenant:org_id": {"L": [{"S": "cJtK9SsnJhKPyxESe7g3DG"}, {"S": "edp8njvgQuzNkLx2ySNfAD"}, {"S": "8TVSi5oACLxTiT8ycKPmaQ"}]}}
{"email_verified": {"BOOL": true}, "update_date": {"S": "2024-02-08T16:42:33.776409-03:00"}, "create_date": {"S": "2019-03-25T00:00:00-03:00"}, "sk": {"S": "emails#sergio@somosbeta.com.br"}, "id": {"S": "5OxmMjL-ujoR5IMGegQz"}, "email_primary": {"BOOL": true}, "mx_record_exists": {"BOOL": true}, "update_date": {"S": "2023-11-09T12:13:04.308986-03:00"}}
{"id": {"S": "5OxmMjL-ujoR5IMGegQz"}, "sk": {"S": "0"}, "update_date": {"S": "2024-02-08T16:42:33.776409-03:00"}, "create_date": {"S": "2019-03-25T00:00:00-03:00"}, "email_verified": {"BOOL": true}, "cognito:sub": {"S": "58efed8d-d276-41a8-8502-4ab8b5a6415e"}, "cpf": {"S": "07879819908"}, "email": {"S": "sergio@somosbeta.com.br"}, "name": {"S": "S\u00e9rgio Rafael de Siqueira"}, "last_login": {"S": "2024-02-08T20:53:45.818126-03:00"}, "tenant:org_id": {"L": [{"S": "cJtK9SsnJhKPyxESe7g3DG"}, {"S": "edp8njvgQuzNkLx2ySNfAD"}, {"S": "8TVSi5oACLxTiT8ycKPmaQ"}]}}
{"id": {"S": "5OxmMjL-ujoR5IMGegQz"}, "sk": {"S": "emails#sergio@somosbeta.com.br"}, "email_verified": {"BOOL": true}, "update_date": {"S": "2024-02-08T16:42:33.776409-03:00"}, "create_date": {"S": "2019-03-25T00:00:00-03:00"}, "email_primary": {"BOOL": true}, "mx_record_exists": {"BOOL": true}, "update_date": {"S": "2023-11-09T12:13:04.308986-03:00"}}
{"id": {"S": "logs#5OxmMjL-ujoR5IMGegQz"}, "sk": {"S": "2024-02-08T16:42:33.776409-03:00"}, "action": {"S": "OPEN_EMAIL"}}
{"id": {"S": "logs#5OxmMjL-ujoR5IMGegQz"}, "sk": {"S": "2019-03-25T00:00:00-03:00"}, "action": {"S": "CLICK_EMAIL"}}

View File

@@ -6,10 +6,11 @@ from botocore.exceptions import ClientError
from layercake.dateutils import ttl
from layercake.dynamodb import (
ComposeKey,
DynamoDBCollection,
DynamoDBPersistenceLayer,
Key,
KeyPair,
PartitionKey,
TransactItems,
serialize,
)
@@ -31,12 +32,22 @@ def test_serialize():
}
def test_key():
assert Key(('122', 'abc'), prefix='schedules') == 'schedules#122#abc'
def test_composekey():
assert ComposeKey(('122', 'abc'), prefix='schedules') == 'schedules#122#abc'
assert ComposeKey(('122', 'abc')) == '122#abc'
assert ComposeKey('122') == '122'
def test_partitionkey():
assert PartitionKey('123') == {'id': '123'}
assert PartitionKey('123').expr_attr_name() == {'#pk': 'id'}
assert PartitionKey('123').expr_attr_values() == {':pk': '123'}
def test_keypair():
assert KeyPair('123', 'abc') == {'id': '123', 'sk': 'abc'}
assert KeyPair('123', 'abc').expr_attr_name() == {'#pk': 'id', '#sk': 'sk'}
assert KeyPair('123', 'abc').expr_attr_values() == {':pk': '123', ':sk': 'abc'}
def test_transact_write_items(
@@ -65,7 +76,7 @@ def test_collection_get_item(
):
collect = DynamoDBCollection(dynamodb_persistence_layer)
data_notfound = collect.get_item(
key=KeyPair(
KeyPair(
pk='5OxmMjL-ujoR5IMGegQz',
sk='tenant',
),
@@ -76,9 +87,9 @@ def test_collection_get_item(
# This item was added from seeds
data = collect.get_item(
key=KeyPair(
KeyPair(
pk='5OxmMjL-ujoR5IMGegQz',
sk=Key('sergio@somosbeta.com.br', prefix='emails'),
sk=ComposeKey('sergio@somosbeta.com.br', prefix='emails'),
),
default={},
)
@@ -102,18 +113,18 @@ def test_collection_put_item(
collect = DynamoDBCollection(dynamodb_persistence_layer)
assert collect.put_item(
key=KeyPair(
KeyPair(
'5OxmMjL-ujoR5IMGegQz',
Key('6d1044d5-18c5-437c-9219-fc2ace7e5ebc', prefix='orgs'),
ComposeKey('6d1044d5-18c5-437c-9219-fc2ace7e5ebc', prefix='orgs'),
),
name='Beta Educação',
ttl=ttl(days=3),
)
data = collect.get_item(
key=KeyPair(
KeyPair(
pk='5OxmMjL-ujoR5IMGegQz',
sk=Key('6d1044d5-18c5-437c-9219-fc2ace7e5ebc', prefix='orgs'),
sk=ComposeKey('6d1044d5-18c5-437c-9219-fc2ace7e5ebc', prefix='orgs'),
),
)
@@ -131,8 +142,23 @@ def test_collection_delete_item(
# This item was added from seeds
assert collect.delete_item(
key=KeyPair(
KeyPair(
'5OxmMjL-ujoR5IMGegQz',
Key('sergio@somsbeta.com.br', prefix='emails'),
ComposeKey('sergio@somsbeta.com.br', prefix='emails'),
)
)
def test_collection_get_items(
dynamodb_seeds,
dynamodb_persistence_layer: DynamoDBPersistenceLayer,
):
collect = DynamoDBCollection(dynamodb_persistence_layer)
# This item was added from seeds
data = collect.get_items(
PartitionKey(
ComposeKey('5OxmMjL-ujoR5IMGegQz', prefix='logs'),
),
)
assert len(data['items']) == 2

14
layercake/uv.lock generated
View File

@@ -403,6 +403,7 @@ dependencies = [
dev = [
{ name = "pytest" },
{ name = "pytest-cov" },
{ name = "pytest-env" },
{ name = "ruff" },
]
@@ -426,6 +427,7 @@ requires-dist = [
dev = [
{ name = "pytest", specifier = ">=8.3.5" },
{ name = "pytest-cov", specifier = ">=6.0.0" },
{ name = "pytest-env", specifier = ">=1.1.5" },
{ name = "ruff", specifier = ">=0.11.1" },
]
@@ -620,6 +622,18 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/36/3b/48e79f2cd6a61dbbd4807b4ed46cb564b4fd50a76166b1c4ea5c1d9e2371/pytest_cov-6.0.0-py3-none-any.whl", hash = "sha256:eee6f1b9e61008bd34975a4d5bab25801eb31898b032dd55addc93e96fcaaa35", size = 22949 },
]
[[package]]
name = "pytest-env"
version = "1.1.5"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "pytest" },
]
sdist = { url = "https://files.pythonhosted.org/packages/1f/31/27f28431a16b83cab7a636dce59cf397517807d247caa38ee67d65e71ef8/pytest_env-1.1.5.tar.gz", hash = "sha256:91209840aa0e43385073ac464a554ad2947cc2fd663a9debf88d03b01e0cc1cf", size = 8911 }
wheels = [
{ url = "https://files.pythonhosted.org/packages/de/b8/87cfb16045c9d4092cfcf526135d73b88101aac83bc1adcf82dfb5fd3833/pytest_env-1.1.5-py3-none-any.whl", hash = "sha256:ce90cf8772878515c24b31cd97c7fa1f4481cd68d588419fd45f10ecaee6bc30", size = 6141 },
]
[[package]]
name = "python-dateutil"
version = "2.9.0.post0"