8000 More useful DataRows · postgres-haskell/postgres-wire@7681fe2 · GitHub
[go: up one dir, main page]
Skip to content
8000

Commit 7681fe2

Browse files
VyacheslavHashovVyacheslavHashov
committed
More useful DataRows
1 parent 5bd2b34 commit 7681fe2

File tree

8 files changed

+88
-44
lines changed

8 files changed

+88
-44
lines changed

‎cbits/include/pw_utils.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,4 +9,4 @@
99
#define NEED_MORE_INPUT 0x01
1010
#define OTHER_HEADER 0x02
1111

12-
size_t scan_datarows(char *buffer, size_t len, int *reason);
12+
size_t scan_datarows(char *buffer, size_t len, unsigned long *p_count, int *reason);

‎cbits/src/pw_utils.c

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,16 @@
11
#include <pw_utils.h>
22

3-
/* pointer to the next data message
4-
length of buffer
5-
ptr where result reasond would be put
6-
returns offset of buffer when datarows end.
7-
*/
8-
size_t scan_datarows(char *buffer, size_t len, int *reason)
3+
/* returns offset of buffer when datarows end. */
4+
size_t scan_datarows(
5+
char *buffer, /* pointer to the next data message */
6+
size_t len, /* length of the buffer */
7+
unsigned long *p_count, /* ptr where count of datarows will be put */
8+
int *reason /* ptr where result reasond will be put */
9+
)
910
{
1011
size_t offset = 0;
1112
uint32_t message_len = 0;
13+
unsigned long count = 0;
1214

1315
while (1)
1416
{
@@ -27,6 +29,8 @@ size_t scan_datarows(char *buffer, size_t len, int *reason)
2729
break;
2830
}
2931
offset = offset + HEADER_TYPE_SIZE + (size_t)message_len;
32+
count = count + 1;
3033
}
34+
*p_count = count;
3135
return offset;
3236
}

‎postgres-wire.cabal

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ library
2929
, Database.PostgreSQL.Protocol.Encoders
3030
, Database.PostgreSQL.Protocol.Decoders
3131
, Database.PostgreSQL.Protocol.Parsers
32-
, Database.PostgreSQL.Protocol.ExtractDataRows
32+
, Database.PostgreSQL.Protocol.DataRows
3333
, Database.PostgreSQL.Protocol.Store.Encode
3434
, Database.PostgreSQL.Protocol.Store.Decode
3535
, Database.PostgreSQL.Protocol.Codecs.Decoders

‎src/Database/PostgreSQL/Driver/Connection.hs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ import qualified Data.ByteString.Char8 as BS(pack, unpack)
4040
import Database.PostgreSQL.Protocol.Encoders
4141
import Database.PostgreSQL.Protocol.Decoders
4242
import Database.PostgreSQL.Protocol.Parsers
43-
import Database.PostgreSQL.Protocol.ExtractDataRows
43+
import Database.PostgreSQL.Protocol.DataRows
4444
import Database.PostgreSQL.Protocol.Types
4545
import Database.PostgreSQL.Protocol.Store.Encode (runEncode, Encode)
4646
import Database.PostgreSQL.Protocol.Store.Decode (runDecode)

‎src/Database/PostgreSQL/Protocol/ExtractDataRows.hs renamed to ‎src/Database/PostgreSQL/Protocol/DataRows.hs

Lines changed: 46 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,14 @@
1-
module Database.PostgreSQL.Protocol.ExtractDataRows
1+
module Database.PostgreSQL.Protocol.DataRows
22
( loopExtractDataRows
3+
, countDataRows
4+
, flattenDataRows
35
) where
46

7+
import Data.Monoid ((<>))
58
import Data.Word (Word8, byteSwap32)
69
import Foreign (peek, peekByteOff, castPtr)
710
import qualified Data.ByteString as B
811
import qualified Data.ByteString.Unsafe as B
9-
import qualified Data.ByteString.Lazy as BL
10-
import qualified Data.ByteString.Lazy.Internal as BL
1112

1213
import Database.PostgreSQL.Driver.Error
1314
import Database.PostgreSQL.Protocol.Types
@@ -23,16 +24,16 @@ loopExtractDataRows
2324
-- Will be called on every DataMessage.
2425
-> (DataMessage -> IO ())
2526
-> IO ()
26-
loopExtractDataRows readMoreAction callback = go "" ""
27+
loopExtractDataRows readMoreAction callback = go "" Empty
2728
where
28-
go :: B.ByteString -> BL.ByteString -> IO ()
29+
-- Note that DataRows go in reverse order.
30+
go :: B.ByteString -> DataRows -> IO ()
2931
go bs acc
3032
| B.length bs < headerSize = readMoreAndGo bs acc
3133
| otherwise = do
3234
ScanRowResult ch rest r <- scanDataRows bs
3335
-- We should force accumulator
34-
-- Note: `BL.chunk` should not prepend empty bytestring as chunk.
35-
let !newAcc = BL.chunk ch acc
36+
let !newAcc = chunk ch acc
3637

3738
case r of
3839
-- Following happened:
@@ -47,27 +48,25 @@ loopExtractDataRows readMoreAction callback = go "" ""
4748
dispatchHeader mt len (B.drop headerSize rest) newAcc
4849

4950
{-# INLINE dispatchHeader #-}
50-
dispatchHeader :: Word8 -> Int -> B.ByteString -> BL.ByteString -> IO ()
51+
dispatchHeader :: Word8 -> Int -> B.ByteString -> DataRows -> IO ()
5152
dispatchHeader mt len bs acc = case mt of
5253
-- 'C' - CommandComplete.
5354
-- Command is completed, return the result.
5455
67 -> do
55-
callback $
56-
DataMessage . DataRows $
57-
BL.foldlChunks (flip BL.chunk) BL.empty acc
56+
callback $
57+
DataMessage $ reverseDataRows acc
5858

5959
newBs <- skipBytes bs len
60-
go newBs BL.empty
60+
go newBs Empty
6161

6262
-- 'I' - EmptyQueryResponse.
6363
-- PostgreSQL sends this if query string was empty and datarows
6464
-- should be empty, but anyway we return data collected in `acc`.
6565
73 -> do
6666
callback $
67-
DataMessage . DataRows $
68-
BL.foldlChunks (flip BL.chunk) BL.empty acc
67+
DataMessage $ reverseDataRows acc
6968

70-
go bs BL.empty
69+
go bs Empty
7170

7271
-- 'E' - ErrorResponse.
7372
-- On ErrorResponse we should discard all the collected datarows.
@@ -76,7 +75,7 @@ loopExtractDataRows readMoreAction callback = go "" ""
7675
desc <- eitherToProtocolEx $ parseErrorDesc b
7776
callback (DataError desc)
7877

79-
go newBs BL.empty
78+
go newBs Empty
8079

8180
-- 'Z' - ReadyForQuery.
8281
-- To know when command processing is finished
@@ -92,7 +91,7 @@ loopExtractDataRows readMoreAction callback = go "" ""
9291
go newBs acc
9392

9493
{-# INLINE readMoreAndGo #-}
95-
readMoreAndGo :: B.ByteString -> BL.ByteString -> IO ()
94+
readMoreAndGo :: B.ByteString -> DataRows -> IO ()
9695
readMoreAndGo bs acc = do
9796
newBs <- readMoreAction bs
9897
go newBs acc
@@ -123,3 +122,33 @@ loopExtractDataRows readMoreAction callback = go "" ""
123122
b <- peek (castPtr ptr)
124123
w <- byteSwap32 <$> peekByteOff (castPtr ptr) 1
125124
pure $ Header b $ fromIntegral (w - 4)
125+
126+
---
127+
-- Utils
128+
--
129+
130+
{-# INLINE chunk #-}
131+
chunk :: DataChunk -> DataRows -> DataRows
132+
chunk ch@(DataChunk len bs) dr
133+
| len == 0 = dr
134+
| otherwise = DataRows ch dr
135+
136+
{-# INLINE foldlDataRows #-}
137+
foldlDataRows :: (a -> DataChunk -> a) -> a -> DataRows -> a
138+
foldlDataRows f z = go z
139+
where
140+
go a Empty = a
141+
go a (DataRows ch dr) = let !na = f a ch in go na dr
142+
143+
{-# INLINE reverseDataRows #-}
144+
reverseDataRows :: DataRows -> DataRows
145+
reverseDataRows = foldlDataRows (flip chunk) Empty
146+
147+
{-# INLINE countDataRows #-}
148+
countDataRows :: DataRows -> Word
149+
countDataRows = foldlDataRows (\acc (DataChunk c _) -> acc + c) 0
150+
151+
-- | For testing only
152+
{-# INLINE flattenDataRows #-}
153+
flattenDataRows :: DataRows -> B.ByteString
154+
flattenDataRows = foldlDataRows (\acc (DataChunk _ bs) -> acc <> bs) ""

‎src/Database/PostgreSQL/Protocol/Types.hs

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -44,10 +44,15 @@ newtype ServerSecretKey = ServerSecretKey Int32 deriving (Show)
4444
data ServerVersion = ServerVersion !Word8 !Word8 !Word8 !ByteString
4545
deriving (Eq, Show)
4646

47+
-- A chunk of DataRows.
48+
-- It is guaranted that a `ByteString` contains integer number of `DataRow`s.
49+
data DataChunk = DataChunk
50+
{-# UNPACK #-} !Word -- ^ Count of DataRows in ByteString
51+
{-# UNPACK #-} !B.ByteString
52+
deriving (Show, Eq)
53+
4754
-- | Helper types that contains only raw DataRows messages.
48-
-- It is guaranted that a single strict chunk of the `ByteString`
49-
-- contains integer number of `DataRow`s.
50-
newtype DataRows = DataRows BL.ByteString
55+
data DataRows = Empty | DataRows {-# UNPACK #-} DataChunk DataRows
5156
deriving (Show, Eq)
5257

5358
-- | Ad-hoc type only for data rows.

‎src/Database/PostgreSQL/Protocol/Utils.hs

Lines changed: 14 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,28 +1,32 @@
11
{-# language ForeignFunctionInterface #-}
22
module Database.PostgreSQL.Protocol.Utils where
33

4-
import Foreign.C.Types (CInt, CSize(..), CChar)
4+
import Foreign.C.Types (CInt, CSize(..), CChar, CULong)
55
import Foreign (Ptr, peek, alloca)
66
import qualified Data.ByteString as B
77
import qualified Data.ByteString.Unsafe as B
88

9+
import Database.PostgreSQL.Protocol.Types (DataChunk(..))
10+
911
data ScanRowResult = ScanRowResult
10-
{-# UNPACK #-} !B.ByteString -- chunk of datarows, may be empty
12+
{-# UNPACK #-} !DataChunk -- chunk of datarows, may be empty
1113
{-# UNPACK #-} !B.ByteString -- the rest of string
1214
{-# UNPACK #-} !Int -- reason code
1315

1416
-- | Scans `ByteString` for a chunk of `DataRow`s.
1517
{-# INLINE scanDataRows #-}
1618
scanDataRows :: B.ByteString -> IO ScanRowResult
1719
scanDataRows bs =
18-
alloca $ \reasonPtr ->
19-
B.unsafeUseAsCStringLen bs $ \(ptr, len) -> do
20-
offset <- fromIntegral <$>
21-
c_scan_datarows ptr (fromIntegral len) reasonPtr
22-
reason <- peek reasonPtr
23-
let (ch, rest) = B.splitAt offset bs
24-
pure $ ScanRowResult ch rest $ fromIntegral reason
20+
alloca $ \countPtr ->
21+
alloca $ \reasonPtr ->
22+
B.unsafeUseAsCStringLen bs $ \(ptr, len) -> do
23+
offset <- fromIntegral <$>
24+
c_scan_datarows ptr (fromIntegral len) countPtr reasonPtr
25+
reason <- fromIntegral <$> peek reasonPtr
26+
count <- fromIntegral <$> peek countPtr
27+
let (ch, rest) = B.splitAt offset bs
28+
pure $ ScanRowResult (DataChunk count ch) rest reason
2529

2630
foreign import ccall unsafe "static pw_utils.h scan_datarows" c_scan_datarows
27-
:: Ptr CChar -> CSize -> Ptr CInt -> IO CSize
31+
:: Ptr CChar -> CSize -> Ptr CULong -> Ptr CInt -> IO CSize
2832

‎tests/Driver.hs

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import Database.PostgreSQL.Driver.Connection
1717
import Database.PostgreSQL.Driver.StatementStorage
1818
import Database.PostgreSQL.Driver.Query
1919
import Database.PostgreSQL.Protocol.Types
20+
import Database.PostgreSQL.Protocol.DataRows
2021
import Database.PostgreSQL.Protocol.Store.Decode
2122
import Database.PostgreSQL.Protocol.Decoders
2223

@@ -54,8 +55,8 @@ fromRight _ = error "fromRight"
5455
fromMessage :: Either e DataRows -> B.ByteString
5556
-- TODO
5657
-- 5 bytes -header, 2 bytes -count, 4 bytes - length
57-
fromMessage (Right (DataRows bs)) = B.drop 11 $ BL.toStrict bs
58-
fromMessage _ = error "from message"
58+
fromMessage (Right rows) = B.drop 11 $ flattenDataRows rows
59+
fromMessage _ = error "from message"
5960

6061
-- | Single batch.
6162
testBatch :: IO ()
@@ -116,7 +117,7 @@ assertQueryNoData q = withConnection $ \c -> do
116117
sendBatchAndSync c [q]
117118
r <- fromRight <$> readNextData c
118119
waitReadyForQuery c
119-
DataRows "" @=? r
120+
Empty @=? r
120121

121122
-- | Asserts that all the received data messages are in form (Right _)
122123
checkRightResult :: Connection -> Int -> Assertion
@@ -226,9 +227,10 @@ testCorrectDatarows = withConnection $ \c -> do
226227
r <- readNextData c
227228
case r of
228229
Left e -> error $ show e
229-
Right (DataRows rows) -> do
230-
let bs = BL.toStrict rows
230+
Right rows -> do
231+
let bs = flattenDataRows rows
231232
map (BS.pack . show ) [1 .. 1000] @=? go bs
233+
countDataRows rows @=? 1000
232234
where
233235
go bs | B.null bs = []
234236
| otherwise = let (rest, v) = runDecode decodeDataRow bs

0 commit comments

Comments
 (0)
0