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]