feat: add native storage support to DuckDB store with JSON column option

- Add use_json_column parameter (defaults to True) for native JSON vs TEXT storage
- Update schema to use native DuckDB types (JSON, TIMESTAMP) for better queryability
- Store value data separately from metadata to eliminate duplication
- Metadata (created_at, ttl, expires_at) in native columns for efficient SQL queries
- Add comprehensive tests for both JSON and TEXT modes
- Add test_native_sql_queryability() demonstrating direct SQL queries on stored data
- Fix timezone handling to ensure consistency with ManagedEntry expectations
- All 411 tests passing (404 passed, 7 skipped)

This enables users to query the DuckDB database directly with SQL for analytics
and data exploration while maintaining full compatibility with the key-value API.

Co-authored-by: William Easton <strawgate@users.noreply.github.com>

Author: github-actions[bot]

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

@@ -20,6 +20,9 @@ class DuckDBStore(BaseContextManagerStore, BaseStore):
     for analytical workloads while supporting standard SQL operations. This store
     can operate in memory-only mode or persist data to disk.
 
+    The store uses native DuckDB types (JSON, TIMESTAMP) to enable efficient SQL queries
+    on stored data. Users can query the database directly for analytics or data exploration.
+
     Note on connection ownership: When you provide an existing connection, the store
     will take ownership and close it when the store is closed or garbage collected.
     If you need to reuse a connection, create separate DuckDB connections for each store.
@@ -28,6 +31,7 @@ class DuckDBStore(BaseContextManagerStore, BaseStore):
     _connection: duckdb.DuckDBPyConnection
     _is_closed: bool
     _owns_connection: bool
+    _use_json_column: bool
 
     @overload
     def __init__(
@@ -36,6 +40,7 @@ def __init__(
         connection: duckdb.DuckDBPyConnection,
         default_collection: str | None = None,
         seed: SEED_DATA_TYPE | None = None,
+        use_json_column: bool = True,
     ) -> None:
         """Initialize the DuckDB store with an existing connection.
 
@@ -47,6 +52,8 @@ def __init__(
             connection: An existing DuckDB connection to use.
             default_collection: The default collection to use if no collection is provided.
             seed: Optional seed data to pre-populate the store.
+            use_json_column: If True, use native JSON column type; if False, use TEXT.
+                Default is True for better queryability and native type support.
         """
 
     @overload
@@ -56,13 +63,16 @@ def __init__(
         database_path: Path | str | None = None,
         default_collection: str | None = None,
         seed: SEED_DATA_TYPE | None = None,
+        use_json_column: bool = True,
     ) -> None:
         """Initialize the DuckDB store with a database path.
 
         Args:
             database_path: Path to the database file. If None or ':memory:', uses in-memory database.
             default_collection: The default collection to use if no collection is provided.
             seed: Optional seed data to pre-populate the store.
+            use_json_column: If True, use native JSON column type; if False, use TEXT.
+                Default is True for better queryability and native type support.
         """
 
     def __init__(
@@ -72,6 +82,7 @@ def __init__(
         database_path: Path | str | None = None,
         default_collection: str | None = None,
         seed: SEED_DATA_TYPE | None = None,
+        use_json_column: bool = True,
     ) -> None:
         """Initialize the DuckDB store.
 
@@ -80,6 +91,8 @@ def __init__(
             database_path: Path to the database file. If None or ':memory:', uses in-memory database.
             default_collection: The default collection to use if no collection is provided.
             seed: Optional seed data to pre-populate the store.
+            use_json_column: If True, use native JSON column type; if False, use TEXT.
+                Default is True for better queryability and native type support.
         """
         if connection is not None and database_path is not None:
             msg = "Provide only one of connection or database_path"
@@ -101,6 +114,7 @@ def __init__(
             self._owns_connection = True
 
         self._is_closed = False
+        self._use_json_column = use_json_column
         self._stable_api = False
 
         super().__init__(default_collection=default_collection, seed=seed)
@@ -109,25 +123,30 @@ def __init__(
     async def _setup(self) -> None:
         """Initialize the database schema for key-value storage.
 
-        Note: The schema stores created_at, ttl, and expires_at as separate columns
-        in addition to the serialized ManagedEntry in value_json. This duplication
-        is intentional for future features:
-        - The expires_at column with its index enables efficient expiration-based
-          cleanup queries (e.g., DELETE FROM kv_entries WHERE expires_at < now())
-        - The separate columns allow for metadata queries without deserializing JSON
-        - Currently, only value_json is read during _get_managed_entry
-
-        This design trades storage space for query flexibility and future extensibility.
+        The schema uses native DuckDB types for efficient querying:
+        - value: JSON or TEXT column storing the actual value data (not full ManagedEntry)
+        - created_at: TIMESTAMP for native datetime operations
+        - ttl: DOUBLE for time-to-live in seconds
+        - expires_at: TIMESTAMP for native expiration queries
+
+        This design enables:
+        - Direct SQL queries on the database for analytics
+        - Efficient expiration cleanup: DELETE FROM kv_entries WHERE expires_at < now()
+        - Metadata queries without JSON deserialization
+        - No data duplication (metadata in columns, value in JSON/TEXT)
         """
+        # Determine column type based on use_json_column setting
+        value_column_type = "JSON" if self._use_json_column else "TEXT"
+
         # Create the main table for storing key-value entries
-        self._connection.execute("""
+        self._connection.execute(f"""
             CREATE TABLE IF NOT EXISTS kv_entries (
                 collection VARCHAR NOT NULL,
                 key VARCHAR NOT NULL,
-                value_json TEXT NOT NULL,
-                created_at DOUBLE,
+                value {value_column_type} NOT NULL,
+                created_at TIMESTAMP,
                 ttl DOUBLE,
-                expires_at DOUBLE,
+                expires_at TIMESTAMP,
                 PRIMARY KEY (collection, key)
             )
         """)
@@ -138,29 +157,58 @@ async def _setup(self) -> None:
             ON kv_entries(collection)
         """)
 
-        # Create index for expiration-based queries (for future cleanup features)
+        # Create index for expiration-based queries
         self._connection.execute("""
             CREATE INDEX IF NOT EXISTS idx_kv_expires_at
             ON kv_entries(expires_at)
         """)
 
     @override
     async def _get_managed_entry(self, *, key: str, collection: str) -> ManagedEntry | None:
-        """Retrieve a managed entry by key from the specified collection."""
+        """Retrieve a managed entry by key from the specified collection.
+
+        Reconstructs the ManagedEntry from the value column and metadata columns.
+        The value column contains only the value data (not the full ManagedEntry),
+        and metadata (created_at, ttl, expires_at) is stored in separate columns.
+        """
         if self._is_closed:
             msg = "Cannot operate on closed DuckDBStore"
             raise RuntimeError(msg)
 
         result = self._connection.execute(
-            "SELECT value_json FROM kv_entries WHERE collection = ? AND key = ?",
+            "SELECT value, created_at, ttl, expires_at FROM kv_entries WHERE collection = ? AND key = ?",
             [collection, key],
         ).fetchone()
 
         if result is None:
             return None
 
-        value_json = result[0]
-        return ManagedEntry.from_json(json_str=value_json)
+        value_data, created_at, ttl, expires_at = result
+
+        # Convert value from JSON/TEXT to dict
+        # If it's already a dict (from JSON column), use it; otherwise parse from string
+        if isinstance(value_data, str):
+            import json
+            value = json.loads(value_data)
+        else:
+            value = value_data
+
+        # DuckDB always returns naive timestamps, but ManagedEntry expects timezone-aware ones
+        # Convert to timezone-aware UTC timestamps
+        from datetime import timezone
+
+        if created_at is not None and created_at.tzinfo is None:
+            created_at = created_at.replace(tzinfo=timezone.utc)
+        if expires_at is not None and expires_at.tzinfo is None:
+            expires_at = expires_at.replace(tzinfo=timezone.utc)
+
+        # Reconstruct ManagedEntry with metadata from columns
+        return ManagedEntry(
+            value=value,
+            created_at=created_at,
+            ttl=ttl,
+            expires_at=expires_at,
+        )
 
     @override
     async def _put_managed_entry(
@@ -170,25 +218,44 @@ async def _put_managed_entry(
         collection: str,
         managed_entry: ManagedEntry,
     ) -> None:
-        """Store a managed entry by key in the specified collection."""
+        """Store a managed entry by key in the specified collection.
+
+        Stores the value and metadata separately:
+        - value: JSON string of just the value data (not full ManagedEntry)
+        - created_at, ttl, expires_at: Stored in native columns for efficient querying
+        """
         if self._is_closed:
             msg = "Cannot operate on closed DuckDBStore"
             raise RuntimeError(msg)
 
-        # Insert or replace the entry
+        # Get just the value as JSON (not the full ManagedEntry)
+        value_json = managed_entry.value_as_json
+
+        # Ensure timestamps are timezone-aware (convert naive to UTC if needed)
+        from datetime import timezone
+
+        created_at = managed_entry.created_at
+        if created_at is not None and created_at.tzinfo is None:
+            created_at = created_at.replace(tzinfo=timezone.utc)
+
+        expires_at = managed_entry.expires_at
+        if expires_at is not None and expires_at.tzinfo is None:
+            expires_at = expires_at.replace(tzinfo=timezone.utc)
+
+        # Insert or replace the entry with metadata in separate columns
         self._connection.execute(
             """
             INSERT OR REPLACE INTO kv_entries
-            (collection, key, value_json, created_at, ttl, expires_at)
+            (collection, key, value, created_at, ttl, expires_at)
             VALUES (?, ?, ?, ?, ?, ?)
         """,
             [
                 collection,
                 key,
-                managed_entry.to_json(),
-                managed_entry.created_at.timestamp() if managed_entry.created_at else None,
+                value_json,
+                created_at,
                 managed_entry.ttl,
-                managed_entry.expires_at.timestamp() if managed_entry.expires_at else None,
+                expires_at,
             ],
         )
 

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

@@ -38,6 +38,21 @@ async def store(self) -> AsyncGenerator[DuckDBStore, None]:
     async def test_not_unbounded(self, store: BaseStore): ...
 
 
+class TestDuckDBStoreTextMode(ContextManagerStoreTestMixin, BaseStoreTests):
+    """Test DuckDB store with TEXT column mode instead of JSON."""
+
+    @override
+    @pytest.fixture
+    async def store(self) -> AsyncGenerator[DuckDBStore, None]:
+        """Test with in-memory DuckDB database using TEXT column."""
+        duckdb_store = DuckDBStore(use_json_column=False)
+        yield duckdb_store
+        await duckdb_store.close()
+
+    @pytest.mark.skip(reason="Local disk stores are unbounded")
+    async def test_not_unbounded(self, store: BaseStore): ...
+
+
 class TestDuckDBStoreSpecific:
     """Test DuckDB-specific functionality."""
 
@@ -48,6 +63,76 @@ async def store(self) -> AsyncGenerator[DuckDBStore, None]:
         yield duckdb_store
         await duckdb_store.close()
 
+    async def test_native_sql_queryability(self):
+        """Test that users can query the database directly with SQL."""
+        store = DuckDBStore(use_json_column=True)
+
+        # Store some test data with known metadata
+        await store.put(collection="products", key="item1", value={"name": "Widget", "price": 10.99}, ttl=3600)
+        await store.put(collection="products", key="item2", value={"name": "Gadget", "price": 25.50}, ttl=7200)
+        await store.put(collection="orders", key="order1", value={"total": 100.00, "items": 3})
+
+        # Query directly via SQL to verify native storage
+        # Check that value is stored as JSON (can extract fields)
+        result = store._connection.execute("""
+            SELECT key, value->'name' as name, value->'price' as price
+            FROM kv_entries
+            WHERE collection = 'products'
+            ORDER BY key
+        """).fetchall()
+
+        assert len(result) == 2
+        assert result[0][0] == "item1"
+        assert result[0][1] == '"Widget"'  # JSON strings are quoted
+        assert result[1][0] == "item2"
+
+        # Query by expiration timestamp
+        result = store._connection.execute("""
+            SELECT COUNT(*)
+            FROM kv_entries
+            WHERE expires_at > now() OR expires_at IS NULL
+        """).fetchone()
+
+        assert result[0] == 3  # All 3 entries should not be expired
+
+        # Query metadata columns directly
+        result = store._connection.execute("""
+            SELECT key, ttl, created_at IS NOT NULL as has_created
+            FROM kv_entries
+            WHERE collection = 'products' AND ttl > 3600
+        """).fetchall()
+
+        assert len(result) == 1  # Only item2 has ttl > 3600
+        assert result[0][0] == "item2"
+        assert result[0][1] == 7200
+        assert result[0][2] is True  # has_created
+
+        await store.close()
+
+    async def test_text_mode_storage(self):
+        """Test that TEXT mode stores value as string instead of native JSON."""
+        store = DuckDBStore(use_json_column=False)
+
+        await store.put(collection="test", key="key1", value={"data": "value"})
+
+        # Query to check column type - in TEXT mode, value should be a string
+        result = store._connection.execute("""
+            SELECT value, typeof(value) as value_type
+            FROM kv_entries
+            WHERE collection = 'test' AND key = 'key1'
+        """).fetchone()
+
+        assert result is not None
+        value_str, value_type = result
+
+        # In TEXT mode, value should be stored as VARCHAR/TEXT
+        assert value_type in ("VARCHAR", "TEXT")
+        # Value should be a JSON string
+        assert isinstance(value_str, str)
+        assert "data" in value_str
+
+        await store.close()
+
     async def test_database_path_initialization(self):
         """Test that store can be initialized with different database path options."""
         # In-memory (default)