HandlerSocketClient-0.0.4: Database/HandlerSocketClient.hs

--
-- Module : HSClient
-- Author : Wu Xingbo
-- Copyright (c) 2011 Wu Xingbo (wuxb45@gmail.com)
-- New BSD License
--
-- |
-- Haskell implementation of a HandlerSocket client (API).
-- HandlerSocket is a MySQL plugin with better performance then SQL.
-- HandlerSocket site: https://github.com/ahiguti/HandlerSocket-Plugin-for-MySQL
--


module Database.HandlerSocketClient (
  -- * Basic Type
  HandSock (..),
  -- * Connection
  connectTo,
  closeHS,
  -- * Do Generic Request
  runReqNoRespData,
  runReqGeneric,
  -- * Do Special Request
  runReqOpen,
  runReqFind,
  runReqFindN,
  runReqInsert,
  runReqUpdate,
  runReqDelete
  ) where
import Data.Word
import Data.Char
import Data.List
import Network.Socket hiding (send,sendTo,recv,recvFrom)
import Network.Socket.ByteString
import System.IO
import qualified Data.ByteString as B

-- | basic data structure for one Socket and one Handle (or Index).
data HandSock = HandSock {getSocket :: Socket, getHost :: String, getPort :: String}

recvOne :: Socket -> Int -> IO B.ByteString
recvOne sock maxlen = do
  s <- recvOneR sock maxlen
  if B.elem 10 s
    then return s
    else return B.empty
  
recvOneR :: Socket -> Int -> IO B.ByteString
recvOneR sock maxlen = do
  s1 <- recv sock maxlen
  if B.null s1
    then return s1
    else
      if B.elem 10 s1
        then return s1
        else fmap (B.append s1) $ recvOneR sock maxlen
recvN :: Socket -> Int -> Int -> IO [B.ByteString]
recvN sock maxlen count = do
  all <- recvNR sock maxlen count
  return $ init $ B.split 10 all

recvNR :: Socket -> Int -> Int -> IO B.ByteString
recvNR sock maxlen todo
  | todo <= 0 = return B.empty
  | otherwise = do
      s1 <- recv sock maxlen
      if B.null s1
        then return s1
        else fmap (B.append s1) $ recvNR sock maxlen $ todo - (B.count 10 s1)

escapeW :: [Word8] -> [Word8]
escapeW [] = []
escapeW (x:xs) = if x < 16 then 1:(0x40+x):(escapeW xs) else x:(escapeW xs)

escape :: String -> [Word8]
escape str = escapeW $ map (fromIntegral . ord) str

unescapeW :: [Word8] -> [Word8]
unescapeW [] = []
unescapeW (x:xs) = 
  if x == 1 
  then case xs of
    (y:ys) -> (y-0x40):(unescapeW ys)
    _ -> [] -- error
  else x:(unescapeW xs)

unescape :: [Word8] -> String
unescape xs = map (chr . fromIntegral) $ unescapeW xs

decodeResp :: B.ByteString -> [String]
decodeResp = map unescape . mysplit . B.unpack
  where
    isspace' = (\x -> (x == 9) || (x == 10)) -- \t \n
    mysplit s = case dropWhile isspace' s of -- drop leading space'
      [] -> []
      (10:_) -> []
      s1 -> w : (mysplit s2) -- first : rest.
            where
              (w,s2) = break isspace' s1

makeReq :: [String] -> B.ByteString
makeReq = B.pack . (++ [10]) . intercalate [9] . map escape

makeReqOpen :: String -> String -> String -> String -> B.ByteString
makeReqOpen db tb ind cs = makeReq ["P","0",db,tb,ind,cs]

makeReqFind :: (Integral a) => String -> String -> a -> B.ByteString
makeReqFind op v lim = makeReq ["0",op,"1",v,show lim,"0"]

makeReqInsert :: [String] -> B.ByteString
makeReqInsert vs = makeReq $ ["0","+",show (length vs)] ++ vs

makeReqUpdate :: String -> [String] -> B.ByteString
makeReqUpdate key newvs = makeReq $ ["0","=","1",key,"1","0","U"] ++ newvs

makeReqDelete :: String -> B.ByteString
makeReqDelete key = makeReq ["0","=","1",key,"1","0","D"]

-- | Connect to specific address & port.
-- Returns the 'HandSock' for later request or close.
connectTo :: String -> String -> IO HandSock
connectTo host port = do
  addrinfos <- getAddrInfo Nothing (Just host) (Just port)
  let addr = head addrinfos
  let addr' = addr {addrFamily = AF_INET, addrSocketType = Stream, addrProtocol = 6}
  sock <- socket AF_INET Stream 6
  connect sock (addrAddress addr')
  return $ HandSock sock host port

runReqBool :: HandSock -> B.ByteString -> IO Bool
runReqBool hs req = do
  sendAll (getSocket hs) req
  resp <- recvOne (getSocket hs) 1024
  if B.null resp
    then return False
    else return $ 48 == (B.head resp)

-- | Do Generic Request with given Strings. Return True or False, discarding all other messages.
runReqNoRespData :: HandSock -> [String] -> IO Bool
runReqNoRespData hs vs = runReqBool hs $ makeReq vs

-- | Do Generic Request, Returns all values received, \"0\" in the head means SUCCESS.
runReqGeneric :: HandSock -> [String] -> IO [String]
runReqGeneric hs vs = do
  let req = makeReq vs
  sendAll (getSocket hs) req
  fmap decodeResp $ recvOne (getSocket hs) 4096

-- | Do \'open_index\' request.
runReqOpen :: HandSock -> String -> String -> String -> String -> IO Bool
runReqOpen hs db tb ind cs = runReqBool hs $ makeReqOpen db tb ind cs

mygroup :: Int -> [a] -> [[a]]
mygroup _ [] = []
mygroup i xs = l:(mygroup i r) 
  where (l,r) = splitAt i xs

splitFindResult :: [String] -> [[String]]
splitFindResult ss =
  case ss of
    ("0":i:xs) -> mygroup (read i) xs
    _ -> []
    
-- | Do \'find\' request. Returns at most lim results.
runReqFind :: (Integral a) => HandSock -> String -> String -> a -> IO [[String]]
runReqFind hs op v lim = do
  let req = makeReqFind op v lim
  sendAll (getSocket hs) req
  fmap (splitFindResult . decodeResp) $ recvOne (getSocket hs) 4096

runReqFindN :: (Integral a) => HandSock -> String -> [String] -> a -> IO [[[String]]]
runReqFindN hs op vs lim = do
  mapM_ (\v -> let req = makeReqFind op v lim in sendAll (getSocket hs) req) vs
  map (splitFindResult . decodeResp) `fmap` recvN (getSocket hs) 4096 (length vs)
  
-- | Do \'insert\' request, insert one record at a time.
runReqInsert :: HandSock -> [String] -> IO Bool
runReqInsert hs vs = runReqBool hs $ makeReqInsert vs

-- | Do \'update\' request, give one value to be compared with the KEY, then provide the new data.
-- Only ONE record will be updated.
runReqUpdate :: HandSock -> String -> [String] -> IO Bool
runReqUpdate hs v0 vs = runReqBool hs $ makeReqUpdate v0 vs

-- | Do \'delete\' request, give one value to be compared with the KEY. 
-- Only ONE record will be deleted.
runReqDelete :: HandSock -> String -> IO Bool
runReqDelete hs key = runReqBool hs $ makeReqDelete key

-- | Close the connection.
closeHS :: HandSock -> IO ()
closeHS hs = sClose (getSocket hs)

reConnect :: HandSock -> IO HandSock
reConnect hs = do
  closeHS hs
  connectTo (getHost hs) (getPort hs)