HandlerSocketClient-0.0.1: Database/HandlerSocket/HSClient.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.HandlerSocket.HSClient (
  -- * Basic Type
  HandSock (..),
  -- * Connection
  connectTo,
  closeHS,
  -- * Do Generic Request
  runReqNoRespData,
  runReqGeneric,
  -- * Do Special Request
  runReqOpen,
  runReqFind,
  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, getHandle :: Int}

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
      [] -> []
      (10:_) -> []
      s1 -> w : (mysplit s2)
            where
              (w,s2) = break isspace' s1

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

makeReqOpen :: (Integral a) => a -> String -> String -> String -> String -> B.ByteString
makeReqOpen hand db tb ind cs = makeReq ["P",show hand,db,tb,ind,cs]

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

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

makeReqUpdate :: (Integral a) => a -> String -> [String] -> B.ByteString
makeReqUpdate hand key newvs = makeReq $ [show hand,"=","1",key,"1","0","U"] ++ newvs

makeReqDelete :: (Integral a) => a -> String -> B.ByteString
makeReqDelete hand key = makeReq [show hand,"=","1",key,"1","0","D"]

-- | Connect to specific address & port with a given handle-id (or index).
-- Returns the 'HandSock' for later request or close.
connectTo :: String -> String -> Int -> IO HandSock
connectTo host port hand = 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 hand

runReqBool :: HandSock -> B.ByteString -> IO Bool
runReqBool (HandSock s h) req = do
  sendAll s req
  resp <- recv s 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 (HandSock s h) vs = do
  let req = makeReq vs
  sendAll s req
  fmap decodeResp $ recv s 4096

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

-- | Do \'find\' request. Returns 
runReqFind :: (Integral a) => HandSock -> String -> String -> a -> IO [[String]]
runReqFind (HandSock s h) op v lim = do
  let req = makeReqFind h op v lim
  sendAll s req
  resp <- fmap decodeResp $ recv s 4096
  case resp of
    (o:i:xs) -> return $ mygroup (read i :: Int) xs
    _ -> return []
  where
    mygroup _ [] = []
    mygroup i xs = l:(mygroup i r) 
      where (l,r) = splitAt i xs

-- | Do \'insert\' request, insert one record at a time.
runReqInsert :: HandSock -> [String] -> IO Bool
runReqInsert hs vs = runReqBool hs $ makeReqInsert (getHandle hs) 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 (getHandle hs) 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 (getHandle hs) key

-- | Close the connection.
closeHS :: HandSock -> IO ()
closeHS hs = sClose (getSocket hs)
{-
main :: IO ()
main = do
  hs <- connectTo "127.0.0.1" "9999" 0
  r1 <- runReqOpen hs "test" "t1" "PRIMARY" "k,v"
  putStrLn $ show r1
  r2 <- runReqFind hs ">" "1" 100
  putStrLn $ show r2
  r3 <- runReqInsert hs ["99", "77"]
  putStrLn $ show r3
  r4 <- runReqUpdate hs "99" ["101", "77"]
  putStrLn $ show r4
  r5 <- runReqDelete hs "101"
  putStrLn $ show r5
  r6 <- runReqGeneric hs ["0","+","2","555","999"]
  putStrLn $ show r6
  return ()
-}