Merge remote-tracking branch 'upstream/main' into improving-ci-coverage-and-configuring-.coveragerc

Author: leseb

src/ogx/core/storage/kvstore/kvstore.py

@@ -13,7 +13,7 @@
 
 import asyncio
 from collections import defaultdict
-from datetime import datetime
+from datetime import UTC, datetime
 from typing import cast
 
 from ogx.core.storage.datatypes import (
@@ -43,26 +43,66 @@ def kvstore_dependencies() -> list[str]:
 class InmemoryKVStoreImpl(KVStore):
     """In-memory key-value store implementation for testing and ephemeral usage."""
 
-    def __init__(self) -> None:
-        self._store: dict[str, str] = {}
+    def __init__(self, namespace: str | None = None) -> None:
+        self._store: dict[str, tuple[str, datetime | None]] = {}
+        self._namespace = namespace
+
+    def _namespaced_key(self, key: str) -> str:
+        if not self._namespace:
+            return key
+        return f"{self._namespace}:{key}"
+
+    def _strip_namespace(self, key: str) -> str:
+        if self._namespace and key.startswith(f"{self._namespace}:"):
+            return key[len(self._namespace) + 1 :]
+        return key
+
+    def _is_expired(self, expiration: datetime | None) -> bool:
+        if expiration is None:
+            return False
+        return datetime.now(tz=UTC) >= expiration
 
     async def initialize(self) -> None:
         pass
 
     async def get(self, key: str) -> str | None:
-        return self._store.get(key)
+        key = self._namespaced_key(key)
+        entry = self._store.get(key)
+        if entry is None:
+            return None
+        value, expiration = entry
+        if self._is_expired(expiration):
+            return None
+        return value
 
     async def set(self, key: str, value: str, expiration: datetime | None = None) -> None:
-        self._store[key] = value
+        key = self._namespaced_key(key)
+        self._store[key] = (value, expiration)
 
     async def values_in_range(self, start_key: str, end_key: str) -> list[str]:
-        return [self._store[key] for key in self._store.keys() if key >= start_key and key < end_key]
+        start_key = self._namespaced_key(start_key)
+        end_key = self._namespaced_key(end_key)
+        result = []
+        for key in sorted(self._store.keys()):
+            if key >= start_key and key < end_key:
+                value, expiration = self._store[key]
+                if not self._is_expired(expiration):
+                    result.append(value)
+        return result
 
     async def keys_in_range(self, start_key: str, end_key: str) -> list[str]:
-        """Get all keys in the given range."""
-        return [key for key in self._store.keys() if key >= start_key and key < end_key]
+        start_key = self._namespaced_key(start_key)
+        end_key = self._namespaced_key(end_key)
+        result = []
+        for key in sorted(self._store.keys()):
+            if key >= start_key and key < end_key:
+                _, expiration = self._store[key]
+                if not self._is_expired(expiration):
+                    result.append(self._strip_namespace(key))
+        return result
 
     async def delete(self, key: str) -> None:
+        key = self._namespaced_key(key)
         self._store.pop(key, None)
 
     async def shutdown(self) -> None:

src/ogx/core/storage/kvstore/postgres/postgres.py

@@ -64,6 +64,11 @@ def _namespaced_key(self, key: str) -> str:
             return key
         return f"{self.config.namespace}:{key}"
 
+    def _strip_namespace(self, key: str) -> str:
+        if self.config.namespace and key.startswith(f"{self.config.namespace}:"):
+            return key[len(self.config.namespace) + 1 :]
+        return key
+
     async def set(self, key: str, value: str, expiration: datetime | None = None) -> None:
         key = self._namespaced_key(key)
         cursor = self._cursor_or_raise()
@@ -129,7 +134,7 @@ async def keys_in_range(self, start_key: str, end_key: str) -> list[str]:
             """,
             (start_key, end_key),
         )
-        return [row[0] for row in cursor.fetchall()]
+        return [self._strip_namespace(row[0]) for row in cursor.fetchall()]
 
     async def shutdown(self) -> None:
         if self._cursor:

src/ogx/core/storage/kvstore/sqlite/sqlite.py

@@ -5,7 +5,7 @@
 # the root directory of this source tree.
 
 import os
-from datetime import datetime
+from datetime import UTC, datetime
 
 import aiosqlite
 
@@ -23,6 +23,7 @@ class SqliteKVStoreImpl(KVStore):
     def __init__(self, config: SqliteKVStoreConfig) -> None:
         self.db_path = config.db_path
         self.table_name = "kvstore"
+        self._namespace = config.namespace
         self._conn: aiosqlite.Connection | None = None
 
     def __str__(self) -> str:
@@ -32,6 +33,16 @@ def _is_memory_db(self) -> bool:
         """Check if this is an in-memory database."""
         return self.db_path == ":memory:" or "mode=memory" in self.db_path
 
+    def _namespaced_key(self, key: str) -> str:
+        if not self._namespace:
+            return key
+        return f"{self._namespace}:{key}"
+
+    def _strip_namespace(self, key: str) -> str:
+        if self._namespace and key.startswith(f"{self._namespace}:"):
+            return key[len(self._namespace) + 1 :]
+        return key
+
     async def initialize(self) -> None:
         # Skip directory creation for in-memory databases and file: URIs
         if not self._is_memory_db() and not self.db_path.startswith("file:"):
@@ -74,103 +85,90 @@ async def shutdown(self) -> None:
             self._conn = None
 
     async def set(self, key: str, value: str, expiration: datetime | None = None) -> None:
+        key = self._namespaced_key(key)
+        exp_str = expiration.isoformat() if expiration else None
         if self._conn:
-            # In-memory database with persistent connection
             await self._conn.execute(
                 f"INSERT OR REPLACE INTO {self.table_name} (key, value, expiration) VALUES (?, ?, ?)",
-                (key, value, expiration),
+                (key, value, exp_str),
             )
             await self._conn.commit()
         else:
-            # File-based database with connection per operation
             async with aiosqlite.connect(self.db_path) as db:
                 await db.execute(
                     f"INSERT OR REPLACE INTO {self.table_name} (key, value, expiration) VALUES (?, ?, ?)",
-                    (key, value, expiration),
+                    (key, value, exp_str),
                 )
                 await db.commit()
 
     async def get(self, key: str) -> str | None:
+        key = self._namespaced_key(key)
+        now = datetime.now(tz=UTC).isoformat()
+        query = f"SELECT value FROM {self.table_name} WHERE key = ? AND (expiration IS NULL OR expiration > ?)"
         if self._conn:
-            # In-memory database with persistent connection
-            async with self._conn.execute(
-                f"SELECT value, expiration FROM {self.table_name} WHERE key = ?", (key,)
-            ) as cursor:
+            async with self._conn.execute(query, (key, now)) as cursor:
                 row = await cursor.fetchone()
                 if row is None:
                     return None
-                value, expiration = row
+                value = row[0]
                 if not isinstance(value, str):
                     logger.warning("Expected string value for key, returning None", key=key, value_type=type(value))
                     return None
                 return value
         else:
-            # File-based database with connection per operation
             async with aiosqlite.connect(self.db_path) as db:
-                async with db.execute(
-                    f"SELECT value, expiration FROM {self.table_name} WHERE key = ?", (key,)
-                ) as cursor:
+                async with db.execute(query, (key, now)) as cursor:
                     row = await cursor.fetchone()
                     if row is None:
                         return None
-                    value, expiration = row
+                    value = row[0]
                     if not isinstance(value, str):
                         logger.warning("Expected string value for key, returning None", key=key, value_type=type(value))
                         return None
                     return value
 
     async def delete(self, key: str) -> None:
+        key = self._namespaced_key(key)
         if self._conn:
-            # In-memory database with persistent connection
             await self._conn.execute(f"DELETE FROM {self.table_name} WHERE key = ?", (key,))
             await self._conn.commit()
         else:
-            # File-based database with connection per operation
             async with aiosqlite.connect(self.db_path) as db:
                 await db.execute(f"DELETE FROM {self.table_name} WHERE key = ?", (key,))
                 await db.commit()
 
     async def values_in_range(self, start_key: str, end_key: str) -> list[str]:
+        start_key = self._namespaced_key(start_key)
+        end_key = self._namespaced_key(end_key)
+        now = datetime.now(tz=UTC).isoformat()
+        query = (
+            f"SELECT value FROM {self.table_name} "
+            f"WHERE key >= ? AND key < ? AND (expiration IS NULL OR expiration > ?) "
+            f"ORDER BY key"
+        )
         if self._conn:
-            # In-memory database with persistent connection
-            async with self._conn.execute(
-                f"SELECT key, value, expiration FROM {self.table_name} WHERE key >= ? AND key <= ?",
-                (start_key, end_key),
-            ) as cursor:
-                result = []
-                async for row in cursor:
-                    _, value, _ = row
-                    result.append(value)
-                return result
+            async with self._conn.execute(query, (start_key, end_key, now)) as cursor:
+                return [row[0] async for row in cursor]
         else:
-            # File-based database with connection per operation
             async with aiosqlite.connect(self.db_path) as db:
-                async with db.execute(
-                    f"SELECT key, value, expiration FROM {self.table_name} WHERE key >= ? AND key <= ?",
-                    (start_key, end_key),
-                ) as cursor:
-                    result = []
-                    async for row in cursor:
-                        _, value, _ = row
-                        result.append(value)
-                    return result
+                async with db.execute(query, (start_key, end_key, now)) as cursor:
+                    return [row[0] async for row in cursor]
 
     async def keys_in_range(self, start_key: str, end_key: str) -> list[str]:
-        """Get all keys in the given range."""
+        start_key = self._namespaced_key(start_key)
+        end_key = self._namespaced_key(end_key)
+        now = datetime.now(tz=UTC).isoformat()
+        query = (
+            f"SELECT key FROM {self.table_name} "
+            f"WHERE key >= ? AND key < ? AND (expiration IS NULL OR expiration > ?) "
+            f"ORDER BY key"
+        )
         if self._conn:
-            # In-memory database with persistent connection
-            cursor = await self._conn.execute(
-                f"SELECT key FROM {self.table_name} WHERE key >= ? AND key <= ?",
-                (start_key, end_key),
-            )
+            cursor = await self._conn.execute(query, (start_key, end_key, now))
             rows = await cursor.fetchall()
-            return [row[0] for row in rows]
+            return [self._strip_namespace(row[0]) for row in rows]
         else:
-            # File-based database with connection per operation
             async with aiosqlite.connect(self.db_path) as db:
-                cursor = await db.execute(
-                    f"SELECT key FROM {self.table_name} WHERE key >= ? AND key <= ?",
-                    (start_key, end_key),
-                )
+                cursor = await db.execute(query, (start_key, end_key, now))
                 rows = await cursor.fetchall()
-                return [row[0] for row in rows]
+                return [self._strip_namespace(row[0]) for row in rows]

tests/unit/registry/test_registry.py

@@ -8,8 +8,8 @@
 import pytest
 
 from ogx.core.datatypes import VectorStoreWithOwner
-from ogx.core.storage.datatypes import KVStoreReference, SqliteKVStoreConfig
-from ogx.core.storage.kvstore import kvstore_impl, register_kvstore_backends
+from ogx.core.storage.datatypes import SqliteKVStoreConfig
+from ogx.core.storage.kvstore.sqlite.sqlite import SqliteKVStoreImpl
 from ogx.core.store.registry import (
     KEY_FORMAT,
     CachedDiskDistributionRegistry,
@@ -69,14 +69,13 @@ async def test_cached_registry_initialization(sqlite_kvstore, sample_vector_stor
     await disk_registry.register(sample_vector_store)
     await disk_registry.register(sample_model)
 
-    # Test cached version loads from disk
+    # Test cached version loads from disk via a fresh KVStore pointing at the same DB
     db_path = sqlite_kvstore.db_path
-    backend_name = "kv_cached_test"
-    register_kvstore_backends({backend_name: SqliteKVStoreConfig(db_path=db_path)})
+    fresh_config = SqliteKVStoreConfig(db_path=db_path)
+    fresh_kvstore = SqliteKVStoreImpl(fresh_config)
+    await fresh_kvstore.initialize()
     # Use cache_ttl_seconds=0 for tests to ensure immediate synchronization
-    cached_registry = CachedDiskDistributionRegistry(
-        await kvstore_impl(KVStoreReference(backend=backend_name, namespace="registry")), cache_ttl_seconds=0
-    )
+    cached_registry = CachedDiskDistributionRegistry(fresh_kvstore, cache_ttl_seconds=0)
     await cached_registry.initialize()
 
     result_vector_store = await cached_registry.get("vector_store", "test_vector_store")
@@ -103,13 +102,12 @@ async def test_cached_registry_updates(cached_disk_dist_registry):
     assert result_vector_store.identifier == new_vector_store.identifier
     assert result_vector_store.provider_id == new_vector_store.provider_id
 
-    # Verify persisted to disk
+    # Verify persisted to disk via a fresh KVStore pointing at the same DB
     db_path = cached_disk_dist_registry.kvstore.db_path
-    backend_name = "kv_cached_new"
-    register_kvstore_backends({backend_name: SqliteKVStoreConfig(db_path=db_path)})
-    new_registry = DiskDistributionRegistry(
-        await kvstore_impl(KVStoreReference(backend=backend_name, namespace="registry"))
-    )
+    fresh_config = SqliteKVStoreConfig(db_path=db_path)
+    fresh_kvstore = SqliteKVStoreImpl(fresh_config)
+    await fresh_kvstore.initialize()
+    new_registry = DiskDistributionRegistry(fresh_kvstore)
     await new_registry.initialize()
     result_vector_store = await new_registry.get("vector_store", "test_vector_store_2")
     assert result_vector_store is not None