Merge pull request #84 from Soostone/request-hook

Add Request hook

bloodhound.cabal

@@ -27,6 +27,7 @@ library
                        Database.Bloodhound.Client
                        Database.Bloodhound.Types
                        Database.Bloodhound.Types.Class
+                       Database.Bloodhound.Types.Internal
   hs-source-dirs:      src
   build-depends:       base             >= 4.3     && <5,
                        bytestring       >= 0.10.0  && <0.11,

src/Database/Bloodhound/Client.hs

@@ -58,6 +58,8 @@ module Database.Bloodhound.Client
        , getStatus
        , encodeBulkOperations
        , encodeBulkOperation
+       -- * Authentication
+       , basicAuthHook
        -- * Reply-handling tools
        , isVersionConflict
        , isSuccess
@@ -175,10 +177,11 @@ dispatch :: MonadBH m => Method -> Text -> Maybe L.ByteString
             -> m Reply
 dispatch dMethod url body = do
   initReq <- liftIO $ parseUrl' url
+  reqHook <- bhRequestHook <$> getBHEnv
   let reqBody = RequestBodyLBS $ fromMaybe emptyBody body
-  let req = initReq { method = dMethod
-                    , requestBody = reqBody
-                    , checkStatus = \_ _ _ -> Nothing}
+  req <- liftIO $ reqHook $ initReq { method = dMethod
+                                    , requestBody = reqBody
+                                    , checkStatus = \_ _ _ -> Nothing}
   mgr <- bhManager <$> getBHEnv
   liftIO $ httpLbs req mgr
 
@@ -222,8 +225,7 @@ bindM2 f ma mb = join (f <$> ma <*> mb)
 withBH :: ManagerSettings -> Server -> BH IO a -> IO a
 withBH ms s f = do
   mgr <- newManager ms
-  let env = BHEnv { bhServer  = s
-                  , bhManager = mgr }
+  let env = mkBHEnv s mgr
   runBH env f
 
 -- Shortcut functions for HTTP methods
@@ -249,11 +251,9 @@ post   = dispatch NHTM.methodPost
 -- Just 200
 getStatus :: MonadBH m => m (Maybe Status)
 getStatus = do
-  url <- joinPath []
-  request <- liftIO $ parseUrl' url
-  mgr <- bhManager <$> getBHEnv
-  response <- liftIO $ httpLbs request mgr
+  response <- get =<< url
   return $ decode (responseBody response)
+  where url = joinPath []
 
 -- | 'createIndex' will create an index given a 'Server', 'IndexSettings', and an 'IndexName'.
 --
@@ -785,3 +785,14 @@ isCreated = statusCheck (== 201)
 
 statusCheck :: (Int -> Bool) -> Reply -> Bool
 statusCheck prd = prd . NHTS.statusCode . responseStatus
+
+-- | This is a hook that can be set via the 'bhRequestHook' function
+-- that will authenticate all requests using an HTTP Basic
+-- Authentication header. Note that it is *strongly* recommended that
+-- this option only be used over an SSL connection.
+--
+-- >> (mkBHEnv myServer myManager) { bhRequestHook = basicAuthHook (EsUsername "myuser") (EsPassword "mypass") }
+basicAuthHook :: Monad m => EsUsername -> EsPassword -> Request -> m Request
+basicAuthHook (EsUsername u) (EsPassword p) = return . applyBasicAuth u' p'
+  where u' = T.encodeUtf8 u
+        p' = T.encodeUtf8 p

src/Database/Bloodhound/Types.hs

@@ -51,7 +51,11 @@ module Database.Bloodhound.Types
        , omitNulls
        , BH(..)
        , runBH
-       , BHEnv(..)
+       , BHEnv
+       , bhServer
+       , bhManager
+       , bhRequestHook
+       , mkBHEnv
        , MonadBH(..)
        , Version(..)
        , Status(..)
@@ -258,6 +262,9 @@ module Database.Bloodhound.Types
        , TermsResult(..)
        , DateHistogramResult(..)
        , DateRangeResult(..)
+
+       , EsUsername(..)
+       , EsPassword(..)
          ) where
 
 import           Control.Applicative
@@ -267,30 +274,31 @@ import           Control.Monad.Reader
 import           Control.Monad.State
 import           Control.Monad.Writer
 import           Data.Aeson
-import           Data.Aeson.Types                (Pair, Parser, emptyObject,
-                                                  parseMaybe)
-import qualified Data.ByteString.Lazy.Char8      as L
+import           Data.Aeson.Types                   (Pair, Parser, emptyObject,
+                                                     parseMaybe)
+import qualified Data.ByteString.Lazy.Char8         as L
 import           Data.Char
-import           Data.Hashable                   (Hashable)
-import qualified Data.HashMap.Strict             as HM
-import           Data.List                       (foldl', nub)
-import           Data.List.NonEmpty              (NonEmpty (..), toList)
-import qualified Data.Map.Strict                 as M
+import           Data.Hashable                      (Hashable)
+import qualified Data.HashMap.Strict                as HM
+import           Data.List                          (foldl', nub)
+import           Data.List.NonEmpty                 (NonEmpty (..), toList)
+import qualified Data.Map.Strict                    as M
 import           Data.Maybe
-import           Data.Text                       (Text)
-import qualified Data.Text                       as T
+import           Data.Text                          (Text)
+import qualified Data.Text                          as T
 import           Data.Time.Calendar
-import           Data.Time.Clock                 (NominalDiffTime, UTCTime)
+import           Data.Time.Clock                    (NominalDiffTime, UTCTime)
 import           Data.Time.Clock.POSIX
-import qualified Data.Traversable                as DT
-import           Data.Typeable                   (Typeable)
-import qualified Data.Vector                     as V
+import qualified Data.Traversable                   as DT
+import           Data.Typeable                      (Typeable)
+import qualified Data.Vector                        as V
 import           GHC.Enum
-import           GHC.Generics                    (Generic)
+import           GHC.Generics                       (Generic)
 import           Network.HTTP.Client
-import qualified Network.HTTP.Types.Method       as NHTM
+import qualified Network.HTTP.Types.Method          as NHTM
 
 import           Database.Bloodhound.Types.Class
+import           Database.Bloodhound.Types.Internal
 
 -- $setup
 -- >>> :set -XOverloadedStrings
@@ -303,20 +311,12 @@ import           Database.Bloodhound.Types.Class
 -- defaultIndexSettings is exported by Database.Bloodhound as well
 -- no trailing slashes in servers, library handles building the path.
 
-{-| Common environment for Elasticsearch calls. Connections will be
-    pipelined according to the provided HTTP connection manager.
--}
-data BHEnv = BHEnv { bhServer  :: Server
-                   , bhManager :: Manager
-                   }
-
-{-| All API calls to Elasticsearch operate within
-    MonadBH. The idea is that it can be easily embedded in your
-    own monad transformer stack. A default instance for a ReaderT and
-    alias 'BH' is provided for the simple case.
--}
-class (Functor m, Applicative m, MonadIO m) => MonadBH m where
-  getBHEnv :: m BHEnv
+-- | Create a 'BHEnv' with all optional fields defaulted. HTTP hook
+-- will be a noop. You can use the exported fields to customize it further, e.g.:
+--
+-- >> (mkBHEnv myServer myManager) { bhRequestHook = customHook }
+mkBHEnv :: Server -> Manager -> BHEnv
+mkBHEnv s m = BHEnv s m return
 
 newtype BH m a = BH {
       unBH :: ReaderT BHEnv m a
@@ -345,10 +345,6 @@ instance (MonadReader r m) => MonadReader r (BH m) where
 instance (Functor m, Applicative m, MonadIO m) => MonadBH (BH m) where
   getBHEnv = BH getBHEnv
 
-
-instance (Functor m, Applicative m, MonadIO m) => MonadBH (ReaderT BHEnv m) where
-  getBHEnv = ask
-
 runBH :: BHEnv -> BH m a -> m a
 runBH e f = runReaderT (unBH f) e
 
@@ -479,9 +475,9 @@ data CompoundFormat = CompoundFileFormat Bool
 
 newtype NominalDiffTimeJSON = NominalDiffTimeJSON { ndtJSON ::  NominalDiffTime }
 
-data IndexSettingsSummary = IndexSettingsSummary { sSummaryIndexName :: IndexName
+data IndexSettingsSummary = IndexSettingsSummary { sSummaryIndexName     :: IndexName
                                                  , sSummaryFixedSettings :: IndexSettings
-                                                 , sSummaryUpdateable :: [UpdatableIndexSetting]}
+                                                 , sSummaryUpdateable    :: [UpdatableIndexSetting]}
                                                  deriving (Eq, Show, Generic, Typeable)
 
 {-| 'Reply' and 'Method' are type synonyms from 'Network.HTTP.Types.Method.Method' -}
@@ -764,10 +760,6 @@ newtype ShardCount = ShardCount Int deriving (Eq, Show, Generic, ToJSON, Typeabl
 -}
 newtype ReplicaCount = ReplicaCount Int deriving (Eq, Show, Generic, ToJSON, Typeable)
 
-{-| 'Server' is used with the client functions to point at the ES instance
--}
-newtype Server = Server Text deriving (Eq, Show, Generic, Typeable)
-
 {-| 'IndexName' is used to describe which index to query/create/delete
 -}
 newtype IndexName = IndexName Text deriving (Eq, Generic, Show, ToJSON, FromJSON, Typeable)
@@ -3544,3 +3536,9 @@ instance Enum DocVersion where
   fromEnum = docVersionNumber
   enumFrom = boundedEnumFrom
   enumFromThen = boundedEnumFromThen
+
+-- | Username type used for HTTP Basic authentication. See 'basicAuthHook'.
+newtype EsUsername = EsUsername { esUsername :: Text } deriving (Show, Eq)
+
+-- | Password type used for HTTP Basic authentication. See 'basicAuthHook'.
+newtype EsPassword = EsPassword { esPassword :: Text } deriving (Show, Eq)

src/Database/Bloodhound/Types/Internal.hs

@@ -0,0 +1,56 @@
+{-# LANGUAGE DeriveDataTypeable         #-}
+{-# LANGUAGE DeriveGeneric              #-}
+{-# LANGUAGE FlexibleInstances          #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+
+-------------------------------------------------------------------------------
+-- |
+-- Module : Database.Bloodhound.Types.Internal
+-- Copyright : (C) 2014 Chris Allen
+-- License : BSD-style (see the file LICENSE)
+-- Maintainer : Chris Allen <cma@bitemyapp.com>
+-- Stability : provisional
+-- Portability : DeriveGeneric, RecordWildCards
+--
+-- Internal data types for Bloodhound. These types may change without
+-- notice so import at your own risk.
+-------------------------------------------------------------------------------
+module Database.Bloodhound.Types.Internal
+    ( BHEnv(..)
+    , Server(..)
+    , MonadBH(..)
+    ) where
+
+
+import           Control.Applicative
+import           Control.Monad.Reader
+import           Data.Text            (Text)
+import           Data.Typeable        (Typeable)
+import           GHC.Generics         (Generic)
+import           Network.HTTP.Client
+
+{-| Common environment for Elasticsearch calls. Connections will be
+    pipelined according to the provided HTTP connection manager.
+-}
+data BHEnv = BHEnv { bhServer      :: Server
+                   , bhManager     :: Manager
+                   , bhRequestHook :: Request -> IO Request
+                   -- ^ Low-level hook that is run before every request is sent. Used to implement custom authentication strategies. Defaults to 'return' with 'mkBHEnv'.
+                   }
+
+instance (Functor m, Applicative m, MonadIO m) => MonadBH (ReaderT BHEnv m) where
+  getBHEnv = ask
+
+{-| 'Server' is used with the client functions to point at the ES instance
+-}
+newtype Server = Server Text deriving (Eq, Show, Generic, Typeable)
+
+{-| All API calls to Elasticsearch operate within
+    MonadBH
+    . The idea is that it can be easily embedded in your
+    own monad transformer stack. A default instance for a ReaderT and
+    alias 'BH' is provided for the simple case.
+-}
+class (Functor m, Applicative m, MonadIO m) => MonadBH m where
+  getBHEnv :: m BHEnv
+

tests/tests.hs

@@ -259,7 +259,7 @@ searchExpectNoResults search = do
 
 searchExpectAggs :: Search -> BH IO ()
 searchExpectAggs search = do
-  reply <- searchAll search
+  reply <- searchByIndex testIndex search
   let isEmpty x = return (M.null x)
   let result = decode (responseBody reply) :: Maybe (SearchResult Tweet)
   liftIO $
@@ -268,7 +268,7 @@ searchExpectAggs search = do
 searchValidBucketAgg :: (BucketAggregation a, FromJSON a, Show a) =>
                         Search -> Text -> (Text -> AggregationResults -> Maybe (Bucket a)) -> BH IO ()
 searchValidBucketAgg search aggKey extractor = do
-  reply <- searchAll search
+  reply <- searchByIndex testIndex search
   let bucketDocs = docCount . head . buckets
   let result = decode (responseBody reply) :: Maybe (SearchResult Tweet)
   let count = result >>= aggregations >>= extractor aggKey >>= \x -> return (bucketDocs x)
@@ -293,8 +293,8 @@ searchExpectSource src expected = do
   _ <- insertData
   let query = QueryMatchQuery $ mkMatchQuery (FieldName "_all") (QueryString "haskell")
   let search = (mkSearch (Just query) Nothing) { source = Just src }
-  reply <- searchAll search
-  result <- parseEsResponse reply--  :: Either EsError (SearchResult Value)
+  reply <- searchByIndex testIndex search
+  result <- parseEsResponse reply
   let value = grabFirst result
   liftIO $
     value `shouldBe` expected