diff --git a/README.md b/README.md index dc4ea24..312c3d6 100644 --- a/README.md +++ b/README.md @@ -91,6 +91,9 @@ end ### Registering data sources +For most use cases, prefer plain string paths. This keeps registration simple +and readable, and is the recommended default. + ```julia # Single registration register!(ctx, "customers", df_customers) # DataFrame @@ -110,6 +113,55 @@ list_sources(ctx) # DataFrame with name / type / info columns deregister!(ctx, "events") ``` +When you need format-specific read options, use typed source wrappers. + +```julia +using QuackSQL: ParquetSource, CsvSource + +# Parquet with read options +register!(ctx, "logs", ParquetSource( + "logs/*.parquet"; + union_by_name=true, + filename=true, + hive_partitioning=true, +)) + +# CSV with read options +register!(ctx, "events_csv", CsvSource( + "events/*.csv"; + header=true, + delim='|', + nullstr="NA", + sample_size=20_000, +)) +``` + +#### Typed source options + +`ParquetSource(path; ...)` + +- `union_by_name::Bool=false` +- `filename::Bool=false` +- `hive_partitioning::Union{Bool,Nothing}=nothing` +- `compression::Union{Symbol,Nothing}=nothing` + +`compression` is accepted for forward compatibility with future write/export +APIs. It is currently **not** a `read_parquet` option, so setting it during +`register!` raises an `ArgumentError`. + +`CsvSource(path; ...)` + +- `header::Union{Bool,Nothing}=nothing` +- `delim::Union{Char,Nothing}=nothing` +- `quotechar::Union{Char,Nothing}=nothing` +- `escape::Union{Char,Nothing}=nothing` +- `nullstr::Union{String,Nothing}=nothing` +- `auto_detect::Bool=true` +- `sample_size::Union{Int,Nothing}=nothing` + +If you do not need these options, continue to use string paths as the preferred +interface. + --- ### Executing queries diff --git a/src/QuackSQL.jl b/src/QuackSQL.jl index 583589f..c68613f 100644 --- a/src/QuackSQL.jl +++ b/src/QuackSQL.jl @@ -96,7 +96,7 @@ export ConnectionPool, acquire!, release!, with_connection, close! export QueryContext, with_context, close! # Source management -export register!, deregister!, list_sources +export register!, deregister!, list_sources, ParquetSource, CsvSource # Query execution export execute, execute!, query, transaction, stream, explain diff --git a/src/sources.jl b/src/sources.jl index db35720..5f9130c 100644 --- a/src/sources.jl +++ b/src/sources.jl @@ -11,19 +11,82 @@ const PARQUET_EXTS = (".parquet", ".pq") const CSV_EXTS = (".csv", ".tsv", ".csv.gz") const DUCKDB_EXTS = (".duckdb", ".db") +abstract type AbstractPathSource end + """ - ParquetSource(path; union_by_name=false) + ParquetSource(path; union_by_name=false, filename=false, + hive_partitioning=nothing, compression=nothing) + +Wrapper for a Parquet path (or glob) plus `read_parquet` options. -Internal wrapper that pairs a Parquet path (or glob) with read options. Created -automatically by `register!` when keyword options are supplied. +`compression` is accepted for forward compatibility with write/export APIs, +but is not supported by `read_parquet` and therefore causes `register!` to +throw an `ArgumentError` when set. """ -struct ParquetSource +struct ParquetSource <: AbstractPathSource path::String union_by_name::Bool + filename::Bool + hive_partitioning::Union{Bool, Nothing} + compression::Union{Symbol, Nothing} +end + +const PARQUET_COMPRESSION_CODECS = Set([ + :snappy, :zstd, :gzip, :brotli, :lz4, :uncompressed +]) + +function ParquetSource( + path::String; + union_by_name::Bool=false, + filename::Bool=false, + hive_partitioning::Union{Bool, Nothing}=nothing, + compression::Union{Symbol, Nothing}=nothing +) + if compression !== nothing && !(compression in PARQUET_COMPRESSION_CODECS) + throw(ArgumentError( + "Invalid compression codec: $compression. " * + "Expected one of: $(collect(PARQUET_COMPRESSION_CODECS))" + )) + end + return ParquetSource(path, union_by_name, filename, hive_partitioning, compression) +end + +# Backward-compatible positional constructor used by existing tests/users. +ParquetSource(path::String, union_by_name::Bool) = + ParquetSource(path; union_by_name=union_by_name) + +""" + CsvSource(path; header=nothing, delim=nothing, quotechar=nothing, escape=nothing, + nullstr=nothing, auto_detect=true, sample_size=nothing) + +Wrapper for a CSV/TSV path (or glob) plus `read_csv_auto` options. +""" +struct CsvSource <: AbstractPathSource + path::String + header::Union{Bool, Nothing} + delim::Union{Char, Nothing} + quotechar::Union{Char, Nothing} + escape::Union{Char, Nothing} + nullstr::Union{String, Nothing} + auto_detect::Bool + sample_size::Union{Int, Nothing} end -ParquetSource(path::String; union_by_name::Bool=false) = - ParquetSource(path, union_by_name) +function CsvSource( + path::String; + header::Union{Bool, Nothing}=nothing, + delim::Union{Char, Nothing}=nothing, + quotechar::Union{Char, Nothing}=nothing, + escape::Union{Char, Nothing}=nothing, + nullstr::Union{String, Nothing}=nothing, + auto_detect::Bool=true, + sample_size::Union{Int, Nothing}=nothing +) + if sample_size !== nothing && sample_size <= 0 + throw(ArgumentError("sample_size must be positive when provided, got: $sample_size")) + end + return CsvSource(path, header, delim, quotechar, escape, nullstr, auto_detect, sample_size) +end """Return true when `s` looks like a parquet source (extension or glob).""" _looks_like_parquet(s::String) = @@ -245,7 +308,20 @@ Apply a single source to an open DuckDB connection. """ function _register_source!(conn::DuckDB.DB, name::String, source) if source isa ParquetSource - _register_parquet_view!(conn, name, source.path; union_by_name=source.union_by_name) + source.compression === nothing || throw(ArgumentError( + "Parquet compression is a write-time option and is not supported by read_parquet/register!. " * + "Got compression=$(source.compression)." + )) + _register_parquet_view!( + conn, + name, + source.path; + union_by_name=source.union_by_name, + filename=source.filename, + hive_partitioning=source.hive_partitioning, + ) + elseif source isa CsvSource + _register_csv_view!(conn, name, source) elseif source isa DataFrame _register_dataframe!(conn, name, source) elseif source isa String @@ -367,11 +443,37 @@ function _register_csv_view!(conn::DuckDB.DB, name::String, path::String) @debug "CSV view registered" name=name path=path end +function _register_csv_view!(conn::DuckDB.DB, name::String, src::CsvSource) + opts = String[] + src.header !== nothing && push!(opts, "header=$(src.header)") + src.delim !== nothing && push!(opts, "delim='$(escape_sql_string(string(src.delim)))'") + src.quotechar !== nothing && push!(opts, "quote='$(escape_sql_string(string(src.quotechar)))'") + src.escape !== nothing && push!(opts, "escape='$(escape_sql_string(string(src.escape)))'") + src.nullstr !== nothing && push!(opts, "nullstr='$(escape_sql_string(src.nullstr))'") + src.auto_detect != true && push!(opts, "auto_detect=$(src.auto_detect)") + src.sample_size !== nothing && push!(opts, "sample_size=$(src.sample_size)") + + sql_opts = isempty(opts) ? "" : ", " * join(opts, ", ") + DuckDB.execute( + conn, + "CREATE OR REPLACE VIEW \"$(escape_identifier(name))\" AS " * + "SELECT * FROM read_csv_auto('$(escape_sql_string(src.path))'$sql_opts)" + ) + @debug "CSV source registered" name=name path=src.path options=opts +end + function _register_parquet_view!(conn::DuckDB.DB, name::String, path::String; - union_by_name::Bool=false) - opts = union_by_name ? ", union_by_name=true" : "" - DuckDB.execute(conn, "CREATE OR REPLACE VIEW \"$(escape_identifier(name))\" AS SELECT * FROM read_parquet('$(escape_sql_string(path))'$opts)") - @debug "Parquet view registered" name=name path=path union_by_name=union_by_name + union_by_name::Bool=false, + filename::Bool=false, + hive_partitioning::Union{Bool, Nothing}=nothing) + opts = String[] + union_by_name && push!(opts, "union_by_name=true") + filename && push!(opts, "filename=true") + hive_partitioning !== nothing && push!(opts, "hive_partitioning=$(hive_partitioning)") + + sql_opts = isempty(opts) ? "" : ", " * join(opts, ", ") + DuckDB.execute(conn, "CREATE OR REPLACE VIEW \"$(escape_identifier(name))\" AS SELECT * FROM read_parquet('$(escape_sql_string(path))'$sql_opts)") + @debug "Parquet view registered" name=name path=path union_by_name=union_by_name filename=filename hive_partitioning=hive_partitioning end function _attach_database!(conn::DuckDB.DB, name::String, path::String) @@ -434,6 +536,7 @@ escape_sql_string(s::String) = replace(s, "'" => "''") function _source_type_label(src)::String src isa DataFrame && return "DataFrame" src isa ParquetSource && return "Parquet" + src isa CsvSource && return "CSV" src isa String && any(endswith(lowercase(src), e) for e in DUCKDB_EXTS) && return "DuckDB" src isa String && any(endswith(lowercase(src), e) for e in PARQUET_EXTS) && return "Parquet" src isa String && any(endswith(lowercase(src), e) for e in CSV_EXTS) && return "CSV" @@ -443,7 +546,25 @@ end function _source_info(src)::String src isa DataFrame && return "$(nrow(src)) rows × $(ncol(src)) cols" - src isa ParquetSource && return "$(src.path)" * (src.union_by_name ? " (union_by_name=true)" : "") + if src isa ParquetSource + opts = String[] + src.union_by_name && push!(opts, "union_by_name=true") + src.filename && push!(opts, "filename=true") + src.hive_partitioning !== nothing && push!(opts, "hive_partitioning=$(src.hive_partitioning)") + src.compression !== nothing && push!(opts, "compression=$(src.compression)") + return isempty(opts) ? src.path : "$(src.path) (" * join(opts, ", ") * ")" + end + if src isa CsvSource + opts = String[] + src.header !== nothing && push!(opts, "header=$(src.header)") + src.delim !== nothing && push!(opts, "delim=$(src.delim)") + src.quotechar !== nothing && push!(opts, "quote=$(src.quotechar)") + src.escape !== nothing && push!(opts, "escape=$(src.escape)") + src.nullstr !== nothing && push!(opts, "nullstr=$(src.nullstr)") + src.auto_detect != true && push!(opts, "auto_detect=$(src.auto_detect)") + src.sample_size !== nothing && push!(opts, "sample_size=$(src.sample_size)") + return isempty(opts) ? src.path : "$(src.path) (" * join(opts, ", ") * ")" + end src isa String && return src return string(src) end diff --git a/test/runtests.jl b/test/runtests.jl index 6afde53..c19e6eb 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -651,6 +651,139 @@ global_logger(ConsoleLogger(stderr, Logging.Warn)) end end + # ── Phase 1 typed file sources ─────────────────────────────────────────── + @testset "ParquetSource extended options" begin + using QuackSQL: ParquetSource + + src = ParquetSource( + "data/*.parquet"; + union_by_name=true, + filename=true, + hive_partitioning=true, + compression=:zstd, + ) + @test src.path == "data/*.parquet" + @test src.union_by_name == true + @test src.filename == true + @test src.hive_partitioning == true + @test src.compression == :zstd + + @test_throws ArgumentError ParquetSource("x.parquet"; compression=:not_a_codec) + end + + @testset "ParquetSource compression rejected on register!" begin + tmp_dir = mktempdir() + try + pq = joinpath(tmp_dir, "data.parquet") + conn_tmp = DuckDB.DB(":memory:") + DuckDB.execute(conn_tmp, "COPY (SELECT 1 AS id) TO '$pq'") + DuckDB.close(conn_tmp) + + ctx = QueryContext() + src = ParquetSource(pq; compression=:zstd) + @test_throws ArgumentError register!(ctx, "bad_parquet", src) + close!(ctx) + finally + rm(tmp_dir; recursive=true) + end + end + + @testset "ParquetSource filename/hive_partitioning options" begin + using QuackSQL: ParquetSource + + tmp_dir = mktempdir() + try + part_dir = joinpath(tmp_dir, "country=US") + mkpath(part_dir) + pq = joinpath(part_dir, "part.parquet") + + conn_tmp = DuckDB.DB(":memory:") + DuckDB.execute(conn_tmp, "COPY (SELECT 10 AS id) TO '$pq'") + DuckDB.close(conn_tmp) + + ctx = QueryContext() + src = ParquetSource(joinpath(tmp_dir, "**", "*.parquet"); filename=true, hive_partitioning=true) + register!(ctx, "pq_ext", src) + + df = execute(ctx, "SELECT * FROM pq_ext") + @test nrow(df) == 1 + @test "filename" in names(df) + @test "country" in names(df) + + meta = execute(ctx, """ + SELECT view_definition + FROM information_schema.views + WHERE table_name = 'pq_ext' + """) + @test nrow(meta) == 1 + @test occursin("filename", meta[1, :view_definition]) + @test occursin("hive_partitioning", meta[1, :view_definition]) + close!(ctx) + finally + rm(tmp_dir; recursive=true) + end + end + + @testset "CsvSource basic options" begin + using QuackSQL: CsvSource + + tmp_dir = mktempdir() + try + csv_file = joinpath(tmp_dir, "people.csv") + open(csv_file, "w") do io + write(io, "id|name|score\n") + write(io, "1|alice|10\n") + write(io, "2|bob|20\n") + end + + ctx = QueryContext() + src = CsvSource(csv_file; delim='|', header=true) + register!(ctx, "people", src) + + df = execute(ctx, "SELECT * FROM people ORDER BY id") + @test nrow(df) == 2 + @test df[1, :name] == "alice" + @test df[2, :score] == 20 + + sources = list_sources(ctx) + row = sources[sources.name .== "people", :] + @test nrow(row) == 1 + @test row[1, :type] == "CSV" + @test occursin("delim=|", row[1, :info]) + close!(ctx) + finally + rm(tmp_dir; recursive=true) + end + end + + @testset "CsvSource nullstr and sample_size validation" begin + using QuackSQL: CsvSource + + @test_throws ArgumentError CsvSource("x.csv"; sample_size=0) + + tmp_dir = mktempdir() + try + csv_file = joinpath(tmp_dir, "metrics.csv") + open(csv_file, "w") do io + write(io, "id,val\n") + write(io, "1,NA\n") + write(io, "2,7\n") + end + + ctx = QueryContext() + src = CsvSource(csv_file; header=true, nullstr="NA") + register!(ctx, "metrics", src) + + df = execute(ctx, "SELECT * FROM metrics ORDER BY id") + @test nrow(df) == 2 + @test ismissing(df[1, :val]) + @test df[2, :val] == 7 + close!(ctx) + finally + rm(tmp_dir; recursive=true) + end + end + # ── _needs_sanitization helper ──────────────────────────────────────────── @testset "_needs_sanitization" begin using QuackSQL: _needs_sanitization