Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Return table description from create table #177

Merged
merged 1 commit into from
Jan 19, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 11 additions & 55 deletions src/aiodynamo/client.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
from __future__ import annotations

import asyncio
import datetime
import json
from dataclasses import dataclass
from typing import (
Expand Down Expand Up @@ -54,8 +53,6 @@
BatchWriteResult,
GlobalSecondaryIndex,
KeySchema,
KeySpec,
KeyType,
LocalSecondaryIndex,
Page,
PayPerRequest,
Expand Down Expand Up @@ -115,7 +112,7 @@ async def create(
gsis: Optional[List[GlobalSecondaryIndex]] = None,
stream: Optional[StreamSpecification] = None,
wait_for_active: Union[bool, RetryConfig] = False,
) -> None:
) -> TableDescription:
return await self.client.create_table(
self.name,
throughput,
Expand Down Expand Up @@ -393,7 +390,7 @@ async def create_table(
gsis: Optional[List[GlobalSecondaryIndex]] = None,
stream: Optional[StreamSpecification] = None,
wait_for_active: Union[bool, RetryConfig] = False,
) -> None:
) -> TableDescription:
attributes = keys.to_attributes()
if lsis is not None:
for index in lsis:
Expand All @@ -418,9 +415,15 @@ async def create_table(
if stream:
payload["StreamSpecification"] = stream.encode()

await self.send_request(action="CreateTable", payload=payload)
response = await self.send_request(action="CreateTable", payload=payload)

description = TableDescription.from_response(response["TableDescription"])

if description.status == TableStatus.active:
return description

async def check() -> bool:
nonlocal description
try:
description = await self.describe_table(name)
return description.status == TableStatus.active
Expand All @@ -429,6 +432,7 @@ async def check() -> bool:

if not await wait(wait_for_active, check):
raise TableDidNotBecomeActive()
return description

async def enable_time_to_live(
self,
Expand Down Expand Up @@ -499,55 +503,7 @@ async def describe_table(self, name: TableName) -> TableDescription:
action="DescribeTable", payload={"TableName": name}
)

description = response["Table"]
attributes: Optional[Dict[str, KeyType]]
if "AttributeDefinitions" in description:
attributes = {
attribute["AttributeName"]: KeyType(attribute["AttributeType"])
for attribute in description["AttributeDefinitions"]
}
else:
attributes = None
creation_time: Optional[datetime.datetime]
if "CreationDateTime" in description:
creation_time = datetime.datetime.fromtimestamp(
description["CreationDateTime"], datetime.timezone.utc
)
else:
creation_time = None
key_schema: Optional[KeySchema]
if attributes and "KeySchema" in description:
key_schema = KeySchema(
*[
KeySpec(
name=key["AttributeName"], type=attributes[key["AttributeName"]]
)
for key in description["KeySchema"]
]
)
else:
key_schema = None
throughput: Optional[ThroughputType]
if (
"BillingModeSummary" in description
and description["BillingModeSummary"]["BillingMode"] == PayPerRequest.MODE
):
throughput = PayPerRequest()
elif "ProvisionedThroughput" in description:
throughput = Throughput(
read=description["ProvisionedThroughput"]["ReadCapacityUnits"],
write=description["ProvisionedThroughput"]["WriteCapacityUnits"],
)
else:
throughput = None
return TableDescription(
attributes=attributes,
created=creation_time,
item_count=description.get("ItemCount", None),
key_schema=key_schema,
throughput=throughput,
status=TableStatus(description["TableStatus"]),
)
return TableDescription.from_response(response["Table"])

async def delete_item(
self,
Expand Down
51 changes: 51 additions & 0 deletions src/aiodynamo/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -206,6 +206,57 @@ class TableDescription:
throughput: Optional[ThroughputType]
status: TableStatus

@classmethod
def from_response(cls, description: Dict[str, Any]) -> "TableDescription":
attributes: Optional[Dict[str, KeyType]]
if "AttributeDefinitions" in description:
attributes = {
attribute["AttributeName"]: KeyType(attribute["AttributeType"])
for attribute in description["AttributeDefinitions"]
}
else:
attributes = None
creation_time: Optional[datetime.datetime]
if "CreationDateTime" in description:
creation_time = datetime.datetime.fromtimestamp(
description["CreationDateTime"], datetime.timezone.utc
)
else:
creation_time = None
key_schema: Optional[KeySchema]
if attributes and "KeySchema" in description:
key_schema = KeySchema(
*[
KeySpec(
name=key["AttributeName"], type=attributes[key["AttributeName"]]
)
for key in description["KeySchema"]
]
)
else:
key_schema = None
throughput: Optional[ThroughputType]
if (
"BillingModeSummary" in description
and description["BillingModeSummary"]["BillingMode"] == PayPerRequest.MODE
):
throughput = PayPerRequest()
elif "ProvisionedThroughput" in description:
throughput = Throughput(
read=description["ProvisionedThroughput"]["ReadCapacityUnits"],
write=description["ProvisionedThroughput"]["WriteCapacityUnits"],
)
else:
throughput = None
return TableDescription(
attributes=attributes,
created=creation_time,
item_count=description.get("ItemCount", None),
key_schema=key_schema,
throughput=throughput,
status=TableStatus(description["TableStatus"]),
)


class Select(Enum):
all_attributes = "ALL_ATTRIBUTES"
Expand Down