agitprop.hs

   1 {-# LANGUAGE OverloadedStrings #-}
   2
   3 import Control.Concurrent (threadDelay)
   4 import Control.Monad (forever)
   5 import qualified Data.ByteString.Lazy.Char8 as BL
   6 import qualified Data.Text as T
   7 import Data.Version (showVersion)
   8 import Data.Word (Word64)
   9 import Network.AMQP
  10 import Network.AMQP.Utils.Connection
  11 import Network.AMQP.Utils.Helpers
  12 import Network.AMQP.Utils.Options
  13 import Paths_amqp_utils (version)
  14 import System.Environment
  15 import System.INotify
  16 import qualified System.Posix.Files as F
  17
  18 main :: IO ()
  19 main = do
  20   hr "starting"
  21   args <- getArgs >>= parseargs "agitprop"
  22   printparam' "client version" $ "amqp-utils " ++ (showVersion version)
  23   printparam' "routing key" $ rKey args
  24   isDir <- F.getFileStatus (inputFile args) >>= return . F.isDirectory
  25   if isDir
  26     then printparam' "hotfolder" $ inputFile args
  27     else printparam' "input file" $
  28          (inputFile args) ++
  29          if (lineMode args)
  30            then " (line-by-line)"
  31            else ""
  32   (conn, chan) <- connect args
  33   printparam' "confirm mode" $ show $ confirm args
  34   if (confirm args)
  35     then do
  36       confirmSelect chan False
  37       addConfirmationListener chan confirmCallback
  38     else return ()
  39   let publishOneMsg f = do
  40         r <-
  41           publishMsg
  42             chan
  43             (T.pack $ currentExchange args)
  44             (T.pack $ rKey args)
  45             newMsg {msgBody = f, msgDeliveryMode = Just Persistent}
  46         printparam "sent" $ fmap show r
  47   if isDir
  48     then do
  49       inotify <- initINotify
  50       wd <-
  51         addWatch
  52           inotify
  53           [Close, MoveIn]
  54           (inputFile args)
  55           (handleEvent publishOneMsg)
  56       hr (inputFile args)
  57       _ <- forever $ threadDelay 1000000
  58       removeWatch wd
  59     else do
  60       hr (inputFile args)
  61       messageFile <- BL.readFile (inputFile args)
  62       if (lineMode args)
  63         then mapM_ publishOneMsg (BL.lines messageFile)
  64         else publishOneMsg messageFile
  65   -- all done. wait and close.
  66   if (confirm args)
  67     then waitForConfirms chan >>= return . show >> return ()
  68     else return ()
  69   closeConnection conn
  70
  71 -- | The handler for publisher confirms
  72 confirmCallback :: (Word64, Bool, AckType) -> IO ()
  73 confirmCallback (deliveryTag, isAll, ackType) =
  74   printparam'
  75     "confirmed"
  76     ((show deliveryTag) ++
  77      (if isAll
  78         then " all "
  79         else " this ") ++
  80      (show ackType))
  81
  82 -- | hotfolder event handler
  83 handleEvent :: (BL.ByteString -> IO ()) -> Event -> IO ()
  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 :: (BL.ByteString -> IO ()) -> FilePath -> IO ()
  90 handleFile _ ('.':_) = return ()
  91 handleFile f x = hr x >> BL.readFile x >>= f