agitprop.hs

   1 {-# LANGUAGE OverloadedStrings #-}
   2
   3 import Control.Concurrent (threadDelay)
   4 import qualified Control.Exception as X
   5 import Control.Monad (forever)
   6 import qualified Data.ByteString.Lazy.Char8 as BL
   7 import qualified Data.Text as T
   8 import Data.Version (showVersion)
   9 import Data.Word (Word64)
  10 import Network.AMQP
  11 import Network.AMQP.Utils.Connection
  12 import Network.AMQP.Utils.Helpers
  13 import Network.AMQP.Utils.Options
  14 import Paths_amqp_utils (version)
  15 import System.Environment
  16 import System.INotify
  17 import qualified System.Posix.Files as F
  18
  19 main :: IO ()
  20 main = do
  21   hr "starting"
  22   args <- getArgs >>= parseargs "agitprop"
  23   printparam' "client version" $ "amqp-utils " ++ (showVersion version)
  24   printparam' "routing key" $ rKey args
  25   isDir <- F.getFileStatus (inputFile args) >>= return . F.isDirectory
  26   if isDir
  27     then printparam' "hotfolder" $ inputFile args
  28     else printparam' "input file" $
  29          (inputFile args) ++
  30          if (lineMode args)
  31            then " (line-by-line)"
  32            else ""
  33   (conn, chan) <- connect args
  34   printparam' "confirm mode" $ show $ confirm args
  35   if (confirm args)
  36     then do
  37       confirmSelect chan False
  38       addConfirmationListener chan confirmCallback
  39     else return ()
  40   let publishOneMsg = publishOneMsg' chan args
  41   X.catch
  42     (if isDir
  43        then do
  44          inotify <- initINotify
  45          wd <-
  46            addWatch
  47              inotify
  48              [CloseWrite, MoveIn]
  49              (inputFile args)
  50              (handleEvent publishOneMsg)
  51          hr $ "watching " ++ (inputFile args)
  52          _ <- forever $ threadDelay 1000000
  53          removeWatch wd
  54        else do
  55          hr $ "sending " ++ (inputFile args)
  56          messageFile <- BL.readFile (inputFile args)
  57          if (lineMode args)
  58            then mapM_ publishOneMsg (BL.lines messageFile)
  59            else publishOneMsg messageFile)
  60     (\exception -> printparam' "exception" $ show (exception :: X.SomeException))
  61   -- all done. wait and close.
  62   if (confirm args)
  63     then waitForConfirms chan >>= return . show >> return ()
  64     else return ()
  65   closeConnection conn
  66
  67 -- | The handler for publisher confirms
  68 confirmCallback :: (Word64, Bool, AckType) -> IO ()
  69 confirmCallback (deliveryTag, isAll, ackType) =
  70   printparam'
  71     "confirmed"
  72     ((show deliveryTag) ++
  73      (if isAll
  74         then " all "
  75         else " this ") ++
  76      (show ackType))
  77
  78 -- | Hotfolder event handler
  79 handleEvent :: (BL.ByteString -> IO ()) -> Event -> IO ()
  80 -- just handle closewrite and movedin events
  81 handleEvent f (Closed False (Just x) True) = handleFile f x
  82 handleEvent f (MovedIn False x _) = handleFile f x
  83 handleEvent _ _ = return ()
  84
  85 -- | Hotfolder file handler
  86 handleFile :: (BL.ByteString -> IO ()) -> FilePath -> IO ()
  87 handleFile _ ('.':_) = return () -- ignore hidden files
  88 handleFile f x =
  89   X.catch
  90     (hr ("sending " ++ x) >> BL.readFile x >>= f)
  91     (\exception ->
  92        printparam' "exception in handleFile" $
  93        show (exception :: X.SomeException))
  94
  95 -- | Publish one message with our settings
  96 publishOneMsg' :: Channel -> Args -> BL.ByteString -> IO ()
  97 publishOneMsg' c a f = do
  98       r <-
  99         publishMsg
 100           c
 101           (T.pack $ currentExchange a)
 102           (T.pack $ rKey a)
 103           newMsg { msgBody = f
 104                  , msgDeliveryMode = Just Persistent
 105                  , msgTimestamp = msgtimestamp a
 106                  , msgID = msgid a
 107                  , msgType = msgtype a
 108                  , msgUserID = msguserid a
 109                  , msgApplicationID = msgappid a
 110                  , msgClusterID = msgclusterid a
 111                  , msgContentType = msgcontenttype a
 112                  , msgContentEncoding = msgcontentencoding a
 113                  , msgReplyTo = msgreplyto a
 114                  , msgPriority = msgprio a
 115                  , msgCorrelationID = msgcorrid a
 116                  , msgExpiration = msgexp a
 117                  -- , msgHeaders =
 118                  }
 119       printparam "sent" $ fmap show r