-- | Several statements in one round trip, each with its own result.
--
-- A 'BatchQuery' pairs each statement with what to do with its results,
-- and 'runBatchQuery' sends them all in one request and reads the
-- results in order:
--
-- > (updated, Only total, items) <- runBatchQuery conn $ (,,)
-- >   <$> batchExecute ("UPDATE item SET price = price * 1.1 WHERE category = " <> param cat)
-- >   <*> batchQuery1 ("SELECT SUM(price) FROM item WHERE category = " <> param cat)
-- >   <*> batchQuery ("SELECT id, name FROM item WHERE category = " <> param cat)
-- > -- updated :: Int64, total :: Double, items :: [(Int32, Text)]
--
-- Each statement is its own @sp_executesql@ call, all of them in one RPC
-- message (as 'executeStatements' sends them): the server compiles and runs
-- them in order, and the response carries each one's results between its
-- own DONEPROC markers, so a statement's consumer sees exactly its own
-- result sets and counts, whatever the statements before it returned.  A
-- statement with no parameters is an @sp_executesql@ call too, not a SQL
-- batch, so a @SET@ option it sets or a @#temp@ table it creates lasts
-- only until it returns; statements that must share such state go in one
-- 'Query' (its text joined, as 'Query' is a 'Monoid') to
-- 'queryResultSets' instead.  @'executeStatements' conn qs@ is
-- @'runBatchQuery' conn ('traverse' 'batchExecute' qs)@ with the
-- counting done in bulk, and 'executeStatementsOutcomes' the way to see which
-- statements failed rather than the first that did.
--
-- The consumers mirror "Hmssql.Query": 'batchExecute' for 'execute',
-- 'batchQuery' for 'query' (with 'batchQuery1', 'batchQueryMay',
-- 'batchQueryFold', 'batchQueryS' and the @With@ variants),
-- 'batchQueryMulti' for 'queryMulti', 'batchQueryResultSets' for
-- 'queryResultSets'.  A statement's results are read as its consumer asks
-- for them, and what it leaves unread is skipped before the next
-- statement's, so a 'BatchQuery' holds no more memory than its results;
-- 'batchQueryS' streams the last statement's rows after 'runBatchQuery'
-- returns, as 'queryS' does.
--
-- = Errors
--
-- The server runs the statements in order, and __statements after a
-- failed one may have run__: an error that only ends its statement (a
-- constraint violation, say) lets the server go on to the next; an error
-- that aborts the batch (@THROW@, a conversion error, a deadlock, any
-- error under @SET XACT_ABORT ON@) ends the request there.  Outside a
-- transaction each statement is committed as it completes.  The first
-- error is thrown as 'SqlServerError' with its statement's text, once the
-- whole response (and the messages after it, for a batch over about 1
-- MiB, as 'executeStatements' explains) has been read, so the connection stays
-- usable.  For all-or-nothing, run the batch inside
-- 'Hmssql.Transaction.withTransaction' and let the exception roll it
-- back.
--
-- A statement that can't be sent (a value that can't be encoded, a
-- placeholder with no parameter) throws 'EncodeError', naming its
-- position, before its message is sent; so does a 'batchQueryS' that
-- isn't the last statement.
--
-- The shape is hpgsql's @Hpgsql.Pipeline@.  It is called a batch because
-- that is what TDS has: several statements in one request, run in order,
-- each committed on its own, later ones running after an earlier one's
-- error.
--
-- Only the last statement may stream.  'batchQueryS' gives a stream to read after
-- 'runBatchQuery' returns, so a stream in the middle would leave the statements
-- after it as actions you would have to run later, in order, on one thread.  With
-- only the last one streaming, there is no such order to get wrong.
module Hmssql.BatchQuery
  ( BatchQuery
  , runBatchQuery

    -- * Statements
  , batchExecute
  , batchExecute_

    -- * Queries
  , batchQuery
  , batchQueryWith
  , batchQueryMWith
  , batchQuery1
  , batchQuery1With
  , batchQueryMay
  , batchQueryMayWith

    -- * Folding queries
  , batchQueryFold
  , batchQueryFoldWith

    -- * Streaming queries
  , batchQueryS
  , batchQuerySWith
  , batchQuerySMWith

    -- * Multiple result sets
  , batchQueryMulti
  , batchQueryMultiWith
  , batchQueryMultiMWith
  , batchQueryResultSets
  ) where

import Control.Exception
import Control.Monad (forM_, unless, void, when)
import Data.Int (Int64)
import Data.IORef
import Data.Text (Text)
import qualified Data.Text as T
import Data.Void (Void)
import Streaming (Of, Stream)

import Hmssql.Connection.Internal (Connection, arResultSet, abandonResponse, connectionContext, drainActiveResponse, readingToEnd, startRequestFor, withConnectionLock)
import Hmssql.Decoding
import Hmssql.Errors
import Hmssql.Query
import Hmssql.Query.Internal
import Hmssql.Query.Types
import Hmssql.SqlValue (SqlRow)
import Hmssql.Tds.Context (TdsContext)
import Hmssql.Tds.Token (ServerMessage)

-- | Several statements sent in one request, each paired with what to do
-- with its results: see "Hmssql.BatchQuery".
--
-- An 'Applicative', not a 'Monad': every statement is known before the
-- batch is sent, so a later statement can't depend on an earlier one's
-- rows (the 'ResultSets' monad of one statement can, within it).
data BatchQuery a = BatchQuery
  ([BatchItem] -> [BatchItem])
    -- ^ The statements, in order (a difference list).
  (BatchRun -> IO a)
    -- ^ Read each statement's results in turn ('nextStatement').

instance Functor BatchQuery where
  fmap :: forall a b. (a -> b) -> BatchQuery a -> BatchQuery b
fmap a -> b
f (BatchQuery [BatchItem] -> [BatchItem]
items BatchRun -> IO a
run) = ([BatchItem] -> [BatchItem]) -> (BatchRun -> IO b) -> BatchQuery b
forall a.
([BatchItem] -> [BatchItem]) -> (BatchRun -> IO a) -> BatchQuery a
BatchQuery [BatchItem] -> [BatchItem]
items ((a -> b) -> IO a -> IO b
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f (IO a -> IO b) -> (BatchRun -> IO a) -> BatchRun -> IO b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BatchRun -> IO a
run)

instance Applicative BatchQuery where
  pure :: forall a. a -> BatchQuery a
pure a
a = ([BatchItem] -> [BatchItem]) -> (BatchRun -> IO a) -> BatchQuery a
forall a.
([BatchItem] -> [BatchItem]) -> (BatchRun -> IO a) -> BatchQuery a
BatchQuery [BatchItem] -> [BatchItem]
forall a. a -> a
id (\BatchRun
_ -> a -> IO a
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a)
  BatchQuery [BatchItem] -> [BatchItem]
items BatchRun -> IO (a -> b)
f <*> :: forall a b. BatchQuery (a -> b) -> BatchQuery a -> BatchQuery b
<*> BatchQuery [BatchItem] -> [BatchItem]
items' BatchRun -> IO a
x = ([BatchItem] -> [BatchItem]) -> (BatchRun -> IO b) -> BatchQuery b
forall a.
([BatchItem] -> [BatchItem]) -> (BatchRun -> IO a) -> BatchQuery a
BatchQuery ([BatchItem] -> [BatchItem]
items ([BatchItem] -> [BatchItem])
-> ([BatchItem] -> [BatchItem]) -> [BatchItem] -> [BatchItem]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [BatchItem] -> [BatchItem]
items') ((BatchRun -> IO b) -> BatchQuery b)
-> (BatchRun -> IO b) -> BatchQuery b
forall a b. (a -> b) -> a -> b
$ \BatchRun
br -> do
    g <- BatchRun -> IO (a -> b)
f BatchRun
br
    g <$> x br

-- | One statement of a batch.
data BatchItem = BatchItem
  { BatchItem -> Query
itemQuery :: !Query
  , BatchItem -> Bool
itemStreams :: !Bool
    -- ^ Its rows are read after 'runBatchQuery' returns (a
    -- 'batchQueryS' stream), so it must be the last.
  }

-- | A batch being run: the connection, its state, and the scope the
-- statement being read is bounded by.
data BatchRun = BatchRun
  { BatchRun -> Connection
brConn :: !Connection
  , BatchRun -> IORef BatchState
brState :: !(IORef BatchState)
  , BatchRun -> BatchScope
brScope :: !BatchScope
  }

-- | Where a batch is: the reader of the message being read, if one has
-- been sent; how many statements have been opened (the position of the
-- next, from 0); the texts of the message's statements not opened yet;
-- the queries after the message; and whether a statement is open (was
-- given to its consumer, and hasn't been skipped to its end).
data BatchState = BatchState
  { BatchState -> Maybe ResponseReader
bsReader :: !(Maybe ResponseReader)
  , BatchState -> Int
bsIndex :: !Int
  , BatchState -> [Text]
bsInMessage :: [Text]
  , BatchState -> [Query]
bsRest :: [Query]
  , BatchState -> Bool
bsOpen :: !Bool
  }

-- | Run the batch: send the statements, read each one's results as its
-- consumer asks, and read the response to its end.  Nothing is sent for
-- a batch of no statements.  Statements that don't fit one message go
-- in the next, sent once the previous response has been read; a
-- statement's consumer runs when its results are reached.
runBatchQuery :: Connection -> BatchQuery a -> IO a
runBatchQuery :: forall a. Connection -> BatchQuery a -> IO a
runBatchQuery Connection
conn (BatchQuery [BatchItem] -> [BatchItem]
mkItems BatchRun -> IO a
run) = case [BatchItem] -> [BatchItem]
mkItems [] of
  [] -> do
    br <- [Query] -> IO BatchRun
newRun []
    run br
  [BatchItem]
items -> Connection -> IO a -> IO a
forall a. Connection -> IO a -> IO a
withConnectionLock Connection
conn (IO a -> IO a) -> IO a -> IO a
forall a b. (a -> b) -> a -> b
$ do
    [(Int, BatchItem)] -> ((Int, BatchItem) -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Int] -> [BatchItem] -> [(Int, BatchItem)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 :: Int ..] ([BatchItem] -> [BatchItem]
forall a. HasCallStack => [a] -> [a]
init [BatchItem]
items)) (((Int, BatchItem) -> IO ()) -> IO ())
-> ((Int, BatchItem) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(Int
i, BatchItem
item) -> Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (BatchItem -> Bool
itemStreams BatchItem
item) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ EncodeError -> IO ()
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO (EncodeError -> IO ()) -> EncodeError -> IO ()
forall a b. (a -> b) -> a -> b
$ [Char] -> EncodeError
EncodeError ([Char] -> EncodeError) -> [Char] -> EncodeError
forall a b. (a -> b) -> a -> b
$
      [Char]
"statement " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
i [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" (from 0) streams its rows (batchQueryS), so it must be the last of the batch"
    br <- [Query] -> IO BatchRun
newRun ((BatchItem -> Query) -> [BatchItem] -> [Query]
forall a b. (a -> b) -> [a] -> [b]
map BatchItem -> Query
itemQuery [BatchItem]
items)
    let body = do
          a <- BatchRun -> IO a
run BatchRun
br
          unless (itemStreams (last items)) $ finishBatch br
          pure a
    body `catch` \(SomeException
e :: SomeException) -> do
      case SomeException -> Maybe SomeAsyncException
forall e. Exception e => SomeException -> Maybe e
fromException SomeException
e of
        -- Interrupted: the message's response is marked abandoned and its
        -- statement canceled at once, and the next request drains it, as
        -- for 'query' ('Hmssql.Connection.Internal.readingToEnd').
        Just (SomeAsyncException
_ :: SomeAsyncException) -> BatchRun -> IO ()
abandonCurrent BatchRun
br
        Maybe SomeAsyncException
Nothing -> BatchRun -> IO ()
abandonBatch BatchRun
br
      SomeException -> IO a
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO SomeException
e
  where
    newRun :: [Query] -> IO BatchRun
newRun [Query]
qs = do
      st <- BatchState -> IO (IORef BatchState)
forall a. a -> IO (IORef a)
newIORef (Maybe ResponseReader
-> Int -> [Text] -> [Query] -> Bool -> BatchState
BatchState Maybe ResponseReader
forall a. Maybe a
Nothing Int
0 [] [Query]
qs Bool
False)
      ended <- newIORef False
      stmt <- newIORef T.empty
      let br = Connection -> IORef BatchState -> BatchScope -> BatchRun
BatchRun Connection
conn IORef BatchState
st (IORef Bool
-> IORef Text -> (ServerMessage -> IO Void) -> BatchScope
BatchScope IORef Bool
ended IORef Text
stmt (BatchRun -> ServerMessage -> IO Void
failBatch BatchRun
br))
      pure br

-- | After an interruption: the response of the message being read, if it
-- hasn't ended, is canceled at once and left to the next request to
-- drain ('abandonResponse').
abandonCurrent :: BatchRun -> IO ()
abandonCurrent :: BatchRun -> IO ()
abandonCurrent BatchRun{IORef BatchState
Connection
BatchScope
brConn :: BatchRun -> Connection
brState :: BatchRun -> IORef BatchState
brScope :: BatchRun -> BatchScope
brConn :: Connection
brState :: IORef BatchState
brScope :: BatchScope
..} = do
  st <- IORef BatchState -> IO BatchState
forall a. IORef a -> IO a
readIORef IORef BatchState
brState
  forM_ (bsReader st) $ \ResponseReader
rr -> Connection -> ActiveResponse -> IO ()
abandonResponse Connection
brConn (ResponseReader -> ActiveResponse
respActive ResponseReader
rr)

-- | Run the action; a synchronous exception runs the handler (an
-- asynchronous one is rethrown as it is, after the Attention 'nextStep'
-- has sent).
catchSyncException :: IO a -> (SomeException -> IO a) -> IO a
catchSyncException :: forall a. IO a -> (SomeException -> IO a) -> IO a
catchSyncException IO a
act SomeException -> IO a
handler = IO a
act IO a -> (SomeException -> IO a) -> IO a
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` \SomeException
e -> case SomeException -> Maybe SomeAsyncException
forall e. Exception e => SomeException -> Maybe e
fromException SomeException
e of
  Just (SomeAsyncException
_ :: SomeAsyncException) -> SomeException -> IO a
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO SomeException
e
  Maybe SomeAsyncException
Nothing -> SomeException -> IO a
handler SomeException
e

-- | The reader of the next statement's results, bounded to them: the
-- statement before it is read to its end first, and the next message
-- sent when the current one's statements are all read.
nextStatement :: BatchRun -> IO ResponseReader
nextStatement :: BatchRun -> IO ResponseReader
nextStatement br :: BatchRun
br@BatchRun{IORef BatchState
Connection
BatchScope
brConn :: BatchRun -> Connection
brState :: BatchRun -> IORef BatchState
brScope :: BatchRun -> BatchScope
brConn :: Connection
brState :: IORef BatchState
brScope :: BatchScope
..} = do
  st <- IORef BatchState -> IO BatchState
forall a. IORef a -> IO a
readIORef IORef BatchState
brState
  when (bsOpen st) $ closeStatement br st
  case (bsReader st, bsInMessage st) of
    (Just ResponseReader
rr, Text
stmt : [Text]
more) -> do
      IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (BatchScope -> IORef Bool
scopeEnded BatchScope
brScope) Bool
False
      IORef Text -> Text -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (BatchScope -> IORef Text
scopeStatement BatchScope
brScope) Text
stmt
      IORef BatchState -> BatchState -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef BatchState
brState BatchState
st { bsInMessage = more, bsIndex = bsIndex st + 1, bsOpen = True }
      ResponseReader -> IO ResponseReader
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ResponseReader
rr
    (Maybe ResponseReader
_, Text
_ : [Text]
_) -> ProtocolError -> IO ResponseReader
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO (ProtocolError -> IO ResponseReader)
-> ProtocolError -> IO ResponseReader
forall a b. (a -> b) -> a -> b
$ [Char] -> ProtocolError
ProtocolError [Char]
"batch: statements to read, but no response"
    (Maybe ResponseReader
mrr, []) -> do
      Maybe ResponseReader -> (ResponseReader -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ Maybe ResponseReader
mrr (BatchRun -> ResponseReader -> IO ()
finishMessage BatchRun
br)
      case BatchState -> [Query]
bsRest BatchState
st of
        [] -> ProtocolError -> IO ResponseReader
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO (ProtocolError -> IO ResponseReader)
-> ProtocolError -> IO ResponseReader
forall a b. (a -> b) -> a -> b
$ [Char] -> ProtocolError
ProtocolError [Char]
"batch: more statements read than were sent"
        [Query]
qs -> do
          ctx <- Connection -> IO TdsContext
connectionContext Connection
brConn
          (stmts, _, rest, ()) <- either (throwIO . EncodeError) pure (takeStatements queryStatement ctx () (bsIndex st) qs)
          texts <- mapM (evaluate . fst) stmts
          Request{..} <- statementsRequest stmts
          -- Masked from the send to recording the reader, so that an
          -- interruption finds the response ('abandonCurrent').
          mask_ $ do
            ar <- startRequestFor brConn requestStatement requestPacketType requestBody
            rr <- (\ResponseReader
r -> ResponseReader
r { respBatch = Just brScope }) <$> responseReader brConn ar Nothing
            writeIORef brState st { bsReader = Just rr, bsInMessage = texts, bsRest = rest, bsOpen = False }
          nextStatement br

-- | Read the open statement to its end (its consumer may have stopped
-- before it), and check it ended: a response that ended before its
-- DONEPROC without an error is a protocol error.
closeStatement :: BatchRun -> BatchState -> IO ()
closeStatement :: BatchRun -> BatchState -> IO ()
closeStatement BatchRun{IORef BatchState
Connection
BatchScope
brConn :: BatchRun -> Connection
brState :: BatchRun -> IORef BatchState
brScope :: BatchRun -> BatchScope
brConn :: Connection
brState :: IORef BatchState
brScope :: BatchScope
..} BatchState
st = Maybe ResponseReader -> (ResponseReader -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (BatchState -> Maybe ResponseReader
bsReader BatchState
st) ((ResponseReader -> IO ()) -> IO ())
-> (ResponseReader -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ResponseReader
rr -> do
  let skip :: IO ()
skip = ResponseReader -> RowRead () -> IO (ResponseStep ())
forall r. ResponseReader -> RowRead r -> IO (ResponseStep r)
readStep ResponseReader
rr RowRead ()
ignoreRows IO (ResponseStep ()) -> (ResponseStep () -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        ResponseStep ()
StepEnd -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        ResponseStep ()
_ -> IO ()
skip
  IO ()
skip
  ended <- IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef (BatchScope -> IORef Bool
scopeEnded BatchScope
brScope)
  unless ended $ throwIO $ ProtocolError $
    "batch: the response ended before statement " ++ show (bsIndex st - 1) ++ " (from 0) completed, and no error was reported"
  -- Each statement supplies its own metadata. Clear at DONEPROC, not when
  -- opening a statement: reset replay can already hold its first COLMETADATA.
  writeIORef (arResultSet (respActive rr)) Nothing
  writeIORef brState st { bsOpen = False }

-- | Read the message's response to its end, past the last statement's
-- DONEPROC: a failed DONE is thrown here.
finishMessage :: BatchRun -> ResponseReader -> IO ()
finishMessage :: BatchRun -> ResponseReader -> IO ()
finishMessage BatchRun{IORef BatchState
Connection
BatchScope
brConn :: BatchRun -> Connection
brState :: BatchRun -> IORef BatchState
brScope :: BatchRun -> BatchScope
brConn :: Connection
brState :: IORef BatchState
brScope :: BatchScope
..} ResponseReader
rr = do
  stmt <- IORef Text -> IO Text
forall a. IORef a -> IO a
readIORef (BatchScope -> IORef Text
scopeStatement BatchScope
brScope)
  drainResponse rr { respBatch = Nothing, respStatement = Just stmt }

-- | After the last consumer: read the open statement to its end and the
-- response after it.  Every statement must have been read.
finishBatch :: BatchRun -> IO ()
finishBatch :: BatchRun -> IO ()
finishBatch br :: BatchRun
br@BatchRun{IORef BatchState
Connection
BatchScope
brConn :: BatchRun -> Connection
brState :: BatchRun -> IORef BatchState
brScope :: BatchRun -> BatchScope
brConn :: Connection
brState :: IORef BatchState
brScope :: BatchScope
..} = do
  st <- IORef BatchState -> IO BatchState
forall a. IORef a -> IO a
readIORef IORef BatchState
brState
  when (bsOpen st) $ closeStatement br st
  unless (null (bsInMessage st) && null (bsRest st)) $
    throwIO $ ProtocolError "batch: statements were sent but their results not read"
  forM_ (bsReader st) (finishMessage br)

-- | An ERROR token arrived while a statement was being read: the rest of
-- the response is read, the messages after it are sent and read (their
-- statements run, as they would have in one message), and the error is
-- thrown for its statement.  A later message that can't be encoded isn't
-- sent, nor those after it; a connection the server closes after the
-- error (severity 20 and up) throws the error too.
failBatch :: BatchRun -> ServerMessage -> IO Void
failBatch :: BatchRun -> ServerMessage -> IO Void
failBatch BatchRun{IORef BatchState
Connection
BatchScope
brConn :: BatchRun -> Connection
brState :: BatchRun -> IORef BatchState
brScope :: BatchRun -> BatchScope
brConn :: Connection
brState :: IORef BatchState
brScope :: BatchScope
..} ServerMessage
msg = do
  stmt <- IORef Text -> IO Text
forall a. IORef a -> IO a
readIORef (BatchScope -> IORef Text
scopeStatement BatchScope
brScope)
  st <- readIORef brState
  let failure = ServerMessage -> Maybe Text -> SqlServerError
SqlServerError ServerMessage
msg (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
stmt)
      drain ActiveResponse
ar = Connection -> ActiveResponse -> IO ()
drainActiveResponse Connection
brConn ActiveResponse
ar IO () -> (ConnectionClosed -> IO ()) -> IO ()
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` \ConnectionClosed
e -> case ConnectionClosed
e of
        ConnectionClosedBy CloseReason
reason | Bool -> Bool
not (CloseReason -> Bool
byClient CloseReason
reason) -> SqlServerError -> IO ()
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO SqlServerError
failure
        ConnectionClosed
_ -> ConnectionClosed -> IO ()
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO ConnectionClosed
e
  forM_ (bsReader st) $ \ResponseReader
rr -> do
    ActiveResponse -> IO ()
drain (ResponseReader -> ActiveResponse
respActive ResponseReader
rr)
    IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (ResponseReader -> IORef Bool
respEnded ResponseReader
rr) Bool
True
  ctx <- connectionContext brConn
  let go !Int
i [Query]
qs = case [Query]
qs of
        [] -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        [Query]
_ -> case StatementEncoder () Query
-> TdsContext
-> ()
-> Int
-> [Query]
-> Either [Char] ([(Text, RpcRequest)], Int, [Query], ())
forall s a.
StatementEncoder s a
-> TdsContext
-> s
-> Int
-> [a]
-> Either [Char] ([(Text, RpcRequest)], Int, [a], s)
takeStatements StatementEncoder () Query
queryStatement TdsContext
ctx () Int
i [Query]
qs of
          Left [Char]
_ -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          Right ([(Text, RpcRequest)]
stmts, Int
n, [Query]
rest, ()) -> do
            Request{..} <- [(Text, RpcRequest)] -> IO Request
statementsRequest [(Text, RpcRequest)]
stmts
            mask $ \forall a. IO a -> IO a
restore -> do
              ar <- Connection
-> Maybe Text
-> PacketType
-> (TdsContext -> PacketWriter -> IO ())
-> IO ActiveResponse
startRequestFor Connection
brConn Maybe Text
requestStatement PacketType
requestPacketType TdsContext -> PacketWriter -> IO ()
requestBody
              readingToEnd brConn ar (restore (drain ar))
            go (i + n) rest
  go (bsIndex st + length (bsInMessage st)) (bsRest st)
  throwIO failure
  where
    byClient :: CloseReason -> Bool
byClient = \case
      CloseReason
ClosedByClient -> Bool
True
      CloseReason
_ -> Bool
False

-- | A consumer threw (a row that doesn't decode, a shape the item didn't
-- expect, a fold's step): read the rest of the response so the
-- connection isn't left mid-message (as 'query' does before throwing a
-- 'DecodeError'), quietly.  The messages after it aren't sent.
abandonBatch :: BatchRun -> IO ()
abandonBatch :: BatchRun -> IO ()
abandonBatch BatchRun{IORef BatchState
Connection
BatchScope
brConn :: BatchRun -> Connection
brState :: BatchRun -> IORef BatchState
brScope :: BatchRun -> BatchScope
brConn :: Connection
brState :: IORef BatchState
brScope :: BatchScope
..} = do
  st <- IORef BatchState -> IO BatchState
forall a. IORef a -> IO a
readIORef IORef BatchState
brState
  forM_ (bsReader st) $ \ResponseReader
rr -> IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef (ResponseReader -> IORef Bool
respEnded ResponseReader
rr) IO Bool -> (Bool -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Bool
True -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    Bool
False -> do
      Connection -> ActiveResponse -> IO ()
drainActiveResponse Connection
brConn (ResponseReader -> ActiveResponse
respActive ResponseReader
rr) IO () -> (SomeException -> IO ()) -> IO ()
forall a. IO a -> (SomeException -> IO a) -> IO a
`catchSyncException` \SomeException
_ -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (ResponseReader -> IORef Bool
respEnded ResponseReader
rr) Bool
True

-- | A statement of a batch and its consumer, run over the reader bounded
-- to the statement's results.
batchItem :: Query -> (TdsContext -> ResponseReader -> IO a) -> BatchQuery a
batchItem :: forall a.
Query -> (TdsContext -> ResponseReader -> IO a) -> BatchQuery a
batchItem Query
q TdsContext -> ResponseReader -> IO a
consume = ([BatchItem] -> [BatchItem]) -> (BatchRun -> IO a) -> BatchQuery a
forall a.
([BatchItem] -> [BatchItem]) -> (BatchRun -> IO a) -> BatchQuery a
BatchQuery (Query -> Bool -> BatchItem
BatchItem Query
q Bool
False BatchItem -> [BatchItem] -> [BatchItem]
forall a. a -> [a] -> [a]
:) ((BatchRun -> IO a) -> BatchQuery a)
-> (BatchRun -> IO a) -> BatchQuery a
forall a b. (a -> b) -> a -> b
$ \BatchRun
br -> do
  rr <- BatchRun -> IO ResponseReader
nextStatement BatchRun
br
  ctx <- connectionContext (brConn br)
  consume ctx rr

-- | The statement's rows affected: what 'execute' returns for it.
batchExecute :: Query -> BatchQuery Int64
batchExecute :: Query -> BatchQuery Int64
batchExecute Query
q = Query
-> (TdsContext -> ResponseReader -> IO Int64) -> BatchQuery Int64
forall a.
Query -> (TdsContext -> ResponseReader -> IO a) -> BatchQuery a
batchItem Query
q (\TdsContext
_ ResponseReader
rr -> ResponseReader -> IO Int64
affectedRows ResponseReader
rr)

-- | 'batchExecute' without the count.
batchExecute_ :: Query -> BatchQuery ()
batchExecute_ :: Query -> BatchQuery ()
batchExecute_ = BatchQuery Int64 -> BatchQuery ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (BatchQuery Int64 -> BatchQuery ())
-> (Query -> BatchQuery Int64) -> Query -> BatchQuery ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Query -> BatchQuery Int64
batchExecute

-- | The rows of the statement's first result set: what 'query' returns
-- for it (@[]@ for a statement with no result set).
batchQuery :: FromSqlRow r => Query -> BatchQuery [r]
batchQuery :: forall r. FromSqlRow r => Query -> BatchQuery [r]
batchQuery = RowDecoder r -> Query -> BatchQuery [r]
forall r. RowDecoder r -> Query -> BatchQuery [r]
batchQueryWith RowDecoder r
forall a. FromSqlRow a => RowDecoder a
rowDecoder

-- | 'batchQuery' with the given row decoder, as 'queryWith'.
batchQueryWith :: RowDecoder r -> Query -> BatchQuery [r]
batchQueryWith :: forall r. RowDecoder r -> Query -> BatchQuery [r]
batchQueryWith RowDecoder r
dec Query
q = Query -> (TdsContext -> ResponseReader -> IO [r]) -> BatchQuery [r]
forall a.
Query -> (TdsContext -> ResponseReader -> IO a) -> BatchQuery a
batchItem Query
q (\TdsContext
_ ResponseReader
rr -> RowCompiler r -> ResponseReader -> IO [r]
forall r. RowCompiler r -> ResponseReader -> IO [r]
readFirstResultSet (RowDecoder r -> RowCompiler r
forall r. RowDecoder r -> RowCompiler r
compilerFor RowDecoder r
dec) ResponseReader
rr)

-- | 'batchQuery' with a 'RowDecoderMonadic', as 'queryMWith'.
batchQueryMWith :: RowDecoderMonadic r -> Query -> BatchQuery [r]
batchQueryMWith :: forall r. RowDecoderMonadic r -> Query -> BatchQuery [r]
batchQueryMWith RowDecoderMonadic r
dec Query
q = Query -> (TdsContext -> ResponseReader -> IO [r]) -> BatchQuery [r]
forall a.
Query -> (TdsContext -> ResponseReader -> IO a) -> BatchQuery a
batchItem Query
q (\TdsContext
_ ResponseReader
rr -> RowCompiler r -> ResponseReader -> IO [r]
forall r. RowCompiler r -> ResponseReader -> IO [r]
readFirstResultSet (RowDecoderMonadic r -> RowCompiler r
forall r. RowDecoderMonadic r -> RowCompiler r
compilerForM RowDecoderMonadic r
dec) ResponseReader
rr)

-- | Exactly one row, as 'query1'.
batchQuery1 :: FromSqlRow r => Query -> BatchQuery r
batchQuery1 :: forall r. FromSqlRow r => Query -> BatchQuery r
batchQuery1 = RowDecoder r -> Query -> BatchQuery r
forall r. RowDecoder r -> Query -> BatchQuery r
batchQuery1With RowDecoder r
forall a. FromSqlRow a => RowDecoder a
rowDecoder

-- | 'batchQuery1' with the given row decoder, as 'query1With'.
batchQuery1With :: RowDecoder r -> Query -> BatchQuery r
batchQuery1With :: forall r. RowDecoder r -> Query -> BatchQuery r
batchQuery1With RowDecoder r
dec Query
q = Query -> (TdsContext -> ResponseReader -> IO r) -> BatchQuery r
forall a.
Query -> (TdsContext -> ResponseReader -> IO a) -> BatchQuery a
batchItem Query
q (\TdsContext
_ ResponseReader
rr -> RowCompiler r -> ResponseReader -> IO [r]
forall r. RowCompiler r -> ResponseReader -> IO [r]
readFirstResultSet (RowDecoder r -> RowCompiler r
forall r. RowDecoder r -> RowCompiler r
compilerFor RowDecoder r
dec) ResponseReader
rr IO [r] -> ([r] -> IO r) -> IO r
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= [r] -> IO r
forall r. [r] -> IO r
exactlyOne)

-- | Zero or one row, as 'queryMay'.
batchQueryMay :: FromSqlRow r => Query -> BatchQuery (Maybe r)
batchQueryMay :: forall r. FromSqlRow r => Query -> BatchQuery (Maybe r)
batchQueryMay = RowDecoder r -> Query -> BatchQuery (Maybe r)
forall r. RowDecoder r -> Query -> BatchQuery (Maybe r)
batchQueryMayWith RowDecoder r
forall a. FromSqlRow a => RowDecoder a
rowDecoder

-- | 'batchQueryMay' with the given row decoder, as 'queryMayWith'.
batchQueryMayWith :: RowDecoder r -> Query -> BatchQuery (Maybe r)
batchQueryMayWith :: forall r. RowDecoder r -> Query -> BatchQuery (Maybe r)
batchQueryMayWith RowDecoder r
dec Query
q = Query
-> (TdsContext -> ResponseReader -> IO (Maybe r))
-> BatchQuery (Maybe r)
forall a.
Query -> (TdsContext -> ResponseReader -> IO a) -> BatchQuery a
batchItem Query
q (\TdsContext
_ ResponseReader
rr -> RowCompiler r -> ResponseReader -> IO [r]
forall r. RowCompiler r -> ResponseReader -> IO [r]
readFirstResultSet (RowDecoder r -> RowCompiler r
forall r. RowDecoder r -> RowCompiler r
compilerFor RowDecoder r
dec) ResponseReader
rr IO [r] -> ([r] -> IO (Maybe r)) -> IO (Maybe r)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= [r] -> IO (Maybe r)
forall r. [r] -> IO (Maybe r)
atMostOne)

-- | Fold over the rows of the statement's first result set as they
-- arrive, as 'queryFold': the step runs with the connection lock held,
-- while the batch is being read, so it can't use the connection.
batchQueryFold :: FromSqlRow r => Query -> a -> (a -> r -> IO a) -> BatchQuery a
batchQueryFold :: forall r a.
FromSqlRow r =>
Query -> a -> (a -> r -> IO a) -> BatchQuery a
batchQueryFold = RowDecoder r -> Query -> a -> (a -> r -> IO a) -> BatchQuery a
forall r a.
RowDecoder r -> Query -> a -> (a -> r -> IO a) -> BatchQuery a
batchQueryFoldWith RowDecoder r
forall a. FromSqlRow a => RowDecoder a
rowDecoder

-- | 'batchQueryFold' with the given row decoder, as 'queryFoldWith'.
batchQueryFoldWith :: RowDecoder r -> Query -> a -> (a -> r -> IO a) -> BatchQuery a
batchQueryFoldWith :: forall r a.
RowDecoder r -> Query -> a -> (a -> r -> IO a) -> BatchQuery a
batchQueryFoldWith RowDecoder r
dec Query
q a
z a -> r -> IO a
step = Query -> (TdsContext -> ResponseReader -> IO a) -> BatchQuery a
forall a.
Query -> (TdsContext -> ResponseReader -> IO a) -> BatchQuery a
batchItem Query
q (\TdsContext
_ ResponseReader
rr -> RowCompiler r -> ResponseReader -> a -> (a -> r -> IO a) -> IO a
forall r a.
RowCompiler r -> ResponseReader -> a -> (a -> r -> IO a) -> IO a
foldFirstResultSet (RowDecoder r -> RowCompiler r
forall r. RowDecoder r -> RowCompiler r
compilerFor RowDecoder r
dec) ResponseReader
rr a
z a -> r -> IO a
step)

-- | Stream the rows of the statement's first result set, as 'queryS':
-- they are read from the socket as the stream is consumed, after
-- 'runBatchQuery' has returned.  So the statement must be the last of
-- its batch ('runBatchQuery' throws 'EncodeError' otherwise, before
-- sending anything), and its stream must be consumed before the next
-- request on the connection, as any 'queryS' stream must (on a MARS
-- connection it stays open beside later requests instead, as a 'queryS'
-- stream does).  An error the statement reports is thrown by the stream;
-- one an earlier statement reports, by 'runBatchQuery'.
batchQueryS :: FromSqlRow r => Query -> BatchQuery (Stream (Of r) IO ())
batchQueryS :: forall r. FromSqlRow r => Query -> BatchQuery (Stream (Of r) IO ())
batchQueryS = RowDecoder r -> Query -> BatchQuery (Stream (Of r) IO ())
forall r. RowDecoder r -> Query -> BatchQuery (Stream (Of r) IO ())
batchQuerySWith RowDecoder r
forall a. FromSqlRow a => RowDecoder a
rowDecoder

-- | 'batchQueryS' with the given row decoder, as 'querySWith'.
batchQuerySWith :: RowDecoder r -> Query -> BatchQuery (Stream (Of r) IO ())
batchQuerySWith :: forall r. RowDecoder r -> Query -> BatchQuery (Stream (Of r) IO ())
batchQuerySWith RowDecoder r
dec = RowCompiler r -> Query -> BatchQuery (Stream (Of r) IO ())
forall r.
RowCompiler r -> Query -> BatchQuery (Stream (Of r) IO ())
batchStream (RowDecoder r -> RowCompiler r
forall r. RowDecoder r -> RowCompiler r
compilerFor RowDecoder r
dec)

-- | 'batchQueryS' with a 'RowDecoderMonadic', as 'querySMWith'.
batchQuerySMWith :: RowDecoderMonadic r -> Query -> BatchQuery (Stream (Of r) IO ())
batchQuerySMWith :: forall r.
RowDecoderMonadic r -> Query -> BatchQuery (Stream (Of r) IO ())
batchQuerySMWith RowDecoderMonadic r
dec = RowCompiler r -> Query -> BatchQuery (Stream (Of r) IO ())
forall r.
RowCompiler r -> Query -> BatchQuery (Stream (Of r) IO ())
batchStream (RowDecoderMonadic r -> RowCompiler r
forall r. RowDecoderMonadic r -> RowCompiler r
compilerForM RowDecoderMonadic r
dec)

batchStream :: RowCompiler r -> Query -> BatchQuery (Stream (Of r) IO ())
batchStream :: forall r.
RowCompiler r -> Query -> BatchQuery (Stream (Of r) IO ())
batchStream RowCompiler r
compile Query
q = ([BatchItem] -> [BatchItem])
-> (BatchRun -> IO (Stream (Of r) IO ()))
-> BatchQuery (Stream (Of r) IO ())
forall a.
([BatchItem] -> [BatchItem]) -> (BatchRun -> IO a) -> BatchQuery a
BatchQuery (Query -> Bool -> BatchItem
BatchItem Query
q Bool
True BatchItem -> [BatchItem] -> [BatchItem]
forall a. a -> [a] -> [a]
:) ((BatchRun -> IO (Stream (Of r) IO ()))
 -> BatchQuery (Stream (Of r) IO ()))
-> (BatchRun -> IO (Stream (Of r) IO ()))
-> BatchQuery (Stream (Of r) IO ())
forall a b. (a -> b) -> a -> b
$ \BatchRun
br -> do
  rr <- BatchRun -> IO ResponseReader
nextStatement BatchRun
br
  pure (streamFirstResultSet compile rr (finishBatch br))

-- | Every result set of the statement, as 'queryMulti'.
batchQueryMulti :: Query -> BatchQuery [ResultSet SqlRow]
batchQueryMulti :: Query -> BatchQuery [ResultSet SqlRow]
batchQueryMulti Query
q = Query
-> (TdsContext -> ResponseReader -> IO [ResultSet SqlRow])
-> BatchQuery [ResultSet SqlRow]
forall a.
Query -> (TdsContext -> ResponseReader -> IO a) -> BatchQuery a
batchItem Query
q (\TdsContext
_ ResponseReader
rr -> RowCompiler SqlRow -> ResponseReader -> IO [ResultSet SqlRow]
forall r. RowCompiler r -> ResponseReader -> IO [ResultSet r]
readResultSets RowCompiler SqlRow
sqlRowCompiler ResponseReader
rr)

-- | 'batchQueryMulti' with the given row decoder, as 'queryMultiWith'.
batchQueryMultiWith :: RowDecoder r -> Query -> BatchQuery [ResultSet r]
batchQueryMultiWith :: forall r. RowDecoder r -> Query -> BatchQuery [ResultSet r]
batchQueryMultiWith RowDecoder r
dec Query
q = Query
-> (TdsContext -> ResponseReader -> IO [ResultSet r])
-> BatchQuery [ResultSet r]
forall a.
Query -> (TdsContext -> ResponseReader -> IO a) -> BatchQuery a
batchItem Query
q (\TdsContext
_ ResponseReader
rr -> RowCompiler r -> ResponseReader -> IO [ResultSet r]
forall r. RowCompiler r -> ResponseReader -> IO [ResultSet r]
readResultSets (RowDecoder r -> RowCompiler r
forall r. RowDecoder r -> RowCompiler r
compilerFor RowDecoder r
dec) ResponseReader
rr)

-- | 'batchQueryMulti' with a 'RowDecoderMonadic', as 'queryMultiMWith'.
batchQueryMultiMWith :: RowDecoderMonadic r -> Query -> BatchQuery [ResultSet r]
batchQueryMultiMWith :: forall r. RowDecoderMonadic r -> Query -> BatchQuery [ResultSet r]
batchQueryMultiMWith RowDecoderMonadic r
dec Query
q = Query
-> (TdsContext -> ResponseReader -> IO [ResultSet r])
-> BatchQuery [ResultSet r]
forall a.
Query -> (TdsContext -> ResponseReader -> IO a) -> BatchQuery a
batchItem Query
q (\TdsContext
_ ResponseReader
rr -> RowCompiler r -> ResponseReader -> IO [ResultSet r]
forall r. RowCompiler r -> ResponseReader -> IO [ResultSet r]
readResultSets (RowDecoderMonadic r -> RowCompiler r
forall r. RowDecoderMonadic r -> RowCompiler r
compilerForM RowDecoderMonadic r
dec) ResponseReader
rr)

-- | The statement's result sets decoded one at a time by a 'ResultSets'
-- action, as 'queryResultSets': the result sets are the statement's own,
-- and asking for one past its last throws 'UnexpectedResult'.
batchQueryResultSets :: ResultSets a -> Query -> BatchQuery a
batchQueryResultSets :: forall a. ResultSets a -> Query -> BatchQuery a
batchQueryResultSets ResultSets a
sets Query
q = Query -> (TdsContext -> ResponseReader -> IO a) -> BatchQuery a
forall a.
Query -> (TdsContext -> ResponseReader -> IO a) -> BatchQuery a
batchItem Query
q (\TdsContext
_ ResponseReader
rr -> ResultSets a -> ResponseReader -> IO a
forall a. ResultSets a -> ResponseReader -> IO a
readResultSetsWith ResultSets a
sets ResponseReader
rr)