diff --git a/Project.toml b/Project.toml index dce5a9c..09cd63e 100644 --- a/Project.toml +++ b/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"] diff --git a/src/Postgres.jl b/src/Postgres.jl index cbba5ba..aaab198 100644 --- a/src/Postgres.jl +++ b/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 diff --git a/test/runtests.jl b/test/runtests.jl index 7f05410..239dd79 100644 --- a/test/runtests.jl +++ b/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