diff --git a/Data/Conduit/ElasticSearch.hs b/Data/Conduit/ElasticSearch.hs --- a/Data/Conduit/ElasticSearch.hs +++ b/Data/Conduit/ElasticSearch.hs @@ -1,28 +1,109 @@ +{-# LANGUAGE RankNTypes #-} {-| This module exports "Conduit" interfaces to ElasticSearch. It is totally experimental. -} -module Data.Conduit.ElasticSearch (esConduit) where +module Data.Conduit.ElasticSearch (esConduit, esSearchSource) where import Prelude hiding (catch) +import Data.Maybe (fromMaybe) import Control.Exception import Data.Conduit import qualified Data.Conduit.List as CL import Network.HTTP.Conduit import Data.Aeson +import Control.Applicative import qualified Data.ByteString as BS import qualified Data.ByteString.Lazy.Char8 as BSL import Data.Time import qualified Data.Text.Lazy.Encoding as E import Data.Text.Format (format,left) import Logstash.Message +import Control.Monad import Control.Monad.IO.Class import Control.Concurrent (threadDelay) import Data.Either (partitionEithers) +import Data.Monoid +import Network.HTTP.Types import qualified Data.HashMap.Strict as HM import qualified Data.Vector as V import qualified Data.Text as T +data SearchObject = SearchObject { _source :: LogstashMessage + , _index :: T.Text + , _score :: Double + , _id :: BS.ByteString + , _type :: T.Text + } deriving (Show) + +instance FromJSON SearchObject where + parseJSON (Object v) = SearchObject + <$> v .: "_source" + <*> v .: "_index" + <*> v .: "_score" + <*> v .: "_id" + <*> v .: "_type" + parseJSON _ = mzero + +data Hits = Hits { _rhits :: [SearchObject] + , _maxScore :: Double + , _total :: Int + } deriving (Show) + +instance FromJSON Hits where + parseJSON (Object v) = Hits + <$> v .: "hits" + <*> v .: "max_score" + <*> v .: "total" + parseJSON _ = mzero + +data SearchResponse = SearchResponse { _timedOut :: Bool + , _hits :: Hits + , _took :: Integer + } deriving (Show) + +instance FromJSON SearchResponse where + parseJSON (Object v) = SearchResponse + <$> v .: "timed_out" + <*> v .: "hits" + <*> v .: "took" + parseJSON _ = mzero + +esSearchSource :: (MonadResource m) => Maybe (Request m) -- ^ Defaults parameters for the http request to ElasticSearch. Use "Nothing" for defaults. + -> BS.ByteString -- ^ Hostname of the ElasticSearch server + -> Int -- ^ Port of the HTTP interface (usually 9200) + -> BS.ByteString -- ^ Prefix of the index (usually logstash) + -> Value -- ^ Request + -> Int -- ^ Maximum size of each response + -> Int -- ^ start + -> Producer m (Either Value [LogstashMessage]) +esSearchSource r h p prefix req maxsize start = self start + where + defR1 = fromMaybe def r + defR2 = defR1 { host = h + , port = p + , path = prefix <> "/_search" + , method = "POST" + , checkStatus = \_ _ _ -> Nothing + } + self :: (MonadResource m) => Int -> Producer m (Either Value [LogstashMessage]) + self i = do + let body = encode $ object [ ("query", req), "from" .= i, "size" .= maxsize ] + request = defR2 { requestBody = RequestBodyLBS body } + resp <- liftIO (safeQuery request) + let respbody = decode (responseBody resp) :: Maybe SearchResponse + errbody = fromMaybe Null (decode (responseBody resp)) :: Value + code = statusCode $ responseStatus resp + case (code, respbody, errbody) of + (200, Just x, _) -> do + let hits = _hits x + total = _total hits + messages = map _source (_rhits hits) + nexti = i + maxsize + yield (Right messages) + when (nexti < total) (self nexti) + _ -> yield (Left errbody) + safeQuery :: Request (ResourceT IO) -> IO (Response BSL.ByteString) safeQuery req = catch (withManager $ httpLbs req) (\e -> print (e :: SomeException) >> threadDelay 500000 >> safeQuery req) @@ -32,17 +113,16 @@ esConduit :: (MonadResource m) => Maybe (Request m) -- ^ Defaults parameters for the http request to ElasticSearch. Use "Nothing" for defaults. -> BS.ByteString -- ^ Hostname of the ElasticSearch server -> Int -- ^ Port of the HTTP interface (usually 9200) + -> T.Text -- ^ Prefix of the index (usually logstash) -> Conduit [LogstashMessage] m [Either (LogstashMessage, Value) Value] -esConduit r h p = CL.map (map prepareBS) =$= CL.mapM sendBulk +esConduit r h p prefix = CL.map (map prepareBS) =$= CL.mapM sendBulk where - defR1 = case r of - Just x -> x - Nothing -> def + defR1 = fromMaybe def r defR2 = defR1 { host = h , port = p , path = "/_bulk" , method = "POST" - , checkStatus = (\_ _ -> Nothing) + , checkStatus = (\_ _ _ -> Nothing) } prepareBS :: LogstashMessage -> Either (LogstashMessage, Value) (LogstashMessage, Value, Value) prepareBS input = @@ -50,7 +130,7 @@ Nothing -> Left (input, object [ "error" .= String "Time was not supplied" ]) Just (UTCTime day _) -> let (y,m,d) = toGregorian day - index = BSL.toStrict (E.encodeUtf8 (format "logstash-{}.{}.{}" (y, left 2 '0' m, left 2 '0' d))) + index = BSL.toStrict (E.encodeUtf8 (format "{}-{}.{}.{}" (prefix, y, left 2 '0' m, left 2 '0' d))) in Right (input, object [ "index" .= object [ "_index" .= index, "_type" .= logstashType input ] ], toJSON input) sendBulk :: (MonadResource m) => [Either (LogstashMessage, Value) (LogstashMessage, Value, Value)] -> m [Either (LogstashMessage, Value) Value] sendBulk input = @@ -59,8 +139,8 @@ body = BSL.unlines $ concatMap (\(_,x,y) -> [encode x, encode y]) tosend req = defR2 { requestBody = RequestBodyLBS body } in do - res <- fmap (responseBody) $ liftIO $ safeQuery req - let genericError er = return (Left (emptyLSMessage "error", object [ "error" .= (T.pack er), "data" .= res ]) : lerrors) + res <- fmap responseBody $ liftIO $ safeQuery req + let genericError er = return (Left (emptyLSMessage "error", object [ "error" .= T.pack er, "data" .= res ]) : lerrors) getObject st (Object hh) = HM.lookup st hh getObject _ _ = Nothing fst3 (a,_,_) = a diff --git a/Data/Conduit/Misc.hs b/Data/Conduit/Misc.hs --- a/Data/Conduit/Misc.hs +++ b/Data/Conduit/Misc.hs @@ -13,9 +13,9 @@ concatFlush :: (Monad m) => Integer -> Conduit [a] m (Flush a) concatFlush mx = concatFlush' 0 where - concatFlush' x = awaitE >>= either return (\input -> if null input - then yield Flush >> concatFlush' 0 - else foldM sendf x input >>= concatFlush' + concatFlush' x = await >>= maybe (return ()) (\input -> if null input + then yield Flush >> concatFlush' 0 + else foldM sendf x input >>= concatFlush' ) sendf curx ev = do yield (Chunk ev) @@ -28,7 +28,7 @@ groupFlush :: (Monad m) => Conduit (Flush a) m [a] groupFlush = grouper [] where - grouper lst = awaitE >>= either return (handle lst) + grouper lst = await >>= maybe (return ()) (handle lst) handle lst Flush = do unless (null lst) (yield (reverse lst)) grouper [] diff --git a/Data/Conduit/Network/Retry.hs b/Data/Conduit/Network/Retry.hs --- a/Data/Conduit/Network/Retry.hs +++ b/Data/Conduit/Network/Retry.hs @@ -23,7 +23,7 @@ This is used to send a full JSON object to Logstash. -} -sinkSocketRetry :: MonadResource m => IO Socket -> Int -> IO () -> GInfSink ByteString m +sinkSocketRetry :: MonadResource m => IO Socket -> Int -> IO () -> Consumer ByteString m () sinkSocketRetry mkSocket delay exeptionCallback = let safeMkSocket :: IO Socket @@ -36,11 +36,11 @@ safeMkSocket >>= putMVar s threadDelay delay safeSend s o - push :: MonadResource m => MVar Socket -> GInfSink ByteString m - push s = awaitE >>= either return (\bs -> lift (liftIO $ safeSend s bs) >> push s) + push :: MonadResource m => MVar Socket -> Consumer ByteString m () + push s = await >>= maybe (return ()) (\bs -> lift (liftIO $ safeSend s bs) >> push s) in bracketP (safeMkSocket >>= newMVar) (takeMVar >=> close) push -- | A specialization of the previous Sink that opens a TCP connection. -tcpSinkRetry :: MonadResource m => ByteString -> Int -> Int -> IO () -> GInfSink ByteString m +tcpSinkRetry :: MonadResource m => ByteString -> Int -> Int -> IO () -> Consumer ByteString m () tcpSinkRetry host port = sinkSocketRetry (fmap fst (getSocket host port)) diff --git a/Logstash/Message.hs b/Logstash/Message.hs --- a/Logstash/Message.hs +++ b/Logstash/Message.hs @@ -123,3 +123,7 @@ Nothing -> do curtime <- getCurrentTime return msg { logstashTime = Just curtime } + +addLogstashTag :: T.Text -> LogstashMessage -> LogstashMessage +addLogstashTag tag msg = msg { logstashTags = tag : logstashTags msg } + diff --git a/hslogstash.cabal b/hslogstash.cabal --- a/hslogstash.cabal +++ b/hslogstash.cabal @@ -2,9 +2,9 @@ -- documentation, see http://haskell.org/cabal/users-guide/ name: hslogstash -version: 0.2.3 +version: 0.2.3.1 synopsis: A library to work with, or as, a logstash server -description: This library contains a few modules that let you work with Logstash messages, read them from a Redis list, store them into Elasticsearch, and more. +description: This library contains a few modules that let you work with Logstash messages, read them from a Redis list, store them in or get them from Elasticsearch, and more. license: BSD3 license-file: LICENSE author: Simon Marechal @@ -23,4 +23,4 @@ extensions: OverloadedStrings, BangPatterns ghc-options: -Wall -- other-modules: - build-depends: base <5, aeson, network, bytestring, text, vector, unordered-containers, time, text-format, attoparsec, hedis, conduit, transformers, http-conduit, iconv, network-conduit, containers, parallel-io, stm-conduit + build-depends: base <5, aeson, network, bytestring, text, vector, unordered-containers, time, text-format, attoparsec, hedis, conduit >= 1.0.0, transformers, http-conduit, iconv, network-conduit, containers, parallel-io, stm-conduit, http-types