<commit_before>{-# LANGUAGE FlexibleContexts  #-}
{-# LANGUAGE OverloadedStrings #-}

module Database.RethinkDB
    ( Handle
    , defaultPort, newHandle
    , run, nextChunk, collect, stop, wait

    , Error(..)

      -- * The Datum type
    , Datum(..), Array, Object, ToDatum(..), FromDatum(..)
    , (.=), (.:), (.:?), object

      -- The Exp type
    , Exp(..), SomeExp(..)
    , Bound(..), Order(..)
    , Sequence
    , Table, Database, SingleSelection
    , Res, Result, FromResponse
    , emptyOptions
    , lift
    , call1, call2

    , IsDatum, IsObject, IsSequence
    ) where


import           Data.Monoid      ((<>))

import           Data.Text        (Text)
import qualified Data.Text        as T

import qualified Data.Vector      as V
import qualified Data.Aeson.Types as A

import           Network.Socket   (Socket)
import           Data.IORef

import           Database.RethinkDB.Types
import           Database.RethinkDB.Types.Datum
import           Database.RethinkDB.Messages


------------------------------------------------------------------------------
-- Handle

data Handle = Handle
    { hSocket   :: !Socket
    , hTokenRef :: !(IORef Token)
    }


-- | The default port where RethinkDB accepts client driver connections.
defaultPort :: Int
defaultPort = 28015


-- | Create a new handle to the RethinkDB server.
newHandle :: Text -> Int -> Maybe Text -> IO Handle
newHandle host port mbAuth = do
    sock <- createSocket host port

    -- Do the handshake dance. Note that we currently ignore the reply and
    -- assume it is "SUCCESS".
    sendMessage sock (handshakeMessage mbAuth)
    _reply <- recvMessage sock handshakeReplyParser

    -- RethinkDB seems to expect the token to never be null. So we start with
    -- one and then count up.
    ref <- newIORef 1

    return $ Handle sock ref


-- | Start a new query and wait for its (first) result. If the result is an
-- single value ('Datum'), then three will be no further results. If it is
-- a sequence, then you must consume results until the sequence ends.
run :: (FromResponse (Result a))
    => Handle -> Exp a -> IO (Res a)
run handle expr = do
    _token <- start handle expr
    reply <- getResponse handle

    case reply of
        Left e -> return $ Left e
        Right response -> case responseType response of
            ClientErrorType  -> mkError response ClientError
            CompileErrorType -> mkError response CompileError
            RuntimeErrorType -> mkError response RuntimeError
            _                -> return $ parseMessage parseResponse response Right

parseMessage :: (a -> A.Parser b) -> a -> (b -> Either Error c) -> Either Error c
parseMessage parser value f = case A.parseEither parser value of
    Left  e -> Left $ ProtocolError $ T.pack e
    Right v -> f v

mkError :: Response -> (T.Text -> Error) -> IO (Either Error a)
mkError r e = return $ case V.toList (responseResult r) of
    [a] -> parseMessage A.parseJSON a (Left . e)
    _   -> Left $ ProtocolError $ "mkError: Could not parse error" <> T.pack (show (responseResult r))



-- | Collect all the values in a sequence and make them available as
-- a 'Vector a'.
collect :: (FromResponse (Sequence a))
        => Handle -> Sequence a -> IO (Either Error (V.Vector a))
collect _        (Done      x) = return $ Right x
collect handle s@(Partial _ x) = do
    chunk <- nextChunk handle s
    case chunk of
        Left e -> return $ Left e
        Right r -> do
            vals <- collect handle r
            case vals of
                Left ve -> return $ Left ve
                Right v -> return $ Right $ x <> v



-- | Start a new query. Returns the 'Token' which can be used to track its
-- progress.
start :: Handle -> Exp a -> IO Token
start handle term = do
    token <- atomicModifyIORef (hTokenRef handle) (\x -> (x + 1, x))
    sendMessage (hSocket handle) (queryMessage token msg)
    return token

  where
    msg = compileTerm $ do
        term'    <- toTerm term
        options' <- toTerm emptyOptions
        return $ A.Array $ V.fromList
            [ A.Number 1
            , term'
            , A.toJSON $ options'
            ]
<commit_msg>Less strict class prerequisites for collect<commit_after>{-# LANGUAGE FlexibleContexts  #-}
{-# LANGUAGE OverloadedStrings #-}

module Database.RethinkDB
    ( Handle
    , defaultPort, newHandle
    , run, nextChunk, collect, stop, wait

    , Error(..)

      -- * The Datum type
    , Datum(..), Array, Object, ToDatum(..), FromDatum(..)
    , (.=), (.:), (.:?), object

      -- The Exp type
    , Exp(..), SomeExp(..)
    , Bound(..), Order(..)
    , Sequence
    , Table, Database, SingleSelection
    , Res, Result, FromResponse
    , emptyOptions
    , lift
    , call1, call2

    , IsDatum, IsObject, IsSequence
    ) where


import           Data.Monoid      ((<>))

import           Data.Text        (Text)
import qualified Data.Text        as T

import qualified Data.Vector      as V
import qualified Data.Aeson.Types as A

import           Network.Socket   (Socket)
import           Data.IORef

import           Database.RethinkDB.Types
import           Database.RethinkDB.Types.Datum
import           Database.RethinkDB.Messages


------------------------------------------------------------------------------
-- Handle

data Handle = Handle
    { hSocket   :: !Socket
    , hTokenRef :: !(IORef Token)
    }


-- | The default port where RethinkDB accepts client driver connections.
defaultPort :: Int
defaultPort = 28015


-- | Create a new handle to the RethinkDB server.
newHandle :: Text -> Int -> Maybe Text -> IO Handle
newHandle host port mbAuth = do
    sock <- createSocket host port

    -- Do the handshake dance. Note that we currently ignore the reply and
    -- assume it is "SUCCESS".
    sendMessage sock (handshakeMessage mbAuth)
    _reply <- recvMessage sock handshakeReplyParser

    -- RethinkDB seems to expect the token to never be null. So we start with
    -- one and then count up.
    ref <- newIORef 1

    return $ Handle sock ref


-- | Start a new query and wait for its (first) result. If the result is an
-- single value ('Datum'), then three will be no further results. If it is
-- a sequence, then you must consume results until the sequence ends.
run :: (FromResponse (Result a))
    => Handle -> Exp a -> IO (Res a)
run handle expr = do
    _token <- start handle expr
    reply <- getResponse handle

    case reply of
        Left e -> return $ Left e
        Right response -> case responseType response of
            ClientErrorType  -> mkError response ClientError
            CompileErrorType -> mkError response CompileError
            RuntimeErrorType -> mkError response RuntimeError
            _                -> return $ parseMessage parseResponse response Right

parseMessage :: (a -> A.Parser b) -> a -> (b -> Either Error c) -> Either Error c
parseMessage parser value f = case A.parseEither parser value of
    Left  e -> Left $ ProtocolError $ T.pack e
    Right v -> f v

mkError :: Response -> (T.Text -> Error) -> IO (Either Error a)
mkError r e = return $ case V.toList (responseResult r) of
    [a] -> parseMessage A.parseJSON a (Left . e)
    _   -> Left $ ProtocolError $ "mkError: Could not parse error" <> T.pack (show (responseResult r))



-- | Collect all the values in a sequence and make them available as
-- a 'Vector a'.
collect :: (FromDatum a) => Handle -> Sequence a -> IO (Either Error (V.Vector a))
collect _        (Done      x) = return $ Right x
collect handle s@(Partial _ x) = do
    chunk <- nextChunk handle s
    case chunk of
        Left e -> return $ Left e
        Right r -> do
            vals <- collect handle r
            case vals of
                Left ve -> return $ Left ve
                Right v -> return $ Right $ x <> v



-- | Start a new query. Returns the 'Token' which can be used to track its
-- progress.
start :: Handle -> Exp a -> IO Token
start handle term = do
    token <- atomicModifyIORef (hTokenRef handle) (\x -> (x + 1, x))
    sendMessage (hSocket handle) (queryMessage token msg)
    return token

  where
    msg = compileTerm $ do
        term'    <- toTerm term
        options' <- toTerm emptyOptions
        return $ A.Array $ V.fromList
            [ A.Number 1
            , term'
            , A.toJSON $ options'
            ]

{'fraction_non_alphanumeric': 0.1159837047687515, 'fraction_numerical': 0.0023963575365444525, 'mean_word_length': 2.0231800072437522, 'pattern_counts': {'":': 0, '<': 29, '<?xml version=': 0, '>': 77, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2068568', 'n_tokens_mistral': 2386, 'n_tokens_neox': 2275, 'n_words': 924}