{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-
    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.IO.Class (MonadIO, liftIO)
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           Network.Wai.Handler.Warp (setPort, setHost, defaultSettings)
import           System.IO (stderr, hPutStrLn)
import           System.Random (randomRIO)
import           Web.Scotty (ScottyM, Options(..), scottyOpts, post, text)
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
      , connectPort = cfgPort
      , connectUser = cfgUser
      , connectPassword = cfgPassword
      , connectDatabase = unquotedIdentifier cfgDatabase
      }

-- 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@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 * 1000000
      dropDatabase configuration databaseId
  where
    createSql = fromString $
      "CREATE DATABASE " ++ sqlDatabaseId ++
      "  WITH TEMPLATE " ++ (sqlIdentifier cfgTemplateDatabase) ++
      "          OWNER \"" ++ cfgTestUser ++ "\""

    sqlDatabaseId = sqlIdentifier databaseId

dropDatabase :: Configuration -> DatabaseId -> IO ()
dropDatabase configuration@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

-- Handle a create request, returns the database ID.
handleCreateRequest :: MonadIO m => Configuration -> Mutex -> m DatabaseId
handleCreateRequest configuration mutex = do
  -- Generate a name for temporary database.
  liftIO mkTemporaryDatabaseId >>= \case
    Left err ->
      error err -- We're generating the name ourselves, so no error
                -- handling needed.
    Right databaseId -> do
      -- Create the temporary database.
      liftIO $ withMutex mutex $ createTemporaryDatabase configuration databaseId
      return databaseId

-- REST interface
routes :: Configuration -> Mutex -> ScottyM ()
routes configuration@Configuration{..} mutex = do
  -- Add all the routes.
  post "/" $ do
    databaseId <- handleCreateRequest configuration mutex
    text $ TL.unlines
      [ TL.pack $ cfgTestUser
      , TL.pack $ cfgTestPassword
      , TL.pack $ cfgHost
      , TL.pack $ show $ cfgPort
      , TL.pack $ unquotedIdentifier databaseId
      ]

main :: IO ()
main = do
  -- Load configuration from file.
  getDataFileName "pg-harness.ini" >>= loadConfiguration >>= \case
    Left msg -> hPutStrLn stderr msg
    Right configuration -> do
      -- Build options for Warp.
      let warpSettings =
            setPort (cfgListenPort configuration) $
            setHost (fromString $ cfgListenHost configuration) $
            defaultSettings
          options = Options { verbose = 1
                            , settings = warpSettings
                            }
      -- Mutex to prevent multiple "create" requests from being
      -- processed simultaneously; PostgreSQL cannot handle
      -- "cloning" a template concurrently.
      mutex <- mkMutex
      -- Force immediate creation of a single temporary database; we
      -- do this so that any configuration issues will become obvious
      -- on startup.
      void $ handleCreateRequest configuration mutex
      -- Start the web serving thread
      putStrLn $ "Starting with configuration: " ++ show configuration
      scottyOpts options $ routes configuration mutex