Add execute_script for running multi-statement SQL

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

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.

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!, " *

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

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.

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}

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