-- | Internal module which implements starting a @hipsql-server@ from Haskell code.
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
module Hipsql.Server.Internal
  ( -- * Disclaimer
    -- $disclaimer

    -- ** Internals
    module Hipsql.Server.Internal
  ) where

import Control.Concurrent (MVar, newEmptyMVar, putMVar, takeMVar)
import Control.Concurrent.Async (race_)
import Control.Exception (Exception, SomeException, catch, fromException, throwIO)
import Control.Monad (mfilter)
import Control.Monad.Except (ExceptT(ExceptT))
import Control.Monad.IO.Class (MonadIO(liftIO))
import Data.ByteString (ByteString)
import Data.IORef (IORef, atomicModifyIORef', newIORef, readIORef)
import Data.Maybe (fromMaybe, listToMaybe)
import Data.String (IsString(fromString))
import Data.Traversable (for)
import GHC.Stack (SrcLoc, prettySrcLoc)
import Hipsql.API (HipsqlRoutes(HipsqlRoutes, eval, getVersion), HipsqlAPI, theHipsqlAPI, theHipsqlApiVersion)
import Hipsql.API.Internal (lookupHipsqlPort)
import Servant.Server
  ( Handler(Handler), HasServer(ServerT), ServerError(errBody), Application, Server, err400, err500
  , hoistServer, serve
  )
import Servant.Server.Generic (genericServerT)
import System.IO (hPutStrLn, stderr)
import qualified Data.ByteString.Char8 as Char8
import qualified Data.ByteString.Lazy as Lazy
import qualified Data.ByteString.Lazy.Char8 as Lazy.Char8
import qualified Data.List as List
import qualified Database.PostgreSQL.LibPQ as LibPQ
import qualified Network.Wai.Handler.Warp as Warp

-- | Runtime environment of the @hipsql@ server.
data ServerEnv = ServerEnv
  { ServerEnv -> Connection
conn :: LibPQ.Connection
  , ServerEnv -> MVar ()
killswitch :: MVar ()
  , ServerEnv -> IORef ServerState
state :: IORef ServerState
  }

-- | Runtime state of the @hipsql@ server.
newtype ServerState = ServerState
  { ServerState -> Bool
extendedDisplay :: Bool
  }

-- | Create an initial 'ServerEnv'.
newServerEnv :: LibPQ.Connection -> IO ServerEnv
newServerEnv :: Connection -> IO ServerEnv
newServerEnv conn :: Connection
conn = do
  MVar ()
killswitch <- IO (MVar ())
forall a. IO (MVar a)
newEmptyMVar
  IORef ServerState
state <- ServerState -> IO (IORef ServerState)
forall a. a -> IO (IORef a)
newIORef ServerState :: Bool -> ServerState
ServerState
    { extendedDisplay :: Bool
extendedDisplay = Bool
False
    }
  ServerEnv -> IO ServerEnv
forall (f :: * -> *) a. Applicative f => a -> f a
pure ServerEnv :: Connection -> MVar () -> IORef ServerState -> ServerEnv
ServerEnv { Connection
conn :: Connection
conn :: Connection
conn, MVar ()
killswitch :: MVar ()
killswitch :: MVar ()
killswitch, IORef ServerState
state :: IORef ServerState
state :: IORef ServerState
state }

-- | Lifts the 'IO' action to a @servant@ 'Handler', ensuring exceptions
-- are handled accordingly.
toHandler :: IO a -> Handler a
toHandler :: IO a -> Handler a
toHandler = ExceptT ServerError IO a -> Handler a
forall a. ExceptT ServerError IO a -> Handler a
Handler (ExceptT ServerError IO a -> Handler a)
-> (IO a -> ExceptT ServerError IO a) -> IO a -> Handler a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO (Either ServerError a) -> ExceptT ServerError IO a
forall e (m :: * -> *) a. m (Either e a) -> ExceptT e m a
ExceptT (IO (Either ServerError a) -> ExceptT ServerError IO a)
-> (IO a -> IO (Either ServerError a))
-> IO a
-> ExceptT ServerError IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO a -> IO (Either ServerError a)
forall a. IO a -> IO (Either ServerError a)
handleServantErr
  where
  handleServantErr :: IO a -> IO (Either ServerError a)
  handleServantErr :: IO a -> IO (Either ServerError a)
handleServantErr x :: IO a
x =
    IO (Either ServerError a)
-> (SomeException -> IO (Either ServerError a))
-> IO (Either ServerError a)
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
catch
      ((a -> Either ServerError a) -> IO a -> IO (Either ServerError a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> Either ServerError a
forall a b. b -> Either a b
Right IO a
x)
      (Either ServerError a -> IO (Either ServerError a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either ServerError a -> IO (Either ServerError a))
-> (SomeException -> Either ServerError a)
-> SomeException
-> IO (Either ServerError a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ServerError -> Either ServerError a
forall a b. a -> Either a b
Left (ServerError -> Either ServerError a)
-> (SomeException -> ServerError)
-> SomeException
-> Either ServerError a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SomeException -> ServerError
toServantError)

  toServantError :: SomeException -> ServerError
  toServantError :: SomeException -> ServerError
toServantError e :: SomeException
e = case SomeException -> Maybe QueryError
forall e. Exception e => SomeException -> Maybe e
fromException SomeException
e of
    Just (QueryError m :: ByteString
m) -> ServerError
err400 { errBody :: ByteString
errBody = ByteString -> ByteString
Lazy.Char8.fromStrict ByteString
m }
    Nothing -> ServerError
err500 { errBody :: ByteString
errBody = String -> ByteString
forall a. IsString a => String -> a
fromString (String -> ByteString) -> String -> ByteString
forall a b. (a -> b) -> a -> b
$ SomeException -> String
forall a. Show a => a -> String
show SomeException
e }

-- | Create a @servant@ 'Server' of the 'HipsqlAPI'.
server :: ServerEnv -> Server HipsqlAPI
server :: ServerEnv -> Server HipsqlAPI
server env :: ServerEnv
env = Proxy
  ((Summary "Gets the current server version"
    :> ("version" :> Get '[JSON] Version))
   :<|> (Summary "Evaluate a psql expression"
         :> ("eval"
             :> (ReqBody '[OctetStream] ByteString
                 :> Post '[OctetStream] ByteString))))
-> (forall x. IO x -> Handler x)
-> ServerT
     ((Summary "Gets the current server version"
       :> ("version" :> Get '[JSON] Version))
      :<|> (Summary "Evaluate a psql expression"
            :> ("eval"
                :> (ReqBody '[OctetStream] ByteString
                    :> Post '[OctetStream] ByteString))))
     IO
-> ServerT
     ((Summary "Gets the current server version"
       :> ("version" :> Get '[JSON] Version))
      :<|> (Summary "Evaluate a psql expression"
            :> ("eval"
                :> (ReqBody '[OctetStream] ByteString
                    :> Post '[OctetStream] ByteString))))
     Handler
forall api (m :: * -> *) (n :: * -> *).
HasServer api '[] =>
Proxy api
-> (forall x. m x -> n x) -> ServerT api m -> ServerT api n
hoistServer Proxy
  ((Summary "Gets the current server version"
    :> ("version" :> Get '[JSON] Version))
   :<|> (Summary "Evaluate a psql expression"
         :> ("eval"
             :> (ReqBody '[OctetStream] ByteString
                 :> Post '[OctetStream] ByteString))))
Proxy HipsqlAPI
theHipsqlAPI forall x. IO x -> Handler x
toHandler ServerT
  ((Summary "Gets the current server version"
    :> ("version" :> Get '[JSON] Version))
   :<|> (Summary "Evaluate a psql expression"
         :> ("eval"
             :> (ReqBody '[OctetStream] ByteString
                 :> Post '[OctetStream] ByteString))))
  IO
ServerT HipsqlAPI IO
ioServer
  where
  ioServer :: ServerT HipsqlAPI IO
  ioServer :: ServerT HipsqlAPI IO
ioServer = HipsqlRoutes (AsServerT IO)
-> ToServant HipsqlRoutes (AsServerT IO)
forall (routes :: * -> *) (m :: * -> *).
GenericServant routes (AsServerT m) =>
routes (AsServerT m) -> ToServant routes (AsServerT m)
genericServerT HipsqlRoutes :: forall route.
(route
 :- (Summary "Gets the current server version"
     :> ("version" :> Get '[JSON] Version)))
-> (route
    :- (Summary "Evaluate a psql expression"
        :> ("eval"
            :> (ReqBody '[OctetStream] ByteString
                :> Post '[OctetStream] ByteString))))
-> HipsqlRoutes route
HipsqlRoutes
    { getVersion :: AsServerT IO
:- (Summary "Gets the current server version"
    :> ("version" :> Get '[JSON] Version))
getVersion = Version -> IO Version
forall (f :: * -> *) a. Applicative f => a -> f a
pure Version
theHipsqlApiVersion
    , AsServerT IO
:- (Summary "Evaluate a psql expression"
    :> ("eval"
        :> (ReqBody '[OctetStream] ByteString
            :> Post '[OctetStream] ByteString)))
ByteString -> IO ByteString
eval :: ByteString -> IO ByteString
eval :: AsServerT IO
:- (Summary "Evaluate a psql expression"
    :> ("eval"
        :> (ReqBody '[OctetStream] ByteString
            :> Post '[OctetStream] ByteString)))
eval
    }

  eval :: ByteString -> IO ByteString
eval input :: ByteString
input = case ByteString -> Maybe (Char, ByteString)
Lazy.Char8.uncons ByteString
input of
    Just ('\\', c :: ByteString
c) -> ByteString -> IO ByteString
runCommand ByteString
c
    _ | ByteString
input ByteString -> [ByteString] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` ["quit", "exit"] -> IO ByteString
quit
    _ -> ByteString -> IO ByteString
runQuery ByteString
input

  quit :: IO ByteString
quit = do
    MVar () -> () -> IO ()
forall a. MVar a -> a -> IO ()
putMVar (ServerEnv -> MVar ()
killswitch ServerEnv
env) ()
    IO ByteString
forall a. Monoid a => a
mempty

  runCommand :: ByteString -> IO ByteString
runCommand = \case
    "x" -> IO ByteString
runToggleExtendedDisplay
    "?" -> IO ByteString
runHelp
    s :: ByteString
s | ByteString
s ByteString -> [ByteString] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` ["q", "quit"] -> IO ByteString
quit
    s :: ByteString
s -> ByteString -> IO ByteString
forall (f :: * -> *) a.
(Applicative f, Semigroup a, IsString a) =>
a -> f a
invalidCommand ByteString
s

  runHelp :: IO ByteString
runHelp = ByteString -> IO ByteString
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
Lazy.fromStrict ByteString
helpMessage

  invalidCommand :: a -> f a
invalidCommand s :: a
s = do
    a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> f a) -> a -> f a
forall a b. (a -> b) -> a -> b
$ "invalid command \\" a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
s a -> a -> a
forall a. Semigroup a => a -> a -> a
<> "\nTry \\? for help."

  runToggleExtendedDisplay :: IO ByteString
runToggleExtendedDisplay = do
    Bool
x <- IO Bool
toggleExtendedDisplay
    ByteString -> IO ByteString
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ "Extended display is " ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> (if Bool
x then "on" else "off") ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> "."

  modify :: (ServerState -> ServerState) -> IO ServerState
modify f :: ServerState -> ServerState
f = do
    IORef ServerState
-> (ServerState -> (ServerState, ServerState)) -> IO ServerState
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' (ServerEnv -> IORef ServerState
state ServerEnv
env) \s :: ServerState
s -> let s' :: ServerState
s' = ServerState -> ServerState
f ServerState
s in (ServerState
s', ServerState
s')

  toggleExtendedDisplay :: IO Bool
toggleExtendedDisplay = do
    ServerState
s <- (ServerState -> ServerState) -> IO ServerState
modify \s :: ServerState
s -> ServerState
s { extendedDisplay :: Bool
extendedDisplay = Bool -> Bool
not (ServerState -> Bool
extendedDisplay ServerState
s) }
    Bool -> IO Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> IO Bool) -> Bool -> IO Bool
forall a b. (a -> b) -> a -> b
$ ServerState -> Bool
extendedDisplay ServerState
s

  runQuery :: ByteString -> IO ByteString
runQuery q :: ByteString
q = do
    if ByteString -> Bool
Lazy.Char8.null ByteString
q then do
      IO ByteString
forall a. Monoid a => a
mempty
    else do
      ServerEnv -> QueryResponse -> IO ByteString
renderQueryResult ServerEnv
env (QueryResponse -> IO ByteString)
-> IO QueryResponse -> IO ByteString
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ServerEnv -> ByteString -> IO QueryResponse
rawQuery ServerEnv
env (ByteString -> ByteString
Lazy.toStrict ByteString
q)

-- | Render a 'QueryResponse' received from issuing a query as human
-- readable. Respects the current 'extendedDisplay' state.
renderQueryResult :: ServerEnv -> QueryResponse -> IO Lazy.ByteString
renderQueryResult :: ServerEnv -> QueryResponse -> IO ByteString
renderQueryResult env :: ServerEnv
env QueryResponse { [Maybe ByteString]
columnNames :: QueryResponse -> [Maybe ByteString]
columnNames :: [Maybe ByteString]
columnNames, [[Maybe ByteString]]
resultRows :: QueryResponse -> [[Maybe ByteString]]
resultRows :: [[Maybe ByteString]]
resultRows } = do
  ServerState { Bool
extendedDisplay :: Bool
extendedDisplay :: ServerState -> Bool
extendedDisplay } <- IORef ServerState -> IO ServerState
forall a. IORef a -> IO a
readIORef (ServerEnv -> IORef ServerState
state ServerEnv
env)
  let rendered :: ByteString
rendered =
        if Bool
extendedDisplay then
          [ByteString] -> [[ByteString]] -> ByteString
renderXTable [ByteString]
renderedColNames [[ByteString]]
renderedValues
        else
          Maybe [ByteString] -> [[ByteString]] -> ByteString
renderTable ([ByteString] -> Maybe [ByteString]
forall a. a -> Maybe a
Just [ByteString]
renderedColNames) [[ByteString]]
renderedValues
  ByteString -> IO ByteString
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ ByteString
rendered ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> "\n" ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
renderedRowCount ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> "\n"
  where
  renderedColNames :: [ByteString]
renderedColNames = (Maybe ByteString -> ByteString)
-> [Maybe ByteString] -> [ByteString]
forall a b. (a -> b) -> [a] -> [b]
map (ByteString -> Maybe ByteString -> ByteString
forall a. a -> Maybe a -> a
fromMaybe "?") [Maybe ByteString]
columnNames
  renderedValues :: [[ByteString]]
renderedValues = ([Maybe ByteString] -> [ByteString])
-> [[Maybe ByteString]] -> [[ByteString]]
forall a b. (a -> b) -> [a] -> [b]
map ((Maybe ByteString -> ByteString)
-> [Maybe ByteString] -> [ByteString]
forall a b. (a -> b) -> [a] -> [b]
map (ByteString -> Maybe ByteString -> ByteString
forall a. a -> Maybe a -> a
fromMaybe "null")) [[Maybe ByteString]]
resultRows
  renderedRowCount :: ByteString
renderedRowCount = case [[Maybe ByteString]] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Maybe ByteString]]
resultRows of
    1 -> "(1 row)"
    n :: Int
n -> "(" ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> String -> ByteString
forall a. IsString a => String -> a
fromString (Int -> String
forall a. Show a => a -> String
show Int
n) ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> " rows)"

-- | Special case of 'renderTable' which renders the supplied header
-- and values via 'extendedDisplay'.
renderXTable :: [ByteString] -> [[ByteString]] -> Lazy.ByteString
renderXTable :: [ByteString] -> [[ByteString]] -> ByteString
renderXTable hs :: [ByteString]
hs = ByteString -> [ByteString] -> ByteString
Lazy.Char8.intercalate "\n" ([ByteString] -> ByteString)
-> ([[ByteString]] -> [ByteString]) -> [[ByteString]] -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> [ByteString] -> ByteString)
-> [Int] -> [[ByteString]] -> [ByteString]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> [ByteString] -> ByteString
go [1..]
  where
  go :: Int -> [ByteString] -> Lazy.ByteString
  go :: Int -> [ByteString] -> ByteString
go i :: Int
i rs :: [ByteString]
rs =
       "-[ RECORD " ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> String -> ByteString
forall a. IsString a => String -> a
fromString (Int -> String
forall a. Show a => a -> String
show Int
i) ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> " ]\n"
    ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> Maybe [ByteString] -> [[ByteString]] -> ByteString
renderTable Maybe [ByteString]
forall a. Maybe a
Nothing ([ByteString] -> [ByteString] -> [[ByteString]]
forall a. [a] -> [a] -> [[a]]
zipL [ByteString]
hs [ByteString]
rs)

  zipL :: [a] -> [a] -> [[a]]
  zipL :: [a] -> [a] -> [[a]]
zipL = (a -> a -> [a]) -> [a] -> [a] -> [[a]]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith \a1 :: a
a1 a2 :: a
a2 -> [a
a1, a
a2]

-- | Render the supplied optional header and values as a table for printing in
-- the console. Attempts to match closely to real psql output.
renderTable :: Maybe [ByteString] -> [[ByteString]] -> Lazy.ByteString
renderTable :: Maybe [ByteString] -> [[ByteString]] -> ByteString
renderTable maybeHeader :: Maybe [ByteString]
maybeHeader rows :: [[ByteString]]
rows = ByteString
renderedHeader ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
renderedTable
  where
  maxLens :: [Int]
maxLens =
    ([ByteString] -> Int) -> [[ByteString]] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map
      (\col :: [ByteString]
col -> if [ByteString] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [ByteString]
col then 0 else [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum ((ByteString -> Int) -> [ByteString] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map ByteString -> Int
Char8.length [ByteString]
col))
      ([[ByteString]] -> [[ByteString]]
forall a. [[a]] -> [[a]]
List.transpose ([ByteString] -> Maybe [ByteString] -> [ByteString]
forall a. a -> Maybe a -> a
fromMaybe [] Maybe [ByteString]
maybeHeader [ByteString] -> [[ByteString]] -> [[ByteString]]
forall a. a -> [a] -> [a]
: [[ByteString]]
rows))

  -- We don't care about the last column since it never needs to be padded,
  -- so 'safeInit' to omit it.
  maxPadAt :: Int -> Int
  maxPadAt :: Int -> Int
maxPadAt i :: Int
i = Int -> (Int -> Int) -> Maybe Int -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe 0 (1Int -> Int -> Int
forall a. Num a => a -> a -> a
+) (Maybe Int -> Int) -> Maybe Int -> Int
forall a b. (a -> b) -> a -> b
$ [Int] -> Maybe Int
forall a. [a] -> Maybe a
listToMaybe ([Int] -> Maybe Int) -> [Int] -> Maybe Int
forall a b. (a -> b) -> a -> b
$ Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
drop Int
i ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ [Int] -> [Int]
forall a. [a] -> [a]
safeInit [Int]
maxLens

  lineLenAt :: Int -> Int
  lineLenAt :: Int -> Int
lineLenAt i :: Int
i = Int -> (Int -> Int) -> Maybe Int -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe 0 (2Int -> Int -> Int
forall a. Num a => a -> a -> a
+) (Maybe Int -> Int) -> Maybe Int -> Int
forall a b. (a -> b) -> a -> b
$ [Int] -> Maybe Int
forall a. [a] -> Maybe a
listToMaybe ([Int] -> Maybe Int) -> [Int] -> Maybe Int
forall a b. (a -> b) -> a -> b
$ Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
drop Int
i [Int]
maxLens

  renderCell :: Int -> ByteString -> ByteString
renderCell i :: Int
i s :: ByteString
s =
    " " ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString -> ByteString
Lazy.Char8.fromStrict ByteString
s ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<>
      Int64 -> Char -> ByteString
Lazy.Char8.replicate
        (Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Int
maxPadAt Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- ByteString -> Int
Char8.length ByteString
s))
        ' '

  renderRow :: [ByteString] -> [ByteString]
renderRow = (Int -> ByteString -> ByteString)
-> [Int] -> [ByteString] -> [ByteString]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> ByteString -> ByteString
renderCell [0..]

  renderedHeader :: ByteString
renderedHeader = case Maybe [ByteString]
maybeHeader of
    Nothing -> ""
    Just header :: [ByteString]
header ->
      ByteString -> [ByteString] -> ByteString
Lazy.Char8.intercalate "|" ([ByteString] -> [ByteString]
renderRow [ByteString]
header) ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> "\n"
        ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString -> [ByteString] -> ByteString
Lazy.Char8.intercalate "+"
            ( (Int -> ByteString) -> [Int] -> [ByteString]
forall a b. (a -> b) -> [a] -> [b]
map (\i :: Int
i -> Int64 -> Char -> ByteString
Lazy.Char8.replicate (Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Int
lineLenAt Int
i)) '-') [0..([ByteString] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ByteString]
header Int -> Int -> Int
forall a. Num a => a -> a -> a
- 1)]
            ) ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> "\n"

  renderedTable :: ByteString
renderedTable =
    ByteString -> [ByteString] -> ByteString
Lazy.Char8.intercalate "\n"
      ([ByteString] -> ByteString) -> [ByteString] -> ByteString
forall a b. (a -> b) -> a -> b
$ ([ByteString] -> ByteString) -> [[ByteString]] -> [ByteString]
forall a b. (a -> b) -> [a] -> [b]
map (ByteString -> [ByteString] -> ByteString
Lazy.Char8.intercalate "|" ([ByteString] -> ByteString)
-> ([ByteString] -> [ByteString]) -> [ByteString] -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [ByteString] -> [ByteString]
renderRow) [[ByteString]]
rows

-- | Same as 'init' except returns an empty list if the supplied list is empty.
safeInit :: [a] -> [a]
safeInit :: [a] -> [a]
safeInit xs :: [a]
xs = if [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [a]
xs then [a]
xs else [a] -> [a]
forall a. [a] -> [a]
init [a]
xs

-- | Execute the supplied query and return its response.
rawQuery :: ServerEnv -> ByteString -> IO QueryResponse
rawQuery :: ServerEnv -> ByteString -> IO QueryResponse
rawQuery env :: ServerEnv
env q :: ByteString
q = do
  IO QueryResponse -> IO QueryResponse
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO do
    Connection -> ByteString -> IO (Maybe Result)
LibPQ.exec (ServerEnv -> Connection
conn ServerEnv
env) ByteString
q IO (Maybe Result)
-> (Maybe Result -> IO QueryResponse) -> IO QueryResponse
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      Just result :: Result
result -> do
        Maybe ByteString
maybeMessage <- Result -> IO (Maybe ByteString)
LibPQ.resultErrorMessage Result
result
        case (ByteString -> Bool) -> Maybe ByteString -> Maybe ByteString
forall (m :: * -> *) a. MonadPlus m => (a -> Bool) -> m a -> m a
mfilter (Bool -> Bool
not (Bool -> Bool) -> (ByteString -> Bool) -> ByteString -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Bool
Char8.null) Maybe ByteString
maybeMessage of
          Nothing -> Result -> IO QueryResponse
mkQueryResponse Result
result
          Just message :: ByteString
message -> QueryError -> IO QueryResponse
forall e a. Exception e => e -> IO a
throwIO (QueryError -> IO QueryResponse) -> QueryError -> IO QueryResponse
forall a b. (a -> b) -> a -> b
$ ByteString -> QueryError
QueryError ByteString
message
      Nothing -> do
        Maybe ByteString
maybeMessage <- Connection -> IO (Maybe ByteString)
LibPQ.errorMessage (ServerEnv -> Connection
conn ServerEnv
env)
        QueryError -> IO QueryResponse
forall e a. Exception e => e -> IO a
throwIO (QueryError -> IO QueryResponse) -> QueryError -> IO QueryResponse
forall a b. (a -> b) -> a -> b
$ ByteString -> QueryError
QueryError (ByteString -> QueryError) -> ByteString -> QueryError
forall a b. (a -> b) -> a -> b
$ ByteString -> Maybe ByteString -> ByteString
forall a. a -> Maybe a -> a
fromMaybe "ERROR: (no message)" Maybe ByteString
maybeMessage

-- | Construct a 'QueryResponse' from a 'LibPQ.Result'.
-- NOTE: This is dangerous as it does not validate that an
-- error did not occur. Prefer 'rawQuery' instead.
mkQueryResponse :: LibPQ.Result -> IO QueryResponse
mkQueryResponse :: Result -> IO QueryResponse
mkQueryResponse result :: Result
result = do
  Row
numRows <- Result -> IO Row
LibPQ.ntuples Result
result
  Column
numCols <- Result -> IO Column
LibPQ.nfields Result
result
  [Maybe ByteString]
columnNames <-
    [Column]
-> (Column -> IO (Maybe ByteString)) -> IO [Maybe ByteString]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for [0..(Column
numCols Column -> Column -> Column
forall a. Num a => a -> a -> a
- 1)] \c :: Column
c -> do
      Result -> Column -> IO (Maybe ByteString)
LibPQ.fname Result
result Column
c
  [[Maybe ByteString]]
resultRows <-
    [Row] -> (Row -> IO [Maybe ByteString]) -> IO [[Maybe ByteString]]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for [0..(Row
numRows Row -> Row -> Row
forall a. Num a => a -> a -> a
- 1)] \r :: Row
r -> do
      [Column]
-> (Column -> IO (Maybe ByteString)) -> IO [Maybe ByteString]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for [0..(Column
numCols Column -> Column -> Column
forall a. Num a => a -> a -> a
- 1)] \c :: Column
c -> do
        Result -> Row -> Column -> IO (Maybe ByteString)
LibPQ.getvalue' Result
result Row
r Column
c
  QueryResponse -> IO QueryResponse
forall (f :: * -> *) a. Applicative f => a -> f a
pure QueryResponse :: [Maybe ByteString] -> [[Maybe ByteString]] -> QueryResponse
QueryResponse
    { [Maybe ByteString]
columnNames :: [Maybe ByteString]
columnNames :: [Maybe ByteString]
columnNames
    , [[Maybe ByteString]]
resultRows :: [[Maybe ByteString]]
resultRows :: [[Maybe ByteString]]
resultRows
    }

-- | An error response from issuing a query.
newtype QueryError = QueryError ByteString
  deriving stock (Int -> QueryError -> ShowS
[QueryError] -> ShowS
QueryError -> String
(Int -> QueryError -> ShowS)
-> (QueryError -> String)
-> ([QueryError] -> ShowS)
-> Show QueryError
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [QueryError] -> ShowS
$cshowList :: [QueryError] -> ShowS
show :: QueryError -> String
$cshow :: QueryError -> String
showsPrec :: Int -> QueryError -> ShowS
$cshowsPrec :: Int -> QueryError -> ShowS
Show)
  deriving anyclass (Show QueryError
Typeable QueryError
(Typeable QueryError, Show QueryError) =>
(QueryError -> SomeException)
-> (SomeException -> Maybe QueryError)
-> (QueryError -> String)
-> Exception QueryError
SomeException -> Maybe QueryError
QueryError -> String
QueryError -> SomeException
forall e.
(Typeable e, Show e) =>
(e -> SomeException)
-> (SomeException -> Maybe e) -> (e -> String) -> Exception e
displayException :: QueryError -> String
$cdisplayException :: QueryError -> String
fromException :: SomeException -> Maybe QueryError
$cfromException :: SomeException -> Maybe QueryError
toException :: QueryError -> SomeException
$ctoException :: QueryError -> SomeException
$cp2Exception :: Show QueryError
$cp1Exception :: Typeable QueryError
Exception)

-- | A successful response from issuing a query.
data QueryResponse = QueryResponse
  { QueryResponse -> [Maybe ByteString]
columnNames :: [Maybe ByteString]
  , QueryResponse -> [[Maybe ByteString]]
resultRows :: [[Maybe ByteString]]
  }

-- | Help message displayed when the user gives the @\?@ command.
helpMessage :: ByteString
helpMessage :: ByteString
helpMessage =
  ByteString -> [ByteString] -> ByteString
Char8.intercalate "\n"
    [ "General"
    , "  \\q        quit psql"
    , "Help"
    , "  \\?        show this help message"
    , ""
    , "Formatting"
    , "  \\x        toggle expanded output"
    ]

-- | A @servant@ 'Application' of the 'HipsqlAPI'.
application :: ServerEnv -> Application
application :: ServerEnv -> Application
application env :: ServerEnv
env = Proxy
  ((Summary "Gets the current server version"
    :> ("version" :> Get '[JSON] Version))
   :<|> (Summary "Evaluate a psql expression"
         :> ("eval"
             :> (ReqBody '[OctetStream] ByteString
                 :> Post '[OctetStream] ByteString))))
-> ServerT
     ((Summary "Gets the current server version"
       :> ("version" :> Get '[JSON] Version))
      :<|> (Summary "Evaluate a psql expression"
            :> ("eval"
                :> (ReqBody '[OctetStream] ByteString
                    :> Post '[OctetStream] ByteString))))
     Handler
-> Application
forall api.
HasServer api '[] =>
Proxy api -> Server api -> Application
serve Proxy
  ((Summary "Gets the current server version"
    :> ("version" :> Get '[JSON] Version))
   :<|> (Summary "Evaluate a psql expression"
         :> ("eval"
             :> (ReqBody '[OctetStream] ByteString
                 :> Post '[OctetStream] ByteString))))
Proxy HipsqlAPI
theHipsqlAPI (ServerEnv -> Server HipsqlAPI
server ServerEnv
env)

-- | Start a @hipsql@ session with the given 'LibPQ.Connection'.
startHipsql :: Maybe SrcLoc -> LibPQ.Connection -> IO ()
startHipsql :: Maybe SrcLoc -> Connection -> IO ()
startHipsql loc :: Maybe SrcLoc
loc conn :: Connection
conn = do
  Config
config <- IO Config
getDefaultConfig
  Deps
deps <- IO Deps
getDefaultDeps
  Maybe SrcLoc -> Config -> Deps -> Connection -> IO ()
startHipsql' Maybe SrcLoc
loc Config
config Deps
deps Connection
conn

-- | Same as 'startHipsql' but allows you to specify the 'Config' and 'Deps' directly.
startHipsql' :: Maybe SrcLoc -> Config -> Deps -> LibPQ.Connection -> IO ()
startHipsql' :: Maybe SrcLoc -> Config -> Deps -> Connection -> IO ()
startHipsql' loc :: Maybe SrcLoc
loc config :: Config
config deps :: Deps
deps conn :: Connection
conn = do
  ServerEnv
env <- Connection -> IO ServerEnv
newServerEnv Connection
conn
  String -> IO ()
logger (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$
    "Starting hipsql server on port "
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
port
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> "; called at "
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String -> (SrcLoc -> String) -> Maybe SrcLoc -> String
forall b a. b -> (a -> b) -> Maybe a -> b
maybe "<unknown>" SrcLoc -> String
prettySrcLoc Maybe SrcLoc
loc
  IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO ()
race_ (ServerEnv -> IO ()
waitForKillswitch ServerEnv
env) (Int -> Application -> IO ()
Warp.run Int
port (ServerEnv -> Application
application ServerEnv
env))
  where
  Config { Int
port :: Config -> Int
port :: Int
port } = Config
config

  Deps { String -> IO ()
logger :: Deps -> String -> IO ()
logger :: String -> IO ()
logger } = Deps
deps

  waitForKillswitch :: ServerEnv -> IO ()
waitForKillswitch env :: ServerEnv
env = do
    MVar () -> IO ()
forall a. MVar a -> IO a
takeMVar (ServerEnv -> MVar ()
killswitch ServerEnv
env)
    String -> IO ()
logger "Shutting down hipsql server"

-- | Same as 'startHipsql' except uses a 'LibPQ.Connection' acquiring function.
-- Useful when integrating with libraries like @postgresql-simple@ which
-- give you exclusive access to the 'LibPQ.Connection' via such a function.
startHipsqlWith :: Maybe SrcLoc -> ((LibPQ.Connection -> IO ()) -> IO ()) -> IO ()
startHipsqlWith :: Maybe SrcLoc -> ((Connection -> IO ()) -> IO ()) -> IO ()
startHipsqlWith loc :: Maybe SrcLoc
loc f :: (Connection -> IO ()) -> IO ()
f = (Connection -> IO ()) -> IO ()
f (Maybe SrcLoc -> Connection -> IO ()
startHipsql Maybe SrcLoc
loc)

-- | Same as 'startHipsqlWith' but allows you to specify the 'Config' and 'Deps' directly.
startHipsqlWith' :: Maybe SrcLoc -> Config -> Deps -> ((LibPQ.Connection -> IO ()) -> IO ()) -> IO ()
startHipsqlWith' :: Maybe SrcLoc
-> Config -> Deps -> ((Connection -> IO ()) -> IO ()) -> IO ()
startHipsqlWith' loc :: Maybe SrcLoc
loc config :: Config
config deps :: Deps
deps f :: (Connection -> IO ()) -> IO ()
f = (Connection -> IO ()) -> IO ()
f (Maybe SrcLoc -> Config -> Deps -> Connection -> IO ()
startHipsql' Maybe SrcLoc
loc Config
config Deps
deps)

-- | Configuration required for starting a @hipsql@ server.
newtype Config = Config
  { Config -> Int
port :: Int
  }

-- | Gets the default 'Config' used by the @hipsql@ server.
getDefaultConfig :: IO Config
getDefaultConfig :: IO Config
getDefaultConfig = do
  IO (Either String Int)
lookupHipsqlPort IO (Either String Int)
-> (Either String Int -> IO Config) -> IO Config
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Left message :: String
message -> do
      String -> IO Config
forall a. HasCallStack => String -> a
error (String -> IO Config) -> String -> IO Config
forall a b. (a -> b) -> a -> b
$ "Failed to start hipsql server; could not parse port: " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
message
    Right port :: Int
port -> do
      Config -> IO Config
forall (f :: * -> *) a. Applicative f => a -> f a
pure Config :: Int -> Config
Config
        { Int
port :: Int
port :: Int
port
        }

-- | Dependencies required for starting a @hipsql@ server.
newtype Deps = Deps
  { Deps -> String -> IO ()
logger :: String -> IO ()
  }

-- | Gets the default 'Deps' used by the @hipsql@ server.
getDefaultDeps :: IO Deps
getDefaultDeps :: IO Deps
getDefaultDeps = do
  Deps -> IO Deps
forall (f :: * -> *) a. Applicative f => a -> f a
pure Deps :: (String -> IO ()) -> Deps
Deps
    { logger :: String -> IO ()
logger = Handle -> String -> IO ()
hPutStrLn Handle
stderr
    }

-- $disclaimer
--
-- Changes to this module will not be reflected in the library's version
-- updates.