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