Skip to main content

tiberius/
client.rs

1mod auth;
2mod config;
3mod connection;
4
5mod tls;
6#[cfg(any(
7    feature = "rustls",
8    feature = "native-tls",
9    feature = "vendored-openssl"
10))]
11mod tls_stream;
12
13pub use auth::*;
14pub use config::*;
15pub(crate) use connection::*;
16
17use crate::tds::codec::RpcValue;
18use crate::tds::stream::ReceivedToken;
19use crate::{
20    result::ExecuteResult,
21    tds::{
22        codec::{self, IteratorJoin},
23        stream::{QueryStream, TokenStream},
24    },
25    BulkLoadRequest, ColumnFlag, MetaDataColumn, SqlReadBytes, ToSql,
26};
27use codec::{
28    BatchRequest, ColumnData, IsolationLevel, PacketHeader, RpcParam, RpcProcId, TokenRpcRequest,
29    TransactionManagerRequest,
30};
31use enumflags2::BitFlags;
32use futures_util::io::{AsyncRead, AsyncWrite};
33use futures_util::stream::TryStreamExt;
34use std::{borrow::Cow, fmt::Debug};
35
36/// `Client` is the main entry point to the SQL Server, providing query
37/// execution capabilities.
38///
39/// A `Client` is created using the [`Config`], defining the needed
40/// connection options and capabilities.
41///
42/// # Example
43///
44/// ```no_run
45/// # use tiberius::{Config, AuthMethod};
46/// use tokio_util::compat::TokioAsyncWriteCompatExt;
47///
48/// # #[tokio::main]
49/// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
50/// let mut config = Config::new();
51///
52/// config.host("0.0.0.0");
53/// config.port(1433);
54/// config.authentication(AuthMethod::sql_server("SA", "<Mys3cureP4ssW0rD>"));
55///
56/// let tcp = tokio::net::TcpStream::connect(config.get_addr()).await?;
57/// tcp.set_nodelay(true)?;
58/// // Client is ready to use.
59/// let client = tiberius::Client::connect(config, tcp.compat_write()).await?;
60/// # Ok(())
61/// # }
62/// ```
63///
64/// # Cancellation safety
65///
66/// A single [`Client`] drives one connection and one request at a time. If a
67/// `query`/`execute`/`simple_query` future — or the result stream it returns —
68/// is dropped before the request has been sent in full and the response fully
69/// consumed (for example under a `tokio::time::timeout` or a `select!` branch
70/// that loses the race), the connection may be left mid-message and out of sync
71/// with the server. A cancelled *write* is detected and any further use of that
72/// connection fails cleanly; a result stream dropped mid-response cannot be
73/// recovered. In both cases the safe course is to drop the `Client` and open a
74/// new connection (a connection pool should discard the connection on error)
75/// rather than reuse it.
76///
77/// [`Config`]: struct.Config.html
78#[derive(Debug)]
79pub struct Client<S: AsyncRead + AsyncWrite + Unpin + Send> {
80    pub(crate) connection: Connection<S>,
81}
82
83impl<S: AsyncRead + AsyncWrite + Unpin + Send> Client<S> {
84    /// Uses an instance of [`Config`] to specify the connection
85    /// options required to connect to the database using an established
86    /// tcp connection
87    ///
88    /// Note: `tcp_stream` is a connected stream, so some parts of the `Config`
89    /// (such as multi-subnet failover, which selects between resolved
90    /// addresses) must be handled while establishing that stream, outside of
91    /// this constructor.
92    ///
93    /// [`Config`]: struct.Config.html
94    pub async fn connect(config: Config, tcp_stream: S) -> crate::Result<Client<S>> {
95        Ok(Client {
96            connection: Connection::connect(config, tcp_stream).await?,
97        })
98    }
99
100    /// Executes SQL statements in the SQL Server, returning the number rows
101    /// affected. Useful for `INSERT`, `UPDATE` and `DELETE` statements. The
102    /// `query` can define the parameter placement by annotating them with
103    /// `@PN`, where N is the index of the parameter, starting from `1`. If
104    /// executing multiple queries at a time, delimit them with `;` and refer to
105    /// [`ExecuteResult`] how to get results for the separate queries.
106    ///
107    /// For mapping of Rust types when writing, see the documentation for
108    /// [`ToSql`]. For reading data from the database, see the documentation for
109    /// [`FromSql`].
110    ///
111    /// This API is not quite suitable for dynamic query parameters. In these
112    /// cases using a [`Query`] object might be easier.
113    ///
114    /// # Example
115    ///
116    /// ```no_run
117    /// # use tiberius::Config;
118    /// # use tokio_util::compat::TokioAsyncWriteCompatExt;
119    /// # use std::env;
120    /// # #[tokio::main]
121    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
122    /// # let c_str = env::var("TIBERIUS_TEST_CONNECTION_STRING").unwrap_or(
123    /// #     "server=tcp:localhost,1433;integratedSecurity=true;TrustServerCertificate=true".to_owned(),
124    /// # );
125    /// # let config = Config::from_ado_string(&c_str)?;
126    /// # let tcp = tokio::net::TcpStream::connect(config.get_addr()).await?;
127    /// # tcp.set_nodelay(true)?;
128    /// # let mut client = tiberius::Client::connect(config, tcp.compat_write()).await?;
129    /// let results = client
130    ///     .execute(
131    ///         "INSERT INTO ##Test (id) VALUES (@P1), (@P2), (@P3)",
132    ///         &[&1i32, &2i32, &3i32],
133    ///     )
134    ///     .await?;
135    /// # Ok(())
136    /// # }
137    /// ```
138    ///
139    /// [`ExecuteResult`]: struct.ExecuteResult.html
140    /// [`ToSql`]: trait.ToSql.html
141    /// [`FromSql`]: trait.FromSql.html
142    /// [`Query`]: struct.Query.html
143    pub async fn execute<'a>(
144        &mut self,
145        query: impl Into<Cow<'a, str>>,
146        params: &[&dyn ToSql],
147    ) -> crate::Result<ExecuteResult> {
148        self.connection.flush_stream().await?;
149        let rpc_params = Self::rpc_params(query);
150
151        let params = params.iter().map(|s| s.to_sql());
152        self.rpc_perform_query(RpcProcId::ExecuteSQL, rpc_params, params)
153            .await?;
154
155        ExecuteResult::new(&mut self.connection).await
156    }
157
158    /// Executes SQL statements in the SQL Server, returning resulting rows.
159    /// Useful for `SELECT` statements. The `query` can define the parameter
160    /// placement by annotating them with `@PN`, where N is the index of the
161    /// parameter, starting from `1`. If executing multiple queries at a time,
162    /// delimit them with `;` and refer to [`QueryStream`] on proper stream
163    /// handling.
164    ///
165    /// For mapping of Rust types when writing, see the documentation for
166    /// [`ToSql`]. For reading data from the database, see the documentation for
167    /// [`FromSql`].
168    ///
169    /// This API can be cumbersome for dynamic query parameters. In these cases,
170    /// if fighting too much with the compiler, using a [`Query`] object might be
171    /// easier.
172    ///
173    /// # Example
174    ///
175    /// ```
176    /// # use tiberius::Config;
177    /// # use tokio_util::compat::TokioAsyncWriteCompatExt;
178    /// # use std::env;
179    /// # #[tokio::main]
180    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
181    /// # let c_str = env::var("TIBERIUS_TEST_CONNECTION_STRING").unwrap_or(
182    /// #     "server=tcp:localhost,1433;integratedSecurity=true;TrustServerCertificate=true".to_owned(),
183    /// # );
184    /// # let config = Config::from_ado_string(&c_str)?;
185    /// # let tcp = tokio::net::TcpStream::connect(config.get_addr()).await?;
186    /// # tcp.set_nodelay(true)?;
187    /// # let mut client = tiberius::Client::connect(config, tcp.compat_write()).await?;
188    /// let stream = client
189    ///     .query(
190    ///         "SELECT @P1, @P2, @P3",
191    ///         &[&1i32, &2i32, &3i32],
192    ///     )
193    ///     .await?;
194    /// # Ok(())
195    /// # }
196    /// ```
197    ///
198    /// [`QueryStream`]: struct.QueryStream.html
199    /// [`Query`]: struct.Query.html
200    /// [`ToSql`]: trait.ToSql.html
201    /// [`FromSql`]: trait.FromSql.html
202    pub async fn query<'a, 'b>(
203        &'a mut self,
204        query: impl Into<Cow<'b, str>>,
205        params: &'b [&'b dyn ToSql],
206    ) -> crate::Result<QueryStream<'a>>
207    where
208        'a: 'b,
209    {
210        self.connection.flush_stream().await?;
211        let rpc_params = Self::rpc_params(query);
212
213        let params = params.iter().map(|p| p.to_sql());
214        self.rpc_perform_query(RpcProcId::ExecuteSQL, rpc_params, params)
215            .await?;
216
217        let ts = TokenStream::new(&mut self.connection);
218        let mut result = QueryStream::new(ts.try_unfold());
219        result.forward_to_metadata().await?;
220
221        Ok(result)
222    }
223
224    /// Execute multiple queries, delimited with `;` and return multiple result
225    /// sets; one for each query.
226    ///
227    /// # Example
228    ///
229    /// ```
230    /// # use tiberius::Config;
231    /// # use tokio_util::compat::TokioAsyncWriteCompatExt;
232    /// # use std::env;
233    /// # #[tokio::main]
234    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
235    /// # let c_str = env::var("TIBERIUS_TEST_CONNECTION_STRING").unwrap_or(
236    /// #     "server=tcp:localhost,1433;integratedSecurity=true;TrustServerCertificate=true".to_owned(),
237    /// # );
238    /// # let config = Config::from_ado_string(&c_str)?;
239    /// # let tcp = tokio::net::TcpStream::connect(config.get_addr()).await?;
240    /// # tcp.set_nodelay(true)?;
241    /// # let mut client = tiberius::Client::connect(config, tcp.compat_write()).await?;
242    /// let row = client.simple_query("SELECT 1 AS col").await?.into_row().await?.unwrap();
243    /// assert_eq!(Some(1i32), row.get("col"));
244    /// # Ok(())
245    /// # }
246    /// ```
247    ///
248    /// # Warning
249    ///
250    /// Do not use this with any user specified input. Please resort to prepared
251    /// statements using the [`query`] method.
252    ///
253    /// [`query`]: #method.query
254    pub async fn simple_query<'a, 'b>(
255        &'a mut self,
256        query: impl Into<Cow<'b, str>>,
257    ) -> crate::Result<QueryStream<'a>>
258    where
259        'a: 'b,
260    {
261        self.connection.flush_stream().await?;
262
263        let req = BatchRequest::new(query, self.connection.context().transaction_descriptor());
264
265        let id = self.connection.context_mut().next_packet_id();
266        self.connection.send(PacketHeader::batch(id), req).await?;
267
268        let ts = TokenStream::new(&mut self.connection);
269
270        let mut result = QueryStream::new(ts.try_unfold());
271        result.forward_to_metadata().await?;
272
273        Ok(result)
274    }
275
276    /// Execute a `BULK INSERT` statement, efficiently storing a large number of
277    /// rows to a specified table. Note: make sure the input row follows the same
278    /// schema as the table, otherwise calling `send()` will return an error.
279    ///
280    /// This is equivalent to calling `bulk_insert("table_name", &["*"])` to merge
281    /// all of a tables columns.
282    ///
283    /// # Example
284    ///
285    /// ```
286    /// # use tiberius::{Config, IntoRow};
287    /// # use tokio_util::compat::TokioAsyncWriteCompatExt;
288    /// # use std::env;
289    /// # #[tokio::main]
290    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
291    /// # let c_str = env::var("TIBERIUS_TEST_CONNECTION_STRING").unwrap_or(
292    /// #     "server=tcp:localhost,1433;integratedSecurity=true;TrustServerCertificate=true".to_owned(),
293    /// # );
294    /// # let config = Config::from_ado_string(&c_str)?;
295    /// # let tcp = tokio::net::TcpStream::connect(config.get_addr()).await?;
296    /// # tcp.set_nodelay(true)?;
297    /// # let mut client = tiberius::Client::connect(config, tcp.compat_write()).await?;
298    /// let create_table = r#"
299    ///     CREATE TABLE ##bulk_test (
300    ///         id INT IDENTITY PRIMARY KEY,
301    ///         val INT NOT NULL
302    ///     )
303    /// "#;
304    ///
305    /// client.simple_query(create_table).await?;
306    ///
307    /// // Start the bulk insert with the client.
308    /// let mut req = client.bulk_insert("##bulk_test").await?;
309    ///
310    /// for i in [0i32, 1i32, 2i32] {
311    ///     let row = (i).into_row();
312    ///
313    ///     // The request will handle flushing to the wire in an optimal way,
314    ///     // balancing between memory usage and IO performance.
315    ///     req.send(row).await?;
316    /// }
317    ///
318    /// // The request must be finalized.
319    /// let res = req.finalize().await?;
320    /// assert_eq!(3, res.total());
321    /// # Ok(())
322    /// # }
323    /// ```
324    pub async fn bulk_insert<'a>(
325        &'a mut self,
326        table: &'a str,
327    ) -> crate::Result<BulkLoadRequest<'a, S>> {
328        self.bulk_insert_columns(table, &["*"]).await
329    }
330
331    /// Execute a `BULK INSERT` statement, efficiently storing a large number of
332    /// rows to a specified table. Note: make sure the input row follows the same
333    /// schema as the column list, otherwise calling `send()` will return an error.
334    ///
335    /// # Example
336    ///
337    /// ```
338    /// # use tiberius::{Config, IntoRow};
339    /// # use tokio_util::compat::TokioAsyncWriteCompatExt;
340    /// # use std::env;
341    /// # #[tokio::main]
342    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
343    /// # let c_str = env::var("TIBERIUS_TEST_CONNECTION_STRING").unwrap_or(
344    /// #     "server=tcp:localhost,1433;integratedSecurity=true;TrustServerCertificate=true".to_owned(),
345    /// # );
346    /// # let config = Config::from_ado_string(&c_str)?;
347    /// # let tcp = tokio::net::TcpStream::connect(config.get_addr()).await?;
348    /// # tcp.set_nodelay(true)?;
349    /// # let mut client = tiberius::Client::connect(config, tcp.compat_write()).await?;
350    /// let create_table = r#"
351    ///     CREATE TABLE ##bulk_test_columns (
352    ///         id INT IDENTITY PRIMARY KEY,
353    ///         foo INT NOT NULL,
354    ///         bar FLOAT NOT NULL
355    ///     )
356    /// "#;
357    ///
358    /// client.simple_query(create_table).await?;
359    ///
360    /// // Start the bulk insert with the client.
361    /// let mut req = client.bulk_insert_columns("##bulk_test_columns", &["foo", "bar"]).await?;
362    ///
363    /// for (i, j) in [(0i32, 0f64), (1i32, 1f64), (2i32, 2f64)] {
364    ///     let row = (i, j).into_row();
365    ///
366    ///     // The request will handle flushing to the wire in an optimal way,
367    ///     // balancing between memory usage and IO performance.
368    ///     req.send(row).await?;
369    /// }
370    ///
371    /// // The request must be finalized.
372    /// let res = req.finalize().await?;
373    /// assert_eq!(3, res.total());
374    /// # Ok(())
375    /// # }
376    /// ```
377    pub async fn bulk_insert_columns<'a>(
378        &'a mut self,
379        table: &'a str,
380        columns: &'a [&'a str],
381    ) -> crate::Result<BulkLoadRequest<'a, S>> {
382        // Retrieve column metadata from the server, keeping only the updateable
383        // columns as bulk targets (identity/computed columns are skipped).
384        let columns: Vec<_> = self
385            .column_metadata(table, columns)
386            .await?
387            .into_iter()
388            .filter(|column| column.base.flags.contains(ColumnFlag::Updateable))
389            .collect();
390
391        // now start bulk upload
392        self.connection.flush_stream().await?;
393        let col_data = columns.iter().map(|c| format!("{}", c)).join(", ");
394        let query = format!("INSERT BULK {} ({})", table, col_data);
395
396        let req = BatchRequest::new(query, self.connection.context().transaction_descriptor());
397        let id = self.connection.context_mut().next_packet_id();
398
399        self.connection.send(PacketHeader::batch(id), req).await?;
400
401        let ts = TokenStream::new(&mut self.connection);
402        ts.flush_done().await?;
403
404        BulkLoadRequest::new(&mut self.connection, columns)
405    }
406
407    /// Retrieve the column metadata for a set of columns of a table, including
408    /// the column names, types (with their size, precision and scale) and flags
409    /// such as nullability and whether a column is an identity column.
410    ///
411    /// Pass `&["*"]` as `columns` to return the metadata for every column of the
412    /// table.
413    ///
414    /// ```no_run
415    /// # use tiberius::Config;
416    /// # use tokio_util::compat::TokioAsyncWriteCompatExt;
417    /// # use std::env;
418    /// # #[tokio::main]
419    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
420    /// # let c_str = env::var("TIBERIUS_TEST_CONNECTION_STRING").unwrap_or(
421    /// #     "server=tcp:localhost,1433;integratedSecurity=true;TrustServerCertificate=true".to_owned(),
422    /// # );
423    /// # let config = Config::from_ado_string(&c_str)?;
424    /// # let tcp = tokio::net::TcpStream::connect(config.get_addr()).await?;
425    /// # tcp.set_nodelay(true)?;
426    /// # let mut client = tiberius::Client::connect(config, tcp.compat_write()).await?;
427    /// let meta = client.column_metadata("some_table", &["*"]).await?;
428    /// assert!(meta[0].base().is_identity());
429    /// # Ok(())
430    /// # }
431    /// ```
432    pub async fn column_metadata(
433        &mut self,
434        table: &str,
435        columns: &[&str],
436    ) -> crate::Result<Vec<MetaDataColumn<'static>>> {
437        self.connection.flush_stream().await?;
438
439        // Ask the server for the column layout without returning any rows.
440        let columns = columns.join(", ");
441        let query = format!("SELECT TOP 0 {columns} FROM {table}");
442
443        let req = BatchRequest::new(query, self.connection.context().transaction_descriptor());
444        let id = self.connection.context_mut().next_packet_id();
445        self.connection.send(PacketHeader::batch(id), req).await?;
446
447        let token_stream = TokenStream::new(&mut self.connection).try_unfold();
448
449        let columns = token_stream
450            .try_fold(None, |mut columns, token| async move {
451                if let ReceivedToken::NewResultset(metadata) = token {
452                    columns = Some(metadata.columns.clone());
453                };
454
455                Ok(columns)
456            })
457            .await?;
458
459        let columns = columns.ok_or_else(|| {
460            crate::Error::Protocol("expecting column metadata from query but not found".into())
461        })?;
462
463        // Own the column names so the returned metadata is not tied to the
464        // lifetime of the token stream.
465        Ok(columns
466            .into_iter()
467            .map(|c| MetaDataColumn {
468                base: c.base,
469                col_name: std::borrow::Cow::Owned(c.col_name.into_owned()),
470            })
471            .collect())
472    }
473
474    /// Sends a TDS Attention signal to the server (packet type `0x06`,
475    /// MS-TDS section 2.2.1.6) to cancel the request that is currently in
476    /// flight on this connection, and drains the acknowledging token stream so
477    /// the connection can be reused for further queries.
478    ///
479    /// The server responds to the Attention signal by aborting the running
480    /// batch or RPC and returning a `DONE` token with the `DONE_ATTN` status
481    /// bit set. This method waits for that acknowledgement before returning,
482    /// discarding any remaining rows or tokens from the cancelled request.
483    ///
484    /// # Query cancellation and futures
485    ///
486    /// Dropping a [`query`], [`execute`] or [`simple_query`] future (for
487    /// example when a `tokio::time::timeout` elapses or a `select!` branch is
488    /// cancelled) stops the client from polling the stream, but it does *not*
489    /// tell the server to stop working on the request. To actually cancel the
490    /// in-flight work on the server, keep the [`Client`] and call
491    /// `cancel_query` on it. Because `cancel_query` borrows the client
492    /// mutably, it can only be issued once the borrowing result stream has
493    /// been dropped — typically from a separate task holding the client, or
494    /// after a cancelled/timed-out future has released its borrow.
495    ///
496    /// [`query`]: #method.query
497    /// [`execute`]: #method.execute
498    /// [`simple_query`]: #method.simple_query
499    pub async fn cancel_query(&mut self) -> crate::Result<()> {
500        self.connection.cancel_request().await?;
501        Ok(())
502    }
503
504    /// Closes this database connection explicitly.
505    pub async fn close(self) -> crate::Result<()> {
506        self.connection.close().await
507    }
508
509    /// Begins a new transaction using a Transaction Manager request
510    /// (`TM_BEGIN_XACT`, MS-TDS 2.2.6.8) instead of a `BEGIN TRAN` T-SQL
511    /// batch.
512    ///
513    /// On success the server replies with a `BeginTransaction` environment
514    /// change token whose descriptor is stored in the connection context and
515    /// automatically attached to subsequent requests, scoping them to the
516    /// transaction. Commit the work with [`commit_transaction`] or discard it
517    /// with [`rollback_transaction`].
518    ///
519    /// The transaction uses the server's default isolation level. Use
520    /// [`begin_transaction_with_isolation`] to request a specific one.
521    ///
522    /// [`commit_transaction`]: #method.commit_transaction
523    /// [`rollback_transaction`]: #method.rollback_transaction
524    /// [`begin_transaction_with_isolation`]: #method.begin_transaction_with_isolation
525    pub async fn begin_transaction(&mut self) -> crate::Result<()> {
526        self.begin_transaction_with_isolation(IsolationLevel::Unspecified)
527            .await
528    }
529
530    /// Begins a new transaction with an explicit isolation level using a
531    /// Transaction Manager request (`TM_BEGIN_XACT`, MS-TDS 2.2.6.8).
532    ///
533    /// See [`begin_transaction`] for details on transaction scoping.
534    ///
535    /// [`begin_transaction`]: #method.begin_transaction
536    pub async fn begin_transaction_with_isolation(
537        &mut self,
538        isolation_level: IsolationLevel,
539    ) -> crate::Result<()> {
540        let req = TransactionManagerRequest::begin(
541            self.connection.context().transaction_descriptor(),
542            isolation_level,
543            "",
544        );
545
546        self.send_transaction_manager_request(req).await
547    }
548
549    /// Commits the active transaction using a Transaction Manager request
550    /// (`TM_COMMIT_XACT`, MS-TDS 2.2.6.8).
551    ///
552    /// After a successful commit the connection is no longer scoped to a
553    /// transaction.
554    pub async fn commit_transaction(&mut self) -> crate::Result<()> {
555        let req = TransactionManagerRequest::commit(
556            self.connection.context().transaction_descriptor(),
557            "",
558        );
559
560        self.send_transaction_manager_request(req).await
561    }
562
563    /// Rolls back the active transaction using a Transaction Manager request
564    /// (`TM_ROLLBACK_XACT`, MS-TDS 2.2.6.8).
565    ///
566    /// After a successful rollback the connection is no longer scoped to a
567    /// transaction.
568    pub async fn rollback_transaction(&mut self) -> crate::Result<()> {
569        let req = TransactionManagerRequest::rollback(
570            self.connection.context().transaction_descriptor(),
571            "",
572        );
573
574        self.send_transaction_manager_request(req).await
575    }
576
577    /// Creates a named savepoint in the active transaction using a Transaction
578    /// Manager request (`TM_SAVE_XACT`, MS-TDS 2.2.6.8).
579    ///
580    /// The savepoint can later be targeted by a T-SQL `ROLLBACK TRANSACTION
581    /// <name>` to undo work performed after it while keeping the surrounding
582    /// transaction open.
583    pub async fn save_transaction<'a>(
584        &mut self,
585        name: impl Into<Cow<'a, str>>,
586    ) -> crate::Result<()> {
587        let req = TransactionManagerRequest::save(
588            self.connection.context().transaction_descriptor(),
589            name,
590        );
591
592        self.send_transaction_manager_request(req).await
593    }
594
595    async fn send_transaction_manager_request(
596        &mut self,
597        req: TransactionManagerRequest<'_>,
598    ) -> crate::Result<()> {
599        self.connection.flush_stream().await?;
600
601        let id = self.connection.context_mut().next_packet_id();
602        self.connection
603            .send(PacketHeader::transaction_manager(id), req)
604            .await?;
605
606        // The server responds with a DONE token (plus an ENVCHANGE token that
607        // the token stream applies to the connection context, updating the
608        // active transaction descriptor).
609        TokenStream::new(&mut self.connection).flush_done().await?;
610
611        Ok(())
612    }
613
614    pub(crate) fn rpc_params<'a>(query: impl Into<Cow<'a, str>>) -> Vec<RpcParam<'a>> {
615        vec![
616            RpcParam {
617                name: Cow::Borrowed("stmt"),
618                flags: BitFlags::empty(),
619                value: RpcValue::Scalar(ColumnData::String(Some(query.into()))),
620            },
621            RpcParam {
622                name: Cow::Borrowed("params"),
623                flags: BitFlags::empty(),
624                value: RpcValue::Scalar(ColumnData::I32(Some(0))),
625            },
626        ]
627    }
628
629    pub(crate) async fn rpc_perform_query<'a, 'b>(
630        &'a mut self,
631        proc_id: RpcProcId,
632        mut rpc_params: Vec<RpcParam<'b>>,
633        params: impl Iterator<Item = ColumnData<'b>>,
634    ) -> crate::Result<()>
635    where
636        'a: 'b,
637    {
638        let mut param_str = String::new();
639
640        for (i, param) in params.enumerate() {
641            if i > 0 {
642                param_str.push(',')
643            }
644            param_str.push_str(&format!("@P{} ", i + 1));
645            param_str.push_str(&param.type_name());
646
647            rpc_params.push(RpcParam {
648                name: Cow::Owned(format!("@P{}", i + 1)),
649                flags: BitFlags::empty(),
650                value: RpcValue::Scalar(param),
651            });
652        }
653
654        if let Some(params) = rpc_params.iter_mut().find(|x| x.name == "params") {
655            params.value = RpcValue::Scalar(ColumnData::String(Some(param_str.into())));
656        }
657
658        let req = TokenRpcRequest::new(
659            proc_id,
660            rpc_params,
661            self.connection.context().transaction_descriptor(),
662        );
663
664        let id = self.connection.context_mut().next_packet_id();
665        self.connection.send(PacketHeader::rpc(id), req).await?;
666
667        Ok(())
668    }
669
670    /// Sends a named-procedure RPC request with the given parameters. The caller
671    /// is responsible for flushing the connection beforehand and for consuming
672    /// the resulting token stream.
673    pub(crate) async fn rpc_run_command<'a, 'b>(
674        &'a mut self,
675        command_name: Cow<'b, str>,
676        rpc_params: Vec<RpcParam<'b>>,
677    ) -> crate::Result<()>
678    where
679        'a: 'b,
680    {
681        let req = TokenRpcRequest::new(
682            command_name,
683            rpc_params,
684            self.connection.context().transaction_descriptor(),
685        );
686
687        let id = self.connection.context_mut().next_packet_id();
688        self.connection.send(PacketHeader::rpc(id), req).await?;
689
690        Ok(())
691    }
692
693    /// Runs a batch query solely to retrieve its column metadata. Used to
694    /// resolve the column layout of a table-valued parameter type.
695    pub(crate) async fn query_run_for_metadata<'b>(
696        &mut self,
697        query: String,
698    ) -> crate::Result<Option<Vec<MetaDataColumn<'b>>>> {
699        self.connection.flush_stream().await?;
700
701        let req = BatchRequest::new(query, self.connection.context().transaction_descriptor());
702
703        let id = self.connection.context_mut().next_packet_id();
704        self.connection.send(PacketHeader::batch(id), req).await?;
705
706        let token_stream = TokenStream::new(&mut self.connection).try_unfold();
707
708        let columns = token_stream
709            .try_fold(None, |mut columns, token| async move {
710                if let ReceivedToken::NewResultset(metadata) = token {
711                    columns = Some(metadata.columns.clone());
712                };
713
714                Ok(columns)
715            })
716            .await?;
717
718        Ok(columns)
719    }
720}