@@ -5,7 +5,7 @@ use crate::io::{PortalId, StatementId};
5
5
use crate :: logger:: QueryLogger ;
6
6
use crate :: message:: {
7
7
self , BackendMessageFormat , Bind , Close , CommandComplete , DataRow , ParameterDescription , Parse ,
8
- ParseComplete , Query , RowDescription ,
8
+ ParseComplete , RowDescription ,
9
9
} ;
10
10
use crate :: statement:: PgStatementMetadata ;
11
11
use crate :: {
@@ -20,6 +20,8 @@ use sqlx_core::arguments::Arguments;
20
20
use sqlx_core:: Either ;
21
21
use std:: { borrow:: Cow , pin:: pin, sync:: Arc } ;
22
22
23
+ use super :: pipe:: Pipe ;
24
+
23
25
async fn prepare (
24
26
conn : & mut PgConnection ,
25
27
sql : & str ,
@@ -46,54 +48,47 @@ async fn prepare(
46
48
param_types. push ( conn. resolve_type_id ( & ty. 0 ) . await ?) ;
47
49
}
48
50
49
- // flush and wait until we are re-ready
50
- conn. wait_until_ready ( ) . await ?;
51
+ let mut pipe = conn. pipe ( |buf| {
52
+ // next we send the PARSE command to the server
53
+ buf. write_msg ( Parse {
54
+ param_types : & param_types,
55
+ query : sql,
56
+ statement : id,
57
+ } ) ?;
58
+
59
+ if metadata. is_none ( ) {
60
+ // get the statement columns and parameters
61
+ buf. write_msg ( message:: Describe :: Statement ( id) ) ?;
62
+ }
51
63
52
- // next we send the PARSE command to the server
53
- conn. inner . stream . write_msg ( Parse {
54
- param_types : & param_types,
55
- query : sql,
56
- statement : id,
64
+ // we ask for the server to immediately send us the result of the PARSE command
65
+ buf. write_sync ( ) ;
66
+ Ok ( ( ) )
57
67
} ) ?;
58
68
59
- if metadata. is_none ( ) {
60
- // get the statement columns and parameters
61
- conn. inner
62
- . stream
63
- . write_msg ( message:: Describe :: Statement ( id) ) ?;
64
- }
65
-
66
- // we ask for the server to immediately send us the result of the PARSE command
67
- conn. write_sync ( ) ;
68
- conn. inner . stream . flush ( ) . await ?;
69
-
70
69
// indicates that the SQL query string is now successfully parsed and has semantic validity
71
- conn . inner . stream . recv_expect :: < ParseComplete > ( ) . await ?;
70
+ pipe . recv_expect :: < ParseComplete > ( ) . await ?;
72
71
73
72
let metadata = if let Some ( metadata) = metadata {
74
73
// each SYNC produces one READY FOR QUERY
75
- conn . recv_ready_for_query ( ) . await ?;
74
+ pipe . recv_ready_for_query ( ) . await ?;
76
75
77
76
// we already have metadata
78
77
metadata
79
78
} else {
80
- let parameters = recv_desc_params ( conn ) . await ?;
79
+ let parameters = recv_desc_params ( & mut pipe ) . await ?;
81
80
82
- let rows = recv_desc_rows ( conn ) . await ?;
81
+ let rows = recv_desc_rows ( & mut pipe ) . await ?;
83
82
84
83
// each SYNC produces one READY FOR QUERY
85
- conn . recv_ready_for_query ( ) . await ?;
84
+ pipe . recv_ready_for_query ( ) . await ?;
86
85
87
86
let parameters = conn. handle_parameter_description ( parameters) . await ?;
88
87
89
88
let ( columns, column_names) = conn
90
89
. handle_row_description ( rows, true , fetch_column_origin)
91
90
. await ?;
92
91
93
- // ensure that if we did fetch custom data, we wait until we are fully ready before
94
- // continuing
95
- conn. wait_until_ready ( ) . await ?;
96
-
97
92
Arc :: new ( PgStatementMetadata {
98
93
parameters,
99
94
columns,
@@ -104,12 +99,12 @@ async fn prepare(
104
99
Ok ( ( id, metadata) )
105
100
}
106
101
107
- async fn recv_desc_params ( conn : & mut PgConnection ) -> Result < ParameterDescription , Error > {
108
- conn . inner . stream . recv_expect ( ) . await
102
+ async fn recv_desc_params ( pipe : & mut Pipe ) -> Result < ParameterDescription , Error > {
103
+ pipe . recv_expect ( ) . await
109
104
}
110
105
111
- async fn recv_desc_rows ( conn : & mut PgConnection ) -> Result < Option < RowDescription > , Error > {
112
- let rows: Option < RowDescription > = match conn . inner . stream . recv ( ) . await ? {
106
+ async fn recv_desc_rows ( pipe : & mut Pipe ) -> Result < Option < RowDescription > , Error > {
107
+ let rows: Option < RowDescription > = match pipe . recv ( ) . await ? {
113
108
// describes the rows that will be returned when the statement is eventually executed
114
109
message if message. format == BackendMessageFormat :: RowDescription => {
115
110
Some ( message. decode ( ) ?)
@@ -130,44 +125,6 @@ async fn recv_desc_rows(conn: &mut PgConnection) -> Result<Option<RowDescription
130
125
}
131
126
132
127
impl PgConnection {
133
- // wait for CloseComplete to indicate a statement was closed
134
- pub ( super ) async fn wait_for_close_complete ( & mut self , mut count : usize ) -> Result < ( ) , Error > {
135
- // we need to wait for the [CloseComplete] to be returned from the server
136
- while count > 0 {
137
- match self . inner . stream . recv ( ) . await ? {
138
- message if message. format == BackendMessageFormat :: PortalSuspended => {
139
- // there was an open portal
140
- // this can happen if the last time a statement was used it was not fully executed
141
- }
142
-
143
- message if message. format == BackendMessageFormat :: CloseComplete => {
144
- // successfully closed the statement (and freed up the server resources)
145
- count -= 1 ;
146
- }
147
-
148
- message => {
149
- return Err ( err_protocol ! (
150
- "expecting PortalSuspended or CloseComplete but received {:?}" ,
151
- message. format
152
- ) ) ;
153
- }
154
- }
155
- }
156
-
157
- Ok ( ( ) )
158
- }
159
-
160
- #[ inline( always) ]
161
- pub ( crate ) fn write_sync ( & mut self ) {
162
- self . inner
163
- . stream
164
- . write_msg ( message:: Sync )
165
- . expect ( "BUG: Sync should not be too big for protocol" ) ;
166
-
167
- // all SYNC messages will return a ReadyForQuery
168
- self . inner . pending_ready_for_query_count += 1 ;
169
- }
170
-
171
128
async fn get_or_prepare (
172
129
& mut self ,
173
130
sql : & str ,
@@ -194,13 +151,14 @@ impl PgConnection {
194
151
195
152
if persistent && self . inner . cache_statement . is_enabled ( ) {
196
153
if let Some ( ( id, _) ) = self . inner . cache_statement . insert ( sql, statement. clone ( ) ) {
197
- self . inner . stream . write_msg ( Close :: Statement ( id) ) ?;
198
- self . write_sync ( ) ;
199
-
200
- self . inner . stream . flush ( ) . await ?;
201
-
202
- self . wait_for_close_complete ( 1 ) . await ?;
203
- self . recv_ready_for_query ( ) . await ?;
154
+ let mut pipe = self . pipe ( |buf| {
155
+ buf. write_msg ( Close :: Statement ( id) ) ?;
156
+ buf. write_sync ( ) ;
157
+ Ok ( ( ) )
158
+ } ) ?;
159
+
160
+ pipe. wait_for_close_complete ( 1 ) . await ?;
161
+ pipe. recv_ready_for_query ( ) . await ?;
204
162
}
205
163
}
206
164
@@ -216,10 +174,8 @@ impl PgConnection {
216
174
) -> Result < impl Stream < Item = Result < Either < PgQueryResult , PgRow > , Error > > + ' e , Error > {
217
175
let mut logger = QueryLogger :: new ( query, self . inner . log_settings . clone ( ) ) ;
218
176
219
- // before we continue, wait until we are "ready" to accept more queries
220
- self . wait_until_ready ( ) . await ?;
221
-
222
177
let mut metadata: Arc < PgStatementMetadata > ;
178
+ let mut pipe: Pipe ;
223
179
224
180
let format = if let Some ( mut arguments) = arguments {
225
181
// Check this before we write anything to the stream.
@@ -246,53 +202,50 @@ impl PgConnection {
246
202
// patch holes created during encoding
247
203
arguments. apply_patches ( self , & metadata. parameters ) . await ?;
248
204
249
- // consume messages till `ReadyForQuery` before bind and execute
250
- self . wait_until_ready ( ) . await ?;
251
-
252
- // bind to attach the arguments to the statement and create a portal
253
- self . inner . stream . write_msg ( Bind {
254
- portal : PortalId :: UNNAMED ,
255
- statement,
256
- formats : & [ PgValueFormat :: Binary ] ,
257
- num_params,
258
- params : & arguments. buffer ,
259
- result_formats : & [ PgValueFormat :: Binary ] ,
260
- } ) ?;
261
-
262
- // executes the portal up to the passed limit
263
- // the protocol-level limit acts nearly identically to the `LIMIT` in SQL
264
- self . inner . stream . write_msg ( message:: Execute {
265
- portal : PortalId :: UNNAMED ,
266
- // Non-zero limits cause query plan pessimization by disabling parallel workers:
267
- // https://github.com/launchbadge/sqlx/issues/3673
268
- limit : 0 ,
205
+ pipe = self . pipe ( |buf| {
206
+ // bind to attach the arguments to the statement and create a portal
207
+ buf. write_msg ( Bind {
208
+ portal : PortalId :: UNNAMED ,
209
+ statement,
210
+ formats : & [ PgValueFormat :: Binary ] ,
211
+ num_params,
212
+ params : & arguments. buffer ,
213
+ result_formats : & [ PgValueFormat :: Binary ] ,
214
+ } ) ?;
215
+
216
+ // executes the portal up to the passed limit
217
+ // the protocol-level limit acts nearly identically to the `LIMIT` in SQL
218
+ buf. write_msg ( message:: Execute {
219
+ portal : PortalId :: UNNAMED ,
220
+ // Non-zero limits cause query plan pessimization by disabling parallel workers:
221
+ // https://github.com/launchbadge/sqlx/issues/3673
222
+ limit : 0 ,
223
+ } ) ?;
224
+ // From https://www.postgresql.org/docs/current/protocol-flow.html:
225
+ //
226
+ // "An unnamed portal is destroyed at the end of the transaction, or as
227
+ // soon as the next Bind statement specifying the unnamed portal as
228
+ // destination is issued. (Note that a simple Query message also
229
+ // destroys the unnamed portal."
230
+
231
+ // we ask the database server to close the unnamed portal and free the associated resources
232
+ // earlier - after the execution of the current query.
233
+ buf. write_msg ( Close :: Portal ( PortalId :: UNNAMED ) ) ?;
234
+
235
+ // finally, [Sync] asks postgres to process the messages that we sent and respond with
236
+ // a [ReadyForQuery] message when it's completely done. Theoretically, we could send
237
+ // dozens of queries before a [Sync] and postgres can handle that. Execution on the server
238
+ // is still serial but it would reduce round-trips. Some kind of builder pattern that is
239
+ // termed batching might suit this.
240
+ buf. write_sync ( ) ;
241
+ Ok ( ( ) )
269
242
} ) ?;
270
- // From https://www.postgresql.org/docs/current/protocol-flow.html:
271
- //
272
- // "An unnamed portal is destroyed at the end of the transaction, or as
273
- // soon as the next Bind statement specifying the unnamed portal as
274
- // destination is issued. (Note that a simple Query message also
275
- // destroys the unnamed portal."
276
-
277
- // we ask the database server to close the unnamed portal and free the associated resources
278
- // earlier - after the execution of the current query.
279
- self . inner
280
- . stream
281
- . write_msg ( Close :: Portal ( PortalId :: UNNAMED ) ) ?;
282
-
283
- // finally, [Sync] asks postgres to process the messages that we sent and respond with
284
- // a [ReadyForQuery] message when it's completely done. Theoretically, we could send
285
- // dozens of queries before a [Sync] and postgres can handle that. Execution on the server
286
- // is still serial but it would reduce round-trips. Some kind of builder pattern that is
287
- // termed batching might suit this.
288
- self . write_sync ( ) ;
289
243
290
244
// prepared statements are binary
291
245
PgValueFormat :: Binary
292
246
} else {
293
247
// Query will trigger a ReadyForQuery
294
- self . inner . stream . write_msg ( Query ( query) ) ?;
295
- self . inner . pending_ready_for_query_count += 1 ;
248
+ pipe = self . queue_simple_query ( query) ?;
296
249
297
250
// metadata starts out as "nothing"
298
251
metadata = Arc :: new ( PgStatementMetadata :: default ( ) ) ;
@@ -301,11 +254,9 @@ impl PgConnection {
301
254
PgValueFormat :: Text
302
255
} ;
303
256
304
- self . inner . stream . flush ( ) . await ?;
305
-
306
257
Ok ( try_stream ! {
307
258
loop {
308
- let message = self . inner . stream . recv( ) . await ?;
259
+ let message = pipe . recv( ) . await ?;
309
260
310
261
match message. format {
311
262
BackendMessageFormat :: BindComplete
@@ -463,8 +414,6 @@ impl<'c> Executor<'c> for &'c mut PgConnection {
463
414
' c : ' e ,
464
415
{
465
416
Box :: pin ( async move {
466
- self . wait_until_ready ( ) . await ?;
467
-
468
417
let ( _, metadata) = self
469
418
. get_or_prepare ( sql, parameters, true , None , true )
470
419
. await ?;
@@ -484,8 +433,6 @@ impl<'c> Executor<'c> for &'c mut PgConnection {
484
433
' c : ' e ,
485
434
{
486
435
Box :: pin ( async move {
487
- self . wait_until_ready ( ) . await ?;
488
-
489
436
let ( stmt_id, metadata) = self . get_or_prepare ( sql, & [ ] , true , None , true ) . await ?;
490
437
491
438
let nullable = self . get_nullable_for_columns ( stmt_id, & metadata) . await ?;
0 commit comments