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 :: forall a. ByteString -> (Connection -> IO a) -> IO a
withCompileTimeConnection ByteString
connStr =
  IO Connection
-> (Connection -> IO ()) -> (Connection -> IO a) -> IO a
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket IO Connection
acquire Connection -> IO ()
release
  where
    acquire :: IO PQ.Connection
    acquire :: IO Connection
acquire = do
      conn <- ByteString -> IO Connection
PQ.connectdb ByteString
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 :: Connection -> IO ()
release = Connection -> IO ()
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
      { ConnectionInfo -> Maybe [Char]
pgHost     :: Maybe String
      , ConnectionInfo -> Maybe [Char]
pgPort     :: Maybe String
      , ConnectionInfo -> Maybe [Char]
pgUser     :: Maybe String
      , ConnectionInfo -> Maybe [Char]
pgPassword :: Maybe String
      , ConnectionInfo -> Maybe [Char]
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 :: ConnectionInfo
def = ConnectionInfo
libpqDefaults

libpqDefaults :: ConnectionInfo
libpqDefaults :: ConnectionInfo
libpqDefaults = Maybe [Char]
-> Maybe [Char]
-> Maybe [Char]
-> Maybe [Char]
-> Maybe [Char]
-> ConnectionInfo
PgSimpleInfo Maybe [Char]
forall a. Maybe a
Nothing Maybe [Char]
forall a. Maybe a
Nothing Maybe [Char]
forall a. Maybe a
Nothing Maybe [Char]
forall a. Maybe a
Nothing Maybe [Char]
forall a. Maybe a
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 :: ConnectionInfo -> ByteString
toConnString (PgConnectionString [Char]
s) = [Char] -> ByteString
BS8.pack [Char]
s
toConnString (PgSimpleInfo Maybe [Char]
host Maybe [Char]
port Maybe [Char]
user Maybe [Char]
pass Maybe [Char]
db) =
  let pairs :: [[Char]]
pairs =
        [Maybe [Char]] -> [[Char]]
forall a. [Maybe a] -> [a]
catMaybes
          [ [Char] -> Maybe [Char] -> Maybe [Char]
kvPair [Char]
"host" Maybe [Char]
host
          , [Char] -> Maybe [Char] -> Maybe [Char]
kvPair [Char]
"port" Maybe [Char]
port
          , [Char] -> Maybe [Char] -> Maybe [Char]
kvPair [Char]
"user" Maybe [Char]
user
          , [Char] -> Maybe [Char] -> Maybe [Char]
kvPair [Char]
"password" Maybe [Char]
pass
          , [Char] -> Maybe [Char] -> Maybe [Char]
kvPair [Char]
"dbname" Maybe [Char]
db
          ]
   in if [[Char]] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [[Char]]
pairs
        then ByteString
BS.empty
        else [Char] -> ByteString
BS8.pack ([[Char]] -> [Char]
unwords [[Char]]
pairs)
  where
    kvPair :: String -> Maybe String -> Maybe String
    kvPair :: [Char] -> Maybe [Char] -> Maybe [Char]
kvPair [Char]
_key Maybe [Char]
Nothing   = Maybe [Char]
forall a. Maybe a
Nothing
    kvPair [Char]
key (Just [Char]
val) = [Char] -> Maybe [Char]
forall a. a -> Maybe a
Just ([Char]
key [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
"=" [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char] -> [Char]
escapeValue [Char]
val)

    escapeValue :: String -> String
    escapeValue :: [Char] -> [Char]
escapeValue [Char]
v
      | [Char] -> Bool
needsQuoting [Char]
v = [Char]
"'" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (Char -> [Char]) -> [Char] -> [Char]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Char -> [Char]
escapeChar [Char]
v [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"'"
      | Bool
otherwise = [Char]
v

    needsQuoting :: String -> Bool
    needsQuoting :: [Char] -> Bool
needsQuoting = (Char -> Bool) -> [Char] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\Char
c -> Char
c Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
' ' Bool -> Bool -> Bool
|| Char
c Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'\'' Bool -> Bool -> Bool
|| Char
c Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'\\')

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