module IHP.Test.Database where

import IHP.Prelude
import qualified Database.PostgreSQL.Simple as PG
import qualified Database.PostgreSQL.Simple.Types as PG
import qualified Data.UUID.V4 as UUID
import qualified Data.UUID as UUID
import qualified Data.Text as Text
import qualified Data.ByteString as ByteString
import qualified IHP.LibDir as LibDir
import qualified Control.Exception as Exception

import qualified System.Process as Process
import qualified System.Directory as Directory

data TestDatabase = TestDatabase
    { TestDatabase -> Text
name :: Text
    , TestDatabase -> ByteString
url :: ByteString
    }

-- | Given a Postgres Database URL it creates a new randomly named database on the database server. Returns a database url to the freshly created database
--
-- >>> createTestDatabase "postgresql:///app?host=/myapp/build/db"
-- TestDatabase { name = "test-7d3bd463-4cce-413f-a272-ac52e5d93739", url = "postgresql:///test-7d3bd463-4cce-413f-a272-ac52e5d93739?host=/myapp/build/db" }
--
createTestDatabase :: ByteString -> IO TestDatabase
createTestDatabase :: ByteString -> IO TestDatabase
createTestDatabase ByteString
databaseUrl = do
--    currentDir <- Directory.getCurrentDirectory
    UUID
databaseId <- IO UUID
UUID.nextRandom
    let databaseName :: Text
databaseName = Text
"test-" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> UUID -> Text
UUID.toText UUID
databaseId

    ByteString -> (Connection -> IO Int64) -> IO Int64
forall {c}. ByteString -> (Connection -> IO c) -> IO c
withConnection ByteString
databaseUrl \Connection
connection -> do
        Connection -> Query -> [Identifier] -> IO Int64
forall q. ToRow q => Connection -> Query -> q -> IO Int64
PG.execute Connection
connection Query
"CREATE DATABASE ?" [Text -> Identifier
PG.Identifier Text
databaseName]

    let ByteString
newUrl :: ByteString = ByteString
databaseUrl
            ByteString -> (ByteString -> Text) -> Text
forall {t1} {t2}. t1 -> (t1 -> t2) -> t2
|> ByteString -> Text
forall a b. ConvertibleStrings a b => a -> b
cs
            Text -> (Text -> Text) -> Text
forall {t1} {t2}. t1 -> (t1 -> t2) -> t2
|> HasCallStack => Text -> Text -> Text -> Text
Text -> Text -> Text -> Text
Text.replace Text
"postgresql:///app" (Text
"postgresql:///" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
databaseName)
            Text -> (Text -> ByteString) -> ByteString
forall {t1} {t2}. t1 -> (t1 -> t2) -> t2
|> Text -> ByteString
forall a b. ConvertibleStrings a b => a -> b
cs

    Text
libDir <- IO Text
LibDir.findLibDirectory

    -- We use the system psql to handle the initial Schema Import as it can handle
    -- complex Schema including variations in formatting, custom types, functions, and table definitions.
    let importSql :: String -> IO ()
importSql String
file = String -> IO ()
Process.callCommand (String
"psql " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> (ByteString -> String
forall a b. ConvertibleStrings a b => a -> b
cs ByteString
newUrl) String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" < " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
file)

    String -> IO ()
importSql (Text -> String
forall a b. ConvertibleStrings a b => a -> b
cs Text
libDir String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"/IHPSchema.sql")
    String -> IO ()
importSql String
"Application/Schema.sql"

    TestDatabase -> IO TestDatabase
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TestDatabase { $sel:name:TestDatabase :: Text
name = Text
databaseName, $sel:url:TestDatabase :: ByteString
url = ByteString
newUrl }

-- | Given the master connection url and the open test database, this will clean up the test database
--
-- >>> deleteDatabase "postgresql:///app?host=/myapp/build/db" TestDatabase { name = "test-7d3bd463-4cce-413f-a272-ac52e5d93739", url = "postgresql:///test-7d3bd463-4cce-413f-a272-ac52e5d93739?host=/myapp/build/db" }
--
-- The master connection url needs to be passed as we cannot drop the database we're currently connected to, and therefore
-- we cannot use the test database itself.
--
deleteDatabase :: ByteString -> TestDatabase -> IO ()
deleteDatabase :: ByteString -> TestDatabase -> IO ()
deleteDatabase ByteString
masterDatabaseUrl TestDatabase
testDatabase = do
    ByteString -> (Connection -> IO Int64) -> IO Int64
forall {c}. ByteString -> (Connection -> IO c) -> IO c
withConnection ByteString
masterDatabaseUrl \Connection
connection -> do
        -- The WITH FORCE is required to force close open connections
        -- Otherwise the DROP DATABASE takes a few seconds to execute
        Connection -> Query -> [Identifier] -> IO Int64
forall q. ToRow q => Connection -> Query -> q -> IO Int64
PG.execute Connection
connection Query
"DROP DATABASE ? WITH (FORCE)" [Text -> Identifier
PG.Identifier (TestDatabase
testDatabase.name)]
    () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()


withConnection :: ByteString -> (Connection -> IO c) -> IO c
withConnection ByteString
databaseUrl = IO Connection
-> (Connection -> IO ()) -> (Connection -> IO c) -> IO c
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
Exception.bracket (ByteString -> IO Connection
PG.connectPostgreSQL ByteString
databaseUrl) Connection -> IO ()
PG.close