sdb.hs

#!/usr/bin/env stack
-- stack --install-ghc runghc --package turtle --package shake -- -rtsopts -with-rtsopts=-I0 -O0

{-# LANGUAGE OverloadedStrings #-}

-- | sdb.hs - a standalone script to manage a local development postgres cluster.

-- Code notes:
--
-- * Original idea taken from https://git.gnu.io/snippets/3
-- * underscore_names are used for functions that mimic shell commands.
-- * I intentionally did not use Turtle.Options nor optparse-applicative.
--   The autogenerated help text is no good.

import Prelude hiding (FilePath)

import Control.Exception.Base (bracket)
import Development.Shake
import GHC.IO.Handle (hDuplicate, hDuplicateTo)
import System.Environment (getProgName, getArgs)
import System.IO (openFile, IOMode(..))
import Turtle hiding (need, opt)
import qualified Data.Text as T
import qualified Filesystem.Path.CurrentOS as P
import qualified System.IO as H
import qualified Turtle

pgWorkDir :: FilePath
pgWorkDir = ".postgres-work"

usageText :: Text -> Shell ()
usageText this = mapM_ err
    [ this <> ": a wrapper to set up environment variables and run various"
    , "commands for hacking on the Snowdrift.coop website."
    , ""
    , "Usage:"
    , ""
    , "    " <> this <> " ACTION [ARGS]"
    , ""
    , "Where ACTION may be one of:"
    , ""
    , "    test              run 'stack test'"
    , "    devel             start 'yesod devel'"
    , "    ghci              start 'ghci' with the db set up for using app/DevelMain.hs"
    , "    clean             \"rm -rf\" the whole cluster"
    , "    export            create data dumps with pg_dump"
    , "    help              print this text"
    , ""
    , "    # Expert commands:"
    , "    env               print export commands for PGHOST and PGDATA"
    , "                             e.g. 'source $(" <> this <> " env)'"
    , "    start             start the cluster (normally done automatically)"
    , "    stop              stop  the cluster (ditto)"
    , "    pg_ctl            run Postgres' pg_ctl(1) utility"
    , "    psql              connect to lypaste_development with psql"
    ]

dbRunning, dbCluster :: H.FilePath
dbRunning = ".postgres-work/data/postmaster.pid"
dbCluster = ".postgres-work/data/postgresql.conf"

main :: IO ()
main = sh $ do
    (dbdir, pghost, pgdata) <- initEnv

    -- An escape hatch before getting down to shake
    args <- liftIO getArgs
    case args of
        ["env"] -> do
            echo ("export PGHOST=" <> toText_ pghost)
            echo ("export PGDATA=" <> toText_ pgdata)
        ("pg_ctl":as') -> procs "pg_ctl" (map T.pack as') empty
        _ -> liftIO (shakeit dbdir pghost pgdata)

shakeit :: FilePath -> FilePath -> FilePath -> IO ()
shakeit dbdir pghost pgdata = shakeArgs shakeOptions $ do
    want ["help"]

    -- Very basic

    phony "help" $ actsh (usageText . T.pack =<< liftIO getProgName)

    -- Basic

    phony "test" $ do
        need [dbRunning]
        command [] "stack" ["test"]

    phony "devel" $ do
        need [dbRunning]
        command [Cwd "."] "stack" ["exec", "yesod", "devel"]

    phony "ghci" $ do
        need [dbRunning]
        command [Cwd "."] "stack" ["ghci", "--package", "foreign-store", "--test"]

    phony "clean" $ do
        need ["stop"]
        removeFilesAfter ".postgres-work" ["//*"]

    phony "export" $ actsh $ do
        y <- testfile (fromText (T.pack dbRunning))
        if y
            then exportDb dbdir
            else err "Is your database running? (Then you should probably catch it!)"

    -- Advanced

    phony "start" (need [dbRunning])

    phony "stop"
        (actsh (shell "pg_ctl status" empty .&&. shell "pg_ctl stop" empty))

    dbRunning %> const (do
        need [dbCluster]
        actsh (shell "pg_ctl status" empty .||. shell "pg_ctl -w start" empty)
        )

    dbCluster %> const (actsh (initCluster pghost pgdata))

    phony "psql" $ do
        need [dbRunning]
        command [] "psql" ["lypaste_development"]

exportDb :: FilePath -> Shell ()
exportDb dbdir = do
    step "Dumping to devDB.sql..."
    mktree (dbdir </> "dev")
    output (dbdir </> "dev" </> "devDB.sql") $ pgDump "lypaste_development"
  where
    pgDump db = inproc "pg_dump" ["--no-owner", "--no-privileges", "--create", db] empty

initCluster :: FilePath -> FilePath -> Shell ()
initCluster pghost pgdata = do
    step "Creating directories..."
    mktree pghost
    mktree pgdata

    step "Initializing cluster..."
    hush $ procs "pg_ctl" ["initdb", "-o", "--nosync --auth=peer", "-D", pgdata'] empty

    step "Updating cluster configuration file..."
    setPgConfigOpts (pgdata </> "postgresql.conf")
        [ -- set the unix socket directory because pg_ctl start doesn't
          -- pay attention to PGHOST (wat.)
          ("unix_socket_directory", pghost')
        , ("unix_socket_directories", pghost')
        , ("archive_mode", "off")
        , ("fsync", "off")
        , ("wal_level", "minimal")
          -- don't bother listening on a port, just a socket.
        , ("listen_addresses", "''")
        ]

    step "Starting database server..."
    procs "pg_ctl" ["start", "-w"] empty

    step "Creating databases..."
    procs "createdb" ["lypaste_development"] empty
    procs "createdb" ["lypaste_test"] empty

    step "Success."

  where

    pghost'  = "'" <> toText_ pghost <> "'"
    pgdata'  = toText_ pgdata
    setPgConfigOpts config opts =
        inplace_ (choice (patterns opts)) config
    patterns opts =
        map (fmap Just . uncurry optSettingPattern) opts
        <> [ commentOrEmpty >> pure Nothing ]
    commentOrEmpty :: Pattern Text
    commentOrEmpty =
        contains (begins spaces >> (T.singleton <$> (char '#' <|> newline)))
    optSettingPattern opt value = do
        -- match the line with the option
        _ <- contains $ do
            void $ begins (star (oneOf " #"))
            void $ text opt
            once (oneOf " =")
        -- replace it with 'opt = value'
        return (opt <> " = " <> value)

-- | Create and export some env variables
initEnv :: Shell (FilePath, FilePath, FilePath)
initEnv = do
    dbdir <- getProjectRoot
    Just path <- Turtle.need "PATH"
    pgPath <- inshell "pg_config --bindir" ""

    let pghost = dbdir </> pgWorkDir </> "sockets"
        pgdata = dbdir </> pgWorkDir </> "data"
    export "PGHOST" (toText_ pghost)
    export "PGDATA" (toText_ pgdata)
    export "PATH" (format (s%":"%s) path pgPath)
    return (dbdir, pghost, pgdata)
  where
    getProjectRoot =
        realpath =<< (directory . P.decodeString <$> liftIO getProgName)


-- ##
-- ## Helper functions/additions to underlying libs
-- ##

-- | Use a 'Shell a' as an 'Action ()'
actsh :: Shell a -> Action ()
actsh = liftIO . sh

-- | Print a header for a step
step :: Text -> Shell ()
step msg = err ("## " <> msg)

-- | Seeing as I use this everywhere
toText_ :: FilePath -> Text
toText_ = format fp

-- | inplace with filtering.
inplace_ :: MonadIO io => Pattern (Maybe Text) -> FilePath -> io ()
inplace_ pat file = liftIO (runManaged (do
    here              <- pwd
    (tmpfile, handle) <- mktemp here "turtle"
    outhandle handle (sed_ pat (input file))
    liftIO (H.hClose handle)
    mv tmpfile file ))

-- | sed with filtering.
sed_ :: Pattern (Maybe Text) -> Shell Text -> Shell Text
sed_ pat orig = flatten $ do
    when (matchesEmpty pat) (die message)
    let pat' = fmap mconcat
            (many (pat <|> fmap (Just . T.singleton) anyChar))
    txt    <- orig
    txt':_ <- return (match pat' txt)
    return txt'
  where
    message = "sed: the given pattern matches the empty string"
    matchesEmpty = not . null . flip match ""
    flatten my = do
        Just y <- my
        return y

-- | Run a shell and send stdout/stderr to nowhere
hush :: Shell () -> Shell ()
hush act = liftIO $ bracket
    (do
        saveErr <- hDuplicate H.stderr
        saveOut <- hDuplicate H.stdout
        h <- openFile "/dev/null" WriteMode
        hDuplicateTo h H.stderr
        hDuplicateTo h H.stdout
        return (saveErr, saveOut)
    )
    (\(saveErr, saveOut) -> do
        hDuplicateTo saveErr H.stderr
        hDuplicateTo saveOut H.stdout)
    (const (sh act))