pg_proto/
resources.rs

1//! Branded prepared-statement and portal resources with proxy name rewriting.
2
3use std::{
4    cell::Cell, collections::HashMap, error::Error, fmt, future::Future, io, marker::PhantomData,
5    pin::Pin,
6};
7
8use bytes::Bytes;
9
10use crate::{
11    Conn, Dirty,
12    auth::Ready,
13    codec::{Bind, Close, Describe, DescribeTarget, Execute, Frame, Parse, TransactionStatus},
14    demux::SessionItem,
15    session::{
16        AwaitingReady, AwaitingReadyTransition, BoundBuilding, Building, Draining,
17        DrainingTransition, ErrorResponse, ReadyState, SimpleQuery,
18    },
19};
20
21/// Runs an operation with a fresh resource brand which cannot escape the closure.
22pub fn with_resources<R>(operation: impl for<'id> FnOnce(ResourceScope<'id>) -> R) -> R {
23    operation(ResourceScope::new())
24}
25
26/// Runs an extended-query operation with one brand shared by its connection
27/// and resource namespace.
28pub fn with_connection_resources<S, P, C, R>(
29    conn: Conn<S, P, C>,
30    operation: impl for<'id> FnOnce(ResourceConnection<'id, S, P, C>) -> R,
31) -> R {
32    operation(ResourceConnection {
33        conn,
34        resources: ResourceScope::new(),
35    })
36}
37
38/// Runs an asynchronous operation with one brand shared by its connection and
39/// resource namespace.
40///
41/// The boxed future may retain branded tokens across await points, while its
42/// output cannot contain the generative lifetime.
43pub async fn with_connection_resources_async<S, P, C, R>(
44    conn: Conn<S, P, C>,
45    operation: impl for<'id> FnOnce(
46        ResourceConnection<'id, S, P, C>,
47    ) -> Pin<Box<dyn Future<Output = R> + 'id>>,
48) -> R {
49    operation(ResourceConnection {
50        conn,
51        resources: ResourceScope::new(),
52    })
53    .await
54}
55
56/// A connection paired with its generative statement and portal namespace.
57#[derive(Debug)]
58pub struct ResourceConnection<'id, S, P, C> {
59    conn: Conn<S, P, C>,
60    resources: ResourceScope<'id>,
61}
62
63impl<'id, S, P, C> ResourceConnection<'id, S, P, C> {
64    /// Borrows the typed connection for transport-only operations such as
65    /// buffering a frame returned by this wrapper.
66    pub const fn connection(&self) -> &Conn<S, P, C> {
67        &self.conn
68    }
69
70    /// Mutably borrows the typed connection for transport-only operations such
71    /// as buffering and flushing returned frames.
72    pub const fn connection_mut(&mut self) -> &mut Conn<S, P, C> {
73        &mut self.conn
74    }
75
76    /// Reports whether a prepared-statement token is still live in this
77    /// connection's namespace.
78    #[must_use]
79    pub fn statement_is_live(&self, statement: &PreparedStatement<'id>) -> bool {
80        self.resources.statements.get(&statement.upstream_name) == Some(&statement.generation)
81    }
82
83    /// Reports whether a portal token is still live in this connection's
84    /// namespace.
85    #[must_use]
86    pub fn portal_is_live(&self, portal: &Portal<'id>) -> bool {
87        self.resources.portals.get(&portal.upstream_name) == Some(&portal.generation)
88    }
89
90    /// Deliberately leaves resource-aware handling while retaining typestate.
91    pub fn into_connection(self) -> Conn<S, P, C> {
92        self.conn
93    }
94}
95
96/// Result of preparing a statement while building an extended-query pipeline.
97pub type PrepareResult<'id, S> = Result<
98    (
99        ResourceConnection<'id, S, Building, Dirty>,
100        PreparedStatement<'id>,
101        Frame,
102    ),
103    ResourceProtocolError,
104>;
105
106/// Result of binding a portal for the first time in a pipeline.
107pub type BindResult<'id, S> = Result<
108    (
109        ResourceConnection<'id, S, BoundBuilding, Dirty>,
110        Portal<'id>,
111        Frame,
112    ),
113    ResourceProtocolError,
114>;
115
116/// Result of preparing another statement after a portal has been bound.
117pub type BoundPrepareResult<'id, S> = Result<
118    (
119        ResourceConnection<'id, S, BoundBuilding, Dirty>,
120        PreparedStatement<'id>,
121        Frame,
122    ),
123    ResourceProtocolError,
124>;
125
126/// Result of binding another portal after the pipeline has become executable.
127pub type RebindResult<'id, S> = Result<
128    (
129        ResourceConnection<'id, S, BoundBuilding, Dirty>,
130        Portal<'id>,
131        Frame,
132    ),
133    ResourceProtocolError,
134>;
135
136#[derive(Debug)]
137/// Readiness projected while retaining the connection's resource brand.
138pub enum ResourceReadyState<'id, S, C> {
139    /// The connection retained its existing cleanliness index.
140    Clean(ResourceConnection<'id, S, Ready, C>),
141    /// Transaction or parameter evidence made the connection dirty.
142    Dirty {
143        /// Ready, dirty connection and its resource namespace.
144        conn: ResourceConnection<'id, S, Ready, Dirty>,
145        /// Transaction status reported by `ReadyForQuery`.
146        status: TransactionStatus,
147        /// Whether reported parameters differ from their startup values.
148        parameters_changed: bool,
149    },
150}
151
152#[derive(Debug)]
153/// Projection while awaiting readiness after an extended-query result.
154pub enum ResourceAwaitingTransition<'id, S, C> {
155    /// A non-terminal item was consumed; continue waiting.
156    Continue(ResourceConnection<'id, S, AwaitingReady, C>, SessionItem),
157    /// `ReadyForQuery` completed the cycle.
158    Ready(ResourceReadyState<'id, S, C>),
159    /// An error entered the drain-until-ready recovery phase.
160    Error(ResourceConnection<'id, S, Draining, C>, ErrorResponse),
161}
162
163#[derive(Debug)]
164/// Projection while draining an errored resource-aware pipeline.
165pub enum ResourceDrainingTransition<'id, S, C> {
166    /// A non-terminal item was consumed; continue draining.
167    Continue(ResourceConnection<'id, S, Draining, C>, SessionItem),
168    /// `ReadyForQuery` completed recovery.
169    Ready(ResourceReadyState<'id, S, C>),
170}
171
172/// Connection-local statement and portal namespaces.
173#[derive(Debug)]
174pub struct ResourceScope<'id> {
175    statements: HashMap<Bytes, u64>,
176    client_statements: HashMap<Bytes, (Bytes, u64)>,
177    portals: HashMap<Bytes, u64>,
178    client_portals: HashMap<Bytes, (Bytes, u64)>,
179    next_generation: u64,
180    _brand: PhantomData<Cell<&'id ()>>,
181}
182
183/// A prepared statement tied to one generative connection brand.
184#[derive(Debug)]
185pub struct PreparedStatement<'id> {
186    client_name: Bytes,
187    upstream_name: Bytes,
188    generation: u64,
189    _brand: PhantomData<Cell<&'id ()>>,
190}
191
192/// A bound portal tied to the same brand as its statement.
193#[derive(Debug)]
194pub struct Portal<'id> {
195    client_name: Bytes,
196    upstream_name: Bytes,
197    generation: u64,
198    _brand: PhantomData<Cell<&'id ()>>,
199}
200
201#[derive(Clone, Debug, Eq, PartialEq)]
202/// Failure to resolve or allocate a branded protocol resource.
203pub enum ResourceError {
204    /// A live prepared statement already uses the requested name.
205    StatementNameCollision,
206    /// A live portal already uses the requested name.
207    PortalNameCollision,
208    /// The prepared-statement token or client name is unknown or stale.
209    UnknownStatement,
210    /// The portal token or client name is unknown or stale.
211    UnknownPortal,
212}
213
214impl fmt::Display for ResourceError {
215    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
216        formatter.write_str(match self {
217            Self::StatementNameCollision => "prepared statement name collision",
218            Self::PortalNameCollision => "portal name collision",
219            Self::UnknownStatement => "unknown or stale prepared statement",
220            Self::UnknownPortal => "unknown or stale portal",
221        })
222    }
223}
224
225impl Error for ResourceError {}
226
227#[derive(Debug)]
228/// A resource-namespace or wire-encoding failure.
229pub enum ResourceProtocolError {
230    /// Resource identity or lifetime validation failed.
231    Resource(ResourceError),
232    /// Reconstruction of the wire message failed.
233    Wire(io::Error),
234}
235
236impl fmt::Display for ResourceProtocolError {
237    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
238        match self {
239            Self::Resource(error) => error.fmt(formatter),
240            Self::Wire(error) => error.fmt(formatter),
241        }
242    }
243}
244
245impl Error for ResourceProtocolError {
246    fn source(&self) -> Option<&(dyn Error + 'static)> {
247        match self {
248            Self::Resource(error) => Some(error),
249            Self::Wire(error) => Some(error),
250        }
251    }
252}
253
254impl From<ResourceError> for ResourceProtocolError {
255    fn from(error: ResourceError) -> Self {
256        Self::Resource(error)
257    }
258}
259
260impl From<io::Error> for ResourceProtocolError {
261    fn from(error: io::Error) -> Self {
262        Self::Wire(error)
263    }
264}
265
266impl<'id> ResourceScope<'id> {
267    fn new() -> Self {
268        Self {
269            statements: HashMap::new(),
270            client_statements: HashMap::new(),
271            portals: HashMap::new(),
272            client_portals: HashMap::new(),
273            next_generation: 0,
274            _brand: PhantomData,
275        }
276    }
277
278    /// Records a simple-query boundary, which destroys the unnamed statement.
279    pub fn simple_query_boundary(&mut self) {
280        self.statements.remove(b"".as_slice());
281        self.client_statements
282            .retain(|_, (upstream, _)| !upstream.is_empty());
283    }
284
285    /// Records transaction end, which destroys the unnamed portal.
286    pub fn transaction_ended(&mut self) {
287        self.portals.remove(b"".as_slice());
288        self.client_portals
289            .retain(|_, (upstream, _)| !upstream.is_empty());
290    }
291
292    /// Allocates a statement token and reconstructable upstream `Parse` message.
293    ///
294    /// # Errors
295    ///
296    /// Rejects duplicate upstream statement names.
297    pub fn prepare(
298        &mut self,
299        client_name: Bytes,
300        upstream_name: Bytes,
301        query: Bytes,
302        parameter_types: Vec<u32>,
303    ) -> Result<(PreparedStatement<'id>, Parse), ResourceError> {
304        let generation = self.allocate(&client_name, &upstream_name, true)?;
305        if self
306            .statements
307            .insert(upstream_name.clone(), generation)
308            .is_some()
309            && !upstream_name.is_empty()
310        {
311            return Err(ResourceError::StatementNameCollision);
312        }
313        let statement = PreparedStatement {
314            client_name: client_name.clone(),
315            upstream_name: upstream_name.clone(),
316            generation,
317            _brand: PhantomData,
318        };
319        self.client_statements
320            .insert(client_name, (upstream_name.clone(), generation));
321        let message = Parse {
322            statement: upstream_name,
323            query,
324            parameter_types,
325        };
326        Ok((statement, message))
327    }
328
329    /// Binds a branded statement into a rewritten portal namespace.
330    ///
331    /// # Errors
332    ///
333    /// Rejects statements not present in this scope and duplicate portal names.
334    pub fn bind(
335        &mut self,
336        statement: &PreparedStatement<'id>,
337        client_name: Bytes,
338        upstream_name: Bytes,
339        parameter_formats: Vec<i16>,
340        parameters: Vec<Option<Bytes>>,
341        result_formats: Vec<i16>,
342    ) -> Result<(Portal<'id>, Bind), ResourceError> {
343        if self.statements.get(&statement.upstream_name) != Some(&statement.generation) {
344            return Err(ResourceError::UnknownStatement);
345        }
346        let generation = self.allocate(&client_name, &upstream_name, false)?;
347        if self
348            .portals
349            .insert(upstream_name.clone(), generation)
350            .is_some()
351            && !upstream_name.is_empty()
352        {
353            return Err(ResourceError::PortalNameCollision);
354        }
355        let portal = Portal {
356            client_name: client_name.clone(),
357            upstream_name: upstream_name.clone(),
358            generation,
359            _brand: PhantomData,
360        };
361        self.client_portals
362            .insert(client_name, (upstream_name.clone(), generation));
363        let message = Bind {
364            portal: upstream_name,
365            statement: statement.upstream_name.clone(),
366            parameter_formats,
367            parameters,
368            result_formats,
369        };
370        Ok((portal, message))
371    }
372
373    /// Closes a statement and removes its upstream name from this scope.
374    ///
375    /// # Errors
376    ///
377    /// Rejects a token which has already been closed.
378    pub fn close_statement(
379        &mut self,
380        statement: PreparedStatement<'id>,
381    ) -> Result<Close, ResourceError> {
382        if self.statements.get(&statement.upstream_name) != Some(&statement.generation) {
383            return Err(ResourceError::UnknownStatement);
384        }
385        self.statements.remove(&statement.upstream_name);
386        self.client_statements.retain(|_, (upstream, generation)| {
387            upstream != &statement.upstream_name || *generation != statement.generation
388        });
389        Ok(Close {
390            target: DescribeTarget::Statement,
391            name: statement.upstream_name,
392        })
393    }
394
395    /// Closes a portal and removes its upstream name from this scope.
396    ///
397    /// # Errors
398    ///
399    /// Rejects a token which has already been closed.
400    pub fn close_portal(&mut self, portal: Portal<'id>) -> Result<Close, ResourceError> {
401        if self.portals.get(&portal.upstream_name) != Some(&portal.generation) {
402            return Err(ResourceError::UnknownPortal);
403        }
404        self.portals.remove(&portal.upstream_name);
405        self.client_portals.retain(|_, (upstream, generation)| {
406            upstream != &portal.upstream_name || *generation != portal.generation
407        });
408        Ok(Close {
409            target: DescribeTarget::Portal,
410            name: portal.upstream_name,
411        })
412    }
413
414    fn execute(&self, portal: &Portal<'id>, max_rows: i32) -> Result<Execute, ResourceError> {
415        if self.portals.get(&portal.upstream_name) != Some(&portal.generation) {
416            return Err(ResourceError::UnknownPortal);
417        }
418        Ok(portal.execute(max_rows))
419    }
420
421    fn describe_portal(&self, portal: &Portal<'id>) -> Result<Describe, ResourceError> {
422        if self.portals.get(&portal.upstream_name) != Some(&portal.generation) {
423            return Err(ResourceError::UnknownPortal);
424        }
425        Ok(portal.describe())
426    }
427
428    fn describe_statement(
429        &self,
430        statement: &PreparedStatement<'id>,
431    ) -> Result<Describe, ResourceError> {
432        if self.statements.get(&statement.upstream_name) != Some(&statement.generation) {
433            return Err(ResourceError::UnknownStatement);
434        }
435        Ok(statement.describe())
436    }
437
438    /// Resolves a client-visible statement name to its branded upstream token.
439    #[must_use]
440    pub fn statement(&self, client_name: &[u8]) -> Option<PreparedStatement<'id>> {
441        let (upstream_name, generation) = self.client_statements.get(client_name)?;
442        Some(PreparedStatement {
443            client_name: Bytes::copy_from_slice(client_name),
444            upstream_name: upstream_name.clone(),
445            generation: *generation,
446            _brand: PhantomData,
447        })
448    }
449
450    /// Resolves a client-visible portal name to its branded upstream token.
451    #[must_use]
452    pub fn portal(&self, client_name: &[u8]) -> Option<Portal<'id>> {
453        let (upstream_name, generation) = self.client_portals.get(client_name)?;
454        Some(Portal {
455            client_name: Bytes::copy_from_slice(client_name),
456            upstream_name: upstream_name.clone(),
457            generation: *generation,
458            _brand: PhantomData,
459        })
460    }
461}
462
463impl<'id, S, C> ResourceConnection<'id, S, Building, C> {
464    /// Creates and sends a branded prepared statement on this connection.
465    ///
466    /// # Errors
467    ///
468    /// Returns namespace or wire reconstruction errors.
469    pub fn prepare(
470        self,
471        client_name: Bytes,
472        upstream_name: Bytes,
473        query: Bytes,
474        parameter_types: Vec<u32>,
475    ) -> PrepareResult<'id, S> {
476        let Self {
477            conn,
478            mut resources,
479        } = self;
480        let (statement, message) =
481            resources.prepare(client_name, upstream_name, query, parameter_types)?;
482        let (conn, frame) = conn.push_parse(&message)?;
483        Ok((ResourceConnection { conn, resources }, statement, frame))
484    }
485
486    /// Creates and sends a branded portal on this connection.
487    ///
488    /// # Errors
489    ///
490    /// Returns namespace or wire reconstruction errors.
491    #[allow(clippy::too_many_arguments)]
492    pub fn bind(
493        self,
494        statement: &PreparedStatement<'id>,
495        client_name: Bytes,
496        upstream_name: Bytes,
497        parameter_formats: Vec<i16>,
498        parameters: Vec<Option<Bytes>>,
499        result_formats: Vec<i16>,
500    ) -> BindResult<'id, S> {
501        let Self {
502            conn,
503            mut resources,
504        } = self;
505        let (portal, message) = resources.bind(
506            statement,
507            client_name,
508            upstream_name,
509            parameter_formats,
510            parameters,
511            result_formats,
512        )?;
513        let (conn, frame) = conn.push_bind(&message)?;
514        Ok((ResourceConnection { conn, resources }, portal, frame))
515    }
516
517    /// Describes only a live statement from this connection.
518    ///
519    /// # Errors
520    ///
521    /// Returns an error for a stale statement or invalid wire value.
522    pub fn describe_statement(
523        self,
524        statement: &PreparedStatement<'id>,
525    ) -> Result<(Self, Frame), ResourceProtocolError> {
526        let message = self.resources.describe_statement(statement)?;
527        let Self { conn, resources } = self;
528        let (conn, frame) = conn.push_describe(&message)?;
529        Ok((Self { conn, resources }, frame))
530    }
531
532    /// Closes a live statement and invalidates its token.
533    ///
534    /// # Errors
535    ///
536    /// Returns an error for a stale statement or invalid wire value.
537    pub fn close_statement(
538        self,
539        statement: PreparedStatement<'id>,
540    ) -> Result<(Self, Frame), ResourceProtocolError> {
541        let Self {
542            conn,
543            mut resources,
544        } = self;
545        let message = resources.close_statement(statement)?;
546        let (conn, frame) = conn.push_close(&message)?;
547        Ok((Self { conn, resources }, frame))
548    }
549
550    /// Closes a live portal and invalidates its token.
551    ///
552    /// # Errors
553    ///
554    /// Returns an error for a stale portal or invalid wire value.
555    pub fn close_portal(self, portal: Portal<'id>) -> Result<(Self, Frame), ResourceProtocolError> {
556        let Self {
557            conn,
558            mut resources,
559        } = self;
560        let message = resources.close_portal(portal)?;
561        let (conn, frame) = conn.push_close(&message)?;
562        Ok((Self { conn, resources }, frame))
563    }
564
565    /// Emits Flush without changing resource or phase evidence.
566    #[must_use]
567    pub fn flush(self) -> (Self, Frame) {
568        let Self { conn, resources } = self;
569        let (conn, frame) = conn.push_flush();
570        (Self { conn, resources }, frame)
571    }
572
573    /// Emits Sync while retaining the namespace through response consumption.
574    #[must_use]
575    pub fn sync(self) -> (ResourceConnection<'id, S, AwaitingReady, C>, Frame) {
576        let Self { conn, resources } = self;
577        let (conn, frame) = conn.push_sync();
578        (ResourceConnection { conn, resources }, frame)
579    }
580}
581
582impl<'id, S, C> ResourceConnection<'id, S, Ready, C> {
583    /// Begins another extended-query cycle with the same resource namespace.
584    #[must_use]
585    pub fn begin_extended(self) -> ResourceConnection<'id, S, Building, C> {
586        let Self { conn, resources } = self;
587        ResourceConnection {
588            conn: conn.begin_extended(),
589            resources,
590        }
591    }
592
593    /// Begins a simple query and invalidates the unnamed prepared statement.
594    ///
595    /// # Errors
596    ///
597    /// Returns an error if the query cannot be reconstructed on the wire.
598    pub fn query(
599        self,
600        query: &[u8],
601    ) -> Result<(ResourceConnection<'id, S, SimpleQuery, Dirty>, Frame), ResourceProtocolError>
602    {
603        let Self {
604            conn,
605            mut resources,
606        } = self;
607        resources.simple_query_boundary();
608        let (conn, frame) = conn.push_query(query)?;
609        Ok((ResourceConnection { conn, resources }, frame))
610    }
611}
612
613impl<'id, S, C> ResourceConnection<'id, S, BoundBuilding, C> {
614    /// Creates another prepared statement while retaining executable portals.
615    ///
616    /// # Errors
617    ///
618    /// Returns namespace or wire reconstruction errors.
619    pub fn prepare(
620        self,
621        client_name: Bytes,
622        upstream_name: Bytes,
623        query: Bytes,
624        parameter_types: Vec<u32>,
625    ) -> BoundPrepareResult<'id, S> {
626        let Self {
627            conn,
628            mut resources,
629        } = self;
630        let (statement, message) =
631            resources.prepare(client_name, upstream_name, query, parameter_types)?;
632        let (conn, frame) = conn.push_parse(&message)?;
633        Ok((ResourceConnection { conn, resources }, statement, frame))
634    }
635
636    /// Creates another portal while retaining prior live portals.
637    ///
638    /// # Errors
639    ///
640    /// Returns namespace or wire reconstruction errors.
641    #[allow(clippy::too_many_arguments)]
642    pub fn bind(
643        self,
644        statement: &PreparedStatement<'id>,
645        client_name: Bytes,
646        upstream_name: Bytes,
647        parameter_formats: Vec<i16>,
648        parameters: Vec<Option<Bytes>>,
649        result_formats: Vec<i16>,
650    ) -> RebindResult<'id, S> {
651        let Self {
652            conn,
653            mut resources,
654        } = self;
655        let (portal, message) = resources.bind(
656            statement,
657            client_name,
658            upstream_name,
659            parameter_formats,
660            parameters,
661            result_formats,
662        )?;
663        let (conn, frame) = conn.push_bind(&message)?;
664        Ok((ResourceConnection { conn, resources }, portal, frame))
665    }
666
667    /// Describes only a live statement from this connection.
668    ///
669    /// # Errors
670    ///
671    /// Returns an error for a stale statement or invalid wire value.
672    pub fn describe_statement(
673        self,
674        statement: &PreparedStatement<'id>,
675    ) -> Result<(Self, Frame), ResourceProtocolError> {
676        let message = self.resources.describe_statement(statement)?;
677        let Self { conn, resources } = self;
678        let (conn, frame) = conn.push_describe(&message)?;
679        Ok((Self { conn, resources }, frame))
680    }
681
682    /// Describes only a live portal from this connection.
683    ///
684    /// # Errors
685    ///
686    /// Returns an error for a stale portal or invalid wire value.
687    pub fn describe_portal(
688        self,
689        portal: &Portal<'id>,
690    ) -> Result<(Self, Frame), ResourceProtocolError> {
691        let message = self.resources.describe_portal(portal)?;
692        let Self { conn, resources } = self;
693        let (conn, frame) = conn.push_describe(&message)?;
694        Ok((Self { conn, resources }, frame))
695    }
696
697    /// Sends an execute which can name only a live portal from this connection.
698    ///
699    /// # Errors
700    ///
701    /// Returns an error for a stale portal or invalid wire value.
702    pub fn execute(
703        self,
704        portal: &Portal<'id>,
705        max_rows: i32,
706    ) -> Result<(Self, Frame), ResourceProtocolError> {
707        let message = self.resources.execute(portal, max_rows)?;
708        let Self { conn, resources } = self;
709        let (conn, frame) = conn.push_execute(&message)?;
710        Ok((Self { conn, resources }, frame))
711    }
712
713    /// Closes a live statement and invalidates its token.
714    ///
715    /// # Errors
716    ///
717    /// Returns an error for a stale statement or invalid wire value.
718    pub fn close_statement(
719        self,
720        statement: PreparedStatement<'id>,
721    ) -> Result<(Self, Frame), ResourceProtocolError> {
722        let Self {
723            conn,
724            mut resources,
725        } = self;
726        let message = resources.close_statement(statement)?;
727        let (conn, frame) = conn.push_close(&message)?;
728        Ok((Self { conn, resources }, frame))
729    }
730
731    /// Closes a live portal and invalidates its token.
732    ///
733    /// # Errors
734    ///
735    /// Returns an error for a stale portal or invalid wire value.
736    pub fn close_portal(self, portal: Portal<'id>) -> Result<(Self, Frame), ResourceProtocolError> {
737        let Self {
738            conn,
739            mut resources,
740        } = self;
741        let message = resources.close_portal(portal)?;
742        let (conn, frame) = conn.push_close(&message)?;
743        Ok((Self { conn, resources }, frame))
744    }
745
746    /// Emits Flush without changing resource or phase evidence.
747    #[must_use]
748    pub fn flush(self) -> (Self, Frame) {
749        let Self { conn, resources } = self;
750        let (conn, frame) = conn.push_flush();
751        (Self { conn, resources }, frame)
752    }
753
754    /// Emits Sync while retaining the namespace through response consumption.
755    #[must_use]
756    pub fn sync(self) -> (ResourceConnection<'id, S, AwaitingReady, C>, Frame) {
757        let Self { conn, resources } = self;
758        let (conn, frame) = conn.push_sync();
759        (ResourceConnection { conn, resources }, frame)
760    }
761}
762
763impl<'id, S, C> ResourceConnection<'id, S, AwaitingReady, C> {
764    /// Consumes one backend response while retaining the branded namespace.
765    #[must_use]
766    pub fn offer(self, item: SessionItem) -> ResourceAwaitingTransition<'id, S, C> {
767        let Self { conn, resources } = self;
768        match conn.offer(item) {
769            AwaitingReadyTransition::Continue(conn, item) => {
770                ResourceAwaitingTransition::Continue(Self { conn, resources }, item)
771            }
772            AwaitingReadyTransition::Ready(ready) => {
773                ResourceAwaitingTransition::Ready(resource_ready(resources, ready))
774            }
775            AwaitingReadyTransition::Error(conn, error) => {
776                ResourceAwaitingTransition::Error(ResourceConnection { conn, resources }, error)
777            }
778        }
779    }
780}
781
782impl<'id, S, C> ResourceConnection<'id, S, Draining, C> {
783    /// Drains one backend response after an error while retaining resources.
784    #[must_use]
785    pub fn offer(self, item: SessionItem) -> ResourceDrainingTransition<'id, S, C> {
786        let Self { conn, resources } = self;
787        match conn.offer(item) {
788            DrainingTransition::Continue(conn, item) => {
789                ResourceDrainingTransition::Continue(Self { conn, resources }, item)
790            }
791            DrainingTransition::Ready(ready) => {
792                ResourceDrainingTransition::Ready(resource_ready(resources, ready))
793            }
794        }
795    }
796}
797
798fn resource_ready<S, C>(
799    mut resources: ResourceScope<'_>,
800    ready: ReadyState<S, C>,
801) -> ResourceReadyState<'_, S, C> {
802    match ready {
803        ReadyState::Clean(conn) => {
804            resources.transaction_ended();
805            ResourceReadyState::Clean(ResourceConnection { conn, resources })
806        }
807        ReadyState::Dirty {
808            conn,
809            status,
810            parameters_changed,
811        } => {
812            if status == TransactionStatus::Idle {
813                resources.transaction_ended();
814            }
815            ResourceReadyState::Dirty {
816                conn: ResourceConnection { conn, resources },
817                status,
818                parameters_changed,
819            }
820        }
821    }
822}
823
824impl ResourceScope<'_> {
825    fn allocate(
826        &mut self,
827        client_name: &Bytes,
828        upstream_name: &Bytes,
829        statement: bool,
830    ) -> Result<u64, ResourceError> {
831        let (resources, client_resources) = if statement {
832            (&self.statements, &self.client_statements)
833        } else {
834            (&self.portals, &self.client_portals)
835        };
836        if (!upstream_name.is_empty() && resources.contains_key(upstream_name))
837            || (!client_name.is_empty() && client_resources.contains_key(client_name))
838        {
839            return Err(if statement {
840                ResourceError::StatementNameCollision
841            } else {
842                ResourceError::PortalNameCollision
843            });
844        }
845        let generation = self.next_generation;
846        self.next_generation = self.next_generation.saturating_add(1);
847        Ok(generation)
848    }
849}
850
851impl PreparedStatement<'_> {
852    /// Returns the statement name presented by the client.
853    #[must_use]
854    pub fn client_name(&self) -> &[u8] {
855        &self.client_name
856    }
857
858    /// Returns the rewritten statement name sent upstream.
859    #[must_use]
860    pub fn upstream_name(&self) -> &[u8] {
861        &self.upstream_name
862    }
863
864    /// Constructs a `Describe` message using the rewritten statement name.
865    #[must_use]
866    pub fn describe(&self) -> Describe {
867        Describe {
868            target: DescribeTarget::Statement,
869            name: self.upstream_name.clone(),
870        }
871    }
872}
873
874impl Portal<'_> {
875    /// Returns the portal name presented by the client.
876    #[must_use]
877    pub fn client_name(&self) -> &[u8] {
878        &self.client_name
879    }
880
881    /// Returns the rewritten portal name sent upstream.
882    #[must_use]
883    pub fn upstream_name(&self) -> &[u8] {
884        &self.upstream_name
885    }
886
887    /// Constructs a `Describe` message using the rewritten portal name.
888    #[must_use]
889    pub fn describe(&self) -> Describe {
890        Describe {
891            target: DescribeTarget::Portal,
892            name: self.upstream_name.clone(),
893        }
894    }
895
896    /// Constructs an `Execute` message using the rewritten portal name.
897    #[must_use]
898    pub fn execute(&self, max_rows: i32) -> Execute {
899        Execute {
900            portal: self.upstream_name.clone(),
901            max_rows,
902        }
903    }
904}
905
906#[cfg(test)]
907mod tests {
908    use super::*;
909
910    #[test]
911    fn resource_connection_sends_only_its_own_branded_portals() {
912        let ready: Conn<(), crate::auth::Ready> = Conn::new(()).transition();
913        with_connection_resources(ready.begin_extended(), |connection| {
914            let (connection, statement, parse) = connection
915                .prepare(
916                    Bytes::from_static(b"client_statement"),
917                    Bytes::from_static(b"proxy_statement"),
918                    Bytes::from_static(b"select $1::int4"),
919                    vec![23],
920                )
921                .unwrap();
922            assert_eq!(parse.tag, b'P');
923            let (connection, portal, bind) = connection
924                .bind(
925                    &statement,
926                    Bytes::from_static(b"client_portal"),
927                    Bytes::from_static(b"proxy_portal"),
928                    vec![1],
929                    vec![Some(Bytes::from_static(b"\0\0\0*"))],
930                    vec![1],
931                )
932                .unwrap();
933            assert_eq!(bind.tag, b'B');
934            let (connection, describe) = connection.describe_portal(&portal).unwrap();
935            assert_eq!(describe.tag, b'D');
936            let (connection, execute) = connection.execute(&portal, 0).unwrap();
937            assert_eq!(execute.tag, b'E');
938            let (connection, second_statement, parse) = connection
939                .prepare(
940                    Bytes::from_static(b"client_statement_2"),
941                    Bytes::from_static(b"proxy_statement_2"),
942                    Bytes::from_static(b"select 2"),
943                    vec![],
944                )
945                .unwrap();
946            assert_eq!(parse.tag, b'P');
947            let (connection, second_portal, bind) = connection
948                .bind(
949                    &second_statement,
950                    Bytes::from_static(b"client_portal_2"),
951                    Bytes::from_static(b"proxy_portal_2"),
952                    vec![],
953                    vec![],
954                    vec![],
955                )
956                .unwrap();
957            assert_eq!(bind.tag, b'B');
958            let (connection, describe) = connection.describe_statement(&second_statement).unwrap();
959            assert_eq!(describe.tag, b'D');
960            let (connection, flush) = connection.flush();
961            assert_eq!(flush.tag, b'H');
962            let (connection, close) = connection.close_portal(second_portal).unwrap();
963            assert_eq!(close.tag, b'C');
964            let (connection, close) = connection.close_portal(portal).unwrap();
965            assert_eq!(close.tag, b'C');
966            let (connection, close) = connection.close_statement(second_statement).unwrap();
967            assert_eq!(close.tag, b'C');
968            let (connection, close) = connection.close_statement(statement).unwrap();
969            assert_eq!(close.tag, b'C');
970            let (awaiting, sync) = connection.sync();
971            assert_eq!(sync.tag, b'S');
972            let ResourceAwaitingTransition::Continue(awaiting, _) = awaiting.offer(
973                SessionItem::Message(crate::codec::BackendMessage::ParseComplete),
974            ) else {
975                panic!("ParseComplete should retain the awaiting phase")
976            };
977            let ResourceAwaitingTransition::Ready(ResourceReadyState::Clean(ready)) = awaiting
978                .offer(SessionItem::ReadyForQuery {
979                    status: TransactionStatus::Idle,
980                    parameters_changed: false,
981                })
982            else {
983                panic!("idle readiness should complete the extended cycle")
984            };
985            ready.begin_extended().into_connection().into_transport();
986        });
987    }
988
989    #[test]
990    fn namespace_resolves_client_names_to_rewritten_resources() {
991        with_resources(|mut resources| {
992            let (statement, _) = resources
993                .prepare(
994                    Bytes::from_static(b"client-statement"),
995                    Bytes::from_static(b"upstream-statement-42"),
996                    Bytes::from_static(b"select $1"),
997                    vec![25],
998                )
999                .unwrap();
1000            let (portal, _) = resources
1001                .bind(
1002                    &statement,
1003                    Bytes::from_static(b"client-portal"),
1004                    Bytes::from_static(b"upstream-portal-42"),
1005                    vec![0],
1006                    vec![Some(Bytes::from_static(b"value"))],
1007                    vec![1],
1008                )
1009                .unwrap();
1010
1011            assert_eq!(
1012                resources
1013                    .statement(b"client-statement")
1014                    .unwrap()
1015                    .upstream_name(),
1016                b"upstream-statement-42"
1017            );
1018            assert_eq!(
1019                resources.portal(b"client-portal").unwrap().upstream_name(),
1020                b"upstream-portal-42"
1021            );
1022
1023            resources.close_portal(portal).unwrap();
1024            resources.close_statement(statement).unwrap();
1025            assert!(resources.portal(b"client-portal").is_none());
1026            assert!(resources.statement(b"client-statement").is_none());
1027        });
1028    }
1029
1030    #[test]
1031    fn idle_readiness_invalidates_only_the_unnamed_portal() {
1032        let ready: Conn<(), crate::auth::Ready> = Conn::new(()).transition();
1033        with_connection_resources(ready.begin_extended(), |connection| {
1034            let (connection, statement, _) = connection
1035                .prepare(
1036                    Bytes::new(),
1037                    Bytes::new(),
1038                    Bytes::from_static(b"select 1"),
1039                    vec![],
1040                )
1041                .unwrap();
1042            let (connection, portal, _) = connection
1043                .bind(
1044                    &statement,
1045                    Bytes::new(),
1046                    Bytes::new(),
1047                    vec![],
1048                    vec![],
1049                    vec![],
1050                )
1051                .unwrap();
1052            assert!(connection.statement_is_live(&statement));
1053            assert!(connection.portal_is_live(&portal));
1054            let (awaiting, _) = connection.sync();
1055            let ResourceAwaitingTransition::Ready(ResourceReadyState::Clean(ready)) = awaiting
1056                .offer(SessionItem::ReadyForQuery {
1057                    status: TransactionStatus::Idle,
1058                    parameters_changed: false,
1059                })
1060            else {
1061                panic!("idle readiness should complete the extended cycle")
1062            };
1063
1064            assert!(ready.statement_is_live(&statement));
1065            assert!(!ready.portal_is_live(&portal));
1066            ready.into_connection().into_transport();
1067        });
1068    }
1069
1070    #[test]
1071    fn simple_query_invalidates_only_the_unnamed_statement() {
1072        let ready: Conn<(), crate::auth::Ready> = Conn::new(()).transition();
1073        with_connection_resources(ready.begin_extended(), |connection| {
1074            let (connection, unnamed, _) = connection
1075                .prepare(
1076                    Bytes::new(),
1077                    Bytes::new(),
1078                    Bytes::from_static(b"select 1"),
1079                    vec![],
1080                )
1081                .unwrap();
1082            let (connection, named, _) = connection
1083                .prepare(
1084                    Bytes::from_static(b"client_named"),
1085                    Bytes::from_static(b"proxy_named"),
1086                    Bytes::from_static(b"select 2"),
1087                    vec![],
1088                )
1089                .unwrap();
1090            let (awaiting, _) = connection.sync();
1091            let ResourceAwaitingTransition::Ready(ResourceReadyState::Clean(ready)) = awaiting
1092                .offer(SessionItem::ReadyForQuery {
1093                    status: TransactionStatus::Idle,
1094                    parameters_changed: false,
1095                })
1096            else {
1097                panic!("idle readiness should complete the extended cycle")
1098            };
1099            assert!(ready.statement_is_live(&unnamed));
1100            assert!(ready.statement_is_live(&named));
1101
1102            let (query, frame) = ready.query(b"select 3").unwrap();
1103            assert_eq!(frame.tag, b'Q');
1104            assert!(!query.statement_is_live(&unnamed));
1105            assert!(query.statement_is_live(&named));
1106            query.into_connection().into_transport();
1107        });
1108    }
1109
1110    #[test]
1111    fn branded_resources_rewrite_names_without_losing_bind_details() {
1112        with_resources(|mut resources| {
1113            let (statement, parse) = resources
1114                .prepare(
1115                    Bytes::from_static(b"client_statement"),
1116                    Bytes::from_static(b"proxy_7_statement"),
1117                    Bytes::from_static(b"select $1::int4"),
1118                    vec![23],
1119                )
1120                .unwrap();
1121            assert_eq!(statement.client_name(), b"client_statement");
1122            assert_eq!(parse.statement, Bytes::from_static(b"proxy_7_statement"));
1123
1124            let (portal, bind) = resources
1125                .bind(
1126                    &statement,
1127                    Bytes::from_static(b"client_portal"),
1128                    Bytes::from_static(b"proxy_7_portal"),
1129                    vec![1],
1130                    vec![Some(Bytes::from_static(b"\0\0\0*"))],
1131                    vec![1],
1132                )
1133                .unwrap();
1134            assert_eq!(bind.statement, Bytes::from_static(b"proxy_7_statement"));
1135            assert_eq!(bind.portal, Bytes::from_static(b"proxy_7_portal"));
1136            assert_eq!(bind.parameter_formats, [1]);
1137            assert_eq!(bind.result_formats, [1]);
1138            assert_eq!(portal.execute(0).portal, bind.portal);
1139
1140            resources.close_portal(portal).unwrap();
1141            resources.close_statement(statement).unwrap();
1142        });
1143    }
1144
1145    #[test]
1146    fn unnamed_resources_replace_the_previous_unnamed_resource() {
1147        with_resources(|mut resources| {
1148            let (obsolete, _) = resources
1149                .prepare(
1150                    Bytes::new(),
1151                    Bytes::new(),
1152                    Bytes::from_static(b"select 1"),
1153                    vec![],
1154                )
1155                .unwrap();
1156            let (replacement, _) = resources
1157                .prepare(
1158                    Bytes::new(),
1159                    Bytes::new(),
1160                    Bytes::from_static(b"select 2"),
1161                    vec![],
1162                )
1163                .expect("unnamed Parse replaces the prior unnamed statement");
1164            assert_eq!(
1165                resources
1166                    .bind(
1167                        &obsolete,
1168                        Bytes::new(),
1169                        Bytes::new(),
1170                        vec![],
1171                        vec![],
1172                        vec![],
1173                    )
1174                    .unwrap_err(),
1175                ResourceError::UnknownStatement
1176            );
1177
1178            let (_, _) = resources
1179                .bind(
1180                    &replacement,
1181                    Bytes::new(),
1182                    Bytes::new(),
1183                    vec![],
1184                    vec![],
1185                    vec![],
1186                )
1187                .unwrap();
1188            resources
1189                .bind(
1190                    &replacement,
1191                    Bytes::new(),
1192                    Bytes::new(),
1193                    vec![],
1194                    vec![],
1195                    vec![],
1196                )
1197                .expect("unnamed Bind replaces the prior unnamed portal");
1198        });
1199    }
1200
1201    #[test]
1202    fn protocol_boundaries_invalidate_unnamed_resource_tokens() {
1203        with_resources(|mut resources| {
1204            let (statement, _) = resources
1205                .prepare(
1206                    Bytes::new(),
1207                    Bytes::new(),
1208                    Bytes::from_static(b"select 1"),
1209                    vec![],
1210                )
1211                .unwrap();
1212            let (portal, _) = resources
1213                .bind(
1214                    &statement,
1215                    Bytes::new(),
1216                    Bytes::new(),
1217                    vec![],
1218                    vec![],
1219                    vec![],
1220                )
1221                .unwrap();
1222
1223            resources.simple_query_boundary();
1224            assert_eq!(
1225                resources
1226                    .bind(
1227                        &statement,
1228                        Bytes::from_static(b"p"),
1229                        Bytes::from_static(b"p"),
1230                        vec![],
1231                        vec![],
1232                        vec![],
1233                    )
1234                    .unwrap_err(),
1235                ResourceError::UnknownStatement
1236            );
1237
1238            resources.transaction_ended();
1239            assert_eq!(
1240                resources.close_portal(portal).unwrap_err(),
1241                ResourceError::UnknownPortal
1242            );
1243        });
1244    }
1245}
pg_proto/resources.rs

pg_proto/
resources.rs
