From 164fc414ac1a23ebb753daa9f5d21432c4486c1c Mon Sep 17 00:00:00 2001
From: Jonas Obrist <jonas.obrist@hennge.com>
Date: Thu, 18 Jan 2024 17:10:51 +0900
Subject: [PATCH] Return table description from create table

---
 src/aiodynamo/client.py | 66 +++++++----------------------------------
 src/aiodynamo/models.py | 51 +++++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+), 55 deletions(-)

diff --git a/src/aiodynamo/client.py b/src/aiodynamo/client.py
index c0a6b71..cc66865 100644
--- a/src/aiodynamo/client.py
+++ b/src/aiodynamo/client.py
@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import asyncio
-import datetime
 import json
 from dataclasses import dataclass
 from typing import (
@@ -54,8 +53,6 @@
     BatchWriteResult,
     GlobalSecondaryIndex,
     KeySchema,
-    KeySpec,
-    KeyType,
     LocalSecondaryIndex,
     Page,
     PayPerRequest,
@@ -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,
@@ -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:
@@ -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
@@ -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,
@@ -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,
diff --git a/src/aiodynamo/models.py b/src/aiodynamo/models.py
index 4fee350..388829c 100644
--- a/src/aiodynamo/models.py
+++ b/src/aiodynamo/models.py
@@ -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"