fix: address DuckDB store critical issues

- Fix import error: use public duckdb module instead of private _duckdb
- Fix naive timestamp handling: use replace(tzinfo=UTC) for naive datetimes
- Add SQL injection protection: validate table names with regex
- Change JSON storage to use dict directly instead of json.dumps()

Fixes all critical issues identified in code review:
- Resolves ModuleNotFoundError that was causing all tests to fail
- Prevents ValueError when DuckDB returns naive timestamps
- Protects against SQL injection via malicious table names
- Improves JSON queryability in DuckDB

Co-authored-by: William Easton <[email protected]>

Author: github-actions[bot]

key-value/key-value-aio/src/key_value/aio/stores/duckdb/store.py

@@ -1,4 +1,5 @@
 import json
+import re
 from datetime import datetime, timezone
 from pathlib import Path
 from typing import Any, cast, overload
@@ -52,9 +53,9 @@ def prepare_dump(self, data: dict[str, Any]) -> dict[str, Any]:
         if collection_name is not None:
             json_document["collection"] = collection_name
 
-        # Store as JSON string for DuckDB's JSON column
-        # DuckDB will parse it and store it as native JSON
-        data["value_dict"] = json.dumps(json_document)
+        # Store the dict directly for DuckDB's JSON column
+        # DuckDB's Python API accepts dict objects and stores them as queryable JSON
+        data["value_dict"] = json_document
 
         return data
 
@@ -103,12 +104,18 @@ def _parse_json_column(self, value_dict: Any) -> dict[str, Any]:
     def _convert_timestamps_to_utc(self, data: dict[str, Any]) -> None:
         """Convert naive timestamps to UTC timezone-aware timestamps."""
         created_at = data.get("created_at")
-        if created_at is not None and isinstance(created_at, datetime) and created_at.tzinfo is None:
-            data["created_at"] = created_at.astimezone(tz=timezone.utc)
+        if created_at is not None and isinstance(created_at, datetime):
+            if created_at.tzinfo is None:
+                data["created_at"] = created_at.replace(tzinfo=timezone.utc)
+            else:
+                data["created_at"] = created_at.astimezone(tz=timezone.utc)
 
         expires_at = data.get("expires_at")
-        if expires_at is not None and isinstance(expires_at, datetime) and expires_at.tzinfo is None:
-            data["expires_at"] = expires_at.astimezone(tz=timezone.utc)
+        if expires_at is not None and isinstance(expires_at, datetime):
+            if expires_at.tzinfo is None:
+                data["expires_at"] = expires_at.replace(tzinfo=timezone.utc)
+            else:
+                data["expires_at"] = expires_at.astimezone(tz=timezone.utc)
 
 
 class DuckDBStore(BaseContextManagerStore, BaseStore):
@@ -214,6 +221,11 @@ def __init__(
 
         self._is_closed = False
         self._adapter = DuckDBSerializationAdapter()
+
+        # Validate table name to prevent SQL injection
+        if not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", table_name):
+            msg = "Table name must start with a letter or underscore and contain only letters, digits, or underscores"
+            raise ValueError(msg)
         self._table_name = table_name
         self._stable_api = False
 
@@ -345,9 +357,15 @@ async def _get_managed_entry(self, *, key: str, collection: str) -> ManagedEntry
         }
 
         if created_at is not None and isinstance(created_at, datetime):
-            document["created_at"] = created_at.astimezone(tz=timezone.utc)
+            if created_at.tzinfo is None:
+                document["created_at"] = created_at.replace(tzinfo=timezone.utc)
+            else:
+                document["created_at"] = created_at.astimezone(tz=timezone.utc)
         if expires_at is not None and isinstance(expires_at, datetime):
-            document["expires_at"] = expires_at.astimezone(tz=timezone.utc)
+            if expires_at.tzinfo is None:
+                document["expires_at"] = expires_at.replace(tzinfo=timezone.utc)
+            else:
+                document["expires_at"] = expires_at.astimezone(tz=timezone.utc)
 
         try:
             return self._adapter.load_dict(data=document)

key-value/key-value-aio/tests/stores/duckdb/test_duckdb.py

@@ -3,8 +3,7 @@
 from tempfile import TemporaryDirectory
 
 import pytest
-from _duckdb import DuckDBPyConnection
-from duckdb import CatalogException
+from duckdb import CatalogException, DuckDBPyConnection
 from inline_snapshot import snapshot
 from typing_extensions import override
 

key-value/key-value-sync/src/key_value/sync/code_gen/stores/duckdb/store.py

@@ -2,6 +2,7 @@
 # from the original file 'store.py'
 # DO NOT CHANGE! Change the original file instead.
 import json
+import re
 from datetime import datetime, timezone
 from pathlib import Path
 from typing import Any, cast, overload
@@ -55,9 +56,9 @@ def prepare_dump(self, data: dict[str, Any]) -> dict[str, Any]:
         if collection_name is not None:
             json_document["collection"] = collection_name
 
-        # Store as JSON string for DuckDB's JSON column
-        # DuckDB will parse it and store it as native JSON
-        data["value_dict"] = json.dumps(json_document)
+        # Store the dict directly for DuckDB's JSON column
+        # DuckDB's Python API accepts dict objects and stores them as queryable JSON
+        data["value_dict"] = json_document
 
         return data
 
@@ -106,12 +107,18 @@ def _parse_json_column(self, value_dict: Any) -> dict[str, Any]:
     def _convert_timestamps_to_utc(self, data: dict[str, Any]) -> None:
         """Convert naive timestamps to UTC timezone-aware timestamps."""
         created_at = data.get("created_at")
-        if created_at is not None and isinstance(created_at, datetime) and (created_at.tzinfo is None):
-            data["created_at"] = created_at.astimezone(tz=timezone.utc)
+        if created_at is not None and isinstance(created_at, datetime):
+            if created_at.tzinfo is None:
+                data["created_at"] = created_at.replace(tzinfo=timezone.utc)
+            else:
+                data["created_at"] = created_at.astimezone(tz=timezone.utc)
 
         expires_at = data.get("expires_at")
-        if expires_at is not None and isinstance(expires_at, datetime) and (expires_at.tzinfo is None):
-            data["expires_at"] = expires_at.astimezone(tz=timezone.utc)
+        if expires_at is not None and isinstance(expires_at, datetime):
+            if expires_at.tzinfo is None:
+                data["expires_at"] = expires_at.replace(tzinfo=timezone.utc)
+            else:
+                data["expires_at"] = expires_at.astimezone(tz=timezone.utc)
 
 
 class DuckDBStore(BaseContextManagerStore, BaseStore):
@@ -216,6 +223,11 @@ def __init__(
 
         self._is_closed = False
         self._adapter = DuckDBSerializationAdapter()
+
+        # Validate table name to prevent SQL injection
+        if not re.fullmatch("[A-Za-z_][A-Za-z0-9_]*", table_name):
+            msg = "Table name must start with a letter or underscore and contain only letters, digits, or underscores"
+            raise ValueError(msg)
         self._table_name = table_name
         self._stable_api = False
 
@@ -315,9 +327,15 @@ def _get_managed_entry(self, *, key: str, collection: str) -> ManagedEntry | Non
         document: dict[str, Any] = {"value_dict": value_dict}
 
         if created_at is not None and isinstance(created_at, datetime):
-            document["created_at"] = created_at.astimezone(tz=timezone.utc)
+            if created_at.tzinfo is None:
+                document["created_at"] = created_at.replace(tzinfo=timezone.utc)
+            else:
+                document["created_at"] = created_at.astimezone(tz=timezone.utc)
         if expires_at is not None and isinstance(expires_at, datetime):
-            document["expires_at"] = expires_at.astimezone(tz=timezone.utc)
+            if expires_at.tzinfo is None:
+                document["expires_at"] = expires_at.replace(tzinfo=timezone.utc)
+            else:
+                document["expires_at"] = expires_at.astimezone(tz=timezone.utc)
 
         try:
             return self._adapter.load_dict(data=document)

key-value/key-value-sync/tests/code_gen/stores/duckdb/test_duckdb.py

@@ -6,8 +6,7 @@
 from tempfile import TemporaryDirectory
 
 import pytest
-from _duckdb import DuckDBPyConnection
-from duckdb import CatalogException
+from duckdb import CatalogException, DuckDBPyConnection
 from inline_snapshot import snapshot
 from typing_extensions import override