{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
module Main (main) where
import Control.Exception (SomeException, bracket)
import Control.Monad
import Control.Monad.Catch (handle, throwM)
import Control.Monad.IO.Class (MonadIO, liftIO)
import qualified Data.Attoparsec.ByteString.Char8 as Atto
import Data.Bits ((.&.), shiftR)
import Data.ByteString (ByteString)
import qualified Data.ByteString.Char8 as B
import Data.Char (ord)
import Data.Foldable (foldMap)
import Data.Maybe (fromMaybe)
import qualified Data.Text as T
import qualified Data.Text.Encoding as T
import qualified Data.Text.IO as T
import Data.Time
import Data.Time.Clock.POSIX
import qualified Data.Vector as V
import Data.Word (Word16, Word32, Word64)
import qualified Data.Configurator as C
import Data.Configurator.Types (Worth(..))
import System.Environment (getEnv)
import System.FilePath.Posix
import System.IO
import System.IO.Unsafe (unsafePerformIO)
import System.Locale (defaultTimeLocale)
import Text.PrettyPrint.Boxes hiding ((<>), (//))
import Data.Hadoop.Configuration (getHadoopConfig)
import Data.Hadoop.Types
import Network.Hadoop.Hdfs hiding (runHdfs)
import Network.Hadoop.Read
import Options.Applicative hiding (Success)
import Data.Version (showVersion)
import Paths_hadoop_tools (version)
------------------------------------------------------------------------
main :: IO ()
main = do
cmd <- execParser optsParser
case cmd of
SubIO io -> io
SubHdfs hdfs -> handle printError (runHdfs hdfs)
where
optsParser = info (helper <*> options)
(fullDesc <> header "hh - Blazing fast interaction with HDFS")
runHdfs :: Hdfs a -> IO a
runHdfs hdfs = do
config <- getConfig
runHdfs' config hdfs
getConfig :: IO HadoopConfig
getConfig = do
hdfsUser <- getHdfsUser
nameNode <- getNameNode
socksProxy <- getSocksProxy
liftM ( set hdfsUser (\c x -> c { hcUser = x })
. set nameNode (\c x -> c { hcNameNodes = [x] })
. set socksProxy (\c x -> c { hcProxy = Just x })
) getHadoopConfig
where
set :: Maybe a -> (b -> a -> b) -> b -> b
set m f c = maybe c (f c) m
------------------------------------------------------------------------
configPath :: FilePath
configPath = unsafePerformIO $ do
home <- getEnv "HOME"
return (home </> ".hh")
{-# NOINLINE configPath #-}
getHdfsUser :: IO (Maybe User)
getHdfsUser = C.load [Optional configPath] >>= flip C.lookup "hdfs.user"
getNameNode :: IO (Maybe NameNode)
getNameNode = do
cfg <- C.load [Optional configPath]
host <- C.lookup cfg "namenode.host"
port <- C.lookupDefault 8020 cfg "namenode.port"
return (Endpoint <$> host <*> pure port)
getSocksProxy :: IO (Maybe SocksProxy)
getSocksProxy = do
cfg <- C.load [Optional configPath]
mhost <- C.lookup cfg "proxy.host"
case mhost of
Nothing -> return Nothing
Just host -> Just . Endpoint host <$> C.lookupDefault 1080 cfg "proxy.port"
------------------------------------------------------------------------
workingDirConfigPath :: FilePath
workingDirConfigPath = unsafePerformIO $ do
home <- getEnv "HOME"
return (home </> ".hhwd")
{-# NOINLINE workingDirConfigPath #-}
getDefaultWorkingDir :: MonadIO m => m HdfsPath
getDefaultWorkingDir = liftIO $ (("/user" //) . T.encodeUtf8 . hcUser) <$> getConfig
getWorkingDir :: MonadIO m => m HdfsPath
getWorkingDir = liftIO $ handle onError
$ B.takeWhile (/= '\n')
<$> B.readFile workingDirConfigPath
where
onError :: SomeException -> IO HdfsPath
onError = const getDefaultWorkingDir
setWorkingDir :: MonadIO m => HdfsPath -> m ()
setWorkingDir path = liftIO $ B.writeFile workingDirConfigPath
$ path <> "\n"
getAbsolute :: MonadIO m => HdfsPath -> m HdfsPath
getAbsolute path = liftIO (normalizePath <$> getPath)
where
getPath = if "/" `B.isPrefixOf` path
then return path
else getWorkingDir >>= \pwd -> return (pwd // path)
normalizePath :: HdfsPath -> HdfsPath
normalizePath = B.intercalate "/" . dropAbsParentDir . B.split '/'
dropAbsParentDir :: [HdfsPath] -> [HdfsPath]
dropAbsParentDir [] = error "dropAbsParentDir: not an absolute path"
dropAbsParentDir (p : ps) = p : reverse (fst $ go [] ps)
where
go [] (".." : ys) = go [] ys
go (_ : xs) (".." : ys) = go xs ys
go xs (y : ys) = go (y : xs) ys
go xs [] = (xs, [])
------------------------------------------------------------------------
data SubCommand = SubCommand
{ subName :: String
, subDescription :: String
, subMethod :: Parser SubMethod
}
data SubMethod = SubIO (IO ())
| SubHdfs (Hdfs ())
sub :: SubCommand -> Mod CommandFields SubMethod
sub SubCommand{..} = command subName (info subMethod $ progDesc subDescription)
options :: Parser SubMethod
options = subparser (foldMap sub allSubCommands)
allSubCommands :: [SubCommand]
allSubCommands =
[ subCat
, subChDir
, subChMod
, subDiskUsage
-- , subFind
, subGet
, subList
, subMkDir
, subPwd
, subRemove
, subRename
, subVersion
]
completePath :: Mod ArgumentFields a
completePath = completer (fileCompletion (const True)) <> metavar "PATH"
completeDir :: Mod ArgumentFields a
completeDir = completer (fileCompletion (== Dir)) <> metavar "DIRECTORY"
bstr :: Monad m => String -> m ByteString
bstr x = B.pack `liftM` str x
subCat :: SubCommand
subCat = SubCommand "cat" "Print the contents of a file to stdout" go
where
go = cat <$> many (argument bstr (completePath <> help "the file to cat"))
cat paths = SubHdfs $ mapM_ (hdfsCat <=< getAbsolute) paths
subChDir :: SubCommand
subChDir = SubCommand "cd" "Change working directory" go
where
go = cd <$> optional (argument bstr (completeDir <> help "the directory to change to"))
cd mpath = SubHdfs $ do
path <- getAbsolute =<< maybe getDefaultWorkingDir return mpath
_ <- getListingOrFail path
setWorkingDir path
subChMod :: SubCommand
subChMod = SubCommand "chmod" "Change permissions" go
where
go = chmod <$> argument bstr (help "permissions mode")
<*> argument bstr (completeDir <> help "the file/directory to chmod")
chmod modeS path = either
(\_ -> error $ "Unknown mode" ++ B.unpack modeS)
(\mode -> SubHdfs $ modifyPerms mode path)
(Atto.parseOnly parseMode modeS)
parseMode = octal
octal :: Atto.Parser Word16
octal = B.foldl' step 0 `fmap` Atto.takeWhile1 isDig
where
isDig w = w >= '0' && w <= '7'
step a w = a * 8 + fromIntegral (ord w - 48)
modifyPerms :: Word16 -> HdfsPath -> Hdfs ()
modifyPerms mode path = do
absPath <- getAbsolute path
info <- getFileInfo absPath
case info of
Nothing -> fail $ unwords ["No such file", B.unpack absPath]
Just FileStatus{..} -> do
{-
liftIO . putStrLn . unwords $ ["Setting perms on", B.unpack absPath]
liftIO . putStrLn . unwords $ ["OLD:", formatMode fsFileType fsPermission]
liftIO . putStrLn . unwords $ ["NEW:", formatMode fsFileType mode]
-}
setPermissions (fromIntegral mode) absPath
subDiskUsage :: SubCommand
subDiskUsage = SubCommand "du" "Show the amount of space used by file or directory" go
where
go = du <$> optional (argument bstr (completePath <> help "the file/directory to check the usage of"))
du path = SubHdfs $ printDiskUsage =<< getAbsolute (fromMaybe "" path)
subFind :: SubCommand
subFind = SubCommand "find" "Recursively search a directory tree" go
where
go = find <$> (argument bstr (completeDir <> help "the path to recursively search"))
<*> (optional (option bstr (long "name" <> metavar "FILENAME"
<> help "the file name to match exactly")))
find path mexpr = SubHdfs $ do
absPath <- getAbsolute path
printFindResults absPath $ maybe (const True) feq mexpr
feq expr FileStatus{..} = expr == fsPath
subGet :: SubCommand
subGet = SubCommand "get" "Get a file" go
where
go = get <$> argument bstr (completePath <> help "source file")
<*> optional (argument str (completePath <> help "destination file"))
get src mdst = SubHdfs $ do
let dst = fromMaybe (takeFileName $ B.unpack src) mdst
absSrc <- getAbsolute src
mReadHandle <- openRead absSrc
let doRead readHandle = liftIO $ bracket
(openFile dst WriteMode)
(hClose)
(\writeHandle -> hdfsMapM_ (B.hPut writeHandle) readHandle)
maybe (return ()) doRead mReadHandle
subList :: SubCommand
subList = SubCommand "ls" "List the contents of a directory" go
where
go = ls <$> optional (argument bstr (completeDir <> help "the directory to list"))
ls path = SubHdfs $ printListing =<< getAbsolute (fromMaybe "" path)
subMkDir :: SubCommand
subMkDir = SubCommand "mkdir" "Create a directory in the specified location" go
where
go = mkdir <$> argument bstr (completeDir <> help "the directory to create")
<*> switch (short 'p' <> help "create intermediate directories")
mkdir path parent = SubHdfs $ do
absPath <- getAbsolute path
ok <- mkdirs parent absPath
unless ok $ liftIO . B.putStrLn $ "Failed to create: " <> absPath
subPwd :: SubCommand
subPwd = SubCommand "pwd" "Print working directory" go
where
go = pure $ SubIO $ B.putStrLn =<< getWorkingDir
subRemove :: SubCommand
subRemove = SubCommand "rm" "Delete a file or directory" go
where
go = rm <$> argument bstr (completePath <> help "the file/directory to remove")
<*> switch (short 'r' <> help "recursively remove the whole file hierarchy")
rm path recursive = SubHdfs $ do
absPath <- getAbsolute path
ok <- delete recursive absPath
unless ok $ liftIO . B.putStrLn $ "Failed to remove: " <> absPath
subRename :: SubCommand
subRename = SubCommand "mv" "Rename a file or directory" go
where
go = mv <$> argument bstr (completePath <> help "source file/directory")
<*> argument bstr (completePath <> help "destination file/directory")
<*> switch (short 'f' <> help "overwrite destination if it exists")
mv src dst force = SubHdfs $ do
absSrc <- getAbsolute src
absDst <- getAbsolute dst
rename force absSrc absDst
subVersion :: SubCommand
subVersion = SubCommand "version" "Show version information" go
where
go = pure $ SubIO $ putStrLn $ "hh version " <> showVersion version
------------------------------------------------------------------------
fileCompletion :: (FileType -> Bool) -> Completer
fileCompletion p = mkCompleter $ \spath -> handle ignore $ runHdfs $ do
let dir = B.pack $ fst $ splitFileName' spath
path = B.pack spath
ls <- getListing' =<< getAbsolute dir
return $ V.toList
. V.map B.unpack
. V.filter (path `B.isPrefixOf`)
. V.map (displayPath dir)
. V.filter (p . fsFileType)
$ ls
where
ignore (RemoteError _ _) = return []
splitFileName' x = case splitFileName x of
("./", f) -> ("", f)
(d, f) -> (d, f)
displayPath :: HdfsPath -> FileStatus -> HdfsPath
displayPath parent file = parent // fsPath file <> suffix
where
suffix = case fsFileType file of
Dir -> "/"
_ -> ""
------------------------------------------------------------------------
printDiskUsage :: HdfsPath -> Hdfs ()
printDiskUsage path = do
ls <- getListingOrFail path
let files = V.map (displayPath path) ls
css <- V.zip files <$> V.mapM getDirSize files
let col a f = vcat a (map (text . f) (V.toList css))
liftIO $ printBox $ col right snd
<+> col left (B.unpack . fst)
where
getDirSize f = handle (\e -> if isAccessDenied e then return "-" else throwM e)
(formatSize . csLength <$> getContentSummary f)
printListing :: HdfsPath -> Hdfs ()
printListing path = do
ls <- getListingOrFail path
let hdfs2utc ms = posixSecondsToUTCTime (fromIntegral ms / 1000)
getModTime = hdfs2utc . fsModificationTime
col a f = vcat a (map (text . f) (V.toList ls))
liftIO $ do
putStrLn $ "Found " <> show (V.length ls) <> " items"
printBox $ col left (\x -> formatMode (fsFileType x) (fsPermission x))
<+> col right (formatBlockReplication . fsBlockReplication)
<+> col left (T.unpack . fsOwner)
<+> col left (T.unpack . fsGroup)
<+> col right (formatSize . fsLength)
<+> col right (formatUTC . getModTime)
<+> col left (T.unpack . T.decodeUtf8 . fsPath)
printFindResults :: HdfsPath -> (FileStatus -> Bool) -> Hdfs ()
printFindResults path cond = handle (liftIO . printError) $ do
ls <- getListingOrFail path
V.mapM_ printMatch ls
where
printMatch :: FileStatus -> Hdfs ()
printMatch fs@FileStatus{..} = do
let path' = displayPath path fs
when (cond fs) (liftIO $ B.putStrLn path')
case fsFileType of
Dir -> printFindResults path' cond
_ -> return ()
------------------------------------------------------------------------
formatSize :: Word64 -> String
formatSize b | b <= 0 = "0"
| b < 1000 = show b <> "B"
| b < 1000000 = show (b `div` 1000) <> "K"
| b < 1000000000 = show (b `div` 1000000) <> "M"
| b < 1000000000000 = show (b `div` 1000000000) <> "G"
| otherwise = show (b `div` 1000000000000) <> "T"
formatBlockReplication :: Word16 -> String
formatBlockReplication x | x == 0 = "-"
| otherwise = show x
formatUTC :: UTCTime -> String
formatUTC = formatTime defaultTimeLocale "%Y-%m-%d %H:%M"
formatMode :: FileType -> Permission -> String
formatMode File = ("-" <>) . formatPermission
formatMode Dir = ("d" <>) . formatPermission
formatMode SymLink = ("l" <>) . formatPermission
formatPermission :: Permission -> String
formatPermission perms = format (perms `shiftR` 6)
<> format (perms `shiftR` 3)
<> format perms
where
format p = conv 0x4 "r" p
<> conv 0x2 "w" p
<> conv 0x1 "x" p
conv bit rep p | (p .&. bit) /= 0 = rep
| otherwise = "-"
------------------------------------------------------------------------
printError :: RemoteError -> IO ()
printError (RemoteError subject body)
| oneLiner = T.putStrLn firstLine
| T.null body = T.putStrLn subject
| otherwise = T.putStrLn subject >> T.putStrLn body
where
oneLiner = subject `elem` [ "org.apache.hadoop.security.AccessControlException"
, "org.apache.hadoop.fs.FileAlreadyExistsException"
, "java.io.FileNotFoundException" ]
firstLine = T.takeWhile (/= '\n') body
isAccessDenied :: RemoteError -> Bool
isAccessDenied (RemoteError s _) = s == "org.apache.hadoop.security.AccessControlException"
------------------------------------------------------------------------
infixr 5 //
(//) :: HdfsPath -> HdfsPath -> HdfsPath
(//) xs ys | B.null xs = ys
| B.null ys = xs
| B.head ys == '/' = ys
| B.last xs == '/' = xs <> ys
| otherwise = xs <> "/" <> ys
trimEnd :: Char -> ByteString -> ByteString
trimEnd b bs | B.null bs = B.empty
| B.last bs == b = trimEnd b (B.init bs)
| otherwise = bs
------------------------------------------------------------------------
getListingOrFail :: HdfsPath -> Hdfs (V.Vector FileStatus)
getListingOrFail path = do
mls <- getListing path
case mls of
Nothing -> throwM $ RemoteError ("File/directory does not exist: " <> T.decodeUtf8 path) T.empty
Just ls -> return ls