From 425f7138ec87674d7cf18d3e795ebd7d1f1c62af Mon Sep 17 00:00:00 2001 From: PingoLee Date: Mon, 5 Oct 2026 11:29:32 -0300 Subject: [PATCH] Add execute_script for running multi-statement SQL DBInterface.execute always uses the extended protocol, so a string holding several statements fails with 42601 "cannot insert multiple commands into a prepared statement". libpq runs the same string with PQexec over the simple-query protocol, which is how migration and schema scripts are usually sent (#23). Postgres.execute_script(conn, sql) sends the string as one Query message and returns the command tag of each statement, discarding rows. The first server error is thrown after draining to ReadyForQuery, so the connection stays usable, and the transaction status is refreshed either way. COPY in a script is rejected: copy-in is aborted with CopyFail, copy-out is drained. The query logger reports it as :execute_script. DBInterface.execute and the internal execute_simple are unchanged. Co-Authored-By: Claude Opus 5.5 --- README.md | 14 +++ docs/src/manual.md | 19 ++++ src/Postgres.jl | 2 +- src/api/API.jl | 67 +++++++++++++ src/api/types.jl | 6 +- src/execute.jl | 51 ++++++++++ test/execute_script_fake_server.jl | 149 +++++++++++++++++++++++++++++ test/runtests.jl | 54 ++++++++++- 8 files changed, 357 insertions(+), 5 deletions(-) create mode 100644 test/execute_script_fake_server.jl diff --git a/README.md b/README.md index 15877e8..06c0f86 100644 --- a/README.md +++ b/README.md @@ -125,6 +125,20 @@ Postgres.clear_statement_cache!(conn) DBInterface.close!(conn) ``` +## SQL scripts + +`DBInterface.execute` runs one statement per call. `Postgres.execute_script` runs a string of +several statements, in one implicit transaction unless the script contains its own `BEGIN`/`COMMIT`, +and returns each command tag. + +```julia +using Postgres, DBInterface +conn = DBInterface.connect(Postgres.Connection, "host=127.0.0.1;user=postgres;password=postgres;dbname=postgres") +Postgres.execute_script(conn, "CREATE TEMP TABLE t (a int); COMMENT ON TABLE t IS 'one; two'") +# ["CREATE TABLE", "COMMENT"] +DBInterface.close!(conn) +``` + ## Transactions ```julia diff --git a/docs/src/manual.md b/docs/src/manual.md index 0124c35..dda02c9 100644 --- a/docs/src/manual.md +++ b/docs/src/manual.md @@ -221,6 +221,25 @@ finally end ``` +## SQL Scripts + +`DBInterface.execute` runs one statement per call. `Postgres.execute_script` +runs a string of several statements, such as a migration or a schema file, in +one round trip and returns each statement's command tag. + +```julia +Postgres.execute_script(conn, """ + CREATE TEMP TABLE script_demo (a int); + ALTER TABLE script_demo ALTER COLUMN a SET NOT NULL; + COMMENT ON TABLE script_demo IS 'one; two'; +""") +# ["CREATE TABLE", "ALTER TABLE", "COMMENT"] +``` + +The statements run in one implicit transaction unless the script contains its +own `BEGIN`/`COMMIT`: if one fails, the earlier ones are rolled back and the +error is thrown. Scripts take no parameters and discard result rows. + ## COPY Use `Postgres.copy_from` and `Postgres.copy_to` for PostgreSQL's COPY protocol. diff --git a/src/Postgres.jl b/src/Postgres.jl index e4a8b81..21059a6 100644 --- a/src/Postgres.jl +++ b/src/Postgres.jl @@ -1507,7 +1507,7 @@ end "Error, Notification, PostgresRange, AbstractPostgresStyle, PostgresStyle, " * "query_logging_enabled, query_logger, notice_callback, notification_callback, parse_dsn, " * "transaction, @transaction, start_transaction, commit, rollback, in_transaction, " * - "cursor, copy_from, copy_to, listen!, unlisten!, notify!, wait_for_notification, " * + "cursor, execute_script, copy_from, copy_to, listen!, unlisten!, notify!, wait_for_notification, " * "register_type!, register_enum!, register_composite!, register_range!, " * "command_tag, rows_affected, cancel_query!, escape_identifier, escape_literal, " * "get_cached_statements, clear_statement_cache!, set_statement_cache_maxsize!, " * diff --git a/src/api/API.jl b/src/api/API.jl index f4e0332..ca67ed7 100644 --- a/src/api/API.jl +++ b/src/api/API.jl @@ -1345,6 +1345,73 @@ end exec(socket::BufferedConn, query::String, debug::Bool) = exec(PostgresStyle(), socket, query, debug) +# Run a query string that may hold several statements as one simple-query +# message, like libpq's PQexec, and return the command tag of each statement +# that completed. Row data is skipped. After an ErrorResponse the server skips +# the rest of the string; the first error is kept and surfaced once the stream +# reaches ReadyForQuery, so the connection stays usable. +function exec_script(style::S, socket::BufferedConn, query::String, debug::Bool, + tx_status_ref::Base.RefValue{UInt8}, + server_parameters::Dict{String, String}) where {S <: AbstractPostgresStyle} + writemessage(socket, debug, 'Q', query) + tags = String[] + server_error = nothing + copy_in_statement = false + copy_out_statement = false + try + while true + mt, len = readheader(socket, debug) + if mt == UInt8('C') + push!(tags, commandComplete(len, socket)) + elseif mt == UInt8('E') + server_error = something(server_error, errorResponse(len, socket, debug)) + elseif mt == UInt8('Z') + # publish through the Ref before any throw below: the + # ReadyForQuery status is authoritative either way + tx_status_ref[] = read_ready_status(socket, len) + break + elseif mt == UInt8('T') || mt == UInt8('D') || mt == UInt8('I') + # row description / data row / empty query response + skipbytes!(socket, len) + elseif mt == UInt8('G') + # CopyInResponse: the server is waiting on us for data. Abort + # the copy with CopyFail (no Sync in the simple protocol) so the + # stream returns to ready instead of deadlocking; a clear client + # error is thrown below. + skipbytes!(socket, len) + copy_in_statement = true + writemessage(socket, debug, 'f', "COPY FROM STDIN is not supported via execute_script") + elseif mt == UInt8('H') || mt == UInt8('d') || mt == UInt8('c') + # CopyOutResponse/CopyData/CopyDone: drain the copy-out stream + # through ReadyForQuery; a clear client error is thrown below + mt == UInt8('H') && (copy_out_statement = true) + skipbytes!(socket, len) + elseif mt == UInt8('N') + notice_callback(style, noticeResponse(len, socket)) + elseif mt == UInt8('A') + notification_callback(style, notificationResponse(len, socket)) + elseif mt == UInt8('S') + parameterStatus!(server_parameters, len, socket) + else + close_and_throw(socket, Error("unexpected message type '$(Char(mt))' from server; connection protocol state is corrupted")) + end + end + catch + close(socket) + server_error === nothing || throw(server_error) + rethrow() + end + # same precedence as applyeach(::Exec): for copy-in the server error is + # just the CopyFail artifact, so the client error wins; for copy-out a + # server error is a genuine failure and is more informative than the + # misuse error. A rejected copy-out was still run, with the rest of the + # script: the simple protocol has no way to abort it. + copy_in_statement && throw(PostgresInterfaceError("COPY ... FROM STDIN is not supported via execute_script; use Postgres.copy_from")) + server_error === nothing || throw(server_error) + copy_out_statement && throw(PostgresInterfaceError("COPY ... TO STDOUT is not supported via execute_script; use Postgres.copy_to")) + return tags +end + function copy_in(style::S, socket, query::String, source::IO, debug::Bool) where {S <: AbstractPostgresStyle} writemessage(socket, debug, 'Q', query) error_msg = nothing diff --git a/src/api/types.jl b/src/api/types.jl index 8341800..0e60113 100644 --- a/src/api/types.jl +++ b/src/api/types.jl @@ -43,9 +43,9 @@ query_logging_enabled(::AbstractPostgresStyle) = false Called after each query when [`query_logging_enabled`](@ref Postgres.API.query_logging_enabled) is true for -`style`. `event` is `:execute`, `:copy_from`, or `:copy_to`; `info` carries -`sql`, `duration_ns`, `success`, the bound `params` (for `:execute`), and -`error` when the query failed. +`style`. `event` is `:execute`, `:execute_script`, `:copy_from`, or +`:copy_to`; `info` carries `sql`, `duration_ns`, `success`, the bound `params` +(for `:execute`), and `error` when the query failed. `info.params` holds the query's parameter values, so a logger that records them will record whatever sensitive data those queries carry. diff --git a/src/execute.jl b/src/execute.jl index 6668b71..5d2a4c9 100644 --- a/src/execute.jl +++ b/src/execute.jl @@ -560,6 +560,57 @@ function DBInterface.execute(conn::Connection, sql::AbstractString, params=nothi end end +""" + Postgres.execute_script(conn, sql) -> Vector{String} + +Run `sql`, which may hold several `;`-separated statements, as one simple-query +message, like libpq's `PQexec`. Returns the command tag of each statement, e.g. +`["CREATE TABLE", "ALTER TABLE"]`; result rows are discarded, so use +`DBInterface.execute` for queries. + +The statements run in one implicit transaction unless `sql` contains its own +transaction control, so commands that refuse to run inside a transaction block +(`VACUUM`, `CREATE INDEX CONCURRENTLY`, ...) must be sent alone. On an error +PostgreSQL skips the remaining statements and the error is thrown; the +connection stays usable. The simple-query protocol has no parameters: use +[`escape_literal`](@ref Postgres.escape_literal) or `DBInterface.execute` for +values. + +COPY belongs to [`copy_from`](@ref Postgres.copy_from) and +[`copy_to`](@ref Postgres.copy_to). A `COPY ... FROM STDIN` in a script is +aborted, which fails the script like any other error (inside an explicit +`BEGIN` the transaction is left failed until `ROLLBACK`). A `COPY ... TO +STDOUT` is rejected only after the rest of the script has run: its data is +discarded and, unless a later statement fails, the other statements' effects +stand. +""" +function execute_script(conn::Connection, sql::AbstractString; debug::Bool=false) + sql_str = String(sql) + style = conn.style + log_enabled = API.query_logging_enabled(style) + start_ns = log_enabled ? time_ns() : 0 + tags = String[] + try + @lock conn.lock begin + checkconn(conn) + status_ref = Ref{UInt8}(UInt8('I')) + try + tags = API.exec_script(style, conn.socket, sql_str, debug || conn.debug, + status_ref, conn.server_parameters) + finally + # as in execute_simple: a failed statement still drained to + # ReadyForQuery and its status is authoritative + conn.server_in_transaction = API.in_transaction_status(status_ref[]) + end + end + log_enabled && query_log_safely(style, :execute_script, (sql=sql_str, duration_ns=time_ns() - start_ns, success=true)) + return tags + catch err + log_enabled && query_log_safely(style, :execute_script, (sql=sql_str, duration_ns=time_ns() - start_ns, success=false, error=err)) + rethrow() + end +end + # Only native values in ordinary owned columns may be converted ahead of the # server. Custom indexing/conversion can have observable per-row side effects. function batch_parameter_type(::Type{T}) where {T} diff --git a/test/execute_script_fake_server.jl b/test/execute_script_fake_server.jl new file mode 100644 index 0000000..aaab928 --- /dev/null +++ b/test/execute_script_fake_server.jl @@ -0,0 +1,149 @@ +# execute_script against a scripted server: simple-query responses a live +# PostgreSQL produces only for particular statements (COPY in either direction, +# CopyBoth) or that need a precise interleaving (notices and rows between +# command tags, an error after earlier statements completed). Reuses the +# fake-server helpers from gssapi.jl and isvalid_fake_server.jl. + +struct ScriptLogStyle <: Postgres.API.AbstractPostgresStyle + events::Vector{Any} +end +Postgres.API.query_logging_enabled(::ScriptLogStyle) = true +Postgres.API.query_logger(s::ScriptLogStyle, event::Symbol, info::NamedTuple) = (push!(s.events, (event, info)); nothing) + +tag_msg(tag) = pgmsg('C', vcat(Vector{UInt8}(tag), 0x00)) +ready_msg(status::Char) = pgmsg('Z', [UInt8(status)]) +# CopyInResponse / CopyOutResponse / CopyBothResponse: text format, no columns +copy_response(code::Char) = pgmsg(code, UInt8[0x00, 0x00, 0x00]) + +# read the client's simple-query message and check its text +function read_query(sock, sql) + code, body = read_message(sock) + code == 'Q' || error("fake server: expected Query, got '$code'") + String(body[1:end-1]) == sql || error("fake server: unexpected query text") + return +end + +function test_execute_script_fake_server() +@testset "execute_script against a scripted server" begin + @testset "returns each command tag; rows skipped, async messages dispatched" begin + sql = "CREATE TABLE t (a int); COMMENT ON TABLE t IS 'one; two'; SELECT 1" + style = RecordingStyle() + with_fake_server(serve_idle(sock -> begin + read_query(sock, sql) + send(sock, vcat(tag_msg("CREATE TABLE"), notice_msg("NOTICE", "00000", "script notice"), + param_status("TimeZone", "UTC"), notification_msg(7, "chan", "payload"), + tag_msg("COMMENT"), pgmsg('T', UInt8[0x00, 0x00]), pgmsg('D', UInt8[0x00, 0x00]), + tag_msg("SELECT 1"), ready_msg('I'))) + drain(sock) + end)) do port, accepted + conn = fake_connection(port, style) + @test Postgres.execute_script(conn, sql) == ["CREATE TABLE", "COMMENT", "SELECT 1"] + @test length(style.notices) == 1 + @test any(==("script notice"), values(style.notices[1])) + @test Postgres.get_server_parameter(conn, "TimeZone") == "UTC" + @test length(style.notifications) == 1 + @test style.notifications[1].channel == "chan" + @test isopen(conn) + @test !conn.server_in_transaction + close(conn) + end + end + + @testset "an empty script returns no tags" begin + with_fake_server(serve_idle(sock -> begin + read_query(sock, "") + send(sock, vcat(pgmsg('I', UInt8[]), ready_msg('I'))) + drain(sock) + end)) do port, accepted + conn = fake_connection(port, Postgres.PostgresStyle()) + @test Postgres.execute_script(conn, "") == String[] + close(conn) + end + end + + @testset "a server error is thrown after draining; transaction status tracked" begin + style = ScriptLogStyle(Any[]) + with_fake_server(serve_idle(sock -> begin + read_query(sock, "BEGIN; SELECT 1; SELECT * FROM missing") + send(sock, vcat(tag_msg("BEGIN"), tag_msg("SELECT 1"), + error_response("ERROR", "42P01", "relation \"missing\" does not exist"), ready_msg('E'))) + read_query(sock, "ROLLBACK") + send(sock, vcat(tag_msg("ROLLBACK"), ready_msg('I'))) + drain(sock) + end)) do port, accepted + conn = fake_connection(port, style) + err = connect_err(() -> Postgres.execute_script(conn, "BEGIN; SELECT 1; SELECT * FROM missing")) + @test err isa Postgres.API.Error + @test err.code == "42P01" + @test isopen(conn) + @test conn.server_in_transaction + @test Postgres.execute_script(conn, "ROLLBACK") == ["ROLLBACK"] + @test !conn.server_in_transaction + @test [(event, info.success) for (event, info) in style.events] == + [(:execute_script, false), (:execute_script, true)] + @test style.events[1][2].error === err + @test style.events[2][2].sql == "ROLLBACK" + @test accepted[] == 1 + close(conn) + end + end + + @testset "COPY FROM STDIN is aborted with CopyFail; connection stays usable" begin + with_fake_server(serve_idle(sock -> begin + read_query(sock, "SELECT 1; COPY t FROM STDIN") + send(sock, vcat(tag_msg("SELECT 1"), copy_response('G'))) + code, _ = read_message(sock) + code == 'f' || error("fake server: expected CopyFail, got '$code'") + send(sock, vcat(error_response("ERROR", "57014", "COPY from stdin failed"), ready_msg('I'))) + read_query(sock, "SELECT 1") + send(sock, vcat(tag_msg("SELECT 1"), ready_msg('I'))) + drain(sock) + end)) do port, accepted + conn = fake_connection(port, Postgres.PostgresStyle()) + err = connect_err(() -> Postgres.execute_script(conn, "SELECT 1; COPY t FROM STDIN")) + @test err isa Postgres.PostgresInterfaceError + @test occursin("copy_from", errtext(err)) + @test isopen(conn) + @test Postgres.execute_script(conn, "SELECT 1") == ["SELECT 1"] + close(conn) + end + end + + @testset "COPY TO STDOUT is drained and rejected; a server error wins" begin + with_fake_server(serve_idle(sock -> begin + read_query(sock, "COPY t TO STDOUT; SELECT 1") + send(sock, vcat(copy_response('H'), pgmsg('d', Vector{UInt8}("1\n")), pgmsg('c', UInt8[]), + tag_msg("COPY 1"), tag_msg("SELECT 1"), ready_msg('I'))) + read_query(sock, "COPY (SELECT 1/0) TO STDOUT") + send(sock, vcat(copy_response('H'), + error_response("ERROR", "22012", "division by zero"), ready_msg('I'))) + drain(sock) + end)) do port, accepted + conn = fake_connection(port, Postgres.PostgresStyle()) + err = connect_err(() -> Postgres.execute_script(conn, "COPY t TO STDOUT; SELECT 1")) + @test err isa Postgres.PostgresInterfaceError + @test occursin("copy_to", errtext(err)) + @test isopen(conn) + err = connect_err(() -> Postgres.execute_script(conn, "COPY (SELECT 1/0) TO STDOUT")) + @test err isa Postgres.API.Error + @test err.code == "22012" + @test isopen(conn) + close(conn) + end + end + + @testset "CopyBothResponse closes the connection" begin + with_fake_server(serve_idle(sock -> begin + read_query(sock, "START_REPLICATION") + send(sock, copy_response('W')) + drain(sock) + end)) do port, accepted + conn = fake_connection(port, Postgres.PostgresStyle()) + err = connect_err(() -> Postgres.execute_script(conn, "START_REPLICATION")) + @test err isa Postgres.API.Error + @test occursin("unexpected message type 'W'", errtext(err)) + @test !isopen(conn) + end + end +end +end diff --git a/test/runtests.jl b/test/runtests.jl index 71a4f5a..3904799 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -498,6 +498,7 @@ include("result_decoding.jl") include("gssapi.jl") include("isvalid_fake_server.jl") include("execute_fake_server.jl") +include("execute_script_fake_server.jl") include("executemany.jl") include("notification_deadlines.jl") @@ -515,7 +516,7 @@ include("notification_deadlines.jl") :Error, :Notification, :PostgresRange, :AbstractPostgresStyle, :PostgresStyle, :query_logging_enabled, :query_logger, :notice_callback, :notification_callback, :parse_dsn, :transaction, Symbol("@transaction"), :start_transaction, :commit, :rollback, :in_transaction, - :cursor, :copy_from, :copy_to, :listen!, :unlisten!, :notify!, :wait_for_notification, + :cursor, :execute_script, :copy_from, :copy_to, :listen!, :unlisten!, :notify!, :wait_for_notification, :register_type!, :register_enum!, :register_composite!, :register_range!, :command_tag, :rows_affected, :cancel_query!, :escape_identifier, :escape_literal, :get_cached_statements, :clear_statement_cache!, :set_statement_cache_maxsize!, @@ -1036,6 +1037,7 @@ include("notification_deadlines.jl") test_isvalid_fake_server() test_isvalid_fragmentation() test_execute_fake_server() + test_execute_script_fake_server() test_executemany_protocol() test_notification_deadlines() test_notification_tls_deadlines() @@ -2536,6 +2538,56 @@ include("notification_deadlines.jl") @test Tables.rowtable(DBInterface.execute(conn, "SELECT 4 AS a"))[1].a == 4 end + @testset "Execute Script" begin + # several statements in one string, with a ';' inside a + # literal and a dollar-quoted body; rows are discarded + tags = Postgres.execute_script(conn, """ + CREATE TEMP TABLE script_test (a int); + ALTER TABLE script_test ALTER COLUMN a SET NOT NULL; + COMMENT ON TABLE script_test IS 'one; two'; + CREATE FUNCTION pg_temp.script_fn() RETURNS int AS \$\$ SELECT 1; \$\$ LANGUAGE sql; + INSERT INTO script_test VALUES (1), (2); + SELECT * FROM script_test; + """) + @test tags == ["CREATE TABLE", "ALTER TABLE", "COMMENT", "CREATE FUNCTION", "INSERT 0 2", "SELECT 2"] + @test Postgres.execute_script(conn, "") == String[] + + # a failing statement skips the rest and rolls back the + # implicit transaction; the connection stays usable + err = try + Postgres.execute_script(conn, "CREATE TEMP TABLE script_rollback (a int); SELECT 1/0; CREATE TEMP TABLE script_never (a int)") + nothing + catch e + e + end + @test err isa Postgres.API.Error + @test err.code == "22012" + @test !(@lock conn.lock conn.server_in_transaction) + @test Tables.rowtable(DBInterface.execute(conn, "SELECT to_regclass('script_rollback') IS NULL AS gone"))[1].gone + + # commands that refuse a transaction block fail inside a + # multi-statement string, as with libpq's PQexec + err = try + Postgres.execute_script(conn, "SELECT 1; VACUUM script_test") + nothing + catch e + e + end + @test err isa Postgres.API.Error + @test err.code == "25001" + + # transaction control inside a script is tracked + @test Postgres.execute_script(conn, "BEGIN; SELECT 1") == ["BEGIN", "SELECT 1"] + @test @lock conn.lock conn.server_in_transaction + @test Postgres.execute_script(conn, "COMMIT") == ["COMMIT"] + @test !(@lock conn.lock conn.server_in_transaction) + + # COPY is rejected without desyncing the connection + @test_throws Postgres.PostgresInterfaceError Postgres.execute_script(conn, "SELECT 1; COPY script_test FROM STDIN") + @test_throws Postgres.PostgresInterfaceError Postgres.execute_script(conn, "COPY script_test TO STDOUT") + @test Tables.rowtable(DBInterface.execute(conn, "SELECT 6 AS a"))[1].a == 6 + end + @testset "Cursor Streaming" begin cur = Postgres.cursor(conn, "SELECT generate_series(1, 5) AS n"; fetchsize=2) values = [row.n for row in cur]