From 2976dedd3c87543b116469eef436affb7c1e4429 Mon Sep 17 00:00:00 2001 From: Matthew Jackson <1085847+MattJackson@users.noreply.github.com> Date: Sat, 29 Aug 2026 14:07:55 -0700 Subject: [PATCH] feat(rpc): named-procedure RPC with OUT params and table-valued parameters Implement calling a stored procedure by name, resolving the `todo!()` that previously blocked named-procedure RPC requests. - Encode `RpcProcIdValue::Name` as a US_VARCHAR proc name in the RPC request, alongside the existing by-id path used by execute/query. - Introduce `RpcValue` (Scalar/Table) so RPC params can carry either a scalar `ColumnData` or a table-valued parameter, and thread it through the existing execute/query call sites. - Add `TypeInfoTvp` to encode TVP type info and rows (MS-TDS 2.2.5.5.5), rewriting fixed-length column types to their nullable var-len variants. - Add a `Command` public API (`bind_param`, `bind_out_param`, `bind_table`, `bind_table_with_dbtype`, `exec`) plus `CommandStream`, `CommandItem`, `CommandResult` and `CommandReturnValue` for reading OUT parameter values, return codes and result sets. - Add the `tvp-macro` crate providing `#[derive(TableValueRow)]`. - Encode nullable scalar values correctly when no TypeInfo is supplied. Unit tests cover named-proc request encoding, TVP type-info encoding and the derive macro output. Integration tests requiring a live SQL Server are added under tests/command.rs. Resolves #275. --- Cargo.toml | 9 +- src/client.rs | 60 ++++- src/command.rs | 296 +++++++++++++++++++++++++ src/lib.rs | 4 +- src/result.rs | 101 ++++++++- src/tds/codec.rs | 2 + src/tds/codec/column_data.rs | 96 +++++--- src/tds/codec/rpc_request.rs | 132 ++++++++++- src/tds/codec/type_info_tvp.rs | 206 ++++++++++++++++++ src/tds/stream.rs | 2 + src/tds/stream/command.rs | 337 +++++++++++++++++++++++++++++ src/tds/stream/query.rs | 4 +- tests/command.rs | 229 ++++++++++++++++++++ tvp-macro/Cargo.toml | 15 ++ tvp-macro/src/attr.rs | 45 ++++ tvp-macro/src/lib.rs | 46 ++++ tvp-macro/src/table_value_param.rs | 167 ++++++++++++++ 17 files changed, 1705 insertions(+), 46 deletions(-) create mode 100644 src/command.rs create mode 100644 src/tds/codec/type_info_tvp.rs create mode 100644 src/tds/stream/command.rs create mode 100644 tests/command.rs create mode 100644 tvp-macro/Cargo.toml create mode 100644 tvp-macro/src/attr.rs create mode 100644 tvp-macro/src/lib.rs create mode 100644 tvp-macro/src/table_value_param.rs diff --git a/Cargo.toml b/Cargo.toml index 0caaac8..bc37b71 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,12 +16,16 @@ repository = "https://github.com/prisma/tiberius" version = "0.12.3" [workspace] -members = ["runtimes-macro"] +members = ["runtimes-macro", "tvp-macro"] [[test]] path = "tests/query.rs" name = "query" +[[test]] +path = "tests/command.rs" +name = "command" + [[test]] path = "tests/named-instance-async.rs" name = "named-instance-async" @@ -162,6 +166,9 @@ version = "1" [dev-dependencies.runtimes-macro] path = "./runtimes-macro" +[dependencies.tvp-macro] +path = "./tvp-macro" + [dev-dependencies] names = "0.14" anyhow = "1" diff --git a/src/client.rs b/src/client.rs index 688721d..d44088a 100644 --- a/src/client.rs +++ b/src/client.rs @@ -14,6 +14,7 @@ pub use auth::*; pub use config::*; pub(crate) use connection::*; +use crate::tds::codec::{MetaDataColumn, RpcValue}; use crate::tds::stream::ReceivedToken; use crate::{ result::ExecuteResult, @@ -357,12 +358,12 @@ impl Client { RpcParam { name: Cow::Borrowed("stmt"), flags: BitFlags::empty(), - value: ColumnData::String(Some(query.into())), + value: RpcValue::Scalar(ColumnData::String(Some(query.into()))), }, RpcParam { name: Cow::Borrowed("params"), flags: BitFlags::empty(), - value: ColumnData::I32(Some(0)), + value: RpcValue::Scalar(ColumnData::I32(Some(0))), }, ] } @@ -388,12 +389,12 @@ impl Client { rpc_params.push(RpcParam { name: Cow::Owned(format!("@P{}", i + 1)), flags: BitFlags::empty(), - value: param, + value: RpcValue::Scalar(param), }); } if let Some(params) = rpc_params.iter_mut().find(|x| x.name == "params") { - params.value = ColumnData::String(Some(param_str.into())); + params.value = RpcValue::Scalar(ColumnData::String(Some(param_str.into()))); } let req = TokenRpcRequest::new( @@ -407,4 +408,55 @@ impl Client { Ok(()) } + + /// Sends a named-procedure RPC request with the given parameters. The caller + /// is responsible for flushing the connection beforehand and for consuming + /// the resulting token stream. + pub(crate) async fn rpc_run_command<'a, 'b>( + &'a mut self, + command_name: Cow<'b, str>, + rpc_params: Vec>, + ) -> crate::Result<()> + where + 'a: 'b, + { + let req = TokenRpcRequest::new( + command_name, + rpc_params, + self.connection.context().transaction_descriptor(), + ); + + let id = self.connection.context_mut().next_packet_id(); + self.connection.send(PacketHeader::rpc(id), req).await?; + + Ok(()) + } + + /// Runs a batch query solely to retrieve its column metadata. Used to + /// resolve the column layout of a table-valued parameter type. + pub(crate) async fn query_run_for_metadata<'b>( + &mut self, + query: String, + ) -> crate::Result>>> { + self.connection.flush_stream().await?; + + let req = BatchRequest::new(query, self.connection.context().transaction_descriptor()); + + let id = self.connection.context_mut().next_packet_id(); + self.connection.send(PacketHeader::batch(id), req).await?; + + let token_stream = TokenStream::new(&mut self.connection).try_unfold(); + + let columns = token_stream + .try_fold(None, |mut columns, token| async move { + if let ReceivedToken::NewResultset(metadata) = token { + columns = Some(metadata.columns.clone()); + }; + + Ok(columns) + }) + .await?; + + Ok(columns) + } } diff --git a/src/command.rs b/src/command.rs new file mode 100644 index 0000000..34340dc --- /dev/null +++ b/src/command.rs @@ -0,0 +1,296 @@ +use std::borrow::Cow; + +use enumflags2::BitFlags; +use futures_util::io::{AsyncRead, AsyncWrite}; + +use crate::{ + tds::{ + codec::{RpcParam, RpcStatus::ByRefValue, RpcValue, TypeInfoTvp}, + stream::{CommandStream, TokenStream}, + }, + Client, ColumnData, IntoSql, +}; + +#[doc(inline)] +pub use tvp_macro::TableValueRow; + +/// A structure that represents a single row of a table-valued parameter (TVP) +/// implements this trait. +/// +/// It can be derived with `#[derive(TableValueRow)]` for structs with named +/// fields. +pub trait TableValueRow<'a> { + /// Binds this row's field values. Called by [`Command`] before making the + /// call to the server; implementations must call + /// [`SqlTableDataRow::add_field`] once per column, in column order. + fn bind_fields(&self, data_row: &mut SqlTableDataRow<'a>); + /// The database type name that represents this TVP, e.g. `dbo.MyType`. + fn get_db_type() -> &'static str; +} + +/// A collection of [`TableValueRow`] values that can be bound as a +/// table-valued parameter. Implemented for any `IntoIterator` of rows. +pub trait TableValue<'a> { + /// Converts this collection into the internal table data representation. + fn into_sql(self) -> SqlTableData<'a>; +} + +impl<'a, R, C> TableValue<'a> for C +where + R: TableValueRow<'a> + 'a, + C: IntoIterator, +{ + fn into_sql(self) -> SqlTableData<'a> { + let mut data = Vec::new(); + for row in self.into_iter() { + let mut data_row = SqlTableDataRow::new(); + row.bind_fields(&mut data_row); + data.push(data_row); + } + + SqlTableData { + rows: data, + db_type: R::get_db_type(), + } + } +} + +/// A remote command (stored procedure or user-defined function) with bound +/// parameters, executed by name via an RPC request. +#[derive(Debug)] +pub struct Command<'a> { + name: Cow<'a, str>, + // The server rejects repeated parameter names, so uniqueness is not checked here. + params: Vec>, +} + +#[derive(Debug)] +struct CommandParam<'a> { + name: Cow<'a, str>, + out: bool, + data: CommandParamData<'a>, +} + +#[derive(Debug)] +enum CommandParamData<'a> { + Scalar(ColumnData<'a>), + Table(SqlTableData<'a>), +} + +/// The internal representation of a table-valued parameter's data. +#[derive(Debug)] +pub struct SqlTableData<'a> { + rows: Vec>, + db_type: &'a str, +} + +/// A single row of a table-valued parameter, used by [`TableValueRow`] +/// implementations to bind column values. +#[derive(Debug)] +pub struct SqlTableDataRow<'a> { + col_data: Vec>, +} + +impl<'a> SqlTableDataRow<'a> { + fn new() -> SqlTableDataRow<'a> { + SqlTableDataRow { + col_data: Vec::new(), + } + } + + /// Adds a field value to this TVP row. Must be called once per column; the + /// values are sent to the server in call order. + pub fn add_field(&mut self, data: impl IntoSql<'a> + 'a) { + self.col_data.push(data.into_sql()); + } +} + +impl<'a> Command<'a> { + /// Constructs a new command with the given procedure or function name. + pub fn new(proc_name: impl Into>) -> Self { + Self { + name: proc_name.into(), + params: Vec::new(), + } + } + + /// Binds a scalar input parameter with the given name. + pub fn bind_param(&mut self, name: impl Into>, data: impl IntoSql<'a> + 'a) { + self.params.push(CommandParam { + name: name.into(), + out: false, + data: CommandParamData::Scalar(data.into_sql()), + }); + } + + /// Binds a by-ref (OUT) scalar parameter. The returned value can be found by + /// the same name in the [`CommandResult`] returned values. + /// + /// [`CommandResult`]: crate::CommandResult + pub fn bind_out_param(&mut self, name: impl Into>, data: impl IntoSql<'a> + 'a) { + self.params.push(CommandParam { + name: name.into(), + out: true, + data: CommandParamData::Scalar(data.into_sql()), + }); + } + + /// Binds a table-valued parameter. The provided argument must implement + /// [`TableValue`]. + /// + /// # Example + /// + /// ```no_run + /// # use std::env; + /// # use tiberius::Config; + /// # use tiberius::{numeric::Numeric, Command, TableValueRow}; + /// # use tokio_util::compat::TokioAsyncWriteCompatExt; + /// #[derive(TableValueRow)] + /// struct SomeGeoList { + /// eid: i32, + /// lat: Numeric, + /// lon: Numeric, + /// } + /// # #[tokio::main] + /// # async fn main() -> Result<(), Box> { + /// # let c_str = env::var("TIBERIUS_TEST_CONNECTION_STRING").unwrap_or( + /// # "server=tcp:localhost,1433;integratedSecurity=true;TrustServerCertificate=true".to_owned(), + /// # ); + /// # let config = Config::from_ado_string(&c_str)?; + /// # let tcp = tokio::net::TcpStream::connect(config.get_addr()).await?; + /// # tcp.set_nodelay(true)?; + /// # let client = tiberius::Client::connect(config, tcp.compat_write()).await?; + /// let r1 = SomeGeoList { + /// eid: 1, + /// lat: Numeric::new_with_scale(10, 6), + /// lon: Numeric::new_with_scale(14, 6), + /// }; + /// let r2 = SomeGeoList { + /// eid: 4, + /// lat: Numeric::new_with_scale(101, 6), + /// lon: Numeric::new_with_scale(142, 6), + /// }; + /// + /// let tbl = vec![r1, r2]; + /// + /// let mut cmd = Command::new("dbo.usp_TheGeoProcedure"); + /// cmd.bind_table("@table", tbl); + /// # Ok(()) + /// # } + /// ``` + pub fn bind_table(&mut self, name: impl Into>, data: impl TableValue<'a> + 'a) { + self.params.push(CommandParam { + name: name.into(), + out: false, + data: CommandParamData::Table(data.into_sql()), + }); + } + + /// The same as [`bind_table`](Self::bind_table), but overrides the database + /// type name used for the TVP. + pub fn bind_table_with_dbtype( + &mut self, + name: impl Into>, + db_type: &'a str, + data: impl TableValue<'a> + 'a, + ) { + self.params.push(CommandParam { + name: name.into(), + out: false, + data: CommandParamData::Table(SqlTableData { + db_type, + ..data.into_sql() + }), + }); + } + + /// Executes the command on the server, returning a [`CommandStream`] that + /// can be collected into a [`CommandResult`] for convenience. + /// + /// [`CommandResult`]: crate::CommandResult + /// + /// # Example + /// + /// ```no_run + /// # use tiberius::{Config, Command}; + /// # use tokio_util::compat::TokioAsyncWriteCompatExt; + /// # use std::env; + /// # #[tokio::main] + /// # async fn main() -> Result<(), Box> { + /// # let c_str = env::var("TIBERIUS_TEST_CONNECTION_STRING").unwrap_or( + /// # "server=tcp:localhost,1433;integratedSecurity=true;TrustServerCertificate=true".to_owned(), + /// # ); + /// # let config = Config::from_ado_string(&c_str)?; + /// # let tcp = tokio::net::TcpStream::connect(config.get_addr()).await?; + /// # tcp.set_nodelay(true)?; + /// # let mut client = tiberius::Client::connect(config, tcp.compat_write()).await?; + /// let mut cmd = Command::new("dbo.usp_SomeStoredProc"); + /// + /// cmd.bind_param("@foo", 34i32); + /// cmd.bind_out_param("@bar", "bar"); + /// let res = cmd.exec(&mut client).await?.into_command_result().await?; + /// + /// let rv: Option<&str> = res.try_return_value("@bar")?; + /// let rc = res.return_code(); + /// # Ok(()) + /// # } + /// ``` + pub async fn exec<'b, S>(self, client: &'b mut Client) -> crate::Result> + where + S: AsyncRead + AsyncWrite + Unpin + Send, + { + let rpc_params = Command::build_rpc_params(self.params, client).await?; + + client.connection.flush_stream().await?; + client.rpc_run_command(self.name, rpc_params).await?; + + let ts = TokenStream::new(&mut client.connection); + let result = CommandStream::new(ts.try_unfold()); + + Ok(result) + } + + async fn build_rpc_params<'b, S>( + cmd_params: Vec>, + client: &'b mut Client, + ) -> crate::Result>> + where + S: AsyncRead + AsyncWrite + Unpin + Send, + { + let mut rpc_params = Vec::new(); + for p in cmd_params.into_iter() { + let rpc_val = match p.data { + CommandParamData::Scalar(col) => RpcValue::Scalar(col), + CommandParamData::Table(t) => { + let type_info_tvp = TypeInfoTvp::new( + t.db_type, + t.rows.into_iter().map(|r| r.col_data).collect(), + ); + // Resolve the TVP column layout from the server. + let cols_metadata = client + .query_run_for_metadata(format!( + "DECLARE @P AS {};SELECT TOP 0 * FROM @P", + t.db_type + )) + .await?; + RpcValue::Table(if let Some(cm) = cols_metadata { + type_info_tvp.with_metadata(cm) + } else { + type_info_tvp + }) + } + }; + let rpc_param = RpcParam { + name: p.name, + flags: if p.out { + BitFlags::from_flag(ByRefValue) + } else { + BitFlags::empty() + }, + value: rpc_val, + }; + rpc_params.push(rpc_param); + } + Ok(rpc_params) + } +} diff --git a/src/lib.rs b/src/lib.rs index 882f5ad..2ea83b9 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -257,6 +257,7 @@ pub(crate) extern crate bigdecimal_ as bigdecimal; mod macros; mod client; +mod command; mod from_sql; mod query; mod sql_read_bytes; @@ -270,6 +271,7 @@ mod tds; mod sql_browser; pub use client::{AuthMethod, Client, Config}; +pub use command::{Command, SqlTableData, SqlTableDataRow, TableValue, TableValueRow}; pub(crate) use error::Error; pub use from_sql::{FromSql, FromSqlOwned}; pub use query::Query; @@ -279,7 +281,7 @@ pub use sql_browser::SqlBrowser; pub use tds::{ codec::{BulkLoadRequest, ColumnData, ColumnFlag, IntoRow, TokenRow, TypeLength}, numeric, - stream::QueryStream, + stream::{CommandReturnValue, CommandStream, QueryStream}, time, xml, EncryptionLevel, }; pub use to_sql::{IntoSql, ToSql}; diff --git a/src/result.rs b/src/result.rs index 19ba6fa..95d2beb 100644 --- a/src/result.rs +++ b/src/result.rs @@ -1,7 +1,9 @@ -pub use crate::tds::stream::{QueryItem, ResultMetadata}; +pub use crate::tds::stream::{CommandItem, QueryItem, ResultMetadata}; use crate::{ client::Connection, - tds::stream::{ReceivedToken, TokenStream}, + error::Error, + tds::stream::{CommandReturnValue, ReceivedToken, TokenStream}, + FromSql, Row, }; use futures_util::io::{AsyncRead, AsyncWrite}; use futures_util::stream::TryStreamExt; @@ -113,3 +115,98 @@ impl IntoIterator for ExecuteResult { self.rows_affected.into_iter() } } + +/// A materialized result from executing a [`Command`], carrying the number of +/// affected rows, the return code, the values of any OUT parameters and any +/// record sets returned by the command. +/// +/// [`Command`]: crate::Command +/// +/// # Example +/// +/// ```no_run +/// # use tiberius::{Config, Command}; +/// # use tokio_util::compat::TokioAsyncWriteCompatExt; +/// # use std::env; +/// # #[tokio::main] +/// # async fn main() -> Result<(), Box> { +/// # let c_str = env::var("TIBERIUS_TEST_CONNECTION_STRING").unwrap_or( +/// # "server=tcp:localhost,1433;integratedSecurity=true;TrustServerCertificate=true".to_owned(), +/// # ); +/// # let config = Config::from_ado_string(&c_str)?; +/// # let tcp = tokio::net::TcpStream::connect(config.get_addr()).await?; +/// # tcp.set_nodelay(true)?; +/// # let mut client = tiberius::Client::connect(config, tcp.compat_write()).await?; +/// let mut cmd = Command::new("dbo.usp_SomeStoredProc"); +/// +/// cmd.bind_param("@foo", 34i32); +/// cmd.bind_out_param("@bar", "bar"); +/// let res = cmd.exec(&mut client).await?.into_command_result().await?; +/// +/// let rv: Option<&str> = res.try_return_value("@bar")?; +/// let rc = res.return_code(); +/// let ra = res.rows_affected(); +/// +/// let rs0 = res.to_query_result(0); +/// # Ok(()) +/// # } +/// ``` +/// +#[derive(Debug)] +pub struct CommandResult { + pub(crate) rows_affected: Vec, + pub(crate) return_code: u32, + pub(crate) return_values: Vec, + pub(crate) query_results: Vec>, +} + +impl<'a> CommandResult { + /// A slice of the numbers of rows affected, in the same order as the + /// statements ran by the command. + pub fn rows_affected(&self) -> &[u64] { + self.rows_affected.as_slice() + } + + /// The return code of the command, as returned by the server. + pub fn return_code(&self) -> u32 { + self.return_code + } + + /// The number of returned values (OUT parameters) available. + pub fn return_values_len(&self) -> usize { + self.return_values.len() + } + + /// Gets a returned value by its OUT parameter name, converting it to `T`. + /// Returns `None` if the value is `NULL`, and an error if no OUT parameter + /// with the given name was returned. + pub fn try_return_value(&'a self, name: &str) -> crate::Result> + where + T: FromSql<'a>, + { + let col_data = self + .return_values + .iter() + .find(|p| p.name.eq(name)) + .ok_or_else(|| { + Error::Conversion(format!("Could not find return value {}", name).into()) + })?; + + T::from_sql(&col_data.data) + } + + /// Gets a returned record set by its zero-based index. Returns `None` if the + /// index is out of range. + pub fn to_query_result(&self, idx: usize) -> Option<&Vec> { + self.query_results.get(idx) + } +} + +impl IntoIterator for CommandResult { + type Item = Vec; + type IntoIter = std::vec::IntoIter; + + fn into_iter(self) -> Self::IntoIter { + self.query_results.into_iter() + } +} diff --git a/src/tds/codec.rs b/src/tds/codec.rs index 07f1331..ee0b863 100644 --- a/src/tds/codec.rs +++ b/src/tds/codec.rs @@ -12,6 +12,7 @@ mod pre_login; mod rpc_request; mod token; mod type_info; +mod type_info_tvp; pub use batch_request::*; pub use bulk_load::*; @@ -28,6 +29,7 @@ pub use pre_login::*; pub use rpc_request::*; pub use token::*; pub use type_info::*; +pub use type_info_tvp::*; const HEADER_BYTES: usize = 8; const ALL_HEADERS_LEN_TX: usize = 22; diff --git a/src/tds/codec/column_data.rs b/src/tds/codec/column_data.rs index fecd83f..df86500 100644 --- a/src/tds/codec/column_data.rs +++ b/src/tds/codec/column_data.rs @@ -158,13 +158,18 @@ impl<'a> Encode> for ColumnData<'a> { (ColumnData::Bit(Some(val)), Some(TypeInfo::FixedLen(FixedLenType::Bit))) => { dst.put_u8(val as u8); } - (ColumnData::Bit(Some(val)), None) => { + (ColumnData::Bit(opt), None) => { // if TypeInfo was not given, encode a TypeInfo // the first 1 is part of TYPE_INFO - // the second 1 is part of TYPE_VARBYTE - let header = [VarLenType::Bitn as u8, 1, 1]; + let header = [VarLenType::Bitn as u8, 1]; dst.extend_from_slice(&header); - dst.put_u8(val as u8); + if let Some(val) = opt { + // the second 1 is part of TYPE_VARBYTE + dst.put_u8(1); + dst.put_u8(val as u8); + } else { + dst.put_u8(0); + } } (ColumnData::U8(opt), Some(TypeInfo::VarLenSized(vlc))) if vlc.r#type() == VarLenType::Intn => @@ -179,10 +184,15 @@ impl<'a> Encode> for ColumnData<'a> { (ColumnData::U8(Some(val)), Some(TypeInfo::FixedLen(FixedLenType::Int1))) => { dst.put_u8(val); } - (ColumnData::U8(Some(val)), None) => { - let header = [VarLenType::Intn as u8, 1, 1]; + (ColumnData::U8(opt), None) => { + let header = [VarLenType::Intn as u8, 1]; dst.extend_from_slice(&header); - dst.put_u8(val); + if let Some(val) = opt { + dst.put_u8(1); + dst.put_u8(val); + } else { + dst.put_u8(0); + } } (ColumnData::I16(Some(val)), Some(TypeInfo::FixedLen(FixedLenType::Int2))) => { dst.put_i16_le(val); @@ -197,11 +207,15 @@ impl<'a> Encode> for ColumnData<'a> { dst.put_u8(0); } } - (ColumnData::I16(Some(val)), None) => { - let header = [VarLenType::Intn as u8, 2, 2]; + (ColumnData::I16(opt), None) => { + let header = [VarLenType::Intn as u8, 2]; dst.extend_from_slice(&header); - - dst.put_i16_le(val); + if let Some(val) = opt { + dst.put_u8(2); + dst.put_i16_le(val); + } else { + dst.put_u8(0); + } } (ColumnData::I32(Some(val)), Some(TypeInfo::FixedLen(FixedLenType::Int4))) => { dst.put_i32_le(val); @@ -216,10 +230,15 @@ impl<'a> Encode> for ColumnData<'a> { dst.put_u8(0); } } - (ColumnData::I32(Some(val)), None) => { - let header = [VarLenType::Intn as u8, 4, 4]; + (ColumnData::I32(opt), None) => { + let header = [VarLenType::Intn as u8, 4]; dst.extend_from_slice(&header); - dst.put_i32_le(val); + if let Some(val) = opt { + dst.put_u8(4); + dst.put_i32_le(val); + } else { + dst.put_u8(0); + } } (ColumnData::I64(Some(val)), Some(TypeInfo::FixedLen(FixedLenType::Int8))) => { dst.put_i64_le(val); @@ -234,10 +253,15 @@ impl<'a> Encode> for ColumnData<'a> { dst.put_u8(0); } } - (ColumnData::I64(Some(val)), None) => { - let header = [VarLenType::Intn as u8, 8, 8]; + (ColumnData::I64(opt), None) => { + let header = [VarLenType::Intn as u8, 8]; dst.extend_from_slice(&header); - dst.put_i64_le(val); + if let Some(val) = opt { + dst.put_u8(8); + dst.put_i64_le(val); + } else { + dst.put_u8(0); + } } (ColumnData::F32(Some(val)), Some(TypeInfo::FixedLen(FixedLenType::Float4))) => { dst.put_f32_le(val); @@ -252,10 +276,15 @@ impl<'a> Encode> for ColumnData<'a> { dst.put_u8(0); } } - (ColumnData::F32(Some(val)), None) => { - let header = [VarLenType::Floatn as u8, 4, 4]; + (ColumnData::F32(opt), None) => { + let header = [VarLenType::Floatn as u8, 4]; dst.extend_from_slice(&header); - dst.put_f32_le(val); + if let Some(val) = opt { + dst.put_u8(4); + dst.put_f32_le(val); + } else { + dst.put_u8(0); + } } (ColumnData::F64(Some(val)), Some(TypeInfo::FixedLen(FixedLenType::Float8))) => { dst.put_f64_le(val); @@ -270,10 +299,15 @@ impl<'a> Encode> for ColumnData<'a> { dst.put_u8(0); } } - (ColumnData::F64(Some(val)), None) => { - let header = [VarLenType::Floatn as u8, 8, 8]; + (ColumnData::F64(opt), None) => { + let header = [VarLenType::Floatn as u8, 8]; dst.extend_from_slice(&header); - dst.put_f64_le(val); + if let Some(val) = opt { + dst.put_u8(8); + dst.put_f64_le(val); + } else { + dst.put_u8(0); + } } (ColumnData::Guid(opt), Some(TypeInfo::VarLenSized(vlc))) if vlc.r#type() == VarLenType::Guid => @@ -288,13 +322,17 @@ impl<'a> Encode> for ColumnData<'a> { dst.put_u8(0); } } - (ColumnData::Guid(Some(uuid)), None) => { - let header = [VarLenType::Guid as u8, 16, 16]; + (ColumnData::Guid(opt), None) => { + let header = [VarLenType::Guid as u8, 16]; dst.extend_from_slice(&header); - - let mut data = *uuid.as_bytes(); - super::guid::reorder_bytes(&mut data); - dst.extend_from_slice(&data); + if let Some(uuid) = opt { + dst.put_u8(16); + let mut data = *uuid.as_bytes(); + super::guid::reorder_bytes(&mut data); + dst.extend_from_slice(&data); + } else { + dst.put_u8(0); + } } (ColumnData::String(opt), Some(TypeInfo::VarLenSized(vlc))) if vlc.r#type() == VarLenType::BigChar diff --git a/src/tds/codec/rpc_request.rs b/src/tds/codec/rpc_request.rs index 368cbd3..97da34e 100644 --- a/src/tds/codec/rpc_request.rs +++ b/src/tds/codec/rpc_request.rs @@ -1,3 +1,4 @@ +use super::TypeInfoTvp; use super::{AllHeaderTy, Encode, ALL_HEADERS_LEN_TX}; use crate::{tds::codec::ColumnData, BytesMutWithTypeInfo, Result}; use bytes::{BufMut, BytesMut}; @@ -46,11 +47,22 @@ impl<'a> TokenRpcRequest<'a> { } } +/// The value carried by an [`RpcParam`]. A scalar column value, or a +/// table-valued parameter (TVP). +#[derive(Debug)] +pub enum RpcValue<'a> { + /// An ordinary scalar parameter value. + Scalar(ColumnData<'a>), + /// A table-valued parameter. As per the TDS grammar, `TYPE_INFO_TVP` + /// carries both the type metadata and the data rows. + Table(TypeInfoTvp<'a>), +} + #[derive(Debug)] pub struct RpcParam<'a> { pub name: Cow<'a, str>, pub flags: BitFlags, - pub value: ColumnData<'a>, + pub value: RpcValue<'a>, } /// 2.2.6.6 RPC Request @@ -103,10 +115,21 @@ impl<'a> Encode for TokenRpcRequest<'a> { let val = (0xffff_u32) | ((*id as u16) as u32) << 16; dst.put_u32_le(val); } - RpcProcIdValue::Name(ref _name) => { - //let (left_bytes, _) = try!(write_varchar::(&mut cursor, name, 0)); - //assert_eq!(left_bytes, 0); - todo!() + RpcProcIdValue::Name(ref name) => { + // ProcName is a US_VARCHAR: a u16 little-endian character count + // followed by that many UTF-16 code units. + let len_pos = dst.len(); + dst.put_u16_le(0u16); + let mut length = 0_u16; + + for chr in name.encode_utf16() { + dst.put_u16_le(chr); + length += 1; + } + + let dst: &mut [u8] = dst.borrow_mut(); + let mut dst = &mut dst[len_pos..]; + dst.put_u16_le(length); } } @@ -134,8 +157,13 @@ impl<'a> Encode for RpcParam<'a> { dst.put_u8(self.flags.bits()); - let mut dst_fi = BytesMutWithTypeInfo::new(dst); - self.value.encode(&mut dst_fi)?; + match self.value { + RpcValue::Scalar(value) => { + let mut dst_ti = BytesMutWithTypeInfo::new(dst); + value.encode(&mut dst_ti)?; + } + RpcValue::Table(value) => value.encode(dst)?, + } let dst: &mut [u8] = dst.borrow_mut(); dst[len_pos] = length; @@ -143,3 +171,93 @@ impl<'a> Encode for RpcParam<'a> { Ok(()) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::tds::codec::ColumnData; + + fn scalar(value: ColumnData<'static>) -> RpcValue<'static> { + RpcValue::Scalar(value) + } + + #[test] + fn encodes_named_proc_header() { + let req = TokenRpcRequest::new( + "dbo.usp_MyProc", + vec![RpcParam { + name: Cow::Borrowed("@id"), + flags: BitFlags::empty(), + value: scalar(ColumnData::I32(Some(1))), + }], + [0u8; 8], + ); + + let mut buf = BytesMut::new(); + req.encode(&mut buf).unwrap(); + + // Skip the ALL_HEADERS block, positioned right at the ProcName. + let name_pos = ALL_HEADERS_LEN_TX; + let len = u16::from_le_bytes([buf[name_pos], buf[name_pos + 1]]); + assert_eq!(len as usize, "dbo.usp_MyProc".encode_utf16().count()); + + // Verify the UTF-16 payload matches the proc name. + let mut chars = Vec::new(); + let mut off = name_pos + 2; + for _ in 0..len { + chars.push(u16::from_le_bytes([buf[off], buf[off + 1]])); + off += 2; + } + assert_eq!(String::from_utf16(&chars).unwrap(), "dbo.usp_MyProc"); + + // Option flags (u16) follow the name. + let flags = u16::from_le_bytes([buf[off], buf[off + 1]]); + assert_eq!(flags, 0); + } + + #[test] + fn named_and_by_id_differ_only_in_proc_slot() { + let by_id = { + let req = TokenRpcRequest::new(RpcProcId::ExecuteSQL, vec![], [0u8; 8]); + let mut buf = BytesMut::new(); + req.encode(&mut buf).unwrap(); + buf + }; + + // By-id encodes 0xFFFF followed by the proc id in the high word. + let val = u32::from_le_bytes([ + by_id[ALL_HEADERS_LEN_TX], + by_id[ALL_HEADERS_LEN_TX + 1], + by_id[ALL_HEADERS_LEN_TX + 2], + by_id[ALL_HEADERS_LEN_TX + 3], + ]); + assert_eq!(val & 0xffff, 0xffff); + assert_eq!((val >> 16) as u16, RpcProcId::ExecuteSQL as u16); + } + + #[test] + fn encodes_param_name_and_by_ref_flag() { + let param = RpcParam { + name: Cow::Borrowed("@out"), + flags: BitFlags::from_flag(RpcStatus::ByRefValue), + value: scalar(ColumnData::I32(Some(7))), + }; + + let mut buf = BytesMut::new(); + param.encode(&mut buf).unwrap(); + + // First byte is the param-name length (in UTF-16 code units). + assert_eq!(buf[0] as usize, "@out".encode_utf16().count()); + + let mut chars = Vec::new(); + let mut off = 1usize; + for _ in 0..buf[0] { + chars.push(u16::from_le_bytes([buf[off], buf[off + 1]])); + off += 2; + } + assert_eq!(String::from_utf16(&chars).unwrap(), "@out"); + + // Status flags byte carries the ByRefValue bit. + assert_eq!(buf[off], RpcStatus::ByRefValue as u8); + } +} diff --git a/src/tds/codec/type_info_tvp.rs b/src/tds/codec/type_info_tvp.rs new file mode 100644 index 0000000..a6d9b71 --- /dev/null +++ b/src/tds/codec/type_info_tvp.rs @@ -0,0 +1,206 @@ +use std::borrow::BorrowMut; + +use asynchronous_codec::BytesMut; +use bytes::BufMut; + +use crate::ColumnData; + +use super::{BytesMutWithTypeInfo, Encode, FixedLenType, MetaDataColumn, TypeInfo, VarLenContext}; + +const TVPTYPE: u8 = 0xF3; + +/// A table-valued parameter (TVP), as described by section 2.2.5.5.5 of MS-TDS. +/// +/// A TVP is passed to a stored procedure as an `RpcParam` and carries both the +/// column metadata (optionally resolved from the server) and the data rows. +#[derive(Debug)] +pub struct TypeInfoTvp<'a> { + schema_name: &'a str, + db_type_name: &'a str, + columns: Option>>, + data: Vec>>, +} + +impl<'a> Encode for TypeInfoTvp<'a> { + fn encode(self, dst: &mut BytesMut) -> crate::Result<()> { + // TVPTYPE = %xF3 + // TVP_TYPE_INFO = TVPTYPE + // TVP_TYPENAME + // TVP_COLMETADATA + // [TVP_ORDER_UNIQUE] + // [TVP_COLUMN_ORDERING] + // TVP_END_TOKEN + + dst.put_u8(TVPTYPE); + put_b_varchar("", dst); // DB name (unused) + put_b_varchar(self.schema_name, dst); + put_b_varchar(self.db_type_name, dst); + + if let Some(ref columns_metadata) = self.columns { + dst.put_u16_le(columns_metadata.len() as u16); + for col in columns_metadata { + // TvpColumnMetaData = UserType + // Flags + // TYPE_INFO + // ColName + dst.put_u32_le(0_u32); // UserType + col.base.clone().encode(dst)?; // Flags + TYPE_INFO + // 2.2.5.5.5.1: ColName MUST be a zero-length string in the TVP. + put_b_varchar("", dst); + } + } else { + // TVP_NULL_TOKEN: the server is expected to know the type. + dst.put_u16_le(0xFFFF_u16); + } + + dst.put_u8(0_u8); // end of TVP_COLMETADATA / optional metadata + + // TVP_ROW_TOKEN = %x01 ; A row as defined by TVP_COLMETADATA follows + // TvpColumnData = TYPE_VARBYTE ; Actual value must match column metadata + // AllColumnData = *TvpColumnData + // TVP_ROW = TVP_ROW_TOKEN AllColumnData + for row in self.data.into_iter() { + dst.put_u8(0x01u8); // TVP_ROW_TOKEN + for (i, col) in row.into_iter().enumerate() { + let mut dst_ti = BytesMutWithTypeInfo::new(dst); + if let Some(ref metadata) = self.columns { + dst_ti = dst_ti.with_type_info(&metadata[i].base.ty); + } + col.encode(&mut dst_ti)?; + } + } + + dst.put_u8(0_u8); // TVP_END_TOKEN + + Ok(()) + } +} + +/// Writes a `B_VARCHAR`: a single-byte code-unit count followed by that many +/// UTF-16 code units. +fn put_b_varchar>(s: T, dst: &mut BytesMut) { + let len_pos = dst.len(); + dst.put_u8(0u8); + let mut length = 0_u8; + + for chr in s.as_ref().encode_utf16() { + dst.put_u16_le(chr); + length += 1; + } + let dst: &mut [u8] = dst.borrow_mut(); + dst[len_pos] = length; +} + +impl<'a> TypeInfoTvp<'a> { + /// Creates a new TVP for the given database type name and data rows. The + /// type name may be qualified with a schema (e.g. `dbo.MyType`). + pub fn new(type_name: &'a str, rows: Vec>>) -> TypeInfoTvp<'a> { + let (schema_name, db_type_name) = if let Some((s, t)) = type_name.split_once('.') { + (s, t) + } else { + ("", type_name) + }; + TypeInfoTvp { + schema_name, + db_type_name, + columns: None, + data: rows, + } + } + + /// Attaches column metadata resolved from the server. Fixed-length column + /// types are rewritten to their nullable variable-length equivalents, as + /// required for TVP column metadata (2.2.5.5.5.3). + pub fn with_metadata(self, metadata: Vec>) -> TypeInfoTvp<'a> { + let mut metadata = metadata; + for mdc in metadata.iter_mut() { + let ty_replace = match mdc.base.ty { + TypeInfo::FixedLen(ref ty) => fixed_to_var_len(*ty), + _ => None, + }; + if let Some(ty) = ty_replace { + mdc.base.ty = ty; + } + } + TypeInfoTvp { + columns: Some(metadata), + ..self + } + } +} + +/// Maps a fixed-length column type to the nullable variable-length equivalent +/// that must be used in TVP column metadata. +fn fixed_to_var_len(ty: FixedLenType) -> Option { + use super::VarLenType; + + let ctx = match ty { + FixedLenType::Int1 => VarLenContext::new(VarLenType::Intn, 1, None), + FixedLenType::Bit => VarLenContext::new(VarLenType::Bitn, 1, None), + FixedLenType::Int2 => VarLenContext::new(VarLenType::Intn, 2, None), + FixedLenType::Int4 => VarLenContext::new(VarLenType::Intn, 4, None), + FixedLenType::Datetime4 => VarLenContext::new(VarLenType::Datetimen, 4, None), + FixedLenType::Float4 => VarLenContext::new(VarLenType::Floatn, 4, None), + FixedLenType::Money => VarLenContext::new(VarLenType::Money, 8, None), + FixedLenType::Datetime => VarLenContext::new(VarLenType::Datetimen, 8, None), + FixedLenType::Float8 => VarLenContext::new(VarLenType::Floatn, 8, None), + FixedLenType::Money4 => VarLenContext::new(VarLenType::Money, 4, None), + FixedLenType::Int8 => VarLenContext::new(VarLenType::Intn, 8, None), + _ => return None, + }; + Some(TypeInfo::VarLenSized(ctx)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn splits_schema_qualified_type_name() { + let tvp = TypeInfoTvp::new("dbo.MyType", Vec::new()); + assert_eq!(tvp.schema_name, "dbo"); + assert_eq!(tvp.db_type_name, "MyType"); + } + + #[test] + fn unqualified_type_name_has_empty_schema() { + let tvp = TypeInfoTvp::new("MyType", Vec::new()); + assert_eq!(tvp.schema_name, ""); + assert_eq!(tvp.db_type_name, "MyType"); + } + + #[test] + fn encodes_tvp_header_and_null_metadata() { + let tvp = TypeInfoTvp::new("dbo.MyType", Vec::new()); + let mut buf = BytesMut::new(); + tvp.encode(&mut buf).unwrap(); + + // TVPTYPE marker. + assert_eq!(buf[0], TVPTYPE); + // DB name is a zero-length B_VARCHAR. + assert_eq!(buf[1], 0); + // Schema name "dbo": length 3 followed by 3 UTF-16 code units. + assert_eq!(buf[2], 3); + + // With no metadata, the column count must be the TVP_NULL_TOKEN. + // Layout: F3, 00 (db), 03 + 6 bytes (schema), 06 + 12 bytes (type), + // then the u16 null token. + let null_token_pos = 1 + 1 + (1 + 6) + (1 + 12); + let token = u16::from_le_bytes([buf[null_token_pos], buf[null_token_pos + 1]]); + assert_eq!(token, 0xFFFF); + } + + #[test] + fn rewrites_fixed_len_to_nullable_var_len() { + use super::super::VarLenType; + assert!(matches!( + fixed_to_var_len(FixedLenType::Int4), + Some(TypeInfo::VarLenSized(ctx)) if ctx.r#type() == VarLenType::Intn + )); + assert!(matches!( + fixed_to_var_len(FixedLenType::Bit), + Some(TypeInfo::VarLenSized(ctx)) if ctx.r#type() == VarLenType::Bitn + )); + assert!(fixed_to_var_len(FixedLenType::Null).is_none()); + } +} diff --git a/src/tds/stream.rs b/src/tds/stream.rs index e645487..d01e6b9 100644 --- a/src/tds/stream.rs +++ b/src/tds/stream.rs @@ -1,5 +1,7 @@ +mod command; mod query; mod token; +pub use command::*; pub use query::*; pub use token::*; diff --git a/src/tds/stream/command.rs b/src/tds/stream/command.rs new file mode 100644 index 0000000..160b7cf --- /dev/null +++ b/src/tds/stream/command.rs @@ -0,0 +1,337 @@ +use crate::tds::stream::ReceivedToken; +use crate::{row::ColumnType, Column, Row}; +use crate::{ColumnData, CommandResult, ResultMetadata}; +use futures_util::{ + ready, + stream::{BoxStream, Peekable, Stream, StreamExt, TryStreamExt}, +}; +use std::{ + fmt::Debug, + pin::Pin, + sync::Arc, + task::{self, Poll}, +}; + +/// A `Stream` of [`CommandItem`] values produced by executing a [`Command`]. +/// +/// Items can be result metadata, rows, a return status, return values (OUT +/// parameters) or a rows-affected count. +/// +/// [`Command`]: crate::Command +/// +/// # Example +/// +/// ```no_run +/// # use std::env; +/// # use tiberius::Config; +/// # use tiberius::{Command, CommandItem}; +/// # use futures_util::TryStreamExt; +/// # use tokio_util::compat::TokioAsyncWriteCompatExt; +/// # #[tokio::main] +/// # async fn main() -> Result<(), Box> { +/// # let c_str = env::var("TIBERIUS_TEST_CONNECTION_STRING").unwrap_or( +/// # "server=tcp:localhost,1433;integratedSecurity=true;TrustServerCertificate=true".to_owned(), +/// # ); +/// # let config = Config::from_ado_string(&c_str)?; +/// # let tcp = tokio::net::TcpStream::connect(config.get_addr()).await?; +/// # tcp.set_nodelay(true)?; +/// # let mut client = tiberius::Client::connect(config, tcp.compat_write()).await?; +/// let mut cmd = Command::new("dbo.usp_SomeStoredProc"); +/// +/// cmd.bind_param("@foo", 34i32); +/// cmd.bind_param("@zoo", "the zoo string prm"); +/// cmd.bind_out_param("@bar", "bar"); +/// let mut stream = cmd.exec(&mut client).await?; +/// +/// while let Some(item) = stream.try_next().await? { +/// match item { +/// // our first item is the column data always +/// CommandItem::Metadata(meta) if meta.result_index() == 0 => { +/// // the first result column info can be handled here +/// } +/// // ... and from there on from 0..N rows +/// CommandItem::Row(row) if row.result_index() == 0 => { +/// let var: Option = row.get(0); +/// } +/// // the second result set returns first another metadata item +/// CommandItem::Metadata(meta) => { +/// // .. handling +/// } +/// // ...and, again, we get rows from the second resultset +/// CommandItem::Row(row) => { +/// let var: Option = row.get(0); +/// } +/// // check return status (returned always) +/// CommandItem::ReturnStatus(rs) => { +/// // .... do something +/// } +/// // collect OUT parameter values +/// CommandItem::ReturnValue(rv) => { +/// // .... do something, like push to a collection +/// } +/// // get affected row count +/// CommandItem::RowsAffected(ra) => { +/// // .... do something, like push to a collection +/// } +/// } +/// } +/// # Ok(()) +/// # } +/// ``` +/// +pub struct CommandStream<'a> { + token_stream: Peekable>>, + columns: Option>>, + result_set_index: Option, +} + +impl<'a> Debug for CommandStream<'a> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CommandStream") + .field( + "token_stream", + &"BoxStream<'a, crate::Result>", + ) + .finish() + } +} + +impl<'a> CommandStream<'a> { + pub(crate) fn new(token_stream: BoxStream<'a, crate::Result>) -> Self { + Self { + token_stream: token_stream.peekable(), + columns: None, + result_set_index: None, + } + } + + /// Collects all results from the command into memory, in the order they + /// were produced by the server. + pub async fn into_command_result(mut self) -> crate::Result { + let mut results: Vec> = Vec::new(); + let mut result: Option> = None; + let mut return_status = 0; + let mut return_values = Vec::new(); + let mut rows_affected = Vec::new(); + + while let Some(item) = self.try_next().await? { + match (item, &mut result) { + (CommandItem::Row(row), None) => { + result = Some(vec![row]); + } + (CommandItem::Row(row), Some(ref mut result)) => result.push(row), + (CommandItem::Metadata(_), None) => { + result = Some(Vec::new()); + } + (CommandItem::Metadata(_), ref mut previous_result) => { + results.push(previous_result.take().unwrap()); + result = None; + } + (CommandItem::ReturnStatus(rs), _) => return_status = rs, + (CommandItem::ReturnValue(rv), _) => return_values.push(rv), + (CommandItem::RowsAffected(rows), _) => rows_affected.push(rows), + } + } + + if let Some(result) = result { + results.push(result); + } + + Ok(CommandResult { + return_code: return_status, + return_values, + query_results: results, + rows_affected, + }) + } + + /// Converts the stream into a stream of rows, dropping all other items. + pub fn into_row_stream(self) -> BoxStream<'a, crate::Result> { + let s = self.try_filter_map(|item| async { + match item { + CommandItem::Row(row) => Ok(Some(row)), + _ => Ok(None), + } + }); + + Box::pin(s) + } +} + +/// A single OUT parameter value returned by a [`Command`]. +/// +/// [`Command`]: crate::Command +#[derive(Debug)] +pub struct CommandReturnValue { + pub(crate) name: String, + pub(crate) ord: u16, + pub(crate) data: ColumnData<'static>, +} + +impl CommandReturnValue { + /// The name of the OUT parameter this value corresponds to. + pub fn name(&self) -> &str { + &self.name + } + + /// The ordinal of the OUT parameter as returned by the server. + pub fn ordinal(&self) -> u16 { + self.ord + } + + /// A reference to the raw column data of the returned value. + pub fn data(&self) -> &ColumnData<'static> { + &self.data + } +} + +/// An item produced by a [`CommandStream`]. +#[derive(Debug)] +pub enum CommandItem { + /// A single row of data. + Row(Row), + /// Metadata describing the upcoming rows. + Metadata(ResultMetadata), + /// The return status from the server. + ReturnStatus(u32), + /// A return value, matching an OUT parameter. + ReturnValue(CommandReturnValue), + /// The number of rows affected by one of the statements ran on the server. + RowsAffected(u64), +} + +impl CommandItem { + pub(crate) fn metadata(columns: Arc>, result_index: usize) -> Self { + Self::Metadata(ResultMetadata { + columns, + result_index, + }) + } + + /// Returns a reference to the metadata, if the item is of a correct variant. + pub fn as_metadata(&self) -> Option<&ResultMetadata> { + match self { + CommandItem::Metadata(ref metadata) => Some(metadata), + _ => None, + } + } + + /// Returns a reference to the row, if the item is of a correct variant. + pub fn as_row(&self) -> Option<&Row> { + match self { + CommandItem::Row(ref row) => Some(row), + _ => None, + } + } + + /// Returns the metadata, if the item is of a correct variant. + pub fn into_metadata(self) -> Option { + match self { + CommandItem::Metadata(metadata) => Some(metadata), + _ => None, + } + } + + /// Returns the row, if the item is of a correct variant. + pub fn into_row(self) -> Option { + match self { + CommandItem::Row(row) => Some(row), + _ => None, + } + } + + /// Returns the return status, if the item is of a correct variant. + pub fn as_return_status(&self) -> Option { + match self { + CommandItem::ReturnStatus(rs) => Some(*rs), + _ => None, + } + } + + /// Returns a reference to the return value, if the item is of a correct variant. + pub fn as_return_value(&self) -> Option<&CommandReturnValue> { + match self { + CommandItem::ReturnValue(rv) => Some(rv), + _ => None, + } + } + + /// Returns the return value, if the item is of a correct variant. + pub fn into_return_value(self) -> Option { + match self { + CommandItem::ReturnValue(rv) => Some(rv), + _ => None, + } + } +} + +impl<'a> Stream for CommandStream<'a> { + type Item = crate::Result; + + fn poll_next(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll> { + let this = self.get_mut(); + + loop { + let token = match ready!(this.token_stream.poll_next_unpin(cx)) { + Some(res) => res?, + None => return Poll::Ready(None), + }; + + return match token { + ReceivedToken::NewResultset(meta) => { + let column_meta = meta + .columns + .iter() + .map(|x| Column { + name: x.col_name.to_string(), + column_type: ColumnType::from(&x.base.ty), + }) + .collect::>(); + + let column_meta = Arc::new(column_meta); + this.columns = Some(column_meta.clone()); + + this.result_set_index = this.result_set_index.map(|i| i + 1); + + let query_item = + CommandItem::metadata(column_meta, *this.result_set_index.get_or_insert(0)); + + Poll::Ready(Some(Ok(query_item))) + } + ReceivedToken::Row(data) => { + let columns = this.columns.as_ref().unwrap().clone(); + let result_index = this.result_set_index.unwrap(); + + let row = Row { + columns, + data, + result_index, + }; + + Poll::Ready(Some(Ok(CommandItem::Row(row)))) + } + ReceivedToken::ReturnStatus(rs) => { + Poll::Ready(Some(Ok(CommandItem::ReturnStatus(rs)))) + } + ReceivedToken::ReturnValue(rv) => { + Poll::Ready(Some(Ok(CommandItem::ReturnValue(CommandReturnValue { + name: rv.param_name, + ord: rv.param_ordinal, + data: rv.value, + })))) + } + ReceivedToken::DoneProc(done) if done.is_final() => continue, + ReceivedToken::DoneProc(done) => { + Poll::Ready(Some(Ok(CommandItem::RowsAffected(done.rows())))) + } + ReceivedToken::DoneInProc(done) => { + Poll::Ready(Some(Ok(CommandItem::RowsAffected(done.rows())))) + } + ReceivedToken::Done(done) => { + Poll::Ready(Some(Ok(CommandItem::RowsAffected(done.rows())))) + } + _ => continue, + }; + } + } +} diff --git a/src/tds/stream/query.rs b/src/tds/stream/query.rs index 0dc6947..59eef42 100644 --- a/src/tds/stream/query.rs +++ b/src/tds/stream/query.rs @@ -280,8 +280,8 @@ impl<'a> QueryStream<'a> { /// Info about the following stream of rows. #[derive(Debug, Clone)] pub struct ResultMetadata { - columns: Arc>, - result_index: usize, + pub(crate) columns: Arc>, + pub(crate) result_index: usize, } impl ResultMetadata { diff --git a/tests/command.rs b/tests/command.rs new file mode 100644 index 0000000..29ea0f1 --- /dev/null +++ b/tests/command.rs @@ -0,0 +1,229 @@ +use futures_util::io::{AsyncRead, AsyncWrite}; +// use futures_util::stream::TryStreamExt; +use names::{Generator, Name}; +use once_cell::sync::Lazy; +use std::cell::RefCell; +use std::env; +use std::sync::Once; + +use tiberius::{numeric::Numeric, Command, Result, TableValueRow}; + +use runtimes_macro::test_on_runtimes; + +// This is used in the testing macro :) +#[allow(dead_code)] +static LOGGER_SETUP: Once = Once::new(); + +static CONN_STR: Lazy = Lazy::new(|| { + env::var("TIBERIUS_TEST_CONNECTION_STRING").unwrap_or_else(|_| { + "server=tcp:localhost,1433;user=SA;password=;IntegratedSecurity=true;TrustServerCertificate=true".to_owned() + }) +}); + +thread_local! { + static NAMES: RefCell>> = + RefCell::new(None); +} + +async fn random_table() -> String { + NAMES.with(|maybe_generator| { + maybe_generator + .borrow_mut() + .get_or_insert_with(|| Generator::with_naming(Name::Plain)) + .next() + .unwrap() + .replace('-', "") + }) +} + +#[test_on_runtimes] +async fn basic_proc_exec(mut conn: tiberius::Client) -> Result<()> +where + S: AsyncRead + AsyncWrite + Unpin + Send, +{ + let table = random_table().await; + let proc = random_table().await; + + conn.simple_query(format!( + r#" + create table ##{} ( + id int identity(1,1), + other varchar(50), + ) + "#, + table + )) + .await?; + + conn.simple_query(format!( + r#" + create or alter procedure {} + @Param1 varchar(50) + as + insert into ##{} (other) + values (@Param1) + + return scope_identity() + "#, + proc, table, + )) + .await?; + + let mut ins_cmd = Command::new(&proc); + ins_cmd.bind_param("@Param1", "some text"); + + let result = ins_cmd.exec(&mut conn).await?.into_command_result().await?; + assert_eq!(1, result.return_code()); + + let mut ins_cmd = Command::new(&proc); + ins_cmd.bind_param("@Param1", "another text"); + + let result = ins_cmd.exec(&mut conn).await?.into_command_result().await?; + assert_eq!(2, result.return_code()); + + Ok(()) +} + +struct GeoTest { + id: i32, + lat: Numeric, + lon: Numeric, +} + +impl<'a> TableValueRow<'a> for GeoTest { + fn bind_fields(&self, data_row: &mut tiberius::SqlTableDataRow<'a>) { + data_row.add_field(self.id); + data_row.add_field(self.lat); + data_row.add_field(self.lon); + } + + fn get_db_type() -> &'static str { + "GeoTest" + } +} + +#[test_on_runtimes] +async fn tvp_proc_exec(mut conn: tiberius::Client) -> Result<()> +where + S: AsyncRead + AsyncWrite + Unpin + Send, +{ + let table = random_table().await; + let proc_tvp = random_table().await; + let proc_ins = random_table().await; + let proc_get = random_table().await; + let db_type = random_table().await; + + conn.simple_query(format!( + r#" + create type dbo.[{}] as table + ( + [ID] int not null, + [lat] decimal(9,6), + [lon] decimal(9,6) + ) + "#, + db_type + )) + .await?; + + conn.simple_query(format!( + r#" + create table ##{} ( + id int not null, + lat decimal(9,6), + lon decimal(9,6), + n varchar(50) + ) + "#, + table + )) + .await?; + + conn.simple_query(format!( + r#" + create or alter procedure {} + @id int, + @geo dbo.[{}] readonly + as + update t set + [lat] = g.[lat], + [lon] = g.[lon] + from + ##{} t + inner join @geo g on g.id = t.id + + "#, + proc_tvp, db_type, table, + )) + .await?; + + conn.simple_query(format!( + r#" + create or alter procedure {} + @id int, + @name varchar(50) + as + insert into ##{} (id, n) + values (@id, @name) + + "#, + proc_ins, table, + )) + .await?; + + conn.simple_query(format!( + r#" + create or alter procedure {} + @id int, + @count int out + as + set @count = (select count(*) from ##{}) + select * from ##{} + "#, + proc_get, table, table + )) + .await?; + + let mut ins_cmd = Command::new(&proc_ins); + ins_cmd.bind_param("@id", 23); + ins_cmd.bind_param("@name", "the twenty three"); + + let result = ins_cmd.exec(&mut conn).await?.into_command_result().await?; + assert_eq!(0, result.return_code()); + + let g1 = GeoTest { + id: 23, + lon: Numeric::new_with_scale(141, 6), + lat: Numeric::new_with_scale(192, 6), + }; + let g2 = GeoTest { + id: 78, + lon: Numeric::new_with_scale(1141, 6), + lat: Numeric::new_with_scale(8192, 6), + }; + let tbl = vec![g1, g2]; + let mut tvp_cmd = Command::new(&proc_tvp); + tvp_cmd.bind_param("@id", 23); + tvp_cmd.bind_table_with_dbtype("@geo", &db_type, tbl); + + let result = tvp_cmd.exec(&mut conn).await?.into_command_result().await?; + assert_eq!(0, result.return_code()); + + let count = 0; + let mut get_cmd = Command::new(&proc_get); + get_cmd.bind_param("@id", 23); + get_cmd.bind_out_param("@count", count); + + let result = get_cmd.exec(&mut conn).await?.into_command_result().await?; + assert_eq!(0, result.return_code()); + let count: i32 = result.try_return_value("@count")?.unwrap(); + assert_eq!(1, count); + + let rows = result.to_query_result(0).unwrap(); + let lat: Numeric = rows[0].get("lat").unwrap(); + let lon: Numeric = rows[0].get("lon").unwrap(); + assert_eq!(Numeric::new_with_scale(141, 6), lon); + assert_eq!(Numeric::new_with_scale(192, 6), lat); + + Ok(()) +} diff --git a/tvp-macro/Cargo.toml b/tvp-macro/Cargo.toml new file mode 100644 index 0000000..117d5e2 --- /dev/null +++ b/tvp-macro/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "tvp-macro" +version = "0.1.0" +edition = "2021" +license = "MIT/Apache-2.0" +description = "Derive macro for tiberius table-valued parameter rows." +repository = "https://github.com/prisma/tiberius" + +[dependencies] +syn = "^2" +quote = "^1" +proc-macro2 = "^1" + +[lib] +proc-macro = true diff --git a/tvp-macro/src/attr.rs b/tvp-macro/src/attr.rs new file mode 100644 index 0000000..f038be6 --- /dev/null +++ b/tvp-macro/src/attr.rs @@ -0,0 +1,45 @@ +pub(crate) struct FieldAttr { + pub colname: Option, +} + +impl FieldAttr { + pub(crate) fn parse(attrs: &[syn::Attribute]) -> Option { + let mut result = None; + for attr in attrs.iter() { + match attr.style { + syn::AttrStyle::Outer => {} + _ => continue, + } + let last_attr_path = attr + .path() + .segments + .last() + .expect("Expected at least one segment where #[segment[::segment*](..)]"); + if last_attr_path.ident != "colname" { + continue; + } + let kv = match attr.meta { + syn::Meta::NameValue(ref kv) => kv, + _ if attr.path().is_ident("colname") => { + panic!("Invalid #[colname] attribute, expected #[colname = \"SomeColName\"]") + } + _ => continue, + }; + if result.is_some() { + panic!("Expected at most one #[colname] attribute"); + } + if let syn::Expr::Lit(syn::ExprLit { + lit: syn::Lit::Str(ref s), + .. + }) = kv.value + { + result = Some(FieldAttr { + colname: Some(s.value()), + }); + } else { + panic!("Non-string literal value in #[colname] attribute"); + } + } + result + } +} diff --git a/tvp-macro/src/lib.rs b/tvp-macro/src/lib.rs new file mode 100644 index 0000000..f053033 --- /dev/null +++ b/tvp-macro/src/lib.rs @@ -0,0 +1,46 @@ +//! A utility proc-macro crate that generates trivial trait implementations used +//! in the Rust-to-SQL data exchange for tiberius table-valued parameters. +extern crate proc_macro; + +#[macro_use] +extern crate quote; +#[macro_use] +extern crate syn; + +use proc_macro::TokenStream; + +macro_rules! sp_quote { + ($($t:tt)*) => (quote_spanned!(proc_macro2::Span::call_site() => $($t)*)) +} + +mod attr; +mod table_value_param; + +/// Generates a trivial implementation of the `TableValueRow` trait. +/// +/// # Applications +/// Apply to structs that represent rows of a table-valued parameter. +/// +/// # Example +/// ```rust,ignore +/// # use tiberius::*; +/// #[derive(TableValueRow)] +/// pub struct SomeGeoList { +/// #[colname = "SomeID"] +/// pub id: i32, +/// #[colname = "LastSyncIPGeoLat"] +/// pub lat: Numeric, +/// #[colname = "LastSyncIPGeoLong"] +/// pub lon: Numeric, +/// } +/// ``` +#[proc_macro_derive(TableValueRow, attributes(colname))] +pub fn table_value_param(input: TokenStream) -> TokenStream { + let ast: syn::DeriveInput = syn::parse(input).expect("Couldn't parse item"); + let result = match ast.data { + syn::Data::Enum(_) => panic!("n/a for enums, makes sense for structs only"), + syn::Data::Struct(ref s) => table_value_param::for_struct(&ast, &s.fields), + syn::Data::Union(_) => panic!("doesn't work with unions"), + }; + result.into() +} diff --git a/tvp-macro/src/table_value_param.rs b/tvp-macro/src/table_value_param.rs new file mode 100644 index 0000000..4bc9ab2 --- /dev/null +++ b/tvp-macro/src/table_value_param.rs @@ -0,0 +1,167 @@ +use proc_macro2::TokenStream; +use syn::punctuated::Punctuated; + +use crate::attr::FieldAttr; + +pub(crate) fn for_struct(ast: &syn::DeriveInput, fields: &syn::Fields) -> TokenStream { + match *fields { + syn::Fields::Named(ref fields) => table_value_param_impl(ast, Some(&fields.named)), + _ => panic!("Only named fields are supported so far"), + } +} + +fn table_value_param_impl( + ast: &syn::DeriveInput, + fields: Option<&Punctuated>, +) -> TokenStream { + let name = &ast.ident; + let (lt_impl, lt_struct) = { + let mut lifetimes: Vec<&syn::Ident> = Vec::new(); + for gp in ast.generics.params.iter() { + if let syn::GenericParam::Lifetime(ltp) = gp { + lifetimes.push(<p.lifetime.ident); + } + } + if lifetimes.len() > 1 { + panic!( + "Only one lifetime specifier is supported for the structure. Found multiple: {}", + lifetimes + .iter() + .map(|lt| lt.to_string()) + .collect::>() + .join(", ") + ); + } + if lifetimes.is_empty() { + (sp_quote!(<'query>), sp_quote!()) + } else { + let lt = lifetimes[0]; + let ts: proc_macro2::TokenStream = format!("< '{} >", lt).parse().unwrap(); + (ts.clone(), ts) + } + }; + let empty = Default::default(); + let fields: Vec<_> = fields.unwrap_or(&empty).iter().map(FieldExt::new).collect(); + let col_names: Vec<_> = fields.iter().map(|f| f.get_col_name()).collect(); + // Column names are not needed for stored procedures, but will be required + // once TVPs are supported for ad-hoc queries. + let _col_names = sp_quote!( #(#col_names),* ); + let col_binds: Vec<_> = fields.iter().map(|f| f.as_bind()).collect(); + let col_binds = sp_quote!( #(#col_binds);*); + sp_quote! { + impl #lt_impl tiberius::TableValueRow #lt_impl for #name #lt_struct { + fn get_db_type() -> &'static str { + stringify!{ #name } + } + + fn bind_fields(&self, data_row: &mut tiberius::SqlTableDataRow #lt_impl) { + #col_binds; + } + } + } +} + +struct FieldExt { + attr: Option, + ident: syn::Ident, +} + +impl FieldExt { + pub fn new(field: &syn::Field) -> FieldExt { + if let Some(ident) = field.ident.clone() { + FieldExt { + attr: FieldAttr::parse(&field.attrs), + ident, + } + } else { + panic!("Field ident is required"); + } + } + pub(crate) fn get_col_name(&self) -> String { + if let Some(attr) = self.attr.as_ref() { + if let Some(colname) = attr.colname.as_ref() { + return colname.to_string(); + } + } + self.ident.to_string() + } + pub(crate) fn as_bind(&self) -> TokenStream { + let name = &self.ident; + sp_quote!(data_row.add_field(self.#name)) + } +} + +#[cfg(test)] +mod tests { + use super::for_struct; + + #[test] + fn basic_nolifetime() { + let ast: syn::DeriveInput = syn::parse_str( + r#" + pub struct SomeGeoList { + #[colname = "SomeID"] + pub id: i32, + #[colname = "LastSyncIPGeoLat"] + pub lat: Numeric, // decimal(9,6) + #[colname = "LastSyncIPGeoLong"] + pub lon: Numeric, // decimal(9,6) + } + "#, + ) + .unwrap(); + let result = match ast.data { + syn::Data::Enum(_) => panic!("n/a for enums, makes sense for structs only"), + syn::Data::Struct(ref s) => for_struct(&ast, &s.fields), + syn::Data::Union(_) => panic!("doesn't work with unions"), + }; + let etalon = sp_quote!( + impl<'query> tiberius::TableValueRow<'query> for SomeGeoList { + fn get_db_type() -> &'static str { + stringify! { SomeGeoList } + } + fn bind_fields(&self, data_row: &mut tiberius::SqlTableDataRow<'query>) { + data_row.add_field(self.id); + data_row.add_field(self.lat); + data_row.add_field(self.lon); + } + } + ); + + assert_eq!(result.to_string(), etalon.to_string()); + } + + #[test] + fn basic_lifetime() { + let ast: syn::DeriveInput = syn::parse_str( + r#" + pub struct AnotherGeoList<'e> { + #[colname = "SomeID"] + pub id: i32, + #[colname = "SomeStr"] + pub s: &'e str, + } + "#, + ) + .unwrap(); + let result = match ast.data { + syn::Data::Enum(_) => panic!("n/a for enums, makes sense for structs only"), + syn::Data::Struct(ref s) => for_struct(&ast, &s.fields), + syn::Data::Union(_) => panic!("doesn't work with unions"), + }; + + let etalon = sp_quote!( + impl<'e> tiberius::TableValueRow<'e> for AnotherGeoList<'e> { + fn get_db_type() -> &'static str { + stringify! { AnotherGeoList } + } + fn bind_fields(&self, data_row: &mut tiberius::SqlTableDataRow<'e>) { + data_row.add_field(self.id); + data_row.add_field(self.s); + } + } + ); + + assert_eq!(result.to_string(), etalon.to_string()); + } +}