module ClickHaskell.Connection where

-- Internal
import ClickHaskell.Primitive

-- GHC included
import Control.Concurrent (MVar)
import Control.Exception (throwIO)
import Data.Binary.Builder (Builder, toLazyByteString)
import Data.Binary.Get
import Data.ByteString as BS (ByteString, length)
import Data.ByteString.Lazy as BSL (ByteString)
import Data.IORef (IORef, atomicModifyIORef, atomicWriteIORef, newIORef, readIORef)
import Data.Text (Text)
import GHC.Exception (Exception)
import Prelude hiding (liftA2)

-- External
import Network.Socket hiding (SocketOption(..))
import Network.Socket qualified as Sock (SocketOption(..))
import Network.Socket.ByteString (recv)
import Network.Socket.ByteString.Lazy (sendAll)



-- * Connection

{- |
  Errors occured on connection operations
-}
data ConnectionError
  = NoAdressResolved
  -- ^ Occurs when 'getAddrInfo' returns an empty result
  | EstablishTimeout
  -- ^ Occurs on 'socket' connection timeout
  deriving (Int -> ConnectionError -> ShowS
[ConnectionError] -> ShowS
ConnectionError -> String
(Int -> ConnectionError -> ShowS)
-> (ConnectionError -> String)
-> ([ConnectionError] -> ShowS)
-> Show ConnectionError
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> ConnectionError -> ShowS
showsPrec :: Int -> ConnectionError -> ShowS
$cshow :: ConnectionError -> String
show :: ConnectionError -> String
$cshowList :: [ConnectionError] -> ShowS
showList :: [ConnectionError] -> ShowS
Show, Show ConnectionError
Typeable ConnectionError
(Typeable ConnectionError, Show ConnectionError) =>
(ConnectionError -> SomeException)
-> (SomeException -> Maybe ConnectionError)
-> (ConnectionError -> String)
-> Exception ConnectionError
SomeException -> Maybe ConnectionError
ConnectionError -> String
ConnectionError -> SomeException
forall e.
(Typeable e, Show e) =>
(e -> SomeException)
-> (SomeException -> Maybe e) -> (e -> String) -> Exception e
$ctoException :: ConnectionError -> SomeException
toException :: ConnectionError -> SomeException
$cfromException :: SomeException -> Maybe ConnectionError
fromException :: SomeException -> Maybe ConnectionError
$cdisplayException :: ConnectionError -> String
displayException :: ConnectionError -> String
Exception)

{- |
  These exceptions might indicate internal bugs.

  If you encounter one, please report it.
-}
data InternalError
  = UnexpectedPacketType UVarInt
  | DeserializationError String
  deriving (Int -> InternalError -> ShowS
[InternalError] -> ShowS
InternalError -> String
(Int -> InternalError -> ShowS)
-> (InternalError -> String)
-> ([InternalError] -> ShowS)
-> Show InternalError
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> InternalError -> ShowS
showsPrec :: Int -> InternalError -> ShowS
$cshow :: InternalError -> String
show :: InternalError -> String
$cshowList :: [InternalError] -> ShowS
showList :: [InternalError] -> ShowS
Show, Show InternalError
Typeable InternalError
(Typeable InternalError, Show InternalError) =>
(InternalError -> SomeException)
-> (SomeException -> Maybe InternalError)
-> (InternalError -> String)
-> Exception InternalError
SomeException -> Maybe InternalError
InternalError -> String
InternalError -> SomeException
forall e.
(Typeable e, Show e) =>
(e -> SomeException)
-> (SomeException -> Maybe e) -> (e -> String) -> Exception e
$ctoException :: InternalError -> SomeException
toException :: InternalError -> SomeException
$cfromException :: SomeException -> Maybe InternalError
fromException :: SomeException -> Maybe InternalError
$cdisplayException :: InternalError -> String
displayException :: InternalError -> String
Exception)

writeToConnection :: Serializable packet => ConnectionState -> packet -> IO ()
writeToConnection :: forall packet.
Serializable packet =>
ConnectionState -> packet -> IO ()
writeToConnection MkConnectionState{ProtocolRevision
revision :: ProtocolRevision
revision :: ConnectionState -> ProtocolRevision
revision, Buffer
buffer :: Buffer
buffer :: ConnectionState -> Buffer
buffer} packet
packet =
  (Buffer -> ByteString -> IO ()
writeSock Buffer
buffer (ByteString -> IO ()) -> (packet -> ByteString) -> packet -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Builder -> ByteString
toLazyByteString (Builder -> ByteString)
-> (packet -> Builder) -> packet -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ProtocolRevision -> packet -> Builder
forall chType.
Serializable chType =>
ProtocolRevision -> chType -> Builder
serialize ProtocolRevision
revision) packet
packet

writeToConnectionEncode :: ConnectionState -> (ProtocolRevision -> Builder) -> IO ()
writeToConnectionEncode :: ConnectionState -> (ProtocolRevision -> Builder) -> IO ()
writeToConnectionEncode MkConnectionState{ProtocolRevision
revision :: ConnectionState -> ProtocolRevision
revision :: ProtocolRevision
revision, Buffer
buffer :: ConnectionState -> Buffer
buffer :: Buffer
buffer} ProtocolRevision -> Builder
serializer =
  (Buffer -> ByteString -> IO ()
writeSock Buffer
buffer (ByteString -> IO ())
-> (Builder -> ByteString) -> Builder -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Builder -> ByteString
toLazyByteString) (ProtocolRevision -> Builder
serializer ProtocolRevision
revision)

data Connection where MkConnection :: (MVar ConnectionState) -> Connection

data ConnectionState = MkConnectionState
  { ConnectionState -> ChString
user     :: ChString
  , ConnectionState -> ChString
hostname :: ChString
  , ConnectionState -> ChString
os_user  :: ChString
  , ConnectionState -> Buffer
buffer   :: Buffer
  , ConnectionState -> ProtocolRevision
revision :: ProtocolRevision
  , ConnectionState -> ConnectionArgs
creds    :: ConnectionArgs
  }

data Buffer = MkBuffer
  { Buffer -> IO ByteString
readSock :: IO BS.ByteString
  , Buffer -> ByteString -> IO ()
writeSock :: BSL.ByteString -> IO ()
  , Buffer -> IO ()
closeSock :: IO ()
  , Buffer -> IORef ByteString
buff :: IORef BS.ByteString
  }

flushBuffer :: Buffer -> IO ()
flushBuffer :: Buffer -> IO ()
flushBuffer MkBuffer{IORef ByteString
buff :: Buffer -> IORef ByteString
buff :: IORef ByteString
buff} = IORef ByteString -> ByteString -> IO ()
forall a. IORef a -> a -> IO ()
atomicWriteIORef IORef ByteString
buff ByteString
""

rawBufferRead :: Buffer -> Get packet -> IO packet
rawBufferRead :: forall packet. Buffer -> Get packet -> IO packet
rawBufferRead buffer :: Buffer
buffer@MkBuffer{IO ()
IO ByteString
IORef ByteString
ByteString -> IO ()
writeSock :: Buffer -> ByteString -> IO ()
readSock :: Buffer -> IO ByteString
closeSock :: Buffer -> IO ()
buff :: Buffer -> IORef ByteString
readSock :: IO ByteString
writeSock :: ByteString -> IO ()
closeSock :: IO ()
buff :: IORef ByteString
..} Get packet
parser = Decoder packet -> IO packet
forall packet. Decoder packet -> IO packet
runBufferReader (Get packet -> Decoder packet
forall a. Get a -> Decoder a
runGetIncremental Get packet
parser)
  where
  runBufferReader :: Decoder packet -> IO packet
  runBufferReader :: forall packet. Decoder packet -> IO packet
runBufferReader = \case
    (Partial Maybe ByteString -> Decoder packet
decoder) -> IO ByteString
readBuffer IO ByteString -> (ByteString -> IO packet) -> IO packet
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Decoder packet -> IO packet
forall packet. Decoder packet -> IO packet
runBufferReader (Decoder packet -> IO packet)
-> (ByteString -> Decoder packet) -> ByteString -> IO packet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe ByteString -> Decoder packet
decoder (Maybe ByteString -> Decoder packet)
-> (ByteString -> Maybe ByteString) -> ByteString -> Decoder packet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just
    (Done ByteString
leftover ByteOffset
_consumed packet
packet) -> packet
packet packet -> IO ByteString -> IO packet
forall a b. a -> IO b -> IO a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ IORef ByteString
-> (ByteString -> (ByteString, ByteString)) -> IO ByteString
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef IORef ByteString
buff (ByteString
leftover,)
    (Fail ByteString
_leftover ByteOffset
_consumed String
msg) -> InternalError -> IO packet
forall e a. Exception e => e -> IO a
throwIO  (String -> InternalError
DeserializationError String
msg)

  readBuffer :: IO BS.ByteString
  readBuffer :: IO ByteString
readBuffer =
    IORef ByteString -> IO ByteString
forall a. IORef a -> IO a
readIORef IORef ByteString
buff
      IO ByteString -> (ByteString -> IO ByteString) -> IO ByteString
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (\ByteString
currentBuffer ->
        case ByteString -> Int
BS.length ByteString
currentBuffer of
          Int
0 -> IO ByteString
readSock
          Int
_ -> Buffer -> IO ()
flushBuffer Buffer
buffer IO () -> IO ByteString -> IO ByteString
forall a b. IO a -> IO b -> IO b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
currentBuffer
      )




-- * Initialization

{- |
  See `defaultConnectionArgs` for documentation
-}
data ConnectionArgs = MkConnectionArgs
  { ConnectionArgs -> Text
user :: Text
  , ConnectionArgs -> Text
pass :: Text
  , ConnectionArgs -> Text
db   :: Text
  , ConnectionArgs -> String
host :: HostName
  , ConnectionArgs -> Maybe String
mPort :: Maybe ServiceName
  , ConnectionArgs -> Bool
isTLS :: Bool
  , ConnectionArgs -> String -> SockAddr -> Socket -> IO Buffer
initBuffer :: HostName -> SockAddr -> Socket -> IO Buffer
  }

{- |
  Default connection settings which follows __clickhouse-client__ defaults

  Use `setUser`, `setPassword`, `setHost`, `setPort`, `setDatabase`
  to modify connection defaults.

  Or 'setSecure', 'overrideTLS' to configure TLS connection
-}
defaultConnectionArgs :: ConnectionArgs
defaultConnectionArgs :: ConnectionArgs
defaultConnectionArgs = MkConnectionArgs
  { user :: Text
user = Text
"default"
  , pass :: Text
pass = Text
""
  , host :: String
host = String
"localhost"
  , db :: Text
db   = Text
"default"
  , isTLS :: Bool
isTLS = Bool
False
  , mPort :: Maybe String
mPort = Maybe String
forall a. Maybe a
Nothing
  , initBuffer :: String -> SockAddr -> Socket -> IO Buffer
initBuffer  = \String
_hostname SockAddr
addrAddress Socket
sock -> do
      Socket -> SocketOption -> Int -> IO ()
setSocketOption Socket
sock SocketOption
Sock.NoDelay Int
1
      Socket -> SocketOption -> Int -> IO ()
setSocketOption Socket
sock SocketOption
Sock.KeepAlive Int
1
      Socket -> SockAddr -> IO ()
connect Socket
sock SockAddr
addrAddress
      IORef ByteString
buff <- ByteString -> IO (IORef ByteString)
forall a. a -> IO (IORef a)
newIORef ByteString
""
      Buffer -> IO Buffer
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
        MkBuffer
          { writeSock :: ByteString -> IO ()
writeSock = \ByteString
bs -> Socket -> ByteString -> IO ()
sendAll Socket
sock ByteString
bs
          , readSock :: IO ByteString
readSock  = Socket -> Int -> IO ByteString
recv Socket
sock Int
4096
          , closeSock :: IO ()
closeSock = Socket -> IO ()
close Socket
sock
          , IORef ByteString
buff :: IORef ByteString
buff :: IORef ByteString
buff
          }
  }


{- |
  Overrides default user __"default"__
-}
setUser :: Text -> ConnectionArgs -> ConnectionArgs
setUser :: Text -> ConnectionArgs -> ConnectionArgs
setUser Text
new MkConnectionArgs{Bool
String
Maybe String
Text
String -> SockAddr -> Socket -> IO Buffer
user :: ConnectionArgs -> Text
pass :: ConnectionArgs -> Text
db :: ConnectionArgs -> Text
host :: ConnectionArgs -> String
mPort :: ConnectionArgs -> Maybe String
isTLS :: ConnectionArgs -> Bool
initBuffer :: ConnectionArgs -> String -> SockAddr -> Socket -> IO Buffer
user :: Text
pass :: Text
db :: Text
host :: String
mPort :: Maybe String
isTLS :: Bool
initBuffer :: String -> SockAddr -> Socket -> IO Buffer
..} = MkConnectionArgs{user :: Text
user=Text
new, Bool
String
Maybe String
Text
String -> SockAddr -> Socket -> IO Buffer
pass :: Text
db :: Text
host :: String
mPort :: Maybe String
isTLS :: Bool
initBuffer :: String -> SockAddr -> Socket -> IO Buffer
pass :: Text
db :: Text
host :: String
mPort :: Maybe String
isTLS :: Bool
initBuffer :: String -> SockAddr -> Socket -> IO Buffer
..}

{- |
  Overrides default password __""__
-}
setPassword :: Text -> ConnectionArgs -> ConnectionArgs
setPassword :: Text -> ConnectionArgs -> ConnectionArgs
setPassword Text
new MkConnectionArgs{Bool
String
Maybe String
Text
String -> SockAddr -> Socket -> IO Buffer
user :: ConnectionArgs -> Text
pass :: ConnectionArgs -> Text
db :: ConnectionArgs -> Text
host :: ConnectionArgs -> String
mPort :: ConnectionArgs -> Maybe String
isTLS :: ConnectionArgs -> Bool
initBuffer :: ConnectionArgs -> String -> SockAddr -> Socket -> IO Buffer
user :: Text
pass :: Text
db :: Text
host :: String
mPort :: Maybe String
isTLS :: Bool
initBuffer :: String -> SockAddr -> Socket -> IO Buffer
..} = MkConnectionArgs{pass :: Text
pass=Text
new, Bool
String
Maybe String
Text
String -> SockAddr -> Socket -> IO Buffer
user :: Text
db :: Text
host :: String
mPort :: Maybe String
isTLS :: Bool
initBuffer :: String -> SockAddr -> Socket -> IO Buffer
user :: Text
db :: Text
host :: String
mPort :: Maybe String
isTLS :: Bool
initBuffer :: String -> SockAddr -> Socket -> IO Buffer
..}

{- |
  Overrides default hostname __"localhost"__
-}
setHost :: HostName -> ConnectionArgs -> ConnectionArgs
setHost :: String -> ConnectionArgs -> ConnectionArgs
setHost String
new MkConnectionArgs{Bool
String
Maybe String
Text
String -> SockAddr -> Socket -> IO Buffer
user :: ConnectionArgs -> Text
pass :: ConnectionArgs -> Text
db :: ConnectionArgs -> Text
host :: ConnectionArgs -> String
mPort :: ConnectionArgs -> Maybe String
isTLS :: ConnectionArgs -> Bool
initBuffer :: ConnectionArgs -> String -> SockAddr -> Socket -> IO Buffer
user :: Text
pass :: Text
db :: Text
host :: String
mPort :: Maybe String
isTLS :: Bool
initBuffer :: String -> SockAddr -> Socket -> IO Buffer
..} = MkConnectionArgs{host :: String
host=String
new, Bool
Maybe String
Text
String -> SockAddr -> Socket -> IO Buffer
user :: Text
pass :: Text
db :: Text
mPort :: Maybe String
isTLS :: Bool
initBuffer :: String -> SockAddr -> Socket -> IO Buffer
user :: Text
pass :: Text
db :: Text
mPort :: Maybe String
isTLS :: Bool
initBuffer :: String -> SockAddr -> Socket -> IO Buffer
..}

{- |
  Overrides default port __9000__ (or __9443__ for TLS)
-}
setPort :: ServiceName -> ConnectionArgs -> ConnectionArgs
setPort :: String -> ConnectionArgs -> ConnectionArgs
setPort String
new MkConnectionArgs{Bool
String
Maybe String
Text
String -> SockAddr -> Socket -> IO Buffer
user :: ConnectionArgs -> Text
pass :: ConnectionArgs -> Text
db :: ConnectionArgs -> Text
host :: ConnectionArgs -> String
mPort :: ConnectionArgs -> Maybe String
isTLS :: ConnectionArgs -> Bool
initBuffer :: ConnectionArgs -> String -> SockAddr -> Socket -> IO Buffer
user :: Text
pass :: Text
db :: Text
host :: String
mPort :: Maybe String
isTLS :: Bool
initBuffer :: String -> SockAddr -> Socket -> IO Buffer
..} = MkConnectionArgs{mPort :: Maybe String
mPort=String -> Maybe String
forall a. a -> Maybe a
Just String
new, Bool
String
Text
String -> SockAddr -> Socket -> IO Buffer
user :: Text
pass :: Text
db :: Text
host :: String
isTLS :: Bool
initBuffer :: String -> SockAddr -> Socket -> IO Buffer
user :: Text
pass :: Text
db :: Text
host :: String
isTLS :: Bool
initBuffer :: String -> SockAddr -> Socket -> IO Buffer
..} 

{- |
  Overrides default database __"default"__
-}
setDatabase :: Text -> ConnectionArgs -> ConnectionArgs
setDatabase :: Text -> ConnectionArgs -> ConnectionArgs
setDatabase Text
new MkConnectionArgs{Bool
String
Maybe String
Text
String -> SockAddr -> Socket -> IO Buffer
user :: ConnectionArgs -> Text
pass :: ConnectionArgs -> Text
db :: ConnectionArgs -> Text
host :: ConnectionArgs -> String
mPort :: ConnectionArgs -> Maybe String
isTLS :: ConnectionArgs -> Bool
initBuffer :: ConnectionArgs -> String -> SockAddr -> Socket -> IO Buffer
user :: Text
pass :: Text
db :: Text
host :: String
mPort :: Maybe String
isTLS :: Bool
initBuffer :: String -> SockAddr -> Socket -> IO Buffer
..} = MkConnectionArgs{db :: Text
db=Text
new, Bool
String
Maybe String
Text
String -> SockAddr -> Socket -> IO Buffer
user :: Text
pass :: Text
host :: String
mPort :: Maybe String
isTLS :: Bool
initBuffer :: String -> SockAddr -> Socket -> IO Buffer
user :: Text
pass :: Text
host :: String
mPort :: Maybe String
isTLS :: Bool
initBuffer :: String -> SockAddr -> Socket -> IO Buffer
..}

overrideNetwork
  :: Bool
  -> (HostName -> SockAddr -> Socket -> IO Buffer)
  -> (ConnectionArgs -> ConnectionArgs)
overrideNetwork :: Bool
-> (String -> SockAddr -> Socket -> IO Buffer)
-> ConnectionArgs
-> ConnectionArgs
overrideNetwork Bool
new String -> SockAddr -> Socket -> IO Buffer
new2 = \MkConnectionArgs{Bool
String
Maybe String
Text
String -> SockAddr -> Socket -> IO Buffer
user :: ConnectionArgs -> Text
pass :: ConnectionArgs -> Text
db :: ConnectionArgs -> Text
host :: ConnectionArgs -> String
mPort :: ConnectionArgs -> Maybe String
isTLS :: ConnectionArgs -> Bool
initBuffer :: ConnectionArgs -> String -> SockAddr -> Socket -> IO Buffer
user :: Text
pass :: Text
db :: Text
host :: String
mPort :: Maybe String
isTLS :: Bool
initBuffer :: String -> SockAddr -> Socket -> IO Buffer
..} ->
  MkConnectionArgs{isTLS :: Bool
isTLS=Bool
new, initBuffer :: String -> SockAddr -> Socket -> IO Buffer
initBuffer=String -> SockAddr -> Socket -> IO Buffer
new2, String
Maybe String
Text
user :: Text
pass :: Text
db :: Text
host :: String
mPort :: Maybe String
user :: Text
pass :: Text
db :: Text
host :: String
mPort :: Maybe String
..}