diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,33 @@ # Revision history for beam-duckdb +## 0.3.1.0 -- unreleased + +* Added support for `beam-migrate`: the new `Database.Beam.DuckDB.Migrate` + module exposes a `migrationBackend :: BeamMigrationBackend DuckDB DuckDBM` + value, along with other utilities. This brings DuckDB on par with `beam-postgres` + and `beam-sqlite` for schema verification and migration script generation. +* Added file-mode `COPY ... TO 'file'` / `COPY ... FROM 'file'` support + via the new `MonadBeamCopyTo` / `MonadBeamCopyFrom` instances on + `DuckDBM`. Smart constructors `copyToCSV` / `copyToParquet` / + `copyToJSON` (and the symmetric `copyFrom*` plus `*With` variants) + build the per-format options records. See + `Database.Beam.DuckDB.Syntax.Extensions.Copy`. +* Added instances for `MonadBeamInsertReturning` / `MonadBeamUpdateReturning` / `MonadBeamDeleteReturning`; +* Added a `BeamHasInsertOnConflict DuckDB` instance, exposing + `insertOnConflict` (and the `anyConflict` / `conflictingFields` / + `onConflictDoNothing` / `onConflictUpdateSet` / + `onConflictUpdateSetWhere` / `onConflictUpdateAll` / + `onConflictUpdateInstead` combinators) for `INSERT ... ON CONFLICT` + on DuckDB. DuckDB does not support partial-index conflict targets + ([INSERT docs](https://duckdb.org/docs/stable/sql/statements/insert), + so `conflictingFieldsWhere` is *not* re-exported from + `Database.Beam.DuckDB`; the `conflictingFieldsWhere` in that module + is a shadowing shim that produces a compile-time `TypeError` + pointing users at `onConflictUpdateSetWhere`. Importing + `conflictingFieldsWhere` directly from + `Database.Beam.Backend.SQL.BeamExtensions` still type-checks but + raises an error at runtime. + ## 0.3.0.0 -- 2026-04-28 * Added a `weekField` definition to the `IsSql92ExtractFieldSyntax` instance diff --git a/beam-duckdb.cabal b/beam-duckdb.cabal --- a/beam-duckdb.cabal +++ b/beam-duckdb.cabal @@ -1,6 +1,6 @@ cabal-version: 3.0 name: beam-duckdb -version: 0.3.0.0 +version: 0.3.1.0 synopsis: DuckDB backend for Beam description: Beam driver for DuckDB, an analytics-focused open-source in-process database. license: MIT @@ -36,15 +36,22 @@ -Wredundant-constraints -fhide-source-paths -Wpartial-fields + -Wunused-packages library import: common exposed-modules: Database.Beam.DuckDB + Database.Beam.DuckDB.Migrate other-modules: Database.Beam.DuckDB.Backend + Database.Beam.DuckDB.Connection Database.Beam.DuckDB.Syntax Database.Beam.DuckDB.Syntax.Builder + Database.Beam.DuckDB.Syntax.Extensions + Database.Beam.DuckDB.Syntax.Extensions.Copy Database.Beam.DuckDB.Syntax.Extensions.DataSource - build-depends: base >=4.11 && <5 - , beam-core >=0.11 && <0.12 + Database.Beam.DuckDB.Syntax.Extensions.InsertOnConflict + build-depends: aeson >=1.0 && <2.3 + , base >=4.11 && <5 + , beam-core >=0.11.1 && <0.12 , beam-migrate ^>=0.6 , bytestring >=0.10 && <0.13 -- Version 0.1.5.0 contained a broken change (RowParser constructor not exposed) @@ -52,9 +59,10 @@ , duckdb-simple (>=0.1.2 && <0.1.5) || (>=0.1.5.1 && <0.1.6) , dlist >=0.8 && <1.1 , free >=4.12 && <5.3 + , hashable >=1.1 && <1.6 , scientific ^>=0.3 , text >=1.0 && <2.2 - , time >=1.6 && <1.16 + , time >=1.6 && <1.17 , transformers >=0.3 && <0.7 , uuid-types ^>=1.0 hs-source-dirs: src @@ -62,6 +70,11 @@ test-suite beam-duckdb-test import: common other-modules: Database.Beam.DuckDB.Test.Extensions + Database.Beam.DuckDB.Test.Extensions.Copy + Database.Beam.DuckDB.Test.Extensions.DataSource + Database.Beam.DuckDB.Test.Extensions.InsertOnConflict + Database.Beam.DuckDB.Test.Extensions.Returning + Database.Beam.DuckDB.Test.Migrate Database.Beam.DuckDB.Test.Query type: exitcode-stdio-1.0 hs-source-dirs: tests @@ -69,10 +82,14 @@ build-depends: base >=4.11 && <5 , beam-core , beam-duckdb + , beam-migrate + , bytestring , duckdb-simple , hedgehog , tasty , tasty-hedgehog , tasty-hunit + , temporary , text , time + , uuid-types diff --git a/src/Database/Beam/DuckDB.hs b/src/Database/Beam/DuckDB.hs --- a/src/Database/Beam/DuckDB.hs +++ b/src/Database/Beam/DuckDB.hs @@ -1,13 +1,7 @@ -{-# LANGUAGE DerivingStrategies #-} -{-# LANGUAGE ExplicitForAll #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE GeneralizedNewtypeDeriving #-} -{-# LANGUAGE LambdaCase #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE UndecidableInstances #-} -- | DuckDB is a powerful in-process database specialized in analytical processing workloads. -- @@ -29,6 +23,35 @@ DuckDB, DuckDBM, + -- * @RETURNING@ support + + -- | The 'MonadBeamInsertReturning', 'MonadBeamUpdateReturning', and + -- 'MonadBeamDeleteReturning' instances let DuckDB users capture the + -- newly-inserted, updated, or deleted rows of a DML statement, which is + -- useful for retrieving defaulted/serial column values. + MonadBeamInsertReturning (..), + runInsertReturningList, + MonadBeamUpdateReturning (..), + runUpdateReturningList, + MonadBeamDeleteReturning (..), + runDeleteReturningList, + + -- * @INSERT ... ON CONFLICT@ support + + -- | The 'BeamHasInsertOnConflict' instance for DuckDB exposes + -- @INSERT ... ON CONFLICT@ via 'insertOnConflict', allowing rows that + -- violate a uniqueness constraint to be ignored + -- ('onConflictDoNothing') or replaced with new values + -- ('onConflictUpdateSet', 'onConflictUpdateAll', + -- 'onConflictUpdateInstead'). + -- + -- DuckDB does not support partial-index conflict targets, so the + -- 'Database.Beam.Backend.SQL.BeamExtensions.conflictingFieldsWhere' + -- class method is /not/ re-exported here. The 'conflictingFieldsWhere' + -- in scope is a compile-time-blocked shim that produces a 'TypeError' + -- when used — pointing users at 'onConflictUpdateSetWhere' instead. + module Database.Beam.DuckDB.Syntax.Extensions.InsertOnConflict, + -- * DuckDB-specific functionality -- ** Data sources @@ -52,175 +75,90 @@ csvWith, CSVOptions (..), defaultCSVOptions, + + -- ** COPY support + DuckDBCopyToOptions, + DuckDBCopyFromOptions, + + -- *** CSV + copyToCSV, + copyToCSVWith, + DuckDBCSVCopyToOptions (..), + defaultDuckDBCSVCopyToOptions, + copyFromCSV, + copyFromCSVWith, + DuckDBCSVCopyFromOptions (..), + defaultDuckDBCSVCopyFromOptions, + + -- *** Parquet + copyToParquet, + copyToParquetWith, + DuckDBParquetCopyToOptions (..), + ParquetCompression (..), + defaultDuckDBParquetCopyToOptions, + copyFromParquet, + copyFromParquetWith, + DuckDBParquetCopyFromOptions (..), + defaultDuckDBParquetCopyFromOptions, + + -- *** JSON + copyToJSON, + copyToJSONWith, + DuckDBJSONCopyToOptions (..), + JSONCompression (..), + defaultDuckDBJSONCopyToOptions, + copyFromJSON, + copyFromJSONWith, + DuckDBJSONCopyFromOptions (..), + defaultDuckDBJSONCopyFromOptions, ) where -import Control.Exception (SomeException (..)) -import Control.Monad (void) -import Control.Monad.Free.Church (runF) -import Control.Monad.IO.Class (MonadIO (..)) -import Control.Monad.Trans.Class (lift) -import Control.Monad.Trans.Reader (ReaderT (..), ask) -import Control.Monad.Trans.State.Strict (StateT (..), get, put) -import qualified Data.DList as DL -import Data.Data (cast) -import Data.Maybe (mapMaybe) -import Data.Text (Text) -import qualified Data.Text as Text -import qualified Data.Text.Lazy as Text.Lazy -import qualified Data.Text.Lazy.Builder as Builder -import Database.Beam.Backend - ( BeamRowReadError (..), - ColumnParseError (..), - FromBackendRow, - FromBackendRowF (..), - FromBackendRowM (..), - MonadBeam, +import Database.Beam.Backend.SQL.BeamExtensions + ( MonadBeamDeleteReturning (..), + MonadBeamInsertReturning (..), + MonadBeamUpdateReturning (..), + runDeleteReturningList, + runInsertReturningList, + runUpdateReturningList, ) -import Database.Beam.Backend.SQL (FromBackendRow (..), MonadBeam (..)) import Database.Beam.DuckDB.Backend (DuckDB) -import Database.Beam.DuckDB.Syntax - ( DuckDBCommandSyntax (..), +import Database.Beam.DuckDB.Connection + ( DuckDBM, + runBeamDuckDB, + runBeamDuckDBDebug, + runBeamDuckDBDebugString, ) -import Database.Beam.DuckDB.Syntax.Builder - ( DuckDBSyntax (..), - SomeField (..), - withPlaceholder, +import Database.Beam.DuckDB.Syntax.Extensions () +import Database.Beam.DuckDB.Syntax.Extensions.Copy + ( DuckDBCSVCopyFromOptions (..), + DuckDBCSVCopyToOptions (..), + DuckDBCopyFromOptions, + DuckDBCopyToOptions, + DuckDBJSONCopyFromOptions (..), + DuckDBJSONCopyToOptions (..), + DuckDBParquetCopyFromOptions (..), + DuckDBParquetCopyToOptions (..), + JSONCompression (..), + ParquetCompression (..), + copyFromCSV, + copyFromCSVWith, + copyFromJSON, + copyFromJSONWith, + copyFromParquet, + copyFromParquetWith, + copyToCSV, + copyToCSVWith, + copyToJSON, + copyToJSONWith, + copyToParquet, + copyToParquetWith, + defaultDuckDBCSVCopyFromOptions, + defaultDuckDBCSVCopyToOptions, + defaultDuckDBJSONCopyFromOptions, + defaultDuckDBJSONCopyToOptions, + defaultDuckDBParquetCopyFromOptions, + defaultDuckDBParquetCopyToOptions, ) import Database.Beam.DuckDB.Syntax.Extensions.DataSource -import Database.DuckDB.Simple (Connection, FromRow, Query (Query), ResultError (..), RowParser, ToField (toField), ToRow (toRow), bind, execute, nextRow, withStatement) -import Database.DuckDB.Simple.FromRow (FromRow (..), RowParser (..), field) -import Database.DuckDB.Simple.Ok (Ok (..)) - --- | 'MonadBeam' instance inside which DuckDB queries are run. See the --- for more information -newtype DuckDBM a - = DuckDBM - { -- | Run an IO action with access to a DuckDB connection and a debug logging - -- function, called or each query submitted on the connection. - runDuckDBM :: ReaderT (Text -> IO (), Connection) IO a - } - deriving (Monad, Functor, Applicative, MonadIO, MonadFail) - --- | Execute queries on a DuckDB 'Connection'. --- --- To trace the exact SQL statements being generated, --- use 'runBeamDuckDBDebug' instead. -runBeamDuckDB :: Connection -> DuckDBM a -> IO a -runBeamDuckDB = runBeamDuckDBDebug (\_ -> pure ()) - -runBeamDuckDBDebug :: (Text -> IO ()) -> Connection -> DuckDBM a -> IO a -runBeamDuckDBDebug debug conn action = - runReaderT (runDuckDBM action) (debug, conn) - --- | Like 'runBeamDuckDBDebug', but accepts a 'String' argument instead of 'Text'. --- --- This is provided for compatibility with other backends -runBeamDuckDBDebugString :: (String -> IO ()) -> Connection -> DuckDBM a -> IO a -runBeamDuckDBDebugString debug = runBeamDuckDBDebug (debug . Text.unpack) - -newtype BeamDuckDBParams = BeamDuckDBParams [SomeField] - -instance ToRow BeamDuckDBParams where - toRow (BeamDuckDBParams x) = map (\(SomeField f) -> toField f) x - -instance MonadBeam DuckDB DuckDBM where - runNoReturn (DuckDBCommandSyntax (DuckDBSyntax cmd vals)) = - DuckDBM $ do - (logger, conn) <- ask - let cmdString = Text.Lazy.toStrict (Builder.toLazyText (withPlaceholder cmd)) - liftIO (logger (cmdString <> ";\n-- With values: " <> Text.pack (show (DL.toList vals)))) - liftIO (void $ execute conn (Query cmdString) (BeamDuckDBParams (DL.toList vals))) - - runReturningMany (DuckDBCommandSyntax (DuckDBSyntax cmd vals)) action = - DuckDBM $ do - (logger, conn) <- ask - let cmdString = Text.Lazy.toStrict (Builder.toLazyText (withPlaceholder cmd)) - liftIO $ do - logger (cmdString <> ";\n-- With values: " <> Text.pack (show (DL.toList vals))) - withStatement conn (Query cmdString) $ \stmt -> - do - bind stmt (map toField (DL.toList vals)) - let nextRow' = - liftIO (nextRow stmt) >>= \case - Nothing -> pure Nothing - Just (BeamDuckDBRow row) -> pure row - runReaderT (runDuckDBM (action nextRow')) (logger, conn) - -newtype BeamDuckDBRow a = BeamDuckDBRow a - -instance (FromBackendRow DuckDB a) => FromRow (BeamDuckDBRow a) where - fromRow = BeamDuckDBRow <$> runF fromBackendRow' finish step - where - FromBackendRowM fromBackendRow' = fromBackendRow - - translateErrors :: Maybe Int -> SomeException -> Maybe SomeException - translateErrors col (SomeException e) = - case cast e of - Just - ( ConversionFailed - { errSQLType = typeString, - errHaskellType = hsString, - errMessage = msg - } - ) -> - Just - ( SomeException - ( BeamRowReadError - col - ( ColumnTypeMismatch - (Text.unpack hsString) - (Text.unpack typeString) - ("conversion failed: " ++ Text.unpack msg) - ) - ) - ) - Just (UnexpectedNull {}) -> - Just (SomeException (BeamRowReadError {brreColumn = col, brreError = ColumnUnexpectedNull})) - Just - ( Incompatible - { errSQLType = typeString, - errHaskellType = hsString, - errMessage = msg - } - ) -> - Just - ( SomeException - ( BeamRowReadError - col - ( ColumnTypeMismatch - (Text.unpack hsString) - (Text.unpack typeString) - ("incompatible: " ++ Text.unpack msg) - ) - ) - ) - Nothing -> Nothing - - finish = pure - - step :: forall a'. FromBackendRowF DuckDB (RowParser a') -> RowParser a' - step (ParseOneField next) = - RowParser $ ReaderT $ \ro -> StateT $ \st@(col, _) -> - case runStateT (runReaderT (runRowParser field) ro) st of - Ok (x, st') -> runStateT (runReaderT (runRowParser (next x)) ro) st' - Errors errs -> Errors (mapMaybe (translateErrors (Just col)) errs) - step (Alt (FromBackendRowM a) (FromBackendRowM b) next) = do - RowParser $ do - let RowParser a' = runF a finish step - RowParser b' = runF b finish step - - st <- lift get - ro <- ask - case runStateT (runReaderT a' ro) st of - Ok (ra, st') -> do - lift $ put st' - runRowParser (next ra) - Errors aErrs -> - case runStateT (runReaderT b' ro) st of - Ok (rb, st') -> do - lift $ put st' - runRowParser (next rb) - Errors bErrs -> - lift (lift (Errors (aErrs ++ bErrs))) - step (FailParseWith err) = RowParser (lift (lift (Errors [SomeException err]))) +import Database.Beam.DuckDB.Syntax.Extensions.InsertOnConflict diff --git a/src/Database/Beam/DuckDB/Backend.hs b/src/Database/Beam/DuckDB/Backend.hs --- a/src/Database/Beam/DuckDB/Backend.hs +++ b/src/Database/Beam/DuckDB/Backend.hs @@ -17,6 +17,7 @@ import Data.Int (Int16, Int32, Int64, Int8) import Data.Scientific (Scientific, scientific) import Data.Text (Text) +import qualified Data.Text as Text import Data.Time (Day, LocalTime, TimeOfDay, UTCTime) import Data.UUID.Types (UUID) import Data.Word (Word16, Word32, Word64, Word8) @@ -31,6 +32,7 @@ parseOneField, ) import Database.Beam.Backend.SQL (FromBackendRow (..)) +import Database.Beam.Backend.SQL.BeamExtensions (BeamSqlBackendCopyFromSyntax, BeamSqlBackendCopyToSyntax) import Database.Beam.DuckDB.Syntax ( DuckDBCommandSyntax, DuckDBExpressionSyntax (..), @@ -40,6 +42,8 @@ emit, parens, ) +import Database.Beam.DuckDB.Syntax.Extensions.Copy (DuckDBCopyFromSyntax, DuckDBCopyToSyntax) +import Database.Beam.Migrate.SQL (BeamMigrateOnlySqlBackend) import Database.Beam.Query.SQL92 (buildSql92Query') import Database.Beam.Query.Types (HasQBuilder (..)) import Database.DuckDB.Simple (Null) @@ -51,6 +55,12 @@ instance BeamSqlBackend DuckDB +instance BeamMigrateOnlySqlBackend DuckDB + +type instance BeamSqlBackendCopyToSyntax DuckDB = DuckDBCopyToSyntax + +type instance BeamSqlBackendCopyFromSyntax DuckDB = DuckDBCopyFromSyntax + instance HasQBuilder DuckDB where buildSqlQuery = buildSql92Query' True @@ -102,6 +112,13 @@ instance FromBackendRow DuckDB Word64 instance FromBackendRow DuckDB Text + +instance FromBackendRow DuckDB Char where + fromBackendRow = do + t <- fromBackendRow @DuckDB @Text + case Text.uncons t of + Just (c, _) -> pure c + Nothing -> fail "Need string of size one to parse Char" instance FromBackendRow DuckDB ByteString diff --git a/src/Database/Beam/DuckDB/Connection.hs b/src/Database/Beam/DuckDB/Connection.hs new file mode 100644 --- /dev/null +++ b/src/Database/Beam/DuckDB/Connection.hs @@ -0,0 +1,210 @@ +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE ExplicitForAll #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | Internal module providing the 'DuckDBM' monad and core 'MonadBeam' +-- instances for the DuckDB backend. Re-exported from "Database.Beam.DuckDB". +module Database.Beam.DuckDB.Connection + ( DuckDBM (..), + runBeamDuckDB, + runBeamDuckDBDebug, + runBeamDuckDBDebugString, + liftIOWithHandle, + BeamDuckDBParams (..), + BeamDuckDBRow (..), + ) +where + +import Control.Exception (SomeException (..)) +import Control.Monad (void) +import Control.Monad.Free.Church (runF) +import Control.Monad.IO.Class (MonadIO (..)) +import Control.Monad.Trans.Class (lift) +import Control.Monad.Trans.Reader (ReaderT (..), ask) +import Control.Monad.Trans.State.Strict (StateT (..), get, put) +import Data.Coerce (coerce) +import qualified Data.DList as DL +import Data.Data (cast) +import Data.Maybe (mapMaybe) +import Data.Text (Text) +import qualified Data.Text as Text +import qualified Data.Text.Lazy as Text.Lazy +import qualified Data.Text.Lazy.Builder as Builder +import Database.Beam.Backend + ( BeamRowReadError (..), + ColumnParseError (..), + FromBackendRow, + FromBackendRowF (..), + FromBackendRowM (..), + MonadBeam, + ) +import Database.Beam.Backend.SQL (FromBackendRow (..), MonadBeam (..)) +import Database.Beam.Backend.SQL.BeamExtensions (MonadBeamCopyFrom (..), MonadBeamCopyTo (..), SqlCopyFrom (..), SqlCopyTo (..)) +import Database.Beam.DuckDB.Backend (DuckDB) +import Database.Beam.DuckDB.Syntax (DuckDBCommandSyntax (..)) +import Database.Beam.DuckDB.Syntax.Builder + ( DuckDBSyntax (..), + SomeField (..), + withPlaceholder, + ) +import Database.Beam.DuckDB.Syntax.Extensions.Copy + ( DuckDBCopyFromSyntax (..), + DuckDBCopyToSyntax (..), + ) +import Database.DuckDB.Simple (Connection, FromRow, Query (Query), ResultError (..), RowParser, ToField (toField), ToRow (toRow), bind, execute, nextRow, withStatement) +import Database.DuckDB.Simple.FromRow (FromRow (..), RowParser (..), field) +import Database.DuckDB.Simple.Ok (Ok (..)) + +-- | 'MonadBeam' instance inside which DuckDB queries are run. See the +-- for more information +newtype DuckDBM a + = DuckDBM + { -- | Run an IO action with access to a DuckDB connection and a debug logging + -- function, called or each query submitted on the connection. + runDuckDBM :: ReaderT (Text -> IO (), Connection) IO a + } + deriving (Monad, Functor, Applicative, MonadIO, MonadFail) + +-- | Execute queries on a DuckDB 'Connection'. +-- +-- To trace the exact SQL statements being generated, +-- use 'runBeamDuckDBDebug' instead. +runBeamDuckDB :: Connection -> DuckDBM a -> IO a +runBeamDuckDB = runBeamDuckDBDebug (\_ -> pure ()) + +runBeamDuckDBDebug :: (Text -> IO ()) -> Connection -> DuckDBM a -> IO a +runBeamDuckDBDebug debug conn action = + runReaderT (runDuckDBM action) (debug, conn) + +-- | Like 'runBeamDuckDBDebug', but accepts a 'String' argument instead of 'Text'. +-- +-- This is provided for compatibility with other backends +runBeamDuckDBDebugString :: (String -> IO ()) -> Connection -> DuckDBM a -> IO a +runBeamDuckDBDebugString debug = runBeamDuckDBDebug (debug . Text.unpack) + +-- | Run an IO action with access to the underlying DuckDB 'Connection'. +-- This is mostly useful for migrate-backend implementors who need to issue +-- raw queries against the connection. +liftIOWithHandle :: (Connection -> IO a) -> DuckDBM a +liftIOWithHandle f = DuckDBM $ ReaderT $ \(_, conn) -> f conn + +newtype BeamDuckDBParams = BeamDuckDBParams [SomeField] + +instance ToRow BeamDuckDBParams where + toRow (BeamDuckDBParams x) = map (\(SomeField f) -> toField f) x + +instance MonadBeam DuckDB DuckDBM where + runNoReturn (DuckDBCommandSyntax (DuckDBSyntax cmd vals)) = + DuckDBM $ do + (logger, conn) <- ask + let cmdString = Text.Lazy.toStrict (Builder.toLazyText (withPlaceholder cmd)) + liftIO (logger (cmdString <> ";\n-- With values: " <> Text.pack (show (DL.toList vals)))) + liftIO (void $ execute conn (Query cmdString) (BeamDuckDBParams (DL.toList vals))) + + runReturningMany (DuckDBCommandSyntax (DuckDBSyntax cmd vals)) action = + DuckDBM $ do + (logger, conn) <- ask + let cmdString = Text.Lazy.toStrict (Builder.toLazyText (withPlaceholder cmd)) + liftIO $ do + logger (cmdString <> ";\n-- With values: " <> Text.pack (show (DL.toList vals))) + withStatement conn (Query cmdString) $ \stmt -> + do + bind stmt (map toField (DL.toList vals)) + let nextRow' = + liftIO (nextRow stmt) >>= \case + Nothing -> pure Nothing + Just (BeamDuckDBRow row) -> pure row + runReaderT (runDuckDBM (action nextRow')) (logger, conn) + +instance MonadBeamCopyTo DuckDB DuckDBM where + runCopyTo SqlCopyToNoColumns = pure () + runCopyTo (SqlCopyTo syntax) = runNoReturn (coerce syntax) + +instance MonadBeamCopyFrom DuckDB DuckDBM where + runCopyFrom SqlCopyFromNoColumns = pure () + runCopyFrom (SqlCopyFrom syntax) = runNoReturn (coerce syntax) + +newtype BeamDuckDBRow a = BeamDuckDBRow a + +instance (FromBackendRow DuckDB a) => FromRow (BeamDuckDBRow a) where + fromRow = BeamDuckDBRow <$> runF fromBackendRow' finish step + where + FromBackendRowM fromBackendRow' = fromBackendRow + + translateErrors :: Maybe Int -> SomeException -> Maybe SomeException + translateErrors col (SomeException e) = + case cast e of + Just + ( ConversionFailed + { errSQLType = typeString, + errHaskellType = hsString, + errMessage = msg + } + ) -> + Just + ( SomeException + ( BeamRowReadError + col + ( ColumnTypeMismatch + (Text.unpack hsString) + (Text.unpack typeString) + ("conversion failed: " ++ Text.unpack msg) + ) + ) + ) + Just (UnexpectedNull {}) -> + Just (SomeException (BeamRowReadError {brreColumn = col, brreError = ColumnUnexpectedNull})) + Just + ( Incompatible + { errSQLType = typeString, + errHaskellType = hsString, + errMessage = msg + } + ) -> + Just + ( SomeException + ( BeamRowReadError + col + ( ColumnTypeMismatch + (Text.unpack hsString) + (Text.unpack typeString) + ("incompatible: " ++ Text.unpack msg) + ) + ) + ) + Nothing -> Nothing + + finish = pure + + step :: forall a'. FromBackendRowF DuckDB (RowParser a') -> RowParser a' + step (ParseOneField next) = + RowParser $ ReaderT $ \ro -> StateT $ \st@(col, _) -> + case runStateT (runReaderT (runRowParser field) ro) st of + Ok (x, st') -> runStateT (runReaderT (runRowParser (next x)) ro) st' + Errors errs -> Errors (mapMaybe (translateErrors (Just col)) errs) + step (Alt (FromBackendRowM a) (FromBackendRowM b) next) = do + RowParser $ do + let RowParser a' = runF a finish step + RowParser b' = runF b finish step + + st <- lift get + ro <- ask + case runStateT (runReaderT a' ro) st of + Ok (ra, st') -> do + lift $ put st' + runRowParser (next ra) + Errors aErrs -> + case runStateT (runReaderT b' ro) st of + Ok (rb, st') -> do + lift $ put st' + runRowParser (next rb) + Errors bErrs -> + lift (lift (Errors (aErrs ++ bErrs))) + step (FailParseWith err) = RowParser (lift (lift (Errors [SomeException err]))) diff --git a/src/Database/Beam/DuckDB/Migrate.hs b/src/Database/Beam/DuckDB/Migrate.hs new file mode 100644 --- /dev/null +++ b/src/Database/Beam/DuckDB/Migrate.hs @@ -0,0 +1,621 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE TypeApplications #-} + +-- | Migrations support for @beam-duckdb@. See "Database.Beam.Migrate" for +-- more information on beam migrations. +module Database.Beam.DuckDB.Migrate + ( -- * Top-level @beam-migrate@ backend + migrationBackend, + DuckDBCommandSyntax, + + -- * DuckDB-specific data types + + -- + -- DuckDB has a number of data types that aren't part of standard SQL or + -- that DuckDB names differently. These smart constructors produce + -- 'Db.DataType's tagged with @DuckDB@ so they can be used in + -- @beam-migrate@ schemas just like 'Database.Beam.Migrate.int' or + -- 'Database.Beam.Migrate.smallint'. + text, + blob, + tinyint, + uuid, + utinyint, + usmallint, + uinteger, + ubigint, + + -- * Conversions and utilities + getDbConstraints, + getDbConstraintsForSchemas, + migrateScript, + ) +where + +import Control.Exception (SomeException (..), catch) +import Control.Monad (forM, void) +import Control.Monad.Trans.Reader (ReaderT (..)) +import Data.Aeson (object, withText, (.=)) +import Data.ByteString (ByteString) +import qualified Data.ByteString.Lazy as BL +import Data.Char (isDigit) +import Data.Int (Int8) +import qualified Data.List.NonEmpty as NE +import Data.Maybe (fromMaybe, mapMaybe) +import Data.Monoid (Endo (..)) +import Data.Text (Text) +import qualified Data.Text as T +import qualified Data.Text.Encoding as TE +import Data.UUID.Types (UUID) +import Data.Word (Word16, Word32, Word64, Word8) +import Database.Beam.Backend.SQL +import Database.Beam.DuckDB.Backend (DuckDB) +import Database.Beam.DuckDB.Connection + ( DuckDBM (..), + liftIOWithHandle, + ) +import Database.Beam.DuckDB.Syntax + ( DuckDBCommandSyntax (..), + DuckDBDataTypeSyntax (..), + ) +import Database.Beam.DuckDB.Syntax.Builder + ( duckDBRenderSyntaxScript, + emit, + ) +import Database.Beam.Haskell.Syntax +import Database.Beam.Migrate.Actions + ( createIndexActionProvider, + defaultActionProvider, + defaultSchemaActionProvider, + dropIndexActionProvider, + ) +import qualified Database.Beam.Migrate.Backend as Tool +import qualified Database.Beam.Migrate.Checks as Db +import qualified Database.Beam.Migrate.SQL as Db +import qualified Database.Beam.Migrate.Serialization as Db +import qualified Database.Beam.Migrate.Types as Db +import qualified Database.Beam.Query.DataTypes as Db +import qualified Database.DuckDB.Simple as DuckDB +import Database.DuckDB.Simple.Types (Only (..)) + +-- | Top-level migration backend for use by @beam-migrate@ tools. +-- +-- @since 0.3.1.0 +migrationBackend :: Tool.BeamMigrationBackend DuckDB DuckDBM +migrationBackend = + Tool.BeamMigrationBackend + { Tool.backendName = "duckdb", + Tool.backendConnStringExplanation = + unlines + [ "For beam-duckdb, this is the path to a DuckDB database file, or :memory: for an in-memory database.", + "", + "See for more information." + ], + Tool.backendGetDbConstraints = getDbConstraints, + Tool.backendPredicateParsers = + Db.sql92Deserializers + <> Db.sql99DataTypeDeserializers + <> Db.sql2008BigIntDataTypeDeserializers + <> duckDBDataTypeDeserializers + <> Db.beamCheckDeserializers, + Tool.backendRenderSyntax = renderDuckDBCommand, + Tool.backendFileExtension = "duckdb.sql", + Tool.backendConvertToHaskell = duckDBPredConverter, + Tool.backendActionProvider = + mconcat + [ defaultActionProvider, + defaultSchemaActionProvider, + createIndexActionProvider, + dropIndexActionProvider + ], + Tool.backendRunSqlScript = \t -> liftIOWithHandle $ \conn -> + void $ DuckDB.execute_ conn (DuckDB.Query t), + Tool.backendWithTransaction = duckDBMWithTransaction, + Tool.backendConnect = \options -> do + conn <- DuckDB.open options + pure + Tool.BeamMigrateConnection + { Tool.backendRun = \(DuckDBM action) -> + catch + (Right <$> runReaderT action (\_ -> pure (), conn)) + (\e -> pure (Left (show (e :: SomeException)))), + Tool.backendClose = DuckDB.close conn + } + } + where + duckDBMWithTransaction :: forall a. DuckDBM a -> DuckDBM a + duckDBMWithTransaction (DuckDBM (ReaderT k)) = + DuckDBM $ ReaderT $ \ctx@(_, conn) -> + DuckDB.withTransaction conn (k ctx) + + renderDuckDBCommand :: DuckDBCommandSyntax -> String + renderDuckDBCommand (DuckDBCommandSyntax inner) = + T.unpack (duckDBRenderSyntaxScript inner) ++ ";" + + duckDBPredConverter :: Tool.HaskellPredicateConverter + duckDBPredConverter = + Tool.sql92HsPredicateConverters @DuckDB duckDBTypeToHs + <> Tool.hsPredicateConverter duckDBHasColumnConstraint + where + duckDBHasColumnConstraint + ( Db.TableColumnHasConstraint tblNm colNm c :: + Db.TableColumnHasConstraint DuckDB + ) + | c == Db.constraintDefinitionSyntax Nothing Db.notNullConstraintSyntax Nothing = + Just + ( Db.SomeDatabasePredicate + ( Db.TableColumnHasConstraint + tblNm + colNm + (Db.constraintDefinitionSyntax Nothing Db.notNullConstraintSyntax Nothing) :: + Db.TableColumnHasConstraint HsMigrateBackend + ) + ) + | otherwise = Nothing + + duckDBTypeToHs :: DuckDBDataTypeSyntax -> Maybe HsDataType + duckDBTypeToHs = Just . hsErrorType . T.unpack . duckDBRenderSyntaxScript . duckDBDataType + +-- | Build a backend-specific 'Db.BeamSerializedDataType' for DuckDB, tagged +-- with the given keyword. Used to give DuckDB-only types a stable serialized +-- representation that round-trips through @beam-migrate@'s predicate cache. +duckDBDataTypeJSON :: Text -> Db.BeamSerializedDataType +duckDBDataTypeJSON tag = + Db.BeamSerializedDataType (Db.beamSerializeJSON "duckdb" (object ["data-type" .= tag])) + +duckDBTinyIntType :: DuckDBDataTypeSyntax +duckDBTinyIntType = DuckDBDataTypeSyntax (emit "TINYINT") (duckDBDataTypeJSON "tinyint") + +duckDBHugeIntType :: DuckDBDataTypeSyntax +duckDBHugeIntType = DuckDBDataTypeSyntax (emit "HUGEINT") (duckDBDataTypeJSON "hugeint") + +duckDBUuidType :: DuckDBDataTypeSyntax +duckDBUuidType = DuckDBDataTypeSyntax (emit "UUID") (duckDBDataTypeJSON "uuid") + +duckDBUTinyIntType :: DuckDBDataTypeSyntax +duckDBUTinyIntType = DuckDBDataTypeSyntax (emit "UTINYINT") (duckDBDataTypeJSON "utinyint") + +duckDBUSmallIntType :: DuckDBDataTypeSyntax +duckDBUSmallIntType = DuckDBDataTypeSyntax (emit "USMALLINT") (duckDBDataTypeJSON "usmallint") + +duckDBUIntegerType :: DuckDBDataTypeSyntax +duckDBUIntegerType = DuckDBDataTypeSyntax (emit "UINTEGER") (duckDBDataTypeJSON "uinteger") + +duckDBUBigIntType :: DuckDBDataTypeSyntax +duckDBUBigIntType = DuckDBDataTypeSyntax (emit "UBIGINT") (duckDBDataTypeJSON "ubigint") + +duckDBUHugeIntType :: DuckDBDataTypeSyntax +duckDBUHugeIntType = DuckDBDataTypeSyntax (emit "UHUGEINT") (duckDBDataTypeJSON "uhugeint") + +-- | DuckDB @TEXT@. Equivalent to 'Database.Beam.Migrate.characterLargeObject' +-- and to @VARCHAR@ without a length limit; DuckDB reports all three as +-- @VARCHAR@ in @information_schema@. +-- +-- @since 0.3.1.0 +text :: Db.DataType DuckDB Text +text = Db.DataType (varCharType Nothing Nothing) + +-- | DuckDB @BLOB@ for arbitrary binary data. Equivalent to +-- 'Database.Beam.Migrate.binaryLargeObject' (DuckDB accepts both @BLOB@ and +-- @BYTEA@ in DDL but reports @BLOB@ in @information_schema@). +-- +-- @since 0.3.1.0 +blob :: Db.DataType DuckDB ByteString +blob = Db.DataType binaryLargeObjectType + +-- | DuckDB @TINYINT@ — a signed 8-bit integer. Not part of SQL standard +-- nor of beam-core (which only exposes 'smallint' upwards). +-- +-- @since 0.3.1.0 +tinyint :: Db.DataType DuckDB Int8 +tinyint = Db.DataType duckDBTinyIntType + +-- | DuckDB @UUID@. Maps to 'Data.UUID.Types.UUID' from @uuid-types@. +-- +-- @since 0.3.1.0 +uuid :: Db.DataType DuckDB UUID +uuid = Db.DataType duckDBUuidType + +-- | DuckDB @UTINYINT@ — an unsigned 8-bit integer. +-- +-- @since 0.3.1.0 +utinyint :: Db.DataType DuckDB Word8 +utinyint = Db.DataType duckDBUTinyIntType + +-- | DuckDB @USMALLINT@ — an unsigned 16-bit integer. +-- +-- @since 0.3.1.0 +usmallint :: Db.DataType DuckDB Word16 +usmallint = Db.DataType duckDBUSmallIntType + +-- | DuckDB @UINTEGER@ — an unsigned 32-bit integer. +-- +-- @since 0.3.1.0 +uinteger :: Db.DataType DuckDB Word32 +uinteger = Db.DataType duckDBUIntegerType + +-- | DuckDB @UBIGINT@ — an unsigned 64-bit integer. +-- +-- @since 0.3.1.0 +ubigint :: Db.DataType DuckDB Word64 +ubigint = Db.DataType duckDBUBigIntType + +-- | 'Db.BeamDeserializers' for DuckDB-specific data types not covered by +-- the standard SQL92/SQL99/SQL2008 deserializers. Wired into +-- 'migrationBackend' by default. +-- +-- @since 0.3.1.0 +duckDBDataTypeDeserializers :: Db.BeamDeserializers DuckDB +duckDBDataTypeDeserializers = + Db.beamDeserializer $ \_ -> + withText "DuckDB data type" $ \tag -> + case tag of + "tinyint" -> pure duckDBTinyIntType + "hugeint" -> pure duckDBHugeIntType + "uuid" -> pure duckDBUuidType + "utinyint" -> pure duckDBUTinyIntType + "usmallint" -> pure duckDBUSmallIntType + "uinteger" -> pure duckDBUIntegerType + "ubigint" -> pure duckDBUBigIntType + "uhugeint" -> pure duckDBUHugeIntType + _ -> fail $ "Unknown DuckDB data type with tag " <> T.unpack tag + +-- | Render a series of 'Db.MigrationSteps' for DuckDB as a list of lazy +-- 'BL.ByteString's suitable for assembling a migration script. +-- +-- @since 0.3.1.0 +migrateScript :: Db.MigrationSteps DuckDB () a -> [BL.ByteString] +migrateScript steps = + "-- Generated by beam-duckdb beam-migrate backend\n" + : "\n" + : appEndo (Db.migrateScript renderHeader renderCommand steps) [] + where + renderHeader nm = + Endo (("-- " <> BL.fromStrict (TE.encodeUtf8 nm) <> "\n") :) + renderCommand (DuckDBCommandSyntax inner) = + Endo ((BL.fromStrict (TE.encodeUtf8 (duckDBRenderSyntaxScript inner)) <> ";\n") :) + +-- | Convert a DuckDB type name (as reported by @information_schema@) into a +-- 'DuckDBDataTypeSyntax', or 'Nothing' if the name is not recognised. +-- +-- Handles common DuckDB-printed forms such as @INTEGER@, @VARCHAR(10)@, +-- @DECIMAL(18,3)@, @TIMESTAMP WITH TIME ZONE@, etc. +duckDBDataTypeFromText :: Text -> Maybe DuckDBDataTypeSyntax +duckDBDataTypeFromText t0 = + let t = T.toUpper (T.strip t0) + (base, parens) = splitParens t + in case base of + "BOOLEAN" -> Just booleanType + "BOOL" -> Just booleanType + "TINYINT" -> Just duckDBTinyIntType + "INT1" -> Just duckDBTinyIntType + "SMALLINT" -> Just smallIntType + "INT2" -> Just smallIntType + "INTEGER" -> Just intType + "INT" -> Just intType + "INT4" -> Just intType + "BIGINT" -> Just bigIntType + "INT8" -> Just bigIntType + "HUGEINT" -> Just duckDBHugeIntType + "UTINYINT" -> Just duckDBUTinyIntType + "USMALLINT" -> Just duckDBUSmallIntType + "UINTEGER" -> Just duckDBUIntegerType + "UBIGINT" -> Just duckDBUBigIntType + "UHUGEINT" -> Just duckDBUHugeIntType + "REAL" -> Just realType + "FLOAT" -> Just (floatType (parsePrec parens)) + "FLOAT4" -> Just realType + "DOUBLE" -> Just doubleType + "FLOAT8" -> Just doubleType + "DOUBLE PRECISION" -> Just doubleType + "NUMERIC" -> Just (numericType (parseNumericPrec parens)) + "DECIMAL" -> Just (decimalType (parseNumericPrec parens)) + "VARCHAR" -> Just (varCharType (parsePrec parens) Nothing) + "CHAR" -> Just (charType (parsePrec parens) Nothing) + "CHARACTER" -> Just (charType (parsePrec parens) Nothing) + "CHARACTER VARYING" -> Just (varCharType (parsePrec parens) Nothing) + "TEXT" -> Just (varCharType Nothing Nothing) + "STRING" -> Just (varCharType Nothing Nothing) + "BLOB" -> Just binaryLargeObjectType + "BYTEA" -> Just binaryLargeObjectType + "UUID" -> Just duckDBUuidType + "DATE" -> Just dateType + "TIME" -> Just (timeType (parsePrec parens) False) + "TIME WITH TIME ZONE" -> Just (timeType (parsePrec parens) True) + "TIMETZ" -> Just (timeType (parsePrec parens) True) + "TIMESTAMP" -> Just (timestampType (parsePrec parens) False) + "TIMESTAMP WITH TIME ZONE" -> Just (timestampType (parsePrec parens) True) + "TIMESTAMPTZ" -> Just (timestampType (parsePrec parens) True) + _ -> Nothing + where + splitParens :: Text -> (Text, Maybe Text) + splitParens t = + case T.breakOn "(" t of + (base, "") -> (T.strip base, Nothing) + (base, rest) -> + let inside = T.dropWhile (== '(') (T.dropWhileEnd (== ')') rest) + in (T.strip base, Just (T.strip inside)) + + parsePrec :: Maybe Text -> Maybe Word + parsePrec Nothing = Nothing + parsePrec (Just t) + | T.null digits = Nothing + | otherwise = Just (read (T.unpack digits)) + where + digits = T.takeWhile isDigit (T.strip t) + + parseNumericPrec :: Maybe Text -> Maybe (Word, Maybe Word) + parseNumericPrec Nothing = Nothing + parseNumericPrec (Just t) = + case T.splitOn "," t of + [p] -> fmap (,Nothing) (readMaybeWord (T.strip p)) + [p, d] -> case (readMaybeWord (T.strip p), readMaybeWord (T.strip d)) of + (Just pw, Just dw) -> Just (pw, Just dw) + (Just pw, Nothing) -> Just (pw, Nothing) + _ -> Nothing + _ -> Nothing + + readMaybeWord :: Text -> Maybe Word + readMaybeWord t + | T.null t || T.any (not . isDigit) t = Nothing + | otherwise = Just (read (T.unpack t)) + +-- | Discover the set of 'Db.SomeDatabasePredicate's describing the current +-- DuckDB database. Includes schema, table, column, primary-key, secondary +-- index, and foreign-key constraint predicates. +-- +-- @since 0.3.1.0 +getDbConstraints :: DuckDBM [Db.SomeDatabasePredicate] +getDbConstraints = do + schemas <- + liftIOWithHandle $ \conn -> + DuckDB.query_ + conn + "SELECT schema_name FROM information_schema.schemata \ + \WHERE schema_name NOT IN ('information_schema', 'pg_catalog', 'main') \ + \AND schema_name NOT LIKE 'pg\\_%' ESCAPE '\\'" + let userSchemas = map (\(Only s) -> s) schemas :: [Text] + case userSchemas of + [] -> getDbConstraintsForSchemas Nothing + ss -> + (++) + <$> getDbConstraintsForSchemas (Just ss) + <*> getDbConstraintsForSchemas Nothing + +-- | Like 'getDbConstraints', but optionally restricted to a list of schemas. +-- If 'Nothing', the @main@ schema (the DuckDB default) is queried. +-- +-- @since 0.3.1.0 +getDbConstraintsForSchemas :: Maybe [Text] -> DuckDBM [Db.SomeDatabasePredicate] +getDbConstraintsForSchemas mSchemas = do + let schemaFilter = case mSchemas of + Nothing -> "table_schema = '" <> implicitSchema <> "'" + Just ss -> + "table_schema IN (" + <> T.intercalate + "," + (map (\s -> "'" <> s <> "'") ss) + <> ")" + schemaPreds = + case mSchemas of + Nothing -> [] + Just ss -> map (Db.SomeDatabasePredicate . Db.SchemaExistsPredicate) ss + + tbls <- + liftIOWithHandle $ \conn -> + DuckDB.query_ conn $ + DuckDB.Query $ + T.unlines + [ "SELECT table_schema, table_name FROM information_schema.tables", + "WHERE table_type='BASE TABLE' AND " <> schemaFilter + ] + let tblExistsPreds = + map + ( \(s, t) -> + Db.SomeDatabasePredicate + ( Db.TableExistsPredicate + (Db.QualifiedName (qualifySchema s mSchemas) t) + ) + ) + tbls + + columnsAndConstraints <- + fmap mconcat . forM tbls $ \(s, t) -> collectColumnInfo s t mSchemas + + primaryKeyPreds <- collectPrimaryKeys schemaFilter mSchemas + fkPreds <- collectForeignKeys schemaFilter mSchemas + + pure $ + concat + [ schemaPreds, + tblExistsPreds, + columnsAndConstraints, + primaryKeyPreds, + fkPreds + ] + +-- | The implicit schema in case it's not specified +-- +-- https://duckdb.org/docs/lts/sql/statements/use +implicitSchema :: Text +implicitSchema = "main" + +-- | If we are looking at the 'main' schema, qualify columns/tables with +-- 'Nothing' (the implicit default). Otherwise qualify with the actual schema. +qualifySchema :: Text -> Maybe [Text] -> Maybe Text +qualifySchema s Nothing + | s == implicitSchema = Nothing + | otherwise = Just s +qualifySchema s _ = Just s + +collectColumnInfo :: + Text -> + Text -> + Maybe [Text] -> + DuckDBM [Db.SomeDatabasePredicate] +collectColumnInfo schemaNm tblNm mSchemas = do + columns <- + liftIOWithHandle $ \conn -> + DuckDB.query + conn + "SELECT column_name, data_type, is_nullable \ + \FROM information_schema.columns \ + \WHERE table_schema = ? AND table_name = ? \ + \ORDER BY ordinal_position" + (schemaNm, tblNm) + let qual = Db.QualifiedName (qualifySchema schemaNm mSchemas) tblNm + mkPreds (colNm, typ, isNullable) = + let dt = fromMaybe (unknownDuckDBType typ) (duckDBDataTypeFromText typ) + colPred = + Db.SomeDatabasePredicate + ( Db.TableHasColumn qual colNm dt :: + Db.TableHasColumn DuckDB + ) + notNullPred = + ( [ Db.SomeDatabasePredicate + ( Db.TableColumnHasConstraint + qual + colNm + ( Db.constraintDefinitionSyntax + Nothing + Db.notNullConstraintSyntax + Nothing + ) :: + Db.TableColumnHasConstraint DuckDB + ) + | isNullable == ("NO" :: Text) + ] + ) + in colPred : notNullPred + pure $ concatMap mkPreds (columns :: [(Text, Text, Text)]) + where + unknownDuckDBType :: Text -> DuckDBDataTypeSyntax + unknownDuckDBType = domainType + +collectPrimaryKeys :: + Text -> + Maybe [Text] -> + DuckDBM [Db.SomeDatabasePredicate] +collectPrimaryKeys schemaFilter mSchemas = do + rows <- + liftIOWithHandle $ \conn -> + DuckDB.query_ conn $ + DuckDB.Query $ + T.unlines + [ "SELECT tc.table_schema, tc.table_name, kcu.column_name, kcu.ordinal_position", + "FROM information_schema.table_constraints tc", + "JOIN information_schema.key_column_usage kcu", + " ON tc.constraint_name = kcu.constraint_name", + " AND tc.constraint_schema = kcu.constraint_schema", + " AND tc.table_name = kcu.table_name", + "WHERE tc.constraint_type = 'PRIMARY KEY' AND tc." <> schemaFilter, + "ORDER BY tc.table_schema, tc.table_name, kcu.ordinal_position" + ] + let grouped = groupByPk (rows :: [(Text, Text, Text, Int)]) + pure $ + map + ( \((sch, tbl), cols) -> + Db.SomeDatabasePredicate + ( Db.TableHasPrimaryKey + (Db.QualifiedName (qualifySchema sch mSchemas) tbl) + cols + ) + ) + grouped + +groupByPk :: + [(Text, Text, Text, Int)] -> + [((Text, Text), [Text])] +groupByPk = go [] + where + go acc [] = reverse acc + go acc rs@((s, t, _, _) : _) = + let (here, rest) = span (\(s', t', _, _) -> s == s' && t == t') rs + cols = map (\(_, _, c, _) -> c) here + in go (((s, t), cols) : acc) rest + +collectForeignKeys :: + Text -> + Maybe [Text] -> + DuckDBM [Db.SomeDatabasePredicate] +collectForeignKeys schemaFilter mSchemas = do + rows <- + liftIOWithHandle $ \conn -> + DuckDB.query_ conn $ + DuckDB.Query $ + T.unlines + [ "SELECT tc.table_schema, tc.table_name, kcu.column_name,", + " ccu.table_schema, ccu.table_name, ccu.column_name,", + " rc.update_rule, rc.delete_rule,", + " kcu.ordinal_position", + "FROM information_schema.table_constraints tc", + "JOIN information_schema.key_column_usage kcu", + " ON tc.constraint_name = kcu.constraint_name", + " AND tc.constraint_schema = kcu.constraint_schema", + "JOIN information_schema.referential_constraints rc", + " ON rc.constraint_name = tc.constraint_name", + " AND rc.constraint_schema = tc.constraint_schema", + "JOIN information_schema.constraint_column_usage ccu", + " ON ccu.constraint_name = rc.unique_constraint_name", + " AND ccu.constraint_schema = rc.unique_constraint_schema", + "WHERE tc.constraint_type = 'FOREIGN KEY' AND tc." <> schemaFilter, + "ORDER BY tc.constraint_schema, tc.constraint_name, kcu.ordinal_position" + ] + pure $ + mapMaybe mkFk $ + groupByFk + ( rows :: + [ ( Text, + Text, + Text, + Text, + Text, + Text, + Text, + Text, + Int + ) + ] + ) + where + groupByFk = go [] + where + go acc [] = reverse acc + go acc rs@((s, t, _, fs, ft, _, _, _, _) : _) = + let (here, rest) = + span + ( \(s', t', _, fs', ft', _, _, _, _) -> + s == s' && t == t' && fs == fs' && ft == ft' + ) + rs + in go (here : acc) rest + + mkFk [] = Nothing + mkFk rows@((srcSchema, srcTbl, _, refSchema, refTbl, _, upd, del, _) : _) = + let localCols = map (\(_, _, c, _, _, _, _, _, _) -> c) rows + refCols = map (\(_, _, _, _, _, c, _, _, _) -> c) rows + in case (NE.nonEmpty localCols, NE.nonEmpty refCols) of + (Just lc, Just rc) -> + Just $ + Db.SomeDatabasePredicate + ( Db.TableHasForeignKey + (Db.QualifiedName (qualifySchema srcSchema mSchemas) srcTbl) + lc + (Db.QualifiedName (qualifySchema refSchema mSchemas) refTbl) + rc + (parseDuckDBForeignKeyAction upd) + (parseDuckDBForeignKeyAction del) + ) + _ -> Nothing + + parseDuckDBForeignKeyAction :: Text -> Db.ForeignKeyAction + parseDuckDBForeignKeyAction t = + case T.toUpper t of + "CASCADE" -> Db.ForeignKeyActionCascade + "SET NULL" -> Db.ForeignKeyActionSetNull + "SET DEFAULT" -> Db.ForeignKeyActionSetDefault + "RESTRICT" -> Db.ForeignKeyActionRestrict + "NO ACTION" -> Db.ForeignKeyNoAction + _ -> Db.ForeignKeyNoAction diff --git a/src/Database/Beam/DuckDB/Syntax.hs b/src/Database/Beam/DuckDB/Syntax.hs --- a/src/Database/Beam/DuckDB/Syntax.hs +++ b/src/Database/Beam/DuckDB/Syntax.hs @@ -1,10 +1,10 @@ {-# LANGUAGE DerivingStrategies #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE TypeFamilies #-} --- TODO: clean up unused top binds {-# OPTIONS_GHC -Wno-unused-top-binds #-} module Database.Beam.DuckDB.Syntax @@ -20,12 +20,35 @@ DuckDBDeleteSyntax (..), DuckDBTableNameSyntax (..), DuckDBOnConflictSyntax (..), + DuckDBFieldNameSyntax (..), + + -- * DDL syntaxes + DuckDBDataTypeSyntax (..), + DuckDBCreateTableSyntax (..), + DuckDBDropTableSyntax (..), + DuckDBAlterTableSyntax (..), + DuckDBAlterTableActionSyntax (..), + DuckDBAlterColumnActionSyntax (..), + DuckDBCreateSchemaSyntax (..), + DuckDBDropSchemaSyntax (..), + DuckDBSchemaNameSyntax (..), + DuckDBColumnSchemaSyntax (..), + DuckDBColumnConstraintDefinitionSyntax (..), + DuckDBColumnConstraintSyntax (..), + DuckDBTableConstraintSyntax (..), + DuckDBMatchTypeSyntax (..), + DuckDBReferentialActionSyntax (..), + DuckDBTableOptionsSyntax (..), + DuckDBIndexOptions (..), ) where +import Data.Aeson (object, withObject, (.:), (.=)) import Data.Bifunctor (second) import Data.Coerce (coerce) +import Data.Hashable (Hashable (..)) import Data.Int (Int16, Int32, Int64, Int8) +import qualified Data.List.NonEmpty as NE import Data.Scientific (Scientific) import Data.String (fromString) import Data.Text (Text) @@ -63,6 +86,7 @@ IsSql92OrderingSyntax (..), IsSql92ProjectionSyntax (..), IsSql92QuantifierSyntax (..), + IsSql92SchemaNameSyntax (..), IsSql92SelectSyntax (..), IsSql92SelectTableSyntax (..), IsSql92Syntax (..), @@ -79,8 +103,46 @@ IsSql99RecursiveCommonTableExpressionSelectSyntax (..), SqlNull (..), ) -import Database.Beam.DuckDB.Syntax.Builder (DuckDBSyntax, commas, emit, emitChar, emitIntegral, emitRealFloat, emitScientific, emitValue, parens, quotedIdentifier, sepBy, spaces) -import Database.Beam.Migrate.Serialization (BeamSerializedDataType) +import Database.Beam.Backend.SQL + ( IsSql2003BinaryAndVarBinaryDataTypeSyntax (..), + IsSql2008BigIntDataTypeSyntax (..), + Sql92DisplaySyntax (..), + ) +import Database.Beam.DuckDB.Syntax.Builder (DuckDBSyntax, commas, duckDBRenderSyntaxScript, emit, emitChar, emitIntegral, emitRealFloat, emitScientific, emitValue, parens, quotedIdentifier, sepBy, spaces) +import Database.Beam.Migrate.Checks (HasDataTypeCreatedCheck (..)) +import Database.Beam.Migrate.SQL + ( IsSql92AlterColumnActionSyntax (..), + IsSql92AlterTableActionSyntax (..), + IsSql92AlterTableSyntax (..), + IsSql92ColumnConstraintDefinitionSyntax (..), + IsSql92ColumnConstraintSyntax (..), + IsSql92ColumnSchemaSyntax (..), + IsSql92CreateDropIndexSyntax (..), + IsSql92CreateSchemaSyntax (..), + IsSql92CreateTableSyntax (..), + IsSql92DdlCommandSyntax (..), + IsSql92DdlSchemaCommandSyntax (..), + IsSql92DropSchemaSyntax (..), + IsSql92DropTableSyntax (..), + IsSql92MatchTypeSyntax (..), + IsSql92ReferentialActionSyntax (..), + IsSql92TableConstraintSyntax (..), + IsSql92UniqueIndexSyntax (..), + Sql92SerializableConstraintDefinitionSyntax (..), + Sql92SerializableDataTypeSyntax (..), + ) +import Database.Beam.Migrate.SQL.Builder (ConstraintAttributeTiming (..), SqlConstraintAttributesBuilder (..), sqlConstraintAttributesSerialized) +import Database.Beam.Migrate.SQL.SQL92 (ForeignKeyAction (..)) +import Database.Beam.Migrate.Serialization + ( BeamSerializedConstraint, + BeamSerializedConstraintDefinition, + BeamSerializedDataType (..), + BeamSerializedExpression (..), + BeamSerializedMatchType, + BeamSerializedReferentialAction, + fromBeamSerializedConstraintDefinition, + fromBeamSerializedDataType, + ) import Database.DuckDB.Simple (Null (Null)) newtype DuckDBCommandSyntax = DuckDBCommandSyntax {fromDuckDBSyntax :: DuckDBSyntax} @@ -97,6 +159,22 @@ updateCmd = DuckDBCommandSyntax . fromDuckDBUpdate deleteCmd = DuckDBCommandSyntax . fromDuckDBDelete +instance IsSql92DdlCommandSyntax DuckDBCommandSyntax where + type Sql92DdlCommandCreateTableSyntax DuckDBCommandSyntax = DuckDBCreateTableSyntax + type Sql92DdlCommandDropTableSyntax DuckDBCommandSyntax = DuckDBDropTableSyntax + type Sql92DdlCommandAlterTableSyntax DuckDBCommandSyntax = DuckDBAlterTableSyntax + + createTableCmd = DuckDBCommandSyntax . fromDuckDBCreateTable + dropTableCmd = DuckDBCommandSyntax . fromDuckDBDropTable + alterTableCmd = DuckDBCommandSyntax . fromDuckDBAlterTable + +instance IsSql92DdlSchemaCommandSyntax DuckDBCommandSyntax where + type Sql92DdlCommandCreateSchemaSyntax DuckDBCommandSyntax = DuckDBCreateSchemaSyntax + type Sql92DdlCommandDropSchemaSyntax DuckDBCommandSyntax = DuckDBDropSchemaSyntax + + createSchemaCmd = DuckDBCommandSyntax . fromDuckDBCreateSchema + dropSchemaCmd = DuckDBCommandSyntax . fromDuckDBDropSchema + newtype DuckDBTableNameSyntax = DuckDBTableNameSyntax {fromDuckDBTableName :: DuckDBSyntax} -- | DuckDB @SELECT@ syntax @@ -124,7 +202,58 @@ { duckDBDataType :: DuckDBSyntax, duckDBDataTypeSerialized :: BeamSerializedDataType } + deriving (Show) +-- DDL syntax newtypes used by 'IsSql92DdlCommandSyntax' and friends. + +newtype DuckDBCreateTableSyntax = DuckDBCreateTableSyntax {fromDuckDBCreateTable :: DuckDBSyntax} + +data DuckDBTableOptionsSyntax = DuckDBTableOptionsSyntax DuckDBSyntax DuckDBSyntax + +newtype DuckDBDropTableSyntax = DuckDBDropTableSyntax {fromDuckDBDropTable :: DuckDBSyntax} + +newtype DuckDBAlterTableSyntax = DuckDBAlterTableSyntax {fromDuckDBAlterTable :: DuckDBSyntax} + +newtype DuckDBAlterTableActionSyntax = DuckDBAlterTableActionSyntax {fromDuckDBAlterTableAction :: DuckDBSyntax} + +newtype DuckDBAlterColumnActionSyntax = DuckDBAlterColumnActionSyntax {fromDuckDBAlterColumnAction :: DuckDBSyntax} + +newtype DuckDBCreateSchemaSyntax = DuckDBCreateSchemaSyntax {fromDuckDBCreateSchema :: DuckDBSyntax} + +newtype DuckDBDropSchemaSyntax = DuckDBDropSchemaSyntax {fromDuckDBDropSchema :: DuckDBSyntax} + +newtype DuckDBSchemaNameSyntax = DuckDBSchemaNameSyntax {fromDuckDBSchemaName :: DuckDBSyntax} + +newtype DuckDBColumnSchemaSyntax = DuckDBColumnSchemaSyntax {fromDuckDBColumnSchema :: DuckDBSyntax} + deriving (Show, Eq) + +data DuckDBColumnConstraintDefinitionSyntax = DuckDBColumnConstraintDefinitionSyntax + { fromDuckDBColumnConstraintDefinition :: DuckDBSyntax, + duckDBColumnConstraintDefinitionSerialized :: BeamSerializedConstraintDefinition + } + deriving (Show) + +data DuckDBColumnConstraintSyntax = DuckDBColumnConstraintSyntax + { fromDuckDBColumnConstraint :: DuckDBSyntax, + duckDBColumnConstraintSerialized :: BeamSerializedConstraint + } + +newtype DuckDBTableConstraintSyntax = DuckDBTableConstraintSyntax {fromDuckDBTableConstraint :: DuckDBSyntax} + +data DuckDBMatchTypeSyntax = DuckDBMatchTypeSyntax + { fromDuckDBMatchType :: DuckDBSyntax, + duckDBMatchTypeSerialized :: BeamSerializedMatchType + } + +data DuckDBReferentialActionSyntax = DuckDBReferentialActionSyntax + { fromDuckDBReferentialAction :: DuckDBSyntax, + duckDBReferentialActionSerialized :: BeamSerializedReferentialAction + } + +newtype DuckDBIndexOptions = DuckDBIndexOptions {duckDBIndexUnique :: Bool} + deriving stock (Show, Eq) + deriving newtype (Hashable) + newtype DuckDBExtractFieldSyntax = DuckDBExtractFieldSyntax {fromDuckDBExtractField :: DuckDBSyntax} newtype DuckDBGroupingSyntax = DuckDBGroupingSyntax {fromDuckDBGrouping :: DuckDBSyntax} @@ -349,6 +478,293 @@ <> emitChar ')', duckDBDataTypeSerialized = rowType (map (second duckDBDataTypeSerialized) fields) } + +instance IsSql2008BigIntDataTypeSyntax DuckDBDataTypeSyntax where + bigIntType = DuckDBDataTypeSyntax (emit "BIGINT") bigIntType + +instance IsSql2003BinaryAndVarBinaryDataTypeSyntax DuckDBDataTypeSyntax where + binaryType sz = + DuckDBDataTypeSyntax + (emit "BLOB" <> optPrec sz) + (binaryType sz) + varBinaryType sz = + DuckDBDataTypeSyntax + (emit "BLOB" <> optPrec sz) + (varBinaryType sz) + +-- | Two 'DuckDBDataTypeSyntax' values are equal when their rendered SQL is +-- equal. This is good enough for migration purposes, where the serialized +-- JSON form is used as the canonical comparison key by @beam-migrate@. +instance Eq DuckDBDataTypeSyntax where + DuckDBDataTypeSyntax a _ == DuckDBDataTypeSyntax b _ = a == b + +instance Hashable DuckDBDataTypeSyntax where + hashWithSalt salt (DuckDBDataTypeSyntax a _) = hashWithSalt salt a + +instance Sql92DisplaySyntax DuckDBDataTypeSyntax where + displaySyntax = displaySyntax . duckDBDataType + +instance Sql92SerializableDataTypeSyntax DuckDBDataTypeSyntax where + serializeDataType = fromBeamSerializedDataType . duckDBDataTypeSerialized + +-- | DuckDB doesn't allow user-defined domain types, so every data type is +-- always available. +instance HasDataTypeCreatedCheck DuckDBDataTypeSyntax where + dataTypeHasBeenCreated _ _ = True + +instance Sql92DisplaySyntax DuckDBColumnSchemaSyntax where + displaySyntax = displaySyntax . fromDuckDBColumnSchema + +instance Hashable DuckDBColumnSchemaSyntax where + hashWithSalt salt = hashWithSalt salt . fromDuckDBColumnSchema + +instance Eq DuckDBColumnConstraintDefinitionSyntax where + DuckDBColumnConstraintDefinitionSyntax a _ == DuckDBColumnConstraintDefinitionSyntax b _ = a == b + +instance Hashable DuckDBColumnConstraintDefinitionSyntax where + hashWithSalt salt = hashWithSalt salt . fromDuckDBColumnConstraintDefinition + +instance Sql92DisplaySyntax DuckDBColumnConstraintDefinitionSyntax where + displaySyntax = displaySyntax . fromDuckDBColumnConstraintDefinition + +instance Sql92SerializableConstraintDefinitionSyntax DuckDBColumnConstraintDefinitionSyntax where + serializeConstraint = fromBeamSerializedConstraintDefinition . duckDBColumnConstraintDefinitionSerialized + +instance IsSql92SchemaNameSyntax DuckDBSchemaNameSyntax where + schemaName s = DuckDBSchemaNameSyntax (quotedIdentifier s) + +instance IsSql92CreateSchemaSyntax DuckDBCreateSchemaSyntax where + type Sql92CreateSchemaSchemaNameSyntax DuckDBCreateSchemaSyntax = DuckDBSchemaNameSyntax + + createSchemaSyntax sNm = + DuckDBCreateSchemaSyntax (emit "CREATE SCHEMA " <> fromDuckDBSchemaName sNm) + +instance IsSql92DropSchemaSyntax DuckDBDropSchemaSyntax where + type Sql92DropSchemaSchemaNameSyntax DuckDBDropSchemaSyntax = DuckDBSchemaNameSyntax + + dropSchemaSyntax sNm = + DuckDBDropSchemaSyntax (emit "DROP SCHEMA " <> fromDuckDBSchemaName sNm) + +instance IsSql92DropTableSyntax DuckDBDropTableSyntax where + type Sql92DropTableTableNameSyntax DuckDBDropTableSyntax = DuckDBTableNameSyntax + + dropTableSyntax tblNm = + DuckDBDropTableSyntax (emit "DROP TABLE " <> fromDuckDBTableName tblNm) + +instance IsSql92AlterTableSyntax DuckDBAlterTableSyntax where + type Sql92AlterTableAlterTableActionSyntax DuckDBAlterTableSyntax = DuckDBAlterTableActionSyntax + type Sql92AlterTableTableNameSyntax DuckDBAlterTableSyntax = DuckDBTableNameSyntax + + alterTableSyntax tblNm action = + DuckDBAlterTableSyntax $ + emit "ALTER TABLE " + <> fromDuckDBTableName tblNm + <> emit " " + <> fromDuckDBAlterTableAction action + +instance IsSql92AlterTableActionSyntax DuckDBAlterTableActionSyntax where + type Sql92AlterTableAlterColumnActionSyntax DuckDBAlterTableActionSyntax = DuckDBAlterColumnActionSyntax + type Sql92AlterTableColumnSchemaSyntax DuckDBAlterTableActionSyntax = DuckDBColumnSchemaSyntax + + alterColumnSyntax colNm action = + DuckDBAlterTableActionSyntax $ + emit "ALTER COLUMN " + <> quotedIdentifier colNm + <> emit " " + <> fromDuckDBAlterColumnAction action + + addColumnSyntax colNm schema = + DuckDBAlterTableActionSyntax $ + emit "ADD COLUMN " + <> quotedIdentifier colNm + <> emit " " + <> fromDuckDBColumnSchema schema + + dropColumnSyntax colNm = + DuckDBAlterTableActionSyntax (emit "DROP COLUMN " <> quotedIdentifier colNm) + + renameTableToSyntax newNm = + DuckDBAlterTableActionSyntax (emit "RENAME TO " <> quotedIdentifier newNm) + + renameColumnToSyntax oldNm newNm = + DuckDBAlterTableActionSyntax $ + emit "RENAME COLUMN " + <> quotedIdentifier oldNm + <> emit " TO " + <> quotedIdentifier newNm + +instance IsSql92AlterColumnActionSyntax DuckDBAlterColumnActionSyntax where + setNullSyntax = DuckDBAlterColumnActionSyntax (emit "DROP NOT NULL") + setNotNullSyntax = DuckDBAlterColumnActionSyntax (emit "SET NOT NULL") + +instance IsSql92CreateTableSyntax DuckDBCreateTableSyntax where + type Sql92CreateTableTableNameSyntax DuckDBCreateTableSyntax = DuckDBTableNameSyntax + type Sql92CreateTableColumnSchemaSyntax DuckDBCreateTableSyntax = DuckDBColumnSchemaSyntax + type Sql92CreateTableTableConstraintSyntax DuckDBCreateTableSyntax = DuckDBTableConstraintSyntax + type Sql92CreateTableOptionsSyntax DuckDBCreateTableSyntax = DuckDBTableOptionsSyntax + + createTableSyntax options tblNm fieldTypes constraints = + let (beforeOptions, afterOptions) = + case options of + Nothing -> (emit " ", emit " ") + Just (DuckDBTableOptionsSyntax before after) -> + (emit " " <> before <> emit " ", emit " " <> after <> emit " ") + in DuckDBCreateTableSyntax $ + emit "CREATE" + <> beforeOptions + <> emit "TABLE " + <> fromDuckDBTableName tblNm + <> emit " (" + <> sepBy + (emit ", ") + ( map (\(nm, ty) -> quotedIdentifier nm <> emit " " <> fromDuckDBColumnSchema ty) fieldTypes + <> map fromDuckDBTableConstraint constraints + ) + <> emit ")" + <> afterOptions + +duckDBForeignKeyAction :: ForeignKeyAction -> DuckDBSyntax +duckDBForeignKeyAction ForeignKeyActionCascade = emit "CASCADE" +duckDBForeignKeyAction ForeignKeyActionSetNull = emit "SET NULL" +duckDBForeignKeyAction ForeignKeyActionSetDefault = emit "SET DEFAULT" +duckDBForeignKeyAction ForeignKeyActionRestrict = emit "RESTRICT" +duckDBForeignKeyAction ForeignKeyNoAction = emit "NO ACTION" + +instance IsSql92TableConstraintSyntax DuckDBTableConstraintSyntax where + primaryKeyConstraintSyntax fieldNames = + DuckDBTableConstraintSyntax $ + emit "PRIMARY KEY(" + <> sepBy (emit ", ") (map quotedIdentifier (NE.toList fieldNames)) + <> emit ")" + foreignKeyConstraintSyntax localCols refTbl refCols onUpdate onDelete = + DuckDBTableConstraintSyntax $ + emit "FOREIGN KEY(" + <> sepBy (emit ", ") (map quotedIdentifier (NE.toList localCols)) + <> emit ") REFERENCES " + <> quotedIdentifier refTbl + <> emit "(" + <> sepBy (emit ", ") (map quotedIdentifier (NE.toList refCols)) + <> emit ")" + <> emit " ON UPDATE " + <> duckDBForeignKeyAction onUpdate + <> emit " ON DELETE " + <> duckDBForeignKeyAction onDelete + +instance IsSql92ColumnSchemaSyntax DuckDBColumnSchemaSyntax where + type Sql92ColumnSchemaColumnTypeSyntax DuckDBColumnSchemaSyntax = DuckDBDataTypeSyntax + type Sql92ColumnSchemaExpressionSyntax DuckDBColumnSchemaSyntax = DuckDBExpressionSyntax + type Sql92ColumnSchemaColumnConstraintDefinitionSyntax DuckDBColumnSchemaSyntax = DuckDBColumnConstraintDefinitionSyntax + + columnSchemaSyntax colType defaultClause constraints collation = + DuckDBColumnSchemaSyntax $ + duckDBDataType colType + <> maybe mempty (\d -> emit " DEFAULT " <> fromDuckDBExpression d) defaultClause + <> ( case constraints of + [] -> mempty + _ -> foldMap (\c -> emit " " <> fromDuckDBColumnConstraintDefinition c) constraints + ) + <> maybe mempty (\nm -> emit " COLLATE " <> quotedIdentifier nm) collation + +instance IsSql92MatchTypeSyntax DuckDBMatchTypeSyntax where + fullMatchSyntax = DuckDBMatchTypeSyntax (emit "FULL") fullMatchSyntax + partialMatchSyntax = DuckDBMatchTypeSyntax (emit "PARTIAL") partialMatchSyntax + +instance IsSql92ReferentialActionSyntax DuckDBReferentialActionSyntax where + referentialActionCascadeSyntax = DuckDBReferentialActionSyntax (emit "CASCADE") referentialActionCascadeSyntax + referentialActionNoActionSyntax = DuckDBReferentialActionSyntax (emit "NO ACTION") referentialActionNoActionSyntax + referentialActionSetDefaultSyntax = DuckDBReferentialActionSyntax (emit "SET DEFAULT") referentialActionSetDefaultSyntax + referentialActionSetNullSyntax = DuckDBReferentialActionSyntax (emit "SET NULL") referentialActionSetNullSyntax + +instance IsSql92ColumnConstraintDefinitionSyntax DuckDBColumnConstraintDefinitionSyntax where + type Sql92ColumnConstraintDefinitionConstraintSyntax DuckDBColumnConstraintDefinitionSyntax = DuckDBColumnConstraintSyntax + type Sql92ColumnConstraintDefinitionAttributesSyntax DuckDBColumnConstraintDefinitionSyntax = SqlConstraintAttributesBuilder + + constraintDefinitionSyntax nm constraint attrs = + DuckDBColumnConstraintDefinitionSyntax syntax serialized + where + syntax = + maybe mempty (\n -> emit "CONSTRAINT " <> quotedIdentifier n <> emit " ") nm + <> fromDuckDBColumnConstraint constraint + <> maybe mempty (\a -> emit " " <> duckDBSqlConstraintAttributes a) attrs + serialized = + constraintDefinitionSyntax nm (duckDBColumnConstraintSerialized constraint) (fmap sqlConstraintAttributesSerialized attrs) + +-- | Emit a 'SqlConstraintAttributesBuilder' as DuckDB syntax. DuckDB does not +-- support @DEFERRABLE@/@INITIALLY DEFERRED@ at the time of writing, but we +-- still emit the corresponding SQL keywords so that user-supplied constraint +-- attributes are not silently dropped. If DuckDB rejects them, the user will +-- see an obvious error rather than a subtle behavioural difference. +duckDBSqlConstraintAttributes :: SqlConstraintAttributesBuilder -> DuckDBSyntax +duckDBSqlConstraintAttributes (SqlConstraintAttributesBuilder timing deferrable) = + maybe mempty timingBuilder timing <> maybe mempty deferrableBuilder deferrable + where + timingBuilder InitiallyDeferred = emit "INITIALLY DEFERRED" + timingBuilder InitiallyImmediate = emit "INITIALLY IMMEDIATE" + deferrableBuilder False = emit "NOT DEFERRABLE" + deferrableBuilder True = emit "DEFERRABLE" + +instance IsSql92ColumnConstraintSyntax DuckDBColumnConstraintSyntax where + type Sql92ColumnConstraintMatchTypeSyntax DuckDBColumnConstraintSyntax = DuckDBMatchTypeSyntax + type Sql92ColumnConstraintReferentialActionSyntax DuckDBColumnConstraintSyntax = DuckDBReferentialActionSyntax + type Sql92ColumnConstraintExpressionSyntax DuckDBColumnConstraintSyntax = DuckDBExpressionSyntax + + notNullConstraintSyntax = DuckDBColumnConstraintSyntax (emit "NOT NULL") notNullConstraintSyntax + uniqueColumnConstraintSyntax = DuckDBColumnConstraintSyntax (emit "UNIQUE") uniqueColumnConstraintSyntax + primaryKeyColumnConstraintSyntax = DuckDBColumnConstraintSyntax (emit "PRIMARY KEY") primaryKeyColumnConstraintSyntax + checkColumnConstraintSyntax expr = + DuckDBColumnConstraintSyntax + (emit "CHECK(" <> fromDuckDBExpression expr <> emit ")") + ( checkColumnConstraintSyntax + ( BeamSerializedExpression + (duckDBRenderSyntaxScript (fromDuckDBExpression expr)) + ) + ) + referencesConstraintSyntax tbl fields matchType onUpdate onDelete = + DuckDBColumnConstraintSyntax syntax serialized + where + syntax = + emit "REFERENCES " + <> quotedIdentifier tbl + <> emit "(" + <> sepBy (emit ", ") (map quotedIdentifier fields) + <> emit ")" + <> maybe mempty (\m -> emit " " <> fromDuckDBMatchType m) matchType + <> maybe mempty (\a -> emit " ON UPDATE " <> fromDuckDBReferentialAction a) onUpdate + <> maybe mempty (\a -> emit " ON DELETE " <> fromDuckDBReferentialAction a) onDelete + serialized = + referencesConstraintSyntax + tbl + fields + (fmap duckDBMatchTypeSerialized matchType) + (fmap duckDBReferentialActionSerialized onUpdate) + (fmap duckDBReferentialActionSerialized onDelete) + +instance IsSql92CreateDropIndexSyntax DuckDBCommandSyntax where + type Sql92CreateIndexOptionsSyntax DuckDBCommandSyntax = DuckDBIndexOptions + + defaultIndexOptions = DuckDBIndexOptions {duckDBIndexUnique = False} + + createIndexCmd idxNm tblNm cols opts = + DuckDBCommandSyntax $ + emit (if duckDBIndexUnique opts then "CREATE UNIQUE INDEX " else "CREATE INDEX ") + <> quotedIdentifier idxNm + <> emit " ON " + <> fromDuckDBTableName tblNm + <> parens (sepBy (emit ", ") (NE.toList (fmap quotedIdentifier cols))) + + dropIndexCmd idxNm = + DuckDBCommandSyntax (emit "DROP INDEX " <> quotedIdentifier idxNm) + + serializeIndexOptions opts = + object ["unique" .= duckDBIndexUnique opts] + + deserializeIndexOptions = + withObject "DuckDBIndexOptions" $ \v -> + DuckDBIndexOptions <$> v .: "unique" + +instance IsSql92UniqueIndexSyntax DuckDBCommandSyntax where + setUniqueIndexOptions u opts = opts {duckDBIndexUnique = u} + indexIsUnique = duckDBIndexUnique optPrec :: Maybe Word -> DuckDBSyntax optPrec Nothing = mempty diff --git a/src/Database/Beam/DuckDB/Syntax/Builder.hs b/src/Database/Beam/DuckDB/Syntax/Builder.hs --- a/src/Database/Beam/DuckDB/Syntax/Builder.hs +++ b/src/Database/Beam/DuckDB/Syntax/Builder.hs @@ -5,11 +5,13 @@ ( DuckDBSyntax (..), SomeField (..), emitChar, + emitCharLit, emit, emitIntegral, emitRealFloat, emitScientific, emit', + emitTextLit, emitValue, spaces, parens, @@ -17,11 +19,14 @@ commas, quotedIdentifier, withPlaceholder, + duckDBRenderSyntaxScript, ) where import Data.DList (DList) import qualified Data.DList as DL +import Data.Hashable (Hashable (..)) +import Data.Scientific (Scientific) import Data.Text (Text) import qualified Data.Text as Text import qualified Data.Text.Lazy as Text.Lazy @@ -29,11 +34,10 @@ import qualified Data.Text.Lazy.Builder as Builder import qualified Data.Text.Lazy.Builder.Int as Builder import qualified Data.Text.Lazy.Builder.RealFloat as Builder +import qualified Data.Text.Lazy.Builder.Scientific as Builder.Scientific import Database.Beam.Backend (Sql92DisplaySyntax (..)) import Database.DuckDB.Simple (ToField (toField)) import Database.DuckDB.Simple.ToField (renderFieldBinding) -import Data.Scientific (Scientific) -import qualified Data.Text.Lazy.Builder.Scientific as Builder.Scientific data SomeField = forall a. (ToField a, Eq a) => SomeField a @@ -66,6 +70,22 @@ emit' :: (Show a) => a -> DuckDBSyntax emit' s = DuckDBSyntax (const (Builder.fromString (show s))) mempty +-- | Render a 'Text' as a single-quoted SQL string literal, e.g. @|@ -> @'|'@. +-- DuckDB requires this for option values like @DELIMITER@, @NULLSTR@, etc. +emitTextLit :: Text -> DuckDBSyntax +emitTextLit t = emitChar '\'' <> emit t <> emitChar '\'' + +-- | Render a 'Char' as a single-quoted SQL string literal, e.g. @"@ -> @'"'@. +-- The apostrophe character is escaped per the SQL standard by doubling, +-- so @\'@ renders as @''''@. +-- +-- The emitter places no constraint on the input 'Char'; format-specific +-- runtime constraints (e.g. DuckDB's CSV @QUOTE@ option requiring exactly +-- 1 byte in UTF-8) are enforced by DuckDB at execute time, not here. +emitCharLit :: Char -> DuckDBSyntax +emitCharLit '\'' = emit "''''" +emitCharLit c = emitChar '\'' <> emitChar c <> emitChar '\'' + -- | Emit a properly escaped value into the syntax emitValue :: (ToField a, Eq a) => a -> DuckDBSyntax emitValue v = @@ -104,6 +124,10 @@ -- TODO: is there a cleaner way to do this? Is it even important to have this instance? show vx == show vy && withPlaceholder sx == withPlaceholder sy + +instance Hashable DuckDBSyntax where + hashWithSalt salt (DuckDBSyntax s vs) = + hashWithSalt salt (Builder.toLazyText (withPlaceholder s), show vs) instance Semigroup DuckDBSyntax where DuckDBSyntax sx vx <> DuckDBSyntax sy vy = diff --git a/src/Database/Beam/DuckDB/Syntax/Extensions.hs b/src/Database/Beam/DuckDB/Syntax/Extensions.hs new file mode 100644 --- /dev/null +++ b/src/Database/Beam/DuckDB/Syntax/Extensions.hs @@ -0,0 +1,81 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} + +module Database.Beam.DuckDB.Syntax.Extensions () where + +import Data.Proxy (Proxy (..)) +import Database.Beam (Beamable, TableSettings) +import Database.Beam.Backend.SQL + ( IsSql92ExpressionSyntax (..), + IsSql92FieldNameSyntax (..), + MonadBeam (runReturningList), + ) +import Database.Beam.Backend.SQL.BeamExtensions + ( MonadBeamDeleteReturning (..), + MonadBeamInsertReturning (..), + MonadBeamUpdateReturning (..), + ) +import Database.Beam.DuckDB.Backend (DuckDB) +import Database.Beam.DuckDB.Connection (DuckDBM) +import Database.Beam.DuckDB.Syntax + ( DuckDBCommandSyntax (..), + DuckDBDeleteSyntax (..), + DuckDBExpressionSyntax (..), + DuckDBInsertSyntax (..), + DuckDBUpdateSyntax (..), + ) +import Database.Beam.DuckDB.Syntax.Builder (DuckDBSyntax, commas, emit) +import Database.Beam.DuckDB.Syntax.Extensions.InsertOnConflict () +import Database.Beam.Query (SqlDelete (..), SqlInsert (..), SqlUpdate (..)) +import Database.Beam.Query.Internal (Projectible, QExpr, QGenExpr (..), project) +import Database.Beam.Schema.Tables (Columnar' (..), changeBeamRep, _fieldName) + +instance MonadBeamInsertReturning DuckDB DuckDBM where + -- \| @since 0.3.1.0 + runInsertReturningListWith SqlInsertNoRows _ = pure [] + runInsertReturningListWith (SqlInsert tblSettings (DuckDBInsertSyntax syntax)) mkProjection = + runReturningList + (DuckDBCommandSyntax (syntax <> returningClause tblSettings mkProjection)) + +instance MonadBeamUpdateReturning DuckDB DuckDBM where + -- \| @since 0.3.1.0 + runUpdateReturningListWith SqlIdentityUpdate _ = pure [] + runUpdateReturningListWith (SqlUpdate tblSettings (DuckDBUpdateSyntax syntax)) mkProjection = + runReturningList + (DuckDBCommandSyntax (syntax <> returningClause tblSettings mkProjection)) + +instance MonadBeamDeleteReturning DuckDB DuckDBM where + -- \| @since 0.3.1.0 + runDeleteReturningListWith (SqlDelete tblSettings (DuckDBDeleteSyntax syntax)) mkProjection = + runReturningList + (DuckDBCommandSyntax (syntax <> returningClause tblSettings mkProjection)) + +-- | Build @RETURNING e1, e2, ...@ for a projection over the given table's +-- fields, using unqualified field references (DuckDB does not require a +-- table prefix in @RETURNING@ projections). +returningClause :: + forall table a. + (Beamable table, Projectible DuckDB a) => + TableSettings table -> + (table (QExpr DuckDB ()) -> a) -> + DuckDBSyntax +returningClause tblSettings mkProjection = + emit " RETURNING " + <> commas + ( map + fromDuckDBExpression + (project (Proxy @DuckDB) (mkProjection tblQ) "t") + ) + where + tblQ :: table (QExpr DuckDB ()) + tblQ = + changeBeamRep + ( \(Columnar' f) -> + Columnar' (QExpr (\_ -> fieldE (unqualifiedField (_fieldName f)))) + ) + tblSettings diff --git a/src/Database/Beam/DuckDB/Syntax/Extensions/Copy.hs b/src/Database/Beam/DuckDB/Syntax/Extensions/Copy.hs new file mode 100644 --- /dev/null +++ b/src/Database/Beam/DuckDB/Syntax/Extensions/Copy.hs @@ -0,0 +1,500 @@ +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeFamilies #-} +{-# OPTIONS_GHC -Wno-unused-top-binds #-} + +module Database.Beam.DuckDB.Syntax.Extensions.Copy + ( DuckDBCopyToSyntax (..), + DuckDBCopyFromSyntax (..), + + -- * COPY options + DuckDBCopyToOptions, + DuckDBCopyFromOptions, + + -- ** CSV + copyToCSV, + copyToCSVWith, + DuckDBCSVCopyToOptions (..), + defaultDuckDBCSVCopyToOptions, + copyFromCSV, + copyFromCSVWith, + DuckDBCSVCopyFromOptions (..), + defaultDuckDBCSVCopyFromOptions, + + -- ** Parquet + copyToParquet, + copyToParquetWith, + DuckDBParquetCopyToOptions (..), + ParquetCompression (..), + defaultDuckDBParquetCopyToOptions, + copyFromParquet, + copyFromParquetWith, + DuckDBParquetCopyFromOptions (..), + defaultDuckDBParquetCopyFromOptions, + + -- ** JSON + copyToJSON, + copyToJSONWith, + DuckDBJSONCopyToOptions (..), + JSONCompression (..), + defaultDuckDBJSONCopyToOptions, + copyFromJSON, + copyFromJSONWith, + DuckDBJSONCopyFromOptions (..), + defaultDuckDBJSONCopyFromOptions, + ) +where + +import Data.Coerce (coerce) +import Data.List.NonEmpty (NonEmpty) +import qualified Data.List.NonEmpty as NE +import Data.Maybe (catMaybes) +import Data.Text (Text) +import qualified Data.Text as Text +import Database.Beam.Backend.SQL.BeamExtensions +import Database.Beam.DuckDB.Syntax (DuckDBSelectSyntax (..)) +import Database.Beam.DuckDB.Syntax.Builder (DuckDBSyntax, commas, emit, emit', emitChar, emitCharLit, emitIntegral, emitTextLit, parens, quotedIdentifier, sepBy) + +-- | DuckDB @COPY ... TO@ source syntax. Wraps the table name (with +-- optional column list) or a @SELECT@ subquery. +newtype DuckDBCopyToSourceSyntax = DuckDBCopyToSourceSyntax {fromDuckDBCopyToSource :: DuckDBSyntax} + +-- | DuckDB @COPY ... TO@ statement syntax. +newtype DuckDBCopyToSyntax = DuckDBCopyToSyntax {fromDuckDBCopyTo :: DuckDBSyntax} + +instance IsSqlCopyToSourceSyntax DuckDBCopyToSourceSyntax where + type SqlCopyToSourceSelectSyntax DuckDBCopyToSourceSyntax = DuckDBSelectSyntax + + copyTableToSyntax mSchema tableName' mColumns = + DuckDBCopyToSourceSyntax $ + maybe (emit tableName') (\schema -> emit schema <> emitChar '.' <> emit tableName') mSchema + <> case mColumns of + Nothing -> mempty + Just cols -> parens $ commas (map quotedIdentifier (NE.toList cols)) + + copySelectToSyntax (DuckDBSelectSyntax select) = DuckDBCopyToSourceSyntax $ parens select + +instance IsSqlCopyToSyntax DuckDBCopyToSyntax where + type SqlCopyToSourceSyntax DuckDBCopyToSyntax = DuckDBCopyToSourceSyntax + type SqlCopyToParams DuckDBCopyToSyntax = DuckDBCopyToOptions + + copyToStmt source options = + let (filepath, optionsSyntax) = emitDuckDBCopyToOptions options + in DuckDBCopyToSyntax $ + mconcat + [ emit "COPY ", + coerce source, + emit " TO ", + emitChar '\'', + emit (Text.pack filepath), + emitChar '\'', + optionsSyntax + ] + +-- | DuckDB @COPY ... FROM@ source syntax. Just the table name plus optional +-- column list. +newtype DuckDBCopyFromSourceSyntax = DuckDBCopyFromSourceSyntax {fromDuckDBCopyFromSource :: DuckDBSyntax} + +-- | DuckDB @COPY ... FROM@ statement syntax. +newtype DuckDBCopyFromSyntax = DuckDBCopyFromSyntax {fromDuckDBCopyFrom :: DuckDBSyntax} + +instance IsSqlCopyFromSourceSyntax DuckDBCopyFromSourceSyntax where + copyTableFromSyntax mSchema tableName' mColumns = + DuckDBCopyFromSourceSyntax $ + maybe (emit tableName') (\schema -> emit schema <> emitChar '.' <> emit tableName') mSchema + <> case mColumns of + Nothing -> mempty + Just cols -> parens $ commas (map quotedIdentifier (NE.toList cols)) + +instance IsSqlCopyFromSyntax DuckDBCopyFromSyntax where + type SqlCopyFromSourceSyntax DuckDBCopyFromSyntax = DuckDBCopyFromSourceSyntax + type SqlCopyFromParams DuckDBCopyFromSyntax = DuckDBCopyFromOptions + + copyFromStmt source options = + let (filepath, optionsSyntax) = emitDuckDBCopyFromOptions options + in DuckDBCopyFromSyntax $ + mconcat + [ emit "COPY ", + coerce source, + emit " FROM ", + emitChar '\'', + emit (Text.pack filepath), + emitChar '\'', + optionsSyntax + ] + +-- | Options for DuckDB's @COPY ... TO@ statement, tagged by the output format. +-- +-- Use 'copyToCSV', 'copyToParquet', and 'copyToJSON' to construct values of type 'DuckDBCopyToOptions'. +-- +-- See +-- for the full set of options each format supports. +-- +-- @since 0.3.1.0 +data DuckDBCopyToOptions + = DuckDBCopyToCSV !FilePath !DuckDBCSVCopyToOptions + | DuckDBCopyToParquet !FilePath !DuckDBParquetCopyToOptions + | DuckDBCopyToJSON !FilePath !DuckDBJSONCopyToOptions + deriving (Eq, Show) + +-- | Copy to a CSV file at the given path with default options. +-- Use 'copyToCSVWith' to override. +-- +-- @since 0.3.1.0 +copyToCSV :: FilePath -> DuckDBCopyToOptions +copyToCSV path = DuckDBCopyToCSV path defaultDuckDBCSVCopyToOptions + +-- | Copy to a CSV file at the given path with the given options. +-- +-- @since 0.3.1.0 +copyToCSVWith :: FilePath -> DuckDBCSVCopyToOptions -> DuckDBCopyToOptions +copyToCSVWith = DuckDBCopyToCSV + +-- | Copy to a Parquet file at the given path with default options. +-- Use 'copyToParquetWith' to override. +-- +-- @since 0.3.1.0 +copyToParquet :: FilePath -> DuckDBCopyToOptions +copyToParquet path = DuckDBCopyToParquet path defaultDuckDBParquetCopyToOptions + +-- | Copy to a Parquet file at the given path with the given options. +-- +-- @since 0.3.1.0 +copyToParquetWith :: FilePath -> DuckDBParquetCopyToOptions -> DuckDBCopyToOptions +copyToParquetWith = DuckDBCopyToParquet + +-- | Copy to a JSON file at the given path with default options. +-- Use 'copyToJSONWith' to override. +-- +-- @since 0.3.1.0 +copyToJSON :: FilePath -> DuckDBCopyToOptions +copyToJSON path = DuckDBCopyToJSON path defaultDuckDBJSONCopyToOptions + +-- | Copy to a JSON file at the given path with the given options. +-- +-- @since 0.3.1.0 +copyToJSONWith :: FilePath -> DuckDBJSONCopyToOptions -> DuckDBCopyToOptions +copyToJSONWith = DuckDBCopyToJSON + +-- | CSV-specific options for @COPY ... TO@. +-- +-- @since 0.3.1.0 +data DuckDBCSVCopyToOptions = DuckDBCSVCopyToOptions + { csvCopyToHeader :: Maybe Bool, + csvCopyToDelimiter :: Maybe Text, + csvCopyToQuote :: Maybe Char, + csvCopyToEscape :: Maybe Char, + csvCopyToNullStr :: Maybe Text, + csvCopyToDateformat :: Maybe Text, + csvCopyToTimestampformat :: Maybe Text, + csvCopyToOverwriteOrIgnore :: Maybe Bool + } + deriving (Eq, Show) + +-- | All-'Nothing' 'DuckDBCSVCopyToOptions'; a sensible starting point for +-- record-update overrides. +-- +-- @since 0.3.1.0 +defaultDuckDBCSVCopyToOptions :: DuckDBCSVCopyToOptions +defaultDuckDBCSVCopyToOptions = + DuckDBCSVCopyToOptions + { csvCopyToHeader = Nothing, + csvCopyToDelimiter = Nothing, + csvCopyToQuote = Nothing, + csvCopyToEscape = Nothing, + csvCopyToNullStr = Nothing, + csvCopyToDateformat = Nothing, + csvCopyToTimestampformat = Nothing, + csvCopyToOverwriteOrIgnore = Nothing + } + +-- | Parquet-specific options for @COPY ... TO@. +-- +-- @since 0.3.1.0 +data DuckDBParquetCopyToOptions = DuckDBParquetCopyToOptions + { parquetCopyToCompression :: Maybe ParquetCompression, + parquetCopyToOverwriteOrIgnore :: Maybe Bool + } + deriving (Eq, Show) + +-- | All-'Nothing' 'DuckDBParquetCopyToOptions'; a sensible starting point +-- for record-update overrides. +-- +-- @since 0.3.1.0 +defaultDuckDBParquetCopyToOptions :: DuckDBParquetCopyToOptions +defaultDuckDBParquetCopyToOptions = + DuckDBParquetCopyToOptions + { parquetCopyToCompression = Nothing, + parquetCopyToOverwriteOrIgnore = Nothing + } + +-- | JSON-specific options for @COPY ... TO@. +-- +-- @since 0.3.1.0 +data DuckDBJSONCopyToOptions = DuckDBJSONCopyToOptions + { jsonCopyToCompression :: Maybe JSONCompression, + jsonCopyToOverwriteOrIgnore :: Maybe Bool + } + deriving (Eq, Show) + +-- | All-'Nothing' 'DuckDBJSONCopyToOptions'; a sensible starting point for +-- record-update overrides. +-- +-- @since 0.3.1.0 +defaultDuckDBJSONCopyToOptions :: DuckDBJSONCopyToOptions +defaultDuckDBJSONCopyToOptions = + DuckDBJSONCopyToOptions + { jsonCopyToCompression = Nothing, + jsonCopyToOverwriteOrIgnore = Nothing + } + +-- | Options for DuckDB's @COPY ... FROM@ statement, tagged by the input format. +-- +-- Use 'copyFromCSV', 'copyFromParquet', and 'copyFromJSON' to construct values of type 'DuckDBCopyFromOptions'. +-- +-- See +-- for the full set of options each format supports. +-- +-- @since 0.3.1.0 +data DuckDBCopyFromOptions + = DuckDBCopyFromCSV !FilePath !DuckDBCSVCopyFromOptions + | DuckDBCopyFromParquet !FilePath !DuckDBParquetCopyFromOptions + | DuckDBCopyFromJSON !FilePath !DuckDBJSONCopyFromOptions + deriving (Eq, Show) + +-- | Copy from a CSV file at the given path with default options. +-- Use 'copyFromCSVWith' to override. +-- +-- @since 0.3.1.0 +copyFromCSV :: FilePath -> DuckDBCopyFromOptions +copyFromCSV path = DuckDBCopyFromCSV path defaultDuckDBCSVCopyFromOptions + +-- | Copy from a CSV file at the given path with the given options. +-- +-- @since 0.3.1.0 +copyFromCSVWith :: FilePath -> DuckDBCSVCopyFromOptions -> DuckDBCopyFromOptions +copyFromCSVWith = DuckDBCopyFromCSV + +-- | Copy from a Parquet file at the given path with default options. +-- Use 'copyFromParquetWith' to override. +-- +-- @since 0.3.1.0 +copyFromParquet :: FilePath -> DuckDBCopyFromOptions +copyFromParquet path = DuckDBCopyFromParquet path defaultDuckDBParquetCopyFromOptions + +-- | Copy from a Parquet file at the given path with the given options. +-- +-- @since 0.3.1.0 +copyFromParquetWith :: FilePath -> DuckDBParquetCopyFromOptions -> DuckDBCopyFromOptions +copyFromParquetWith = DuckDBCopyFromParquet + +-- | Copy from a JSON file at the given path with default options. +-- Use 'copyFromJSONWith' to override. +-- +-- @since 0.3.1.0 +copyFromJSON :: FilePath -> DuckDBCopyFromOptions +copyFromJSON path = DuckDBCopyFromJSON path defaultDuckDBJSONCopyFromOptions + +-- | Copy from a JSON file at the given path with the given options. +-- +-- @since 0.3.1.0 +copyFromJSONWith :: FilePath -> DuckDBJSONCopyFromOptions -> DuckDBCopyFromOptions +copyFromJSONWith = DuckDBCopyFromJSON + +-- | CSV-specific options for @COPY ... FROM@. +-- +-- @since 0.3.1.0 +data DuckDBCSVCopyFromOptions = DuckDBCSVCopyFromOptions + { csvCopyFromHeader :: Maybe Bool, + csvCopyFromDelimiter :: Maybe Text, + csvCopyFromQuote :: Maybe Char, + csvCopyFromEscape :: Maybe Char, + csvCopyFromNullStr :: Maybe Text, + csvCopyFromDateformat :: Maybe Text, + csvCopyFromTimestampformat :: Maybe Text, + csvCopyFromAutoDetect :: Maybe Bool, + csvCopyFromIgnoreErrors :: Maybe Bool, + csvCopyFromSkip :: Maybe Int, + csvCopyFromAllVarchar :: Maybe Bool + } + deriving (Eq, Show) + +-- | All-'Nothing' 'DuckDBCSVCopyFromOptions'; a sensible starting point for +-- record-update overrides. +-- +-- @since 0.3.1.0 +defaultDuckDBCSVCopyFromOptions :: DuckDBCSVCopyFromOptions +defaultDuckDBCSVCopyFromOptions = + DuckDBCSVCopyFromOptions + { csvCopyFromHeader = Nothing, + csvCopyFromDelimiter = Nothing, + csvCopyFromQuote = Nothing, + csvCopyFromEscape = Nothing, + csvCopyFromNullStr = Nothing, + csvCopyFromDateformat = Nothing, + csvCopyFromTimestampformat = Nothing, + csvCopyFromAutoDetect = Nothing, + csvCopyFromIgnoreErrors = Nothing, + csvCopyFromSkip = Nothing, + csvCopyFromAllVarchar = Nothing + } + +-- | Parquet-specific options for @COPY ... FROM@. DuckDB's Parquet reader has +-- almost no read-time options; this is here so the Parquet format can be +-- selected with consistent defaults. +-- +-- @since 0.3.1.0 +data DuckDBParquetCopyFromOptions = DuckDBParquetCopyFromOptions + { + } + deriving (Eq, Show) + +-- | The (only) 'DuckDBParquetCopyFromOptions' value. +-- +-- @since 0.3.1.0 +defaultDuckDBParquetCopyFromOptions :: DuckDBParquetCopyFromOptions +defaultDuckDBParquetCopyFromOptions = DuckDBParquetCopyFromOptions + +-- | JSON-specific options for @COPY ... FROM@. +-- +-- @since 0.3.1.0 +data DuckDBJSONCopyFromOptions = DuckDBJSONCopyFromOptions + { jsonCopyFromAutoDetect :: Maybe Bool, + -- | @COMPRESSION@: how the input file is compressed. DuckDB's JSON + -- reader normally auto-detects this from the file extension, but + -- when the path doesn't carry the usual @.gz@ / @.zst@ suffix the + -- compression has to be specified explicitly so it can be + -- decompressed on read. Round-tripping a compressed JSON write + -- requires this to match the matching @TO@-side compression. + jsonCopyFromCompression :: Maybe JSONCompression + } + deriving (Eq, Show) + +-- | All-'Nothing' 'DuckDBJSONCopyFromOptions'; a sensible starting point for +-- record-update overrides. +-- +-- @since 0.3.1.0 +defaultDuckDBJSONCopyFromOptions :: DuckDBJSONCopyFromOptions +defaultDuckDBJSONCopyFromOptions = + DuckDBJSONCopyFromOptions + { jsonCopyFromAutoDetect = Nothing, + jsonCopyFromCompression = Nothing + } + +-- | Compression algorithms accepted by DuckDB's Parquet writer (and reader). +-- +-- See ; +-- the JSON format accepts a strictly smaller set of values, modelled +-- separately by 'JSONCompression'. +-- +-- @since 0.3.1.0 +data ParquetCompression + = ParquetUncompressed + | ParquetSnappy + | ParquetGzip + | ParquetZstd + deriving (Eq, Show, Ord, Enum, Bounded) + +-- | Compression algorithms accepted by DuckDB's JSON writer (and reader). +-- +-- DuckDB's JSON format does not support @SNAPPY@, so 'JSONCompression' +-- omits it. See . +-- +-- @since 0.3.1.0 +data JSONCompression + = JSONUncompressed + | JSONGzip + | JSONZstd + deriving (Eq, Show, Ord, Enum, Bounded) + +emitDuckDBCopyToOptions :: DuckDBCopyToOptions -> (FilePath, DuckDBSyntax) +emitDuckDBCopyToOptions (DuckDBCopyToCSV fp o) = (fp, emitOptionList (emit "CSV") (csvCopyToFields o)) +emitDuckDBCopyToOptions (DuckDBCopyToParquet fp o) = (fp, emitOptionList (emit "PARQUET") (parquetCopyToFields o)) +emitDuckDBCopyToOptions (DuckDBCopyToJSON fp o) = (fp, emitOptionList (emit "JSON") (jsonCopyToFields o)) + +emitDuckDBCopyFromOptions :: DuckDBCopyFromOptions -> (FilePath, DuckDBSyntax) +emitDuckDBCopyFromOptions (DuckDBCopyFromCSV fp o) = (fp, emitOptionList (emit "CSV") (csvCopyFromFields o)) +emitDuckDBCopyFromOptions (DuckDBCopyFromParquet fp o) = (fp, emitOptionList (emit "PARQUET") (parquetCopyFromFields o)) +emitDuckDBCopyFromOptions (DuckDBCopyFromJSON fp o) = (fp, emitOptionList (emit "JSON") (jsonCopyFromFields o)) + +-- | Render a non-empty list of options like @ (FORMAT FMT, OPT1 …, OPT2 …)@. +-- The @FORMAT@ entry is always present, since the constructor of the options +-- value pins it. +emitOptionList :: DuckDBSyntax -> [DuckDBSyntax] -> DuckDBSyntax +emitOptionList fmt items = + emit " (" + <> sepBy (emit ", ") ((emit "FORMAT " <> fmt) : items) + <> emitChar ')' + +csvCopyToFields :: DuckDBCSVCopyToOptions -> [DuckDBSyntax] +csvCopyToFields opts = + catMaybes + [ fmap (\x -> emit "HEADER " <> emit' x) (csvCopyToHeader opts), + fmap (\x -> emit "DELIMITER " <> emitTextLit x) (csvCopyToDelimiter opts), + fmap (\x -> emit "QUOTE " <> emitCharLit x) (csvCopyToQuote opts), + fmap (\x -> emit "ESCAPE " <> emitCharLit x) (csvCopyToEscape opts), + fmap (\x -> emit "NULLSTR " <> emitTextLit x) (csvCopyToNullStr opts), + fmap (\x -> emit "DATEFORMAT " <> emitTextLit x) (csvCopyToDateformat opts), + fmap (\x -> emit "TIMESTAMPFORMAT " <> emitTextLit x) (csvCopyToTimestampformat opts), + fmap (\x -> emit "OVERWRITE_OR_IGNORE " <> emit' x) (csvCopyToOverwriteOrIgnore opts) + ] + +parquetCopyToFields :: DuckDBParquetCopyToOptions -> [DuckDBSyntax] +parquetCopyToFields opts = + catMaybes + [ fmap (\x -> emit "COMPRESSION " <> emitParquetCompression x) (parquetCopyToCompression opts), + fmap (\x -> emit "OVERWRITE_OR_IGNORE " <> emit' x) (parquetCopyToOverwriteOrIgnore opts) + ] + +jsonCopyToFields :: DuckDBJSONCopyToOptions -> [DuckDBSyntax] +jsonCopyToFields opts = + catMaybes + [ fmap (\x -> emit "COMPRESSION " <> emitJSONCompression x) (jsonCopyToCompression opts), + fmap (\x -> emit "OVERWRITE_OR_IGNORE " <> emit' x) (jsonCopyToOverwriteOrIgnore opts) + ] + +csvCopyFromFields :: DuckDBCSVCopyFromOptions -> [DuckDBSyntax] +csvCopyFromFields opts = + catMaybes + [ fmap (\x -> emit "HEADER " <> emit' x) (csvCopyFromHeader opts), + fmap (\x -> emit "DELIMITER " <> emitTextLit x) (csvCopyFromDelimiter opts), + fmap (\x -> emit "QUOTE " <> emitCharLit x) (csvCopyFromQuote opts), + fmap (\x -> emit "ESCAPE " <> emitCharLit x) (csvCopyFromEscape opts), + fmap (\x -> emit "NULLSTR " <> emitTextLit x) (csvCopyFromNullStr opts), + fmap (\x -> emit "DATEFORMAT " <> emitTextLit x) (csvCopyFromDateformat opts), + fmap (\x -> emit "TIMESTAMPFORMAT " <> emitTextLit x) (csvCopyFromTimestampformat opts), + fmap (\x -> emit "AUTO_DETECT " <> emit' x) (csvCopyFromAutoDetect opts), + fmap (\x -> emit "IGNORE_ERRORS " <> emit' x) (csvCopyFromIgnoreErrors opts), + fmap (\x -> emit "SKIP " <> emitIntegral x) (csvCopyFromSkip opts), + fmap (\x -> emit "ALL_VARCHAR " <> emit' x) (csvCopyFromAllVarchar opts) + ] + +parquetCopyFromFields :: DuckDBParquetCopyFromOptions -> [DuckDBSyntax] +parquetCopyFromFields _ = [] + +jsonCopyFromFields :: DuckDBJSONCopyFromOptions -> [DuckDBSyntax] +jsonCopyFromFields opts = + catMaybes + [ fmap (\x -> emit "AUTO_DETECT " <> emit' x) (jsonCopyFromAutoDetect opts), + fmap (\x -> emit "COMPRESSION " <> emitJSONCompression x) (jsonCopyFromCompression opts) + ] + +emitPartitionBy :: NonEmpty Text -> DuckDBSyntax +emitPartitionBy cols = + emit "PARTITION_BY (" + <> commas (map quotedIdentifier (NE.toList cols)) + <> emitChar ')' + +emitParquetCompression :: ParquetCompression -> DuckDBSyntax +emitParquetCompression ParquetUncompressed = emit "UNCOMPRESSED" +emitParquetCompression ParquetSnappy = emit "SNAPPY" +emitParquetCompression ParquetGzip = emit "GZIP" +emitParquetCompression ParquetZstd = emit "ZSTD" + +emitJSONCompression :: JSONCompression -> DuckDBSyntax +emitJSONCompression JSONUncompressed = emit "UNCOMPRESSED" +emitJSONCompression JSONGzip = emit "GZIP" +emitJSONCompression JSONZstd = emit "ZSTD" diff --git a/src/Database/Beam/DuckDB/Syntax/Extensions/DataSource.hs b/src/Database/Beam/DuckDB/Syntax/Extensions/DataSource.hs --- a/src/Database/Beam/DuckDB/Syntax/Extensions/DataSource.hs +++ b/src/Database/Beam/DuckDB/Syntax/Extensions/DataSource.hs @@ -96,7 +96,7 @@ -- | Declare a CSV file (or files) as the source of data for a database. -- -- See 'csv' if you want to use default options. --- +-- -- Note: Be careful to sanitize any user input that is then used as -- a data source. csvWith :: @@ -138,7 +138,7 @@ -- | Define an Apache Iceberg table with default table options. -- -- See 'icebergTableWith' if you want to change the default options. --- +-- -- Note: Be careful to sanitize any user input that is then used as -- a data source. icebergTable :: @@ -150,7 +150,7 @@ -- | Define an Apache Iceberg table with options. -- -- See 'icebergTable' if you want to use default options. --- +-- -- Note: Be careful to sanitize any user input that is then used as -- a data source. icebergTableWith :: @@ -176,7 +176,7 @@ defaultIcebergTableOptions = IcebergTableOptions Nothing Nothing -- | Declare a Parquet file(s) or glob(s) as the source of data for a database. --- +-- -- Note: Be careful to sanitize any user input that is then used as -- a data source. parquet :: diff --git a/src/Database/Beam/DuckDB/Syntax/Extensions/InsertOnConflict.hs b/src/Database/Beam/DuckDB/Syntax/Extensions/InsertOnConflict.hs new file mode 100644 --- /dev/null +++ b/src/Database/Beam/DuckDB/Syntax/Extensions/InsertOnConflict.hs @@ -0,0 +1,227 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE InstanceSigs #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} + +module Database.Beam.DuckDB.Syntax.Extensions.InsertOnConflict + ( BeamHasInsertOnConflict + ( SqlConflictTarget, + SqlConflictAction, + insertOnConflict, + anyConflict, + conflictingFields, + onConflictDoNothing, + onConflictUpdateSet, + onConflictUpdateSetWhere + ), + conflictingFieldsWhere, + onConflictUpdateAll, + onConflictUpdateInstead, + UnsupportedConflictingFieldsWhereOnDuckDB, + ) +where + +import Data.Proxy (Proxy (..)) +import Database.Beam (Beamable, TableEntity) +import Database.Beam.Backend.SQL + ( IsSql92ExpressionSyntax (..), + IsSql92FieldNameSyntax (..), + IsSql92InsertSyntax (..), + IsSql92TableNameSyntax (..), + ) +import Database.Beam.Backend.SQL.BeamExtensions + ( BeamHasInsertOnConflict (SqlConflictAction, SqlConflictTarget), + onConflictUpdateAll, + onConflictUpdateInstead, + ) +import qualified Database.Beam.Backend.SQL.BeamExtensions as Beam +import Database.Beam.DuckDB.Backend (DuckDB) +import Database.Beam.DuckDB.Syntax + ( DuckDBExpressionSyntax (..), + DuckDBFieldNameSyntax (..), + DuckDBInsertSyntax (..), + ) +import Database.Beam.DuckDB.Syntax.Builder + ( DuckDBSyntax, + commas, + emit, + parens, + ) +import Database.Beam.Query (SqlInsert (..), SqlInsertValues (..)) +import Database.Beam.Query.Internal + ( Projectible, + QAssignment (..), + QExpr, + QField (..), + QGenExpr (..), + QInternal, + project, + ) +import Database.Beam.Schema.Tables + ( Columnar' (..), + DatabaseEntity (..), + DatabaseEntityDescriptor (..), + TableField, + allBeamValues, + changeBeamRep, + _fieldName, + ) +import GHC.TypeLits (ErrorMessage (..), TypeError) + +-- | @since 0.3.1.0 +instance Beam.BeamHasInsertOnConflict DuckDB where + newtype SqlConflictTarget DuckDB table = DuckDBConflictTarget + {unDuckDBConflictTarget :: table (QExpr DuckDB QInternal) -> DuckDBSyntax} + newtype SqlConflictAction DuckDB table = DuckDBConflictAction + {unDuckDBConflictAction :: forall s. table (QField s) -> DuckDBSyntax} + + insertOnConflict :: + forall table db s. + (Beamable table) => + DatabaseEntity DuckDB db (TableEntity table) -> + SqlInsertValues DuckDB (table (QExpr DuckDB s)) -> + SqlConflictTarget DuckDB table -> + SqlConflictAction DuckDB table -> + SqlInsert DuckDB table + insertOnConflict (DatabaseEntity dt) values target action = case values of + SqlInsertValuesEmpty -> SqlInsertNoRows + SqlInsertValues vs -> + let tblSettings = dbTableSettings dt + mkField :: + forall a. + Columnar' (TableField table) a -> + Columnar' (QField QInternal) a + mkField (Columnar' fd) = Columnar' $ QField False (dbTableCurrentName dt) (_fieldName fd) + tblFields = changeBeamRep mkField tblSettings + tblExprs = + changeBeamRep + ( \(Columnar' (QField _ _ nm)) -> + Columnar' (QExpr (\_ -> fieldE (unqualifiedField nm))) + ) + tblFields + fieldList = allBeamValues (\(Columnar' f) -> _fieldName f) tblSettings + DuckDBInsertSyntax baseInsert = + insertStmt (tableName (dbTableSchema dt) (dbTableCurrentName dt)) fieldList vs + in SqlInsert tblSettings $ + DuckDBInsertSyntax $ + baseInsert + <> emit " ON CONFLICT " + <> unDuckDBConflictTarget target tblExprs + <> unDuckDBConflictAction action tblFields + + anyConflict = DuckDBConflictTarget (const mempty) + + conflictingFields makeProjection = + DuckDBConflictTarget $ \tbl -> + parens + ( commas $ + map fromDuckDBExpression $ + project (Proxy @DuckDB) (makeProjection tbl) "t" + ) + <> emit " " + + -- \| DuckDB does not support partial-index conflict targets, so calling this is always a bug. + -- The supported way to invoke this method from @Database.Beam.DuckDB@ goes through the shadowing + -- 'conflictingFieldsWhere' below, which produces a /compile-time/ error pointing users at 'onConflictUpdateSetWhere'. + conflictingFieldsWhere _ _ = + error + "Database.Beam.DuckDB: `conflictingFieldsWhere` is not supported on DuckDB. \ + \DuckDB does not support partial-index conflict targets — use \ + \`onConflictUpdateSetWhere` to filter the rows that are updated instead." + + onConflictDoNothing = DuckDBConflictAction (const (emit "DO NOTHING")) + + onConflictUpdateSet mkAssignments = + DuckDBConflictAction $ \tbl -> + let QAssignment assignments = mkAssignments tbl (excluded tbl) + assignmentSyntaxes = + [ fromDuckDBFieldName fieldNm + <> emit "=" + <> parens (fromDuckDBExpression expr) + | (fieldNm, expr) <- assignments + ] + in emit "DO UPDATE SET " <> commas assignmentSyntaxes + + onConflictUpdateSetWhere mkAssignments where_ = + DuckDBConflictAction $ \tbl -> + let QAssignment assignments = mkAssignments tbl (excluded tbl) + QExpr where_' = where_ tbl (excluded tbl) + assignmentSyntaxes = + [ fromDuckDBFieldName fieldNm + <> emit "=" + <> parens (fromDuckDBExpression expr) + | (fieldNm, expr) <- assignments + ] + in emit "DO UPDATE SET " + <> commas assignmentSyntaxes + <> emit " WHERE " + <> fromDuckDBExpression (where_' "t") + +-- | Build a table reference whose fields point to the special @excluded@ +-- pseudo-table available inside DuckDB's @ON CONFLICT DO UPDATE@ — i.e. the +-- row that would have been inserted had there been no conflict. +excluded :: + (Beamable table) => + table (QField s) -> + table (QExpr DuckDB s) +excluded = + changeBeamRep + ( \(Columnar' (QField _ _ nm)) -> + Columnar' (QExpr (\_ -> fieldE (qualifiedField "excluded" nm))) + ) + +-- | A class with a single 'TypeError'-guarded instance, used purely to +-- defer firing of the type error until a /use site/ of +-- 'conflictingFieldsWhere'. The instance context is only discharged +-- when GHC actually selects this instance, which means defining a +-- method body in terms of it (as @'conflictingFieldsWhere' = …@) does +-- not trip the error at module load time. +-- +-- = References +-- +-- DuckDB's binder unconditionally rejects @ON CONFLICT (cols) WHERE …@: +-- see @GenerateMergeInto@ in +-- , +-- which throws @"ON CONFLICT WHERE clause is only supported in DO UPDATE +-- SET ... WHERE ... The WHERE clause after the conflict columns is used +-- for partial indexes which are not supported."@. The user-facing +-- +-- describe both @WHERE@ positions, but only the trailing-update form +-- actually binds; the index-predicate form (used in PostgreSQL to +-- arbitrate between /partial/ unique indexes) has nothing to match +-- against in DuckDB. +-- +-- @since 0.3.1.0 +class UnsupportedConflictingFieldsWhereOnDuckDB where + -- | A 'conflictingFieldsWhere' that always fails to compile when + -- used against DuckDB. Re-exported from "Database.Beam.DuckDB" so + -- that it shadows the class method of the same name from + -- "Database.Beam.Backend.SQL.BeamExtensions". + -- + -- @since 0.3.1.0 + conflictingFieldsWhere :: + forall table proj. + (Projectible DuckDB proj, Beamable table) => + (table (QExpr DuckDB QInternal) -> proj) -> + (forall s. table (QExpr DuckDB s) -> QExpr DuckDB s Bool) -> + SqlConflictTarget DuckDB table + +instance + (TypeError DuckDBConflictingFieldsWhereError) => + UnsupportedConflictingFieldsWhereOnDuckDB + where + conflictingFieldsWhere _ _ = error "unreachable: TypeError fires first" + +type DuckDBConflictingFieldsWhereError = + 'Text "`conflictingFieldsWhere` is not supported on the DuckDB backend." + ':$$: 'Text "DuckDB does not support partial-index conflict targets" + ':$$: 'Text "(the WHERE clause between the conflict columns and DO)." + ':$$: 'Text "Use `onConflictUpdateSetWhere` to filter the rows that are updated instead." diff --git a/tests/Database/Beam/DuckDB/Test/Extensions.hs b/tests/Database/Beam/DuckDB/Test/Extensions.hs --- a/tests/Database/Beam/DuckDB/Test/Extensions.hs +++ b/tests/Database/Beam/DuckDB/Test/Extensions.hs @@ -1,365 +1,17 @@ -{-# LANGUAGE DeriveAnyClass #-} -{-# LANGUAGE DeriveGeneric #-} -{-# LANGUAGE ExplicitForAll #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE StandaloneDeriving #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeFamilies #-} - module Database.Beam.DuckDB.Test.Extensions (tests) where -import Data.Int (Int32) -import Data.List.NonEmpty (NonEmpty ((:|))) -import Data.Text (Text) -import Data.Time (Day, fromGregorian) -import Database.Beam - ( Beamable, - Columnar, - Database, - DatabaseSettings, - Generic, - Identity, - Table (..), - aggregate_, - as_, - countAll_, - dbModification, - defaultDbSettings, - guard_, - min_, - runSelectReturningList, - runSelectReturningOne, - select, - tableModification, - val_, - withDbModification, - (==.), - ) -import Database.Beam.DuckDB - ( CSVOptions (..), - DataSourceEntity, - DuckDB, - allFromDataSource_, - allowMovedPaths, - csvWith, - dataSource, - defaultCSVOptions, - defaultIcebergTableOptions, - icebergTableWith, - modifyDataSourceFields, - parquet, - runBeamDuckDB, - ) -import Database.DuckDB.Simple (Connection, withConnection) +import qualified Database.Beam.DuckDB.Test.Extensions.Copy as Copy +import qualified Database.Beam.DuckDB.Test.Extensions.DataSource as DataSource +import qualified Database.Beam.DuckDB.Test.Extensions.InsertOnConflict as InsertOnConflict +import qualified Database.Beam.DuckDB.Test.Extensions.Returning as Returning import Test.Tasty (TestTree, testGroup) -import Test.Tasty.HUnit (testCase, (@?=)) tests :: TestTree tests = testGroup "Extensions" - [ testGroup - "Parquet" - [ testCountFromParquet, - testQueryFromParquet, - testCountFromMultipleParquetFiles, - testQueryFromMultipleParquetFiles - ], - testGroup - "Apache Iceberg" - [ testCountFromIceberg, - testParseDateFromIceberg - ], - testGroup - "CSV" - [ testCountFromCSV, - testQueryFromCSV - ] + [ Copy.tests, + DataSource.tests, + InsertOnConflict.tests, + Returning.tests ] - -testCountFromParquet :: TestTree -testCountFromParquet = testCase "Counting records from a Parquet file" $ do - results <- withTestDb $ \conn -> - runBeamDuckDB conn $ - runSelectReturningList $ - select $ do - exam <- allFromDataSource_ (_dbExams testDb) - pure (_examName exam) - results @?= ["alice", "bob", "carol", "dave"] - -testQueryFromParquet :: TestTree -testQueryFromParquet = testCase "Query from a Parquet file" $ do - results <- withTestDb $ \conn -> - runBeamDuckDB conn $ - runSelectReturningOne $ - select $ do - exam <- allFromDataSource_ (_dbExams testDb) - guard_ (_examId exam ==. 1) - pure (_examName exam) - results @?= Just "alice" - -testCountFromMultipleParquetFiles :: TestTree -testCountFromMultipleParquetFiles = testCase "Counting records from multiple Parquet file" $ do - results <- withTestDb $ \conn -> - runBeamDuckDB conn $ - runSelectReturningList $ - select $ do - exam <- allFromDataSource_ (_dbExamsMulti testDb) - pure (_examName exam) - results @?= ["alice", "bob", "carol", "dave", "erika", "francis", "genevieve", "hugo"] - -testQueryFromMultipleParquetFiles :: TestTree -testQueryFromMultipleParquetFiles = testCase "Query from multiple Parquet files" $ do - results <- withTestDb $ \conn -> - runBeamDuckDB conn $ - runSelectReturningOne $ - select $ do - exam <- allFromDataSource_ (_dbExamsMulti testDb) - guard_ (_examId exam ==. 5) - pure (_examName exam) - results @?= Just "erika" - -testCountFromIceberg :: TestTree -testCountFromIceberg = testCase "Counting records from an Apache Iceberg table" $ do - results <- withTestDb $ \conn -> - runBeamDuckDB conn $ - runSelectReturningOne $ - select $ - aggregate_ (\_ -> as_ @Int32 countAll_) (allFromDataSource_ (_dbLineItems testDb)) - results @?= Just 51793 -- From DuckDB's documentation - -testParseDateFromIceberg :: TestTree -testParseDateFromIceberg = testCase "Test parsing date columns from Apache Iceberg table" $ do - results <- withTestDb $ \conn -> - runBeamDuckDB conn $ - runSelectReturningOne $ - select $ - aggregate_ (min_ . _lineitemShipdate) (allFromDataSource_ (_dbLineItems testDb)) - - results @?= Just (Just (fromGregorian 1992 01 04)) - -testCountFromCSV :: TestTree -testCountFromCSV = testCase "Counting records from a CSV file" $ do - results <- withTestDb $ \conn -> - runBeamDuckDB conn $ - runSelectReturningOne $ - select $ - aggregate_ (\_ -> as_ @Int32 countAll_) (allFromDataSource_ (_dbFlights testDb)) - results @?= Just 3 - -testQueryFromCSV :: TestTree -testQueryFromCSV = testCase "Query from a CSV file" $ do - results <- withTestDb $ \conn -> - runBeamDuckDB conn $ - runSelectReturningOne $ - select $ do - flight <- allFromDataSource_ (_dbFlights testDb) - guard_ (_flightDate flight ==. val_ (fromGregorian 1988 01 01)) - pure flight - results - @?= Just - ( Flight - { _flightDate = fromGregorian 1988 01 01, - _flightUniqueCarrier = "AA", - _flightOriginCity = "New York, NY", - _flightDestCity = "Los Angeles, CA" - } - ) - -data ExamT f = Exam - { _examId :: Columnar f Int32, - _examName :: Columnar f Text, - _examScore :: Columnar f Double, - _examDate :: Columnar f Day - } - deriving (Generic) - -type Exam = ExamT Identity - -type ExamId = PrimaryKey ExamT Identity - -deriving instance Show ExamId - -deriving instance Eq ExamId - -deriving instance Ord ExamId - -deriving instance Show Exam - -deriving instance Eq Exam - -deriving instance Ord Exam - -instance Beamable ExamT - -instance Table ExamT where - data PrimaryKey ExamT f = ExamId (Columnar f Int32) - deriving (Generic) - primaryKey = ExamId . _examId - -instance Beamable (PrimaryKey ExamT) - -data LineItemT f = Lineitem - { _lineitemOrderkey :: Columnar f Int32, - _lineitemPartkey :: Columnar f Int32, - _lineitemSuppkey :: Columnar f Int32, - _lineitemLinenumber :: Columnar f Int32, - _lineitemQuantity :: Columnar f Double, - _lineitemExtendedprice :: Columnar f Double, - _lineitemDiscount :: Columnar f Double, - _lineitemTax :: Columnar f Double, - _lineitemReturnflag :: Columnar f Text, - _lineitemLinestatus :: Columnar f Text, - _lineitemShipdate :: Columnar f Day, - _lineitemCommitdate :: Columnar f Day, - _lineitemReceiptdate :: Columnar f Day, - _lineitemShipinstruct :: Columnar f Text, - _lineitemShipmode :: Columnar f Text, - _lineitemComment :: Columnar f Text - } - deriving (Generic) - -type Lineitem = LineItemT Identity - -deriving instance Show Lineitem - -deriving instance Eq Lineitem - -instance Beamable LineItemT - -instance Table LineItemT where - data PrimaryKey LineItemT f - = LineItemKey (Columnar f Int32) (Columnar f Int32) - deriving (Generic) - primaryKey li = - LineItemKey (_lineitemOrderkey li) (_lineitemLinenumber li) - -instance Beamable (PrimaryKey LineItemT) - -deriving instance Show (PrimaryKey LineItemT Identity) - -deriving instance Eq (PrimaryKey LineItemT Identity) - -data FlightT f = Flight - { _flightDate :: Columnar f Day, - _flightUniqueCarrier :: Columnar f Text, - _flightOriginCity :: Columnar f Text, - _flightDestCity :: Columnar f Text - } - deriving (Generic) - -type Flight = FlightT Identity - -deriving instance Show Flight - -deriving instance Eq Flight - -instance Beamable FlightT - -instance Table FlightT where - data PrimaryKey FlightT f - = FlightKey - (Columnar f Day) - (Columnar f Text) - (Columnar f Text) - (Columnar f Text) - deriving (Generic) - primaryKey fl = - FlightKey - (_flightDate fl) - (_flightUniqueCarrier fl) - (_flightOriginCity fl) - (_flightDestCity fl) - -instance Beamable (PrimaryKey FlightT) - -deriving instance Show (PrimaryKey FlightT Identity) - -deriving instance Eq (PrimaryKey FlightT Identity) - -data TestDB f = TestDB - { _dbExams :: f (DataSourceEntity ExamT), - _dbExamsMulti :: f (DataSourceEntity ExamT), -- set up with multiple parquet files - _dbLineItems :: f (DataSourceEntity LineItemT), - _dbFlights :: f (DataSourceEntity FlightT) - } - deriving (Generic, Database DuckDB) - -testDb :: DatabaseSettings DuckDB TestDB -testDb = - defaultDbSettings - `withDbModification` (dbModification @_ @DuckDB) - { _dbExams = - dataSource - (parquet ("tests/data/test1.parquet" :| [])) - <> modifyDataSourceFields - tableModification - { _examId = "id", - _examName = "name", - _examScore = "score", - _examDate = "exam_date" - }, - _dbExamsMulti = - dataSource - ( parquet - ("tests/data/test1.parquet" :| ["tests/data/test2.parquet"]) - ) - <> modifyDataSourceFields - tableModification - { _examId = "id", - _examName = "name", - _examScore = "score", - _examDate = "exam_date" - }, - _dbLineItems = - dataSource - ( icebergTableWith - "tests/data/lineitem_iceberg" - (defaultIcebergTableOptions {allowMovedPaths = Just True}) - ) - <> modifyDataSourceFields - tableModification - { _lineitemOrderkey = "l_orderkey", - _lineitemPartkey = "l_partkey", - _lineitemSuppkey = "l_suppkey", - _lineitemLinenumber = "l_linenumber", - _lineitemQuantity = "l_quantity", - _lineitemExtendedprice = "l_extendedprice", - _lineitemDiscount = "l_discount", - _lineitemTax = "l_tax", - _lineitemReturnflag = "l_returnflag", - _lineitemLinestatus = "l_linestatus", - _lineitemShipdate = "l_shipdate", - _lineitemCommitdate = "l_commitdate", - _lineitemReceiptdate = "l_receiptdate", - _lineitemShipinstruct = "l_shipinstruct", - _lineitemShipmode = "l_shipmode", - _lineitemComment = "l_comment" - }, - _dbFlights = - dataSource - ( csvWith - ("tests/data/flights.csv" :| []) - ( defaultCSVOptions - { header = Just True, - comment = Just '#', - delim = Just "|", - ignoreErrors = Just False - } - ) - ) - <> modifyDataSourceFields - tableModification - { _flightDate = "FlightDate", - _flightUniqueCarrier = "UniqueCarrier", - _flightOriginCity = "OriginCityName", - _flightDestCity = "DestCityName" - } - } - -withTestDb :: - (Connection -> IO a) -> - IO a -withTestDb = - withConnection ":memory:" diff --git a/tests/Database/Beam/DuckDB/Test/Extensions/Copy.hs b/tests/Database/Beam/DuckDB/Test/Extensions/Copy.hs new file mode 100644 --- /dev/null +++ b/tests/Database/Beam/DuckDB/Test/Extensions/Copy.hs @@ -0,0 +1,586 @@ +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TypeFamilies #-} + +module Database.Beam.DuckDB.Test.Extensions.Copy (tests) where + +import Control.Monad (guard, void) +import Control.Monad.IO.Class (liftIO) +import qualified Data.Char as Char +import Data.Int (Int32) +import qualified Data.List as List +import Data.Text (Text) +import qualified Data.Text as Text +import qualified Data.Text.IO as Text.IO +import Database.Beam + ( Beamable, + Columnar, + Database, + DatabaseSettings, + Generic, + Identity, + Table (..), + TableEntity, + all_, + asc_, + dbModification, + defaultDbSettings, + guard_, + insert, + insertValues, + modifyTableFields, + orderBy_, + runInsert, + runSelectReturningList, + select, + tableModification, + val_, + withDbModification, + (>.), + ) +import Database.Beam.Backend.SQL.BeamExtensions + ( copySelectTo, + copyTableFrom, + copyTableTo, + runCopyFrom, + runCopyTo, + ) +import Database.Beam.DuckDB + ( DuckDB, + DuckDBCSVCopyFromOptions (..), + DuckDBCSVCopyToOptions (..), + DuckDBJSONCopyFromOptions (..), + DuckDBJSONCopyToOptions (..), + DuckDBParquetCopyToOptions (..), + JSONCompression (..), + ParquetCompression (..), + copyFromCSV, + copyFromCSVWith, + copyFromJSONWith, + copyFromParquet, + copyFromParquetWith, + copyToCSV, + copyToCSVWith, + copyToJSONWith, + copyToParquetWith, + defaultDuckDBCSVCopyFromOptions, + defaultDuckDBCSVCopyToOptions, + defaultDuckDBJSONCopyFromOptions, + defaultDuckDBJSONCopyToOptions, + defaultDuckDBParquetCopyFromOptions, + defaultDuckDBParquetCopyToOptions, + runBeamDuckDB, + ) +import Database.DuckDB.Simple (Connection, execute_, withConnection) +import Hedgehog (Gen, forAll, property, (===)) +import qualified Hedgehog.Gen as Gen +import System.IO (hClose) +import System.IO.Temp (withSystemTempFile) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (testCase, (@?=)) +import Test.Tasty.Hedgehog (testProperty) + +tests :: TestTree +tests = + testGroup + "COPY statements" + [ testGroup + "COPY ... TO" + [ testCopyAllColumnsToCsv, + testCopyMultiColumnProjectionToCsv, + testCopySingleColumnProjectionToCsv, + testCopySelectToCsv + ], + testGroup + "COPY ... FROM" + [ testCopyFromCsvAllColumns, + testCopyFromCsvMultiColumnProjection, + testCopyFromCsvSingleColumnProjection + ], + testGroup + "options" + [ testCopyToWithCustomDelimiterAndNoHeader, + testCopyFromWithCustomDelimiter, + testRoundTripParquetWithCompression + ], + testGroup + "property-based round-trip" + [ propRoundTripCsv, + propRoundTripParquet, + propRoundTripJson + ] + ] + +testCopyAllColumnsToCsv :: TestTree +testCopyAllColumnsToCsv = testCase "Copy all columns of a table to a CSV file" $ + withCopyTarget $ \conn path -> do + runBeamDuckDB conn $ + runCopyTo $ + copyTableTo (_dbWidgets testDb) id (copyToCSV path) + + (header, rows) <- readCsv path + header @?= "id,name,price" + rows + @?= List.sort + [ "1,Widget,9.99", + "2,Sprocket,4.5", + "3,Cog,1.25" + ] + +testCopyMultiColumnProjectionToCsv :: TestTree +testCopyMultiColumnProjectionToCsv = + testCase "Copy a multi-column projection of a table to a CSV file" $ + withCopyTarget $ \conn path -> do + runBeamDuckDB conn $ + runCopyTo $ + copyTableTo + (_dbWidgets testDb) + (\w -> (_widgetId w, _widgetName w)) + (copyToCSV path) + + (header, rows) <- readCsv path + header @?= "id,name" + rows + @?= List.sort + [ "1,Widget", + "2,Sprocket", + "3,Cog" + ] + +testCopySingleColumnProjectionToCsv :: TestTree +testCopySingleColumnProjectionToCsv = + testCase "Copy a single column of a table to a CSV file" $ + withCopyTarget $ \conn path -> do + runBeamDuckDB conn $ + runCopyTo $ + copyTableTo (_dbWidgets testDb) _widgetName (copyToCSV path) + + (header, rows) <- readCsv path + header @?= "name" + rows @?= List.sort ["Widget", "Sprocket", "Cog"] + +testCopySelectToCsv :: TestTree +testCopySelectToCsv = + testCase "Copy the result of a SELECT query to a CSV file" $ + withCopyTarget $ \conn path -> do + runBeamDuckDB conn $ + runCopyTo $ + copySelectTo + ( select $ do + w <- all_ (_dbWidgets testDb) + guard_ (_widgetPrice w >. val_ 4.0) + pure (_widgetName w) + ) + (copyToCSV path) + + (_header, rows) <- readCsv path + -- Only Widget (9.99) and Sprocket (4.5) have price > 4.0; Cog (1.25) is excluded. + rows @?= List.sort ["Widget", "Sprocket"] + +testCopyFromCsvAllColumns :: TestTree +testCopyFromCsvAllColumns = + testCase "Copy all columns of a table from a CSV file" $ + withCopyFromSource csvAllColumns $ \conn path -> do + runBeamDuckDB conn $ + runCopyFrom $ + copyTableFrom (_dbWidgets testDb) id (copyFromCSV path) + + rows <- queryAllWidgets conn + rows + @?= [ Widget 1 "Widget" 9.99, + Widget 2 "Sprocket" 4.5, + Widget 3 "Cog" 1.25 + ] + where + csvAllColumns = + "id,name,price\n\ + \1,Widget,9.99\n\ + \2,Sprocket,4.5\n\ + \3,Cog,1.25\n" + +testCopyFromCsvMultiColumnProjection :: TestTree +testCopyFromCsvMultiColumnProjection = + testCase "Copy a multi-column projection of a table from a CSV file" $ + withCopyFromSource csvIdName $ \conn path -> do + runBeamDuckDB conn $ + runCopyFrom $ + copyTableFrom + (_dbWidgets testDb) + (\w -> (_widgetId w, _widgetName w)) + (copyFromCSV path) + + rows <- queryAllWidgets conn + -- 'price' was not in the CSV, so it gets the table's DEFAULT value. + rows + @?= [ Widget 1 "Widget" 0, + Widget 2 "Sprocket" 0, + Widget 3 "Cog" 0 + ] + where + csvIdName = + "id,name\n\ + \1,Widget\n\ + \2,Sprocket\n\ + \3,Cog\n" + +testCopyFromCsvSingleColumnProjection :: TestTree +testCopyFromCsvSingleColumnProjection = + testCase "Copy a single column of a table from a CSV file" $ + withCopyFromSource csvIdOnly $ \conn path -> do + runBeamDuckDB conn $ + runCopyFrom $ + copyTableFrom (_dbWidgets testDb) _widgetId (copyFromCSV path) + + rows <- queryAllWidgets conn + -- Only 'id' was in the CSV, so 'name' and 'price' get table DEFAULTs. + rows + @?= [ Widget 1 "" 0, + Widget 2 "" 0, + Widget 3 "" 0 + ] + where + csvIdOnly = + "id\n\ + \1\n\ + \2\n\ + \3\n" + +testCopyToWithCustomDelimiterAndNoHeader :: TestTree +testCopyToWithCustomDelimiterAndNoHeader = + testCase "Copy a table to a CSV file with custom delimiter and no header" $ + withCopyTarget $ \conn path -> do + let opts = + copyToCSVWith + path + defaultDuckDBCSVCopyToOptions + { csvCopyToHeader = Just False, + csvCopyToDelimiter = Just "|" + } + runBeamDuckDB conn $ + runCopyTo $ + copyTableTo (_dbWidgets testDb) id opts + + contents <- Text.IO.readFile path + let dataLines = List.sort (filter (not . Text.null) (Text.lines contents)) + dataLines + @?= List.sort + [ "1|Widget|9.99", + "2|Sprocket|4.5", + "3|Cog|1.25" + ] + +testCopyFromWithCustomDelimiter :: TestTree +testCopyFromWithCustomDelimiter = + testCase "Copy from a CSV file with a custom delimiter" $ + withCopyFromSource pipeDelimited $ \conn path -> do + let opts = + copyFromCSVWith + path + defaultDuckDBCSVCopyFromOptions + { csvCopyFromDelimiter = Just "|" + } + runBeamDuckDB conn $ + runCopyFrom $ + copyTableFrom (_dbWidgets testDb) id opts + + rows <- queryAllWidgets conn + rows + @?= [ Widget 1 "Widget" 9.99, + Widget 2 "Sprocket" 4.5, + Widget 3 "Cog" 1.25 + ] + where + pipeDelimited = + "id|name|price\n\ + \1|Widget|9.99\n\ + \2|Sprocket|4.5\n\ + \3|Cog|1.25\n" + +testRoundTripParquetWithCompression :: TestTree +testRoundTripParquetWithCompression = + testCase "Round-trip data through a ZSTD-compressed Parquet file" $ + withTestDb widgetData $ \conn -> + withSystemTempFile "beam-duckdb-copy-.parquet" $ \path h -> do + hClose h + let toOpts = + copyToParquetWith + path + defaultDuckDBParquetCopyToOptions + { parquetCopyToCompression = Just ParquetZstd + } + fromOpts = copyFromParquet path + + runBeamDuckDB conn $ + runCopyTo $ + copyTableTo (_dbWidgets testDb) id toOpts + + -- Wipe the seeded rows, then re-import from the Parquet file we wrote. + void $ execute_ conn "DELETE FROM widgets" + + runBeamDuckDB conn $ + runCopyFrom $ + copyTableFrom (_dbWidgets testDb) id fromOpts + + rows <- queryAllWidgets conn + rows + @?= [ Widget 1 "Widget" 9.99, + Widget 2 "Sprocket" 4.5, + Widget 3 "Cog" 1.25 + ] + +propRoundTripCsv :: TestTree +propRoundTripCsv = + testProperty "CSV options round-trip" $ + property $ do + toOpts <- forAll genCsvToOpts + let fromOpts = matchedCsvFromOpts toOpts + rows <- liftIO $ csvRoundTrip toOpts fromOpts + rows === widgetData + where + csvRoundTrip toOpts fromOpts = + withTestDb widgetData $ \conn -> + withSystemTempFile "beam-duckdb-prop-csv-.csv" $ \path h -> do + hClose h + runBeamDuckDB conn $ + runCopyTo $ + copyTableTo (_dbWidgets testDb) id (copyToCSVWith path toOpts) + void $ execute_ conn "DELETE FROM widgets" + runBeamDuckDB conn $ + runCopyFrom $ + copyTableFrom (_dbWidgets testDb) id (copyFromCSVWith path fromOpts) + queryAllWidgets conn + + genCsvToOpts :: Hedgehog.Gen DuckDBCSVCopyToOptions + genCsvToOpts = do + hdr <- Gen.maybe Gen.bool + delim <- Gen.maybe (Gen.element [",", ";", "|", "\t"]) + nullStr <- Gen.maybe (Gen.element ["NULL", "~~~"]) + quote <- Gen.maybe genCsvSpecialChar + escape <- Gen.maybe genCsvSpecialChar + + -- quote and escape can't overlap with nullStr and delim, + -- otherwise the CSV data may be unparseable. + guard (maybe True not $ Text.elem <$> quote <*> delim) + guard (quote /= Just ',') -- the default delimiter + guard (maybe True not $ Text.elem <$> quote <*> nullStr) + guard (maybe True not $ Text.elem <$> escape <*> delim) + guard (escape /= Just ',') -- the default delimiter + guard (maybe True not $ Text.elem <$> escape <*> nullStr) + + pure + defaultDuckDBCSVCopyToOptions + { csvCopyToHeader = hdr, + csvCopyToDelimiter = delim, + csvCopyToQuote = quote, + csvCopyToEscape = escape, + csvCopyToNullStr = nullStr, + -- The TO-side @overwrite_or_ignore@ field is intentionally fixed to + -- @Just True@: a property that replays through the same temp file 25 + -- times needs the second iteration onwards to overwrite, and the field + -- itself is independent of the readable contents we round-trip. + csvCopyToOverwriteOrIgnore = Just True + } + + matchedCsvFromOpts :: DuckDBCSVCopyToOptions -> DuckDBCSVCopyFromOptions + matchedCsvFromOpts o = + defaultDuckDBCSVCopyFromOptions + { csvCopyFromHeader = csvCopyToHeader o, + csvCopyFromDelimiter = csvCopyToDelimiter o, + csvCopyFromQuote = csvCopyToQuote o, + csvCopyFromEscape = csvCopyToEscape o, + csvCopyFromNullStr = csvCopyToNullStr o + } + + -- Printable single-byte ASCII characters that don't collide with + -- the seeded 'widgetData' (which is alphanumeric plus '.'), and the delimiter + genCsvSpecialChar :: Hedgehog.Gen Char + genCsvSpecialChar = + Gen.element + [ c + | c <- enumFromTo minBound maxBound, + Char.isPrint c, + Char.isAscii c, + not (Char.isAlphaNum c), + c /= '.' + ] + +propRoundTripParquet :: TestTree +propRoundTripParquet = + testProperty "Parquet options round-trip" $ + property $ do + toOpts <- forAll genParquetToOpts + let fromOpts = defaultDuckDBParquetCopyFromOptions + rows <- liftIO $ parquetRoundTrip toOpts fromOpts + rows === widgetData + where + parquetRoundTrip toOpts fromOpts = + withTestDb widgetData $ \conn -> + withSystemTempFile "beam-duckdb-prop-parquet-.parquet" $ \path h -> do + hClose h + runBeamDuckDB conn $ + runCopyTo $ + copyTableTo (_dbWidgets testDb) id (copyToParquetWith path toOpts) + void $ execute_ conn "DELETE FROM widgets" + runBeamDuckDB conn $ + runCopyFrom $ + copyTableFrom (_dbWidgets testDb) id (copyFromParquetWith path fromOpts) + queryAllWidgets conn + + genParquetToOpts :: Hedgehog.Gen DuckDBParquetCopyToOptions + genParquetToOpts = do + comp <- Gen.maybe Gen.enumBounded + pure + defaultDuckDBParquetCopyToOptions + { parquetCopyToCompression = comp, + parquetCopyToOverwriteOrIgnore = Just True + } + +propRoundTripJson :: TestTree +propRoundTripJson = + testProperty "JSON options round-trip" $ + property $ do + toOpts <- forAll genJsonToOpts + fromOpts <- forAll (genJsonFromOpts (jsonCopyToCompression toOpts)) + rows <- liftIO $ jsonRoundTrip toOpts fromOpts + rows === widgetData + where + jsonRoundTrip toOpts fromOpts = + withTestDb widgetData $ \conn -> + withSystemTempFile "beam-duckdb-prop-json-.json" $ \path h -> do + hClose h + runBeamDuckDB conn $ + runCopyTo $ + copyTableTo (_dbWidgets testDb) id (copyToJSONWith path toOpts) + void $ execute_ conn "DELETE FROM widgets" + runBeamDuckDB conn $ + runCopyFrom $ + copyTableFrom (_dbWidgets testDb) id (copyFromJSONWith path fromOpts) + queryAllWidgets conn + + genJsonToOpts :: Hedgehog.Gen DuckDBJSONCopyToOptions + genJsonToOpts = do + comp <- Gen.maybe Gen.enumBounded + pure + defaultDuckDBJSONCopyToOptions + { jsonCopyToCompression = comp, + jsonCopyToOverwriteOrIgnore = Just True + } + + genJsonFromOpts :: Maybe JSONCompression -> Hedgehog.Gen DuckDBJSONCopyFromOptions + genJsonFromOpts comp = do + auto <- Gen.maybe Gen.bool + pure + defaultDuckDBJSONCopyFromOptions + { jsonCopyFromAutoDetect = auto, + jsonCopyFromCompression = comp + } + +withCopyTarget :: (Connection -> FilePath -> IO a) -> IO a +withCopyTarget action = + withTestDb widgetData $ \conn -> + withSystemTempFile "beam-duckdb-copy-.csv" $ \path h -> do + -- DuckDB's COPY TO opens the file itself; close our handle first + -- so it isn't held open while DuckDB writes. + hClose h + action conn path + +withCopyFromSource :: Text -> (Connection -> FilePath -> IO a) -> IO a +withCopyFromSource csvContent action = + withTestDb [] $ \conn -> + withSystemTempFile "beam-duckdb-copy-from-.csv" $ \path h -> do + Text.IO.hPutStr h csvContent + hClose h + action conn path + +readCsv :: FilePath -> IO (Text, [Text]) +readCsv path = do + contents <- Text.IO.readFile path + case Text.lines contents of + [] -> fail ("expected at least a header line in " <> path) + (h : ds) -> pure (h, List.sort ds) + +queryAllWidgets :: Connection -> IO [Widget] +queryAllWidgets conn = + runBeamDuckDB conn $ + runSelectReturningList $ + select $ + orderBy_ (asc_ . _widgetId) $ + all_ (_dbWidgets testDb) + +widgetData :: [Widget] +widgetData = + [ Widget 1 "Widget" 9.99, + Widget 2 "Sprocket" 4.50, + Widget 3 "Cog" 1.25 + ] + +data WidgetT f = Widget + { _widgetId :: Columnar f Int32, + _widgetName :: Columnar f Text, + _widgetPrice :: Columnar f Double + } + deriving (Generic) + +type Widget = WidgetT Identity + +deriving instance Show Widget + +deriving instance Eq Widget + +instance Beamable WidgetT + +instance Table WidgetT where + data PrimaryKey WidgetT f = WidgetId (Columnar f Int32) + deriving (Generic) + primaryKey = WidgetId . _widgetId + +instance Beamable (PrimaryKey WidgetT) + +newtype TestDB f = TestDB + { _dbWidgets :: f (TableEntity WidgetT) + } + deriving (Generic, Database be) + +testDb :: DatabaseSettings DuckDB TestDB +testDb = + defaultDbSettings + `withDbModification` dbModification + { _dbWidgets = + modifyTableFields + tableModification + { _widgetId = "id", + _widgetName = "name", + _widgetPrice = "price" + } + } + +createTables :: Connection -> IO () +createTables conn = + void $ + execute_ + conn + "CREATE TABLE widgets (\ + \ id INTEGER PRIMARY KEY,\ + \ name TEXT NOT NULL DEFAULT '',\ + \ price DOUBLE NOT NULL DEFAULT 0\ + \)" + +seedData :: Connection -> [Widget] -> IO () +seedData conn widgets = + runBeamDuckDB conn $ + runInsert $ + insert (_dbWidgets testDb) $ + insertValues widgets + +withTestDb :: + [Widget] -> + (Connection -> IO a) -> + IO a +withTestDb widgets action = + withConnection ":memory:" $ \conn -> do + createTables conn + seedData conn widgets + action conn diff --git a/tests/Database/Beam/DuckDB/Test/Extensions/DataSource.hs b/tests/Database/Beam/DuckDB/Test/Extensions/DataSource.hs new file mode 100644 --- /dev/null +++ b/tests/Database/Beam/DuckDB/Test/Extensions/DataSource.hs @@ -0,0 +1,365 @@ +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE ExplicitForAll #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} + +module Database.Beam.DuckDB.Test.Extensions.DataSource (tests) where + +import Data.Int (Int32) +import Data.List.NonEmpty (NonEmpty ((:|))) +import Data.Text (Text) +import Data.Time (Day, fromGregorian) +import Database.Beam + ( Beamable, + Columnar, + Database, + DatabaseSettings, + Generic, + Identity, + Table (..), + aggregate_, + as_, + countAll_, + dbModification, + defaultDbSettings, + guard_, + min_, + runSelectReturningList, + runSelectReturningOne, + select, + tableModification, + val_, + withDbModification, + (==.), + ) +import Database.Beam.DuckDB + ( CSVOptions (..), + DataSourceEntity, + DuckDB, + allFromDataSource_, + allowMovedPaths, + csvWith, + dataSource, + defaultCSVOptions, + defaultIcebergTableOptions, + icebergTableWith, + modifyDataSourceFields, + parquet, + runBeamDuckDB, + ) +import Database.DuckDB.Simple (Connection, withConnection) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (testCase, (@?=)) + +tests :: TestTree +tests = + testGroup + "DataSource" + [ testGroup + "Parquet" + [ testCountFromParquet, + testQueryFromParquet, + testCountFromMultipleParquetFiles, + testQueryFromMultipleParquetFiles + ], + testGroup + "Apache Iceberg" + [ testCountFromIceberg, + testParseDateFromIceberg + ], + testGroup + "CSV" + [ testCountFromCSV, + testQueryFromCSV + ] + ] + +testCountFromParquet :: TestTree +testCountFromParquet = testCase "Counting records from a Parquet file" $ do + results <- withTestDb $ \conn -> + runBeamDuckDB conn $ + runSelectReturningList $ + select $ do + exam <- allFromDataSource_ (_dbExams testDb) + pure (_examName exam) + results @?= ["alice", "bob", "carol", "dave"] + +testQueryFromParquet :: TestTree +testQueryFromParquet = testCase "Query from a Parquet file" $ do + results <- withTestDb $ \conn -> + runBeamDuckDB conn $ + runSelectReturningOne $ + select $ do + exam <- allFromDataSource_ (_dbExams testDb) + guard_ (_examId exam ==. 1) + pure (_examName exam) + results @?= Just "alice" + +testCountFromMultipleParquetFiles :: TestTree +testCountFromMultipleParquetFiles = testCase "Counting records from multiple Parquet file" $ do + results <- withTestDb $ \conn -> + runBeamDuckDB conn $ + runSelectReturningList $ + select $ do + exam <- allFromDataSource_ (_dbExamsMulti testDb) + pure (_examName exam) + results @?= ["alice", "bob", "carol", "dave", "erika", "francis", "genevieve", "hugo"] + +testQueryFromMultipleParquetFiles :: TestTree +testQueryFromMultipleParquetFiles = testCase "Query from multiple Parquet files" $ do + results <- withTestDb $ \conn -> + runBeamDuckDB conn $ + runSelectReturningOne $ + select $ do + exam <- allFromDataSource_ (_dbExamsMulti testDb) + guard_ (_examId exam ==. 5) + pure (_examName exam) + results @?= Just "erika" + +testCountFromIceberg :: TestTree +testCountFromIceberg = testCase "Counting records from an Apache Iceberg table" $ do + results <- withTestDb $ \conn -> + runBeamDuckDB conn $ + runSelectReturningOne $ + select $ + aggregate_ (\_ -> as_ @Int32 countAll_) (allFromDataSource_ (_dbLineItems testDb)) + results @?= Just 51793 -- From DuckDB's documentation + +testParseDateFromIceberg :: TestTree +testParseDateFromIceberg = testCase "Test parsing date columns from Apache Iceberg table" $ do + results <- withTestDb $ \conn -> + runBeamDuckDB conn $ + runSelectReturningOne $ + select $ + aggregate_ (min_ . _lineitemShipdate) (allFromDataSource_ (_dbLineItems testDb)) + + results @?= Just (Just (fromGregorian 1992 01 04)) + +testCountFromCSV :: TestTree +testCountFromCSV = testCase "Counting records from a CSV file" $ do + results <- withTestDb $ \conn -> + runBeamDuckDB conn $ + runSelectReturningOne $ + select $ + aggregate_ (\_ -> as_ @Int32 countAll_) (allFromDataSource_ (_dbFlights testDb)) + results @?= Just 3 + +testQueryFromCSV :: TestTree +testQueryFromCSV = testCase "Query from a CSV file" $ do + results <- withTestDb $ \conn -> + runBeamDuckDB conn $ + runSelectReturningOne $ + select $ do + flight <- allFromDataSource_ (_dbFlights testDb) + guard_ (_flightDate flight ==. val_ (fromGregorian 1988 01 01)) + pure flight + results + @?= Just + ( Flight + { _flightDate = fromGregorian 1988 01 01, + _flightUniqueCarrier = "AA", + _flightOriginCity = "New York, NY", + _flightDestCity = "Los Angeles, CA" + } + ) + +data ExamT f = Exam + { _examId :: Columnar f Int32, + _examName :: Columnar f Text, + _examScore :: Columnar f Double, + _examDate :: Columnar f Day + } + deriving (Generic) + +type Exam = ExamT Identity + +type ExamId = PrimaryKey ExamT Identity + +deriving instance Show ExamId + +deriving instance Eq ExamId + +deriving instance Ord ExamId + +deriving instance Show Exam + +deriving instance Eq Exam + +deriving instance Ord Exam + +instance Beamable ExamT + +instance Table ExamT where + data PrimaryKey ExamT f = ExamId (Columnar f Int32) + deriving (Generic) + primaryKey = ExamId . _examId + +instance Beamable (PrimaryKey ExamT) + +data LineItemT f = Lineitem + { _lineitemOrderkey :: Columnar f Int32, + _lineitemPartkey :: Columnar f Int32, + _lineitemSuppkey :: Columnar f Int32, + _lineitemLinenumber :: Columnar f Int32, + _lineitemQuantity :: Columnar f Double, + _lineitemExtendedprice :: Columnar f Double, + _lineitemDiscount :: Columnar f Double, + _lineitemTax :: Columnar f Double, + _lineitemReturnflag :: Columnar f Text, + _lineitemLinestatus :: Columnar f Text, + _lineitemShipdate :: Columnar f Day, + _lineitemCommitdate :: Columnar f Day, + _lineitemReceiptdate :: Columnar f Day, + _lineitemShipinstruct :: Columnar f Text, + _lineitemShipmode :: Columnar f Text, + _lineitemComment :: Columnar f Text + } + deriving (Generic) + +type Lineitem = LineItemT Identity + +deriving instance Show Lineitem + +deriving instance Eq Lineitem + +instance Beamable LineItemT + +instance Table LineItemT where + data PrimaryKey LineItemT f + = LineItemKey (Columnar f Int32) (Columnar f Int32) + deriving (Generic) + primaryKey li = + LineItemKey (_lineitemOrderkey li) (_lineitemLinenumber li) + +instance Beamable (PrimaryKey LineItemT) + +deriving instance Show (PrimaryKey LineItemT Identity) + +deriving instance Eq (PrimaryKey LineItemT Identity) + +data FlightT f = Flight + { _flightDate :: Columnar f Day, + _flightUniqueCarrier :: Columnar f Text, + _flightOriginCity :: Columnar f Text, + _flightDestCity :: Columnar f Text + } + deriving (Generic) + +type Flight = FlightT Identity + +deriving instance Show Flight + +deriving instance Eq Flight + +instance Beamable FlightT + +instance Table FlightT where + data PrimaryKey FlightT f + = FlightKey + (Columnar f Day) + (Columnar f Text) + (Columnar f Text) + (Columnar f Text) + deriving (Generic) + primaryKey fl = + FlightKey + (_flightDate fl) + (_flightUniqueCarrier fl) + (_flightOriginCity fl) + (_flightDestCity fl) + +instance Beamable (PrimaryKey FlightT) + +deriving instance Show (PrimaryKey FlightT Identity) + +deriving instance Eq (PrimaryKey FlightT Identity) + +data TestDB f = TestDB + { _dbExams :: f (DataSourceEntity ExamT), + _dbExamsMulti :: f (DataSourceEntity ExamT), -- set up with multiple parquet files + _dbLineItems :: f (DataSourceEntity LineItemT), + _dbFlights :: f (DataSourceEntity FlightT) + } + deriving (Generic, Database DuckDB) + +testDb :: DatabaseSettings DuckDB TestDB +testDb = + defaultDbSettings + `withDbModification` (dbModification @_ @DuckDB) + { _dbExams = + dataSource + (parquet ("tests/data/test1.parquet" :| [])) + <> modifyDataSourceFields + tableModification + { _examId = "id", + _examName = "name", + _examScore = "score", + _examDate = "exam_date" + }, + _dbExamsMulti = + dataSource + ( parquet + ("tests/data/test1.parquet" :| ["tests/data/test2.parquet"]) + ) + <> modifyDataSourceFields + tableModification + { _examId = "id", + _examName = "name", + _examScore = "score", + _examDate = "exam_date" + }, + _dbLineItems = + dataSource + ( icebergTableWith + "tests/data/lineitem_iceberg" + (defaultIcebergTableOptions {allowMovedPaths = Just True}) + ) + <> modifyDataSourceFields + tableModification + { _lineitemOrderkey = "l_orderkey", + _lineitemPartkey = "l_partkey", + _lineitemSuppkey = "l_suppkey", + _lineitemLinenumber = "l_linenumber", + _lineitemQuantity = "l_quantity", + _lineitemExtendedprice = "l_extendedprice", + _lineitemDiscount = "l_discount", + _lineitemTax = "l_tax", + _lineitemReturnflag = "l_returnflag", + _lineitemLinestatus = "l_linestatus", + _lineitemShipdate = "l_shipdate", + _lineitemCommitdate = "l_commitdate", + _lineitemReceiptdate = "l_receiptdate", + _lineitemShipinstruct = "l_shipinstruct", + _lineitemShipmode = "l_shipmode", + _lineitemComment = "l_comment" + }, + _dbFlights = + dataSource + ( csvWith + ("tests/data/flights.csv" :| []) + ( defaultCSVOptions + { header = Just True, + comment = Just '#', + delim = Just "|", + ignoreErrors = Just False + } + ) + ) + <> modifyDataSourceFields + tableModification + { _flightDate = "FlightDate", + _flightUniqueCarrier = "UniqueCarrier", + _flightOriginCity = "OriginCityName", + _flightDestCity = "DestCityName" + } + } + +withTestDb :: + (Connection -> IO a) -> + IO a +withTestDb = + withConnection ":memory:" diff --git a/tests/Database/Beam/DuckDB/Test/Extensions/InsertOnConflict.hs b/tests/Database/Beam/DuckDB/Test/Extensions/InsertOnConflict.hs new file mode 100644 --- /dev/null +++ b/tests/Database/Beam/DuckDB/Test/Extensions/InsertOnConflict.hs @@ -0,0 +1,306 @@ +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TypeFamilies #-} + +module Database.Beam.DuckDB.Test.Extensions.InsertOnConflict (tests) where + +import Control.Monad (void) +import Data.Int (Int32) +import qualified Data.List as List +import Data.Text (Text) +import Database.Beam + ( Beamable, + Columnar, + Database, + DatabaseSettings, + Generic, + Identity, + SqlOrd ((>.)), + Table (..), + TableEntity, + all_, + asc_, + current_, + dbModification, + defaultDbSettings, + insert, + insertValues, + modifyTableFields, + orderBy_, + runInsert, + runSelectReturningList, + select, + tableModification, + val_, + withDbModification, + (<-.), + ) +import Database.Beam.DuckDB + ( BeamHasInsertOnConflict + ( anyConflict, + conflictingFields, + insertOnConflict, + onConflictDoNothing, + onConflictUpdateSet, + onConflictUpdateSetWhere + ), + DuckDB, + onConflictUpdateAll, + onConflictUpdateInstead, + runBeamDuckDB, + runInsertReturningList, + ) +import Database.DuckDB.Simple (Connection, execute_, withConnection) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (testCase, (@?=)) + +tests :: TestTree +tests = + testGroup + "INSERT ... ON CONFLICT" + [ testAnyConflictDoNothing, + testConflictingFieldsDoNothing, + testConflictingFieldsUpdateAll, + testConflictingFieldsUpdateInstead, + testConflictingFieldsUpdateSet, + testConflictingFieldsUpdateSetWhere, + testInsertOnConflictReturning + ] + +testAnyConflictDoNothing :: TestTree +testAnyConflictDoNothing = + testCase "anyConflict + onConflictDoNothing leaves existing rows untouched" $ + withTestDb widgetData $ \conn -> do + runBeamDuckDB conn $ + runInsert $ + insertOnConflict + (_dbWidgets testDb) + (insertValues [Widget 2 "Sprocket-collision" 0.0]) + anyConflict + onConflictDoNothing + remaining <- queryAllWidgets conn + remaining @?= List.sortOn _widgetId widgetData + +testConflictingFieldsDoNothing :: TestTree +testConflictingFieldsDoNothing = + testCase "conflictingFields + onConflictDoNothing inserts non-conflicting rows" $ + withTestDb widgetData $ \conn -> do + runBeamDuckDB conn $ + runInsert $ + insertOnConflict + (_dbWidgets testDb) + ( insertValues + [ Widget 2 "Sprocket-collision" 0.0, + Widget 4 "Gear" 7.75 + ] + ) + (conflictingFields _widgetId) + onConflictDoNothing + remaining <- queryAllWidgets conn + remaining + @?= [ Widget 1 "Widget" 9.99, + Widget 2 "Sprocket" 4.50, + Widget 3 "Cog" 1.25, + Widget 4 "Gear" 7.75 + ] + +testConflictingFieldsUpdateAll :: TestTree +testConflictingFieldsUpdateAll = + testCase "conflictingFields + onConflictUpdateAll overwrites every column" $ + withTestDb widgetData $ \conn -> do + runBeamDuckDB conn $ + runInsert $ + insertOnConflict + (_dbWidgets testDb) + ( insertValues + [ Widget 2 "Sprocket-v2" 6.00, + Widget 4 "Gear" 7.75 + ] + ) + (conflictingFields _widgetId) + onConflictUpdateAll + remaining <- queryAllWidgets conn + remaining + @?= [ Widget 1 "Widget" 9.99, + Widget 2 "Sprocket-v2" 6.00, + Widget 3 "Cog" 1.25, + Widget 4 "Gear" 7.75 + ] + +testConflictingFieldsUpdateInstead :: TestTree +testConflictingFieldsUpdateInstead = + testCase "conflictingFields + onConflictUpdateInstead overwrites a subset" $ + withTestDb widgetData $ \conn -> do + runBeamDuckDB conn $ + runInsert $ + insertOnConflict + (_dbWidgets testDb) + (insertValues [Widget 2 "Sprocket-v3" 100.0]) + (conflictingFields _widgetId) + (onConflictUpdateInstead _widgetName) + remaining <- queryAllWidgets conn + remaining + @?= [ Widget 1 "Widget" 9.99, + Widget 2 "Sprocket-v3" 4.50, + Widget 3 "Cog" 1.25 + ] + +testConflictingFieldsUpdateSet :: TestTree +testConflictingFieldsUpdateSet = + testCase "conflictingFields + onConflictUpdateSet applies a custom assignment" $ + withTestDb widgetData $ \conn -> do + runBeamDuckDB conn $ + runInsert $ + insertOnConflict + (_dbWidgets testDb) + (insertValues [Widget 2 "Sprocket-discount" 0.0]) + (conflictingFields _widgetId) + ( onConflictUpdateSet + ( \fields _excluded -> + _widgetPrice fields <-. current_ (_widgetPrice fields) - val_ 0.50 + ) + ) + remaining <- queryAllWidgets conn + remaining + @?= [ Widget 1 "Widget" 9.99, + Widget 2 "Sprocket" 4.00, + Widget 3 "Cog" 1.25 + ] + +testConflictingFieldsUpdateSetWhere :: TestTree +testConflictingFieldsUpdateSetWhere = + testCase "conflictingFields + onConflictUpdateSetWhere only updates when predicate holds" $ + withTestDb widgetData $ \conn -> do + runBeamDuckDB conn $ + runInsert $ + insertOnConflict + (_dbWidgets testDb) + ( insertValues + [ Widget 1 "Widget-replacement" 1.00, + Widget 3 "Cog-replacement" 50.00 + ] + ) + (conflictingFields _widgetId) + ( onConflictUpdateSetWhere + ( \fields excludedRow -> + _widgetName fields <-. _widgetName excludedRow + ) + -- Only overwrite when the incoming price is higher than the existing one + ( \fields excludedRow -> + _widgetPrice excludedRow >. current_ (_widgetPrice fields) + ) + ) + remaining <- queryAllWidgets conn + remaining + @?= [ Widget 1 "Widget" 9.99, + Widget 2 "Sprocket" 4.50, + Widget 3 "Cog-replacement" 1.25 + ] + +testInsertOnConflictReturning :: TestTree +testInsertOnConflictReturning = + testCase "INSERT ON CONFLICT ... RETURNING returns inserted-or-updated rows" $ + withTestDb widgetData $ \conn -> do + returned <- + runBeamDuckDB conn $ + runInsertReturningList $ + insertOnConflict + (_dbWidgets testDb) + ( insertValues + [ Widget 2 "Sprocket-v2" 6.00, + Widget 4 "Gear" 7.75 + ] + ) + (conflictingFields _widgetId) + onConflictUpdateAll + List.sort returned + @?= List.sort + [ Widget 2 "Sprocket-v2" 6.00, + Widget 4 "Gear" 7.75 + ] + +data WidgetT f = Widget + { _widgetId :: Columnar f Int32, + _widgetName :: Columnar f Text, + _widgetPrice :: Columnar f Double + } + deriving (Generic) + +type Widget = WidgetT Identity + +deriving instance Show Widget + +deriving instance Eq Widget + +deriving instance Ord Widget + +instance Beamable WidgetT + +instance Table WidgetT where + data PrimaryKey WidgetT f = WidgetId (Columnar f Int32) + deriving (Generic) + primaryKey = WidgetId . _widgetId + +instance Beamable (PrimaryKey WidgetT) + +newtype TestDB f = TestDB + { _dbWidgets :: f (TableEntity WidgetT) + } + deriving (Generic, Database be) + +testDb :: DatabaseSettings DuckDB TestDB +testDb = + defaultDbSettings + `withDbModification` dbModification + { _dbWidgets = + modifyTableFields + tableModification + { _widgetId = "id", + _widgetName = "name", + _widgetPrice = "price" + } + } + +createTables :: Connection -> IO () +createTables conn = + void $ + execute_ + conn + "CREATE TABLE widgets (\ + \ id INTEGER PRIMARY KEY,\ + \ name TEXT NOT NULL DEFAULT '',\ + \ price DOUBLE NOT NULL DEFAULT 0\ + \)" + +seedData :: Connection -> [Widget] -> IO () +seedData _ [] = pure () +seedData conn widgets = + void $ + runBeamDuckDB conn $ + runInsertReturningList $ + insert (_dbWidgets testDb) (insertValues widgets) + +withTestDb :: [Widget] -> (Connection -> IO a) -> IO a +withTestDb widgets action = + withConnection ":memory:" $ \conn -> do + createTables conn + seedData conn widgets + action conn + +queryAllWidgets :: Connection -> IO [Widget] +queryAllWidgets conn = + runBeamDuckDB conn $ + runSelectReturningList $ + select $ + orderBy_ (asc_ . _widgetId) $ + all_ (_dbWidgets testDb) + +widgetData :: [Widget] +widgetData = + [ Widget 1 "Widget" 9.99, + Widget 2 "Sprocket" 4.50, + Widget 3 "Cog" 1.25 + ] diff --git a/tests/Database/Beam/DuckDB/Test/Extensions/Returning.hs b/tests/Database/Beam/DuckDB/Test/Extensions/Returning.hs new file mode 100644 --- /dev/null +++ b/tests/Database/Beam/DuckDB/Test/Extensions/Returning.hs @@ -0,0 +1,239 @@ +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TypeFamilies #-} + +module Database.Beam.DuckDB.Test.Extensions.Returning (tests) where + +import Control.Monad (void) +import Data.Int (Int32) +import qualified Data.List as List +import Data.Text (Text) +import Database.Beam + ( Beamable, + Columnar, + Database, + DatabaseSettings, + Generic, + Identity, + Table (..), + TableEntity, + all_, + asc_, + current_, + dbModification, + defaultDbSettings, + delete, + insert, + insertValues, + modifyTableFields, + orderBy_, + runSelectReturningList, + save, + select, + tableModification, + update, + val_, + withDbModification, + (<-.), + (==.), + ) +import Database.Beam.DuckDB + ( DuckDB, + runBeamDuckDB, + runDeleteReturningList, + runInsertReturningList, + runUpdateReturningList, + ) +import Database.DuckDB.Simple (Connection, execute_, withConnection) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (testCase, (@?=)) + +tests :: TestTree +tests = + testGroup + "RETURNING" + [ testGroup + "INSERT ... RETURNING" + [ testInsertReturningAll, + testInsertReturningEmpty + ], + testGroup + "UPDATE ... RETURNING" + [ testUpdateReturningAll, + testUpdateReturningOnlyMatched + ], + testGroup + "DELETE ... RETURNING" + [ testDeleteReturningAll, + testDeleteReturningOnlyMatched + ] + ] + +testInsertReturningAll :: TestTree +testInsertReturningAll = + testCase "INSERT returns the inserted rows" $ + withTestDb [] $ \conn -> do + returned <- + runBeamDuckDB conn $ + runInsertReturningList $ + insert (_dbWidgets testDb) (insertValues widgetData) + List.sort returned @?= List.sort widgetData + +testInsertReturningEmpty :: TestTree +testInsertReturningEmpty = + testCase "INSERT with no rows returns []" $ + withTestDb widgetData $ \conn -> do + returned <- + runBeamDuckDB conn $ + runInsertReturningList $ + insert (_dbWidgets testDb) (insertValues []) + returned @?= [] + remaining <- queryAllWidgets conn + remaining @?= List.sortOn _widgetId widgetData + +testUpdateReturningAll :: TestTree +testUpdateReturningAll = + testCase "UPDATE returns the post-update rows" $ + withTestDb widgetData $ \conn -> do + returned <- + runBeamDuckDB conn $ + runUpdateReturningList $ + update + (_dbWidgets testDb) + (\w -> _widgetPrice w <-. current_ (_widgetPrice w) + val_ 1.0) + (\_ -> val_ True) + List.sort returned + @?= List.sort + [ Widget 1 "Widget" 10.99, + Widget 2 "Sprocket" 5.50, + Widget 3 "Cog" 2.25 + ] + +testUpdateReturningOnlyMatched :: TestTree +testUpdateReturningOnlyMatched = + testCase "UPDATE with WHERE returns only the matched rows" $ + withTestDb widgetData $ \conn -> do + returned <- + runBeamDuckDB conn $ + runUpdateReturningList $ + save (_dbWidgets testDb) (Widget 2 "Sprocket-v2" 4.50) + returned @?= [Widget 2 "Sprocket-v2" 4.50] + remaining <- queryAllWidgets conn + remaining + @?= [ Widget 1 "Widget" 9.99, + Widget 2 "Sprocket-v2" 4.50, + Widget 3 "Cog" 1.25 + ] + +testDeleteReturningAll :: TestTree +testDeleteReturningAll = + testCase "DELETE returns the deleted rows" $ + withTestDb widgetData $ \conn -> do + returned <- + runBeamDuckDB conn $ + runDeleteReturningList $ + delete (_dbWidgets testDb) (\_ -> val_ True) + List.sort returned @?= List.sort widgetData + remaining <- queryAllWidgets conn + remaining @?= [] + +testDeleteReturningOnlyMatched :: TestTree +testDeleteReturningOnlyMatched = + testCase "DELETE with WHERE returns only the matched rows" $ + withTestDb widgetData $ \conn -> do + returned <- + runBeamDuckDB conn $ + runDeleteReturningList $ + delete (_dbWidgets testDb) (\w -> _widgetId w ==. val_ 2) + returned @?= [Widget 2 "Sprocket" 4.50] + remaining <- queryAllWidgets conn + remaining + @?= [ Widget 1 "Widget" 9.99, + Widget 3 "Cog" 1.25 + ] + +data WidgetT f = Widget + { _widgetId :: Columnar f Int32, + _widgetName :: Columnar f Text, + _widgetPrice :: Columnar f Double + } + deriving (Generic) + +type Widget = WidgetT Identity + +deriving instance Show Widget + +deriving instance Eq Widget + +deriving instance Ord Widget + +instance Beamable WidgetT + +instance Table WidgetT where + data PrimaryKey WidgetT f = WidgetId (Columnar f Int32) + deriving (Generic) + primaryKey = WidgetId . _widgetId + +instance Beamable (PrimaryKey WidgetT) + +newtype TestDB f = TestDB + { _dbWidgets :: f (TableEntity WidgetT) + } + deriving (Generic, Database be) + +testDb :: DatabaseSettings DuckDB TestDB +testDb = + defaultDbSettings + `withDbModification` dbModification + { _dbWidgets = + modifyTableFields + tableModification + { _widgetId = "id", + _widgetName = "name", + _widgetPrice = "price" + } + } + +createTables :: Connection -> IO () +createTables conn = + void $ + execute_ + conn + "CREATE TABLE widgets (\ + \ id INTEGER PRIMARY KEY,\ + \ name TEXT NOT NULL DEFAULT '',\ + \ price DOUBLE NOT NULL DEFAULT 0\ + \)" + +seedData :: Connection -> [Widget] -> IO () +seedData _ [] = pure () +seedData conn widgets = + void $ + runBeamDuckDB conn $ + runInsertReturningList $ + insert (_dbWidgets testDb) (insertValues widgets) + +withTestDb :: [Widget] -> (Connection -> IO a) -> IO a +withTestDb widgets action = + withConnection ":memory:" $ \conn -> do + createTables conn + seedData conn widgets + action conn + +queryAllWidgets :: Connection -> IO [Widget] +queryAllWidgets conn = + runBeamDuckDB conn $ + runSelectReturningList $ + select $ + orderBy_ (asc_ . _widgetId) $ + all_ (_dbWidgets testDb) + +widgetData :: [Widget] +widgetData = + [ Widget 1 "Widget" 9.99, + Widget 2 "Sprocket" 4.50, + Widget 3 "Cog" 1.25 + ] diff --git a/tests/Database/Beam/DuckDB/Test/Migrate.hs b/tests/Database/Beam/DuckDB/Test/Migrate.hs new file mode 100644 --- /dev/null +++ b/tests/Database/Beam/DuckDB/Test/Migrate.hs @@ -0,0 +1,483 @@ +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} + +module Database.Beam.DuckDB.Test.Migrate (tests) where + +import Control.Arrow ((>>>)) +import Control.Exception (bracket) +import Control.Monad (void) +import Data.ByteString (ByteString) +import Data.Functor ((<&>)) +import Data.Int (Int32, Int8) +import qualified Data.Text as T +import Data.UUID.Types (UUID) +import Data.Word (Word16, Word32, Word64, Word8) +import Database.Beam + ( Beamable, + C, + Database, + Generic, + MonadBeam (runNoReturn), + Table (..), + TableEntity, + ) +import Database.Beam.DuckDB (DuckDB, runBeamDuckDB) +import Database.Beam.DuckDB.Migrate + ( blob, + getDbConstraints, + migrationBackend, + text, + tinyint, + ubigint, + uinteger, + usmallint, + utinyint, + uuid, + ) +import Database.Beam.Migrate + ( CheckedDatabaseSettings, + DatabasePredicate (..), + Migration, + SomeDatabasePredicate (..), + addColumn, + alterTable, + createTable, + defaultMigratableDbSettings, + dropColumn, + executeMigration, + field, + migrationStep, + notNull, + runMigrationSteps, + ) +import Database.Beam.Migrate.Simple + ( VerificationResult (..), + autoMigrate, + verifySchema, + ) +import Database.Beam.Query.DataTypes (double, int, maybeType, varchar) +import Database.DuckDB.Simple (Connection, execute_) +import qualified Database.DuckDB.Simple as DuckDB +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (Assertion, assertBool, assertEqual, assertFailure, testCase, testCaseSteps) + +tests :: TestTree +tests = + testGroup + "Migration tests" + [ verifiesPrimaryKey, + verifiesNotNull, + verifiesForeignKey, + autoMigrateCreatesSchema, + autoMigrateIsIdempotent, + multiStepMigration, + testGroup + "Data types" + [ duckDBSpecificTypesRoundTrip, + duckDBUnsignedIntegersRoundTrip + ] + ] + +withInMemory :: (Connection -> IO a) -> IO a +withInMemory = bracket (DuckDB.open ":memory:") DuckDB.close + +newtype WithPkT f = WithPkT + { _with_pk_value :: C f Int32 + } + deriving (Generic, Beamable) + +instance Table WithPkT where + newtype PrimaryKey WithPkT f = WithPkPk (C f Int32) + deriving (Generic, Beamable) + + primaryKey = WithPkPk . _with_pk_value + +newtype WithPkDb entity = WithPkDb + { _with_pk :: entity (TableEntity WithPkT) + } + deriving (Generic, Database DuckDB) + +withPkDbChecked :: CheckedDatabaseSettings DuckDB WithPkDb +withPkDbChecked = defaultMigratableDbSettings + +verifiesPrimaryKey :: TestTree +verifiesPrimaryKey = + testCase "verifySchema correctly detects primary key" $ + withInMemory $ \conn -> do + void $ execute_ conn "CREATE TABLE with_pk (with_pk_value INTEGER NOT NULL PRIMARY KEY)" + runBeamDuckDB conn (verifySchema migrationBackend withPkDbChecked) + >>= \case + VerificationSucceeded -> pure () + VerificationFailed failures -> + assertFailure $ "Verification failed: " ++ show failures + +newtype NotNullT f = NotNullT + { _nn_value :: C f T.Text + } + deriving (Generic, Beamable) + +instance Table NotNullT where + newtype PrimaryKey NotNullT f = NotNullPk (C f T.Text) + deriving (Generic, Beamable) + + primaryKey = NotNullPk . _nn_value + +newtype NotNullDb entity = NotNullDb + { _nn_tbl :: entity (TableEntity NotNullT) + } + deriving (Generic, Database DuckDB) + +notNullDbChecked :: CheckedDatabaseSettings DuckDB NotNullDb +notNullDbChecked = defaultMigratableDbSettings + +verifiesNotNull :: TestTree +verifiesNotNull = + testCase "verifySchema correctly detects NOT NULL columns" $ + withInMemory $ \conn -> do + void $ + execute_ + conn + "CREATE TABLE nn_tbl (nn_value VARCHAR NOT NULL PRIMARY KEY)" + runBeamDuckDB conn (verifySchema migrationBackend notNullDbChecked) + >>= \case + VerificationSucceeded -> pure () + VerificationFailed failures -> + assertFailure $ "Verification failed: " ++ show failures + +newtype FkParentT f = FkParentT + { _fk_parent_id :: C f Int32 + } + deriving (Generic, Beamable) + +instance Table FkParentT where + newtype PrimaryKey FkParentT f = FkParentPk (C f Int32) + deriving (Generic, Beamable) + primaryKey = FkParentPk . _fk_parent_id + +data FkChildT f = FkChildT + { _fk_child_id :: C f Int32, + _fk_child_parent_id :: PrimaryKey FkParentT f + } + deriving (Generic, Beamable) + +instance Table FkChildT where + newtype PrimaryKey FkChildT f = FkChildPk (C f Int32) + deriving (Generic, Beamable) + primaryKey = FkChildPk . _fk_child_id + +data FkDb entity = FkDb + { _fk_parent :: entity (TableEntity FkParentT), + _fk_child :: entity (TableEntity FkChildT) + } + deriving (Generic, Database DuckDB) + +verifiesForeignKey :: TestTree +verifiesForeignKey = + testCase "getDbConstraints sees both parent and child tables" $ + withInMemory $ \conn -> do + void $ + execute_ + conn + "CREATE TABLE fk_parent (fk_parent_id INTEGER NOT NULL PRIMARY KEY)" + void $ + execute_ + conn + "CREATE TABLE fk_child (\ + \ fk_child_id INTEGER NOT NULL PRIMARY KEY,\ + \ fk_child_parent_id INTEGER NOT NULL,\ + \ FOREIGN KEY (fk_child_parent_id) REFERENCES fk_parent (fk_parent_id))" + preds <- runBeamDuckDB conn getDbConstraints + assertBool "should see fk_parent in predicates" $ + any (predicateNamesTable "fk_parent") preds + assertBool "should see fk_child in predicates" $ + any (predicateNamesTable "fk_child") preds + where + predicateNamesTable :: T.Text -> SomeDatabasePredicate -> Bool + predicateNamesTable name (SomeDatabasePredicate p) = + name `T.isInfixOf` T.pack (englishDescription p) + +autoMigrateCreatesSchema :: TestTree +autoMigrateCreatesSchema = + testCase "autoMigrate creates a schema that round-trips through verifySchema" $ + withInMemory $ \conn -> do + runBeamDuckDB conn $ autoMigrate migrationBackend withPkDbChecked + runBeamDuckDB conn (verifySchema migrationBackend withPkDbChecked) + >>= \case + VerificationSucceeded -> pure () + VerificationFailed failures -> + assertFailure $ "Verification failed: " ++ show failures + preds <- runBeamDuckDB conn getDbConstraints + assertBool "should see with_pk table after autoMigrate" $ + any (predicateNamesTable "with_pk") preds + where + predicateNamesTable :: T.Text -> SomeDatabasePredicate -> Bool + predicateNamesTable name (SomeDatabasePredicate p) = + name `T.isInfixOf` T.pack (englishDescription p) + +autoMigrateIsIdempotent :: TestTree +autoMigrateIsIdempotent = + testCase "autoMigrate is idempotent" $ + withInMemory $ \conn -> do + runBeamDuckDB conn $ autoMigrate migrationBackend withPkDbChecked + runBeamDuckDB conn $ autoMigrate migrationBackend withPkDbChecked + runBeamDuckDB conn (verifySchema migrationBackend withPkDbChecked) >>= \case + VerificationSucceeded -> pure () + VerificationFailed failures -> + assertFailure $ "Verification failed: " ++ show failures + +data TypesT f = TypesT + { _types_pk :: C f Int32, + _types_text :: C f T.Text, + _types_blob :: C f ByteString, + _types_tinyint :: C f Int8, + _types_uuid :: C f UUID + } + deriving (Generic, Beamable) + +instance Table TypesT where + newtype PrimaryKey TypesT f = TypesPk (C f Int32) + deriving (Generic, Beamable) + primaryKey = TypesPk . _types_pk + +newtype TypesDb entity = TypesDb + { _types_tbl :: entity (TableEntity TypesT) + } + deriving (Generic, Database DuckDB) + +duckDBSpecificTypesRoundTrip :: TestTree +duckDBSpecificTypesRoundTrip = + testCase "verifySchema round-trips schemas that use text/blob/tinyint/uuid" $ + withInMemory $ \conn -> do + -- Create the table directly with the DuckDB-specific keywords. The + -- schema we'll verify against, defined below via 'createTable' + 'field', + -- should round-trip cleanly against information_schema. + void $ + execute_ + conn + "CREATE TABLE types_tbl (\ + \ types_pk INTEGER NOT NULL PRIMARY KEY,\ + \ types_text TEXT NOT NULL,\ + \ types_blob BLOB NOT NULL,\ + \ types_tinyint TINYINT NOT NULL,\ + \ types_uuid UUID NOT NULL)" + let dsl = do + typesTbl <- + createTable + "types_tbl" + TypesT + { _types_pk = field "types_pk" int notNull, + _types_text = field "types_text" text notNull, + _types_blob = field "types_blob" blob notNull, + _types_tinyint = field "types_tinyint" tinyint notNull, + _types_uuid = field "types_uuid" uuid notNull + } + pure (TypesDb typesTbl) + -- Drop the table so executeMigration can re-create it from the DSL, + -- then verify it matches. This is the standard 'beam-migrate' flow. + void $ execute_ conn "DROP TABLE types_tbl" + result <- runBeamDuckDB conn $ do + db <- executeMigration runNoReturn dsl + verifySchema migrationBackend db + case result of + VerificationSucceeded -> pure () + VerificationFailed failures -> + assertFailure $ "Verification failed: " ++ show failures + +data UIntsT f = UIntsT + { _uints_pk :: C f Int32, + _uints_utiny :: C f Word8, + _uints_usmall :: C f Word16, + _uints_uint :: C f Word32, + _uints_ubig :: C f Word64 + } + deriving (Generic, Beamable) + +instance Table UIntsT where + newtype PrimaryKey UIntsT f = UIntsPk (C f Int32) + deriving (Generic, Beamable) + primaryKey = UIntsPk . _uints_pk + +newtype UIntsDb entity = UIntsDb + { _uints_tbl :: entity (TableEntity UIntsT) + } + deriving (Generic, Database DuckDB) + +duckDBUnsignedIntegersRoundTrip :: TestTree +duckDBUnsignedIntegersRoundTrip = + testCase "verifySchema round-trips schemas that use utinyint/usmallint/uinteger/ubigint" $ + withInMemory $ \conn -> do + let dsl = do + uintsTbl <- + createTable + "uints_tbl" + UIntsT + { _uints_pk = field "uints_pk" int notNull, + _uints_utiny = field "uints_utiny" utinyint notNull, + _uints_usmall = field "uints_usmall" usmallint notNull, + _uints_uint = field "uints_uint" uinteger notNull, + _uints_ubig = field "uints_ubig" ubigint notNull + } + pure (UIntsDb uintsTbl) + result <- runBeamDuckDB conn $ do + db <- executeMigration runNoReturn dsl + verifySchema migrationBackend db + case result of + VerificationSucceeded -> pure () + VerificationFailed failures -> + assertFailure $ "Verification failed: " ++ show failures + +data WidgetsV1T f = WidgetsV1T + { _v1_id :: C f Int32, + _v1_name :: C f T.Text + } + deriving (Generic, Beamable) + +instance Table WidgetsV1T where + newtype PrimaryKey WidgetsV1T f = WidgetsV1Pk (C f Int32) + deriving (Generic, Beamable) + primaryKey = WidgetsV1Pk . _v1_id + +data WidgetsV2T f = WidgetsV2T + { _v2_id :: C f Int32, + _v2_name :: C f T.Text, + -- widgets gains a nullable 'price' column. DuckDB rejects + -- @ALTER TABLE ... ADD COLUMN ... @ with the parser error + -- "Adding columns with constraints not yet supported", so the added + -- column is nullable here. + _v2_price :: C f (Maybe Double) + } + deriving (Generic, Beamable) + +instance Table WidgetsV2T where + newtype PrimaryKey WidgetsV2T f = WidgetsV2Pk (C f Int32) + deriving (Generic, Beamable) + primaryKey = WidgetsV2Pk . _v2_id + +newtype WidgetsV1Db entity = WidgetsV1Db + { _v1_widgets :: entity (TableEntity WidgetsV1T) + } + deriving (Generic, Database DuckDB) + +newtype WidgetsV2Db entity = WidgetsV2Db + { _v2_widgets :: entity (TableEntity WidgetsV2T) + } + deriving (Generic, Database DuckDB) + +stepV1 :: () -> Migration DuckDB (CheckedDatabaseSettings DuckDB WidgetsV1Db) +stepV1 () = + WidgetsV1Db + <$> createTable + "widgets" + WidgetsV1T + { _v1_id = field "id" int notNull, + _v1_name = field "name" (varchar Nothing) notNull + } + +stepV2 :: + CheckedDatabaseSettings DuckDB WidgetsV1Db -> + Migration DuckDB (CheckedDatabaseSettings DuckDB WidgetsV2Db) +stepV2 oldDb = + WidgetsV2Db + <$> alterTable + (_v1_widgets oldDb) + ( \oldTbl -> do + price <- addColumn (field "price" (maybeType double)) + pure + WidgetsV2T + { _v2_id = _v1_id oldTbl, + _v2_name = _v1_name oldTbl, + _v2_price = price + } + ) + +stepV3 :: + CheckedDatabaseSettings DuckDB WidgetsV2Db -> + Migration DuckDB (CheckedDatabaseSettings DuckDB WidgetsV1Db) +stepV3 (WidgetsV2Db v2Table) = + WidgetsV1Db + <$> alterTable + v2Table + ( \v2Table' -> do + dropColumn (_v2_price v2Table') + pure (WidgetsV1T (_v2_id v2Table') (_v2_name v2Table')) + ) + +runVerifySchema :: + (Database DuckDB db) => + Connection -> + CheckedDatabaseSettings DuckDB db -> + Assertion +runVerifySchema conn db = + runBeamDuckDB conn (verifySchema migrationBackend db) >>= \case + VerificationSucceeded -> pure () + VerificationFailed failures -> + assertFailure $ "Verification failed: " ++ show failures + +multiStepMigration :: TestTree +multiStepMigration = + testCaseSteps "two-step migration: create + alter table verifies against final schema" $ \testStep -> + withInMemory $ \conn -> do + let getColumns = + ( DuckDB.query_ + conn + "SELECT column_name FROM information_schema.columns \ + \WHERE table_name = 'widgets' \ + \ORDER BY ordinal_position" :: + IO [DuckDB.Only T.Text] + ) + <&> map (\(DuckDB.Only c) -> c) + + testStep "Initial migration" + v1Db <- runBeamDuckDB conn $ + runMigrationSteps 0 Nothing (migrationStep "V1" stepV1) $ \_ _ step -> + executeMigration runNoReturn step + runVerifySchema conn v1Db + + getColumns >>= \cols -> + assertEqual + ("expected 'price' in columns " ++ show cols) + ["id", "name"] + cols + + testStep "Migration across V1 and V2" + v2Db <- runBeamDuckDB conn + $ runMigrationSteps + 1 + Nothing + ( migrationStep "V1" stepV1 + >>> migrationStep "V2" stepV2 + ) + $ \_ _ step -> + executeMigration runNoReturn step + runVerifySchema conn v2Db + + getColumns >>= \cols -> + assertEqual + ("expected 'price' in columns " ++ show cols) + ["id", "name", "price"] + cols + + testStep "Migration back to V1" + v3Db <- runBeamDuckDB conn + $ runMigrationSteps + 2 + Nothing + ( migrationStep "V1" stepV1 + >>> migrationStep "V2" stepV2 + >>> migrationStep "V3" stepV3 + ) + $ \_ _ step -> + executeMigration runNoReturn step + runVerifySchema conn v3Db + + getColumns >>= \cols -> + assertEqual + ("expected price column dropped, in columns " ++ show cols) + ["id", "name"] + cols diff --git a/tests/Database/Beam/DuckDB/Test/Query.hs b/tests/Database/Beam/DuckDB/Test/Query.hs --- a/tests/Database/Beam/DuckDB/Test/Query.hs +++ b/tests/Database/Beam/DuckDB/Test/Query.hs @@ -55,12 +55,14 @@ isTrue_, lagWithDefault_, lastValue_, + lead1_, leftJoin_, modifyTableFields, noBounds_, noOrder_, noPartition_, nrows_, + nub_, orderBy_, orderPartitionBy_, over_, @@ -148,7 +150,10 @@ ], testGroup "LAG/LEAD" - [testLag, testLead], + [ testLag, + testLead, + testLeadOverNub + ], testGroup "FILTER" [ testFilterWhereCountAll, @@ -779,6 +784,34 @@ -- Second leads to third's qty (1) -- Third has no successor → default 0 rows @?= [0, 2, 1] + +-- Regression test for #746 +testLeadOverNub :: TestTree +testLeadOverNub = testCase "LEAD over nub_ (issue #746)" $ do + let users = [User 1 "Alice" 30] + products = [Product 1 "Widget" 999] + orders = + [ Order 1 (UserId 1) (ProductId 1) 10, + Order 2 (UserId 1) (ProductId 1) 10, + Order 3 (UserId 1) (ProductId 1) 20, + Order 4 (UserId 1) (ProductId 1) 20, + Order 5 (UserId 1) (ProductId 1) 30, + Order 6 (UserId 1) (ProductId 1) 30 + ] + withTestDb users products orders $ \conn -> do + rows <- + runBeamDuckDB conn $ + runSelectReturningList $ + select $ + withWindow_ + (\q -> frame_ noPartition_ (orderPartitionBy_ (asc_ q)) noBounds_) + (\q w -> (q, lead1_ q `over_` w)) + (nub_ (_orderQuantity <$> all_ (_dbOrders testDb))) + rows + @?= [ (10 :: Int32, Just 20), + (20, Just 30), + (30, Nothing) + ] testFilterWhereCountAll :: TestTree testFilterWhereCountAll = testCase "COUNT(*) FILTER (WHERE ... )" $ do diff --git a/tests/Main.hs b/tests/Main.hs --- a/tests/Main.hs +++ b/tests/Main.hs @@ -1,14 +1,16 @@ module Main (main) where import qualified Database.Beam.DuckDB.Test.Extensions (tests) +import qualified Database.Beam.DuckDB.Test.Migrate (tests) import qualified Database.Beam.DuckDB.Test.Query (tests) import Test.Tasty (defaultMain, testGroup) main :: IO () main = - defaultMain $ - testGroup - "beam-duckdb tests" - [ Database.Beam.DuckDB.Test.Query.tests - , Database.Beam.DuckDB.Test.Extensions.tests - ] + defaultMain $ + testGroup + "beam-duckdb tests" + [ Database.Beam.DuckDB.Test.Query.tests, + Database.Beam.DuckDB.Test.Extensions.tests, + Database.Beam.DuckDB.Test.Migrate.tests + ]