Merge pull request #241 from finos/fix/sql_strategy

fix(storage): repair SQL/DynamoDB strategies + add cross-backend contract tests

Author: kirillsc

src/orb/domain/base/domain_interfaces.py

@@ -5,7 +5,7 @@
 from abc import ABC, abstractmethod
 from typing import Any, Generic, Optional, Protocol, TypeVar
 
-from .entity import AggregateRoot, Entity
+from .entity import AggregateRoot
 
 T = TypeVar("T")  # Generic type for domain entities/aggregates
 A = TypeVar("A", bound=AggregateRoot)
@@ -90,18 +90,6 @@ def __exit__(self, exc_type, exc_val, exc_tb) -> None:
         """Exit the unit of work context."""
         ...
 
-    def register_new(self, entity: Entity) -> None:
-        """Register a new entity."""
-        ...
-
-    def register_dirty(self, entity: Entity) -> None:
-        """Register a dirty entity."""
-        ...
-
-    def register_removed(self, entity: Entity) -> None:
-        """Register a removed entity."""
-        ...
-
     @abstractmethod
     def begin(self) -> None:
         """Begin a new transaction."""

src/orb/infrastructure/storage/components/sql_query_builder.py

@@ -33,6 +33,7 @@ def __init__(self, table_name: str, columns: dict[str, str]) -> None:
             table_name: Name of the database table
             columns: Dictionary of column names and types
         """
+        super().__init__()
         self.table_name = table_name
         self.columns = columns
         self.logger = get_logger(__name__)
@@ -42,6 +43,64 @@ def __init__(self, table_name: str, columns: dict[str, str]) -> None:
         for column in columns:
             self._validate_identifier(column)
 
+    def build_query(self, query_spec: dict[str, Any]) -> str:
+        """
+        Build a SQL query from specification (implements QueryManager interface).
+
+        Dispatches to existing build_* methods based on query_spec["type"].
+        """
+        query_type = query_spec.get("type", "").upper()
+        if query_type == "CREATE_TABLE":
+            return self.build_create_table()
+        if query_type == "SELECT_ALL":
+            return self.build_select_all()
+        if query_type == "SELECT_BY_ID":
+            id_column = query_spec.get("id_column", "id")
+            sql, _ = self.build_select_by_id(id_column)
+            return sql
+        if query_type == "INSERT":
+            sql, _ = self.build_insert(query_spec.get("data", {}))
+            return sql
+        if query_type == "UPDATE":
+            sql, _ = self.build_update(
+                query_spec.get("data", {}),
+                query_spec.get("id_column", "id"),
+                query_spec.get("entity_id", ""),
+            )
+            return sql
+        if query_type == "DELETE":
+            sql, _ = self.build_delete(query_spec.get("id_column", "id"))
+            return sql
+        raise ValueError(f"Unknown query type: {query_type}")
+
+    def execute_query(self, query: str, parameters: Optional[dict[str, Any]] = None) -> Any:
+        """
+        Reject direct execution — execution belongs to SQLConnectionManager.
+
+        SQLQueryBuilder is build-only; this method exists to satisfy the
+        QueryManager interface but should not be used. Callers must execute
+        the built query via SQLConnectionManager.execute_query.
+        """
+        raise NotImplementedError(
+            "SQLQueryBuilder does not execute queries; use SQLConnectionManager.execute_query"
+        )
+
+    def validate_query(self, query: str) -> bool:
+        """
+        Validate a SQL query string (implements QueryManager interface).
+
+        Performs basic structural validation: non-empty, contains a known
+        SQL keyword, and no unbalanced quotes.
+        """
+        if not query or not query.strip():
+            return False
+        keywords = ("SELECT", "INSERT", "UPDATE", "DELETE", "CREATE")
+        if not any(kw in query.upper() for kw in keywords):
+            return False
+        if query.count("'") % 2 != 0:
+            return False
+        return True
+
     def _validate_identifier(self, identifier: str) -> None:
         """
         Validate SQL identifier against whitelist pattern.

src/orb/infrastructure/storage/sql/registration.py

@@ -118,12 +118,12 @@ def create_sql_unit_of_work(config: Any) -> Any:
             echo=getattr(sql_config, "echo", False),
         )
 
-        return SQLUnitOfWork(engine)  # type: ignore[abstract]
+        return SQLUnitOfWork(engine)
     else:
         # For testing or other scenarios - assume it's a dict with connection info
         connection_string = config.get("connection_string", "sqlite:///data/test.db")
         engine = create_engine(connection_string)
-        return SQLUnitOfWork(engine)  # type: ignore[abstract]
+        return SQLUnitOfWork(engine)
 
 
 def register_sql_storage() -> None:

src/orb/infrastructure/storage/sql/strategy.py

@@ -43,7 +43,7 @@ def __init__(self, config: dict[str, Any], table_name: str, columns: dict[str, s
 
         # Initialize components
         self.connection_manager = SQLConnectionManager(config)
-        self.query_builder = SQLQueryBuilder(table_name, columns)  # type: ignore[abstract]
+        self.query_builder = SQLQueryBuilder(table_name, columns)
         self.serializer = SQLSerializer(id_column=self._get_id_column())
         self.lock_manager = LockManager("simple")  # Simple lock for SQL
 
@@ -265,7 +265,7 @@ def save_batch(self, entities: dict[str, dict[str, Any]]) -> None:
 
                 with self.connection_manager.get_session() as session:
                     for serialized_data in serialized_list:
-                        session.execute(query, serialized_data)
+                        session.execute(text(query), serialized_data)
                     session.commit()
 
                 self.logger.debug("Saved batch of %s entities", len(entities))

src/orb/providers/aws/storage/components/dynamodb_transaction_manager.py

@@ -3,6 +3,7 @@
 from contextlib import contextmanager
 from typing import Any, Callable, Optional
 
+from boto3.dynamodb.types import TypeSerializer
 from botocore.exceptions import ClientError
 
 from orb.infrastructure.storage.components.transaction_manager import (
@@ -30,6 +31,11 @@ def __init__(self, client_manager) -> None:
         self.client_manager = client_manager
         self.transaction_items: list[dict[str, Any]] = []
         self.max_transaction_items = 25  # DynamoDB limit
+        # transact_write_items uses the low-level Client API, which requires
+        # AttributeValue-tagged dicts (e.g. {"id": {"S": "abc"}}). The
+        # higher-level Resource API (used elsewhere) accepts plain dicts;
+        # this serializer converts at the boundary.
+        self._serializer = TypeSerializer()
 
     def begin_transaction(self) -> None:
         """Begin a new DynamoDB transaction."""
@@ -60,7 +66,8 @@ def add_put_item(
         if len(self.transaction_items) >= self.max_transaction_items:
             raise RuntimeError(f"Transaction cannot exceed {self.max_transaction_items} items")
 
-        put_request = {"Put": {"TableName": table_name, "Item": item}}
+        serialized_item = {k: self._serializer.serialize(v) for k, v in item.items()}
+        put_request = {"Put": {"TableName": table_name, "Item": serialized_item}}
 
         if condition_expression:
             put_request["Put"]["ConditionExpression"] = condition_expression
@@ -92,12 +99,16 @@ def add_update_item(
         if len(self.transaction_items) >= self.max_transaction_items:
             raise RuntimeError(f"Transaction cannot exceed {self.max_transaction_items} items")
 
+        serialized_key = {k: self._serializer.serialize(v) for k, v in key.items()}
+        serialized_values = {
+            k: self._serializer.serialize(v) for k, v in expression_attribute_values.items()
+        }
         update_request = {
             "Update": {
                 "TableName": table_name,
-                "Key": key,
+                "Key": serialized_key,
                 "UpdateExpression": update_expression,
-                "ExpressionAttributeValues": expression_attribute_values,
+                "ExpressionAttributeValues": serialized_values,
             }
         }
 
@@ -127,7 +138,8 @@ def add_delete_item(
         if len(self.transaction_items) >= self.max_transaction_items:
             raise RuntimeError(f"Transaction cannot exceed {self.max_transaction_items} items")
 
-        delete_request = {"Delete": {"TableName": table_name, "Key": key}}
+        serialized_key = {k: self._serializer.serialize(v) for k, v in key.items()}
+        delete_request = {"Delete": {"TableName": table_name, "Key": serialized_key}}
 
         if condition_expression:
             delete_request["Delete"]["ConditionExpression"] = condition_expression

tests/integration/storage/__init__.py

@@ -0,0 +1,137 @@
+"""Fixtures for storage strategy contract tests.
+
+Provides parameterised `storage_strategy` and `unit_of_work` fixtures so
+each backend (JSON, SQL/SQLite, DynamoDB/moto) runs the same contract tests.
+"""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import Iterator
+
+import pytest
+
+try:
+    from moto import mock_aws as _moto_mock_aws
+
+    HAS_MOTO = True
+except ImportError:
+    HAS_MOTO = False
+    _moto_mock_aws = None
+
+
+@contextmanager
+def _mock_aws():
+    if _moto_mock_aws is None:
+        yield
+        return
+    with _moto_mock_aws():
+        yield
+
+
+_ENTITY_TABLE = "entities"
+_ENTITY_COLUMNS = {"id": "TEXT PRIMARY KEY", "data": "TEXT", "name": "TEXT"}
+
+
+# ---------------------------------------------------------------------------
+# Strategy fixtures (low-level: SQLStorageStrategy / JSONStorageStrategy / DynamoDB)
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def json_strategy(tmp_path):
+    from orb.infrastructure.storage.json.strategy import JSONStorageStrategy
+
+    file_path = tmp_path / "entities.json"
+    return JSONStorageStrategy(file_path=str(file_path), entity_type="entities")
+
+
+@pytest.fixture
+def sql_strategy():
+    from orb.infrastructure.storage.sql.strategy import SQLStorageStrategy
+
+    return SQLStorageStrategy(
+        config={"type": "sqlite", "name": ":memory:"},
+        table_name=_ENTITY_TABLE,
+        columns=_ENTITY_COLUMNS,
+    )
+
+
+@pytest.fixture
+def dynamodb_strategy() -> Iterator:
+    if not HAS_MOTO:
+        pytest.skip("moto not installed")
+
+    from orb.providers.aws.storage.strategy import DynamoDBStorageStrategy
+
+    with _mock_aws():
+        # aws_client=None forces internal boto3 session, which moto intercepts.
+        from orb.infrastructure.adapters.logging_adapter import LoggingAdapter
+
+        strategy = DynamoDBStorageStrategy(
+            logger=LoggingAdapter("test.dynamo"),
+            aws_client=None,
+            region="us-east-1",
+            table_name=_ENTITY_TABLE,
+        )
+        yield strategy
+
+
+@pytest.fixture(params=["json", "sql", "dynamodb"])
+def storage_strategy(request, json_strategy, sql_strategy, dynamodb_strategy):
+    """Parameterised strategy fixture used by contract tests."""
+    return {
+        "json": json_strategy,
+        "sql": sql_strategy,
+        "dynamodb": dynamodb_strategy,
+    }[request.param]
+
+
+# ---------------------------------------------------------------------------
+# UnitOfWork fixtures
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def json_uow(tmp_path):
+    from orb.infrastructure.storage.json.unit_of_work import JSONUnitOfWork
+
+    return JSONUnitOfWork(data_dir=str(tmp_path))
+
+
+@pytest.fixture
+def sql_uow():
+    from sqlalchemy import create_engine
+
+    from orb.infrastructure.storage.sql.unit_of_work import SQLUnitOfWork
+
+    engine = create_engine("sqlite:///:memory:")
+    return SQLUnitOfWork(engine)
+
+
+@pytest.fixture
+def dynamodb_uow():
+    if not HAS_MOTO:
+        pytest.skip("moto not installed")
+
+    from orb.providers.aws.storage.unit_of_work import DynamoDBUnitOfWork
+
+    from orb.infrastructure.adapters.logging_adapter import LoggingAdapter
+
+    with _mock_aws():
+        uow = DynamoDBUnitOfWork(
+            aws_client=None,
+            logger=LoggingAdapter("test.dynamo.uow"),
+            region="us-east-1",
+        )
+        yield uow
+
+
+@pytest.fixture(params=["json", "sql", "dynamodb"])
+def unit_of_work(request, json_uow, sql_uow, dynamodb_uow):
+    """Parameterised UoW fixture used by UoW contract tests."""
+    return {
+        "json": json_uow,
+        "sql": sql_uow,
+        "dynamodb": dynamodb_uow,
+    }[request.param]

tests/integration/storage/conftest.py

@@ -0,0 +1,32 @@
+"""Aurora SQL storage factory tests.
+
+Aurora's storage strategy is SQLAlchemy-based; the strategy and UoW
+themselves are exercised by SQL contract tests against sqlite. These
+tests verify the Aurora factory wires its connection string and engine
+into a working SQLUnitOfWork without contacting AWS.
+"""
+
+import pytest
+
+
+@pytest.mark.integration
+class TestAuroraFactory:
+    def test_create_aurora_strategy_with_simple_config(self):
+        from orb.providers.aws.storage.registration import create_aurora_strategy
+        from orb.infrastructure.storage.sql.strategy import SQLStorageStrategy
+
+        class _Cfg:
+            connection_string = "sqlite:///:memory:"
+
+        strategy = create_aurora_strategy(_Cfg())
+        assert isinstance(strategy, SQLStorageStrategy)
+
+    def test_create_aurora_unit_of_work_with_dict_config(self):
+        from orb.providers.aws.storage.registration import create_aurora_unit_of_work
+        from orb.infrastructure.storage.sql.unit_of_work import SQLUnitOfWork
+
+        uow = create_aurora_unit_of_work({"connection_string": "sqlite:///:memory:"})
+        assert isinstance(uow, SQLUnitOfWork)
+        assert uow.machines is not None
+        assert uow.requests is not None
+        assert uow.templates is not None