{-
    pg-harness, REST service for creating temporary PostgreSQL databases.
    Copyright (C) 2014, 2015 Bardur Arantsson

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU Affero General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU Affero General Public License for more details.

    You should have received a copy of the GNU Affero General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
-}
module Main ( main ) where

import           Control.Concurrent (threadDelay)
import           Control.Concurrent.Async (async, withAsync, waitCatch)
import           Control.Exception (bracket)
import           Control.Monad (void)
import           Control.Monad.Trans.Class (lift)
import           Data.String (fromString)
import qualified Data.Text.Lazy as TL
import           Database.PostgreSQL.Simple (Connection, ConnectInfo(..), Only(..))
import qualified Database.PostgreSQL.Simple as P
import           System.IO (stderr, hPutStrLn)
import           System.Random (randomRIO)
import           Web.Scotty (scotty, post, text, ScottyM, raise)
import           Paths_pg_harness_server (getDataFileName)
import           PgHarness.Mutex
import           PgHarness.Configuration
import           PgHarness.DatabaseId

-- Perform an operation with an open connection to the database.
-- The connection will automatically be closed after the operation
-- completes, regardless of whether it completes successfully or
-- not (e.g. if there's an exception).
withConnection :: Configuration -> (Connection -> IO a) -> IO a
withConnection configuration action = bracket (P.connect connectInfo) P.close action
  where
    connectInfo = ConnectInfo
      { connectHost = cfgHost configuration
      , connectPort = cfgPort configuration
      , connectUser = cfgUser configuration
      , connectPassword = cfgPassword configuration
      , connectDatabase = unquotedIdentifier (cfgDatabase configuration)
      }

-- Create a random valid PostgreSQL identifier.
mkRandomIdent :: IO String
mkRandomIdent = do
  h <- chooseElement letters
  t <- sequence $ replicate 32 $ chooseElement lettersAndDigits
  return (h:t)
  where
    letters = "abcdefhijklmnopqrstuvwxyz"
    digits = "0123456789"
    lettersAndDigits = letters ++ digits
    chooseElement t = do
      i <- randomRIO (0, length t - 1)
      return $ t !! i

-- Create temporary database ID.
mkTemporaryDatabaseId :: IO (Either String DatabaseId)
mkTemporaryDatabaseId = fmap mkDatabaseId $ fmap ("temp_" ++) mkRandomIdent

-- Create temporary database and return its name.
createTemporaryDatabase :: Configuration -> DatabaseId -> IO ()
createTemporaryDatabase configuration databaseId = do
  -- Connect to the administrative database.
  withConnection configuration $ \connection -> do
    -- Create the temporary database.
    void $ P.execute_ connection $ createSql
    putStrLn $ "Created temporary database: " ++ sqlDatabaseId
    -- Spawn a thread to kill the database after a certain delay.
    void $ async $ do
      threadDelay $ (cfgDurationSeconds configuration) * 1000000
      dropDatabase configuration databaseId
  where
    createSql = fromString $
      "CREATE DATABASE " ++ sqlDatabaseId ++
      "  WITH TEMPLATE " ++ (sqlIdentifier $ cfgTemplateDatabase configuration) ++
      "          OWNER \"" ++ (cfgTestUser configuration) ++ "\""

    sqlDatabaseId = sqlIdentifier databaseId

dropDatabase :: Configuration -> DatabaseId -> IO ()
dropDatabase configuration databaseId =
  -- Since this is running in a separate thread we need explicit
  -- exception handling to avoid silent exceptions.
  withAsync doDropDatabase waitCatch >>= \case
    Left exc -> hPutStrLn stderr $ "Exception occurred while creating temporary database: " ++ show exc
    Right _  -> return ()
  where
    doDropDatabase = withConnection configuration $ \connection -> do
      -- We need to block new connections to the temporary database; otherwise
      -- a reconnect during the destruction could foil our attempt to drop.
      void $ P.execute_ connection blockConnectionsSql
      -- Now we can kill the backends, i.e. terminate all the connections. No
      -- new connections can be created because of the previous block.
      P.forEach_ connection terminateConnectionsSql $ \(Only (_::Bool)) -> return ()
      -- Finally, we can drop.
      void $ P.execute_ connection dropSql
      -- Logging
      putStrLn $ "Dropped temporary database: " ++ sqlDatabaseId

    dropSql = fromString $
        "DROP DATABASE " ++ sqlDatabaseId

    blockConnectionsSql = fromString $
        "UPDATE pg_database \
        \   SET datallowconn = FALSE \
        \ WHERE datname = '" ++ sqlDatabaseId ++ "'"

    terminateConnectionsSql = fromString $
        "SELECT pg_terminate_backend(pid) \
        \  FROM pg_stat_activity \
        \ WHERE pid <> pg_backend_pid() \
        \   AND datname = '" ++ sqlDatabaseId ++ "'"

    sqlDatabaseId = sqlIdentifier databaseId

-- REST interface
routes :: Configuration -> ScottyM ()
routes configuration = do
  -- Mutex to prevent multiple "create" requests from being
  -- processed simultaneously; PostgreSQL cannot handle
  -- "cloning" a template concurrently.
  mutex <- lift $ mkMutex
  -- Add all the routes.
  post "/" $ do
    -- Generate a name for temporary database.
    lift mkTemporaryDatabaseId >>= \case
      Left err ->
        raise $ TL.pack err
      Right databaseId -> do
        -- Create the temporary database.
        lift $ mutex $ createTemporaryDatabase configuration databaseId
        -- Return a string with the username/password and database name.
        text $ TL.pack $
          cfgTestUser configuration ++ ":" ++ cfgTestPassword configuration ++
          "@" ++
          cfgHost configuration ++ ":" ++ (show $ cfgPort configuration) ++
          "/" ++
          (unquotedIdentifier databaseId)

main :: IO ()
main = do
  -- Load configuration from file.
  getDataFileName "pg-harness.ini" >>= loadConfiguration >>= \case
    Left msg -> hPutStrLn stderr msg
    Right configuration -> do
      -- Start the web serving thread
      putStrLn $ "Starting with configuration: " ++ show configuration
      scotty (cfgListenPort configuration) $ routes configuration