refactor: improved connection state

.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
+      - 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

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

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"(?<![:\w]):([A-Za-z_]\w*)" => 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

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.

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