{-# LANGUAGE DataKinds #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeOperators #-}

module Server.Dump (DumpAPI, dumpHandler) where

import Control.Monad.IO.Class (liftIO)
import qualified Data.ByteString.Lazy.Char8 as BL
import Data.OpenApi (NamedSchema (NamedSchema), ToSchema (declareNamedSchema))
import Data.OpenApi.ParamSchema (binarySchema)
import Network.HTTP.Media ((//))
import Servant
import Servant.Types.SourceT
import System.Environment (getEnv)
import System.Process
    ( CreateProcess (env, std_out)
    , StdStream (CreatePipe)
    , createProcess
    , proc
    )

newtype SQLBytes
    = SQLBytes
    { SQLBytes -> ByteString
unSQLBytes :: BL.ByteString
    }

instance ToSchema SQLBytes where
    declareNamedSchema :: Proxy SQLBytes -> Declare (Definitions Schema) NamedSchema
declareNamedSchema Proxy SQLBytes
_ = NamedSchema -> Declare (Definitions Schema) NamedSchema
forall a. a -> DeclareT (Definitions Schema) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (NamedSchema -> Declare (Definitions Schema) NamedSchema)
-> NamedSchema -> Declare (Definitions Schema) NamedSchema
forall a b. (a -> b) -> a -> b
$ Maybe Text -> Schema -> NamedSchema
NamedSchema (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"SQL Bytes") Schema
binarySchema

instance ToSourceIO BL.ByteString SQLBytes where
    toSourceIO :: SQLBytes -> SourceIO ByteString
toSourceIO (SQLBytes ByteString
bs) = StepT Identity ByteString -> SourceIO ByteString
forall (f :: * -> *) a (m :: * -> *).
Foldable f =>
f a -> SourceT m a
source (StepT Identity ByteString -> SourceIO ByteString)
-> StepT Identity ByteString -> SourceIO ByteString
forall a b. (a -> b) -> a -> b
$ ByteString
-> StepT Identity ByteString -> StepT Identity ByteString
forall (m :: * -> *) a. a -> StepT m a -> StepT m a
Yield ByteString
bs StepT Identity ByteString
forall (m :: * -> *) a. StepT m a
Stop

data SQL

instance Accept SQL where
    contentType :: Proxy SQL -> MediaType
contentType Proxy SQL
_ = ByteString
"text" ByteString -> ByteString -> MediaType
// ByteString
"plain"

instance MimeRender SQL SQLBytes where
    mimeRender :: Proxy SQL -> SQLBytes -> ByteString
mimeRender Proxy SQL
_ = SQLBytes -> ByteString
unSQLBytes

instance MimeRender SQL BL.ByteString where
    mimeRender :: Proxy SQL -> ByteString -> ByteString
mimeRender Proxy SQL
_ = ByteString -> ByteString
forall a. a -> a
id

type DumpAPI =
    "dump"
        :> StreamGet
            NewlineFraming
            SQL
            SQLBytes

dumpHandler
    :: Handler SQLBytes
dumpHandler :: Handler SQLBytes
dumpHandler = IO SQLBytes -> Handler SQLBytes
forall a. IO a -> Handler a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO SQLBytes -> Handler SQLBytes)
-> IO SQLBytes -> Handler SQLBytes
forall a b. (a -> b) -> a -> b
$ do
    String
host <- String -> IO String
getEnv String
"POSTGRES_HOST"
    String
port <- String -> IO String
getEnv String
"POSTGRES_PORT"
    String
user <- String -> IO String
getEnv String
"POSTGRES_USER"
    String
password <- String -> IO String
getEnv String
"POSTGRES_PASSWORD"
    String
db <- String -> IO String
getEnv String
"POSTGRES_DB"

    (Maybe Handle
_, Maybe Handle
mHout, Maybe Handle
_, ProcessHandle
_) <-
        CreateProcess
-> IO (Maybe Handle, Maybe Handle, Maybe Handle, ProcessHandle)
createProcess
            (String -> [String] -> CreateProcess
proc String
"pg_dump" [String
"-h", String
host, String
"-p", String
port, String
"-U", String
user, String
"-d", String
db])
                { std_out = CreatePipe
                , env = Just [("PGPASSWORD", password)]
                }

    Handle
hout <- case Maybe Handle
mHout of
        Just Handle
h -> Handle -> IO Handle
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Handle
h
        Maybe Handle
Nothing -> String -> IO Handle
forall a. HasCallStack => String -> a
error String
"Failed to create stdout pipe for pg_dump"

    ByteString
dump <- Handle -> IO ByteString
BL.hGetContents Handle
hout

    SQLBytes -> IO SQLBytes
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (SQLBytes -> IO SQLBytes) -> SQLBytes -> IO SQLBytes
forall a b. (a -> b) -> a -> b
$ ByteString -> SQLBytes
SQLBytes ByteString
dump