packages feed

bein-0.3: Bein/Minion/Protocol.hs

module Bein.Minion.Protocol where

import System.Directory (doesFileExist)
import Prelude hiding (log)
import Text.ParserCombinators.Parsec
    ( CharParser,
      anyChar,
      char,
      newline,
      noneOf,
      oneOf,
      spaces,
      string,
      eof,
      (<|>),
      many,
      parse,
      try )
import Control.Monad ( MonadPlus(mzero) )
import Control.Monad.Trans ( liftIO )
import Database.HDBC ( fromSql, toSql, IConnection(commit, run) )
import System.Directory ( copyFile )
import System.FilePath.Posix ( joinPath )
import System.Posix.Files ( createSymbolicLink, rename )
import Bein.Minion.Types
    ( Configuration(file_repository), BeinM, configField, State )
import Bein.Minion.Commands
    ( encodeResponse,
      catchR,
      scratchDir,
      execution,
      database,
      query,
      update,
      receiveOneRow, 
      maybeRowQuery )

minionProtocol :: String -> BeinM State String                    
minionProtocol received = do
  case parse command "socket" received of
    Left e -> return $ "200 unknown command\n" ++ show e ++ "\n.\n"
    Right c -> executeCommand c

command :: CharParser st MinionCommand
command = foldr (<|>) mzero commands
  where commands = map try [log,get,put,completed,failed]

executeCommand :: MinionCommand -> BeinM State String
executeCommand c = executeCommand' c `catchR` 
                   (\e -> return $ "15 database access error\n" ++ show e ++ "\n.\n")

executeCommand' :: MinionCommand -> BeinM State String
executeCommand' (Log body) = do 
  mapM_ writeLog $ lines body
  database >>= liftIO . commit
  return "0 ok\n.\n"
executeCommand' (Get lbl format) = do
  getInput lbl format
executeCommand' (Put lbl format body) = do
  putOutput lbl format body
executeCommand' Completed = do
  ex <- execution
  update "update executions set status='complete' where id=?" [toSql ex]
  return "0 ok\n.\n"
executeCommand' Failed = do
  ex <- execution
  update "update executions set status='failed' where id=?" [toSql ex]
  return "0 ok\n.\n"

writeLog :: String -> BeinM State Integer
writeLog line = do
  ex <- execution
  database >>= \conn -> liftIO $ run conn queryString [toSql ex, toSql line]
    where queryString = "insert into execution_logs (id,log_time,log_message) " ++ 
                        "values (?,now(),?)"

data MinionCommand = Log String
                   | Get Label GetFormat
                   | Put Label PutFormat String
                   | Completed
                   | Failed
                   deriving (Eq,Show,Read)
                            
type Label = String                            

data GetFormat = GetCopy | GetReadOnlyLink | GetNoFormat deriving (Eq,Show,Read)

data PutFormat = PutMove | PutCopy | PutNoFormat deriving (Eq,Show,Read)

log :: CharParser st MinionCommand
log = do
  string "10 log"
  newline
  ls <- many anyChar
  eof
  return $ Log ls
  
quotedString :: CharParser st String
quotedString = do
  char '"'
  str <- many escapedChar
  char '"'
  return str
    where escapedChar = (try (string "\\\"") >> return '"') <|> noneOf "\""

getFormat :: CharParser st GetFormat
getFormat = (try (string "copy") >> return GetCopy) <|>
            (string "read-only-link" >> return GetReadOnlyLink)

get :: CharParser st MinionCommand
get = do
  string "20 get"
  spaces
  lbl <- quotedString
  many $ oneOf " \t"
  format <- try getFormat <|> return GetNoFormat
  many $ oneOf " \t"
  newline
  eof
  return $ Get lbl format

putFormat :: CharParser st PutFormat
putFormat = (try (string "copy") >> return PutCopy) <|>
            (try (string "move") >> return PutMove) <|>
            return PutNoFormat

put :: CharParser st MinionCommand
put = do
  string "30 put"
  spaces
  lbl <- quotedString
  many $ oneOf " \t"
  format <- try putFormat <|> return PutNoFormat
  many $ oneOf " \t"
  newline
  ls <- many anyChar
  eof
  return $ Put lbl format ls

completed :: CharParser st MinionCommand
completed = bodylessCommand "100 completed" Completed

failed :: CharParser st MinionCommand
failed = bodylessCommand "200 failed" Failed

bodylessCommand :: String -> a -> CharParser st a
bodylessCommand cmd val = do string cmd
                             newline
                             eof
                             return val


getInput :: String -> GetFormat -> BeinM State String
getInput lbl format = do
  ex <- execution
  t <- maybeRowQuery "select type from execution_inputs where id=? and label=?" [toSql ex, toSql lbl]
  case fmap (fromSql . head) t of
    Just "sequence" -> return "500 unimplemented command\n.\n"
    Just "number" -> getNumberInput lbl
    Just "string" -> getStringInput lbl
    Just "file" -> getFileInput lbl format
    Nothing -> return $ encodeResponse "32 no such label" ""
    _ -> error "Invalid type in getInput."
    
getStringInput :: String -> BeinM State String
getStringInput lbl = do
  ex <- execution
  [v] <- receiveOneRow $ query "select value from execution_string_inputs where id=? and label=?" [toSql ex, toSql lbl]
  let v' = fromSql v
  return $ encodeResponse "0 ok" v'    
  
getNumberInput :: String -> BeinM State String
getNumberInput lbl = do
  ex <- execution
  [v] <- receiveOneRow $ query "select value from execution_number_inputs where id=? and label=?" [toSql ex, toSql lbl]
  let v' = fromSql v
  return $ encodeResponse "0 ok" v'
  
getFileInput :: String -> GetFormat -> BeinM State String
getFileInput lbl GetNoFormat = getFileInput lbl GetCopy
getFileInput lbl format = do
  ex <- execution
  sd <- scratchDir
  [fname] <- receiveOneRow $ query ("select f.stored_as from files as f right join " ++ 
    "execution_object_inputs as e on e.value=f.id where e.label=? and e.id=?")
    [toSql lbl, toSql ex]
  let fname' = fromSql fname
  filePath <- configField file_repository
  liftIO $ (if format == GetCopy then copyFile else createSymbolicLink) (joinPath [filePath, fname']) (joinPath [sd, fname'])
  return $ encodeResponse "0 ok" fname'
  
  
putOutput :: String -> PutFormat -> String -> BeinM State String        
putOutput lbl format body = do
  ex <- execution
  r <- maybeRowQuery "select target,type from execution_outputs where id=? and label=?" [toSql ex, toSql lbl]
  case fmap (\[a,b] -> (fromSql a, fromSql b)) r of
    Just (_,"sequence") -> return "500 unimplemented command\n.\n"
    Just (target,"file") -> putFileOutput lbl format body target
    Nothing -> return $ encodeResponse "32 no such label" ""
    _ -> return $ encodeResponse "1000 unknown error" ""

putFileOutput :: String -> PutFormat -> String -> Integer -> BeinM State String
putFileOutput _lbl format body target = do
  if length (lines body) > 1 || ' ' `elem` body 
    then return $ encodeResponse "36 invalid file name" "File name cannot contain newlines or spaces"
    else do let [fname] = lines body
            sd <- scratchDir
            filePath <- configField file_repository
            [targetName] <- receiveOneRow $ query "select unique_name(in_repository(''),50)" []
            let targetName' = fromSql targetName
            q <- liftIO $ doesFileExist (joinPath [sd,fname])
            if q 
              then do liftIO $ (if format == PutNoFormat || format == PutCopy then copyFile else rename) 
                        (joinPath [sd,fname]) (joinPath [filePath,targetName'])
                      update "insert into files (id,user_filename,stored_as) values (?,?,?)" [toSql target, toSql fname, toSql targetName']
                      return $ encodeResponse "0 ok" ""
              else return $ encodeResponse "35 file does not exist" ""