diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 5a67cf2..c4eeb0b 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -15,6 +15,7 @@ jobs: test: name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ github.event_name }} runs-on: ${{ matrix.os }} + timeout-minutes: 20 strategy: fail-fast: false matrix: @@ -49,6 +50,7 @@ jobs: auth-modes: name: Auth ${{ matrix.auth }} - Julia 1 - ubuntu-latest runs-on: ubuntu-latest + timeout-minutes: 20 strategy: fail-fast: false matrix: diff --git a/Project.toml b/Project.toml index 2d29a60..165d80d 100644 --- a/Project.toml +++ b/Project.toml @@ -34,8 +34,9 @@ julia = "1.10" [extras] Harbor = "af79dbb9-1a80-47ad-8928-192a4af69376" +JuliaC = "acedd4c2-ced6-4a15-accc-2607eb759ba2" Sockets = "6462fe0b-24de-5631-8697-dd941f90decc" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Harbor", "Sockets", "Test"] +test = ["Harbor", "JuliaC", "Sockets", "Test"] diff --git a/docs/src/manual.md b/docs/src/manual.md index 62cb9a2..a0d010f 100644 --- a/docs/src/manual.md +++ b/docs/src/manual.md @@ -108,6 +108,20 @@ profiles = DBInterface.execute(conn, """ The `postgres=(name=:column_name,)` tag is only needed when a column should map to a differently named field. Columns such as `id` or `name` can be left untagged because they already match the Julia field name. +### Driver Styles + +Connection behavior such as query logging, server notices, and asynchronous +notifications is selected by a concrete driver style. Subtype +`Postgres.AbstractPostgresStyle`, overload the documented behavior hooks for +that style, and pass an instance with the `style` connection keyword. The +default `Postgres.PostgresStyle` keeps query logging disabled and reports +server notices through Julia's logger. + +```@docs +Postgres.AbstractPostgresStyle +Postgres.PostgresStyle +``` + ## Parameters And Prepared Statements Use PostgreSQL placeholders (`$1`, `$2`, ...) and pass a tuple or other iterable of parameter values. diff --git a/src/Postgres.jl b/src/Postgres.jl index d025314..0c1944d 100644 --- a/src/Postgres.jl +++ b/src/Postgres.jl @@ -16,10 +16,9 @@ include("connection_string.jl") using .ConnectionString const Pools = ConcurrentUtilities.Pools -const NOOP_QUERY_LOGGER = (event, info) -> nothing const ReseauConn = Union{Reseau.TCP.Conn, Reseau.TLS.Conn} -mutable struct Connection{T} <: DBInterface.Connection +mutable struct Connection{T, S <: API.AbstractPostgresStyle} <: DBInterface.Connection const lock::ReentrantLock socket::ReseauConn const host::String @@ -47,14 +46,15 @@ mutable struct Connection{T} <: DBInterface.Connection closed::Bool # if explicitly closed by user; guarded by lock const reconnect::Bool debug::Bool - notice_callback::Function # callback for NOTICE messages - notification_callback::Function # callback for NOTIFY messages - query_logger::Function # callback for query/copy events + # style-dispatched behavior (query_logger / notice_callback / notification_callback + # overloads on a custom AbstractPostgresStyle) — Function-typed callback fields are + # dynamic dispatch at every use, unresolvable under `juliac --trim` + const style::S in_transaction::Bool # track transaction state transaction_depth::Int # track nested transactions (SAVEPOINTs) generation::Int # increment on reconnect to invalidate statements - function Connection(; host::AbstractString="", user::AbstractString="", password::Union{AbstractString, Nothing}=nothing, dbname::AbstractString="", port::Integer=5432, debug::Bool=false, reconnect::Bool=false, application_name::Union{AbstractString, Nothing}=nothing, connect_timeout::Union{Integer, Nothing}=nothing, sslmode::Union{AbstractString, Nothing}=nothing, sslrootcert::Union{AbstractString, Nothing}=nothing, sslcert::Union{AbstractString, Nothing}=nothing, sslkey::Union{AbstractString, Nothing}=nothing, sslcapath::Union{AbstractString, Nothing}=nothing, statement_timeout::Union{Integer, Nothing}=nothing, statement_cache_maxsize::Integer=100, sslservername::Union{AbstractString, Nothing}=nothing) + function Connection(; host::AbstractString="", user::AbstractString="", password::Union{AbstractString, Nothing}=nothing, dbname::AbstractString="", port::Integer=5432, debug::Bool=false, reconnect::Bool=false, application_name::Union{AbstractString, Nothing}=nothing, connect_timeout::Union{Integer, Nothing}=nothing, sslmode::Union{AbstractString, Nothing}=nothing, sslrootcert::Union{AbstractString, Nothing}=nothing, sslcert::Union{AbstractString, Nothing}=nothing, sslkey::Union{AbstractString, Nothing}=nothing, sslcapath::Union{AbstractString, Nothing}=nothing, sslservername::Union{AbstractString, Nothing}=nothing, statement_timeout::Union{Integer, Nothing}=nothing, statement_cache_maxsize::Integer=100, style::API.AbstractPostgresStyle=PostgresStyle()) host = String(host) user = String(user) dbname = String(dbname) @@ -72,19 +72,14 @@ mutable struct Connection{T} <: DBInterface.Connection maxsize = max(0, Int(statement_cache_maxsize)) #TODO: if values have spaces, need to single-quote them # also need to escape single quotes/backslahes then with backslashes - socket, pid, skey, server_params = API.connect(host, port, dbname, user, password, debug, app_name, timeout, sslmode_val, sslrootcert_val, sslcert_val, sslkey_val, sslcapath_val, statement_timeout_val; sslservername = sslservername_val) + socket, pid, skey, server_params = API.connect(host, port, dbname, user, password, debug, app_name, timeout, sslmode_val, sslrootcert_val, sslcert_val, sslkey_val, sslcapath_val, sslservername_val, statement_timeout_val) registry = Dict(API.DEFAULT_TYPE_REGISTRY) - default_notice_callback = notice -> begin - msg = get(notice, "M", "") - !isempty(msg) && @warn msg - return - end - default_notification_callback = notification -> nothing - default_query_logger = NOOP_QUERY_LOGGER - return new{Statement}(ReentrantLock(), socket, host, user, password, dbname, port, app_name, timeout, sslmode_val, sslrootcert_val, sslcert_val, sslkey_val, sslcapath_val, sslservername_val, statement_timeout_val, pid, skey, Dict{String, Statement}(), maxsize, 0, server_params, registry, false, reconnect, debug, default_notice_callback, default_notification_callback, default_query_logger, false, 0, 1) + return new{Statement{typeof(style)}, typeof(style)}(ReentrantLock(), socket, host, user, password, dbname, port, app_name, timeout, sslmode_val, sslrootcert_val, sslcert_val, sslkey_val, sslcapath_val, sslservername_val, statement_timeout_val, pid, skey, Dict{String, Statement{typeof(style)}}(), maxsize, 0, server_params, registry, false, reconnect, debug, style, false, 0, 1) end end +_style_type(::Connection{T, S}) where {T, S} = S + Base.isopen(conn::Connection) = @lock conn.lock isopen(conn.socket) function Base.show(io::IO, conn::Connection) @@ -126,11 +121,6 @@ function evict_lru_statement!(conn::Connection) return end -function log_query(logger::Function, event::Symbol, info::NamedTuple) - logger(event, info) - return -end - function get_cached_statements(conn::Connection) @lock conn.lock copy(conn.statements) end @@ -167,32 +157,9 @@ get_server_parameter(conn::Connection, param::String) = @lock conn.lock get(conn get_server_parameters(conn::Connection) = @lock conn.lock copy(conn.server_parameters) -set_notice_callback!(conn::Connection, f::Function) = @lock conn.lock begin - conn.notice_callback = f - return conn -end - -function get_notice_callback(conn::Connection) - return @lock conn.lock conn.notice_callback -end - -set_notification_callback!(conn::Connection, f::Function) = @lock conn.lock begin - conn.notification_callback = f - return conn -end - -function get_notification_callback(conn::Connection) - return @lock conn.lock conn.notification_callback -end - -set_query_logger!(conn::Connection, f::Function) = @lock conn.lock begin - conn.query_logger = f - return conn -end - -function get_query_logger(conn::Connection) - return @lock conn.lock conn.query_logger -end +# NOTE: runtime callback setters are gone — customize behavior by passing a custom +# AbstractPostgresStyle to Connection(; style=...) and overloading the style-first +# interface methods (query_logger / notice_callback / notification_callback). function get_statement_timeout(conn::Connection) return @lock conn.lock conn.statement_timeout @@ -279,11 +246,11 @@ function wait_for_notification(conn::Connection; timeout::Union{Real, Nothing}=n mt, len = API.readheader(conn.socket, conn.debug) if mt == UInt8('A') notification = API.notificationResponse(len, conn.socket) - conn.notification_callback(notification) + API.notification_callback(conn.style, notification) return notification elseif mt == UInt8('N') notice = API.noticeResponse(len, conn.socket) - conn.notice_callback(notice) + API.notice_callback(conn.style, notice) elseif mt == UInt8('S') buf = read(conn.socket, len) update_server_parameters!(conn, buf) @@ -303,18 +270,17 @@ function wait_for_notification(conn::Connection; timeout::Union{Real, Nothing}=n end function copy_from(conn::Connection, sql::AbstractString, data::IO; debug::Bool=false) - logger = conn.query_logger - log_enabled = logger !== NOOP_QUERY_LOGGER + log_enabled = API.query_logging_enabled(conn.style) start_ns = log_enabled ? time_ns() : 0 sql_str = String(sql) try @lock conn.lock begin checkconn(conn) - API.copy_in(conn.socket, sql_str, data, debug || conn.debug, conn.notice_callback, conn.notification_callback) + API.copy_in(conn.style, conn.socket, sql_str, data, debug || conn.debug) end - log_enabled && log_query(logger, :copy_from, (sql=sql_str, duration_ns=time_ns() - start_ns, success=true)) + log_enabled && API.query_logger(conn.style, :copy_from, (sql=sql_str, duration_ns=time_ns() - start_ns, success=true)) catch err - log_enabled && log_query(logger, :copy_from, (sql=sql_str, duration_ns=time_ns() - start_ns, success=false, error=err)) + log_enabled && API.query_logger(conn.style, :copy_from, (sql=sql_str, duration_ns=time_ns() - start_ns, success=false, error=err)) rethrow() end return conn @@ -326,18 +292,17 @@ function copy_from(conn::Connection, sql::AbstractString, data; debug::Bool=fals end function copy_to(conn::Connection, sql::AbstractString, dest::IO; debug::Bool=false) - logger = conn.query_logger - log_enabled = logger !== NOOP_QUERY_LOGGER + log_enabled = API.query_logging_enabled(conn.style) start_ns = log_enabled ? time_ns() : 0 sql_str = String(sql) try @lock conn.lock begin checkconn(conn) - API.copy_out(conn.socket, sql_str, dest, debug || conn.debug, conn.notice_callback, conn.notification_callback) + API.copy_out(conn.style, conn.socket, sql_str, dest, debug || conn.debug) end - log_enabled && log_query(logger, :copy_to, (sql=sql_str, duration_ns=time_ns() - start_ns, success=true)) + log_enabled && API.query_logger(conn.style, :copy_to, (sql=sql_str, duration_ns=time_ns() - start_ns, success=true)) catch err - log_enabled && log_query(logger, :copy_to, (sql=sql_str, duration_ns=time_ns() - start_ns, success=false, error=err)) + log_enabled && API.query_logger(conn.style, :copy_to, (sql=sql_str, duration_ns=time_ns() - start_ns, success=false, error=err)) rethrow() end return dest @@ -446,7 +411,7 @@ function checkconn(conn::Connection) # connection is closed, but not explicitly, reconnect conn.in_transaction && throw(PostgresInterfaceError("postgres connection has been closed or disconnected; reconnect disabled during transaction")) conn.reconnect || throw(PostgresInterfaceError("postgres connection has been closed or disconnected; reconnect disabled")) - conn.socket, conn.pid, conn.skey, server_params = API.connect(conn.host, conn.port, conn.dbname, conn.user, conn.password, conn.debug, conn.application_name, conn.connect_timeout, conn.sslmode, conn.sslrootcert, conn.sslcert, conn.sslkey, conn.sslcapath, conn.statement_timeout; sslservername = conn.sslservername) + conn.socket, conn.pid, conn.skey, server_params = API.connect(conn.host, conn.port, conn.dbname, conn.user, conn.password, conn.debug, conn.application_name, conn.connect_timeout, conn.sslmode, conn.sslrootcert, conn.sslcert, conn.sslkey, conn.sslcapath, conn.sslservername, conn.statement_timeout) empty!(conn.statements) conn.in_transaction = false conn.transaction_depth = 0 @@ -458,8 +423,8 @@ function checkconn(conn::Connection) return end -function DBInterface.connect(::Type{Connection}, host::AbstractString, user::AbstractString, passwd::Union{AbstractString, Nothing}; dbname::AbstractString="", port::Integer=5432, debug::Bool=false, reconnect::Bool=false, application_name::Union{AbstractString, Nothing}=nothing, connect_timeout::Union{Integer, Nothing}=nothing, sslmode::Union{AbstractString, Nothing}=nothing, sslrootcert::Union{AbstractString, Nothing}=nothing, sslcert::Union{AbstractString, Nothing}=nothing, sslkey::Union{AbstractString, Nothing}=nothing, sslcapath::Union{AbstractString, Nothing}=nothing, statement_timeout::Union{Integer, Nothing}=nothing, statement_cache_maxsize::Integer=100, sslservername::Union{AbstractString, Nothing}=nothing) - Connection(host=host, user=user, password=passwd, dbname=dbname, port=port, debug=debug, reconnect=reconnect, application_name=application_name, connect_timeout=connect_timeout, sslmode=sslmode, sslrootcert=sslrootcert, sslcert=sslcert, sslkey=sslkey, sslcapath=sslcapath, statement_timeout=statement_timeout, statement_cache_maxsize=statement_cache_maxsize, sslservername=sslservername) +function DBInterface.connect(::Type{Connection}, host::AbstractString, user::AbstractString, passwd::Union{AbstractString, Nothing}; dbname::AbstractString="", port::Integer=5432, debug::Bool=false, reconnect::Bool=false, application_name::Union{AbstractString, Nothing}=nothing, connect_timeout::Union{Integer, Nothing}=nothing, sslmode::Union{AbstractString, Nothing}=nothing, sslrootcert::Union{AbstractString, Nothing}=nothing, sslcert::Union{AbstractString, Nothing}=nothing, sslkey::Union{AbstractString, Nothing}=nothing, sslcapath::Union{AbstractString, Nothing}=nothing, sslservername::Union{AbstractString, Nothing}=nothing, statement_timeout::Union{Integer, Nothing}=nothing, statement_cache_maxsize::Integer=100, style::API.AbstractPostgresStyle=PostgresStyle()) + Connection(host=host, user=user, password=passwd, dbname=dbname, port=port, debug=debug, reconnect=reconnect, application_name=application_name, connect_timeout=connect_timeout, sslmode=sslmode, sslrootcert=sslrootcert, sslcert=sslcert, sslkey=sslkey, sslcapath=sslcapath, sslservername=sslservername, statement_timeout=statement_timeout, statement_cache_maxsize=statement_cache_maxsize, style=style) end function DBInterface.connect(::Type{Connection}, dsn::String; debug::Bool=false, reconnect::Bool=false, statement_cache_maxsize::Union{Integer, Nothing}=nothing) @@ -503,8 +468,8 @@ function ConnectionPool(connector::Function; limit::Integer=10) return ConnectionPool(pool, connector) end -function ConnectionPool(::Type{Connection}, host::AbstractString, user::AbstractString, passwd::Union{AbstractString, Nothing}; dbname::AbstractString="", port::Integer=5432, debug::Bool=false, reconnect::Bool=false, application_name::Union{AbstractString, Nothing}=nothing, connect_timeout::Union{Integer, Nothing}=nothing, sslmode::Union{AbstractString, Nothing}=nothing, sslrootcert::Union{AbstractString, Nothing}=nothing, sslcert::Union{AbstractString, Nothing}=nothing, sslkey::Union{AbstractString, Nothing}=nothing, sslcapath::Union{AbstractString, Nothing}=nothing, statement_timeout::Union{Integer, Nothing}=nothing, statement_cache_maxsize::Integer=100, limit::Integer=10) - connector = () -> DBInterface.connect(Connection, host, user, passwd; dbname=dbname, port=port, debug=debug, reconnect=reconnect, application_name=application_name, connect_timeout=connect_timeout, sslmode=sslmode, sslrootcert=sslrootcert, sslcert=sslcert, sslkey=sslkey, sslcapath=sslcapath, statement_timeout=statement_timeout, statement_cache_maxsize=statement_cache_maxsize) +function ConnectionPool(::Type{Connection}, host::AbstractString, user::AbstractString, passwd::Union{AbstractString, Nothing}; dbname::AbstractString="", port::Integer=5432, debug::Bool=false, reconnect::Bool=false, application_name::Union{AbstractString, Nothing}=nothing, connect_timeout::Union{Integer, Nothing}=nothing, sslmode::Union{AbstractString, Nothing}=nothing, sslrootcert::Union{AbstractString, Nothing}=nothing, sslcert::Union{AbstractString, Nothing}=nothing, sslkey::Union{AbstractString, Nothing}=nothing, sslcapath::Union{AbstractString, Nothing}=nothing, statement_timeout::Union{Integer, Nothing}=nothing, statement_cache_maxsize::Integer=100, limit::Integer=10, style::API.AbstractPostgresStyle=PostgresStyle()) + connector = () -> DBInterface.connect(Connection, host, user, passwd; dbname=dbname, port=port, debug=debug, reconnect=reconnect, application_name=application_name, connect_timeout=connect_timeout, sslmode=sslmode, sslrootcert=sslrootcert, sslcert=sslcert, sslkey=sslkey, sslcapath=sslcapath, statement_timeout=statement_timeout, statement_cache_maxsize=statement_cache_maxsize, style=style) return ConnectionPool(connector; limit=limit) end @@ -574,7 +539,7 @@ include("execute.jl") # does not exist"). Also one network round trip instead of three. Callers # must hold conn.lock. function execute_simple(conn::Connection, sql::String) - API.exec(conn.socket, sql, conn.debug) + API.exec(conn.style, conn.socket, sql, conn.debug) return conn end @@ -632,7 +597,7 @@ function rollback(conn::Connection) return conn end -function transaction(f::Function, conn::Connection) +function transaction(f::F, conn::Connection) where {F} start_transaction(conn) try result = f(conn) @@ -644,7 +609,7 @@ function transaction(f::Function, conn::Connection) end end -function DBInterface.transaction(f::Function, conn::Connection) +function DBInterface.transaction(f::F, conn::Connection) where {F} start_transaction(conn) try result = f() diff --git a/src/api/API.jl b/src/api/API.jl index 0847a2c..9a07b5d 100644 --- a/src/api/API.jl +++ b/src/api/API.jl @@ -1,8 +1,9 @@ module API using UUIDs, Dates, Reseau, SASLAuth, MD5, Parsers, StructUtils, Logging, JSON, Random +import ..PostgresInterfaceError -export PostgresStyle, Error, Notification, Numeric, PostgresRange +export PostgresStyle, AbstractPostgresStyle, query_logging_enabled, query_logger, notice_callback, notification_callback, Error, Notification, Numeric, PostgresRange const ReseauConn = Union{Reseau.TCP.Conn, Reseau.TLS.Conn} const SKIP_BUFFER_SIZE = 8192 @@ -179,6 +180,9 @@ msgsizeof(x) = sizeof(x) msgsizeof(x::Tuple{String, String}) = sizeof(x[1]) + 1 + sizeof(x[2]) + 1 msgsizeof(x::Params) = sum(4 + (ismissing(p) ? 0 : sizeof(p)) for p in x.params; init=0) +_msgsizeof_parts(::Tuple{}) = 0 +_msgsizeof_parts(parts::Tuple) = msgsizeof(first(parts)) + _msgsizeof_parts(Base.tail(parts)) + writepart(io, x) = write(io, x) function writepart(io, x::String) write(io, x) @@ -203,36 +207,91 @@ function writepart(io, x::Params) end end -function writemessage(socket, debug, code::Char, parts...) +_writeparts(io, ::Tuple{}) = nothing +function _writeparts(io, parts::Tuple) + writepart(io, first(parts)) + _writeparts(io, Base.tail(parts)) + return nothing +end + +function _write_message_to_buffer(buf::IOBuffer, debug::Bool, msg::Tuple) + code = first(msg)::Char + parts = Base.tail(msg) + debug && @info "sending message: $code, $parts" + len = Int32(4 + _msgsizeof_parts(parts)) + code != '\0' && write(buf, UInt8(code)) + write(buf, hton(len)) + _writeparts(buf, parts) + return nothing +end + +function _writemessage_parts(socket, debug::Bool, code::Char, parts::Tuple)::Nothing debug && @info "sending message: $code, $parts" - len = Int32(4 + sum(msgsizeof(x) for x in parts; init=0)) + len = Int32(4 + _msgsizeof_parts(parts)) buf = IOBuffer(Vector{UInt8}(undef, len + 1); write=true) code != '\0' && write(buf, UInt8(code)) write(buf, hton(len)) - for part in parts - writepart(buf, part) - end + _writeparts(buf, parts) write(socket, take!(buf)) flush(socket) - return + return nothing +end + +writemessage(socket, debug::Bool, code::Char) = _writemessage_parts(socket, debug, code, ()) +writemessage(socket, debug::Bool, code::Char, part1) = _writemessage_parts(socket, debug, code, (part1,)) +writemessage(socket, debug::Bool, code::Char, part1, part2) = _writemessage_parts(socket, debug, code, (part1, part2)) +writemessage(socket, debug::Bool, code::Char, part1, part2, part3) = _writemessage_parts(socket, debug, code, (part1, part2, part3)) +writemessage(socket, debug::Bool, code::Char, parts...) = _writemessage_parts(socket, debug, code, parts) + +function _write_startup_param(buf::IOBuffer, key::String, value::String)::Nothing + writepart(buf, (key, value)) + return nothing end -function writemessages(socket, debug, msgs...) +function writestartupmessage( + socket, + debug::Bool, + user::String, + dbname::String, + application_name::Union{Nothing, String}, + statement_timeout::Union{Nothing, Int}, +)::Nothing + timeout_options = statement_timeout === nothing ? nothing : string("-c statement_timeout=", statement_timeout) + len = 8 + msgsizeof(("user", user)) + msgsizeof(("database", dbname)) + 1 + application_name !== nothing && (len += msgsizeof(("application_name", application_name))) + timeout_options !== nothing && (len += msgsizeof(("options", timeout_options))) + debug && @info "sending startup message" + buf = IOBuffer(Vector{UInt8}(undef, len); write=true) + write(buf, hton(Int32(len))) + write(buf, hton(Int32(196608))) + _write_startup_param(buf, "user", user) + _write_startup_param(buf, "database", dbname) + application_name !== nothing && _write_startup_param(buf, "application_name", application_name) + timeout_options !== nothing && _write_startup_param(buf, "options", timeout_options) + write(buf, UInt8(0)) + write(socket, take!(buf)) + flush(socket) + return nothing +end + +function writemessages(socket, debug::Bool, msgs::Vararg{Tuple, N}) where {N} buf = IOBuffer() - for (code, parts...) in msgs - debug && @info "sending message: $code, $parts" - len = Int32(4 + sum(msgsizeof(x) for x in parts; init=0)) - code != '\0' && write(buf, UInt8(code)) - write(buf, hton(len)) - for part in parts - writepart(buf, part) - end + for msg in msgs + _write_message_to_buffer(buf, debug, msg) end write(socket, take!(buf)) flush(socket) return end +_sum_codes(::Tuple{}) = 0 +_sum_codes(codes::Tuple) = UInt8(first(codes)) + _sum_codes(Base.tail(codes)) + +_contains_code(::UInt8, ::Tuple{}) = false +function _contains_code(mt::UInt8, codes::Tuple) + return mt == UInt8(first(codes)) || _contains_code(mt, Base.tail(codes)) +end + function skipbytes!(io::IO, n::Integer) remaining = Int(n) remaining <= 0 && return nothing @@ -309,10 +368,10 @@ function expect_auth_message(socket, debug, mt, len) end # wait for code, then ready -function waitfor(socket, debug, codes...) +function waitfor(socket, debug::Bool, codes::Vararg{Char, N}) where {N} error = false error_msg = nothing - found = sum(UInt8, codes) + found = _sum_codes(codes) pid = skey = Int32(0) server_params = Dict{String, String}() debug && @info "waitfor: $codes" @@ -342,19 +401,17 @@ function waitfor(socket, debug, codes...) server_params[key] = val i = j + 1 end - elseif Char(mt) in codes + elseif _contains_code(mt, codes) # found found -= mt if mt == UInt8('K') pid = ntoh(read(socket, Int32)) skey = ntoh(read(socket, Int32)) else + # read off message skipbytes!(socket, len) end found == 0 && break - else - # read off message - skipbytes!(socket, len) end end catch @@ -441,7 +498,7 @@ function authRequest(debug, len, socket, user, password, client::Union{Nothing, mechanisms = split(data, '\0'; keepempty=false) if "SCRAM-SHA-256" ∉ mechanisms - close_and_throw(socket, Error("no supported SASL mechanisms: $mechanisms")) + close_and_throw(socket, Error("no supported SASL mechanisms")) end client = SASLAuth.SCRAMSHA256Client(user, password) msg, _ = SASLAuth.step!(client, nothing) @@ -473,7 +530,30 @@ function authRequest(debug, len, socket, user, password, client::Union{Nothing, end end -function connectsocket(host::AbstractString, port::Integer; connect_timeout::Union{Int, Nothing}=nothing) +function authRequest(debug, len, socket, user, ::Nothing, client::Nothing=nothing) + auth_code = ntoh(read(socket, Int32)) + debug && @info "auth code: $auth_code" + if auth_code == 0 + return socket + elseif auth_code == 2 + close_and_throw(socket, Error("kerberos v5 authentication not supported")) + elseif auth_code == 3 || auth_code == 5 || auth_code == 10 || auth_code == 11 || auth_code == 12 + close_and_throw(socket, Error("server requested password authentication but no password was provided")) + elseif auth_code == 7 + close_and_throw(socket, Error("GSSAPI authentication not supported")) + elseif auth_code == 8 + close_and_throw(socket, Error("GSSAPI/SSPI continuation not supported")) + elseif auth_code == 9 + close_and_throw(socket, Error("SSPI authentication not supported")) + else + close_and_throw(socket, Error("unknown authentication code: $auth_code")) + end +end + +connectsocket(host::AbstractString, port::Integer; connect_timeout::Union{Int, Nothing}=nothing) = + connectsocket(host, port, connect_timeout) + +function connectsocket(host::AbstractString, port::Integer, @nospecialize(connect_timeout::Union{Int, Nothing})) address = string(host, ":", Int(port)) return if connect_timeout === nothing Reseau.TCP.connect(address) @@ -493,20 +573,30 @@ function tlsupgrade( ssl_cacert::Union{String, Nothing}=nothing, ssl_capath::Union{String, Nothing}=nothing, ) + return tlsupgrade(socket, connect_timeout, server_name, verify_peer, + ssl_cert, ssl_key, ssl_cacert, ssl_capath) +end + +function tlsupgrade(socket::Reseau.TCP.Conn, @nospecialize(connect_timeout::Union{Int, Nothing}), + @nospecialize(server_name::Union{String, Nothing}), verify_peer::Bool, + @nospecialize(ssl_cert::Union{String, Nothing}), @nospecialize(ssl_key::Union{String, Nothing}), + @nospecialize(ssl_cacert::Union{String, Nothing}), @nospecialize(ssl_capath::Union{String, Nothing})) ca_file = ssl_cacert === nothing ? ssl_capath : ssl_cacert handshake_timeout_ns = connect_timeout === nothing ? Int64(0) : Int64(connect_timeout) * 1_000_000_000 - tls_conn = Reseau.TLS.client( - socket, - Reseau.TLS.Config( - ; - server_name, - verify_peer, - cert_file=ssl_cert, - key_file=ssl_key, - ca_file, - handshake_timeout_ns, - ), - ) + sni = server_name isa String ? server_name : nothing + + # positional Config, split on the cert/key pair: the kwargs form (and >2 + # Union-valued args at once) is unresolvable dynamic dispatch under --trim + config = if ssl_cert === nothing && ssl_key === nothing + Reseau.TLS.Config(sni, verify_peer, verify_peer, Reseau.TLS.ClientAuthMode.NoClientCert, + nothing, nothing, ca_file, nothing, String[], UInt16[], + handshake_timeout_ns, Reseau.TLS.TLS1_2_VERSION, nothing, false) + else + Reseau.TLS.Config(sni, verify_peer, verify_peer, Reseau.TLS.ClientAuthMode.NoClientCert, + ssl_cert::String, ssl_key::String, ca_file, nothing, String[], UInt16[], + handshake_timeout_ns, Reseau.TLS.TLS1_2_VERSION, nothing, false) + end + tls_conn = Reseau.TLS.client(socket, config) try Reseau.TLS.handshake!(tls_conn) return tls_conn @@ -519,9 +609,23 @@ end # sslservername: TLS SNI override for when `host` is a pre-resolved address — # SNI-routed servers (e.g. Neon) need the hostname on the TLS handshake even # when the TCP dial goes to an IP. -function connect(host::String, port::Integer, dbname::String, user::String, password::Union{String, Nothing}, debug::Bool, application_name::Union{String, Nothing}, connect_timeout::Union{Int, Nothing}, sslmode::Union{String, Nothing}, sslrootcert::Union{String, Nothing}, sslcert::Union{String, Nothing}, sslkey::Union{String, Nothing}, sslcapath::Union{String, Nothing}, statement_timeout::Union{Int, Nothing}; sslservername::Union{String, Nothing}=nothing) - socket = connectsocket(host, port; connect_timeout) - sslmode_str = sslmode === nothing ? "prefer" : lowercase(String(sslmode)) +function connect(host::String, port::Integer, dbname::String, user::String, @nospecialize(password::Union{String, Nothing}), debug::Bool, @nospecialize(application_name::Union{String, Nothing}), @nospecialize(connect_timeout::Union{Int, Nothing}), @nospecialize(sslmode::Union{String, Nothing}), @nospecialize(sslrootcert::Union{String, Nothing}), @nospecialize(sslcert::Union{String, Nothing}), @nospecialize(sslkey::Union{String, Nothing}), @nospecialize(sslcapath::Union{String, Nothing}), @nospecialize(sslservername::Union{String, Nothing}), @nospecialize(statement_timeout::Union{Int, Nothing})) + # re-assert the @nospecialize'd params to their declared unions: the asserts give + # inference the (static) union types without re-introducing per-argument + # specialization, so the kwarg NamedTuples below have static types instead of + # runtime apply_type — which `juliac --trim` can't resolve + password_v = password::Union{String, Nothing} + application_name_v = application_name::Union{String, Nothing} + connect_timeout_v = connect_timeout::Union{Int, Nothing} + sslmode_v = sslmode::Union{String, Nothing} + sslrootcert_v = sslrootcert::Union{String, Nothing} + sslcert_v = sslcert::Union{String, Nothing} + sslkey_v = sslkey::Union{String, Nothing} + sslcapath_v = sslcapath::Union{String, Nothing} + sslservername_v = sslservername::Union{String, Nothing} + statement_timeout_v = statement_timeout::Union{Int, Nothing} + socket = connectsocket(host, port, connect_timeout_v) + sslmode_str = sslmode_v === nothing ? "prefer" : lowercase(String(sslmode_v)) sslmode_str == "disable" || sslmode_str == "prefer" || sslmode_str == "require" || sslmode_str == "verify-full" || throw(Error("invalid sslmode: $sslmode_str")) if sslmode_str != "disable" # send SSLRequest @@ -529,16 +633,10 @@ function connect(host::String, port::Integer, dbname::String, user::String, pass mt = read(socket, UInt8) if mt == UInt8('S') # upgrade socket to tls and do handshake - socket = tlsupgrade( - socket; - connect_timeout, - server_name=something(sslservername, host), - verify_peer=sslmode_str == "verify-full", - ssl_cert=sslcert, - ssl_key=sslkey, - ssl_cacert=sslrootcert, - ssl_capath=sslcapath, - ) + socket = tlsupgrade(socket, connect_timeout_v, + sslservername_v isa String ? sslservername_v : host, + sslmode_str == "verify-full", + sslcert_v, sslkey_v, sslrootcert_v, sslcapath_v) elseif mt == UInt8('N') (sslmode_str == "require" || sslmode_str == "verify-full") && throw(Error("server does not support SSL")) elseif mt == UInt8('E') @@ -549,22 +647,19 @@ function connect(host::String, port::Integer, dbname::String, user::String, pass close_and_throw(socket, Error("unexpected response to SSLRequest: $(Char(mt))")) end end - # Build startup parameters - params = [("user", user), ("database", dbname)] - if !isnothing(application_name) - push!(params, ("application_name", application_name)) - end - if !isnothing(statement_timeout) - push!(params, ("options", "-c statement_timeout=$(statement_timeout)")) + # socket-union isa split (post-TLS-upgrade φ) so the call resolves under --trim + if socket isa Reseau.TCP.Conn + writestartupmessage(socket::Reseau.TCP.Conn, debug, user, dbname, application_name_v, statement_timeout_v) + else + writestartupmessage(socket::Reseau.TLS.Conn, debug, user, dbname, application_name_v, statement_timeout_v) end - writemessage(socket, debug, '\0', Int32(196608), params..., UInt8(0)) # read initial response mt, len = readheader(socket, debug) if mt == UInt8('E') # error close_and_throw_error_response(socket, len, debug) elseif mt == UInt8('R') - authRequest(debug, len, socket, user, password) + authRequest(debug, len, socket, user, password_v) elseif mt == UInt8('v') # server version too old close_and_throw(socket, Error("server version too old")) @@ -640,15 +735,9 @@ struct DataRow type_registry::Dict{Int, TypeInfo} end -struct PostgresStyle <: StructUtils.StructStyle end - -StructUtils.fieldtagkey(::PostgresStyle) = :postgres -StructUtils.structlike(::PostgresStyle, ::Type{<:Number}) = false -StructUtils.structlike(::PostgresStyle, ::Type{<:JSON.LazyValue}) = false -StructUtils.lift(::PostgresStyle, ::Type{T}, x::T) where {T<:JSON.LazyValue} = x, nothing -StructUtils.lift(::PostgresStyle, ::Type{T}, x::T, tags) where {T<:JSON.LazyValue} = x, nothing +# (style types + behavior interface live in types.jl, included before this point) -function StructUtils.applyeach(::PostgresStyle, f, dr::DataRow) +function StructUtils.applyeach(::AbstractPostgresStyle, f, dr::DataRow) buf = dr.buf GC.@preserve buf begin ncols = Int(ntoh(unsafe_load(Ptr{Int16}(pointer(buf))))) @@ -669,14 +758,13 @@ function StructUtils.applyeach(::PostgresStyle, f, dr::DataRow) return end -struct Exec{N, A} +struct Exec{S <: AbstractPostgresStyle} + style::S socket::ReseauConn names::Vector{Symbol} typeIds::Vector{Int} type_registry::Dict{Int, TypeInfo} debug::Bool - notice_callback::N - notification_callback::A command_tag::Base.RefValue{Union{Nothing, String}} rows_affected::Base.RefValue{Union{Nothing, Int}} end @@ -697,7 +785,7 @@ function rows_affected_from_command_tag(tag::String) end end -function StructUtils.applyeach(::PostgresStyle, f, e::Exec) +function StructUtils.applyeach(::AbstractPostgresStyle, f, e::Exec) nrows = 0 server_error = nothing consumer_error = nothing @@ -736,11 +824,11 @@ function StructUtils.applyeach(::PostgresStyle, f, e::Exec) elseif mt == UInt8('A') # notification response notification = notificationResponse(len, e.socket) - e.notification_callback(notification) + notification_callback(e.style, notification) elseif mt == UInt8('N') # notice response notice = noticeResponse(len, e.socket) - e.notice_callback(notice) + notice_callback(e.style, notice) else throw(Error("unexpected message type '$(Char(mt))' from server; connection protocol state is corrupted")) end @@ -762,24 +850,53 @@ function StructUtils.applyeach(::PostgresStyle, f, e::Exec) return end -function exec(socket::ReseauConn, stmtname::String, params::Vector{Union{String, Missing}}, names, typeIds, type_registry::Dict{Int, TypeInfo}, debug::Bool, rowlimit::Int=0, notice_callback::N=(notice)->nothing, notification_callback::A=(notification)->nothing) where {N, A} +function exec(style::S, socket::ReseauConn, stmtname::String, params::Vector{Union{String, Missing}}, names, typeIds, type_registry::Dict{Int, TypeInfo}, debug::Bool, rowlimit::Int=0) where {S <: AbstractPostgresStyle} #TODO: support binary format: here and in applycast npformats = Int16(0) # all params use text format nparams = Int16(length(params)) # bind, then execute, then sync writemessages(socket, debug, ('B', "", stmtname, npformats, nparams, Params(params), Int16(0)), ('E', "", Int32(rowlimit)), ('S',)) waitfor(socket, debug, '2') - return Exec(socket, names, typeIds, type_registry, debug, notice_callback, notification_callback, Ref{Union{Nothing, String}}(nothing), Ref{Union{Nothing, Int}}(nothing)) + return Exec{S}(style, socket, names, typeIds, type_registry, debug, Ref{Union{Nothing, String}}(nothing), Ref{Union{Nothing, Int}}(nothing)) end -function exec(socket, query::String, debug::Bool) +function exec(style::S, socket::ReseauConn, query::String, debug::Bool) where {S <: AbstractPostgresStyle} writemessages(socket, debug, ('Q', query)) - waitfor(socket, debug, 'Z') - #TODO: handle all the various response message types, like applyeach above + describeprepared + server_error = nothing + try + while true + mt, len = readheader(socket, debug) + if mt == UInt8('E') + # Keep draining through ReadyForQuery before surfacing the + # server error so the connection remains reusable. + server_error = errorResponse(len, socket, debug) + elseif mt == UInt8('Z') + skipbytes!(socket, len) + break + elseif mt == UInt8('N') + notice_callback(style, noticeResponse(len, socket)) + elseif mt == UInt8('A') + notification_callback(style, notificationResponse(len, socket)) + elseif mt == UInt8('C') || mt == UInt8('T') || mt == UInt8('D') || + mt == UInt8('I') || mt == UInt8('S') + # CommandComplete and any incidental simple-query result data. + skipbytes!(socket, len) + 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 + server_error === nothing || throw(server_error) return end -function copy_in(socket, query::String, source::IO, debug::Bool, notice_callback::Function, notification_callback::Function) +exec(socket::ReseauConn, query::String, debug::Bool) = exec(PostgresStyle(), socket, query, debug) + +function copy_in(style::S, socket, query::String, source::IO, debug::Bool) where {S <: AbstractPostgresStyle} writemessage(socket, debug, 'Q', query) error_msg = nothing while true @@ -791,10 +908,10 @@ function copy_in(socket, query::String, source::IO, debug::Bool, notice_callback error_msg = errorResponse(len, socket, debug) elseif mt == UInt8('N') notice = noticeResponse(len, socket) - notice_callback(notice) + notice_callback(style, notice) elseif mt == UInt8('A') notification = notificationResponse(len, socket) - notification_callback(notification) + notification_callback(style, notification) else skipbytes!(socket, len) end @@ -816,10 +933,10 @@ function copy_in(socket, query::String, source::IO, debug::Bool, notice_callback skipbytes!(socket, len) elseif mt == UInt8('N') notice = noticeResponse(len, socket) - notice_callback(notice) + notice_callback(style, notice) elseif mt == UInt8('A') notification = notificationResponse(len, socket) - notification_callback(notification) + notification_callback(style, notification) elseif mt == UInt8('Z') skipbytes!(socket, len) break @@ -831,7 +948,7 @@ function copy_in(socket, query::String, source::IO, debug::Bool, notice_callback return end -function copy_out(socket, query::String, dest::IO, debug::Bool, notice_callback::Function, notification_callback::Function) +function copy_out(style::S, socket, query::String, dest::IO, debug::Bool) where {S <: AbstractPostgresStyle} writemessage(socket, debug, 'Q', query) error_msg = nothing while true @@ -848,10 +965,10 @@ function copy_out(socket, query::String, dest::IO, debug::Bool, notice_callback: error_msg = errorResponse(len, socket, debug) elseif mt == UInt8('N') notice = noticeResponse(len, socket) - notice_callback(notice) + notice_callback(style, notice) elseif mt == UInt8('A') notification = notificationResponse(len, socket) - notification_callback(notification) + notification_callback(style, notification) elseif mt == UInt8('Z') skipbytes!(socket, len) break @@ -899,4 +1016,9 @@ export cancel_request include("../array_parsing.jl") using .ArrayParsing +function __init__() + _populate_default_type_registry!() + return nothing +end + end diff --git a/src/api/types.jl b/src/api/types.jl index 00ad1d2..a2cc1c1 100644 --- a/src/api/types.jl +++ b/src/api/types.jl @@ -1,9 +1,50 @@ const JSONType = typeof(JSON.lazy("1")) +""" + AbstractPostgresStyle <: StructUtils.StructStyle + +Style hierarchy for customizing driver behavior, following the StructUtils "style" +pattern. Pass a custom style to `Connection(; style=MyStyle())` and overload the +behavior interface on it: + + Postgres.query_logging_enabled(::MyStyle) = true + Postgres.query_logger(::MyStyle, event::Symbol, info::NamedTuple) = ... + Postgres.notice_callback(::MyStyle, notice) = ... + Postgres.notification_callback(::MyStyle, notification) = ... + +Custom styles inherit the default row-materialization traits (lift/structlike/...), +which dispatch on `AbstractPostgresStyle`, and are used as the StructUtils style when +materializing query results — so `StructUtils.lift` overloads on a custom style apply +to row values too. Static dispatch on the style (rather than `Function`-typed callback +fields) also keeps the driver compilable under `juliac --trim`. +""" +abstract type AbstractPostgresStyle <: StructUtils.StructStyle end + +"The default style: no query logging, NOTICE messages surface as `@warn`." +struct PostgresStyle <: AbstractPostgresStyle end + +# behavior interface (style-first; overload on your own style) +query_logging_enabled(::AbstractPostgresStyle) = false +query_logger(::AbstractPostgresStyle, event::Symbol, info::NamedTuple) = nothing +function notice_callback(::AbstractPostgresStyle, notice) + msg = get(notice, "M", "") + !isempty(msg) && @warn msg + return nothing +end +notification_callback(::AbstractPostgresStyle, notification) = nothing + +StructUtils.fieldtagkey(::AbstractPostgresStyle) = :postgres +StructUtils.structlike(::AbstractPostgresStyle, ::Type{<:Number}) = false +StructUtils.structlike(::AbstractPostgresStyle, ::Type{<:JSON.LazyValue}) = false +StructUtils.lift(::AbstractPostgresStyle, ::Type{T}, x::T) where {T<:JSON.LazyValue} = x, nothing +StructUtils.lift(::AbstractPostgresStyle, ::Type{T}, x::T, tags) where {T<:JSON.LazyValue} = x, nothing + + struct Numeric coeff::BigInt scale::Int end +StructUtils.structlike(::AbstractPostgresStyle, ::Type{Numeric}) = false Base.:(==)(a::Numeric, b::Numeric) = a.coeff == b.coeff && a.scale == b.scale @@ -15,14 +56,73 @@ struct PostgresRange{T} empty::Bool end +# ── CastFn: trim-safe type-erased value caster (the Reseau TaskFn pattern) ── +# A per-callable-type @generated @cfunction whose first C argument is the callable +# (Ref{F}), so the parser call inside the trampoline is concretely dispatched; the +# registry invocation is a ccall through the stored pointer — statically resolvable +# under `juliac --trim`, where a `Function`-typed field call is open-set dynamic +# dispatch. The C signature is primitive-only (the verifier rejects boxed-Any +# cfunction signatures): value/registry/result cross as raw pointers to +# caller-GC.@preserve'd objects. `_root` keeps the callable alive. + +struct _CastCallWrapper <: Function end + +@generated function _cast_gen_fptr(::Type{F}) where F + quote + @cfunction($(_CastCallWrapper()), Cvoid, (Ref{$F}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid})) + end +end + +struct CastFn + ptr::Ptr{Cvoid} # @cfunction pointer (specialized per callable type F) + objptr::Ptr{Cvoid} # pointer to the callable object + _root::Any # GC root — prevents collection, never dispatched on +end + +function CastFn(callable::F) where F + ptr = _cast_gen_fptr(F) + objref = Base.cconvert(Ref{F}, callable) + objptr = Ptr{Cvoid}(Base.unsafe_convert(Ref{F}, objref)) + return CastFn(ptr, objptr, objref) +end + +CastFn(callable::CastFn) = callable + struct TypeInfo julia_type::Type - parser::Union{Function, Nothing} + parser::Union{CastFn, Nothing} +end + +# convenience: registry entries and register_type! keep passing plain functions +TypeInfo(julia_type::Type, parser::Function) = TypeInfo(julia_type, CastFn(parser)) + +function (::_CastCallWrapper)(f::F, valptr::Ptr{Cvoid}, regptr::Ptr{Cvoid}, outptr::Ptr{Cvoid}) where {F} + valref = unsafe_pointer_to_objref(valptr)::Base.RefValue{Any} + out = unsafe_pointer_to_objref(outptr)::Base.RefValue{Any} + registry = unsafe_pointer_to_objref(regptr)::Dict{Int, TypeInfo} + out[] = f(valref[]::String, registry) + return nothing +end + +@inline function (c::CastFn)(val::String, registry::Dict{Int, TypeInfo}) + valref = Ref{Any}(val) + out = Ref{Any}(nothing) + GC.@preserve valref out registry begin + ccall(c.ptr, Cvoid, (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}), c.objptr, pointer_from_objref(valref), pointer_from_objref(registry), pointer_from_objref(out)) + end + return out[] end const IntervalType = Union{Dates.Period, Dates.CompoundPeriod} -const DEFAULT_TYPE_REGISTRY = Dict{Int, TypeInfo}( +# Populated in `_populate_default_type_registry!` (called from API.__init__): the +# entries hold CastFn @cfunction pointers, which must be created at runtime — a +# precompile-time-built const would serialize stale pointers into the image +# (bus error on first use). +const DEFAULT_TYPE_REGISTRY = Dict{Int, TypeInfo}() + +function _populate_default_type_registry!() + merge!(DEFAULT_TYPE_REGISTRY, Dict{Int, TypeInfo}( 23 => TypeInfo(Int32, nothing), 26 => TypeInfo(Cuint, nothing), 20 => TypeInfo(Int64, nothing), @@ -72,7 +172,9 @@ const DEFAULT_TYPE_REGISTRY = Dict{Int, TypeInfo}( 3908 => TypeInfo(PostgresRange{DateTime}, (val, registry) -> parse_range(val, 1114, registry)), 3910 => TypeInfo(PostgresRange{DateTime}, (val, registry) -> parse_range(val, 1184, registry)), 3912 => TypeInfo(PostgresRange{Date}, (val, registry) -> parse_range(val, 1082, registry)), -) + )) + return nothing +end default_type_info(oid::Int) = get(DEFAULT_TYPE_REGISTRY, oid, TypeInfo(String, nothing)) @@ -100,14 +202,82 @@ const DATETIME_OPTIONS = Parsers.Options(dateformat=dateformat"yyyy-mm-dd HH:MM: return sign * (hours * 3600 + mins * 60) end +# ── hand-rolled postgres text-format date/time parsing ────────────────────── +# The wire formats are fixed-layout ("YYYY-MM-DD", "HH:MM:SS[.ffffff]", +# "YYYY-MM-DD HH:MM:SS[.ffffff][±TZ]"), so direct digit extraction is both faster +# than the generic Dates machinery and — decisive under `juliac --trim` — fully +# static: Parsers' dateformat path iterates a type-erased Vector{AbstractDateToken}, +# which is dynamic dispatch per token. Fractional seconds beyond millisecond +# precision are truncated (DateTime/Time storage precision). +@inline _pg_digit(b::UInt8)::Int = Int(b - UInt8('0')) +@inline _pg_isdigit(b::UInt8)::Bool = UInt8('0') <= b <= UInt8('9') + +@inline function _pg_date_at(c, o::Int)::Date + y = _pg_digit(c[o]) * 1000 + _pg_digit(c[o+1]) * 100 + _pg_digit(c[o+2]) * 10 + _pg_digit(c[o+3]) + m = _pg_digit(c[o+5]) * 10 + _pg_digit(c[o+6]) + d = _pg_digit(c[o+8]) * 10 + _pg_digit(c[o+9]) + return Date(y, m, d) +end + +@inline function _pg_hms_at(c, o::Int) + h = _pg_digit(c[o]) * 10 + _pg_digit(c[o+1]) + mi = _pg_digit(c[o+3]) * 10 + _pg_digit(c[o+4]) + se = _pg_digit(c[o+6]) * 10 + _pg_digit(c[o+7]) + ms = 0 + i = o + 8 + if i <= length(c) && c[i] == UInt8('.') + i += 1 + mult = 100 + while i <= length(c) && _pg_isdigit(c[i]) + if mult > 0 + ms += _pg_digit(c[i]) * mult + mult ÷= 10 + end + i += 1 + end + end + return h, mi, se, ms +end + +function pg_parse_date(s::AbstractString)::Date + c = codeunits(s) + length(c) >= 10 || throw(ArgumentError("invalid postgres date")) + return _pg_date_at(c, 1) +end + +function pg_parse_time(s::AbstractString)::Time + c = codeunits(s) + length(c) >= 8 || throw(ArgumentError("invalid postgres time")) + h, mi, se, ms = _pg_hms_at(c, 1) + return Time(h, mi, se, ms) +end + +function pg_parse_datetime(s::AbstractString)::DateTime + c = codeunits(s) + length(c) >= 19 || throw(ArgumentError("invalid postgres timestamp")) + d = _pg_date_at(c, 1) + h, mi, se, ms = _pg_hms_at(c, 12) + return DateTime(Dates.year(d), Dates.month(d), Dates.day(d), h, mi, se, ms) +end + +# timestamptz column into a DateTime field: sniff a trailing offset/Z +function pg_parse_datetime_any(s::AbstractString)::DateTime + isempty(s) && throw(ArgumentError("invalid postgres timestamp")) + ch = s[end] + if ch == 'Z' || (length(s) >= 20 && (any(isequal('+'), SubString(s, 20)) || any(isequal('-'), SubString(s, 20)))) + return parse_timestamptz(String(s)) + end + return pg_parse_datetime(s) +end + @inline function parse_timestamptz(val::String) - lastindex(val) == 0 && return Parsers.parse(DateTime, val, DATETIME_OPTIONS) + lastindex(val) == 0 && throw(ArgumentError("invalid postgres timestamptz")) if val[end] == 'Z' ts = SubString(val, 1, prevind(val, lastindex(val))) - return Parsers.parse(DateTime, ts, DATETIME_OPTIONS) + return pg_parse_datetime(ts) end space_idx = findfirst(isequal(' '), val) - space_idx === nothing && return Parsers.parse(DateTime, val, DATETIME_OPTIONS) + space_idx === nothing && return pg_parse_datetime(val) offset_idx = nothing i = lastindex(val) while i > space_idx @@ -118,9 +288,9 @@ end end i = prevind(val, i) end - offset_idx === nothing && return Parsers.parse(DateTime, val, DATETIME_OPTIONS) + offset_idx === nothing && return pg_parse_datetime(val) ts = SubString(val, 1, prevind(val, offset_idx)) - dt = Parsers.parse(DateTime, ts, DATETIME_OPTIONS) + dt = pg_parse_datetime(ts) offset = SubString(val, offset_idx) seconds = tzoffset_seconds(offset) return dt - Dates.Second(seconds) @@ -185,8 +355,8 @@ function parse_interval_time(token::AbstractString) if occursin('.', seconds_part) whole, frac = split(seconds_part, '.'; limit=2) seconds = parse(Int, whole) - frac = rpad(frac[1:min(end, 3)], 3, '0') - milliseconds = parse(Int, frac) + fs = frac[1:min(end, 3)] + milliseconds = parse(Int, fs) * 10^(3 - length(fs)) else seconds = parse(Int, seconds_part) end @@ -229,7 +399,12 @@ function parse_interval(val::String) end isempty(periods) && return Dates.Millisecond(0) length(periods) == 1 && return only(periods) - return Dates.CompoundPeriod(periods...) + # n=0 construction skips CompoundPeriod's canonicalize loop (whose Period + + # merge is dynamic dispatch under --trim); pg interval text is already + # canonical: unique units, descending order, zero units omitted + cp = Dates.CompoundPeriod(Dates.Period[]) + append!(cp.periods, periods) + return cp end function split_range_values(val::String) @@ -261,21 +436,61 @@ function parse_range_value(token::String, typeId::Int, registry::Dict{Int, TypeI return parse_value(typeId, token, registry) end +# construct over the standard range element types explicitly: PostgresRange{T} +# from a runtime `julia_type` field is a runtime apply_type under --trim +@noinline function _make_range(::Type{T}, @nospecialize(lower), @nospecialize(upper), li::Bool, ui::Bool, empty::Bool) where {T} + l = lower === missing ? missing : (lower::T) + u = upper === missing ? missing : (upper::T) + return PostgresRange{T}(l, u, li, ui, empty) +end + +function _range_typed(@nospecialize(T), @nospecialize(lower), @nospecialize(upper), li::Bool, ui::Bool, empty::Bool) + T === Int32 && return _make_range(Int32, lower, upper, li, ui, empty) + T === Int64 && return _make_range(Int64, lower, upper, li, ui, empty) + T === Float64 && return _make_range(Float64, lower, upper, li, ui, empty) + T === Numeric && return _make_range(Numeric, lower, upper, li, ui, empty) + T === Date && return _make_range(Date, lower, upper, li, ui, empty) + T === DateTime && return _make_range(DateTime, lower, upper, li, ui, empty) + throw(PostgresInterfaceError("no trim-safe range constructor registered for this element type; " * + "register a parser function for the range type")) +end + function parse_range(val::String, typeId::Int, registry::Dict{Int, TypeInfo}) - lowercase(val) == "empty" && return PostgresRange{type_info(registry, typeId).julia_type}(missing, missing, false, false, true) + T = type_info(registry, typeId).julia_type + lowercase(val) == "empty" && return _range_typed(T, missing, missing, false, false, true) lower_inclusive = val[1] == '[' upper_inclusive = val[end] == ']' inner = val[2:end - 1] left, right = split_range_values(inner) lower = parse_range_value(left, typeId, registry) upper = parse_range_value(right, typeId, registry) - T = type_info(registry, typeId).julia_type - return PostgresRange{T}(lower, upper, lower_inclusive, upper_inclusive, false) + return _range_typed(T, lower, upper, lower_inclusive, upper_inclusive, false) end parse_array_scalar(typeId::Int, registry::Dict{Int, TypeInfo}, value::Missing) = missing -parse_array_scalar(typeId::Int, registry::Dict{Int, TypeInfo}, value::AbstractVector) = [parse_array_scalar(typeId, registry, v) for v in value] parse_array_scalar(typeId::Int, registry::Dict{Int, TypeInfo}, value::AbstractString) = parse_value(typeId, String(value), registry) +# explicit two-level walk: the self-recursive AbstractVector method is an +# unresolved invoke under --trim; deeper nesting throws (register a parser +# function for exotic multidimensional array types) +function parse_array_scalar(typeId::Int, registry::Dict{Int, TypeInfo}, value::AbstractVector) + return Any[v === missing ? missing : + v isa String ? parse_value(typeId, v, registry) : + v isa SubString{String} ? parse_value(typeId, String(v), registry) : + v isa Vector{Any} ? _parse_array_level2(typeId, registry, v) : + v isa Vector{String} ? _parse_array_level2(typeId, registry, v) : + v isa Vector{Union{Missing, String}} ? _parse_array_level2(typeId, registry, v) : + v + for v in value] +end +function _parse_array_level2(typeId::Int, registry::Dict{Int, TypeInfo}, + value::Union{Vector{Any}, Vector{String}, Vector{Union{Missing, String}}}) + return Any[v === missing ? missing : + v isa String ? parse_value(typeId, v, registry) : + v isa SubString{String} ? parse_value(typeId, String(v), registry) : + v isa AbstractVector ? throw(PostgresInterfaceError("arrays nested deeper than two dimensions are not supported on the untyped parse path")) : + v + for v in value] +end function parse_array_by_oid(val::String, typeId::Int, registry::Dict{Int, TypeInfo}) parsed = parse_array(val, String) @@ -422,13 +637,31 @@ function parse_value(typeId::Int, val::String, registry::Dict{Int, TypeInfo}) if typeId == 1184 return parse_timestamptz(val) end - return Parsers.parse(T, val, DATETIME_OPTIONS) + return pg_parse_datetime(val) elseif T == UUID return UUID(val) elseif T == Numeric return parse_numeric(val) + elseif T == Int16 + return Parsers.parse(Int16, val) + elseif T == Int32 + return Parsers.parse(Int32, val) + elseif T == Int64 + return Parsers.parse(Int64, val) + elseif T == Cuint + return Parsers.parse(Cuint, val) + elseif T == Float32 + return Parsers.parse(Float32, val) + elseif T == Float64 + return Parsers.parse(Float64, val) + elseif T == Date + return pg_parse_date(val) + elseif T == Time + return pg_parse_time(val) elseif T <: Number || T <: Dates.TimeType - return Parsers.parse(T, val) + # custom numeric/time registrations must supply a parser function: a + # runtime-Type Parsers.parse here is unresolvable under --trim + throw(PostgresInterfaceError("no parser registered for this type; pass a parser function when registering the type")) elseif T == Vector{UInt8} return decode_bytea(val) elseif T == JSONType @@ -458,3 +691,50 @@ end f(name, parse_value(typeId, val, registry)) return end + +# ── typed-struct materialization: parse by the declared field type ─────────── +@static if isdefined(StructUtils, :InterpClosure) && isdefined(StructUtils, :HotStructClosure) + # Typed targets know each field type, so let their PostgresStyle lift parse + # the wire string directly. Untyped destinations keep the OID parser above. + @inline function applycast(f::Union{StructUtils.InterpClosure, StructUtils.HotStructClosure}, name, typeId, val::String, registry::Dict{Int, TypeInfo}) + f(name, val) + return + end +end + +StructUtils.lift(::AbstractPostgresStyle, ::Type{Int8}, s::String) = Parsers.parse(Int8, s), nothing +StructUtils.lift(::AbstractPostgresStyle, ::Type{Bool}, s::String) = (s == "t" || s == "1"), nothing +StructUtils.lift(::AbstractPostgresStyle, ::Type{Char}, s::String) = s[1], nothing +StructUtils.lift(::AbstractPostgresStyle, ::Type{Int16}, s::String) = Parsers.parse(Int16, s), nothing +StructUtils.lift(::AbstractPostgresStyle, ::Type{Int32}, s::String) = Parsers.parse(Int32, s), nothing +StructUtils.lift(::AbstractPostgresStyle, ::Type{Int64}, s::String) = Parsers.parse(Int64, s), nothing +StructUtils.lift(::AbstractPostgresStyle, ::Type{Cuint}, s::String) = Parsers.parse(Cuint, s), nothing +StructUtils.lift(::AbstractPostgresStyle, ::Type{Float32}, s::String) = Parsers.parse(Float32, s), nothing +StructUtils.lift(::AbstractPostgresStyle, ::Type{Float64}, s::String) = Parsers.parse(Float64, s), nothing +StructUtils.lift(::AbstractPostgresStyle, ::Type{Date}, s::String) = pg_parse_date(s), nothing +StructUtils.lift(::AbstractPostgresStyle, ::Type{Time}, s::String) = pg_parse_time(s), nothing +StructUtils.lift(::AbstractPostgresStyle, ::Type{DateTime}, s::String) = pg_parse_datetime_any(s), nothing +StructUtils.lift(::AbstractPostgresStyle, ::Type{UUID}, s::String) = UUID(s), nothing +StructUtils.lift(::AbstractPostgresStyle, ::Type{Numeric}, s::String) = parse_numeric(s), nothing +StructUtils.lift(::AbstractPostgresStyle, ::Type{IntervalType}, s::String) = parse_interval(s), nothing +StructUtils.lift(::AbstractPostgresStyle, ::Type{Vector{UInt8}}, s::String) = decode_bytea(s), nothing +StructUtils.lift(::AbstractPostgresStyle, ::Type{JSONType}, s::String) = JSON.lazy(s), nothing +StructUtils.lift(::AbstractPostgresStyle, ::Type{Vector{String}}, s::String) = parse_array(s, String), nothing +StructUtils.lift(::AbstractPostgresStyle, ::Type{Vector{Int16}}, s::String) = parse_array(s, Int16), nothing +StructUtils.lift(::AbstractPostgresStyle, ::Type{Vector{Int32}}, s::String) = parse_array(s, Int32), nothing +StructUtils.lift(::AbstractPostgresStyle, ::Type{Vector{Int64}}, s::String) = parse_array(s, Int64), nothing +StructUtils.lift(::AbstractPostgresStyle, ::Type{Vector{Float32}}, s::String) = parse_array(s, Float32), nothing +StructUtils.lift(::AbstractPostgresStyle, ::Type{Vector{Float64}}, s::String) = parse_array(s, Float64), nothing +StructUtils.lift(::AbstractPostgresStyle, ::Type{Vector{Bool}}, s::String) = parse_array(s, Bool), nothing +StructUtils.lift(::AbstractPostgresStyle, ::Type{Vector{Date}}, s::String) = parse_array(s, Date), nothing +StructUtils.lift(::AbstractPostgresStyle, ::Type{Vector{Time}}, s::String) = parse_array(s, Time), nothing +StructUtils.lift(::AbstractPostgresStyle, ::Type{Vector{DateTime}}, s::String) = parse_array(s, DateTime), nothing +StructUtils.lift(::AbstractPostgresStyle, ::Type{Vector{UUID}}, s::String) = parse_array(s, UUID), nothing +StructUtils.lift(::AbstractPostgresStyle, ::Type{Vector{Numeric}}, s::String) = parse_array(s, Numeric), nothing + +# For array-typed fields the generic `make` takes its arraylike branch (applyeach +# over the source) before consulting lifts — but our source is the wire STRING, +# which must be parsed, not iterated. Route (String source → AbstractVector target) +# through the lifts above. +StructUtils.make(st::AbstractPostgresStyle, ::Type{T}, s::String) where {T <: AbstractVector} = StructUtils.lift(st, T, s) +StructUtils.make(st::AbstractPostgresStyle, ::Type{T}, s::String, tags) where {T <: AbstractVector} = StructUtils.lift(st, T, s) diff --git a/src/array_parsing.jl b/src/array_parsing.jl index 164fd43..d7b5d67 100644 --- a/src/array_parsing.jl +++ b/src/array_parsing.jl @@ -1,6 +1,7 @@ module ArrayParsing -using Parsers +using Parsers, Dates, UUIDs +import ..pg_parse_date, ..pg_parse_time, ..pg_parse_datetime_any, ..parse_numeric, ..Numeric const BRACKET_OPEN = UInt8('[') const BRACKET_CLOSE = UInt8(']') @@ -75,7 +76,7 @@ function parse_bool_token(token::String) throw(ArgumentError("invalid boolean token: $token")) end -function parse_scalar(token::String, inner_type::DataType, quoted::Bool) +function parse_scalar(token::String, inner_type::Type{T}, quoted::Bool) where {T} !quoted && token == NULL_STR && return missing inner_type === String && return token inner_type === Bool && return parse_bool_token(token) @@ -84,10 +85,15 @@ function parse_scalar(token::String, inner_type::DataType, quoted::Bool) inner_type === Int64 && return Parsers.parse(Int64, token) inner_type === Float32 && return Parsers.parse(Float32, token) inner_type === Float64 && return Parsers.parse(Float64, token) + inner_type === Date && return pg_parse_date(token) + inner_type === Time && return pg_parse_time(token) + inner_type === DateTime && return pg_parse_datetime_any(token) + inner_type === UUID && return UUID(token) + inner_type === Numeric && return parse_numeric(token) return token end -function coerce_array(values::Vector{Any}, inner_type::DataType) +function coerce_array(values::Vector{Any}, inner_type::Type{T}) where {T} isempty(values) && return inner_type[] has_nested = false has_missing = false @@ -119,7 +125,7 @@ function coerce_array(values::Vector{Any}, inner_type::DataType) return dest end -function parse_array_value(code::Base.CodeUnits{UInt8, String}, pos::Ref{Int}, inner_type::DataType) +function parse_array_value(code::Base.CodeUnits{UInt8, String}, pos::Ref{Int}, inner_type::Type{T}) where {T} c = code[pos[]] if c == BRACE_OPEN || c == BRACKET_OPEN return parse_array_inner(code, pos, inner_type) @@ -132,7 +138,7 @@ function parse_array_value(code::Base.CodeUnits{UInt8, String}, pos::Ref{Int}, i end end -function parse_array_inner(code::Base.CodeUnits{UInt8, String}, pos::Ref{Int}, inner_type::DataType) +function parse_array_inner(code::Base.CodeUnits{UInt8, String}, pos::Ref{Int}, inner_type::Type{T}) where {T} values = Any[] pos[] += 1 while pos[] <= length(code) @@ -153,7 +159,7 @@ function parse_array_inner(code::Base.CodeUnits{UInt8, String}, pos::Ref{Int}, i return coerce_array(values, inner_type) end -function parse_array(str::String, inner_type::DataType) +function parse_array(str::String, inner_type::Type{T}) where {T} code = codeunits(str) pos = Ref{Int}(1) skip_ws(code, pos) diff --git a/src/execute.jl b/src/execute.jl index f2e090c..93f9715 100644 --- a/src/execute.jl +++ b/src/execute.jl @@ -42,8 +42,11 @@ function DBInterface.close!(::Result) return end -mutable struct Statement <: DBInterface.Statement - const conn::Connection +# parametric on the connection's style so `conn` stays a concrete type — with the +# 2-parameter Connection, a bare `Connection{Statement}` field is a UnionAll, making +# every stmt.conn access (and everything downstream) dynamic under `juliac --trim` +mutable struct Statement{S <: API.AbstractPostgresStyle} <: DBInterface.Statement + const conn::Connection{Statement{S}, S} name::String const sql::String const nfields::Int @@ -60,8 +63,9 @@ end DBInterface.getconnection(stmt::Statement) = stmt.conn -mutable struct Cursor - const conn::Connection +# parametric on the connection's style, same rationale as Statement{S} +mutable struct Cursor{S <: API.AbstractPostgresStyle} + const conn::Connection{Statement{S}, S} const portal::String const names::Vector{Symbol} const typeIds::Vector{Int} @@ -75,9 +79,9 @@ mutable struct Cursor owns_transaction::Bool end -Base.IteratorSize(::Type{Cursor}) = Base.SizeUnknown() -Base.IteratorEltype(::Type{Cursor}) = Base.HasEltype() -Base.eltype(::Type{Cursor}) = ResultRow +Base.IteratorSize(::Type{<:Cursor}) = Base.SizeUnknown() +Base.IteratorEltype(::Type{<:Cursor}) = Base.HasEltype() +Base.eltype(::Type{<:Cursor}) = ResultRow function Base.show(io::IO, stmt::Statement) println(io, "Postgres.Statement:") @@ -118,7 +122,7 @@ function DBInterface.prepare(conn::Connection, sql::AbstractString; debug::Bool= nparams, names, types = API.describeprepared(conn.socket, name, debug) params = Union{String, Missing}[missing for _ = 1:nparams] last_used = next_statement_clock!(conn) - return Statement(conn, name, sql_str, length(names), names, types, nparams, params, false, false, conn.generation, last_used) + return Statement{_style_type(conn)}(conn, name, sql_str, length(names), names, types, nparams, params, false, false, conn.generation, last_used) end # evict if at max size while length(conn.statements) >= conn.statement_cache_maxsize @@ -129,7 +133,7 @@ function DBInterface.prepare(conn::Connection, sql::AbstractString; debug::Bool= nparams, names, types = API.describeprepared(conn.socket, name, debug) params = Union{String, Missing}[missing for _ = 1:nparams] last_used = next_statement_clock!(conn) - stmt = Statement(conn, name, sql_str, length(names), names, types, nparams, params, false, true, conn.generation, last_used) + stmt = Statement{_style_type(conn)}(conn, name, sql_str, length(names), names, types, nparams, params, false, true, conn.generation, last_used) conn.statements[sql_str] = stmt return stmt end @@ -143,7 +147,6 @@ function DBInterface.close!(stmt::Statement) stmt.closed = true return end - checkconn(stmt.conn) stmt.cached && haskey(stmt.conn.statements, stmt.sql) && stmt.conn.statements[stmt.sql] === stmt && delete!(stmt.conn.statements, stmt.sql) API.close_statement(stmt.conn.socket, stmt.name, stmt.conn.debug) stmt.closed = true @@ -166,22 +169,49 @@ function DBInterface.close!(cursor::Cursor) end Base.close(cursor::Cursor) = DBInterface.close!(cursor) -_param(x::AbstractString) = String(x) -_param(x) = string(x) +_param(x::AbstractString)::String = String(x) +_param(x)::String = string(x) _param(x::Missing) = x _param(::Nothing) = missing -_param(x::AbstractVector{UInt8}) = string("\\x", bytes2hex(x)) +_param(x::AbstractVector{UInt8})::String = string("\\x", bytes2hex(x)) # convert to postgres array literal syntax: { x, y, z } # strings must be double-quoted and double quotes and backslashes escaped # missing values are NULL -_aparam(x::AbstractString) = string("\"", replace(x, r"([\"\\])" => s"\\\1"), "\"") -_aparam(::Missing) = "NULL" -_aparam(::Nothing) = "NULL" -_aparam(x) = _param(x) -_param(x::AbstractVector) = string("{", join([_aparam(y) for y in x], ", "), "}") +_aparam(x::AbstractString)::String = string("\"", replace(x, r"([\"\\])" => s"\\\1"), "\"") +_aparam(::Missing)::String = "NULL" +_aparam(::Nothing)::String = "NULL" +_aparam(x)::String = _param(x) +function _param(x::AbstractVector)::String + io = IOBuffer() + write(io, '{') + first_item = true + for y in x + if first_item + first_item = false + else + write(io, ", ") + end + write(io, _aparam(y)) + end + write(io, '}') + return String(take!(io)) +end @noinline param_mismatch(sql, nparams, n) = throw(PostgresInterfaceError("number of parameters provided ($n) does not match number of placeholders ($nparams) in sql: $sql")) +@generated function bind_tuple_params!(dest::Vector{Union{String, Missing}}, params::Tuple{Vararg{Any, N}}) where {N} + assigns = [:(dest[$i] = _param(params[$i])) for i in 1:N] + return Expr(:block, assigns..., :(return nothing)) +end + +function bind_params!(dest::Vector{Union{String, Missing}}, params::Tuple, sql::AbstractString) + nparams = length(params) + nparams > length(dest) && param_mismatch(sql, length(dest), nparams) + nparams == length(dest) || param_mismatch(sql, length(dest), nparams) + bind_tuple_params!(dest, params) + return +end + function bind_params!(dest::Vector{Union{String, Missing}}, params, sql::AbstractString) nparams = 0 if params !== nothing @@ -207,13 +237,16 @@ mutable struct RowClosure i::Int end -@inline function (f::RowClosure)(k, v) +# @nospecialize(v): parse_value's return is Any by nature (OID-driven); the value +# lands in a Vector{Any}, so one instance suffices and the applycast call site +# stays statically resolvable under --trim +@inline function (f::RowClosure)(k, @nospecialize(v)) if v === nothing # translate nothing -> missing for Tables.jl @inbounds f.types[f.i] = Union{f.types[f.i], Missing} @inbounds f.data[f.i] = missing else - if v isa AbstractVector && Missing <: eltype(v) + if v isa AbstractVector{>:Missing} @inbounds f.types[f.i] = typeof(v) end @inbounds f.data[f.i] = v @@ -227,9 +260,9 @@ function makeresult(e::API.Exec) types = Type[API.juliatype(x -> x, i, e.type_registry) for i in typeIds] lookup = Dict(x => i for (i, x) in enumerate(names)) rows = ResultRow[] - StructUtils.applyeach(PostgresStyle(), e) do i, row + StructUtils.applyeach(e.style, e) do i, row data = Vector{Any}(undef, length(names)) - StructUtils.applyeach(PostgresStyle(), RowClosure(data, types, 1), row) + StructUtils.applyeach(e.style, RowClosure(data, types, 1), row) push!(rows, ResultRow(data, names, types, lookup, i)) end return Result(names, types, rows, e.command_tag[], e.rows_affected[]) @@ -250,7 +283,7 @@ function read_portal_batch!(cursor::Cursor) row = API.DataRow(read(conn.socket, len), cursor.names, cursor.typeIds, conn.type_registry) try data = Vector{Any}(undef, length(cursor.names)) - StructUtils.applyeach(PostgresStyle(), RowClosure(data, cursor.types, 1), row) + StructUtils.applyeach(conn.style, RowClosure(data, cursor.types, 1), row) push!(rows, ResultRow(data, cursor.names, cursor.types, cursor.lookup, cursor.rowcount)) catch err # value conversion failed; keep reading through @@ -268,10 +301,10 @@ function read_portal_batch!(cursor::Cursor) done = true elseif mt == UInt8('N') notice = API.noticeResponse(len, conn.socket) - conn.notice_callback(notice) + API.notice_callback(conn.style, notice) elseif mt == UInt8('A') notification = API.notificationResponse(len, conn.socket) - conn.notification_callback(notification) + API.notification_callback(conn.style, notification) elseif mt == UInt8('E') error_msg = API.errorResponse(len, conn.socket, conn.debug) elseif mt == UInt8('Z') @@ -317,8 +350,8 @@ function Base.iterate(cursor::Cursor, state=nothing) end function DBInterface.execute(stmt::Statement, params=nothing, ::Type{T}=Any; debug::Bool=false, binary::Bool=false) where {T} - logger = stmt.conn.query_logger - log_enabled = logger !== NOOP_QUERY_LOGGER + style = stmt.conn.style + log_enabled = API.query_logging_enabled(style) start_ns = log_enabled ? time_ns() : 0 result = nothing try @@ -326,21 +359,29 @@ function DBInterface.execute(stmt::Statement, params=nothing, ::Type{T}=Any; deb # check that connection/statement are ok checkstmt(stmt) bind_params!(stmt.params, params, stmt.sql) - e = API.exec(stmt.conn.socket, stmt.name, stmt.params, stmt.names, stmt.typeIds, stmt.conn.type_registry, debug, 0, stmt.conn.notice_callback, stmt.conn.notification_callback) - result = T === Any ? makeresult(e) : StructUtils.arraylike(T) ? StructUtils.make(T, e, PostgresStyle()) : only(StructUtils.make(Vector{T}, e, PostgresStyle())) + # isa-split the socket union with per-branch typeasserts (identical calls in + # both branches get tail-merged back into one dynamic call by the optimizer), + # so the exec call resolves statically under `juliac --trim` + socket = stmt.conn.socket + e = if socket isa Reseau.TCP.Conn + API.exec(style, socket::Reseau.TCP.Conn, stmt.name, stmt.params, stmt.names, stmt.typeIds, stmt.conn.type_registry, debug, 0) + else + API.exec(style, socket::Reseau.TLS.Conn, stmt.name, stmt.params, stmt.names, stmt.typeIds, stmt.conn.type_registry, debug, 0) + end + result = T === Any ? makeresult(e) : StructUtils.arraylike(T) ? StructUtils.make(T, e, style) : only(StructUtils.make(Vector{T}, e, style)) end - log_enabled && log_query(logger, :execute, (sql=stmt.sql, params=params, duration_ns=time_ns() - start_ns, success=true)) + log_enabled && API.query_logger(style, :execute, (sql=stmt.sql, params=params, duration_ns=time_ns() - start_ns, success=true)) return result catch err - log_enabled && log_query(logger, :execute, (sql=stmt.sql, params=params, duration_ns=time_ns() - start_ns, success=false, error=err)) + log_enabled && API.query_logger(style, :execute, (sql=stmt.sql, params=params, duration_ns=time_ns() - start_ns, success=false, error=err)) rethrow() end end function DBInterface.execute(conn::Connection, sql::AbstractString, params=nothing, ::Type{T}=Any; debug::Bool=false) where {T} sql_str = String(sql) - logger = conn.query_logger - log_enabled = logger !== NOOP_QUERY_LOGGER + style = conn.style + log_enabled = API.query_logging_enabled(style) start_ns = log_enabled ? time_ns() : 0 result = nothing try @@ -349,13 +390,19 @@ function DBInterface.execute(conn::Connection, sql::AbstractString, params=nothi stmtname = API.prepare(conn.socket, sql_str, debug; name="") nparams, names, types = API.describeprepared(conn.socket, stmtname, debug) params_vec = build_params(params, nparams, sql_str) - e = API.exec(conn.socket, stmtname, params_vec, names, types, conn.type_registry, debug, 0, conn.notice_callback, conn.notification_callback) - result = T === Any ? makeresult(e) : StructUtils.arraylike(T) ? StructUtils.make(T, e, PostgresStyle()) : only(StructUtils.make(Vector{T}, e, PostgresStyle())) + # see the statement-execute method: socket union isa-split for --trim + socket = conn.socket + e = if socket isa Reseau.TCP.Conn + API.exec(style, socket::Reseau.TCP.Conn, stmtname, params_vec, names, types, conn.type_registry, debug, 0) + else + API.exec(style, socket::Reseau.TLS.Conn, stmtname, params_vec, names, types, conn.type_registry, debug, 0) + end + result = T === Any ? makeresult(e) : StructUtils.arraylike(T) ? StructUtils.make(T, e, style) : only(StructUtils.make(Vector{T}, e, style)) end - log_enabled && log_query(logger, :execute, (sql=sql_str, params=params, duration_ns=time_ns() - start_ns, success=true)) + log_enabled && API.query_logger(style, :execute, (sql=sql_str, params=params, duration_ns=time_ns() - start_ns, success=true)) return result catch err - log_enabled && log_query(logger, :execute, (sql=sql_str, params=params, duration_ns=time_ns() - start_ns, success=false, error=err)) + log_enabled && API.query_logger(style, :execute, (sql=sql_str, params=params, duration_ns=time_ns() - start_ns, success=false, error=err)) rethrow() end end @@ -368,7 +415,7 @@ function cursor(stmt::Statement, params=nothing; fetchsize::Integer=1000, owns_t portal = string(UUIDs.uuid4()) types = Type[API.juliatype(x -> x, i, conn.type_registry) for i in stmt.typeIds] lookup = Dict(x => i for (i, x) in enumerate(stmt.names)) - cursor = Cursor(conn, portal, stmt.names, stmt.typeIds, types, lookup, max(1, Int(fetchsize)), ResultRow[], 1, false, 0, owns_transaction) + cursor = Cursor{_style_type(conn)}(conn, portal, stmt.names, stmt.typeIds, types, lookup, max(1, Int(fetchsize)), ResultRow[], 1, false, 0, owns_transaction) API.writemessages(conn.socket, conn.debug, ('B', portal, stmt.name, Int16(0), Int16(length(stmt.params)), API.Params(stmt.params), Int16(0)), ('E', portal, Int32(cursor.fetchsize)), ('S',)) read_portal_batch!(cursor) return cursor diff --git a/test/postgres_trim_queries.jl b/test/postgres_trim_queries.jl new file mode 100644 index 0000000..560502a --- /dev/null +++ b/test/postgres_trim_queries.jl @@ -0,0 +1,176 @@ +using Dates +using DBInterface +using Postgres +using StructUtils +using UUIDs + +struct TrimId + profile_id::Int32 +end + +struct TrimCount + count::Int32 +end + +struct TrimName + display_name::String +end + +StructUtils.@tags struct TrimProfile + profileId::Int32 &(postgres=(name=:profile_id,),) + displayName::String &(postgres=(name=:display_name,),) + createdAt::DateTime &(postgres=(name=:created_at,),) + active::Bool + score::Union{Missing, Int32} + uid::UUID + flags::Vector{Int32} +end + +function _postgres_trim_connect() + host = get(ENV, "POSTGRES_TRIM_HOST", "127.0.0.1") + port = parse(Int, get(ENV, "POSTGRES_TRIM_PORT", "5432")) + user = get(ENV, "POSTGRES_TRIM_USER", "postgres") + dbname = get(ENV, "POSTGRES_TRIM_DBNAME", "postgres") + return DBInterface.connect( + Postgres.Connection, + host, + user, + nothing; + dbname=dbname, + port=port, + sslmode="disable", + connect_timeout=2, + application_name="postgres_trim", + statement_cache_maxsize=4, + ) +end + +function _assert_trim_profile(profile::TrimProfile, id::Int32, name::String)::Nothing + profile.profileId == id || error("unexpected profile id") + profile.displayName == name || error("unexpected profile name") + profile.active || error("expected active profile") + profile.uid == UUID("12345678-1234-5678-1234-567812345678") || id != 1 || error("unexpected UUID") + !isempty(profile.flags) || error("expected non-empty flags array") + return nothing +end + +function run_postgres_trim_queries()::Nothing + conn = _postgres_trim_connect() + try + DBInterface.execute(conn, """ + CREATE TEMP TABLE trim_compile_profiles ( + profile_id integer PRIMARY KEY, + display_name text NOT NULL, + created_at timestamp NOT NULL, + active boolean NOT NULL, + score integer, + uid uuid NOT NULL, + flags integer[] NOT NULL + ) + """) + + insert_sql = raw""" + INSERT INTO trim_compile_profiles ( + profile_id, + display_name, + created_at, + active, + score, + uid, + flags + ) VALUES ( + $1, + $2, + $3, + $4, + $5, + $6, + $7::integer[] + ) + RETURNING profile_id + """ + + first_id = DBInterface.execute( + conn, + insert_sql, + ( + Int32(1), + "Ada", + DateTime(2024, 1, 2, 3, 4, 5), + true, + Int32(99), + UUID("12345678-1234-5678-1234-567812345678"), + Int32[1, 2, 3], + ), + TrimId, + ) + first_id.profile_id == 1 || error("unexpected inserted id") + + DBInterface.transaction(conn) do + DBInterface.execute( + conn, + insert_sql, + ( + Int32(2), + "Grace", + DateTime(2024, 1, 3, 4, 5, 6), + true, + missing, + UUID("87654321-4321-8765-4321-876543218765"), + Int32[4, 5], + ), + ) + end + + stmt = DBInterface.prepare(conn, raw""" + SELECT profile_id, display_name, created_at, active, score, uid, flags + FROM trim_compile_profiles + WHERE profile_id = $1 + """) + try + profile = DBInterface.execute(stmt, (Int32(1),), TrimProfile) + _assert_trim_profile(profile, Int32(1), "Ada") + finally + DBInterface.close!(stmt) + end + + profiles = DBInterface.execute(conn, """ + SELECT profile_id, display_name, created_at, active, score, uid, flags + FROM trim_compile_profiles + ORDER BY profile_id + """, (), Vector{TrimProfile}) + length(profiles) == 2 || error("expected two profiles") + _assert_trim_profile(profiles[2], Int32(2), "Grace") + + count_row = DBInterface.execute( + conn, + "SELECT count(*)::integer AS count FROM trim_compile_profiles", + (), + TrimCount, + ) + count_row.count == 2 || error("unexpected typed count") + + name_row = DBInterface.execute( + conn, + "SELECT display_name FROM trim_compile_profiles WHERE profile_id = 1", + (), + TrimName, + ) + name_row.display_name == "Ada" || error("unexpected typed name") + + update_result = DBInterface.execute(conn, "UPDATE trim_compile_profiles SET score = coalesce(score, 0) + 1") + Postgres.rows_affected(update_result) == 2 || error("unexpected rows affected") + occursin("UPDATE", Postgres.command_tag(update_result)) || error("unexpected command tag") + finally + DBInterface.close!(conn) + end + return nothing +end + +function @main(args::Vector{String})::Cint + _ = args + run_postgres_trim_queries() + return 0 +end + +Base.Experimental.entrypoint(main, (Vector{String},)) diff --git a/test/runtests.jl b/test/runtests.jl index 0c3366c..76609c5 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -10,6 +10,16 @@ using Postgres using Sockets using Random +# Style-based customization (the runtime callback setters are gone): overload the +# behavior interface on a custom AbstractPostgresStyle and pass it at connect time. +const LOGGED_EVENTS = NamedTuple[] +const NOTICE_SEEN = Ref(false) +struct LoggingStyle <: Postgres.API.AbstractPostgresStyle end +Postgres.API.query_logging_enabled(::LoggingStyle) = true +Postgres.API.query_logger(::LoggingStyle, event::Symbol, info::NamedTuple) = (push!(LOGGED_EVENTS, (event=event, info=info)); nothing) +Postgres.API.notice_callback(::LoggingStyle, notice) = (NOTICE_SEEN[] = true; nothing) + + # Integration tests for Postgres.jl protocol and API behavior. const JSONType = typeof(JSON.lazy("{}")) const IMAGE_REF = get(ENV, "POSTGRES_IMAGE", "postgres:16") @@ -548,6 +558,8 @@ end end end + include("trim_compile_tests.jl") + if !docker_available() @info "Docker not available; skipping Postgres integration tests." @test true @@ -612,7 +624,7 @@ end # a consumer exception mid-result drains the remaining rows so # the connection stays usable stmt = DBInterface.prepare(connp, "SELECT i AS x, repeat('y', 10) AS s FROM generate_series(1, 200) i") - ex = Postgres.API.exec(connp.socket, stmt.name, Union{String, Missing}[], stmt.names, stmt.typeIds, connp.type_registry, false) + ex = Postgres.API.exec(connp.style, connp.socket, stmt.name, Union{String, Missing}[], stmt.names, stmt.typeIds, connp.type_registry, false) err = try StructUtils.applyeach(Postgres.API.PostgresStyle(), ex) do i, row i == 3 && error("consumer abort") @@ -791,6 +803,16 @@ end Postgres.start_transaction(conn) @test_throws Postgres.API.Error DBInterface.execute(conn, "INVALID SQL") Postgres.rollback(conn) + + # Simple-query protocol messages include CommandComplete + # before ReadyForQuery. Verify the driver consumes both and + # routes asynchronous notices through the connection style. + NOTICE_SEEN[] = false + simple_conn = DBInterface.connect(Postgres.Connection, cfg.host, cfg.user, cfg.password; dbname=cfg.dbname, port=cfg.port, style=LoggingStyle()) + Postgres.API.exec(simple_conn.style, simple_conn.socket, raw"DO $$ BEGIN RAISE NOTICE 'simple query'; END $$;", false) + @test NOTICE_SEEN[] + @test Tables.rowtable(DBInterface.execute(simple_conn, "SELECT 7 AS a"))[1].a == 7 + close(simple_conn) end @testset "Transaction Macro" begin @@ -939,21 +961,17 @@ end @test count_row.count == 3 end - @testset "Query Logger" begin - events = NamedTuple[] - previous = Postgres.get_query_logger(conn) - Postgres.set_query_logger!(conn, (event, info) -> begin - push!(events, (event=event, info=info)) - return - end) - rows = Tables.rowtable(DBInterface.execute(conn, "SELECT 1 AS a")) + @testset "Query Logger (style)" begin + log_conn = DBInterface.connect(Postgres.Connection, cfg.host, cfg.user, cfg.password; dbname=cfg.dbname, port=cfg.port, style=LoggingStyle()) + empty!(LOGGED_EVENTS) + rows = Tables.rowtable(DBInterface.execute(log_conn, "SELECT 1 AS a")) @test rows[1].a == 1 - @test !isempty(events) - @test events[end].event == :execute - @test events[end].info.success - @test_throws Postgres.API.Error DBInterface.execute(conn, "INVALID SQL") - @test !events[end].info.success - Postgres.set_query_logger!(conn, previous) + @test !isempty(LOGGED_EVENTS) + @test LOGGED_EVENTS[end].event == :execute + @test LOGGED_EVENTS[end].info.success + @test_throws Postgres.API.Error DBInterface.execute(log_conn, "INVALID SQL") + @test !LOGGED_EVENTS[end].info.success + close(log_conn) end @testset "Connection Pool" begin @@ -1039,16 +1057,12 @@ end DBInterface.close!(cur) end - @testset "Notice Callback" begin - notice_seen = Ref(false) - previous = Postgres.get_notice_callback(conn) - Postgres.set_notice_callback!(conn, notice -> begin - notice_seen[] = true - return - end) - DBInterface.execute(conn, raw"DO $$ BEGIN RAISE NOTICE 'hello'; END $$;") - Postgres.set_notice_callback!(conn, previous) - @test notice_seen[] + @testset "Notice Callback (style)" begin + NOTICE_SEEN[] = false + notice_conn = DBInterface.connect(Postgres.Connection, cfg.host, cfg.user, cfg.password; dbname=cfg.dbname, port=cfg.port, style=LoggingStyle()) + DBInterface.execute(notice_conn, raw"DO $$ BEGIN RAISE NOTICE 'hello'; END $$;") + close(notice_conn) + @test NOTICE_SEEN[] end @testset "Cancel Request" begin @@ -1075,6 +1089,8 @@ end complex_interval = only(Tables.rowtable(DBInterface.execute(conn, "SELECT '1 year 2 mons 3 days 04:05:06.789'::interval AS interval_col"))) @test complex_interval.interval_col == Dates.CompoundPeriod(Dates.Year(1), Dates.Month(2), Dates.Day(3), Dates.Hour(4), Dates.Minute(5), Dates.Second(6), Dates.Millisecond(789)) end + + run_postgres_trim_compile_tests(cfg) finally isopen(conn) && DBInterface.close!(conn) end diff --git a/test/trim_compile_tests.jl b/test/trim_compile_tests.jl new file mode 100644 index 0000000..bec75b3 --- /dev/null +++ b/test/trim_compile_tests.jl @@ -0,0 +1,263 @@ +using Test + +const _POSTGRES_TRIM_SUPPORTED = VERSION >= v"1.12.0-rc1" +const _POSTGRES_TRIM_PRE_RELEASE = !isempty(VERSION.prerelease) +const _POSTGRES_JULIAC_ENTRYPOINT_EXPR = "using JuliaC; if isdefined(JuliaC, :main); JuliaC.main(ARGS); else JuliaC._main_cli(ARGS); end" + +function _postgres_trim_compile_timeout_s()::Float64 + default = Sys.iswindows() ? "1200.0" : "180.0" + return parse(Float64, get(ENV, "POSTGRES_TRIM_COMPILE_TIMEOUT_S", default)) +end + +function _postgres_trim_error_budget()::Int + # StructUtils 2.8.2 still takes its generic construction path. The + # companion trim branch removes these errors; keep released-dependency CI + # bounded so compiler regressions are visible in the meantime. + return parse(Int, get(ENV, "POSTGRES_TRIM_ERROR_BUDGET", "92")) +end + +function _postgres_trim_project_path()::String + active_project = Base.active_project() + if active_project !== nothing && isfile(active_project) + return dirname(active_project) + end + return normpath(joinpath(@__DIR__, "..")) +end + +function _postgres_trim_env(cfg) + run_env = copy(ENV) + run_env["POSTGRES_TRIM_HOST"] = cfg.host + run_env["POSTGRES_TRIM_PORT"] = string(cfg.port) + run_env["POSTGRES_TRIM_USER"] = cfg.user + run_env["POSTGRES_TRIM_DBNAME"] = cfg.dbname + return run_env +end + +function _run_postgres_trim_compile(project_path::String, script_path::String, output_name::String; timeout_s::Float64 = _postgres_trim_compile_timeout_s(), bundle_dir::Union{Nothing, String} = nothing) + julia_exe = joinpath(Sys.BINDIR, Base.julia_exename()) + cmd = if bundle_dir === nothing + `$julia_exe --startup-file=no --history-file=no --code-coverage=none --project=$project_path -e $(_POSTGRES_JULIAC_ENTRYPOINT_EXPR) -- --output-exe $output_name --project=$project_path --experimental --trim=safe $script_path` + else + `$julia_exe --startup-file=no --history-file=no --code-coverage=none --project=$project_path -e $(_POSTGRES_JULIAC_ENTRYPOINT_EXPR) -- --output-exe $output_name --bundle $bundle_dir --project=$project_path --experimental --trim=safe $script_path` + end + compile_env = copy(ENV) + compile_env["RESEAU_PRECOMPILE_ONLY"] = get(ENV, "POSTGRES_TRIM_RESEAU_PRECOMPILE_ONLY", "tcp") + cmd = setenv(cmd, compile_env) + return _run_postgres_trim_command_with_timeout(cmd; timeout_s = timeout_s, log_label = "compile") +end + +function _run_postgres_trim_executable(run_cmd::Cmd; timeout_s::Float64 = 30.0) + return _run_postgres_trim_command_with_timeout(run_cmd; timeout_s = timeout_s, log_label = "run") +end + +function _run_postgres_trim_command_with_timeout(cmd::Cmd; timeout_s::Float64, log_label::String) + output_path = tempname() + out = open(output_path, "w") + exit_code = -1 + timed_out = false + try + proc = run(pipeline(ignorestatus(cmd), stdout = out, stderr = out); wait = false) + timed_out = _wait_postgres_trim_process_with_timeout!(proc; timeout_s = timeout_s, log_label = log_label) + exit_code = something(proc.exitcode, -1) + finally + close(out) + end + output = try + read(output_path, String) + catch + "" + finally + rm(output_path; force = true) + end + return exit_code, output, timed_out +end + +function _wait_postgres_trim_process_with_timeout!(proc::Base.Process; timeout_s::Float64, log_label::String) + started_at = time() + next_log_at = started_at + 10.0 + timed_out = false + while Base.process_running(proc) + now = time() + if now - started_at >= timeout_s + timed_out = true + try + kill(proc) + catch + end + _kill_postgres_trim_windows_process_tree!(proc) + _wait_postgres_trim_process_exit_after_kill!(proc; timeout_s = 5.0, log_label = log_label) + return timed_out + end + if now >= next_log_at + elapsed = round(now - started_at; digits = 1) + println("[trim] $(log_label) WAIT $(elapsed)s") + flush(stdout) + next_log_at = now + 10.0 + end + sleep(0.1) + end + if !Base.process_running(proc) + try + wait(proc) + catch + end + end + return timed_out +end + +function _kill_postgres_trim_windows_process_tree!(proc::Base.Process)::Nothing + Sys.iswindows() || return nothing + pid = try + getpid(proc) + catch + return nothing + end + try + run(ignorestatus(`taskkill /PID $pid /T /F`)) + catch + end + return nothing +end + +function _wait_postgres_trim_process_exit_after_kill!(proc::Base.Process; timeout_s::Float64, log_label::String)::Nothing + deadline = time() + timeout_s + while Base.process_running(proc) && time() < deadline + sleep(0.1) + end + if Base.process_running(proc) + println("[trim] $(log_label) process still running after kill; continuing after timeout") + flush(stdout) + end + return nothing +end + +function _postgres_trim_timeout_error(kind::String, script_file::String, output::String = "") + msg = "trim $kind timed out for $(script_file)" + if !isempty(output) + msg = string(msg, "\n---- captured output ----\n", output, "\n---- end captured output ----") + end + throw(ArgumentError(msg)) +end + +function _maybe_print_postgres_trim_output(header::String, output::String) + isempty(output) && return nothing + println(header) + println(output) + println("---- end output ----") + return nothing +end + +function _postgres_trim_executable_timeout_s()::Float64 + default = Sys.iswindows() ? "180.0" : "30.0" + return parse(Float64, get(ENV, "POSTGRES_TRIM_EXE_TIMEOUT_S", default)) +end + +function _postgres_trim_selected_workloads(workloads::Vector{Tuple{String, String}})::Vector{Tuple{String, String}} + only = strip(get(ENV, "POSTGRES_TRIM_ONLY", "")) + isempty(only) && return workloads + selected = Tuple{String, String}[] + for workload in workloads + workload[1] == only && push!(selected, workload) + end + isempty(selected) && throw(ArgumentError("unknown POSTGRES_TRIM_ONLY workload: $(only)")) + return selected +end + +function _postgres_trim_use_bundle()::Bool + return get(ENV, "POSTGRES_TRIM_BUNDLE", "0") == "1" +end + +function _parse_postgres_trim_verify_totals(output::String) + m = match(r"Trim verify finished with\s+(\d+)\s+errors,\s+(\d+)\s+warnings\.", output) + m === nothing && return nothing + return parse(Int, m.captures[1]), parse(Int, m.captures[2]) +end + +function _count_postgres_trim_verify_messages(output::String)::Tuple{Int,Int} + errors = length(collect(eachmatch(r"Verifier error #\d+:", output))) + warnings = length(collect(eachmatch(r"Verifier warning #\d+:", output))) + return errors, warnings +end + +function _run_postgres_trim_case(cfg, project_path::String, script_file::String, output_name::String) + script_path = joinpath(@__DIR__, script_file) + @test isfile(script_path) + println("[trim] compile START $(script_file)") + start_t = time() + mktempdir() do tmpdir + cd(tmpdir) do + bundle_dir = _postgres_trim_use_bundle() ? joinpath(tmpdir, "bundle") : nothing + exit_code, output, timed_out = _run_postgres_trim_compile(project_path, script_path, output_name; bundle_dir = bundle_dir) + if timed_out + _postgres_trim_timeout_error("compile", script_file, output) + end + totals = _parse_postgres_trim_verify_totals(output) + trim_errors, trim_warnings = if totals === nothing + fallback = _count_postgres_trim_verify_messages(output) + if exit_code != 0 && fallback == (0, 0) + error("failed to parse trim verifier summary:\n$output") + end + fallback + else + totals + end + trim_error_budget = _postgres_trim_error_budget() + println("[trim] verifier $(script_file): errors=$(trim_errors) warnings=$(trim_warnings) budget=$(trim_error_budget)") + if get(ENV, "POSTGRES_TRIM_PRINT_OUTPUT", "0") == "1" || trim_errors > trim_error_budget || trim_warnings > 0 + _maybe_print_postgres_trim_output("---- trim compile output ($(script_file)) ----", output) + end + @test trim_errors <= trim_error_budget + @test trim_warnings == 0 + if trim_errors > 0 + @test exit_code != 0 + println("[trim] executable skipped for $(script_file): verifier errors remain") + return nothing + end + output_path = Sys.iswindows() ? "$(output_name).exe" : output_name + run_path = bundle_dir === nothing ? output_path : joinpath(bundle_dir, "bin", output_path) + @test exit_code == 0 + @test isfile(run_path) + run_cmd = setenv(`$(abspath(run_path))`, _postgres_trim_env(cfg)) + run_timeout_s = _postgres_trim_executable_timeout_s() + run_exit, run_output, run_timed_out = _run_postgres_trim_executable(run_cmd; timeout_s = run_timeout_s) + if run_timed_out + _postgres_trim_timeout_error("executable run", script_file, run_output) + end + if run_exit != 0 + _maybe_print_postgres_trim_output("---- trim executable output ($(script_file)) ----", run_output) + end + @test run_exit == 0 + end + end + println("[trim] compile DONE $(script_file) ($(round(time() - start_t; digits = 2))s)") + return nothing +end + +function run_postgres_trim_compile_tests(cfg)::Nothing + @testset "Trim Compile" begin + if Sys.iswindows() + println("[trim] skip Windows: JuliaC trim compilation is currently too slow or stalls on Windows CI") + @test true + elseif !_POSTGRES_TRIM_SUPPORTED + println("[trim] skip Julia < 1.12: JuliaC trim compilation is unavailable") + @test true + elseif _POSTGRES_TRIM_PRE_RELEASE + println("[trim] skip prerelease Julia: trim verifier behavior is not stable yet") + @test true + elseif !(occursin("trust", DEFAULT_AUTH) || occursin("trust", DEFAULT_INITDB_ARGS)) + println("[trim] skip non-trust auth mode: main CI covers trim once, auth-mode jobs focus on authentication") + @test true + else + project_path = _postgres_trim_project_path() + println("[trim] project $(project_path)") + trim_workloads = [ + ("postgres_trim_queries.jl", "postgres_trim_queries"), + ] + trim_workloads = _postgres_trim_selected_workloads(trim_workloads) + for (script_file, output_name) in trim_workloads + _run_postgres_trim_case(cfg, project_path, script_file, output_name) + end + end + end + return nothing +end