Repository: PrefectHQ/prefect
Pull request: Allow child transactions to access parent values set after child transaction opens
<!-- 
Thanks for opening a pull request to Prefect! 
If this is your first contribution, please make sure to review our contribution guidelines: https://docs.prefect.io/latest/contributing/overview/
-->

<!-- Include an overview of the proposed changes here -->

This PR enables child transactions to access parent stored values set after the child transaction opened. This behavior is achieved by having child transactions delegate `.get` calls to their parent if they don't have a stored value for the provided key. The value from the parent transaction is copied and stored on the child transaction to avoid accidental mutation.

These changes resolve the issue in https://github.com/PrefectHQ/prefect/issues/15309 where a value set on a parent was accessed in a rollback hook since the rollback hook gets the task's transaction. These changes also mean that we don't have to copy all stored values from parent to child, which should reduce memory usage in some scenarios.

Closes https://github.com/PrefectHQ/prefect/issues/15309

### Checklist
<!-- These boxes may be checked after opening the pull request. -->

- [x] This pull request references any related issue by including "closes `<link to issue>`"
  - If no issue exists and your change is not a small fix, please [create an issue](https://github.com/PrefectHQ/prefect/issues/new/choose) first.
- [x] If this pull request adds new functionality, it includes unit tests that cover the changes
- [ ] If this pull request removes docs files, it includes redirect settings in `mint.json`.
- [ ] If this pull request adds functions or classes, it includes helpful docstrings.


Problem statement:
`transaction.get` functionality depends on `transaction.set` location within flow
### Bug summary

Accessing a transaction key-value pair that is set within the transaction context block without calling `get_transaction` is sensitive to where the `.set` method is called.

MRE where the rollback hook fails. Move the `transaction.set` above the call to `write_file` and then the rollback hook succeeds. 

```python
import os
from time import sleep

from prefect import task, flow
from prefect.transactions import get_transaction, transaction


@task
def write_file(contents: str):
    "Writes to a file."
    with open("side-effect.txt", "w") as f:
        f.write(contents)


@write_file.on_rollback
def del_file(txn):
    "Deletes file."
    os.unlink("side-effect.txt")  # get the transaction object
    print(txn.get("contents"))  # prints "hello world"


@task
def quality_test():
    "Checks contents of file."
    with open("side-effect.txt", "r") as f:
        data = f.readlines()

    if len(data) < 2:
        raise ValueError(f"Not enough data! ")
        # prints "hello world" as part of the error message


@flow
def pipeline(contents: str):
    with transaction() as txn:
        write_file(contents)
        sleep(2)
        txn.set(
            "contents", contents
        )  # set the contents key to the value of the contents variable
        quality_test()


if __name__ == "__main__":
    pipeline(contents="hello world")
```

Explicitly calling `get_transaction` also works, regardless of where the transaction is set.

cc: @cicdw 

### Version info (`prefect version` output)

```Text
Version:             3.0.1
API version:         0.8.4
Python version:      3.12.2
Git commit:          c6b2ffe1
Built:               Fri, Sep 6, 2024 10:05 AM
OS/Arch:             darwin/arm64
Profile:             sandbox-jeff
Server type:         cloud
Pydantic version:    2.8.2
```


### Additional context

_No response_

Hints from issue discussion:

Commit message: Allow child transactions to access parent values set after child transaction opens

Patch to review:
diff --git a/src/prefect/transactions.py b/src/prefect/transactions.py
index 138b215e56b5..ff4a9fb0bbe5 100644
--- a/src/prefect/transactions.py
+++ b/src/prefect/transactions.py
@@ -81,11 +81,22 @@ def set(self, name: str, value: Any) -> None:
         self._stored_values[name] = value
 
     def get(self, name: str, default: Any = NotSet) -> Any:
-        if name not in self._stored_values:
+        value = self._stored_values.get(name, NotSet)
+        if value is NotSet:
+            parent = self.get_parent()
+            if parent is not None:
+                try:
+                    value = copy.deepcopy(parent.get(name))
+                    self._stored_values[name] = value
+                except ValueError:
+                    pass
+
+        if value is NotSet:
             if default is not NotSet:
-                return default
-            raise ValueError(f"Could not retrieve value for unknown key: {name}")
-        return self._stored_values.get(name)
+                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 +119,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
diff --git a/tests/test_transactions.py b/tests/test_transactions.py
index 858154b50963..72853d11b77a 100644
--- a/tests/test_transactions.py
+++ b/tests/test_transactions.py
@@ -608,6 +608,34 @@ 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])
+
+            # child can't modify parent's values
+            assert child_transaction.get("list") == [1, 2, 3]
+            child_transaction.get("list").append(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):


Code review comments:
diff_hunk: @@ -81,11 +81,22 @@ def set(self, name: str, value: Any) -> None:
         self._stored_values[name] = value
 
     def get(self, name: str, default: Any = NotSet) -> Any:
-        if name not in self._stored_values:
+        value = self._stored_values.get(name, NotSet)
+        if value is NotSet:
+            parent = self.get_parent()
+            if parent is not None:
+                try:
+                    value = copy.deepcopy(parent.get(name))
+                    self._stored_values[name] = value
path: src/prefect/transactions.py
text: @user1:
this feels like it could lead to some edge cases -- I get that this serves as a cache but is there a functional / behaviorial reason to update this name on `self`?

@user1:
(the `deepcopy` makes sense to avoid mutations)

@author:
Yeah, you can see the use case I was thinking of in the test. If I get and mutate a value (e.g. append to a list) I'd expect to see that update the next time I get the value. That makes storing the value on the child transaction necessary.

@user1:
ok interesting, but there's an analogous situation that could be equally argued for in which the parent mutates the value after the child calls `get`; only sharing that because the lack of obvious priority there suggests to me that there's a better interface for getting data into transaction hooks.  Should we consider allowing for the same signature as state change hooks? 

@user1:
in the meantime, i'm OK with the version you have here

@author:
That's a good point about the order still mattering. What if said that you can only update a stored value in a transaction by calling `.set()`? Mutations to values returned by `.get` wouldn't propagate the stored values. That seems like it'd be more predictable.

@user1:
Yea I really like that; in that case we'd always return a `deepcopy` from `get`, right?

@author:
Yeah that'd be necessary to make the stored values "immutable". I can make that update an see how it feels.

@author:
@user1 I've made the update to return copies of stored values and it feels pretty good. Let me know what you think!