From 701b00a6b96f4f1f3936edbc4b756245be3afd5e Mon Sep 17 00:00:00 2001
From: Michael Terry <michael.terry@childrens.harvard.edu>
Date: Wed, 10 Jan 2024 15:02:31 -0500
Subject: [PATCH] Fix duckdb's from_iso8601_timestamp() to handle just-year
 timestamps

i.e. "2024" -> "2024-01-01"
---
 cumulus_library/databases.py |  8 ++++++++
 tests/test_duckdb.py         | 27 ++++++++++++++++++++++++++-
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/cumulus_library/databases.py b/cumulus_library/databases.py
index 75b2c489..db97b7a7 100644
--- a/cumulus_library/databases.py
+++ b/cumulus_library/databases.py
@@ -176,6 +176,14 @@ def _compat_from_iso8601_timestamp(
     ) -> Optional[datetime.datetime]:
         if value is None:
             return None
+        if (
+            len(value) < 10
+        ):  # handle partial dates like 1970 or 1980-12 (which spec allows)
+            pieces = value.split("-")
+            if len(pieces) == 1:
+                return datetime.datetime(int(pieces[0]), 1, 1)
+            else:
+                return datetime.datetime(int(pieces[0]), int(pieces[1]), 1)
         return datetime.datetime.fromisoformat(value)
 
     def cursor(self) -> duckdb.DuckDBPyConnection:
diff --git a/tests/test_duckdb.py b/tests/test_duckdb.py
index 45c3b498..09fe41e4 100644
--- a/tests/test_duckdb.py
+++ b/tests/test_duckdb.py
@@ -4,10 +4,13 @@
 import os
 import tempfile
 
+from datetime import datetime, timedelta, timezone
 from unittest import mock
 from pathlib import Path
 
-from cumulus_library import cli
+import pytest
+
+from cumulus_library import cli, databases
 
 
 @mock.patch.dict(
@@ -48,3 +51,25 @@ def test_duckdb_core_build_and_export():
             ) as f:
                 expected = f.read().strip()
             assert generated == expected, basename
+
+
+@pytest.mark.parametrize(
+    "timestamp,expected",
+    [
+        ("2021", datetime(2021, 1, 1, tzinfo=timezone.utc)),
+        ("2019-10", datetime(2019, 10, 1, tzinfo=timezone.utc)),
+        ("1923-01-23", datetime(1923, 1, 23, tzinfo=timezone.utc)),
+        (
+            "2023-01-16T07:55:25-05:00",
+            datetime(2023, 1, 16, 7, 55, 25, tzinfo=timezone(timedelta(hours=-5))),
+        ),
+    ],
+)
+def test_duckdb_from_iso8601_timestamp(timestamp, expected):
+    db = databases.DuckDatabaseBackend(":memory:")
+    parsed = (
+        db.cursor()
+        .execute(f"select from_iso8601_timestamp('{timestamp}')")
+        .fetchone()[0]
+    )
+    assert parsed, expected