Make stream more general

Author: Johannes

app/Main.hs

@@ -39,7 +39,7 @@ repl c = do
       hSetBuffering stdout LineBuffering
       catch
         (catch
-           (do (_, count) <- ODBC.stream c input output (False, 0 :: Int)
+           (do (_, count) <- ODBC.stream c input output (\_ -> pure (False, 0 :: Int))
                putStrLn ("Rows: " ++ show count))
            (\case
               UserInterrupt -> pure ()
@@ -54,7 +54,7 @@ piped c = do
   hSetBuffering stdout LineBuffering
   catch
     (catch
-       (do (_, count) <- ODBC.stream c input output (False, 0 :: Int)
+       (do (_, count) <- ODBC.stream c input output (\_ -> pure (False, 0 :: Int))
            putStrLn ("Rows: " ++ show count))
        (\case
           UserInterrupt -> pure ()

src/Database/ODBC/Internal.hs

@@ -334,7 +334,7 @@ stream ::
   -> (state -> [(Column, Value)] -> m (Step state))
   -- ^ A stepping function that gets as input the current @state@ and
   -- a row, returning either a new @state@ or a final @result@.
-  -> state
+  -> ([Column] -> m state)
   -- ^ A state that you can use for the computation. Strictly
   -- evaluated each iteration.
   -> m state
@@ -352,7 +352,7 @@ streamWithParams ::
   -> (state -> [(Column, Value)] -> m (Step state))
   -- ^ A stepping function that gets as input the current @state@ and
   -- a row, returning either a new @state@ or a final @result@.
-  -> state
+  -> ([Column] -> m state)
   -- ^ A state that you can use for the computation. Strictly
   -- evaluated each iteration.
   -> m state
@@ -496,7 +496,7 @@ fetchIterator ::
      Ptr EnvAndDbc
   -> UnliftIO m
   -> (state -> [(Column, Value)] -> m (Step state))
-  -> state
+  -> ([Column] -> m state)
   -> SQLHSTMT s
   -> IO state
 fetchIterator dbc (UnliftIO runInIO) step state0 stmt = do
@@ -533,9 +533,8 @@ fetchIterator dbc (UnliftIO runInIO) step state0 stmt = do
                      (coerce retcode0)
                      "Unexpected return code"
                      sqlState)
-  if cols > 0
-    then loop state0
-    else pure state0
+
+  (if cols > 0 then loop else pure) =<< runInIO (state0 types)
 
 -- | Fetch all results from possible multiple statements.
 fetchAllResults :: Ptr EnvAndDbc -> SQLHSTMT s -> IO ()

src/Database/ODBC/SQLServer.hs

@@ -455,7 +455,7 @@ stream ::
   -> (state -> row -> m (Internal.Step state))
   -- ^ A stepping function that gets as input the current @state@ and
   -- a row, returning either a new @state@ or a final @result@.
-  -> state
+  -> ([Internal.Column] -> m state)
   -- ^ A state that you can use for the computation. Strictly
   -- evaluated each iteration.
   -> m state

test/Main.hs

@@ -231,7 +231,7 @@ dataRetrieval = do
             c
             "DROP TABLE IF EXISTS no_such_table"
             (\s _ -> pure (Stop s))
-            []
+            (\_ -> pure [])
         shouldBe (map (map snd) (rows1 ++ rows2)) [])
   quickCheckInternalRoundtrip
     "Int"