Connection.hs

module Hasql.Generate.Connection
    ( ConnectionInfo (..)
    , libpqDefaults
    , toConnString
    , withCompileTimeConnection
    ) where

----------------------------------------------------------------------------------------------------

import           Control.Exception         ( bracket, throwIO )

import           Data.Bool                 ( Bool (..), otherwise, (||) )
import qualified Data.ByteString           as BS
import qualified Data.ByteString.Char8     as BS8
import           Data.Char                 ( Char )
import           Data.Default.Class        ( Default (..) )
import           Data.Eq                   ( (/=), (==) )
import           Data.Functor              ( (<$>) )
import           Data.List
    ( any
    , concatMap
    , null
    , unwords
    , (++)
    )
import           Data.Maybe                ( Maybe (..), catMaybes, maybe )
import           Data.Semigroup            ( (<>) )
import           Data.String               ( String )

import qualified Database.PostgreSQL.LibPQ as PQ

import           Prelude                   ( Applicative (pure), userError )

import           System.IO                 ( IO )

----------------------------------------------------------------------------------------------------

{-  Acquire a libpq connection for compile-time schema introspection, run the
    provided action, and ensure the connection is closed afterward.

    The 'Data.ByteString.ByteString' argument is passed directly to
    @PQ.connectdb@. An empty string causes libpq to read connection
    parameters from standard PostgreSQL environment variables (@PGHOST@,
    @PGPORT@, @PGUSER@, @PGPASSWORD@, @PGDATABASE@, etc.).

    Throws an 'IOError' if the connection cannot be established.
-}
withCompileTimeConnection :: BS8.ByteString -> (PQ.Connection -> IO a) -> IO a
withCompileTimeConnection connStr =
  bracket acquire release
  where
    acquire :: IO PQ.Connection
    acquire = do
      conn <- PQ.connectdb connStr
      status <- PQ.status conn
      if status /= PQ.ConnectionOk
        then do
          msg <- maybe "unknown error" BS8.unpack <$> PQ.errorMessage conn
          PQ.finish conn
          throwIO (userError ("hasql-generate: compile-time DB connection failed: " <> msg))
        else
          pure conn

    release :: PQ.Connection -> IO ()
    release = PQ.finish

----------------------------------------------------------------------------------------------------

{-  Connection configuration for compile-time PostgreSQL introspection.

    @PgSimpleInfo@ lets you specify individual connection fields. Any field
    left as @Nothing@ is omitted from the connection string, so libpq falls
    back to its standard environment variables (@PGHOST@, @PGPORT@, etc.).

    @PgConnectionString@ passes a raw libpq connection string through as-is,
    useful for advanced options like @sslmode=verify-full@.
-}
data ConnectionInfo
    = PgSimpleInfo
      { pgHost     :: Maybe String
      , pgPort     :: Maybe String
      , pgUser     :: Maybe String
      , pgPassword :: Maybe String
      , pgDatabase :: Maybe String
      }
    | PgConnectionString String

{-  All @Nothing@ — produces an empty connection string so libpq reads
    standard environment variables.
-}
instance Default ConnectionInfo where
  def :: ConnectionInfo
  def = libpqDefaults

libpqDefaults :: ConnectionInfo
libpqDefaults = PgSimpleInfo Nothing Nothing Nothing Nothing Nothing

{-  Assemble a 'ConnectionInfo' into a libpq connection 'BS.ByteString'.

    For @PgSimpleInfo@, builds a @key=value@ string from the fields that
    are @Just@, omitting @Nothing@ fields. For @PgConnectionString@, passes
    the raw string through.
-}
toConnString :: ConnectionInfo -> BS.ByteString
toConnString (PgConnectionString s) = BS8.pack s
toConnString (PgSimpleInfo host port user pass db) =
  let pairs =
        catMaybes
          [ kvPair "host" host
          , kvPair "port" port
          , kvPair "user" user
          , kvPair "password" pass
          , kvPair "dbname" db
          ]
   in if null pairs
        then BS.empty
        else BS8.pack (unwords pairs)
  where
    kvPair :: String -> Maybe String -> Maybe String
    kvPair _key Nothing   = Nothing
    kvPair key (Just val) = Just (key <> "=" <> escapeValue val)

    escapeValue :: String -> String
    escapeValue v
      | needsQuoting v = "'" ++ concatMap escapeChar v ++ "'"
      | otherwise = v

    needsQuoting :: String -> Bool
    needsQuoting = any (\c -> c == ' ' || c == '\'' || c == '\\')

    escapeChar :: Char -> String
    escapeChar '\'' = "\\'"
    escapeChar '\\' = "\\\\"
    escapeChar c    = [c]