diff --git a/ChangeLog.md b/ChangeLog.md --- a/ChangeLog.md +++ b/ChangeLog.md @@ -1,3 +1,31 @@ +# 0.6.1.0 + +## Added features + +* Added file-mode `COPY ... TO 'file'` / `COPY ... FROM 'file'` support + via the new `MonadBeamCopyTo` / `MonadBeamCopyFrom` instances on `Pg`. + Smart constructors `copyToText` / `copyToCSV` (and `copyFromText` / + `copyFromCSV`, plus `*With` variants) build the per-format options + records. Note that this requires the `pg_write_server_files` / + `pg_read_server_files` role (or superuser) on the connecting role — + see `Database.Beam.Postgres.Extensions.Copy.File`. +* Added streaming `COPY ... TO STDOUT` / `COPY ... FROM STDIN` support + via the new `MonadBeamCopyToStream` / `MonadBeamCopyFromStream` + instances on `Pg`. Smart constructors `copyToTextStream` / + `copyToCSVStream` / `copyFromTextStream` / `copyFromCSVStream` build + the per-format options. Streaming COPY does not require any special + role attribute and is the appropriate choice when the client and + server are on different hosts — see + `Database.Beam.Postgres.Extensions.Copy.Stream`. + +## Bug fixes + +* Fixed an issue where a window function applied over the result of `nub_` (or + the Postgres-specific `pgNubBy_`) would emit `DISTINCT` and the window + expression in the same `SELECT`, causing the window to evaluate against the + pre-deduplicated rows. The inner select is now materialised as a subquery + whenever it carries `DISTINCT`, `GROUP BY`, or `HAVING` (#756). + # 0.6.0.0 ## Interface changes diff --git a/Database/Beam/Postgres.hs b/Database/Beam/Postgres.hs --- a/Database/Beam/Postgres.hs +++ b/Database/Beam/Postgres.hs @@ -16,74 +16,202 @@ -- -- For examples on how to use @beam-postgres@ usage, see -- . - module Database.Beam.Postgres - ( -- * Beam Postgres backend - Postgres(..), Pg, liftIOWithHandle + ( -- * Beam Postgres backend + Postgres (..), + Pg, + liftIOWithHandle, -- ** Executing actions against the backend - , runBeamPostgres, runBeamPostgresDebug + runBeamPostgres, + runBeamPostgresDebug, -- ** Postgres syntax - , PgCommandSyntax, PgSyntax - , PgSelectSyntax, PgInsertSyntax - , PgUpdateSyntax, PgDeleteSyntax + PgCommandSyntax, + PgSyntax, + PgSelectSyntax, + PgInsertSyntax, + PgUpdateSyntax, + PgDeleteSyntax, - -- * Beam URI support - , postgresUriSyntax + -- ** COPY support - -- * Postgres-specific features - -- ** Postgres-specific data types + -- + -- File-mode @COPY@ support for Postgres. Use 'copyToText' / 'copyToCSV' (or 'copyFromText' / + -- 'copyFromCSV') for default options, and the @*With@ variants to override + -- format-specific fields. - , json, jsonb, uuid, money - , tsquery, tsvector, text, bytea - , unboundedArray + -- *** text format + copyToText, + copyToTextWith, + PgTextCopyToOptions (..), + defaultPgTextCopyToOptions, + copyFromText, + copyFromTextWith, + PgTextCopyFromOptions (..), + defaultPgTextCopyFromOptions, - -- *** @SERIAL@ support - , smallserial, serial, bigserial + -- *** CSV format + copyToCSV, + copyToCSVWith, + PgCSVCopyToOptions (..), + defaultPgCSVCopyToOptions, + copyFromCSV, + copyFromCSVWith, + PgCSVCopyFromOptions (..), + defaultPgCSVCopyFromOptions, - , module Database.Beam.Postgres.PgSpecific - , module Database.Beam.Postgres.TempTable + -- *** Top-level options sums + PgCopyToOptions, + PgCopyFromOptions, - -- ** Postgres extension support - , PgExtensionEntity, IsPgExtension(..) - , pgCreateExtension, pgDropExtension - , getPgExtension + -- ** Streaming COPY support - -- ** Utilities for defining custom instances - , fromPgIntegral - , fromPgScientificOrIntegral + -- + -- Streaming @COPY ... TO STDOUT@ / @COPY ... FROM STDIN@. + -- Use 'copyToTextStream' / 'copyToCSVStream' (or 'copyFromTextStream' / + -- 'copyFromCSVStream') with the 'runCopyToStream' / 'runCopyFromStream' + -- runners from "Database.Beam.Backend.SQL.BeamExtensions". - -- ** Debug support + -- *** text format + copyToTextStream, + copyToTextStreamWith, + copyFromTextStream, + copyFromTextStreamWith, - , PgDebugStmt - , pgTraceStmtIO, pgTraceStmtIO' - , pgTraceStmt + -- *** CSV format + copyToCSVStream, + copyToCSVStreamWith, + copyFromCSVStream, + copyFromCSVStreamWith, - -- * @postgresql-simple@ re-exports + -- *** Top-level options sums + PgCopyToStreamOptions, + PgCopyFromStreamOptions, - , Pg.ResultError(..), Pg.SqlError(..) + -- * Beam URI support + postgresUriSyntax, - , Pg.Connection, Pg.ConnectInfo(..) - , Pg.defaultConnectInfo + -- * Postgres-specific features - , Pg.connectPostgreSQL, Pg.connect - , Pg.close + -- ** Postgres-specific data types + json, + jsonb, + uuid, + money, + tsquery, + tsvector, + text, + bytea, + unboundedArray, - ) where + -- *** @SERIAL@ support + smallserial, + serial, + bigserial, + module Database.Beam.Postgres.PgSpecific, + module Database.Beam.Postgres.TempTable, + -- ** Postgres extension support + PgExtensionEntity, + IsPgExtension (..), + pgCreateExtension, + pgDropExtension, + getPgExtension, + + -- ** Utilities for defining custom instances + fromPgIntegral, + fromPgScientificOrIntegral, + + -- ** Debug support + PgDebugStmt, + pgTraceStmtIO, + pgTraceStmtIO', + pgTraceStmt, + + -- * @postgresql-simple@ re-exports + Pg.ResultError (..), + Pg.SqlError (..), + Pg.Connection, + Pg.ConnectInfo (..), + Pg.defaultConnectInfo, + Pg.connectPostgreSQL, + Pg.connect, + Pg.close, + ) +where + import Database.Beam.Postgres.Connection -import Database.Beam.Postgres.Full () -- for BeamHasInsertOnConflict instance -import Database.Beam.Postgres.Syntax -import Database.Beam.Postgres.Types -import Database.Beam.Postgres.PgSpecific -import Database.Beam.Postgres.Migrate ( tsquery, tsvector, text, bytea, unboundedArray - , json, jsonb, uuid, money, smallserial, serial - , bigserial) -import Database.Beam.Postgres.Extensions ( PgExtensionEntity, IsPgExtension(..) - , pgCreateExtension, pgDropExtension - , getPgExtension ) +-- for BeamHasInsertOnConflict instance + import Database.Beam.Postgres.Debug +import Database.Beam.Postgres.Extensions + ( IsPgExtension (..), + PgExtensionEntity, + getPgExtension, + pgCreateExtension, + pgDropExtension, + ) +import Database.Beam.Postgres.Extensions.Copy.File + ( PgCSVCopyFromOptions (..), + PgCSVCopyToOptions (..), + PgCopyFromOptions, + PgCopyToOptions, + PgTextCopyFromOptions (..), + PgTextCopyToOptions (..), + copyFromCSV, + copyFromCSVWith, + copyFromText, + copyFromTextWith, + copyToCSV, + copyToCSVWith, + copyToText, + copyToTextWith, + defaultPgCSVCopyFromOptions, + defaultPgCSVCopyToOptions, + defaultPgTextCopyFromOptions, + defaultPgTextCopyToOptions, + ) +import Database.Beam.Postgres.Extensions.Copy.Stream + ( PgCopyFromStreamOptions, + PgCopyToStreamOptions, + copyFromCSVStream, + copyFromCSVStreamWith, + copyFromTextStream, + copyFromTextStreamWith, + copyToCSVStream, + copyToCSVStreamWith, + copyToTextStream, + copyToTextStreamWith, + ) +import Database.Beam.Postgres.Full () +import Database.Beam.Postgres.Migrate + ( bigserial, + bytea, + json, + jsonb, + money, + serial, + smallserial, + text, + tsquery, + tsvector, + unboundedArray, + uuid, + ) +import Database.Beam.Postgres.PgSpecific +import Database.Beam.Postgres.Syntax + ( PgCommandSyntax, + PgDeleteSyntax, + PgInsertSyntax, + PgSelectSyntax, + PgSyntax, + PgUpdateSyntax, + ) import Database.Beam.Postgres.TempTable - +import Database.Beam.Postgres.Types + ( Postgres (..), + fromPgIntegral, + fromPgScientificOrIntegral, + ) import qualified Database.PostgreSQL.Simple as Pg diff --git a/Database/Beam/Postgres/Connection.hs b/Database/Beam/Postgres/Connection.hs --- a/Database/Beam/Postgres/Connection.hs +++ b/Database/Beam/Postgres/Connection.hs @@ -25,7 +25,8 @@ , postgresUriSyntax ) where -import Control.Exception (SomeException(..), throwIO) +import Control.Exception (SomeException(..), throwIO, onException, catch) +import Control.Monad (void) import Control.Monad.Base (MonadBase(..)) import Control.Monad.Free.Church import Control.Monad.IO.Class @@ -41,6 +42,10 @@ import Database.Beam.Backend.URI import Database.Beam.Schema.Tables +import Database.Beam.Postgres.Extensions.Copy.File + ( PgCopyFromSyntax(..), PgCopyToSyntax(..) ) +import Database.Beam.Postgres.Extensions.Copy.Stream + ( PgCopyFromStreamSyntax(..), PgCopyToStreamSyntax(..) ) import Database.Beam.Postgres.Syntax import Database.Beam.Postgres.Full import Database.Beam.Postgres.Types @@ -48,6 +53,7 @@ import qualified Database.PostgreSQL.LibPQ as Pg hiding (Connection, escapeStringConn, escapeIdentifier, escapeByteaConn, exec) import qualified Database.PostgreSQL.Simple as Pg +import qualified Database.PostgreSQL.Simple.Copy as PgCopy import qualified Database.PostgreSQL.Simple.FromField as Pg import qualified Database.PostgreSQL.Simple.Internal as Pg ( Field(..), RowParser(..) @@ -406,6 +412,50 @@ Nothing -> pure (acc []) Just x -> collectM (acc . (x:)) in collectM id + +instance MonadBeamCopyTo Postgres Pg where + runCopyTo SqlCopyToNoColumns = pure () + runCopyTo (SqlCopyTo (PgCopyToSyntax syntax)) = + runNoReturn (PgCommandSyntax PgCommandTypeDataUpdate syntax) + +instance MonadBeamCopyFrom Postgres Pg where + runCopyFrom SqlCopyFromNoColumns = pure () + runCopyFrom (SqlCopyFrom (PgCopyFromSyntax syntax)) = + runNoReturn (PgCommandSyntax PgCommandTypeDataUpdate syntax) + +instance MonadBeamCopyToStream Postgres Pg where + -- | `runCopyToStream` is exception-safe; if the output stream + -- produces an exception, the stream will be drained before the exception + -- is re-thrown + runCopyToStream SqlCopyToStreamNoColumns _ = pure () + runCopyToStream (SqlCopyToStream (PgCopyToStreamSyntax syntax)) sink = + liftIOWithHandle $ \conn -> do + query <- pgRenderSyntax conn syntax + PgCopy.copy_ conn (Pg.Query query) + let loop onRow = do + PgCopy.getCopyData conn >>= \case + PgCopy.CopyOutRow chunk -> onRow chunk >> loop onRow + PgCopy.CopyOutDone _ -> pure () + -- Like 'loop', but doesn't use the sink at all. This is used + -- to drain elements in the COPY stream before re-throwing an exception + drain = loop (const (pure ())) + loop sink `catch` (\(e::SomeException) -> drain >> throwIO e) + + +instance MonadBeamCopyFromStream Postgres Pg where + -- | `runCopyFromStream` is exception-safe; if the input stream + -- produces an exception, the connection will be reset to a safe state, + -- aborting the COPY operation. + runCopyFromStream SqlCopyFromStreamNoColumns _ = pure () + runCopyFromStream (SqlCopyFromStream (PgCopyFromStreamSyntax syntax)) producer = + liftIOWithHandle $ \conn -> do + query <- pgRenderSyntax conn syntax + let loop = + producer >>= \case + Just chunk -> PgCopy.putCopyData conn chunk >> loop + Nothing -> void $ PgCopy.putCopyEnd conn + PgCopy.copy_ conn (Pg.Query query) + loop `onException` PgCopy.putCopyError conn mempty instance MonadBeamInsertReturning Postgres Pg where runInsertReturningListWith i mkProjection = do diff --git a/Database/Beam/Postgres/Extensions/Copy/File.hs b/Database/Beam/Postgres/Extensions/Copy/File.hs new file mode 100644 --- /dev/null +++ b/Database/Beam/Postgres/Extensions/Copy/File.hs @@ -0,0 +1,421 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeFamilies #-} + +module Database.Beam.Postgres.Extensions.Copy.File + ( PgCopyToSyntax (..), + PgCopyToSourceSyntax (..), + PgCopyFromSyntax (..), + PgCopyFromSourceSyntax (..), + + -- * COPY options + PgCopyToOptions, + PgCopyFromOptions, + + -- ** text + copyToText, + copyToTextWith, + PgTextCopyToOptions (..), + defaultPgTextCopyToOptions, + copyFromText, + copyFromTextWith, + PgTextCopyFromOptions (..), + defaultPgTextCopyFromOptions, + + -- ** CSV + copyToCSV, + copyToCSVWith, + PgCSVCopyToOptions (..), + defaultPgCSVCopyToOptions, + copyFromCSV, + copyFromCSVWith, + PgCSVCopyFromOptions (..), + defaultPgCSVCopyFromOptions, + + -- * Internal — exposed for use by sibling modules + emitOptionList, + textCopyToFields, + csvCopyToFields, + textCopyFromFields, + csvCopyFromFields, + ) +where + +import qualified Data.List.NonEmpty as NE +import Data.Maybe (catMaybes) +import Data.Text (Text) +import qualified Data.Text as T +import Database.Beam.Backend.SQL.BeamExtensions + ( IsSqlCopyFromSourceSyntax (..), + IsSqlCopyFromSyntax (..), + IsSqlCopyToSourceSyntax (..), + IsSqlCopyToSyntax (..), + ) +import Database.Beam.Postgres.Syntax + ( PgSelectSyntax (..), + PgSyntax, + emit, + pgBoolLit, + pgCharLit, + pgParens, + pgQuotedIdentifier, + pgSepBy, + pgStringLit, + ) + +-- | PostgreSQL @COPY ... TO@ source syntax. Wraps the table name (with +-- optional column list) or a @SELECT@ subquery. +-- +-- @since 0.6.1.0 +newtype PgCopyToSourceSyntax = PgCopyToSourceSyntax {fromPgCopyToSource :: PgSyntax} + +-- | PostgreSQL @COPY ... TO@ statement syntax. +-- +-- @since 0.6.1.0 +newtype PgCopyToSyntax = PgCopyToSyntax {fromPgCopyTo :: PgSyntax} + +-- | PostgreSQL @COPY ... FROM@ source syntax. Just the table name plus +-- optional column list — Postgres @COPY ... FROM@ does not accept a SELECT. +-- +-- @since 0.6.1.0 +newtype PgCopyFromSourceSyntax = PgCopyFromSourceSyntax {fromPgCopyFromSource :: PgSyntax} + +-- | PostgreSQL @COPY ... FROM@ statement syntax. +-- +-- @since 0.6.1.0 +newtype PgCopyFromSyntax = PgCopyFromSyntax {fromPgCopyFrom :: PgSyntax} + +instance IsSqlCopyToSourceSyntax PgCopyToSourceSyntax where + type SqlCopyToSourceSelectSyntax PgCopyToSourceSyntax = PgSelectSyntax + + copyTableToSyntax mSchema tableName' mColumns = + PgCopyToSourceSyntax $ + maybe mempty (\schema -> pgQuotedIdentifier schema <> emit ".") mSchema + <> pgQuotedIdentifier tableName' + <> case mColumns of + Nothing -> mempty + Just cols -> pgParens $ pgSepBy (emit ", ") (map pgQuotedIdentifier (NE.toList cols)) + + copySelectToSyntax (PgSelectSyntax select) = PgCopyToSourceSyntax $ pgParens select + +instance IsSqlCopyToSyntax PgCopyToSyntax where + type SqlCopyToSourceSyntax PgCopyToSyntax = PgCopyToSourceSyntax + type SqlCopyToParams PgCopyToSyntax = PgCopyToOptions + + copyToStmt (PgCopyToSourceSyntax source) options = + let (filepath, optionsSyntax) = emitPgCopyToOptions options + in PgCopyToSyntax $ + mconcat + [ emit "COPY ", + source, + emit " TO ", + pgStringLit (T.pack filepath), + optionsSyntax + ] + +instance IsSqlCopyFromSourceSyntax PgCopyFromSourceSyntax where + copyTableFromSyntax mSchema tableName' mColumns = + PgCopyFromSourceSyntax $ + maybe mempty (\schema -> pgQuotedIdentifier schema <> emit ".") mSchema + <> pgQuotedIdentifier tableName' + <> case mColumns of + Nothing -> mempty + Just cols -> pgParens $ pgSepBy (emit ", ") (map pgQuotedIdentifier (NE.toList cols)) + +instance IsSqlCopyFromSyntax PgCopyFromSyntax where + type SqlCopyFromSourceSyntax PgCopyFromSyntax = PgCopyFromSourceSyntax + type SqlCopyFromParams PgCopyFromSyntax = PgCopyFromOptions + + copyFromStmt (PgCopyFromSourceSyntax source) options = + let (filepath, optionsSyntax) = emitPgCopyFromOptions options + in PgCopyFromSyntax $ + mconcat + [ emit "COPY ", + source, + emit " FROM ", + pgStringLit (T.pack filepath), + optionsSyntax + ] + +-- | Options for PostgreSQL's @COPY ... TO@ statement, tagged by the output +-- format. Use 'copyToText' or 'copyToCSV' to construct. +-- +-- See +-- for the full set of options each format supports. +-- +-- @since 0.6.1.0 +data PgCopyToOptions + = PgCopyToText !FilePath !PgTextCopyToOptions + | PgCopyToCSV !FilePath !PgCSVCopyToOptions + deriving (Eq, Show) + +-- | Copy to a text file at the given path with default options. +-- Use 'copyToTextWith' to override. +-- +-- @since 0.6.1.0 +copyToText :: FilePath -> PgCopyToOptions +copyToText path = PgCopyToText path defaultPgTextCopyToOptions + +-- | Copy to a text file at the given path with the given options. +-- +-- @since 0.6.1.0 +copyToTextWith :: FilePath -> PgTextCopyToOptions -> PgCopyToOptions +copyToTextWith = PgCopyToText + +-- | Copy to a CSV file at the given path with default options. +-- Use 'copyToCSVWith' to override. +-- +-- @since 0.6.1.0 +copyToCSV :: FilePath -> PgCopyToOptions +copyToCSV path = PgCopyToCSV path defaultPgCSVCopyToOptions + +-- | Copy to a CSV file at the given path with the given options. +-- +-- @since 0.6.1.0 +copyToCSVWith :: FilePath -> PgCSVCopyToOptions -> PgCopyToOptions +copyToCSVWith = PgCopyToCSV + +-- | Options for PostgreSQL's @COPY ... FROM@ statement, tagged by the input +-- format. Use 'copyFromText' or 'copyFromCSV' to construct. +-- +-- @since 0.6.1.0 +data PgCopyFromOptions + = PgCopyFromText !FilePath !PgTextCopyFromOptions + | PgCopyFromCSV !FilePath !PgCSVCopyFromOptions + deriving (Eq, Show) + +-- | Copy from a text file at the given path with default options. +-- Use 'copyFromTextWith' to override. +-- +-- @since 0.6.1.0 +copyFromText :: FilePath -> PgCopyFromOptions +copyFromText path = PgCopyFromText path defaultPgTextCopyFromOptions + +-- | Copy from a text file at the given path with the given options. +-- +-- @since 0.6.1.0 +copyFromTextWith :: FilePath -> PgTextCopyFromOptions -> PgCopyFromOptions +copyFromTextWith = PgCopyFromText + +-- | Copy from a CSV file at the given path with default options. +-- Use 'copyFromCSVWith' to override. +-- +-- @since 0.6.1.0 +copyFromCSV :: FilePath -> PgCopyFromOptions +copyFromCSV path = PgCopyFromCSV path defaultPgCSVCopyFromOptions + +-- | Copy from a CSV file at the given path with the given options. +-- +-- @since 0.6.1.0 +copyFromCSVWith :: FilePath -> PgCSVCopyFromOptions -> PgCopyFromOptions +copyFromCSVWith = PgCopyFromCSV + +-- | Options for the @text@ format on @COPY ... TO@. +-- +-- @since 0.6.1.0 +data PgTextCopyToOptions = PgTextCopyToOptions + { -- | @DELIMITER@: single-character column delimiter. Defaults to TAB. + pgTextCopyToDelimiter :: Maybe Char, + -- | @NULL@: token written for SQL @NULL@ values. Defaults to @\\N@. + pgTextCopyToNullStr :: Maybe Text, + -- | @HEADER@: emit a header row with column names. + pgTextCopyToHeader :: Maybe Bool, + -- | @ENCODING@: client-encoding override for the file. + pgTextCopyToEncoding :: Maybe Text + } + deriving (Eq, Show) + +-- | All-'Nothing' 'PgTextCopyToOptions'; a sensible starting point for +-- record-update overrides. +-- +-- @since 0.6.1.0 +defaultPgTextCopyToOptions :: PgTextCopyToOptions +defaultPgTextCopyToOptions = + PgTextCopyToOptions + { pgTextCopyToDelimiter = Nothing, + pgTextCopyToNullStr = Nothing, + pgTextCopyToHeader = Nothing, + pgTextCopyToEncoding = Nothing + } + +-- | Options for the @csv@ format on @COPY ... TO@. +-- +-- @since 0.6.1.0 +data PgCSVCopyToOptions = PgCSVCopyToOptions + { -- | @DELIMITER@: single-character column delimiter. Defaults to a comma. + pgCsvCopyToDelimiter :: Maybe Char, + -- | @NULL@: token written for SQL @NULL@ values. Defaults to the empty + -- string. + pgCsvCopyToNullStr :: Maybe Text, + -- | @HEADER@: emit a header row with column names. + pgCsvCopyToHeader :: Maybe Bool, + -- | @QUOTE@: single-character quote character. Defaults to @\"@. + pgCsvCopyToQuote :: Maybe Char, + -- | @ESCAPE@: single-character escape character. Defaults to the + -- value of 'pgCsvCopyToQuote'. + pgCsvCopyToEscape :: Maybe Char, + -- | @ENCODING@: client-encoding override for the file. + pgCsvCopyToEncoding :: Maybe Text + } + deriving (Eq, Show) + +-- | All-'Nothing' 'PgCSVCopyToOptions'. +-- +-- @since 0.6.1.0 +defaultPgCSVCopyToOptions :: PgCSVCopyToOptions +defaultPgCSVCopyToOptions = + PgCSVCopyToOptions + { pgCsvCopyToDelimiter = Nothing, + pgCsvCopyToNullStr = Nothing, + pgCsvCopyToHeader = Nothing, + pgCsvCopyToQuote = Nothing, + pgCsvCopyToEscape = Nothing, + pgCsvCopyToEncoding = Nothing + } + +-- | Options for the @text@ format on @COPY ... FROM@. +-- +-- @since 0.6.1.0 +data PgTextCopyFromOptions = PgTextCopyFromOptions + { -- | @DELIMITER@: single-character column delimiter. Defaults to TAB. + pgTextCopyFromDelimiter :: Maybe Char, + -- | @NULL@: token to interpret as SQL @NULL@. Defaults to @\\N@. + pgTextCopyFromNullStr :: Maybe Text, + -- | @HEADER@: skip the first line as a header row. + pgTextCopyFromHeader :: Maybe Bool, + -- | @ENCODING@: client-encoding override for the file. + pgTextCopyFromEncoding :: Maybe Text, + -- | @FREEZE@: skip WAL when loading; only legal under specific + -- conditions (see the PostgreSQL @COPY@ documentation). + pgTextCopyFromFreeze :: Maybe Bool + } + deriving (Eq, Show) + +-- | All-'Nothing' 'PgTextCopyFromOptions'. +-- +-- @since 0.6.1.0 +defaultPgTextCopyFromOptions :: PgTextCopyFromOptions +defaultPgTextCopyFromOptions = + PgTextCopyFromOptions + { pgTextCopyFromDelimiter = Nothing, + pgTextCopyFromNullStr = Nothing, + pgTextCopyFromHeader = Nothing, + pgTextCopyFromEncoding = Nothing, + pgTextCopyFromFreeze = Nothing + } + +-- | Options for the @csv@ format on @COPY ... FROM@. +-- +-- @since 0.6.1.0 +data PgCSVCopyFromOptions = PgCSVCopyFromOptions + { -- | @DELIMITER@: single-character column delimiter. Defaults to a comma. + pgCsvCopyFromDelimiter :: Maybe Char, + -- | @NULL@: token to interpret as SQL @NULL@. Defaults to the empty + -- string. + pgCsvCopyFromNullStr :: Maybe Text, + -- | @HEADER@: skip the first line as a header row. + pgCsvCopyFromHeader :: Maybe Bool, + -- | @QUOTE@: single-character quote character. Defaults to @\"@. + pgCsvCopyFromQuote :: Maybe Char, + -- | @ESCAPE@: single-character escape character. Defaults to the + -- value of 'pgCsvCopyFromQuote'. + pgCsvCopyFromEscape :: Maybe Char, + -- | @ENCODING@: client-encoding override for the file. + pgCsvCopyFromEncoding :: Maybe Text, + -- | @FREEZE@: skip WAL when loading; only legal under specific + -- conditions (see the PostgreSQL @COPY@ documentation). + pgCsvCopyFromFreeze :: Maybe Bool + } + deriving (Eq, Show) + +-- | All-'Nothing' 'PgCSVCopyFromOptions'; a sensible starting point for +-- record-update overrides. +-- +-- @since 0.6.1.0 +defaultPgCSVCopyFromOptions :: PgCSVCopyFromOptions +defaultPgCSVCopyFromOptions = + PgCSVCopyFromOptions + { pgCsvCopyFromDelimiter = Nothing, + pgCsvCopyFromNullStr = Nothing, + pgCsvCopyFromHeader = Nothing, + pgCsvCopyFromQuote = Nothing, + pgCsvCopyFromEscape = Nothing, + pgCsvCopyFromEncoding = Nothing, + pgCsvCopyFromFreeze = Nothing + } + +-- * Emission + +emitPgCopyToOptions :: PgCopyToOptions -> (FilePath, PgSyntax) +emitPgCopyToOptions (PgCopyToText fp o) = (fp, emitOptionList (emit "text") (textCopyToFields o)) +emitPgCopyToOptions (PgCopyToCSV fp o) = (fp, emitOptionList (emit "csv") (csvCopyToFields o)) + +emitPgCopyFromOptions :: PgCopyFromOptions -> (FilePath, PgSyntax) +emitPgCopyFromOptions (PgCopyFromText fp o) = (fp, emitOptionList (emit "text") (textCopyFromFields o)) +emitPgCopyFromOptions (PgCopyFromCSV fp o) = (fp, emitOptionList (emit "csv") (csvCopyFromFields o)) + +-- | Render @ (FORMAT FMT, OPT1 …, OPT2 …)@. The @FORMAT@ entry is always +-- present since the constructor of the options value pins it. +-- +-- @since 0.6.1.0 +emitOptionList :: PgSyntax -> [PgSyntax] -> PgSyntax +emitOptionList fmt items = + emit " (" + <> pgSepBy (emit ", ") ((emit "FORMAT " <> fmt) : items) + <> emit ")" + +-- | Render the per-option items for a 'PgTextCopyToOptions' value. +-- +-- @since 0.6.1.0 +textCopyToFields :: PgTextCopyToOptions -> [PgSyntax] +textCopyToFields o = + catMaybes + [ fmap (\c -> emit "DELIMITER " <> pgCharLit c) (pgTextCopyToDelimiter o), + fmap (\s -> emit "NULL " <> pgStringLit s) (pgTextCopyToNullStr o), + fmap (\b -> emit "HEADER " <> pgBoolLit b) (pgTextCopyToHeader o), + fmap (\s -> emit "ENCODING " <> pgStringLit s) (pgTextCopyToEncoding o) + ] + +-- | Render the per-option items for a 'PgCSVCopyToOptions' value. +-- +-- @since 0.6.1.0 +csvCopyToFields :: PgCSVCopyToOptions -> [PgSyntax] +csvCopyToFields o = + catMaybes + [ fmap (\c -> emit "DELIMITER " <> pgCharLit c) (pgCsvCopyToDelimiter o), + fmap (\s -> emit "NULL " <> pgStringLit s) (pgCsvCopyToNullStr o), + fmap (\b -> emit "HEADER " <> pgBoolLit b) (pgCsvCopyToHeader o), + fmap (\c -> emit "QUOTE " <> pgCharLit c) (pgCsvCopyToQuote o), + fmap (\c -> emit "ESCAPE " <> pgCharLit c) (pgCsvCopyToEscape o), + fmap (\s -> emit "ENCODING " <> pgStringLit s) (pgCsvCopyToEncoding o) + ] + +-- | Render the per-option items for a 'PgTextCopyFromOptions' value. +-- +-- @since 0.6.1.0 +textCopyFromFields :: PgTextCopyFromOptions -> [PgSyntax] +textCopyFromFields o = + catMaybes + [ fmap (\c -> emit "DELIMITER " <> pgCharLit c) (pgTextCopyFromDelimiter o), + fmap (\s -> emit "NULL " <> pgStringLit s) (pgTextCopyFromNullStr o), + fmap (\b -> emit "HEADER " <> pgBoolLit b) (pgTextCopyFromHeader o), + fmap (\s -> emit "ENCODING " <> pgStringLit s) (pgTextCopyFromEncoding o), + fmap (\b -> emit "FREEZE " <> pgBoolLit b) (pgTextCopyFromFreeze o) + ] + +-- | Render the per-option items for a 'PgCSVCopyFromOptions' value. +-- +-- @since 0.6.1.0 +csvCopyFromFields :: PgCSVCopyFromOptions -> [PgSyntax] +csvCopyFromFields o = + catMaybes + [ fmap (\c -> emit "DELIMITER " <> pgCharLit c) (pgCsvCopyFromDelimiter o), + fmap (\s -> emit "NULL " <> pgStringLit s) (pgCsvCopyFromNullStr o), + fmap (\b -> emit "HEADER " <> pgBoolLit b) (pgCsvCopyFromHeader o), + fmap (\c -> emit "QUOTE " <> pgCharLit c) (pgCsvCopyFromQuote o), + fmap (\c -> emit "ESCAPE " <> pgCharLit c) (pgCsvCopyFromEscape o), + fmap (\s -> emit "ENCODING " <> pgStringLit s) (pgCsvCopyFromEncoding o), + fmap (\b -> emit "FREEZE " <> pgBoolLit b) (pgCsvCopyFromFreeze o) + ] diff --git a/Database/Beam/Postgres/Extensions/Copy/Stream.hs b/Database/Beam/Postgres/Extensions/Copy/Stream.hs new file mode 100644 --- /dev/null +++ b/Database/Beam/Postgres/Extensions/Copy/Stream.hs @@ -0,0 +1,175 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeFamilies #-} + +module Database.Beam.Postgres.Extensions.Copy.Stream + ( PgCopyToStreamSyntax (..), + PgCopyFromStreamSyntax (..), + + -- * COPY options + PgCopyToStreamOptions, + PgCopyFromStreamOptions, + + -- ** text + copyToTextStream, + copyToTextStreamWith, + copyFromTextStream, + copyFromTextStreamWith, + + -- ** CSV + copyToCSVStream, + copyToCSVStreamWith, + copyFromCSVStream, + copyFromCSVStreamWith, + ) +where + +import Database.Beam.Backend.SQL.BeamExtensions + ( IsSqlCopyFromStreamSyntax (..), + IsSqlCopyToStreamSyntax (..), + ) +import Database.Beam.Postgres.Extensions.Copy.File + ( PgCSVCopyFromOptions (..), + PgCSVCopyToOptions (..), + PgCopyFromSourceSyntax (..), + PgCopyToSourceSyntax (..), + PgTextCopyFromOptions (..), + PgTextCopyToOptions (..), + csvCopyFromFields, + csvCopyToFields, + defaultPgCSVCopyFromOptions, + defaultPgCSVCopyToOptions, + defaultPgTextCopyFromOptions, + defaultPgTextCopyToOptions, + emitOptionList, + textCopyFromFields, + textCopyToFields, + ) +import Database.Beam.Postgres.Syntax (PgSyntax, emit) + +-- | PostgreSQL streaming @COPY ... TO STDOUT@ statement syntax. +-- +-- @since 0.6.1.0 +newtype PgCopyToStreamSyntax = PgCopyToStreamSyntax {fromPgCopyToStream :: PgSyntax} + +-- | PostgreSQL streaming @COPY ... FROM STDIN@ statement syntax. +-- +-- @since 0.6.1.0 +newtype PgCopyFromStreamSyntax = PgCopyFromStreamSyntax {fromPgCopyFromStream :: PgSyntax} + +instance IsSqlCopyToStreamSyntax PgCopyToStreamSyntax where + type SqlCopyToStreamSourceSyntax PgCopyToStreamSyntax = PgCopyToSourceSyntax + type SqlCopyToStreamParams PgCopyToStreamSyntax = PgCopyToStreamOptions + + copyToStreamStmt (PgCopyToSourceSyntax source) options = + PgCopyToStreamSyntax $ + mconcat + [ emit "COPY ", + source, + emit " TO STDOUT", + emitPgCopyToStreamOptions options + ] + +instance IsSqlCopyFromStreamSyntax PgCopyFromStreamSyntax where + type SqlCopyFromStreamSourceSyntax PgCopyFromStreamSyntax = PgCopyFromSourceSyntax + type SqlCopyFromStreamParams PgCopyFromStreamSyntax = PgCopyFromStreamOptions + + copyFromStreamStmt (PgCopyFromSourceSyntax source) options = + PgCopyFromStreamSyntax $ + mconcat + [ emit "COPY ", + source, + emit " FROM STDIN", + emitPgCopyFromStreamOptions options + ] + +-- | Options for PostgreSQL's streaming @COPY ... TO STDOUT@ statement, +-- tagged by the output format. +-- +-- See +-- for the full set of options each format supports. +-- +-- @since 0.6.1.0 +data PgCopyToStreamOptions + = PgCopyToStreamText !PgTextCopyToOptions + | PgCopyToStreamCSV !PgCSVCopyToOptions + deriving (Eq, Show) + +-- | Stream out using the @text@ format with default options. +-- Use 'copyToTextStreamWith' to override. +-- +-- @since 0.6.1.0 +copyToTextStream :: PgCopyToStreamOptions +copyToTextStream = PgCopyToStreamText defaultPgTextCopyToOptions + +-- | Stream out using the @text@ format with the given options. +-- +-- @since 0.6.1.0 +copyToTextStreamWith :: PgTextCopyToOptions -> PgCopyToStreamOptions +copyToTextStreamWith = PgCopyToStreamText + +-- | Stream out using the @csv@ format with default options. +-- Use 'copyToCSVStreamWith' to override. +-- +-- @since 0.6.1.0 +copyToCSVStream :: PgCopyToStreamOptions +copyToCSVStream = PgCopyToStreamCSV defaultPgCSVCopyToOptions + +-- | Stream out using the @csv@ format with the given options. +-- +-- @since 0.6.1.0 +copyToCSVStreamWith :: PgCSVCopyToOptions -> PgCopyToStreamOptions +copyToCSVStreamWith = PgCopyToStreamCSV + +-- | Options for PostgreSQL's streaming @COPY ... FROM STDIN@ statement, +-- tagged by the input format. +-- +-- @since 0.6.1.0 +data PgCopyFromStreamOptions + = PgCopyFromStreamText !PgTextCopyFromOptions + | PgCopyFromStreamCSV !PgCSVCopyFromOptions + deriving (Eq, Show) + +-- | Stream in using the @text@ format with default options. +-- Use 'copyFromTextStreamWith' to override. +-- +-- @since 0.6.1.0 +copyFromTextStream :: PgCopyFromStreamOptions +copyFromTextStream = PgCopyFromStreamText defaultPgTextCopyFromOptions + +-- | Stream in using the @text@ format with the given options. +-- +-- @since 0.6.1.0 +copyFromTextStreamWith :: PgTextCopyFromOptions -> PgCopyFromStreamOptions +copyFromTextStreamWith = PgCopyFromStreamText + +-- | Stream in using the @csv@ format with default options. +-- Use 'copyFromCSVStreamWith' to override. +-- +-- @since 0.6.1.0 +copyFromCSVStream :: PgCopyFromStreamOptions +copyFromCSVStream = PgCopyFromStreamCSV defaultPgCSVCopyFromOptions + +-- | Stream in using the @csv@ format with the given options. +-- +-- @since 0.6.1.0 +copyFromCSVStreamWith :: PgCSVCopyFromOptions -> PgCopyFromStreamOptions +copyFromCSVStreamWith = PgCopyFromStreamCSV + +-- | Render the @WITH (FORMAT ..., ...)@ tail of a streaming +-- @COPY ... TO STDOUT@ statement. +-- +-- @since 0.6.1.0 +emitPgCopyToStreamOptions :: PgCopyToStreamOptions -> PgSyntax +emitPgCopyToStreamOptions (PgCopyToStreamText o) = emitOptionList (emit "text") (textCopyToFields o) +emitPgCopyToStreamOptions (PgCopyToStreamCSV o) = emitOptionList (emit "csv") (csvCopyToFields o) + +-- | Render the @WITH (FORMAT ..., ...)@ tail of a streaming +-- @COPY ... FROM STDIN@ statement. +-- +-- @since 0.6.1.0 +emitPgCopyFromStreamOptions :: PgCopyFromStreamOptions -> PgSyntax +emitPgCopyFromStreamOptions (PgCopyFromStreamText o) = emitOptionList (emit "text") (textCopyFromFields o) +emitPgCopyFromStreamOptions (PgCopyFromStreamCSV o) = emitOptionList (emit "csv") (csvCopyFromFields o) diff --git a/Database/Beam/Postgres/Syntax.hs b/Database/Beam/Postgres/Syntax.hs --- a/Database/Beam/Postgres/Syntax.hs +++ b/Database/Beam/Postgres/Syntax.hs @@ -22,7 +22,7 @@ , emit, emitBuilder, escapeString , escapeBytea, escapeIdentifier , pgParens - + , pgStringLit, pgCharLit, pgBoolLit , nextSyntaxStep , PgCommandSyntax(..), PgCommandType(..) @@ -131,6 +131,8 @@ import qualified Database.PostgreSQL.Simple.Types as Pg (Oid(..), Binary(..), Null(..)) import qualified Database.PostgreSQL.Simple.Time as Pg (Date, LocalTimestamp, UTCTimestamp) import qualified Database.PostgreSQL.Simple.HStore as Pg (HStoreList, HStoreMap, HStoreBuilder) +import Data.Text (Text) +import qualified Data.ByteString.Char8 as BS8 data PostgresInaccessible @@ -1389,6 +1391,27 @@ pgParens :: PgSyntax -> PgSyntax pgParens a = emit "(" <> a <> emit ")" + +-- | Render a 'Text' as a single-quoted SQL string literal, with proper +-- Postgres escaping. The surrounding quotes are added here; 'escapeString' +-- only handles the escaping of the contents. +-- +-- @since 0.6.1.0 +pgStringLit :: Text -> PgSyntax +pgStringLit t = emit "'" <> escapeString (TE.encodeUtf8 t) <> emit "'" + +-- | Render a 'Char' as a single-quoted SQL string literal. +-- +-- @since 0.6.1.0 +pgCharLit :: Char -> PgSyntax +pgCharLit c = emit "'" <> escapeString (BS8.singleton c) <> emit "'" + +-- | Render a boolean as the literal @TRUE@ or @FALSE@. +-- +-- @since 0.6.1.0 +pgBoolLit :: Bool -> PgSyntax +pgBoolLit True = emit "TRUE" +pgBoolLit False = emit "FALSE" pgTableOp :: ByteString -> PgSelectTableSyntax -> PgSelectTableSyntax -> PgSelectTableSyntax diff --git a/Database/Beam/Postgres/Types.hs b/Database/Beam/Postgres/Types.hs --- a/Database/Beam/Postgres/Types.hs +++ b/Database/Beam/Postgres/Types.hs @@ -18,8 +18,18 @@ import Database.Beam import Database.Beam.Backend import Database.Beam.Backend.Internal.Compat +import Database.Beam.Backend.SQL.BeamExtensions + ( BeamSqlBackendCopyFromStreamSyntax + , BeamSqlBackendCopyFromSyntax + , BeamSqlBackendCopyToStreamSyntax + , BeamSqlBackendCopyToSyntax + ) import Database.Beam.Migrate.Generics import Database.Beam.Migrate.SQL (BeamMigrateOnlySqlBackend) +import Database.Beam.Postgres.Extensions.Copy.File + ( PgCopyFromSyntax, PgCopyToSyntax ) +import Database.Beam.Postgres.Extensions.Copy.Stream + ( PgCopyFromStreamSyntax, PgCopyToStreamSyntax ) import Database.Beam.Postgres.Syntax import Database.Beam.Query.SQL92 @@ -165,6 +175,12 @@ instance BeamSqlBackend Postgres instance BeamMigrateOnlySqlBackend Postgres type instance BeamSqlBackendSyntax Postgres = PgCommandSyntax + +type instance BeamSqlBackendCopyToSyntax Postgres = PgCopyToSyntax +type instance BeamSqlBackendCopyFromSyntax Postgres = PgCopyFromSyntax + +type instance BeamSqlBackendCopyToStreamSyntax Postgres = PgCopyToStreamSyntax +type instance BeamSqlBackendCopyFromStreamSyntax Postgres = PgCopyFromStreamSyntax instance BeamSqlBackendIsString Postgres String instance BeamSqlBackendIsString Postgres Text diff --git a/beam-postgres.cabal b/beam-postgres.cabal --- a/beam-postgres.cabal +++ b/beam-postgres.cabal @@ -1,5 +1,5 @@ name: beam-postgres -version: 0.6.0.0 +version: 0.6.1.0 synopsis: Connection layer between beam and postgres description: Beam driver for , an advanced open-source RDBMS homepage: https://haskell-beam.github.io/beam/user-guide/backends/beam-postgres @@ -30,12 +30,14 @@ other-modules: Database.Beam.Postgres.Connection Database.Beam.Postgres.Debug Database.Beam.Postgres.Extensions + Database.Beam.Postgres.Extensions.Copy.File + Database.Beam.Postgres.Extensions.Copy.Stream Database.Beam.Postgres.Extensions.Internal Database.Beam.Postgres.PgSpecific Database.Beam.Postgres.Types build-depends: base >=4.11 && <5.0, - beam-core >=0.11 && <0.12, + beam-core >=0.11.1 && <0.12, beam-migrate >=0.6 && <0.7, postgresql-libpq >=0.8 && <0.12, @@ -48,7 +50,7 @@ hashable >=1.1 && <1.6, lifted-base >=0.2 && <0.3, free >=4.12 && <5.3, - time >=1.6 && <1.15, + time >=1.6 && <1.17, monad-control >=1.0 && <1.1, mtl >=2.1 && <2.4, conduit >=1.2 && <1.4, @@ -80,8 +82,10 @@ hs-source-dirs: test main-is: Main.hs other-modules: Database.Beam.Postgres.Test, + Database.Beam.Postgres.Test.Copy, Database.Beam.Postgres.Test.Marshal, Database.Beam.Postgres.Test.Select, + Database.Beam.Postgres.Test.Select.PgNubBy, Database.Beam.Postgres.Test.DataTypes, Database.Beam.Postgres.Test.Migrate, Database.Beam.Postgres.Test.TempTable, @@ -99,6 +103,7 @@ tasty, text, testcontainers, + time, uuid, vector default-language: Haskell2010 diff --git a/test/Database/Beam/Postgres/Test/Copy.hs b/test/Database/Beam/Postgres/Test/Copy.hs new file mode 100644 --- /dev/null +++ b/test/Database/Beam/Postgres/Test/Copy.hs @@ -0,0 +1,444 @@ +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TypeFamilies #-} + +module Database.Beam.Postgres.Test.Copy (tests) where + +import Control.Monad (void) +import Control.Monad.IO.Class (liftIO) +import Data.ByteString (ByteString) +import qualified Data.ByteString as BS +import Data.IORef (modifyIORef', newIORef, readIORef, writeIORef) +import Data.Int (Int32) +import Data.List (sort) +import Data.Text (Text) +import qualified Data.Text as T +import Database.Beam +import Database.Beam.Backend.SQL.BeamExtensions + ( MonadBeamCopyFrom (..), + MonadBeamCopyFromStream (..), + MonadBeamCopyTo (..), + MonadBeamCopyToStream (..), + copySelectTo, + copySelectToStream, + copyTableFrom, + copyTableFromStream, + copyTableTo, + copyTableToStream, + ) +import Database.Beam.Postgres +import Database.Beam.Postgres.Test (withTestPostgres) +import Database.PostgreSQL.Simple (Connection, execute_) +import Hedgehog ((===)) +import qualified Hedgehog +import qualified Hedgehog.Gen as Gen +import qualified Hedgehog.Range as Range +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (Assertion, assertBool, assertEqual, testCase) + +tests :: IO ByteString -> TestTree +tests getConn = + testGroup + "COPY statements" + [ testGroup + "COPY ... TO / ... FROM round-trip" + [ testRoundtripCsvAllColumns getConn, + testRoundtripCsvMultiColumnProjection getConn, + testRoundtripCsvSingleColumnProjection getConn, + testRoundtripCsvFromSelect getConn + ], + testGroup + "options" + [ testCustomDelimiter getConn, + testNoHeader getConn, + testTextFormat getConn + ], + testGroup + "streaming COPY" + [ testStreamRoundtripCsv getConn, + testStreamFromSelect getConn + ], + testGroup + "property-based round-trip" + [ testPropCsvRoundtrip getConn, + testPropTextRoundtrip getConn + ] + ] + +-- | The full table is exported to a CSV, the table is truncated, and the CSV +-- is re-imported. The final rows must match the original ones. +testRoundtripCsvAllColumns :: IO ByteString -> TestTree +testRoundtripCsvAllColumns getConn = + testCopy getConn "copy_all_columns_csv" "/tmp/copy_all_columns.csv" $ \conn path -> do + seedWidgets conn widgetData + runBeamPostgres conn $ do + runCopyTo $ copyTableTo (_dbWidgets testDb) id (copyToCSV path) + truncateWidgets conn + runBeamPostgres conn $ do + runCopyFrom $ copyTableFrom (_dbWidgets testDb) id (copyFromCSV path) + rows <- queryAllWidgets conn + assertEqual "round-tripped rows match seed" widgetData rows + +-- | Project a subset of columns; the unprojected column gets its DEFAULT on +-- import. +testRoundtripCsvMultiColumnProjection :: IO ByteString -> TestTree +testRoundtripCsvMultiColumnProjection getConn = + testCopy getConn "copy_multi_column_csv" "/tmp/copy_multi_column.csv" $ \conn path -> do + seedWidgets conn widgetData + runBeamPostgres conn $ do + runCopyTo $ + copyTableTo + (_dbWidgets testDb) + (\w -> (_widgetId w, _widgetName w)) + (copyToCSV path) + truncateWidgets conn + runBeamPostgres conn $ do + runCopyFrom $ + copyTableFrom + (_dbWidgets testDb) + (\w -> (_widgetId w, _widgetName w)) + (copyFromCSV path) + rows <- queryAllWidgets conn + -- 'price' was not in the projection, so it gets its DEFAULT (0). + let expected = [w {_widgetPrice = 0} | w <- widgetData] + assertEqual "id+name round-tripped, price defaulted" expected rows + +-- | Single-column projection. +testRoundtripCsvSingleColumnProjection :: IO ByteString -> TestTree +testRoundtripCsvSingleColumnProjection getConn = + testCopy getConn "copy_single_column_csv" "/tmp/copy_single_column.csv" $ \conn path -> do + seedWidgets conn widgetData + runBeamPostgres conn $ do + runCopyTo $ copyTableTo (_dbWidgets testDb) _widgetId (copyToCSV path) + truncateWidgets conn + runBeamPostgres conn $ do + runCopyFrom $ copyTableFrom (_dbWidgets testDb) _widgetId (copyFromCSV path) + rows <- queryAllWidgets conn + let expected = [w {_widgetName = "", _widgetPrice = 0} | w <- widgetData] + assertEqual "id-only round-tripped, name+price defaulted" expected rows + +-- | Use copySelectTo with a filtered query, then COPY FROM the result back +-- into the table. +testRoundtripCsvFromSelect :: IO ByteString -> TestTree +testRoundtripCsvFromSelect getConn = + testCopy getConn "copy_select_csv" "/tmp/copy_select.csv" $ \conn path -> do + seedWidgets conn widgetData + runBeamPostgres conn $ do + runCopyTo $ + copySelectTo + ( select $ do + w <- all_ (_dbWidgets testDb) + guard_ (_widgetPrice w >. val_ 4.0) + pure w + ) + (copyToCSV path) + truncateWidgets conn + runBeamPostgres conn $ do + runCopyFrom $ copyTableFrom (_dbWidgets testDb) id (copyFromCSV path) + rows <- queryAllWidgets conn + -- Only Widget (9.99) and Sprocket (4.5) have price > 4.0; Cog is excluded. + let expected = sort [w | w <- widgetData, _widgetPrice w > 4.0] + assertEqual "select-filtered round-trip" expected rows + +-- | Use a custom delimiter on both sides of the round-trip. +testCustomDelimiter :: IO ByteString -> TestTree +testCustomDelimiter getConn = + testCopy getConn "copy_custom_delim_csv" "/tmp/copy_custom_delim.csv" $ \conn path -> do + seedWidgets conn widgetData + let toOpts = copyToCSVWith path defaultPgCSVCopyToOptions {pgCsvCopyToDelimiter = Just '|'} + fromOpts = copyFromCSVWith path defaultPgCSVCopyFromOptions {pgCsvCopyFromDelimiter = Just '|'} + runBeamPostgres conn $ do + runCopyTo $ copyTableTo (_dbWidgets testDb) id toOpts + truncateWidgets conn + runBeamPostgres conn $ do + runCopyFrom $ copyTableFrom (_dbWidgets testDb) id fromOpts + rows <- queryAllWidgets conn + assertEqual "round-trip with '|' delimiter" widgetData rows + +-- | Toggle HEADER off on TO and on on FROM; round-trip should still work +-- (as long as both sides agree). +testNoHeader :: IO ByteString -> TestTree +testNoHeader getConn = + testCopy getConn "copy_no_header_csv" "/tmp/copy_no_header.csv" $ \conn path -> do + seedWidgets conn widgetData + let toOpts = copyToCSVWith path defaultPgCSVCopyToOptions {pgCsvCopyToHeader = Just False} + fromOpts = copyFromCSVWith path defaultPgCSVCopyFromOptions {pgCsvCopyFromHeader = Just False} + runBeamPostgres conn $ do + runCopyTo $ copyTableTo (_dbWidgets testDb) id toOpts + truncateWidgets conn + runBeamPostgres conn $ do + runCopyFrom $ copyTableFrom (_dbWidgets testDb) id fromOpts + rows <- queryAllWidgets conn + assertEqual "round-trip without header line" widgetData rows + +-- | The 'text' format with custom delimiter, round-tripped. +testTextFormat :: IO ByteString -> TestTree +testTextFormat getConn = + testCopy getConn "copy_text_format" "/tmp/copy_text.txt" $ \conn path -> do + seedWidgets conn widgetData + let toOpts = copyToTextWith path defaultPgTextCopyToOptions {pgTextCopyToDelimiter = Just '\t'} + fromOpts = copyFromTextWith path defaultPgTextCopyFromOptions {pgTextCopyFromDelimiter = Just '\t'} + runBeamPostgres conn $ do + runCopyTo $ copyTableTo (_dbWidgets testDb) id toOpts + truncateWidgets conn + runBeamPostgres conn $ do + runCopyFrom $ copyTableFrom (_dbWidgets testDb) id fromOpts + rows <- queryAllWidgets conn + assertEqual "round-trip via text format" widgetData rows + +-- | Round-trip widgets through streaming @COPY ... TO STDOUT@ then +-- @COPY ... FROM STDIN@. The bytes pass through the client connection only; +-- no server-side file is touched. +testStreamRoundtripCsv :: IO ByteString -> TestTree +testStreamRoundtripCsv getConn = + testCopy getConn "stream_roundtrip_csv" "" $ \conn _ -> do + seedWidgets conn widgetData + -- Drain the stream into an IORef. + chunksRef <- newIORef [] + runBeamPostgres conn $ + runCopyToStream + (copyTableToStream (_dbWidgets testDb) id copyToCSVStream) + (\chunk -> modifyIORef' chunksRef (chunk :)) + payload <- BS.concat . reverse <$> readIORef chunksRef + assertBool "stream produced non-empty payload" (not (BS.null payload)) + -- Replay the captured bytes back into the table. + truncateWidgets conn + sourceRef <- newIORef (Just payload) + runBeamPostgres conn $ + runCopyFromStream + (copyTableFromStream (_dbWidgets testDb) id copyFromCSVStream) + ( do + mchunk <- readIORef sourceRef + writeIORef sourceRef Nothing + pure mchunk + ) + rows <- queryAllWidgets conn + assertEqual "round-tripped rows match seed" widgetData rows + +-- | @COPY (SELECT ...) TO STDOUT@ via 'copySelectToStream'. +testStreamFromSelect :: IO ByteString -> TestTree +testStreamFromSelect getConn = + testCopy getConn "stream_from_select" "" $ \conn _ -> do + seedWidgets conn widgetData + chunksRef <- newIORef [] + runBeamPostgres conn $ + runCopyToStream + ( copySelectToStream + ( select $ do + w <- all_ (_dbWidgets testDb) + guard_ (_widgetPrice w >. val_ 4.0) + pure w + ) + copyToCSVStream + ) + (\chunk -> modifyIORef' chunksRef (chunk :)) + payload <- BS.concat . reverse <$> readIORef chunksRef + -- The payload should mention the two widgets with price > 4.0 and not + -- the one priced 1.25. + assertBool "Widget present" ("Widget" `BS.isInfixOf` payload) + assertBool "Sprocket present" ("Sprocket" `BS.isInfixOf` payload) + assertBool "Cog absent" (not ("Cog" `BS.isInfixOf` payload)) + +testPropCsvRoundtrip :: IO ByteString -> TestTree +testPropCsvRoundtrip getConn = + testCopy getConn "prop_csv_roundtrip" "/tmp/prop_csv.csv" $ \conn path -> do + passes <- + Hedgehog.check . Hedgehog.property $ do + toOpts <- Hedgehog.forAll genCsvToOpts + let fromOpts = matchedCsvFromOpts toOpts + rows <- + liftIO $ + roundtripWidgets conn $ \c -> do + runBeamPostgres c $ + runCopyTo $ + copyTableTo (_dbWidgets testDb) id (copyToCSVWith path toOpts) + truncateWidgets c + runBeamPostgres c $ + runCopyFrom $ + copyTableFrom (_dbWidgets testDb) id (copyFromCSVWith path fromOpts) + rows === widgetData + assertBool "Hedgehog property failed" passes + where + genCsvToOpts = do + delim <- Gen.maybe genDelimiterChar + hdr <- Gen.maybe Gen.bool + quote <- Gen.maybe genQuoteChar + escape <- Gen.maybe genEscapeChar + nullStr <- Gen.maybe genNullStr + pure + defaultPgCSVCopyToOptions + { pgCsvCopyToDelimiter = delim, + pgCsvCopyToHeader = hdr, + pgCsvCopyToQuote = quote, + pgCsvCopyToEscape = escape, + pgCsvCopyToNullStr = nullStr + } + + matchedCsvFromOpts :: PgCSVCopyToOptions -> PgCSVCopyFromOptions + matchedCsvFromOpts o = + defaultPgCSVCopyFromOptions + { pgCsvCopyFromDelimiter = pgCsvCopyToDelimiter o, + pgCsvCopyFromHeader = pgCsvCopyToHeader o, + pgCsvCopyFromQuote = pgCsvCopyToQuote o, + pgCsvCopyFromEscape = pgCsvCopyToEscape o, + pgCsvCopyFromNullStr = pgCsvCopyToNullStr o + } + +testPropTextRoundtrip :: IO ByteString -> TestTree +testPropTextRoundtrip getConn = + testCopy getConn "prop_text_roundtrip" "/tmp/prop_text.txt" $ \conn path -> do + passes <- + Hedgehog.check . Hedgehog.property $ do + toOpts <- Hedgehog.forAll genTextToOpts + let fromOpts = matchedTextFromOpts toOpts + rows <- + liftIO $ + roundtripWidgets conn $ \c -> do + runBeamPostgres c $ + runCopyTo $ + copyTableTo (_dbWidgets testDb) id (copyToTextWith path toOpts) + truncateWidgets c + runBeamPostgres c $ + runCopyFrom $ + copyTableFrom (_dbWidgets testDb) id (copyFromTextWith path fromOpts) + rows === widgetData + assertBool "Hedgehog property failed" passes + where + genTextToOpts = do + delim <- Gen.maybe genDelimiterChar + hdr <- Gen.maybe Gen.bool + -- Avoid generating a NULL string that contains the chosen delimiter. + nullStr <- Gen.maybe (Gen.filter (not . T.any (== ',')) genNullStr) + pure + defaultPgTextCopyToOptions + { pgTextCopyToDelimiter = delim, + pgTextCopyToHeader = hdr, + pgTextCopyToNullStr = nullStr + } + + -- \| Mirror of 'matchedCsvFromOpts' for the @text@ format. 'pgTextCopyFromFreeze' + -- is intentionally left at @Nothing@: PostgreSQL only accepts @FREEZE@ when + -- the target table has not been touched in the current (sub)transaction, but + -- our property test truncates before the FROM, which violates that + -- precondition. + matchedTextFromOpts :: PgTextCopyToOptions -> PgTextCopyFromOptions + matchedTextFromOpts o = + defaultPgTextCopyFromOptions + { pgTextCopyFromDelimiter = pgTextCopyToDelimiter o, + pgTextCopyFromHeader = pgTextCopyToHeader o, + pgTextCopyFromNullStr = pgTextCopyToNullStr o + } + +-- Single-byte delimiters that never appear in 'widgetData'. Avoiding @.@ +-- matters for the @text@ format because the data contains floating-point +-- prices like @9.99@. +genDelimiterChar :: Hedgehog.Gen Char +genDelimiterChar = Gen.element [',', ';', '|', '\t'] + +genQuoteChar :: Hedgehog.Gen Char +genQuoteChar = Gen.element ['"', '\''] + +genEscapeChar :: Hedgehog.Gen Char +genEscapeChar = Gen.element ['"', '\'', '\\'] + +genNullStr :: Hedgehog.Gen Text +genNullStr = Gen.element ["NULL", "~~~"] + +roundtripWidgets :: Connection -> (Connection -> IO ()) -> IO [Widget] +roundtripWidgets conn action = do + truncateWidgets conn + seedWidgets conn widgetData + action conn + queryAllWidgets conn + +testCopy :: + IO ByteString -> + String -> + -- | Server-side file path inside the Postgres container. + FilePath -> + (Connection -> FilePath -> Assertion) -> + TestTree +testCopy getConn name path action = testCase name $ + withTestPostgres name getConn $ \conn -> do + createWidgetsTable conn + action conn path + +queryAllWidgets :: Connection -> IO [Widget] +queryAllWidgets conn = + runBeamPostgres conn $ + runSelectReturningList $ + select $ + orderBy_ (asc_ . _widgetId) $ + all_ (_dbWidgets testDb) + +createWidgetsTable :: Connection -> IO () +createWidgetsTable conn = + void $ + execute_ + conn + "CREATE TABLE widgets (\ + \ id INTEGER PRIMARY KEY,\ + \ name TEXT NOT NULL DEFAULT '',\ + \ price DOUBLE PRECISION NOT NULL DEFAULT 0\ + \)" + +truncateWidgets :: Connection -> IO () +truncateWidgets conn = void $ execute_ conn "TRUNCATE TABLE widgets" + +seedWidgets :: Connection -> [Widget] -> IO () +seedWidgets conn widgets = + runBeamPostgres conn $ + runInsert $ + insert (_dbWidgets testDb) (insertValues widgets) + +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 + +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 Postgres TestDB +testDb = + defaultDbSettings + `withDbModification` dbModification + { _dbWidgets = + modifyTableFields + tableModification + { _widgetId = "id", + _widgetName = "name", + _widgetPrice = "price" + } + } diff --git a/test/Database/Beam/Postgres/Test/Select/PgNubBy.hs b/test/Database/Beam/Postgres/Test/Select/PgNubBy.hs new file mode 100644 --- /dev/null +++ b/test/Database/Beam/Postgres/Test/Select/PgNubBy.hs @@ -0,0 +1,119 @@ +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE StandaloneDeriving #-} + +module Database.Beam.Postgres.Test.Select.PgNubBy (tests) where + +import Control.Monad (void) +import Data.ByteString (ByteString) +import Data.Int (Int32) +import Data.Text (Text) +import Data.Time.Calendar (Day, fromGregorian) +import Database.Beam +import Database.Beam.Migrate (defaultMigratableDbSettings) +import Database.Beam.Migrate.Simple (autoMigrate) +import Database.Beam.Postgres +import Database.Beam.Postgres.Migrate (migrationBackend) +import Database.Beam.Postgres.Test (withTestPostgres) +import Test.Tasty +import Test.Tasty.HUnit + +tests :: IO ByteString -> TestTree +tests getConn = + testGroup + "pgNubBy_ / nub_ with window functions (issue #746)" + [ testPgNubByWithLead getConn, + testNubWithLead getConn + ] + +-- Reproducer for issue #746 +testPgNubByWithLead :: IO ByteString -> TestTree +testPgNubByWithLead getConn = testCase "pgNubBy_ feeding lead1_" $ + withTestPostgres "issue_746_pg_nub_by" getConn $ \conn -> do + setupDb conn + results <- + runBeamPostgres conn $ + runSelectReturningList $ + select $ + withWindow_ + (\vf -> frame_ noPartition_ (orderPartitionBy_ (asc_ vf)) noBounds_) + (\vf w -> (vf, lead1_ vf `over_` w)) + (pgNubBy_ id (validFrom <$> all_ (persons db))) + + let expected = + [ (day 2025 1 1, Just (day 2025 1 2)), + (day 2025 1 2, Just (day 2025 1 3)), + (day 2025 1 3, Nothing) + ] + assertEqual "lead1_ over pgNubBy_ should pair each distinct date with the next" expected results + +-- Reproducer for issue #746 with @nub_@ +testNubWithLead :: IO ByteString -> TestTree +testNubWithLead getConn = testCase "nub_ feeding lead1_" $ + withTestPostgres "issue_746_nub" getConn $ \conn -> do + setupDb conn + results <- + runBeamPostgres conn $ + runSelectReturningList $ + select $ + withWindow_ + (\vf -> frame_ noPartition_ (orderPartitionBy_ (asc_ vf)) noBounds_) + (\vf w -> (vf, lead1_ vf `over_` w)) + (nub_ (validFrom <$> all_ (persons db))) + + let expected = + [ (day 2025 1 1, Just (day 2025 1 2)), + (day 2025 1 2, Just (day 2025 1 3)), + (day 2025 1 3, Nothing) + ] + assertEqual "lead1_ over nub_ should pair each distinct date with the next" expected results + +data PersonT f = Person + { name :: C f Text, + validFrom :: C f Day, + idx :: C f Int32 + } + deriving (Generic) + +type Person = PersonT Identity + +deriving instance Show Person + +deriving instance Eq Person + +instance Beamable PersonT + +instance Table PersonT where + data PrimaryKey PersonT f = PersonKey (C f Text) + deriving stock (Generic) + deriving anyclass (Beamable) + + primaryKey Person {name} = PersonKey name + +newtype Db f = Db + { persons :: f (TableEntity PersonT) + } + deriving (Generic) + +instance Database Postgres Db + +db :: DatabaseSettings Postgres Db +db = defaultDbSettings + +day :: Integer -> Int -> Int -> Day +day = fromGregorian + +seedRows :: [Person] +seedRows = + [ Person "A" (day 2025 1 1) 1, + Person "B" (day 2025 1 1) 1, + Person "C" (day 2025 1 2) 2, + Person "D" (day 2025 1 2) 2, + Person "E" (day 2025 1 3) 3, + Person "F" (day 2025 1 3) 3 + ] + +setupDb :: Connection -> IO () +setupDb conn = runBeamPostgres conn $ do + void $ autoMigrate migrationBackend (defaultMigratableDbSettings @Postgres @Db) + runInsert $ insert (persons db) $ insertValues seedRows diff --git a/test/Main.hs b/test/Main.hs --- a/test/Main.hs +++ b/test/Main.hs @@ -7,10 +7,12 @@ import Test.Tasty import qualified TestContainers.Tasty as TC -import qualified Database.Beam.Postgres.Test.Select as Select -import qualified Database.Beam.Postgres.Test.Marshal as Marshal +import qualified Database.Beam.Postgres.Test.Copy as Copy import qualified Database.Beam.Postgres.Test.DataTypes as DataType +import qualified Database.Beam.Postgres.Test.Marshal as Marshal import qualified Database.Beam.Postgres.Test.Migrate as Migrate +import qualified Database.Beam.Postgres.Test.Select as Select +import qualified Database.Beam.Postgres.Test.Select.PgNubBy as Select.PgNubBy import qualified Database.Beam.Postgres.Test.TempTable as TempTable import qualified Database.Beam.Postgres.Test.Windowing as Windowing import Database.PostgreSQL.Simple ( ConnectInfo(..), defaultConnectInfo ) @@ -23,10 +25,12 @@ testGroup "beam-postgres tests" [ Marshal.tests getConnStr , Select.tests getConnStr + , Select.PgNubBy.tests getConnStr , DataType.tests getConnStr , Migrate.tests getConnStr , TempTable.tests getConnStr , Windowing.tests getConnStr + , Copy.tests getConnStr ]