From 7674d74666db8c2200fd288dc169d420ef52d1b4 Mon Sep 17 00:00:00 2001
From: Andrew Clemons <andrew.clemons@gmail.com>
Date: Fri, 6 Dec 2024 17:13:44 +0900
Subject: [PATCH] Allow aiodynamo to fetch credentials for snapstart lambdas.

For snapstart, the lambda boots once to do any static init.
Subsequently, this snapshotted version is used for handling subsequent
requests.

Normally for lambdas, the AWS credentials are in the environment. Since
snapstart will reuse a snapshot of the running application, it can't
pass in variables in this way since they need to be refreshed.

Instead, the `AWS_CONTAINER_CREDENTIALS_FULL_URI` env var is set and the
`ContainerMetadataCredentials` in aiodynamo should be used to get
credentials if using `Credentials.auto()` to initialise.

Unfortunately, it looks like AWS formats the expiration timestamps
slightly differently for this API on lambda, including the microseconds
now. The `parse_amazon_timestamp` function fails to parse the value,
resulting in `ContainerMetadataCredentials` being rejected as a
`ChainCredentials` candidate and we end up with no credentials at all.

I can't really find a definitive documentation of what the format should
be so I can point to it, but obviously we know the code as it is works
on ECS/EC2 etc so we must continue to be able to parse those. I've
simply added a fallback to use microseconds if the string has a
full-stop in it.

It seems botocore is using the `dateutil` package to handle their
parsing:

https://github.com/boto/botocore/blob/f49ead849aa5a4ea428d9f378de14db6f4c6d645/botocore/utils.py#L950
---
 src/aiodynamo/utils.py   |  9 ++++++---
 tests/unit/test_utils.py | 24 +++++++++++++++++++++++-
 2 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/src/aiodynamo/utils.py b/src/aiodynamo/utils.py
index 9a70fe9..b02154f 100644
--- a/src/aiodynamo/utils.py
+++ b/src/aiodynamo/utils.py
@@ -168,9 +168,12 @@ def serialize_dict(value: Mapping[str, Any]) -> Dict[str, Dict[str, Any]]:
 
 
 def parse_amazon_timestamp(timestamp: str) -> datetime.datetime:
-    return datetime.datetime.strptime(timestamp, "%Y-%m-%dT%H:%M:%SZ").replace(
-        tzinfo=datetime.timezone.utc
-    )
+    if "." in timestamp:
+        value = datetime.datetime.strptime(timestamp, "%Y-%m-%dT%H:%M:%S.%fZ")
+    else:
+        value = datetime.datetime.strptime(timestamp, "%Y-%m-%dT%H:%M:%SZ")
+
+    return value.replace(tzinfo=datetime.timezone.utc)
 
 
 async def wait(
diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py
index 9734ef9..01e9818 100644
--- a/tests/unit/test_utils.py
+++ b/tests/unit/test_utils.py
@@ -1,4 +1,5 @@
 import base64
+import datetime
 from decimal import Decimal
 from functools import partial
 from typing import Any, Callable, Dict
@@ -10,7 +11,28 @@
 )
 
 from aiodynamo.types import NumericTypeConverter
-from aiodynamo.utils import deserialize, dy2py
+from aiodynamo.utils import deserialize, dy2py, parse_amazon_timestamp
+
+
+@pytest.mark.parametrize(
+    ("amazon_timestamp", "expected"),
+    [
+        (
+            "2020-03-12T15:37:51Z",
+            datetime.datetime(2020, 3, 12, 15, 37, 51, tzinfo=datetime.timezone.utc),
+        ),
+        (
+            "2024-12-06T08:03:52.192266Z",
+            datetime.datetime(
+                2024, 12, 6, 8, 3, 52, 192266, tzinfo=datetime.timezone.utc
+            ),
+        ),
+    ],
+)
+def test_parse_amazon_timestamp(
+    amazon_timestamp: str, expected: datetime.datetime
+) -> None:
+    assert parse_amazon_timestamp(amazon_timestamp) == expected
 
 
 def test_binary_decode() -> None: