8000 Raw connection · postgres-haskell/postgres-wire@bd4abcb · GitHub
[go: up one dir, main page]
Skip to content

Commit bd4abcb

Browse files
VyacheslavHashovVyacheslavHashov
committed
Raw connection
1 parent b48fbca commit bd4abcb

File tree

3 files changed

+35
-20
lines changed

3 files changed

+35
-20
lines changed

‎postgres-wire.cabal

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,6 @@ library
2020
, Database.PostgreSQL.Settings
2121
, Database.PostgreSQL.StatementStorage
2222
, Database.PostgreSQL.Types
23-
, Database.PostgreSQL.Session
2423

2524
, Database.PostgreSQL.Protocol.Types
2625
, Database.PostgreSQL.Protocol.Encoders

‎src/Database/PostgreSQL/Connection.hs

Lines changed: 32 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,8 @@ import Data.Binary.Get ( runGetIncremental, pushChunk)
2121
import qualified Data.Binary.Get as BG (Decoder(..))
2222
import Data.Maybe (fromJust)
2323
import qualified Data.Vector as V
24-
import System.Socket hiding (connect, close, Error(..))
25-
import qualified System.Socket as Socket (connect, close)
24+
import System.Socket (Socket, socket)
25+
import qualified System.Socket as Socket (connect, close, send, receive)
2626
import System.Socket.Family.Inet6
2727
import System.Socket.Type.Stream
2828
import System.Socket.Protocol.TCP
@@ -41,7 +41,7 @@ import Database.PostgreSQL.Types
4141
type UnixSocket = Socket Unix Stream Unix
4242
-- data Connection = Connection (Socket Inet6 Stream TCP)
4343
data Connection = Connection
44-
{ connSocket :: UnixSocket
44+
{ connRawConnection :: RawConnection
4545
, connReceiverThread :: ThreadId
4646
-- channel only for Data messages
4747
, connOutDataChan :: OutChan (Either Error DataMessage)
@@ -64,6 +64,13 @@ data Error
6464
data DataMessage = DataMessage [V.Vector B.ByteString]
6565
deriving (Show)
6666

67+
-- | Abstraction over raw socket connection or tls connection
68+
data RawConnection = RawConnection
69+
{ rFlush :: IO ()
70+
, rClose :: IO ()
71+
, rSend :: B.ByteString -> IO ()
72+
, rReceive :: Int -> IO B.ByteString
73+
}
6774

6875
address :: SocketAddress Unix
6976
address = fromJust $ socketAddressUnixPath "/var/run/postgresql/.s.PGSQL.5432"
@@ -73,15 +80,21 @@ connect settings = do
7380
s <- socket
7481
Socket.connect s address
7582
sendStartMessage s $ consStartupMessage settings
76-
r <- receive s 4096 mempty
83+
r <- Socket.receive s 4096 mempty
7784
readAuthMessage r
7885

7986
(inDataChan, outDataChan) <- newChan
8087
(inAllChan, outAllChan) <- newChan
81-
tid <- forkIO $ receiverThread s inDataChan inAllChan
88+
let rawConnection = RawConnection
89+
{ rFlush = pure ()
90+
, rClose = Socket.close s
91+
, rSend = \msg -> void $ Socket.send s msg mempty
92+
, rReceive = \n -> Socket.receive s n mempty
93+
}
94+
tid <- forkIO $ receiverThread rawConnection inDataChan inAllChan
8295
storage <- newStatementStorage
8396
pure Connection
84-
{ connSocket = s
97+
{ connRawConnection = rawConnection
8598
, connReceiverThread = tid
8699
, connOutDataChan = outDataChan
87100
, connOutAllChan = outAllChan
@@ -96,7 +109,7 @@ connect settings = do
96109
close :: Connection -> IO ()
97110
close conn = do
98111
killThread $ connReceiverThread conn
99-
Socket.close $ connSocket conn
112+
rClose $ connRawConnection conn
100113

101114
consStartupMessage :: ConnectionSettings -> StartMessage
102115
consStartupMessage stg = StartupMessage
@@ -105,12 +118,12 @@ consStartupMessage stg = StartupMessage
105118
sendStartMessage :: UnixSocket -> StartMessage -> IO ()
106119
sendStartMessage sock msg = void $ do
107120
let smsg = toStrict . toLazyByteString $ encodeStartMessage msg
108-
send sock smsg mempty
121+
Socket.send sock smsg mempty
109122

110-
sendMessage :: UnixSocket -> ClientMessage -> IO ()
111-
sendMessage sock msg = void $ do
123+
sendMessage :: RawConnection -> ClientMessage -> IO ()
124+
sendMessage rawConn msg = void $ do
112125
let smsg = toStrict . toLazyByteString $ encodeClientMessage msg
113-
send sock smsg mempty
126+
rSend rawConn smsg
114127

115128
readAuthMessage :: B.ByteString -> IO ()
116129
readAuthMessage s =
@@ -121,15 +134,15 @@ readAuthMessage s =
121134
f -> error $ show s
122135

123136
receiverThread
124-
:: UnixSocket
137+
:: RawConnection
125138
-> InChan (Either Error DataMessage)
126139
-> InChan ServerMessage
127140
-> IO ()
128-
receiverThread sock dataChan allChan = receiveLoop []
141+
receiverThread rawConn dataChan allChan = receiveLoop []
129142
where
130143
receiveLoop :: [V.Vector B.ByteString] -> IO()
131144
receiveLoop acc = do
132-
r <- receive sock 4096 mempty
145+
r <- rReceive rawConn 4096
133146
-- print r
134147
go r acc >>= receiveLoop
135148

@@ -220,7 +233,7 @@ data Query = Query
220233
sendBatch :: Connection -> [Query] -> IO ()
221234
sendBatch conn = traverse_ sendSingle
222235
where
223-
s = connSocket conn
236+
s = connRawConnection conn
224237
sname = StatementName ""
225238
pname = PortalName ""
226239
sendSingle q = do
@@ -230,10 +243,10 @@ sendBatch conn = traverse_ sendSingle
230243
sendMessage s $ Execute pname noLimitToReceive
231244

232245
sendSync :: Connection -> IO ()
233-
sendSync conn = sendMessage (connSocket conn) Sync
246+
sendSync conn = sendMessage (connRawConnection conn) Sync
234247

235248
sendFlush :: Connection -> IO ()
236-
sendFlush conn = sendMessage (connSocket conn) Flush
249+
sendFlush conn = sendMessage (connRawConnection conn) Flush
237250

238251
readNextData :: Connection -> IO (Either Error DataMessage)
239252
readNextData conn = readChan $ connOutDataChan conn
@@ -269,7 +282,7 @@ describeStatement conn stmt = do
269282
sendMessage s Sync
270283
parseMessages <$> waitReadyForQueryCollect conn
271284
where
272-
s = connSocket conn
285+
s = connRawConnection conn
273286
sname = StatementName ""
274287
parseMessages msgs = case msgs of
275288
[ParameterDescription params, NoData]
@@ -289,6 +302,7 @@ test :: IO ()
289302
test = do
290303
c <- connect defaultConnectionSettings
291304
sendBatch c queries
305+
sendSync c
292306
readResults c $ length queries
293307
readReadyForQuery c >>= print
294308
close c

‎src/Database/PostgreSQL/Settings.hs

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,10 @@ module Database.PostgreSQL.Settings where
55
import Data.Word (Word16)
66
import Data.ByteString (ByteString)
77

8+
-- | Connection settings to PostgreSQL
89
data ConnectionSettings = ConnectionSettings
9-
{ settingsHost :: ByteString
10+
{ -- On empty string a Unix socket will be used.
11+
settingsHost :: ByteString
1012
, settingsPort :: Word16
1113
, settingsDatabase :: ByteString
1214
, settingsUser :: ByteString

0 commit comments

Comments
 (0)
0