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