diff --git a/ChangeLog.md b/ChangeLog.md --- a/ChangeLog.md +++ b/ChangeLog.md @@ -1,3 +1,18 @@ +# 0.5.6.0 + +## Added features + +* Support for temporary tables. + +* `getDbConstraintsForSchemas` now discovers foreign key constraints + via `pg_constraint`, including `ON DELETE` / `ON UPDATE` actions. + +## Performance improvements + +* By minimizing redundant work in the hot loop, the performance of `beam-postgres` + when fetching data from a database has been improved by 30%. The performance + of `beam-postgres` is now within 10% of raw queries via `postgresql-simple` (#797). + # 0.5.5.0 ## Added features diff --git a/Database/Beam/Postgres.hs b/Database/Beam/Postgres.hs --- a/Database/Beam/Postgres.hs +++ b/Database/Beam/Postgres.hs @@ -43,6 +43,7 @@ , smallserial, serial, bigserial , module Database.Beam.Postgres.PgSpecific + , module Database.Beam.Postgres.TempTable -- ** Postgres extension support , PgExtensionEntity, IsPgExtension(..) @@ -83,5 +84,6 @@ , pgCreateExtension, pgDropExtension , getPgExtension ) import Database.Beam.Postgres.Debug +import Database.Beam.Postgres.TempTable import qualified Database.PostgreSQL.Simple as Pg diff --git a/Database/Beam/Postgres/Conduit.hs b/Database/Beam/Postgres/Conduit.hs --- a/Database/Beam/Postgres/Conduit.hs +++ b/Database/Beam/Postgres/Conduit.hs @@ -48,6 +48,7 @@ import qualified Conduit as C import Data.Int (Int64) import Data.Maybe (fromMaybe) +import qualified Data.Vector as V import qualified Control.Monad.Fail as Fail @@ -217,7 +218,7 @@ liftIO (Pg.resultStatus row) >>= \case Pg.SingleTuple -> - do fields' <- liftIO (maybe (getFields row) pure fields) + do fields' <- liftIO (maybe (getFields row) (pure . V.fromList) fields) parsedRow <- liftIO $ bracket (putMVar connectionHandle conn') (\_ -> takeMVar connectionHandle) @@ -226,7 +227,7 @@ Left err -> liftIO (bailEarly conn' row ("Could not read row: " <> show err)) Right parsedRow' -> do C.yield parsedRow' - streamResults conn conn' (Just fields') + streamResults conn conn' (Just (V.toList fields')) Pg.TuplesOk -> liftIO (finishQuery conn') Pg.EmptyQuery -> Fail.fail "No query" Pg.CommandOk -> pure () 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 @@ -5,12 +5,12 @@ {-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE DeriveFunctor #-} -{-# LANGUAGE StandaloneDeriving #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE CPP #-} +{-# LANGUAGE BangPatterns #-} module Database.Beam.Postgres.Connection ( Pg(..), PgF(..) @@ -26,6 +26,9 @@ , postgresUriSyntax ) where import Control.Exception (SomeException(..), throwIO) +import Data.IORef (newIORef, readIORef, writeIORef) +import Data.Vector (Vector) +import qualified Data.Vector as V import Control.Monad.Base (MonadBase(..)) import Control.Monad.Free.Church import Control.Monad.IO.Class @@ -118,29 +121,34 @@ -- * Run row readers -getFields :: Pg.Result -> IO [Pg.Field] +getFields :: Pg.Result -> IO (Vector Pg.Field) getFields res = do Pg.Col colCount <- Pg.nfields res - - let getField col = - Pg.Field res (Pg.Col col) <$> Pg.ftype res (Pg.Col col) - - mapM getField [0..colCount - 1] + V.generateM (fromIntegral colCount) $ \i -> + let col = Pg.Col (fromIntegral i) + in Pg.Field res col <$> Pg.ftype res col runPgRowReader :: - Pg.Connection -> Pg.Row -> Pg.Result -> [Pg.Field] -> FromBackendRowM Postgres a -> IO (Either BeamRowReadError a) + Pg.Connection -> Pg.Row -> Pg.Result -> Vector Pg.Field -> FromBackendRowM Postgres a -> IO (Either BeamRowReadError a) runPgRowReader conn rowIdx res fields (FromBackendRowM readRow) = - Pg.nfields res >>= \(Pg.Col colCount) -> - runF readRow finish step 0 colCount fields + -- 'colCount' and 'fields' are both invariant for the duration of one + -- 'runPgRowReader' call, so we capture them in the closure of 'step' + -- rather than threading them through the free-monad result type. + -- That makes the per-step closure smaller and avoids an O(n) 'length' + -- per row (Vector stores its length, so 'V.length' is O(1)). + runF readRow finish step 0 where + !colCount = fromIntegral (V.length fields) :: CInt - step :: forall x. FromBackendRowF Postgres (CInt -> CInt -> [PgI.Field] -> IO (Either BeamRowReadError x)) - -> CInt -> CInt -> [PgI.Field] -> IO (Either BeamRowReadError x) - step (ParseOneField _) curCol colCount [] = pure (Left (BeamRowReadError (Just (fromIntegral curCol)) (ColumnNotEnoughColumns (fromIntegral colCount)))) - step (ParseOneField _) curCol colCount _ - | curCol >= colCount = pure (Left (BeamRowReadError (Just (fromIntegral curCol)) (ColumnNotEnoughColumns (fromIntegral colCount)))) - step (ParseOneField (next' :: next -> _)) curCol colCount (field:remainingFields) = - do fieldValue <- Pg.getvalue' res rowIdx (Pg.Col curCol) + step :: forall x. FromBackendRowF Postgres (CInt -> IO (Either BeamRowReadError x)) + -> CInt -> IO (Either BeamRowReadError x) + step (ParseOneField _) curCol + | curCol >= colCount = + pure (Left (BeamRowReadError (Just (fromIntegral curCol)) + (ColumnNotEnoughColumns (fromIntegral colCount)))) + step (ParseOneField (next' :: next -> _)) curCol = + do let field = V.unsafeIndex fields (fromIntegral curCol) + fieldValue <- Pg.getvalue' res rowIdx (Pg.Col curCol) res' <- Pg.runConversion (Pg.fromField field fieldValue) conn case res' of Pg.Errors errs -> @@ -162,26 +170,46 @@ Pg.UnexpectedNull {} -> pure ColumnUnexpectedNull in pure (Left (BeamRowReadError (Just (fromIntegral curCol)) err)) - Pg.Ok x -> next' x (curCol + 1) colCount remainingFields + Pg.Ok x -> next' x (curCol + 1) - step (Alt (FromBackendRowM a) (FromBackendRowM b) next) curCol colCount cols = - do aRes <- runF a (\x curCol' colCount' cols' -> pure (Right (next x curCol' colCount' cols'))) step curCol colCount cols + step (Alt (FromBackendRowM a) (FromBackendRowM b) next) curCol = + do aRes <- runF a (\x curCol' -> pure (Right (next x curCol'))) step curCol case aRes of Right next' -> next' Left aErr -> do - bRes <- runF b (\x curCol' colCount' cols' -> pure (Right (next x curCol' colCount' cols'))) step curCol colCount cols + bRes <- runF b (\x curCol' -> pure (Right (next x curCol'))) step curCol case bRes of Right next' -> next' Left {} -> pure (Left aErr) - step (FailParseWith err) _ _ _ = + step (FailParseWith err) _ = pure (Left err) - finish x _ _ _ = pure (Right x) + finish x _ = pure (Right x) withPgDebug :: (String -> IO ()) -> Pg.Connection -> Pg a -> IO (Either BeamRowReadError a) -withPgDebug dbg conn (Pg action) = - let finish x = pure (Right x) +withPgDebug dbg conn (Pg action) = do + -- One-entry cache for the cursor-batch path: 'Pg.Result' is constant + -- within a batch but changes between batches. Caching by 'Pg.Result' + -- equality (a 'ForeignPtr' comparison, ~free) avoids the redundant + -- 'getFields' / 'nfields' / 'ftype' calls within each batch. + -- + -- Default batch size is 256 rows, set by 'postgresql-simple' + -- 'defaultFoldOptions' (FetchQuantity = Automatic, which resolves to + -- 256 in 'Database.PostgreSQL.Simple'). + fieldsCache <- newIORef (Nothing :: Maybe (Vector Pg.Field)) + + let cachedGetFields :: Pg.Result -> IO (Vector Pg.Field) + cachedGetFields res = do + cached <- readIORef fieldsCache + case cached of + Just fs -> pure fs + _ -> do + fs <- getFields res + writeIORef fieldsCache (Just fs) + pure fs + + finish x = pure (Right x) step (PgLiftIO io next) = io >>= next step (PgLiftWithHandle withConn next) = withConn dbg conn >>= next step (PgFetchNext next) = next Nothing @@ -227,8 +255,14 @@ sts <- Pg.resultStatus res case sts of Pg.TuplesOk -> do + -- Hoist per-query metadata out of the per-row loop: the + -- same 'Pg.Result' is used for every row, so 'fields' + -- and 'rowCount' are loop-invariant. + fields <- cachedGetFields res + Pg.Row rowCount <- Pg.ntuples res let Pg process = mkProcess (Pg (liftF (PgFetchNext id))) - runF process (\x _ -> Pg.unsafeFreeResult res >> next x) (stepReturningList res) 0 + runF process (\x _ -> Pg.unsafeFreeResult res >> next x) + (stepReturningList fields rowCount res) 0 _ -> Pg.throwResultError errMsg res sts stepReturningNone :: forall a. PgF (IO (Either BeamRowReadError a)) -> IO (Either BeamRowReadError a) @@ -237,18 +271,18 @@ stepReturningNone (PgFetchNext next) = next Nothing stepReturningNone (PgRunReturning {}) = pure (Left (BeamRowReadError Nothing (ColumnErrorInternal "Nested queries not allowed"))) - stepReturningList :: forall a. Pg.Result -> PgF (CInt -> IO (Either BeamRowReadError a)) -> CInt -> IO (Either BeamRowReadError a) - stepReturningList _ (PgLiftIO action' next) rowIdx = action' >>= \x -> next x rowIdx - stepReturningList res (PgFetchNext next) rowIdx = - do fields <- getFields res - Pg.Row rowCount <- Pg.ntuples res - if rowIdx >= rowCount - then next Nothing rowIdx - else runPgRowReader conn (Pg.Row rowIdx) res fields fromBackendRow >>= \case - Left err -> pure (Left err) - Right r -> next (Just r) (rowIdx + 1) - stepReturningList _ (PgRunReturning {}) _ = pure (Left (BeamRowReadError Nothing (ColumnErrorInternal "Nested queries not allowed"))) - stepReturningList _ (PgLiftWithHandle {}) _ = pure (Left (BeamRowReadError Nothing (ColumnErrorInternal "Nested queries not allowed"))) + stepReturningList :: forall a. Vector Pg.Field -> CInt -> Pg.Result + -> PgF (CInt -> IO (Either BeamRowReadError a)) + -> CInt -> IO (Either BeamRowReadError a) + stepReturningList _ _ _ (PgLiftIO action' next) rowIdx = action' >>= \x -> next x rowIdx + stepReturningList fields rowCount res (PgFetchNext next) rowIdx = + if rowIdx >= rowCount + then next Nothing rowIdx + else runPgRowReader conn (Pg.Row rowIdx) res fields fromBackendRow >>= \case + Left err -> pure (Left err) + Right r -> next (Just r) (rowIdx + 1) + stepReturningList _ _ _ (PgRunReturning {}) _ = pure (Left (BeamRowReadError Nothing (ColumnErrorInternal "Nested queries not allowed"))) + stepReturningList _ _ _ (PgLiftWithHandle {}) _ = pure (Left (BeamRowReadError Nothing (ColumnErrorInternal "Nested queries not allowed"))) finishProcess :: forall a. a -> Maybe PgI.Row -> IO (PgStream a) finishProcess x _ = pure (PgStreamDone (Right x)) @@ -260,12 +294,12 @@ case res of Nothing -> next Nothing Nothing Just (PgI.Row rowIdx res') -> - getFields res' >>= \fields -> + cachedGetFields res' >>= \fields -> runPgRowReader conn rowIdx res' fields fromBackendRow >>= \case Left err -> pure (PgStreamDone (Left err)) Right r -> next (Just r) Nothing stepProcess (PgFetchNext next) (Just (PgI.Row rowIdx res)) = - getFields res >>= \fields -> + cachedGetFields res >>= \fields -> runPgRowReader conn rowIdx res fields fromBackendRow >>= \case Left err -> pure (PgStreamDone (Left err)) Right r -> pure (PgStreamContinue (next (Just r))) @@ -275,7 +309,8 @@ runConsumer :: forall a. PgStream a -> PgI.Row -> IO (PgStream a) runConsumer s@(PgStreamDone {}) _ = pure s runConsumer (PgStreamContinue next) row = next (Just row) - in runF action finish step + + runF action finish step -- * Beam Monad class diff --git a/Database/Beam/Postgres/Migrate.hs b/Database/Beam/Postgres/Migrate.hs --- a/Database/Beam/Postgres/Migrate.hs +++ b/Database/Beam/Postgres/Migrate.hs @@ -31,7 +31,8 @@ ) where import Database.Beam.Backend.SQL -import Database.Beam.Migrate.Actions (defaultActionProvider, defaultSchemaActionProvider, +import Database.Beam.Migrate.Actions (defaultActionProvider, + defaultSchemaActionProvider, createIndexActionProvider, dropIndexActionProvider) import qualified Database.Beam.Migrate.Backend as Tool import qualified Database.Beam.Migrate.Checks as Db @@ -450,6 +451,53 @@ , " AND NOT EXISTS (SELECT 1 FROM unnest(ix.indkey) AS k(attnum) WHERE k.attnum = 0)" , "GROUP BY ns.nspname, c.relname, i.relname, ix.indisunique" ])) + -- Collect foreign key constraints via pg_constraint. + let fkBaseQuery = unlines + [ -- NULL out 'public' since it is the implicit default schema in Postgres + "SELECT NULLIF(ns.nspname, 'public')," + , " c.relname," + -- re-aggregate local and referenced column names in key order (see ORDINALITY below) + , " array_agg(a.attname ORDER BY k.n)," + , " NULLIF(fns.nspname, 'public')," + , " fc.relname," + , " array_agg(fa.attname ORDER BY k.n)," + , " con.confupdtype," + , " con.confdeltype" + , "FROM pg_constraint con" + , "JOIN pg_class c ON c.oid = con.conrelid" + , "JOIN pg_namespace ns ON ns.oid = c.relnamespace" + , "JOIN pg_class fc ON fc.oid = con.confrelid" + , "JOIN pg_namespace fns ON fns.oid = fc.relnamespace" + -- ORDINALITY retains column ordering for both local and referenced sides + , "CROSS JOIN unnest(con.conkey, con.confkey)" + , " WITH ORDINALITY AS k(attnum, fattnum, n)" + , "JOIN pg_attribute a ON a.attnum = k.attnum AND a.attrelid = con.conrelid" + , "JOIN pg_attribute fa ON fa.attnum = k.fattnum AND fa.attrelid = con.confrelid" + -- retain only foreign key constraints + , "WHERE con.contype = 'f'" ] + mkFkPred (srcSchema, srcTbl, localCols, refSchema, refTbl, refCols, updCode, delCode) = + case (NE.nonEmpty (V.toList localCols), NE.nonEmpty (V.toList refCols)) of + (Just lcNE, Just rcNE) -> + Just $ Db.SomeDatabasePredicate + (Db.TableHasForeignKey (Db.QualifiedName srcSchema srcTbl) lcNE + (Db.QualifiedName refSchema refTbl) rcNE + (parsePgForeignKeyAction updCode) + (parsePgForeignKeyAction delCode)) + _ -> Nothing + foreignKeyChecks <- mapMaybe mkFkPred <$> + case subschemas of + Nothing -> + Pg.query_ conn (fromString (fkBaseQuery <> + " AND ns.nspname = any(current_schemas(false))\n" <> + "GROUP BY ns.nspname, c.relname, con.oid, fns.nspname, fc.relname,\n" <> + " con.confupdtype, con.confdeltype")) + Just ss -> + Pg.query conn (fromString (fkBaseQuery <> + " AND ns.nspname IN ?\n" <> + "GROUP BY ns.nspname, c.relname, con.oid, fns.nspname, fc.relname,\n" <> + " con.confupdtype, con.confdeltype")) + (Pg.Only (Pg.In ss)) + let enumerations = map (\(enumNm, _, options) -> Db.SomeDatabasePredicate (PgHasEnum enumNm (V.toList options))) enumerationData @@ -457,7 +505,25 @@ map (\(Pg.Only extname) -> Db.SomeDatabasePredicate (PgHasExtension extname)) <$> Pg.query_ conn "SELECT extname from pg_extension" - pure (tblsExist ++ columnChecks ++ primaryKeys ++ secondaryIndices ++ enumerations ++ extensions) + pure $ + concat + [ tblsExist + , columnChecks + , primaryKeys + , secondaryIndices + , foreignKeyChecks + , enumerations + , extensions + ] + +parsePgForeignKeyAction :: Char -> Db.ForeignKeyAction +parsePgForeignKeyAction 'c' = Db.ForeignKeyActionCascade +parsePgForeignKeyAction 'n' = Db.ForeignKeyActionSetNull +parsePgForeignKeyAction 'd' = Db.ForeignKeyActionSetDefault +parsePgForeignKeyAction 'r' = Db.ForeignKeyActionRestrict +parsePgForeignKeyAction 'a' = Db.ForeignKeyNoAction +parsePgForeignKeyAction c = error $ + "parsePgForeignKeyAction: unrecognised foreign key action code: " ++ show c -- * Postgres-specific data types 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 @@ -1179,10 +1179,24 @@ map fromPgTableConstraint constraints) <> emit ")" <> afterOptions +pgForeignKeyAction :: ForeignKeyAction -> PgSyntax +pgForeignKeyAction ForeignKeyActionCascade = emit "CASCADE" +pgForeignKeyAction ForeignKeyActionSetNull = emit "SET NULL" +pgForeignKeyAction ForeignKeyActionSetDefault = emit "SET DEFAULT" +pgForeignKeyAction ForeignKeyActionRestrict = emit "RESTRICT" +pgForeignKeyAction ForeignKeyNoAction = emit "NO ACTION" + instance IsSql92TableConstraintSyntax PgTableConstraintSyntax where primaryKeyConstraintSyntax fieldNames = PgTableConstraintSyntax $ - emit "PRIMARY KEY(" <> pgSepBy (emit ", ") (map pgQuotedIdentifier fieldNames) <> emit ")" + emit "PRIMARY KEY(" <> pgSepBy (emit ", ") (map pgQuotedIdentifier $ NE.toList fieldNames) <> emit ")" + foreignKeyConstraintSyntax localCols refTbl refCols onUpdate onDelete = + PgTableConstraintSyntax $ + emit "FOREIGN KEY(" <> pgSepBy (emit ", ") (map pgQuotedIdentifier $ NE.toList localCols) <> emit ")" <> + emit " REFERENCES " <> pgQuotedIdentifier refTbl <> + emit "(" <> pgSepBy (emit ", ") (map pgQuotedIdentifier $ NE.toList refCols) <> emit ")" <> + emit " ON UPDATE " <> pgForeignKeyAction onUpdate <> + emit " ON DELETE " <> pgForeignKeyAction onDelete instance Hashable PgColumnSchemaSyntax where hashWithSalt salt = hashWithSalt salt . fromPgColumnSchema diff --git a/Database/Beam/Postgres/TempTable.hs b/Database/Beam/Postgres/TempTable.hs new file mode 100644 --- /dev/null +++ b/Database/Beam/Postgres/TempTable.hs @@ -0,0 +1,62 @@ +{-# LANGUAGE ConstraintKinds #-} + +{-# OPTIONS_GHC -fno-warn-orphans #-} + +module Database.Beam.Postgres.TempTable + ( -- * Temporary tables + + -- ** @CREATE TEMPORARY TABLE@ + runCreateTempTable + + -- ** Options + , TempTableCreateMode(..) + , TempTableOptions(..) + , defaultTempTableOptions + + -- ** Schema constraint + , HasDefaultPostgresTempTableSchema + + ) where + +import Database.Beam.Migrate.TempTable + ( BeamHasTempTables(..) + , HasDefaultTempTableSchema + , TempTableCreateMode(..) + , TempTableOptions(..) + , defaultTempTableOptions + , runCreateTempTable + ) +import Database.Beam.Postgres.Types (Postgres) +import Database.Beam.Postgres.Syntax + ( PgCommandSyntax(..), PgCommandType(..) + , fromPgColumnSchema + , emit, pgQuotedIdentifier, pgSepBy, pgParens + ) + +-------------------------------------------------------------------------------- + +-- | Tables for which @CREATE TEMPORARY TABLE@ statements can be emitted +-- automatically. +type HasDefaultPostgresTempTableSchema tbl = HasDefaultTempTableSchema Postgres tbl + +instance BeamHasTempTables Postgres where + createTempTableCmd nm colDefs pkCols ifNotExists = + PgCommandSyntax PgCommandTypeDdl $ + emit "CREATE TEMPORARY TABLE " + <> (if ifNotExists then emit "IF NOT EXISTS " else mempty) + <> pgQuotedIdentifier nm + <> pgParens (pgSepBy (emit ", ") (colDefPieces ++ pkPiece)) + where + colDefPieces = + map (\(n, s) -> pgQuotedIdentifier n <> emit " " <> fromPgColumnSchema s) + colDefs + + pkPiece + | null pkCols = [] + | otherwise = + [ emit "PRIMARY KEY " + <> pgParens (pgSepBy (emit ", ") (map pgQuotedIdentifier pkCols)) ] + + dropTempTableIfExistsCmd nm = + PgCommandSyntax PgCommandTypeDdl $ + emit "DROP TABLE IF EXISTS " <> pgQuotedIdentifier nm 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.5.5.0 +version: 0.5.6.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 @@ -25,6 +25,8 @@ Database.Beam.Postgres.PgCrypto Database.Beam.Postgres.Extensions.UuidOssp + Database.Beam.Postgres.TempTable + other-modules: Database.Beam.Postgres.Connection Database.Beam.Postgres.Debug Database.Beam.Postgres.Extensions @@ -81,7 +83,8 @@ Database.Beam.Postgres.Test.Marshal, Database.Beam.Postgres.Test.Select, Database.Beam.Postgres.Test.DataTypes, - Database.Beam.Postgres.Test.Migrate + Database.Beam.Postgres.Test.Migrate, + Database.Beam.Postgres.Test.TempTable build-depends: aeson, base, diff --git a/test/Database/Beam/Postgres/Test/Migrate.hs b/test/Database/Beam/Postgres/Test/Migrate.hs --- a/test/Database/Beam/Postgres/Test/Migrate.hs +++ b/test/Database/Beam/Postgres/Test/Migrate.hs @@ -31,6 +31,9 @@ , dropSchemaWorks postgresConn , indexVerification postgresConn , uniqueIndexVerification postgresConn + , foreignKeyVerification postgresConn + , foreignKeyOnDeleteCascadeVerification postgresConn + , foreignKeyActionsWork postgresConn ] data CharT f @@ -189,6 +192,134 @@ (dbModification @_ @Postgres) { _idx_tbl = addTableIndex "idx_tbl_value_uniq" idxOpts (\t -> selectorColumnName _idx_value t NE.:| []) } + runBeamPostgres conn (verifySchema migrationBackend db) >>= \case + VerificationSucceeded -> return () + VerificationFailed failures -> fail ("Verification failed: " ++ show failures) + +-- Foreign key test tables + +data 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 Postgres) + +-- | Verifies that 'verifySchema' correctly detects a plain foreign key +foreignKeyVerification :: IO ByteString -> TestTree +foreignKeyVerification pgConn = + testCase "verifySchema correctly detects a plain foreign key" $ + withTestPostgres "db_fk" pgConn $ \conn -> do + Pg.execute_ conn "CREATE TABLE fk_parent (fk_parent_id integer NOT NULL PRIMARY KEY)" + Pg.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))" + let db :: CheckedDatabaseSettings Postgres FkDb + db = defaultMigratableDbSettings `withDbModification` + (dbModification @_ @Postgres) + { _fk_child = + addTableForeignKey (_fk_parent db) + (foreignKeyColumns _fk_child_parent_id) + primaryKeyColumns + ForeignKeyNoAction + ForeignKeyNoAction + <> modifyCheckedTable id + (FkChildT { _fk_child_id = "fk_child_id" + , _fk_child_parent_id = FkParentPk "fk_child_parent_id" }) } + runBeamPostgres conn (verifySchema migrationBackend db) >>= \case + VerificationSucceeded -> return () + VerificationFailed failures -> fail ("Verification failed: " ++ show failures) + +-- | Verifies that foreign key actions are enforced at runtime. +foreignKeyActionsWork :: IO ByteString -> TestTree +foreignKeyActionsWork pgConn = + testCase "cascading foreign key actions" $ + withTestPostgres "db_fk_actions" pgConn $ \conn -> do + let db :: CheckedDatabaseSettings Postgres FkDb + db = defaultMigratableDbSettings `withDbModification` + (dbModification @_ @Postgres) + { _fk_child = + addTableForeignKey (_fk_parent db) + (foreignKeyColumns _fk_child_parent_id) + primaryKeyColumns + ForeignKeyActionCascade + ForeignKeyActionCascade + } + unc = unCheckDatabase db + runBeamPostgres conn $ autoMigrate migrationBackend db + + -- Insert two parents and three children (two for parent 1, one for parent 2). + runBeamPostgres conn $ do + runInsert $ insert (_fk_parent unc) $ insertValues + [ FkParentT 1, FkParentT 2 ] + runInsert $ insert (_fk_child unc) $ insertValues + [ FkChildT 1 (FkParentPk 1), FkChildT 2 (FkParentPk 1), FkChildT 3 (FkParentPk 2) ] + + -- ON UPDATE CASCADE: changing fk_parent_id 1 → 10 should cascade to child rows. + runBeamPostgres conn $ + runUpdate $ update (_fk_parent unc) + (\p -> _fk_parent_id p <-. val_ 10) + (\p -> _fk_parent_id p ==. val_ 1) + childrenOf10 <- runBeamPostgres conn $ runSelectReturningList $ select $ + filter_ (\c -> let FkParentPk pid = _fk_child_parent_id c in pid ==. val_ 10) $ all_ (_fk_child unc) + assertEqual "two children should now reference updated parent id 10" + 2 (length childrenOf10) + childrenOf1 <- runBeamPostgres conn $ runSelectReturningList $ select $ + filter_ (\c -> let FkParentPk pid = _fk_child_parent_id c in pid ==. val_ 1) $ all_ (_fk_child unc) + assertEqual "no children should still reference old parent id 1" + 0 (length childrenOf1) + + -- ON DELETE CASCADE: deleting parent 2 should remove its child row. + runBeamPostgres conn $ + runDelete $ delete (_fk_parent unc) (\p -> _fk_parent_id p ==. val_ 2) + childrenOf2 <- runBeamPostgres conn $ runSelectReturningList $ select $ + filter_ (\c -> let FkParentPk pid = _fk_child_parent_id c in pid ==. val_ 2) $ all_ (_fk_child unc) + assertEqual "child of deleted parent 2 should be removed" + 0 (length childrenOf2) + allChildren <- runBeamPostgres conn $ runSelectReturningList $ select $ + all_ (_fk_child unc) + assertEqual "only the two children of parent 10 should remain" + 2 (length allChildren) + +-- | Verifies that 'verifySchema' correctly detects a foreign key with ON DELETE CASCADE +foreignKeyOnDeleteCascadeVerification :: IO ByteString -> TestTree +foreignKeyOnDeleteCascadeVerification pgConn = + testCase "verifySchema correctly detects a foreign key with ON DELETE CASCADE" $ + withTestPostgres "db_fk_cascade" pgConn $ \conn -> do + Pg.execute_ conn "CREATE TABLE fk_parent (fk_parent_id integer NOT NULL PRIMARY KEY)" + Pg.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) \ + \ON DELETE CASCADE)" + let db :: CheckedDatabaseSettings Postgres FkDb + db = defaultMigratableDbSettings `withDbModification` + (dbModification @_ @Postgres) + { _fk_child = + addTableForeignKey (_fk_parent db) + (foreignKeyColumns _fk_child_parent_id) + primaryKeyColumns + ForeignKeyNoAction + ForeignKeyActionCascade + <> modifyCheckedTable id + (FkChildT { _fk_child_id = "fk_child_id" + , _fk_child_parent_id = FkParentPk "fk_child_parent_id" }) } runBeamPostgres conn (verifySchema migrationBackend db) >>= \case VerificationSucceeded -> return () VerificationFailed failures -> fail ("Verification failed: " ++ show failures) diff --git a/test/Database/Beam/Postgres/Test/TempTable.hs b/test/Database/Beam/Postgres/Test/TempTable.hs new file mode 100644 --- /dev/null +++ b/test/Database/Beam/Postgres/Test/TempTable.hs @@ -0,0 +1,87 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE StandaloneDeriving #-} + +module Database.Beam.Postgres.Test.TempTable (tests) where + +import Data.ByteString (ByteString) +import Data.Int (Int32) +import Data.Kind (Type) +import GHC.Exts (Any) + +import Database.Beam +import Database.Beam.Postgres +import Database.PostgreSQL.Simple (execute_) + +import Test.Tasty +import Test.Tasty.HUnit + +import Database.Beam.Postgres.Test + +tests :: IO ByteString -> TestTree +tests getConn = testGroup "Temporary table tests" + [ testStageFilteredRows getConn + , testDropAndCreate getConn + ] + +-- | A permanent table of items. +data ItemT f = Item + { itemId :: C f Int32 + , itemValue :: C f Int32 + } deriving (Generic, Beamable) + +deriving instance Show (ItemT Identity) +deriving instance Eq (ItemT Identity) + +instance Table ItemT where + data PrimaryKey ItemT f = ItemKey (C f Int32) + deriving (Generic, Beamable) + primaryKey = ItemKey . itemId + +data ItemDb entity = ItemDb + { dbItems :: entity (TableEntity ItemT) + } deriving (Generic, Database Postgres) + +itemDb :: DatabaseSettings be ItemDb +itemDb = defaultDbSettings + +-- | Stage a filtered subset of a permanent table in a temp table using +-- INSERT INTO temp SELECT ... FROM permanent WHERE ..., then read it back. +testStageFilteredRows :: IO ByteString -> TestTree +testStageFilteredRows getConn = testCase "sanity test" $ + withTestPostgres "temp_table_stage" getConn $ \conn -> do + execute_ conn "CREATE TABLE items (id INT PRIMARY KEY, value INT NOT NULL)" + result <- runBeamPostgres conn $ do + runInsert $ insert (dbItems itemDb) $ insertValues + [ Item 1 5, Item 2 10, Item 3 15, Item 4 20, Item 5 25 ] + scratch <- runCreateTempTable @_ @Any @ItemT + defaultTempTableOptions "high_value_items" + runInsert $ insert scratch $ insertFrom $ + filter_ (\r -> itemValue r >. val_ 15) $ + all_ (dbItems itemDb) + runSelectReturningList $ select $ + orderBy_ (asc_ . itemId) $ all_ scratch + assertEqual "high-value items staged" [Item 4 20, Item 5 25] result + +testDropAndCreate :: IO ByteString -> TestTree +testDropAndCreate getConn = testCase "drop-and-create" $ + withTestPostgres "temp_table_drop_and_create" getConn $ \conn -> do + execute_ conn "CREATE TABLE items (id INT PRIMARY KEY, value INT NOT NULL)" + result <- runBeamPostgres conn $ do + runInsert $ insert (dbItems itemDb) $ insertValues + [ Item 1 5, Item 2 10, Item 3 15, Item 4 20, Item 5 25 ] + -- First pass: stage low-value items. + scratch <- runCreateTempTable @_ @Any @ItemT + defaultTempTableOptions "scratch" + runInsert $ insert scratch $ insertFrom $ + filter_ (\r -> itemValue r <=. val_ 15) $ + all_ (dbItems itemDb) + -- Second pass: drop and recreate, then stage high-value items only. + scratch' <- runCreateTempTable @_ @Any @ItemT + (defaultTempTableOptions { tempTableCreateMode = DropAndCreate }) + "scratch" + runInsert $ insert scratch' $ insertFrom $ + filter_ (\r -> itemValue r >. val_ 15) $ + all_ (dbItems itemDb) + runSelectReturningList $ select $ + orderBy_ (asc_ . itemId) $ all_ scratch' + assertEqual "only high-value items after reset" [Item 4 20, Item 5 25] result diff --git a/test/Main.hs b/test/Main.hs --- a/test/Main.hs +++ b/test/Main.hs @@ -11,6 +11,7 @@ import qualified Database.Beam.Postgres.Test.Marshal as Marshal import qualified Database.Beam.Postgres.Test.DataTypes as DataType import qualified Database.Beam.Postgres.Test.Migrate as Migrate +import qualified Database.Beam.Postgres.Test.TempTable as TempTable import Database.PostgreSQL.Simple ( ConnectInfo(..), defaultConnectInfo ) import qualified Database.PostgreSQL.Simple as Postgres @@ -23,6 +24,7 @@ , Select.tests getConnStr , DataType.tests getConnStr , Migrate.tests getConnStr + , TempTable.tests getConnStr ]