diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,4 @@ +# Changelog + +## 0.1.0.0 +- Initial release diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,373 @@ +Mozilla Public License Version 2.0 +================================== + +1. Definitions +-------------- + +1.1. "Contributor" + means each individual or legal entity that creates, contributes to + the creation of, or owns Covered Software. + +1.2. "Contributor Version" + means the combination of the Contributions of others (if any) used + by a Contributor and that particular Contributor's Contribution. + +1.3. "Contribution" + means Covered Software of a particular Contributor. + +1.4. "Covered Software" + means Source Code Form to which the initial Contributor has attached + the notice in Exhibit A, the Executable Form of such Source Code + Form, and Modifications of such Source Code Form, in each case + including portions thereof. + +1.5. "Incompatible With Secondary Licenses" + means + + (a) that the initial Contributor has attached the notice described + in Exhibit B to the Covered Software; or + + (b) that the Covered Software was made available under the terms of + version 1.1 or earlier of the License, but not also under the + terms of a Secondary License. + +1.6. "Executable Form" + means any form of the work other than Source Code Form. + +1.7. "Larger Work" + means a work that combines Covered Software with other material, in + a separate file or files, that is not Covered Software. + +1.8. "License" + means this document. + +1.9. "Licensable" + means having the right to grant, to the maximum extent possible, + whether at the time of the initial grant or subsequently, any and + all of the rights conveyed by this License. + +1.10. "Modifications" + means any of the following: + + (a) any file in Source Code Form that results from an addition to, + deletion from, or modification of the contents of Covered + Software; or + + (b) any new file in Source Code Form that contains any Covered + Software. + +1.11. "Patent Claims" of a Contributor + means any patent claim(s), including without limitation, method, + process, and apparatus claims, in any patent Licensable by such + Contributor that would be infringed, but for the grant of the + License, by the making, using, selling, offering for sale, having + made, import, or transfer of either its Contributions or its + Contributor Version. + +1.12. "Secondary License" + means either the GNU General Public License, Version 2.0, the GNU + Lesser General Public License, Version 2.1, the GNU Affero General + Public License, Version 3.0, or any later versions of those + licenses. + +1.13. "Source Code Form" + means the form of the work preferred for making modifications. + +1.14. "You" (or "Your") + means an individual or a legal entity exercising rights under this + License. For legal entities, "You" includes any entity that + controls, is controlled by, or is under common control with You. For + purposes of this definition, "control" means (a) the power, direct + or indirect, to cause the direction or management of such entity, + whether by contract or otherwise, or (b) ownership of more than + fifty percent (50%) of the outstanding shares or beneficial + ownership of such entity. + +2. License Grants and Conditions +-------------------------------- + +2.1. Grants + +Each Contributor hereby grants You a world-wide, royalty-free, +non-exclusive license: + +(a) under intellectual property rights (other than patent or trademark) + Licensable by such Contributor to use, reproduce, make available, + modify, display, perform, distribute, and otherwise exploit its + Contributions, either on an unmodified basis, with Modifications, or + as part of a Larger Work; and + +(b) under Patent Claims of such Contributor to make, use, sell, offer + for sale, have made, import, and otherwise transfer either its + Contributions or its Contributor Version. + +2.2. Effective Date + +The licenses granted in Section 2.1 with respect to any Contribution +become effective for each Contribution on the date the Contributor first +distributes such Contribution. + +2.3. Limitations on Grant Scope + +The licenses granted in this Section 2 are the only rights granted under +this License. No additional rights or licenses will be implied from the +distribution or licensing of Covered Software under this License. +Notwithstanding Section 2.1(b) above, no patent license is granted by a +Contributor: + +(a) for any code that a Contributor has removed from Covered Software; + or + +(b) for infringements caused by: (i) Your and any other third party's + modifications of Covered Software, or (ii) the combination of its + Contributions with other software (except as part of its Contributor + Version); or + +(c) under Patent Claims infringed by Covered Software in the absence of + its Contributions. + +This License does not grant any rights in the trademarks, service marks, +or logos of any Contributor (except as may be necessary to comply with +the notice requirements in Section 3.4). + +2.4. Subsequent Licenses + +No Contributor makes additional grants as a result of Your choice to +distribute the Covered Software under a subsequent version of this +License (see Section 10.2) or under the terms of a Secondary License (if +permitted under the terms of Section 3.3). + +2.5. Representation + +Each Contributor represents that the Contributor believes its +Contributions are its original creation(s) or it has sufficient rights +to grant the rights to its Contributions conveyed by this License. + +2.6. Fair Use + +This License is not intended to limit any rights You have under +applicable copyright doctrines of fair use, fair dealing, or other +equivalents. + +2.7. Conditions + +Sections 3.1, 3.2, 3.3, and 3.4 are conditions of the licenses granted +in Section 2.1. + +3. Responsibilities +------------------- + +3.1. Distribution of Source Form + +All distribution of Covered Software in Source Code Form, including any +Modifications that You create or to which You contribute, must be under +the terms of this License. You must inform recipients that the Source +Code Form of the Covered Software is governed by the terms of this +License, and how they can obtain a copy of this License. You may not +attempt to alter or restrict the recipients' rights in the Source Code +Form. + +3.2. Distribution of Executable Form + +If You distribute Covered Software in Executable Form then: + +(a) such Covered Software must also be made available in Source Code + Form, as described in Section 3.1, and You must inform recipients of + the Executable Form how they can obtain a copy of such Source Code + Form by reasonable means in a timely manner, at a charge no more + than the cost of distribution to the recipient; and + +(b) You may distribute such Executable Form under the terms of this + License, or sublicense it under different terms, provided that the + license for the Executable Form does not attempt to limit or alter + the recipients' rights in the Source Code Form under this License. + +3.3. Distribution of a Larger Work + +You may create and distribute a Larger Work under terms of Your choice, +provided that You also comply with the requirements of this License for +the Covered Software. If the Larger Work is a combination of Covered +Software with a work governed by one or more Secondary Licenses, and the +Covered Software is not Incompatible With Secondary Licenses, this +License permits You to additionally distribute such Covered Software +under the terms of such Secondary License(s), so that the recipient of +the Larger Work may, at their option, further distribute the Covered +Software under the terms of either this License or such Secondary +License(s). + +3.4. Notices + +You may not remove or alter the substance of any license notices +(including copyright notices, patent notices, disclaimers of warranty, +or limitations of liability) contained within the Source Code Form of +the Covered Software, except that You may alter any license notices to +the extent required to remedy known factual inaccuracies. + +3.5. Application of Additional Terms + +You may choose to offer, and to charge a fee for, warranty, support, +indemnity or liability obligations to one or more recipients of Covered +Software. However, You may do so only on Your own behalf, and not on +behalf of any Contributor. You must make it absolutely clear that any +such warranty, support, indemnity, or liability obligation is offered by +You alone, and You hereby agree to indemnify every Contributor for any +liability incurred by such Contributor as a result of warranty, support, +indemnity or liability terms You offer. You may include additional +disclaimers of warranty and limitations of liability specific to any +jurisdiction. + +4. Inability to Comply Due to Statute or Regulation +--------------------------------------------------- + +If it is impossible for You to comply with any of the terms of this +License with respect to some or all of the Covered Software due to +statute, judicial order, or regulation then You must: (a) comply with +the terms of this License to the maximum extent possible; and (b) +describe the limitations and the code they affect. Such description must +be placed in a text file included with all distributions of the Covered +Software under this License. Except to the extent prohibited by statute +or regulation, such description must be sufficiently detailed for a +recipient of ordinary skill to be able to understand it. + +5. Termination +-------------- + +5.1. The rights granted under this License will terminate automatically +if You fail to comply with any of its terms. However, if You become +compliant, then the rights granted under this License from a particular +Contributor are reinstated (a) provisionally, unless and until such +Contributor explicitly and finally terminates Your grants, and (b) on an +ongoing basis, if such Contributor fails to notify You of the +non-compliance by some reasonable means prior to 60 days after You have +come back into compliance. Moreover, Your grants from a particular +Contributor are reinstated on an ongoing basis if such Contributor +notifies You of the non-compliance by some reasonable means, this is the +first time You have received notice of non-compliance with this License +from such Contributor, and You become compliant prior to 30 days after +Your receipt of the notice. + +5.2. If You initiate litigation against any entity by asserting a patent +infringement claim (excluding declaratory judgment actions, +counter-claims, and cross-claims) alleging that a Contributor Version +directly or indirectly infringes any patent, then the rights granted to +You by any and all Contributors for the Covered Software under Section +2.1 of this License shall terminate. + +5.3. In the event of termination under Sections 5.1 or 5.2 above, all +end user license agreements (excluding distributors and resellers) which +have been validly granted by You or Your distributors under this License +prior to termination shall survive termination. + +************************************************************************ +* * +* 6. Disclaimer of Warranty * +* ------------------------- * +* * +* Covered Software is provided under this License on an "as is" * +* basis, without warranty of any kind, either expressed, implied, or * +* statutory, including, without limitation, warranties that the * +* Covered Software is free of defects, merchantable, fit for a * +* particular purpose or non-infringing. The entire risk as to the * +* quality and performance of the Covered Software is with You. * +* Should any Covered Software prove defective in any respect, You * +* (not any Contributor) assume the cost of any necessary servicing, * +* repair, or correction. This disclaimer of warranty constitutes an * +* essential part of this License. No use of any Covered Software is * +* authorized under this License except under this disclaimer. * +* * +************************************************************************ + +************************************************************************ +* * +* 7. Limitation of Liability * +* -------------------------- * +* * +* Under no circumstances and under no legal theory, whether tort * +* (including negligence), contract, or otherwise, shall any * +* Contributor, or anyone who distributes Covered Software as * +* permitted above, be liable to You for any direct, indirect, * +* special, incidental, or consequential damages of any character * +* including, without limitation, damages for lost profits, loss of * +* goodwill, work stoppage, computer failure or malfunction, or any * +* and all other commercial damages or losses, even if such party * +* shall have been informed of the possibility of such damages. This * +* limitation of liability shall not apply to liability for death or * +* personal injury resulting from such party's negligence to the * +* extent applicable law prohibits such limitation. Some * +* jurisdictions do not allow the exclusion or limitation of * +* incidental or consequential damages, so this exclusion and * +* limitation may not apply to You. * +* * +************************************************************************ + +8. Litigation +------------- + +Any litigation relating to this License may be brought only in the +courts of a jurisdiction where the defendant maintains its principal +place of business and such litigation shall be governed by laws of that +jurisdiction, without reference to its conflict-of-law provisions. +Nothing in this Section shall prevent a party's ability to bring +cross-claims or counter-claims. + +9. Miscellaneous +---------------- + +This License represents the complete agreement concerning the subject +matter hereof. If any provision of this License is held to be +unenforceable, such provision shall be reformed only to the extent +necessary to make it enforceable. Any law or regulation which provides +that the language of a contract shall be construed against the drafter +shall not be used to construe this License against a Contributor. + +10. Versions of the License +--------------------------- + +10.1. New Versions + +Mozilla Foundation is the license steward. Except as provided in Section +10.3, no one other than the license steward has the right to modify or +publish new versions of this License. Each version will be given a +distinguishing version number. + +10.2. Effect of New Versions + +You may distribute the Covered Software under the terms of the version +of the License under which You originally received the Covered Software, +or under the terms of any subsequent version published by the license +steward. + +10.3. Modified Versions + +If you create software not governed by this License, and you want to +create a new license for such software, you may create and use a +modified version of this License if you rename the license and remove +any references to the name of the license steward (except to note that +such modified license differs from this License). + +10.4. Distributing Source Code Form that is Incompatible With Secondary +Licenses + +If You choose to distribute Source Code Form that is Incompatible With +Secondary Licenses under the terms of this version of the License, the +notice described in Exhibit B of this License must be attached. + +Exhibit A - Source Code Form License Notice +------------------------------------------- + + This Source Code Form is subject to the terms of the Mozilla Public + License, v. 2.0. If a copy of the MPL was not distributed with this + file, You can obtain one at https://mozilla.org/MPL/2.0/. + +If it is not possible or desirable to put the notice in a particular +file, then You may include the notice in a location (such as a LICENSE +file in a relevant directory) where a recipient would be likely to look +for such a notice. + +You may add additional accurate notices of copyright ownership. + +Exhibit B - "Incompatible With Secondary Licenses" Notice +--------------------------------------------------------- + + This Source Code Form is "Incompatible With Secondary Licenses", as + defined by the Mozilla Public License, v. 2.0. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,206 @@ +# duckdb-simple + +`duckdb-simple` provides a high-level Haskell interface to DuckDB inspired by +the ergonomics of [`sqlite-simple`](https://hackage.haskell.org/package/sqlite-simple). +It builds on the low-level bindings exposed by [`duckdb-ffi`](../duckdb-ffi) and +provides a small, focused API for opening connections, running queries, binding +parameters, and decoding typed results. + +## Getting Started + +```haskell +{-# LANGUAGE OverloadedStrings #-} + +import Database.DuckDB.Simple +import Database.DuckDB.Simple.Types (Only (..)) + +main :: IO () +main = + withConnection ":memory:" \conn -> do + _ <- execute_ conn "CREATE TABLE items (id INTEGER, name TEXT)" + _ <- execute conn "INSERT INTO items VALUES (?, ?)" (1 :: Int, "banana" :: String) + rows <- query_ conn "SELECT id, name FROM items ORDER BY id" + mapM_ print (rows :: [(Int, String)]) +``` + +### Key Modules + +- `Database.DuckDB.Simple` – primary API: connections, statements, execution, + queries, and error handling. +- `Database.DuckDB.Simple.ToField` / `ToRow` – typeclasses for preparing + parameters that can be passed to `execute`/`query`. +- `Database.DuckDB.Simple.FromField` / `FromRow` – typeclasses for decoding + query results returned by `query`/`query_`. +- `Database.DuckDB.Simple.Types` – common utility types (`Query`, `Null`, + `Only`, `(:.)`, `SQLError`). +- `Database.DuckDB.Simple.Function` – register scalar Haskell functions that + can be invoked directly from SQL. + +## Querying Data + +```haskell +import Database.DuckDB.Simple +import Database.DuckDB.Simple.Types (Only (..)) + +fetchNames :: Connection -> IO [Maybe String] +fetchNames conn = do + _ <- execute_ conn "CREATE TABLE names (value TEXT)" + _ <- executeMany conn "INSERT INTO names VALUES (?)" + [Only (Just "Alice"), Only (Nothing :: Maybe String)] + fmap fromOnly <$> query_ conn "SELECT value FROM names ORDER BY value IS NULL, value" +``` + +The execution helpers return the number of affected rows (`Int`) so callers can +assert on data changes when needed. + +## Named Parameters + +duckdb-simple supports both positional (`?`) and named parameters. Named +parameters are bound with the `(:=)` helper exported from +`Database.DuckDB.Simple.ToField`. + +```haskell +import Database.DuckDB.Simple +import Database.DuckDB.Simple (NamedParam ((:=))) + +insertNamed :: Connection -> IO Int +insertNamed conn = + executeNamed conn + "INSERT INTO events VALUES ($kind, $payload)" + ["$kind" := ("metric" :: String), "$payload" := ("ok" :: String)] +``` + +DuckDB currently does not allow mixing positional and named placeholders within +the same SQL statement; the library preserves DuckDB’s error message in that +situation. Savepoints are also unavailable in DuckDB at the moment, so +`withSavepoint` throws an `SQLError` detailing the limitation. + +If the number of supplied parameters does not match the statement’s declared +placeholders—or if you attempt to bind named arguments to a positional-only +statement—`duckdb-simple` raises a `FormatError` before executing the query. + +### Decoding rows + +`FromRow` is powered by a `RowParser`, which means instances can be written in a +monadic/Applicative style and even derived generically for product types: + +```haskell +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DeriveGeneric #-} + +import Database.DuckDB.Simple +import GHC.Generics (Generic) + +data Person = Person + { personId :: Int + , personName :: Text + } + deriving stock (Show, Generic) + deriving anyclass (FromRow) + +fetchPeople :: Connection -> IO [Person] +fetchPeople conn = query_ conn "SELECT id, name FROM person ORDER BY id" +``` + +Helper combinators such as `field`, `fieldWith`, and `numFieldsRemaining` are +available when a custom instance needs fine-grained control. + +### Resource Management + +- `withConnection` and `withStatement` wrap the open/close lifecycle and guard + against exceptions; use them whenever possible to avoid leaking C handles. +- All intermediate DuckDB objects (results, prepared statements, values) are + released immediately after use. Long queries still materialise their result + sets when using the eager helpers; reach for `fold`/`fold_`/`foldNamed` (or + the lower-level `nextRow`) to stream results in constant space. +- `execute`/`query` variants reset statement bindings each run so prepared + statements can be reused safely. + +### Metadata helpers + +- `columnCount` and `columnName` expose prepared-statement metadata so you can + inspect result shapes before executing a query. +- `rowsChanged` tracks the number of rows affected by the most recent mutation + on a connection. DuckDB does not offer a `lastInsertRowId`; prefer SQL + `RETURNING` clauses when you need generated identifiers. +### Streaming Results + +`fold`, `fold_`, and `foldNamed` expose DuckDB’s chunked result API, letting you +aggregate or stream rows without materialising the entire result set: + +```haskell +import Database.DuckDB.Simple.Types (Only (..)) + +sumValues :: Connection -> IO Int +sumValues conn = + fold_ conn "SELECT n FROM stream_fold ORDER BY n" 0 $ \acc (Only n) -> + pure (acc + n) +``` + +For manual cursor-style iteration, use `nextRow`/`nextRowWith` on an open +`Statement` to pull rows one at a time and decide when to stop. + +### Temporal & Binary Types + +`duckdb-simple` now maps DuckDB temporal columns directly onto familiar +`Data.Time` types (`DATE` → `Day`, `TIME` → `TimeOfDay`, `TIMESTAMP` → +`LocalTime`/`UTCTime`). Binary blobs surface as strict `ByteString` values. Casting +logic plugs into the existing `ToField`/`FromField` classes, so round-tripping values +through prepared statements works just like the numeric and text helpers shown earlier. + +### Feature Coverage & Roadmap + +Included today: + +- Connections, prepared statements, positional/named parameter binding. +- High-level execution (`execute*`) and eager queries (`query*`, `queryNamed`). +- Streaming/folding helpers (`fold`, `foldNamed`, `fold_`, `nextRow`). +- Temporal (`Day`, `TimeOfDay`, `LocalTime`, `UTCTime`) and blob (`ByteString`) + round-trips via `FromField`/`ToField` instances. +- Row decoding via `FromField`/`FromRow`, including generic deriving support. +- Basic transaction helpers (`withTransaction`, `withSavepoint` fallback). +- Metadata helpers: `columnCount`, `columnName`, and connection-level + `rowsChanged`. + +Planned for a future release: + +- Broader type coverage for structured/list/decimal families, including UUID-friendly APIs. +- Native savepoints once DuckDB exposes the required support. + +## User-Defined Functions + +Scalar Haskell functions can be registered with DuckDB connections and used in +SQL expressions. Argument and result types reuse the existing `FromField` and +`FunctionResult` machinery, so `Maybe` values and `IO` actions work out of the +box. + +```haskell +import Data.Int (Int64) +import Database.DuckDB.Simple +import Database.DuckDB.Simple.Function (createFunction, deleteFunction) +import Database.DuckDB.Simple.Types (Only (..)) + +registerAndUse :: Connection -> IO [Only Int64] +registerAndUse conn = do + createFunction conn "hs_times_two" (\(x :: Int64) -> x * 2) + result <- query_ conn "SELECT hs_times_two(21)" :: IO [Only Int64] + deleteFunction conn "hs_times_two" + pure result +``` + +Exceptions raised while the function executes are propagated back to DuckDB as +`SQLError` values, and `deleteFunction` issues a `DROP FUNCTION IF EXISTS` +statement to remove the registration. Current DuckDB releases mark C API +registrations as internal, so the drop operation reports an error instead of +removing the function; duckdb-simple surfaces that limitation as an +`SQLError`. + +## Tests + +The test suite is built with [tasty](https://hackage.haskell.org/package/tasty) +and covers connection management, statement lifecycle, parameter binding, and +query execution. + +``` +cabal test duckdb-simple-test --test-show-details=direct +``` diff --git a/duckdb-simple.cabal b/duckdb-simple.cabal new file mode 100644 --- /dev/null +++ b/duckdb-simple.cabal @@ -0,0 +1,62 @@ +cabal-version: 3.4 +name: duckdb-simple +version: 0.1.0.0 +license: MPL-2.0 +license-file: LICENSE +author: Matthias Pall Gissurarson +maintainer: mpg@mpg.is +category: Database +build-type: Simple +extra-doc-files: README.md +homepage: https://github.com/Tritlo/duckdb-haskell +bug-reports: https://github.com/Tritlo/duckdb-haskell/issues +synopsis: Haskell FFI bindings for DuckDB +description: + This library provides a mid-level interface for interacting with DuckDB, + in the style of other "simple" libraries such as sqlite-simple and + postgresql-simple. + . + Tested with DuckDB version 1.4.0, and 1.4.0. +extra-doc-files: CHANGELOG.md, README.md + +source-repository head + type: git + location: https://github.com/Tritlo/duckdb-haskell.git + subdir: duckdb-simple + + + +library + exposed-modules: + Database.DuckDB.Simple + Database.DuckDB.Simple.Function + Database.DuckDB.Simple.FromField + Database.DuckDB.Simple.FromRow + Database.DuckDB.Simple.Internal + Database.DuckDB.Simple.Ok + Database.DuckDB.Simple.ToField + Database.DuckDB.Simple.ToRow + Database.DuckDB.Simple.Types + hs-source-dirs: src + default-language: Haskell2010 + build-depends: + base >=4.14 && <5 + , bytestring >= 0.11 && <0.12 + , duckdb-ffi >= 0.1.4 && <0.1.5 + , text >= 2.0 && <2.1 + , transformers >= 0.6 && <0.7 + , time >= 1.12 && <1.13 + +test-suite duckdb-simple-test + type: exitcode-stdio-1.0 + hs-source-dirs: test + main-is: Spec.hs + default-language: Haskell2010 + build-depends: + base >=4.14 && <5 + , duckdb-simple + , bytestring + , text + , tasty >=1.4 && <1.6 + , tasty-hunit >=0.10 && <0.12 + , time diff --git a/src/Database/DuckDB/Simple.hs b/src/Database/DuckDB/Simple.hs new file mode 100644 --- /dev/null +++ b/src/Database/DuckDB/Simple.hs @@ -0,0 +1,1226 @@ +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE TupleSections #-} + +{- | +Module : Database.DuckDB.Simple +Description : High-level DuckDB API in the duckdb-simple style. + +The API mirrors the ergonomics of @sqlite-simple@ while being backed by the +DuckDB C API. It supports connection management, parameter binding, execution, +and typed result decoding. See 'README.md' for usage examples. +-} +module Database.DuckDB.Simple ( + -- * Connections + Connection, + open, + close, + withConnection, + + -- * Queries and statements + Query (..), + Statement, + openStatement, + closeStatement, + withStatement, + clearStatementBindings, + namedParameterIndex, + columnCount, + columnName, + executeStatement, + execute, + executeMany, + execute_, + bind, + bindNamed, + executeNamed, + queryNamed, + fold, + fold_, + foldNamed, + withTransaction, + query, + queryWith, + query_, + queryWith_, + nextRow, + nextRowWith, + + -- * Errors and conversions + SQLError (..), + FormatError (..), + ResultError (..), + FieldParser, + FromField (..), + FromRow (..), + RowParser, + field, + fieldWith, + numFieldsRemaining, + -- Re-export parameter helper types. + ToField (..), + ToRow (..), + FieldBinding, + NamedParam (..), + Null (..), + Only (..), + (:.) (..), + + -- * User-defined scalar functions + Function, + createFunction, + deleteFunction, +) where + +import Control.Monad (forM, foldM, join, void, when, zipWithM, zipWithM_) +import Data.IORef (IORef, atomicModifyIORef', mkWeakIORef, newIORef, readIORef, writeIORef) +import Control.Exception (SomeException, bracket, finally, mask, onException, throwIO, try) +import qualified Data.ByteString as BS +import Data.Int (Int16, Int32, Int64, Int8) +import Data.Maybe (isJust, isNothing) +import Data.Text (Text) +import qualified Data.Text as Text +import qualified Data.Text.Foreign as TextForeign +import qualified Data.Text.Encoding as TextEncoding +import Data.Time.Calendar (Day, fromGregorian) +import Data.Time.LocalTime (LocalTime (..), TimeOfDay (..)) +import Data.Word (Word16, Word32, Word64, Word8) +import Data.Ratio ((%)) +import Database.DuckDB.FFI +import Database.DuckDB.Simple.FromField ( + Field (..), + FieldValue (..), + FieldParser, + FromField (..), + ResultError (..), + ) +import Database.DuckDB.Simple.FromRow ( + FromRow (..), + RowParser, + field, + fieldWith, + numFieldsRemaining, + parseRow, + rowErrorsToSqlError + ) +import Database.DuckDB.Simple.Function (Function, createFunction, deleteFunction) +import Database.DuckDB.Simple.Internal ( + Connection (..), + ConnectionState (..), + Query (..), + SQLError (..), + Statement (..), + StatementState (..), + StatementStream (..), + StatementStreamChunk (..), + StatementStreamChunkVector (..), + StatementStreamColumn (..), + StatementStreamState (..), + withConnectionHandle, + withQueryCString, + withStatementHandle + ) +import Database.DuckDB.Simple.ToField (FieldBinding, NamedParam (..), ToField (..), bindFieldBinding, renderFieldBinding) +import Database.DuckDB.Simple.ToRow (ToRow (..)) +import Database.DuckDB.Simple.Types (FormatError (..), Null (..), Only (..), (:.) (..)) +import Foreign.C.String (CString, peekCString, withCString) +import Foreign.C.Types (CBool (..), CDouble (..)) +import Foreign.Marshal.Alloc (alloca, free, malloc) +import Foreign.Ptr (Ptr, castPtr, nullPtr, plusPtr) +import Foreign.Storable (peek, peekElemOff, poke) +import Database.DuckDB.Simple.Ok (Ok(..)) + +-- | Open a DuckDB database located at the supplied path. +open :: FilePath -> IO Connection +open path = + mask \restore -> do + db <- restore (openDatabase path) + conn <- + restore (connectDatabase db) + `onException` closeDatabaseHandle db + createConnection db conn + `onException` do + closeConnectionHandle conn + closeDatabaseHandle db + +-- | Close a connection. The operation is idempotent. +close :: Connection -> IO () +close Connection{connectionState} = + void $ + atomicModifyIORef' connectionState \case + ConnectionClosed -> (ConnectionClosed, pure ()) + openState@(ConnectionOpen{}) -> + (ConnectionClosed, closeHandles openState) + +-- | Run an action with a freshly opened connection, closing it afterwards. +withConnection :: FilePath -> (Connection -> IO a) -> IO a +withConnection path = bracket (open path) close + +-- | Prepare a SQL statement for execution. +openStatement :: Connection -> Query -> IO Statement +openStatement conn queryText = + mask \restore -> do + handle <- + restore $ + withConnectionHandle conn \connPtr -> + withQueryCString queryText \sql -> + alloca \stmtPtr -> do + rc <- c_duckdb_prepare connPtr sql stmtPtr + stmt <- peek stmtPtr + if rc == DuckDBSuccess + then pure stmt + else do + errMsg <- fetchPrepareError stmt + c_duckdb_destroy_prepare stmtPtr + throwIO $ mkPrepareError queryText errMsg + createStatement conn handle queryText + `onException` destroyPrepared handle + +-- | Finalise a prepared statement. The operation is idempotent. +closeStatement :: Statement -> IO () +closeStatement stmt@Statement{statementState} = do + resetStatementStream stmt + void $ + atomicModifyIORef' statementState \case + StatementClosed -> (StatementClosed, pure ()) + StatementOpen{statementHandle} -> + (StatementClosed, destroyPrepared statementHandle) + +-- | Run an action with a prepared statement, closing it afterwards. +withStatement :: Connection -> Query -> (Statement -> IO a) -> IO a +withStatement conn sql = bracket (openStatement conn sql) closeStatement + +-- | Bind positional parameters to a prepared statement, replacing any previous bindings. +bind :: Statement -> [FieldBinding] -> IO () +bind stmt fields = do + resetStatementStream stmt + withStatementHandle stmt \handle -> do + let actual = length fields + expected <- fmap fromIntegral (c_duckdb_nparams handle) + when (actual /= expected) $ + throwFormatErrorBindings stmt (parameterCountMessage expected actual) fields + parameterNames <- fetchParameterNames handle expected + when (any isJust parameterNames) $ + throwFormatErrorBindings stmt (Text.pack "duckdb-simple: statement defines named parameters; use executeNamed or bindNamed") fields + clearStatementBindings stmt + zipWithM_ apply [1 ..] fields + where + parameterCountMessage expected actual = + Text.pack $ + "duckdb-simple: SQL query contains " + <> show expected + <> " parameter(s), but " + <> show actual + <> " argument(s) were supplied" + + apply :: Int -> FieldBinding -> IO () + apply idx = bindFieldBinding stmt (fromIntegral idx :: DuckDBIdx) + +-- | Bind named parameters to a prepared statement, preserving any positional bindings. +bindNamed :: Statement -> [NamedParam] -> IO () +bindNamed stmt params = + let bindings = fmap (\(name := value) -> (name, toField value)) params + parameterCountMessage expected actual = + Text.pack $ + "duckdb-simple: SQL query contains " + <> show expected + <> " named parameter(s), but " + <> show actual + <> " argument(s) were supplied" + unknownNameMessage name = + Text.concat + [ Text.pack "duckdb-simple: unknown named parameter " + , name + ] + apply (name, binding) = do + mIdx <- namedParameterIndex stmt name + case mIdx of + Nothing -> + throwFormatErrorNamed stmt (unknownNameMessage name) bindings + Just idx -> bindFieldBinding stmt (fromIntegral idx :: DuckDBIdx) binding + in do + resetStatementStream stmt + withStatementHandle stmt \handle -> do + let actual = length bindings + expected <- fmap fromIntegral (c_duckdb_nparams handle) + when (actual /= expected) $ + throwFormatErrorNamed stmt (parameterCountMessage expected actual) bindings + parameterNames <- fetchParameterNames handle expected + when (all isNothing parameterNames && expected > 0) $ + throwFormatErrorNamed stmt (Text.pack "duckdb-simple: statement does not define named parameters; use positional bindings or adjust the SQL") bindings + clearStatementBindings stmt + mapM_ apply bindings + +fetchParameterNames :: DuckDBPreparedStatement -> Int -> IO [Maybe Text] +fetchParameterNames handle count = + forM [1 .. count] \idx -> do + namePtr <- c_duckdb_parameter_name handle (fromIntegral idx) + if namePtr == nullPtr + then pure Nothing + else do + name <- Text.pack <$> peekCString namePtr + c_duckdb_free (castPtr namePtr) + let normalized = normalizeName name + if normalized == Text.pack (show idx) + then pure Nothing + else pure (Just name) + +-- | Remove all parameter bindings associated with a prepared statement. +clearStatementBindings :: Statement -> IO () +clearStatementBindings stmt = + withStatementHandle stmt \handle -> do + rc <- c_duckdb_clear_bindings handle + when (rc /= DuckDBSuccess) $ do + err <- fetchPrepareError handle + throwIO $ mkPrepareError (statementQuery stmt) err + +-- | Look up the 1-based index of a named placeholder. +namedParameterIndex :: Statement -> Text -> IO (Maybe Int) +namedParameterIndex stmt name = + withStatementHandle stmt \handle -> + let normalized = normalizeName name + in TextForeign.withCString normalized \cName -> + alloca \idxPtr -> do + rc <- c_duckdb_bind_parameter_index handle idxPtr cName + if rc == DuckDBSuccess + then do + idx <- peek idxPtr + if idx == 0 + then pure Nothing + else pure (Just (fromIntegral idx)) + else pure Nothing + +-- | Retrieve the number of columns produced by the supplied prepared statement. +columnCount :: Statement -> IO Int +columnCount stmt = + withStatementHandle stmt $ \handle -> + fmap fromIntegral (c_duckdb_prepared_statement_column_count handle) + +-- | Look up the zero-based column name exposed by a prepared statement result. +columnName :: Statement -> Int -> IO Text +columnName stmt columnIndex + | columnIndex < 0 = throwIO (columnIndexError stmt columnIndex Nothing) + | otherwise = + withStatementHandle stmt \handle -> do + total <- fmap fromIntegral (c_duckdb_prepared_statement_column_count handle) + when (columnIndex >= total) $ + throwIO (columnIndexError stmt columnIndex (Just total)) + namePtr <- c_duckdb_prepared_statement_column_name handle (fromIntegral columnIndex) + if namePtr == nullPtr + then throwIO (columnNameUnavailableError stmt columnIndex) + else do + name <- Text.pack <$> peekCString namePtr + c_duckdb_free (castPtr namePtr) + pure name + +-- | Execute a prepared statement and return the number of affected rows. +-- Resets any active result stream before running and raises a 'SqlError' +-- if DuckDB reports a failure. +executeStatement :: Statement -> IO Int +executeStatement stmt = + withStatementHandle stmt \handle -> do + resetStatementStream stmt + alloca \resPtr -> do + rc <- c_duckdb_execute_prepared handle resPtr + if rc == DuckDBSuccess + then do + changed <- resultRowsChanged resPtr + c_duckdb_destroy_result resPtr + pure changed + else do + (errMsg, _) <- fetchResultError resPtr + c_duckdb_destroy_result resPtr + throwIO $ mkPrepareError (statementQuery stmt) errMsg + +-- | Execute a query with positional parameters and return the affected row count. +execute :: (ToRow q) => Connection -> Query -> q -> IO Int +execute conn queryText params = + withStatement conn queryText \stmt -> do + bind stmt (toRow params) + executeStatement stmt + +-- | Execute the same query multiple times with different parameter sets. +executeMany :: (ToRow q) => Connection -> Query -> [q] -> IO Int +executeMany conn queryText rows = + withStatement conn queryText \stmt -> do + sum <$> mapM (\row -> bind stmt (toRow row) >> executeStatement stmt) rows + +-- | Execute an ad-hoc query without parameters and return the affected row count. +execute_ :: Connection -> Query -> IO Int +execute_ conn queryText = + withConnectionHandle conn \connPtr -> + withQueryCString queryText \sql -> + alloca \resPtr -> do + rc <- c_duckdb_query connPtr sql resPtr + if rc == DuckDBSuccess + then do + changed <- resultRowsChanged resPtr + c_duckdb_destroy_result resPtr + pure changed + else do + (errMsg, errType) <- fetchResultError resPtr + c_duckdb_destroy_result resPtr + throwIO $ mkExecuteError queryText errMsg errType + +-- | Execute a query that uses named parameters. +executeNamed :: Connection -> Query -> [NamedParam] -> IO Int +executeNamed conn queryText params = + withStatement conn queryText \stmt -> do + bindNamed stmt params + executeStatement stmt + +-- | Run a parameterised query and decode every resulting row eagerly. +query :: (ToRow q, FromRow r) => Connection -> Query -> q -> IO [r] +query = queryWith fromRow + +-- | Run a parameterised query with a custom row parser. +queryWith :: (ToRow q) => RowParser r -> Connection -> Query -> q -> IO [r] +queryWith parser conn queryText params = + withStatement conn queryText \stmt -> do + bind stmt (toRow params) + withStatementHandle stmt \handle -> + alloca \resPtr -> do + rc <- c_duckdb_execute_prepared handle resPtr + if rc == DuckDBSuccess + then do + rows <- collectRows resPtr + c_duckdb_destroy_result resPtr + convertRowsWith parser queryText rows + else do + (errMsg, errType) <- fetchResultError resPtr + c_duckdb_destroy_result resPtr + throwIO $ mkExecuteError queryText errMsg errType + +-- | Run a query that uses named parameters and decode all rows eagerly. +queryNamed :: (FromRow r) => Connection -> Query -> [NamedParam] -> IO [r] +queryNamed conn queryText params = + withStatement conn queryText \stmt -> do + bindNamed stmt params + withStatementHandle stmt \handle -> + alloca \resPtr -> do + rc <- c_duckdb_execute_prepared handle resPtr + if rc == DuckDBSuccess + then do + rows <- collectRows resPtr + c_duckdb_destroy_result resPtr + convertRows queryText rows + else do + (errMsg, errType) <- fetchResultError resPtr + c_duckdb_destroy_result resPtr + throwIO $ mkExecuteError queryText errMsg errType + +-- | Run a query without supplying parameters and decode all rows eagerly. +query_ :: (FromRow r) => Connection -> Query -> IO [r] +query_ = queryWith_ fromRow + +-- | Run a query without parameters using a custom row parser. +queryWith_ :: RowParser r -> Connection -> Query -> IO [r] +queryWith_ parser conn queryText = + withConnectionHandle conn \connPtr -> + withQueryCString queryText \sql -> + alloca \resPtr -> do + rc <- c_duckdb_query connPtr sql resPtr + if rc == DuckDBSuccess + then do + rows <- collectRows resPtr + c_duckdb_destroy_result resPtr + convertRowsWith parser queryText rows + else do + (errMsg, errType) <- fetchResultError resPtr + c_duckdb_destroy_result resPtr + throwIO $ mkExecuteError queryText errMsg errType + +-- Streaming folds ----------------------------------------------------------- + +-- | Stream a parameterised query through an accumulator without loading all rows. +-- Bind the supplied parameters, start a streaming result, and apply the step +-- function row by row to produce a final accumulator value. +fold :: (FromRow row, ToRow params) => Connection -> Query -> params -> a -> (a -> row -> IO a) -> IO a +fold conn queryText params initial step = + withStatement conn queryText \stmt -> do + resetStatementStream stmt + bind stmt (toRow params) + foldStatementWith fromRow stmt initial step + +-- | Stream a parameterless query through an accumulator without loading all rows. +fold_ :: (FromRow row) => Connection -> Query -> a -> (a -> row -> IO a) -> IO a +fold_ conn queryText initial step = + withStatement conn queryText \stmt -> do + resetStatementStream stmt + foldStatementWith fromRow stmt initial step + +-- | Stream a query that uses named parameters through an accumulator. +foldNamed :: (FromRow row) => Connection -> Query -> [NamedParam] -> a -> (a -> row -> IO a) -> IO a +foldNamed conn queryText params initial step = + withStatement conn queryText \stmt -> do + resetStatementStream stmt + bindNamed stmt params + foldStatementWith fromRow stmt initial step + +foldStatementWith :: RowParser row -> Statement -> a -> (a -> row -> IO a) -> IO a +foldStatementWith parser stmt initial step = + let loop acc = do + nextVal <- nextRowWith parser stmt + case nextVal of + Nothing -> pure acc + Just row -> do + acc' <- step acc row + acc' `seq` loop acc' + in loop initial `finally` resetStatementStream stmt + +-- | Fetch the next row from a streaming statement, stopping when no rows remain. +nextRow :: (FromRow r) => Statement -> IO (Maybe r) +nextRow = nextRowWith fromRow + +-- | Fetch the next row using a custom parser, returning 'Nothing' once exhausted. +nextRowWith :: RowParser r -> Statement -> IO (Maybe r) +nextRowWith parser stmt@Statement{statementStream} = + mask \restore -> do + state <- readIORef statementStream + case state of + StatementStreamIdle -> do + newStream <- restore (startStatementStream stmt) + case newStream of + Nothing -> pure Nothing + Just stream -> restore (consumeStream statementStream parser stmt stream) + StatementStreamActive stream -> + restore (consumeStream statementStream parser stmt stream) + +resetStatementStream :: Statement -> IO () +resetStatementStream Statement{statementStream} = + cleanupStatementStreamRef statementStream + +consumeStream :: IORef StatementStreamState -> RowParser r -> Statement -> StatementStream -> IO (Maybe r) +consumeStream streamRef parser stmt stream = do + result <- + ( try (streamNextRow (statementQuery stmt) stream) + :: IO (Either SomeException (Maybe [Field], StatementStream)) + ) + case result of + Left err -> do + finalizeStream stream + writeIORef streamRef StatementStreamIdle + throwIO err + Right (maybeFields, updatedStream) -> + case maybeFields of + Nothing -> do + finalizeStream updatedStream + writeIORef streamRef StatementStreamIdle + pure Nothing + Just fields -> + case parseRow parser fields of + Errors rowErr -> do + finalizeStream updatedStream + writeIORef streamRef StatementStreamIdle + throwIO $ rowErrorsToSqlError (statementQuery stmt) rowErr + Ok value -> do + writeIORef streamRef (StatementStreamActive updatedStream) + pure (Just value) + +startStatementStream :: Statement -> IO (Maybe StatementStream) +startStatementStream stmt = + withStatementHandle stmt \handle -> do + columns <- collectStreamColumns handle + resultPtr <- malloc + rc <- c_duckdb_execute_prepared_streaming handle resultPtr + if rc /= DuckDBSuccess + then do + (errMsg, errType) <- fetchResultError resultPtr + c_duckdb_destroy_result resultPtr + free resultPtr + throwIO $ mkExecuteError (statementQuery stmt) errMsg errType + else do + resultType <- c_duckdb_result_return_type resultPtr + if resultType /= DuckDBResultTypeQueryResult + then do + c_duckdb_destroy_result resultPtr + free resultPtr + pure Nothing + else pure (Just (StatementStream resultPtr columns Nothing)) + +collectStreamColumns :: DuckDBPreparedStatement -> IO [StatementStreamColumn] +collectStreamColumns handle = do + rawCount <- c_duckdb_prepared_statement_column_count handle + let cc = fromIntegral rawCount :: Int + forM [0 .. cc - 1] \idx -> do + namePtr <- c_duckdb_prepared_statement_column_name handle (fromIntegral idx) + name <- + if namePtr == nullPtr + then pure (Text.pack ("column" <> show idx)) + else Text.pack <$> peekCString namePtr + dtype <- c_duckdb_prepared_statement_column_type handle (fromIntegral idx) + pure + StatementStreamColumn + { statementStreamColumnIndex = idx + , statementStreamColumnName = name + , statementStreamColumnType = dtype + } + +streamNextRow :: Query -> StatementStream -> IO (Maybe [Field], StatementStream) +streamNextRow queryText stream@StatementStream{statementStreamChunk = Nothing} = do + refreshed <- fetchChunk stream + case statementStreamChunk refreshed of + Nothing -> pure (Nothing, refreshed) + Just chunk -> emitRow queryText refreshed chunk +streamNextRow queryText stream@StatementStream{statementStreamChunk = Just chunk} = + emitRow queryText stream chunk + +fetchChunk :: StatementStream -> IO StatementStream +fetchChunk stream@StatementStream{statementStreamResult} = do + chunk <- c_duckdb_stream_fetch_chunk statementStreamResult + if chunk == nullPtr + then pure stream + else do + rawSize <- c_duckdb_data_chunk_get_size chunk + let rowCount = fromIntegral rawSize :: Int + if rowCount <= 0 + then do + destroyDataChunk chunk + fetchChunk stream + else do + vectors <- prepareChunkVectors chunk (statementStreamColumns stream) + let chunkState = + StatementStreamChunk + { statementStreamChunkPtr = chunk + , statementStreamChunkSize = rowCount + , statementStreamChunkIndex = 0 + , statementStreamChunkVectors = vectors + } + 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 + } + +emitRow :: Query -> StatementStream -> StatementStreamChunk -> IO (Maybe [Field], StatementStream) +emitRow queryText stream chunk@StatementStreamChunk{statementStreamChunkIndex, statementStreamChunkSize} = do + fields <- + buildRow + queryText + (statementStreamColumns stream) + (statementStreamChunkVectors chunk) + statementStreamChunkIndex + let nextIndex = statementStreamChunkIndex + 1 + if nextIndex < statementStreamChunkSize + then + let updatedChunk = chunk{statementStreamChunkIndex = nextIndex} + in pure (Just fields, stream{statementStreamChunk = Just updatedChunk}) + else do + destroyDataChunk (statementStreamChunkPtr chunk) + pure (Just fields, stream{statementStreamChunk = Nothing}) + +buildRow :: Query -> [StatementStreamColumn] -> [StatementStreamChunkVector] -> Int -> IO [Field] +buildRow queryText columns vectors rowIdx = + zipWithM (buildField queryText rowIdx) columns vectors + +buildField :: Query -> Int -> StatementStreamColumn -> StatementStreamChunkVector -> IO Field +buildField queryText rowIdx column StatementStreamChunkVector{statementStreamChunkVectorData, statementStreamChunkVectorValidity} = do + let duckIdx = fromIntegral rowIdx + dtype = statementStreamColumnType column + valid <- chunkIsRowValid statementStreamChunkVectorValidity duckIdx + value <- + if not valid + then pure FieldNull + else case dtype of + DuckDBTypeBoolean -> do + raw <- peekElemOff (castPtr statementStreamChunkVectorData :: Ptr Word8) rowIdx + pure (FieldBool (raw /= 0)) + DuckDBTypeTinyInt -> do + raw <- peekElemOff (castPtr statementStreamChunkVectorData :: Ptr Int8) rowIdx + pure (FieldInt8 (fromIntegral raw)) + DuckDBTypeSmallInt -> do + raw <- peekElemOff (castPtr statementStreamChunkVectorData :: Ptr Int16) rowIdx + pure (FieldInt16 (fromIntegral raw)) + DuckDBTypeInteger -> do + raw <- peekElemOff (castPtr statementStreamChunkVectorData :: Ptr Int32) rowIdx + pure (FieldInt32 (fromIntegral raw)) + DuckDBTypeBigInt -> do + raw <- peekElemOff (castPtr statementStreamChunkVectorData :: Ptr Int64) rowIdx + pure (FieldInt64 raw) + DuckDBTypeUTinyInt -> do + raw <- peekElemOff (castPtr statementStreamChunkVectorData :: Ptr Word8) rowIdx + pure (FieldWord8 (fromIntegral raw)) + DuckDBTypeUSmallInt -> do + raw <- peekElemOff (castPtr statementStreamChunkVectorData :: Ptr Word16) rowIdx + pure (FieldWord16 (fromIntegral raw)) + DuckDBTypeUInteger -> do + raw <- peekElemOff (castPtr statementStreamChunkVectorData :: Ptr Word32) rowIdx + pure (FieldWord32 (fromIntegral raw)) + DuckDBTypeUBigInt -> do + raw <- peekElemOff (castPtr statementStreamChunkVectorData :: Ptr Word64) rowIdx + pure (FieldWord64 (fromIntegral raw)) + DuckDBTypeBlob -> do + blob <- chunkDecodeBlob statementStreamChunkVectorData duckIdx + pure (FieldBlob blob) + DuckDBTypeDate -> do + raw <- peekElemOff (castPtr statementStreamChunkVectorData :: Ptr Int32) rowIdx + day <- decodeDuckDBDate (DuckDBDate raw) + pure (FieldDate day) + DuckDBTypeTime -> do + raw <- peekElemOff (castPtr statementStreamChunkVectorData :: Ptr DuckDBTime) rowIdx + tod <- decodeDuckDBTime raw + pure (FieldTime tod) + DuckDBTypeTimestamp -> do + raw <- peekElemOff (castPtr statementStreamChunkVectorData :: Ptr DuckDBTimestamp) rowIdx + ts <- decodeDuckDBTimestamp raw + pure (FieldTimestamp ts) + DuckDBTypeFloat -> do + raw <- peekElemOff (castPtr statementStreamChunkVectorData :: Ptr Float) rowIdx + pure (FieldDouble (realToFrac raw)) + DuckDBTypeDouble -> do + raw <- peekElemOff (castPtr statementStreamChunkVectorData :: Ptr Double) rowIdx + pure (FieldDouble raw) + DuckDBTypeVarchar -> FieldText <$> chunkDecodeText statementStreamChunkVectorData duckIdx + DuckDBTypeHugeInt -> error "duckdb-simple: HUGEINT is not supported" + DuckDBTypeUHugeInt -> error "duckdb-simple: HUGEINT is not supported" + _ -> + throwIO (streamingUnsupportedTypeError queryText column) + pure + Field + { fieldName = statementStreamColumnName column + , fieldIndex = statementStreamColumnIndex column + , fieldValue = value + } + +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) + +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 + +cleanupStatementStreamRef :: IORef StatementStreamState -> IO () +cleanupStatementStreamRef ref = do + state <- atomicModifyIORef' ref (StatementStreamIdle,) + finalizeStreamState state + +finalizeStreamState :: StatementStreamState -> IO () +finalizeStreamState = \case + StatementStreamIdle -> pure () + StatementStreamActive stream -> finalizeStream stream + +finalizeStream :: StatementStream -> IO () +finalizeStream StatementStream{statementStreamResult, statementStreamChunk} = do + maybe (pure ()) finalizeChunk statementStreamChunk + c_duckdb_destroy_result statementStreamResult + free statementStreamResult + +finalizeChunk :: StatementStreamChunk -> IO () +finalizeChunk StatementStreamChunk{statementStreamChunkPtr} = + destroyDataChunk statementStreamChunkPtr + +destroyDataChunk :: DuckDBDataChunk -> IO () +destroyDataChunk chunk = + alloca \ptr -> do + poke ptr chunk + c_duckdb_destroy_data_chunk ptr + +streamingUnsupportedTypeError :: Query -> StatementStreamColumn -> SQLError +streamingUnsupportedTypeError queryText StatementStreamColumn{statementStreamColumnName, statementStreamColumnType} = + SQLError + { sqlErrorMessage = + Text.concat + [ Text.pack "duckdb-simple: streaming does not yet support column " + , statementStreamColumnName + , Text.pack " with DuckDB type " + , Text.pack (show statementStreamColumnType) + ] + , sqlErrorType = Nothing + , sqlErrorQuery = Just queryText + } + +-- | Run an action inside a transaction. +withTransaction :: Connection -> IO a -> IO a +withTransaction conn action = + mask \restore -> do + void (execute_ conn begin) + let rollbackAction = void (execute_ conn rollback) + result <- restore action `onException` rollbackAction + void (execute_ conn commit) + pure result + + where begin = Query (Text.pack "BEGIN TRANSACTION") + commit = Query (Text.pack "COMMIT") + rollback = Query (Text.pack "ROLLBACK") + + +-- Internal helpers ----------------------------------------------------------- + +createConnection :: DuckDBDatabase -> DuckDBConnection -> IO Connection +createConnection db conn = do + ref <- newIORef (ConnectionOpen db conn) + _ <- + mkWeakIORef ref $ + void $ + atomicModifyIORef' ref \case + ConnectionClosed -> (ConnectionClosed, pure ()) + openState@(ConnectionOpen{}) -> + (ConnectionClosed, closeHandles openState) + pure Connection{connectionState = ref} + +createStatement :: Connection -> DuckDBPreparedStatement -> Query -> IO Statement +createStatement parent handle queryText = do + ref <- newIORef (StatementOpen handle) + streamRef <- newIORef StatementStreamIdle + _ <- + mkWeakIORef ref $ + do join $ + atomicModifyIORef' ref $ \case + StatementClosed -> (StatementClosed, pure ()) + StatementOpen{statementHandle} -> + ( StatementClosed + , do + cleanupStatementStreamRef streamRef + destroyPrepared statementHandle + ) + pure + Statement + { statementState = ref + , statementConnection = parent + , statementQuery = queryText + , statementStream = streamRef + } + +openDatabase :: FilePath -> IO DuckDBDatabase +openDatabase path = + alloca \dbPtr -> + alloca \errPtr -> do + poke errPtr nullPtr + withCString path \cPath -> do + rc <- c_duckdb_open_ext cPath dbPtr nullPtr errPtr + if rc == DuckDBSuccess + then do + db <- peek dbPtr + maybeFreeErr errPtr + pure db + else do + errMsg <- peekError errPtr + maybeFreeErr errPtr + throwIO $ mkOpenError errMsg + +connectDatabase :: DuckDBDatabase -> IO DuckDBConnection +connectDatabase db = + alloca \connPtr -> do + rc <- c_duckdb_connect db connPtr + if rc == DuckDBSuccess + then peek connPtr + else throwIO mkConnectError + +closeHandles :: ConnectionState -> IO () +closeHandles ConnectionClosed = pure () +closeHandles ConnectionOpen{connectionDatabase, connectionHandle} = do + closeConnectionHandle connectionHandle + closeDatabaseHandle connectionDatabase + +closeConnectionHandle :: DuckDBConnection -> IO () +closeConnectionHandle conn = + alloca \ptr -> poke ptr conn >> c_duckdb_disconnect ptr + +closeDatabaseHandle :: DuckDBDatabase -> IO () +closeDatabaseHandle db = + alloca \ptr -> poke ptr db >> c_duckdb_close ptr + +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 Text.pack <$> peekCString msgPtr + +fetchResultError :: Ptr DuckDBResult -> IO (Text, Maybe DuckDBErrorType) +fetchResultError resultPtr = do + msgPtr <- c_duckdb_result_error resultPtr + msg <- + if msgPtr == nullPtr + then pure (Text.pack "duckdb-simple: query failed") + else Text.pack <$> peekCString msgPtr + errType <- c_duckdb_result_error_type resultPtr + let classified = + if errType == DuckDBErrorInvalid + then Nothing + else Just errType + pure (msg, classified) + +mkOpenError :: Text -> SQLError +mkOpenError msg = + SQLError + { sqlErrorMessage = msg + , sqlErrorType = Nothing + , sqlErrorQuery = Nothing + } + +mkConnectError :: SQLError +mkConnectError = + SQLError + { sqlErrorMessage = Text.pack "duckdb-simple: failed to create connection handle" + , sqlErrorType = Nothing + , sqlErrorQuery = Nothing + } + +mkPrepareError :: Query -> Text -> SQLError +mkPrepareError queryText msg = + SQLError + { sqlErrorMessage = msg + , sqlErrorType = Nothing + , sqlErrorQuery = Just queryText + } + +mkExecuteError :: Query -> Text -> Maybe DuckDBErrorType -> SQLError +mkExecuteError queryText msg errType = + SQLError + { sqlErrorMessage = msg + , sqlErrorType = errType + , sqlErrorQuery = Just queryText + } + +throwFormatError :: Statement -> Text -> [String] -> IO a +throwFormatError Statement{statementQuery} message params = + throwIO + FormatError + { formatErrorMessage = message + , formatErrorQuery = statementQuery + , formatErrorParams = params + } + +throwFormatErrorBindings :: Statement -> Text -> [FieldBinding] -> IO a +throwFormatErrorBindings stmt message bindings = + throwFormatError stmt message (map renderFieldBinding bindings) + +throwFormatErrorNamed :: Statement -> Text -> [(Text, FieldBinding)] -> IO a +throwFormatErrorNamed stmt message bindings = + throwFormatError stmt message (map renderNamed bindings) + where + renderNamed (name, binding) = + Text.unpack name <> " := " <> renderFieldBinding binding + +columnIndexError :: Statement -> Int -> Maybe Int -> SQLError +columnIndexError stmt idx total = + let base = + Text.concat + [ Text.pack "duckdb-simple: column index " + , Text.pack (show idx) + , Text.pack " out of bounds" + ] + message = + case total of + Nothing -> base + Just count -> + Text.concat + [ base + , Text.pack " (column count: " + , Text.pack (show count) + , Text.pack ")" + ] + in SQLError + { sqlErrorMessage = message + , sqlErrorType = Nothing + , sqlErrorQuery = Just (statementQuery stmt) + } + +columnNameUnavailableError :: Statement -> Int -> SQLError +columnNameUnavailableError stmt idx = + SQLError + { sqlErrorMessage = + Text.concat + [ Text.pack "duckdb-simple: column name unavailable for index " + , Text.pack (show idx) + ] + , sqlErrorType = Nothing + , sqlErrorQuery = Just (statementQuery stmt) + } + + +normalizeName :: Text -> Text +normalizeName name = + case Text.uncons name of + Just (prefix, rest) + | prefix == ':' || prefix == '$' || prefix == '@' -> rest + _ -> name + +resultRowsChanged :: Ptr DuckDBResult -> IO Int +resultRowsChanged resPtr = fromIntegral <$> c_duckdb_rows_changed resPtr + +convertRows :: (FromRow r) => Query -> [[Field]] -> IO [r] +convertRows = convertRowsWith fromRow + +convertRowsWith :: RowParser r -> Query -> [[Field]] -> IO [r] +convertRowsWith parser queryText rows = + case traverse (parseRow parser) rows of + Errors err -> throwIO (rowErrorsToSqlError queryText err) + Ok ok -> pure ok + +collectRows :: Ptr DuckDBResult -> IO [[Field]] +collectRows resPtr = do + columns <- collectResultColumns resPtr + chunkCount <- (fromIntegral <$> c_duckdb_result_chunk_count resPtr) :: IO Int + if null columns || chunkCount <= 0 + then pure [] + else do + (_, chunks) <- + foldM + (\(rowBase, acc) chunkIdx -> do + (nextBase, chunkRows) <- collectChunkRows resPtr columns chunkIdx rowBase + pure (nextBase, chunkRows : acc) + ) + (0, []) + [0 .. chunkCount - 1] + pure (concat (reverse chunks)) + +collectResultColumns :: Ptr DuckDBResult -> IO [StatementStreamColumn] +collectResultColumns resPtr = do + rawCount <- c_duckdb_column_count resPtr + let cc = fromIntegral rawCount :: Int + forM [0 .. cc - 1] \columnIndex -> do + namePtr <- c_duckdb_column_name resPtr (fromIntegral columnIndex) + name <- + if namePtr == nullPtr + then pure (Text.pack ("column" <> show columnIndex)) + else Text.pack <$> peekCString namePtr + dtype <- c_duckdb_column_type resPtr (fromIntegral columnIndex) + pure + StatementStreamColumn + { statementStreamColumnIndex = columnIndex + , statementStreamColumnName = name + , statementStreamColumnType = dtype + } + +collectChunkRows :: Ptr DuckDBResult -> [StatementStreamColumn] -> Int -> Int -> IO (Int, [[Field]]) +collectChunkRows resPtr columns chunkIndex rowBase = do + chunk <- c_duckdb_result_get_chunk resPtr (fromIntegral chunkIndex) + if chunk == nullPtr + then pure (rowBase, []) + else + finally + (do + rawSize <- c_duckdb_data_chunk_get_size chunk + let rowCount = fromIntegral rawSize :: Int + when (rowCount > 0 && not (null columns)) $ + void $ + fetchFieldValue + resPtr + (statementStreamColumnType (head columns)) + (statementStreamColumnIndex (head columns)) + rowBase + if rowCount <= 0 + then pure (rowBase, []) + else do + vectors <- prepareChunkVectors chunk columns + rows <- mapM (buildMaterializedRow columns vectors) [0 .. rowCount - 1] + pure (rowBase + rowCount, rows) + ) + (destroyDataChunk chunk) + +buildMaterializedRow :: [StatementStreamColumn] -> [StatementStreamChunkVector] -> Int -> IO [Field] +buildMaterializedRow columns vectors rowIdx = + zipWithM (buildMaterializedField rowIdx) columns vectors + +buildMaterializedField :: Int -> StatementStreamColumn -> StatementStreamChunkVector -> IO Field +buildMaterializedField rowIdx column StatementStreamChunkVector{statementStreamChunkVectorData, statementStreamChunkVectorValidity} = do + value <- + materializedValueFromPointers + (statementStreamColumnType column) + statementStreamChunkVectorData + statementStreamChunkVectorValidity + rowIdx + pure + Field + { fieldName = statementStreamColumnName column + , fieldIndex = statementStreamColumnIndex column + , fieldValue = value + } + +fetchFieldValue :: Ptr DuckDBResult -> DuckDBType -> Int -> Int -> IO FieldValue +fetchFieldValue resPtr dtype columnIndex rowIndex = + withChunkForRow resPtr rowIndex \chunk localRow -> do + vector <- c_duckdb_data_chunk_get_vector chunk (fromIntegral columnIndex) + dataPtr <- c_duckdb_vector_get_data vector + validity <- c_duckdb_vector_get_validity vector + materializedValueFromPointers dtype dataPtr validity localRow + +withChunkForRow :: Ptr DuckDBResult -> Int -> (DuckDBDataChunk -> Int -> IO a) -> IO a +withChunkForRow resPtr targetRow action = do + chunkCount <- (fromIntegral <$> c_duckdb_result_chunk_count resPtr) :: IO Int + let seek chunkIdx remaining + | chunkIdx >= chunkCount = + throwIO (userError "duckdb-simple: row index out of bounds while materialising result") + | otherwise = do + chunk <- c_duckdb_result_get_chunk resPtr (fromIntegral chunkIdx) + if chunk == nullPtr + then seek (chunkIdx + 1) remaining + else do + rawSize <- c_duckdb_data_chunk_get_size chunk + let rowCount = fromIntegral rawSize :: Int + if rowCount <= 0 + then do + destroyDataChunk chunk + seek (chunkIdx + 1) remaining + else if remaining < rowCount + then + finally + (action chunk remaining) + (destroyDataChunk chunk) + else do + destroyDataChunk chunk + seek (chunkIdx + 1) (remaining - rowCount) + seek 0 targetRow + +materializedValueFromPointers :: DuckDBType -> Ptr () -> Ptr Word64 -> Int -> IO FieldValue +materializedValueFromPointers dtype dataPtr validity rowIdx = do + let duckIdx = fromIntegral rowIdx :: DuckDBIdx + valid <- chunkIsRowValid validity duckIdx + 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 -> FieldText <$> chunkDecodeText dataPtr duckIdx + 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 + DuckDBTypeTimestamp -> do + raw <- peekElemOff (castPtr dataPtr :: Ptr DuckDBTimestamp) rowIdx + FieldTimestamp <$> decodeDuckDBTimestamp raw + DuckDBTypeDecimal -> do + decimal <- peekElemOff (castPtr dataPtr :: Ptr DuckDBDecimal) rowIdx + alloca \ptr -> do + poke ptr decimal + CDouble value <- c_duckdb_decimal_to_double ptr + pure (FieldDouble (realToFrac value)) + DuckDBTypeHugeInt -> + error "duckdb-simple: HugeInt type not supported" + DuckDBTypeUHugeInt -> + error "duckdb-simple: UHugeInt type not supported" + other -> + error ("duckdb-simple: unsupported DuckDB type in eager result: " <> show other) + +decodeDuckDBDate :: DuckDBDate -> IO Day +decodeDuckDBDate raw = + alloca \ptr -> do + c_duckdb_from_date raw ptr + dateStruct <- peek ptr + pure (dateStructToDay dateStruct) + +decodeDuckDBTime :: DuckDBTime -> IO TimeOfDay +decodeDuckDBTime raw = + alloca \ptr -> do + c_duckdb_from_time raw ptr + timeStruct <- peek ptr + pure (timeStructToTimeOfDay timeStruct) + +decodeDuckDBTimestamp :: DuckDBTimestamp -> IO LocalTime +decodeDuckDBTimestamp raw = + alloca \ptr -> do + c_duckdb_from_timestamp raw ptr + DuckDBTimestampStruct{duckDBTimestampStructDate = dateStruct, duckDBTimestampStructTime = timeStruct} <- peek ptr + pure + LocalTime + { localDay = dateStructToDay dateStruct + , localTimeOfDay = timeStructToTimeOfDay timeStruct + } + +dateStructToDay :: DuckDBDateStruct -> Day +dateStructToDay DuckDBDateStruct{duckDBDateStructYear, duckDBDateStructMonth, duckDBDateStructDay} = + fromGregorian (fromIntegral duckDBDateStructYear) (fromIntegral duckDBDateStructMonth) (fromIntegral duckDBDateStructDay) + +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 + +peekError :: Ptr CString -> IO Text +peekError ptr = do + errPtr <- peek ptr + if errPtr == nullPtr + then pure (Text.pack "duckdb-simple: failed to open database") + else do + message <- peekCString errPtr + pure (Text.pack message) + +maybeFreeErr :: Ptr CString -> IO () +maybeFreeErr ptr = do + errPtr <- peek ptr + when (errPtr /= nullPtr) $ c_duckdb_free (castPtr errPtr) diff --git a/src/Database/DuckDB/Simple/FromField.hs b/src/Database/DuckDB/Simple/FromField.hs new file mode 100644 --- /dev/null +++ b/src/Database/DuckDB/Simple/FromField.hs @@ -0,0 +1,328 @@ +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE ViewPatterns #-} +{-# LANGUAGE StrictData #-} +{-# LANGUAGE OverloadedStrings #-} + +{- | +Module : Database.DuckDB.Simple.FromField +Description : Conversion from DuckDB column values to Haskell types. +-} +module Database.DuckDB.Simple.FromField ( + Field (..), + FieldValue (..), + ResultError (..), + FieldParser, + FromField (..), + left +) where + +import Control.Exception (Exception, SomeException (..)) +import qualified Data.ByteString as BS +import Data.Int (Int16, Int32, Int64, Int8) +import Data.Text (Text) +import qualified Data.Text as Text +import qualified Data.Text.Encoding as TextEncoding +import Data.Time.Calendar (Day) +import Data.Time.Clock (UTCTime (..)) +import Data.Time.LocalTime (LocalTime (..), TimeOfDay (..), localTimeToUTC, utc) +import Data.Word (Word16, Word32, Word64, Word8) +import Database.DuckDB.Simple.Types (Null (..)) +import Database.DuckDB.Simple.Ok +import Data.Data (Typeable, Proxy (..), typeRep) + +-- | Internal representation of a column value. +data FieldValue + = FieldNull + | FieldInt8 Int8 + | FieldInt16 Int16 + | FieldInt32 Int32 + | FieldInt64 Int64 + | FieldWord8 Word8 + | FieldWord16 Word16 + | FieldWord32 Word32 + | FieldWord64 Word64 + | FieldFloat Float + | FieldDouble Double + | FieldText Text + | FieldBool Bool + | FieldBlob BS.ByteString + | FieldDate Day + | FieldTime TimeOfDay + | FieldTimestamp LocalTime + -- TODO: HugeInt and UHugeInt support + + deriving (-- | FieldInteger Integer + -- | FieldNatural Natural + Eq, Show) + +-- | Pattern synonym to make it easier to match on any integral type. +pattern FieldInt :: Int -> FieldValue +pattern FieldInt i <- (fieldValueToInt -> Just i) + where + FieldInt i = FieldInt64 (fromIntegral i) + +fieldValueToInt :: FieldValue -> Maybe Int +fieldValueToInt (FieldInt8 i) = Just (fromIntegral i) +fieldValueToInt (FieldInt16 i) = Just (fromIntegral i) +fieldValueToInt (FieldInt32 i) = Just (fromIntegral i) +fieldValueToInt (FieldInt64 i) = Just (fromIntegral i) +fieldValueToInt _ = Nothing + +-- | Pattern synonym to make it easier to match on any word size +pattern FieldWord :: Word -> FieldValue +pattern FieldWord i <- (fieldValueToWord -> Just i) + where + FieldWord i = FieldWord64 (fromIntegral i) + +fieldValueToWord :: FieldValue -> Maybe Word +fieldValueToWord (FieldWord8 i) = Just (fromIntegral i) +fieldValueToWord (FieldWord16 i) = Just (fromIntegral i) +fieldValueToWord (FieldWord32 i) = Just (fromIntegral i) +fieldValueToWord (FieldWord64 i) = Just (fromIntegral i) +fieldValueToWord _ = Nothing + +-- | Metadata for a single column in a row. +data Field = Field + { fieldName :: Text + , fieldIndex :: Int + , fieldValue :: FieldValue + } + deriving (Eq, Show) + +-- | Exception thrown if conversion from a SQL value to a Haskell +-- value fails. +data ResultError = Incompatible { errSQLType :: Text + , errHaskellType :: Text + , errMessage :: Text } + -- ^ The SQL and Haskell types are not compatible. + | UnexpectedNull { errSQLType :: Text + , errHaskellType :: Text + , errMessage :: Text } + -- ^ A SQL @NULL@ was encountered when the Haskell + -- type did not permit it. + | ConversionFailed { errSQLType :: Text + , errHaskellType :: Text + , errMessage :: Text } + -- ^ The SQL value could not be parsed, or could not + -- be represented as a valid Haskell value, or an + -- unexpected low-level error occurred (e.g. mismatch + -- between metadata and actual data in a row). + deriving (Eq, Show, Typeable) + +instance Exception ResultError + +-- From sqlite-simple: +left :: Exception a => a -> Ok b +left = Errors . (:[]) . SomeException + + +-- | Parser used by 'FromField' instances and utilities such as +-- 'Database.DuckDB.Simple.FromRow.fieldWith'. The supplied 'Field' contains +-- column metadata and an already-decoded 'FieldValue'; callers should return +-- 'Ok' on success or 'Errors' (typically wrapping a 'ResultError') when the +-- conversion fails. +type FieldParser a = Field -> Ok a + +-- | Types that can be constructed from a DuckDB column. +class FromField a where + fromField :: FieldParser a + +instance FromField FieldValue where + fromField Field{fieldValue} = Ok fieldValue + +instance FromField Null where + fromField Field{fieldValue} = + case fieldValue of + FieldNull -> Ok Null + other -> left (Incompatible (fieldValueTypeName other) "Null" "expected NULL") + +instance FromField Bool where + fromField Field{fieldValue} = + case fieldValue of + FieldBool b -> Ok b + FieldInt i -> Ok (i /= 0) + other -> left (Incompatible (fieldValueTypeName other) "Bool" "invalid type for boolean") + +instance FromField Int8 where + fromField Field{fieldValue} = + case fieldValue of + FieldInt i -> Ok (fromIntegral i) + other -> left (Incompatible (fieldValueTypeName other) "Int8" "invalid type for Int8") + +instance FromField Int64 where + fromField Field{fieldValue} = + case fieldValue of + FieldInt i -> Ok (fromIntegral i) + other -> left (Incompatible (fieldValueTypeName other) "Int64" "invalid type for Int64") + +instance FromField Int32 where + fromField field@Field{fieldValue} = + case fieldValue of + FieldInt i -> boundedIntegral field i + other -> left (Incompatible (fieldValueTypeName other) "Int32" "invalid type for integer") + +instance FromField Int16 where + fromField field@Field{fieldValue} = + case fieldValue of + FieldInt i -> boundedIntegral field i + other -> left (Incompatible (fieldValueTypeName other) "Int16" "invalid type for integer") + +instance FromField Int where + fromField field@Field{fieldValue} = + case fieldValue of + FieldInt i -> boundedIntegral field i + other -> left (Incompatible (fieldValueTypeName other) "Int" "invalid type for integer") + +instance FromField Word64 where + fromField Field{fieldValue} = + case fieldValue of + FieldInt i + | i >= 0 -> Ok (fromIntegral i) + | otherwise -> + left $ + ConversionFailed (fieldValueTypeName fieldValue) "Word64" $ + Text.pack "negative value cannot be converted to unsigned integer" + FieldWord w -> Ok (fromIntegral w) + other -> left (Incompatible (fieldValueTypeName other) "Word64" "invalid type for Word64") + +instance FromField Word32 where + fromField field@Field{fieldValue} = + case fieldValue of + FieldInt i + | i >= 0 -> boundedIntegral field i + | otherwise -> + left $ + ConversionFailed (fieldValueTypeName fieldValue) "Word32" $ + Text.pack "negative value cannot be converted to unsigned integer" + FieldWord w -> Ok (fromIntegral w) + other -> left (Incompatible (fieldValueTypeName other) "Word32" "invalid type for Word32") + +instance FromField Word16 where + fromField field@Field{fieldValue} = + case fieldValue of + FieldInt i + | i >= 0 -> boundedIntegral field i + | otherwise -> + left $ + ConversionFailed (fieldValueTypeName fieldValue) "Word16" $ + Text.pack "negative value cannot be converted to unsigned integer" + FieldWord w -> Ok (fromIntegral w) + other -> left (Incompatible (fieldValueTypeName other) "Word16" "invalid type for Word16") + +instance FromField Word8 where + fromField field@Field{fieldValue} = + case fieldValue of + FieldInt i + | i >= 0 -> boundedIntegral field i + | otherwise -> + left $ + ConversionFailed (fieldValueTypeName fieldValue) "Word8" $ + Text.pack "negative value cannot be converted to unsigned integer" + FieldWord w -> Ok (fromIntegral w) + other -> left (Incompatible (fieldValueTypeName other) "Word8" "invalid type for Word8") + +instance FromField Double where + fromField Field{fieldValue} = + case fieldValue of + FieldDouble d -> Ok d + FieldInt i -> Ok (fromIntegral i) + other -> left (Incompatible (fieldValueTypeName other) "Double" "invalid type for double") + +instance FromField Float where + fromField field = + case (fromField field :: Ok Double) of + Errors err -> Errors err + Ok d -> Ok (realToFrac d) + +instance FromField Text where + fromField Field{fieldValue} = + case fieldValue of + FieldText t -> Ok t + FieldInt i -> Ok (Text.pack (show i)) + FieldDouble d -> Ok (Text.pack (show d)) + FieldBool b -> Ok (if b then Text.pack "1" else Text.pack "0") + FieldNull -> + left $ UnexpectedNull (fieldValueTypeName fieldValue) "Text" $ Text.pack "unexpected null value" + other -> left (Incompatible (fieldValueTypeName other) "Text" "invalid type for text") + +instance FromField String where + fromField field = Text.unpack <$> fromField field + +instance FromField BS.ByteString where + fromField Field{fieldValue} = + case fieldValue of + FieldBlob bs -> Ok bs + FieldText t -> Ok (TextEncoding.encodeUtf8 t) + FieldNull -> + left $ UnexpectedNull (fieldValueTypeName fieldValue) "ByteString" $ Text.pack "unexpected null value" + other -> left (Incompatible (fieldValueTypeName other) "ByteString" "invalid type for blob") + +instance FromField Day where + fromField Field{fieldValue} = + case fieldValue of + FieldDate day -> Ok day + FieldTimestamp LocalTime{localDay} -> Ok localDay + other -> left (Incompatible (fieldValueTypeName other) "Day" "invalid type for date") + +instance FromField TimeOfDay where + fromField Field{fieldValue} = + case fieldValue of + FieldTime tod -> Ok tod + FieldTimestamp LocalTime{localTimeOfDay} -> Ok localTimeOfDay + other -> left (Incompatible (fieldValueTypeName other) "TimeOfDay" "invalid type for time") + +instance FromField LocalTime where + fromField Field{ fieldValue} = + case fieldValue of + FieldTimestamp ts -> Ok ts + FieldDate day -> Ok (LocalTime day midnight) + other -> left (Incompatible (fieldValueTypeName other) "LocalTime" "invalid type for timestamp") + where + midnight = TimeOfDay 0 0 0 + +instance FromField UTCTime where + fromField field = + case fromField field of + Ok (timestamp :: LocalTime) -> Ok (localTimeToUTC utc timestamp) + Errors errs -> Errors errs + +instance (FromField a) => FromField (Maybe a) where + fromField Field{fieldValue = FieldNull} = Ok Nothing + fromField field = Just <$> fromField field + +-- | Helper for bounded integral conversions. +boundedIntegral :: forall a. (Integral a, Bounded a, Typeable a) => Field -> Int -> Ok a +boundedIntegral Field{fieldValue} i + | toInteger i < toInteger (minBound :: a) = + left $ ConversionFailed (fieldValueTypeName fieldValue) (Text.pack $ show (typeRep (Proxy :: Proxy a))) $ Text.pack "integer value out of bounds" + | toInteger i > toInteger (maxBound :: a) = + left $ ConversionFailed (fieldValueTypeName fieldValue) (Text.pack $ show (typeRep (Proxy :: Proxy a))) $ Text.pack "integer value out of bounds" + | otherwise = Ok (fromIntegral i) + +fieldValueTypeName :: FieldValue -> Text +fieldValueTypeName = \case + FieldNull -> "NULL" + FieldInt8{} -> "INT1" + FieldInt16{} -> "INT2" + FieldInt32{} -> "INT4" + FieldInt64{} -> "INT8" + FieldWord8{} -> "UTINYINT" + FieldWord16{} -> "USMALLINT" + FieldWord32{} -> "UINTEGER" + FieldWord64{} -> "UBIGINT" + FieldFloat{} -> "FLOAT" + FieldDouble{} -> "DOUBLE" + FieldText{} -> "TEXT" + FieldBool{} -> "BOOLEAN" + FieldBlob{} -> "BLOB" + FieldDate{} -> "DATE" + FieldTime{} -> "TIME" + FieldTimestamp{} -> "TIMESTAMP" + +-- TODO: Not supported yet +-- FieldInteger{} -> "HUGEINT" +-- FieldNatural{} -> "UHUGEINT" diff --git a/src/Database/DuckDB/Simple/FromRow.hs b/src/Database/DuckDB/Simple/FromRow.hs new file mode 100644 --- /dev/null +++ b/src/Database/DuckDB/Simple/FromRow.hs @@ -0,0 +1,272 @@ +{-# LANGUAGE DefaultSignatures #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE OverloadedStrings #-} + +{- | +Module : Database.DuckDB.Simple.FromRow +Description : Convert rows of 'Field's into Haskell values using a parser-style interface. +-} +module Database.DuckDB.Simple.FromRow ( + -- * Row parsing + RowParser (..), + field, + fieldWith, + numFieldsRemaining, + parseRow, + + -- * Generic derivation + GFromRow (..), + FromRow (..), + + -- * Error conversion + resultErrorToSqlError, + rowErrorsToSqlError, +) where + +import Control.Applicative (Alternative (..)) +import Control.Exception (Exception, SomeException, fromException, toException) +import Control.Monad (MonadPlus, replicateM) +import Control.Monad.Trans.Class (lift) +import Control.Monad.Trans.Reader (ReaderT, ask, runReaderT) +import Control.Monad.Trans.State.Strict (StateT, get, put, runStateT) +import Data.Maybe (listToMaybe, mapMaybe) +import Data.Text (Text) +import qualified Data.Text as Text +import GHC.Generics + +import Database.DuckDB.Simple.FromField +import Database.DuckDB.Simple.Internal (SQLError (..)) +import Database.DuckDB.Simple.Ok (Ok (..)) +import Database.DuckDB.Simple.Types (Only (..), Query, (:.) (..)) + +-- | Row parsing environment (read-only data available to the parser). +newtype RowParseRO = RowParseRO + { rowParseColumnCount :: Int + } + +-- | Column-out-of-bounds sentinel used internally to map parser failures. +newtype ColumnOutOfBounds = ColumnOutOfBounds {columnOutOfBoundsIndex :: Int} + deriving stock (Eq, Show) + +instance Exception ColumnOutOfBounds + +-- | Parser used by 'FromRow' implementations. +newtype RowParser a = RowParser + { runRowParser :: ReaderT RowParseRO (StateT (Int, [Field]) Ok) a + } + deriving stock (Functor) + deriving newtype (Applicative, Alternative, Monad, MonadPlus) + +-- | Generic derivation helper mirroring @sqlite-simple@. +class GFromRow f where + gFromRow :: RowParser (f p) + +instance GFromRow U1 where + gFromRow = pure U1 + +instance (GFromRow a) => GFromRow (M1 i c a) where + gFromRow = M1 <$> gFromRow + +instance (FromField a) => GFromRow (K1 i a) where + gFromRow = K1 <$> field + +instance (GFromRow a, GFromRow b) => GFromRow (a :*: b) where + gFromRow = (:*:) <$> gFromRow <*> gFromRow + +-- | Types that can be constructed from database rows. +class FromRow a where + fromRow :: RowParser a + default fromRow :: (Generic a, GFromRow (Rep a)) => RowParser a + fromRow = to <$> gFromRow + +-- | Pull the next field using the provided 'FieldParser'. +fieldWith :: FieldParser a -> RowParser a +fieldWith fieldParser = RowParser $ do + RowParseRO{rowParseColumnCount} <- ask + (columnIndex, remaining) <- lift get + case remaining of + [] -> do + lift (put (columnIndex + 1, [])) + lift (lift (Errors [toException (ColumnOutOfBounds (columnIndex + 1))])) + (f : rest) -> do + lift (put (columnIndex + 1, rest)) + if columnIndex >= rowParseColumnCount + then lift (lift (Errors [toException (ColumnOutOfBounds (columnIndex + 1))])) + else case fieldParser f of + Errors err -> lift $ lift $ Errors err + Ok value -> pure value + +-- | Pull the next field and parse it using its 'FromField' instance. +field :: (FromField a) => RowParser a +field = fieldWith fromField + +-- | Report how many columns remain unread in the current row. +numFieldsRemaining :: RowParser Int +numFieldsRemaining = RowParser $ do + RowParseRO{rowParseColumnCount} <- ask + (columnIndex, _) <- lift get + pure (rowParseColumnCount - columnIndex) + +-- | Execute a 'RowParser' against the provided row. +parseRow :: RowParser a -> [Field] -> Ok a +parseRow parser fields = + let context = RowParseRO (length fields) + initialState = (0, fields) + in case runStateT (runReaderT (runRowParser parser) context) initialState of + Ok (value, (columnCount, _)) + | columnCount == length fields -> Ok value + | otherwise -> left (ColumnOutOfBounds (columnCount + 1)) + Errors errs -> Errors errs + +instance FromRow () where + fromRow = pure () + +instance (FromField a) => FromRow (Only a) where + fromRow = Only <$> field + +instance (FromRow a, FromRow b) => FromRow (a :. b) where + fromRow = (:.) <$> fromRow <*> fromRow + +instance (FromField a, FromField b) => FromRow (a, b) where + fromRow = (,) <$> field <*> field + +instance (FromField a, FromField b, FromField c) => FromRow (a, b, c) where + fromRow = (,,) <$> field <*> field <*> field + +instance (FromField a, FromField b, FromField c, FromField d) => FromRow (a, b, c, d) where + fromRow = (,,,) <$> field <*> field <*> field <*> field + +instance (FromField a, FromField b, FromField c, FromField d, FromField e) => FromRow (a, b, c, d, e) where + fromRow = (,,,,) <$> field <*> field <*> field <*> field <*> field + +instance (FromField a, FromField b, FromField c, FromField d, FromField e, FromField f) => FromRow (a, b, c, d, e, f) where + fromRow = (,,,,,) <$> field <*> field <*> field <*> field <*> field <*> field + +instance + (FromField a, FromField b, FromField c, FromField d, FromField e, FromField f, FromField g) => + FromRow (a, b, c, d, e, f, g) + where + fromRow = (,,,,,,) <$> field <*> field <*> field <*> field <*> field <*> field <*> field + +instance + (FromField a, FromField b, FromField c, FromField d, FromField e, FromField f, FromField g, FromField h) => + FromRow (a, b, c, d, e, f, g, h) + where + fromRow = (,,,,,,,) <$> field <*> field <*> field <*> field <*> field <*> field <*> field <*> field + +instance + ( FromField a + , FromField b + , FromField c + , FromField d + , FromField e + , FromField f + , FromField g + , FromField h + , FromField i + ) => + FromRow (a, b, c, d, e, f, g, h, i) + where + fromRow = + (,,,,,,,,) + <$> field + <*> field + <*> field + <*> field + <*> field + <*> field + <*> field + <*> field + <*> field + +instance + ( FromField a + , FromField b + , FromField c + , FromField d + , FromField e + , FromField f + , FromField g + , FromField h + , FromField i + , FromField j + ) => + FromRow (a, b, c, d, e, f, g, h, i, j) + where + fromRow = + (,,,,,,,,,) + <$> field + <*> field + <*> field + <*> field + <*> field + <*> field + <*> field + <*> field + <*> field + <*> field + +instance (FromField a) => FromRow [a] where + fromRow = do + remaining <- numFieldsRemaining + replicateM remaining field + +-- | Convert a 'ResultError' into a user-facing 'SQLError'. +resultErrorToSqlError :: Query -> ResultError -> SQLError +resultErrorToSqlError query err = + SQLError + { sqlErrorMessage = renderError err + , sqlErrorType = Nothing + , sqlErrorQuery = Just query + } + +-- | Collapse parser failure diagnostics into an 'SQLError' while preserving the query. +rowErrorsToSqlError :: Query -> [SomeException] -> SQLError +rowErrorsToSqlError query errs = + case listToMaybe (mapMaybe (fromException :: SomeException -> Maybe ResultError) errs) of + Just resultErr -> resultErrorToSqlError query resultErr + Nothing -> + case listToMaybe (mapMaybe (fromException :: SomeException -> Maybe ColumnOutOfBounds) errs) of + Just (ColumnOutOfBounds idx) -> + SQLError + { sqlErrorMessage = + Text.concat + [ "duckdb-simple: column index " + , Text.pack (show idx) + , " out of bounds" + ] + , sqlErrorType = Nothing + , sqlErrorQuery = Just query + } + Nothing -> + SQLError + { sqlErrorMessage = + Text.pack $ "duckdb-simple: row-parsing failed:" <> show errs + , sqlErrorType = Nothing + , sqlErrorQuery = Just query + } + +renderError :: ResultError -> Text +renderError = \case + Incompatible{errSQLType, errHaskellType} -> + Text.concat + [ "duckdb-simple: column " + , " has type " + , errSQLType + , " but expected " + , errHaskellType + ] + UnexpectedNull{errSQLType} -> + Text.concat + [ "duckdb-simple: column " + , " is NULL but expected " + , errSQLType + ] + ConversionFailed{errMessage} -> + errMessage diff --git a/src/Database/DuckDB/Simple/Function.hs b/src/Database/DuckDB/Simple/Function.hs new file mode 100644 --- /dev/null +++ b/src/Database/DuckDB/Simple/Function.hs @@ -0,0 +1,538 @@ +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE UndecidableInstances #-} + +{- | +Module : Database.DuckDB.Simple.Function +Description : Register scalar Haskell functions with DuckDB connections. + +This module mirrors the high-level API provided by @sqlite-simple@ for user +defined functions, adapted to DuckDB's chunked execution model. It allows +pure and 'IO'-based Haskell functions to be exposed to SQL while reusing the +existing 'FromField' and 'ToField'-style typeclass machinery for argument +decoding and result marshalling. +-} +module Database.DuckDB.Simple.Function ( + Function, + createFunction, + deleteFunction, +) where + +import Control.Exception ( + SomeException, + bracket, + displayException, + onException, + throwIO, + try, + ) +import Control.Monad (forM, forM_, when) +import qualified Data.ByteString as BS +import Data.Int (Int16, Int32, Int64, Int8) +import Data.Proxy (Proxy (..)) +import Data.Text (Text) +import qualified Data.Text as Text +import qualified Data.Text.Encoding as TE +import qualified Data.Text.Foreign as TextForeign +import Data.Word (Word16, Word32, Word64, Word8) +import Database.DuckDB.FFI +import Database.DuckDB.Simple.FromField ( + Field (..), + FieldValue (..), + FromField (..), + ) +import Database.DuckDB.Simple.Internal ( + Connection, + Query (..), + SQLError (..), + withConnectionHandle, + withQueryCString, + ) +import Foreign.C.String (peekCString) +import Foreign.C.Types (CBool (..)) +import Foreign.Marshal.Alloc (alloca) +import Foreign.Ptr (Ptr, castPtr, freeHaskellFunPtr, nullPtr, plusPtr) +import Foreign.StablePtr (StablePtr, castPtrToStablePtr, castStablePtrToPtr, deRefStablePtr, freeStablePtr, newStablePtr) +import Foreign.Storable (peekElemOff, poke, pokeElemOff) +import Database.DuckDB.Simple.Ok (Ok(..)) + +-- | Tag DuckDB logical types we support for scalar return values. +data ScalarType + = ScalarTypeBoolean + | ScalarTypeBigInt + | ScalarTypeDouble + | ScalarTypeVarchar + +-- | Runtime representation of values returned to DuckDB. +data ScalarValue + = ScalarNull + | ScalarBoolean !Bool + | ScalarInteger !Int64 + | ScalarDouble !Double + | ScalarText !Text + +-- | Class of scalar results that can be produced by user-defined functions. +class FunctionResult a where + scalarReturnType :: Proxy a -> ScalarType + toScalarValue :: a -> IO ScalarValue + +instance FunctionResult Int where + scalarReturnType _ = ScalarTypeBigInt + toScalarValue value = pure (ScalarInteger (fromIntegral value)) + +instance FunctionResult Int16 where + scalarReturnType _ = ScalarTypeBigInt + toScalarValue value = pure (ScalarInteger (fromIntegral value)) + +instance FunctionResult Int32 where + scalarReturnType _ = ScalarTypeBigInt + toScalarValue value = pure (ScalarInteger (fromIntegral value)) + +instance FunctionResult Int64 where + scalarReturnType _ = ScalarTypeBigInt + toScalarValue value = pure (ScalarInteger value) + +instance FunctionResult Word where + scalarReturnType _ = ScalarTypeBigInt + toScalarValue value = pure (ScalarInteger (fromIntegral value)) + +instance FunctionResult Word16 where + scalarReturnType _ = ScalarTypeBigInt + toScalarValue value = pure (ScalarInteger (fromIntegral value)) + +instance FunctionResult Word32 where + scalarReturnType _ = ScalarTypeBigInt + toScalarValue value = pure (ScalarInteger (fromIntegral value)) + +instance FunctionResult Word64 where + scalarReturnType _ = ScalarTypeBigInt + toScalarValue value = pure (ScalarInteger (fromIntegral value)) + +instance FunctionResult Double where + scalarReturnType _ = ScalarTypeDouble + toScalarValue value = pure (ScalarDouble value) + +instance FunctionResult Float where + scalarReturnType _ = ScalarTypeDouble + toScalarValue value = pure (ScalarDouble (realToFrac value)) + +instance FunctionResult Bool where + scalarReturnType _ = ScalarTypeBoolean + toScalarValue value = pure (ScalarBoolean value) + +instance FunctionResult Text where + scalarReturnType _ = ScalarTypeVarchar + toScalarValue value = pure (ScalarText value) + +instance FunctionResult String where + scalarReturnType _ = ScalarTypeVarchar + toScalarValue value = pure (ScalarText (Text.pack value)) + +instance (FunctionResult a) => FunctionResult (Maybe a) where + scalarReturnType _ = scalarReturnType (Proxy :: Proxy a) + toScalarValue Nothing = pure ScalarNull + toScalarValue (Just value) = toScalarValue value + +-- | Argument types supported by the scalar function machinery. +class FunctionArg a where + argumentType :: Proxy a -> DuckDBType + +instance FunctionArg Int where + argumentType _ = DuckDBTypeBigInt + +instance FunctionArg Int16 where + argumentType _ = DuckDBTypeSmallInt + +instance FunctionArg Int32 where + argumentType _ = DuckDBTypeInteger + +instance FunctionArg Int64 where + argumentType _ = DuckDBTypeBigInt + +instance FunctionArg Word where + argumentType _ = DuckDBTypeUBigInt + +instance FunctionArg Word16 where + argumentType _ = DuckDBTypeUSmallInt + +instance FunctionArg Word32 where + argumentType _ = DuckDBTypeUInteger + +instance FunctionArg Word64 where + argumentType _ = DuckDBTypeUBigInt + +instance FunctionArg Double where + argumentType _ = DuckDBTypeDouble + +instance FunctionArg Float where + argumentType _ = DuckDBTypeFloat + +instance FunctionArg Bool where + argumentType _ = DuckDBTypeBoolean + +instance FunctionArg Text where + argumentType _ = DuckDBTypeVarchar + +instance FunctionArg String where + argumentType _ = DuckDBTypeVarchar + +instance (FunctionArg a) => FunctionArg (Maybe a) where + argumentType _ = argumentType (Proxy :: Proxy a) + +-- | Typeclass describing Haskell functions that can be exposed to DuckDB. +class Function a where + argumentTypes :: Proxy a -> [DuckDBType] + returnType :: Proxy a -> ScalarType + isVolatile :: Proxy a -> Bool + applyFunction :: [Field] -> a -> IO ScalarValue + +instance {-# OVERLAPPABLE #-} (FunctionResult a) => Function a where + argumentTypes _ = [] + returnType _ = scalarReturnType (Proxy :: Proxy a) + isVolatile _ = False + applyFunction [] value = toScalarValue value + applyFunction _ _ = throwIO (functionInvocationError (Text.pack "unexpected arguments supplied")) + +instance {-# OVERLAPPING #-} (FunctionResult a) => Function (IO a) where + argumentTypes _ = [] + returnType _ = scalarReturnType (Proxy :: Proxy a) + isVolatile _ = True + applyFunction [] action = action >>= toScalarValue + applyFunction _ _ = throwIO (functionInvocationError (Text.pack "unexpected arguments supplied")) + +instance {-# OVERLAPPABLE #-} (FromField a, FunctionArg a, Function r) => Function (a -> r) where + argumentTypes _ = argumentType (Proxy :: Proxy a) : argumentTypes (Proxy :: Proxy r) + returnType _ = returnType (Proxy :: Proxy r) + isVolatile _ = isVolatile (Proxy :: Proxy r) + applyFunction [] _ = + throwIO (functionInvocationError (Text.pack "insufficient arguments supplied")) + applyFunction (field : rest) fn = + case fromField field of + Errors err -> throwIO (argumentConversionError (fieldIndex field) err) + Ok value -> applyFunction rest (fn value) + +-- | Register a Haskell function under the supplied name. +createFunction :: forall f. (Function f) => Connection -> Text -> f -> IO () +createFunction conn name fn = do + funPtr <- mkScalarFun (scalarFunctionHandler fn) + funPtrStable <- newStablePtr funPtr + destroyCb <- mkDeleteCallback releaseFunctionPtr + let release = destroyFunctionResources funPtr funPtrStable destroyCb + bracket c_duckdb_create_scalar_function cleanupScalarFunction \scalarFun -> + (`onException` release) $ do + TextForeign.withCString name $ \cName -> + c_duckdb_scalar_function_set_name scalarFun cName + forM_ (argumentTypes (Proxy :: Proxy f)) \dtype -> + withLogicalType dtype $ \logical -> + c_duckdb_scalar_function_add_parameter scalarFun logical + withLogicalType (duckTypeForScalar (returnType (Proxy :: Proxy f))) $ \logical -> + c_duckdb_scalar_function_set_return_type scalarFun logical + when (isVolatile (Proxy :: Proxy f)) $ + c_duckdb_scalar_function_set_volatile scalarFun + c_duckdb_scalar_function_set_function scalarFun funPtr + c_duckdb_scalar_function_set_extra_info scalarFun (castStablePtrToPtr funPtrStable) destroyCb + withConnectionHandle conn \connPtr -> do + rc <- c_duckdb_register_scalar_function connPtr scalarFun + if rc == DuckDBSuccess + then pure () + else throwIO (functionInvocationError (Text.pack "duckdb-simple: registering function failed")) + +-- | Drop a previously registered scalar function by issuing a DROP FUNCTION statement. +deleteFunction :: Connection -> Text -> IO () +deleteFunction conn name = + do + outcome <- + try $ + withConnectionHandle conn \connPtr -> do + let dropQuery = + Query $ + Text.concat + [ Text.pack "DROP FUNCTION IF EXISTS " + , qualifyIdentifier name + ] + withQueryCString dropQuery \sql -> + alloca \resPtr -> do + rc <- c_duckdb_query connPtr sql resPtr + if rc == DuckDBSuccess + then c_duckdb_destroy_result resPtr + else do + errMsg <- fetchResultError resPtr + c_duckdb_destroy_result resPtr + throwIO + SQLError + { sqlErrorMessage = errMsg + , sqlErrorType = Nothing + , sqlErrorQuery = Just dropQuery + } + case outcome of + Right () -> pure () + Left err + -- DuckDB does not allow dropping scalar functions registered via the C API, + -- so we ignore that specific error here. + -- TODO: Update this when DuckDB adds support for dropping such functions. + | Text.isInfixOf (Text.pack "Cannot drop internal catalog entry") (sqlErrorMessage err) -> return () + | otherwise -> throwIO err + +cleanupScalarFunction :: DuckDBScalarFunction -> IO () +cleanupScalarFunction scalarFun = + alloca \ptr -> do + poke ptr scalarFun + c_duckdb_destroy_scalar_function ptr + +destroyFunctionResources :: DuckDBScalarFunctionFun -> StablePtr DuckDBScalarFunctionFun -> DuckDBDeleteCallback -> IO () +destroyFunctionResources funPtr funPtrStable destroyCb = do + freeHaskellFunPtr funPtr + freeStablePtr funPtrStable + freeHaskellFunPtr destroyCb + +releaseFunctionPtr :: Ptr () -> IO () +releaseFunctionPtr rawPtr = + when (rawPtr /= nullPtr) $ do + let stablePtr = castPtrToStablePtr rawPtr :: StablePtr DuckDBScalarFunctionFun + funPtr <- deRefStablePtr stablePtr + freeHaskellFunPtr funPtr + freeStablePtr stablePtr + +withLogicalType :: DuckDBType -> (DuckDBLogicalType -> IO a) -> IO a +withLogicalType dtype = + bracket + ( do + logical <- c_duckdb_create_logical_type dtype + when (logical == nullPtr) $ + throwIO $ + functionInvocationError (Text.pack "duckdb-simple: failed to allocate logical type") + pure logical + ) + destroyLogicalType + +destroyLogicalType :: DuckDBLogicalType -> IO () +destroyLogicalType logicalType = + alloca \ptr -> do + poke ptr logicalType + c_duckdb_destroy_logical_type ptr + +duckTypeForScalar :: ScalarType -> DuckDBType +duckTypeForScalar = \case + ScalarTypeBoolean -> DuckDBTypeBoolean + ScalarTypeBigInt -> DuckDBTypeBigInt + ScalarTypeDouble -> DuckDBTypeDouble + ScalarTypeVarchar -> DuckDBTypeVarchar + +scalarFunctionHandler :: forall f. (Function f) => f -> DuckDBFunctionInfo -> DuckDBDataChunk -> DuckDBVector -> IO () +scalarFunctionHandler fn info chunk outVec = do + result <- + try do + rawColumnCount <- c_duckdb_data_chunk_get_column_count chunk + let columnCount = fromIntegral rawColumnCount :: Int + expected = length (argumentTypes (Proxy :: Proxy f)) + when (columnCount /= expected) $ + throwIO $ + functionInvocationError $ + Text.concat + [ Text.pack "duckdb-simple: function expected " + , Text.pack (show expected) + , Text.pack " arguments but received " + , Text.pack (show columnCount) + ] + rawRowCount <- c_duckdb_data_chunk_get_size chunk + let rowCount = fromIntegral rawRowCount :: Int + readers <- mapM (makeColumnReader chunk) [0 .. expected - 1] + rows <- + forM [0 .. rowCount - 1] \row -> + forM readers \reader -> + reader (fromIntegral row) + results <- mapM (`applyFunction` fn) rows + writeResults (returnType (Proxy :: Proxy f)) results outVec + c_duckdb_data_chunk_set_size chunk (fromIntegral rowCount) + case result of + Left (err :: SomeException) -> do + c_duckdb_data_chunk_set_size chunk 0 + let message = Text.pack (displayException err) + TextForeign.withCString message $ \cMsg -> + c_duckdb_scalar_function_set_error info cMsg + Right () -> pure () + +type ColumnReader = DuckDBIdx -> IO Field + +makeColumnReader :: DuckDBDataChunk -> Int -> IO ColumnReader +makeColumnReader chunk columnIndex = do + vector <- c_duckdb_data_chunk_get_vector chunk (fromIntegral columnIndex) + logical <- c_duckdb_vector_get_column_type vector + dtype <- c_duckdb_get_type_id logical + destroyLogicalType logical + dataPtr <- c_duckdb_vector_get_data vector + validity <- c_duckdb_vector_get_validity vector + let name = Text.pack ("arg" <> show columnIndex) + pure \rowIdx -> do + valid <- isRowValid validity rowIdx + value <- + if valid + then fetchValue dtype dataPtr rowIdx + else pure FieldNull + pure + Field + { fieldName = name + , fieldIndex = columnIndex + , fieldValue = value + } + +isRowValid :: Ptr Word64 -> DuckDBIdx -> IO Bool +isRowValid validity rowIdx + | validity == nullPtr = pure True + | otherwise = do + CBool flag <- c_duckdb_validity_row_is_valid validity rowIdx + pure (flag /= 0) + +fetchValue :: DuckDBType -> Ptr () -> DuckDBIdx -> IO FieldValue +fetchValue dtype dataPtr rowIdx + | dtype == DuckDBTypeBoolean = do + value <- peekElemOff (castPtr dataPtr :: Ptr Word8) (fromIntegral rowIdx) + pure (FieldBool (value /= 0)) + | dtype == DuckDBTypeTinyInt = do + value <- peekElemOff (castPtr dataPtr :: Ptr Int8) (fromIntegral rowIdx) + pure (FieldInt8 (fromIntegral value)) + | dtype == DuckDBTypeSmallInt = do + value <- peekElemOff (castPtr dataPtr :: Ptr Int16) (fromIntegral rowIdx) + pure (FieldInt16 (fromIntegral value)) + | dtype == DuckDBTypeInteger = do + value <- peekElemOff (castPtr dataPtr :: Ptr Int32) (fromIntegral rowIdx) + pure (FieldInt32 (fromIntegral value)) + | dtype == DuckDBTypeBigInt || dtype == DuckDBTypeHugeInt = do + value <- peekElemOff (castPtr dataPtr :: Ptr Int64) (fromIntegral rowIdx) + pure (FieldInt64 value) + | dtype == DuckDBTypeUTinyInt = do + value <- peekElemOff (castPtr dataPtr :: Ptr Word8) (fromIntegral rowIdx) + pure (FieldWord8 (fromIntegral value)) + | dtype == DuckDBTypeUSmallInt = do + value <- peekElemOff (castPtr dataPtr :: Ptr Word16) (fromIntegral rowIdx) + pure (FieldWord16 (fromIntegral value)) + | dtype == DuckDBTypeUInteger = do + value <- peekElemOff (castPtr dataPtr :: Ptr Word32) (fromIntegral rowIdx) + pure (FieldWord32 (fromIntegral value)) + | dtype == DuckDBTypeUBigInt || dtype == DuckDBTypeUHugeInt = do + value <- peekElemOff (castPtr dataPtr :: Ptr Word64) (fromIntegral rowIdx) + pure (FieldWord64 (fromIntegral value)) + | dtype == DuckDBTypeFloat = do + value <- peekElemOff (castPtr dataPtr :: Ptr Float) (fromIntegral rowIdx) + pure (FieldDouble (realToFrac value)) + | dtype == DuckDBTypeDouble = do + value <- peekElemOff (castPtr dataPtr :: Ptr Double) (fromIntegral rowIdx) + pure (FieldDouble value) + | dtype == DuckDBTypeVarchar = do + textValue <- decodeDuckText dataPtr rowIdx + pure (FieldText textValue) + | otherwise = + throwIO $ + functionInvocationError $ + Text.concat + [ Text.pack "duckdb-simple: unsupported argument type " + , Text.pack (show dtype) + ] + +decodeDuckText :: Ptr () -> DuckDBIdx -> IO Text +decodeDuckText dataPtr rowIdx = do + let bytePtr = castPtr dataPtr :: Ptr Word8 + offset = fromIntegral rowIdx * duckdbStringTSize + stringPtr = castPtr (bytePtr `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 (TE.decodeUtf8 bytes) + +duckdbStringTSize :: Int +duckdbStringTSize = 16 + +writeResults :: ScalarType -> [ScalarValue] -> DuckDBVector -> IO () +writeResults resultType values outVec = do + let hasNulls = any isNullValue values + when hasNulls $ + c_duckdb_vector_ensure_validity_writable outVec + dataPtr <- c_duckdb_vector_get_data outVec + validityPtr <- c_duckdb_vector_get_validity outVec + forM_ (zip [0 ..] values) \(idx, val) -> + case (resultType, val) of + (_, ScalarNull) -> + markInvalid validityPtr idx + (ScalarTypeBoolean, ScalarBoolean flag) -> do + markValid validityPtr idx + pokeElemOff (castPtr dataPtr :: Ptr Word8) idx (if flag then 1 else 0) + (ScalarTypeBigInt, ScalarInteger intval) -> do + markValid validityPtr idx + pokeElemOff (castPtr dataPtr :: Ptr Int64) idx intval + (ScalarTypeDouble, ScalarDouble dbl) -> do + markValid validityPtr idx + pokeElemOff (castPtr dataPtr :: Ptr Double) idx dbl + (ScalarTypeVarchar, ScalarText txt) -> do + markValid validityPtr idx + TextForeign.withCStringLen txt \(ptr, len) -> + c_duckdb_vector_assign_string_element_len outVec (fromIntegral idx) ptr (fromIntegral len) + _ -> + throwIO $ + functionInvocationError $ + Text.pack "duckdb-simple: result type mismatch when materialising scalar function output" + +markInvalid :: Ptr Word64 -> Int -> IO () +markInvalid validity idx + | validity == nullPtr = pure () + | otherwise = c_duckdb_validity_set_row_invalid validity (fromIntegral idx) + +markValid :: Ptr Word64 -> Int -> IO () +markValid validity idx + | validity == nullPtr = pure () + | otherwise = c_duckdb_validity_set_row_valid validity (fromIntegral idx) + +isNullValue :: ScalarValue -> Bool +isNullValue = \case + ScalarNull -> True + _ -> False + +argumentConversionError :: Int -> [SomeException] -> SQLError +argumentConversionError idx err = + let message = + Text.concat + [ Text.pack "duckdb-simple: unable to convert argument #" + , Text.pack (show (idx + 1)) + , Text.pack ": " + , Text.pack (show err) + ] + in functionInvocationError message + +functionInvocationError :: Text -> SQLError +functionInvocationError message = + SQLError + { sqlErrorMessage = message + , sqlErrorType = Nothing + , sqlErrorQuery = Nothing + } + +fetchResultError :: Ptr DuckDBResult -> IO Text +fetchResultError resPtr = do + msgPtr <- c_duckdb_result_error resPtr + if msgPtr == nullPtr + then pure (Text.pack "duckdb-simple: DROP FUNCTION failed") + else Text.pack <$> peekCString msgPtr + +qualifyIdentifier :: Text -> Text +qualifyIdentifier rawName = + let parts = Text.splitOn "." rawName + in Text.intercalate (Text.pack ".") (map quoteIdent parts) + +quoteIdent :: Text -> Text +quoteIdent ident = + Text.concat + [ Text.pack "\"" + , Text.replace (Text.pack "\"") (Text.pack "\"\"") ident + , Text.pack "\"" + ] + +foreign import ccall "wrapper" + mkScalarFun :: (DuckDBFunctionInfo -> DuckDBDataChunk -> DuckDBVector -> IO ()) -> IO DuckDBScalarFunctionFun + +foreign import ccall "wrapper" + mkDeleteCallback :: (Ptr () -> IO ()) -> IO DuckDBDeleteCallback diff --git a/src/Database/DuckDB/Simple/Internal.hs b/src/Database/DuckDB/Simple/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Database/DuckDB/Simple/Internal.hs @@ -0,0 +1,182 @@ +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE StrictData #-} + +{- | +Module : Database.DuckDB.Simple.Internal +Description : Internal machinery backing the duckdb-simple API surface. + +This module provides access to the opaque data constructors and helper +utilities required by the high-level API. It is not part of the supported +public interface; consumers should depend on 'Database.DuckDB.Simple' instead. +-} +module Database.DuckDB.Simple.Internal ( + -- * Data constructors (internal use only) + Query (..), + Connection (..), + ConnectionState (..), + Statement (..), + StatementState (..), + StatementStreamState (..), + StatementStream (..), + StatementStreamColumn (..), + StatementStreamChunk (..), + StatementStreamChunkVector (..), + SQLError (..), + toSQLError, + + -- * Helpers + connectionClosedError, + statementClosedError, + withConnectionHandle, + withStatementHandle, + withQueryCString, +) where + +import Control.Exception (Exception, throwIO) +import Data.IORef (IORef, readIORef) +import Data.String (IsString (..)) +import Data.Text (Text) +import qualified Data.Text as Text +import qualified Data.Text.Foreign as TextForeign +import Data.Word (Word64) +import Database.DuckDB.FFI ( + DuckDBDataChunk, + DuckDBConnection, + DuckDBDatabase, + DuckDBErrorType, + DuckDBPreparedStatement, + DuckDBResult, + DuckDBType, + DuckDBVector, + ) +import Foreign.C.String (CString) +import Foreign.Ptr (Ptr) + +-- | Represents a textual SQL query with UTF-8 encoding semantics. +newtype Query = Query + { fromQuery :: Text + -- ^ Extract the underlying textual representation of the query. + } + deriving stock (Eq, Ord, Show) + +instance Semigroup Query where + Query a <> Query b = Query (a <> b) + +instance IsString Query where + fromString = Query . Text.pack + +-- | Tracks the lifetime of a DuckDB database and connection pair. +newtype Connection = Connection { connectionState :: IORef ConnectionState } + +-- | Internal connection lifecycle state. +data ConnectionState + = ConnectionClosed + | ConnectionOpen + { connectionDatabase :: DuckDBDatabase + , connectionHandle :: DuckDBConnection + } + +-- | Tracks the lifetime of a prepared statement. +data Statement = Statement + { statementState :: IORef StatementState + , statementConnection :: Connection + , statementQuery :: Query + , statementStream :: IORef StatementStreamState + } + +-- | Internal statement lifecycle state. +data StatementState + = StatementClosed + | StatementOpen + { statementHandle :: DuckDBPreparedStatement + } + +-- | Streaming execution state for prepared statements. +data StatementStreamState + = StatementStreamIdle + | StatementStreamActive !StatementStream + +-- | Streaming cursor backing an active result set. +data StatementStream = StatementStream + { statementStreamResult :: Ptr DuckDBResult + , statementStreamColumns :: [StatementStreamColumn] + , statementStreamChunk :: Maybe StatementStreamChunk + } + +-- | Metadata describing a result column surfaced through streaming. +data StatementStreamColumn = StatementStreamColumn + { statementStreamColumnIndex :: Int + , statementStreamColumnName :: Text + , statementStreamColumnType :: DuckDBType + } + +-- | Currently loaded data chunk plus iteration cursor. +data StatementStreamChunk = StatementStreamChunk + { 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 + } + +-- | Represents an error reported by DuckDB or by duckdb-simple itself. +data SQLError = SQLError + { sqlErrorMessage :: Text + , sqlErrorType :: Maybe DuckDBErrorType + , sqlErrorQuery :: Maybe Query + } + deriving stock (Eq, Show) + +instance Exception SQLError + +toSQLError :: Exception e => e -> SQLError +toSQLError ex = SQLError + { sqlErrorMessage = Text.pack (show ex) + , sqlErrorType = Nothing + , sqlErrorQuery = Nothing + } + +-- | Shared error value used when an operation targets a closed connection. +connectionClosedError :: SQLError +connectionClosedError = + SQLError + { sqlErrorMessage = Text.pack "duckdb-simple: connection is closed" + , sqlErrorType = Nothing + , sqlErrorQuery = Nothing + } + +-- | Shared error value used when an operation targets a closed statement. +statementClosedError :: Statement -> SQLError +statementClosedError Statement{statementQuery} = + SQLError + { sqlErrorMessage = Text.pack "duckdb-simple: statement is closed" + , sqlErrorType = Nothing + , sqlErrorQuery = Just statementQuery + } + +-- | Provide a UTF-8 encoded C string view of the query text. +withQueryCString :: Query -> (CString -> IO a) -> IO a +withQueryCString (Query txt) = TextForeign.withCString txt + +-- | Internal helper for safely accessing the underlying prepared statement. +withStatementHandle :: Statement -> (DuckDBPreparedStatement -> IO a) -> IO a +withStatementHandle stmt@Statement{statementState} action = do + state <- readIORef statementState + case state of + StatementClosed -> throwIO (statementClosedError stmt) + StatementOpen{statementHandle} -> action statementHandle + +-- | Internal helper for safely accessing the underlying connection handle. +withConnectionHandle :: Connection -> (DuckDBConnection -> IO a) -> IO a +withConnectionHandle Connection{connectionState} action = do + state <- readIORef connectionState + case state of + ConnectionClosed -> throwIO connectionClosedError + ConnectionOpen{connectionHandle} -> action connectionHandle diff --git a/src/Database/DuckDB/Simple/Ok.hs b/src/Database/DuckDB/Simple/Ok.hs new file mode 100644 --- /dev/null +++ b/src/Database/DuckDB/Simple/Ok.hs @@ -0,0 +1,64 @@ +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DerivingStrategies #-} + +{- | +Module : Database.DuckDB.Simple.Ok +Description : Lightweight error accumulation for row parsing. + +The 'Ok' type mirrors the helper used by @sqlite-simple@: it behaves like an +`Either [SomeException]` with sensible 'Alternative' semantics that accumulate +failure reasons instead of discarding them. This underpins the RowParser +machinery in 'Database.DuckDB.Simple.FromRow'. +-} +module Database.DuckDB.Simple.Ok ( + Ok (..), + ManyErrors (..), +) where + +import Control.Applicative (Alternative (..)) +import Control.Exception (Exception, SomeException, toException) +import Control.Monad (MonadPlus (..)) +import qualified Control.Monad.Fail as Fail + +-- | Simple success-or-errors container with error accumulation. +data Ok a + = Errors [SomeException] + | Ok !a + deriving stock (Show) + deriving stock (Functor) + +-- | Two failures are considered equal regardless of their payload. +instance (Eq a) => Eq (Ok a) where + Errors _ == Errors _ = True + Ok a == Ok b = a == b + _ == _ = False + +instance Applicative Ok where + pure = Ok + Ok f <*> Ok a = Ok (f a) + Errors es <*> Ok _ = Errors es + Ok _ <*> Errors es = Errors es + Errors es1 <*> Errors es2 = Errors (es1 <> es2) + +instance Alternative Ok where + empty = Errors [] + Ok x <|> _ = Ok x + Errors _ <|> Ok y = Ok y + Errors xs <|> Errors ys = Errors (xs <> ys) + +instance Monad Ok where + Ok a >>= f = f a + Errors es >>= _ = Errors es + +instance MonadPlus Ok where + mzero = empty + mplus = (<|>) + +instance Fail.MonadFail Ok where + fail msg = Errors [toException (userError msg)] + +-- | Bundle multiple underlying exceptions into a single throwable value. +newtype ManyErrors = ManyErrors [SomeException] + deriving stock (Show) + +instance Exception ManyErrors diff --git a/src/Database/DuckDB/Simple/ToField.hs b/src/Database/DuckDB/Simple/ToField.hs new file mode 100644 --- /dev/null +++ b/src/Database/DuckDB/Simple/ToField.hs @@ -0,0 +1,284 @@ +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE OverloadedStrings #-} + +{- | +Module : Database.DuckDB.Simple.ToField +Description : Convert Haskell parameters into DuckDB bindable values. + +The 'ToField' class mirrors the interface provided by @sqlite-simple@ while +delegating to the DuckDB C API under the hood. +-} +module Database.DuckDB.Simple.ToField ( + FieldBinding, + ToField (..), + NamedParam (..), + bindFieldBinding, + renderFieldBinding, +) where + +import Control.Exception (bracket, throwIO) +import Control.Monad (when) +import qualified Data.ByteString as BS +import Data.Int (Int16, Int32, Int64) +import Data.Text (Text) +import qualified Data.Text as Text +import qualified Data.Text.Foreign as TextForeign +import Data.Time.Calendar (Day, toGregorian) +import Data.Time.Clock (UTCTime (..), diffTimeToPicoseconds) +import Data.Time.LocalTime (LocalTime (..), TimeOfDay (..), timeOfDayToTime, utc, utcToLocalTime) +import Data.Word (Word16, Word32, Word64, Word8) +import Database.DuckDB.FFI +import Database.DuckDB.Simple.Internal ( + SQLError (..), + Statement (..), + withStatementHandle, + ) +import Database.DuckDB.Simple.Types (Null (..)) +import Foreign.C.String (peekCString) +import Foreign.C.Types (CDouble (..)) +import Foreign.Marshal.Alloc (alloca) +import Foreign.Ptr (Ptr, castPtr, nullPtr) +import Foreign.Storable (poke) + +-- | Represents a named parameter binding using the @:=@ operator. +data NamedParam where + (:=) :: (ToField a) => Text -> a -> NamedParam + +infixr 3 := + +-- | 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 ()) + , fieldBindingDisplay :: !String + } + +-- | Apply a 'FieldBinding' to the given statement/index. +bindFieldBinding :: Statement -> DuckDBIdx -> FieldBinding -> IO () +bindFieldBinding stmt idx FieldBinding{fieldBindingAction} = fieldBindingAction stmt idx + +-- | 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 + +instance ToField Null where + toField Null = nullBinding "NULL" + +instance ToField Bool where + toField value = + mkFieldBinding (show value) $ \stmt idx -> + bindDuckValue stmt idx (c_duckdb_create_bool (if value then 1 else 0)) + +instance ToField Int where + toField = intBinding . (fromIntegral :: Int -> Int64) + +instance ToField Int16 where + toField = intBinding . (fromIntegral :: Int16 -> Int64) + +instance ToField Int32 where + toField = intBinding . (fromIntegral :: Int32 -> Int64) + +instance ToField Int64 where + toField = intBinding + +instance ToField Word where + toField value = intBinding (fromIntegral value :: Int64) + +instance ToField Word16 where + toField value = intBinding (fromIntegral value :: Int64) + +instance ToField Word32 where + toField value = intBinding (fromIntegral value :: Int64) + +instance ToField Word64 where + toField value = + mkFieldBinding + (show value) + \stmt idx -> + bindDuckValue stmt idx (c_duckdb_create_int64 (fromIntegral value)) + +instance ToField Double where + toField value = + mkFieldBinding + (show value) + \stmt idx -> + bindDuckValue stmt idx (c_duckdb_create_double (CDouble value)) + +instance ToField Float where + toField value = + mkFieldBinding + (show value) + \stmt idx -> + bindDuckValue stmt idx (c_duckdb_create_double (CDouble (realToFrac value))) + +instance ToField Text where + toField txt = + mkFieldBinding + (show txt) + \stmt idx -> + TextForeign.withCString txt $ \cstr -> + bindDuckValue stmt idx (c_duckdb_create_varchar cstr) + +instance ToField String where + toField str = + mkFieldBinding + (show str) + \stmt idx -> + TextForeign.withCString (Text.pack str) $ \cstr -> + bindDuckValue stmt idx (c_duckdb_create_varchar cstr) + +instance ToField BS.ByteString where + toField bs = + mkFieldBinding + ("") + \stmt idx -> + BS.useAsCStringLen bs \(ptr, len) -> + bindDuckValue stmt idx (c_duckdb_create_blob (castPtr ptr :: Ptr Word8) (fromIntegral len)) + +instance ToField Day where + toField day = + mkFieldBinding + (show day) + \stmt idx -> + bindDuckValue stmt idx $ do + duckDate <- encodeDay day + c_duckdb_create_date duckDate + +instance ToField TimeOfDay where + toField tod = + mkFieldBinding + (show tod) + \stmt idx -> + bindDuckValue stmt idx $ do + duckTime <- encodeTimeOfDay tod + c_duckdb_create_time duckTime + +instance ToField LocalTime where + toField ts = + mkFieldBinding + (show ts) + \stmt idx -> + bindDuckValue stmt idx $ do + duckTimestamp <- encodeLocalTime ts + c_duckdb_create_timestamp duckTimestamp + +instance ToField UTCTime where + toField utcTime = + let FieldBinding{fieldBindingAction = action} = toField (utcToLocalTime utc utcTime) + in FieldBinding + { fieldBindingAction = action + , fieldBindingDisplay = show utcTime + } + +instance (ToField a) => ToField (Maybe a) where + toField Nothing = nullBinding "Nothing" + toField (Just value) = + let binding = toField value + in binding + { fieldBindingDisplay = "Just " <> renderFieldBinding binding + } + +-- | Helper for binding 'Null' values. +nullBinding :: String -> FieldBinding +nullBinding repr = + mkFieldBinding + repr + \stmt idx -> + bindDuckValue stmt idx c_duckdb_create_null_value + +intBinding :: Int64 -> FieldBinding +intBinding value = + mkFieldBinding + (show value) + \stmt idx -> + bindDuckValue stmt idx (c_duckdb_create_int64 value) + +encodeDay :: Day -> IO DuckDBDate +encodeDay day = + alloca \ptr -> do + poke ptr (dayToDateStruct day) + c_duckdb_to_date ptr + +encodeTimeOfDay :: TimeOfDay -> IO DuckDBTime +encodeTimeOfDay tod = + alloca \ptr -> do + poke ptr (timeOfDayToStruct tod) + c_duckdb_to_time ptr + +encodeLocalTime :: LocalTime -> IO DuckDBTimestamp +encodeLocalTime LocalTime{localDay, localTimeOfDay} = + alloca \ptr -> do + poke + ptr + DuckDBTimestampStruct + { duckDBTimestampStructDate = dayToDateStruct localDay + , duckDBTimestampStructTime = timeOfDayToStruct localTimeOfDay + } + c_duckdb_to_timestamp ptr + +dayToDateStruct :: Day -> DuckDBDateStruct +dayToDateStruct day = + let (year, month, dayOfMonth) = toGregorian day + in DuckDBDateStruct + { duckDBDateStructYear = fromIntegral year + , duckDBDateStructMonth = fromIntegral month + , duckDBDateStructDay = fromIntegral dayOfMonth + } + +timeOfDayToStruct :: TimeOfDay -> DuckDBTimeStruct +timeOfDayToStruct tod = + let totalPicoseconds = diffTimeToPicoseconds (timeOfDayToTime tod) + totalMicros = totalPicoseconds `div` 1000000 + (hours, remHour) = totalMicros `divMod` (60 * 60 * 1000000) + (minutes, remMinute) = remHour `divMod` (60 * 1000000) + (seconds, micros) = remMinute `divMod` 1000000 + in DuckDBTimeStruct + { duckDBTimeStructHour = fromIntegral hours + , duckDBTimeStructMinute = fromIntegral minutes + , duckDBTimeStructSecond = fromIntegral seconds + , duckDBTimeStructMicros = fromIntegral micros + } + +bindDuckValue :: Statement -> DuckDBIdx -> IO DuckDBValue -> IO () +bindDuckValue stmt idx makeValue = + withStatementHandle stmt \handle -> + bracket makeValue destroyValue \value -> do + rc <- c_duckdb_bind_value handle idx value + when (rc /= DuckDBSuccess) $ do + err <- fetchPrepareError 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 = + throwIO + SQLError + { sqlErrorMessage = msg + , sqlErrorType = Nothing + , sqlErrorQuery = Just statementQuery + } diff --git a/src/Database/DuckDB/Simple/ToRow.hs b/src/Database/DuckDB/Simple/ToRow.hs new file mode 100644 --- /dev/null +++ b/src/Database/DuckDB/Simple/ToRow.hs @@ -0,0 +1,65 @@ +{-# LANGUAGE DefaultSignatures #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeOperators #-} + +{- | +Module : Database.DuckDB.Simple.ToRow +Description : Rendering Haskell collections into parameter rows. + +The 'ToRow' class converts Haskell values into lists of 'FieldBinding's that +can be consumed by statement binding functions. +-} +module Database.DuckDB.Simple.ToRow ( + ToRow (..), + GToRow (..), +) where + +import Database.DuckDB.Simple.ToField (FieldBinding, ToField (..)) +import Database.DuckDB.Simple.Types (Only (..), (:.) (..)) +import GHC.Generics + +-- | Types that can be transformed into parameter bindings. +class ToRow a where + toRow :: a -> [FieldBinding] + default toRow :: (Generic a, GToRow (Rep a)) => a -> [FieldBinding] + toRow = gtoRow . from + +-- | Generic helper for deriving 'ToRow'. +class GToRow f where + gtoRow :: f p -> [FieldBinding] + +instance GToRow U1 where + gtoRow _ = [] + +instance (ToField a) => GToRow (K1 i a) where + gtoRow (K1 a) = [toField a] + +instance (GToRow a, GToRow b) => GToRow (a :*: b) where + gtoRow (a :*: b) = gtoRow a ++ gtoRow b + +instance (GToRow a) => GToRow (M1 i c a) where + gtoRow (M1 a) = gtoRow a + +instance ToRow () where + toRow () = [] + +instance (ToField a) => ToRow (Only a) where + toRow (Only a) = [toField a] + +instance (ToField a, ToField b) => ToRow (a, b) where + toRow (a, b) = [toField a, toField b] + +instance (ToField a, ToField b, ToField c) => ToRow (a, b, c) where + toRow (a, b, c) = [toField a, toField b, toField c] + +instance (ToField a, ToField b, ToField c, ToField d) => ToRow (a, b, c, d) where + toRow (a, b, c, d) = [toField a, toField b, toField c, toField d] + +instance (ToField a, ToField b, ToField c, ToField d, ToField e) => ToRow (a, b, c, d, e) where + toRow (a, b, c, d, e) = [toField a, toField b, toField c, toField d, toField e] + +instance (ToField a) => ToRow [a] where + toRow = fmap toField + +instance (ToRow a, ToRow b) => ToRow (a :. b) where + toRow (a :. b) = toRow a ++ toRow b diff --git a/src/Database/DuckDB/Simple/Types.hs b/src/Database/DuckDB/Simple/Types.hs new file mode 100644 --- /dev/null +++ b/src/Database/DuckDB/Simple/Types.hs @@ -0,0 +1,58 @@ +{-# LANGUAGE TypeOperators #-} + +{- | +Module : Database.DuckDB.Simple.Types +Description : Shared data types for the duckdb-simple public surface. + +The datatypes in this module are intentionally kept lightweight. The main +`Database.DuckDB.Simple` module exposes the types without their constructors so +that callers interact with them through the high-level API. The actual +definitions live in 'Database.DuckDB.Simple.Internal'. +-} +module Database.DuckDB.Simple.Types ( + Connection, + Statement, + Query (..), + SQLError (..), + FormatError (..), + Null (..), + Only (..), + (:.) (..), +) where + +import Control.Exception (Exception) +import Data.Text (Text) + +import Database.DuckDB.Simple.Internal ( + Connection, + Query (..), + SQLError (..), + Statement, + ) + +-- | Placeholder representing SQL @NULL@. +data Null = Null + deriving (Eq, Ord, Show, Read) + +-- | Wrapper used for single-column rows. +newtype Only a = Only {fromOnly :: a} + deriving (Eq, Ord, Show, Read) + +-- | Convenience product type for combining 'FromRow'/'ToRow' instances. +data h :. t = h :. t + deriving (Eq, Ord, Show, Read) + +infixr 3 :. + +-- | Raised when parameter formatting fails before a statement is executed. +data FormatError = FormatError + { formatErrorMessage :: !Text + -- ^ Human-readable description of the mismatch (e.g. parameter counts). + , formatErrorQuery :: !Query + -- ^ Query that triggered the formatting failure. + , formatErrorParams :: ![String] + -- ^ Rendered parameter values supplied by the caller (used for diagnostics). + } + deriving (Eq, Show) + +instance Exception FormatError diff --git a/test/Spec.hs b/test/Spec.hs new file mode 100644 --- /dev/null +++ b/test/Spec.hs @@ -0,0 +1,463 @@ +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeOperators #-} + +-- | Tasty-based test suite for duckdb-simple. +module Main (main) where + +import Control.Applicative ((<|>)) +import Control.Exception (ErrorCall, Exception, try) +import Control.Monad (replicateM_) +import qualified Data.ByteString as BS +import Data.IORef (atomicModifyIORef', newIORef) +import Data.Int (Int64) +import qualified Data.Text as Text +import Data.Time.Calendar (fromGregorian) +import Data.Time.Clock (UTCTime) +import Data.Time.LocalTime (LocalTime (..), TimeOfDay (..), localTimeToUTC, utc) +import Database.DuckDB.Simple +import Database.DuckDB.Simple.FromField (Field (..), left) +import GHC.Generics (Generic) +import Test.Tasty (TestTree, defaultMain, testGroup) +import Test.Tasty.HUnit +import Database.DuckDB.Simple.Ok (Ok(..)) + +data Person = Person + { personId :: Int + , personName :: Text.Text + } + deriving stock (Eq, Show, Generic) + deriving anyclass (FromRow) + +data WithRemaining = WithRemaining Int Int + deriving (Eq, Show) + +instance FromRow WithRemaining where + fromRow = do + firstVal <- field + remaining <- numFieldsRemaining + replicateM_ remaining (fieldWith (const (Ok ()))) + pure (WithRemaining firstVal remaining) + +newtype YesNo = YesNo Bool + deriving (Eq, Show) + +instance FromRow YesNo where + fromRow = parseYes <|> parseNo + where + parseYes = YesNo True <$ fieldWith (match (Text.pack "yes")) + parseNo = YesNo False <$ fieldWith (match (Text.pack "no")) + + match expected fld@Field{} = + case fromField fld :: Ok Text.Text of + Errors err -> Errors err + Ok txt -> + let normalized = Text.toLower txt + in if normalized == expected + then Ok () + else left $ + ConversionFailed + (Text.pack "TEXT") + (Text.pack "YesNo") + "failed to match exact string" + +newtype NonEmptyText = NonEmptyText Text.Text + deriving (Eq, Show) + +nonEmptyTextParser :: FieldParser NonEmptyText +nonEmptyTextParser f@Field{} = + case fromField f of + Errors err -> Errors err + Ok txt + | Text.null txt -> + left $ + ConversionFailed + (Text.pack "TEXT") + (Text.pack "NonEmptyText") + (Text.pack "NonEmptyText requires a non-empty string") + | otherwise -> Ok (NonEmptyText txt) + +instance FromField NonEmptyText where + fromField = nonEmptyTextParser +main :: IO () +main = defaultMain tests + +tests :: TestTree +tests = + testGroup + "duckdb-simple" + [ connectionTests + , withConnectionTests + , statementTests + , typeTests + , streamingTests + , functionsTests + , transactionTests + ] + +connectionTests :: TestTree +connectionTests = + testGroup + "open/close" + [ testCase "opens and closes an in-memory database" $ do + conn <- open ":memory:" + close conn + , testCase "allows closing the same connection twice" $ do + conn <- open ":memory:" + close conn + close conn + ] + +withConnectionTests :: TestTree +withConnectionTests = + testGroup + "withConnection" + [ testCase "returns the action result" $ do + result <- withConnection ":memory:" \_ -> pure (21 :: Int) + assertEqual "action result" 21 result + , testCase "propagates exceptions from the action" $ + assertThrowsErrorCall $ + withConnection ":memory:" (\_ -> error "boom" :: IO ()) + ] + +statementTests :: TestTree +statementTests = + testGroup + "statements" + [ testCase "prepares and closes a simple statement" $ + withConnection ":memory:" \conn -> do + stmt <- openStatement conn "SELECT 42" + closeStatement stmt + , testCase "allows closing a statement twice" $ + withConnection ":memory:" \conn -> do + stmt <- openStatement conn "SELECT 1" + closeStatement stmt + closeStatement stmt + , testCase "rejects preparing statements on a closed connection" $ do + conn <- open ":memory:" + close conn + assertThrowsSQLError $ + openStatement conn "SELECT 1" + , testCase "throws SQLError for invalid SQL" $ + withConnection ":memory:" \conn -> + assertThrowsSQLError $ + openStatement conn "THIS IS NOT VALID SQL" + , testCase "withStatement closes the statement automatically" $ + withConnection ":memory:" \conn -> do + result <- withStatement conn "SELECT 1" \_ -> pure ("done" :: String) + assertEqual "withStatement result" "done" result + , testCase "withStatement closes statements even when the action fails" $ + withConnection ":memory:" \conn -> do + assertThrowsErrorCall $ + withStatement conn "SELECT 1" (\_ -> error "boom" :: IO ()) + stmt <- openStatement conn "SELECT 1" + closeStatement stmt + , testCase "execute returns affected row count" $ + withConnection ":memory:" \conn -> do + _ <- execute_ conn "CREATE TABLE test_exec (x INTEGER)" + count <- execute conn "INSERT INTO test_exec VALUES (?)" (Only (1 :: Int)) + assertEqual "rows affected" 1 count + , testCase "execute runs with positional parameters" $ + withConnection ":memory:" \conn -> do + _ <- execute_ conn "CREATE TABLE exec_params (a INTEGER, b TEXT)" + count <- execute conn "INSERT INTO exec_params VALUES (?, ?)" (5 :: Int, "hi" :: String) + assertEqual "rows affected" 1 count + , testCase "rejects invalid direct execution" $ + withConnection ":memory:" \conn -> + assertThrowsSQLError $ + execute_ conn "THIS IS NOT SQL" + , testCase "binds positional parameters" $ + withConnection ":memory:" \conn -> do + _ <- execute_ conn "CREATE TABLE params (a INTEGER, b TEXT)" + withStatement conn "INSERT INTO params VALUES (?, ?)" \stmt -> do + bind stmt [toField (5 :: Int), toField ("hi" :: String)] + changed <- executeStatement stmt + assertEqual "rows affected" 1 changed + , testCase "executeMany reuses prepared statements" $ + withConnection ":memory:" \conn -> do + _ <- execute_ conn "CREATE TABLE params_many (a INTEGER, b TEXT)" + total <- executeMany conn "INSERT INTO params_many VALUES (?, ?)" [(1 :: Int, "x" :: String), (2 :: Int, "y" :: String)] + assertEqual "rows affected" 2 total + , testCase "executeNamed binds named parameters" $ + withConnection ":memory:" \conn -> do + _ <- execute_ conn "CREATE TABLE named_params (a INTEGER, b TEXT)" + count <- executeNamed conn "INSERT INTO named_params VALUES ($a, $b)" ["$a" := (1 :: Int), "$b" := ("named" :: String)] + assertEqual "rows affected" 1 count + rows <- queryNamed conn "SELECT a FROM named_params WHERE b = $label" ["$label" := ("named" :: String)] + assertEqual "named query" [Only (1 :: Int)] rows + , testCase "rejects incorrect positional argument counts" $ + withConnection ":memory:" \conn -> do + _ <- execute_ conn "CREATE TABLE param_count (a INTEGER, b INTEGER)" + assertThrowsFormatError + (execute conn "INSERT INTO param_count VALUES (?, ?)" (Only (1 :: Int))) + (Text.isInfixOf "parameter(s)" . formatErrorMessage) + , testCase "rejects positional bindings on named statements" $ + withConnection ":memory:" \conn -> do + _ <- execute_ conn "CREATE TABLE mixed_mode (a INTEGER)" + assertThrowsFormatError + (execute conn "INSERT INTO mixed_mode VALUES ($a)" (Only (1 :: Int))) + (Text.isInfixOf "named parameters" . formatErrorMessage) + , testCase "reports error when mixing positional and named parameters" $ + withConnection ":memory:" \conn -> do + _ <- execute_ conn "CREATE TABLE mixed_params (a INTEGER, b TEXT)" + assertThrows + ( withStatement conn "INSERT INTO mixed_params VALUES (?, $label)" \stmt -> do + bind stmt [toField (1 :: Int)] + bindNamed stmt ["$label" := ("combo" :: String)] + _ <- executeStatement stmt + pure () + ) + ( \(err :: SQLError) -> + Text.isInfixOf "Mixing named and positional parameters" (sqlErrorMessage err) + ) + , testCase "generic FromRow derivation works" $ + withConnection ":memory:" \conn -> do + _ <- execute_ conn "CREATE TABLE person (id INTEGER, name TEXT)" + _ <- executeMany conn "INSERT INTO person VALUES (?, ?)" [(1 :: Int, "Alice" :: String), (2 :: Int, "Bob" :: String)] + people <- query_ conn "SELECT id, name FROM person ORDER BY id" :: IO [Person] + assertEqual "person rows" [Person 1 (Text.pack "Alice"), Person 2 (Text.pack "Bob")] people + , testCase "(:.) composes row parsing and parameter encoding" $ + withConnection ":memory:" \conn -> do + _ <- execute_ conn "CREATE TABLE dot_pair (a INTEGER, b INTEGER, label TEXT)" + let payload = Only (7 :: Int) :. (8 :: Int, Text.pack "hi") + _ <- execute conn "INSERT INTO dot_pair VALUES (?, ?, ?)" payload + rows <- query_ conn "SELECT a, b, label FROM dot_pair" :: IO [Only Int :. (Int, Text.Text)] + case rows of + [Only a :. (b, label)] -> do + assertEqual "(:.) first" 7 a + assertEqual "(:.) second" 8 b + assertEqual "(:.) label" (Text.pack "hi") label + other -> assertFailure ("unexpected (:.) rows: " <> show other) + , testCase "query_ fetches rows" $ + withConnection ":memory:" \conn -> do + _ <- execute_ conn "CREATE TABLE query_rows (a INTEGER, b TEXT)" + _ <- executeMany conn "INSERT INTO query_rows VALUES (?, ?)" [(1 :: Int, "x" :: String), (2 :: Int, "y" :: String)] + rows <- query_ conn "SELECT a, b FROM query_rows ORDER BY a" + assertEqual "query rows" [(1 :: Int, "x" :: String), (2 :: Int, "y" :: String)] rows + , testCase "query decodes NULL as Maybe" $ + withConnection ":memory:" \conn -> do + _ <- execute_ conn "CREATE TABLE maybe_vals (a TEXT)" + _ <- execute conn "INSERT INTO maybe_vals VALUES (?)" (Only (Just ("present" :: String))) + _ <- execute conn "INSERT INTO maybe_vals VALUES (?)" (Only (Nothing :: Maybe String)) + rows <- query_ conn "SELECT a FROM maybe_vals ORDER BY a IS NULL, a" + assertEqual "maybe decoding" [Only (Just ("present" :: String)), Only Nothing] rows + , testCase "query fetches rows with parameters" $ + withConnection ":memory:" \conn -> do + _ <- execute_ conn "CREATE TABLE query_params (a INTEGER, b TEXT)" + _ <- executeMany conn "INSERT INTO query_params VALUES (?, ?)" [(1 :: Int, "x" :: String), (2 :: Int, "y" :: String)] + rows <- query conn "SELECT a FROM query_params WHERE b = ?" (Only ("y" :: String)) + assertEqual "query result" [Only (2 :: Int)] rows + , testCase "column mismatch surfaces as SQLError" $ + withConnection ":memory:" \conn -> do + _ <- execute_ conn "CREATE TABLE mismatch (a INTEGER, b INTEGER)" + _ <- execute_ conn "INSERT INTO mismatch VALUES (1, 2)" + assertThrows + (query_ conn "SELECT a, b FROM mismatch" :: IO [Only Int]) + (Text.isInfixOf "column index" . sqlErrorMessage) + , testCase "clears statement bindings without error" $ + withConnection ":memory:" \conn -> do + stmt <- openStatement conn "SELECT ?" + clearStatementBindings stmt + closeStatement stmt + , testCase "resolves named parameter indices" $ + withConnection ":memory:" \conn -> do + stmt <- openStatement conn "SELECT $named_param" + idx <- namedParameterIndex stmt "named_param" + assertEqual "named parameter index" (Just 1) idx + idxWithPrefix <- namedParameterIndex stmt "$named_param" + assertEqual "named parameter index with prefix" (Just 1) idxWithPrefix + colonIdx <- namedParameterIndex stmt ":named_param" + assertEqual "named parameter index with colon" (Just 1) colonIdx + none <- namedParameterIndex stmt "missing" + assertEqual "missing parameter index" Nothing none + closeStatement stmt + , testCase "reports column metadata for statements" $ + withConnection ":memory:" \conn -> do + stmt <- openStatement conn "SELECT 1 AS a, 2 AS b" + count <- columnCount stmt + assertEqual "column count" 2 count + name0 <- columnName stmt 0 + name1 <- columnName stmt 1 + assertEqual "column names" [Text.pack "a", Text.pack "b"] [name0, name1] + assertThrows + (columnName stmt 2) + (Text.isInfixOf "out of bounds" . sqlErrorMessage) + closeStatement stmt + , testCase "bindNamed rejects statements without named placeholders" $ + withConnection ":memory:" \conn -> do + stmt <- openStatement conn "SELECT ?" + assertThrowsFormatError + (bindNamed stmt [":value" := (1 :: Int)]) + (Text.isInfixOf "does not define named parameters" . formatErrorMessage) + closeStatement stmt + , testCase "bindNamed rejects unknown parameters" $ + withConnection ":memory:" \conn -> do + stmt <- openStatement conn "SELECT $known" + assertThrowsFormatError + (bindNamed stmt ["$missing" := (1 :: Int)]) + (Text.isInfixOf "unknown named parameter" . formatErrorMessage) + closeStatement stmt + , testCase "numFieldsRemaining reports remaining columns" $ + withConnection ":memory:" \conn -> do + _ <- execute_ conn "CREATE TABLE remaining (a INTEGER, b INTEGER, c INTEGER)" + _ <- execute conn "INSERT INTO remaining VALUES (?, ?, ?)" (1 :: Int, 2 :: Int, 3 :: Int) + rows <- query_ conn "SELECT a, b, c FROM remaining" :: IO [WithRemaining] + assertEqual "remaining count" [WithRemaining 1 2] rows + , testCase "RowParser alternatives fall back" $ + withConnection ":memory:" \conn -> do + _ <- execute_ conn "CREATE TABLE yesno (answer TEXT)" + _ <- executeMany conn "INSERT INTO yesno VALUES (?)" [Only ("yes" :: String), Only ("no" :: String)] + rows <- query_ conn "SELECT answer FROM yesno ORDER BY answer" :: IO [YesNo] + assertEqual "yes/no parsing" [YesNo False, YesNo True] rows + ] + +typeTests :: TestTree +typeTests = + testGroup + "types" + [ testCase "round-trips date/time/timestamp" $ + withConnection ":memory:" \conn -> do + _ <- execute_ conn "CREATE TABLE temporals (d DATE, t TIME, ts TIMESTAMP)" + let dayVal = fromGregorian 2024 10 12 + timeVal = TimeOfDay 14 30 15.123456 + tsVal = LocalTime dayVal timeVal + _ <- execute conn "INSERT INTO temporals VALUES (?, ?, ?)" (dayVal, timeVal, tsVal) + [(dRes, tRes, tsRes)] <- query_ conn "SELECT d, t, ts FROM temporals" + assertEqual "date round-trip" dayVal dRes + assertEqual "time round-trip" timeVal tRes + assertEqual "timestamp round-trip" tsVal tsRes + [Only utcRes] <- query_ conn "SELECT ts FROM temporals" :: IO [Only UTCTime] + assertEqual "timestamp as UTC" (localTimeToUTC utc tsVal) utcRes + , testCase "round-trips blob payloads" $ + withConnection ":memory:" \conn -> do + _ <- execute_ conn "CREATE TABLE blobs (payload BLOB)" + let payload = BS.pack [0, 1, 2, 3, 255] + _ <- execute conn "INSERT INTO blobs VALUES (?)" (Only payload) + [Only blobOut] <- query_ conn "SELECT payload FROM blobs" + assertEqual "blob round-trip" payload blobOut + , testCase "custom FieldParser enforces invariants" $ + withConnection ":memory:" \conn -> do + _ <- execute_ conn "CREATE TABLE nonempty (name TEXT)" + _ <- execute conn "INSERT INTO nonempty VALUES (?)" (Only (Text.pack "okay")) + rows <- query_ conn "SELECT name FROM nonempty" :: IO [Only NonEmptyText] + assertEqual "non-empty success" [Only (NonEmptyText (Text.pack "okay"))] rows + _ <- execute conn "INSERT INTO nonempty VALUES (?)" (Only (Text.pack "")) + assertThrows + (query_ conn "SELECT name FROM nonempty" :: IO [Only NonEmptyText]) + (Text.isInfixOf "NonEmptyText requires a non-empty string" . sqlErrorMessage) + ] + +streamingTests :: TestTree +streamingTests = + testGroup + "streaming" + [ testCase "fold_ accumulates large result sets" $ + withConnection ":memory:" \conn -> do + _ <- execute_ conn "CREATE TABLE stream_fold (n INTEGER)" + let rows = fmap Only [1 .. 1000 :: Int] + _ <- executeMany conn "INSERT INTO stream_fold VALUES (?)" rows + total <- fold_ conn "SELECT n FROM stream_fold ORDER BY n" 0 \acc (Only n) -> pure (acc + n) + assertEqual "folded sum" (sum ([1 .. 1000] :: [Int])) total + , testCase "fold and foldNamed respect parameters" $ + withConnection ":memory:" \conn -> do + _ <- execute_ conn "CREATE TABLE stream_filter (n INTEGER)" + let values = fmap Only [1 .. 50 :: Int] + _ <- executeMany conn "INSERT INTO stream_filter VALUES (?)" values + gtTotal <- + fold conn "SELECT n FROM stream_filter WHERE n > ?" (Only (40 :: Int)) 0 $ + \acc (Only n) -> pure (acc + n) + let expected = sum ([41 .. 50] :: [Int]) + assertEqual "filtered fold sum" expected gtTotal + leCount <- + foldNamed conn "SELECT n FROM stream_filter WHERE n <= $limit" ["$limit" := (10 :: Int)] 0 $ + \acc (Only (_ :: Int)) -> pure (acc + 1) + assertEqual "foldNamed count" (10 :: Int) leCount + , testCase "nextRow streams rows sequentially" $ + withConnection ":memory:" \conn -> do + _ <- execute_ conn "CREATE TABLE stream_cursor (n INTEGER)" + _ <- executeMany conn "INSERT INTO stream_cursor VALUES (?)" (fmap Only [1 .. 3 :: Int]) + withStatement conn "SELECT n FROM stream_cursor ORDER BY n" \stmt -> do + first <- nextRow stmt + second <- nextRow stmt + third <- nextRow stmt + done <- nextRow stmt + assertEqual + "cursor iteration" + [Just (Only (1 :: Int)), Just (Only 2), Just (Only 3), Nothing] + [first, second, third, done] + ] + +functionsTests :: TestTree +functionsTests = + testGroup + "user-defined functions" + [ testCase "registers pure scalar function" $ + withConnection ":memory:" \conn -> do + createFunction conn "hs_times_two" (\(x :: Int64) -> x * 2) + result <- query_ conn "SELECT hs_times_two(21)" :: IO [Only Int64] + assertEqual "times_two result" [Only 42] result + , testCase "handles nullable arguments and results" $ + withConnection ":memory:" \conn -> do + createFunction conn "hs_optional" (\(mx :: Maybe Int64) -> fmap (+ 1) mx) + rows <- + query_ conn "SELECT hs_optional(x) FROM (VALUES (NULL), (41)) AS t(x)" :: + IO [Only (Maybe Int64)] + assertEqual "optional results" [Only Nothing, Only (Just 42)] rows + , testCase "supports IO-based functions" $ + withConnection ":memory:" \conn -> do + ref <- newIORef (0 :: Int) + createFunction conn "hs_counter" $ do + atomicModifyIORef' ref \n -> + let next = n + 1 + in (next, next) + first <- query_ conn "SELECT hs_counter()" :: IO [Only Int] + second <- query_ conn "SELECT hs_counter()" :: IO [Only Int] + assertEqual "first counter call" [Only 1] first + assertEqual "second counter call" [Only 2] second + ] + +transactionTests :: TestTree +transactionTests = + testGroup + "transactions" + [ testCase "withTransaction commits on success" $ + withConnection ":memory:" \conn -> do + _ <- execute_ conn "CREATE TABLE tx_commit (x INTEGER)" + withTransaction conn $ do + _ <- execute conn "INSERT INTO tx_commit VALUES (?)" (Only (1 :: Int)) + pure () + rows <- query_ conn "SELECT COUNT(*) FROM tx_commit" :: IO [Only Int] + assertEqual "commit" [Only 1] rows + , testCase "withTransaction rolls back on exception" $ + withConnection ":memory:" \conn -> do + _ <- execute_ conn "CREATE TABLE tx_rollback (x INTEGER)" + assertThrowsErrorCall $ + withTransaction conn $ do + _ <- execute conn "INSERT INTO tx_rollback VALUES (?)" (Only (1 :: Int)) + error "boom" + rows <- query_ conn "SELECT COUNT(*) FROM tx_rollback" :: IO [Only Int] + assertEqual "rollback" [Only 0] rows + ] + +assertThrows :: forall e a. (Exception e) => IO a -> (e -> Bool) -> Assertion +assertThrows action predicate = do + outcome <- try action + case outcome of + Left (err :: e) -> assertBool "unexpected exception" (predicate err) + Right _ -> assertFailure "expected exception, but action succeeded" + +assertThrowsSQLError :: IO a -> Assertion +assertThrowsSQLError action = + assertThrows action (const True :: SQLError -> Bool) + +assertThrowsErrorCall :: IO a -> Assertion +assertThrowsErrorCall action = + assertThrows action (const True :: ErrorCall -> Bool) + +assertThrowsFormatError :: IO a -> (FormatError -> Bool) -> Assertion +assertThrowsFormatError = assertThrows