SelectAPI.hs

--
-- MinIO Haskell SDK, (C) 2017-2019 MinIO, Inc.
--
-- Licensed under the Apache License, Version 2.0 (the "License");
-- you may not use this file except in compliance with the License.
-- You may obtain a copy of the License at
--
--     http://www.apache.org/licenses/LICENSE-2.0
--
-- Unless required by applicable law or agreed to in writing, software
-- distributed under the License is distributed on an "AS IS" BASIS,
-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-- See the License for the specific language governing permissions and
-- limitations under the License.
--

module Network.Minio.SelectAPI
  ( -- | The `selectObjectContent` allows querying CSV, JSON or Parquet
    -- format objects in AWS S3 and in MinIO using SQL Select
    -- statements. This allows significant reduction of data transfer
    -- from object storage for computation-intensive tasks, as relevant
    -- data is filtered close to the storage.
    selectObjectContent,
    SelectRequest,
    selectRequest,

    -- *** Input Serialization
    InputSerialization,
    defaultCsvInput,
    linesJsonInput,
    documentJsonInput,
    defaultParquetInput,
    setInputCSVProps,
    CompressionType (..),
    setInputCompressionType,

    -- *** CSV Format details

    -- | CSV format options such as delimiters and quote characters are
    -- specified using using the functions below. Options are combined
    -- monoidally.
    CSVProp,
    recordDelimiter,
    fieldDelimiter,
    quoteCharacter,
    quoteEscapeCharacter,
    commentCharacter,
    allowQuotedRecordDelimiter,
    FileHeaderInfo (..),
    fileHeaderInfo,
    QuoteFields (..),
    quoteFields,

    -- *** Output Serialization
    OutputSerialization,
    defaultCsvOutput,
    defaultJsonOutput,
    outputCSVFromProps,
    outputJSONFromRecordDelimiter,

    -- *** Progress messages
    setRequestProgressEnabled,

    -- *** Interpreting Select output

    -- | The conduit returned by `selectObjectContent` returns values of
    -- the `EventMessage` data type. This returns the query output
    -- messages formatted according to the chosen output serialization,
    -- interleaved with progress messages (if enabled by
    -- `setRequestProgressEnabled`), and at the end a statistics
    -- message.
    --
    -- If the application is interested in only the payload, then
    -- `getPayloadBytes` can be used. For example to simply print the
    -- payload to stdout:
    --
    -- > resultConduit <- selectObjectContent bucket object mySelectRequest
    -- > runConduit $ resultConduit .| getPayloadBytes .| stdoutC
    --
    -- Note that runConduit, the connect operator (.|) and stdoutC are
    -- all from the "conduit" package.
    getPayloadBytes,
    EventMessage (..),
    Progress (..),
    Stats,
  )
where

import Conduit ((.|))
import qualified Conduit as C
import qualified Data.Binary as Bin
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as LB
import Data.Digest.CRC32 (crc32, crc32Update)
import Lib.Prelude
import qualified Network.HTTP.Conduit as NC
import qualified Network.HTTP.Types as HT
import Network.Minio.API
import Network.Minio.Data
import Network.Minio.Errors
import Network.Minio.Utils
import Network.Minio.XmlGenerator
import Network.Minio.XmlParser
import UnliftIO (MonadUnliftIO)

data EventStreamException
  = ESEPreludeCRCFailed
  | ESEMessageCRCFailed
  | ESEUnexpectedEndOfStream
  | ESEDecodeFail [Char]
  | ESEInvalidHeaderType
  | ESEInvalidHeaderValueType
  | ESEInvalidMessageType
  deriving stock (Eq, Show)

instance Exception EventStreamException

-- chunkSize in bytes is 32KiB
chunkSize :: Int
chunkSize = 32 * 1024

parseBinary :: Bin.Binary a => ByteString -> IO a
parseBinary b = do
  case Bin.decodeOrFail $ LB.fromStrict b of
    Left (_, _, msg) -> throwIO $ ESEDecodeFail msg
    Right (_, _, r) -> return r

bytesToHeaderName :: Text -> IO MsgHeaderName
bytesToHeaderName t = case t of
  ":message-type" -> return MessageType
  ":event-type" -> return EventType
  ":content-type" -> return ContentType
  ":error-code" -> return ErrorCode
  ":error-message" -> return ErrorMessage
  _ -> throwIO ESEInvalidHeaderType

parseHeaders ::
  MonadUnliftIO m =>
  Word32 ->
  C.ConduitM ByteString a m [MessageHeader]
parseHeaders 0 = return []
parseHeaders hdrLen = do
  bs1 <- readNBytes 1
  n :: Word8 <- liftIO $ parseBinary bs1

  headerKeyBytes <- readNBytes $ fromIntegral n
  let headerKey = decodeUtf8Lenient headerKeyBytes
  headerName <- liftIO $ bytesToHeaderName headerKey

  bs2 <- readNBytes 1
  headerValueType :: Word8 <- liftIO $ parseBinary bs2
  when (headerValueType /= 7) $ throwIO ESEInvalidHeaderValueType

  bs3 <- readNBytes 2
  vLen :: Word16 <- liftIO $ parseBinary bs3
  headerValueBytes <- readNBytes $ fromIntegral vLen
  let headerValue = decodeUtf8Lenient headerValueBytes
      m = (headerName, headerValue)
      k = 1 + fromIntegral n + 1 + 2 + fromIntegral vLen

  ms <- parseHeaders (hdrLen - k)
  return (m : ms)

-- readNBytes returns N bytes read from the string and throws an
-- exception if N bytes are not present on the stream.
readNBytes :: MonadUnliftIO m => Int -> C.ConduitM ByteString a m ByteString
readNBytes n = do
  b <- LB.toStrict <$> (C.takeCE n .| C.sinkLazy)
  if B.length b /= n
    then throwIO ESEUnexpectedEndOfStream
    else return b

crcCheck ::
  MonadUnliftIO m =>
  C.ConduitM ByteString ByteString m ()
crcCheck = do
  b <- readNBytes 12
  n :: Word32 <- liftIO $ parseBinary $ B.take 4 b
  preludeCRC :: Word32 <- liftIO $ parseBinary $ B.drop 8 b
  when (crc32 (B.take 8 b) /= preludeCRC) $
    throwIO ESEPreludeCRCFailed

  -- we do not yield the checksum
  C.yield $ B.take 8 b

  -- 12 bytes have been read off the current message. Now read the
  -- next (n-12)-4 bytes and accumulate the checksum, and yield it.
  let startCrc = crc32 b
  finalCrc <- accumulateYield (fromIntegral n - 16) startCrc

  bs <- readNBytes 4
  expectedCrc :: Word32 <- liftIO $ parseBinary bs

  when (finalCrc /= expectedCrc) $
    throwIO ESEMessageCRCFailed

  -- we unconditionally recurse - downstream figures out when to
  -- quit reading the stream
  crcCheck
  where
    accumulateYield n checkSum = do
      let toRead = min n chunkSize
      b <- readNBytes toRead
      let c' = crc32Update checkSum b
          n' = n - B.length b
      C.yield b
      if n' > 0
        then accumulateYield n' c'
        else return c'

handleMessage :: MonadUnliftIO m => C.ConduitT ByteString EventMessage m ()
handleMessage = do
  b1 <- readNBytes 4
  msgLen :: Word32 <- liftIO $ parseBinary b1

  b2 <- readNBytes 4
  hdrLen :: Word32 <- liftIO $ parseBinary b2

  hs <- parseHeaders hdrLen

  let payloadLen = msgLen - hdrLen - 16
      getHdrVal h = fmap snd . find ((h ==) . fst)
      eventHdrValue = getHdrVal EventType hs
      msgHdrValue = getHdrVal MessageType hs
      errCode = getHdrVal ErrorCode hs
      errMsg = getHdrVal ErrorMessage hs

  case msgHdrValue of
    Just "event" -> do
      case eventHdrValue of
        Just "Records" -> passThrough $ fromIntegral payloadLen
        Just "Cont" -> return ()
        Just "Progress" -> do
          bs <- readNBytes $ fromIntegral payloadLen
          progress <- parseSelectProgress bs
          C.yield $ ProgressEventMessage progress
        Just "Stats" -> do
          bs <- readNBytes $ fromIntegral payloadLen
          stats <- parseSelectProgress bs
          C.yield $ StatsEventMessage stats
        Just "End" -> return ()
        _ -> throwIO ESEInvalidMessageType
      when (eventHdrValue /= Just "End") handleMessage
    Just "error" -> do
      let reqMsgMay = RequestLevelErrorMessage <$> errCode <*> errMsg
      maybe (throwIO ESEInvalidMessageType) C.yield reqMsgMay
    _ -> throwIO ESEInvalidMessageType
  where
    passThrough 0 = return ()
    passThrough n = do
      let c = min n chunkSize
      b <- readNBytes c
      C.yield $ RecordPayloadEventMessage b
      passThrough $ n - B.length b

selectProtoConduit ::
  MonadUnliftIO m =>
  C.ConduitT ByteString EventMessage m ()
selectProtoConduit = crcCheck .| handleMessage

-- | selectObjectContent calls the SelectRequest on the given
-- object. It returns a Conduit of event messages that can be consumed
-- by the client.
selectObjectContent ::
  Bucket ->
  Object ->
  SelectRequest ->
  Minio (C.ConduitT () EventMessage Minio ())
selectObjectContent b o r = do
  let reqInfo =
        defaultS3ReqInfo
          { riMethod = HT.methodPost,
            riBucket = Just b,
            riObject = Just o,
            riPayload = PayloadBS $ mkSelectRequest r,
            riNeedsLocation = False,
            riQueryParams = [("select", Nothing), ("select-type", Just "2")]
          }
  -- print $ mkSelectRequest r
  resp <- mkStreamRequest reqInfo
  return $ NC.responseBody resp .| selectProtoConduit

-- | A helper conduit that returns only the record payload bytes.
getPayloadBytes :: MonadIO m => C.ConduitT EventMessage ByteString m ()
getPayloadBytes = do
  evM <- C.await
  case evM of
    Just v -> do
      case v of
        RecordPayloadEventMessage b -> C.yield b
        RequestLevelErrorMessage c m -> liftIO $ throwIO $ SelectErr c m
        _ -> return ()
      getPayloadBytes
    Nothing -> return ()