1.0 hardening: fix @transaction early-return composition

Project.toml

@@ -22,6 +22,7 @@ Aqua = "0.8"
 ConcurrentUtilities = "2.1"
 DBInterface = "2.5"
 Dates = "1.10"
+Distributed = "1.10"
 Harbor = "1"
 JSON = "1"
 MD5 = "0.2"
@@ -39,9 +40,10 @@ julia = "1.10"
 
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 Harbor = "af79dbb9-1a80-47ad-8928-192a4af69376"
 Sockets = "6462fe0b-24de-5631-8697-dd941f90decc"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "Harbor", "Sockets", "Test"]
+test = ["Aqua", "Distributed", "Harbor", "Sockets", "Test"]

src/Postgres.jl

@@ -1218,57 +1218,60 @@ function DBInterface.transaction(f::F, conn::Connection) where {F}
     end
 end
 
-struct TransactionReturn{T} <: Exception
-    value::T
-end
-
-function rewrite_transaction_returns(expr)
-    expr isa Expr || return expr
-    if expr.head === :return
-        value = isempty(expr.args) ? nothing : rewrite_transaction_returns(only(expr.args))
-        marker = GlobalRef(@__MODULE__, :TransactionReturn)
-        return Expr(:call, GlobalRef(Core, :throw), Expr(:call, marker, value))
-    elseif expr.head === :function || expr.head === :(->) || expr.head === :quote
-        # A return in a nested function belongs to that function, not to the
-        # scope that contains this transaction macro.
-        return expr
-    end
-    return Expr(expr.head, map(rewrite_transaction_returns, expr.args)...)
-end
-
 """
     Postgres.@transaction conn expr
 
-Run `expr` inside a transaction: committed if it completes, rolled back if it
-throws. Evaluates to `expr`'s value.
+Run `expr` inside a transaction. Any non-exceptional exit commits: normal
+completion, `return` (which then returns from the enclosing function),
+`break`, or `continue`. Only a thrown exception rolls back. Evaluates to
+`expr`'s value. Nested `@transaction` blocks use savepoints, and an early
+`return` commits every enclosing level on its way out.
+
+The body keeps plain Julia semantics: a `return` inside a nested function,
+closure, `do`-block, or any task-forming macro (`Threads.@spawn`, `@async`,
+`Distributed.@spawnat`, third-party equivalents) belongs to that function or
+task, exactly as it would outside the macro.
 """
 macro transaction(conn, expr)
-    body = rewrite_transaction_returns(expr)
     quote
         # bind once: the connection expression may have side effects
         # (`@transaction acquire(pool) ...` would otherwise take a different
         # connection for the BEGIN, the COMMIT and the ROLLBACK)
         local c = $(esc(conn))
         local success = false
+        local completed = false
         start_transaction(c)
         try
-            local result
-            try
-                result = $(esc(body))
-            catch err
-                if err isa TransactionReturn
-                    commit(c)
-                    success = true
-                    return err.value
-                end
-                rethrow()
-            end
+            local result = $(esc(expr))
             commit(c)
             success = true
+            completed = true
             result
         catch
-            !success && rollback_for_failed_transaction!(c)
+            if !success
+                rollback_for_failed_transaction!(c)
+            end
+            completed = true
             rethrow()
+        finally
+            # A non-exceptional, non-local exit — return, break, continue —
+            # reaches here without passing the commit above or the catch:
+            # commit this level on the way out. `return` unwinds through every
+            # enclosing expansion's finally, so each level commits exactly
+            # once, innermost first; no AST rewriting is needed, and returns
+            # inside closures or task-forming macros keep their plain-Julia
+            # meaning untouched. If the commit fails, roll back THIS level
+            # before propagating (commit at savepoint depth leaves the depth
+            # unchanged on failure), so every enclosing level — macro
+            # expansion or plain catch — can then unwind its own.
+            if !completed
+                try
+                    commit(c)
+                catch
+                    rollback_for_failed_transaction!(c)
+                    rethrow()
+                end
+            end
         end
     end
 end

test/runtests.jl

@@ -1,6 +1,7 @@
 using Test
 using Aqua
 using Dates
+using Distributed
 using UUIDs
 using DBInterface
 using Tables
@@ -70,6 +71,17 @@ struct Int8Row
     s::String
 end
 
+# A third-party-style task macro the driver has never heard of: it wraps its
+# body in a Task and fetches it. @transaction must leave the body's `return`
+# with its plain meaning (the task's result) — no allowlist involved.
+macro local_task(body)
+    quote
+        local t = Task(() -> $(esc(body)))
+        schedule(t)
+        fetch(t)
+    end
+end
+
 # user-IO failure injection for the COPY hardening tests
 struct ThrowingSource <: IO end
 Base.eof(::ThrowingSource) = false
@@ -1494,6 +1506,195 @@ end
                     @test early_return(conn) === :early
                     @test !Postgres.in_transaction(conn)
                     @test only(Tables.rowtable(DBInterface.execute(conn, "SELECT count(*)::int AS n FROM macro_test"))).n == 3
+
+                    # An early return from a NESTED @transaction must commit
+                    # every level — the inner expansion must not intercept the
+                    # outer's marker and skip the outer commit — and must not
+                    # leave the connection inside a transaction.
+                    nested_return = function(c)
+                        Postgres.@transaction c begin
+                            DBInterface.execute(c, "INSERT INTO macro_test (value) VALUES (10)")
+                            Postgres.@transaction c begin
+                                DBInterface.execute(c, "INSERT INTO macro_test (value) VALUES (11)")
+                                return :nested_early
+                            end
+                        end
+                        return :late
+                    end
+                    @test nested_return(conn) === :nested_early
+                    @test !Postgres.in_transaction(conn)
+                    @test only(Tables.rowtable(DBInterface.execute(conn, "SELECT count(*)::int AS n FROM macro_test WHERE value IN (10, 11)"))).n == 2
+
+                    # A user catch inside the body must not intercept the
+                    # return marker and turn the return into its own value.
+                    catch_return = function(c)
+                        Postgres.@transaction c begin
+                            try
+                                DBInterface.execute(c, "INSERT INTO macro_test (value) VALUES (12)")
+                                return :from_try
+                            catch
+                                return :from_catch
+                            end
+                        end
+                        return :late
+                    end
+                    @test catch_return(conn) === :from_try
+                    @test !Postgres.in_transaction(conn)
+                    catch_var_return = function(c)
+                        Postgres.@transaction c begin
+                            try
+                                return :from_try2
+                            catch err
+                                return err
+                            end
+                        end
+                    end
+                    @test catch_var_return(conn) === :from_try2
+
+                    # A return inside a task-forming macro is that task's
+                    # result, not a transaction return.
+                    spawn_return = function(c)
+                        Postgres.@transaction c begin
+                            t = Threads.@spawn begin
+                                return :task_value
+                            end
+                            fetch(t)
+                        end
+                    end
+                    @test spawn_return(conn) === :task_value
+                    @test !Postgres.in_transaction(conn)
+
+                    # break and continue are deliberate non-exceptional exits:
+                    # they commit, like return, and leave no transaction open.
+                    for _ in 1:1
+                        Postgres.@transaction conn begin
+                            DBInterface.execute(conn, "INSERT INTO macro_test (value) VALUES (13)")
+                            break
+                        end
+                    end
+                    @test !Postgres.in_transaction(conn)
+                    @test only(Tables.rowtable(DBInterface.execute(conn, "SELECT count(*)::int AS n FROM macro_test WHERE value = 13"))).n == 1
+                    for _ in 1:2
+                        Postgres.@transaction conn begin
+                            DBInterface.execute(conn, "INSERT INTO macro_test (value) VALUES (14)")
+                            continue
+                        end
+                    end
+                    @test !Postgres.in_transaction(conn)
+                    @test only(Tables.rowtable(DBInterface.execute(conn, "SELECT count(*)::int AS n FROM macro_test WHERE value = 14"))).n == 2
+
+                    # A return inside a short-form function defined in the body
+                    # belongs to that function: it must not early-return the
+                    # enclosing function, and the helper must stay callable
+                    # after the block without leaking the marker.
+                    local escaped_helper
+                    shortform_result = (function(c)
+                        Postgres.@transaction c begin
+                            helper(x) = (x < 0 && return :neg; :pos)
+                            escaped_helper = helper
+                            (helper(-1), helper(1))
+                        end
+                    end)(conn)
+                    @test shortform_result == (:neg, :pos)
+                    @test !Postgres.in_transaction(conn)
+                    @test escaped_helper(-5) === :neg
+
+                    # The body keeps plain Julia semantics with no AST rewrite,
+                    # so a return inside ANY closure-forming construct behaves
+                    # exactly as it does outside the macro — including
+                    # third-party task macros no allowlist could cover.
+                    local_task_result = (function(c)
+                        Postgres.@transaction c begin
+                            DBInterface.execute(c, "INSERT INTO macro_test (value) VALUES (20)")
+                            @local_task begin
+                                return :local_task_value
+                            end
+                        end
+                    end)(conn)
+                    @test local_task_result === :local_task_value
+                    @test !Postgres.in_transaction(conn)
+                    @test only(Tables.rowtable(DBInterface.execute(conn, "SELECT count(*)::int AS n FROM macro_test WHERE value = 20"))).n == 1
+
+                    # Distributed task macros: with no workers added, worker 1
+                    # is this process, so these run locally end to end.
+                    for distributed_case in (
+                        (c) -> Postgres.@transaction(c, fetch(Distributed.@spawnat 1 begin
+                            return :spawnat_value
+                        end)),
+                        (c) -> Postgres.@transaction(c, Distributed.@fetchfrom 1 begin
+                            return :fetchfrom_value
+                        end),
+                        (c) -> Postgres.@transaction(c, Distributed.@fetch begin
+                            return :fetch_value
+                        end),
+                    )
+                        val = distributed_case(conn)
+                        @test val in (:spawnat_value, :fetchfrom_value, :fetch_value)
+                        @test !Postgres.in_transaction(conn)
+                    end
+
+                    # A return in a flattened-iterator expression is legal
+                    # plain Julia: it returns from the generated per-element
+                    # closure, so its value becomes the inner iterator (an Int
+                    # yields itself once) and the enclosing function continues.
+                    # Wrapped in @transaction the behavior must be identical,
+                    # and the block commits on normal completion.
+                    plain_flatten = function()
+                        vals = [x for x in 1:2 for y in (return x)]
+                        (:reached, vals)
+                    end
+                    wrapped_flatten = function(c)
+                        vals = Postgres.@transaction c begin
+                            DBInterface.execute(c, "INSERT INTO macro_test (value) VALUES (21)")
+                            [x for x in 1:2 for y in (return x)]
+                        end
+                        (:reached, vals)
+                    end
+                    @test plain_flatten() == (:reached, [1, 2])
+                    @test wrapped_flatten(conn) == plain_flatten()
+                    @test !Postgres.in_transaction(conn)
+                    @test only(Tables.rowtable(DBInterface.execute(conn, "SELECT count(*)::int AS n FROM macro_test WHERE value = 21"))).n == 1
+
+                    # A commit that fails in the finally (break out of a nested
+                    # level whose savepoint was aborted by a swallowed server
+                    # error) must roll back ITS level before propagating, so
+                    # every enclosing level can unwind its own — nothing may be
+                    # left open, client- or server-side.
+                    nested_break_err = try
+                        for _ in 1:1
+                            Postgres.@transaction conn begin
+                                DBInterface.execute(conn, "INSERT INTO macro_test (value) VALUES (22)")
+                                Postgres.@transaction conn begin
+                                    try
+                                        DBInterface.execute(conn, "SELECT 1/0")
+                                    catch
+                                        # swallowed: the savepoint is now aborted
+                                    end
+                                    break
+                                end
+                            end
+                        end
+                        nothing
+                    catch e
+                        e
+                    end
+                    @test nested_break_err isa Postgres.API.Error
+                    @test !Postgres.in_transaction(conn)
+                    @test !conn.server_in_transaction
+                    @test isempty(Tables.rowtable(DBInterface.execute(conn, "SELECT * FROM macro_test WHERE value = 22")))
+
+                    # Recursion re-enters the SAME expansion: an inner frame's
+                    # return exits only that frame, and each frame commits.
+                    recursive_txn = function f(c, n)
+                        Postgres.@transaction c begin
+                            DBInterface.execute(c, raw"INSERT INTO macro_test (value) VALUES ($1)", (100 + n,))
+                            n == 0 && return :bottom
+                            f(c, n - 1)
+                        end
+                    end
+                    @test recursive_txn(conn, 2) === :bottom
+                    @test !Postgres.in_transaction(conn)
+                    @test only(Tables.rowtable(DBInterface.execute(conn, "SELECT count(*)::int AS n FROM macro_test WHERE value IN (100, 101, 102)"))).n == 3
                 end
 
                 @testset "Nested Transactions" begin