Allow asserting non-null columns to avoid missing

src/LibPQ.jl

@@ -578,6 +578,9 @@ mutable struct Result <: Data.Source
     "Julia types for each column in the result"
     column_types::Vector{Type}
 
+    "Whether to expect NULL for each column (whether output data can have `missing`)"
+    not_null::Vector{Bool}
+
     "Conversions from PostgreSQL data to Julia types for each column in the result"
     column_funcs::Vector{Base.Callable}
 
@@ -589,6 +592,7 @@ mutable struct Result <: Data.Source
         column_types::AbstractDict=ColumnTypeMap(),
         type_map::AbstractDict=PQTypeMap(),
         conversions::AbstractDict=PQConversions(),
+        not_null=false,
     )
         jl_result = new(result, cleared)
 
@@ -627,6 +631,33 @@ mutable struct Result <: Data.Source
             func_lookup[(oid, typ)]
         end)
 
+        # figure out which columns the user says may contain nulls
+        if not_null isa Bool
+            jl_result.not_null = fill(not_null, size(col_types))
+        elseif not_null isa AbstractArray
+            if eltype(not_null) === Bool
+                if length(not_null) != length(col_types)
+                    throw(ArgumentError(
+                        "The length of keyword argument not_null, when an array, must be equal to the number of columns"
+                    ))
+                end
+
+                jl_result.not_null = not_null
+            else
+                # assume array of column names
+                jl_result.not_null = fill(false, size(col_types))
+
+                for col_name in not_null
+                    col_num = column_number(jl_result, col_name)
+                    jl_result.not_null[col_num] = true
+                end
+            end
+        else
+            throw(ArgumentError(
+                "Unsupported type $(typeof(not_null)) for keyword argument not_null"
+            ))
+        end
+
         return jl_result
     end
 end

src/datastreams.jl

@@ -5,8 +5,12 @@ function Data.schema(jl_result::Result)
 end
 
 function Data.schema(jl_result::Result, ::Type{Data.Field})
+    types = map(jl_result.not_null, column_types(jl_result)) do not_null, col_type
+        not_null ? col_type : Union{col_type, Missing}
+    end
+
     return Data.Schema(
-        Type[Union{T, Missing} for T in column_types(jl_result)],
+        types,
         column_names(jl_result),
         num_rows(jl_result),
     )
@@ -23,29 +27,56 @@ end
 Data.streamtype(::Type{Result}, ::Type{Data.Field}) = true
 Data.accesspattern(jl_result::Result) = RandomAccess()
 
+# function Data.streamfrom(
+#     jl_result::Result,
+#     ::Type{Data.Field},
+#     ::Type{Union{T, Missing}},
+#     row::Int,
+#     col::Int,
+# )::Union{T, Missing} where T
+#     if libpq_c.PQgetisnull(jl_result.result, row - 1, col - 1) == 1
+#         return missing
+#     else
+#         oid = jl_result.column_oids[col]
+#         return jl_result.column_funcs[col](PQValue{oid}(jl_result, row, col))::T
+#     end
+# end
+
+@inline _non_null_type(::Type{Union{T, Missing}}) where {T} = T
+@inline _non_null_type(::Type{T}) where {T} = T
+
+# allow types that aren't just unions to handle nulls
 function Data.streamfrom(
     jl_result::Result,
     ::Type{Data.Field},
-    ::Type{Union{T, Missing}},
+    ::Type{T},
     row::Int,
     col::Int,
-)::Union{T, Missing} where T
+)::T where T>:Missing
+    NNT = _non_null_type(T)
+
     if libpq_c.PQgetisnull(jl_result.result, row - 1, col - 1) == 1
         return missing
     else
         oid = jl_result.column_oids[col]
-        return jl_result.column_funcs[col](PQValue{oid}(jl_result, row, col))::T
+        return jl_result.column_funcs[col](PQValue{oid}(jl_result, row, col))::NNT
     end
 end
 
+# if a user says they don't want Missing, error on NULL
 function Data.streamfrom(
     jl_result::Result,
     ::Type{Data.Field},
-    ::Type{String},
+    ::Type{T},
     row::Int,
     col::Int,
-)::String
-    unsafe_string(libpq_c.PQgetvalue(jl_result.result, row - 1, col - 1))
+)::T where T
+    if libpq_c.PQgetisnull(jl_result.result, row - 1, col - 1) == 1
+        error("Unexpected NULL at column $col row $row")
+    end
+
+    oid = jl_result.column_oids[col]
+    return jl_result.column_funcs[col](PQValue{oid}(jl_result, row, col))
 end
 
 ### SOURCE END

test/runtests.jl

@@ -315,6 +315,85 @@ end
             close(conn)
         end
 
+        @testset "Not Nulls" begin
+            conn = Connection("dbname=postgres user=$DATABASE_USER"; throw_error=true)
+
+            result = execute(conn, "SELECT NULL"; not_null=[false], throw_error=true)
+            @test status(result) == LibPQ.libpq_c.PGRES_TUPLES_OK
+            @test LibPQ.num_columns(result) == 1
+            @test LibPQ.num_rows(result) == 1
+
+            data = Data.stream!(result, NamedTuple)
+
+            @test data[1][1] === missing
+
+            clear!(result)
+
+            result = execute(conn, "SELECT NULL"; not_null=true, throw_error=true)
+            @test status(result) == LibPQ.libpq_c.PGRES_TUPLES_OK
+            @test LibPQ.num_columns(result) == 1
+            @test LibPQ.num_rows(result) == 1
+
+            @test_throws ErrorException Data.stream!(result, NamedTuple)
+
+            clear!(result)
+
+            result = execute(conn, "SELECT NULL"; not_null=[true], throw_error=true)
+            @test status(result) == LibPQ.libpq_c.PGRES_TUPLES_OK
+            @test LibPQ.num_columns(result) == 1
+            @test LibPQ.num_rows(result) == 1
+
+            @test_throws ErrorException Data.stream!(result, NamedTuple)
+
+            clear!(result)
+
+            result = execute(conn, """
+                SELECT no_nulls, yes_nulls FROM (
+                    VALUES ('foo', 'bar'), ('baz', NULL)
+                ) AS temp (no_nulls, yes_nulls)
+                ORDER BY no_nulls DESC;
+                """;
+                not_null=[true, false],
+                throw_error=true,
+            )
+            @test status(result) == LibPQ.libpq_c.PGRES_TUPLES_OK
+            @test LibPQ.num_rows(result) == 2
+            @test LibPQ.num_columns(result) == 2
+
+            data = Data.stream!(result, NamedTuple)
+
+            @test data[:no_nulls] == ["foo", "baz"]
+            @test data[:no_nulls] isa Vector{String}
+            @test data[:yes_nulls][1] == "bar"
+            @test data[:yes_nulls][2] === missing
+            @test data[:yes_nulls] isa Vector{Union{String, Missing}}
+
+            clear!(result)
+
+            result = execute(conn, """
+                SELECT no_nulls, yes_nulls FROM (
+                    VALUES ('foo', 'bar'), ('baz', NULL)
+                ) AS temp (no_nulls, yes_nulls)
+                ORDER BY no_nulls DESC;
+                """;
+                not_null=false,
+                throw_error=true,
+            )
+            @test status(result) == LibPQ.libpq_c.PGRES_TUPLES_OK
+            @test LibPQ.num_rows(result) == 2
+            @test LibPQ.num_columns(result) == 2
+
+            data = Data.stream!(result, NamedTuple)
+
+            @test data[:no_nulls] == ["foo", "baz"]
+            @test data[:no_nulls] isa Vector{Union{String, Missing}}
+            @test data[:yes_nulls][1] == "bar"
+            @test data[:yes_nulls][2] === missing
+            @test data[:yes_nulls] isa Vector{Union{String, Missing}}
+
+            clear!(result)
+        end
+
         @testset "Type Conversions" begin
             @testset "Automatic" begin
                 conn = Connection("dbname=postgres user=$DATABASE_USER"; throw_error=true)