diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,74 @@ # Changelog +## 0.3.0.0 + +- Use native DuckDB 1.5.6 by default. Keep native support for DuckDB >= 1.5.3 + and < 1.6. Test native versions 1.5.3 through 1.5.6. +- Require `duckdb-ffi >= 1.5.6.0` for the updated native bindings. +- Decode VARIANT results to the `FieldValue` of the stored value, so existing + `FromField` instances read them. Objects decode to `FieldStruct` values with + VARIANT fields, in entry order. The decoder checks DuckDB's private 1.5 + format and payload bounds, and rejects cyclic child references. It has no + depth limit for acyclic values. +- Add `Variant`, a `FieldValue` that binds as a VARIANT, and `variantObject`, + which builds an object payload. Scalars keep their native type. Object + parameters reject duplicate, empty, and NUL-containing keys because the + native constructors cannot represent all of these names. + `Variant` has no `ToDuckValue` instance, and `logicalTypeFromRep` raises an + error for VARIANT. +- A VARIANT timestamp payload outside the microsecond range binds as a + millisecond or second timestamp, so it keeps its value. +- Read the VARIANT type and GEOMETRY types for a list of CRS definitions, + because the C API cannot create them. A connection reads them with one + query on a separate connection, the first time a parameter needs one. + Parameters use these types. Add `ConnectionOptions`, `openWithOptions`, and + `withConnectionWithOptions` to set the CRS list, which defaults to + `OGC:CRS84`. +- Array parameters use the element's `ToField` instance. Elements require + `ToField` and `DuckDBColumnType`. Add a `ToDuckValue` instance for arrays + whose elements have `ToDuckValue`. It does not need a connection. +- Array parameters of STRUCT values, UNION values, generic records, and + arrays bind. They take the element type of the first element that is not + NULL. The other non-NULL elements must have the same type, including field + names, decimal precision and scale, and nested types. Different types raise + an error, so DuckDB cannot cast away fields or round values. An empty or + all-NULL array of these elements raises an error. Scalar elements keep the + type of their column type name. +- A GEOMETRY payload decodes to `FieldGeometry` with raw WKB. Import those + bytes with `ST_GeomFromWKB(?)::VARIANT`. +- Add parameter and result instances for `Data.Geometry.Geometry` from + `geometry-simple`. It provides decoded shapes, unboxed coordinate vectors, + and runtime coordinate layouts. This type does not store CRS metadata. + `RawGeometry` retains WKB and CRS without decoding coordinates. Existing + `ByteString` results still return WKB. Empty points remain distinct from SQL NULL. +- Use `geometry-simple >= 0.1.1.0` for pure geometry types and codecs. + Structured parameters use its WKT writer and DuckDB's native cast. + `RawGeometry` has no parameter instance. Import its WKB with + `ST_GeomFromWKB` and apply its CRS with `ST_SetCRS`. This preserves mixed + layouts, dimensional empties, and native NaN payloads without WKT conversion. +- Add `FieldGeometry` and `LogicalTypeGeometry` to the public value and type + representations. Update exhaustive matches when upgrading. + Nested result decoding retains raw WKB and CRS. Generic parameters with + non-NULL `FieldGeometry` values raise an error instead of converting their + bytes through WKT. Construct those nested values with explicit SQL import. +- Keep CRS metadata in this package. The standalone shape has no CRS. + Both forms own their memory and remain usable after the connection closes. + CRS text can hold an identifier, a custom name, or a full WKT2/PROJJSON + definition. +- Composite parameters keep GEOMETRY CRS metadata for the CRS definitions + that the connection read. Other CRS definitions bind without a CRS, and + `logicalTypeFromRep` creates `GEOMETRY` with no CRS. Insert the value into a + column with a CRS, or cast it in SQL, to apply the CRS. +- Read STRUCT and UNION type metadata once for each result vector instead of + once for each row. Prepare their direct child readers at the same time. + A local benchmark with a CRS-tagged GEOMETRY field ran about ten times faster. + LIST, ARRAY, and MAP decoders can still read child metadata for each parent row. +- `Database.DuckDB.Simple.Internal` changes. `StatementStreamChunk` keeps one + reader for each column in `statementStreamChunkReaders`, and + `StatementStreamChunkVector` is removed. The module also exports + `withTypeCache`, `fetchPrepareError`, `duckDBTypeFromName`, and + `duckDBTypeToName`. + ## 0.2.0.0 ### Query execution and resource lifetime diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -183,17 +183,29 @@ storeArray :: Connection -> IO [Array Int Int] storeArray conn = do _ <- execute_ conn "CREATE TABLE arrays (vals INTEGER[3])" - let arr = listArray (0, 2) [1, 2, 3] + let arr = listArray (0, 2) [1, 2, 3] :: Array Int Int _ <- execute conn "INSERT INTO arrays VALUES (?)" (Only arr) fmap fromOnly <$> query_ conn "SELECT vals FROM arrays" storeList :: Connection -> IO [[Int]] storeList conn = do _ <- execute_ conn "CREATE TABLE lists (vals INTEGER[])" - _ <- execute conn "INSERT INTO lists VALUES (?)" (Only [1, 2, 3]) + let arr = listArray (0, 2) [1, 2, 3] :: Array Int Int + _ <- execute conn "INSERT INTO lists VALUES (?)" (Only arr) fmap fromOnly <$> query_ conn "SELECT vals FROM lists" ``` +A Haskell list reads a LIST result. Lists have no parameter instance, so bind +an `Array`. DuckDB casts the array to the LIST column type. + +An array parameter with scalar elements takes its element type from the +column type name of the element, so an empty array keeps its type. An array +of STRUCT values, UNION values, generic records, or arrays takes the type of +its first element that is not NULL. All other non-NULL elements must have the +same type, including field names, decimal precision and scale, and nested +types. Different types raise an error before DuckDB can cast away fields or +round values. An empty or all-NULL array of such elements raises an error. + ### Infinite dates and timestamps Use `Database.DuckDB.Simple.Time` when a column can contain temporal infinity. @@ -237,6 +249,7 @@ [(s, u)] <- query_ conn "SELECT {'a': 1, 'b': 2}, \ \CAST(union_value(x := 42) AS UNION(x INT, y VARCHAR))" + :: IO [(StructValue FieldValue, UnionValue FieldValue)] _ <- execute conn "INSERT INTO composite VALUES (?, ?)" (s, u) query_ conn "SELECT s, u FROM composite" ``` @@ -298,7 +311,8 @@ `Statement` to pull rows one at a time and decide when to stop. Cursors support the same column types as eager queries, including STRUCT -and UNION values with nested collections and NULLs. +and UNION values with nested collections and NULLs. VARIANT and GEOMETRY +work with eager queries, cursors, and folds. #### Optional native streaming @@ -365,6 +379,167 @@ result size. The older Arrow query and scan bindings remain available through `Database.DuckDB.FFI.Deprecated` and emit deprecation warnings. +### VARIANT + +A VARIANT result decodes to the `FieldValue` of the stored value, with its +native type. The usual `FromField` instances read VARIANT columns, for example +as `Int64`, `Text`, `[a]`, or a generic record. Arrays decode to `FieldList`. +Objects decode to `FieldStruct` values whose fields have the VARIANT type, in +entry order. `variantObject` builds such a payload from a list of entries. +SQL NULL decodes to `FieldNull`. + +`Variant` from `Database.DuckDB.Simple.Variant` wraps a `FieldValue`. Its +`FromField` instance reads any column. Its `ToField` instance binds the +payload as a VARIANT, so `SELECT ?` returns a VARIANT: + +```haskell +query conn "SELECT ?" (Only (Variant (FieldList [FieldInt8 1, FieldText "two"]))) +``` + +A scalar payload keeps its native type, such as `TINYINT` or `DECIMAL(4,2)`. +Time and timestamp payloads bind as microsecond types, or as nanosecond types +when they have sub-microsecond digits. So a `TIMESTAMP_S` result binds back as +a `TIMESTAMP` when its value fits. Wider timestamp values use milliseconds or +seconds without losing digits. Lists, arrays, and STRUCT fields bind as +VARIANT values. MAP and ENUM payloads raise an error. Object parameters reject +duplicate keys, empty keys, and keys that contain NUL. You can also bind a +plain value and cast it in SQL, as in `?::VARIANT`. + +The C API cannot create a usable VARIANT type +([#27](https://github.com/Tritlo/duckdb-haskell/issues/27)). Each connection +reads the VARIANT type the first time a parameter needs it, and its +parameters use that type. The connection reads the type with a query on a +separate connection. So the query does not run in your transaction, and it +does not end a streaming result. +`Variant` has no `ToDuckValue` instance, and `logicalTypeFromRep` raises an +error for VARIANT, because neither has a connection. + +A GEOMETRY payload decodes to `FieldGeometry` with raw WKB and no CRS. Import +these bytes with `ST_GeomFromWKB(?)::VARIANT`. A `Variant` parameter that +contains `FieldGeometry` raises an error, also inside arrays and objects. A +TIMETZ payload with an offset in seconds raises an error, as a TIMETZ column +does. Results can contain empty or NUL object keys. Text and blob payloads can +contain NUL. + +DuckDB 1.5 has no C API for reading VARIANT values. The decoder checks the +native version and physical schema before it reads the internal representation. +It checks payload bounds and rejects unknown tags. This format dependency is +limited to the supported DuckDB 1.5 line. + +Persistent VARIANT columns require storage format `v1.5.0` or later. For a new +database, pass `[("storage_compatibility_version", "v1.5.0")]` to +`openWithConfig` or `withConnectionWithConfig`. The library does not change an +existing database's storage compatibility setting. + +### GEOMETRY + +Use `Geometry` from `Data.Geometry` for decoded shapes. The `geometry-simple` +package supplies the type and unboxed coordinate vectors. `duckdb-simple` +supplies its parameter and result instances: + +```haskell +import qualified Data.Geometry as G +import qualified Data.Vector.Unboxed as U +import Database.DuckDB.Simple + +let line = G.LineString (G.CoordinatesXY (U.fromList [G.XY 0 0, G.XY 1 2, G.XY 3 4])) +rows <- query conn "SELECT ?" (Only line) :: IO [Only G.Geometry] +``` + +Points and coordinate sequences store their XY, XYZ, XYM, or XYZM layout. +`G.EmptyPoint G.DimXYZ` is an empty XYZ point. Use `Maybe G.Geometry` for SQL NULL. +Use `decodeWKT` from `Data.Geometry.WKT` to parse text without a database. + +`G.Geometry` has no CRS metadata. Reading it returns the coordinates without +their CRS label. Binding it creates a `GEOMETRY` with no CRS. To attach a label, +use `ST_SetCRS(?, ?)` with a shape and CRS text. This does not transform coordinates. + +Use `RawGeometry` from `Database.DuckDB.Simple.Geometry` to read WKB and CRS +without decoding coordinates. Its `rawGeometryWKB` and `rawGeometryCRS` fields +own their data. They remain usable after the connection closes. Existing +`ByteString` result decoding also returns WKB. + +`RawGeometry` has no `ToField` instance. Import its bytes explicitly: + +```haskell +import Database.DuckDB.Simple.Geometry + +[Only raw] <- query_ conn "SELECT 'POINT Z (1 2 3)'::GEOMETRY('OGC:CRS84')" +rows <- (case rawGeometryCRS raw of + Nothing -> query conn + "SELECT system.main.ST_GeomFromWKB(?)" + (Only (rawGeometryWKB raw)) + Just crs -> query conn + "SELECT system.main.ST_SetCRS(system.main.ST_GeomFromWKB(?), ?)" + (rawGeometryWKB raw, crs) + ) :: IO [Only RawGeometry] +``` + +Use `Nothing` for no CRS and omit `ST_SetCRS` in that case. Passing SQL NULL +to `ST_SetCRS` returns a NULL geometry. CRS text can contain `OGC:CRS84`, +a custom name, or a full WKT2/PROJJSON definition. DuckDB can reduce a known +CRS definition to its registered identifier. It can also normalize WKB byte +order. The qualified function names select DuckDB's built-ins even if a user +macro has the same name. + +`fromRawGeometry` decodes the shape and drops CRS metadata. `toRawGeometry` +encodes a shape with no CRS. These helpers follow the geometry-simple contracts. +NaN in both WKB point X and Y denotes an empty point. Empty multi-geometries +and collections have no stored layout tag in the decoded representation. +Keep the raw bytes when these details must survive. + +Structured parameters use `encodeWKT` and DuckDB's native cast. DuckDB 1.5 +has no WKB value constructor in its C API. The WKT writer combines layouts +in multi-geometries and polygon rings. It fills absent Z or M with NaN. +DuckDB's WKT parser limits nesting to 16 levels and rejects empty polygon rings +and mixed collection layouts. Explicit WKB import preserves mixed member +layouts and native NaN payload bits that WKT cannot retain. + +`FieldGeometry` contains a raw result. Generic STRUCT and UNION parameters +that contain non-NULL raw geometry raise an error. Construct those values in +SQL with explicit WKB import. Their result decoding retains WKB and CRS, +including inside LIST, ARRAY, MAP, STRUCT, and UNION values. + +`LogicalTypeGeometry` describes CRS metadata. The C API cannot create a +GEOMETRY type with a CRS. So each connection reads a GEOMETRY type for each +CRS in a list, together with the VARIANT type, the first time a parameter +needs one of them. The default list holds +`OGC:CRS84`. Composite parameters use these types, so typed NULLs, empty +collections, and inactive UNION members keep their CRS. Set the list in +`ConnectionOptions`: + +```haskell +let options = defaultConnectionOptions{connectionGeometryCRS = ["OGC:CRS84", "EPSG:3857"]} +withConnectionWithOptions "shapes.duckdb" options \conn -> ... +``` + +A list entry can be an identifier, a custom name, or a full WKT2 or PROJJSON +definition. A composite parameter with a CRS that is not in the list binds its +geometry members without a CRS. `logicalTypeFromRep` has no connection, so it +creates `GEOMETRY` without a CRS. To apply a CRS in these cases, insert the +value into a column with that CRS, or cast it in SQL: + +```haskell +query conn "SELECT ?::UNION(number BIGINT, shape GEOMETRY('OGC:CRS84'))" (Only value) +``` + +The CRS in a cast must be a constant. DuckDB rejects a parameter as a type +modifier, so `?::GEOMETRY(?)` is not valid. To use a CRS that is known only at +run time, write it into the query text as a SQL string literal, and double each +single quote. A cast accepts a CRS that DuckDB recognizes, such as `OGC:CRS84`, +or a full WKT2 or PROJJSON definition. Unless an extension recognizes it, +DuckDB rejects other identifiers, such as `EPSG:4326`, and custom names. +`ST_SetCRS` also accepts custom names, but it returns `GEOMETRY` with no CRS +for a NULL input. + +A bound composite without a CRS also changes the type of expressions that +combine it with CRS data. `UNION ALL` and `COALESCE` of `GEOMETRY` and +`GEOMETRY('OGC:CRS84')` give `GEOMETRY` with no CRS for all rows. Cast the +parameter before you combine it with other values. + +See [geometry-simple](https://github.com/Tritlo/geometry-simple) for the seven +supported families, construction checks, and codec normalization rules. + ### Feature Coverage - Connections, prepared statements, positional/named parameter binding. @@ -375,7 +550,7 @@ decimals (with width/scale), intervals, precise and timezone-aware temporals, enums, bit strings, blobs, bignums, and UUIDs. - Composite types: STRUCTs, UNIONs, LISTs, fixed-length ARRAYs, and MAPs with - full encoding/decoding support. + typed parameters and results. - Generic encoding/decoding: automatic STRUCT/UNION mapping for Haskell ADTs via GHC generics and the `ViaDuckDB` deriving-via helper. - Row decoding via `FromField`/`FromRow`, with generic deriving for product types. diff --git a/bench/Main.hs b/bench/Main.hs --- a/bench/Main.hs +++ b/bench/Main.hs @@ -6,12 +6,17 @@ import Control.Exception (evaluate) import Control.Monad (forM_, replicateM, unless) +import qualified Data.ByteString as BS +import qualified Data.Geometry as G import Data.Int (Int64) import qualified Data.List as List import qualified Data.Text as Text import Data.Time.Clock.POSIX (utcTimeToPOSIXSeconds) import Data.Time.LocalTime (LocalTime, localTimeToUTC, utc) import Database.DuckDB.Simple +import Database.DuckDB.Simple.FromField (FieldValue (..)) +import Database.DuckDB.Simple.Geometry (RawGeometry (..)) +import Database.DuckDB.Simple.Variant (Variant (..)) import GHC.Clock (getMonotonicTimeNSec) import System.Environment (getArgs) import System.Mem (performMajorGC) @@ -26,6 +31,8 @@ expected = count * (count - 1) `div` 2 withConnectionWithConfig ":memory:" [("threads", "1")] $ \conn -> do createFunction conn "bench_identity" (id :: Int64 -> Int64) + [Only geometry] <- query_ conn "SELECT 'POINT (1 2)'::GEOMETRY('OGC:CRS84')" :: IO [Only RawGeometry] + let typedGeometry = G.PointGeometry (G.PointXY (G.XY 1 2)) let sql = Query ("SELECT i FROM range(" <> Text.pack (show count) <> ") t(i)") action = case workload of "eager" -> do @@ -44,9 +51,32 @@ "parameters" -> do rows <- replicateM (fromIntegral count) (query conn "SELECT ?::BIGINT" (Only (1 :: Int64))) evaluate (sum [n | [Only n] <- rows]) - _ -> fail "Expected eager, fold, scalar, text, timestamp, or parameters" + "variant" -> do + rows <- query_ conn (Query ("SELECT i::VARIANT FROM range(" <> Text.pack (show count) <> ") t(i)")) + evaluate (List.foldl' (\acc (Only value) -> acc + variantNumber value) 0 rows) + "variant-parameters" -> do + rows <- replicateM (fromIntegral count) (query conn "SELECT ?" (Only (Variant (FieldInt64 1)))) + evaluate (sum [variantNumber value | [Only value] <- rows]) + "geometry" -> do + rows <- query_ conn (Query ("SELECT 'POINT (1 2)'::GEOMETRY('OGC:CRS84') FROM range(" <> Text.pack (show count) <> ")")) + evaluate (sum [fromIntegral (BS.length (rawGeometryWKB value)) | Only value <- rows]) + "geometry-parameters" -> do + rows <- replicateM (fromIntegral count) (query conn "SELECT system.main.ST_SetCRS(system.main.ST_GeomFromWKB(?), ?)" (rawGeometryWKB geometry, rawGeometryCRS geometry)) + evaluate (sum [fromIntegral (BS.length (rawGeometryWKB value)) | [Only value] <- rows]) + "geometry-typed" -> do + rows <- query_ conn (Query ("SELECT 'POINT (1 2)'::GEOMETRY('OGC:CRS84') FROM range(" <> Text.pack (show count) <> ")")) + evaluate (sum [pointChecksum value | Only value <- rows]) + "geometry-typed-parameters" -> do + rows <- replicateM (fromIntegral count) (query conn "SELECT ?" (Only typedGeometry)) + evaluate (sum [pointChecksum value | [Only value] <- rows]) + _ -> fail "Expected eager, fold, scalar, text, timestamp, parameters, variant, variant-parameters, geometry, geometry-parameters, geometry-typed, or geometry-typed-parameters" expectedResult = case workload of "parameters" -> count + "variant-parameters" -> count + "geometry" -> count * 21 + "geometry-parameters" -> count * 21 + "geometry-typed" -> count * 3 + "geometry-typed-parameters" -> count * 3 "text" -> count * fromIntegral (Text.length (Text.replicate 4 "duckdb λ text")) "timestamp" -> count * 946684800 + expected _ -> expected @@ -60,3 +90,13 @@ end <- getMonotonicTimeNSec check result printf "%s,%d,%d,%.3f,%d\n" workload count run (fromIntegral (end - start) / 1000000 :: Double) result + +-- | Fail if the VARIANT benchmark changes its scalar type. +variantNumber :: Variant -> Int64 +variantNumber (Variant (FieldInt64 value)) = value +variantNumber value = error ("unexpected VARIANT benchmark value: " <> show value) + +-- | Force coordinate decoding and check the point shape. +pointChecksum :: G.Geometry -> Int64 +pointChecksum (G.PointGeometry (G.PointXY (G.XY x y))) = round (x + y) +pointChecksum value = error ("unexpected geometry benchmark value: " <> show value) diff --git a/duckdb-simple.cabal b/duckdb-simple.cabal --- a/duckdb-simple.cabal +++ b/duckdb-simple.cabal @@ -1,6 +1,6 @@ cabal-version: 3.4 name: duckdb-simple -version: 0.2.0.0 +version: 0.3.0.0 license: MPL-2.0 license-file: LICENSE author: Matthias Pall Gissurarson @@ -22,7 +22,7 @@ prepared statements, chunked folds, transactions, and Haskell scalar functions. The API follows the style of sqlite-simple and postgresql-simple. . - Supports native DuckDB >= 1.5.3 and < 1.6. Tested with version 1.5.3. + Supports native DuckDB >= 1.5.3 and < 1.6. Tested with versions 1.5.3 through 1.5.6. extra-doc-files: CHANGELOG.md @@ -51,6 +51,7 @@ Database.DuckDB.Simple.FromRow Database.DuckDB.Simple.Function Database.DuckDB.Simple.Generic + Database.DuckDB.Simple.Geometry Database.DuckDB.Simple.Internal Database.DuckDB.Simple.Logging Database.DuckDB.Simple.LogicalRep @@ -59,12 +60,16 @@ Database.DuckDB.Simple.ToField Database.DuckDB.Simple.ToRow Database.DuckDB.Simple.Types + Database.DuckDB.Simple.Variant other-modules: Database.DuckDB.Simple.Arrow.Internal Database.DuckDB.Simple.Callback + Database.DuckDB.Simple.Element Database.DuckDB.Simple.Materialize Database.DuckDB.Simple.Result + Database.DuckDB.Simple.TypeCache + Database.DuckDB.Simple.VariantCodec hs-source-dirs: src default-language: Haskell2010 @@ -73,7 +78,8 @@ base >=4.14 && <5, bytestring >=0.11 && <0.13, containers >=0.6 && <0.9, - duckdb-ffi >=1.5.3.0 && <1.6, + duckdb-ffi >=1.5.6.0 && <1.6, + geometry-simple >=0.1.1.0 && <0.2, text >=2.0 && <2.2, time >=1.12 && <1.16, transformers >=0.6 && <0.7, @@ -88,10 +94,13 @@ CancellationTests CoreRegressionTests ExtensionRegressionTests + GeometryRegressionTests Properties StreamingTests + TestUtils TimeTests ValueRegressionTests + VariantRegressionTests default-language: Haskell2010 ghc-options: @@ -104,8 +113,9 @@ bytestring, containers >=0.6 && <0.9, directory >=1.3 && <1.4, - duckdb-ffi >=1.5.3.0 && <1.6, + duckdb-ffi >=1.5.6.0 && <1.6, duckdb-simple, + geometry-simple >=0.1.1.0 && <0.2, QuickCheck >=2.14 && <2.18, tasty >=1.4 && <1.6, tasty-expected-failure >=0.12 && <0.13, @@ -114,7 +124,42 @@ text, time, uuid >=1.3 && <1.4, + vector >=0.13 && <0.14, +test-suite duckdb-simple-codec-test + type: exitcode-stdio-1.0 + hs-source-dirs: + test-codec + src + + main-is: Main.hs + other-modules: + Database.DuckDB.Simple.Element + Database.DuckDB.Simple.FromField + Database.DuckDB.Simple.Geometry + Database.DuckDB.Simple.Internal + Database.DuckDB.Simple.LogicalRep + Database.DuckDB.Simple.Ok + Database.DuckDB.Simple.Time + Database.DuckDB.Simple.TypeCache + Database.DuckDB.Simple.Variant + Database.DuckDB.Simple.VariantCodec + + default-language: Haskell2010 + build-depends: + array >=0.5 && <0.6, + base >=4.14 && <5, + bytestring >=0.11 && <0.13, + containers >=0.6 && <0.9, + duckdb-ffi >=1.5.6.0 && <1.6, + geometry-simple >=0.1.1.0 && <0.2, + tasty >=1.4 && <1.6, + tasty-hunit >=0.10 && <0.12, + text >=2.0 && <2.2, + time >=1.12 && <1.16, + transformers >=0.6 && <0.7, + uuid >=1.3 && <1.4, + test-suite duckdb-simple-leak-test type: exitcode-stdio-1.0 hs-source-dirs: leaktest @@ -123,7 +168,9 @@ ghc-options: -threaded build-depends: base >=4.14 && <5, + bytestring, duckdb-simple, + geometry-simple >=0.1.1.0 && <0.2, test-suite duckdb-simple-dataframe-test type: exitcode-stdio-1.0 @@ -138,7 +185,7 @@ base >=4.14 && <5, dataframe-arrow-bridge >=1.0 && <1.1, dataframe-core >=2.5 && <2.6, - duckdb-ffi >=1.5.3.0 && <1.6, + duckdb-ffi >=1.5.6.0 && <1.6, duckdb-simple, tasty >=1.4 && <1.6, tasty-hunit >=0.10 && <0.12, @@ -156,6 +203,8 @@ build-depends: base >=4.14 && <5, + bytestring, duckdb-simple, + geometry-simple >=0.1.1.0 && <0.2, text, time, diff --git a/leaktest/Main.hs b/leaktest/Main.hs --- a/leaktest/Main.hs +++ b/leaktest/Main.hs @@ -20,6 +20,8 @@ import Control.Concurrent (forkFinally, killThread, newEmptyMVar, putMVar, takeMVar, threadDelay, tryPutMVar) import Control.Exception (AsyncException (ThreadKilled), IOException, SomeException, evaluate, fromException, try) import Control.Monad (forM_, replicateM, unless, void, when) +import qualified Data.ByteString as BS +import qualified Data.Geometry as G import Data.IORef (atomicModifyIORef', atomicWriteIORef, mkWeakIORef, newIORef, readIORef) import Data.Int (Int64) import Data.List (isPrefixOf) @@ -27,8 +29,10 @@ import Database.DuckDB.Simple import qualified Database.DuckDB.Simple.Copy as Copy import qualified Database.DuckDB.Simple.Deprecated.Streaming as Streaming -import Database.DuckDB.Simple.FromField (FieldValue) +import Database.DuckDB.Simple.FromField (DecimalValue (..), FieldValue (..)) +import Database.DuckDB.Simple.Geometry (RawGeometry (..), fromRawGeometry) import qualified Database.DuckDB.Simple.Logging as Logging +import Database.DuckDB.Simple.Variant (Variant (..)) import System.Environment (getArgs, lookupEnv) import System.Exit (exitFailure) import System.Mem (performMajorGC) @@ -81,12 +85,14 @@ stmt <- openStatement conn "SELECT 42" closeStatement stmt _ <- query_ conn "SELECT 42" :: IO [Only Int64] - -- VARIANT decoding fails after DuckDB allocates the materialized result. + -- The first VARIANT parameter fills the type cache. Close must destroy it. + _ <- query conn "SELECT ?" (Only (Variant (FieldInt64 42))) :: IO [Only Variant] + -- A TIMETZ offset with seconds fails during result materialization. -- The result and connection must still be destroyed. - rejected <- try (query_ conn "SELECT i::VARIANT FROM range(100000) t(i)") :: IO (Either SomeException [Only FieldValue]) + rejected <- try (query_ conn "SELECT {'xs': [i, i + 1, i + 2, i + 3], 'bad': '12:00:00+01:23:45'::TIMETZ} FROM range(100000) t(i)") :: IO (Either SomeException [Only FieldValue]) case rejected of Left _ -> pure () - Right _ -> fail "expected unsupported VARIANT conversion" + Right _ -> fail "expected TIMETZ offset conversion failure" close conn main :: IO () @@ -97,14 +103,38 @@ [name] -> (name, 10) [name, n] | Just count <- readMaybe n, count > 0 -> (name, count) _ -> ("invalid", 0) - unless (mode `elem` ["all", "open-close", "long-lived", "callbacks", "cancel", "decode-failure"]) $ - fail "Expected [all|open-close|long-lived|callbacks|cancel|decode-failure] [positive batch count]" + unless (mode `elem` ["all", "open-close", "long-lived", "callbacks", "cancel", "decode-failure", "new-types"]) $ + fail "Expected [all|open-close|long-lived|callbacks|cancel|decode-failure|new-types] [positive batch count]" when (mode `elem` ["all", "open-close"]) checkOpenClose when (mode `elem` ["all", "long-lived"]) (checkLongLived batches) when (mode `elem` ["all", "callbacks"]) (checkCallbacks batches) when (mode `elem` ["all", "cancel"]) (checkCancellation batches) when (mode == "decode-failure") checkDecodeFailure + when (mode `elem` ["all", "new-types"]) (checkNewTypes batches) +-- | Import raw WKB and release native values after failed nested construction. +checkNewTypes :: Int -> IO () +checkNewTypes batches = + withConnectionWithConfig ":memory:" [("threads", "1")] \conn -> do + [Only geometry] <- query_ conn "SELECT 'POINT ZM (1 2 3 4)'::GEOMETRY('OGC:CRS84')" :: IO [Only RawGeometry] + typed <- either fail pure (fromRawGeometry geometry) :: IO G.Geometry + let variant = Variant (FieldList [FieldInt64 42, FieldList [FieldNull, FieldText "before\0after"]]) + batch = forM_ [1 .. 100 :: Int] \_ -> do + rows <- query conn "SELECT ?, system.main.ST_SetCRS(system.main.ST_GeomFromWKB(?), ?), ?" (variant, rawGeometryWKB geometry, rawGeometryCRS geometry, typed) + unless (rows == [(variant, geometry, typed)]) (fail "new type round trip failed") + expectFailure (query conn "SELECT ?, system.main.ST_GeomFromWKB(?)" (typed, BS.pack [1, 1, 0, 0, 0]) :: IO [(G.Geometry, RawGeometry)]) + expectFailure (query conn "SELECT system.main.ST_GeomFromWKB(?)" (Only (rawGeometryWKB geometry)) :: IO [Only Int64]) + expectFailure (query conn "SELECT ?" (Only (Variant (FieldList [variantPayload variant, FieldDecimal (DecimalValue 0 0 0)]))) :: IO [Only Variant]) + expectFailure (query conn "SELECT ?" (Only variant) :: IO [Only Int64]) + batch + performMajorGC + before <- readUsage + forM_ [1 .. batches] \n -> do + batch + performMajorGC + after <- readUsage + reportOptional ("open connection, " <> show (n * 100) <> " new type cycles") before after + -- | Check native results and database handles across connection lifetimes. checkOpenClose :: IO () checkOpenClose = do @@ -129,7 +159,7 @@ let cycleQuery = do rows <- query_ conn "SELECT sum(i)::BIGINT FROM range(10000) t(i)" unless (rows == [Only (49995000 :: Int64)]) (fail "wrong query result") - expectFailure (query_ conn "SELECT i::VARIANT FROM range(10000) t(i)" :: IO [Only FieldValue]) + expectFailure (query_ conn "SELECT {'xs': [i, i + 1, i + 2, i + 3], 'bad': '12:00:00+01:23:45'::TIMETZ} FROM range(10000) t(i)" :: IO [Only FieldValue]) expectFailure (query_ conn "SELECT CAST('bad' AS BIGINT)" :: IO [Only Int64]) bind stmt [toField (42 :: Int64)] nextRow stmt >>= \row -> unless (row == Just (Only (42 :: Int64))) (fail "wrong cursor result") @@ -149,7 +179,7 @@ checkDecodeFailure :: IO () checkDecodeFailure = withConnectionWithConfig ":memory:" [("threads", "1")] \conn -> do - let rejected = expectFailure (query_ conn "SELECT i::VARIANT FROM range(100000) t(i)" :: IO [Only FieldValue]) + let rejected = expectFailure (query_ conn "SELECT {'xs': [i, i + 1, i + 2, i + 3], 'bad': '12:00:00+01:23:45'::TIMETZ} FROM range(100000) t(i)" :: IO [Only FieldValue]) rejected performMajorGC before <- readUsage @@ -271,6 +301,12 @@ if arrow then Streaming.foldArrow_ conn sql () (\() _ _ -> delivered) else Streaming.fold_ conn sql () (\() (Only (_ :: Int64)) -> delivered) + cancel \signal -> do + signal + void (query_ conn "SELECT {'x': i, 'values': [i, i + 1]}::VARIANT FROM range(100000) t(i)" :: IO [Only Variant]) + cancel \signal -> do + blocked <- newEmptyMVar + void (fold_ conn "SELECT 'POINT (1 2)'::GEOMETRY('OGC:CRS84') FROM range(100000)" (0 :: Int64) (\n (Only (_ :: RawGeometry)) -> signal >> takeMVar blocked >> pure (n + 1))) batch performMajorGC before <- readUsage @@ -278,7 +314,7 @@ batch performMajorGC after <- readUsage - reportOptional ("open connection, " <> show (n * 50) <> " cancellations") before after + reportOptional ("open connection, " <> show (n * 70) <> " cancellations") before after -- | Report native counters when the operating system provides them. reportOptional :: String -> Maybe Usage -> Maybe Usage -> IO () diff --git a/src/Database/DuckDB/Simple.hs b/src/Database/DuckDB/Simple.hs --- a/src/Database/DuckDB/Simple.hs +++ b/src/Database/DuckDB/Simple.hs @@ -16,9 +16,13 @@ Connection, open, openWithConfig, + ConnectionOptions (..), + defaultConnectionOptions, + openWithOptions, close, withConnection, withConnectionWithConfig, + withConnectionWithOptions, -- * Queries and statements Query (..), @@ -112,6 +116,7 @@ Statement (..), StatementState (..), StatementStreamState (..), + fetchPrepareError, keepAlive, peekUtf8CString, runInterruptibleQuery, @@ -125,6 +130,7 @@ import qualified Database.DuckDB.Simple.Result as Result import Database.DuckDB.Simple.ToField (DuckDBColumnType (..), FieldBinding, NamedParam (..), ToField (..), bindFieldBinding, duckdbColumnType, renderFieldBinding) import Database.DuckDB.Simple.ToRow (ToRow (..)) +import Database.DuckDB.Simple.TypeCache (TypeCache, createTypeCache, defaultGeometryCRS, destroyTypeCache) import Database.DuckDB.Simple.Types (FormatError (..), Null (..), Only (..), (:.) (..)) import Foreign.C.String (CString) import Foreign.Marshal.Alloc (alloca) @@ -137,16 +143,45 @@ -- | Open a DuckDB database with configuration flags applied before startup. openWithConfig :: FilePath -> [(Text, Text)] -> IO Connection -openWithConfig path settings = +openWithConfig path settings = openWithOptions path defaultConnectionOptions{connectionConfig = settings} + +-- | Settings that duckdb-simple applies when it opens a connection. +data ConnectionOptions = ConnectionOptions + { connectionConfig :: [(Text, Text)] + -- ^ DuckDB configuration flags, as for 'openWithConfig'. + , connectionGeometryCRS :: [Text] + {- ^ CRS definitions. The connection reads a GEOMETRY type for each CRS + the first time a parameter needs a VARIANT type or a GEOMETRY type with a + CRS. Parameters can then contain GEOMETRY types with these CRSs. + -} + } + deriving (Eq, Show) + +-- | No configuration flags, and the CRS @OGC:CRS84@. +defaultConnectionOptions :: ConnectionOptions +defaultConnectionOptions = + ConnectionOptions + { connectionConfig = [] + , connectionGeometryCRS = defaultGeometryCRS + } + +{- | Open a DuckDB database with options. The connection reads the VARIANT type +and the GEOMETRY types for the configured CRSs the first time a parameter +needs one of them. It reads them with one query on a separate connection. +-} +openWithOptions :: FilePath -> ConnectionOptions -> IO Connection +openWithOptions path ConnectionOptions{connectionConfig, connectionGeometryCRS} = mask_ do - db <- openDatabaseWithConfig path settings + db <- openDatabaseWithConfig path connectionConfig conn <- connectDatabase db `onException` closeDatabaseHandle db - createConnection db conn - `onException` do - closeConnectionHandle conn - closeDatabaseHandle db + let closeBoth = closeConnectionHandle conn >> closeDatabaseHandle db + cache <- + createTypeCache db connectionGeometryCRS + `onException` closeBoth + createConnection db conn cache + `onException` (destroyTypeCache cache >> closeBoth) -- | Close a connection. The operation is idempotent. close :: Connection -> IO () @@ -166,6 +201,10 @@ withConnectionWithConfig :: FilePath -> [(Text, Text)] -> (Connection -> IO a) -> IO a withConnectionWithConfig path settings = bracket (openWithConfig path settings) close +-- | Run an action with a connection opened with options, closing it afterwards. +withConnectionWithOptions :: FilePath -> ConnectionOptions -> (Connection -> IO a) -> IO a +withConnectionWithOptions path options = bracket (openWithOptions path options) close + -- | Prepare a SQL statement for execution. openStatement :: Connection -> Query -> IO Statement openStatement conn queryText = @@ -181,7 +220,7 @@ if rc == DuckDBSuccess then pure stmt else do - errMsg <- fetchPrepareError stmt + errMsg <- fetchPrepareError (Text.pack "duckdb-simple: prepare failed") stmt throwIO $ mkPrepareError queryText errMsg createStatement conn handle queryText `onException` destroyPrepared handle @@ -284,7 +323,7 @@ resetStatementStream stmt rc <- c_duckdb_clear_bindings handle when (rc /= DuckDBSuccess) $ do - err <- fetchPrepareError handle + err <- fetchPrepareError (Text.pack "duckdb-simple: prepare failed") handle throwIO $ mkPrepareError (statementQuery stmt) err -- | Look up the 1-based index of a named placeholder. @@ -448,9 +487,9 @@ -- Internal helpers ----------------------------------------------------------- -createConnection :: DuckDBDatabase -> DuckDBConnection -> IO Connection -createConnection db conn = do - ref <- newIORef (ConnectionOpen db conn) +createConnection :: DuckDBDatabase -> DuckDBConnection -> TypeCache -> IO Connection +createConnection db conn cache = do + ref <- newIORef (ConnectionOpen db conn cache) _ <- mkWeakIORef ref $ join $ @@ -534,7 +573,8 @@ closeHandles :: ConnectionState -> IO () closeHandles ConnectionClosed = pure () -closeHandles ConnectionOpen{connectionDatabase, connectionHandle} = do +closeHandles ConnectionOpen{connectionDatabase, connectionHandle, connectionTypeCache} = do + destroyTypeCache connectionTypeCache closeConnectionHandle connectionHandle closeDatabaseHandle connectionDatabase @@ -549,13 +589,6 @@ destroyPrepared :: DuckDBPreparedStatement -> IO () destroyPrepared stmt = alloca \ptr -> poke ptr stmt >> c_duckdb_destroy_prepare ptr - -fetchPrepareError :: DuckDBPreparedStatement -> IO Text -fetchPrepareError stmt = do - msgPtr <- c_duckdb_prepare_error stmt - if msgPtr == nullPtr - then pure (Text.pack "duckdb-simple: prepare failed") - else peekUtf8CString msgPtr mkOpenError :: Text -> SQLError mkOpenError msg = diff --git a/src/Database/DuckDB/Simple/Copy.hs b/src/Database/DuckDB/Simple/Copy.hs --- a/src/Database/DuckDB/Simple/Copy.hs +++ b/src/Database/DuckDB/Simple/Copy.hs @@ -23,7 +23,7 @@ import Database.DuckDB.Simple.Callback (runCallback, transferCallbackState, withCallbackResources) import Database.DuckDB.Simple.FromField (Field (..)) import Database.DuckDB.Simple.Internal (Connection, destroyLogicalType, peekUtf8CString, throwRegistrationError, withConnectionHandle) -import Database.DuckDB.Simple.Materialize (materializeValue) +import Database.DuckDB.Simple.Materialize (prepareVectorReader) import Foreign.Marshal.Alloc (alloca) import Foreign.Ptr (Ptr, nullPtr) import Foreign.StablePtr (StablePtr, castPtrToStablePtr, deRefStablePtr) @@ -180,13 +180,10 @@ makeColumnReader :: DuckDBDataChunk -> Int -> IO ColumnReader makeColumnReader chunk columnIndex = do - vector <- c_duckdb_data_chunk_get_vector chunk (fromIntegral columnIndex) - dtype <- bracket (c_duckdb_vector_get_column_type vector) destroyLogicalType c_duckdb_get_type_id - dataPtr <- c_duckdb_vector_get_data vector - validity <- c_duckdb_vector_get_validity vector + readValue <- c_duckdb_data_chunk_get_vector chunk (fromIntegral columnIndex) >>= prepareVectorReader let name = Text.pack ("column" <> show columnIndex) pure \rowIdx -> do - fieldValue <- materializeValue dtype vector dataPtr validity (fromIntegral rowIdx) + fieldValue <- readValue (fromIntegral rowIdx) pure Field{fieldName = name, fieldIndex = columnIndex, fieldValue} destroyCopyFunction :: DuckDBCopyFunction -> IO () diff --git a/src/Database/DuckDB/Simple/Element.hs b/src/Database/DuckDB/Simple/Element.hs new file mode 100644 --- /dev/null +++ b/src/Database/DuckDB/Simple/Element.hs @@ -0,0 +1,228 @@ +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE NamedFieldPuns #-} + +{- | Decoders for one native DuckDB element in vector memory. Result decoding +uses them for the types whose element needs no logical type metadata. +-} +module Database.DuckDB.Simple.Element ( + chunkIsRowValid, + bitStringFromBytes, + chunkDecodeBlob, + decodeElement, + duckDBHugeIntToInteger, + vectorElementType, + withVectorType, +) where + +import Control.Exception (bracket, throwIO) +import Control.Monad (when) +import Data.Bits (clearBit, shiftL, xor, (.|.)) +import qualified Data.ByteString as BS +import Data.Int (Int64) +import Data.Ratio ((%)) +import qualified Data.Text.Encoding as TextEncoding +import Data.Time.Calendar (addDays, fromGregorian) +import Data.Time.Clock.POSIX (posixSecondsToUTCTime) +import Data.Time.LocalTime (TimeOfDay (..), minutesToTimeZone, utc, utcToLocalTime) +import qualified Data.UUID as UUID +import Data.Word (Word64, Word8) +import Database.DuckDB.FFI +import Database.DuckDB.Simple.FromField ( + BigNum (..), + BitString (..), + FieldValue (..), + IntervalValue (..), + TimeWithZone (..), + fromBigNumBytes, + ) +import Database.DuckDB.Simple.LogicalRep (destroyLogicalType) +import Database.DuckDB.Simple.Time (Date, LocalTimestamp, UTCTimestamp, Unbounded (..)) +import Foreign.C.Types (CBool (..)) +import Foreign.Marshal.Alloc (alloca) +import Foreign.Ptr (Ptr, castPtr, nullPtr, plusPtr) +import Foreign.Storable (Storable (..), peekElemOff) + +chunkIsRowValid :: Ptr Word64 -> DuckDBIdx -> IO Bool +chunkIsRowValid validity rowIdx + | validity == nullPtr = pure True + | otherwise = do + CBool flag <- c_duckdb_validity_row_is_valid validity rowIdx + pure (flag /= 0) + +chunkDecodeBlob :: Ptr () -> DuckDBIdx -> IO BS.ByteString +chunkDecodeBlob dataPtr rowIdx = do + let base = castPtr dataPtr :: Ptr Word8 + offset = fromIntegral rowIdx * duckdbStringTSize + stringPtr = castPtr (base `plusPtr` offset) :: Ptr DuckDBStringT + len <- c_duckdb_string_t_length stringPtr + if len == 0 + then pure BS.empty + else do + ptr <- c_duckdb_string_t_data stringPtr + BS.packCStringLen (ptr, fromIntegral len) + +duckdbStringTSize :: Int +duckdbStringTSize = 16 + +{- | Decode DuckDB's BIT bytes: a padding count, then the data bytes. Clear the +unused high bits of the first data byte, which DuckDB sets. +-} +bitStringFromBytes :: BS.ByteString -> BitString +bitStringFromBytes bytes = case BS.unpack bytes of + [] -> BitString 0 BS.empty + [padding] -> BitString padding BS.empty + padding : first : rest -> + BitString padding (BS.pack (foldl clearBit first [8 - fromIntegral padding .. 7] : rest)) + +-- | Borrow the logical type of a vector. The type is destroyed after the action. +withVectorType :: DuckDBVector -> (DuckDBLogicalType -> IO a) -> IO a +withVectorType vector = bracket (c_duckdb_vector_get_column_type vector) destroyLogicalType + +vectorElementType :: DuckDBVector -> IO DuckDBType +vectorElementType vec = + bracket (c_duckdb_vector_get_column_type vec) destroyLogicalType c_duckdb_get_type_id + +{- | Decode the element at an index of a vector's data. The type must not need +logical type metadata: DECIMAL, ENUM, GEOMETRY, and nested types use their own +decoders. The caller checks validity. +-} +decodeElement :: DuckDBType -> Ptr () -> Int -> IO FieldValue +decodeElement dtype dataPtr rowIdx = case dtype of + DuckDBTypeBoolean -> do + raw <- peekElemOff (castPtr dataPtr :: Ptr Word8) rowIdx + pure (FieldBool (raw /= 0)) + DuckDBTypeTinyInt -> FieldInt8 <$> peekElemOff (castPtr dataPtr) rowIdx + DuckDBTypeSmallInt -> FieldInt16 <$> peekElemOff (castPtr dataPtr) rowIdx + DuckDBTypeInteger -> FieldInt32 <$> peekElemOff (castPtr dataPtr) rowIdx + DuckDBTypeBigInt -> FieldInt64 <$> peekElemOff (castPtr dataPtr) rowIdx + DuckDBTypeUTinyInt -> FieldWord8 <$> peekElemOff (castPtr dataPtr) rowIdx + DuckDBTypeUSmallInt -> FieldWord16 <$> peekElemOff (castPtr dataPtr) rowIdx + DuckDBTypeUInteger -> FieldWord32 <$> peekElemOff (castPtr dataPtr) rowIdx + DuckDBTypeUBigInt -> FieldWord64 <$> peekElemOff (castPtr dataPtr) rowIdx + DuckDBTypeFloat -> FieldFloat <$> peekElemOff (castPtr dataPtr) rowIdx + DuckDBTypeDouble -> FieldDouble <$> peekElemOff (castPtr dataPtr) rowIdx + DuckDBTypeVarchar -> FieldText . TextEncoding.decodeUtf8 <$> chunkDecodeBlob dataPtr index + DuckDBTypeStringLiteral -> FieldText . TextEncoding.decodeUtf8 <$> chunkDecodeBlob dataPtr index + DuckDBTypeBlob -> FieldBlob <$> chunkDecodeBlob dataPtr index + DuckDBTypeUUID -> do + DuckDBUHugeInt lower upperBiased <- peekElemOff (castPtr dataPtr :: Ptr DuckDBUHugeInt) rowIdx + let upper = upperBiased `xor` (0x8000000000000000 :: Word64) + pure (FieldUUID (UUID.fromWords64 (fromIntegral upper) lower)) + DuckDBTypeDate -> peekElemOff (castPtr dataPtr) rowIdx >>= fmap FieldDate . decodeDuckDBDate + DuckDBTypeTime -> peekElemOff (castPtr dataPtr) rowIdx >>= fmap FieldTime . decodeDuckDBTime + DuckDBTypeTimeNs -> FieldTime . decodeDuckDBTimeNs <$> peekElemOff (castPtr dataPtr) rowIdx + DuckDBTypeTimeTz -> peekElemOff (castPtr dataPtr) rowIdx >>= fmap FieldTimeTZ . decodeDuckDBTimeTz + DuckDBTypeTimestamp -> peekElemOff (castPtr dataPtr) rowIdx >>= fmap FieldTimestamp . decodeDuckDBTimestamp + DuckDBTypeTimestampS -> peekElemOff (castPtr dataPtr) rowIdx >>= fmap FieldTimestamp . decodeDuckDBTimestampSeconds + DuckDBTypeTimestampMs -> peekElemOff (castPtr dataPtr) rowIdx >>= fmap FieldTimestamp . decodeDuckDBTimestampMilliseconds + DuckDBTypeTimestampNs -> peekElemOff (castPtr dataPtr) rowIdx >>= fmap FieldTimestamp . decodeDuckDBTimestampNanoseconds + DuckDBTypeTimestampTz -> peekElemOff (castPtr dataPtr) rowIdx >>= fmap FieldTimestampTZ . decodeDuckDBTimestampUTCTime + DuckDBTypeInterval -> FieldInterval . intervalValueFromDuckDB <$> peekElemOff (castPtr dataPtr) rowIdx + DuckDBTypeHugeInt -> FieldHugeInt . duckDBHugeIntToInteger <$> peekElemOff (castPtr dataPtr) rowIdx + DuckDBTypeUHugeInt -> FieldUHugeInt . duckDBUHugeIntToInteger <$> peekElemOff (castPtr dataPtr) rowIdx + DuckDBTypeBit -> FieldBit . bitStringFromBytes <$> chunkDecodeBlob dataPtr index + DuckDBTypeBigNum -> do + bytes <- chunkDecodeBlob dataPtr index + pure (FieldBigNum (BigNum (if BS.length bytes < 3 then 0 else fromBigNumBytes (BS.unpack bytes)))) + DuckDBTypeIntegerLiteral -> FieldInt64 <$> peekElemOff (castPtr dataPtr) rowIdx + DuckDBTypeInvalid -> + error "duckdb-simple: INVALID type in eager result" + DuckDBTypeAny -> + error "duckdb-simple: ANY columns should not appear in results" + other -> + error ("duckdb-simple: UNKNOWN type in eager result: " <> show other) + where + index = fromIntegral rowIdx :: DuckDBIdx + +-- | Decode dates with exact epoch arithmetic and preserve infinity. +decodeDuckDBDate :: DuckDBDate -> IO Date +decodeDuckDBDate (DuckDBDate days) = + pure (decodeUnbounded (\value -> addDays (toInteger value) (fromGregorian 1970 1 1)) days) + +decodeDuckDBTime :: DuckDBTime -> IO TimeOfDay +decodeDuckDBTime raw = + alloca $ \ptr -> do + c_duckdb_from_time raw ptr + timeStruct <- peek ptr + pure (timeStructToTimeOfDay timeStruct) + +decodeDuckDBTimestamp :: DuckDBTimestamp -> IO LocalTimestamp +decodeDuckDBTimestamp (DuckDBTimestamp micros) = decodeTimestampUnits 1000000 micros + +-- | Interpret native infinity sentinels before converting a finite payload. +decodeUnbounded :: (Integral a, Bounded a) => (a -> b) -> a -> Unbounded b +decodeUnbounded decode value + | value == maxBound = PosInfinity + | value == negate maxBound = NegInfinity + | otherwise = Finite (decode value) + +-- | Decode timestamp units without overflowing an intermediate Int64. +decodeTimestampUnits :: Integer -> Int64 -> IO LocalTimestamp +decodeTimestampUnits units = + pure . decodeUnbounded (utcToLocalTime utc . posixSecondsToUTCTime . fromRational . (% units) . toInteger) + +decodeDuckDBTimeNs :: DuckDBTimeNs -> TimeOfDay +decodeDuckDBTimeNs (DuckDBTimeNs nanos) = + let (hours, remainderHours) = nanos `divMod` (60 * 60 * 1000000000) + (minutes, remainderMinutes) = remainderHours `divMod` (60 * 1000000000) + (seconds, fractionalNanos) = remainderMinutes `divMod` 1000000000 + fractional = fromRational (toInteger fractionalNanos % 1000000000) + totalSeconds = fromIntegral seconds + fractional + in TimeOfDay + (fromIntegral hours) + (fromIntegral minutes) + totalSeconds + +decodeDuckDBTimeTz :: DuckDBTimeTz -> IO TimeWithZone +decodeDuckDBTimeTz raw = + alloca $ \ptr -> do + c_duckdb_from_time_tz raw ptr + DuckDBTimeTzStruct{duckDBTimeTzStructTime = timeStruct, duckDBTimeTzStructOffset = offset} <- peek ptr + when (offset `rem` 60 /= 0) $ + throwIO (userError "duckdb-simple: TIMETZ offset cannot be represented in whole minutes") + let timeOfDay = timeStructToTimeOfDay timeStruct + minutes = fromIntegral offset `div` 60 + zone = minutesToTimeZone minutes + pure TimeWithZone{timeWithZoneTime = timeOfDay, timeWithZoneZone = zone} + +decodeDuckDBTimestampSeconds :: DuckDBTimestampS -> IO LocalTimestamp +decodeDuckDBTimestampSeconds (DuckDBTimestampS seconds) = + decodeTimestampUnits 1 seconds + +decodeDuckDBTimestampMilliseconds :: DuckDBTimestampMs -> IO LocalTimestamp +decodeDuckDBTimestampMilliseconds (DuckDBTimestampMs millis) = + decodeTimestampUnits 1000 millis + +decodeDuckDBTimestampNanoseconds :: DuckDBTimestampNs -> IO LocalTimestamp +decodeDuckDBTimestampNanoseconds (DuckDBTimestampNs nanos) = decodeTimestampUnits 1000000000 nanos + +decodeDuckDBTimestampUTCTime :: DuckDBTimestamp -> IO UTCTimestamp +decodeDuckDBTimestampUTCTime (DuckDBTimestamp micros) = + pure (decodeUnbounded (posixSecondsToUTCTime . fromRational . (% 1000000) . toInteger) micros) + +intervalValueFromDuckDB :: DuckDBInterval -> IntervalValue +intervalValueFromDuckDB DuckDBInterval{duckDBIntervalMonths, duckDBIntervalDays, duckDBIntervalMicros} = + IntervalValue + { intervalMonths = duckDBIntervalMonths + , intervalDays = duckDBIntervalDays + , intervalMicros = duckDBIntervalMicros + } + +duckDBHugeIntToInteger :: DuckDBHugeInt -> Integer +duckDBHugeIntToInteger DuckDBHugeInt{duckDBHugeIntLower, duckDBHugeIntUpper} = + (fromIntegral duckDBHugeIntUpper `shiftL` 64) .|. fromIntegral duckDBHugeIntLower + +duckDBUHugeIntToInteger :: DuckDBUHugeInt -> Integer +duckDBUHugeIntToInteger DuckDBUHugeInt{duckDBUHugeIntLower, duckDBUHugeIntUpper} = + (fromIntegral duckDBUHugeIntUpper `shiftL` 64) .|. fromIntegral duckDBUHugeIntLower + +timeStructToTimeOfDay :: DuckDBTimeStruct -> TimeOfDay +timeStructToTimeOfDay DuckDBTimeStruct{duckDBTimeStructHour, duckDBTimeStructMinute, duckDBTimeStructSecond, duckDBTimeStructMicros} = + let secondsInt = fromIntegral duckDBTimeStructSecond :: Integer + micros = fromIntegral duckDBTimeStructMicros :: Integer + fractional = fromRational (micros % 1000000) + totalSeconds = fromInteger secondsInt + fractional + in TimeOfDay + (fromIntegral duckDBTimeStructHour) + (fromIntegral duckDBTimeStructMinute) + totalSeconds diff --git a/src/Database/DuckDB/Simple/FromField.hs b/src/Database/DuckDB/Simple/FromField.hs --- a/src/Database/DuckDB/Simple/FromField.hs +++ b/src/Database/DuckDB/Simple/FromField.hs @@ -15,6 +15,7 @@ module Database.DuckDB.Simple.FromField ( Field (..), FieldValue (..), + RawGeometry (..), StructField (..), StructValue (..), UnionMemberType (..), @@ -39,6 +40,7 @@ import Data.Bits (Bits (..), finiteBitSize) import qualified Data.ByteString as BS import Data.Data (Typeable, typeRep) +import qualified Data.Geometry as G import Data.Int (Int16, Int32, Int64, Int8) import Data.Map.Strict (Map) import qualified Data.Map.Strict as Map @@ -58,6 +60,7 @@ ) import qualified Data.UUID as UUID import Data.Word (Word16, Word32, Word64, Word8) +import Database.DuckDB.Simple.Geometry (RawGeometry (..), fromRawGeometry) import Database.DuckDB.Simple.LogicalRep ( LogicalTypeRep (..), StructField (..), @@ -67,12 +70,14 @@ ) import Database.DuckDB.Simple.Ok import Database.DuckDB.Simple.Time (Date, LocalTimestamp, UTCTimestamp, Unbounded (..)) -import Database.DuckDB.Simple.Types (Null (..)) import GHC.Float (double2Float, float2Double) import GHC.Num.Integer (integerFromWordList) import Numeric.Natural (Natural) --- | Internal representation of a column value. +{- | Internal representation of a column value. +'FieldGeometry' contains raw WKB and CRS metadata. Generic parameters that +contain a non-NULL geometry require explicit SQL import. +-} data FieldValue = FieldNull | FieldInt8 Int8 @@ -89,6 +94,7 @@ | FieldText Text | FieldBool Bool | FieldBlob BS.ByteString + | FieldGeometry RawGeometry | FieldDate Date | FieldTime TimeOfDay | FieldTimestamp LocalTimestamp @@ -305,6 +311,20 @@ instance FromField FieldValue where fromField Field{fieldValue} = Ok fieldValue +instance FromField RawGeometry where + fromField f@Field{fieldValue} = + case fieldValue of + FieldGeometry value -> Ok value + FieldNull -> returnError UnexpectedNull f "" + _ -> returnError Incompatible f "expected GEOMETRY" + +-- | Decode the shape without CRS metadata. Use t'RawGeometry' to retain the CRS. +instance FromField G.Geometry where + fromField f@Field{fieldValue} = case fieldValue of + FieldGeometry raw -> either (returnError ConversionFailed f . Text.pack) pure (fromRawGeometry raw) + FieldNull -> returnError UnexpectedNull f "" + _ -> returnError Incompatible f "expected GEOMETRY" + instance FromField (StructValue FieldValue) where fromField f@Field{fieldValue} = case fieldValue of @@ -319,12 +339,6 @@ FieldNull -> returnError UnexpectedNull f "" _ -> returnError Incompatible f "expected UNION" -instance FromField Null where - fromField f@Field{fieldValue} = - case fieldValue of - FieldNull -> Ok Null - _ -> returnError Incompatible f "expected NULL" - instance FromField UUID.UUID where fromField f@Field{fieldValue} = case fieldValue of @@ -549,6 +563,7 @@ fromField f@Field{fieldValue} = case fieldValue of FieldBlob bs -> Ok bs + FieldGeometry RawGeometry{rawGeometryWKB} -> Ok rawGeometryWKB FieldText t -> Ok (TextEncoding.encodeUtf8 t) FieldBit (BitString _ bits) -> Ok bits FieldNull -> returnError UnexpectedNull f "" @@ -777,6 +792,7 @@ FieldText{} -> "TEXT" FieldBool{} -> "BOOLEAN" FieldBlob{} -> "BLOB" + FieldGeometry{} -> "GEOMETRY" FieldDate{} -> "DATE" FieldTime{} -> "TIME" FieldTimestamp{} -> "TIMESTAMP" diff --git a/src/Database/DuckDB/Simple/Function.hs b/src/Database/DuckDB/Simple/Function.hs --- a/src/Database/DuckDB/Simple/Function.hs +++ b/src/Database/DuckDB/Simple/Function.hs @@ -56,7 +56,7 @@ withQueryCString, withResult, ) -import Database.DuckDB.Simple.Materialize (materializeValue) +import Database.DuckDB.Simple.Materialize (prepareVectorReader) import Database.DuckDB.Simple.Ok (Ok (..)) import Foreign.Marshal.Alloc (alloca) import Foreign.Ptr (FunPtr, Ptr, castPtr, nullPtr) @@ -360,13 +360,10 @@ makeColumnReader :: DuckDBDataChunk -> Int -> IO ColumnReader makeColumnReader chunk columnIndex = do - vector <- c_duckdb_data_chunk_get_vector chunk (fromIntegral columnIndex) - dtype <- bracket (c_duckdb_vector_get_column_type vector) destroyLogicalType c_duckdb_get_type_id - dataPtr <- c_duckdb_vector_get_data vector - validity <- c_duckdb_vector_get_validity vector + readValue <- c_duckdb_data_chunk_get_vector chunk (fromIntegral columnIndex) >>= prepareVectorReader let name = Text.pack ("arg" <> show columnIndex) pure \rowIdx -> do - value <- materializeValue dtype vector dataPtr validity (fromIntegral rowIdx) + value <- readValue (fromIntegral rowIdx) pure Field { fieldName = name diff --git a/src/Database/DuckDB/Simple/Generic.hs b/src/Database/DuckDB/Simple/Generic.hs --- a/src/Database/DuckDB/Simple/Generic.hs +++ b/src/Database/DuckDB/Simple/Generic.hs @@ -121,6 +121,7 @@ TimeWithZone (..), returnError, ) +import Database.DuckDB.Simple.Internal (duckDBTypeToName) import Database.DuckDB.Simple.LogicalRep ( LogicalTypeRep (..), StructField (..), @@ -131,6 +132,7 @@ import Database.DuckDB.Simple.Ok (Ok (..)) import Database.DuckDB.Simple.Time (Unbounded (..)) import Database.DuckDB.Simple.ToField (DuckDBColumnType (..), ToField (..)) +import Database.DuckDB.Simple.Variant (Variant (..)) -------------------------------------------------------------------------------- -- DuckValue: bridge between Haskell scalars and FieldValue/LogicalTypeRep @@ -158,6 +160,10 @@ Ok x -> Right x Errors errs -> Left (unlines (map displayException errs)) +instance DuckValue Variant where + duckToField = variantPayload + duckLogicalType _ = LogicalTypeScalar DuckDBTypeVariant + instance DuckValue Bool where duckToField = FieldBool duckLogicalType _ = LogicalTypeScalar DuckDBTypeBoolean @@ -446,10 +452,7 @@ instance (GSum f) => GToField' 'True (M1 D meta f) where gToField' _ (M1 value) = let members = gSumMembers (Proxy :: Proxy (f p)) - membersArray = - case members of - [] -> listArray (0, -1) [] - _ -> listArray (0, length members - 1) members + membersArray = listArray (0, length members - 1) members (idx, payload) = gSumEncode value label = unionMemberName (members !! idx) in EncodedUnion @@ -461,10 +464,7 @@ } gLogicalType' _ _ = let members = gSumMembers (Proxy :: Proxy (f p)) - membersArray = - case members of - [] -> listArray (0, -1) [] - _ -> listArray (0, length members - 1) members + membersArray = listArray (0, length members - 1) members in LogicalTypeUnion membersArray -------------------------------------------------------------------------------- @@ -491,12 +491,9 @@ -} listArrayFrom :: [Text] -> [b] -> Array Int (StructField b) listArrayFrom names values = - case values of - [] -> listArray (0, -1) [] - _ -> - listArray - (0, length values - 1) - (zipWith (\n v -> StructField{structFieldName = n, structFieldValue = v}) names values) + listArray + (0, length values - 1) + (zipWith (\n v -> StructField{structFieldName = n, structFieldValue = v}) names values) {- | Collect the components (values and types) of a product. Implementations produce parallel lists so we can zip them during encoding and decoding. @@ -707,9 +704,10 @@ LogicalTypeList{} -> Text.pack "LIST" LogicalTypeArray{} -> Text.pack "ARRAY" LogicalTypeMap{} -> Text.pack "MAP" - LogicalTypeScalar dtype -> duckdbTypeToName dtype + LogicalTypeScalar dtype -> duckDBTypeToName dtype LogicalTypeDecimal{} -> Text.pack "DECIMAL" LogicalTypeEnum{} -> Text.pack "ENUM" + LogicalTypeGeometry{} -> Text.pack "GEOMETRY" {- | Deriving-via @ToField@ instance. We reuse the helpers above to decide whether the top-level representation is a union, struct, or scalar and then @@ -738,37 +736,6 @@ Right value -> pure (ViaDuckDB value) Left err -> returnError ConversionFailed f (Text.pack err) - -duckdbTypeToName :: DuckDBType -> Text - -{- | Translate a @DuckDBType@ into a textual label for diagnostics and -documentation. This mirrors the naming used in "Database.DuckDB.Simple.ToField". --} -duckdbTypeToName dtype - | dtype == DuckDBTypeBoolean = Text.pack "BOOLEAN" - | dtype == DuckDBTypeTinyInt = Text.pack "TINYINT" - | dtype == DuckDBTypeSmallInt = Text.pack "SMALLINT" - | dtype == DuckDBTypeInteger = Text.pack "INTEGER" - | dtype == DuckDBTypeBigInt = Text.pack "BIGINT" - | dtype == DuckDBTypeUTinyInt = Text.pack "UTINYINT" - | dtype == DuckDBTypeUSmallInt = Text.pack "USMALLINT" - | dtype == DuckDBTypeUInteger = Text.pack "UINTEGER" - | dtype == DuckDBTypeUBigInt = Text.pack "UBIGINT" - | dtype == DuckDBTypeFloat = Text.pack "FLOAT" - | dtype == DuckDBTypeDouble = Text.pack "DOUBLE" - | dtype == DuckDBTypeVarchar = Text.pack "VARCHAR" - | dtype == DuckDBTypeBlob = Text.pack "BLOB" - | dtype == DuckDBTypeDate = Text.pack "DATE" - | dtype == DuckDBTypeTime = Text.pack "TIME" - | dtype == DuckDBTypeTimestamp = Text.pack "TIMESTAMP" - | dtype == DuckDBTypeTimestampTz = Text.pack "TIMESTAMP_TZ" - | dtype == DuckDBTypeUUID = Text.pack "UUID" - | dtype == DuckDBTypeInterval = Text.pack "INTERVAL" - | dtype == DuckDBTypeHugeInt = Text.pack "HUGEINT" - | dtype == DuckDBTypeUHugeInt = Text.pack "UHUGEINT" - | dtype == DuckDBTypeBigNum = Text.pack "BIGNUM" - | dtype == DuckDBTypeTimeTz = Text.pack "TIME_TZ" - | otherwise = Text.pack (show dtype) -------------------------------------------------------------------------------- -- DuckDB type constructors (re-exported patterns) diff --git a/src/Database/DuckDB/Simple/Geometry.hs b/src/Database/DuckDB/Simple/Geometry.hs new file mode 100644 --- /dev/null +++ b/src/Database/DuckDB/Simple/Geometry.hs @@ -0,0 +1,40 @@ +-- | Raw geometry values and conversion to 'G.Geometry' from @geometry-simple@. +module Database.DuckDB.Simple.Geometry ( + RawGeometry (..), + fromRawGeometry, + toRawGeometry, +) where + +import Data.ByteString (ByteString) +import qualified Data.Geometry as G +import Data.Geometry.WKB (decodeWKB, encodeWKB) +import Data.Text (Text) + +{- | ISO WKB bytes with an optional CRS definition. +The definition can be an identifier, a custom name, or WKT2\/PROJJSON text. +Reading this type copies the bytes without decoding the coordinates. +This type has no @ToField@ instance. Bind 'rawGeometryWKB' with +@ST_GeomFromWKB(?)@. Apply 'rawGeometryCRS' with @ST_SetCRS(..., ?)@. +Omit @ST_SetCRS@ when the CRS is 'Nothing'. SQL NULL propagates to the geometry. +DuckDB can normalize the byte order and CRS definition during import. +Derived equality compares the bytes and CRS, not spatial equivalence. +-} +data RawGeometry = RawGeometry + { rawGeometryWKB :: !ByteString + , rawGeometryCRS :: !(Maybe Text) + } + deriving (Eq, Show, Read) + +{- | Decode owned WKB with the standalone codec. The shape does not store CRS +metadata. Empty multi-geometries and collections have no layout tag in the +decoded representation. Keep the raw value when these details are needed. +-} +fromRawGeometry :: RawGeometry -> Either String G.Geometry +fromRawGeometry = decodeWKB . rawGeometryWKB + +{- | Encode a shape as ISO WKB with no CRS. Set 'rawGeometryCRS' on the result +to supply a CRS definition for explicit SQL import. This does not transform +the coordinates. +-} +toRawGeometry :: G.Geometry -> Either String RawGeometry +toRawGeometry shape = (`RawGeometry` Nothing) <$> encodeWKB shape diff --git a/src/Database/DuckDB/Simple/Internal.hs b/src/Database/DuckDB/Simple/Internal.hs --- a/src/Database/DuckDB/Simple/Internal.hs +++ b/src/Database/DuckDB/Simple/Internal.hs @@ -25,7 +25,6 @@ ResultMode (..), StatementStreamColumn (..), StatementStreamChunk (..), - StatementStreamChunkVector (..), SQLError (..), toSQLError, @@ -35,9 +34,13 @@ keepAlive, withDatabaseHandle, withConnectionHandle, + withTypeCache, withStatementHandle, withQueryCString, peekUtf8CString, + fetchPrepareError, + duckDBTypeFromName, + duckDBTypeToName, withResult, runInterruptibleQuery, executePreparedResult, @@ -58,40 +61,42 @@ import Control.Monad (when) import qualified Data.ByteString as BS import Data.IORef (IORef, readIORef) +import Data.List (find) import Data.String (IsString (..)) import Data.Text (Text) import qualified Data.Text as Text import qualified Data.Text.Encoding as TextEncoding import qualified Data.Text.Foreign as TextForeign -import Data.Word (Word64) import Database.DuckDB.FFI ( DuckDBClientContext, DuckDBConnection, DuckDBDataChunk, DuckDBDatabase, DuckDBErrorType, - DuckDBLogicalType, DuckDBPreparedStatement, DuckDBResult, DuckDBState, DuckDBType, DuckDBValue, - DuckDBVector, c_duckdb_connection_get_client_context, c_duckdb_destroy_client_context, c_duckdb_destroy_data_chunk, - c_duckdb_destroy_logical_type, c_duckdb_destroy_result, c_duckdb_destroy_value, c_duckdb_execute_prepared, c_duckdb_fetch_chunk, c_duckdb_interrupt, + c_duckdb_prepare_error, c_duckdb_result_error, c_duckdb_result_error_type, pattern DuckDBErrorInvalid, pattern DuckDBSuccess, ) +import qualified Database.DuckDB.FFI as FFI import Database.DuckDB.FFI.Deprecated (c_duckdb_execute_prepared_streaming) +import Database.DuckDB.Simple.FromField (FieldValue) +import Database.DuckDB.Simple.LogicalRep (destroyLogicalType) +import Database.DuckDB.Simple.TypeCache (TypeCache) import Foreign.C.String (CString) import Foreign.Marshal.Alloc (alloca) import Foreign.Marshal.Utils (fillBytes) @@ -122,6 +127,7 @@ | ConnectionOpen { connectionDatabase :: DuckDBDatabase , connectionHandle :: DuckDBConnection + , connectionTypeCache :: TypeCache } -- | Tracks the lifetime of a prepared statement. @@ -168,14 +174,8 @@ { statementStreamChunkPtr :: DuckDBDataChunk , statementStreamChunkSize :: Int , statementStreamChunkIndex :: Int - , statementStreamChunkVectors :: [StatementStreamChunkVector] - } - --- | Raw vector pointers backing a chunk column. -data StatementStreamChunkVector = StatementStreamChunkVector - { statementStreamChunkVectorHandle :: DuckDBVector - , statementStreamChunkVectorData :: Ptr () - , statementStreamChunkVectorValidity :: Ptr Word64 + , statementStreamChunkReaders :: [Int -> IO FieldValue] + -- ^ One reader for each column. A reader must not outlive its chunk. } -- | Represents an error reported by DuckDB or by duckdb-simple itself. @@ -300,6 +300,60 @@ interruptRetryDelayMicros :: Int interruptRetryDelayMicros = 10 * 1000 +{- | The column type names that 'Database.DuckDB.Simple.ToField.DuckDBColumnType' +instances use, and the types they denote. Each type has one name. +-} +duckDBTypeNames :: [(Text, DuckDBType)] +duckDBTypeNames = + map + (\(name, dtype) -> (Text.pack name, dtype)) + [ ("BOOLEAN", FFI.DuckDBTypeBoolean) + , ("TINYINT", FFI.DuckDBTypeTinyInt) + , ("SMALLINT", FFI.DuckDBTypeSmallInt) + , ("INTEGER", FFI.DuckDBTypeInteger) + , ("BIGINT", FFI.DuckDBTypeBigInt) + , ("HUGEINT", FFI.DuckDBTypeHugeInt) + , ("UTINYINT", FFI.DuckDBTypeUTinyInt) + , ("USMALLINT", FFI.DuckDBTypeUSmallInt) + , ("UINTEGER", FFI.DuckDBTypeUInteger) + , ("UBIGINT", FFI.DuckDBTypeUBigInt) + , ("UHUGEINT", FFI.DuckDBTypeUHugeInt) + , ("FLOAT", FFI.DuckDBTypeFloat) + , ("DOUBLE", FFI.DuckDBTypeDouble) + , ("DATE", FFI.DuckDBTypeDate) + , ("TIME", FFI.DuckDBTypeTime) + , ("TIMETZ", FFI.DuckDBTypeTimeTz) + , ("TIMESTAMP", FFI.DuckDBTypeTimestamp) + , ("TIMESTAMPTZ", FFI.DuckDBTypeTimestampTz) + , ("INTERVAL", FFI.DuckDBTypeInterval) + , ("TEXT", FFI.DuckDBTypeVarchar) + , ("BLOB", FFI.DuckDBTypeBlob) + , ("GEOMETRY", FFI.DuckDBTypeGeometry) + , ("VARIANT", FFI.DuckDBTypeVariant) + , ("UUID", FFI.DuckDBTypeUUID) + , ("BIT", FFI.DuckDBTypeBit) + , ("BIGNUM", FFI.DuckDBTypeBigNum) + , -- NULL gives an element type to Maybe values without data. + ("NULL", FFI.DuckDBTypeSQLNull) + ] + +-- | Find the type that a column type name denotes. +duckDBTypeFromName :: Text -> Maybe DuckDBType +duckDBTypeFromName name = lookup name duckDBTypeNames + +-- | Find the column type name of a type. Other types use their 'Show' text. +duckDBTypeToName :: DuckDBType -> Text +duckDBTypeToName dtype = + maybe (Text.pack (show dtype)) fst (find ((== dtype) . snd) duckDBTypeNames) + +{- | Read the error message of a prepared statement as UTF-8. Use the fallback +when DuckDB reports no message. +-} +fetchPrepareError :: Text -> DuckDBPreparedStatement -> IO Text +fetchPrepareError fallback statement = do + messagePtr <- c_duckdb_prepare_error statement + if messagePtr == nullPtr then pure fallback else peekUtf8CString messagePtr + -- | Copy a result error while its native result remains alive. fetchResultError :: Ptr DuckDBResult -> IO (Text, Maybe DuckDBErrorType) fetchResultError resultPtr = do @@ -351,6 +405,15 @@ ConnectionClosed -> throwIO connectionClosedError ConnectionOpen{connectionHandle} -> action connectionHandle +-- | Borrow the type cache of an open connection. +withTypeCache :: Connection -> (TypeCache -> IO a) -> IO a +withTypeCache conn@Connection{connectionState} action = + keepAlive conn $ do + state <- readIORef connectionState + case state of + ConnectionClosed -> throwIO connectionClosedError + ConnectionOpen{connectionTypeCache} -> action connectionTypeCache + -- | Internal helper for safely accessing the underlying database handle. withDatabaseHandle :: Connection -> (DuckDBDatabase -> IO a) -> IO a withDatabaseHandle conn@Connection{connectionState} action = @@ -378,11 +441,6 @@ destroyValue :: DuckDBValue -> IO () destroyValue value = alloca $ \ptr -> poke ptr value >> c_duckdb_destroy_value ptr - --- | Destroy a logical type handle. -destroyLogicalType :: DuckDBLogicalType -> IO () -destroyLogicalType logicalType = - alloca $ \ptr -> poke ptr logicalType >> c_duckdb_destroy_logical_type ptr -- | Throw a standardised registration error. throwRegistrationError :: String -> IO a diff --git a/src/Database/DuckDB/Simple/LogicalRep.hs b/src/Database/DuckDB/Simple/LogicalRep.hs --- a/src/Database/DuckDB/Simple/LogicalRep.hs +++ b/src/Database/DuckDB/Simple/LogicalRep.hs @@ -17,6 +17,7 @@ unionValueTypeRep, logicalTypeToRep, logicalTypeFromRep, + logicalTypeFromRepWith, destroyLogicalType, ) where @@ -41,6 +42,7 @@ data LogicalTypeRep = LogicalTypeScalar DuckDBType | LogicalTypeDecimal !Word8 !Word8 + | LogicalTypeGeometry !(Maybe Text) | LogicalTypeList LogicalTypeRep | LogicalTypeArray LogicalTypeRep !Word64 | LogicalTypeMap LogicalTypeRep LogicalTypeRep @@ -100,6 +102,9 @@ logicalTypeToRep logical = do dtype <- c_duckdb_get_type_id logical case dtype of + DuckDBTypeGeometry -> + bracket (c_duckdb_geometry_type_get_crs logical) (c_duckdb_free . castPtr) \ptr -> + LogicalTypeGeometry <$> if ptr == nullPtr then pure Nothing else Just . TextEncoding.decodeUtf8 <$> BS.packCString ptr DuckDBTypeStruct -> do childCountRaw <- c_duckdb_struct_type_child_count logical childCount <- word64ToInt (Text.pack "struct child count") childCountRaw @@ -114,10 +119,7 @@ pure StructField{structFieldName = name, structFieldValue = childRep} pure $ LogicalTypeStruct - ( if childCount <= 0 - then listArray (0, -1) [] - else listArray (0, childCount - 1) fields - ) + (listArray (0, childCount - 1) fields) DuckDBTypeUnion -> do memberCountRaw <- c_duckdb_union_type_member_count logical memberCount <- word64ToInt (Text.pack "union member count") memberCountRaw @@ -132,10 +134,7 @@ pure UnionMemberType{unionMemberName = name, unionMemberType = memberRep} pure $ LogicalTypeUnion - ( if memberCount <= 0 - then listArray (0, -1) [] - else listArray (0, memberCount - 1) members - ) + (listArray (0, memberCount - 1) members) DuckDBTypeList -> do childRep <- bracket (c_duckdb_list_type_child_type logical) destroyLogicalType logicalTypeToRep pure (LogicalTypeList childRep) @@ -163,46 +162,64 @@ pure entry pure $ LogicalTypeEnum - ( if count <= 0 - then listArray (0, -1) [] - else listArray (0, count - 1) entries - ) + (listArray (0, count - 1) entries) _ -> pure (LogicalTypeScalar dtype) --- | Materialize a DuckDB logical type handle from a @LogicalTypeRep@ tree. +{- | Materialize a DuckDB logical type handle from a @LogicalTypeRep@ tree. +The C API cannot create a usable VARIANT type or a GEOMETRY type with a CRS. +This function raises an error for VARIANT. It creates GEOMETRY without a CRS +and ignores the CRS of 'LogicalTypeGeometry'. Parameter binding uses the +types that a connection reads the first time a parameter needs one. +-} logicalTypeFromRep :: LogicalTypeRep -> IO DuckDBLogicalType -logicalTypeFromRep rep = do +logicalTypeFromRep = logicalTypeFromRepWith \case + -- TODO: improve this when this becomes available in the C API. + -- See https://github.com/Tritlo/duckdb-haskell/issues/26 and + -- https://github.com/Tritlo/duckdb-haskell/issues/27. + LogicalTypeScalar DuckDBTypeVariant -> + throwIO (userError "duckdb-simple: a VARIANT type needs the type cache of a connection; bind the value as a parameter or cast a plain value with ?::VARIANT") + _ -> c_duckdb_create_logical_type DuckDBTypeGeometry + +{- | Materialize a type tree. The function argument creates the leaves that the +C API cannot create: VARIANT, and GEOMETRY with a CRS. The caller must +destroy the result. +-} +logicalTypeFromRepWith :: (LogicalTypeRep -> IO DuckDBLogicalType) -> LogicalTypeRep -> IO DuckDBLogicalType +logicalTypeFromRepWith resolve rep = do logical <- create rep when (logical == nullPtr) $ throwIO (userError "duckdb-simple: DuckDB logical type construction failed") pure logical where create = \case + leaf@(LogicalTypeScalar DuckDBTypeVariant) -> resolve leaf LogicalTypeScalar dtype -> c_duckdb_create_logical_type dtype + leaf@(LogicalTypeGeometry (Just _)) -> resolve leaf + LogicalTypeGeometry Nothing -> c_duckdb_create_logical_type DuckDBTypeGeometry LogicalTypeDecimal width scale -> do when (width < 1 || width > 38 || scale > width) $ throwIO (userError "duckdb-simple: invalid DECIMAL width or scale") c_duckdb_create_decimal_type width scale LogicalTypeList elemRep -> - bracket (logicalTypeFromRep elemRep) destroyLogicalType c_duckdb_create_list_type + bracket (logicalTypeFromRepWith resolve elemRep) destroyLogicalType c_duckdb_create_list_type LogicalTypeArray elemRep size -> - bracket (logicalTypeFromRep elemRep) destroyLogicalType $ + bracket (logicalTypeFromRepWith resolve elemRep) destroyLogicalType $ flip c_duckdb_create_array_type size LogicalTypeMap keyRep valueRep -> - bracket (logicalTypeFromRep keyRep) destroyLogicalType \keyType -> - bracket (logicalTypeFromRep valueRep) destroyLogicalType $ + bracket (logicalTypeFromRepWith resolve keyRep) destroyLogicalType \keyType -> + bracket (logicalTypeFromRepWith resolve valueRep) destroyLogicalType $ c_duckdb_create_map_type keyType LogicalTypeStruct fieldArray -> do let fields = elems fieldArray - withMany (\field -> bracket (logicalTypeFromRep (structFieldValue field)) destroyLogicalType) fields \childTypes -> + withMany (\field -> bracket (logicalTypeFromRepWith resolve (structFieldValue field)) destroyLogicalType) fields \childTypes -> withMany withTypeName (map structFieldName fields) \names -> withArray names \nameArray -> withArray childTypes \typeArray -> c_duckdb_create_struct_type typeArray nameArray (fromIntegral (length fields)) LogicalTypeUnion memberArray -> do let members = elems memberArray - withMany (\member -> bracket (logicalTypeFromRep (unionMemberType member)) destroyLogicalType) members \memberTypes -> + withMany (\member -> bracket (logicalTypeFromRepWith resolve (unionMemberType member)) destroyLogicalType) members \memberTypes -> withMany withTypeName (map unionMemberName members) \names -> withArray names \nameArray -> withArray memberTypes \typeArray -> diff --git a/src/Database/DuckDB/Simple/Materialize.hs b/src/Database/DuckDB/Simple/Materialize.hs --- a/src/Database/DuckDB/Simple/Materialize.hs +++ b/src/Database/DuckDB/Simple/Materialize.hs @@ -1,40 +1,33 @@ {-# LANGUAGE BlockArguments #-} +{-# LANGUAGE LambdaCase #-} {-# LANGUAGE NamedFieldPuns #-} module Database.DuckDB.Simple.Materialize ( - materializeValue, + prepareValueReader, + prepareVectorReader, ) where -import Control.Exception (bracket, throwIO) +import Control.Exception (throwIO) import Control.Monad (forM, when) import Data.Array (Array, elems, listArray) -import Data.Bits (clearBit, shiftL, xor, (.|.)) -import qualified Data.ByteString as BS -import Data.Int (Int16, Int32, Int64, Int8) +import Data.Int (Int16, Int32, Int64) import qualified Data.Map.Strict as Map -import Data.Ratio ((%)) import Data.Text (Text) import qualified Data.Text as Text -import qualified Data.Text.Encoding as TextEncoding -import Data.Time.Calendar (addDays, fromGregorian) -import Data.Time.Clock.POSIX (posixSecondsToUTCTime) -import Data.Time.LocalTime ( - TimeOfDay (..), - minutesToTimeZone, - utc, - utcToLocalTime, - ) -import qualified Data.UUID as UUID import Data.Word (Word16, Word32, Word64, Word8) import Database.DuckDB.FFI +import Database.DuckDB.Simple.Element ( + chunkDecodeBlob, + chunkIsRowValid, + decodeElement, + duckDBHugeIntToInteger, + vectorElementType, + withVectorType, + ) import Database.DuckDB.Simple.FromField ( - BigNum (..), - BitString (..), DecimalValue (..), FieldValue (..), - IntervalValue (..), - TimeWithZone (..), - fromBigNumBytes, + RawGeometry (..), ) import Database.DuckDB.Simple.LogicalRep ( LogicalTypeRep (..), @@ -44,234 +37,106 @@ UnionValue (..), logicalTypeToRep, ) -import Database.DuckDB.Simple.Time (Date, LocalTimestamp, UTCTimestamp, Unbounded (..)) -import Foreign.C.Types (CBool (..)) -import Foreign.Marshal.Alloc (alloca) -import Foreign.Ptr (Ptr, castPtr, nullPtr, plusPtr) +import Database.DuckDB.Simple.VariantCodec (decodeVariant, prepareVariantDecoder) +import Foreign.Ptr (Ptr, castPtr, nullPtr) import Foreign.Storable (Storable (..), peekElemOff) -chunkIsRowValid :: Ptr Word64 -> DuckDBIdx -> IO Bool -chunkIsRowValid validity rowIdx - | validity == nullPtr = pure True - | otherwise = do - CBool flag <- c_duckdb_validity_row_is_valid validity rowIdx - pure (flag /= 0) - -chunkDecodeText :: Ptr () -> DuckDBIdx -> IO Text -chunkDecodeText dataPtr rowIdx = do - let base = castPtr dataPtr :: Ptr Word8 - offset = fromIntegral rowIdx * duckdbStringTSize - stringPtr = castPtr (base `plusPtr` offset) :: Ptr DuckDBStringT - len <- c_duckdb_string_t_length stringPtr - if len == 0 - then pure Text.empty - else do - cstr <- c_duckdb_string_t_data stringPtr - bytes <- BS.packCStringLen (cstr, fromIntegral len) - pure (TextEncoding.decodeUtf8 bytes) +-- | Read the type, data, and validity of a vector, and prepare its reader. +prepareVectorReader :: DuckDBVector -> IO (Int -> IO FieldValue) +prepareVectorReader vector = do + dtype <- vectorElementType vector + dataPtr <- c_duckdb_vector_get_data vector + validity <- c_duckdb_vector_get_validity vector + prepareValueReader dtype vector dataPtr validity -chunkDecodeBlob :: Ptr () -> DuckDBIdx -> IO BS.ByteString -chunkDecodeBlob dataPtr rowIdx = do - let base = castPtr dataPtr :: Ptr Word8 - offset = fromIntegral rowIdx * duckdbStringTSize - stringPtr = castPtr (base `plusPtr` offset) :: Ptr DuckDBStringT - len <- c_duckdb_string_t_length stringPtr - if len == 0 - then pure BS.empty - else do - ptr <- c_duckdb_string_t_data stringPtr - BS.packCStringLen (ptr, fromIntegral len) +-- | Prepare metadata once for a vector. The reader must not outlive its chunk. +prepareValueReader :: DuckDBType -> DuckDBVector -> Ptr () -> Ptr Word64 -> IO (Int -> IO FieldValue) +prepareValueReader dtype vector dataPtr validity = case dtype of + DuckDBTypeVariant -> prepareVariantDecoder vector + DuckDBTypeGeometry -> whenValid FieldGeometry <$> prepareGeometryDecoder vector dataPtr + DuckDBTypeStruct -> whenValid FieldStruct <$> prepareStructDecoder vector + DuckDBTypeUnion -> whenValid FieldUnion <$> prepareUnionDecoder vector + _ -> pure (materializeValue dtype vector dataPtr validity) + where + whenValid wrap decode row = do + valid <- chunkIsRowValid validity (fromIntegral row) + if valid then wrap <$> decode row else pure FieldNull -duckdbStringTSize :: Int -duckdbStringTSize = 16 +{- | Copy CRS metadata once for a vector. The decoder copies the WKB bytes of +a valid row. It must not outlive its chunk. +-} +prepareGeometryDecoder :: DuckDBVector -> Ptr () -> IO (Int -> IO RawGeometry) +prepareGeometryDecoder vector dataPtr = do + crs <- withVectorType vector \logical -> do + rep <- logicalTypeToRep logical + case rep of + LogicalTypeGeometry value -> pure value + _ -> throwIO (userError "duckdb-simple: invalid GEOMETRY type") + pure \row -> (`RawGeometry` crs) <$> chunkDecodeBlob dataPtr (fromIntegral row) materializeValue :: DuckDBType -> DuckDBVector -> Ptr () -> Ptr Word64 -> Int -> IO FieldValue materializeValue dtype vector dataPtr validity rowIdx = do - let duckIdx = fromIntegral rowIdx :: DuckDBIdx - valid <- chunkIsRowValid validity duckIdx + valid <- chunkIsRowValid validity (fromIntegral rowIdx) if not valid then pure FieldNull else case dtype of - DuckDBTypeBoolean -> do - raw <- peekElemOff (castPtr dataPtr :: Ptr Word8) rowIdx - pure (FieldBool (raw /= 0)) - DuckDBTypeTinyInt -> do - raw <- peekElemOff (castPtr dataPtr :: Ptr Int8) rowIdx - pure (FieldInt8 raw) - DuckDBTypeSmallInt -> do - raw <- peekElemOff (castPtr dataPtr :: Ptr Int16) rowIdx - pure (FieldInt16 raw) - DuckDBTypeInteger -> do - raw <- peekElemOff (castPtr dataPtr :: Ptr Int32) rowIdx - pure (FieldInt32 raw) - DuckDBTypeBigInt -> do - raw <- peekElemOff (castPtr dataPtr :: Ptr Int64) rowIdx - pure (FieldInt64 raw) - DuckDBTypeUTinyInt -> do - raw <- peekElemOff (castPtr dataPtr :: Ptr Word8) rowIdx - pure (FieldWord8 raw) - DuckDBTypeUSmallInt -> do - raw <- peekElemOff (castPtr dataPtr :: Ptr Word16) rowIdx - pure (FieldWord16 raw) - DuckDBTypeUInteger -> do - raw <- peekElemOff (castPtr dataPtr :: Ptr Word32) rowIdx - pure (FieldWord32 raw) - DuckDBTypeUBigInt -> do - raw <- peekElemOff (castPtr dataPtr :: Ptr Word64) rowIdx - pure (FieldWord64 raw) - DuckDBTypeFloat -> do - raw <- peekElemOff (castPtr dataPtr :: Ptr Float) rowIdx - pure (FieldFloat raw) - DuckDBTypeDouble -> do - raw <- peekElemOff (castPtr dataPtr :: Ptr Double) rowIdx - pure (FieldDouble raw) - DuckDBTypeVarchar -> FieldText <$> chunkDecodeText dataPtr duckIdx - DuckDBTypeUUID -> do - DuckDBUHugeInt lower upperBiased <- peekElemOff (castPtr dataPtr :: Ptr DuckDBUHugeInt) rowIdx - let upper = upperBiased `xor` (0x8000000000000000 :: Word64) - pure (FieldUUID (UUID.fromWords64 (fromIntegral upper) lower)) - DuckDBTypeBlob -> FieldBlob <$> chunkDecodeBlob dataPtr duckIdx - DuckDBTypeDate -> do - raw <- peekElemOff (castPtr dataPtr :: Ptr Int32) rowIdx - FieldDate <$> decodeDuckDBDate (DuckDBDate raw) - DuckDBTypeTime -> do - raw <- peekElemOff (castPtr dataPtr :: Ptr DuckDBTime) rowIdx - FieldTime <$> decodeDuckDBTime raw - DuckDBTypeTimeNs -> do - raw <- peekElemOff (castPtr dataPtr :: Ptr DuckDBTimeNs) rowIdx - pure (FieldTime (decodeDuckDBTimeNs raw)) - DuckDBTypeTimeTz -> do - raw <- peekElemOff (castPtr dataPtr :: Ptr DuckDBTimeTz) rowIdx - FieldTimeTZ <$> decodeDuckDBTimeTz raw - DuckDBTypeTimestamp -> do - raw <- peekElemOff (castPtr dataPtr :: Ptr DuckDBTimestamp) rowIdx - FieldTimestamp <$> decodeDuckDBTimestamp raw - DuckDBTypeTimestampS -> do - raw <- peekElemOff (castPtr dataPtr :: Ptr DuckDBTimestampS) rowIdx - FieldTimestamp <$> decodeDuckDBTimestampSeconds raw - DuckDBTypeTimestampMs -> do - raw <- peekElemOff (castPtr dataPtr :: Ptr DuckDBTimestampMs) rowIdx - FieldTimestamp <$> decodeDuckDBTimestampMilliseconds raw - DuckDBTypeTimestampNs -> do - raw <- peekElemOff (castPtr dataPtr :: Ptr DuckDBTimestampNs) rowIdx - FieldTimestamp <$> decodeDuckDBTimestampNanoseconds raw - DuckDBTypeTimestampTz -> do - raw <- peekElemOff (castPtr dataPtr :: Ptr DuckDBTimestamp) rowIdx - FieldTimestampTZ <$> decodeDuckDBTimestampUTCTime raw - DuckDBTypeInterval -> do - raw <- peekElemOff (castPtr dataPtr :: Ptr DuckDBInterval) rowIdx - pure (FieldInterval (intervalValueFromDuckDB raw)) + DuckDBTypeGeometry -> FieldGeometry <$> (prepareGeometryDecoder vector dataPtr >>= ($ rowIdx)) + DuckDBTypeVariant -> decodeVariant vector rowIdx DuckDBTypeDecimal -> - bracket - (c_duckdb_vector_get_column_type vector) - (\lty -> alloca $ \ptr -> poke ptr lty >> c_duckdb_destroy_logical_type ptr) - \logical -> do - width <- c_duckdb_decimal_width logical - scale <- c_duckdb_decimal_scale logical - internalTy <- c_duckdb_decimal_internal_type logical - rawValue <- - case internalTy of - DuckDBTypeSmallInt -> - toInteger <$> peekElemOff (castPtr dataPtr :: Ptr Int16) rowIdx - DuckDBTypeInteger -> - toInteger <$> peekElemOff (castPtr dataPtr :: Ptr Int32) rowIdx - DuckDBTypeBigInt -> - toInteger <$> peekElemOff (castPtr dataPtr :: Ptr Int64) rowIdx - DuckDBTypeHugeInt -> - toInteger . duckDBHugeIntToInteger - <$> peekElemOff (castPtr dataPtr :: Ptr DuckDBHugeInt) rowIdx - _ -> - error "duckdb-simple: unsupported decimal internal storage type" - pure (FieldDecimal (DecimalValue width scale rawValue)) - DuckDBTypeHugeInt -> do - raw <- peekElemOff (castPtr dataPtr :: Ptr DuckDBHugeInt) rowIdx - pure (FieldHugeInt (duckDBHugeIntToInteger raw)) - DuckDBTypeUHugeInt -> do - raw <- peekElemOff (castPtr dataPtr :: Ptr DuckDBUHugeInt) rowIdx - pure (FieldUHugeInt (duckDBUHugeIntToInteger raw)) - DuckDBTypeBit -> do - let base = castPtr dataPtr :: Ptr Word8 - offset = fromIntegral rowIdx * duckdbStringTSize - stringPtr = castPtr (base `plusPtr` offset) :: Ptr DuckDBStringT - len <- c_duckdb_string_t_length stringPtr - ptr <- c_duckdb_string_t_data stringPtr - bs <- BS.unpack <$> BS.packCStringLen (ptr, fromIntegral len) - case bs of - [] -> pure (FieldBit (BitString 0 BS.empty)) - [padding] -> pure (FieldBit (BitString padding BS.empty)) - (paddingByte : b : bits) -> do - let cleared = foldl clearBit b [8 - fromIntegral paddingByte .. 7] - pure (FieldBit (BitString paddingByte (BS.pack (cleared : bits)))) - DuckDBTypeBigNum -> do - let base = castPtr dataPtr :: Ptr Word8 - offset = fromIntegral rowIdx * duckdbStringTSize - stringPtr = castPtr (base `plusPtr` offset) :: Ptr DuckDBStringT - len <- c_duckdb_string_t_length stringPtr - if len < 3 - then pure (FieldBigNum (BigNum 0)) - else do - ptr <- c_duckdb_string_t_data stringPtr - bytes <- BS.unpack <$> BS.packCStringLen (ptr, fromIntegral len) - pure (FieldBigNum (BigNum (fromBigNumBytes bytes))) + withVectorType vector \logical -> do + width <- c_duckdb_decimal_width logical + scale <- c_duckdb_decimal_scale logical + internalTy <- c_duckdb_decimal_internal_type logical + rawValue <- + case internalTy of + DuckDBTypeSmallInt -> + toInteger <$> peekElemOff (castPtr dataPtr :: Ptr Int16) rowIdx + DuckDBTypeInteger -> + toInteger <$> peekElemOff (castPtr dataPtr :: Ptr Int32) rowIdx + DuckDBTypeBigInt -> + toInteger <$> peekElemOff (castPtr dataPtr :: Ptr Int64) rowIdx + DuckDBTypeHugeInt -> + duckDBHugeIntToInteger <$> peekElemOff (castPtr dataPtr :: Ptr DuckDBHugeInt) rowIdx + _ -> + error "duckdb-simple: unsupported decimal internal storage type" + pure (FieldDecimal (DecimalValue width scale rawValue)) DuckDBTypeArray -> FieldArray <$> decodeArrayElements vector rowIdx DuckDBTypeList -> FieldList <$> decodeListElements vector dataPtr rowIdx DuckDBTypeMap -> FieldMap <$> decodeMapPairs vector dataPtr rowIdx - DuckDBTypeStruct -> - FieldStruct <$> decodeStructValue vector rowIdx - DuckDBTypeUnion -> - FieldUnion <$> decodeUnionValue vector dataPtr rowIdx + DuckDBTypeStruct -> FieldStruct <$> (prepareStructDecoder vector >>= ($ rowIdx)) + DuckDBTypeUnion -> FieldUnion <$> (prepareUnionDecoder vector >>= ($ rowIdx)) DuckDBTypeEnum -> - bracket - (c_duckdb_vector_get_column_type vector) - (\lty -> alloca $ \ptr -> poke ptr lty >> c_duckdb_destroy_logical_type ptr) - \logical -> do - enumInternal <- c_duckdb_enum_internal_type logical - case enumInternal of - DuckDBTypeUTinyInt -> - FieldEnum . fromIntegral <$> peekElemOff (castPtr dataPtr :: Ptr Word8) rowIdx - DuckDBTypeUSmallInt -> - FieldEnum . fromIntegral <$> peekElemOff (castPtr dataPtr :: Ptr Word16) rowIdx - DuckDBTypeUInteger -> - FieldEnum <$> peekElemOff (castPtr dataPtr :: Ptr Word32) rowIdx - _ -> - error "duckdb-simple: unsupported enum internal storage type" - DuckDBTypeSQLNull -> - pure FieldNull - DuckDBTypeStringLiteral -> FieldText <$> chunkDecodeText dataPtr duckIdx - DuckDBTypeIntegerLiteral -> do - raw <- peekElemOff (castPtr dataPtr :: Ptr Int64) rowIdx - pure (FieldInt64 raw) - DuckDBTypeInvalid -> - error "duckdb-simple: INVALID type in eager result" - DuckDBTypeAny -> - error "duckdb-simple: ANY columns should not appear in results" - other -> - error ("duckdb-simple: UNKNOWN type in eager result: " <> show other) + withVectorType vector \logical -> do + enumInternal <- c_duckdb_enum_internal_type logical + case enumInternal of + DuckDBTypeUTinyInt -> + FieldEnum . fromIntegral <$> peekElemOff (castPtr dataPtr :: Ptr Word8) rowIdx + DuckDBTypeUSmallInt -> + FieldEnum . fromIntegral <$> peekElemOff (castPtr dataPtr :: Ptr Word16) rowIdx + DuckDBTypeUInteger -> + FieldEnum <$> peekElemOff (castPtr dataPtr :: Ptr Word32) rowIdx + _ -> + error "duckdb-simple: unsupported enum internal storage type" + DuckDBTypeSQLNull -> pure FieldNull + _ -> decodeElement dtype dataPtr rowIdx decodeArrayElements :: DuckDBVector -> Int -> IO (Array Int FieldValue) decodeArrayElements vector rowIdx = do arraySize <- - bracket - (c_duckdb_vector_get_column_type vector) - (\logical -> alloca $ \ptr -> poke ptr logical >> c_duckdb_destroy_logical_type ptr) - \logical -> do - sizeRaw <- c_duckdb_array_type_array_size logical - let sizeWord = fromIntegral sizeRaw :: Word64 - ensureWithinIntRange (Text.pack "array size") sizeWord + withVectorType vector \logical -> do + sizeRaw <- c_duckdb_array_type_array_size logical + let sizeWord = fromIntegral sizeRaw :: Word64 + ensureWithinIntRange (Text.pack "array size") sizeWord childVec <- c_duckdb_array_vector_get_child vector when (childVec == nullPtr) $ throwIO (userError "duckdb-simple: array child vector is null") - childType <- vectorElementType childVec - childData <- c_duckdb_vector_get_data childVec - childValidity <- c_duckdb_vector_get_validity childVec + readChild <- prepareVectorReader childVec let baseIdx = rowIdx * arraySize values <- forM [0 .. arraySize - 1] \delta -> - materializeValue childType childVec childData childValidity (baseIdx + delta) + readChild (baseIdx + delta) pure $ - if arraySize <= 0 - then listArray (0, -1) [] - else listArray (0, arraySize - 1) values + listArray (0, arraySize - 1) values decodeListElements :: DuckDBVector -> Ptr () -> Int -> IO [FieldValue] decodeListElements vector dataPtr rowIdx = do @@ -280,11 +145,9 @@ childVec <- c_duckdb_list_vector_get_child vector when (childVec == nullPtr) $ throwIO (userError "duckdb-simple: list child vector is null") - childType <- vectorElementType childVec - childData <- c_duckdb_vector_get_data childVec - childValidity <- c_duckdb_vector_get_validity childVec + readChild <- prepareVectorReader childVec forM [0 .. len - 1] \delta -> - materializeValue childType childVec childData childValidity (baseIdx + delta) + readChild (baseIdx + delta) decodeMapPairs :: DuckDBVector -> Ptr () -> Int -> IO [(FieldValue, FieldValue)] decodeMapPairs vector dataPtr rowIdx = do @@ -297,53 +160,49 @@ valueVec <- c_duckdb_struct_vector_get_child structVec 1 when (keyVec == nullPtr || valueVec == nullPtr) $ throwIO (userError "duckdb-simple: map child vectors are null") - keyType <- vectorElementType keyVec - valueType <- vectorElementType valueVec - keyData <- c_duckdb_vector_get_data keyVec - valueData <- c_duckdb_vector_get_data valueVec - keyValidity <- c_duckdb_vector_get_validity keyVec - valueValidity <- c_duckdb_vector_get_validity valueVec + readKey <- prepareVectorReader keyVec + readValue <- prepareVectorReader valueVec forM [0 .. len - 1] \delta -> do let childIdx = baseIdx + delta - keyValue <- materializeValue keyType keyVec keyData keyValidity childIdx - valueValue <- materializeValue valueType valueVec valueData valueValidity childIdx + keyValue <- readKey childIdx + valueValue <- readValue childIdx pure (keyValue, valueValue) -decodeStructValue :: DuckDBVector -> Int -> IO (StructValue FieldValue) -decodeStructValue vector rowIdx = - bracket - (c_duckdb_vector_get_column_type vector) - (\logical -> alloca $ \ptr -> poke ptr logical >> c_duckdb_destroy_logical_type ptr) - \logical -> do - structTypeRep <- logicalTypeToRep logical - structFields <- - case structTypeRep of - LogicalTypeStruct typeArray -> pure typeArray - other -> - throwIO - ( userError - ( "duckdb-simple: expected STRUCT logical type, but saw " - <> show other - ) +{- | Read the STRUCT type and prepare a reader for each child once for a vector. +The decoder reads a valid row. It must not outlive its chunk. +-} +prepareStructDecoder :: DuckDBVector -> IO (Int -> IO (StructValue FieldValue)) +prepareStructDecoder vector = + withVectorType vector \logical -> do + structTypeRep <- logicalTypeToRep logical + structFields <- + case structTypeRep of + LogicalTypeStruct typeArray -> pure typeArray + other -> + throwIO + ( userError + ( "duckdb-simple: expected STRUCT logical type, but saw " + <> show other ) - let typeList = elems structFields - count = length typeList + ) + let typeList = elems structFields + count = length typeList + indexMap = + Map.fromList (zip (map structFieldName typeList) [0 ..]) + childReaders <- + forM (zip [0 .. count - 1] typeList) \(childIdx, StructField{structFieldName}) -> do + childVec <- c_duckdb_struct_vector_get_child vector (fromIntegral childIdx) + when (childVec == nullPtr) $ + throwIO (userError "duckdb-simple: struct child vector is null") + readChild <- prepareVectorReader childVec + pure (structFieldName, readChild) + pure \rowIdx -> do valueFields <- - forM (zip [0 .. count - 1] typeList) \(childIdx, StructField{structFieldName}) -> do - childVec <- c_duckdb_struct_vector_get_child vector (fromIntegral childIdx) - when (childVec == nullPtr) $ - throwIO (userError "duckdb-simple: struct child vector is null") - childType <- vectorElementType childVec - childData <- c_duckdb_vector_get_data childVec - childValidity <- c_duckdb_vector_get_validity childVec - value <- materializeValue childType childVec childData childValidity rowIdx - pure StructField{structFieldName, structFieldValue = value} + forM childReaders \(name, readChild) -> do + value <- readChild rowIdx + pure StructField{structFieldName = name, structFieldValue = value} let fieldArray = - if count <= 0 - then listArray (0, -1) [] - else listArray (0, count - 1) valueFields - indexMap = - Map.fromList (zip (map structFieldName typeList) [0 ..]) + listArray (0, count - 1) valueFields pure StructValue { structValueFields = fieldArray @@ -351,89 +210,84 @@ , structValueIndex = indexMap } -decodeUnionValue :: DuckDBVector -> Ptr () -> Int -> IO (UnionValue FieldValue) -decodeUnionValue vector _dataPtr rowIdx = - bracket - (c_duckdb_vector_get_column_type vector) - (\logical -> alloca $ \ptr -> poke ptr logical >> c_duckdb_destroy_logical_type ptr) - \logical -> do - unionTypeRep <- logicalTypeToRep logical - membersArray <- - case unionTypeRep of - LogicalTypeUnion members -> pure members - other -> - throwIO - ( userError - ( "duckdb-simple: expected UNION logical type, but saw " - <> show other - ) - ) - let membersList = elems membersArray - memberCount = length membersList - tagVec <- c_duckdb_struct_vector_get_child vector 0 - when (tagVec == nullPtr) $ - throwIO (userError "duckdb-simple: union tag vector is null") - tagType <- vectorElementType tagVec - tagData <- c_duckdb_vector_get_data tagVec - tagValidity <- c_duckdb_vector_get_validity tagVec - tagValue <- materializeValue tagType tagVec tagData tagValidity rowIdx - memberIdx <- - case tagValue of - FieldWord8 tagWord -> pure (fromIntegral tagWord :: Int) - FieldWord16 tagWord -> pure (fromIntegral tagWord :: Int) - FieldWord32 tagWord -> - if tagWord <= fromIntegral (maxBound :: Word16) - then pure (fromIntegral tagWord) - else throwIO (userError "duckdb-simple: union tag exceeds Word16 range") - FieldWord64 tagWord -> - if tagWord <= fromIntegral (maxBound :: Word16) - then pure (fromIntegral tagWord) - else throwIO (userError "duckdb-simple: union tag exceeds Word16 range") - FieldInt8 tagInt - | tagInt >= 0 -> pure (fromIntegral tagInt) - | otherwise -> throwIO (userError "duckdb-simple: union tag negative") - FieldInt16 tagInt - | tagInt >= 0 -> pure (fromIntegral tagInt) - | otherwise -> throwIO (userError "duckdb-simple: union tag negative") - FieldInt32 tagInt - | tagInt >= 0 && tagInt <= fromIntegral (maxBound :: Word16) -> pure (fromIntegral tagInt) - | tagInt < 0 -> throwIO (userError "duckdb-simple: union tag negative") - | otherwise -> throwIO (userError "duckdb-simple: union tag exceeds Word16 range") - FieldInt64 tagInt - | tagInt >= 0 && tagInt <= fromIntegral (maxBound :: Word16) -> pure (fromIntegral tagInt) - | tagInt < 0 -> throwIO (userError "duckdb-simple: union tag negative") - | otherwise -> throwIO (userError "duckdb-simple: union tag exceeds Word16 range") - FieldNull -> - throwIO (userError "duckdb-simple: encountered NULL union tag") - other -> - throwIO - ( userError - ( "duckdb-simple: unexpected union tag value " - <> show other - ) +{- | Read the UNION type and prepare the tag reader and a reader for each member +once for a vector. The decoder reads a valid row. It must not outlive its chunk. +-} +prepareUnionDecoder :: DuckDBVector -> IO (Int -> IO (UnionValue FieldValue)) +prepareUnionDecoder vector = + withVectorType vector \logical -> do + unionTypeRep <- logicalTypeToRep logical + membersArray <- + case unionTypeRep of + LogicalTypeUnion members -> pure members + other -> + throwIO + ( userError + ( "duckdb-simple: expected UNION logical type, but saw " + <> show other ) + ) + let membersList = elems membersArray + memberCount = length membersList + tagVec <- c_duckdb_struct_vector_get_child vector 0 + when (tagVec == nullPtr) $ + throwIO (userError "duckdb-simple: union tag vector is null") + readTag <- prepareVectorReader tagVec + memberReaders <- + forM [1 .. memberCount] \childIdx -> do + memberVec <- c_duckdb_struct_vector_get_child vector (fromIntegral childIdx) + when (memberVec == nullPtr) $ + throwIO (userError "duckdb-simple: union member vector is null") + prepareVectorReader memberVec + pure \rowIdx -> do + memberIdx <- readTag rowIdx >>= unionTagIndex when (memberIdx < 0 || memberIdx >= memberCount) $ throwIO (userError "duckdb-simple: union tag out of range") - let selectedMember = membersList !! memberIdx - memberLabel = unionMemberName selectedMember - memberVec <- c_duckdb_struct_vector_get_child vector (fromIntegral (memberIdx + 1)) - when (memberVec == nullPtr) $ - throwIO (userError "duckdb-simple: union member vector is null") - memberType <- vectorElementType memberVec - memberData <- c_duckdb_vector_get_data memberVec - memberValidity <- c_duckdb_vector_get_validity memberVec - payload <- materializeValue memberType memberVec memberData memberValidity rowIdx + payload <- (memberReaders !! memberIdx) rowIdx pure UnionValue { unionValueIndex = fromIntegral memberIdx - , unionValueLabel = memberLabel + , unionValueLabel = unionMemberName (membersList !! memberIdx) , unionValuePayload = payload , unionValueMembers = membersArray } -vectorElementType :: DuckDBVector -> IO DuckDBType -vectorElementType vec = - bracket (c_duckdb_vector_get_column_type vec) destroyLogicalType c_duckdb_get_type_id +-- | Convert a decoded UNION tag to a member index. +unionTagIndex :: FieldValue -> IO Int +unionTagIndex = \case + FieldWord8 tagWord -> pure (fromIntegral tagWord :: Int) + FieldWord16 tagWord -> pure (fromIntegral tagWord :: Int) + FieldWord32 tagWord -> + if tagWord <= fromIntegral (maxBound :: Word16) + then pure (fromIntegral tagWord) + else throwIO (userError "duckdb-simple: union tag exceeds Word16 range") + FieldWord64 tagWord -> + if tagWord <= fromIntegral (maxBound :: Word16) + then pure (fromIntegral tagWord) + else throwIO (userError "duckdb-simple: union tag exceeds Word16 range") + FieldInt8 tagInt + | tagInt >= 0 -> pure (fromIntegral tagInt) + | otherwise -> throwIO (userError "duckdb-simple: union tag negative") + FieldInt16 tagInt + | tagInt >= 0 -> pure (fromIntegral tagInt) + | otherwise -> throwIO (userError "duckdb-simple: union tag negative") + FieldInt32 tagInt + | tagInt >= 0 && tagInt <= fromIntegral (maxBound :: Word16) -> pure (fromIntegral tagInt) + | tagInt < 0 -> throwIO (userError "duckdb-simple: union tag negative") + | otherwise -> throwIO (userError "duckdb-simple: union tag exceeds Word16 range") + FieldInt64 tagInt + | tagInt >= 0 && tagInt <= fromIntegral (maxBound :: Word16) -> pure (fromIntegral tagInt) + | tagInt < 0 -> throwIO (userError "duckdb-simple: union tag negative") + | otherwise -> throwIO (userError "duckdb-simple: union tag exceeds Word16 range") + FieldNull -> + throwIO (userError "duckdb-simple: encountered NULL union tag") + other -> + throwIO + ( userError + ( "duckdb-simple: unexpected union tag value " + <> show other + ) + ) listEntryBounds :: Text -> DuckDBListEntry -> IO (Int, Int) listEntryBounds context DuckDBListEntry{duckDBListEntryOffset, duckDBListEntryLength} = do @@ -452,102 +306,3 @@ in if actual <= limit then pure (fromInteger actual) else throwIO (userError ("duckdb-simple: " <> Text.unpack context <> " exceeds Int range")) - --- | Decode dates with exact epoch arithmetic and preserve infinity. -decodeDuckDBDate :: DuckDBDate -> IO Date -decodeDuckDBDate (DuckDBDate days) = - pure (decodeUnbounded (\value -> addDays (toInteger value) (fromGregorian 1970 1 1)) days) - -decodeDuckDBTime :: DuckDBTime -> IO TimeOfDay -decodeDuckDBTime raw = - alloca $ \ptr -> do - c_duckdb_from_time raw ptr - timeStruct <- peek ptr - pure (timeStructToTimeOfDay timeStruct) - -decodeDuckDBTimestamp :: DuckDBTimestamp -> IO LocalTimestamp -decodeDuckDBTimestamp (DuckDBTimestamp micros) = decodeTimestampUnits 1000000 micros - --- | Interpret native infinity sentinels before converting a finite payload. -decodeUnbounded :: (Integral a, Bounded a) => (a -> b) -> a -> Unbounded b -decodeUnbounded decode value - | value == maxBound = PosInfinity - | value == negate maxBound = NegInfinity - | otherwise = Finite (decode value) - --- | Decode timestamp units without overflowing an intermediate Int64. -decodeTimestampUnits :: Integer -> Int64 -> IO LocalTimestamp -decodeTimestampUnits units = - pure . decodeUnbounded (utcToLocalTime utc . posixSecondsToUTCTime . fromRational . (% units) . toInteger) - -decodeDuckDBTimeNs :: DuckDBTimeNs -> TimeOfDay -decodeDuckDBTimeNs (DuckDBTimeNs nanos) = - let (hours, remainderHours) = nanos `divMod` (60 * 60 * 1000000000) - (minutes, remainderMinutes) = remainderHours `divMod` (60 * 1000000000) - (seconds, fractionalNanos) = remainderMinutes `divMod` 1000000000 - fractional = fromRational (toInteger fractionalNanos % 1000000000) - totalSeconds = fromIntegral seconds + fractional - in TimeOfDay - (fromIntegral hours) - (fromIntegral minutes) - totalSeconds - -decodeDuckDBTimeTz :: DuckDBTimeTz -> IO TimeWithZone -decodeDuckDBTimeTz raw = - alloca $ \ptr -> do - c_duckdb_from_time_tz raw ptr - DuckDBTimeTzStruct{duckDBTimeTzStructTime = timeStruct, duckDBTimeTzStructOffset = offset} <- peek ptr - when (offset `rem` 60 /= 0) $ - throwIO (userError "duckdb-simple: TIMETZ offset cannot be represented in whole minutes") - let timeOfDay = timeStructToTimeOfDay timeStruct - minutes = fromIntegral offset `div` 60 - zone = minutesToTimeZone minutes - pure TimeWithZone{timeWithZoneTime = timeOfDay, timeWithZoneZone = zone} - -decodeDuckDBTimestampSeconds :: DuckDBTimestampS -> IO LocalTimestamp -decodeDuckDBTimestampSeconds (DuckDBTimestampS seconds) = - decodeTimestampUnits 1 seconds - -decodeDuckDBTimestampMilliseconds :: DuckDBTimestampMs -> IO LocalTimestamp -decodeDuckDBTimestampMilliseconds (DuckDBTimestampMs millis) = - decodeTimestampUnits 1000 millis - -decodeDuckDBTimestampNanoseconds :: DuckDBTimestampNs -> IO LocalTimestamp -decodeDuckDBTimestampNanoseconds (DuckDBTimestampNs nanos) = decodeTimestampUnits 1000000000 nanos - -decodeDuckDBTimestampUTCTime :: DuckDBTimestamp -> IO UTCTimestamp -decodeDuckDBTimestampUTCTime (DuckDBTimestamp micros) = - pure (decodeUnbounded (posixSecondsToUTCTime . fromRational . (% 1000000) . toInteger) micros) - -intervalValueFromDuckDB :: DuckDBInterval -> IntervalValue -intervalValueFromDuckDB DuckDBInterval{duckDBIntervalMonths, duckDBIntervalDays, duckDBIntervalMicros} = - IntervalValue - { intervalMonths = duckDBIntervalMonths - , intervalDays = duckDBIntervalDays - , intervalMicros = duckDBIntervalMicros - } - -duckDBHugeIntToInteger :: DuckDBHugeInt -> Integer -duckDBHugeIntToInteger DuckDBHugeInt{duckDBHugeIntLower, duckDBHugeIntUpper} = - (fromIntegral duckDBHugeIntUpper `shiftL` 64) .|. fromIntegral duckDBHugeIntLower - -duckDBUHugeIntToInteger :: DuckDBUHugeInt -> Integer -duckDBUHugeIntToInteger DuckDBUHugeInt{duckDBUHugeIntLower, duckDBUHugeIntUpper} = - (fromIntegral duckDBUHugeIntUpper `shiftL` 64) .|. fromIntegral duckDBUHugeIntLower - -destroyLogicalType :: DuckDBLogicalType -> IO () -destroyLogicalType logicalType = - alloca $ \ptr -> do - poke ptr logicalType - c_duckdb_destroy_logical_type ptr - -timeStructToTimeOfDay :: DuckDBTimeStruct -> TimeOfDay -timeStructToTimeOfDay DuckDBTimeStruct{duckDBTimeStructHour, duckDBTimeStructMinute, duckDBTimeStructSecond, duckDBTimeStructMicros} = - let secondsInt = fromIntegral duckDBTimeStructSecond :: Integer - micros = fromIntegral duckDBTimeStructMicros :: Integer - fractional = fromRational (micros % 1000000) - totalSeconds = fromInteger secondsInt + fractional - in TimeOfDay - (fromIntegral duckDBTimeStructHour) - (fromIntegral duckDBTimeStructMinute) - totalSeconds diff --git a/src/Database/DuckDB/Simple/Result.hs b/src/Database/DuckDB/Simple/Result.hs --- a/src/Database/DuckDB/Simple/Result.hs +++ b/src/Database/DuckDB/Simple/Result.hs @@ -17,10 +17,10 @@ import Data.IORef (IORef, atomicModifyIORef', readIORef, writeIORef) import qualified Data.Text as Text import Database.DuckDB.FFI -import Database.DuckDB.Simple.FromField (Field (..)) +import Database.DuckDB.Simple.FromField (Field (..), FieldValue) import Database.DuckDB.Simple.FromRow (RowParser, parseRow, rowErrorsToSqlError) import Database.DuckDB.Simple.Internal -import Database.DuckDB.Simple.Materialize (materializeValue) +import Database.DuckDB.Simple.Materialize (prepareVectorReader) import Database.DuckDB.Simple.Ok (Ok (..)) import Foreign.Marshal.Alloc (free, malloc) import Foreign.Marshal.Utils (fillBytes) @@ -72,7 +72,7 @@ case statementStreamChunk loaded of Nothing -> exhaustStatementStream streamRef >> pure Nothing Just chunk -> do - fields <- restore $ buildMaterializedRow (statementStreamColumns loaded) (statementStreamChunkVectors chunk) (statementStreamChunkIndex chunk) + fields <- restore $ buildMaterializedRow (statementStreamColumns loaded) (statementStreamChunkReaders chunk) (statementStreamChunkIndex chunk) parsed <- restore (evaluate (parseRow parser fields)) case parsed of Errors rowErr -> throwIO $ rowErrorsToSqlError (statementQuery stmt) rowErr @@ -124,30 +124,23 @@ destroyDataChunk chunk fetchChunk conn queryText stream else do - vectors <- - prepareChunkVectors chunk (statementStreamColumns stream) + readers <- + prepareChunkReaders chunk (statementStreamColumns stream) `onException` destroyDataChunk chunk let chunkState = StatementStreamChunk { statementStreamChunkPtr = chunk , statementStreamChunkSize = rowCount , statementStreamChunkIndex = 0 - , statementStreamChunkVectors = vectors + , statementStreamChunkReaders = readers } pure stream{statementStreamChunk = Just chunkState} -prepareChunkVectors :: DuckDBDataChunk -> [StatementStreamColumn] -> IO [StatementStreamChunkVector] -prepareChunkVectors chunk columns = - forM columns \StatementStreamColumn{statementStreamColumnIndex} -> do - vector <- c_duckdb_data_chunk_get_vector chunk (fromIntegral statementStreamColumnIndex) - dataPtr <- c_duckdb_vector_get_data vector - validity <- c_duckdb_vector_get_validity vector - pure - StatementStreamChunkVector - { statementStreamChunkVectorHandle = vector - , statementStreamChunkVectorData = dataPtr - , statementStreamChunkVectorValidity = validity - } +-- | Prepare one reader for each column. The readers must not outlive the chunk. +prepareChunkReaders :: DuckDBDataChunk -> [StatementStreamColumn] -> IO [Int -> IO FieldValue] +prepareChunkReaders chunk columns = + forM columns \StatementStreamColumn{statementStreamColumnIndex} -> + c_duckdb_data_chunk_get_vector chunk (fromIntegral statementStreamColumnIndex) >>= prepareVectorReader -- | Clear cursor ownership before releasing native resources. cleanupStatementStreamRef :: IORef StatementStreamState -> IO () @@ -197,8 +190,8 @@ if null columns then pure (Just (replicate rowCount [])) else do - vectors <- prepareChunkVectors chunk columns - rows <- mapM (buildMaterializedRow columns vectors) [0 .. rowCount - 1] + readers <- prepareChunkReaders chunk columns + rows <- mapM (buildMaterializedRow columns readers) [0 .. rowCount - 1] pure (Just rows) collectResultColumns :: Ptr DuckDBResult -> IO [StatementStreamColumn] @@ -219,19 +212,13 @@ , statementStreamColumnType = dtype } -buildMaterializedRow :: [StatementStreamColumn] -> [StatementStreamChunkVector] -> Int -> IO [Field] -buildMaterializedRow columns vectors rowIdx = - zipWithM (buildMaterializedField rowIdx) columns vectors +buildMaterializedRow :: [StatementStreamColumn] -> [Int -> IO FieldValue] -> Int -> IO [Field] +buildMaterializedRow columns readers rowIdx = + zipWithM (buildMaterializedField rowIdx) columns readers -buildMaterializedField :: Int -> StatementStreamColumn -> StatementStreamChunkVector -> IO Field -buildMaterializedField rowIdx column StatementStreamChunkVector{statementStreamChunkVectorHandle, statementStreamChunkVectorData, statementStreamChunkVectorValidity} = do - value <- - materializeValue - (statementStreamColumnType column) - statementStreamChunkVectorHandle - statementStreamChunkVectorData - statementStreamChunkVectorValidity - rowIdx +buildMaterializedField :: Int -> StatementStreamColumn -> (Int -> IO FieldValue) -> IO Field +buildMaterializedField rowIdx column readValue = do + value <- readValue rowIdx pure Field { fieldName = statementStreamColumnName column diff --git a/src/Database/DuckDB/Simple/ToField.hs b/src/Database/DuckDB/Simple/ToField.hs --- a/src/Database/DuckDB/Simple/ToField.hs +++ b/src/Database/DuckDB/Simple/ToField.hs @@ -25,11 +25,15 @@ ) where import Control.Exception (bracket, throwIO) -import Control.Monad (when) +import Control.Monad (filterM, forM_, when) import Data.Array (Array, elems) import Data.Bits (complement, shiftL, shiftR, (.&.), (.|.)) import qualified Data.ByteString as BS +import Data.Fixed (Pico) +import qualified Data.Geometry as G +import qualified Data.Geometry.WKT as WKT import Data.Int (Int16, Int32, Int64, Int8) +import qualified Data.List as List import Data.Proxy (Proxy (..)) import Data.Text (Text) import qualified Data.Text as Text @@ -44,7 +48,11 @@ import Database.DuckDB.Simple.Internal ( SQLError (..), Statement (..), + destroyValue, + duckDBTypeFromName, + fetchPrepareError, withStatementHandle, + withTypeCache, ) import Database.DuckDB.Simple.LogicalRep ( LogicalTypeRep (..), @@ -52,13 +60,17 @@ StructValue (..), UnionMemberType (..), UnionValue (..), + destroyLogicalType, logicalTypeFromRep, + logicalTypeFromRepWith, + logicalTypeToRep, structValueTypeRep, unionValueTypeRep, ) import Database.DuckDB.Simple.Time (Date, LocalTimestamp, UTCTimestamp, Unbounded (..)) +import Database.DuckDB.Simple.TypeCache (TypeCache, cachedLogicalType) import Database.DuckDB.Simple.Types (Null (..)) -import Foreign.C.String (peekCString) +import Database.DuckDB.Simple.Variant (Variant (..)) import Foreign.C.Types (CDouble (..), CFloat (..)) import Foreign.Marshal (fromBool) import Foreign.Marshal.Alloc (alloca) @@ -76,7 +88,7 @@ -- | Encapsulates the action required to bind a single positional parameter, together with a textual description used in diagnostics. data FieldBinding = FieldBinding - { fieldBindingAction :: !(Statement -> DuckDBIdx -> IO ()) + { fieldBindingValue :: !(TypeCache -> IO DuckDBValue) , fieldBindingDisplay :: !String } @@ -86,10 +98,20 @@ toDuckValue :: a -> IO DuckDBValue valueBinding :: String -> IO DuckDBValue -> FieldBinding -valueBinding display mkValue = - mkFieldBinding display $ \stmt idx -> - bindDuckValue stmt idx mkValue +valueBinding display = cacheValueBinding display . const +-- | Construct a value with the type cache of the statement's connection. +cacheValueBinding :: String -> (TypeCache -> IO DuckDBValue) -> FieldBinding +cacheValueBinding display makeValue = + FieldBinding + { fieldBindingValue = makeValue + , fieldBindingDisplay = display + } + +-- | Build types with the cached types for VARIANT and GEOMETRY with a CRS. +cachedTypeFromRep :: TypeCache -> LogicalTypeRep -> IO DuckDBLogicalType +cachedTypeFromRep = logicalTypeFromRepWith . cachedLogicalType + -- | Types that map to a concrete DuckDB column type when used with @ToField@. class DuckDBColumnType a where duckdbColumnTypeFor :: Proxy a -> Text @@ -100,19 +122,14 @@ -- | Apply a @FieldBinding@ to the given statement/index. bindFieldBinding :: Statement -> DuckDBIdx -> FieldBinding -> IO () -bindFieldBinding stmt idx FieldBinding{fieldBindingAction} = fieldBindingAction stmt idx +bindFieldBinding stmt idx FieldBinding{fieldBindingValue} = + withTypeCache (statementConnection stmt) \cache -> + bindDuckValue stmt idx (fieldBindingValue cache) -- | Render a bound parameter for error reporting. renderFieldBinding :: FieldBinding -> String renderFieldBinding FieldBinding{fieldBindingDisplay} = fieldBindingDisplay -mkFieldBinding :: String -> (Statement -> DuckDBIdx -> IO ()) -> FieldBinding -mkFieldBinding display action = - FieldBinding - { fieldBindingAction = action - , fieldBindingDisplay = display - } - -- | Types that can be used as positional parameters. class ToField a where toField :: a -> FieldBinding @@ -144,6 +161,16 @@ instance ToField Day instance ToField TimeOfDay instance ToField LocalTime + +-- | Bind the shape as @GEOMETRY@ with no CRS. +instance ToField G.Geometry + +-- | Bind the payload as a VARIANT with the type cache of the connection. +instance ToField Variant where + toField value = + cacheValueBinding (show value) \cache -> + variantDuckValue (cachedTypeFromRep cache) (variantPayload value) + instance ToField UTCTime instance ToField (Unbounded Day) instance ToField (Unbounded LocalTime) @@ -154,12 +181,14 @@ instance ToField (StructValue FieldValue) where toField structVal = - valueBinding "" (structValueDuckValue structVal) + cacheValueBinding "" \cache -> + structValueDuckValue (cachedTypeFromRep cache) structVal instance ToField (UnionValue FieldValue) where toField unionVal = let label = Text.unpack (unionValueLabel unionVal) - in valueBinding (" label <> ">") (unionValueDuckValue unionVal) + in cacheValueBinding (" label <> ">") \cache -> + unionValueDuckValue (cachedTypeFromRep cache) unionVal instance DuckDBColumnType BitString where duckdbColumnTypeFor _ = "BIT" @@ -170,11 +199,11 @@ ("") (toDuckValue bs) -instance (DuckDBColumnType a, ToDuckValue a) => ToField (Array Int a) where +instance (DuckDBColumnType a, ToField a) => ToField (Array Int a) where toField arr = - valueBinding + cacheValueBinding ("") - (arrayDuckValue arr) + \cache -> arrayDuckValue (cachedTypeFromRep cache) (\value -> fieldBindingValue (toField value) cache) arr instance (ToField a) => ToField (Maybe a) where toField Nothing = nullBinding "Nothing" @@ -184,6 +213,12 @@ { fieldBindingDisplay = "Just " <> renderFieldBinding binding } +instance DuckDBColumnType G.Geometry where + duckdbColumnTypeFor _ = "GEOMETRY" + +instance DuckDBColumnType Variant where + duckdbColumnTypeFor _ = "VARIANT" + instance DuckDBColumnType Null where duckdbColumnTypeFor _ = "NULL" @@ -420,19 +455,50 @@ utcTimestampDuckValue value = encodeUnbounded (encodeTimestampUnits 1000000 . utcToLocalTime utc) value >>= c_duckdb_create_timestamp_tz . DuckDBTimestamp +{- | Build an array value with a function for each element. A scalar element +type comes from the column type name of the element, so an empty array keeps +it. Other element types, such as STRUCT, UNION, and ARRAY, come from the +first element that is not NULL. All present elements must have the same type. +-} arrayDuckValue :: forall a. - (DuckDBColumnType a, ToDuckValue a) => + (DuckDBColumnType a) => + (LogicalTypeRep -> IO DuckDBLogicalType) -> + (a -> IO DuckDBValue) -> Array Int a -> IO DuckDBValue -arrayDuckValue arr = - bracket (createElementLogicalType (Proxy :: Proxy a)) destroyLogicalType \elementType -> - withCreatedValues (map toDuckValue (elems arr)) \values -> +arrayDuckValue typeFromRep elementValue arr = + withCreatedValues (map elementValue (elems arr)) \values -> + withElementType values \elementType -> withDuckValues values \ptr -> checkedValue (c_duckdb_create_array_value elementType ptr (fromIntegral (length values))) + where + typeName = duckdbColumnType (Proxy :: Proxy a) + withElementType values action = + case duckDBTypeFromName typeName of + Just dtype -> bracket (typeFromRep (LogicalTypeScalar dtype)) destroyLogicalType action + Nothing -> do + present <- filterM (fmap (== 0) . c_duckdb_is_null_value) values + case present of + -- The value owns this type. + value : rest -> do + logical <- c_duckdb_get_value_type value + expected <- logicalTypeToRep logical + forM_ rest \element -> do + actual <- c_duckdb_get_value_type element >>= logicalTypeToRep + when (actual /= expected) $ + throwIO (userError "duckdb-simple: array elements have different logical types") + action logical + [] -> + throwIO + SQLError + { sqlErrorMessage = "duckdb-simple: an empty or all-NULL array of " <> typeName <> " elements has no element type" + , sqlErrorType = Nothing + , sqlErrorQuery = Nothing + } -structValueDuckValue :: StructValue FieldValue -> IO DuckDBValue -structValueDuckValue StructValue{structValueFields, structValueTypes, structValueIndex = _} = do +structValueDuckValue :: (LogicalTypeRep -> IO DuckDBLogicalType) -> StructValue FieldValue -> IO DuckDBValue +structValueDuckValue typeFromRep StructValue{structValueFields, structValueTypes, structValueIndex = _} = do let valueFields = elems structValueFields typeFields = elems structValueTypes typeNames = map structFieldName typeFields @@ -444,17 +510,17 @@ let actions = zipWith ( \StructField{structFieldValue = typeRep} StructField{structFieldValue = fieldVal} -> - fieldValueWithTypeDuckValue typeRep fieldVal + fieldValueWithTypeDuckValue typeFromRep typeRep fieldVal ) typeFields valueFields - bracket (logicalTypeFromRep (LogicalTypeStruct structValueTypes)) destroyLogicalType \structLogical -> + bracket (typeFromRep (LogicalTypeStruct structValueTypes)) destroyLogicalType \structLogical -> withCreatedValues actions \childValues -> withDuckValues childValues $ \ptr -> checkedValue (c_duckdb_create_struct_value structLogical ptr) -unionValueDuckValue :: UnionValue FieldValue -> IO DuckDBValue -unionValueDuckValue UnionValue{unionValueIndex, unionValueLabel, unionValuePayload, unionValueMembers} = do +unionValueDuckValue :: (LogicalTypeRep -> IO DuckDBLogicalType) -> UnionValue FieldValue -> IO DuckDBValue +unionValueDuckValue typeFromRep UnionValue{unionValueIndex, unionValueLabel, unionValuePayload, unionValueMembers} = do let membersList = elems unionValueMembers idx = fromIntegral unionValueIndex :: Int memberCount = length membersList @@ -463,20 +529,25 @@ let UnionMemberType{unionMemberName, unionMemberType = memberType} = membersList !! idx when (unionValueLabel /= unionMemberName) $ throwIO (userError "duckdb-simple: union tag and member name mismatch") - bracket (logicalTypeFromRep (LogicalTypeUnion unionValueMembers)) destroyLogicalType \unionLogical -> - bracket (checkedValue (fieldValueWithTypeDuckValue memberType unionValuePayload)) destroyValue \payloadValue -> + bracket (typeFromRep (LogicalTypeUnion unionValueMembers)) destroyLogicalType \unionLogical -> + bracket (checkedValue (fieldValueWithTypeDuckValue typeFromRep memberType unionValuePayload)) destroyValue \payloadValue -> checkedValue (c_duckdb_create_union_value unionLogical (fromIntegral unionValueIndex) payloadValue) -fieldValueWithTypeDuckValue :: LogicalTypeRep -> FieldValue -> IO DuckDBValue -fieldValueWithTypeDuckValue typeRep FieldNull = - bracket (logicalTypeFromRep typeRep) destroyLogicalType \logical -> +fieldValueWithTypeDuckValue :: (LogicalTypeRep -> IO DuckDBLogicalType) -> LogicalTypeRep -> FieldValue -> IO DuckDBValue +fieldValueWithTypeDuckValue typeFromRep typeRep FieldNull = + bracket (typeFromRep typeRep) destroyLogicalType \logical -> withCreatedValues [nullDuckValue] \values -> withDuckValues values \ptr -> bracket (checkedValue (c_duckdb_create_list_value logical ptr 1)) destroyValue \list -> checkedValue (c_duckdb_get_list_child list 0) -fieldValueWithTypeDuckValue rep value = +fieldValueWithTypeDuckValue typeFromRep rep value = case rep of + LogicalTypeScalar DuckDBTypeVariant -> variantDuckValue typeFromRep value LogicalTypeScalar dtype -> scalarFieldValueDuckValue dtype value + LogicalTypeGeometry _ -> + case value of + FieldGeometry{} -> unsupportedRawGeometryBinding + other -> typeMismatch "GEOMETRY" other LogicalTypeDecimal width scale -> case value of FieldDecimal decVal@DecimalValue{decimalWidth, decimalScale} @@ -486,8 +557,8 @@ LogicalTypeList elemRep -> case value of FieldList elemsList -> - bracket (logicalTypeFromRep elemRep) destroyLogicalType \childLogical -> - withCreatedValues (map (fieldValueWithTypeDuckValue elemRep) elemsList) \values -> + bracket (typeFromRep elemRep) destroyLogicalType \childLogical -> + withCreatedValues (map (fieldValueWithTypeDuckValue typeFromRep elemRep) elemsList) \values -> withDuckValues values \ptr -> checkedValue (c_duckdb_create_list_value childLogical ptr (fromIntegral (length values))) other -> typeMismatch "LIST" other @@ -498,17 +569,17 @@ actualCount = length elemsList when (fromIntegral actualCount /= size) $ throwIO (userError "duckdb-simple: array length mismatch") - bracket (logicalTypeFromRep elemRep) destroyLogicalType \childLogical -> - withCreatedValues (map (fieldValueWithTypeDuckValue elemRep) elemsList) \values -> + bracket (typeFromRep elemRep) destroyLogicalType \childLogical -> + withCreatedValues (map (fieldValueWithTypeDuckValue typeFromRep elemRep) elemsList) \values -> withDuckValues values \ptr -> checkedValue (c_duckdb_create_array_value childLogical ptr (fromIntegral actualCount)) other -> typeMismatch "ARRAY" other LogicalTypeMap keyRep valueRep -> case value of FieldMap pairs -> - bracket (logicalTypeFromRep (LogicalTypeMap keyRep valueRep)) destroyLogicalType \mapLogical -> - withCreatedValues (map (fieldValueWithTypeDuckValue keyRep . fst) pairs) \keyValues -> - withCreatedValues (map (fieldValueWithTypeDuckValue valueRep . snd) pairs) \valValues -> + bracket (typeFromRep (LogicalTypeMap keyRep valueRep)) destroyLogicalType \mapLogical -> + withCreatedValues (map (fieldValueWithTypeDuckValue typeFromRep keyRep . fst) pairs) \keyValues -> + withCreatedValues (map (fieldValueWithTypeDuckValue typeFromRep valueRep . snd) pairs) \valValues -> withDuckValues keyValues \keyPtr -> withDuckValues valValues \valPtr -> checkedValue (c_duckdb_create_map_value mapLogical keyPtr valPtr (fromIntegral (length pairs))) @@ -516,13 +587,13 @@ LogicalTypeStruct structRep -> case value of FieldStruct structVal - | structValueTypeRep structVal == LogicalTypeStruct structRep -> structValueDuckValue structVal + | structValueTypeRep structVal == LogicalTypeStruct structRep -> structValueDuckValue typeFromRep structVal | otherwise -> throwIO (userError "duckdb-simple: struct value type mismatch") other -> typeMismatch "STRUCT" other LogicalTypeUnion unionRep -> case value of FieldUnion unionVal - | unionValueTypeRep unionVal == LogicalTypeUnion unionRep -> unionValueDuckValue unionVal + | unionValueTypeRep unionVal == LogicalTypeUnion unionRep -> unionValueDuckValue typeFromRep unionVal | otherwise -> throwIO (userError "duckdb-simple: union value type mismatch") other -> typeMismatch "UNION" other LogicalTypeEnum dict -> @@ -546,6 +617,7 @@ (DuckDBTypeDouble, FieldDouble d) -> doubleDuckValue d (DuckDBTypeVarchar, FieldText t) -> textDuckValue t (DuckDBTypeBlob, FieldBlob b) -> blobDuckValue b + (DuckDBTypeGeometry, FieldGeometry{}) -> unsupportedRawGeometryBinding (DuckDBTypeUUID, FieldUUID u) -> uuidDuckValue u (DuckDBTypeBit, FieldBit bits) -> bitDuckValue bits (DuckDBTypeDate, FieldDate d) -> dateDuckValue d @@ -680,57 +752,92 @@ ) ) -createElementLogicalType :: forall a. (DuckDBColumnType a) => Proxy a -> IO DuckDBLogicalType -createElementLogicalType proxy = - let typeName = duckdbColumnType proxy - in case duckDBTypeFromName typeName of - Just dtype -> c_duckdb_create_logical_type dtype - Nothing -> - throwIO - ( SQLError - { sqlErrorMessage = - Text.concat - [ "duckdb-simple: unsupported array element type " - , typeName - ] - , sqlErrorType = Nothing - , sqlErrorQuery = Nothing - } - ) +-- | Reject raw values that the C API cannot bind without format conversion. +unsupportedRawGeometryBinding :: IO a +unsupportedRawGeometryBinding = + throwIO (userError "duckdb-simple: raw GEOMETRY binding requires explicit ST_GeomFromWKB and ST_SetCRS parameters") -duckDBTypeFromName :: Text -> Maybe DuckDBType -duckDBTypeFromName name = - case name of - "BOOLEAN" -> Just DuckDBTypeBoolean - "TINYINT" -> Just DuckDBTypeTinyInt - "SMALLINT" -> Just DuckDBTypeSmallInt - "INTEGER" -> Just DuckDBTypeInteger - "BIGINT" -> Just DuckDBTypeBigInt - "UTINYINT" -> Just DuckDBTypeUTinyInt - "USMALLINT" -> Just DuckDBTypeUSmallInt - "UINTEGER" -> Just DuckDBTypeUInteger - "UBIGINT" -> Just DuckDBTypeUBigInt - "FLOAT" -> Just DuckDBTypeFloat - "DOUBLE" -> Just DuckDBTypeDouble - "DATE" -> Just DuckDBTypeDate - "TIME" -> Just DuckDBTypeTime - "TIMESTAMP" -> Just DuckDBTypeTimestamp - "TIMESTAMPTZ" -> Just DuckDBTypeTimestampTz - "TEXT" -> Just DuckDBTypeVarchar - "BLOB" -> Just DuckDBTypeBlob - "UUID" -> Just DuckDBTypeUUID - "BIT" -> Just DuckDBTypeBit - "BIGNUM" -> Just DuckDBTypeBigNum - -- treat NULL as SQLNULL to provide element type for Maybe values without data - "NULL" -> Just DuckDBTypeSQLNull - _ -> Nothing +{- | Construct an owned VARIANT value. A scalar keeps its native type. Lists, +arrays, and STRUCT fields contain VARIANT values. The C API casts the payload +through a one-element VARIANT list. +-} +variantDuckValue :: (LogicalTypeRep -> IO DuckDBLogicalType) -> FieldValue -> IO DuckDBValue +variantDuckValue typeFromRep value = do + (rep, payload) <- variantPayloadType value + bracket (typeFromRep (LogicalTypeScalar DuckDBTypeVariant)) destroyLogicalType \variantType -> + withCreatedValues [fieldValueWithTypeDuckValue typeFromRep rep payload] \values -> + withDuckValues values \ptr -> + bracket (checkedValue (c_duckdb_create_list_value variantType ptr 1)) destroyValue \list -> + checkedValue (c_duckdb_get_list_child list 0) -destroyLogicalType :: DuckDBLogicalType -> IO () -destroyLogicalType logical = - alloca $ \ptr -> do - poke ptr logical - c_duckdb_destroy_logical_type ptr +{- | Choose the native type of a VARIANT payload. Containers get VARIANT +elements and fields. Time values with sub-microsecond digits use nanosecond +types. Wide timestamps use milliseconds or seconds when these preserve the +value and fit the native range. +-} +variantPayloadType :: FieldValue -> IO (LogicalTypeRep, FieldValue) +variantPayloadType value = case value of + FieldNull -> pure (variant, value) + FieldBool{} -> scalar DuckDBTypeBoolean + FieldInt8{} -> scalar DuckDBTypeTinyInt + FieldInt16{} -> scalar DuckDBTypeSmallInt + FieldInt32{} -> scalar DuckDBTypeInteger + FieldInt64{} -> scalar DuckDBTypeBigInt + FieldWord8{} -> scalar DuckDBTypeUTinyInt + FieldWord16{} -> scalar DuckDBTypeUSmallInt + FieldWord32{} -> scalar DuckDBTypeUInteger + FieldWord64{} -> scalar DuckDBTypeUBigInt + FieldHugeInt{} -> scalar DuckDBTypeHugeInt + FieldUHugeInt{} -> scalar DuckDBTypeUHugeInt + FieldFloat{} -> scalar DuckDBTypeFloat + FieldDouble{} -> scalar DuckDBTypeDouble + FieldDecimal DecimalValue{decimalWidth, decimalScale} -> pure (LogicalTypeDecimal decimalWidth decimalScale, value) + FieldText{} -> scalar DuckDBTypeVarchar + FieldBlob{} -> scalar DuckDBTypeBlob + FieldUUID{} -> scalar DuckDBTypeUUID + FieldDate{} -> scalar DuckDBTypeDate + FieldTime time + | hasNanos time -> scalar DuckDBTypeTimeNs + | otherwise -> scalar DuckDBTypeTime + FieldTimestamp (Finite LocalTime{localDay, localTimeOfDay}) + | hasNanos localTimeOfDay -> scalar DuckDBTypeTimestampNs + | inFiniteRange (minBound :: Int64) maxBound micros -> scalar DuckDBTypeTimestamp + | micros `rem` 1000 == 0 && inFiniteRange (minBound :: Int64) maxBound (micros `div` 1000) -> scalar DuckDBTypeTimestampMs + | micros `rem` 1000000 == 0 && inFiniteRange (minBound :: Int64) maxBound (micros `div` 1000000) -> scalar DuckDBTypeTimestampS + where + micros = diffDays localDay (fromGregorian 1970 1 1) * 86400 * 1000000 + diffTimeToPicoseconds (timeOfDayToTime localTimeOfDay) `div` 1000000 + FieldTimestamp{} -> scalar DuckDBTypeTimestamp + FieldTimestampTZ{} -> scalar DuckDBTypeTimestampTz + FieldTimeTZ{} -> scalar DuckDBTypeTimeTz + FieldInterval{} -> scalar DuckDBTypeInterval + FieldBigNum{} -> scalar DuckDBTypeBigNum + FieldBit{} -> scalar DuckDBTypeBit + FieldList{} -> pure (LogicalTypeList variant, value) + FieldArray items -> pure (LogicalTypeList variant, FieldList (elems items)) + FieldStruct structValue@StructValue{structValueTypes} -> do + let names = map structFieldName (elems structValueTypes) + when (any Text.null names || length names /= length (List.nub names)) $ + throwIO (userError "duckdb-simple: VARIANT objects need unique, nonempty keys") + let types = fmap (\field -> field{structFieldValue = variant}) structValueTypes + pure (LogicalTypeStruct types, FieldStruct structValue{structValueTypes = types}) + FieldUnion unionValue -> pure (unionValueTypeRep unionValue, value) + FieldGeometry{} -> unsupportedRawGeometryBinding + FieldMap{} -> throwIO (userError "duckdb-simple: VARIANT payloads cannot contain MAP values") + FieldEnum{} -> throwIO (userError "duckdb-simple: VARIANT payloads cannot contain ENUM values") + where + variant = LogicalTypeScalar DuckDBTypeVariant + scalar dtype = pure (LogicalTypeScalar dtype, value) + hasNanos time = snd (properFraction (todSec time * 1000000) :: (Integer, Pico)) /= 0 +instance ToDuckValue G.Geometry where + toDuckValue geometry = do + wkt <- either (throwIO . userError) pure (WKT.encodeWKT geometry) + bracket (logicalTypeFromRep (LogicalTypeGeometry Nothing)) destroyLogicalType \logical -> + withCreatedValues [textDuckValue wkt] \values -> + withDuckValues values \ptr -> + bracket (checkedValue (c_duckdb_create_list_value logical ptr 1)) destroyValue \list -> + checkedValue (c_duckdb_get_list_child list 0) + instance ToDuckValue Null where toDuckValue _ = nullDuckValue @@ -819,11 +926,17 @@ toDuckValue = utcTimestampDuckValue instance ToDuckValue (StructValue FieldValue) where - toDuckValue = structValueDuckValue + toDuckValue = structValueDuckValue logicalTypeFromRep instance ToDuckValue (UnionValue FieldValue) where - toDuckValue = unionValueDuckValue + toDuckValue = unionValueDuckValue logicalTypeFromRep +{- | Build an array without a connection. The elements need 'ToDuckValue', so +this instance does not accept t'Variant' elements. 'toField' binds those. +-} +instance (DuckDBColumnType a, ToDuckValue a) => ToDuckValue (Array Int a) where + toDuckValue = arrayDuckValue logicalTypeFromRep toDuckValue + instance (ToDuckValue a) => ToDuckValue (Maybe a) where toDuckValue Nothing = nullDuckValue toDuckValue (Just value) = toDuckValue value @@ -860,9 +973,14 @@ -- | Check storage limits and the two DuckDB infinity sentinels. checkFiniteRange :: (Integral a) => String -> a -> a -> Integer -> IO () checkFiniteRange label lower upper value = - when (value < toInteger lower || value >= toInteger upper || value == negate (toInteger upper)) $ + when (not (inFiniteRange lower upper value)) $ throwIO (userError ("duckdb-simple: " <> label <> " value out of finite range")) +-- | Check storage limits without accepting the two infinity sentinels. +inFiniteRange :: (Integral a) => a -> a -> Integer -> Bool +inFiniteRange lower upper value = + value >= toInteger lower && value < toInteger upper && value /= negate (toInteger upper) + -- | Validate time components and convert to the requested units per second. timeOfDayUnits :: Integer -> TimeOfDay -> IO Integer timeOfDayUnits units tod@(TimeOfDay hours minutes seconds) = do @@ -879,21 +997,8 @@ bracket (checkedValue makeValue) destroyValue \value -> do rc <- c_duckdb_bind_value handle idx value when (rc /= DuckDBSuccess) $ do - err <- fetchPrepareError handle + err <- fetchPrepareError (Text.pack "duckdb-simple: parameter binding failed") handle throwBindError stmt err - -destroyValue :: DuckDBValue -> IO () -destroyValue value = - alloca \ptr -> do - poke ptr value - c_duckdb_destroy_value ptr - -fetchPrepareError :: DuckDBPreparedStatement -> IO Text -fetchPrepareError handle = do - msgPtr <- c_duckdb_prepare_error handle - if msgPtr == nullPtr - then pure (Text.pack "duckdb-simple: parameter binding failed") - else Text.pack <$> peekCString msgPtr throwBindError :: Statement -> Text -> IO a throwBindError Statement{statementQuery} msg = diff --git a/src/Database/DuckDB/Simple/TypeCache.hs b/src/Database/DuckDB/Simple/TypeCache.hs new file mode 100644 --- /dev/null +++ b/src/Database/DuckDB/Simple/TypeCache.hs @@ -0,0 +1,174 @@ +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE OverloadedStrings #-} + +{- | Native types that the C API cannot create. A connection reads them the +first time a parameter needs one, and destroys them when it closes. +-} +module Database.DuckDB.Simple.TypeCache ( + TypeCache, + defaultGeometryCRS, + createTypeCache, + destroyTypeCache, + cachedLogicalType, +) where + +import Control.Concurrent.MVar (MVar, modifyMVarMasked, modifyMVar_, newMVar) +import Control.Exception (bracket, finally, mask_, onException, throwIO) +import Control.Monad (forM_, when) +import qualified Data.ByteString as BS +import Data.List (nub) +import Data.Map.Strict (Map) +import qualified Data.Map.Strict as Map +import Data.Text (Text) +import qualified Data.Text as Text +import qualified Data.Text.Encoding as TextEncoding +import Database.DuckDB.FFI +import Database.DuckDB.Simple.LogicalRep (LogicalTypeRep (..), destroyLogicalType, logicalTypeToRep) +import Foreign.Marshal.Alloc (alloca) +import Foreign.Marshal.Utils (fillBytes) +import Foreign.Ptr (Ptr, nullPtr) +import Foreign.Storable (peek, poke, sizeOf) + +{- | The configured CRSs and the native types of a connection. The types are +absent until a parameter needs one. The connection owns the types. +-} +data TypeCache = TypeCache + { typeCacheDatabase :: !DuckDBDatabase + , typeCacheCRS :: ![Text] + , typeCacheTypes :: !(MVar (Maybe NativeTypes)) + } + +{- | The VARIANT type and the GEOMETRY types for the configured CRSs. Each +GEOMETRY type has two keys: the configured CRS text and the CRS text that +DuckDB reports for the type. +-} +data NativeTypes = NativeTypes + { nativeVariant :: !DuckDBLogicalType + , nativeGeometry :: !(Map Text DuckDBLogicalType) + , nativeOwned :: ![DuckDBLogicalType] + } + +-- | The CRS that a connection reads when the options do not give a list. +defaultGeometryCRS :: [Text] +defaultGeometryCRS = ["OGC:CRS84"] + +{- | Make an empty cache for a database. The caller must destroy the cache. A +CRS must be nonempty and must not contain NUL. +-} +createTypeCache :: DuckDBDatabase -> [Text] -> IO TypeCache +createTypeCache database crsList = do + let crss = nub crsList + when (any Text.null crss || any (Text.any (== '\0')) crss) $ + throwIO (userError "duckdb-simple: a GEOMETRY CRS must be nonempty and must not contain NUL") + TypeCache database crss <$> newMVar Nothing + +-- | Destroy the types in the cache, if the connection read them. +destroyTypeCache :: TypeCache -> IO () +destroyTypeCache TypeCache{typeCacheTypes} = + modifyMVar_ typeCacheTypes \types -> Nothing <$ mapM_ (mapM_ destroyLogicalType . nativeOwned) types + +{- | Copy a cached type for a leaf that the C API cannot create. The caller +must destroy the copy. The first VARIANT leaf or GEOMETRY leaf with a CRS +reads the types. A GEOMETRY CRS that the cache does not hold gives GEOMETRY +without a CRS. +-} +cachedLogicalType :: TypeCache -> LogicalTypeRep -> IO DuckDBLogicalType +cachedLogicalType cache = \case + LogicalTypeScalar DuckDBTypeVariant -> nativeTypes cache >>= copyLogicalType . nativeVariant + LogicalTypeGeometry (Just crs) -> do + NativeTypes{nativeGeometry} <- nativeTypes cache + maybe (c_duckdb_create_logical_type DuckDBTypeGeometry) copyLogicalType (Map.lookup crs nativeGeometry) + LogicalTypeGeometry Nothing -> c_duckdb_create_logical_type DuckDBTypeGeometry + other -> throwIO (userError ("duckdb-simple: the type cache cannot create " <> show other)) + +{- | Get the native types. Read them on first use with a separate connection, +so the query does not run in the transaction of the caller. A failed read +leaves the cache empty. +-} +nativeTypes :: TypeCache -> IO NativeTypes +nativeTypes TypeCache{typeCacheDatabase, typeCacheCRS, typeCacheTypes} = + modifyMVarMasked typeCacheTypes \cached -> do + types <- maybe (withTypeConnection typeCacheDatabase (`readNativeTypes` typeCacheCRS)) pure cached + pure (Just types, types) + +-- | Read the VARIANT type and a GEOMETRY type for each CRS with one query. +readNativeTypes :: DuckDBConnection -> [Text] -> IO NativeTypes +readNativeTypes connection crss = do + let sql = Text.concat ("SELECT NULL::VARIANT" : [", system.main.ST_SetCRS('POINT EMPTY'::GEOMETRY, ?)" | _ <- crss]) + withTypeQuery connection sql crss \result -> mask_ do + owned <- columnTypes result (length crss + 1) + case owned of + [] -> throwIO (userError "duckdb-simple: the type query returned no columns") + variant : geometry -> do + reported <- mapM logicalTypeToRep geometry `onException` mapM_ destroyLogicalType owned + pure + NativeTypes + { nativeVariant = variant + , nativeGeometry = + Map.fromList + ( concat + [ (crs, logical) : [(reportedCRS, logical) | LogicalTypeGeometry (Just reportedCRS) <- [rep]] + | (crs, rep, logical) <- zip3 crss reported geometry + ] + ) + , nativeOwned = owned + } + +-- | Take the types of the first columns. Destroy the taken types on failure. +columnTypes :: Ptr DuckDBResult -> Int -> IO [DuckDBLogicalType] +columnTypes result count = go [] 0 + where + go taken column + | column == count = pure (reverse taken) + | otherwise = do + logical <- c_duckdb_column_logical_type result (fromIntegral column) `onException` mapM_ destroyLogicalType taken + when (logical == nullPtr) do + mapM_ destroyLogicalType taken + throwIO (userError "duckdb-simple: the type query returned no type") + go (logical : taken) (column + 1) + +-- | Copy a type through a LIST type, because the C API has no copy function. +copyLogicalType :: DuckDBLogicalType -> IO DuckDBLogicalType +copyLogicalType logical = + -- TODO: use a copy function when the C API has one. + -- See https://github.com/Tritlo/duckdb-haskell/issues/30 and + -- https://github.com/duckdb/duckdb/issues/26664. + bracket (c_duckdb_create_list_type logical) destroyLogicalType c_duckdb_list_type_child_type + +-- | Run an action with a new connection to the database. Disconnect after it. +withTypeConnection :: DuckDBDatabase -> (DuckDBConnection -> IO a) -> IO a +withTypeConnection database action = + alloca \connectionPtr -> do + connected <- c_duckdb_connect database connectionPtr + when (connected /= DuckDBSuccess) $ + throwIO (userError "duckdb-simple: cannot connect to read the VARIANT and GEOMETRY types") + (peek connectionPtr >>= action) `finally` c_duckdb_disconnect connectionPtr + +-- | Prepare and run a constant query with text parameters, and borrow its result. +withTypeQuery :: DuckDBConnection -> Text -> [Text] -> (Ptr DuckDBResult -> IO a) -> IO a +withTypeQuery connection sql parameters action = + BS.useAsCString (TextEncoding.encodeUtf8 sql) \sqlPtr -> + alloca \statementPtr -> do + poke statementPtr nullPtr + bracket (c_duckdb_prepare connection sqlPtr statementPtr) (const (c_duckdb_destroy_prepare statementPtr)) \prepared -> do + statement <- peek statementPtr + when (prepared /= DuckDBSuccess) do + errorPtr <- c_duckdb_prepare_error statement + message <- if errorPtr == nullPtr then pure "prepare failed" else TextEncoding.decodeUtf8 <$> BS.packCString errorPtr + queryFailed message + forM_ (zip [1 ..] parameters) \(index, parameter) -> + BS.useAsCStringLen (TextEncoding.encodeUtf8 parameter) \(ptr, len) -> do + bound <- c_duckdb_bind_varchar_length statement index ptr (fromIntegral len) + when (bound /= DuckDBSuccess) (queryFailed ("cannot bind CRS " <> parameter)) + alloca \result -> do + fillBytes result 0 (sizeOf (undefined :: DuckDBResult)) + bracket (c_duckdb_execute_prepared statement result) (const (c_duckdb_destroy_result result)) \executed -> do + when (executed /= DuckDBSuccess) do + errorPtr <- c_duckdb_result_error result + message <- if errorPtr == nullPtr then pure "execution failed" else TextEncoding.decodeUtf8 <$> BS.packCString errorPtr + queryFailed message + action result + where + queryFailed message = throwIO (userError ("duckdb-simple: cannot read the VARIANT and GEOMETRY types: " <> Text.unpack message)) diff --git a/src/Database/DuckDB/Simple/Types.hs b/src/Database/DuckDB/Simple/Types.hs --- a/src/Database/DuckDB/Simple/Types.hs +++ b/src/Database/DuckDB/Simple/Types.hs @@ -22,17 +22,28 @@ import Control.Exception (Exception) import Data.Text (Text) +import qualified Data.Text as Text +import Database.DuckDB.Simple.FromField (Field (..), FieldValue (..), FromField (..), ResultError (..), returnError) import Database.DuckDB.Simple.Internal ( Connection, Query (..), SQLError (..), Statement, ) +import Database.DuckDB.Simple.Ok (Ok (..)) -- | Placeholder representing SQL @NULL@. data Null = Null deriving (Eq, Ord, Show, Read) + +{- | This instance is in this module because @Internal@ imports @FromField@. +An import of this module from @FromField@ would cause an import cycle. +-} +instance FromField Null where + fromField f = case fieldValue f of + FieldNull -> Ok Null + _ -> returnError Incompatible f (Text.pack "expected NULL") -- | Wrapper used for single-column rows. newtype Only a = Only {fromOnly :: a} diff --git a/src/Database/DuckDB/Simple/Variant.hs b/src/Database/DuckDB/Simple/Variant.hs new file mode 100644 --- /dev/null +++ b/src/Database/DuckDB/Simple/Variant.hs @@ -0,0 +1,40 @@ +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE PatternSynonyms #-} + +-- | DuckDB VARIANT values. +module Database.DuckDB.Simple.Variant ( + Variant (..), + variantObject, +) where + +import Data.Array (listArray) +import qualified Data.Map.Strict as Map +import Data.Text (Text) +import Database.DuckDB.FFI (pattern DuckDBTypeVariant) +import Database.DuckDB.Simple.FromField (Field (..), FieldValue (..), FromField (..)) +import Database.DuckDB.Simple.LogicalRep (LogicalTypeRep (..), StructField (..), StructValue (..)) +import Database.DuckDB.Simple.Ok (Ok (..)) + +{- | A VARIANT value. The payload is a 'FieldValue' with the native type of the +stored value. Arrays are 'FieldList' values. Objects are 'FieldStruct' values +whose fields have the VARIANT type. SQL NULL is 'FieldNull'. +A t'Variant' parameter binds its payload as a VARIANT. +-} +newtype Variant = Variant {variantPayload :: FieldValue} + deriving (Eq, Show) + +{- | Build the payload of a VARIANT object. Each field has the VARIANT type, +in the order of the entries. +-} +variantObject :: [(Text, FieldValue)] -> FieldValue +variantObject entries = + FieldStruct + StructValue + { structValueFields = listArray (0, length entries - 1) [StructField name value | (name, value) <- entries] + , structValueTypes = listArray (0, length entries - 1) [StructField name (LogicalTypeScalar DuckDBTypeVariant) | (name, _) <- entries] + , structValueIndex = Map.fromList (zip (map fst entries) [0 ..]) + } + +-- | Read any column. A VARIANT column gives its decoded payload. +instance FromField Variant where + fromField Field{fieldValue} = Ok (Variant fieldValue) diff --git a/src/Database/DuckDB/Simple/VariantCodec.hs b/src/Database/DuckDB/Simple/VariantCodec.hs new file mode 100644 --- /dev/null +++ b/src/Database/DuckDB/Simple/VariantCodec.hs @@ -0,0 +1,452 @@ +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} + +{- | Checked access to DuckDB 1.5's private VARIANT payload. +Only public C handles and vector accessors cross the native boundary. +-} +module Database.DuckDB.Simple.VariantCodec ( + decodeVariant, + prepareVariantDecoder, + decodeVariantPayload, +) where + +import Control.Exception (bracket, throwIO) +import Control.Monad (forM, forM_, unless, when) +import Control.Monad.Trans.Class (lift) +import Control.Monad.Trans.State.Strict (StateT, evalStateT, gets, modify') +import Data.Array (Array, bounds, listArray, (!)) +import Data.Bits (complement, finiteBitSize, shiftL, (.&.), (.|.)) +import Data.ByteString (ByteString) +import qualified Data.ByteString as BS +import qualified Data.IntMap.Strict as IntMap +import qualified Data.IntSet as IntSet +import qualified Data.Set as Set +import Data.Text (Text) +import qualified Data.Text as Text +import qualified Data.Text.Encoding as Text +import Data.Word (Word32, Word8) +import Database.DuckDB.FFI +import Database.DuckDB.Simple.Element (bitStringFromBytes, chunkDecodeBlob, chunkIsRowValid, decodeElement) +import Database.DuckDB.Simple.FromField ( + BigNum (..), + DecimalValue (..), + FieldValue (..), + RawGeometry (..), + fromBigNumBytes, + ) +import Database.DuckDB.Simple.Internal (destroyLogicalType) +import Database.DuckDB.Simple.Variant (variantObject) +import Foreign.C.String (peekCString) +import Foreign.Marshal.Alloc (alloca, allocaBytesAligned) +import Foreign.Marshal.Utils (copyBytes) +import Foreign.Ptr (Ptr, castPtr, nullPtr) +import Foreign.Storable (Storable, peek, peekElemOff, poke, sizeOf) +import Text.Read (readMaybe) + +-- | Raise a codec error before an invalid native operation. +codecError :: String -> IO a +codecError message = throwIO (userError ("duckdb-simple: VARIANT: " <> message)) + +-- | Convert a checked pure result to an IO result. +checked :: Either String a -> IO a +checked = either codecError pure + +-- | Reject native versions outside the supported private-format range. +checkVersion :: IO () +checkVersion = do + version <- c_duckdb_library_version >>= peekCString + let parts = Text.splitOn "." (Text.pack version) + patch = case parts of + ["v1", "5", p] -> readMaybe (Text.unpack p) :: Maybe Int + _ -> Nothing + unless (maybe False (>= 3) patch) (codecError ("unsupported native version " <> version)) + +-- | Check the byte order and pointer width required by native scalar loads. +checkPlatform :: IO () +checkPlatform = alloca \ptr -> do + poke ptr (1 :: Word32) + first <- peek (castPtr ptr :: Ptr Word8) + unless (first == 1 && sizeOf (nullPtr :: Ptr ()) == 8) $ + codecError "the private codec requires a 64-bit little-endian host" + +-- | Acquire a non-NULL native handle. +nonNull :: String -> IO (Ptr a) -> IO (Ptr a) +nonNull label action = do + ptr <- action + when (ptr == nullPtr) (codecError (label <> " returned NULL")) + pure ptr + +-- | Check a logical type's tag. +expectType :: DuckDBType -> DuckDBLogicalType -> IO () +expectType expected logical = do + actual <- c_duckdb_get_type_id logical + unless (actual == expected) (codecError "unexpected physical child type") + +-- | Check a STRUCT's names and owned child descriptors. +checkStruct :: DuckDBLogicalType -> [(Text, DuckDBLogicalType -> IO ())] -> IO () +checkStruct logical fields = do + count <- c_duckdb_struct_type_child_count logical + unless (count == fromIntegral (length fields)) (codecError "unexpected physical child count") + sequence_ + [ do + bracket (nonNull "child name" (c_duckdb_struct_type_child_name logical index)) (c_duckdb_free . castPtr) \name -> do + bytes <- BS.packCString name + unless (bytes == Text.encodeUtf8 expectedName) (codecError "unexpected physical child name") + bracket (nonNull "child type" (c_duckdb_struct_type_child_type logical index)) destroyLogicalType checkChild + | (index, (expectedName, checkChild)) <- zip [0 ..] fields + ] + +-- | Check a LIST and its owned element descriptor. +checkList :: (DuckDBLogicalType -> IO ()) -> DuckDBLogicalType -> IO () +checkList checkChild logical = do + expectType DuckDBTypeList logical + bracket (nonNull "list child type" (c_duckdb_list_type_child_type logical)) destroyLogicalType checkChild + +-- | Check the complete, unshredded four-child VARIANT schema. +checkSchema :: DuckDBLogicalType -> IO () +checkSchema logical = do + expectType DuckDBTypeVariant logical + checkStruct + logical + [ ("keys", checkList (expectType DuckDBTypeVarchar)) + , + ( "children" + , checkList \child -> do + expectType DuckDBTypeStruct child + checkStruct child [("keys_index", expectType DuckDBTypeUInteger), ("values_index", expectType DuckDBTypeUInteger)] + ) + , + ( "values" + , checkList \child -> do + expectType DuckDBTypeStruct child + checkStruct child [("type_id", expectType DuckDBTypeUTinyInt), ("byte_offset", expectType DuckDBTypeUInteger)] + ) + , ("data", expectType DuckDBTypeBlob) + ] + +-- | Check an element index before native pointer arithmetic. +checkIndex :: Int -> Int -> IO () +checkIndex width index = + when (index < 0 || index > maxBound `div` width) (codecError "native element index exceeds Int range") + +-- | Prepare validity access for a chunk. The caller checks row bounds. +prepareValidity :: DuckDBVector -> IO (Int -> IO Bool) +prepareValidity vector = do + validity <- c_duckdb_vector_get_validity vector + pure (chunkIsRowValid validity . fromIntegral) + +-- | Borrow a fixed-width buffer until its chunk is destroyed. +prepareElementReader :: (Storable a) => DuckDBVector -> IO (Int -> IO a) +prepareElementReader vector = do + ptr <- c_duckdb_vector_get_data vector + valid <- prepareValidity vector + pure (readAt valid (castPtr ptr)) + where + readAt :: (Storable a) => (Int -> IO Bool) -> Ptr a -> Int -> IO a + readAt valid ptr index = do + checkIndex (sizeOfElement ptr) index + when (ptr == nullPtr) (codecError "NULL vector data") + present <- valid index + unless present (codecError "NULL physical payload element") + peekElemOff ptr index + sizeOfElement :: (Storable a) => Ptr a -> Int + sizeOfElement ptr = sizeOf (undefined `asTypeOf` element ptr) + element :: Ptr a -> a + element _ = undefined + +-- | Borrow a string buffer and copy each requested value into Haskell memory. +prepareBytesReader :: DuckDBVector -> IO (Int -> IO ByteString) +prepareBytesReader vector = do + base <- c_duckdb_vector_get_data vector + valid <- prepareValidity vector + pure \index -> do + checkIndex 16 index + present <- valid index + unless present (codecError "NULL physical string element") + when (base == nullPtr) (codecError "NULL string vector data") + chunkDecodeBlob base (fromIntegral index) + +-- | Convert a nonnegative bounded integer to Int. +checkedInt :: Integer -> IO Int +checkedInt n + | n < 0 || n > toInteger (maxBound :: Int) = codecError "payload size exceeds Int range" + | otherwise = pure (fromInteger n) + +-- | Prepare LIST bounds checks against the chunk's fixed child size. +prepareListBounds :: DuckDBVector -> IO (Int -> IO (Int, Int)) +prepareListBounds vector = do + readEntry <- prepareElementReader vector + size <- c_duckdb_list_vector_get_size vector + pure \row -> do + DuckDBListEntry offset count <- readEntry row + unless (offset <= size && count <= size - offset) (codecError "LIST bounds exceed child size") + start <- checkedInt (toInteger offset) + len <- checkedInt (toInteger count) + _ <- checkedInt (toInteger offset + toInteger count) + pure (start, len) + +-- | Decode one row while its flattened result chunk remains alive. +decodeVariant :: DuckDBVector -> Int -> IO FieldValue +decodeVariant vector row = prepareVariantDecoder vector >>= ($ row) + +{- | Check the format once and borrow buffers for a flattened result chunk. +The returned reader must not outlive the chunk. The caller supplies row indices +within that chunk. DuckDB result Fetch flattens nested vectors in 1.5. +Each read copies its payload into Haskell memory, including referenced keys. +-} +prepareVariantDecoder :: DuckDBVector -> IO (Int -> IO FieldValue) +prepareVariantDecoder vector = do + checkVersion + checkPlatform + when (vector == nullPtr) (codecError "NULL vector") + bracket (nonNull "vector type" (c_duckdb_vector_get_column_type vector)) destroyLogicalType checkSchema + valid <- prepareValidity vector + keys <- child vector 0 + children <- child vector 1 + values <- child vector 2 + blob <- child vector 3 + keyBounds <- prepareListBounds keys + childBounds <- prepareListBounds children + valueBounds <- prepareListBounds values + keyVector <- nonNull "keys vector" (c_duckdb_list_vector_get_child keys) + childVector <- nonNull "children vector" (c_duckdb_list_vector_get_child children) + valueVector <- nonNull "values vector" (c_duckdb_list_vector_get_child values) + keyIndices <- child childVector 0 + valueIndices <- child childVector 1 + tags <- child valueVector 0 + offsets <- child valueVector 1 + readTag <- prepareElementReader tags + readOffset <- prepareElementReader offsets + hasKey <- prepareValidity keyIndices + readKey <- prepareElementReader keyIndices + readValue <- prepareElementReader valueIndices + readKeyBytes <- prepareBytesReader keyVector + readBlob <- prepareBytesReader blob + pure \row -> do + checkIndex 16 row + present <- valid row + if not present + then pure FieldNull + else do + (keyStart, keyCount) <- keyBounds row + (childStart, childCount) <- childBounds row + (valueStart, valueCount) <- valueBounds row + valueRows <- forM [valueStart .. valueStart + valueCount - 1] \index -> + (,) <$> readTag index <*> readOffset index + childRows <- forM [childStart .. childStart + childCount - 1] \index -> do + keyed <- hasKey index + key <- if keyed then Just <$> readKey index else pure Nothing + value <- readValue index + when (toInteger value >= toInteger valueCount) (codecError "child value index out of bounds") + case key of + Just k | toInteger k >= toInteger keyCount -> codecError "child key index out of bounds" + _ -> pure () + pure (key, value) + let usedKeys = Set.toList (Set.fromList [k | (Just k, _) <- childRows]) + keyRows <- forM usedKeys \index -> do + bytes <- readKeyBytes (keyStart + fromIntegral index) + text <- checked (either (Left . show) Right (Text.decodeUtf8' bytes)) + pure (index, text) + bytes <- readBlob row + decodeVariantPayload valueRows childRows keyRows bytes + where + child parent index = nonNull "STRUCT vector child" (c_duckdb_struct_vector_get_child parent index) + +{- | Decode copied 1.5 payload data for one non-NULL row. +Values are (tag, byte offset). Children are (optional key index, value index). +Keys pair an index with its text. Indices are relative to this row's LISTs. +The root value has index zero. This helper checks all metadata before decoding. +It raises an error for cycles. Shared values use a memo table. +-} +decodeVariantPayload :: [(Word8, Word32)] -> [(Maybe Word32, Word32)] -> [(Word32, Text)] -> ByteString -> IO FieldValue +decodeVariantPayload valueRows childRows keyRows bytes = do + when (finiteBitSize (0 :: Int) < 64) (codecError "the private codec requires a 64-bit host") + let values = listArray (0, length valueRows - 1) valueRows + children = listArray (0, length childRows - 1) childRows + keys = IntMap.fromList [(fromIntegral k, t) | (k, t) <- keyRows] + valueCount = length valueRows + childCount = length childRows + when (null valueRows) (codecError "missing root value") + unless (IntMap.size keys == length keyRows) (codecError "duplicate key dictionary index") + forM_ valueRows \(tag, offset) -> do + when (tag > 33) (codecError "unknown payload tag") + when (toInteger offset > toInteger (BS.length bytes)) (codecError "byte offset exceeds data size") + forM_ childRows \(key, value) -> do + when (toInteger value >= toInteger valueCount) (codecError "child value index out of bounds") + case key of + Just k -> unless (IntMap.member (fromIntegral k) keys) (codecError "missing child key") + Nothing -> pure () + evalStateT (visit values children keys childCount IntSet.empty 0) IntMap.empty + where + visit :: Array Int (Word8, Word32) -> Array Int (Maybe Word32, Word32) -> IntMap.IntMap Text -> Int -> IntSet.IntSet -> Int -> StateT (IntMap.IntMap FieldValue) IO FieldValue + visit values children keys childCount ancestors index = do + when (IntSet.member index ancestors) (lift (codecError "cyclic child reference")) + cached <- gets (IntMap.lookup index) + case cached of + Just result -> pure result + Nothing -> do + (tag, offset) <- lift (checked (arrayElement values index)) + let payload = BS.drop (fromIntegral offset) bytes + result <- case tag of + 29 -> nested True payload + 30 -> nested False payload + _ -> lift (decodeScalar tag payload) + modify' (IntMap.insert index result) + pure result + where + nested object payload = do + (count, rest) <- lift (checked (readVarint payload)) + start <- if count == 0 then pure 0 else fst <$> lift (checked (readVarint rest)) + unless (toInteger start + toInteger count <= toInteger childCount) $ + lift (codecError "container child range out of bounds") + entries <- forM [fromIntegral start .. fromIntegral start + fromIntegral count - 1] \childIndex -> do + (key, childIndexValue) <- lift (checked (arrayElement children childIndex)) + name <- case (object, key) of + (True, Just k) -> case IntMap.lookup (fromIntegral k) keys of + Just text -> pure text + Nothing -> lift (codecError "missing object key") + (False, Nothing) -> pure Text.empty + _ -> lift (codecError "container key validity does not match its tag") + item <- visit values children keys childCount (IntSet.insert index ancestors) (fromIntegral childIndexValue) + pure (name, item) + if object + then do + unless (Set.size (Set.fromList (map fst entries)) == length entries) $ + lift (codecError "duplicate object key") + pure (variantObject entries) + else pure (FieldList (map snd entries)) + +-- | Read an array element after checking both bounds. +arrayElement :: Array Int a -> Int -> Either String a +arrayElement array index + | index < lower || index > upper = Left "value index out of bounds" + | otherwise = Right (array ! index) + where + (lower, upper) = bounds array + +-- | Read an unsigned base-128 uint32 without overflow or truncation. +readVarint :: ByteString -> Either String (Word32, ByteString) +readVarint = go 0 0 + where + go shift value input = case BS.uncons input of + Nothing -> Left "truncated varint" + Just (byte, rest) + | shift == 28 && byte > 15 -> Left "varint exceeds uint32" + | otherwise -> + let result = value .|. (fromIntegral (byte .&. 127) `shiftL` shift) + in if byte .&. 128 == 0 then Right (result, rest) else go (shift + 7) result rest + +-- | Read a bounded fixed-size payload in little-endian order. +readUnsigned :: Int -> ByteString -> Either String (Integer, ByteString) +readUnsigned count bytes + | BS.length bytes < count = Left "truncated scalar payload" + | otherwise = + let (part, rest) = BS.splitAt count bytes + in Right (BS.foldr (\byte value -> value `shiftL` 8 .|. toInteger byte) 0 part, rest) + +-- | Read a two's-complement integer with a checked payload length. +readSigned :: Int -> ByteString -> Either String Integer +readSigned count bytes = do + (value, _) <- readUnsigned count bytes + pure (if value >= 2 ^ (count * 8 - 1) then value - 2 ^ (count * 8) else value) + +-- | Read a length-prefixed string or blob. +readString :: ByteString -> Either String ByteString +readString bytes = do + (count, rest) <- readVarint bytes + if toInteger count > toInteger (BS.length rest) + then Left "string length exceeds payload size" + else Right (BS.take (fromIntegral count) rest) + +-- | Decode the scalar tags from VariantLogicalType in DuckDB 1.5. +decodeScalar :: Word8 -> ByteString -> IO FieldValue +decodeScalar tag bytes = case tag of + 0 -> pure FieldNull + 1 -> pure (FieldBool True) + 2 -> pure (FieldBool False) + 15 -> checked do + (precision, rest) <- readVarint bytes + (scale, digits) <- readVarint rest + when (precision < 1 || precision > 38 || scale > precision) (Left "invalid decimal metadata") + let count = if precision <= 4 then 2 else if precision <= 9 then 4 else if precision <= 18 then 8 else 16 + value <- readSigned count digits + when (abs value >= 10 ^ precision) (Left "unscaled decimal exceeds its precision") + pure (FieldDecimal (DecimalValue (fromIntegral precision) (fromIntegral scale) value)) + 16 -> checked do + string <- readString bytes + FieldText <$> either (Left . show) Right (Text.decodeUtf8' string) + 17 -> FieldBlob <$> checked (readString bytes) + 31 -> FieldBigNum . BigNum <$> checked (readString bytes >>= decodeBigNum) + 32 -> checked do + bitBytes <- readString bytes + case BS.unpack (BS.take 2 bitBytes) of + [padding, first] -> do + when (padding > 7) (Left "BIT padding exceeds seven") + let maskBits = paddingMask padding + unless (first .&. maskBits == maskBits) (Left "invalid native BIT padding") + pure (FieldBit (bitStringFromBytes bitBytes)) + _ -> Left "BIT requires a padding byte and nonempty data" + 33 -> FieldGeometry . (`RawGeometry` Nothing) <$> checked (readString bytes) + _ -> case fixedWidthTag tag of + Just (dtype, size) -> decodeFixedWidth dtype size bytes + Nothing -> codecError "unknown scalar tag" + +{- | The type and the payload size of each fixed-width scalar tag. These +payloads have the memory layout of one vector element of the type. +-} +fixedWidthTag :: Word8 -> Maybe (DuckDBType, Int) +fixedWidthTag = \case + 3 -> Just (DuckDBTypeTinyInt, 1) + 4 -> Just (DuckDBTypeSmallInt, 2) + 5 -> Just (DuckDBTypeInteger, 4) + 6 -> Just (DuckDBTypeBigInt, 8) + 7 -> Just (DuckDBTypeHugeInt, 16) + 8 -> Just (DuckDBTypeUTinyInt, 1) + 9 -> Just (DuckDBTypeUSmallInt, 2) + 10 -> Just (DuckDBTypeUInteger, 4) + 11 -> Just (DuckDBTypeUBigInt, 8) + 12 -> Just (DuckDBTypeUHugeInt, 16) + 13 -> Just (DuckDBTypeFloat, 4) + 14 -> Just (DuckDBTypeDouble, 8) + 18 -> Just (DuckDBTypeUUID, 16) + 19 -> Just (DuckDBTypeDate, 4) + 20 -> Just (DuckDBTypeTime, 8) + 21 -> Just (DuckDBTypeTimeNs, 8) + 22 -> Just (DuckDBTypeTimestampS, 8) + 23 -> Just (DuckDBTypeTimestampMs, 8) + 24 -> Just (DuckDBTypeTimestamp, 8) + 25 -> Just (DuckDBTypeTimestampNs, 8) + 26 -> Just (DuckDBTypeTimeTz, 8) + 27 -> Just (DuckDBTypeTimestampTz, 8) + 28 -> Just (DuckDBTypeInterval, 16) + _ -> Nothing + +-- | Copy a fixed-width payload to aligned memory and decode it as one element. +decodeFixedWidth :: DuckDBType -> Int -> ByteString -> IO FieldValue +decodeFixedWidth dtype size bytes = do + when (BS.length bytes < size) (codecError "truncated scalar payload") + allocaBytesAligned size 16 \buffer -> do + BS.useAsCStringLen bytes \(source, _) -> copyBytes buffer (castPtr source) size + decodeElement dtype (castPtr buffer) 0 + +-- | Get the high-bit mask used for native BIT padding. +paddingMask :: Word8 -> Word8 +paddingMask padding = complement ((1 `shiftL` (8 - fromIntegral padding)) - 1) + +-- | Check BIGNUM's sign, three-byte length header, and magnitude, then decode it. +decodeBigNum :: ByteString -> Either String Integer +decodeBigNum bytes = do + when (BS.length bytes < 4) (Left "truncated BIGNUM header or magnitude") + let (header, raw) = BS.splitAt 3 bytes + encoded = BS.foldl' (\n byte -> n `shiftL` 8 .|. fromIntegral byte) (0 :: Word32) header + negative = encoded .&. 0x800000 == 0 + decoded = if negative then complement encoded .&. 0xffffff else encoded + count = decoded .&. 0x7fffff + magnitude = if negative then BS.map complement raw else raw + unless (toInteger count == toInteger (BS.length magnitude)) (Left "BIGNUM header length mismatch") + case BS.uncons magnitude of + Just (0, rest) | not (BS.null rest) || negative -> Left "noncanonical BIGNUM magnitude" + _ -> Right () + pure (fromBigNumBytes (BS.unpack bytes)) diff --git a/test-codec/Main.hs b/test-codec/Main.hs new file mode 100644 --- /dev/null +++ b/test-codec/Main.hs @@ -0,0 +1,88 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} + +-- | Check the private codec's validation without constructing invalid native vectors. +module Main (main) where + +import Control.Exception (IOException, try) +import Control.Monad (forM_) +import Data.Bits (shiftR, (.&.), (.|.)) +import qualified Data.ByteString as BS +import Data.Text (Text) +import Data.Word (Word32, Word8) +import Database.DuckDB.Simple.FromField (FieldValue (..)) +import Database.DuckDB.Simple.Variant (variantObject) +import Database.DuckDB.Simple.VariantCodec (decodeVariantPayload) +import Test.Tasty (TestTree, defaultMain, testGroup) +import Test.Tasty.HUnit + +-- | Run malformed payload checks and boundary cases for the native format. +main :: IO () +main = + defaultMain $ + testGroup + "VARIANT payload validation" + [ rejected "missing root" [] [] [] [] + , rejected "unknown tag" [(34, 0)] [] [] [] + , rejected "offset outside blob" [(1, 1)] [] [] [] + , rejected "truncated BIGINT" [(6, 0)] [] [] (replicate 7 0) + , rejected "truncated varint" [(16, 0)] [] [] [128] + , rejected "varint overflows uint32" [(16, 0)] [] [] [255, 255, 255, 255, 16] + , rejected "unterminated varint" [(16, 0)] [] [] [255, 255, 255, 255, 255, 0] + , rejected "string outside blob" [(16, 0)] [] [] [3, 97, 98] + , rejected "invalid UTF-8" [(16, 0)] [] [] [1, 255] + , rejected "decimal scale exceeds precision" [(15, 0)] [] [] [2, 3, 1, 0] + , rejected "decimal value exceeds precision" [(15, 0)] [] [] [1, 0, 100, 0] + , rejected "BIGNUM header is truncated" [(31, 0)] [] [] [2, 128, 0] + , rejected "BIGNUM length mismatch" [(31, 0)] [] [] [4, 128, 0, 2, 1] + , rejected "BIT padding exceeds seven" [(32, 0)] [] [] [2, 8, 255] + , rejected "BIT padding bits are absent" [(32, 0)] [] [] [2, 3, 21] + , rejected "child value out of range" [(30, 0)] [(Nothing, 1)] [] [1, 0] + , rejected "container range out of bounds" [(30, 0)] [] [] [1, 0] + , rejected "object key is missing" [(29, 0), (1, 2)] [(Just 0, 1)] [] [1, 0] + , rejected "array child has a key" [(30, 0), (1, 2)] [(Just 0, 1)] [(0, "x")] [1, 0] + , rejected "object child has no key" [(29, 0), (1, 2)] [(Nothing, 1)] [] [1, 0] + , rejected "duplicate dictionary index" [(1, 0)] [] [(0, "x"), (0, "y")] [] + , rejected "duplicate object key" [(29, 0), (1, 2)] [(Just 0, 1), (Just 1, 1)] [(0, "x"), (1, "x")] [2, 0] + , rejected "cyclic child reference" [(30, 0)] [(Nothing, 0)] [] [1, 0] + , testCase "empty values need no payload bytes" $ + decodeVariantPayload [(30, 0)] [] [] (BS.singleton 0) >>= (@?= FieldList []) + , testCase "shared children decode without duplicate traversal" $ + decodeVariantPayload [(30, 0), (3, 2)] [(Nothing, 1), (Nothing, 1)] [] (BS.pack [2, 0, 7]) + >>= (@?= FieldList [FieldInt8 7, FieldInt8 7]) + , testCase "length-aware keys preserve embedded NUL" $ + decodeVariantPayload [(29, 0), (1, 2)] [(Just 0, 1)] [(0, "a\0b")] (BS.pack [1, 0]) + >>= (@?= variantObject [("a\0b", FieldBool True)]) + , testCase "deep acyclic containers are accepted" $ + forM_ [128, 129, 257, 1025] $ \count -> + let (values, children, bytes) = nested count + in decodeVariantPayload values children [] bytes >>= (@?= foldr (const (FieldList . pure)) (FieldInt8 7) [1 .. count - 1]) + ] + +-- | Require a validation error for copied native payload data. +rejected :: String -> [(Word8, Word32)] -> [(Maybe Word32, Word32)] -> [(Word32, Text)] -> [Word8] -> TestTree +rejected label values children keys bytes = + testCase label $ assertRejected (decodeVariantPayload values children keys (BS.pack bytes)) + +-- | Require a codec error. +assertRejected :: IO FieldValue -> Assertion +assertRejected action = do + result <- try action + case result of + Left (_ :: IOException) -> pure () + Right value -> assertFailure ("expected payload rejection, got " <> show value) + +-- | Construct a chain of arrays with one scalar leaf. +nested :: Int -> ([(Word8, Word32)], [(Maybe Word32, Word32)], BS.ByteString) +nested count = + ( zip (replicate (count - 1) 30 <> [3]) offsets + , [(Nothing, fromIntegral (n + 1)) | n <- [0 .. count - 2]] + , BS.pack (concat payloads) + ) + where + payloads = [1 : varint (fromIntegral n) | n <- [0 .. count - 2]] <> [[7]] + offsets = scanl (+) 0 (map (fromIntegral . length) payloads) + varint :: Word32 -> [Word8] + varint value + | value < 128 = [fromIntegral value] + | otherwise = fromIntegral (value .&. 127 .|. 128) : varint (value `shiftR` 7) diff --git a/test/GeometryRegressionTests.hs b/test/GeometryRegressionTests.hs new file mode 100644 --- /dev/null +++ b/test/GeometryRegressionTests.hs @@ -0,0 +1,404 @@ +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE OverloadedStrings #-} +{-# OPTIONS_GHC -Wno-deprecations #-} + +module GeometryRegressionTests (tests) where + +import Control.Exception (SomeException, bracket, displayException, try) +import Control.Monad (forM_) +import Data.Array (Array, listArray) +import qualified Data.ByteString as BS +import qualified Data.Geometry as G +import qualified Data.Geometry.WKB as WKB +import qualified Data.Geometry.WKT as WKT +import Data.Int (Int64) +import Data.List (isInfixOf) +import qualified Data.Map.Strict as Map +import Data.Text (Text) +import qualified Data.Text as Text +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as U +import Database.DuckDB.Simple +import qualified Database.DuckDB.Simple.Deprecated.Streaming as Streaming +import Database.DuckDB.Simple.FromField (FieldValue (..), StructValue (..), UnionValue) +import Database.DuckDB.Simple.Geometry (RawGeometry (..), fromRawGeometry, toRawGeometry) +import Database.DuckDB.Simple.LogicalRep (LogicalTypeRep (..), StructField (..), destroyLogicalType, logicalTypeFromRep, logicalTypeToRep) +import GHC.Float (castWord64ToDouble) +import System.Mem (performMajorGC) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit +import TestUtils (assertFailureIO) + +-- | Check native construction, metadata, nesting, and streaming. +tests :: TestTree +tests = + testGroup + "geometry integration" + [ testCase "Haskell shapes bind directly in every coordinate layout" $ + withConnection ":memory:" \conn -> do + shapeRoundTrips conn G.DimXY G.PointXY G.CoordinatesXY (G.XY 1 2) (G.XY 3 4) + shapeRoundTrips conn G.DimXYZ G.PointXYZ G.CoordinatesXYZ (G.XYZ 1 2 3) (G.XYZ 4 5 6) + shapeRoundTrips conn G.DimXYM G.PointXYM G.CoordinatesXYM (G.XYM 1 2 3) (G.XYM 4 5 6) + shapeRoundTrips conn G.DimXYZM G.PointXYZM G.CoordinatesXYZM (G.XYZM 1 2 3 4) (G.XYZM 5 6 7 8) + , testCase "large coordinate buffers survive native round trips" $ + withConnection ":memory:" \conn -> do + let coordinates = U.generate 10000 (\i -> G.XY (fromIntegral i) (fromIntegral (i * 2))) + points = U.imap (\i coordinate -> if i `mod` 17 == 0 then G.EmptyPoint G.DimXY else G.PointXY coordinate) coordinates + forM_ [G.LineString (G.CoordinatesXY coordinates), G.MultiPoint points] \shape -> do + (query conn "SELECT ?" (Only shape)) >>= (@?= [Only shape]) + raw <- either assertFailure pure (toRawGeometry shape) + (importRaw conn raw) >>= (@?= [Only raw]) + , testCase "WKT binding preserves exact finite coordinate bits" $ + withConnection ":memory:" \conn -> + forM_ [0, 0x8000000000000000, 1, 0x8000000000000001, 0x000fffffffffffff, 0x0010000000000000, 0x3fb999999999999b, 0x3ff0000000000001, 0x44b52d02c7e14af6, 0x7fefffffffffffff, 0xffefffffffffffff] \bits -> do + let shape = G.PointGeometry (G.PointXY (G.XY (castWord64ToDouble bits) (-0.0))) + expected <- either assertFailure pure (WKB.encodeWKB shape) + [(actual, wkt)] <- query conn "SELECT ST_AsWKB(?), ST_AsText(?)" (shape, shape) :: IO [(BS.ByteString, Text)] + actual @?= expected + (WKT.decodeWKT wkt >>= WKB.encodeWKB) @?= Right expected + (query conn "SELECT system.main.ST_AsWKB(system.main.ST_GeomFromWKB(?))" (Only expected) :: IO [Only BS.ByteString]) + >>= (@?= [Only expected]) + , testCase "pure WKT decoding agrees with native geometry results" $ + withConnection ":memory:" \conn -> + forM_ shapes \wkt -> do + [(native, rendered)] <- query conn "SELECT ?::GEOMETRY, ST_AsText(?::GEOMETRY)" (wkt, wkt) :: IO [(G.Geometry, Text)] + WKT.decodeWKT wkt @?= Right native + WKT.decodeWKT rendered @?= Right native + , testCase "native bare multipoint EMPTY syntax parses as WKT" $ + withConnection ":memory:" \conn -> do + [(native, bytes, rendered)] <- + query_ conn "SELECT g, ST_AsWKB(g), ST_AsText(g) FROM (SELECT 'MULTIPOINT ((1 2), EMPTY, (3 4))'::GEOMETRY AS g)" :: IO [(G.Geometry, BS.ByteString, Text)] + native @?= G.MultiPoint (U.fromList [G.PointXY (G.XY 1 2), G.EmptyPoint G.DimXY, G.PointXY (G.XY 3 4)]) + WKB.decodeWKB bytes @?= Right native + WKT.decodeWKT rendered @?= Right native + (WKT.encodeWKT native >>= WKT.decodeWKT) @?= Right native + , testCase "standalone WKT syntax binds through native geometry parameters" $ + withConnection ":memory:" \conn -> + forM_ ["point(+.5 -1.e+2)", "MULTIPOINT (EMPTY, (1 2), (3 4))", "MULTIPOINT (EMPTY, 1 2, EMPTY, 3 4, EMPTY)"] \wkt -> do + shape <- either assertFailure pure (WKT.decodeWKT wkt) + (query conn "SELECT ?" (Only shape)) >>= (@?= [Only shape]) + , testCase "native geometry depth limit produces a recoverable error" $ + withConnection ":memory:" \conn -> do + let nested = iterate (G.GeometryCollection . V.singleton) (G.PointGeometry (G.PointXY (G.XY 1 2))) + accepted = nested !! 15 + rejected = nested !! 16 + (query conn "SELECT ?" (Only accepted)) >>= (@?= [Only accepted]) + assertFailureIO (query conn "SELECT ?" (Only rejected) :: IO [Only G.Geometry]) + (query_ conn "SELECT 42" :: IO [Only Int64]) >>= (@?= [Only 42]) + , testCase "decoded native shapes round trip through parameters and WKB" $ + withConnection ":memory:" \conn -> + forM_ shapes \wkt -> do + [Only geometry] <- query conn "SELECT ?::GEOMETRY" (Only wkt) :: IO [Only G.Geometry] + (query conn "SELECT ?" (Only geometry) :: IO [Only G.Geometry]) >>= (@?= [Only geometry]) + raw <- either assertFailure pure (toRawGeometry geometry) + fromRawGeometry raw @?= Right geometry + , testCase "decoded rows retain point and sequence layouts, including empties" $ + withConnection ":memory:" \conn -> + forM_ + [ ("POINT Z EMPTY", G.PointGeometry (G.EmptyPoint G.DimXYZ)) + , ("LINESTRING M EMPTY", G.LineString (G.CoordinatesXYM U.empty)) + , ("POINT Z (1 2 3)", G.PointGeometry (G.PointXYZ (G.XYZ 1 2 3))) + , ("POINT EMPTY", G.PointGeometry (G.EmptyPoint G.DimXY)) + ] + \(wkt, expected) -> + (query conn "SELECT ?::GEOMETRY" (Only (wkt :: Text)) :: IO [Only G.Geometry]) + >>= (@?= [Only expected]) + , testCase "decoded arrays preserve shapes and NULL elements without CRS" $ + withConnection ":memory:" \conn -> do + let geometry = G.PointGeometry (G.PointXYM (G.XYM 1 2 3)) + values = listArray (0, 2) [Nothing, Just geometry, Just (G.PointGeometry (G.EmptyPoint G.DimXYM))] + (query conn "SELECT ?, ST_CRS((?)[2])" (values, values) :: IO [(Array Int (Maybe G.Geometry), Maybe Text)]) + >>= (@?= [(values, Nothing)]) + , testCase "only raw geometry retains CRS metadata" $ + withConnection ":memory:" \conn -> do + [(raw, shape)] <- query_ conn "SELECT g, g FROM (SELECT 'POINT ZM (1 2 3 4)'::GEOMETRY('OGC:CRS84') AS g)" :: IO [(RawGeometry, G.Geometry)] + rawGeometryCRS raw @?= Just "OGC:CRS84" + shape @?= G.PointGeometry (G.PointXYZM (G.XYZM 1 2 3 4)) + fromRawGeometry raw @?= Right shape + unannotated <- either assertFailure pure (toRawGeometry shape) + rawGeometryCRS unannotated @?= Nothing + rawGeometryWKB unannotated @?= rawGeometryWKB raw + (query conn "SELECT system.main.ST_CRS(system.main.ST_SetCRS(system.main.ST_GeomFromWKB(?), ?))" (rawParameters raw) :: IO [Only (Maybe Text)]) + >>= (@?= [Only (Just "OGC:CRS84")]) + (query conn "SELECT ST_CRS(?), ?" (shape, shape) :: IO [(Maybe Text, G.Geometry)]) + >>= (@?= [(Nothing, shape)]) + , testCase "decoded coordinates remain usable after closing their connection" $ do + geometry <- withConnection ":memory:" \conn -> do + [Only value] <- query_ conn "SELECT 'LINESTRING (1 2, 3 4, 5 6)'::GEOMETRY('OGC:CRS84')" :: IO [Only G.Geometry] + pure value + performMajorGC + case geometry of + G.LineString (G.CoordinatesXY coords) -> U.foldl' (\acc (G.XY x y) -> acc + x + y) 0 coords @?= 21 + other -> assertFailure ("unexpected geometry: " <> show other) + , testCase "native WKB round trips through explicit binary import" $ + withConnection ":memory:" \conn -> + forM_ shapes \wkt -> do + [Only geometry] <- query conn "SELECT ?::GEOMETRY" (Only wkt) + (query conn "SELECT typeof(system.main.ST_GeomFromWKB(?)), system.main.ST_GeomFromWKB(?)" (rawGeometryWKB geometry, rawGeometryWKB geometry) :: IO [(Text, RawGeometry)]) + >>= (@?= [("GEOMETRY", geometry)]) + , testCase "raw empty containers retain layouts that decoded values do not store" $ + withConnection ":memory:" \conn -> + forM_ + ( [ (family <> " " <> dimensions <> " EMPTY", family <> " EMPTY") + | family <- ["MULTIPOINT", "MULTILINESTRING", "MULTIPOLYGON", "GEOMETRYCOLLECTION"] + , dimensions <- ["Z", "M", "ZM"] + ] + ++ [("GEOMETRYCOLLECTION ZM (MULTIPOINT ZM EMPTY, GEOMETRYCOLLECTION ZM EMPTY)", "GEOMETRYCOLLECTION (MULTIPOINT EMPTY, GEOMETRYCOLLECTION EMPTY)")] + ) + \(wkt, normalizedWKT) -> do + [Only raw] <- query conn "SELECT ?::GEOMETRY('OGC:CRS84')" (Only (wkt :: Text)) :: IO [Only RawGeometry] + (importRaw conn raw) >>= (@?= [Only raw]) + decoded <- either assertFailure pure (fromRawGeometry raw) + WKT.encodeWKT decoded @?= Right normalizedWKT + normalized <- either assertFailure pure (toRawGeometry decoded) + rawGeometryCRS normalized @?= Nothing + assertBool "decoded empty containers must not retain an unstored dimension tag" (rawGeometryWKB normalized /= rawGeometryWKB raw) + , testCase "mixed-layout collections use binary import without WKT normalization" $ + withConnection ":memory:" \conn -> do + let geometry = G.GeometryCollection (V.fromList [G.PointGeometry (G.PointXY (G.XY 1 2)), G.PointGeometry (G.PointXYZ (G.XYZ 3 4 5))]) + raw <- either assertFailure pure (toRawGeometry geometry) + assertFailureIO (query conn "SELECT ?" (Only geometry) :: IO [Only G.Geometry]) + (importRaw conn raw) >>= (@?= [Only raw]) + (query conn "SELECT ST_GeomFromWKB(?)" (Only (rawGeometryWKB raw)) :: IO [Only G.Geometry]) + >>= (@?= [Only geometry]) + (query_ conn "SELECT 42" :: IO [Only Int64]) >>= (@?= [Only 42]) + , testCase "mixed multipoint layouts promote missing ordinates to NaN" $ + withConnection ":memory:" \conn -> do + let geometry = G.MultiPoint (U.fromList [G.PointXY (G.XY 1 2), G.PointXYZ (G.XYZ 3 4 5)]) + [Only actual] <- query conn "SELECT ?" (Only geometry) :: IO [Only G.Geometry] + case actual of + G.MultiPoint points -> do + U.length points @?= 2 + case points U.! 0 of + G.PointXYZ (G.XYZ x y z) -> do + (x, y) @?= (1, 2) + assertBool "promoted Z must be NaN" (isNaN z) + other -> assertFailure ("unexpected promoted point: " <> show other) + points U.! 1 @?= G.PointXYZ (G.XYZ 3 4 5) + other -> assertFailure ("unexpected geometry: " <> show other) + , testCase "raw NaN points retain ordinates that decoded emptiness normalizes" $ + withConnection ":memory:" \conn -> do + [Only raw] <- query_ conn "SELECT 'POINT Z (NaN NaN 7)'::GEOMETRY" :: IO [Only RawGeometry] + (importRaw conn raw) >>= (@?= [Only raw]) + decoded <- either assertFailure pure (fromRawGeometry raw) + decoded @?= G.PointGeometry (G.EmptyPoint G.DimXYZ) + normalized <- either assertFailure pure (toRawGeometry decoded) + assertBool "decoding an XY-NaN point normalizes its extra ordinates" (rawGeometryWKB normalized /= rawGeometryWKB raw) + (query conn "SELECT ?" (Only decoded) :: IO [Only G.Geometry]) >>= (@?= [Only decoded]) + , testCase "binary import preserves distinct NaN payload bits" $ + withConnection ":memory:" \conn -> do + let shape = G.PointGeometry (G.PointXYZ (G.XYZ (castWord64ToDouble 0x7ff8000000000001) (castWord64ToDouble 0x7ff8000000000002) 7)) + raw <- either assertFailure pure (toRawGeometry shape) + (importRaw conn raw) >>= (@?= [Only raw]) + fromRawGeometry raw @?= Right (G.PointGeometry (G.EmptyPoint G.DimXYZ)) + , testCase "binary import preserves mixed empty and nonempty member layouts" $ + withConnection ":memory:" \conn -> do + let shape = G.GeometryCollection (V.fromList [G.PointGeometry (G.EmptyPoint G.DimXY), G.PointGeometry (G.PointXYZ (G.XYZ 1 2 3))]) + raw <- either assertFailure pure (toRawGeometry shape) + (query conn "SELECT system.main.ST_GeomFromWKB(?)" (Only (rawGeometryWKB raw)) :: IO [Only RawGeometry]) >>= (@?= [Only raw]) + (query conn "SELECT system.main.ST_GeomFromWKB(?)" (Only (rawGeometryWKB raw)) :: IO [Only G.Geometry]) >>= (@?= [Only shape]) + , testCase "decoded polygon writing normalizes all-empty rings" $ + withConnection ":memory:" \conn -> do + let emptyRing = G.CoordinatesXYZ U.empty + geometry = G.Polygon (G.PolygonRings emptyRing (V.singleton emptyRing)) + expected = G.Polygon (G.PolygonRings emptyRing V.empty) + (query conn "SELECT ?" (Only geometry) :: IO [Only G.Geometry]) >>= (@?= [Only expected]) + , testCase "invalid decoded construction fails without poisoning the connection" $ + withConnection ":memory:" \conn -> do + let closedRing = G.CoordinatesXY (U.fromList [G.XY 0 0, G.XY 1 0, G.XY 1 1, G.XY 0 0]) + openRing = G.CoordinatesXY (U.fromList [G.XY 0 0, G.XY 1 0, G.XY 1 1]) + emptyRing = G.CoordinatesXY U.empty + forM_ + [ G.LineString (G.CoordinatesXY (U.singleton (G.XY 1 2))) + , G.Polygon (G.PolygonRings openRing V.empty) + , G.Polygon (G.PolygonRings emptyRing (V.singleton closedRing)) + ] + \shape -> + assertFailureIO (query conn "SELECT ?" (Only shape) :: IO [Only G.Geometry]) + [Only raw] <- query_ conn "SELECT 'LINESTRING (1 2)'::GEOMETRY" :: IO [Only RawGeometry] + assertBool "raw result bytes remain available" (not (BS.null (rawGeometryWKB raw))) + assertFailureIO (query_ conn "SELECT 'LINESTRING (1 2)'::GEOMETRY" :: IO [Only G.Geometry]) + (query_ conn "SELECT 42" :: IO [Only Int64]) >>= (@?= [Only 42]) + , testCase "CRS survives explicit binary import and table insertion" $ + withConnection ":memory:" \conn -> do + [Only geometry] <- query_ conn "SELECT 'POINT ZM (1 2 3 4)'::GEOMETRY('OGC:CRS84')" + rawGeometryCRS geometry @?= Just "OGC:CRS84" + (importRaw conn geometry) + >>= (@?= [Only geometry]) + _ <- execute_ conn "CREATE TABLE shapes (shape GEOMETRY('OGC:CRS84'))" + _ <- execute conn "INSERT INTO shapes VALUES (system.main.ST_SetCRS(system.main.ST_GeomFromWKB(?), ?))" (rawParameters geometry) + (query_ conn "SELECT shape FROM shapes" :: IO [Only RawGeometry]) >>= (@?= [Only geometry]) + , testCase "CRS parameters preserve custom names and full definitions" $ + withConnection ":memory:" \conn -> do + [Only geometry] <- query_ conn "SELECT 'POINT (1 2)'::GEOMETRY" + forM_ + [ "local' íslenska λ" + , "ENGCRS[\"Local grid\",EDATUM[\"Local datum\"],CS[Cartesian,2],AXIS[\"x\",east,ORDER[1],LENGTHUNIT[\"metre\",1]],AXIS[\"y\",north,ORDER[2],LENGTHUNIT[\"metre\",1]]]" + , "{\"type\":\"EngineeringCRS\",\"name\":\"Local grid\",\"datum\":{\"type\":\"EngineeringDatum\",\"name\":\"Local datum\"},\"coordinate_system\":{\"subtype\":\"Cartesian\",\"axis\":[{\"name\":\"x\",\"abbreviation\":\"x\",\"direction\":\"east\",\"unit\":\"metre\"},{\"name\":\"y\",\"abbreviation\":\"y\",\"direction\":\"north\",\"unit\":\"metre\"}]}}" + ] + \crs -> do + let annotated = geometry{rawGeometryCRS = Just crs} + (importRaw conn annotated) >>= (@?= [Only annotated]) + , testCase "geometry construction ignores user macros with built-in names" $ + withConnection ":memory:" \conn -> do + [Only raw] <- query_ conn "SELECT 'POINT (1 2)'::GEOMETRY('OGC:CRS84')" :: IO [Only RawGeometry] + decoded <- either assertFailure pure (fromRawGeometry raw) + _ <- execute_ conn "CREATE MACRO ST_AsText(x) AS 42" + _ <- execute_ conn "CREATE MACRO ST_GeomFromWKB(x) AS 'POINT (9 9)'::GEOMETRY" + _ <- execute_ conn "CREATE MACRO ST_SetCRS(x, crs) AS 42" + (query conn "SELECT system.main.ST_SetCRS(system.main.ST_GeomFromWKB(?), ?), ?" (rawGeometryWKB raw, rawGeometryCRS raw, decoded) :: IO [(RawGeometry, G.Geometry)]) + >>= (@?= [(raw, decoded)]) + , testCase "NULL and empty geometry stay distinct" $ + withConnection ":memory:" \conn -> do + [Only empty] <- query_ conn "SELECT 'POINT EMPTY'::GEOMETRY" :: IO [Only RawGeometry] + (query conn "SELECT system.main.ST_GeomFromWKB(?), system.main.ST_GeomFromWKB(?)" (Nothing :: Maybe BS.ByteString, Just (rawGeometryWKB empty)) :: IO [(Maybe RawGeometry, Maybe RawGeometry)]) + >>= (@?= [(Nothing, Just empty)]) + , testCase "binary import constructs arrays with a common CRS and NULL elements" $ + withConnection ":memory:" \conn -> do + [Only geometry] <- query_ conn "SELECT 'POINT (1 2)'::GEOMETRY('OGC:CRS84')" :: IO [Only RawGeometry] + let values = listArray (0, 2) [Nothing, Just geometry, Nothing] + (query conn "SELECT [NULL, system.main.ST_SetCRS(system.main.ST_GeomFromWKB(?), ?), NULL]::GEOMETRY('OGC:CRS84')[3]" (rawParameters geometry) :: IO [Only (Array Int (Maybe RawGeometry))]) >>= (@?= [Only values]) + , testCase "nested LIST, ARRAY and MAP results preserve CRS and reject implicit raw binding" $ + withConnection ":memory:" \conn -> do + [Only value] <- + query_ + conn + "SELECT {'xs': ['POINT (1 2)'::GEOMETRY('OGC:CRS84'), NULL], 'fixed': ['POINT EMPTY'::GEOMETRY('OGC:CRS84')]::GEOMETRY('OGC:CRS84')[1], 'map': MAP {'one': 'POINT (3 4)'::GEOMETRY('OGC:CRS84')}}" :: + IO [Only (StructValue FieldValue)] + assertRawBindingFailure (query conn "SELECT ?" (Only value) :: IO [Only (StructValue FieldValue)]) + (query_ conn "SELECT 42" :: IO [Only Int64]) >>= (@?= [Only 42]) + , testCase "UNION geometry results retain CRS and reject implicit raw binding" $ + withConnection ":memory:" \conn -> do + [Only value] <- query_ conn "SELECT union_value(shape := 'POINT (1 2)'::GEOMETRY('OGC:CRS84'))" :: IO [Only (UnionValue FieldValue)] + assertRawBindingFailure (query conn "SELECT ?" (Only value) :: IO [Only (UnionValue FieldValue)]) + [Only nullValue] <- query_ conn "SELECT union_value(shape := NULL::GEOMETRY('OGC:CRS84'))" :: IO [Only (UnionValue FieldValue)] + (query conn "SELECT ?" (Only nullValue) :: IO [Only (UnionValue FieldValue)]) >>= (@?= [Only nullValue]) + (query_ conn "SELECT 42" :: IO [Only Int64]) >>= (@?= [Only 42]) + , testCase "inactive UNION geometry members keep the CRS of the type cache" $ + withConnection ":memory:" \conn -> do + [Only value] <- query_ conn "SELECT union_value(number := 42::BIGINT)::UNION(number BIGINT, shape GEOMETRY('OGC:CRS84'))" :: IO [Only (UnionValue FieldValue)] + (query conn "SELECT typeof(?)" (Only value) :: IO [Only Text]) >>= (@?= [Only "UNION(number BIGINT, shape GEOMETRY('OGC:CRS84'))"]) + (query conn "SELECT ?" (Only value) :: IO [Only (UnionValue FieldValue)]) >>= (@?= [Only value]) + , testCase "NULL and empty nested geometry collections keep the CRS of the type cache" $ + withConnection ":memory:" \conn -> do + _ <- + execute_ + conn + "CREATE TABLE nested AS SELECT {'null_list': NULL::GEOMETRY('OGC:CRS84')[], 'empty_list': []::GEOMETRY('OGC:CRS84')[], 'null_array': NULL::GEOMETRY('OGC:CRS84')[2], 'array': [NULL, NULL]::GEOMETRY('OGC:CRS84')[2], 'null_map': NULL::MAP(VARCHAR, GEOMETRY('OGC:CRS84')), 'empty_map': map([], [])::MAP(VARCHAR, GEOMETRY('OGC:CRS84')), 'map': map(['one'], [NULL::GEOMETRY('OGC:CRS84')])} AS v" + [Only value] <- query_ conn "SELECT v FROM nested" :: IO [Only (StructValue FieldValue)] + (query conn "SELECT ?" (Only value) :: IO [Only (StructValue FieldValue)]) >>= (@?= [Only value]) + _ <- execute conn "INSERT INTO nested VALUES (?)" (Only value) + (query_ conn "SELECT v FROM nested" :: IO [Only (StructValue FieldValue)]) >>= (@?= [Only value, Only value]) + , testCase "configured CRS definitions bind; other CRS definitions bind without a CRS" $ do + let crs = "local' íslenska λ" + withConnectionWithOptions ":memory:" defaultConnectionOptions{connectionGeometryCRS = [crs]} \conn -> do + [Only typed] <- query conn "SELECT typeof(?)" (Only (nullShape crs)) :: IO [Only Text] + assertBool (Text.unpack typed) ("íslenska λ" `Text.isInfixOf` typed) + (query conn "SELECT ?" (Only (nullShape crs)) :: IO [Only (StructValue FieldValue)]) >>= (@?= [Only (nullShape crs)]) + (query conn "SELECT typeof(?)" (Only (nullShape "OGC:CRS84")) :: IO [Only Text]) >>= (@?= [Only "STRUCT(shape GEOMETRY)"]) + , testCase "connection options validate CRS definitions and pass configuration flags" $ do + forM_ ["", "OGC:CRS84\0bad"] \crs -> + assertFailureIO (withConnectionWithOptions ":memory:" defaultConnectionOptions{connectionGeometryCRS = [crs]} (const (pure ()))) + let options = defaultConnectionOptions{connectionConfig = [("threads", "1")], connectionGeometryCRS = ["local grid", "OGC:CRS84", "OGC:CRS84"]} + withConnectionWithOptions ":memory:" options \conn -> + (query_ conn "SELECT current_setting('threads')" :: IO [Only Int64]) >>= (@?= [Only 1]) + , testCase "logical type construction drops CRS" $ do + let logical = LogicalTypeList (LogicalTypeGeometry (Just "OGC:CRS84")) + bracket (logicalTypeFromRep logical) destroyLogicalType logicalTypeToRep + >>= (@?= LogicalTypeList (LogicalTypeGeometry Nothing)) + , testGroup + "folds cross chunk boundaries and preserve CRS" + [ testCase mode $ + withConnectionWithConfig ":memory:" [("threads", "1")] \conn -> do + [Only expected] <- query_ conn "SELECT 'POINT (1 2)'::GEOMETRY('OGC:CRS84')" :: IO [Only RawGeometry] + count <- foldRows conn "SELECT 'POINT (1 2)'::GEOMETRY('OGC:CRS84') FROM range(5000)" (0 :: Int64) \n (Only actual) -> do + actual @?= expected + pure (n + 1) + count @?= 5000 + | (mode, foldRows) <- [("materialized", fold_), ("deprecated streaming", Streaming.fold_)] + ] + , testCase "invalid binary import leaves the connection usable" $ + withConnection ":memory:" \conn -> do + assertFailureIO (query conn "SELECT system.main.ST_GeomFromWKB(?)" (Only (BS.pack [1, 1, 0, 0, 0])) :: IO [Only RawGeometry]) + (query_ conn "SELECT 42" :: IO [Only Int64]) >>= (@?= [Only 42]) + ] + +-- | A one-field STRUCT with a NULL GEOMETRY of the given CRS. +nullShape :: Text -> StructValue FieldValue +nullShape crs = + StructValue + { structValueFields = listArray (0, 0) [StructField "shape" FieldNull] + , structValueTypes = listArray (0, 0) [StructField "shape" (LogicalTypeGeometry (Just crs))] + , structValueIndex = Map.singleton "shape" 0 + } + +-- | Exercise all coordinate layouts with explicit point and sequence constructors. +shapeRoundTrips :: (G.Coordinate coord) => Connection -> G.Dimensions -> (coord -> G.Point) -> (U.Vector coord -> G.Coordinates) -> coord -> coord -> Assertion +shapeRoundTrips conn dimensions point coordinates first second = do + let line = coordinates (U.fromList [first, second]) + ring = coordinates (U.fromList [first, second, first, first]) + emptyLine = coordinates U.empty + emptyPoint = G.EmptyPoint dimensions + polygon = G.PolygonRings ring V.empty + emptyPolygon = G.PolygonRings emptyLine V.empty + values = + [ G.PointGeometry (point first) + , G.PointGeometry emptyPoint + , G.LineString line + , G.LineString emptyLine + , G.Polygon polygon + , G.Polygon emptyPolygon + , G.MultiPoint (U.fromList [point first, emptyPoint, point second]) + , G.MultiPoint U.empty + , G.MultiLineString (V.fromList [line, emptyLine]) + , G.MultiPolygon (V.fromList [polygon, emptyPolygon]) + , G.GeometryCollection (V.fromList [G.PointGeometry emptyPoint, G.LineString line, G.GeometryCollection V.empty]) + ] + forM_ values \shape -> do + rows <- query conn "SELECT ?, ST_CRS(?)" (shape, shape) + rows @?= [(shape, Nothing :: Maybe Text)] + raw <- either assertFailure pure (toRawGeometry shape) + fromRawGeometry raw @?= Right shape + let annotated = raw{rawGeometryCRS = Just "local' íslenska λ"} + (importRaw conn annotated) + >>= (@?= [Only annotated]) + +-- | Import WKB without converting a missing CRS to SQL NULL. +importRaw :: Connection -> RawGeometry -> IO [Only RawGeometry] +importRaw conn raw = case rawGeometryCRS raw of + Nothing -> query conn "SELECT system.main.ST_GeomFromWKB(?)" (Only (rawGeometryWKB raw)) + Just crs -> query conn "SELECT system.main.ST_SetCRS(system.main.ST_GeomFromWKB(?), ?)" (rawGeometryWKB raw, crs) + +-- | Bind raw bytes and CRS as separate SQL parameters. +rawParameters :: RawGeometry -> (BS.ByteString, Maybe Text) +rawParameters raw = (rawGeometryWKB raw, rawGeometryCRS raw) + +-- | Check that generic values cannot silently normalize raw geometry. +assertRawBindingFailure :: IO a -> Assertion +assertRawBindingFailure action = do + result <- try (action >> pure ()) :: IO (Either SomeException ()) + case result of + Left err -> assertBool (displayException err) ("raw GEOMETRY binding requires explicit" `isInfixOf` displayException err) + Right () -> assertFailure "expected raw geometry binding rejection" + +-- | Native shapes cover all families, dimensions, empty children, and exact doubles. +shapes :: [Text] +shapes = + [ "POINT (0.10000000000000002 -0.0)" + , "POINT Z (1 2 3)" + , "POINT M (1 2 3)" + , "POINT ZM (1 2 3 4)" + , "POINT EMPTY" + , "LINESTRING EMPTY" + , "LINESTRING (1 2, 3 4)" + , "POLYGON ((0 0, 1 0, 1 1, 0 0))" + , "POLYGON ((0 0, 10 0, 10 10, 0 10, 0 0), (2 2, 2 4, 4 4, 4 2, 2 2))" + , "MULTIPOINT ((1 2), EMPTY, (3 4))" + , "MULTILINESTRING ((1 2, 3 4), EMPTY)" + , "MULTIPOLYGON (((0 0, 1 0, 1 1, 0 0)), EMPTY)" + , "GEOMETRYCOLLECTION (POINT (1 2), LINESTRING (0 0, 1 1))" + , "GEOMETRYCOLLECTION ZM (POINT ZM (1 2 3 4), POINT ZM EMPTY)" + , "POINT Z EMPTY" + , "LINESTRING M EMPTY" + , "GEOMETRYCOLLECTION ZM EMPTY" + ] diff --git a/test/Spec.hs b/test/Spec.hs --- a/test/Spec.hs +++ b/test/Spec.hs @@ -82,6 +82,7 @@ import Database.DuckDB.Simple.Time (Unbounded (..)) import ExtensionRegressionTests (extensionRegressionTests) import GHC.Generics (Generic) +import qualified GeometryRegressionTests import Numeric.Natural (Natural) import Properties (roundTripTests) import StreamingTests (nativeStreamingTests) @@ -92,6 +93,7 @@ import Test.Tasty.QuickCheck (testProperty, (===)) import TimeTests (timeTests) import ValueRegressionTests (valueRegressionTests) +import qualified VariantRegressionTests data Person = Person { personId :: Int @@ -209,6 +211,8 @@ , cancellationTests , nativeStreamingTests , extensionRegressionTests + , GeometryRegressionTests.tests + , VariantRegressionTests.tests , valueRegressionTests , timeTests , withConnectionTests diff --git a/test/TestUtils.hs b/test/TestUtils.hs new file mode 100644 --- /dev/null +++ b/test/TestUtils.hs @@ -0,0 +1,13 @@ +-- | Assertions that several test modules share. +module TestUtils (assertFailureIO) where + +import Control.Exception (SomeException, try) +import Test.Tasty.HUnit (Assertion, assertFailure) + +-- | Require an exception without depending on native error text. +assertFailureIO :: IO a -> Assertion +assertFailureIO action = do + result <- try (action >> pure ()) :: IO (Either SomeException ()) + case result of + Left _ -> pure () + Right () -> assertFailure "expected an exception" diff --git a/test/ValueRegressionTests.hs b/test/ValueRegressionTests.hs --- a/test/ValueRegressionTests.hs +++ b/test/ValueRegressionTests.hs @@ -26,7 +26,7 @@ import Database.DuckDB.Simple import Database.DuckDB.Simple.FromField (BitString (..), DecimalValue (..), FieldValue (..), TimeWithZone (..), bsFromBool) import Database.DuckDB.Simple.Generic (ViaDuckDB (..), genericFromFieldValue, genericToStructValue) -import Database.DuckDB.Simple.Internal (withConnectionHandle) +import Database.DuckDB.Simple.Internal (destroyValue, withConnectionHandle) import Database.DuckDB.Simple.LogicalRep import Database.DuckDB.Simple.Time (Unbounded (..)) import Database.DuckDB.Simple.ToField (ToDuckValue (..)) @@ -37,6 +37,7 @@ import GHC.Generics (Generic) import Test.Tasty (TestTree, defaultMain, testGroup) import Test.Tasty.HUnit +import TestUtils (assertFailureIO) -- | Native timestamps used to test the full storage range. data NativeTimestamp = Seconds Int64 | Milliseconds Int64 | Microseconds Int64 | Nanoseconds Int64 @@ -53,6 +54,27 @@ instance ToField NativeTimestamp +-- | An array element with a ToField instance and no ToDuckValue instance. +newtype DelegatedInteger = DelegatedInteger Int64 + +instance DuckDBColumnType DelegatedInteger where + duckdbColumnTypeFor _ = "BIGINT" + +instance ToField DelegatedInteger where + toField (DelegatedInteger value) = toField value + +-- | Bind an array through its connection-free value. +newtype Scores = Scores (Array Int Int64) + deriving (Show) + +instance DuckDBColumnType Scores where + duckdbColumnTypeFor _ = "BIGINT[]" + +instance ToDuckValue Scores where + toDuckValue (Scores values) = toDuckValue values + +instance ToField Scores + -- | Record with identical field types to detect positional decoding. data NamedRecord = NamedRecord {firstValue :: Int64, secondValue :: Int64} deriving (Eq, Show, Generic) @@ -61,6 +83,11 @@ data NullableSum = EmptyMember | DataMember Int64 deriving (Eq, Show, Generic) +-- | A generic record for array parameters. +data Pair = Pair {pairName :: Text, pairValue :: Int64} + deriving stock (Eq, Show, Generic) + deriving (DuckDBColumnType, ToField, FromField) via (ViaDuckDB Pair) + -- | Generic nullary constructors retain their existing UNION schema. data Colour = Red | Blue deriving stock (Eq, Show, Generic) @@ -86,6 +113,47 @@ let members = listArray (0, 1) [UnionMemberType "number" (LogicalTypeScalar DuckDBTypeBigInt), UnionMemberType "text" (LogicalTypeScalar DuckDBTypeVarchar)] original = UnionValue 0 "number" FieldNull members (query conn "SELECT ?" (Only original) :: IO [Only (UnionValue FieldValue)]) >>= (@?= [Only original]) + , testCase "array elements need only a ToField instance" $ withConnection ":memory:" \conn -> do + let values = listArray (0 :: Int, 2) [Nothing, Just (DelegatedInteger 41), Just (DelegatedInteger 99)] + expected = listArray (0, 2) [Nothing, Just 41, Just 99] :: Array Int (Maybe Int64) + (query conn "SELECT ?" (Only values) :: IO [Only (Array Int (Maybe Int64))]) >>= (@?= [Only expected]) + , testCase "array children use default ToDuckValue instances" $ withConnection ":memory:" \conn -> do + let units = [-1, 0, 1234567] + values = listArray (0 :: Int, 2) (map Microseconds units) + expected = listArray (0, 2) (map (utcToLocalTime utc . posixSecondsToUTCTime . fromRational . (% 1000000) . toInteger) units) + (query conn "SELECT ?" (Only values) :: IO [Only (Array Int LocalTime)]) >>= (@?= [Only expected]) + , testCase "arrays bind through ToDuckValue without a connection" $ withConnection ":memory:" \conn -> do + let values = listArray (0, 2) [1, -2, 3] + (query conn "SELECT typeof(?), ?" (Scores values, Scores values) :: IO [(Text, Array Int Int64)]) >>= (@?= [("BIGINT[3]", values)]) + , testCase "arrays of STRUCT, record, and ARRAY elements take the type of a present element" $ withConnection ":memory:" \conn -> do + [Only struct] <- query_ conn "SELECT {'a': 1, 'b': 'x'}" :: IO [Only (StructValue FieldValue)] + let structs = listArray (0, 1) [struct, struct] + (query conn "SELECT typeof(?), ?" (structs, structs) :: IO [(Text, Array Int (StructValue FieldValue))]) >>= (@?= [("STRUCT(a INTEGER, b VARCHAR)[2]", structs)]) + let pairs = listArray (0, 2) [Nothing, Just (Pair "x" 1), Just (Pair "y" 2)] + (query conn "SELECT typeof(?), ?" (pairs, pairs) :: IO [(Text, Array Int (Maybe Pair))]) >>= (@?= [("STRUCT(pairName VARCHAR, pairValue BIGINT)[3]", pairs)]) + let nested = listArray (0, 1) [listArray (0, 1) [1, 2], listArray (0, 1) [3, 4]] :: Array Int (Array Int Int64) + (query conn "SELECT typeof(?), ?" (nested, nested) :: IO [(Text, Array Int (Array Int Int64))]) >>= (@?= [("BIGINT[2][2]", nested)]) + , testCase "empty arrays of record elements have no element type" $ withConnection ":memory:" \conn -> + assertFailureIO (query conn "SELECT ?" (Only (listArray (0, -1) [] :: Array Int Pair)) :: IO [Only Text]) + , testCase "composite arrays reject schemas that would discard fields or digits" $ withConnection ":memory:" \conn -> do + forM_ + [ ("SELECT {'a': 1::INTEGER}", "SELECT {'a': 2::INTEGER, 'b': 3::INTEGER}") + , ("SELECT {'a': 1::INTEGER}", "SELECT {'a': 1.6::DOUBLE}") + , ("SELECT {'a': 1.2::DECIMAL(4,1)}", "SELECT {'a': 1.26::DECIMAL(4,2)}") + , ("SELECT {'a': [1.2::DECIMAL(4,1)]}", "SELECT {'a': [1.26::DECIMAL(4,2)]}") + ] + \(firstQuery, secondQuery) -> do + [Only first] <- query_ conn firstQuery :: IO [Only (StructValue FieldValue)] + [Only second] <- query_ conn secondQuery :: IO [Only (StructValue FieldValue)] + forM_ [[Nothing, Just first, Just second], [Just second, Nothing, Just first]] \elements -> do + let values = listArray (0 :: Int, 2) elements + assertFailureIO (query conn "SELECT ?" (Only values) :: IO [Only FieldValue]) + assertFailureIO (bracket (toDuckValue values) destroyValue (const (pure ()))) + (query_ conn "SELECT 42" :: IO [Only Int64]) >>= (@?= [Only 42]) + , testCase "GEOMETRY decodes as well-known binary" $ withConnection ":memory:" \conn -> do + [Only bytes] <- query_ conn "SELECT 'POINT(1 2)'::GEOMETRY" :: IO [Only BS.ByteString] + BS.length bytes @?= 21 + (query conn "SELECT ST_AsText(ST_GeomFromWKB(?))" (Only bytes) :: IO [Only Text]) >>= (@?= [Only "POINT (1 2)"]) , testCase "Float parameter retains FLOAT type" $ withConnection ":memory:" \conn -> do (query conn "SELECT typeof(?), ?" (1.25 :: Float, 1.25 :: Float) :: IO [(Text, Float)]) >>= (@?= [("FLOAT", 1.25)]) , testCase "Float special values survive decoding" $ withConnection ":memory:" \conn -> do diff --git a/test/VariantRegressionTests.hs b/test/VariantRegressionTests.hs new file mode 100644 --- /dev/null +++ b/test/VariantRegressionTests.hs @@ -0,0 +1,342 @@ +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE DerivingVia #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE PatternSynonyms #-} +{-# OPTIONS_GHC -Wno-deprecations #-} + +module VariantRegressionTests (tests) where + +import Control.Exception (SomeException, bracket, displayException, try) +import Control.Monad (forM_, void, when) +import Data.Array (Array, listArray) +import qualified Data.ByteString as BS +import qualified Data.Geometry as G +import Data.Int (Int64) +import Data.List (isInfixOf) +import Data.Text (Text) +import qualified Data.Text as Text +import qualified Data.Vector as V +import Database.DuckDB.FFI +import Database.DuckDB.Simple +import qualified Database.DuckDB.Simple.Deprecated.Streaming as Streaming +import Database.DuckDB.Simple.FromField (BitString (..), DecimalValue (..), FieldValue (..), StructValue (..), UnionValue) +import Database.DuckDB.Simple.Generic (ViaDuckDB (..)) +import Database.DuckDB.Simple.Geometry (RawGeometry (..), toRawGeometry) +import Database.DuckDB.Simple.Internal (destroyValue, withConnectionHandle) +import Database.DuckDB.Simple.LogicalRep (LogicalTypeRep (..), destroyLogicalType, logicalTypeFromRep) +import Database.DuckDB.Simple.Variant +import Foreign.C.String (withCString) +import Foreign.Marshal.Alloc (alloca) +import Foreign.Ptr (nullPtr) +import Foreign.Storable (peek, poke) +import GHC.Float (castDoubleToWord64, castFloatToWord32, castWord64ToDouble) +import GHC.Generics (Generic) +import System.Directory (doesFileExist, getTemporaryDirectory, removeFile) +import System.IO (hClose, openBinaryTempFile) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit +import TestUtils (assertFailureIO) + +-- | A generic record checks the bridge to composite type metadata. +newtype VariantRecord = VariantRecord {payload :: Variant} + deriving stock (Eq, Show, Generic) + deriving (DuckDBColumnType, ToField, FromField) via (ViaDuckDB VariantRecord) + +-- | Check VARIANT payloads against native decoding and parameter round trips. +tests :: TestTree +tests = + testGroup + "variant integration" + [ testGroup + "scalar payloads decode like native values" + [ testCase (Text.unpack sql) $ withConnection ":memory:" \conn -> do + [(Variant actual, native)] <- query_ conn (Query ("SELECT (" <> sql <> ")::VARIANT, " <> sql)) :: IO [(Variant, FieldValue)] + actual @?= native + roundTrip conn native + | sql <- scalarCases + ] + , testCase "explicit casts bind plain parameters" $ withConnection ":memory:" \conn -> do + [Only record] <- query_ conn "SELECT {'a': 1, 'b': 'x'}" :: IO [Only (StructValue FieldValue)] + (query conn "SELECT ?::VARIANT, ?::VARIANT, ?::VARIANT, ?::VARIANT" (42 :: Int64, "two" :: Text, record, listArray (0, 1) [1, 2] :: Array Int Int64) :: IO [(Variant, Variant, Variant, Variant)]) + >>= (@?= [(Variant (FieldInt64 42), Variant (FieldText "two"), Variant (variantObject [("a", FieldInt32 1), ("b", FieldText "x")]), Variant (FieldList [FieldInt64 1, FieldInt64 2]))]) + (query conn "SELECT [NULL::VARIANT, ?::VARIANT, ?::VARIANT]" (42 :: Int64, "two" :: Text) :: IO [Only [Variant]]) + >>= (@?= [Only [Variant FieldNull, Variant (FieldInt64 42), Variant (FieldText "two")]]) + _ <- execute_ conn "CREATE TABLE variants (v VARIANT)" + _ <- executeMany conn "INSERT INTO variants VALUES (?)" [Only (7 :: Int64), Only 8] + _ <- execute conn "INSERT INTO variants VALUES (?)" (Only record) + (query_ conn "SELECT v FROM variants" :: IO [Only Variant]) + >>= (@?= [Only (Variant (FieldInt64 7)), Only (Variant (FieldInt64 8)), Only (Variant (variantObject [("a", FieldInt32 1), ("b", FieldText "x")]))]) + , testCase "existing FromField instances read VARIANT payloads" $ withConnection ":memory:" \conn -> do + (query_ conn "SELECT 42::BIGINT::VARIANT, 'x'::VARIANT, [1, 2, 3]::VARIANT" :: IO [(Int64, Text, [Int64])]) + >>= (@?= [(42, "x", [1, 2, 3])]) + (query_ conn "SELECT {'payload': {'x': 18446744073709551615::UBIGINT}::VARIANT}" :: IO [Only VariantRecord]) + >>= (@?= [Only (VariantRecord (Variant (variantObject [("x", FieldWord64 maxBound)])))]) + , testCase "bound payloads keep integer widths and decimal scale" $ withConnection ":memory:" \conn -> + forM_ [("'-128'::TINYINT", FieldInt8 minBound), ("12.34::DECIMAL(4,2)", FieldDecimal (DecimalValue 4 2 1234)), ("255::UTINYINT", FieldWord8 maxBound)] \(sql, value) -> + (query conn (Query ("SELECT variant_typeof(?) = variant_typeof((" <> sql <> ")::VARIANT)")) (Only (Variant value)) :: IO [Only Bool]) + >>= (@?= [Only True]) + , testCase "floating special values retain type and bits" $ withConnection ":memory:" \conn -> do + forM_ [0, -0.0, 1 / 0, -1 / 0, 0 / 0] \value -> do + [Only (Variant actual)] <- query conn "SELECT ?" (Only (Variant (FieldDouble value))) + case actual of + FieldDouble decoded + | isNaN value -> assertBool "expected Double NaN" (isNaN decoded) + | otherwise -> castDoubleToWord64 decoded @?= castDoubleToWord64 value + _ -> assertFailure (show actual) + forM_ [0, -0.0, 1 / 0, -1 / 0, 0 / 0] \value -> do + [Only (Variant actual)] <- query conn "SELECT ?" (Only (Variant (FieldFloat value))) + case actual of + FieldFloat decoded + | isNaN value -> assertBool "expected Float NaN" (isNaN decoded) + | otherwise -> castFloatToWord32 decoded @?= castFloatToWord32 value + _ -> assertFailure (show actual) + , testCase "TIMETZ offsets decode like native values" $ withConnection ":memory:" \conn -> do + forM_ ["'12:00:00+01:23'::TIMETZ", "'23:59:59.999999-15:59'::TIMETZ", "'00:00:00+00'::TIMETZ"] \sql -> do + [(Variant actual, native)] <- query_ conn (Query ("SELECT (" <> sql <> ")::VARIANT, " <> sql)) :: IO [(Variant, FieldValue)] + actual @?= native + roundTrip conn native + assertFailureIO (query_ conn "SELECT '12:00:00+01:23:45'::TIMETZ::VARIANT" :: IO [Only Variant]) + , testCase "wide finite timestamp payloads round trip without losing digits" $ withConnection ":memory:" \conn -> do + void (execute_ conn "CREATE TABLE wide_timestamps (seconds TIMESTAMP_S, millis TIMESTAMP_MS)") + appendWideTimestamps conn [(30000000000000, 30000000000000001), (-30000000000000, -30000000000000001), (maxBound - 1, maxBound - 1), (minBound, minBound)] + rows <- query_ conn "SELECT seconds::VARIANT, seconds, millis::VARIANT, millis FROM wide_timestamps" :: IO [(Variant, FieldValue, Variant, FieldValue)] + length rows @?= 4 + forM_ rows \(seconds, nativeSeconds, millis, nativeMillis) -> + forM_ [(seconds, nativeSeconds), (millis, nativeMillis)] \(Variant actual, native) -> do + actual @?= native + roundTrip conn actual + , testCase "text and binary values preserve embedded NUL" $ withConnection ":memory:" \conn -> do + roundTrip conn (FieldText "before\0after íslenska λ 😀") + roundTrip conn (FieldBlob (BS.pack [0, 1, 127, 128, 255, 0])) + , testCase "objects preserve case-sensitive names and heterogeneous children" $ withConnection ":memory:" \conn -> do + let value = + variantObject + [ ("Case", FieldInt64 9007199254740993) + , ("case", FieldDecimal (DecimalValue 38 9 1234567890123456789)) + , ("quote' íslenska λ", FieldList [FieldNull, FieldText "x", FieldList [], variantObject []]) + ] + roundTrip conn value + (query conn "SELECT variant_extract(?, 'Case')::BIGINT, variant_extract(?, 'case')::DECIMAL(38,9)::VARCHAR" (Variant value, Variant value) :: IO [(Int64, Text)]) + >>= (@?= [(9007199254740993, "1234567890.123456789")]) + , testCase "empty containers and NULL stay distinct" $ withConnection ":memory:" \conn -> do + forM_ [FieldNull, FieldList [], variantObject [], FieldList [FieldNull], variantObject [("x", FieldNull)]] (roundTrip conn) + (query_ conn "SELECT NULL::VARIANT" :: IO [Only (Maybe Variant)]) >>= (@?= [Only Nothing]) + (query conn "SELECT ?" (Only (Variant FieldNull)) :: IO [Only FieldValue]) >>= (@?= [Only FieldNull]) + , testCase "deep native arrays decode without a library depth limit" $ withConnection ":memory:" \conn -> + forM_ [128, 256 :: Int] \depth -> do + let expected = Variant (foldr (const (FieldList . pure)) (FieldWord64 7) [1 .. depth]) + (query conn "SELECT (repeat('[', ?) || '7' || repeat(']', ?))::JSON::VARIANT" (depth, depth) :: IO [Only Variant]) + >>= (@?= [Only expected]) + , testCase "ARRAY parameters and generic records contain VARIANT" $ withConnection ":memory:" \conn -> do + let array = listArray (0, 2) [Variant (FieldInt8 1), Variant (FieldText "two"), Variant FieldNull] + (query conn "SELECT ?" (Only array) :: IO [Only (Array Int Variant)]) >>= (@?= [Only array]) + let record = VariantRecord (Variant (variantObject [("x", FieldWord64 maxBound)])) + (query conn "SELECT ?" (Only record) :: IO [Only VariantRecord]) >>= (@?= [Only record]) + , testCase "nullable ARRAY parameters preserve heterogeneous VARIANT values" $ withConnection ":memory:" \conn -> do + let values = listArray (0, 3) [Nothing, Just (Variant (FieldInt64 42)), Just (Variant (FieldText "two")), Just (Variant (variantObject [("xs", FieldList [FieldNull, FieldBool True])]))] + (query conn "SELECT typeof(?), ?" (values, values) :: IO [(Text, Array Int (Maybe Variant))]) >>= (@?= [("VARIANT[4]", values)]) + , testCase "empty ARRAY parameters retain VARIANT element type" $ withConnection ":memory:" \conn -> do + let values = listArray (0, -1) [] :: Array Int (Maybe Variant) + (query conn "SELECT typeof(?)" (Only values) :: IO [Only Text]) >>= (@?= [Only "VARIANT[ANY]"]) + , testCase "parameter binding supplies a complete VARIANT type" $ withConnection ":memory:" \conn -> do + let value = variantObject [("xs", FieldList [FieldNull, FieldInt64 42])] + roundTrip conn value + (query conn "SELECT typeof(?)" (Only (Variant value)) :: IO [Only Text]) >>= (@?= [Only "VARIANT"]) + , testCase "the first VARIANT parameter does not end a streaming result" $ withConnection ":memory:" \conn -> + withStatement conn "SELECT ?" \stmt -> do + total <- Streaming.fold_ conn "SELECT i FROM range(100000) t(i)" 0 \acc (Only i) -> do + when (i == 0) (bind stmt [toField (Variant (FieldInt64 42))]) + pure (acc + i) + total @?= (sum [0 .. 99999] :: Int64) + , testCase "VARIANT type construction raises an error" $ do + result <- try (bracket (logicalTypeFromRep (LogicalTypeScalar DuckDBTypeVariant)) destroyLogicalType (const (pure ()))) + case result of + Left err -> assertBool (displayException (err :: SomeException)) ("?::VARIANT" `isInfixOf` displayException err) + Right () -> assertFailure "expected VARIANT type rejection" + , testCase "VARIANT payloads and GEOMETRY CRS metadata bind together" $ withConnection ":memory:" \conn -> do + [Only value] <- query_ conn "SELECT {'payload': 42::VARIANT, 'shape': NULL::GEOMETRY('OGC:CRS84')}" :: IO [Only (StructValue FieldValue)] + (query conn "SELECT typeof(?)" (Only value) :: IO [Only Text]) >>= (@?= [Only "STRUCT(payload VARIANT, shape GEOMETRY('OGC:CRS84'))"]) + (query conn "SELECT ?" (Only value) :: IO [Only (StructValue FieldValue)]) >>= (@?= [Only value]) + , testCase "inactive VARIANT UNION members have complete types" $ withConnection ":memory:" \conn -> do + [Only value] <- query_ conn "SELECT union_value(number := 42::BIGINT)::UNION(number BIGINT, payload VARIANT)" :: IO [Only (UnionValue FieldValue)] + (query conn "SELECT ?" (Only value) :: IO [Only (UnionValue FieldValue)]) >>= (@?= [Only value]) + , testCase "native LIST, ARRAY, MAP and STRUCT metadata round trip" $ withConnection ":memory:" \conn -> do + [Only value] <- + query_ + conn + "SELECT {'xs': [1::VARIANT, 'two'::VARIANT, NULL], 'fixed': [42::VARIANT]::VARIANT[1], 'map': MAP {'x': {'a': 9}::VARIANT}}" :: + IO [Only (StructValue FieldValue)] + (query conn "SELECT ?" (Only value) :: IO [Only (StructValue FieldValue)]) >>= (@?= [Only value]) + , testCase "UNION payloads retain VARIANT type, including NULL" $ withConnection ":memory:" \conn -> + forM_ ["SELECT union_value(v := {'a': 42}::VARIANT)", "SELECT union_value(v := NULL::VARIANT)"] \sql -> do + [Only value] <- query_ conn sql :: IO [Only (UnionValue FieldValue)] + (query conn "SELECT ?" (Only value) :: IO [Only (UnionValue FieldValue)]) >>= (@?= [Only value]) + , testGroup + "heterogeneous folds cross chunk boundaries" + [ testCase mode $ withConnectionWithConfig ":memory:" [("threads", "1")] \conn -> do + count <- foldRows + conn + "SELECT CASE WHEN i % 3 = 0 THEN i::VARIANT WHEN i % 3 = 1 THEN 'text'::VARIANT ELSE {'n': i, 'xs': [NULL, i]}::VARIANT END FROM range(5000) t(i)" + (0 :: Int64) + \n (Only (Variant actual)) -> do + let expected = case n `mod` 3 of + 0 -> FieldInt64 n + 1 -> FieldText "text" + _ -> variantObject [("n", FieldInt64 n), ("xs", FieldList [FieldNull, FieldInt64 n])] + actual @?= expected + pure (n + 1) + count @?= 5000 + | (mode, foldRows) <- [("materialized", fold_), ("deprecated streaming", Streaming.fold_)] + ] + , testCase "filtered and reordered rows use the right child offsets" $ withConnection ":memory:" \conn -> do + rows <- query_ conn "SELECT {'n': i, 'xs': [i, i + 1]}::VARIANT FROM range(10000) t(i) WHERE i % 97 = 0 ORDER BY i DESC LIMIT 40" + let expected = [Only (Variant (variantObject [("n", FieldInt64 n), ("xs", FieldList [FieldInt64 n, FieldInt64 (n + 1)])])) | n <- take 40 (reverse [0, 97 .. 9999])] + rows @?= expected + , testCase "file-backed values survive checkpoint and reopen" $ + bracket newDatabase removeDatabase \path -> do + withConnectionWithConfig path [("storage_compatibility_version", "v1.5.0")] \conn -> do + void (execute_ conn "CREATE TABLE stored AS SELECT i, {'n': i, 'xs': [i, NULL], 'text': i::VARCHAR}::VARIANT AS value, 'POINT (1 2)'::GEOMETRY('OGC:CRS84') AS geometry FROM range(10000) t(i)") + void (execute_ conn "CHECKPOINT") + withConnectionWithConfig path [("storage_compatibility_version", "v1.5.0")] \conn -> do + rows <- query_ conn "SELECT value FROM stored WHERE i % 97 = 0 ORDER BY i DESC LIMIT 40" :: IO [Only Variant] + let expected = + [ Only (Variant (variantObject [("n", FieldInt64 n), ("xs", FieldList [FieldInt64 n, FieldNull]), ("text", FieldText (Text.pack (show n)))])) + | n <- take 40 (reverse [0, 97 .. 9999]) + ] + rows @?= expected + [Only geometry] <- query_ conn "SELECT geometry FROM stored LIMIT 1" :: IO [Only RawGeometry] + rawGeometryCRS geometry @?= Just "OGC:CRS84" + BS.length (rawGeometryWKB geometry) @?= 21 + , testCase "geometry payload keeps WKB; native VARIANT does not retain CRS" $ withConnection ":memory:" \conn -> do + [Only geometry] <- query_ conn "SELECT 'POINT ZM (1 2 3 4)'::GEOMETRY('OGC:CRS84')" :: IO [Only RawGeometry] + [Only value] <- query_ conn "SELECT 'POINT ZM (1 2 3 4)'::GEOMETRY('OGC:CRS84')::VARIANT" :: IO [Only Variant] + value @?= geometryPayload (rawGeometryWKB geometry) + (query conn "SELECT system.main.ST_GeomFromWKB(?)::VARIANT" (Only (rawGeometryWKB geometry)) :: IO [Only Variant]) >>= (@?= [Only value]) + (query conn "SELECT system.main.ST_CRS((system.main.ST_GeomFromWKB(?)::VARIANT)::GEOMETRY)" (Only (rawGeometryWKB geometry)) :: IO [Only (Maybe Text)]) >>= (@?= [Only Nothing]) + forM_ [variantPayload value, FieldList [variantPayload value], variantObject [("shape", variantPayload value)]] \input -> do + result <- try (query conn "SELECT ?" (Only (Variant input)) :: IO [Only Variant]) :: IO (Either SomeException [Only Variant]) + case result of + Left err -> assertBool (displayException err) ("raw GEOMETRY binding requires explicit" `isInfixOf` displayException err) + Right _ -> assertFailure "expected raw geometry binding rejection" + (query_ conn "SELECT 42" :: IO [Only Int64]) >>= (@?= [Only 42]) + , testCase "geometry payload retains empty layout tags and native NaN points" $ withConnection ":memory:" \conn -> + forM_ ["GEOMETRYCOLLECTION ZM EMPTY", "MULTIPOLYGON M EMPTY", "POINT Z (NaN NaN 7)"] \wkt -> do + [Only geometry] <- query conn "SELECT ?::GEOMETRY" (Only (wkt :: Text)) :: IO [Only RawGeometry] + (query conn "SELECT system.main.ST_GeomFromWKB(?)::VARIANT" (Only (rawGeometryWKB geometry)) :: IO [Only Variant]) >>= (@?= [Only (geometryPayload (rawGeometryWKB geometry))]) + , testCase "geometry payload conversion ignores user macros with built-in names" $ withConnection ":memory:" \conn -> do + [Only geometry] <- query_ conn "SELECT 'POINT (1 2)'::GEOMETRY" :: IO [Only RawGeometry] + _ <- execute_ conn "CREATE MACRO ST_AsText(x) AS 42" + _ <- execute_ conn "CREATE MACRO ST_GeomFromWKB(x) AS 'POINT (9 9)'::GEOMETRY" + (query conn "SELECT system.main.ST_GeomFromWKB(?)::VARIANT" (Only (rawGeometryWKB geometry)) :: IO [Only Variant]) >>= (@?= [Only (geometryPayload (rawGeometryWKB geometry))]) + , testCase "binary geometry payloads retain NaN bits and mixed member layouts" $ withConnection ":memory:" \conn -> + forM_ + [ G.PointGeometry (G.PointXYZ (G.XYZ (castWord64ToDouble 0x7ff8000000000001) (castWord64ToDouble 0x7ff8000000000002) 7)) + , G.GeometryCollection (V.fromList [G.PointGeometry (G.PointXY (G.XY 1 2)), G.PointGeometry (G.PointXYZ (G.XYZ 3 4 5))]) + ] + \shape -> do + raw <- either assertFailure pure (toRawGeometry shape) + (query conn "SELECT system.main.ST_GeomFromWKB(?)::VARIANT" (Only (rawGeometryWKB raw)) :: IO [Only Variant]) >>= (@?= [Only (geometryPayload (rawGeometryWKB raw))]) + , testCase "invalid values fail before binding and connection reuse succeeds" $ withConnection ":memory:" \conn -> do + forM_ + [ FieldHugeInt (2 ^ (127 :: Int)) + , FieldUHugeInt (-1) + , FieldUHugeInt (2 ^ (128 :: Int)) + , FieldDecimal (DecimalValue 0 0 0) + , FieldDecimal (DecimalValue 5 6 1) + , FieldDecimal (DecimalValue 3 0 1000) + , FieldBit (BitString 8 (BS.singleton 0)) + , variantObject [("same", FieldNull), ("same", FieldBool True)] + , variantObject [("", FieldBool True)] + , variantObject [("first", FieldBool False), ("", FieldBool True)] + , variantObject [("before\0after", FieldNull)] + ] + \value -> assertFailureIO (query conn "SELECT ?" (Only (Variant value)) :: IO [Only Variant]) + (query_ conn "SELECT 42" :: IO [Only Int64]) >>= (@?= [Only 42]) + ] + +-- | Bind a payload as VARIANT without a cast and read it back. +roundTrip :: Connection -> FieldValue -> Assertion +roundTrip conn expected = + (query conn "SELECT typeof(?) AS kind, ? AS value" (Variant expected, Variant expected) :: IO [(Text, Variant)]) >>= (@?= [("VARIANT", Variant expected)]) + +{- | Insert native timestamp units without the parameter API, which converts +TIMESTAMP_S and TIMESTAMP_MS parameters to microseconds before execution. +-} +appendWideTimestamps :: Connection -> [(Int64, Int64)] -> IO () +appendWideTimestamps conn rows = + withConnectionHandle conn \native -> + withCString "wide_timestamps" \table -> + alloca \appenderPtr -> do + poke appenderPtr nullPtr + bracket (c_duckdb_appender_create native nullPtr table appenderPtr) (const (void (c_duckdb_appender_destroy appenderPtr))) \created -> do + created @?= DuckDBSuccess + appender <- peek appenderPtr + forM_ rows \(seconds, millis) -> do + bracket (c_duckdb_create_timestamp_s (DuckDBTimestampS seconds)) destroyValue (c_duckdb_append_value appender) >>= (@?= DuckDBSuccess) + bracket (c_duckdb_create_timestamp_ms (DuckDBTimestampMs millis)) destroyValue (c_duckdb_append_value appender) >>= (@?= DuckDBSuccess) + c_duckdb_appender_end_row appender >>= (@?= DuckDBSuccess) + c_duckdb_appender_flush appender >>= (@?= DuckDBSuccess) + +-- | A raw geometry payload has no CRS inside a VARIANT. +geometryPayload :: BS.ByteString -> Variant +geometryPayload wkb = Variant (FieldGeometry (RawGeometry wkb Nothing)) + +-- | SQL constructors for each scalar payload tag. +scalarCases :: [Text] +scalarCases = + [ "NULL" + , "TRUE" + , "FALSE" + , "'-128'::TINYINT" + , "'-32768'::SMALLINT" + , "'-2147483648'::INTEGER" + , "9007199254740993::BIGINT" + , "'-170141183460469231731687303715884105728'::HUGEINT" + , "255::UTINYINT" + , "65535::USMALLINT" + , "4294967295::UINTEGER" + , "18446744073709551615::UBIGINT" + , "'340282366920938463463374607431768211455'::UHUGEINT" + , "1.25::FLOAT" + , "1.25::DOUBLE" + , "12.34::DECIMAL(4,2)" + , "-1234.56::DECIMAL(9,2)" + , "123456789012.345::DECIMAL(18,3)" + , "1234567890.123456789::DECIMAL(38,9)" + , "'text λ'::VARCHAR" + , "'\\x00\\xFF'::BLOB" + , "'01234567-89ab-cdef-fedc-ba9876543210'::UUID" + , "DATE '1970-01-02'" + , "DATE 'infinity'" + , "DATE '-infinity'" + , "TIME '01:02:03.456789'" + , "'01:02:03.456789012'::TIME_NS" + , "'1970-01-01 00:00:01'::TIMESTAMP_S" + , "'1969-12-31 23:59:59.999'::TIMESTAMP_MS" + , "'1970-01-01 00:00:01.234567'::TIMESTAMP" + , "'1970-01-01 00:00:01.234567891'::TIMESTAMP_NS" + , "'infinity'::TIMESTAMP_NS" + , "'1970-01-01 00:00:01.234567+00'::TIMESTAMPTZ" + , "INTERVAL '-4 MONTHS 8 DAYS 123456789 MICROSECONDS'" + , "'12345678901234567890123456789012345678901234567890'::BIGNUM" + , "'-12345678901234567890123456789012345678901234567890'::BIGNUM" + , "0::BIGNUM" + , "'10101'::BIT" + ] + +-- | Reserve a unique path and let DuckDB create its file. +newDatabase :: IO FilePath +newDatabase = do + directory <- getTemporaryDirectory + (path, handle) <- openBinaryTempFile directory "duckdb-variant" + hClose handle + removeFile path + pure path + +-- | Remove files created by this test after DuckDB closes them. +removeDatabase :: FilePath -> IO () +removeDatabase path = + forM_ [path, path <> ".wal"] \file -> do + exists <- doesFileExist file + when exists (removeFile file)