Allow child transactions to access parent values set after child tran…

…saction opens (PrefectHQ#15342)

src/prefect/transactions.py

@@ -78,14 +78,85 @@ class Transaction(ContextModel):
     __var__: ContextVar = ContextVar("transaction")
 
     def set(self, name: str, value: Any) -> None:
+        """
+        Set a stored value in the transaction.
+
+        Args:
+            name: The name of the value to set
+            value: The value to set
+
+        Examples:
+            Set a value for use later in the transaction:
+            ```python
+            with transaction() as txn:
+                txn.set("key", "value")
+                ...
+                assert txn.get("key") == "value"
+            ```
+        """
         self._stored_values[name] = value
 
     def get(self, name: str, default: Any = NotSet) -> Any:
-        if name not in self._stored_values:
-            if default is not NotSet:
-                return default
-            raise ValueError(f"Could not retrieve value for unknown key: {name}")
-        return self._stored_values.get(name)
+        """
+        Get a stored value from the transaction.
+
+        Child transactions will return values from their parents unless a value with
+        the same name is set in the child transaction.
+
+        Direct changes to returned values will not update the stored value. To update the
+        stored value, use the `set` method.
+
+        Args:
+            name: The name of the value to get
+            default: The default value to return if the value is not found
+
+        Returns:
+            The value from the transaction
+
+        Examples:
+            Get a value from the transaction:
+            ```python
+            with transaction() as txn:
+                txn.set("key", "value")
+                ...
+                assert txn.get("key") == "value"
+            ```
+
+            Get a value from a parent transaction:
+            ```python
+            with transaction() as parent:
+                parent.set("key", "parent_value")
+                with transaction() as child:
+                    assert child.get("key") == "parent_value"
+            ```
+
+            Update a stored value:
+            ```python
+            with transaction() as txn:
+                txn.set("key", [1, 2, 3])
+                value = txn.get("key")
+                value.append(4)
+                # Stored value is not updated until `.set` is called
+                assert value == [1, 2, 3, 4]
+                assert txn.get("key") == [1, 2, 3]
+
+                txn.set("key", value)
+                assert txn.get("key") == [1, 2, 3, 4]
+            ```
+        """
+        # deepcopy to prevent mutation of stored values
+        value = copy.deepcopy(self._stored_values.get(name, NotSet))
+        if value is NotSet:
+            # if there's a parent transaction, get the value from the parent
+            parent = self.get_parent()
+            if parent is not None:
+                value = parent.get(name, default)
+            # if there's no parent transaction, use the default
+            elif default is not NotSet:
+                value = default
+            else:
+                raise ValueError(f"Could not retrieve value for unknown key: {name}")
+        return value
 
     def is_committed(self) -> bool:
         return self.state == TransactionState.COMMITTED
@@ -108,8 +179,6 @@ def __enter__(self):
                 "Context already entered. Context enter calls cannot be nested."
             )
         parent = get_transaction()
-        if parent:
-            self._stored_values = copy.deepcopy(parent._stored_values)
         # set default commit behavior; either inherit from parent or set a default of eager
         if self.commit_mode is None:
             self.commit_mode = parent.commit_mode if parent else CommitMode.LAZY

tests/test_transactions.py

@@ -598,7 +598,9 @@ def test_get_and_set_data_doesnt_mutate_parent(self):
         with transaction(key="test") as top:
             top.set("key", {"x": [42]})
             with transaction(key="nested") as inner:
-                inner.get("key")["x"].append(43)
+                inner_value = inner.get("key")
+                inner_value["x"].append(43)
+                inner.set("key", inner_value)
                 assert inner.get("key") == {"x": [42, 43]}
                 assert top.get("key") == {"x": [42]}
             assert top.get("key") == {"x": [42]}
@@ -610,6 +612,35 @@ def test_get_raises_on_unknown_but_allows_default(self):
             assert txn.get("foobar", None) is None
             assert txn.get("foobar", "string") == "string"
 
+    def test_parent_values_set_after_child_open_are_available(self):
+        parent_transaction = Transaction()
+        child_transaction = Transaction()
+
+        parent_transaction.__enter__()
+        child_transaction.__enter__()
+
+        try:
+            parent_transaction.set("key", "value")
+
+            # child can access parent's values
+            assert child_transaction.get("key") == "value"
+
+            parent_transaction.set("list", [1, 2, 3])
+            assert child_transaction.get("list") == [1, 2, 3]
+
+            # Mutating the value doesn't update the stored value
+            child_transaction.get("list").append(4)
+            assert child_transaction.get("list") == [1, 2, 3]
+            child_transaction.set("list", [1, 2, 3, 4])
+            assert child_transaction.get("list") == [1, 2, 3, 4]
+
+            # parent transaction isn't affected by child's modifications
+            assert parent_transaction.get("list") == [1, 2, 3]
+
+        finally:
+            child_transaction.__exit__(None, None, None)
+            parent_transaction.__exit__(None, None, None)
+
 
 class TestIsolationLevel:
     def test_default_isolation_level(self):