{-# LANGUAGE OverloadedStrings #-}

module Database.Beam.AutoMigrate.BenchUtil
  ( SpineStrict (..)
  , predictableSchemas
  , connInfo
  , setupDatabase
  , cleanDatabase
  , tearDownDatabase
  )
where

import Control.DeepSeq
import Control.Exception (finally)
import Data.ByteString (ByteString)
import Database.Beam.AutoMigrate
import Database.Beam.AutoMigrate.Schema.Gen (genSimilarSchemas)
import qualified Database.PostgreSQL.Simple as Pg
import System.Random.SplitMix (mkSMGen)
import Test.QuickCheck.Gen
import Test.QuickCheck.Random

newtype SpineStrict = SS {SpineStrict -> Diff
unSS :: Diff}

-- For us is enough to make the list of edits spine-strict.
instance NFData SpineStrict where
  rnf :: SpineStrict -> ()
rnf (SS (Left DiffError
e)) = DiffError -> ()
forall a. NFData a => a -> ()
rnf DiffError
e
  rnf (SS (Right [WithPriority Edit]
edits)) = [WithPriority Edit] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [WithPriority Edit]
edits Int -> () -> ()
forall a b. NFData a => a -> b -> b
`deepseq` ()

predictableSchemas :: Int -> IO (Schema, Schema)
predictableSchemas :: Int -> IO (Schema, Schema)
predictableSchemas Int
tableNum = do
  let g :: QCGen -> Int -> (Schema, Schema)
g = Gen (Schema, Schema) -> QCGen -> Int -> (Schema, Schema)
forall a. Gen a -> QCGen -> Int -> a
unGen Gen (Schema, Schema)
genSimilarSchemas
  let r :: QCGen
r = SMGen -> QCGen
QCGen (Word64 -> SMGen
mkSMGen Word64
42)
  (Schema, Schema) -> IO (Schema, Schema)
forall (m :: * -> *) a. Monad m => a -> m a
return (QCGen -> Int -> (Schema, Schema)
g QCGen
r Int
tableNum)

connInfo :: ByteString
connInfo :: ByteString
connInfo = ByteString
"host=localhost port=5432 dbname=beam-migrate-prototype-bench"

setupDatabase :: Schema -> IO Pg.Connection
setupDatabase :: Schema -> IO Connection
setupDatabase Schema
dbSchema = do
  Connection
conn <- ByteString -> IO Connection
Pg.connectPostgreSQL ByteString
connInfo
  let mig :: Migration Pg
mig = Diff -> Migration Pg
forall (m :: * -> *). Monad m => Diff -> Migration m
createMigration (Schema -> Schema -> Diff
forall a. Diffable a => a -> a -> Diff
diff Schema
dbSchema Schema
noSchema)
  MonadBeam Postgres Pg => Connection -> Migration Pg -> IO ()
Connection -> Migration Pg -> IO ()
runMigrationUnsafe Connection
conn Migration Pg
mig -- At this point the DB contains the full schema.
  Connection -> IO Connection
forall (f :: * -> *) a. Applicative f => a -> f a
pure Connection
conn

cleanDatabase :: Pg.Connection -> IO ()
cleanDatabase :: Connection -> IO ()
cleanDatabase Connection
conn = do
  -- Delete all tables to start from a clean slate
  Int64
_ <- Connection -> Query -> IO Int64
Pg.execute_ Connection
conn Query
"DROP SCHEMA public CASCADE"
  Int64
_ <- Connection -> Query -> IO Int64
Pg.execute_ Connection
conn Query
"CREATE SCHEMA public"
  Int64
_ <- Connection -> Query -> IO Int64
Pg.execute_ Connection
conn Query
"GRANT USAGE ON SCHEMA public TO public"
  Int64
_ <- Connection -> Query -> IO Int64
Pg.execute_ Connection
conn Query
"GRANT CREATE ON SCHEMA public TO public"
  () -> IO ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

tearDownDatabase :: Pg.Connection -> IO ()
tearDownDatabase :: Connection -> IO ()
tearDownDatabase Connection
conn = Connection -> IO ()
cleanDatabase Connection
conn IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO a
`finally` Connection -> IO ()
Pg.close Connection
conn