diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index bb137e7..7975f75 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -37,4 +37,11 @@ jobs: arch: ${{ matrix.arch }} - uses: julia-actions/cache@v2 - uses: julia-actions/julia-buildpkg@v1 - - uses: julia-actions/julia-runtest@v1 \ No newline at end of file + - uses: julia-actions/julia-runtest@v1 + with: + coverage: true + - uses: julia-actions/julia-processcoverage@v1 + - uses: actions/upload-artifact@v4 + with: + name: coverage-lcov-${{ matrix.version }}-${{ matrix.os }}-${{ matrix.arch }} + path: lcov.info \ No newline at end of file diff --git a/src/context.jl b/src/context.jl index fd838cf..02f7374 100644 --- a/src/context.jl +++ b/src/context.jl @@ -126,10 +126,13 @@ function _with_conn(f::Function, ctx::QueryContext) # Sync pool's source registry with context's registry # (register! on ctx also adds to pool, but belt-and-suspenders) for (name, src) in ctx.sources - if !(name in get(ctx._pool._applied, conn, Set{String}())) + needs_apply = lock(ctx._pool._lock) do + !(name in get(ctx._pool._applied, conn, Set{String}())) + end + if needs_apply _register_source!(conn, name, src) lock(ctx._pool._lock) do - push!(ctx._pool._applied[conn], name) + push!(get!(ctx._pool._applied, conn, Set{String}()), name) end end end diff --git a/src/params.jl b/src/params.jl index dcb7bb1..7cca416 100644 --- a/src/params.jl +++ b/src/params.jl @@ -52,19 +52,156 @@ keyword argument. function _bind_named(sql::String, named_params::Base.Pairs)::Tuple{String, Vector} values = Any[] order = String[] - # Match :word — but not :: (cast operator) and not inside string literals - sql_out = replace(sql, r"(? function(m) - param_name = Symbol(m[2:end]) # strip leading ':' - if !haskey(named_params, param_name) - throw(QueryError( - "Named parameter ':$(param_name)' not provided.", - sql, Dict(named_params), nothing - )) + io = IOBuffer() + + is_word_char(c::Char) = isletter(c) || isdigit(c) || c == '_' + + i = firstindex(sql) + in_single = false + in_double = false + in_line_comment = false + in_block_comment = false + + while i <= lastindex(sql) + c = sql[i] + ni = i < lastindex(sql) ? nextind(sql, i) : nothing + nc = ni === nothing ? '\0' : sql[ni] + + if in_line_comment + write(io, c) + if c == '\n' + in_line_comment = false + end + i = ni === nothing ? lastindex(sql) + 1 : ni + continue end - push!(values, named_params[param_name]) - push!(order, string(param_name)) - "?" - end) + + if in_block_comment + write(io, c) + if c == '*' && ni !== nothing && nc == '/' + write(io, nc) + i = ni < lastindex(sql) ? nextind(sql, ni) : lastindex(sql) + 1 + in_block_comment = false + else + i = ni === nothing ? lastindex(sql) + 1 : ni + end + continue + end + + if in_single + write(io, c) + if c == '\'' + if ni !== nothing && nc == '\'' + write(io, nc) + i = ni < lastindex(sql) ? nextind(sql, ni) : lastindex(sql) + 1 + else + in_single = false + i = ni === nothing ? lastindex(sql) + 1 : ni + end + else + i = ni === nothing ? lastindex(sql) + 1 : ni + end + continue + end + + if in_double + write(io, c) + if c == '"' + if ni !== nothing && nc == '"' + write(io, nc) + i = ni < lastindex(sql) ? nextind(sql, ni) : lastindex(sql) + 1 + else + in_double = false + i = ni === nothing ? lastindex(sql) + 1 : ni + end + else + i = ni === nothing ? lastindex(sql) + 1 : ni + end + continue + end + + if c == '-' && ni !== nothing && nc == '-' + write(io, c) + write(io, nc) + in_line_comment = true + i = ni < lastindex(sql) ? nextind(sql, ni) : lastindex(sql) + 1 + continue + end + + if c == '/' && ni !== nothing && nc == '*' + write(io, c) + write(io, nc) + in_block_comment = true + i = ni < lastindex(sql) ? nextind(sql, ni) : lastindex(sql) + 1 + continue + end + + if c == '\'' + in_single = true + write(io, c) + i = ni === nothing ? lastindex(sql) + 1 : ni + continue + end + + if c == '"' + in_double = true + write(io, c) + i = ni === nothing ? lastindex(sql) + 1 : ni + continue + end + + if c == ':' + prev_i = i > firstindex(sql) ? prevind(sql, i) : nothing + prev_c = prev_i === nothing ? '\0' : sql[prev_i] + + # Skip cast operator (::) and ensure we're not in a word token. + if ni !== nothing && nc == ':' + write(io, c) + write(io, nc) + i = ni < lastindex(sql) ? nextind(sql, ni) : lastindex(sql) + 1 + continue + elseif prev_i !== nothing && (prev_c == ':' || is_word_char(prev_c)) + write(io, c) + i = ni === nothing ? lastindex(sql) + 1 : ni + continue + end + + # Parse :identifier + if ni !== nothing + first_name_c = sql[ni] + if isletter(first_name_c) || first_name_c == '_' + j = ni + while true + jj = j < lastindex(sql) ? nextind(sql, j) : nothing + if jj === nothing + break + end + cj = sql[jj] + is_word_char(cj) || break + j = jj + end + name_str = sql[ni:j] + param_name = Symbol(name_str) + if !haskey(named_params, param_name) + throw(QueryError( + "Named parameter ':$(param_name)' not provided.", + sql, Dict(named_params), nothing + )) + end + push!(values, named_params[param_name]) + push!(order, string(param_name)) + write(io, '?') + i = j < lastindex(sql) ? nextind(sql, j) : lastindex(sql) + 1 + continue + end + end + end + + write(io, c) + i = ni === nothing ? lastindex(sql) + 1 : ni + end + + sql_out = String(take!(io)) @debug "Bound named params" order=order values=values return (sql_out, values) end diff --git a/src/pool.jl b/src/pool.jl index 5d0dbf0..8eb942b 100644 --- a/src/pool.jl +++ b/src/pool.jl @@ -38,10 +38,13 @@ mutable struct ConnectionPool config::QueryConfig channel::Channel{DuckDB.DB} size::Int + _closed::Bool # Per-connection set of source names already registered on that connection. # Guarded by _lock so it is safe to read/write from multiple tasks. _applied::IdDict{DuckDB.DB, Set{String}} + _in_use::IdDict{DuckDB.DB, Bool} + _pending_drops::IdDict{DuckDB.DB, Dict{String, Any}} _lock::ReentrantLock # Shared source registry (name → raw source object). @@ -58,8 +61,9 @@ mutable struct ConnectionPool _validate_config(config) ch = Channel{DuckDB.DB}(size) pool = new( - db_path, config, ch, size, - IdDict{DuckDB.DB, Set{String}}(), ReentrantLock(), + db_path, config, ch, size, false, + IdDict{DuckDB.DB, Set{String}}(), IdDict{DuckDB.DB, Bool}(), + IdDict{DuckDB.DB, Dict{String, Any}}(), ReentrantLock(), Dict{String, Any}() ) # Pre-warm all connections @@ -98,6 +102,26 @@ function _ensure_sources_applied!(pool::ConnectionPool, conn::DuckDB.DB) end end +function _apply_pending_drops!(pool::ConnectionPool, conn::DuckDB.DB) + pending = lock(pool._lock) do + haskey(pool._pending_drops, conn) || return Pair{String, Any}[] + queued = collect(pool._pending_drops[conn]) + delete!(pool._pending_drops, conn) + queued + end + + for (name, src) in pending + try + _deregister_source!(conn, name, src) + catch e + @warn "Pool: failed pending source deregistration" name=name exception=e + end + lock(pool._lock) do + delete!(get!(pool._applied, conn, Set{String}()), name) + end + end +end + function _finalizer_close!(pool::ConnectionPool) try close!(pool) catch end # never throw from a finalizer end @@ -111,10 +135,34 @@ Borrow a connection from the pool, blocking until one is available. You **must** call `release!` when done, or preferably use `with_connection`. """ function acquire!(pool::ConnectionPool)::DuckDB.DB + lock(pool._lock) do + pool._closed && throw(QueryError("ConnectionPool has been closed.")) + end + conn = take!(pool.channel) - _ensure_sources_applied!(pool, conn) - @debug "Pool: acquired connection" pool_size=pool.size - return conn + try + lock(pool._lock) do + if pool._closed + throw(QueryError("ConnectionPool has been closed.")) + end + pool._in_use[conn] = true + end + _apply_pending_drops!(pool, conn) + _ensure_sources_applied!(pool, conn) + @debug "Pool: acquired connection" pool_size=pool.size + return conn + catch + should_close = lock(pool._lock) do + haskey(pool._in_use, conn) && delete!(pool._in_use, conn) + pool._closed + end + if should_close + try close(conn) catch end + else + put!(pool.channel, conn) + end + rethrow() + end end """ @@ -123,6 +171,21 @@ end Return a previously acquired connection to the pool. """ function release!(pool::ConnectionPool, conn::DuckDB.DB) + was_in_use, is_closed = lock(pool._lock) do + in_use = haskey(pool._in_use, conn) + in_use && delete!(pool._in_use, conn) + in_use, pool._closed + end + + if is_closed + try close(conn) catch end + @debug "Pool: released connection after close (closed instead of re-pooled)" + return + end + + was_in_use || throw(QueryError("Attempted to release a connection that is not checked out from this pool.")) + _apply_pending_drops!(pool, conn) + _ensure_sources_applied!(pool, conn) put!(pool.channel, conn) @debug "Pool: released connection" end @@ -153,12 +216,22 @@ end Drain the pool channel and close every connection. """ function close!(pool::ConnectionPool) + borrowed_count = lock(pool._lock) do + pool._closed && return 0 + pool._closed = true + length(pool._in_use) + end + # Drain and close all connections currently in the channel while isready(pool.channel) conn = take!(pool.channel) try close(conn) catch end + lock(pool._lock) do + delete!(pool._applied, conn) + delete!(pool._pending_drops, conn) + end end - @debug "Pool: closed all connections" + @debug "Pool: closed all available connections" borrowed=borrowed_count end """ @@ -170,6 +243,10 @@ The source will be lazily applied to each connection on first use. See `register!` on `QueryContext` for supported source types. """ function register!(pool::ConnectionPool, name::String, source) + lock(pool._lock) do + pool._closed && throw(QueryError("ConnectionPool has been closed.")) + end + pool.sources[name] = source # Invalidate per-connection tracking so all existing connections # will re-apply the full source list on next acquisition. diff --git a/src/query.jl b/src/query.jl index 2ef5401..bc4f97f 100644 --- a/src/query.jl +++ b/src/query.jl @@ -269,6 +269,7 @@ function explain(ctx::QueryContext, sql::String; analyze::Bool=false)::String prefix = analyze ? "EXPLAIN ANALYZE " : "EXPLAIN " _with_conn(ctx) do conn df, _ = _run(conn, ctx.config, "$prefix$sql") + (ncol(df) == 0 || nrow(df) == 0) && return "" join(df[!, end], "\n") end end diff --git a/src/sources.jl b/src/sources.jl index 34b8fee..db35720 100644 --- a/src/sources.jl +++ b/src/sources.jl @@ -72,6 +72,11 @@ register!(ctx, ``` """ function register!(ctx::QueryContext, name::String, source; union_by_name::Bool=false) + ctx._closed && throw(QueryError("QueryContext has been closed.")) + + had_previous = haskey(ctx.sources, name) + previous_source = had_previous ? ctx.sources[name] : nothing + actual_source = if union_by_name && source isa String _looks_like_parquet(source) || throw(ArgumentError( "union_by_name is only supported for Parquet/glob sources, got: $source" @@ -81,8 +86,69 @@ function register!(ctx::QueryContext, name::String, source; union_by_name::Bool= source end ctx.sources[name] = actual_source + + if ctx._pool !== nothing + pool = ctx._pool + + lock(pool._lock) do + pool.sources[name] = actual_source + end + + available = DuckDB.DB[] + while isready(pool.channel) + push!(available, take!(pool.channel)) + end + + try + # Apply changes immediately to currently idle connections. + for conn in available + if had_previous + try + _deregister_source!(conn, name, previous_source) + catch e + @warn "Failed to deregister replaced pooled source" name=name exception=e + end + lock(pool._lock) do + delete!(get!(pool._applied, conn, Set{String}()), name) + end + end + + _register_source!(conn, name, actual_source) + lock(pool._lock) do + push!(get!(pool._applied, conn, Set{String}()), name) + end + end + + # For checked-out connections, queue replacement cleanup and force + # source re-application when they are next released/acquired. + lock(pool._lock) do + for conn in keys(pool._in_use) + if had_previous + pending = get!(pool._pending_drops, conn, Dict{String, Any}()) + pending[name] = previous_source + end + delete!(get!(pool._applied, conn, Set{String}()), name) + end + end + finally + for conn in available + put!(pool.channel, conn) + end + end + + @debug "Source registered (pooled context)" name=name type=typeof(actual_source) + return + end + # Eagerly apply to the live single connection (if any) if ctx._conn !== nothing + if had_previous + try + _deregister_source!(ctx._conn, name, previous_source) + catch e + @warn "Failed to deregister replaced source" name=name exception=e + end + end _register_source!(ctx._conn, name, actual_source) end @debug "Source registered" name=name type=typeof(actual_source) @@ -102,11 +168,53 @@ Remove a previously registered source. For DataFrames and views this drops the VIEW; for attached databases this issues DETACH. """ function deregister!(ctx::QueryContext, name::String) + ctx._closed && throw(QueryError("QueryContext has been closed.")) + if !haskey(ctx.sources, name) @warn "Attempted to deregister unknown source" name=name return end src = pop!(ctx.sources, name) + + if ctx._pool !== nothing + pool = ctx._pool + + lock(pool._lock) do + haskey(pool.sources, name) && delete!(pool.sources, name) + end + + available = DuckDB.DB[] + while isready(pool.channel) + push!(available, take!(pool.channel)) + end + + try + # Apply deregistration immediately to idle pooled connections. + for conn in available + _deregister_source!(conn, name, src) + lock(pool._lock) do + delete!(get!(pool._applied, conn, Set{String}()), name) + end + end + + # Queue deregistration for currently checked-out connections. + lock(pool._lock) do + for conn in keys(pool._in_use) + pending = get!(pool._pending_drops, conn, Dict{String, Any}()) + pending[name] = src + delete!(get!(pool._applied, conn, Set{String}()), name) + end + end + finally + for conn in available + put!(pool.channel, conn) + end + end + + @debug "Source deregistered (pooled context)" name=name + return + end + if ctx._conn !== nothing _deregister_source!(ctx._conn, name, src) end diff --git a/test/runtests.jl b/test/runtests.jl index c5dfdfc..6afde53 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -226,6 +226,16 @@ global_logger(ConsoleLogger(stderr, Logging.Warn)) close!(ctx) end + @testset "on_error = :missing" begin + ctx = QueryContext(on_error=:missing) + df = execute(ctx, "SELECT * FROM nonexistent_table_xyz") + @test df isa DataFrame + @test names(df) == ["result"] + @test nrow(df) == 1 + @test ismissing(df[1, :result]) + close!(ctx) + end + # ── explain ─────────────────────────────────────────────────────────────── @testset "explain" begin ctx = QueryContext() @@ -235,6 +245,14 @@ global_logger(ConsoleLogger(stderr, Logging.Warn)) close!(ctx) end + @testset "explain with non-throw on_error" begin + ctx = QueryContext(on_error=:empty) + plan = explain(ctx, "SELECT * FROM nonexistent_table_xyz") + @test plan isa String + @test isempty(plan) + close!(ctx) + end + # ── QueryResult – Tables.jl interface ──────────────────────────────────── @testset "QueryResult Tables.jl interface" begin ctx = QueryContext() @@ -264,6 +282,116 @@ global_logger(ConsoleLogger(stderr, Logging.Warn)) close!(pool) end + @testset "Pooled context source deregister/register propagation" begin + ctx = QueryContext(":memory:"; pool_size=2) + pool = ctx._pool + + df_v1 = DataFrame(id=1:2, v=[10, 20]) + register!(ctx, "t", df_v1) + + # Hold all pooled connections so updates must be queued for in-use conns. + conn1 = acquire!(pool) + conn2 = acquire!(pool) + + @test nrow(execute(conn1, ctx.config, "SELECT * FROM t")) == 2 + @test nrow(execute(conn2, ctx.config, "SELECT * FROM t")) == 2 + + # Deregister while both conns are checked out. + deregister!(ctx, "t") + release!(pool, conn1) + release!(pool, conn2) + + # Both pooled connections should observe the source removal. + conn1 = acquire!(pool) + conn2 = acquire!(pool) + @test_throws QueryError execute(conn1, ctx.config, "SELECT * FROM t") + @test_throws QueryError execute(conn2, ctx.config, "SELECT * FROM t") + release!(pool, conn1) + release!(pool, conn2) + + # Replace source while both conns are checked out. + register!(ctx, "t", df_v1) + conn1 = acquire!(pool) + conn2 = acquire!(pool) + df_v2 = DataFrame(id=1:3, v=[100, 200, 300]) + register!(ctx, "t", df_v2) + release!(pool, conn1) + release!(pool, conn2) + + # Both connections should now see replacement data, not stale data. + conn1 = acquire!(pool) + conn2 = acquire!(pool) + @test execute(conn1, ctx.config, "SELECT SUM(v) AS s FROM t")[1, :s] == 600 + @test execute(conn2, ctx.config, "SELECT COUNT(*) AS n FROM t")[1, :n] == 3 + release!(pool, conn1) + release!(pool, conn2) + + close!(ctx) + end + + @testset "Pooled source sync concurrency stress" begin + ctx = QueryContext(":memory:"; pool_size=4) + register!(ctx, "t", DataFrame(id=1:2, v=[10, 20])) + + err_lock = ReentrantLock() + errors = Any[] + + readers = [Threads.@spawn begin + for _ in 1:40 + try + df = execute(ctx, "SELECT COUNT(*) AS n FROM t") + if !(df[1, :n] in (2, 3)) + lock(err_lock) do + push!(errors, ArgumentError("Unexpected row count $(df[1, :n]) in concurrency stress test")) + end + end + catch e + lock(err_lock) do + push!(errors, e) + end + end + yield() + end + end for _ in 1:4] + + writer = Threads.@spawn begin + for i in 1:60 + if isodd(i) + register!(ctx, "t", DataFrame(id=1:2, v=[10, 20])) + else + register!(ctx, "t", DataFrame(id=1:3, v=[10, 20, 30])) + end + yield() + end + end + + foreach(wait, readers) + wait(writer) + + @test isempty(errors) + close!(ctx) + end + + @testset "ConnectionPool close lifecycle" begin + pool = ConnectionPool(":memory:"; size=1) + conn = acquire!(pool) + + # Closing with a borrowed connection should mark the pool closed, + # reject new acquisitions, and close borrowed handles on release. + close!(pool) + @test_throws QueryError acquire!(pool) + @test_throws QueryError with_connection(pool) do c + DuckDB.execute(c, "SELECT 1") + end + + @test_nowarn release!(pool, conn) + @test !haskey(pool._in_use, conn) + @test !isready(pool.channel) # closed pool does not re-pool released connections + + # Idempotent close + @test_nowarn close!(pool) + end + # ── File persistence ────────────────────────────────────────────────────── @testset "File-based database" begin db_file = tempname() * ".duckdb" @@ -283,6 +411,24 @@ global_logger(ConsoleLogger(stderr, Logging.Warn)) end end + @testset "Readonly context blocks writes" begin + db_file = tempname() * ".duckdb" + try + ctx_rw = QueryContext(db_file) + execute!(ctx_rw, "CREATE TABLE ro_items (id INT)") + execute!(ctx_rw, "INSERT INTO ro_items VALUES (1)") + close!(ctx_rw) + + ctx_ro = QueryContext(db_file; readonly=true) + df = execute(ctx_ro, "SELECT COUNT(*) AS n FROM ro_items") + @test df[1, :n] == 1 + @test_throws QueryError execute!(ctx_ro, "INSERT INTO ro_items VALUES (2)") + close!(ctx_ro) + finally + isfile(db_file) && rm(db_file) + end + end + # ── Closed context guard ───────────────────────────────────────────────── @testset "Error on closed context" begin ctx = QueryContext() @@ -521,4 +667,28 @@ global_logger(ConsoleLogger(stderr, Logging.Warn)) @test !_needs_sanitization(Union{String, Missing}) end + @testset "Named parameter parser edge cases" begin + ctx = QueryContext() + + # Placeholder-like token inside a string literal must not be bound. + df_lit = execute(ctx, "SELECT ':not_a_param' AS txt, :x AS x", x=7) + @test df_lit[1, :txt] == ":not_a_param" + @test df_lit[1, :x] == 7 + + # Cast operator should not be parsed as a named parameter. + df_cast = execute(ctx, "SELECT 1::INT AS one, :y AS y", y=9) + @test df_cast[1, :one] == 1 + @test df_cast[1, :y] == 9 + + # Comment text containing :tokens must not be bound. + df_comment = execute(ctx, """ + -- :ignored_line_comment + /* :ignored_block_comment */ + SELECT :z AS z + """, z=11) + @test df_comment[1, :z] == 11 + + close!(ctx) + end + end # @testset "QuackSQL.jl"