opium/lib/Database/PostgreSQL/Opium/FromField.hs

{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeApplications #-}

module Database.PostgreSQL.Opium.FromField
  ( -- * Decoding data from @libpq@
    FromField (..)
  , fromField
    -- * Utility types
  , RawField (..)
  ) where

import Data.Attoparsec.ByteString (Parser)
import Data.Bits (Bits (..))
import Data.ByteString (ByteString)
import Data.Functor (($>))
import Data.Int (Int32)
import Data.Proxy (Proxy (..))
import Data.Time
  ( Day (..)
  , DiffTime
  , TimeOfDay
  , UTCTime (..)
  , addDays
  , fromGregorian
  , picosecondsToDiffTime
  , timeToTimeOfDay
  , toGregorian
  )
import Data.Text (Text)
import Data.Word (Word32, Word64)
import Database.PostgreSQL.LibPQ (Oid)
import Unsafe.Coerce (unsafeCoerce)

import qualified Data.Attoparsec.ByteString as AP
import qualified Data.ByteString as BS
import qualified Data.Text as Text
import qualified Data.Text.Encoding as Encoding

import qualified Database.PostgreSQL.Opium.Oid as Oid

(\/) :: (a -> Bool) -> (a -> Bool) -> a -> Bool
p \/ q = \x -> p x || q x

fromField :: FromField a => ByteString -> Either String a
fromField =
  AP.parseOnly parseField

class FromField a where
  validOid :: Proxy a -> Oid -> Bool
  parseField :: Parser a

-- | See https://www.postgresql.org/docs/current/datatype-binary.html.
-- Accepts @bytea@.
instance FromField ByteString where
  validOid Proxy = Oid.bytea
  parseField = AP.takeByteString

-- | See https://www.postgresql.org/docs/current/datatype-character.html.
-- Accepts @text@, @character@ and @character varying@.
instance FromField Text where
  validOid Proxy = Oid.text \/ Oid.character \/ Oid.characterVarying
  parseField = Encoding.decodeUtf8 <$> AP.takeByteString

-- Accepts @text@, @character@ and @character varying@.
-- | See https://www.postgresql.org/docs/current/datatype-character.html.
instance FromField String where
  validOid Proxy = validOid @Text Proxy
  parseField = Text.unpack <$> parseField

-- | See https://www.postgresql.org/docs/current/datatype-character.html.
-- This instance accepts all character types but fails to decode fields that are not exactly one character.
instance FromField Char where
  validOid Proxy = validOid @Text Proxy
  parseField = do
    str <- parseField
    case str of
      [c] -> pure c
      _ -> fail "Char accepts single characters only"

intParser :: (Bits a, Num a) => Parser a
intParser = readBigEndian <$> AP.takeByteString
  where
    readBigEndian = BS.foldl' (\x b -> x `shiftL` 8 .|. fromIntegral b) 0

-- | See https://www.postgresql.org/docs/current/datatype-numeric.html.
-- We assume that 'Int' has 64 bits. This is not guaranteed but reasonable enough.
instance FromField Int where
  validOid Proxy = Oid.smallint \/ Oid.integer \/ Oid.bigint
  parseField = intParser

-- | See https://www.postgresql.org/docs/current/datatype-numeric.html.
instance FromField Integer where
  validOid Proxy = Oid.smallint \/ Oid.integer \/ Oid.bigint
  parseField = intParser

instance FromField Word where
  validOid Proxy = Oid.smallint \/ Oid.integer \/ Oid.bigint
  parseField = intParser

-- | See https://www.postgresql.org/docs/current/datatype-numeric.html.
-- Accepts only @real@ fields, not @double precision@.
instance FromField Float where
  validOid Proxy = Oid.real
  -- Afaict there's no cleaner (@base@) way to access the underlying bits.
  -- In C we'd do
  --
  -- union { float a; uint32_t b; } x;
  -- x.b = ...;
  -- return x.a;
  parseField = unsafeCoerce <$> intParser @Word32

-- | See https://www.postgresql.org/docs/current/datatype-numeric.html.
-- Accepts only @double precision@ fields, not @real@.
instance FromField Double where
  validOid Proxy = Oid.doublePrecision
  parseField = unsafeCoerce <$> intParser @Word64

boolParser :: Parser Bool
boolParser = AP.choice
  [ AP.word8 1 $> True
  , AP.word8 0 $> False
  ]

-- | See https://www.postgresql.org/docs/current/datatype-boolean.html.
instance FromField Bool where
  validOid Proxy = Oid.boolean
  parseField = boolParser

postgresEpoch :: Day
postgresEpoch = fromGregorian 2000 1 1

fromPostgresJulian :: Integer -> Day
fromPostgresJulian x
  | year <= 0 = fromGregorian (year - 1) month dayOfMonth
  | otherwise = day
  where
    day = addDays (fromIntegral x) postgresEpoch
    (year, month, dayOfMonth) = toGregorian day

-- | See https://www.postgresql.org/docs/current/datatype-datetime.html.
-- Relevant as well: https://git.postgresql.org/gitweb/?p=postgresql.git;a=blob;f=src/backend/utils/adt/datetime.c;h=267dfd37b2e8b9bc63797c69b9ca2e45e6bfde61;hb=HEAD#l267.
instance FromField Day where
  validOid Proxy = Oid.date
  parseField = fromPostgresJulian . fromIntegral <$> intParser @Int32

-- | See https://www.postgresql.org/docs/current/datatype-datetime.html.
-- Binary format: https://git.postgresql.org/gitweb/?p=postgresql.git;a=blob;f=src/backend/utils/adt/date.c;h=ae0f24de2c3c54eb6d0405cdb212597c2407238e;hb=HEAD#l1542.
-- Accepts @time@.
instance FromField DiffTime where
  validOid Proxy = Oid.time
  parseField = microsecondsToDiffTime <$> intParser
    where
      microsecondsToDiffTime :: Integer -> DiffTime
      microsecondsToDiffTime ms = picosecondsToDiffTime $ ms * 1000000

-- | See https://www.postgresql.org/docs/current/datatype-datetime.html.
-- Binary format: https://git.postgresql.org/gitweb/?p=postgresql.git;a=blob;f=src/backend/utils/adt/date.c;h=ae0f24de2c3c54eb6d0405cdb212597c2407238e;hb=HEAD#l1542.
-- Accepts @time@.
instance FromField TimeOfDay where
  validOid Proxy = Oid.time
  parseField = timeToTimeOfDay <$> parseField @DiffTime

instance FromField UTCTime where
  validOid Proxy = Oid.timestamp
  parseField = fromPostgresTimestamp <$> intParser @Int
    where
      fromPostgresTimestamp :: Int -> UTCTime
      fromPostgresTimestamp ts =
        let
          (days, microseconds) = ts `divMod` (86400 * 1000000)
          day = fromPostgresJulian $ fromIntegral days
          time = picosecondsToDiffTime $ fromIntegral microseconds * 1000000
        in
          UTCTime day time

newtype RawField a = RawField a
  deriving (Eq, Show)

instance FromField a => FromField (RawField a) where
  validOid Proxy = const True
  parseField = RawField <$> parseField