Tmp.hs

{-| Create temporary postgresql databases.

The main usecase for this are tests where you don’t want to assume that a certain database exists.
-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Database.PostgreSQL.Tmp
  ( defaultDB
  , DBInfo(..)
  , withTmpDB
  , withTmpDB'
  -- * Low level APIs
  , createTmpDB
  , dropTmpDB
  , newRole
  , dropRole
  , newDB
  , dropDB
  ) where

import           Control.Applicative (pure)
import           Control.Exception
import           Data.ByteString (ByteString)
import           Data.Coerce
import           Data.Int
import           Data.Monoid
import qualified Data.Text as T
import           Database.PostgreSQL.Simple
import           Database.PostgreSQL.Simple.Types

-- | Connection string for the @postgres@ database owned by the
-- @postgres@ user
defaultDB :: ByteString
defaultDB = "dbname='postgres' user='postgres'"

-- | The data necessary to connect to the temporary database
data DBInfo =
  DBInfo {dbName :: T.Text
         ,roleName :: T.Text} deriving (Show,Read,Eq,Ord)

-- | Convenience wrapper for 'withTmpDB'' using 'defaultDB'
withTmpDB :: (DBInfo -> IO a) -> IO a
withTmpDB = withTmpDB' defaultDB

-- | Create a temporary database and a temporary role that the
-- callback can operate on. After the action has finished the database
-- and the role are destroyed.
--
-- This is a `bracket`-style wrapper around 'createTmpDB' and 'dropTmpDB'
--
-- This function assumes that the connection string points to a
-- database containing the tables called @pg_roles@ and @pg_database@
-- and that the user has the @CREATEDB@ and @CREATEROLE@ privileges.
withTmpDB' :: ByteString -> (DBInfo -> IO a) -> IO a
withTmpDB' conStr f = bracket (createTmpDB conStr) dropTmpDB (\(_,dbInfo) -> f dbInfo)

-- | Create a temporary database and a temporary role.
--
-- To destroy the database and the role use `dropTmpDB`.
--
-- This function assumes that the connection string points to a
-- database containing the tables called @pg_roles@ and @pg_database@
-- and that the user has the @CREATEDB@ and @CREATEROLE@ privileges.
createTmpDB :: ByteString -> IO (Connection, DBInfo)
createTmpDB conStr = do
  conn <- connectPostgreSQL conStr
  role <- newRole conn
  db <- newDB conn role
  pure (conn, DBInfo {dbName = db,roleName = role})

-- | Destroy the database and the role created by `createTmpDB`.
dropTmpDB :: (Connection, DBInfo) -> IO ()
dropTmpDB (conn, DBInfo db role) = do
  _ <- dropDB conn db
  _ <- dropRole conn role
  close conn

-- | Create a new role that does not already exist and return its name.
--
-- The new role does not have a password and has the @CREATEDB@
-- privilege. The database that the connection points to is assumed to
-- contain a table called @pg_roles@ with a @rolname@ column.
newRole :: Connection -> IO T.Text
newRole conn =
  do (roles :: [Only T.Text]) <- query_ conn "SELECT rolname FROM pg_roles"
     let newName = freshName "tmp" (coerce roles)
     _ <- execute conn "CREATE USER ? WITH CREATEDB" (Only (Identifier newName))
     pure newName

-- | Drop the role.
dropRole :: Connection -> T.Text -> IO Int64
dropRole conn name = execute conn "DROP ROLE ?" (Only (Identifier name))

-- | Create a new database that is owned by the user.
newDB :: Connection -> T.Text -> IO T.Text
newDB conn role =
  do (dbNames :: [Only T.Text]) <- query_ conn "SELECT datname FROM pg_database"
     let newName = freshName "tmp" (coerce dbNames)
     _ <- execute conn "CREATE DATABASE ? OWNER ?" (Identifier newName,Identifier role)
     pure newName

-- | Drop the database.
dropDB :: Connection -> T.Text -> IO Int64
dropDB conn name =
  execute conn "DROP DATABASE ?" (Only (Identifier name))

-- | Create a fresh name that is not in the list of already existing names.
--
-- The fresh name is generated by appending a number to the supplied
-- template.
freshName :: T.Text -> [T.Text] -> T.Text
freshName template existingNames = loop 0
-- We could use a Set here to speed up the lookup, however the
-- construction of that Set is linear as well so it would only pay off
-- if at least one of the lookups fails.
  where loop :: Int -> T.Text
        loop i =
          if (template <> T.pack (show i)) `elem` existingNames
             then loop (i + 1)
             else (template <> T.pack (show i))