diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright (c)2011, Ozgun Ataman + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of Ozgun Ataman nor the names of other + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/cassy.cabal b/cassy.cabal new file mode 100644 --- /dev/null +++ b/cassy.cabal @@ -0,0 +1,79 @@ +Name: cassy +Version: 0.2.0.1 +Synopsis: A high level driver for the Cassandra datastore +License: BSD3 +License-file: LICENSE +Author: Ozgun Ataman +Maintainer: ozataman@gmail.com +Category: Data +Build-type: Simple +description: + The objective is to completely isolate away the thrift layer, providing + a more idiomatic Haskell experience working with Cassandra. + . + Certain parts of the API was inspired by pycassa (Python client) and + hscassandra (on Hackage). + . + A brief explanation of modules: + . + * /Database.Cassandra.Basic/: Contains a low level, simple + implementation of Cassandra interaction using the thrift API + underneath. + . + * /Database.Cassandra.JSON/: A higher level API that operates on + values with ToJSON and FromJSON isntances from the /aeson/ + library. This module has in part been inspired by Bryan + O\'Sullivan\'s /riak/ client for Haskell. + . + * /Database.Cassandra.Pool/: Handles a /pool/ of connections to + multiple servers in a cluster, splitting the load among them. + . + * /Database.Cassandra.Types/: A common set of types used everywhere. + . + Potential TODOs include: + . + * Support for counters and batch mutators + . + * Support for database admin operations + . + * Support for composite column types + +-- Extra-source-files: + +Cabal-version: >=1.2 + + + +Library + hs-source-dirs: src + Exposed-modules: + Database.Cassandra.Basic + Database.Cassandra.JSON + Database.Cassandra.Pool + Database.Cassandra.Types + + -- Packages needed in order to build this package. + Build-depends: + base >= 4 && < 5 + , bytestring + , containers + , network + , time + , transformers + , stm + , syb + , text + , attoparsec >= 0.10 && < 0.11 + , aeson + , Thrift >= 0.6 + , cassandra-thrift >= 0.8 + + Extensions: + TypeSynonymInstances + + -- Modules not exported by this package. + -- Other-modules: + + -- Extra tools (e.g. alex, hsc2hs, ...) needed to build the source. + -- Build-tools: + diff --git a/src/Database/Cassandra/Basic.hs b/src/Database/Cassandra/Basic.hs new file mode 100644 --- /dev/null +++ b/src/Database/Cassandra/Basic.hs @@ -0,0 +1,259 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE PatternGuards, NamedFieldPuns, RecordWildCards #-} + + +module Database.Cassandra.Basic + +( + -- * Basic Types + ColumnFamily(..) + , Key(..) + , ColumnName(..) + , Value(..) + , Column(..) + , col + , Row(..) + , ConsistencyLevel(..) + + -- * Filtering + , Selector(..) + , Order(..) + , KeySelector(..) + , KeyRangeType(..) + + -- * Exceptions + , CassandraException(..) + + -- * Connection + , CPool + , Server(..) + , defServer + , defServers + , KeySpace(..) + , createCassandraPool + + -- * Cassandra Operations + , getCol + , get + , getMulti + , insert + , delete + + -- * Utility + , getTime + , throwing +) where + + +import Control.Exception +import Control.Monad +import Data.ByteString.Lazy (ByteString) +import qualified Database.Cassandra.Thrift.Cassandra_Client as C +import qualified Database.Cassandra.Thrift.Cassandra_Types as T +import Database.Cassandra.Thrift.Cassandra_Types + (ConsistencyLevel(..)) +import Data.Map (Map) +import qualified Data.Map as M +import Data.Maybe (mapMaybe) +import Network +import Prelude hiding (catch) + +import Database.Cassandra.Pool +import Database.Cassandra.Types + + +test = do + pool <- createCassandraPool [("127.0.0.1", PortNumber 9160)] 3 300 "Keyspace1" + withPool pool $ \ Cassandra{..} -> do + let cp = T.ColumnParent (Just "CF1") Nothing + let sr = Just $ T.SliceRange (Just "") (Just "") (Just False) (Just 100) + let ks = Just ["eben"] + let sp = T.SlicePredicate Nothing sr + C.get_slice (cProto, cProto) "darak" cp sp ONE + get pool "darak" "CF1" All ONE + getCol pool "CF1" "darak" "eben" ONE + insert pool "CF1" "test1" ONE [col "col1" "val1", col "col2" "val2"] + get pool "test1" "CF1" All ONE >>= putStrLn . show + delete pool "test1" "CF1" (ColNames ["col2"]) ONE + get pool "test1" "CF1" (Range Nothing Nothing Reversed 1) ONE >>= putStrLn . show + + +------------------------------------------------------------------------------ +-- | Get a single key-column value +getCol + :: CPool + -> ColumnFamily + -> Key + -- ^ Row key + -> ColumnName + -- ^ Column/SuperColumn name + -> ConsistencyLevel + -- ^ Read quorum + -> IO (Either CassandraException Column) +getCol p cf k cn cl = do + c <- get p cf k (ColNames [cn]) cl + case c of + Left e -> return $ Left e + Right [] -> return $ Left NotFoundException + Right (x:_) -> return $ Right x + + +------------------------------------------------------------------------------ +-- | An arbitrary get operation - slice with 'Selector' +get + :: CPool + -> ColumnFamily + -- ^ in ColumnFamily + -> Key + -- ^ Row key to get + -> Selector + -- ^ Slice columns with selector + -> ConsistencyLevel + -> IO (Either CassandraException Row) +get p cf k s cl = withPool p $ \ Cassandra{..} -> do + res <- wrapException $ C.get_slice (cProto, cProto) k cp (mkPredicate s) cl + case res of + Left e -> return $ Left e + Right xs -> return $ do + cs <- mapM castColumn xs + case cs of + [] -> Left NotFoundException + _ -> Right $ cs + where + cp = T.ColumnParent (Just cf) Nothing + + +------------------------------------------------------------------------------ +-- | Do multiple 'get's in one DB hit +getMulti + :: CPool + -> ColumnFamily + -> KeySelector + -- ^ A selection of rows to fetch in one hit + -> Selector + -- ^ Subject to column selector conditions + -> ConsistencyLevel + -> IO (Either CassandraException (Map ByteString Row)) + -- ^ A Map from Row keys to 'Row's is returned +getMulti p cf ks s cl = withPool p $ \ Cassandra{..} -> do + case ks of + Keys xs -> do + res <- wrapException $ C.multiget_slice (cProto, cProto) xs cp (mkPredicate s) cl + case res of + Left e -> return $ Left e + Right m -> return . Right $ M.mapMaybe f m + KeyRange {} -> do + res <- wrapException $ + C.get_range_slices (cProto, cProto) cp (mkPredicate s) (mkKeyRange ks) cl + case res of + Left e -> return $ Left e + Right res' -> return . Right $ collectKeySlices res' + where + collectKeySlices :: [T.KeySlice] -> Map ByteString Row + collectKeySlices ks = M.fromList $ mapMaybe collectKeySlice ks + where + f (k, Just x) = True + f _ = False + g (k, Just x) = (k,x) + + + collectKeySlice (T.KeySlice (Just k) (Just xs)) = + case mapM castColumn xs of + Left _ -> Nothing + Right xs' -> Just (k, xs') + collectKeySlice _ = Nothing + + cp = T.ColumnParent (Just cf) Nothing + f xs = + case mapM castColumn xs of + Left _ -> Nothing + Right xs' -> Just xs' + + +------------------------------------------------------------------------------ +-- | Insert an entire row into the db. +-- +-- This will do as many round-trips as necessary to insert the full row. +insert + :: CPool + -> ColumnFamily + -> Key + -> ConsistencyLevel + -> Row + -> IO (Either CassandraException ()) +insert p cf k cl row = withPool p $ \ Cassandra{..} -> do + let iOne c = do + c' <- mkThriftCol c + wrapException $ C.insert (cProto, cProto) k cp c' cl + res <- sequenceE $ map iOne row + return $ res >> return () + where + cp = T.ColumnParent (Just cf) Nothing + + +sequenceE :: [IO (Either a b)] -> IO (Either a [b]) +sequenceE [] = return (return []) +sequenceE (a:as) = do + r1 <- a + rr <- sequenceE as + return $ liftM2 (:) r1 rr + + +------------------------------------------------------------------------------ +-- | Delete an entire row, specific columns or a specific sub-set of columns +-- within a SuperColumn. +delete + :: CPool + -> ColumnFamily + -- ^ In 'ColumnFamily' + -> Key + -- ^ Key to be deleted + -> Selector + -- ^ Columns to be deleted + -> ConsistencyLevel + -> IO (Either CassandraException ()) +delete p cf k s cl = withPool p $ \ Cassandra {..} -> do + now <- getTime + wrapException $ case s of + All -> C.remove (cProto, cProto) k cpAll now cl + ColNames cs -> forM_ cs $ \c -> do + C.remove (cProto, cProto) k (cpCol c) now cl + SupNames sn cs -> forM_ cs $ \c -> do + C.remove (cProto, cProto) k (cpSCol sn c) now cl + where + -- wipe out the entire row + cpAll = T.ColumnPath (Just cf) Nothing Nothing + + -- just a single column + cpCol name = T.ColumnPath (Just cf) Nothing (Just name) + + -- scope column by supercol + cpSCol sc name = T.ColumnPath (Just cf) (Just sc) (Just name) + + + +------------------------------------------------------------------------------ +-- | Wrap exceptions into an explicit type +wrapException :: IO a -> IO (Either CassandraException a) +wrapException a = + (a >>= return . Right) + `catch` (\(T.NotFoundException) -> return $ Left NotFoundException) + `catch` (\(T.InvalidRequestException e) -> + return . Left . InvalidRequestException $ maybe "" id e) + `catch` (\T.UnavailableException -> return $ Left UnavailableException) + `catch` (\T.TimedOutException -> return $ Left TimedOutException) + `catch` (\(T.AuthenticationException e) -> + return . Left . AuthenticationException $ maybe "" id e) + `catch` (\(T.AuthorizationException e) -> + return . Left . AuthorizationException $ maybe "" id e) + `catch` (\T.SchemaDisagreementException -> return $ Left SchemaDisagreementException) + + +------------------------------------------------------------------------------- +-- | Make exceptions implicit +throwing :: IO (Either CassandraException a) -> IO a +throwing f = do + res <- f + case res of + Left e -> throw e + Right a -> return a diff --git a/src/Database/Cassandra/JSON.hs b/src/Database/Cassandra/JSON.hs new file mode 100644 --- /dev/null +++ b/src/Database/Cassandra/JSON.hs @@ -0,0 +1,308 @@ +{-# LANGUAGE PatternGuards, + NamedFieldPuns, + OverloadedStrings, + TypeSynonymInstances, + ScopedTypeVariables, + RecordWildCards #-} + +{-| + A higher level module for working with Cassandra. + + Row and Column keys can be any string-like type implementing the + CKey typeclass. You can add your own types by defining new instances + + Serialization and de-serialization of Column values are taken care of + automatically using the ToJSON and FromJSON typeclasses. + + Also, this module currently attempts to reduce verbosity by + throwing errors instead of returning Either types as in the + 'Database.Cassandra.Basic' module. + +-} + +module Database.Cassandra.JSON +( + + -- * Connection + CPool + , Server(..) + , defServer + , defServers + , KeySpace(..) + , createCassandraPool + + + -- * Necessary Types + , CKey(..) + , ModifyOperation(..) + , ColumnFamily(..) + , ConsistencyLevel(..) + , CassandraException(..) + + -- * Cassandra Operations + , get + , getCol + , insertCol + , modify + , modify_ + , delete + + +) where + +import Control.Exception +import Control.Monad +import Data.Aeson as A +import Data.Aeson.Parser (value) +import qualified Data.Attoparsec as Atto (IResult(..), parse) +import qualified Data.ByteString.Char8 as B +import Data.ByteString.Lazy.Char8 (ByteString) +import qualified Data.ByteString.Lazy.Char8 as LB +import Data.Map (Map) +import qualified Data.Map as M +import qualified Data.Text as T +import qualified Data.Text.Encoding as T +import qualified Data.Text.Lazy as LT +import qualified Data.Text.Lazy.Encoding as LT +import Network +import Prelude hiding (catch) + +import Database.Cassandra.Basic hiding (get, getCol, delete) +import qualified Database.Cassandra.Basic as CB +import Database.Cassandra.Pool +import Database.Cassandra.Types + + +------------------------------------------------------------------------------- +---- CKey Typeclass +------------------------------------------------------------------------------- + + +------------------------------------------------------------------------------ +-- | A typeclass to enable using any string-like type for row and column keys +class CKey a where + toBS :: a -> ByteString + fromBS :: ByteString -> a + +instance CKey String where + toBS = LB.pack + fromBS = LB.unpack + +instance CKey LT.Text where + toBS = LT.encodeUtf8 + fromBS = LT.decodeUtf8 + +instance CKey T.Text where + toBS = toBS . LT.fromChunks . return + fromBS = T.concat . LT.toChunks . fromBS + +instance CKey B.ByteString where + toBS = LB.fromChunks . return + fromBS = B.concat . LB.toChunks . fromBS + +instance CKey ByteString where + toBS = id + fromBS = id + + + +------------------------------------------------------------------------------- +-- | Convert regular column to a key-value pair +col2val :: (CKey colKey, FromJSON a) => Column -> (colKey, a) +col2val (Column nm val _ _) = (fromBS nm, maybe err id $ unMarshallJSON' val) + where err = error "Value can't be parsed from JSON." +col2val _ = error "col2val is not implemented for SuperColumns" + + + +------------------------------------------------------------------------------ +-- | Possible outcomes of a modify operation +data ModifyOperation a = + Update a + | Delete + | DoNothing + deriving (Eq,Show,Ord,Read) + + +------------------------------------------------------------------------------ +-- | A modify function that will fetch a specific column, apply modification +-- function on it and save results back to Cassandra. +-- +-- A 'b' side value is returned for computational convenience. +-- +-- This is intended to be a workhorse function, in that you should be +-- able to do all kinds of relatively straightforward operations just +-- using this function. +-- +-- This method may throw a 'CassandraException' for all exceptions other than +-- 'NotFoundException'. +modify + :: (CKey rowKey, CKey colKey, ToJSON a, FromJSON a) + => CPool + -> ColumnFamily + -> rowKey + -> colKey + -> ConsistencyLevel + -- ^ Read quorum + -> ConsistencyLevel + -- ^ Write quorum + -> (Maybe a -> IO (ModifyOperation a, b)) + -- ^ Modification function. Called with 'Just' the value if present, + -- 'Nothing' otherwise. + -> IO b + -- ^ Return the decided 'ModifyOperation' and its execution outcome +modify cp cf k cn rcl wcl f = + let + k' = toBS k + cn' = toBS cn + execF prev = do + (fres, b) <- f prev + dbres <- case fres of + (Update a) -> + insert cp cf k' wcl [col cn' (marshallJSON' a)] + (Delete) -> + CB.delete cp cf k' (ColNames [cn']) wcl + (DoNothing) -> return $ Right () + case dbres of + Left e -> throw e -- Modify op returned error; throw it + Right _ -> return b + in do + res <- CB.getCol cp cf k' cn' rcl + case res of + Left NotFoundException -> execF Nothing + Left e -> throw e + Right Column{..} -> execF (unMarshallJSON' colVal) + Right SuperColumn{..} -> throw $ + OperationNotSupported "modify not implemented for SuperColumn" + + +------------------------------------------------------------------------------ +-- | Same as 'modify' but does not offer a side value. +-- +-- This method may throw a 'CassandraException' for all exceptions other than +-- 'NotFoundException'. +modify_ + :: (CKey rowKey, CKey colKey, ToJSON a, FromJSON a) + => CPool + -> ColumnFamily + -> rowKey + -> colKey + -> ConsistencyLevel + -- ^ Read quorum + -> ConsistencyLevel + -- ^ Write quorum + -> (Maybe a -> IO (ModifyOperation a)) + -- ^ Modification function. Called with 'Just' the value if present, + -- 'Nothing' otherwise. + -> IO () +modify_ cp cf k cn rcl wcl f = + let + f' prev = do + op <- f prev + return (op, ()) + in do + modify cp cf k cn rcl wcl f' + return () + + +------------------------------------------------------------------------------- +-- Simple insertion function making use of typeclasses +insertCol + :: (CKey rowKey, CKey colKey, ToJSON a) + => CPool -> ColumnFamily + -> rowKey + -> colKey + -> ConsistencyLevel + -> a -- ^ Content + -> IO () +insertCol cp cf k cn cl a = + throwing $ insert cp cf (toBS k) cl [col (toBS cn) (marshallJSON' a)] + + +------------------------------------------------------------------------------ +-- | An arbitrary get operation - slice with 'Selector'. +-- +-- Internally based on Basic.get. Table is assumed to be a regular +-- ColumnFamily and contents of returned columns are cast into the +-- target type. +get + :: (CKey rowKey, CKey colKey, FromJSON a) + => CPool -> ColumnFamily + -> rowKey + -> Selector + -> ConsistencyLevel + -> IO [(colKey, a)] +get cp cf k s cl = do + res <- throwing $ CB.get cp cf (toBS k) s cl + return $ map col2val res + + +------------------------------------------------------------------------------- +-- | Get a single column from a single row +getCol + :: (CKey rowKey, CKey colKey, FromJSON a) + => CPool -> ColumnFamily + -> rowKey + -> colKey + -> ConsistencyLevel + -> IO (Maybe a) +getCol cp cf rk ck cl = do + res <- CB.getCol cp cf (toBS rk) (toBS ck) cl + case res of + Left NotFoundException -> return Nothing + Left e -> throw e + Right a -> + let (_ :: ByteString, x) = col2val a + in return $ Just x + + + +------------------------------------------------------------------------------ +-- | Same as the 'delete' in the 'Cassandra.Basic' module, except that +-- it throws an exception rather than returning an explicit Either +-- value. +delete + :: (CKey rowKey) + => CPool + -- ^ Cassandra connection + -> ColumnFamily + -- ^ In 'ColumnFamily' + -> rowKey + -- ^ Key to be deleted + -> Selector + -- ^ Columns to be deleted + -> ConsistencyLevel + -> IO () +delete p cf k s cl = throwing $ CB.delete p cf (toBS k) s cl + + +------------------------------------------------------------------------------ +-- | Lazy 'marshallJSON' +marshallJSON' :: ToJSON a => a -> ByteString +marshallJSON' = LB.fromChunks . return . marshallJSON + + +------------------------------------------------------------------------------ +-- | Encode JSON +marshallJSON :: ToJSON a => a -> B.ByteString +marshallJSON = B.concat . LB.toChunks . A.encode + + +------------------------------------------------------------------------------ +-- | Lazy 'unMarshallJSON' +unMarshallJSON' :: FromJSON a => ByteString -> Maybe a +unMarshallJSON' = unMarshallJSON . B.concat . LB.toChunks + +------------------------------------------------------------------------------ +-- | Decode JSON +unMarshallJSON :: FromJSON a => B.ByteString -> Maybe a +unMarshallJSON = pJson + where + pJson bs = val + where + js = Atto.parse value bs + val = case js of + Atto.Done _ r -> case fromJSON r of + Error e -> error $ "JSON err: " ++ show e + Success a -> a + _ -> Nothing diff --git a/src/Database/Cassandra/Pool.hs b/src/Database/Cassandra/Pool.hs new file mode 100644 --- /dev/null +++ b/src/Database/Cassandra/Pool.hs @@ -0,0 +1,206 @@ +{-# LANGUAGE PatternGuards, NamedFieldPuns, RecordWildCards #-} + + +module Database.Cassandra.Pool where + + + +import Control.Applicative ((<$>)) +import Control.Concurrent.STM +import Control.Exception (SomeException, catch, onException) +import Control.Monad (forM_, forever, join, liftM2, unless, when) +import Control.Monad.IO.Class (liftIO) +import Data.ByteString (ByteString) +import Data.List (partition) +import Data.Time.Clock (NominalDiffTime, UTCTime, diffUTCTime, getCurrentTime) +import Prelude hiding (catch) +import System.Mem.Weak (addFinalizer) +import System.IO (hClose, Handle(..)) + + +import qualified Database.Cassandra.Thrift.Cassandra_Client as C +import Thrift.Transport +import Thrift.Transport.Handle +import Thrift.Transport.Framed +import Thrift.Protocol.Binary +import Network + +------------------------------------------------------------------------------ +-- | A round-robin pool of cassandra connections +type CPool = Pool Cassandra Server + + +type Server = (HostName, PortID) + + +-- | A localhost server with default configuration +defServer :: Server +defServer = ("127.0.0.1", PortNumber 9160) + + +-- | A single localhost server with default configuration +defServers :: [Server] +defServers = [defServer] + + +type KeySpace = String + + +data Cassandra = Cassandra { + cHandle :: Handle + , cFramed :: FramedTransport Handle + , cProto :: BinaryProtocol (FramedTransport Handle) +} + + + +-- | Create a pool of connections to a cluster of Cassandra boxes +-- +-- Each box in the cluster will get up to n connections. The pool will send +-- queries in round-robin fashion to balance load on each box in the cluster. +createCassandraPool + :: [Server] + -- ^ List of servers to connect to + -> Int + -- ^ Max connections per server (n) + -> NominalDiffTime + -- ^ Kill each connection after this many seconds + -> KeySpace + -- ^ Each pool operates on a single KeySpace + -> IO CPool +createCassandraPool servers n maxIdle ks = createPool cr dest n maxIdle servers + where + cr :: Server -> IO Cassandra + cr (host, p) = do + h <- hOpen (host, p) + ft <- openFramedTransport h + let p = BinaryProtocol ft + C.set_keyspace (p,p) ks + return $ Cassandra h ft p + dest h = hClose $ cHandle h + + +------------------------------------------------------------------------------ +-- Generic pool functionality - might want to factor out one day +-- +------------------------------------------------------------------------------ + +newtype Pool a s = Pool { stripes :: TVar (Ring (Stripe a s)) } + + +createPool cr dest n maxIdle servers = do + when (maxIdle < 0.5) $ + modError "pool " $ "invalid idle time " ++ show maxIdle + when (n < 1) $ + modError "pool " $ "invalid maximum resource count " ++ show n + stripes' <- mapM (createStripe cr dest n maxIdle) servers + -- reaperId <- forkIO $ reaper destroy idleTime localPools + -- addFinalizer p $ killThread reaperId + tv <- atomically $ newTVar (mkRing stripes') + return $ Pool tv + + + +withPool :: Pool a s -> (a -> IO b) -> IO b +withPool Pool{..} f = do + Ring{..} <- atomically $ do + r <- readTVar stripes + writeTVar stripes $ next r + return r + withStripe current f + + +data Ring a = Ring { + current :: !a + , used :: [a] + , upcoming :: [a] + } + + +mkRing [] = error "Can't make a ring from empty list" +mkRing (a:as) = Ring a [] as + + +next :: Ring a -> Ring a +next Ring{..} + | (n:rest) <- upcoming + = Ring n (current : used) rest +next Ring{..} + | (n:rest) <- reverse (current : used) + = Ring n [] rest + + +data Stripe a s = Stripe { + idle :: TVar [Connection a] + -- ^ FIFO buffer of idle connections + , inUse :: TVar Int + -- ^ Set of in-use connections + , server :: s + -- ^ Server this strip is connected to + , create :: s -> IO a + -- ^ Create action + , destroy :: (a -> IO ()) + -- ^ Destroy action + , cxns :: Int + -- ^ Max connections + , ttl :: NominalDiffTime + -- ^ TTL for each connection + } + + +createStripe + :: (s -> IO a) + -> (a -> IO ()) + -> Int + -> NominalDiffTime + -> s + -> IO (Stripe a s) +createStripe cr dest n maxIdle s = atomically $ do + idles <- newTVar [] + used <- newTVar 0 + return $ Stripe { + idle = idles + , inUse = used + , server = s + , create = cr + , destroy = dest + , cxns = n + , ttl = maxIdle + } + + +withStripe :: Stripe a s -> (a -> IO b) -> IO b +withStripe Stripe{..} f = do + res <- join . atomically $ do + cs <- readTVar idle + case cs of + (Connection{..}:rest) -> writeTVar idle rest >> return (return cxn) + [] -> do + used <- readTVar inUse + when (used == cxns) retry + writeTVar inUse $! used + 1 + return $ create server + `onException` atomically (modifyTVar_ inUse (subtract 1)) + ret <- f res `onException` (destroy res `onException` return ()) + now <- getCurrentTime + atomically $ modifyTVar_ idle (Connection res now : ) + return ret + + + +data Connection a = Connection { + cxn :: a + , lastUse :: UTCTime + } + + + +modifyTVar_ :: TVar a -> (a -> a) -> STM () +modifyTVar_ v f = readTVar v >>= \a -> writeTVar v $! f a + + +modError :: String -> String -> a +modError func msg = + error $ "Data.Pool." ++ func ++ ": " ++ msg + + diff --git a/src/Database/Cassandra/Types.hs b/src/Database/Cassandra/Types.hs new file mode 100644 --- /dev/null +++ b/src/Database/Cassandra/Types.hs @@ -0,0 +1,178 @@ +{-# LANGUAGE TypeSynonymInstances #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE PatternGuards, NamedFieldPuns, RecordWildCards #-} + +module Database.Cassandra.Types where + +import Control.Monad +import Control.Exception +import Data.Generics +import Data.ByteString.Lazy (ByteString) +import qualified Data.ByteString.Char8 as B +import qualified Data.ByteString.Lazy.Char8 as LB +import Data.Int (Int32, Int64) +import Data.Time +import Data.Time.Clock.POSIX + +import qualified Database.Cassandra.Thrift.Cassandra_Types as C + + +-- | A 'Key' range selector to use with 'getMulti'. +data KeySelector = + Keys [Key] + -- ^ Just a list of keys to get + | KeyRange KeyRangeType Key Key Int32 + -- ^ A range of keys to get. Remember that RandomPartitioner ranges may not + -- mean much as keys are randomly assigned to nodes. + deriving (Show) + + +-- | Encodes the Key vs. Token options in the thrift API. +-- +-- 'InclusiveRange' ranges are just plain intuitive range queries. +-- 'WrapAround' ranges are also inclusive, but they wrap around the ring. +data KeyRangeType = InclusiveRange | WrapAround + deriving (Show) + + +mkKeyRange (KeyRange ty st end cnt) = case ty of + InclusiveRange -> C.KeyRange (Just st) (Just end) Nothing Nothing (Just cnt) + WrapAround -> C.KeyRange Nothing Nothing (Just $ LB.unpack st) (Just $ LB.unpack end) (Just cnt) + +-- | A column selector/filter statement for queries. +-- +-- Remember that SuperColumns are always fully deserialized, so we don't offer +-- a way to filter columns within a 'SuperColumn'. +data Selector = + All + -- ^ Return everything in 'Row' + | ColNames [ColumnName] + -- ^ Return specific columns or super-columns depending on the 'ColumnFamily' + | SupNames ColumnName [ColumnName] + -- ^ When deleting specific columns in a super column + | Range (Maybe ColumnName) (Maybe ColumnName) Order Int32 + -- ^ Return a range of columns or super-columns + deriving (Show) + + +mkPredicate :: Selector -> C.SlicePredicate +mkPredicate s = + let + allRange = C.SliceRange (Just "") (Just "") (Just False) (Just 100) + in case s of + All -> C.SlicePredicate Nothing (Just allRange) + ColNames ks -> C.SlicePredicate (Just ks) Nothing + Range st end ord cnt -> + let + st' = st `mplus` Just "" + end' = end `mplus` Just "" + in C.SlicePredicate Nothing + (Just (C.SliceRange st' end' (Just $ renderOrd ord) (Just cnt))) + + + +------------------------------------------------------------------------------ +-- | Order in a range query +data Order = Regular | Reversed + deriving (Show) + +renderOrd Regular = False +renderOrd Reversed = True + + +type ColumnFamily = String + + +type Key = ByteString + + +type ColumnName = ByteString + + +type Value = ByteString + + +------------------------------------------------------------------------------ +-- | A Column is either a single key-value pair or a SuperColumn with an +-- arbitrary number of key-value pairs +data Column = + SuperColumn ColumnName [Column] + | Column { + colKey :: ColumnName + , colVal :: Value + , colTS :: Maybe Int64 + -- ^ Last update timestamp; will be overridden during write/update ops + , colTTL :: Maybe Int32 + -- ^ A TTL after which Cassandra will erase the column + } + deriving (Eq,Show,Read,Ord) + + +------------------------------------------------------------------------------ +-- | A full row is simply a sequence of columns +type Row = [Column] + + +------------------------------------------------------------------------------ +-- | A short-hand for creating key-value 'Column' values +col :: ByteString -> ByteString -> Column +col k v = Column k v Nothing Nothing + + +mkThriftCol :: Column -> IO C.Column +mkThriftCol Column{..} = do + now <- getTime + return $ C.Column (Just colKey) (Just colVal) (Just now) colTTL + + +castColumn :: C.ColumnOrSuperColumn -> Either CassandraException Column +castColumn x | Just c <- C.f_ColumnOrSuperColumn_column x = castCol c + | Just c <- C.f_ColumnOrSuperColumn_super_column x = castSuperCol c +castColumn _ = + Left $ ConversionException "Unsupported/unexpected ColumnOrSuperColumn type" + + +castCol :: C.Column -> Either CassandraException Column +castCol c + | Just nm <- C.f_Column_name c + , Just val <- C.f_Column_value c + , Just ts <- C.f_Column_timestamp c + , ttl <- C.f_Column_ttl c + = Right $ Column nm val (Just ts) ttl +castCol _ = Left $ ConversionException "Can't parse Column" + + +castSuperCol :: C.SuperColumn -> Either CassandraException Column +castSuperCol c + | Just nm <- C.f_SuperColumn_name c + , Just cols <- C.f_SuperColumn_columns c + , Right cols' <- mapM castCol cols + = Right $ SuperColumn nm cols' +castSuperCol _ = Left $ ConversionException "Can't parse SuperColumn" + + +data CassandraException = + NotFoundException + | InvalidRequestException String + | UnavailableException + | TimedOutException + | AuthenticationException String + | AuthorizationException String + | SchemaDisagreementException + | ConversionException String + | OperationNotSupported String + deriving (Eq,Show,Read,Ord,Data,Typeable) + + +instance Exception CassandraException + + +------------------------------------------------------------------------------ +-- | Cassandra is VERY sensitive to its timestamp values. As a convention, +-- timestamps are always in microseconds +getTime :: IO Int64 +getTime = do + t <- getPOSIXTime + return . fromIntegral . floor $ t * 1000000 +