hasql-generate-1.0.0: src/Hasql/Generate/Internal/Introspect.hs

module Hasql.Generate.Internal.Introspect
    ( ColumnInfo (..)
    , introspectColumns
    , introspectEnumLabels
    , introspectPrimaryKey
    ) where

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

import           Control.Exception         ( throwIO )
import           Control.Monad             ( return )

import           Data.Bool                 ( Bool (..), otherwise )
import qualified Data.ByteString           as BS
import qualified Data.ByteString.Char8     as BS8
import           Data.Eq                   ( (==) )
import           Data.Function             ( ($) )
import           Data.Functor              ( (<$>) )
import           Data.Maybe                ( Maybe (..), maybe )
import           Data.Semigroup            ( (<>) )
import           Data.String               ( String )

import qualified Database.PostgreSQL.LibPQ as PQ

import           Prelude
    ( Applicative (pure)
    , userError
    , (+)
    )

import           System.IO                 ( IO )

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

data ColumnInfo
    = ColumnInfo
      { colName       :: String
      , colPgSchema   :: String
      , colPgType     :: String
      , colIsEnum     :: Bool
      , colNotNull    :: Bool
      , colHasDefault :: Bool
      }

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

{-  Query @pg_catalog@ for all visible user columns of the given table, including
    both base types (@typtype = 'b'@) and enum types (@typtype = 'e'@), excluding
    composite, pseudo, and range types. Results are ordered by column position.
-}
introspectColumns :: PQ.Connection -> String -> String -> IO [ColumnInfo]
introspectColumns conn schema table = do
  result <- PQ.execParams conn columnSQL [textParam schema, textParam table] PQ.Text
  withQueryResult "column" result parseColumnRows

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

{-  Query @pg_catalog@ for the primary key column names of the given table,
    ordered by their position within the index.
-}
introspectPrimaryKey :: PQ.Connection -> String -> String -> IO [String]
introspectPrimaryKey conn schema table = do
  result <- PQ.execParams conn primaryKeySQL [textParam schema, textParam table] PQ.Text
  withQueryResult "primary key" result parsePkRows

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

{-  Query @pg_catalog@ for the labels of a PostgreSQL enum type in the given
    schema, ordered by their sort position (@enumsortorder@).
-}
introspectEnumLabels :: PQ.Connection -> String -> String -> IO [String]
introspectEnumLabels conn schema typeName = do
  result <- PQ.execParams conn enumLabelSQL [textParam schema, textParam typeName] PQ.Text
  withQueryResult "enum label" result parseEnumRows

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

columnSQL :: BS.ByteString
columnSQL =
  "SELECT tn.nspname, a.attname, t.typname, t.typtype, a.attnotnull, a.atthasdef \
  \FROM pg_catalog.pg_attribute a \
  \JOIN pg_catalog.pg_class c ON c.oid = a.attrelid \
  \JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \
  \JOIN pg_catalog.pg_type t ON t.oid = a.atttypid \
  \JOIN pg_catalog.pg_namespace tn ON tn.oid = t.typnamespace \
  \WHERE n.nspname = $1 \
  \  AND c.relname = $2 \
  \  AND a.attnum > 0 \
  \  AND NOT a.attisdropped \
  \  AND t.typtype IN ('b', 'e') \
  \  AND t.typcategory NOT IN ('C', 'P', 'X') \
  \ORDER BY a.attnum"

primaryKeySQL :: BS.ByteString
primaryKeySQL =
  "SELECT a.attname \
  \FROM pg_catalog.pg_index i \
  \JOIN pg_catalog.pg_class c ON c.oid = i.indrelid \
  \JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \
  \JOIN pg_catalog.pg_attribute a ON a.attrelid = i.indrelid \
  \  AND a.attnum = ANY(i.indkey) \
  \WHERE n.nspname = $1 \
  \  AND c.relname = $2 \
  \  AND i.indisprimary \
  \ORDER BY array_position(i.indkey, a.attnum)"

enumLabelSQL :: BS.ByteString
enumLabelSQL =
  "SELECT e.enumlabel \
  \FROM pg_catalog.pg_enum e \
  \JOIN pg_catalog.pg_type t ON t.oid = e.enumtypid \
  \JOIN pg_catalog.pg_namespace n ON n.oid = t.typnamespace \
  \WHERE n.nspname = $1 \
  \  AND t.typname = $2 \
  \ORDER BY e.enumsortorder"

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

parseColumnRows :: PQ.Result -> IO [ColumnInfo]
parseColumnRows res = do
  nRows <- PQ.ntuples res
  mapRows nRows $ \row -> do
    schemaVal <- PQ.getvalue res row 0
    nameVal <- PQ.getvalue res row 1
    typeVal <- PQ.getvalue res row 2
    typtypeVal <- PQ.getvalue res row 3
    notNullVal <- PQ.getvalue res row 4
    defVal <- PQ.getvalue res row 5
    let schema = maybe "" BS8.unpack schemaVal
        name = maybe "" BS8.unpack nameVal
        typ = maybe "" BS8.unpack typeVal
        isEnum = Just "e" == typtypeVal
        notNull = maybe False parseBool notNullVal
        hasDef = maybe False parseBool defVal
    return (ColumnInfo name schema typ isEnum notNull hasDef)

parsePkRows :: PQ.Result -> IO [String]
parsePkRows res = do
  nRows <- PQ.ntuples res
  mapRows nRows $ \row -> do
    nameVal <- PQ.getvalue res row 0
    return (maybe "" BS8.unpack nameVal)

parseEnumRows :: PQ.Result -> IO [String]
parseEnumRows res = do
  nRows <- PQ.ntuples res
  mapRows nRows $ \row -> do
    labelVal <- PQ.getvalue res row 0
    return (maybe "" BS8.unpack labelVal)

parseBool :: BS.ByteString -> Bool
parseBool bs
  | bs == "t" = True
  | otherwise = False

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

-- | Wrap a 'String' as a libpq text parameter (oid 25 = @text@).
textParam :: String -> Maybe (PQ.Oid, BS.ByteString, PQ.Format)
textParam str = Just (PQ.Oid 25, BS8.pack str, PQ.Text)

{-  Unwrap and validate a libpq query result, passing it to a row-parser on
    success. The @label@ argument names the introspection step for error
    messages.
-}
withQueryResult :: String -> Maybe PQ.Result -> (PQ.Result -> IO a) -> IO a
withQueryResult label result parse = case result of
  Nothing -> throwIO (userError ("hasql-generate: " <> label <> " introspection query returned no result"))
  Just res -> do
    status <- PQ.resultStatus res
    if status == PQ.TuplesOk
      then parse res
      else do
        msg <- maybe "unknown error" BS8.unpack <$> PQ.resultErrorMessage res
        throwIO (userError ("hasql-generate: " <> label <> " introspection failed: " <> msg))

-- |  Iterate over row indices [0 .. nRows-1] collecting results.
mapRows :: PQ.Row -> (PQ.Row -> IO a) -> IO [a]
mapRows nRows fn = go (PQ.Row 0)
  where
    go i
      | i == nRows = pure []
      | otherwise = do
          x <- fn i
          xs <- go (nextRow i)
          pure (x : xs)

    nextRow :: PQ.Row -> PQ.Row
    nextRow (PQ.Row n) = PQ.Row (n + 1)